diff --git a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts index 770da1939848..c38c35d0b826 100644 --- a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts +++ b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts @@ -21,6 +21,7 @@ import { ProviderInstanceId, ProviderSessionId, ProviderThreadId, + ProviderTurnId, RunAttemptId, RunId, ThreadId, @@ -924,6 +925,31 @@ 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, + providerTurnId = "provider-turn:foundation-compact", + ): OrchestrationV2DomainEvent => ({ + id: EventId.make(`event:foundation-compact:${suffix}`), + type: "provider-turn.updated", + threadId, + occurredAt: now, + payload: { + id: ProviderTurnId.make(providerTurnId), + providerThreadId: ProviderThreadId.make("provider-thread:foundation-compact"), + nodeId, + runAttemptId: null, + nativeTurnRef: null, + ordinal: providerTurnId === "provider-turn:foundation-compact" ? 1 : 2, + 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 }), @@ -933,6 +959,7 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { nodeEvent("node-1", "running"), itemEvent("item-1", "streaming"), itemEvent("item-middle", "still streaming"), + providerTurnEvent("turn-1", null), ...Array.from({ length: 501 }, (_, index) => threadStateEvent(`history-${index}`, "thread.visited"), ), @@ -940,8 +967,18 @@ 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("other-turn", 30, "provider-turn:foundation-compact:other"), ], }); + const providerCommandId = CommandId.make("command:foundation-compact:provider-turn"); + yield* eventSink.write({ + commandId: providerCommandId, + events: [providerTurnEvent("turn-commanded", 15)], + }); + yield* eventSink.write({ + events: [providerTurnEvent("turn-3", 20), providerTurnEvent("turn-4", null)], + }); const beforeCompaction = yield* projections.getThreadProjection(threadId); const importedV1ThreadId = "thread:foundation-compact-v1-imported"; @@ -980,7 +1017,7 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { `; const summary = yield* maintenance.compactEventStore(); - assert.isAtLeast(summary.deletedEventCount, 504); + assert.isAtLeast(summary.deletedEventCount, 505); assert.isAtLeast(summary.deletedReceiptCount, 1); const remaining = yield* sql<{ readonly event_id: string }>` @@ -998,10 +1035,15 @@ 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:other-turn", + "event:foundation-compact:turn-commanded", + "event:foundation-compact:turn-3", + "event:foundation-compact:turn-4", ], ); @@ -1026,6 +1068,13 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { ["command:foundation-compact:pending"], ); + const commandEvents = yield* eventSink + .readByCommandId({ commandId: providerCommandId }) + .pipe(Stream.runCollect); + assert.deepEqual( + commandEvents.map((stored) => stored.event.id), + ["event:foundation-compact:turn-commanded"], + ); assert.isTrue((yield* maintenance.verify).valid); assert.isTrue((yield* maintenance.rebuild).valid); assert.deepEqual(yield* projections.getThreadProjection(threadId), beforeCompaction); diff --git a/apps/server/src/orchestration-v2/ProjectionMaintenance.ts b/apps/server/src/orchestration-v2/ProjectionMaintenance.ts index aad9d355924b..2641f9c7d52e 100644 --- a/apps/server/src/orchestration-v2/ProjectionMaintenance.ts +++ b/apps/server/src/orchestration-v2/ProjectionMaintenance.ts @@ -245,6 +245,7 @@ export const layer: Layer.Layer< const retainedReadStateThreadIds = new Set(); // Retain each key's oldest candidate, deleting it only when an older event appears. const oldestCandidateByEntityKey = new Map(); + const providerTurnKeysWithRetainedUsage = new Set(); let deletedEventCount = 0; let deletedReceiptCount = 0; @@ -257,6 +258,7 @@ export const layer: Layer.Layer< readonly event_type: string; readonly entity_id: string | null; readonly imported_legacy_thread: number; + readonly carries_token_usage: number; }>` SELECT event.sequence, @@ -268,6 +270,10 @@ export const layer: Layer.Layer< WHEN event.application_event_version = 2 AND event.event_type IN ('message.updated', 'node.updated') THEN json_extract(event.payload_json, '$.id') + WHEN event.application_event_version = 2 + AND event.event_type = 'provider-turn.updated' + AND event.command_id IS NULL + THEN json_extract(event.payload_json, '$.id') WHEN event.application_event_version = 2 AND event.event_type = 'turn-item.updated' AND event.command_id IS NULL @@ -280,6 +286,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 ( @@ -326,11 +338,18 @@ export const layer: Layer.Layer< row.event_type === "thread.pull-request-link-updated" || row.event_type === "message.updated" || row.event_type === "node.updated" || + (row.event_type === "provider-turn.updated" && row.entity_id !== null) || (row.event_type === "turn-item.updated" && row.entity_id !== null) ) { 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);