diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index 42b7caf7c93..b86285bc49a 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -778,6 +778,52 @@ State updates processed by the loop come from: The runner is responsible for subscribing to the `chat:update:{chat_id}` pubsub channel. During bootstrap, it must first subscribe to the channel and then fetch the initial state of the chat from the database to avoid missing any updates. +### Lifecycle tracing + +The runner owns a `chat_turn` trace span for the turn it is running, implemented by `runnerTurnSpan` in `turn_trace.go`. The span and the stages inside it are emitted through `chatloop.StageTracer`, which produces an OpenTelemetry span and an observation on the `coderd_chatd_stage_duration_seconds{stage, scope, chat_kind, model}` histogram from a single `End` call, so wherever both exist the trace and metric durations cannot disagree. Tracing is enabled by the `TracerProvider` server option. A nil provider disables spans without disabling the histogram. The `--chat-stage-metrics` option (`off`, `basic`, `full`; default `off`) filters the histogram without affecting spans: `off` registers none of the stage families, `basic` observes only the wait, connect, and model-call stages, and `full` observes every stage. `coderd_chatd_stage_metrics_level` reports the configured level. + +`chat_turn` is a standalone trace root. The HTTP request that triggered the turn ran on a different goroutine, and often a different replica, from the worker that runs it, and no trace context is persisted with the message, so there is nothing to parent the span to. + +#### Turn span lifecycle + +One `chat_turn` span covers one prompt, not one runner. A runner keeps ownership of a chat across queued-message promotions, and runners are also spawned for abandon, interrupt, and timeout tasks that run no turn, so the span is started lazily by the first generation task and replaced when the turn finishes. + +- Start: every iteration of the generation loop calls `Ensure`, which returns a context parented to the open turn and a `turnToken` identifying it. When no turn is open, `Ensure` starts one with its start timestamp backdated to the trigger message's `created_at` and records an `acquisition` stage from that instant to now, covering the time between the message landing in history and a worker picking the chat up. +- Complete: after the `FinishTurn` transition commits, the generation step calls `Complete`. This marks the turn finished but leaves the span open. If the transition promoted a queued message, `FinishTurnResult.PromotedQueuedAt` carries that message's `created_at` to `Complete`. +- Settle: when the step returns to the generation loop, after the step's own `generation_step` stage has ended, the loop calls `Settle`, which closes the span. Closing in two steps ensures the finishing step is counted inside the turn. If `Complete` recorded a promotion, `Settle` immediately opens the next turn anchored at the moment the promoted message was queued and records a `queue_wait` stage from that instant to the promotion. A turn opened this way records no `acquisition` stage, since the two windows would overlap. +- Next prompt: if a new prompt starts a generation task while a finished turn is still open, `Ensure` settles the old turn first and then opens a new one. +- Runner exit: the runner ends whatever span is still open when it shuts down. + +The runner cancels the active task and spawns its replacement without waiting for the old goroutine to exit (see [Event processing](#event-processing)), so an old task can still be unwinding while the new one calls `Ensure` and rotates the turn. Each task carries the `turnToken` returned by its own `Ensure` call, and `Complete` and `Settle` do nothing when the token does not identify the open turn. A stale task therefore cannot close the turn that replaced its own. + +Queued messages can also be promoted outside a generation step, through `PromoteQueued`. That path records `queue_wait` as a standalone stage, since no turn exists yet to attach it to. + +Work detached from the turn, such as title, summary, and status label generation, runs on a context with the span context stripped and `scope=background`, so its stages start their own trace roots and are separable from turn-scoped stages in the histogram. + +#### Stages + +Every stage carries `scope` (`turn` or `background`) and `chat_kind` (`root` or `subagent`), and once the model is resolved, `model`. Spans additionally carry `reasoning_effort`; it is not a metric label because it multiplies series per model. Stages that are not tied to a model call (`acquisition`, `queue_wait`, `mcp_connect`, `retry_backoff`, `commit`) carry an empty `model` label. `prepare` is stamped with the model once preparation resolves it. + +At `--chat-stage-metrics=basic` the histogram observes `chat_turn`, `queue_wait`, `acquisition`, `mcp_connect`, `stream`, `time_to_first_token`, `provider_attempt`, `tool_call`, `commit`, and `retry_backoff`. `generation_step`, `prepare`, `thinking`, and `compaction` are span-only at that level. + +Live stages wrap a section of code and end when it returns: + +- `generation_step`: one iteration of the generation loop, from loading state to applying a transition. It carries `generation_attempt` and `generation_action`. +- `prepare`: generation preparation, including model resolution and tool assembly. +- `mcp_connect`: connecting to the configured MCP servers, inside `prepare`. +- `provider_attempt`: one HTTP round trip to the model provider, emitted by the transport, so a retried request produces one stage per attempt. It ends when response headers arrive and is marked errored for HTTP status 400 and above. +- `stream`: the provider stream, from opening the request to consuming the last part. +- `time_to_first_token`: nested in `stream`, from opening the request to the first streamed part. A stream that ends before any part arrives closes this span with an error and records no histogram observation. +- `retry_backoff`: the wait before retrying a failed LLM API call. +- `commit`: the `CommitStep` transaction. +- `compaction`: a compaction pass. + +Reconstructed stages are recorded after the fact from timestamps captured elsewhere: + +- `acquisition` and `queue_wait`: described above. +- `thinking`: one per reasoning part, from the part's start to its completion timestamp in the persisted step. +- `tool_call`: one per local tool call, from the tool billing recorder's start and completion stamps. + ### Event shape Every event that the runner loop processes has the following shape: diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index aaf238e45f8..4e0982c5e26 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -21,6 +21,7 @@ import ( "github.com/google/uuid" "github.com/prometheus/client_golang/prometheus" "github.com/sqlc-dev/pqtype" + "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" "golang.org/x/sync/errgroup" "golang.org/x/xerrors" @@ -2145,9 +2146,10 @@ func (p *Server) PromoteQueued( } var ( - result PromoteQueuedResult - refreshChat database.Chat - refreshedOK bool + result PromoteQueuedResult + refreshChat database.Chat + refreshedOK bool + promotedQueuedAt time.Time ) machine := p.newChatMachine(opts.ChatID) updateErr := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { @@ -2167,6 +2169,7 @@ func (p *Server) PromoteQueued( } if promoteResult.InsertedMessage != nil { result.PromotedMessage = *promoteResult.InsertedMessage + promotedQueuedAt = promoteResult.QueuedMessage.CreatedAt } // Capture the chat inside the transaction so the watch event // published below uses the snapshot bump and status change @@ -2186,6 +2189,13 @@ func (p *Server) PromoteQueued( if refreshedOK { p.publishChatPubsubEvent(refreshChat, codersdk.ChatWatchEventKindStatusChange, nil) } + if !promotedQueuedAt.IsZero() { + var chatKind string + if refreshedOK { + chatKind = chatKindAttr(refreshChat) + } + p.recordQueueWait(ctx, opts.ChatID, chatKind, promotedQueuedAt, p.stages.Now()) + } return result, nil } @@ -4612,7 +4622,7 @@ func (p *Server) finalizeSuccessfulTurnStatusLabelWithAfterFunc( logger slog.Logger, afterFinalize func(context.Context, string), ) { - finalizeCtx, stopFinalizeCtx := p.inflightContext(ctx) + finalizeCtx, stopFinalizeCtx := p.inflightChatContext(ctx, chat) if err := p.goInflight(func() { defer stopFinalizeCtx() statusLabel := p.generateFinalTurnStatusLabel(finalizeCtx, chat, status, runResult, logger) @@ -4698,7 +4708,7 @@ func (p *Server) setLastTurnSummaryAsync( if chat.LastTurnSummary.Valid && strings.TrimSpace(chat.LastTurnSummary.String) == summary { return } - updateCtx, stopUpdateCtx := p.inflightContext(ctx) + updateCtx, stopUpdateCtx := p.inflightChatContext(ctx, chat) if err := p.goInflight(func() { defer stopUpdateCtx() p.updateLastTurnSummary(updateCtx, chat, chat.HistoryVersion, summary, logger) @@ -4718,7 +4728,7 @@ func (p *Server) clearLastTurnSummaryAsync( chat database.Chat, logger slog.Logger, ) { - clearCtx, stopClearCtx := p.inflightContext(ctx) + clearCtx, stopClearCtx := p.inflightChatContext(ctx, chat) if err := p.goInflight(func() { defer stopClearCtx() p.updateLastTurnSummary(clearCtx, chat, chat.HistoryVersion, "", logger) @@ -4813,7 +4823,7 @@ func (p *Server) maybeGenerateChatSummaryAsync( if chat.ParentChatID.Valid { return } - ctx, cancel := p.inflightContext(ctx) + ctx, cancel := p.inflightChatContext(ctx, chat) if err := p.goInflight(func() { defer cancel() p.generateAndStoreChatSummary(ctx, logger, chat) @@ -4994,7 +5004,7 @@ func (p *Server) storeSubagentReportSummaryAsync( chat database.Chat, logger slog.Logger, ) { - summaryCtx, stopSummaryCtx := p.inflightContext(ctx) + summaryCtx, stopSummaryCtx := p.inflightChatContext(ctx, chat) if err := p.goInflight(func() { defer stopSummaryCtx() p.storeSubagentReportSummary(summaryCtx, chat, logger) @@ -5090,7 +5100,13 @@ func (p *Server) Close() error { // must be called once the work completes to release the shutdown hook. // The caller is responsible for providing their own timeout. func (p *Server) inflightContext(reqCtx context.Context) (context.Context, func()) { - ctx, cancel := context.WithCancel(context.WithoutCancel(reqCtx)) + // Inflight work outlives the caller, so the caller's span and stage + // scope are stripped from the context: spans started on this context + // become their own roots instead of children that end after their + // parent, and their stages are recorded as background work. + detached := trace.ContextWithSpanContext(context.WithoutCancel(reqCtx), trace.SpanContext{}) + detached = chatloop.ContextWithScope(detached, chatloop.ScopeBackground) + ctx, cancel := context.WithCancel(detached) stop := context.AfterFunc(p.ctx, cancel) return ctx, func() { stop() @@ -5098,6 +5114,29 @@ func (p *Server) inflightContext(reqCtx context.Context) (context.Context, func( } } +// recordQueueWait emits the queue_wait stage for a message that sat +// queued from queuedAt until promotedAt. The span context is stripped +// from ctx so the stage is a standalone span rather than a child of +// the span in ctx, and the scope and chat kind are set explicitly +// because ctx does not carry the turn's. An empty chatKind records the +// stage without one. +func (p *Server) recordQueueWait(ctx context.Context, chatID uuid.UUID, chatKind string, queuedAt, promotedAt time.Time) { + standalone := trace.ContextWithSpanContext(ctx, trace.SpanContext{}) + standalone = chatloop.ContextWithChatKind(standalone, chatKind) + p.stages.RecordAs(standalone, chatloop.StageQueueWait, chatloop.ScopeTurn, + chatloop.StageModel{}, queuedAt, promotedAt, nil, + attribute.String(chatloop.AttrChatID, chatID.String()), + ) +} + +// inflightChatContext is inflightContext for work that belongs to a +// known chat. The chat kind is set on the returned context so the +// stages of the detached work carry it. +func (p *Server) inflightChatContext(reqCtx context.Context, chat database.Chat) (context.Context, func()) { + ctx, stop := p.inflightContext(reqCtx) + return chatloop.ContextWithChatKind(ctx, chatKindAttr(chat)), stop +} + func (p *Server) goInflight(f func()) error { if p.inflightClosed.Load() { return errInflightClosed diff --git a/coderd/x/chatd/chatstate/transitions.go b/coderd/x/chatd/chatstate/transitions.go index 6a607597a6f..a0ed1e40cc4 100644 --- a/coderd/x/chatd/chatstate/transitions.go +++ b/coderd/x/chatd/chatstate/transitions.go @@ -1448,6 +1448,9 @@ type FinishTurnInput struct{} type FinishTurnResult struct { Chat database.Chat PromotedMessage *database.ChatMessage + // PromotedQueuedAt is the queued row's creation time when this + // transition promoted a queue head, and the zero time otherwise. + PromotedQueuedAt time.Time } // FinishTurn completes a running turn. @@ -1509,8 +1512,9 @@ func (tx *Tx) FinishTurn(_ FinishTurnInput) (FinishTurnResult, error) { promoted = &inserted[len(inserted)-1] } return FinishTurnResult{ - Chat: updated, - PromotedMessage: promoted, + Chat: updated, + PromotedMessage: promoted, + PromotedQueuedAt: head.CreatedAt, }, nil } diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 11a0163fd2a..e2aff186525 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -5,11 +5,13 @@ import ( "encoding/json" "errors" "strings" + "sync" "time" "charm.land/fantasy" "github.com/google/uuid" "github.com/sqlc-dev/pqtype" + "go.opentelemetry.io/otel/attribute" "golang.org/x/xerrors" "cdr.dev/slog/v3" @@ -68,6 +70,8 @@ type generationPrepared struct { // user-facing errors. See chatloop.GenerateAssistantOptions.ErrorProvider. ResolvedProvider string + StageModel chatloop.StageModel + ModelConfigID uuid.UUID CallTemplate fantasy.Call ContextLimitFallback int64 @@ -441,11 +445,16 @@ func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskS if err != nil { return xerrors.Errorf("load generation state: %w", err) } + turnCtx, turnToken := input.Turn.Ensure(ctx, chat, triggerMessageTime(messages)) + input.TurnToken = turnToken var again bool - input, again, err = s.runGenerationStep(ctx, machine, input, chat, messages) + input, again, err = s.runGenerationStep(turnCtx, machine, input, chat, messages) if again { continue } + // The step's stage has ended by now, so a turn the step finished + // closes with that stage counted. + input.Turn.Settle(ctx, input.TurnToken) return err } } @@ -462,6 +471,12 @@ func (s *taskStarter) runGenerationStep( chat database.Chat, messages []database.ChatMessage, ) (next chatWorkerTaskStartInput, again bool, err error) { + ctx, stepSpan := s.server.stages.Start(ctx, chatloop.StageGenerationStep, + attribute.String(chatloop.AttrChatID, input.ChatID.String()), + attribute.Int64(chatloop.AttrGenerationAttempt, input.GenerationAttempt), + ) + defer func() { stepSpan.End(err) }() + if s.server.hooks.Enabled() { result, dispatched, err := s.startGenerationSession(ctx, machine, input, chat, messages) if err != nil { @@ -480,9 +495,18 @@ func (s *taskStarter) runGenerationStep( Messages: messages, RecordMCPConnectSummaries: input.DebugTurn.RecordMCPConnectSummaries, } - prepared, err := retryGenerationPhase(ctx, s, "prepare", func() (generationPrepared, error) { - return s.server.prepareGeneration(ctx, prepareInput) + prepareCtx, prepareSpan := s.server.stages.Start(ctx, chatloop.StagePrepare) + prepared, err := retryGenerationPhase(prepareCtx, s, "prepare", func() (generationPrepared, error) { + return s.server.prepareGeneration(prepareCtx, prepareInput) }) + if err == nil { + providerAttr := attribute.String(chatloop.AttrProvider, prepared.ResolvedProvider) + prepareSpan.SetAttributes(providerAttr) + prepareSpan.SetModel(prepared.StageModel) + stepSpan.SetAttributes(providerAttr) + stepSpan.SetModel(prepared.StageModel) + } + prepareSpan.End(err) if err != nil { if errors.Is(err, errTaskExpectedExit) || errors.Is(err, errTaskRetryable) { return input, false, xerrors.Errorf("prepare generation: %w", err) @@ -526,6 +550,7 @@ func (s *taskStarter) runGenerationStep( return input, false, s.finishGenerationError(ctx, machine, input, err, generationAttemptNotRequired) } + stepSpan.SetAttributes(attribute.String(chatloop.AttrGenerationAction, string(decision.kind))) var actionErr error switch decision.kind { case generationActionEnterRequiresAction: @@ -668,13 +693,17 @@ func (*taskStarter) recordGenerationRetry( } func (s *taskStarter) waitGenerationRetry(ctx context.Context, delay time.Duration) error { + _, span := s.server.stages.Start(ctx, chatloop.StageRetryBackoff) timer := s.opts.Clock.NewTimer(delay, "chatworker", "generation-retry") defer timer.Stop() select { case <-timer.C: + span.End(nil) return nil case <-ctx.Done(): - return errors.Join(errTaskExpectedExit, xerrors.Errorf("wait generation retry: %w", ctx.Err())) + err := errors.Join(errTaskExpectedExit, xerrors.Errorf("wait generation retry: %w", ctx.Err())) + span.End(err) + return err } } @@ -773,10 +802,13 @@ func (s *taskStarter) generateAssistant( Logger: s.opts.Logger, Clock: s.opts.Clock, Metrics: s.server.metrics, + Stages: s.server.stages, + StageModel: prepared.StageModel, }) if err != nil { return xerrors.Errorf("generate assistant: %w", err) } + s.recordThinkingStages(runCtx, prepared, outcome.Step) if len(outcome.Step.Content) == 0 { return s.finishGenerationTurn(ctx, machine, input, generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, requireGenerationAttempt(attempt.number)) } @@ -803,6 +835,31 @@ func (s *taskStarter) generateAssistant( return s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionGenerateAssistant, messages, generationCommitHooks{}) } +// recordThinkingStages emits one thinking stage per reasoning part of +// a step, bounded by the part's own start and completion timestamps. +// Parts are paired by index; a part without a completion timestamp +// ends the sweep because later indexes cannot be paired either. +func (s *taskStarter) recordThinkingStages( + ctx context.Context, + prepared generationPrepared, + step chatloop.PersistedStep, +) { + for index, startedAt := range step.ReasoningStartedAt { + if index >= len(step.ReasoningCompletedAt) { + return + } + s.server.stages.Record( + ctx, + chatloop.StageThinking, + prepared.StageModel, + startedAt, + step.ReasoningCompletedAt[index], + nil, + attribute.String(chatloop.AttrProvider, prepared.ResolvedProvider), + ) + } +} + func (s *taskStarter) admitStepToolCalls( ctx context.Context, input chatWorkerTaskStartInput, @@ -881,6 +938,41 @@ func (r *bufferToolBillingRecorder) RecordComplete(dispatchIndex int, completedA r.recordComplete(r.allowedIndexes[dispatchIndex], completedAt) } +// toolCallStageRecorder emits one tool_call stage per local tool call +// while forwarding every callback to the billing recorder it wraps. +// Completions arrive on the tool goroutines, so the pending starts are +// mutex guarded. +type toolCallStageRecorder struct { + inner chatloop.ToolBillingRecorder + record func(dispatchIndex int, startedAt, completedAt time.Time) + + mu sync.Mutex + starts map[int]time.Time +} + +func (r *toolCallStageRecorder) RecordStart(dispatchIndex int, startedAt time.Time) { + if r.inner != nil { + r.inner.RecordStart(dispatchIndex, startedAt) + } + r.mu.Lock() + r.starts[dispatchIndex] = startedAt + r.mu.Unlock() +} + +func (r *toolCallStageRecorder) RecordComplete(dispatchIndex int, completedAt time.Time) { + if r.inner != nil { + r.inner.RecordComplete(dispatchIndex, completedAt) + } + r.mu.Lock() + startedAt, started := r.starts[dispatchIndex] + delete(r.starts, dispatchIndex) + r.mu.Unlock() + if !started { + return + } + r.record(dispatchIndex, startedAt, completedAt) +} + func (s *taskStarter) executeLocalTools( ctx context.Context, machine *chatstate.ChatMachine, @@ -928,6 +1020,20 @@ func (s *taskStarter) executeLocalTools( recordComplete: attempt.recordToolCompletion, } } + toolStages := &toolCallStageRecorder{ + inner: billingRecorder, + starts: make(map[int]time.Time, len(allowed)), + record: func(dispatchIndex int, startedAt, completedAt time.Time) { + toolName := "" + if dispatchIndex >= 0 && dispatchIndex < len(allowed) { + toolName = allowed[dispatchIndex].ToolName + } + s.server.stages.Record(ctx, chatloop.StageToolCall, prepared.StageModel, startedAt, completedAt, nil, + attribute.String(chatloop.AttrToolName, toolName), + attribute.String(chatloop.AttrProvider, provider), + ) + }, + } outcome, err = chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ Tools: prepared.Tools, ActiveTools: prepared.ActiveTools, @@ -942,7 +1048,7 @@ func (s *taskStarter) executeLocalTools( ContextLimit: prepared.ContextLimitFallback, ToolNameAliases: subagentToolNameAliases, UnbilledToolNames: unbilledSubagentToolNames, - BillingRecorder: billingRecorder, + BillingRecorder: toolStages, PublishMessagePart: attempt.publish, Logger: s.opts.Logger, Metrics: s.server.metrics, @@ -1072,7 +1178,13 @@ func (s *taskStarter) generateCompaction( // debug run; without it startCompactionDebugRun finds no parent and // skips debug instrumentation entirely. runCtx := input.DebugTurn.Ensure(ctx, prepared.Chat, prepared.Debug) - outcome, err := chatloop.GenerateCompaction(runCtx, compactionOpts) + compactionCtx, compactionSpan := s.server.stages.Start(runCtx, chatloop.StageCompaction, + attribute.String(chatloop.AttrProvider, metricProvider), + attribute.String(chatloop.AttrCompactionSource, string(source)), + ) + compactionSpan.SetModel(compactionStageModel(prepared, metricModel)) + outcome, err := chatloop.GenerateCompaction(compactionCtx, compactionOpts) + compactionSpan.End(err) if err != nil { s.server.metrics.RecordCompaction(metricProvider, metricModel, false, err) return xerrors.Errorf("generate compaction: %w", err) @@ -1128,6 +1240,18 @@ func (s *taskStarter) generateCompaction( return nil } +// compactionStageModel labels the compaction stage with the summary +// model. The effort carries over only when the chat model runs the +// summary; an override resolves its own call options, which this call +// site does not see. +func compactionStageModel(prepared generationPrepared, summaryModel string) chatloop.StageModel { + model := chatloop.StageModel{Model: summaryModel} + if prepared.Compaction != nil && prepared.Compaction.Override == nil { + model.Effort = prepared.StageModel.Effort + } + return model +} + // compactionMetricIdentity returns the provider/model labels for compaction // metrics. Override labels come from prepare-time resolution so events // recorded before the override client is built (still-over-limit) match @@ -1266,7 +1390,11 @@ func (s *taskStarter) commitGenerationStep( } var committed database.Chat insertedMessages := []runnerActionMessage{} - err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + commitCtx, commitSpan := s.server.stages.Start(ctx, chatloop.StageCommit, + attribute.String(chatloop.AttrGenerationAction, string(kind)), + attribute.Int64(chatloop.AttrGenerationAttempt, attempt), + ) + err := machine.Update(commitCtx, func(tx *chatstate.Tx, store database.Store) error { if _, err := loadChatForGeneration(ctx, store, input, requireGenerationAttempt(attempt)); err != nil { return xerrors.Errorf("load chat for generation: %w", err) } @@ -1297,6 +1425,7 @@ func (s *taskStarter) commitGenerationStep( committed = loadedChat return nil }) + commitSpan.End(err) if err != nil { return normalizeTaskTransitionError(err, "commit generation step") } @@ -1436,6 +1565,7 @@ func (s *taskStarter) finishGenerationTurnWithoutHook( fence generationAttemptFence, ) error { var committed database.Chat + var promotedQueuedAt time.Time err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { if _, err := loadChatForGeneration(ctx, store, input, fence); err != nil { return xerrors.Errorf("load chat for generation: %w", err) @@ -1446,6 +1576,7 @@ func (s *taskStarter) finishGenerationTurnWithoutHook( } if finishResult.PromotedMessage != nil { decision.promotedMessageID = finishResult.PromotedMessage.ID + promotedQueuedAt = finishResult.PromotedQueuedAt } committed = finishResult.Chat return nil @@ -1455,6 +1586,7 @@ func (s *taskStarter) finishGenerationTurnWithoutHook( recordGenerationFinishFailure(input.DebugTurn, err) return err } + input.Turn.Complete(input.TurnToken, promotedQueuedAt) return s.completeGenerationTurn(ctx, input, committed, decision.promotedMessageID) } @@ -1503,6 +1635,7 @@ func (s *taskStarter) finishGenerationTurn( continueTurn := strings.TrimSpace(response.GetModelContext()) != "" && input.StopNudges.claim(nudgeKey) var committed database.Chat + var promotedQueuedAt time.Time err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { if _, err := loadChatForGeneration(ctx, store, input, fence); err != nil { return xerrors.Errorf("load chat for generation: %w", err) @@ -1519,6 +1652,7 @@ func (s *taskStarter) finishGenerationTurn( } if finishResult.PromotedMessage != nil { decision.promotedMessageID = finishResult.PromotedMessage.ID + promotedQueuedAt = finishResult.PromotedQueuedAt } committed = finishResult.Chat return nil @@ -1545,6 +1679,7 @@ func (s *taskStarter) finishGenerationTurn( Kind: runnerActionKind(generationActionGenerateAssistant), }) } + input.Turn.Complete(input.TurnToken, promotedQueuedAt) return s.completeGenerationTurn(ctx, input, committed, decision.promotedMessageID) } diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index ff49843deb6..2e01ebff20f 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -369,8 +369,9 @@ func (server *Server) prepareGeneration( logger.Warn(ctx, "failed to load MCP user tokens", slog.Error(tokenErr)) } mcpTokens = server.refreshExpiredMCPTokens(ctx, logger, mcpConnectConfigs, mcpTokens) + connectCtx, connectSpan := server.stages.Start(ctx, chatloop.StageMCPConnect) mcpTools, mcpSummaries, mcpCleanup = mcpclient.ConnectAll( - ctx, + connectCtx, logger, mcpConnectConfigs, mcpTokens, @@ -378,6 +379,7 @@ func (server *Server) prepareGeneration( server.oidcTokenSource, chatprovider.CoderHeaders(chat), ) + connectSpan.End(nil) return nil }) } @@ -738,6 +740,7 @@ func (server *Server) prepareGeneration( ModelRoute: modelRoute, ModelBuildOptions: modelOpts, ResolvedProvider: resolved.resolvedProvider, + StageModel: resolved.stageModel(), ModelConfigID: modelConfig.ID, CallTemplate: resolved.newCall(), ContextLimitFallback: modelConfig.ContextLimit, diff --git a/coderd/x/chatd/options.go b/coderd/x/chatd/options.go index b887b928baf..c15acfcb6a3 100644 --- a/coderd/x/chatd/options.go +++ b/coderd/x/chatd/options.go @@ -68,8 +68,12 @@ type chatWorkerTaskStartInput struct { Status database.ChatStatus RequiresActionDeadlineAt sql.NullTime DebugTurn *runnerDebugTurn - SessionStart *sessionStartTracker - StopNudges *stopNudgeTracker + Turn *runnerTurnSpan + // TurnToken identifies the turn this task's steps run in. The zero + // token identifies no turn. + TurnToken turnToken + SessionStart *sessionStartTracker + StopNudges *stopNudgeTracker } func (i chatWorkerTaskStartInput) hookTurnID() *uuid.UUID { diff --git a/coderd/x/chatd/runner.go b/coderd/x/chatd/runner.go index 72b916a38a6..121452021a9 100644 --- a/coderd/x/chatd/runner.go +++ b/coderd/x/chatd/runner.go @@ -57,6 +57,7 @@ type runner struct { tasksByIndex map[taskIndexKey]taskInstanceID localLocks *localLockSet debugTurn *runnerDebugTurn + turnSpan *runnerTurnSpan sessionStart sessionStartTracker stopNudges stopNudgeTracker } @@ -71,6 +72,7 @@ func newRunner(ctx context.Context, mgr *runnerManager, rec *runnerRecord, opts tasksByIndex: make(map[taskIndexKey]taskInstanceID), localLocks: newLocalLockSet(), debugTurn: newRunnerDebugTurn(ctx, opts.Logger), + turnSpan: newRunnerTurnSpan(mgr.server.stages), } } @@ -86,6 +88,7 @@ func (r *runner) run() { r.cancelActiveTask() r.waitForTasks() r.closeDebugTurn() + r.turnSpan.End(nil) return } } @@ -229,6 +232,7 @@ func (r *runner) spawnTaskIfNeeded(kind taskKind, state runnerStateUpdate) { Status: state.Status, RequiresActionDeadlineAt: state.RequiresActionDeadlineAt, DebugTurn: r.debugTurn, + Turn: r.turnSpan, SessionStart: &r.sessionStart, StopNudges: &r.stopNudges, } diff --git a/coderd/x/chatd/stage_internal_test.go b/coderd/x/chatd/stage_internal_test.go index c5e7a15befc..4ec04fcc2b0 100644 --- a/coderd/x/chatd/stage_internal_test.go +++ b/coderd/x/chatd/stage_internal_test.go @@ -4,9 +4,12 @@ import ( "context" "io" "net/http" + "sort" "strings" "testing" + "time" + "github.com/google/uuid" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -15,6 +18,7 @@ import ( "go.opentelemetry.io/otel/trace" "golang.org/x/xerrors" + "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/x/chatd/chatloop" ) @@ -182,3 +186,235 @@ func TestStageSpanRoundTripperModel(t *testing.T) { require.Contains(t, ended[0].Attributes(), attribute.String(chatloop.AttrReasoningEffort, model.Effort)) } + +func TestRunnerTurnSpanStartsAtTriggerMessage(t *testing.T) { + t.Parallel() + tracer, recorder := newStageTestTracer(t) + turn := newRunnerTurnSpan(tracer) + chat := database.Chat{ID: uuid.New()} + triggerAt := time.Now().Add(-2 * time.Second) + + turnCtx, _ := turn.Ensure(t.Context(), chat, triggerAt) + turn.End(nil) + + ended := recorder.Ended() + require.Len(t, ended, 2) + acquisition, chatTurn := ended[0], ended[1] + require.Equal(t, chatloop.StageAcquisition, acquisition.Name()) + require.Equal(t, chatloop.StageChatTurn, chatTurn.Name()) + require.Equal(t, triggerAt.UTC(), chatTurn.StartTime().UTC()) + require.Equal(t, triggerAt.UTC(), acquisition.StartTime().UTC()) + require.False(t, acquisition.StartTime().Before(chatTurn.StartTime())) + require.False(t, acquisition.EndTime().After(chatTurn.EndTime())) + require.Equal(t, chatTurn.SpanContext().SpanID(), acquisition.Parent().SpanID()) + require.Equal(t, chatTurn.SpanContext().TraceID(), trace.SpanContextFromContext(turnCtx).TraceID()) +} + +func TestRunnerTurnSpanParentsRecordedStages(t *testing.T) { + t.Parallel() + tracer, recorder := newStageTestTracer(t) + turn := newRunnerTurnSpan(tracer) + chat := database.Chat{ID: uuid.New()} + + turnCtx, _ := turn.Ensure(t.Context(), chat, time.Now().Add(-time.Second)) + stepCtx, step := tracer.Start(turnCtx, chatloop.StageGenerationStep) + step.End(nil) + + // The step has already ended, so the queue wait is recorded + // against the turn context the runner keeps. + tracer.Record(turn.Context(stepCtx), chatloop.StageQueueWait, chatloop.StageModel{}, + time.Now().Add(-500*time.Millisecond), time.Now(), nil) + turn.End(nil) + + var queueWait, chatTurn, generationStep sdktrace.ReadOnlySpan + for _, span := range recorder.Ended() { + switch span.Name() { + case chatloop.StageQueueWait: + queueWait = span + case chatloop.StageChatTurn: + chatTurn = span + case chatloop.StageGenerationStep: + generationStep = span + } + } + require.NotNil(t, queueWait) + require.NotNil(t, chatTurn) + require.NotNil(t, generationStep) + require.Equal(t, chatTurn.SpanContext().SpanID(), queueWait.Parent().SpanID()) + require.NotEqual(t, generationStep.SpanContext().SpanID(), queueWait.Parent().SpanID()) +} + +func TestServerRecordQueueWaitIsStandalone(t *testing.T) { + t.Parallel() + tracer, recorder := newStageTestTracer(t) + server := &Server{stages: tracer} + + // The promoting request has its own span, which the queue wait must + // not join. + requestCtx, requestSpan := tracer.Start(t.Context(), chatloop.StageCommit) + queuedAt := time.Now().Add(-30 * time.Second) + server.recordQueueWait(requestCtx, uuid.New(), chatloop.ChatKindSubagent, queuedAt, queuedAt.Add(20*time.Second)) + requestSpan.End(nil) + + var queueWait, request sdktrace.ReadOnlySpan + for _, span := range recorder.Ended() { + switch span.Name() { + case chatloop.StageQueueWait: + queueWait = span + case chatloop.StageCommit: + request = span + } + } + require.NotNil(t, queueWait) + require.NotNil(t, request) + require.False(t, queueWait.Parent().IsValid()) + require.NotEqual(t, request.SpanContext().TraceID(), queueWait.SpanContext().TraceID()) + require.Contains(t, queueWait.Attributes(), + attribute.String(chatloop.AttrScope, chatloop.ScopeTurn)) + require.Contains(t, queueWait.Attributes(), + attribute.String(chatloop.AttrChatKind, chatloop.ChatKindSubagent)) +} + +func TestServerInflightContextIsBackgroundScoped(t *testing.T) { + t.Parallel() + tracer, recorder := newStageTestTracer(t) + serverCtx, serverCancel := context.WithCancel(context.Background()) + t.Cleanup(serverCancel) + server := &Server{ctx: serverCtx, stages: tracer} + + chat := database.Chat{ID: uuid.New()} + turn := newRunnerTurnSpan(tracer) + turnCtx, _ := turn.Ensure(t.Context(), chat, time.Now().Add(-time.Second)) + inflightCtx, stop := server.inflightChatContext(turnCtx, chat) + t.Cleanup(stop) + + _, span := tracer.Start(inflightCtx, chatloop.StageGenerationStep) + span.End(nil) + turn.End(nil) + + for _, ended := range recorder.Ended() { + if ended.Name() != chatloop.StageGenerationStep { + continue + } + require.False(t, ended.Parent().IsValid()) + require.Contains(t, ended.Attributes(), + attribute.String(chatloop.AttrScope, chatloop.ScopeBackground)) + // Detached work keeps the chat it belongs to, so background + // stages stay attributable to a chat kind. + require.Contains(t, ended.Attributes(), + attribute.String(chatloop.AttrChatKind, chatloop.ChatKindRoot)) + } +} + +func TestRunnerTurnSpanCarriesChatKind(t *testing.T) { + t.Parallel() + tracer, recorder := newStageTestTracer(t) + turn := newRunnerTurnSpan(tracer) + chat := database.Chat{ID: uuid.New(), ParentChatID: uuid.NullUUID{UUID: uuid.New(), Valid: true}} + + turnCtx, _ := turn.Ensure(t.Context(), chat, time.Now().Add(-time.Second)) + _, step := tracer.Start(turnCtx, chatloop.StageGenerationStep) + step.End(nil) + turn.End(nil) + + require.NotEmpty(t, recorder.Ended()) + for _, span := range recorder.Ended() { + require.Contains(t, span.Attributes(), + attribute.String(chatloop.AttrChatKind, chatloop.ChatKindSubagent), + "stage %s must carry the turn's chat kind", span.Name()) + } +} + +// turnSpansByStart returns the chat_turn spans the recorder saw, +// ordered by start time. +func turnSpansByStart(t *testing.T, recorder *tracetest.SpanRecorder) []sdktrace.ReadOnlySpan { + t.Helper() + var turns []sdktrace.ReadOnlySpan + for _, span := range recorder.Ended() { + if span.Name() == chatloop.StageChatTurn { + turns = append(turns, span) + } + } + sort.Slice(turns, func(i, j int) bool { + return turns[i].StartTime().Before(turns[j].StartTime()) + }) + return turns +} + +func TestRunnerTurnSpanSettleRotatesOnPromotion(t *testing.T) { + t.Parallel() + tracer, recorder := newStageTestTracer(t) + turn := newRunnerTurnSpan(tracer) + chat := database.Chat{ID: uuid.New()} + + turnCtx, token := turn.Ensure(t.Context(), chat, time.Now().Add(-time.Minute)) + _, step := tracer.Start(turnCtx, chatloop.StageGenerationStep) + queuedAt := time.Now().Add(-30 * time.Second) + turn.Complete(token, queuedAt) + promotedBy := time.Now() + step.End(nil) + turn.Settle(t.Context(), token) + + // The next turn is open, anchored at the promoted message. + _, nextToken := turn.Ensure(t.Context(), chat, queuedAt) + require.NotEqual(t, token, nextToken) + turn.End(nil) + + turns := turnSpansByStart(t, recorder) + require.Len(t, turns, 2) + require.Equal(t, queuedAt.UTC(), turns[1].StartTime().UTC()) + require.NotEqual(t, turns[0].SpanContext().TraceID(), turns[1].SpanContext().TraceID()) + require.Contains(t, turns[1].Attributes(), + attribute.String(chatloop.AttrChatKind, chatloop.ChatKindRoot)) + + // The promoted message's wait belongs to the turn it opens and ends + // when the promotion happened, not when the turn settled. + var queueWait sdktrace.ReadOnlySpan + for _, span := range recorder.Ended() { + if span.Name() == chatloop.StageQueueWait { + queueWait = span + } + } + require.NotNil(t, queueWait) + require.Equal(t, turns[1].SpanContext().SpanID(), queueWait.Parent().SpanID()) + require.Equal(t, queuedAt.UTC(), queueWait.StartTime().UTC()) + require.False(t, queueWait.EndTime().After(promotedBy)) + + // The rotated turn's head is the queue wait, so it records no + // acquisition: both windows start at the same instant and both + // would count as scheduling time. + var acquisitions int + for _, span := range recorder.Ended() { + if span.Name() == chatloop.StageAcquisition { + acquisitions++ + require.Equal(t, turns[0].SpanContext().SpanID(), span.Parent().SpanID()) + } + } + require.Equal(t, 1, acquisitions) +} + +func TestRunnerTurnSpanEnsureOpensTurnPerPrompt(t *testing.T) { + t.Parallel() + tracer, recorder := newStageTestTracer(t) + turn := newRunnerTurnSpan(tracer) + chat := database.Chat{ID: uuid.New()} + + firstTrigger := time.Now().Add(-2 * time.Minute) + _, first := turn.Ensure(t.Context(), chat, firstTrigger) + // A second prompt on the same runner reuses the open turn until it + // finishes. + _, again := turn.Ensure(t.Context(), chat, firstTrigger) + require.Equal(t, first, again) + require.Len(t, turnSpansByStart(t, recorder), 0) + + turn.Complete(first, time.Time{}) + secondTrigger := time.Now().Add(-time.Minute) + _, second := turn.Ensure(t.Context(), chat, secondTrigger) + require.NotEqual(t, first, second) + turn.End(nil) + + turns := turnSpansByStart(t, recorder) + require.Len(t, turns, 2) + require.Equal(t, firstTrigger.UTC(), turns[0].StartTime().UTC()) + require.Equal(t, secondTrigger.UTC(), turns[1].StartTime().UTC()) +} diff --git a/coderd/x/chatd/turn_trace.go b/coderd/x/chatd/turn_trace.go new file mode 100644 index 00000000000..b26b56a555b --- /dev/null +++ b/coderd/x/chatd/turn_trace.go @@ -0,0 +1,258 @@ +package chatd + +import ( + "context" + "sync" + "time" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatloop" +) + +// turnToken identifies one turn. Methods that take a token act only +// while that turn is the open one, so a holder of a token for a turn +// that has since been replaced cannot finish or invalidate the +// replacement. +type turnToken uint64 + +// runnerTurnSpan owns the chat_turn span of one runner. A runner is +// created when a chat is acquired and torn down when the chat leaves +// the running states, so the runner's lifetime bounds the turns it +// runs. +// +// The span is created on the first generation task rather than at +// runner construction: runners are also spawned to abandon or time +// out chats, which run no turn. One runner can run several turns in +// sequence: a finished turn is replaced by a new span, either when a +// queued message is promoted or when the next prompt starts a task. +// +// A turn closes in two steps: Complete marks it finished and Settle +// closes the span, so stages still open at Complete end inside the +// turn's span if they end before Settle. +type runnerTurnSpan struct { + stages *chatloop.StageTracer + + mu sync.Mutex + span *chatloop.StageSpan + spanCtx trace.SpanContext + chatID string + chatKind string + // token is the identity of the open turn. It advances every time a + // turn span is started. + token turnToken + // open is true while a turn span exists that has not been closed. + open bool + ended bool + // finished marks an open turn that reached a terminal transition + // and is waiting for Settle to close it. + finished bool + // pendingPromotion is the queued message the finishing transition + // promoted, when it promoted one. Settle opens the next turn from + // it. + pendingPromotion *turnPromotion +} + +// turnPromotion records a queued message promoted by the transition +// that finished a turn: when the message was queued and when the +// promotion happened. +type turnPromotion struct { + queuedAt time.Time + promotedAt time.Time +} + +func newRunnerTurnSpan(stages *chatloop.StageTracer) *runnerTurnSpan { + return &runnerTurnSpan{stages: stages} +} + +// Ensure returns a context parented to the open chat_turn span and the +// token of that turn, starting the span when none is open. triggerAt is +// the trigger message's creation time and becomes the span's start +// timestamp, so the acquisition stage reconstructed from the same +// instant falls inside the turn. The acquisition stage carries no model +// identity: the turn's model is not resolved until preparation runs. +// +// A turn that already reached a terminal transition is settled first: +// the prompt this call runs is a new turn, and folding it into the old +// span would report the two as one. +// +// The span is a standalone trace root. The request that triggered the +// turn is handled by a different goroutine, and often a different +// replica, than the worker that runs it, so no inbound span context +// is available here to link. +func (t *runnerTurnSpan) Ensure(ctx context.Context, chat database.Chat, triggerAt time.Time) (context.Context, turnToken) { + if t == nil { + return ctx, 0 + } + t.mu.Lock() + defer t.mu.Unlock() + if t.ended { + return ctx, 0 + } + if t.open && t.finished { + t.settleLocked(ctx) + } + if t.open { + return t.contextLocked(ctx), t.token + } + t.chatID = chat.ID.String() + t.chatKind = chatKindAttr(chat) + turnCtx := t.startLocked(ctx, triggerAt) + // The window between the trigger message landing in history and a + // worker picking the chat up is the acquisition. A turn opened by a + // promotion has no acquisition: its head is the queue wait of the + // message that opened it, and recording both would count that + // window twice. + t.stages.Record(turnCtx, chatloop.StageAcquisition, chatloop.StageModel{}, + triggerAt, t.stages.Now(), nil, + attribute.String(chatloop.AttrChatID, t.chatID)) + return t.contextLocked(ctx), t.token +} + +// startLocked opens a chat_turn span and returns the context parented +// to it. +func (t *runnerTurnSpan) startLocked(ctx context.Context, startAt time.Time) context.Context { + t.token++ + t.open = true + t.finished = false + t.pendingPromotion = nil + + // The chat kind rides on the context so every stage of the turn + // carries the kind. + ctx = chatloop.ContextWithChatKind(ctx, t.chatKind) + turnCtx, span := t.stages.StartRootAt(ctx, chatloop.StageChatTurn, startAt, nil, + attribute.String(chatloop.AttrChatID, t.chatID)) + t.span = span + t.spanCtx = span.SpanContext() + return turnCtx +} + +// Context returns ctx parented to the chat_turn span, or ctx +// unchanged while no turn span is open. +func (t *runnerTurnSpan) Context(ctx context.Context) context.Context { + if t == nil { + return ctx + } + t.mu.Lock() + defer t.mu.Unlock() + return t.contextLocked(ctx) +} + +func (t *runnerTurnSpan) contextLocked(ctx context.Context) context.Context { + if !t.open || t.ended { + return ctx + } + // The scope and chat kind are set independently of the span context + // so stages run on this context keep them when tracing is not + // recording. + ctx = chatloop.ContextWithScope(ctx, chatloop.ScopeTurn) + ctx = chatloop.ContextWithChatKind(ctx, t.chatKind) + if !t.spanCtx.IsValid() { + return ctx + } + return trace.ContextWithSpanContext(ctx, t.spanCtx) +} + +// Complete marks the turn identified by token as finished normally. +// The span stays open until Settle. +// +// A non-zero queuedAt is the creation time of a queued message the +// finishing transition promoted. Settle opens the next turn anchored +// at it, because the wait that message served and the work it causes +// belong to the turn it starts rather than the one that released it. +func (t *runnerTurnSpan) Complete(token turnToken, queuedAt time.Time) { + if t == nil { + return + } + t.mu.Lock() + defer t.mu.Unlock() + if !t.ownsLocked(token) { + return + } + t.finished = true + if !queuedAt.IsZero() { + t.pendingPromotion = &turnPromotion{queuedAt: queuedAt, promotedAt: t.stages.Now()} + } +} + +// Settle closes the turn identified by token if Complete marked it +// finished, and opens the next turn when the finishing transition +// promoted a queued message. A turn that is not finished is left open. +func (t *runnerTurnSpan) Settle(ctx context.Context, token turnToken) { + if t == nil { + return + } + t.mu.Lock() + defer t.mu.Unlock() + if !t.ownsLocked(token) || !t.finished { + return + } + t.settleLocked(ctx) +} + +// ownsLocked reports whether token identifies the open turn. +func (t *runnerTurnSpan) ownsLocked(token turnToken) bool { + return t.open && !t.ended && token == t.token +} + +// settleLocked closes the open, finished turn. When the finishing +// transition promoted a queued message, it opens the next turn anchored +// at the moment the message was queued and records that message's +// queue wait against the new turn. +func (t *runnerTurnSpan) settleLocked(ctx context.Context) { + promotion := t.pendingPromotion + t.closeLocked(nil) + if promotion == nil { + return + } + turnCtx := t.startLocked(ctx, promotion.queuedAt) + t.stages.Record(turnCtx, chatloop.StageQueueWait, chatloop.StageModel{}, + promotion.queuedAt, promotion.promotedAt, nil, + attribute.String(chatloop.AttrChatID, t.chatID)) +} + +// End closes the chat_turn span. Later calls are ignored. +func (t *runnerTurnSpan) End(err error) { + if t == nil { + return + } + t.mu.Lock() + defer t.mu.Unlock() + if t.ended || !t.open { + t.ended = true + return + } + t.ended = true + t.closeLocked(err) +} + +// closeLocked ends the open turn span. +func (t *runnerTurnSpan) closeLocked(err error) { + t.span.End(err) + t.span = nil + t.spanCtx = trace.SpanContext{} + t.open = false + t.finished = false + t.pendingPromotion = nil +} + +// chatKindAttr labels a chat as a subagent or a top-level chat. +func chatKindAttr(chat database.Chat) string { + if chat.ParentChatID.Valid { + return chatloop.ChatKindSubagent + } + return chatloop.ChatKindRoot +} + +// triggerMessageTime returns the creation time of the message that +// triggered the turn, which is the last user prompt in history. It +// returns the zero time when the history holds no user prompt. +func triggerMessageTime(messages []database.ChatMessage) time.Time { + index := lastUserPromptIndex(messages) + if index == -1 { + return time.Time{} + } + return messages[index].CreatedAt +}