package tools
import (
"bytes"
"context"
"encoding/json"
"errors"
"strings"
"testing"
"harness/internal/metrics"
)
type echoArgs struct {
Text string `json:"text" jsonschema:"description=Text to echo back"`
Times int `json:"times,omitempty" jsonschema:"description=How many times to repeat it"`
}
func echoTool(fail bool) Tool {
return New("echo", "Echo text back", func(_ context.Context, in echoArgs) (string, error) {
if fail {
return "", errors.New("echo refused")
}
n := in.Times
if n == 0 {
n = 1
}
return strings.Repeat(in.Text, n), nil
})
}
func TestSchemaForReflectsStructTags(t *testing.T) {
var got struct {
Type string `json:"type"`
Required []string `json:"required"`
Properties map[string]struct {
Type string `json:"type"`
Description string `json:"description"`
} `json:"properties"`
AdditionalProperties *bool `json:"additionalProperties"`
Schema json.RawMessage `json:"$schema"`
}
if err := json.Unmarshal(SchemaFor[echoArgs](), &got); err != nil {
t.Fatalf("schema is not valid JSON: %v", err)
}
if got.Type != "object" {
t.Errorf("type = %q, want object", got.Type)
}
if len(got.Required) != 1 || got.Required[0] != "text" {
t.Errorf("required = %v, want [text] (omitempty fields must be optional)", got.Required)
}
if got.Properties["times"].Type != "integer" {
t.Errorf("times type = %q, want integer", got.Properties["times"].Type)
}
if got.Properties["text"].Description != "Text to echo back" {
t.Errorf("description not carried over: %q", got.Properties["text"].Description)
}
if got.Schema != nil {
t.Errorf("$schema should be stripped, got %s", got.Schema)
}
}
func TestSchemaForNoFields(t *testing.T) {
type empty struct{}
var got map[string]any
if err := json.Unmarshal(SchemaFor[empty](), &got); err != nil {
t.Fatalf("invalid JSON: %v", err)
}
if got["type"] != "object" {
t.Errorf("type = %v, want object", got["type"])
}
if _, ok := got["properties"]; !ok {
t.Error("properties must be present even when empty")
}
}
func TestTypedExecute(t *testing.T) {
tests := []struct {
name string
raw string
want string
wantErr string
}{
{name: "decodes typed args", raw: `{"text":"ab","times":3}`, want: "ababab"},
{name: "omitted optional field", raw: `{"text":"x"}`, want: "x"},
{name: "empty input treated as {}", raw: ``, want: ""},
{name: "malformed JSON", raw: `{"text":`, wantErr: "invalid arguments"},
{name: "wrong field type", raw: `{"text":"x","times":"three"}`, wantErr: "invalid arguments"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
got, err := echoTool(false).Execute(context.Background(), json.RawMessage(tc.raw))
switch {
case tc.wantErr != "" && err == nil:
t.Fatalf("want error containing %q, got none (result %q)", tc.wantErr, got)
case tc.wantErr != "":
if !strings.Contains(err.Error(), tc.wantErr) {
t.Fatalf("error = %v, want it to contain %q", err, tc.wantErr)
}
case err != nil:
t.Fatalf("unexpected error: %v", err)
case got != tc.want:
t.Fatalf("result = %q, want %q", got, tc.want)
}
})
}
}
func TestRegistryDispatch(t *testing.T) {
r := NewRegistry(nil).Register(echoTool(false))
got, err := r.Execute(context.Background(), "echo", json.RawMessage(`{"text":"hi"}`))
if err != nil || got != "hi" {
t.Fatalf("Execute = (%q, %v), want (\"hi\", nil)", got, err)
}
if _, err := r.Execute(context.Background(), "nope", nil); err == nil {
t.Error("unknown tool must return an error the model can read")
} else if !strings.Contains(err.Error(), "nope") {
t.Errorf("error should name the tool, got %v", err)
}
}
func TestRegistryToolErrorIsReturnedNotSwallowed(t *testing.T) {
r := NewRegistry(nil).Register(echoTool(true))
out, err := r.Execute(context.Background(), "echo", json.RawMessage(`{"text":"hi"}`))
if err == nil {
t.Fatal("want the tool's error to surface")
}
if out != "" {
t.Errorf("output = %q, want empty on failure", out)
}
if err.Error() != "echo refused" {
t.Errorf("error = %q, want it passed through verbatim", err)
}
}
func TestRegistryListIsStable(t *testing.T) {
r := NewRegistry(nil).Register(
New("zulu", "", func(context.Context, echoArgs) (string, error) { return "", nil }),
New("alpha", "", func(context.Context, echoArgs) (string, error) { return "", nil }),
)
var names []string
for _, tl := range r.List() {
names = append(names, tl.Name())
}
if len(names) != 2 || names[0] != "alpha" || names[1] != "zulu" {
t.Errorf("List = %v, want sorted [alpha zulu] for a byte-stable prompt prefix", names)
}
}
func TestRegistryDuplicatePanics(t *testing.T) {
defer func() {
if recover() == nil {
t.Error("registering the same name twice must panic at wiring time")
}
}()
NewRegistry(nil).Register(echoTool(false), echoTool(false))
}
func TestExecuteRespectsContext(t *testing.T) {
blocked := New("block", "", func(ctx context.Context, _ echoArgs) (string, error) {
<-ctx.Done()
return "", ctx.Err()
})
ctx, cancel := context.WithCancel(context.Background())
cancel()
if _, err := NewRegistry(nil).Register(blocked).Execute(ctx, "block", nil); !errors.Is(err, context.Canceled) {
t.Errorf("error = %v, want context.Canceled", err)
}
}
func TestRegistryLogsEveryDispatch(t *testing.T) {
var buf bytes.Buffer
r := NewRegistry(metrics.NewWriter(&buf, "sess")).Register(echoTool(false))
r.Execute(context.Background(), "echo", json.RawMessage(`{"text":"hi"}`))
r.Execute(context.Background(), "missing", nil)
lines := strings.Split(strings.TrimRight(buf.String(), "\n"), "\n")
if len(lines) != 2 {
t.Fatalf("got %d log lines, want one per dispatch:\n%s", len(lines), buf.String())
}
var ok struct {
Tool string `json:"tool"`
OK bool `json:"ok"`
Out int `json:"output_bytes"`
}
json.Unmarshal([]byte(lines[0]), &ok)
if ok.Tool != "echo" || !ok.OK || ok.Out != 2 {
t.Errorf("success event = %+v, want echo/ok/2 bytes", ok)
}
if !strings.Contains(lines[1], `"ok":false`) {
t.Errorf("unknown-tool dispatch should log a failure, got %s", lines[1])
}
}