From eeacd5554cb74b844b53d928d515322abebb4577 Mon Sep 17 00:00:00 2001 From: Bhavesh Patel Date: Thu, 27 Aug 2026 23:21:43 +0530 Subject: [PATCH 1/4] Update session metrics on createTurn and terminal writes --- .changeset/20260827210200-session-metrics.md | 2 +- .../src/agent-session/store/ISessionStore.ts | 14 +- .../store/InMemorySessionStore.ts | 14 ++ .../agent-session/store/storeContractSuite.ts | 159 ++++++++++++++++++ .../postgres/session-store/queries/turns.ts | 58 ++++++- .../db/sqlite/session-store/queries/turns.ts | 43 ++++- 6 files changed, 277 insertions(+), 13 deletions(-) diff --git a/.changeset/20260827210200-session-metrics.md b/.changeset/20260827210200-session-metrics.md index b5f1befbd..8cc3536e4 100644 --- a/.changeset/20260827210200-session-metrics.md +++ b/.changeset/20260827210200-session-metrics.md @@ -3,4 +3,4 @@ '@truefoundry/trueforge': patch --- -Persist zero-initialized metrics on agent sessions. +Persist zero-initialized session metrics and fold 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..06f69207f 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. @@ -282,9 +283,10 @@ export interface ISessionStore< /** * 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}. + * did cancel, insert the caller-built `turn_done_event` and add cost plus + * `completed_at − created_at` into `session.metrics` — already-terminal + * turns skip the event insert and the metrics fold; (c) return the now-immutable + * turn record. Missing turn → {@link TurnNotFoundError}. */ freezeAndGetTurn(input: FreezeAndGetTurnInput): Promise>; @@ -297,7 +299,9 @@ export interface ISessionStore< ): Promise<{ data: TurnRecordWithoutSnapshot[]; pagination: TokenPagination }>; /** - * Writes the terminal state and `turn_done_event` atomically. Store contract — + * Writes the terminal state and `turn_done_event` atomically, and adds + * `metrics.total_cost_in_usd` plus `completed_at − created_at` into + * `session.metrics` in that same unit. Store contract — * **first terminal write wins**: * - Allowed: `running` → `done` | `cancelled` | `error`. * - Rejected with **409** (or equivalent conflict): any write when status is already 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..d90f2625d 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,115 @@ 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('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..b38071a6d 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,39 @@ 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; created_at: Date; state: TerminalTurnState }, +): Promise { + const elapsed_ms = Date.parse(input.state.completed_at) - input.created_at.getTime(); + const total_cost_in_usd = input.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)`, + ), + sql`'{total_duration_ms}'`, + sql`to_jsonb((metrics->>'total_duration_ms')::int + ${total_duration_ms}::int)`, + ), + }) + .where('session_id', '=', input.session_id) + .execute(); +} + function terminalTurnState(state: TurnState, turn_id: string): TerminalTurnState { switch (state.status) { case 'running': @@ -384,6 +418,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 +654,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 +668,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, + created_at: updateResult.created_at, + state: cancelledState, + }); } const record = await assembleTurnRecord(trx, input); @@ -706,10 +749,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 +767,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, + created_at: result.created_at, + 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..9415a2e03 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; created_at: Date; state: TerminalTurnState }, +): Promise { + const elapsed_ms = Date.parse(input.state.completed_at) - input.created_at.getTime(); + const total_cost_in_usd = input.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, + created_at: new Date(updateResult.created_at), + 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, + created_at: new Date(result.created_at), + state: input.state, + }); + await trx .insertInto('session_event') .values({ From 08b4f41342a63494c5ebe1f08a0292ed4f16ed93 Mon Sep 17 00:00:00 2001 From: Srajan Asthana Date: Mon, 31 Aug 2026 14:02:46 +0530 Subject: [PATCH 2/4] fix: duration fold overflowing postgres int --- .../trueforge/src/db/postgres/session-store/queries/turns.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 b38071a6d..7600be953 100644 --- a/packages/trueforge/src/db/postgres/session-store/queries/turns.ts +++ b/packages/trueforge/src/db/postgres/session-store/queries/turns.ts @@ -136,8 +136,9 @@ async function addSessionCostAndDuration( 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')::int + ${total_duration_ms}::int)`, + sql`to_jsonb((metrics->>'total_duration_ms')::bigint + ${total_duration_ms}::bigint)`, ), }) .where('session_id', '=', input.session_id) From b0aaefb710f816445835d5f0784ce3f091afb885 Mon Sep 17 00:00:00 2001 From: Srajan Asthana Date: Mon, 31 Aug 2026 14:09:00 +0530 Subject: [PATCH 3/4] update changeset --- .changeset/20260827210200-session-metrics.md | 2 +- .changeset/20260831083812-session-metrics-fold.md | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) create mode 100644 .changeset/20260831083812-session-metrics-fold.md diff --git a/.changeset/20260827210200-session-metrics.md b/.changeset/20260827210200-session-metrics.md index 8cc3536e4..b5f1befbd 100644 --- a/.changeset/20260827210200-session-metrics.md +++ b/.changeset/20260827210200-session-metrics.md @@ -3,4 +3,4 @@ '@truefoundry/trueforge': patch --- -Persist zero-initialized session metrics and fold totals on createTurn and terminal writes. +Persist zero-initialized metrics on agent sessions. 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. From de3ba3e96aa5221d4d107e7491946f5abfe8b298 Mon Sep 17 00:00:00 2001 From: Srajan Asthana Date: Mon, 31 Aug 2026 16:22:51 +0530 Subject: [PATCH 4/4] resolve comments --- .../src/agent-session/store/ISessionStore.ts | 21 ++------ .../agent-session/store/storeContractSuite.ts | 48 +++++++++++++++++++ .../postgres/session-store/queries/turns.ts | 14 +++--- .../db/sqlite/session-store/queries/turns.ts | 14 +++--- 4 files changed, 67 insertions(+), 30 deletions(-) diff --git a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts index 06f69207f..87222662a 100644 --- a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts @@ -282,11 +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` and add cost plus - * `completed_at − created_at` into `session.metrics` — already-terminal - * turns skip the event insert and the metrics fold; (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>; @@ -299,17 +296,9 @@ export interface ISessionStore< ): Promise<{ data: TurnRecordWithoutSnapshot[]; pagination: TokenPagination }>; /** - * Writes the terminal state and `turn_done_event` atomically, and adds - * `metrics.total_cost_in_usd` plus `completed_at − created_at` into - * `session.metrics` in that same unit. 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/tests/agent-session/store/storeContractSuite.ts b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts index d90f2625d..243f9f537 100644 --- a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts +++ b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts @@ -1619,6 +1619,54 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { }); }); + 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 7600be953..f333b8e5e 100644 --- a/packages/trueforge/src/db/postgres/session-store/queries/turns.ts +++ b/packages/trueforge/src/db/postgres/session-store/queries/turns.ts @@ -122,10 +122,10 @@ function incrementSessionTotalTurns(): RawBuilder { /** Same tx as the turn flip; the session row lock serializes concurrent terminal folds. */ async function addSessionCostAndDuration( trx: Transaction, - input: { session_id: string; created_at: Date; state: TerminalTurnState }, + input: { session_id: string; turn_created_at: Date; turn_state: TerminalTurnState }, ): Promise { - const elapsed_ms = Date.parse(input.state.completed_at) - input.created_at.getTime(); - const total_cost_in_usd = input.state.metrics?.total_cost_in_usd ?? 0; + 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') @@ -672,8 +672,8 @@ export async function freezeAndGetTurn(db: Kysely, input: FreezeAndGet // Only the winning cancel folds; a freeze of an already-terminal turn is a read. await addSessionCostAndDuration(trx, { session_id: input.session_id, - created_at: updateResult.created_at, - state: cancelledState, + turn_created_at: updateResult.created_at, + turn_state: cancelledState, }); } @@ -770,8 +770,8 @@ export async function updateTurnState(db: Kysely, input: UpdateTurnSta await addSessionCostAndDuration(trx, { session_id: input.session_id, - created_at: result.created_at, - state: input.state, + turn_created_at: result.created_at, + turn_state: input.state, }); await trx 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 9415a2e03..fc1fe30b4 100644 --- a/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts +++ b/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts @@ -112,10 +112,10 @@ 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; created_at: Date; state: TerminalTurnState }, + input: { session_id: string; turn_created_at: Date; turn_state: TerminalTurnState }, ): Promise { - const elapsed_ms = Date.parse(input.state.completed_at) - input.created_at.getTime(); - const total_cost_in_usd = input.state.metrics?.total_cost_in_usd ?? 0; + 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') @@ -709,8 +709,8 @@ export async function freezeAndGetTurn(db: Kysely, input: FreezeAndGet // Only the winning cancel folds; a freeze of an already-terminal turn is a read. await addSessionCostAndDuration(trx, { session_id: input.session_id, - created_at: new Date(updateResult.created_at), - state: cancelledState, + turn_created_at: new Date(updateResult.created_at), + turn_state: cancelledState, }); } @@ -816,8 +816,8 @@ export async function updateTurnState(db: Kysely, input: UpdateTurnSta await addSessionCostAndDuration(trx, { session_id: input.session_id, - created_at: new Date(result.created_at), - state: input.state, + turn_created_at: new Date(result.created_at), + turn_state: input.state, }); await trx