Agent-Native/runtime/tools/tools_test.go

package tools

import (
	"bytes"
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestSchemaRejectsMissingAndWrongType(t *testing.T) {
	dir := t.TempDir()
	reg := WorkspaceTools(dir)
	missing := reg.Execute(context.Background(), "read_file", []byte(`{}`))
	if !missing.IsError || !strings.Contains(missing.Content, "path") {
		t.Fatalf("missing %#v", missing)
	}
	wrong := reg.Execute(context.Background(), "read_file", []byte(`{"path":1}`))
	if !wrong.IsError || !strings.Contains(wrong.Content, "string") {
		t.Fatalf("type %#v", wrong)
	}
}

func TestPathEscapeAndSymlink(t *testing.T) {
	dir := t.TempDir()
	outside := t.TempDir()
	secret := filepath.Join(outside, "secret.txt")
	if err := os.WriteFile(secret, []byte("hidden"), 0o644); err != nil {
		t.Fatal(err)
	}
	link := filepath.Join(dir, "leak")
	if err := os.Symlink(secret, link); err != nil {
		t.Fatal(err)
	}
	reg := WorkspaceTools(dir)
	up := reg.Execute(context.Background(), "read_file", []byte(`{"path":"../secret.txt"}`))
	if !up.IsError {
		t.Fatalf("parent escape succeeded: %s", up.Content)
	}
	via := reg.Execute(context.Background(), "read_file", []byte(`{"path":"leak"}`))
	if !via.IsError || strings.Contains(via.Content, "hidden") {
		t.Fatalf("symlink escape: %#v", via)
	}
	absOutside := reg.Execute(context.Background(), "write_file", []byte(`{"path":"`+secret+`","content":"no"}`))
	if !absOutside.IsError {
		t.Fatal("absolute outside write succeeded")
	}
}

func TestWriteReadEditGrep(t *testing.T) {
	dir := t.TempDir()
	reg := WorkspaceTools(dir)
	wrote := reg.Execute(context.Background(), "write_file", []byte(`{"path":"docs/a.md","content":"alpha beta"}`))
	if wrote.IsError {
		t.Fatal(wrote.Content)
	}
	got := reg.Execute(context.Background(), "read_file", []byte(`{"path":"docs/a.md"}`))
	if got.IsError || got.Content != "alpha beta" {
		t.Fatalf("read %#v", got)
	}
	edited := reg.Execute(context.Background(), "edit_file", []byte(`{"path":"docs/a.md","old":"beta","new":"gamma"}`))
	if edited.IsError {
		t.Fatal(edited.Content)
	}
	found := reg.Execute(context.Background(), "grep", []byte(`{"pattern":"gamma"}`))
	if found.IsError || !strings.Contains(found.Content, "docs/a.md:1:") {
		t.Fatalf("grep %#v", found)
	}
	listed := reg.Execute(context.Background(), "list_dir", []byte(`{"path":"docs"}`))
	if !strings.Contains(listed.Content, "a.md") {
		t.Fatalf("list %#v", listed)
	}
}

func TestBeforeHookBlocksAndAudits(t *testing.T) {
	dir := t.TempDir()
	reg := WorkspaceTools(dir)
	var audit bytes.Buffer
	reg.Audit = &audit
	reg.Before = func(ctx context.Context, name string, args json.RawMessage) (bool, string) {
		return true, "policy"
	}
	res := reg.Execute(context.Background(), "list_dir", []byte(`{}`))
	if !res.IsError || res.Content != "policy" {
		t.Fatalf("hook %#v", res)
	}
	if !strings.Contains(audit.String(), "blocked=true") {
		t.Fatalf("audit %s", audit.String())
	}
}

type slowHookTool struct{}

func (slowHookTool) Name() string                { return "slow_hook" }
func (slowHookTool) Description() string         { return "noop" }
func (slowHookTool) Parameters() json.RawMessage { return json.RawMessage(`{"type":"object","properties":{}}`) }
func (slowHookTool) Execute(context.Context, json.RawMessage) (Result, error) {
	return Result{Content: "ok"}, nil
}

func TestHooksDoNotRunConcurrently(t *testing.T) {
	reg := NewRegistry(t.TempDir())
	reg.Register(slowHookTool{})
	var mu sync.Mutex
	inFlight, maxInFlight := 0, 0
	reg.Before = func(context.Context, string, json.RawMessage) (bool, string) {
		mu.Lock()
		inFlight++
		if inFlight > maxInFlight {
			maxInFlight = inFlight
		}
		mu.Unlock()
		time.Sleep(20 * time.Millisecond)
		mu.Lock()
		inFlight--
		mu.Unlock()
		return false, ""
	}
	var wg sync.WaitGroup
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			reg.Execute(context.Background(), "slow_hook", []byte(`{}`))
		}()
	}
	wg.Wait()
	if maxInFlight != 1 {
		t.Fatalf("hooks overlapped: max in flight %d", maxInFlight)
	}
}

func TestAfterHookRewrites(t *testing.T) {
	dir := t.TempDir()
	reg := WorkspaceTools(dir)
	reg.After = func(ctx context.Context, name string, args json.RawMessage, result Result) Result {
		result.Content = "rewritten"
		return result
	}
	res := reg.Execute(context.Background(), "list_dir", []byte(`{}`))
	if res.Content != "rewritten" {
		t.Fatalf("%#v", res)
	}
}