opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-21 20:51:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: track step runtime_ms on chat messages (#23219)

Browse Source 
## Summary

Adds a `runtime_ms` column to `chat_messages` that records the
wall-clock duration (in milliseconds) of each LLM step. This covers LLM
streaming, tool execution, and retries — the full time the agent is
"alive" for a step.

This is the foundation for billing by agent alive time. The column
follows the same pattern as `total_cost_micros`: stored per assistant
message, aggregatable with `SUM()` over time periods by user.

## Changes

- **Migration**: adds nullable `runtime_ms bigint` to `chat_messages`.
- **chatloop**: adds `Runtime time.Duration` field to `PersistedStep`,
measures `time.Since(stepStart)` at the beginning of each step (covering
stream + tool execution + retries).
- **chatd**: passes `step.Runtime.Milliseconds()` to the assistant
message `InsertChatMessage` call; all other message types (system, user,
tool) get `NULL`.
- **Tests**: adds `runtime > 0` assertion in chatloop tests.

## Billing query pattern

Once ready, aggregation mirrors the existing cost queries:

```sql
SELECT COALESCE(SUM(cm.runtime_ms), 0)::bigint AS total_runtime_ms
FROM chat_messages cm
JOIN chats c ON c.id = cm.chat_id
WHERE c.owner_id = @user_id
  AND cm.created_at >= @start_time
  AND cm.created_at < @end_time
  AND cm.runtime_ms IS NOT NULL;
```
This commit is contained in:
Kyle Carberry
2026-03-18 10:57:35 -04:00
committed by GitHub
parent 3bcb7de7c0
commit 4dd8531f37
10 changed files with 59 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/chatd/chatd.go
+13
View File
@@ -486,6 +486,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C
ContextLimit: sql.NullInt64{},
Compressed: sql.NullBool{},
TotalCostMicros: sql.NullInt64{},
RuntimeMs: sql.NullInt64{},
})
if err != nil {
return xerrors.Errorf("insert system message: %w", err)
@@ -524,6 +525,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C
ContextLimit: sql.NullInt64{},
Compressed: sql.NullBool{},
TotalCostMicros: sql.NullInt64{},
RuntimeMs: sql.NullInt64{},
})
if err != nil {
return xerrors.Errorf("insert workspace awareness message: %w", err)
@@ -552,6 +554,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
TotalCostMicros: sql.NullInt64{},
RuntimeMs: sql.NullInt64{},
Compressed: sql.NullBool{},
})
if err != nil {
@@ -1170,6 +1173,7 @@ func insertUserMessageAndSetPending(
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
TotalCostMicros: sql.NullInt64{},
RuntimeMs: sql.NullInt64{},
Compressed: sql.NullBool{},
})
if err != nil {
@@ -2147,6 +2151,7 @@ func (p *Server) tryAutoPromoteQueuedMessage(
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
TotalCostMicros: sql.NullInt64{},
RuntimeMs: sql.NullInt64{},
Compressed: sql.NullBool{},
})
if err != nil {
@@ -2709,6 +2714,10 @@ func (p *Server) runChat(
ContextLimit: step.ContextLimit,
Compressed: sql.NullBool{},
TotalCostMicros: usageNullInt64Ptr(totalCostMicros),
RuntimeMs: sql.NullInt64{
Int64: step.Runtime.Milliseconds(),
Valid: step.Runtime > 0,
},
})
if insertErr != nil {
return xerrors.Errorf("insert assistant message: %w", insertErr)
@@ -2733,6 +2742,7 @@ func (p *Server) runChat(
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
TotalCostMicros: sql.NullInt64{},
RuntimeMs: sql.NullInt64{},
Compressed: sql.NullBool{},
})
if insertErr != nil {
@@ -3117,6 +3127,7 @@ func (p *Server) persistChatContextSummary(
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
TotalCostMicros: sql.NullInt64{},
RuntimeMs: sql.NullInt64{},
})
if txErr != nil {
return xerrors.Errorf("insert hidden summary message: %w", txErr)
@@ -3142,6 +3153,7 @@ func (p *Server) persistChatContextSummary(
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
TotalCostMicros: sql.NullInt64{},
RuntimeMs: sql.NullInt64{},
})
if txErr != nil {
return xerrors.Errorf("insert summary tool call message: %w", txErr)
@@ -3168,6 +3180,7 @@ func (p *Server) persistChatContextSummary(
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
TotalCostMicros: sql.NullInt64{},
RuntimeMs: sql.NullInt64{},
})
if txErr != nil {
return xerrors.Errorf("insert summary tool result message: %w", txErr)
coderd/chatd/chatd_test.go
+7
View File
@@ -726,6 +726,7 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) {
ContextLimit: sql.NullInt64{},
Compressed: sql.NullBool{},
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
RuntimeMs: sql.NullInt64{},
})
require.NoError(t, err)
@@ -824,6 +825,7 @@ func TestPromoteQueuedAllowsAlreadyQueuedMessageWhenUsageLimitReached(t *testing
ContextLimit: sql.NullInt64{},
Compressed: sql.NullBool{},
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
RuntimeMs: sql.NullInt64{},
})
require.NoError(t, err)
@@ -1012,6 +1014,7 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
ContextLimit: sql.NullInt64{},
Compressed: sql.NullBool{},
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
RuntimeMs: sql.NullInt64{},
})
require.NoError(t, err)
@@ -1107,6 +1110,7 @@ func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) {
ContextLimit: sql.NullInt64{},
Compressed: sql.NullBool{},
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
RuntimeMs: sql.NullInt64{},
})
require.NoError(t, err)
@@ -1196,6 +1200,7 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) {
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
Compressed: sql.NullBool{},
RuntimeMs: sql.NullInt64{},
})
require.NoError(t, err)
@@ -1544,6 +1549,7 @@ func TestSubscribeAfterMessageID(t *testing.T) {
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
Compressed: sql.NullBool{},
RuntimeMs: sql.NullInt64{},
})
require.NoError(t, err)
@@ -1567,6 +1573,7 @@ func TestSubscribeAfterMessageID(t *testing.T) {
CacheReadTokens: sql.NullInt64{},
ContextLimit: sql.NullInt64{},
Compressed: sql.NullBool{},
RuntimeMs: sql.NullInt64{},
})
require.NoError(t, err)
coderd/chatd/chatloop/chatloop.go
+7
View File
@@ -42,6 +42,11 @@ type PersistedStep struct {
Content []fantasy.Content
Usage fantasy.Usage
ContextLimit sql.NullInt64
// Runtime is the wall-clock duration of this step,
// covering LLM streaming, tool execution, and retries.
// Zero indicates the duration was not measured (e.g.
// interrupted steps).
Runtime time.Duration
}
// RunOptions configures a single streaming chat loop run.
@@ -260,6 +265,7 @@ func Run(ctx context.Context, opts RunOptions) error {
for step := 0; totalSteps < opts.MaxSteps; step++ {
totalSteps++
stepStart := time.Now()
// Copy messages so that provider-specific caching
// mutations don't leak back to the caller's slice.
// copy copies Message structs by value, so field
@@ -365,6 +371,7 @@ func Run(ctx context.Context, opts RunOptions) error {
Content: result.content,
Usage: result.usage,
ContextLimit: contextLimit,
Runtime: time.Since(stepStart),
}); err != nil {
if errors.Is(err, ErrInterrupted) {
persistInterruptedStep(ctx, opts, &result)
coderd/chatd/chatloop/chatloop_test.go
+3
View File
@@ -7,6 +7,7 @@ import (
"strings"
"sync"
"testing"
"time"
"charm.land/fantasy"
fantasyanthropic "charm.land/fantasy/providers/anthropic"
@@ -64,6 +65,8 @@ func TestRun_ActiveToolsPrepareBehavior(t *testing.T) {
require.Equal(t, 1, persistStepCalls)
require.True(t, persistedStep.ContextLimit.Valid)
require.Equal(t, int64(4096), persistedStep.ContextLimit.Int64)
require.Greater(t, persistedStep.Runtime, time.Duration(0),
"step runtime should be positive")
require.NotEmpty(t, capturedCall.Prompt)
require.False(t, containsPromptSentinel(capturedCall.Prompt))
coderd/database/dump.sql
Generated
+2 -1
View File
@@ -1289,7 +1289,8 @@ CREATE TABLE chat_messages (
compressed boolean DEFAULT false NOT NULL,
created_by uuid,
content_version smallint NOT NULL,
total_cost_micros bigint
total_cost_micros bigint,
runtime_ms bigint
);
CREATE SEQUENCE chat_messages_id_seq
coderd/database/migrations/000444_chat_message_runtime_ms.down.sql
+1
View File
@@ -0,0 +1 @@
ALTER TABLE chat_messages DROP COLUMN runtime_ms;
coderd/database/migrations/000444_chat_message_runtime_ms.up.sql
+1
View File
@@ -0,0 +1 @@
ALTER TABLE chat_messages ADD COLUMN runtime_ms bigint;
coderd/database/models.go
+1
View File
@@ -4224,6 +4224,7 @@ type ChatMessage struct {
CreatedBy uuid.NullUUID `db:"created_by" json:"created_by"`
ContentVersion int16 `db:"content_version" json:"content_version"`
TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"`
RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"`
}
type ChatModelConfig struct {
coderd/database/queries.sql.go
+20 -9
View File
@@ -3866,7 +3866,7 @@ func (q *sqlQuerier) GetChatDiffStatusesByChatIDs(ctx context.Context, chatIds [
const getChatMessageByID = `-- name: GetChatMessageByID :one
SELECT
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms
FROM
chat_messages
WHERE
@@ -3895,13 +3895,14 @@ func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMess
&i.CreatedBy,
&i.ContentVersion,
&i.TotalCostMicros,
&i.RuntimeMs,
)
return i, err
}
const getChatMessagesByChatID = `-- name: GetChatMessagesByChatID :many
SELECT
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms
FROM
chat_messages
WHERE
@@ -3945,6 +3946,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
&i.CreatedBy,
&i.ContentVersion,
&i.TotalCostMicros,
&i.RuntimeMs,
); err != nil {
return nil, err
}
@@ -3961,7 +3963,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
const getChatMessagesByChatIDDescPaginated = `-- name: GetChatMessagesByChatIDDescPaginated :many
SELECT
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms
FROM
chat_messages
WHERE
@@ -4011,6 +4013,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a
&i.CreatedBy,
&i.ContentVersion,
&i.TotalCostMicros,
&i.RuntimeMs,
); err != nil {
return nil, err
}
@@ -4042,7 +4045,7 @@ WITH latest_compressed_summary AS (
1
)
SELECT
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms
FROM
chat_messages
WHERE
@@ -4110,6 +4113,7 @@ func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatI
&i.CreatedBy,
&i.ContentVersion,
&i.TotalCostMicros,
&i.RuntimeMs,
); err != nil {
return nil, err
}
@@ -4314,7 +4318,7 @@ func (q *sqlQuerier) GetChats(ctx context.Context, arg GetChatsParams) ([]Chat,
const getLastChatMessageByRole = `-- name: GetLastChatMessageByRole :one
SELECT
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms
FROM
chat_messages
WHERE
@@ -4353,6 +4357,7 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh
&i.CreatedBy,
&i.ContentVersion,
&i.TotalCostMicros,
&i.RuntimeMs,
)
return i, err
}
@@ -4540,7 +4545,8 @@ INSERT INTO chat_messages (
cache_read_tokens,
context_limit,
compressed,
total_cost_micros
total_cost_micros,
runtime_ms
) VALUES (
$1::uuid,
$2::uuid,
@@ -4557,10 +4563,11 @@ INSERT INTO chat_messages (
$13::bigint,
$14::bigint,
COALESCE($15::boolean, FALSE),
$16::bigint
$16::bigint,
$17::bigint
)
RETURNING
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms
`
type InsertChatMessageParams struct {
@@ -4580,6 +4587,7 @@ type InsertChatMessageParams struct {
ContextLimit sql.NullInt64 `db:"context_limit" json:"context_limit"`
Compressed sql.NullBool `db:"compressed" json:"compressed"`
TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"`
RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"`
}
func (q *sqlQuerier) InsertChatMessage(ctx context.Context, arg InsertChatMessageParams) (ChatMessage, error) {
@@ -4600,6 +4608,7 @@ func (q *sqlQuerier) InsertChatMessage(ctx context.Context, arg InsertChatMessag
arg.ContextLimit,
arg.Compressed,
arg.TotalCostMicros,
arg.RuntimeMs,
)
var i ChatMessage
err := row.Scan(
@@ -4621,6 +4630,7 @@ func (q *sqlQuerier) InsertChatMessage(ctx context.Context, arg InsertChatMessag
&i.CreatedBy,
&i.ContentVersion,
&i.TotalCostMicros,
&i.RuntimeMs,
)
return i, err
}
@@ -4892,7 +4902,7 @@ SET
WHERE
id = $3::bigint
RETURNING
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms
`
type UpdateChatMessageByIDParams struct {
@@ -4923,6 +4933,7 @@ func (q *sqlQuerier) UpdateChatMessageByID(ctx context.Context, arg UpdateChatMe
&i.CreatedBy,
&i.ContentVersion,
&i.TotalCostMicros,
&i.RuntimeMs,
)
return i, err
}
coderd/database/queries/chats.sql
+4 -2
View File
@@ -205,7 +205,8 @@ INSERT INTO chat_messages (
cache_read_tokens,
context_limit,
compressed,
total_cost_micros
total_cost_micros,
runtime_ms
) VALUES (
@chat_id::uuid,
sqlc.narg('created_by')::uuid,
@@ -222,7 +223,8 @@ INSERT INTO chat_messages (
sqlc.narg('cache_read_tokens')::bigint,
sqlc.narg('context_limit')::bigint,
COALESCE(sqlc.narg('compressed')::boolean, FALSE),
sqlc.narg('total_cost_micros')::bigint
sqlc.narg('total_cost_micros')::bigint,
sqlc.narg('runtime_ms')::bigint
)
RETURNING
*;