Files
zero/internal/agent/deferred_loop_test.go
KRATOS 2ea36cf138 Deferred (lazy) tool loading: tool_search + auto threshold for large MCP tool sets (#148)
* Deferred (lazy) tool loading: tool_search + auto threshold for large MCP tool sets

Adds deferred tool loading so a large MCP tool set no longer sends every tool's full JSON schema each turn. When the number of visible MCP tools crosses a threshold (config tools.deferThreshold, default 10; 0 disables), their schemas are withheld and advertised compactly in a per-turn system-reminder; the model pulls a tool's full schema on demand via a new tool_search tool. Strictly additive: below the threshold behaviour is byte-identical to today. No new module dependencies.

- Eligibility: MCP tools opt in via an optional Deferred() bool interface (tools.IsDeferred); built-ins stay eager. No change to the core Tool interface.
- Formatting (internal/tools/deferred.go): compact one-liners wrapped in a <system-reminder> by BuildDeferredReminder; server label via MCPServerName().
- tool_search (internal/tools/tool_search.go): SideEffectNone/PermissionAllow/AdvertiseInAuto; select:Name1,Name2 (exact) or keyword (ranked); returns full schemas and signals the loop via Meta["load_tools"]. Operator-filter-aware (never lists tools hidden by --enabled-tools/--disabled-tools).
- Agent loop: partitionTools is the single source of truth; eligible counts only ToolVisible deferred tools; Meta["load_tools"] lifts to ToolResult.LoadedTools and unions into a per-run loaded-set advertised next turn. Reminder appended only to the per-turn request copy (never persisted), including both reactive-compaction retry paths. tool_search stays reachable when deferral is active (exempt from the EnabledTools allowlist).
- Config + wiring: tools.deferThreshold (default 10; 0 disables; negative rejected); tool_search registered and DeferThreshold threaded in exec and TUI, gated on the same visible-deferred count the loop uses.

Recreated on top of current main; resolves the SideEffectNone clash with #146 (escalate_model). Supersedes #147.

* deferred-tools: activate deferral only when tool_search is runnable; drop dead deferredTools()

Addresses #148 review:
- partitionTools gates 'active' on a registered, non-disabled, mode-advertised
  tool_search (mirrors the executeToolCall dispatch gate). When the loader is
  unavailable (explicitly disabled, missing, or not advertised e.g. spec-draft),
  fall back to the eager/inactive path instead of hiding deferred tools behind a
  loader the model can never call (the dead-end).
- Flip TestDisabledToolSearch... to assert the eager fallback; register a real
  tool_search in the active-path tests that previously relied on activation
  without a usable loader.
- Remove unused toolSearchTool.deferredTools().

* deferred-tools: disable deferral in exec when tool_search disabled; drop unused toolDefinitions shim

Further #148 review:
- exec computes one effective DeferThreshold, forced to 0 when tool_search is in
  --disabled-tools, and feeds it to both registerToolSearchIfEligible and
  agent.Options.DeferThreshold — so the registration gate and the loop partition
  gate agree, and a disabled loader never enables deferral.
- Remove the unused toolDefinitions shim (golangci-lint unused; no callers).
2026-06-09 12:20:22 +05:30

769 lines
32 KiB
Go

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 <system-reminder> 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)
}
}