feat(agent): persist per-session token and cost accounting
Adds the agent_sessions table (migration) and a session_store module that owns its lifecycle: - ensure_session: idempotently creates the session row on first chat. - record_turn: accumulates input/output/total tokens, cost and turn count after each completed turn, bumping updated_at. - get_session: user-scoped fetch so a tenant can only read its own session. Backed by SQLite via SQLx, matching the existing service style. Unit-tested against an in-memory database (idempotency, accumulation, user scoping).
This commit is contained in:
@@ -0,0 +1,28 @@
|
||||
-- Agent Workbench sessions: per-session token / cost accounting and metadata.
|
||||
--
|
||||
-- One row per agent conversation session. Token counts and cost accumulate
|
||||
-- across the turns of a session; `updated_at` advances on every `/api/agent/chat`
|
||||
-- round-trip. `user_id` is nullable so admin (global-token) sessions are allowed
|
||||
-- without a user row, but every session is still isolated by its own sandbox
|
||||
-- directory keyed on `id`.
|
||||
CREATE TABLE IF NOT EXISTS agent_sessions (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
user_id INTEGER,
|
||||
title TEXT,
|
||||
model TEXT NOT NULL DEFAULT '',
|
||||
permission_mode TEXT NOT NULL DEFAULT 'default',
|
||||
input_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
output_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
cost_usd REAL NOT NULL DEFAULT 0,
|
||||
num_turns INTEGER NOT NULL DEFAULT 0,
|
||||
message_count INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS agent_sessions_user_id_idx
|
||||
ON agent_sessions (user_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS agent_sessions_updated_at_idx
|
||||
ON agent_sessions (updated_at);
|
||||
@@ -6,6 +6,7 @@
|
||||
//! - [`permission`]: tool risk classification + the per-mode decision matrix.
|
||||
//! - [`sandbox`]: per-session isolated working directory + path validation.
|
||||
//! - [`approval`]: process-global registry bridging the SSE approval round-trip.
|
||||
//! - [`session_store`]: persistence of per-session token/cost accounting.
|
||||
//!
|
||||
//! The HTTP surface (SSE chat, approval callback, status) lives in
|
||||
//! `crate::routes::agent`.
|
||||
@@ -13,3 +14,4 @@
|
||||
pub mod approval;
|
||||
pub mod permission;
|
||||
pub mod sandbox;
|
||||
pub mod session_store;
|
||||
|
||||
@@ -0,0 +1,201 @@
|
||||
//! Persistence for Agent Workbench session accounting.
|
||||
//!
|
||||
//! Each agent session has one row in `agent_sessions` tracking cumulative token
|
||||
//! usage and cost across its turns. A session is created (idempotently) on its
|
||||
//! first chat request and its counters are advanced after every completed turn,
|
||||
//! using the usage/cost reported by the SDK `Result` event.
|
||||
|
||||
use crate::error::{AppError, AppResult};
|
||||
use serde::Serialize;
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
/// One agent session's persisted accounting state.
|
||||
#[derive(Debug, Clone, Serialize, sqlx::FromRow)]
|
||||
pub struct AgentSession {
|
||||
pub id: String,
|
||||
pub user_id: Option<i64>,
|
||||
pub title: Option<String>,
|
||||
pub model: String,
|
||||
pub permission_mode: String,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub total_tokens: i64,
|
||||
pub cost_usd: f64,
|
||||
pub num_turns: i64,
|
||||
pub message_count: i64,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
/// Token/cost delta produced by one completed agent turn.
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct UsageDelta {
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cost_usd: f64,
|
||||
pub num_turns: i64,
|
||||
}
|
||||
|
||||
/// Create the session row if it does not exist yet. Idempotent: an existing
|
||||
/// session keeps its counters and metadata untouched.
|
||||
pub async fn ensure_session(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
user_id: Option<i64>,
|
||||
model: &str,
|
||||
permission_mode: &str,
|
||||
) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"INSERT INTO agent_sessions (id, user_id, model, permission_mode) \
|
||||
VALUES (?, ?, ?, ?) \
|
||||
ON CONFLICT(id) DO NOTHING",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(user_id)
|
||||
.bind(model)
|
||||
.bind(permission_mode)
|
||||
.execute(db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Advance a session's cumulative counters by one turn's usage delta and bump
|
||||
/// `updated_at`. The user-visible message count is incremented by `messages`
|
||||
/// (typically 2: one user prompt + one assistant reply).
|
||||
pub async fn record_turn(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
delta: UsageDelta,
|
||||
messages: i64,
|
||||
) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE agent_sessions SET \
|
||||
input_tokens = input_tokens + ?, \
|
||||
output_tokens = output_tokens + ?, \
|
||||
total_tokens = total_tokens + ?, \
|
||||
cost_usd = cost_usd + ?, \
|
||||
num_turns = num_turns + ?, \
|
||||
message_count = message_count + ?, \
|
||||
updated_at = datetime('now') \
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(delta.input_tokens)
|
||||
.bind(delta.output_tokens)
|
||||
.bind(delta.input_tokens + delta.output_tokens)
|
||||
.bind(delta.cost_usd)
|
||||
.bind(delta.num_turns)
|
||||
.bind(messages)
|
||||
.bind(session_id)
|
||||
.execute(db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fetch a single session by id (scoped to `user_id` when provided, so a user
|
||||
/// cannot read another tenant's session). Admin callers pass `None`.
|
||||
pub async fn get_session(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
user_id: Option<i64>,
|
||||
) -> AppResult<Option<AgentSession>> {
|
||||
let row = match user_id {
|
||||
Some(uid) => {
|
||||
sqlx::query_as::<_, AgentSession>(
|
||||
"SELECT * FROM agent_sessions WHERE id = ? AND user_id = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(uid)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
}
|
||||
None => {
|
||||
sqlx::query_as::<_, AgentSession>("SELECT * FROM agent_sessions WHERE id = ?")
|
||||
.bind(session_id)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
}
|
||||
}
|
||||
.map_err(AppError::Database)?;
|
||||
Ok(row)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
async fn pool() -> SqlitePool {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(include_str!(
|
||||
"../../../migrations/20260622000001_agent_sessions.sql"
|
||||
))
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ensure_is_idempotent_and_accumulates() {
|
||||
let db = pool().await;
|
||||
ensure_session(&db, "s1", Some(7), "gpt-4o", "default")
|
||||
.await
|
||||
.unwrap();
|
||||
// Second ensure must not reset counters.
|
||||
record_turn(
|
||||
&db,
|
||||
"s1",
|
||||
UsageDelta { input_tokens: 100, output_tokens: 40, cost_usd: 0.5, num_turns: 3 },
|
||||
2,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
ensure_session(&db, "s1", Some(7), "gpt-4o", "default")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let s = get_session(&db, "s1", Some(7)).await.unwrap().unwrap();
|
||||
assert_eq!(s.input_tokens, 100);
|
||||
assert_eq!(s.output_tokens, 40);
|
||||
assert_eq!(s.total_tokens, 140);
|
||||
assert!((s.cost_usd - 0.5).abs() < 1e-9);
|
||||
assert_eq!(s.num_turns, 3);
|
||||
assert_eq!(s.message_count, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn record_turn_adds_across_turns() {
|
||||
let db = pool().await;
|
||||
ensure_session(&db, "s2", None, "sonnet", "default").await.unwrap();
|
||||
for _ in 0..3 {
|
||||
record_turn(
|
||||
&db,
|
||||
"s2",
|
||||
UsageDelta { input_tokens: 10, output_tokens: 5, cost_usd: 0.01, num_turns: 1 },
|
||||
2,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let s = get_session(&db, "s2", None).await.unwrap().unwrap();
|
||||
assert_eq!(s.total_tokens, 45);
|
||||
assert_eq!(s.num_turns, 3);
|
||||
assert_eq!(s.message_count, 6);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_session_is_user_scoped() {
|
||||
let db = pool().await;
|
||||
ensure_session(&db, "s3", Some(1), "m", "default").await.unwrap();
|
||||
// Wrong user cannot see it.
|
||||
assert!(get_session(&db, "s3", Some(2)).await.unwrap().is_none());
|
||||
// Owner can.
|
||||
assert!(get_session(&db, "s3", Some(1)).await.unwrap().is_some());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user