diff --git a/coderd/x/chatd/chaterror/classify.go b/coderd/x/chatd/chaterror/classify.go index 4bf28efd4f..44527822ff 100644 --- a/coderd/x/chatd/chaterror/classify.go +++ b/coderd/x/chatd/chaterror/classify.go @@ -7,10 +7,15 @@ import ( "time" "golang.org/x/net/http2" + "golang.org/x/xerrors" "github.com/coder/coder/v2/codersdk" ) +// ErrProviderTransportReset identifies provider stream cancellations that +// occur while the caller-owned chat context is still alive. +var ErrProviderTransportReset = xerrors.New("provider transport reset") + // ClassifiedError is the normalized, user-facing view of an // underlying provider or runtime error. type ClassifiedError struct { @@ -147,9 +152,10 @@ func Classify(err error) ClassifiedError { statusCode = extractStatusCode(lower) } provider := detectProvider(lower) - canceled := errors.Is(err, context.Canceled) || strings.Contains(lower, "context canceled") + canceled := errors.Is(err, context.Canceled) + providerTransportReset := errors.Is(err, ErrProviderTransportReset) interrupted := containsAny(lower, interruptedPatterns...) - if canceled || interrupted { + if interrupted { return normalizeClassification(ClassifiedError{ Message: "The request was canceled before it completed.", Detail: structured.detail, @@ -209,9 +215,11 @@ func Classify(err error) ClassifiedError { // over broader string fallbacks so protocol bugs do not retry. timeoutPatternMatch = false } - timeoutMatch := deadline || statusCode == 408 || statusCode == 502 || - statusCode == 503 || statusCode == 504 || - retryableHTTP2StreamReset || timeoutPatternMatch + providerTransportResetMatch := providerTransportReset && statusCode == 0 + timeoutMatch := providerTransportResetMatch || deadline || + statusCode == 408 || statusCode == 502 || statusCode == 503 || + statusCode == 504 || retryableHTTP2StreamReset || + timeoutPatternMatch genericRetryableMatch := statusCode == 500 || containsAny(lower, genericRetryablePatterns...) // Config signals should beat ambiguous wrapper signals so @@ -289,6 +297,17 @@ func Classify(err error) ClassifiedError { }) } + if canceled { + return normalizeClassification(ClassifiedError{ + Message: "The request was canceled before it completed.", + Detail: structured.detail, + Kind: codersdk.ChatErrorKindGeneric, + Provider: provider, + StatusCode: statusCode, + RetryAfter: structured.retryAfter, + }) + } + return normalizeClassification(ClassifiedError{ Detail: structured.detail, Kind: codersdk.ChatErrorKindGeneric, diff --git a/coderd/x/chatd/chaterror/classify_test.go b/coderd/x/chatd/chaterror/classify_test.go index 0e2e008bb8..a35ee865b9 100644 --- a/coderd/x/chatd/chaterror/classify_test.go +++ b/coderd/x/chatd/chaterror/classify_test.go @@ -2,6 +2,7 @@ package chaterror_test import ( "context" + "errors" "fmt" "io" "net/http" @@ -219,6 +220,57 @@ func TestClassify(t *testing.T) { StatusCode: 0, }, }, + { + name: "ProviderTransportResetIsRetryable", + err: errors.Join(chaterror.ErrProviderTransportReset, context.Canceled), + want: chaterror.ClassifiedError{ + Message: "The AI provider is temporarily unavailable.", + Kind: codersdk.ChatErrorKindTimeout, + Provider: "", + Retryable: true, + StatusCode: 0, + }, + }, + { + name: "BareContextCanceledStaysNonRetryable", + err: context.Canceled, + want: chaterror.ClassifiedError{ + Message: "The request was canceled before it completed.", + Kind: codersdk.ChatErrorKindGeneric, + Provider: "", + Retryable: false, + StatusCode: 0, + }, + }, + { + name: "Status500ContextCanceledClassifiesAsRetryable", + err: xerrors.Errorf("received status 500 from upstream: %w", context.Canceled), + want: chaterror.ClassifiedError{ + Message: "The AI provider returned an unexpected error.", + Kind: codersdk.ChatErrorKindGeneric, + Provider: "", + Retryable: true, + StatusCode: http.StatusInternalServerError, + }, + }, + { + name: "ProviderStatus500ContextCanceledClassifiesAsRetryable", + err: xerrors.Errorf("provider stream closed: %w", errors.Join( + context.Canceled, + &fantasy.ProviderError{ + Message: "context canceled", + StatusCode: http.StatusInternalServerError, + }, + )), + want: chaterror.ClassifiedError{ + Message: "The AI provider returned an unexpected error.", + Detail: "context canceled", + Kind: codersdk.ChatErrorKindGeneric, + Provider: "", + Retryable: true, + StatusCode: http.StatusInternalServerError, + }, + }, // The next cases model the error that fantasy produces // when aibridge's disabledProviderHandler returns a 503 // plain-text sentinel. Fantasy sets Title from the HTTP diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 64b1d8f97c..2ba435d9ca 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -795,6 +795,74 @@ func TestRun_HTTP2TransportErrorClassifiedAsRetryableTimeout(t *testing.T) { } } +func TestRun_RetriesProviderContextCanceledStreamError(t *testing.T) { + t.Parallel() + + attempts := 0 + retryErrs := make(chan error, chatretry.MaxAttempts) + retries := make(chan chatretry.ClassifiedError, chatretry.MaxAttempts) + var persisted []fantasy.Content + ctx := testutil.Context(t, testutil.WaitShort) + model := &chattest.FakeModel{ + ProviderName: "openai", + StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + attempts++ + if attempts == 1 { + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "partial"}, + {Type: fantasy.StreamPartTypeError, Error: context.Canceled}, + }), nil + } + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-2"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-2", Delta: "done"}, + {Type: fantasy.StreamPartTypeTextEnd, ID: "text-2"}, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, + }), nil + }, + } + + err := Run(ctx, RunOptions{ + Model: model, + MaxSteps: 1, + ContextLimitFallback: 4096, + PersistStep: func(_ context.Context, step PersistedStep) error { + persisted = append([]fantasy.Content(nil), step.Content...) + return nil + }, + OnRetry: func( + _ int, + retryErr error, + classified chatretry.ClassifiedError, + _ time.Duration, + ) { + retryErrs <- retryErr + retries <- classified + }, + }) + require.NoError(t, err) + require.Equal(t, 2, attempts) + require.Len(t, retryErrs, 1) + require.Len(t, retries, 1) + retryErr := testutil.RequireReceive(ctx, t, retryErrs) + classified := testutil.RequireReceive(ctx, t, retries) + require.ErrorIs(t, retryErr, chaterror.ErrProviderTransportReset) + require.ErrorIs(t, retryErr, context.Canceled) + require.Equal(t, codersdk.ChatErrorKindTimeout, classified.Kind) + require.True(t, classified.Retryable) + require.Equal(t, "openai", classified.Provider) + require.Equal(t, "OpenAI is temporarily unavailable.", classified.Message) + + text := requireTextContent(t, persisted, "done") + require.Equal(t, "done", text.Text) + for _, block := range persisted { + if text, ok := fantasy.AsContentType[fantasy.TextContent](block); ok { + require.NotContains(t, text.Text, "partial") + } + } +} + func TestRun_RetriesSilenceTimeoutBeforeFirstPart(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatloop/metrics_test.go b/coderd/x/chatd/chatloop/metrics_test.go index c0c86deacc..adc23ef291 100644 --- a/coderd/x/chatd/chatloop/metrics_test.go +++ b/coderd/x/chatd/chatloop/metrics_test.go @@ -577,24 +577,30 @@ func TestRun_StreamRetry_RecordsMetric(t *testing.T) { }) } -// TestRun_StreamRetry_CanceledDoesNotIncrement pins the invariant -// that canceled streams never increment stream_retries_total. -// chaterror.Classify routes context.Canceled to -// ClassifiedError{Retryable: false}, so chatretry.Retry returns -// immediately without calling onRetry. This test guards against -// future classification changes that could silently introduce -// misleading retry samples. -func TestRun_StreamRetry_CanceledDoesNotIncrement(t *testing.T) { +// TestRun_StreamRetry_ContextCanceledTransportResetIncrements pins the +// invariant that provider-originated context cancellation is counted as +// a retryable transport reset when the chat context is still alive. +func TestRun_StreamRetry_ContextCanceledTransportResetIncrements(t *testing.T) { t.Parallel() reg := prometheus.NewRegistry() metrics := chatloop.NewMetrics(reg) + attempts := 0 model := &chattest.FakeModel{ ProviderName: "test-provider", ModelName: "test-model", StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return nil, context.Canceled + attempts++ + if attempts == 1 { + return nil, context.Canceled + } + return func(yield func(fantasy.StreamPart) bool) { + _ = yield(fantasy.StreamPart{ + Type: fantasy.StreamPartTypeFinish, + FinishReason: fantasy.FinishReasonStop, + }) + }, nil }, } @@ -607,19 +613,15 @@ func TestRun_StreamRetry_CanceledDoesNotIncrement(t *testing.T) { }, Metrics: metrics, }) - // Expect an error (the stream failed); we don't care which error - // kind as long as no retry was recorded. - require.Error(t, err) - - families, err := reg.Gather() require.NoError(t, err) + require.Equal(t, 2, attempts) - for _, f := range families { - if f.GetName() == "coderd_chatd_stream_retries_total" { - assert.Empty(t, f.GetMetric(), - "stream_retries_total should have no samples after a canceled stream") - } - } + requireCounter(t, reg, "coderd_chatd_stream_retries_total", 1, map[string]string{ + "provider": "test-provider", + "model": "test-model", + "kind": string(codersdk.ChatErrorKindTimeout), + "chain_broken": "false", + }) } func TestRun_ToolError_RecordsMetric(t *testing.T) { diff --git a/coderd/x/chatd/chatretry/chatretry.go b/coderd/x/chatd/chatretry/chatretry.go index 10e2d7e806..c7833369a7 100644 --- a/coderd/x/chatd/chatretry/chatretry.go +++ b/coderd/x/chatd/chatretry/chatretry.go @@ -5,6 +5,7 @@ package chatretry import ( "context" + "errors" "time" "golang.org/x/xerrors" @@ -30,8 +31,8 @@ const ( type ClassifiedError = chaterror.ClassifiedError -// IsRetryable determines whether an error from an LLM provider is -// transient and worth retrying. +// IsRetryable reports whether err is retryable. Unlike Retry, it does not +// reclassify bare context.Canceled as a transport reset. func IsRetryable(err error) bool { return chaterror.Classify(err).Retryable } @@ -60,6 +61,29 @@ func effectiveDelay(attempt int, classified ClassifiedError) time.Duration { return delay } +func contextError(ctx context.Context) error { + if cause := context.Cause(ctx); cause != nil { + return cause + } + return ctx.Err() +} + +// classifyProviderAttemptError must be called after the caller's context +// has been checked. Provider clients can surface remote stream resets as +// bare context.Canceled, which this converts into a retryable transport reset. +func classifyProviderAttemptError(err error) (ClassifiedError, error) { + classified := chaterror.Classify(err) + if classified.Retryable || classified.StatusCode != 0 || !errors.Is(err, context.Canceled) { + return classified, err + } + wrapped := errors.Join(chaterror.ErrProviderTransportReset, err) + reclassified := chaterror.Classify(wrapped) + if !reclassified.Retryable { + return classified, err + } + return reclassified, wrapped +} + // RetryFn is the function to retry. It receives a context and returns // an error. The context may be a child of the original with adjusted // deadlines for individual attempts. @@ -75,26 +99,33 @@ type OnRetryFn func(attempt int, err error, classified ClassifiedError, delay ti // Retries use exponential backoff capped at MaxDelay, unless the // normalized error includes a longer provider Retry-After hint. // +// When fn returns bare context.Canceled while ctx is still alive, Retry +// treats it as a provider transport reset and retries it. +// // The onRetry callback (if non-nil) is called before each retry // attempt, giving the caller a chance to reset state, log, or // publish status events. func Retry(ctx context.Context, fn RetryFn, onRetry OnRetryFn) error { var attempt int for { + if ctxErr := contextError(ctx); ctxErr != nil { + return ctxErr + } + err := fn(ctx) if err == nil { return nil } - classified := chaterror.Classify(err) - if !classified.Retryable { - return chaterror.WithClassification(err, classified) + // fn runs with ctx. If it canceled the caller's context, that cause + // wins over the provider error returned from fn. + if ctxErr := contextError(ctx); ctxErr != nil { + return ctxErr } - // If the caller's context is already done, return the - // context error so cancellation propagates cleanly. - if ctx.Err() != nil { - return ctx.Err() + classified, err := classifyProviderAttemptError(err) + if !classified.Retryable { + return chaterror.WithClassification(err, classified) } attempt++ @@ -115,7 +146,7 @@ func Retry(ctx context.Context, fn RetryFn, onRetry OnRetryFn) error { select { case <-ctx.Done(): timer.Stop() - return ctx.Err() + return contextError(ctx) case <-timer.C: } } diff --git a/coderd/x/chatd/chatretry/chatretry_test.go b/coderd/x/chatd/chatretry/chatretry_test.go index d17774d2f4..61fdb047bb 100644 --- a/coderd/x/chatd/chatretry/chatretry_test.go +++ b/coderd/x/chatd/chatretry/chatretry_test.go @@ -15,6 +15,7 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chaterror" "github.com/coder/coder/v2/coderd/x/chatd/chatretry" + "github.com/coder/coder/v2/codersdk" ) func TestIsRetryableDelegatesToClassification(t *testing.T) { @@ -162,6 +163,130 @@ func TestRetry_MultipleTransientThenSuccess(t *testing.T) { require.Equal(t, 4, calls) } +func TestRetry_ContextCanceledStatus500ThenSuccess(t *testing.T) { + t.Parallel() + + calls := 0 + err := chatretry.Retry(context.Background(), func(_ context.Context) error { + calls++ + if calls == 1 { + return xerrors.Errorf("received status 500 from upstream: %w", context.Canceled) + } + return nil + }, nil) + require.NoError(t, err) + require.Equal(t, 2, calls) +} + +func TestRetry_ContextCanceledNonRetryableDoesNotWrapAsTransportReset(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + wantKind codersdk.ChatErrorKind + wantStatus int + }{ + { + name: "Status401", + err: xerrors.Errorf("received status 401 from upstream: %w", context.Canceled), + wantKind: codersdk.ChatErrorKindAuth, + wantStatus: 401, + }, + { + name: "QuotaNoStatus", + err: xerrors.Errorf("insufficient_quota: %w", context.Canceled), + wantKind: codersdk.ChatErrorKindUsageLimit, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + calls := 0 + err := chatretry.Retry(context.Background(), func(_ context.Context) error { + calls++ + return tt.err + }, nil) + require.Error(t, err) + require.ErrorIs(t, err, context.Canceled) + require.NotErrorIs(t, err, chaterror.ErrProviderTransportReset) + require.Equal(t, 1, calls) + classified := chaterror.Classify(err) + require.Equal(t, tt.wantKind, classified.Kind) + require.False(t, classified.Retryable) + require.Equal(t, tt.wantStatus, classified.StatusCode) + }) + } +} + +func TestRetry_ContextCanceledFromAttemptWithHealthyParentRetries(t *testing.T) { + t.Parallel() + + calls := 0 + var retryErr error + var retryClassified chatretry.ClassifiedError + err := chatretry.Retry(context.Background(), func(_ context.Context) error { + calls++ + if calls == 1 { + return context.Canceled + } + return nil + }, func( + _ int, + err error, + classified chatretry.ClassifiedError, + _ time.Duration, + ) { + retryErr = err + retryClassified = classified + }) + require.NoError(t, err) + require.Equal(t, 2, calls) + require.ErrorIs(t, retryErr, chaterror.ErrProviderTransportReset) + require.ErrorIs(t, retryErr, context.Canceled) + require.Equal(t, chaterror.ClassifiedError{ + Message: "The AI provider is temporarily unavailable.", + Kind: codersdk.ChatErrorKindTimeout, + Retryable: true, + StatusCode: 0, + }, retryClassified) +} + +func TestRetry_ContextCanceledFromParentDoesNotRetry(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + + calls := 0 + err := chatretry.Retry(ctx, func(_ context.Context) error { + calls++ + cancel() + return context.Canceled + }, nil) + require.ErrorIs(t, err, context.Canceled) + require.NotErrorIs(t, err, chaterror.ErrProviderTransportReset) + require.Equal(t, 1, calls) +} + +func TestRetry_ParentCancelCauseIsPreserved(t *testing.T) { + t.Parallel() + + cause := xerrors.New("retry parent stopped") + ctx, cancel := context.WithCancelCause(context.Background()) + + calls := 0 + err := chatretry.Retry(ctx, func(_ context.Context) error { + calls++ + cancel(cause) + return context.Canceled + }, nil) + require.ErrorIs(t, err, cause) + require.NotErrorIs(t, err, chaterror.ErrProviderTransportReset) + require.Equal(t, 1, calls) +} + func TestRetry_NonRetryableError(t *testing.T) { t.Parallel()