diff --git a/internal/app/mcp_runtime.go b/internal/app/mcp_runtime.go index 778fdc40..479e5b53 100644 --- a/internal/app/mcp_runtime.go +++ b/internal/app/mcp_runtime.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "sort" + "strings" "github.com/usewhale/whale/internal/core" whalemcp "github.com/usewhale/whale/internal/mcp" @@ -103,6 +104,38 @@ func (a *App) setupDeferredToolSearchLocked(catalog *whalemcp.DeferredToolCatalo a.toolset.SetDeferredToolSearch(catAdapter, promoter, renderer) } +const availableDeferredToolsMaxChars = 4000 + +// renderDeferredToolsBlock returns the block for +// injection into the system prompt. Returns empty string when no MCP tools are +// registered. The output is capped at availableDeferredToolsMaxChars to avoid +// blowing up the system prompt when many MCP tools are registered. +func (a *App) renderDeferredToolsBlock() string { + if a.mcpManager == nil { + return "" + } + catalog := a.mcpManager.BuildDeferredCatalog() + block := whalemcp.RenderAvailableDeferredTools(catalog) + if len(block) <= availableDeferredToolsMaxChars { + return block + } + // Truncate at a newline boundary near the limit. + truncated := block[:availableDeferredToolsMaxChars] + if idx := strings.LastIndex(truncated, "\n"); idx > 0 { + truncated = truncated[:idx] + } + // Count how many tools we omitted. + allTools := catalog.Names() + shownCount := 0 + for _, name := range allTools { + if strings.Contains(truncated, name) { + shownCount++ + } + } + omitted := len(allTools) - shownCount + return truncated + fmt.Sprintf("\n... %d more tool(s) omitted\n", omitted) +} + // makeDeferredPromoter returns a function that builds full Tool objects for given names, // adds them to registries, and returns their specs. func (a *App) makeDeferredPromoter() tools.DeferredToolPromoter { diff --git a/internal/app/mcp_runtime_test.go b/internal/app/mcp_runtime_test.go index 6e322158..ea0177cd 100644 --- a/internal/app/mcp_runtime_test.go +++ b/internal/app/mcp_runtime_test.go @@ -295,3 +295,86 @@ func TestRestorePromotedToolsBuildToolsError(t *testing.T) { t.Fatal("expected error when BuildTools fails for non-existent tool, got nil") } } + +// --- renderDeferredToolsBlock tests --- + +func TestRenderDeferredToolsBlockNilManager(t *testing.T) { + app := &App{} // mcpManager is nil + if got := app.renderDeferredToolsBlock(); got != "" { + t.Fatalf("expected empty string for nil mcpManager, got %q", got) + } +} + +func TestRenderDeferredToolsBlockWithTools(t *testing.T) { + mgr := newMCPRuntimeTestManager(t, "echoes a message") + app := &App{mcpManager: mgr} + + block := app.renderDeferredToolsBlock() + if block == "" { + t.Fatal("expected non-empty block for populated catalog") + } + if !strings.Contains(block, "") { + t.Fatalf("expected opening tag, got %q", block) + } + if !strings.Contains(block, "") { + t.Fatalf("expected closing tag, got %q", block) + } + if !strings.Contains(block, "mcp__runtime__echo") { + t.Fatalf("expected tool name in block, got %q", block) + } +} + +func TestRenderDeferredToolsBlockFormat(t *testing.T) { + mgr := newMCPRuntimeTestManager(t, "echoes a message") + app := &App{mcpManager: mgr} + + block := app.renderDeferredToolsBlock() + if block == "" { + t.Fatal("expected non-empty block") + } + if !strings.HasPrefix(block, "\n") { + t.Fatalf("expected block to start with opening tag on its own line, got %q", block) + } + if !strings.HasSuffix(block, "") { + t.Fatalf("expected block to end with closing tag, got %q", block) + } + if !strings.Contains(block, "[server: ") { + t.Fatalf("expected server section, got %q", block) + } + if !strings.Contains(block, " — ") { + t.Fatalf("expected tool name/description separator, got %q", block) + } +} + +func TestRenderDeferredToolsBlockTruncation(t *testing.T) { + // Simulate truncation by verifying that a block exceeding the limit + // is handled correctly. We test the truncation behaviour directly + // rather than through the full MCP pipeline (which has description + // length limits in the protocol). + longLine := strings.Repeat("x", availableDeferredToolsMaxChars+100) + block := "\n[server: test]\n mcp__test__tool — " + longLine + "\n" + + // Apply the same truncation logic as renderDeferredToolsBlock. + truncated := block + if len(truncated) > availableDeferredToolsMaxChars { + truncated = block[:availableDeferredToolsMaxChars] + if idx := strings.LastIndex(truncated, "\n"); idx > 0 { + truncated = truncated[:idx] + } + truncated += "\n... more tool(s) omitted\n" + } + + if !strings.HasPrefix(truncated, "") { + t.Fatalf("truncated block should start with opening tag, got %q", truncated) + } + if !strings.Contains(truncated, "more tool(s) omitted") { + t.Fatalf("truncated block should contain omission notice, got %q", truncated) + } + if !strings.HasSuffix(truncated, "") { + t.Fatalf("truncated block should end with closing tag, got %q", truncated) + } + // Truncation should have removed content. + if strings.Contains(truncated, "xxx") { + t.Fatal("truncated block should not contain the overflow content") + } +} diff --git a/internal/app/runtime.go b/internal/app/runtime.go index db617a47..c67d8dd2 100644 --- a/internal/app/runtime.go +++ b/internal/app/runtime.go @@ -116,6 +116,7 @@ func (a *App) ensureAgent() (*agent.Agent, error) { agent.WithHookRunner(a.hookRunner), agent.WithExtraSystemBlocks(pluginBlocks...), agent.WithDynamicSystemBlocksForTurn(a.workflowDynamicSystemBlock), + agent.WithDynamicSystemBlocks(func() string { return a.renderDeferredToolsBlock() }), agent.WithProjectMemory(a.cfg.MemoryEnabled, a.cfg.MemoryMaxChars, parseCSVList(a.cfg.MemoryFileOrder), a.workspaceRoot), agent.WithWorktreeContext(a.worktree.Path, a.worktree.OriginalWorkspace), agent.WithMaxParallelSubagents(a.cfg.MaxParallelSubagents), diff --git a/internal/mcp/deferred.go b/internal/mcp/deferred.go index 2a94c75c..3712ed81 100644 --- a/internal/mcp/deferred.go +++ b/internal/mcp/deferred.go @@ -129,9 +129,6 @@ func (c *DeferredToolCatalog) selectSearch(names string) []DeferredToolMeta { for _, t := range c.tools { if wanted[t.Name] { results = append(results, t) - if len(results) >= maxSearchResults { - break - } } } return results diff --git a/internal/mcp/deferred_test.go b/internal/mcp/deferred_test.go index 70461dfa..8c26b009 100644 --- a/internal/mcp/deferred_test.go +++ b/internal/mcp/deferred_test.go @@ -165,6 +165,29 @@ func TestSearchSelectFormEmptyNames(t *testing.T) { } } +func TestSearchSelectFormNoCap(t *testing.T) { + // select: queries should not be capped at maxSearchResults — the user + // explicitly requested those names, so all should be returned. + tools := make([]DeferredToolMeta, 0, maxSearchResults+5) + for i := range maxSearchResults + 5 { + tools = append(tools, testCatalogMeta( + "mcp__s__tool_"+string(rune('a'+i%26))+string(rune('a'+i/26)), + "s", + "", + )) + } + c := NewDeferredToolCatalog(tools) + // Build a select: query with all tool names. + names := make([]string, len(tools)) + for i, t := range tools { + names[i] = t.Name + } + results := c.Search("select:" + strings.Join(names, ",")) + if len(results) != len(tools) { + t.Fatalf("select should return all %d requested tools, got %d", len(tools), len(results)) + } +} + func TestSearchMustHaveForm(t *testing.T) { c := NewDeferredToolCatalog([]DeferredToolMeta{ testCatalogMeta("mcp__github__search_code", "github", "search code in repos"), diff --git a/internal/tools/catalog_mcp.go b/internal/tools/catalog_mcp.go index 4a1f4b43..8d3ebaa9 100644 --- a/internal/tools/catalog_mcp.go +++ b/internal/tools/catalog_mcp.go @@ -108,11 +108,7 @@ Returns up to 5 matching tools and activates them for subsequent calls.`, names := catalog.Names() hint := "" if len(names) > 0 { - sample := names - if len(sample) > 5 { - sample = sample[:5] - } - hint = fmt.Sprintf(" No match. Available deferred tools (showing first 5 of %d): %s", len(names), strings.Join(sample, ", ")) + hint = fmt.Sprintf(" No match. Available deferred tools (%d total): %s", len(names), strings.Join(names, ", ")) } return core.ToolResult{ ToolCallID: call.ID, diff --git a/internal/tools/catalog_mcp_test.go b/internal/tools/catalog_mcp_test.go index 9043dc56..e628abdb 100644 --- a/internal/tools/catalog_mcp_test.go +++ b/internal/tools/catalog_mcp_test.go @@ -186,6 +186,9 @@ func TestToolSearchNoMatch(t *testing.T) { if !strings.Contains(result.ModelText, "mcp__s__alpha") { t.Fatalf("expected hint with available tool names, got %q", result.ModelText) } + if !strings.Contains(result.ModelText, "(1 total)") { + t.Fatalf("expected hint with total count, got %q", result.ModelText) + } } func TestToolSearchWithoutPromoterReturnsInfo(t *testing.T) {