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/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index 62bf8922a5..95789e59f2 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -11075,7 +11075,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. Resume with from_event_id or after_event_id.", "parameters": [ { "description": "Session identifier.", @@ -11101,18 +11101,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": "Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.", "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": "Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.", "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": { @@ -11164,10 +11186,10 @@ } } }, - "description": "Cannot subscribe — the live stream no longer exists." + "description": "The turn stream is no longer available or resumable." } }, - "summary": "Subscribe to a running turn", + "summary": "Subscribe to a turn stream", "tags": [ "Agent Sessions" ], @@ -11177,7 +11199,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 62bf8922a5..95789e59f2 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -11075,7 +11075,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. Resume with from_event_id or after_event_id.", "parameters": [ { "description": "Session identifier.", @@ -11101,18 +11101,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": "Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.", "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": "Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.", "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": { @@ -11164,10 +11186,10 @@ } } }, - "description": "Cannot subscribe — the live stream no longer exists." + "description": "The turn stream is no longer available or resumable." } }, - "summary": "Subscribe to a running turn", + "summary": "Subscribe to a turn stream", "tags": [ "Agent Sessions" ], @@ -11177,7 +11199,7 @@ "x-fern-sdk-method-name": "subscribe_to_turn", "x-fern-streaming": { "format": "sse", - "resumable": true + "resumable": false } } }, diff --git a/packages/assistant-ui-runtime/src/convertTurnMessages.ts b/packages/assistant-ui-runtime/src/convertTurnMessages.ts index c24918fd1e..a314571414 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, @@ -77,7 +78,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', @@ -416,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; } @@ -494,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); } } @@ -745,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 ?? []); @@ -1024,7 +1032,7 @@ async function ingestTurnEventsIntoFold( ...(pageToken != null ? { pageToken } : {}), }), ); - ingestCollectedEventsIntoFold(foldState, events); + ingestCollectedEventsIntoFold(foldState, events, turnId); } async function fetchTurnEvents( @@ -1047,7 +1055,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 || @@ -1056,7 +1068,7 @@ function ingestCollectedEventsIntoFold(foldState: PeerThreadFoldState, events: P ) { continue; } - ingestTurnEvent(foldState, event); + ingestTurnEvent(foldState, event, turnId); } } @@ -1571,7 +1583,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( @@ -1852,8 +1864,11 @@ export async function* streamTurnEvents( groupRootBaseline?: readonly string[], 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; @@ -1906,24 +1921,45 @@ 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 (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 || + 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, earliestOpenModelMessageId(foldState, activeTurnId))); + } + } if (event.type === EVENT_TYPE.TURN_CREATED) { onTurnIdAvailable?.(event.turnId); @@ -1942,29 +1978,37 @@ export async function* streamTurnEvents( if (event.type === EVENT_TYPE.USER_MCP_AUTH_CONTINUE) { pendingMcpAuth = undefined; - ingestStreamEvent(foldState, event); + ingestStreamEvent(foldState, event, activeTurnId); continue; } 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; } @@ -1979,13 +2023,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; } @@ -1995,16 +2043,20 @@ 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; } - 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 0da2b8d997..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,16 +447,62 @@ function compactionPartFromStarted({ }; } -export function ingestStreamEvent(state: PeerThreadFoldState, message: TurnStreamingEvent): boolean { +/** 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; + } + if (earliest === undefined || id < earliest) { + earliest = id; + } + } + } + return earliest; +} + +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)) { - return false; + // 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 { + state.ingestedEventIds.add(message.id); } - state.ingestedEventIds.add(message.id); } if (isFoldNeutralInboundEvent(message)) { @@ -490,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/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 7c520d8826..e83ed1edf3 100644 --- a/packages/assistant-ui-runtime/src/server/types.ts +++ b/packages/assistant-ui-runtime/src/server/types.ts @@ -424,6 +424,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..f49a97f8aa 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 } from './turnStreamUpdate.js'; +import { earliestOpenModelMessageId, PeerThreadFoldState } from './foldPeerThreads.js'; +import { + resumeCursorCoveringOpenModelMessages, + type TurnStreamUpdate, + type TurnSubscribeCursor, +} from './turnStreamUpdate.js'; export interface StreamTurnOptions { userMessage?: UserMessageContent; @@ -31,6 +35,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 +61,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 +76,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 +96,16 @@ export async function* resumeTurnStream( server.subscribeToTurn({ sessionId, turnId, - ...(afterSequenceNumber != null ? { afterSequenceNumber } : {}), + ...(resumeCursorCoveringOpenModelMessages(subscribeCursor, earliestOpenModelMessageId(foldState, turnId)) ?? + {}), abortSignal, }), foldState, groupRootBaseline, 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 9db3271305..572e82abb8 100644 --- a/packages/assistant-ui-runtime/src/turnStreamUpdate.ts +++ b/packages/assistant-ui-runtime/src/turnStreamUpdate.ts @@ -1,8 +1,108 @@ 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. 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 }, + 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 }; +} + +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?: unknown }[], +): 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 && isRunningState(event.state)); + 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 + * 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 2252585afb..c3c123c193 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 type { TurnStreamUpdate } from './turnStreamUpdate.js'; +import { + createSubscribeReplayProgress, + cursorFromTurnEvents, + type TurnStreamUpdate, + type TurnSubscribeCursor, +} from './turnStreamUpdate.js'; export interface UseTrueForgeAgentMessagesOptions { server: AgentChatServer; @@ -101,6 +106,27 @@ 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; + } + const events = []; + for (let index = history.length - 1; index >= 0; index -= 1) { + const item = history[index]; + if (item?.turnId === turnId && 'id' in item.event && typeof item.event.id === 'string') { + events.push(item.event); + } + } + return cursorFromTurnEvents(events); } interface RunStreamReconnect { @@ -243,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]; } @@ -581,14 +611,26 @@ 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; + // Inclusive resume replays the message shell and every delta already seen. + // 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 => { lastSequenceNumber = sequenceNumber; consecutiveFailures = 0; }, + onSubscribeCursor: cursor => { + subscribeCursor = cursor; + if (replayProgress.observe(cursor)) { + consecutiveFailures = 0; + } + }, }; const flushPendingStreamUpdate = () => { @@ -637,6 +679,7 @@ export function useTrueForgeAgentMessages({ update, segmentStatus, ...(updateSequenceNumber != null ? { lastSequenceNumber: updateSequenceNumber } : {}), + ...(subscribeCursor != null ? { subscribeCursor } : {}), }, pendingMcpAuth, ...(runningTurn != null ? { runningTurn } : {}), @@ -663,6 +706,7 @@ export function useTrueForgeAgentMessages({ }; const readStream = (signal: AbortSignal): AsyncGenerator => { + replayProgress.beginAttempt(); if (!useSubscribe) { return createStream(signal, track); } @@ -672,9 +716,10 @@ export function useTrueForgeAgentMessages({ turnIdRef.current, snapshotRef.current.fold, signal, - lastSequenceNumber, + subscribeCursor, snapshotRef.current.groupRootBaseline, track.onSequenceNumber, + track.onSubscribeCursor, ); }; @@ -756,6 +801,7 @@ export function useTrueForgeAgentMessages({ ...prev.activeStream, segmentStatus, ...(lastSequenceNumber != null ? { lastSequenceNumber } : {}), + ...(subscribeCursor != null ? { subscribeCursor } : {}), }, }); return commitActiveStream(marked); @@ -881,9 +927,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 +1130,7 @@ export function useTrueForgeAgentMessages({ groupRootBaseline, handleGatewayTurnId, track.onSequenceNumber, + track.onSubscribeCursor, ), turnIdRef, { @@ -1136,8 +1184,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 +1194,10 @@ export function useTrueForgeAgentMessages({ turnId, snapshotRef.current.fold, signal, - afterSequenceNumber, + cursorForTurn(snapshotRef.current, turnId), snapshotRef.current.groupRootBaseline, track.onSequenceNumber, + track.onSubscribeCursor, ), { current: turnId }, { @@ -1283,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 @@ -1429,7 +1476,6 @@ export function useTrueForgeAgentMessages({ if (resumeSessionId == null) { return; } - const active = snapshotRef.current.activeStream; await runStream( (signal, track) => resumeTurnStream( @@ -1438,9 +1484,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..bad00fe5cf 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'; @@ -1879,6 +1879,102 @@ 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('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( @@ -2289,6 +2385,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( @@ -3922,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/assistant-ui-runtime/test/foldPeerThreads.test.ts b/packages/assistant-ui-runtime/test/foldPeerThreads.test.ts index ab9ab85c73..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 = { @@ -55,6 +88,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/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/turnStreamUpdate.test.ts b/packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts new file mode 100644 index 0000000000..d2c05017cd --- /dev/null +++ b/packages/assistant-ui-runtime/test/turnStreamUpdate.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, it } from 'vitest'; +import { + createSubscribeReplayProgress, + cursorFromTurnEvents, + 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('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(); + 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/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-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/agent-session/store/InMemorySessionStore.ts b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts index 5d71550e97..1cdd4a0339 100644 --- a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts @@ -68,7 +68,6 @@ import { TurnNotRunningError, TurnOwnershipLostError, } from './SessionStoreErrors'; - /* 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; 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/store/storeContractSuite.ts b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts index 20b4e243e3..b9f1f3b114 100644 --- a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts +++ b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts @@ -5,6 +5,7 @@ import { EventType } from '../../../src/agent-session/schemas/events'; import { CancellationReason, type NonTerminalTurnState, + type TerminalTurnState, type TurnNonTerminalMetrics, } from '../../../src/agent-session/schemas/turn'; import type { ISessionStore } from '../../../src/agent-session/store/ISessionStore'; @@ -1418,6 +1419,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, @@ -1485,6 +1497,99 @@ 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), + metrics: createEmptyAgentThreadMetrics(), + }); + 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), + turn_metrics: turnMetricsFor(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), + metrics: createEmptyAgentThreadMetrics(), + }), + ).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), + turn_metrics: turnMetricsFor(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); 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-sdk/reference.md b/packages/trueforge-sdk/reference.md index 9e7f9c6cc3..c9789f1397 100644 --- a/packages/trueforge-sdk/reference.md +++ b/packages/trueforge-sdk/reference.md @@ -2992,7 +2992,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. 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 8449695d7f..649a80a80b 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. 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( @@ -2189,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({ @@ -2234,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 03af7554c8..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,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. */ + /** 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; + /** 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 3b1858b11c..44af448431 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/.env.example b/packages/trueforge/.env.example index 21761a0af2..3e4c462158 100644 --- a/packages/trueforge/.env.example +++ b/packages/trueforge/.env.example @@ -120,6 +120,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 @@ -327,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/apis/schedules.ts b/packages/trueforge/src/apis/schedules.ts index fe61e4bc1e..16203cb267 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 when a schedule run supersedes a turn that is not in this process. */ - sessionStore: Pick, 'getTurn'>; + /** Previous-turn cancel, and rebuilding a missing event stream. */ + sessionStore: Pick, 'getTurn' | 'listTurnEvents'>; webSearchProviderStore: IWebSearchProviderStore; } diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 17d9b529ab..93d29c7053 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -86,7 +86,15 @@ import { subscribeTurnRoute, } from '../routes/turnRoutes'; import type { ActiveTurnRegistry } from '../runtime/activeTurns'; -import { StreamGoneError, type EventSubscription, type EventSubscriptionRegistry } from '../runtime/event-subscription'; +import { + eventIdTimestampMs, + isLegacySequenceEntryId, + LEGACY_SEQUENCE_ID_MAX, + StreamGoneError, + type EventSubscription, + type EventSubscriptionPollOptions, + type EventSubscriptionRegistry, +} from '../runtime/event-subscription'; import { validateSandboxFilePath } from '../runtime/sandboxFilePath'; import { buildTurnSandbox, @@ -95,6 +103,7 @@ import { resolveSandboxEnvironment, resolveSandboxProvider, } from '../runtime/sessionResources'; +import { restoreTurnEventLog } from '../runtime/turnEventLogRestore'; import { MAX_SESSION_TITLE_LENGTH } from '../schemas/session'; import { assertGatewayMetadataRequestHeaders } from '../truefoundry/gatewayMetadata'; import { newId } from '../utils/id'; @@ -229,8 +238,8 @@ export type BeginTurnExecutionDeps = Pick< skillStore: Pick; /** Undefined for sessions that never use memory (scheduled runs). */ memory: TenantMemoryDeps | undefined; - /** Loads the previous turn when create must cancel a run that is not in this process. */ - sessionStore: Pick; + /** Previous-turn cancel, and rebuilding a missing event stream. */ + sessionStore: Pick; redis?: RedisClient | undefined; }; @@ -477,9 +486,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), }; } @@ -552,16 +561,14 @@ export async function drainTurnEvents(input: { executionTimeout: TurnExecutionTimeout; logger: Logger; pausedKeepalive: TurnPausedKeepalive; - onEvent?: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; + onEvent?: (event: TurnStreamingEvent) => Promise; }): Promise { 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. - const sequenceNumber = await turnEventStream.put(event, { - streamTTLSeconds: streamTTLSecondsFor(event), - }); + await turnEventStream.put(event, { streamTTLSeconds: streamTTLSecondsFor(event) }); if (event.type === EventType.TURN_UPDATE) { if (event.state.status === 'paused') { pausedKeepalive.start(); @@ -571,11 +578,14 @@ export async function drainTurnEvents(input: { executionTimeout.start(); } } - await onEvent?.(event, sequenceNumber); + await onEvent?.(event); } } 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) { + 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), @@ -732,7 +742,7 @@ export async function beginTurnExecution( turn, }); - // Held for the whole turn; the stream's sequence counter dies with it. + // 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 { @@ -791,8 +801,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, restore a missing event stream from the + * durable log, and drain in the background. Inbox replay happens inside rebuild. */ async function resumePausedTurn(input: { session: SessionHandle; @@ -860,8 +870,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 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, + }); } catch (error) { input.deps.logger.warn('Paused turn stream resume failed', extractErrorLogFields(error)); executionTimeout.clear(); @@ -1895,11 +1911,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) { @@ -2002,7 +2018,21 @@ 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 (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({ @@ -2026,11 +2056,22 @@ export function createTurnsRouter(deps: TurnsRouterDeps(deps: TurnsRouterDeps LEGACY_SEQUENCE_ID_MAX) { + throw new HTTPException(400, { + message: 'after_sequence_number is not supported above 50000.', + }); + } + includeSequenceId = true; + } // One lifecycle controller: server-side timeout, client disconnect, and // normal teardown all abort it, which ends poll even while it is parked // waiting for events. @@ -2053,16 +2117,25 @@ export function createTurnsRouter(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(turnEventSsePayload(event, includeSequenceId ? sequenceNumber : undefined)); if (event.type === EventType.TURN_DONE) { break; } diff --git a/packages/trueforge/src/config.ts b/packages/trueforge/src/config.ts index 0469906f93..5508fa843d 100644 --- a/packages/trueforge/src/config.ts +++ b/packages/trueforge/src/config.ts @@ -809,6 +809,12 @@ export interface SharedServerConfiguration { * Env: `TURN_PAUSED_KEEPALIVE_SECONDS`. Default 300. */ TURN_PAUSED_KEEPALIVE_SECONDS: number; + /** + * Max ms to wait when acquiring the per-turn lock. + * `0` waits until the holder releases it. + * Env: `TURN_LOCK_ACQUIRE_TIMEOUT_MS`. Default 0. + */ + TURN_LOCK_ACQUIRE_TIMEOUT_MS: number; /** * Max ms to wait for one create-turn SSE write before dropping that client * and keeping the turn drain running. Env: `TURN_SSE_STREAM_WRITE_TIMEOUT_MS`. Default 60000. @@ -1122,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; @@ -1241,6 +1247,11 @@ const shared: SharedServerConfiguration = { raw: getEnv('TURN_PAUSED_KEEPALIVE_SECONDS'), defaultValue: 300, }), + TURN_LOCK_ACQUIRE_TIMEOUT_MS: parseNonNegativeInt({ + envKey: 'TURN_LOCK_ACQUIRE_TIMEOUT_MS', + raw: getEnv('TURN_LOCK_ACQUIRE_TIMEOUT_MS'), + defaultValue: 0, + }), TURN_SSE_STREAM_WRITE_TIMEOUT_MS: parsePositiveInt({ envKey: 'TURN_SSE_STREAM_WRITE_TIMEOUT_MS', raw: getEnv('TURN_SSE_STREAM_WRITE_TIMEOUT_MS'), 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 9b70adc1c1..1535abd07f 100644 --- a/packages/trueforge/src/db/memoryIngestStore.ts +++ b/packages/trueforge/src/db/memoryIngestStore.ts @@ -59,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; } @@ -73,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 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)`; 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 => 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..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 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)`; /** 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 => diff --git a/packages/trueforge/src/main.ts b/packages/trueforge/src/main.ts index 58ce33da69..ea91204ce4 100644 --- a/packages/trueforge/src/main.ts +++ b/packages/trueforge/src/main.ts @@ -562,6 +562,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/memory/exchanges.ts b/packages/trueforge/src/memory/exchanges.ts index b1eed5c9a9..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 finished 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; } @@ -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'; @@ -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.state), continuation: !hasUserMessage(turn), }; exchanges.push(current); } else { current.turns.push(turn); current.status = statusAfter(turn.state); - current.completedAt = completedAt(turn.state); } if (current.status === 'closed') { current = undefined; @@ -130,37 +127,6 @@ 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; -} - -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/routes/turnRoutes.ts b/packages/trueforge/src/routes/turnRoutes.ts index a9a1a9b05b..01c26c7d42 100644 --- a/packages/trueforge/src/routes/turnRoutes.ts +++ b/packages/trueforge/src/routes/turnRoutes.ts @@ -294,12 +294,11 @@ export const subscribeTurnRoute = createRoute({ method: 'get', path: '/{session_id}/turns/{turn_id}/subscribe', tags: [OpenApiTag.AGENT_SESSIONS], - summary: 'Subscribe to a running turn', - 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).', + summary: 'Subscribe to a turn stream', + 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: true }, + 'x-fern-streaming': { format: 'sse', resumable: false }, request: { params: TurnIdParamsSchema, query: SubscribeTurnQuerySchema, @@ -327,7 +326,7 @@ export const subscribeTurnRoute = createRoute({ }, 412: { content: { 'application/json': { schema: RequestErrorResponseSchema } }, - description: 'Cannot subscribe — the live stream no longer exists.', + description: 'The turn stream is no longer available or resumable.', }, }, }); diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index d8670f0397..331a5cdc4f 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -5,11 +5,13 @@ * `track()` owns registration and cleanup around the stream lifecycle. */ import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; -import { getOperationMetrics } from '@truefoundry/trueforge-core/core'; -import { Mutex } from 'async-mutex'; +import { extractErrorLogFields, getOperationMetrics } from '@truefoundry/trueforge-core/core'; +import { E_TIMEOUT, Mutex, withTimeout, type MutexInterface } 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; @@ -29,6 +31,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 => { @@ -60,11 +74,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' }); } @@ -73,6 +90,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); @@ -81,9 +100,20 @@ export class ActiveTurnRegistry { if (!existing) { this.turnLocks.set(key, mutex); } + // 0 waits forever. withTimeout(0) would reject on the next timer turn. + let guarded: MutexInterface = mutex; + if (this.lockAcquireTimeoutMs > 0) { + guarded = withTimeout( + mutex, + this.lockAcquireTimeoutMs, + new TurnLockAcquireTimeoutError({ + timeoutMs: this.lockAcquireTimeoutMs, + cause: E_TIMEOUT, + }), + ); + } try { - // TODO (chiragjn): support acquire timeout - return await mutex.runExclusive(fn); + return await guarded.runExclusive(fn); } finally { if (!mutex.isLocked()) { this.turnLocks.delete(key); @@ -267,7 +297,7 @@ export class ActiveTurnRegistry { } 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; @@ -280,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/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 328aff64dc..f5ddaf0478 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 (dense - * sequences, whole-stream TTL, poll-forever generators, same errors). - * 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 { @@ -12,20 +11,34 @@ import { type EventSubscriptionPutOptions, type SequencedEvent, } from '.'; +import { + entryIdMilliseconds, + EventIdSeek, + eventIdTimestampMs, + isEntryAfterCursor, + isLegacySequenceEntryId, +} from './entryId'; +import { applyEventLogRestore, type RestoreLogOptions, type RestoreLogOutcome } from './restore'; + +function ttlDeadline(streamTTLSeconds: number | undefined): number | undefined { + return streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? Date.now() + streamTTLSeconds * 1_000 : undefined; +} -interface InMemoryStream { - /** Dense log: the event at index i has sequence_number i + 1 (1-indexed). */ - events: SequencedEvent[]; - /** Undefined = no TTL set yet (mirrors a Redis key without EXPIRE). */ +interface StoredEvent { + entryId: string; + event: T; +} + +interface InMemoryStream { + /** Appended in entry-id order. */ + events: StoredEvent[]; + /** Undefined until the first put/restore 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. - */ -export class InMemoryEventStreamStore { +/** Owns per-stream logs and TTL timers for this process. */ +export class InMemoryEventStreamStore { private readonly streams = new Map>(); /** Rings per stream id on append and expiry so parked pollers wake instantly. */ private readonly changes = new EventEmitter(); @@ -35,29 +48,71 @@ 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); + /** + * 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; 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); + } - const sequenceNumber = stream.events.length + 1; - // Deep-copy so the log is isolated from caller mutations (mirrors JSON round-trip in Redis). - stream.events.push({ ...structuredClone(event), sequence_number: sequenceNumber }); + /** + * 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; + } + 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); + return true; + } + + readTailEntryId(streamId: string): string | undefined { + const stream = this.getLiveStream(streamId); + return stream?.events[stream.events.length - 1]?.entryId; + } + + firstEntryId(streamId: string): string | undefined { + return this.getLiveStream(streamId)?.events[0]?.entryId; + } - if (streamTTLSeconds && streamTTLSeconds > 0) { - stream.expiresAtMs = Date.now() + streamTTLSeconds * 1_000; - this.scheduleExpiry(streamId, stream); + hasEventId({ streamId, eventId }: { streamId: string; eventId: string }): boolean { + return this.getLiveStream(streamId)?.events.some(stored => stored.event.id === eventId) ?? false; + } + + /** Drops the log the way a TTL expiry does. */ + evictStream(streamId: string): void { + const stream = this.streams.get(streamId); + if (stream === undefined) { + return; } - this.changes.emit(streamId); - return sequenceNumber; + this.dropStream(streamId, stream); } - /** Refreshes TTL on a live stream. Sequence numbers already follow the log length. */ - refreshStreamTtl(streamId: string, streamTTLSeconds?: number): void { + refreshStreamTtl(streamId: string, streamTTLSeconds: number | undefined): void { + const deadline = ttlDeadline(streamTTLSeconds); + if (deadline === undefined) { + return; + } 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); } @@ -114,44 +169,152 @@ export class InMemoryEventStreamStore { ); stream.expiryTimer.unref(); } + + 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. + */ + 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 { +export class InMemoryEventSubscription implements EventSubscription { constructor( private readonly store: InMemoryEventStreamStore, private readonly streamId: string, ) {} - resumeProducer(options?: EventSubscriptionPutOptions): Promise { - this.store.refreshStreamTtl(this.streamId, options?.streamTTLSeconds); + restoreLog(options: RestoreLogOptions): Promise { + return applyEventLogRestore({ + readPages: options.readPages, + readTailEntryId: () => Promise.resolve(this.store.readTailEntryId(this.streamId)), + 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(); + }, + }); + } + + put(event: T, options?: EventSubscriptionPutOptions): Promise { + this.store.appendEvent({ + streamId: this.streamId, + event, + streamTTLSeconds: options?.streamTTLSeconds, + }); return Promise.resolve(); } - put(event: T, options?: EventSubscriptionPutOptions): Promise { - return Promise.resolve(this.store.append(this.streamId, event, options?.streamTTLSeconds)); + hasEventId(eventId: string): Promise { + return Promise.resolve(this.store.hasEventId({ streamId: this.streamId, eventId })); + } + + firstEntryId(): Promise { + return Promise.resolve(this.store.firstEntryId(this.streamId)); } 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. + // Same admission window as the shared subscribe threshold. 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(); } - 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". 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; + const readCursor = + afterSequenceNumber === undefined || afterSequenceNumber === 0 ? '0-0' : `${String(afterSequenceNumber + 1)}-0`; + let index = 0; + 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) { + 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; + } + 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); + let index = 0; + let scanning = true; for (;;) { if (signal?.aborted) { return; @@ -161,7 +324,16 @@ export class InMemoryEventSubscription implements EventSubscri throw new StreamGoneError(this.streamId); } const end = live.events.length; - if (cursor >= end) { + 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) { @@ -172,16 +344,24 @@ 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)}`); + 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)}`); + } + for (const event of seek.push(this.polled(stored))) { + yield event; } - cursor += 1; - yield structuredClone(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 ab03654f23..676671407f 100644 --- a/packages/trueforge/src/runtime/event-subscription/index.ts +++ b/packages/trueforge/src/runtime/event-subscription/index.ts @@ -1,33 +1,67 @@ import type { RedisClient } from '@truefoundry/trueforge-core/request-reply'; import { InMemoryEventStreamStore, InMemoryEventSubscription } from './inMemory'; import { RedisEventSubscription } from './redis'; +import type { RestoreLogOptions, RestoreLogOutcome } from './restore'; + +export type { RestoreLogOptions, RestoreLogOutcome } from './restore'; export interface EventSubscriptionPutOptions { /** Sets the lifetime of the whole stream after this append. */ streamTTLSeconds?: number | 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 { eventIdTimestampMs, isLegacySequenceEntryId, LEGACY_SEQUENCE_ID_MAX } from './entryId'; -export interface EventSubscriptionPollOptions { +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 { /** - * 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 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. */ - resumeProducer(options?: EventSubscriptionPutOptions): Promise; + restoreLog(options: RestoreLogOptions): 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; + + /** 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 @@ -36,25 +70,24 @@ 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(); + 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) { diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index 3cf26c7c32..a6eac50e4d 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, @@ -8,58 +9,197 @@ import { type EventSubscriptionPutOptions, type SequencedEvent, } from '.'; +import { entryIdMilliseconds, eventIdSeekCursor, eventIdTimestampMs, isLegacySequenceEntryId } from './entryId'; +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). */ const SUBSCRIBE_STREAM_POLL_SLEEP_INTERVAL_MS = 1_000; -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; -} +/** 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 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]=payload id ARGV[2]=exclusive start ARGV[3]=batch + * ARGV[4]='1' on the first step + * + * 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_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[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]} + end + end + return {'missing', cursor} +end +local previous = cursor +for i = 1, #entries do + if payloadId(entries[i][2]) == targetId then + return {'found', previous} + 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 +end +for i = 1, #entries do + if payloadId(entries[i][2]) == targetId then + return {'found', entries[i][1], tip} + end +end +if #entries < batch then + if tip == '' then + return {'missing', '0-0', tip} + end + return {'missing', tip, tip} +end +return {'more', entries[#entries][1], tip} +`; /** - * 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. + * Live put. KEYS[1]=stream ARGV[1]=ms ARGV[2]=event id ARGV[3]=payload ARGV[4]=ttl + * + * 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. */ -export class RedisEventSubscription implements EventSubscription { - private nextSequenceNumber = 1; +const PUT_STREAM_ENTRY = ` +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[1], ms .. '-*', 'id', ARGV[2], 'data', ARGV[3]) +local ttl = tonumber(ARGV[4]) +if ttl > 0 then + redis.call('EXPIRE', KEYS[1], ttl) +end +return 'ok' +`; + +function ttlArgument(streamTTLSeconds: number | undefined): string { + return String(streamTTLSeconds !== undefined && streamTTLSeconds > 0 ? streamTTLSeconds : 0); +} + +interface SeekStep { + state: 'more' | 'found' | 'missing'; + cursor: string; + /** Only set by the backward step. */ + tip: string; +} +function parseSeekReply(value: unknown): SeekStep { + if (Array.isArray(value) && (value.length === 2 || value.length === 3)) { + const state: unknown = value[0]; + const cursor: unknown = value[1]; + const tip: unknown = value[2] ?? ''; + if ( + (state === 'more' || state === 'found' || state === 'missing') && + typeof cursor === 'string' && + typeof tip === 'string' + ) { + return { state, cursor, tip }; + } + } + throw new Error('Unexpected event-id seek reply'); +} + +/** 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 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); - } + restoreLog(options: RestoreLogOptions): Promise { + return applyEventLogRestore({ + readPages: options.readPages, + readTailEntryId: () => this.readTailEntryId(), + installIfAbsent: pages => this.installIfAbsent(pages, options.streamTTLSeconds), + refreshTtl: () => this.refreshTtl(options.streamTTLSeconds), + }); } - 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), + async put(event: T, options?: EventSubscriptionPutOptions): Promise { + const reply: unknown = await this.redis.eval(PUT_STREAM_ENTRY, { + keys: [this.streamId], + arguments: [ + this.entryMillisecondFor(event), + event.id, + JSON.stringify(event), + ttlArgument(options?.streamTTLSeconds), + ], }); - if (options?.streamTTLSeconds && options.streamTTLSeconds > 0) { - multi.expire(this.streamId, options.streamTTLSeconds); + if (reply !== 'ok') { + throw new Error(`Unexpected stream put reply on ${this.streamId}`); } - await multi.exec(); - return sequenceNumber; + } + + async hasEventId(eventId: string): Promise { + const { found } = await this.seekEventIdCursor({ mode: 'from', eventId }); + return found; + } + + async firstEntryId(): Promise { + const entries = await this.redis.xRange(this.streamId, '-', '+', { COUNT: 1 }); + return entries[0]?.id; } async assertSubscribable(): Promise { @@ -74,10 +214,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 @@ -85,46 +233,229 @@ 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); + } + } + } + + /** + * 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, + ): 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, signal }); + for (;;) { + const messages = await this.readAfter(cursor, signal); + if (messages === undefined) { + return; + } + if (messages.length === 0) { + if (!(await this.waitForEntries(signal))) { + return; } + 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; + } + } + } + + /** + * 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 }> { + const floor = eventIdSeekCursor({ + eventId: input.eventId, + firstEntryId: await this.firstEntryId(), + }); + 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 = + 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; + tip = step.tip; + first = false; + } + return { cursor: cursor === '' ? floor : cursor, found: false }; + } + + 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 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 readTailEntryId(): Promise { + const entries = await this.redis.xRevRange(this.streamId, '+', '-', { COUNT: 1 }); + return entries[0]?.id; + } + + private async refreshTtl(streamTTLSeconds: number | undefined): Promise { + if (streamTTLSeconds !== undefined && streamTTLSeconds > 0) { + await this.redis.expire(this.streamId, streamTTLSeconds); + } + } - 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, - }); + /** + * 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) })); } - yield { ...event, sequence_number: sequenceNumberFromEntryId(this.streamId, id) }; + 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 { + return String(eventIdTimestampMs(event.id) ?? Date.now()); + } } 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..e749f61d5d --- /dev/null +++ b/packages/trueforge/src/runtime/event-subscription/restore.ts @@ -0,0 +1,39 @@ +/** + * 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, in the order to append. */ + readPages: () => AsyncIterable; + streamTTLSeconds?: number | undefined; +} + +/** + * `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: { + readPages: () => AsyncIterable; + readTailEntryId: () => Promise; + installIfAbsent: (pages: AsyncIterable) => Promise; + refreshTtl: () => Promise; +}): Promise { + if ((await input.readTailEntryId()) !== undefined) { + await input.refreshTtl(); + return 'present'; + } + const installed = await input.installIfAbsent(input.readPages()); + 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 new file mode 100644 index 0000000000..8ce4288f1b --- /dev/null +++ b/packages/trueforge/src/runtime/turnEventLogRestore.ts @@ -0,0 +1,100 @@ +import type { ISessionStore, TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session'; +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. + */ +export async function restoreTurnEventLog(input: { + sessionStore: Pick, 'listTurnEvents'>; + subscription: EventSubscription; + logger: Logger; + sessionId: string; + turnId: string; + streamTTLSeconds: number; +}): Promise { + 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, + limit: RESTORE_PAGE_LIMIT, + page_token: pageToken, + order: 'asc', + }); + 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/src/schemas/turn.ts b/packages/trueforge/src/schemas/turn.ts index b40a1796cf..47d3663497 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('Deprecated. Rejected on new turns. Use from_event_id or after_event_id instead.'), ), + 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/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 22d007a2cb..fdcf11cff0 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'; @@ -28,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 { @@ -101,6 +102,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); @@ -165,35 +184,12 @@ export function runMemoryIngestStoreContractSuite( expect(await candidateIds()).toEqual(['ui']); }); - it('lists only sessions paused on the user from the paused range', async () => { + it('reports a paused turn as paused', 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 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']); + await addPausedTurn({ session_id: 'paused', turn_id: 't1' }); + + const [turn] = await getEnv().store.listSessionTurns({ session_id: 'paused' }); + expect(turn?.state).toEqual({ status: 'paused' }); }); it('includes sessions on a registry agent only when its memory is enabled, even from the UI', async () => { diff --git a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts index 6960f7924c..9384ede144 100644 --- a/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts +++ b/packages/trueforge/tests/runtime/event-subscription/contractSuite.ts @@ -1,7 +1,12 @@ -import type { EventSubscriptionRegistry, SequencedEvent } from '../../../src/runtime/event-subscription'; - +import { + StreamGoneError, + 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; type: string; payload: { index: number; text: string }; } @@ -15,8 +20,16 @@ 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)}` } }; +} + +/** 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). */ @@ -24,10 +37,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) { @@ -44,21 +57,29 @@ 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 () => { - const registry = createRegistry(); +export function runEventSubscriptionContractSuite( + setup: () => { + registry: EventSubscriptionRegistry; + /** Drops the log the way a TTL expiry does. */ + evictStream: (streamId: string) => Promise; + }, +) { + it('put appends events in order without a sequence number', async () => { + const { registry } = setup(); 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 () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); @@ -68,11 +89,11 @@ 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 () => { - const registry = createRegistry(); + const { registry } = setup(); const streamId = nextStreamId(); const producer = registry.get(streamId); @@ -82,46 +103,239 @@ 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 () => { - const registry = createRegistry(); + const { registry } = setup(); 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 }); + } + + const received = await collect(registry, streamId, 3, { afterEventId: 'evt-1' }); + 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 }); } - // 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 }, - ]); + 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(); + 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 } = setup(); + 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('restores every page of a missing stream in order', async () => { + const { registry } = setup(); + const streamId = nextStreamId(); + 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, + }); + + 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 = createRegistry(); + const { registry } = setup(); 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 the event log and continues after the tail', async () => { + 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 }); + await producer.put(second, { streamTTLSeconds: STREAM_TTL_SECONDS }); + await evictStream(streamId); + await expect(producer.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); + + await producer.restoreLog({ + readPages: () => durablePages([first, second]), + streamTTLSeconds: STREAM_TTL_SECONDS, + }); + + const third = makeEvent(3); + 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, 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 restore = () => + producer.restoreLog({ + readPages: () => durablePages([first, second, durableOnly]), + streamTTLSeconds: STREAM_TTL_SECONDS, + }); + await restore(); + await restore(); + + const fourth = makeEvent(4); + await producer.put(fourth, { streamTTLSeconds: STREAM_TTL_SECONDS }); + const received = await collect(registry, streamId, 3); + expect(received).toEqual([first, second, fourth]); + }); + + it('keeps consecutive identical payloads, like two deltas with the same text', async () => { + const { registry } = setup(); + const streamId = nextStreamId(); + 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([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); + }); + + 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(); + 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 new file mode 100644 index 0000000000..5ae41f56c6 --- /dev/null +++ b/packages/trueforge/tests/runtime/event-subscription/entryId.test.ts @@ -0,0 +1,24 @@ +import { eventIdSeekCursor, eventIdTimestampMs } from '../../../src/runtime/event-subscription/entryId'; + +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`, + ); + }); +}); 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..26d5e1c76b 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,71 @@ describeRedis('RedisEventSubscription (EventSubscription contract)', () => { await client.close(); }); - runEventSubscriptionContractSuite(() => new EventSubscriptionRegistry(client)); + runEventSubscriptionContractSuite(() => ({ + registry: new EventSubscriptionRegistry(client), + evictStream: async streamId => { + 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) }]); + 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 () => { + 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/restore.test.ts b/packages/trueforge/tests/runtime/event-subscription/restore.test.ts new file mode 100644 index 0000000000..9eeaca6a40 --- /dev/null +++ b/packages/trueforge/tests/runtime/event-subscription/restore.test.ts @@ -0,0 +1,83 @@ +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('passes every page to the install when the stream is missing', async () => { + const pages: Row[][] = []; + let refreshed = 0; + + const outcome = await applyEventLogRestore({ + readPages: () => pagesOf([created, paused], [], [running, done]), + readTailEntryId: () => Promise.resolve(undefined), + installIfAbsent: async events => { + for await (const page of events) { + pages.push([...page]); + } + return true; + }, + refreshTtl: () => { + refreshed += 1; + return Promise.resolve(); + }, + }); + + 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; + + const outcome = await applyEventLogRestore({ + readPages: () => pagesOf([created, paused, running]), + readTailEntryId: () => Promise.resolve('2-0'), + installIfAbsent: () => { + installed += 1; + return Promise.resolve(true); + }, + refreshTtl: () => { + refreshed += 1; + return Promise.resolve(); + }, + }); + + 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; + + const outcome = await applyEventLogRestore({ + readPages: () => pagesOf([created]), + readTailEntryId: () => Promise.resolve(undefined), + installIfAbsent: () => Promise.resolve(false), + refreshTtl: () => { + refreshed += 1; + return Promise.resolve(); + }, + }); + + 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 new file mode 100644 index 0000000000..fc695b0248 --- /dev/null +++ b/packages/trueforge/tests/runtime/event-subscription/turnEventLogRestore.test.ts @@ -0,0 +1,259 @@ +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, + makeModelMessageEvent, + makeTurnCreatedEvent, + TEST_ACTIVE_EXECUTOR_ID, +} from '../../../../trueforge-core/tests/agent-session/testHelpers'; +import { + EventSubscriptionRegistry, + StreamGoneError, + type SequencedEvent, +} from '../../../src/runtime/event-subscription'; +import { InMemoryEventStreamStore } from '../../../src/runtime/event-subscription/inMemory'; +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, + streamId: string, + count: number, +): Promise[]> { + const abort = new AbortController(); + const generator = registry.get(streamId).poll({ 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; +} + +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 = await createPausedTurnStore(); + const created = makeTurnCreatedEvent(TURN_ID); + const message = makeModelMessageEvent(); + await store.appendToEvents({ + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + session_id: SESSION_ID, + turn_id: TURN_ID, + events: [created, message], + }); + + const streams = new InMemoryEventStreamStore(); + const registry = new EventSubscriptionRegistry(undefined, streams); + const streamId = 'turn-stream'; + const subscription = registry.get(streamId); + const { logger } = recordingLogger(); + const restore = () => + restoreTurnEventLog({ + sessionStore: store, + subscription, + logger, + sessionId: SESSION_ID, + turnId: TURN_ID, + streamTTLSeconds: STREAM_TTL_SECONDS, + }); + + await expect(subscription.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); + + await restore(); + expect((await collect(registry, streamId, 2)).map(event => event.id)).toEqual([created.id, message.id]); + + await restore(); + expect((await collect(registry, streamId, 2)).map(event => event.id)).toEqual([created.id, message.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, + session_id: SESSION_ID, + turn_id: TURN_ID, + events: [later], + }); + 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); + await expect(subscription.assertSubscribable()).rejects.toBeInstanceOf(StreamGoneError); + + await restore(); + const appended = makeModelMessageEvent(); + await subscription.put(appended, { streamTTLSeconds: STREAM_TTL_SECONDS }); + expect((await collect(registry, streamId, 4)).map(event => event.id)).toEqual([ + created.id, + message.id, + later.id, + 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']); + }); +}); diff --git a/packages/trueforge/tests/unit/apis/turns.test.ts b/packages/trueforge/tests/unit/apis/turns.test.ts index 5acc14d3fe..d347fea80e 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,17 @@ 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, + makeModelMessageEvent, + 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'; @@ -883,4 +895,163 @@ 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 eventSubscriptions = new EventSubscriptionRegistry(undefined); + const app = new OpenAPIHono(); + app.route( + '/', + createTurnsRouter({ + activeTurns, + authorizer: new TrueForgeAuthorizer(), + eventSubscriptions, + 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), + }), + ); + + // 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(); + + // 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); + } + + 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(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 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); + }); + }); }); diff --git a/packages/trueforge/tests/unit/controller/memoryIngest.test.ts b/packages/trueforge/tests/unit/controller/memoryIngest.test.ts index 05b381662f..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 === 'done' && tip.state.required_actions.length > 0; - 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,38 +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({ completedAt: at(-12 * 60 * MINUTE + MINUTE), toolCallId: 'call-1' }), - createdAt: at(-12 * 60 * 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 () => { @@ -485,43 +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({ completedAt: at(-12.5 * 60 * MINUTE), toolCallId: 'call-1' }), - createdAt: at(-12.5 * 60 * MINUTE - 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 () => { @@ -599,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 24bb11ab7f..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() { @@ -20,7 +18,7 @@ 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), }), turn({ @@ -98,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', () => { @@ -116,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', () => { @@ -167,7 +163,7 @@ describe('renderExchange', () => { ], }, ], - state: pausedState({ completedAt: at(MINUTE), toolCallId: 'other-tool' }), + state: pausedState(), }), turn({ id: 'a2', @@ -194,7 +190,7 @@ 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(), }), ]), ); @@ -208,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 8b5e6fa564..729328046e 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' }; diff --git a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts index c7026627e7..f414b1307f 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,91 @@ 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('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(); + 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'); + }); }); 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 = { diff --git a/python/trueforge_sdk/reference.md b/python/trueforge_sdk/reference.md index 54e0b8a7c0..4a436572d8 100644 --- a/python/trueforge_sdk/reference.md +++ b/python/trueforge_sdk/reference.md @@ -3550,7 +3550,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. Resume with from_event_id or after_event_id.
@@ -3607,7 +3607,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]` — 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 + +
+
+ +
+
+ +**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 f07077f346..3b29158db2 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 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. 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. + 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 + + 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 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. 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. + 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 + + 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 e9572d6db8..126486597b 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 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. 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. + 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 + + 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 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. 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. + 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 + + 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: