diff --git a/.changeset/20260831083812-session-metrics-fold.md b/.changeset/20260831083812-session-metrics-fold.md new file mode 100644 index 000000000..2ea79a557 --- /dev/null +++ b/.changeset/20260831083812-session-metrics-fold.md @@ -0,0 +1,6 @@ +--- +'@truefoundry/trueforge-core': patch +'@truefoundry/trueforge': patch +--- + +Fold session metrics totals on createTurn and terminal writes. diff --git a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts index 30cfa6bf6..87222662a 100644 --- a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts @@ -262,7 +262,8 @@ export interface ISessionStore< * row or leave `last_turn_id` pointing at a turn that was never created. * The implementation supplies the mechanism (session lock, row lock/tx, …). * - * Also bumps `session.last_activity_timestamp_ms` in that same atomic unit. + * Also bumps `session.last_activity_timestamp_ms` and increments + * `session.metrics.total_turns` in that same atomic unit. * * Fork semantics for `turn.previous_turn_id`: * - `null` — new root turn (no parent); always allowed. @@ -281,10 +282,8 @@ export interface ISessionStore< createTurn(input: CreateTurnInput): Promise; /** - * One tx: (a) conditionally cancel a running turn with `reason`; (b) if it - * did cancel, insert the caller-built `turn_done_event` — already-terminal - * turns skip the event insert; (c) return the now-immutable turn record. - * Missing turn → {@link TurnNotFoundError}. + * Cancel if still running (persist `turn_done` and fold cost/duration into + * `session.metrics`); already-terminal turns are a read. Missing → {@link TurnNotFoundError}. */ freezeAndGetTurn(input: FreezeAndGetTurnInput): Promise>; @@ -297,15 +296,9 @@ export interface ISessionStore< ): Promise<{ data: TurnRecordWithoutSnapshot[]; pagination: TokenPagination }>; /** - * Writes the terminal state and `turn_done_event` atomically. Store contract — - * **first terminal write wins**: - * - Allowed: `running` → `done` | `cancelled` | `error`. - * - Rejected with **409** (or equivalent conflict): any write when status is already - * terminal — including done→cancelled, cancelled→done, error→*, terminal→running. - * - Missing turn → 404 / not-found. - * - * Check under the same concurrency control as other turn mutations (lock/CAS) — - * not a racy read-then-write outside the critical section. + * First terminal write wins (`running` → done/cancelled/error); otherwise 409. + * Winning write also folds cost/duration into `session.metrics`. Missing → 404. + * Must use the same lock/CAS as other turn mutations. */ updateTurnState(input: UpdateTurnStateInput): Promise; diff --git a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts index 5ebeacaf0..095f32c47 100644 --- a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts @@ -337,6 +337,7 @@ export class InMemorySessionStore< this.events.set(tKey, []); stored.turnIds.push(input.turn.turn_id); stored.record.last_turn_id = input.turn.turn_id; + stored.record.metrics.total_turns += 1; stored.record.last_activity_timestamp_ms = Date.now(); stored.record.updated_at = new Date(); if (input.update_session_title_if_not_exist !== null && stored.record.title === null) { @@ -360,6 +361,7 @@ export class InMemorySessionStore< if (list) { list.push(deepCopy(input.turn_done_event)); } + this.addTerminalSessionMetrics(input.session_id, turn.created_at, cancelledState); } return deepCopy(turn); @@ -402,6 +404,7 @@ export class InMemorySessionStore< if (list) { list.push(deepCopy(input.turn_done_event)); } + this.addTerminalSessionMetrics(input.session_id, turn.created_at, input.state); } async appendToEvents(input: AppendToEventsInput): Promise { @@ -415,6 +418,17 @@ export class InMemorySessionStore< return; } + /** Cost from turn metrics; duration is completed_at − created_at, floored at 0. */ + private addTerminalSessionMetrics(sessionId: string, created_at: Date, state: TerminalTurnState): void { + const stored = this.sessions.get(sessionKey(sessionId)); + if (!stored) { + throw new SessionNotFoundError(sessionId); + } + const elapsed_ms = Date.parse(state.completed_at) - created_at.getTime(); + stored.record.metrics.total_cost_in_usd += state.metrics?.total_cost_in_usd ?? 0; + stored.record.metrics.total_duration_ms += elapsed_ms > 0 ? Math.trunc(elapsed_ms) : 0; + } + private requireTurn(sessionId: string, turnId: string): TurnRecord { const turn = this.turns.get(turnKey({ session_id: sessionId, turn_id: turnId })); if (!turn) { diff --git a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts index 30ec42212..243f9f537 100644 --- a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts +++ b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts @@ -840,6 +840,24 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { expect(mustGet(after).last_activity_timestamp_ms).toBeGreaterThan(mustGet(before).last_activity_timestamp_ms); }); + it('increments session.metrics.total_turns without cost or duration', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + const afterFirst = await store.getSession({ tenant_id: tenant, session_id: sessionId }); + expect(mustGet(afterFirst).metrics).toEqual({ + total_cost_in_usd: 0, + total_duration_ms: 0, + total_turns: 1, + }); + await finishTurn(store, 'turn-1'); + await store.createTurn( + makeCreateTurnInput({ sessionId, turnId: 'turn-2', previousTurnId: 'turn-1', firstTurnId: 'turn-1' }), + ); + const afterSecond = await store.getSession({ tenant_id: tenant, session_id: sessionId }); + expect(mustGet(afterSecond).metrics.total_turns).toBe(2); + }); + it('update_session_title_if_not_exist sets once and never overwrites', async () => { const store = createStore(); await seedSession(store); @@ -1322,6 +1340,38 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { }); }); + it('folds duration into session.metrics when cancel applies, not on a second freeze', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + const cancelledState = makeCancelledTurnState(CancellationReason.CancelledForNextTurn); + const record = await store.freezeAndGetTurn({ + session_id: sessionId, + turn_id: 'turn-1', + reason: CancellationReason.CancelledForNextTurn, + turn_done_event: makeTurnDoneEvent(cancelledState), + }); + if (record.state.status !== 'cancelled') { + throw new Error(`expected cancelled turn, got ${record.state.status}`); + } + const afterCancel = await store.getSession({ tenant_id: tenant, session_id: sessionId }); + const elapsed_ms = Date.parse(record.state.completed_at) - record.created_at.getTime(); + expect(mustGet(afterCancel).metrics).toEqual({ + total_cost_in_usd: 0, + total_duration_ms: elapsed_ms > 0 ? Math.trunc(elapsed_ms) : 0, + total_turns: 1, + }); + + await store.freezeAndGetTurn({ + session_id: sessionId, + turn_id: 'turn-1', + reason: CancellationReason.CancelledForNextTurn, + turn_done_event: makeTurnDoneEvent(cancelledState), + }); + const afterSecond = await store.getSession({ tenant_id: tenant, session_id: sessionId }); + expect(mustGet(afterSecond).metrics).toEqual(mustGet(afterCancel).metrics); + }); + it('on an already-terminal turn is a plain read without duplicating turn.done', async () => { const store = createStore(); await seedSession(store); @@ -1460,6 +1510,163 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { expect(doneEvents[0]).toEqual(turnDone); }); + it('adds cost and duration into session.metrics on running → terminal', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + const turn = await store.getTurn({ session_id: sessionId, turn_id: 'turn-1' }); + const createdAt = mustGet(turn).created_at; + const completedAt = new Date(createdAt.getTime() + 1500).toISOString(); + const state = { + ...makeDoneTurnState(), + completed_at: completedAt, + metrics: { total_cost_in_usd: 1.25 }, + }; + await store.updateTurnState({ + session_id: sessionId, + turn_id: 'turn-1', + state, + turn_done_event: makeTurnDoneEvent(state), + }); + const session = await store.getSession({ tenant_id: tenant, session_id: sessionId }); + expect(mustGet(session).metrics).toEqual({ + total_cost_in_usd: 1.25, + total_duration_ms: 1500, + total_turns: 1, + }); + }); + + it('does not add session.metrics again on a losing terminal write', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + const turn = await store.getTurn({ session_id: sessionId, turn_id: 'turn-1' }); + const createdAt = mustGet(turn).created_at; + const doneState = { + ...makeDoneTurnState(), + completed_at: new Date(createdAt.getTime() + 1500).toISOString(), + metrics: { total_cost_in_usd: 1.25 }, + }; + await store.updateTurnState({ + session_id: sessionId, + turn_id: 'turn-1', + state: doneState, + turn_done_event: makeTurnDoneEvent(doneState), + }); + const afterFirst = mustGet(await store.getSession({ tenant_id: tenant, session_id: sessionId })).metrics; + + const losingState = { + ...makeCancelledTurnState(CancellationReason.ClientCancelled), + completed_at: new Date(createdAt.getTime() + 8000).toISOString(), + metrics: { total_cost_in_usd: 9.99 }, + }; + await expect( + store.updateTurnState({ + session_id: sessionId, + turn_id: 'turn-1', + state: losingState, + turn_done_event: makeTurnDoneEvent(losingState), + }), + ).rejects.toBeInstanceOf(SessionStoreConflictError); + + const afterSecond = await store.getSession({ tenant_id: tenant, session_id: sessionId }); + expect(mustGet(afterSecond).metrics).toEqual(afterFirst); + expect(afterFirst).toEqual({ + total_cost_in_usd: 1.25, + total_duration_ms: 1500, + total_turns: 1, + }); + }); + + it('adds cost and duration from a second done turn onto existing session.metrics', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + const turn1 = await store.getTurn({ session_id: sessionId, turn_id: 'turn-1' }); + const turn1Done = { + ...makeDoneTurnState(), + completed_at: new Date(mustGet(turn1).created_at.getTime() + 1500).toISOString(), + metrics: { total_cost_in_usd: 1.25 }, + }; + await store.updateTurnState({ + session_id: sessionId, + turn_id: 'turn-1', + state: turn1Done, + turn_done_event: makeTurnDoneEvent(turn1Done), + }); + + await store.createTurn( + makeCreateTurnInput({ sessionId, turnId: 'turn-2', previousTurnId: 'turn-1', firstTurnId: 'turn-1' }), + ); + const turn2 = await store.getTurn({ session_id: sessionId, turn_id: 'turn-2' }); + const turn2Done = { + ...makeDoneTurnState(), + completed_at: new Date(mustGet(turn2).created_at.getTime() + 800).toISOString(), + metrics: { total_cost_in_usd: 0.5 }, + }; + await store.updateTurnState({ + session_id: sessionId, + turn_id: 'turn-2', + state: turn2Done, + turn_done_event: makeTurnDoneEvent(turn2Done), + }); + + const session = await store.getSession({ tenant_id: tenant, session_id: sessionId }); + expect(mustGet(session).metrics).toEqual({ + total_cost_in_usd: 1.75, + total_duration_ms: 2300, + total_turns: 2, + }); + }); + + it('accumulates session.metrics across turn1 done, turn2 cancel, and turn3 create', async () => { + const store = createStore(); + await seedSession(store); + + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + const turn1 = await store.getTurn({ session_id: sessionId, turn_id: 'turn-1' }); + const turn1Done = { + ...makeDoneTurnState(), + completed_at: new Date(mustGet(turn1).created_at.getTime() + 1000).toISOString(), + metrics: { total_cost_in_usd: 1.0 }, + }; + await store.updateTurnState({ + session_id: sessionId, + turn_id: 'turn-1', + state: turn1Done, + turn_done_event: makeTurnDoneEvent(turn1Done), + }); + + await store.createTurn( + makeCreateTurnInput({ sessionId, turnId: 'turn-2', previousTurnId: 'turn-1', firstTurnId: 'turn-1' }), + ); + const cancelledState = makeCancelledTurnState(CancellationReason.CancelledForNextTurn); + const turn2 = await store.freezeAndGetTurn({ + session_id: sessionId, + turn_id: 'turn-2', + reason: CancellationReason.CancelledForNextTurn, + turn_done_event: makeTurnDoneEvent(cancelledState), + }); + if (turn2.state.status !== 'cancelled') { + throw new Error(`expected cancelled turn, got ${turn2.state.status}`); + } + const turn2DurationMs = Math.max( + 0, + Math.trunc(Date.parse(turn2.state.completed_at) - turn2.created_at.getTime()), + ); + + await store.createTurn( + makeCreateTurnInput({ sessionId, turnId: 'turn-3', previousTurnId: 'turn-2', firstTurnId: 'turn-1' }), + ); + + const session = await store.getSession({ tenant_id: tenant, session_id: sessionId }); + expect(mustGet(session).metrics).toEqual({ + total_cost_in_usd: 1.0, + total_duration_ms: 1000 + turn2DurationMs, + total_turns: 3, + }); + }); + it('missing turn → not found', async () => { const store = createStore(); await seedSession(store); diff --git a/packages/trueforge/src/db/postgres/session-store/queries/turns.ts b/packages/trueforge/src/db/postgres/session-store/queries/turns.ts index f89176048..f333b8e5e 100644 --- a/packages/trueforge/src/db/postgres/session-store/queries/turns.ts +++ b/packages/trueforge/src/db/postgres/session-store/queries/turns.ts @@ -1,3 +1,4 @@ +import type { SessionMetrics } from '@truefoundry/trueforge-core/agent-session'; import type { TurnRecord, TurnSnapshot } from '@truefoundry/trueforge-core/agent-session/models/TurnRecord'; import { type TerminalTurnState, @@ -28,7 +29,7 @@ import { getEmptyCurrentContextUsage } from '@truefoundry/trueforge-core/core/ru import type { SandboxInfo } from '@truefoundry/trueforge-core/core/sandbox/Sandbox'; import { sql, type Kysely, type QueryCreator, type RawBuilder, type Transaction } from 'kysely'; import { isUniqueViolation } from '../../client'; -import { json } from '../../sqlExpressions'; +import { json, jsonbSet } from '../../sqlExpressions'; import type { Database, TurnCheckpoint, TurnThreadCheckpoint } from '../../types'; import { lateralUnnestBigintArrayWithOrdinality } from '../sqlExpressions'; @@ -110,6 +111,40 @@ export interface ListTurnsResult { type DbOrTrx = Kysely | Transaction; type TurnFenceDb = DbOrTrx | QueryCreator; +function incrementSessionTotalTurns(): RawBuilder { + return jsonbSet( + sql`metrics`, + sql`'{total_turns}'`, + sql`to_jsonb((metrics->>'total_turns')::int + 1)`, + ); +} + +/** Same tx as the turn flip; the session row lock serializes concurrent terminal folds. */ +async function addSessionCostAndDuration( + trx: Transaction, + input: { session_id: string; turn_created_at: Date; turn_state: TerminalTurnState }, +): Promise { + const elapsed_ms = Date.parse(input.turn_state.completed_at) - input.turn_created_at.getTime(); + const total_cost_in_usd = input.turn_state.metrics?.total_cost_in_usd ?? 0; + const total_duration_ms = elapsed_ms > 0 ? Math.trunc(elapsed_ms) : 0; + await trx + .updateTable('session') + .set({ + metrics: jsonbSet( + jsonbSet( + sql`metrics`, + sql`'{total_cost_in_usd}'`, + sql`to_jsonb((metrics->>'total_cost_in_usd')::double precision + ${total_cost_in_usd}::double precision)`, + ), + // bigint: ::int overflows at ~24.8 days of summed ms and would roll back the terminal tx. + sql`'{total_duration_ms}'`, + sql`to_jsonb((metrics->>'total_duration_ms')::bigint + ${total_duration_ms}::bigint)`, + ), + }) + .where('session_id', '=', input.session_id) + .execute(); +} + function terminalTurnState(state: TurnState, turn_id: string): TerminalTurnState { switch (state.status) { case 'running': @@ -384,6 +419,8 @@ export async function createTurn(db: Kysely, input: CreateTurnInput): last_turn_id: input.turn.turn_id, updated_at: sql`now()`, last_activity_timestamp_ms: input.last_activity_timestamp_ms, + // total_turns rides the same tip UPDATE so a later failure in this tx rolls it back. + metrics: incrementSessionTotalTurns(), ...(input.update_session_title_if_not_exist !== null ? { title: sql`COALESCE(title, ${input.update_session_title_if_not_exist})`, @@ -618,9 +655,10 @@ export async function freezeAndGetTurn(db: Kysely, input: FreezeAndGet .where('session_id', '=', input.session_id) .where('turn_id', '=', input.turn_id) .where(sql`state->>'status' = 'running'`) + .returning(['created_at']) .executeTakeFirst(); - if (Number(updateResult.numUpdatedRows) > 0) { + if (updateResult !== undefined) { await trx .insertInto('session_event') .values({ @@ -631,6 +669,12 @@ export async function freezeAndGetTurn(db: Kysely, input: FreezeAndGet created_at: new Date(input.turn_done_event.created_at), }) .execute(); + // Only the winning cancel folds; a freeze of an already-terminal turn is a read. + await addSessionCostAndDuration(trx, { + session_id: input.session_id, + turn_created_at: updateResult.created_at, + turn_state: cancelledState, + }); } const record = await assembleTurnRecord(trx, input); @@ -706,10 +750,11 @@ export async function updateTurnState(db: Kysely, input: UpdateTurnSta .where('session_id', '=', input.session_id) .where('turn_id', '=', input.turn_id) .where(sql`state->>'status' = 'running'`) + .returning(['created_at']) .executeTakeFirst(); - const numUpdated = Number(result.numUpdatedRows); - if (numUpdated === 0) { + // No RETURNING row: UPDATE matched 0 running turns. + if (result === undefined) { const existing = await trx .selectFrom('turn') .select('state') @@ -723,6 +768,12 @@ export async function updateTurnState(db: Kysely, input: UpdateTurnSta throw new TurnNotRunningError(input.turn_id, terminalTurnState(existing.state, input.turn_id)); } + await addSessionCostAndDuration(trx, { + session_id: input.session_id, + turn_created_at: result.created_at, + turn_state: input.state, + }); + await trx .insertInto('session_event') .values({ diff --git a/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts b/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts index 0c5416127..fc1fe30b4 100644 --- a/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts +++ b/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts @@ -109,6 +109,27 @@ export interface ListTurnsResult { type DbOrTrx = Kysely | Transaction; +/** Same tx as the turn flip; BEGIN IMMEDIATE serializes concurrent terminal folds. */ +async function addSessionCostAndDuration( + trx: Transaction, + input: { session_id: string; turn_created_at: Date; turn_state: TerminalTurnState }, +): Promise { + const elapsed_ms = Date.parse(input.turn_state.completed_at) - input.turn_created_at.getTime(); + const total_cost_in_usd = input.turn_state.metrics?.total_cost_in_usd ?? 0; + const total_duration_ms = elapsed_ms > 0 ? Math.trunc(elapsed_ms) : 0; + await trx + .updateTable('session') + .set({ + metrics: sql`jsonb_set( + jsonb_set(metrics, '$.total_cost_in_usd', jsonb((metrics->>'total_cost_in_usd') + ${total_cost_in_usd})), + '$.total_duration_ms', + jsonb((metrics->>'total_duration_ms') + ${total_duration_ms}) + )`, + }) + .where('session_id', '=', input.session_id) + .execute(); +} + function terminalTurnState(state: TurnState, turn_id: string): TerminalTurnState { switch (state.status) { case 'running': @@ -334,6 +355,8 @@ export async function createTurn(db: Kysely, input: CreateTurnInput): last_turn_id: input.turn.turn_id, updated_at: nowIso(), last_activity_timestamp_ms: input.last_activity_timestamp_ms, + // total_turns rides the same tip UPDATE so a later failure in this tx rolls it back. + metrics: sql`jsonb_set(metrics, '$.total_turns', jsonb((metrics->>'total_turns') + 1))`, }) .where('session_id', '=', input.session_id); @@ -669,9 +692,10 @@ export async function freezeAndGetTurn(db: Kysely, input: FreezeAndGet .where('session_id', '=', input.session_id) .where('turn_id', '=', input.turn_id) .where(sql`state->>'status' = 'running'`) + .returning(['created_at']) .executeTakeFirst(); - if (Number(updateResult.numUpdatedRows) > 0) { + if (updateResult !== undefined) { await trx .insertInto('session_event') .values({ @@ -682,6 +706,12 @@ export async function freezeAndGetTurn(db: Kysely, input: FreezeAndGet created_at: input.turn_done_event.created_at, }) .execute(); + // Only the winning cancel folds; a freeze of an already-terminal turn is a read. + await addSessionCostAndDuration(trx, { + session_id: input.session_id, + turn_created_at: new Date(updateResult.created_at), + turn_state: cancelledState, + }); } const record = await assembleTurnRecord(trx, input); @@ -766,10 +796,11 @@ export async function updateTurnState(db: Kysely, input: UpdateTurnSta .where('session_id', '=', input.session_id) .where('turn_id', '=', input.turn_id) .where(sql`state->>'status' = 'running'`) + .returning(['created_at']) .executeTakeFirst(); - const numUpdated = Number(result.numUpdatedRows); - if (numUpdated === 0) { + // No RETURNING row: UPDATE matched 0 running turns. + if (result === undefined) { const existing = await trx .selectFrom('turn') .select([jsonText(sql.ref('state')).as('state')]) @@ -783,6 +814,12 @@ export async function updateTurnState(db: Kysely, input: UpdateTurnSta throw new TurnNotRunningError(input.turn_id, terminalTurnState(existing.state, input.turn_id)); } + await addSessionCostAndDuration(trx, { + session_id: input.session_id, + turn_created_at: new Date(result.created_at), + turn_state: input.state, + }); + await trx .insertInto('session_event') .values({