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]) } }