From b318089b489540c4c406dd5771f163df9306f4f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Fri, 28 Aug 2026 10:39:40 +0000 Subject: [PATCH] fix: add bounded jittered backoff to StartableToolSet retry path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes issue #4060: RAG semantic-embeddings indexing triggered a rate-limit retry storm — repeated toolset-start attempts had no pacing after a 429 from the embedding provider. Implementation: - modelerrors.RetryableHTTPStatus(err): HTTP-status classifier that recognises 429, 408, and 5xx via *StatusError first, then falls back to statusCodeRegex. The toolset gate pre-filters to *StatusError via errors.As before calling it, so port numbers and chunk counts in plain error strings cannot arm the gate. - pkg/tools/startable_backoff.go: bounded exponential backoff with additive 0-20% jitter (base=15s, cap=5min, delay∈[d,1.2d]). - Gate in tryStartLocked (TryStart/TryStartWithTimeout only): blocking Start() bypasses it so mcpcatalog enable and skill startup are immediate. - Gate adopts a live StartReporter after /toolset-restart without waiting for the window to expire. - Wrap embedding errors in WrapHTTPError at openai/client.go and dmr/embed.go so a 429 from the embedding provider surfaces as *StatusError and correctly arms the gate. - WithStartRetryJitter / WithStartRetryClock options via variadic NewStartable for deterministic test control. - Stale 'retry on next turn' log messages updated in agent.go/mcp.go. - Partial-start exemption documented (code-mode composites remain unpaced; follow-up at issue #4067). Tests (same commit, covering the above): - startable_backoff_test.go: unit tests for the gate (gate fires on 429/408/5xx StatusError, not on plain text / context errors, blocking Start() ungated, concurrency, jitter bounds) - startable_backoff_regression_test.go: consumer-shaped regression suite (RAG/MCP/LSP error shapes, no-goroutine/timer leak, latch) - rag_backoff_test.go: real-toolset integration test via rag.New + fake clock Docs: - docs/tools/rag/index.md: 'Indexing failures, retries and backoff' section with trigger table, parameters, and troubleshooting. - docs/tools/mcp/index.md, docs/tools/lsp/index.md: lifecycle notes confirming local startup failures fail fast. Scope: DefaultStartTimeout (30s) unchanged — deferred. --- docs/community/troubleshooting/index.md | 2 + docs/tools/lsp/index.md | 2 + docs/tools/mcp/index.md | 2 + docs/tools/rag/index.md | 80 +++ pkg/agent/agent.go | 4 +- pkg/model/provider/dmr/embed.go | 3 +- pkg/model/provider/openai/client.go | 2 +- pkg/modelerrors/modelerrors.go | 15 + pkg/modelerrors/modelerrors_test.go | 54 ++ pkg/tools/builtin/rag/rag_backoff_test.go | 129 ++++ pkg/tools/export_test.go | 18 +- pkg/tools/mcp/mcp.go | 2 +- pkg/tools/startable.go | 106 +++- pkg/tools/startable_backoff.go | 57 ++ .../startable_backoff_regression_test.go | 301 +++++++++ pkg/tools/startable_backoff_test.go | 597 ++++++++++++++++++ 16 files changed, 1361 insertions(+), 13 deletions(-) create mode 100644 pkg/tools/builtin/rag/rag_backoff_test.go create mode 100644 pkg/tools/startable_backoff.go create mode 100644 pkg/tools/startable_backoff_regression_test.go create mode 100644 pkg/tools/startable_backoff_test.go diff --git a/docs/community/troubleshooting/index.md b/docs/community/troubleshooting/index.md index 642e3a1453..8350b7f923 100644 --- a/docs/community/troubleshooting/index.md +++ b/docs/community/troubleshooting/index.md @@ -199,6 +199,8 @@ MCP tools using stdio transport must complete the initialization handshake befor If a toolset keeps crashing in a tight loop, tune the [`lifecycle`](../../configuration/tools/index.md#toolset-lifecycle) block on the toolset (e.g. raise `backoff.initial`, lower `max_restarts`, or switch to the `best-effort` profile) so a flaky dependency does not amplify into a restart storm. +If a **RAG knowledge base** is failing to index because the embedding provider is rate-limiting requests (HTTP 429), Docker Agent automatically backs off and retries — see [Indexing failures, retries and backoff](../../tools/rag/index.md#indexing-failures-retries-and-backoff) for the retry schedule and the `max_indexing_concurrency` / `max_embedding_concurrency` knobs that control how much concurrent load is generated. + ## Configuration Errors ### YAML syntax issues diff --git a/docs/tools/lsp/index.md b/docs/tools/lsp/index.md index b2d97b7227..916d8efb38 100644 --- a/docs/tools/lsp/index.md +++ b/docs/tools/lsp/index.md @@ -197,6 +197,8 @@ Available Capabilities: LSP toolsets are managed by the same supervisor as MCP toolsets, so a crashed `gopls` (or any other language server) is reconnected automatically with exponential backoff. Use the [`lifecycle`](../../configuration/tools/index.md#toolset-lifecycle) block to tune the policy per toolset — for example, mark `gopls` as `strict` if your CI flow requires it to be available, or use `/toolset-restart gopls` from the TUI to force a reconnect when the server gets stuck. +**Startup failure behaviour:** local process failures such as a missing binary or server-unavailable error fail fast — each turn retries immediately with no artificial delay. The backoff described in [Indexing failures, retries and backoff](../rag/index.md#indexing-failures-retries-and-backoff) applies only to structured HTTP rate-limit and server-error responses from model providers, which local LSP process errors are not. + ```yaml toolsets: - type: lsp diff --git a/docs/tools/mcp/index.md b/docs/tools/mcp/index.md index e8984c6f18..ed29f8425f 100644 --- a/docs/tools/mcp/index.md +++ b/docs/tools/mcp/index.md @@ -243,6 +243,8 @@ See [Per-Toolset Model Routing](../../configuration/tools/index.md#per-toolset-m Local stdio and remote MCP servers are supervised: crashed servers reconnect automatically with exponential backoff. **Remote** MCP servers (Streamable HTTP / SSE) also reconnect after idle/clean connection closes — services like Notion and Linear periodically close idle connections, and Docker Agent reconnects transparently. Tune the policy with the `lifecycle` block: +**Startup failure behaviour:** local process failures (missing binary, connection refused, authentication error) fail fast — each turn retries immediately with no artificial delay added. Only structured HTTP rate-limit (HTTP 429) and server error (HTTP 5xx) responses from a backing model provider trigger a backoff window. See [Indexing failures, retries and backoff](../rag/index.md#indexing-failures-retries-and-backoff) for the full policy. + ```yaml toolsets: - type: mcp diff --git a/docs/tools/rag/index.md b/docs/tools/rag/index.md index 0aa77a9935..6419a77d7a 100644 --- a/docs/tools/rag/index.md +++ b/docs/tools/rag/index.md @@ -174,6 +174,85 @@ chunking: > > Currently supports Go (`.go`) files. More languages will be added. Falls back to plain text chunking for unsupported file types. +## Indexing failures, retries and backoff + +When a knowledge-base fails to start — because the embedding provider is rate-limiting +your requests or returning a transient server error — Docker Agent spaces out retry +attempts with bounded exponential backoff instead of hammering the provider on +every agent turn. + +### What triggers backoff + +Backoff applies only to **structured HTTP rate-limit (429), request timeout (408), and server error (5xx)** +responses from the embedding or model provider. These are the signals that mean +"slow down", and spacing them out gives the provider room to recover: + +| Failure kind | Behaviour | +|---|---| +| HTTP 429 (rate limit) | Backoff: next attempt delayed | +| HTTP 5xx / overload (500, 502, 503, 504, 529) | Backoff: next attempt delayed | +| HTTP 408 (request timeout) | Backoff: next attempt delayed | +| Bad config, invalid model, wrong API key (other 4xx errors) | Fail fast: retried every turn with no added delay | +| Context cancellation or agent shutdown | Immediate: no delay | + +### Retry policy and parameters + +The backoff is **bounded exponential with additive jitter**: + +- **Base delay**: 15 seconds +- **Maximum delay**: 5 minutes +- **Growth**: doubles after each consecutive retryable failure (15s → 30s → 1m → 2m → 4m → 5m) +- **Jitter**: each wait is a random value in `[nominal, 1.2×nominal]` (additive 0–20%) + so concurrent knowledge-base sources spread their retries and avoid + hammering the provider together + +The gate is a lightweight wall-clock check — it creates no background threads or +timers. A Stop command or agent shutdown takes effect immediately regardless of +how much of the backoff window remains. + +### Operational impact + +**Before**: a rate-limited knowledge base was re-indexed on every agent turn — +`max_indexing_concurrency × max_embedding_concurrency` concurrent provider calls +could relaunch within milliseconds, easily tripping rate limits for both the +knowledge base and the agent's own model calls. + +**After**: retries are spaced out and jittered so the provider has room to recover +before the next attempt. The agent continues working with any other toolsets that +are not affected. + +### What you will see + +- Docker Agent logs a single warning when a knowledge base first fails to start. + Repeated failures in between are logged at debug level only, so you are not + flooded with alerts on every turn. Recovery is intentionally silent — the + tool appearing in the agent's tool list is the signal that indexing succeeded. +- The knowledge-base tool does not appear in the agent's tool list until indexing + succeeds. A successful start is silent — the tool is listed and the agent uses it. + +### Troubleshooting repeated 429 errors + +If you see persistent `429` errors in the logs: + +1. **Check provider rate limits.** Your embedding API key may have a low requests-per-minute + quota. Upgrading the plan or using a different API key can help. +2. **Reduce concurrency.** The chunked-embeddings and semantic-embeddings strategies + accept `max_indexing_concurrency` (default `3`) and `max_embedding_concurrency` + (default `3`) parameters. Lowering these reduces simultaneous requests: + + ```yaml + rag: + docs: + docs: [./knowledge-base] + strategies: + - type: chunked-embeddings + max_indexing_concurrency: 1 + max_embedding_concurrency: 1 + ``` + +3. **Use a model with a higher quota.** Some providers offer higher rate limits on + specific embedding model tiers. + ## Debugging RAG Enable debug logging to see retrieval details: @@ -218,6 +297,7 @@ Look for log tags: `[RAG Manager]`, `[Chunked-Embeddings Strategy]`, `[BM25 Stra | `limit` | int | `5` | Max results from this strategy | | `embedding_batch_size` | int | `50` | Chunks per embedding request | | `max_embedding_concurrency` | int | `3` | Max concurrent embedding requests | +| `max_indexing_concurrency` | int | `3` | Max concurrent file-indexing tasks | | `chunking.size` | int | `1500` | Chunk size in characters (`4000` when `code_aware` is set) | | `chunking.overlap` | int | `75` | Overlap between chunks in characters | | `chunking.code_aware` | bool | `false` | AST-based chunking (Go files only) | diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 9198ef7560..6538aa2663 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -645,10 +645,10 @@ func (a *Agent) ensureToolSetsAreStarted(ctx context.Context) { continue } if toolSet.ShouldReportFailure() { - slog.WarnContext(ctx, "Toolset start failed; will retry on next turn", "agent", a.Name(), "toolset", desc, "error", err) + slog.WarnContext(ctx, "Toolset start failed; will retry (backoff may apply)", "agent", a.Name(), "toolset", desc, "error", err) a.AddToolWarning(fmt.Sprintf("%s start failed: %v", desc, err)) } else { - slog.DebugContext(ctx, "Toolset still unavailable; retrying next turn", "agent", a.Name(), "toolset", desc, "error", err) + slog.DebugContext(ctx, "Toolset still unavailable; will retry (backoff may apply)", "agent", a.Name(), "toolset", desc, "error", err) } } } diff --git a/pkg/model/provider/dmr/embed.go b/pkg/model/provider/dmr/embed.go index aa38185d09..1de7f8fdd0 100644 --- a/pkg/model/provider/dmr/embed.go +++ b/pkg/model/provider/dmr/embed.go @@ -17,6 +17,7 @@ import ( "github.com/openai/openai-go/v3" "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/model/provider/oaistream" "github.com/docker/docker-agent/pkg/rag/types" ) @@ -50,7 +51,7 @@ func (c *Client) CreateBatchEmbedding(ctx context.Context, texts []string) (*bas Model: c.ModelConfig.Model, }) if err != nil { - return nil, fmt.Errorf("failed to create embeddings: %w", err) + return nil, fmt.Errorf("failed to create embeddings: %w", oaistream.WrapOpenAIError(err)) } if len(response.Data) != len(texts) { diff --git a/pkg/model/provider/openai/client.go b/pkg/model/provider/openai/client.go index 5b37c58cd7..ae1e81baab 100644 --- a/pkg/model/provider/openai/client.go +++ b/pkg/model/provider/openai/client.go @@ -1165,7 +1165,7 @@ func (c *Client) CreateBatchEmbedding(ctx context.Context, texts []string) (*bas response, err := client.Embeddings.New(ctx, params) if err != nil { slog.ErrorContext(ctx, "OpenAI batch embedding request failed", "error", err) - return nil, fmt.Errorf("failed to create batch embeddings: %w", err) + return nil, fmt.Errorf("failed to create batch embeddings: %w", oaistream.WrapOpenAIError(err)) } if len(response.Data) != len(texts) { diff --git a/pkg/modelerrors/modelerrors.go b/pkg/modelerrors/modelerrors.go index 428465654a..7e7c250927 100644 --- a/pkg/modelerrors/modelerrors.go +++ b/pkg/modelerrors/modelerrors.go @@ -781,3 +781,18 @@ func scalarString(v any) string { return fmt.Sprint(v) } } + +// RetryableHTTPStatus reports whether err contains an HTTP status that warrants +// backoff (429 rate-limit, 408 request-timeout, or a 5xx server error). It +// checks for a *StatusError in the chain first; if none is found it falls back +// to matching \b([45]\d{2})\b in the error message, which can produce false +// positives for port numbers or similar numeric patterns. Callers that need +// strict StatusError-only classification should pre-filter with errors.As. +// No context-error handling is performed. +func RetryableHTTPStatus(err error) bool { + code := extractHTTPStatusCode(err) + if code == 0 { + return false + } + return code == http.StatusTooManyRequests || isRetryableStatusCode(code) +} diff --git a/pkg/modelerrors/modelerrors_test.go b/pkg/modelerrors/modelerrors_test.go index b9e43bc043..61443dcb3d 100644 --- a/pkg/modelerrors/modelerrors_test.go +++ b/pkg/modelerrors/modelerrors_test.go @@ -944,3 +944,57 @@ func TestScalarStringEdgeCases(t *testing.T) { }) } } + +func TestRetryableHTTPStatus(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + expected bool + }{ + // Retryable HTTP codes via *StatusError. + {name: "429 rate-limit StatusError", err: &StatusError{StatusCode: 429, Err: errors.New("rate limited")}, expected: true}, + {name: "408 request timeout StatusError", err: &StatusError{StatusCode: 408, Err: errors.New("timeout")}, expected: true}, + {name: "500 server error StatusError", err: &StatusError{StatusCode: 500, Err: errors.New("internal server error")}, expected: true}, + {name: "503 unavailable StatusError", err: &StatusError{StatusCode: 503, Err: errors.New("service unavailable")}, expected: true}, + {name: "529 overloaded StatusError", err: &StatusError{StatusCode: 529, Err: errors.New("overloaded")}, expected: true}, + // Non-retryable HTTP codes via *StatusError. + {name: "400 bad request StatusError", err: &StatusError{StatusCode: 400, Err: errors.New("bad request")}, expected: false}, + {name: "401 unauthorized StatusError", err: &StatusError{StatusCode: 401, Err: errors.New("unauthorized")}, expected: false}, + {name: "403 forbidden StatusError", err: &StatusError{StatusCode: 403, Err: errors.New("forbidden")}, expected: false}, + {name: "404 not found StatusError", err: &StatusError{StatusCode: 404, Err: errors.New("not found")}, expected: false}, + // Plain string errors without HTTP codes: must stay non-retryable so + // MCP/LSP "connection refused" errors are not paced. + {name: "connection refused", err: errors.New("connection refused: dial tcp 127.0.0.1:9999"), expected: false}, + {name: "no such host", err: errors.New("no such host: example.invalid"), expected: false}, + // Plain text that DOES contain a retryable HTTP code: the regex fallback + // in extractHTTPStatusCode finds it. NOTE: RetryableHTTPStatus returns + // true here, but startBackoffRetryable (the toolset gate classifier) + // requires a *StatusError and returns false for the same input — the + // narrowing is deliberate to prevent port-number/chunk-count false positives. + {name: "503 in plain text", err: errors.New("upstream: 503 Service Unavailable"), expected: true}, + {name: "429 in plain text", err: errors.New("provider said: 429 Too Many Requests"), expected: true}, + // HTTP-status precedence: a StatusError{429} wrapped alongside + // context.DeadlineExceeded must return true — the HTTP signal wins. + { + name: "429 StatusError + DeadlineExceeded in chain", + err: fmt.Errorf("start budget exceeded: %w, provider said: %w", + context.DeadlineExceeded, + &StatusError{StatusCode: 429, Err: errors.New("rate limited")}), + expected: true, + }, + // Bare context errors: no HTTP code, must return false. + {name: "bare DeadlineExceeded", err: context.DeadlineExceeded, expected: false}, + {name: "bare Canceled", err: context.Canceled, expected: false}, + {name: "nil error", err: nil, expected: false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got := RetryableHTTPStatus(tc.err) + assert.Equal(t, tc.expected, got) + }) + } +} diff --git a/pkg/tools/builtin/rag/rag_backoff_test.go b/pkg/tools/builtin/rag/rag_backoff_test.go new file mode 100644 index 0000000000..4010e1eef9 --- /dev/null +++ b/pkg/tools/builtin/rag/rag_backoff_test.go @@ -0,0 +1,129 @@ +package rag + +// Regression tests proving that the real RAG ToolSet, when wrapped in +// tools.NewStartable, correctly engages the backoff gate for structured +// HTTP-status failures (e.g. a 429 rate-limit from the embedding provider) +// and does NOT throttle plain (non-StatusError) failures. +// +// Gate enforcement is in tryStartLocked (TryStart paths only); all gate +// assertions here use TryStart, not Start. A frozen fake clock supplied via +// WithStartRetryClock gives deterministic windows without needing synctest. + +import ( + "context" + "fmt" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/modelerrors" + "github.com/docker/docker-agent/pkg/rag" + "github.com/docker/docker-agent/pkg/rag/strategy" + "github.com/docker/docker-agent/pkg/tools" +) + +// countingStatusErrStrategy is a mockStrategy variant whose Initialize +// returns a *modelerrors.StatusError and counts every invocation. +type countingStatusErrStrategy struct { + mockStrategy + + statusCode int + calls atomic.Int32 +} + +func (s *countingStatusErrStrategy) Initialize(_ context.Context, _ []string, _ strategy.ChunkingConfig) error { + s.calls.Add(1) + return &modelerrors.StatusError{ + StatusCode: s.statusCode, + Err: fmt.Errorf("HTTP %d error from provider", s.statusCode), + } +} + +// countingPlainErrStrategy always returns a plain (non-StatusError) error +// so the backoff gate must never fire. +type countingPlainErrStrategy struct { + mockStrategy + + calls atomic.Int32 +} + +func (s *countingPlainErrStrategy) Initialize(_ context.Context, _ []string, _ strategy.ChunkingConfig) error { + s.calls.Add(1) + return assert.AnError +} + +func buildRAGToolSet(t *testing.T, impl strategy.Strategy) *ToolSet { + t.Helper() + cfg := rag.Config{ + StrategyConfigs: []strategy.Config{ + {Name: "test-strategy", Strategy: impl}, + }, + } + mgr, err := rag.New(t.Context(), "test-rag", cfg, nil) + require.NoError(t, err) + return &ToolSet{manager: mgr, toolName: "test-rag"} +} + +// TestRAGStartableBackoff_StatusErrorEngagesGate proves that the real RAG +// ToolSet, when it fails with a *modelerrors.StatusError, arms the backoff +// gate in TryStart: the immediate retry is gated and the underlying is not +// called again until the window expires. +// +// A frozen fake clock (WithStartRetryClock) and identity jitter give exact, +// deterministic windows without synctest complications from the RAG manager's +// internal goroutines. +func TestRAGStartableBackoff_StatusErrorEngagesGate(t *testing.T) { + t.Parallel() + + counting := &countingStatusErrStrategy{statusCode: 429} + toolset := buildRAGToolSet(t, counting) + + // Frozen clock: window = base + 0% jitter = exactly base. + // We advance the clock manually to control expiry. + now := time.Unix(1_000_000, 0) + s := tools.NewStartable(toolset, + tools.WithStartRetryJitter(func(d time.Duration) time.Duration { return d }), // identity + tools.WithStartRetryClock(func() time.Time { return now }), + ) + + // Attempt 1: TryStart invokes Initialize and arms the gate. + _, err := s.TryStart(t.Context()) + require.Error(t, err, "expected failure on attempt 1") + assert.Equal(t, int32(1), counting.calls.Load(), "Initialize must be called on attempt 1") + + // Immediate TryStart: gate must suppress it (clock not advanced). + _, err = s.TryStart(t.Context()) + require.Error(t, err) + assert.Equal(t, int32(1), counting.calls.Load(), + "gate must suppress TryStart within the window") + + // Advance clock past base window (base × 1.0 with identity jitter). + // Use 6 minutes to exceed any possible jittered window for any attempt. + now = now.Add(6 * time.Minute) + + // Gate expired: TryStart must invoke Initialize again. + _, err = s.TryStart(t.Context()) + require.Error(t, err, "still failing — expected error") + assert.Equal(t, int32(2), counting.calls.Load(), + "Initialize must be called again once the backoff window expires") +} + +// TestRAGStartableBackoff_PlainErrorNoGate proves that a plain (non-StatusError) +// failure does NOT arm the gate: every TryStart reaches Initialize. +func TestRAGStartableBackoff_PlainErrorNoGate(t *testing.T) { + t.Parallel() + + counting := &countingPlainErrStrategy{} + toolset := buildRAGToolSet(t, counting) + s := tools.NewStartable(toolset) + + for i := int32(1); i <= 3; i++ { + _, err := s.TryStart(t.Context()) + require.Error(t, err) + assert.Equal(t, i, counting.calls.Load(), + "plain error must not gate subsequent TryStart calls (attempt %d)", i) + } +} diff --git a/pkg/tools/export_test.go b/pkg/tools/export_test.go index 0a25624d48..b38f123111 100644 --- a/pkg/tools/export_test.go +++ b/pkg/tools/export_test.go @@ -1,6 +1,9 @@ package tools -import "context" +import ( + "context" + "time" +) // ExportedPublishStopRequest exposes the request-publication half of // StopIfStarted for tests in the _test package: it leaves a pending stop @@ -12,3 +15,16 @@ func (s *StartableToolSet) ExportedPublishStopRequest(ctx context.Context) { s.stopRequested = true s.stopRequestCtx = ctx } + +// ExportedComputeStartBackoff exposes computeStartBackoff for unit tests +// that verify the bounds and cap behaviour without going through a full +// Start cycle. +func ExportedComputeStartBackoff(attempt int, jitterFn func(time.Duration) time.Duration) time.Duration { + return computeStartBackoff(attempt, jitterFn) +} + +// Exported backoff bound constants for test assertions. +const ( + ExportedStartBackoffBase = startBackoffBase + ExportedStartBackoffMax = startBackoffMax +) diff --git a/pkg/tools/mcp/mcp.go b/pkg/tools/mcp/mcp.go index c5116335db..6df4ba602b 100644 --- a/pkg/tools/mcp/mcp.go +++ b/pkg/tools/mcp/mcp.go @@ -96,7 +96,7 @@ func CreateToolSet(ctx context.Context, toolset latest.Toolset, runConfig *confi case toolset.Command != "": resolvedCommand, err := toolinstall.EnsureCommand(ctx, toolset.Command, toolset.Version) if err != nil { - slog.WarnContext(ctx, "MCP command not yet available, will retry on next turn", + slog.WarnContext(ctx, "MCP command not yet available, will retry", "command", toolset.Command, "error", err) resolvedCommand = toolset.Command } diff --git a/pkg/tools/startable.go b/pkg/tools/startable.go index fd76877745..93c04a686d 100644 --- a/pkg/tools/startable.go +++ b/pkg/tools/startable.go @@ -294,11 +294,44 @@ type StartableToolSet struct { // emit a different, more targeted message (e.g. "needs re-auth" vs // "start failed") for the recovery case. recoveryStreak failureStreak + + // startBackoff throttles retryable start failures via a wall-clock gate + // enforced in tryStartLocked (#4060). All fields guarded by mu. + startBackoffUntil time.Time // zero means no active window + startBackoffAttempt int // consecutive retryable failures + startBackoffErr error // retained cause returned while throttled + startJitter func(time.Duration) time.Duration // nil uses additiveJitter; tests override via WithStartRetryJitter + now func() time.Time // nil uses time.Now; tests override via WithStartRetryClock +} + +// StartableOption is a functional option for NewStartable. +type StartableOption func(*StartableToolSet) + +// WithStartRetryJitter sets the jitter function for backoff delays; for testing. +func WithStartRetryJitter(fn func(time.Duration) time.Duration) StartableOption { + return func(s *StartableToolSet) { s.startJitter = fn } +} + +// WithStartRetryClock sets the clock used by the backoff gate; for testing. +func WithStartRetryClock(now func() time.Time) StartableOption { + return func(s *StartableToolSet) { s.now = now } } // NewStartable wraps a ToolSet for lazy initialization. -func NewStartable(ts ToolSet) *StartableToolSet { - return &StartableToolSet{ToolSet: ts} +func NewStartable(ts ToolSet, opts ...StartableOption) *StartableToolSet { + s := &StartableToolSet{ToolSet: ts} + for _, opt := range opts { + opt(s) + } + return s +} + +// nowFn returns s.now if set, otherwise time.Now. +func (s *StartableToolSet) nowFn() time.Time { + if s.now != nil { + return s.now() + } + return time.Now() } // IsStarted returns whether the toolset has been successfully started. @@ -367,7 +400,7 @@ func (s *StartableToolSet) TryStart(ctx context.Context) (started bool, err erro started = s.TryIsStarted() } }() - if err := s.startLocked(ctx); err != nil { + if err := s.tryStartLocked(ctx); err != nil { return false, err } return s.started, nil @@ -419,6 +452,51 @@ func (s *StartableToolSet) TryStartWithTimeout(ctx context.Context, timeout time } } +// resetStartBackoff clears all backoff state. s.mu must be held. +func (s *StartableToolSet) resetStartBackoff() { + s.startBackoffUntil = time.Time{} + s.startBackoffAttempt = 0 + s.startBackoffErr = nil +} + +// setStartBackoff records a retryable start failure and arms the cooldown +// gate. Non-retryable errors clear the window so the next attempt runs +// immediately (preserving today’s prompt-fail behaviour for auth/config +// errors). s.mu must be held. +func (s *StartableToolSet) setStartBackoff(err error) { + if !startBackoffRetryable(err) { + s.resetStartBackoff() + return + } + s.startBackoffAttempt++ + delay := computeStartBackoff(s.startBackoffAttempt, s.startJitter) + s.startBackoffUntil = s.nowFn().Add(delay) + s.startBackoffErr = err +} + +// tryStartLocked enforces the backoff gate and adopts any external recovery +// a live StartReporter reports. Only called from TryStart — blocking Start() +// must never be gated. s.mu must be held. +func (s *StartableToolSet) tryStartLocked(ctx context.Context) error { + // A live reporter while started==false and a window is armed means an + // external restart (e.g. /toolset-restart) cleared the failure; adopt it. + // Guard with the window check to leave non-gated TryStart semantics unchanged. + if !s.started && !s.startBackoffUntil.IsZero() { + if reporter, ok := As[StartReporter](s.ToolSet); ok && reporter.IsStarted() { + s.started = true + s.startStreak.reset() + s.recoveryStreak.reset() + s.resetStartBackoff() + return nil + } + } + // Gate: only the non-blocking TryStart paths enforce the schedule. + if !s.startBackoffUntil.IsZero() && s.nowFn().Before(s.startBackoffUntil) { + return s.startBackoffErr + } + return s.startLocked(ctx) +} + // startLocked implements the start sequence shared by Start and TryStart. // s.mu must be held. func (s *StartableToolSet) startLocked(ctx context.Context) (err error) { @@ -431,6 +509,9 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) { recovering = true } + // Gate is in tryStartLocked (TryStart only); blocking Start() always + // attempts the underlying to avoid delaying explicit starts. + // Span the toolset startup — MCP handshake, OAuth probes, // tool discovery, etc. can take seconds to minutes and the // "tools loading…" UI was previously unattributable. Only @@ -463,6 +544,7 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) { if err := restarter.Restart(ctx); err != nil { s.startStreak.fail() s.recoveryStreak.fail() + s.setStartBackoff(err) return err } } else if startable, ok := As[Startable](s.ToolSet); ok { @@ -482,12 +564,17 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) { if err := startable.Start(ctx); err != nil { s.startStreak.fail() var partial *PartialStartError - if errors.As(err, &partial) { + switch { + case 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 + // 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() // 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 @@ -497,23 +584,27 @@ func (s *StartableToolSet) startLocked(ctx context.Context) (err error) { if partial.LostAfterStart { s.recoveryStreak.fail() } - } else if recovering { + case 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() + s.setStartBackoff(err) + default: + s.setStartBackoff(err) } return err } } - // Successful start: clear the streak so any future failure is reported - // as fresh. This is the recovery path — it is intentionally silent. + // Successful start: clear streaks and backoff so any future failure is + // reported as fresh. This is the recovery path — it is intentionally silent. s.started = true s.startStreak.reset() s.recoveryStreak.reset() + s.resetStartBackoff() return nil } @@ -592,6 +683,7 @@ func (s *StartableToolSet) stopLocked(ctx context.Context) error { s.startStreak.reset() s.listStreak.reset() s.recoveryStreak.reset() + s.resetStartBackoff() if startable, ok := As[Startable](s.ToolSet); ok { return startable.Stop(ctx) } diff --git a/pkg/tools/startable_backoff.go b/pkg/tools/startable_backoff.go new file mode 100644 index 0000000000..1b99f2f089 --- /dev/null +++ b/pkg/tools/startable_backoff.go @@ -0,0 +1,57 @@ +package tools + +import ( + "errors" + "math/rand/v2" + "time" + + "github.com/docker/docker-agent/pkg/modelerrors" +) + +const ( + startBackoffBase = 15 * time.Second + startBackoffMax = 5 * time.Minute +) + +// startBackoffRetryable reports whether err carries a retryable HTTP status +// (429, 408, or 5xx). Requires a *modelerrors.StatusError anywhere in the +// chain; plain network errors and the regex fallback in RetryableHTTPStatus +// are intentionally excluded to avoid arming the gate on port numbers or +// chunk counters that match the \b[45]\d{2}\b pattern. +// A StatusError wins even when context.DeadlineExceeded is also in the chain. +func startBackoffRetryable(err error) bool { + var se *modelerrors.StatusError + if !errors.As(err, &se) { + return false + } + return modelerrors.RetryableHTTPStatus(se) +} + +// computeStartBackoff returns the exponential backoff delay for attempt +// (1-indexed), capped at startBackoffMax with additive 0–20% jitter. +// jitterFn overrides the default additiveJitter when non-nil. +func computeStartBackoff(attempt int, jitterFn func(time.Duration) time.Duration) time.Duration { + if attempt <= 0 { + attempt = 1 + } + + // Double until cap; overflow-safe because we stop at startBackoffMax. + nominal := startBackoffBase + for i := 1; i < attempt; i++ { + nominal *= 2 + if nominal >= startBackoffMax { + nominal = startBackoffMax + break + } + } + + if jitterFn != nil { + return jitterFn(nominal) + } + return additiveJitter(nominal) +} + +// additiveJitter returns d + rand.N([0, d/5]), giving delay ∈ [d, 1.2d]. +func additiveJitter(d time.Duration) time.Duration { + return d + time.Duration(rand.N(int64(d/5)+1)) +} diff --git a/pkg/tools/startable_backoff_regression_test.go b/pkg/tools/startable_backoff_regression_test.go new file mode 100644 index 0000000000..af5690e5e7 --- /dev/null +++ b/pkg/tools/startable_backoff_regression_test.go @@ -0,0 +1,301 @@ +package tools_test + +// Regression suite for the StartableToolSet backoff gate (issue #4060). +// +// Gate enforcement is in tryStartLocked, called only from TryStart/TryStartWithTimeout. +// Blocking Start() bypasses the gate and always reaches the underlying toolset — +// that is intentional (mcpcatalog enable, skill sub-session startup must be immediate). +// Gate-assertion tests therefore drive TryStart; compatibility tests use Start or +// TryStart interchangeably (the gate never fires for non-StatusError errors). +// +// The helpers startErrToolSet / retryableErr / rateLimitErr / nonRetryableErr / +// identityJitter / newThrottledStartable are defined in startable_backoff_test.go +// (same package) and are reused directly here. + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "testing/synctest" + "time" + + "gotest.tools/v3/assert" + is "gotest.tools/v3/assert/cmp" + + "github.com/docker/docker-agent/pkg/modelerrors" + "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tools/lifecycle" +) + +// ragWrappedStatusErr returns a StatusError wrapped through two levels of +// fmt.Errorf, approximating the real RAG toolset error chain +// (Initialize → Manager → ToolSet.Start). +func ragWrappedStatusErr(code int) error { + base := &modelerrors.StatusError{ + StatusCode: code, + Err: fmt.Errorf("HTTP %d from provider", code), + } + aborted := fmt.Errorf("indexing aborted due to non-retryable model error: %w", base) + return fmt.Errorf("failed to initialize RAG manager %q: %w", "knowledge-base", aborted) +} + +// TestBackoffRegression_RAGShapedFailureAndRecovery verifies the full +// backoff-then-recovery cycle for an error shaped like a real RAG toolset +// hitting a rate-limit (429) during indexing. +func TestBackoffRegression_RAGShapedFailureAndRecovery(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + inner := &startErrToolSet{} + inner.setErr(ragWrappedStatusErr(429)) + s := newThrottledStartable(inner) + + // Attempt 1: underlying Start runs via TryStart and fails. + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil, "expected failure on attempt 1") + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), "attempt 1 must invoke the underlying") + + // Immediate TryStart retry: gate must suppress it (within the base window). + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), "immediate TryStart must be gated within the window") + + // Advance fake clock past the base window (identity jitter ⇒ window == base exactly). + time.Sleep(tools.ExportedStartBackoffBase + time.Millisecond) //nolint:forbidigo // synctest bubble + + // Gate expired: TryStart must invoke the underlying again. + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil, "still failing — expected error") + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), "TryStart must re-run after the window expires") + + // Recovery: advance past the second window (base×2 with identity jitter) + // then clear the error so the next TryStart succeeds. + time.Sleep(2*tools.ExportedStartBackoffBase + time.Millisecond) //nolint:forbidigo // synctest bubble + inner.clearErr() + started, err := s.TryStart(t.Context()) + assert.NilError(t, err, "recovery TryStart must succeed") + assert.Check(t, started) + assert.Check(t, is.Equal(s.IsStarted(), true)) + + // After recovery, stop and verify the next failure starts a fresh base window. + assert.NilError(t, s.Stop(t.Context())) + inner.setErr(ragWrappedStatusErr(503)) + callsBefore := inner.starts.Load() + _, _ = s.TryStart(t.Context()) // arms fresh gate at base + // Immediate TryStart must be gated (fresh base window after reset). + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), callsBefore+1), + "after recovery, the next retryable failure must start a fresh base-delay window") + }) +} + +// TestBackoffRegression_MCPShapedCompatibility verifies that Start failures +// shaped like a real MCP toolset (plain connection error, no StatusError) +// do NOT engage the backoff gate. Every call reaches the underlying. +func TestBackoffRegression_MCPShapedCompatibility(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + inner.setErr(errors.New("connection refused: dial tcp 127.0.0.1:9999")) + s := newThrottledStartable(inner) + + // Three consecutive TryStart calls: each must invoke the underlying (no gate). + for i := int32(1); i <= 3; i++ { + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), i), + "MCP connection-error failure must not engage backoff (attempt %d)", i) + } + + inner.clearErr() + assert.NilError(t, s.Start(t.Context()), "recovery must work without a backoff window") + assert.Check(t, is.Equal(s.IsStarted(), true)) +} + +// TestBackoffRegression_LSPShapedCompatibility mirrors the MCP test for the +// LSP toolset error shape (lifecycle.ErrServerUnavailable, not a StatusError). +func TestBackoffRegression_LSPShapedCompatibility(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + inner.setErr(fmt.Errorf("binary not found: %w", lifecycle.ErrServerUnavailable)) + s := newThrottledStartable(inner) + + for i := int32(1); i <= 3; i++ { + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), i), + "LSP ErrServerUnavailable must not engage backoff (attempt %d)", i) + } + + inner.clearErr() + assert.NilError(t, s.Start(t.Context())) + assert.Check(t, is.Equal(s.IsStarted(), true)) +} + +// TestBackoffRegression_BlockingStartIsNeverGated verifies that blocking +// Start() always reaches the underlying toolset even while a TryStart backoff +// window is active. This is intentional: mcpcatalog enable and skill +// sub-session startup use Start() and must be immediate. +func TestBackoffRegression_BlockingStartIsNeverGated(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + inner.setErr(retryableErr()) + + // Freeze the clock so the window stays active for the test duration. + now := time.Unix(1_000_000, 0) + s := tools.NewStartable(inner, + tools.WithStartRetryJitter(identityJitter), + tools.WithStartRetryClock(func() time.Time { return now }), + ) + + // Arm the gate via TryStart. + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // Blocking Start() must bypass the gate and invoke the underlying. + err = s.Start(t.Context()) + assert.Check(t, err != nil, "blocking Start must fail (retryableErr still set)") + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), + "blocking Start must bypass the gate and reach the underlying") + + // TryStart within the (re-armed) window must still be gated. + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), + "TryStart within window must be gated even after a blocking Start") +} + +// TestBackoffRegression_ConcurrentStartsNoMultiplication verifies that many +// concurrent TryStart calls against a failing toolset do not multiply the +// underlying retry activity. The single-flight TryLock serialises them: the +// first goroutine runs and arms the gate; the rest either skip the lock or +// hit the gate — in both cases without invoking the underlying. +func TestBackoffRegression_ConcurrentStartsNoMultiplication(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + inner.setErr(retryableErr()) + s := newThrottledStartable(inner) // identity jitter → base window stays active + + const goroutines = 20 + var wg sync.WaitGroup + wg.Add(goroutines) + for range goroutines { + go func() { + defer wg.Done() + _, _ = s.TryStart(t.Context()) + }() + } + wg.Wait() + + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), + "concurrent TryStart calls must invoke the underlying exactly once (single-flight + gate)") +} + +// TestBackoffRegression_GateSpawnsNoTimersOrGoroutines verifies that the +// backoff gate is purely a wall-clock check with no background resources. +func TestBackoffRegression_GateSpawnsNoTimersOrGoroutines(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + inner := &startErrToolSet{} + inner.setErr(retryableErr()) + s := newThrottledStartable(inner) + + // Arm the gate via TryStart. + _, _ = s.TryStart(t.Context()) + + // Several gated TryStart retries — if any created a timer the bubble stalls. + for range 5 { + _, _ = s.TryStart(t.Context()) + } + // Bubble settling proves no leaked goroutines or timers. + }) +} + +// TestBackoffRegression_CancellationNoWindowSet verifies that context +// cancellation does not arm the backoff gate. +func TestBackoffRegression_CancellationNoWindowSet(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + inner := &startErrToolSet{} + inner.setErr(context.Canceled) + s := newThrottledStartable(inner) + + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // No window set: next TryStart must reach the underlying immediately. + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), + "context cancellation must not arm the backoff gate") + }) +} + +// TestBackoffRegression_JitterDeSynchronizesRetries verifies that additive +// jitter ([nominal, 1.2×nominal]) produces a spread of distinct durations, +// preventing concurrent toolsets from synchronising retries into a burst. +func TestBackoffRegression_JitterDeSynchronizesRetries(t *testing.T) { + t.Parallel() + + const samples = 100 + attempt := 1 + nominal := tools.ExportedStartBackoffBase // 15s + + seen := make(map[time.Duration]bool, samples) + for range samples { + d := tools.ExportedComputeStartBackoff(attempt, nil) // real (random) additive jitter + assert.Check(t, d >= nominal, + "additive jitter floor must be the nominal itself, got %s < %s", d, nominal) + assert.Check(t, d <= nominal+nominal/5, + "additive jitter ceiling must be 1.2×nominal, got %s > %s", d, nominal+nominal/5) + seen[d] = true + } + // With 100 samples over a 3s range in 1ns increments, identical values + // would be astronomically improbable — assert meaningful spread. + assert.Check(t, len(seen) > 5, + "additive jitter must produce meaningful spread; got only %d distinct values", len(seen)) +} + +// TestBackoffRegression_AlreadyStartedNotRestarted pins that a latched +// toolset is not restarted by subsequent Start or TryStart calls. +func TestBackoffRegression_AlreadyStartedNotRestarted(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + s := newThrottledStartable(inner) + + assert.NilError(t, s.Start(t.Context())) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + assert.Check(t, is.Equal(s.IsStarted(), true)) + + for range 3 { + assert.NilError(t, s.Start(t.Context())) + } + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), + "a latched toolset must not be restarted by subsequent Start calls") +} + +func http408Err() error { + return &modelerrors.StatusError{StatusCode: 408, Err: errors.New("request timeout")} +} + +// TestBackoffRegression_HTTP408EngagesGate verifies that HTTP 408 (Request +// Timeout) arms the backoff gate via TryStart. +func TestBackoffRegression_HTTP408EngagesGate(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + inner := &startErrToolSet{} + inner.setErr(http408Err()) + s := newThrottledStartable(inner) + + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil, "expected failure on attempt 1") + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // Immediate TryStart retry: gate must suppress it. + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), + "HTTP 408 must arm the backoff gate: TryStart must not reach the underlying") + }) +} diff --git a/pkg/tools/startable_backoff_test.go b/pkg/tools/startable_backoff_test.go new file mode 100644 index 0000000000..23b58f629e --- /dev/null +++ b/pkg/tools/startable_backoff_test.go @@ -0,0 +1,597 @@ +package tools_test + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "gotest.tools/v3/assert" + is "gotest.tools/v3/assert/cmp" + + "github.com/docker/docker-agent/pkg/modelerrors" + "github.com/docker/docker-agent/pkg/tools" +) + +// startErrToolSet is a minimal Startable whose error field controls the +// return value of Start. Setting err to nil makes Start succeed. starts +// tracks every invocation so tests can assert the gate is not bypassed. +type startErrToolSet struct { + err atomic.Pointer[error] + starts atomic.Int32 +} + +func (s *startErrToolSet) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } + +func (s *startErrToolSet) Start(context.Context) error { + s.starts.Add(1) + if p := s.err.Load(); p != nil { + return *p + } + return nil +} + +func (s *startErrToolSet) Stop(context.Context) error { return nil } + +func (s *startErrToolSet) setErr(err error) { s.err.Store(&err) } +func (s *startErrToolSet) clearErr() { s.err.Store(nil) } + +// retryableErr returns a *modelerrors.StatusError for status 503 (Service +// Unavailable), which RetryableHTTPStatus classifies as retryable. +func retryableErr() error { + return &modelerrors.StatusError{StatusCode: 503, Err: errors.New("service unavailable")} +} + +// rateLimitErr returns a *modelerrors.StatusError for status 429 (Too Many +// Requests), which RetryableHTTPStatus classifies as rate-limited. +func rateLimitErr() error { + return &modelerrors.StatusError{StatusCode: 429, Err: errors.New("rate limited")} +} + +// nonRetryableErr returns a *modelerrors.StatusError for status 401 +// (Unauthorized), which RetryableHTTPStatus classifies as non-retryable. +func nonRetryableErr() error { + return &modelerrors.StatusError{StatusCode: 401, Err: errors.New("unauthorized")} +} + +// identityJitter is a deterministic jitter function that returns the nominal +// delay unchanged, making backoff windows exactly predictable in tests. +// Valid under the additive-jitter contract because nominal ∈ [nominal, 1.2×nominal]. +func identityJitter(d time.Duration) time.Duration { return d } + +// newThrottledStartable returns a StartableToolSet wrapping inner with +// identity jitter so fake-clock or synctest can drive fake time precisely. +func newThrottledStartable(inner tools.ToolSet) *tools.StartableToolSet { + return tools.NewStartable(inner, tools.WithStartRetryJitter(identityJitter)) +} + +// TestStartableToolSet_RetryableFailureBacksOff verifies the core gate +// contract (via TryStart): the first attempt runs normally; within the +// backoff window the underlying Start is not invoked again; once the +// window expires the next attempt runs. +func TestStartableToolSet_RetryableFailureBacksOff(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + inner := &startErrToolSet{} + inner.setErr(retryableErr()) + s := newThrottledStartable(inner) + + // Attempt 1: underlying Start runs and fails. + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil, "expected error on first attempt") + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), "underlying Start must be called once") + + // Immediate retry via TryStart: gate must block it (still within window). + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil, "expected error within window") + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), "gate must not invoke underlying Start within window") + + // Advance fake clock past the base (15s) window (identity jitter → window == base). + time.Sleep(tools.ExportedStartBackoffBase + time.Millisecond) //nolint:forbidigo // inside synctest bubble: Sleep advances fake time + + // After the window expires the next TryStart runs the underlying attempt. + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil, "still failing — expected error after window") + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), "underlying Start must be called again after window") + }) +} + +// TestStartableToolSet_SuccessfulStartClearsBackoff verifies that a +// successful TryStart resets all backoff state so the next failure starts a +// fresh (base-delay) window rather than a longer accumulated one. +func TestStartableToolSet_SuccessfulStartClearsBackoff(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + inner := &startErrToolSet{} + inner.setErr(retryableErr()) + s := newThrottledStartable(inner) + + // Two retryable failures to advance the attempt counter. + _, _ = s.TryStart(t.Context()) + time.Sleep(tools.ExportedStartBackoffBase + time.Millisecond) //nolint:forbidigo // inside synctest bubble + _, _ = s.TryStart(t.Context()) + + // Advance past the second window (base×2 = 30s with identity jitter). + time.Sleep(2*tools.ExportedStartBackoffBase + time.Millisecond) //nolint:forbidigo // inside synctest bubble + + // Now let the toolset succeed. + inner.clearErr() + started, err := s.TryStart(t.Context()) + assert.NilError(t, err) + assert.Check(t, started, "toolset must be started after success") + assert.Check(t, is.Equal(s.IsStarted(), true)) + + // Stop so we can drive a fresh failure. + assert.NilError(t, s.Stop(t.Context())) + + // A fresh retryable failure starts a new base-delay window (attempt reset). + inner.setErr(retryableErr()) + callsBefore := inner.starts.Load() + _, _ = s.TryStart(t.Context()) + // Immediate retry must be gated (base window = 15s after reset). + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), callsBefore+1), + "after success + stop: fresh failure must start a new base-delay window") + }) +} + +// TestStartableToolSet_NoBackoffForNonRetryableFailures verifies that +// non-retryable errors (e.g. HTTP 401) do not set a backoff window — +// every TryStart reaches the underlying, preserving retry-every-turn +// behaviour for permanent config/auth errors. +func TestStartableToolSet_NoBackoffForNonRetryableFailures(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + inner.setErr(nonRetryableErr()) + s := newThrottledStartable(inner) + + for i := 1; i <= 3; i++ { + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil, "attempt %d: expected error", i) + assert.Check(t, is.Equal(inner.starts.Load(), int32(i)), + "attempt %d: non-retryable error must not gate subsequent starts", i) + } +} + +// TestStartableToolSet_NoBackoffOnContextCancellation verifies that a Start +// failure caused by context cancellation does not set a backoff window. +// Context errors are shutdown signals and must never delay future attempts. +func TestStartableToolSet_NoBackoffOnContextCancellation(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + inner.setErr(context.Canceled) + s := newThrottledStartable(inner) + + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // Immediate retry must proceed — no gate from a context error. + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), + "context cancellation must not set a backoff window") +} + +// TestStartableToolSet_RateLimitBacksOff verifies that HTTP 429 (rate +// limit) — the primary storm signal from issue #4060 — triggers the gate. +func TestStartableToolSet_RateLimitBacksOff(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + inner := &startErrToolSet{} + inner.setErr(rateLimitErr()) + s := newThrottledStartable(inner) + + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // Immediate retry must be gated. + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), + "rate-limit (429) must trigger the backoff gate in TryStart") + }) +} + +// TestStartableToolSet_StopClearsBackoff verifies that an explicit Stop +// resets the backoff window so the next TryStart runs immediately. +func TestStartableToolSet_StopClearsBackoff(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + inner := &startErrToolSet{} + inner.setErr(retryableErr()) + s := newThrottledStartable(inner) + + // Arm the backoff gate. + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // Immediate TryStart is gated. + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // Stop must reset the gate. + assert.NilError(t, s.Stop(t.Context())) + + // The next TryStart must run the underlying attempt immediately. + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), + "Stop must clear the backoff window so the next TryStart runs") + }) +} + +// TestStartableToolSet_BackoffNoDoubleStartWithinWindow verifies that many +// concurrent TryStart calls against a failing toolset do not multiply the +// underlying retry activity. The single-flight TryLock serialises them: the +// first goroutine runs the underlying Start (failing), arms the gate, then +// releases the lock. All subsequent callers either skip the lock (TryLock +// already held → return (false,nil)) or hit the gate and receive the cached +// error — in both cases without invoking the underlying Start. +func TestStartableToolSet_BackoffNoDoubleStartWithinWindow(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + inner.setErr(retryableErr()) + s := newThrottledStartable(inner) // identity jitter → base window == 15s + + const goroutines = 20 + var wg sync.WaitGroup + wg.Add(goroutines) + for range goroutines { + go func() { + defer wg.Done() + _, _ = s.TryStart(t.Context()) + }() + } + wg.Wait() + + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), + "concurrent TryStart calls must invoke the underlying Start exactly once") +} + +// TestStartableToolSet_BackoffAtBoundary pins the window-expiry contract +// using a manual fake clock: a TryStart strictly before expiry is gated, while +// one at or after expiry is allowed through. +func TestStartableToolSet_BackoffAtBoundary(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + inner.setErr(retryableErr()) + + var now time.Time + now = time.Unix(1_000_000, 0) // arbitrary fixed start + clock := func() time.Time { return now } + + s := tools.NewStartable(inner, + tools.WithStartRetryJitter(identityJitter), + tools.WithStartRetryClock(clock), + ) + + // Arm the gate: first TryStart fails, window = [now, now+base]. + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // One nanosecond before expiry: gate must hold. + now = now.Add(tools.ExportedStartBackoffBase - time.Nanosecond) + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1)), + "TryStart 1ns before window expiry must still be gated") + + // Exactly at expiry (clock == backoffUntil → not-before → gate open). + now = now.Add(time.Nanosecond) + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil, "still failing — expected error") + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), + "TryStart at expiry must be allowed through the gate") +} + +// TestStartBackoffBoundedAndJittered validates the pure mathematical +// properties of computeStartBackoff without a full Start cycle: +// - delay is always in [nominal, nominal+nominal/5] (additive 0–20% jitter) +// - delay is always <= startBackoffMax +// - no overflow or panic for large attempt counts +// - attempt=1 uses the base delay as nominal +func TestStartBackoffBoundedAndJittered(t *testing.T) { + t.Parallel() + + base := tools.ExportedStartBackoffBase + backoffMax := tools.ExportedStartBackoffMax + + // Verify bounds across a range of attempt counts including a huge one. + attempts := []int{1, 2, 3, 4, 5, 6, 7, 8, 100, 1000} + for _, attempt := range attempts { + // nominal is base * 2^(attempt-1), capped at max. + nominal := base + for i := 1; i < attempt; i++ { + nominal *= 2 + if nominal >= backoffMax { + nominal = backoffMax + break + } + } + + // Call the pure helper multiple times to sample the jitter distribution. + for range 20 { + d := tools.ExportedComputeStartBackoff(attempt, nil) + assert.Check(t, d > 0, "attempt %d: delay must be positive, got %s", attempt, d) + // Additive jitter adds 0–20% of nominal; nominal is capped at + // backoffMax, so the jittered delay is bounded by 1.2×backoffMax. + assert.Check(t, d <= backoffMax+backoffMax/5, + "attempt %d: delay %s exceeds 1.2×cap %s", attempt, d, backoffMax+backoffMax/5) + assert.Check(t, d >= nominal, + "attempt %d: delay %s is below additive-jitter floor %s", attempt, d, nominal) + assert.Check(t, d <= nominal+nominal/5, + "attempt %d: delay %s exceeds 1.2×nominal %s", attempt, d, nominal+nominal/5) + } + } + + // Identity jitter must return the nominal exactly. + d := tools.ExportedComputeStartBackoff(1, identityJitter) + assert.Check(t, is.Equal(d, base), "attempt 1 with identity jitter must equal base") + + d = tools.ExportedComputeStartBackoff(100, identityJitter) + assert.Check(t, is.Equal(d, backoffMax), "large attempt with identity jitter must equal max") +} + +// TestStartBackoffExponentialGrowth verifies that the nominal delay doubles +// with each attempt up to the cap (using identity jitter for precision). +func TestStartBackoffExponentialGrowth(t *testing.T) { + t.Parallel() + + base := tools.ExportedStartBackoffBase + backoffMax := tools.ExportedStartBackoffMax + + prev := time.Duration(0) + for attempt := 1; attempt <= 8; attempt++ { + d := tools.ExportedComputeStartBackoff(attempt, identityJitter) + assert.Check(t, d >= prev, "attempt %d: delay must not decrease (got %s < %s)", attempt, d, prev) + assert.Check(t, d <= backoffMax, "attempt %d: delay %s exceeds cap", attempt, d) + if prev > 0 && prev < backoffMax { + assert.Check(t, d == prev*2 || d == backoffMax, + "attempt %d: expected double (%s) or cap (%s), got %s", attempt, prev*2, backoffMax, d) + } + prev = d + if d == backoffMax { + break // further attempts stay capped; no need to keep checking + } + } + + // Base case. + assert.Check(t, is.Equal(tools.ExportedComputeStartBackoff(1, identityJitter), base)) +} + +// TestStartableToolSet_BackoffDormantForPlainErrors confirms that plain +// network errors (errors.New("boom") style, with no HTTP status code) never +// engage the backoff gate — the gate requires a retryable HTTP status code. +func TestStartableToolSet_BackoffDormantForPlainErrors(t *testing.T) { + t.Parallel() + + // errors.New("boom") has no HTTP status code → no gate. + inner := &startErrToolSet{} + inner.setErr(errors.New("boom")) + s := newThrottledStartable(inner) + + for i := 1; i <= 5; i++ { + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(i)), + "plain error must not engage backoff: attempt %d should have run", i) + } +} + +// TestStartableToolSet_PlainTextStatusShapeDoesNotArmGate pins that a plain +// error whose message *contains* a status-shaped number (e.g. "503" in a +// URL path or "429 of 812" in a progress log) does NOT arm the gate. +// RetryableHTTPStatus may return true for these via its regex fallback, but +// startBackoffRetryable pre-filters to *modelerrors.StatusError, so the +// regex path never reaches the gate. This guards against accidental +// loosening of the classifier. +func TestStartableToolSet_PlainTextStatusShapeDoesNotArmGate(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + err error + }{ + {"status text in message", errors.New("upstream: 503 Service Unavailable")}, + {"status in port number", errors.New("Post \"http://localhost:503/mcp\": connection refused")}, + {"status in progress counter", errors.New("chunk 429 of 812 failed")}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + inner := &startErrToolSet{} + inner.setErr(tc.err) + s := newThrottledStartable(inner) + for i := int32(1); i <= 3; i++ { + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), i), + "status-shaped plain error must not arm gate (attempt %d)", i) + } + }) + } +} + +// 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. +// +// Scenario: window expires → partial start latches and clears gate → +// immediate next TryStart must invoke the underlying without delay. +func TestStartableToolSet_PartialStartClearsBackoffGate(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + inner := &partialGateClearToolSet{} + + // Step 1: retryable failure arms the gate. + inner.setErr(retryableErr()) + s := newThrottledStartable(inner) + + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // Step 2: advance past the window, then produce a partial start. + time.Sleep(tools.ExportedStartBackoffBase + time.Millisecond) //nolint:forbidigo // inside synctest bubble + partialErr := &tools.PartialStartError{Err: errors.New("inner-b broken")} + inner.setErr(partialErr) + + _, 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(2))) + assert.Check(t, is.Equal(s.IsStarted(), true), "partial start must latch the wrapper") + + // Step 3: gate must now be cleared — a non-retryable error must + // invoke the underlying immediately (no residual backoff window). + inner.setErr(nonRetryableErr()) + _, _ = s.TryStart(t.Context()) + assert.Check(t, is.Equal(inner.starts.Load(), int32(3)), + "after partial start clears the gate, next TryStart must reach the underlying") + }) +} + +// 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. +type partialGateClearToolSet struct { + err atomic.Pointer[error] + starts atomic.Int32 + stable atomic.Bool +} + +func (p *partialGateClearToolSet) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } + +func (p *partialGateClearToolSet) Start(context.Context) error { + p.starts.Add(1) + p.stable.Store(false) // degraded by default; only fully healthy on nil error + if ptr := p.err.Load(); ptr != nil { + return *ptr // PartialStartError keeps stable=false (still degraded) + } + p.stable.Store(true) + return nil +} + +func (p *partialGateClearToolSet) Stop(context.Context) error { + p.stable.Store(false) + return nil +} + +func (p *partialGateClearToolSet) IsStarted() bool { return p.stable.Load() } +func (p *partialGateClearToolSet) setErr(err error) { p.err.Store(&err) } + +// reporterRecoveryToolSet is a minimal Startable + StartReporter whose +// Start() always increments a counter and returns the configured error. +// IsStarted() returns whatever the caller sets in the started field, so a +// test can simulate an external restart (e.g. via /toolset-restart) by +// setting started = true while the wrapper's s.started is still false. +type reporterRecoveryToolSet struct { + starts atomic.Int32 + startErr atomic.Pointer[error] + isStarted atomic.Bool +} + +func (r *reporterRecoveryToolSet) Tools(context.Context) ([]tools.Tool, error) { return nil, nil } +func (r *reporterRecoveryToolSet) Start(context.Context) error { + r.starts.Add(1) + if p := r.startErr.Load(); p != nil { + return *p + } + r.isStarted.Store(true) + return nil +} + +func (r *reporterRecoveryToolSet) Stop(context.Context) error { + r.isStarted.Store(false) + return nil +} +func (r *reporterRecoveryToolSet) IsStarted() bool { return r.isStarted.Load() } +func (r *reporterRecoveryToolSet) setStartErr(err error) { r.startErr.Store(&err) } +func (r *reporterRecoveryToolSet) clearStartErr() { r.startErr.Store(nil) } + +// TestStartableToolSet_ExternalRecoveryClearsBackoffGate verifies that a +// live StartReporter clears a stale backoff window in tryStartLocked so a +// successful /toolset-restart is immediately visible via the next TryStart, +// without waiting for the window to expire. +// +// Scenario: +// 1. TryStart fails with a retryable error (gate armed). +// 2. External restart succeeds: inner toolset's IsStarted() flips to true +// while the wrapper's started flag is still false. +// 3. TryStart sees the live reporter, adopts the state, clears the window, +// and returns (true, nil) WITHOUT invoking the underlying Start. +func TestStartableToolSet_ExternalRecoveryClearsBackoffGate(t *testing.T) { + t.Parallel() + + inner := &reporterRecoveryToolSet{} + inner.setStartErr(retryableErr()) + s := newThrottledStartable(inner) + + // Step 1: arm the gate — TryStart fails, window = base (15 s with identity jitter). + started, err := s.TryStart(t.Context()) + assert.Check(t, err != nil, "expected error arming the gate") + assert.Check(t, !started) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // Step 2: simulate an external restart (e.g. via /toolset-restart that + // called the inner Restartable directly). The inner reports IsStarted=true + // while the wrapper still has started=false and an active backoff window. + inner.clearStartErr() // no more start error + inner.isStarted.Store(true) // inner reports live + + // Step 3: TryStart must detect the live reporter, adopt the state, + // clear the gate, and return (true, nil) — without calling underlying Start. + callsBefore := inner.starts.Load() + started, err = s.TryStart(t.Context()) + assert.NilError(t, err, "TryStart must succeed after external recovery") + assert.Check(t, started, "wrapper must be latched as started after reporter-based recovery") + assert.Check(t, is.Equal(s.IsStarted(), true), "IsStarted must report started after recovery") + assert.Check(t, is.Equal(inner.starts.Load(), callsBefore), + "external recovery must not invoke underlying Start: the wrapper adopted the reporter's state") + + // Subsequent TryStart must short-circuit (latched healthy) without another Start. + started, err = s.TryStart(t.Context()) + assert.NilError(t, err) + assert.Check(t, started) + assert.Check(t, is.Equal(inner.starts.Load(), callsBefore), + "latched toolset must not be restarted on subsequent TryStart") +} + +// TestStartableToolSet_BlockingStartSkipsGate pins that blocking Start() is +// never gated — it always reaches the underlying toolset — while a concurrent +// TryStart() within the same window is correctly throttled. This invariant +// keeps mcpcatalog enable and skill sub-session startup responsive regardless +// of backoff state. +func TestStartableToolSet_BlockingStartSkipsGate(t *testing.T) { + t.Parallel() + + inner := &startErrToolSet{} + inner.setErr(retryableErr()) + + // Freeze the clock so the window stays active for the duration of the test. + now := time.Unix(1_000_000, 0) + s := tools.NewStartable(inner, + tools.WithStartRetryJitter(identityJitter), + tools.WithStartRetryClock(func() time.Time { return now }), + ) + + // Arm the gate via TryStart. + _, err := s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(1))) + + // Blocking Start must invoke the underlying even while the gate is active. + err = s.Start(t.Context()) + assert.Check(t, err != nil, "blocking Start must fail (retryableErr still set)") + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), + "blocking Start must bypass the gate and reach the underlying") + + // TryStart within the (now re-armed) window must still be gated. + _, err = s.TryStart(t.Context()) + assert.Check(t, err != nil) + assert.Check(t, is.Equal(inner.starts.Load(), int32(2)), + "TryStart within window must be gated: underlying must not be invoked again") +}