mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 01:09:46 -10:00
Compare commits
31
Commits
v0.3.0
...
deferred-tools
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
27e250c7b1 | ||
|
|
3c82c20d20 | ||
|
|
f7c74ef784 | ||
|
|
ae4ce72815 | ||
|
|
99cb28d1a3 | ||
|
|
410173c701 | ||
|
|
71da864a74 | ||
|
|
c8fa54c675 | ||
|
|
c3ece4b198 | ||
|
|
4960e40f1b | ||
|
|
12add88691 | ||
|
|
25ae21952d | ||
|
|
0d35d568f2 | ||
|
|
d967b26817 | ||
|
|
980f7d7f41 | ||
|
|
4721ce727c | ||
|
|
fdf11d6a0a | ||
|
|
d8cf8d3923 | ||
|
|
8774620e09 | ||
|
|
44998d4745 | ||
|
|
3a453baa7c | ||
|
|
ad34419ee5 | ||
|
|
9afcd9651c | ||
|
|
f9416f3283 | ||
|
|
624d7f798b | ||
|
|
adeda876b7 | ||
|
|
a1feeb37e2 | ||
|
|
b13529cb14 | ||
|
|
2c227a3395 | ||
|
|
4c7c3a1f48 | ||
|
|
2f525c2f5f |
No files matched your search
@@ -0,0 +1,735 @@
|
||||
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"})
|
||||
|
||||
// 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"})
|
||||
|
||||
loaded := map[string]bool{"mcp__srv__alpha": true, "mcp__srv__beta": true}
|
||||
exposed, reminder := partitionTools(registry, PermissionModeAuto, Options{DeferThreshold: 2}, loaded)
|
||||
|
||||
if len(exposed) != 2 {
|
||||
t.Fatalf("expected both loaded deferred tools exposed, 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"})
|
||||
|
||||
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"})
|
||||
|
||||
// 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"})
|
||||
|
||||
// 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 TestDisabledToolSearchStaysHiddenAndRejectedWhenActive(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"},
|
||||
}
|
||||
|
||||
exposed, _ := partitionTools(registry, PermissionModeAuto, options, map[string]bool{})
|
||||
for _, def := range exposed {
|
||||
if def.Name == "tool_search" {
|
||||
t.Fatalf("tool_search must stay hidden when explicitly disabled, got %#v", exposed)
|
||||
}
|
||||
}
|
||||
|
||||
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.StatusError || result.DenialReason != DenialFiltered {
|
||||
t.Fatalf("a disabled tool_search call must be filtered out, got status=%s denial=%v output=%q", result.Status, result.DenialReason, result.Output)
|
||||
}
|
||||
}
|
||||
+161
-7
@@ -61,6 +61,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
|
||||
@@ -70,9 +75,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
|
||||
@@ -91,9 +106,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)
|
||||
}
|
||||
@@ -117,9 +146,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 {
|
||||
@@ -201,6 +243,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
|
||||
}
|
||||
messages = append(messages, zeroruntime.Message{
|
||||
Role: zeroruntime.MessageRoleTool,
|
||||
Content: toolResult.Output,
|
||||
@@ -315,7 +363,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,
|
||||
@@ -395,6 +449,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.
|
||||
})
|
||||
@@ -417,6 +475,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
Redacted: result.Redacted,
|
||||
ChangedFiles: result.ChangedFiles,
|
||||
Display: result.Display,
|
||||
LoadedTools: loadedToolsFromResult(result.Meta),
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -799,24 +858,100 @@ func permissionActionFromSandbox(action sandbox.Action) PermissionAction {
|
||||
}
|
||||
}
|
||||
|
||||
// toolDefinitions returns the full-schema tool list for a turn. It now delegates
|
||||
// to partitionTools with no loaded set; with deferral inactive (the only state
|
||||
// the bare toolDefinitions is used in) the exposed slice is byte-identical to the
|
||||
// pre-deferral output. Retained so existing callers/tests keep their signature.
|
||||
func toolDefinitions(registry *tools.Registry, permissionMode PermissionMode, options Options) []zeroruntime.ToolDefinition {
|
||||
exposed, _ := partitionTools(registry, permissionMode, options, map[string]bool{})
|
||||
return exposed
|
||||
}
|
||||
|
||||
// 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++
|
||||
}
|
||||
}
|
||||
|
||||
active := options.DeferThreshold > 0 && eligible >= options.DeferThreshold
|
||||
|
||||
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 the gateway to the allowlisted deferred
|
||||
// tools, 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 unless an explicit DisabledTools entry
|
||||
// turns the loader itself off. This never runs on the inactive path, so the
|
||||
// byte-identical below-threshold output is preserved.
|
||||
if active && !exposedNames[tools.ToolSearchToolName] && !containsToolName(options.DisabledTools, tools.ToolSearchToolName) {
|
||||
if loader, ok := registry.Get(tools.ToolSearchToolName); ok {
|
||||
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 {
|
||||
@@ -937,6 +1072,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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
// DenialCategory classifies why a tool call was blocked before it executed.
|
||||
@@ -125,6 +129,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
|
||||
|
||||
+22
-10
@@ -326,6 +326,27 @@ func runInteractiveTUIWithSkin(stderr io.Writer, deps appDeps, skin string, perm
|
||||
return writeAppError(stderr, err.Error(), 1)
|
||||
}
|
||||
defer closeMCPRuntime(stderr, mcpRuntime)
|
||||
// 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
|
||||
// files nor run shell. Ask advertises them and routes each through the existing
|
||||
// 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)
|
||||
@@ -336,16 +357,6 @@ func runInteractiveTUIWithSkin(stderr io.Writer, deps appDeps, skin string, perm
|
||||
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
|
||||
// files nor run shell. Ask advertises them and routes each through the existing
|
||||
// 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).
|
||||
if permissionMode == "" {
|
||||
permissionMode = agent.PermissionModeAsk
|
||||
}
|
||||
return deps.runTUI(context.Background(), tui.Options{
|
||||
Cwd: workspaceRoot,
|
||||
ProviderName: resolved.Provider.Name,
|
||||
@@ -362,6 +373,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,
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/sessions"
|
||||
"github.com/Gitlawb/zero/internal/specialist"
|
||||
"github.com/Gitlawb/zero/internal/streamjson"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/worktrees"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
@@ -196,6 +197,14 @@ 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.
|
||||
registerToolSearchIfEligible(registry, resolved.Tools.DeferThreshold, permissionMode, options.enabledTools, options.disabledTools)
|
||||
images, err := resolveExecImages(options.imagePaths, workspaceRoot)
|
||||
if err != nil {
|
||||
return writeExecFormatUsageError(stdout, stderr, options.outputFormat, err.Error())
|
||||
@@ -311,6 +320,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: resolved.Tools.DeferThreshold,
|
||||
SessionID: preparedSession.Session.SessionID,
|
||||
CallingSessionID: options.callingSessionID,
|
||||
CallingToolUseID: options.callingToolUseID,
|
||||
@@ -414,6 +424,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 {
|
||||
|
||||
@@ -34,7 +34,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)}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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{}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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{} }
|
||||
|
||||
@@ -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_<server>_<tool>"). 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: <short-desc> | server: <X> | <input-hint>". The
|
||||
// "server: <X>" 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 = "<system-reminder>"
|
||||
systemReminderEnd = "</system-reminder>"
|
||||
)
|
||||
|
||||
// BuildDeferredReminder wraps the given deferred-tool advertisement lines in a
|
||||
// <system-reminder> 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()
|
||||
}
|
||||
@@ -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, "<system-reminder>") {
|
||||
t.Fatalf("missing opening tag: %q", got)
|
||||
}
|
||||
if !strings.HasSuffix(got, "</system-reminder>") {
|
||||
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_<server>_<tool>"), 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,259 @@
|
||||
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, ",")},
|
||||
}
|
||||
}
|
||||
|
||||
// deferredTools returns the registry's deferred-eligible tools sorted by name so
|
||||
// keyword ranking and listings are deterministic.
|
||||
func (tool toolSearchTool) deferredTools() []Tool {
|
||||
return tool.visibleDeferredTools(nil, nil)
|
||||
}
|
||||
|
||||
// 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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,9 @@ type Permission string
|
||||
type Status string
|
||||
|
||||
const (
|
||||
// SideEffectNone marks tools that neither read nor mutate state (e.g.
|
||||
// tool_search, which only reports already-registered tool schemas).
|
||||
SideEffectNone SideEffect = "none"
|
||||
SideEffectRead SideEffect = "read"
|
||||
SideEffectWrite SideEffect = "write"
|
||||
SideEffectShell SideEffect = "shell"
|
||||
|
||||
Reference in new issue
Block a user