Function TestRun_StreamRetry_RecordsMetric

coderd/x/chatd/chatloop/metrics_test.go:513–578

TestRun_StreamRetry_RecordsMetric exercises the end-to-end retry path: a retryable error on the first Stream call, success on the second. Asserts both the metric and the back-compat OnRetry callback fire. Note: chatretry.Retry uses time.NewTimer (not quartz.Clock), so this test pays chatretry.Initi

(t *testing.T)

// this test pays chatretry.InitialDelay (1s) of real wall-clock
// time per retry. Keep it to one retry.
func TestRun_StreamRetry_RecordsMetric(t *testing.T) {
	t.Parallel()

	reg := prometheus.NewRegistry()
	metrics := chatloop.NewMetrics(reg)

	type retryCall struct {
		attempt    int
		classified chatretry.ClassifiedError
	}
	var retries []retryCall

	calls := 0
	model := &chattest.FakeModel{
		ProviderName: "test-provider",
		ModelName:    "test-model",
		StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) {
			calls++
			if calls == 1 {
				return nil, xerrors.New("received status 429 from upstream")
			}
			return func(yield func(fantasy.StreamPart) bool) {
				yield(fantasy.StreamPart{
					Type:         fantasy.StreamPartTypeFinish,
					FinishReason: fantasy.FinishReasonStop,
				})
			}, nil
		},
	}

	err := chatloop.Run(context.Background(), chatloop.RunOptions{
		Model:                model,
		MaxSteps:             1,
		ContextLimitFallback: 4096,
		PersistStep: func(_ context.Context, _ chatloop.PersistedStep) error {
			return nil
		},
		Metrics: metrics,
		OnRetry: func(
			attempt int,
			_ error,
			classified chatretry.ClassifiedError,
			_ time.Duration,
		) {
			retries = append(retries, retryCall{
				attempt:    attempt,
				classified: classified,
			})
		},
	})
	require.NoError(t, err)

	// Back-compat: OnRetry still fires with classified error.
	require.Len(t, retries, 1)
	assert.Equal(t, 1, retries[0].attempt)
	assert.Equal(t, codersdk.ChatErrorKindRateLimit, retries[0].classified.Kind)
	assert.Equal(t, "test-provider", retries[0].classified.Provider)