diff --git a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts index 172378b9a9ef1..4e7bd34fd1f88 100644 --- a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts +++ b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts @@ -22,6 +22,7 @@ import { ProviderInstanceId, ProviderSessionId, ProviderThreadId, + ProviderTurnId, RunAttemptId, RunId, ThreadId, @@ -907,6 +908,30 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { }, }); + // Updates without usage inherit it from an earlier update, so the + // newest update carrying usage must outlive compaction. + const providerTurnEvent = ( + suffix: string, + usedTokens: number | null, + ): OrchestrationV2DomainEvent => ({ + id: EventId.make(`event:foundation-compact:${suffix}`), + type: "provider-turn.updated", + threadId, + occurredAt: now, + payload: { + id: ProviderTurnId.make("provider-turn:foundation-compact"), + providerThreadId: ProviderThreadId.make("provider-thread:foundation-compact"), + nodeId, + runAttemptId: null, + nativeTurnRef: null, + ordinal: 1, + status: usedTokens === null ? "completed" : "running", + startedAt: now, + completedAt: null, + ...(usedTokens === null ? {} : { tokenUsage: { usedTokens, updatedAt: nowIso } }), + }, + }); + yield* eventSink.write({ events: [ threadCreatedEvent({ id: "event:foundation-compact:create", thread, now }), @@ -915,6 +940,7 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { messageEvent("message-1", "streaming"), nodeEvent("node-1", "running"), itemEvent("item-1", "streaming"), + providerTurnEvent("turn-1", null), ...Array.from({ length: 501 }, (_, index) => threadStateEvent(`history-${index}`, "thread.visited"), ), @@ -922,6 +948,9 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { messageEvent("message-2", "final"), nodeEvent("node-2", "completed"), itemEvent("item-2", "final"), + providerTurnEvent("turn-2", 10), + providerTurnEvent("turn-3", 20), + providerTurnEvent("turn-4", null), ], }); const beforeCompaction = yield* projections.getThreadProjection(threadId); @@ -963,9 +992,10 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { const summary = yield* maintenance.compactEventStore; // Superseded state spans several discovery pages. Both turn-item updates - // stay, and so do the first message and node events (replay order). The + // stay, and so do the first message, node, and provider-turn events + // (replay order) plus the newest provider-turn usage. The // rename stays too: visits are read state and only supersede each other. - assert.isAtLeast(summary.deletedEventCount, 504); + assert.isAtLeast(summary.deletedEventCount, 505); assert.isAtLeast(summary.deletedReceiptCount, 1); const remaining = yield* sql<{ readonly event_id: string }>` @@ -983,10 +1013,13 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { "event:foundation-compact:message-1", "event:foundation-compact:node-1", "event:foundation-compact:item-1", + "event:foundation-compact:turn-1", "event:foundation-compact:visit-2", "event:foundation-compact:message-2", "event:foundation-compact:node-2", "event:foundation-compact:item-2", + "event:foundation-compact:turn-3", + "event:foundation-compact:turn-4", ], ); diff --git a/apps/server/src/orchestration-v2/ProjectionMaintenance.ts b/apps/server/src/orchestration-v2/ProjectionMaintenance.ts index 6c5835d8f40ec..4502df612a9db 100644 --- a/apps/server/src/orchestration-v2/ProjectionMaintenance.ts +++ b/apps/server/src/orchestration-v2/ProjectionMaintenance.ts @@ -232,10 +232,12 @@ export const layer: Layer.Layer< /** * Scan newest first to retain the newest state for each entity. Keyed - * entities (messages, nodes, PR links) also keep their oldest event: client - * reducers append an entity where they first see it, so a replay from any - * cursor must either carry that first event or find the entity already - * present. Only the events in between are deleted. + * entities (messages, nodes, provider turns, PR links) also keep their + * oldest event: client reducers append an entity where they first see it, + * so a replay from any cursor must either carry that first event or find + * the entity already present. Only the events in between are deleted. + * Provider turns also keep their newest event with token usage, because + * reducers carry usage forward from earlier updates that omit it. * Page every event, including non-candidates: filtering before LIMIT could * still scan the entire history when superseded events are sparse. * turn-item.updated stays intact because replay assigns positions on first write. @@ -256,6 +258,7 @@ export const layer: Layer.Layer< // Per key: the lowest sequence seen so far after the newest. It is // deleted once an even older event turns up and kept otherwise. const oldestCandidateByEntityKey = new Map(); + const providerTurnKeysWithRetainedUsage = new Set(); let deletedEventCount = 0; let deletedReceiptCount = 0; @@ -267,6 +270,7 @@ export const layer: Layer.Layer< readonly stream_id: string; readonly event_type: string; readonly entity_id: string | null; + readonly carries_token_usage: number; readonly imported_legacy_thread: number; }>` SELECT @@ -277,7 +281,7 @@ export const layer: Layer.Layer< event.event_type, CASE WHEN event.application_event_version = 2 - AND event.event_type IN ('message.updated', 'node.updated') + AND event.event_type IN ('message.updated', 'node.updated', 'provider-turn.updated') THEN json_extract(event.payload_json, '$.id') WHEN event.application_event_version = 2 AND event.event_type = 'thread.pull-request-link-updated' @@ -286,6 +290,12 @@ export const layer: Layer.Layer< json_extract(event.payload_json, '$.number')) ELSE NULL END AS entity_id, + CASE + WHEN event.application_event_version = 2 + AND event.event_type = 'provider-turn.updated' + THEN json_extract(event.payload_json, '$.tokenUsage') IS NOT NULL + ELSE 0 + END AS carries_token_usage, CASE WHEN event.application_event_version = 1 AND event.aggregate_kind = 'thread' THEN EXISTS ( @@ -330,11 +340,18 @@ export const layer: Layer.Layer< } else if ( row.event_type === "thread.pull-request-link-updated" || row.event_type === "message.updated" || - row.event_type === "node.updated" + row.event_type === "node.updated" || + row.event_type === "provider-turn.updated" ) { const key = encodeEntityKey([row.event_type, row.stream_id, row.entity_id]); const candidate = oldestCandidateByEntityKey.get(key); - if (candidate === undefined) oldestCandidateByEntityKey.set(key, null); + const retainsUsage = + row.carries_token_usage === 1 && !providerTurnKeysWithRetainedUsage.has(key); + if (retainsUsage) { + providerTurnKeysWithRetainedUsage.add(key); + if (candidate !== undefined && candidate !== null) obsolete.push(candidate); + oldestCandidateByEntityKey.set(key, null); + } else if (candidate === undefined) oldestCandidateByEntityKey.set(key, null); else { if (candidate !== null) obsolete.push(candidate); oldestCandidateByEntityKey.set(key, row.sequence);