fix: pace retryable partial-start failures in code-mode composites

docs/features/code-mode/index.md

@@ -53,7 +53,7 @@ It is not a general-purpose replacement for direct tool calls: for an agent that
 - **Failures are diagnosable.** If the script throws or returns unexpectedly, the response includes the tool calls it made before failing (name, arguments, and result or error), so the model can see what happened and adjust the script on the next attempt.
 - **Not every tool is wrapped.** Tools in the `todo` category are excluded from the script environment and stay directly callable as ordinary tools — Code Mode does not replace them.
 - **The script runs in an embedded, sandboxed JS engine** ([goja](https://github.com/dop251/goja)), not Node.js or a browser: there is no filesystem, network, or process access beyond the tool functions injected into it.
-- **Partial startup is supported.** If one toolset fails to initialize (for example, an MCP server that won't connect), Code Mode remains available with the successfully loaded toolsets. The failed toolset is omitted from the JavaScript environment and retried on subsequent turns. A warning is emitted once when the failure first occurs, but not on subsequent turns while the failure persists. When the toolset recovers, its tools silently reappear in the environment.
+- **Partial startup is supported.** If one toolset fails to initialize (for example, an MCP server that won't connect), Code Mode remains available with the successfully loaded toolsets. The failed toolset is omitted from the JavaScript environment and retried on subsequent turns. A warning is emitted once when the failure first occurs, but not on subsequent turns while the failure persists. When the toolset recovers, its tools silently reappear in the environment. When the failed toolset's cause is retryable (e.g. an MCP server or RAG knowledge base hitting rate limits), its retries are paced by the same [backoff gate](../../tools/mcp/index.md#lifecycle-auto-restart-profiles) used outside code mode, instead of retrying on every turn.
 
 ## Interaction With Permissions and Tool Approval
 

docs/tools/mcp/index.md

@@ -258,7 +258,7 @@ toolsets:
 
 See [Toolset Lifecycle](../../configuration/tools/index.md#toolset-lifecycle) for all profiles and tuning knobs, and [`/toolset-restart`](../../features/tui/index.md) to force a reconnect from the TUI.
 
-**Startup failure behaviour:** local MCP failures (missing binary, connection refused, bad auth) fail fast — each turn retries immediately with no artificial delay. Remote MCP servers (Streamable HTTP / SSE) that respond with one of a fixed set of retryable HTTP statuses — 429 Too Many Requests, 408 Request Timeout, 500/502/503/504, or 529 (Anthropic-style "overloaded") — are paced by the same [bounded exponential backoff gate](../rag/index.md#indexing-failures-retries-and-backoff) that RAG embedding calls use, so a temporarily-overloaded remote MCP server does not trigger a new connect attempt on every agent turn. This is a fixed enumeration, not a full 5xx range: less-common codes such as 501, 505, or the Cloudflare 520–527 family do not arm the gate. Note MCP's trigger set is broader than RAG's current 429-only pacing (see the linked page and [#4097](https://github.com/docker/docker-agent/issues/4097)). Not yet applied when toolsets are wrapped in code mode — see [#4067](https://github.com/docker/docker-agent/issues/4067).
+**Startup failure behaviour:** local MCP failures (missing binary, connection refused, bad auth) fail fast — each turn retries immediately with no artificial delay. Remote MCP servers (Streamable HTTP / SSE) that respond with one of a fixed set of retryable HTTP statuses — 429 Too Many Requests, 408 Request Timeout, 500/502/503/504, or 529 (Anthropic-style "overloaded") — are paced by the same [bounded exponential backoff gate](../rag/index.md#indexing-failures-retries-and-backoff) that RAG embedding calls use, so a temporarily-overloaded remote MCP server does not trigger a new connect attempt on every agent turn. This is a fixed enumeration, not a full 5xx range: less-common codes such as 501, 505, or the Cloudflare 520–527 family do not arm the gate. Note MCP's trigger set is broader than RAG's current 429-only pacing (see the linked page and [#4097](https://github.com/docker/docker-agent/issues/4097)). This pacing also applies when toolsets are wrapped in [code mode](../../features/code-mode/index.md#limits--security-notes): a retryable failure in the degraded subset paces that subset's retry the same way, while the composite's healthy tools stay available.
 
 ## Combined Example
 

pkg/tools/codemode/codemode_test.go

@@ -3,6 +3,8 @@ package codemode
 import (
 	"context"
 	"encoding/json"
+	"errors"
+	"net/http"
 	"sync"
 	"testing"
 	"time"
@@ -11,6 +13,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/docker/docker-agent/pkg/modelerrors"
 	"github.com/docker/docker-agent/pkg/tools"
 )
 
@@ -971,3 +974,62 @@ func TestCodeModeTool_FailureIncludesToolArguments(t *testing.T) {
 	assert.Equal(t, map[string]any{"value": "test123"}, scriptResult.ToolCalls[0].Arguments)
 	assert.Equal(t, "result", scriptResult.ToolCalls[0].Result)
 }
+
+// TestCodeModeTool_RateLimitedInnerPacesRetry is the codemode-level
+// regression for #4067: when the composite degrades with a retryable inner
+// cause (e.g. a RAG toolset hitting 429s), the outer StartableToolSet gate
+// must pace the composite's own retry of the failed subset instead of
+// re-invoking every inner toolset's Start on every turn, while the healthy
+// subset's declarations stay listed throughout.
+func TestCodeModeTool_RateLimitedInnerPacesRetry(t *testing.T) {
+	t.Parallel()
+
+	healthy := &testToolSet{tools: []tools.Tool{{Name: "fetch_url"}}}
+	rateLimited := &testToolSet{
+		startErr: &modelerrors.StatusError{StatusCode: http.StatusTooManyRequests, Err: errors.New("rate limited")},
+		tools:    []tools.Tool{{Name: "rag_search"}},
+	}
+
+	composite := Wrap(healthy, rateLimited)
+
+	now := time.Now()
+	identityJitter := func(d time.Duration) time.Duration { return d }
+	s := tools.NewStartable(composite,
+		tools.WithStartRetryClock(func() time.Time { return now }),
+		tools.WithStartRetryJitter(identityJitter),
+	)
+
+	// Turn 1: partial start latches the wrapper (healthy subset stays
+	// listed) and arms the gate on the retryable inner cause.
+	_, err := s.TryStart(t.Context())
+	require.True(t, tools.IsPartialStart(err))
+	assert.True(t, s.IsStarted(), "healthy subset must be listed after a partial start")
+	assert.Equal(t, 1, healthy.start)
+	assert.Equal(t, 1, rateLimited.start)
+
+	allTools, terr := composite.Tools(t.Context())
+	require.NoError(t, terr)
+	require.Len(t, allTools, 1)
+	assert.Contains(t, allTools[0].Description, "FetchUrl")
+	assert.NotContains(t, allTools[0].Description, "RagSearch")
+
+	// Turn 2 (same instant): gated — neither inner is retried.
+	_, err = s.TryStart(t.Context())
+	require.Error(t, err)
+	assert.True(t, s.IsStarted(), "healthy subset must stay listed while gated")
+	assert.Equal(t, 1, healthy.start, "healthy inner must not be re-started while gated")
+	assert.Equal(t, 1, rateLimited.start, "degraded inner's retry must be paced, not re-attempted every turn")
+
+	// Advance comfortably past the (5-minute-capped) window: the next turn
+	// retries the degraded subset, which now recovers.
+	now = now.Add(6 * time.Minute)
+	rateLimited.startErr = nil
+	started, err := s.TryStart(t.Context())
+	require.NoError(t, err)
+	assert.True(t, started)
+	assert.Equal(t, 2, rateLimited.start, "degraded inner must be retried once the window elapses")
+
+	allTools, terr = composite.Tools(t.Context())
+	require.NoError(t, terr)
+	assert.Contains(t, allTools[0].Description, "RagSearch", "recovered inner's declarations must reappear")
+}

pkg/tools/startable.go

@@ -581,10 +581,13 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) {
 				// its StartReporter keeps returning false while degraded,
 				// so the failed subset is retried on the next Start.
 				s.started = true
-				// Known limitation: the failed subset's per-turn retry is
-				// not paced by the gate (code-mode composites are the main
-				// affected path). See issue #4067 for the follow-up fix.
-				s.resetStartBackoff()
+				// setStartBackoff arms the gate when the aggregated cause is
+				// retryable (errors.As walks the errors.Join tree, so any one
+				// retryable inner cause is enough) and otherwise resets it,
+				// so a degraded subset (e.g. a RAG toolset hitting 429s) is
+				// paced like any other retryable failure (#4067) without
+				// slowing down a non-retryable one.
+				s.setStartBackoff(err)
 				// 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

pkg/tools/startable_backoff_test.go

@@ -418,14 +418,17 @@ func TestStartableToolSet_PlainTextStatusShapeDoesNotArmGate(t *testing.T) {
 	}
 }
 
-// TestStartableToolSet_PartialStartClearsBackoffGate pins that a
-// PartialStartError (composite partially healthy) clears any active backoff
-// window — the toolset is (partially) up, so the cold-start gate must not
-// suppress the composite's next recovery attempt for the failed subset.
+// TestStartableToolSet_NonRetryablePartialStartClearsBackoffGate pins that a
+// PartialStartError whose aggregated cause is NOT retryable (composite
+// partially healthy) clears any active backoff window — the toolset is
+// (partially) up, and its failed subset isn't a pacing candidate, so the
+// cold-start gate must not suppress the composite's next recovery attempt.
 //
-// Scenario: window expires → partial start latches and clears gate →
-// immediate next TryStart must invoke the underlying without delay.
-func TestStartableToolSet_PartialStartClearsBackoffGate(t *testing.T) {
+// Scenario: window expires → partial start (plain error cause) latches and
+// clears gate → immediate next TryStart must invoke the underlying without
+// delay. A partial start whose cause IS retryable instead arms the gate —
+// see TestStartableToolSet_RetryablePartialStartArmsBackoffGate.
+func TestStartableToolSet_NonRetryablePartialStartClearsBackoffGate(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		inner := &partialGateClearToolSet{}
 
@@ -455,10 +458,68 @@ func TestStartableToolSet_PartialStartClearsBackoffGate(t *testing.T) {
 	})
 }
 
-// partialGateClearToolSet is a minimal Startable + StartReporter for
-// TestStartableToolSet_PartialStartClearsBackoffGate: IsStarted is true
-// only when the last Start returned nil, matching the composite-toolset
-// contract that drives the recovery path.
+// TestStartableToolSet_RetryablePartialStartArmsBackoffGate is the fix for
+// #4067: a PartialStartError whose aggregated cause IS retryable must arm
+// the gate exactly like a total failure would, even though the wrapper
+// stays latched as started (s.started=true) so the healthy subset keeps
+// listing. Without this, a degraded code-mode composite retries its failed
+// inner subset (e.g. a RAG toolset hitting 429s) unpaced on every turn.
+func TestStartableToolSet_RetryablePartialStartArmsBackoffGate(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		inner := &partialGateClearToolSet{}
+		inner.setErr(tools.NewPartialStartError(rateLimitErr()))
+		s := newThrottledStartable(inner)
+
+		// Attempt 1: underlying Start runs, partial failure latches started.
+		_, err := s.TryStart(t.Context())
+		assert.Check(t, tools.IsPartialStart(err), "expected partial start, got: %v", err)
+		assert.Check(t, is.Equal(inner.starts.Load(), int32(1)))
+		assert.Check(t, is.Equal(s.IsStarted(), true), "partial start must latch the wrapper")
+
+		// Immediate retry via TryStart: gate must block it (still within window).
+		_, err = s.TryStart(t.Context())
+		assert.Check(t, tools.IsPartialStart(err), "gate must still report the retained partial error")
+		assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), "gate must not invoke underlying Start within window")
+		assert.Check(t, is.Equal(s.IsStarted(), true), "healthy subset must stay listed while gated")
+
+		// Advance the fake clock past the base window; the next TryStart must
+		// reach the underlying again.
+		time.Sleep(tools.ExportedStartBackoffBase + time.Millisecond) //nolint:forbidigo // inside synctest bubble
+		_, err = s.TryStart(t.Context())
+		assert.Check(t, tools.IsPartialStart(err), "still degraded — expected partial start after window")
+		assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), "underlying Start must be called again after window")
+	})
+}
+
+// TestStartableToolSet_MixedAuthPartialStartArmsBackoffGate pins the
+// ANY-cause semantics of the #4067 fix: a PartialStartError joining one
+// authorization-required cause with one retryable cause must still arm the
+// gate (errors.As walks the whole errors.Join tree, so one retryable cause
+// is enough), even though the batch stays classified as NOT auth-only —
+// mirroring the ALL-causes semantics IsAuthorizationRequired already uses
+// for AuthOnly.
+func TestStartableToolSet_MixedAuthPartialStartArmsBackoffGate(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		authErr := &tools.AuthorizationRequiredError{URL: "https://example.test/mcp"}
+		inner := &partialGateClearToolSet{}
+		inner.setErr(tools.NewPartialStartError(authErr, rateLimitErr()))
+		s := newThrottledStartable(inner)
+
+		_, err := s.TryStart(t.Context())
+		assert.Check(t, tools.IsPartialStart(err))
+		assert.Check(t, !tools.IsAuthorizationRequired(err), "mixed batch must not be classified auth-only")
+		assert.Check(t, is.Equal(inner.starts.Load(), int32(1)))
+
+		_, err = s.TryStart(t.Context())
+		assert.Check(t, err != nil)
+		assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), "the retryable cause in the mix must still arm the gate")
+	})
+}
+
+// partialGateClearToolSet is a minimal Startable + StartReporter for the
+// partial-start backoff-gate tests above: IsStarted is true only when the
+// last Start returned nil, matching the composite-toolset contract that
+// drives the recovery path.
 type partialGateClearToolSet struct {
 	err    atomic.Pointer[error]
 	starts atomic.Int32