coder · jscottmiller · Sep 4, 2026
diff --git 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,37 +893,65 @@ 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,
 	clock quartz.Clock,
 	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
@@ -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")
+	})
+}