diff --git a/internal/agent/deferred_loop_test.go b/internal/agent/deferred_loop_test.go new file mode 100644 index 00000000..174178ab --- /dev/null +++ b/internal/agent/deferred_loop_test.go @@ -0,0 +1,768 @@ +package agent + +import ( + "context" + "errors" + "reflect" + "sort" + "strings" + "testing" + + "github.com/Gitlawb/zero/internal/tools" + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +func TestOptionsDeferThresholdFieldExists(t *testing.T) { + options := Options{DeferThreshold: 10} + if options.DeferThreshold != 10 { + t.Fatalf("expected DeferThreshold 10, got %d", options.DeferThreshold) + } +} + +func TestToolResultLoadedToolsField(t *testing.T) { + result := ToolResult{LoadedTools: []string{"Alpha", "Beta"}} + if len(result.LoadedTools) != 2 || result.LoadedTools[0] != "Alpha" || result.LoadedTools[1] != "Beta" { + t.Fatalf("expected LoadedTools [Alpha Beta], got %#v", result.LoadedTools) + } + // Default zero value is nil for an ordinary result. + if (ToolResult{}).LoadedTools != nil { + t.Fatalf("expected nil LoadedTools by default") + } +} + +// loadSignalTool returns Meta["load_tools"] like tool_search does, so we can +// assert executeToolCall lifts it into ToolResult.LoadedTools. +type loadSignalTool struct{ value string } + +func (t loadSignalTool) Name() string { return "load_signal" } +func (t loadSignalTool) Description() string { return "emits a load_tools signal" } +func (t loadSignalTool) Parameters() tools.Schema { + return tools.Schema{Type: "object", AdditionalProperties: false} +} +func (t loadSignalTool) Safety() tools.Safety { + return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow} +} +func (t loadSignalTool) Run(_ context.Context, _ map[string]any) tools.Result { + return tools.Result{Status: tools.StatusOK, Output: "ok"} +} +func (t loadSignalTool) RunWithOptions(_ context.Context, _ map[string]any, _ tools.RunOptions) tools.Result { + return tools.Result{Status: tools.StatusOK, Output: "ok", Meta: map[string]string{"load_tools": t.value}} +} + +func TestExecuteToolCallLiftsLoadTools(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(loadSignalTool{value: " Alpha , Beta ,, "}) + + result, abortErr := executeToolCall( + context.Background(), + registry, + ToolCall{ID: "c1", Name: "load_signal", Arguments: ""}, + PermissionModeAuto, + Options{}, + ) + if abortErr != nil { + t.Fatalf("unexpected abort error: %v", abortErr) + } + want := []string{"Alpha", "Beta"} + if len(result.LoadedTools) != len(want) { + t.Fatalf("expected LoadedTools %#v, got %#v", want, result.LoadedTools) + } + for i := range want { + if result.LoadedTools[i] != want[i] { + t.Fatalf("expected LoadedTools %#v, got %#v", want, result.LoadedTools) + } + } +} + +func TestExecuteToolCallNoLoadToolsMetaLeavesNil(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(secretEmittingTool{output: "plain"}) + + result, abortErr := executeToolCall( + context.Background(), + registry, + ToolCall{ID: "c1", Name: "leak", Arguments: ""}, + PermissionModeAuto, + Options{}, + ) + if abortErr != nil { + t.Fatalf("unexpected abort error: %v", abortErr) + } + if result.LoadedTools != nil { + t.Fatalf("expected nil LoadedTools for a tool with no load_tools meta, got %#v", result.LoadedTools) + } +} + +// fakeDeferredTool is deferred-eligible (implements Deferred() bool) like an MCP +// tool wrapper, so partitionTools counts and (when active) hides it. +type fakeDeferredTool struct { + name string + desc string +} + +func (t fakeDeferredTool) Name() string { return t.name } +func (t fakeDeferredTool) Description() string { return t.desc } +func (t fakeDeferredTool) Parameters() tools.Schema { + return tools.Schema{Type: "object", AdditionalProperties: false} +} +func (t fakeDeferredTool) Safety() tools.Safety { + return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow} +} +func (t fakeDeferredTool) Run(_ context.Context, _ map[string]any) tools.Result { + return tools.Result{Status: tools.StatusOK, Output: "ok"} +} +func (t fakeDeferredTool) Deferred() bool { return true } + +// fakeToolSearchTool stands in for component D's tool_search (a non-deferred +// builtin) so the inactive path can assert it is dropped. +type fakeToolSearchTool struct{} + +func (fakeToolSearchTool) Name() string { return "tool_search" } +func (fakeToolSearchTool) Description() string { return "load deferred tool schemas" } +func (fakeToolSearchTool) Parameters() tools.Schema { + return tools.Schema{Type: "object", AdditionalProperties: false} +} +func (fakeToolSearchTool) Safety() tools.Safety { + return tools.Safety{SideEffect: tools.SideEffectNone, Permission: tools.PermissionAllow, AdvertiseInAuto: true} +} +func (fakeToolSearchTool) Run(_ context.Context, _ map[string]any) tools.Result { + return tools.Result{Status: tools.StatusOK, Output: "ok"} +} + +func TestPartitionToolsInactiveIsByteIdenticalAndDropsToolSearch(t *testing.T) { + root := t.TempDir() + registry := tools.NewRegistry() + registry.Register(tools.NewReadFileTool(root)) + registry.Register(fakeDeferredTool{name: "mcp__srv__a", desc: "tool a"}) + registry.Register(fakeToolSearchTool{}) + + options := Options{DeferThreshold: 0} // 0 => deferral disabled => inactive path. + + // DeferThreshold 0 => deferral disabled => inactive path. + exposed, reminder := partitionTools(registry, PermissionModeAuto, options, map[string]bool{}) + + if reminder != "" { + t.Fatalf("expected empty reminder on inactive path, got %q", reminder) + } + + // Strong byte-identity assertion: build the reference the LEGACY way (the old + // toolDefinitions construction — every visible/advertised tool's full schema, + // EXCEPT tool_search, alpha-sorted by name) and require an exact DeepEqual. This + // pins that the inactive partition produces the pre-deferral output verbatim, + // not merely the same set of names. + reference := legacyToolDefinitions(registry, PermissionModeAuto, options) + if !reflect.DeepEqual(exposed, reference) { + t.Fatalf("inactive partition not byte-identical to legacy toolDefinitions:\n got %#v\nwant %#v", exposed, reference) + } + + // Belt-and-suspenders: tool_search is dropped, the deferred tool keeps its full + // schema, and only the expected names are present. + for _, def := range exposed { + if def.Name == "tool_search" { + t.Fatalf("tool_search must be dropped on inactive path, got %#v", exposed) + } + } + wantNames := map[string]bool{"read_file": true, "mcp__srv__a": true} + if len(exposed) != len(wantNames) { + t.Fatalf("expected %d exposed tools, got %d: %#v", len(wantNames), len(exposed), exposed) + } + for _, def := range exposed { + if !wantNames[def.Name] { + t.Fatalf("unexpected exposed tool %q", def.Name) + } + if def.Name == "mcp__srv__a" && def.Parameters["type"] != "object" { + t.Fatalf("expected full schema for deferred tool on inactive path, got %#v", def.Parameters) + } + } +} + +// legacyToolDefinitions reconstructs the PRE-deferral tool-list builder: every +// tool that is visible (passes the operator filters) and advertised for the mode, +// EXCEPT tool_search, rendered with its full schema and alpha-sorted by name. It +// is the byte-identity reference for the inactive partition path. +func legacyToolDefinitions(registry *tools.Registry, permissionMode PermissionMode, options Options) []zeroruntime.ToolDefinition { + definitions := make([]zeroruntime.ToolDefinition, 0) + for _, tool := range registry.All() { + if !ToolVisible(tool, permissionMode, options.EnabledTools, options.DisabledTools) { + continue + } + if tool.Name() == tools.ToolSearchToolName { + continue + } + definitions = append(definitions, zeroruntime.ToolDefinition{ + Name: tool.Name(), + Description: tool.Description(), + Parameters: schemaToRuntimeMap(tool.Parameters()), + }) + } + sort.Slice(definitions, func(left, right int) bool { + return definitions[left].Name < definitions[right].Name + }) + return definitions +} + +// Below-threshold-but-eligible (count < threshold) is also inactive. +func TestPartitionToolsBelowThresholdInactive(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(fakeDeferredTool{name: "mcp__srv__a", desc: "a"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__b", desc: "b"}) + + exposed, reminder := partitionTools(registry, PermissionModeAuto, Options{DeferThreshold: 10}, map[string]bool{}) + if reminder != "" { + t.Fatalf("expected empty reminder below threshold, got %q", reminder) + } + if len(exposed) != 2 { + t.Fatalf("expected both deferred tools exposed below threshold, got %#v", exposed) + } +} + +// FIX 6(a): a deferred tool removed by DisabledTools must NOT count toward the +// eligible total. With N deferred registered and threshold == N, disabling one +// drops the surviving eligible count to N-1 < threshold, so the partition takes +// the INACTIVE path (empty reminder, all VISIBLE tools exposed with full schemas, +// the disabled one filtered out entirely). +func TestPartitionToolsDisabledDeferredDropsBelowThresholdInactive(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta"}) + + exposed, reminder := partitionTools(registry, PermissionModeAuto, Options{ + DeferThreshold: 2, + DisabledTools: []string{"mcp__srv__beta"}, + }, map[string]bool{}) + + // Eligible drops to 1 (< threshold 2) => inactive: empty reminder. + if reminder != "" { + t.Fatalf("expected inactive path (empty reminder) when a disable drops eligible below threshold, got %q", reminder) + } + // Only the surviving visible tool is exposed, with its FULL schema; the disabled + // tool is filtered out entirely. + if len(exposed) != 1 || exposed[0].Name != "mcp__srv__alpha" { + t.Fatalf("expected only mcp__srv__alpha exposed on inactive path, got %#v", exposed) + } + if exposed[0].Parameters["type"] != "object" { + t.Fatalf("surviving deferred tool must keep its full schema on inactive path, got %#v", exposed[0].Parameters) + } +} + +// FIX 6(b): with deferral ACTIVE, a DisabledTools-hidden deferred tool must never +// appear in the reminder NOR be exposed — it is filtered out before the partition +// even considers it. +func TestPartitionToolsActiveExcludesDisabledDeferredFromReminderAndExposed(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__gamma", desc: "gamma"}) + registry.Register(fakeToolSearchTool{}) // usable loader => deferral can activate + + // 3 deferred, disable beta => 2 surviving eligible, threshold 2 => active. + exposed, reminder := partitionTools(registry, PermissionModeAuto, Options{ + DeferThreshold: 2, + DisabledTools: []string{"mcp__srv__beta"}, + }, map[string]bool{}) + + if reminder == "" { + t.Fatalf("expected active path with a non-empty reminder for the unloaded tools") + } + if strings.Contains(reminder, "mcp__srv__beta") { + t.Fatalf("disabled deferred tool must not appear in the reminder, got %q", reminder) + } + for _, def := range exposed { + if def.Name == "mcp__srv__beta" { + t.Fatalf("disabled deferred tool must not be exposed, got %#v", exposed) + } + } + // The two surviving deferred tools are hidden (unloaded) and named in the reminder. + if !strings.Contains(reminder, "mcp__srv__alpha") || !strings.Contains(reminder, "mcp__srv__gamma") { + t.Fatalf("reminder must list the surviving unloaded deferred tools, got %q", reminder) + } +} + +func TestPartitionToolsActiveHidesUnloadedExposesLoaded(t *testing.T) { + root := t.TempDir() + registry := tools.NewRegistry() + registry.Register(tools.NewReadFileTool(root)) // non-deferred builtin + registry.Register(fakeToolSearchTool{}) // non-deferred, must stay exposed + registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha tool"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta tool"}) + + loaded := map[string]bool{"mcp__srv__alpha": true} + + // 2 eligible deferred tools, threshold 2 => active. + exposed, reminder := partitionTools(registry, PermissionModeAuto, Options{DeferThreshold: 2}, loaded) + + exposedNames := map[string]bool{} + for _, def := range exposed { + exposedNames[def.Name] = true + } + if !exposedNames["read_file"] { + t.Fatalf("expected builtin read_file exposed, got %#v", exposed) + } + if !exposedNames["tool_search"] { + t.Fatalf("expected tool_search exposed on active path, got %#v", exposed) + } + if !exposedNames["mcp__srv__alpha"] { + t.Fatalf("expected loaded deferred tool exposed, got %#v", exposed) + } + if exposedNames["mcp__srv__beta"] { + t.Fatalf("unloaded deferred tool must be hidden from exposed, got %#v", exposed) + } + if reminder == "" { + t.Fatalf("expected a non-empty reminder for the hidden tool") + } + if !strings.Contains(reminder, "mcp__srv__beta") { + t.Fatalf("expected reminder to list the hidden tool, got %q", reminder) + } + if strings.Contains(reminder, "mcp__srv__alpha") { + t.Fatalf("loaded tool must not appear in the reminder, got %q", reminder) + } +} + +func TestPartitionToolsActiveNothingHiddenEmptyReminder(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta"}) + registry.Register(fakeToolSearchTool{}) // usable loader => deferral can activate + + loaded := map[string]bool{"mcp__srv__alpha": true, "mcp__srv__beta": true} + exposed, reminder := partitionTools(registry, PermissionModeAuto, Options{DeferThreshold: 2}, loaded) + + exposedNames := map[string]bool{} + for _, def := range exposed { + exposedNames[def.Name] = true + } + if !exposedNames["mcp__srv__alpha"] || !exposedNames["mcp__srv__beta"] { + t.Fatalf("expected both loaded deferred tools exposed, got %#v", exposed) + } + if !exposedNames["tool_search"] { + t.Fatalf("expected tool_search exposed on active path, got %#v", exposed) + } + // BuildDeferredReminder returns "" for no hidden lines. + if reminder != "" { + t.Fatalf("expected empty reminder when nothing is hidden, got %q", reminder) + } +} + +func TestRunLoadsDeferredToolThenAdvertisesNextTurn(t *testing.T) { + registry := tools.NewRegistry() + // load_signal asks the loop to load mcp__srv__alpha next turn. + registry.Register(loadSignalTool{value: "mcp__srv__alpha"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha tool"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta tool"}) + // A real, usable tool_search must be registered for deferral to activate + // (otherwise the loop falls back to eager so it never strands the loader). + registry.Register(tools.NewToolSearchTool(registry)) + + provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{ + { // turn 1: call load_signal + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "c1", ToolName: "load_signal"}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "c1"}, + {Type: zeroruntime.StreamEventDone}, + }, + { // turn 2: final answer + {Type: zeroruntime.StreamEventText, Content: "done"}, + {Type: zeroruntime.StreamEventDone}, + }, + }} + + result, err := Run(context.Background(), "go", provider, Options{ + Registry: registry, + DeferThreshold: 2, // 2 deferred tools => active + }) + if err != nil { + t.Fatal(err) + } + if len(provider.requests) != 2 { + t.Fatalf("expected two turns, got %d", len(provider.requests)) + } + + // Turn 1: alpha is hidden (not advertised) and the reminder is a trailing user msg. + turn1Tools := map[string]bool{} + for _, def := range provider.requests[0].Tools { + turn1Tools[def.Name] = true + } + if turn1Tools["mcp__srv__alpha"] { + t.Fatalf("turn 1 must not advertise the not-yet-loaded deferred tool") + } + last1 := provider.requests[0].Messages[len(provider.requests[0].Messages)-1] + if last1.Role != zeroruntime.MessageRoleUser || !strings.Contains(last1.Content, "tool_search") { + t.Fatalf("turn 1 request must end with the deferred-tools reminder, got role=%s content=%q", last1.Role, last1.Content) + } + + // Turn 2: alpha is now loaded and advertised with a full schema. + turn2Tools := map[string]bool{} + for _, def := range provider.requests[1].Tools { + turn2Tools[def.Name] = true + } + if !turn2Tools["mcp__srv__alpha"] { + t.Fatalf("turn 2 must advertise the loaded deferred tool, got %#v", provider.requests[1].Tools) + } + if turn2Tools["mcp__srv__beta"] { + t.Fatalf("beta was never loaded; it must stay hidden in turn 2") + } + + // The reminder must NOT persist into the returned message history. + for _, m := range result.Messages { + if m.Role == zeroruntime.MessageRoleUser && strings.Contains(m.Content, "Call tool_search") { + t.Fatalf("the deferred-tools reminder must not be persisted in result.Messages") + } + } +} + +// TestRunReactiveRetryKeepsLoadedDeferredToolAndReminder drives a mid-run +// context-limit error that triggers reactive compaction+retry while deferral is +// ACTIVE and a deferred tool is already loaded. The retried turn must NOT be +// degraded to the empty-loaded/no-reminder state: it must still advertise the +// loaded tool's FULL schema and carry the for the hidden tool. +func TestRunReactiveRetryKeepsLoadedDeferredToolAndReminder(t *testing.T) { + registry := tools.NewRegistry() + // load_signal asks the loop to load mcp__srv__alpha next turn. + registry.Register(loadSignalTool{value: "mcp__srv__alpha"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha tool"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta tool"}) + // A real, usable tool_search must be registered for deferral to activate + // (otherwise the loop falls back to eager so it never strands the loader). + registry.Register(tools.NewToolSearchTool(registry)) + + // Request indices (mockProvider plays one turn per request, in order): + // 0: turn 1 — calls load_signal (loads mcp__srv__alpha for later turns) + // 1: turn 2 — emits a context-limit error MID-stream -> reactive recover + // 2: the summarize call inside Compact (must return non-empty text) + // 3: the RETRY request rebuilt after compaction — asserted below + provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{ + { // turn 1: call load_signal + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "c1", ToolName: "load_signal"}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "c1"}, + {Type: zeroruntime.StreamEventDone}, + }, + { // turn 2: mid-stream context-limit error -> reactive compaction + {Type: zeroruntime.StreamEventError, Error: "prompt is too long: 250000 tokens > 200000 maximum"}, + }, + { // summarize call inside Compact + {Type: zeroruntime.StreamEventText, Content: "SUMMARY"}, + {Type: zeroruntime.StreamEventDone}, + }, + { // retry of turn 2 after compaction: final answer + {Type: zeroruntime.StreamEventText, Content: "done"}, + {Type: zeroruntime.StreamEventDone}, + }, + }} + + // A large user prompt so the elided middle is big enough that reactive + // compaction actually shrinks the history (recover only retries when it can + // shrink). ContextWindow is set high enough that PROACTIVE compaction (which + // fires at 0.8 * window at the top of a turn) does NOT trigger first — so the + // only summarize call is the reactive one, keeping the request sequence + // predictable: turn1, errored turn2, summarize, retry. + bigPrompt := strings.Repeat("work on this task. ", 2000) + result, err := Run(context.Background(), bigPrompt, provider, Options{ + Registry: registry, + DeferThreshold: 2, // 2 deferred tools => active + ContextWindow: 1_000_000, + CompactionPreserveLast: 2, + }) + if err != nil { + t.Fatalf("run returned error: %v", err) + } + if result.FinalAnswer != "done" { + t.Fatalf("expected final answer from the retried turn, got %q", result.FinalAnswer) + } + // 4 provider calls: turn1, errored turn2, summarize, retry. + if len(provider.requests) != 4 { + t.Fatalf("expected 4 provider requests (turn1, errored turn2, summarize, retry), got %d", len(provider.requests)) + } + + retry := provider.requests[3] + + // The retry must advertise the loaded deferred tool with its FULL schema — + // NOT re-hide it as the empty-loaded partition would. + var alpha *zeroruntime.ToolDefinition + for i := range retry.Tools { + if retry.Tools[i].Name == "mcp__srv__alpha" { + alpha = &retry.Tools[i] + } + if retry.Tools[i].Name == "mcp__srv__beta" { + t.Fatalf("retry must keep the never-loaded deferred tool hidden, got %#v", retry.Tools) + } + } + if alpha == nil { + t.Fatalf("retry must still advertise the loaded deferred tool mcp__srv__alpha, got %#v", retry.Tools) + } + if alpha.Parameters["type"] != "object" { + t.Fatalf("retry must advertise the loaded tool's FULL schema, got %#v", alpha.Parameters) + } + + // The retry must carry the deferred-tools reminder as its trailing user + // message (the hidden tool is mcp__srv__beta) — not the no-reminder state. + last := retry.Messages[len(retry.Messages)-1] + if last.Role != zeroruntime.MessageRoleUser || !strings.Contains(last.Content, "tool_search") { + t.Fatalf("retry request must end with the deferred-tools reminder, got role=%s content=%q", last.Role, last.Content) + } + if !strings.Contains(last.Content, "mcp__srv__beta") { + t.Fatalf("retry reminder must list the still-hidden tool mcp__srv__beta, got %q", last.Content) + } + if strings.Contains(last.Content, "mcp__srv__alpha") { + t.Fatalf("loaded tool must not appear in the retry reminder, got %q", last.Content) + } + + // The reminder must NOT persist into the returned message history, even on + // the reactive-retry path. + for _, m := range result.Messages { + if m.Role == zeroruntime.MessageRoleUser && strings.Contains(m.Content, "Call tool_search") { + t.Fatalf("the deferred-tools reminder must not be persisted in result.Messages") + } + } +} + +// connectErrorProvider returns a connect-time error (StreamCompletion itself +// returns a non-nil error) on the request at index errAtRequest, exercising the +// FIRST reactive-recovery block in the loop (loop.go:99). Every other request +// streams the corresponding turn's events. It mirrors mockProvider's recording so +// the rebuilt retry request can be asserted. +type connectErrorProvider struct { + turns [][]zeroruntime.StreamEvent + errAtRequest int + errText string + requests []zeroruntime.CompletionRequest +} + +func (provider *connectErrorProvider) StreamCompletion(_ context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) { + provider.requests = append(provider.requests, request) + index := len(provider.requests) - 1 + if index == provider.errAtRequest { + return nil, errors.New(provider.errText) + } + events := []zeroruntime.StreamEvent{{Type: zeroruntime.StreamEventDone}} + if index < len(provider.turns) { + events = provider.turns[index] + } + ch := make(chan zeroruntime.StreamEvent, len(events)) + for _, event := range events { + ch <- event + } + close(ch) + return ch, nil +} + +// TestRunConnectTimeReactiveRetryKeepsLoadedDeferredToolAndReminder mirrors +// TestRunReactiveRetryKeepsLoadedDeferredToolAndReminder but drives the +// CONNECT-TIME error path: StreamCompletion itself returns a non-nil +// context-limit error (rather than a mid-stream StreamEventError). With deferral +// ACTIVE and a deferred tool already loaded, the rebuilt retry request must still +// advertise the loaded tool's FULL schema AND carry the deferred-tools reminder — +// i.e. the first reactive block reuses exposed/reminder, not the empty-loaded +// partition. +func TestRunConnectTimeReactiveRetryKeepsLoadedDeferredToolAndReminder(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(loadSignalTool{value: "mcp__srv__alpha"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha tool"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta tool"}) + // A real, usable tool_search must be registered for deferral to activate + // (otherwise the loop falls back to eager so it never strands the loader). + registry.Register(tools.NewToolSearchTool(registry)) + + // Request indices: + // 0: turn 1 — calls load_signal (loads mcp__srv__alpha for later turns) + // 1: turn 2 — StreamCompletion returns a connect-time context-limit error + // 2: the summarize call inside Compact (must return non-empty text) + // 3: the RETRY request rebuilt after compaction — asserted below + provider := &connectErrorProvider{ + errAtRequest: 1, + errText: "prompt is too long: 250000 tokens > 200000 maximum", + turns: [][]zeroruntime.StreamEvent{ + { // turn 1: call load_signal + {Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "c1", ToolName: "load_signal"}, + {Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "c1"}, + {Type: zeroruntime.StreamEventDone}, + }, + nil, // index 1: replaced by the connect-time error + { // index 2: summarize call inside Compact + {Type: zeroruntime.StreamEventText, Content: "SUMMARY"}, + {Type: zeroruntime.StreamEventDone}, + }, + { // index 3: retry of turn 2 after compaction: final answer + {Type: zeroruntime.StreamEventText, Content: "done"}, + {Type: zeroruntime.StreamEventDone}, + }, + }, + } + + bigPrompt := strings.Repeat("work on this task. ", 2000) + result, err := Run(context.Background(), bigPrompt, provider, Options{ + Registry: registry, + DeferThreshold: 2, // 2 deferred tools => active + ContextWindow: 1_000_000, + CompactionPreserveLast: 2, + }) + if err != nil { + t.Fatalf("run returned error: %v", err) + } + if result.FinalAnswer != "done" { + t.Fatalf("expected final answer from the retried turn, got %q", result.FinalAnswer) + } + // 4 provider calls: turn1, errored-connect turn2, summarize, retry. + if len(provider.requests) != 4 { + t.Fatalf("expected 4 provider requests (turn1, errored turn2, summarize, retry), got %d", len(provider.requests)) + } + + retry := provider.requests[3] + + // The retry must advertise the loaded deferred tool with its FULL schema and + // keep the never-loaded one hidden. + var alpha *zeroruntime.ToolDefinition + for i := range retry.Tools { + if retry.Tools[i].Name == "mcp__srv__alpha" { + alpha = &retry.Tools[i] + } + if retry.Tools[i].Name == "mcp__srv__beta" { + t.Fatalf("retry must keep the never-loaded deferred tool hidden, got %#v", retry.Tools) + } + } + if alpha == nil { + t.Fatalf("retry must still advertise the loaded deferred tool mcp__srv__alpha, got %#v", retry.Tools) + } + if alpha.Parameters["type"] != "object" { + t.Fatalf("retry must advertise the loaded tool's FULL schema, got %#v", alpha.Parameters) + } + + // The retry must carry the deferred-tools reminder as its trailing user message. + last := retry.Messages[len(retry.Messages)-1] + if last.Role != zeroruntime.MessageRoleUser || !strings.Contains(last.Content, "tool_search") { + t.Fatalf("retry request must end with the deferred-tools reminder, got role=%s content=%q", last.Role, last.Content) + } + if !strings.Contains(last.Content, "mcp__srv__beta") { + t.Fatalf("retry reminder must list the still-hidden tool mcp__srv__beta, got %q", last.Content) + } + if strings.Contains(last.Content, "mcp__srv__alpha") { + t.Fatalf("loaded tool must not appear in the retry reminder, got %q", last.Content) + } + + // The reminder must NOT persist into the returned message history. + for _, m := range result.Messages { + if m.Role == zeroruntime.MessageRoleUser && strings.Contains(m.Content, "Call tool_search") { + t.Fatalf("the deferred-tools reminder must not be persisted in result.Messages") + } + } +} + +// TestAllowlistedDeferredToolsKeepToolSearchReachable is the FIX 1+2 dead-end +// regression: N deferred tools are registered alongside tool_search, the operator +// allowlists ONLY the N deferred names (NOT tool_search), and DeferThreshold == N. +// Deferral must ACTIVATE, the partition must STILL expose tool_search (the gateway +// to the allowlisted tools), and a tool_search call must be DISPATCH-ALLOWED — so +// the reminder never points the model at an unreachable tool. +func TestAllowlistedDeferredToolsKeepToolSearchReachable(t *testing.T) { + registry := tools.NewRegistry() + deferredNames := []string{"mcp__srv__alpha", "mcp__srv__beta"} + for _, name := range deferredNames { + registry.Register(fakeDeferredTool{name: name, desc: name + " tool"}) + } + // The REAL tool_search (so partitionTools can look it up by name and dispatch + // can run it through registry.RunWithOptions). + registry.Register(tools.NewToolSearchTool(registry)) + + options := Options{ + DeferThreshold: len(deferredNames), // threshold == N => active + // Allowlist the deferred tools but NOT tool_search. + EnabledTools: append([]string{}, deferredNames...), + } + + exposed, reminder := partitionTools(registry, PermissionModeAuto, options, map[string]bool{}) + + // Deferral active => non-empty reminder advertising the hidden deferred tools. + if reminder == "" { + t.Fatalf("expected deferral active (non-empty reminder) at threshold N, got empty") + } + if !strings.Contains(reminder, "tool_search") { + t.Fatalf("reminder must instruct the model to call tool_search, got %q", reminder) + } + + // FIX 2(a): tool_search is exposed even though the allowlist omits it. + exposedNames := map[string]bool{} + for _, def := range exposed { + exposedNames[def.Name] = true + } + if !exposedNames["tool_search"] { + t.Fatalf("tool_search must be exposed on the active path despite the allowlist omitting it, got %#v", exposed) + } + // The allowlisted-but-unloaded deferred tools are hidden (not exposed). + for _, name := range deferredNames { + if exposedNames[name] { + t.Fatalf("unloaded deferred tool %q must be hidden when active, got %#v", name, exposed) + } + } + + // FIX 2(b): a tool_search call must be DISPATCH-ALLOWED (not rejected by the + // allowlist that omits it). It runs through the registry and reports a load. + result, abortErr := executeToolCall( + context.Background(), + registry, + ToolCall{ID: "c1", Name: "tool_search", Arguments: `{"query":"select:mcp__srv__alpha"}`}, + PermissionModeAuto, + options, + ) + if abortErr != nil { + t.Fatalf("unexpected abort error: %v", abortErr) + } + if result.Status != tools.StatusOK { + t.Fatalf("tool_search call was rejected: status=%s output=%q", result.Status, result.Output) + } + if result.DenialReason == DenialFiltered { + t.Fatalf("tool_search must not be denied by the allowlist, got DenialFiltered: %q", result.Output) + } + if len(result.LoadedTools) != 1 || result.LoadedTools[0] != "mcp__srv__alpha" { + t.Fatalf("tool_search call must load mcp__srv__alpha, got LoadedTools=%#v", result.LoadedTools) + } +} + +// TestDisabledToolSearchStaysHiddenAndRejectedWhenActive verifies an explicit +// DisabledTools entry for tool_search is STILL honored on the active path: the +// loader is not exposed and a call to it is rejected (FIX 2 exempts the allowlist +// only, never the denylist). +func TestDisabledToolSearchFallsBackToEager(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha"}) + registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta"}) + registry.Register(tools.NewToolSearchTool(registry)) + + options := Options{ + DeferThreshold: 2, + DisabledTools: []string{"tool_search"}, + } + + // tool_search is explicitly disabled, so it can never run. Deferral must NOT + // activate — otherwise the loop would hide the deferred tools behind a loader + // the dispatch gate rejects (an inescapable dead-end). Expect the eager / + // inactive fallback: every deferred tool exposed with its full schema, no + // tool_search definition, and an empty reminder. + exposed, reminder := partitionTools(registry, PermissionModeAuto, options, map[string]bool{}) + if reminder != "" { + t.Fatalf("inactive fallback must emit no reminder, got %q", reminder) + } + exposedNames := make(map[string]bool, len(exposed)) + for _, def := range exposed { + exposedNames[def.Name] = true + } + if exposedNames["tool_search"] { + t.Fatalf("tool_search must not be advertised when disabled, got %#v", exposed) + } + if !exposedNames["mcp__srv__alpha"] || !exposedNames["mcp__srv__beta"] { + t.Fatalf("deferred tools must be exposed eagerly when deferral is inactive, got %#v", exposed) + } + + // The deferred tools themselves stay directly callable (not stranded behind a + // disabled loader). + result, abortErr := executeToolCall( + context.Background(), + registry, + ToolCall{ID: "c1", Name: "mcp__srv__alpha", Arguments: `{}`}, + PermissionModeAuto, + options, + ) + if abortErr != nil { + t.Fatalf("unexpected abort error: %v", abortErr) + } + if result.Status != tools.StatusOK { + t.Fatalf("deferred tool must be callable under eager fallback, got status=%s output=%q", result.Status, result.Output) + } +} diff --git a/internal/agent/loop.go b/internal/agent/loop.go index ed70810d..f4309e79 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -67,6 +67,11 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) guards := newGuardState() compactor := newCompactionState(options) + // loaded tracks deferred-eligible tools the model has pulled via tool_search + // during THIS run. It is consulted by partitionTools each turn to expose a + // loaded tool's full schema; it lives only for the run (v1 within-run scope). + loaded := map[string]bool{} + result := Result{Messages: copyMessages(messages)} for turn := 0; turn < maxTurns; turn++ { result.Turns = turn + 1 @@ -76,9 +81,19 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // request. A no-op when ContextWindow == 0 (compaction disabled). messages = compactor.maybeCompact(ctx, provider, messages) + exposed, reminder := partitionTools(registry, permissionMode, options, loaded) request := zeroruntime.CompletionRequest{ Messages: copyMessages(messages), - Tools: toolDefinitions(registry, permissionMode, options), + Tools: exposed, + } + if reminder != "" { + // Append to the per-turn request copy only — NEVER to persistent + // messages — so the reminder refreshes each turn and never accumulates + // in the saved transcript. + request.Messages = append(request.Messages, zeroruntime.Message{ + Role: zeroruntime.MessageRoleUser, + Content: reminder, + }) } // Report the per-category context budget for this turn so a surface can @@ -97,9 +112,23 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) result.Messages = copyMessages(messages) return result, retryErr } + // Rebuild from the compacted messages but reuse the SAME active-mode + // partition (exposed) and reminder computed for this turn: they depend + // on registry+loaded, not on the messages, so they stay valid after + // compaction. Using the bare toolDefinitions here would route through an + // empty-loaded partition, re-hiding every already-loaded deferred tool + // and dropping the reminder when deferral is active. request = zeroruntime.CompletionRequest{ Messages: copyMessages(messages), - Tools: toolDefinitions(registry, permissionMode, options), + Tools: exposed, + } + if reminder != "" { + // Append to the per-turn retry copy only — NEVER to persistent + // messages — matching the main path's reminder-not-persisted invariant. + request.Messages = append(request.Messages, zeroruntime.Message{ + Role: zeroruntime.MessageRoleUser, + Content: reminder, + }) } stream, err = provider.StreamCompletion(ctx, request) } @@ -123,9 +152,22 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) result.Messages = copyMessages(messages) return result, retryErr } + // Reuse the SAME active-mode partition (exposed) and reminder from this + // turn rather than the bare toolDefinitions: exposed/reminder depend on + // registry+loaded (not the messages), so they stay valid after compaction. + // Routing through an empty-loaded partition here would re-hide every + // already-loaded deferred tool and drop the reminder when deferral is active. retryRequest := zeroruntime.CompletionRequest{ Messages: copyMessages(messages), - Tools: toolDefinitions(registry, permissionMode, options), + Tools: exposed, + } + if reminder != "" { + // Append to the per-turn retry copy only — NEVER to persistent + // messages — matching the main path's reminder-not-persisted invariant. + retryRequest.Messages = append(retryRequest.Messages, zeroruntime.Message{ + Role: zeroruntime.MessageRoleUser, + Content: reminder, + }) } retryStream, retryStreamErr := provider.StreamCompletion(ctx, retryRequest) if retryStreamErr != nil { @@ -212,6 +254,12 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) if options.OnToolResult != nil { options.OnToolResult(toolResult) } + // Union the deferred tools this result asked to load into the per-run + // set BEFORE any abort/stop/guard branch, so a load that coincides with + // a turn-ending result is still recorded for the next turn's partition. + for _, name := range toolResult.LoadedTools { + loaded[name] = true + } if turnRequestedModel == "" && toolResult.RequestedModel != "" { turnRequestedModel = toolResult.RequestedModel } @@ -355,7 +403,13 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal }, nil } } - if !ToolAllowedByFilters(call.Name, options.EnabledTools, options.DisabledTools) { + // tool_search is the gateway to the allowlisted deferred tools, so a non-empty + // EnabledTools allowlist that omits it must NOT reject the call — otherwise the + // reminder points the model at a tool the dispatch gate rejects (an inescapable + // dead-end). The allowlist is exempted; an explicit DisabledTools entry for + // tool_search is still honored (only the allowlist is exempted, not the denylist). + toolSearchAllowed := call.Name == tools.ToolSearchToolName && !containsToolName(options.DisabledTools, tools.ToolSearchToolName) + if !toolSearchAllowed && !ToolAllowedByFilters(call.Name, options.EnabledTools, options.DisabledTools) { return ToolResult{ ToolCallID: call.ID, Name: call.Name, @@ -435,6 +489,10 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal ReasoningEffort: options.ReasoningEffort, Depth: options.Depth, Cwd: options.Cwd, + // Forward the run's operator tool filters so a filter-aware tool + // (tool_search) never discloses or loads an operator-hidden deferred tool. + EnabledTools: options.EnabledTools, + DisabledTools: options.DisabledTools, // The sandbox decision (if any) is returned synchronously on the Result and // used here for permission event building. }) @@ -457,6 +515,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal Redacted: result.Redacted, ChangedFiles: result.ChangedFiles, Display: result.Display, + LoadedTools: loadedToolsFromResult(result.Meta), // A tool may signal a mid-run model escalation by carrying the target id // in Meta["escalate_to_model"]. Lift it into the typed loop-level field; // the Run turn loop performs the actual provider switch. Empty for every @@ -844,24 +903,102 @@ func permissionActionFromSandbox(action sandbox.Action) PermissionAction { } } -func toolDefinitions(registry *tools.Registry, permissionMode PermissionMode, options Options) []zeroruntime.ToolDefinition { +// partitionTools builds the per-turn advertised tool list and an optional +// deferred-tools reminder. INACTIVE (DeferThreshold <= 0 or the eligible count is +// below it): every visible tool is exposed with its full schema EXCEPT tool_search +// (dropped so it is never advertised when it cannot help), and the reminder is +// empty — byte-identical to the pre-deferral output. ACTIVE: a deferred-eligible +// tool is exposed only when loaded[name]; otherwise it is hidden and its compact +// line goes into the reminder. Non-deferred tools (including tool_search) are +// always exposed. The exposed slice is alpha-sorted by name, matching the legacy +// order so the inactive path is stable. +func partitionTools(registry *tools.Registry, permissionMode PermissionMode, options Options, loaded map[string]bool) ([]zeroruntime.ToolDefinition, string) { registeredTools := registry.All() - definitions := make([]zeroruntime.ToolDefinition, 0, len(registeredTools)) + + visible := make([]tools.Tool, 0, len(registeredTools)) + eligible := 0 for _, tool := range registeredTools { if !ToolVisible(tool, permissionMode, options.EnabledTools, options.DisabledTools) { continue } + visible = append(visible, tool) + if tools.IsDeferred(tool) { + eligible++ + } + } + + // Deferral may activate only when tool_search is actually runnable; otherwise + // the loop would hide deferred tools behind a loader the dispatch gate rejects + // — an inescapable dead-end. "Runnable" mirrors executeToolCall's gate: + // registered, not in DisabledTools, and advertised in the current permission + // mode (e.g. not spec-draft, where tool_search is not advertised). The + // EnabledTools allowlist is intentionally NOT checked here — tool_search is + // exempt from the allowlist at dispatch, so an allowlist that omits it must + // not disable deferral. + loader, loaderFound := registry.Get(tools.ToolSearchToolName) + loaderUsable := loaderFound && + !containsToolName(options.DisabledTools, tools.ToolSearchToolName) && + ToolAdvertised(loader, permissionMode) + + active := options.DeferThreshold > 0 && eligible >= options.DeferThreshold && loaderUsable + + definitions := make([]zeroruntime.ToolDefinition, 0, len(visible)) + exposedNames := make(map[string]bool, len(visible)) + var hiddenLines []string + for _, tool := range visible { + name := tool.Name() + deferred := tools.IsDeferred(tool) + + if !active { + // Inactive: byte-identical to legacy, but tool_search is never advertised. + if name == tools.ToolSearchToolName { + continue + } + definitions = append(definitions, zeroruntime.ToolDefinition{ + Name: name, + Description: tool.Description(), + Parameters: schemaToRuntimeMap(tool.Parameters()), + }) + exposedNames[name] = true + continue + } + + // Active path. + if deferred && !loaded[name] { + hiddenLines = append(hiddenLines, tools.DeferredLine(tool)) + continue + } definitions = append(definitions, zeroruntime.ToolDefinition{ - Name: tool.Name(), + Name: name, Description: tool.Description(), Parameters: schemaToRuntimeMap(tool.Parameters()), }) + exposedNames[name] = true + } + + // On the ACTIVE path tool_search is guaranteed runnable (active implies + // loaderUsable), so it must ALWAYS be reachable — even when a non-empty + // EnabledTools allowlist omits it (the operator allowlisted the deferred + // tools, not the loader). Expose its full definition whenever it is not + // already in the exposed set. This never runs on the inactive path, so the + // byte-identical below-threshold output is preserved. + if active && !exposedNames[tools.ToolSearchToolName] { + definitions = append(definitions, zeroruntime.ToolDefinition{ + Name: loader.Name(), + Description: loader.Description(), + Parameters: schemaToRuntimeMap(loader.Parameters()), + }) } sort.Slice(definitions, func(left int, right int) bool { return definitions[left].Name < definitions[right].Name }) - return definitions + + reminder := "" + if active { + reminder = tools.BuildDeferredReminder(hiddenLines) + } + return definitions, reminder } func ToolVisible(tool tools.Tool, permissionMode PermissionMode, enabledTools []string, disabledTools []string) bool { @@ -982,6 +1119,25 @@ func stopReasonFromToolResult(result ToolResult) StopReason { return "" } +// loadedToolsFromResult extracts the deferred-tool names a tool (tool_search) +// asked the loop to expose next turn from Meta["load_tools"] (comma-separated). +// It trims each name and drops empties; returns nil when the key is absent or +// yields no names, so an ordinary result keeps a nil LoadedTools. Mirrors the +// Meta-driven control signal read by stopReasonFromToolResult. +func loadedToolsFromResult(meta map[string]string) []string { + raw := meta["load_tools"] + if strings.TrimSpace(raw) == "" { + return nil + } + var names []string + for _, part := range strings.Split(raw, ",") { + if trimmed := strings.TrimSpace(part); trimmed != "" { + names = append(names, trimmed) + } + } + return names +} + // appendAbortedToolResults adds a placeholder tool-result message for each of // the given (unexecuted) tool calls, so every advertised tool_use keeps a // matching tool_result when the loop halts a turn before all calls have run. diff --git a/internal/agent/types.go b/internal/agent/types.go index c71cac7b..dbe68b52 100644 --- a/internal/agent/types.go +++ b/internal/agent/types.go @@ -54,6 +54,10 @@ type ToolResult struct { // DenialReason categorizes why a tool call was blocked (empty when it ran). // It lets a surface distinguish the cause precisely instead of parsing Output. DenialReason DenialCategory + // LoadedTools carries the deferred-tool names a tool_search call asked the + // loop to expose next turn (lifted from Meta["load_tools"]). nil for every + // ordinary tool result; only tool_search populates it. + LoadedTools []string // RequestedModel is the model id a tool asked the loop to switch to for the // rest of the run (lifted from the tool's Meta["escalate_to_model"]). Empty // for every normal tool result; the Run loop performs the switch when it is @@ -130,6 +134,12 @@ type AskUserResponse struct { type Options struct { MaxTurns int + // DeferThreshold activates deferred MCP-tool loading: when the number of + // deferred-eligible visible tools is >= this value (and it is > 0), their + // full schemas are withheld and advertised as compact lines via tool_search. + // 0 (or below the eligible count) keeps every tool eager — byte-identical to + // the pre-deferral behavior. + DeferThreshold int // Specialist/sub-agent metadata is carried through exec now and consumed by // the specialist runtime in later slices. SessionID string diff --git a/internal/cli/app.go b/internal/cli/app.go index 822e78d3..af9f0daf 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -328,16 +328,6 @@ func runInteractiveTUIWithSkin(stderr io.Writer, deps appDeps, skin string, perm return writeAppError(stderr, err.Error(), 1) } defer closeMCPRuntime(stderr, mcpRuntime) - sandboxStore, err := deps.newSandboxStore() - if err != nil { - return writeAppError(stderr, "failed to initialize sandbox grants: "+err.Error(), 1) - } - sandboxEngine := sandbox.NewEngine(sandbox.EngineOptions{ - WorkspaceRoot: workspaceRoot, - Policy: applyConfiguredAutonomyCeiling(sandbox.DefaultPolicy(), resolved.Sandbox.MaxAutonomy), - Store: sandboxStore, - Backend: deps.selectSandboxBackend(sandbox.BackendOptions{}), - }) // Ask (not Auto) is the interactive default: in Auto, ToolAdvertised exposes // only PermissionAllow tools, so prompt-gated tools (write_file/edit_file/bash/ // apply_patch) would never be offered to the model — the TUI could neither edit @@ -345,9 +335,30 @@ func runInteractiveTUIWithSkin(stderr io.Writer, deps appDeps, skin string, perm // OnPermissionRequest flow; shift+tab lets the user switch modes live. An // explicit --skip-permissions-unsafe launch overrides this to unsafe (the only // way to reach unsafe, since shift+tab deliberately cycles auto↔ask only). + // + // Resolve the effective mode BEFORE the deferral gate below so the registration + // count uses the SAME permission mode the agent loop's partition will use; an + // empty mode here would mis-gate prompt-advertised deferred tools. if permissionMode == "" { permissionMode = agent.PermissionModeAsk } + // Activate deferred MCP-tool loading for the interactive run only when the + // VISIBLE deferred-eligible count meets the resolved threshold, matching exec. + // The registry is complete (core + specialist + MCP) here, so the count is + // accurate; below threshold this is a no-op and the surface is unchanged. The + // interactive surface applies no operator tool filters, so enabled/disabled are + // nil — matching the AgentOptions below. + registerToolSearchIfEligible(registry, resolved.Tools.DeferThreshold, permissionMode, nil, nil) + sandboxStore, err := deps.newSandboxStore() + if err != nil { + return writeAppError(stderr, "failed to initialize sandbox grants: "+err.Error(), 1) + } + sandboxEngine := sandbox.NewEngine(sandbox.EngineOptions{ + WorkspaceRoot: workspaceRoot, + Policy: applyConfiguredAutonomyCeiling(sandbox.DefaultPolicy(), resolved.Sandbox.MaxAutonomy), + Store: sandboxStore, + Backend: deps.selectSandboxBackend(sandbox.BackendOptions{}), + }) return deps.runTUI(context.Background(), tui.Options{ Cwd: workspaceRoot, ProviderName: resolved.Provider.Name, @@ -364,6 +375,7 @@ func runInteractiveTUIWithSkin(stderr io.Writer, deps appDeps, skin string, perm PermissionMode: permissionMode, Autonomy: string(sandbox.AutonomyLow), Sandbox: sandboxEngine, + DeferThreshold: resolved.Tools.DeferThreshold, }, PermissionMode: permissionMode, Skin: skin, diff --git a/internal/cli/deferred_wiring_test.go b/internal/cli/deferred_wiring_test.go new file mode 100644 index 00000000..868b4e92 --- /dev/null +++ b/internal/cli/deferred_wiring_test.go @@ -0,0 +1,221 @@ +package cli + +import ( + "bytes" + "context" + "errors" + "io" + "strings" + "testing" + + "github.com/Gitlawb/zero/internal/agent" + "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/mcp" + "github.com/Gitlawb/zero/internal/tools" + "github.com/Gitlawb/zero/internal/tui" +) + +// cliFakeDeferredTool is deferred-eligible (implements Deferred() bool), mirroring +// an MCP registry tool, so it counts toward the deferral threshold. +type cliFakeDeferredTool struct { + name string +} + +func (t cliFakeDeferredTool) Name() string { return t.name } +func (t cliFakeDeferredTool) Description() string { return "fake deferred tool" } +func (t cliFakeDeferredTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} } +func (t cliFakeDeferredTool) Safety() tools.Safety { + return tools.Safety{SideEffect: tools.SideEffectNetwork, Permission: tools.PermissionAllow} +} +func (t cliFakeDeferredTool) Run(context.Context, map[string]any) tools.Result { + return tools.Result{Status: tools.StatusOK, Output: "ok"} +} +func (t cliFakeDeferredTool) Deferred() bool { return true } + +func registryHasToolSearch(registry *tools.Registry) bool { + _, ok := registry.Get("tool_search") + return ok +} + +func TestRegisterToolSearchIfEligibleRegistersAtThreshold(t *testing.T) { + registry := tools.NewRegistry() + for i := 0; i < 3; i++ { + registry.Register(cliFakeDeferredTool{name: "mcp_srv_t" + string(rune('a'+i))}) + } + + registerToolSearchIfEligible(registry, 3, agent.PermissionModeAuto, nil, nil) + + if !registryHasToolSearch(registry) { + t.Fatal("expected tool_search registered when eligible count == threshold") + } +} + +func TestRegisterToolSearchIfEligibleSkipsBelowThreshold(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(cliFakeDeferredTool{name: "mcp_srv_ta"}) + registry.Register(cliFakeDeferredTool{name: "mcp_srv_tb"}) + // A plain (non-deferred) MCP-named tool must NOT count toward eligibility. + registry.Register(cliFakeMCPRegistryTool{}) + + registerToolSearchIfEligible(registry, 3, agent.PermissionModeAuto, nil, nil) + + if registryHasToolSearch(registry) { + t.Fatal("expected no tool_search when eligible count (2) < threshold (3)") + } +} + +func TestRegisterToolSearchIfEligibleSkipsWhenThresholdZero(t *testing.T) { + registry := tools.NewRegistry() + for i := 0; i < 5; i++ { + registry.Register(cliFakeDeferredTool{name: "mcp_srv_t" + string(rune('a'+i))}) + } + + registerToolSearchIfEligible(registry, 0, agent.PermissionModeAuto, nil, nil) + + if registryHasToolSearch(registry) { + t.Fatal("expected no tool_search when threshold is 0 (disabled)") + } +} + +func TestDeferredEligibleCountIgnoresCoreTools(t *testing.T) { + registry := newCoreRegistry(t.TempDir()) + // newCoreRegistry holds only built-ins; none implement Deferred(). + if got := deferredEligibleCount(registry, agent.PermissionModeAuto, nil, nil); got != 0 { + t.Fatalf("deferredEligibleCount(core) = %d, want 0", got) + } +} + +// FIX 1: a deferred tool the operator hid via --disabled-tools must NOT count +// toward the visible-deferred total, so registration agrees with the loop's +// activation gate. Two deferred + threshold 2 normally registers tool_search; +// disabling one drops the visible count to 1 (< 2) so it must NOT register. +func TestRegisterToolSearchSkipsWhenDisabledDropsVisibleBelowThreshold(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(cliFakeDeferredTool{name: "mcp_srv_ta"}) + registry.Register(cliFakeDeferredTool{name: "mcp_srv_tb"}) + + registerToolSearchIfEligible(registry, 2, agent.PermissionModeAuto, nil, []string{"mcp_srv_tb"}) + + if registryHasToolSearch(registry) { + t.Fatal("expected no tool_search: a disabled deferred tool must not count toward the visible-deferred total") + } +} + +// FIX 3: validateExecToolFilters must treat tool_search as always-valid even +// though it is not registered yet at validation time (it is registered later only +// when deferral activates). Listing it in --enabled-tools/--disabled-tools must +// not raise "Unknown tool". +func TestValidateExecToolFiltersAllowsToolSearch(t *testing.T) { + registry := newCoreRegistry(t.TempDir()) + // tool_search is NOT registered in this registry — it would be added later. + if _, present := registry.Get(tools.ToolSearchToolName); present { + t.Fatalf("precondition: tool_search must not be registered yet") + } + + if err := validateExecToolFilters(execOptions{enabledTools: []string{tools.ToolSearchToolName}}, registry); err != nil { + t.Fatalf("--enabled-tools tool_search must validate, got error: %v", err) + } + if err := validateExecToolFilters(execOptions{disabledTools: []string{tools.ToolSearchToolName}}, registry); err != nil { + t.Fatalf("--disabled-tools tool_search must validate, got error: %v", err) + } + // A genuinely unknown tool still errors. + if err := validateExecToolFilters(execOptions{enabledTools: []string{"definitely_not_a_tool"}}, registry); err == nil { + t.Fatal("expected an Unknown tool error for an unregistered, non-tool_search name") + } +} + +// TestRunExecListToolsAdvertisesMCPToolsWithoutToolSearch verifies that +// `exec --list-tools` lists the core + MCP tools WITHOUT tool_search. This is +// independent of the deferral threshold: --list-tools short-circuits and returns +// before resolveConfig and registerToolSearchIfEligible ever run (see exec.go), +// so tool_search is never registered on this path. The threshold gate itself is +// exercised by TestRegisterToolSearchIfEligible{RegistersAtThreshold, +// SkipsBelowThreshold,SkipsWhenThresholdZero} and the end-to-end at-threshold +// path in TestTUIRunThreadsDeferThresholdAndRegistersToolSearch. +func TestRunExecListToolsAdvertisesMCPToolsWithoutToolSearch(t *testing.T) { + cwd := t.TempDir() + var stdout bytes.Buffer + var stderr bytes.Buffer + + exitCode := runWithDeps([]string{"exec", "--list-tools"}, &stdout, &stderr, appDeps{ + getwd: func() (string, error) { return cwd, nil }, + resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) { + return config.ResolvedConfig{}, errors.New("provider should not be resolved for --list-tools") + }, + resolveMCPConfig: func(string) (config.MCPConfig, error) { + return config.MCPConfig{Servers: map[string]config.MCPServerConfig{ + "docs": {Type: "stdio", Command: "docs-mcp"}, + }}, nil + }, + newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil }, + registerMCPTools: func(_ context.Context, registry *tools.Registry, _ config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) { + registry.Register(cliFakeMCPRegistryTool{}) + return closeFunc(func() error { return nil }), nil + }, + }) + + if exitCode != exitSuccess { + t.Fatalf("expected exit %d, got %d: %s", exitSuccess, exitCode, stderr.String()) + } + out := stdout.String() + if !strings.Contains(out, "mcp_docs_lookup") { + t.Fatalf("expected MCP tool advertised by --list-tools, got %q", out) + } + // --list-tools never registers tool_search (it returns before the threshold + // gate runs); the threshold logic is verified by the unit/E2E tests named above. + if strings.Contains(out, "tool_search") { + t.Fatalf("expected NO tool_search from --list-tools, got %q", out) + } +} + +func TestTUIRunThreadsDeferThresholdAndRegistersToolSearch(t *testing.T) { + cwd := t.TempDir() + var captured agent.Options + var capturedRegistry *tools.Registry + + // Empty args route runWithDeps to the interactive TUI (runInteractiveTUIWithSkin). + // There is no "--tui" flag; that arg would hit the unknown-command path. + exitCode := runWithDeps([]string{}, io.Discard, io.Discard, appDeps{ + getwd: func() (string, error) { return cwd, nil }, + resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) { + return config.ResolvedConfig{ + Provider: config.ProviderProfile{ + Name: "p", + ProviderKind: config.ProviderKindOpenAICompatible, + BaseURL: "http://127.0.0.1/v1", + Model: "m", + }, + Tools: config.ToolsConfig{DeferThreshold: 2}, + }, nil + }, + resolveMCPConfig: func(string) (config.MCPConfig, error) { + return config.MCPConfig{Servers: map[string]config.MCPServerConfig{ + "docs": {Type: "stdio", Command: "docs-mcp"}, + }}, nil + }, + newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil }, + registerMCPTools: func(_ context.Context, registry *tools.Registry, _ config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) { + registry.Register(cliFakeDeferredTool{name: "mcp_docs_ta"}) + registry.Register(cliFakeDeferredTool{name: "mcp_docs_tb"}) + return closeFunc(func() error { return nil }), nil + }, + runTUI: func(_ context.Context, options tui.Options) int { + captured = options.AgentOptions + capturedRegistry = options.Registry + return exitSuccess + }, + }) + + if exitCode != exitSuccess { + t.Fatalf("expected exit %d, got %d", exitSuccess, exitCode) + } + if captured.DeferThreshold != 2 { + t.Fatalf("AgentOptions.DeferThreshold = %d, want 2", captured.DeferThreshold) + } + if capturedRegistry == nil { + t.Fatal("expected registry passed to TUI") + } + if _, ok := capturedRegistry.Get("tool_search"); !ok { + t.Fatal("expected tool_search registered for TUI run at/above threshold") + } +} diff --git a/internal/cli/exec.go b/internal/cli/exec.go index 348cac41..1a17d221 100644 --- a/internal/cli/exec.go +++ b/internal/cli/exec.go @@ -224,6 +224,23 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in if !config.HasProviderProfile(resolved.Provider) { return writeExecProviderError(stdout, stderr, options.outputFormat, "provider_error", "No provider configured. Set OPENAI_MODEL/OPENAI_API_KEY or add .zero/config.json.") } + // Activate deferred MCP-tool loading for this run only when the VISIBLE + // deferred-eligible count meets the resolved threshold; below threshold this + // is a no-op and tool advertising stays byte-identical. The registry is + // already complete (core + MCP) at this point, so the count is accurate. The + // permission mode and operator tool filters MUST match the values passed to + // agent.Run below so this registration gate counts the same population the + // loop's partition gate counts. + // tool_search is the only way to reach a hidden deferred tool, so if the + // operator explicitly disables it, deferral must not activate at all — + // otherwise a positive threshold would hide tools behind a loader the run + // rejects (a dead-end). Force the effective threshold to 0 so this + // registration gate and agent.Run's partition gate agree the run is inactive. + effectiveDeferThreshold := resolved.Tools.DeferThreshold + if toolListContains(options.disabledTools, tools.ToolSearchToolName) { + effectiveDeferThreshold = 0 + } + registerToolSearchIfEligible(registry, effectiveDeferThreshold, permissionMode, options.enabledTools, options.disabledTools) images, err := resolveExecImages(options.imagePaths, workspaceRoot) if err != nil { return writeExecFormatUsageError(stdout, stderr, options.outputFormat, err.Error()) @@ -388,6 +405,7 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in result, err := agent.Run(runCtx, agentPrompt, provider, agent.Options{ MaxTurns: resolved.MaxTurns, ContextWindow: modelContextWindow(modelRegistry, resolved.Provider.Model), + DeferThreshold: effectiveDeferThreshold, SessionID: preparedSession.Session.SessionID, CallingSessionID: options.callingSessionID, CallingToolUseID: options.callingToolUseID, @@ -501,6 +519,46 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in return exitSuccess } +// deferredEligibleCount returns the number of registered tools that are +// deferred-eligible (MCP tools) AND visible to the model for THIS run — i.e. they +// pass the same agent.ToolVisible gate (permission-mode advertising + operator +// allow/deny filters) that the agent loop's partitionTools applies when it +// decides whether deferral activates. Counting the SAME visible-deferred +// population here keeps registration and activation in agreement: tool_search is +// registered iff the partition will actually go active. Built-ins never implement +// the Deferred interface, so they never count. +func deferredEligibleCount(registry *tools.Registry, permissionMode agent.PermissionMode, enabledTools []string, disabledTools []string) int { + count := 0 + for _, tool := range registry.All() { + if !tools.IsDeferred(tool) { + continue + } + if !agent.ToolVisible(tool, permissionMode, enabledTools, disabledTools) { + continue + } + count++ + } + return count +} + +// registerToolSearchIfEligible registers the tool_search tool only when deferral +// is active for this run: the visible-deferred count (the same population the +// agent loop's partition counts) meets the (positive) threshold. Below threshold +// or with a zero/negative threshold, tool_search is never registered, so the +// agent loop's partition stays inactive and tool advertising is byte-identical to +// today. The permissionMode + enabled/disabled filters MUST match the values the +// run passes to agent.Run so the registration gate and the activation gate count +// the same tools. +func registerToolSearchIfEligible(registry *tools.Registry, deferThreshold int, permissionMode agent.PermissionMode, enabledTools []string, disabledTools []string) { + if deferThreshold <= 0 { + return + } + if deferredEligibleCount(registry, permissionMode, enabledTools, disabledTools) < deferThreshold { + return + } + registry.Register(tools.NewToolSearchTool(registry)) +} + func buildExecSandboxEngine(workspaceRoot string, resolved config.ResolvedConfig, deps appDeps) (*sandbox.Engine, error) { store, err := deps.newSandboxStore() if err != nil { diff --git a/internal/cli/exec_tools.go b/internal/cli/exec_tools.go index 9be7f249..7cda6f73 100644 --- a/internal/cli/exec_tools.go +++ b/internal/cli/exec_tools.go @@ -35,7 +35,11 @@ func validateExecToolFilters(options execOptions, registry *tools.Registry) erro if !toolNamePattern.MatchString(name) { return execUsageError{fmt.Sprintf("Invalid tool name %q.", name)} } - if _, ok := registry.Get(name); !ok { + // tool_search is registered later (only when deferral activates), so it is + // not in the registry yet at validation time. Treat it as always-valid so an + // operator can list it in --enabled-tools / --disabled-tools without a + // spurious "Unknown tool" error. + if _, ok := registry.Get(name); !ok && name != tools.ToolSearchToolName { return execUsageError{fmt.Sprintf("Unknown tool: %s", name)} } } diff --git a/internal/config/resolver.go b/internal/config/resolver.go index d887945a..3fa5c3a5 100644 --- a/internal/config/resolver.go +++ b/internal/config/resolver.go @@ -14,6 +14,8 @@ import ( const defaultMaxTurns = 12 +const defaultDeferThreshold = 10 + func Resolve(options ResolveOptions) (ResolvedConfig, error) { cfg := FileConfig{ MaxTurns: defaultMaxTurns, @@ -48,6 +50,13 @@ func Resolve(options ResolveOptions) (ResolvedConfig, error) { applyOverrides(&cfg, options.Overrides) + if !cfg.Tools.deferThresholdSet && cfg.Tools.DeferThreshold == 0 { + cfg.Tools.DeferThreshold = defaultDeferThreshold + } + if cfg.Tools.DeferThreshold < 0 { + return ResolvedConfig{}, fmt.Errorf("invalid tools.deferThreshold %d: must be >= 0", cfg.Tools.DeferThreshold) + } + if maxAutonomy := strings.TrimSpace(cfg.Sandbox.MaxAutonomy); maxAutonomy != "" { // Fail loud on an invalid ceiling. An unvalidated typo (e.g. "moderate") // would otherwise survive Resolve untouched, reach the sandbox bridge, @@ -70,6 +79,7 @@ func Resolve(options ResolveOptions) (ResolvedConfig, error) { MaxTurns: cfg.MaxTurns, MCP: cfg.MCP, Sandbox: cfg.Sandbox, + Tools: cfg.Tools, }, nil } @@ -117,6 +127,10 @@ func mergeConfig(dst *FileConfig, src FileConfig) { if maxAutonomy := strings.TrimSpace(src.Sandbox.MaxAutonomy); maxAutonomy != "" { dst.Sandbox.MaxAutonomy = maxAutonomy } + if src.Tools.deferThresholdSet { + dst.Tools.DeferThreshold = src.Tools.DeferThreshold + dst.Tools.deferThresholdSet = true + } } func mergeProjectConfig(dst *FileConfig, src FileConfig) error { @@ -137,6 +151,10 @@ func mergeProjectConfig(dst *FileConfig, src FileConfig) error { if maxAutonomy := strings.TrimSpace(src.Sandbox.MaxAutonomy); maxAutonomy != "" { dst.Sandbox.MaxAutonomy = maxAutonomy } + if src.Tools.deferThresholdSet { + dst.Tools.DeferThreshold = src.Tools.DeferThreshold + dst.Tools.deferThresholdSet = true + } return nil } @@ -441,6 +459,10 @@ func applyOverrides(cfg *FileConfig, overrides Overrides) { if maxAutonomy := strings.TrimSpace(overrides.Sandbox.MaxAutonomy); maxAutonomy != "" { cfg.Sandbox.MaxAutonomy = maxAutonomy } + if overrides.Tools.deferThresholdSet || overrides.Tools.DeferThreshold != 0 { + cfg.Tools.DeferThreshold = overrides.Tools.DeferThreshold + cfg.Tools.deferThresholdSet = true + } for _, provider := range overrides.Providers { mergeProvider(cfg, provider) } diff --git a/internal/config/resolver_test.go b/internal/config/resolver_test.go index f92bb769..9f040df4 100644 --- a/internal/config/resolver_test.go +++ b/internal/config/resolver_test.go @@ -1305,3 +1305,153 @@ func providerByName(t *testing.T, providers []ProviderProfile, name string) Prov t.Fatalf("provider %q not found in %#v", name, providers) return ProviderProfile{} } + +func TestResolveDefaultsDeferThresholdWhenUnset(t *testing.T) { + path := writeConfig(t, `{ + "activeProvider": "p", + "providers": [{"name": "p", "provider": "openai", "apiKey": "sk", "model": "m"}] + }`) + + resolved, err := Resolve(ResolveOptions{ProjectConfigPath: path, Env: map[string]string{}}) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if resolved.Tools.DeferThreshold != defaultDeferThreshold { + t.Fatalf("Tools.DeferThreshold = %d, want %d", resolved.Tools.DeferThreshold, defaultDeferThreshold) + } +} + +func TestResolveDeferThresholdFromFile(t *testing.T) { + path := writeConfig(t, `{ + "activeProvider": "p", + "providers": [{"name": "p", "provider": "openai", "apiKey": "sk", "model": "m"}], + "tools": {"deferThreshold": 4} + }`) + + resolved, err := Resolve(ResolveOptions{ProjectConfigPath: path, Env: map[string]string{}}) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if resolved.Tools.DeferThreshold != 4 { + t.Fatalf("Tools.DeferThreshold = %d, want 4", resolved.Tools.DeferThreshold) + } +} + +func TestResolveDeferThresholdZeroDisablesViaFile(t *testing.T) { + // 0 is a valid, meaningful value (disabled). It must survive Resolve and NOT + // be re-defaulted to 10, so a user can explicitly disable deferral. + path := writeConfig(t, `{ + "activeProvider": "p", + "providers": [{"name": "p", "provider": "openai", "apiKey": "sk", "model": "m"}], + "tools": {"deferThreshold": 0} + }`) + + resolved, err := Resolve(ResolveOptions{ProjectConfigPath: path, Env: map[string]string{}}) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if resolved.Tools.DeferThreshold != 0 { + t.Fatalf("Tools.DeferThreshold = %d, want 0 (explicit disable preserved)", resolved.Tools.DeferThreshold) + } +} + +func TestResolveDeferThresholdOverrideWins(t *testing.T) { + path := writeConfig(t, `{ + "activeProvider": "p", + "providers": [{"name": "p", "provider": "openai", "apiKey": "sk", "model": "m"}], + "tools": {"deferThreshold": 4} + }`) + + resolved, err := Resolve(ResolveOptions{ + ProjectConfigPath: path, + Env: map[string]string{}, + Overrides: Overrides{Tools: ToolsConfig{DeferThreshold: 20}}, + }) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if resolved.Tools.DeferThreshold != 20 { + t.Fatalf("Tools.DeferThreshold = %d, want 20 (override wins)", resolved.Tools.DeferThreshold) + } +} + +func TestResolveToolsOverrideDisablesDeferralOverNonZeroBase(t *testing.T) { + // A programmatic Override built via ToolsOverride(0) carries the presence flag, + // so it must override an explicit non-zero base threshold down to 0 (disabled). + // A bare ToolsConfig{DeferThreshold: 0} could not do this — it is indistinguishable + // from "unset" — which is exactly the trap ToolsOverride exists to avoid. + path := writeConfig(t, `{ + "activeProvider": "p", + "providers": [{"name": "p", "provider": "openai", "apiKey": "sk", "model": "m"}], + "tools": {"deferThreshold": 4} + }`) + + resolved, err := Resolve(ResolveOptions{ + ProjectConfigPath: path, + Env: map[string]string{}, + Overrides: Overrides{Tools: ToolsOverride(0)}, + }) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if resolved.Tools.DeferThreshold != 0 { + t.Fatalf("Tools.DeferThreshold = %d, want 0 (ToolsOverride(0) disables deferral)", resolved.Tools.DeferThreshold) + } +} + +func TestResolveToolsOverrideSetsNonZeroOverNonZeroBase(t *testing.T) { + // ToolsOverride(7) over an explicit non-zero base must win with the new value. + path := writeConfig(t, `{ + "activeProvider": "p", + "providers": [{"name": "p", "provider": "openai", "apiKey": "sk", "model": "m"}], + "tools": {"deferThreshold": 4} + }`) + + resolved, err := Resolve(ResolveOptions{ + ProjectConfigPath: path, + Env: map[string]string{}, + Overrides: Overrides{Tools: ToolsOverride(7)}, + }) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if resolved.Tools.DeferThreshold != 7 { + t.Fatalf("Tools.DeferThreshold = %d, want 7 (ToolsOverride(7) wins)", resolved.Tools.DeferThreshold) + } +} + +func TestApplyOverridesToolsOverrideZeroOverridesNonZero(t *testing.T) { + // Unit-level check on applyOverrides directly: ToolsOverride(0) must flip an + // explicit non-zero base to 0 (deferThresholdSet honored), while a bare + // non-zero ToolsConfig still overrides via the != 0 branch. + cfg := FileConfig{Tools: ToolsOverride(4)} + applyOverrides(&cfg, Overrides{Tools: ToolsOverride(0)}) + if cfg.Tools.DeferThreshold != 0 { + t.Fatalf("applyOverrides ToolsOverride(0): DeferThreshold = %d, want 0", cfg.Tools.DeferThreshold) + } + if !cfg.Tools.deferThresholdSet { + t.Fatal("applyOverrides ToolsOverride(0): deferThresholdSet = false, want true") + } + + cfg = FileConfig{Tools: ToolsOverride(4)} + applyOverrides(&cfg, Overrides{Tools: ToolsConfig{DeferThreshold: 5}}) + if cfg.Tools.DeferThreshold != 5 { + t.Fatalf("applyOverrides bare non-zero: DeferThreshold = %d, want 5", cfg.Tools.DeferThreshold) + } +} + +func TestResolveRejectsNegativeDeferThreshold(t *testing.T) { + path := writeConfig(t, `{ + "activeProvider": "p", + "providers": [{"name": "p", "provider": "openai", "apiKey": "sk", "model": "m"}], + "tools": {"deferThreshold": -1} + }`) + + _, err := Resolve(ResolveOptions{ProjectConfigPath: path, Env: map[string]string{}}) + if err == nil { + t.Fatal("Resolve() error = nil for negative deferThreshold, want failure") + } + if !strings.Contains(err.Error(), "tools.deferThreshold") { + t.Fatalf("Resolve() error = %v, want tools.deferThreshold rejection", err) + } +} diff --git a/internal/config/types.go b/internal/config/types.go index 0dde20b0..d29b0de1 100644 --- a/internal/config/types.go +++ b/internal/config/types.go @@ -58,12 +58,43 @@ type SandboxConfig struct { MaxAutonomy string `json:"maxAutonomy,omitempty"` } +type ToolsConfig struct { + DeferThreshold int `json:"deferThreshold,omitempty"` + deferThresholdSet bool +} + +// ToolsOverride builds a ToolsConfig that explicitly overrides the deferred-tool +// threshold (including to 0, which disables deferral). Use this for programmatic +// Overrides — a bare ToolsConfig{DeferThreshold: 0} is indistinguishable from +// "unset" and will not override. +func ToolsOverride(deferThreshold int) ToolsConfig { + return ToolsConfig{DeferThreshold: deferThreshold, deferThresholdSet: true} +} + +func (cfg *ToolsConfig) UnmarshalJSON(data []byte) error { + type rawTools struct { + DeferThreshold *int `json:"deferThreshold"` + } + var raw rawTools + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + cfg.DeferThreshold = 0 + cfg.deferThresholdSet = false + if raw.DeferThreshold != nil { + cfg.DeferThreshold = *raw.DeferThreshold + cfg.deferThresholdSet = true + } + return nil +} + type FileConfig struct { ActiveProvider string `json:"activeProvider,omitempty"` Providers []ProviderProfile `json:"providers,omitempty"` MaxTurns int `json:"maxTurns,omitempty"` MCP MCPConfig `json:"mcp,omitempty"` Sandbox SandboxConfig `json:"sandbox,omitempty"` + Tools ToolsConfig `json:"tools,omitempty"` } type ResolveOptions struct { @@ -81,6 +112,7 @@ type Overrides struct { MaxTurns int MCP MCPConfig Sandbox SandboxConfig + Tools ToolsConfig } type ResolvedConfig struct { @@ -90,6 +122,7 @@ type ResolvedConfig struct { MaxTurns int MCP MCPConfig Sandbox SandboxConfig + Tools ToolsConfig } type MCPConfig struct { @@ -114,6 +147,7 @@ func (cfg *FileConfig) UnmarshalJSON(data []byte) error { MaxTurns int `json:"maxTurns"` MCP MCPConfig `json:"mcp"` Sandbox SandboxConfig `json:"sandbox"` + Tools ToolsConfig `json:"tools"` MCPServers map[string]MCPServerConfig `json:"mcpServers"` MCPServersSnake map[string]MCPServerConfig `json:"mcp_servers"` } @@ -127,6 +161,7 @@ func (cfg *FileConfig) UnmarshalJSON(data []byte) error { cfg.MaxTurns = raw.MaxTurns cfg.MCP = raw.MCP cfg.Sandbox = raw.Sandbox + cfg.Tools = raw.Tools if cfg.MCP.Servers == nil && (len(raw.MCPServers) > 0 || len(raw.MCPServersSnake) > 0) { cfg.MCP.Servers = map[string]MCPServerConfig{} } diff --git a/internal/config/types_test.go b/internal/config/types_test.go new file mode 100644 index 00000000..372ca989 --- /dev/null +++ b/internal/config/types_test.go @@ -0,0 +1,45 @@ +package config + +import ( + "encoding/json" + "testing" +) + +func TestToolsConfigJSONRoundTrip(t *testing.T) { + var cfg FileConfig + if err := json.Unmarshal([]byte(`{"tools":{"deferThreshold":25}}`), &cfg); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if cfg.Tools.DeferThreshold != 25 { + t.Fatalf("Tools.DeferThreshold = %d, want 25", cfg.Tools.DeferThreshold) + } + + encoded, err := json.Marshal(ToolsConfig{DeferThreshold: 7}) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + if string(encoded) != `{"deferThreshold":7}` { + t.Fatalf("Marshal() = %s, want {\"deferThreshold\":7}", encoded) + } + + // omitempty: a zero value must not emit the field. + emptyEncoded, err := json.Marshal(ToolsConfig{}) + if err != nil { + t.Fatalf("Marshal(empty) error = %v", err) + } + if string(emptyEncoded) != `{}` { + t.Fatalf("Marshal(empty) = %s, want {}", emptyEncoded) + } +} + +func TestToolsConfigPresentOnOverridesAndResolved(t *testing.T) { + // Compile-time guard that Overrides and ResolvedConfig carry the field too. + overrides := Overrides{Tools: ToolsConfig{DeferThreshold: 3}} + resolved := ResolvedConfig{Tools: ToolsConfig{DeferThreshold: 4}} + if overrides.Tools.DeferThreshold != 3 { + t.Fatalf("Overrides.Tools.DeferThreshold = %d, want 3", overrides.Tools.DeferThreshold) + } + if resolved.Tools.DeferThreshold != 4 { + t.Fatalf("ResolvedConfig.Tools.DeferThreshold = %d, want 4", resolved.Tools.DeferThreshold) + } +} diff --git a/internal/mcp/registry.go b/internal/mcp/registry.go index 68055130..1f2a9350 100644 --- a/internal/mcp/registry.go +++ b/internal/mcp/registry.go @@ -156,6 +156,22 @@ func (tool registryTool) Safety() tools.Safety { return tool.safety } +// Deferred marks every MCP tool as deferred-eligible: when many MCP tools are +// registered the agent loop may withhold their full schema and advertise them +// via tool_search. Built-in tools do not implement this interface and stay +// eager. +func (tool registryTool) Deferred() bool { + return true +} + +// MCPServerName reports the tool's originating MCP server name so the deferred- +// tools reminder labels it correctly, even when the sanitized server token in the +// synthesized tool name contains an underscore (which the name-only parser would +// truncate). It returns the true configured server name, not the sanitized token. +func (tool registryTool) MCPServerName() string { + return tool.server.Name +} + func (tool registryTool) Run(ctx context.Context, args map[string]any) tools.Result { result, err := tool.client.CallTool(ctx, tool.remote.Name, args) if err != nil { diff --git a/internal/mcp/registry_test.go b/internal/mcp/registry_test.go index 51193701..f7270162 100644 --- a/internal/mcp/registry_test.go +++ b/internal/mcp/registry_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "path/filepath" + "strings" "testing" "time" @@ -171,11 +172,75 @@ func TestRegisterToolsPreservesPriorRegistryStateOnFailure(t *testing.T) { } } +func TestRegistryToolIsDeferredEligible(t *testing.T) { + client := &fakeToolClient{} + tool := newRegistryTool( + Server{Name: "docs", Type: "stdio"}, + RemoteTool{Name: "lookup", Description: "Lookup documentation"}, + client, + RegisterOptions{}, + ) + + // Sanity: the wrapper synthesizes the expected sanitized name so the test + // exercises the real production path, not a hand-built struct. + if tool.Name() != "mcp_docs_lookup" { + t.Fatalf("registryTool.Name() = %q, want mcp_docs_lookup", tool.Name()) + } + + if !tool.Deferred() { + t.Fatal("registryTool.Deferred() = false, want true (all MCP tools are deferred-eligible)") + } + + // The exported helper in the tools package must agree via the optional + // interface, since the agent loop partitions tools through tools.IsDeferred. + if !tools.IsDeferred(tool) { + t.Fatal("tools.IsDeferred(registryTool) = false, want true") + } +} + +// TestRegistryToolReportsMCPServerName verifies the registryTool reports its true +// configured server name (not the sanitized tool-name token) so the deferred-tools +// reminder labels a multi-token server correctly via tools.DeferredLine. +func TestRegistryToolReportsMCPServerName(t *testing.T) { + client := &fakeToolClient{} + // A server name that sanitizes to a token containing an underscore ("git_hub"): + // the name-only parser would truncate the label to "git". + tool := newRegistryTool( + Server{Name: "git hub", Type: "stdio"}, + RemoteTool{Name: "create_issue", Description: "Create a GitHub issue."}, + client, + RegisterOptions{}, + ) + + if tool.MCPServerName() != "git hub" { + t.Fatalf("MCPServerName() = %q, want %q", tool.MCPServerName(), "git hub") + } + + // DeferredLine must prefer the reported server name over the name-derived token. + line := tools.DeferredLine(tool) + if !strings.Contains(line, "server: git hub") { + t.Fatalf("DeferredLine = %q, want it to label server as %q via MCPServerName()", line, "git hub") + } + // The truncated token-only label ("git") must NOT be the server segment. + if strings.Contains(line, "server: git |") { + t.Fatalf("DeferredLine = %q, mislabeled multi-token server with the truncated token", line) + } +} + +// TestRegistryToolNameRoundTripsToServerToken pins the synthesized name format +// for a single-token server so the fallback label parser (mcpServerFromToolName, +// exercised in the tools package) recovers the right server token. +func TestRegistryToolNameRoundTripsToServerToken(t *testing.T) { + if name := registryToolName("docs", "lookup"); name != "mcp_docs_lookup" { + t.Fatalf("registryToolName(docs, lookup) = %q, want mcp_docs_lookup", name) + } +} + type fakePreexistingTool struct { name string } -func (t *fakePreexistingTool) Name() string { return t.name } +func (t *fakePreexistingTool) Name() string { return t.name } func (t *fakePreexistingTool) Description() string { return "preexisting tool" } func (t *fakePreexistingTool) Parameters() tools.Schema { return tools.Schema{} } func (t *fakePreexistingTool) Safety() tools.Safety { return tools.Safety{} } diff --git a/internal/tools/deferred.go b/internal/tools/deferred.go new file mode 100644 index 00000000..f2f2f629 --- /dev/null +++ b/internal/tools/deferred.go @@ -0,0 +1,209 @@ +package tools + +import ( + "fmt" + "regexp" + "sort" + "strings" +) + +const ( + defaultShortenMax = 200 + maxSchemaHintParams = 4 + maxSchemaHintLen = 360 +) + +var ( + headingPrefix = regexp.MustCompile(`^#{1,6}\s+(.+)$`) + genericHeading = regexp.MustCompile(`(?i)^(overview|description|summary)$`) + collapseWhitespace = regexp.MustCompile(`\s+`) +) + +// normalizeDescriptionLine trims a line and unwraps a leading markdown heading. +func normalizeDescriptionLine(line string) string { + trimmed := strings.TrimSpace(line) + if m := headingPrefix.FindStringSubmatch(trimmed); m != nil { + return strings.TrimSpace(m[1]) + } + return trimmed +} + +func isGenericDescriptionHeading(line string) bool { + return genericHeading.MatchString(line) +} + +// truncateDescription clips desc to at most max runes, preferring a word +// boundary and appending a single-rune ellipsis when it had to cut. +func truncateDescription(desc string, max int) string { + runes := []rune(desc) + if max <= 0 || len(runes) <= max { + return desc + } + cut := string(runes[:max-1]) + if idx := strings.LastIndexByte(cut, ' '); idx > 0 { + cut = cut[:idx] + } + return strings.TrimRight(cut, " ") + "…" +} + +// shortenDescription reduces desc to a single meaningful line, collapses +// whitespace, and truncates to at most max runes with an ellipsis. +func shortenDescription(desc string, max int) string { + if desc == "" { + return "" + } + if max <= 0 { + max = defaultShortenMax + } + var lines []string + for _, raw := range strings.Split(desc, "\n") { + if line := normalizeDescriptionLine(raw); line != "" { + lines = append(lines, collapseWhitespace.ReplaceAllString(line, " ")) + } + } + if len(lines) == 0 { + return "" + } + meaningful := lines[0] + if isGenericDescriptionHeading(meaningful) && len(lines) > 1 { + meaningful = meaningful + " — " + lines[1] + } + return truncateDescription(meaningful, max) +} + +// formatInputSchemaHint renders a one-line summary of a tool's input schema, +// e.g. "inputs (* required): a (string)*, b (number); +N more". Property names +// are sorted for deterministic output (Schema.Properties is a map). Returns +// "(none)" when the schema declares no properties. At most maxSchemaHintParams +// params are shown; the rest are summarized as "; +N more". +func formatInputSchemaHint(schema Schema) string { + if len(schema.Properties) == 0 { + return "(none)" + } + + names := make([]string, 0, len(schema.Properties)) + for name := range schema.Properties { + names = append(names, name) + } + sort.Strings(names) + + required := make(map[string]bool, len(schema.Required)) + for _, name := range schema.Required { + required[name] = true + } + + shown := names + if len(shown) > maxSchemaHintParams { + shown = shown[:maxSchemaHintParams] + } + + parts := make([]string, 0, len(shown)) + for _, name := range shown { + prop := schema.Properties[name] + marker := "" + if required[name] { + marker = "*" + } + typePart := "" + if t := strings.TrimSpace(prop.Type); t != "" { + typePart = " (" + t + ")" + } + parts = append(parts, name+typePart+marker) + } + + more := "" + if len(names) > maxSchemaHintParams { + more = fmt.Sprintf("; +%d more", len(names)-maxSchemaHintParams) + } + + hint := "inputs (* required): " + strings.Join(parts, ", ") + more + return shortenDescription(hint, maxSchemaHintLen) +} + +// mcpToolNamePrefix mirrors the prefix used by mcp.registryToolName. +const mcpToolNamePrefix = "mcp_" + +// mcpServerFromToolName extracts the server token from a synthesized MCP tool +// name produced by mcp.registryToolName ("mcp__"). It returns "" +// for non-MCP names and for names that lack both a server and a tool segment. +func mcpServerFromToolName(name string) string { + rest, ok := strings.CutPrefix(name, mcpToolNamePrefix) + if !ok { + return "" + } + sep := strings.IndexByte(rest, '_') + if sep <= 0 || sep == len(rest)-1 { + // No server token, or nothing after the server token (no tool part). + return "" + } + return rest[:sep] +} + +// formatDeferredToolLine renders a single compact advertisement line for a +// deferred tool: "name: | server: | ". The +// "server: " segment is omitted when server is empty (non-MCP tools). +func formatDeferredToolLine(name, description, server string, schema Schema) string { + desc := shortenDescription(description, defaultShortenMax) + if desc == "" { + desc = "No description provided" + } + parts := []string{name + ": " + desc} + if server != "" { + parts = append(parts, "server: "+server) + } + parts = append(parts, formatInputSchemaHint(schema)) + return strings.Join(parts, " | ") +} + +// mcpServerNamed is an optional interface a deferred MCP tool implements to +// report its true (un-sanitized-token) server name for the reminder label. When +// a tool provides it, DeferredLine prefers it over the name-derived token, which +// would mislabel a server whose sanitized name itself contains an underscore +// (e.g. "git_hub" → "git"). It affects the cosmetic reminder label only; tool +// resolution never depends on this. +type mcpServerNamed interface { + MCPServerName() string +} + +// DeferredLine renders the compact advertisement line for a single deferred +// tool, deriving the MCP server label from the tool's reported server name when +// available, falling back to the token parsed from the tool's name. It is the +// exported entry point the agent loop uses to build each line of the +// deferred-tools reminder, so callers in other packages never touch the +// unexported formatters. +func DeferredLine(t Tool) string { + server := mcpServerFromToolName(t.Name()) + if named, ok := t.(mcpServerNamed); ok { + if reported := strings.TrimSpace(named.MCPServerName()); reported != "" { + server = reported + } + } + return formatDeferredToolLine(t.Name(), t.Description(), server, t.Parameters()) +} + +const ( + systemReminderStart = "" + systemReminderEnd = "" +) + +// BuildDeferredReminder wraps the given deferred-tool advertisement lines in a +// block instructing the model to load a tool's full schema +// via tool_search (using a `select:Name1,Name2` exact query or keywords) +// before calling it. Returns "" when there are no lines. +func BuildDeferredReminder(lines []string) string { + if len(lines) == 0 { + return "" + } + var b strings.Builder + b.WriteString(systemReminderStart) + b.WriteString("\n") + b.WriteString("The tools listed below are available in this environment, but their full schemas are omitted from the current tool list to save context.\n") + b.WriteString(`Call tool_search with query "select:Name1,Name2" (the exact names before each colon) or with keywords to load a tool's full schema before calling it. Calling an omitted tool directly will fail with an input validation error.` + "\n") + b.WriteString("Do not call tool_search for tools already present in the current tool list, and do not guess or invent tool names.\n") + b.WriteString("\n") + b.WriteString("Deferred tools:\n") + b.WriteString(strings.Join(lines, "\n")) + b.WriteString("\n") + b.WriteString(systemReminderEnd) + return b.String() +} diff --git a/internal/tools/deferred_test.go b/internal/tools/deferred_test.go new file mode 100644 index 00000000..9c0dca99 --- /dev/null +++ b/internal/tools/deferred_test.go @@ -0,0 +1,340 @@ +package tools + +import ( + "context" + "strings" + "testing" +) + +func TestShortenDescription(t *testing.T) { + tests := []struct { + name string + desc string + max int + want string + }{ + {name: "empty", desc: "", max: 200, want: ""}, + {name: "single short line", desc: "Reads a file from disk.", max: 200, want: "Reads a file from disk."}, + { + name: "takes first non-empty line", + desc: "\n\nReads a file.\nSecond line ignored.", + max: 200, + want: "Reads a file.", + }, + { + name: "strips markdown heading", + desc: "## Read tool\nbody", + max: 200, + want: "Read tool", + }, + { + name: "generic heading joins next line", + desc: "Overview\nFetches a URL over HTTP.", + max: 200, + want: "Overview — Fetches a URL over HTTP.", + }, + { + name: "collapses internal whitespace", + desc: "Runs a\tshell command", + max: 200, + want: "Runs a shell command", + }, + { + name: "truncates on word boundary with ellipsis", + desc: "alpha beta gamma delta epsilon zeta", + max: 20, + want: "alpha beta gamma…", + }, + { + name: "truncates mid-word when no boundary", + desc: "supercalifragilistic", + max: 10, + want: "supercali…", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := shortenDescription(tc.desc, tc.max) + if got != tc.want { + t.Fatalf("shortenDescription(%q, %d) = %q, want %q", tc.desc, tc.max, got, tc.want) + } + if len([]rune(got)) > tc.max && tc.max > 0 { + t.Fatalf("result %q exceeds max %d runes", got, tc.max) + } + if strings.Contains(got, "\n") { + t.Fatalf("result %q contains a newline", got) + } + }) + } +} + +func TestFormatInputSchemaHint(t *testing.T) { + tests := []struct { + name string + schema Schema + want string + }{ + { + name: "no properties", + schema: Schema{Type: "object"}, + want: "(none)", + }, + { + name: "required marked and sorted", + schema: Schema{ + Type: "object", + Properties: map[string]PropertySchema{ + "a": {Type: "string"}, + "b": {Type: "number"}, + }, + Required: []string{"a"}, + }, + want: "inputs (* required): a (string)*, b (number)", + }, + { + name: "untyped property omits type parens", + schema: Schema{ + Type: "object", + Properties: map[string]PropertySchema{ + "q": {}, + }, + Required: []string{"q"}, + }, + want: "inputs (* required): q*", + }, + { + name: "caps at four and reports remainder", + schema: Schema{ + Type: "object", + Properties: map[string]PropertySchema{ + "a": {Type: "string"}, + "b": {Type: "string"}, + "c": {Type: "string"}, + "d": {Type: "string"}, + "e": {Type: "string"}, + "f": {Type: "string"}, + }, + Required: []string{"a"}, + }, + want: "inputs (* required): a (string)*, b (string), c (string), d (string); +2 more", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := formatInputSchemaHint(tc.schema); got != tc.want { + t.Fatalf("formatInputSchemaHint() = %q, want %q", got, tc.want) + } + }) + } +} + +func TestMCPServerFromToolName(t *testing.T) { + tests := []struct { + name string + toolName string + want string + }{ + {name: "standard mcp name", toolName: "mcp_github_create_issue", want: "github"}, + {name: "single tool segment", toolName: "mcp_weather_forecast", want: "weather"}, + {name: "server with no tool part", toolName: "mcp_github_tool", want: "github"}, + {name: "non-mcp builtin", toolName: "read_file", want: ""}, + {name: "tool_search itself", toolName: "tool_search", want: ""}, + {name: "prefix only no server", toolName: "mcp_", want: ""}, + {name: "prefix and server but no tool", toolName: "mcp_github", want: ""}, + {name: "empty", toolName: "", want: ""}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := mcpServerFromToolName(tc.toolName); got != tc.want { + t.Fatalf("mcpServerFromToolName(%q) = %q, want %q", tc.toolName, got, tc.want) + } + }) + } +} + +func TestFormatDeferredToolLine(t *testing.T) { + schema := Schema{ + Type: "object", + Properties: map[string]PropertySchema{ + "url": {Type: "string"}, + }, + Required: []string{"url"}, + } + + t.Run("with server", func(t *testing.T) { + got := formatDeferredToolLine("mcp_github_fetch", "Fetches a URL.", "github", schema) + want := "mcp_github_fetch: Fetches a URL. | server: github | inputs (* required): url (string)*" + if got != want { + t.Fatalf("got %q, want %q", got, want) + } + }) + + t.Run("omits server when empty", func(t *testing.T) { + got := formatDeferredToolLine("read_file", "Reads a file.", "", schema) + want := "read_file: Reads a file. | inputs (* required): url (string)*" + if got != want { + t.Fatalf("got %q, want %q", got, want) + } + if strings.Contains(got, "server:") { + t.Fatalf("server part leaked into %q", got) + } + }) + + t.Run("empty schema shows none", func(t *testing.T) { + got := formatDeferredToolLine("ping", "Pings.", "", Schema{Type: "object"}) + want := "ping: Pings. | (none)" + if got != want { + t.Fatalf("got %q, want %q", got, want) + } + }) + + t.Run("blank description gets placeholder", func(t *testing.T) { + got := formatDeferredToolLine("noop", "", "", Schema{Type: "object"}) + want := "noop: No description provided | (none)" + if got != want { + t.Fatalf("got %q, want %q", got, want) + } + }) + + t.Run("long description truncated", func(t *testing.T) { + long := strings.Repeat("word ", 80) // 400 chars + got := formatDeferredToolLine("big", long, "", Schema{Type: "object"}) + if !strings.Contains(got, "…") { + t.Fatalf("expected truncation ellipsis in %q", got) + } + descPart := strings.SplitN(strings.TrimPrefix(got, "big: "), " | ", 2)[0] + if len([]rune(descPart)) > defaultShortenMax { + t.Fatalf("description part %d runes exceeds %d", len([]rune(descPart)), defaultShortenMax) + } + }) +} + +func TestBuildDeferredReminder(t *testing.T) { + t.Run("no lines returns empty", func(t *testing.T) { + if got := BuildDeferredReminder(nil); got != "" { + t.Fatalf("BuildDeferredReminder(nil) = %q, want empty", got) + } + if got := BuildDeferredReminder([]string{}); got != "" { + t.Fatalf("BuildDeferredReminder([]) = %q, want empty", got) + } + }) + + t.Run("wraps lines in system-reminder", func(t *testing.T) { + lines := []string{ + "mcp_github_fetch: Fetches a URL. | server: github | inputs (* required): url (string)*", + "mcp_weather_now: Current weather. | server: weather | (none)", + } + got := BuildDeferredReminder(lines) + + if !strings.HasPrefix(got, "") { + t.Fatalf("missing opening tag: %q", got) + } + if !strings.HasSuffix(got, "") { + t.Fatalf("missing closing tag: %q", got) + } + if !strings.Contains(got, "tool_search") { + t.Fatalf("reminder must name the tool_search tool: %q", got) + } + if !strings.Contains(got, `select:`) { + t.Fatalf("reminder must explain select: query form: %q", got) + } + for _, line := range lines { + if !strings.Contains(got, line) { + t.Fatalf("reminder missing line %q in %q", line, got) + } + } + // Lines are newline-joined, not concatenated. + if !strings.Contains(got, lines[0]+"\n"+lines[1]) { + t.Fatalf("lines not newline-joined in %q", got) + } + }) +} + +func TestDeferredLineUsesToolFieldsAndServer(t *testing.T) { + tool := fakeDeferredTool{ + baseTool: baseTool{ + name: "mcp_docs_lookup", + description: "Look up documentation for a symbol.", + parameters: Schema{ + Type: "object", + Properties: map[string]PropertySchema{"symbol": {Type: "string"}}, + Required: []string{"symbol"}, + }, + }, + deferred: true, + } + got := DeferredLine(tool) + for _, want := range []string{"mcp_docs_lookup: Look up documentation", "server: docs", "inputs (* required): symbol (string)*"} { + if !strings.Contains(got, want) { + t.Fatalf("DeferredLine = %q, missing %q", got, want) + } + } +} + +// serverNamedTool implements the optional mcpServerNamed interface so DeferredLine +// can prefer the reported server name over the name-derived token. +type serverNamedTool struct { + baseTool + server string +} + +func (t serverNamedTool) Run(context.Context, map[string]any) Result { return okResult("ok") } +func (t serverNamedTool) Deferred() bool { return true } +func (t serverNamedTool) MCPServerName() string { return t.server } + +func TestDeferredLinePrefersMCPServerName(t *testing.T) { + // The tool name sanitizes to a server token containing an underscore + // ("git_hub"): mcpServerFromToolName would truncate the label to "git", so the + // reported server name must win. + tool := serverNamedTool{ + baseTool: baseTool{ + name: "mcp_git_hub_create_issue", + description: "Create a GitHub issue.", + parameters: Schema{Type: "object"}, + }, + server: "git hub", + } + got := DeferredLine(tool) + if !strings.Contains(got, "server: git hub") { + t.Fatalf("DeferredLine = %q, want server label from MCPServerName() (%q)", got, "git hub") + } + if strings.Contains(got, "server: git |") { + t.Fatalf("DeferredLine = %q, used truncated name-derived token instead of MCPServerName()", got) + } +} + +// A blank/whitespace MCPServerName() must fall back to the name-derived token so a +// tool that reports nothing useful is still labeled from its synthesized name. +func TestDeferredLineFallsBackWhenMCPServerNameBlank(t *testing.T) { + tool := serverNamedTool{ + baseTool: baseTool{ + name: "mcp_docs_lookup", + description: "Look up docs.", + parameters: Schema{Type: "object"}, + }, + server: " ", + } + got := DeferredLine(tool) + if !strings.Contains(got, "server: docs") { + t.Fatalf("DeferredLine = %q, want fallback server token %q", got, "docs") + } +} + +// TestMCPServerFromToolNameRoundTrip pins the fallback parser against the name +// format produced for a single-token MCP server ("mcp__"), the same +// format mcp.registryToolName synthesizes. +func TestMCPServerFromToolNameRoundTrip(t *testing.T) { + cases := map[string]string{ + "mcp_docs_lookup": "docs", + "mcp_weather_forecast": "weather", + "mcp_github_create_issue": "github", + } + for name, wantServer := range cases { + if got := mcpServerFromToolName(name); got != wantServer { + t.Fatalf("mcpServerFromToolName(%q) = %q, want %q", name, got, wantServer) + } + } +} diff --git a/internal/tools/registry.go b/internal/tools/registry.go index d856faf9..db5313a7 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -23,6 +23,13 @@ type RunOptions struct { ReasoningEffort string Depth int Cwd string + // EnabledTools / DisabledTools carry the run's operator tool filters so a + // filter-aware tool (tool_search) never discloses or loads an operator-hidden + // tool. They use the same allow/deny semantics as the agent's filter gate: + // denied if listed in DisabledTools; if EnabledTools is non-empty, allowed + // only when listed in it. + EnabledTools []string + DisabledTools []string } type sandboxAwareTool interface { @@ -33,10 +40,25 @@ type optionsAwareTool interface { RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result } +// deferredTool is an optional interface a tool implements to mark itself +// deferred-eligible: when many such tools are registered, the agent loop may +// withhold their full schema and advertise them via tool_search instead. +type deferredTool interface { + Deferred() bool +} + func NewRegistry() *Registry { return &Registry{tools: make(map[string]Tool)} } +// IsDeferred reports whether a tool opts into deferred loading. A tool is +// deferred-eligible only if it implements deferredTool and its Deferred() +// method returns true; tools that do not implement the interface are eager. +func IsDeferred(t Tool) bool { + d, ok := t.(deferredTool) + return ok && d.Deferred() +} + func (registry *Registry) Register(tool Tool) { registry.tools[tool.Name()] = tool } diff --git a/internal/tools/registry_test.go b/internal/tools/registry_test.go index 418d7796..9183b1b1 100644 --- a/internal/tools/registry_test.go +++ b/internal/tools/registry_test.go @@ -329,3 +329,47 @@ func TestRunWithOptionsLeavesCleanOutputUnchanged(t *testing.T) { t.Fatalf("clean output altered: redacted=%v output=%q", res.Redacted, res.Output) } } + +// fakeDeferredTool implements the optional Deferred() interface and reports +// itself as deferred-eligible; fakePlainTool does not implement it. +type fakeDeferredTool struct { + baseTool + deferred bool +} + +func (tool fakeDeferredTool) Run(context.Context, map[string]any) Result { + return okResult("ok") +} + +func (tool fakeDeferredTool) Deferred() bool { return tool.deferred } + +type fakePlainTool struct { + baseTool +} + +func (tool fakePlainTool) Run(context.Context, map[string]any) Result { + return okResult("ok") +} + +func TestIsDeferredReportsOptionalInterface(t *testing.T) { + eligible := fakeDeferredTool{ + baseTool: baseTool{name: "eligible", description: "deferred-eligible tool"}, + deferred: true, + } + if !IsDeferred(eligible) { + t.Fatal("IsDeferred(eligible) = false, want true for a tool whose Deferred() returns true") + } + + notEligible := fakeDeferredTool{ + baseTool: baseTool{name: "not_eligible", description: "implements interface but opts out"}, + deferred: false, + } + if IsDeferred(notEligible) { + t.Fatal("IsDeferred(notEligible) = true, want false when Deferred() returns false") + } + + plain := fakePlainTool{baseTool: baseTool{name: "plain", description: "no deferred interface"}} + if IsDeferred(plain) { + t.Fatal("IsDeferred(plain) = true, want false for a tool that does not implement deferredTool") + } +} diff --git a/internal/tools/tool_search.go b/internal/tools/tool_search.go new file mode 100644 index 00000000..18d074a8 --- /dev/null +++ b/internal/tools/tool_search.go @@ -0,0 +1,253 @@ +package tools + +import ( + "context" + "encoding/json" + "sort" + "strconv" + "strings" +) + +// ToolSearchToolName is the canonical registry name of the deferred-tool loader. +// It is exported so the agent loop, CLI filter validation, and the partition can +// reference the loader without re-stating the literal (keeping all the gates that +// must treat it specially in agreement). +const ToolSearchToolName = "tool_search" + +// toolSearchMaxKeywordMatches caps how many deferred tools a bare-keyword query +// loads in one call, keeping the returned schemas bounded. +const toolSearchMaxKeywordMatches = 10 + +// toolSearchTool lets the model pull a deferred tool's full schema on demand. +// It holds the live registry (like escalate_model holds the model registry) so +// it can resolve names against the currently registered deferred-eligible tools. +type toolSearchTool struct { + baseTool + registry *Registry +} + +// NewToolSearchTool builds the tool_search tool over the given registry. The +// tool is informational/no-side-effect and is advertised even in auto mode so +// the model can always discover withheld tools. +func NewToolSearchTool(registry *Registry) Tool { + return toolSearchTool{ + baseTool: baseTool{ + name: ToolSearchToolName, + description: "Search for and load deferred tools by exact name or keyword. Use query \"select:Name1,Name2\" to load specific tools, or keywords to find matching tools. Loading a tool returns its full schema so you can call it on the next turn.", + parameters: Schema{ + Type: "object", + Properties: map[string]PropertySchema{ + "query": { + Type: "string", + Description: "Either \"select:Name1,Name2\" for exact tool names, or space-separated keywords to match tool names and descriptions.", + }, + }, + Required: []string{"query"}, + AdditionalProperties: false, + }, + safety: Safety{ + SideEffect: SideEffectNone, + Permission: PermissionAllow, + Reason: "Lists and loads already-registered tool schemas; performs no side effects.", + AdvertiseInAuto: true, + }, + }, + registry: registry, + } +} + +// Run satisfies the Tool interface; actual dispatch goes through RunWithOptions. +func (tool toolSearchTool) Run(ctx context.Context, args map[string]any) Result { + return tool.RunWithOptions(ctx, args, RunOptions{}) +} + +func (tool toolSearchTool) RunWithOptions(_ context.Context, args map[string]any, options RunOptions) Result { + query, err := stringArg(args, "query", "", true) + if err != nil { + return errorResult("Error: Invalid arguments for tool_search: " + err.Error()) + } + query = strings.TrimSpace(query) + + // Honor the run's operator filters so an operator-hidden deferred tool is + // invisible to tool_search: it never resolves via select:, never ranks for a + // keyword query, and is omitted from the no-match listing. + deferred := tool.visibleDeferredTools(options.EnabledTools, options.DisabledTools) + + var matches []Tool + if rest, ok := strings.CutPrefix(query, "select:"); ok { + matches = tool.resolveExact(rest, deferred) + } else { + matches = tool.rankByKeyword(query, deferred) + } + + if len(matches) == 0 { + return okResult(tool.noMatchMessage(query, deferred)) + } + + names := make([]string, 0, len(matches)) + for _, match := range matches { + names = append(names, match.Name()) + } + return Result{ + Status: StatusOK, + Output: renderLoadedTools(matches), + Meta: map[string]string{"load_tools": strings.Join(names, ",")}, + } +} + +// visibleDeferredTools returns the registry's deferred-eligible tools that pass +// the operator allow/deny filters, sorted by name so keyword ranking and +// listings stay deterministic. A nil/empty filter pair admits every deferred +// tool (the pre-filter behavior). The allow/deny semantics mirror the agent's +// ToolAllowedByFilters: denied if in disabled; if enabled is non-empty, the tool +// must be listed in it. +func (tool toolSearchTool) visibleDeferredTools(enabled []string, disabled []string) []Tool { + var deferred []Tool + if tool.registry != nil { + for _, candidate := range tool.registry.All() { + if !IsDeferred(candidate) { + continue + } + if !toolAllowedByFilters(candidate.Name(), enabled, disabled) { + continue + } + deferred = append(deferred, candidate) + } + } + sort.Slice(deferred, func(left, right int) bool { + return deferred[left].Name() < deferred[right].Name() + }) + return deferred +} + +// toolAllowedByFilters mirrors agent.ToolAllowedByFilters (kept here to avoid an +// import cycle: the agent package imports tools, not the other way around). A +// name is denied when it appears in disabled; when enabled is non-empty, the +// name must appear in it to be allowed. +func toolAllowedByFilters(name string, enabled []string, disabled []string) bool { + if len(enabled) > 0 && !containsName(enabled, name) { + return false + } + return !containsName(disabled, name) +} + +func containsName(names []string, name string) bool { + for _, candidate := range names { + if candidate == name { + return true + } + } + return false +} + +// resolveExact maps a comma-separated name list (the part after "select:") to +// the matching deferred tools, preserving the model's order and skipping blanks +// and unknown names. +func (tool toolSearchTool) resolveExact(list string, deferred []Tool) []Tool { + byName := make(map[string]Tool, len(deferred)) + for _, candidate := range deferred { + byName[candidate.Name()] = candidate + } + var matches []Tool + seen := make(map[string]bool) + for _, raw := range strings.Split(list, ",") { + name := strings.TrimSpace(raw) + if name == "" || seen[name] { + continue + } + if candidate, ok := byName[name]; ok { + seen[name] = true + matches = append(matches, candidate) + } + } + return matches +} + +// rankByKeyword scores deferred tools by case-insensitive substring match on the +// name (weighted higher) then the description, and returns the top matches. +func (tool toolSearchTool) rankByKeyword(query string, deferred []Tool) []Tool { + keywords := strings.Fields(strings.ToLower(query)) + if len(keywords) == 0 { + return nil + } + type scored struct { + tool Tool + score int + order int + } + var ranked []scored + for index, candidate := range deferred { + name := strings.ToLower(candidate.Name()) + desc := strings.ToLower(candidate.Description()) + score := 0 + for _, keyword := range keywords { + if strings.Contains(name, keyword) { + score += 2 + } + if strings.Contains(desc, keyword) { + score++ + } + } + if score > 0 { + ranked = append(ranked, scored{tool: candidate, score: score, order: index}) + } + } + sort.SliceStable(ranked, func(left, right int) bool { + if ranked[left].score != ranked[right].score { + return ranked[left].score > ranked[right].score + } + return ranked[left].order < ranked[right].order + }) + if len(ranked) > toolSearchMaxKeywordMatches { + ranked = ranked[:toolSearchMaxKeywordMatches] + } + matches := make([]Tool, 0, len(ranked)) + for _, entry := range ranked { + matches = append(matches, entry.tool) + } + return matches +} + +// noMatchMessage reports that nothing loaded and names the available deferred +// tools so the model can retry with a valid select: query. +func (tool toolSearchTool) noMatchMessage(query string, deferred []Tool) string { + if len(deferred) == 0 { + return "No deferred tools are available to load." + } + names := make([]string, 0, len(deferred)) + for _, candidate := range deferred { + names = append(names, candidate.Name()) + } + return "No tools matched \"" + query + "\". Available tools: " + strings.Join(names, ", ") + + `. Retry with query "select:Name1,Name2" using exact names.` +} + +// renderLoadedTools lists each loaded tool's name, full description, and full +// input schema (pretty-printed JSON) so the model has the complete spec inline. +func renderLoadedTools(matches []Tool) string { + var builder strings.Builder + builder.WriteString("Loaded ") + if len(matches) == 1 { + builder.WriteString("1 tool") + } else { + builder.WriteString(strconv.Itoa(len(matches))) + builder.WriteString(" tools") + } + builder.WriteString(". Full schemas follow; call them on the next turn.\n") + for _, match := range matches { + builder.WriteString("\n## ") + builder.WriteString(match.Name()) + builder.WriteString("\n") + builder.WriteString(match.Description()) + builder.WriteString("\n") + schemaJSON, err := json.MarshalIndent(match.Parameters(), "", " ") + if err != nil { + builder.WriteString("(schema unavailable)\n") + continue + } + builder.WriteString("input schema:\n") + builder.Write(schemaJSON) + builder.WriteString("\n") + } + return builder.String() +} diff --git a/internal/tools/tool_search_test.go b/internal/tools/tool_search_test.go new file mode 100644 index 00000000..a4270a31 --- /dev/null +++ b/internal/tools/tool_search_test.go @@ -0,0 +1,299 @@ +package tools + +import ( + "context" + "strings" + "testing" +) + +// searchFakeTool is a minimal deferred-eligible tool for tool_search tests: +// it implements Tool plus the optional Deferred() bool that IsDeferred checks. +type searchFakeTool struct { + name string + description string + parameters Schema +} + +func (t searchFakeTool) Name() string { return t.name } +func (t searchFakeTool) Description() string { return t.description } +func (t searchFakeTool) Parameters() Schema { return t.parameters } +func (t searchFakeTool) Safety() Safety { + return Safety{SideEffect: SideEffectRead, Permission: PermissionAllow} +} +func (t searchFakeTool) Run(context.Context, map[string]any) Result { + return Result{Status: StatusOK} +} +func (t searchFakeTool) Deferred() bool { return true } + +func newDeferredFixtureRegistry() *Registry { + reg := NewRegistry() + reg.Register(searchFakeTool{ + name: "weather_lookup", + description: "Look up the current weather for a city.", + parameters: Schema{ + Type: "object", + Properties: map[string]PropertySchema{ + "city": {Type: "string", Description: "City name to look up."}, + }, + Required: []string{"city"}, + AdditionalProperties: false, + }, + }) + reg.Register(searchFakeTool{ + name: "stock_quote", + description: "Fetch a stock quote for a ticker symbol.", + parameters: Schema{ + Type: "object", + Properties: map[string]PropertySchema{"ticker": {Type: "string"}}, + Required: []string{"ticker"}, + AdditionalProperties: false, + }, + }) + return reg +} + +func TestToolSearchExposesExpectedSafetyAndSchema(t *testing.T) { + tool := NewToolSearchTool(NewRegistry()) + + if tool.Name() != "tool_search" { + t.Fatalf("name = %q, want tool_search", tool.Name()) + } + if tool.Description() == "" { + t.Fatal("tool_search must have a description") + } + + safety := tool.Safety() + if safety.SideEffect != SideEffectNone { + t.Fatalf("side effect = %s, want none", safety.SideEffect) + } + if safety.Permission != PermissionAllow { + t.Fatalf("permission = %s, want allow", safety.Permission) + } + if !safety.AdvertiseInAuto { + t.Fatal("tool_search must be AdvertiseInAuto") + } + + schema := tool.Parameters() + if schema.Type != "object" || schema.AdditionalProperties { + t.Fatalf("unexpected schema header: %#v", schema) + } + queryProp, ok := schema.Properties["query"] + if !ok { + t.Fatal("schema missing query property") + } + if queryProp.Type != "string" { + t.Fatalf("query type = %s, want string", queryProp.Type) + } + if len(schema.Required) != 1 || schema.Required[0] != "query" { + t.Fatalf("required = %#v, want [query]", schema.Required) + } +} + +// tool_search must run through the registry's optionsAwareTool dispatch and be +// allowed without a permission grant (SideEffectNone + PermissionAllow). +func TestToolSearchRunsThroughRegistryWithoutPermission(t *testing.T) { + reg := newDeferredFixtureRegistry() + reg.Register(NewToolSearchTool(reg)) + + result := reg.Run(context.Background(), "tool_search", map[string]any{ + "query": "select:weather_lookup", + }) + + if result.Status != StatusOK { + t.Fatalf("status = %s, want ok; output=%q", result.Status, result.Output) + } + if result.Meta["load_tools"] != "weather_lookup" { + t.Fatalf("Meta[load_tools] = %q, want weather_lookup", result.Meta["load_tools"]) + } +} + +func TestToolSearchMissingQueryArgIsError(t *testing.T) { + tool := NewToolSearchTool(NewRegistry()).(optionsAwareTool) + result := tool.RunWithOptions(context.Background(), map[string]any{}, RunOptions{}) + if result.Status != StatusError { + t.Fatalf("status = %s, want error for missing query", result.Status) + } + if !strings.Contains(result.Output, "query is required") { + t.Fatalf("unexpected error output: %q", result.Output) + } +} + +func TestToolSearchUnknownQueryReturnsNoMeta(t *testing.T) { + reg := newDeferredFixtureRegistry() + tool := NewToolSearchTool(reg).(optionsAwareTool) + + for _, query := range []string{"select:does_not_exist", "select:", "zzznomatch"} { + result := tool.RunWithOptions(context.Background(), + map[string]any{"query": query}, RunOptions{}) + + if result.Status != StatusOK { + t.Fatalf("query %q: status = %s, want ok (informational)", query, result.Status) + } + if _, present := result.Meta["load_tools"]; present { + t.Fatalf("query %q: must NOT set load_tools, got %q", query, result.Meta["load_tools"]) + } + // Informational message should name the available tools so the model can retry. + if !strings.Contains(result.Output, "weather_lookup") || !strings.Contains(result.Output, "stock_quote") { + t.Fatalf("query %q: message must name available tools, got %q", query, result.Output) + } + } +} + +func TestToolSearchEmptyRegistryReportsNothingAvailable(t *testing.T) { + tool := NewToolSearchTool(NewRegistry()).(optionsAwareTool) + + result := tool.RunWithOptions(context.Background(), + map[string]any{"query": "select:anything"}, RunOptions{}) + + if result.Status != StatusOK { + t.Fatalf("status = %s, want ok", result.Status) + } + if _, present := result.Meta["load_tools"]; present { + t.Fatalf("empty registry must not set load_tools, got %q", result.Meta["load_tools"]) + } + if !strings.Contains(result.Output, "No deferred tools are available") { + t.Fatalf("unexpected message: %q", result.Output) + } +} + +func TestToolSearchKeywordRanksByNameThenDescription(t *testing.T) { + reg := NewRegistry() + // name match should outrank a description-only match. + reg.Register(searchFakeTool{ + name: "weather_lookup", + description: "Look up the current weather for a city.", + parameters: Schema{Type: "object", AdditionalProperties: false}, + }) + reg.Register(searchFakeTool{ + name: "forecast_report", + description: "Generates a multi-day weather outlook.", + parameters: Schema{Type: "object", AdditionalProperties: false}, + }) + tool := NewToolSearchTool(reg).(optionsAwareTool) + + result := tool.RunWithOptions(context.Background(), + map[string]any{"query": "weather"}, RunOptions{}) + + if result.Status != StatusOK { + t.Fatalf("status = %s, want ok", result.Status) + } + loaded := result.Meta["load_tools"] + // Both match "weather"; the name match (weather_lookup) must come first. + if loaded != "weather_lookup,forecast_report" { + t.Fatalf("load_tools = %q, want weather_lookup,forecast_report (name match ranked first)", loaded) + } +} + +func TestToolSearchKeywordExcludesNonMatches(t *testing.T) { + reg := newDeferredFixtureRegistry() + tool := NewToolSearchTool(reg).(optionsAwareTool) + + result := tool.RunWithOptions(context.Background(), + map[string]any{"query": "stock"}, RunOptions{}) + + if got := result.Meta["load_tools"]; got != "stock_quote" { + t.Fatalf("load_tools = %q, want only stock_quote", got) + } + if strings.Contains(result.Output, "weather_lookup") { + t.Fatalf("non-matching tool weather_lookup leaked into output: %q", result.Output) + } +} + +func TestToolSearchSelectLoadsExactNames(t *testing.T) { + reg := newDeferredFixtureRegistry() + tool := NewToolSearchTool(reg) + + optioned, ok := tool.(optionsAwareTool) + if !ok { + t.Fatalf("tool_search must implement optionsAwareTool") + } + + result := optioned.RunWithOptions(context.Background(), + map[string]any{"query": "select:weather_lookup,stock_quote"}, RunOptions{}) + + if result.Status != StatusOK { + t.Fatalf("status = %s, want ok; output=%q", result.Status, result.Output) + } + if got := result.Meta["load_tools"]; got != "weather_lookup,stock_quote" { + t.Fatalf("Meta[load_tools] = %q, want %q", got, "weather_lookup,stock_quote") + } + // Output must carry the FULL description and full schema (not a compact line). + if !strings.Contains(result.Output, "Look up the current weather for a city.") { + t.Fatalf("output missing full description: %q", result.Output) + } + if !strings.Contains(result.Output, "weather_lookup") || !strings.Contains(result.Output, "\"city\"") { + t.Fatalf("output missing full schema for weather_lookup: %q", result.Output) + } + if !strings.Contains(result.Output, "stock_quote") || !strings.Contains(result.Output, "\"ticker\"") { + t.Fatalf("output missing full schema for stock_quote: %q", result.Output) + } +} + +// A deferred tool the operator hid via DisabledTools must be invisible to +// tool_search: it never resolves via select:, it cannot be the operator-hidden +// candidate, and it must not appear in the no-match listing. +func TestToolSearchExcludesDisabledDeferredTool(t *testing.T) { + reg := newDeferredFixtureRegistry() + tool := NewToolSearchTool(reg).(optionsAwareTool) + disabled := RunOptions{DisabledTools: []string{"stock_quote"}} + + // select: a hidden tool must NOT resolve (no load_tools, informational message). + // Query "select:stock_quote_x" avoids echoing the disabled name verbatim so the + // listing-omission assertion below is not confused by the query echo. + selectResult := tool.RunWithOptions(context.Background(), + map[string]any{"query": "select:stock_quote_x"}, disabled) + if _, present := selectResult.Meta["load_tools"]; present { + t.Fatalf("disabled tool must not resolve via select:, got load_tools=%q", selectResult.Meta["load_tools"]) + } + // The available-tools listing (the segment after "Available tools: ") must omit + // the disabled tool while still naming the allowed one. + listing := selectResult.Output + if idx := strings.Index(listing, "Available tools: "); idx >= 0 { + listing = listing[idx:] + } + if strings.Contains(listing, "stock_quote") { + t.Fatalf("disabled tool leaked into no-match listing: %q", selectResult.Output) + } + if !strings.Contains(listing, "weather_lookup") { + t.Fatalf("no-match listing must still name the allowed tool, got %q", selectResult.Output) + } + + // keyword: a hidden tool must NOT rank. + keywordResult := tool.RunWithOptions(context.Background(), + map[string]any{"query": "stock"}, disabled) + if got := keywordResult.Meta["load_tools"]; got != "" { + t.Fatalf("disabled tool must not rank for a keyword query, got load_tools=%q", got) + } + if strings.Contains(keywordResult.Output, "stock_quote") { + t.Fatalf("disabled tool leaked into keyword output: %q", keywordResult.Output) + } +} + +// With a non-empty EnabledTools allowlist, a deferred tool absent from it is +// invisible to tool_search even though it is registered and deferred-eligible. +func TestToolSearchHonorsEnabledAllowlist(t *testing.T) { + reg := newDeferredFixtureRegistry() + tool := NewToolSearchTool(reg).(optionsAwareTool) + // Allowlist admits only weather_lookup; stock_quote is excluded. + allow := RunOptions{EnabledTools: []string{"weather_lookup"}} + + // The excluded tool must not resolve, even by exact name. + selectResult := tool.RunWithOptions(context.Background(), + map[string]any{"query": "select:weather_lookup,stock_quote"}, allow) + if got := selectResult.Meta["load_tools"]; got != "weather_lookup" { + t.Fatalf("allowlist must load only weather_lookup, got load_tools=%q", got) + } + if strings.Contains(selectResult.Output, "ticker") { + t.Fatalf("excluded tool's schema leaked into output: %q", selectResult.Output) + } + + // A keyword for the excluded tool finds nothing and the listing omits it. + keywordResult := tool.RunWithOptions(context.Background(), + map[string]any{"query": "stock"}, allow) + if _, present := keywordResult.Meta["load_tools"]; present { + t.Fatalf("excluded tool must not rank under the allowlist, got %q", keywordResult.Meta["load_tools"]) + } + if strings.Contains(keywordResult.Output, "stock_quote") { + t.Fatalf("excluded tool leaked into keyword no-match listing: %q", keywordResult.Output) + } +} diff --git a/internal/tools/types.go b/internal/tools/types.go index 4fb4d9f4..fe4ea7cd 100644 --- a/internal/tools/types.go +++ b/internal/tools/types.go @@ -11,15 +11,16 @@ type Permission string type Status string const ( + // SideEffectNone marks a control-only tool that neither reads nor mutates + // state — no read/write/shell/network/out-of-workspace effect. Examples: + // tool_search (only reports already-registered tool schemas) and + // escalate_model (only requests a loop-level model switch). + SideEffectNone SideEffect = "none" SideEffectRead SideEffect = "read" SideEffectWrite SideEffect = "write" SideEffectShell SideEffect = "shell" SideEffectNetwork SideEffect = "network" SideEffectOutOfWorkspace SideEffect = "out_of_workspace" - // SideEffectNone marks a control-only tool that performs no read/write/shell/ - // network/out-of-workspace effect (e.g. escalate_model, whose only effect is a - // loop-level model switch the agent loop performs). - SideEffectNone SideEffect = "none" ) const (