mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: defer MCP tool schemas behind a find_tools search (#28225)
## Summary When the `mcp-tool-search` experiment is enabled, chatd stops inlining connected MCP tool schemas into every generation. It instead exposes a built-in `find_tools` tool whose description carries a compact catalog of the deferred tools, and only ships full JSON schemas for tools the model has activated by searching or by calling them directly. Closes [CODAGT-760](https://linear.app/coder/issue/CODAGT-760). ## Problem Tool-heavy agent configurations (GitHub, Linear, Notion, and dev-tooling MCP servers) inline over 100k tokens of tool schema definitions into every generation. Initial uncached requests reached ~216k tokens with time-to-first-token close to nine minutes, while the model typically invokes only a handful of tools per turn. ## How it works - `decideMCPToolSearch` defers external and workspace `.mcp.json` MCP tools whenever the experiment is enabled. Native, dynamic, provider, skill, and transport tools are never deferred. - `find_tools` embeds a server-grouped catalog in its tool description (degrading to names-only, then counts-only, then a constant-size summary past a context-scaled size cap) and scores keyword matches across tool names, descriptions, parameter schemas, and server metadata. Queries can scope to one server with a `server:` prefix, and exact `names` arguments always activate. - Activation state is ephemeral: it is re-derived each generation from surviving chat history (`find_tools` results and direct calls to deferred tools), so activations naturally lapse when compaction summarizes them away. Aggregate activated schema weight is capped at 10% of the context window, shedding the least recently activated schemas first; `find_tools` shares that budget across parallel calls in one step. No new persistence. - Deferred tools stay registered for execution, so the model can call a cataloged tool directly without searching first; the schema is activated for subsequent steps. - Fail-open: the experiment being disabled, an empty candidate set, or an MCP tool named `find_tools` all disable deferral, leaving today's behavior byte-identical on the wire. - Prometheus counters/histograms track `find_tools` calls, matches, activations, and deferred token weight. - The conversation timeline renders `find_tools` calls with a collapsed search summary and expandable match list, falling back to the generic renderer on malformed payloads. ## Validation - Unit tests for the catalog, matcher, experiment-gated decision, and activation derivation; end-to-end chatd generation tests covering search-then-call, direct-call activation, experiment-off wire parity, compaction lapse, and subagent tool gating. - Storybook interaction tests for the timeline rendering and malformed-payload fallback. - Remote dogfood UAT on dev.coder.com passed: deferral with a real MCP server and Anthropic model, direct calls without prior search, activation persistence across turns, experiment-off parity, and clean UI/console. > Disclosure: Mux (AI agent) authored this PR on Mike's behalf.
This commit is contained in:
@@ -238,10 +238,16 @@ type AssistantOutcome struct {
|
||||
|
||||
// ExecuteLocalToolsOptions configures one local tool execution batch.
|
||||
type ExecuteLocalToolsOptions struct {
|
||||
Tools []fantasy.AgentTool
|
||||
ActiveTools []string
|
||||
ProviderTools []ProviderTool
|
||||
ToolCalls []fantasy.ToolCallContent
|
||||
Tools []fantasy.AgentTool
|
||||
ActiveTools []string
|
||||
AllowInactiveTools map[string]bool
|
||||
ProviderTools []ProviderTool
|
||||
ToolCalls []fantasy.ToolCallContent
|
||||
// ObservedToolCalls optionally carries the step's full assistant
|
||||
// tool-call batch, including calls denied before execution, so
|
||||
// step observers account for denied siblings that derivation will
|
||||
// still count. Defaults to ToolCalls.
|
||||
ObservedToolCalls []fantasy.ToolCallContent
|
||||
|
||||
ExclusiveToolNames map[string]bool
|
||||
BuiltinToolNames map[string]bool
|
||||
@@ -594,8 +600,10 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool
|
||||
opts.Clock,
|
||||
opts.Tools,
|
||||
opts.ActiveTools,
|
||||
opts.AllowInactiveTools,
|
||||
opts.ProviderTools,
|
||||
localCalls,
|
||||
opts.ObservedToolCalls,
|
||||
opts.Metrics,
|
||||
opts.Logger,
|
||||
provider,
|
||||
@@ -1057,8 +1065,10 @@ func executeTools(
|
||||
clock quartz.Clock,
|
||||
allTools []fantasy.AgentTool,
|
||||
activeTools []string,
|
||||
allowInactiveTools map[string]bool,
|
||||
providerTools []ProviderTool,
|
||||
toolCalls []fantasy.ToolCallContent,
|
||||
observedToolCalls []fantasy.ToolCallContent,
|
||||
metrics *Metrics,
|
||||
logger slog.Logger,
|
||||
provider, model string,
|
||||
@@ -1109,47 +1119,83 @@ func executeTools(
|
||||
}
|
||||
}
|
||||
|
||||
observed := observedToolCalls
|
||||
if observed == nil {
|
||||
observed = localToolCalls
|
||||
}
|
||||
notifyStepToolCallObservers(toolMap, toolNameAliases, observed)
|
||||
|
||||
results := make([]fantasy.ToolResultContent, len(localToolCalls))
|
||||
completedAt := make([]time.Time, len(localToolCalls))
|
||||
runCall := func(i int, tc fantasy.ToolCallContent) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
results[i] = fantasy.ToolResultContent{
|
||||
ToolCallID: tc.ToolCallID,
|
||||
ToolName: tc.ToolName,
|
||||
Result: fantasy.ToolResultOutputContentError{
|
||||
Error: xerrors.Errorf("tool panicked: %v", r),
|
||||
},
|
||||
}
|
||||
}
|
||||
// Record when this tool completed (or panicked).
|
||||
// Captured per call so parallel tools get
|
||||
// accurate individual completion times.
|
||||
completedAt[i] = clockNow(clock)
|
||||
}()
|
||||
results[i] = executeSingleTool(
|
||||
ctx,
|
||||
toolMap,
|
||||
tc,
|
||||
metrics,
|
||||
logger,
|
||||
provider,
|
||||
model,
|
||||
builtinToolNames,
|
||||
activeTools,
|
||||
allowInactiveTools,
|
||||
providerRunnerNames,
|
||||
resultProviderMetadata,
|
||||
maxResultBytes,
|
||||
toolNameAliases,
|
||||
)
|
||||
}
|
||||
// Calls to tools that opt in via SerialToolCalls run in tool-call
|
||||
// order after every concurrent sibling has settled. The step waits
|
||||
// for all calls anyway, so sequencing them last costs nothing, and
|
||||
// order-sensitive shared state (for example the find_tools
|
||||
// activation budget) is claimed deterministically after sibling
|
||||
// outcomes are known. All other calls stay concurrent.
|
||||
var serialIndexes []int
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(len(localToolCalls))
|
||||
for i, tc := range localToolCalls {
|
||||
if isSerialToolCall(toolMap, toolNameAliases, tc.ToolName) {
|
||||
serialIndexes = append(serialIndexes, i)
|
||||
continue
|
||||
}
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
results[i] = fantasy.ToolResultContent{
|
||||
ToolCallID: tc.ToolCallID,
|
||||
ToolName: tc.ToolName,
|
||||
Result: fantasy.ToolResultOutputContentError{
|
||||
Error: xerrors.Errorf("tool panicked: %v", r),
|
||||
},
|
||||
}
|
||||
}
|
||||
// Record when this tool completed (or panicked).
|
||||
// Captured per-goroutine so parallel tools get
|
||||
// accurate individual completion times.
|
||||
completedAt[i] = clockNow(clock)
|
||||
}()
|
||||
results[i] = executeSingleTool(
|
||||
ctx,
|
||||
toolMap,
|
||||
tc,
|
||||
metrics,
|
||||
logger,
|
||||
provider,
|
||||
model,
|
||||
builtinToolNames,
|
||||
activeTools,
|
||||
providerRunnerNames,
|
||||
resultProviderMetadata,
|
||||
maxResultBytes,
|
||||
toolNameAliases,
|
||||
)
|
||||
runCall(i, tc)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// Reconcile settled sibling outcomes before serial tools run, so
|
||||
// for example find_tools refunds reservations of errored direct
|
||||
// calls before its searches admit activations.
|
||||
settled := make([]fantasy.ToolResultContent, 0, len(results))
|
||||
for i := range results {
|
||||
if !slices.Contains(serialIndexes, i) {
|
||||
settled = append(settled, results[i])
|
||||
}
|
||||
}
|
||||
notifyStepToolResultObservers(toolMap, toolNameAliases, localToolCalls, observed, settled)
|
||||
|
||||
for _, i := range serialIndexes {
|
||||
runCall(i, localToolCalls[i])
|
||||
}
|
||||
|
||||
// Publish results in the original tool-call order so SSE
|
||||
// subscribers see a deterministic event sequence.
|
||||
if onResult != nil {
|
||||
@@ -1262,6 +1308,7 @@ func executeSingleTool(
|
||||
provider, model string,
|
||||
builtinToolNames map[string]bool,
|
||||
activeTools []string,
|
||||
allowInactiveTools map[string]bool,
|
||||
providerRunnerNames map[string]struct{},
|
||||
resultProviderMetadata map[string]func(fantasy.ToolResponse) fantasy.ProviderMetadata,
|
||||
maxResultBytes int,
|
||||
@@ -1294,7 +1341,7 @@ func executeSingleTool(
|
||||
}
|
||||
|
||||
_, isProviderRunner := providerRunnerNames[resolvedName]
|
||||
if !isProviderRunner && !isToolActive(resolvedName, activeTools) {
|
||||
if !isProviderRunner && !isToolActive(resolvedName, activeTools) && !allowInactiveTools[resolvedName] {
|
||||
result.Result = fantasy.ToolResultOutputContentError{
|
||||
Error: xerrors.New("Tool not active in this turn: " + resolvedName),
|
||||
}
|
||||
@@ -1455,6 +1502,119 @@ func isToolActive(name string, activeTools []string) bool {
|
||||
return len(activeTools) == 0 || slices.Contains(activeTools, name)
|
||||
}
|
||||
|
||||
// serialToolCaller is implemented by tools whose calls within one step
|
||||
// must execute in tool-call order because they claim from shared state.
|
||||
type serialToolCaller interface{ SerialToolCalls() bool }
|
||||
|
||||
// stepToolCallObserver is implemented by tools that need to see every
|
||||
// tool-call name in the step before any call executes, for example so
|
||||
// find_tools can charge same-step direct calls against its budget.
|
||||
type stepToolCallObserver interface{ ObserveStepToolCalls(names []string) }
|
||||
|
||||
// notifyStepToolCallObservers passes the step's resolved tool-call
|
||||
// names to each distinct called tool that observes them.
|
||||
func notifyStepToolCallObservers(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, calls []fantasy.ToolCallContent) {
|
||||
names := make([]string, 0, len(calls))
|
||||
for _, tc := range calls {
|
||||
name := tc.ToolName
|
||||
if alias, ok := toolNameAliases[name]; ok {
|
||||
name = alias
|
||||
}
|
||||
names = append(names, name)
|
||||
}
|
||||
notified := make(map[string]struct{}, len(names))
|
||||
for _, name := range names {
|
||||
if _, dup := notified[name]; dup {
|
||||
continue
|
||||
}
|
||||
notified[name] = struct{}{}
|
||||
tool, ok := toolMap[name]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if observer, ok := tool.(stepToolCallObserver); ok {
|
||||
observer.ObserveStepToolCalls(names)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// stepToolResultObserver is implemented by tools that need the step's
|
||||
// per-call execution outcomes, for example so find_tools can refund
|
||||
// budget it reserved for a direct call whose execution errored. names
|
||||
// and errored are parallel slices in the observed tool-call order;
|
||||
// outcomes are kept per call because one tool can be called several
|
||||
// times in a step with different results.
|
||||
type stepToolResultObserver interface {
|
||||
ObserveStepToolResults(names []string, errored []bool)
|
||||
}
|
||||
|
||||
// notifyStepToolResultObservers passes the settled sibling outcomes to
|
||||
// each distinct called tool that observes them, per call in observed
|
||||
// order. Observed calls missing from the executed batch were rejected
|
||||
// before execution (for example malformed JSON partitioned into
|
||||
// synthetic denials) and settle as errored, since their persisted
|
||||
// results always carry IsError. Serial calls have not run yet, so
|
||||
// their own outcomes are reported as not errored; observers only need
|
||||
// the concurrent siblings they share state with.
|
||||
func notifyStepToolResultObservers(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, calls, observed []fantasy.ToolCallContent, settled []fantasy.ToolResultContent) {
|
||||
resolve := func(name string) string {
|
||||
if alias, ok := toolNameAliases[name]; ok {
|
||||
return alias
|
||||
}
|
||||
return name
|
||||
}
|
||||
erroredByID := make(map[string]bool, len(settled))
|
||||
for _, tr := range settled {
|
||||
_, isErr := tr.Result.(fantasy.ToolResultOutputContentError)
|
||||
erroredByID[tr.ToolCallID] = isErr
|
||||
}
|
||||
executedIDs := make(map[string]struct{}, len(calls))
|
||||
for _, tc := range calls {
|
||||
executedIDs[tc.ToolCallID] = struct{}{}
|
||||
}
|
||||
names := make([]string, 0, len(observed))
|
||||
errored := make([]bool, 0, len(observed))
|
||||
for _, tc := range observed {
|
||||
names = append(names, resolve(tc.ToolName))
|
||||
if isErr, ok := erroredByID[tc.ToolCallID]; ok {
|
||||
errored = append(errored, isErr)
|
||||
continue
|
||||
}
|
||||
_, executed := executedIDs[tc.ToolCallID]
|
||||
errored = append(errored, !executed)
|
||||
}
|
||||
notified := make(map[string]struct{}, len(calls))
|
||||
for _, tc := range calls {
|
||||
name := tc.ToolName
|
||||
if alias, ok := toolNameAliases[name]; ok {
|
||||
name = alias
|
||||
}
|
||||
if _, dup := notified[name]; dup {
|
||||
continue
|
||||
}
|
||||
notified[name] = struct{}{}
|
||||
tool, ok := toolMap[name]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if observer, ok := tool.(stepToolResultObserver); ok {
|
||||
observer.ObserveStepToolResults(names, errored)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func isSerialToolCall(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, name string) bool {
|
||||
if alias, ok := toolNameAliases[name]; ok {
|
||||
name = alias
|
||||
}
|
||||
tool, ok := toolMap[name]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
serial, ok := tool.(serialToolCaller)
|
||||
return ok && serial.SerialToolCalls()
|
||||
}
|
||||
|
||||
// buildToolDefinitions converts AgentTool definitions into the
|
||||
// fantasy.Tool slice expected by fantasy.Call. When activeTools
|
||||
// is non-empty, only function tools whose name appears in the
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"iter"
|
||||
"runtime"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
@@ -872,6 +874,313 @@ func TestSanitizeAnthropicProviderToolContent(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
type serialMarkerTool struct{ fantasy.AgentTool }
|
||||
|
||||
func (serialMarkerTool) SerialToolCalls() bool { return true }
|
||||
|
||||
type observerMarkerTool struct {
|
||||
fantasy.AgentTool
|
||||
observed func(names []string)
|
||||
}
|
||||
|
||||
func (t observerMarkerTool) ObserveStepToolCalls(names []string) { t.observed(names) }
|
||||
|
||||
func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var mu sync.Mutex
|
||||
var observedNames []string
|
||||
observedBeforeRun := false
|
||||
observer := observerMarkerTool{
|
||||
AgentTool: fantasy.NewAgentTool(
|
||||
"observer_tool",
|
||||
"records sibling calls",
|
||||
func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
mu.Lock()
|
||||
observedBeforeRun = observedNames != nil
|
||||
mu.Unlock()
|
||||
return fantasy.NewTextResponse("ok"), nil
|
||||
},
|
||||
),
|
||||
observed: func(names []string) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
observedNames = append([]string{}, names...)
|
||||
},
|
||||
}
|
||||
var uncalledObserved atomic.Bool
|
||||
uncalledObserver := observerMarkerTool{
|
||||
AgentTool: fantasy.NewAgentTool(
|
||||
"uncalled_observer",
|
||||
"never called this step",
|
||||
func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
return fantasy.NewTextResponse("ok"), nil
|
||||
},
|
||||
),
|
||||
observed: func([]string) { uncalledObserved.Store(true) },
|
||||
}
|
||||
other := fantasy.NewAgentTool(
|
||||
"other_tool",
|
||||
"plain tool",
|
||||
func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
return fantasy.NewTextResponse("ok"), nil
|
||||
},
|
||||
)
|
||||
|
||||
executeTools(
|
||||
context.Background(),
|
||||
quartz.NewReal(),
|
||||
[]fantasy.AgentTool{observer, uncalledObserver, other},
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
[]fantasy.ToolCallContent{
|
||||
{ToolCallID: "1", ToolName: "observer_alias", Input: "{}"},
|
||||
{ToolCallID: "2", ToolName: "other_tool", Input: "{}"},
|
||||
},
|
||||
[]fantasy.ToolCallContent{
|
||||
{ToolCallID: "1", ToolName: "observer_alias", Input: "{}"},
|
||||
{ToolCallID: "2", ToolName: "other_tool", Input: "{}"},
|
||||
{ToolCallID: "3", ToolName: "denied_tool", Input: "{}"},
|
||||
},
|
||||
NewMetrics(prometheus.NewRegistry()),
|
||||
slog.Make(),
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
defaultToolResultBytes,
|
||||
map[string]string{"observer_alias": "observer_tool"},
|
||||
nil,
|
||||
)
|
||||
|
||||
require.Equal(t, []string{"observer_tool", "other_tool", "denied_tool"}, observedNames,
|
||||
"a called observer sees every observed tool-call name, including calls denied before execution")
|
||||
require.True(t, observedBeforeRun, "observers are notified before any tool call executes")
|
||||
require.False(t, uncalledObserved.Load(), "tools not called this step are not notified")
|
||||
}
|
||||
|
||||
type resultObserverMarkerTool struct {
|
||||
fantasy.AgentTool
|
||||
observedResults func(names []string, errored []bool)
|
||||
}
|
||||
|
||||
func (t resultObserverMarkerTool) ObserveStepToolResults(names []string, errored []bool) {
|
||||
t.observedResults(names, errored)
|
||||
}
|
||||
|
||||
func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var mu sync.Mutex
|
||||
var gotNames []string
|
||||
var gotErrored []bool
|
||||
notifications := 0
|
||||
observer := resultObserverMarkerTool{
|
||||
AgentTool: fantasy.NewAgentTool(
|
||||
"observer_tool",
|
||||
"records sibling outcomes",
|
||||
func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
return fantasy.NewTextResponse("ok"), nil
|
||||
},
|
||||
),
|
||||
observedResults: func(names []string, errored []bool) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
notifications++
|
||||
gotNames = append([]string{}, names...)
|
||||
gotErrored = append([]bool{}, errored...)
|
||||
},
|
||||
}
|
||||
failing := fantasy.NewAgentTool(
|
||||
"failing_tool",
|
||||
"returns an error result",
|
||||
func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
return fantasy.NewTextErrorResponse("remote error"), nil
|
||||
},
|
||||
)
|
||||
|
||||
executed := []fantasy.ToolCallContent{
|
||||
{ToolCallID: "1", ToolName: "observer_alias", Input: "{}"},
|
||||
{ToolCallID: "2", ToolName: "failing_tool", Input: "{}"},
|
||||
{ToolCallID: "3", ToolName: "missing_tool", Input: "{}"},
|
||||
}
|
||||
executeTools(
|
||||
context.Background(),
|
||||
quartz.NewReal(),
|
||||
[]fantasy.AgentTool{observer, failing},
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
executed,
|
||||
append(slices.Clone(executed), fantasy.ToolCallContent{
|
||||
ToolCallID: "4", ToolName: "rejected_tool", Input: "{not json",
|
||||
}),
|
||||
NewMetrics(prometheus.NewRegistry()),
|
||||
slog.Make(),
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
defaultToolResultBytes,
|
||||
map[string]string{"observer_alias": "observer_tool"},
|
||||
nil,
|
||||
)
|
||||
|
||||
require.Equal(t, 1, notifications, "each called observer is notified once per step")
|
||||
require.Equal(t, []string{"observer_tool", "failing_tool", "missing_tool", "rejected_tool"}, gotNames,
|
||||
"outcomes are reported per call in observed order with aliases resolved")
|
||||
require.Equal(t, []bool{false, true, true, true}, gotErrored,
|
||||
"error results, unresolvable tools, and observed calls rejected before execution all settle as errored outcomes")
|
||||
}
|
||||
|
||||
type serialResultObserverTool struct {
|
||||
fantasy.AgentTool
|
||||
observedResults func(names []string, errored []bool)
|
||||
}
|
||||
|
||||
func (serialResultObserverTool) SerialToolCalls() bool { return true }
|
||||
|
||||
func (t serialResultObserverTool) ObserveStepToolResults(names []string, errored []bool) {
|
||||
t.observedResults(names, errored)
|
||||
}
|
||||
|
||||
func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var mu sync.Mutex
|
||||
var erroredAtNotify []string
|
||||
var erroredAtRun []string
|
||||
notified := false
|
||||
serial := serialResultObserverTool{
|
||||
AgentTool: fantasy.NewAgentTool(
|
||||
"serial_observer",
|
||||
"observes sibling outcomes before running",
|
||||
func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
mu.Lock()
|
||||
erroredAtRun = append([]string{}, erroredAtNotify...)
|
||||
mu.Unlock()
|
||||
return fantasy.NewTextResponse("ok"), nil
|
||||
},
|
||||
),
|
||||
observedResults: func(names []string, errored []bool) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
notified = true
|
||||
erroredAtNotify = nil
|
||||
for i, name := range names {
|
||||
if errored[i] {
|
||||
erroredAtNotify = append(erroredAtNotify, name)
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
failing := fantasy.NewAgentTool(
|
||||
"failing_tool",
|
||||
"returns an error result",
|
||||
func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
return fantasy.NewTextErrorResponse("remote error"), nil
|
||||
},
|
||||
)
|
||||
|
||||
results := executeTools(
|
||||
context.Background(),
|
||||
quartz.NewReal(),
|
||||
[]fantasy.AgentTool{serial, failing},
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
[]fantasy.ToolCallContent{
|
||||
{ToolCallID: "1", ToolName: "serial_observer", Input: "{}"},
|
||||
{ToolCallID: "2", ToolName: "failing_tool", Input: "{}"},
|
||||
},
|
||||
nil,
|
||||
NewMetrics(prometheus.NewRegistry()),
|
||||
slog.Make(),
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
defaultToolResultBytes,
|
||||
nil,
|
||||
nil,
|
||||
)
|
||||
|
||||
require.True(t, notified)
|
||||
require.Equal(t, []string{"failing_tool"}, erroredAtRun,
|
||||
"a serial tool must see settled sibling outcomes before it executes")
|
||||
require.Len(t, results, 2)
|
||||
require.Equal(t, "1", results[0].ToolCallID, "results keep original call order")
|
||||
}
|
||||
|
||||
func TestExecuteToolsSerialToolCallOrder(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var mu sync.Mutex
|
||||
var events []string
|
||||
inFlight := 0
|
||||
maxInFlight := 0
|
||||
record := func(event string, delta int) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
inFlight += delta
|
||||
if inFlight > maxInFlight {
|
||||
maxInFlight = inFlight
|
||||
}
|
||||
events = append(events, event)
|
||||
}
|
||||
serial := serialMarkerTool{AgentTool: fantasy.NewAgentTool(
|
||||
"serial_tool",
|
||||
"records call order",
|
||||
func(_ context.Context, _ struct{}, call fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
record(call.ID+":start", 1)
|
||||
runtime.Gosched()
|
||||
record(call.ID+":end", -1)
|
||||
return fantasy.NewTextResponse("ok"), nil
|
||||
},
|
||||
)}
|
||||
parallelRan := make(chan struct{})
|
||||
parallel := fantasy.NewAgentTool(
|
||||
"parallel_tool",
|
||||
"plain tool",
|
||||
func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
close(parallelRan)
|
||||
return fantasy.NewTextResponse("ok"), nil
|
||||
},
|
||||
)
|
||||
|
||||
calls := []fantasy.ToolCallContent{
|
||||
{ToolCallID: "a", ToolName: "serial_tool", Input: "{}"},
|
||||
{ToolCallID: "p", ToolName: "parallel_tool", Input: "{}"},
|
||||
{ToolCallID: "b", ToolName: "serial_tool", Input: "{}"},
|
||||
{ToolCallID: "c", ToolName: "serial_tool", Input: "{}"},
|
||||
}
|
||||
results := executeTools(
|
||||
context.Background(),
|
||||
quartz.NewReal(),
|
||||
[]fantasy.AgentTool{serial, parallel},
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
calls,
|
||||
nil,
|
||||
NewMetrics(prometheus.NewRegistry()),
|
||||
slog.Make(),
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
defaultToolResultBytes,
|
||||
nil,
|
||||
nil,
|
||||
)
|
||||
|
||||
require.Equal(t, []string{"a:start", "a:end", "b:start", "b:end", "c:start", "c:end"}, events,
|
||||
"serial tool calls must run one at a time in tool-call order")
|
||||
require.Equal(t, 1, maxInFlight)
|
||||
select {
|
||||
case <-parallelRan:
|
||||
default:
|
||||
t.Fatal("parallel tool call did not run")
|
||||
}
|
||||
require.Len(t, results, len(calls))
|
||||
for i, tc := range calls {
|
||||
require.Equal(t, tc.ToolCallID, results[i].ToolCallID, "results keep original call order")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -912,6 +1221,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) {
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
[]string{"screenshot"},
|
||||
nil,
|
||||
map[string]struct{}{},
|
||||
nil,
|
||||
defaultToolResultBytes,
|
||||
@@ -961,6 +1271,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) {
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
[]string{"screenshot"},
|
||||
nil,
|
||||
map[string]struct{}{},
|
||||
nil,
|
||||
defaultToolResultBytes,
|
||||
@@ -1005,6 +1316,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) {
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
[]string{"echo"},
|
||||
nil,
|
||||
map[string]struct{}{},
|
||||
nil,
|
||||
defaultToolResultBytes,
|
||||
@@ -1053,6 +1365,7 @@ func TestExecuteSingleTool_ResolvesToolNameAlias(t *testing.T) {
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
[]string{"interrupt_agent"},
|
||||
nil,
|
||||
map[string]struct{}{},
|
||||
nil,
|
||||
defaultToolResultBytes,
|
||||
@@ -1093,6 +1406,7 @@ func TestExecuteSingleTool_UnknownAliasFallsThrough(t *testing.T) {
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
[]string{"interrupt_agent"},
|
||||
nil,
|
||||
map[string]struct{}{},
|
||||
nil,
|
||||
defaultToolResultBytes,
|
||||
@@ -1103,3 +1417,32 @@ func TestExecuteSingleTool_UnknownAliasFallsThrough(t *testing.T) {
|
||||
require.True(t, ok, "expected error output, got %T", result.Result)
|
||||
require.Contains(t, errOutput.Error.Error(), "close_agent")
|
||||
}
|
||||
|
||||
func TestExecuteSingleTool_AllowsDeferredDirectCall(t *testing.T) {
|
||||
t.Parallel()
|
||||
tool := fantasy.NewAgentTool(
|
||||
"server__direct",
|
||||
"direct",
|
||||
func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
return fantasy.NewTextResponse("ok"), nil
|
||||
},
|
||||
)
|
||||
result := executeSingleTool(
|
||||
context.Background(),
|
||||
map[string]fantasy.AgentTool{"server__direct": tool},
|
||||
fantasy.ToolCallContent{ToolCallID: "call-direct", ToolName: "server__direct", Input: "{}"},
|
||||
NewMetrics(prometheus.NewRegistry()),
|
||||
slog.Make(),
|
||||
"fake", "fake-model",
|
||||
map[string]bool{},
|
||||
[]string{"find_tools"},
|
||||
map[string]bool{"server__direct": true},
|
||||
map[string]struct{}{},
|
||||
nil,
|
||||
defaultToolResultBytes,
|
||||
nil,
|
||||
)
|
||||
text, ok := result.Result.(fantasy.ToolResultOutputContentText)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "ok", text.Text)
|
||||
}
|
||||
|
||||
@@ -27,17 +27,21 @@ const (
|
||||
|
||||
// Metrics holds Prometheus metrics for the chatd subsystem.
|
||||
type Metrics struct {
|
||||
Chats *prometheus.GaugeVec
|
||||
MessageCount *prometheus.HistogramVec
|
||||
PromptSizeBytes *prometheus.HistogramVec
|
||||
ToolResultSizeBytes *prometheus.HistogramVec
|
||||
ToolResultTruncatedTotal *prometheus.CounterVec
|
||||
ToolErrorsTotal *prometheus.CounterVec
|
||||
TTFTSeconds *prometheus.HistogramVec
|
||||
CompactionTotal *prometheus.CounterVec
|
||||
StepsTotal *prometheus.CounterVec
|
||||
StreamRetriesTotal *prometheus.CounterVec
|
||||
StreamBufferDroppedTotal prometheus.Counter
|
||||
Chats *prometheus.GaugeVec
|
||||
MessageCount *prometheus.HistogramVec
|
||||
PromptSizeBytes *prometheus.HistogramVec
|
||||
ToolResultSizeBytes *prometheus.HistogramVec
|
||||
ToolResultTruncatedTotal *prometheus.CounterVec
|
||||
ToolErrorsTotal *prometheus.CounterVec
|
||||
TTFTSeconds *prometheus.HistogramVec
|
||||
CompactionTotal *prometheus.CounterVec
|
||||
StepsTotal *prometheus.CounterVec
|
||||
StreamRetriesTotal *prometheus.CounterVec
|
||||
StreamBufferDroppedTotal prometheus.Counter
|
||||
FindToolsCallsTotal prometheus.Counter
|
||||
FindToolsEmptyTotal prometheus.Counter
|
||||
FindToolsMatchCount prometheus.Histogram
|
||||
FindToolsActivationsTotal prometheus.Counter
|
||||
}
|
||||
|
||||
// NewMetrics creates a new Metrics instance registered with the
|
||||
@@ -109,6 +113,31 @@ func NewMetrics(reg prometheus.Registerer) *Metrics {
|
||||
Name: "stream_retries_total",
|
||||
Help: "Total LLM stream retries.",
|
||||
}, []string{"provider", "model", "kind"}),
|
||||
FindToolsCallsTotal: factory.NewCounter(prometheus.CounterOpts{
|
||||
Namespace: metricsNamespace,
|
||||
Subsystem: metricsSubsystem,
|
||||
Name: "find_tools_calls_total",
|
||||
Help: "Total find_tools calls.",
|
||||
}),
|
||||
FindToolsEmptyTotal: factory.NewCounter(prometheus.CounterOpts{
|
||||
Namespace: metricsNamespace,
|
||||
Subsystem: metricsSubsystem,
|
||||
Name: "find_tools_empty_total",
|
||||
Help: "Total find_tools calls with no matches.",
|
||||
}),
|
||||
FindToolsMatchCount: factory.NewHistogram(prometheus.HistogramOpts{
|
||||
Namespace: metricsNamespace,
|
||||
Subsystem: metricsSubsystem,
|
||||
Name: "find_tools_match_count",
|
||||
Help: "Number of matches returned by find_tools calls.",
|
||||
Buckets: prometheus.LinearBuckets(0, 2, 11),
|
||||
}),
|
||||
FindToolsActivationsTotal: factory.NewCounter(prometheus.CounterOpts{
|
||||
Namespace: metricsNamespace,
|
||||
Subsystem: metricsSubsystem,
|
||||
Name: "find_tools_activations_total",
|
||||
Help: "Total deferred tool activations returned by find_tools.",
|
||||
}),
|
||||
StreamBufferDroppedTotal: factory.NewCounter(prometheus.CounterOpts{
|
||||
Namespace: metricsNamespace,
|
||||
Subsystem: metricsSubsystem,
|
||||
|
||||
Reference in New Issue
Block a user