mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-12 08:41:32 -10:00
* 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).
769 lines
32 KiB
Go
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)
|
|
}
|
|
}
|