From bf4b80a1e1d33c91d3967042d5c47a16ce7517c8 Mon Sep 17 00:00:00 2001 From: Scott Miller Date: Fri, 4 Sep 2026 03:41:19 +0000 Subject: [PATCH] feat(coderd/x/chatd/chatloop): emit stream and time_to_first_token stages GenerateAssistant takes a StageTracer and StageModel and opens a stream stage around the guarded model stream. guardedStream opens a time_to_first_token stage that finishTTFT closes exactly once, on the first streamed part or when the attempt is released without one; only windows a part closed feed the TTFT and stage histograms. The attempt is released before the stream stage ends so the TTFT window always falls inside the stream that contains it. --- coderd/x/chatd/chatloop/chatloop.go | 61 +++++++- .../x/chatd/chatloop/stage_internal_test.go | 139 ++++++++++++++++++ 2 files changed, 193 insertions(+), 7 deletions(-) create mode 100644 coderd/x/chatd/chatloop/stage_internal_test.go diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index f22cbc8d64b78..8e3fa8499a25d 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -18,6 +18,7 @@ import ( fantasyanthropic "charm.land/fantasy/providers/anthropic" "charm.land/fantasy/schema" "github.com/google/uuid" + "go.opentelemetry.io/otel/attribute" "golang.org/x/xerrors" "cdr.dev/slog/v3" @@ -234,6 +235,8 @@ type GenerateAssistantOptions struct { OnModelStreamStart func() Logger slog.Logger Metrics *Metrics + Stages *StageTracer + StageModel StageModel } // AssistantOutcome is the durable assistant-side result from one model call. @@ -437,8 +440,12 @@ func GenerateAssistant(ctx context.Context, opts GenerateAssistantOptions) (Assi opts.OnModelStreamStart() } stepCtx := chatdebug.ReuseStep(ctx) + streamCtx, streamSpan := opts.Stages.Start(stepCtx, StageStream, + attribute.String(AttrProvider, provider), + ) + streamSpan.SetModel(opts.StageModel) attempt, streamErr := guardedStream( - stepCtx, + streamCtx, provider, modelName, opts.Clock, @@ -447,8 +454,11 @@ func GenerateAssistant(ctx context.Context, opts GenerateAssistantOptions) (Assi return opts.Model.Stream(attemptCtx, call) }, opts.Metrics, + opts.Stages, + opts.StageModel, ) if streamErr != nil { + streamSpan.End(streamErr) wrappedErr := wrapProviderStreamError(errorProvider, streamErr) classified := chaterror.Classify(wrappedErr).WithProvider(errorProvider) if classified.Retryable { @@ -457,9 +467,17 @@ func GenerateAssistant(ctx context.Context, opts GenerateAssistantOptions) (Assi return AssistantOutcome{}, wrappedErr } defer attempt.release() + // Releasing the attempt closes the time_to_first_token window, so it + // must happen before the stream stage ends: a window still open then + // would be counted outside the stream that contains it. + endStream := func(err error) { + attempt.release() + streamSpan.End(err) + } result, processErr := processStepStream(attempt.ctx, attempt.stream, opts.Clock, publishMessagePart) if err := attempt.finish(processErr); err != nil { + endStream(err) if errors.Is(err, ErrInterrupted) { return AssistantOutcome{}, ErrInterrupted } @@ -472,6 +490,7 @@ func GenerateAssistant(ctx context.Context, opts GenerateAssistantOptions) (Assi } contextLimit := extractContextLimitWithFallback(result.providerMetadata, opts.ContextLimitFallback) + endStream(nil) result.content = chatsanitize.SanitizeAnthropicProviderToolStepContent( ctx, opts.Logger, provider, modelName, "assistant_helper", 0, result.finishReason, result.content, @@ -874,6 +893,10 @@ func classifyStreamSilenceTimeout( }) } +// errNoFirstToken marks a time_to_first_token window that ended when +// the attempt was released or failed before any part streamed. +var errNoFirstToken = xerrors.New("stream ended before the first token") + func guardedStream( parent context.Context, provider, model string, @@ -881,30 +904,54 @@ func guardedStream( timeout time.Duration, openStream func(context.Context) (fantasy.StreamResponse, error), metrics *Metrics, + stages *StageTracer, + stageModel StageModel, ) (guardedAttempt, error) { attemptCtx, cancelAttempt := context.WithCancelCause(parent) guard := newStreamSilenceGuard(clock, timeout, cancelAttempt) + streamStart := clock.Now() + _, ttftSpan := stages.Start(parent, StageTimeToFirstToken, + attribute.String(AttrProvider, provider), + ) + ttftSpan.SetModel(stageModel) + var ttftOnce sync.Once + // finishTTFT closes the time_to_first_token window exactly once, + // either on the first streamed part or when the attempt is released + // without one. The TTFT histogram and the stage histogram only count + // windows that a part actually closed; a window cut short by a + // failure measures the failure, not the model's latency. + finishTTFT := func(err error) { + ttftOnce.Do(func() { + if err == nil { + metrics.TTFTSeconds.WithLabelValues(provider, model).Observe( + clock.Since(streamStart).Seconds(), + ) + ttftSpan.End(nil) + return + } + ttftSpan.EndWithoutObservation(err) + }) + } var releaseOnce sync.Once release := func() { releaseOnce.Do(func() { guard.Disarm() cancelAttempt(nil) + finishTTFT(errNoFirstToken) }) } - streamStart := clock.Now() stream, err := openStream(attemptCtx) if err != nil { err = classifyStreamSilenceTimeout(attemptCtx, provider, err) + finishTTFT(err) release() return guardedAttempt{}, err } - recordTTFT := sync.OnceFunc(func() { - metrics.TTFTSeconds.WithLabelValues(provider, model).Observe( - clock.Since(streamStart).Seconds(), - ) - }) + recordTTFT := func() { + finishTTFT(nil) + } return guardedAttempt{ ctx: attemptCtx, stream: fantasy.StreamResponse(func(yield func(fantasy.StreamPart) bool) { diff --git a/coderd/x/chatd/chatloop/stage_internal_test.go b/coderd/x/chatd/chatloop/stage_internal_test.go new file mode 100644 index 0000000000000..7aaf7d73641d6 --- /dev/null +++ b/coderd/x/chatd/chatloop/stage_internal_test.go @@ -0,0 +1,139 @@ +package chatloop + +import ( + "context" + "testing" + "time" + + "charm.land/fantasy" + "github.com/prometheus/client_golang/prometheus" + dto "github.com/prometheus/client_model/go" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/codes" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "golang.org/x/xerrors" + + "github.com/coder/quartz" +) + +// ttftStageCount returns the observation count of the +// time_to_first_token series on stage_duration_seconds, and whether the +// series exists at all. +func ttftStageCount(t *testing.T, registry *prometheus.Registry) (uint64, bool) { + t.Helper() + families, err := registry.Gather() + require.NoError(t, err) + for _, family := range families { + if family.GetName() != "coderd_chatd_stage_duration_seconds" { + continue + } + for _, metric := range family.GetMetric() { + if stageLabel(metric) == StageTimeToFirstToken { + return metric.GetHistogram().GetSampleCount(), true + } + } + } + return 0, false +} + +func stageLabel(metric *dto.Metric) string { + for _, label := range metric.GetLabel() { + if label.GetName() == "stage" { + return label.GetValue() + } + } + return "" +} + +func ttftSpanStatus(t *testing.T, spans *tracetest.SpanRecorder) sdktrace.ReadOnlySpan { + t.Helper() + for _, span := range spans.Ended() { + if span.Name() == StageTimeToFirstToken { + return span + } + } + t.Fatalf("no %s span was recorded", StageTimeToFirstToken) + return nil +} + +func TestGuardedStreamTTFTStage(t *testing.T) { + t.Parallel() + + newFixture := func(t *testing.T) (*StageTracer, *tracetest.SpanRecorder, *prometheus.Registry, *Metrics) { + t.Helper() + recorder := tracetest.NewSpanRecorder() + provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder)) + t.Cleanup(func() { + require.NoError(t, provider.Shutdown(context.Background())) + }) + registry := prometheus.NewRegistry() + metrics := NewMetrics(registry) + return NewStageTracer(provider, metrics), recorder, registry, metrics + } + + t.Run("FirstPartObservesStage", func(t *testing.T) { + t.Parallel() + stages, spans, registry, metrics := newFixture(t) + + attempt, err := guardedStream( + t.Context(), "anthropic", "claude", quartz.NewMock(t), time.Minute, + func(context.Context) (fantasy.StreamResponse, error) { + return fantasy.StreamResponse(func(yield func(fantasy.StreamPart) bool) { + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, Delta: "hi"}) + }), nil + }, + metrics, stages, StageModel{Model: "claude", Effort: "high"}, + ) + require.NoError(t, err) + parts := 0 + for range attempt.stream { + parts++ + } + require.Equal(t, 1, parts) + attempt.release() + + count, ok := ttftStageCount(t, registry) + require.True(t, ok) + require.Equal(t, uint64(1), count) + require.Equal(t, codes.Unset, ttftSpanStatus(t, spans).Status().Code) + }) + + t.Run("OpenFailureRecordsSpanWithoutObservation", func(t *testing.T) { + t.Parallel() + stages, spans, registry, metrics := newFixture(t) + + openErr := xerrors.New("provider refused the request") + _, err := guardedStream( + t.Context(), "anthropic", "claude", quartz.NewMock(t), time.Minute, + func(context.Context) (fantasy.StreamResponse, error) { + return nil, openErr + }, + metrics, stages, StageModel{Model: "claude"}, + ) + require.ErrorIs(t, err, openErr) + + require.Equal(t, codes.Error, ttftSpanStatus(t, spans).Status().Code) + _, ok := ttftStageCount(t, registry) + require.False(t, ok, "a failed window must not be observed") + }) + + t.Run("ReleaseWithoutPartRecordsSpanWithoutObservation", func(t *testing.T) { + t.Parallel() + stages, spans, registry, metrics := newFixture(t) + + attempt, err := guardedStream( + t.Context(), "anthropic", "claude", quartz.NewMock(t), time.Minute, + func(context.Context) (fantasy.StreamResponse, error) { + return fantasy.StreamResponse(func(func(fantasy.StreamPart) bool) {}), nil + }, + metrics, stages, StageModel{Model: "claude"}, + ) + require.NoError(t, err) + attempt.release() + + require.Equal(t, codes.Error, ttftSpanStatus(t, spans).Status().Code) + _, ok := ttftStageCount(t, registry) + require.False(t, ok, "a window closed without a token must not be observed") + }) +}