diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index 2d2a88a8b..aadcfc672 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -19,6 +19,7 @@ import ( "github.com/docker/docker-agent/pkg/model/provider/base" "github.com/docker/docker-agent/pkg/modelsdev" "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tools/codemode" ) // Safety returns the author-declared default set via WithSafety, empty @@ -765,3 +766,322 @@ func TestAgentToolsRecoversWhenUnderlyingToolsetDies(t *testing.T) { assert.Equal(t, 1, stub.startCalls) assert.Equal(t, 1, stub.restartCalls) } + +// countingToolSet is a Startable ToolSet with a stable description, a +// mutable start error and start/stop counters, used to simulate healthy and +// failing MCP-like toolsets aggregated inside a code-mode wrapper. +type countingToolSet struct { + desc string + startErr error + start int + stop int + stubs []tools.Tool +} + +var ( + _ tools.ToolSet = (*countingToolSet)(nil) + _ tools.Startable = (*countingToolSet)(nil) + _ tools.Describer = (*countingToolSet)(nil) +) + +func (c *countingToolSet) Describe() string { return c.desc } + +func (c *countingToolSet) Start(context.Context) error { + c.start++ + return c.startErr +} + +func (c *countingToolSet) Stop(context.Context) error { + c.stop++ + return nil +} + +func (c *countingToolSet) Tools(context.Context) ([]tools.Tool, error) { return c.stubs, nil } + +// TestAgentToolsCodeModePartialStartKeepsCodeModeAvailable is the regression +// test for #3978: with code_mode_tools enabled every toolset is aggregated +// into a single codemode wrapper, and one failing MCP server used to take +// down the whole wrapper — run_tools_with_javascript disappeared and the +// healthy toolsets were rolled back. A failing inner toolset must instead +// degrade the wrapper: healthy declarations stay available, the failure is +// warned about once per streak, and the failed toolset is retried on +// subsequent turns. +func TestAgentToolsCodeModePartialStartKeepsCodeModeAvailable(t *testing.T) { + t.Parallel() + + healthy := &countingToolSet{ + desc: "fetch built-in", + stubs: []tools.Tool{{Name: "fetch_url", Parameters: map[string]any{}}}, + } + failing := &countingToolSet{ + desc: "mcp(ref=broken)", + startErr: errors.New("connection refused"), + stubs: []tools.Tool{{Name: "broken_tool", Parameters: map[string]any{}}}, + } + a := New("root", "test", WithToolSets(codemode.Wrap(healthy, failing))) + + // Turn 1: code mode stays available with the healthy toolset only. + got, err := a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1, "turn 1: run_tools_with_javascript must survive a failing inner toolset") + assert.Equal(t, "run_tools_with_javascript", got[0].Name) + assert.Contains(t, got[0].Description, "FetchUrl", "healthy toolset must stay declared") + assert.NotContains(t, got[0].Description, "BrokenTool", "failed toolset must be omitted") + assert.Equal(t, 0, healthy.stop, "healthy toolset must not be rolled back on a peer's failure") + + warnings := a.DrainWarnings() + require.Len(t, warnings, 1, "turn 1: the inner failure must still be surfaced") + assert.Contains(t, warnings[0], "start failed") + assert.Contains(t, warnings[0], "mcp(ref=broken)") + assert.Contains(t, warnings[0], "connection refused") + + // Turn 2: the failed toolset is retried, without duplicate warnings and + // without restarting its healthy peer. + got, err = a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1, "turn 2: code mode must stay available while the failure persists") + assert.Equal(t, 2, failing.start, "turn 2: failed toolset must be retried") + assert.Equal(t, 1, healthy.start, "turn 2: healthy toolset must not be restarted") + assert.Empty(t, a.DrainWarnings(), "turn 2: no duplicate warning on repeated failure") + + // Turn 3: recovery — the toolset's declarations reappear, silently. + failing.startErr = nil + got, err = a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Contains(t, got[0].Description, "BrokenTool", "recovered toolset must be declared again") + assert.Empty(t, a.DrainWarnings(), "recovery must be silent") +} + +// dyingToolSet extends countingToolSet with live lifecycle reporting +// (tools.StartReporter) and in-place recovery (tools.Restartable), +// mimicking a supervisor-backed MCP toolset whose session can die in the +// background after a successful start. +type dyingToolSet struct { + countingToolSet + + started bool + restarts int + restartErr error +} + +var ( + _ tools.StartReporter = (*dyingToolSet)(nil) + _ tools.Restartable = (*dyingToolSet)(nil) +) + +func (d *dyingToolSet) Start(ctx context.Context) error { + if err := d.countingToolSet.Start(ctx); err != nil { + return err + } + d.started = true + return nil +} + +func (d *dyingToolSet) Restart(context.Context) error { + d.restarts++ + if d.restartErr != nil { + return d.restartErr + } + d.started = true + return nil +} + +func (d *dyingToolSet) IsStarted() bool { return d.started } + +// TestAgentToolsCodeModeInnerDiesAfterStart covers the full agent-level arc +// of an MCP-like inner toolset inside the codemode wrapper that starts +// successfully and later dies (e.g. background session loss): the composite +// must detect the death through the inner's StartReporter, degrade — the +// healthy peer and run_tools_with_javascript stay available — warn once, +// recover the dead inner via Restart (not a blind Start), and restore its +// declarations silently once the recovery succeeds. +func TestAgentToolsCodeModeInnerDiesAfterStart(t *testing.T) { + t.Parallel() + + healthy := &countingToolSet{ + desc: "fetch built-in", + stubs: []tools.Tool{{Name: "fetch_url", Parameters: map[string]any{}}}, + } + dying := &dyingToolSet{countingToolSet: countingToolSet{ + desc: "mcp(ref=github)", + stubs: []tools.Tool{{Name: "list_issues", Parameters: map[string]any{}}}, + }} + a := New("root", "test", WithToolSets(codemode.Wrap(healthy, dying))) + + // Turn 1: initial success — both toolsets declared, no warnings. + got, err := a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Contains(t, got[0].Description, "FetchUrl") + assert.Contains(t, got[0].Description, "ListIssues") + assert.Empty(t, a.DrainWarnings()) + assert.Equal(t, 1, dying.start) + + // Background death: the inner reports dead and can't come back yet. + dying.started = false + dying.restartErr = errors.New("session lost") + + // Turn 2: degraded — recovery goes through Restart, the failure is + // warned once, the healthy peer and the wrapper stay available. + got, err = a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1, "turn 2: run_tools_with_javascript must survive an inner death") + assert.Contains(t, got[0].Description, "FetchUrl", "healthy toolset must stay declared") + assert.NotContains(t, got[0].Description, "ListIssues", "dead toolset must be omitted") + assert.Equal(t, 1, dying.restarts, "turn 2: dead inner must be recovered via Restart") + assert.Equal(t, 1, dying.start, "turn 2: dead inner must not be blindly re-Started") + assert.Equal(t, 1, healthy.start, "turn 2: healthy peer must not be restarted") + + warnings := a.DrainWarnings() + require.Len(t, warnings, 1, "turn 2: the death must be surfaced once") + assert.Contains(t, warnings[0], "mcp(ref=github)") + assert.Contains(t, warnings[0], "session lost") + + // Turn 3: still dead — retried via Restart, no duplicate warning. + got, err = a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Equal(t, 2, dying.restarts, "turn 3: recovery retries must keep using Restart") + assert.Equal(t, 1, dying.start) + assert.Empty(t, a.DrainWarnings(), "turn 3: no duplicate warning on repeated failure") + + // Turn 4: Restart succeeds — declarations reappear, silently. + dying.restartErr = nil + got, err = a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Contains(t, got[0].Description, "ListIssues", "recovered toolset must be declared again") + assert.Equal(t, 3, dying.restarts) + assert.Equal(t, 1, dying.start, "recovery must never fall back to a blind Start") + assert.Empty(t, a.DrainWarnings(), "recovery must be silent") +} + +// TestAgentToolsCodeModeInitialAuthDeferralStaysSilent is the regression +// test for the initial-OAuth-deferral arc inside the codemode wrapper: an +// inner toolset that defers on authorization at startup (it never worked +// before) must stay silent across turns, even though the composite latches +// started on the first partial start and every later Start takes the +// recovery path. Before the LostAfterStart classification, turn 2 misread +// the retried deferral as a formerly-healthy toolset dying and emitted the +// "needs re-authentication" notice. +func TestAgentToolsCodeModeInitialAuthDeferralStaysSilent(t *testing.T) { + t.Parallel() + + healthy := &countingToolSet{ + desc: "fetch built-in", + stubs: []tools.Tool{{Name: "fetch_url", Parameters: map[string]any{}}}, + } + unauthorized := &countingToolSet{ + desc: "mcp(ref=notion)", + startErr: &tools.AuthorizationRequiredError{URL: "https://example.test/mcp"}, + stubs: []tools.Tool{{Name: "search_pages", Parameters: map[string]any{}}}, + } + a := New("root", "test", WithToolSets(codemode.Wrap(healthy, unauthorized))) + + // Turns 1-3: the deferral is retried every turn and must never warn — + // the OAuth dialog appears naturally on the first interactive turn. + for turn := 1; turn <= 3; turn++ { + got, err := a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1, "turn %d: code mode must stay available", turn) + assert.Contains(t, got[0].Description, "FetchUrl", "turn %d: healthy toolset must stay declared", turn) + assert.NotContains(t, got[0].Description, "SearchPages", "turn %d: deferred toolset must be omitted", turn) + assert.Empty(t, a.DrainWarnings(), "turn %d: initial OAuth deferral must stay silent", turn) + assert.Equal(t, turn, unauthorized.start, "turn %d: deferred toolset must be retried", turn) + assert.Equal(t, 1, healthy.start, "turn %d: healthy peer must not be restarted", turn) + } + + // Authorization completes: declarations appear, still silently. + unauthorized.startErr = nil + got, err := a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Contains(t, got[0].Description, "SearchPages", "authorized toolset must be declared") + assert.Empty(t, a.DrainWarnings(), "successful authorization must be silent") +} + +// TestAgentToolsCodeModeTotalStartFailureNotExposed verifies that when +// every inner toolset fails and none is available, the codemode wrapper +// fails like a regular toolset: it is not latched, so +// run_tools_with_javascript is not exposed with an empty function list, +// the failure is warned once per streak, and a cold-start retry runs every +// turn until recovery. +func TestAgentToolsCodeModeTotalStartFailureNotExposed(t *testing.T) { + t.Parallel() + + failingA := &countingToolSet{ + desc: "mcp(ref=broken-a)", + startErr: errors.New("connection refused"), + stubs: []tools.Tool{{Name: "tool_a", Parameters: map[string]any{}}}, + } + failingB := &countingToolSet{ + desc: "mcp(ref=broken-b)", + startErr: errors.New("no such host"), + stubs: []tools.Tool{{Name: "tool_b", Parameters: map[string]any{}}}, + } + a := New("root", "test", WithToolSets(codemode.Wrap(failingA, failingB))) + + // Turn 1: nothing usable — code mode must not be listed at all. + got, err := a.Tools(t.Context()) + require.NoError(t, err) + assert.Empty(t, got, "turn 1: code mode must not be exposed with an empty function list") + + warnings := a.DrainWarnings() + require.Len(t, warnings, 1, "turn 1: the total failure must be surfaced once") + assert.Contains(t, warnings[0], "start failed") + assert.Contains(t, warnings[0], "mcp(ref=broken-a)") + assert.Contains(t, warnings[0], "mcp(ref=broken-b)") + + // Turn 2: cold retry, no duplicate warning. + got, err = a.Tools(t.Context()) + require.NoError(t, err) + assert.Empty(t, got, "turn 2: code mode must stay unlisted while the total failure persists") + assert.Equal(t, 2, failingA.start, "turn 2: total failure must keep the cold-start retry") + assert.Equal(t, 2, failingB.start, "turn 2: total failure must keep the cold-start retry") + assert.Empty(t, a.DrainWarnings(), "turn 2: no duplicate warning on repeated failure") + + // Turn 3: recovery — code mode appears with both declarations, silently. + failingA.startErr = nil + failingB.startErr = nil + got, err = a.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Contains(t, got[0].Description, "ToolA", "recovered toolset must be declared") + assert.Contains(t, got[0].Description, "ToolB", "recovered toolset must be declared") + assert.Empty(t, a.DrainWarnings(), "recovery must be silent") +} + +// TestAgentToolsCodeModeTotalMixedFailureWarnsRealCause pins the mixed +// total failure (one inner deferred on OAuth, one failed for real): the +// real failure must be warned about like any other start failure — the +// deferral next to it must not reclassify the whole batch as a silent +// auth-required wait — and code mode stays unlisted since nothing is +// usable. +func TestAgentToolsCodeModeTotalMixedFailureWarnsRealCause(t *testing.T) { + t.Parallel() + + unauthorized := &countingToolSet{ + desc: "mcp(ref=notion)", + startErr: &tools.AuthorizationRequiredError{URL: "https://example.test/mcp"}, + stubs: []tools.Tool{{Name: "search_pages", Parameters: map[string]any{}}}, + } + failing := &countingToolSet{ + desc: "mcp(ref=broken)", + startErr: errors.New("connection refused"), + stubs: []tools.Tool{{Name: "broken_tool", Parameters: map[string]any{}}}, + } + a := New("root", "test", WithToolSets(codemode.Wrap(unauthorized, failing))) + + got, err := a.Tools(t.Context()) + require.NoError(t, err) + assert.Empty(t, got, "code mode must not be exposed while every inner toolset is down") + + warnings := a.DrainWarnings() + require.Len(t, warnings, 1, "the real failure must be surfaced despite the OAuth deferral next to it") + assert.Contains(t, warnings[0], "start failed") + assert.Contains(t, warnings[0], "mcp(ref=broken)") + assert.Contains(t, warnings[0], "connection refused") + assert.NotContains(t, warnings[0], "re-authentication", + "an initial mixed failure must not read as a re-auth notice") +} diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index 2909a74e3..75a984f60 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -1812,7 +1812,8 @@ func (r *LocalRuntime) emitToolsProgressively(ctx context.Context, a *agent.Agen // failure-reported flag here would suppress the *real* // failure (e.g. server 4xx on the eventual interactive // retry) that the user actually needs to see. - if tools.IsAuthorizationRequired(err) { + switch { + case tools.IsAuthorizationRequired(err): // Two cases: // 1. Initial startup deferral (toolset never ran): the // OAuth dialog will appear naturally on the first user @@ -1828,8 +1829,6 @@ func (r *LocalRuntime) emitToolsProgressively(ctx context.Context, a *agent.Agen } else { slog.DebugContext(ctx, "Toolset deferred until first message", "agent", a.Name(), "toolset", desc, "reason", err) } - continue - } // Route real failures through the agent's warning // channel so the TUI surfaces a persistent, // user-visible notice that includes the actual @@ -1838,13 +1837,19 @@ func (r *LocalRuntime) emitToolsProgressively(ctx context.Context, a *agent.Agen // once-per-streak guard as ensureToolSetsAreStarted // so a failing toolset doesn't flood the UI with a // new warning every time the agent is restarted. - if !startable.ShouldReportFailure() { + case startable.ShouldReportFailure(): + slog.WarnContext(ctx, "Toolset start failed; skipping", "agent", a.Name(), "toolset", desc, "error", err) + a.AddToolWarning(fmt.Sprintf("%s start failed: %v", desc, err)) + default: slog.DebugContext(ctx, "Toolset still unavailable; skipping", "agent", a.Name(), "toolset", desc, "error", err) + } + // A partial start leaves the composite latched and usable: + // fall through so its healthy subset (and the composite's + // own wrapper tool) is still listed and counted. Fully + // failed toolsets have nothing to list — skip them. + if !tools.IsPartialStart(err) { continue } - slog.WarnContext(ctx, "Toolset start failed; skipping", "agent", a.Name(), "toolset", desc, "error", err) - a.AddToolWarning(fmt.Sprintf("%s start failed: %v", desc, err)) - continue } } diff --git a/pkg/runtime/runtime_test.go b/pkg/runtime/runtime_test.go index dcb589e0e..d5637a504 100644 --- a/pkg/runtime/runtime_test.go +++ b/pkg/runtime/runtime_test.go @@ -1671,6 +1671,55 @@ func TestEmitStartupInfo_RecoveryAuthNoticeEmittedOnce(t *testing.T) { require.Len(t, noticesPhase4, 1, "fresh failure after streak reset must emit a new notice") } +// TestEmitStartupInfo_PartialStartStillCountsCodeModeTools is the runtime +// regression test for #3978 degraded startup: when the codemode composite +// starts partially (one inner MCP-like toolset fails), the composite is +// latched as started and emitToolsProgressively must keep listing it — the +// run_tools_with_javascript wrapper is counted in the sidebar's ToolsetInfo +// — while the inner failure still surfaces as a warning. +func TestEmitStartupInfo_PartialStartStillCountsCodeModeTools(t *testing.T) { + t.Parallel() + + prov := &mockProvider{id: "test/startup-model", stream: &mockStream{}} + + healthy := newStubToolSet(nil, []tools.Tool{{Name: "fetch_url", Parameters: map[string]any{}}}, nil) + failing := newStubToolSet(errors.New("connection refused"), nil, nil) + + root := agent.New("root", "agent", + agent.WithModel(prov), + agent.WithToolSets(codemode.Wrap(healthy, failing)), + ) + tm := team.New(team.WithAgents(root)) + + rt, err := NewLocalRuntime(t.Context(), tm, WithCurrentAgent("root"), WithModelStore(mockModelStore{})) + require.NoError(t, err) + + events := make(chan Event, 32) + rt.EmitStartupInfo(t.Context(), nil, NewChannelSink(events)) + close(events) + + var toolsetInfos []*ToolsetInfoEvent + var warning *WarningEvent + for e := range events { + switch ev := e.(type) { + case *ToolsetInfoEvent: + toolsetInfos = append(toolsetInfos, ev) + case *WarningEvent: + warning = ev + } + } + + require.NotEmpty(t, toolsetInfos, "expected at least one ToolsetInfo event") + last := toolsetInfos[len(toolsetInfos)-1] + assert.False(t, last.Loading, "final ToolsetInfo must report Loading=false") + assert.Equal(t, 1, last.AvailableTools, + "the degraded codemode wrapper must still be listed and counted (run_tools_with_javascript)") + + require.NotNil(t, warning, "the failed inner toolset must still surface a warning") + assert.Contains(t, warning.Message, "start failed") + assert.Contains(t, warning.Message, "connection refused") +} + // TestConfigureToolsetHandlers_ReachesThroughCodeModeWrapper is the // regression test for the sigma MCP OAuth bug: an agent with // code_mode_tools:true wraps all its toolsets in a single codemode diff --git a/pkg/tools/capabilities.go b/pkg/tools/capabilities.go index 511ceed22..feb6d17a2 100644 --- a/pkg/tools/capabilities.go +++ b/pkg/tools/capabilities.go @@ -13,6 +13,16 @@ type Startable interface { Stop(ctx context.Context) error } +// StartReporter is implemented by toolsets whose live lifecycle state can +// be queried independently of the StartableToolSet wrapper's latched state +// (e.g. an MCP toolset whose supervisor lost the session in the background). +// The wrapper consults it on Start to decide whether a recovery is needed, +// and composite toolsets consult their inner toolsets' reporters to detect +// an inner that started successfully and later died. +type StartReporter interface { + IsStarted() bool +} + // PeerDependent is implemented by toolsets whose Start reads from the // agent's other toolsets (e.g. the deferred aggregator lists its source // toolsets' tools). Callers that start an agent's toolsets concurrently diff --git a/pkg/tools/codemode/codemode.go b/pkg/tools/codemode/codemode.go index f3f87dc31..bf5cd8fe3 100644 --- a/pkg/tools/codemode/codemode.go +++ b/pkg/tools/codemode/codemode.go @@ -5,7 +5,9 @@ import ( "encoding/json" "errors" "fmt" + "slices" "strings" + "sync" "github.com/docker/docker-agent/pkg/tools" ) @@ -31,17 +33,48 @@ Available tools/functions: func Wrap(toolsets ...tools.ToolSet) tools.ToolSet { return &codeModeTool{ toolsets: toolsets, + states: make([]innerState, len(toolsets)), } } +// innerState tracks the lifecycle of one inner toolset. Never-started +// toolsets (innerIdle) are still listed so Tools() keeps working for +// direct, unstarted use; toolsets whose last Start attempt failed +// (innerFailed) and previously-running toolsets that died or failed to +// recover (innerLost) are omitted until a later recovery succeeds. +type innerState int8 + +const ( + innerIdle innerState = iota + innerStarted + innerFailed + innerLost +) + type codeModeTool struct { toolsets []tools.ToolSet + + // lifecycleMu serializes Start and Stop against each other so their + // state transitions cannot interleave (the outer StartableToolSet + // wrapper already single-flights them, but direct use of Wrap must + // stay safe). It is held across inner Start/Restart/Stop calls, so + // Tools/IsStarted/runJavascript must never take it. + lifecycleMu sync.Mutex + + // mu guards states and is only held for short in-memory sections: + // no inner-toolset method (Start, Restart, Stop, IsStarted, Tools) + // is ever called with mu held, so a wedged inner blocking on I/O + // cannot stall concurrent Tools/IsStarted/runJavascript. Lock order: + // lifecycleMu before mu; mu is a leaf. + mu sync.Mutex + states []innerState } // Verify interface compliance var ( _ tools.ToolSet = (*codeModeTool)(nil) _ tools.Startable = (*codeModeTool)(nil) + _ tools.StartReporter = (*codeModeTool)(nil) _ tools.Named = (*codeModeTool)(nil) _ tools.Elicitable = (*codeModeTool)(nil) _ tools.Sampleable = (*codeModeTool)(nil) @@ -63,13 +96,66 @@ func isExcludedTool(tool tools.Tool) bool { return tool.Category == "todo" } +// availableToolsets returns the inner toolsets whose tools may be exposed: +// every toolset except those whose last Start attempt failed (innerFailed, +// innerLost) and those that started successfully but report dead since (a +// dead inner cannot list its tools — e.g. an MCP toolset without a live +// session errors out — and would take run_tools_with_javascript down with +// it). +func (c *codeModeTool) availableToolsets() []tools.ToolSet { + states := c.snapshotStates() + + available := make([]tools.ToolSet, 0, len(c.toolsets)) + for i, t := range c.toolsets { + switch states[i] { + case innerFailed: + continue + case innerStarted, innerLost: + if innerDied(t) { + continue + } + } + available = append(available, t) + } + return available +} + +// snapshotStates returns a copy of the per-toolset states for lock-free +// inspection. The snapshot may lag an in-flight Start/Stop; that is the +// point — readers must not wait on inner lifecycle I/O. +func (c *codeModeTool) snapshotStates() []innerState { + c.mu.Lock() + defer c.mu.Unlock() + return slices.Clone(c.states) +} + +func (c *codeModeTool) state(i int) innerState { + c.mu.Lock() + defer c.mu.Unlock() + return c.states[i] +} + +func (c *codeModeTool) setState(i int, s innerState) { + c.mu.Lock() + defer c.mu.Unlock() + c.states[i] = s +} + +// innerDied reports whether a previously-started inner toolset says it is no +// longer running (e.g. its MCP session was lost in the background). Toolsets +// without a tools.StartReporter cannot report death and count as alive. +func innerDied(t tools.ToolSet) bool { + reporter, ok := tools.As[tools.StartReporter](t) + return ok && !reporter.IsStarted() +} + func (c *codeModeTool) Tools(ctx context.Context) ([]tools.Tool, error) { var ( functionsDoc []string excludedTools []tools.Tool ) - for _, toolset := range c.toolsets { + for _, toolset := range c.availableToolsets() { allTools, err := toolset.Tools(ctx) if err != nil { return nil, err @@ -113,30 +199,126 @@ func (c *codeModeTool) Tools(ctx context.Context) ([]tools.Tool, error) { return allTools, nil } +// Start brings every inner toolset up: cold starts (innerIdle), retries of +// toolsets whose initial start failed (innerFailed), and recoveries of +// toolsets that started successfully but died since — detected live through +// their tools.StartReporter, then remembered as innerLost. Recoveries of a +// Restartable inner go through Restart, never blindly through Start: an MCP +// supervisor's Start can be a no-op while it still holds the dead session, +// and Restart also waits for an in-flight background reconnect instead of +// racing it. This mirrors StartableToolSet's own recovery dispatch. +// +// A failing toolset does not abort its peers: healthy toolsets stay started +// and their functions remain exposed, while failed ones are omitted from +// Tools() until a later Start succeeds. Failures are reported through +// tools.PartialStartError so the StartableToolSet wrapper keeps +// run_tools_with_javascript available, still surfaces the warning, and — +// because IsStarted() reports false while any inner toolset is down — +// retries the failed subset on the next turn. When every inner toolset +// fails there is no healthy subset worth exposing: Start returns a +// tools.TotalStartError instead, so the wrapper stays unlatched — +// run_tools_with_javascript is not listed with an empty function list — +// and the next turn retries from cold. func (c *codeModeTool) Start(ctx context.Context) error { - var started []tools.Startable - var errs []error - for _, t := range c.toolsets { - if s, ok := tools.As[tools.Startable](t); ok { - if err := s.Start(ctx); err != nil { - errs = append(errs, err) - } else { - started = append(started, s) + c.lifecycleMu.Lock() + defer c.lifecycleMu.Unlock() + + var ( + errs []error + lost bool + ) + for i, t := range c.toolsets { + // The attempt runs without c.mu held so a wedged inner cannot + // stall Tools/IsStarted; its outcome is committed under c.mu. + next, err := startInner(ctx, t, c.state(i)) + c.setState(i, next) + if err != nil { + errs = append(errs, fmt.Errorf("%s: %w", tools.DescribeToolSet(t), err)) + if next == innerLost { + lost = true } } } - if len(errs) > 0 { - // Roll back successfully-started toolsets so we don't leave - // the system in a partially-started state. - for _, s := range started { - errs = append(errs, s.Stop(ctx)) + switch { + case len(errs) == 0: + return nil + case len(errs) == len(c.toolsets): + // Total failure: nothing is available, so degraded mode has no + // healthy subset to preserve. The dedicated non-partial type keeps + // the all-causes auth classification (a bare errors.Join would let + // one OAuth deferral hide a real failure). See the doc comment above. + return tools.NewTotalStartError(errs...) + default: + err := tools.NewPartialStartError(errs...) + // Only post-start loss (an inner that was running and died) may + // trigger the wrapper's recovery notice; retried initial failures + // must stay silent. + err.LostAfterStart = lost + return err + } +} + +// startInner runs one inner toolset's start/recovery attempt and returns +// the resulting state. It must be called without c.mu held: Start, Restart +// and the IsStarted death probe can block on I/O or take the inner's own +// locks. +func startInner(ctx context.Context, t tools.ToolSet, state innerState) (innerState, error) { + s, ok := tools.As[tools.Startable](t) + if !ok { + return innerStarted, nil + } + if state == innerStarted || state == innerLost { + if !innerDied(t) { + // Running, or a lost toolset whose supervisor already + // reconnected in the background: nothing to do. + return innerStarted, nil + } + state = innerLost + } + recovering := state == innerLost + var err error + if restarter, ok := tools.As[tools.Restartable](t); recovering && ok { + err = restarter.Restart(ctx) + } else { + err = s.Start(ctx) + } + if err != nil { + if recovering { + return innerLost, err } - return errors.Join(errs...) + return innerFailed, err } - return nil + return innerStarted, nil +} + +// IsStarted implements tools.StartReporter for the StartableToolSet wrapper: +// it reports false while any inner toolset is not running — never started, +// failed, or started-then-died per its own reporter — so the wrapper +// re-invokes Start on the next turn and the degraded subset is retried. +func (c *codeModeTool) IsStarted() bool { + states := c.snapshotStates() + + for i, t := range c.toolsets { + if states[i] != innerStarted { + return false + } + if innerDied(t) { + return false + } + } + return true } func (c *codeModeTool) Stop(ctx context.Context) error { + c.lifecycleMu.Lock() + defer c.lifecycleMu.Unlock() + + c.mu.Lock() + for i := range c.states { + c.states[i] = innerIdle + } + c.mu.Unlock() + var errs []error for _, t := range c.toolsets { if s, ok := tools.As[tools.Startable](t); ok { diff --git a/pkg/tools/codemode/codemode_test.go b/pkg/tools/codemode/codemode_test.go index d35413477..f7c6e48ef 100644 --- a/pkg/tools/codemode/codemode_test.go +++ b/pkg/tools/codemode/codemode_test.go @@ -3,7 +3,9 @@ package codemode import ( "context" "encoding/json" + "sync" "testing" + "time" "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/assert" @@ -245,9 +247,12 @@ func TestCodeModeTool_CallEcho(t *testing.T) { require.Empty(t, scriptResult.StdOut) } -// TestCodeModeTool_StartRollsBackOnError verifies that when one toolset fails -// to start, all successfully-started toolsets are stopped (rolled back). -func TestCodeModeTool_StartRollsBackOnError(t *testing.T) { +// TestCodeModeTool_StartKeepsHealthyToolsetsOnError verifies that when one +// toolset fails to start, its peers stay started (no rollback) and the +// failure is reported as a tools.PartialStartError so the StartableToolSet +// wrapper keeps run_tools_with_javascript available while still surfacing +// the error (#3978). +func TestCodeModeTool_StartKeepsHealthyToolsetsOnError(t *testing.T) { t.Parallel() failing := &testToolSet{startErr: assert.AnError} healthy := &testToolSet{} @@ -256,9 +261,380 @@ func TestCodeModeTool_StartRollsBackOnError(t *testing.T) { err := tool.Start(t.Context()) require.ErrorIs(t, err, assert.AnError) + require.True(t, tools.IsPartialStart(err), "partial failure must be reported as PartialStartError") assert.Equal(t, 1, failing.start, "failing toolset should have attempted start") assert.Equal(t, 1, healthy.start, "healthy toolset should have attempted start") - assert.Equal(t, 1, healthy.stop, "healthy toolset should be rolled back after failure") + assert.Equal(t, 0, healthy.stop, "healthy toolset must not be rolled back on a peer's failure") +} + +// TestCodeModeTool_PartialStartExposesHealthyTools verifies the degraded-mode +// contract: after a partial start, Tools() still returns +// run_tools_with_javascript with the healthy toolsets' declarations (the +// failed toolset's are omitted), scripts can call the healthy tools, the +// failed toolset alone is retried on the next Start, and a successful retry +// restores its declarations. +func TestCodeModeTool_PartialStartExposesHealthyTools(t *testing.T) { + t.Parallel() + healthy := &testToolSet{ + tools: []tools.Tool{ + { + Name: "fetch_url", + Handler: tools.NewHandler(func(ctx context.Context, args map[string]any) (*tools.ToolCallResult, error) { + return tools.ResultSuccess("fetched"), nil + }), + }, + {Name: "todo_write", Category: "todo"}, + }, + } + failing := &testToolSet{ + startErr: assert.AnError, + tools: []tools.Tool{{Name: "broken_tool"}}, + } + + tool := Wrap(healthy, failing) + startable := tool.(tools.Startable) + reporter := tool.(tools.StartReporter) + + require.Error(t, startable.Start(t.Context())) + assert.False(t, reporter.IsStarted(), "degraded wrapper must report unstarted so the failed toolset is retried") + + allTools, err := tool.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, allTools, 2) + assert.Equal(t, "run_tools_with_javascript", allTools[0].Name) + assert.Contains(t, allTools[0].Description, "declare function FetchUrl", "healthy toolset must stay declared") + assert.NotContains(t, allTools[0].Description, "BrokenTool", "failed toolset must be omitted") + assert.Equal(t, "todo_write", allTools[1].Name, "todo exclusion must be preserved in degraded mode") + + result, err := allTools[0].Handler(t.Context(), tools.ToolCall{ + Function: tools.FunctionCall{ + Arguments: `{"script":"return fetch_url();"}`, + }, + }, tools.NopRuntime{}) + require.NoError(t, err) + var scriptResult ScriptResult + require.NoError(t, json.Unmarshal([]byte(result.Output), &scriptResult)) + assert.Equal(t, "fetched", scriptResult.Value, "healthy tools must stay callable from scripts") + + // Retry: only the failed toolset is started again. + require.Error(t, startable.Start(t.Context())) + assert.Equal(t, 1, healthy.start, "healthy toolset must not be restarted on retry") + assert.Equal(t, 2, failing.start, "failed toolset must be retried") + + // Recovery: the toolset comes back and its declarations reappear. + failing.startErr = nil + require.NoError(t, startable.Start(t.Context())) + assert.True(t, reporter.IsStarted()) + + allTools, err = tool.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, allTools, 2) + assert.Contains(t, allTools[0].Description, "declare function BrokenTool", "recovered toolset must be declared again") +} + +// reportingToolSet is a testToolSet that also reports its live lifecycle +// state (tools.StartReporter) and supports in-place recovery +// (tools.Restartable), like a supervisor-backed MCP toolset. A background +// session death is simulated by flipping started to false. +type reportingToolSet struct { + testToolSet + + started bool + restarts int + restartErr error +} + +var ( + _ tools.StartReporter = (*reportingToolSet)(nil) + _ tools.Restartable = (*reportingToolSet)(nil) +) + +func (r *reportingToolSet) Start(ctx context.Context) error { + if err := r.testToolSet.Start(ctx); err != nil { + return err + } + r.started = true + return nil +} + +func (r *reportingToolSet) Restart(context.Context) error { + r.restarts++ + if r.restartErr != nil { + return r.restartErr + } + r.started = true + return nil +} + +func (r *reportingToolSet) IsStarted() bool { return r.started } + +// TestCodeModeTool_InnerDeathIsDetectedAndRecoveredViaRestart covers the +// "started successfully, then died" arc for a supervisor-backed inner +// toolset (e.g. MCP): the composite must detect the death through the +// inner's tools.StartReporter, degrade (omit the dead inner while keeping +// the healthy one listed), and recover it via Restart — not Start, which +// can be a no-op on a supervisor still holding the dead session. A failed +// recovery surfaces as a PartialStartError and is retried. +func TestCodeModeTool_InnerDeathIsDetectedAndRecoveredViaRestart(t *testing.T) { + t.Parallel() + healthy := &testToolSet{tools: []tools.Tool{{Name: "fetch_url"}}} + flaky := &reportingToolSet{testToolSet: testToolSet{tools: []tools.Tool{{Name: "flaky_tool"}}}} + + tool := Wrap(healthy, flaky) + startable := tool.(tools.Startable) + reporter := tool.(tools.StartReporter) + + // Initial start: everything up. + require.NoError(t, startable.Start(t.Context())) + assert.True(t, reporter.IsStarted()) + assert.Equal(t, 1, flaky.start) + + // The inner dies in the background: the composite reports degraded and + // omits the dead inner's declarations so listing keeps working. + flaky.started = false + assert.False(t, reporter.IsStarted(), "a dead inner must degrade the composite") + + allTools, err := tool.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, allTools, 1) + assert.Contains(t, allTools[0].Description, "FetchUrl", "healthy toolset must stay declared") + assert.NotContains(t, allTools[0].Description, "FlakyTool", "dead inner must be omitted") + + // Failed recovery: the dead inner is recovered via Restart, its peers + // are left alone, and the composite stays degraded. + flaky.restartErr = assert.AnError + err = startable.Start(t.Context()) + require.ErrorIs(t, err, assert.AnError) + assert.True(t, tools.IsPartialStart(err)) + assert.Equal(t, 1, flaky.restarts, "dead inner must be recovered via Restart") + assert.Equal(t, 1, flaky.start, "dead inner must not be blindly re-Started") + assert.Equal(t, 1, healthy.start, "healthy peer must not be restarted") + assert.False(t, reporter.IsStarted()) + + // Successful Restart on the next recovery attempt brings the inner back. + flaky.restartErr = nil + flaky.startErr = assert.AnError // Start must not be used while recovering + require.NoError(t, startable.Start(t.Context())) + assert.Equal(t, 2, flaky.restarts) + assert.True(t, reporter.IsStarted()) + + allTools, err = tool.Tools(t.Context()) + require.NoError(t, err) + assert.Contains(t, allTools[0].Description, "FlakyTool", "recovered inner must be declared again") +} + +// blockingToolSet is an inner toolset whose Start wedges: it ignores ctx +// and blocks until release is closed, like an unresponsive MCP server. +// entered is closed once Start is inside the blocking section. +type blockingToolSet struct { + entered chan struct{} + release chan struct{} +} + +var ( + _ tools.ToolSet = (*blockingToolSet)(nil) + _ tools.Startable = (*blockingToolSet)(nil) +) + +func newBlockingToolSet() *blockingToolSet { + return &blockingToolSet{ + entered: make(chan struct{}), + release: make(chan struct{}), + } +} + +func (b *blockingToolSet) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } + +func (b *blockingToolSet) Start(context.Context) error { + close(b.entered) + <-b.release + return nil +} + +func (b *blockingToolSet) Stop(context.Context) error { return nil } + +// TestCodeModeTool_ToolsNotBlockedByWedgedInnerStart pins the short-lock +// contract behind #3978: c.mu is never held across inner lifecycle calls, +// so while one inner's Start is wedged (ignoring its context), Tools() and +// IsStarted() return promptly and run_tools_with_javascript keeps exposing +// the healthy peer instead of queueing behind the mutex. +func TestCodeModeTool_ToolsNotBlockedByWedgedInnerStart(t *testing.T) { + t.Parallel() + healthy := &testToolSet{tools: []tools.Tool{{Name: "fetch_url"}}} + wedged := newBlockingToolSet() + releaseWedged := sync.OnceFunc(func() { close(wedged.release) }) + defer releaseWedged() // unblock the Start goroutine even if an assertion fails first + + tool := Wrap(healthy, wedged) + + // Capture the test context here: goroutines below may outlive a t.Fatal + // path and must not call t.Context() themselves. + ctx := t.Context() + startDone := make(chan error, 1) + go func() { startDone <- tool.(tools.Startable).Start(ctx) }() + + select { + case <-wedged.entered: + case <-time.After(5 * time.Second): + t.Fatal("wedged inner Start was never entered") + } + + type listing struct { + started bool + tools []tools.Tool + err error + } + listed := make(chan listing, 1) + go func() { + started := tool.(tools.StartReporter).IsStarted() + allTools, err := tool.Tools(ctx) + listed <- listing{started: started, tools: allTools, err: err} + }() + + select { + case res := <-listed: + require.NoError(t, res.err) + assert.False(t, res.started, "composite must report unstarted while an inner Start is in flight") + require.NotEmpty(t, res.tools) + assert.Equal(t, "run_tools_with_javascript", res.tools[0].Name) + assert.Contains(t, res.tools[0].Description, "FetchUrl", "healthy peer must stay declared") + case <-time.After(5 * time.Second): + t.Fatal("Tools() blocked behind a wedged inner Start") + } + + releaseWedged() + select { + case err := <-startDone: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("Start did not return after the wedged inner was released") + } +} + +// TestCodeModeTool_PartialStartAuthClassification pins how a partial start +// is classified for the authorization special-handling: a batch where every +// failed inner deferred on OAuth is auth-only (silently deferrable), while a +// mixed batch (auth + real failure) must NOT satisfy IsAuthorizationRequired +// — otherwise the real failure would be hidden behind the silent deferral. +// A healthy peer keeps each batch genuinely partial; total failures are +// tools.TotalStartError (see TestCodeModeTool_TotalStartFailureIsNotPartial). +func TestCodeModeTool_PartialStartAuthClassification(t *testing.T) { + t.Parallel() + authErr := &tools.AuthorizationRequiredError{URL: "https://example.test/mcp"} + + authOnly := Wrap(&testToolSet{}, &testToolSet{startErr: authErr}).(tools.Startable) + err := authOnly.Start(t.Context()) + require.True(t, tools.IsPartialStart(err)) + assert.True(t, tools.IsAuthorizationRequired(err), + "an auth-only partial start must keep the silent OAuth-deferral handling") + + mixed := Wrap(&testToolSet{}, &testToolSet{startErr: authErr}, &testToolSet{startErr: assert.AnError}).(tools.Startable) + err = mixed.Start(t.Context()) + require.True(t, tools.IsPartialStart(err)) + assert.False(t, tools.IsAuthorizationRequired(err), + "a mixed batch must not be classified auth-only: the non-auth failure needs surfacing") + require.ErrorIs(t, err, assert.AnError, "the non-auth cause must stay reachable via errors.Is") +} + +// TestCodeModeTool_PartialStartLostAfterStartClassification pins how the +// composite classifies partial failures for the wrapper's recovery notice: +// an inner that never came up (initial failure) leaves LostAfterStart +// false, while an inner that started successfully and was lost since sets +// it, so only real post-start losses can mark the wrapper's recovery streak. +func TestCodeModeTool_PartialStartLostAfterStartClassification(t *testing.T) { + t.Parallel() + healthy := &testToolSet{} + flaky := &reportingToolSet{testToolSet: testToolSet{startErr: assert.AnError}} + + startable := Wrap(healthy, flaky).(tools.Startable) + + // Initial failure: the inner never started, so nothing was lost. + var partial *tools.PartialStartError + require.ErrorAs(t, startable.Start(t.Context()), &partial) + assert.False(t, partial.LostAfterStart, "an initial inner failure is not a post-start loss") + + // The inner recovers, then dies in the background and fails to restart. + flaky.startErr = nil + require.NoError(t, startable.Start(t.Context())) + flaky.started = false + flaky.restartErr = assert.AnError + + require.ErrorAs(t, startable.Start(t.Context()), &partial) + assert.True(t, partial.LostAfterStart, "a started-then-lost inner must be classified as post-start loss") +} + +// TestCodeModeTool_TotalStartFailureIsNotPartial pins the total-failure +// contract: when every inner toolset fails and none is available, Start +// returns a tools.TotalStartError instead of a tools.PartialStartError. A +// PartialStartError would latch the StartableToolSet wrapper as started and +// expose run_tools_with_javascript with an empty function list; the +// non-partial error keeps the wrapper unlatched so the next Start is a cold +// retry. +func TestCodeModeTool_TotalStartFailureIsNotPartial(t *testing.T) { + t.Parallel() + failingA := &testToolSet{startErr: assert.AnError, tools: []tools.Tool{{Name: "tool_a"}}} + failingB := &testToolSet{startErr: assert.AnError, tools: []tools.Tool{{Name: "tool_b"}}} + + tool := Wrap(failingA, failingB) + wrapper := tools.NewStartable(tool) + + err := wrapper.Start(t.Context()) + require.ErrorIs(t, err, assert.AnError) + assert.False(t, tools.IsPartialStart(err), "total failure must not be reported as partial") + assert.False(t, wrapper.IsStarted(), "total failure must not latch the wrapper") + + // Cold retry: every inner is started again, and a successful retry + // exposes all declarations. + failingA.startErr = nil + failingB.startErr = nil + require.NoError(t, wrapper.Start(t.Context())) + assert.True(t, wrapper.IsStarted()) + assert.Equal(t, 2, failingA.start) + assert.Equal(t, 2, failingB.start) + + allTools, err := tool.Tools(t.Context()) + require.NoError(t, err) + require.Len(t, allTools, 1) + assert.Contains(t, allTools[0].Description, "ToolA") + assert.Contains(t, allTools[0].Description, "ToolB") +} + +// TestCodeModeTool_TotalMixedFailureIsNotAuthRequired pins the auth +// classification of a mixed total failure (one inner deferred on OAuth, one +// failed for real): it must be non-partial — so the wrapper does not latch +// — and must NOT satisfy IsAuthorizationRequired, otherwise the real +// failure would be suppressed behind the silent auth-deferral handling. +// Both causes stay reachable via errors.Is/errors.As. +func TestCodeModeTool_TotalMixedFailureIsNotAuthRequired(t *testing.T) { + t.Parallel() + authErr := &tools.AuthorizationRequiredError{URL: "https://example.test/mcp"} + + tool := Wrap(&testToolSet{startErr: authErr}, &testToolSet{startErr: assert.AnError}).(tools.Startable) + err := tool.Start(t.Context()) + require.Error(t, err) + assert.False(t, tools.IsPartialStart(err), "total failure must not be reported as partial") + assert.False(t, tools.IsAuthorizationRequired(err), + "a mixed total failure must not be classified auth-only: the real failure needs surfacing") + require.ErrorIs(t, err, assert.AnError, "the non-auth cause must stay reachable via errors.Is") + var target *tools.AuthorizationRequiredError + require.ErrorAs(t, err, &target, "the auth cause must stay reachable via errors.As") +} + +// TestCodeModeTool_TotalAuthDeferralKeepsAuthClassification pins that a +// total failure where every inner deferred on OAuth is still classified +// authorization-required even though it is no longer a PartialStartError, +// so the initial deferral stays silent and the dialog appears naturally on +// the first interactive turn. +func TestCodeModeTool_TotalAuthDeferralKeepsAuthClassification(t *testing.T) { + t.Parallel() + authErr := &tools.AuthorizationRequiredError{URL: "https://example.test/mcp"} + + tool := Wrap(&testToolSet{startErr: authErr}).(tools.Startable) + err := tool.Start(t.Context()) + require.Error(t, err) + assert.False(t, tools.IsPartialStart(err), "total failure must not be reported as partial") + assert.True(t, tools.IsAuthorizationRequired(err), + "an all-auth total failure must keep the silent OAuth-deferral handling") } // TestCodeModeTool_StartStopWrappedToolSet verifies that Start/Stop find diff --git a/pkg/tools/codemode/exec.go b/pkg/tools/codemode/exec.go index 9808b68c2..899c60641 100644 --- a/pkg/tools/codemode/exec.go +++ b/pkg/tools/codemode/exec.go @@ -75,8 +75,10 @@ func (c *codeModeTool) runJavascript(ctx context.Context, rt tools.Runtime, scri ) _ = vm.Set("console", console(&stdOut, &stdErr)) - // Inject every tool as a javascript function. - for _, toolset := range c.toolsets { + // Inject every available tool as a javascript function. Toolsets whose + // start failed or that died since are omitted, matching the declarations + // listed by Tools(). + for _, toolset := range c.availableToolsets() { allTools, err := toolset.Tools(ctx) if err != nil { return ScriptResult{}, err diff --git a/pkg/tools/interactive.go b/pkg/tools/interactive.go index 300011c53..f0a87d3c6 100644 --- a/pkg/tools/interactive.go +++ b/pkg/tools/interactive.go @@ -59,7 +59,21 @@ func (e *AuthorizationRequiredError) Error() string { // signals that the toolset failed to start because OAuth is needed and the // caller chose to defer the prompt. Callers can use this to render a softer, // "needs auth" notice instead of a red error. +// +// PartialStartError and TotalStartError batch several inner-toolset +// failures and only count as authorization-required when all of their +// causes are (see the AuthOnly fields): matching a mixed batch through +// plain errors.As would hide the non-auth causes behind the silent +// auth-deferral handling. func IsAuthorizationRequired(err error) bool { + var partial *PartialStartError + if errors.As(err, &partial) { + return partial.AuthOnly + } + var total *TotalStartError + if errors.As(err, &total) { + return total.AuthOnly + } var target *AuthorizationRequiredError return errors.As(err, &target) } diff --git a/pkg/tools/startable.go b/pkg/tools/startable.go index a819b63ba..360f1c0e9 100644 --- a/pkg/tools/startable.go +++ b/pkg/tools/startable.go @@ -2,6 +2,7 @@ package tools import ( "context" + "errors" "fmt" "sync" @@ -45,12 +46,6 @@ type failureStreak struct { pending bool // true if the current streak's first failure is unreported } -// startReporter is implemented by toolsets whose live lifecycle state can be -// queried independently of the StartableToolSet wrapper's latched state. -type startReporter interface { - IsStarted() bool -} - func (f *failureStreak) fail() { if !f.active { f.active = true @@ -71,6 +66,128 @@ func (f *failureStreak) shouldReport() bool { return true } +// PartialStartError is returned from Start by composite toolsets (toolsets +// aggregating several inner toolsets, e.g. Code Mode) when only part of +// their inner toolsets came up. StartableToolSet treats it specially: the +// wrapper is latched as started so the healthy subset stays listed and +// usable, while the error still propagates so callers can warn about the +// failed subset. The composite must also implement StartReporter, returning +// false while degraded, so the next Start call takes the recovery path and +// retries the failed inner toolsets. +// +// PartialStartError is only for genuinely partial outcomes: a composite +// whose inner toolsets all failed should return a TotalStartError instead, +// so the wrapper stays unlatched and the next Start is a cold retry. +// +// Use NewPartialStartError so AuthOnly is classified from the causes. +type PartialStartError struct { + // Err aggregates the individual inner-toolset failures (usually an + // errors.Join of one error per failed toolset). + Err error + // AuthOnly is true when every individual cause is an + // authorization-required deferral (see IsAuthorizationRequired). A mixed + // batch (auth + real failure) must stay false so the non-auth cause is + // surfaced as a failure instead of being hidden behind the silent + // auth-deferral handling. + AuthOnly bool + // LostAfterStart is true when at least one failed inner toolset had + // previously started successfully and was lost since (died in the + // background or failed to recover). The composite sets it so + // StartableToolSet can tell a real recovery failure — worth the targeted + // re-auth notice — from a retried initial failure (e.g. an OAuth + // deferral), which must stay silent across turns even though the + // wrapper latched started on an earlier partial start. + LostAfterStart bool +} + +// allCausesAuthorizationRequired reports whether there is at least one +// non-nil cause and every one of them classifies as authorization-required. +// Aggregate start errors must be classified with these all-causes semantics: +// matching the batch through plain errors.As (ANY semantics) would let a +// single OAuth deferral hide the real failures joined next to it behind the +// silent auth-deferral handling. +func allCausesAuthorizationRequired(causes []error) bool { + authOnly := false + for _, cause := range causes { + if cause == nil { + continue + } + if !IsAuthorizationRequired(cause) { + return false + } + authOnly = true + } + return authOnly +} + +// NewPartialStartError joins the given per-toolset start failures and +// classifies the batch as auth-only when every (non-nil) cause reports +// IsAuthorizationRequired. LostAfterStart is left false: the composite must +// set it afterwards when a failed inner had already started successfully. +func NewPartialStartError(causes ...error) *PartialStartError { + return &PartialStartError{Err: errors.Join(causes...), AuthOnly: allCausesAuthorizationRequired(causes)} +} + +func (e *PartialStartError) Error() string { + if e == nil || e.Err == nil { + return "partial toolset start failure" + } + return e.Err.Error() +} + +// Unwrap exposes the underlying error(s) to errors.Is/errors.As, e.g. so +// IsAuthorizationRequired can detect a deferred-OAuth inner toolset. +func (e *PartialStartError) Unwrap() error { return e.Err } + +// IsPartialStart reports whether err (or any error wrapped by it) signals +// that a composite toolset started with only part of its inner toolsets +// available. +func IsPartialStart(err error) bool { + var target *PartialStartError + return errors.As(err, &target) +} + +// TotalStartError is the total-failure counterpart of PartialStartError, +// returned from Start by composite toolsets when every inner toolset +// failed. It is deliberately not a PartialStartError: with no healthy +// subset to preserve, StartableToolSet must not latch the wrapper as +// started, so the composite's own tool is not listed and the next Start +// retries from cold. +// +// The dedicated type exists so IsAuthorizationRequired applies the same +// all-causes classification as for partial failures: a bare errors.Join +// would match a single OAuth deferral via errors.As (ANY semantics) and +// silently suppress the real failures joined next to it. +// +// Use NewTotalStartError so AuthOnly is classified from the causes. +type TotalStartError struct { + // Err aggregates the individual inner-toolset failures (usually an + // errors.Join of one error per failed toolset). + Err error + // AuthOnly is true when every individual cause is an + // authorization-required deferral — same contract as + // PartialStartError.AuthOnly. + AuthOnly bool +} + +// NewTotalStartError joins the given per-toolset start failures and +// classifies the batch as auth-only when every (non-nil) cause reports +// IsAuthorizationRequired. +func NewTotalStartError(causes ...error) *TotalStartError { + return &TotalStartError{Err: errors.Join(causes...), AuthOnly: allCausesAuthorizationRequired(causes)} +} + +func (e *TotalStartError) Error() string { + if e == nil || e.Err == nil { + return "total toolset start failure" + } + return e.Err.Error() +} + +// Unwrap exposes the underlying error(s) to errors.Is/errors.As so the +// individual causes stay reachable. +func (e *TotalStartError) Unwrap() error { return e.Err } + // StartableToolSet wraps a ToolSet with lazy, single-flight start semantics. // This is the canonical way to manage toolset lifecycle. // @@ -120,7 +237,7 @@ func (s *StartableToolSet) Start(ctx context.Context) (err error) { recovering := false if s.started { - if reporter, ok := As[startReporter](s.ToolSet); !ok || reporter.IsStarted() { + if reporter, ok := As[StartReporter](s.ToolSet); !ok || reporter.IsStarted() { return nil } s.started = false @@ -177,6 +294,30 @@ func (s *StartableToolSet) Start(ctx context.Context) (err error) { }() if err := startable.Start(ctx); err != nil { s.startStreak.fail() + var partial *PartialStartError + if errors.As(err, &partial) { + // A partial start still latches started: the composite's + // healthy inner toolsets must stay listed and usable, and + // its StartReporter keeps returning false while degraded, + // so the failed subset is retried on the next Start. + s.started = true + // The latch makes every later Start a recovery run, so + // recovering alone cannot tell an inner that was started + // and lost from one that never came up (e.g. an initial + // OAuth deferral retried each turn, which must stay + // silent). Only actual post-start loss, as classified by + // the composite, marks the recovery streak. + if partial.LostAfterStart { + s.recoveryStreak.fail() + } + } else if recovering { + // A failed recovery marks the recovery streak here too, not + // only in the Restartable branch above: toolsets recovering + // through plain Start (a StartReporter without Restartable, + // or a composite whose inner toolsets all went down) need + // the targeted re-auth notice as well. + s.recoveryStreak.fail() + } return err } } diff --git a/pkg/tools/startable_test.go b/pkg/tools/startable_test.go index de4dfb3ce..cc134f474 100644 --- a/pkg/tools/startable_test.go +++ b/pkg/tools/startable_test.go @@ -3,6 +3,7 @@ package tools_test import ( "context" "errors" + "fmt" "testing" "gotest.tools/v3/assert" @@ -457,3 +458,257 @@ func TestStartableToolSet_ShouldReportRecoveryFailure_ResetsOnStop(t *testing.T) assert.Check(t, s.Start(t.Context()) != nil) assert.Check(t, is.Equal(s.ShouldReportRecoveryFailure(), true), "fresh recovery after Stop must report again") } + +// partialFlappyToolSet simulates a composite toolset (e.g. Code Mode) whose +// Start returns a scripted sequence of errors — typically +// *tools.PartialStartError values — and whose IsStarted reports true only +// after a fully-successful Start, matching the tools.StartReporter contract +// composites must implement for partial starts. +type partialFlappyToolSet struct { + errs []error + callIdx int + starts int + healthy bool +} + +func (p *partialFlappyToolSet) Tools(context.Context) ([]tools.Tool, error) { + return []tools.Tool{{Name: "partial_tool"}}, nil +} + +func (p *partialFlappyToolSet) Start(context.Context) error { + p.starts++ + if p.callIdx < len(p.errs) { + err := p.errs[p.callIdx] + p.callIdx++ + if err != nil { + p.healthy = false + return err + } + } + p.healthy = true + return nil +} + +func (p *partialFlappyToolSet) Stop(context.Context) error { + p.healthy = false + return nil +} + +func (p *partialFlappyToolSet) IsStarted() bool { return p.healthy } + +// TestStartableToolSet_PartialStartLatchesStartedAndRetries pins the +// PartialStartError contract (#3978): a partial start latches the wrapper as +// started (the healthy subset stays listed via Tools), still returns the +// error with once-per-streak reporting, keeps retrying the inner Start on +// subsequent calls while the composite reports unstarted, and a +// fully-successful retry resets the streak. +func TestStartableToolSet_PartialStartLatchesStartedAndRetries(t *testing.T) { + t.Parallel() + + partialErr := &tools.PartialStartError{Err: errors.New("mcp(ref=broken): boom")} + inner := &partialFlappyToolSet{errs: []error{partialErr, partialErr, nil}} + s := tools.NewStartable(inner) + + // Turn 1: partial failure — wrapper must be started AND the error reported. + err := s.Start(t.Context()) + assert.Check(t, err != nil, "partial start must still return the error") + assert.Check(t, is.Equal(tools.IsPartialStart(err), true)) + assert.Check(t, is.Equal(s.IsStarted(), true), "partial start must latch the wrapper as started") + assert.Check(t, is.Equal(s.ShouldReportFailure(), true), "first partial failure must be reported") + + // The healthy subset stays listed. + ta, listErr := s.Tools(t.Context()) + assert.NilError(t, listErr) + assert.Check(t, is.Len(ta, 1)) + + // Turn 2: still partial — retried, still started, no duplicate report. + assert.Check(t, s.Start(t.Context()) != nil, "expected error on turn 2") + assert.Check(t, is.Equal(inner.starts, 2), "degraded composite must be retried") + assert.Check(t, is.Equal(s.IsStarted(), true)) + assert.Check(t, is.Equal(s.ShouldReportFailure(), false), "duplicate partial failure must not report") + + // Turn 3: full success — streak resets silently, no more retries needed. + assert.NilError(t, s.Start(t.Context())) + assert.Check(t, is.Equal(inner.starts, 3)) + assert.NilError(t, s.Start(t.Context())) + assert.Check(t, is.Equal(inner.starts, 3), "fully-started composite must not be re-started") + + // A fresh partial failure after recovery starts a new streak. + inner.healthy = false + inner.errs = append(inner.errs, partialErr) + assert.Check(t, s.Start(t.Context()) != nil) + assert.Check(t, is.Equal(s.ShouldReportFailure(), true), "fresh failure after recovery must warn again") +} + +// TestStartableToolSet_NonPartialFailureStaysUnstarted pins that the partial +// latch does not weaken the plain-failure contract: a regular Start error +// leaves the wrapper unstarted. +func TestStartableToolSet_NonPartialFailureStaysUnstarted(t *testing.T) { + t.Parallel() + + f := &flappyToolSet{errs: []error{errors.New("boom")}} + s := tools.NewStartable(f) + + assert.Check(t, s.Start(t.Context()) != nil) + assert.Check(t, is.Equal(s.IsStarted(), false), "plain start failure must not latch started") +} + +// TestStartableToolSet_RecoveryFailureViaStartMarksRecoveryStreak pins that +// a failed recovery through the plain Startable branch — a toolset with a +// StartReporter but no Restartable, like a composite that dispatches +// Restart per inner toolset itself — marks the recovery streak, so the +// targeted re-auth notice (ShouldReportRecoveryFailure) fires for +// composites too, not only for Restartable toolsets. +func TestStartableToolSet_RecoveryFailureViaStartMarksRecoveryStreak(t *testing.T) { + t.Parallel() + + inner := &partialFlappyToolSet{errs: []error{nil, errors.New("session lost")}} + s := tools.NewStartable(inner) + + assert.NilError(t, s.Start(t.Context())) + assert.Check(t, is.Equal(s.ShouldReportRecoveryFailure(), false), "no recovery failure yet") + + // Background death: the toolset reports dead, so the next Start takes + // the recovery path; the inner is not Restartable, so recovery goes + // through Start, which fails. + inner.healthy = false + + assert.Check(t, s.Start(t.Context()) != nil, "expected error on recovery") + assert.Check(t, is.Equal(s.ShouldReportRecoveryFailure(), true), + "a recovery failure through the Startable branch must mark the recovery streak") + assert.Check(t, is.Equal(s.ShouldReportRecoveryFailure(), false), "second call in same streak must be false") +} + +// TestStartableToolSet_PartialStartInitialFailureNeverMarksRecoveryStreak +// pins that the started latch of a partial start does not turn retried +// initial failures into recovery failures: a composite reporting a partial +// failure without post-start loss (LostAfterStart=false, e.g. an inner +// deferring OAuth on every turn) must stay silent on the recovery channel +// even though every Start after the latch takes the recovery path. +func TestStartableToolSet_PartialStartInitialFailureNeverMarksRecoveryStreak(t *testing.T) { + t.Parallel() + + partialErr := &tools.PartialStartError{Err: errors.New("mcp(a): authorization deferred"), AuthOnly: true} + inner := &partialFlappyToolSet{errs: []error{partialErr, partialErr, partialErr}} + s := tools.NewStartable(inner) + + for turn := 1; turn <= 3; turn++ { + assert.Check(t, s.Start(t.Context()) != nil, "turn %d: expected partial failure", turn) + assert.Check(t, is.Equal(s.ShouldReportRecoveryFailure(), false), + "turn %d: a retried initial failure must not mark the recovery streak", turn) + } + assert.Check(t, is.Equal(s.IsStarted(), true), "partial start must still latch the wrapper") +} + +// TestStartableToolSet_PartialStartLostAfterStartMarksRecoveryStreak pins +// the counterpart: a partial failure that includes an inner lost after a +// successful start (LostAfterStart=true) marks the recovery streak so the +// targeted re-auth notice fires, once per streak. +func TestStartableToolSet_PartialStartLostAfterStartMarksRecoveryStreak(t *testing.T) { + t.Parallel() + + lostErr := &tools.PartialStartError{Err: errors.New("mcp(a): session lost"), LostAfterStart: true} + inner := &partialFlappyToolSet{errs: []error{nil, lostErr}} + s := tools.NewStartable(inner) + + assert.NilError(t, s.Start(t.Context())) + + // Background death: the next Start is a recovery run and reports an + // inner that was started and lost. + inner.healthy = false + + assert.Check(t, s.Start(t.Context()) != nil, "expected partial failure on recovery") + assert.Check(t, is.Equal(s.ShouldReportRecoveryFailure(), true), + "post-start loss inside a partial failure must mark the recovery streak") + assert.Check(t, is.Equal(s.ShouldReportRecoveryFailure(), false), "second call in same streak must be false") +} + +// TestPartialStartError_NilSafety pins that Error() never panics on a +// zero-value or nil PartialStartError (e.g. one constructed directly +// without NewPartialStartError). +func TestPartialStartError_NilSafety(t *testing.T) { + t.Parallel() + + assert.Check(t, is.Equal((&tools.PartialStartError{}).Error(), "partial toolset start failure")) + var nilErr *tools.PartialStartError + assert.Check(t, is.Equal(nilErr.Error(), "partial toolset start failure")) +} + +// TestTotalStartError_NilSafety mirrors TestPartialStartError_NilSafety for +// the total-failure counterpart. +func TestTotalStartError_NilSafety(t *testing.T) { + t.Parallel() + + assert.Check(t, is.Equal((&tools.TotalStartError{}).Error(), "total toolset start failure")) + var nilErr *tools.TotalStartError + assert.Check(t, is.Equal(nilErr.Error(), "total toolset start failure")) +} + +// TestIsAuthorizationRequired_PartialStartClassification pins how partial +// starts interact with the OAuth-deferral special handling: a batch is only +// authorization-required when ALL of its causes are. A mixed batch (auth + +// real failure) must not match — otherwise the non-auth failure would be +// hidden behind the silent auth-deferral path and never warned about. +func TestIsAuthorizationRequired_PartialStartClassification(t *testing.T) { + t.Parallel() + + authErr := &tools.AuthorizationRequiredError{URL: "https://example.test/mcp"} + connErr := errors.New("mcp(b): connection refused") + + authOnly := tools.NewPartialStartError(fmt.Errorf("mcp(a): %w", authErr)) + assert.Check(t, is.Equal(tools.IsPartialStart(authOnly), true)) + assert.Check(t, is.Equal(tools.IsAuthorizationRequired(authOnly), true), + "auth-only batch must keep the silent deferral handling") + + mixed := tools.NewPartialStartError(fmt.Errorf("mcp(a): %w", authErr), connErr) + assert.Check(t, is.Equal(tools.IsPartialStart(mixed), true)) + assert.Check(t, is.Equal(tools.IsAuthorizationRequired(mixed), false), + "mixed batch must not be classified auth-only") + assert.Check(t, is.Equal(tools.IsAuthorizationRequired(fmt.Errorf("starting: %w", mixed)), false), + "classification must hold through further wrapping") + + // errors.Is/As stay intact: both causes remain reachable. + assert.Check(t, errors.Is(mixed, connErr), "non-auth cause must stay reachable via errors.Is") + var target *tools.AuthorizationRequiredError + assert.Check(t, errors.As(mixed, &target), "auth cause must stay reachable via errors.As") + + nonAuth := tools.NewPartialStartError(connErr) + assert.Check(t, is.Equal(tools.IsAuthorizationRequired(nonAuth), false)) + + // Plain (non-partial) auth errors keep matching as before. + assert.Check(t, is.Equal(tools.IsAuthorizationRequired(fmt.Errorf("x: %w", authErr)), true)) +} + +// TestIsAuthorizationRequired_TotalStartClassification mirrors the partial +// classification for TotalStartError, the total-failure aggregate: the +// batch is only authorization-required when ALL of its causes are, so a +// mixed total failure cannot suppress its real causes behind the silent +// auth-deferral handling, while remaining non-partial so StartableToolSet +// never latches it. +func TestIsAuthorizationRequired_TotalStartClassification(t *testing.T) { + t.Parallel() + + authErr := &tools.AuthorizationRequiredError{URL: "https://example.test/mcp"} + connErr := errors.New("mcp(b): connection refused") + + authOnly := tools.NewTotalStartError(fmt.Errorf("mcp(a): %w", authErr)) + assert.Check(t, is.Equal(tools.IsPartialStart(authOnly), false), + "a total failure must never classify as partial") + assert.Check(t, is.Equal(tools.IsAuthorizationRequired(authOnly), true), + "all-auth total failure must keep the silent deferral handling") + + mixed := tools.NewTotalStartError(fmt.Errorf("mcp(a): %w", authErr), connErr) + assert.Check(t, is.Equal(tools.IsPartialStart(mixed), false)) + assert.Check(t, is.Equal(tools.IsAuthorizationRequired(mixed), false), + "mixed total failure must not be classified auth-only") + assert.Check(t, is.Equal(tools.IsAuthorizationRequired(fmt.Errorf("starting: %w", mixed)), false), + "classification must hold through further wrapping") + + // errors.Is/As stay intact: both causes remain reachable. + assert.Check(t, errors.Is(mixed, connErr), "non-auth cause must stay reachable via errors.Is") + var target *tools.AuthorizationRequiredError + assert.Check(t, errors.As(mixed, &target), "auth cause must stay reachable via errors.As") + + nonAuth := tools.NewTotalStartError(connErr) + assert.Check(t, is.Equal(tools.IsAuthorizationRequired(nonAuth), false)) +}