diff --git a/.changeset/turn-offload-steal.md b/.changeset/turn-offload-steal.md new file mode 100644 index 000000000..6b56c477d --- /dev/null +++ b/.changeset/turn-offload-steal.md @@ -0,0 +1,6 @@ +--- +'@truefoundry/trueforge': patch +'@truefoundry/trueforge-core': patch +--- + +Rebuild a paused turn on its owner, offload it after a keepalive, and steal it only when the owner's heartbeat is gone. diff --git a/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index f43c59d2c..ec9f7375e 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -10055,7 +10055,17 @@ } } }, - "description": "Turn is terminal, or an event id already exists." + "description": "Turn is terminal, running and not accepting events, or an event id already exists." + }, + "500": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RequestErrorResponse" + } + } + }, + "description": "The executor ran the request and failed while applying the events." } }, "summary": "Create turn events", diff --git a/docs/openapi.json b/docs/openapi.json index f43c59d2c..ec9f7375e 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -10055,7 +10055,17 @@ } } }, - "description": "Turn is terminal, or an event id already exists." + "description": "Turn is terminal, running and not accepting events, or an event id already exists." + }, + "500": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RequestErrorResponse" + } + } + }, + "description": "The executor ran the request and failed while applying the events." } }, "summary": "Create turn events", diff --git a/packages/assistant-ui-runtime/src/convertTurnMessages.ts b/packages/assistant-ui-runtime/src/convertTurnMessages.ts index c546ea7ec..048af6d75 100644 --- a/packages/assistant-ui-runtime/src/convertTurnMessages.ts +++ b/packages/assistant-ui-runtime/src/convertTurnMessages.ts @@ -653,58 +653,52 @@ async function resolveSessionTip(options: { sessionId: string; itemsAsc: readonly SessionEventItem[]; }): Promise { - const { server, sessionId, itemsAsc } = options; - const open = findOpenTurnCreated(itemsAsc); - if (open != null) { - const continuationInput = open.event.input ?? []; - if (typeof server.getTurn === 'function') { - try { - const turn = await server.getTurn({ - sessionId, - turnId: open.turnId, - }); - if (turn.state.status === TURN_STATUS.RUNNING || turn.state.status === TURN_STATUS.PAUSED) { - return { - runningTurn: - open.update?.state.status === TURN_STATUS.PAUSED && turn.state.status === TURN_STATUS.RUNNING - ? turnFromCreatedEvent({ - sessionId, - turnId: open.turnId, - event: open.event, - update: open.update, - }) - : turn, - ...(open.sandboxId != null ? { sandboxId: open.sandboxId } : {}), - continuationInput: turn.input ?? continuationInput, - }; - } - // Tip finished between listEvents and getTurn: nothing to - // resume, but its answers still belong in the fold, and the turn - // itself must be committed so the painted user message stays. - return { finishedTurn: turn, continuationInput: turn.input ?? continuationInput }; - } catch { - // getTurn failed; synthesize from the open turn.created below. - } - } - return { - runningTurn: turnFromCreatedEvent({ + const open = findOpenTurnCreated(options.itemsAsc); + if (open == null) { + return { continuationInput: [] }; + } + + const { server, sessionId } = options; + const continuationInput = open.event.input ?? []; + if (typeof server.getTurn === 'function') { + try { + const turn = await server.getTurn({ sessionId, turnId: open.turnId, - event: open.event, - ...(open.update != null ? { update: open.update } : {}), - }), - ...(open.sandboxId != null ? { sandboxId: open.sandboxId } : {}), - continuationInput, - }; + }); + if (turn.state.status === TURN_STATUS.RUNNING || turn.state.status === TURN_STATUS.PAUSED) { + return { + runningTurn: + open.update?.state.status === TURN_STATUS.PAUSED && turn.state.status === TURN_STATUS.RUNNING + ? turnFromCreatedEvent({ + sessionId, + turnId: open.turnId, + event: open.event, + update: open.update, + }) + : turn, + ...(open.sandboxId != null ? { sandboxId: open.sandboxId } : {}), + continuationInput: turn.input ?? continuationInput, + }; + } + // Tip finished between listEvents and getTurn: nothing to + // resume, but its answers still belong in the fold, and the turn + // itself must be committed so the painted user message stays. + return { finishedTurn: turn, continuationInput: turn.input ?? continuationInput }; + } catch { + // getTurn failed; synthesize from the open turn.created below. + } } - - // Hosts that exclude the running turn from listEvents still surface it as - // the first row of listTurns when that API is tip-first. - const turnsPage = await server.listTurns({ sessionId, limit: 1 }); - const tip = turnsPage.data[0]; - return tip?.state.status === TURN_STATUS.RUNNING || tip?.state.status === TURN_STATUS.PAUSED - ? { runningTurn: tip, continuationInput: tip.input ?? [] } - : { continuationInput: [] }; + return { + runningTurn: turnFromCreatedEvent({ + sessionId, + turnId: open.turnId, + event: open.event, + ...(open.update != null ? { update: open.update } : {}), + }), + ...(open.sandboxId != null ? { sandboxId: open.sandboxId } : {}), + continuationInput, + }; } /** diff --git a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts index add165673..fed4fe3dc 100644 --- a/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts +++ b/packages/assistant-ui-runtime/src/useTrueForgeAgentMessages.ts @@ -999,8 +999,7 @@ export function useTrueForgeAgentMessages({ if (!runStreamStarted) { // A supersede already aborted the live stream. A send that fails // before it installs a message must not cancel the painted open tip. - const hasIncompleteStream = - snapshotRef.current.activeStream != null && snapshotRef.current.activeStream.streamComplete !== true; + const hasIncompleteStream = isIncompleteActiveStream(snapshotRef.current); if ('userMessage' in options && (pendingUserWasSet || hasIncompleteStream)) { const abandoned = abandonInFlightClientTurn(commitActiveStream(snapshotRef.current)); snapshotRef.current = abandoned; diff --git a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts index 2bb3182e2..f4dddc18a 100644 --- a/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts +++ b/packages/assistant-ui-runtime/test/convertTurnMessages.test.ts @@ -1699,6 +1699,7 @@ describe('convertTurnMessages', () => { status: { type: 'incomplete', reason: 'cancelled' }, metadata: { custom: { cancellationReason: 'abandoned' } }, sequenceNumber: 1, + turnState: { status: 'cancelled', reason: 'abandoned', completedAt: createdAt }, }, ]); }); diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index f2567c7f1..c2505f86c 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -1,7 +1,7 @@ /** * Bound session handle: starts turns via {@link SessionHandle.createTurn}. */ -import { newEventId, type UserMessage } from '../core/events/schema'; +import { newEventId, type TurnUserEvent, type UserMessage } from '../core/events/schema'; import type { AgentDefinition } from '../core/runtime/AgentDefinition'; import { AgentThread } from '../core/runtime/AgentThread'; import type { AgentThreadAppendContext, AgentThreadSnapshot } from '../core/runtime/AgentThread.types'; @@ -14,18 +14,22 @@ import type { ITurnResourceResolver, ResolvedAgentDefinition } from './ITurnReso import type { SessionRecord } from './models/SessionRecord'; import { MAIN_THREAD_ID, type TurnRecord, type TurnSnapshot } from './models/TurnRecord'; import type { AgentSpec } from './schemas/agentSpec'; -import type { SessionEventItem } from './schemas/events'; +import type { PersistedTurnEvent, SessionEventItem } from './schemas/events'; import { EventType, type TurnDoneEvent } from './schemas/events'; import type { TokenPagination } from './schemas/pagination'; import type { SessionAgent } from './schemas/session'; import { CancellationReason } from './schemas/turn'; import type { ISessionStore, NewThreadInit, TurnContextAppend, TurnRecordWithoutSnapshot } from './store/ISessionStore'; +import { TurnNotFoundError, TurnNotRunningError, TurnOwnershipLostError } from './store/SessionStoreErrors'; import { TurnHandle } from './TurnHandle'; /** How many recent ancestors SessionHandle persists on each turn. Store-local. */ const MAX_TURN_ANCESTORS = 20; -function resolvePreviousTurnId(requested: string | undefined, lastTurnId: string | null | undefined): string | null { +export function resolvePreviousTurnId( + requested: string | undefined, + lastTurnId: string | null | undefined, +): string | null { if (requested === 'none') { return null; } @@ -187,56 +191,11 @@ export class SessionHandle< // and rethrow — same pattern as the gateway's sessionTurnHandler, which // wraps getAgent/handler.init and calls gatewayStore.dispose() on failure. try { - const tracing = input.resolver.createTracing(); - const agent = this.session.agent; - const spec = agent.type === 'inline' ? agent.spec : await input.resolver.resolveAgentSpec({ agent_id: agent.id }); - const sandbox = await input.resolver.resolveSandbox({ - spec, - existing: previous?.snapshot.sandbox_info ?? undefined, - previousTurn: previous, - signal: input.abortController.signal, - tracing, - }); - - const definitionsByThreadId = await this.resolveAgentDefinitions({ - previous, - resolver: input.resolver, - spec, - tracing, - signal: input.abortController.signal, - }); - - // Wire Code Mode once from main toolSets before threads share the sandbox. - const mainDefinition = definitionsByThreadId.get(MAIN_THREAD_ID); - if (!mainDefinition) { - throw new Error('Unreachable: missing resolved agent definition for main thread'); - } - if (sandbox) { - sandbox.configureCodeMode(mainDefinition.definition.toolSets ?? []); - } - - const agentThreads = this.buildThreads({ + const { orchestrator, agentThreads } = await this.assembleTurn({ previous, resolver: input.resolver, - spec, - sandbox, - tracing, - definitionsByThreadId, - }); - - const createDynamicSubAgentThread = this.makeCreateDynamicSubAgentThread({ - resolver: input.resolver, - previous, - spec, - sandbox, - tracing, - }); - - const orchestrator = new AgentThreadOrchestrator({ - agentThreads, - createDynamicSubAgentThread, - tracing, - logger: input.resolver.logger, + abortController: input.abortController, + pendingTurnEvents: [], }); const new_context_appends = await collectContextAppends(orchestrator.send(input.input ?? [])); @@ -317,6 +276,70 @@ export class SessionHandle< } } + /** + * Rebuild a paused turn from its own snapshot. Does not create a row, freeze a + * predecessor, or re-send the original input. The caller must already own + * `active_executor_id`. Unconsumed inbox rows are placed on the orchestrator + * before the handle is returned, so a later send() lands after them. + */ + async rebuildTurn(input: { + turn_id: string; + active_executor_id: string; + resolver: ITurnResourceResolver; + abortController: AbortController; + }): Promise> { + try { + const turn = await this.store.getTurn({ + session_id: this.session.session_id, + turn_id: input.turn_id, + }); + if (!turn) { + throw new TurnNotFoundError(input.turn_id); + } + if (turn.active_executor_id !== input.active_executor_id) { + throw new TurnOwnershipLostError({ + turn_id: input.turn_id, + active_executor_id: input.active_executor_id, + }); + } + if (turn.state.status === 'running') { + throw new Error(`Turn ${input.turn_id} must be paused to rebuild`); + } + if (turn.state.status !== 'paused') { + throw new TurnNotRunningError(input.turn_id, turn.state); + } + + const pendingTurnEvents = await this.store.listUnconsumedTurnInboundEvents({ + session_id: this.session.session_id, + turn_id: input.turn_id, + }); + input.resolver.logger.info('Rebuilding paused turn', { + pendingEventCount: pendingTurnEvents.length, + }); + const { orchestrator } = await this.assembleTurn({ + previous: turn, + resolver: input.resolver, + abortController: input.abortController, + pendingTurnEvents, + }); + + return new TurnHandle({ + store: this.store, + turn, + orchestrator, + resolver: input.resolver, + abortController: input.abortController, + }); + } catch (error) { + await input.resolver.close().catch((closeError: unknown) => { + input.resolver.logger.warn('TurnResourceResolver.close() failed after rebuildTurn() error', { + err: closeError, + }); + }); + throw error; + } + } + /** * Conditionally cancel a running or paused turn in this session and return the * now-immutable record. Already-terminal turns are returned unchanged. @@ -331,15 +354,33 @@ export class SessionHandle< }); } - /** Returns the turn handle (store-backed; not executable), or undefined if not found. */ - async getTurn(turn_id: string): Promise | undefined> { - return TurnHandle.get({ - store: this.store, + /** The stored turn, including its snapshot. Undefined when this session has no such turn. */ + async getTurn(turn_id: string): Promise | undefined> { + return this.store.getTurn({ session_id: this.session.session_id, turn_id, }); } + /** Paginated events for one turn in this session. */ + async listTurnEvents(input: { + turn_id: string; + limit: number; + page_token?: string | undefined; + order?: 'asc' | 'desc' | undefined; + }): Promise<{ + data: PersistedTurnEvent[]; + pagination: TokenPagination; + }> { + return this.store.listTurnEvents({ + session_id: this.session.session_id, + turn_id: input.turn_id, + limit: input.limit, + page_token: input.page_token, + order: input.order, + }); + } + /** Paginated list of this session's turn rows (no snapshot). */ async listTurns(input: { limit: number; @@ -373,6 +414,75 @@ export class SessionHandle< }); } + /** + * Threads and orchestrator for one turn. `previous` is the frozen predecessor, + * absent for a new root, or the paused turn itself when rebuilding. + */ + private async assembleTurn(input: { + previous: TurnRecord | undefined; + resolver: ITurnResourceResolver; + abortController: AbortController; + pendingTurnEvents: TurnUserEvent[]; + }): Promise<{ orchestrator: AgentThreadOrchestrator; agentThreads: Map }> { + const tracing = input.resolver.createTracing(); + const agent = this.session.agent; + const spec = agent.type === 'inline' ? agent.spec : await input.resolver.resolveAgentSpec({ agent_id: agent.id }); + const sandbox = await input.resolver.resolveSandbox({ + spec, + existing: input.previous?.snapshot.sandbox_info ?? undefined, + previousTurn: input.previous, + signal: input.abortController.signal, + tracing, + }); + + const definitionsByThreadId = await this.resolveAgentDefinitions({ + previous: input.previous, + resolver: input.resolver, + spec, + tracing, + signal: input.abortController.signal, + }); + + // Wire Code Mode once from main toolSets before threads share the sandbox. + const mainDefinition = definitionsByThreadId.get(MAIN_THREAD_ID); + if (!mainDefinition) { + throw new Error('Unreachable: missing resolved agent definition for main thread'); + } + if (sandbox) { + sandbox.configureCodeMode(mainDefinition.definition.toolSets ?? []); + } + + const agentThreads = this.buildThreads({ + previous: input.previous, + resolver: input.resolver, + spec, + sandbox, + tracing, + definitionsByThreadId, + }); + + const createDynamicSubAgentThread = this.makeCreateDynamicSubAgentThread({ + resolver: input.resolver, + previous: input.previous, + spec, + sandbox, + tracing, + }); + + const previousMcpServers = input.previous?.snapshot.mcp_servers; + const orchestrator = new AgentThreadOrchestrator({ + agentThreads, + createDynamicSubAgentThread, + tracing, + logger: input.resolver.logger, + pendingTurnEvents: input.pendingTurnEvents, + mcpServers: + previousMcpServers !== null && previousMcpServers !== undefined ? Object.values(previousMcpServers) : [], + }); + + return { orchestrator, agentThreads }; + } + private async resolveAgentDefinitions(input: { previous: TurnRecord | undefined; resolver: ITurnResourceResolver; diff --git a/packages/trueforge-core/src/agent-session/TurnHandle.ts b/packages/trueforge-core/src/agent-session/TurnHandle.ts index bf2a758b8..214fa9c69 100644 --- a/packages/trueforge-core/src/agent-session/TurnHandle.ts +++ b/packages/trueforge-core/src/agent-session/TurnHandle.ts @@ -47,6 +47,16 @@ import { TurnNotRunningError, TurnOwnershipLostError } from './store/SessionStor /** Streaming yield union — deltas pass through; never persisted. No sequence_number. */ export type TurnStreamingEvent = PersistedTurnEvent | ModelMessageDeltaEvent; +/** + * Abort reason for keepalive offload and for dropping a stale local handle. + * The generator stops and the resolver closes. The turn row is left as it is. + */ +export const TurnOffloadReason = 'turn-offload'; + +export function isTurnOffloadReason(reason: unknown): boolean { + return reason === TurnOffloadReason; +} + function abortOnReturn(generator: AsyncGenerator, abortController: AbortController): AsyncGenerator { const returnGenerator = generator.return.bind(generator); generator.return = value => { @@ -168,49 +178,23 @@ function resolveTerminalState(input: { export class TurnHandle> { private readonly store: ISessionStore; private turn: TurnRecord; - private readonly orchestrator: AgentThreadOrchestrator | undefined; - private readonly resolver: ITurnResourceResolver | undefined; + private readonly orchestrator: AgentThreadOrchestrator; + private readonly resolver: ITurnResourceResolver; private readonly abortController: AbortController; private streamStarted = false; constructor(options: { store: ISessionStore; turn: TurnRecord; - orchestrator?: AgentThreadOrchestrator | undefined; - resolver?: ITurnResourceResolver | undefined; - abortController?: AbortController | undefined; + orchestrator: AgentThreadOrchestrator; + resolver: ITurnResourceResolver; + abortController: AbortController; }) { this.store = options.store; this.turn = options.turn; this.orchestrator = options.orchestrator; this.resolver = options.resolver; - this.abortController = options.abortController ?? new AbortController(); - } - - /** Store-only handle (e.g. from {@link SessionHandle.getTurn}) — stream() is not available. */ - static fromRecord>(options: { - store: ISessionStore; - turn: TurnRecord; - }): TurnHandle { - return new TurnHandle(options); - } - - static async get>(options: { - store: ISessionStore; - session_id: string; - turn_id: string; - }): Promise | undefined> { - const turn = await options.store.getTurn({ - session_id: options.session_id, - turn_id: options.turn_id, - }); - if (!turn) { - return undefined; - } - return TurnHandle.fromRecord({ - store: options.store, - turn, - }); + this.abortController = options.abortController; } get id(): string { @@ -246,17 +230,8 @@ export class TurnHandle> { } async *send(events: TurnUserEvent[]): AsyncGenerator { - const orchestrator = this.requireLiveOrchestrator('send'); - const accepted = orchestrator.send(events); - yield* accepted; - orchestrator.wake(); - } - - private requireLiveOrchestrator(method: string): AgentThreadOrchestrator { - if (!this.orchestrator) { - throw new Error(`TurnHandle.${method}() is only available on a live turn from SessionHandle.createTurn()`); - } - return this.orchestrator; + yield* this.orchestrator.send(events); + this.orchestrator.wake(); } /** @@ -293,12 +268,7 @@ export class TurnHandle> { throw new Error('TurnHandle.stream() is single-use and was already called'); } this.streamStarted = true; - const orchestrator = this.orchestrator; - const resolver = this.resolver; const signal = this.abortController.signal; - if (!orchestrator || !resolver) { - throw new Error('TurnHandle.stream() is only available on turns returned from SessionHandle.createTurn()'); - } // A consumer can leave early (break, or iterator.return()) while execute is parked // on wakeSignal.wait() at a pause. return() does not cancel that await, so the @@ -307,7 +277,7 @@ export class TurnHandle> { // the latch. The loop then sees the signal and returns. A turn that finishes on // its own does not call return(), so it is not aborted. An abort that already // happened (shutdown, timeout, barge-in) keeps its reason. - return abortOnReturn(this.streamExecution(orchestrator, resolver, signal), this.abortController); + return abortOnReturn(this.streamExecution(this.orchestrator, this.resolver, signal), this.abortController); } private async *streamExecution( @@ -315,57 +285,91 @@ export class TurnHandle> { resolver: ITurnResourceResolver, signal: AbortSignal, ): AsyncGenerator { + if (this.turn.state.status !== 'running' && this.turn.state.status !== 'paused') { + throw new Error(`Cannot execute a ${this.turn.state.status} turn`); + } + const turnStatus = this.turn.state.status; let caughtError: Error | undefined; let executeResult: AgentThreadExecutionResult | undefined; - let ownershipLost: TurnOwnershipLostError | undefined; try { - yield await this.persistTurnCreated(); - executeResult = yield* this.executeAndPersist(orchestrator, signal); + if (turnStatus === 'running') { + yield await this.persistTurnCreated(); + } + executeResult = yield* this.executeAndPersist(orchestrator, signal, turnStatus); } catch (error) { caughtError = error instanceof Error ? error : new Error(String(error)); - if (caughtError instanceof TurnOwnershipLostError) { - // executeAndPersist already waited for abort. Skip the terminal write. - // The row belongs to the other replica. - ownershipLost = caughtError; - } } finally { - // Terminal write stays in finally so a consumer break still closes the turn. - let turnDone: TurnDoneEvent | undefined; - if (!ownershipLost) { - try { - // TurnNotRunningError: another writer already terminated the turn — defer to its state. - if (caughtError instanceof TurnNotRunningError) { - turnDone = this.turnDoneFromStoreConflict(caughtError, orchestrator); - } else { - turnDone = await this.persistTurnTerminal({ - signal, - caughtError, - executeResult, - orchestrator, - }); - } - } catch (error) { - if (error instanceof TurnOwnershipLostError) { - ownershipLost = error; - } else if (error instanceof TurnNotRunningError) { - turnDone = this.turnDoneFromStoreConflict(error, orchestrator); - } else { - await this.closeResolver(resolver); - // eslint-disable-next-line no-unsafe-finally -- terminal-state write failed; reject the stream - throw error; - } - } + // A consumer break still enters finally, so the turn closes on that path too. + yield* this.closeTurn({ signal, caughtError, executeResult, orchestrator, resolver }); + } + } + + /** + * Stop the in-memory run. Offload and a lost owner skip the terminal write. + * Every other exit writes `turn.done`, or emits the state another writer already stored. + */ + private async *closeTurn(input: { + signal: AbortSignal; + caughtError: Error | undefined; + executeResult: AgentThreadExecutionResult | undefined; + orchestrator: AgentThreadOrchestrator; + resolver: ITurnResourceResolver; + }): AsyncGenerator { + let ownershipLost = input.caughtError instanceof TurnOwnershipLostError ? input.caughtError : undefined; + let turnDone: TurnDoneEvent | undefined; + let closeError: Error | undefined; + try { + // Offload leaves the row paused. A lost owner must not write the other replica's row. + if (ownershipLost === undefined && !isTurnOffloadReason(input.signal.reason)) { + turnDone = await this.turnDoneForClose(input); } - await this.closeResolver(resolver); - if (turnDone) { - yield turnDone; + } catch (error) { + // updateTurnTerminalState lost active_executor_id. Another replica claimed this + // paused turn after the last event write, so this run must not emit turn.done. + if (error instanceof TurnOwnershipLostError) { + ownershipLost = error; + } else { + closeError = error instanceof Error ? error : new Error(String(error), { cause: error }); } + } finally { + await this.closeResolver(input.resolver); + } + if (closeError !== undefined) { + throw closeError; + } + if (isTurnOffloadReason(input.signal.reason)) { + input.resolver.logger.info('Offloaded paused turn without a terminal write'); + } else if (ownershipLost !== undefined) { + input.resolver.logger.info('Lost turn ownership; skipped the terminal write'); + } + if (turnDone !== undefined) { + yield turnDone; } - if (ownershipLost) { + if (ownershipLost !== undefined) { throw ownershipLost; } } + /** Write this run's terminal state. If the row is already terminal, emit the stored state. */ + private async turnDoneForClose(input: { + signal: AbortSignal; + caughtError: Error | undefined; + executeResult: AgentThreadExecutionResult | undefined; + orchestrator: AgentThreadOrchestrator; + }): Promise { + if (input.caughtError instanceof TurnNotRunningError) { + return this.turnDoneFromStoredState(input.caughtError, input.orchestrator); + } + try { + return await this.persistTurnTerminal(input); + } catch (error) { + if (error instanceof TurnNotRunningError) { + return this.turnDoneFromStoredState(error, input.orchestrator); + } + throw error; + } + } + private async persistTurnCreated(): Promise { const turnCreated: TurnCreatedEvent = { type: EventType.TURN_CREATED, @@ -389,8 +393,9 @@ export class TurnHandle> { private async *executeAndPersist( orchestrator: AgentThreadOrchestrator, signal: AbortSignal, + turnStatus: NonTerminalTurnState['status'], ): AsyncGenerator { - const generator = orchestrator.execute({ signal }); + const generator = orchestrator.execute({ signal, turnStatus }); try { let iterResult = await generator.next(); while (!iterResult.done) { @@ -447,11 +452,12 @@ export class TurnHandle> { return turnDone; } - private turnDoneFromStoreConflict( - error: TurnNotRunningError, - orchestrator: AgentThreadOrchestrator, - updatedAt: Date = new Date(), - ): TurnDoneEvent { + /** + * Another writer already stored the terminal row. Emit that state with this run's metrics. + * There is no second write. + */ + private turnDoneFromStoredState(error: TurnNotRunningError, orchestrator: AgentThreadOrchestrator): TurnDoneEvent { + const updatedAt = new Date(); const state: TerminalTurnState = { ...error.state, metrics: turnMetricsFromAgentThreadMetrics(orchestrator.getMetrics()), diff --git a/packages/trueforge-core/src/agent-session/index.ts b/packages/trueforge-core/src/agent-session/index.ts index 7b2b78b30..d73b5a9cb 100644 --- a/packages/trueforge-core/src/agent-session/index.ts +++ b/packages/trueforge-core/src/agent-session/index.ts @@ -133,7 +133,7 @@ export type { ITurnResourceResolver } from './ITurnResourceResolver'; export { TurnResourceResolver } from './TurnResourceResolver'; export type { TurnSandboxFactory } from './TurnResourceResolver'; -export { SessionHandle } from './SessionHandle'; +export { SessionHandle, resolvePreviousTurnId } from './SessionHandle'; export { Sessions } from './Sessions'; -export { TurnHandle } from './TurnHandle'; +export { TurnHandle, TurnOffloadReason, isTurnOffloadReason } from './TurnHandle'; export type { TurnStreamingEvent } from './TurnHandle'; diff --git a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts index 72c90e7d6..8bb6defc7 100644 --- a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts @@ -185,9 +185,7 @@ export interface AppendToEventsInput extends TurnProgressKeys { events: PersistedTurnEvent[]; } -export interface InsertTurnInboundEventsInput { - session_id: string; - turn_id: string; +export interface InsertTurnInboundEventsInput extends TurnProgressKeys { /** Empty array is a no-op. */ events: TurnUserEvent[]; } @@ -398,7 +396,8 @@ export interface ISessionStore< /** * Durable inbound send-event inbox for a tip the caller still owns. Tip must be * non-terminal (`running` or `paused`) — terminal tip → - * {@link TurnNotRunningError}. Missing session → {@link SessionNotFoundError}; + * {@link TurnNotRunningError}. A different `active_executor_id` → + * {@link TurnOwnershipLostError}. Missing session → {@link SessionNotFoundError}; * unknown turn → {@link TurnNotFoundError}. Duplicate `event_id` on that tip → * {@link TurnEventAlreadyExistsError}. */ diff --git a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts index 8b98ee00f..31a7968f0 100644 --- a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts @@ -568,6 +568,7 @@ export class InMemorySessionStore< return; } this.requireSession(input.session_id); + this.requireOwnedTurn(input); this.requireNonTerminalTurn(input); const tKey = turnKey(input); let list = this.inboundEvents.get(tKey); diff --git a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts index cdb122ab1..bf44470a8 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts @@ -182,6 +182,13 @@ export interface AgentThreadOrchestratorInput { createDynamicSubAgentThread: CreateDynamicSubAgentThread; tracing: AgentTracing; logger: Logger; + /** Inbox rows not yet in the snapshot. `execute` applies them. A rebuild passes these. */ + pendingTurnEvents: TurnUserEvent[]; + /** + * MCP servers already initialized on an earlier turn. A later allow-session policy + * is patched onto these records. This turn often does not emit `mcp.initialize` again. + */ + mcpServers: MCPServerInitInfo[]; } export class AgentThreadOrchestrator { @@ -192,7 +199,7 @@ export class AgentThreadOrchestrator { // Finished sub-agents removed from `agentThreads`; kept so totals still include them. private finishedSubAgentMetrics: AgentThreadMetrics = createEmptyAgentThreadMetrics(); private readonly wakeSignal: Signalable = signalable(); - private readonly pendingTurnEvents: TurnUserEvent[] = []; + private readonly pendingTurnEvents: TurnUserEvent[]; // Last-seen MCP server init records (by id)/ private readonly mcpServerInitInfoById = new Map(); @@ -201,6 +208,10 @@ export class AgentThreadOrchestrator { this.createDynamicSubAgentThread = params.createDynamicSubAgentThread; this.tracing = params.tracing; this.logger = params.logger.child({ module: 'AgentThreadOrchestrator' }); + this.pendingTurnEvents = [...params.pendingTurnEvents]; + for (const server of params.mcpServers) { + this.mcpServerInitInfoById.set(server.id, server); + } } public wake(): void { @@ -566,8 +577,14 @@ export class AgentThreadOrchestrator { public async *execute({ signal, + turnStatus = 'running', }: { signal: AbortSignal; + /** + * Status of the turn being executed. A rebuild passes `paused` so the first + * runnable batch persists `running`. A fresh execute starts `running`. + */ + turnStatus?: 'running' | 'paused' | undefined; }): AsyncGenerator { const { agentThreads } = this; @@ -575,7 +592,10 @@ export class AgentThreadOrchestrator { let output: ModelMessageEvent | null = null; let rootAgentError: AgentThreadExecutionResult['root_agent_error']; let rootFinished = false; - let turnPaused = false; + let turnPaused = turnStatus === 'paused'; + if (turnPaused) { + this.logger.info('Executing paused turn', { pendingEventCount: this.pendingTurnEvents.length }); + } const mainThread = this.getMainThread(); const rootSpan = createRootAgentSpan(mainThread, this.tracing); @@ -602,7 +622,8 @@ export class AgentThreadOrchestrator { return done(); } - // Apply accepted user events in the same order they entered the turn-level queue. + // splice returns the current items and leaves the queue empty before the body runs. + // A send during a yield pushes onto that empty queue and waits for the next loop. for (const event of this.pendingTurnEvents.splice(0)) { switch (event.type) { case EventType.USER_TOOL_APPROVAL_POLICY: diff --git a/packages/trueforge-core/src/request-reply/client.ts b/packages/trueforge-core/src/request-reply/client.ts index b580928bf..2de08eacc 100644 --- a/packages/trueforge-core/src/request-reply/client.ts +++ b/packages/trueforge-core/src/request-reply/client.ts @@ -1,6 +1,11 @@ import { randomUUID } from 'node:crypto'; import { performance } from 'node:perf_hooks'; -import { NoResponderError, RequestTimeoutError } from './errors'; +import { + NO_RESPONDER_HEARTBEAT_LOST_DETAIL, + NO_RESPONDER_PUBLISH_DETAIL, + NoResponderError, + RequestTimeoutError, +} from './errors'; import type { RedisClient } from './redisClient'; import type { JSONReply, JSONValue, PublishedRequest, RequestEnvelope } from './types'; import { jsonReplySchema } from './types'; @@ -71,7 +76,7 @@ export async function redisRequest({ const serializedPayload = JSON.stringify(payload); const subscriberCount = await redisClient.publish(requestChannel(executorId), serializedPayload); if (subscriberCount === 0) { - throw new NoResponderError(executorId, 'No executor received the published request'); + throw new NoResponderError(executorId, NO_RESPONDER_PUBLISH_DETAIL); } const replyDeadline = performance.now() + replyTimeoutMs; @@ -85,7 +90,7 @@ export async function redisRequest({ // We already checked that the executor was alive before sending, but it may have // died while we wait. Recheck here so we fail soon instead of waiting until the timeout. if ((await redisClient.exists(aliveK)) === 0) { - throw new NoResponderError(executorId); + throw new NoResponderError(executorId, NO_RESPONDER_HEARTBEAT_LOST_DETAIL); } if (performance.now() >= replyDeadline) { throw new RequestTimeoutError(replyTimeoutMs); diff --git a/packages/trueforge-core/src/request-reply/errors.ts b/packages/trueforge-core/src/request-reply/errors.ts index 05229d549..181cc3521 100644 --- a/packages/trueforge-core/src/request-reply/errors.ts +++ b/packages/trueforge-core/src/request-reply/errors.ts @@ -1,9 +1,16 @@ +/** Heartbeat key was absent before the request was published. */ +export const NO_RESPONDER_HEARTBEAT_DETAIL = 'heartbeat-not-present'; +/** PUBLISH returned 0 while the heartbeat key still existed. */ +export const NO_RESPONDER_PUBLISH_DETAIL = 'no-executor-subscribed'; +/** Heartbeat key disappeared after the request was published, while waiting for the reply. */ +export const NO_RESPONDER_HEARTBEAT_LOST_DETAIL = 'heartbeat-lost-during-wait'; + /** The target executor's heartbeat is missing or nobody is subscribed (map to HTTP 412). */ export class NoResponderError extends Error { readonly code = 'NO_RESPONDER' as const; constructor( readonly executorId: string, - readonly detail = 'heartbeat not present', + readonly detail = NO_RESPONDER_HEARTBEAT_DETAIL, ) { super(`No responder for executor: ${executorId} (${detail}).`); this.name = 'NoResponderError'; diff --git a/packages/trueforge-core/src/request-reply/index.ts b/packages/trueforge-core/src/request-reply/index.ts index bf5cfd5cf..27c7f70a8 100644 --- a/packages/trueforge-core/src/request-reply/index.ts +++ b/packages/trueforge-core/src/request-reply/index.ts @@ -7,11 +7,18 @@ export { redisRequest } from './client'; export type { SendRequestOptions } from './client'; -export { NoResponderError, ReplyError, RequestTimeoutError } from './errors'; +export { + NO_RESPONDER_HEARTBEAT_DETAIL, + NO_RESPONDER_HEARTBEAT_LOST_DETAIL, + NO_RESPONDER_PUBLISH_DETAIL, + NoResponderError, + ReplyError, + RequestTimeoutError, +} from './errors'; export { RequestReplyExecutor } from './executor'; export type { RequestReplyErrorHandler, RunExecutorOptions } from './executor'; export type { RedisClient } from './redisClient'; export { RequestReplyRouter } from './router'; export type { RouteHandler } from './router'; -export { jsonReplySchema } from './types'; +export { jsonReplySchema, jsonValueFrom } from './types'; export type { JSONReply, JSONValue, RequestEnvelope, RequestHandler } from './types'; diff --git a/packages/trueforge-core/src/request-reply/types.ts b/packages/trueforge-core/src/request-reply/types.ts index 2b0f35904..98fdba054 100644 --- a/packages/trueforge-core/src/request-reply/types.ts +++ b/packages/trueforge-core/src/request-reply/types.ts @@ -15,6 +15,16 @@ const jsonSchema: z.ZodType = z.lazy(() => { ]); }); +/** JSON.stringify drops `undefined`, which JSONValue cannot carry. */ +export function jsonValueFrom(value: unknown): JSONValue { + const serialized = JSON.stringify(value); + if (typeof serialized !== 'string') { + throw new Error('Value is not JSON'); + } + const parsed: unknown = JSON.parse(serialized); + return jsonSchema.parse(parsed); +} + /** Request payload as seen by the caller of redisRequest and by route handlers. */ export interface RequestEnvelope { body: T; diff --git a/packages/trueforge-core/tests/agent-session/sessions.test.ts b/packages/trueforge-core/tests/agent-session/sessions.test.ts index 7e55ebaf3..17dbf59cb 100644 --- a/packages/trueforge-core/tests/agent-session/sessions.test.ts +++ b/packages/trueforge-core/tests/agent-session/sessions.test.ts @@ -3,7 +3,6 @@ import { CancellationReason } from '../../src/agent-session/schemas/turn'; import { Sessions } from '../../src/agent-session/Sessions'; import { InMemorySessionStore } from '../../src/agent-session/store/InMemorySessionStore'; import { TurnNotFoundError } from '../../src/agent-session/store/SessionStoreErrors'; -import { TurnHandle } from '../../src/agent-session/TurnHandle'; import { makeAgentSpec, makeTestResolver, mintTestTurnId, TEST_ACTIVE_EXECUTOR_ID } from './testHelpers'; describe('Sessions / SessionHandle / TurnHandle (storage + createTurn)', () => { @@ -178,22 +177,10 @@ describe('Sessions / SessionHandle / TurnHandle (storage + createTurn)', () => { }); const getSession = jest.spyOn(store, 'getSession'); - const direct = await TurnHandle.get({ - store, - session_id: 's1', - turn_id: created.id, - }); const throughSession = await session.getTurn(created.id); - const missingDirect = await TurnHandle.get({ - store, - session_id: 's1', - turn_id: 'missing-turn', - }); const missingThroughSession = await session.getTurn('missing-turn'); - expect(direct?.id).toBe(created.id); - expect(throughSession?.id).toBe(created.id); - expect(missingDirect).toBeUndefined(); + expect(throughSession?.turn_id).toBe(created.id); expect(missingThroughSession).toBeUndefined(); expect(getSession).not.toHaveBeenCalled(); }); diff --git a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts index 56f2ef9e4..0a9625bd6 100644 --- a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts +++ b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts @@ -781,6 +781,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [ { id: newEventId(), @@ -2966,6 +2967,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { await store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [later, earlier], }); @@ -2973,6 +2975,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [later], }), ).rejects.toMatchObject({ @@ -2993,6 +2996,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [fresh, later], }), ).rejects.toMatchObject({ @@ -3005,6 +3009,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [ { id: dupId, @@ -3035,6 +3040,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [ { id: 'evt-after-done', @@ -3049,6 +3055,29 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { ).rejects.toBeInstanceOf(TurnNotRunningError); }); + it('turn_inbound_events: a stale executor cannot insert', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + await expect( + store.insertTurnInboundEvents({ + session_id: sessionId, + turn_id: 'turn-1', + active_executor_id: 'someone-else', + events: [ + { + id: 'evt-stale', + type: 'user.tool_approval' as const, + thread_id: 'main', + tool_call_id: 'tc-stale', + approval: { status: 'allow' as const }, + created_at: new Date().toISOString(), + }, + ], + }), + ).rejects.toBeInstanceOf(TurnOwnershipLostError); + }); + it('turn_inbound_events: list unconsumed and mark consumed', async () => { const store = createStore(); await seedSession(store); @@ -3074,6 +3103,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { await store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [evtB, evtA], }); @@ -3203,6 +3233,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [inbound], }), ]); @@ -3223,6 +3254,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { await store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [ { id: newEventId(), @@ -3239,6 +3271,7 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { store.insertTurnInboundEvents({ session_id: sessionId, turn_id: 'turn-1', + active_executor_id: TEST_ACTIVE_EXECUTOR_ID, events: [ { id: newEventId(), diff --git a/packages/trueforge-core/tests/agent-session/turnResume.test.ts b/packages/trueforge-core/tests/agent-session/turnResume.test.ts index 376a7e0aa..85aeeba71 100644 --- a/packages/trueforge-core/tests/agent-session/turnResume.test.ts +++ b/packages/trueforge-core/tests/agent-session/turnResume.test.ts @@ -30,6 +30,7 @@ async function sendWithInbox(turn: TurnHandle, store: InMemorySessionStore, even await store.insertTurnInboundEvents({ session_id: turn.session_id, turn_id: turn.id, + active_executor_id: turn.record.active_executor_id, events: batch, }); } @@ -288,7 +289,9 @@ describe('TurnHandle.send() full-approval resume (agent-session e2e)', () => { const pending = iterator.next(); pending.catch(() => undefined); - await sendWithInbox(turn, store, [ + // Stale executor cannot insert inbox rows. Wake the parked generator so the + // next fenced persist hits TurnOwnershipLostError and stays suspended. + for await (const _batch of turn.send([ { type: EventType.USER_TOOL_APPROVAL, id: newEventId(), @@ -297,7 +300,9 @@ describe('TurnHandle.send() full-approval resume (agent-session e2e)', () => { tool_call_id: WRITE_NOTE_CALL_ID, approval: { status: 'allow' }, }, - ]); + ])) { + // intentionally not inserting + } try { const parked = await Promise.race([ diff --git a/packages/trueforge-core/tests/orchestration/orchestration.test.ts b/packages/trueforge-core/tests/orchestration/orchestration.test.ts index 8b2812ad9..41643e66c 100644 --- a/packages/trueforge-core/tests/orchestration/orchestration.test.ts +++ b/packages/trueforge-core/tests/orchestration/orchestration.test.ts @@ -74,6 +74,8 @@ describe('orchestration: mocked LLM and no tools', () => { createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent in no-tool test')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + pendingTurnEvents: [], + mcpServers: [], }); for await (const _event of orchestrator.send([{ type: EventType.USER_MESSAGE, content: 'hello' }])) { diff --git a/packages/trueforge-core/tests/orchestration/orchestrationApproval.test.ts b/packages/trueforge-core/tests/orchestration/orchestrationApproval.test.ts index 05373e1b9..1d3b0d4a9 100644 --- a/packages/trueforge-core/tests/orchestration/orchestrationApproval.test.ts +++ b/packages/trueforge-core/tests/orchestration/orchestrationApproval.test.ts @@ -1,5 +1,5 @@ /** Pause on write_note approval, then resume after allow or deny. */ -import { EventType, newEventId } from '../../src/core/events/schema'; +import { EventType, newEventId, type MCPServerInitInfo } from '../../src/core/events/schema'; import type { IToolSet, ToolSource } from '../../src/core/mcp/IMCPServer'; import { toolResultResponse } from '../../src/core/mcp/IMCPServer'; import { ToolSet } from '../../src/core/mcp/ToolSet'; @@ -423,6 +423,54 @@ describe('orchestration: a policy that lands mid-pause resolves an existing pend expect(callTool).toHaveBeenCalledTimes(1); }); + it('patches the policy onto the MCP server record from the previous turn', async () => { + const { orchestrator, callTool } = makeApprovalHarnessWithUserToolSet(ROOT_FINAL, [ + { + id: 'notes', + name: POLICY_SERVER_NAME, + session_id: 'mcp-session', + transport_type: 'streamable-http', + }, + ]); + + for await (const _event of orchestrator.send([{ type: EventType.USER_MESSAGE, content: 'hello' }])) { + void _event; + } + const iterator = orchestrator.execute({ signal: new AbortController().signal }); + expect((await driveUntilPauseOrDone(iterator)).kind).toBe('paused'); + + for await (const _batch of orchestrator.send([ + { + type: EventType.USER_TOOL_APPROVAL_POLICY, + id: newEventId(), + created_at: new Date().toISOString(), + policies: [{ server_name: POLICY_SERVER_NAME, name: WRITE_NOTE_TOOL_NAME, policy: { type: 'allow_session' } }], + }, + ])) { + void _batch; + } + orchestrator.wake(); + + const resumed = await driveUntilPauseOrDone(iterator); + expect(resumed.kind).toBe('done'); + if (resumed.kind !== 'done') { + throw new Error('expected turn to finish after the policy landed'); + } + expect(resumed.events[0]).toMatchObject({ + type: InternalEventType.APPROVAL_POLICY_APPLY, + mcp_servers_patches: [ + { + id: 'notes', + name: POLICY_SERVER_NAME, + session_id: 'mcp-session', + transport_type: 'streamable-http', + approval_policies: { [WRITE_NOTE_TOOL_NAME]: { type: 'allow_session' } }, + }, + ], + }); + expect(callTool).toHaveBeenCalledTimes(1); + }); + it('applies a decision and policy in accepted order', async () => { const { orchestrator, callTool } = makeApprovalHarnessWithUserToolSet(ROOT_FINAL); @@ -598,6 +646,8 @@ describe('AgentThreadOrchestrator.send: approval policy validation', () => { createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent in policy test')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + pendingTurnEvents: [], + mcpServers: [], }); }); @@ -787,11 +837,16 @@ function makeApprovalHarness(finalReply: string): { createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent in approval test')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + pendingTurnEvents: [], + mcpServers: [], }); return { orchestrator, thread, callTool }; } -function makeApprovalHarnessWithUserToolSet(finalReply: string): { +function makeApprovalHarnessWithUserToolSet( + finalReply: string, + mcpServers?: MCPServerInitInfo[], +): { orchestrator: AgentThreadOrchestrator; thread: AgentThread; callTool: jest.Mock; @@ -835,6 +890,8 @@ function makeApprovalHarnessWithUserToolSet(finalReply: string): { createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent in approval policy test')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + pendingTurnEvents: [], + mcpServers: mcpServers ?? [], }); return { orchestrator, thread, callTool, toolSet }; } diff --git a/packages/trueforge-core/tests/orchestration/orchestrationMcpAuth.test.ts b/packages/trueforge-core/tests/orchestration/orchestrationMcpAuth.test.ts index c38a22b0e..973e51e01 100644 --- a/packages/trueforge-core/tests/orchestration/orchestrationMcpAuth.test.ts +++ b/packages/trueforge-core/tests/orchestration/orchestrationMcpAuth.test.ts @@ -70,6 +70,8 @@ describe('orchestration: MCP auth continuation', () => { createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + pendingTurnEvents: [], + mcpServers: [], }); const iterator = orchestrator.execute({ signal: new AbortController().signal }); diff --git a/packages/trueforge-core/tests/orchestration/orchestrationUserMessageReset.test.ts b/packages/trueforge-core/tests/orchestration/orchestrationUserMessageReset.test.ts index 87c5e239d..a8ad433ed 100644 --- a/packages/trueforge-core/tests/orchestration/orchestrationUserMessageReset.test.ts +++ b/packages/trueforge-core/tests/orchestration/orchestrationUserMessageReset.test.ts @@ -131,6 +131,8 @@ describe('orchestration: user-message reset', () => { createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + pendingTurnEvents: [], + mcpServers: [], }); const sendEvents = []; @@ -220,6 +222,8 @@ describe('orchestration: user-message reset', () => { createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + pendingTurnEvents: [], + mcpServers: [], }); const sendEvents: unknown[] = []; @@ -304,6 +308,8 @@ describe('orchestration: user-message reset', () => { createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + pendingTurnEvents: [], + mcpServers: [], }); await expect( diff --git a/packages/trueforge-core/tests/orchestration/orchestrationWithTools.test.ts b/packages/trueforge-core/tests/orchestration/orchestrationWithTools.test.ts index 76de4ab7f..3901012c1 100644 --- a/packages/trueforge-core/tests/orchestration/orchestrationWithTools.test.ts +++ b/packages/trueforge-core/tests/orchestration/orchestrationWithTools.test.ts @@ -191,6 +191,8 @@ describe('orchestration: dynamic sub-agent', () => { createDynamicSubAgentThread: createSubAgentThread, tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + pendingTurnEvents: [], + mcpServers: [], }; const orchestrator = new AgentThreadOrchestrator(orchestratorInput); 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 80b16ac83..ba4f60323 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts @@ -1754,6 +1754,7 @@ export class SessionsClient { * @throws {@link TrueForge.ForbiddenError} * @throws {@link TrueForge.NotFoundError} * @throws {@link TrueForge.ConflictError} + * @throws {@link TrueForge.InternalServerError} * @throws {@link errors.TrueForgeError} * @throws {@link errors.TrueForgeTimeoutError} * @@ -1872,6 +1873,17 @@ export class SessionsClient { }), _response.rawResponse, ); + case 500: + throw new TrueForge.InternalServerError( + serializers.RequestErrorResponse.parseOrThrow(_response.error.body, { + unrecognizedObjectKeys: "passthrough", + allowUnrecognizedUnionMembers: true, + allowUnrecognizedEnumValues: true, + skipValidation: true, + breadcrumbsPrefix: ["response"], + }), + _response.rawResponse, + ); default: throw new errors.TrueForgeError({ statusCode: _response.error.statusCode, diff --git a/packages/trueforge-sdk/tests/wire/sessions.test.ts b/packages/trueforge-sdk/tests/wire/sessions.test.ts index 81a0d6795..0a7996c6e 100644 --- a/packages/trueforge-sdk/tests/wire/sessions.test.ts +++ b/packages/trueforge-sdk/tests/wire/sessions.test.ts @@ -1555,6 +1555,35 @@ describe("SessionsClient", () => { }).rejects.toThrow(TrueForgeTypes.ConflictError); }); + test("create_turn_events (6)", async () => { + const server = mockServerPool.createServer(); + const client = new TrueForge({ maxRetries: 0, token: "test", baseUrl: server.baseUrl }); + const rawRequestBody = { events: [{ type: "user.mcp_auth_continue" }, { type: "user.mcp_auth_continue" }] }; + const rawResponseBody = { error: { message: "message" } }; + + server + .mockEndpoint() + .post("/api/v1/sessions/session_id/turns/turn_id/events") + .jsonBody(rawRequestBody) + .respondWith() + .statusCode(500) + .jsonBody(rawResponseBody) + .build(); + + await expect(async () => { + return await client.sessions.createTurnEvents("session_id", "turn_id", { + events: [ + { + type: "user.mcp_auth_continue", + }, + { + type: "user.mcp_auth_continue", + }, + ], + }); + }).rejects.toThrow(TrueForgeTypes.InternalServerError); + }); + test("subscribe_to_turn (1)", async () => { const server = mockServerPool.createServer(); const client = new TrueForge({ maxRetries: 0, token: "test", baseUrl: server.baseUrl }); diff --git a/packages/trueforge/.env.example b/packages/trueforge/.env.example index 684ed8de2..afa73b235 100644 --- a/packages/trueforge/.env.example +++ b/packages/trueforge/.env.example @@ -115,6 +115,8 @@ REDIS_URL=redis://localhost:6379 # TURN_STREAM_POST_COMPLETION_TTL_SECONDS=300 ## A subscribe request stays open for 10 minutes by default. # 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 for one create-turn SSE write; after this the client is dropped but the turn keeps running. Default 60000. # TURN_SSE_STREAM_WRITE_TIMEOUT_MS=60000 diff --git a/packages/trueforge/scripts/write-openapi.ts b/packages/trueforge/scripts/write-openapi.ts index 0ccac01ae..bb6a520bf 100644 --- a/packages/trueforge/scripts/write-openapi.ts +++ b/packages/trueforge/scripts/write-openapi.ts @@ -84,20 +84,20 @@ const app = createServerApp({ }), resolveSkillStore: () => skillStore, resolveTurnSkillsStore: () => skillStore, - resolveSandboxProviderStore: () => sandboxProviderStore, - resolveWebSearchProviderStore: () => webSearchProviderStore, resolveAgentStore: () => agentStore, agentStore, sandboxEnvironmentStore, sandboxProviderStore, + webSearchProviderStore, withTransaction: callback => db.transaction().execute(callback), scheduleStore: new SqliteScheduleStore(db), tokenStore, sessionStore, sessionMetricsStore: new SqliteSessionMetricsStore(db), sessions: new Sessions({ sessionStore }), - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: winston.createLogger({ silent: true }) }), requestReplyRouter: new RequestReplyRouter(), + resolveTurnExecutionStores: () => Promise.reject(new Error('turn stores are not used while writing OpenAPI')), eventSubscriptions: new EventSubscriptionRegistry(undefined), logger: winston.createLogger({ silent: true }), oidcClient: undefined, diff --git a/packages/trueforge/src/apis/agents.ts b/packages/trueforge/src/apis/agents.ts index b083460c1..37a782207 100644 --- a/packages/trueforge/src/apis/agents.ts +++ b/packages/trueforge/src/apis/agents.ts @@ -40,9 +40,9 @@ export interface AgentsRouterDeps { resolveModelProviderStore: (c: Context) => IModelProviderStore; resolveMcpServerStore: (c: Context) => IMcpServerStore; resolveSkillStore: ResolveSkillStore; - resolveSandboxProviderStore: (c: Context) => ISandboxProviderStore; + sandboxProviderStore: ISandboxProviderStore; sandboxEnvironmentStore: ISandboxEnvironmentStore; - resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; + webSearchProviderStore: IWebSearchProviderStore; withTransaction: WithTransaction; resolveRequestContext: ResolveRequestContext; authorizer: Authorizer; @@ -125,9 +125,9 @@ export function createAgentsRouter(deps: AgentsRouterDeps(deps: AgentsRouterDeps(deps: { sandboxEnvironmentStore: ISandboxEnvironmentStore; - resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; + webSearchProviderStore: IWebSearchProviderStore; withTransaction: WithTransaction; logger: Logger; resolveRequestContext: ResolveRequestContext; @@ -46,7 +45,7 @@ export function createCapabilitiesRouter(deps: { const settingsEnabled = hasAdminRole(requestContext); const webSearchEnabled = await hasConfiguredWebSearchProvider({ tenant_id: requestContext.tenant_id, - store: deps.resolveWebSearchProviderStore(c), + store: deps.webSearchProviderStore, }); return c.json( { diff --git a/packages/trueforge/src/apis/sandboxEnvironments.ts b/packages/trueforge/src/apis/sandboxEnvironments.ts index fd7f20e85..bc172ed6d 100644 --- a/packages/trueforge/src/apis/sandboxEnvironments.ts +++ b/packages/trueforge/src/apis/sandboxEnvironments.ts @@ -44,7 +44,7 @@ import { MissingStoredSecretError } from '../utils/secretRedaction'; export interface SandboxEnvironmentsRouterDeps { sandboxEnvironmentStore: ISandboxEnvironmentStore; resolveAgentStore: (c: Context) => IAgentStore; - resolveSandboxProviderStore: (c: Context) => ISandboxProviderStore; + sandboxProviderStore: ISandboxProviderStore; resolveRequestContext: ResolveRequestContext; logger: Logger; } @@ -146,10 +146,7 @@ export function createSandboxEnvironmentsRouter( const created_by_subject = createdBySubjectFromRequestContext(requestContext); try { - const provider = await resolveSandboxProviderRecord( - deps.resolveSandboxProviderStore(c), - requestContext.tenant_id, - ); + const provider = await resolveSandboxProviderRecord(deps.sandboxProviderStore, requestContext.tenant_id); if (provider === undefined) { return { kind: 'error', status: 422, message: 'No sandbox provider configured' }; } @@ -274,7 +271,7 @@ export function createSandboxEnvironmentsRouter( try { const secrets = await store.listSecretsByEnvironment({ environment_id: existing.environment.id }); if (secrets.length > 0) { - const provider = await resolveSandboxProviderRecord(deps.resolveSandboxProviderStore(c), tenant_id); + const provider = await resolveSandboxProviderRecord(deps.sandboxProviderStore, tenant_id); if (provider === undefined) { return c.json({ error: { message: 'No sandbox provider configured' } }, 422); } diff --git a/packages/trueforge/src/apis/sandboxProviders.ts b/packages/trueforge/src/apis/sandboxProviders.ts index 6227767a7..f7c96d7ec 100644 --- a/packages/trueforge/src/apis/sandboxProviders.ts +++ b/packages/trueforge/src/apis/sandboxProviders.ts @@ -1,5 +1,4 @@ import { OpenAPIHono, type RouteHandler } from '@hono/zod-openapi'; -import type { Context } from 'hono'; import { HTTPException } from 'hono/http-exception'; import type { Logger } from 'winston'; import { createdBySubjectFromRequestContext, type ResolveRequestContext } from '../auth/identity'; @@ -13,7 +12,7 @@ import type { SandboxProviderManifest, UpdateSandboxProviderRequest } from '../s import { MissingStoredSecretError, resolveStoredSecretValue, toRedactedSecretValue } from '../utils/secretRedaction'; export interface SandboxProvidersRouterDeps { - resolveSandboxProviderStore: (c: Context) => ISandboxProviderStore; + sandboxProviderStore: ISandboxProviderStore; sandboxEnvironmentStore: ISandboxEnvironmentStore; withTransaction: WithTransaction; logger: Logger; @@ -40,7 +39,7 @@ function storedApiKey(record: SandboxProviderRecord | undefined): string | undef export function createSandboxProvidersRouter(deps: SandboxProvidersRouterDeps) { const getHandler: RouteHandler = async c => { const requestContext = deps.resolveRequestContext(c); - const store = deps.resolveSandboxProviderStore(c); + const store = deps.sandboxProviderStore; const record = await store.getSandboxProvider(requestContext.tenant_id); if (record?.manifest.type !== 'daytona') { return c.json({ error: { message: 'No sandbox provider configured' } }, 404); @@ -58,7 +57,7 @@ export function createSandboxProvidersRouter(deps: SandboxProvider const putHandler: RouteHandler = async c => { const body: UpdateSandboxProviderRequest = c.req.valid('json'); const requestContext = deps.resolveRequestContext(c); - const store = deps.resolveSandboxProviderStore(c); + const store = deps.sandboxProviderStore; const incoming = body.manifest; const resolveManifest = (existing: SandboxProviderRecord | undefined): SandboxProviderManifest => ({ ...incoming, diff --git a/packages/trueforge/src/apis/schedules.ts b/packages/trueforge/src/apis/schedules.ts index ab255f7a8..b2fd7a8f7 100644 --- a/packages/trueforge/src/apis/schedules.ts +++ b/packages/trueforge/src/apis/schedules.ts @@ -4,6 +4,7 @@ import { OpenAPIHono, type RouteHandler } from '@hono/zod-openapi'; import { InvalidPageTokenError, + type ISessionStore, type Sessions, type TurnStreamingEvent, } from '@truefoundry/trueforge-core/agent-session'; @@ -69,20 +70,24 @@ import { getTurnExecutionError, startTurnInProcess } from './turns'; /** Runtime + Context store resolvers needed to start a schedule turn. */ export interface ScheduleTurnExecutionDeps { - scheduleStore: IScheduleStore; - sessions: Sessions; activeTurns: ActiveTurnRegistry; + /** Persistence agent store (schedule agent binding is not caller-scoped). */ + agentStore: IAgentStore; eventSubscriptions: EventSubscriptionRegistry; logger: Logger; - resolveModelProviderStore: (input: { c: Context; runAsAgent?: AgentRecord }) => IModelProviderStore; + /** Needed to resolve MCP servers from the agent spec. */ resolveMcpServerStore: (input: { c: Context; runAsAgent?: AgentRecord }) => IMcpServerWithAuthStore; - resolveSandboxProviderStore: (c: Context) => ISandboxProviderStore; - sandboxEnvironmentStore: ISandboxEnvironmentStore; - resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; - /** Persistence agent store (schedule agent binding is not caller-scoped). */ - agentStore: IAgentStore; - /** Call with `{}` — schedules have no caller `x-tfg-skills`. */ + /** Needed to resolve models from the agent spec. */ + resolveModelProviderStore: (input: { c: Context; runAsAgent?: AgentRecord }) => IModelProviderStore; + /** Needed to resolve skills from the agent spec. */ resolveTurnSkillsStore: (input: ResolveTurnSkillsStoreInput) => Pick; + sandboxEnvironmentStore: ISandboxEnvironmentStore; + sandboxProviderStore: ISandboxProviderStore; + scheduleStore: IScheduleStore; + sessions: Sessions; + /** Previous-turn cancel when a schedule run supersedes a turn that is not in this process. */ + sessionStore: Pick; + webSearchProviderStore: IWebSearchProviderStore; } export interface SchedulesRouterDeps extends ScheduleTurnExecutionDeps { @@ -122,11 +127,12 @@ export async function startScheduleRunOnRequest(params: { agentStore: deps.agentStore, modelProviderStore: deps.resolveModelProviderStore({ c, runAsAgent: prepared.agent }), mcpServerStore: deps.resolveMcpServerStore({ c, runAsAgent: prepared.agent }), - sandboxProviderStore: deps.resolveSandboxProviderStore(c), + sandboxProviderStore: deps.sandboxProviderStore, sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - webSearchProviderStore: deps.resolveWebSearchProviderStore(c), + webSearchProviderStore: deps.webSearchProviderStore, skillStore: deps.resolveTurnSkillsStore({}), logger: deps.logger, + sessionStore: deps.sessionStore, }, }); } @@ -209,6 +215,7 @@ export function createScheduleExecutionRouter(deps: ScheduleTurnEx requestContextFromCreatedBySubject({ tenant_id: item.schedule.tenant_id, created_by_subject: item.schedule.created_by_subject, + user_credential: null, }), ); updateLogContext({ tenant_id: item.schedule.tenant_id }); diff --git a/packages/trueforge/src/apis/sessions.ts b/packages/trueforge/src/apis/sessions.ts index f8f825f7a..39b8e4276 100644 --- a/packages/trueforge/src/apis/sessions.ts +++ b/packages/trueforge/src/apis/sessions.ts @@ -29,7 +29,6 @@ import type { Context } from 'hono'; import { HTTPException } from 'hono/http-exception'; import type { Logger } from 'winston'; import { z } from 'zod'; -import type { ResolveSkillStoreInput } from '../app'; import type { Authorizer } from '../auth/authorizer'; import { createdBySubjectFromRequestContext, type RequestContext, type ResolveRequestContext } from '../auth/identity'; import configuration from '../config'; @@ -85,50 +84,51 @@ export function toWireSession(record: SessionRecord): Session { } export interface SessionsRouterDeps { - sessions: Sessions; - sessionStore: ISessionStore; activeTurns: ActiveTurnRegistry; - resolveModelProviderStore: (input: { c: Context; allowInlineResources?: boolean }) => IModelProviderStore; - resolveMcpServerStore: (input: { c: Context; allowInlineResources?: boolean }) => IMcpServerStore; - resolveSkillStore: (input: ResolveSkillStoreInput) => ISkillStore; - resolveAgentStore: (c: Context) => IAgentStore; - resolveSandboxProviderStore: (c: Context) => ISandboxProviderStore; - sandboxEnvironmentStore: ISandboxEnvironmentStore; + authorizer: Authorizer; + logger: Logger; + /** + * Publishes to another executor: + * - Creating the next turn -> redis request `sessions/cancel` + * when the previous run is not in this process. + * Unset in standalone. + */ redis?: RedisClient | undefined; - resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; + /** + * Used to register routes + * - `sessions/cancel` -> {@link cancelSessionTurnPeerHandler} + */ requestReplyRouter: RequestReplyRouter; + resolveAgentStore: (c: Context) => IAgentStore; + resolveMcpServerStore: (input: { c: Context; allowInlineResources?: boolean }) => IMcpServerStore; + resolveModelProviderStore: (input: { c: Context; allowInlineResources?: boolean }) => IModelProviderStore; resolveRequestContext: ResolveRequestContext; - logger: Logger; - authorizer: Authorizer; -} - -function cancelTurnOnThisExecutor( - activeTurns: ActiveTurnRegistry, - input: { sessionId: string; turnId: string; reason: CancellationReason }, -): boolean { - return activeTurns.cancelIfRunning({ - sessionId: input.sessionId, - turnId: input.turnId, - abortReason: input.reason, - }); + resolveSkillStore: (input: { c: Context; allowInlineResources?: boolean }) => ISkillStore; + sandboxEnvironmentStore: ISandboxEnvironmentStore; + sandboxProviderStore: ISandboxProviderStore; + sessions: Sessions; + sessionStore: ISessionStore; + webSearchProviderStore: IWebSearchProviderStore; } /** - * Peer-facing cancel handler: aborts the turn if it runs in this process. - * 200 = abort fired, 412 = not running here (treated by callers as a no-op). + * Peer-facing cancel handler. 200 means this process found the run and signaled abort. + * The drain writes `turn.done` afterward. 412 means it is not running here. */ export function cancelSessionTurnPeerHandler(activeTurns: ActiveTurnRegistry): RequestReplyRouteHandler { - // Synchronous by nature; the transport expects a Promise and require-await - // forbids an async fn without awaits. return request => { const parsed = CancelPeerBodySchema.safeParse(request.body); if (!parsed.success) { return Promise.resolve({ status: 400, body: { message: 'Invalid sessions/cancel payload' } }); } - const found = cancelTurnOnThisExecutor(activeTurns, { + // Note that we do not wait for the abort to complete. + // This means a turn.done can arrive delayed till any in progress step is completed + // if they do not honor the abort signal. + // This is intentional. + const found = activeTurns.cancelIfRunning({ sessionId: parsed.data.session_id, turnId: parsed.data.turn_id, - reason: parsed.data.reason, + abortReason: parsed.data.reason, }); if (!found) { return Promise.resolve({ status: 412, body: { message: 'Turn is not running on this executor' } }); @@ -151,9 +151,9 @@ export interface CancelTurnDeps { * owning peer over Redis request-reply. Callers state the motive; default is a * plain client cancel. * - * A confirmed abort (this process, or peer HTTP 200) lets TurnHandle persist - * the terminal state. If abort cannot be confirmed, this replica freezes the - * turn in the store so the session is not stuck `running` or `paused`. + * A found run (this process, or peer HTTP 200) is aborted and not frozen. The drain + * writes the terminal state afterward. A stuck call can leave the row non-terminal. + * If the run was not found, this replica freezes the turn in the store. * * Redis timeout and Redis/transport failures are not a clean cancellation — * the owning replica may still be executing — but the turn is still frozen. @@ -198,6 +198,7 @@ export async function cancelSessionTurn( } } catch (error) { const fields = { + session_id: sessionId, turn_id: turnId, owner, ...extractErrorLogFields(error), @@ -212,7 +213,11 @@ export async function cancelSessionTurn( return; } - const aborted = cancelTurnOnThisExecutor(deps.activeTurns, { sessionId, turnId, reason }); + const aborted = deps.activeTurns.cancelIfRunning({ + sessionId, + turnId, + abortReason: reason, + }); if (!aborted) { await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); } @@ -252,12 +257,12 @@ function isSessionOwner({ function validateInlineSessionSpec( deps: Pick< SessionsRouterDeps, - | 'resolveModelProviderStore' | 'resolveMcpServerStore' + | 'resolveModelProviderStore' | 'resolveSkillStore' - | 'resolveSandboxProviderStore' | 'sandboxEnvironmentStore' - | 'resolveWebSearchProviderStore' + | 'sandboxProviderStore' + | 'webSearchProviderStore' >, input: { c: Context; spec: AgentSpec; requestContext: RequestContext }, ): Promise { @@ -269,24 +274,24 @@ function validateInlineSessionSpec( modelProviderStore: deps.resolveModelProviderStore({ c, allowInlineResources: true }), mcpServerStore: deps.resolveMcpServerStore({ c, allowInlineResources: true }), skillStore: deps.resolveSkillStore({ c, allowInlineResources: true }), - sandboxProviderStore: deps.resolveSandboxProviderStore(c), + sandboxProviderStore: deps.sandboxProviderStore, sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - webSearchProviderStore: deps.resolveWebSearchProviderStore(c), + webSearchProviderStore: deps.webSearchProviderStore, }); } type InternalSessionsRouterDeps = Pick< SessionsRouterDeps, - | 'sessions' - | 'resolveModelProviderStore' + | 'authorizer' + | 'resolveAgentStore' | 'resolveMcpServerStore' + | 'resolveModelProviderStore' + | 'resolveRequestContext' | 'resolveSkillStore' - | 'resolveAgentStore' - | 'resolveSandboxProviderStore' | 'sandboxEnvironmentStore' - | 'resolveWebSearchProviderStore' - | 'resolveRequestContext' - | 'authorizer' + | 'sandboxProviderStore' + | 'sessions' + | 'webSearchProviderStore' >; function createGetOrCreateSessionByExternalIdHandler( @@ -502,9 +507,9 @@ export function createSessionsRouter(deps: SessionsRouterDeps) { modelProviderStore: deps.resolveModelProviderStore({ c }), mcpServerStore: deps.resolveMcpServerStore({ c }), skillStore: deps.resolveSkillStore({ c }), - sandboxProviderStore: deps.resolveSandboxProviderStore(c), + sandboxProviderStore: deps.sandboxProviderStore, sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - webSearchProviderStore: deps.resolveWebSearchProviderStore(c), + webSearchProviderStore: deps.webSearchProviderStore, }); } } diff --git a/packages/trueforge/src/apis/settings.ts b/packages/trueforge/src/apis/settings.ts index 5de65847f..a6681ea32 100644 --- a/packages/trueforge/src/apis/settings.ts +++ b/packages/trueforge/src/apis/settings.ts @@ -26,9 +26,9 @@ export interface SettingsRouterDeps { resolveMcpServerStore: (input: { c: Context }) => IMcpServerWithAuthStore; tokenStore: IOAuthTokenStore; resolveSkillStore: ResolveSkillStore; - resolveSandboxProviderStore: (c: Context) => ISandboxProviderStore; + sandboxProviderStore: ISandboxProviderStore; sandboxEnvironmentStore: ISandboxEnvironmentStore; - resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; + webSearchProviderStore: IWebSearchProviderStore; /** Delete routes refuse to remove a catalog entry an agent still references. */ resolveAgentStore: (c: Context) => IAgentStore; withTransaction: WithTransaction; @@ -70,7 +70,7 @@ export function createSettingsRouter(deps: SettingsRouterDeps(deps: SettingsRouterDeps | undefined; + /** Saved agent when the session references one. The caller loads it. */ + runAsAgent?: AgentRecord; +} + +export interface TurnExecutionStores { + agentStore: IAgentStore; + modelProviderStore: IModelProviderStore; + mcpServerStore: IMcpServerWithAuthStore; + skillStore: Pick, 'resolveTurnSkills'>; + sandboxProviderStore: ISandboxProviderStore; + webSearchProviderStore: IWebSearchProviderStore; +} + export interface TurnsRouterDeps { - sessions: Sessions; - sessionStore: ISessionStore; activeTurns: ActiveTurnRegistry; - resolveModelProviderStore: (input: ResolveModelProviderStoreInput) => IModelProviderStore; - resolveMcpServerStore: (input: ResolveMcpServerStoreInput) => IMcpServerWithAuthStore; - resolveTurnSkillsStore: (input: ResolveTurnSkillsStoreInput) => Pick; - resolveAgentStore: (c: Context) => IAgentStore; + authorizer: Authorizer; /** Resumable live turn-event transport: create-turn writes, subscribe polls. */ eventSubscriptions: EventSubscriptionRegistry; - resolveSandboxProviderStore: (c: Context) => ISandboxProviderStore; - sandboxEnvironmentStore: ISandboxEnvironmentStore; - resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; logger: Logger; + /** + * Publishes to another executor: + * - POST /events -> redis request `turns/send-events` + * - Creating the next turn -> redis request `sessions/cancel` + * when the previous run is not in this process. + * Unset in standalone. + */ + redis?: RedisClient | undefined; + /** + * Used to register routes + * - `turns/send-events` -> {@link handlePeerTurnEvents} + */ + requestReplyRouter: RequestReplyRouter; + /** List, get, list events, subscribe, download, and create (referenced-agent check). */ + resolveAgentStore: (requestContext: RequestContext) => IAgentStore; resolveRequestContext: ResolveRequestContext; - authorizer: Authorizer; + /** + * Model, MCP, skill, agent, sandbox, and web-search stores for a turn. + * Create, POST /events, and a peered `turns/send-events` all call this with the caller's + * request context and headers. + */ + resolveTurnExecutionStores: (input: ResolveTurnExecutionStoresInput) => Promise; + /** Create, POST /events, and peered send-events, via the turn execution deps. */ + sandboxEnvironmentStore: ISandboxEnvironmentStore; + /** Sandbox file download. */ + sandboxProviderStore: ISandboxProviderStore; + sessions: Sessions; + sessionStore: ISessionStore; } /** @@ -148,8 +214,43 @@ export type BeginTurnExecutionDeps = Pick< sandboxProviderStore: ISandboxProviderStore; webSearchProviderStore: IWebSearchProviderStore; skillStore: Pick; + /** Loads the previous turn when create must cancel a run that is not in this process. */ + sessionStore: Pick; + redis?: RedisClient | undefined; }; +async function referencedRunAsAgent(input: { + session: SessionHandle; + agentStore: IAgentStore; +}): Promise { + if (input.session.record.agent.type !== 'reference') { + return undefined; + } + return input.agentStore.getAgent({ + tenant_id: input.session.tenant_id, + id: input.session.record.agent.id, + }); +} + +async function executionDepsForCaller(input: { + deps: TurnsRouterDeps; + session: SessionHandle; + requestContext: RequestContext; + requestHeaders?: Record | undefined; + runAsAgent?: AgentRecord; +}): Promise { + const stores = await input.deps.resolveTurnExecutionStores(input); + return { + activeTurns: input.deps.activeTurns, + eventSubscriptions: input.deps.eventSubscriptions, + logger: input.deps.logger, + sandboxEnvironmentStore: input.deps.sandboxEnvironmentStore, + sessionStore: input.deps.sessionStore, + redis: input.deps.redis, + ...stores, + }; +} + interface BeginTurnExecutionParams { session: SessionHandle; input: UserMessage[] | undefined; @@ -379,26 +480,50 @@ export function streamTTLSecondsFor(event: TurnStreamingEvent): number | undefin export function turnStreamId(tenantId: string, sessionId: string, turnId: string): string { return redisKey('agent', 'turn', tenantId, sessionId, turnId, 'stream'); } - /** * Dual-write turn events to the resumable subscription registry, then optionally * forward each sequenced event (SSE path). Shared by streaming and non-streaming * create-turn; the HTTP response lifecycle does not own execution. */ +/** Start and clear the pause timer for the one turn this drain is consuming. */ +interface TurnPausedKeepalive { + start: () => void; + clear: () => void; +} + +/** + * Running-only execution deadline for one turn. + * Each entry to `running` starts a fresh window. Pause clears it. + */ +interface TurnExecutionTimeout { + start: () => void; + clear: () => void; +} + export async function drainTurnEvents(input: { trackedStream: AsyncIterable; turnEventStream: EventSubscription; - maxExecutionTimer: NodeJS.Timeout; + executionTimeout: TurnExecutionTimeout; logger: Logger; + pausedKeepalive: TurnPausedKeepalive; onEvent?: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; }): Promise { - const { trackedStream, turnEventStream, maxExecutionTimer, logger, onEvent } = input; + const { trackedStream, turnEventStream, executionTimeout, logger, pausedKeepalive, onEvent } = input; try { for await (const event of trackedStream) { // Dual-write before any client sink so subscribers can resume after disconnect. const sequenceNumber = await turnEventStream.put(event, { streamTTLSeconds: streamTTLSecondsFor(event), }); + if (event.type === EventType.TURN_UPDATE) { + if (event.state.status === 'paused') { + pausedKeepalive.start(); + executionTimeout.clear(); + } else { + pausedKeepalive.clear(); + executionTimeout.start(); + } + } await onEvent?.(event, sequenceNumber); } } catch (error) { @@ -408,7 +533,7 @@ export async function drainTurnEvents(input: { logger.error('Unexpected error in turn event drain', extractErrorLogFields(error)); } } finally { - clearTimeout(maxExecutionTimer); + executionTimeout.clear(); } } @@ -416,8 +541,79 @@ export async function drainTurnEvents(input: { export interface TurnEventDrainInput { trackedStream: AsyncIterable; turnEventStream: EventSubscription; - maxExecutionTimer: NodeJS.Timeout; + executionTimeout: TurnExecutionTimeout; logger: Logger; + pausedKeepalive: TurnPausedKeepalive; +} + +/** + * Header check, abort, execution timeout, web search, and the per-turn resolver. + * Create and resume each keep their own persist and drain steps. + */ +async function prepareTurnExecution(input: { + session: SessionHandle; + turnId: string; + userRef: string; + requestHeaders: Record | undefined; + deps: BeginTurnExecutionDeps; +}): Promise<{ + abortController: AbortController; + resolver: TurnResourceResolver; + executionTimeout: TurnExecutionTimeout; +}> { + assertGatewayMetadataRequestHeaders(input.requestHeaders); + const abortController = new AbortController(); + const webSearchProvider = await resolveWebSearchProvider({ + tenant_id: input.session.tenant_id, + store: input.deps.webSearchProviderStore, + }); + const sessionAgent = input.session.record.agent; + const resolver = createTurnResolver({ + mcpServerStore: input.deps.mcpServerStore, + skillStore: input.deps.skillStore, + sandboxProviderStore: input.deps.sandboxProviderStore, + sandboxEnvironmentStore: input.deps.sandboxEnvironmentStore, + agentStore: input.deps.agentStore, + modelProviderStore: input.deps.modelProviderStore, + webSearchProvider, + logger: input.deps.logger, + signal: abortController.signal, + userRef: input.userRef, + session: input.session, + turnMetadata: { + sessionId: input.session.session_id, + turnId: input.turnId, + ...(sessionAgent.type === 'reference' ? { agent: { id: sessionAgent.id, name: sessionAgent.name } } : {}), + ...(input.requestHeaders === undefined ? {} : { requestHeaders: input.requestHeaders }), + }, + }); + let executionTimer: NodeJS.Timeout | undefined; + const executionTimeout: TurnExecutionTimeout = { + start: () => { + if (executionTimer !== undefined) { + clearTimeout(executionTimer); + } + executionTimer = setTimeout(() => { + executionTimer = undefined; + if (!abortController.signal.aborted) { + abortController.abort(CancellationReason.ServerExecutionTimeout); + } + }, configuration.SERVER_EXECUTION_TIMEOUT_SECONDS * 1000); + executionTimer.unref(); + }, + clear: () => { + if (executionTimer === undefined) { + return; + } + clearTimeout(executionTimer); + executionTimer = undefined; + }, + }; + return { + abortController, + resolver, + executionTimeout, + }; } /** @@ -428,8 +624,6 @@ export async function beginTurnExecution( params: BeginTurnExecutionParams, ): Promise<{ turn: TurnHandle; drainInput: TurnEventDrainInput }> { const { session, input, previous_turn_id: previousTurnId, userRef, requestHeaders, deps } = params; - // Fail closed on a bad inbound header before minting a turn id / starting execution. - assertGatewayMetadataRequestHeaders(requestHeaders); const sessionId = session.session_id; const turnId = newId(); const tenantId = session.tenant_id; @@ -439,54 +633,33 @@ export async function beginTurnExecution( session_id: sessionId, turn_id: turnId, }); - const sessionAgent = session.record.agent; - - const abortController = new AbortController(); - const webSearchProvider = await resolveWebSearchProvider({ - tenant_id: tenantId, - store: deps.webSearchProviderStore, - }); - const resolver = createTurnResolver({ - mcpServerStore: deps.mcpServerStore, - skillStore: deps.skillStore, - sandboxProviderStore: deps.sandboxProviderStore, - sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - agentStore: deps.agentStore, - modelProviderStore: deps.modelProviderStore, - webSearchProvider, - logger: deps.logger, - signal: abortController.signal, - userRef, + const { abortController, resolver, executionTimeout } = await prepareTurnExecution({ session, - turnMetadata: { - sessionId: session.session_id, - turnId, - ...(sessionAgent.type === 'reference' ? { agent: { id: sessionAgent.id, name: sessionAgent.name } } : {}), - ...(requestHeaders === undefined ? {} : { requestHeaders }), - }, + turnId, + userRef, + requestHeaders, + deps, }); // First turn only: derive the title from the first user message. The store // never overwrites an existing title. const title = session.record.last_turn_id ? undefined : deriveSessionTitle(input); + // Freeze the predecessor in the store and leave its generator parked. A paused previous + // turn stays in memory until the keepalive timer offloads it. That is deliberate. + // TODO: freezing drops the previous turn's token cost. Fix when paused turns persist metrics. const turn = await session.createTurn({ turn_id: turnId, active_executor_id: configuration.EXECUTOR_ID, resolver, - abortController: abortController, + abortController, previous_turn_id: previousTurnId, input, update_session_title_if_not_exist: title, }); - const maxExecutionTimer = setTimeout(() => { - if (!abortController.signal.aborted) { - abortController.abort(CancellationReason.ServerExecutionTimeout); - } - }, configuration.SERVER_EXECUTION_TIMEOUT_SECONDS * 1000); - maxExecutionTimer.unref(); - + // A new turn is running immediately. Pause and the later resume are handled from the drain. + executionTimeout.start(); const trackedStream = deps.activeTurns.track({ abortController, stream: turn.stream(), @@ -501,8 +674,16 @@ export async function beginTurnExecution( drainInput: { trackedStream, turnEventStream, - maxExecutionTimer, + executionTimeout, logger: deps.logger, + pausedKeepalive: { + start: () => { + deps.activeTurns.startPausedKeepalive({ sessionId, turnId: turn.id }); + }, + clear: () => { + deps.activeTurns.clearPausedKeepalive({ sessionId, turnId: turn.id }); + }, + }, }, }; } @@ -529,6 +710,146 @@ export async function startTurnInProcess(params: BeginTurnExecutionParams): Prom return turn; } +function cancellationReason(reason: unknown): CancellationReason { + switch (reason) { + case CancellationReason.ServerExecutionTimeout: + case CancellationReason.CancelledForNextTurn: + case CancellationReason.Abandoned: + case CancellationReason.ClientCancelled: + return reason; + default: + return CancellationReason.ClientCancelled; + } +} + +/** + * 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. + */ +async function resumePausedTurn(input: { + session: SessionHandle; + turnId: string; + requestContext: RequestContext; + requestHeaders?: Record | undefined; + deps: BeginTurnExecutionDeps; +}): Promise<{ turn: TurnHandle; signal: AbortSignal; startDrain: () => void }> { + // Same chain as create: model, MCP, sandbox, and the detached drain keep these ids. + updateLogContext({ + tenant_id: input.session.tenant_id, + session_id: input.session.session_id, + turn_id: input.turnId, + }); + const { abortController, resolver, executionTimeout } = await prepareTurnExecution({ + session: input.session, + turnId: input.turnId, + userRef: input.requestContext.subject.id, + requestHeaders: input.requestHeaders, + deps: input.deps, + }); + const sessionId = input.session.session_id; + // Visible to cancel and shutdown before rebuildTurn returns a handle. + input.deps.activeTurns.reservePausedTurn({ + sessionId, + turnId: input.turnId, + abortController, + }); + + let turn: TurnHandle; + try { + turn = await input.session.rebuildTurn({ + turn_id: input.turnId, + active_executor_id: configuration.EXECUTOR_ID, + resolver, + abortController, + }); + } catch (error) { + executionTimeout.clear(); + // Cancel aborted this rebuilding run and did not freeze. If the stream never starts, + // freeze here so the row still becomes terminal. Offload leaves the row paused. + if ( + abortController.signal.aborted && + !isTurnOffloadReason(abortController.signal.reason) && + !(error instanceof TurnOwnershipLostError) + ) { + try { + await input.session.freezeTurn({ + turn_id: input.turnId, + reason: cancellationReason(abortController.signal.reason), + }); + } catch (freezeError) { + if (!(freezeError instanceof TurnNotFoundError)) { + input.deps.activeTurns.releaseReservedTurn({ sessionId, turnId: input.turnId }); + throw new Error(`Failed to freeze turn ${input.turnId} after rebuild failed`, { cause: freezeError }); + } + } + } + input.deps.activeTurns.releaseReservedTurn({ sessionId, turnId: input.turnId }); + throw error; + } + + const turnEventStream = input.deps.eventSubscriptions.get( + 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 }); + } catch (error) { + input.deps.logger.warn('Paused turn stream resume failed', extractErrorLogFields(error)); + executionTimeout.clear(); + abortController.abort(TurnOffloadReason); + const tracked = input.deps.activeTurns.track({ + abortController, + stream: turn.stream(), + turn, + }); + await tracked.return(undefined); + throw error; + } + + let drainStarted = false; + return { + turn, + signal: abortController.signal, + startDrain: () => { + if (drainStarted) { + return; + } + drainStarted = true; + const trackedStream = input.deps.activeTurns.track({ + abortController, + stream: turn.stream(), + turn, + }); + // Resume starts execute already paused, so no new paused `turn.update` is yielded + // unless the turn runs and pauses again. Arm here or this run never offloads. + input.deps.activeTurns.startPausedKeepalive({ + sessionId: input.session.session_id, + turnId: turn.id, + }); + void drainTurnEvents({ + trackedStream, + turnEventStream, + executionTimeout, + logger: input.deps.logger, + pausedKeepalive: { + start: () => { + input.deps.activeTurns.startPausedKeepalive({ + sessionId: input.session.session_id, + turnId: turn.id, + }); + }, + clear: () => { + input.deps.activeTurns.clearPausedKeepalive({ + sessionId: input.session.session_id, + turnId: turn.id, + }); + }, + }, + }); + }, + }; +} + /** Mapped client error from turn execution; undefined means rethrow. */ export interface TurnExecutionError { status: 400 | 404 | 422; @@ -594,8 +915,642 @@ function isSessionOwner({ return created_by_subject.subject_id === subject_id; } -const FORBIDDEN_SESSION_ACCESS = 'Only the session creator can access this session'; -const FORBIDDEN_CREATE_TURN = 'Only the session creator can create turns'; +/** Request-reply path the owner serves for inbound turn events. */ +const TURNS_SEND_EVENTS_PATH = 'turns/send-events'; + +const PeerRequestContextSchema = z.object({ + tenant_id: z.string(), + subject: z.object({ + id: z.string(), + type: z.string(), + display_name: z.string(), + }), + roles: z.array(z.string()), + user_credential: z.string().nullable(), +}); + +const SendEventsPeerBodySchema = z.object({ + tenant_id: z.string(), + session_id: z.string(), + turn_id: z.string(), + events: z.array(TurnUserEventSchema), + request_context: PeerRequestContextSchema, + request_headers: z.record(z.string(), z.string()).optional(), +}); + +const PeerEventsReplySchema = z.object({ + events: z.array(TurnUserEventSchema), +}); + +const PeerMessageSchema = z.object({ + message: z.string(), +}); + +type SendEventsResult = + | { type: 'applied'; events: TurnUserEvent[] } // the events were applied successfully + | { type: 'invalid'; message: string } // the events were invalid + | { type: 'conflict'; message: string } // the turn is not accepting events + | { type: 'failed'; message: string }; // this executor ran the request and failed while applying the events + +export enum PeerOutcome { + /** The owner did not receive this request: heartbeat missing, or publish reached nobody. */ + NoResponder = 'no_responder', + /** The peer replied that it does not own the row. The batch was not applied. */ + NotOwner = 'not_owner', + /** The request may already have been applied but we did not receive the reply. Do not send it again. */ + OwnerNotGone = 'owner_not_gone', +} + +type ForwardedSendResult = SendEventsResult | { type: 'peer'; peerOutcome: PeerOutcome }; + +export enum TurnOwnershipAction { + /** Handle on the local ActiveTurn. */ + Run = 'run', + /** This replica owns the turn but has no ActiveTurn: rebuild it, then handle. */ + Rebuild = 'rebuild', + /** The turn cannot accept this request (terminal, or running with its executor gone). */ + Reject = 'reject', +} + +/** Send a batch to the live handle and persist each yielded inbox batch. */ +async function applyTurnEvents(input: { + handle: TurnHandle; + events: TurnUserEvent[]; + sessionStore: Pick; + sessionId: string; + turnId: string; +}): Promise { + try { + for await (const batch of input.handle.send(input.events)) { + await input.sessionStore.insertTurnInboundEvents({ + session_id: input.sessionId, + turn_id: input.turnId, + active_executor_id: input.handle.record.active_executor_id, + events: batch, + }); + } + } catch (error) { + if (error instanceof AgentHarnessError && error.code === 'invalid_send_input') { + return { type: 'invalid', message: error.message }; + } + if (error instanceof TurnEventAlreadyExistsError) { + return { type: 'conflict', message: error.message }; + } + throw error; + } + return { type: 'applied', events: input.events }; +} + +/** + * Run or rebuild a turn this executor owns, then apply the batch. + * An HTTP send that already owns the row, a peered send, and a steal that just claimed the row all come here. + * If the row is claimed away before the batch is applied, return `'retry'`. + * The peered send turns that into 412. A peered request does not steal. + */ +async function rebuildOptionallyAndRun(input: { + session: SessionHandle; + turn: TurnRecord; + events: TurnUserEvent[]; + sessionStore: Pick; + requestContext: RequestContext; + requestHeaders: Record | undefined; + deps: BeginTurnExecutionDeps; +}): Promise { + const sessionId = input.session.session_id; + const turnId = input.turn.turn_id; + try { + const handle = input.deps.activeTurns.getTurnHandle({ sessionId, turnId }); + const hasLocalActiveTurn = handle !== undefined; + let action: TurnOwnershipAction; + if (!isNonTerminalTurnState(input.turn.state)) { + // done, cancelled, and error cannot accept a send + action = TurnOwnershipAction.Reject; + } else if (hasLocalActiveTurn) { + // we already have a local active turn + action = TurnOwnershipAction.Run; + } else if (input.turn.state.status === 'paused') { + // a paused turn, but we don't have a local active turn + action = TurnOwnershipAction.Rebuild; + } else { + // a running state turn, either + // - cancelled but hasn't reached turn.done + // - stuck in running state because the executor died suddenly + action = TurnOwnershipAction.Reject; + } + + const status = input.turn.state.status; + input.deps.logger.info('Owned turn send', { action, status }); + + switch (action) { + case TurnOwnershipAction.Run: { + if (handle === undefined) { + throw new Error('Unreachable: run requires a local turn'); + } + // Drop the pause timer before wake. If it fired while this handler held the lock, + // the callback sees the timer was replaced and does not offload. A send that leaves + // the turn paused rearms; a later `running` update clears that timer from the drain. + input.deps.activeTurns.clearPausedKeepalive({ sessionId, turnId }); + try { + const applied = await applyTurnEvents({ + handle, + events: input.events, + sessionStore: input.sessionStore, + sessionId, + turnId, + }); + input.deps.logger.info('Owned turn send finished', { action, status, reason: 'run', result: applied.type }); + return applied; + } finally { + if (handle.state.status === 'paused') { + input.deps.activeTurns.startPausedKeepalive({ sessionId, turnId }); + } + } + } + case TurnOwnershipAction.Rebuild: { + // A cancel on paused turn aborts the local run and leaves the turn status paused in store + // until the drain writes turn.done + // getTurnHandle hides that run. Reject instead of rebuilding over it. + // we will not be able to apply events on cancelled turn anyway. + if (input.deps.activeTurns.isCancelled({ sessionId, turnId })) { + input.deps.logger.info('Owned turn send finished', { + action, + status, + reason: 'cancelled-locally', + result: 'conflict', + }); + return { type: 'conflict', message: 'Turn is cancelled, cannot apply events' }; + } + const resumed = await resumePausedTurn({ + session: input.session, + turnId, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + deps: input.deps, + }); + + // when we start resumePausedTurn, we first stage the turn in ActiveTurns while it is being rebuilt + // Case 1: + // Happy path: the turn is resumed and events are applied + // resumed? -> applyTurnEvents -> startDrain -> return applied + // Case 2: + // While we are rebuilding, the turn can be cancelled or server shutdown can happen + // We startDrain regardless so the loop exits (turn.done or offloaded) + // aborted? -> startDrain (without await) -> return conflict + // Case 3: + // Applying events fails after the turn is resumed + // We startDrain so turn gets parked and sticks around in memory + // apply? -> startDrain (without await) -> return conflict + try { + if (resumed.signal.aborted) { + input.deps.logger.info('Owned turn send finished', { + action, + status, + reason: 'aborted-during-rebuild', + result: 'conflict', + }); + return { type: 'conflict', message: 'Turn is cancelled or unable to accept events' }; + } + const applied = await applyTurnEvents({ + handle: resumed.turn, + events: input.events, + sessionStore: input.sessionStore, + sessionId, + turnId, + }); + input.deps.logger.info('Owned turn send finished', { + action, + status, + reason: 'rebuild', + result: applied.type, + }); + return applied; + } finally { + resumed.startDrain(); + } + } + case TurnOwnershipAction.Reject: { + const reason = input.turn.state.status === 'running' ? 'running-without-generator' : 'not-accepting'; + input.deps.logger.info('Owned turn send finished', { action, status, reason, result: 'conflict' }); + return { + type: 'conflict', + message: + input.turn.state.status === 'running' + ? 'Turn is running but its executor is gone' + : 'Turn is cancelled or unable to accept events', + }; + } + } + } catch (error) { + if (error instanceof TurnOwnershipLostError) { + input.deps.logger.info('Owned turn send finished', { + status: input.turn.state.status, + reason: 'ownership-lost', + result: 'retry', + }); + return 'retry'; + } + input.deps.logger.error('Failed to apply events to an owned turn', extractErrorLogFields(error)); + const cause = error instanceof Error ? error.message : String(error); + return { type: 'failed', message: `Failed to apply events: ${cause}` }; + } +} + +/** Ask the owning replica to apply the events. A peer outcome means this replica must decide again. */ +async function forwardTurnEvents(input: { + redis: RedisClient | undefined; + activeExecutorId: string; + sessionId: string; + turnId: string; + events: TurnUserEvent[]; + requestContext: RequestContext; + requestHeaders?: Record | undefined; + logger: Logger; +}): Promise { + if (!input.redis) { + // unreachable code, we should always have a redis client + input.logger.error('Failed to apply turn events: no redis client'); + return { type: 'failed', message: 'Unable to reach the turn executor' }; + } + const { requestContext } = input; + const body = SendEventsPeerBodySchema.parse({ + tenant_id: requestContext.tenant_id, + session_id: input.sessionId, + turn_id: input.turnId, + events: input.events, + request_context: { + tenant_id: requestContext.tenant_id, + subject: requestContext.subject, + roles: requestContext.roles, + user_credential: requestContext.user_credential, + }, + ...(input.requestHeaders === undefined ? {} : { request_headers: input.requestHeaders }), + }); + try { + const reply = await redisRequest({ + redis: input.redis, + executorId: input.activeExecutorId, + path: TURNS_SEND_EVENTS_PATH, + request: { body: jsonValueFrom(body) }, + options: { + replyTimeoutMs: configuration.REDIS_REQUEST_REPLY_TIMEOUT_MS, + pollIntervalMs: configuration.REDIS_REQUEST_REPLY_POLL_INTERVAL_MS, + }, + }); + + input.logger.info('Turn send peer replied', { activeExecutorId: input.activeExecutorId, status: reply.status }); + switch (reply.status) { + case 200: { + const parsed = PeerEventsReplySchema.safeParse(reply.body); + if (!parsed.success) { + // ideally unreachable because peer should send a valid response + // this can cause a bad state for e.g. if peer accepted the batch, applied approvals, but then sent an invalid response + input.logger.warn('Turn send peer reply was invalid', { activeExecutorId: input.activeExecutorId }); + return { type: 'failed', message: 'Turn executor sent an invalid response' }; + } + return { type: 'applied', events: parsed.data.events }; + } + // The peer replied the events were invalid + case 400: { + const parsed = PeerMessageSchema.safeParse(reply.body); + return { + type: 'invalid', + message: parsed.success ? parsed.data.message : `Turn rejected the events`, + }; + } + // The peer replied turn is not accepting events + case 409: { + const parsed = PeerMessageSchema.safeParse(reply.body); + return { + type: 'conflict', + message: parsed.success ? parsed.data.message : `Turn is not accepting events`, + }; + } + // The peer replied that it does not own the row. The batch was not applied. + case 412: + return { type: 'peer', peerOutcome: PeerOutcome.NotOwner }; + case 500: + default: { + const parsed = PeerMessageSchema.safeParse(reply.body); + return { + type: 'failed', + message: parsed.success ? parsed.data.message : `Failed to apply events`, + }; + } + } + } catch (error) { + if (error instanceof NoResponderError) { + input.logger.info('Turn send peer missed', { activeExecutorId: input.activeExecutorId, detail: error.detail }); + // The request was published, then the heartbeat disappeared while waiting for the reply. + // The owner may have applied the batch. Do not send it again. + if (error.detail === NO_RESPONDER_HEARTBEAT_LOST_DETAIL) { + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } + // Heartbeat was already missing, or publish reached nobody. This batch was not delivered. + return { type: 'peer', peerOutcome: PeerOutcome.NoResponder }; + } else if (error instanceof RequestTimeoutError) { + input.logger.info('Turn send peer missed', { activeExecutorId: input.activeExecutorId, detail: 'timeout' }); + // The owner may have applied the batch before the reply wait ended. + // Do not send it again; the client retries. + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } + input.logger.warn('Turn send peer missed', { + activeExecutorId: input.activeExecutorId, + detail: 'transport', + ...extractErrorLogFields(error), + }); + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } +} + +async function handleParsedPeerTurnEvents(deps: TurnsRouterDeps, body: z.infer) { + const session = await deps.sessions.get({ + tenant_id: body.tenant_id, + session_id: body.session_id, + }); + if (!session) { + deps.logger.info('Peered turn send finished', { result: 'not-found', reason: 'session-not-found' }); + return { status: 404, body: { message: `Session not found` } }; + } + const result = await deps.activeTurns.withTurnLock( + { sessionId: body.session_id, turnId: body.turn_id }, + async (): Promise => { + const requestContext = requestContextFromWire(body.request_context); + const runAsAgent = await referencedRunAsAgent({ + session, + agentStore: deps.resolveAgentStore(requestContext), + }); + const executionDeps = await executionDepsForCaller({ + deps, + session, + requestContext, + requestHeaders: body.request_headers, + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + }); + // Same run as an owner. A peered request does not steal, so a lost row is 412. + const turn = await session.getTurn(body.turn_id); + if (!turn) { + deps.logger.info('Peered turn send finished', { result: 'conflict', reason: 'turn-not-found' }); + return { type: 'conflict', message: `Turn not found` }; + } + if (turn.active_executor_id !== configuration.EXECUTOR_ID) { + deps.logger.info('Peered turn send finished', { + result: 'not-owner', + reason: 'executor-mismatch', + activeExecutorId: turn.active_executor_id, + }); + return { type: 'not_owner' }; + } + const owned = await rebuildOptionallyAndRun({ + session, + turn, + events: body.events, + sessionStore: deps.sessionStore, + requestContext, + requestHeaders: body.request_headers, + deps: executionDeps, + }); + if (owned === 'retry') { + deps.logger.info('Peered turn send finished', { result: 'not-owner', reason: 'ownership-lost' }); + return { type: 'not_owner' }; + } + return owned; + }, + ); + switch (result.type) { + case 'applied': + return { status: 200, body: jsonValueFrom({ events: result.events }) }; + case 'invalid': + return { status: 400, body: { message: result.message } }; + case 'not_owner': + return { status: 412, body: { message: 'Turn is no longer owned by this executor' } }; + case 'conflict': + return { status: 409, body: { message: result.message } }; + case 'failed': + return { status: 500, body: { message: result.message } }; + default: { + const _exhaustive: never = result; + throw new Error(`Unreachable send result: ${JSON.stringify(_exhaustive)}`); + } + } +} + +function handlePeerTurnEvents(deps: TurnsRouterDeps): RequestReplyRouteHandler { + return async request => { + const parsed = SendEventsPeerBodySchema.safeParse(request.body); + if (!parsed.success) { + return { status: 400, body: { message: 'Invalid turns/send-events payload' } }; + } + // No HTTP request context on this path. Resume and its drain inherit these ids. + return runWithLogContext( + { + tenant_id: parsed.data.tenant_id, + session_id: parsed.data.session_id, + turn_id: parsed.data.turn_id, + }, + () => handleParsedPeerTurnEvents(deps, parsed.data), + ); + }; +} + +/** + * A user send ends in `rebuildOptionallyAndRun` by one of three paths: + * this process already owns the row and runs it; + * this process forwards to the owner, and that peer runs it; + * the owner never received the batch, this process steals the paused row, and then runs it. + */ + +/** + * One try at delivering a user send. + * We own the turn: run or rebuild it. If the turn is claimed away before the batch is applied, try again. + * The peer says it is not the owner: the batch was not applied, so load the turn and send it to whoever owns it now. + * The batch never reached the owner: steal it when the turn is still paused under that same owner. A running turn is a conflict. Another executor means try again. + * The owner may already have applied the batch: stop, and let the client retry. + */ +async function sendTurnEvents(input: { + session: SessionHandle; + sessionStore: Pick; + turnId: string; + events: TurnUserEvent[]; + requestContext: RequestContext; + requestHeaders: Record | undefined; + deps: BeginTurnExecutionDeps; +}): Promise { + const sessionId = input.session.session_id; + const turn = await input.session.getTurn(input.turnId); + if (!turn) { + input.deps.logger.info('Turn send finished', { reason: 'turn-not-found', result: 'conflict' }); + return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; + } + // done, cancelled, and error cannot accept a send. + if (!isNonTerminalTurnState(turn.state)) { + input.deps.logger.info('Turn send finished', { + reason: 'turn-status-terminal', + result: 'conflict', + status: turn.state.status, + }); + return { type: 'conflict', message: `Turn is not accepting events: ${turn.turn_id}` }; + } + if (turn.active_executor_id === configuration.EXECUTOR_ID) { + return rebuildOptionallyAndRun({ + session: input.session, + turn, + events: input.events, + sessionStore: input.sessionStore, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + deps: input.deps, + }); + } + + // If the turn got paused and parked and then stolen, we need to drop the local turn + await input.deps.activeTurns.dropLocalTurn({ sessionId, turnId: input.turnId }); + + input.deps.logger.info('Turn send forwarding', { + activeExecutorId: turn.active_executor_id, + status: turn.state.status, + }); + // Try to forward the events to the owner + const forwarded = await forwardTurnEvents({ + redis: input.deps.redis, + activeExecutorId: turn.active_executor_id, + sessionId, + turnId: input.turnId, + events: input.events, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + logger: input.deps.logger, + }); + + // If the reply anything other than peering call status, return the result + // that means owner either accepted the events or rejected them + if (forwarded.type !== 'peer') { + input.deps.logger.info('Turn send finished', { reason: 'peer-reply', result: forwarded.type }); + return forwarded; + } + + const refetchedTurn = await input.session.getTurn(input.turnId); + if (refetchedTurn === undefined) { + // Unreachable. Turns are not deleted, and the load above found this row. + input.deps.logger.info('Turn send finished', { reason: 'turn-not-found', result: 'conflict' }); + return { type: 'conflict', message: `Turn not found` }; + } + + // 412: the peer did not apply the batch. Load the row again and send to the current owner. + // Timeout, No Responder (heartbeat lost after publish), and transport errors may already have been + // applied, so this pass does not send again. + if (forwarded.peerOutcome === PeerOutcome.NotOwner || forwarded.peerOutcome === PeerOutcome.OwnerNotGone) { + if (!isNonTerminalTurnState(refetchedTurn.state)) { + input.deps.logger.info('Turn send finished', { + reason: 'turn-status-terminal', + result: 'conflict', + status: refetchedTurn.state.status, + peerOutcome: forwarded.peerOutcome, + }); + return { type: 'conflict', message: 'Turn is not accepting events' }; + } + if (forwarded.peerOutcome === PeerOutcome.NotOwner) { + input.deps.logger.info('Turn send retrying', { + reason: 'peer-not-owner', + result: 'retry', + owner: refetchedTurn.active_executor_id, + }); + return 'retry'; + } + input.deps.logger.info('Turn send finished', { + reason: 'owner-not-gone', + result: 'failed', + peerOutcome: forwarded.peerOutcome, + owner: refetchedTurn.active_executor_id, + status: refetchedTurn.state.status, + }); + return { type: 'failed', message: 'Turn executor did not acknowledge the events' }; + } + + // No Responder (heartbeat missing or zero subscribers) case + // This batch was not delivered. A finished row is a conflict. A new owner has not + // seen it, so send there even if that owner already moved the row to running. + // The same owner with a running row will miss the forward again; do not steal it. + if (!isNonTerminalTurnState(refetchedTurn.state)) { + input.deps.logger.info('Turn send finished', { + reason: 'turn-status-terminal', + result: 'conflict', + status: refetchedTurn.state.status, + peerOutcome: forwarded.peerOutcome, + }); + return { type: 'conflict', message: 'Turn is in terminal state, cannot apply events' }; + } + if (refetchedTurn.active_executor_id !== turn.active_executor_id) { + input.deps.logger.info('Turn send retrying', { + reason: 'owner-changed', + result: 'retry', + fromExecutor: turn.active_executor_id, + owner: refetchedTurn.active_executor_id, + status: refetchedTurn.state.status, + }); + return 'retry'; + } + if (refetchedTurn.state.status === 'running') { + // TODO: this is the case where turn is stuck in running state because the executor died suddenly + // we cannot steal reliably yet - we will handle this in upcoming changes + input.deps.logger.info('Turn send finished', { + reason: 'running-not-delivered', + result: 'conflict', + owner: refetchedTurn.active_executor_id, + }); + return { type: 'conflict', message: 'Turn is running and its executor is gone' }; + } + + // Paused, still owned by the executor that did not receive this batch. + input.deps.logger.info('Turn send stealing', { + fromExecutor: refetchedTurn.active_executor_id, + }); + const won = await input.sessionStore.claimTurnExecutor({ + session_id: sessionId, + turn_id: input.turnId, + expected_executor_id: refetchedTurn.active_executor_id, + new_executor_id: configuration.EXECUTOR_ID, + }); + if (!won) { + input.deps.logger.info('Turn send retrying', { + reason: 'steal-lost', + result: 'retry', + fromExecutor: refetchedTurn.active_executor_id, + }); + return 'retry'; + } + input.deps.logger.info('Turn send steal claimed the row', { + executorId: configuration.EXECUTOR_ID, + }); + return rebuildOptionallyAndRun({ + session: input.session, + turn: refetchedTurn, + events: input.events, + sessionStore: input.sessionStore, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + deps: input.deps, + }); +} + +/** User send. Forwards to the owner and may steal a paused turn whose heartbeat is gone. */ +async function sendTurnEventsWithRetries(input: { + session: SessionHandle; + sessionStore: Pick; + turnId: string; + events: TurnUserEvent[]; + requestContext: RequestContext; + requestHeaders: Record | undefined; + deps: BeginTurnExecutionDeps; +}): Promise { + for (let attempt = 0; attempt < MAX_OWNERSHIP_ATTEMPTS; attempt++) { + const result = await sendTurnEvents(input); + if (result !== 'retry') { + return result; + } + } + input.deps.logger.error('Turn send ownership did not settle', { attempts: MAX_OWNERSHIP_ATTEMPTS }); + return { type: 'failed', message: 'Unable to determine the executor for the turn' }; +} /** DB-backed turns (mounted at /api/v1/sessions). */ export function createTurnsRouter(deps: TurnsRouterDeps) { @@ -612,7 +1567,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { } const allowed = await canReadSession({ shared: session.record.shared, - store: deps.resolveAgentStore(c), + store: deps.resolveAgentStore(requestContext), context: requestContext, authorizer: deps.authorizer, agent_id: session.record.agent.type === 'reference' ? session.record.agent.id : undefined, @@ -647,7 +1602,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { } const allowed = await canReadSession({ shared: session.record.shared, - store: deps.resolveAgentStore(c), + store: deps.resolveAgentStore(requestContext), context: requestContext, authorizer: deps.authorizer, agent_id: session.record.agent.type === 'reference' ? session.record.agent.id : undefined, @@ -660,7 +1615,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { if (!turn) { return c.json({ error: { message: `Turn not found: ${turnId}` } }, 404); } - return c.json({ data: toWireTurn(turn.record) }, 200); + return c.json({ data: toWireTurn(turn) }, 200); }; const downloadSandboxFileHandler: RouteHandler = async c => { @@ -695,14 +1650,14 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { if (!turn) { return c.json({ error: { message: `Turn not found: ${turnId}` } }, 404); } - sandboxId = turn.record.snapshot.sandbox_info?.sandbox_id; + sandboxId = turn.snapshot.sandbox_info?.sandbox_id; if (sandboxId === undefined) { return c.json({ error: { message: `Turn has no sandbox: ${turnId}` } }, 412); } const provider = await resolveSandboxProvider({ tenant_id: requestContext.tenant_id, - store: deps.resolveSandboxProviderStore(c), + store: deps.sandboxProviderStore, logger: deps.logger, sessionId, }); @@ -745,7 +1700,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { } const allowed = await canReadSession({ shared: session.record.shared, - store: deps.resolveAgentStore(c), + store: deps.resolveAgentStore(requestContext), context: requestContext, authorizer: deps.authorizer, agent_id: session.record.agent.type === 'reference' ? session.record.agent.id : undefined, @@ -759,7 +1714,8 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { return c.json({ error: { message: `Turn not found: ${turnId}` } }, 404); } try { - const { data, pagination } = await turn.listEvents({ + const { data, pagination } = await session.listTurnEvents({ + turn_id: turnId, limit: query.limit, page_token: query.page_token, order: query.order, @@ -797,7 +1753,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { let referencedAgent: AgentRecord | undefined; if (session.record.agent.type === 'reference') { const agentId = session.record.agent.id; - const agent = await deps.resolveAgentStore(c).getAgent({ + const agent = await deps.resolveAgentStore(requestContext).getAgent({ tenant_id: requestContext.tenant_id, id: agentId, }); @@ -816,31 +1772,19 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { } // x-tfg-skills / x-tfg-mcp / x-tfg-models on turns only when the session agent is inline. - const allowInlineResources = session.record.agent.type === 'inline'; const turnParams: BeginTurnExecutionParams = { session, input: body.input, previous_turn_id: body.previous_turn_id, userRef: requestContext.subject.id, requestHeaders: c.req.header(), - deps: { - ...deps, - modelProviderStore: deps.resolveModelProviderStore({ - c, - allowInlineResources, - ...(referencedAgent !== undefined ? { runAsAgent: referencedAgent } : {}), - }), - mcpServerStore: deps.resolveMcpServerStore({ - c, - allowInlineResources, - ...(referencedAgent !== undefined ? { runAsAgent: referencedAgent } : {}), - }), - skillStore: deps.resolveTurnSkillsStore({ c, allowInlineResources }), - agentStore: deps.resolveAgentStore(c), - sandboxProviderStore: deps.resolveSandboxProviderStore(c), - sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - webSearchProviderStore: deps.resolveWebSearchProviderStore(c), - }, + deps: await executionDepsForCaller({ + deps, + session, + requestContext, + requestHeaders: c.req.header(), + ...(referencedAgent !== undefined ? { runAsAgent: referencedAgent } : {}), + }), }; try { @@ -902,6 +1846,76 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { } }; + const postTurnEventsHandler: RouteHandler = async c => { + const { session_id: sessionId, turn_id: turnId } = c.req.valid('param'); + const body = c.req.valid('json'); + const requestContext = deps.resolveRequestContext(c); + const session = await deps.sessions.get({ + tenant_id: requestContext.tenant_id, + session_id: sessionId, + }); + if (!session) { + return c.json({ error: { message: `Session not found: ${sessionId}` } }, 404); + } + if ( + !isSessionOwner({ + subject_id: requestContext.subject.id, + created_by_subject: session.record.created_by_subject, + }) + ) { + return c.json({ error: { message: FORBIDDEN_SESSION_ACCESS } }, 403); + } + const turn = await session.getTurn(turnId); + if (!turn) { + return c.json({ error: { message: `Turn not found: ${turnId}` } }, 404); + } + + const createdAt = new Date().toISOString(); + const events: TurnUserEvent[] = body.events.map(payload => ({ + ...payload, + id: newEventId(), + created_at: createdAt, + })); + const requestHeaders = c.req.header(); + + return deps.activeTurns.withTurnLock({ sessionId, turnId }, async () => { + const runAsAgent = await referencedRunAsAgent({ + session, + agentStore: deps.resolveAgentStore(requestContext), + }); + const executionDeps = await executionDepsForCaller({ + deps, + session, + requestContext, + requestHeaders, + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + }); + const result = await sendTurnEventsWithRetries({ + session, + sessionStore: deps.sessionStore, + turnId, + events, + requestContext, + requestHeaders, + deps: executionDeps, + }); + switch (result.type) { + case 'applied': + return c.json({ data: result.events }, 201); + case 'invalid': + return c.json({ error: { message: `Failed to apply events: ${result.message}` } }, 400); + case 'conflict': + return c.json({ error: { message: `Failed to apply events: ${result.message}` } }, 409); + case 'failed': + return c.json({ error: { message: `Failed to apply events: ${result.message}` } }, 500); + default: { + const _exhaustive: never = result; + throw new Error(`Unreachable send result: ${JSON.stringify(_exhaustive)}`); + } + } + }); + }; + const subscribeTurnHandler: RouteHandler = async c => { const { session_id: sessionId, turn_id: turnId } = c.req.valid('param'); const query = c.req.valid('query'); @@ -916,7 +1930,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { return c.json({ error: { message: `Session not found: ${sessionId}` } }, 404); } const allowed = await canReadAgentBoundResource({ - store: deps.resolveAgentStore(c), + store: deps.resolveAgentStore(requestContext), context: requestContext, authorizer: deps.authorizer, agent_id: session.record.agent.type === 'reference' ? session.record.agent.id : undefined, @@ -983,63 +1997,6 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { }); }; - const postTurnEventsHandler: RouteHandler = async c => { - const { session_id: sessionId, turn_id: turnId } = c.req.valid('param'); - const body = c.req.valid('json'); - const requestContext = deps.resolveRequestContext(c); - const session = await deps.sessions.get({ - tenant_id: requestContext.tenant_id, - session_id: sessionId, - }); - if (!session) { - return c.json({ error: { message: `Session not found: ${sessionId}` } }, 404); - } - if ( - !isSessionOwner({ - subject_id: requestContext.subject.id, - created_by_subject: session.record.created_by_subject, - }) - ) { - return c.json({ error: { message: FORBIDDEN_SESSION_ACCESS } }, 403); - } - const turn = await session.getTurn(turnId); - if (!turn) { - return c.json({ error: { message: `Turn not found: ${turnId}` } }, 404); - } - - return deps.activeTurns.withTurnLock({ sessionId, turnId }, async () => { - const turnHandle = deps.activeTurns.getTurnHandle({ sessionId, turnId }); - if (!turnHandle) { - return c.json({ error: { message: `Turn is not running on this server: ${turnId}` } }, 409); - } - - const createdAt = new Date().toISOString(); - const events: TurnUserEvent[] = body.events.map(payload => ({ - ...payload, - id: newEventId(), - created_at: createdAt, - })); - - try { - // Insert each accepted batch before resuming send(), which enqueues it and wakes execution. - for await (const batch of turnHandle.send(events)) { - await deps.sessionStore.insertTurnInboundEvents({ - session_id: sessionId, - turn_id: turnId, - events: batch, - }); - } - } catch (error) { - if (error instanceof AgentHarnessError && error.code === 'invalid_send_input') { - return c.json({ error: { message: error.message } }, 400); - } - throw error; - } - - return c.json({ data: events }, 201); - }); - }; - const router = new OpenAPIHono(); router.openapi(createAndExecuteTurnRoute, createAndExecuteTurnHandler); router.openapi(listTurnsRoute, listTurnsHandler); @@ -1048,5 +2005,6 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { router.openapi(listTurnEventsRoute, listTurnEventsHandler); router.openapi(postTurnEventsRoute, postTurnEventsHandler); router.openapi(subscribeTurnRoute, subscribeTurnHandler); + deps.requestReplyRouter.registerRoute(TURNS_SEND_EVENTS_PATH, handlePeerTurnEvents(deps)); return router; } diff --git a/packages/trueforge/src/apis/webSearchProviders.ts b/packages/trueforge/src/apis/webSearchProviders.ts index 824a3e467..2345c743e 100644 --- a/packages/trueforge/src/apis/webSearchProviders.ts +++ b/packages/trueforge/src/apis/webSearchProviders.ts @@ -1,5 +1,4 @@ import { OpenAPIHono, type RouteHandler } from '@hono/zod-openapi'; -import type { Context } from 'hono'; import { HTTPException } from 'hono/http-exception'; import type { ResolveRequestContext } from '../auth/identity'; import type { IWebSearchProviderStore, WebSearchProviderRecord } from '../db/webSearchProviderStore'; @@ -13,7 +12,7 @@ import { import { MissingStoredSecretError, resolveStoredSecretValue, toRedactedSecretValue } from '../utils/secretRedaction'; export interface WebSearchProvidersRouterDeps { - resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; + webSearchProviderStore: IWebSearchProviderStore; resolveRequestContext: ResolveRequestContext; } @@ -52,7 +51,7 @@ function toWireProvider(record: WebSearchProviderRecord): ConfiguredWebSearchPro export function createWebSearchProvidersRouter(deps: WebSearchProvidersRouterDeps) { const getHandler: RouteHandler = async c => { const requestContext = deps.resolveRequestContext(c); - const record = await deps.resolveWebSearchProviderStore(c).getProvider(requestContext.tenant_id); + const record = await deps.webSearchProviderStore.getProvider(requestContext.tenant_id); if (!record) { return c.json({ error: { message: 'No web search provider configured' } }, 404); } @@ -60,7 +59,7 @@ export function createWebSearchProvidersRouter(deps: WebSearchProvidersRouterDep }; const putHandler: RouteHandler = async c => { - const store = deps.resolveWebSearchProviderStore(c); + const store = deps.webSearchProviderStore; const body: UpdateWebSearchProviderRequest = c.req.valid('json'); const requestContext = deps.resolveRequestContext(c); const provider = body.manifest; diff --git a/packages/trueforge/src/app.ts b/packages/trueforge/src/app.ts index b66bbab55..18024dbcb 100644 --- a/packages/trueforge/src/app.ts +++ b/packages/trueforge/src/app.ts @@ -25,10 +25,10 @@ import { createInternalMetricsRouter } from './apis/sessionMetrics'; import { createInternalSessionsRouter, createSessionsRouter } from './apis/sessions'; import { createSettingsRouter } from './apis/settings'; import { createAvailableSkillsRouter } from './apis/skills'; -import { createTurnsRouter } from './apis/turns'; +import { createTurnsRouter, type ResolveTurnExecutionStoresInput, type TurnExecutionStores } from './apis/turns'; import type { Authenticator } from './auth/authenticator'; import type { Authorizer } from './auth/authorizer'; -import { resolveRequestContext } from './auth/identity'; +import { resolveRequestContext, type RequestContext } from './auth/identity'; import { createAdminAuthMiddleware, createApiKeyAuthMiddleware, @@ -269,62 +269,60 @@ export interface ResolveSkillStoreInput { } export interface ServerDeps { - modelCatalog: ModelCatalog; + activeTurns: ActiveTurnRegistry; + /** Persistence agent store (schedule runs resolve the bound agent without an HTTP caller). */ + agentStore: IAgentStore; + /** Startup-selected authenticator; middleware is built from this once per app. */ + authenticator: Authenticator; + /** Startup-selected agent authorization policy. */ + authorizer: Authorizer; + /** Hands out each turn's resumable event stream to the create and subscribe handlers. */ + eventSubscriptions: EventSubscriptionRegistry; + logger: Logger; mcpCatalog: McpCatalog; - skillCatalog: SkillCatalog; - sandboxCatalog: SandboxCatalog; - webSearchCatalog: WebSearchCatalog; - /** Per-request store: DB singleton, or a token-bound TrueFoundry store in TrueFoundry mode. */ - resolveModelProviderStore: (input: ResolveModelProviderStoreInput) => IModelProviderStore; + modelCatalog: ModelCatalog; + /** Discovered openid-client configuration; undefined when browser login is disabled. */ + oidcClient: Configuration | undefined; + /** Redis client used for redis request-reply */ + redis?: RedisClient | undefined; + /** Redis request-reply router for registering peer handlers. */ + requestReplyRouter: RequestReplyRouter; + /** Per-request store: DB singleton, or a token-bound TrueFoundry decorator in TrueFoundry mode. */ + resolveAgentStore: (requestContext: RequestContext) => IAgentStore; /** * Per-request store: DB singleton, or a token-bound TrueFoundry store in TrueFoundry mode. * The unauthenticated OAuth callback omits `c` and gets the DB persistence store. */ resolveMcpServerStore: (input: ResolveMcpServerStoreInput) => IMcpServerWithAuthStore; - /** Per-request store: DB singleton, or a token-bound TrueFoundry decorator in TrueFoundry mode. */ - resolveAgentStore: (c: Context) => IAgentStore; - /** - * Per-request store: DB singleton, or the env-backed shared store in TrueFoundry mode - * (`TRUEFOUNDRY_SANDBOX_*` + static SETTINGS JSON). - */ - resolveSandboxProviderStore: (c: Context) => ISandboxProviderStore; - resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; + /** Per-request store: DB singleton, or a token-bound TrueFoundry store in TrueFoundry mode. */ + resolveModelProviderStore: (input: ResolveModelProviderStoreInput) => IModelProviderStore; /** Per-request store: DB git skills, or TrueFoundry registry catalog in TrueFoundry mode. */ resolveSkillStore: (input: ResolveSkillStoreInput) => ISkillStore; + /** + * Stores for running a turn. HTTP create and POST /events pass the request. + * A peered `turns/send-events` passes the forwarded context and headers. + */ + resolveTurnExecutionStores: (input: ResolveTurnExecutionStoresInput) => Promise>; /** * Turn skill mounts. HTTP turns may apply `x-tfg-skills` when allowInlineResources is set; * schedules call with `{}` (registry/DB only). */ resolveTurnSkillsStore: (input: ResolveTurnSkillsStoreInput) => Pick, 'resolveTurnSkills'>; - withTransaction: WithTransaction; - tokenStore: IOAuthTokenStore; - scheduleStore: IScheduleStore; - sessionStore: ISessionStore; - sessionMetricsStore: ISessionMetricsStore; - /** Persistence agent store (schedule runs resolve the bound agent without an HTTP caller). */ - agentStore: IAgentStore; - /** Sandbox environment parent + version persistence (no TrueFoundry dual-write). */ + sandboxCatalog: SandboxCatalog; + /** Sandbox environment parent + version persistence, or the shared TrueFoundry store. */ sandboxEnvironmentStore: ISandboxEnvironmentStore; - /** - * Context-free sandbox provider store for internal build progress (persistence or TFY). - * Same backing store as resolveSandboxProviderStore, without request context. - */ + /** DB singleton, or the env-backed shared store in TrueFoundry mode. */ sandboxProviderStore: ISandboxProviderStore; + scheduleStore: IScheduleStore; + sessionMetricsStore: ISessionMetricsStore; sessions: Sessions; - activeTurns: ActiveTurnRegistry; - /** Primary Redis client (server-owned); undefined in standalone mode. */ - redis?: RedisClient | undefined; - /** Request-reply dispatch table served by this replica's executor. */ - requestReplyRouter: RequestReplyRouter; - /** Hands out each turn's resumable event stream to the create and subscribe handlers. */ - eventSubscriptions: EventSubscriptionRegistry; - logger: Logger; - /** Discovered openid-client configuration; undefined when browser login is disabled. */ - oidcClient: Configuration | undefined; - /** Startup-selected authenticator; middleware is built from this once per app. */ - authenticator: Authenticator; - /** Startup-selected agent authorization policy. */ - authorizer: Authorizer; + sessionStore: ISessionStore; + skillCatalog: SkillCatalog; + tokenStore: IOAuthTokenStore; + webSearchCatalog: WebSearchCatalog; + /** DB singleton, or the env-backed shared store in TrueFoundry mode. */ + webSearchProviderStore: IWebSearchProviderStore; + withTransaction: WithTransaction; } export function createServerApp(deps: ServerDeps) { @@ -333,21 +331,7 @@ export function createServerApp(deps: ServerDeps) { const adminAuthMiddleware = createAdminAuthMiddleware(deps.authenticator); const scheduleExecutionAuthMiddleware = createApiKeyAuthMiddleware(configuration.TRUEFORGE_API_KEY); const authEnabled = getTrueForgeAuthMode() !== TrueForgeAuthMode.Standalone; - const scheduleTurnDeps = { - scheduleStore: deps.scheduleStore, - sessions: deps.sessions, - agentStore: deps.agentStore, - eventSubscriptions: deps.eventSubscriptions, - logger: deps.logger, - resolveModelProviderStore: deps.resolveModelProviderStore, - resolveMcpServerStore: deps.resolveMcpServerStore, - resolveSandboxProviderStore: deps.resolveSandboxProviderStore, - sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - resolveWebSearchProviderStore: deps.resolveWebSearchProviderStore, - activeTurns: deps.activeTurns, - resolveTurnSkillsStore: deps.resolveTurnSkillsStore, - }; - + const resolveHttpAgentStore = (c: Context) => deps.resolveAgentStore(resolveRequestContext(c)); const httpMetrics = createHttpMetrics(); app.use('*', createRequestLogContextMiddleware()); app.use('/api/v1/sessions/:session_id/*', createSessionPathLogContextMiddleware()); @@ -361,10 +345,8 @@ export function createServerApp(deps: ServerDeps) { app.use('*', createClientCertificateMiddleware(deps.logger)); } app.use('*', createRequestBodyLimitMiddleware(configuration.MAX_REQUEST_BODY_BYTES)); - app.get('/healthz', c => c.json({ status: 'ok', version: PACKAGE_VERSION })); app.get('/metrics', httpMetrics.handler); - app.route( '/api/v1/auth', createAuthRouter({ @@ -378,7 +360,7 @@ export function createServerApp(deps: ServerDeps) { withAuth( createCapabilitiesRouter({ sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - resolveWebSearchProviderStore: deps.resolveWebSearchProviderStore, + webSearchProviderStore: deps.webSearchProviderStore, withTransaction: deps.withTransaction, logger: deps.logger, resolveRequestContext, @@ -449,13 +431,13 @@ export function createServerApp(deps: ServerDeps) { '/api/v1/agents', withAuth( createAgentsRouter({ - resolveAgentStore: deps.resolveAgentStore, + resolveAgentStore: resolveHttpAgentStore, resolveModelProviderStore: c => deps.resolveModelProviderStore({ c }), resolveMcpServerStore: c => deps.resolveMcpServerStore({ c }), resolveSkillStore: c => deps.resolveSkillStore({ c }), - resolveSandboxProviderStore: deps.resolveSandboxProviderStore, + sandboxProviderStore: deps.sandboxProviderStore, sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - resolveWebSearchProviderStore: deps.resolveWebSearchProviderStore, + webSearchProviderStore: deps.webSearchProviderStore, withTransaction: deps.withTransaction, resolveRequestContext, authorizer: deps.authorizer, @@ -468,17 +450,32 @@ export function createServerApp(deps: ServerDeps) { withAuth( createSandboxEnvironmentsRouter({ sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - resolveAgentStore: deps.resolveAgentStore, - resolveSandboxProviderStore: deps.resolveSandboxProviderStore, + resolveAgentStore: resolveHttpAgentStore, + sandboxProviderStore: deps.sandboxProviderStore, resolveRequestContext, logger: deps.logger, }), authMiddleware, ), ); + const scheduleTurnExecutionRouterDeps = { + activeTurns: deps.activeTurns, + agentStore: deps.agentStore, + eventSubscriptions: deps.eventSubscriptions, + logger: deps.logger, + resolveMcpServerStore: deps.resolveMcpServerStore, + resolveModelProviderStore: deps.resolveModelProviderStore, + resolveTurnSkillsStore: deps.resolveTurnSkillsStore, + sandboxEnvironmentStore: deps.sandboxEnvironmentStore, + sandboxProviderStore: deps.sandboxProviderStore, + scheduleStore: deps.scheduleStore, + sessions: deps.sessions, + sessionStore: deps.sessionStore, + webSearchProviderStore: deps.webSearchProviderStore, + }; app.route( '/api/internal/schedules', - withAuth(createScheduleExecutionRouter(scheduleTurnDeps), scheduleExecutionAuthMiddleware), + withAuth(createScheduleExecutionRouter(scheduleTurnExecutionRouterDeps), scheduleExecutionAuthMiddleware), ); app.route( '/api/internal/sandbox-environments', @@ -495,11 +492,11 @@ export function createServerApp(deps: ServerDeps) { '/api/v1/schedules', withAuth( createSchedulesRouter({ - ...scheduleTurnDeps, - resolveAgentStore: deps.resolveAgentStore, - withTransaction: deps.withTransaction, - resolveRequestContext, + ...scheduleTurnExecutionRouterDeps, authorizer: deps.authorizer, + resolveAgentStore: resolveHttpAgentStore, + resolveRequestContext, + withTransaction: deps.withTransaction, }), authMiddleware, ), @@ -512,10 +509,10 @@ export function createServerApp(deps: ServerDeps) { resolveMcpServerStore: deps.resolveMcpServerStore, tokenStore: deps.tokenStore, resolveSkillStore: c => deps.resolveSkillStore({ c }), - resolveSandboxProviderStore: deps.resolveSandboxProviderStore, + sandboxProviderStore: deps.sandboxProviderStore, sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - resolveWebSearchProviderStore: deps.resolveWebSearchProviderStore, - resolveAgentStore: deps.resolveAgentStore, + webSearchProviderStore: deps.webSearchProviderStore, + resolveAgentStore: resolveHttpAgentStore, withTransaction: deps.withTransaction, logger: deps.logger, resolveRequestContext, @@ -527,16 +524,16 @@ export function createServerApp(deps: ServerDeps) { '/api/internal/sessions', withAuth( createInternalSessionsRouter({ - sessions: deps.sessions, - resolveModelProviderStore: deps.resolveModelProviderStore, + authorizer: deps.authorizer, + resolveAgentStore: resolveHttpAgentStore, resolveMcpServerStore: deps.resolveMcpServerStore, + resolveModelProviderStore: deps.resolveModelProviderStore, + resolveRequestContext, resolveSkillStore: deps.resolveSkillStore, - resolveAgentStore: deps.resolveAgentStore, - resolveSandboxProviderStore: deps.resolveSandboxProviderStore, sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - resolveWebSearchProviderStore: deps.resolveWebSearchProviderStore, - resolveRequestContext, - authorizer: deps.authorizer, + sandboxProviderStore: deps.sandboxProviderStore, + sessions: deps.sessions, + webSearchProviderStore: deps.webSearchProviderStore, }), authMiddleware, ), @@ -547,7 +544,7 @@ export function createServerApp(deps: ServerDeps) { createInternalMetricsRouter({ sessionMetricsStore: deps.sessionMetricsStore, resolveRequestContext, - resolveAgentStore: deps.resolveAgentStore, + resolveAgentStore: resolveHttpAgentStore, authorizer: deps.authorizer, }), authMiddleware, @@ -558,7 +555,7 @@ export function createServerApp(deps: ServerDeps) { withAuth( createPermissionsRouter({ authorizer: deps.authorizer, - resolveAgentStore: deps.resolveAgentStore, + resolveAgentStore: resolveHttpAgentStore, scheduleStore: deps.scheduleStore, sessionStore: deps.sessionStore, resolveRequestContext, @@ -570,21 +567,21 @@ export function createServerApp(deps: ServerDeps) { '/api/v1/sessions', withAuth( createSessionsRouter({ - sessions: deps.sessions, - sessionStore: deps.sessionStore, activeTurns: deps.activeTurns, - resolveModelProviderStore: deps.resolveModelProviderStore, - resolveMcpServerStore: deps.resolveMcpServerStore, - resolveSkillStore: deps.resolveSkillStore, - resolveAgentStore: deps.resolveAgentStore, - resolveSandboxProviderStore: deps.resolveSandboxProviderStore, - sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - resolveWebSearchProviderStore: deps.resolveWebSearchProviderStore, + authorizer: deps.authorizer, + logger: deps.logger, redis: deps.redis, requestReplyRouter: deps.requestReplyRouter, + resolveAgentStore: resolveHttpAgentStore, + resolveMcpServerStore: deps.resolveMcpServerStore, + resolveModelProviderStore: deps.resolveModelProviderStore, resolveRequestContext, - logger: deps.logger, - authorizer: deps.authorizer, + resolveSkillStore: deps.resolveSkillStore, + sandboxEnvironmentStore: deps.sandboxEnvironmentStore, + sandboxProviderStore: deps.sandboxProviderStore, + sessions: deps.sessions, + sessionStore: deps.sessionStore, + webSearchProviderStore: deps.webSearchProviderStore, }), authMiddleware, ), @@ -593,20 +590,19 @@ export function createServerApp(deps: ServerDeps) { '/api/v1/sessions', withAuth( createTurnsRouter({ - sessions: deps.sessions, - sessionStore: deps.sessionStore, activeTurns: deps.activeTurns, - resolveModelProviderStore: deps.resolveModelProviderStore, - resolveMcpServerStore: deps.resolveMcpServerStore, - resolveTurnSkillsStore: deps.resolveTurnSkillsStore, - resolveAgentStore: deps.resolveAgentStore, + authorizer: deps.authorizer, eventSubscriptions: deps.eventSubscriptions, - resolveSandboxProviderStore: deps.resolveSandboxProviderStore, - sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - resolveWebSearchProviderStore: deps.resolveWebSearchProviderStore, logger: deps.logger, + redis: deps.redis, + requestReplyRouter: deps.requestReplyRouter, + resolveAgentStore: deps.resolveAgentStore, resolveRequestContext, - authorizer: deps.authorizer, + resolveTurnExecutionStores: deps.resolveTurnExecutionStores, + sandboxEnvironmentStore: deps.sandboxEnvironmentStore, + sandboxProviderStore: deps.sandboxProviderStore, + sessions: deps.sessions, + sessionStore: deps.sessionStore, }), authMiddleware, ), diff --git a/packages/trueforge/src/auth/identity.ts b/packages/trueforge/src/auth/identity.ts index 3ff6d9745..2613bfce0 100644 --- a/packages/trueforge/src/auth/identity.ts +++ b/packages/trueforge/src/auth/identity.ts @@ -93,16 +93,25 @@ export function requestSubjectFromCreatedBySubject(subject: CreatedBySubject): R export function requestContextFromCreatedBySubject(params: { tenant_id: string; created_by_subject: CreatedBySubject; + user_credential: string | null; }): RequestContext { const base: RequestContext = { tenant_id: params.tenant_id, subject: requestSubjectFromCreatedBySubject(params.created_by_subject), roles: [], - user_credential: null, + user_credential: params.user_credential, }; return isTrueFoundryModeEnabled(configuration) ? createTrueFoundryRequestContext(base) : base; } +/** + * A context that arrived as JSON has no TrueFoundry token cache. Wrap a fresh one + * so model and MCP stores can vend on this replica. + */ +export function requestContextFromWire(context: RequestContext): RequestContext { + return isTrueFoundryModeEnabled(configuration) ? createTrueFoundryRequestContext(context) : context; +} + /** Persistable creator snapshot derived from the authenticated request. */ export function createdBySubjectFromRequestContext(ctx: RequestContext): CreatedBySubject { return { diff --git a/packages/trueforge/src/config.ts b/packages/trueforge/src/config.ts index 2913f3551..0a4a0cb55 100644 --- a/packages/trueforge/src/config.ts +++ b/packages/trueforge/src/config.ts @@ -761,8 +761,10 @@ export interface SharedServerConfiguration { */ GRACEFUL_TIMEOUT_SECONDS: number; /** - * Max seconds a single turn may execute before it is cancelled with - * `server-execution-timeout`. Env: `SERVER_EXECUTION_TIMEOUT_SECONDS`. Default 600 (10 minutes). + * Max seconds a turn may stay `running` before it is cancelled with + * `server-execution-timeout`. Paused time is not counted: the window clears on + * pause and starts again on the next `running`. Env: `SERVER_EXECUTION_TIMEOUT_SECONDS`. + * Default 3600 (60 minutes). */ SERVER_EXECUTION_TIMEOUT_SECONDS: number; /** @@ -780,6 +782,12 @@ export interface SharedServerConfiguration { * Env: `TURN_SUBSCRIBE_TIMEOUT_MS`. Default 600000. */ TURN_SUBSCRIBE_TIMEOUT_MS: number; + /** + * How long a paused turn stays in memory before it is offloaded. + * The row stays paused. The next send on the owner rebuilds it. + * Env: `TURN_PAUSED_KEEPALIVE_SECONDS`. Default 300. + */ + TURN_PAUSED_KEEPALIVE_SECONDS: 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. @@ -1171,6 +1179,11 @@ const shared: SharedServerConfiguration = { raw: getEnv('TURN_SUBSCRIBE_TIMEOUT_MS'), defaultValue: 600_000, }), + TURN_PAUSED_KEEPALIVE_SECONDS: parsePositiveInt({ + envKey: 'TURN_PAUSED_KEEPALIVE_SECONDS', + raw: getEnv('TURN_PAUSED_KEEPALIVE_SECONDS'), + defaultValue: 300, + }), 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/db/postgres/session-store/queries/inboundEvents.ts b/packages/trueforge/src/db/postgres/session-store/queries/inboundEvents.ts index 79cf6991e..e3a65054c 100644 --- a/packages/trueforge/src/db/postgres/session-store/queries/inboundEvents.ts +++ b/packages/trueforge/src/db/postgres/session-store/queries/inboundEvents.ts @@ -16,7 +16,7 @@ import { toTurnUserEvent, toTurnUserEventMessage } from '../../../turnInboundEve import { isUniqueViolation } from '../../client'; import { json } from '../../sqlExpressions'; import type { Database } from '../../types'; -import { assertTurnNonTerminal } from './turns'; +import { assertOwnedNonTerminalTurn, assertTurnNonTerminal } from './turns'; async function requireSession(db: Kysely, sessionId: string): Promise { const row = await db @@ -67,12 +67,16 @@ export async function insertTurnInboundEvents( }); } - const keys = { session_id: input.session_id, turn_id: input.turn_id }; + const keys = { + session_id: input.session_id, + turn_id: input.turn_id, + active_executor_id: input.active_executor_id, + }; try { - // Ensure the turn does not stop before we insert events. + // The row must still be non-terminal and owned by this executor. await db.transaction().execute(async trx => { - await assertTurnNonTerminal(trx, keys); + await assertOwnedNonTerminalTurn(trx, keys); await trx .insertInto('turn_inbound_events') .values( diff --git a/packages/trueforge/src/db/sqlite/session-store/queries/inboundEvents.ts b/packages/trueforge/src/db/sqlite/session-store/queries/inboundEvents.ts index 31798bb51..fd42d22d1 100644 --- a/packages/trueforge/src/db/sqlite/session-store/queries/inboundEvents.ts +++ b/packages/trueforge/src/db/sqlite/session-store/queries/inboundEvents.ts @@ -15,7 +15,7 @@ import { toTurnUserEvent, toTurnUserEventMessage } from '../../../turnInboundEve import { isUniqueViolation } from '../../client'; import { jsonbBind, jsonText } from '../../sqlExpressions'; import type { Database } from '../../types'; -import { assertTurnNonTerminal } from './turns'; +import { assertOwnedNonTerminalTurn, assertTurnNonTerminal } from './turns'; async function requireSession(db: Kysely, sessionId: string): Promise { const row = await db @@ -66,12 +66,16 @@ export async function insertTurnInboundEvents( }); } - const keys = { session_id: input.session_id, turn_id: input.turn_id }; + const keys = { + session_id: input.session_id, + turn_id: input.turn_id, + active_executor_id: input.active_executor_id, + }; try { - // Ensure the turn does not stop before we insert events. + // The row must still be non-terminal and owned by this executor. await db.transaction().execute(async trx => { - await assertTurnNonTerminal(trx, keys); + await assertOwnedNonTerminalTurn(trx, keys); await trx .insertInto('turn_inbound_events') .values( diff --git a/packages/trueforge/src/main.ts b/packages/trueforge/src/main.ts index 609848285..f92fb95bd 100644 --- a/packages/trueforge/src/main.ts +++ b/packages/trueforge/src/main.ts @@ -14,7 +14,6 @@ import { configureOutboundUrlGuard, extractErrorLogFields, } from '@truefoundry/trueforge-core/core'; -import type { Context } from 'hono'; import { mkdir } from 'node:fs/promises'; import path from 'node:path'; import { @@ -84,19 +83,13 @@ import { type TurnStreamingEvent, } from '@truefoundry/trueforge-core/agent-session'; import { RequestReplyExecutor, RequestReplyRouter, type RedisClient } from '@truefoundry/trueforge-core/request-reply'; -import type { Kysely, Transaction } from 'kysely'; +import type { Transaction } from 'kysely'; import type { Logger } from 'winston'; -import { - createServerApp, - type ResolveMcpServerStoreInput, - type ResolveModelProviderStoreInput, - type ResolveSkillStoreInput, - type ResolveTurnSkillsStoreInput, -} from './app'; +import { createServerApp } from './app'; import { TrueForgeAuthorizer, type Authorizer } from './auth/authorizer'; import { createAuthenticator } from './auth/createAuthenticator'; -import { resolveRequestContext, STANDALONE_REQUEST_CONTEXT, type RequestContext } from './auth/identity'; +import type { RequestContext } from './auth/identity'; import { initOidc } from './auth/oidc'; import { McpCatalog } from './catalog/McpCatalog'; import { ModelCatalog } from './catalog/ModelCatalog'; @@ -105,11 +98,9 @@ import { SkillCatalog } from './catalog/SkillCatalog'; import { WebSearchCatalog } from './catalog/WebSearchCatalog'; import { type DistributedServerConfiguration } from './config'; import { createController } from './controller'; -import type { AgentRecord, IAgentStore } from './db/agentStore'; -import type { IMcpServerStore, IMcpServerWithAuthStore } from './db/mcpServerStore'; +import type { IAgentStore } from './db/agentStore'; +import type { IMcpServerWithAuthStore } from './db/mcpServerStore'; import { McpServerWithAuthStore } from './db/McpServerWithAuthStore'; -import type { IModelProviderStore } from './db/modelProviderStore'; -import type { PostgresAgentStore } from './db/postgres/agent-store/PostgresAgentStore'; import type { Database as PostgresDatabase } from './db/postgres/types'; import type { ISandboxEnvironmentStore } from './db/sandboxEnvironmentStore'; import type { ISandboxProviderStore } from './db/sandboxProviderStore'; @@ -124,88 +115,53 @@ import { serverTlsServeOptions } from './http/tls'; import { createServerLogger, shouldColorize } from './logger'; import type { IOAuthTokenStore } from './mcp/auth/types'; import { PACKAGE_VERSION } from './packageVersion'; +import { + buildResolveAgentStore, + buildResolveMcpServerStore, + buildResolveModelProviderStore, + buildResolveSkillStore, + buildResolveTurnSkillsStore, + createRequestStoreResolvers, + type ResolveMcpServerStore, + type ResolveModelProviderStore, + type ResolveTurnSkillsStore, +} from './resolvers'; import { ActiveTurnRegistry } from './runtime/activeTurns'; import { EventSubscriptionRegistry } from './runtime/event-subscription'; import type { ConnectedRedis } from './runtime/redis'; import { initSentry } from './sentry'; import { printStandaloneStartupBanner } from './startupBanner'; -import { InlineMcpServerStore } from './truefoundry/InlineMcpServerStore'; -import { InlineModelProviderStore } from './truefoundry/InlineModelProviderStore'; -import { - assertInlineResourceHeadersAllowed, - parseInlineMcpServers, - parseInlineModelProviders, - parseInlineSkills, - X_TFG_MCP, - X_TFG_MODELS, - X_TFG_SKILLS, - type InlineMcpServers, - type InlineModelProviders, - type InlineSkills, -} from './truefoundry/inlineResources'; -import { InlineSkillStore } from './truefoundry/InlineSkillStore'; -import { - parsePerServerMcpHeaders, - X_TFG_MCP_HEADERS, - type PerServerMcpHeaders, -} from './truefoundry/perServerMcpHeaders'; -import { TrueFoundryAgentStore } from './truefoundry/TrueFoundryAgentStore'; import { TrueFoundryAuthorizer } from './truefoundry/TrueFoundryAuthorizer'; -import { TrueFoundryMcpServerStore } from './truefoundry/TrueFoundryMcpServerStore'; -import { TrueFoundryModelProviderStore } from './truefoundry/TrueFoundryModelProviderStore'; import { TrueFoundrySandboxEnvironmentStore } from './truefoundry/TrueFoundrySandboxEnvironmentStore'; import { TrueFoundrySandboxProviderStore } from './truefoundry/TrueFoundrySandboxProviderStore'; import { TrueFoundryServiceFoundryServerClient } from './truefoundry/TrueFoundryServiceFoundryServerClient'; -import { TrueFoundryAdminSkillStore, TrueFoundrySkillStore } from './truefoundry/TrueFoundrySkillStore'; import { TrueFoundryWebSearchProviderStore } from './truefoundry/TrueFoundryWebSearchProviderStore'; -/** Persistence-layer model store resolve (RequestContext, not Hono Context). */ -type ResolveModelProviderStore = (input: { - rc: RequestContext; - runAsAgent?: AgentRecord; - inlineModelProviders?: InlineModelProviders; -}) => IModelProviderStore; - -/** Persistence-layer MCP store resolve (RequestContext, not Hono Context). */ -type ResolveMcpServerStore = (input: { - rc: RequestContext; - runAsAgent?: AgentRecord; - perServerHeaders?: PerServerMcpHeaders; - inlineMcpServers?: InlineMcpServers; -}) => IMcpServerWithAuthStore; - -/** Turn skill mounts: closed-set `x-tfg-skills`, else the registry/DB store. */ -type ResolveTurnSkillsStore = (input: { - inlineSkills?: InlineSkills; -}) => Pick, 'resolveTurnSkills'>; - /** Persistence + optional Redis wired for the selected topology. */ interface ServerPersistence { - withTransaction: WithTransaction; - sessionStore: ISessionStore; - sessionMetricsStore: ISessionMetricsStore; - tokenStore: IOAuthTokenStore; - scheduleStore: IScheduleStore; + /** extra pre-resolved stores for scheduled runs */ + agentStore: IAgentStore; + destroyDb: () => Promise; mcpOAuthStore: IMcpServerWithAuthStore; - resolveModelProviderStore: ResolveModelProviderStore; + /** Connected Redis (client + mode) for distributed peering; undefined in standalone. */ + redis: ConnectedRedis | undefined; + resolveAgentStore: (rc: RequestContext) => IAgentStore; resolveMcpServerStore: ResolveMcpServerStore; - resolveSandboxProviderStore: (rc: RequestContext) => ISandboxProviderStore; - resolveWebSearchProviderStore: (rc: RequestContext) => IWebSearchProviderStore; + resolveModelProviderStore: ResolveModelProviderStore; /** Per-request store: DB git skills, or TrueFoundry registry catalog (caller JWT). */ resolveSkillStore: (rc: RequestContext) => ISkillStore; - resolveAgentStore: (rc: RequestContext) => IAgentStore; - /** extra pre-resolved stores for scheduled runs */ - agentStore: IAgentStore; - sandboxEnvironmentStore: ISandboxEnvironmentStore; - /** Context-free provider store for internal sandbox-env build progress. */ - sandboxProviderStore: ISandboxProviderStore; /** Turn skill mounts: registry (API key) or closed-set `x-tfg-skills`. */ resolveTurnSkillsStore: ResolveTurnSkillsStore; - destroyDb: () => Promise; - /** Connected Redis (client + mode) for distributed peering; undefined in standalone. */ - redis: ConnectedRedis | undefined; + sandboxEnvironmentStore: ISandboxEnvironmentStore; + sandboxProviderStore: ISandboxProviderStore; + scheduleStore: IScheduleStore; /** One shared client for TrueFoundry store resolvers + auth; undefined when TrueFoundry mode is off. */ serviceFoundryClient: TrueFoundryServiceFoundryServerClient | undefined; + sessionMetricsStore: ISessionMetricsStore; + sessionStore: ISessionStore; + tokenStore: IOAuthTokenStore; + webSearchProviderStore: IWebSearchProviderStore; + withTransaction: WithTransaction; } /** Shared ServiceFoundry HTTP client when TrueFoundry mode is on; otherwise undefined. */ @@ -229,139 +185,26 @@ function createServiceFoundryServerClient(logger: Logger): TrueFoundryServiceFou }); } -/** Per-request SFY skill catalog; otherwise {@link persistenceStore}. Caller JWT for list/validate. */ -function buildResolveSkillStore(options: { - persistenceStore: ISkillStore; - client: TrueFoundryServiceFoundryServerClient | undefined; -}): (requestContext: RequestContext) => ISkillStore { - const { persistenceStore, client } = options; - if (client === undefined) { - return () => persistenceStore; - } - return requestContext => - new TrueFoundrySkillStore({ - client, - context: requestContext, - }); -} - -/** - * Turn skill mounts. Registry path uses the service API key; `inlineSkills` is a closed set from - * `x-tfg-skills` with no registry lookup. - */ -function buildResolveTurnSkillsStore(options: { - persistenceStore: ISkillStore; - client: TrueFoundryServiceFoundryServerClient | undefined; -}): ResolveTurnSkillsStore { - const { persistenceStore, client } = options; - const registry: Pick, 'resolveTurnSkills'> = client - ? new TrueFoundryAdminSkillStore({ client }) - : persistenceStore; - return input => { - if (input.inlineSkills !== undefined) { - return new InlineSkillStore({ inline: input.inlineSkills }); - } - return registry; - }; -} - -/** - * Per-request model-provider store over a shared ServiceFoundry client, else {@link persistenceStore}. - * `runAsAgent` is set only when a turn should use a saved agent's token. - */ -function buildResolveModelProviderStore(options: { - persistenceStore: IModelProviderStore; - client: TrueFoundryServiceFoundryServerClient | undefined; - logger: Logger; -}): ResolveModelProviderStore { - return input => { - const { persistenceStore, client } = options; - if (client) { - const store = new TrueFoundryModelProviderStore({ - client, - requestContext: input.rc, - agent: input.runAsAgent, - logger: options.logger, - }); - if (input.inlineModelProviders !== undefined) { - return new InlineModelProviderStore({ inner: store, inline: input.inlineModelProviders }); - } - return store; - } - return persistenceStore; - }; -} - -/** - * Per-request MCP store over ServiceFoundry. `runAsAgent` selects a saved agent's token for turns. - */ -function buildResolveMcpServerStore(options: { - persistenceStore: IMcpServerStore; - tokenStore: IOAuthTokenStore; - client: TrueFoundryServiceFoundryServerClient | undefined; - logger: Logger; -}): ResolveMcpServerStore { - const withAuthMcpPersistenceStore = new McpServerWithAuthStore({ - store: options.persistenceStore, - tokenStore: options.tokenStore, - clientName: configuration.MCP_DCR_OAUTH_CLIENT_NAME, - }); - return input => { - const { client } = options; - if (client) { - const store = new TrueFoundryMcpServerStore({ - client, - requestContext: input.rc, - agent: input.runAsAgent, - logger: options.logger, - ...(input.perServerHeaders === undefined ? {} : { perServerHeaders: input.perServerHeaders }), - }); - if (input.inlineMcpServers !== undefined) { - return new InlineMcpServerStore({ inner: store, inline: input.inlineMcpServers }); - } - return store; - } - return withAuthMcpPersistenceStore; - }; -} - -/** Per-request agent store decorator over DB persistence, else {@link persistenceStore}. */ -function buildResolveAgentStore(options: { - persistenceStore: PostgresAgentStore; - db: Kysely; - client: TrueFoundryServiceFoundryServerClient | undefined; -}): (rc: RequestContext) => IAgentStore> { - const { persistenceStore, client, db } = options; - return rc => { - if (client) { - return new TrueFoundryAgentStore({ - inner: persistenceStore, - client, - context: rc, - db, - }); - } - return persistenceStore; - }; -} - /** * Sandbox-provider store resolver. In TrueFoundry mode the shared env-backed store is reused; * otherwise the persistence store is reused as-is. */ -function buildResolveSandboxProviderStore(options: { - persistenceStore: ISandboxProviderStore; -}): (rc: RequestContext) => ISandboxProviderStore { +function buildSandboxProviderStore(options: { persistenceStore: ISandboxProviderStore }): ISandboxProviderStore { const { persistenceStore } = options; if (isTrueFoundryModeEnabled(configuration)) { - return () => new TrueFoundrySandboxProviderStore(); + return new TrueFoundrySandboxProviderStore(); } - return () => persistenceStore; + return persistenceStore; } -function wrapSandboxEnvironmentStore( - persistenceStore: ISandboxEnvironmentStore, -): ISandboxEnvironmentStore { +/** + * Sandbox environment store resolver. In TrueFoundry mode the shared env-backed store is reused; + * otherwise the persistence store is reused as-is. + */ +function buildSandboxEnvironmentStore(options: { + persistenceStore: ISandboxEnvironmentStore; +}): ISandboxEnvironmentStore { + const { persistenceStore } = options; if (isTrueFoundryModeEnabled(configuration)) { // Discriminator only — full resolve validates secrets/settings used elsewhere. return new TrueFoundrySandboxEnvironmentStore(persistenceStore, { @@ -371,14 +214,16 @@ function wrapSandboxEnvironmentStore( return persistenceStore; } -function buildResolveWebSearchProviderStore(options: { - persistenceStore: IWebSearchProviderStore; -}): (rc: RequestContext) => IWebSearchProviderStore { +/** + * Web search provider store resolver. In TrueFoundry mode the shared web search provider store is reused; + * otherwise the persistence store is reused as-is. + */ +function buildWebSearchProviderStore(options: { persistenceStore: IWebSearchProviderStore }): IWebSearchProviderStore { const { persistenceStore } = options; if (isTrueFoundryModeEnabled(configuration)) { - return () => new TrueFoundryWebSearchProviderStore(); + return new TrueFoundryWebSearchProviderStore(); } - return () => persistenceStore; + return persistenceStore; } /** SQLite stores; Redis unused (executor peering disabled). */ @@ -426,39 +271,44 @@ async function createStandalonePersistence(options: { const tokenStore = new SqliteOAuthTokenStore(db); const agentStore = new SqliteAgentStore(db); - const sandboxEnvironmentStore = wrapSandboxEnvironmentStore(new SqliteSandboxEnvironmentStore(db)); const modelProviderStore = new SqliteModelProviderStore(db); const mcpServerStore = new McpServerWithAuthStore({ store: new SqliteMcpServerStore(db), tokenStore, clientName: configuration.MCP_DCR_OAUTH_CLIENT_NAME, }); - const sandboxProviderStore = new SqliteSandboxProviderStore(db); - const webSearchProviderStore = new SqliteWebSearchProviderStore(db); const skillStore = new SqliteSkillStore(db); + const sandboxProviderStore = buildSandboxProviderStore({ + persistenceStore: new SqliteSandboxProviderStore(db), + }); + const sandboxEnvironmentStore = buildSandboxEnvironmentStore({ + persistenceStore: new SqliteSandboxEnvironmentStore(db), + }); + const webSearchProviderStore = buildWebSearchProviderStore({ + persistenceStore: new SqliteWebSearchProviderStore(db), + }); return { - withTransaction: callback => db.transaction().execute(callback), - sessionStore: new SqliteSessionStore(db), - sessionMetricsStore: new SqliteSessionMetricsStore(db), + agentStore, + destroyDb: () => db.destroy(), mcpOAuthStore: mcpServerStore, - tokenStore, - scheduleStore: new SqliteScheduleStore(db), - resolveModelProviderStore: () => modelProviderStore, + redis: undefined, + resolveAgentStore: () => agentStore, resolveMcpServerStore: () => mcpServerStore, - resolveSandboxProviderStore: () => sandboxProviderStore, - resolveWebSearchProviderStore: () => webSearchProviderStore, + resolveModelProviderStore: () => modelProviderStore, resolveSkillStore: () => skillStore, - resolveAgentStore: () => agentStore, - agentStore, - sandboxEnvironmentStore, - sandboxProviderStore, resolveTurnSkillsStore: buildResolveTurnSkillsStore({ persistenceStore: skillStore, client: undefined, }), - destroyDb: () => db.destroy(), - redis: undefined, + sandboxEnvironmentStore, + sandboxProviderStore, + scheduleStore: new SqliteScheduleStore(db), serviceFoundryClient: undefined, + sessionMetricsStore: new SqliteSessionMetricsStore(db), + sessionStore: new SqliteSessionStore(db), + tokenStore, + webSearchProviderStore, + withTransaction: callback => db.transaction().execute(callback), }; } @@ -537,6 +387,7 @@ async function createDistributedPersistence(options: { } const serviceFoundryClient = createServiceFoundryServerClient(logger); const tokenStore = new PostgresOAuthTokenStore(db); + const agentStore = new PostgresAgentStore(db); const modelProviderStore = new PostgresModelProviderStore(db); const mcpServerStore = new PostgresMcpServerStore(db); const mcpServerWithAuthStore = new McpServerWithAuthStore({ @@ -544,14 +395,15 @@ async function createDistributedPersistence(options: { tokenStore, clientName: configuration.MCP_DCR_OAUTH_CLIENT_NAME, }); - const sandboxProviderStore = new PostgresSandboxProviderStore(db); - const webSearchProviderStore = new PostgresWebSearchProviderStore(db); const skillStore = new PostgresSkillStore(db); - const agentStore = new PostgresAgentStore(db); - const sandboxEnvironmentStore = wrapSandboxEnvironmentStore(new PostgresSandboxEnvironmentStore(db)); - const resolveTurnSkillsStore = buildResolveTurnSkillsStore({ - persistenceStore: skillStore, - client: serviceFoundryClient, + const sandboxProviderStore = buildSandboxProviderStore({ + persistenceStore: new PostgresSandboxProviderStore(db), + }); + const sandboxEnvironmentStore = buildSandboxEnvironmentStore({ + persistenceStore: new PostgresSandboxEnvironmentStore(db), + }); + const webSearchProviderStore = buildWebSearchProviderStore({ + persistenceStore: new PostgresWebSearchProviderStore(db), }); const resolveModelProviderStore = buildResolveModelProviderStore({ persistenceStore: modelProviderStore, @@ -564,39 +416,23 @@ async function createDistributedPersistence(options: { client: serviceFoundryClient, logger, }); + const resolveSkillStore = buildResolveSkillStore({ + persistenceStore: skillStore, + client: serviceFoundryClient, + }); const resolveAgentStore = buildResolveAgentStore({ persistenceStore: agentStore, db, client: serviceFoundryClient, }); - const resolveSandboxProviderStore = buildResolveSandboxProviderStore({ - persistenceStore: sandboxProviderStore, - }); - const resolveWebSearchProviderStore = buildResolveWebSearchProviderStore({ - persistenceStore: webSearchProviderStore, - }); - const resolveSkillStore = buildResolveSkillStore({ + const resolveTurnSkillsStore = buildResolveTurnSkillsStore({ persistenceStore: skillStore, client: serviceFoundryClient, }); return { - withTransaction: callback => db.transaction().execute(callback), - sessionStore: new PostgresSessionStore(db), - sessionMetricsStore: new PostgresSessionMetricsStore(db), - mcpOAuthStore: mcpServerWithAuthStore, - tokenStore, - scheduleStore: new PostgresScheduleStore(db), - resolveModelProviderStore, - resolveMcpServerStore, - resolveSandboxProviderStore, - resolveWebSearchProviderStore, - resolveSkillStore, - resolveAgentStore, agentStore, - sandboxEnvironmentStore, - sandboxProviderStore: resolveSandboxProviderStore(STANDALONE_REQUEST_CONTEXT), - resolveTurnSkillsStore, destroyDb: () => db.destroy(), + mcpOAuthStore: mcpServerWithAuthStore, redis: await connectRedis({ connection: redisConnection, connectTimeoutMs: redisConnectTimeoutMs, @@ -612,28 +448,45 @@ async function createDistributedPersistence(options: { keyPassphrase: redisTlsKeyPassphrase, }, }), + resolveAgentStore, + resolveMcpServerStore, + resolveModelProviderStore, + resolveSkillStore, + resolveTurnSkillsStore, + sandboxEnvironmentStore, + sandboxProviderStore, + scheduleStore: new PostgresScheduleStore(db), serviceFoundryClient, + sessionMetricsStore: new PostgresSessionMetricsStore(db), + sessionStore: new PostgresSessionStore(db), + tokenStore, + webSearchProviderStore, + withTransaction: callback => db.transaction().execute(callback), }; } /** Keeps `TTransaction` concrete when wiring a single persistence topology into the app. */ async function createServerRuntime(persistence: ServerPersistence, logger: Logger) { const { - withTransaction, - sessionStore, - sessionMetricsStore, - tokenStore, - scheduleStore, - mcpOAuthStore, agentStore, - sandboxEnvironmentStore, - resolveTurnSkillsStore, destroyDb, + mcpOAuthStore, redis, + sandboxEnvironmentStore, + sandboxProviderStore, + scheduleStore, serviceFoundryClient, + sessionMetricsStore, + sessionStore, + tokenStore, + webSearchProviderStore, + withTransaction, } = persistence; - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ + pausedKeepaliveMs: configuration.TURN_PAUSED_KEEPALIVE_SECONDS * 1000, + logger, + }); const requestReplyRouter = new RequestReplyRouter(); const eventSubscriptions = new EventSubscriptionRegistry(redis?.client); const sessions = new Sessions({ sessionStore }); @@ -676,120 +529,53 @@ async function createServerRuntime(persistence: ServerPersistence< }) : undefined; - // Hono handlers get Context; persistence resolvers take RequestContext. - // x-tfg-skills / x-tfg-mcp / x-tfg-models are rejected unless allowInlineResources - // (inline session create or spec replace, or a turn for an inline session). - // x-tfg-mcp-headers is separate (per-server invoke headers). - const isTrueFoundryMode = isTrueFoundryModeEnabled(configuration); - const resolveModelProviderStore = (input: ResolveModelProviderStoreInput) => { - const { c, runAsAgent, allowInlineResources = false } = input; - const rawInlineModelsHeader = c.req.header(X_TFG_MODELS); - if (rawInlineModelsHeader !== undefined) { - assertInlineResourceHeadersAllowed({ - headerNames: [X_TFG_MODELS], - isTrueFoundryMode, - allowInlineResources, - }); - } - return persistence.resolveModelProviderStore({ - rc: resolveRequestContext(c), - ...(runAsAgent !== undefined ? { runAsAgent } : {}), - ...(rawInlineModelsHeader !== undefined - ? { inlineModelProviders: parseInlineModelProviders(rawInlineModelsHeader) } - : {}), - }); - }; - const resolveMcpServerStore = (input: ResolveMcpServerStoreInput) => { - const { c, runAsAgent, allowInlineResources = false } = input; - if (c === undefined) { - return mcpOAuthStore; - } - const rawPerServerHeaders = c.req.header(X_TFG_MCP_HEADERS); - const rawInlineMcpServersHeader = c.req.header(X_TFG_MCP); - if (rawInlineMcpServersHeader !== undefined) { - assertInlineResourceHeadersAllowed({ - headerNames: [X_TFG_MCP], - isTrueFoundryMode, - allowInlineResources, - }); - } - return persistence.resolveMcpServerStore({ - rc: resolveRequestContext(c), - ...(runAsAgent !== undefined ? { runAsAgent } : {}), - ...(rawPerServerHeaders !== undefined ? { perServerHeaders: parsePerServerMcpHeaders(rawPerServerHeaders) } : {}), - ...(rawInlineMcpServersHeader !== undefined - ? { inlineMcpServers: parseInlineMcpServers(rawInlineMcpServersHeader) } - : {}), - }); - }; - const resolveAgentStore = (c: Context) => persistence.resolveAgentStore(resolveRequestContext(c)); - const resolveSandboxProviderStore = (c: Context) => persistence.resolveSandboxProviderStore(resolveRequestContext(c)); - const resolveWebSearchProviderStore = (c: Context) => - persistence.resolveWebSearchProviderStore(resolveRequestContext(c)); - // Catalog and agent saves reject x-tfg-skills. Inline session create and spec replace opt in and get a - // closed-set store for validation only; the header is not persisted. - const resolveSkillStore = (input: ResolveSkillStoreInput) => { - const { c, allowInlineResources = false } = input; - const rawInlineSkillsHeader = c.req.header(X_TFG_SKILLS); - if (rawInlineSkillsHeader !== undefined) { - assertInlineResourceHeadersAllowed({ - headerNames: [X_TFG_SKILLS], - isTrueFoundryMode, - allowInlineResources, - }); - return new InlineSkillStore({ inline: parseInlineSkills(rawInlineSkillsHeader) }); - } - return persistence.resolveSkillStore(resolveRequestContext(c)); - }; - // Turn mounts: registry (API key), or closed-set header skills when `c` + allowInlineResources. - // Schedules omit `c` and always get the registry/DB store. - const resolveTurnSkillsStoreForRequest = (input: ResolveTurnSkillsStoreInput) => { - const { c, allowInlineResources = false } = input; - if (c === undefined) { - return resolveTurnSkillsStore({}); - } - const rawInlineSkillsHeader = c.req.header(X_TFG_SKILLS); - if (rawInlineSkillsHeader !== undefined) { - assertInlineResourceHeadersAllowed({ - headerNames: [X_TFG_SKILLS], - isTrueFoundryMode, - allowInlineResources, - }); - } - return resolveTurnSkillsStore( - rawInlineSkillsHeader === undefined ? {} : { inlineSkills: parseInlineSkills(rawInlineSkillsHeader) }, - ); - }; + const { + resolveMcpServerStore, + resolveModelProviderStore, + resolveSkillStore, + resolveTurnExecutionStores, + resolveTurnSkillsStore, + } = createRequestStoreResolvers({ + isTrueFoundryMode: isTrueFoundryModeEnabled(configuration), + mcpOAuthStore, + resolveAgentStore: persistence.resolveAgentStore, + resolveMcpServerStore: persistence.resolveMcpServerStore, + resolveModelProviderStore: persistence.resolveModelProviderStore, + resolveSkillStore: persistence.resolveSkillStore, + resolveTurnSkillsStore: persistence.resolveTurnSkillsStore, + sandboxProviderStore, + webSearchProviderStore, + }); const app = createServerApp({ - modelCatalog: ModelCatalog.load(), + activeTurns, + agentStore, + authenticator, + authorizer, + eventSubscriptions, + logger, mcpCatalog: McpCatalog.load(), - skillCatalog: SkillCatalog.load(), - sandboxCatalog: SandboxCatalog.load(), - webSearchCatalog: WebSearchCatalog.load(), - resolveModelProviderStore, + modelCatalog: ModelCatalog.load(), + oidcClient, + redis: redis?.client, + requestReplyRouter, + resolveAgentStore: persistence.resolveAgentStore, resolveMcpServerStore, - resolveAgentStore, - resolveSandboxProviderStore, - resolveWebSearchProviderStore, + resolveModelProviderStore, resolveSkillStore, - resolveTurnSkillsStore: resolveTurnSkillsStoreForRequest, - withTransaction, - tokenStore, - scheduleStore, - agentStore, + resolveTurnExecutionStores, + resolveTurnSkillsStore, + sandboxCatalog: SandboxCatalog.load(), sandboxEnvironmentStore, - sandboxProviderStore: persistence.sandboxProviderStore, - sessionStore, + sandboxProviderStore, + scheduleStore, sessionMetricsStore, sessions, - activeTurns, - redis: redis?.client, - requestReplyRouter, - eventSubscriptions, - logger, - oidcClient, - authenticator, - authorizer, + sessionStore, + skillCatalog: SkillCatalog.load(), + tokenStore, + webSearchCatalog: WebSearchCatalog.load(), + webSearchProviderStore, + withTransaction, }); return { activeTurns, app, controller, destroyDb, redis, requestReplyRouter }; diff --git a/packages/trueforge/src/resolvers.ts b/packages/trueforge/src/resolvers.ts new file mode 100644 index 000000000..58b28ce1d --- /dev/null +++ b/packages/trueforge/src/resolvers.ts @@ -0,0 +1,389 @@ +import type { Kysely, Transaction } from 'kysely'; +import type { Logger } from 'winston'; +import type { ResolveTurnExecutionStoresInput, TurnExecutionStores } from './apis/turns'; +import { + type ResolveMcpServerStoreInput, + type ResolveModelProviderStoreInput, + type ResolveSkillStoreInput, + type ResolveTurnSkillsStoreInput, +} from './app'; +import { resolveRequestContext, type RequestContext } from './auth/identity'; +import configuration from './config'; +import type { AgentRecord, IAgentStore } from './db/agentStore'; +import type { IMcpServerStore, IMcpServerWithAuthStore } from './db/mcpServerStore'; +import { McpServerWithAuthStore } from './db/McpServerWithAuthStore'; +import type { IModelProviderStore } from './db/modelProviderStore'; +import type { PostgresAgentStore } from './db/postgres/agent-store/PostgresAgentStore'; +import type { Database as PostgresDatabase } from './db/postgres/types'; +import type { ISandboxProviderStore } from './db/sandboxProviderStore'; +import type { ISkillStore } from './db/skillStore'; +import type { IWebSearchProviderStore } from './db/webSearchProviderStore'; +import type { IOAuthTokenStore } from './mcp/auth/types'; +import { headerValue } from './truefoundry/gatewayMetadata'; +import { InlineMcpServerStore } from './truefoundry/InlineMcpServerStore'; +import { InlineModelProviderStore } from './truefoundry/InlineModelProviderStore'; +import { + assertInlineResourceHeadersAllowed, + parseInlineMcpServers, + parseInlineModelProviders, + parseInlineSkills, + X_TFG_MCP, + X_TFG_MODELS, + X_TFG_SKILLS, + type InlineMcpServers, + type InlineModelProviders, + type InlineSkills, +} from './truefoundry/inlineResources'; +import { InlineSkillStore } from './truefoundry/InlineSkillStore'; +import { + parsePerServerMcpHeaders, + X_TFG_MCP_HEADERS, + type PerServerMcpHeaders, +} from './truefoundry/perServerMcpHeaders'; +import { TrueFoundryAgentStore } from './truefoundry/TrueFoundryAgentStore'; +import { TrueFoundryMcpServerStore } from './truefoundry/TrueFoundryMcpServerStore'; +import { TrueFoundryModelProviderStore } from './truefoundry/TrueFoundryModelProviderStore'; +import type { TrueFoundryServiceFoundryServerClient } from './truefoundry/TrueFoundryServiceFoundryServerClient'; +import { TrueFoundryAdminSkillStore, TrueFoundrySkillStore } from './truefoundry/TrueFoundrySkillStore'; + +/** ------------------------------------------------------------------------ + * Persistence-layer model store resolve (RequestContext, not Hono Context) + * ------------------------------------------------------------------------- + */ +export type ResolveModelProviderStore = (input: { + rc: RequestContext; + runAsAgent?: AgentRecord; + inlineModelProviders?: InlineModelProviders; +}) => IModelProviderStore; + +/** Persistence-layer MCP store resolve (RequestContext, not Hono Context). */ +export type ResolveMcpServerStore = (input: { + rc: RequestContext; + runAsAgent?: AgentRecord; + perServerHeaders?: PerServerMcpHeaders; + inlineMcpServers?: InlineMcpServers; +}) => IMcpServerWithAuthStore; + +/** Turn skill mounts: closed-set `x-tfg-skills`, else the registry/DB store. */ +export type ResolveSkillStore = (rc: RequestContext) => ISkillStore; + +/** Turn skill mounts: closed-set `x-tfg-skills`, else the registry/DB store. */ +export type ResolveTurnSkillsStore = (input: { + inlineSkills?: InlineSkills; +}) => Pick, 'resolveTurnSkills'>; + +/** Per-request agent store decorator over DB persistence, else the persistence store. */ +export type ResolveAgentStore = (rc: RequestContext) => IAgentStore; + +/** + * Per-request model-provider store over a shared ServiceFoundry client, else the persistence store. + * `runAsAgent` is set only when a turn should use a saved agent's token. + */ +export function buildResolveModelProviderStore(options: { + persistenceStore: IModelProviderStore; + client: TrueFoundryServiceFoundryServerClient | undefined; + logger: Logger; +}): ResolveModelProviderStore { + return input => { + const { persistenceStore, client } = options; + if (client) { + const store = new TrueFoundryModelProviderStore({ + client, + requestContext: input.rc, + agent: input.runAsAgent, + logger: options.logger, + }); + if (input.inlineModelProviders) { + return new InlineModelProviderStore({ inner: store, inline: input.inlineModelProviders }); + } + return store; + } + return persistenceStore; + }; +} + +/** + * Per-request MCP store over ServiceFoundry. `runAsAgent` selects a saved agent's token for turns. + */ +export function buildResolveMcpServerStore(options: { + persistenceStore: IMcpServerStore; + tokenStore: IOAuthTokenStore; + client: TrueFoundryServiceFoundryServerClient | undefined; + logger: Logger; +}): ResolveMcpServerStore { + const withAuthMcpPersistenceStore = new McpServerWithAuthStore({ + store: options.persistenceStore, + tokenStore: options.tokenStore, + clientName: configuration.MCP_DCR_OAUTH_CLIENT_NAME, + }); + return input => { + const { client } = options; + if (client) { + const store = new TrueFoundryMcpServerStore({ + client, + requestContext: input.rc, + agent: input.runAsAgent, + logger: options.logger, + ...(input.perServerHeaders === undefined ? {} : { perServerHeaders: input.perServerHeaders }), + }); + if (input.inlineMcpServers) { + return new InlineMcpServerStore({ inner: store, inline: input.inlineMcpServers }); + } + return store; + } + return withAuthMcpPersistenceStore; + }; +} + +/** Per-request SFY skill catalog; otherwise the persistence store. Caller JWT for list/validate. */ +export function buildResolveSkillStore(options: { + persistenceStore: ISkillStore; + client: TrueFoundryServiceFoundryServerClient | undefined; +}): ResolveSkillStore { + const { persistenceStore, client } = options; + if (client) { + return requestContext => + new TrueFoundrySkillStore({ + client, + requestContext, + }); + } + return () => persistenceStore; +} + +/** + * Turn skill mounts. Registry path uses the service API key; `inlineSkills` is a closed set from + * `x-tfg-skills` with no registry lookup. + */ +export function buildResolveTurnSkillsStore(options: { + persistenceStore: ISkillStore; + client: TrueFoundryServiceFoundryServerClient | undefined; +}): ResolveTurnSkillsStore { + const { persistenceStore, client } = options; + return input => { + if (input.inlineSkills !== undefined) { + return new InlineSkillStore({ inline: input.inlineSkills }); + } + if (client) { + return new TrueFoundryAdminSkillStore({ client }); + } + return persistenceStore; + }; +} + +/** Per-request agent store decorator over DB persistence, else the persistence store. */ +export function buildResolveAgentStore(options: { + persistenceStore: PostgresAgentStore; + db: Kysely; + client: TrueFoundryServiceFoundryServerClient | undefined; +}): ResolveAgentStore> { + const { persistenceStore, client, db } = options; + return requestContext => { + if (client) { + return new TrueFoundryAgentStore({ + inner: persistenceStore, + client, + requestContext, + db, + }); + } + return persistenceStore; + }; +} + +/** Catalog skill store: caller JWT, plus a closed `x-tfg-skills` set when inline resources are allowed. */ +interface SkillStoreForInput { + requestContext: RequestContext; + headers?: Record | undefined; + allowInlineResources?: boolean; +} + +/** Turn skill mounts. No request context: registry, or a closed `x-tfg-skills` set. */ +interface TurnSkillStoreForInput { + headers?: Record | undefined; + allowInlineResources?: boolean; +} + +/** + * HTTP and peer entrypoints share these. HTTP passes Hono `Context`; a peered send passes the + * forwarded request context and headers. + */ +export function createRequestStoreResolvers(input: { + isTrueFoundryMode: boolean; + mcpOAuthStore: IMcpServerWithAuthStore; + resolveAgentStore: ResolveAgentStore; + resolveMcpServerStore: ResolveMcpServerStore; + resolveModelProviderStore: ResolveModelProviderStore; + resolveSkillStore: ResolveSkillStore; + resolveTurnSkillsStore: ResolveTurnSkillsStore; + sandboxProviderStore: ISandboxProviderStore; + webSearchProviderStore: IWebSearchProviderStore; +}): { + resolveMcpServerStore: (mcpInput: ResolveMcpServerStoreInput) => IMcpServerWithAuthStore; + resolveModelProviderStore: (modelInput: ResolveModelProviderStoreInput) => IModelProviderStore; + resolveSkillStore: (skillInput: ResolveSkillStoreInput) => ISkillStore; + resolveTurnExecutionStores: ( + turnInput: ResolveTurnExecutionStoresInput, + ) => Promise>; + resolveTurnSkillsStore: ( + skillsInput: ResolveTurnSkillsStoreInput, + ) => Pick, 'resolveTurnSkills'>; +} { + const { isTrueFoundryMode, mcpOAuthStore, sandboxProviderStore, webSearchProviderStore } = input; + + // x-tfg-skills / x-tfg-mcp / x-tfg-models are rejected unless allowInlineResources + // (inline session create or spec replace, or a turn for an inline session). + // x-tfg-mcp-headers is separate (per-server invoke headers). + const inlineHeader = (headerInput: { + headers: Record | undefined; + name: string; + allowInlineResources: boolean; + }): string | undefined => { + const raw = headerValue(headerInput.headers, headerInput.name); + if (raw === undefined) { + return undefined; + } + assertInlineResourceHeadersAllowed({ + headerNames: [headerInput.name], + isTrueFoundryMode, + allowInlineResources: headerInput.allowInlineResources, + }); + return raw; + }; + const modelProviderStoreFor = (storeInput: { + requestContext: RequestContext; + headers?: Record | undefined; + allowInlineResources?: boolean; + runAsAgent?: AgentRecord; + }) => { + const { requestContext, headers, allowInlineResources, runAsAgent } = storeInput; + const rawInlineModelsHeader = inlineHeader({ + headers, + name: X_TFG_MODELS, + allowInlineResources: allowInlineResources ?? false, + }); + return input.resolveModelProviderStore({ + rc: requestContext, + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + ...(rawInlineModelsHeader !== undefined + ? { inlineModelProviders: parseInlineModelProviders(rawInlineModelsHeader) } + : {}), + }); + }; + const mcpServerStoreFor = (storeInput: { + requestContext: RequestContext; + headers?: Record | undefined; + allowInlineResources?: boolean; + runAsAgent?: AgentRecord; + }) => { + const { requestContext, headers, allowInlineResources, runAsAgent } = storeInput; + const rawInlineMcpHeader = inlineHeader({ + headers, + name: X_TFG_MCP, + allowInlineResources: allowInlineResources ?? false, + }); + const rawPerServerHeaders = headerValue(headers, X_TFG_MCP_HEADERS); + return input.resolveMcpServerStore({ + rc: requestContext, + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + ...(rawPerServerHeaders !== undefined ? { perServerHeaders: parsePerServerMcpHeaders(rawPerServerHeaders) } : {}), + ...(rawInlineMcpHeader !== undefined ? { inlineMcpServers: parseInlineMcpServers(rawInlineMcpHeader) } : {}), + }); + }; + const skillStoreFor = (storeInput: SkillStoreForInput) => { + const { requestContext, headers, allowInlineResources } = storeInput; + const rawInlineSkillsHeader = inlineHeader({ + headers, + name: X_TFG_SKILLS, + allowInlineResources: allowInlineResources ?? false, + }); + if (rawInlineSkillsHeader === undefined) { + return input.resolveSkillStore(requestContext); + } + return new InlineSkillStore({ inline: parseInlineSkills(rawInlineSkillsHeader) }); + }; + const turnSkillStoreFor = (storeInput: TurnSkillStoreForInput) => { + const { headers, allowInlineResources } = storeInput; + const rawInlineSkillsHeader = inlineHeader({ + headers, + name: X_TFG_SKILLS, + allowInlineResources: allowInlineResources ?? false, + }); + return input.resolveTurnSkillsStore( + rawInlineSkillsHeader === undefined ? {} : { inlineSkills: parseInlineSkills(rawInlineSkillsHeader) }, + ); + }; + const resolveModelProviderStore = (modelInput: ResolveModelProviderStoreInput) => { + const { c, allowInlineResources, runAsAgent } = modelInput; + return modelProviderStoreFor({ + requestContext: resolveRequestContext(c), + headers: c.req.header(), + ...(allowInlineResources !== undefined ? { allowInlineResources } : {}), + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + }); + }; + const resolveMcpServerStore = (mcpInput: ResolveMcpServerStoreInput) => { + const { c, allowInlineResources, runAsAgent } = mcpInput; + if (c === undefined) { + return mcpOAuthStore; + } + return mcpServerStoreFor({ + requestContext: resolveRequestContext(c), + headers: c.req.header(), + ...(allowInlineResources !== undefined ? { allowInlineResources } : {}), + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + }); + }; + // Catalog and agent saves reject x-tfg-skills. Inline session create and spec replace opt in and get a + // closed-set store for validation only; the header is not persisted. + const resolveSkillStore = (skillInput: ResolveSkillStoreInput) => { + const { c, allowInlineResources } = skillInput; + return skillStoreFor({ + requestContext: resolveRequestContext(c), + headers: c.req.header(), + ...(allowInlineResources !== undefined ? { allowInlineResources } : {}), + }); + }; + // Turn mounts: registry (API key), or closed-set header skills when `c` + allowInlineResources. + // Schedules omit `c` and always get the registry/DB store. + const resolveTurnSkillsStore = (skillsInput: ResolveTurnSkillsStoreInput) => { + const { c, allowInlineResources } = skillsInput; + if (c === undefined) { + return input.resolveTurnSkillsStore({}); + } + return turnSkillStoreFor({ + headers: c.req.header(), + ...(allowInlineResources !== undefined ? { allowInlineResources } : {}), + }); + }; + const resolveTurnExecutionStores = (turnInput: ResolveTurnExecutionStoresInput) => { + const { runAsAgent } = turnInput; + const allowInlineResources = turnInput.session.record.agent.type === 'inline'; + return Promise.resolve({ + agentStore: input.resolveAgentStore(turnInput.requestContext), + modelProviderStore: modelProviderStoreFor({ + requestContext: turnInput.requestContext, + headers: turnInput.requestHeaders, + allowInlineResources, + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + }), + mcpServerStore: mcpServerStoreFor({ + requestContext: turnInput.requestContext, + headers: turnInput.requestHeaders, + allowInlineResources, + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + }), + skillStore: turnSkillStoreFor({ headers: turnInput.requestHeaders, allowInlineResources }), + sandboxProviderStore, + webSearchProviderStore, + }); + }; + + return { + resolveMcpServerStore, + resolveModelProviderStore, + resolveSkillStore, + resolveTurnExecutionStores, + resolveTurnSkillsStore, + }; +} diff --git a/packages/trueforge/src/routes/turnRoutes.ts b/packages/trueforge/src/routes/turnRoutes.ts index a4fa5d827..a9a1a9b05 100644 --- a/packages/trueforge/src/routes/turnRoutes.ts +++ b/packages/trueforge/src/routes/turnRoutes.ts @@ -207,7 +207,11 @@ export const postTurnEventsRoute = createRoute({ }, 409: { content: { 'application/json': { schema: RequestErrorResponseSchema } }, - description: 'Turn is terminal, or an event id already exists.', + description: 'Turn is terminal, running and not accepting events, or an event id already exists.', + }, + 500: { + content: { 'application/json': { schema: RequestErrorResponseSchema } }, + description: 'The executor ran the request and failed while applying the events.', }, }, }); diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 1586a5fe2..d8ff64680 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -4,26 +4,69 @@ * TurnHandle.stream() writes the terminal state when the signal fires. * `track()` owns registration and cleanup around the stream lifecycle. */ -import { CancellationReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; +import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; import { Mutex } from 'async-mutex'; +import type { Logger } from 'winston'; -interface ActiveTurnRun { - turn: TurnHandle; +const DEFAULT_PAUSED_KEEPALIVE_MS = 300_000; + +interface TurnRunCompletion { abortController: AbortController; waitUntilCompleted: Promise; markCompleted: () => void; } +/** Paused rebuild before `rebuildTurn` returns a handle. Shutdown offloads this phase. */ +interface RebuildingTurnRun extends TurnRunCompletion { + phase: 'rebuilding'; +} + +interface LiveTurnRun extends TurnRunCompletion { + phase: 'live'; + turn: TurnHandle; +} + +type ActiveTurnRun = RebuildingTurnRun | LiveTurnRun; + +function setupTurnRun(abortController: AbortController): TurnRunCompletion { + const { promise: waitUntilCompleted, resolve } = Promise.withResolvers(); + const markCompleted = (): void => { + resolve(undefined); + }; + return { abortController, waitUntilCompleted, markCompleted }; +} + function activeTurnKey(sessionId: string, turnId: string): string { return `${sessionId}:${turnId}`; } +function isPausedRun(run: ActiveTurnRun | undefined): run is ActiveTurnRun { + if (!run) { + return false; + } + if (run.phase === 'rebuilding') { + return true; + } + return run.turn.state.status === 'paused'; +} + export class ActiveTurnRegistry { private readonly runs = new Map(); /** Per-turn mutex. Dropped once unlocked, which means no waiter is still queued. */ private readonly turnLocks = new Map(); + // active turns keepalive timers, whenever turn is paused, we start a timer to keep the turn alive for a certain amount of time + // if the timer expires, we offload the turn + private readonly keepaliveTimers = new Map(); + // the time to keep the turn alive for, in milliseconds + private readonly pausedKeepaliveMs: number; + private readonly logger: Logger; private alreadyShutDownAbortReason: CancellationReason | undefined; + constructor(options: { logger: Logger; pausedKeepaliveMs?: number }) { + this.pausedKeepaliveMs = options.pausedKeepaliveMs ?? DEFAULT_PAUSED_KEEPALIVE_MS; + this.logger = options.logger.child({ module: 'ActiveTurnRegistry' }); + } + /** * Runs `fn` while holding the per-turn lock, after every earlier holder of * the same turn has released it. Different turns never wait on each other. @@ -50,32 +93,48 @@ export class ActiveTurnRegistry { /** * Registers the run immediately, then returns a generator that forwards * `stream` and removes the run when the stream completes (or the consumer - * exits early). If shutdown has already begun, aborts the controller with - * the shutdown reason so the turn ends as abandoned. + * exits early). If shutdown has already begun, a paused turn is offloaded and + * any other turn is aborted with the shutdown reason. */ track(input: { abortController: AbortController; stream: AsyncIterable; turn: TurnHandle }): AsyncGenerator { const key = activeTurnKey(input.turn.session_id, input.turn.id); - const { promise: waitUntilCompleted, resolve } = Promise.withResolvers(); - const markCompleted = (): void => { - resolve(undefined); - }; - const run: ActiveTurnRun = { - abortController: input.abortController, - waitUntilCompleted, - markCompleted, - turn: input.turn, - }; - this.runs.set(key, run); + const existing = this.runs.get(key); + let run: LiveTurnRun; + if (existing) { + // Keep the rebuilding run's completion promise. Shutdown may already be waiting on it + // and does not re-read this map entry. + if (existing.phase !== 'rebuilding' || existing.abortController !== input.abortController) { + throw new Error(`Turn ${input.turn.id} is already registered in this process`); + } + run = { + phase: 'live', + turn: input.turn, + abortController: existing.abortController, + waitUntilCompleted: existing.waitUntilCompleted, + markCompleted: existing.markCompleted, + }; + this.runs.set(key, run); + } else { + run = { + phase: 'live', + turn: input.turn, + ...setupTurnRun(input.abortController), + }; + this.runs.set(key, run); + } - // Accepted HTTP may still call track() after shutdownAndWait snapshotted; - // abort immediately so the turn ends as abandoned instead of running past exit. - if (this.alreadyShutDownAbortReason !== undefined) { - if (!input.abortController.signal.aborted) { - input.abortController.abort(this.alreadyShutDownAbortReason); + // Accepted HTTP may still call track() after shutdownAndWait snapshotted. + // Offload a paused turn so its row survives. Abort any other turn with the shutdown reason. + if (this.alreadyShutDownAbortReason !== undefined && !input.abortController.signal.aborted) { + const offload = isPausedRun(run); + if (offload) { + this.logger.info('Offloading paused turn', { cause: 'shutdown' }); } + input.abortController.abort(offload ? TurnOffloadReason : this.alreadyShutDownAbortReason); } const complete = () => { + this.clearPausedKeepalive({ sessionId: input.turn.session_id, turnId: input.turn.id }); run.markCompleted(); if (this.runs.get(key) === run) { this.runs.delete(key); @@ -110,14 +169,18 @@ export class ActiveTurnRegistry { /** * The live turn handle for a turn executing in this process, or undefined when it is not - * resumable here — never started, already terminal (the stream removed its run), aborting, or - * registered without a handle. + * resumable here — never started, already terminal (the stream removed its run), aborting, + * or still in the rebuilding phase. + * An aborted run stays registered until its drain finishes. `isCancelled` tells a + * send to reject instead of rebuilding over that drain. */ getTurnHandle(input: { sessionId: string; turnId: string }): TurnHandle | undefined { const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); - if (!run) { + if (run?.phase !== 'live') { return undefined; } + // we hide the turn from the client if it was already aborted + // so that the client does not try to send events to it if (run.abortController.signal.aborted) { return undefined; } @@ -125,21 +188,170 @@ export class ActiveTurnRegistry { } /** - * Enter shutdown mode, abort every run with `abortReason`, and wait until the - * registry is empty. Late `track()` calls (in-flight HTTP that registered after - * the first snapshot) are aborted immediately as `abortReason` and included in - * subsequent wait iterations. + * True when a local run was aborted with a cancellation reason. + * Offload is not a cancellation, so the caller may rebuild. + * Cancel does not take the turn lock and does not wait for the drain, and `getTurnHandle` + * hides the aborted run. The row can still look paused. The caller rejects the send + * instead of starting a second generator with the same executor id. + */ + isCancelled(input: { sessionId: string; turnId: string }): boolean { + const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); + if (!run?.abortController.signal.aborted) { + return false; + } + const { signal } = run.abortController; + return ( + signal.reason === CancellationReason.ServerExecutionTimeout || + signal.reason === CancellationReason.CancelledForNextTurn || + signal.reason === CancellationReason.Abandoned || + signal.reason === CancellationReason.ClientCancelled + ); + } + + /** + * Insert a rebuilding run before its stream exists. + * Cancel and shutdown can see the controller. `track()` replaces it with a live run + * that keeps this completion promise. + * Call `releaseReservedTurn` when the stream will never start. + */ + reservePausedTurn(input: { sessionId: string; turnId: string; abortController: AbortController }): void { + const key = activeTurnKey(input.sessionId, input.turnId); + if (this.runs.has(key)) { + throw new Error(`Turn ${input.turnId} is already registered in this process`); + } + const run: RebuildingTurnRun = { + phase: 'rebuilding', + ...setupTurnRun(input.abortController), + }; + this.runs.set(key, run); + // Resume can insert this run after shutdownAndWait snapshotted `runs` or already returned. + // Abort now so rebuilding does not start a stream that loop will never run. + if (this.alreadyShutDownAbortReason !== undefined && !input.abortController.signal.aborted) { + this.logger.info('Offloading paused turn', { cause: 'shutdown' }); + input.abortController.abort(TurnOffloadReason); + } + } + + /** Drop a rebuilding run whose stream never started, and unblock shutdown waiting on it. */ + releaseReservedTurn(input: { sessionId: string; turnId: string }): void { + const key = activeTurnKey(input.sessionId, input.turnId); + const run = this.runs.get(key); + if (run?.phase !== 'rebuilding') { + return; + } + this.runs.delete(key); + run.markCompleted(); + } + + /** + * After a turn pauses, keep the run for `pausedKeepaliveMs`. On expiry, offload + * under the turn lock: the generator stops without a terminal write. + * A later send on this owner rebuilds. `0` disables the timer. + */ + startPausedKeepalive(input: { sessionId: string; turnId: string }): void { + const replaced = this.discardPausedKeepalive(input); + if (this.pausedKeepaliveMs <= 0 || this.alreadyShutDownAbortReason !== undefined) { + if (replaced) { + this.logger.info('Paused keepalive cleared', { pausedKeepaliveMs: this.pausedKeepaliveMs }); + } + return; + } + const key = activeTurnKey(input.sessionId, input.turnId); + const timer = setTimeout(() => { + void 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; + } + this.keepaliveTimers.delete(key); + const run = this.runs.get(key); + if (isPausedRun(run)) { + this.logger.info('Offloading paused turn', { cause: 'keepalive' }); + await this.finishRun(run, TurnOffloadReason); + return; + } + this.logger.info('Paused keepalive expired; turn is no longer paused'); + }); + }, this.pausedKeepaliveMs); + timer.unref(); + this.keepaliveTimers.set(key, timer); + this.logger.info(replaced ? 'Paused keepalive restarted' : 'Paused keepalive started', { + pausedKeepaliveMs: this.pausedKeepaliveMs, + }); + } + + clearPausedKeepalive(input: { sessionId: string; turnId: string }): void { + if (!this.discardPausedKeepalive(input)) { + return; + } + this.logger.info('Paused keepalive cleared'); + } + + private discardPausedKeepalive(input: { sessionId: string; turnId: string }): boolean { + const key = activeTurnKey(input.sessionId, input.turnId); + const timer = this.keepaliveTimers.get(key); + if (!timer) { + return false; + } + clearTimeout(timer); + this.keepaliveTimers.delete(key); + return true; + } + + /** + * Stop a local generator left behind after the row's executor moved. + * The abort reason tells the drain not to write that row. + */ + async dropLocalTurn(input: { sessionId: string; turnId: string }): Promise { + const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); + if (!run) { + return; + } + if (!run.abortController.signal.aborted) { + this.logger.info('Offloading paused turn', { cause: 'not-owner' }); + } + await this.finishRun(run, TurnOffloadReason); + } + + /** + * Enter shutdown mode and wait until the registry is empty. + * Does not take the per-turn lock. + * A send holds it across rebuild and `dropLocalTurn`; + * Keepalive holds it across `finishRun`. + * Waiting for the lock would defer the abort until those calls returned, + * so it could not interrupt them. + * Paused live runs and rebuilding runs are offloaded so their approval state survives. */ async shutdownAndWait(abortReason: CancellationReason): Promise { this.alreadyShutDownAbortReason = abortReason; + const timerCount = this.keepaliveTimers.size; + for (const timer of this.keepaliveTimers.values()) { + clearTimeout(timer); + } + this.keepaliveTimers.clear(); + if (timerCount > 0) { + this.logger.info('Cleared paused keepalive timers for shutdown', { timerCount }); + } while (this.runs.size > 0) { const pending = Array.from(this.runs.values()); for (const run of pending) { - if (!run.abortController.signal.aborted) { - run.abortController.abort(abortReason); + if (run.abortController.signal.aborted) { + continue; + } + const offload = isPausedRun(run); + if (offload) { + this.logger.info('Offloading paused turn', { cause: 'shutdown' }); } + run.abortController.abort(offload ? TurnOffloadReason : abortReason); } await Promise.allSettled(pending.map(run => run.waitUntilCompleted)); } } + + private async finishRun(run: ActiveTurnRun, reason: CancellationReason | typeof TurnOffloadReason): Promise { + if (!run.abortController.signal.aborted) { + run.abortController.abort(reason); + } + await run.waitUntilCompleted; + } } diff --git a/packages/trueforge/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index 350f51222..328aff64d 100644 --- a/packages/trueforge/src/runtime/event-subscription/inMemory.ts +++ b/packages/trueforge/src/runtime/event-subscription/inMemory.ts @@ -51,6 +51,16 @@ export class InMemoryEventStreamStore { return sequenceNumber; } + /** Refreshes TTL on a live stream. Sequence numbers already follow the log length. */ + refreshStreamTtl(streamId: string, streamTTLSeconds?: number): void { + const stream = this.getLiveStream(streamId); + if (!stream || !streamTTLSeconds || streamTTLSeconds <= 0) { + return; + } + stream.expiresAtMs = Date.now() + streamTTLSeconds * 1_000; + this.scheduleExpiry(streamId, stream); + } + /** Returns the stream if it exists and has not passed its TTL; drops it lazily otherwise. */ getLiveStream(streamId: string): InMemoryStream | undefined { const stream = this.streams.get(streamId); @@ -113,6 +123,11 @@ export class InMemoryEventSubscription implements EventSubscri private readonly streamId: string, ) {} + resumeProducer(options?: EventSubscriptionPutOptions): Promise { + this.store.refreshStreamTtl(this.streamId, options?.streamTTLSeconds); + return Promise.resolve(); + } + put(event: T, options?: EventSubscriptionPutOptions): Promise { return Promise.resolve(this.store.append(this.streamId, event, options?.streamTTLSeconds)); } diff --git a/packages/trueforge/src/runtime/event-subscription/index.ts b/packages/trueforge/src/runtime/event-subscription/index.ts index 1fe701545..ab03654f2 100644 --- a/packages/trueforge/src/runtime/event-subscription/index.ts +++ b/packages/trueforge/src/runtime/event-subscription/index.ts @@ -19,6 +19,13 @@ export interface EventSubscriptionPollOptions { * Per-stream state (e.g. the sequence counter) lives on the instance and is freed with it. */ 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. + */ + resumeProducer(options?: EventSubscriptionPutOptions): Promise; + /** Appends an event, assigns its sequence number, and returns that number. */ put(event: T, options?: EventSubscriptionPutOptions): Promise; diff --git a/packages/trueforge/src/runtime/event-subscription/redis.ts b/packages/trueforge/src/runtime/event-subscription/redis.ts index f0ee5e8fb..3cf26c7c3 100644 --- a/packages/trueforge/src/runtime/event-subscription/redis.ts +++ b/packages/trueforge/src/runtime/event-subscription/redis.ts @@ -35,6 +35,18 @@ export class RedisEventSubscription implements EventSubscripti private readonly streamId: string, ) {} + async resumeProducer(options?: EventSubscriptionPutOptions): Promise { + const entries = await this.redis.xRevRange(this.streamId, '+', '-', { COUNT: 1 }); + const latest = entries[0]; + if (latest) { + this.nextSequenceNumber = sequenceNumberFromEntryId(this.streamId, latest.id) + 1; + } + const streamTTLSeconds = options?.streamTTLSeconds; + if (streamTTLSeconds && streamTTLSeconds > 0) { + await this.redis.expire(this.streamId, streamTTLSeconds); + } + } + async put(event: T, options?: EventSubscriptionPutOptions): Promise { const sequenceNumber = this.nextSequenceNumber; this.nextSequenceNumber += 1; diff --git a/packages/trueforge/src/runtime/turnOwnership.ts b/packages/trueforge/src/runtime/turnOwnership.ts deleted file mode 100644 index 5d7c5e4a5..000000000 --- a/packages/trueforge/src/runtime/turnOwnership.ts +++ /dev/null @@ -1,74 +0,0 @@ -import type { TurnState } from '@truefoundry/trueforge-core/agent-session'; - -/** Result of asking the owning replica to handle a turn request. */ -export enum PeerOutcome { - /** Peer replied 200 and handled the request. */ - Ok = 'ok', - /** Heartbeat missing or nobody subscribed: the owner is gone. */ - NoResponder = 'no_responder', - /** Heartbeat present but no reply in budget: the owner may still be applying. */ - Timeout = 'timeout', - /** Peer is alive but has no ActiveTurn (412); it should rebuild, not us. */ - NotRunningOnPeer = 'not_running_on_peer', - /** Redis/transport failed on this replica, so peer liveness is unknown. */ - TransportError = 'transport_error', -} - -export enum TurnOwnershipAction { - /** Handle on the local ActiveTurn. */ - Run = 'run', - /** This replica owns the turn but has no ActiveTurn: rebuild it, then handle. */ - Rebuild = 'rebuild', - /** Ask the owning replica to handle it. */ - Forward = 'forward', - /** The owning replica handled it; nothing left to do here. */ - Forwarded = 'forwarded', - /** Claim the paused turn from a dead owner; only the CAS winner rebuilds. */ - Steal = 'steal', - /** The turn cannot accept this request (terminal, or running with its executor gone). */ - Reject = 'reject', - /** Ownership cannot be resolved right now; the caller should retry later. */ - Retry = 'retry', -} - -export interface TurnOwnershipInput { - status: TurnState['status']; - ownedByThisExecutor: boolean; - hasLocalActiveTurn: boolean; - /** `null` until the owning peer has been asked. */ - peerOutcome: PeerOutcome | null; -} - -/** - * Picks what to do with a send or wake request for a turn. Only picks: it does - * not lock, claim, or call peers. Steals only a `paused` turn whose owner has - * no responder, because a `running` owner may still be executing tools. - */ -export function resolveOwnershipAction(input: TurnOwnershipInput): TurnOwnershipAction { - const { status, ownedByThisExecutor, hasLocalActiveTurn, peerOutcome } = input; - - if (status !== 'running' && status !== 'paused') { - return TurnOwnershipAction.Reject; - } - - if (ownedByThisExecutor) { - if (hasLocalActiveTurn) { - return TurnOwnershipAction.Run; - } - // A running turn without an ActiveTurn lost its executor; it is never resumed. - return status === 'paused' ? TurnOwnershipAction.Rebuild : TurnOwnershipAction.Reject; - } - - switch (peerOutcome) { - case null: - return TurnOwnershipAction.Forward; - case PeerOutcome.Ok: - return TurnOwnershipAction.Forwarded; - case PeerOutcome.NoResponder: - return status === 'paused' ? TurnOwnershipAction.Steal : TurnOwnershipAction.Retry; - case PeerOutcome.Timeout: - case PeerOutcome.NotRunningOnPeer: - case PeerOutcome.TransportError: - return TurnOwnershipAction.Retry; - } -} diff --git a/packages/trueforge/src/truefoundry/TrueFoundryAgentStore.ts b/packages/trueforge/src/truefoundry/TrueFoundryAgentStore.ts index c6b7464ad..b553865b7 100644 --- a/packages/trueforge/src/truefoundry/TrueFoundryAgentStore.ts +++ b/packages/trueforge/src/truefoundry/TrueFoundryAgentStore.ts @@ -84,12 +84,12 @@ export class TrueFoundryAgentStore implements IAgentStore> constructor(input: { inner: PostgresAgentStore; client: TrueFoundryServiceFoundryServerClient; - context: RequestContext; + requestContext: RequestContext; db: Kysely; }) { this.#inner = input.inner; this.#client = input.client; - this.#resolveAccessToken = callerAccessToken(input.context); + this.#resolveAccessToken = callerAccessToken(input.requestContext); this.#db = input.db; } diff --git a/packages/trueforge/src/truefoundry/TrueFoundrySkillStore.ts b/packages/trueforge/src/truefoundry/TrueFoundrySkillStore.ts index fc19eb641..d85d019b7 100644 --- a/packages/trueforge/src/truefoundry/TrueFoundrySkillStore.ts +++ b/packages/trueforge/src/truefoundry/TrueFoundrySkillStore.ts @@ -113,10 +113,10 @@ export class TrueFoundrySkillStore readonly #client: TrueFoundrySkillApiClient; readonly #resolveAccessToken: ResolveServiceFoundryAuthorization; - constructor(input: { client: TrueFoundrySkillApiClient; context: RequestContext }) { + constructor(input: { client: TrueFoundrySkillApiClient; requestContext: RequestContext }) { super(input); this.#client = input.client; - this.#resolveAccessToken = callerAccessToken(input.context); + this.#resolveAccessToken = callerAccessToken(input.requestContext); } async listSkills(input: ListSkillsInput, transaction?: TTransaction): Promise { diff --git a/packages/trueforge/src/truefoundry/gatewayMetadata.ts b/packages/trueforge/src/truefoundry/gatewayMetadata.ts index 213f31ea1..34456cb48 100644 --- a/packages/trueforge/src/truefoundry/gatewayMetadata.ts +++ b/packages/trueforge/src/truefoundry/gatewayMetadata.ts @@ -103,7 +103,7 @@ export function gatewayMetadataHeadersForTurn(turnMetadata: TurnMetadata | undef ); } -function headerValue(headers: Record | undefined, name: string): string | undefined { +export function headerValue(headers: Record | undefined, name: string): string | undefined { if (headers === undefined) { return undefined; } diff --git a/packages/trueforge/tests/unit/apis/agents.test.ts b/packages/trueforge/tests/unit/apis/agents.test.ts index 450ceb4c5..af5b2b1e1 100644 --- a/packages/trueforge/tests/unit/apis/agents.test.ts +++ b/packages/trueforge/tests/unit/apis/agents.test.ts @@ -105,9 +105,9 @@ describe('agents router', () => { resolveModelProviderStore: () => modelProviderStore, resolveMcpServerStore: () => new SqliteMcpServerStore(db), resolveSkillStore: () => new SqliteSkillStore(db), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), + webSearchProviderStore: new SqliteWebSearchProviderStore(db), withTransaction: callback => db.transaction().execute(callback), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, authorizer: new TrueForgeAuthorizer(), @@ -117,9 +117,9 @@ describe('agents router', () => { resolveModelProviderStore: () => modelProviderStore, resolveMcpServerStore: () => new SqliteMcpServerStore(db), resolveSkillStore: () => new SqliteSkillStore(db), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), + webSearchProviderStore: new SqliteWebSearchProviderStore(db), withTransaction: callback => db.transaction().execute(callback), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, authorizer: denyAllAuthorizer, diff --git a/packages/trueforge/tests/unit/apis/capabilities.test.ts b/packages/trueforge/tests/unit/apis/capabilities.test.ts index b095194d2..8cd1a3dc0 100644 --- a/packages/trueforge/tests/unit/apis/capabilities.test.ts +++ b/packages/trueforge/tests/unit/apis/capabilities.test.ts @@ -85,7 +85,7 @@ describe('capabilities routers', () => { return withAuth( createCapabilitiesRouter({ sandboxEnvironmentStore: mockDefaultEnvStore(), - resolveWebSearchProviderStore: () => webSearchStore ?? emptyWebSearchStore, + webSearchProviderStore: webSearchStore ?? emptyWebSearchStore, withTransaction: (() => { const run = (callback: (trx: never) => Promise): Promise => db.transaction().execute(callback as never); @@ -298,7 +298,7 @@ describe('capabilities routers', () => { const router = withAuth( createCapabilitiesRouter({ sandboxEnvironmentStore: mockDefaultEnvStore(), - resolveWebSearchProviderStore: () => emptyWebSearchStore, + webSearchProviderStore: emptyWebSearchStore, withTransaction: (() => { const run = (callback: (trx: never) => Promise): Promise => db.transaction().execute(callback as never); diff --git a/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts b/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts index 17409f077..81d408666 100644 --- a/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts +++ b/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts @@ -29,7 +29,7 @@ describe('public CRUD after session deletion', () => { await migrateSqliteToLatest(db); const sessionStore = new SqliteSessionStore(db); const sessions = new Sessions({ sessionStore }); - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }); const modelProviderStore = new SqliteModelProviderStore(db); const tokenStore = new SqliteOAuthTokenStore(db); const mcpServerStore = new McpServerWithAuthStore({ @@ -47,40 +47,46 @@ describe('public CRUD after session deletion', () => { app.route( '/', createSessionsRouter({ - sessions, - sessionStore, activeTurns, - resolveModelProviderStore: () => modelProviderStore, - resolveMcpServerStore: () => mcpServerStore, - resolveSkillStore: () => skillStore, - resolveAgentStore: () => agentStore, - resolveSandboxProviderStore: () => sandboxProviderStore, - sandboxEnvironmentStore, - resolveWebSearchProviderStore: () => webSearchProviderStore, + authorizer: new TrueForgeAuthorizer(), + logger: createLogger({ silent: true }), redis: createClient(), requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => agentStore, + resolveMcpServerStore: () => mcpServerStore, + resolveModelProviderStore: () => modelProviderStore, resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - logger: createLogger({ silent: true }), - authorizer: new TrueForgeAuthorizer(), + resolveSkillStore: () => skillStore, + sandboxEnvironmentStore, + sandboxProviderStore, + sessions, + sessionStore, + webSearchProviderStore, }), ); app.route( '/', createTurnsRouter({ - sessions, - sessionStore, activeTurns, - resolveModelProviderStore: () => modelProviderStore, - resolveMcpServerStore: () => mcpServerStore, - resolveTurnSkillsStore: () => skillStore, - resolveAgentStore: () => agentStore, + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => sandboxProviderStore, - sandboxEnvironmentStore, - resolveWebSearchProviderStore: () => webSearchProviderStore, logger: createLogger({ silent: true }), + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => agentStore, resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer: new TrueForgeAuthorizer(), + resolveTurnExecutionStores: () => + Promise.resolve({ + agentStore, + modelProviderStore, + mcpServerStore, + skillStore, + sandboxProviderStore, + webSearchProviderStore, + }), + sandboxEnvironmentStore, + sandboxProviderStore, + sessions, + sessionStore, }), ); diff --git a/packages/trueforge/tests/unit/apis/modelProviders.test.ts b/packages/trueforge/tests/unit/apis/modelProviders.test.ts index 334791203..ec32cad32 100644 --- a/packages/trueforge/tests/unit/apis/modelProviders.test.ts +++ b/packages/trueforge/tests/unit/apis/modelProviders.test.ts @@ -138,9 +138,9 @@ async function createRouters(): Promise<{ }), tokenStore, resolveSkillStore: () => new SqliteSkillStore(db), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), + webSearchProviderStore: new SqliteWebSearchProviderStore(db), resolveAgentStore: () => agentStore, withTransaction: callback => db.transaction().execute(callback), logger: winston.createLogger({ silent: true }), diff --git a/packages/trueforge/tests/unit/apis/sandboxEnvironments.test.ts b/packages/trueforge/tests/unit/apis/sandboxEnvironments.test.ts index c51953cec..3d4ecf3b1 100644 --- a/packages/trueforge/tests/unit/apis/sandboxEnvironments.test.ts +++ b/packages/trueforge/tests/unit/apis/sandboxEnvironments.test.ts @@ -52,7 +52,7 @@ describe('sandbox environments API → build controller path', () => { ({ listAgentNamesUsingSandboxEnvironment: jest.fn().mockResolvedValue([]), }) as never, - resolveSandboxProviderStore: () => sandboxProviderStore, + sandboxProviderStore, resolveRequestContext: () => requestContext, logger, }); diff --git a/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts b/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts index 2bbda5453..cb3cb88e8 100644 --- a/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts +++ b/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts @@ -1,6 +1,7 @@ import { OpenAPIHono } from '@hono/zod-openapi'; import type { AgentSpec } from '@truefoundry/trueforge-core/agent-session'; import { AgentSpecSchema, Sessions } from '@truefoundry/trueforge-core/agent-session'; +import { RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; import { createLogger } from 'winston'; import { createTurnsRouter, toContentDisposition } from '../../../src/apis/turns'; import { TrueForgeAuthorizer } from '../../../src/auth/authorizer'; @@ -39,25 +40,30 @@ async function buildApp() { app.route( '/', createTurnsRouter({ - sessions, - sessionStore, - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore: () => new SqliteModelProviderStore(db), - resolveMcpServerStore: () => - new McpServerWithAuthStore({ - store: new SqliteMcpServerStore(db), - tokenStore, - clientName: 'test-client', - }), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), - resolveAgentStore: () => new SqliteAgentStore(db), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => new SqliteAgentStore(db), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer: new TrueForgeAuthorizer(), + resolveTurnExecutionStores: () => + Promise.resolve({ + agentStore: new SqliteAgentStore(db), + modelProviderStore: new SqliteModelProviderStore(db), + mcpServerStore: new McpServerWithAuthStore({ + store: new SqliteMcpServerStore(db), + tokenStore, + clientName: 'test-client', + }), + skillStore: new SqliteSkillStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + webSearchProviderStore: new SqliteWebSearchProviderStore(db), + }), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore, }), ); diff --git a/packages/trueforge/tests/unit/apis/sandboxProviders.test.ts b/packages/trueforge/tests/unit/apis/sandboxProviders.test.ts index e3ff236cb..d2e55c501 100644 --- a/packages/trueforge/tests/unit/apis/sandboxProviders.test.ts +++ b/packages/trueforge/tests/unit/apis/sandboxProviders.test.ts @@ -63,7 +63,7 @@ async function createRouters(): Promise<{ const sandboxProviderStore = new SqliteSandboxProviderStore(db); return { settingsRouter: createSandboxProvidersRouter({ - resolveSandboxProviderStore: () => sandboxProviderStore, + sandboxProviderStore, sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), withTransaction: callback => db.transaction().execute(callback), logger, @@ -83,7 +83,7 @@ describe('sandboxProviders router', () => { await migrateSqliteToLatest(db); sandboxProviderStore = new SqliteSandboxProviderStore(db); settingsRouter = createSandboxProvidersRouter({ - resolveSandboxProviderStore: () => sandboxProviderStore, + sandboxProviderStore, sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), withTransaction: callback => db.transaction().execute(callback), logger, diff --git a/packages/trueforge/tests/unit/apis/schedules.test.ts b/packages/trueforge/tests/unit/apis/schedules.test.ts index 49194728a..d1f0a4d6e 100644 --- a/packages/trueforge/tests/unit/apis/schedules.test.ts +++ b/packages/trueforge/tests/unit/apis/schedules.test.ts @@ -62,18 +62,19 @@ const scheduleBody = { function stubTurnExecutionDeps(agentStore: SqliteAgentStore, scheduleStore: SqliteScheduleStore) { const sessionStore = new InMemorySessionStore(); return { - scheduleStore, - sessions: new Sessions({ sessionStore }), + activeTurns: new ActiveTurnRegistry({ logger: winston.createLogger({ silent: true }) }), agentStore, - activeTurns: new ActiveTurnRegistry(), eventSubscriptions: new EventSubscriptionRegistry(undefined), logger: winston.createLogger({ silent: true }), - resolveModelProviderStore: () => ({}) as never, resolveMcpServerStore: () => ({}) as never, + resolveModelProviderStore: () => ({}) as never, resolveTurnSkillsStore: () => ({ resolveTurnSkills: async () => [] }), - resolveSandboxProviderStore: () => ({}) as never, sandboxEnvironmentStore: {} as never, - resolveWebSearchProviderStore: () => ({}) as never, + sandboxProviderStore: {} as never, + scheduleStore, + sessions: new Sessions({ sessionStore }), + sessionStore, + webSearchProviderStore: {} as never, }; } diff --git a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts index c3b63eef1..15673b631 100644 --- a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts +++ b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts @@ -13,6 +13,7 @@ import { import { createTurnsRouter } from '../../../src/apis/turns'; import { TrueForgeAuthorizer, type Authorizer } from '../../../src/auth/authorizer'; import { STANDALONE_REQUEST_CONTEXT } from '../../../src/auth/identity'; +import configuration from '../../../src/config'; import { McpServerWithAuthStore } from '../../../src/db/McpServerWithAuthStore'; import { migrateSqliteToLatest } from '../../../src/db/migrateSqlite'; import { SqliteAgentStore } from '../../../src/db/sqlite/agent-store/SqliteAgentStore'; @@ -115,21 +116,21 @@ describe('sessions HTTP agent binding', () => { resolveMcpServerStore = jest.fn(() => mcpServerStore); resolveSkillStore = jest.fn(() => skillStore); const deps: SessionsRouterDeps = { - sessions, - sessionStore, - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore, - resolveMcpServerStore, - resolveSkillStore, - resolveAgentStore: () => agentStore, - resolveSandboxProviderStore: () => sandboxProviderStore, - sandboxEnvironmentStore, - resolveWebSearchProviderStore: () => webSearchProviderStore, + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer: new TrueForgeAuthorizer(), + logger: createLogger({ silent: true }), redis: createClient(), requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => agentStore, + resolveMcpServerStore, + resolveModelProviderStore, resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - logger: createLogger({ silent: true }), - authorizer: new TrueForgeAuthorizer(), + resolveSkillStore, + sandboxEnvironmentStore, + sandboxProviderStore, + sessions, + sessionStore, + webSearchProviderStore, }; sessionDeps = deps; app = new OpenAPIHono(); @@ -138,25 +139,30 @@ describe('sessions HTTP agent binding', () => { app.route( '/', createTurnsRouter({ - sessions, - sessionStore, activeTurns: deps.activeTurns, - resolveModelProviderStore: () => modelProviderStore, - resolveMcpServerStore: () => - new McpServerWithAuthStore({ - store: mcpServerStore, - tokenStore, - clientName: 'test-client', - }), - resolveTurnSkillsStore: () => skillStore, - resolveAgentStore: () => agentStore, + authorizer: deps.authorizer, eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => sandboxProviderStore, - sandboxEnvironmentStore, - resolveWebSearchProviderStore: () => webSearchProviderStore, logger: deps.logger, + requestReplyRouter: deps.requestReplyRouter, + resolveAgentStore: () => agentStore, resolveRequestContext: deps.resolveRequestContext, - authorizer: deps.authorizer, + resolveTurnExecutionStores: () => + Promise.resolve({ + agentStore, + modelProviderStore, + mcpServerStore: new McpServerWithAuthStore({ + store: mcpServerStore, + tokenStore, + clientName: 'test-client', + }), + skillStore, + sandboxProviderStore, + webSearchProviderStore, + }), + sandboxEnvironmentStore, + sandboxProviderStore, + sessions, + sessionStore, }), ); app.route('/api/internal/sessions', createInternalSessionsRouter(deps)); @@ -486,11 +492,17 @@ describe('sessions HTTP agent binding', () => { const created = await app.request('/', jsonInit('POST', { agent: { spec: inlineSpec } })); expect(created.status).toBe(201); const { data: session } = (await created.json()) as { data: { id: string } }; - await sessionStore.createTurn(makeCreateTurnInput({ sessionId: session.id, turnId: 'tip-1' })); + await sessionStore.createTurn( + makeCreateTurnInput({ sessionId: session.id, turnId: 'tip-1', activeExecutorId: configuration.EXECUTOR_ID }), + ); const send = jest.fn().mockImplementation(async function* (events: Parameters[0]) { yield events; }); - jest.spyOn(sessionDeps.activeTurns, 'getTurnHandle').mockReturnValue({ send } as unknown as TurnHandle); + jest.spyOn(sessionDeps.activeTurns, 'getTurnHandle').mockReturnValue({ + send, + state: { status: 'running' }, + record: { active_executor_id: configuration.EXECUTOR_ID }, + } as unknown as TurnHandle); const res = await app.request( `/${session.id}/turns/tip-1/events`, @@ -559,7 +571,9 @@ describe('sessions HTTP agent binding', () => { const created = await app.request('/', jsonInit('POST', { agent: { spec: inlineSpec } })); expect(created.status).toBe(201); const { data: session } = (await created.json()) as { data: { id: string } }; - await sessionStore.createTurn(makeCreateTurnInput({ sessionId: session.id, turnId: 'tip-1' })); + await sessionStore.createTurn( + makeCreateTurnInput({ sessionId: session.id, turnId: 'tip-1', activeExecutorId: configuration.EXECUTOR_ID }), + ); const registry = sessionDeps.activeTurns; const gate = Promise.withResolvers(); diff --git a/packages/trueforge/tests/unit/apis/sessions.test.ts b/packages/trueforge/tests/unit/apis/sessions.test.ts index eb84dbc76..6a4e483bf 100644 --- a/packages/trueforge/tests/unit/apis/sessions.test.ts +++ b/packages/trueforge/tests/unit/apis/sessions.test.ts @@ -8,10 +8,13 @@ import type { import { CancellationReason, TurnNotFoundError } from '@truefoundry/trueforge-core/agent-session'; import { NoResponderError, redisRequest, RequestTimeoutError } from '@truefoundry/trueforge-core/request-reply'; import type { RedisClientType } from 'redis'; +import { createLogger } from 'winston'; import { cancelSessionTurn } from '../../../src/apis/sessions'; import configuration from '../../../src/config'; import { ActiveTurnRegistry } from '../../../src/runtime/activeTurns'; +const logger = createLogger({ silent: true }); + jest.mock('@truefoundry/trueforge-core/request-reply', () => { const actual = jest.requireActual( '@truefoundry/trueforge-core/request-reply', @@ -86,7 +89,7 @@ function cancelDeps(input: { function trackRun(registry: ActiveTurnRegistry, turnId: string): AbortController { const abortController = new AbortController(); registry.track({ - turn: { session_id: SESSION_ID, id: turnId } as TurnHandle, + turn: { session_id: SESSION_ID, id: turnId, state: { status: 'running' } } as TurnHandle, abortController, stream: (async function* () { await new Promise(() => undefined); @@ -102,7 +105,7 @@ describe('cancelSessionTurn', () => { }); it('aborts a turn running in this process without freezing', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-local-running'; const abortController = trackRun(activeTurns, turnId); const session = sessionHandle(); @@ -120,7 +123,7 @@ describe('cancelSessionTurn', () => { }); it('freezes when this executor owns the turn but the run is gone', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-local-gone'; const session = sessionHandle(); @@ -138,12 +141,12 @@ describe('cancelSessionTurn', () => { }); it('aborts a paused turn parked in this process, then drops it from the registry', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-local-paused'; const abortController = new AbortController(); const parked = Promise.withResolvers(); const tracked = activeTurns.track({ - turn: { session_id: SESSION_ID, id: turnId } as TurnHandle, + turn: { session_id: SESSION_ID, id: turnId, state: { status: 'paused' } } as TurnHandle, abortController, stream: (async function* () { yield 'paused'; @@ -183,7 +186,7 @@ describe('cancelSessionTurn', () => { }); it('freezes a paused turn this executor owns when the run is gone', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-paused-gone'; const session = sessionHandle(); @@ -199,7 +202,7 @@ describe('cancelSessionTurn', () => { }); it('asks the owning peer to cancel a paused turn and does not freeze when the peer aborts', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-paused-remote'; const session = sessionHandle(); redisRequestMock.mockResolvedValue({ status: 200, body: {} }); @@ -219,7 +222,7 @@ describe('cancelSessionTurn', () => { }); it('freezes when the run is not in this process and there is no Redis client', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-remote-no-redis'; const session = sessionHandle(); @@ -239,7 +242,7 @@ describe('cancelSessionTurn', () => { }); it('does not freeze when the owning peer aborts', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-remote-ok'; const session = sessionHandle(); redisRequestMock.mockResolvedValue({ status: 200, body: {} }); @@ -263,7 +266,7 @@ describe('cancelSessionTurn', () => { }); it('freezes when the owning peer is alive but the turn is not running there', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-remote-412'; const session = sessionHandle(); redisRequestMock.mockResolvedValue({ status: 412, body: { message: 'Turn is not running on this executor' } }); @@ -282,10 +285,10 @@ describe('cancelSessionTurn', () => { }); it('freezes when the owning executor is unreachable', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-remote-no-responder'; const session = sessionHandle(); - const logger = silentLogger(); + const cancelLogger = silentLogger(); redisRequestMock.mockRejectedValue(new NoResponderError(REMOTE_EXECUTOR)); await expect( @@ -295,23 +298,23 @@ describe('cancelSessionTurn', () => { turn: turnRecord({ turnId: turnId, state: { status: 'running' }, activeExecutorId: REMOTE_EXECUTOR }), session, redis: REDIS, - logger, + logger: cancelLogger, }), { turnId }, ), ).resolves.toBeUndefined(); expect(session.freezeTurn).toHaveBeenCalled(); - expect(logger.warn).toHaveBeenCalledWith( + expect(cancelLogger.warn).toHaveBeenCalledWith( 'Failed to reach owning executor over Redis; freezing the turn', expect.objectContaining({ turn_id: turnId, owner: REMOTE_EXECUTOR }), ); }); it('freezes when waiting for the owning executor times out', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-remote-timeout'; const session = sessionHandle(); - const logger = silentLogger(); + const cancelLogger = silentLogger(); redisRequestMock.mockRejectedValue(new RequestTimeoutError(60_000)); await expect( @@ -321,23 +324,23 @@ describe('cancelSessionTurn', () => { turn: turnRecord({ turnId: turnId, state: { status: 'running' }, activeExecutorId: REMOTE_EXECUTOR }), session, redis: REDIS, - logger, + logger: cancelLogger, }), { turnId }, ), ).resolves.toBeUndefined(); expect(session.freezeTurn).toHaveBeenCalled(); - expect(logger.warn).toHaveBeenCalledWith( + expect(cancelLogger.warn).toHaveBeenCalledWith( 'Timed out waiting for owning executor to cancel; freezing the turn', expect.objectContaining({ turn_id: turnId, owner: REMOTE_EXECUTOR }), ); }); it('freezes when Redis itself cannot be reached', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-remote-redis-down'; const session = sessionHandle(); - const logger = silentLogger(); + const cancelLogger = silentLogger(); redisRequestMock.mockRejectedValue(new Error('Redis connection closed')); await expect( @@ -347,13 +350,13 @@ describe('cancelSessionTurn', () => { turn: turnRecord({ turnId: turnId, state: { status: 'running' }, activeExecutorId: REMOTE_EXECUTOR }), session, redis: REDIS, - logger, + logger: cancelLogger, }), { turnId }, ), ).resolves.toBeUndefined(); expect(session.freezeTurn).toHaveBeenCalled(); - expect(logger.warn).toHaveBeenCalledWith( + expect(cancelLogger.warn).toHaveBeenCalledWith( 'Failed to reach owning executor over Redis; freezing the turn', expect.objectContaining({ turn_id: turnId, @@ -364,7 +367,7 @@ describe('cancelSessionTurn', () => { }); it('treats a missing turn as a successful cancel', async () => { - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ logger }); const turnId = 'turn-missing'; const session = sessionHandle(); session.freezeTurn = jest.fn().mockRejectedValue(new TurnNotFoundError(turnId)); diff --git a/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts b/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts index 026e609db..c5b25220e 100644 --- a/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts +++ b/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts @@ -1,6 +1,7 @@ import { OpenAPIHono } from '@hono/zod-openapi'; import type { Sessions } from '@truefoundry/trueforge-core/agent-session'; import { AgentHarnessError } from '@truefoundry/trueforge-core/core'; +import { RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; import { createLogger } from 'winston'; import { createTurnsRouter } from '../../../src/apis/turns'; import { createAppErrorHandler } from '../../../src/app'; @@ -68,25 +69,30 @@ async function postTurnRejectingWith(error: AgentHarnessError): Promise modelProviderStore, - resolveMcpServerStore: () => - new McpServerWithAuthStore({ - store: new SqliteMcpServerStore(db), - tokenStore, - clientName: 'test-client', - }), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), - resolveAgentStore: () => new SqliteAgentStore(db), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => new SqliteAgentStore(db), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer: new TrueForgeAuthorizer(), + resolveTurnExecutionStores: () => + Promise.resolve({ + agentStore: new SqliteAgentStore(db), + modelProviderStore, + mcpServerStore: new McpServerWithAuthStore({ + store: new SqliteMcpServerStore(db), + tokenStore, + clientName: 'test-client', + }), + skillStore: new SqliteSkillStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + webSearchProviderStore: new SqliteWebSearchProviderStore(db), + }), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore: new SqliteSessionStore(db), }), ); diff --git a/packages/trueforge/tests/unit/apis/turns.test.ts b/packages/trueforge/tests/unit/apis/turns.test.ts index 961cededf..dd732d052 100644 --- a/packages/trueforge/tests/unit/apis/turns.test.ts +++ b/packages/trueforge/tests/unit/apis/turns.test.ts @@ -3,9 +3,10 @@ import { AgentSpecSchema, Sessions, TurnNotFoundError, - type TurnInputItem, type TurnStreamingEvent, + type UserMessage, } from '@truefoundry/trueforge-core/agent-session'; +import { RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; import type { Kysely } from 'kysely'; import { createLogger } from 'winston'; import { createTurnsRouter, deriveSessionTitle, turnStreamId } from '../../../src/apis/turns'; @@ -36,13 +37,24 @@ function mcpServerStoreWithAuth(db: Kysely, tokenStore: SqliteOAuthTok }); } +function turnExecutionStores(db: Kysely, tokenStore: SqliteOAuthTokenStore) { + return { + agentStore: new SqliteAgentStore(db), + modelProviderStore: new SqliteModelProviderStore(db), + mcpServerStore: mcpServerStoreWithAuth(db, tokenStore), + skillStore: new SqliteSkillStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + webSearchProviderStore: new SqliteWebSearchProviderStore(db), + }; +} + describe('turns', () => { it('namespaces turn stream ids under tfg', () => { expect(turnStreamId('ten', 'sess', 'turn1')).toBe('tfg:agent:turn:ten:sess:turn1:stream'); }); describe('deriveSessionTitle', () => { - function userMessage(content: string): TurnInputItem[] { + function userMessage(content: string): UserMessage[] { return [{ type: 'user.message', content }]; } @@ -119,20 +131,18 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - sessions, - sessionStore, - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore: () => new SqliteModelProviderStore(db), - resolveMcpServerStore: () => mcpServerStoreWithAuth(db, tokenStore), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), - resolveAgentStore: () => new SqliteAgentStore(db), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => new SqliteAgentStore(db), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer: new TrueForgeAuthorizer(), + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, tokenStore)), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore, }), ); @@ -215,19 +225,7 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - sessions: new Sessions({ sessionStore }), - sessionStore, - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore: () => new SqliteModelProviderStore(db), - resolveMcpServerStore: () => mcpServerStoreWithAuth(db, new SqliteOAuthTokenStore(db)), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), - resolveAgentStore: () => agentStore, - eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), - logger: createLogger({ silent: true }), - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer: { listAgentAccess: input => Promise.resolve( @@ -241,6 +239,16 @@ describe('turns', () => { permissions: Object.fromEntries(resourceIds.map(id => [id, []])), }), }, + eventSubscriptions: new EventSubscriptionRegistry(undefined), + logger: createLogger({ silent: true }), + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => agentStore, + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, new SqliteOAuthTokenStore(db))), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions: new Sessions({ sessionStore }), + sessionStore, }), ); @@ -317,20 +325,18 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - sessions: new Sessions({ sessionStore }), - sessionStore, - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore: () => new SqliteModelProviderStore(db), - resolveMcpServerStore: () => mcpServerStoreWithAuth(db, new SqliteOAuthTokenStore(db)), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), - resolveAgentStore: () => new SqliteAgentStore(db), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => new SqliteAgentStore(db), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer: new TrueForgeAuthorizer(), + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, new SqliteOAuthTokenStore(db))), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions: new Sessions({ sessionStore }), + sessionStore, }), ); @@ -388,20 +394,18 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - sessions, - sessionStore, - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore: () => new SqliteModelProviderStore(db), - resolveMcpServerStore: () => mcpServerStoreWithAuth(db, new SqliteOAuthTokenStore(db)), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), - resolveAgentStore: () => new SqliteAgentStore(db), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => new SqliteAgentStore(db), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer: new TrueForgeAuthorizer(), + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, new SqliteOAuthTokenStore(db))), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore, }), ); @@ -496,20 +500,18 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - sessions, - sessionStore: new SqliteSessionStore(db), - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore: () => modelProviderStore, - resolveMcpServerStore: () => mcpServerStoreWithAuth(db, tokenStore), - resolveAgentStore: () => new SqliteAgentStore(db), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions, - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger, + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => new SqliteAgentStore(db), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer: new TrueForgeAuthorizer(), + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, tokenStore)), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore: new SqliteSessionStore(db), }), ); @@ -627,20 +629,18 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - sessions, - sessionStore: new SqliteSessionStore(db), - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore: () => modelProviderStore, - resolveMcpServerStore: () => mcpServerStoreWithAuth(db, tokenStore), - resolveAgentStore: () => new SqliteAgentStore(db), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger, + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => new SqliteAgentStore(db), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer: new TrueForgeAuthorizer(), + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, tokenStore)), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore: new SqliteSessionStore(db), }), ); @@ -746,20 +746,18 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - sessions, - sessionStore: new SqliteSessionStore(db), - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore: () => modelProviderStore, - resolveMcpServerStore: () => mcpServerStoreWithAuth(db, tokenStore), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), - resolveAgentStore: () => new SqliteAgentStore(db), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger, + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => new SqliteAgentStore(db), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer: new TrueForgeAuthorizer(), + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, tokenStore)), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore: new SqliteSessionStore(db), }), ); @@ -833,20 +831,18 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - sessions, - sessionStore, - activeTurns: new ActiveTurnRegistry(), - resolveModelProviderStore: () => new SqliteModelProviderStore(db), - resolveMcpServerStore: () => mcpServerStoreWithAuth(db, tokenStore), - resolveTurnSkillsStore: () => new SqliteSkillStore(db), - resolveAgentStore: () => agentStore, + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), + authorizer, eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), + requestReplyRouter: new RequestReplyRouter(), + resolveAgentStore: () => agentStore, resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, - authorizer, + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, tokenStore)), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore, }), ); return { app, agent, agentStore }; diff --git a/packages/trueforge/tests/unit/apis/webSearchProviders.test.ts b/packages/trueforge/tests/unit/apis/webSearchProviders.test.ts index ddb050c2c..7957c6ba4 100644 --- a/packages/trueforge/tests/unit/apis/webSearchProviders.test.ts +++ b/packages/trueforge/tests/unit/apis/webSearchProviders.test.ts @@ -18,7 +18,7 @@ async function createRouter() { await migrateSqliteToLatest(db); const store = new SqliteWebSearchProviderStore(db); const router = createWebSearchProvidersRouter({ - resolveWebSearchProviderStore: () => store, + webSearchProviderStore: store, resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, }); return { router, store }; diff --git a/packages/trueforge/tests/unit/auth/identity.test.ts b/packages/trueforge/tests/unit/auth/identity.test.ts index 12a4062fb..4c6f43135 100644 --- a/packages/trueforge/tests/unit/auth/identity.test.ts +++ b/packages/trueforge/tests/unit/auth/identity.test.ts @@ -30,6 +30,7 @@ describe('requestContextFromCreatedBySubject', () => { subject_type: 'user', subject_display_name: 'Alice', }, + user_credential: null, }); expect(rc.tenant_id).toBe('acme'); expect(rc.subject).toEqual({ id: 'alice', type: 'user', display_name: 'Alice' }); diff --git a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts index 6e78520ce..c7026627e 100644 --- a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts +++ b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts @@ -1,8 +1,11 @@ -import { CancellationReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; +import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; +import { createLogger } from 'winston'; import { ActiveTurnRegistry } from '../../../src/runtime/activeTurns'; -function turn(sessionId: string, turnId: string): TurnHandle { - return { session_id: sessionId, id: turnId } as TurnHandle; +const logger = createLogger({ silent: true }); + +function turn(sessionId: string, turnId: string, status: 'running' | 'paused' = 'running'): TurnHandle { + return { session_id: sessionId, id: turnId, state: { status } } as TurnHandle; } async function* values(items: T[]): AsyncGenerator { @@ -32,7 +35,7 @@ async function* gateOnAbort(signal: AbortSignal): AsyncGenerator { describe('ActiveTurnRegistry', () => { it('track passes events through and removes the run when the stream completes', async () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); const abortController = new AbortController(); const tracked = registry.track({ turn: turn('s1', 't1'), @@ -55,7 +58,7 @@ describe('ActiveTurnRegistry', () => { }); it('track cleans up when the consumer breaks early', async () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); const abortController = new AbortController(); const tracked = registry.track({ turn: turn('s1', 't1'), @@ -77,7 +80,7 @@ describe('ActiveTurnRegistry', () => { }); it('track cleans up when the stream throws', async () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); const abortController = new AbortController(); const tracked = registry.track({ turn: turn('s1', 't1'), @@ -100,7 +103,7 @@ describe('ActiveTurnRegistry', () => { }); it('cancelIfRunning aborts with the given reason and returns true', () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); const abortController = new AbortController(); void registry.track({ turn: turn('s1', 't1'), @@ -117,10 +120,24 @@ describe('ActiveTurnRegistry', () => { ).toBe(true); expect(abortController.signal.aborted).toBe(true); expect(abortController.signal.reason).toBe(CancellationReason.ClientCancelled); + expect(registry.isCancelled({ sessionId: 's1', turnId: 't1' })).toBe(true); + }); + + it('isCancelled ignores an offloaded run', () => { + const registry = new ActiveTurnRegistry({ logger }); + const abortController = new AbortController(); + void registry.track({ + turn: turn('s1', 't1', 'paused'), + abortController, + stream: values([1]), + }); + abortController.abort(TurnOffloadReason); + + expect(registry.isCancelled({ sessionId: 's1', turnId: 't1' })).toBe(false); }); it('cancelIfRunning returns false for unknown ids', () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); expect( registry.cancelIfRunning({ sessionId: 'missing', @@ -131,7 +148,7 @@ describe('ActiveTurnRegistry', () => { }); it('cancelIfRunning does not re-abort an already-aborted controller', () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); const abortController = new AbortController(); void registry.track({ turn: turn('s1', 't1'), @@ -151,7 +168,7 @@ describe('ActiveTurnRegistry', () => { }); it('shutdownAndWait aborts runs and waits until tracked streams finish', async () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); const abortController = new AbortController(); const tracked = registry.track({ turn: turn('s1', 't1'), @@ -179,7 +196,7 @@ describe('ActiveTurnRegistry', () => { }); it('late track after shutdownAndWait aborts immediately with the shutdown reason', async () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); await registry.shutdownAndWait(CancellationReason.Abandoned); const abortController = new AbortController(); @@ -196,4 +213,23 @@ describe('ActiveTurnRegistry', () => { void value; } }); + + it('late track after shutdownAndWait offloads a paused turn', async () => { + const registry = new ActiveTurnRegistry({ logger }); + await registry.shutdownAndWait(CancellationReason.Abandoned); + + const abortController = new AbortController(); + const tracked = registry.track({ + turn: turn('s1', 'late-paused', 'paused'), + abortController, + stream: values(['x']), + }); + + expect(abortController.signal.aborted).toBe(true); + expect(abortController.signal.reason).toBe(TurnOffloadReason); + + for await (const value of tracked) { + void value; + } + }); }); diff --git a/packages/trueforge/tests/unit/runtime/turnOwnership.test.ts b/packages/trueforge/tests/unit/runtime/turnOwnership.test.ts deleted file mode 100644 index 04672786f..000000000 --- a/packages/trueforge/tests/unit/runtime/turnOwnership.test.ts +++ /dev/null @@ -1,139 +0,0 @@ -import type { TurnState } from '@truefoundry/trueforge-core/agent-session'; -import { - PeerOutcome, - resolveOwnershipAction, - TurnOwnershipAction, - type TurnOwnershipInput, -} from '../../../src/runtime/turnOwnership'; - -const TERMINAL_STATUSES: TurnState['status'][] = ['done', 'cancelled', 'error']; -const NON_TERMINAL_STATUSES: TurnState['status'][] = ['running', 'paused']; -const PEER_OUTCOMES: (PeerOutcome | null)[] = [null, ...Object.values(PeerOutcome)]; - -interface Row { - name: string; - input: TurnOwnershipInput; - expected: TurnOwnershipAction; -} - -const ROWS: Row[] = [ - { - name: 'owner with a live ActiveTurn runs locally (running)', - input: { status: 'running', ownedByThisExecutor: true, hasLocalActiveTurn: true, peerOutcome: null }, - expected: TurnOwnershipAction.Run, - }, - { - name: 'owner with a live ActiveTurn runs locally (paused, inside keepalive)', - input: { status: 'paused', ownedByThisExecutor: true, hasLocalActiveTurn: true, peerOutcome: null }, - expected: TurnOwnershipAction.Run, - }, - { - name: 'owner without an ActiveTurn rebuilds a paused turn (not a steal)', - input: { status: 'paused', ownedByThisExecutor: true, hasLocalActiveTurn: false, peerOutcome: null }, - expected: TurnOwnershipAction.Rebuild, - }, - { - name: 'owner without an ActiveTurn rejects a running turn (execute gone)', - input: { status: 'running', ownedByThisExecutor: true, hasLocalActiveTurn: false, peerOutcome: null }, - expected: TurnOwnershipAction.Reject, - }, - { - name: 'non-owner forwards a paused turn before asking the peer', - input: { status: 'paused', ownedByThisExecutor: false, hasLocalActiveTurn: false, peerOutcome: null }, - expected: TurnOwnershipAction.Forward, - }, - { - name: 'non-owner forwards a running turn before asking the peer', - input: { status: 'running', ownedByThisExecutor: false, hasLocalActiveTurn: false, peerOutcome: null }, - expected: TurnOwnershipAction.Forward, - }, - { - name: 'peer 200 means the owner handled it', - input: { status: 'paused', ownedByThisExecutor: false, hasLocalActiveTurn: false, peerOutcome: PeerOutcome.Ok }, - expected: TurnOwnershipAction.Forwarded, - }, - { - name: 'paused + no responder steals', - input: { - status: 'paused', - ownedByThisExecutor: false, - hasLocalActiveTurn: false, - peerOutcome: PeerOutcome.NoResponder, - }, - expected: TurnOwnershipAction.Steal, - }, - { - name: 'running + no responder never steals', - input: { - status: 'running', - ownedByThisExecutor: false, - hasLocalActiveTurn: false, - peerOutcome: PeerOutcome.NoResponder, - }, - expected: TurnOwnershipAction.Retry, - }, - { - name: 'paused + peer timeout does not steal (owner may still be applying)', - input: { - status: 'paused', - ownedByThisExecutor: false, - hasLocalActiveTurn: false, - peerOutcome: PeerOutcome.Timeout, - }, - expected: TurnOwnershipAction.Retry, - }, - { - name: 'paused + peer 412 does not steal (empty registry is not a dead executor)', - input: { - status: 'paused', - ownedByThisExecutor: false, - hasLocalActiveTurn: false, - peerOutcome: PeerOutcome.NotRunningOnPeer, - }, - expected: TurnOwnershipAction.Retry, - }, - { - name: 'paused + local transport error does not steal (cannot see heartbeats)', - input: { - status: 'paused', - ownedByThisExecutor: false, - hasLocalActiveTurn: false, - peerOutcome: PeerOutcome.TransportError, - }, - expected: TurnOwnershipAction.Retry, - }, -]; - -describe('resolveOwnershipAction', () => { - it.each(ROWS)('$name', ({ input, expected }) => { - expect(resolveOwnershipAction(input)).toBe(expected); - }); - - it('rejects every terminal turn regardless of ownership or peer outcome', () => { - for (const status of TERMINAL_STATUSES) { - for (const ownedByThisExecutor of [true, false]) { - for (const hasLocalActiveTurn of [true, false]) { - for (const peerOutcome of PEER_OUTCOMES) { - expect(resolveOwnershipAction({ status, ownedByThisExecutor, hasLocalActiveTurn, peerOutcome })).toBe( - TurnOwnershipAction.Reject, - ); - } - } - } - } - }); - - it('steals only for paused + no responder', () => { - for (const status of NON_TERMINAL_STATUSES) { - for (const ownedByThisExecutor of [true, false]) { - for (const hasLocalActiveTurn of [true, false]) { - for (const peerOutcome of PEER_OUTCOMES) { - const action = resolveOwnershipAction({ status, ownedByThisExecutor, hasLocalActiveTurn, peerOutcome }); - const shouldSteal = status === 'paused' && !ownedByThisExecutor && peerOutcome === PeerOutcome.NoResponder; - expect(action === TurnOwnershipAction.Steal).toBe(shouldSteal); - } - } - } - } - }); -}); diff --git a/packages/trueforge/tests/unit/sandbox/sandboxEnvironmentSecrets.test.ts b/packages/trueforge/tests/unit/sandbox/sandboxEnvironmentSecrets.test.ts index 4848adb15..607132edd 100644 --- a/packages/trueforge/tests/unit/sandbox/sandboxEnvironmentSecrets.test.ts +++ b/packages/trueforge/tests/unit/sandbox/sandboxEnvironmentSecrets.test.ts @@ -48,7 +48,7 @@ describe('sandbox environment secrets', () => { createSandboxEnvironmentsRouter({ sandboxEnvironmentStore, resolveAgentStore: () => ({ listAgentNamesUsingSandboxEnvironment: jest.fn().mockResolvedValue([]) }) as never, - resolveSandboxProviderStore: () => sandboxProviderStore, + sandboxProviderStore, resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, logger, }), diff --git a/packages/trueforge/tests/unit/truefoundry/TrueFoundryAgentStore.test.ts b/packages/trueforge/tests/unit/truefoundry/TrueFoundryAgentStore.test.ts index 9d66285fe..3e6c78ca7 100644 --- a/packages/trueforge/tests/unit/truefoundry/TrueFoundryAgentStore.test.ts +++ b/packages/trueforge/tests/unit/truefoundry/TrueFoundryAgentStore.test.ts @@ -99,7 +99,7 @@ function tfStore(input: { return new TrueFoundryAgentStore({ inner: input.inner, client: input.client, - context: { + requestContext: { tenant_id: TENANT, subject: { id: 'user-1', type: 'user', display_name: 'user-1' }, roles: [], diff --git a/packages/trueforge/tests/unit/truefoundry/TrueFoundrySkillStore.test.ts b/packages/trueforge/tests/unit/truefoundry/TrueFoundrySkillStore.test.ts index ebea908f8..4893a01b7 100644 --- a/packages/trueforge/tests/unit/truefoundry/TrueFoundrySkillStore.test.ts +++ b/packages/trueforge/tests/unit/truefoundry/TrueFoundrySkillStore.test.ts @@ -53,7 +53,7 @@ function createStore( }; const store = new TrueFoundrySkillStore({ client, - context: { + requestContext: { tenant_id: TENANT, subject: { id: 'user-1', type: 'user', display_name: 'user-1' }, roles: [], diff --git a/packages/trueforge/tests/unit/truefoundry/trueFoundryNaming.test.ts b/packages/trueforge/tests/unit/truefoundry/trueFoundryNaming.test.ts index f5c626a68..2b7400e19 100644 --- a/packages/trueforge/tests/unit/truefoundry/trueFoundryNaming.test.ts +++ b/packages/trueforge/tests/unit/truefoundry/trueFoundryNaming.test.ts @@ -177,7 +177,7 @@ describe('TrueFoundry naming vs NameSchema', () => { .mockResolvedValue([{ fqn: SKILL_FQN, name: 'my.echo', description: 'Echo' }]), apiKey: 'tfy-api-key', }; - const store = new TrueFoundrySkillStore({ client, context: requestContext }); + const store = new TrueFoundrySkillStore({ client, requestContext }); const listed = await store.listSkills({ tenant_id: 'default', names: undefined }); expect(listed).toHaveLength(1); 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 7d18c4749..8cfda1477 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py @@ -25,6 +25,7 @@ from ..errors.failed_dependency_error import FailedDependencyError from ..errors.forbidden_error import ForbiddenError from ..errors.gone_error import GoneError +from ..errors.internal_server_error import InternalServerError from ..errors.not_found_error import NotFoundError from ..errors.precondition_failed_error import PreconditionFailedError from ..errors.unauthorized_error import UnauthorizedError @@ -1542,6 +1543,17 @@ def create_turn_events( ), ), ) + if _response.status_code == 500: + raise InternalServerError( + headers=dict(_response.headers), + body=typing.cast( + RequestErrorResponse, + construct_type( + type_=RequestErrorResponse, # type: ignore + object_=_response.json(), + ), + ), + ) _response_json = _response.json() except JSONDecodeError: raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text) @@ -3189,6 +3201,17 @@ async def create_turn_events( ), ), ) + if _response.status_code == 500: + raise InternalServerError( + headers=dict(_response.headers), + body=typing.cast( + RequestErrorResponse, + construct_type( + type_=RequestErrorResponse, # type: ignore + object_=_response.json(), + ), + ), + ) _response_json = _response.json() except JSONDecodeError: raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)