Files
coder/coderd/x/chatd/message_conversion.go
T
4b7494be72 feat: harden chat generation runtime instrumentation for billing (#27451)
Closes CODAGT-835

## Summary

`chat_messages.runtime_ms` becomes the billing source of truth for Coder
Agents runtime (summed hourly by #27312), but it was built for
debugging: the June refactor (#26270) silently stopped recording
tool-step runtime, compaction was never measured, and interrupted turns
lost their partial runtime entirely. This PR defines the billable
metric, closes the paths that dropped it, and documents the definition
where the data lives.

## The billable definition

**`runtime_ms` is the wall-clock duration of the model invocation that
produced the persisted message content**, measured from just before the
provider stream opens until it is fully consumed.

What counts:

- Assistant generation steps, in top-level and sub-agent chats
(sub-agents are ordinary chats on the same generation path).
- Compaction summarization calls, persisted on the compaction assistant
message (**new**).
- Interrupted attempts: the message-part episode's lifetime is persisted
on the partial assistant message committed by `FinishInterruption`, so
partial generation time survives interruption (**new**; measured via a
new `Buffer.EpisodeDuration`, which works even though the generation
goroutine and the interrupt task are different tasks).

What deliberately does not count (each is documented in code and docs):

- **Local tool execution.** Tool wall time includes idle waits, most
importantly `wait_agent` polling a sub-agent chat that already bills its
own model invocations; billing the batch would double count, and
excluding one tool from a concurrent batch's wall time is ill-defined.
Pre-refactor instrumentation did include tool time; this makes the
exclusion an explicit product definition instead of a silent regression.
- **Failed model calls whose output is discarded** (retried attempts,
terminal errors, content-filter refusals). They persist no content, so
they bill nothing; billing errs toward undercounting. Notably a
stream-silence timeout can burn 10 idle minutes before a retry, which
should not be billable "active generation". If product later wants
failed attempts billed, that needs a place to persist runtime on error
turns (`FinishError` inserts no rows today) and is a deliberate
follow-up, not instrumentation drift.
- **Ancillary calls that produce no chat messages** (title generation,
advisor, turn summaries) and all idle/parked time (`requires_action`,
queueing).

The definition is documented as `COMMENT ON COLUMN
chat_messages.runtime_ms` (migration 000551, surfacing as a Go doc
comment on `ChatMessage.RuntimeMs`), on
`chatloop.PersistedStep.Runtime`, in the chatd architecture doc, and in
the Spend Management docs page.

## Index for the hourly scan

None needed: `GetTotalChatMessageRuntimeMsInRange` (#27312) filters an
hour-wide `created_at` range, which the existing
`idx_chat_messages_created_at` b-tree already serves; the residual
`runtime_ms IS NOT NULL` filter applies to one hour of rows. A partial
index would add permanent write amplification for a query that runs once
an hour.

> [!NOTE]
> Migration 000551 is also claimed by #27312; whichever merges second
renumbers via `fix_migration_numbers.sh`.

## Tests

- End-to-end: the existing full-server generation test now asserts
`RuntimeMs.Valid` on the committed assistant row (it previously read
`.Int64` without checking `.Valid`, so it passed on NULL).
- Interrupted turn: full task-level test (real DB, mock clock) asserting
the partial assistant message persists the attempt's runtime.
- Errored stream: asserts a failed invocation yields no step and no
runtime.
- Tool-using turn: asserts runtime lands on the assistant row only and
tool rows stay NULL.
- Compaction: asserts the summarization call duration is recorded and
lands on the compaction assistant message only.
- `messagepartbuffer.EpisodeDuration` unit coverage.

Blocks: CODAGT-843 (B3), CODAGT-838 (D8).

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Hugo Dutka <hugo@coder.com>
2026-08-06 16:09:38 +07:00

923 lines
29 KiB
Go

package chatd
import (
"cmp"
"context"
"database/sql"
"encoding/json"
"slices"
"strings"
"time"
"charm.land/fantasy"
"github.com/google/uuid"
"github.com/sqlc-dev/pqtype"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/x/chatd/chatloop"
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
"github.com/coder/coder/v2/coderd/x/chatd/chatstate"
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
"github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer"
"github.com/coder/coder/v2/codersdk"
)
const interruptedToolResultErrorMessage = "tool call was interrupted before it produced a result"
type buildCommitStepMessagesInput struct {
modelConfigID uuid.UUID
step stepData
toolNameToConfigID map[string]uuid.UUID
logger slog.Logger
contentVersion int16
hookRewrittenToolCalls map[string]json.RawMessage
}
type stepMessagesForCommit struct {
Messages []chatstate.Message
VisibleIndexes []int
// ConsumeCompactionRequest clears the manual compaction marker
// atomically with the commit. Set on compaction commits.
ConsumeCompactionRequest bool
}
func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesForCommit, error) {
contentVersion := input.contentVersion
if contentVersion == 0 {
contentVersion = chatprompt.CurrentContentVersion
}
assistantBlocks, toolResults := splitStepContent(input.step.Content)
assistantParts := buildAssistantParts(input.logger, assistantBlocks, toolResults, input.step, input.toolNameToConfigID, input.hookRewrittenToolCalls)
messages := make([]chatstate.Message, 0, 1+len(toolResults))
if len(assistantParts) > 0 {
assistantContent, err := chatprompt.MarshalParts(assistantParts)
if err != nil {
return stepMessagesForCommit{}, xerrors.Errorf("marshal assistant content: %w", err)
}
messages = append(messages, assistantMessage(input.modelConfigID, contentVersion, assistantContent, input.step))
}
for _, toolResult := range toolResults {
part := chatprompt.PartFromContentWithLogger(context.Background(), input.logger, toolResult)
applyToolMetadata(&part, input.toolNameToConfigID)
if part.ToolCallID != "" && input.step.ToolResultCreatedAt != nil {
if ts, ok := input.step.ToolResultCreatedAt[part.ToolCallID]; ok {
part.CreatedAt = &ts
}
}
content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{part})
if err != nil {
return stepMessagesForCommit{}, xerrors.Errorf("marshal tool result: %w", err)
}
messages = append(messages, baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content))
}
return stepMessagesForCommit{
Messages: messages,
VisibleIndexes: visibleMessageIndexes(messages),
}, nil
}
func splitStepContent(content []fantasy.Content) ([]fantasy.Content, []fantasy.ToolResultContent) {
assistantBlocks := make([]fantasy.Content, 0, len(content))
toolResults := make([]fantasy.ToolResultContent, 0)
for _, block := range content {
if tr, ok := asToolResultContent(block); ok && !tr.ProviderExecuted {
toolResults = append(toolResults, tr)
continue
}
assistantBlocks = append(assistantBlocks, block)
}
return assistantBlocks, toolResults
}
func asToolResultContent(block fantasy.Content) (fantasy.ToolResultContent, bool) {
if tr, ok := fantasy.AsContentType[fantasy.ToolResultContent](block); ok {
return tr, true
}
if tr, ok := fantasy.AsContentType[*fantasy.ToolResultContent](block); ok && tr != nil {
return *tr, true
}
return fantasy.ToolResultContent{}, false
}
func buildAssistantParts(
logger slog.Logger,
assistantBlocks []fantasy.Content,
toolResults []fantasy.ToolResultContent,
step stepData,
toolNameToConfigID map[string]uuid.UUID,
hookRewrittenToolCalls map[string]json.RawMessage,
) []codersdk.ChatMessagePart {
parts := make([]codersdk.ChatMessagePart, 0, len(assistantBlocks)+len(toolResults))
reasoningIdx := 0
for _, block := range assistantBlocks {
part := chatprompt.PartFromContentWithLogger(context.Background(), logger, block)
applyToolMetadata(&part, toolNameToConfigID)
switch part.Type {
case codersdk.ChatMessagePartTypeToolCall:
if part.ToolCallID != "" && step.ToolCallCreatedAt != nil {
if ts, ok := step.ToolCallCreatedAt[part.ToolCallID]; ok {
part.CreatedAt = &ts
}
}
// Hooks never see provider-executed calls, so such a call must not
// inherit attribution from an ordinary call that reused its ID.
if part.ToolCallID != "" && !part.ProviderExecuted {
_, part.HookRewritten = hookRewrittenToolCalls[part.ToolCallID]
}
case codersdk.ChatMessagePartTypeToolResult:
if part.ToolCallID != "" && step.ToolResultCreatedAt != nil {
if ts, ok := step.ToolResultCreatedAt[part.ToolCallID]; ok {
part.CreatedAt = &ts
}
}
case codersdk.ChatMessagePartTypeReasoning:
if reasoningIdx < len(step.ReasoningStartedAt) {
if ts := step.ReasoningStartedAt[reasoningIdx]; !ts.IsZero() {
part.CreatedAt = &ts
}
}
if reasoningIdx < len(step.ReasoningCompletedAt) {
if ts := step.ReasoningCompletedAt[reasoningIdx]; !ts.IsZero() {
part.CompletedAt = &ts
}
}
reasoningIdx++
}
if part.Type != "" {
parts = append(parts, part)
}
}
for _, tr := range toolResults {
attachments, err := chattool.AttachmentsFromMetadata(tr.ClientMetadata)
if err != nil {
logger.Warn(context.Background(), "skipping malformed tool attachment metadata",
slog.F("tool_name", tr.ToolName),
slog.F("tool_call_id", tr.ToolCallID),
slog.Error(err),
)
continue
}
for _, attachment := range attachments {
parts = append(parts, codersdk.ChatMessageFile(attachment.FileID, attachment.MediaType, attachment.Name))
}
}
return parts
}
func applyToolMetadata(part *codersdk.ChatMessagePart, toolNameToConfigID map[string]uuid.UUID) {
if part.ToolName == "" || len(toolNameToConfigID) == 0 {
return
}
if configID, ok := toolNameToConfigID[part.ToolName]; ok {
part.MCPServerConfigID = uuid.NullUUID{UUID: configID, Valid: true}
}
}
func assistantMessage(
modelConfigID uuid.UUID,
contentVersion int16,
content pqtype.NullRawMessage,
step stepData,
) chatstate.Message {
msg := baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, modelConfigID, contentVersion, content)
if step.Usage != (fantasy.Usage{}) {
msg.InputTokens = nullInt64IfNonZero(step.Usage.InputTokens)
msg.OutputTokens = nullInt64IfNonZero(step.Usage.OutputTokens)
msg.TotalTokens = nullInt64IfNonZero(step.Usage.TotalTokens)
msg.ReasoningTokens = nullInt64IfNonZero(step.Usage.ReasoningTokens)
msg.CacheCreationTokens = nullInt64IfNonZero(step.Usage.CacheCreationTokens)
msg.CacheReadTokens = nullInt64IfNonZero(step.Usage.CacheReadTokens)
}
msg.ContextLimit = step.ContextLimit
// InsertChatMessages maps a zero runtime to NULL, so a model
// invocation shorter than a millisecond persists the same way an
// unmeasured one does.
msg.RuntimeMs = nullInt64IfNonZero(step.Runtime.Milliseconds())
return msg
}
func baseMessage(
role database.ChatMessageRole,
visibility database.ChatMessageVisibility,
modelConfigID uuid.UUID,
contentVersion int16,
content pqtype.NullRawMessage,
) chatstate.Message {
return chatstate.Message{
Role: role,
Content: content,
Visibility: visibility,
ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil},
ContentVersion: contentVersion,
}
}
func nullInt64IfNonZero(value int64) sql.NullInt64 {
if value == 0 {
return sql.NullInt64{}
}
return sql.NullInt64{Int64: value, Valid: true}
}
func visibleMessageIndexes(messages []chatstate.Message) []int {
indexes := make([]int, 0, len(messages))
for i, msg := range messages {
if msg.Visibility == database.ChatMessageVisibilityBoth || msg.Visibility == database.ChatMessageVisibilityUser {
indexes = append(indexes, i)
}
}
return indexes
}
func textFromParts(parts []codersdk.ChatMessagePart) string {
var builder strings.Builder
for _, part := range parts {
if part.Type == codersdk.ChatMessagePartTypeText {
_, _ = builder.WriteString(part.Text)
}
}
return builder.String()
}
type buildCompactionMessagesInput struct {
modelConfigID uuid.UUID
toolCallID string
toolName string
compaction compactionOutcome
contentVersion int16
}
type compactionMessagesForCommit struct {
Messages []chatstate.Message
HiddenCount int
}
func buildCompactionMessages(input buildCompactionMessagesInput) (compactionMessagesForCommit, error) {
contentVersion := input.contentVersion
if contentVersion == 0 {
contentVersion = chatprompt.CurrentContentVersion
}
toolName := input.toolName
if toolName == "" {
toolName = "chat_summarized"
}
systemContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(input.compaction.SystemSummary)})
if err != nil {
return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction system summary: %w", err)
}
source := input.compaction.Source
if source == "" {
source = chatloop.CompactionSourceAutomatic
}
args, err := json.Marshal(map[string]any{
"source": source,
"threshold_percent": input.compaction.ThresholdPercent,
})
if err != nil {
return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction args: %w", err)
}
assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{
codersdk.ChatMessageToolCall(input.toolCallID, toolName, args),
})
if err != nil {
return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction tool call: %w", err)
}
summaryResult, err := json.Marshal(map[string]any{
"summary": input.compaction.SummaryReport,
"source": source,
"threshold_percent": input.compaction.ThresholdPercent,
"usage_percent": input.compaction.UsagePercent,
"context_tokens": input.compaction.ContextTokens,
"context_limit_tokens": input.compaction.ContextLimit,
})
if err != nil {
return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction result: %w", err)
}
toolContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{
codersdk.ChatMessageToolResult(input.toolCallID, toolName, summaryResult, false, false),
})
if err != nil {
return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction tool result: %w", err)
}
assistantMsg := baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, input.modelConfigID, contentVersion, assistantContent)
assistantMsg.RuntimeMs = nullInt64IfNonZero(input.compaction.Runtime.Milliseconds())
messages := []chatstate.Message{
{
Role: database.ChatMessageRoleUser,
Content: systemContent,
Visibility: database.ChatMessageVisibilityModel,
ModelConfigID: uuid.NullUUID{UUID: input.modelConfigID, Valid: input.modelConfigID != uuid.Nil},
ContentVersion: contentVersion,
},
assistantMsg,
baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, toolContent),
}
for i := range messages {
messages[i].Compressed = true
}
return compactionMessagesForCommit{Messages: messages, HiddenCount: 1}, nil
}
// Hook model-context messages use the user role but must not reset
// per-turn guards.
func lastUserPromptIndex(messages []database.ChatMessage) int {
index := -1
for i, msg := range messages {
if msg.Deleted || msg.Compressed {
continue
}
if msg.Role == database.ChatMessageRoleUser && msg.Visibility != database.ChatMessageVisibilityModel {
index = i
}
}
return index
}
func currentTurnStartIndex(messages []database.ChatMessage) int {
return lastUserPromptIndex(messages) + 1
}
func currentTurnStepCount(messages []database.ChatMessage) int {
count := 0
for _, msg := range messages[currentTurnStartIndex(messages):] {
if msg.Deleted || msg.Compressed {
continue
}
if msg.Role == database.ChatMessageRoleAssistant {
count++
}
}
return count
}
type compactionRequirement int
const (
compactionRequirementNotNeeded compactionRequirement = iota
compactionRequirementNeeded
)
func compactionStatusFromHistory(
messages []database.ChatMessage,
requirement compactionRequirement,
thresholdPercent int32,
contextLimit int64,
) compactionStatus {
boundaryIndex := latestCompactionBoundaryIndex(messages)
if requirement == compactionRequirementNeeded {
if boundaryIndex == -1 {
return compactionStatusNeeded
}
// The first assistant response after the previously compacted summary.
// Messages with role ChatMessageRoleAssistant carry context usage.
// Looking at ChatMessageRoleAssistant is enough - ChatMessageRoleTool
// does not carry context usage, and is always preceded by an assistant
// message.
if assistant, ok := firstUncompressedAssistantAfter(messages, boundaryIndex); ok &&
postCompactionAssistantOverLimit(assistant, thresholdPercent, contextLimit) {
return compactionStatusStillOverLimit
}
if hasUncompressedMessageAfter(messages, boundaryIndex) {
return compactionStatusNeeded
}
return compactionStatusAfterCompaction
}
if boundaryIndex != -1 && !hasUncompressedMessageAfter(messages, boundaryIndex) {
return compactionStatusAfterCompaction
}
return compactionStatusNotNeeded
}
func latestCompactionBoundaryIndex(messages []database.ChatMessage) int {
for i := len(messages) - 1; i >= 0; i-- {
if isCompactionBoundaryMessage(messages[i]) {
return i
}
}
return -1
}
func isCompactionBoundaryMessage(msg database.ChatMessage) bool {
if msg.Deleted || !msg.Compressed {
return false
}
parts, err := chatprompt.ParseContent(msg)
if err != nil {
return false
}
for _, part := range parts {
if part.ToolName == "chat_summarized" &&
(part.Type == codersdk.ChatMessagePartTypeToolCall || part.Type == codersdk.ChatMessagePartTypeToolResult) {
return true
}
}
return false
}
func firstUncompressedAssistantAfter(messages []database.ChatMessage, index int) (database.ChatMessage, bool) {
for i := index + 1; i < len(messages); i++ {
msg := messages[i]
if msg.Deleted || msg.Compressed {
continue
}
if msg.Role == database.ChatMessageRoleAssistant {
return msg, true
}
}
return database.ChatMessage{}, false
}
func hasUncompressedMessageAfter(messages []database.ChatMessage, index int) bool {
for i := index + 1; i < len(messages); i++ {
msg := messages[i]
if !msg.Deleted && !msg.Compressed {
return true
}
}
return false
}
func postCompactionAssistantOverLimit(msg database.ChatMessage, thresholdPercent int32, contextLimit int64) bool {
return shouldCompactPromptUsage(usageFromMessage(msg), contextLimit, thresholdPercent)
}
func usageFromMessage(msg database.ChatMessage) fantasy.Usage {
var usage fantasy.Usage
if msg.InputTokens.Valid {
usage.InputTokens = msg.InputTokens.Int64
}
if msg.OutputTokens.Valid {
usage.OutputTokens = msg.OutputTokens.Int64
}
if msg.TotalTokens.Valid {
usage.TotalTokens = msg.TotalTokens.Int64
}
if msg.ReasoningTokens.Valid {
usage.ReasoningTokens = msg.ReasoningTokens.Int64
}
if msg.CacheCreationTokens.Valid {
usage.CacheCreationTokens = msg.CacheCreationTokens.Int64
}
if msg.CacheReadTokens.Valid {
usage.CacheReadTokens = msg.CacheReadTokens.Int64
}
return usage
}
func historyHasStopAfterToolResult(messages []database.ChatMessage, stopAfterTools map[string]struct{}) (bool, error) {
if len(stopAfterTools) == 0 {
return false, nil
}
for _, msg := range messages[currentTurnStartIndex(messages):] {
if msg.Deleted || msg.Compressed || msg.Role != database.ChatMessageRoleTool {
continue
}
parts, err := chatprompt.ParseContent(msg)
if err != nil {
return false, xerrors.Errorf("parse tool message: %w", err)
}
for _, part := range parts {
if part.Type != codersdk.ChatMessagePartTypeToolResult || part.IsError {
continue
}
if _, ok := stopAfterTools[part.ToolName]; ok {
return true, nil
}
}
}
return false, nil
}
func currentHistoryComplete(messages []database.ChatMessage) (bool, error) {
idx := lastMessageIndex(messages, func(database.ChatMessage) bool { return true })
if idx == -1 || messages[idx].Role != database.ChatMessageRoleAssistant {
return false, nil
}
parts, err := chatprompt.ParseContent(messages[idx])
if err != nil {
return false, xerrors.Errorf("parse latest assistant message: %w", err)
}
for _, part := range parts {
if part.Type == codersdk.ChatMessagePartTypeToolCall && !part.ProviderExecuted {
return false, nil
}
}
return true, nil
}
func lastMessageIndex(messages []database.ChatMessage, accept func(database.ChatMessage) bool) int {
for i := len(messages) - 1; i >= 0; i-- {
if messages[i].Deleted || messages[i].Compressed {
continue
}
if accept(messages[i]) {
return i
}
}
return -1
}
func handledToolCallIDs(messages []database.ChatMessage) (map[string]bool, error) {
handled := make(map[string]bool)
for _, msg := range messages {
if msg.Deleted || msg.Compressed || msg.Role != database.ChatMessageRoleTool {
continue
}
parts, err := chatprompt.ParseContent(msg)
if err != nil {
return nil, xerrors.Errorf("parse tool message: %w", err)
}
for _, part := range parts {
if part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolCallID != "" {
handled[part.ToolCallID] = true
}
}
}
return handled, nil
}
type bufferedPartsToPartialMessagesInput struct {
parts []messagepartbuffer.Part
modelConfigID uuid.UUID
contentVersion int16
logger slog.Logger
interruptedAt time.Time
// attemptRuntime is the interrupted attempt's billable model
// invocation window: the span from the provider stream opening to
// the interrupt closing its buffer episode. It is persisted as
// runtime_ms on the first partial assistant message when the
// attempt streamed model-generated assistant content.
attemptRuntime time.Duration
}
type partialToolCall struct {
part codersdk.ChatMessagePart
index int
argsDelta strings.Builder
valid bool
durable bool
}
type partialToolResult struct {
part codersdk.ChatMessagePart
resultDelta strings.Builder
completed bool
}
func bufferedPartsToPartialMessages(input bufferedPartsToPartialMessagesInput) ([]chatstate.Message, error) {
contentVersion := input.contentVersion
if contentVersion == 0 {
contentVersion = chatprompt.CurrentContentVersion
}
parts := slices.Clone(input.parts)
slices.SortFunc(parts, func(a, b messagepartbuffer.Part) int {
return cmp.Compare(a.Seq, b.Seq)
})
state := partialMessageConversionState{
input: input,
contentVersion: contentVersion,
toolCalls: make(map[string]*partialToolCall),
toolResults: make(map[string]*partialToolResult),
answered: make(map[string]bool),
}
for _, buffered := range parts {
if err := state.consume(buffered); err != nil {
return nil, err
}
}
if err := state.finalizeToolCallPlaceholders(); err != nil {
return nil, err
}
if err := state.flushAssistant(); err != nil {
return nil, err
}
if err := state.flushAccumulatedToolResults(); err != nil {
return nil, err
}
if err := state.appendSyntheticInterruptionResults(); err != nil {
return nil, err
}
if input.attemptRuntime > 0 && state.modelStreamedAssistant {
// Usage reporting sums runtime_ms across rows, so placing the
// whole span on the first assistant message is sufficient.
for i := range state.messages {
if state.messages[i].Role != database.ChatMessageRoleAssistant {
continue
}
state.messages[i].RuntimeMs = nullInt64IfNonZero(input.attemptRuntime.Milliseconds())
break
}
}
return state.messages, nil
}
type partialMessageConversionState struct {
input bufferedPartsToPartialMessagesInput
contentVersion int16
messages []chatstate.Message
assistantParts []codersdk.ChatMessagePart
toolCalls map[string]*partialToolCall
toolCallOrder []string
toolResults map[string]*partialToolResult
toolResultOrder []string
answered map[string]bool
// modelStreamedAssistant records whether any assistant part came
// from the model stream itself (text, reasoning, tool calls,
// sources). Tool execution also publishes assistant-role file
// parts for attachments; those alone must not attract the
// attempt's runtime, because tool batches are not billable. The
// buffer episode only carries a runtime when a provider stream
// was opened, so this is a second gate rather than the only one.
modelStreamedAssistant bool
}
func (s *partialMessageConversionState) consume(buffered messagepartbuffer.Part) error {
switch buffered.Role {
case codersdk.ChatMessageRoleAssistant:
s.consumeAssistantPart(buffered)
case codersdk.ChatMessageRoleTool:
return s.consumeToolPart(buffered)
default:
s.logSkippedPart(buffered, "unsupported buffered part role")
}
return nil
}
func (s *partialMessageConversionState) consumeAssistantPart(buffered messagepartbuffer.Part) {
part := buffered.MessagePart
if part.Type == "" {
s.logSkippedPart(buffered, "empty buffered assistant part type")
return
}
if part.Type != codersdk.ChatMessagePartTypeFile {
s.modelStreamedAssistant = true
}
if part.Type != codersdk.ChatMessagePartTypeToolCall {
if part.Type == codersdk.ChatMessagePartTypeReasoning &&
!s.input.interruptedAt.IsZero() {
interruptedAt := s.input.interruptedAt
if part.CreatedAt == nil {
part.CreatedAt = &interruptedAt
}
if part.CompletedAt == nil {
part.CompletedAt = &interruptedAt
}
}
s.assistantParts = append(s.assistantParts, part)
return
}
if part.ToolCallID == "" {
s.logSkippedPart(buffered, "tool call part missing tool call ID")
return
}
call := s.toolCall(part.ToolCallID)
call.part.Type = codersdk.ChatMessagePartTypeToolCall
call.part.ToolCallID = part.ToolCallID
if part.ToolName != "" {
call.part.ToolName = part.ToolName
}
if part.MCPServerConfigID.Valid {
call.part.MCPServerConfigID = part.MCPServerConfigID
}
if part.CreatedAt != nil {
call.part.CreatedAt = part.CreatedAt
}
call.part.ProviderExecuted = call.part.ProviderExecuted || part.ProviderExecuted
if part.ArgsDelta != "" {
if call.durable {
s.logSkippedPart(buffered, "tool call args delta arrived after full tool call")
return
}
_, _ = call.argsDelta.WriteString(part.ArgsDelta)
return
}
durable := part
durable.ArgsDelta = ""
if len(durable.Args) > 0 && !json.Valid(durable.Args) {
call.valid = false
s.assistantParts[call.index] = codersdk.ChatMessagePart{}
s.logSkippedPart(buffered, "tool call part has invalid durable args")
return
}
if call.durable {
s.logSkippedPart(buffered, "duplicate durable tool call part")
}
call.part = durable
call.valid = true
call.durable = true
s.assistantParts[call.index] = durable
}
func (s *partialMessageConversionState) consumeToolPart(buffered messagepartbuffer.Part) error {
part := buffered.MessagePart
if part.Type != codersdk.ChatMessagePartTypeToolResult {
s.logSkippedPart(buffered, "non tool-result part with tool role")
return nil
}
if part.ToolCallID == "" {
s.logSkippedPart(buffered, "tool result part missing tool call ID")
return nil
}
if part.ResultReset {
result := s.toolResult(part.ToolCallID)
result.part.ToolCallID = part.ToolCallID
result.part.ToolName = part.ToolName
result.resultDelta.Reset()
s.logSkippedPart(buffered, "streaming tool result reset is not durable")
return nil
}
if part.ResultDelta != "" {
result := s.toolResult(part.ToolCallID)
result.part.ToolCallID = part.ToolCallID
if part.ToolName != "" {
result.part.ToolName = part.ToolName
}
if part.MCPServerConfigID.Valid {
result.part.MCPServerConfigID = part.MCPServerConfigID
}
if part.CreatedAt != nil {
result.part.CreatedAt = part.CreatedAt
}
result.part.ProviderExecuted = result.part.ProviderExecuted || part.ProviderExecuted
_, _ = result.resultDelta.WriteString(part.ResultDelta)
return nil
}
if err := s.finalizeToolCallPlaceholders(); err != nil {
return err
}
if !s.toolCallDurable(part.ToolCallID) {
s.logSkippedPart(buffered, "tool result has no matching durable tool call")
return nil
}
if len(part.Result) == 0 || !json.Valid(part.Result) {
s.logSkippedPart(buffered, "tool result part has invalid durable result")
return nil
}
if s.answered[part.ToolCallID] {
s.logSkippedPart(buffered, "duplicate durable tool result part")
return nil
}
part.ResultDelta = ""
part.ResultReset = false
if err := s.flushAssistant(); err != nil {
return err
}
if err := s.appendToolResult(part); err != nil {
return err
}
s.answered[part.ToolCallID] = true
return nil
}
func (s *partialMessageConversionState) toolCall(id string) *partialToolCall {
call := s.toolCalls[id]
if call != nil {
return call
}
call = &partialToolCall{index: len(s.assistantParts), valid: true}
s.toolCalls[id] = call
s.toolCallOrder = append(s.toolCallOrder, id)
s.assistantParts = append(s.assistantParts, codersdk.ChatMessagePart{})
return call
}
func (s *partialMessageConversionState) toolResult(id string) *partialToolResult {
result := s.toolResults[id]
if result != nil {
return result
}
result = &partialToolResult{}
s.toolResults[id] = result
s.toolResultOrder = append(s.toolResultOrder, id)
return result
}
func (s *partialMessageConversionState) finalizeToolCallPlaceholders() error {
for _, id := range s.toolCallOrder {
call := s.toolCalls[id]
if call == nil || call.durable || !call.valid {
continue
}
args := json.RawMessage(call.argsDelta.String())
if len(args) == 0 || !json.Valid(args) {
s.assistantParts[call.index] = codersdk.ChatMessagePart{}
call.valid = false
s.logSkippedPart(messagepartbuffer.Part{
Role: codersdk.ChatMessageRoleAssistant,
MessagePart: call.part,
}, "tool call args delta did not form durable JSON")
continue
}
call.part.Args = args
call.part.ArgsDelta = ""
call.durable = true
s.assistantParts[call.index] = call.part
}
return nil
}
func (s *partialMessageConversionState) flushAssistant() error {
if len(s.assistantParts) == 0 {
return nil
}
durable := make([]codersdk.ChatMessagePart, 0, len(s.assistantParts))
for _, part := range s.assistantParts {
if part.Type == "" {
continue
}
part.ArgsDelta = ""
part.ResultDelta = ""
part.ResultReset = false
durable = append(durable, part)
}
s.assistantParts = nil
if len(durable) == 0 {
return nil
}
content, err := chatprompt.MarshalParts(durable)
if err != nil {
return xerrors.Errorf("marshal partial assistant: %w", err)
}
s.messages = append(s.messages, baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, s.input.modelConfigID, s.contentVersion, content))
return nil
}
func (s *partialMessageConversionState) flushAccumulatedToolResults() error {
for _, id := range s.toolResultOrder {
if s.answered[id] {
continue
}
result := s.toolResults[id]
if result == nil || result.completed {
continue
}
if result.resultDelta.Len() == 0 {
continue
}
s.logSkippedPart(messagepartbuffer.Part{Role: codersdk.ChatMessageRoleTool, MessagePart: result.part}, "streaming tool result delta is not durable")
}
return nil
}
func (s *partialMessageConversionState) appendToolResult(part codersdk.ChatMessagePart) error {
content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{part})
if err != nil {
return xerrors.Errorf("marshal partial tool result: %w", err)
}
s.messages = append(s.messages, baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, s.input.modelConfigID, s.contentVersion, content))
return nil
}
func (s *partialMessageConversionState) appendSyntheticInterruptionResults() error {
for _, id := range s.toolCallOrder {
if s.answered[id] {
continue
}
call := s.toolCalls[id]
if call == nil || !call.valid || !call.durable || call.part.ProviderExecuted {
continue
}
result, err := json.Marshal(map[string]string{"error": interruptedToolResultErrorMessage})
if err != nil {
return xerrors.Errorf("marshal synthetic interruption result: %w", err)
}
part := codersdk.ChatMessageToolResult(call.part.ToolCallID, call.part.ToolName, result, true, false)
part.MCPServerConfigID = call.part.MCPServerConfigID
if !s.input.interruptedAt.IsZero() {
part.CreatedAt = &s.input.interruptedAt
}
if err := s.appendToolResult(part); err != nil {
return xerrors.Errorf("marshal synthetic interruption message: %w", err)
}
s.answered[id] = true
}
return nil
}
func (s *partialMessageConversionState) toolCallDurable(id string) bool {
call := s.toolCalls[id]
return call != nil && call.valid && call.durable
}
func (s *partialMessageConversionState) logSkippedPart(buffered messagepartbuffer.Part, reason string) {
s.input.logger.Warn(context.Background(), "skipping buffered chat message part",
slog.F("reason", reason),
slog.F("role", buffered.Role),
slog.F("part_type", buffered.MessagePart.Type),
slog.F("tool_call_id", buffered.MessagePart.ToolCallID),
slog.F("tool_name", buffered.MessagePart.ToolName),
)
}