From 118638a2caf48bb290ec55ff927710ed35703e36 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Mon, 26 Oct 2026 10:39:34 +0530 Subject: [PATCH 01/21] Rebuild events stream for subscribe on paused turns --- .../trueforge-core/src/agent-session/index.ts | 4 + .../src/agent-session/store/ISessionStore.ts | 40 ++++ .../store/InMemorySessionStore.ts | 56 ++++++ .../agent-session/store/turnEventSequence.ts | 8 + .../agent-session/store/storeContractSuite.ts | 176 ++++++++++++++++- packages/trueforge/.env.example | 2 + packages/trueforge/src/apis/schedules.ts | 4 +- packages/trueforge/src/apis/turns.ts | 108 ++++++++-- packages/trueforge/src/config.ts | 11 ++ ...24_000001_session_event_sequence_number.ts | 18 ++ .../session-store/PostgresSessionStore.ts | 16 ++ .../postgres/session-store/queries/events.ts | 124 ++++++++++++ packages/trueforge/src/db/postgres/types.ts | 5 + ...24_000001_session_event_sequence_number.ts | 20 ++ .../session-store/SqliteSessionStore.ts | 16 ++ .../db/sqlite/session-store/queries/events.ts | 122 ++++++++++++ packages/trueforge/src/db/sqlite/types.ts | 5 + packages/trueforge/src/main.ts | 1 + packages/trueforge/src/runtime/activeTurns.ts | 36 +++- .../runtime/event-subscription/inMemory.ts | 159 ++++++++++++--- .../src/runtime/event-subscription/index.ts | 39 +++- .../src/runtime/event-subscription/redis.ts | 185 ++++++++++++++++-- .../src/runtime/event-subscription/repair.ts | 153 +++++++++++++++ .../src/runtime/turnEventLogRepair.ts | 43 ++++ .../event-subscription/contractSuite.ts | 113 ++++++++++- .../turnEventLogRepair.test.ts | 127 ++++++++++++ .../tests/unit/runtime/activeTurns.test.ts | 70 ++++++- 27 files changed, 1588 insertions(+), 73 deletions(-) create mode 100644 packages/trueforge-core/src/agent-session/store/turnEventSequence.ts create mode 100644 packages/trueforge/src/db/postgres/migrations/20261024_000001_session_event_sequence_number.ts create mode 100644 packages/trueforge/src/db/sqlite/migrations/20261024_000001_session_event_sequence_number.ts create mode 100644 packages/trueforge/src/runtime/event-subscription/repair.ts create mode 100644 packages/trueforge/src/runtime/turnEventLogRepair.ts create mode 100644 packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts diff --git a/packages/trueforge-core/src/agent-session/index.ts b/packages/trueforge-core/src/agent-session/index.ts index 6fda6fc6d4..264fe5da4e 100644 --- a/packages/trueforge-core/src/agent-session/index.ts +++ b/packages/trueforge-core/src/agent-session/index.ts @@ -86,6 +86,7 @@ export type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, + AssignTurnEventSequenceNumbersInput, ClaimTurnExecutorInput, CreateSessionInput, CreateTurnInput, @@ -99,6 +100,7 @@ export type { InsertTurnInboundEventsInput, ListSessionEventsInput, ListSessionsInput, + ListTurnEventLogInput, ListTurnEventsInput, ListTurnsInput, ListUnconsumedTurnInboundEventsInput, @@ -110,6 +112,8 @@ export type { PatchThreadCapabilityStateInput, RemoveThreadsInput, TurnContextAppend, + TurnEventLogEntry, + TurnEventSequenceAssignment, TurnRecordWithoutSnapshot, UpdateSessionInput, UpdateTurnNonTerminalStateInput, diff --git a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts index 7502af8233..7554784d59 100644 --- a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts @@ -185,6 +185,32 @@ export interface AppendToEventsInput extends TurnProgressKeys { events: PersistedTurnEvent[]; } +/** One durable event plus the stream sequence assigned when it was first appended. */ +export interface TurnEventLogEntry { + event: PersistedTurnEvent; + /** Null until a producer streams the event and records the number. */ + sequence_number: number | null; +} + +export interface ListTurnEventLogInput { + session_id: string; + turn_id: string; +} + +/** Sequence written onto an existing event row. The event payload is not rewritten. */ +export interface TurnEventSequenceAssignment { + event_id: string; + sequence_number: number; +} + +export interface AssignTurnEventSequenceNumbersInput extends TurnProgressKeys { + /** + * Empty still checks `active_executor_id`. Producer repair uses that so a stale + * executor fails before it claims or appends the live stream. + */ + assignments: TurnEventSequenceAssignment[]; +} + export interface InsertTurnInboundEventsInput extends TurnProgressKeys { /** Empty array is a no-op. */ events: TurnUserEvent[]; @@ -400,6 +426,20 @@ export interface ISessionStore< /** Appends events while the caller owns a `running` or `paused` turn; terminal turns are immutable. */ appendToEvents(input: AppendToEventsInput, transaction?: TTransaction): Promise; + /** + * Full turn event log in persist order (`event_id` ascending), with the stream + * sequence stored beside each payload. Missing turn throws TurnNotFoundError. + */ + listTurnEventLog(input: ListTurnEventLogInput): Promise; + + /** + * Records sequence numbers for events that do not have one yet. The caller must + * own a non-terminal turn. Already-stored numbers that match are a no-op; a + * different stored number is an invariant failure. Unknown `event_id` is too. + * An empty assignment still fences on `active_executor_id`. + */ + assignTurnEventSequenceNumbers(input: AssignTurnEventSequenceNumbersInput, transaction?: TTransaction): Promise; + /** * Durable inbound send-event inbox for a tip the caller still owns. Tip must be * non-terminal (`running` or `paused`) — terminal tip → diff --git a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts index 93e606dfde..0c0b276d64 100644 --- a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts @@ -12,6 +12,7 @@ import type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, + AssignTurnEventSequenceNumbersInput, ClaimTurnExecutorInput, CreateSessionInput, CreateTurnInput, @@ -25,6 +26,7 @@ import type { ISessionStore, ListSessionEventsInput, ListSessionsInput, + ListTurnEventLogInput, ListTurnEventsInput, ListTurnsInput, ListUnconsumedTurnInboundEventsInput, @@ -37,6 +39,7 @@ import type { PatchThreadsMCPAuthInput, RemoveThreadsInput, TurnContextAppend, + TurnEventLogEntry, TurnRecordWithoutSnapshot, UpdateSessionInput, UpdateTurnNonTerminalStateInput, @@ -61,6 +64,7 @@ import { TurnNotRunningError, TurnOwnershipLostError, } from './SessionStoreErrors'; +import { assertTurnEventSequenceNumber } from './turnEventSequence'; /* eslint-disable @typescript-eslint/require-await -- in-memory store is synchronous; methods stay async so thrown SessionStore*Error reject as Promises for ISessionStore callers */ @@ -195,6 +199,8 @@ export class InMemorySessionStore< private readonly sessions = new Map>(); private readonly turns = new Map>(); private readonly events = new Map(); + /** Stream sequence per event id, filled when the event is first streamed. */ + private readonly eventSequences = new Map>(); /** session_id:turn_id → inbound send-event inbox */ private readonly inboundEvents = new Map(); @@ -575,6 +581,56 @@ export class InMemorySessionStore< return; } + async listTurnEventLog(input: ListTurnEventLogInput): Promise { + this.requireTurn(input.session_id, input.turn_id); + const tKey = turnKey(input); + const list = this.events.get(tKey); + if (!list) { + throw new TurnNotFoundError(input.turn_id); + } + const sequences = this.eventSequences.get(tKey); + return [...list] + .sort((left, right) => left.id.localeCompare(right.id)) + .map(event => ({ + event: deepCopy(event), + sequence_number: sequences?.get(event.id) ?? null, + })); + } + + async assignTurnEventSequenceNumbers(input: AssignTurnEventSequenceNumbersInput): Promise { + this.requireOwnedTurn(input); + this.requireNonTerminalTurn(input); + if (input.assignments.length === 0) { + return; + } + const tKey = turnKey(input); + const list = this.events.get(tKey); + if (!list) { + throw new TurnNotFoundError(input.turn_id); + } + let sequences = this.eventSequences.get(tKey); + if (!sequences) { + sequences = new Map(); + this.eventSequences.set(tKey, sequences); + } + for (const assignment of input.assignments) { + assertTurnEventSequenceNumber(assignment.sequence_number); + if (!list.some(event => event.id === assignment.event_id)) { + throw new SessionStoreInvariantError( + `event ${assignment.event_id} is not in the log for turn ${input.turn_id}`, + ); + } + const current = sequences.get(assignment.event_id); + if (current === undefined) { + sequences.set(assignment.event_id, assignment.sequence_number); + } else if (current !== assignment.sequence_number) { + throw new SessionStoreInvariantError( + `event ${assignment.event_id} already has sequence_number ${String(current)}`, + ); + } + } + } + async patchThreadsMCPAuth(input: PatchThreadsMCPAuthInput): Promise { this.requireOwnedTurn(input); const turn = this.requireNonTerminalTurn(input); diff --git a/packages/trueforge-core/src/agent-session/store/turnEventSequence.ts b/packages/trueforge-core/src/agent-session/store/turnEventSequence.ts new file mode 100644 index 0000000000..fb10f88d46 --- /dev/null +++ b/packages/trueforge-core/src/agent-session/store/turnEventSequence.ts @@ -0,0 +1,8 @@ +import { SessionStoreInvariantError } from './SessionStoreErrors'; + +/** Stream sequences are positive safe integers. The SSE id is this number. */ +export function assertTurnEventSequenceNumber(sequenceNumber: number): void { + if (!Number.isSafeInteger(sequenceNumber) || sequenceNumber < 1) { + throw new SessionStoreInvariantError('sequence_number must be a positive integer'); + } +} diff --git a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts index e99ab30d2f..bf185106da 100644 --- a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts +++ b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts @@ -2,7 +2,11 @@ import { z } from 'zod'; import { MAIN_THREAD_ID } from '../../../src/agent-session/models/TurnRecord'; import type { PersistedTurnEvent } from '../../../src/agent-session/schemas/events'; import { EventType } from '../../../src/agent-session/schemas/events'; -import { CancellationReason, type NonTerminalTurnState } from '../../../src/agent-session/schemas/turn'; +import { + CancellationReason, + type NonTerminalTurnState, + type TerminalTurnState, +} from '../../../src/agent-session/schemas/turn'; import type { ISessionStore } from '../../../src/agent-session/store/ISessionStore'; import { decodeSessionEventPageToken } from '../../../src/agent-session/store/SessionEventPageToken'; import { @@ -1439,6 +1443,17 @@ export function runStoreContractSuite( }); } + function terminalState(status: 'done' | 'cancelled' | 'error'): TerminalTurnState { + switch (status) { + case 'done': + return makeDoneTurnState(); + case 'cancelled': + return makeCancelledTurnState(CancellationReason.ClientCancelled); + case 'error': + return { status: 'error', message: 'failed', completed_at: new Date().toISOString() }; + } + } + async function freezeTurn(store: Store, turnId: string) { await store.freezeAndGetTurn({ session_id: sessionId, @@ -1507,6 +1522,95 @@ export function runStoreContractSuite( expect(await lastTurnDetails(store)).toEqual({ id: 'turn-3', state: { status: 'running' } }); }); + it.each(['done', 'cancelled', 'error'] as const)( + 'a late non-terminal update does not overwrite terminal %s', + async status => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + const terminal = terminalState(status); + await store.updateTurnTerminalState({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: 'turn-1', + state: terminal, + turn_done_event: makeTurnDoneEvent(terminal), + }); + const expected = { id: 'turn-1', state: { status } }; + + await expect(setNonTerminalState(store, 'turn-1', makePausedTurnState())).rejects.toBeInstanceOf( + TurnNotRunningError, + ); + await expect(setNonTerminalState(store, 'turn-1', { status: 'running' })).rejects.toBeInstanceOf( + TurnNotRunningError, + ); + expect(await lastTurnDetails(store)).toEqual(expected); + }, + ); + + it('a stale executor does not update last_turn_details, and the new owner can resume it', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + await setNonTerminalState(store, 'turn-1', makePausedTurnState()); + const won = await store.claimTurnExecutor({ + session_id: sessionId, + turn_id: 'turn-1', + expected_executor_id: TEST_ACTIVE_EXECUTOR_ID, + new_executor_id: 'executor-b', + }); + expect(won).toBe(true); + const paused = { id: 'turn-1', state: { status: 'paused' as const } }; + expect(await lastTurnDetails(store)).toEqual(paused); + + const runningState = { status: 'running' as const }; + await expect( + store.updateTurnNonTerminalState({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: 'turn-1', + state: runningState, + turn_update_event: makeTurnUpdateEvent(runningState), + }), + ).rejects.toBeInstanceOf(TurnOwnershipLostError); + const doneState = makeDoneTurnState(); + await expect( + store.updateTurnTerminalState({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: 'turn-1', + state: doneState, + turn_done_event: makeTurnDoneEvent(doneState), + }), + ).rejects.toBeInstanceOf(TurnOwnershipLostError); + expect(await lastTurnDetails(store)).toEqual(paused); + + await store.updateTurnNonTerminalState({ + active_executor_id: 'executor-b', + session_id: sessionId, + turn_id: 'turn-1', + state: runningState, + turn_update_event: makeTurnUpdateEvent(runningState), + }); + expect(await lastTurnDetails(store)).toEqual({ id: 'turn-1', state: { status: 'running' } }); + }); + + it('points at the new tip after a paused predecessor is frozen', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + await setNonTerminalState(store, 'turn-1', makePausedTurnState()); + expect(await lastTurnDetails(store)).toEqual({ id: 'turn-1', state: { status: 'paused' } }); + + await freezeTurn(store, 'turn-1'); + expect(await lastTurnDetails(store)).toEqual({ id: 'turn-1', state: { status: 'cancelled' } }); + + await store.createTurn( + makeCreateTurnInput({ sessionId, turnId: 'turn-2', previousTurnId: 'turn-1', firstTurnId: 'turn-1' }), + ); + expect(await lastTurnDetails(store)).toEqual({ id: 'turn-2', state: { status: 'running' } }); + }); + it('a superseded turn finishing still adds its cost and duration to session metrics', async () => { const store = createStore(); await seedSession(store); @@ -4415,6 +4519,76 @@ export function runStoreContractSuite( }); } + describe('turn event sequence numbers', () => { + it('stores sequence_number beside the payload and rejects a stale owner', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + const created = makeTurnCreatedEvent('turn-1'); + const message = makeModelMessageEvent(); + await store.appendToEvents({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: 'turn-1', + events: [created, message], + }); + + const unsequenced = await store.listTurnEventLog({ session_id: sessionId, turn_id: 'turn-1' }); + expect(unsequenced.map(row => row.sequence_number)).toEqual([null, null]); + + await expect( + store.assignTurnEventSequenceNumbers({ + active_executor_id: 'other-executor', + session_id: sessionId, + turn_id: 'turn-1', + assignments: [{ event_id: created.id, sequence_number: 1 }], + }), + ).rejects.toBeInstanceOf(TurnOwnershipLostError); + expect( + (await store.listTurnEventLog({ session_id: sessionId, turn_id: 'turn-1' })).map(row => row.sequence_number), + ).toEqual([null, null]); + + await store.assignTurnEventSequenceNumbers({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: 'turn-1', + assignments: [ + { event_id: created.id, sequence_number: 1 }, + { event_id: message.id, sequence_number: 2 }, + ], + }); + await store.assignTurnEventSequenceNumbers({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: 'turn-1', + assignments: [{ event_id: created.id, sequence_number: 1 }], + }); + + const sequenced = await store.listTurnEventLog({ session_id: sessionId, turn_id: 'turn-1' }); + expect(sequenced.map(row => ({ id: row.event.id, sequence_number: row.sequence_number }))).toEqual([ + { id: created.id, sequence_number: 1 }, + { id: message.id, sequence_number: 2 }, + ]); + const listed = await store.listTurnEvents({ + session_id: sessionId, + turn_id: 'turn-1', + limit: 10, + page_token: undefined, + order: 'asc', + }); + expect(listed.data[0]).not.toHaveProperty('sequence_number'); + + await expect( + store.assignTurnEventSequenceNumbers({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: 'turn-1', + assignments: [{ event_id: created.id, sequence_number: 9 }], + }), + ).rejects.toBeInstanceOf(SessionStoreInvariantError); + }); + }); + describe('no SSE surface', () => { it('ISessionStore has no subscribe* members', () => { const store = createStore(); diff --git a/packages/trueforge/.env.example b/packages/trueforge/.env.example index afa73b2351..8d3e999e1d 100644 --- a/packages/trueforge/.env.example +++ b/packages/trueforge/.env.example @@ -117,6 +117,8 @@ REDIS_URL=redis://localhost:6379 # TURN_SUBSCRIBE_TIMEOUT_MS=600000 ## Seconds a paused turn stays in memory before offload. The row stays paused. Default 300. # TURN_PAUSED_KEEPALIVE_SECONDS=300 +## Max ms to wait for the per-turn lock. 0 waits until the holder releases it. Default 0. +# TURN_LOCK_ACQUIRE_TIMEOUT_MS=0 ## Max ms for one create-turn SSE write; after this the client is dropped but the turn keeps running. Default 60000. # TURN_SSE_STREAM_WRITE_TIMEOUT_MS=60000 diff --git a/packages/trueforge/src/apis/schedules.ts b/packages/trueforge/src/apis/schedules.ts index afbb51454c..bd3910ede9 100644 --- a/packages/trueforge/src/apis/schedules.ts +++ b/packages/trueforge/src/apis/schedules.ts @@ -86,8 +86,8 @@ export interface ScheduleTurnExecutionDeps { sandboxProviderStore: ISandboxProviderStore; scheduleStore: IScheduleStore; sessions: Sessions; - /** Previous-turn cancel when a schedule run supersedes a turn that is not in this process. */ - sessionStore: Pick, 'getTurn'>; + /** Previous-turn cancel, and sequence numbers for the turn event stream. */ + sessionStore: Pick, 'getTurn' | 'listTurnEventLog' | 'assignTurnEventSequenceNumbers'>; webSearchProviderStore: IWebSearchProviderStore; } diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 01a4c0f4b9..763ac5d6f1 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -80,7 +80,12 @@ import { subscribeTurnRoute, } from '../routes/turnRoutes'; import type { ActiveTurnRegistry } from '../runtime/activeTurns'; -import { StreamGoneError, type EventSubscription, type EventSubscriptionRegistry } from '../runtime/event-subscription'; +import { + StreamGoneError, + StreamOwnerMismatchError, + type EventSubscription, + type EventSubscriptionRegistry, +} from '../runtime/event-subscription'; import { validateSandboxFilePath } from '../runtime/sandboxFilePath'; import { buildTurnSandbox, @@ -89,6 +94,7 @@ import { resolveSandboxEnvironment, resolveSandboxProvider, } from '../runtime/sessionResources'; +import { repairTurnEventLog } from '../runtime/turnEventLogRepair'; import { DEFAULT_SANDBOX_ENVIRONMENT_NAME } from '../schemas/sandboxEnvironment'; import { MAX_SESSION_TITLE_LENGTH } from '../schemas/session'; import { assertGatewayMetadataRequestHeaders } from '../truefoundry/gatewayMetadata'; @@ -216,8 +222,10 @@ export type BeginTurnExecutionDeps = Pick< sandboxProviderStore: ISandboxProviderStore; webSearchProviderStore: IWebSearchProviderStore; skillStore: Pick; - /** Loads the previous turn when create must cancel a run that is not in this process. */ - sessionStore: Pick; + /** + * Previous-turn cancel, durable sequence numbers, and rebuilding a missing event stream. + */ + sessionStore: Pick; redis?: RedisClient | undefined; }; @@ -499,21 +507,54 @@ interface TurnExecutionTimeout { clear: () => void; } +async function recordPersistedEventSequence(input: { + sessionStore: Pick; + sessionId: string; + turnId: string; + executorId: string; + event: TurnStreamingEvent; + sequenceNumber: number; +}): Promise { + if (input.event.type === EventType.MODEL_MESSAGE_DELTA) { + return; + } + await input.sessionStore.assignTurnEventSequenceNumbers({ + session_id: input.sessionId, + turn_id: input.turnId, + active_executor_id: input.executorId, + assignments: [{ event_id: input.event.id, sequence_number: input.sequenceNumber }], + }); +} + export async function drainTurnEvents(input: { trackedStream: AsyncIterable; turnEventStream: EventSubscription; executionTimeout: TurnExecutionTimeout; logger: Logger; pausedKeepalive: TurnPausedKeepalive; + executorId: string; + /** Writes the sequence put just assigned onto the durable event row. */ + recordSequence: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; onEvent?: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; }): Promise { - const { trackedStream, turnEventStream, executionTimeout, logger, pausedKeepalive, onEvent } = input; + const { + trackedStream, + turnEventStream, + executionTimeout, + logger, + pausedKeepalive, + executorId, + recordSequence, + onEvent, + } = input; try { for await (const event of trackedStream) { // Dual-write before any client sink so subscribers can resume after disconnect. const sequenceNumber = await turnEventStream.put(event, { streamTTLSeconds: streamTTLSecondsFor(event), + executorId, }); + await recordSequence(event, sequenceNumber); if (event.type === EventType.TURN_UPDATE) { if (event.state.status === 'paused') { pausedKeepalive.start(); @@ -526,8 +567,15 @@ export async function drainTurnEvents(input: { await onEvent?.(event, sequenceNumber); } } catch (error) { - if (error instanceof SessionStoreNotFoundError) { - logger.warn('Turn stream ended after session/turn was removed', extractErrorLogFields(error)); + if ( + error instanceof SessionStoreNotFoundError || + error instanceof TurnOwnershipLostError || + error instanceof StreamOwnerMismatchError + ) { + logger.warn( + 'Turn stream ended after session/turn was removed or ownership was lost', + extractErrorLogFields(error), + ); } else { logger.error('Unexpected error in turn event drain', extractErrorLogFields(error)); } @@ -543,6 +591,8 @@ export interface TurnEventDrainInput { executionTimeout: TurnExecutionTimeout; logger: Logger; pausedKeepalive: TurnPausedKeepalive; + executorId: string; + recordSequence: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; } /** @@ -665,7 +715,7 @@ export async function beginTurnExecution( turn, }); - // Held for the whole turn; the stream's sequence counter dies with it. + // Held for the whole turn. Repair reseeds the counter when this replica resumes. const turnEventStream = deps.eventSubscriptions.get(turnStreamId(tenantId, sessionId, turn.id)); return { @@ -683,6 +733,16 @@ export async function beginTurnExecution( deps.activeTurns.clearPausedKeepalive({ sessionId, turnId: turn.id }); }, }, + executorId: configuration.EXECUTOR_ID, + recordSequence: (event, sequenceNumber) => + recordPersistedEventSequence({ + sessionStore: deps.sessionStore, + sessionId, + turnId: turn.id, + executorId: configuration.EXECUTOR_ID, + event, + sequenceNumber, + }), }, }; } @@ -724,8 +784,8 @@ function cancellationReason(reason: unknown): CancellationReason { } /** - * Rebuild a paused turn this replica owns, reseed the event-stream counter from - * the existing tail, and drain in the background. Inbox replay happens inside rebuild. + * Rebuild a paused turn this replica owns, repair the event stream from the + * durable log, and drain in the background. Inbox replay happens inside rebuild. */ async function resumePausedTurn(input: { session: SessionHandle; @@ -792,8 +852,14 @@ async function resumePausedTurn(input: { turnStreamId(input.session.tenant_id, input.session.session_id, turn.id), ); try { - // Counter only. A missing stream stays at 1; log repair from the TurnRecord is later. - await turnEventStream.resumeProducer({ streamTTLSeconds: configuration.TURN_STREAM_TTL_SECONDS }); + await repairTurnEventLog({ + sessionStore: input.deps.sessionStore, + subscription: turnEventStream, + sessionId: input.session.session_id, + turnId: turn.id, + streamTTLSeconds: configuration.TURN_STREAM_TTL_SECONDS, + producerExecutorId: configuration.EXECUTOR_ID, + }); } catch (error) { input.deps.logger.warn('Paused turn stream resume failed', extractErrorLogFields(error)); executionTimeout.clear(); @@ -832,6 +898,16 @@ async function resumePausedTurn(input: { turnEventStream, executionTimeout, logger: input.deps.logger, + executorId: configuration.EXECUTOR_ID, + recordSequence: (event, sequenceNumber) => + recordPersistedEventSequence({ + sessionStore: input.deps.sessionStore, + sessionId: input.session.session_id, + turnId: turn.id, + executorId: configuration.EXECUTOR_ID, + event, + sequenceNumber, + }), pausedKeepalive: { start: () => { input.deps.activeTurns.startPausedKeepalive({ @@ -1950,8 +2026,16 @@ export function createTurnsRouter(deps: TurnsRouterDeps): Promise { + await sql` + SET LOCAL lock_timeout = '5s'; + ALTER TABLE session_event + ADD COLUMN sequence_number integer; + `.execute(db); +} + +export async function down(db: Kysely): Promise { + await sql` + SET LOCAL lock_timeout = '5s'; + ALTER TABLE session_event + DROP COLUMN IF EXISTS sequence_number; + `.execute(db); +} diff --git a/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts b/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts index dc08267e73..85559ba84b 100644 --- a/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts +++ b/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts @@ -7,6 +7,7 @@ import type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, + AssignTurnEventSequenceNumbersInput, ClaimTurnExecutorInput, CreateSessionInput, CreateTurnInput, @@ -20,6 +21,7 @@ import type { ISessionStore, ListSessionEventsInput, ListSessionsInput, + ListTurnEventLogInput, ListTurnEventsInput, ListTurnsInput, ListUnconsumedTurnInboundEventsInput, @@ -30,6 +32,7 @@ import type { PatchThreadCapabilityStateInput, PatchThreadsMCPAuthInput, RemoveThreadsInput, + TurnEventLogEntry, TurnRecordWithoutSnapshot, UpdateSessionInput, UpdateTurnNonTerminalStateInput, @@ -44,7 +47,9 @@ import type { Database } from '../types'; import { patchThreadCapabilityState as patchThreadCapabilityStateQuery } from './queries/capabilities'; import { appendToEvents as appendToEventsQuery, + assignTurnEventSequenceNumbers as assignTurnEventSequenceNumbersQuery, listSessionEvents as listSessionEventsQuery, + listTurnEventLog as listTurnEventLogQuery, listTurnEvents as listTurnEventsQuery, } from './queries/events'; import { @@ -222,6 +227,17 @@ export class PostgresSessionStore implements ISessionStore return appendToEventsQuery(transaction ?? this.db, input); } + listTurnEventLog(input: ListTurnEventLogInput): Promise { + return listTurnEventLogQuery(this.db, input); + } + + assignTurnEventSequenceNumbers( + input: AssignTurnEventSequenceNumbersInput, + transaction?: Transaction, + ): Promise { + return assignTurnEventSequenceNumbersQuery(transaction ?? this.db, input); + } + insertTurnInboundEvents(input: InsertTurnInboundEventsInput): Promise { return insertTurnInboundEventsQuery(this.db, input); } diff --git a/packages/trueforge/src/db/postgres/session-store/queries/events.ts b/packages/trueforge/src/db/postgres/session-store/queries/events.ts index f33dfd0a83..5e547cd672 100644 --- a/packages/trueforge/src/db/postgres/session-store/queries/events.ts +++ b/packages/trueforge/src/db/postgres/session-store/queries/events.ts @@ -2,8 +2,11 @@ import type { PersistedTurnEvent, SessionEventItem } from '@truefoundry/trueforg import type { TokenPagination } from '@truefoundry/trueforge-core/agent-session/schemas/pagination'; import type { AppendToEventsInput, + AssignTurnEventSequenceNumbersInput, ListSessionEventsInput, + ListTurnEventLogInput, ListTurnEventsInput, + TurnEventLogEntry, } from '@truefoundry/trueforge-core/agent-session/store/ISessionStore'; import { decodeOffsetPageToken, @@ -16,10 +19,13 @@ import { } from '@truefoundry/trueforge-core/agent-session/store/SessionEventPageToken'; import { SessionNotFoundError, + SessionStoreInvariantError, TurnNotFoundError, } from '@truefoundry/trueforge-core/agent-session/store/SessionStoreErrors'; +import { assertTurnEventSequenceNumber } from '@truefoundry/trueforge-core/agent-session/store/turnEventSequence'; import type { Kysely, Transaction } from 'kysely'; import { sql } from 'kysely'; +import { isTransaction } from '../../../transaction'; import { json } from '../../sqlExpressions'; import type { Database } from '../../types'; import { unnestWithOrdinality, values } from '../sqlExpressions'; @@ -70,6 +76,124 @@ export async function appendToEvents( } } +export async function listTurnEventLog( + db: Kysely, + input: ListTurnEventLogInput, +): Promise { + const turn = await db + .selectFrom('turn') + .select('turn_id') + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .executeTakeFirst(); + if (!turn) { + throw new TurnNotFoundError(input.turn_id); + } + + const rows = await db + .selectFrom('session_event') + .select(['event', 'sequence_number']) + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .orderBy('event_id', 'asc') + .execute(); + + return rows.map(row => ({ + event: row.event, + sequence_number: row.sequence_number, + })); +} + +export async function assignTurnEventSequenceNumbers( + db: Kysely | Transaction, + input: AssignTurnEventSequenceNumbersInput, +): Promise { + if (isTransaction(db)) { + await assignTurnEventSequenceNumbersInTransaction(db, input); + return; + } + await db.transaction().execute(trx => assignTurnEventSequenceNumbersInTransaction(trx, input)); +} + +async function assignTurnEventSequenceNumbersInTransaction( + db: Transaction, + input: AssignTurnEventSequenceNumbersInput, +): Promise { + for (const assignment of input.assignments) { + assertTurnEventSequenceNumber(assignment.sequence_number); + } + + const keys = { + session_id: input.session_id, + turn_id: input.turn_id, + active_executor_id: input.active_executor_id, + }; + const fence = await db + .selectFrom('turn') + .select(sql`1`.as('one')) + .where('session_id', '=', keys.session_id) + .where('turn_id', '=', keys.turn_id) + .where('active_executor_id', '=', keys.active_executor_id) + .where(sql`state->>'status' IN ('running', 'paused')`) + .forShare() + .executeTakeFirst(); + if (!fence) { + await classifyTurnFenceWriteFailure(db, keys); + } + if (input.assignments.length === 0) { + return; + } + + const rows = await db + .selectFrom('session_event') + .select(['event_id', 'sequence_number']) + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .where( + 'event_id', + 'in', + input.assignments.map(assignment => assignment.event_id), + ) + .execute(); + const stored = new Map(rows.map(row => [row.event_id, row.sequence_number])); + + for (const assignment of input.assignments) { + if (!stored.has(assignment.event_id)) { + throw new SessionStoreInvariantError(`event ${assignment.event_id} is not in the log for turn ${input.turn_id}`); + } + const current = stored.get(assignment.event_id); + if (current === null || current === undefined) { + const updated = await db + .updateTable('session_event') + .set({ sequence_number: assignment.sequence_number }) + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .where('event_id', '=', assignment.event_id) + .where('sequence_number', 'is', null) + .executeTakeFirst(); + if (updated.numUpdatedRows === 0n) { + const again = await db + .selectFrom('session_event') + .select('sequence_number') + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .where('event_id', '=', assignment.event_id) + .executeTakeFirst(); + if (again?.sequence_number !== assignment.sequence_number) { + throw new SessionStoreInvariantError( + `event ${assignment.event_id} already has sequence_number ${String(again?.sequence_number ?? 'none')}`, + ); + } + } + stored.set(assignment.event_id, assignment.sequence_number); + } else if (current !== assignment.sequence_number) { + throw new SessionStoreInvariantError( + `event ${assignment.event_id} already has sequence_number ${String(current)}`, + ); + } + } +} + /** * listTurnEvents — one statement: drive from turn, left-join a page of events. * Missing turn → 0 rows → TurnNotFoundError. Empty log → one null-event sentinel. diff --git a/packages/trueforge/src/db/postgres/types.ts b/packages/trueforge/src/db/postgres/types.ts index ed62e68466..0f2cfe6ea2 100644 --- a/packages/trueforge/src/db/postgres/types.ts +++ b/packages/trueforge/src/db/postgres/types.ts @@ -261,6 +261,11 @@ export interface SessionEventTable { * Ordering still uses event_id. */ created_at: Date; + /** + * Stream sequence assigned when the event is first appended to the live log. + * Null until then. The event jsonb stays the client payload. + */ + sequence_number: ColumnType; } /** diff --git a/packages/trueforge/src/db/sqlite/migrations/20261024_000001_session_event_sequence_number.ts b/packages/trueforge/src/db/sqlite/migrations/20261024_000001_session_event_sequence_number.ts new file mode 100644 index 0000000000..c423ed01a5 --- /dev/null +++ b/packages/trueforge/src/db/sqlite/migrations/20261024_000001_session_event_sequence_number.ts @@ -0,0 +1,20 @@ +import { sql, type Kysely } from 'kysely'; + +/** + * Stream sequence lives beside the event payload so a missing stream can be rebuilt. + * Mirrors db/postgres/migrations/20261024_000001_session_event_sequence_number.ts. + */ +export async function up(db: Kysely): Promise { + await db.transaction().execute(async trx => { + await sql` + ALTER TABLE session_event + ADD COLUMN sequence_number INTEGER + `.execute(trx); + }); +} + +export async function down(db: Kysely): Promise { + await db.transaction().execute(async trx => { + await sql`ALTER TABLE session_event DROP COLUMN sequence_number`.execute(trx); + }); +} diff --git a/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts b/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts index b6798bc359..086630d06c 100644 --- a/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts +++ b/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts @@ -7,6 +7,7 @@ import type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, + AssignTurnEventSequenceNumbersInput, ClaimTurnExecutorInput, CreateSessionInput, CreateTurnInput, @@ -20,6 +21,7 @@ import type { ISessionStore, ListSessionEventsInput, ListSessionsInput, + ListTurnEventLogInput, ListTurnEventsInput, ListTurnsInput, ListUnconsumedTurnInboundEventsInput, @@ -30,6 +32,7 @@ import type { PatchThreadCapabilityStateInput, PatchThreadsMCPAuthInput, RemoveThreadsInput, + TurnEventLogEntry, TurnRecordWithoutSnapshot, UpdateSessionInput, UpdateTurnNonTerminalStateInput, @@ -44,7 +47,9 @@ import type { Database } from '../types'; import { patchThreadCapabilityState as patchThreadCapabilityStateQuery } from './queries/capabilities'; import { appendToEvents as appendToEventsQuery, + assignTurnEventSequenceNumbers as assignTurnEventSequenceNumbersQuery, listSessionEvents as listSessionEventsQuery, + listTurnEventLog as listTurnEventLogQuery, listTurnEvents as listTurnEventsQuery, } from './queries/events'; import { @@ -220,6 +225,17 @@ export class SqliteSessionStore implements ISessionStore, return appendToEventsQuery(transaction ?? this.db, input); } + listTurnEventLog(input: ListTurnEventLogInput): Promise { + return listTurnEventLogQuery(this.db, input); + } + + assignTurnEventSequenceNumbers( + input: AssignTurnEventSequenceNumbersInput, + transaction?: Transaction, + ): Promise { + return assignTurnEventSequenceNumbersQuery(transaction ?? this.db, input); + } + insertTurnInboundEvents(input: InsertTurnInboundEventsInput): Promise { return insertTurnInboundEventsQuery(this.db, input); } diff --git a/packages/trueforge/src/db/sqlite/session-store/queries/events.ts b/packages/trueforge/src/db/sqlite/session-store/queries/events.ts index a53874d80f..05a906e569 100644 --- a/packages/trueforge/src/db/sqlite/session-store/queries/events.ts +++ b/packages/trueforge/src/db/sqlite/session-store/queries/events.ts @@ -2,8 +2,11 @@ import type { PersistedTurnEvent, SessionEventItem } from '@truefoundry/trueforg import type { TokenPagination } from '@truefoundry/trueforge-core/agent-session/schemas/pagination'; import type { AppendToEventsInput, + AssignTurnEventSequenceNumbersInput, ListSessionEventsInput, + ListTurnEventLogInput, ListTurnEventsInput, + TurnEventLogEntry, } from '@truefoundry/trueforge-core/agent-session/store/ISessionStore'; import { decodeOffsetPageToken, @@ -16,8 +19,10 @@ import { } from '@truefoundry/trueforge-core/agent-session/store/SessionEventPageToken'; import { SessionNotFoundError, + SessionStoreInvariantError, TurnNotFoundError, } from '@truefoundry/trueforge-core/agent-session/store/SessionStoreErrors'; +import { assertTurnEventSequenceNumber } from '@truefoundry/trueforge-core/agent-session/store/turnEventSequence'; import { sql, type Kysely, type Transaction } from 'kysely'; import { isTransaction } from '../../../transaction'; import { jsonbBind, jsonText } from '../../sqlExpressions'; @@ -69,6 +74,123 @@ async function appendToEventsInTransaction(db: Transaction, input: App await db.insertInto('session_event').values(eventRows).execute(); } +export async function listTurnEventLog( + db: Kysely, + input: ListTurnEventLogInput, +): Promise { + const turn = await db + .selectFrom('turn') + .select('turn_id') + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .executeTakeFirst(); + if (!turn) { + throw new TurnNotFoundError(input.turn_id); + } + + const rows = await db + .selectFrom('session_event') + .select(['sequence_number', jsonText(sql.ref('event')).as('event')]) + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .orderBy('event_id', 'asc') + .execute(); + + return rows.map(row => ({ + event: row.event, + sequence_number: row.sequence_number, + })); +} + +export async function assignTurnEventSequenceNumbers( + db: DbOrTrx, + input: AssignTurnEventSequenceNumbersInput, +): Promise { + if (isTransaction(db)) { + await assignTurnEventSequenceNumbersInTransaction(db, input); + return; + } + await db.transaction().execute(trx => assignTurnEventSequenceNumbersInTransaction(trx, input)); +} + +async function assignTurnEventSequenceNumbersInTransaction( + db: Transaction, + input: AssignTurnEventSequenceNumbersInput, +): Promise { + for (const assignment of input.assignments) { + assertTurnEventSequenceNumber(assignment.sequence_number); + } + + const keys: TurnKeys = { + session_id: input.session_id, + turn_id: input.turn_id, + active_executor_id: input.active_executor_id, + }; + const fenceRow = await db + .selectFrom('turn') + .select(sql`1`.as('one')) + .where('session_id', '=', keys.session_id) + .where('turn_id', '=', keys.turn_id) + .where('active_executor_id', '=', keys.active_executor_id) + .where(sql`state->>'status' IN ('running', 'paused')`) + .executeTakeFirst(); + if (!fenceRow) { + await classifyTurnFenceWriteFailure(db, keys); + } + if (input.assignments.length === 0) { + return; + } + + const rows = await db + .selectFrom('session_event') + .select(['event_id', 'sequence_number']) + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .where( + 'event_id', + 'in', + input.assignments.map(assignment => assignment.event_id), + ) + .execute(); + const stored = new Map(rows.map(row => [row.event_id, row.sequence_number])); + + for (const assignment of input.assignments) { + if (!stored.has(assignment.event_id)) { + throw new SessionStoreInvariantError(`event ${assignment.event_id} is not in the log for turn ${input.turn_id}`); + } + const current = stored.get(assignment.event_id); + if (current === null || current === undefined) { + const updated = await db + .updateTable('session_event') + .set({ sequence_number: assignment.sequence_number }) + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .where('event_id', '=', assignment.event_id) + .where('sequence_number', 'is', null) + .executeTakeFirst(); + if (updated.numUpdatedRows === 0n) { + const again = await db + .selectFrom('session_event') + .select('sequence_number') + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .where('event_id', '=', assignment.event_id) + .executeTakeFirst(); + if (again?.sequence_number !== assignment.sequence_number) { + throw new SessionStoreInvariantError( + `event ${assignment.event_id} already has sequence_number ${String(again?.sequence_number ?? 'none')}`, + ); + } + } + stored.set(assignment.event_id, assignment.sequence_number); + } else if (current !== assignment.sequence_number) { + throw new SessionStoreInvariantError( + `event ${assignment.event_id} already has sequence_number ${String(current)}`, + ); + } + } +} + /** * listTurnEvents — one statement: drive from turn, left-join a page of events. * Missing turn → 0 rows → TurnNotFoundError. Empty log → one null-event sentinel. diff --git a/packages/trueforge/src/db/sqlite/types.ts b/packages/trueforge/src/db/sqlite/types.ts index 27aa02fcab..daea815d97 100644 --- a/packages/trueforge/src/db/sqlite/types.ts +++ b/packages/trueforge/src/db/sqlite/types.ts @@ -155,6 +155,11 @@ export interface SessionEventTable { event_id: string; event: JsonbColumn; created_at: string; + /** + * Stream sequence assigned when the event is first appended to the live log. + * Null until then. The event blob stays the client payload. + */ + sequence_number: ColumnType; } /** diff --git a/packages/trueforge/src/main.ts b/packages/trueforge/src/main.ts index 74fab831f5..479fc8f1fc 100644 --- a/packages/trueforge/src/main.ts +++ b/packages/trueforge/src/main.ts @@ -485,6 +485,7 @@ async function createServerRuntime(persistence: ServerPersistence< const activeTurns = new ActiveTurnRegistry({ pausedKeepaliveMs: configuration.TURN_PAUSED_KEEPALIVE_SECONDS * 1000, + lockAcquireTimeoutMs: configuration.TURN_LOCK_ACQUIRE_TIMEOUT_MS, logger, }); const requestReplyRouter = new RequestReplyRouter(); diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 601e84f115..ac375cce8e 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -5,10 +5,12 @@ * `track()` owns registration and cleanup around the stream lifecycle. */ import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; -import { Mutex } from 'async-mutex'; +import { E_TIMEOUT, Mutex, withTimeout } from 'async-mutex'; import type { Logger } from 'winston'; const DEFAULT_PAUSED_KEEPALIVE_MS = 300_000; +/** `0` waits until the lock is free. A positive value is the acquire budget in milliseconds. */ +const DEFAULT_LOCK_ACQUIRE_TIMEOUT_MS = 0; interface TurnRunCompletion { abortController: AbortController; @@ -28,6 +30,18 @@ interface LiveTurnRun extends TurnRunCompletion { type ActiveTurnRun = RebuildingTurnRun | LiveTurnRun; +/** `withTurnLock` did not acquire the per-turn lock before `timeoutMs`. */ +export class TurnLockAcquireTimeoutError extends Error { + readonly timeoutMs: number; + + constructor(options: { timeoutMs: number; cause: Error }) { + super(`Timed out acquiring the turn lock after ${String(options.timeoutMs)}ms`, { cause: options.cause }); + this.name = 'TurnLockAcquireTimeoutError'; + this.timeoutMs = options.timeoutMs; + Object.setPrototypeOf(this, new.target.prototype); + } +} + function setupTurnRun(abortController: AbortController): TurnRunCompletion { const { promise: waitUntilCompleted, resolve } = Promise.withResolvers(); const markCompleted = (): void => { @@ -59,11 +73,14 @@ export class ActiveTurnRegistry { private readonly keepaliveTimers = new Map(); // the time to keep the turn alive for, in milliseconds private readonly pausedKeepaliveMs: number; + /** Milliseconds to wait for the per-turn lock. `0` waits without a timeout. */ + private readonly lockAcquireTimeoutMs: number; private readonly logger: Logger; private alreadyShutDownAbortReason: CancellationReason | undefined; - constructor(options: { logger: Logger; pausedKeepaliveMs?: number }) { + constructor(options: { logger: Logger; pausedKeepaliveMs?: number; lockAcquireTimeoutMs?: number }) { this.pausedKeepaliveMs = options.pausedKeepaliveMs ?? DEFAULT_PAUSED_KEEPALIVE_MS; + this.lockAcquireTimeoutMs = options.lockAcquireTimeoutMs ?? DEFAULT_LOCK_ACQUIRE_TIMEOUT_MS; this.logger = options.logger.child({ module: 'ActiveTurnRegistry' }); } @@ -72,6 +89,8 @@ export class ActiveTurnRegistry { * the same turn has released it. Different turns never wait on each other. * Callers MUST re-read registry state inside `fn` instead of trusting a * lookup made before acquiring the lock. + * `lockAcquireTimeoutMs` of 0 waits until the lock is free. A positive value + * throws `TurnLockAcquireTimeoutError` when the wait exceeds it. */ async withTurnLock(input: { sessionId: string; turnId: string }, fn: () => Promise): Promise { const key = activeTurnKey(input.sessionId, input.turnId); @@ -80,9 +99,18 @@ export class ActiveTurnRegistry { if (!existing) { this.turnLocks.set(key, mutex); } + // 0 waits forever. withTimeout(0) would reject on the next timer turn. + const guarded = this.lockAcquireTimeoutMs > 0 ? withTimeout(mutex, this.lockAcquireTimeoutMs) : mutex; try { - // TODO (chiragjn): support acquire timeout - return await mutex.runExclusive(fn); + return await guarded.runExclusive(fn); + } catch (caught) { + if (caught === E_TIMEOUT) { + throw new TurnLockAcquireTimeoutError({ + timeoutMs: this.lockAcquireTimeoutMs, + cause: caught, + }); + } + throw caught; } finally { if (!mutex.isLocked()) { this.turnLocks.delete(key); diff --git a/packages/trueforge/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index 328aff64dc..c29fc89bc3 100644 --- a/packages/trueforge/src/runtime/event-subscription/inMemory.ts +++ b/packages/trueforge/src/runtime/event-subscription/inMemory.ts @@ -6,15 +6,17 @@ import { EventEmitter, once } from 'node:events'; import { StreamGoneError, + StreamOwnerMismatchError, SUBSCRIBE_STREAM_THRESHOLD_MS, type EventSubscription, type EventSubscriptionPollOptions, type EventSubscriptionPutOptions, type SequencedEvent, } from '.'; +import { applyEventLogRepair, eventIdOf, type RepairLogOptions } from './repair'; interface InMemoryStream { - /** Dense log: the event at index i has sequence_number i + 1 (1-indexed). */ + /** Appended in increasing sequence_number order. */ events: SequencedEvent[]; /** Undefined = no TTL set yet (mirrors a Redis key without EXPIRE). */ expiresAtMs?: number | undefined; @@ -27,6 +29,8 @@ interface InMemoryStream { */ export class InMemoryEventStreamStore { private readonly streams = new Map>(); + /** Survives stream eviction so a stale producer cannot recreate a missing log. */ + private readonly owners = new Map(); /** Rings per stream id on append and expiry so parked pollers wake instantly. */ private readonly changes = new EventEmitter(); @@ -35,20 +39,63 @@ export class InMemoryEventStreamStore { this.changes.setMaxListeners(0); } - append(streamId: string, event: T, streamTTLSeconds?: number): number { - const stream = this.getLiveStream(streamId) ?? { events: [] }; - this.streams.set(streamId, stream); + claimOwner(streamId: string, executorId: string): void { + this.owners.set(streamId, executorId); + } - const sequenceNumber = stream.events.length + 1; + /** + * Appends at an explicit sequence. A sequence at or below the tail is a duplicate + * when the tail already has that event id, and a conflict otherwise. + */ + appendExplicit(input: { + streamId: string; + event: T; + sequenceNumber: number; + eventId: string | undefined; + executorId: string | undefined; + checkOwner: boolean; + streamTTLSeconds: number | undefined; + }): 'ok' | 'duplicate' | 'owner_mismatch' | 'conflict' { + const owner = this.owners.get(input.streamId); + if (input.checkOwner && owner !== undefined && owner !== input.executorId) { + return 'owner_mismatch'; + } + const stream = this.getLiveStream(input.streamId) ?? { events: [] }; + const last = stream.events[stream.events.length - 1]; + if (last !== undefined && input.sequenceNumber <= last.sequence_number) { + const tailId = eventIdOf(last); + if (input.eventId !== undefined && tailId === input.eventId) { + return 'duplicate'; + } + return 'conflict'; + } + this.streams.set(input.streamId, stream); // Deep-copy so the log is isolated from caller mutations (mirrors JSON round-trip in Redis). - stream.events.push({ ...structuredClone(event), sequence_number: sequenceNumber }); + stream.events.push({ ...structuredClone(input.event), sequence_number: input.sequenceNumber }); + if (input.streamTTLSeconds !== undefined && input.streamTTLSeconds > 0) { + stream.expiresAtMs = Date.now() + input.streamTTLSeconds * 1_000; + this.scheduleExpiry(input.streamId, stream); + } + this.changes.emit(input.streamId); + return 'ok'; + } - if (streamTTLSeconds && streamTTLSeconds > 0) { - stream.expiresAtMs = Date.now() + streamTTLSeconds * 1_000; - this.scheduleExpiry(streamId, stream); + readTail(streamId: string): { eventId: string; sequenceNumber: number } | undefined { + const stream = this.getLiveStream(streamId); + const last = stream?.events[stream.events.length - 1]; + if (last === undefined) { + return undefined; } - this.changes.emit(streamId); - return sequenceNumber; + return { eventId: eventIdOf(last) ?? '', sequenceNumber: last.sequence_number }; + } + + /** Drops the log and keeps the owner, matching Redis stream-key expiry. */ + evictStream(streamId: string): void { + const stream = this.streams.get(streamId); + if (stream === undefined) { + return; + } + this.dropStream(streamId, stream); } /** Refreshes TTL on a live stream. Sequence numbers already follow the log length. */ @@ -118,28 +165,88 @@ export class InMemoryEventStreamStore { /** Stream-scoped view over the shared store; implements the transport contract. */ export class InMemoryEventSubscription implements EventSubscription { + private nextSequenceNumber = 1; + constructor( private readonly store: InMemoryEventStreamStore, private readonly streamId: string, ) {} - resumeProducer(options?: EventSubscriptionPutOptions): Promise { - this.store.refreshStreamTtl(this.streamId, options?.streamTTLSeconds); - return Promise.resolve(); + repairLog(options: RepairLogOptions): Promise { + const executorId = options.executorId; + return applyEventLogRepair({ + entries: options.entries, + persistSequences: options.persistSequences, + claimOwner: + executorId === undefined + ? undefined + : () => { + this.store.claimOwner(this.streamId, executorId); + return Promise.resolve(); + }, + readTail: () => Promise.resolve(this.store.readTail(this.streamId)), + append: entry => { + const outcome = this.store.appendExplicit({ + streamId: this.streamId, + event: entry.event, + sequenceNumber: entry.sequence_number, + eventId: entry.id, + executorId, + checkOwner: false, + streamTTLSeconds: options.streamTTLSeconds, + }); + if (outcome === 'ok') { + return Promise.resolve('appended'); + } + if (outcome === 'duplicate') { + return Promise.resolve('duplicate'); + } + if (outcome === 'owner_mismatch') { + return Promise.reject(new StreamOwnerMismatchError(this.streamId)); + } + return Promise.reject( + new Error(`explicit stream id ${String(entry.sequence_number)} conflicts on ${this.streamId}`), + ); + }, + refreshTtl: () => { + this.store.refreshStreamTtl(this.streamId, options.streamTTLSeconds); + return Promise.resolve(); + }, + refreshTtlOnUntouchedStream: executorId !== undefined, + }).then(next => { + this.nextSequenceNumber = next; + }); } put(event: T, options?: EventSubscriptionPutOptions): Promise { - return Promise.resolve(this.store.append(this.streamId, event, options?.streamTTLSeconds)); + const sequenceNumber = this.nextSequenceNumber; + this.nextSequenceNumber += 1; + const outcome = this.store.appendExplicit({ + streamId: this.streamId, + event, + sequenceNumber, + eventId: eventIdOf(event), + executorId: options?.executorId, + checkOwner: true, + streamTTLSeconds: options?.streamTTLSeconds, + }); + if (outcome === 'owner_mismatch') { + return Promise.reject(new StreamOwnerMismatchError(this.streamId)); + } + if (outcome === 'conflict') { + return Promise.reject(new Error(`explicit stream id ${String(sequenceNumber)} conflicts on ${this.streamId}`)); + } + return Promise.resolve(sequenceNumber); } assertSubscribable(): Promise { const stream = this.store.getLiveStream(this.streamId); if (!stream) { - throw new StreamGoneError(this.streamId); + return Promise.reject(new StreamGoneError(this.streamId)); } // Expiring within the threshold is treated as already gone, mirroring the Redis backend. if (stream.expiresAtMs !== undefined && stream.expiresAtMs - Date.now() < SUBSCRIBE_STREAM_THRESHOLD_MS) { - throw new StreamGoneError(this.streamId); + return Promise.reject(new StreamGoneError(this.streamId)); } return Promise.resolve(); } @@ -149,9 +256,8 @@ export class InMemoryEventSubscription implements EventSubscri options?: EventSubscriptionPollOptions, ): AsyncGenerator, void, unknown> { const signal = options?.signal; - // 0 and omitted both mean "from the start". With 1-indexed sequences stored at - // index seq-1, "strictly after N" starts at array index N. - let cursor = afterSequenceNumber === undefined || afterSequenceNumber === 0 ? 0 : afterSequenceNumber; + // 0 and omitted both mean "from the start". Yield events whose sequence is strictly greater. + let after = afterSequenceNumber === undefined || afterSequenceNumber === 0 ? 0 : afterSequenceNumber; for (;;) { if (signal?.aborted) { return; @@ -160,8 +266,8 @@ export class InMemoryEventSubscription implements EventSubscri if (!live) { throw new StreamGoneError(this.streamId); } - const end = live.events.length; - if (cursor >= end) { + const pending = live.events.filter(event => event.sequence_number > after); + if (pending.length === 0) { try { await this.store.waitForChange(this.streamId, signal); } catch (error) { @@ -173,13 +279,8 @@ export class InMemoryEventSubscription implements EventSubscri continue; } // Deep-copy each yielded event so consumers cannot mutate the stored log. - // Snapshot `end` so appends during yield are left for the next pass. - while (cursor < end) { - const event = live.events[cursor]; - if (event === undefined) { - throw new Error(`Corrupt stream ${this.streamId}: missing event at index ${String(cursor)}`); - } - cursor += 1; + for (const event of pending) { + after = event.sequence_number; yield structuredClone(event); } } diff --git a/packages/trueforge/src/runtime/event-subscription/index.ts b/packages/trueforge/src/runtime/event-subscription/index.ts index ab03654f23..a7d130182b 100644 --- a/packages/trueforge/src/runtime/event-subscription/index.ts +++ b/packages/trueforge/src/runtime/event-subscription/index.ts @@ -1,10 +1,18 @@ import type { RedisClient } from '@truefoundry/trueforge-core/request-reply'; import { InMemoryEventStreamStore, InMemoryEventSubscription } from './inMemory'; import { RedisEventSubscription } from './redis'; +import type { RepairLogOptions } from './repair'; + +export type { EventLogEntry, RepairLogOptions } from './repair'; export interface EventSubscriptionPutOptions { /** Sets the lifetime of the whole stream after this append. */ streamTTLSeconds?: number | undefined; + /** + * XADD runs only when the stream-owner key is unset or equals this id. + * A missing stream still rejects a caller that does not own the key. + */ + executorId?: string | undefined; } export type SequencedEvent = T & { sequence_number: number }; @@ -20,11 +28,11 @@ export interface EventSubscriptionPollOptions { */ export interface EventSubscription { /** - * Seeds the producer sequence from the last stored entry and refreshes the stream TTL. - * A missing stream stays at 1. Call this on a new producer before put. - * Repairing a missing or partial log from the TurnRecord is a later phase. + * Rebuilds a missing or behind stream from the durable log, then reseeds the + * producer counter. An existing stream whose tail event id is the last log id + * is left in place. Pass `executorId` only for the database owner that will put. */ - resumeProducer(options?: EventSubscriptionPutOptions): Promise; + repairLog(options: RepairLogOptions): Promise; /** Appends an event, assigns its sequence number, and returns that number. */ put(event: T, options?: EventSubscriptionPutOptions): Promise; @@ -62,6 +70,18 @@ export class EventSubscriptionRegistry { } return new InMemoryEventSubscription(this.memoryStore, streamId); } + + /** + * Drops the live log the way a TTL expiry does, and keeps the stream-owner key + * so a stale producer still cannot recreate it. + */ + evictStream(streamId: string): Promise { + if (this.redis) { + return this.redis.del(streamId).then(() => undefined); + } + this.memoryStore.evictStream(streamId); + return Promise.resolve(); + } } /** Reject new subscriptions when the stream expires within this window. */ @@ -77,3 +97,14 @@ export class StreamGoneError extends Error { Object.setPrototypeOf(this, new.target.prototype); } } + +/** put refused because the stream-owner key belongs to another executor. */ +export class StreamOwnerMismatchError extends Error { + readonly code = 'STREAM_OWNER_MISMATCH' as const; + + constructor(readonly streamId: string) { + super(`Cannot append to stream ${streamId}: caller is not the stream owner`); + this.name = 'StreamOwnerMismatchError'; + Object.setPrototypeOf(this, new.target.prototype); + } +} diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index 3cf26c7c32..d45fcc9828 100644 --- a/packages/trueforge/src/runtime/event-subscription/redis.ts +++ b/packages/trueforge/src/runtime/event-subscription/redis.ts @@ -2,17 +2,59 @@ import type { RedisClient } from '@truefoundry/trueforge-core/request-reply'; import { setTimeout as sleep } from 'node:timers/promises'; import { StreamGoneError, + StreamOwnerMismatchError, SUBSCRIBE_STREAM_THRESHOLD_MS, type EventSubscription, type EventSubscriptionPollOptions, type EventSubscriptionPutOptions, type SequencedEvent, } from '.'; +import { applyEventLogRepair, eventIdFromStreamData, eventIdOf, type RepairLogOptions } from './repair'; const SUBSCRIBE_STREAM_POLL_ITEMS_COUNT = 100; /** Delay between XREAD polls while a stream has no new events (a blocking XREAD would monopolize the shared connection). */ const SUBSCRIBE_STREAM_POLL_SLEEP_INTERVAL_MS = 1_000; +/** + * One script for put and for suffix backfill. Put checks the stream-owner key. + * Repair passes check-owner 0 so a subscribe can append without claiming. + * An explicit id at or below the tail is a conflict; the caller treats it as a + * no-op when that tail already holds the same event id. + */ +const APPEND_STREAM_ENTRY = ` +local checkOwner = ARGV[1] +if checkOwner == '1' then + local owner = redis.call('GET', KEYS[1]) + if owner ~= false and owner ~= ARGV[2] then + return {'owner_mismatch', ''} + end +end +local added = redis.pcall('XADD', KEYS[2], ARGV[3], 'data', ARGV[4]) +if type(added) == 'table' and added['err'] ~= nil then + local err = added['err'] + if string.find(err, 'equal or smaller', 1, true) then + local tail = redis.call('XREVRANGE', KEYS[2], '+', '-', 'COUNT', 1) + local data = '' + if tail[1] ~= nil then + local fields = tail[1][2] + for i = 1, #fields, 2 do + if fields[i] == 'data' then + data = fields[i + 1] + break + end + end + end + return {'conflict', data} + end + return {'error', err} +end +local ttl = tonumber(ARGV[5]) +if ttl ~= nil and ttl > 0 then + redis.call('EXPIRE', KEYS[2], ttl) +end +return {'ok', ''} +`; + function sequenceNumberFromEntryId(streamId: string, entryId: string): number { const separator = entryId.indexOf('-'); const raw = separator === -1 ? entryId : entryId.slice(0, separator); @@ -23,9 +65,31 @@ function sequenceNumberFromEntryId(streamId: string, entryId: string): number { return sequenceNumber; } +function streamOwnerKey(streamId: string): string { + return `${streamId}:owner`; +} + +type AppendStatus = 'ok' | 'conflict' | 'owner_mismatch' | 'error'; + +function isAppendStatus(value: unknown): value is AppendStatus { + return value === 'ok' || value === 'conflict' || value === 'owner_mismatch' || value === 'error'; +} + +function parseAppendReply(value: unknown): { status: AppendStatus; detail: string } { + if (!Array.isArray(value) || value.length !== 2) { + throw new Error('Unexpected stream append reply'); + } + const status: unknown = Object.prototype.hasOwnProperty.call(value, 0) ? Reflect.get(value, 0) : undefined; + const detail: unknown = Object.prototype.hasOwnProperty.call(value, 1) ? Reflect.get(value, 1) : undefined; + if (isAppendStatus(status) && typeof detail === 'string') { + return { status, detail }; + } + throw new Error('Unexpected stream append reply'); +} + /** - * One Redis Stream with locally minted 1-indexed sequence numbers. Assumes a - * single producer per stream; the counter dies with the instance, the data via key TTL. + * One Redis Stream. Sequence numbers are explicit ids (`n-1`). The owner key is + * separate from the stream so a missing stream still rejects a stale producer. */ export class RedisEventSubscription implements EventSubscription { private nextSequenceNumber = 1; @@ -35,30 +99,45 @@ export class RedisEventSubscription implements EventSubscripti private readonly streamId: string, ) {} - async resumeProducer(options?: EventSubscriptionPutOptions): Promise { - const entries = await this.redis.xRevRange(this.streamId, '+', '-', { COUNT: 1 }); - const latest = entries[0]; - if (latest) { - this.nextSequenceNumber = sequenceNumberFromEntryId(this.streamId, latest.id) + 1; - } - const streamTTLSeconds = options?.streamTTLSeconds; - if (streamTTLSeconds && streamTTLSeconds > 0) { - await this.redis.expire(this.streamId, streamTTLSeconds); - } + async repairLog(options: RepairLogOptions): Promise { + const executorId = options.executorId; + this.nextSequenceNumber = await applyEventLogRepair({ + entries: options.entries, + persistSequences: options.persistSequences, + claimOwner: + executorId === undefined + ? undefined + : () => this.redis.set(streamOwnerKey(this.streamId), executorId).then(() => undefined), + readTail: () => this.readTail(), + append: async entry => { + const reply = await this.appendEntry({ + sequenceNumber: entry.sequence_number, + payload: JSON.stringify({ ...entry.event, sequence_number: entry.sequence_number }), + checkOwner: false, + executorId: '', + streamTTLSeconds: options.streamTTLSeconds, + }); + return this.outcomeForEvent({ reply, eventId: entry.id, sequenceNumber: entry.sequence_number }); + }, + refreshTtl: () => this.refreshTtl(options.streamTTLSeconds), + refreshTtlOnUntouchedStream: executorId !== undefined, + }); } async put(event: T, options?: EventSubscriptionPutOptions): Promise { const sequenceNumber = this.nextSequenceNumber; this.nextSequenceNumber += 1; - - const sequencedEvent: SequencedEvent = { ...event, sequence_number: sequenceNumber }; - const multi = this.redis.multi().xAdd(this.streamId, `${String(sequenceNumber)}-1`, { - data: JSON.stringify(sequencedEvent), + const reply = await this.appendEntry({ + sequenceNumber, + payload: JSON.stringify({ ...event, sequence_number: sequenceNumber }), + checkOwner: true, + executorId: options?.executorId ?? '', + streamTTLSeconds: options?.streamTTLSeconds, }); - if (options?.streamTTLSeconds && options.streamTTLSeconds > 0) { - multi.expire(this.streamId, options.streamTTLSeconds); + if (reply.status === 'owner_mismatch') { + throw new StreamOwnerMismatchError(this.streamId); } - await multi.exec(); + await this.outcomeForEvent({ reply, eventId: eventIdOf(event), sequenceNumber }); return sequenceNumber; } @@ -127,4 +206,72 @@ export class RedisEventSubscription implements EventSubscripti } } } + + private async readTail(): Promise<{ eventId: string; sequenceNumber: number } | undefined> { + const entries = await this.redis.xRevRange(this.streamId, '+', '-', { COUNT: 1 }); + const latest = entries[0]; + if (!latest) { + return undefined; + } + const data = latest.message['data']; + if (data === undefined) { + throw new Error(`Corrupt stream entry ${latest.id} on ${this.streamId}: missing data field`); + } + return { + eventId: eventIdFromStreamData(data) ?? '', + sequenceNumber: sequenceNumberFromEntryId(this.streamId, latest.id), + }; + } + + private async refreshTtl(streamTTLSeconds: number | undefined): Promise { + if (streamTTLSeconds !== undefined && streamTTLSeconds > 0) { + await this.redis.expire(this.streamId, streamTTLSeconds); + } + } + + private async appendEntry(input: { + sequenceNumber: number; + payload: string; + checkOwner: boolean; + executorId: string; + streamTTLSeconds: number | undefined; + }): Promise<{ status: AppendStatus; detail: string }> { + const ttl = input.streamTTLSeconds !== undefined && input.streamTTLSeconds > 0 ? input.streamTTLSeconds : 0; + const reply: unknown = await this.redis.eval(APPEND_STREAM_ENTRY, { + keys: [streamOwnerKey(this.streamId), this.streamId], + arguments: [ + input.checkOwner ? '1' : '0', + input.executorId, + `${String(input.sequenceNumber)}-1`, + input.payload, + String(ttl), + ], + }); + return parseAppendReply(reply); + } + + private outcomeForEvent(input: { + reply: { status: AppendStatus; detail: string }; + eventId: string | undefined; + sequenceNumber: number; + }): Promise<'appended' | 'duplicate'> { + if (input.reply.status === 'ok') { + return Promise.resolve('appended'); + } + if (input.reply.status === 'owner_mismatch') { + return Promise.reject(new StreamOwnerMismatchError(this.streamId)); + } + if (input.reply.status === 'error') { + return Promise.reject(new Error(`Redis stream append failed on ${this.streamId}: ${input.reply.detail}`)); + } + const tailId = eventIdFromStreamData(input.reply.detail); + if (tailId !== undefined && input.eventId !== undefined && tailId === input.eventId) { + return Promise.resolve('duplicate'); + } + return Promise.reject( + new Error( + `explicit stream id ${String(input.sequenceNumber)} conflicts on ${this.streamId} and the tail event is ${tailId ?? 'missing'}`, + ), + ); + } } diff --git a/packages/trueforge/src/runtime/event-subscription/repair.ts b/packages/trueforge/src/runtime/event-subscription/repair.ts new file mode 100644 index 0000000000..199a267e47 --- /dev/null +++ b/packages/trueforge/src/runtime/event-subscription/repair.ts @@ -0,0 +1,153 @@ +/** + * Rebuild a turn's live event stream from the durable log. + * Sequence numbers already stored are written back as explicit stream ids. + * An event with no number yet takes max(stored) + 1 in persist order. + */ + +export interface EventLogEntry { + /** Payload id. Compared with the stream tail. Redis stream ids are sequence numbers, not this. */ + id: string; + event: T; + sequence_number?: number | undefined; +} + +interface SequencedLogEntry { + id: string; + event: T; + sequence_number: number; +} + +export interface RepairLogOptions { + entries: readonly EventLogEntry[]; + streamTTLSeconds?: number | undefined; + /** + * Set by the database owner that is about to produce. Claims the stream-owner + * key. A subscribe repair omits this and does not claim. + */ + executorId?: string | undefined; + /** + * Persist numbers assigned to events that did not have one. Called even when + * the list is empty so the owner fence runs before the stream is touched. + */ + persistSequences?: (assignments: { event_id: string; sequence_number: number }[]) => Promise; +} + +interface StreamTail { + eventId: string; + sequenceNumber: number; +} + +export function eventIdOf(event: object): string | undefined { + if (!('id' in event)) { + return undefined; + } + const id = event.id; + return typeof id === 'string' ? id : undefined; +} + +export function eventIdFromStreamData(data: string): string | undefined { + let parsed: unknown; + try { + parsed = JSON.parse(data); + } catch { + return undefined; + } + if (typeof parsed !== 'object' || parsed === null) { + return undefined; + } + return eventIdOf(parsed); +} + +function planSequences( + entries: readonly EventLogEntry[], +): { + sequenced: SequencedLogEntry[]; + newlyAssigned: { event_id: string; sequence_number: number }[]; +} { + let maxStored = 0; + for (const entry of entries) { + const sequence = entry.sequence_number; + if (sequence !== undefined && sequence > maxStored) { + maxStored = sequence; + } + } + const sequenced: SequencedLogEntry[] = []; + const newlyAssigned: { event_id: string; sequence_number: number }[] = []; + let next = maxStored; + for (const entry of entries) { + if (entry.sequence_number === undefined) { + next += 1; + sequenced.push({ id: entry.id, event: entry.event, sequence_number: next }); + newlyAssigned.push({ event_id: entry.id, sequence_number: next }); + } else { + sequenced.push({ id: entry.id, event: entry.event, sequence_number: entry.sequence_number }); + } + } + return { sequenced, newlyAssigned }; +} + +/** + * Suffix the stream does not already have. Redis explicit ids must increase, so + * the suffix is appended in sequence order. That matches persist order when + * numbers were assigned as events were streamed. + */ +function suffixAfterTail( + sequenced: readonly SequencedLogEntry[], + tail: StreamTail | undefined, +): SequencedLogEntry[] { + if (tail === undefined) { + return [...sequenced].sort((left, right) => left.sequence_number - right.sequence_number); + } + const last = sequenced[sequenced.length - 1]; + if (last?.id === tail.eventId) { + return []; + } + const tailIndex = sequenced.findIndex(entry => entry.id === tail.eventId); + const after = + tailIndex >= 0 + ? sequenced.slice(tailIndex + 1) + : sequenced.filter(entry => entry.sequence_number > tail.sequenceNumber); + return [...after].sort((left, right) => left.sequence_number - right.sequence_number); +} + +function nextSequence(sequenced: readonly { sequence_number: number }[], tail: StreamTail | undefined): number { + let max = tail?.sequenceNumber ?? 0; + for (const entry of sequenced) { + if (entry.sequence_number > max) { + max = entry.sequence_number; + } + } + return max + 1; +} + +export async function applyEventLogRepair(input: { + entries: readonly EventLogEntry[]; + persistSequences?: RepairLogOptions['persistSequences']; + claimOwner?: (() => Promise) | undefined; + readTail: () => Promise; + append: (entry: SequencedLogEntry) => Promise<'appended' | 'duplicate'>; + refreshTtl: () => Promise; + /** Producer resume refreshes TTL even when the existing stream is already complete. */ + refreshTtlOnUntouchedStream: boolean; +}): Promise { + const plan = planSequences(input.entries); + if (input.persistSequences !== undefined) { + await input.persistSequences(plan.newlyAssigned); + } + if (input.claimOwner !== undefined) { + await input.claimOwner(); + } + const tail = await input.readTail(); + const suffix = suffixAfterTail(plan.sequenced, tail); + let appended = false; + for (const entry of suffix) { + const outcome = await input.append(entry); + if (outcome === 'appended') { + appended = true; + } + } + if (appended || input.refreshTtlOnUntouchedStream) { + await input.refreshTtl(); + } + return nextSequence(plan.sequenced, tail); +} diff --git a/packages/trueforge/src/runtime/turnEventLogRepair.ts b/packages/trueforge/src/runtime/turnEventLogRepair.ts new file mode 100644 index 0000000000..b1b05235e3 --- /dev/null +++ b/packages/trueforge/src/runtime/turnEventLogRepair.ts @@ -0,0 +1,43 @@ +import type { ISessionStore, TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session'; +import type { EventLogEntry, EventSubscription } from './event-subscription'; + +/** + * Rebuild the live turn stream from the durable event log. + * The producing owner records any missing sequence numbers and claims the + * stream-owner key. Subscribe omits `producerExecutorId` and only appends. + */ +export async function repairTurnEventLog(input: { + sessionStore: Pick, 'listTurnEventLog' | 'assignTurnEventSequenceNumbers'>; + subscription: EventSubscription; + sessionId: string; + turnId: string; + streamTTLSeconds: number; + producerExecutorId?: string | undefined; +}): Promise { + const rows = await input.sessionStore.listTurnEventLog({ + session_id: input.sessionId, + turn_id: input.turnId, + }); + const entries: EventLogEntry[] = rows.map(row => ({ + id: row.event.id, + event: row.event, + ...(row.sequence_number === null ? {} : { sequence_number: row.sequence_number }), + })); + const producerExecutorId = input.producerExecutorId; + await input.subscription.repairLog({ + entries, + streamTTLSeconds: input.streamTTLSeconds, + ...(producerExecutorId === undefined + ? {} + : { + executorId: producerExecutorId, + persistSequences: assignments => + input.sessionStore.assignTurnEventSequenceNumbers({ + session_id: input.sessionId, + turn_id: input.turnId, + active_executor_id: producerExecutorId, + assignments, + }), + }), + }); +} diff --git a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts index 6960f7924c..f8caeda7e2 100644 --- a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts +++ b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts @@ -1,7 +1,13 @@ -import type { EventSubscriptionRegistry, SequencedEvent } from '../../../src/runtime/event-subscription'; +import { + StreamGoneError, + StreamOwnerMismatchError, + type EventSubscriptionRegistry, + type SequencedEvent, +} from '../../../src/runtime/event-subscription'; /** Payload shape used by the suite; distinct per event so round-trips are provable. */ export interface TestEvent { + id?: string; type: string; payload: { index: number; text: string }; } @@ -124,4 +130,109 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub { ...makeEvent(2), sequence_number: 3 }, ]); }); + + it('rebuilds a missing stream from stored sequences and continues after the tail', async () => { + const registry = createRegistry(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + const first = { id: 'evt-1', ...makeEvent(0) }; + const second = { id: 'evt-2', ...makeEvent(1) }; + await producer.put(first, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); + await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); + await registry.evictStream(streamId); + + await producer.repairLog({ + entries: [ + { id: 'evt-1', event: first, sequence_number: 1 }, + { id: 'evt-2', event: second, sequence_number: 2 }, + ], + streamTTLSeconds: STREAM_TTL_SECONDS, + executorId: 'owner', + persistSequences: () => Promise.resolve(), + }); + + const third = { id: 'evt-3', ...makeEvent(2) }; + const next = await producer.put(third, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); + expect(next).toBe(3); + const received = await collect(registry, streamId, 3); + expect(received).toEqual([ + { ...first, sequence_number: 1 }, + { ...second, sequence_number: 2 }, + { ...third, sequence_number: 3 }, + ]); + }); + + it('leaves an existing stream untouched and a second repair does not append again', async () => { + const registry = createRegistry(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + const first = { id: 'evt-1', ...makeEvent(0) }; + const second = { id: 'evt-2', ...makeEvent(1) }; + await producer.put(first, { streamTTLSeconds: STREAM_TTL_SECONDS }); + await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS }); + const repair = () => + producer.repairLog({ + entries: [ + { id: 'evt-1', event: first, sequence_number: 1 }, + { id: 'evt-2', event: second, sequence_number: 2 }, + ], + streamTTLSeconds: STREAM_TTL_SECONDS, + executorId: 'owner', + persistSequences: () => Promise.resolve(), + }); + await repair(); + await repair(); + + const third = { id: 'evt-3', ...makeEvent(2) }; + const next = await producer.put(third, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); + expect(next).toBe(3); + const received = await collect(registry, streamId, 3); + expect(received).toEqual([ + { ...first, sequence_number: 1 }, + { ...second, sequence_number: 2 }, + { ...third, sequence_number: 3 }, + ]); + }); + + it('rejects a stale producer when the stream is missing', async () => { + const registry = createRegistry(); + const streamId = nextStreamId(); + const owner = registry.get(streamId); + const event = { id: 'evt-1', ...makeEvent(0) }; + await owner.repairLog({ + entries: [{ id: 'evt-1', event, sequence_number: 1 }], + streamTTLSeconds: STREAM_TTL_SECONDS, + executorId: 'owner', + persistSequences: () => Promise.resolve(), + }); + await registry.evictStream(streamId); + + const stale = registry.get(streamId); + await expect( + stale.put({ id: 'evt-late', ...makeEvent(9) }, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'stale' }), + ).rejects.toBeInstanceOf(StreamOwnerMismatchError); + await expect(stale.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); + }); + + it('treats an explicit-id conflict as a no-op when the event id is already the tail', async () => { + const registry = createRegistry(); + const streamId = nextStreamId(); + const left = registry.get(streamId); + const right = registry.get(streamId); + const event = { id: 'evt-same', ...makeEvent(0) }; + const sequences = await Promise.all([ + left.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }), + right.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }), + ]); + expect(sequences).toEqual([1, 1]); + + const follow = { id: 'evt-next', ...makeEvent(1) }; + const next = await left.put(follow, { streamTTLSeconds: STREAM_TTL_SECONDS }); + expect(next).toBe(2); + const received = await collect(registry, streamId, 2); + expect(received).toEqual([ + { ...event, sequence_number: 1 }, + { ...follow, sequence_number: 2 }, + ]); + }); } diff --git a/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts new file mode 100644 index 0000000000..99fc92caf3 --- /dev/null +++ b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts @@ -0,0 +1,127 @@ +import type { TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session/TurnHandle'; +import { InMemorySessionStore } from '@truefoundry/trueforge-core/agent-session/store/InMemorySessionStore'; +import { TurnOwnershipLostError } from '@truefoundry/trueforge-core/agent-session/store/SessionStoreErrors'; +import { + makeAgentSpec, + makeCreateTurnInput, + makeModelMessageEvent, + makeTurnCreatedEvent, + TEST_ACTIVE_EXECUTOR_ID, +} from '../../../../trueforge-core/tests/agent-session/testHelpers'; +import { + EventSubscriptionRegistry, + StreamGoneError, + StreamOwnerMismatchError, + type SequencedEvent, +} from '../../../src/runtime/event-subscription'; +import { repairTurnEventLog } from '../../../src/runtime/turnEventLogRepair'; + +const STREAM_TTL_SECONDS = 300; + +async function collect( + registry: EventSubscriptionRegistry, + streamId: string, + count: number, +): Promise[]> { + const abort = new AbortController(); + const generator = registry.get(streamId).poll(undefined, { signal: abort.signal }); + const received: SequencedEvent[] = []; + try { + for await (const event of generator) { + received.push(event); + if (received.length >= count) { + break; + } + } + } finally { + abort.abort(); + await generator.return(undefined); + } + return received; +} + +describe('repairTurnEventLog', () => { + it('rebuilds a missing stream from the durable log and keeps the sequence', async () => { + const store = new InMemorySessionStore(); + const sessionId = 's1'; + const turnId = 'turn-1'; + await store.createSession({ + tenant_id: 't1', + session_id: sessionId, + created_by_subject: { subject_id: 'user-1', subject_type: 'user', subject_display_name: 'user-1' }, + agent: { type: 'inline', spec: makeAgentSpec() }, + custom: null, + metadata: {}, + external_id: null, + source: null, + }); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId })); + const created = makeTurnCreatedEvent(turnId); + const message = makeModelMessageEvent(); + await store.appendToEvents({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: turnId, + events: [created, message], + }); + + const registry = new EventSubscriptionRegistry(undefined); + const streamId = 'turn-stream'; + const subscription = registry.get(streamId); + const repair = (producerExecutorId?: string) => + repairTurnEventLog({ + sessionStore: store, + subscription, + sessionId, + turnId, + streamTTLSeconds: STREAM_TTL_SECONDS, + ...(producerExecutorId === undefined ? {} : { producerExecutorId }), + }); + + await expect(repair('other-executor')).rejects.toBeInstanceOf(TurnOwnershipLostError); + await expect(subscription.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); + + await repair(); + expect( + (await store.listTurnEventLog({ session_id: sessionId, turn_id: turnId })).map(row => row.sequence_number), + ).toEqual([null, null]); + expect((await collect(registry, streamId, 2)).map(event => event.sequence_number)).toEqual([1, 2]); + + await repair(); + expect((await collect(registry, streamId, 2)).map(event => event.sequence_number)).toEqual([1, 2]); + + await repair(TEST_ACTIVE_EXECUTOR_ID); + expect( + (await store.listTurnEventLog({ session_id: sessionId, turn_id: turnId })).map(row => row.sequence_number), + ).toEqual([1, 2]); + expect((await collect(registry, streamId, 2)).map(event => event.sequence_number)).toEqual([1, 2]); + + const later = makeModelMessageEvent(); + await store.appendToEvents({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: turnId, + events: [later], + }); + await repair(TEST_ACTIVE_EXECUTOR_ID); + expect( + (await store.listTurnEventLog({ session_id: sessionId, turn_id: turnId })).map(row => row.sequence_number), + ).toEqual([1, 2, 3]); + expect((await collect(registry, streamId, 3)).map(event => event.sequence_number)).toEqual([1, 2, 3]); + + await registry.evictStream(streamId); + const stale = registry.get(streamId); + await expect( + stale.put(makeModelMessageEvent(), { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'other-executor' }), + ).rejects.toBeInstanceOf(StreamOwnerMismatchError); + await expect(stale.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); + + await repair(TEST_ACTIVE_EXECUTOR_ID); + const next = await subscription.put(makeModelMessageEvent(), { + streamTTLSeconds: STREAM_TTL_SECONDS, + executorId: TEST_ACTIVE_EXECUTOR_ID, + }); + expect(next).toBe(4); + expect((await collect(registry, streamId, 4)).map(event => event.sequence_number)).toEqual([1, 2, 3, 4]); + }); +}); diff --git a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts index c7026627e7..780d029624 100644 --- a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts +++ b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts @@ -1,6 +1,7 @@ import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; +import { E_TIMEOUT } from 'async-mutex'; import { createLogger } from 'winston'; -import { ActiveTurnRegistry } from '../../../src/runtime/activeTurns'; +import { ActiveTurnRegistry, TurnLockAcquireTimeoutError } from '../../../src/runtime/activeTurns'; const logger = createLogger({ silent: true }); @@ -33,6 +34,13 @@ async function* gateOnAbort(signal: AbortSignal): AsyncGenerator { }); } +function requireTurnLockAcquireTimeout(caught: unknown): TurnLockAcquireTimeoutError { + if (!(caught instanceof TurnLockAcquireTimeoutError)) { + throw new Error('expected TurnLockAcquireTimeoutError', { cause: caught }); + } + return caught; +} + describe('ActiveTurnRegistry', () => { it('track passes events through and removes the run when the stream completes', async () => { const registry = new ActiveTurnRegistry({ logger }); @@ -232,4 +240,64 @@ describe('ActiveTurnRegistry', () => { void value; } }); + + it('throws TurnLockAcquireTimeoutError when the lock is not acquired in time', async () => { + const timeoutMs = 40; + const registry = new ActiveTurnRegistry({ logger, lockAcquireTimeoutMs: timeoutMs }); + const gate = Promise.withResolvers(); + const holder = registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, () => gate.promise); + const waiter = registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, async () => 'acquired'); + + let caught: unknown; + try { + await waiter; + } catch (error) { + caught = error; + } + expect(caught).toBeInstanceOf(TurnLockAcquireTimeoutError); + const error = requireTurnLockAcquireTimeout(caught); + expect(error.timeoutMs).toBe(timeoutMs); + expect(error.cause).toBe(E_TIMEOUT); + + gate.resolve(undefined); + await holder; + }); + + it('does not throw when the holder releases the lock before the acquire timeout', async () => { + const registry = new ActiveTurnRegistry({ logger, lockAcquireTimeoutMs: 300 }); + const gate = Promise.withResolvers(); + const holder = registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, () => gate.promise); + let entered = false; + const waiter = registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, async () => { + entered = true; + return 'acquired'; + }); + + await new Promise(resolve => { + setTimeout(resolve, 30); + }); + expect(entered).toBe(false); + gate.resolve(undefined); + await holder; + await expect(waiter).resolves.toBe('acquired'); + }); + + it('waits without a timeout when lockAcquireTimeoutMs is 0', async () => { + const registry = new ActiveTurnRegistry({ logger, lockAcquireTimeoutMs: 0 }); + const gate = Promise.withResolvers(); + const holder = registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, () => gate.promise); + let entered = false; + const waiter = registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, async () => { + entered = true; + return 'acquired'; + }); + + await new Promise(resolve => { + setTimeout(resolve, 80); + }); + expect(entered).toBe(false); + gate.resolve(undefined); + await holder; + await expect(waiter).resolves.toBe('acquired'); + }); }); From aa7918b4a27ac7df0beaa2020c15d3aa893181f1 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 27 Oct 2026 18:39:28 +0530 Subject: [PATCH 02/21] fix: date session event sequence migration after pause migrations Stream migrations must run after the pause migrations dated 20261028. Co-authored-by: Cursor --- ...mber.ts => 20261028_000003_session_event_sequence_number.ts} | 0 ...mber.ts => 20261028_000003_session_event_sequence_number.ts} | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename packages/trueforge/src/db/postgres/migrations/{20261024_000001_session_event_sequence_number.ts => 20261028_000003_session_event_sequence_number.ts} (100%) rename packages/trueforge/src/db/sqlite/migrations/{20261024_000001_session_event_sequence_number.ts => 20261028_000003_session_event_sequence_number.ts} (90%) diff --git a/packages/trueforge/src/db/postgres/migrations/20261024_000001_session_event_sequence_number.ts b/packages/trueforge/src/db/postgres/migrations/20261028_000003_session_event_sequence_number.ts similarity index 100% rename from packages/trueforge/src/db/postgres/migrations/20261024_000001_session_event_sequence_number.ts rename to packages/trueforge/src/db/postgres/migrations/20261028_000003_session_event_sequence_number.ts diff --git a/packages/trueforge/src/db/sqlite/migrations/20261024_000001_session_event_sequence_number.ts b/packages/trueforge/src/db/sqlite/migrations/20261028_000003_session_event_sequence_number.ts similarity index 90% rename from packages/trueforge/src/db/sqlite/migrations/20261024_000001_session_event_sequence_number.ts rename to packages/trueforge/src/db/sqlite/migrations/20261028_000003_session_event_sequence_number.ts index c423ed01a5..38bfb13084 100644 --- a/packages/trueforge/src/db/sqlite/migrations/20261024_000001_session_event_sequence_number.ts +++ b/packages/trueforge/src/db/sqlite/migrations/20261028_000003_session_event_sequence_number.ts @@ -2,7 +2,7 @@ import { sql, type Kysely } from 'kysely'; /** * Stream sequence lives beside the event payload so a missing stream can be rebuilt. - * Mirrors db/postgres/migrations/20261024_000001_session_event_sequence_number.ts. + * Mirrors db/postgres/migrations/20261028_000003_session_event_sequence_number.ts. */ export async function up(db: Kysely): Promise { await db.transaction().execute(async trx => { From 304b0a0ef8d9e80ac31b75c2ef89951732519e19 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Fri, 30 Oct 2026 11:01:34 +0530 Subject: [PATCH 03/21] Restore an expired turn stream only while the turn is paused. Running and finished turns keep the existing stream or fall back to list events, and subscribe does not rebuild the turn. Co-authored-by: Cursor --- .changeset/paused-subscribe-stream-restore.md | 5 + packages/trueforge/src/apis/turns.ts | 21 +-- packages/trueforge/src/routes/turnRoutes.ts | 6 +- .../src/runtime/turnEventLogRepair.ts | 6 +- .../trueforge/tests/unit/apis/turns.test.ts | 128 ++++++++++++++++++ 5 files changed, 151 insertions(+), 15 deletions(-) create mode 100644 .changeset/paused-subscribe-stream-restore.md diff --git a/.changeset/paused-subscribe-stream-restore.md b/.changeset/paused-subscribe-stream-restore.md new file mode 100644 index 0000000000..d859a67513 --- /dev/null +++ b/.changeset/paused-subscribe-stream-restore.md @@ -0,0 +1,5 @@ +--- +'@truefoundry/trueforge': patch +--- + +Restore an expired live stream only for a paused turn. Running and finished turns keep the existing stream or use list events. diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index e2c2b3cfe0..513f6b2004 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -2077,16 +2077,19 @@ export function createTurnsRouter(deps: TurnsRouterDeps(input: { sessionStore: Pick, 'listTurnEventLog' | 'assignTurnEventSequenceNumbers'>; diff --git a/packages/trueforge/tests/unit/apis/turns.test.ts b/packages/trueforge/tests/unit/apis/turns.test.ts index 2fab66dc70..a5474d8641 100644 --- a/packages/trueforge/tests/unit/apis/turns.test.ts +++ b/packages/trueforge/tests/unit/apis/turns.test.ts @@ -1,6 +1,7 @@ import { OpenAPIHono } from '@hono/zod-openapi'; import { AgentSpecSchema, + CancellationReason, Sessions, TurnNotFoundError, type TurnStreamingEvent, @@ -9,6 +10,16 @@ import { import { RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; import type { Kysely } from 'kysely'; import { createLogger } from 'winston'; +import { createEmptyAgentThreadMetrics } from '../../../../trueforge-core/src/core/runtime/metrics'; +import { + makeCancelledTurnState, + makeCreateTurnInput, + makePausedTurnState, + makeTurnCreatedEvent, + makeTurnDoneEvent, + makeTurnUpdateEvent, + TEST_ACTIVE_EXECUTOR_ID, +} from '../../../../trueforge-core/tests/agent-session/testHelpers'; import { createTurnsRouter, deriveSessionTitle, turnStreamId } from '../../../src/apis/turns'; import { TrueForgeAuthorizer, type Authorizer } from '../../../src/auth/authorizer'; import { STANDALONE_REQUEST_CONTEXT } from '../../../src/auth/identity'; @@ -882,4 +893,121 @@ describe('turns', () => { expect(deniedCanAccessAgent.mock.calls.map(([input]) => input.action)).toEqual(['use']); }); }); + + describe('subscribe stream restore', () => { + it('restores an expired stream only while the turn is paused', async () => { + const db = createSqliteDb(':memory:'); + await migrateSqliteToLatest(db); + const sessionStore = new SqliteSessionStore(db); + const sessions = new Sessions({ sessionStore }); + const sessionId = 's1'; + await sessionStore.createSession({ + tenant_id: 'default', + session_id: sessionId, + created_by_subject: { + subject_id: STANDALONE_REQUEST_CONTEXT.subject.id, + subject_type: STANDALONE_REQUEST_CONTEXT.subject.type, + subject_display_name: STANDALONE_REQUEST_CONTEXT.subject.display_name, + }, + agent: { + type: 'inline', + spec: AgentSpecSchema.parse({ + model: { name: 'test-provider/test-model' }, + instructions: 'test', + }), + }, + custom: null, + metadata: {}, + external_id: null, + source: null, + }); + const activeTurns = new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }); + const tokenStore = new SqliteOAuthTokenStore(db); + const app = new OpenAPIHono(); + app.route( + '/', + createTurnsRouter({ + activeTurns, + authorizer: new TrueForgeAuthorizer(), + eventSubscriptions: new EventSubscriptionRegistry(undefined), + logger: createLogger({ silent: true }), + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => new SqliteAgentStore(db), + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, tokenStore)), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore, + accessControlStore: new SqliteAccessControlStore(db), + }), + ); + + const runningTurnId = 'running-turn'; + await sessionStore.createTurn(makeCreateTurnInput({ sessionId, turnId: runningTurnId })); + const runningEvent = makeTurnCreatedEvent(runningTurnId); + await sessionStore.appendToEvents({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: runningTurnId, + events: [runningEvent], + }); + const runningSubscribe = await app.request(`/${sessionId}/turns/${runningTurnId}/subscribe`); + expect(runningSubscribe.status).toBe(412); + + const cancelled = makeCancelledTurnState(CancellationReason.ClientCancelled); + await sessionStore.freezeAndGetTurn({ + session_id: sessionId, + turn_id: runningTurnId, + reason: CancellationReason.ClientCancelled, + turn_done_event: makeTurnDoneEvent(cancelled), + }); + const finishedSubscribe = await app.request(`/${sessionId}/turns/${runningTurnId}/subscribe`); + expect(finishedSubscribe.status).toBe(412); + + const pausedTurnId = 'paused-turn'; + await sessionStore.createTurn(makeCreateTurnInput({ sessionId, turnId: pausedTurnId })); + const paused = makePausedTurnState(); + await sessionStore.updateTurnNonTerminalState({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: pausedTurnId, + state: paused, + turn_update_event: makeTurnUpdateEvent(paused), + turn_metrics: { + status: 'paused', + observed_at: new Date().toISOString(), + metrics: createEmptyAgentThreadMetrics(), + }, + }); + const pausedEvent = makeTurnCreatedEvent(pausedTurnId); + await sessionStore.appendToEvents({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: sessionId, + turn_id: pausedTurnId, + events: [pausedEvent], + }); + const pausedSubscribe = await app.request(`/${sessionId}/turns/${pausedTurnId}/subscribe`); + expect(pausedSubscribe.status).toBe(200); + const reader = pausedSubscribe.body?.getReader(); + if (reader === undefined) { + throw new Error('subscribe response has no body'); + } + const decoder = new TextDecoder(); + let replayed = ''; + try { + while (!replayed.includes(pausedEvent.id)) { + const { value, done } = await reader.read(); + if (done || value === undefined) { + break; + } + replayed += decoder.decode(value, { stream: true }); + } + } finally { + await reader.cancel(); + } + expect(replayed).toContain(pausedEvent.id); + expect(activeTurns.getTurnHandle({ sessionId, turnId: pausedTurnId })).toBeUndefined(); + }); + }); }); From 6f1045f8fd725a0d2440b1e3dddb6adf37c128a0 Mon Sep 17 00:00:00 2001 From: "trueforge-dev-bot[bot]" Date: Fri, 30 Oct 2026 05:34:26 +0000 Subject: [PATCH 04/21] Regenerate OpenAPI document and SDKs --- .github/fern/openapi/openapi.json | 6 +++--- docs/openapi.json | 6 +++--- packages/trueforge-sdk/reference.md | 2 +- .../src/api/resources/sessions/client/Client.ts | 2 +- python/trueforge_sdk/reference.md | 2 +- python/trueforge_sdk/src/trueforge_sdk/sessions/client.py | 4 ++-- .../trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py | 4 ++-- 7 files changed, 13 insertions(+), 13 deletions(-) diff --git a/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index 43f809f95d..05ba080394 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -10610,7 +10610,7 @@ }, "/api/v1/sessions/{session_id}/turns/{turn_id}/subscribe": { "get": { - "description": "Subscribe to the live SSE stream for a turn. Only the session creator may subscribe. Pass `after_sequence_number` to resume after a disconnect (exclusive — events after this sequence number are replayed).", + "description": "Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence.", "parameters": [ { "description": "Session identifier.", @@ -10699,10 +10699,10 @@ } } }, - "description": "Cannot subscribe — the live stream no longer exists." + "description": "The live stream does not exist and this turn is not paused, so it was not restored." } }, - "summary": "Subscribe to a running turn", + "summary": "Subscribe to a turn stream", "tags": [ "Agent Sessions" ], diff --git a/docs/openapi.json b/docs/openapi.json index 43f809f95d..05ba080394 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -10610,7 +10610,7 @@ }, "/api/v1/sessions/{session_id}/turns/{turn_id}/subscribe": { "get": { - "description": "Subscribe to the live SSE stream for a turn. Only the session creator may subscribe. Pass `after_sequence_number` to resume after a disconnect (exclusive — events after this sequence number are replayed).", + "description": "Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence.", "parameters": [ { "description": "Session identifier.", @@ -10699,10 +10699,10 @@ } } }, - "description": "Cannot subscribe — the live stream no longer exists." + "description": "The live stream does not exist and this turn is not paused, so it was not restored." } }, - "summary": "Subscribe to a running turn", + "summary": "Subscribe to a turn stream", "tags": [ "Agent Sessions" ], diff --git a/packages/trueforge-sdk/reference.md b/packages/trueforge-sdk/reference.md index 2bb8a1e3b7..17c468fbfc 100644 --- a/packages/trueforge-sdk/reference.md +++ b/packages/trueforge-sdk/reference.md @@ -2921,7 +2921,7 @@ await client.sessions.createTurnEvents("session_id", "turn_id", {
-Subscribe to the live SSE stream for a turn. Only the session creator may subscribe. Pass `after_sequence_number` to resume after a disconnect (exclusive — events after this sequence number are replayed). +Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence.
diff --git a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts index 7d67890b87..44d0a87324 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts @@ -2140,7 +2140,7 @@ export class SessionsClient { } /** - * Subscribe to the live SSE stream for a turn. Only the session creator may subscribe. Pass `after_sequence_number` to resume after a disconnect (exclusive — events after this sequence number are replayed). + * Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. */ public subscribeToTurn( session_id: string, diff --git a/python/trueforge_sdk/reference.md b/python/trueforge_sdk/reference.md index b8fb9e73af..c03fb3076c 100644 --- a/python/trueforge_sdk/reference.md +++ b/python/trueforge_sdk/reference.md @@ -3462,7 +3462,7 @@ client.sessions.create_turn_events(
-Subscribe to the live SSE stream for a turn. Only the session creator may subscribe. Pass `after_sequence_number` to resume after a disconnect (exclusive — events after this sequence number are replayed). +Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence.
diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py index 404bc64b85..a26183cc54 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py @@ -837,7 +837,7 @@ def subscribe_to_turn( request_options: typing.Optional[RequestOptions] = None, ) -> Stream[TurnStreamingEvent]: """ - Subscribe to the live SSE stream for a turn. Only the session creator may subscribe. Pass `after_sequence_number` to resume after a disconnect (exclusive — events after this sequence number are replayed). + Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. Parameters ---------- @@ -1824,7 +1824,7 @@ def subscribe_to_turn( request_options: typing.Optional[RequestOptions] = None, ) -> AsyncStream[TurnStreamingEvent]: """ - Subscribe to the live SSE stream for a turn. Only the session creator may subscribe. Pass `after_sequence_number` to resume after a disconnect (exclusive — events after this sequence number are replayed). + Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. Parameters ---------- diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py index b8a825d02c..d1d46972f3 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py @@ -1750,7 +1750,7 @@ def subscribe_to_turn( request_options: typing.Optional[RequestOptions] = None, ) -> typing.Iterator[HttpResponse[Stream[TurnStreamingEvent]]]: """ - Subscribe to the live SSE stream for a turn. Only the session creator may subscribe. Pass `after_sequence_number` to resume after a disconnect (exclusive — events after this sequence number are replayed). + Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. Parameters ---------- @@ -3583,7 +3583,7 @@ async def subscribe_to_turn( request_options: typing.Optional[RequestOptions] = None, ) -> typing.AsyncIterator[AsyncHttpResponse[AsyncStream[TurnStreamingEvent]]]: """ - Subscribe to the live SSE stream for a turn. Only the session creator may subscribe. Pass `after_sequence_number` to resume after a disconnect (exclusive — events after this sequence number are replayed). + Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. Parameters ---------- From 5284d4eedccf23f4c501497b1248c0ea98582f37 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 3 Nov 2026 16:29:04 +0530 Subject: [PATCH 05/21] Subscribe to turn streams by event id (#1242) * Subscribe and repair turn streams by event id. New live and repaired streams use Redis ids {ulidMillisecond}-{n} and do not assign sequence numbers. Subscribe accepts from_event_id or after_event_id. after_sequence_number remains only for older integer streams. * Fix event-id resume on legacy streams, repair, and reconnect. Scan leftover sequence streams from the start, retry a distinct shared-id payload after a timestamp conflict, and match repair rows without depending on JSON key order. Replay mcp.auth_required, and do not treat that inclusive shell as reconnect progress. * Minor fixes from self review * Regenerate OpenAPI document and SDKs * Simplify redis lua scripts --------- Co-authored-by: Cursor Agent Co-authored-by: trueforge-dev-bot[bot] --- .changeset/event-id-subscribe-cursor.md | 8 + .github/fern/openapi/openapi.json | 28 +- docs/openapi.json | 28 +- .../src/convertTurnMessages.ts | 93 ++-- .../assistant-ui-runtime/src/server/events.ts | 3 +- .../assistant-ui-runtime/src/server/types.ts | 4 + .../src/sessionSnapshot.ts | 4 +- .../assistant-ui-runtime/src/streamTurn.ts | 11 +- .../src/turnStreamUpdate.ts | 20 + .../src/useTrueForgeAgentMessages.ts | 60 ++- .../test/convertTurnMessages.test.ts | 47 ++ .../test/streamTurn.test.ts | 12 +- .../test/useTrueForgeAgentMessages.test.tsx | 41 +- packages/trueforge-sdk/reference.md | 2 +- .../api/resources/sessions/client/Client.ts | 6 +- .../SubscribeToTurnSessionsRequest.ts | 6 +- .../chatServer.ts | 37 +- .../harnessServer.test.ts | 33 ++ packages/trueforge/src/apis/turns.ts | 119 +++--- packages/trueforge/src/routes/turnRoutes.ts | 2 +- packages/trueforge/src/runtime/activeTurns.ts | 2 +- .../src/runtime/event-subscription/entryId.ts | 134 ++++++ .../runtime/event-subscription/inMemory.ts | 223 +++++++--- .../src/runtime/event-subscription/index.ts | 58 ++- .../src/runtime/event-subscription/redis.ts | 401 +++++++++++++----- .../src/runtime/event-subscription/repair.ts | 150 +++---- .../src/runtime/turnEventLogRepair.ts | 26 +- packages/trueforge/src/schemas/turn.ts | 25 +- .../event-subscription/contractSuite.ts | 187 ++++---- .../event-subscription/entryId.test.ts | 41 ++ .../turnEventLogRepair.test.ts | 26 +- .../trueforge/tests/unit/apis/turns.test.ts | 17 + python/trueforge_sdk/reference.md | 20 +- .../src/trueforge_sdk/sessions/client.py | 28 +- .../src/trueforge_sdk/sessions/raw_client.py | 28 +- 35 files changed, 1393 insertions(+), 537 deletions(-) create mode 100644 .changeset/event-id-subscribe-cursor.md create mode 100644 packages/trueforge/src/runtime/event-subscription/entryId.ts create mode 100644 packages/trueforge/tests/runtime/event-subscription/entryId.test.ts diff --git a/.changeset/event-id-subscribe-cursor.md b/.changeset/event-id-subscribe-cursor.md new file mode 100644 index 0000000000..a378d48bb3 --- /dev/null +++ b/.changeset/event-id-subscribe-cursor.md @@ -0,0 +1,8 @@ +--- +'@truefoundry/trueforge': patch +'@truefoundry/trueforge-assistant-ui-runtime': patch +'@truefoundry/trueforge-ui': patch +'@truefoundry/trueforge-sdk': patch +--- + +Subscribe with from_event_id or after_event_id. New turn streams store and repair by event id and omit SSE sequence ids. after_sequence_number remains for older integer streams. Regenerate SDK from the updated OpenAPI spec. diff --git a/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index 0c9a3925b2..c2071411ca 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -10725,7 +10725,7 @@ }, "/api/v1/sessions/{session_id}/turns/{turn_id}/subscribe": { "get": { - "description": "Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence.", + "description": "Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id", "parameters": [ { "description": "Session identifier.", @@ -10751,18 +10751,40 @@ } }, { - "description": "Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer.", + "description": "Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.", "in": "query", "name": "after_sequence_number", "required": false, "schema": { - "description": "Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer.", + "description": "Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.", "minimum": 0, "type": [ "integer", "null" ] } + }, + { + "description": "Inclusive cursor. Replay starts at the entry with this event id", + "in": "query", + "name": "from_event_id", + "required": false, + "schema": { + "description": "Inclusive cursor. Replay starts at the entry with this event id", + "minLength": 1, + "type": "string" + } + }, + { + "description": "Exclusive cursor. Replay starts after the last entry with this event id.", + "in": "query", + "name": "after_event_id", + "required": false, + "schema": { + "description": "Exclusive cursor. Replay starts after the last entry with this event id.", + "minLength": 1, + "type": "string" + } } ], "responses": { diff --git a/docs/openapi.json b/docs/openapi.json index 0c9a3925b2..c2071411ca 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -10725,7 +10725,7 @@ }, "/api/v1/sessions/{session_id}/turns/{turn_id}/subscribe": { "get": { - "description": "Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence.", + "description": "Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id", "parameters": [ { "description": "Session identifier.", @@ -10751,18 +10751,40 @@ } }, { - "description": "Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer.", + "description": "Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.", "in": "query", "name": "after_sequence_number", "required": false, "schema": { - "description": "Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer.", + "description": "Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.", "minimum": 0, "type": [ "integer", "null" ] } + }, + { + "description": "Inclusive cursor. Replay starts at the entry with this event id", + "in": "query", + "name": "from_event_id", + "required": false, + "schema": { + "description": "Inclusive cursor. Replay starts at the entry with this event id", + "minLength": 1, + "type": "string" + } + }, + { + "description": "Exclusive cursor. Replay starts after the last entry with this event id.", + "in": "query", + "name": "after_event_id", + "required": false, + "schema": { + "description": "Exclusive cursor. Replay starts after the last entry with this event id.", + "minLength": 1, + "type": "string" + } } ], "responses": { diff --git a/packages/assistant-ui-runtime/src/convertTurnMessages.ts b/packages/assistant-ui-runtime/src/convertTurnMessages.ts index 0f305b6937..13b79d90da 100644 --- a/packages/assistant-ui-runtime/src/convertTurnMessages.ts +++ b/packages/assistant-ui-runtime/src/convertTurnMessages.ts @@ -77,7 +77,7 @@ import { toolResponseStatus, } from './toolResponse.js'; import { appendMcpAuthToTurnContent, appendToolApprovalToTurnContent } from './turnEventHelpers.js'; -import type { TurnStreamUpdate } from './turnStreamUpdate.js'; +import { subscribeCursorForEvent, type TurnStreamUpdate, type TurnSubscribeCursor } from './turnStreamUpdate.js'; const MESSAGE_ID_SUFFIX = { ASSISTANT: '-assistant', @@ -1857,6 +1857,7 @@ export async function* streamTurnEvents( groupRootBaseline?: readonly string[], onTurnIdAvailable?: (turnId: string) => void, onSequenceNumber?: (sequenceNumber: number) => void, + onSubscribeCursor?: (cursor: TurnSubscribeCursor) => void, ): AsyncGenerator { let pendingMcpAuth: McpAuthRequiredEvent | undefined; let sandboxId: string | undefined; @@ -1911,24 +1912,34 @@ export async function* streamTurnEvents( findFirstPendingApprovalThreadId(foldState) != null || findFirstPendingResponseThreadId(foldState) != null; - const yieldContentOrRequiredAction = (sequenceNumber: number): TurnStreamUpdate => { + const withSequence = (update: TurnStreamUpdate, sequenceNumber: number | undefined): TurnStreamUpdate => + sequenceNumber === undefined ? update : { ...update, sequenceNumber }; + + const yieldContentOrRequiredAction = (sequenceNumber: number | undefined): TurnStreamUpdate => { if (foldStillNeedsUserInput()) { - return { - ...buildRequiredActionUpdate(), + return withSequence( + { + ...buildRequiredActionUpdate(), + turnState: { status: TURN_STATUS.PAUSED }, + }, sequenceNumber, - turnState: { status: TURN_STATUS.PAUSED }, - }; + ); } - return { content: yieldContent() ?? [], sequenceNumber }; + return withSequence({ content: yieldContent() ?? [] }, sequenceNumber); }; let yieldedPaused = false; let lastSequenceNumber: number | undefined; for await (const data of stream) { - lastSequenceNumber = data.sequenceNumber; - onSequenceNumber?.(data.sequenceNumber); + if (data.sequenceNumber !== undefined) { + lastSequenceNumber = data.sequenceNumber; + onSequenceNumber?.(data.sequenceNumber); + } const event = data.event; + if ('id' in event && typeof event.id === 'string') { + onSubscribeCursor?.(subscribeCursorForEvent(event)); + } if (event.type === EVENT_TYPE.TURN_CREATED) { onTurnIdAvailable?.(event.turnId); @@ -1954,22 +1965,30 @@ export async function* streamTurnEvents( if (event.type === EVENT_TYPE.TURN_UPDATE) { if (event.state.status === TURN_STATUS.RUNNING) { pendingMcpAuth = undefined; - yield withSandbox({ - content: yieldContent() ?? [], - status: runningAssistantStatus(), - sequenceNumber: data.sequenceNumber, - turnState: { status: TURN_STATUS.RUNNING }, - }); + yield withSandbox( + withSequence( + { + content: yieldContent() ?? [], + status: runningAssistantStatus(), + turnState: { status: TURN_STATUS.RUNNING }, + }, + data.sequenceNumber, + ), + ); continue; } // Pause is in-stream: the same SSE continues through apply-only echoes, // turn.update running, and turn.done. Do not detach the consumer. - yield withSandbox({ - ...buildRequiredActionUpdate(), - sequenceNumber: data.sequenceNumber, - turnState: { status: TURN_STATUS.PAUSED }, - }); + yield withSandbox( + withSequence( + { + ...buildRequiredActionUpdate(), + turnState: { status: TURN_STATUS.PAUSED }, + }, + data.sequenceNumber, + ), + ); yieldedPaused = true; continue; } @@ -1984,13 +2003,17 @@ export async function* streamTurnEvents( const previousParentTerminal = foldState.parentTerminalStatus; foldState.parentTerminalStatus = 'cancelled'; try { - yield withSandbox({ - content: buildRootAssistantContentForIds(foldState, ids), - status: assistantStatusFromTurnState(event.state), - metadata: { custom: { [CANCELLATION_REASON_CUSTOM_KEY]: event.state.reason } }, - sequenceNumber: data.sequenceNumber, - turnState: event.state, - }); + yield withSandbox( + withSequence( + { + content: buildRootAssistantContentForIds(foldState, ids), + status: assistantStatusFromTurnState(event.state), + metadata: { custom: { [CANCELLATION_REASON_CUSTOM_KEY]: event.state.reason } }, + turnState: event.state, + }, + data.sequenceNumber, + ), + ); } finally { foldState.parentTerminalStatus = previousParentTerminal; } @@ -2000,12 +2023,16 @@ export async function* streamTurnEvents( // the transport body to close because resumable subscriptions may linger. // Yield error/complete state (HITL) rather than throwing — reconnect must // not subscribe-retry a terminal turn; callers read `turnState` / status. - yield withSandbox({ - content: yieldContent() ?? [], - status: assistantStatusFromTurnState(event.state), - sequenceNumber: data.sequenceNumber, - turnState: event.state, - }); + yield withSandbox( + withSequence( + { + content: yieldContent() ?? [], + status: assistantStatusFromTurnState(event.state), + turnState: event.state, + }, + data.sequenceNumber, + ), + ); return; } diff --git a/packages/assistant-ui-runtime/src/server/events.ts b/packages/assistant-ui-runtime/src/server/events.ts index 4063e8e7c1..f033c64857 100644 --- a/packages/assistant-ui-runtime/src/server/events.ts +++ b/packages/assistant-ui-runtime/src/server/events.ts @@ -320,7 +320,8 @@ export type ActionRequiredEvent = ToolApprovalRequiredEvent | ToolResponseRequir export type PersistedTurnEvent = TurnCreatedEvent | TurnUpdateEvent | TurnDoneEvent | TurnEvent; export interface TurnStreamData { - sequenceNumber: number; + /** Present only when the server wrote an SSE id for an older integer stream. */ + sequenceNumber?: number; event: TStreamEvent; } diff --git a/packages/assistant-ui-runtime/src/server/types.ts b/packages/assistant-ui-runtime/src/server/types.ts index 1b8220c4ac..46127bf9da 100644 --- a/packages/assistant-ui-runtime/src/server/types.ts +++ b/packages/assistant-ui-runtime/src/server/types.ts @@ -422,6 +422,10 @@ export interface SubscribeToTurnRequest { sessionId: string; turnId: string; afterSequenceNumber?: number; + /** Inclusive cursor. Replay starts at this event. */ + fromEventId?: string; + /** Exclusive cursor. Replay starts after the last entry with this event id. */ + afterEventId?: string; abortSignal?: AbortSignal; } diff --git a/packages/assistant-ui-runtime/src/sessionSnapshot.ts b/packages/assistant-ui-runtime/src/sessionSnapshot.ts index 7bdc068ba3..5416cacc2b 100644 --- a/packages/assistant-ui-runtime/src/sessionSnapshot.ts +++ b/packages/assistant-ui-runtime/src/sessionSnapshot.ts @@ -13,7 +13,7 @@ import { extractTurnUserText } from './extractTurnUserText.js'; import { PeerThreadFoldState } from './foldPeerThreads.js'; import type { StoredApprovalDecision } from './toolApproval.js'; import type { StoredToolResponse } from './toolResponse.js'; -import type { TurnStreamUpdate } from './turnStreamUpdate.js'; +import type { TurnStreamUpdate, TurnSubscribeCursor } from './turnStreamUpdate.js'; export const STREAM_SEGMENT_STATUS = { DISCONNECTED: 'disconnected', @@ -79,6 +79,8 @@ export interface ActiveStreamState { */ segmentStatus: StreamSegmentStatus; lastSequenceNumber?: number | undefined; + /** Event-id resume cursor. New streams do not use `lastSequenceNumber`. */ + subscribeCursor?: TurnSubscribeCursor | undefined; } export interface PendingUserMessage { diff --git a/packages/assistant-ui-runtime/src/streamTurn.ts b/packages/assistant-ui-runtime/src/streamTurn.ts index 790b7ad109..6e6d8bc1be 100644 --- a/packages/assistant-ui-runtime/src/streamTurn.ts +++ b/packages/assistant-ui-runtime/src/streamTurn.ts @@ -4,7 +4,7 @@ import type { AgentChatServer, PreviousTurnIdInput, UserMessageContent } from '. import { streamTurnEvents } from './convertTurnMessages.js'; import { PeerThreadFoldState } from './foldPeerThreads.js'; -import type { TurnStreamUpdate } from './turnStreamUpdate.js'; +import type { TurnStreamUpdate, TurnSubscribeCursor } from './turnStreamUpdate.js'; export interface StreamTurnOptions { userMessage?: UserMessageContent; @@ -31,6 +31,7 @@ export async function* streamTurnContent( */ onTurnIdAvailable?: (turnId: string) => void, onSequenceNumber?: (sequenceNumber: number) => void, + onSubscribeCursor?: (cursor: TurnSubscribeCursor) => void, ): AsyncGenerator { // Aborting only detaches this client from the run; the turn keeps running on // the backend so switching sessions (or remounting) can reattach via @@ -56,7 +57,7 @@ export async function* streamTurnContent( }); try { - yield* streamTurnEvents(stream, foldState, groupRootBaseline, notifyTurnId, onSequenceNumber); + yield* streamTurnEvents(stream, foldState, groupRootBaseline, notifyTurnId, onSequenceNumber, onSubscribeCursor); } catch (error) { if (error instanceof Error && error.name === 'AbortError') { return; @@ -71,9 +72,10 @@ export async function* resumeTurnStream( turnId: string, foldState: PeerThreadFoldState, abortSignal: AbortSignal, - afterSequenceNumber?: number, + subscribeCursor?: TurnSubscribeCursor, groupRootBaseline?: readonly string[], onSequenceNumber?: (sequenceNumber: number) => void, + onSubscribeCursor?: (cursor: TurnSubscribeCursor) => void, ): AsyncGenerator { // Optional on custom backends. Callers detect the gap and report it, so an // empty stream here is safer than throwing mid-render. @@ -90,13 +92,14 @@ export async function* resumeTurnStream( server.subscribeToTurn({ sessionId, turnId, - ...(afterSequenceNumber != null ? { afterSequenceNumber } : {}), + ...(subscribeCursor ?? {}), abortSignal, }), foldState, groupRootBaseline, undefined, onSequenceNumber, + onSubscribeCursor, ); } catch (error) { if (error instanceof Error && error.name === 'AbortError') { diff --git a/packages/assistant-ui-runtime/src/turnStreamUpdate.ts b/packages/assistant-ui-runtime/src/turnStreamUpdate.ts index 9db3271305..12bc2aaf32 100644 --- a/packages/assistant-ui-runtime/src/turnStreamUpdate.ts +++ b/packages/assistant-ui-runtime/src/turnStreamUpdate.ts @@ -1,8 +1,28 @@ import type { MessageStatus } from '@assistant-ui/core'; import type { AssistantContentPart } from './modelMessageContent.js'; +import { EVENT_TYPE } from './server/events.js'; import type { TurnState } from './server/index.js'; +/** Resume cursor chosen from the last event the client applied. */ +export type TurnSubscribeCursor = { fromEventId: string } | { afterEventId: string }; + +/** + * A model message replays from its shell. `mcp.auth_required` is replayed too: + * the connect prompt lives only on that event, and the next `turn.update` + * does not repeat it. Every other event resumes after itself. + */ +export function subscribeCursorForEvent(event: { type: string; id: string }): TurnSubscribeCursor { + if ( + event.type === EVENT_TYPE.MODEL_MESSAGE || + event.type === EVENT_TYPE.MODEL_MESSAGE_DELTA || + event.type === EVENT_TYPE.MCP_AUTH_REQUIRED + ) { + return { fromEventId: event.id }; + } + return { afterEventId: event.id }; +} + export interface TurnStreamUpdate { content: AssistantContentPart[]; status?: MessageStatus; diff --git a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts index 2252585afb..f5bb263ff1 100644 --- a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts +++ b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts @@ -50,7 +50,7 @@ import { delayReconnect, isAbortError, STREAM_RECONNECT_MAX_ATTEMPTS } from './s import { resumeTurnStream, streamTurnContent } from './streamTurn.js'; import { mapApprovalDecision, type RespondToToolApprovalOptions } from './toolApproval.js'; import { type RespondToToolResponseOptions } from './toolResponse.js'; -import type { TurnStreamUpdate } from './turnStreamUpdate.js'; +import { subscribeCursorForEvent, type TurnStreamUpdate, type TurnSubscribeCursor } from './turnStreamUpdate.js'; export interface UseTrueForgeAgentMessagesOptions { server: AgentChatServer; @@ -101,6 +101,26 @@ interface QueuedSubscription { interface StreamSequenceTrack { onSequenceNumber: (sequenceNumber: number) => void; + onSubscribeCursor: (cursor: TurnSubscribeCursor) => void; +} + +/** Live cursor wins. A refresh falls back to the last stored event for the turn. */ +function cursorForTurn(snapshot: SessionSnapshot, turnId: string): TurnSubscribeCursor | undefined { + const active = snapshot.activeStream; + if (active?.turnId === turnId && active.subscribeCursor !== undefined) { + return active.subscribeCursor; + } + const history = snapshot.historyEvents; + if (history === undefined) { + return undefined; + } + for (let index = history.length - 1; index >= 0; index -= 1) { + const item = history[index]; + if (item?.turnId === turnId) { + return subscribeCursorForEvent(item.event); + } + } + return undefined; } interface RunStreamReconnect { @@ -581,14 +601,32 @@ export function useTrueForgeAgentMessages({ const activeAtStart = snapshotRef.current.activeStream; let lastSequenceNumber: number | undefined = activeAtStart?.turnId === turnIdRef.current ? activeAtStart.lastSequenceNumber : undefined; + let subscribeCursor = cursorForTurn(snapshotRef.current, turnIdRef.current); let useSubscribe = false; let consecutiveFailures = 0; + let eventsThisAttempt = 0; + let attemptCursor: TurnSubscribeCursor | undefined = subscribeCursor; const track: StreamSequenceTrack = { onSequenceNumber: sequenceNumber => { lastSequenceNumber = sequenceNumber; consecutiveFailures = 0; }, + onSubscribeCursor: cursor => { + eventsThisAttempt += 1; + // Inclusive replay delivers the shell we already have. That is not + // new progress, so a stream that dies on the shell still counts. + const replayedShell = + eventsThisAttempt === 1 && + attemptCursor !== undefined && + 'fromEventId' in attemptCursor && + 'fromEventId' in cursor && + attemptCursor.fromEventId === cursor.fromEventId; + subscribeCursor = cursor; + if (!replayedShell) { + consecutiveFailures = 0; + } + }, }; const flushPendingStreamUpdate = () => { @@ -637,6 +675,7 @@ export function useTrueForgeAgentMessages({ update, segmentStatus, ...(updateSequenceNumber != null ? { lastSequenceNumber: updateSequenceNumber } : {}), + ...(subscribeCursor != null ? { subscribeCursor } : {}), }, pendingMcpAuth, ...(runningTurn != null ? { runningTurn } : {}), @@ -663,6 +702,8 @@ export function useTrueForgeAgentMessages({ }; const readStream = (signal: AbortSignal): AsyncGenerator => { + eventsThisAttempt = 0; + attemptCursor = subscribeCursor; if (!useSubscribe) { return createStream(signal, track); } @@ -672,9 +713,10 @@ export function useTrueForgeAgentMessages({ turnIdRef.current, snapshotRef.current.fold, signal, - lastSequenceNumber, + subscribeCursor, snapshotRef.current.groupRootBaseline, track.onSequenceNumber, + track.onSubscribeCursor, ); }; @@ -756,6 +798,7 @@ export function useTrueForgeAgentMessages({ ...prev.activeStream, segmentStatus, ...(lastSequenceNumber != null ? { lastSequenceNumber } : {}), + ...(subscribeCursor != null ? { subscribeCursor } : {}), }, }); return commitActiveStream(marked); @@ -881,9 +924,10 @@ export function useTrueForgeAgentMessages({ turn.id, loadedSnapshot.fold, signal, - loadedSnapshot.activeStream?.lastSequenceNumber, + cursorForTurn(loadedSnapshot, turn.id), loadedSnapshot.groupRootBaseline, track.onSequenceNumber, + track.onSubscribeCursor, ), { current: turn.id }, { @@ -1083,6 +1127,7 @@ export function useTrueForgeAgentMessages({ groupRootBaseline, handleGatewayTurnId, track.onSequenceNumber, + track.onSubscribeCursor, ), turnIdRef, { @@ -1136,8 +1181,6 @@ export function useTrueForgeAgentMessages({ } const startSubscription = () => { - const latest = snapshotRef.current.activeStream; - const afterSequenceNumber = latest?.turnId === turnId ? latest.lastSequenceNumber : undefined; // Subscribe first so either valid server ordering (subscribe-before-send // or a concurrently arriving POST) cannot leave a replay gap. void runStream( @@ -1148,9 +1191,10 @@ export function useTrueForgeAgentMessages({ turnId, snapshotRef.current.fold, signal, - afterSequenceNumber, + cursorForTurn(snapshotRef.current, turnId), snapshotRef.current.groupRootBaseline, track.onSequenceNumber, + track.onSubscribeCursor, ), { current: turnId }, { @@ -1429,7 +1473,6 @@ export function useTrueForgeAgentMessages({ if (resumeSessionId == null) { return; } - const active = snapshotRef.current.activeStream; await runStream( (signal, track) => resumeTurnStream( @@ -1438,9 +1481,10 @@ export function useTrueForgeAgentMessages({ turn.id, snapshotRef.current.fold, signal, - active?.turnId === turn.id ? active.lastSequenceNumber : undefined, + cursorForTurn(snapshotRef.current, turn.id), snapshotRef.current.groupRootBaseline, track.onSequenceNumber, + track.onSubscribeCursor, ), { current: turn.id }, { diff --git a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts index 84194847d3..af4e105b21 100644 --- a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts +++ b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts @@ -1879,6 +1879,53 @@ describe('convertTurnMessages', () => { expect(onSequenceNumber).toHaveBeenCalledTimes(3); }); + it('resumes a model message from its id and every other event after its id', async () => { + const foldState = new PeerThreadFoldState(); + const cursors: { fromEventId?: string; afterEventId?: string }[] = []; + + await collectStream( + streamTurnEvents( + (async function* () { + yield { + event: { + type: 'model.message' as const, + createdAt, + id: 'm1', + threadId: ROOT_THREAD_ID, + content: 'hi', + }, + }; + yield { + event: { + type: 'mcp.auth_required' as const, + id: 'auth-1', + createdAt, + threadId: null, + mcpServers: [{ id: 'github', name: 'GitHub', authUrl: 'https://example.com/auth' }], + }, + }; + yield { + event: { + type: 'turn.done' as const, + id: 'turn-done', + createdAt, + state: { status: 'done' as const, completedAt: createdAt }, + }, + }; + })(), + foldState, + undefined, + undefined, + undefined, + cursor => { + cursors.push(cursor); + }, + ), + ); + + expect(cursors).toEqual([{ fromEventId: 'm1' }, { fromEventId: 'auth-1' }, { afterEventId: 'turn-done' }]); + }); + it('yields folded content after each ingested stream event', async () => { const foldState = new PeerThreadFoldState(); const updates = await collectStream( diff --git a/packages/assistant-ui-runtime/test/streamTurn.test.ts b/packages/assistant-ui-runtime/test/streamTurn.test.ts index 232e9e3035..13806d2f36 100644 --- a/packages/assistant-ui-runtime/test/streamTurn.test.ts +++ b/packages/assistant-ui-runtime/test/streamTurn.test.ts @@ -300,7 +300,7 @@ describe('streamTurn', () => { }); describe('resumeTurnStream', () => { - it('reconnects with afterSequenceNumber and yields updates', async () => { + it('reconnects with an event-id cursor and yields updates', async () => { const foldState = new PeerThreadFoldState(); const subscribeToTurn = vi.fn(async function* () { yield streamData(2, { @@ -323,12 +323,14 @@ describe('streamTurn', () => { }); const updates = await collectUpdates( - resumeTurnStream(server, SESSION_ID, 'turn-1', foldState, new AbortController().signal, 1), + resumeTurnStream(server, SESSION_ID, 'turn-1', foldState, new AbortController().signal, { + fromEventId: 'm1', + }), ); expect(subscribeToTurn).toHaveBeenCalledWith({ sessionId: SESSION_ID, turnId: 'turn-1', - afterSequenceNumber: 1, + fromEventId: 'm1', abortSignal: expect.any(AbortSignal), }); expect(updates[0]).toEqual({ content: [{ type: 'text', text: 'resumed' }], sequenceNumber: 2 }); @@ -384,7 +386,9 @@ describe('streamTurn', () => { const server = mockServer({ subscribeToTurn }); const updates = await collectUpdates( - resumeTurnStream(server, SESSION_ID, 'turn-1', new PeerThreadFoldState(), new AbortController().signal, 1), + resumeTurnStream(server, SESSION_ID, 'turn-1', new PeerThreadFoldState(), new AbortController().signal, { + afterEventId: 'prior', + }), ); expect(updates).toMatchObject([ diff --git a/packages/assistant-ui-runtime/test/useTrueForgeAgentMessages.test.tsx b/packages/assistant-ui-runtime/test/useTrueForgeAgentMessages.test.tsx index ca06e82db3..7b09dbb94e 100644 --- a/packages/assistant-ui-runtime/test/useTrueForgeAgentMessages.test.tsx +++ b/packages/assistant-ui-runtime/test/useTrueForgeAgentMessages.test.tsx @@ -498,6 +498,7 @@ describe('useTrueForgeAgentMessages', () => { expect.any(Array), expect.any(Function), expect.any(Function), + expect.any(Function), ); await act(async () => { @@ -514,6 +515,7 @@ describe('useTrueForgeAgentMessages', () => { expect.any(Array), expect.any(Function), expect.any(Function), + expect.any(Function), ); }); @@ -1958,7 +1960,7 @@ describe('useTrueForgeAgentMessages', () => { }); }); - it('uses the last ingested sequence when subscribing before an action', async () => { + it('uses the stored event-id cursor when subscribing before an action', async () => { const base = snapshotWithAssistantMessage(assistantMessageWithPendingApproval()); const active = base.activeStream; if (active == null) { @@ -1966,7 +1968,7 @@ describe('useTrueForgeAgentMessages', () => { } vi.mocked(loadSessionSnapshot).mockResolvedValue( replaceSessionSnapshot(base, { - activeStream: { ...active, lastSequenceNumber: 7 }, + activeStream: { ...active, subscribeCursor: { afterEventId: 'approval-required-1' } }, }), ); @@ -1983,18 +1985,30 @@ describe('useTrueForgeAgentMessages', () => { 'turn-1', expect.any(PeerThreadFoldState), expect.any(AbortSignal), - 7, + { afterEventId: 'approval-required-1' }, undefined, expect.any(Function), + expect.any(Function), ); }); it('queues a subscription when an action arrives while the paused segment is finishing', async () => { let releasePausedSegment: (() => void) | undefined; vi.mocked(streamTurnContent).mockImplementation( - (_server, _sessionId, _fold, _options, _signal, _baseline, onTurnIdAvailable) => + ( + _server, + _sessionId, + _fold, + _options, + _signal, + _baseline, + onTurnIdAvailable, + _onSequenceNumber, + onSubscribeCursor, + ) => (async function* () { onTurnIdAvailable?.('turn-paused'); + onSubscribeCursor?.({ afterEventId: 'approval-required-1' }); yield { content: [...assistantMessageWithPendingApproval().content], status: { type: 'requires-action', reason: 'tool-calls' }, @@ -2031,9 +2045,10 @@ describe('useTrueForgeAgentMessages', () => { 'turn-paused', expect.any(PeerThreadFoldState), expect.any(AbortSignal), - 4, + { afterEventId: 'approval-required-1' }, expect.any(Array), expect.any(Function), + expect.any(Function), ); }); @@ -2546,8 +2561,19 @@ describe('useTrueForgeAgentMessages', () => { const onError = vi.fn(); const onPreTurnFailure = vi.fn(); vi.mocked(streamTurnContent).mockImplementation( - async function* (_server, _sessionId, _fold, _options, _signal, _baseline, onTurnIdAvailable) { + async function* ( + _server, + _sessionId, + _fold, + _options, + _signal, + _baseline, + onTurnIdAvailable, + _onSequenceNumber, + onSubscribeCursor, + ) { onTurnIdAvailable?.('gateway-turn-123'); + onSubscribeCursor?.({ fromEventId: 'model-partial' }); yield { content: [{ type: 'text' as const, text: 'partial' }], sequenceNumber: 4 }; throw new Error('Mid-stream error'); }, @@ -2585,9 +2611,10 @@ describe('useTrueForgeAgentMessages', () => { 'gateway-turn-123', expect.anything(), expect.any(AbortSignal), - 4, + { fromEventId: 'model-partial' }, expect.anything(), expect.any(Function), + expect.any(Function), ); }); }); diff --git a/packages/trueforge-sdk/reference.md b/packages/trueforge-sdk/reference.md index c55e49ae59..888f671391 100644 --- a/packages/trueforge-sdk/reference.md +++ b/packages/trueforge-sdk/reference.md @@ -2921,7 +2921,7 @@ await client.sessions.createTurnEvents("session_id", "turn_id", {
-Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. +Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id
diff --git a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts index 44d0a87324..eaadc33a03 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts @@ -2140,7 +2140,7 @@ export class SessionsClient { } /** - * Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. + * Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id */ public subscribeToTurn( session_id: string, @@ -2159,9 +2159,11 @@ export class SessionsClient { request: TrueForge.SubscribeToTurnSessionsRequest = {}, requestOptions?: SessionsClient.RequestOptions, ): Promise>> { - const { afterSequenceNumber } = request; + const { afterSequenceNumber, fromEventId, afterEventId } = request; const _queryParams: Record = { after_sequence_number: afterSequenceNumber, + from_event_id: fromEventId, + after_event_id: afterEventId, }; const _authRequest: core.AuthRequest = await this._options.authProvider.getAuthRequest(); const _headers: core.Fetcher.Args["headers"] = mergeHeaders( diff --git a/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts b/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts index 03af7554c8..4aa5bc30e7 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts @@ -5,6 +5,10 @@ * {} */ export interface SubscribeToTurnSessionsRequest { - /** Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer. */ + /** Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. */ afterSequenceNumber?: number | null; + /** Inclusive cursor. Replay starts at the entry with this event id */ + fromEventId?: string; + /** Exclusive cursor. Replay starts after the last entry with this event id. */ + afterEventId?: string; } diff --git a/packages/trueforge-ui/src/plugins/trueforge-agent-server-adapter/chatServer.ts b/packages/trueforge-ui/src/plugins/trueforge-agent-server-adapter/chatServer.ts index e080ada06f..0a4b9054f1 100644 --- a/packages/trueforge-ui/src/plugins/trueforge-agent-server-adapter/chatServer.ts +++ b/packages/trueforge-ui/src/plugins/trueforge-agent-server-adapter/chatServer.ts @@ -157,10 +157,11 @@ export function toListResult( }; } -function sequenceNumber(id: string | undefined, fallback: number): number { - if (id === undefined) return fallback; +/** Integer SSE ids belong to older sequence streams. A missing id is not a cursor. */ +function sequenceNumberFromSseId(id: string | undefined): number | undefined { + if (id === undefined) return undefined; const parsed = Number(id); - return Number.isSafeInteger(parsed) && parsed >= 0 ? parsed : fallback; + return Number.isSafeInteger(parsed) && parsed >= 0 ? parsed : undefined; } function toHarnessContent(content: UserMessageContent): TrueForgeApi.UserMessageContent { @@ -255,16 +256,13 @@ export function createHarnessChatServer( ...(input === undefined ? {} : { input: toHarnessInput(input) }), ...(previousTurnId === undefined ? {} : { previousTurnId }), }); - let fallbackSequence = 0; for await (const item of stream.withMetadata()) { const event = toUiStreamingEvent(item.data); - if (event !== undefined) { - yield { - sequenceNumber: sequenceNumber(item.id, fallbackSequence), - event, - }; + if (event === undefined) { + continue; } - fallbackSequence += 1; + const sequenceNumber = sequenceNumberFromSseId(item.id); + yield sequenceNumber === undefined ? { event } : { sequenceNumber, event }; } }, @@ -275,29 +273,32 @@ export function createHarnessChatServer( return response.data.map(toUiInboundEvent); }, - /** Resume a live turn; omitted/0 `afterSequenceNumber` replays from the start. */ + /** Resume a live turn from an event id, or a legacy sequence on an older stream. */ async *subscribeToTurn({ sessionId, turnId, afterSequenceNumber, + fromEventId, + afterEventId, }: { sessionId: string; turnId: string; afterSequenceNumber?: number; + fromEventId?: string; + afterEventId?: string; }) { const stream = await client.sessions.subscribeToTurn(sessionId, turnId, { ...(afterSequenceNumber === undefined ? {} : { afterSequenceNumber }), + ...(fromEventId === undefined ? {} : { fromEventId }), + ...(afterEventId === undefined ? {} : { afterEventId }), }); - let fallbackSequence = 0; for await (const item of stream.withMetadata()) { const event = toUiStreamingEvent(item.data); - if (event !== undefined) { - yield { - sequenceNumber: sequenceNumber(item.id, fallbackSequence), - event, - }; + if (event === undefined) { + continue; } - fallbackSequence += 1; + const sequenceNumber = sequenceNumberFromSseId(item.id); + yield sequenceNumber === undefined ? { event } : { sequenceNumber, event }; } }, diff --git a/packages/trueforge-ui/test/plugins/trueforge-agent-server-adapter/harnessServer.test.ts b/packages/trueforge-ui/test/plugins/trueforge-agent-server-adapter/harnessServer.test.ts index 8f6cbddd93..552b55f8e6 100644 --- a/packages/trueforge-ui/test/plugins/trueforge-agent-server-adapter/harnessServer.test.ts +++ b/packages/trueforge-ui/test/plugins/trueforge-agent-server-adapter/harnessServer.test.ts @@ -417,6 +417,39 @@ describe('createHarnessChatServer', () => { assert.deepEqual(deletedSessions, ['ses_1']); }); + it('omits sequenceNumber when the SSE id is absent and forwards fromEventId', async () => { + const queries: string[] = []; + const eventId = '01arz3ndektsv4rrffq69g5fav'; + const fetchCursor: typeof fetch = async input => { + const url = input instanceof Request ? input.url : String(input); + if (url.includes('/api/v1/sessions/ses_1/turns/trn_1/subscribe')) { + queries.push(new URL(url, 'http://test.local').search); + return new Response( + `data: {"type":"model.message.delta","id":"${eventId}","thread_id":"main","content":"world"}\n\n`, + { headers: { 'content-type': 'text/event-stream' } }, + ); + } + return new Response(`Unexpected request: ${url}`, { status: 500 }); + }; + const server = createHarnessChatServer({ fetch: fetchCursor }); + if (server.subscribeToTurn === undefined) { + throw new Error('Expected subscribeToTurn'); + } + + const events = []; + for await (const event of server.subscribeToTurn({ + sessionId: 'ses_1', + turnId: 'trn_1', + fromEventId: eventId, + })) { + events.push(event); + } + + assert.match(queries.at(-1) ?? '', /from_event_id=01arz3ndektsv4rrffq69g5fav/); + assert.equal(events[0]?.sequenceNumber, undefined); + assert.equal(events[0]?.event.id, eventId); + }); + it('subscribes to a turn and forwards afterSequenceNumber', async () => { const server = createHarnessChatServer({ fetch: fetchMock }); assert.ok(server.subscribeToTurn); diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 513f6b2004..c184e93a15 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -84,11 +84,15 @@ import { } from '../routes/turnRoutes'; import type { ActiveTurnRegistry } from '../runtime/activeTurns'; import { + isLegacySequenceEntryId, + LEGACY_SEQUENCE_ID_MAX, StreamGoneError, StreamOwnerMismatchError, type EventSubscription, + type EventSubscriptionPollOptions, type EventSubscriptionRegistry, } from '../runtime/event-subscription'; +import { eventIdTimestampMs } from '../runtime/event-subscription/entryId'; import { validateSandboxFilePath } from '../runtime/sandboxFilePath'; import { buildTurnSandbox, @@ -473,9 +477,9 @@ export function deriveSessionTitle(input: UserMessage[] | undefined): string | u * sequence number; the event body itself is not numbered (yield order is * persist order, so the transport boundary stamps). */ -export function turnEventSsePayload(event: TurnStreamingEvent, sequenceNumber: number): { id: string; data: string } { +export function turnEventSsePayload(event: TurnStreamingEvent, sequenceNumber?: number): { id?: string; data: string } { return { - id: String(sequenceNumber), + ...(sequenceNumber === undefined ? {} : { id: String(sequenceNumber) }), data: JSON.stringify(event), }; } @@ -542,25 +546,6 @@ interface TurnExecutionTimeout { clear: () => void; } -async function recordPersistedEventSequence(input: { - sessionStore: Pick; - sessionId: string; - turnId: string; - executorId: string; - event: TurnStreamingEvent; - sequenceNumber: number; -}): Promise { - if (input.event.type === EventType.MODEL_MESSAGE_DELTA) { - return; - } - await input.sessionStore.assignTurnEventSequenceNumbers({ - session_id: input.sessionId, - turn_id: input.turnId, - active_executor_id: input.executorId, - assignments: [{ event_id: input.event.id, sequence_number: input.sequenceNumber }], - }); -} - export async function drainTurnEvents(input: { trackedStream: AsyncIterable; turnEventStream: EventSubscription; @@ -568,29 +553,17 @@ export async function drainTurnEvents(input: { logger: Logger; pausedKeepalive: TurnPausedKeepalive; executorId: string; - /** Writes the sequence put just assigned onto the durable event row. */ - recordSequence: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; - onEvent?: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; + onEvent?: (event: TurnStreamingEvent) => Promise; }): Promise { - const { - trackedStream, - turnEventStream, - executionTimeout, - logger, - pausedKeepalive, - executorId, - recordSequence, - onEvent, - } = input; + const { trackedStream, turnEventStream, executionTimeout, logger, pausedKeepalive, executorId, onEvent } = input; try { for await (const event of trackedStream) { reportUnexpectedTurnDone({ event, logger }); // Dual-write before any client sink so subscribers can resume after disconnect. - const sequenceNumber = await turnEventStream.put(event, { + await turnEventStream.put(event, { streamTTLSeconds: streamTTLSecondsFor(event), executorId, }); - await recordSequence(event, sequenceNumber); if (event.type === EventType.TURN_UPDATE) { if (event.state.status === 'paused') { pausedKeepalive.start(); @@ -600,7 +573,7 @@ export async function drainTurnEvents(input: { executionTimeout.start(); } } - await onEvent?.(event, sequenceNumber); + await onEvent?.(event); } } catch (error) { if ( @@ -631,7 +604,6 @@ export interface TurnEventDrainInput { logger: Logger; pausedKeepalive: TurnPausedKeepalive; executorId: string; - recordSequence: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; } /** @@ -755,7 +727,7 @@ export async function beginTurnExecution( turn, }); - // Held for the whole turn. Repair reseeds the counter when this replica resumes. + // Held for the whole turn. Repair appends any missing log suffix with event-id stream ids. const turnEventStream = deps.eventSubscriptions.get(turnStreamId(tenantId, sessionId, turn.id)); return { @@ -774,15 +746,6 @@ export async function beginTurnExecution( }, }, executorId: configuration.EXECUTOR_ID, - recordSequence: (event, sequenceNumber) => - recordPersistedEventSequence({ - sessionStore: deps.sessionStore, - sessionId, - turnId: turn.id, - executorId: configuration.EXECUTOR_ID, - event, - sequenceNumber, - }), }, }; } @@ -939,15 +902,6 @@ async function resumePausedTurn(input: { executionTimeout, logger: input.deps.logger, executorId: configuration.EXECUTOR_ID, - recordSequence: (event, sequenceNumber) => - recordPersistedEventSequence({ - sessionStore: input.deps.sessionStore, - sessionId: input.session.session_id, - turnId: turn.id, - executorId: configuration.EXECUTOR_ID, - event, - sequenceNumber, - }), pausedKeepalive: { start: () => { input.deps.activeTurns.startPausedKeepalive({ @@ -1943,11 +1897,11 @@ export function createTurnsRouter(deps: TurnsRouterDeps { + onEvent: async event => { if (!stream.closed && !stream.aborted && shouldWriteToSSEStream) { try { await withTimeout( - stream.writeSSE(turnEventSsePayload(event, sequenceNumber)), + stream.writeSSE(turnEventSsePayload(event)), configuration.TURN_SSE_STREAM_WRITE_TIMEOUT_MS, ); } catch (error) { @@ -2050,7 +2004,26 @@ export function createTurnsRouter(deps: TurnsRouterDeps = async c => { const { session_id: sessionId, turn_id: turnId } = c.req.valid('param'); const query = c.req.valid('query'); + const fromEventId = query.from_event_id; + const afterEventId = query.after_event_id; const afterSequenceNumber = resolveAfterSequenceNumber(c, query.after_sequence_number); + // we recheck here outside of zod because we support Last-Event-ID header for legacy clients + if (afterSequenceNumber !== undefined && (fromEventId !== undefined || afterEventId !== undefined)) { + throw new HTTPException(400, { + message: 'Pass only one of after_sequence_number, from_event_id, or after_event_id.', + }); + } + if (afterSequenceNumber !== undefined && afterSequenceNumber > LEGACY_SEQUENCE_ID_MAX) { + throw new HTTPException(400, { + message: 'after_sequence_number is not supported above 50000.', + }); + } + if (fromEventId !== undefined && eventIdTimestampMs(fromEventId) === undefined) { + throw new HTTPException(400, { message: 'Invalid event id.' }); + } + if (afterEventId !== undefined && eventIdTimestampMs(afterEventId) === undefined) { + throw new HTTPException(400, { message: 'Invalid event id.' }); + } const requestContext = deps.resolveRequestContext(c); const session = await deps.sessions.get({ @@ -2098,6 +2071,17 @@ export function createTurnsRouter(deps: TurnsRouterDeps(deps: TurnsRouterDeps { stream.onAbort(() => { getOperationMetrics()?.recordTurnSseDisconnect({ path: 'subscribe' }); subscribeAbort.abort(new Error('client-disconnected')); }); - const generator = turnEventStream.poll(afterSequenceNumber, { signal: subscribeAbort.signal }); + const generator = turnEventStream.poll(pollOptions); try { for await (const { sequence_number: sequenceNumber, ...event } of generator) { - await stream.writeSSE(turnEventSsePayload(event, sequenceNumber)); + await stream.writeSSE( + includeSequenceId && sequenceNumber !== undefined + ? turnEventSsePayload(event, sequenceNumber) + : turnEventSsePayload(event), + ); if (event.type === EventType.TURN_DONE) { break; } diff --git a/packages/trueforge/src/routes/turnRoutes.ts b/packages/trueforge/src/routes/turnRoutes.ts index 082d628fd1..7bfb4e6ee5 100644 --- a/packages/trueforge/src/routes/turnRoutes.ts +++ b/packages/trueforge/src/routes/turnRoutes.ts @@ -296,7 +296,7 @@ export const subscribeTurnRoute = createRoute({ tags: [OpenApiTag.AGENT_SESSIONS], summary: 'Subscribe to a turn stream', description: - 'Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence.', + 'Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id', 'x-fern-sdk-group-name': ['sessions'], 'x-fern-sdk-method-name': 'subscribe_to_turn', 'x-fern-streaming': { format: 'sse', resumable: true }, diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 82072be54b..0531553593 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -108,7 +108,7 @@ export class ActiveTurnRegistry { if (caught === E_TIMEOUT) { throw new TurnLockAcquireTimeoutError({ timeoutMs: this.lockAcquireTimeoutMs, - cause: caught, + cause: E_TIMEOUT, }); } throw caught; diff --git a/packages/trueforge/src/runtime/event-subscription/entryId.ts b/packages/trueforge/src/runtime/event-subscription/entryId.ts new file mode 100644 index 0000000000..5393d31f86 --- /dev/null +++ b/packages/trueforge/src/runtime/event-subscription/entryId.ts @@ -0,0 +1,134 @@ +import { decodeTime } from 'ulid'; + +/** + * Old streams used `{sequence}-1` with a small sequence. A unix-millisecond id + * is far above this. The cutoff is a temporary compat check, not an event limit. + */ +export const LEGACY_SEQUENCE_ID_MAX = 50_000; + +/** Redis stream sequence is an unsigned 64-bit integer. The id just before `{ms}-0` is `{ms-1}` plus this. */ +const REDIS_MAX_SEQUENCE = '18446744073709551615'; + +/** Millisecond from a ULID event id. Other ids have no timestamp. */ +export function eventIdTimestampMs(eventId: string): number | undefined { + try { + return decodeTime(eventId); + } catch { + return undefined; + } +} + +export function entryIdMilliseconds(entryId: string): number { + const separator = entryId.indexOf('-'); + const raw = separator === -1 ? entryId : entryId.slice(0, separator); + const milliseconds = Number(raw); + if (!Number.isSafeInteger(milliseconds) || milliseconds < 0) { + throw new Error(`Corrupt stream entry id ${entryId}`); + } + return milliseconds; +} + +export function isLegacySequenceEntryId(entryId: string): boolean { + return entryIdMilliseconds(entryId) <= LEGACY_SEQUENCE_ID_MAX; +} + +/** + * Where an event-id scan starts. A leftover `{sequence}-1` stream is ordered by + * sequence, not by the payload ULID, so the scan has to walk it from the start. + */ +export function eventIdSeekCursor(input: { eventId: string; firstEntryId: string | undefined }): string { + if (input.firstEntryId !== undefined && isLegacySequenceEntryId(input.firstEntryId)) { + return '0-0'; + } + const timestamp = eventIdTimestampMs(input.eventId); + if (timestamp === undefined) { + return '0-0'; + } + return exclusiveCursorBeforeMillisecond(timestamp); +} + +/** Exclusive XREAD cursor that includes every entry at `milliseconds` and later. */ +export function exclusiveCursorBeforeMillisecond(milliseconds: number): string { + if (milliseconds <= 0) { + return '0-0'; + } + return `${String(milliseconds - 1)}-${REDIS_MAX_SEQUENCE}`; +} + +function entryIdSequence(entryId: string): bigint { + const separator = entryId.indexOf('-'); + const raw = separator === -1 ? '0' : entryId.slice(separator + 1); + return BigInt(raw); +} + +/** Redis stream order: millisecond, then the per-millisecond counter. */ +export function compareEntryIds(left: string, right: string): number { + const leftMs = entryIdMilliseconds(left); + const rightMs = entryIdMilliseconds(right); + if (leftMs !== rightMs) { + return leftMs < rightMs ? -1 : 1; + } + const leftSequence = entryIdSequence(left); + const rightSequence = entryIdSequence(right); + if (leftSequence === rightSequence) { + return 0; + } + return leftSequence < rightSequence ? -1 : 1; +} + +export function isEntryAfterCursor(entryId: string, cursor: string): boolean { + return compareEntryIds(entryId, cursor) > 0; +} + +/** + * `from` yields at the first payload id match. `after` yields after the last + * match already buffered, then skips a later copy of that id. An id missing + * from the buffer is dropped and later entries are yielded. + */ +export class EventIdSeek { + private phase: 'scan' | 'yield' = 'scan'; + private readonly buffered: T[] = []; + + constructor( + private readonly mode: 'from' | 'after', + private readonly eventId: string, + ) {} + + push(event: T): T[] { + if (this.phase === 'yield') { + if (this.mode === 'after' && event.id === this.eventId) { + return []; + } + return [event]; + } + this.buffered.push(event); + if (this.mode === 'from' && event.id === this.eventId) { + this.phase = 'yield'; + const start = this.buffered.findIndex(item => item.id === this.eventId); + const emit = this.buffered.slice(start); + this.buffered.length = 0; + return emit; + } + return []; + } + + caughtUp(): T[] { + if (this.phase === 'yield') { + return []; + } + this.phase = 'yield'; + if (this.mode === 'from') { + this.buffered.length = 0; + return []; + } + let lastMatch = -1; + for (let index = 0; index < this.buffered.length; index += 1) { + if (this.buffered[index]?.id === this.eventId) { + lastMatch = index; + } + } + const emit = lastMatch === -1 ? [] : this.buffered.slice(lastMatch + 1); + this.buffered.length = 0; + return emit; + } +} diff --git a/packages/trueforge/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index c29fc89bc3..ca9d5bae05 100644 --- a/packages/trueforge/src/runtime/event-subscription/inMemory.ts +++ b/packages/trueforge/src/runtime/event-subscription/inMemory.ts @@ -1,6 +1,6 @@ /** - * Single-process backend mirroring RedisEventSubscription semantics (dense - * sequences, whole-stream TTL, poll-forever generators, same errors). + * Single-process backend mirroring RedisEventSubscription semantics + * (ULID-millisecond entry ids, whole-stream TTL, poll-forever generators). * Resume only works within one replica. */ import { EventEmitter, once } from 'node:events'; @@ -13,11 +13,24 @@ import { type EventSubscriptionPutOptions, type SequencedEvent, } from '.'; -import { applyEventLogRepair, eventIdOf, type RepairLogOptions } from './repair'; +import { + entryIdMilliseconds, + EventIdSeek, + eventIdSeekCursor, + eventIdTimestampMs, + isEntryAfterCursor, + isLegacySequenceEntryId, +} from './entryId'; +import { applyEventLogRepair, type RepairLogOptions } from './repair'; + +interface StoredEvent { + entryId: string; + event: T; +} -interface InMemoryStream { - /** Appended in increasing sequence_number order. */ - events: SequencedEvent[]; +interface InMemoryStream { + /** Appended in Redis stream-id order. */ + events: StoredEvent[]; /** Undefined = no TTL set yet (mirrors a Redis key without EXPIRE). */ expiresAtMs?: number | undefined; expiryTimer?: NodeJS.Timeout | undefined; @@ -27,7 +40,7 @@ interface InMemoryStream { * Owns the stream logs and their TTLs; the expiry timer is the in-process * equivalent of Redis key eviction. */ -export class InMemoryEventStreamStore { +export class InMemoryEventStreamStore { private readonly streams = new Map>(); /** Survives stream eviction so a stale producer cannot recreate a missing log. */ private readonly owners = new Map(); @@ -44,34 +57,32 @@ export class InMemoryEventStreamStore { } /** - * Appends at an explicit sequence. A sequence at or below the tail is a duplicate - * when the tail already has that event id, and a conflict otherwise. + * Appends with a `{millisecond}-{n}` id. The same payload already at the tail + * is a duplicate. A millisecond older than the tail reuses the tail millisecond. */ - appendExplicit(input: { + appendEvent(input: { streamId: string; event: T; - sequenceNumber: number; - eventId: string | undefined; executorId: string | undefined; checkOwner: boolean; streamTTLSeconds: number | undefined; - }): 'ok' | 'duplicate' | 'owner_mismatch' | 'conflict' { + }): 'ok' | 'duplicate' | 'owner_mismatch' { const owner = this.owners.get(input.streamId); if (input.checkOwner && owner !== undefined && owner !== input.executorId) { return 'owner_mismatch'; } const stream = this.getLiveStream(input.streamId) ?? { events: [] }; + const event = structuredClone(input.event); + const payload = JSON.stringify(event); const last = stream.events[stream.events.length - 1]; - if (last !== undefined && input.sequenceNumber <= last.sequence_number) { - const tailId = eventIdOf(last); - if (input.eventId !== undefined && tailId === input.eventId) { - return 'duplicate'; - } - return 'conflict'; + if (last !== undefined && JSON.stringify(last.event) === payload) { + return 'duplicate'; } this.streams.set(input.streamId, stream); - // Deep-copy so the log is isolated from caller mutations (mirrors JSON round-trip in Redis). - stream.events.push({ ...structuredClone(input.event), sequence_number: input.sequenceNumber }); + stream.events.push({ + entryId: this.nextEntryId(stream, eventIdTimestampMs(event.id) ?? Date.now()), + event, + }); if (input.streamTTLSeconds !== undefined && input.streamTTLSeconds > 0) { stream.expiresAtMs = Date.now() + input.streamTTLSeconds * 1_000; this.scheduleExpiry(input.streamId, stream); @@ -80,13 +91,17 @@ export class InMemoryEventStreamStore { return 'ok'; } - readTail(streamId: string): { eventId: string; sequenceNumber: number } | undefined { + readTail(streamId: string): { eventId: string; payload: string } | undefined { const stream = this.getLiveStream(streamId); const last = stream?.events[stream.events.length - 1]; if (last === undefined) { return undefined; } - return { eventId: eventIdOf(last) ?? '', sequenceNumber: last.sequence_number }; + return { eventId: last.event.id, payload: JSON.stringify(last.event) }; + } + + firstEntryId(streamId: string): string | undefined { + return this.getLiveStream(streamId)?.events[0]?.entryId; } /** Drops the log and keeps the owner, matching Redis stream-key expiry. */ @@ -161,12 +176,28 @@ export class InMemoryEventStreamStore { ); stream.expiryTimer.unref(); } + + /** + * Redis `XADD -*` : counter 0 when the millisecond is newer than the last + * id, otherwise the last millisecond's counter plus one. + */ + private nextEntryId(stream: InMemoryStream, milliseconds: number): string { + const last = stream.events[stream.events.length - 1]; + if (last === undefined) { + return `${String(milliseconds)}-0`; + } + const lastMilliseconds = entryIdMilliseconds(last.entryId); + if (milliseconds > lastMilliseconds) { + return `${String(milliseconds)}-0`; + } + const separator = last.entryId.indexOf('-'); + const lastCounter = Number(last.entryId.slice(separator + 1)); + return `${String(lastMilliseconds)}-${String(lastCounter + 1)}`; + } } /** Stream-scoped view over the shared store; implements the transport contract. */ -export class InMemoryEventSubscription implements EventSubscription { - private nextSequenceNumber = 1; - +export class InMemoryEventSubscription implements EventSubscription { constructor( private readonly store: InMemoryEventStreamStore, private readonly streamId: string, @@ -176,7 +207,6 @@ export class InMemoryEventSubscription implements EventSubscri const executorId = options.executorId; return applyEventLogRepair({ entries: options.entries, - persistSequences: options.persistSequences, claimOwner: executorId === undefined ? undefined @@ -186,46 +216,30 @@ export class InMemoryEventSubscription implements EventSubscri }, readTail: () => Promise.resolve(this.store.readTail(this.streamId)), append: entry => { - const outcome = this.store.appendExplicit({ + const outcome = this.store.appendEvent({ streamId: this.streamId, event: entry.event, - sequenceNumber: entry.sequence_number, - eventId: entry.id, executorId, checkOwner: false, streamTTLSeconds: options.streamTTLSeconds, }); - if (outcome === 'ok') { - return Promise.resolve('appended'); - } - if (outcome === 'duplicate') { - return Promise.resolve('duplicate'); - } if (outcome === 'owner_mismatch') { return Promise.reject(new StreamOwnerMismatchError(this.streamId)); } - return Promise.reject( - new Error(`explicit stream id ${String(entry.sequence_number)} conflicts on ${this.streamId}`), - ); + return Promise.resolve(outcome === 'duplicate' ? 'duplicate' : 'appended'); }, refreshTtl: () => { this.store.refreshStreamTtl(this.streamId, options.streamTTLSeconds); return Promise.resolve(); }, refreshTtlOnUntouchedStream: executorId !== undefined, - }).then(next => { - this.nextSequenceNumber = next; }); } - put(event: T, options?: EventSubscriptionPutOptions): Promise { - const sequenceNumber = this.nextSequenceNumber; - this.nextSequenceNumber += 1; - const outcome = this.store.appendExplicit({ + put(event: T, options?: EventSubscriptionPutOptions): Promise { + const outcome = this.store.appendEvent({ streamId: this.streamId, event, - sequenceNumber, - eventId: eventIdOf(event), executorId: options?.executorId, checkOwner: true, streamTTLSeconds: options?.streamTTLSeconds, @@ -233,10 +247,11 @@ export class InMemoryEventSubscription implements EventSubscri if (outcome === 'owner_mismatch') { return Promise.reject(new StreamOwnerMismatchError(this.streamId)); } - if (outcome === 'conflict') { - return Promise.reject(new Error(`explicit stream id ${String(sequenceNumber)} conflicts on ${this.streamId}`)); - } - return Promise.resolve(sequenceNumber); + return Promise.resolve(); + } + + firstEntryId(): Promise { + return Promise.resolve(this.store.firstEntryId(this.streamId)); } assertSubscribable(): Promise { @@ -251,13 +266,22 @@ export class InMemoryEventSubscription implements EventSubscri return Promise.resolve(); } - async *poll( - afterSequenceNumber?: number, - options?: EventSubscriptionPollOptions, - ): AsyncGenerator, void, unknown> { + async *poll(options?: EventSubscriptionPollOptions): AsyncGenerator, void, unknown> { + const afterSequenceNumber = options?.afterSequenceNumber; + const fromEventId = options?.fromEventId; + const afterEventId = options?.afterEventId; + const cursors = [afterSequenceNumber, fromEventId, afterEventId].filter(value => value !== undefined); + if (cursors.length > 1) { + throw new Error('Pass only one of afterSequenceNumber, fromEventId, or afterEventId'); + } + if (fromEventId !== undefined || afterEventId !== undefined) { + yield* this.pollByEventId(options ?? {}); + return; + } const signal = options?.signal; - // 0 and omitted both mean "from the start". Yield events whose sequence is strictly greater. - let after = afterSequenceNumber === undefined || afterSequenceNumber === 0 ? 0 : afterSequenceNumber; + const readCursor = + afterSequenceNumber === undefined || afterSequenceNumber === 0 ? '0-0' : `${String(afterSequenceNumber + 1)}-0`; + let index = 0; for (;;) { if (signal?.aborted) { return; @@ -266,8 +290,8 @@ export class InMemoryEventSubscription implements EventSubscri if (!live) { throw new StreamGoneError(this.streamId); } - const pending = live.events.filter(event => event.sequence_number > after); - if (pending.length === 0) { + const end = live.events.length; + if (index >= end) { try { await this.store.waitForChange(this.streamId, signal); } catch (error) { @@ -278,11 +302,84 @@ export class InMemoryEventSubscription implements EventSubscri } continue; } - // Deep-copy each yielded event so consumers cannot mutate the stored log. - for (const event of pending) { - after = event.sequence_number; - yield structuredClone(event); + while (index < end) { + const stored = live.events[index]; + index += 1; + if (stored === undefined) { + throw new Error(`Corrupt stream ${this.streamId}: missing event at index ${String(index - 1)}`); + } + if (!isEntryAfterCursor(stored.entryId, readCursor)) { + continue; + } + yield this.polled(stored); } } } + + private async *pollByEventId( + options: EventSubscriptionPollOptions, + ): AsyncGenerator, void, unknown> { + const fromEventId = options.fromEventId; + const afterEventId = options.afterEventId; + const eventId = fromEventId ?? afterEventId; + if (eventId === undefined) { + throw new Error('Pass fromEventId or afterEventId'); + } + const signal = options.signal; + const seek = new EventIdSeek>(fromEventId !== undefined ? 'from' : 'after', eventId); + const readCursor = eventIdSeekCursor({ eventId, firstEntryId: this.store.firstEntryId(this.streamId) }); + let index = 0; + let scanning = true; + for (;;) { + if (signal?.aborted) { + return; + } + const live = this.store.getLiveStream(this.streamId); + if (!live) { + throw new StreamGoneError(this.streamId); + } + const end = live.events.length; + if (index >= end) { + if (scanning) { + scanning = false; + for (const event of seek.caughtUp()) { + yield structuredClone(event); + } + // An append during the replay yield happened before this poller was + // waiting. Look at the log again before parking. + continue; + } + try { + await this.store.waitForChange(this.streamId, signal); + } catch (error) { + if (signal?.aborted) { + return; + } + throw error; + } + continue; + } + while (index < end) { + const stored = live.events[index]; + index += 1; + if (stored === undefined) { + throw new Error(`Corrupt stream ${this.streamId}: missing event at index ${String(index - 1)}`); + } + if (!isEntryAfterCursor(stored.entryId, readCursor)) { + continue; + } + for (const event of seek.push(this.polled(stored))) { + yield event; + } + } + } + } + + private polled(stored: StoredEvent): SequencedEvent { + const event = structuredClone(stored.event); + if (!isLegacySequenceEntryId(stored.entryId)) { + return event; + } + return { ...event, sequence_number: entryIdMilliseconds(stored.entryId) }; + } } diff --git a/packages/trueforge/src/runtime/event-subscription/index.ts b/packages/trueforge/src/runtime/event-subscription/index.ts index a7d130182b..5d1adaa468 100644 --- a/packages/trueforge/src/runtime/event-subscription/index.ts +++ b/packages/trueforge/src/runtime/event-subscription/index.ts @@ -15,27 +15,55 @@ export interface EventSubscriptionPutOptions { executorId?: string | undefined; } -export type SequencedEvent = T & { sequence_number: number }; +/** `sequence_number` is set only when replaying an old `{sequence}-1` entry. */ +export type SequencedEvent = T & { sequence_number?: number }; -export interface EventSubscriptionPollOptions { +export { isLegacySequenceEntryId, LEGACY_SEQUENCE_ID_MAX } from './entryId'; + +export type EventSubscriptionPollOptions = { /** Ends the generator promptly on abort, releasing any parked idle wait. */ signal?: AbortSignal | undefined; -} +} & ( + | { + /** + * Legacy exclusive cursor. Replay events after this sequence (1-indexed). + * Omit or pass 0 to start from the first event. + */ + afterSequenceNumber?: number | undefined; + fromEventId?: never; + afterEventId?: never; + } + | { + /** Inclusive payload cursor. Replay starts at this event id. */ + fromEventId: string; + afterSequenceNumber?: never; + afterEventId?: never; + } + | { + /** Exclusive payload cursor. Replay starts after the last entry with this event id. */ + afterEventId: string; + afterSequenceNumber?: never; + fromEventId?: never; + } +); /** * One turn's resumable event stream, obtained from {@link EventSubscriptionRegistry.get}. * Per-stream state (e.g. the sequence counter) lives on the instance and is freed with it. */ -export interface EventSubscription { +export interface EventSubscription { /** - * Rebuilds a missing or behind stream from the durable log, then reseeds the - * producer counter. An existing stream whose tail event id is the last log id - * is left in place. Pass `executorId` only for the database owner that will put. + * Rebuilds a missing or behind stream from the durable log using event-id + * stream ids. An existing stream whose tail is already the log tail is left + * in place. Pass `executorId` only for the database owner that will put. */ repairLog(options: RepairLogOptions): Promise; - /** Appends an event, assigns its sequence number, and returns that number. */ - put(event: T, options?: EventSubscriptionPutOptions): Promise; + /** Appends an event. The stream id is `{ulidMillisecond}-{n}` when the event id is a ULID. */ + put(event: T, options?: EventSubscriptionPutOptions): Promise; + + /** First entry id, or undefined when the stream has no entries yet. */ + firstEntryId(): Promise; /** * Subscription admission check, run before SSE headers are sent: rejects with @@ -44,21 +72,15 @@ export interface EventSubscription { */ assertSubscribable(): Promise; - /** - * Yields events strictly after the supplied sequence number. - * Sequences are 1-indexed; omit or pass 0 to start from the first event. - */ - poll( - afterSequenceNumber?: number, - options?: EventSubscriptionPollOptions, - ): AsyncGenerator, void, unknown>; + /** Yields events from the chosen cursor, or from the start when none is set. */ + poll(options?: EventSubscriptionPollOptions): AsyncGenerator, void, unknown>; } /** * Hands out stream-scoped {@link EventSubscription}s: Redis-backed when a * client is supplied, otherwise views over one shared in-process store. */ -export class EventSubscriptionRegistry { +export class EventSubscriptionRegistry { /** One store for the whole process so producers and subscribers share streams. */ private readonly memoryStore = new InMemoryEventStreamStore(); diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index d45fcc9828..e13a924ddd 100644 --- a/packages/trueforge/src/runtime/event-subscription/redis.ts +++ b/packages/trueforge/src/runtime/event-subscription/redis.ts @@ -9,17 +9,66 @@ import { type EventSubscriptionPutOptions, type SequencedEvent, } from '.'; -import { applyEventLogRepair, eventIdFromStreamData, eventIdOf, type RepairLogOptions } from './repair'; +import { entryIdMilliseconds, eventIdSeekCursor, eventIdTimestampMs, isLegacySequenceEntryId } from './entryId'; +import { applyEventLogRepair, eventIdFromStreamData, type RepairLogOptions } from './repair'; const SUBSCRIBE_STREAM_POLL_ITEMS_COUNT = 100; /** Delay between XREAD polls while a stream has no new events (a blocking XREAD would monopolize the shared connection). */ const SUBSCRIBE_STREAM_POLL_SLEEP_INTERVAL_MS = 1_000; +/** + * Jump from ARGV[3], walk forward, return an exclusive XREAD cursor. + * TS does the batch XREADs separately. + * + * KEYS[1]=stream ARGV[1]=from|after ARGV[2]=payload id ARGV[3]=start ARGV[4]=batch + * + * from → stop at the first match; return the prior entry (XREAD is exclusive). + * after → walk the contiguous shared-id group (message + deltas); return its last entry. + * missing → tip (live-only). + */ +const SEEK_EVENT_ID_CURSOR = ` +local mode, targetId, cursor, batch = ARGV[1], ARGV[2], ARGV[3], tonumber(ARGV[4]) +local startBound = '(' .. cursor +local previous, lastMatch = nil, nil +while true do + local entries = redis.call('XRANGE', KEYS[1], startBound, '+', 'COUNT', batch) + if #entries == 0 then break end + for i = 1, #entries do + local entryId, fields = entries[i][1], entries[i][2] + local data = nil + for j = 1, #fields, 2 do + if fields[j] == 'data' then data = fields[j + 1]; break end + end + local payloadId = nil + if data ~= nil then + local ok, obj = pcall(cjson.decode, data) + if ok and type(obj) == 'table' and type(obj['id']) == 'string' then + payloadId = obj['id'] + end + end + if mode == 'from' then + if payloadId == targetId then + return previous or cursor + end + elseif payloadId == targetId then + lastMatch = entryId + elseif lastMatch ~= nil then + return lastMatch + end + previous = entryId + end + startBound = '(' .. entries[#entries][1] + if #entries < batch then break end +end +return lastMatch or previous or cursor +`; + /** * One script for put and for suffix backfill. Put checks the stream-owner key. * Repair passes check-owner 0 so a subscribe can append without claiming. - * An explicit id at or below the tail is a conflict; the caller treats it as a - * no-op when that tail already holds the same event id. + * The id is `{millisecond}-*`. Redis 7 assigns the counter. The same payload + * already at the tail is a duplicate. A millisecond older than the tail is a + * conflict; the caller retries with the tail millisecond. */ const APPEND_STREAM_ENTRY = ` local checkOwner = ARGV[1] @@ -29,6 +78,15 @@ if checkOwner == '1' then return {'owner_mismatch', ''} end end +local tail = redis.call('XREVRANGE', KEYS[2], '+', '-', 'COUNT', 1) +if tail[1] ~= nil then + local fields = tail[1][2] + for i = 1, #fields, 2 do + if fields[i] == 'data' and fields[i + 1] == ARGV[4] then + return {'duplicate', ''} + end + end +end local added = redis.pcall('XADD', KEYS[2], ARGV[3], 'data', ARGV[4]) if type(added) == 'table' and added['err'] ~= nil then local err = added['err'] @@ -55,24 +113,16 @@ end return {'ok', ''} `; -function sequenceNumberFromEntryId(streamId: string, entryId: string): number { - const separator = entryId.indexOf('-'); - const raw = separator === -1 ? entryId : entryId.slice(0, separator); - const sequenceNumber = Number(raw); - if (!Number.isSafeInteger(sequenceNumber) || sequenceNumber < 0) { - throw new Error(`Corrupt stream entry ${entryId} on ${streamId}: invalid sequence number`); - } - return sequenceNumber; -} - function streamOwnerKey(streamId: string): string { return `${streamId}:owner`; } -type AppendStatus = 'ok' | 'conflict' | 'owner_mismatch' | 'error'; +type AppendStatus = 'ok' | 'duplicate' | 'conflict' | 'owner_mismatch' | 'error'; function isAppendStatus(value: unknown): value is AppendStatus { - return value === 'ok' || value === 'conflict' || value === 'owner_mismatch' || value === 'error'; + return ( + value === 'ok' || value === 'duplicate' || value === 'conflict' || value === 'owner_mismatch' || value === 'error' + ); } function parseAppendReply(value: unknown): { status: AppendStatus; detail: string } { @@ -88,12 +138,10 @@ function parseAppendReply(value: unknown): { status: AppendStatus; detail: strin } /** - * One Redis Stream. Sequence numbers are explicit ids (`n-1`). The owner key is - * separate from the stream so a missing stream still rejects a stale producer. + * One Redis Stream. Entry ids are `{ulidMillisecond}-{n}` from `XADD`. The owner + * key is separate from the stream so a missing stream still rejects a stale producer. */ -export class RedisEventSubscription implements EventSubscription { - private nextSequenceNumber = 1; - +export class RedisEventSubscription implements EventSubscription { constructor( private readonly redis: RedisClient, private readonly streamId: string, @@ -101,44 +149,37 @@ export class RedisEventSubscription implements EventSubscripti async repairLog(options: RepairLogOptions): Promise { const executorId = options.executorId; - this.nextSequenceNumber = await applyEventLogRepair({ + await applyEventLogRepair({ entries: options.entries, - persistSequences: options.persistSequences, claimOwner: executorId === undefined ? undefined : () => this.redis.set(streamOwnerKey(this.streamId), executorId).then(() => undefined), readTail: () => this.readTail(), - append: async entry => { - const reply = await this.appendEntry({ - sequenceNumber: entry.sequence_number, - payload: JSON.stringify({ ...entry.event, sequence_number: entry.sequence_number }), + append: entry => + this.appendEvent({ + event: entry.event, checkOwner: false, executorId: '', streamTTLSeconds: options.streamTTLSeconds, - }); - return this.outcomeForEvent({ reply, eventId: entry.id, sequenceNumber: entry.sequence_number }); - }, + }), refreshTtl: () => this.refreshTtl(options.streamTTLSeconds), refreshTtlOnUntouchedStream: executorId !== undefined, }); } - async put(event: T, options?: EventSubscriptionPutOptions): Promise { - const sequenceNumber = this.nextSequenceNumber; - this.nextSequenceNumber += 1; - const reply = await this.appendEntry({ - sequenceNumber, - payload: JSON.stringify({ ...event, sequence_number: sequenceNumber }), + async put(event: T, options?: EventSubscriptionPutOptions): Promise { + await this.appendEvent({ + event, checkOwner: true, executorId: options?.executorId ?? '', streamTTLSeconds: options?.streamTTLSeconds, }); - if (reply.status === 'owner_mismatch') { - throw new StreamOwnerMismatchError(this.streamId); - } - await this.outcomeForEvent({ reply, eventId: eventIdOf(event), sequenceNumber }); - return sequenceNumber; + } + + async firstEntryId(): Promise { + const entries = await this.redis.xRange(this.streamId, '-', '+', { COUNT: 1 }); + return entries[0]?.id; } async assertSubscribable(): Promise { @@ -153,10 +194,18 @@ export class RedisEventSubscription implements EventSubscripti } } - async *poll( - afterSequenceNumber?: number, - options?: EventSubscriptionPollOptions, - ): AsyncGenerator, void, unknown> { + async *poll(options?: EventSubscriptionPollOptions): AsyncGenerator, void, unknown> { + const afterSequenceNumber = options?.afterSequenceNumber; + const fromEventId = options?.fromEventId; + const afterEventId = options?.afterEventId; + const cursors = [afterSequenceNumber, fromEventId, afterEventId].filter(value => value !== undefined); + if (cursors.length > 1) { + throw new Error('Pass only one of afterSequenceNumber, fromEventId, or afterEventId'); + } + if (fromEventId !== undefined || afterEventId !== undefined) { + yield* this.pollByEventId(options ?? {}); + return; + } const signal = options?.signal; // 0 and omitted both mean "from the start" (sequences are 1-indexed). // Entries are `${seq}-1`. XREAD returns IDs greater than `id`, so @@ -164,50 +213,146 @@ export class RedisEventSubscription implements EventSubscripti let cursor = afterSequenceNumber === undefined || afterSequenceNumber === 0 ? '1-0' : `${String(afterSequenceNumber + 1)}-0`; for (;;) { - if (signal?.aborted) { + const messages = await this.readAfter(cursor, signal); + if (messages === undefined) { return; } - if ((await this.redis.exists(this.streamId)) === 0) { - throw new StreamGoneError(this.streamId); - } - const reply = await this.redis.xRead([{ key: this.streamId, id: cursor }], { - COUNT: SUBSCRIBE_STREAM_POLL_ITEMS_COUNT, - }); - const messages = reply?.[0]?.messages ?? []; if (messages.length === 0) { - try { - await sleep(SUBSCRIBE_STREAM_POLL_SLEEP_INTERVAL_MS, undefined, { signal }); - } catch (error) { - if (signal?.aborted) { - return; - } - throw error; + if (!(await this.waitForEntries(signal))) { + return; } continue; } for (const { id, message } of messages) { cursor = id; - const data = message['data']; - if (data === undefined) { - throw new Error(`Corrupt stream entry ${id} on ${this.streamId}: missing data field`); - } + yield this.readEvent(id, message); + } + } + } - let event: T; - try { - // Own writes read back; trusted without validation, as in the gateway. - event = JSON.parse(data) as T; - } catch (error) { - throw new Error(`Corrupt stream entry ${id} on ${this.streamId}: data is not valid JSON`, { - cause: error, - }); + /** + * One Lua seek for the XREAD cursor, then the same forever-poll as sequence + * resume. Seek and batch reads are separate Redis round-trips on purpose. + */ + private async *pollByEventId( + options: EventSubscriptionPollOptions, + ): AsyncGenerator, void, unknown> { + const fromEventId = options.fromEventId; + const afterEventId = options.afterEventId; + const eventId = fromEventId ?? afterEventId; + if (eventId === undefined) { + throw new Error('Pass fromEventId or afterEventId'); + } + const mode = fromEventId !== undefined ? 'from' : 'after'; + const signal = options.signal; + let cursor = await this.seekEventIdCursor({ mode, eventId }); + for (;;) { + const messages = await this.readAfter(cursor, signal); + if (messages === undefined) { + return; + } + if (messages.length === 0) { + if (!(await this.waitForEntries(signal))) { + return; } - yield { ...event, sequence_number: sequenceNumberFromEntryId(this.streamId, id) }; + continue; + } + for (const { id, message } of messages) { + cursor = id; + const event = this.readEvent(id, message); + // after: drop live rows that reuse the seek id (another delta). + if (mode === 'after' && event.id === eventId) { + continue; + } + yield event; + } + } + } + + /** + * Jump near the ULID millisecond (or stream start), walk forward in Lua until + * the payload id, return exclusive XREAD cursor. See SEEK_EVENT_ID_CURSOR. + */ + private async seekEventIdCursor(input: { mode: 'from' | 'after'; eventId: string }): Promise { + const startCursor = eventIdSeekCursor({ + eventId: input.eventId, + firstEntryId: await this.firstEntryId(), + }); + const reply: unknown = await this.redis.eval(SEEK_EVENT_ID_CURSOR, { + keys: [this.streamId], + arguments: [input.mode, input.eventId, startCursor, String(SUBSCRIBE_STREAM_POLL_ITEMS_COUNT)], + }); + if (typeof reply !== 'string') { + throw new Error('Unexpected event-id seek reply'); + } + return reply; + } + + private async readAfter( + cursor: string, + signal: AbortSignal | undefined, + ): Promise<{ id: string; message: Record }[] | undefined> { + if (signal?.aborted) { + return undefined; + } + if ((await this.redis.exists(this.streamId)) === 0) { + throw new StreamGoneError(this.streamId); + } + const reply = await this.redis.xRead([{ key: this.streamId, id: cursor }], { + COUNT: SUBSCRIBE_STREAM_POLL_ITEMS_COUNT, + }); + return reply?.[0]?.messages ?? []; + } + + private async waitForEntries(signal: AbortSignal | undefined): Promise { + try { + await sleep(SUBSCRIBE_STREAM_POLL_SLEEP_INTERVAL_MS, undefined, { signal }); + } catch (error) { + if (signal?.aborted) { + return false; } + throw error; } + return true; } - private async readTail(): Promise<{ eventId: string; sequenceNumber: number } | undefined> { + private readEvent(id: string, message: Record): SequencedEvent { + const data = message['data']; + if (data === undefined) { + throw new Error(`Corrupt stream entry ${id} on ${this.streamId}: missing data field`); + } + let parsed: unknown; + try { + parsed = JSON.parse(data); + } catch (error) { + throw new Error(`Corrupt stream entry ${id} on ${this.streamId}: data is not valid JSON`, { cause: error }); + } + if (!this.isStoredEvent(parsed)) { + throw new Error(`Corrupt stream entry ${id} on ${this.streamId}: data is not an event`); + } + if (!isLegacySequenceEntryId(id)) { + if (Object.hasOwn(parsed, 'sequence_number')) { + const copy: T = { ...parsed }; + Reflect.deleteProperty(copy, 'sequence_number'); + return copy; + } + return parsed; + } + return { ...parsed, sequence_number: entryIdMilliseconds(id) }; + } + + private isStoredEvent(value: unknown): value is T { + return ( + typeof value === 'object' && + value !== null && + !Array.isArray(value) && + 'id' in value && + typeof value.id === 'string' + ); + } + + private async readTail(): Promise<{ eventId: string; payload: string } | undefined> { const entries = await this.redis.xRevRange(this.streamId, '+', '-', { COUNT: 1 }); const latest = entries[0]; if (!latest) { @@ -217,10 +362,16 @@ export class RedisEventSubscription implements EventSubscripti if (data === undefined) { throw new Error(`Corrupt stream entry ${latest.id} on ${this.streamId}: missing data field`); } - return { - eventId: eventIdFromStreamData(data) ?? '', - sequenceNumber: sequenceNumberFromEntryId(this.streamId, latest.id), - }; + const eventId = eventIdFromStreamData(data); + if (eventId === undefined) { + throw new Error(`Corrupt stream entry ${latest.id} on ${this.streamId}: missing event id`); + } + return { eventId, payload: data }; + } + + private async lastEntryId(): Promise { + const entries = await this.redis.xRevRange(this.streamId, '+', '-', { COUNT: 1 }); + return entries[0]?.id; } private async refreshTtl(streamTTLSeconds: number | undefined): Promise { @@ -229,8 +380,69 @@ export class RedisEventSubscription implements EventSubscripti } } + private async appendEvent(input: { + event: T; + checkOwner: boolean; + executorId: string; + streamTTLSeconds: number | undefined; + }): Promise<'appended' | 'duplicate'> { + const payload = JSON.stringify(input.event); + const reply = await this.appendEntry({ + streamId: this.streamIdFor(input.event), + payload, + checkOwner: input.checkOwner, + executorId: input.executorId, + streamTTLSeconds: input.streamTTLSeconds, + }); + if (reply.status === 'ok') { + return 'appended'; + } + if (reply.status === 'duplicate') { + return 'duplicate'; + } + if (reply.status === 'owner_mismatch') { + throw new StreamOwnerMismatchError(this.streamId); + } + if (reply.status === 'error') { + throw new Error(`Redis stream append failed on ${this.streamId}: ${reply.detail}`); + } + // The tail id can match a delta of the same message. Identical payloads + // already returned `duplicate` above; a distinct payload must be retried + // at the tail millisecond. + const lastId = await this.lastEntryId(); + if (lastId === undefined) { + throw new Error(`Redis stream append conflicted on an empty stream ${this.streamId}`); + } + let retry: { status: AppendStatus; detail: string }; + try { + retry = await this.appendEntry({ + streamId: `${String(entryIdMilliseconds(lastId))}-*`, + payload, + checkOwner: input.checkOwner, + executorId: input.executorId, + streamTTLSeconds: input.streamTTLSeconds, + }); + } catch (error) { + throw new Error(`Failed to append event on ${this.streamId}`, { cause: error }); + } + if (retry.status === 'ok') { + return 'appended'; + } + if (retry.status === 'duplicate') { + return 'duplicate'; + } + if (retry.status === 'owner_mismatch') { + throw new StreamOwnerMismatchError(this.streamId); + } + throw new Error(`Redis stream append failed on ${this.streamId}: ${retry.detail}`, { cause: reply.detail }); + } + + private streamIdFor(event: T): string { + return `${String(eventIdTimestampMs(event.id) ?? Date.now())}-*`; + } + private async appendEntry(input: { - sequenceNumber: number; + streamId: string; payload: string; checkOwner: boolean; executorId: string; @@ -239,39 +451,8 @@ export class RedisEventSubscription implements EventSubscripti const ttl = input.streamTTLSeconds !== undefined && input.streamTTLSeconds > 0 ? input.streamTTLSeconds : 0; const reply: unknown = await this.redis.eval(APPEND_STREAM_ENTRY, { keys: [streamOwnerKey(this.streamId), this.streamId], - arguments: [ - input.checkOwner ? '1' : '0', - input.executorId, - `${String(input.sequenceNumber)}-1`, - input.payload, - String(ttl), - ], + arguments: [input.checkOwner ? '1' : '0', input.executorId, input.streamId, input.payload, String(ttl)], }); return parseAppendReply(reply); } - - private outcomeForEvent(input: { - reply: { status: AppendStatus; detail: string }; - eventId: string | undefined; - sequenceNumber: number; - }): Promise<'appended' | 'duplicate'> { - if (input.reply.status === 'ok') { - return Promise.resolve('appended'); - } - if (input.reply.status === 'owner_mismatch') { - return Promise.reject(new StreamOwnerMismatchError(this.streamId)); - } - if (input.reply.status === 'error') { - return Promise.reject(new Error(`Redis stream append failed on ${this.streamId}: ${input.reply.detail}`)); - } - const tailId = eventIdFromStreamData(input.reply.detail); - if (tailId !== undefined && input.eventId !== undefined && tailId === input.eventId) { - return Promise.resolve('duplicate'); - } - return Promise.reject( - new Error( - `explicit stream id ${String(input.sequenceNumber)} conflicts on ${this.streamId} and the tail event is ${tailId ?? 'missing'}`, - ), - ); - } } diff --git a/packages/trueforge/src/runtime/event-subscription/repair.ts b/packages/trueforge/src/runtime/event-subscription/repair.ts index 199a267e47..f194b0e7ef 100644 --- a/packages/trueforge/src/runtime/event-subscription/repair.ts +++ b/packages/trueforge/src/runtime/event-subscription/repair.ts @@ -1,23 +1,16 @@ /** * Rebuild a turn's live event stream from the durable log. - * Sequence numbers already stored are written back as explicit stream ids. - * An event with no number yet takes max(stored) + 1 in persist order. + * Append uses the same event-id stream ids as a live put. The durable row + * order is the order. No sequence number is assigned. */ -export interface EventLogEntry { - /** Payload id. Compared with the stream tail. Redis stream ids are sequence numbers, not this. */ +export interface EventLogEntry { + /** Payload id. The Redis stream id is not this value. */ id: string; event: T; - sequence_number?: number | undefined; } -interface SequencedLogEntry { - id: string; - event: T; - sequence_number: number; -} - -export interface RepairLogOptions { +export interface RepairLogOptions { entries: readonly EventLogEntry[]; streamTTLSeconds?: number | undefined; /** @@ -25,26 +18,50 @@ export interface RepairLogOptions { * key. A subscribe repair omits this and does not claim. */ executorId?: string | undefined; - /** - * Persist numbers assigned to events that did not have one. Called even when - * the list is empty so the owner fence runs before the stream is touched. - */ - persistSequences?: (assignments: { event_id: string; sequence_number: number }[]) => Promise; } interface StreamTail { eventId: string; - sequenceNumber: number; + /** JSON written on the stream, used to tell shared-id rows apart. */ + payload: string; } -export function eventIdOf(event: object): string | undefined { - if (!('id' in event)) { - return undefined; +function isPlainObject(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +/** Key order and a legacy `sequence_number` field are not part of event identity. */ +function canonicalize(value: unknown): unknown { + if (Array.isArray(value)) { + return value.map(item => canonicalize(item)); + } + if (!isPlainObject(value)) { + return value; + } + const normalized: Record = {}; + for (const key of Object.keys(value).sort()) { + if (key === 'sequence_number') { + continue; + } + const child = value[key]; + if (child !== undefined) { + normalized[key] = canonicalize(child); + } + } + return normalized; +} + +export function eventPayloadMatches(event: object, payload: string): boolean { + let parsed: unknown; + try { + parsed = JSON.parse(payload); + } catch { + return false; } - const id = event.id; - return typeof id === 'string' ? id : undefined; + return JSON.stringify(canonicalize(event)) === JSON.stringify(canonicalize(parsed)); } +/** Payload id from a stream data field. Undefined when JSON is corrupt or id is missing. */ export function eventIdFromStreamData(data: string): string | undefined { let parsed: unknown; try { @@ -52,93 +69,51 @@ export function eventIdFromStreamData(data: string): string | undefined { } catch { return undefined; } - if (typeof parsed !== 'object' || parsed === null) { + if (typeof parsed !== 'object' || parsed === null || !('id' in parsed)) { return undefined; } - return eventIdOf(parsed); -} - -function planSequences( - entries: readonly EventLogEntry[], -): { - sequenced: SequencedLogEntry[]; - newlyAssigned: { event_id: string; sequence_number: number }[]; -} { - let maxStored = 0; - for (const entry of entries) { - const sequence = entry.sequence_number; - if (sequence !== undefined && sequence > maxStored) { - maxStored = sequence; - } - } - const sequenced: SequencedLogEntry[] = []; - const newlyAssigned: { event_id: string; sequence_number: number }[] = []; - let next = maxStored; - for (const entry of entries) { - if (entry.sequence_number === undefined) { - next += 1; - sequenced.push({ id: entry.id, event: entry.event, sequence_number: next }); - newlyAssigned.push({ event_id: entry.id, sequence_number: next }); - } else { - sequenced.push({ id: entry.id, event: entry.event, sequence_number: entry.sequence_number }); - } - } - return { sequenced, newlyAssigned }; + return typeof parsed.id === 'string' ? parsed.id : undefined; } /** - * Suffix the stream does not already have. Redis explicit ids must increase, so - * the suffix is appended in sequence order. That matches persist order when - * numbers were assigned as events were streamed. + * Rows after the stream tail. Rows that share an id (a message and its deltas) + * are distinguished by the stored payload. A tail the log does not contain is + * left alone. */ -function suffixAfterTail( - sequenced: readonly SequencedLogEntry[], +function suffixAfterTail( + entries: readonly EventLogEntry[], tail: StreamTail | undefined, -): SequencedLogEntry[] { +): EventLogEntry[] { if (tail === undefined) { - return [...sequenced].sort((left, right) => left.sequence_number - right.sequence_number); - } - const last = sequenced[sequenced.length - 1]; - if (last?.id === tail.eventId) { - return []; + return [...entries]; } - const tailIndex = sequenced.findIndex(entry => entry.id === tail.eventId); - const after = - tailIndex >= 0 - ? sequenced.slice(tailIndex + 1) - : sequenced.filter(entry => entry.sequence_number > tail.sequenceNumber); - return [...after].sort((left, right) => left.sequence_number - right.sequence_number); -} - -function nextSequence(sequenced: readonly { sequence_number: number }[], tail: StreamTail | undefined): number { - let max = tail?.sequenceNumber ?? 0; - for (const entry of sequenced) { - if (entry.sequence_number > max) { - max = entry.sequence_number; + let match = -1; + for (let index = 0; index < entries.length; index += 1) { + const entry = entries[index]; + if (entry?.id === tail.eventId && eventPayloadMatches(entry.event, tail.payload)) { + match = index; } } - return max + 1; + if (match === -1) { + return []; + } + return entries.slice(match + 1); } -export async function applyEventLogRepair(input: { +export async function applyEventLogRepair(input: { entries: readonly EventLogEntry[]; - persistSequences?: RepairLogOptions['persistSequences']; claimOwner?: (() => Promise) | undefined; readTail: () => Promise; - append: (entry: SequencedLogEntry) => Promise<'appended' | 'duplicate'>; + append: (entry: EventLogEntry) => Promise<'appended' | 'duplicate'>; refreshTtl: () => Promise; /** Producer resume refreshes TTL even when the existing stream is already complete. */ refreshTtlOnUntouchedStream: boolean; -}): Promise { - const plan = planSequences(input.entries); - if (input.persistSequences !== undefined) { - await input.persistSequences(plan.newlyAssigned); - } +}): Promise { if (input.claimOwner !== undefined) { await input.claimOwner(); } const tail = await input.readTail(); - const suffix = suffixAfterTail(plan.sequenced, tail); + const suffix = suffixAfterTail(input.entries, tail); let appended = false; for (const entry of suffix) { const outcome = await input.append(entry); @@ -149,5 +124,4 @@ export async function applyEventLogRepair(input: { if (appended || input.refreshTtlOnUntouchedStream) { await input.refreshTtl(); } - return nextSequence(plan.sequenced, tail); } diff --git a/packages/trueforge/src/runtime/turnEventLogRepair.ts b/packages/trueforge/src/runtime/turnEventLogRepair.ts index 5190aabaff..baf0294b72 100644 --- a/packages/trueforge/src/runtime/turnEventLogRepair.ts +++ b/packages/trueforge/src/runtime/turnEventLogRepair.ts @@ -4,7 +4,8 @@ import type { EventLogEntry, EventSubscription } from './event-subscription'; /** * Rebuild a paused turn's live stream from the durable log. * Subscribe omits `producerExecutorId` and only appends. - * The producing resume records any missing sequence numbers and claims the stream-owner key. + * The producing resume claims the stream-owner key. Rows are appended with the + * same event-id stream ids as a live put. */ export async function repairTurnEventLog(input: { sessionStore: Pick, 'listTurnEventLog' | 'assignTurnEventSequenceNumbers'>; @@ -21,23 +22,20 @@ export async function repairTurnEventLog(input: { const entries: EventLogEntry[] = rows.map(row => ({ id: row.event.id, event: row.event, - ...(row.sequence_number === null ? {} : { sequence_number: row.sequence_number }), })); const producerExecutorId = input.producerExecutorId; + if (producerExecutorId !== undefined) { + // Empty assignment still takes the turn-owner fence. It does not write sequence numbers. + await input.sessionStore.assignTurnEventSequenceNumbers({ + session_id: input.sessionId, + turn_id: input.turnId, + active_executor_id: producerExecutorId, + assignments: [], + }); + } await input.subscription.repairLog({ entries, streamTTLSeconds: input.streamTTLSeconds, - ...(producerExecutorId === undefined - ? {} - : { - executorId: producerExecutorId, - persistSequences: assignments => - input.sessionStore.assignTurnEventSequenceNumbers({ - session_id: input.sessionId, - turn_id: input.turnId, - active_executor_id: producerExecutorId, - assignments, - }), - }), + ...(producerExecutorId === undefined ? {} : { executorId: producerExecutorId }), }); } diff --git a/packages/trueforge/src/schemas/turn.ts b/packages/trueforge/src/schemas/turn.ts index e771902f12..cf22f2a73e 100644 --- a/packages/trueforge/src/schemas/turn.ts +++ b/packages/trueforge/src/schemas/turn.ts @@ -17,10 +17,29 @@ export const SubscribeTurnQuerySchema = z .int() .nonnegative() .optional() - .describe( - 'Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer.', - ), + .describe('Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.'), ), + from_event_id: z + .string() + .min(1) + .optional() + .describe('Inclusive cursor. Replay starts at the entry with this event id'), + after_event_id: z + .string() + .min(1) + .optional() + .describe('Exclusive cursor. Replay starts after the last entry with this event id.'), + }) + .superRefine((query, ctx) => { + const cursors = [query.after_sequence_number, query.from_event_id, query.after_event_id].filter( + value => value !== undefined, + ); + if (cursors.length > 1) { + ctx.addIssue({ + code: 'custom', + message: 'Pass only one of after_sequence_number, from_event_id, or after_event_id.', + }); + } }) .openapi('SubscribeTurnQuery'); diff --git a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts index f8caeda7e2..24d6bc1af5 100644 --- a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts +++ b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts @@ -1,13 +1,14 @@ import { StreamGoneError, StreamOwnerMismatchError, + type EventSubscriptionPollOptions, type EventSubscriptionRegistry, type SequencedEvent, } from '../../../src/runtime/event-subscription'; /** Payload shape used by the suite; distinct per event so round-trips are provable. */ export interface TestEvent { - id?: string; + id: string; type: string; payload: { index: number; text: string }; } @@ -21,8 +22,8 @@ function nextStreamId(): string { return `eventsub-contract:${String(process.pid)}:${String(Date.now())}:${String(streamCounter)}`; } -function makeEvent(index: number): TestEvent { - return { type: 'test.event', payload: { index, text: `event-${String(index)}` } }; +function makeEvent(index: number, id = `evt-${String(index)}`): TestEvent { + return { id, type: 'test.event', payload: { index, text: `event-${String(index)}` } }; } /** Drains `count` events from a fresh poll, then aborts it (poll loops forever by design). */ @@ -30,10 +31,10 @@ async function collect( registry: EventSubscriptionRegistry, streamId: string, count: number, - afterSequenceNumber?: number, + cursor?: Omit, ): Promise[]> { const abort = new AbortController(); - const generator = registry.get(streamId).poll(afterSequenceNumber, { signal: abort.signal }); + const generator = registry.get(streamId).poll({ signal: abort.signal, ...cursor }); const received: SequencedEvent[] = []; try { for await (const event of generator) { @@ -51,16 +52,18 @@ async function collect( /** EventSubscription contract suite — factory-injected so both backends run it. */ export function runEventSubscriptionContractSuite(createRegistry: () => EventSubscriptionRegistry) { - it('put assigns dense 1-indexed sequence numbers', async () => { + it('put appends events in order without a sequence number', async () => { const registry = createRegistry(); const streamId = nextStreamId(); const producer = registry.get(streamId); - const sequenceNumbers = []; - for (let i = 0; i < 3; i += 1) { - sequenceNumbers.push(await producer.put(makeEvent(i), { streamTTLSeconds: STREAM_TTL_SECONDS })); + const sent = [makeEvent(0), makeEvent(1), makeEvent(2)]; + for (const event of sent) { + await producer.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }); } - expect(sequenceNumbers).toEqual([1, 2, 3]); + const received = await collect(registry, streamId, sent.length); + expect(received).toEqual(sent); + expect(received.every(event => !Object.hasOwn(event, 'sequence_number'))).toBe(true); }); it('poll from the start yields every event sent, in order, with identical payloads', async () => { @@ -74,7 +77,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub } const received = await collect(registry, streamId, sent.length); - expect(received).toEqual(sent.map((event, index) => ({ ...event, sequence_number: index + 1 }))); + expect(received).toEqual(sent); }); it('poll with afterSequenceNumber 0 matches an omitted cursor (from the start)', async () => { @@ -88,9 +91,8 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub } const fromOmitted = await collect(registry, streamId, sent.length); - const fromZero = await collect(registry, streamId, sent.length, 0); + const fromZero = await collect(registry, streamId, sent.length, { afterSequenceNumber: 0 }); expect(fromZero).toEqual(fromOmitted); - expect(fromZero[0]?.sequence_number).toBe(1); }); it('poll from a mid-stream cursor yields only events strictly after it', async () => { @@ -98,18 +100,80 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub const streamId = nextStreamId(); const producer = registry.get(streamId); - const sent = [makeEvent(0), makeEvent(1), makeEvent(2), makeEvent(3), makeEvent(4)]; + const sent = [0, 1, 2, 3, 4].map(index => makeEvent(index)); for (const event of sent) { await producer.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }); } - // after 1 → skip seq 1; replay seq 2, 3, 4 - const received = await collect(registry, streamId, 3, 1); - expect(received).toEqual([ - { ...sent[1], sequence_number: 2 }, - { ...sent[2], sequence_number: 3 }, - { ...sent[3], sequence_number: 4 }, - ]); + const received = await collect(registry, streamId, 3, { afterEventId: 'evt-1' }); + expect(received.map(event => event.id)).toEqual(['evt-2', 'evt-3', 'evt-4']); + }); + + it('keeps an event appended while an event-id replay is delivered', async () => { + const registry = createRegistry(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + const sent = [0, 1, 2].map(index => makeEvent(index)); + for (const event of sent) { + await producer.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }); + } + + const abort = new AbortController(); + const generator = registry.get(streamId).poll({ + signal: abort.signal, + afterEventId: 'evt-1', + }); + try { + const replayed = await generator.next(); + expect(replayed.value?.id).toBe('evt-2'); + const extra = makeEvent(9, 'evt-extra'); + await producer.put(extra, { streamTTLSeconds: STREAM_TTL_SECONDS }); + const appended = await generator.next(); + expect(appended.value).toEqual(extra); + } finally { + abort.abort(); + await generator.return(undefined); + } + }); + + it('appends a distinct shared-id payload when its timestamp is behind the tail', async () => { + const { ulid } = await import('ulid'); + const registry = createRegistry(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + const later = makeEvent(0, ulid(2_000_000_000_000).toLowerCase()); + const shell = makeEvent(1, ulid(1_000_000_000_000).toLowerCase()); + const delta = { ...shell, payload: { index: 2, text: 'delta' } }; + await producer.put(later, { streamTTLSeconds: STREAM_TTL_SECONDS }); + await producer.put(shell, { streamTTLSeconds: STREAM_TTL_SECONDS }); + await producer.put(delta, { streamTTLSeconds: STREAM_TTL_SECONDS }); + + const received = await collect(registry, streamId, 3); + expect(received).toEqual([later, shell, delta]); + }); + + it('repairs a suffix when the durable row key order differs from the stream payload', async () => { + const registry = createRegistry(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + const stored = makeEvent(0); + await producer.put(stored, { streamTTLSeconds: STREAM_TTL_SECONDS }); + const fromDatabase = { + id: 'evt-0', + type: stored.type, + payload: { text: stored.payload.text, index: stored.payload.index }, + }; + const third = makeEvent(2); + await producer.repairLog({ + entries: [ + { id: 'evt-0', event: fromDatabase }, + { id: 'evt-2', event: third }, + ], + streamTTLSeconds: STREAM_TTL_SECONDS, + }); + + const received = await collect(registry, streamId, 2); + expect(received.map(event => event.id)).toEqual(['evt-0', 'evt-2']); }); it('poll parked at the tip receives events put after it subscribed', async () => { @@ -117,99 +181,85 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub const streamId = nextStreamId(); const producer = registry.get(streamId); - await producer.put(makeEvent(0), { streamTTLSeconds: STREAM_TTL_SECONDS }); + const first = makeEvent(0); + await producer.put(first, { streamTTLSeconds: STREAM_TTL_SECONDS }); - // after 1 (= tip): nothing to replay, so the poller parks for live events. - const pending = collect(registry, streamId, 2, 1); - await producer.put(makeEvent(1)); - await producer.put(makeEvent(2)); + const pending = collect(registry, streamId, 2, { afterEventId: 'evt-0' }); + const second = makeEvent(1); + const third = makeEvent(2); + await producer.put(second); + await producer.put(third); const received = await pending; - expect(received).toEqual([ - { ...makeEvent(1), sequence_number: 2 }, - { ...makeEvent(2), sequence_number: 3 }, - ]); + expect(received).toEqual([second, third]); }); - it('rebuilds a missing stream from stored sequences and continues after the tail', async () => { + it('rebuilds a missing stream from the event log and continues after the tail', async () => { const registry = createRegistry(); const streamId = nextStreamId(); const producer = registry.get(streamId); - const first = { id: 'evt-1', ...makeEvent(0) }; - const second = { id: 'evt-2', ...makeEvent(1) }; + const first = makeEvent(1); + const second = makeEvent(2); await producer.put(first, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); await registry.evictStream(streamId); await producer.repairLog({ entries: [ - { id: 'evt-1', event: first, sequence_number: 1 }, - { id: 'evt-2', event: second, sequence_number: 2 }, + { id: first.id, event: first }, + { id: second.id, event: second }, ], streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner', - persistSequences: () => Promise.resolve(), }); - const third = { id: 'evt-3', ...makeEvent(2) }; - const next = await producer.put(third, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); - expect(next).toBe(3); + const third = makeEvent(3); + await producer.put(third, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); const received = await collect(registry, streamId, 3); - expect(received).toEqual([ - { ...first, sequence_number: 1 }, - { ...second, sequence_number: 2 }, - { ...third, sequence_number: 3 }, - ]); + expect(received).toEqual([first, second, third]); }); it('leaves an existing stream untouched and a second repair does not append again', async () => { const registry = createRegistry(); const streamId = nextStreamId(); const producer = registry.get(streamId); - const first = { id: 'evt-1', ...makeEvent(0) }; - const second = { id: 'evt-2', ...makeEvent(1) }; + const first = makeEvent(1); + const second = makeEvent(2); await producer.put(first, { streamTTLSeconds: STREAM_TTL_SECONDS }); await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS }); const repair = () => producer.repairLog({ entries: [ - { id: 'evt-1', event: first, sequence_number: 1 }, - { id: 'evt-2', event: second, sequence_number: 2 }, + { id: first.id, event: first }, + { id: second.id, event: second }, ], streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner', - persistSequences: () => Promise.resolve(), }); await repair(); await repair(); - const third = { id: 'evt-3', ...makeEvent(2) }; - const next = await producer.put(third, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); - expect(next).toBe(3); + const third = makeEvent(3); + await producer.put(third, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); const received = await collect(registry, streamId, 3); - expect(received).toEqual([ - { ...first, sequence_number: 1 }, - { ...second, sequence_number: 2 }, - { ...third, sequence_number: 3 }, - ]); + expect(received).toEqual([first, second, third]); }); it('rejects a stale producer when the stream is missing', async () => { const registry = createRegistry(); const streamId = nextStreamId(); const owner = registry.get(streamId); - const event = { id: 'evt-1', ...makeEvent(0) }; + const event = makeEvent(1); await owner.repairLog({ - entries: [{ id: 'evt-1', event, sequence_number: 1 }], + entries: [{ id: event.id, event }], streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner', - persistSequences: () => Promise.resolve(), }); await registry.evictStream(streamId); const stale = registry.get(streamId); await expect( - stale.put({ id: 'evt-late', ...makeEvent(9) }, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'stale' }), + stale.put(makeEvent(9, 'evt-late'), { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'stale' }), ).rejects.toBeInstanceOf(StreamOwnerMismatchError); await expect(stale.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); }); @@ -219,20 +269,15 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub const streamId = nextStreamId(); const left = registry.get(streamId); const right = registry.get(streamId); - const event = { id: 'evt-same', ...makeEvent(0) }; - const sequences = await Promise.all([ + const event = makeEvent(0, 'evt-same'); + await Promise.all([ left.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }), right.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }), ]); - expect(sequences).toEqual([1, 1]); - const follow = { id: 'evt-next', ...makeEvent(1) }; - const next = await left.put(follow, { streamTTLSeconds: STREAM_TTL_SECONDS }); - expect(next).toBe(2); + const follow = makeEvent(1, 'evt-next'); + await left.put(follow, { streamTTLSeconds: STREAM_TTL_SECONDS }); const received = await collect(registry, streamId, 2); - expect(received).toEqual([ - { ...event, sequence_number: 1 }, - { ...follow, sequence_number: 2 }, - ]); + expect(received).toEqual([event, follow]); }); } diff --git a/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts b/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts new file mode 100644 index 0000000000..447bd288dc --- /dev/null +++ b/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts @@ -0,0 +1,41 @@ +import { eventIdSeekCursor, eventIdTimestampMs } from '../../../src/runtime/event-subscription/entryId'; +import { eventPayloadMatches } from '../../../src/runtime/event-subscription/repair'; + +const EVENT_ID = '01arz3ndektsv4rrffq69g5fav'; + +describe('event id cursors', () => { + it('accepts the lowercase ids the server mints', () => { + expect(eventIdTimestampMs(EVENT_ID)).toBe(eventIdTimestampMs(EVENT_ID.toUpperCase())); + }); + + it('scans a legacy sequence stream from the start', () => { + expect(eventIdSeekCursor({ eventId: EVENT_ID, firstEntryId: '12-1' })).toBe('0-0'); + }); + + it('seeks a ULID-millisecond stream at that millisecond', () => { + const timestamp = eventIdTimestampMs(EVENT_ID); + expect(timestamp).toBeDefined(); + if (timestamp === undefined) { + return; + } + expect(eventIdSeekCursor({ eventId: EVENT_ID, firstEntryId: `${String(timestamp)}-0` })).toBe( + `${String(timestamp - 1)}-18446744073709551615`, + ); + }); + + it('matches a durable row when key order or a legacy sequence field differs', () => { + const event = { id: 'evt-1', type: 'model.message', payload: { index: 1, text: 'hi' } }; + expect(eventPayloadMatches(event, '{"payload":{"text":"hi","index":1},"type":"model.message","id":"evt-1"}')).toBe( + true, + ); + expect( + eventPayloadMatches( + event, + '{"id":"evt-1","type":"model.message","payload":{"index":1,"text":"hi"},"sequence_number":4}', + ), + ).toBe(true); + expect(eventPayloadMatches(event, '{"id":"evt-1","type":"model.message","payload":{"index":2,"text":"hi"}}')).toBe( + false, + ); + }); +}); diff --git a/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts index 99fc92caf3..47159a4983 100644 --- a/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts @@ -24,7 +24,7 @@ async function collect( count: number, ): Promise[]> { const abort = new AbortController(); - const generator = registry.get(streamId).poll(undefined, { signal: abort.signal }); + const generator = registry.get(streamId).poll({ signal: abort.signal }); const received: SequencedEvent[] = []; try { for await (const event of generator) { @@ -41,7 +41,7 @@ async function collect( } describe('repairTurnEventLog', () => { - it('rebuilds a missing stream from the durable log and keeps the sequence', async () => { + it('rebuilds a missing stream from the durable log without assigning sequence numbers', async () => { const store = new InMemorySessionStore(); const sessionId = 's1'; const turnId = 'turn-1'; @@ -85,16 +85,15 @@ describe('repairTurnEventLog', () => { expect( (await store.listTurnEventLog({ session_id: sessionId, turn_id: turnId })).map(row => row.sequence_number), ).toEqual([null, null]); - expect((await collect(registry, streamId, 2)).map(event => event.sequence_number)).toEqual([1, 2]); + expect((await collect(registry, streamId, 2)).map(event => event.id)).toEqual([created.id, message.id]); await repair(); - expect((await collect(registry, streamId, 2)).map(event => event.sequence_number)).toEqual([1, 2]); + expect((await collect(registry, streamId, 2)).map(event => event.id)).toEqual([created.id, message.id]); await repair(TEST_ACTIVE_EXECUTOR_ID); expect( (await store.listTurnEventLog({ session_id: sessionId, turn_id: turnId })).map(row => row.sequence_number), - ).toEqual([1, 2]); - expect((await collect(registry, streamId, 2)).map(event => event.sequence_number)).toEqual([1, 2]); + ).toEqual([null, null]); const later = makeModelMessageEvent(); await store.appendToEvents({ @@ -104,10 +103,10 @@ describe('repairTurnEventLog', () => { events: [later], }); await repair(TEST_ACTIVE_EXECUTOR_ID); + expect((await collect(registry, streamId, 3)).map(event => event.id)).toEqual([created.id, message.id, later.id]); expect( (await store.listTurnEventLog({ session_id: sessionId, turn_id: turnId })).map(row => row.sequence_number), - ).toEqual([1, 2, 3]); - expect((await collect(registry, streamId, 3)).map(event => event.sequence_number)).toEqual([1, 2, 3]); + ).toEqual([null, null, null]); await registry.evictStream(streamId); const stale = registry.get(streamId); @@ -117,11 +116,16 @@ describe('repairTurnEventLog', () => { await expect(stale.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); await repair(TEST_ACTIVE_EXECUTOR_ID); - const next = await subscription.put(makeModelMessageEvent(), { + const appended = makeModelMessageEvent(); + await subscription.put(appended, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: TEST_ACTIVE_EXECUTOR_ID, }); - expect(next).toBe(4); - expect((await collect(registry, streamId, 4)).map(event => event.sequence_number)).toEqual([1, 2, 3, 4]); + expect((await collect(registry, streamId, 4)).map(event => event.id)).toEqual([ + created.id, + message.id, + later.id, + appended.id, + ]); }); }); diff --git a/packages/trueforge/tests/unit/apis/turns.test.ts b/packages/trueforge/tests/unit/apis/turns.test.ts index a5474d8641..7923762e9b 100644 --- a/packages/trueforge/tests/unit/apis/turns.test.ts +++ b/packages/trueforge/tests/unit/apis/turns.test.ts @@ -1007,7 +1007,24 @@ describe('turns', () => { await reader.cancel(); } expect(replayed).toContain(pausedEvent.id); + expect(replayed.split('\n').some(line => line.startsWith('id:'))).toBe(false); expect(activeTurns.getTurnHandle({ sessionId, turnId: pausedTurnId })).toBeUndefined(); + + const mixed = await app.request( + `/${sessionId}/turns/${pausedTurnId}/subscribe?after_sequence_number=1&from_event_id=${pausedEvent.id}`, + ); + expect(mixed.status).toBe(400); + + const huge = await app.request(`/${sessionId}/turns/missing-turn/subscribe?after_sequence_number=100001`); + expect(huge.status).toBe(400); + + const invalid = await app.request(`/${sessionId}/turns/${pausedTurnId}/subscribe?from_event_id=not-a-ulid`); + expect(invalid.status).toBe(400); + + const sequenceOnEventIdStream = await app.request( + `/${sessionId}/turns/${pausedTurnId}/subscribe?after_sequence_number=1`, + ); + expect(sequenceOnEventIdStream.status).toBe(400); }); }); }); diff --git a/python/trueforge_sdk/reference.md b/python/trueforge_sdk/reference.md index b6033df628..256a69e2d0 100644 --- a/python/trueforge_sdk/reference.md +++ b/python/trueforge_sdk/reference.md @@ -3462,7 +3462,7 @@ client.sessions.create_turn_events(
-Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. +Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id
@@ -3519,7 +3519,23 @@ client.sessions.subscribe_to_turn(
-**after_sequence_number:** `typing.Optional[int]` — Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer. +**after_sequence_number:** `typing.Optional[int]` — Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. + +
+
+ +
+
+ +**from_event_id:** `typing.Optional[str]` — Inclusive cursor. Replay starts at the entry with this event id + +
+
+ +
+
+ +**after_event_id:** `typing.Optional[str]` — Exclusive cursor. Replay starts after the last entry with this event id.
diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py index a26183cc54..8198004487 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py @@ -834,10 +834,12 @@ def subscribe_to_turn( session_id: str, turn_id: str, after_sequence_number: typing.Optional[int] = None, + from_event_id: typing.Optional[str] = None, + after_event_id: typing.Optional[str] = None, request_options: typing.Optional[RequestOptions] = None, ) -> Stream[TurnStreamingEvent]: """ - Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. + Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id Parameters ---------- @@ -848,7 +850,13 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer. + Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. + + from_event_id : typing.Optional[str] + Inclusive cursor. Replay starts at the entry with this event id + + after_event_id : typing.Optional[str] + Exclusive cursor. Replay starts after the last entry with this event id. request_options : typing.Optional[RequestOptions] Request-specific configuration. @@ -879,6 +887,8 @@ def _events() -> typing.Generator[StreamEvent[TurnStreamingEvent], None, None]: session_id=session_id, turn_id=turn_id, after_sequence_number=after_sequence_number, + from_event_id=from_event_id, + after_event_id=after_event_id, request_options=request_options, ) as r: yield from r.data.with_metadata() @@ -1821,10 +1831,12 @@ def subscribe_to_turn( session_id: str, turn_id: str, after_sequence_number: typing.Optional[int] = None, + from_event_id: typing.Optional[str] = None, + after_event_id: typing.Optional[str] = None, request_options: typing.Optional[RequestOptions] = None, ) -> AsyncStream[TurnStreamingEvent]: """ - Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. + Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id Parameters ---------- @@ -1835,7 +1847,13 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer. + Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. + + from_event_id : typing.Optional[str] + Inclusive cursor. Replay starts at the entry with this event id + + after_event_id : typing.Optional[str] + Exclusive cursor. Replay starts after the last entry with this event id. request_options : typing.Optional[RequestOptions] Request-specific configuration. @@ -1874,6 +1892,8 @@ async def _events() -> typing.AsyncGenerator[StreamEvent[TurnStreamingEvent], No session_id=session_id, turn_id=turn_id, after_sequence_number=after_sequence_number, + from_event_id=from_event_id, + after_event_id=after_event_id, request_options=request_options, ) as r: async for _event in r.data.with_metadata(): diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py index d1d46972f3..dbe1c8ba9b 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py @@ -1747,10 +1747,12 @@ def subscribe_to_turn( session_id: str, turn_id: str, after_sequence_number: typing.Optional[int] = None, + from_event_id: typing.Optional[str] = None, + after_event_id: typing.Optional[str] = None, request_options: typing.Optional[RequestOptions] = None, ) -> typing.Iterator[HttpResponse[Stream[TurnStreamingEvent]]]: """ - Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. + Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id Parameters ---------- @@ -1761,7 +1763,13 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer. + Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. + + from_event_id : typing.Optional[str] + Inclusive cursor. Replay starts at the entry with this event id + + after_event_id : typing.Optional[str] + Exclusive cursor. Replay starts after the last entry with this event id. request_options : typing.Optional[RequestOptions] Request-specific configuration. @@ -1776,6 +1784,8 @@ def subscribe_to_turn( method="GET", params={ "after_sequence_number": after_sequence_number, + "from_event_id": from_event_id, + "after_event_id": after_event_id, }, request_options=request_options, ) as _response: @@ -3580,10 +3590,12 @@ async def subscribe_to_turn( session_id: str, turn_id: str, after_sequence_number: typing.Optional[int] = None, + from_event_id: typing.Optional[str] = None, + after_event_id: typing.Optional[str] = None, request_options: typing.Optional[RequestOptions] = None, ) -> typing.AsyncIterator[AsyncHttpResponse[AsyncStream[TurnStreamingEvent]]]: """ - Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. A running or finished turn is not restored. Pass `after_sequence_number` to resume after that sequence. + Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id Parameters ---------- @@ -3594,7 +3606,13 @@ async def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Exclusive resume cursor: replay only events with a sequence number greater than this value. Omit to start from the beginning of the live buffer. + Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. + + from_event_id : typing.Optional[str] + Inclusive cursor. Replay starts at the entry with this event id + + after_event_id : typing.Optional[str] + Exclusive cursor. Replay starts after the last entry with this event id. request_options : typing.Optional[RequestOptions] Request-specific configuration. @@ -3609,6 +3627,8 @@ async def subscribe_to_turn( method="GET", params={ "after_sequence_number": after_sequence_number, + "from_event_id": from_event_id, + "after_event_id": after_event_id, }, request_options=request_options, ) as _response: From a3a6ba48d60425c92b98209dde711f6e8f6f01a7 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 3 Nov 2026 16:32:53 +0530 Subject: [PATCH 06/21] Address some review comments --- .../src/foldPeerThreads.ts | 9 +++-- .../test/foldPeerThreads.test.ts | 25 +++++++++++++ packages/trueforge/src/apis/turns.ts | 12 +++++++ .../runtime/event-subscription/inMemory.ts | 22 +++++------- .../src/runtime/event-subscription/redis.ts | 12 ++++++- .../src/runtime/event-subscription/repair.ts | 11 ++++++ .../event-subscription/entryId.test.ts | 36 +++++++++++++++++++ .../trueforge/tests/unit/apis/turns.test.ts | 6 ++++ 8 files changed, 117 insertions(+), 16 deletions(-) diff --git a/packages/assistant-ui-runtime/src/foldPeerThreads.ts b/packages/assistant-ui-runtime/src/foldPeerThreads.ts index 0da2b8d997..6e51df167f 100644 --- a/packages/assistant-ui-runtime/src/foldPeerThreads.ts +++ b/packages/assistant-ui-runtime/src/foldPeerThreads.ts @@ -452,9 +452,14 @@ export function ingestStreamEvent(state: PeerThreadFoldState, message: TurnStrea if (!isEventDelta(message)) { if (state.ingestedEventIds.has(message.id)) { - return false; + // Inclusive resume replays the model.message shell so later deltas rebuild + // from a clean base instead of stacking on a partial in-memory fold. + if (message.type !== EVENT_TYPE.MODEL_MESSAGE) { + return false; + } + } else { + state.ingestedEventIds.add(message.id); } - state.ingestedEventIds.add(message.id); } if (isFoldNeutralInboundEvent(message)) { diff --git a/packages/assistant-ui-runtime/test/foldPeerThreads.test.ts b/packages/assistant-ui-runtime/test/foldPeerThreads.test.ts index ab9ab85c73..1ab7acebfb 100644 --- a/packages/assistant-ui-runtime/test/foldPeerThreads.test.ts +++ b/packages/assistant-ui-runtime/test/foldPeerThreads.test.ts @@ -55,6 +55,31 @@ describe('foldPeerThreads', () => { expect(buildRootAssistantContent(state)).toEqual([{ type: 'text', text: 'hello' }]); }); + it('replaces an already-ingested model.message shell on inclusive replay', () => { + const state = new PeerThreadFoldState(); + ingestStreamEvent( + state, + modelMessage({ + id: 'm1', + threadId: ROOT_THREAD_ID, + content: 'partial', + }), + ); + expect(buildRootAssistantContent(state)).toEqual([{ type: 'text', text: 'partial' }]); + + expect( + ingestStreamEvent( + state, + modelMessage({ + id: 'm1', + threadId: ROOT_THREAD_ID, + content: 'fresh', + }), + ), + ).toBe(true); + expect(buildRootAssistantContent(state)).toEqual([{ type: 'text', text: 'fresh' }]); + }); + it('captures sub-agent title from thread.created and tags nested messages', () => { const state = new PeerThreadFoldState(); diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index c184e93a15..120f2c7b9e 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -2047,6 +2047,18 @@ export function createTurnsRouter(deps: TurnsRouterDeps row.event.id === resumeEventId)) { + throw new HTTPException(400, { message: 'Event id is not part of this turn.' }); + } + } const turnEventStream = deps.eventSubscriptions.get(turnStreamId(requestContext.tenant_id, sessionId, turnId)); diff --git a/packages/trueforge/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index ca9d5bae05..0efb83435d 100644 --- a/packages/trueforge/src/runtime/event-subscription/inMemory.ts +++ b/packages/trueforge/src/runtime/event-subscription/inMemory.ts @@ -1,7 +1,6 @@ /** - * Single-process backend mirroring RedisEventSubscription semantics - * (ULID-millisecond entry ids, whole-stream TTL, poll-forever generators). - * Resume only works within one replica. + * Process-local EventSubscription backend for tests and single-replica runs. + * Entry ids are `{millisecond}-{n}`; streams expire after an optional TTL. */ import { EventEmitter, once } from 'node:events'; import { @@ -29,17 +28,14 @@ interface StoredEvent { } interface InMemoryStream { - /** Appended in Redis stream-id order. */ + /** Appended in entry-id order. */ events: StoredEvent[]; - /** Undefined = no TTL set yet (mirrors a Redis key without EXPIRE). */ + /** Undefined until the first put/repair sets a TTL. */ expiresAtMs?: number | undefined; expiryTimer?: NodeJS.Timeout | undefined; } -/** - * Owns the stream logs and their TTLs; the expiry timer is the in-process - * equivalent of Redis key eviction. - */ +/** Owns per-stream logs, owners, and TTL timers for this process. */ export class InMemoryEventStreamStore { private readonly streams = new Map>(); /** Survives stream eviction so a stale producer cannot recreate a missing log. */ @@ -104,7 +100,7 @@ export class InMemoryEventStreamStore { return this.getLiveStream(streamId)?.events[0]?.entryId; } - /** Drops the log and keeps the owner, matching Redis stream-key expiry. */ + /** Drops the log and keeps the owner so a stale producer cannot recreate it. */ evictStream(streamId: string): void { const stream = this.streams.get(streamId); if (stream === undefined) { @@ -178,8 +174,8 @@ export class InMemoryEventStreamStore { } /** - * Redis `XADD -*` : counter 0 when the millisecond is newer than the last - * id, otherwise the last millisecond's counter plus one. + * Next `{ms}-{n}` id: counter resets when `milliseconds` moves forward, + * otherwise increments within the current millisecond. */ private nextEntryId(stream: InMemoryStream, milliseconds: number): string { const last = stream.events[stream.events.length - 1]; @@ -259,7 +255,7 @@ export class InMemoryEventSubscription implements Even if (!stream) { return Promise.reject(new StreamGoneError(this.streamId)); } - // Expiring within the threshold is treated as already gone, mirroring the Redis backend. + // Same admission window as the shared subscribe threshold. if (stream.expiresAtMs !== undefined && stream.expiresAtMs - Date.now() < SUBSCRIBE_STREAM_THRESHOLD_MS) { return Promise.reject(new StreamGoneError(this.streamId)); } diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index e13a924ddd..148e8b777d 100644 --- a/packages/trueforge/src/runtime/event-subscription/redis.ts +++ b/packages/trueforge/src/runtime/event-subscription/redis.ts @@ -32,7 +32,17 @@ local startBound = '(' .. cursor local previous, lastMatch = nil, nil while true do local entries = redis.call('XRANGE', KEYS[1], startBound, '+', 'COUNT', batch) - if #entries == 0 then break end + if #entries == 0 then + -- Start was past the tip (e.g. a future ULID ms): park at the real tip so + -- XREAD still sees later live appends instead of waiting above them. + if previous == nil then + local tip = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) + if tip[1] ~= nil then + return tip[1][1] + end + end + break + end for i = 1, #entries do local entryId, fields = entries[i][1], entries[i][2] local data = nil diff --git a/packages/trueforge/src/runtime/event-subscription/repair.ts b/packages/trueforge/src/runtime/event-subscription/repair.ts index f194b0e7ef..2da4f7e8dd 100644 --- a/packages/trueforge/src/runtime/event-subscription/repair.ts +++ b/packages/trueforge/src/runtime/event-subscription/repair.ts @@ -87,6 +87,9 @@ function suffixAfterTail( if (tail === undefined) { return [...entries]; } + // Prefer an exact payload match. The live tip may be a model.message.delta + // that was never durable; fall back to the last durable row with the same id + // (the message shell) so later durable events still append. let match = -1; for (let index = 0; index < entries.length; index += 1) { const entry = entries[index]; @@ -94,6 +97,14 @@ function suffixAfterTail( match = index; } } + if (match === -1) { + for (let index = entries.length - 1; index >= 0; index -= 1) { + if (entries[index]?.id === tail.eventId) { + match = index; + break; + } + } + } if (match === -1) { return []; } diff --git a/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts b/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts index 447bd288dc..ce01d29d61 100644 --- a/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts @@ -39,3 +39,39 @@ describe('event id cursors', () => { ); }); }); + +describe('suffixAfterTail via repairLog', () => { + it('appends durable rows after a non-durable shared-id stream tip', async () => { + const { EventSubscriptionRegistry } = await import('../../../src/runtime/event-subscription'); + const registry = new EventSubscriptionRegistry<{ id: string; type: string; payload: { text: string } }>(undefined); + const streamId = `repair-delta-tip:${String(process.pid)}`; + const producer = registry.get(streamId); + const shell = { id: 'evt-1', type: 'model.message', payload: { text: 'shell' } }; + const deltaTip = { id: 'evt-1', type: 'model.message.delta', payload: { text: 'delta' } }; + const next = { id: 'evt-2', type: 'turn.done', payload: { text: 'done' } }; + await producer.put(shell, { streamTTLSeconds: 300 }); + await producer.put(deltaTip, { streamTTLSeconds: 300 }); + await producer.repairLog({ + entries: [ + { id: 'evt-1', event: shell }, + { id: 'evt-2', event: next }, + ], + streamTTLSeconds: 300, + }); + const abort = new AbortController(); + const generator = producer.poll({ signal: abort.signal }); + const received: string[] = []; + try { + for await (const event of generator) { + received.push(event.id); + if (received.length >= 3) { + break; + } + } + } finally { + abort.abort(); + await generator.return(undefined); + } + expect(received).toEqual(['evt-1', 'evt-1', 'evt-2']); + }); +}); diff --git a/packages/trueforge/tests/unit/apis/turns.test.ts b/packages/trueforge/tests/unit/apis/turns.test.ts index 7923762e9b..5abee524c0 100644 --- a/packages/trueforge/tests/unit/apis/turns.test.ts +++ b/packages/trueforge/tests/unit/apis/turns.test.ts @@ -1021,6 +1021,12 @@ describe('turns', () => { const invalid = await app.request(`/${sessionId}/turns/${pausedTurnId}/subscribe?from_event_id=not-a-ulid`); expect(invalid.status).toBe(400); + const { ulid } = await import('ulid'); + const foreign = await app.request( + `/${sessionId}/turns/${pausedTurnId}/subscribe?from_event_id=${ulid().toLowerCase()}`, + ); + expect(foreign.status).toBe(400); + const sequenceOnEventIdStream = await app.request( `/${sessionId}/turns/${pausedTurnId}/subscribe?after_sequence_number=1`, ); From 1c11b995529a21ec8f7997d17a0f3a366f207d8f Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 3 Nov 2026 17:17:51 +0530 Subject: [PATCH 07/21] Fix turn stream resume dropping deltas and appending a stale repair. Identical payloads were skipped, an in-progress event id was rejected, and a subscribe repair could land behind a newer tail. Sequence numbers are no longer stored, and an older client can still resume an event-id stream by the count it already applied. Co-authored-by: Cursor --- .changeset/event-id-subscribe-cursor.md | 1 + .../src/useTrueForgeAgentMessages.ts | 30 +- .../trueforge-core/src/agent-session/index.ts | 3 - .../src/agent-session/store/ISessionStore.ts | 35 +-- .../store/InMemorySessionStore.ts | 53 +--- .../agent-session/store/turnEventSequence.ts | 8 - .../agent-session/store/storeContractSuite.ts | 62 +--- packages/trueforge/src/apis/schedules.ts | 4 +- packages/trueforge/src/apis/turns.ts | 61 ++-- ...28_000003_session_event_sequence_number.ts | 18 -- .../session-store/PostgresSessionStore.ts | 12 +- .../postgres/session-store/queries/events.ts | 104 +------ packages/trueforge/src/db/postgres/types.ts | 5 - ...28_000003_session_event_sequence_number.ts | 20 -- .../session-store/SqliteSessionStore.ts | 12 +- .../db/sqlite/session-store/queries/events.ts | 102 +------ packages/trueforge/src/db/sqlite/types.ts | 5 - .../runtime/event-subscription/inMemory.ts | 138 ++++++--- .../src/runtime/event-subscription/index.ts | 28 +- .../src/runtime/event-subscription/redis.ts | 285 ++++++++---------- .../src/runtime/event-subscription/repair.ts | 39 ++- .../src/runtime/turnEventLogRepair.ts | 27 +- packages/trueforge/src/schemas/turn.ts | 4 +- .../event-subscription/contractSuite.ts | 66 ++-- .../inMemory.contract.test.ts | 12 +- .../event-subscription/redis.contract.test.ts | 7 +- .../runtime/event-subscription/repair.test.ts | 49 +++ .../turnEventLogRepair.test.ts | 19 +- .../trueforge/tests/unit/apis/turns.test.ts | 60 +++- 29 files changed, 505 insertions(+), 764 deletions(-) delete mode 100644 packages/trueforge-core/src/agent-session/store/turnEventSequence.ts delete mode 100644 packages/trueforge/src/db/postgres/migrations/20261028_000003_session_event_sequence_number.ts delete mode 100644 packages/trueforge/src/db/sqlite/migrations/20261028_000003_session_event_sequence_number.ts create mode 100644 packages/trueforge/tests/runtime/event-subscription/repair.test.ts diff --git a/.changeset/event-id-subscribe-cursor.md b/.changeset/event-id-subscribe-cursor.md index a378d48bb3..78176a73e4 100644 --- a/.changeset/event-id-subscribe-cursor.md +++ b/.changeset/event-id-subscribe-cursor.md @@ -1,5 +1,6 @@ --- '@truefoundry/trueforge': patch +'@truefoundry/trueforge-core': patch '@truefoundry/trueforge-assistant-ui-runtime': patch '@truefoundry/trueforge-ui': patch '@truefoundry/trueforge-sdk': patch diff --git a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts index f5bb263ff1..fbaba8be6f 100644 --- a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts +++ b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts @@ -604,8 +604,12 @@ export function useTrueForgeAgentMessages({ let subscribeCursor = cursorForTurn(snapshotRef.current, turnIdRef.current); let useSubscribe = false; let consecutiveFailures = 0; - let eventsThisAttempt = 0; - let attemptCursor: TurnSubscribeCursor | undefined = subscribeCursor; + // Inclusive resume replays the message shell and every delta already seen. + // Only events past the most this cursor ever delivered count as progress, + // so a stream that keeps dying mid-replay still exhausts the retry budget. + let progressCursorKey: string | undefined; + let mostDeliveredAtCursor = 0; + let deliveredAtCursorThisAttempt = 0; const track: StreamSequenceTrack = { onSequenceNumber: sequenceNumber => { @@ -613,17 +617,16 @@ export function useTrueForgeAgentMessages({ consecutiveFailures = 0; }, onSubscribeCursor: cursor => { - eventsThisAttempt += 1; - // Inclusive replay delivers the shell we already have. That is not - // new progress, so a stream that dies on the shell still counts. - const replayedShell = - eventsThisAttempt === 1 && - attemptCursor !== undefined && - 'fromEventId' in attemptCursor && - 'fromEventId' in cursor && - attemptCursor.fromEventId === cursor.fromEventId; subscribeCursor = cursor; - if (!replayedShell) { + const key = 'fromEventId' in cursor ? `from:${cursor.fromEventId}` : `after:${cursor.afterEventId}`; + if (key !== progressCursorKey) { + progressCursorKey = key; + mostDeliveredAtCursor = 0; + deliveredAtCursorThisAttempt = 0; + } + deliveredAtCursorThisAttempt += 1; + if (deliveredAtCursorThisAttempt > mostDeliveredAtCursor) { + mostDeliveredAtCursor = deliveredAtCursorThisAttempt; consecutiveFailures = 0; } }, @@ -702,8 +705,7 @@ export function useTrueForgeAgentMessages({ }; const readStream = (signal: AbortSignal): AsyncGenerator => { - eventsThisAttempt = 0; - attemptCursor = subscribeCursor; + deliveredAtCursorThisAttempt = 0; if (!useSubscribe) { return createStream(signal, track); } diff --git a/packages/trueforge-core/src/agent-session/index.ts b/packages/trueforge-core/src/agent-session/index.ts index 264fe5da4e..924b702509 100644 --- a/packages/trueforge-core/src/agent-session/index.ts +++ b/packages/trueforge-core/src/agent-session/index.ts @@ -86,7 +86,6 @@ export type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, - AssignTurnEventSequenceNumbersInput, ClaimTurnExecutorInput, CreateSessionInput, CreateTurnInput, @@ -112,8 +111,6 @@ export type { PatchThreadCapabilityStateInput, RemoveThreadsInput, TurnContextAppend, - TurnEventLogEntry, - TurnEventSequenceAssignment, TurnRecordWithoutSnapshot, UpdateSessionInput, UpdateTurnNonTerminalStateInput, diff --git a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts index 730588df8f..a9a1a8794f 100644 --- a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts @@ -193,32 +193,11 @@ export interface AppendToEventsInput extends TurnProgressKeys { events: PersistedTurnEvent[]; } -/** One durable event plus the stream sequence assigned when it was first appended. */ -export interface TurnEventLogEntry { - event: PersistedTurnEvent; - /** Null until a producer streams the event and records the number. */ - sequence_number: number | null; -} - export interface ListTurnEventLogInput { session_id: string; turn_id: string; } -/** Sequence written onto an existing event row. The event payload is not rewritten. */ -export interface TurnEventSequenceAssignment { - event_id: string; - sequence_number: number; -} - -export interface AssignTurnEventSequenceNumbersInput extends TurnProgressKeys { - /** - * Empty still checks `active_executor_id`. Producer repair uses that so a stale - * executor fails before it claims or appends the live stream. - */ - assignments: TurnEventSequenceAssignment[]; -} - export interface InsertTurnInboundEventsInput extends TurnProgressKeys { /** Empty array is a no-op. */ events: TurnUserEvent[]; @@ -439,18 +418,10 @@ export interface ISessionStore< appendToEvents(input: AppendToEventsInput, transaction?: TTransaction): Promise; /** - * Full turn event log in persist order (`event_id` ascending), with the stream - * sequence stored beside each payload. Missing turn throws TurnNotFoundError. - */ - listTurnEventLog(input: ListTurnEventLogInput): Promise; - - /** - * Records sequence numbers for events that do not have one yet. The caller must - * own a non-terminal turn. Already-stored numbers that match are a no-op; a - * different stored number is an invariant failure. Unknown `event_id` is too. - * An empty assignment still fences on `active_executor_id`. + * Full turn event log in persist order (`event_id` ascending). Missing turn + * throws TurnNotFoundError. */ - assignTurnEventSequenceNumbers(input: AssignTurnEventSequenceNumbersInput, transaction?: TTransaction): Promise; + listTurnEventLog(input: ListTurnEventLogInput): Promise; /** * Durable inbound send-event inbox for a tip the caller still owns. Tip must be diff --git a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts index 62fc4087ac..ff3034a46e 100644 --- a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts @@ -19,7 +19,6 @@ import type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, - AssignTurnEventSequenceNumbersInput, ClaimTurnExecutorInput, CreateSessionInput, CreateTurnInput, @@ -46,7 +45,6 @@ import type { PatchThreadsMCPAuthInput, RemoveThreadsInput, TurnContextAppend, - TurnEventLogEntry, TurnRecordWithoutSnapshot, UpdateSessionInput, UpdateTurnNonTerminalStateInput, @@ -71,8 +69,6 @@ import { TurnNotRunningError, TurnOwnershipLostError, } from './SessionStoreErrors'; -import { assertTurnEventSequenceNumber } from './turnEventSequence'; - /* eslint-disable @typescript-eslint/require-await -- in-memory store is synchronous; methods stay async so thrown SessionStore*Error reject as Promises for ISessionStore callers */ type StoredEvent = PersistedTurnEvent; @@ -208,8 +204,6 @@ export class InMemorySessionStore< private readonly sessions = new Map>(); private readonly turns = new Map>(); private readonly events = new Map(); - /** Stream sequence per event id, filled when the event is first streamed. */ - private readonly eventSequences = new Map>(); /** session_id:turn_id → inbound send-event inbox */ private readonly inboundEvents = new Map(); @@ -638,54 +632,13 @@ export class InMemorySessionStore< return; } - async listTurnEventLog(input: ListTurnEventLogInput): Promise { + async listTurnEventLog(input: ListTurnEventLogInput): Promise { this.requireTurn(input.session_id, input.turn_id); - const tKey = turnKey(input); - const list = this.events.get(tKey); - if (!list) { - throw new TurnNotFoundError(input.turn_id); - } - const sequences = this.eventSequences.get(tKey); - return [...list] - .sort((left, right) => left.id.localeCompare(right.id)) - .map(event => ({ - event: deepCopy(event), - sequence_number: sequences?.get(event.id) ?? null, - })); - } - - async assignTurnEventSequenceNumbers(input: AssignTurnEventSequenceNumbersInput): Promise { - this.requireOwnedTurn(input); - this.requireNonTerminalTurn(input); - if (input.assignments.length === 0) { - return; - } - const tKey = turnKey(input); - const list = this.events.get(tKey); + const list = this.events.get(turnKey(input)); if (!list) { throw new TurnNotFoundError(input.turn_id); } - let sequences = this.eventSequences.get(tKey); - if (!sequences) { - sequences = new Map(); - this.eventSequences.set(tKey, sequences); - } - for (const assignment of input.assignments) { - assertTurnEventSequenceNumber(assignment.sequence_number); - if (!list.some(event => event.id === assignment.event_id)) { - throw new SessionStoreInvariantError( - `event ${assignment.event_id} is not in the log for turn ${input.turn_id}`, - ); - } - const current = sequences.get(assignment.event_id); - if (current === undefined) { - sequences.set(assignment.event_id, assignment.sequence_number); - } else if (current !== assignment.sequence_number) { - throw new SessionStoreInvariantError( - `event ${assignment.event_id} already has sequence_number ${String(current)}`, - ); - } - } + return deepCopy([...list].sort((left, right) => (left.id < right.id ? -1 : left.id > right.id ? 1 : 0))); } async patchThreadsMCPAuth(input: PatchThreadsMCPAuthInput): Promise { diff --git a/packages/trueforge-core/src/agent-session/store/turnEventSequence.ts b/packages/trueforge-core/src/agent-session/store/turnEventSequence.ts deleted file mode 100644 index fb10f88d46..0000000000 --- a/packages/trueforge-core/src/agent-session/store/turnEventSequence.ts +++ /dev/null @@ -1,8 +0,0 @@ -import { SessionStoreInvariantError } from './SessionStoreErrors'; - -/** Stream sequences are positive safe integers. The SSE id is this number. */ -export function assertTurnEventSequenceNumber(sequenceNumber: number): void { - if (!Number.isSafeInteger(sequenceNumber) || sequenceNumber < 1) { - throw new SessionStoreInvariantError('sequence_number must be a positive integer'); - } -} diff --git a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts index 1da87e64ac..ac5fc9944e 100644 --- a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts +++ b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts @@ -4595,8 +4595,8 @@ export function runStoreContractSuite( }); } - describe('turn event sequence numbers', () => { - it('stores sequence_number beside the payload and rejects a stale owner', async () => { + describe('listTurnEventLog', () => { + it('returns every event of the turn in event_id order and rejects a missing turn', async () => { const store = createStore(); await seedSession(store); await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); @@ -4609,59 +4609,11 @@ export function runStoreContractSuite( events: [created, message], }); - const unsequenced = await store.listTurnEventLog({ session_id: sessionId, turn_id: 'turn-1' }); - expect(unsequenced.map(row => row.sequence_number)).toEqual([null, null]); - - await expect( - store.assignTurnEventSequenceNumbers({ - active_executor_id: 'other-executor', - session_id: sessionId, - turn_id: 'turn-1', - assignments: [{ event_id: created.id, sequence_number: 1 }], - }), - ).rejects.toBeInstanceOf(TurnOwnershipLostError); - expect( - (await store.listTurnEventLog({ session_id: sessionId, turn_id: 'turn-1' })).map(row => row.sequence_number), - ).toEqual([null, null]); - - await store.assignTurnEventSequenceNumbers({ - active_executor_id: TEST_ACTIVE_EXECUTOR_ID, - session_id: sessionId, - turn_id: 'turn-1', - assignments: [ - { event_id: created.id, sequence_number: 1 }, - { event_id: message.id, sequence_number: 2 }, - ], - }); - await store.assignTurnEventSequenceNumbers({ - active_executor_id: TEST_ACTIVE_EXECUTOR_ID, - session_id: sessionId, - turn_id: 'turn-1', - assignments: [{ event_id: created.id, sequence_number: 1 }], - }); - - const sequenced = await store.listTurnEventLog({ session_id: sessionId, turn_id: 'turn-1' }); - expect(sequenced.map(row => ({ id: row.event.id, sequence_number: row.sequence_number }))).toEqual([ - { id: created.id, sequence_number: 1 }, - { id: message.id, sequence_number: 2 }, - ]); - const listed = await store.listTurnEvents({ - session_id: sessionId, - turn_id: 'turn-1', - limit: 10, - page_token: undefined, - order: 'asc', - }); - expect(listed.data[0]).not.toHaveProperty('sequence_number'); - - await expect( - store.assignTurnEventSequenceNumbers({ - active_executor_id: TEST_ACTIVE_EXECUTOR_ID, - session_id: sessionId, - turn_id: 'turn-1', - assignments: [{ event_id: created.id, sequence_number: 9 }], - }), - ).rejects.toBeInstanceOf(SessionStoreInvariantError); + const log = await store.listTurnEventLog({ session_id: sessionId, turn_id: 'turn-1' }); + expect(log).toEqual([created, message].sort((left, right) => (left.id < right.id ? -1 : 1))); + await expect(store.listTurnEventLog({ session_id: sessionId, turn_id: 'missing' })).rejects.toBeInstanceOf( + TurnNotFoundError, + ); }); }); diff --git a/packages/trueforge/src/apis/schedules.ts b/packages/trueforge/src/apis/schedules.ts index 1898221c48..b7ea71a6a8 100644 --- a/packages/trueforge/src/apis/schedules.ts +++ b/packages/trueforge/src/apis/schedules.ts @@ -87,8 +87,8 @@ export interface ScheduleTurnExecutionDeps { sandboxProviderStore: ISandboxProviderStore; scheduleStore: IScheduleStore; sessions: Sessions; - /** Previous-turn cancel, and sequence numbers for the turn event stream. */ - sessionStore: Pick, 'getTurn' | 'listTurnEventLog' | 'assignTurnEventSequenceNumbers'>; + /** Previous-turn cancel, and rebuilding a missing event stream. */ + sessionStore: Pick, 'getTurn' | 'listTurnEventLog'>; webSearchProviderStore: IWebSearchProviderStore; } diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 120f2c7b9e..5ef2590db2 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -84,6 +84,7 @@ import { } from '../routes/turnRoutes'; import type { ActiveTurnRegistry } from '../runtime/activeTurns'; import { + eventIdTimestampMs, isLegacySequenceEntryId, LEGACY_SEQUENCE_ID_MAX, StreamGoneError, @@ -92,7 +93,6 @@ import { type EventSubscriptionPollOptions, type EventSubscriptionRegistry, } from '../runtime/event-subscription'; -import { eventIdTimestampMs } from '../runtime/event-subscription/entryId'; import { validateSandboxFilePath } from '../runtime/sandboxFilePath'; import { buildTurnSandbox, @@ -233,10 +233,8 @@ export type BeginTurnExecutionDeps = Pick< sandboxProviderStore: ISandboxProviderStore; webSearchProviderStore: IWebSearchProviderStore; skillStore: Pick; - /** - * Previous-turn cancel, durable sequence numbers, and rebuilding a missing event stream. - */ - sessionStore: Pick; + /** Previous-turn cancel, and rebuilding a missing event stream. */ + sessionStore: Pick; redis?: RedisClient | undefined; }; @@ -2013,11 +2011,6 @@ export function createTurnsRouter(deps: TurnsRouterDeps LEGACY_SEQUENCE_ID_MAX) { - throw new HTTPException(400, { - message: 'after_sequence_number is not supported above 50000.', - }); - } if (fromEventId !== undefined && eventIdTimestampMs(fromEventId) === undefined) { throw new HTTPException(400, { message: 'Invalid event id.' }); } @@ -2047,19 +2040,6 @@ export function createTurnsRouter(deps: TurnsRouterDeps row.event.id === resumeEventId)) { - throw new HTTPException(400, { message: 'Event id is not part of this turn.' }); - } - } - const turnEventStream = deps.eventSubscriptions.get(turnStreamId(requestContext.tenant_id, sessionId, turnId)); // A paused turn may have been offloaded after its Redis stream expired. @@ -2083,16 +2063,28 @@ export function createTurnsRouter(deps: TurnsRouterDeps LEGACY_SEQUENCE_ID_MAX) { throw new HTTPException(400, { - message: 'after_sequence_number is not supported for this stream.', + message: 'after_sequence_number is not supported above 50000.', }); + } else { + includeSequenceId = true; } - includeSequenceId = true; } // One lifecycle controller: server-side timeout, client disconnect, and // normal teardown all abort it, which ends poll even while it is parked @@ -2113,7 +2105,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps(deps: TurnsRouterDeps 0 ? seen : undefined; + seen += 1; await stream.writeSSE( includeSequenceId && sequenceNumber !== undefined ? turnEventSsePayload(event, sequenceNumber) - : turnEventSsePayload(event), + : compatIndex !== undefined + ? turnEventSsePayload(event, compatIndex) + : turnEventSsePayload(event), ); if (event.type === EventType.TURN_DONE) { break; diff --git a/packages/trueforge/src/db/postgres/migrations/20261028_000003_session_event_sequence_number.ts b/packages/trueforge/src/db/postgres/migrations/20261028_000003_session_event_sequence_number.ts deleted file mode 100644 index 01f99bca11..0000000000 --- a/packages/trueforge/src/db/postgres/migrations/20261028_000003_session_event_sequence_number.ts +++ /dev/null @@ -1,18 +0,0 @@ -import { sql, type Kysely } from 'kysely'; - -/** Stream sequence lives beside the event payload so a missing Redis key can be rebuilt. */ -export async function up(db: Kysely): Promise { - await sql` - SET LOCAL lock_timeout = '5s'; - ALTER TABLE session_event - ADD COLUMN sequence_number integer; - `.execute(db); -} - -export async function down(db: Kysely): Promise { - await sql` - SET LOCAL lock_timeout = '5s'; - ALTER TABLE session_event - DROP COLUMN IF EXISTS sequence_number; - `.execute(db); -} diff --git a/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts b/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts index 85559ba84b..58e7b85e42 100644 --- a/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts +++ b/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts @@ -7,7 +7,6 @@ import type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, - AssignTurnEventSequenceNumbersInput, ClaimTurnExecutorInput, CreateSessionInput, CreateTurnInput, @@ -32,7 +31,6 @@ import type { PatchThreadCapabilityStateInput, PatchThreadsMCPAuthInput, RemoveThreadsInput, - TurnEventLogEntry, TurnRecordWithoutSnapshot, UpdateSessionInput, UpdateTurnNonTerminalStateInput, @@ -47,7 +45,6 @@ import type { Database } from '../types'; import { patchThreadCapabilityState as patchThreadCapabilityStateQuery } from './queries/capabilities'; import { appendToEvents as appendToEventsQuery, - assignTurnEventSequenceNumbers as assignTurnEventSequenceNumbersQuery, listSessionEvents as listSessionEventsQuery, listTurnEventLog as listTurnEventLogQuery, listTurnEvents as listTurnEventsQuery, @@ -227,17 +224,10 @@ export class PostgresSessionStore implements ISessionStore return appendToEventsQuery(transaction ?? this.db, input); } - listTurnEventLog(input: ListTurnEventLogInput): Promise { + listTurnEventLog(input: ListTurnEventLogInput): Promise { return listTurnEventLogQuery(this.db, input); } - assignTurnEventSequenceNumbers( - input: AssignTurnEventSequenceNumbersInput, - transaction?: Transaction, - ): Promise { - return assignTurnEventSequenceNumbersQuery(transaction ?? this.db, input); - } - insertTurnInboundEvents(input: InsertTurnInboundEventsInput): Promise { return insertTurnInboundEventsQuery(this.db, input); } diff --git a/packages/trueforge/src/db/postgres/session-store/queries/events.ts b/packages/trueforge/src/db/postgres/session-store/queries/events.ts index dbf4949a1d..debfb0d478 100644 --- a/packages/trueforge/src/db/postgres/session-store/queries/events.ts +++ b/packages/trueforge/src/db/postgres/session-store/queries/events.ts @@ -2,11 +2,9 @@ import type { PersistedTurnEvent, SessionEventItem } from '@truefoundry/trueforg import type { TokenPagination } from '@truefoundry/trueforge-core/agent-session/schemas/pagination'; import type { AppendToEventsInput, - AssignTurnEventSequenceNumbersInput, ListSessionEventsInput, ListTurnEventLogInput, ListTurnEventsInput, - TurnEventLogEntry, } from '@truefoundry/trueforge-core/agent-session/store/ISessionStore'; import { decodeOffsetPageToken, @@ -19,13 +17,10 @@ import { } from '@truefoundry/trueforge-core/agent-session/store/SessionEventPageToken'; import { SessionNotFoundError, - SessionStoreInvariantError, TurnNotFoundError, } from '@truefoundry/trueforge-core/agent-session/store/SessionStoreErrors'; -import { assertTurnEventSequenceNumber } from '@truefoundry/trueforge-core/agent-session/store/turnEventSequence'; import type { Kysely, Transaction } from 'kysely'; import { sql } from 'kysely'; -import { isTransaction } from '../../../transaction'; import { json } from '../../sqlExpressions'; import type { Database } from '../../types'; import { unnestWithOrdinality, values } from '../sqlExpressions'; @@ -79,7 +74,7 @@ export async function appendToEvents( export async function listTurnEventLog( db: Kysely, input: ListTurnEventLogInput, -): Promise { +): Promise { const turn = await db .selectFrom('turn') .select('turn_id') @@ -92,106 +87,13 @@ export async function listTurnEventLog( const rows = await db .selectFrom('session_event') - .select(['event', 'sequence_number']) + .select('event') .where('session_id', '=', input.session_id) .where('turn_id', '=', input.turn_id) .orderBy('event_id', 'asc') .execute(); - return rows.map(row => ({ - event: row.event, - sequence_number: row.sequence_number, - })); -} - -export async function assignTurnEventSequenceNumbers( - db: Kysely | Transaction, - input: AssignTurnEventSequenceNumbersInput, -): Promise { - if (isTransaction(db)) { - await assignTurnEventSequenceNumbersInTransaction(db, input); - return; - } - await db.transaction().execute(trx => assignTurnEventSequenceNumbersInTransaction(trx, input)); -} - -async function assignTurnEventSequenceNumbersInTransaction( - db: Transaction, - input: AssignTurnEventSequenceNumbersInput, -): Promise { - for (const assignment of input.assignments) { - assertTurnEventSequenceNumber(assignment.sequence_number); - } - - const keys = { - session_id: input.session_id, - turn_id: input.turn_id, - active_executor_id: input.active_executor_id, - }; - const fence = await db - .selectFrom('turn') - .select(sql`1`.as('one')) - .where('session_id', '=', keys.session_id) - .where('turn_id', '=', keys.turn_id) - .where('active_executor_id', '=', keys.active_executor_id) - .where(sql`state->>'status' IN ('running', 'paused')`) - .forShare() - .executeTakeFirst(); - if (!fence) { - await classifyTurnFenceWriteFailure(db, keys); - } - if (input.assignments.length === 0) { - return; - } - - const rows = await db - .selectFrom('session_event') - .select(['event_id', 'sequence_number']) - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .where( - 'event_id', - 'in', - input.assignments.map(assignment => assignment.event_id), - ) - .execute(); - const stored = new Map(rows.map(row => [row.event_id, row.sequence_number])); - - for (const assignment of input.assignments) { - if (!stored.has(assignment.event_id)) { - throw new SessionStoreInvariantError(`event ${assignment.event_id} is not in the log for turn ${input.turn_id}`); - } - const current = stored.get(assignment.event_id); - if (current === null || current === undefined) { - const updated = await db - .updateTable('session_event') - .set({ sequence_number: assignment.sequence_number }) - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .where('event_id', '=', assignment.event_id) - .where('sequence_number', 'is', null) - .executeTakeFirst(); - if (updated.numUpdatedRows === 0n) { - const again = await db - .selectFrom('session_event') - .select('sequence_number') - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .where('event_id', '=', assignment.event_id) - .executeTakeFirst(); - if (again?.sequence_number !== assignment.sequence_number) { - throw new SessionStoreInvariantError( - `event ${assignment.event_id} already has sequence_number ${String(again?.sequence_number ?? 'none')}`, - ); - } - } - stored.set(assignment.event_id, assignment.sequence_number); - } else if (current !== assignment.sequence_number) { - throw new SessionStoreInvariantError( - `event ${assignment.event_id} already has sequence_number ${String(current)}`, - ); - } - } + return rows.map(row => row.event); } /** diff --git a/packages/trueforge/src/db/postgres/types.ts b/packages/trueforge/src/db/postgres/types.ts index 5f976c9c3b..f1b5edc136 100644 --- a/packages/trueforge/src/db/postgres/types.ts +++ b/packages/trueforge/src/db/postgres/types.ts @@ -272,11 +272,6 @@ export interface SessionEventTable { * Ordering still uses event_id. */ created_at: Date; - /** - * Stream sequence assigned when the event is first appended to the live log. - * Null until then. The event jsonb stays the client payload. - */ - sequence_number: ColumnType; } /** diff --git a/packages/trueforge/src/db/sqlite/migrations/20261028_000003_session_event_sequence_number.ts b/packages/trueforge/src/db/sqlite/migrations/20261028_000003_session_event_sequence_number.ts deleted file mode 100644 index 38bfb13084..0000000000 --- a/packages/trueforge/src/db/sqlite/migrations/20261028_000003_session_event_sequence_number.ts +++ /dev/null @@ -1,20 +0,0 @@ -import { sql, type Kysely } from 'kysely'; - -/** - * Stream sequence lives beside the event payload so a missing stream can be rebuilt. - * Mirrors db/postgres/migrations/20261028_000003_session_event_sequence_number.ts. - */ -export async function up(db: Kysely): Promise { - await db.transaction().execute(async trx => { - await sql` - ALTER TABLE session_event - ADD COLUMN sequence_number INTEGER - `.execute(trx); - }); -} - -export async function down(db: Kysely): Promise { - await db.transaction().execute(async trx => { - await sql`ALTER TABLE session_event DROP COLUMN sequence_number`.execute(trx); - }); -} diff --git a/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts b/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts index 086630d06c..fedd48c751 100644 --- a/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts +++ b/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts @@ -7,7 +7,6 @@ import type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, - AssignTurnEventSequenceNumbersInput, ClaimTurnExecutorInput, CreateSessionInput, CreateTurnInput, @@ -32,7 +31,6 @@ import type { PatchThreadCapabilityStateInput, PatchThreadsMCPAuthInput, RemoveThreadsInput, - TurnEventLogEntry, TurnRecordWithoutSnapshot, UpdateSessionInput, UpdateTurnNonTerminalStateInput, @@ -47,7 +45,6 @@ import type { Database } from '../types'; import { patchThreadCapabilityState as patchThreadCapabilityStateQuery } from './queries/capabilities'; import { appendToEvents as appendToEventsQuery, - assignTurnEventSequenceNumbers as assignTurnEventSequenceNumbersQuery, listSessionEvents as listSessionEventsQuery, listTurnEventLog as listTurnEventLogQuery, listTurnEvents as listTurnEventsQuery, @@ -225,17 +222,10 @@ export class SqliteSessionStore implements ISessionStore, return appendToEventsQuery(transaction ?? this.db, input); } - listTurnEventLog(input: ListTurnEventLogInput): Promise { + listTurnEventLog(input: ListTurnEventLogInput): Promise { return listTurnEventLogQuery(this.db, input); } - assignTurnEventSequenceNumbers( - input: AssignTurnEventSequenceNumbersInput, - transaction?: Transaction, - ): Promise { - return assignTurnEventSequenceNumbersQuery(transaction ?? this.db, input); - } - insertTurnInboundEvents(input: InsertTurnInboundEventsInput): Promise { return insertTurnInboundEventsQuery(this.db, input); } diff --git a/packages/trueforge/src/db/sqlite/session-store/queries/events.ts b/packages/trueforge/src/db/sqlite/session-store/queries/events.ts index 99db6a5eee..7b924edf21 100644 --- a/packages/trueforge/src/db/sqlite/session-store/queries/events.ts +++ b/packages/trueforge/src/db/sqlite/session-store/queries/events.ts @@ -2,11 +2,9 @@ import type { PersistedTurnEvent, SessionEventItem } from '@truefoundry/trueforg import type { TokenPagination } from '@truefoundry/trueforge-core/agent-session/schemas/pagination'; import type { AppendToEventsInput, - AssignTurnEventSequenceNumbersInput, ListSessionEventsInput, ListTurnEventLogInput, ListTurnEventsInput, - TurnEventLogEntry, } from '@truefoundry/trueforge-core/agent-session/store/ISessionStore'; import { decodeOffsetPageToken, @@ -19,10 +17,8 @@ import { } from '@truefoundry/trueforge-core/agent-session/store/SessionEventPageToken'; import { SessionNotFoundError, - SessionStoreInvariantError, TurnNotFoundError, } from '@truefoundry/trueforge-core/agent-session/store/SessionStoreErrors'; -import { assertTurnEventSequenceNumber } from '@truefoundry/trueforge-core/agent-session/store/turnEventSequence'; import { sql, type Kysely, type Transaction } from 'kysely'; import { isTransaction } from '../../../transaction'; import { jsonbBind, jsonText } from '../../sqlExpressions'; @@ -77,7 +73,7 @@ async function appendToEventsInTransaction(db: Transaction, input: App export async function listTurnEventLog( db: Kysely, input: ListTurnEventLogInput, -): Promise { +): Promise { const turn = await db .selectFrom('turn') .select('turn_id') @@ -90,105 +86,13 @@ export async function listTurnEventLog( const rows = await db .selectFrom('session_event') - .select(['sequence_number', jsonText(sql.ref('event')).as('event')]) + .select(jsonText(sql.ref('event')).as('event')) .where('session_id', '=', input.session_id) .where('turn_id', '=', input.turn_id) .orderBy('event_id', 'asc') .execute(); - return rows.map(row => ({ - event: row.event, - sequence_number: row.sequence_number, - })); -} - -export async function assignTurnEventSequenceNumbers( - db: DbOrTrx, - input: AssignTurnEventSequenceNumbersInput, -): Promise { - if (isTransaction(db)) { - await assignTurnEventSequenceNumbersInTransaction(db, input); - return; - } - await db.transaction().execute(trx => assignTurnEventSequenceNumbersInTransaction(trx, input)); -} - -async function assignTurnEventSequenceNumbersInTransaction( - db: Transaction, - input: AssignTurnEventSequenceNumbersInput, -): Promise { - for (const assignment of input.assignments) { - assertTurnEventSequenceNumber(assignment.sequence_number); - } - - const keys: TurnKeys = { - session_id: input.session_id, - turn_id: input.turn_id, - active_executor_id: input.active_executor_id, - }; - const fenceRow = await db - .selectFrom('turn') - .select(sql`1`.as('one')) - .where('session_id', '=', keys.session_id) - .where('turn_id', '=', keys.turn_id) - .where('active_executor_id', '=', keys.active_executor_id) - .where(sql`state->>'status' IN ('running', 'paused')`) - .executeTakeFirst(); - if (!fenceRow) { - await classifyTurnFenceWriteFailure(db, keys); - } - if (input.assignments.length === 0) { - return; - } - - const rows = await db - .selectFrom('session_event') - .select(['event_id', 'sequence_number']) - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .where( - 'event_id', - 'in', - input.assignments.map(assignment => assignment.event_id), - ) - .execute(); - const stored = new Map(rows.map(row => [row.event_id, row.sequence_number])); - - for (const assignment of input.assignments) { - if (!stored.has(assignment.event_id)) { - throw new SessionStoreInvariantError(`event ${assignment.event_id} is not in the log for turn ${input.turn_id}`); - } - const current = stored.get(assignment.event_id); - if (current === null || current === undefined) { - const updated = await db - .updateTable('session_event') - .set({ sequence_number: assignment.sequence_number }) - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .where('event_id', '=', assignment.event_id) - .where('sequence_number', 'is', null) - .executeTakeFirst(); - if (updated.numUpdatedRows === 0n) { - const again = await db - .selectFrom('session_event') - .select('sequence_number') - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .where('event_id', '=', assignment.event_id) - .executeTakeFirst(); - if (again?.sequence_number !== assignment.sequence_number) { - throw new SessionStoreInvariantError( - `event ${assignment.event_id} already has sequence_number ${String(again?.sequence_number ?? 'none')}`, - ); - } - } - stored.set(assignment.event_id, assignment.sequence_number); - } else if (current !== assignment.sequence_number) { - throw new SessionStoreInvariantError( - `event ${assignment.event_id} already has sequence_number ${String(current)}`, - ); - } - } + return rows.map(row => row.event); } /** diff --git a/packages/trueforge/src/db/sqlite/types.ts b/packages/trueforge/src/db/sqlite/types.ts index 2f40811eec..9100a6d19e 100644 --- a/packages/trueforge/src/db/sqlite/types.ts +++ b/packages/trueforge/src/db/sqlite/types.ts @@ -169,11 +169,6 @@ export interface SessionEventTable { event_id: string; event: JsonbColumn; created_at: string; - /** - * Stream sequence assigned when the event is first appended to the live log. - * Null until then. The event blob stays the client payload. - */ - sequence_number: ColumnType; } /** diff --git a/packages/trueforge/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index 0efb83435d..d97da3f354 100644 --- a/packages/trueforge/src/runtime/event-subscription/inMemory.ts +++ b/packages/trueforge/src/runtime/event-subscription/inMemory.ts @@ -20,7 +20,11 @@ import { isEntryAfterCursor, isLegacySequenceEntryId, } from './entryId'; -import { applyEventLogRepair, type RepairLogOptions } from './repair'; +import { applyEventLogRepair, type RepairLogOptions, type StreamTail } from './repair'; + +function ttlDeadline(streamTTLSeconds: number | undefined): number | undefined { + return streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? Date.now() + streamTTLSeconds * 1_000 : undefined; +} interface StoredEvent { entryId: string; @@ -38,8 +42,11 @@ interface InMemoryStream { /** Owns per-stream logs, owners, and TTL timers for this process. */ export class InMemoryEventStreamStore { private readonly streams = new Map>(); - /** Survives stream eviction so a stale producer cannot recreate a missing log. */ - private readonly owners = new Map(); + /** + * Survives an explicit eviction so a stale producer cannot recreate a missing + * log. Expires with the stream TTL so idle turns do not accumulate owners. + */ + private readonly owners = new Map(); /** Rings per stream id on append and expiry so parked pollers wake instantly. */ private readonly changes = new EventEmitter(); @@ -48,58 +55,72 @@ export class InMemoryEventStreamStore { this.changes.setMaxListeners(0); } - claimOwner(streamId: string, executorId: string): void { - this.owners.set(streamId, executorId); + claimOwner(input: { streamId: string; executorId: string; streamTTLSeconds: number | undefined }): void { + this.owners.set(input.streamId, { + executorId: input.executorId, + expiresAtMs: ttlDeadline(input.streamTTLSeconds), + }); } /** - * Appends with a `{millisecond}-{n}` id. The same payload already at the tail - * is a duplicate. A millisecond older than the tail reuses the tail millisecond. + * Live put with a `{millisecond}-{n}` id. A millisecond older than the tail + * reuses the tail millisecond. Identical payloads are kept. */ appendEvent(input: { streamId: string; event: T; executorId: string | undefined; - checkOwner: boolean; streamTTLSeconds: number | undefined; - }): 'ok' | 'duplicate' | 'owner_mismatch' { - const owner = this.owners.get(input.streamId); - if (input.checkOwner && owner !== undefined && owner !== input.executorId) { + }): 'ok' | 'owner_mismatch' { + const owner = this.liveOwner(input.streamId); + if (owner !== undefined && owner !== input.executorId) { return 'owner_mismatch'; } const stream = this.getLiveStream(input.streamId) ?? { events: [] }; - const event = structuredClone(input.event); - const payload = JSON.stringify(event); - const last = stream.events[stream.events.length - 1]; - if (last !== undefined && JSON.stringify(last.event) === payload) { - return 'duplicate'; + this.streams.set(input.streamId, stream); + this.push(stream, input.event); + this.refreshStreamTtl(input.streamId, input.streamTTLSeconds); + this.changes.emit(input.streamId); + return 'ok'; + } + + /** Repair backfill: appends only while the tail is still `tailEntryId`. */ + appendAfterTail(input: { + streamId: string; + tailEntryId: string | undefined; + events: T[]; + streamTTLSeconds: number | undefined; + }): 'ok' | 'stale' { + const stream = this.getLiveStream(input.streamId) ?? { events: [] }; + if (stream.events[stream.events.length - 1]?.entryId !== input.tailEntryId) { + return 'stale'; } this.streams.set(input.streamId, stream); - stream.events.push({ - entryId: this.nextEntryId(stream, eventIdTimestampMs(event.id) ?? Date.now()), - event, - }); - if (input.streamTTLSeconds !== undefined && input.streamTTLSeconds > 0) { - stream.expiresAtMs = Date.now() + input.streamTTLSeconds * 1_000; - this.scheduleExpiry(input.streamId, stream); + for (const event of input.events) { + this.push(stream, event); } + this.refreshStreamTtl(input.streamId, input.streamTTLSeconds); this.changes.emit(input.streamId); return 'ok'; } - readTail(streamId: string): { eventId: string; payload: string } | undefined { + readTail(streamId: string): StreamTail | undefined { const stream = this.getLiveStream(streamId); const last = stream?.events[stream.events.length - 1]; if (last === undefined) { return undefined; } - return { eventId: last.event.id, payload: JSON.stringify(last.event) }; + return { entryId: last.entryId, eventId: last.event.id, payload: JSON.stringify(last.event) }; } firstEntryId(streamId: string): string | undefined { return this.getLiveStream(streamId)?.events[0]?.entryId; } + hasEventId(streamId: string, eventId: string): boolean { + return this.getLiveStream(streamId)?.events.some(stored => stored.event.id === eventId) ?? false; + } + /** Drops the log and keeps the owner so a stale producer cannot recreate it. */ evictStream(streamId: string): void { const stream = this.streams.get(streamId); @@ -109,13 +130,21 @@ export class InMemoryEventStreamStore { this.dropStream(streamId, stream); } - /** Refreshes TTL on a live stream. Sequence numbers already follow the log length. */ - refreshStreamTtl(streamId: string, streamTTLSeconds?: number): void { + /** Refreshes TTL on a live stream and its owner. */ + refreshStreamTtl(streamId: string, streamTTLSeconds: number | undefined): void { + const deadline = ttlDeadline(streamTTLSeconds); + if (deadline === undefined) { + return; + } + const owner = this.owners.get(streamId); + if (owner !== undefined) { + owner.expiresAtMs = deadline; + } const stream = this.getLiveStream(streamId); - if (!stream || !streamTTLSeconds || streamTTLSeconds <= 0) { + if (!stream) { return; } - stream.expiresAtMs = Date.now() + streamTTLSeconds * 1_000; + stream.expiresAtMs = deadline; this.scheduleExpiry(streamId, stream); } @@ -127,6 +156,7 @@ export class InMemoryEventStreamStore { } if (stream.expiresAtMs !== undefined && stream.expiresAtMs <= Date.now()) { this.dropStream(streamId, stream); + this.dropOwnerIfExpired(streamId); return undefined; } return stream; @@ -167,12 +197,30 @@ export class InMemoryEventStreamStore { stream.expiryTimer = setTimeout( () => { this.dropStream(streamId, stream); + this.dropOwnerIfExpired(streamId); }, Math.max(0, stream.expiresAtMs - Date.now()), ); stream.expiryTimer.unref(); } + private liveOwner(streamId: string): string | undefined { + this.dropOwnerIfExpired(streamId); + return this.owners.get(streamId)?.executorId; + } + + private dropOwnerIfExpired(streamId: string): void { + const expiresAtMs = this.owners.get(streamId)?.expiresAtMs; + if (expiresAtMs !== undefined && expiresAtMs <= Date.now()) { + this.owners.delete(streamId); + } + } + + private push(stream: InMemoryStream, event: T): void { + const copy = structuredClone(event); + stream.events.push({ entryId: this.nextEntryId(stream, eventIdTimestampMs(copy.id) ?? Date.now()), event: copy }); + } + /** * Next `{ms}-{n}` id: counter resets when `milliseconds` moves forward, * otherwise increments within the current millisecond. @@ -207,23 +255,22 @@ export class InMemoryEventSubscription implements Even executorId === undefined ? undefined : () => { - this.store.claimOwner(this.streamId, executorId); + this.store.claimOwner({ + streamId: this.streamId, + executorId, + streamTTLSeconds: options.streamTTLSeconds, + }); return Promise.resolve(); }, readTail: () => Promise.resolve(this.store.readTail(this.streamId)), - append: entry => { - const outcome = this.store.appendEvent({ - streamId: this.streamId, - event: entry.event, - executorId, - checkOwner: false, - streamTTLSeconds: options.streamTTLSeconds, - }); - if (outcome === 'owner_mismatch') { - return Promise.reject(new StreamOwnerMismatchError(this.streamId)); - } - return Promise.resolve(outcome === 'duplicate' ? 'duplicate' : 'appended'); - }, + appendAfterTail: batch => + Promise.resolve( + this.store.appendAfterTail({ + streamId: this.streamId, + ...batch, + streamTTLSeconds: options.streamTTLSeconds, + }), + ), refreshTtl: () => { this.store.refreshStreamTtl(this.streamId, options.streamTTLSeconds); return Promise.resolve(); @@ -237,7 +284,6 @@ export class InMemoryEventSubscription implements Even streamId: this.streamId, event, executorId: options?.executorId, - checkOwner: true, streamTTLSeconds: options?.streamTTLSeconds, }); if (outcome === 'owner_mismatch') { @@ -246,6 +292,10 @@ export class InMemoryEventSubscription implements Even return Promise.resolve(); } + hasEventId(eventId: string): Promise { + return Promise.resolve(this.store.hasEventId(this.streamId, eventId)); + } + firstEntryId(): Promise { return Promise.resolve(this.store.firstEntryId(this.streamId)); } diff --git a/packages/trueforge/src/runtime/event-subscription/index.ts b/packages/trueforge/src/runtime/event-subscription/index.ts index 5d1adaa468..0e5125ee36 100644 --- a/packages/trueforge/src/runtime/event-subscription/index.ts +++ b/packages/trueforge/src/runtime/event-subscription/index.ts @@ -9,7 +9,7 @@ export interface EventSubscriptionPutOptions { /** Sets the lifetime of the whole stream after this append. */ streamTTLSeconds?: number | undefined; /** - * XADD runs only when the stream-owner key is unset or equals this id. + * The append runs only when the stream-owner key is unset or equals this id. * A missing stream still rejects a caller that does not own the key. */ executorId?: string | undefined; @@ -18,7 +18,7 @@ export interface EventSubscriptionPutOptions { /** `sequence_number` is set only when replaying an old `{sequence}-1` entry. */ export type SequencedEvent = T & { sequence_number?: number }; -export { isLegacySequenceEntryId, LEGACY_SEQUENCE_ID_MAX } from './entryId'; +export { eventIdTimestampMs, isLegacySequenceEntryId, LEGACY_SEQUENCE_ID_MAX } from './entryId'; export type EventSubscriptionPollOptions = { /** Ends the generator promptly on abort, releasing any parked idle wait. */ @@ -65,6 +65,9 @@ export interface EventSubscription { /** First entry id, or undefined when the stream has no entries yet. */ firstEntryId(): Promise; + /** Whether some entry on the live stream carries this payload id. */ + hasEventId(eventId: string): Promise; + /** * Subscription admission check, run before SSE headers are sent: rejects with * {@link StreamGoneError} when the stream is missing or expires within @@ -82,9 +85,14 @@ export interface EventSubscription { */ export class EventSubscriptionRegistry { /** One store for the whole process so producers and subscribers share streams. */ - private readonly memoryStore = new InMemoryEventStreamStore(); + private readonly memoryStore: InMemoryEventStreamStore; - constructor(private readonly redis: RedisClient | undefined) {} + constructor( + private readonly redis: RedisClient | undefined, + memoryStore?: InMemoryEventStreamStore, + ) { + this.memoryStore = memoryStore ?? new InMemoryEventStreamStore(); + } get(streamId: string): EventSubscription { if (this.redis) { @@ -92,18 +100,6 @@ export class EventSubscriptionRegistry { } return new InMemoryEventSubscription(this.memoryStore, streamId); } - - /** - * Drops the live log the way a TTL expiry does, and keeps the stream-owner key - * so a stale producer still cannot recreate it. - */ - evictStream(streamId: string): Promise { - if (this.redis) { - return this.redis.del(streamId).then(() => undefined); - } - this.memoryStore.evictStream(streamId); - return Promise.resolve(); - } } /** Reject new subscriptions when the stream expires within this window. */ diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index 148e8b777d..1f9290d4de 100644 --- a/packages/trueforge/src/runtime/event-subscription/redis.ts +++ b/packages/trueforge/src/runtime/event-subscription/redis.ts @@ -10,7 +10,7 @@ import { type SequencedEvent, } from '.'; import { entryIdMilliseconds, eventIdSeekCursor, eventIdTimestampMs, isLegacySequenceEntryId } from './entryId'; -import { applyEventLogRepair, eventIdFromStreamData, type RepairLogOptions } from './repair'; +import { applyEventLogRepair, eventIdFromStreamData, type RepairLogOptions, type StreamTail } from './repair'; const SUBSCRIBE_STREAM_POLL_ITEMS_COUNT = 100; /** Delay between XREAD polls while a stream has no new events (a blocking XREAD would monopolize the shared connection). */ @@ -22,9 +22,10 @@ const SUBSCRIBE_STREAM_POLL_SLEEP_INTERVAL_MS = 1_000; * * KEYS[1]=stream ARGV[1]=from|after ARGV[2]=payload id ARGV[3]=start ARGV[4]=batch * + * Returns {cursor, found}. * from → stop at the first match; return the prior entry (XREAD is exclusive). * after → walk the contiguous shared-id group (message + deltas); return its last entry. - * missing → tip (live-only). + * missing → tip (live-only), found = 0. */ const SEEK_EVENT_ID_CURSOR = ` local mode, targetId, cursor, batch = ARGV[1], ARGV[2], ARGV[3], tonumber(ARGV[4]) @@ -38,7 +39,7 @@ while true do if previous == nil then local tip = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) if tip[1] ~= nil then - return tip[1][1] + return {tip[1][1], 0} end end break @@ -58,98 +59,110 @@ while true do end if mode == 'from' then if payloadId == targetId then - return previous or cursor + return {previous or cursor, 1} end elseif payloadId == targetId then lastMatch = entryId elseif lastMatch ~= nil then - return lastMatch + return {lastMatch, 1} end previous = entryId end startBound = '(' .. entries[#entries][1] if #entries < batch then break end end -return lastMatch or previous or cursor +if lastMatch ~= nil then + return {lastMatch, 1} +end +return {previous or cursor, 0} `; /** - * One script for put and for suffix backfill. Put checks the stream-owner key. - * Repair passes check-owner 0 so a subscribe can append without claiming. - * The id is `{millisecond}-*`. Redis 7 assigns the counter. The same payload - * already at the tail is a duplicate. A millisecond older than the tail is a - * conflict; the caller retries with the tail millisecond. + * Live put. KEYS[1]=owner KEYS[2]=stream ARGV[1]=executor ARGV[2]=ms ARGV[3]=payload ARGV[4]=ttl + * + * Rejects when another executor owns the stream. The entry id is `{ms}-*`, raised + * to the tail millisecond when the event id is older than the tail (a delta of a + * message whose id was minted before later events). Identical payloads are kept: + * two deltas of the same message can carry the same text. */ -const APPEND_STREAM_ENTRY = ` -local checkOwner = ARGV[1] -if checkOwner == '1' then - local owner = redis.call('GET', KEYS[1]) - if owner ~= false and owner ~= ARGV[2] then - return {'owner_mismatch', ''} - end +const PUT_STREAM_ENTRY = ` +local owner = redis.call('GET', KEYS[1]) +if owner and owner ~= ARGV[1] then + return 'owner_mismatch' end +local ms = ARGV[2] local tail = redis.call('XREVRANGE', KEYS[2], '+', '-', 'COUNT', 1) if tail[1] ~= nil then - local fields = tail[1][2] - for i = 1, #fields, 2 do - if fields[i] == 'data' and fields[i + 1] == ARGV[4] then - return {'duplicate', ''} - end + local tailMs = string.match(tail[1][1], '^%d+') + if tonumber(tailMs) > tonumber(ms) then + ms = tailMs end end -local added = redis.pcall('XADD', KEYS[2], ARGV[3], 'data', ARGV[4]) -if type(added) == 'table' and added['err'] ~= nil then - local err = added['err'] - if string.find(err, 'equal or smaller', 1, true) then - local tail = redis.call('XREVRANGE', KEYS[2], '+', '-', 'COUNT', 1) - local data = '' - if tail[1] ~= nil then - local fields = tail[1][2] - for i = 1, #fields, 2 do - if fields[i] == 'data' then - data = fields[i + 1] - break - end - end - end - return {'conflict', data} +redis.call('XADD', KEYS[2], ms .. '-*', 'data', ARGV[3]) +local ttl = tonumber(ARGV[4]) +if ttl > 0 then + redis.call('EXPIRE', KEYS[2], ttl) + if owner then + redis.call('EXPIRE', KEYS[1], ttl) end - return {'error', err} end -local ttl = tonumber(ARGV[5]) -if ttl ~= nil and ttl > 0 then - redis.call('EXPIRE', KEYS[2], ttl) +return 'ok' +`; + +/** + * Repair backfill. KEYS[1]=stream ARGV[1]=expected tail id ('' = empty) ARGV[2]=ttl + * ARGV[3..]=ms, payload pairs. + * + * Appends only while the tail is still the one the suffix was computed from; + * otherwise returns 'stale' and the caller recomputes. + */ +const APPEND_AFTER_TAIL = ` +local tail = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) +local tailId = '' +if tail[1] ~= nil then + tailId = tail[1][1] +end +if tailId ~= ARGV[1] then + return 'stale' +end +local tailMs = string.match(tailId, '^%d+') or '0' +for i = 3, #ARGV, 2 do + local ms = ARGV[i] + if tonumber(tailMs) > tonumber(ms) then + ms = tailMs + end + local added = redis.call('XADD', KEYS[1], ms .. '-*', 'data', ARGV[i + 1]) + tailMs = string.match(added, '^%d+') +end +local ttl = tonumber(ARGV[2]) +if ttl > 0 then + redis.call('EXPIRE', KEYS[1], ttl) end -return {'ok', ''} +return 'ok' `; function streamOwnerKey(streamId: string): string { return `${streamId}:owner`; } -type AppendStatus = 'ok' | 'duplicate' | 'conflict' | 'owner_mismatch' | 'error'; - -function isAppendStatus(value: unknown): value is AppendStatus { - return ( - value === 'ok' || value === 'duplicate' || value === 'conflict' || value === 'owner_mismatch' || value === 'error' - ); +function ttlArgument(streamTTLSeconds: number | undefined): string { + return String(streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? streamTTLSeconds : 0); } -function parseAppendReply(value: unknown): { status: AppendStatus; detail: string } { - if (!Array.isArray(value) || value.length !== 2) { - throw new Error('Unexpected stream append reply'); - } - const status: unknown = Object.prototype.hasOwnProperty.call(value, 0) ? Reflect.get(value, 0) : undefined; - const detail: unknown = Object.prototype.hasOwnProperty.call(value, 1) ? Reflect.get(value, 1) : undefined; - if (isAppendStatus(status) && typeof detail === 'string') { - return { status, detail }; +function parseSeekReply(value: unknown): { cursor: string; found: boolean } { + if (Array.isArray(value) && value.length === 2) { + const cursor: unknown = value[0]; + const found: unknown = value[1]; + if (typeof cursor === 'string' && (found === 0 || found === 1)) { + return { cursor, found: found === 1 }; + } } - throw new Error('Unexpected stream append reply'); + throw new Error('Unexpected event-id seek reply'); } /** * One Redis Stream. Entry ids are `{ulidMillisecond}-{n}` from `XADD`. The owner - * key is separate from the stream so a missing stream still rejects a stale producer. + * key is separate from the stream so a deleted stream still rejects a stale producer. */ export class RedisEventSubscription implements EventSubscription { constructor( @@ -161,30 +174,35 @@ export class RedisEventSubscription implements EventSu const executorId = options.executorId; await applyEventLogRepair({ entries: options.entries, - claimOwner: - executorId === undefined - ? undefined - : () => this.redis.set(streamOwnerKey(this.streamId), executorId).then(() => undefined), + claimOwner: executorId === undefined ? undefined : () => this.claimOwner(executorId, options.streamTTLSeconds), readTail: () => this.readTail(), - append: entry => - this.appendEvent({ - event: entry.event, - checkOwner: false, - executorId: '', - streamTTLSeconds: options.streamTTLSeconds, - }), + appendAfterTail: batch => this.appendAfterTail({ ...batch, streamTTLSeconds: options.streamTTLSeconds }), refreshTtl: () => this.refreshTtl(options.streamTTLSeconds), refreshTtlOnUntouchedStream: executorId !== undefined, }); } async put(event: T, options?: EventSubscriptionPutOptions): Promise { - await this.appendEvent({ - event, - checkOwner: true, - executorId: options?.executorId ?? '', - streamTTLSeconds: options?.streamTTLSeconds, + const reply: unknown = await this.redis.eval(PUT_STREAM_ENTRY, { + keys: [streamOwnerKey(this.streamId), this.streamId], + arguments: [ + options?.executorId ?? '', + this.entryMillisecondFor(event), + JSON.stringify(event), + ttlArgument(options?.streamTTLSeconds), + ], }); + if (reply === 'owner_mismatch') { + throw new StreamOwnerMismatchError(this.streamId); + } + if (reply !== 'ok') { + throw new Error(`Unexpected stream put reply on ${this.streamId}`); + } + } + + async hasEventId(eventId: string): Promise { + const { found } = await this.seekEventIdCursor({ mode: 'from', eventId }); + return found; } async firstEntryId(): Promise { @@ -256,7 +274,7 @@ export class RedisEventSubscription implements EventSu } const mode = fromEventId !== undefined ? 'from' : 'after'; const signal = options.signal; - let cursor = await this.seekEventIdCursor({ mode, eventId }); + let { cursor } = await this.seekEventIdCursor({ mode, eventId }); for (;;) { const messages = await this.readAfter(cursor, signal); if (messages === undefined) { @@ -284,7 +302,10 @@ export class RedisEventSubscription implements EventSu * Jump near the ULID millisecond (or stream start), walk forward in Lua until * the payload id, return exclusive XREAD cursor. See SEEK_EVENT_ID_CURSOR. */ - private async seekEventIdCursor(input: { mode: 'from' | 'after'; eventId: string }): Promise { + private async seekEventIdCursor(input: { + mode: 'from' | 'after'; + eventId: string; + }): Promise<{ cursor: string; found: boolean }> { const startCursor = eventIdSeekCursor({ eventId: input.eventId, firstEntryId: await this.firstEntryId(), @@ -293,10 +314,7 @@ export class RedisEventSubscription implements EventSu keys: [this.streamId], arguments: [input.mode, input.eventId, startCursor, String(SUBSCRIBE_STREAM_POLL_ITEMS_COUNT)], }); - if (typeof reply !== 'string') { - throw new Error('Unexpected event-id seek reply'); - } - return reply; + return parseSeekReply(reply); } private async readAfter( @@ -362,7 +380,7 @@ export class RedisEventSubscription implements EventSu ); } - private async readTail(): Promise<{ eventId: string; payload: string } | undefined> { + private async readTail(): Promise { const entries = await this.redis.xRevRange(this.streamId, '+', '-', { COUNT: 1 }); const latest = entries[0]; if (!latest) { @@ -376,93 +394,46 @@ export class RedisEventSubscription implements EventSu if (eventId === undefined) { throw new Error(`Corrupt stream entry ${latest.id} on ${this.streamId}: missing event id`); } - return { eventId, payload: data }; + return { entryId: latest.id, eventId, payload: data }; } - private async lastEntryId(): Promise { - const entries = await this.redis.xRevRange(this.streamId, '+', '-', { COUNT: 1 }); - return entries[0]?.id; + private async claimOwner(executorId: string, streamTTLSeconds: number | undefined): Promise { + const key = streamOwnerKey(this.streamId); + const multi = this.redis.multi().set(key, executorId); + if (streamTTLSeconds !== undefined && streamTTLSeconds > 0) { + multi.expire(key, streamTTLSeconds); + } + await multi.exec(); } + /** The owner key shares the stream lifetime so it does not outlive an idle turn. */ private async refreshTtl(streamTTLSeconds: number | undefined): Promise { if (streamTTLSeconds !== undefined && streamTTLSeconds > 0) { - await this.redis.expire(this.streamId, streamTTLSeconds); + await this.redis + .multi() + .expire(this.streamId, streamTTLSeconds) + .expire(streamOwnerKey(this.streamId), streamTTLSeconds) + .exec(); } } - private async appendEvent(input: { - event: T; - checkOwner: boolean; - executorId: string; + private async appendAfterTail(input: { + tailEntryId: string | undefined; + events: T[]; streamTTLSeconds: number | undefined; - }): Promise<'appended' | 'duplicate'> { - const payload = JSON.stringify(input.event); - const reply = await this.appendEntry({ - streamId: this.streamIdFor(input.event), - payload, - checkOwner: input.checkOwner, - executorId: input.executorId, - streamTTLSeconds: input.streamTTLSeconds, + }): Promise<'ok' | 'stale'> { + const pairs = input.events.flatMap(event => [this.entryMillisecondFor(event), JSON.stringify(event)]); + const reply: unknown = await this.redis.eval(APPEND_AFTER_TAIL, { + keys: [this.streamId], + arguments: [input.tailEntryId ?? '', ttlArgument(input.streamTTLSeconds), ...pairs], }); - if (reply.status === 'ok') { - return 'appended'; - } - if (reply.status === 'duplicate') { - return 'duplicate'; - } - if (reply.status === 'owner_mismatch') { - throw new StreamOwnerMismatchError(this.streamId); - } - if (reply.status === 'error') { - throw new Error(`Redis stream append failed on ${this.streamId}: ${reply.detail}`); - } - // The tail id can match a delta of the same message. Identical payloads - // already returned `duplicate` above; a distinct payload must be retried - // at the tail millisecond. - const lastId = await this.lastEntryId(); - if (lastId === undefined) { - throw new Error(`Redis stream append conflicted on an empty stream ${this.streamId}`); + if (reply === 'ok' || reply === 'stale') { + return reply; } - let retry: { status: AppendStatus; detail: string }; - try { - retry = await this.appendEntry({ - streamId: `${String(entryIdMilliseconds(lastId))}-*`, - payload, - checkOwner: input.checkOwner, - executorId: input.executorId, - streamTTLSeconds: input.streamTTLSeconds, - }); - } catch (error) { - throw new Error(`Failed to append event on ${this.streamId}`, { cause: error }); - } - if (retry.status === 'ok') { - return 'appended'; - } - if (retry.status === 'duplicate') { - return 'duplicate'; - } - if (retry.status === 'owner_mismatch') { - throw new StreamOwnerMismatchError(this.streamId); - } - throw new Error(`Redis stream append failed on ${this.streamId}: ${retry.detail}`, { cause: reply.detail }); - } - - private streamIdFor(event: T): string { - return `${String(eventIdTimestampMs(event.id) ?? Date.now())}-*`; + throw new Error(`Unexpected stream repair reply on ${this.streamId}`); } - private async appendEntry(input: { - streamId: string; - payload: string; - checkOwner: boolean; - executorId: string; - streamTTLSeconds: number | undefined; - }): Promise<{ status: AppendStatus; detail: string }> { - const ttl = input.streamTTLSeconds !== undefined && input.streamTTLSeconds > 0 ? input.streamTTLSeconds : 0; - const reply: unknown = await this.redis.eval(APPEND_STREAM_ENTRY, { - keys: [streamOwnerKey(this.streamId), this.streamId], - arguments: [input.checkOwner ? '1' : '0', input.executorId, input.streamId, input.payload, String(ttl)], - }); - return parseAppendReply(reply); + private entryMillisecondFor(event: T): string { + return String(eventIdTimestampMs(event.id) ?? Date.now()); } } diff --git a/packages/trueforge/src/runtime/event-subscription/repair.ts b/packages/trueforge/src/runtime/event-subscription/repair.ts index 2da4f7e8dd..3105e69a54 100644 --- a/packages/trueforge/src/runtime/event-subscription/repair.ts +++ b/packages/trueforge/src/runtime/event-subscription/repair.ts @@ -20,12 +20,16 @@ export interface RepairLogOptions { executorId?: string | undefined; } -interface StreamTail { +export interface StreamTail { + entryId: string; eventId: string; /** JSON written on the stream, used to tell shared-id rows apart. */ payload: string; } +/** Backfill gives up after this many concurrent tail moves. */ +const REPAIR_MAX_ATTEMPTS = 5; + function isPlainObject(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value); } @@ -111,11 +115,16 @@ function suffixAfterTail( return entries.slice(match + 1); } +/** + * Appends the missing suffix only while the stream tail is still the one it was + * computed from. A concurrent writer (live producer or another repair) moves the + * tail; the suffix is then recomputed, so stale rows never land after newer ones. + */ export async function applyEventLogRepair(input: { entries: readonly EventLogEntry[]; claimOwner?: (() => Promise) | undefined; readTail: () => Promise; - append: (entry: EventLogEntry) => Promise<'appended' | 'duplicate'>; + appendAfterTail: (batch: { tailEntryId: string | undefined; events: T[] }) => Promise<'ok' | 'stale'>; refreshTtl: () => Promise; /** Producer resume refreshes TTL even when the existing stream is already complete. */ refreshTtlOnUntouchedStream: boolean; @@ -123,16 +132,22 @@ export async function applyEventLogRepair(input: { if (input.claimOwner !== undefined) { await input.claimOwner(); } - const tail = await input.readTail(); - const suffix = suffixAfterTail(input.entries, tail); - let appended = false; - for (const entry of suffix) { - const outcome = await input.append(entry); - if (outcome === 'appended') { - appended = true; + for (let attempt = 0; attempt < REPAIR_MAX_ATTEMPTS; attempt += 1) { + const tail = await input.readTail(); + const suffix = suffixAfterTail(input.entries, tail); + if (suffix.length === 0) { + if (input.refreshTtlOnUntouchedStream) { + await input.refreshTtl(); + } + return; + } + const outcome = await input.appendAfterTail({ + tailEntryId: tail?.entryId, + events: suffix.map(entry => entry.event), + }); + if (outcome === 'ok') { + return; } } - if (appended || input.refreshTtlOnUntouchedStream) { - await input.refreshTtl(); - } + throw new Error(`Stream repair gave up after ${String(REPAIR_MAX_ATTEMPTS)} concurrent tail changes`); } diff --git a/packages/trueforge/src/runtime/turnEventLogRepair.ts b/packages/trueforge/src/runtime/turnEventLogRepair.ts index baf0294b72..753b3af156 100644 --- a/packages/trueforge/src/runtime/turnEventLogRepair.ts +++ b/packages/trueforge/src/runtime/turnEventLogRepair.ts @@ -1,41 +1,26 @@ import type { ISessionStore, TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session'; -import type { EventLogEntry, EventSubscription } from './event-subscription'; +import type { EventSubscription } from './event-subscription'; /** * Rebuild a paused turn's live stream from the durable log. * Subscribe omits `producerExecutorId` and only appends. - * The producing resume claims the stream-owner key. Rows are appended with the - * same event-id stream ids as a live put. + * The producing resume, which already owns the turn row, claims the stream-owner key. */ export async function repairTurnEventLog(input: { - sessionStore: Pick, 'listTurnEventLog' | 'assignTurnEventSequenceNumbers'>; + sessionStore: Pick, 'listTurnEventLog'>; subscription: EventSubscription; sessionId: string; turnId: string; streamTTLSeconds: number; producerExecutorId?: string | undefined; }): Promise { - const rows = await input.sessionStore.listTurnEventLog({ + const events = await input.sessionStore.listTurnEventLog({ session_id: input.sessionId, turn_id: input.turnId, }); - const entries: EventLogEntry[] = rows.map(row => ({ - id: row.event.id, - event: row.event, - })); - const producerExecutorId = input.producerExecutorId; - if (producerExecutorId !== undefined) { - // Empty assignment still takes the turn-owner fence. It does not write sequence numbers. - await input.sessionStore.assignTurnEventSequenceNumbers({ - session_id: input.sessionId, - turn_id: input.turnId, - active_executor_id: producerExecutorId, - assignments: [], - }); - } await input.subscription.repairLog({ - entries, + entries: events.map(event => ({ id: event.id, event })), streamTTLSeconds: input.streamTTLSeconds, - ...(producerExecutorId === undefined ? {} : { executorId: producerExecutorId }), + ...(input.producerExecutorId === undefined ? {} : { executorId: input.producerExecutorId }), }); } diff --git a/packages/trueforge/src/schemas/turn.ts b/packages/trueforge/src/schemas/turn.ts index cf22f2a73e..382945d146 100644 --- a/packages/trueforge/src/schemas/turn.ts +++ b/packages/trueforge/src/schemas/turn.ts @@ -17,7 +17,9 @@ export const SubscribeTurnQuerySchema = z .int() .nonnegative() .optional() - .describe('Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.'), + .describe( + 'Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.', + ), ), from_event_id: z .string() diff --git a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts index 24d6bc1af5..d6c1301d53 100644 --- a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts +++ b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts @@ -51,9 +51,15 @@ async function collect( } /** EventSubscription contract suite — factory-injected so both backends run it. */ -export function runEventSubscriptionContractSuite(createRegistry: () => EventSubscriptionRegistry) { +export function runEventSubscriptionContractSuite( + setup: () => { + registry: EventSubscriptionRegistry; + /** Drops the log and keeps the owner, the way a TTL expiry does. */ + evictStream: (streamId: string) => Promise; + }, +) { it('put appends events in order without a sequence number', async () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); @@ -67,7 +73,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub }); it('poll from the start yields every event sent, in order, with identical payloads', async () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); @@ -81,7 +87,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub }); it('poll with afterSequenceNumber 0 matches an omitted cursor (from the start)', async () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); @@ -96,7 +102,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub }); it('poll from a mid-stream cursor yields only events strictly after it', async () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); @@ -110,7 +116,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub }); it('keeps an event appended while an event-id replay is delivered', async () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); const sent = [0, 1, 2].map(index => makeEvent(index)); @@ -138,7 +144,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub it('appends a distinct shared-id payload when its timestamp is behind the tail', async () => { const { ulid } = await import('ulid'); - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); const later = makeEvent(0, ulid(2_000_000_000_000).toLowerCase()); @@ -153,7 +159,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub }); it('repairs a suffix when the durable row key order differs from the stream payload', async () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); const stored = makeEvent(0); @@ -177,7 +183,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub }); it('poll parked at the tip receives events put after it subscribed', async () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); @@ -195,14 +201,14 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub }); it('rebuilds a missing stream from the event log and continues after the tail', async () => { - const registry = createRegistry(); + const { registry, evictStream } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); const first = makeEvent(1); const second = makeEvent(2); await producer.put(first, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); - await registry.evictStream(streamId); + await evictStream(streamId); await producer.repairLog({ entries: [ @@ -220,7 +226,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub }); it('leaves an existing stream untouched and a second repair does not append again', async () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); const first = makeEvent(1); @@ -246,7 +252,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub }); it('rejects a stale producer when the stream is missing', async () => { - const registry = createRegistry(); + const { registry, evictStream } = setup(); const streamId = nextStreamId(); const owner = registry.get(streamId); const event = makeEvent(1); @@ -255,7 +261,7 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner', }); - await registry.evictStream(streamId); + await evictStream(streamId); const stale = registry.get(streamId); await expect( @@ -264,20 +270,26 @@ export function runEventSubscriptionContractSuite(createRegistry: () => EventSub await expect(stale.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); }); - it('treats an explicit-id conflict as a no-op when the event id is already the tail', async () => { - const registry = createRegistry(); + it('keeps consecutive identical payloads, like two deltas with the same text', async () => { + const { registry } = setup(); const streamId = nextStreamId(); - const left = registry.get(streamId); - const right = registry.get(streamId); - const event = makeEvent(0, 'evt-same'); - await Promise.all([ - left.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }), - right.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }), - ]); - - const follow = makeEvent(1, 'evt-next'); - await left.put(follow, { streamTTLSeconds: STREAM_TTL_SECONDS }); + const producer = registry.get(streamId); + const delta = makeEvent(0, 'evt-message'); + await producer.put(delta, { streamTTLSeconds: STREAM_TTL_SECONDS }); + await producer.put(delta, { streamTTLSeconds: STREAM_TTL_SECONDS }); + const received = await collect(registry, streamId, 2); - expect(received).toEqual([event, follow]); + expect(received).toEqual([delta, delta]); + }); + + it('reports whether an event id is on the live stream', async () => { + const { registry } = setup(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + await producer.put(makeEvent(0), { streamTTLSeconds: STREAM_TTL_SECONDS }); + await producer.put(makeEvent(1), { streamTTLSeconds: STREAM_TTL_SECONDS }); + + await expect(producer.hasEventId('evt-1')).resolves.toBe(true); + await expect(producer.hasEventId('evt-missing')).resolves.toBe(false); }); } diff --git a/packages/trueforge/tests/runtime/event-subscription/inMemory.contract.test.ts b/packages/trueforge/tests/runtime/event-subscription/inMemory.contract.test.ts index df315befd6..3a475f6da1 100644 --- a/packages/trueforge/tests/runtime/event-subscription/inMemory.contract.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/inMemory.contract.test.ts @@ -1,6 +1,16 @@ import { EventSubscriptionRegistry } from '../../../src/runtime/event-subscription'; +import { InMemoryEventStreamStore } from '../../../src/runtime/event-subscription/inMemory'; import { runEventSubscriptionContractSuite, type TestEvent } from './contractSuite'; describe('InMemoryEventSubscription (EventSubscription contract)', () => { - runEventSubscriptionContractSuite(() => new EventSubscriptionRegistry(undefined)); + runEventSubscriptionContractSuite(() => { + const store = new InMemoryEventStreamStore(); + return { + registry: new EventSubscriptionRegistry(undefined, store), + evictStream: streamId => { + store.evictStream(streamId); + return Promise.resolve(); + }, + }; + }); }); diff --git a/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts b/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts index 43b246d638..63df7d180c 100644 --- a/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts @@ -22,5 +22,10 @@ describeRedis('RedisEventSubscription (EventSubscription contract)', () => { await client.close(); }); - runEventSubscriptionContractSuite(() => new EventSubscriptionRegistry(client)); + runEventSubscriptionContractSuite(() => ({ + registry: new EventSubscriptionRegistry(client), + evictStream: async streamId => { + await client.del(streamId); + }, + })); }); diff --git a/packages/trueforge/tests/runtime/event-subscription/repair.test.ts b/packages/trueforge/tests/runtime/event-subscription/repair.test.ts new file mode 100644 index 0000000000..6809a21799 --- /dev/null +++ b/packages/trueforge/tests/runtime/event-subscription/repair.test.ts @@ -0,0 +1,49 @@ +import { applyEventLogRepair, type StreamTail } from '../../../src/runtime/event-subscription/repair'; + +interface Row { + id: string; + type: string; +} + +function tailOf(row: Row, entryId: string): StreamTail { + return { entryId, eventId: row.id, payload: JSON.stringify(row) }; +} + +describe('applyEventLogRepair', () => { + it('recomputes the suffix when a concurrent writer moves the tail', async () => { + const created: Row = { id: 'evt-1', type: 'turn.created' }; + const paused: Row = { id: 'evt-2', type: 'turn.update' }; + const running: Row = { id: 'evt-3', type: 'turn.update' }; + const entries = [created, paused, running].map(event => ({ id: event.id, event })); + // The first read sees only `created`; by the append the producer has written through `running`. + const tails = [tailOf(created, '1-0'), tailOf(running, '3-0')]; + const batches: { tailEntryId: string | undefined; events: Row[] }[] = []; + + await applyEventLogRepair({ + entries, + readTail: () => Promise.resolve(tails.shift()), + appendAfterTail: batch => { + batches.push(batch); + return Promise.resolve('stale'); + }, + refreshTtl: () => Promise.resolve(), + refreshTtlOnUntouchedStream: false, + }); + + expect(batches).toEqual([{ tailEntryId: '1-0', events: [paused, running] }]); + }); + + it('gives up when the tail keeps moving', async () => { + const created: Row = { id: 'evt-1', type: 'turn.created' }; + const paused: Row = { id: 'evt-2', type: 'turn.update' }; + await expect( + applyEventLogRepair({ + entries: [created, paused].map(event => ({ id: event.id, event })), + readTail: () => Promise.resolve(tailOf(created, '1-0')), + appendAfterTail: () => Promise.resolve('stale'), + refreshTtl: () => Promise.resolve(), + refreshTtlOnUntouchedStream: false, + }), + ).rejects.toThrow('Stream repair gave up'); + }); +}); diff --git a/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts index 47159a4983..80136aaf7f 100644 --- a/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts @@ -1,6 +1,5 @@ import type { TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session/TurnHandle'; import { InMemorySessionStore } from '@truefoundry/trueforge-core/agent-session/store/InMemorySessionStore'; -import { TurnOwnershipLostError } from '@truefoundry/trueforge-core/agent-session/store/SessionStoreErrors'; import { makeAgentSpec, makeCreateTurnInput, @@ -14,6 +13,7 @@ import { StreamOwnerMismatchError, type SequencedEvent, } from '../../../src/runtime/event-subscription'; +import { InMemoryEventStreamStore } from '../../../src/runtime/event-subscription/inMemory'; import { repairTurnEventLog } from '../../../src/runtime/turnEventLogRepair'; const STREAM_TTL_SECONDS = 300; @@ -41,7 +41,7 @@ async function collect( } describe('repairTurnEventLog', () => { - it('rebuilds a missing stream from the durable log without assigning sequence numbers', async () => { + it('rebuilds a missing stream from the durable log', async () => { const store = new InMemorySessionStore(); const sessionId = 's1'; const turnId = 'turn-1'; @@ -65,7 +65,8 @@ describe('repairTurnEventLog', () => { events: [created, message], }); - const registry = new EventSubscriptionRegistry(undefined); + const streams = new InMemoryEventStreamStore(); + const registry = new EventSubscriptionRegistry(undefined, streams); const streamId = 'turn-stream'; const subscription = registry.get(streamId); const repair = (producerExecutorId?: string) => @@ -78,22 +79,15 @@ describe('repairTurnEventLog', () => { ...(producerExecutorId === undefined ? {} : { producerExecutorId }), }); - await expect(repair('other-executor')).rejects.toBeInstanceOf(TurnOwnershipLostError); await expect(subscription.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); await repair(); - expect( - (await store.listTurnEventLog({ session_id: sessionId, turn_id: turnId })).map(row => row.sequence_number), - ).toEqual([null, null]); expect((await collect(registry, streamId, 2)).map(event => event.id)).toEqual([created.id, message.id]); await repair(); expect((await collect(registry, streamId, 2)).map(event => event.id)).toEqual([created.id, message.id]); await repair(TEST_ACTIVE_EXECUTOR_ID); - expect( - (await store.listTurnEventLog({ session_id: sessionId, turn_id: turnId })).map(row => row.sequence_number), - ).toEqual([null, null]); const later = makeModelMessageEvent(); await store.appendToEvents({ @@ -104,11 +98,8 @@ describe('repairTurnEventLog', () => { }); await repair(TEST_ACTIVE_EXECUTOR_ID); expect((await collect(registry, streamId, 3)).map(event => event.id)).toEqual([created.id, message.id, later.id]); - expect( - (await store.listTurnEventLog({ session_id: sessionId, turn_id: turnId })).map(row => row.sequence_number), - ).toEqual([null, null, null]); - await registry.evictStream(streamId); + streams.evictStream(streamId); const stale = registry.get(streamId); await expect( stale.put(makeModelMessageEvent(), { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'other-executor' }), diff --git a/packages/trueforge/tests/unit/apis/turns.test.ts b/packages/trueforge/tests/unit/apis/turns.test.ts index 5abee524c0..2e472dc0a0 100644 --- a/packages/trueforge/tests/unit/apis/turns.test.ts +++ b/packages/trueforge/tests/unit/apis/turns.test.ts @@ -14,6 +14,7 @@ import { createEmptyAgentThreadMetrics } from '../../../../trueforge-core/src/co import { makeCancelledTurnState, makeCreateTurnInput, + makeModelMessageEvent, makePausedTurnState, makeTurnCreatedEvent, makeTurnDoneEvent, @@ -923,13 +924,14 @@ describe('turns', () => { }); const activeTurns = new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }); const tokenStore = new SqliteOAuthTokenStore(db); + const eventSubscriptions = new EventSubscriptionRegistry(undefined); const app = new OpenAPIHono(); app.route( '/', createTurnsRouter({ activeTurns, authorizer: new TrueForgeAuthorizer(), - eventSubscriptions: new EventSubscriptionRegistry(undefined), + eventSubscriptions, logger: createLogger({ silent: true }), requestReplyRouter: new RequestReplyRouter(), resolveAgentStore: () => new SqliteAgentStore(db), @@ -943,6 +945,52 @@ describe('turns', () => { }), ); + // An in-progress model.message shell is stream-only; resuming from it must still work. + const streamingTurnId = 'streaming-turn'; + await sessionStore.createTurn(makeCreateTurnInput({ sessionId, turnId: streamingTurnId })); + const shell = makeModelMessageEvent(); + await eventSubscriptions + .get(turnStreamId('default', sessionId, streamingTurnId)) + .put(shell, { streamTTLSeconds: 300 }); + const midMessage = await app.request( + `/${sessionId}/turns/${streamingTurnId}/subscribe?from_event_id=${shell.id}`, + ); + expect(midMessage.status).toBe(200); + await midMessage.body?.cancel(); + + // An older client counted id-less events from 0. after=1 has indexes 0 and 1. + const countedTurnId = 'counted-turn'; + await sessionStore.createTurn(makeCreateTurnInput({ sessionId, turnId: countedTurnId })); + const firstCounted = makeModelMessageEvent(); + const kept = makeModelMessageEvent(); + const counted = [firstCounted, makeModelMessageEvent(), kept]; + const countedStream = eventSubscriptions.get(turnStreamId('default', sessionId, countedTurnId)); + for (const event of counted) { + await countedStream.put(event, { streamTTLSeconds: 300 }); + } + const countedResume = await app.request(`/${sessionId}/turns/${countedTurnId}/subscribe?after_sequence_number=1`); + expect(countedResume.status).toBe(200); + const countedReader = countedResume.body?.getReader(); + if (countedReader === undefined) { + throw new Error('subscribe response has no body'); + } + const countedDecoder = new TextDecoder(); + let countedBody = ''; + try { + while (!countedBody.includes(kept.id)) { + const { value, done } = await countedReader.read(); + if (done || value === undefined) { + break; + } + countedBody += countedDecoder.decode(value, { stream: true }); + } + } finally { + await countedReader.cancel(); + } + expect(countedBody).toContain(kept.id); + expect(countedBody).not.toContain(firstCounted.id); + expect(countedBody).toContain('id: 2'); + const runningTurnId = 'running-turn'; await sessionStore.createTurn(makeCreateTurnInput({ sessionId, turnId: runningTurnId })); const runningEvent = makeTurnCreatedEvent(runningTurnId); @@ -1015,8 +1063,9 @@ describe('turns', () => { ); expect(mixed.status).toBe(400); - const huge = await app.request(`/${sessionId}/turns/missing-turn/subscribe?after_sequence_number=100001`); - expect(huge.status).toBe(400); + const huge = await app.request(`/${sessionId}/turns/${pausedTurnId}/subscribe?after_sequence_number=100001`); + expect(huge.status).toBe(200); + await huge.body?.cancel(); const invalid = await app.request(`/${sessionId}/turns/${pausedTurnId}/subscribe?from_event_id=not-a-ulid`); expect(invalid.status).toBe(400); @@ -1026,11 +1075,6 @@ describe('turns', () => { `/${sessionId}/turns/${pausedTurnId}/subscribe?from_event_id=${ulid().toLowerCase()}`, ); expect(foreign.status).toBe(400); - - const sequenceOnEventIdStream = await app.request( - `/${sessionId}/turns/${pausedTurnId}/subscribe?after_sequence_number=1`, - ); - expect(sequenceOnEventIdStream.status).toBe(400); }); }); }); From 401c7c3f8b0991c733fc9ddd3537649d4bef433e Mon Sep 17 00:00:00 2001 From: "trueforge-dev-bot[bot]" Date: Tue, 3 Nov 2026 11:50:47 +0000 Subject: [PATCH 08/21] Regenerate OpenAPI document and SDKs --- .github/fern/openapi/openapi.json | 4 ++-- docs/openapi.json | 4 ++-- .../client/requests/SubscribeToTurnSessionsRequest.ts | 2 +- python/trueforge_sdk/reference.md | 2 +- python/trueforge_sdk/src/trueforge_sdk/sessions/client.py | 4 ++-- python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py | 4 ++-- 6 files changed, 10 insertions(+), 10 deletions(-) diff --git a/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index c2071411ca..35f29f6c8e 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -10751,12 +10751,12 @@ } }, { - "description": "Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.", + "description": "Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.", "in": "query", "name": "after_sequence_number", "required": false, "schema": { - "description": "Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.", + "description": "Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.", "minimum": 0, "type": [ "integer", diff --git a/docs/openapi.json b/docs/openapi.json index c2071411ca..35f29f6c8e 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -10751,12 +10751,12 @@ } }, { - "description": "Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.", + "description": "Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.", "in": "query", "name": "after_sequence_number", "required": false, "schema": { - "description": "Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id.", + "description": "Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.", "minimum": 0, "type": [ "integer", diff --git a/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts b/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts index 4aa5bc30e7..9fe7a07bc8 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts @@ -5,7 +5,7 @@ * {} */ export interface SubscribeToTurnSessionsRequest { - /** Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. */ + /** Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. */ afterSequenceNumber?: number | null; /** Inclusive cursor. Replay starts at the entry with this event id */ fromEventId?: string; diff --git a/python/trueforge_sdk/reference.md b/python/trueforge_sdk/reference.md index 256a69e2d0..a3f130867a 100644 --- a/python/trueforge_sdk/reference.md +++ b/python/trueforge_sdk/reference.md @@ -3519,7 +3519,7 @@ client.sessions.subscribe_to_turn(
-**after_sequence_number:** `typing.Optional[int]` — Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. +**after_sequence_number:** `typing.Optional[int]` — Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.
diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py index 8198004487..484a94c3bf 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py @@ -850,7 +850,7 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. + Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id @@ -1847,7 +1847,7 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. + Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py index dbe1c8ba9b..5235647b8f 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py @@ -1763,7 +1763,7 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. + Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id @@ -3606,7 +3606,7 @@ async def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor. Replay events after this sequence. Omit when using an event id. + Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id From bdeceee970ebd499b3ac6000229993f4237c9261 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 3 Nov 2026 17:42:22 +0530 Subject: [PATCH 09/21] Repair turn streams by paging listTurnEvents ascending Drop listTurnEventLog. Repair reads durable rows one page at a time in persist order and appends each page only while the stream tail is the entry id the previous page wrote. Co-authored-by: Cursor --- .../trueforge-core/src/agent-session/index.ts | 1 - .../src/agent-session/store/ISessionStore.ts | 11 --- .../store/InMemorySessionStore.ts | 10 --- .../agent-session/store/storeContractSuite.ts | 22 ----- packages/trueforge/src/apis/schedules.ts | 2 +- packages/trueforge/src/apis/turns.ts | 2 +- .../session-store/PostgresSessionStore.ts | 6 -- .../postgres/session-store/queries/events.ts | 26 ------ .../session-store/SqliteSessionStore.ts | 6 -- .../db/sqlite/session-store/queries/events.ts | 26 ------ .../runtime/event-subscription/inMemory.ts | 14 +-- .../src/runtime/event-subscription/redis.ts | 46 +++++++--- .../src/runtime/event-subscription/repair.ts | 90 +++++++++++++++---- .../src/runtime/turnEventLogRepair.ts | 34 +++++-- .../event-subscription/contractSuite.ts | 15 ++-- .../event-subscription/entryId.test.ts | 6 +- .../runtime/event-subscription/repair.test.ts | 47 ++++++++-- 17 files changed, 201 insertions(+), 163 deletions(-) diff --git a/packages/trueforge-core/src/agent-session/index.ts b/packages/trueforge-core/src/agent-session/index.ts index 924b702509..6fda6fc6d4 100644 --- a/packages/trueforge-core/src/agent-session/index.ts +++ b/packages/trueforge-core/src/agent-session/index.ts @@ -99,7 +99,6 @@ export type { InsertTurnInboundEventsInput, ListSessionEventsInput, ListSessionsInput, - ListTurnEventLogInput, ListTurnEventsInput, ListTurnsInput, ListUnconsumedTurnInboundEventsInput, diff --git a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts index a9a1a8794f..d0d63b8c60 100644 --- a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts @@ -193,11 +193,6 @@ export interface AppendToEventsInput extends TurnProgressKeys { events: PersistedTurnEvent[]; } -export interface ListTurnEventLogInput { - session_id: string; - turn_id: string; -} - export interface InsertTurnInboundEventsInput extends TurnProgressKeys { /** Empty array is a no-op. */ events: TurnUserEvent[]; @@ -417,12 +412,6 @@ export interface ISessionStore< /** Appends events while the caller owns a `running` or `paused` turn; terminal turns are immutable. */ appendToEvents(input: AppendToEventsInput, transaction?: TTransaction): Promise; - /** - * Full turn event log in persist order (`event_id` ascending). Missing turn - * throws TurnNotFoundError. - */ - listTurnEventLog(input: ListTurnEventLogInput): Promise; - /** * Durable inbound send-event inbox for a tip the caller still owns. Tip must be * non-terminal (`running` or `paused`) — terminal tip → diff --git a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts index ff3034a46e..1cdd4a0339 100644 --- a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts @@ -32,7 +32,6 @@ import type { ISessionStore, ListSessionEventsInput, ListSessionsInput, - ListTurnEventLogInput, ListTurnEventsInput, ListTurnsInput, ListUnconsumedTurnInboundEventsInput, @@ -632,15 +631,6 @@ export class InMemorySessionStore< return; } - async listTurnEventLog(input: ListTurnEventLogInput): Promise { - this.requireTurn(input.session_id, input.turn_id); - const list = this.events.get(turnKey(input)); - if (!list) { - throw new TurnNotFoundError(input.turn_id); - } - return deepCopy([...list].sort((left, right) => (left.id < right.id ? -1 : left.id > right.id ? 1 : 0))); - } - async patchThreadsMCPAuth(input: PatchThreadsMCPAuthInput): Promise { this.requireOwnedTurn(input); const turn = this.requireNonTerminalTurn(input); diff --git a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts index ac5fc9944e..b9f1f3b114 100644 --- a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts +++ b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts @@ -4595,28 +4595,6 @@ export function runStoreContractSuite( }); } - describe('listTurnEventLog', () => { - it('returns every event of the turn in event_id order and rejects a missing turn', async () => { - const store = createStore(); - await seedSession(store); - await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); - const created = makeTurnCreatedEvent('turn-1'); - const message = makeModelMessageEvent(); - await store.appendToEvents({ - active_executor_id: TEST_ACTIVE_EXECUTOR_ID, - session_id: sessionId, - turn_id: 'turn-1', - events: [created, message], - }); - - const log = await store.listTurnEventLog({ session_id: sessionId, turn_id: 'turn-1' }); - expect(log).toEqual([created, message].sort((left, right) => (left.id < right.id ? -1 : 1))); - await expect(store.listTurnEventLog({ session_id: sessionId, turn_id: 'missing' })).rejects.toBeInstanceOf( - TurnNotFoundError, - ); - }); - }); - describe('no SSE surface', () => { it('ISessionStore has no subscribe* members', () => { const store = createStore(); diff --git a/packages/trueforge/src/apis/schedules.ts b/packages/trueforge/src/apis/schedules.ts index b7ea71a6a8..7d1376f331 100644 --- a/packages/trueforge/src/apis/schedules.ts +++ b/packages/trueforge/src/apis/schedules.ts @@ -88,7 +88,7 @@ export interface ScheduleTurnExecutionDeps { scheduleStore: IScheduleStore; sessions: Sessions; /** Previous-turn cancel, and rebuilding a missing event stream. */ - sessionStore: Pick, 'getTurn' | 'listTurnEventLog'>; + sessionStore: Pick, 'getTurn' | 'listTurnEvents'>; webSearchProviderStore: IWebSearchProviderStore; } diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 5ef2590db2..e891b18b8f 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -234,7 +234,7 @@ export type BeginTurnExecutionDeps = Pick< webSearchProviderStore: IWebSearchProviderStore; skillStore: Pick; /** Previous-turn cancel, and rebuilding a missing event stream. */ - sessionStore: Pick; + sessionStore: Pick; redis?: RedisClient | undefined; }; diff --git a/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts b/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts index 58e7b85e42..dc08267e73 100644 --- a/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts +++ b/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts @@ -20,7 +20,6 @@ import type { ISessionStore, ListSessionEventsInput, ListSessionsInput, - ListTurnEventLogInput, ListTurnEventsInput, ListTurnsInput, ListUnconsumedTurnInboundEventsInput, @@ -46,7 +45,6 @@ import { patchThreadCapabilityState as patchThreadCapabilityStateQuery } from '. import { appendToEvents as appendToEventsQuery, listSessionEvents as listSessionEventsQuery, - listTurnEventLog as listTurnEventLogQuery, listTurnEvents as listTurnEventsQuery, } from './queries/events'; import { @@ -224,10 +222,6 @@ export class PostgresSessionStore implements ISessionStore return appendToEventsQuery(transaction ?? this.db, input); } - listTurnEventLog(input: ListTurnEventLogInput): Promise { - return listTurnEventLogQuery(this.db, input); - } - insertTurnInboundEvents(input: InsertTurnInboundEventsInput): Promise { return insertTurnInboundEventsQuery(this.db, input); } diff --git a/packages/trueforge/src/db/postgres/session-store/queries/events.ts b/packages/trueforge/src/db/postgres/session-store/queries/events.ts index debfb0d478..8be36bdb75 100644 --- a/packages/trueforge/src/db/postgres/session-store/queries/events.ts +++ b/packages/trueforge/src/db/postgres/session-store/queries/events.ts @@ -3,7 +3,6 @@ import type { TokenPagination } from '@truefoundry/trueforge-core/agent-session/ import type { AppendToEventsInput, ListSessionEventsInput, - ListTurnEventLogInput, ListTurnEventsInput, } from '@truefoundry/trueforge-core/agent-session/store/ISessionStore'; import { @@ -71,31 +70,6 @@ export async function appendToEvents( } } -export async function listTurnEventLog( - db: Kysely, - input: ListTurnEventLogInput, -): Promise { - const turn = await db - .selectFrom('turn') - .select('turn_id') - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .executeTakeFirst(); - if (!turn) { - throw new TurnNotFoundError(input.turn_id); - } - - const rows = await db - .selectFrom('session_event') - .select('event') - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .orderBy('event_id', 'asc') - .execute(); - - return rows.map(row => row.event); -} - /** * listTurnEvents — one statement: drive from turn, left-join a page of events. * Missing turn → 0 rows → TurnNotFoundError. Empty log → one null-event sentinel. diff --git a/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts b/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts index fedd48c751..b6798bc359 100644 --- a/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts +++ b/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts @@ -20,7 +20,6 @@ import type { ISessionStore, ListSessionEventsInput, ListSessionsInput, - ListTurnEventLogInput, ListTurnEventsInput, ListTurnsInput, ListUnconsumedTurnInboundEventsInput, @@ -46,7 +45,6 @@ import { patchThreadCapabilityState as patchThreadCapabilityStateQuery } from '. import { appendToEvents as appendToEventsQuery, listSessionEvents as listSessionEventsQuery, - listTurnEventLog as listTurnEventLogQuery, listTurnEvents as listTurnEventsQuery, } from './queries/events'; import { @@ -222,10 +220,6 @@ export class SqliteSessionStore implements ISessionStore, return appendToEventsQuery(transaction ?? this.db, input); } - listTurnEventLog(input: ListTurnEventLogInput): Promise { - return listTurnEventLogQuery(this.db, input); - } - insertTurnInboundEvents(input: InsertTurnInboundEventsInput): Promise { return insertTurnInboundEventsQuery(this.db, input); } diff --git a/packages/trueforge/src/db/sqlite/session-store/queries/events.ts b/packages/trueforge/src/db/sqlite/session-store/queries/events.ts index 7b924edf21..d4a28ff9a8 100644 --- a/packages/trueforge/src/db/sqlite/session-store/queries/events.ts +++ b/packages/trueforge/src/db/sqlite/session-store/queries/events.ts @@ -3,7 +3,6 @@ import type { TokenPagination } from '@truefoundry/trueforge-core/agent-session/ import type { AppendToEventsInput, ListSessionEventsInput, - ListTurnEventLogInput, ListTurnEventsInput, } from '@truefoundry/trueforge-core/agent-session/store/ISessionStore'; import { @@ -70,31 +69,6 @@ async function appendToEventsInTransaction(db: Transaction, input: App await db.insertInto('session_event').values(eventRows).execute(); } -export async function listTurnEventLog( - db: Kysely, - input: ListTurnEventLogInput, -): Promise { - const turn = await db - .selectFrom('turn') - .select('turn_id') - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .executeTakeFirst(); - if (!turn) { - throw new TurnNotFoundError(input.turn_id); - } - - const rows = await db - .selectFrom('session_event') - .select(jsonText(sql.ref('event')).as('event')) - .where('session_id', '=', input.session_id) - .where('turn_id', '=', input.turn_id) - .orderBy('event_id', 'asc') - .execute(); - - return rows.map(row => row.event); -} - /** * listTurnEvents — one statement: drive from turn, left-join a page of events. * Missing turn → 0 rows → TurnNotFoundError. Empty log → one null-event sentinel. diff --git a/packages/trueforge/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index d97da3f354..7652743575 100644 --- a/packages/trueforge/src/runtime/event-subscription/inMemory.ts +++ b/packages/trueforge/src/runtime/event-subscription/inMemory.ts @@ -20,7 +20,7 @@ import { isEntryAfterCursor, isLegacySequenceEntryId, } from './entryId'; -import { applyEventLogRepair, type RepairLogOptions, type StreamTail } from './repair'; +import { applyEventLogRepair, type AppendAfterTailResult, type RepairLogOptions, type StreamTail } from './repair'; function ttlDeadline(streamTTLSeconds: number | undefined): number | undefined { return streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? Date.now() + streamTTLSeconds * 1_000 : undefined; @@ -90,10 +90,10 @@ export class InMemoryEventStreamStore { tailEntryId: string | undefined; events: T[]; streamTTLSeconds: number | undefined; - }): 'ok' | 'stale' { + }): AppendAfterTailResult { const stream = this.getLiveStream(input.streamId) ?? { events: [] }; if (stream.events[stream.events.length - 1]?.entryId !== input.tailEntryId) { - return 'stale'; + return { status: 'stale' }; } this.streams.set(input.streamId, stream); for (const event of input.events) { @@ -101,7 +101,11 @@ export class InMemoryEventStreamStore { } this.refreshStreamTtl(input.streamId, input.streamTTLSeconds); this.changes.emit(input.streamId); - return 'ok'; + const written = stream.events[stream.events.length - 1]; + if (written === undefined) { + throw new Error('Stream repair append wrote no entry'); + } + return { status: 'ok', tailEntryId: written.entryId }; } readTail(streamId: string): StreamTail | undefined { @@ -250,7 +254,7 @@ export class InMemoryEventSubscription implements Even repairLog(options: RepairLogOptions): Promise { const executorId = options.executorId; return applyEventLogRepair({ - entries: options.entries, + readPages: options.readPages, claimOwner: executorId === undefined ? undefined diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index 1f9290d4de..25e6c063c5 100644 --- a/packages/trueforge/src/runtime/event-subscription/redis.ts +++ b/packages/trueforge/src/runtime/event-subscription/redis.ts @@ -10,7 +10,13 @@ import { type SequencedEvent, } from '.'; import { entryIdMilliseconds, eventIdSeekCursor, eventIdTimestampMs, isLegacySequenceEntryId } from './entryId'; -import { applyEventLogRepair, eventIdFromStreamData, type RepairLogOptions, type StreamTail } from './repair'; +import { + applyEventLogRepair, + eventIdFromStreamData, + type AppendAfterTailResult, + type RepairLogOptions, + type StreamTail, +} from './repair'; const SUBSCRIBE_STREAM_POLL_ITEMS_COUNT = 100; /** Delay between XREAD polls while a stream has no new events (a blocking XREAD would monopolize the shared connection). */ @@ -113,8 +119,9 @@ return 'ok' * Repair backfill. KEYS[1]=stream ARGV[1]=expected tail id ('' = empty) ARGV[2]=ttl * ARGV[3..]=ms, payload pairs. * - * Appends only while the tail is still the one the suffix was computed from; - * otherwise returns 'stale' and the caller recomputes. + * Appends only while the tail is still the one the suffix was computed from. + * Returns {'ok', lastEntryId} or 'stale'. The caller uses lastEntryId as the + * expected tail of the next page. */ const APPEND_AFTER_TAIL = ` local tail = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) @@ -126,19 +133,23 @@ if tailId ~= ARGV[1] then return 'stale' end local tailMs = string.match(tailId, '^%d+') or '0' +local lastAdded = nil for i = 3, #ARGV, 2 do local ms = ARGV[i] if tonumber(tailMs) > tonumber(ms) then ms = tailMs end - local added = redis.call('XADD', KEYS[1], ms .. '-*', 'data', ARGV[i + 1]) - tailMs = string.match(added, '^%d+') + lastAdded = redis.call('XADD', KEYS[1], ms .. '-*', 'data', ARGV[i + 1]) + tailMs = string.match(lastAdded, '^%d+') end local ttl = tonumber(ARGV[2]) if ttl > 0 then redis.call('EXPIRE', KEYS[1], ttl) end -return 'ok' +if lastAdded == nil then + return 'stale' +end +return {'ok', lastAdded} `; function streamOwnerKey(streamId: string): string { @@ -149,6 +160,20 @@ function ttlArgument(streamTTLSeconds: number | undefined): string { return String(streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? streamTTLSeconds : 0); } +function parseAppendReply(value: unknown, streamId: string): AppendAfterTailResult { + if (value === 'stale') { + return { status: 'stale' }; + } + if (Array.isArray(value) && value.length === 2) { + const status: unknown = value[0]; + const tailEntryId: unknown = value[1]; + if (status === 'ok' && typeof tailEntryId === 'string') { + return { status: 'ok', tailEntryId }; + } + } + throw new Error(`Unexpected stream repair reply on ${streamId}`); +} + function parseSeekReply(value: unknown): { cursor: string; found: boolean } { if (Array.isArray(value) && value.length === 2) { const cursor: unknown = value[0]; @@ -173,7 +198,7 @@ export class RedisEventSubscription implements EventSu async repairLog(options: RepairLogOptions): Promise { const executorId = options.executorId; await applyEventLogRepair({ - entries: options.entries, + readPages: options.readPages, claimOwner: executorId === undefined ? undefined : () => this.claimOwner(executorId, options.streamTTLSeconds), readTail: () => this.readTail(), appendAfterTail: batch => this.appendAfterTail({ ...batch, streamTTLSeconds: options.streamTTLSeconds }), @@ -421,16 +446,13 @@ export class RedisEventSubscription implements EventSu tailEntryId: string | undefined; events: T[]; streamTTLSeconds: number | undefined; - }): Promise<'ok' | 'stale'> { + }): Promise { const pairs = input.events.flatMap(event => [this.entryMillisecondFor(event), JSON.stringify(event)]); const reply: unknown = await this.redis.eval(APPEND_AFTER_TAIL, { keys: [this.streamId], arguments: [input.tailEntryId ?? '', ttlArgument(input.streamTTLSeconds), ...pairs], }); - if (reply === 'ok' || reply === 'stale') { - return reply; - } - throw new Error(`Unexpected stream repair reply on ${this.streamId}`); + return parseAppendReply(reply, this.streamId); } private entryMillisecondFor(event: T): string { diff --git a/packages/trueforge/src/runtime/event-subscription/repair.ts b/packages/trueforge/src/runtime/event-subscription/repair.ts index 3105e69a54..b4ccf0d138 100644 --- a/packages/trueforge/src/runtime/event-subscription/repair.ts +++ b/packages/trueforge/src/runtime/event-subscription/repair.ts @@ -11,7 +11,11 @@ export interface EventLogEntry { } export interface RepairLogOptions { - entries: readonly EventLogEntry[]; + /** + * Durable rows in persist order, one page at a time. Called again when a + * concurrent append moves the tail, so the walk restarts from the first page. + */ + readPages: () => AsyncIterable[]>; streamTTLSeconds?: number | undefined; /** * Set by the database owner that is about to produce. Claims the stream-owner @@ -27,6 +31,20 @@ export interface StreamTail { payload: string; } +/** `'stale'` means the tail moved and the caller must recompute. */ +export type AppendAfterTailResult = { status: 'ok'; tailEntryId: string } | { status: 'stale' }; + +/** One page, for a caller that already holds the rows. Each call starts a new walk. */ +export function eventLogPages( + entries: readonly EventLogEntry[], +): () => AsyncIterable[]> { + return () => + (async function* () { + await Promise.resolve(); + yield entries; + })(); +} + /** Backfill gives up after this many concurrent tail moves. */ const REPAIR_MAX_ATTEMPTS = 5; @@ -117,14 +135,15 @@ function suffixAfterTail( /** * Appends the missing suffix only while the stream tail is still the one it was - * computed from. A concurrent writer (live producer or another repair) moves the - * tail; the suffix is then recomputed, so stale rows never land after newer ones. + * computed from. Pages stay in persist order: an empty stream is filled from + * the start, and an existing stream appends only after its tail. A concurrent + * writer moves the tail; the walk restarts, so stale rows never land after newer ones. */ export async function applyEventLogRepair(input: { - entries: readonly EventLogEntry[]; + readPages: () => AsyncIterable[]>; claimOwner?: (() => Promise) | undefined; readTail: () => Promise; - appendAfterTail: (batch: { tailEntryId: string | undefined; events: T[] }) => Promise<'ok' | 'stale'>; + appendAfterTail: (batch: { tailEntryId: string | undefined; events: T[] }) => Promise; refreshTtl: () => Promise; /** Producer resume refreshes TTL even when the existing stream is already complete. */ refreshTtlOnUntouchedStream: boolean; @@ -134,20 +153,61 @@ export async function applyEventLogRepair(input: { } for (let attempt = 0; attempt < REPAIR_MAX_ATTEMPTS; attempt += 1) { const tail = await input.readTail(); - const suffix = suffixAfterTail(input.entries, tail); - if (suffix.length === 0) { - if (input.refreshTtlOnUntouchedStream) { - await input.refreshTtl(); + const outcome = await appendMissingPages({ + readPages: input.readPages, + tail, + appendAfterTail: input.appendAfterTail, + }); + if (outcome === 'stale') { + continue; + } + if (outcome === 'untouched' && input.refreshTtlOnUntouchedStream) { + await input.refreshTtl(); + } + return; + } + throw new Error(`Stream repair gave up after ${String(REPAIR_MAX_ATTEMPTS)} concurrent tail changes`); +} + +/** + * One page is in memory at a time. The next page's expected tail is the entry + * id the previous append wrote, so a producer cannot land between the two. + */ +async function appendMissingPages(input: { + readPages: () => AsyncIterable[]>; + tail: StreamTail | undefined; + appendAfterTail: (batch: { tailEntryId: string | undefined; events: T[] }) => Promise; +}): Promise<'appended' | 'untouched' | 'stale'> { + let expectedTailId = input.tail?.entryId; + let passedTail = input.tail === undefined; + let appended = false; + for await (const page of input.readPages()) { + let rows: readonly EventLogEntry[]; + if (!passedTail) { + const tail = input.tail; + if (tail === undefined || !page.some(entry => entry.id === tail.eventId)) { + continue; } - return; + passedTail = true; + rows = suffixAfterTail(page, tail); + } else { + rows = page; + } + if (rows.length === 0) { + continue; } const outcome = await input.appendAfterTail({ - tailEntryId: tail?.entryId, - events: suffix.map(entry => entry.event), + tailEntryId: expectedTailId, + events: rows.map(entry => entry.event), }); - if (outcome === 'ok') { - return; + if (outcome.status === 'stale') { + return 'stale'; } + expectedTailId = outcome.tailEntryId; + appended = true; } - throw new Error(`Stream repair gave up after ${String(REPAIR_MAX_ATTEMPTS)} concurrent tail changes`); + if (!passedTail) { + return 'untouched'; + } + return appended ? 'appended' : 'untouched'; } diff --git a/packages/trueforge/src/runtime/turnEventLogRepair.ts b/packages/trueforge/src/runtime/turnEventLogRepair.ts index 753b3af156..39604c8cdd 100644 --- a/packages/trueforge/src/runtime/turnEventLogRepair.ts +++ b/packages/trueforge/src/runtime/turnEventLogRepair.ts @@ -1,26 +1,48 @@ import type { ISessionStore, TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session'; +import { EVENTS_PAGE_LIMIT } from '../schemas/common'; import type { EventSubscription } from './event-subscription'; /** * Rebuild a paused turn's live stream from the durable log. * Subscribe omits `producerExecutorId` and only appends. * The producing resume, which already owns the turn row, claims the stream-owner key. + * Pages are read ascending so an empty stream is appended from start to end. */ export async function repairTurnEventLog(input: { - sessionStore: Pick, 'listTurnEventLog'>; + sessionStore: Pick, 'listTurnEvents'>; subscription: EventSubscription; sessionId: string; turnId: string; streamTTLSeconds: number; producerExecutorId?: string | undefined; }): Promise { - const events = await input.sessionStore.listTurnEventLog({ - session_id: input.sessionId, - turn_id: input.turnId, - }); await input.subscription.repairLog({ - entries: events.map(event => ({ id: event.id, event })), + readPages: () => durableEventPages(input), streamTTLSeconds: input.streamTTLSeconds, ...(input.producerExecutorId === undefined ? {} : { executorId: input.producerExecutorId }), }); } + +/** One page in memory. `order: 'asc'` is persist order, which is also append order. */ +async function* durableEventPages(input: { + sessionStore: Pick, 'listTurnEvents'>; + sessionId: string; + turnId: string; +}): AsyncGenerator<{ id: string; event: TurnStreamingEvent }[]> { + let pageToken: string | undefined; + for (;;) { + const page = await input.sessionStore.listTurnEvents({ + session_id: input.sessionId, + turn_id: input.turnId, + limit: EVENTS_PAGE_LIMIT, + page_token: pageToken, + order: 'asc', + }); + yield page.data.map(event => ({ id: event.id, event })); + const nextPageToken = page.pagination.next_page_token; + if (nextPageToken === undefined) { + return; + } + pageToken = nextPageToken; + } +} diff --git a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts index d6c1301d53..a5b6948e33 100644 --- a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts +++ b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts @@ -5,6 +5,7 @@ import { type EventSubscriptionRegistry, type SequencedEvent, } from '../../../src/runtime/event-subscription'; +import { eventLogPages } from '../../../src/runtime/event-subscription/repair'; /** Payload shape used by the suite; distinct per event so round-trips are provable. */ export interface TestEvent { @@ -171,10 +172,10 @@ export function runEventSubscriptionContractSuite( }; const third = makeEvent(2); await producer.repairLog({ - entries: [ + readPages: eventLogPages([ { id: 'evt-0', event: fromDatabase }, { id: 'evt-2', event: third }, - ], + ]), streamTTLSeconds: STREAM_TTL_SECONDS, }); @@ -211,10 +212,10 @@ export function runEventSubscriptionContractSuite( await evictStream(streamId); await producer.repairLog({ - entries: [ + readPages: eventLogPages([ { id: first.id, event: first }, { id: second.id, event: second }, - ], + ]), streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner', }); @@ -235,10 +236,10 @@ export function runEventSubscriptionContractSuite( await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS }); const repair = () => producer.repairLog({ - entries: [ + readPages: eventLogPages([ { id: first.id, event: first }, { id: second.id, event: second }, - ], + ]), streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner', }); @@ -257,7 +258,7 @@ export function runEventSubscriptionContractSuite( const owner = registry.get(streamId); const event = makeEvent(1); await owner.repairLog({ - entries: [{ id: event.id, event }], + readPages: eventLogPages([{ id: event.id, event }]), streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner', }); diff --git a/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts b/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts index ce01d29d61..7cc369616b 100644 --- a/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts @@ -1,5 +1,5 @@ import { eventIdSeekCursor, eventIdTimestampMs } from '../../../src/runtime/event-subscription/entryId'; -import { eventPayloadMatches } from '../../../src/runtime/event-subscription/repair'; +import { eventLogPages, eventPayloadMatches } from '../../../src/runtime/event-subscription/repair'; const EVENT_ID = '01arz3ndektsv4rrffq69g5fav'; @@ -52,10 +52,10 @@ describe('suffixAfterTail via repairLog', () => { await producer.put(shell, { streamTTLSeconds: 300 }); await producer.put(deltaTip, { streamTTLSeconds: 300 }); await producer.repairLog({ - entries: [ + readPages: eventLogPages([ { id: 'evt-1', event: shell }, { id: 'evt-2', event: next }, - ], + ]), streamTTLSeconds: 300, }); const abort = new AbortController(); diff --git a/packages/trueforge/tests/runtime/event-subscription/repair.test.ts b/packages/trueforge/tests/runtime/event-subscription/repair.test.ts index 6809a21799..9d20e39022 100644 --- a/packages/trueforge/tests/runtime/event-subscription/repair.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/repair.test.ts @@ -1,4 +1,4 @@ -import { applyEventLogRepair, type StreamTail } from '../../../src/runtime/event-subscription/repair'; +import { applyEventLogRepair, eventLogPages, type StreamTail } from '../../../src/runtime/event-subscription/repair'; interface Row { id: string; @@ -20,11 +20,11 @@ describe('applyEventLogRepair', () => { const batches: { tailEntryId: string | undefined; events: Row[] }[] = []; await applyEventLogRepair({ - entries, + readPages: eventLogPages(entries), readTail: () => Promise.resolve(tails.shift()), appendAfterTail: batch => { batches.push(batch); - return Promise.resolve('stale'); + return Promise.resolve({ status: 'stale' }); }, refreshTtl: () => Promise.resolve(), refreshTtlOnUntouchedStream: false, @@ -38,12 +38,49 @@ describe('applyEventLogRepair', () => { const paused: Row = { id: 'evt-2', type: 'turn.update' }; await expect( applyEventLogRepair({ - entries: [created, paused].map(event => ({ id: event.id, event })), + readPages: eventLogPages([created, paused].map(event => ({ id: event.id, event }))), readTail: () => Promise.resolve(tailOf(created, '1-0')), - appendAfterTail: () => Promise.resolve('stale'), + appendAfterTail: () => Promise.resolve({ status: 'stale' }), refreshTtl: () => Promise.resolve(), refreshTtlOnUntouchedStream: false, }), ).rejects.toThrow('Stream repair gave up'); }); + + it('appends later pages after the entry id the previous page wrote', async () => { + const created: Row = { id: 'evt-1', type: 'turn.created' }; + const paused: Row = { id: 'evt-2', type: 'turn.update' }; + const running: Row = { id: 'evt-3', type: 'turn.update' }; + const done: Row = { id: 'evt-4', type: 'turn.done' }; + const batches: { tailEntryId: string | undefined; events: Row[] }[] = []; + const writtenTailIds = ['8-0', '9-0']; + + await applyEventLogRepair({ + readPages: () => + (async function* () { + yield [{ id: created.id, event: created }]; + yield [ + { id: paused.id, event: paused }, + { id: running.id, event: running }, + ]; + yield [{ id: done.id, event: done }]; + })(), + readTail: () => Promise.resolve(tailOf(paused, '2-0')), + appendAfterTail: batch => { + batches.push(batch); + const tailEntryId = writtenTailIds.shift(); + if (tailEntryId === undefined) { + return Promise.resolve({ status: 'stale' }); + } + return Promise.resolve({ status: 'ok', tailEntryId }); + }, + refreshTtl: () => Promise.resolve(), + refreshTtlOnUntouchedStream: false, + }); + + expect(batches).toEqual([ + { tailEntryId: '2-0', events: [running] }, + { tailEntryId: '8-0', events: [done] }, + ]); + }); }); From 1071348c69973de316e8d4cd1765c4ebf9d1c82f Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 3 Nov 2026 17:51:09 +0530 Subject: [PATCH 10/21] Throw the turn lock timeout directly from withTimeout Pass TurnLockAcquireTimeoutError as the mutex timeout error so the acquire path no longer catches and rewraps E_TIMEOUT. Co-authored-by: Cursor --- packages/trueforge/src/runtime/activeTurns.ts | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 0531553593..3974ca6529 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -6,7 +6,7 @@ */ import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; import { getOperationMetrics } from '@truefoundry/trueforge-core/core'; -import { E_TIMEOUT, Mutex, withTimeout } from 'async-mutex'; +import { E_TIMEOUT, Mutex, withTimeout, type MutexInterface } from 'async-mutex'; import type { Logger } from 'winston'; const DEFAULT_PAUSED_KEEPALIVE_MS = 300_000; @@ -101,17 +101,19 @@ export class ActiveTurnRegistry { this.turnLocks.set(key, mutex); } // 0 waits forever. withTimeout(0) would reject on the next timer turn. - const guarded = this.lockAcquireTimeoutMs > 0 ? withTimeout(mutex, this.lockAcquireTimeoutMs) : mutex; - try { - return await guarded.runExclusive(fn); - } catch (caught) { - if (caught === E_TIMEOUT) { - throw new TurnLockAcquireTimeoutError({ + let guarded: MutexInterface = mutex; + if (this.lockAcquireTimeoutMs > 0) { + guarded = withTimeout( + mutex, + this.lockAcquireTimeoutMs, + new TurnLockAcquireTimeoutError({ timeoutMs: this.lockAcquireTimeoutMs, cause: E_TIMEOUT, - }); - } - throw caught; + }), + ); + } + try { + return await guarded.runExclusive(fn); } finally { if (!mutex.isLocked()) { this.turnLocks.delete(key); From 5f6d7edbf9024da5ad53c2902faa687d61b349b6 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 4 Nov 2026 10:16:15 +0530 Subject: [PATCH 11/21] Reject after_sequence_number on new-style turn streams Only leftover {sequence}-1 streams accept the legacy cursor; new streams return 400 for after_sequence_number and Last-Event-ID. Drop the skip-count compat path and mark subscribe non-resumable for Fern so the SDK stops reconnecting with Last-Event-ID. Co-authored-by: Cursor --- .changeset/event-id-subscribe-cursor.md | 2 +- packages/trueforge/src/apis/turns.ts | 33 ++++--------- packages/trueforge/src/routes/turnRoutes.ts | 2 +- packages/trueforge/src/schemas/turn.ts | 2 +- .../trueforge/tests/unit/apis/turns.test.ts | 46 +++++-------------- 5 files changed, 22 insertions(+), 63 deletions(-) diff --git a/.changeset/event-id-subscribe-cursor.md b/.changeset/event-id-subscribe-cursor.md index 78176a73e4..7ef905f39a 100644 --- a/.changeset/event-id-subscribe-cursor.md +++ b/.changeset/event-id-subscribe-cursor.md @@ -6,4 +6,4 @@ '@truefoundry/trueforge-sdk': patch --- -Subscribe with from_event_id or after_event_id. New turn streams store and repair by event id and omit SSE sequence ids. after_sequence_number remains for older integer streams. Regenerate SDK from the updated OpenAPI spec. +Subscribe with from_event_id or after_event_id. New turn streams store and repair by event id and omit SSE sequence ids. after_sequence_number remains for older integer streams. The subscribe stream is not Fern-resumable, so the SDK does not reconnect it with Last-Event-ID. Regenerate SDK from the updated OpenAPI spec. diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index e891b18b8f..b4971b4dfe 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -2070,21 +2070,21 @@ export function createTurnsRouter(deps: TurnsRouterDeps LEGACY_SEQUENCE_ID_MAX) { + throw new HTTPException(400, { + message: 'after_sequence_number is not supported on this stream. Use from_event_id or after_event_id.', + }); + } + if (afterSequenceNumber > LEGACY_SEQUENCE_ID_MAX) { throw new HTTPException(400, { message: 'after_sequence_number is not supported above 50000.', }); - } else { - includeSequenceId = true; } + includeSequenceId = true; } // One lifecycle controller: server-side timeout, client disconnect, and // normal teardown all abort it, which ends poll even while it is parked @@ -2117,25 +2117,8 @@ export function createTurnsRouter(deps: TurnsRouterDeps 0 ? seen : undefined; - seen += 1; - await stream.writeSSE( - includeSequenceId && sequenceNumber !== undefined - ? turnEventSsePayload(event, sequenceNumber) - : compatIndex !== undefined - ? turnEventSsePayload(event, compatIndex) - : turnEventSsePayload(event), - ); + await stream.writeSSE(turnEventSsePayload(event, includeSequenceId ? sequenceNumber : undefined)); if (event.type === EventType.TURN_DONE) { break; } diff --git a/packages/trueforge/src/routes/turnRoutes.ts b/packages/trueforge/src/routes/turnRoutes.ts index 7bfb4e6ee5..f1ca71a657 100644 --- a/packages/trueforge/src/routes/turnRoutes.ts +++ b/packages/trueforge/src/routes/turnRoutes.ts @@ -299,7 +299,7 @@ export const subscribeTurnRoute = createRoute({ 'Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id', 'x-fern-sdk-group-name': ['sessions'], 'x-fern-sdk-method-name': 'subscribe_to_turn', - 'x-fern-streaming': { format: 'sse', resumable: true }, + 'x-fern-streaming': { format: 'sse', resumable: false }, request: { params: TurnIdParamsSchema, query: SubscribeTurnQuerySchema, diff --git a/packages/trueforge/src/schemas/turn.ts b/packages/trueforge/src/schemas/turn.ts index 382945d146..d607385825 100644 --- a/packages/trueforge/src/schemas/turn.ts +++ b/packages/trueforge/src/schemas/turn.ts @@ -18,7 +18,7 @@ export const SubscribeTurnQuerySchema = z .nonnegative() .optional() .describe( - 'Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.', + 'Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.', ), ), from_event_id: z diff --git a/packages/trueforge/tests/unit/apis/turns.test.ts b/packages/trueforge/tests/unit/apis/turns.test.ts index 2e472dc0a0..64283529dc 100644 --- a/packages/trueforge/tests/unit/apis/turns.test.ts +++ b/packages/trueforge/tests/unit/apis/turns.test.ts @@ -958,38 +958,18 @@ describe('turns', () => { expect(midMessage.status).toBe(200); await midMessage.body?.cancel(); - // An older client counted id-less events from 0. after=1 has indexes 0 and 1. - const countedTurnId = 'counted-turn'; - await sessionStore.createTurn(makeCreateTurnInput({ sessionId, turnId: countedTurnId })); - const firstCounted = makeModelMessageEvent(); - const kept = makeModelMessageEvent(); - const counted = [firstCounted, makeModelMessageEvent(), kept]; - const countedStream = eventSubscriptions.get(turnStreamId('default', sessionId, countedTurnId)); - for (const event of counted) { - await countedStream.put(event, { streamTTLSeconds: 300 }); - } - const countedResume = await app.request(`/${sessionId}/turns/${countedTurnId}/subscribe?after_sequence_number=1`); - expect(countedResume.status).toBe(200); - const countedReader = countedResume.body?.getReader(); - if (countedReader === undefined) { - throw new Error('subscribe response has no body'); - } - const countedDecoder = new TextDecoder(); - let countedBody = ''; - try { - while (!countedBody.includes(kept.id)) { - const { value, done } = await countedReader.read(); - if (done || value === undefined) { - break; - } - countedBody += countedDecoder.decode(value, { stream: true }); - } - } finally { - await countedReader.cancel(); + // after_sequence_number is only for leftover integer streams. + const eventIdTurnId = 'event-id-turn'; + await sessionStore.createTurn(makeCreateTurnInput({ sessionId, turnId: eventIdTurnId })); + await eventSubscriptions + .get(turnStreamId('default', sessionId, eventIdTurnId)) + .put(makeModelMessageEvent(), { streamTTLSeconds: 300 }); + for (const cursor of [0, 1, 100001]) { + const rejected = await app.request( + `/${sessionId}/turns/${eventIdTurnId}/subscribe?after_sequence_number=${String(cursor)}`, + ); + expect(rejected.status).toBe(400); } - expect(countedBody).toContain(kept.id); - expect(countedBody).not.toContain(firstCounted.id); - expect(countedBody).toContain('id: 2'); const runningTurnId = 'running-turn'; await sessionStore.createTurn(makeCreateTurnInput({ sessionId, turnId: runningTurnId })); @@ -1063,10 +1043,6 @@ describe('turns', () => { ); expect(mixed.status).toBe(400); - const huge = await app.request(`/${sessionId}/turns/${pausedTurnId}/subscribe?after_sequence_number=100001`); - expect(huge.status).toBe(200); - await huge.body?.cancel(); - const invalid = await app.request(`/${sessionId}/turns/${pausedTurnId}/subscribe?from_event_id=not-a-ulid`); expect(invalid.status).toBe(400); From d4f21846fe9d2d2de1f63c01b68fc628d5b39499 Mon Sep 17 00:00:00 2001 From: "trueforge-dev-bot[bot]" Date: Wed, 4 Nov 2026 04:49:01 +0000 Subject: [PATCH 12/21] Regenerate OpenAPI document and SDKs --- .github/fern/openapi/openapi.json | 6 ++-- docs/openapi.json | 6 ++-- .../api/resources/sessions/client/Client.ts | 36 ------------------- .../SubscribeToTurnSessionsRequest.ts | 2 +- python/trueforge_sdk/reference.md | 2 +- .../src/trueforge_sdk/sessions/client.py | 4 +-- .../src/trueforge_sdk/sessions/raw_client.py | 4 +-- 7 files changed, 12 insertions(+), 48 deletions(-) diff --git a/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index cad043038e..49c60ecd73 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -10767,12 +10767,12 @@ } }, { - "description": "Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.", + "description": "Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.", "in": "query", "name": "after_sequence_number", "required": false, "schema": { - "description": "Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.", + "description": "Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.", "minimum": 0, "type": [ "integer", @@ -10865,7 +10865,7 @@ "x-fern-sdk-method-name": "subscribe_to_turn", "x-fern-streaming": { "format": "sse", - "resumable": true + "resumable": false } } }, diff --git a/docs/openapi.json b/docs/openapi.json index cad043038e..49c60ecd73 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -10767,12 +10767,12 @@ } }, { - "description": "Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.", + "description": "Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.", "in": "query", "name": "after_sequence_number", "required": false, "schema": { - "description": "Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id.", + "description": "Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.", "minimum": 0, "type": [ "integer", @@ -10865,7 +10865,7 @@ "x-fern-sdk-method-name": "subscribe_to_turn", "x-fern-streaming": { "format": "sse", - "resumable": true + "resumable": false } } }, diff --git a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts index eaadc33a03..23c30cd16b 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts @@ -2191,35 +2191,6 @@ export class SessionsClient { fetchFn: this._options?.fetch, logging: this._options.logging, }); - const _reconnect = async (lastEventId: string) => { - const _reconnectResponse = await (this._options.fetcher ?? core.fetcher)({ - url: core.url.join( - (await core.Supplier.get(this._options.baseUrl)) ?? - (await core.Supplier.get(this._options.environment)), - `api/v1/sessions/${core.url.encodePathParam(session_id, errors.TrueForgeError)}/turns/${core.url.encodePathParam(turn_id, errors.TrueForgeError)}/subscribe`, - ), - method: "GET", - headers: { ..._headers, "Last-Event-ID": lastEventId }, - queryString: core.url - .queryBuilder() - .addMany(_queryParams) - .mergeAdditional(requestOptions?.queryParams) - .build(), - responseType: "sse", - timeoutMs: (requestOptions?.timeoutInSeconds ?? this._options?.timeoutInSeconds ?? 60) * 1000, - maxRetries: requestOptions?.maxRetries ?? this._options?.maxRetries, - abortSignal: requestOptions?.abortSignal, - fetchFn: this._options?.fetch, - logging: this._options.logging, - }); - if (!_reconnectResponse.ok) { - throw new Error("SSE stream reconnection failed"); - } - if (_reconnectResponse.body == null) { - throw new Error("SSE stream reconnection failed: empty response body"); - } - return _reconnectResponse.body; - }; if (_response.ok) { return { data: new core.Stream({ @@ -2236,14 +2207,7 @@ export class SessionsClient { signal: requestOptions?.abortSignal, eventShape: { type: "sse", - resumable: true, }, - reconnectionEnabled: - requestOptions?.stream?.reconnectionEnabled ?? this._options?.stream?.reconnectionEnabled, - maxReconnectionAttempts: - requestOptions?.stream?.maxReconnectionAttempts ?? - this._options?.stream?.maxReconnectionAttempts, - reconnect: _reconnect, }), rawResponse: _response.rawResponse, }; diff --git a/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts b/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts index 9fe7a07bc8..7a3a016f20 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts @@ -5,7 +5,7 @@ * {} */ export interface SubscribeToTurnSessionsRequest { - /** Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. */ + /** Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. */ afterSequenceNumber?: number | null; /** Inclusive cursor. Replay starts at the entry with this event id */ fromEventId?: string; diff --git a/python/trueforge_sdk/reference.md b/python/trueforge_sdk/reference.md index a3f130867a..b355ff5c50 100644 --- a/python/trueforge_sdk/reference.md +++ b/python/trueforge_sdk/reference.md @@ -3519,7 +3519,7 @@ client.sessions.subscribe_to_turn(
-**after_sequence_number:** `typing.Optional[int]` — Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. +**after_sequence_number:** `typing.Optional[int]` — Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.
diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py index 484a94c3bf..f5e2ee4fbb 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py @@ -850,7 +850,7 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. + Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id @@ -1847,7 +1847,7 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. + Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py index 5235647b8f..2075c59d6b 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py @@ -1763,7 +1763,7 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. + Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id @@ -3606,7 +3606,7 @@ async def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor. On an event-id stream it is the last 0-based index an older client counted. Omit when using an event id. + Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id From 74c1d8bfbf19a19015ff545878bc17281d815ec6 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 4 Nov 2026 10:25:42 +0530 Subject: [PATCH 13/21] Trim subscribe API descriptions to caller contract Co-authored-by: Cursor --- packages/trueforge/src/routes/turnRoutes.ts | 5 ++--- packages/trueforge/src/schemas/turn.ts | 4 +--- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/packages/trueforge/src/routes/turnRoutes.ts b/packages/trueforge/src/routes/turnRoutes.ts index f1ca71a657..01c26c7d42 100644 --- a/packages/trueforge/src/routes/turnRoutes.ts +++ b/packages/trueforge/src/routes/turnRoutes.ts @@ -295,8 +295,7 @@ export const subscribeTurnRoute = createRoute({ path: '/{session_id}/turns/{turn_id}/subscribe', tags: [OpenApiTag.AGENT_SESSIONS], summary: 'Subscribe to a turn stream', - description: - 'Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id', + description: 'Subscribe to the live SSE stream. Resume with from_event_id or after_event_id.', 'x-fern-sdk-group-name': ['sessions'], 'x-fern-sdk-method-name': 'subscribe_to_turn', 'x-fern-streaming': { format: 'sse', resumable: false }, @@ -327,7 +326,7 @@ export const subscribeTurnRoute = createRoute({ }, 412: { content: { 'application/json': { schema: RequestErrorResponseSchema } }, - description: 'The live stream does not exist and this turn is not paused, so it was not restored.', + description: 'The turn stream is no longer available or resumable.', }, }, }); diff --git a/packages/trueforge/src/schemas/turn.ts b/packages/trueforge/src/schemas/turn.ts index d607385825..aea42b14cc 100644 --- a/packages/trueforge/src/schemas/turn.ts +++ b/packages/trueforge/src/schemas/turn.ts @@ -17,9 +17,7 @@ export const SubscribeTurnQuerySchema = z .int() .nonnegative() .optional() - .describe( - 'Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.', - ), + .describe('Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.'), ), from_event_id: z .string() From 21068a53a32f318a63f30412dc3388868c37d86c Mon Sep 17 00:00:00 2001 From: "trueforge-dev-bot[bot]" Date: Wed, 4 Nov 2026 04:58:20 +0000 Subject: [PATCH 14/21] Regenerate OpenAPI document and SDKs --- .github/fern/openapi/openapi.json | 8 ++++---- docs/openapi.json | 8 ++++---- packages/trueforge-sdk/reference.md | 2 +- .../src/api/resources/sessions/client/Client.ts | 2 +- .../client/requests/SubscribeToTurnSessionsRequest.ts | 2 +- python/trueforge_sdk/reference.md | 4 ++-- python/trueforge_sdk/src/trueforge_sdk/sessions/client.py | 8 ++++---- .../src/trueforge_sdk/sessions/raw_client.py | 8 ++++---- 8 files changed, 21 insertions(+), 21 deletions(-) diff --git a/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index 49c60ecd73..1d286b2c1d 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -10741,7 +10741,7 @@ }, "/api/v1/sessions/{session_id}/turns/{turn_id}/subscribe": { "get": { - "description": "Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id", + "description": "Subscribe to the live SSE stream. Resume with from_event_id or after_event_id.", "parameters": [ { "description": "Session identifier.", @@ -10767,12 +10767,12 @@ } }, { - "description": "Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.", + "description": "Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.", "in": "query", "name": "after_sequence_number", "required": false, "schema": { - "description": "Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.", + "description": "Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.", "minimum": 0, "type": [ "integer", @@ -10852,7 +10852,7 @@ } } }, - "description": "The live stream does not exist and this turn is not paused, so it was not restored." + "description": "The turn stream is no longer available or resumable." } }, "summary": "Subscribe to a turn stream", diff --git a/docs/openapi.json b/docs/openapi.json index 49c60ecd73..1d286b2c1d 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -10741,7 +10741,7 @@ }, "/api/v1/sessions/{session_id}/turns/{turn_id}/subscribe": { "get": { - "description": "Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id", + "description": "Subscribe to the live SSE stream. Resume with from_event_id or after_event_id.", "parameters": [ { "description": "Session identifier.", @@ -10767,12 +10767,12 @@ } }, { - "description": "Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.", + "description": "Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.", "in": "query", "name": "after_sequence_number", "required": false, "schema": { - "description": "Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id.", + "description": "Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.", "minimum": 0, "type": [ "integer", @@ -10852,7 +10852,7 @@ } } }, - "description": "The live stream does not exist and this turn is not paused, so it was not restored." + "description": "The turn stream is no longer available or resumable." } }, "summary": "Subscribe to a turn stream", diff --git a/packages/trueforge-sdk/reference.md b/packages/trueforge-sdk/reference.md index 888f671391..aaac034bd3 100644 --- a/packages/trueforge-sdk/reference.md +++ b/packages/trueforge-sdk/reference.md @@ -2921,7 +2921,7 @@ await client.sessions.createTurnEvents("session_id", "turn_id", {
-Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id +Subscribe to the live SSE stream. Resume with from_event_id or after_event_id.
diff --git a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts index 23c30cd16b..3cb42f349c 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts @@ -2140,7 +2140,7 @@ export class SessionsClient { } /** - * Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id + * Subscribe to the live SSE stream. Resume with from_event_id or after_event_id. */ public subscribeToTurn( session_id: string, diff --git a/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts b/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts index 7a3a016f20..a7434aa897 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/requests/SubscribeToTurnSessionsRequest.ts @@ -5,7 +5,7 @@ * {} */ export interface SubscribeToTurnSessionsRequest { - /** Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. */ + /** Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead. */ afterSequenceNumber?: number | null; /** Inclusive cursor. Replay starts at the entry with this event id */ fromEventId?: string; diff --git a/python/trueforge_sdk/reference.md b/python/trueforge_sdk/reference.md index b355ff5c50..6b870fe380 100644 --- a/python/trueforge_sdk/reference.md +++ b/python/trueforge_sdk/reference.md @@ -3462,7 +3462,7 @@ client.sessions.create_turn_events(
-Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id +Subscribe to the live SSE stream. Resume with from_event_id or after_event_id.
@@ -3519,7 +3519,7 @@ client.sessions.subscribe_to_turn(
-**after_sequence_number:** `typing.Optional[int]` — Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. +**after_sequence_number:** `typing.Optional[int]` — Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.
diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py index f5e2ee4fbb..4ba56d3f67 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/client.py @@ -839,7 +839,7 @@ def subscribe_to_turn( request_options: typing.Optional[RequestOptions] = None, ) -> Stream[TurnStreamingEvent]: """ - Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id + Subscribe to the live SSE stream. Resume with from_event_id or after_event_id. Parameters ---------- @@ -850,7 +850,7 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. + Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id @@ -1836,7 +1836,7 @@ def subscribe_to_turn( request_options: typing.Optional[RequestOptions] = None, ) -> AsyncStream[TurnStreamingEvent]: """ - Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id + Subscribe to the live SSE stream. Resume with from_event_id or after_event_id. Parameters ---------- @@ -1847,7 +1847,7 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. + Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id diff --git a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py index 2075c59d6b..361cf60c5d 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py @@ -1752,7 +1752,7 @@ def subscribe_to_turn( request_options: typing.Optional[RequestOptions] = None, ) -> typing.Iterator[HttpResponse[Stream[TurnStreamingEvent]]]: """ - Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id + Subscribe to the live SSE stream. Resume with from_event_id or after_event_id. Parameters ---------- @@ -1763,7 +1763,7 @@ def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. + Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id @@ -3595,7 +3595,7 @@ async def subscribe_to_turn( request_options: typing.Optional[RequestOptions] = None, ) -> typing.AsyncIterator[AsyncHttpResponse[AsyncStream[TurnStreamingEvent]]]: """ - Subscribe to the live SSE stream. A paused turn with an expired stream is restored from stored events. Resume with from_event_id or after_event_id + Subscribe to the live SSE stream. Resume with from_event_id or after_event_id. Parameters ---------- @@ -3606,7 +3606,7 @@ async def subscribe_to_turn( Turn identifier. after_sequence_number : typing.Optional[int] - Legacy exclusive cursor for older integer streams only. Rejected on event-id streams; use from_event_id or after_event_id. + Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead. from_event_id : typing.Optional[str] Inclusive cursor. Replay starts at the entry with this event id From a0a11b7b51a4b0469afeadd10485e2f10fb9b106 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 5 Nov 2026 09:46:12 +0530 Subject: [PATCH 15/21] Restore a missing turn stream from the whole durable log. Partial repair cannot match emit order, so an existing stream is left untouched. Seeks compare a dedicated id field instead of decoding payloads, and the Redis owner key is removed so ownership stays on the turn row. Co-authored-by: Cursor --- .changeset/event-id-subscribe-cursor.md | 9 - .changeset/paused-subscribe-stream-restore.md | 5 - .../src/convertTurnMessages.ts | 10 +- .../test/convertTurnMessages.test.ts | 29 +- packages/trueforge/src/apis/turns.ts | 29 +- packages/trueforge/src/runtime/activeTurns.ts | 13 +- .../runtime/event-subscription/inMemory.ts | 98 +------ .../src/runtime/event-subscription/index.ts | 27 +- .../src/runtime/event-subscription/redis.ts | 262 ++++++++---------- .../src/runtime/event-subscription/repair.ts | 213 -------------- .../src/runtime/event-subscription/restore.ts | 45 +++ ...entLogRepair.ts => turnEventLogRestore.ts} | 24 +- .../event-subscription/contractSuite.ts | 136 ++++----- .../event-subscription/entryId.test.ts | 53 ---- .../event-subscription/redis.contract.test.ts | 42 +++ .../runtime/event-subscription/repair.test.ts | 86 ------ .../event-subscription/restore.test.ts | 88 ++++++ ...ir.test.ts => turnEventLogRestore.test.ts} | 37 +-- .../tests/unit/runtime/activeTurns.test.ts | 27 ++ 19 files changed, 498 insertions(+), 735 deletions(-) delete mode 100644 .changeset/event-id-subscribe-cursor.md delete mode 100644 .changeset/paused-subscribe-stream-restore.md delete mode 100644 packages/trueforge/src/runtime/event-subscription/repair.ts create mode 100644 packages/trueforge/src/runtime/event-subscription/restore.ts rename packages/trueforge/src/runtime/{turnEventLogRepair.ts => turnEventLogRestore.ts} (55%) delete mode 100644 packages/trueforge/tests/runtime/event-subscription/repair.test.ts create mode 100644 packages/trueforge/tests/runtime/event-subscription/restore.test.ts rename packages/trueforge/tests/runtime/event-subscription/{turnEventLogRepair.test.ts => turnEventLogRestore.test.ts} (74%) diff --git a/.changeset/event-id-subscribe-cursor.md b/.changeset/event-id-subscribe-cursor.md deleted file mode 100644 index 7ef905f39a..0000000000 --- a/.changeset/event-id-subscribe-cursor.md +++ /dev/null @@ -1,9 +0,0 @@ ---- -'@truefoundry/trueforge': patch -'@truefoundry/trueforge-core': patch -'@truefoundry/trueforge-assistant-ui-runtime': patch -'@truefoundry/trueforge-ui': patch -'@truefoundry/trueforge-sdk': patch ---- - -Subscribe with from_event_id or after_event_id. New turn streams store and repair by event id and omit SSE sequence ids. after_sequence_number remains for older integer streams. The subscribe stream is not Fern-resumable, so the SDK does not reconnect it with Last-Event-ID. Regenerate SDK from the updated OpenAPI spec. diff --git a/.changeset/paused-subscribe-stream-restore.md b/.changeset/paused-subscribe-stream-restore.md deleted file mode 100644 index d859a67513..0000000000 --- a/.changeset/paused-subscribe-stream-restore.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -'@truefoundry/trueforge': patch ---- - -Restore an expired live stream only for a paused turn. Running and finished turns keep the existing stream or use list events. diff --git a/packages/assistant-ui-runtime/src/convertTurnMessages.ts b/packages/assistant-ui-runtime/src/convertTurnMessages.ts index 13b79d90da..abbc1d70f2 100644 --- a/packages/assistant-ui-runtime/src/convertTurnMessages.ts +++ b/packages/assistant-ui-runtime/src/convertTurnMessages.ts @@ -1938,7 +1938,15 @@ export async function* streamTurnEvents( } const event = data.event; if ('id' in event && typeof event.id === 'string') { - onSubscribeCursor?.(subscribeCursorForEvent(event)); + const endsMcpAuthPrompt = + event.type === EVENT_TYPE.USER_MCP_AUTH_CONTINUE || + event.type === EVENT_TYPE.TURN_DONE || + (event.type === EVENT_TYPE.TURN_UPDATE && event.state.status === TURN_STATUS.RUNNING); + // Resume replays from mcp.auth_required until the prompt is answered. A later + // paused update must not move the cursor past it: a fresh fold would lose the prompt. + if (pendingMcpAuth == null || event.type === EVENT_TYPE.MCP_AUTH_REQUIRED || endsMcpAuthPrompt) { + onSubscribeCursor?.(subscribeCursorForEvent(event)); + } } if (event.type === EVENT_TYPE.TURN_CREATED) { diff --git a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts index af4e105b21..b8e536d868 100644 --- a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts +++ b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts @@ -41,7 +41,7 @@ import { turnToSessionRecord, } from '../src/sessionSnapshot.js'; import { applyUserToolResponsesToFold, toolResponseStatus } from '../src/toolResponse.js'; -import type { TurnStreamUpdate } from '../src/turnStreamUpdate.js'; +import type { TurnStreamUpdate, TurnSubscribeCursor } from '../src/turnStreamUpdate.js'; const createdAt = new Date().toISOString(); const SESSION_ID = 'session-1'; @@ -2336,6 +2336,33 @@ describe('convertTurnMessages', () => { }); }); + it('keeps the resume cursor on mcp.auth_required until a running update answers it', async () => { + const cursors: TurnSubscribeCursor[] = []; + const mcpAuth = { + type: 'mcp.auth_required' as const, + id: 'mcp-auth', + createdAt, + mcpServers: [{ id: 'github-auth', name: 'github', authUrl: 'https://example.com/auth' }], + }; + await collectStream( + streamTurnEvents( + streamFrom([ + mcpAuth, + { type: 'turn.update', id: 'pause-mcp', createdAt, threadId: null, state: { status: 'paused' } }, + { type: 'turn.update', id: 'run-mcp', createdAt, threadId: null, state: { status: 'running' } }, + { type: 'turn.done', id: 'done-mcp', createdAt, state: { status: 'done', completedAt: createdAt } }, + ]), + new PeerThreadFoldState(), + undefined, + undefined, + undefined, + cursor => cursors.push(cursor), + ), + ); + + expect(cursors).toEqual([{ fromEventId: 'mcp-auth' }, { afterEventId: 'run-mcp' }, { afterEventId: 'done-mcp' }]); + }); + it('stamps sandboxId onto every content yield after sandbox.created is seen', async () => { const foldState = new PeerThreadFoldState(); const updates = await collectStream( diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index b4971b4dfe..d18ccd360c 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -88,7 +88,6 @@ import { isLegacySequenceEntryId, LEGACY_SEQUENCE_ID_MAX, StreamGoneError, - StreamOwnerMismatchError, type EventSubscription, type EventSubscriptionPollOptions, type EventSubscriptionRegistry, @@ -101,7 +100,7 @@ import { resolveSandboxEnvironment, resolveSandboxProvider, } from '../runtime/sessionResources'; -import { repairTurnEventLog } from '../runtime/turnEventLogRepair'; +import { restoreTurnEventLog } from '../runtime/turnEventLogRestore'; import { DEFAULT_SANDBOX_ENVIRONMENT_NAME } from '../schemas/sandboxEnvironment'; import { MAX_SESSION_TITLE_LENGTH } from '../schemas/session'; import { assertGatewayMetadataRequestHeaders } from '../truefoundry/gatewayMetadata'; @@ -550,18 +549,14 @@ export async function drainTurnEvents(input: { executionTimeout: TurnExecutionTimeout; logger: Logger; pausedKeepalive: TurnPausedKeepalive; - executorId: string; onEvent?: (event: TurnStreamingEvent) => Promise; }): Promise { - const { trackedStream, turnEventStream, executionTimeout, logger, pausedKeepalive, executorId, onEvent } = input; + const { trackedStream, turnEventStream, executionTimeout, logger, pausedKeepalive, onEvent } = input; try { for await (const event of trackedStream) { reportUnexpectedTurnDone({ event, logger }); // Dual-write before any client sink so subscribers can resume after disconnect. - await turnEventStream.put(event, { - streamTTLSeconds: streamTTLSecondsFor(event), - executorId, - }); + await turnEventStream.put(event, { streamTTLSeconds: streamTTLSecondsFor(event) }); if (event.type === EventType.TURN_UPDATE) { if (event.state.status === 'paused') { pausedKeepalive.start(); @@ -574,11 +569,7 @@ export async function drainTurnEvents(input: { await onEvent?.(event); } } catch (error) { - if ( - error instanceof SessionStoreNotFoundError || - error instanceof TurnOwnershipLostError || - error instanceof StreamOwnerMismatchError - ) { + if (error instanceof SessionStoreNotFoundError || error instanceof TurnOwnershipLostError) { logger.warn( 'Turn stream ended after session/turn was removed or ownership was lost', extractErrorLogFields(error), @@ -601,7 +592,6 @@ export interface TurnEventDrainInput { executionTimeout: TurnExecutionTimeout; logger: Logger; pausedKeepalive: TurnPausedKeepalive; - executorId: string; } /** @@ -725,7 +715,7 @@ export async function beginTurnExecution( turn, }); - // Held for the whole turn. Repair appends any missing log suffix with event-id stream ids. + // Held for the whole turn. A missing stream is restored from the durable log on resume. const turnEventStream = deps.eventSubscriptions.get(turnStreamId(tenantId, sessionId, turn.id)); return { @@ -743,7 +733,6 @@ export async function beginTurnExecution( deps.activeTurns.clearPausedKeepalive({ sessionId, turnId: turn.id }); }, }, - executorId: configuration.EXECUTOR_ID, }, }; } @@ -785,7 +774,7 @@ function cancellationReason(reason: unknown): CancellationReason { } /** - * Rebuild a paused turn this replica owns, repair the event stream from the + * Rebuild a paused turn this replica owns, restore a missing event stream from the * durable log, and drain in the background. Inbox replay happens inside rebuild. */ async function resumePausedTurn(input: { @@ -853,13 +842,12 @@ async function resumePausedTurn(input: { turnStreamId(input.session.tenant_id, input.session.session_id, turn.id), ); try { - await repairTurnEventLog({ + await restoreTurnEventLog({ sessionStore: input.deps.sessionStore, subscription: turnEventStream, sessionId: input.session.session_id, turnId: turn.id, streamTTLSeconds: configuration.TURN_STREAM_TTL_SECONDS, - producerExecutorId: configuration.EXECUTOR_ID, }); } catch (error) { input.deps.logger.warn('Paused turn stream resume failed', extractErrorLogFields(error)); @@ -899,7 +887,6 @@ async function resumePausedTurn(input: { turnEventStream, executionTimeout, logger: input.deps.logger, - executorId: configuration.EXECUTOR_ID, pausedKeepalive: { start: () => { input.deps.activeTurns.startPausedKeepalive({ @@ -2047,7 +2034,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps { } const key = activeTurnKey(input.sessionId, input.turnId); const timer = setTimeout(() => { - void this.withTurnLock(input, async () => { + this.withTurnLock(input, async () => { // A send may have cleared or replaced this timer while the callback waited for the lock. if (this.keepaliveTimers.get(key) !== timer) { return; @@ -310,6 +310,15 @@ export class ActiveTurnRegistry { return; } this.logger.info('Paused keepalive expired; turn is no longer paused'); + }).catch((error: unknown) => { + // A busy turn lock is transient. Restart the window rather than keep a spent timer. + if (error instanceof TurnLockAcquireTimeoutError && this.keepaliveTimers.get(key) === timer) { + this.keepaliveTimers.delete(key); + this.logger.warn('Paused keepalive could not take the turn lock; restarting', extractErrorLogFields(error)); + this.startPausedKeepalive(input); + return; + } + this.logger.error('Paused keepalive offload failed', extractErrorLogFields(error)); }); }, this.pausedKeepaliveMs); timer.unref(); diff --git a/packages/trueforge/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index 7652743575..b1f979c4a2 100644 --- a/packages/trueforge/src/runtime/event-subscription/inMemory.ts +++ b/packages/trueforge/src/runtime/event-subscription/inMemory.ts @@ -5,7 +5,6 @@ import { EventEmitter, once } from 'node:events'; import { StreamGoneError, - StreamOwnerMismatchError, SUBSCRIBE_STREAM_THRESHOLD_MS, type EventSubscription, type EventSubscriptionPollOptions, @@ -15,12 +14,11 @@ import { import { entryIdMilliseconds, EventIdSeek, - eventIdSeekCursor, eventIdTimestampMs, isEntryAfterCursor, isLegacySequenceEntryId, } from './entryId'; -import { applyEventLogRepair, type AppendAfterTailResult, type RepairLogOptions, type StreamTail } from './repair'; +import { applyEventLogRestore, type AppendAfterTailResult, type RestoreLogOptions } from './restore'; function ttlDeadline(streamTTLSeconds: number | undefined): number | undefined { return streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? Date.now() + streamTTLSeconds * 1_000 : undefined; @@ -34,19 +32,14 @@ interface StoredEvent { interface InMemoryStream { /** Appended in entry-id order. */ events: StoredEvent[]; - /** Undefined until the first put/repair sets a TTL. */ + /** Undefined until the first put/restore sets a TTL. */ expiresAtMs?: number | undefined; expiryTimer?: NodeJS.Timeout | undefined; } -/** Owns per-stream logs, owners, and TTL timers for this process. */ +/** Owns per-stream logs and TTL timers for this process. */ export class InMemoryEventStreamStore { private readonly streams = new Map>(); - /** - * Survives an explicit eviction so a stale producer cannot recreate a missing - * log. Expires with the stream TTL so idle turns do not accumulate owners. - */ - private readonly owners = new Map(); /** Rings per stream id on append and expiry so parked pollers wake instantly. */ private readonly changes = new EventEmitter(); @@ -55,36 +48,19 @@ export class InMemoryEventStreamStore { this.changes.setMaxListeners(0); } - claimOwner(input: { streamId: string; executorId: string; streamTTLSeconds: number | undefined }): void { - this.owners.set(input.streamId, { - executorId: input.executorId, - expiresAtMs: ttlDeadline(input.streamTTLSeconds), - }); - } - /** * Live put with a `{millisecond}-{n}` id. A millisecond older than the tail * reuses the tail millisecond. Identical payloads are kept. */ - appendEvent(input: { - streamId: string; - event: T; - executorId: string | undefined; - streamTTLSeconds: number | undefined; - }): 'ok' | 'owner_mismatch' { - const owner = this.liveOwner(input.streamId); - if (owner !== undefined && owner !== input.executorId) { - return 'owner_mismatch'; - } + appendEvent(input: { streamId: string; event: T; streamTTLSeconds: number | undefined }): void { const stream = this.getLiveStream(input.streamId) ?? { events: [] }; this.streams.set(input.streamId, stream); this.push(stream, input.event); this.refreshStreamTtl(input.streamId, input.streamTTLSeconds); this.changes.emit(input.streamId); - return 'ok'; } - /** Repair backfill: appends only while the tail is still `tailEntryId`. */ + /** Restore append: appends only while the tail is still `tailEntryId`. */ appendAfterTail(input: { streamId: string; tailEntryId: string | undefined; @@ -103,29 +79,25 @@ export class InMemoryEventStreamStore { this.changes.emit(input.streamId); const written = stream.events[stream.events.length - 1]; if (written === undefined) { - throw new Error('Stream repair append wrote no entry'); + throw new Error('Stream restore append wrote no entry'); } return { status: 'ok', tailEntryId: written.entryId }; } - readTail(streamId: string): StreamTail | undefined { + readTailEntryId(streamId: string): string | undefined { const stream = this.getLiveStream(streamId); - const last = stream?.events[stream.events.length - 1]; - if (last === undefined) { - return undefined; - } - return { entryId: last.entryId, eventId: last.event.id, payload: JSON.stringify(last.event) }; + return stream?.events[stream.events.length - 1]?.entryId; } firstEntryId(streamId: string): string | undefined { return this.getLiveStream(streamId)?.events[0]?.entryId; } - hasEventId(streamId: string, eventId: string): boolean { + hasEventId({ streamId, eventId }: { streamId: string; eventId: string }): boolean { return this.getLiveStream(streamId)?.events.some(stored => stored.event.id === eventId) ?? false; } - /** Drops the log and keeps the owner so a stale producer cannot recreate it. */ + /** Drops the log the way a TTL expiry does. */ evictStream(streamId: string): void { const stream = this.streams.get(streamId); if (stream === undefined) { @@ -134,16 +106,11 @@ export class InMemoryEventStreamStore { this.dropStream(streamId, stream); } - /** Refreshes TTL on a live stream and its owner. */ refreshStreamTtl(streamId: string, streamTTLSeconds: number | undefined): void { const deadline = ttlDeadline(streamTTLSeconds); if (deadline === undefined) { return; } - const owner = this.owners.get(streamId); - if (owner !== undefined) { - owner.expiresAtMs = deadline; - } const stream = this.getLiveStream(streamId); if (!stream) { return; @@ -160,7 +127,6 @@ export class InMemoryEventStreamStore { } if (stream.expiresAtMs !== undefined && stream.expiresAtMs <= Date.now()) { this.dropStream(streamId, stream); - this.dropOwnerIfExpired(streamId); return undefined; } return stream; @@ -201,25 +167,12 @@ export class InMemoryEventStreamStore { stream.expiryTimer = setTimeout( () => { this.dropStream(streamId, stream); - this.dropOwnerIfExpired(streamId); }, Math.max(0, stream.expiresAtMs - Date.now()), ); stream.expiryTimer.unref(); } - private liveOwner(streamId: string): string | undefined { - this.dropOwnerIfExpired(streamId); - return this.owners.get(streamId)?.executorId; - } - - private dropOwnerIfExpired(streamId: string): void { - const expiresAtMs = this.owners.get(streamId)?.expiresAtMs; - if (expiresAtMs !== undefined && expiresAtMs <= Date.now()) { - this.owners.delete(streamId); - } - } - private push(stream: InMemoryStream, event: T): void { const copy = structuredClone(event); stream.events.push({ entryId: this.nextEntryId(stream, eventIdTimestampMs(copy.id) ?? Date.now()), event: copy }); @@ -251,22 +204,10 @@ export class InMemoryEventSubscription implements Even private readonly streamId: string, ) {} - repairLog(options: RepairLogOptions): Promise { - const executorId = options.executorId; - return applyEventLogRepair({ + restoreLog(options: RestoreLogOptions): Promise { + return applyEventLogRestore({ readPages: options.readPages, - claimOwner: - executorId === undefined - ? undefined - : () => { - this.store.claimOwner({ - streamId: this.streamId, - executorId, - streamTTLSeconds: options.streamTTLSeconds, - }); - return Promise.resolve(); - }, - readTail: () => Promise.resolve(this.store.readTail(this.streamId)), + readTailEntryId: () => Promise.resolve(this.store.readTailEntryId(this.streamId)), appendAfterTail: batch => Promise.resolve( this.store.appendAfterTail({ @@ -279,25 +220,20 @@ export class InMemoryEventSubscription implements Even this.store.refreshStreamTtl(this.streamId, options.streamTTLSeconds); return Promise.resolve(); }, - refreshTtlOnUntouchedStream: executorId !== undefined, }); } put(event: T, options?: EventSubscriptionPutOptions): Promise { - const outcome = this.store.appendEvent({ + this.store.appendEvent({ streamId: this.streamId, event, - executorId: options?.executorId, streamTTLSeconds: options?.streamTTLSeconds, }); - if (outcome === 'owner_mismatch') { - return Promise.reject(new StreamOwnerMismatchError(this.streamId)); - } return Promise.resolve(); } hasEventId(eventId: string): Promise { - return Promise.resolve(this.store.hasEventId(this.streamId, eventId)); + return Promise.resolve(this.store.hasEventId({ streamId: this.streamId, eventId })); } firstEntryId(): Promise { @@ -377,7 +313,6 @@ export class InMemoryEventSubscription implements Even } const signal = options.signal; const seek = new EventIdSeek>(fromEventId !== undefined ? 'from' : 'after', eventId); - const readCursor = eventIdSeekCursor({ eventId, firstEntryId: this.store.firstEntryId(this.streamId) }); let index = 0; let scanning = true; for (;;) { @@ -415,9 +350,6 @@ export class InMemoryEventSubscription implements Even if (stored === undefined) { throw new Error(`Corrupt stream ${this.streamId}: missing event at index ${String(index - 1)}`); } - if (!isEntryAfterCursor(stored.entryId, readCursor)) { - continue; - } for (const event of seek.push(this.polled(stored))) { yield event; } diff --git a/packages/trueforge/src/runtime/event-subscription/index.ts b/packages/trueforge/src/runtime/event-subscription/index.ts index 0e5125ee36..b1725bc1d5 100644 --- a/packages/trueforge/src/runtime/event-subscription/index.ts +++ b/packages/trueforge/src/runtime/event-subscription/index.ts @@ -1,18 +1,13 @@ import type { RedisClient } from '@truefoundry/trueforge-core/request-reply'; import { InMemoryEventStreamStore, InMemoryEventSubscription } from './inMemory'; import { RedisEventSubscription } from './redis'; -import type { RepairLogOptions } from './repair'; +import type { RestoreLogOptions } from './restore'; -export type { EventLogEntry, RepairLogOptions } from './repair'; +export type { RestoreLogOptions } from './restore'; export interface EventSubscriptionPutOptions { /** Sets the lifetime of the whole stream after this append. */ streamTTLSeconds?: number | undefined; - /** - * The append runs only when the stream-owner key is unset or equals this id. - * A missing stream still rejects a caller that does not own the key. - */ - executorId?: string | undefined; } /** `sequence_number` is set only when replaying an old `{sequence}-1` entry. */ @@ -53,11 +48,10 @@ export type EventSubscriptionPollOptions = { */ export interface EventSubscription { /** - * Rebuilds a missing or behind stream from the durable log using event-id - * stream ids. An existing stream whose tail is already the log tail is left - * in place. Pass `executorId` only for the database owner that will put. + * Rebuilds a missing stream from the whole durable log using event-id stream + * ids. An existing stream is left in place and only gets its TTL refreshed. */ - repairLog(options: RepairLogOptions): Promise; + restoreLog(options: RestoreLogOptions): Promise; /** Appends an event. The stream id is `{ulidMillisecond}-{n}` when the event id is a ULID. */ put(event: T, options?: EventSubscriptionPutOptions): Promise; @@ -115,14 +109,3 @@ export class StreamGoneError extends Error { Object.setPrototypeOf(this, new.target.prototype); } } - -/** put refused because the stream-owner key belongs to another executor. */ -export class StreamOwnerMismatchError extends Error { - readonly code = 'STREAM_OWNER_MISMATCH' as const; - - constructor(readonly streamId: string) { - super(`Cannot append to stream ${streamId}: caller is not the stream owner`); - this.name = 'StreamOwnerMismatchError'; - Object.setPrototypeOf(this, new.target.prototype); - } -} diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index 25e6c063c5..b266295158 100644 --- a/packages/trueforge/src/runtime/event-subscription/redis.ts +++ b/packages/trueforge/src/runtime/event-subscription/redis.ts @@ -2,7 +2,6 @@ import type { RedisClient } from '@truefoundry/trueforge-core/request-reply'; import { setTimeout as sleep } from 'node:timers/promises'; import { StreamGoneError, - StreamOwnerMismatchError, SUBSCRIBE_STREAM_THRESHOLD_MS, type EventSubscription, type EventSubscriptionPollOptions, @@ -10,116 +9,105 @@ import { type SequencedEvent, } from '.'; import { entryIdMilliseconds, eventIdSeekCursor, eventIdTimestampMs, isLegacySequenceEntryId } from './entryId'; -import { - applyEventLogRepair, - eventIdFromStreamData, - type AppendAfterTailResult, - type RepairLogOptions, - type StreamTail, -} from './repair'; +import { applyEventLogRestore, type AppendAfterTailResult, type RestoreLogOptions } from './restore'; const SUBSCRIBE_STREAM_POLL_ITEMS_COUNT = 100; /** Delay between XREAD polls while a stream has no new events (a blocking XREAD would monopolize the shared connection). */ const SUBSCRIBE_STREAM_POLL_SLEEP_INTERVAL_MS = 1_000; /** - * Jump from ARGV[3], walk forward, return an exclusive XREAD cursor. - * TS does the batch XREADs separately. + * One bounded step of an event-id seek. The TS side loops, so Redis serves + * other clients between batches. * - * KEYS[1]=stream ARGV[1]=from|after ARGV[2]=payload id ARGV[3]=start ARGV[4]=batch + * KEYS[1]=stream ARGV[1]=from|after ARGV[2]=payload id ARGV[3]=exclusive start + * ARGV[4]=batch ARGV[5]=last match so far ('' = none) ARGV[6]='1' on the first step * - * Returns {cursor, found}. - * from → stop at the first match; return the prior entry (XREAD is exclusive). - * after → walk the contiguous shared-id group (message + deltas); return its last entry. - * missing → tip (live-only), found = 0. + * Returns {state, cursor, lastMatch}. `cursor` is an exclusive XREAD cursor. + * more → keep going from `cursor` carrying `lastMatch`. + * found → from: the entry before the first match. after: the last entry with the id. + * missing → the last entry seen (live-only resume). */ -const SEEK_EVENT_ID_CURSOR = ` +const SEEK_EVENT_ID_STEP = ` local mode, targetId, cursor, batch = ARGV[1], ARGV[2], ARGV[3], tonumber(ARGV[4]) -local startBound = '(' .. cursor -local previous, lastMatch = nil, nil -while true do - local entries = redis.call('XRANGE', KEYS[1], startBound, '+', 'COUNT', batch) - if #entries == 0 then +local lastMatch = ARGV[5] +local entries = redis.call('XRANGE', KEYS[1], '(' .. cursor, '+', 'COUNT', batch) +if #entries == 0 then + if ARGV[6] == '1' then -- Start was past the tip (e.g. a future ULID ms): park at the real tip so -- XREAD still sees later live appends instead of waiting above them. - if previous == nil then - local tip = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) - if tip[1] ~= nil then - return {tip[1][1], 0} - end + local tip = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) + if tip[1] ~= nil then + return {'missing', tip[1][1], ''} end - break end - for i = 1, #entries do - local entryId, fields = entries[i][1], entries[i][2] - local data = nil - for j = 1, #fields, 2 do - if fields[j] == 'data' then data = fields[j + 1]; break end - end - local payloadId = nil - if data ~= nil then - local ok, obj = pcall(cjson.decode, data) - if ok and type(obj) == 'table' and type(obj['id']) == 'string' then - payloadId = obj['id'] - end + if lastMatch ~= '' then + return {'found', lastMatch, ''} + end + return {'missing', cursor, ''} +end +local previous = cursor +for i = 1, #entries do + local entryId, fields = entries[i][1], entries[i][2] + local payloadId, data = nil, nil + for j = 1, #fields, 2 do + if fields[j] == 'id' then payloadId = fields[j + 1]; break end + if fields[j] == 'data' then data = fields[j + 1] end + end + -- Old entries have no id field; decoding is limited to those. + if payloadId == nil and data ~= nil then + local ok, obj = pcall(cjson.decode, data) + if ok and type(obj) == 'table' and type(obj['id']) == 'string' then + payloadId = obj['id'] end + end + if payloadId == targetId then if mode == 'from' then - if payloadId == targetId then - return {previous or cursor, 1} - end - elseif payloadId == targetId then - lastMatch = entryId - elseif lastMatch ~= nil then - return {lastMatch, 1} + return {'found', previous, ''} end - previous = entryId + lastMatch = entryId end - startBound = '(' .. entries[#entries][1] - if #entries < batch then break end + previous = entryId end -if lastMatch ~= nil then - return {lastMatch, 1} +if #entries < batch then + if lastMatch ~= '' then + return {'found', lastMatch, ''} + end + return {'missing', previous, ''} end -return {previous or cursor, 0} +return {'more', previous, lastMatch} `; /** - * Live put. KEYS[1]=owner KEYS[2]=stream ARGV[1]=executor ARGV[2]=ms ARGV[3]=payload ARGV[4]=ttl + * Live put. KEYS[1]=stream ARGV[1]=ms ARGV[2]=event id ARGV[3]=payload ARGV[4]=ttl * - * Rejects when another executor owns the stream. The entry id is `{ms}-*`, raised - * to the tail millisecond when the event id is older than the tail (a delta of a - * message whose id was minted before later events). Identical payloads are kept: - * two deltas of the same message can carry the same text. + * The entry id is `{ms}-*`, raised to the tail millisecond when the event id is + * older than the tail (a delta of a message whose id was minted before later + * events); XADD rejects a smaller id. Identical payloads are kept: two deltas of + * the same message can carry the same text. The event id is its own field so a + * seek compares strings instead of decoding payloads on the Redis thread. */ const PUT_STREAM_ENTRY = ` -local owner = redis.call('GET', KEYS[1]) -if owner and owner ~= ARGV[1] then - return 'owner_mismatch' -end -local ms = ARGV[2] -local tail = redis.call('XREVRANGE', KEYS[2], '+', '-', 'COUNT', 1) +local ms = ARGV[1] +local tail = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) if tail[1] ~= nil then local tailMs = string.match(tail[1][1], '^%d+') if tonumber(tailMs) > tonumber(ms) then ms = tailMs end end -redis.call('XADD', KEYS[2], ms .. '-*', 'data', ARGV[3]) +redis.call('XADD', KEYS[1], ms .. '-*', 'id', ARGV[2], 'data', ARGV[3]) local ttl = tonumber(ARGV[4]) if ttl > 0 then - redis.call('EXPIRE', KEYS[2], ttl) - if owner then - redis.call('EXPIRE', KEYS[1], ttl) - end + redis.call('EXPIRE', KEYS[1], ttl) end return 'ok' `; /** - * Repair backfill. KEYS[1]=stream ARGV[1]=expected tail id ('' = empty) ARGV[2]=ttl - * ARGV[3..]=ms, payload pairs. + * Restore append. KEYS[1]=stream ARGV[1]=expected tail id ('' = empty) ARGV[2]=ttl + * ARGV[3..]=ms, event id, payload triples. * - * Appends only while the tail is still the one the suffix was computed from. + * Appends only while the tail is still the expected one. * Returns {'ok', lastEntryId} or 'stale'. The caller uses lastEntryId as the * expected tail of the next page. */ @@ -134,12 +122,12 @@ if tailId ~= ARGV[1] then end local tailMs = string.match(tailId, '^%d+') or '0' local lastAdded = nil -for i = 3, #ARGV, 2 do +for i = 3, #ARGV, 3 do local ms = ARGV[i] if tonumber(tailMs) > tonumber(ms) then ms = tailMs end - lastAdded = redis.call('XADD', KEYS[1], ms .. '-*', 'data', ARGV[i + 1]) + lastAdded = redis.call('XADD', KEYS[1], ms .. '-*', 'id', ARGV[i + 1], 'data', ARGV[i + 2]) tailMs = string.match(lastAdded, '^%d+') end local ttl = tonumber(ARGV[2]) @@ -152,10 +140,6 @@ end return {'ok', lastAdded} `; -function streamOwnerKey(streamId: string): string { - return `${streamId}:owner`; -} - function ttlArgument(streamTTLSeconds: number | undefined): string { return String(streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? streamTTLSeconds : 0); } @@ -171,55 +155,57 @@ function parseAppendReply(value: unknown, streamId: string): AppendAfterTailResu return { status: 'ok', tailEntryId }; } } - throw new Error(`Unexpected stream repair reply on ${streamId}`); + throw new Error(`Unexpected stream restore reply on ${streamId}`); } -function parseSeekReply(value: unknown): { cursor: string; found: boolean } { - if (Array.isArray(value) && value.length === 2) { - const cursor: unknown = value[0]; - const found: unknown = value[1]; - if (typeof cursor === 'string' && (found === 0 || found === 1)) { - return { cursor, found: found === 1 }; +interface SeekStep { + state: 'more' | 'found' | 'missing'; + cursor: string; + lastMatch: string; +} + +function parseSeekReply(value: unknown): SeekStep { + if (Array.isArray(value) && value.length === 3) { + const state: unknown = value[0]; + const cursor: unknown = value[1]; + const lastMatch: unknown = value[2]; + if ( + (state === 'more' || state === 'found' || state === 'missing') && + typeof cursor === 'string' && + typeof lastMatch === 'string' + ) { + return { state, cursor, lastMatch }; } } throw new Error('Unexpected event-id seek reply'); } -/** - * One Redis Stream. Entry ids are `{ulidMillisecond}-{n}` from `XADD`. The owner - * key is separate from the stream so a deleted stream still rejects a stale producer. - */ +/** One Redis Stream. Entry ids are `{ulidMillisecond}-{n}` from `XADD`. */ export class RedisEventSubscription implements EventSubscription { constructor( private readonly redis: RedisClient, private readonly streamId: string, ) {} - async repairLog(options: RepairLogOptions): Promise { - const executorId = options.executorId; - await applyEventLogRepair({ + async restoreLog(options: RestoreLogOptions): Promise { + await applyEventLogRestore({ readPages: options.readPages, - claimOwner: executorId === undefined ? undefined : () => this.claimOwner(executorId, options.streamTTLSeconds), - readTail: () => this.readTail(), + readTailEntryId: () => this.readTailEntryId(), appendAfterTail: batch => this.appendAfterTail({ ...batch, streamTTLSeconds: options.streamTTLSeconds }), refreshTtl: () => this.refreshTtl(options.streamTTLSeconds), - refreshTtlOnUntouchedStream: executorId !== undefined, }); } async put(event: T, options?: EventSubscriptionPutOptions): Promise { const reply: unknown = await this.redis.eval(PUT_STREAM_ENTRY, { - keys: [streamOwnerKey(this.streamId), this.streamId], + keys: [this.streamId], arguments: [ - options?.executorId ?? '', this.entryMillisecondFor(event), + event.id, JSON.stringify(event), ttlArgument(options?.streamTTLSeconds), ], }); - if (reply === 'owner_mismatch') { - throw new StreamOwnerMismatchError(this.streamId); - } if (reply !== 'ok') { throw new Error(`Unexpected stream put reply on ${this.streamId}`); } @@ -285,8 +271,8 @@ export class RedisEventSubscription implements EventSu } /** - * One Lua seek for the XREAD cursor, then the same forever-poll as sequence - * resume. Seek and batch reads are separate Redis round-trips on purpose. + * A seek for the XREAD cursor, then the same forever-poll as sequence + * resume. Seek steps and batch reads are separate Redis round-trips on purpose. */ private async *pollByEventId( options: EventSubscriptionPollOptions, @@ -299,7 +285,7 @@ export class RedisEventSubscription implements EventSu } const mode = fromEventId !== undefined ? 'from' : 'after'; const signal = options.signal; - let { cursor } = await this.seekEventIdCursor({ mode, eventId }); + let { cursor } = await this.seekEventIdCursor({ mode, eventId, signal }); for (;;) { const messages = await this.readAfter(cursor, signal); if (messages === undefined) { @@ -324,22 +310,42 @@ export class RedisEventSubscription implements EventSu } /** - * Jump near the ULID millisecond (or stream start), walk forward in Lua until - * the payload id, return exclusive XREAD cursor. See SEEK_EVENT_ID_CURSOR. + * Jump near the ULID millisecond (or stream start), then walk forward one + * bounded Lua step at a time until the seek settles. Returns an exclusive + * XREAD cursor. See SEEK_EVENT_ID_STEP. An aborted signal ends the walk early. */ private async seekEventIdCursor(input: { mode: 'from' | 'after'; eventId: string; + signal?: AbortSignal | undefined; }): Promise<{ cursor: string; found: boolean }> { - const startCursor = eventIdSeekCursor({ + let cursor = eventIdSeekCursor({ eventId: input.eventId, firstEntryId: await this.firstEntryId(), }); - const reply: unknown = await this.redis.eval(SEEK_EVENT_ID_CURSOR, { - keys: [this.streamId], - arguments: [input.mode, input.eventId, startCursor, String(SUBSCRIBE_STREAM_POLL_ITEMS_COUNT)], - }); - return parseSeekReply(reply); + let lastMatch = ''; + let first = true; + while (input.signal?.aborted !== true) { + const reply: unknown = await this.redis.eval(SEEK_EVENT_ID_STEP, { + keys: [this.streamId], + arguments: [ + input.mode, + input.eventId, + cursor, + String(SUBSCRIBE_STREAM_POLL_ITEMS_COUNT), + lastMatch, + first ? '1' : '0', + ], + }); + const step = parseSeekReply(reply); + if (step.state !== 'more') { + return { cursor: step.cursor, found: step.state === 'found' }; + } + cursor = step.cursor; + lastMatch = step.lastMatch; + first = false; + } + return { cursor, found: false }; } private async readAfter( @@ -405,40 +411,14 @@ export class RedisEventSubscription implements EventSu ); } - private async readTail(): Promise { + private async readTailEntryId(): Promise { const entries = await this.redis.xRevRange(this.streamId, '+', '-', { COUNT: 1 }); - const latest = entries[0]; - if (!latest) { - return undefined; - } - const data = latest.message['data']; - if (data === undefined) { - throw new Error(`Corrupt stream entry ${latest.id} on ${this.streamId}: missing data field`); - } - const eventId = eventIdFromStreamData(data); - if (eventId === undefined) { - throw new Error(`Corrupt stream entry ${latest.id} on ${this.streamId}: missing event id`); - } - return { entryId: latest.id, eventId, payload: data }; - } - - private async claimOwner(executorId: string, streamTTLSeconds: number | undefined): Promise { - const key = streamOwnerKey(this.streamId); - const multi = this.redis.multi().set(key, executorId); - if (streamTTLSeconds !== undefined && streamTTLSeconds > 0) { - multi.expire(key, streamTTLSeconds); - } - await multi.exec(); + return entries[0]?.id; } - /** The owner key shares the stream lifetime so it does not outlive an idle turn. */ private async refreshTtl(streamTTLSeconds: number | undefined): Promise { if (streamTTLSeconds !== undefined && streamTTLSeconds > 0) { - await this.redis - .multi() - .expire(this.streamId, streamTTLSeconds) - .expire(streamOwnerKey(this.streamId), streamTTLSeconds) - .exec(); + await this.redis.expire(this.streamId, streamTTLSeconds); } } @@ -447,10 +427,10 @@ export class RedisEventSubscription implements EventSu events: T[]; streamTTLSeconds: number | undefined; }): Promise { - const pairs = input.events.flatMap(event => [this.entryMillisecondFor(event), JSON.stringify(event)]); + const triples = input.events.flatMap(event => [this.entryMillisecondFor(event), event.id, JSON.stringify(event)]); const reply: unknown = await this.redis.eval(APPEND_AFTER_TAIL, { keys: [this.streamId], - arguments: [input.tailEntryId ?? '', ttlArgument(input.streamTTLSeconds), ...pairs], + arguments: [input.tailEntryId ?? '', ttlArgument(input.streamTTLSeconds), ...triples], }); return parseAppendReply(reply, this.streamId); } diff --git a/packages/trueforge/src/runtime/event-subscription/repair.ts b/packages/trueforge/src/runtime/event-subscription/repair.ts deleted file mode 100644 index b4ccf0d138..0000000000 --- a/packages/trueforge/src/runtime/event-subscription/repair.ts +++ /dev/null @@ -1,213 +0,0 @@ -/** - * Rebuild a turn's live event stream from the durable log. - * Append uses the same event-id stream ids as a live put. The durable row - * order is the order. No sequence number is assigned. - */ - -export interface EventLogEntry { - /** Payload id. The Redis stream id is not this value. */ - id: string; - event: T; -} - -export interface RepairLogOptions { - /** - * Durable rows in persist order, one page at a time. Called again when a - * concurrent append moves the tail, so the walk restarts from the first page. - */ - readPages: () => AsyncIterable[]>; - streamTTLSeconds?: number | undefined; - /** - * Set by the database owner that is about to produce. Claims the stream-owner - * key. A subscribe repair omits this and does not claim. - */ - executorId?: string | undefined; -} - -export interface StreamTail { - entryId: string; - eventId: string; - /** JSON written on the stream, used to tell shared-id rows apart. */ - payload: string; -} - -/** `'stale'` means the tail moved and the caller must recompute. */ -export type AppendAfterTailResult = { status: 'ok'; tailEntryId: string } | { status: 'stale' }; - -/** One page, for a caller that already holds the rows. Each call starts a new walk. */ -export function eventLogPages( - entries: readonly EventLogEntry[], -): () => AsyncIterable[]> { - return () => - (async function* () { - await Promise.resolve(); - yield entries; - })(); -} - -/** Backfill gives up after this many concurrent tail moves. */ -const REPAIR_MAX_ATTEMPTS = 5; - -function isPlainObject(value: unknown): value is Record { - return typeof value === 'object' && value !== null && !Array.isArray(value); -} - -/** Key order and a legacy `sequence_number` field are not part of event identity. */ -function canonicalize(value: unknown): unknown { - if (Array.isArray(value)) { - return value.map(item => canonicalize(item)); - } - if (!isPlainObject(value)) { - return value; - } - const normalized: Record = {}; - for (const key of Object.keys(value).sort()) { - if (key === 'sequence_number') { - continue; - } - const child = value[key]; - if (child !== undefined) { - normalized[key] = canonicalize(child); - } - } - return normalized; -} - -export function eventPayloadMatches(event: object, payload: string): boolean { - let parsed: unknown; - try { - parsed = JSON.parse(payload); - } catch { - return false; - } - return JSON.stringify(canonicalize(event)) === JSON.stringify(canonicalize(parsed)); -} - -/** Payload id from a stream data field. Undefined when JSON is corrupt or id is missing. */ -export function eventIdFromStreamData(data: string): string | undefined { - let parsed: unknown; - try { - parsed = JSON.parse(data); - } catch { - return undefined; - } - if (typeof parsed !== 'object' || parsed === null || !('id' in parsed)) { - return undefined; - } - return typeof parsed.id === 'string' ? parsed.id : undefined; -} - -/** - * Rows after the stream tail. Rows that share an id (a message and its deltas) - * are distinguished by the stored payload. A tail the log does not contain is - * left alone. - */ -function suffixAfterTail( - entries: readonly EventLogEntry[], - tail: StreamTail | undefined, -): EventLogEntry[] { - if (tail === undefined) { - return [...entries]; - } - // Prefer an exact payload match. The live tip may be a model.message.delta - // that was never durable; fall back to the last durable row with the same id - // (the message shell) so later durable events still append. - let match = -1; - for (let index = 0; index < entries.length; index += 1) { - const entry = entries[index]; - if (entry?.id === tail.eventId && eventPayloadMatches(entry.event, tail.payload)) { - match = index; - } - } - if (match === -1) { - for (let index = entries.length - 1; index >= 0; index -= 1) { - if (entries[index]?.id === tail.eventId) { - match = index; - break; - } - } - } - if (match === -1) { - return []; - } - return entries.slice(match + 1); -} - -/** - * Appends the missing suffix only while the stream tail is still the one it was - * computed from. Pages stay in persist order: an empty stream is filled from - * the start, and an existing stream appends only after its tail. A concurrent - * writer moves the tail; the walk restarts, so stale rows never land after newer ones. - */ -export async function applyEventLogRepair(input: { - readPages: () => AsyncIterable[]>; - claimOwner?: (() => Promise) | undefined; - readTail: () => Promise; - appendAfterTail: (batch: { tailEntryId: string | undefined; events: T[] }) => Promise; - refreshTtl: () => Promise; - /** Producer resume refreshes TTL even when the existing stream is already complete. */ - refreshTtlOnUntouchedStream: boolean; -}): Promise { - if (input.claimOwner !== undefined) { - await input.claimOwner(); - } - for (let attempt = 0; attempt < REPAIR_MAX_ATTEMPTS; attempt += 1) { - const tail = await input.readTail(); - const outcome = await appendMissingPages({ - readPages: input.readPages, - tail, - appendAfterTail: input.appendAfterTail, - }); - if (outcome === 'stale') { - continue; - } - if (outcome === 'untouched' && input.refreshTtlOnUntouchedStream) { - await input.refreshTtl(); - } - return; - } - throw new Error(`Stream repair gave up after ${String(REPAIR_MAX_ATTEMPTS)} concurrent tail changes`); -} - -/** - * One page is in memory at a time. The next page's expected tail is the entry - * id the previous append wrote, so a producer cannot land between the two. - */ -async function appendMissingPages(input: { - readPages: () => AsyncIterable[]>; - tail: StreamTail | undefined; - appendAfterTail: (batch: { tailEntryId: string | undefined; events: T[] }) => Promise; -}): Promise<'appended' | 'untouched' | 'stale'> { - let expectedTailId = input.tail?.entryId; - let passedTail = input.tail === undefined; - let appended = false; - for await (const page of input.readPages()) { - let rows: readonly EventLogEntry[]; - if (!passedTail) { - const tail = input.tail; - if (tail === undefined || !page.some(entry => entry.id === tail.eventId)) { - continue; - } - passedTail = true; - rows = suffixAfterTail(page, tail); - } else { - rows = page; - } - if (rows.length === 0) { - continue; - } - const outcome = await input.appendAfterTail({ - tailEntryId: expectedTailId, - events: rows.map(entry => entry.event), - }); - if (outcome.status === 'stale') { - return 'stale'; - } - expectedTailId = outcome.tailEntryId; - appended = true; - } - if (!passedTail) { - return 'untouched'; - } - return appended ? 'appended' : 'untouched'; -} diff --git a/packages/trueforge/src/runtime/event-subscription/restore.ts b/packages/trueforge/src/runtime/event-subscription/restore.ts new file mode 100644 index 0000000000..4d9c442c24 --- /dev/null +++ b/packages/trueforge/src/runtime/event-subscription/restore.ts @@ -0,0 +1,45 @@ +/** + * Rebuild a turn's live event stream from the durable log, only when the stream + * is gone. Durable persist order can differ from emit order, so a stream that + * still exists is never patched. + */ + +export interface RestoreLogOptions { + /** Durable events, one page at a time. Appended in the order read. */ + readPages: () => AsyncIterable; + streamTTLSeconds?: number | undefined; +} + +/** `'stale'` means the tail moved: another writer created or extended the stream. */ +export type AppendAfterTailResult = { status: 'ok'; tailEntryId: string } | { status: 'stale' }; + +/** + * Appends every page when the stream has no entries. A stream that already has + * entries is left as is and only gets its TTL refreshed, so a near-expiry + * stream stays subscribable. The next page's expected tail is the entry id the + * previous append wrote, so a concurrent writer stops the walk instead of + * landing between pages. + */ +export async function applyEventLogRestore(input: { + readPages: () => AsyncIterable; + readTailEntryId: () => Promise; + appendAfterTail: (batch: { tailEntryId: string | undefined; events: T[] }) => Promise; + refreshTtl: () => Promise; +}): Promise { + if ((await input.readTailEntryId()) !== undefined) { + await input.refreshTtl(); + return; + } + let expectedTailId: string | undefined; + for await (const page of input.readPages()) { + if (page.length === 0) { + continue; + } + const outcome = await input.appendAfterTail({ tailEntryId: expectedTailId, events: [...page] }); + if (outcome.status === 'stale') { + await input.refreshTtl(); + return; + } + expectedTailId = outcome.tailEntryId; + } +} diff --git a/packages/trueforge/src/runtime/turnEventLogRepair.ts b/packages/trueforge/src/runtime/turnEventLogRestore.ts similarity index 55% rename from packages/trueforge/src/runtime/turnEventLogRepair.ts rename to packages/trueforge/src/runtime/turnEventLogRestore.ts index 39604c8cdd..6e239a12d5 100644 --- a/packages/trueforge/src/runtime/turnEventLogRepair.ts +++ b/packages/trueforge/src/runtime/turnEventLogRestore.ts @@ -1,44 +1,42 @@ import type { ISessionStore, TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session'; -import { EVENTS_PAGE_LIMIT } from '../schemas/common'; import type { EventSubscription } from './event-subscription'; +/** Large pages keep a full restore to few round trips. */ +const RESTORE_PAGE_LIMIT = 1_000; + /** - * Rebuild a paused turn's live stream from the durable log. - * Subscribe omits `producerExecutorId` and only appends. - * The producing resume, which already owns the turn row, claims the stream-owner key. - * Pages are read ascending so an empty stream is appended from start to end. + * Restore a paused turn's live stream from the durable log when the stream is gone. + * An existing stream is never patched: persist order can differ from emit order. */ -export async function repairTurnEventLog(input: { +export async function restoreTurnEventLog(input: { sessionStore: Pick, 'listTurnEvents'>; subscription: EventSubscription; sessionId: string; turnId: string; streamTTLSeconds: number; - producerExecutorId?: string | undefined; }): Promise { - await input.subscription.repairLog({ + await input.subscription.restoreLog({ readPages: () => durableEventPages(input), streamTTLSeconds: input.streamTTLSeconds, - ...(input.producerExecutorId === undefined ? {} : { executorId: input.producerExecutorId }), }); } -/** One page in memory. `order: 'asc'` is persist order, which is also append order. */ +/** One page in memory, in event-id order. */ async function* durableEventPages(input: { sessionStore: Pick, 'listTurnEvents'>; sessionId: string; turnId: string; -}): AsyncGenerator<{ id: string; event: TurnStreamingEvent }[]> { +}): AsyncGenerator { let pageToken: string | undefined; for (;;) { const page = await input.sessionStore.listTurnEvents({ session_id: input.sessionId, turn_id: input.turnId, - limit: EVENTS_PAGE_LIMIT, + limit: RESTORE_PAGE_LIMIT, page_token: pageToken, order: 'asc', }); - yield page.data.map(event => ({ id: event.id, event })); + yield page.data; const nextPageToken = page.pagination.next_page_token; if (nextPageToken === undefined) { return; diff --git a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts index a5b6948e33..99d78cc855 100644 --- a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts +++ b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts @@ -1,12 +1,9 @@ import { StreamGoneError, - StreamOwnerMismatchError, type EventSubscriptionPollOptions, type EventSubscriptionRegistry, type SequencedEvent, } from '../../../src/runtime/event-subscription'; -import { eventLogPages } from '../../../src/runtime/event-subscription/repair'; - /** Payload shape used by the suite; distinct per event so round-trips are provable. */ export interface TestEvent { id: string; @@ -27,6 +24,14 @@ function makeEvent(index: number, id = `evt-${String(index)}`): TestEvent { return { id, type: 'test.event', payload: { index, text: `event-${String(index)}` } }; } +/** One durable page per array, in the order given. */ +async function* durablePages(...pages: TestEvent[][]): AsyncGenerator { + await Promise.resolve(); + for (const page of pages) { + yield page; + } +} + /** Drains `count` events from a fresh poll, then aborts it (poll loops forever by design). */ async function collect( registry: EventSubscriptionRegistry, @@ -55,7 +60,7 @@ async function collect( export function runEventSubscriptionContractSuite( setup: () => { registry: EventSubscriptionRegistry; - /** Drops the log and keeps the owner, the way a TTL expiry does. */ + /** Drops the log the way a TTL expiry does. */ evictStream: (streamId: string) => Promise; }, ) { @@ -116,6 +121,35 @@ export function runEventSubscriptionContractSuite( expect(received.map(event => event.id)).toEqual(['evt-2', 'evt-3', 'evt-4']); }); + it('after an event id resumes past the last entry that carries it, even when other events interleave', async () => { + const { registry } = setup(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + const sent = [makeEvent(0, 'evt-a'), makeEvent(1, 'evt-b'), makeEvent(2, 'evt-a'), makeEvent(3, 'evt-c')]; + for (const event of sent) { + await producer.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }); + } + + const received = await collect(registry, streamId, 1, { afterEventId: 'evt-a' }); + expect(received.map(event => event.id)).toEqual(['evt-c']); + }); + + it('seeks an event id that sits beyond the first scan batch', async () => { + const { registry } = setup(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + const sent = Array.from({ length: 250 }, (_, index) => makeEvent(index)); + for (const event of sent) { + await producer.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }); + } + + await expect(producer.hasEventId('evt-230')).resolves.toBe(true); + const from = await collect(registry, streamId, 3, { fromEventId: 'evt-230' }); + expect(from.map(event => event.id)).toEqual(['evt-230', 'evt-231', 'evt-232']); + const after = await collect(registry, streamId, 2, { afterEventId: 'evt-247' }); + expect(after.map(event => event.id)).toEqual(['evt-248', 'evt-249']); + }); + it('keeps an event appended while an event-id replay is delivered', async () => { const { registry } = setup(); const streamId = nextStreamId(); @@ -159,28 +193,17 @@ export function runEventSubscriptionContractSuite( expect(received).toEqual([later, shell, delta]); }); - it('repairs a suffix when the durable row key order differs from the stream payload', async () => { + it('restores every page of a missing stream in order', async () => { const { registry } = setup(); const streamId = nextStreamId(); - const producer = registry.get(streamId); - const stored = makeEvent(0); - await producer.put(stored, { streamTTLSeconds: STREAM_TTL_SECONDS }); - const fromDatabase = { - id: 'evt-0', - type: stored.type, - payload: { text: stored.payload.text, index: stored.payload.index }, - }; - const third = makeEvent(2); - await producer.repairLog({ - readPages: eventLogPages([ - { id: 'evt-0', event: fromDatabase }, - { id: 'evt-2', event: third }, - ]), + const sent = Array.from({ length: 5 }, (_, index) => makeEvent(index)); + + await registry.get(streamId).restoreLog({ + readPages: () => durablePages(sent.slice(0, 2), [], sent.slice(2)), streamTTLSeconds: STREAM_TTL_SECONDS, }); - const received = await collect(registry, streamId, 2); - expect(received.map(event => event.id)).toEqual(['evt-0', 'evt-2']); + expect(await collect(registry, streamId, sent.length)).toEqual(sent); }); it('poll parked at the tip receives events put after it subscribed', async () => { @@ -207,68 +230,43 @@ export function runEventSubscriptionContractSuite( const producer = registry.get(streamId); const first = makeEvent(1); const second = makeEvent(2); - await producer.put(first, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); - await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); + await producer.put(first, { streamTTLSeconds: STREAM_TTL_SECONDS }); + await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS }); await evictStream(streamId); + await expect(producer.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); - await producer.repairLog({ - readPages: eventLogPages([ - { id: first.id, event: first }, - { id: second.id, event: second }, - ]), + await producer.restoreLog({ + readPages: () => durablePages([first, second]), streamTTLSeconds: STREAM_TTL_SECONDS, - executorId: 'owner', }); const third = makeEvent(3); - await producer.put(third, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); + await producer.put(third, { streamTTLSeconds: STREAM_TTL_SECONDS }); const received = await collect(registry, streamId, 3); expect(received).toEqual([first, second, third]); }); - it('leaves an existing stream untouched and a second repair does not append again', async () => { + it('leaves an existing stream untouched, even when the log has more rows', async () => { const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); const first = makeEvent(1); const second = makeEvent(2); + const durableOnly = makeEvent(3); await producer.put(first, { streamTTLSeconds: STREAM_TTL_SECONDS }); await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS }); - const repair = () => - producer.repairLog({ - readPages: eventLogPages([ - { id: first.id, event: first }, - { id: second.id, event: second }, - ]), + const restore = () => + producer.restoreLog({ + readPages: () => durablePages([first, second, durableOnly]), streamTTLSeconds: STREAM_TTL_SECONDS, - executorId: 'owner', }); - await repair(); - await repair(); + await restore(); + await restore(); - const third = makeEvent(3); - await producer.put(third, { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'owner' }); + const fourth = makeEvent(4); + await producer.put(fourth, { streamTTLSeconds: STREAM_TTL_SECONDS }); const received = await collect(registry, streamId, 3); - expect(received).toEqual([first, second, third]); - }); - - it('rejects a stale producer when the stream is missing', async () => { - const { registry, evictStream } = setup(); - const streamId = nextStreamId(); - const owner = registry.get(streamId); - const event = makeEvent(1); - await owner.repairLog({ - readPages: eventLogPages([{ id: event.id, event }]), - streamTTLSeconds: STREAM_TTL_SECONDS, - executorId: 'owner', - }); - await evictStream(streamId); - - const stale = registry.get(streamId); - await expect( - stale.put(makeEvent(9, 'evt-late'), { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'stale' }), - ).rejects.toBeInstanceOf(StreamOwnerMismatchError); - await expect(stale.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); + expect(received).toEqual([first, second, fourth]); }); it('keeps consecutive identical payloads, like two deltas with the same text', async () => { @@ -293,4 +291,18 @@ export function runEventSubscriptionContractSuite( await expect(producer.hasEventId('evt-1')).resolves.toBe(true); await expect(producer.hasEventId('evt-missing')).resolves.toBe(false); }); + + it('seeks past large payloads', async () => { + const { registry } = setup(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + const large: TestEvent = { id: 'evt-large', type: 'test.event', payload: { index: 0, text: 'x'.repeat(512_000) } }; + const target = makeEvent(1); + await producer.put(large, { streamTTLSeconds: STREAM_TTL_SECONDS }); + await producer.put(target, { streamTTLSeconds: STREAM_TTL_SECONDS }); + + await expect(producer.hasEventId(target.id)).resolves.toBe(true); + const received = await collect(registry, streamId, 1, { afterEventId: large.id }); + expect(received).toEqual([target]); + }); } diff --git a/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts b/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts index 7cc369616b..5ae41f56c6 100644 --- a/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts @@ -1,5 +1,4 @@ import { eventIdSeekCursor, eventIdTimestampMs } from '../../../src/runtime/event-subscription/entryId'; -import { eventLogPages, eventPayloadMatches } from '../../../src/runtime/event-subscription/repair'; const EVENT_ID = '01arz3ndektsv4rrffq69g5fav'; @@ -22,56 +21,4 @@ describe('event id cursors', () => { `${String(timestamp - 1)}-18446744073709551615`, ); }); - - it('matches a durable row when key order or a legacy sequence field differs', () => { - const event = { id: 'evt-1', type: 'model.message', payload: { index: 1, text: 'hi' } }; - expect(eventPayloadMatches(event, '{"payload":{"text":"hi","index":1},"type":"model.message","id":"evt-1"}')).toBe( - true, - ); - expect( - eventPayloadMatches( - event, - '{"id":"evt-1","type":"model.message","payload":{"index":1,"text":"hi"},"sequence_number":4}', - ), - ).toBe(true); - expect(eventPayloadMatches(event, '{"id":"evt-1","type":"model.message","payload":{"index":2,"text":"hi"}}')).toBe( - false, - ); - }); -}); - -describe('suffixAfterTail via repairLog', () => { - it('appends durable rows after a non-durable shared-id stream tip', async () => { - const { EventSubscriptionRegistry } = await import('../../../src/runtime/event-subscription'); - const registry = new EventSubscriptionRegistry<{ id: string; type: string; payload: { text: string } }>(undefined); - const streamId = `repair-delta-tip:${String(process.pid)}`; - const producer = registry.get(streamId); - const shell = { id: 'evt-1', type: 'model.message', payload: { text: 'shell' } }; - const deltaTip = { id: 'evt-1', type: 'model.message.delta', payload: { text: 'delta' } }; - const next = { id: 'evt-2', type: 'turn.done', payload: { text: 'done' } }; - await producer.put(shell, { streamTTLSeconds: 300 }); - await producer.put(deltaTip, { streamTTLSeconds: 300 }); - await producer.repairLog({ - readPages: eventLogPages([ - { id: 'evt-1', event: shell }, - { id: 'evt-2', event: next }, - ]), - streamTTLSeconds: 300, - }); - const abort = new AbortController(); - const generator = producer.poll({ signal: abort.signal }); - const received: string[] = []; - try { - for await (const event of generator) { - received.push(event.id); - if (received.length >= 3) { - break; - } - } - } finally { - abort.abort(); - await generator.return(undefined); - } - expect(received).toEqual(['evt-1', 'evt-1', 'evt-2']); - }); }); diff --git a/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts b/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts index 63df7d180c..60a196f91f 100644 --- a/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts @@ -28,4 +28,46 @@ describeRedis('RedisEventSubscription (EventSubscription contract)', () => { await client.del(streamId); }, })); + + describe('entry format', () => { + const streamId = `eventsub-format:${String(process.pid)}:${String(Date.now())}`; + + afterEach(async () => { + await client.del(streamId); + }); + + const event: TestEvent = { id: 'evt-1', type: 'test.event', payload: { index: 1, text: 'one' } }; + + it('put writes the event id as its own field next to the payload', async () => { + const registry = new EventSubscriptionRegistry(client); + await registry.get(streamId).put(event, { streamTTLSeconds: 60 }); + + const entries = await client.xRange(streamId, '-', '+'); + expect(entries.map(entry => entry.message)).toEqual([{ id: 'evt-1', data: JSON.stringify(event) }]); + }); + + it('restoreLog writes the event id as its own field next to the payload', async () => { + const registry = new EventSubscriptionRegistry(client); + await registry.get(streamId).restoreLog({ + readPages: async function* () { + await Promise.resolve(); + yield [event]; + }, + streamTTLSeconds: 60, + }); + + const entries = await client.xRange(streamId, '-', '+'); + expect(entries.map(entry => entry.message)).toEqual([{ id: 'evt-1', data: JSON.stringify(event) }]); + }); + + it('still finds an id in an entry that has no id field', async () => { + await client.xAdd(streamId, '1-1', { + data: JSON.stringify({ id: 'evt-old', type: 'test.event', payload: { index: 0, text: 'old' } }), + }); + + const registry = new EventSubscriptionRegistry(client); + await expect(registry.get(streamId).hasEventId('evt-old')).resolves.toBe(true); + await expect(registry.get(streamId).hasEventId('evt-other')).resolves.toBe(false); + }); + }); }); diff --git a/packages/trueforge/tests/runtime/event-subscription/repair.test.ts b/packages/trueforge/tests/runtime/event-subscription/repair.test.ts deleted file mode 100644 index 9d20e39022..0000000000 --- a/packages/trueforge/tests/runtime/event-subscription/repair.test.ts +++ /dev/null @@ -1,86 +0,0 @@ -import { applyEventLogRepair, eventLogPages, type StreamTail } from '../../../src/runtime/event-subscription/repair'; - -interface Row { - id: string; - type: string; -} - -function tailOf(row: Row, entryId: string): StreamTail { - return { entryId, eventId: row.id, payload: JSON.stringify(row) }; -} - -describe('applyEventLogRepair', () => { - it('recomputes the suffix when a concurrent writer moves the tail', async () => { - const created: Row = { id: 'evt-1', type: 'turn.created' }; - const paused: Row = { id: 'evt-2', type: 'turn.update' }; - const running: Row = { id: 'evt-3', type: 'turn.update' }; - const entries = [created, paused, running].map(event => ({ id: event.id, event })); - // The first read sees only `created`; by the append the producer has written through `running`. - const tails = [tailOf(created, '1-0'), tailOf(running, '3-0')]; - const batches: { tailEntryId: string | undefined; events: Row[] }[] = []; - - await applyEventLogRepair({ - readPages: eventLogPages(entries), - readTail: () => Promise.resolve(tails.shift()), - appendAfterTail: batch => { - batches.push(batch); - return Promise.resolve({ status: 'stale' }); - }, - refreshTtl: () => Promise.resolve(), - refreshTtlOnUntouchedStream: false, - }); - - expect(batches).toEqual([{ tailEntryId: '1-0', events: [paused, running] }]); - }); - - it('gives up when the tail keeps moving', async () => { - const created: Row = { id: 'evt-1', type: 'turn.created' }; - const paused: Row = { id: 'evt-2', type: 'turn.update' }; - await expect( - applyEventLogRepair({ - readPages: eventLogPages([created, paused].map(event => ({ id: event.id, event }))), - readTail: () => Promise.resolve(tailOf(created, '1-0')), - appendAfterTail: () => Promise.resolve({ status: 'stale' }), - refreshTtl: () => Promise.resolve(), - refreshTtlOnUntouchedStream: false, - }), - ).rejects.toThrow('Stream repair gave up'); - }); - - it('appends later pages after the entry id the previous page wrote', async () => { - const created: Row = { id: 'evt-1', type: 'turn.created' }; - const paused: Row = { id: 'evt-2', type: 'turn.update' }; - const running: Row = { id: 'evt-3', type: 'turn.update' }; - const done: Row = { id: 'evt-4', type: 'turn.done' }; - const batches: { tailEntryId: string | undefined; events: Row[] }[] = []; - const writtenTailIds = ['8-0', '9-0']; - - await applyEventLogRepair({ - readPages: () => - (async function* () { - yield [{ id: created.id, event: created }]; - yield [ - { id: paused.id, event: paused }, - { id: running.id, event: running }, - ]; - yield [{ id: done.id, event: done }]; - })(), - readTail: () => Promise.resolve(tailOf(paused, '2-0')), - appendAfterTail: batch => { - batches.push(batch); - const tailEntryId = writtenTailIds.shift(); - if (tailEntryId === undefined) { - return Promise.resolve({ status: 'stale' }); - } - return Promise.resolve({ status: 'ok', tailEntryId }); - }, - refreshTtl: () => Promise.resolve(), - refreshTtlOnUntouchedStream: false, - }); - - expect(batches).toEqual([ - { tailEntryId: '2-0', events: [running] }, - { tailEntryId: '8-0', events: [done] }, - ]); - }); -}); diff --git a/packages/trueforge/tests/runtime/event-subscription/restore.test.ts b/packages/trueforge/tests/runtime/event-subscription/restore.test.ts new file mode 100644 index 0000000000..3b54bd3150 --- /dev/null +++ b/packages/trueforge/tests/runtime/event-subscription/restore.test.ts @@ -0,0 +1,88 @@ +import { applyEventLogRestore } from '../../../src/runtime/event-subscription/restore'; + +interface Row { + id: string; + type: string; +} + +const created: Row = { id: 'evt-1', type: 'turn.created' }; +const paused: Row = { id: 'evt-2', type: 'turn.update' }; +const running: Row = { id: 'evt-3', type: 'turn.update' }; +const done: Row = { id: 'evt-4', type: 'turn.done' }; + +async function* pagesOf(...pages: Row[][]): AsyncGenerator { + await Promise.resolve(); + for (const page of pages) { + yield page; + } +} + +describe('applyEventLogRestore', () => { + it('appends every page of the log to an empty stream, chaining the expected tail', async () => { + const batches: { tailEntryId: string | undefined; events: Row[] }[] = []; + const writtenTailIds = ['2-0', '4-0']; + let refreshed = 0; + + await applyEventLogRestore({ + readPages: () => pagesOf([created, paused], [], [running, done]), + readTailEntryId: () => Promise.resolve(undefined), + appendAfterTail: batch => { + batches.push(batch); + const tailEntryId = writtenTailIds.shift(); + return Promise.resolve(tailEntryId === undefined ? { status: 'stale' } : { status: 'ok', tailEntryId }); + }, + refreshTtl: () => { + refreshed += 1; + return Promise.resolve(); + }, + }); + + expect(batches).toEqual([ + { tailEntryId: undefined, events: [created, paused] }, + { tailEntryId: '2-0', events: [running, done] }, + ]); + expect(refreshed).toBe(0); + }); + + it('leaves an existing stream alone and only refreshes its TTL', async () => { + let appended = 0; + let refreshed = 0; + + await applyEventLogRestore({ + readPages: () => pagesOf([created, paused, running]), + readTailEntryId: () => Promise.resolve('2-0'), + appendAfterTail: () => { + appended += 1; + return Promise.resolve({ status: 'stale' }); + }, + refreshTtl: () => { + refreshed += 1; + return Promise.resolve(); + }, + }); + + expect(appended).toBe(0); + expect(refreshed).toBe(1); + }); + + it('stops when a concurrent writer moves the tail and refreshes the TTL', async () => { + const batches: { tailEntryId: string | undefined; events: Row[] }[] = []; + let refreshed = 0; + + await applyEventLogRestore({ + readPages: () => pagesOf([created], [paused]), + readTailEntryId: () => Promise.resolve(undefined), + appendAfterTail: batch => { + batches.push(batch); + return Promise.resolve({ status: 'stale' }); + }, + refreshTtl: () => { + refreshed += 1; + return Promise.resolve(); + }, + }); + + expect(batches).toEqual([{ tailEntryId: undefined, events: [created] }]); + expect(refreshed).toBe(1); + }); +}); diff --git a/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRestore.test.ts similarity index 74% rename from packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts rename to packages/trueforge/tests/runtime/event-subscription/turnEventLogRestore.test.ts index 80136aaf7f..40080dbcf1 100644 --- a/packages/trueforge/tests/runtime/event-subscription/turnEventLogRepair.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRestore.test.ts @@ -10,11 +10,10 @@ import { import { EventSubscriptionRegistry, StreamGoneError, - StreamOwnerMismatchError, type SequencedEvent, } from '../../../src/runtime/event-subscription'; import { InMemoryEventStreamStore } from '../../../src/runtime/event-subscription/inMemory'; -import { repairTurnEventLog } from '../../../src/runtime/turnEventLogRepair'; +import { restoreTurnEventLog } from '../../../src/runtime/turnEventLogRestore'; const STREAM_TTL_SECONDS = 300; @@ -40,8 +39,8 @@ async function collect( return received; } -describe('repairTurnEventLog', () => { - it('rebuilds a missing stream from the durable log', async () => { +describe('restoreTurnEventLog', () => { + it('rebuilds a missing stream from the whole durable log and never patches an existing one', async () => { const store = new InMemorySessionStore(); const sessionId = 's1'; const turnId = 'turn-1'; @@ -69,26 +68,24 @@ describe('repairTurnEventLog', () => { const registry = new EventSubscriptionRegistry(undefined, streams); const streamId = 'turn-stream'; const subscription = registry.get(streamId); - const repair = (producerExecutorId?: string) => - repairTurnEventLog({ + const restore = () => + restoreTurnEventLog({ sessionStore: store, subscription, sessionId, turnId, streamTTLSeconds: STREAM_TTL_SECONDS, - ...(producerExecutorId === undefined ? {} : { producerExecutorId }), }); await expect(subscription.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); - await repair(); + await restore(); expect((await collect(registry, streamId, 2)).map(event => event.id)).toEqual([created.id, message.id]); - await repair(); + await restore(); expect((await collect(registry, streamId, 2)).map(event => event.id)).toEqual([created.id, message.id]); - await repair(TEST_ACTIVE_EXECUTOR_ID); - + // A durable row the stream lacks is not patched in: only a missing stream is restored. const later = makeModelMessageEvent(); await store.appendToEvents({ active_executor_id: TEST_ACTIVE_EXECUTOR_ID, @@ -96,22 +93,16 @@ describe('repairTurnEventLog', () => { turn_id: turnId, events: [later], }); - await repair(TEST_ACTIVE_EXECUTOR_ID); - expect((await collect(registry, streamId, 3)).map(event => event.id)).toEqual([created.id, message.id, later.id]); + await restore(); + await expect(subscription.hasEventId(later.id)).resolves.toBe(false); + // Once the stream is gone, the next restore rebuilds it from the whole log. streams.evictStream(streamId); - const stale = registry.get(streamId); - await expect( - stale.put(makeModelMessageEvent(), { streamTTLSeconds: STREAM_TTL_SECONDS, executorId: 'other-executor' }), - ).rejects.toBeInstanceOf(StreamOwnerMismatchError); - await expect(stale.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); + await expect(subscription.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); - await repair(TEST_ACTIVE_EXECUTOR_ID); + await restore(); const appended = makeModelMessageEvent(); - await subscription.put(appended, { - streamTTLSeconds: STREAM_TTL_SECONDS, - executorId: TEST_ACTIVE_EXECUTOR_ID, - }); + await subscription.put(appended, { streamTTLSeconds: STREAM_TTL_SECONDS }); expect((await collect(registry, streamId, 4)).map(event => event.id)).toEqual([ created.id, message.id, diff --git a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts index 780d029624..f414b1307f 100644 --- a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts +++ b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts @@ -282,6 +282,33 @@ describe('ActiveTurnRegistry', () => { await expect(waiter).resolves.toBe('acquired'); }); + it('restarts the paused keepalive when the turn lock is busy past the acquire timeout', async () => { + jest.useFakeTimers(); + const warn = jest.spyOn(logger, 'warn'); + const info = jest.spyOn(logger, 'info'); + try { + const registry = new ActiveTurnRegistry({ logger, pausedKeepaliveMs: 100, lockAcquireTimeoutMs: 20 }); + const key = { sessionId: 's1', turnId: 't1' }; + const gate = Promise.withResolvers(); + const holder = registry.withTurnLock(key, () => gate.promise); + + registry.startPausedKeepalive(key); + // The keepalive fires at 100ms, then gives up on the busy lock at 120ms. + await jest.advanceTimersByTimeAsync(120); + expect(warn).toHaveBeenCalledWith('Paused keepalive could not take the turn lock; restarting', expect.anything()); + + gate.resolve(undefined); + await holder; + // The restarted timer takes the free lock and runs its expiry path. + await jest.advanceTimersByTimeAsync(100); + expect(info).toHaveBeenCalledWith('Paused keepalive expired; turn is no longer paused'); + } finally { + jest.useRealTimers(); + warn.mockRestore(); + info.mockRestore(); + } + }); + it('waits without a timeout when lockAcquireTimeoutMs is 0', async () => { const registry = new ActiveTurnRegistry({ logger, lockAcquireTimeoutMs: 0 }); const gate = Promise.withResolvers(); From 7eddf03db56fd7c344d4b19d715e9f26ff1b2a28 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 5 Nov 2026 10:18:55 +0530 Subject: [PATCH 16/21] Install a restored turn stream atomically and seek after-cursors backward. A restore writes the durable log to a private key and renames it onto the stream only while the stream is still absent, so a stream is never left as a prefix. An after_event_id seek walks down from the tip, so a recent cursor settles in one step. Co-authored-by: Cursor --- .../runtime/event-subscription/inMemory.ts | 50 ++-- .../src/runtime/event-subscription/redis.ts | 268 ++++++++++-------- .../src/runtime/event-subscription/restore.ts | 28 +- .../event-subscription/contractSuite.ts | 33 +++ .../event-subscription/redis.contract.test.ts | 19 ++ .../event-subscription/restore.test.ts | 40 ++- 6 files changed, 244 insertions(+), 194 deletions(-) diff --git a/packages/trueforge/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index b1f979c4a2..89f7439f2c 100644 --- a/packages/trueforge/src/runtime/event-subscription/inMemory.ts +++ b/packages/trueforge/src/runtime/event-subscription/inMemory.ts @@ -18,7 +18,7 @@ import { isEntryAfterCursor, isLegacySequenceEntryId, } from './entryId'; -import { applyEventLogRestore, type AppendAfterTailResult, type RestoreLogOptions } from './restore'; +import { applyEventLogRestore, type RestoreLogOptions } from './restore'; function ttlDeadline(streamTTLSeconds: number | undefined): number | undefined { return streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? Date.now() + streamTTLSeconds * 1_000 : undefined; @@ -60,28 +60,25 @@ export class InMemoryEventStreamStore { this.changes.emit(input.streamId); } - /** Restore append: appends only while the tail is still `tailEntryId`. */ - appendAfterTail(input: { - streamId: string; - tailEntryId: string | undefined; - events: T[]; - streamTTLSeconds: number | undefined; - }): AppendAfterTailResult { - const stream = this.getLiveStream(input.streamId) ?? { events: [] }; - if (stream.events[stream.events.length - 1]?.entryId !== input.tailEntryId) { - return { status: 'stale' }; + /** + * Publishes a fully built log only while the stream is still absent. + * The check and the insert are synchronous, so a concurrent put cannot land in between. + */ + installIfAbsent(input: { streamId: string; events: readonly T[]; streamTTLSeconds: number | undefined }): boolean { + if (input.events.length === 0) { + return true; } - this.streams.set(input.streamId, stream); + const stream: InMemoryStream = { events: [] }; for (const event of input.events) { this.push(stream, event); } + if (this.getLiveStream(input.streamId) !== undefined) { + return false; + } + this.streams.set(input.streamId, stream); this.refreshStreamTtl(input.streamId, input.streamTTLSeconds); this.changes.emit(input.streamId); - const written = stream.events[stream.events.length - 1]; - if (written === undefined) { - throw new Error('Stream restore append wrote no entry'); - } - return { status: 'ok', tailEntryId: written.entryId }; + return true; } readTailEntryId(streamId: string): string | undefined { @@ -208,14 +205,17 @@ export class InMemoryEventSubscription implements Even return applyEventLogRestore({ readPages: options.readPages, readTailEntryId: () => Promise.resolve(this.store.readTailEntryId(this.streamId)), - appendAfterTail: batch => - Promise.resolve( - this.store.appendAfterTail({ - streamId: this.streamId, - ...batch, - streamTTLSeconds: options.streamTTLSeconds, - }), - ), + installIfAbsent: async pages => { + const events: T[] = []; + for await (const page of pages) { + events.push(...page); + } + return this.store.installIfAbsent({ + streamId: this.streamId, + events, + streamTTLSeconds: options.streamTTLSeconds, + }); + }, refreshTtl: () => { this.store.refreshStreamTtl(this.streamId, options.streamTTLSeconds); return Promise.resolve(); diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index b266295158..5ee588b767 100644 --- a/packages/trueforge/src/runtime/event-subscription/redis.ts +++ b/packages/trueforge/src/runtime/event-subscription/redis.ts @@ -1,4 +1,5 @@ import type { RedisClient } from '@truefoundry/trueforge-core/request-reply'; +import { randomUUID } from 'node:crypto'; import { setTimeout as sleep } from 'node:timers/promises'; import { StreamGoneError, @@ -9,72 +10,102 @@ import { type SequencedEvent, } from '.'; import { entryIdMilliseconds, eventIdSeekCursor, eventIdTimestampMs, isLegacySequenceEntryId } from './entryId'; -import { applyEventLogRestore, type AppendAfterTailResult, type RestoreLogOptions } from './restore'; +import { applyEventLogRestore, type RestoreLogOptions } from './restore'; const SUBSCRIBE_STREAM_POLL_ITEMS_COUNT = 100; /** Delay between XREAD polls while a stream has no new events (a blocking XREAD would monopolize the shared connection). */ const SUBSCRIBE_STREAM_POLL_SLEEP_INTERVAL_MS = 1_000; +/** Lua helper shared by the seek scripts. Old entries have no id field; decoding is limited to those. */ +const PAYLOAD_ID_LUA = ` +local function payloadId(fields) + local data = nil + for j = 1, #fields, 2 do + if fields[j] == 'id' then return fields[j + 1] end + if fields[j] == 'data' then data = fields[j + 1] end + end + if data == nil then return nil end + local ok, obj = pcall(cjson.decode, data) + if ok and type(obj) == 'table' and type(obj['id']) == 'string' then + return obj['id'] + end + return nil +end +`; + /** - * One bounded step of an event-id seek. The TS side loops, so Redis serves - * other clients between batches. + * One bounded forward step to the first entry with a payload id. The TS side + * loops, so Redis serves other clients between batches. * - * KEYS[1]=stream ARGV[1]=from|after ARGV[2]=payload id ARGV[3]=exclusive start - * ARGV[4]=batch ARGV[5]=last match so far ('' = none) ARGV[6]='1' on the first step + * KEYS[1]=stream ARGV[1]=payload id ARGV[2]=exclusive start ARGV[3]=batch + * ARGV[4]='1' on the first step * - * Returns {state, cursor, lastMatch}. `cursor` is an exclusive XREAD cursor. - * more → keep going from `cursor` carrying `lastMatch`. - * found → from: the entry before the first match. after: the last entry with the id. + * Returns {state, cursor}. `cursor` is an exclusive XREAD cursor. + * more → continue from `cursor`. found → the entry before the first match. * missing → the last entry seen (live-only resume). */ -const SEEK_EVENT_ID_STEP = ` -local mode, targetId, cursor, batch = ARGV[1], ARGV[2], ARGV[3], tonumber(ARGV[4]) -local lastMatch = ARGV[5] +const SEEK_FIRST_STEP = `${PAYLOAD_ID_LUA} +local targetId, cursor, batch = ARGV[1], ARGV[2], tonumber(ARGV[3]) local entries = redis.call('XRANGE', KEYS[1], '(' .. cursor, '+', 'COUNT', batch) if #entries == 0 then - if ARGV[6] == '1' then + if ARGV[4] == '1' then -- Start was past the tip (e.g. a future ULID ms): park at the real tip so -- XREAD still sees later live appends instead of waiting above them. local tip = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) if tip[1] ~= nil then - return {'missing', tip[1][1], ''} + return {'missing', tip[1][1]} end end - if lastMatch ~= '' then - return {'found', lastMatch, ''} - end - return {'missing', cursor, ''} + return {'missing', cursor} end local previous = cursor for i = 1, #entries do - local entryId, fields = entries[i][1], entries[i][2] - local payloadId, data = nil, nil - for j = 1, #fields, 2 do - if fields[j] == 'id' then payloadId = fields[j + 1]; break end - if fields[j] == 'data' then data = fields[j + 1] end + if payloadId(entries[i][2]) == targetId then + return {'found', previous} end - -- Old entries have no id field; decoding is limited to those. - if payloadId == nil and data ~= nil then - local ok, obj = pcall(cjson.decode, data) - if ok and type(obj) == 'table' and type(obj['id']) == 'string' then - payloadId = obj['id'] - end + previous = entries[i][1] +end +if #entries < batch then + return {'missing', previous} +end +return {'more', previous} +`; + +/** + * One bounded backward step to the last entry with a payload id. Walking down + * from the tip, the first match is the last one, so a recent cursor settles in one step. + * + * KEYS[1]=stream ARGV[1]=payload id ARGV[2]=exclusive floor ARGV[3]=batch + * ARGV[4]=exclusive upper bound ('' = the tip) ARGV[5]=tip seen so far ('' = none) + * + * Returns {state, cursor, tip}. found → the last entry with the id (XREAD cursor). + * more → continue below `cursor`. missing → `cursor` is the tip (live-only resume). + */ +const SEEK_LAST_STEP = `${PAYLOAD_ID_LUA} +local targetId, floor, batch, upper, tip = ARGV[1], ARGV[2], tonumber(ARGV[3]), ARGV[4], ARGV[5] +local high = '+' +if upper ~= '' then high = '(' .. upper end +local entries = redis.call('XREVRANGE', KEYS[1], high, '(' .. floor, 'COUNT', batch) +if tip == '' then + if upper == '' and entries[1] ~= nil then + tip = entries[1][1] + else + local last = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) + if last[1] ~= nil then tip = last[1][1] end end - if payloadId == targetId then - if mode == 'from' then - return {'found', previous, ''} - end - lastMatch = entryId +end +for i = 1, #entries do + if payloadId(entries[i][2]) == targetId then + return {'found', entries[i][1], tip} end - previous = entryId end if #entries < batch then - if lastMatch ~= '' then - return {'found', lastMatch, ''} + if tip == '' then + return {'missing', '0-0', tip} end - return {'missing', previous, ''} + return {'missing', tip, tip} end -return {'more', previous, lastMatch} +return {'more', entries[#entries][1], tip} `; /** @@ -103,78 +134,28 @@ end return 'ok' `; -/** - * Restore append. KEYS[1]=stream ARGV[1]=expected tail id ('' = empty) ARGV[2]=ttl - * ARGV[3..]=ms, event id, payload triples. - * - * Appends only while the tail is still the expected one. - * Returns {'ok', lastEntryId} or 'stale'. The caller uses lastEntryId as the - * expected tail of the next page. - */ -const APPEND_AFTER_TAIL = ` -local tail = redis.call('XREVRANGE', KEYS[1], '+', '-', 'COUNT', 1) -local tailId = '' -if tail[1] ~= nil then - tailId = tail[1][1] -end -if tailId ~= ARGV[1] then - return 'stale' -end -local tailMs = string.match(tailId, '^%d+') or '0' -local lastAdded = nil -for i = 3, #ARGV, 3 do - local ms = ARGV[i] - if tonumber(tailMs) > tonumber(ms) then - ms = tailMs - end - lastAdded = redis.call('XADD', KEYS[1], ms .. '-*', 'id', ARGV[i + 1], 'data', ARGV[i + 2]) - tailMs = string.match(lastAdded, '^%d+') -end -local ttl = tonumber(ARGV[2]) -if ttl > 0 then - redis.call('EXPIRE', KEYS[1], ttl) -end -if lastAdded == nil then - return 'stale' -end -return {'ok', lastAdded} -`; - function ttlArgument(streamTTLSeconds: number | undefined): string { return String(streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? streamTTLSeconds : 0); } -function parseAppendReply(value: unknown, streamId: string): AppendAfterTailResult { - if (value === 'stale') { - return { status: 'stale' }; - } - if (Array.isArray(value) && value.length === 2) { - const status: unknown = value[0]; - const tailEntryId: unknown = value[1]; - if (status === 'ok' && typeof tailEntryId === 'string') { - return { status: 'ok', tailEntryId }; - } - } - throw new Error(`Unexpected stream restore reply on ${streamId}`); -} - interface SeekStep { state: 'more' | 'found' | 'missing'; cursor: string; - lastMatch: string; + /** Only set by the backward step. */ + tip: string; } function parseSeekReply(value: unknown): SeekStep { - if (Array.isArray(value) && value.length === 3) { + if (Array.isArray(value) && (value.length === 2 || value.length === 3)) { const state: unknown = value[0]; const cursor: unknown = value[1]; - const lastMatch: unknown = value[2]; + const tip: unknown = value[2] ?? ''; if ( (state === 'more' || state === 'found' || state === 'missing') && typeof cursor === 'string' && - typeof lastMatch === 'string' + typeof tip === 'string' ) { - return { state, cursor, lastMatch }; + return { state, cursor, tip }; } } throw new Error('Unexpected event-id seek reply'); @@ -191,7 +172,7 @@ export class RedisEventSubscription implements EventSu await applyEventLogRestore({ readPages: options.readPages, readTailEntryId: () => this.readTailEntryId(), - appendAfterTail: batch => this.appendAfterTail({ ...batch, streamTTLSeconds: options.streamTTLSeconds }), + installIfAbsent: pages => this.installIfAbsent(pages, options.streamTTLSeconds), refreshTtl: () => this.refreshTtl(options.streamTTLSeconds), }); } @@ -310,42 +291,44 @@ export class RedisEventSubscription implements EventSu } /** - * Jump near the ULID millisecond (or stream start), then walk forward one - * bounded Lua step at a time until the seek settles. Returns an exclusive - * XREAD cursor. See SEEK_EVENT_ID_STEP. An aborted signal ends the walk early. + * Bound the walk at the ULID millisecond (or stream start), then run bounded + * Lua steps until the seek settles. `from` walks up to the first match; `after` + * walks down from the tip to the last one. Returns an exclusive XREAD cursor. + * An aborted signal ends the walk early. */ private async seekEventIdCursor(input: { mode: 'from' | 'after'; eventId: string; signal?: AbortSignal | undefined; }): Promise<{ cursor: string; found: boolean }> { - let cursor = eventIdSeekCursor({ + const floor = eventIdSeekCursor({ eventId: input.eventId, firstEntryId: await this.firstEntryId(), }); - let lastMatch = ''; + const batch = String(SUBSCRIBE_STREAM_POLL_ITEMS_COUNT); + let cursor = input.mode === 'from' ? floor : ''; + let tip = ''; let first = true; while (input.signal?.aborted !== true) { - const reply: unknown = await this.redis.eval(SEEK_EVENT_ID_STEP, { - keys: [this.streamId], - arguments: [ - input.mode, - input.eventId, - cursor, - String(SUBSCRIBE_STREAM_POLL_ITEMS_COUNT), - lastMatch, - first ? '1' : '0', - ], - }); + const reply: unknown = + input.mode === 'from' + ? await this.redis.eval(SEEK_FIRST_STEP, { + keys: [this.streamId], + arguments: [input.eventId, cursor, batch, first ? '1' : '0'], + }) + : await this.redis.eval(SEEK_LAST_STEP, { + keys: [this.streamId], + arguments: [input.eventId, floor, batch, cursor, tip], + }); const step = parseSeekReply(reply); if (step.state !== 'more') { return { cursor: step.cursor, found: step.state === 'found' }; } cursor = step.cursor; - lastMatch = step.lastMatch; + tip = step.tip; first = false; } - return { cursor, found: false }; + return { cursor: cursor === '' ? floor : cursor, found: false }; } private async readAfter( @@ -422,17 +405,54 @@ export class RedisEventSubscription implements EventSu } } - private async appendAfterTail(input: { - tailEntryId: string | undefined; - events: T[]; - streamTTLSeconds: number | undefined; - }): Promise { - const triples = input.events.flatMap(event => [this.entryMillisecondFor(event), event.id, JSON.stringify(event)]); - const reply: unknown = await this.redis.eval(APPEND_AFTER_TAIL, { - keys: [this.streamId], - arguments: [input.tailEntryId ?? '', ttlArgument(input.streamTTLSeconds), ...triples], - }); - return parseAppendReply(reply, this.streamId); + /** + * Writes the log to a private key, then swaps it in only if the stream is still + * absent. A lost swap deletes the copy. The build is pipelined per page; the + * swap is the atomic step, so a reader never sees a prefix. + */ + private async installIfAbsent( + pages: AsyncIterable, + streamTTLSeconds: number | undefined, + ): Promise { + const tempKey = `${this.streamId}:restore:${randomUUID()}`; + let previousMs = 0; + let wrote = false; + try { + for await (const page of pages) { + const writes: Promise[] = []; + for (const event of page) { + const ms = Math.max(Number(this.entryMillisecondFor(event)), previousMs); + previousMs = ms; + writes.push(this.redis.xAdd(tempKey, `${String(ms)}-*`, { id: event.id, data: JSON.stringify(event) })); + } + if (writes.length === 0) { + continue; + } + await Promise.all(writes); + wrote = true; + await this.expireKey(tempKey, streamTTLSeconds); + } + if (!wrote) { + return true; + } + // Restart the clock so a slow read does not expire the key at the moment of the swap. + await this.expireKey(tempKey, streamTTLSeconds); + const renamed = await this.redis.renameNX(tempKey, this.streamId); + if (renamed === 1) { + return true; + } + await this.redis.del(tempKey); + return false; + } catch (error) { + await this.redis.del(tempKey).catch(() => undefined); + throw error; + } + } + + private async expireKey(key: string, streamTTLSeconds: number | undefined): Promise { + if (streamTTLSeconds !== undefined && streamTTLSeconds > 0) { + await this.redis.expire(key, streamTTLSeconds); + } } private entryMillisecondFor(event: T): string { diff --git a/packages/trueforge/src/runtime/event-subscription/restore.ts b/packages/trueforge/src/runtime/event-subscription/restore.ts index 4d9c442c24..b57f119748 100644 --- a/packages/trueforge/src/runtime/event-subscription/restore.ts +++ b/packages/trueforge/src/runtime/event-subscription/restore.ts @@ -5,41 +5,27 @@ */ export interface RestoreLogOptions { - /** Durable events, one page at a time. Appended in the order read. */ + /** Durable events, one page at a time, in the order to append. */ readPages: () => AsyncIterable; streamTTLSeconds?: number | undefined; } -/** `'stale'` means the tail moved: another writer created or extended the stream. */ -export type AppendAfterTailResult = { status: 'ok'; tailEntryId: string } | { status: 'stale' }; - /** - * Appends every page when the stream has no entries. A stream that already has - * entries is left as is and only gets its TTL refreshed, so a near-expiry - * stream stays subscribable. The next page's expected tail is the entry id the - * previous append wrote, so a concurrent writer stops the walk instead of - * landing between pages. + * Installs the whole log when the stream is still absent. Returns false when the + * stream was created first: the caller refreshes its TTL and keeps it as is. */ export async function applyEventLogRestore(input: { readPages: () => AsyncIterable; readTailEntryId: () => Promise; - appendAfterTail: (batch: { tailEntryId: string | undefined; events: T[] }) => Promise; + installIfAbsent: (pages: AsyncIterable) => Promise; refreshTtl: () => Promise; }): Promise { if ((await input.readTailEntryId()) !== undefined) { await input.refreshTtl(); return; } - let expectedTailId: string | undefined; - for await (const page of input.readPages()) { - if (page.length === 0) { - continue; - } - const outcome = await input.appendAfterTail({ tailEntryId: expectedTailId, events: [...page] }); - if (outcome.status === 'stale') { - await input.refreshTtl(); - return; - } - expectedTailId = outcome.tailEntryId; + const installed = await input.installIfAbsent(input.readPages()); + if (!installed) { + await input.refreshTtl(); } } diff --git a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts index 99d78cc855..9384ede144 100644 --- a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts +++ b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts @@ -206,6 +206,22 @@ export function runEventSubscriptionContractSuite( expect(await collect(registry, streamId, sent.length)).toEqual(sent); }); + it('keeps a single copy when two restores run together', async () => { + const { registry } = setup(); + const streamId = nextStreamId(); + const sent = [makeEvent(0), makeEvent(1), makeEvent(2)]; + const restore = () => + registry.get(streamId).restoreLog({ + readPages: () => durablePages(sent), + streamTTLSeconds: STREAM_TTL_SECONDS, + }); + await Promise.all([restore(), restore()]); + + const extra = makeEvent(3); + await registry.get(streamId).put(extra, { streamTTLSeconds: STREAM_TTL_SECONDS }); + expect(await collect(registry, streamId, sent.length + 1)).toEqual([...sent, extra]); + }); + it('poll parked at the tip receives events put after it subscribed', async () => { const { registry } = setup(); const streamId = nextStreamId(); @@ -292,6 +308,23 @@ export function runEventSubscriptionContractSuite( await expect(producer.hasEventId('evt-missing')).resolves.toBe(false); }); + it('after finds the last copy of an id several batches below the tip', async () => { + const { registry } = setup(); + const streamId = nextStreamId(); + const producer = registry.get(streamId); + const shell = makeEvent(0, 'evt-message'); + const delta = { ...shell, payload: { index: 1, text: 'delta' } }; + const later = Array.from({ length: 250 }, (_, index) => makeEvent(index + 2)); + await producer.put(shell, { streamTTLSeconds: STREAM_TTL_SECONDS }); + await producer.put(delta, { streamTTLSeconds: STREAM_TTL_SECONDS }); + for (const event of later) { + await producer.put(event, { streamTTLSeconds: STREAM_TTL_SECONDS }); + } + + const received = await collect(registry, streamId, later.length, { afterEventId: shell.id }); + expect(received).toEqual(later); + }); + it('seeks past large payloads', async () => { const { registry } = setup(); const streamId = nextStreamId(); diff --git a/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts b/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts index 60a196f91f..26d5e1c76b 100644 --- a/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/redis.contract.test.ts @@ -58,6 +58,25 @@ describeRedis('RedisEventSubscription (EventSubscription contract)', () => { const entries = await client.xRange(streamId, '-', '+'); expect(entries.map(entry => entry.message)).toEqual([{ id: 'evt-1', data: JSON.stringify(event) }]); + expect(await client.keys(`${streamId}:restore:*`)).toEqual([]); + }); + + it('does not replace a stream that appears before the swap', async () => { + const registry = new EventSubscriptionRegistry(client); + const live: TestEvent = { id: 'evt-live', type: 'test.event', payload: { index: 0, text: 'live' } }; + await registry.get(streamId).put(live, { streamTTLSeconds: 60 }); + + await registry.get(streamId).restoreLog({ + readPages: async function* () { + await Promise.resolve(); + yield [event]; + }, + streamTTLSeconds: 60, + }); + + const entries = await client.xRange(streamId, '-', '+'); + expect(entries.map(entry => entry.message)).toEqual([{ id: 'evt-live', data: JSON.stringify(live) }]); + expect(await client.keys(`${streamId}:restore:*`)).toEqual([]); }); it('still finds an id in an entry that has no id field', async () => { diff --git a/packages/trueforge/tests/runtime/event-subscription/restore.test.ts b/packages/trueforge/tests/runtime/event-subscription/restore.test.ts index 3b54bd3150..92c7771c14 100644 --- a/packages/trueforge/tests/runtime/event-subscription/restore.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/restore.test.ts @@ -18,18 +18,18 @@ async function* pagesOf(...pages: Row[][]): AsyncGenerator { } describe('applyEventLogRestore', () => { - it('appends every page of the log to an empty stream, chaining the expected tail', async () => { - const batches: { tailEntryId: string | undefined; events: Row[] }[] = []; - const writtenTailIds = ['2-0', '4-0']; + it('passes every page to the install when the stream is missing', async () => { + const pages: Row[][] = []; let refreshed = 0; await applyEventLogRestore({ readPages: () => pagesOf([created, paused], [], [running, done]), readTailEntryId: () => Promise.resolve(undefined), - appendAfterTail: batch => { - batches.push(batch); - const tailEntryId = writtenTailIds.shift(); - return Promise.resolve(tailEntryId === undefined ? { status: 'stale' } : { status: 'ok', tailEntryId }); + installIfAbsent: async events => { + for await (const page of events) { + pages.push([...page]); + } + return true; }, refreshTtl: () => { refreshed += 1; @@ -37,23 +37,20 @@ describe('applyEventLogRestore', () => { }, }); - expect(batches).toEqual([ - { tailEntryId: undefined, events: [created, paused] }, - { tailEntryId: '2-0', events: [running, done] }, - ]); + expect(pages).toEqual([[created, paused], [], [running, done]]); expect(refreshed).toBe(0); }); it('leaves an existing stream alone and only refreshes its TTL', async () => { - let appended = 0; + let installed = 0; let refreshed = 0; await applyEventLogRestore({ readPages: () => pagesOf([created, paused, running]), readTailEntryId: () => Promise.resolve('2-0'), - appendAfterTail: () => { - appended += 1; - return Promise.resolve({ status: 'stale' }); + installIfAbsent: () => { + installed += 1; + return Promise.resolve(true); }, refreshTtl: () => { refreshed += 1; @@ -61,28 +58,23 @@ describe('applyEventLogRestore', () => { }, }); - expect(appended).toBe(0); + expect(installed).toBe(0); expect(refreshed).toBe(1); }); - it('stops when a concurrent writer moves the tail and refreshes the TTL', async () => { - const batches: { tailEntryId: string | undefined; events: Row[] }[] = []; + it('refreshes the TTL when the install loses the race', async () => { let refreshed = 0; await applyEventLogRestore({ - readPages: () => pagesOf([created], [paused]), + readPages: () => pagesOf([created]), readTailEntryId: () => Promise.resolve(undefined), - appendAfterTail: batch => { - batches.push(batch); - return Promise.resolve({ status: 'stale' }); - }, + installIfAbsent: () => Promise.resolve(false), refreshTtl: () => { refreshed += 1; return Promise.resolve(); }, }); - expect(batches).toEqual([{ tailEntryId: undefined, events: [created] }]); expect(refreshed).toBe(1); }); }); From 6527f22f0c4d9f6b3a3982e33f9c7858975db145 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 5 Nov 2026 10:41:21 +0530 Subject: [PATCH 17/21] Log turn stream restore decisions, page queries, and timings. Restore logs the missing-stream rebuild, each durable page query, the install or lost race, and failures. Counts and timings stay in the message text so they are not indexed as log fields. Request-scoped ids come from the log context. Co-authored-by: Cursor --- packages/trueforge/src/apis/turns.ts | 2 + .../runtime/event-subscription/inMemory.ts | 4 +- .../src/runtime/event-subscription/index.ts | 7 +- .../src/runtime/event-subscription/redis.ts | 6 +- .../src/runtime/event-subscription/restore.ts | 18 +- .../src/runtime/turnEventLogRestore.ts | 66 +++++- .../event-subscription/restore.test.ts | 9 +- .../turnEventLogRestore.test.ts | 188 ++++++++++++++++-- 8 files changed, 257 insertions(+), 43 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index d18ccd360c..de8f136de7 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -845,6 +845,7 @@ async function resumePausedTurn(input: { await restoreTurnEventLog({ sessionStore: input.deps.sessionStore, subscription: turnEventStream, + logger: input.deps.logger, sessionId: input.session.session_id, turnId: turn.id, streamTTLSeconds: configuration.TURN_STREAM_TTL_SECONDS, @@ -2037,6 +2038,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps 0 ? Date.now() + streamTTLSeconds * 1_000 : undefined; @@ -201,7 +201,7 @@ export class InMemoryEventSubscription implements Even private readonly streamId: string, ) {} - restoreLog(options: RestoreLogOptions): Promise { + restoreLog(options: RestoreLogOptions): Promise { return applyEventLogRestore({ readPages: options.readPages, readTailEntryId: () => Promise.resolve(this.store.readTailEntryId(this.streamId)), diff --git a/packages/trueforge/src/runtime/event-subscription/index.ts b/packages/trueforge/src/runtime/event-subscription/index.ts index b1725bc1d5..676671407f 100644 --- a/packages/trueforge/src/runtime/event-subscription/index.ts +++ b/packages/trueforge/src/runtime/event-subscription/index.ts @@ -1,9 +1,9 @@ import type { RedisClient } from '@truefoundry/trueforge-core/request-reply'; import { InMemoryEventStreamStore, InMemoryEventSubscription } from './inMemory'; import { RedisEventSubscription } from './redis'; -import type { RestoreLogOptions } from './restore'; +import type { RestoreLogOptions, RestoreLogOutcome } from './restore'; -export type { RestoreLogOptions } from './restore'; +export type { RestoreLogOptions, RestoreLogOutcome } from './restore'; export interface EventSubscriptionPutOptions { /** Sets the lifetime of the whole stream after this append. */ @@ -50,8 +50,9 @@ export interface EventSubscription { /** * Rebuilds a missing stream from the whole durable log using event-id stream * ids. An existing stream is left in place and only gets its TTL refreshed. + * Resolves with how the restore ended. */ - restoreLog(options: RestoreLogOptions): Promise; + restoreLog(options: RestoreLogOptions): Promise; /** Appends an event. The stream id is `{ulidMillisecond}-{n}` when the event id is a ULID. */ put(event: T, options?: EventSubscriptionPutOptions): Promise; diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index 5ee588b767..a6eac50e4d 100644 --- a/packages/trueforge/src/runtime/event-subscription/redis.ts +++ b/packages/trueforge/src/runtime/event-subscription/redis.ts @@ -10,7 +10,7 @@ import { type SequencedEvent, } from '.'; import { entryIdMilliseconds, eventIdSeekCursor, eventIdTimestampMs, isLegacySequenceEntryId } from './entryId'; -import { applyEventLogRestore, type RestoreLogOptions } from './restore'; +import { applyEventLogRestore, type RestoreLogOptions, type RestoreLogOutcome } from './restore'; const SUBSCRIBE_STREAM_POLL_ITEMS_COUNT = 100; /** Delay between XREAD polls while a stream has no new events (a blocking XREAD would monopolize the shared connection). */ @@ -168,8 +168,8 @@ export class RedisEventSubscription implements EventSu private readonly streamId: string, ) {} - async restoreLog(options: RestoreLogOptions): Promise { - await applyEventLogRestore({ + restoreLog(options: RestoreLogOptions): Promise { + return applyEventLogRestore({ readPages: options.readPages, readTailEntryId: () => this.readTailEntryId(), installIfAbsent: pages => this.installIfAbsent(pages, options.streamTTLSeconds), diff --git a/packages/trueforge/src/runtime/event-subscription/restore.ts b/packages/trueforge/src/runtime/event-subscription/restore.ts index b57f119748..e749f61d5d 100644 --- a/packages/trueforge/src/runtime/event-subscription/restore.ts +++ b/packages/trueforge/src/runtime/event-subscription/restore.ts @@ -11,7 +11,13 @@ export interface RestoreLogOptions { } /** - * Installs the whole log when the stream is still absent. Returns false when the + * `present`: the stream already existed. `installed`: the log was written. + * `lost_race`: another writer created the stream while the log was being read. + */ +export type RestoreLogOutcome = 'present' | 'installed' | 'lost_race'; + +/** + * Installs the whole log when the stream is still absent. Reports `lost_race` when the * stream was created first: the caller refreshes its TTL and keeps it as is. */ export async function applyEventLogRestore(input: { @@ -19,13 +25,15 @@ export async function applyEventLogRestore(input: { readTailEntryId: () => Promise; installIfAbsent: (pages: AsyncIterable) => Promise; refreshTtl: () => Promise; -}): Promise { +}): Promise { if ((await input.readTailEntryId()) !== undefined) { await input.refreshTtl(); - return; + return 'present'; } const installed = await input.installIfAbsent(input.readPages()); - if (!installed) { - await input.refreshTtl(); + if (installed) { + return 'installed'; } + await input.refreshTtl(); + return 'lost_race'; } diff --git a/packages/trueforge/src/runtime/turnEventLogRestore.ts b/packages/trueforge/src/runtime/turnEventLogRestore.ts index 6e239a12d5..8ce4288f1b 100644 --- a/packages/trueforge/src/runtime/turnEventLogRestore.ts +++ b/packages/trueforge/src/runtime/turnEventLogRestore.ts @@ -1,9 +1,17 @@ import type { ISessionStore, TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session'; -import type { EventSubscription } from './event-subscription'; +import { extractErrorLogFields } from '@truefoundry/trueforge-core/core/util/errorLogFields'; +import type { Logger } from 'winston'; +import type { EventSubscription, RestoreLogOutcome } from './event-subscription'; /** Large pages keep a full restore to few round trips. */ const RESTORE_PAGE_LIMIT = 1_000; +/** Counts across pages, reported once the restore ends. */ +interface RestoreProgress { + pages: number; + events: number; +} + /** * Restore a paused turn's live stream from the durable log when the stream is gone. * An existing stream is never patched: persist order can differ from emit order. @@ -11,24 +19,50 @@ const RESTORE_PAGE_LIMIT = 1_000; export async function restoreTurnEventLog(input: { sessionStore: Pick, 'listTurnEvents'>; subscription: EventSubscription; + logger: Logger; sessionId: string; turnId: string; streamTTLSeconds: number; }): Promise { - await input.subscription.restoreLog({ - readPages: () => durableEventPages(input), - streamTTLSeconds: input.streamTTLSeconds, - }); + const progress: RestoreProgress = { pages: 0, events: 0 }; + const startedAt = performance.now(); + let outcome: RestoreLogOutcome; + try { + outcome = await input.subscription.restoreLog({ + readPages: () => durableEventPages({ ...input, progress }), + streamTTLSeconds: input.streamTTLSeconds, + }); + } catch (error) { + input.logger.warn( + `Turn stream restore failed (${restoreStats(progress, startedAt)})`, + extractErrorLogFields(error), + ); + throw error; + } + const stats = restoreStats(progress, startedAt); + if (outcome === 'present') { + // Paused resumes of a live stream are the common case; keep them out of info logs. + input.logger.debug(`Turn stream present, TTL refreshed in ${String(elapsedMs(startedAt))}ms`); + } else if (outcome === 'installed') { + input.logger.info(`Turn stream restored from durable log (${stats})`); + } else { + input.logger.info(`Turn stream restore lost a race to another writer (${stats})`); + } } /** One page in memory, in event-id order. */ async function* durableEventPages(input: { sessionStore: Pick, 'listTurnEvents'>; + logger: Logger; sessionId: string; turnId: string; + progress: RestoreProgress; }): AsyncGenerator { + // Pages are requested only when the stream was absent at the tail check. + input.logger.info('Turn stream missing, rebuilding from durable log'); let pageToken: string | undefined; for (;;) { + const queryStartedAt = performance.now(); const page = await input.sessionStore.listTurnEvents({ session_id: input.sessionId, turn_id: input.turnId, @@ -36,11 +70,31 @@ async function* durableEventPages(input: { page_token: pageToken, order: 'asc', }); - yield page.data; const nextPageToken = page.pagination.next_page_token; + input.progress.pages += 1; + input.progress.events += page.data.length; + const more = nextPageToken !== undefined ? ', more remain' : ''; + input.logger.info( + `Turn stream restore page ${String(input.progress.pages)} queried ${String(page.data.length)} rows in ${String(elapsedMs(queryStartedAt))}ms${more}`, + ); + yield page.data; if (nextPageToken === undefined) { return; } pageToken = nextPageToken; } } + +/** Counts and elapsed time, kept in the message so they are not indexed as fields. */ +function restoreStats(progress: RestoreProgress, startedAt: number): string { + return `${countLabel(progress.events, 'event')}, ${countLabel(progress.pages, 'page')}, ${String(elapsedMs(startedAt))}ms`; +} + +function countLabel(count: number, singular: string): string { + return `${String(count)} ${singular}${count === 1 ? '' : 's'}`; +} + +/** Milliseconds with one decimal, for sub-millisecond queries. */ +function elapsedMs(startedAt: number): number { + return Math.round((performance.now() - startedAt) * 10) / 10; +} diff --git a/packages/trueforge/tests/runtime/event-subscription/restore.test.ts b/packages/trueforge/tests/runtime/event-subscription/restore.test.ts index 92c7771c14..9eeaca6a40 100644 --- a/packages/trueforge/tests/runtime/event-subscription/restore.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/restore.test.ts @@ -22,7 +22,7 @@ describe('applyEventLogRestore', () => { const pages: Row[][] = []; let refreshed = 0; - await applyEventLogRestore({ + const outcome = await applyEventLogRestore({ readPages: () => pagesOf([created, paused], [], [running, done]), readTailEntryId: () => Promise.resolve(undefined), installIfAbsent: async events => { @@ -39,13 +39,14 @@ describe('applyEventLogRestore', () => { expect(pages).toEqual([[created, paused], [], [running, done]]); expect(refreshed).toBe(0); + expect(outcome).toBe('installed'); }); it('leaves an existing stream alone and only refreshes its TTL', async () => { let installed = 0; let refreshed = 0; - await applyEventLogRestore({ + const outcome = await applyEventLogRestore({ readPages: () => pagesOf([created, paused, running]), readTailEntryId: () => Promise.resolve('2-0'), installIfAbsent: () => { @@ -60,12 +61,13 @@ describe('applyEventLogRestore', () => { expect(installed).toBe(0); expect(refreshed).toBe(1); + expect(outcome).toBe('present'); }); it('refreshes the TTL when the install loses the race', async () => { let refreshed = 0; - await applyEventLogRestore({ + const outcome = await applyEventLogRestore({ readPages: () => pagesOf([created]), readTailEntryId: () => Promise.resolve(undefined), installIfAbsent: () => Promise.resolve(false), @@ -76,5 +78,6 @@ describe('applyEventLogRestore', () => { }); expect(refreshed).toBe(1); + expect(outcome).toBe('lost_race'); }); }); diff --git a/packages/trueforge/tests/runtime/event-subscription/turnEventLogRestore.test.ts b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRestore.test.ts index 40080dbcf1..fc695b0248 100644 --- a/packages/trueforge/tests/runtime/event-subscription/turnEventLogRestore.test.ts +++ b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRestore.test.ts @@ -1,5 +1,6 @@ import type { TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session/TurnHandle'; import { InMemorySessionStore } from '@truefoundry/trueforge-core/agent-session/store/InMemorySessionStore'; +import winston from 'winston'; import { makeAgentSpec, makeCreateTurnInput, @@ -16,6 +17,9 @@ import { InMemoryEventStreamStore } from '../../../src/runtime/event-subscriptio import { restoreTurnEventLog } from '../../../src/runtime/turnEventLogRestore'; const STREAM_TTL_SECONDS = 300; +const TENANT_ID = 't1'; +const SESSION_ID = 's1'; +const TURN_ID = 'turn-1'; async function collect( registry: EventSubscriptionRegistry, @@ -39,28 +43,52 @@ async function collect( return received; } +async function createPausedTurnStore(): Promise { + const store = new InMemorySessionStore(); + await store.createSession({ + tenant_id: TENANT_ID, + session_id: SESSION_ID, + created_by_subject: { subject_id: 'user-1', subject_type: 'user', subject_display_name: 'user-1' }, + agent: { type: 'inline', spec: makeAgentSpec() }, + custom: null, + metadata: {}, + external_id: null, + source: null, + }); + await store.createTurn(makeCreateTurnInput({ sessionId: SESSION_ID, turnId: TURN_ID })); + return store; +} + +/** Collects log records in memory so tests can assert on what a developer would see. */ +class RecordingTransport extends winston.Transport { + readonly records: winston.Logform.TransformableInfo[] = []; + + log(info: winston.Logform.TransformableInfo, next: () => void): void { + this.records.push(info); + next(); + } +} + +function recordingLogger(): { logger: winston.Logger; records: winston.Logform.TransformableInfo[] } { + const transport = new RecordingTransport(); + const logger = winston.createLogger({ level: 'debug', transports: [transport] }); + return { logger, records: transport.records }; +} + +/** Winston writes through a stream; let queued records reach the transport. */ +async function flushLogs(): Promise { + await new Promise(resolve => setImmediate(resolve)); +} + describe('restoreTurnEventLog', () => { it('rebuilds a missing stream from the whole durable log and never patches an existing one', async () => { - const store = new InMemorySessionStore(); - const sessionId = 's1'; - const turnId = 'turn-1'; - await store.createSession({ - tenant_id: 't1', - session_id: sessionId, - created_by_subject: { subject_id: 'user-1', subject_type: 'user', subject_display_name: 'user-1' }, - agent: { type: 'inline', spec: makeAgentSpec() }, - custom: null, - metadata: {}, - external_id: null, - source: null, - }); - await store.createTurn(makeCreateTurnInput({ sessionId, turnId })); - const created = makeTurnCreatedEvent(turnId); + const store = await createPausedTurnStore(); + const created = makeTurnCreatedEvent(TURN_ID); const message = makeModelMessageEvent(); await store.appendToEvents({ active_executor_id: TEST_ACTIVE_EXECUTOR_ID, - session_id: sessionId, - turn_id: turnId, + session_id: SESSION_ID, + turn_id: TURN_ID, events: [created, message], }); @@ -68,12 +96,14 @@ describe('restoreTurnEventLog', () => { const registry = new EventSubscriptionRegistry(undefined, streams); const streamId = 'turn-stream'; const subscription = registry.get(streamId); + const { logger } = recordingLogger(); const restore = () => restoreTurnEventLog({ sessionStore: store, subscription, - sessionId, - turnId, + logger, + sessionId: SESSION_ID, + turnId: TURN_ID, streamTTLSeconds: STREAM_TTL_SECONDS, }); @@ -89,8 +119,8 @@ describe('restoreTurnEventLog', () => { const later = makeModelMessageEvent(); await store.appendToEvents({ active_executor_id: TEST_ACTIVE_EXECUTOR_ID, - session_id: sessionId, - turn_id: turnId, + session_id: SESSION_ID, + turn_id: TURN_ID, events: [later], }); await restore(); @@ -110,4 +140,120 @@ describe('restoreTurnEventLog', () => { appended.id, ]); }); + + it('logs the missing stream, one line per page with its query time, and the total duration', async () => { + const store = await createPausedTurnStore(); + // More rows than one page, so the log shows a second page query. + const events = [makeTurnCreatedEvent(TURN_ID), ...Array.from({ length: 1_001 }, () => makeModelMessageEvent())]; + await store.appendToEvents({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: SESSION_ID, + turn_id: TURN_ID, + events, + }); + + const streams = new InMemoryEventStreamStore(); + const registry = new EventSubscriptionRegistry(undefined, streams); + const subscription = registry.get('turn-stream'); + const { logger, records } = recordingLogger(); + const restore = () => + restoreTurnEventLog({ + sessionStore: store, + subscription, + logger, + sessionId: SESSION_ID, + turnId: TURN_ID, + streamTTLSeconds: STREAM_TTL_SECONDS, + }); + + await restore(); + await flushLogs(); + const rebuilt = records.filter(record => record.message.startsWith('Turn stream')); + const [missing, firstPage, secondPage, restored] = rebuilt; + expect(missing?.message).toBe('Turn stream missing, rebuilding from durable log'); + expect(firstPage?.message).toMatch(/^Turn stream restore page 1 queried 1000 rows in \d+(\.\d)?ms, more remain$/); + expect(secondPage?.message).toMatch(/^Turn stream restore page 2 queried 2 rows in \d+(\.\d)?ms$/); + expect(restored?.message).toMatch(/^Turn stream restored from durable log \(1002 events, 2 pages, \d+(\.\d)?ms\)$/); + for (const record of rebuilt) { + expect(record?.level).toBe('info'); + expect(Object.keys(record ?? {}).sort()).toEqual(['level', 'message']); + } + + records.length = 0; + await restore(); + await flushLogs(); + expect(records[0]?.message).toMatch(/^Turn stream present, TTL refreshed in \d+(\.\d)?ms$/); + expect(records[0]?.level).toBe('debug'); + expect(Object.keys(records[0] ?? {}).sort()).toEqual(['level', 'message']); + }); + + it('logs a lost race when two restores install the same stream', async () => { + const store = await createPausedTurnStore(); + await store.appendToEvents({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: SESSION_ID, + turn_id: TURN_ID, + events: [makeTurnCreatedEvent(TURN_ID)], + }); + + const streams = new InMemoryEventStreamStore(); + const registry = new EventSubscriptionRegistry(undefined, streams); + const subscription = registry.get('turn-stream'); + const { logger, records } = recordingLogger(); + const restore = () => + restoreTurnEventLog({ + sessionStore: store, + subscription, + logger, + sessionId: SESSION_ID, + turnId: TURN_ID, + streamTTLSeconds: STREAM_TTL_SECONDS, + }); + + await Promise.all([restore(), restore()]); + await flushLogs(); + const outcomes = records.filter( + record => record.message.includes('restore lost a race') || record.message.includes('restored from durable log'), + ); + const messages = outcomes.map(record => record.message); + expect( + messages.filter(message => /^Turn stream restored from durable log \(1 event, 1 page, /.test(message)), + ).toHaveLength(1); + expect( + messages.filter(message => + /^Turn stream restore lost a race to another writer \(1 event, 1 page, /.test(message), + ), + ).toHaveLength(1); + expect(outcomes.every(record => Object.keys(record).sort().join() === 'level,message')).toBe(true); + }); + + it('logs a failed rebuild with its progress and rethrows the error', async () => { + const failingStore = { + listTurnEvents: () => Promise.reject(new Error('database unavailable')), + }; + const streams = new InMemoryEventStreamStore(); + const registry = new EventSubscriptionRegistry(undefined, streams); + const subscription = registry.get('turn-stream'); + const { logger, records } = recordingLogger(); + + await expect( + restoreTurnEventLog({ + sessionStore: failingStore, + subscription, + logger, + sessionId: SESSION_ID, + turnId: TURN_ID, + streamTTLSeconds: STREAM_TTL_SECONDS, + }), + ).rejects.toThrow('database unavailable'); + await flushLogs(); + const failure = records.find(record => record.message.startsWith('Turn stream restore failed')); + expect(failure?.message).toMatch(/^Turn stream restore failed \(0 events, 0 pages, \d+(\.\d)?ms\)$/); + expect(failure).toMatchObject({ + level: 'warn', + error: 'database unavailable', + }); + expect(typeof failure?.['stack']).toBe('string'); + expect(Object.keys(failure ?? {}).sort()).toEqual(['error', 'level', 'message', 'stack']); + }); }); From 7730956a2e857f647d3488a41340923401c6ab3d Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Fri, 6 Nov 2026 10:05:47 +0530 Subject: [PATCH 18/21] Keep interleaved model-message resumes from duplicating text. A resume starts at the oldest open message shell, and reconnect progress is counted per cursor so replaying the same events still exhausts the retry budget. Test fixtures include the memory catalog and configuration switch. Co-authored-by: Cursor --- .../src/convertTurnMessages.ts | 3 +- .../src/foldPeerThreads.ts | 17 +++++ .../assistant-ui-runtime/src/streamTurn.ts | 10 ++- .../src/turnStreamUpdate.ts | 68 ++++++++++++++++--- .../src/useTrueForgeAgentMessages.ts | 28 ++++---- .../test/convertTurnMessages.test.ts | 49 +++++++++++++ .../test/turnStreamUpdate.test.ts | 42 ++++++++++++ .../tests/unit/apis/modelProviders.test.ts | 2 + .../tests/unit/apis/sessionHttp.test.ts | 4 +- 9 files changed, 193 insertions(+), 30 deletions(-) create mode 100644 packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts diff --git a/packages/assistant-ui-runtime/src/convertTurnMessages.ts b/packages/assistant-ui-runtime/src/convertTurnMessages.ts index f7f6b54e66..f6a5419478 100644 --- a/packages/assistant-ui-runtime/src/convertTurnMessages.ts +++ b/packages/assistant-ui-runtime/src/convertTurnMessages.ts @@ -33,6 +33,7 @@ import { extractTurnUserText } from './extractTurnUserText.js'; import { buildRootAssistantContent, buildRootAssistantContentForIds, + earliestOpenModelMessageId, findFirstPendingApprovalThreadId, findFirstPendingResponseThreadId, ingestStreamEvent, @@ -1940,7 +1941,7 @@ export async function* streamTurnEvents( // Resume replays from mcp.auth_required until the prompt is answered. A later // paused update must not move the cursor past it: a fresh fold would lose the prompt. if (pendingMcpAuth == null || event.type === EVENT_TYPE.MCP_AUTH_REQUIRED || endsMcpAuthPrompt) { - onSubscribeCursor?.(subscribeCursorForEvent(event)); + onSubscribeCursor?.(subscribeCursorForEvent(event, earliestOpenModelMessageId(foldState))); } } diff --git a/packages/assistant-ui-runtime/src/foldPeerThreads.ts b/packages/assistant-ui-runtime/src/foldPeerThreads.ts index 6e51df167f..bd5fba0d7b 100644 --- a/packages/assistant-ui-runtime/src/foldPeerThreads.ts +++ b/packages/assistant-ui-runtime/src/foldPeerThreads.ts @@ -445,6 +445,23 @@ function compactionPartFromStarted({ }; } +/** Shell id of the oldest model message that has not finished. ULIDs sort by time. */ +export function earliestOpenModelMessageId(state: PeerThreadFoldState): string | undefined { + let earliest: string | undefined; + for (const bucket of state.threads.values()) { + for (const id of bucket.modelMessageIds) { + const event = bucket.events.get(id); + if (event?.type !== EVENT_TYPE.MODEL_MESSAGE || event.finishReason != null) { + continue; + } + if (earliest === undefined || id < earliest) { + earliest = id; + } + } + } + return earliest; +} + export function ingestStreamEvent(state: PeerThreadFoldState, message: TurnStreamingEvent): boolean { if (message.type === EVENT_TYPE.MCP_AUTH_REQUIRED || isTurnScopedEvent(message)) { return false; diff --git a/packages/assistant-ui-runtime/src/streamTurn.ts b/packages/assistant-ui-runtime/src/streamTurn.ts index 6e6d8bc1be..3352791bfd 100644 --- a/packages/assistant-ui-runtime/src/streamTurn.ts +++ b/packages/assistant-ui-runtime/src/streamTurn.ts @@ -3,8 +3,12 @@ import { EVENT_TYPE } from './server/events.js'; import type { AgentChatServer, PreviousTurnIdInput, UserMessageContent } from './server/types.js'; import { streamTurnEvents } from './convertTurnMessages.js'; -import { PeerThreadFoldState } from './foldPeerThreads.js'; -import type { TurnStreamUpdate, TurnSubscribeCursor } from './turnStreamUpdate.js'; +import { earliestOpenModelMessageId, PeerThreadFoldState } from './foldPeerThreads.js'; +import { + resumeCursorCoveringOpenModelMessages, + type TurnStreamUpdate, + type TurnSubscribeCursor, +} from './turnStreamUpdate.js'; export interface StreamTurnOptions { userMessage?: UserMessageContent; @@ -92,7 +96,7 @@ export async function* resumeTurnStream( server.subscribeToTurn({ sessionId, turnId, - ...(subscribeCursor ?? {}), + ...(resumeCursorCoveringOpenModelMessages(subscribeCursor, earliestOpenModelMessageId(foldState)) ?? {}), abortSignal, }), foldState, diff --git a/packages/assistant-ui-runtime/src/turnStreamUpdate.ts b/packages/assistant-ui-runtime/src/turnStreamUpdate.ts index 12bc2aaf32..71c85edf27 100644 --- a/packages/assistant-ui-runtime/src/turnStreamUpdate.ts +++ b/packages/assistant-ui-runtime/src/turnStreamUpdate.ts @@ -8,21 +8,73 @@ import type { TurnState } from './server/index.js'; export type TurnSubscribeCursor = { fromEventId: string } | { afterEventId: string }; /** - * A model message replays from its shell. `mcp.auth_required` is replayed too: - * the connect prompt lives only on that event, and the next `turn.update` + * A model message replays from its shell. When several messages are still open, + * replay starts at the earliest shell so a later message's deltas are not appended + * onto an earlier message that was already folded. `mcp.auth_required` is replayed + * too: the connect prompt lives only on that event, and the next `turn.update` * does not repeat it. Every other event resumes after itself. */ -export function subscribeCursorForEvent(event: { type: string; id: string }): TurnSubscribeCursor { - if ( - event.type === EVENT_TYPE.MODEL_MESSAGE || - event.type === EVENT_TYPE.MODEL_MESSAGE_DELTA || - event.type === EVENT_TYPE.MCP_AUTH_REQUIRED - ) { +export function subscribeCursorForEvent( + event: { type: string; id: string }, + earliestOpenModelMessageId?: string, +): TurnSubscribeCursor { + if (event.type === EVENT_TYPE.MODEL_MESSAGE || event.type === EVENT_TYPE.MODEL_MESSAGE_DELTA) { + const fromEventId = + earliestOpenModelMessageId !== undefined && earliestOpenModelMessageId < event.id + ? earliestOpenModelMessageId + : event.id; + return { fromEventId }; + } + if (event.type === EVENT_TYPE.MCP_AUTH_REQUIRED) { return { fromEventId: event.id }; } return { afterEventId: event.id }; } +/** + * Resume of one model message replays every event after its shell. If an older + * message is still open, start at that shell instead so its later deltas rebuild + * the message rather than appending to the copy already in the fold. + */ +export function resumeCursorCoveringOpenModelMessages( + cursor: TurnSubscribeCursor | undefined, + earliestOpenModelMessageId: string | undefined, +): TurnSubscribeCursor | undefined { + if (cursor === undefined || !('fromEventId' in cursor)) { + return cursor; + } + if (earliestOpenModelMessageId === undefined || earliestOpenModelMessageId >= cursor.fromEventId) { + return cursor; + } + return { fromEventId: earliestOpenModelMessageId }; +} + +/** Deliveries of one cursor, kept so a replay of the same events is not new progress. */ +export function createSubscribeReplayProgress(): { + beginAttempt: () => void; + observe: (cursor: TurnSubscribeCursor) => boolean; +} { + const mostDelivered = new Map(); + let thisAttempt = new Map(); + return { + beginAttempt() { + thisAttempt = new Map(); + }, + /** True when this attempt delivered more of this cursor than any earlier attempt. */ + observe(cursor) { + const key = 'fromEventId' in cursor ? `from:${cursor.fromEventId}` : `after:${cursor.afterEventId}`; + const seen = (thisAttempt.get(key) ?? 0) + 1; + thisAttempt.set(key, seen); + const most = mostDelivered.get(key) ?? 0; + if (seen <= most) { + return false; + } + mostDelivered.set(key, seen); + return true; + }, + }; +} + export interface TurnStreamUpdate { content: AssistantContentPart[]; status?: MessageStatus; diff --git a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts index fbaba8be6f..5012ef43f1 100644 --- a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts +++ b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts @@ -50,7 +50,12 @@ import { delayReconnect, isAbortError, STREAM_RECONNECT_MAX_ATTEMPTS } from './s import { resumeTurnStream, streamTurnContent } from './streamTurn.js'; import { mapApprovalDecision, type RespondToToolApprovalOptions } from './toolApproval.js'; import { type RespondToToolResponseOptions } from './toolResponse.js'; -import { subscribeCursorForEvent, type TurnStreamUpdate, type TurnSubscribeCursor } from './turnStreamUpdate.js'; +import { + createSubscribeReplayProgress, + subscribeCursorForEvent, + type TurnStreamUpdate, + type TurnSubscribeCursor, +} from './turnStreamUpdate.js'; export interface UseTrueForgeAgentMessagesOptions { server: AgentChatServer; @@ -605,11 +610,10 @@ export function useTrueForgeAgentMessages({ let useSubscribe = false; let consecutiveFailures = 0; // Inclusive resume replays the message shell and every delta already seen. - // Only events past the most this cursor ever delivered count as progress, - // so a stream that keeps dying mid-replay still exhausts the retry budget. - let progressCursorKey: string | undefined; - let mostDeliveredAtCursor = 0; - let deliveredAtCursorThisAttempt = 0; + // Only a cursor delivered more times than on any earlier attempt counts as + // progress, including when two message ids alternate. A stream that keeps + // dying at the same replay position still exhausts the retry budget. + const replayProgress = createSubscribeReplayProgress(); const track: StreamSequenceTrack = { onSequenceNumber: sequenceNumber => { @@ -618,15 +622,7 @@ export function useTrueForgeAgentMessages({ }, onSubscribeCursor: cursor => { subscribeCursor = cursor; - const key = 'fromEventId' in cursor ? `from:${cursor.fromEventId}` : `after:${cursor.afterEventId}`; - if (key !== progressCursorKey) { - progressCursorKey = key; - mostDeliveredAtCursor = 0; - deliveredAtCursorThisAttempt = 0; - } - deliveredAtCursorThisAttempt += 1; - if (deliveredAtCursorThisAttempt > mostDeliveredAtCursor) { - mostDeliveredAtCursor = deliveredAtCursorThisAttempt; + if (replayProgress.observe(cursor)) { consecutiveFailures = 0; } }, @@ -705,7 +701,7 @@ export function useTrueForgeAgentMessages({ }; const readStream = (signal: AbortSignal): AsyncGenerator => { - deliveredAtCursorThisAttempt = 0; + replayProgress.beginAttempt(); if (!useSubscribe) { return createStream(signal, track); } diff --git a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts index b8e536d868..8903add70f 100644 --- a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts +++ b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts @@ -1926,6 +1926,55 @@ describe('convertTurnMessages', () => { expect(cursors).toEqual([{ fromEventId: 'm1' }, { fromEventId: 'auth-1' }, { afterEventId: 'turn-done' }]); }); + it('keeps the resume cursor on the earliest open model message while a later one is streaming', async () => { + const cursors: { fromEventId?: string; afterEventId?: string }[] = []; + await collectStream( + streamTurnEvents( + streamFrom([ + modelMessage({ id: 'm-a', threadId: ROOT_THREAD_ID, content: '' }), + modelMessage({ id: 'm-b', threadId: ROOT_THREAD_ID, content: '' }), + { + type: 'model.message.delta', + id: 'm-b', + createdAt, + threadId: ROOT_THREAD_ID, + content: 'later', + }, + modelMessage({ id: 'm-a', threadId: ROOT_THREAD_ID, content: 'earlier', finishReason: 'stop' }), + { + type: 'model.message.delta', + id: 'm-b', + createdAt, + threadId: ROOT_THREAD_ID, + content: ' still open', + }, + { + type: 'turn.done', + id: 'done-open', + createdAt, + state: { status: 'done', completedAt: createdAt }, + }, + ]), + new PeerThreadFoldState(), + undefined, + undefined, + undefined, + cursor => { + cursors.push(cursor); + }, + ), + ); + + expect(cursors).toEqual([ + { fromEventId: 'm-a' }, + { fromEventId: 'm-a' }, + { fromEventId: 'm-a' }, + { fromEventId: 'm-a' }, + { fromEventId: 'm-b' }, + { afterEventId: 'done-open' }, + ]); + }); + it('yields folded content after each ingested stream event', async () => { const foldState = new PeerThreadFoldState(); const updates = await collectStream( diff --git a/packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts b/packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts new file mode 100644 index 0000000000..bbc582e2eb --- /dev/null +++ b/packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it } from 'vitest'; +import { + createSubscribeReplayProgress, + resumeCursorCoveringOpenModelMessages, + subscribeCursorForEvent, +} from '../src/turnStreamUpdate.js'; + +describe('subscribeCursorForEvent', () => { + it('starts a model-message resume at an older open shell', () => { + expect(subscribeCursorForEvent({ type: 'model.message.delta', id: 'm-b' }, 'm-a')).toEqual({ + fromEventId: 'm-a', + }); + expect(subscribeCursorForEvent({ type: 'model.message.delta', id: 'm-a' }, 'm-b')).toEqual({ + fromEventId: 'm-a', + }); + }); +}); + +describe('resumeCursorCoveringOpenModelMessages', () => { + it('pulls a from cursor back to an older open shell and leaves other cursors alone', () => { + expect(resumeCursorCoveringOpenModelMessages({ fromEventId: 'm-b' }, 'm-a')).toEqual({ fromEventId: 'm-a' }); + expect(resumeCursorCoveringOpenModelMessages({ fromEventId: 'm-a' }, 'm-b')).toEqual({ fromEventId: 'm-a' }); + expect(resumeCursorCoveringOpenModelMessages({ afterEventId: 'pause' }, 'm-a')).toEqual({ + afterEventId: 'pause', + }); + expect(resumeCursorCoveringOpenModelMessages(undefined, 'm-a')).toBeUndefined(); + }); +}); + +describe('createSubscribeReplayProgress', () => { + it('does not treat a replay of interleaved cursors as progress', () => { + const progress = createSubscribeReplayProgress(); + const cursors = [{ fromEventId: 'm-a' }, { fromEventId: 'm-b' }, { fromEventId: 'm-a' }] as const; + + progress.beginAttempt(); + expect(cursors.map(cursor => progress.observe(cursor))).toEqual([true, true, true]); + + progress.beginAttempt(); + expect(cursors.map(cursor => progress.observe(cursor))).toEqual([false, false, false]); + expect(progress.observe({ fromEventId: 'm-b' })).toBe(true); + }); +}); diff --git a/packages/trueforge/tests/unit/apis/modelProviders.test.ts b/packages/trueforge/tests/unit/apis/modelProviders.test.ts index b5e9a9c86d..7d5b0e6773 100644 --- a/packages/trueforge/tests/unit/apis/modelProviders.test.ts +++ b/packages/trueforge/tests/unit/apis/modelProviders.test.ts @@ -6,6 +6,7 @@ import { createModelsRouter } from '../../../src/apis/models'; import { createSettingsRouter } from '../../../src/apis/settings'; import { createdBySubjectFromRequestContext, STANDALONE_REQUEST_CONTEXT } from '../../../src/auth/identity'; import { McpCatalog } from '../../../src/catalog/McpCatalog'; +import { MemoryCatalog } from '../../../src/catalog/MemoryCatalog'; import { ModelCatalog } from '../../../src/catalog/ModelCatalog'; import { SandboxCatalog } from '../../../src/catalog/SandboxCatalog'; import { SkillCatalog } from '../../../src/catalog/SkillCatalog'; @@ -155,6 +156,7 @@ async function createRouters(): Promise<{ skillCatalog: SkillCatalog.load(), sandboxCatalog: SandboxCatalog.load(), webSearchCatalog: WebSearchCatalog.load(), + memoryCatalog: MemoryCatalog.load(), }), modelsRouter: createModelsRouter({ resolveModelProviderStore: () => modelProviderStore, diff --git a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts index 1b0e495a2e..9603001678 100644 --- a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts +++ b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts @@ -952,7 +952,7 @@ describe('sessions HTTP agent binding', () => { await configurationStore.upsertConfiguration({ tenant_id: 'default', - manifest: { web_search: { enabled: false } }, + manifest: { web_search: { enabled: false }, memory: { enabled: false } }, }); try { // The agent was saved while web search was allowed, so its stored spec is grandfathered. @@ -970,7 +970,7 @@ describe('sessions HTTP agent binding', () => { } finally { await configurationStore.upsertConfiguration({ tenant_id: 'default', - manifest: { web_search: { enabled: true } }, + manifest: { web_search: { enabled: true }, memory: { enabled: false } }, }); } From 928f9892e6583a5000c1af7b4680666799d03ae9 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Fri, 6 Nov 2026 10:48:35 +0530 Subject: [PATCH 19/21] Scope resume cursors to the active turn and keep MCP auth cursors across pauses - Pick the earliest open model message only from the current turn so a resume cursor from another turn cannot leak in. - Keep a finished model.message when a replayed shell arrives, so replays do not truncate tool arguments while approvals are pending. - Walk history newest-first for the reconnect cursor and pin it to mcp.auth_required until the auth prompt is answered or the turn moves on. Co-authored-by: Cursor --- .../src/convertTurnMessages.ts | 31 ++++++++++---- .../src/foldPeerThreads.ts | 41 +++++++++++++++---- .../assistant-ui-runtime/src/streamTurn.ts | 4 +- .../src/turnStreamUpdate.ts | 24 +++++++++++ .../src/useTrueForgeAgentMessages.ts | 31 ++++++++------ .../test/foldPeerThreads.test.ts | 33 +++++++++++++++ .../test/turnStreamUpdate.test.ts | 22 ++++++++++ 7 files changed, 157 insertions(+), 29 deletions(-) diff --git a/packages/assistant-ui-runtime/src/convertTurnMessages.ts b/packages/assistant-ui-runtime/src/convertTurnMessages.ts index f6a5419478..8bb4e580fb 100644 --- a/packages/assistant-ui-runtime/src/convertTurnMessages.ts +++ b/packages/assistant-ui-runtime/src/convertTurnMessages.ts @@ -417,6 +417,12 @@ export function prependFoldState(older: PeerThreadFoldState, newer: PeerThreadFo for (const [threadId, link] of newer.threadParents) { result.threadParents.set(threadId, link); } + for (const [eventId, messageTurnId] of older.modelMessageTurnIds) { + result.modelMessageTurnIds.set(eventId, messageTurnId); + } + for (const [eventId, messageTurnId] of newer.modelMessageTurnIds) { + result.modelMessageTurnIds.set(eventId, messageTurnId); + } return result; } @@ -495,7 +501,7 @@ function ingestSessionEventsIntoSnapshot( } else if (event.type === EVENT_TYPE.USER_MCP_AUTH_CONTINUE) { snapshot.pendingMcpAuth = undefined; } - ingestTurnEvent(snapshot.fold, event); + ingestTurnEvent(snapshot.fold, event, currentTurnId); currentContentEvents.push(event); } } @@ -1025,7 +1031,7 @@ async function ingestTurnEventsIntoFold( ...(pageToken != null ? { pageToken } : {}), }), ); - ingestCollectedEventsIntoFold(foldState, events); + ingestCollectedEventsIntoFold(foldState, events, turnId); } async function fetchTurnEvents( @@ -1048,7 +1054,11 @@ async function fetchTurnEvents( return events; } -function ingestCollectedEventsIntoFold(foldState: PeerThreadFoldState, events: PersistedTurnEvent[]): void { +function ingestCollectedEventsIntoFold( + foldState: PeerThreadFoldState, + events: PersistedTurnEvent[], + turnId: string, +): void { for (const event of events) { if ( event.type === EVENT_TYPE.TURN_CREATED || @@ -1057,7 +1067,7 @@ function ingestCollectedEventsIntoFold(foldState: PeerThreadFoldState, events: P ) { continue; } - ingestTurnEvent(foldState, event); + ingestTurnEvent(foldState, event, turnId); } } @@ -1572,7 +1582,7 @@ export async function buildSnapshotFromSession( const eventArrays = await fetchAllTurnEventsWithConcurrency(server, sessionId, [turn], concurrency); const tipEvents = eventArrays[0] ?? []; - ingestCollectedEventsIntoFold(snapshot.fold, tipEvents); + ingestCollectedEventsIntoFold(snapshot.fold, tipEvents, turn.id); applyUserToolResponsesToFold(snapshot.fold, turn.input ?? []); const sandboxEvent = tipEvents.find( @@ -1854,8 +1864,10 @@ export async function* streamTurnEvents( onTurnIdAvailable?: (turnId: string) => void, onSequenceNumber?: (sequenceNumber: number) => void, onSubscribeCursor?: (cursor: TurnSubscribeCursor) => void, + turnId?: string, ): AsyncGenerator { let pendingMcpAuth: McpAuthRequiredEvent | undefined; + let activeTurnId = turnId; let sandboxId: string | undefined; let sandboxIdYielded = false; @@ -1933,6 +1945,9 @@ export async function* streamTurnEvents( onSequenceNumber?.(data.sequenceNumber); } const event = data.event; + if (event.type === EVENT_TYPE.TURN_CREATED) { + activeTurnId = event.turnId; + } if ('id' in event && typeof event.id === 'string') { const endsMcpAuthPrompt = event.type === EVENT_TYPE.USER_MCP_AUTH_CONTINUE || @@ -1941,7 +1956,7 @@ export async function* streamTurnEvents( // Resume replays from mcp.auth_required until the prompt is answered. A later // paused update must not move the cursor past it: a fresh fold would lose the prompt. if (pendingMcpAuth == null || event.type === EVENT_TYPE.MCP_AUTH_REQUIRED || endsMcpAuthPrompt) { - onSubscribeCursor?.(subscribeCursorForEvent(event, earliestOpenModelMessageId(foldState))); + onSubscribeCursor?.(subscribeCursorForEvent(event, earliestOpenModelMessageId(foldState, activeTurnId))); } } @@ -1962,7 +1977,7 @@ export async function* streamTurnEvents( if (event.type === EVENT_TYPE.USER_MCP_AUTH_CONTINUE) { pendingMcpAuth = undefined; - ingestStreamEvent(foldState, event); + ingestStreamEvent(foldState, event, activeTurnId); continue; } @@ -2040,7 +2055,7 @@ export async function* streamTurnEvents( return; } - if (!ingestStreamEvent(foldState, event)) { + if (!ingestStreamEvent(foldState, event, activeTurnId)) { continue; } diff --git a/packages/assistant-ui-runtime/src/foldPeerThreads.ts b/packages/assistant-ui-runtime/src/foldPeerThreads.ts index bd5fba0d7b..981bcddae5 100644 --- a/packages/assistant-ui-runtime/src/foldPeerThreads.ts +++ b/packages/assistant-ui-runtime/src/foldPeerThreads.ts @@ -116,6 +116,8 @@ function assertRootThreadEvent(threadId: string, eventType: string): void { export class PeerThreadFoldState { readonly threads = new Map(); readonly threadParents = new Map(); + /** Which turn produced a model message, so a later turn does not resume from an older shell. */ + readonly modelMessageTurnIds = new Map(); /** * POST responses, live replay, and history can expose the same durable event. * Delta events deliberately bypass this set because many chunks share an id. @@ -445,11 +447,14 @@ function compactionPartFromStarted({ }; } -/** Shell id of the oldest model message that has not finished. ULIDs sort by time. */ -export function earliestOpenModelMessageId(state: PeerThreadFoldState): string | undefined { +/** Shell id of the oldest unfinished model message. Pass `turnId` to ignore other turns. */ +export function earliestOpenModelMessageId(state: PeerThreadFoldState, turnId?: string): string | undefined { let earliest: string | undefined; for (const bucket of state.threads.values()) { for (const id of bucket.modelMessageIds) { + if (turnId !== undefined && state.modelMessageTurnIds.get(id) !== turnId) { + continue; + } const event = bucket.events.get(id); if (event?.type !== EVENT_TYPE.MODEL_MESSAGE || event.finishReason != null) { continue; @@ -462,16 +467,37 @@ export function earliestOpenModelMessageId(state: PeerThreadFoldState): string | return earliest; } -export function ingestStreamEvent(state: PeerThreadFoldState, message: TurnStreamingEvent): boolean { +function rememberModelMessageTurn( + state: PeerThreadFoldState, + event: { type: string; id: string }, + turnId: string | undefined, +): void { + if (turnId === undefined || event.type !== EVENT_TYPE.MODEL_MESSAGE) { + return; + } + state.modelMessageTurnIds.set(event.id, turnId); +} + +/** A finished message already holds its full arguments. Replaying its shell would truncate them. */ +function isFinishedModelMessage(state: PeerThreadFoldState, message: TurnStreamingEvent): boolean { + if (message.type !== EVENT_TYPE.MODEL_MESSAGE) { + return false; + } + const existing = state.threads.get(message.threadId)?.events.get(message.id); + return existing?.type === EVENT_TYPE.MODEL_MESSAGE && existing.finishReason != null; +} + +export function ingestStreamEvent(state: PeerThreadFoldState, message: TurnStreamingEvent, turnId?: string): boolean { + rememberModelMessageTurn(state, message, turnId); if (message.type === EVENT_TYPE.MCP_AUTH_REQUIRED || isTurnScopedEvent(message)) { return false; } if (!isEventDelta(message)) { if (state.ingestedEventIds.has(message.id)) { - // Inclusive resume replays the model.message shell so later deltas rebuild - // from a clean base instead of stacking on a partial in-memory fold. - if (message.type !== EVENT_TYPE.MODEL_MESSAGE) { + // Inclusive resume replays an unfinished shell so later deltas rebuild it. + // A finished message already has its full tool arguments; keep those. + if (message.type !== EVENT_TYPE.MODEL_MESSAGE || isFinishedModelMessage(state, message)) { return false; } } else { @@ -512,7 +538,8 @@ export function ingestStreamEvent(state: PeerThreadFoldState, message: TurnStrea return isContentAffectingEvent(message); } -export function ingestTurnEvent(state: PeerThreadFoldState, event: TurnEvent): void { +export function ingestTurnEvent(state: PeerThreadFoldState, event: TurnEvent, turnId?: string): void { + rememberModelMessageTurn(state, event, turnId); if (state.ingestedEventIds.has(event.id)) { return; } diff --git a/packages/assistant-ui-runtime/src/streamTurn.ts b/packages/assistant-ui-runtime/src/streamTurn.ts index 3352791bfd..f49a97f8aa 100644 --- a/packages/assistant-ui-runtime/src/streamTurn.ts +++ b/packages/assistant-ui-runtime/src/streamTurn.ts @@ -96,7 +96,8 @@ export async function* resumeTurnStream( server.subscribeToTurn({ sessionId, turnId, - ...(resumeCursorCoveringOpenModelMessages(subscribeCursor, earliestOpenModelMessageId(foldState)) ?? {}), + ...(resumeCursorCoveringOpenModelMessages(subscribeCursor, earliestOpenModelMessageId(foldState, turnId)) ?? + {}), abortSignal, }), foldState, @@ -104,6 +105,7 @@ export async function* resumeTurnStream( undefined, onSequenceNumber, onSubscribeCursor, + turnId, ); } catch (error) { if (error instanceof Error && error.name === 'AbortError') { diff --git a/packages/assistant-ui-runtime/src/turnStreamUpdate.ts b/packages/assistant-ui-runtime/src/turnStreamUpdate.ts index 71c85edf27..c6a527d1e4 100644 --- a/packages/assistant-ui-runtime/src/turnStreamUpdate.ts +++ b/packages/assistant-ui-runtime/src/turnStreamUpdate.ts @@ -31,6 +31,30 @@ export function subscribeCursorForEvent( return { afterEventId: event.id }; } +/** + * Newest event first. A pending `mcp.auth_required` stays the cursor even when a + * later paused `turn.update` is the last stored row: the prompt is only on that event. + */ +export function cursorFromTurnEvents( + eventsNewestFirst: readonly { type: string; id: string; state?: { status?: string } }[], +): TurnSubscribeCursor | undefined { + let newest: TurnSubscribeCursor | undefined; + for (const event of eventsNewestFirst) { + newest ??= subscribeCursorForEvent(event); + if (event.type === EVENT_TYPE.MCP_AUTH_REQUIRED) { + return { fromEventId: event.id }; + } + const promptEnded = + event.type === EVENT_TYPE.USER_MCP_AUTH_CONTINUE || + event.type === EVENT_TYPE.TURN_DONE || + (event.type === EVENT_TYPE.TURN_UPDATE && event.state?.status === 'running'); + if (promptEnded) { + return newest; + } + } + return newest; +} + /** * Resume of one model message replays every event after its shell. If an older * message is still open, start at that shell instead so its later deltas rebuild diff --git a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts index 5012ef43f1..c3c123c193 100644 --- a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts +++ b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts @@ -52,7 +52,7 @@ import { mapApprovalDecision, type RespondToToolApprovalOptions } from './toolAp import { type RespondToToolResponseOptions } from './toolResponse.js'; import { createSubscribeReplayProgress, - subscribeCursorForEvent, + cursorFromTurnEvents, type TurnStreamUpdate, type TurnSubscribeCursor, } from './turnStreamUpdate.js'; @@ -119,13 +119,14 @@ function cursorForTurn(snapshot: SessionSnapshot, turnId: string): TurnSubscribe if (history === undefined) { return undefined; } + const events = []; for (let index = history.length - 1; index >= 0; index -= 1) { const item = history[index]; - if (item?.turnId === turnId) { - return subscribeCursorForEvent(item.event); + if (item?.turnId === turnId && 'id' in item.event && typeof item.event.id === 'string') { + events.push(item.event); } } - return undefined; + return cursorFromTurnEvents(events); } interface RunStreamReconnect { @@ -268,14 +269,18 @@ function materializeAbandonedStreamRootIds(options: { }, })); const modelId = `client-abandoned-${options.turnId}`; - ingestTurnEvent(options.fold, { - type: EVENT_TYPE.MODEL_MESSAGE, - id: modelId, - threadId: ROOT_THREAD_ID, - createdAt: new Date().toISOString(), - content: text.length > 0 ? text : null, - ...(toolCalls.length > 0 ? { toolCalls } : {}), - }); + ingestTurnEvent( + options.fold, + { + type: EVENT_TYPE.MODEL_MESSAGE, + id: modelId, + threadId: ROOT_THREAD_ID, + createdAt: new Date().toISOString(), + content: text.length > 0 ? text : null, + ...(toolCalls.length > 0 ? { toolCalls } : {}), + }, + options.turnId, + ); return [modelId]; } @@ -1325,7 +1330,7 @@ export function useTrueForgeAgentMessages({ }); updateSnapshot(previous => { for (const persisted of created) { - ingestTurnEvent(previous.fold, persisted); + ingestTurnEvent(previous.fold, persisted, turnId); } // Keep the optimistic overlay through partial pauses. Some custom // hosts return persisted user events before replaying the source diff --git a/packages/assistant-ui-runtime/test/foldPeerThreads.test.ts b/packages/assistant-ui-runtime/test/foldPeerThreads.test.ts index 1ab7acebfb..4ec1e3fb3d 100644 --- a/packages/assistant-ui-runtime/test/foldPeerThreads.test.ts +++ b/packages/assistant-ui-runtime/test/foldPeerThreads.test.ts @@ -8,6 +8,7 @@ import { applyApprovalPoliciesToFold, buildRootAssistantContent, buildRootAssistantContentForIds, + earliestOpenModelMessageId, ingestStreamEvent, ingestTurnEvent, PeerThreadFoldState, @@ -27,6 +28,38 @@ function threadCreated(event: Omit): T } describe('foldPeerThreads', () => { + it('ignores an unfinished shell from another turn when choosing the resume cursor', () => { + const state = new PeerThreadFoldState(); + ingestTurnEvent(state, modelMessage({ id: 'm-old', threadId: ROOT_THREAD_ID, content: 'stuck' }), 'turn-old'); + ingestTurnEvent(state, modelMessage({ id: 'm-new', threadId: ROOT_THREAD_ID, content: '' }), 'turn-new'); + + expect(earliestOpenModelMessageId(state, 'turn-new')).toBe('m-new'); + }); + + it('keeps a finished model message when its shell is replayed', () => { + const state = new PeerThreadFoldState(); + const finished = modelMessage({ + id: 'm-done', + threadId: ROOT_THREAD_ID, + content: '', + finishReason: 'tool_calls', + toolCalls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'save', arguments: '{"title":"full"}' }, + }, + ], + }); + ingestStreamEvent(state, finished, 'turn-1'); + ingestStreamEvent(state, modelMessage({ id: 'm-done', threadId: ROOT_THREAD_ID, content: '' }), 'turn-1'); + + const stored = state.threads.get(ROOT_THREAD_ID)?.events.get('m-done'); + expect(stored?.type === 'model.message' ? stored.toolCalls?.[0]?.function.arguments : undefined).toBe( + '{"title":"full"}', + ); + }); + it('ignores sandbox.created (no threadId, not foldable) without crashing', () => { const state = new PeerThreadFoldState(); const sandboxEvent = { diff --git a/packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts b/packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts index bbc582e2eb..d2c05017cd 100644 --- a/packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts +++ b/packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest'; import { createSubscribeReplayProgress, + cursorFromTurnEvents, resumeCursorCoveringOpenModelMessages, subscribeCursorForEvent, } from '../src/turnStreamUpdate.js'; @@ -27,6 +28,27 @@ describe('resumeCursorCoveringOpenModelMessages', () => { }); }); +describe('cursorFromTurnEvents', () => { + it('keeps a pending mcp auth event when a later pause is the newest row', () => { + expect( + cursorFromTurnEvents([ + { type: 'turn.update', id: 'pause', state: { status: 'paused' } }, + { type: 'mcp.auth_required', id: 'auth' }, + ]), + ).toEqual({ fromEventId: 'auth' }); + }); + + it('uses the newest event once the auth prompt has been answered', () => { + expect( + cursorFromTurnEvents([ + { type: 'turn.done', id: 'done', state: { status: 'done' } }, + { type: 'turn.update', id: 'running', state: { status: 'running' } }, + { type: 'mcp.auth_required', id: 'auth' }, + ]), + ).toEqual({ afterEventId: 'done' }); + }); +}); + describe('createSubscribeReplayProgress', () => { it('does not treat a replay of interleaved cursors as progress', () => { const progress = createSubscribeReplayProgress(); From 365719ae221beed5749d3724ab232502e48689ab Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Fri, 6 Nov 2026 23:04:39 +0530 Subject: [PATCH 20/21] Fix MCP auth wait after turn rebuild, memory ingest of paused turns, and early sandbox id - Restore pending_mcp_auth onto rebuilt AgentThreads so a turn paused for MCP auth does not run before user.mcp_auth_continue arrives. - Memory ingest treats a tip turn as paused when status is 'paused' and uses turn.updated_at as the pause time, matching the new paused state. - Carry the open tip's sandbox id onto the pre-getTurn progress paint. - Accept any state shape in cursorFromTurnEvents and narrow it locally. Co-authored-by: Cursor --- .../src/convertTurnMessages.ts | 1 + .../src/turnStreamUpdate.ts | 8 +- .../test/convertTurnMessages.test.ts | 56 ++++++++++++++ .../src/agent-session/SessionHandle.ts | 1 + .../src/core/runtime/AgentThread.ts | 1 + .../src/core/runtime/AgentThread.types.ts | 2 + .../tests/agent-session/turnResume.test.ts | 74 +++++++++++++++++++ .../trueforge/src/db/memoryIngestStore.ts | 2 + .../PostgresMemoryIngestStore.ts | 12 ++- .../SqliteMemoryIngestStore.ts | 5 +- packages/trueforge/src/memory/exchanges.ts | 27 ++++--- .../db/memoryIngestStoreContractSuite.ts | 45 +++++++---- .../unit/controller/memoryIngest.test.ts | 8 +- .../tests/unit/memory/exchanges.test.ts | 9 ++- .../tests/unit/memory/turnFixtures.ts | 22 ++---- 15 files changed, 219 insertions(+), 54 deletions(-) diff --git a/packages/assistant-ui-runtime/src/convertTurnMessages.ts b/packages/assistant-ui-runtime/src/convertTurnMessages.ts index 8bb4e580fb..a314571414 100644 --- a/packages/assistant-ui-runtime/src/convertTurnMessages.ts +++ b/packages/assistant-ui-runtime/src/convertTurnMessages.ts @@ -752,6 +752,7 @@ export async function buildSnapshotFromSessionEvents( event: open.event, ...(open.update != null ? { update: open.update } : {}), }), + open.sandboxId, ); if (open != null) { applyUserToolResponsesToFold(paintable.fold, open.event.input ?? []); diff --git a/packages/assistant-ui-runtime/src/turnStreamUpdate.ts b/packages/assistant-ui-runtime/src/turnStreamUpdate.ts index c6a527d1e4..572e82abb8 100644 --- a/packages/assistant-ui-runtime/src/turnStreamUpdate.ts +++ b/packages/assistant-ui-runtime/src/turnStreamUpdate.ts @@ -31,12 +31,16 @@ export function subscribeCursorForEvent( return { afterEventId: event.id }; } +function isRunningState(state: unknown): boolean { + return typeof state === 'object' && state !== null && 'status' in state && state.status === 'running'; +} + /** * Newest event first. A pending `mcp.auth_required` stays the cursor even when a * later paused `turn.update` is the last stored row: the prompt is only on that event. */ export function cursorFromTurnEvents( - eventsNewestFirst: readonly { type: string; id: string; state?: { status?: string } }[], + eventsNewestFirst: readonly { type: string; id: string; state?: unknown }[], ): TurnSubscribeCursor | undefined { let newest: TurnSubscribeCursor | undefined; for (const event of eventsNewestFirst) { @@ -47,7 +51,7 @@ export function cursorFromTurnEvents( const promptEnded = event.type === EVENT_TYPE.USER_MCP_AUTH_CONTINUE || event.type === EVENT_TYPE.TURN_DONE || - (event.type === EVENT_TYPE.TURN_UPDATE && event.state?.status === 'running'); + (event.type === EVENT_TYPE.TURN_UPDATE && isRunningState(event.state)); if (promptEnded) { return newest; } diff --git a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts index 8903add70f..bad00fe5cf 100644 --- a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts +++ b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts @@ -4045,6 +4045,62 @@ describe('buildSnapshotFromSessionEvents', () => { expect(snapshot.pendingUser?.content).toBe('now'); }); + it('carries the open tip sandbox id on the pre-getTurn paint', async () => { + const progress: SessionSnapshot[] = []; + let resolveGetTurn: (turn: Turn) => void = () => undefined; + const getTurn = vi.fn( + () => + new Promise(resolve => { + resolveGetTurn = resolve; + }), + ); + const items: SessionEventItem[] = [ + { + turnId: 't-running', + event: { + type: 'turn.created', + id: 'evt-c-running', + turnId: 't-running', + input: [{ type: 'user.message', content: 'now' }], + state: { status: 'running' }, + createdAt, + }, + }, + { turnId: 't-running', event: sandboxCreated({ id: 'sandbox-evt', sandboxId: 'sbx-123' }) }, + { turnId: 't-running', event: modelMessage({ id: 'm1', threadId: ROOT_THREAD_ID, content: 'working' }) }, + ]; + const server = { + listTurns: vi.fn(async () => ({ data: [] })), + getTurn, + listEvents: sessionEventsPage(items), + listTurnEvents: async () => ({ data: [] }), + } as unknown as AgentChatServer; + + let notifyProgress: () => void = () => undefined; + const progressed = new Promise(resolve => { + notifyProgress = resolve; + }); + const pending = buildSnapshotFromSessionEvents(server, SESSION_ID, snap => { + progress.push(snap); + if (snap.pendingUser != null) { + notifyProgress(); + } + }); + + await progressed; + expect(progress.at(-1)?.activeStream?.update.metadata?.custom?.['sandboxId']).toBe('sbx-123'); + + resolveGetTurn({ + id: 't-running', + sessionId: SESSION_ID, + state: { status: 'running' }, + input: [{ type: 'user.message', content: 'now' }], + createdAt, + }); + const snapshot = await pending; + expect(snapshot.activeStream?.update.metadata?.custom?.['sandboxId']).toBe('sbx-123'); + }); + it('publishes a paused tip with requires-action before getTurn resolves', async () => { const progress: SessionSnapshot[] = []; let resolveGetTurn: (turn: Turn) => void = () => undefined; diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index 1d14a91176..b872b40d26 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -582,6 +582,7 @@ export class SessionHandle< parent: input.previousThreadSnapshot?.parent ?? undefined, agentInfo: input.previousThreadSnapshot?.agent_info ?? undefined, preComputedCompletion: input.previousThreadSnapshot?.completion ?? undefined, + pendingMCPAuth: input.previousThreadSnapshot?.pending_mcp_auth, capabilities, capabilityState: input.previousThreadSnapshot?.capability_state ?? undefined, tracing: input.tracing, diff --git a/packages/trueforge-core/src/core/runtime/AgentThread.ts b/packages/trueforge-core/src/core/runtime/AgentThread.ts index 27550f07f4..bb72c87008 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThread.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThread.ts @@ -516,6 +516,7 @@ export class AgentThread { this.currentContextUsage = input.currentContextUsage ?? getEmptyCurrentContextUsage(); this.agentInfo = input.agentInfo; this.preComputedCompletion = input.preComputedCompletion; + this.pendingMCPAuth = input.pendingMCPAuth ?? false; this.sandbox = input.sandbox; const capabilities = input.capabilities ?? []; diff --git a/packages/trueforge-core/src/core/runtime/AgentThread.types.ts b/packages/trueforge-core/src/core/runtime/AgentThread.types.ts index 74504839d1..7f27a328c7 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThread.types.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThread.types.ts @@ -222,6 +222,8 @@ export interface AgentThreadConstructorInput { context?: ContextMessage[] | undefined; currentContextUsage?: CurrentContextUsage | undefined; preComputedCompletion?: SubAgentCompletion | undefined; + /** Restores a wait on `user.mcp_auth_continue` when the thread is rebuilt from a snapshot. */ + pendingMCPAuth?: boolean | undefined; sandbox?: HarnessSandbox | undefined; capabilities?: readonly AgentCapability[] | undefined; /** diff --git a/packages/trueforge-core/tests/agent-session/turnResume.test.ts b/packages/trueforge-core/tests/agent-session/turnResume.test.ts index 773e728d79..b6ea02d073 100644 --- a/packages/trueforge-core/tests/agent-session/turnResume.test.ts +++ b/packages/trueforge-core/tests/agent-session/turnResume.test.ts @@ -508,6 +508,80 @@ describe('TurnHandle.send() MCP-auth continuation', () => { expect(turn.state.status).toBe('done'); }); + it('keeps waiting for continue when another executor rebuilds the paused turn', async () => { + const store = new InMemorySessionStore(); + const sessions = new Sessions({ sessionStore: store }); + const session = await sessions.create({ + tenant_id: 'tenant-1', + session_id: 's1', + created_by_subject: { subject_id: 'user-1', subject_type: 'user', subject_display_name: 'user-1' }, + agent: { type: 'inline', spec: makeAgentSpec() }, + external_id: null, + }); + const { capability, authorize, listTools } = makeAuthCapability(); + const resolverFor = () => + makeTestResolver({ + extraCapabilities: [capability], + llmCreate: jest.fn(() => textReplyStream('authorized')), + }); + const original = new AbortController(); + const turn = await session.createTurn({ + turn_id: mintTestTurnId(), + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + input: [{ type: EventType.USER_MESSAGE, content: 'hello' }], + previous_turn_id: 'none', + abortController: original, + resolver: resolverFor(), + }); + + const originalIterator = turn.stream(); + let step = await withTimeout(originalIterator.next(), 1_000, 'initial auth event'); + while (!step.done && !(step.value.type === EventType.TURN_UPDATE && step.value.state.status === 'paused')) { + step = await withTimeout(originalIterator.next(), 1_000, 'auth pause'); + } + expect(step.done).toBe(false); + const listToolsCallsAtPause = listTools.mock.calls.length; + + const won = await store.claimTurnExecutor({ + session_id: 's1', + turn_id: turn.id, + expected_executor_id: TEST_ACTIVE_EXECUTOR_ID, + new_executor_id: 'executor-b', + }); + expect(won).toBe(true); + const parkedOriginal = originalIterator.next(); + parkedOriginal.catch(() => undefined); + original.abort(CancellationReason.Abandoned); + + const rebuilt = await session.rebuildTurn({ + turn_id: turn.id, + active_executor_id: 'executor-b', + resolver: resolverFor(), + abortController: new AbortController(), + }); + // A rebuilt paused turn parks until a send arrives; it must not retry initialization on its own. + const iterator = rebuilt.stream(); + const firstEvent = iterator.next(); + await new Promise(resolve => setTimeout(resolve, 50)); + expect(listTools).toHaveBeenCalledTimes(listToolsCallsAtPause); + + authorize(); + await sendWithInbox(rebuilt, store, [ + { type: EventType.USER_MCP_AUTH_CONTINUE, id: newEventId(), created_at: new Date().toISOString() }, + ]); + + const eventTypes: string[] = []; + let next = await withTimeout(firstEvent, 1_000, 'rebuilt stream start'); + while (!next.done) { + eventTypes.push(next.value.type); + next = await withTimeout(iterator.next(), 1_000, 'rebuilt stream'); + } + + expect(eventTypes.filter(type => type === EventType.MCP_AUTH_REQUIRED)).toHaveLength(0); + expect(eventTypes.filter(type => type === EventType.USER_MCP_AUTH_CONTINUE)).toHaveLength(1); + expect(rebuilt.state.status).toBe('done'); + }); + it('continues after tool-execution auth without reinitializing tools', async () => { const store = new InMemorySessionStore(); const sessions = new Sessions({ sessionStore: store }); diff --git a/packages/trueforge/src/db/memoryIngestStore.ts b/packages/trueforge/src/db/memoryIngestStore.ts index 9b70adc1c1..5c57f5808f 100644 --- a/packages/trueforge/src/db/memoryIngestStore.ts +++ b/packages/trueforge/src/db/memoryIngestStore.ts @@ -43,6 +43,8 @@ export interface MemoryIngestTurn { state: TurnState; /** ISO-8601 UTC instant. */ created_at: string; + /** ISO-8601 UTC instant; for a paused turn, when it paused. */ + updated_at: string; } export interface MemoryIngestModelMessage { diff --git a/packages/trueforge/src/db/postgres/memory-ingest-store/PostgresMemoryIngestStore.ts b/packages/trueforge/src/db/postgres/memory-ingest-store/PostgresMemoryIngestStore.ts index 54c86df037..23d1322fba 100644 --- a/packages/trueforge/src/db/postgres/memory-ingest-store/PostgresMemoryIngestStore.ts +++ b/packages/trueforge/src/db/postgres/memory-ingest-store/PostgresMemoryIngestStore.ts @@ -30,8 +30,8 @@ const textOnlyInput = sql`( ORDER BY item_index), '[]'::jsonb) FROM jsonb_array_elements(input) WITH ORDINALITY AS i(item, item_index) )`; -/** The tip turn finished waiting on the user (a question, approval or auth). */ -const tipPaused = sql`(t.state->>'status' = 'done' AND jsonb_array_length(COALESCE(t.state->'required_actions', '[]'::jsonb)) > 0)`; +/** The tip turn is paused waiting on the user (a question, approval or auth). */ +const tipPaused = sql`t.state->>'status' = 'paused'`; export class PostgresMemoryIngestStore implements IMemoryIngestStore { readonly #db: Kysely; @@ -107,12 +107,16 @@ export class PostgresMemoryIngestStore implements IMemoryIngestStore { async listSessionTurns(input: { session_id: string }): Promise { const rows = await this.#db .selectFrom('turn') - .select(['turn_id', 'previous_turn_id', textOnlyInput.as('input'), 'state', 'created_at']) + .select(['turn_id', 'previous_turn_id', textOnlyInput.as('input'), 'state', 'created_at', 'updated_at']) .where('session_id', '=', input.session_id) .orderBy('created_at', 'asc') .orderBy('turn_id', 'asc') .execute(); - return rows.map(row => ({ ...row, created_at: row.created_at.toISOString() })); + return rows.map(row => ({ + ...row, + created_at: row.created_at.toISOString(), + updated_at: row.updated_at.toISOString(), + })); } async listMainThreadModelMessages(input: { diff --git a/packages/trueforge/src/db/sqlite/memory-ingest-store/SqliteMemoryIngestStore.ts b/packages/trueforge/src/db/sqlite/memory-ingest-store/SqliteMemoryIngestStore.ts index c7ca3ed6a6..cde9ad143c 100644 --- a/packages/trueforge/src/db/sqlite/memory-ingest-store/SqliteMemoryIngestStore.ts +++ b/packages/trueforge/src/db/sqlite/memory-ingest-store/SqliteMemoryIngestStore.ts @@ -35,8 +35,8 @@ const textOnlyInput = sql`( ORDER BY i.key) FROM json_each(input) AS i )`; -/** The tip turn finished waiting on the user (a question, approval or auth). */ -const tipPaused = sql`(json_extract(t.state, '$.status') = 'done' AND json_array_length(t.state, '$.required_actions') > 0)`; +/** The tip turn is paused waiting on the user (a question, approval or auth). */ +const tipPaused = sql`json_extract(t.state, '$.status') = 'paused'`; /** Mirrors PostgresMemoryIngestStore; timestamps are ISO-8601 TEXT. */ export class SqliteMemoryIngestStore implements IMemoryIngestStore { readonly #db: Kysely; @@ -124,6 +124,7 @@ export class SqliteMemoryIngestStore implements IMemoryIngestStore { textOnlyInput.as('input'), jsonText(sql.ref('state')).as('state'), 'created_at', + 'updated_at', ]) .where('session_id', '=', input.session_id) .orderBy('created_at', 'asc') diff --git a/packages/trueforge/src/memory/exchanges.ts b/packages/trueforge/src/memory/exchanges.ts index b1eed5c9a9..2784ef65a4 100644 --- a/packages/trueforge/src/memory/exchanges.ts +++ b/packages/trueforge/src/memory/exchanges.ts @@ -9,7 +9,7 @@ const TRUNCATION_MARKER = ' [truncated]'; /** * - `closed`: the last turn is terminal with nothing pending, or a later user message superseded it. - * - `paused`: the last turn finished waiting on the user (a question, approval or auth). + * - `paused`: the last turn is paused waiting on the user (a question, approval or auth). * - `running`: the last turn has not finished. */ export type ExchangeStatus = 'closed' | 'paused' | 'running'; @@ -31,10 +31,10 @@ function hasUserMessage(turn: MemoryIngestTurn): boolean { function statusAfter(state: TurnState): ExchangeStatus { switch (state.status) { case 'running': - case 'paused': return 'running'; + case 'paused': + return 'paused'; case 'done': - return state.required_actions.length === 0 ? 'closed' : 'paused'; case 'error': case 'cancelled': return 'closed'; @@ -78,14 +78,14 @@ export function buildExchanges(turns: readonly MemoryIngestTurn[]): TurnExchange current = { turns: [turn], status: statusAfter(turn.state), - completedAt: completedAt(turn.state), + completedAt: completedAt(turn), continuation: !hasUserMessage(turn), }; exchanges.push(current); } else { current.turns.push(turn); current.status = statusAfter(turn.state); - current.completedAt = completedAt(turn.state); + current.completedAt = completedAt(turn); } if (current.status === 'closed') { current = undefined; @@ -130,10 +130,19 @@ export function resumePoint({ return { startIndex: forkIndex === undefined ? 0 : forkIndex + 1, forked: true }; } -function completedAt(state: TurnState): string | null { - return state.status === 'done' || state.status === 'error' || state.status === 'cancelled' - ? state.completed_at - : null; +/** A paused turn's `updated_at` is when it paused: nothing else writes it while it waits. */ +function completedAt(turn: MemoryIngestTurn): string | null { + const { state } = turn; + switch (state.status) { + case 'running': + return null; + case 'paused': + return turn.updated_at; + case 'done': + case 'error': + case 'cancelled': + return state.completed_at; + } } export type ExchangeReadiness = { ready: true; partial: boolean } | { ready: false }; diff --git a/packages/trueforge/tests/db/memoryIngestStoreContractSuite.ts b/packages/trueforge/tests/db/memoryIngestStoreContractSuite.ts index 22d007a2cb..383e4c97de 100644 --- a/packages/trueforge/tests/db/memoryIngestStoreContractSuite.ts +++ b/packages/trueforge/tests/db/memoryIngestStoreContractSuite.ts @@ -16,7 +16,9 @@ import { createEmptyAgentThreadMetrics } from '../../../trueforge-core/src/core/ import { makeCreateTurnInput, makeDoneTurnState, + makePausedTurnState, makeTurnDoneEvent, + makeTurnUpdateEvent, TEST_ACTIVE_EXECUTOR_ID, } from '../../../trueforge-core/tests/agent-session/testHelpers'; import type { IMemoryIngestStore } from '../../src/db/memoryIngestStore'; @@ -101,6 +103,24 @@ export function runMemoryIngestStoreContractSuite( }); } + async function addPausedTurn({ session_id, turn_id }: { session_id: string; turn_id: string }) { + const { sessionStore } = getEnv(); + await sessionStore.createTurn(makeCreateTurnInput({ sessionId: session_id, turnId: turn_id })); + const state = makePausedTurnState(); + await sessionStore.updateTurnNonTerminalState({ + session_id, + turn_id, + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + state, + turn_update_event: makeTurnUpdateEvent(state), + turn_metrics: { + status: 'paused', + observed_at: new Date().toISOString(), + metrics: createEmptyAgentThreadMetrics(), + }, + }); + } + async function candidateIds(since: Date = HOUR_AGO()): Promise { const candidates = await getEnv().store.listCandidateSessions({ ...recent(since), limit: 100 }); return candidates.map(candidate => candidate.session_id); @@ -167,21 +187,7 @@ export function runMemoryIngestStoreContractSuite( it('lists only sessions paused on the user from the paused range', async () => { await createSession({ session_id: 'paused' }); - const pausedState: TerminalTurnState = { - status: 'done', - output: null, - completed_at: new Date().toISOString(), - required_actions: [ - { - type: 'tool.response_required', - id: 'evt-req', - created_at: new Date().toISOString(), - thread_id: MAIN_THREAD_ID, - tool_calls: [{ id: 'call-1', source_event_id: 'evt-ask' }], - }, - ], - }; - await addTurn({ session_id: 'paused', turn_id: 't1', state: pausedState }); + await addPausedTurn({ session_id: 'paused', turn_id: 't1' }); await createSession({ session_id: 'closed' }); await addTurn({ session_id: 'closed', turn_id: 't1' }); @@ -196,6 +202,15 @@ export function runMemoryIngestStoreContractSuite( expect(candidates.map(c => c.session_id)).toEqual(['paused']); }); + it('reports a paused turn with the time it paused', async () => { + await createSession({ session_id: 'paused' }); + await addPausedTurn({ session_id: 'paused', turn_id: 't1' }); + + const [turn] = await getEnv().store.listSessionTurns({ session_id: 'paused' }); + expect(turn?.state).toEqual({ status: 'paused' }); + expect(turn?.updated_at).toMatch(ISO_UTC); + }); + it('includes sessions on a registry agent only when its memory is enabled, even from the UI', async () => { const { seedAgent } = getEnv(); await seedAgent({ tenant_id: TENANT, agent_id: 'agent-on', memory_enabled: true }); diff --git a/packages/trueforge/tests/unit/controller/memoryIngest.test.ts b/packages/trueforge/tests/unit/controller/memoryIngest.test.ts index 05b381662f..7a07cc2bc5 100644 --- a/packages/trueforge/tests/unit/controller/memoryIngest.test.ts +++ b/packages/trueforge/tests/unit/controller/memoryIngest.test.ts @@ -71,7 +71,7 @@ class FakeStore implements IMemoryIngestStore { // session.updated_at moves when a turn is created. const sessionUpdatedAt = tip.created_at; const updatedMs = Date.parse(sessionUpdatedAt); - const tipPaused = tip.state.status === 'done' && tip.state.required_actions.length > 0; + const tipPaused = tip.state.status === 'paused'; const inRange = updatedMs >= input.recent_since.getTime() || (tipPaused && updatedMs >= input.paused_from.getTime() && updatedMs < input.paused_until.getTime()); @@ -262,8 +262,9 @@ describe('runMemoryIngestPass', () => { id: 'u1', previous: null, input: [userMessage('Plan dinner')], - state: pausedState({ completedAt: at(-12 * 60 * MINUTE + MINUTE), toolCallId: 'call-1' }), + state: pausedState(), createdAt: at(-12 * 60 * MINUTE), + updatedAt: at(-12 * 60 * MINUTE + MINUTE), }), ], }), @@ -491,8 +492,9 @@ describe('runMemoryIngestPass', () => { id: 'u1', previous: null, input: [userMessage('Plan dinner')], - state: pausedState({ completedAt: at(-12.5 * 60 * MINUTE), toolCallId: 'call-1' }), + state: pausedState(), createdAt: at(-12.5 * 60 * MINUTE - MINUTE), + updatedAt: at(-12.5 * 60 * MINUTE), }); const fakeSession = session({ id: 's1', diff --git a/packages/trueforge/tests/unit/memory/exchanges.test.ts b/packages/trueforge/tests/unit/memory/exchanges.test.ts index 24bb11ab7f..6911e05057 100644 --- a/packages/trueforge/tests/unit/memory/exchanges.test.ts +++ b/packages/trueforge/tests/unit/memory/exchanges.test.ts @@ -20,8 +20,9 @@ function pausedThenAnsweredTurns() { id: 'u1', previous: null, input: [userMessage('Dinner ideas? I am vegetarian.')], - state: pausedState({ completedAt: at(MINUTE), toolCallId: 'call-1' }), + state: pausedState(), createdAt: at(0), + updatedAt: at(MINUTE), }), turn({ id: 'a2', @@ -167,7 +168,8 @@ describe('renderExchange', () => { ], }, ], - state: pausedState({ completedAt: at(MINUTE), toolCallId: 'other-tool' }), + state: pausedState(), + updatedAt: at(MINUTE), }), turn({ id: 'a2', @@ -194,7 +196,8 @@ describe('renderExchange', () => { id: 'u1', previous: null, input: [userMessage('x'.repeat(MAX_MEMORY_MESSAGE_CHARS + 500))], - state: pausedState({ completedAt: at(MINUTE), toolCallId: 'call-1' }), + state: pausedState(), + updatedAt: at(MINUTE), }), ]), ); diff --git a/packages/trueforge/tests/unit/memory/turnFixtures.ts b/packages/trueforge/tests/unit/memory/turnFixtures.ts index 8b5e6fa564..453228dcc4 100644 --- a/packages/trueforge/tests/unit/memory/turnFixtures.ts +++ b/packages/trueforge/tests/unit/memory/turnFixtures.ts @@ -46,21 +46,8 @@ export function doneState({ return { status: 'done', output, required_actions: [], completed_at: completedAt }; } -export function pausedState({ completedAt, toolCallId }: { completedAt: string; toolCallId: string }): TurnState { - return { - status: 'done', - output: null, - completed_at: completedAt, - required_actions: [ - { - type: 'tool.response_required', - id: `req-${toolCallId}`, - created_at: completedAt, - thread_id: MAIN_THREAD_ID, - tool_calls: [{ id: toolCallId, source_event_id: `ask-${toolCallId}` }], - }, - ], - }; +export function pausedState(): TurnState { + return { status: 'paused' }; } export const RUNNING: TurnState = { status: 'running' }; @@ -71,14 +58,17 @@ export function turn({ input, state, createdAt = at(0), + updatedAt = createdAt, }: { id: string; previous: string | null; input: LegacyTurnInputItem[]; state: TurnState; createdAt?: string; + /** For a paused turn, when it paused. */ + updatedAt?: string; }): MemoryIngestTurn { - return { turn_id: id, previous_turn_id: previous, input, state, created_at: createdAt }; + return { turn_id: id, previous_turn_id: previous, input, state, created_at: createdAt, updated_at: updatedAt }; } export function userMessage(text: string): LegacyTurnInputItem { From 455dfb15b34476ec24a17bcf105612dd365f0516 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Sat, 7 Nov 2026 00:11:31 +0530 Subject: [PATCH 21/21] Send memory exchanges only once closed A turn paused on the user can resume on the same turn id, so ingesting it early moved the checkpoint onto a turn whose final reply was then skipped. Only closed exchanges are ingested now; the pause timeout and partial sends are removed. Co-authored-by: Cursor --- .../memory-ingest-closed-exchanges-only.md | 5 ++ packages/trueforge/.env.example | 4 +- packages/trueforge/src/config.ts | 2 +- .../trueforge/src/controller/memoryIngest.ts | 22 ++---- .../trueforge/src/db/memoryIngestStore.ts | 10 +-- .../PostgresMemoryIngestStore.ts | 17 +--- .../SqliteMemoryIngestStore.ts | 14 +--- packages/trueforge/src/memory/exchanges.ts | 53 ++----------- .../resolveTrueFoundryMemoryProviderConfig.ts | 9 --- .../db/memoryIngestStoreContractSuite.ts | 25 +----- .../unit/controller/memoryIngest.test.ts | 77 +++++++------------ .../tests/unit/memory/exchanges.test.ts | 52 ++----------- .../tests/unit/memory/turnFixtures.ts | 5 +- ...lveTrueFoundryMemoryProviderConfig.test.ts | 1 - 14 files changed, 61 insertions(+), 235 deletions(-) create mode 100644 .changeset/memory-ingest-closed-exchanges-only.md diff --git a/.changeset/memory-ingest-closed-exchanges-only.md b/.changeset/memory-ingest-closed-exchanges-only.md new file mode 100644 index 0000000000..66adc2d87a --- /dev/null +++ b/.changeset/memory-ingest-closed-exchanges-only.md @@ -0,0 +1,5 @@ +--- +'@truefoundry/trueforge': patch +--- + +Memory ingest sends an exchange only once it is closed, so a turn paused on the user is no longer sent early. The `pause_timeout_ms` setting in `TRUEFOUNDRY_MEMORY_SETTINGS` is removed; remove it from any existing configuration. diff --git a/packages/trueforge/.env.example b/packages/trueforge/.env.example index 577054cf37..3e4c462158 100644 --- a/packages/trueforge/.env.example +++ b/packages/trueforge/.env.example @@ -329,8 +329,8 @@ POSTGRES_PORT=5432 # TRUEFOUNDRY_MEMORY_TENANT_ALLOWLIST=["tenant-a","tenant-b"] ## Static JSON settings for the memory-ingest loop; every field is optional (defaults shown). ## lookback_ms: sessions whose latest turn started longer ago are no longer picked up. -## pause_timeout_ms: an exchange paused on the user (question, approval or auth) is sent as-is after this long. -# TRUEFOUNDRY_MEMORY_SETTINGS={"interval_ms":300000,"max_sessions_per_run":50,"max_exchanges_per_session":20,"lookback_ms":7200000,"pause_timeout_ms":43200000} +## An exchange paused on the user (question, approval or auth) is sent only once it closes. +# TRUEFOUNDRY_MEMORY_SETTINGS={"interval_ms":300000,"max_sessions_per_run":50,"max_exchanges_per_session":20,"lookback_ms":7200000} ## Shared sandbox for all tenants (TrueFoundry mode only). PUT settings is disabled (424). # TRUEFOUNDRY_SANDBOX_ENABLED=false diff --git a/packages/trueforge/src/config.ts b/packages/trueforge/src/config.ts index c8af1fbf92..5508fa843d 100644 --- a/packages/trueforge/src/config.ts +++ b/packages/trueforge/src/config.ts @@ -1128,7 +1128,7 @@ export type DistributedServerConfiguration = SharedServerConfiguration & { TRUEFOUNDRY_MEMORY_TENANT_ALLOWLIST: string[]; /** * Static JSON settings for the memory-ingest loop: `interval_ms`, `max_sessions_per_run`, - * `max_exchanges_per_session`, `lookback_ms`, `pause_timeout_ms` (each optional, with defaults). + * `max_exchanges_per_session`, `lookback_ms` (each optional, with defaults). * Env: `TRUEFOUNDRY_MEMORY_SETTINGS`. */ TRUEFOUNDRY_MEMORY_SETTINGS: string | undefined; diff --git a/packages/trueforge/src/controller/memoryIngest.ts b/packages/trueforge/src/controller/memoryIngest.ts index a491cb9fcc..a5b10ebe8f 100644 --- a/packages/trueforge/src/controller/memoryIngest.ts +++ b/packages/trueforge/src/controller/memoryIngest.ts @@ -9,7 +9,6 @@ import { activePath, askUserQuestionsByToolCallId, buildExchanges, - exchangeReadiness, renderExchange, resumePoint, type TurnExchange, @@ -30,7 +29,6 @@ export interface MemoryIngestConfig { maxSessionsPerRun: number; maxExchangesPerSession: number; lookbackMs: number; - pauseTimeoutMs: number; } export interface MemoryIngestDeps { @@ -63,22 +61,17 @@ function errorMessage(error: unknown): string { return error instanceof Error ? error.message : String(error); } -/** Exchanges that may be sent now, in order, stopping at the first one that is not ready. */ +/** Closed exchanges, in order, stopping at the first open one (running or paused on the user). */ function readyExchanges({ exchanges, - now, config, }: { exchanges: readonly TurnExchange[]; - now: Date; config: MemoryIngestConfig; }): TurnExchange[] { const ready: TurnExchange[] = []; for (const exchange of exchanges) { - if ( - ready.length === config.maxExchangesPerSession || - !exchangeReadiness({ exchange, now, pauseTimeoutMs: config.pauseTimeoutMs }).ready - ) { + if (ready.length === config.maxExchangesPerSession || exchange.status !== 'closed') { break; } ready.push(exchange); @@ -99,7 +92,7 @@ async function ingestSession({ deps: MemoryIngestDeps; signal: AbortSignal; }): Promise { - const { store, logger, config, now } = deps; + const { store, logger, config } = deps; const { cursor } = candidate; const logFields = { session_id: candidate.session_id, @@ -121,7 +114,7 @@ async function ingestSession({ }); } - const ready = readyExchanges({ exchanges: buildExchanges(path.slice(startIndex)), now: now(), config }); + const ready = readyExchanges({ exchanges: buildExchanges(path.slice(startIndex)), config }); if (ready.length === 0) { return { outcome: 'not_ready', exchangesSent: 0, exchangesFailed: 0 }; } @@ -206,8 +199,8 @@ export async function runMemoryIngestPass({ const { store, logger, config, now } = deps; const startedAt = now(); - // Recent sessions within the lookback; paused ones only once their pause has timed out, for one - // lookback-wide range. A paused session the user answers gets a new turn and so becomes recent. + // Sessions whose latest turn started within the lookback. A paused turn is not ingested until the + // user's answer or next message closes its exchange, which creates a new turn. const startedMs = startedAt.getTime(); const providers = await deps.providersByTenant(); const listed = @@ -216,8 +209,6 @@ export async function runMemoryIngestPass({ : await store.listCandidateSessions({ tenant_ids: [...providers.keys()], recent_since: new Date(startedMs - config.lookbackMs), - paused_from: new Date(startedMs - config.pauseTimeoutMs - config.lookbackMs), - paused_until: new Date(startedMs - config.pauseTimeoutMs), limit: config.maxSessionsPerRun + 1, }); const backlog = listed.length > config.maxSessionsPerRun; @@ -334,7 +325,6 @@ export function memoryIngestFromConfig({ maxSessionsPerRun: settings.max_sessions_per_run, maxExchangesPerSession: settings.max_exchanges_per_session, lookbackMs: settings.lookback_ms, - pauseTimeoutMs: settings.pause_timeout_ms, }, }; } diff --git a/packages/trueforge/src/db/memoryIngestStore.ts b/packages/trueforge/src/db/memoryIngestStore.ts index 5c57f5808f..1535abd07f 100644 --- a/packages/trueforge/src/db/memoryIngestStore.ts +++ b/packages/trueforge/src/db/memoryIngestStore.ts @@ -43,8 +43,6 @@ export interface MemoryIngestTurn { state: TurnState; /** ISO-8601 UTC instant. */ created_at: string; - /** ISO-8601 UTC instant; for a paused turn, when it paused. */ - updated_at: string; } export interface MemoryIngestModelMessage { @@ -61,12 +59,6 @@ export interface ListMemoryIngestCandidatesInput { tenant_ids: readonly string[]; /** Any session updated at or after this instant. */ recent_since: Date; - /** - * Plus sessions whose tip turn is paused on the user (a question, approval or auth), updated in - * [`paused_from`, `paused_until`): their pause has timed out, so what exists is sent. - */ - paused_from: Date; - paused_until: Date; limit: number; } @@ -75,7 +67,7 @@ export type UpsertMemoryIngestCursorInput = Omit`( ORDER BY item_index), '[]'::jsonb) FROM jsonb_array_elements(input) WITH ORDINALITY AS i(item, item_index) )`; -/** The tip turn is paused waiting on the user (a question, approval or auth). */ -const tipPaused = sql`t.state->>'status' = 'paused'`; export class PostgresMemoryIngestStore implements IMemoryIngestStore { readonly #db: Kysely; @@ -63,12 +61,7 @@ export class PostgresMemoryIngestStore implements IMemoryIngestStore { 'c.updated_at as cursor_updated_at', ]) .where('s.source', 'is', null) - .where(eb => - eb.or([ - eb('s.updated_at', '>=', input.recent_since), - eb.and([eb('s.updated_at', '>=', input.paused_from), eb('s.updated_at', '<', input.paused_until), tipPaused]), - ]), - ) + .where('s.updated_at', '>=', input.recent_since) // Matches resolveMemoryScope: never agent-builder chats; a bound agent's manifest decides; UI chats without one use the default scope. .where(sql`NOT (${createAgentSession})`) .where(eb => @@ -107,16 +100,12 @@ export class PostgresMemoryIngestStore implements IMemoryIngestStore { async listSessionTurns(input: { session_id: string }): Promise { const rows = await this.#db .selectFrom('turn') - .select(['turn_id', 'previous_turn_id', textOnlyInput.as('input'), 'state', 'created_at', 'updated_at']) + .select(['turn_id', 'previous_turn_id', textOnlyInput.as('input'), 'state', 'created_at']) .where('session_id', '=', input.session_id) .orderBy('created_at', 'asc') .orderBy('turn_id', 'asc') .execute(); - return rows.map(row => ({ - ...row, - created_at: row.created_at.toISOString(), - updated_at: row.updated_at.toISOString(), - })); + return rows.map(row => ({ ...row, created_at: row.created_at.toISOString() })); } async listMainThreadModelMessages(input: { diff --git a/packages/trueforge/src/db/sqlite/memory-ingest-store/SqliteMemoryIngestStore.ts b/packages/trueforge/src/db/sqlite/memory-ingest-store/SqliteMemoryIngestStore.ts index cde9ad143c..654f9749c2 100644 --- a/packages/trueforge/src/db/sqlite/memory-ingest-store/SqliteMemoryIngestStore.ts +++ b/packages/trueforge/src/db/sqlite/memory-ingest-store/SqliteMemoryIngestStore.ts @@ -35,8 +35,6 @@ const textOnlyInput = sql`( ORDER BY i.key) FROM json_each(input) AS i )`; -/** The tip turn is paused waiting on the user (a question, approval or auth). */ -const tipPaused = sql`json_extract(t.state, '$.status') = 'paused'`; /** Mirrors PostgresMemoryIngestStore; timestamps are ISO-8601 TEXT. */ export class SqliteMemoryIngestStore implements IMemoryIngestStore { readonly #db: Kysely; @@ -70,16 +68,7 @@ export class SqliteMemoryIngestStore implements IMemoryIngestStore { 'c.updated_at as cursor_updated_at', ]) .where('s.source', 'is', null) - .where(eb => - eb.or([ - eb('s.updated_at', '>=', input.recent_since.toISOString()), - eb.and([ - eb('s.updated_at', '>=', input.paused_from.toISOString()), - eb('s.updated_at', '<', input.paused_until.toISOString()), - tipPaused, - ]), - ]), - ) + .where('s.updated_at', '>=', input.recent_since.toISOString()) // Matches resolveMemoryScope: never agent-builder chats; a bound agent's manifest decides; UI chats without one use the default scope. .where(sql`NOT (${createAgentSession})`) .where(eb => @@ -124,7 +113,6 @@ export class SqliteMemoryIngestStore implements IMemoryIngestStore { textOnlyInput.as('input'), jsonText(sql.ref('state')).as('state'), 'created_at', - 'updated_at', ]) .where('session_id', '=', input.session_id) .orderBy('created_at', 'asc') diff --git a/packages/trueforge/src/memory/exchanges.ts b/packages/trueforge/src/memory/exchanges.ts index 2784ef65a4..2dba7d2780 100644 --- a/packages/trueforge/src/memory/exchanges.ts +++ b/packages/trueforge/src/memory/exchanges.ts @@ -8,8 +8,10 @@ export const MAX_MEMORY_MESSAGE_CHARS = 4_000; const TRUNCATION_MARKER = ' [truncated]'; /** - * - `closed`: the last turn is terminal with nothing pending, or a later user message superseded it. - * - `paused`: the last turn is paused waiting on the user (a question, approval or auth). + * - `closed`: the last turn is terminal, or a later user message superseded it. Only closed + * exchanges are ingested. + * - `paused`: the last turn is paused waiting on the user (a question, approval or auth). The + * exchange stays open: the answer is a resume turn that continues it. * - `running`: the last turn has not finished. */ export type ExchangeStatus = 'closed' | 'paused' | 'running'; @@ -18,9 +20,7 @@ export type ExchangeStatus = 'closed' | 'paused' | 'running'; export interface TurnExchange { turns: [MemoryIngestTurn, ...MemoryIngestTurn[]]; status: ExchangeStatus; - /** When the exchange ended; null while it is running. */ - completedAt: string | null; - /** Starts with a resume turn (e.g. an answer) whose exchange was partly ingested earlier. */ + /** Starts with a resume turn (e.g. an answer) because the checkpoint forked inside its exchange. */ continuation: boolean; } @@ -71,21 +71,18 @@ export function buildExchanges(turns: readonly MemoryIngestTurn[]): TurnExchange for (const turn of turns) { if (current !== undefined && hasUserMessage(turn)) { current.status = 'closed'; - current.completedAt ??= turn.created_at; current = undefined; } if (current === undefined) { current = { turns: [turn], status: statusAfter(turn.state), - completedAt: completedAt(turn), continuation: !hasUserMessage(turn), }; exchanges.push(current); } else { current.turns.push(turn); current.status = statusAfter(turn.state); - current.completedAt = completedAt(turn); } if (current.status === 'closed') { current = undefined; @@ -130,46 +127,6 @@ export function resumePoint({ return { startIndex: forkIndex === undefined ? 0 : forkIndex + 1, forked: true }; } -/** A paused turn's `updated_at` is when it paused: nothing else writes it while it waits. */ -function completedAt(turn: MemoryIngestTurn): string | null { - const { state } = turn; - switch (state.status) { - case 'running': - return null; - case 'paused': - return turn.updated_at; - case 'done': - case 'error': - case 'cancelled': - return state.completed_at; - } -} - -export type ExchangeReadiness = { ready: true; partial: boolean } | { ready: false }; - -/** - * Closed exchanges are sent as soon as they close; exchanges paused on the user longer than the - * pause timeout are sent as they stand (`partial`), and the later answer follows as a continuation. - */ -export function exchangeReadiness({ - exchange, - now, - pauseTimeoutMs, -}: { - exchange: TurnExchange; - now: Date; - pauseTimeoutMs: number; -}): ExchangeReadiness { - const finishedAt = exchange.completedAt; - if (exchange.status === 'running' || finishedAt === null) { - return { ready: false }; - } - if (exchange.status === 'closed') { - return { ready: true, partial: false }; - } - return now.getTime() - Date.parse(finishedAt) >= pauseTimeoutMs ? { ready: true, partial: true } : { ready: false }; -} - /** `ask_user_question` tool call id → question text, from main-thread model messages. */ export function askUserQuestionsByToolCallId(messages: readonly StoredModelMessage[]): Map { const questions = new Map(); diff --git a/packages/trueforge/src/truefoundry/resolveTrueFoundryMemoryProviderConfig.ts b/packages/trueforge/src/truefoundry/resolveTrueFoundryMemoryProviderConfig.ts index 17047d04ce..178869b5da 100644 --- a/packages/trueforge/src/truefoundry/resolveTrueFoundryMemoryProviderConfig.ts +++ b/packages/trueforge/src/truefoundry/resolveTrueFoundryMemoryProviderConfig.ts @@ -21,15 +21,6 @@ export const TrueFoundryMemorySettingsSchema = z .int() .positive() .default(2 * 60 * 60 * 1000), - /** - * An exchange paused on the user (a question, approval or auth) is sent as-is after this long. - * Paused sessions are looked up again in a lookback-wide range this far back. - */ - pause_timeout_ms: z - .number() - .int() - .positive() - .default(12 * 60 * 60 * 1000), }) .strict(); diff --git a/packages/trueforge/tests/db/memoryIngestStoreContractSuite.ts b/packages/trueforge/tests/db/memoryIngestStoreContractSuite.ts index 383e4c97de..fdcf11cff0 100644 --- a/packages/trueforge/tests/db/memoryIngestStoreContractSuite.ts +++ b/packages/trueforge/tests/db/memoryIngestStoreContractSuite.ts @@ -30,11 +30,10 @@ const UI: SessionMetadata = { 'tfg.source': 'ui' }; const SPEC = AgentSpecSchema.parse({ model: { name: 'test-provider/test-model' }, instructions: 'test' }); const ISO_UTC = /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z$/; const HOUR_AGO = () => new Date(Date.now() - 60 * 60 * 1000); -const IN_A_MINUTE = () => new Date(Date.now() + 60 * 1000); -/** Every session updated since `since`; an empty paused range. */ +/** Every session updated since `since`. */ function recent(since: Date) { - return { tenant_ids: [TENANT], recent_since: since, paused_from: since, paused_until: since }; + return { tenant_ids: [TENANT], recent_since: since }; } export interface MemoryIngestStoreContractEnv { @@ -185,30 +184,12 @@ export function runMemoryIngestStoreContractSuite( expect(await candidateIds()).toEqual(['ui']); }); - it('lists only sessions paused on the user from the paused range', async () => { - await createSession({ session_id: 'paused' }); - await addPausedTurn({ session_id: 'paused', turn_id: 't1' }); - await createSession({ session_id: 'closed' }); - await addTurn({ session_id: 'closed', turn_id: 't1' }); - - // Nothing is recent; both sessions fall inside the paused range. - const candidates = await getEnv().store.listCandidateSessions({ - tenant_ids: [TENANT], - recent_since: IN_A_MINUTE(), - paused_from: HOUR_AGO(), - paused_until: IN_A_MINUTE(), - limit: 10, - }); - expect(candidates.map(c => c.session_id)).toEqual(['paused']); - }); - - it('reports a paused turn with the time it paused', async () => { + it('reports a paused turn as paused', async () => { await createSession({ session_id: 'paused' }); await addPausedTurn({ session_id: 'paused', turn_id: 't1' }); const [turn] = await getEnv().store.listSessionTurns({ session_id: 'paused' }); expect(turn?.state).toEqual({ status: 'paused' }); - expect(turn?.updated_at).toMatch(ISO_UTC); }); it('includes sessions on a registry agent only when its memory is enabled, even from the UI', async () => { diff --git a/packages/trueforge/tests/unit/controller/memoryIngest.test.ts b/packages/trueforge/tests/unit/controller/memoryIngest.test.ts index 7a07cc2bc5..df6f567acf 100644 --- a/packages/trueforge/tests/unit/controller/memoryIngest.test.ts +++ b/packages/trueforge/tests/unit/controller/memoryIngest.test.ts @@ -38,7 +38,6 @@ const CONFIG: MemoryIngestConfig = { maxSessionsPerRun: 50, maxExchangesPerSession: 20, lookbackMs: 120 * MINUTE, - pauseTimeoutMs: 720 * MINUTE, }; interface FakeSession { @@ -71,11 +70,7 @@ class FakeStore implements IMemoryIngestStore { // session.updated_at moves when a turn is created. const sessionUpdatedAt = tip.created_at; const updatedMs = Date.parse(sessionUpdatedAt); - const tipPaused = tip.state.status === 'paused'; - const inRange = - updatedMs >= input.recent_since.getTime() || - (tipPaused && updatedMs >= input.paused_from.getTime() && updatedMs < input.paused_until.getTime()); - if (cursor?.last_ingested_turn_id === tip.turn_id || !inRange) { + if (cursor?.last_ingested_turn_id === tip.turn_id || updatedMs < input.recent_since.getTime()) { continue; } candidates.push({ @@ -248,39 +243,20 @@ describe('runMemoryIngestPass', () => { expect(store.cursors.get('s1')).toMatchObject({ external_thread_id: 'thread-s1', last_error: null }); }); - it('looks back 2 hours for all sessions and 12-14 hours only for paused ones', async () => { + it('looks back 2 hours for sessions', async () => { const store = new FakeStore(); const list = jest.spyOn(store, 'listCandidateSessions'); - // Last turn started 3 hours ago and finished: out of the recent range. + // Last turn started 3 hours ago and finished: out of the lookback. store.add(session({ id: 'stale', turns: [closedTurn('t1', null, -119)] })); - // Paused on a question 13 hours ago: in the paused range, and past the pause timeout. - store.add( - session({ - id: 'paused', - turns: [ - turn({ - id: 'u1', - previous: null, - input: [userMessage('Plan dinner')], - state: pausedState(), - createdAt: at(-12 * 60 * MINUTE), - updatedAt: at(-12 * 60 * MINUTE + MINUTE), - }), - ], - }), - ); + store.add(session({ id: 'recent', turns: [closedTurn('t1', null, 10)] })); const { provider, calls } = fakeProvider(); await runMemoryIngestPass({ deps: deps({ store, provider }), signal: signal() }); expect(list).toHaveBeenCalledWith( - expect.objectContaining({ - recent_since: new Date(NOW.getTime() - 120 * MINUTE), - paused_from: new Date(NOW.getTime() - 14 * 60 * MINUTE), - paused_until: new Date(NOW.getTime() - 12 * 60 * MINUTE), - }), + expect.objectContaining({ recent_since: new Date(NOW.getTime() - 120 * MINUTE) }), ); - expect(calls.map(call => call.exchange.sessionId)).toEqual(['paused']); + expect(calls.map(call => call.exchange.sessionId)).toEqual(['recent']); }); it('sends an exchange that finished just before the run', async () => { @@ -486,44 +462,48 @@ describe('runMemoryIngestPass', () => { ); }); - it('sends a timed-out question as partial, then the answer with its question as a continuation', async () => { + it('does not ingest a paused turn; sends the whole exchange once the answer finishes it', async () => { const store = new FakeStore(); - const paused = turn({ - id: 'u1', - previous: null, - input: [userMessage('Plan dinner')], - state: pausedState(), - createdAt: at(-12.5 * 60 * MINUTE - MINUTE), - updatedAt: at(-12.5 * 60 * MINUTE), - }); const fakeSession = session({ id: 's1', - turns: [paused], + turns: [ + turn({ + id: 'u1', + previous: null, + input: [userMessage('Plan dinner')], + state: pausedState(), + createdAt: at(10 * MINUTE), + }), + ], messages: [{ turn_id: 'u1', message: askCall({ toolCallId: 'call-1', question: 'Which cuisine?' }) }], }); store.add(fakeSession); const { provider, calls } = fakeProvider(); - const run = () => runMemoryIngestPass({ deps: deps({ store, provider }), signal: signal() }); + const logger = fakeLogger(); + const run = () => runMemoryIngestPass({ deps: deps({ store, provider, logger }), signal: signal() }); - await run(); - expect(calls[0]?.exchange.messages.map(m => m.content)).toEqual(['Plan dinner']); - expect(store.cursors.get('s1')?.last_ingested_turn_id).toBe('u1'); + const paused = await run(); + expect(calls).toHaveLength(0); + expect(store.upserts).toHaveLength(0); + expect(paused.sessions_by_outcome).toEqual({ not_ready: 1 }); fakeSession.turns.push( turn({ id: 'a2', previous: 'u1', input: [answer({ toolCallId: 'call-1', content: 'Thai' })], - state: doneState({ completedAt: at(30 * MINUTE), output: reply('Green curry it is') }), - createdAt: at(29 * MINUTE), + state: doneState({ completedAt: at(21 * MINUTE), output: reply('Green curry it is') }), + createdAt: at(20 * MINUTE), }), ); await run(); - expect(calls[1]?.exchange.messages.map(m => m.content)).toEqual([ + expect(calls).toHaveLength(1); + expect(calls[0]?.exchange.lastTurnId).toBe('a2'); + expect(calls[0]?.exchange.messages.map(m => m.content)).toEqual([ + 'Plan dinner', '(Answering "Which cuisine?") Thai', 'Green curry it is', ]); - expect(calls[1]?.externalThreadId).toBe('thread-s1'); }); it('caps exchanges per session and sessions per run; each run starts from the oldest unfinished session', async () => { @@ -601,7 +581,6 @@ describe('runMemoryIngestPass', () => { maxSessionsPerRun: 50, maxExchangesPerSession: 20, lookbackMs: 7_200_000, - pauseTimeoutMs: 43_200_000, }); }); }); diff --git a/packages/trueforge/tests/unit/memory/exchanges.test.ts b/packages/trueforge/tests/unit/memory/exchanges.test.ts index 6911e05057..0e73cf6c67 100644 --- a/packages/trueforge/tests/unit/memory/exchanges.test.ts +++ b/packages/trueforge/tests/unit/memory/exchanges.test.ts @@ -2,16 +2,14 @@ import { activePath, askUserQuestionsByToolCallId, buildExchanges, - exchangeReadiness, MAX_MEMORY_MESSAGE_CHARS, renderExchange, resumePoint, type TurnExchange, } from '../../../src/memory/exchanges'; -import { answer, askCall, at, doneState, pausedState, reply, RUNNING, T0, turn, userMessage } from './turnFixtures'; +import { answer, askCall, at, doneState, pausedState, reply, RUNNING, turn, userMessage } from './turnFixtures'; const MINUTE = 60_000; -const HOUR = 60 * MINUTE; /** u1 asks dinner ideas → agent asks a question (paused) → a2 answers → agent replies. */ function pausedThenAnsweredTurns() { @@ -22,7 +20,6 @@ function pausedThenAnsweredTurns() { input: [userMessage('Dinner ideas? I am vegetarian.')], state: pausedState(), createdAt: at(0), - updatedAt: at(MINUTE), }), turn({ id: 'a2', @@ -99,12 +96,10 @@ describe('buildExchanges', () => { createdAt: at(5 * MINUTE), }); const exchanges = buildExchanges([u1, next]); - expect(exchanges.map(e => [e.turns.map(t => t.turn_id), e.status, e.completedAt])).toEqual([ - [['u1'], 'closed', at(MINUTE)], - [['u2'], 'closed', at(6 * MINUTE)], + expect(exchanges.map(e => [e.turns.map(t => t.turn_id), e.status])).toEqual([ + [['u1'], 'closed'], + [['u2'], 'closed'], ]); - const later = { now: new Date(T0 + 10 * MINUTE), pauseTimeoutMs: 12 * HOUR }; - expect(exchanges.map(exchange => exchangeReadiness({ exchange, ...later }).ready)).toEqual([true, true]); }); it('ends a turn left running (e.g. a crashed executor) when the next user message arrives', () => { @@ -117,7 +112,7 @@ describe('buildExchanges', () => { createdAt: at(2 * MINUTE), }); const [first] = buildExchanges([stale, next]); - expect(first).toMatchObject({ status: 'closed', completedAt: at(2 * MINUTE) }); + expect(first).toMatchObject({ status: 'closed' }); }); it('marks an exchange that starts with a resume turn as a continuation', () => { @@ -169,7 +164,6 @@ describe('renderExchange', () => { }, ], state: pausedState(), - updatedAt: at(MINUTE), }), turn({ id: 'a2', @@ -197,7 +191,6 @@ describe('renderExchange', () => { previous: null, input: [userMessage('x'.repeat(MAX_MEMORY_MESSAGE_CHARS + 500))], state: pausedState(), - updatedAt: at(MINUTE), }), ]), ); @@ -211,41 +204,6 @@ describe('renderExchange', () => { }); }); -describe('exchangeReadiness', () => { - const options = { pauseTimeoutMs: 12 * HOUR }; - - it('sends a closed exchange as soon as it closes', () => { - const exchange = only(buildExchanges(pausedThenAnsweredTurns())); - if (exchange.completedAt === null) { - throw new Error('fixture'); - } - expect(exchangeReadiness({ exchange, now: new Date(exchange.completedAt), ...options })).toEqual({ - ready: true, - partial: false, - }); - }); - - it('sends a paused exchange as partial only after the pause timeout', () => { - const [u1] = pausedThenAnsweredTurns(); - if (u1 === undefined) { - throw new Error('fixture'); - } - const exchange = only(buildExchanges([u1])); - expect(exchangeReadiness({ exchange, now: new Date(T0 + 11 * HOUR), ...options })).toEqual({ ready: false }); - expect(exchangeReadiness({ exchange, now: new Date(T0 + 13 * HOUR), ...options })).toEqual({ - ready: true, - partial: true, - }); - }); - - it('never sends a running exchange', () => { - const exchange = only( - buildExchanges([turn({ id: 'r', previous: null, input: [userMessage('hi')], state: RUNNING })]), - ); - expect(exchangeReadiness({ exchange, now: new Date(T0 + 24 * HOUR), ...options })).toEqual({ ready: false }); - }); -}); - describe('resumePoint', () => { const closed = (id: string, previous: string | null) => turn({ id, previous, input: [userMessage(id)], state: doneState({ completedAt: at(0), output: null }) }); diff --git a/packages/trueforge/tests/unit/memory/turnFixtures.ts b/packages/trueforge/tests/unit/memory/turnFixtures.ts index 453228dcc4..729328046e 100644 --- a/packages/trueforge/tests/unit/memory/turnFixtures.ts +++ b/packages/trueforge/tests/unit/memory/turnFixtures.ts @@ -58,17 +58,14 @@ export function turn({ input, state, createdAt = at(0), - updatedAt = createdAt, }: { id: string; previous: string | null; input: LegacyTurnInputItem[]; state: TurnState; createdAt?: string; - /** For a paused turn, when it paused. */ - updatedAt?: string; }): MemoryIngestTurn { - return { turn_id: id, previous_turn_id: previous, input, state, created_at: createdAt, updated_at: updatedAt }; + return { turn_id: id, previous_turn_id: previous, input, state, created_at: createdAt }; } export function userMessage(text: string): LegacyTurnInputItem { diff --git a/packages/trueforge/tests/unit/truefoundry/resolveTrueFoundryMemoryProviderConfig.test.ts b/packages/trueforge/tests/unit/truefoundry/resolveTrueFoundryMemoryProviderConfig.test.ts index cf492fef89..f7fc3a34a1 100644 --- a/packages/trueforge/tests/unit/truefoundry/resolveTrueFoundryMemoryProviderConfig.test.ts +++ b/packages/trueforge/tests/unit/truefoundry/resolveTrueFoundryMemoryProviderConfig.test.ts @@ -14,7 +14,6 @@ const DEFAULT_SETTINGS = { max_sessions_per_run: 50, max_exchanges_per_session: 20, lookback_ms: 7_200_000, - pause_timeout_ms: 43_200_000, }; const ZEP_ON = {