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

## 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
+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
@@ -0,0 +1 @@
ALTER TABLE chat_messages DROP COLUMN runtime_ms;
@@ -0,0 +1 @@
ALTER TABLE chat_messages ADD COLUMN runtime_ms bigint;
+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 {
+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
}
+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
*;