From 00516af3872bd3d6783dbded2d6f02985da1b15f Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 15 Oct 2026 09:32:06 +0530 Subject: [PATCH 01/32] feat: rebuild a paused turn from its own snapshot Resume skips turn.created and replays the unconsumed inbox. A new producer reseeds its sequence from the existing stream tail. Co-authored-by: Cursor --- .../src/agent-session/SessionHandle.ts | 105 ++++++++++++++++++ .../src/agent-session/TurnHandle.ts | 56 +++++++++- .../trueforge-core/src/agent-session/index.ts | 2 +- .../runtime/event-subscription/inMemory.ts | 15 +++ .../src/runtime/event-subscription/index.ts | 7 ++ .../src/runtime/event-subscription/redis.ts | 12 ++ 6 files changed, 191 insertions(+), 6 deletions(-) diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index f2567c7f1..e4a793734 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -20,6 +20,7 @@ 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. */ @@ -317,6 +318,110 @@ 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 queued before the handle + * is returned, so a later send() lands after them and stream() only executes. + */ + 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 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: turn.snapshot.sandbox_info ?? undefined, + previousTurn: turn, + signal: input.abortController.signal, + tracing, + }); + + const definitionsByThreadId = await this.resolveAgentDefinitions({ + previous: turn, + resolver: input.resolver, + spec, + tracing, + signal: input.abortController.signal, + }); + + 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: turn, + resolver: input.resolver, + spec, + sandbox, + tracing, + definitionsByThreadId, + }); + + const createDynamicSubAgentThread = this.makeCreateDynamicSubAgentThread({ + resolver: input.resolver, + previous: turn, + spec, + sandbox, + tracing, + }); + + const orchestrator = new AgentThreadOrchestrator({ + agentThreads, + createDynamicSubAgentThread, + tracing, + logger: input.resolver.logger, + }); + + const handle = new TurnHandle({ + store: this.store, + turn, + orchestrator, + resolver: input.resolver, + abortController: input.abortController, + resume: true, + }); + await handle.replayUnconsumedInbox(); + return handle; + } 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. diff --git a/packages/trueforge-core/src/agent-session/TurnHandle.ts b/packages/trueforge-core/src/agent-session/TurnHandle.ts index bf2a758b8..5f6e4f9b6 100644 --- a/packages/trueforge-core/src/agent-session/TurnHandle.ts +++ b/packages/trueforge-core/src/agent-session/TurnHandle.ts @@ -1,7 +1,7 @@ /** * Durable turn handle. {@link TurnHandle.stream} is execute-once (persist-before-yield). */ -import { AgentHarnessError } from '../core/errors'; +import { AgentHarnessError, InvalidAgentSendInputError } from '../core/errors'; import type { MCPAuthRequiredEvent, MCPServerInitInfo, @@ -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 => { @@ -172,6 +182,8 @@ export class TurnHandle> { private readonly resolver: ITurnResourceResolver | undefined; private readonly abortController: AbortController; private streamStarted = false; + /** Resume skips turn.created. Inbox replay happens before stream() is pulled. */ + private readonly resume: boolean; constructor(options: { store: ISessionStore; @@ -179,12 +191,14 @@ export class TurnHandle> { orchestrator?: AgentThreadOrchestrator | undefined; resolver?: ITurnResourceResolver | undefined; abortController?: AbortController | undefined; + resume?: boolean | undefined; }) { this.store = options.store; this.turn = options.turn; this.orchestrator = options.orchestrator; this.resolver = options.resolver; this.abortController = options.abortController ?? new AbortController(); + this.resume = options.resume ?? false; } /** Store-only handle (e.g. from {@link SessionHandle.getTurn}) — stream() is not available. */ @@ -254,7 +268,9 @@ export class TurnHandle> { private requireLiveOrchestrator(method: string): AgentThreadOrchestrator { if (!this.orchestrator) { - throw new Error(`TurnHandle.${method}() is only available on a live turn from SessionHandle.createTurn()`); + throw new Error( + `TurnHandle.${method}() is only available on a live turn from SessionHandle.createTurn() or SessionHandle.rebuildTurn()`, + ); } return this.orchestrator; } @@ -297,7 +313,9 @@ export class TurnHandle> { 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()'); + throw new Error( + 'TurnHandle.stream() is only available on turns returned from SessionHandle.createTurn() or SessionHandle.rebuildTurn()', + ); } // A consumer can leave early (break, or iterator.return()) while execute is parked @@ -319,7 +337,9 @@ export class TurnHandle> { let executeResult: AgentThreadExecutionResult | undefined; let ownershipLost: TurnOwnershipLostError | undefined; try { - yield await this.persistTurnCreated(); + if (!this.resume) { + yield await this.persistTurnCreated(); + } executeResult = yield* this.executeAndPersist(orchestrator, signal); } catch (error) { caughtError = error instanceof Error ? error : new Error(String(error)); @@ -331,7 +351,8 @@ export class TurnHandle> { } finally { // Terminal write stays in finally so a consumer break still closes the turn. let turnDone: TurnDoneEvent | undefined; - if (!ownershipLost) { + // Offload ends the in-memory run only. A terminal write would cancel a turn that is still paused. + if (!ownershipLost && !isTurnOffloadReason(signal.reason)) { try { // TurnNotRunningError: another writer already terminated the turn — defer to its state. if (caughtError instanceof TurnNotRunningError) { @@ -366,6 +387,31 @@ export class TurnHandle> { } } + /** + * Queue inbox rows onto the orchestrator before execute starts. + * A row the paused snapshot already applied fails validation; skip it so a + * crash between persist and consume does not apply the same decision twice. + */ + async replayUnconsumedInbox(): Promise { + const orchestrator = this.requireLiveOrchestrator('replayUnconsumedInbox'); + const rows = await this.store.listUnconsumedTurnInboundEvents({ + session_id: this.turn.session_id, + turn_id: this.turn.turn_id, + }); + for (const event of rows) { + try { + for await (const batch of orchestrator.send([event])) { + void batch; + } + } catch (error) { + if (error instanceof InvalidAgentSendInputError) { + continue; + } + throw error; + } + } + } + private async persistTurnCreated(): Promise { const turnCreated: TurnCreatedEvent = { type: EventType.TURN_CREATED, diff --git a/packages/trueforge-core/src/agent-session/index.ts b/packages/trueforge-core/src/agent-session/index.ts index 7b2b78b30..6221c0243 100644 --- a/packages/trueforge-core/src/agent-session/index.ts +++ b/packages/trueforge-core/src/agent-session/index.ts @@ -135,5 +135,5 @@ export type { TurnSandboxFactory } from './TurnResourceResolver'; export { SessionHandle } 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/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index 350f51222..2dbef9ac2 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. */ + rearm(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.rearm(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; From cd79652688759ef596e0271f1fb89d7b4c4c2828 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 15 Oct 2026 09:57:30 +0530 Subject: [PATCH 02/32] feat: offload a paused turn and steal it when the owner is gone Keep a paused run for a keepalive, then stop it without a terminal write. A later send rebuilds on the owner, or claims the turn only when the owner's heartbeat is missing. Co-authored-by: Cursor --- .changeset/turn-offload-steal.md | 6 + .../src/request-reply/client.ts | 4 +- .../src/request-reply/errors.ts | 7 +- .../trueforge-core/src/request-reply/index.ts | 8 +- packages/trueforge/.env.example | 2 + packages/trueforge/scripts/write-openapi.ts | 1 + packages/trueforge/src/apis/turns.ts | 279 +++++++++++++-- packages/trueforge/src/app.ts | 14 +- packages/trueforge/src/config.ts | 11 + packages/trueforge/src/main.ts | 38 +- packages/trueforge/src/routes/turnRoutes.ts | 4 + packages/trueforge/src/runtime/activeTurns.ts | 90 ++++- .../trueforge/src/runtime/turnDelivery.ts | 328 ++++++++++++++++++ .../trueforge/src/runtime/turnOwnership.ts | 25 +- .../tests/unit/apis/sessionHttp.test.ts | 9 +- .../tests/unit/runtime/turnOwnership.test.ts | 29 +- 16 files changed, 766 insertions(+), 89 deletions(-) create mode 100644 .changeset/turn-offload-steal.md create mode 100644 packages/trueforge/src/runtime/turnDelivery.ts 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/packages/trueforge-core/src/request-reply/client.ts b/packages/trueforge-core/src/request-reply/client.ts index b580928bf..187208a50 100644 --- a/packages/trueforge-core/src/request-reply/client.ts +++ b/packages/trueforge-core/src/request-reply/client.ts @@ -1,6 +1,6 @@ import { randomUUID } from 'node:crypto'; import { performance } from 'node:perf_hooks'; -import { NoResponderError, RequestTimeoutError } from './errors'; +import { 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 +71,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; diff --git a/packages/trueforge-core/src/request-reply/errors.ts b/packages/trueforge-core/src/request-reply/errors.ts index 05229d549..4fc455c17 100644 --- a/packages/trueforge-core/src/request-reply/errors.ts +++ b/packages/trueforge-core/src/request-reply/errors.ts @@ -1,9 +1,14 @@ +/** Heartbeat key was absent. */ +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 received the published request'; + /** 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..0b20d04d5 100644 --- a/packages/trueforge-core/src/request-reply/index.ts +++ b/packages/trueforge-core/src/request-reply/index.ts @@ -7,7 +7,13 @@ export { redisRequest } from './client'; export type { SendRequestOptions } from './client'; -export { NoResponderError, ReplyError, RequestTimeoutError } from './errors'; +export { + NO_RESPONDER_HEARTBEAT_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'; diff --git a/packages/trueforge/.env.example b/packages/trueforge/.env.example index 684ed8de2..b68b28e64 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 30. +# TURN_PAUSED_KEEPALIVE_SECONDS=30 ## 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 c8e17f8a1..128d02443 100644 --- a/packages/trueforge/scripts/write-openapi.ts +++ b/packages/trueforge/scripts/write-openapi.ts @@ -98,6 +98,7 @@ const app = createServerApp({ sessions: new Sessions({ sessionStore }), activeTurns: new ActiveTurnRegistry(), requestReplyRouter: new RequestReplyRouter(), + createPeerTurnExecutionDeps: () => Promise.reject(new Error('peer rebuild is not used while writing OpenAPI')), eventSubscriptions: new EventSubscriptionRegistry(undefined), logger: winston.createLogger({ silent: true }), oidcClient: undefined, diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index c9498ac93..5923e1b40 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -8,6 +8,7 @@ import { EventType, SessionStoreConflictError, SessionStoreNotFoundError, + TurnOffloadReason, TurnResourceResolver, type SessionHandle, type TurnHandle, @@ -30,6 +31,7 @@ import { VercelAILLM, withTimeout, } from '@truefoundry/trueforge-core/core'; +import type { JSONReply, RedisClient, RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; import type { Context } from 'hono'; import { HTTPException } from 'hono/http-exception'; import { streamSSE } from 'hono/streaming'; @@ -64,6 +66,12 @@ import { resolveSandboxEnvironment, resolveSandboxProvider, } from '../runtime/sessionResources'; +import { + deliverTurnEvents, + jsonValueFrom, + SendEventsPeerBodySchema, + TURNS_SEND_EVENTS_PATH, +} from '../runtime/turnDelivery'; import { DEFAULT_SANDBOX_ENVIRONMENT_NAME } from '../schemas/sandboxEnvironment'; import { MAX_SESSION_TITLE_LENGTH } from '../schemas/session'; import { assertGatewayMetadataRequestHeaders } from '../truefoundry/gatewayMetadata'; @@ -129,6 +137,16 @@ export interface TurnsRouterDeps { logger: Logger; resolveRequestContext: ResolveRequestContext; authorizer: Authorizer; + /** Present when this replica can forward events to the database owner. */ + redis?: RedisClient | undefined; + /** Present when this replica serves peer turn-event requests. */ + requestReplyRouter?: RequestReplyRouter | undefined; + /** + * Store bundle for a peer rebuild, which has no inbound HTTP context. + * The client-facing handler uses the request's own stores instead. + */ + createPeerTurnExecutionDeps?: + ((input: { session: SessionHandle; userRef: string }) => Promise) | undefined; } /** @@ -363,15 +381,23 @@ export async function drainTurnEvents(input: { turnId: string; maxExecutionTimer: NodeJS.Timeout; logger: Logger; + activeTurns?: ActiveTurnRegistry | undefined; onEvent?: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; }): Promise { - const { trackedStream, turnEventStream, sessionId, turnId, maxExecutionTimer, logger, onEvent } = input; + const { trackedStream, turnEventStream, sessionId, turnId, maxExecutionTimer, logger, activeTurns, 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 (activeTurns && event.type === EventType.TURN_UPDATE) { + if (event.state.status === 'paused') { + activeTurns.armPausedKeepalive({ sessionId, turnId }); + } else { + activeTurns.clearPausedKeepalive({ sessionId, turnId }); + } + } await onEvent?.(event, sequenceNumber); } } catch (error) { @@ -401,6 +427,7 @@ export interface TurnEventDrainInput { turnId: string; maxExecutionTimer: NodeJS.Timeout; logger: Logger; + activeTurns: ActiveTurnRegistry; } /** @@ -482,6 +509,7 @@ export async function beginTurnExecution( turnId: turn.id, maxExecutionTimer, logger: deps.logger, + activeTurns: deps.activeTurns, }, }; } @@ -508,6 +536,103 @@ export async function startTurnInProcess(params: BeginTurnExecutionParams): Prom return turn; } +/** + * 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; + userRef: string; + requestHeaders?: Record | undefined; + deps: BeginTurnExecutionDeps; +}): Promise<{ turn: TurnHandle; startDrain: () => void }> { + assertGatewayMetadataRequestHeaders(input.requestHeaders); + const abortController = new AbortController(); + const maxExecutionTimer = setTimeout(() => { + if (!abortController.signal.aborted) { + abortController.abort(CancellationReason.ServerExecutionTimeout); + } + }, configuration.SERVER_EXECUTION_TIMEOUT_SECONDS * 1000); + maxExecutionTimer.unref(); + + const sessionAgent = input.session.record.agent; + const webSearchProvider = await resolveWebSearchProvider({ + tenant_id: input.session.tenant_id, + store: input.deps.webSearchProviderStore, + }); + 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 turn: TurnHandle; + try { + turn = await input.session.rebuildTurn({ + turn_id: input.turnId, + active_executor_id: configuration.EXECUTOR_ID, + resolver, + abortController, + }); + } catch (error) { + clearTimeout(maxExecutionTimer); + 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) { + clearTimeout(maxExecutionTimer); + abortController.abort(TurnOffloadReason); + await turn.stream().return(undefined); + throw error; + } + + let drainStarted = false; + return { + turn, + startDrain: () => { + if (drainStarted) { + return; + } + drainStarted = true; + const trackedStream = input.deps.activeTurns.track({ + abortController, + stream: turn.stream(), + turn, + }); + void drainTurnEvents({ + trackedStream, + turnEventStream, + sessionId: input.session.session_id, + turnId: turn.id, + maxExecutionTimer, + logger: input.deps.logger, + activeTurns: input.deps.activeTurns, + }); + }, + }; +} + /** Mapped client error from turn execution; undefined means rethrow. */ export interface TurnExecutionError { status: 400 | 404 | 422; @@ -973,39 +1098,73 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { 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 turnHandle = deps.activeTurns.getTurnHandle({ sessionId, turnId }); - if (!turnHandle) { - return c.json({ error: { message: `Turn is not running on this server: ${turnId}` } }, 409); + const executionDeps = await executionDepsForRequest(c, session); + const result = await deliverTurnEvents({ + session, + sessionStore: deps.sessionStore, + activeTurns: deps.activeTurns, + turnId, + events, + redis: deps.redis, + logger: deps.logger, + allowSteal: true, + userRef: requestContext.subject.id, + tenantId: requestContext.tenant_id, + requestHeaders, + resume: () => + resumePausedTurn({ + session, + turnId, + userRef: requestContext.subject.id, + requestHeaders, + deps: executionDeps, + }), + }); + if (result.type === 'applied') { + return c.json({ data: result.events }, 201); } - - 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; + if (result.type === 'invalid') { + return c.json({ error: { message: result.message } }, 400); } - - return c.json({ data: events }, 201); + if (result.type === 'conflict') { + return c.json({ error: { message: result.message } }, 409); + } + const message = result.type === 'not_owner' ? `Turn owner is unavailable: ${turnId}` : result.message; + return c.json({ error: { message } }, 503); }); }; + async function executionDepsForRequest(c: Context, session: SessionHandle): Promise { + let referencedAgent: AgentRecord | undefined; + if (session.record.agent.type === 'reference') { + referencedAgent = await deps.resolveAgentStore(c).getAgent({ + tenant_id: session.tenant_id, + id: session.record.agent.id, + }); + } + return { + activeTurns: deps.activeTurns, + eventSubscriptions: deps.eventSubscriptions, + logger: deps.logger, + sandboxEnvironmentStore: deps.sandboxEnvironmentStore, + agentStore: deps.resolveAgentStore(c), + modelProviderStore: deps.resolveModelProviderStore(c, referencedAgent), + mcpServerStore: deps.resolveMcpServerStore(c, referencedAgent), + skillStore: deps.resolveSkillStore(c), + sandboxProviderStore: deps.resolveSandboxProviderStore(c), + webSearchProviderStore: deps.resolveWebSearchProviderStore(c), + }; + } + const router = new OpenAPIHono(); router.openapi(createAndExecuteTurnRoute, createAndExecuteTurnHandler); router.openapi(listTurnsRoute, listTurnsHandler); @@ -1014,5 +1173,69 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { router.openapi(listTurnEventsRoute, listTurnEventsHandler); router.openapi(postTurnEventsRoute, postTurnEventsHandler); router.openapi(subscribeTurnRoute, subscribeTurnHandler); + if (deps.requestReplyRouter) { + deps.requestReplyRouter.registerRoute(TURNS_SEND_EVENTS_PATH, request => handlePeerTurnEvents(deps, request)); + } return router; } + +async function handlePeerTurnEvents(deps: TurnsRouterDeps, request: { body: unknown }): Promise { + const parsed = SendEventsPeerBodySchema.safeParse(request.body); + if (!parsed.success) { + return { status: 400, body: jsonValueFrom({ message: 'Invalid turns/send-events payload' }) }; + } + const createPeerTurnExecutionDeps = deps.createPeerTurnExecutionDeps; + if (!createPeerTurnExecutionDeps) { + return { status: 503, body: jsonValueFrom({ message: 'Turn owner is unavailable' }) }; + } + const session = await deps.sessions.get({ + tenant_id: parsed.data.tenant_id, + session_id: parsed.data.session_id, + }); + if (!session) { + return { status: 404, body: jsonValueFrom({ message: `Session not found: ${parsed.data.session_id}` }) }; + } + const result = await deps.activeTurns.withTurnLock( + { sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }, + async () => { + const executionDeps = await createPeerTurnExecutionDeps({ + session, + userRef: parsed.data.user_ref, + }); + return deliverTurnEvents({ + session, + sessionStore: deps.sessionStore, + activeTurns: deps.activeTurns, + turnId: parsed.data.turn_id, + events: parsed.data.events, + redis: undefined, + logger: deps.logger, + allowSteal: false, + userRef: parsed.data.user_ref, + tenantId: parsed.data.tenant_id, + requestHeaders: parsed.data.request_headers, + resume: () => + resumePausedTurn({ + session, + turnId: parsed.data.turn_id, + userRef: parsed.data.user_ref, + requestHeaders: parsed.data.request_headers, + deps: executionDeps, + }), + }); + }, + ); + if (result.type === 'applied') { + return { status: 200, body: jsonValueFrom({ events: result.events }) }; + } + if (result.type === 'invalid') { + return { status: 400, body: jsonValueFrom({ message: result.message }) }; + } + if (result.type === 'not_owner') { + return { status: 412, body: jsonValueFrom({ message: 'Turn is not running on this executor' }) }; + } + if (result.type === 'conflict') { + return { status: 409, body: jsonValueFrom({ message: result.message }) }; + } + return { status: 503, body: jsonValueFrom({ message: result.message }) }; +} diff --git a/packages/trueforge/src/app.ts b/packages/trueforge/src/app.ts index 1c0b0533e..99bae1a99 100644 --- a/packages/trueforge/src/app.ts +++ b/packages/trueforge/src/app.ts @@ -1,7 +1,12 @@ /** The API: resource routers, the OpenAPI document and Swagger UI, all under /api/v1. */ import { swaggerUI } from '@hono/swagger-ui'; import { OpenAPIHono, z } from '@hono/zod-openapi'; -import type { ISessionStore, Sessions, TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session'; +import type { + ISessionStore, + SessionHandle, + Sessions, + TurnStreamingEvent, +} from '@truefoundry/trueforge-core/agent-session'; import { extractErrorLogFields } from '@truefoundry/trueforge-core/core'; import type { RedisClient, RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; import type { Context, ErrorHandler, MiddlewareHandler } from 'hono'; @@ -24,7 +29,7 @@ import { createInternalMetricsRouter } from './apis/sessionMetrics'; import { createInternalSessionsRouter, createSessionsRouter } from './apis/sessions'; import { createSettingsRouter } from './apis/settings'; import { createAvailableSkillsRouter, type ResolveSkillStore } from './apis/skills'; -import { createTurnsRouter } from './apis/turns'; +import { createTurnsRouter, type BeginTurnExecutionDeps } from './apis/turns'; import type { Authenticator } from './auth/authenticator'; import type { Authorizer } from './auth/authorizer'; import { resolveRequestContext } from './auth/identity'; @@ -220,6 +225,8 @@ export interface ServerDeps { redis?: RedisClient | undefined; /** Request-reply dispatch table served by this replica's executor. */ requestReplyRouter: RequestReplyRouter; + /** Store bundle for rebuilding a paused turn on a peer request, which has no HTTP context. */ + createPeerTurnExecutionDeps: (input: { session: SessionHandle; userRef: string }) => Promise; /** Hands out each turn's resumable event stream to the create and subscribe handlers. */ eventSubscriptions: EventSubscriptionRegistry; logger: Logger; @@ -504,6 +511,9 @@ export function createServerApp(deps: ServerDeps) { logger: deps.logger, resolveRequestContext, authorizer: deps.authorizer, + redis: deps.redis, + requestReplyRouter: deps.requestReplyRouter, + createPeerTurnExecutionDeps: deps.createPeerTurnExecutionDeps, }), authMiddleware, ), diff --git a/packages/trueforge/src/config.ts b/packages/trueforge/src/config.ts index daacd98fc..ec4b7c0eb 100644 --- a/packages/trueforge/src/config.ts +++ b/packages/trueforge/src/config.ts @@ -772,6 +772,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 30. + */ + 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. @@ -1163,6 +1169,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: 30, + }), 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/main.ts b/packages/trueforge/src/main.ts index a62c53c8c..5decb143d 100644 --- a/packages/trueforge/src/main.ts +++ b/packages/trueforge/src/main.ts @@ -81,6 +81,7 @@ import { CancellationReason, Sessions, type ISessionStore, + type SessionHandle, type TurnStreamingEvent, } from '@truefoundry/trueforge-core/agent-session'; import { RequestReplyExecutor, RequestReplyRouter, type RedisClient } from '@truefoundry/trueforge-core/request-reply'; @@ -90,7 +91,12 @@ import type { Logger } from 'winston'; 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 { + requestContextFromCreatedBySubject, + resolveRequestContext, + STANDALONE_REQUEST_CONTEXT, + type RequestContext, +} from './auth/identity'; import { initOidc } from './auth/oidc'; import { McpCatalog } from './catalog/McpCatalog'; import { ModelCatalog } from './catalog/ModelCatalog'; @@ -590,7 +596,9 @@ async function createServerRuntime(persistence: ServerPersistence< serviceFoundryClient, } = persistence; - const activeTurns = new ActiveTurnRegistry(); + const activeTurns = new ActiveTurnRegistry({ + pausedKeepaliveMs: configuration.TURN_PAUSED_KEEPALIVE_SECONDS * 1000, + }); const requestReplyRouter = new RequestReplyRouter(); const eventSubscriptions = new EventSubscriptionRegistry(redis?.client); const sessions = new Sessions({ sessionStore }); @@ -679,6 +687,31 @@ async function createServerRuntime(persistence: ServerPersistence< } return new InlineSkillStore({ inner: store, inline: parseInlineSkills(rawInline) }); }; + const createPeerTurnExecutionDeps = async (input: { session: SessionHandle; userRef: string }) => { + const requestContext = requestContextFromCreatedBySubject({ + tenant_id: input.session.tenant_id, + created_by_subject: input.session.record.created_by_subject, + }); + let referencedAgent: AgentRecord | undefined; + if (input.session.record.agent.type === 'reference') { + referencedAgent = await agentStore.getAgent({ + tenant_id: input.session.tenant_id, + id: input.session.record.agent.id, + }); + } + return { + activeTurns, + eventSubscriptions, + logger, + sandboxEnvironmentStore, + agentStore, + modelProviderStore: persistence.resolveModelProviderStore(requestContext, referencedAgent), + mcpServerStore: persistence.resolveMcpServerStore(requestContext, referencedAgent), + skillStore: turnSkillsResolverStore, + sandboxProviderStore: persistence.resolveSandboxProviderStore(requestContext), + webSearchProviderStore: persistence.resolveWebSearchProviderStore(requestContext), + }; + }; const app = createServerApp({ modelCatalog: ModelCatalog.load(), mcpCatalog: McpCatalog.load(), @@ -704,6 +737,7 @@ async function createServerRuntime(persistence: ServerPersistence< activeTurns, redis: redis?.client, requestReplyRouter, + createPeerTurnExecutionDeps, eventSubscriptions, logger, oidcClient, diff --git a/packages/trueforge/src/routes/turnRoutes.ts b/packages/trueforge/src/routes/turnRoutes.ts index a4fa5d827..e1b7612a3 100644 --- a/packages/trueforge/src/routes/turnRoutes.ts +++ b/packages/trueforge/src/routes/turnRoutes.ts @@ -209,6 +209,10 @@ export const postTurnEventsRoute = createRoute({ content: { 'application/json': { schema: RequestErrorResponseSchema } }, description: 'Turn is terminal, or an event id already exists.', }, + 503: { + content: { 'application/json': { schema: RequestErrorResponseSchema } }, + description: 'Turn owner is unavailable. Retry.', + }, }, }); diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 1586a5fe2..eac81644c 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -4,9 +4,11 @@ * 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'; +const DEFAULT_PAUSED_KEEPALIVE_MS = 30_000; + interface ActiveTurnRun { turn: TurnHandle; abortController: AbortController; @@ -18,12 +20,26 @@ function activeTurnKey(sessionId: string, turnId: string): string { return `${sessionId}:${turnId}`; } +/** A test double may omit `state`. A real handle always has it. */ +function isPausedRun(run: ActiveTurnRun | undefined): run is ActiveTurnRun { + if (!run || !Object.hasOwn(run.turn, 'state')) { + return false; + } + 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(); + private readonly keepaliveTimers = new Map(); + private readonly pausedKeepaliveMs: number; private alreadyShutDownAbortReason: CancellationReason | undefined; + constructor(options?: { pausedKeepaliveMs?: number }) { + this.pausedKeepaliveMs = options?.pausedKeepaliveMs ?? DEFAULT_PAUSED_KEEPALIVE_MS; + } + /** * 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. @@ -76,6 +92,7 @@ export class ActiveTurnRegistry { } 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); @@ -125,21 +142,80 @@ 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. + * 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. + */ + armPausedKeepalive(input: { sessionId: string; turnId: string }): void { + this.clearPausedKeepalive(input); + if (this.pausedKeepaliveMs <= 0 || this.alreadyShutDownAbortReason !== undefined) { + return; + } + const key = activeTurnKey(input.sessionId, input.turnId); + const timer = setTimeout(() => { + this.keepaliveTimers.delete(key); + void this.withTurnLock(input, async () => { + const run = this.runs.get(key); + if (!isPausedRun(run)) { + return; + } + await this.finishRun(run, TurnOffloadReason); + }); + }, this.pausedKeepaliveMs); + timer.unref(); + this.keepaliveTimers.set(key, timer); + } + + clearPausedKeepalive(input: { sessionId: string; turnId: string }): void { + const key = activeTurnKey(input.sessionId, input.turnId); + const timer = this.keepaliveTimers.get(key); + if (!timer) { + return; + } + clearTimeout(timer); + this.keepaliveTimers.delete(key); + } + + /** + * Abort the local run and wait until it leaves the registry. + * Caller must already hold the per-turn lock. Offload does not write a terminal state. + */ + async dropLocalTurn(input: { sessionId: string; turnId: string }): Promise { + const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); + if (!run) { + return; + } + await this.finishRun(run, TurnOffloadReason); + } + + /** + * Enter shutdown mode and wait until the registry is empty. + * Paused runs are offloaded so their approval state survives. Running runs abort as `abortReason`. + * Late `track()` calls abort immediately as `abortReason`. */ async shutdownAndWait(abortReason: CancellationReason): Promise { this.alreadyShutDownAbortReason = abortReason; + for (const timer of this.keepaliveTimers.values()) { + clearTimeout(timer); + } + this.keepaliveTimers.clear(); 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 reason = isPausedRun(run) ? TurnOffloadReason : abortReason; + run.abortController.abort(reason); } 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/turnDelivery.ts b/packages/trueforge/src/runtime/turnDelivery.ts new file mode 100644 index 000000000..04f5477e2 --- /dev/null +++ b/packages/trueforge/src/runtime/turnDelivery.ts @@ -0,0 +1,328 @@ +/** + * Sends inbound events to the turn that owns them: the local machine, a rebuild + * of a paused turn this replica owns, or a peer. Steal is a conditional claim + * after the owner's heartbeat is gone. + */ +import type { + ISessionStore, + SessionHandle, + TurnHandle, + TurnUserEvent, +} from '@truefoundry/trueforge-core/agent-session'; +import { + TurnEventAlreadyExistsError, + TurnOwnershipLostError, + TurnUserEventSchema, +} from '@truefoundry/trueforge-core/agent-session'; +import { AgentHarnessError } from '@truefoundry/trueforge-core/core'; +import { + NO_RESPONDER_PUBLISH_DETAIL, + NoResponderError, + redisRequest, + RequestTimeoutError, + type JSONValue, + type RedisClient, +} from '@truefoundry/trueforge-core/request-reply'; +import type { Logger } from 'winston'; +import { z } from 'zod'; +import configuration from '../config'; +import type { ActiveTurnRegistry } from './activeTurns'; +import { PeerOutcome, resolveOwnershipAction, TurnOwnershipAction } from './turnOwnership'; + +/** Request-reply path the owner serves for inbound turn events. */ +export const TURNS_SEND_EVENTS_PATH = 'turns/send-events'; + +const MAX_OWNERSHIP_ATTEMPTS = 4; + +export const SendEventsPeerBodySchema = z.object({ + tenant_id: z.string(), + session_id: z.string(), + turn_id: z.string(), + user_ref: z.string(), + events: z.array(TurnUserEventSchema), + request_headers: z.record(z.string(), z.string()).optional(), +}); + +const PeerEventsReplySchema = z.object({ + events: z.array(TurnUserEventSchema), +}); + +const PeerMessageSchema = z.object({ + message: z.string(), +}); + +function isJsonValue(value: unknown): value is JSONValue { + if (value === null || typeof value === 'string' || typeof value === 'boolean') { + return true; + } + if (typeof value === 'number') { + return Number.isFinite(value); + } + if (Array.isArray(value)) { + return value.every(isJsonValue); + } + if (typeof value === 'object') { + return Object.values(value).every(isJsonValue); + } + return false; +} + +/** JSON.stringify drops `undefined`, which the Redis envelope cannot carry. */ +export function jsonValueFrom(value: unknown): JSONValue { + const parsed: unknown = JSON.parse(JSON.stringify(value)); + if (!isJsonValue(parsed)) { + throw new Error('Turn event payload is not JSON'); + } + return parsed; +} + +export type TurnDeliveryResult = + | { type: 'applied'; events: TurnUserEvent[] } + | { type: 'invalid'; message: string } + | { type: 'conflict'; message: string } + | { type: 'unavailable'; message: string } + | { type: 'not_owner' }; + +export async function deliverTurnEvents(input: { + session: SessionHandle; + sessionStore: Pick; + activeTurns: ActiveTurnRegistry; + turnId: string; + events: TurnUserEvent[]; + redis: RedisClient | undefined; + logger: Pick; + /** + * Rebuild the paused turn this replica owns. Do not start execute yet: + * the caller queues this request's events first, then starts the drain. + */ + resume: () => Promise<{ turn: TurnHandle; startDrain: () => void }>; + /** Peer delivery is owner-local. Only the replica that received the client request may steal. */ + allowSteal: boolean; + userRef: string; + tenantId: string; + requestHeaders?: Record | undefined; +}): Promise { + const sessionId = input.session.session_id; + + for (let attempt = 0; attempt < MAX_OWNERSHIP_ATTEMPTS; attempt++) { + const loaded = await input.session.getTurn(input.turnId); + if (!loaded) { + return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; + } + const status = loaded.state.status; + const owner = loaded.record.active_executor_id; + const ownedByThisExecutor = owner === configuration.EXECUTOR_ID; + + if (!ownedByThisExecutor) { + await input.activeTurns.dropLocalTurn({ sessionId, turnId: input.turnId }); + } + + if (!input.allowSteal && !ownedByThisExecutor) { + return { type: 'not_owner' }; + } + + const hasLocalActiveTurn = input.activeTurns.getTurnHandle({ sessionId, turnId: input.turnId }) !== undefined; + let action = resolveOwnershipAction({ + status, + ownedByThisExecutor, + hasLocalActiveTurn, + peerOutcome: null, + }); + + if (action === TurnOwnershipAction.Forward) { + const forwarded = await forwardTurnEvents({ + redis: input.redis, + owner, + sessionId, + turnId: input.turnId, + events: input.events, + userRef: input.userRef, + tenantId: input.tenantId, + requestHeaders: input.requestHeaders, + logger: input.logger, + }); + if (forwarded.type !== 'peer') { + return forwarded; + } + const fresh = await input.session.getTurn(input.turnId); + const freshStatus = fresh?.state.status ?? status; + const freshOwner = fresh?.record.active_executor_id ?? owner; + if (forwarded.peerOutcome === PeerOutcome.NoResponder && (freshStatus !== 'paused' || freshOwner !== owner)) { + continue; + } + action = resolveOwnershipAction({ + status: freshStatus, + ownedByThisExecutor: freshOwner === configuration.EXECUTOR_ID, + hasLocalActiveTurn: false, + peerOutcome: forwarded.peerOutcome, + }); + if ( + action === TurnOwnershipAction.Forward || + action === TurnOwnershipAction.Run || + action === TurnOwnershipAction.Rebuild + ) { + continue; + } + } + + if (action === TurnOwnershipAction.Run) { + const handle = input.activeTurns.getTurnHandle({ sessionId, turnId: input.turnId }); + if (!handle) { + continue; + } + return applyTurnEvents({ + handle, + events: input.events, + sessionStore: input.sessionStore, + sessionId, + turnId: input.turnId, + }); + } + + if (action === TurnOwnershipAction.Rebuild || action === TurnOwnershipAction.Steal) { + if (action === TurnOwnershipAction.Steal) { + const won = await input.sessionStore.claimTurnExecutor({ + session_id: sessionId, + turn_id: input.turnId, + expected_executor_id: owner, + new_executor_id: configuration.EXECUTOR_ID, + }); + if (!won) { + continue; + } + } + try { + const resumed = await input.resume(); + try { + return await applyTurnEvents({ + handle: resumed.turn, + events: input.events, + sessionStore: input.sessionStore, + sessionId, + turnId: input.turnId, + }); + } finally { + resumed.startDrain(); + } + } catch (error) { + if (error instanceof TurnOwnershipLostError) { + continue; + } + throw error; + } + } + + if (action === TurnOwnershipAction.Reject) { + return { + type: 'conflict', + message: + status === 'running' + ? `Turn is running but its executor is gone: ${input.turnId}` + : `Turn is not accepting events: ${input.turnId}`, + }; + } + + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } + + return { type: 'unavailable', message: `Turn ownership did not settle: ${input.turnId}` }; +} + +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, + 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 }; +} + +type ForwardResult = TurnDeliveryResult | { type: 'peer'; peerOutcome: PeerOutcome }; + +async function forwardTurnEvents(input: { + redis: RedisClient | undefined; + owner: string; + sessionId: string; + turnId: string; + events: TurnUserEvent[]; + userRef: string; + tenantId: string; + requestHeaders?: Record | undefined; + logger: Pick; +}): Promise { + if (!input.redis) { + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } + const body = SendEventsPeerBodySchema.parse({ + tenant_id: input.tenantId, + session_id: input.sessionId, + turn_id: input.turnId, + user_ref: input.userRef, + events: input.events, + ...(input.requestHeaders === undefined ? {} : { request_headers: input.requestHeaders }), + }); + try { + const reply = await redisRequest({ + redis: input.redis, + executorId: input.owner, + 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, + }, + }); + if (reply.status === 200) { + const parsed = PeerEventsReplySchema.safeParse(reply.body); + if (!parsed.success) { + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } + return { type: 'applied', events: parsed.data.events }; + } + if (reply.status === 412) { + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } + if (reply.status === 409 || reply.status === 400) { + const parsed = PeerMessageSchema.safeParse(reply.body); + return { + type: reply.status === 400 ? 'invalid' : 'conflict', + message: parsed.success ? parsed.data.message : `Turn is not accepting events: ${input.turnId}`, + }; + } + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } catch (error) { + if (error instanceof NoResponderError) { + if (error.detail === NO_RESPONDER_PUBLISH_DETAIL) { + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } + return { type: 'peer', peerOutcome: PeerOutcome.NoResponder }; + } + if (error instanceof RequestTimeoutError) { + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } + input.logger.warn('Failed to reach owning executor for turn events', { + sessionId: input.sessionId, + turnId: input.turnId, + owner: input.owner, + }); + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } +} diff --git a/packages/trueforge/src/runtime/turnOwnership.ts b/packages/trueforge/src/runtime/turnOwnership.ts index 5d7c5e4a5..07d148c3f 100644 --- a/packages/trueforge/src/runtime/turnOwnership.ts +++ b/packages/trueforge/src/runtime/turnOwnership.ts @@ -1,17 +1,14 @@ import type { TurnState } from '@truefoundry/trueforge-core/agent-session'; -/** Result of asking the owning replica to handle a turn request. */ +/** + * Result of a forward that did not itself finish the request. + * A peer 200, 400, or 409 is returned directly and never becomes an outcome. + */ export enum PeerOutcome { - /** Peer replied 200 and handled the request. */ - Ok = 'ok', - /** Heartbeat missing or nobody subscribed: the owner is gone. */ + /** Heartbeat missing: 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', + /** Publish-zero, timeout, 412, or transport failure. The owner is not known to be gone. */ + OwnerNotGone = 'owner_not_gone', } export enum TurnOwnershipAction { @@ -21,8 +18,6 @@ export enum TurnOwnershipAction { 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). */ @@ -62,13 +57,9 @@ export function resolveOwnershipAction(input: TurnOwnershipInput): TurnOwnership 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: + case PeerOutcome.OwnerNotGone: return TurnOwnershipAction.Retry; } } diff --git a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts index 678ef7619..0351ef498 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'; @@ -480,7 +481,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 send = jest.fn().mockImplementation(async function* (events: Parameters[0]) { yield events; }); @@ -553,7 +556,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/runtime/turnOwnership.test.ts b/packages/trueforge/tests/unit/runtime/turnOwnership.test.ts index 04672786f..8663d02bd 100644 --- a/packages/trueforge/tests/unit/runtime/turnOwnership.test.ts +++ b/packages/trueforge/tests/unit/runtime/turnOwnership.test.ts @@ -47,11 +47,6 @@ const ROWS: Row[] = [ 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: { @@ -73,32 +68,12 @@ const ROWS: Row[] = [ 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)', + name: 'paused + forward that does not show a missing heartbeat does not steal', input: { status: 'paused', ownedByThisExecutor: false, hasLocalActiveTurn: false, - peerOutcome: PeerOutcome.TransportError, + peerOutcome: PeerOutcome.OwnerNotGone, }, expected: TurnOwnershipAction.Retry, }, From c3f9816a94fbd7db444c6c77489ffcf28519215f Mon Sep 17 00:00:00 2001 From: "trueforge-dev-bot[bot]" Date: Thu, 15 Oct 2026 04:30:21 +0000 Subject: [PATCH 03/32] Regenerate OpenAPI document and SDKs --- .github/fern/openapi/openapi.json | 10 +++++++ docs/openapi.json | 10 +++++++ .../src/api/errors/ServiceUnavailableError.ts | 24 +++++++++++++++ .../trueforge-sdk/src/api/errors/index.ts | 1 + .../api/resources/sessions/client/Client.ts | 12 ++++++++ .../trueforge-sdk/tests/wire/sessions.test.ts | 29 +++++++++++++++++++ .../src/trueforge_sdk/__init__.py | 3 ++ .../src/trueforge_sdk/errors/__init__.py | 3 ++ .../errors/service_unavailable_error.py | 11 +++++++ .../src/trueforge_sdk/sessions/raw_client.py | 23 +++++++++++++++ 10 files changed, 126 insertions(+) create mode 100644 packages/trueforge-sdk/src/api/errors/ServiceUnavailableError.ts create mode 100644 python/trueforge_sdk/src/trueforge_sdk/errors/service_unavailable_error.py diff --git a/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index 83a233ac8..f6d12504b 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -10037,6 +10037,16 @@ } }, "description": "Turn is terminal, or an event id already exists." + }, + "503": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RequestErrorResponse" + } + } + }, + "description": "Turn owner is unavailable. Retry." } }, "summary": "Create turn events", diff --git a/docs/openapi.json b/docs/openapi.json index 83a233ac8..f6d12504b 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -10037,6 +10037,16 @@ } }, "description": "Turn is terminal, or an event id already exists." + }, + "503": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RequestErrorResponse" + } + } + }, + "description": "Turn owner is unavailable. Retry." } }, "summary": "Create turn events", diff --git a/packages/trueforge-sdk/src/api/errors/ServiceUnavailableError.ts b/packages/trueforge-sdk/src/api/errors/ServiceUnavailableError.ts new file mode 100644 index 000000000..58e0399d1 --- /dev/null +++ b/packages/trueforge-sdk/src/api/errors/ServiceUnavailableError.ts @@ -0,0 +1,24 @@ +// This file was auto-generated by Fern from our API Definition. + +import type * as core from "../../core/index.js"; +import * as errors from "../../errors/index.js"; +import type * as TrueForge from "../index.js"; + +export class ServiceUnavailableError extends errors.TrueForgeError { + public declare readonly body: TrueForge.RequestErrorResponse; + + constructor(body: TrueForge.RequestErrorResponse, rawResponse?: core.RawResponse) { + super({ + message: "ServiceUnavailableError", + statusCode: 503, + body: body, + rawResponse: rawResponse, + }); + Object.setPrototypeOf(this, new.target.prototype); + if (Error.captureStackTrace) { + Error.captureStackTrace(this, this.constructor); + } + + this.name = "ServiceUnavailableError"; + } +} diff --git a/packages/trueforge-sdk/src/api/errors/index.ts b/packages/trueforge-sdk/src/api/errors/index.ts index 5388d473c..e47bb54e0 100644 --- a/packages/trueforge-sdk/src/api/errors/index.ts +++ b/packages/trueforge-sdk/src/api/errors/index.ts @@ -8,5 +8,6 @@ export * from "./GoneError.js"; export * from "./InternalServerError.js"; export * from "./NotFoundError.js"; export * from "./PreconditionFailedError.js"; +export * from "./ServiceUnavailableError.js"; export * from "./UnauthorizedError.js"; export * from "./UnprocessableEntityError.js"; 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..5894fe5d8 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.ServiceUnavailableError} * @throws {@link errors.TrueForgeError} * @throws {@link errors.TrueForgeTimeoutError} * @@ -1872,6 +1873,17 @@ export class SessionsClient { }), _response.rawResponse, ); + case 503: + throw new TrueForge.ServiceUnavailableError( + 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..173a9feee 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(503) + .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.ServiceUnavailableError); + }); + test("subscribe_to_turn (1)", async () => { const server = mockServerPool.createServer(); const client = new TrueForge({ maxRetries: 0, token: "test", baseUrl: server.baseUrl }); diff --git a/python/trueforge_sdk/src/trueforge_sdk/__init__.py b/python/trueforge_sdk/src/trueforge_sdk/__init__.py index 7801ca100..96ec2727b 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/__init__.py +++ b/python/trueforge_sdk/src/trueforge_sdk/__init__.py @@ -287,6 +287,7 @@ InternalServerError, NotFoundError, PreconditionFailedError, + ServiceUnavailableError, UnauthorizedError, UnprocessableEntityError, ) @@ -506,6 +507,7 @@ "ScheduleRun": ".types", "ScheduleRunStatus": ".types", "ScheduleStatus": ".types", + "ServiceUnavailableError": ".errors", "Session": ".types", "SessionAgent": ".types", "SessionAgentInline": ".types", @@ -828,6 +830,7 @@ def __dir__(): "ScheduleRun", "ScheduleRunStatus", "ScheduleStatus", + "ServiceUnavailableError", "Session", "SessionAgent", "SessionAgentInline", diff --git a/python/trueforge_sdk/src/trueforge_sdk/errors/__init__.py b/python/trueforge_sdk/src/trueforge_sdk/errors/__init__.py index 07f9aa6be..ca21b04f3 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/errors/__init__.py +++ b/python/trueforge_sdk/src/trueforge_sdk/errors/__init__.py @@ -16,6 +16,7 @@ from .internal_server_error import InternalServerError from .not_found_error import NotFoundError from .precondition_failed_error import PreconditionFailedError + from .service_unavailable_error import ServiceUnavailableError from .unauthorized_error import UnauthorizedError from .unprocessable_entity_error import UnprocessableEntityError _dynamic_imports: typing.Dict[str, str] = { @@ -29,6 +30,7 @@ "InternalServerError": ".internal_server_error", "NotFoundError": ".not_found_error", "PreconditionFailedError": ".precondition_failed_error", + "ServiceUnavailableError": ".service_unavailable_error", "UnauthorizedError": ".unauthorized_error", "UnprocessableEntityError": ".unprocessable_entity_error", } @@ -66,6 +68,7 @@ def __dir__(): "InternalServerError", "NotFoundError", "PreconditionFailedError", + "ServiceUnavailableError", "UnauthorizedError", "UnprocessableEntityError", ] diff --git a/python/trueforge_sdk/src/trueforge_sdk/errors/service_unavailable_error.py b/python/trueforge_sdk/src/trueforge_sdk/errors/service_unavailable_error.py new file mode 100644 index 000000000..2cb49bd7c --- /dev/null +++ b/python/trueforge_sdk/src/trueforge_sdk/errors/service_unavailable_error.py @@ -0,0 +1,11 @@ +# This file was auto-generated by Fern from our API Definition. + +import typing + +from ..core.api_error import ApiError +from ..types.request_error_response import RequestErrorResponse + + +class ServiceUnavailableError(ApiError): + def __init__(self, body: RequestErrorResponse, headers: typing.Optional[typing.Dict[str, str]] = None): + super().__init__(status_code=503, headers=headers, body=body) 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..8a0252b1f 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py @@ -27,6 +27,7 @@ from ..errors.gone_error import GoneError from ..errors.not_found_error import NotFoundError from ..errors.precondition_failed_error import PreconditionFailedError +from ..errors.service_unavailable_error import ServiceUnavailableError from ..errors.unauthorized_error import UnauthorizedError from ..errors.unprocessable_entity_error import UnprocessableEntityError from ..types.cancel_session_response import CancelSessionResponse @@ -1542,6 +1543,17 @@ def create_turn_events( ), ), ) + if _response.status_code == 503: + raise ServiceUnavailableError( + 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 == 503: + raise ServiceUnavailableError( + 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) From 6b589c7da4dc3b5f4b5132b4909868717baed3df Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 15 Oct 2026 18:41:51 +0530 Subject: [PATCH 04/32] fix: emit turn.done when the next turn supersedes a live run A rebuilt turn starts paused so the first runnable batch persists running. Peer resume applies inline skills from the forwarded request. Creating the next turn aborts the previous run so its drain writes turn.done. Co-authored-by: Cursor --- .../src/agent-session/SessionHandle.ts | 5 +- .../src/agent-session/TurnHandle.ts | 23 ++- .../trueforge-core/src/agent-session/index.ts | 2 +- .../core/runtime/AgentThreadOrchestrator.ts | 14 +- packages/trueforge/src/apis/sessions.ts | 159 +----------------- packages/trueforge/src/apis/turns.ts | 35 +++- packages/trueforge/src/app.ts | 6 +- packages/trueforge/src/main.ts | 19 ++- packages/trueforge/src/runtime/activeTurns.ts | 17 ++ packages/trueforge/src/runtime/cancelTurn.ts | 149 ++++++++++++++++ .../src/truefoundry/InlineSkillStore.ts | 56 +++++- .../src/truefoundry/gatewayMetadata.ts | 2 +- 12 files changed, 307 insertions(+), 180 deletions(-) create mode 100644 packages/trueforge/src/runtime/cancelTurn.ts diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index e4a793734..36d34061c 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -26,7 +26,10 @@ 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; } diff --git a/packages/trueforge-core/src/agent-session/TurnHandle.ts b/packages/trueforge-core/src/agent-session/TurnHandle.ts index 5f6e4f9b6..f3c7e72dd 100644 --- a/packages/trueforge-core/src/agent-session/TurnHandle.ts +++ b/packages/trueforge-core/src/agent-session/TurnHandle.ts @@ -188,6 +188,7 @@ export class TurnHandle> { constructor(options: { store: ISessionStore; turn: TurnRecord; + // orchestrator and resolver are optional for store-only handles (e.g. from SessionHandle.getTurn) orchestrator?: AgentThreadOrchestrator | undefined; resolver?: ITurnResourceResolver | undefined; abortController?: AbortController | undefined; @@ -260,19 +261,13 @@ 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() or SessionHandle.rebuildTurn()`, + `TurnHandle.send() is only available on a live turn. Perhaps this turn handle was created for read only actions?`, ); } - return this.orchestrator; + yield* this.orchestrator.send(events); + this.orchestrator.wake(); } /** @@ -393,14 +388,18 @@ export class TurnHandle> { * crash between persist and consume does not apply the same decision twice. */ async replayUnconsumedInbox(): Promise { - const orchestrator = this.requireLiveOrchestrator('replayUnconsumedInbox'); + if (!this.orchestrator) { + throw new Error( + `TurnHandle.replayUnconsumedInbox() is only available on a live turn. Perhaps this turn handle was created for read only actions?`, + ); + } const rows = await this.store.listUnconsumedTurnInboundEvents({ session_id: this.turn.session_id, turn_id: this.turn.turn_id, }); for (const event of rows) { try { - for await (const batch of orchestrator.send([event])) { + for await (const batch of this.orchestrator.send([event])) { void batch; } } catch (error) { @@ -436,7 +435,7 @@ export class TurnHandle> { orchestrator: AgentThreadOrchestrator, signal: AbortSignal, ): AsyncGenerator { - const generator = orchestrator.execute({ signal }); + const generator = orchestrator.execute({ signal, resume: this.resume }); try { let iterResult = await generator.next(); while (!iterResult.done) { diff --git a/packages/trueforge-core/src/agent-session/index.ts b/packages/trueforge-core/src/agent-session/index.ts index 6221c0243..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, TurnOffloadReason, isTurnOffloadReason } from './TurnHandle'; export type { TurnStreamingEvent } from './TurnHandle'; diff --git a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts index cdb122ab1..bda49c7ce 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts @@ -566,8 +566,14 @@ export class AgentThreadOrchestrator { public async *execute({ signal, + resume = false, }: { signal: AbortSignal; + /** + * The stored turn is already paused. Start there so the first runnable batch + * persists `running`; a fresh execute would skip that transition. + */ + resume?: boolean; }): AsyncGenerator { const { agentThreads } = this; @@ -575,7 +581,7 @@ export class AgentThreadOrchestrator { let output: ModelMessageEvent | null = null; let rootAgentError: AgentThreadExecutionResult['root_agent_error']; let rootFinished = false; - let turnPaused = false; + let turnPaused = resume; const mainThread = this.getMainThread(); const rootSpan = createRootAgentSpan(mainThread, this.tracing); @@ -602,8 +608,10 @@ export class AgentThreadOrchestrator { return done(); } - // Apply accepted user events in the same order they entered the turn-level queue. - for (const event of this.pendingTurnEvents.splice(0)) { + // Take the queue, then empty it, so a send during this pass waits for the next loop. + const queued = this.pendingTurnEvents.slice(); + this.pendingTurnEvents.length = 0; + for (const event of queued) { switch (event.type) { case EventType.USER_TOOL_APPROVAL_POLICY: yield* this.applyApprovalPolicies(event); diff --git a/packages/trueforge/src/apis/sessions.ts b/packages/trueforge/src/apis/sessions.ts index 0931ab5be..a9b75d6d4 100644 --- a/packages/trueforge/src/apis/sessions.ts +++ b/packages/trueforge/src/apis/sessions.ts @@ -2,29 +2,17 @@ * DB-backed sessions APIs (mounted at /api/v1/sessions and /api/internal/sessions). */ import { OpenAPIHono, type RouteHandler } from '@hono/zod-openapi'; -import type { ISessionStore, SessionHandle, SessionRecord, Sessions } from '@truefoundry/trueforge-core/agent-session'; +import type { ISessionStore, SessionRecord, Sessions } from '@truefoundry/trueforge-core/agent-session'; import { - CancellationReason, - isNonTerminalTurnState, SessionStoreConflictError, SessionStoreInvariantError, SessionStoreNotFoundError, - TurnNotFoundError, } from '@truefoundry/trueforge-core/agent-session'; -import { extractErrorLogFields } from '@truefoundry/trueforge-core/core'; -import { - redisRequest, - RequestTimeoutError, - type RedisClient, - type RouteHandler as RequestReplyRouteHandler, - type RequestReplyRouter, -} from '@truefoundry/trueforge-core/request-reply'; +import { type RedisClient, type RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; import type { Context } from 'hono'; import type { Logger } from 'winston'; -import { z } from 'zod'; import type { Authorizer } from '../auth/authorizer'; import { createdBySubjectFromRequestContext, type ResolveRequestContext } from '../auth/identity'; -import configuration from '../config'; import type { IAgentStore } from '../db/agentStore'; import type { IMcpServerStore } from '../db/mcpServerStore'; import type { IModelProviderStore } from '../db/modelProviderStore'; @@ -42,6 +30,7 @@ import { updateSessionRoute, } from '../routes/sessionRoutes'; import type { ActiveTurnRegistry } from '../runtime/activeTurns'; +import { cancelSessionTurn, cancelSessionTurnPeerHandler, SESSIONS_CANCEL_PATH } from '../runtime/cancelTurn'; import { validateAgentSpec } from '../runtime/sessionResources'; import { honoQueriesToRecord } from '../schemas/deepObjectQuery'; import { isSessionAgentNameRef, parseListSessionsQuery, type Session } from '../schemas/session'; @@ -49,16 +38,8 @@ import { newId } from '../utils/id'; import { agentIfAccessible, canReadAgentBoundResource, canReadSession, resolveManagedAgentIds } from './agentAccess'; import type { ResolveSkillStore } from './skills'; -/** Request-reply path a replica serves to cancel a turn it owns. */ -export const SESSIONS_CANCEL_PATH = 'sessions/cancel'; - -/** Wire body of a peer cancel; validated on receipt (it crosses processes via Redis). */ -const CancelPeerBodySchema = z.object({ - session_id: z.string(), - turn_id: z.string(), - reason: z.enum(CancellationReason), -}); -type CancelPeerBody = z.infer; +export { cancelSessionTurn, cancelSessionTurnPeerHandler, SESSIONS_CANCEL_PATH } from '../runtime/cancelTurn'; +export type { CancelTurnDeps } from '../runtime/cancelTurn'; export function toWireSession(record: SessionRecord): Session { return { @@ -93,136 +74,6 @@ export interface SessionsRouterDeps { 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, - }); -} - -/** - * 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). - */ -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, { - sessionId: parsed.data.session_id, - turnId: parsed.data.turn_id, - reason: parsed.data.reason, - }); - if (!found) { - return Promise.resolve({ status: 412, body: { message: 'Turn is not running on this executor' } }); - } - return Promise.resolve({ status: 200, body: {} }); - }; -} - -/** A registry to abort in, a session to freeze, durable state to read, and a way to reach peers. */ -export interface CancelTurnDeps { - activeTurns: ActiveTurnRegistry; - session: Pick; - sessionStore: Pick; - redis?: RedisClient | undefined; - logger: Pick; -} - -/** - * Cancels a `running` or `paused` turn wherever it lives: locally or on the - * 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`. - * - * Redis timeout and Redis/transport failures are not a clean cancellation — - * the owning replica may still be executing — but the turn is still frozen. - * Later writes from that replica lose to first-terminal-write-wins. - */ -export async function cancelSessionTurn( - deps: CancelTurnDeps, - input: { turnId: string; reason?: CancellationReason }, -): Promise { - const { turnId, reason = CancellationReason.ClientCancelled } = input; - const sessionId = deps.session.session_id; - - const turn = await deps.sessionStore.getTurn({ - session_id: sessionId, - turn_id: turnId, - }); - if (!turn || !isNonTerminalTurnState(turn.state)) { - return; - } - - const owner = turn.active_executor_id; - // Without a Redis client there is no peer to ask, so a different owner falls - // through to the local lookup and freezes if the run is gone. - if (owner !== configuration.EXECUTOR_ID && deps.redis) { - try { - const reply = await redisRequest({ - redis: deps.redis, - executorId: owner, - path: SESSIONS_CANCEL_PATH, - request: { - body: { session_id: sessionId, turn_id: turnId, reason }, - }, - options: { - replyTimeoutMs: configuration.REDIS_REQUEST_REPLY_TIMEOUT_MS, - pollIntervalMs: configuration.REDIS_REQUEST_REPLY_POLL_INTERVAL_MS, - }, - }); - if (reply.status === 200) { - return; - } - } catch (error) { - const fields = { - sessionId, - turnId, - owner, - ...extractErrorLogFields(error), - }; - if (error instanceof RequestTimeoutError) { - deps.logger.warn('Timed out waiting for owning executor to cancel; freezing the turn', fields); - } else { - deps.logger.warn('Failed to reach owning executor over Redis; freezing the turn', fields); - } - } - await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); - return; - } - - const aborted = cancelTurnOnThisExecutor(deps.activeTurns, { sessionId, turnId, reason }); - if (!aborted) { - await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); - } -} - -/** Freeze a non-terminal turn; missing turns are a no-op (already gone). */ -async function freezeTurnIgnoringMissing( - session: Pick, - input: { turnId: string; reason: CancellationReason }, -): Promise { - try { - await session.freezeTurn({ turn_id: input.turnId, reason: input.reason }); - } catch (error) { - if (error instanceof TurnNotFoundError) { - return; - } - throw error; - } -} - const FORBIDDEN_SESSION_ACCESS = 'Only the session creator can access this session'; function isSessionOwner({ diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 5923e1b40..12efc65c3 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -6,6 +6,7 @@ import type { ISessionStore, Sessions, Turn, TurnStreamingEvent } from '@truefou import { CancellationReason, EventType, + resolvePreviousTurnId, SessionStoreConflictError, SessionStoreNotFoundError, TurnOffloadReason, @@ -57,6 +58,7 @@ import { subscribeTurnRoute, } from '../routes/turnRoutes'; import type { ActiveTurnRegistry } from '../runtime/activeTurns'; +import { cancelSessionTurn } from '../runtime/cancelTurn'; import { StreamGoneError, type EventSubscription, type EventSubscriptionRegistry } from '../runtime/event-subscription'; import { validateSandboxFilePath } from '../runtime/sandboxFilePath'; import { @@ -146,7 +148,12 @@ export interface TurnsRouterDeps { * The client-facing handler uses the request's own stores instead. */ createPeerTurnExecutionDeps?: - ((input: { session: SessionHandle; userRef: string }) => Promise) | undefined; + | ((input: { + session: SessionHandle; + userRef: string; + requestHeaders?: Record | undefined; + }) => Promise) + | undefined; } /** @@ -164,6 +171,9 @@ export type BeginTurnExecutionDeps = Pick< sandboxProviderStore: ISandboxProviderStore; webSearchProviderStore: IWebSearchProviderStore; skillStore: Pick; + /** Present on the HTTP create path so a previous turn can be superseded on its owner. */ + sessionStore?: Pick | undefined; + redis?: RedisClient | undefined; }; interface BeginTurnExecutionParams { @@ -474,6 +484,28 @@ export async function beginTurnExecution( // never overwrites an existing title. const title = session.record.last_turn_id ? undefined : deriveSessionTitle(input); + // Unpark a live previous turn so its drain emits turn.done before this turn freezes the row. + const previousId = resolvePreviousTurnId(previousTurnId, session.record.last_turn_id); + if (previousId) { + const hadLocalRun = await deps.activeTurns.abortAndWait({ + sessionId, + turnId: previousId, + abortReason: CancellationReason.CancelledForNextTurn, + }); + if (!hadLocalRun && deps.sessionStore) { + await cancelSessionTurn( + { + activeTurns: deps.activeTurns, + session, + sessionStore: deps.sessionStore, + redis: deps.redis, + logger: deps.logger, + }, + { turnId: previousId, reason: CancellationReason.CancelledForNextTurn }, + ); + } + } + const turn = await session.createTurn({ turn_id: turnId, active_executor_id: configuration.EXECUTOR_ID, @@ -1201,6 +1233,7 @@ async function handlePeerTurnEvents(deps: TurnsRouterDeps, request: { body: unkn const executionDeps = await createPeerTurnExecutionDeps({ session, userRef: parsed.data.user_ref, + requestHeaders: parsed.data.request_headers, }); return deliverTurnEvents({ session, diff --git a/packages/trueforge/src/app.ts b/packages/trueforge/src/app.ts index 99bae1a99..582bea981 100644 --- a/packages/trueforge/src/app.ts +++ b/packages/trueforge/src/app.ts @@ -226,7 +226,11 @@ export interface ServerDeps { /** Request-reply dispatch table served by this replica's executor. */ requestReplyRouter: RequestReplyRouter; /** Store bundle for rebuilding a paused turn on a peer request, which has no HTTP context. */ - createPeerTurnExecutionDeps: (input: { session: SessionHandle; userRef: string }) => Promise; + createPeerTurnExecutionDeps: (input: { + session: SessionHandle; + userRef: string; + requestHeaders?: Record | undefined; + }) => Promise; /** Hands out each turn's resumable event stream to the create and subscribe handlers. */ eventSubscriptions: EventSubscriptionRegistry; logger: Logger; diff --git a/packages/trueforge/src/main.ts b/packages/trueforge/src/main.ts index 5decb143d..1cd9aa26a 100644 --- a/packages/trueforge/src/main.ts +++ b/packages/trueforge/src/main.ts @@ -129,6 +129,7 @@ import { EventSubscriptionRegistry } from './runtime/event-subscription'; import type { ConnectedRedis } from './runtime/redis'; import { initSentry } from './sentry'; import { printStandaloneStartupBanner } from './startupBanner'; +import { headerValue } from './truefoundry/gatewayMetadata'; import { InlineMcpServerStore } from './truefoundry/InlineMcpServerStore'; import { InlineModelProviderStore } from './truefoundry/InlineModelProviderStore'; import { @@ -139,7 +140,7 @@ import { X_TFG_MODELS, X_TFG_SKILLS, } from './truefoundry/inlineResources'; -import { InlineSkillStore } from './truefoundry/InlineSkillStore'; +import { InlineSkillStore, overlayInlineTurnSkills } from './truefoundry/InlineSkillStore'; import { parsePerServerMcpHeaders, X_TFG_MCP_HEADERS, @@ -687,7 +688,11 @@ async function createServerRuntime(persistence: ServerPersistence< } return new InlineSkillStore({ inner: store, inline: parseInlineSkills(rawInline) }); }; - const createPeerTurnExecutionDeps = async (input: { session: SessionHandle; userRef: string }) => { + const createPeerTurnExecutionDeps = async (input: { + session: SessionHandle; + userRef: string; + requestHeaders?: Record | undefined; + }) => { const requestContext = requestContextFromCreatedBySubject({ tenant_id: input.session.tenant_id, created_by_subject: input.session.record.created_by_subject, @@ -699,6 +704,14 @@ async function createServerRuntime(persistence: ServerPersistence< id: input.session.record.agent.id, }); } + const rawInlineSkills = headerValue(input.requestHeaders, X_TFG_SKILLS); + const skillStore = + isTrueFoundryModeEnabled(configuration) && rawInlineSkills !== undefined + ? overlayInlineTurnSkills({ + resolver: turnSkillsResolverStore, + inline: parseInlineSkills(rawInlineSkills), + }) + : turnSkillsResolverStore; return { activeTurns, eventSubscriptions, @@ -707,7 +720,7 @@ async function createServerRuntime(persistence: ServerPersistence< agentStore, modelProviderStore: persistence.resolveModelProviderStore(requestContext, referencedAgent), mcpServerStore: persistence.resolveMcpServerStore(requestContext, referencedAgent), - skillStore: turnSkillsResolverStore, + skillStore, sandboxProviderStore: persistence.resolveSandboxProviderStore(requestContext), webSearchProviderStore: persistence.resolveWebSearchProviderStore(requestContext), }; diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index eac81644c..382c6fe72 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -125,6 +125,23 @@ export class ActiveTurnRegistry { return true; } + /** + * Abort a local run and wait until its drain has finished. + * A paused generator is parked on the abort latch, so this unparks it and lets + * the drain emit `turn.done`. Returns false when this process has no run. + */ + async abortAndWait(input: { sessionId: string; turnId: string; abortReason: CancellationReason }): Promise { + const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); + if (!run) { + return false; + } + if (!run.abortController.signal.aborted) { + run.abortController.abort(input.abortReason); + } + await run.waitUntilCompleted; + return true; + } + /** * 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 diff --git a/packages/trueforge/src/runtime/cancelTurn.ts b/packages/trueforge/src/runtime/cancelTurn.ts new file mode 100644 index 000000000..0554d303a --- /dev/null +++ b/packages/trueforge/src/runtime/cancelTurn.ts @@ -0,0 +1,149 @@ +import type { ISessionStore, SessionHandle } from '@truefoundry/trueforge-core/agent-session'; +import { + CancellationReason, + isNonTerminalTurnState, + TurnNotFoundError, +} from '@truefoundry/trueforge-core/agent-session'; +import { extractErrorLogFields } from '@truefoundry/trueforge-core/core'; +import { + redisRequest, + RequestTimeoutError, + type RedisClient, + type RouteHandler as RequestReplyRouteHandler, +} from '@truefoundry/trueforge-core/request-reply'; +import type { Logger } from 'winston'; +import { z } from 'zod'; +import configuration from '../config'; +import type { ActiveTurnRegistry } from './activeTurns'; + +/** Request-reply path a replica serves to cancel a turn it owns. */ +export const SESSIONS_CANCEL_PATH = 'sessions/cancel'; + +/** Wire body of a peer cancel; validated on receipt (it crosses processes via Redis). */ +const CancelPeerBodySchema = z.object({ + session_id: z.string(), + turn_id: z.string(), + reason: z.enum(CancellationReason), +}); +type CancelPeerBody = z.infer; + +/** + * Peer-facing cancel handler. 200 means the local run has finished, so its drain + * has emitted `turn.done`. 412 means it is not running here. + */ +export function cancelSessionTurnPeerHandler(activeTurns: ActiveTurnRegistry): RequestReplyRouteHandler { + return async request => { + const parsed = CancelPeerBodySchema.safeParse(request.body); + if (!parsed.success) { + return { status: 400, body: { message: 'Invalid sessions/cancel payload' } }; + } + const found = await activeTurns.abortAndWait({ + sessionId: parsed.data.session_id, + turnId: parsed.data.turn_id, + abortReason: parsed.data.reason, + }); + if (!found) { + return { status: 412, body: { message: 'Turn is not running on this executor' } }; + } + return { status: 200, body: {} }; + }; +} + +/** A registry to abort in, a session to freeze, durable state to read, and a way to reach peers. */ +export interface CancelTurnDeps { + activeTurns: ActiveTurnRegistry; + session: Pick; + sessionStore: Pick; + redis?: RedisClient | undefined; + logger: Pick; +} + +/** + * Cancels a `running` or `paused` turn wherever it lives: locally or on the + * 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`. + * + * Redis timeout and Redis/transport failures are not a clean cancellation — + * the owning replica may still be executing — but the turn is still frozen. + * Later writes from that replica lose to first-terminal-write-wins. + */ +export async function cancelSessionTurn( + deps: CancelTurnDeps, + input: { turnId: string; reason?: CancellationReason }, +): Promise { + const { turnId, reason = CancellationReason.ClientCancelled } = input; + const sessionId = deps.session.session_id; + + const turn = await deps.sessionStore.getTurn({ + session_id: sessionId, + turn_id: turnId, + }); + if (!turn || !isNonTerminalTurnState(turn.state)) { + return; + } + + const owner = turn.active_executor_id; + // Without a Redis client there is no peer to ask, so a different owner falls + // through to the local lookup and freezes if the run is gone. + if (owner !== configuration.EXECUTOR_ID && deps.redis) { + try { + const reply = await redisRequest({ + redis: deps.redis, + executorId: owner, + path: SESSIONS_CANCEL_PATH, + request: { + body: { session_id: sessionId, turn_id: turnId, reason }, + }, + options: { + replyTimeoutMs: configuration.REDIS_REQUEST_REPLY_TIMEOUT_MS, + pollIntervalMs: configuration.REDIS_REQUEST_REPLY_POLL_INTERVAL_MS, + }, + }); + if (reply.status === 200) { + return; + } + } catch (error) { + const fields = { + sessionId, + turnId, + owner, + ...extractErrorLogFields(error), + }; + if (error instanceof RequestTimeoutError) { + deps.logger.warn('Timed out waiting for owning executor to cancel; freezing the turn', fields); + } else { + deps.logger.warn('Failed to reach owning executor over Redis; freezing the turn', fields); + } + } + await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); + return; + } + + const aborted = deps.activeTurns.cancelIfRunning({ + sessionId, + turnId, + abortReason: reason, + }); + if (!aborted) { + await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); + } +} + +/** Freeze a non-terminal turn; missing turns are a no-op (already gone). */ +async function freezeTurnIgnoringMissing( + session: Pick, + input: { turnId: string; reason: CancellationReason }, +): Promise { + try { + await session.freezeTurn({ turn_id: input.turnId, reason: input.reason }); + } catch (error) { + if (error instanceof TurnNotFoundError) { + return; + } + throw error; + } +} diff --git a/packages/trueforge/src/truefoundry/InlineSkillStore.ts b/packages/trueforge/src/truefoundry/InlineSkillStore.ts index f37aaca53..f0e0de51d 100644 --- a/packages/trueforge/src/truefoundry/InlineSkillStore.ts +++ b/packages/trueforge/src/truefoundry/InlineSkillStore.ts @@ -13,12 +13,62 @@ import type { import type { SkillVersion } from '../schemas/skill'; import type { InlineSkills } from './inlineResources'; +/** + * Overlay request-scoped skills on a resolver that only implements `resolveTurnSkills`. + * Registry skills still go to that resolver. The other skill APIs are unused on this path. + */ +export function overlayInlineTurnSkills(input: { + resolver: Pick, 'resolveTurnSkills'>; + inline: InlineSkills; +}): Pick, 'resolveTurnSkills'> { + return new InlineSkillStore({ + inner: new ResolveOnlySkillStore(input.resolver), + inline: input.inline, + }); +} + +class ResolveOnlySkillStore implements ISkillStore { + constructor(private readonly resolver: Pick, 'resolveTurnSkills'>) {} + + resolveTurnSkills(input: AgentSkillsInput): Promise { + return this.resolver.resolveTurnSkills(input); + } + + listSkills(): Promise { + return Promise.resolve([]); + } + + createSkill(): Promise { + return unsupportedSkillWrite(); + } + + upsertSkill(): Promise { + return unsupportedSkillWrite(); + } + + deleteSkill(): Promise { + return unsupportedSkillWrite(); + } + + listSkillVersions(): Promise { + return Promise.resolve([]); + } + + validateAgentSkills(): Promise { + return unsupportedSkillWrite(); + } +} + +function unsupportedSkillWrite(): Promise { + return Promise.reject(new Error('Peer skill resolution only resolves turn skills')); +} + /** * Serves the skills a request brought with it, and delegates everything else. * - * Mirrors {@link InlineMcpServerStore}: name-filtered list and validate/resolve are overlaid, an - * unfiltered list passes through so request-scoped skills stay out of the tenant's settings, and - * writes delegate because there is no row to write. + * Name-filtered list and validate/resolve are overlaid, an unfiltered list passes through so + * request-scoped skills stay out of the tenant's settings, and writes delegate because there is + * no row to write. */ export class InlineSkillStore implements ISkillStore { readonly #inner: ISkillStore; 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; } From 9c564d751c30c787413f086a60bca35be0433e84 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Fri, 16 Oct 2026 09:49:54 +0530 Subject: [PATCH 05/32] fix: drop the peer inline-skill overlay A peer rebuild has no caller token, and the admin skill resolver is not a full skill store, so wrapping x-tfg-skills there needed a stub. Leave skills on that resolver until the header wrap is shared with models and MCP. Co-authored-by: Cursor --- packages/trueforge/src/main.ts | 27 ++++------ .../src/truefoundry/InlineSkillStore.ts | 50 ------------------- 2 files changed, 9 insertions(+), 68 deletions(-) diff --git a/packages/trueforge/src/main.ts b/packages/trueforge/src/main.ts index 1cd9aa26a..6617fe355 100644 --- a/packages/trueforge/src/main.ts +++ b/packages/trueforge/src/main.ts @@ -129,7 +129,6 @@ import { EventSubscriptionRegistry } from './runtime/event-subscription'; import type { ConnectedRedis } from './runtime/redis'; import { initSentry } from './sentry'; import { printStandaloneStartupBanner } from './startupBanner'; -import { headerValue } from './truefoundry/gatewayMetadata'; import { InlineMcpServerStore } from './truefoundry/InlineMcpServerStore'; import { InlineModelProviderStore } from './truefoundry/InlineModelProviderStore'; import { @@ -140,7 +139,7 @@ import { X_TFG_MODELS, X_TFG_SKILLS, } from './truefoundry/inlineResources'; -import { InlineSkillStore, overlayInlineTurnSkills } from './truefoundry/InlineSkillStore'; +import { InlineSkillStore } from './truefoundry/InlineSkillStore'; import { parsePerServerMcpHeaders, X_TFG_MCP_HEADERS, @@ -215,14 +214,14 @@ function buildResolveSkillStore(options: { client: TrueFoundryServiceFoundryServerClient | undefined; }): (requestContext: RequestContext) => ISkillStore { const { persistenceStore, client } = options; - if (client === undefined) { - return () => persistenceStore; + if (client) { + return requestContext => + new TrueFoundrySkillStore({ + client, + context: requestContext, + }); } - return requestContext => - new TrueFoundrySkillStore({ - client, - context: requestContext, - }); + return () => persistenceStore; } function buildTurnSkillsResolverStore(options: { @@ -704,14 +703,6 @@ async function createServerRuntime(persistence: ServerPersistence< id: input.session.record.agent.id, }); } - const rawInlineSkills = headerValue(input.requestHeaders, X_TFG_SKILLS); - const skillStore = - isTrueFoundryModeEnabled(configuration) && rawInlineSkills !== undefined - ? overlayInlineTurnSkills({ - resolver: turnSkillsResolverStore, - inline: parseInlineSkills(rawInlineSkills), - }) - : turnSkillsResolverStore; return { activeTurns, eventSubscriptions, @@ -720,7 +711,7 @@ async function createServerRuntime(persistence: ServerPersistence< agentStore, modelProviderStore: persistence.resolveModelProviderStore(requestContext, referencedAgent), mcpServerStore: persistence.resolveMcpServerStore(requestContext, referencedAgent), - skillStore, + skillStore: turnSkillsResolverStore, sandboxProviderStore: persistence.resolveSandboxProviderStore(requestContext), webSearchProviderStore: persistence.resolveWebSearchProviderStore(requestContext), }; diff --git a/packages/trueforge/src/truefoundry/InlineSkillStore.ts b/packages/trueforge/src/truefoundry/InlineSkillStore.ts index f0e0de51d..a54a0e206 100644 --- a/packages/trueforge/src/truefoundry/InlineSkillStore.ts +++ b/packages/trueforge/src/truefoundry/InlineSkillStore.ts @@ -13,56 +13,6 @@ import type { import type { SkillVersion } from '../schemas/skill'; import type { InlineSkills } from './inlineResources'; -/** - * Overlay request-scoped skills on a resolver that only implements `resolveTurnSkills`. - * Registry skills still go to that resolver. The other skill APIs are unused on this path. - */ -export function overlayInlineTurnSkills(input: { - resolver: Pick, 'resolveTurnSkills'>; - inline: InlineSkills; -}): Pick, 'resolveTurnSkills'> { - return new InlineSkillStore({ - inner: new ResolveOnlySkillStore(input.resolver), - inline: input.inline, - }); -} - -class ResolveOnlySkillStore implements ISkillStore { - constructor(private readonly resolver: Pick, 'resolveTurnSkills'>) {} - - resolveTurnSkills(input: AgentSkillsInput): Promise { - return this.resolver.resolveTurnSkills(input); - } - - listSkills(): Promise { - return Promise.resolve([]); - } - - createSkill(): Promise { - return unsupportedSkillWrite(); - } - - upsertSkill(): Promise { - return unsupportedSkillWrite(); - } - - deleteSkill(): Promise { - return unsupportedSkillWrite(); - } - - listSkillVersions(): Promise { - return Promise.resolve([]); - } - - validateAgentSkills(): Promise { - return unsupportedSkillWrite(); - } -} - -function unsupportedSkillWrite(): Promise { - return Promise.reject(new Error('Peer skill resolution only resolves turn skills')); -} - /** * Serves the skills a request brought with it, and delegates everything else. * From 5a4ce8d7c375a547683375946b8572f0441613a8 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Fri, 16 Oct 2026 09:49:57 +0530 Subject: [PATCH 06/32] feat: keep a paused turn in memory for five minutes Thirty seconds offloads a paused run before the next user event typically arrives. Co-authored-by: Cursor --- packages/trueforge/.env.example | 4 ++-- packages/trueforge/src/config.ts | 4 ++-- packages/trueforge/src/runtime/activeTurns.ts | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/trueforge/.env.example b/packages/trueforge/.env.example index b68b28e64..afa73b235 100644 --- a/packages/trueforge/.env.example +++ b/packages/trueforge/.env.example @@ -115,8 +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 30. -# TURN_PAUSED_KEEPALIVE_SECONDS=30 +## 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/src/config.ts b/packages/trueforge/src/config.ts index ec4b7c0eb..8f57d75be 100644 --- a/packages/trueforge/src/config.ts +++ b/packages/trueforge/src/config.ts @@ -775,7 +775,7 @@ export interface SharedServerConfiguration { /** * 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 30. + * Env: `TURN_PAUSED_KEEPALIVE_SECONDS`. Default 300. */ TURN_PAUSED_KEEPALIVE_SECONDS: number; /** @@ -1172,7 +1172,7 @@ const shared: SharedServerConfiguration = { TURN_PAUSED_KEEPALIVE_SECONDS: parsePositiveInt({ envKey: 'TURN_PAUSED_KEEPALIVE_SECONDS', raw: getEnv('TURN_PAUSED_KEEPALIVE_SECONDS'), - defaultValue: 30, + defaultValue: 300, }), TURN_SSE_STREAM_WRITE_TIMEOUT_MS: parsePositiveInt({ envKey: 'TURN_SSE_STREAM_WRITE_TIMEOUT_MS', diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 382c6fe72..4a0aff6ae 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -7,7 +7,7 @@ import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; import { Mutex } from 'async-mutex'; -const DEFAULT_PAUSED_KEEPALIVE_MS = 30_000; +const DEFAULT_PAUSED_KEEPALIVE_MS = 300_000; interface ActiveTurnRun { turn: TurnHandle; From 64e7d3d72b2e4028242dc3e496130c6f798a9c08 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Mon, 19 Oct 2026 22:58:22 +0530 Subject: [PATCH 07/32] fix: send the caller request context when peering a paused turn The owner was rebuilding model and MCP stores as the session creator with no credential. Carry the caller's context on the peer body, and let schedule dispatch pass an explicit null credential. Co-authored-by: Cursor --- packages/trueforge/src/apis/schedules.ts | 1 + packages/trueforge/src/apis/turns.ts | 21 ++++++------ packages/trueforge/src/app.ts | 4 +-- packages/trueforge/src/auth/identity.ts | 11 ++++++- packages/trueforge/src/main.ts | 14 ++------ .../trueforge/src/runtime/turnDelivery.ts | 33 ++++++++++++++----- .../tests/unit/auth/identity.test.ts | 1 + 7 files changed, 51 insertions(+), 34 deletions(-) diff --git a/packages/trueforge/src/apis/schedules.ts b/packages/trueforge/src/apis/schedules.ts index ff70519db..cdd817dbb 100644 --- a/packages/trueforge/src/apis/schedules.ts +++ b/packages/trueforge/src/apis/schedules.ts @@ -206,6 +206,7 @@ export function createScheduleExecutionRouter(deps: ScheduleTurnEx requestContextFromCreatedBySubject({ tenant_id: item.schedule.tenant_id, created_by_subject: item.schedule.created_by_subject, + user_credential: null, }), ); await startScheduleRunOnRequest({ c, item, deps }); diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 42dbc1861..4ffd95426 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -39,7 +39,7 @@ import { streamSSE } from 'hono/streaming'; import type { Logger } from 'winston'; import type { ResolveMcpServerStoreInput, ResolveModelProviderStoreInput, ResolveTurnSkillsStoreInput } from '../app'; import type { Authorizer } from '../auth/authorizer'; -import type { ResolveRequestContext } from '../auth/identity'; +import { requestContextFromWire, type RequestContext, type ResolveRequestContext } from '../auth/identity'; import configuration from '../config'; import type { AgentRecord, IAgentStore } from '../db/agentStore'; import type { IMcpServerWithAuthStore } from '../db/mcpServerStore'; @@ -151,7 +151,7 @@ export interface TurnsRouterDeps { createPeerTurnExecutionDeps?: | ((input: { session: SessionHandle; - userRef: string; + requestContext: RequestContext; requestHeaders?: Record | undefined; }) => Promise) | undefined; @@ -603,7 +603,7 @@ export async function startTurnInProcess(params: BeginTurnExecutionParams): Prom async function resumePausedTurn(input: { session: SessionHandle; turnId: string; - userRef: string; + requestContext: RequestContext; requestHeaders?: Record | undefined; deps: BeginTurnExecutionDeps; }): Promise<{ turn: TurnHandle; startDrain: () => void }> { @@ -631,7 +631,7 @@ async function resumePausedTurn(input: { webSearchProvider, logger: input.deps.logger, signal: abortController.signal, - userRef: input.userRef, + userRef: input.requestContext.subject.id, session: input.session, turnMetadata: { sessionId: input.session.session_id, @@ -1187,14 +1187,13 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { redis: deps.redis, logger: deps.logger, allowSteal: true, - userRef: requestContext.subject.id, - tenantId: requestContext.tenant_id, + requestContext, requestHeaders, resume: () => resumePausedTurn({ session, turnId, - userRef: requestContext.subject.id, + requestContext, requestHeaders, deps: executionDeps, }), @@ -1277,9 +1276,10 @@ async function handlePeerTurnEvents(deps: TurnsRouterDeps, request: { body: unkn const result = await deps.activeTurns.withTurnLock( { sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }, async () => { + const requestContext = requestContextFromWire(parsed.data.request_context); const executionDeps = await createPeerTurnExecutionDeps({ session, - userRef: parsed.data.user_ref, + requestContext, requestHeaders: parsed.data.request_headers, }); return deliverTurnEvents({ @@ -1291,14 +1291,13 @@ async function handlePeerTurnEvents(deps: TurnsRouterDeps, request: { body: unkn redis: undefined, logger: deps.logger, allowSteal: false, - userRef: parsed.data.user_ref, - tenantId: parsed.data.tenant_id, + requestContext, requestHeaders: parsed.data.request_headers, resume: () => resumePausedTurn({ session, turnId: parsed.data.turn_id, - userRef: parsed.data.user_ref, + requestContext, requestHeaders: parsed.data.request_headers, deps: executionDeps, }), diff --git a/packages/trueforge/src/app.ts b/packages/trueforge/src/app.ts index 9dfbc60f1..3f39b003e 100644 --- a/packages/trueforge/src/app.ts +++ b/packages/trueforge/src/app.ts @@ -32,7 +32,7 @@ import { createAvailableSkillsRouter } from './apis/skills'; import { createTurnsRouter, type BeginTurnExecutionDeps } 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, @@ -280,7 +280,7 @@ export interface ServerDeps { /** Store bundle for rebuilding a paused turn on a peer request, which has no HTTP context. */ createPeerTurnExecutionDeps: (input: { session: SessionHandle; - userRef: string; + requestContext: RequestContext; requestHeaders?: Record | undefined; }) => Promise; /** Hands out each turn's resumable event stream to the create and subscribe handlers. */ 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/main.ts b/packages/trueforge/src/main.ts index 4cc30bf02..9c1941828 100644 --- a/packages/trueforge/src/main.ts +++ b/packages/trueforge/src/main.ts @@ -97,12 +97,7 @@ import { } from './app'; import { TrueForgeAuthorizer, type Authorizer } from './auth/authorizer'; import { createAuthenticator } from './auth/createAuthenticator'; -import { - requestContextFromCreatedBySubject, - resolveRequestContext, - STANDALONE_REQUEST_CONTEXT, - type RequestContext, -} from './auth/identity'; +import { resolveRequestContext, STANDALONE_REQUEST_CONTEXT, type RequestContext } from './auth/identity'; import { initOidc } from './auth/oidc'; import { McpCatalog } from './catalog/McpCatalog'; import { ModelCatalog } from './catalog/ModelCatalog'; @@ -771,13 +766,10 @@ async function createServerRuntime(persistence: ServerPersistence< }; const createPeerTurnExecutionDeps = async (input: { session: SessionHandle; - userRef: string; + requestContext: RequestContext; requestHeaders?: Record | undefined; }) => { - const requestContext = requestContextFromCreatedBySubject({ - tenant_id: input.session.tenant_id, - created_by_subject: input.session.record.created_by_subject, - }); + const { requestContext } = input; let referencedAgent: AgentRecord | undefined; if (input.session.record.agent.type === 'reference') { referencedAgent = await agentStore.getAgent({ diff --git a/packages/trueforge/src/runtime/turnDelivery.ts b/packages/trueforge/src/runtime/turnDelivery.ts index 04f5477e2..b921b6e2c 100644 --- a/packages/trueforge/src/runtime/turnDelivery.ts +++ b/packages/trueforge/src/runtime/turnDelivery.ts @@ -25,6 +25,7 @@ import { } from '@truefoundry/trueforge-core/request-reply'; import type { Logger } from 'winston'; import { z } from 'zod'; +import type { RequestContext } from '../auth/identity'; import configuration from '../config'; import type { ActiveTurnRegistry } from './activeTurns'; import { PeerOutcome, resolveOwnershipAction, TurnOwnershipAction } from './turnOwnership'; @@ -34,12 +35,23 @@ export const TURNS_SEND_EVENTS_PATH = 'turns/send-events'; const MAX_OWNERSHIP_ATTEMPTS = 4; +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(), +}); + export const SendEventsPeerBodySchema = z.object({ tenant_id: z.string(), session_id: z.string(), turn_id: z.string(), - user_ref: z.string(), events: z.array(TurnUserEventSchema), + request_context: PeerRequestContextSchema, request_headers: z.record(z.string(), z.string()).optional(), }); @@ -98,8 +110,7 @@ export async function deliverTurnEvents(input: { resume: () => Promise<{ turn: TurnHandle; startDrain: () => void }>; /** Peer delivery is owner-local. Only the replica that received the client request may steal. */ allowSteal: boolean; - userRef: string; - tenantId: string; + requestContext: RequestContext; requestHeaders?: Record | undefined; }): Promise { const sessionId = input.session.session_id; @@ -136,8 +147,7 @@ export async function deliverTurnEvents(input: { sessionId, turnId: input.turnId, events: input.events, - userRef: input.userRef, - tenantId: input.tenantId, + requestContext: input.requestContext, requestHeaders: input.requestHeaders, logger: input.logger, }); @@ -263,20 +273,25 @@ async function forwardTurnEvents(input: { sessionId: string; turnId: string; events: TurnUserEvent[]; - userRef: string; - tenantId: string; + requestContext: RequestContext; requestHeaders?: Record | undefined; logger: Pick; }): Promise { if (!input.redis) { return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } + const { requestContext } = input; const body = SendEventsPeerBodySchema.parse({ - tenant_id: input.tenantId, + tenant_id: requestContext.tenant_id, session_id: input.sessionId, turn_id: input.turnId, - user_ref: input.userRef, 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 { 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' }); From 39160bdb7a2783d51e88fe180f62dda36c670c6c Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Mon, 19 Oct 2026 23:18:30 +0530 Subject: [PATCH 08/32] fix: offload paused turns that the keepalive timer could not see TurnHandle.state is a prototype getter, so paused runs never offloaded and shutdown cancelled them. A rebuilt turn stays paused without a new pause event, and a live send must reset the timer so it cannot abort the run before running is persisted. Co-authored-by: Cursor --- packages/trueforge/src/apis/turns.ts | 6 +++++ packages/trueforge/src/runtime/activeTurns.ts | 13 +++++++--- .../trueforge/src/runtime/turnDelivery.ts | 24 +++++++++++++------ 3 files changed, 33 insertions(+), 10 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 4ffd95426..9c8d00217 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -680,6 +680,12 @@ async function resumePausedTurn(input: { 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.armPausedKeepalive({ + sessionId: input.session.session_id, + turnId: turn.id, + }); void drainTurnEvents({ trackedStream, turnEventStream, diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 4a0aff6ae..c40cd066d 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -20,9 +20,12 @@ function activeTurnKey(sessionId: string, turnId: string): string { return `${sessionId}:${turnId}`; } -/** A test double may omit `state`. A real handle always has it. */ +/** + * `state` is a getter on TurnHandle's prototype, so it is not an own property. + * A test double may omit it; that run is not treated as paused. + */ function isPausedRun(run: ActiveTurnRun | undefined): run is ActiveTurnRun { - if (!run || !Object.hasOwn(run.turn, 'state')) { + if (!run || !('state' in run.turn)) { return false; } return run.turn.state.status === 'paused'; @@ -170,8 +173,12 @@ export class ActiveTurnRegistry { } const key = activeTurnKey(input.sessionId, input.turnId); const timer = setTimeout(() => { - this.keepaliveTimers.delete(key); 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)) { return; diff --git a/packages/trueforge/src/runtime/turnDelivery.ts b/packages/trueforge/src/runtime/turnDelivery.ts index b921b6e2c..e86a909e5 100644 --- a/packages/trueforge/src/runtime/turnDelivery.ts +++ b/packages/trueforge/src/runtime/turnDelivery.ts @@ -180,13 +180,23 @@ export async function deliverTurnEvents(input: { if (!handle) { continue; } - return applyTurnEvents({ - handle, - events: input.events, - sessionStore: input.sessionStore, - sessionId, - turnId: input.turnId, - }); + // 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.activeTurns.clearPausedKeepalive({ sessionId, turnId: input.turnId }); + try { + return await applyTurnEvents({ + handle, + events: input.events, + sessionStore: input.sessionStore, + sessionId, + turnId: input.turnId, + }); + } finally { + if (handle.state.status === 'paused') { + input.activeTurns.armPausedKeepalive({ sessionId, turnId: input.turnId }); + } + } } if (action === TurnOwnershipAction.Rebuild || action === TurnOwnershipAction.Steal) { From 66149c41f4318fb0a66427cb278881854ef34dba Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 20 Oct 2026 09:52:47 +0530 Subject: [PATCH 09/32] fix: require the peer router and session store on turn start Creating the next turn cancels a previous run that is not in this process, including a schedule run, so that path needs the session store. The turns router always registers turns/send-events. Co-authored-by: Cursor --- packages/trueforge/src/apis/schedules.ts | 4 + packages/trueforge/src/apis/sessions.ts | 14 +- packages/trueforge/src/apis/turns.ts | 159 ++++++++++-------- packages/trueforge/src/app.ts | 7 +- packages/trueforge/src/main.ts | 1 + .../unit/apis/deletedSessionCrud.test.ts | 1 + .../unit/apis/sandboxFileDownload.test.ts | 2 + .../tests/unit/apis/schedules.test.ts | 1 + .../tests/unit/apis/sessionHttp.test.ts | 1 + .../unit/apis/turnHarnessErrorStatus.test.ts | 2 + .../trueforge/tests/unit/apis/turns.test.ts | 8 + 11 files changed, 123 insertions(+), 77 deletions(-) diff --git a/packages/trueforge/src/apis/schedules.ts b/packages/trueforge/src/apis/schedules.ts index cdd817dbb..08f9b54c9 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'; @@ -82,6 +83,8 @@ export interface ScheduleTurnExecutionDeps { agentStore: IAgentStore; /** Call with `{}` — schedules have no caller `x-tfg-skills`. */ resolveTurnSkillsStore: (input: ResolveTurnSkillsStoreInput) => Pick; + /** Previous-turn cancel when a schedule run supersedes a turn that is not in this process. */ + sessionStore: Pick; } export interface SchedulesRouterDeps extends ScheduleTurnExecutionDeps { @@ -125,6 +128,7 @@ export async function startScheduleRunOnRequest(params: { webSearchProviderStore: deps.resolveWebSearchProviderStore(c), skillStore: deps.resolveTurnSkillsStore({}), logger: deps.logger, + sessionStore: deps.sessionStore, }, }); } diff --git a/packages/trueforge/src/apis/sessions.ts b/packages/trueforge/src/apis/sessions.ts index 31e3fe851..89cdc69bd 100644 --- a/packages/trueforge/src/apis/sessions.ts +++ b/packages/trueforge/src/apis/sessions.ts @@ -68,10 +68,20 @@ export interface SessionsRouterDeps { resolveAgentStore: (c: Context) => IAgentStore; resolveSandboxProviderStore: (c: Context) => ISandboxProviderStore; sandboxEnvironmentStore: ISandboxEnvironmentStore; - redis?: RedisClient | undefined; resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; - requestReplyRouter: RequestReplyRouter; resolveRequestContext: ResolveRequestContext; + /** + * 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; + /** + * Used to register routes + * - `sessions/cancel` -> {@link cancelSessionTurnPeerHandler} + */ + requestReplyRouter: RequestReplyRouter; logger: Logger; authorizer: Authorizer; } diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 9c8d00217..7bc4e79bb 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -32,7 +32,11 @@ import { VercelAILLM, withTimeout, } from '@truefoundry/trueforge-core/core'; -import type { JSONReply, RedisClient, RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; +import type { + RedisClient, + RouteHandler as RequestReplyRouteHandler, + RequestReplyRouter, +} from '@truefoundry/trueforge-core/request-reply'; import type { Context } from 'hono'; import { HTTPException } from 'hono/http-exception'; import { streamSSE } from 'hono/streaming'; @@ -140,13 +144,23 @@ export interface TurnsRouterDeps { logger: Logger; resolveRequestContext: ResolveRequestContext; authorizer: Authorizer; - /** Present when this replica can forward events to the database owner. */ + /** + * 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; - /** Present when this replica serves peer turn-event requests. */ - requestReplyRouter?: RequestReplyRouter | undefined; /** - * Store bundle for a peer rebuild, which has no inbound HTTP context. - * The client-facing handler uses the request's own stores instead. + * Used to register routes + * - `turns/send-events` -> {@link handlePeerTurnEvents} + */ + requestReplyRouter: RequestReplyRouter; + /** + * Builds model, MCP, and skill stores for that inbound rebuild. The peer body carries + * the caller's request context and headers, and there is no Hono context. POST /events + * on this process builds the same stores from the HTTP request instead. */ createPeerTurnExecutionDeps?: | ((input: { @@ -172,8 +186,8 @@ export type BeginTurnExecutionDeps = Pick< sandboxProviderStore: ISandboxProviderStore; webSearchProviderStore: IWebSearchProviderStore; skillStore: Pick; - /** Present on the HTTP create path so a previous turn can be superseded on its owner. */ - sessionStore?: Pick | undefined; + /** Loads the previous turn when create must cancel a run that is not in this process. */ + sessionStore: Pick; redis?: RedisClient | undefined; }; @@ -520,7 +534,7 @@ export async function beginTurnExecution( turnId: previousId, abortReason: CancellationReason.CancelledForNextTurn, }); - if (!hadLocalRun && deps.sessionStore) { + if (!hadLocalRun) { await cancelSessionTurn( { activeTurns: deps.activeTurns, @@ -1246,6 +1260,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { skillStore: deps.resolveTurnSkillsStore({ c, allowInlineResources }), sandboxProviderStore: deps.resolveSandboxProviderStore(c), webSearchProviderStore: deps.resolveWebSearchProviderStore(c), + sessionStore: deps.sessionStore, }; } @@ -1257,70 +1272,70 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { router.openapi(listTurnEventsRoute, listTurnEventsHandler); router.openapi(postTurnEventsRoute, postTurnEventsHandler); router.openapi(subscribeTurnRoute, subscribeTurnHandler); - if (deps.requestReplyRouter) { - deps.requestReplyRouter.registerRoute(TURNS_SEND_EVENTS_PATH, request => handlePeerTurnEvents(deps, request)); - } + deps.requestReplyRouter.registerRoute(TURNS_SEND_EVENTS_PATH, handlePeerTurnEvents(deps)); return router; } -async function handlePeerTurnEvents(deps: TurnsRouterDeps, request: { body: unknown }): Promise { - const parsed = SendEventsPeerBodySchema.safeParse(request.body); - if (!parsed.success) { - return { status: 400, body: jsonValueFrom({ message: 'Invalid turns/send-events payload' }) }; - } - const createPeerTurnExecutionDeps = deps.createPeerTurnExecutionDeps; - if (!createPeerTurnExecutionDeps) { - return { status: 503, body: jsonValueFrom({ message: 'Turn owner is unavailable' }) }; - } - const session = await deps.sessions.get({ - tenant_id: parsed.data.tenant_id, - session_id: parsed.data.session_id, - }); - if (!session) { - return { status: 404, body: jsonValueFrom({ message: `Session not found: ${parsed.data.session_id}` }) }; - } - const result = await deps.activeTurns.withTurnLock( - { sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }, - async () => { - const requestContext = requestContextFromWire(parsed.data.request_context); - const executionDeps = await createPeerTurnExecutionDeps({ - session, - requestContext, - requestHeaders: parsed.data.request_headers, - }); - return deliverTurnEvents({ - session, - sessionStore: deps.sessionStore, - activeTurns: deps.activeTurns, - turnId: parsed.data.turn_id, - events: parsed.data.events, - redis: undefined, - logger: deps.logger, - allowSteal: false, - requestContext, - requestHeaders: parsed.data.request_headers, - resume: () => - resumePausedTurn({ - session, - turnId: parsed.data.turn_id, - requestContext, - requestHeaders: parsed.data.request_headers, - deps: executionDeps, - }), - }); - }, - ); - if (result.type === 'applied') { - return { status: 200, body: jsonValueFrom({ events: result.events }) }; - } - if (result.type === 'invalid') { - return { status: 400, body: jsonValueFrom({ message: result.message }) }; - } - if (result.type === 'not_owner') { - return { status: 412, body: jsonValueFrom({ message: 'Turn is not running on this executor' }) }; - } - if (result.type === 'conflict') { - return { status: 409, body: jsonValueFrom({ message: result.message }) }; - } - return { status: 503, body: jsonValueFrom({ message: result.message }) }; +function handlePeerTurnEvents(deps: TurnsRouterDeps): RequestReplyRouteHandler { + return async request => { + const parsed = SendEventsPeerBodySchema.safeParse(request.body); + if (!parsed.success) { + return { status: 400, body: jsonValueFrom({ message: 'Invalid turns/send-events payload' }) }; + } + const createPeerTurnExecutionDeps = deps.createPeerTurnExecutionDeps; + if (!createPeerTurnExecutionDeps) { + return { status: 503, body: jsonValueFrom({ message: 'Turn owner is unavailable' }) }; + } + const session = await deps.sessions.get({ + tenant_id: parsed.data.tenant_id, + session_id: parsed.data.session_id, + }); + if (!session) { + return { status: 404, body: jsonValueFrom({ message: `Session not found: ${parsed.data.session_id}` }) }; + } + const result = await deps.activeTurns.withTurnLock( + { sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }, + async () => { + const requestContext = requestContextFromWire(parsed.data.request_context); + const executionDeps = await createPeerTurnExecutionDeps({ + session, + requestContext, + requestHeaders: parsed.data.request_headers, + }); + return deliverTurnEvents({ + session, + sessionStore: deps.sessionStore, + activeTurns: deps.activeTurns, + turnId: parsed.data.turn_id, + events: parsed.data.events, + redis: undefined, + logger: deps.logger, + allowSteal: false, + requestContext, + requestHeaders: parsed.data.request_headers, + resume: () => + resumePausedTurn({ + session, + turnId: parsed.data.turn_id, + requestContext, + requestHeaders: parsed.data.request_headers, + deps: executionDeps, + }), + }); + }, + ); + if (result.type === 'applied') { + return { status: 200, body: jsonValueFrom({ events: result.events }) }; + } + if (result.type === 'invalid') { + return { status: 400, body: jsonValueFrom({ message: result.message }) }; + } + if (result.type === 'not_owner') { + return { status: 412, body: jsonValueFrom({ message: 'Turn is not running on this executor' }) }; + } + if (result.type === 'conflict') { + return { status: 409, body: jsonValueFrom({ message: result.message }) }; + } + return { status: 503, body: jsonValueFrom({ message: result.message }) }; + }; } diff --git a/packages/trueforge/src/app.ts b/packages/trueforge/src/app.ts index 3f39b003e..95b7d37f8 100644 --- a/packages/trueforge/src/app.ts +++ b/packages/trueforge/src/app.ts @@ -273,11 +273,11 @@ export interface ServerDeps { sandboxProviderStore: ISandboxProviderStore; sessions: Sessions; activeTurns: ActiveTurnRegistry; - /** Primary Redis client (server-owned); undefined in standalone mode. */ + /** Redis client used for redis request-reply */ redis?: RedisClient | undefined; - /** Request-reply dispatch table served by this replica's executor. */ + /** Redis request-reply router for registering peer handlers. */ requestReplyRouter: RequestReplyRouter; - /** Store bundle for rebuilding a paused turn on a peer request, which has no HTTP context. */ + /** Store factory for a peered rebuild. The peer body has no Hono context. */ createPeerTurnExecutionDeps: (input: { session: SessionHandle; requestContext: RequestContext; @@ -313,6 +313,7 @@ export function createServerApp(deps: ServerDeps) { resolveWebSearchProviderStore: deps.resolveWebSearchProviderStore, activeTurns: deps.activeTurns, resolveTurnSkillsStore: deps.resolveTurnSkillsStore, + sessionStore: deps.sessionStore, }; const httpMetrics = createHttpMetrics(); diff --git a/packages/trueforge/src/main.ts b/packages/trueforge/src/main.ts index 9c1941828..cf7cf39c9 100644 --- a/packages/trueforge/src/main.ts +++ b/packages/trueforge/src/main.ts @@ -829,6 +829,7 @@ async function createServerRuntime(persistence: ServerPersistence< ), sandboxProviderStore: persistence.resolveSandboxProviderStore(requestContext), webSearchProviderStore: persistence.resolveWebSearchProviderStore(requestContext), + sessionStore, }; }; const app = createServerApp({ diff --git a/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts b/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts index 17409f077..411ef0ab6 100644 --- a/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts +++ b/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts @@ -80,6 +80,7 @@ describe('public CRUD after session deletion', () => { resolveWebSearchProviderStore: () => webSearchProviderStore, logger: createLogger({ silent: true }), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer: new TrueForgeAuthorizer(), }), ); diff --git a/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts b/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts index 2bbda5453..c9c911df5 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'; @@ -57,6 +58,7 @@ async function buildApp() { resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer: new TrueForgeAuthorizer(), }), ); diff --git a/packages/trueforge/tests/unit/apis/schedules.test.ts b/packages/trueforge/tests/unit/apis/schedules.test.ts index 51fa79f64..6f161a6bf 100644 --- a/packages/trueforge/tests/unit/apis/schedules.test.ts +++ b/packages/trueforge/tests/unit/apis/schedules.test.ts @@ -74,6 +74,7 @@ function stubTurnExecutionDeps(agentStore: SqliteAgentStore, scheduleStore: Sqli resolveSandboxProviderStore: () => ({}) as never, sandboxEnvironmentStore: {} as never, resolveWebSearchProviderStore: () => ({}) as never, + sessionStore, }; } diff --git a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts index 8a444f60e..88e701189 100644 --- a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts +++ b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts @@ -157,6 +157,7 @@ describe('sessions HTTP agent binding', () => { resolveWebSearchProviderStore: () => webSearchProviderStore, logger: deps.logger, resolveRequestContext: deps.resolveRequestContext, + requestReplyRouter: deps.requestReplyRouter, authorizer: deps.authorizer, }), ); diff --git a/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts b/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts index 026e609db..4598260ac 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'; @@ -86,6 +87,7 @@ async function postTurnRejectingWith(error: AgentHarnessError): Promise new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer: new TrueForgeAuthorizer(), }), ); diff --git a/packages/trueforge/tests/unit/apis/turns.test.ts b/packages/trueforge/tests/unit/apis/turns.test.ts index 65d707cf0..e28211812 100644 --- a/packages/trueforge/tests/unit/apis/turns.test.ts +++ b/packages/trueforge/tests/unit/apis/turns.test.ts @@ -6,6 +6,7 @@ import { type TurnInputItem, type TurnStreamingEvent, } 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'; @@ -131,6 +132,7 @@ describe('turns', () => { resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer: new TrueForgeAuthorizer(), }), ); @@ -227,6 +229,7 @@ describe('turns', () => { resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer: { listAgentAccess: input => Promise.resolve( @@ -329,6 +332,7 @@ describe('turns', () => { resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer: new TrueForgeAuthorizer(), }), ); @@ -400,6 +404,7 @@ describe('turns', () => { resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer: new TrueForgeAuthorizer(), }), ); @@ -508,6 +513,7 @@ describe('turns', () => { resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger, resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer: new TrueForgeAuthorizer(), }), ); @@ -617,6 +623,7 @@ describe('turns', () => { resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger, resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer: new TrueForgeAuthorizer(), }), ); @@ -704,6 +711,7 @@ describe('turns', () => { resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + requestReplyRouter: new RequestReplyRouter(), authorizer, }), ); From 21ac55ab130e5efa138ff5def6326728c6a0a479 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 20 Oct 2026 19:50:44 +0530 Subject: [PATCH 10/32] Refactor resolvers --- packages/trueforge/scripts/write-openapi.ts | 5 +- packages/trueforge/src/apis/agents.ts | 12 +- packages/trueforge/src/apis/capabilities.ts | 5 +- .../trueforge/src/apis/sandboxEnvironments.ts | 9 +- .../trueforge/src/apis/sandboxProviders.ts | 7 +- packages/trueforge/src/apis/schedules.ts | 24 +- packages/trueforge/src/apis/sessions.ts | 53 +- packages/trueforge/src/apis/settings.ts | 8 +- packages/trueforge/src/apis/turns.ts | 183 +++--- .../trueforge/src/apis/webSearchProviders.ts | 7 +- packages/trueforge/src/app.ts | 215 +++---- packages/trueforge/src/main.ts | 584 +++++------------- packages/trueforge/src/resolvers.ts | 389 ++++++++++++ .../src/truefoundry/TrueFoundryAgentStore.ts | 4 +- .../src/truefoundry/TrueFoundrySkillStore.ts | 4 +- .../trueforge/tests/unit/apis/agents.test.ts | 8 +- .../tests/unit/apis/capabilities.test.ts | 4 +- .../unit/apis/deletedSessionCrud.test.ts | 49 +- .../tests/unit/apis/modelProviders.test.ts | 4 +- .../unit/apis/sandboxEnvironments.test.ts | 2 +- .../unit/apis/sandboxFileDownload.test.ts | 36 +- .../tests/unit/apis/sandboxProviders.test.ts | 4 +- .../tests/unit/apis/schedules.test.ts | 12 +- .../tests/unit/apis/sessionHttp.test.ts | 58 +- .../unit/apis/turnHarnessErrorStatus.test.ts | 36 +- .../trueforge/tests/unit/apis/turns.test.ts | 148 +++-- .../unit/apis/webSearchProviders.test.ts | 2 +- .../sandbox/sandboxEnvironmentSecrets.test.ts | 2 +- 28 files changed, 986 insertions(+), 888 deletions(-) create mode 100644 packages/trueforge/src/resolvers.ts diff --git a/packages/trueforge/scripts/write-openapi.ts b/packages/trueforge/scripts/write-openapi.ts index 862dd6972..832ae5823 100644 --- a/packages/trueforge/scripts/write-openapi.ts +++ b/packages/trueforge/scripts/write-openapi.ts @@ -84,12 +84,11 @@ 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, @@ -98,7 +97,7 @@ const app = createServerApp({ sessions: new Sessions({ sessionStore }), activeTurns: new ActiveTurnRegistry(), requestReplyRouter: new RequestReplyRouter(), - createPeerTurnExecutionDeps: () => Promise.reject(new Error('peer rebuild is not used while writing OpenAPI')), + 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 08f9b54c9..5d770f2b4 100644 --- a/packages/trueforge/src/apis/schedules.ts +++ b/packages/trueforge/src/apis/schedules.ts @@ -69,22 +69,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 { @@ -123,9 +125,9 @@ 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, diff --git a/packages/trueforge/src/apis/sessions.ts b/packages/trueforge/src/apis/sessions.ts index 89cdc69bd..55defe3c8 100644 --- a/packages/trueforge/src/apis/sessions.ts +++ b/packages/trueforge/src/apis/sessions.ts @@ -12,7 +12,6 @@ import { type RedisClient, type RequestReplyRouter } from '@truefoundry/trueforg import type { Context } from 'hono'; import { HTTPException } from 'hono/http-exception'; import type { Logger } from 'winston'; -import type { ResolveSkillStoreInput } from '../app'; import type { Authorizer } from '../auth/authorizer'; import { createdBySubjectFromRequestContext, type RequestContext, type ResolveRequestContext } from '../auth/identity'; import type { IAgentStore } from '../db/agentStore'; @@ -59,17 +58,9 @@ 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; - resolveWebSearchProviderStore: (c: Context) => IWebSearchProviderStore; - resolveRequestContext: ResolveRequestContext; + authorizer: Authorizer; + logger: Logger; /** * Publishes to another executor: * - Creating the next turn -> redis request `sessions/cancel` @@ -82,8 +73,16 @@ export interface SessionsRouterDeps { * - `sessions/cancel` -> {@link cancelSessionTurnPeerHandler} */ requestReplyRouter: RequestReplyRouter; - logger: Logger; - authorizer: Authorizer; + resolveAgentStore: (c: Context) => IAgentStore; + resolveMcpServerStore: (input: { c: Context; allowInlineResources?: boolean }) => IMcpServerStore; + resolveModelProviderStore: (input: { c: Context; allowInlineResources?: boolean }) => IModelProviderStore; + resolveRequestContext: ResolveRequestContext; + resolveSkillStore: (input: { c: Context; allowInlineResources?: boolean }) => ISkillStore; + sandboxEnvironmentStore: ISandboxEnvironmentStore; + sandboxProviderStore: ISandboxProviderStore; + sessions: Sessions; + sessionStore: ISessionStore; + webSearchProviderStore: IWebSearchProviderStore; } const FORBIDDEN_SESSION_ACCESS = 'Only the session creator can access this session'; @@ -105,12 +104,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 { @@ -122,24 +121,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( @@ -354,9 +353,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; - resolveRequestContext: ResolveRequestContext; - authorizer: Authorizer; /** * Publishes to another executor: * - POST /events -> redis request `turns/send-events` @@ -157,18 +167,21 @@ export interface TurnsRouterDeps { * - `turns/send-events` -> {@link handlePeerTurnEvents} */ requestReplyRouter: RequestReplyRouter; + /** List, get, list events, subscribe, download, and create (referenced-agent check). */ + resolveAgentStore: (requestContext: RequestContext) => IAgentStore; + resolveRequestContext: ResolveRequestContext; /** - * Builds model, MCP, and skill stores for that inbound rebuild. The peer body carries - * the caller's request context and headers, and there is no Hono context. POST /events - * on this process builds the same stores from the HTTP request instead. + * 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. */ - createPeerTurnExecutionDeps?: - | ((input: { - session: SessionHandle; - requestContext: RequestContext; - requestHeaders?: Record | undefined; - }) => Promise) - | undefined; + 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; } /** @@ -191,6 +204,38 @@ export type BeginTurnExecutionDeps = 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; @@ -796,7 +841,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, @@ -831,7 +876,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, @@ -886,7 +931,7 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { const provider = await resolveSandboxProvider({ tenant_id: requestContext.tenant_id, - store: deps.resolveSandboxProviderStore(c), + store: deps.sandboxProviderStore, logger: deps.logger, sessionId, }); @@ -931,7 +976,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, @@ -983,7 +1028,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, }); @@ -1002,31 +1047,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 { @@ -1095,7 +1128,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, @@ -1197,7 +1230,17 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { const requestHeaders = c.req.header(); return deps.activeTurns.withTurnLock({ sessionId, turnId }, async () => { - const executionDeps = await executionDepsForRequest(c, session); + const runAsAgent = await referencedRunAsAgent({ + session, + agentStore: deps.resolveAgentStore(requestContext), + }); + const executionDeps = await executionDepsForCaller({ + deps, + session, + requestContext, + requestHeaders, + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + }); const result = await deliverTurnEvents({ session, sessionStore: deps.sessionStore, @@ -1232,38 +1275,6 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { }); }; - async function executionDepsForRequest(c: Context, session: SessionHandle): Promise { - let referencedAgent: AgentRecord | undefined; - if (session.record.agent.type === 'reference') { - referencedAgent = await deps.resolveAgentStore(c).getAgent({ - tenant_id: session.tenant_id, - id: session.record.agent.id, - }); - } - const allowInlineResources = session.record.agent.type === 'inline'; - return { - activeTurns: deps.activeTurns, - eventSubscriptions: deps.eventSubscriptions, - logger: deps.logger, - sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - agentStore: deps.resolveAgentStore(c), - modelProviderStore: deps.resolveModelProviderStore({ - c, - allowInlineResources, - ...(referencedAgent !== undefined ? { runAsAgent: referencedAgent } : {}), - }), - mcpServerStore: deps.resolveMcpServerStore({ - c, - allowInlineResources, - ...(referencedAgent !== undefined ? { runAsAgent: referencedAgent } : {}), - }), - skillStore: deps.resolveTurnSkillsStore({ c, allowInlineResources }), - sandboxProviderStore: deps.resolveSandboxProviderStore(c), - webSearchProviderStore: deps.resolveWebSearchProviderStore(c), - sessionStore: deps.sessionStore, - }; - } - const router = new OpenAPIHono(); router.openapi(createAndExecuteTurnRoute, createAndExecuteTurnHandler); router.openapi(listTurnsRoute, listTurnsHandler); @@ -1282,10 +1293,6 @@ function handlePeerTurnEvents(deps: TurnsRouterDeps): RequestReplyRouteHandler { if (!parsed.success) { return { status: 400, body: jsonValueFrom({ message: 'Invalid turns/send-events payload' }) }; } - const createPeerTurnExecutionDeps = deps.createPeerTurnExecutionDeps; - if (!createPeerTurnExecutionDeps) { - return { status: 503, body: jsonValueFrom({ message: 'Turn owner is unavailable' }) }; - } const session = await deps.sessions.get({ tenant_id: parsed.data.tenant_id, session_id: parsed.data.session_id, @@ -1297,10 +1304,16 @@ function handlePeerTurnEvents(deps: TurnsRouterDeps): RequestReplyRouteHandler { { sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }, async () => { const requestContext = requestContextFromWire(parsed.data.request_context); - const executionDeps = await createPeerTurnExecutionDeps({ + const runAsAgent = await referencedRunAsAgent({ + session, + agentStore: deps.resolveAgentStore(requestContext), + }); + const executionDeps = await executionDepsForCaller({ + deps, session, requestContext, requestHeaders: parsed.data.request_headers, + ...(runAsAgent !== undefined ? { runAsAgent } : {}), }); return deliverTurnEvents({ session, 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 95b7d37f8..a387e25e0 100644 --- a/packages/trueforge/src/app.ts +++ b/packages/trueforge/src/app.ts @@ -1,12 +1,7 @@ /** The API: resource routers, the OpenAPI document and Swagger UI, all under /api/v1. */ import { swaggerUI } from '@hono/swagger-ui'; import { OpenAPIHono, z } from '@hono/zod-openapi'; -import type { - ISessionStore, - SessionHandle, - Sessions, - TurnStreamingEvent, -} from '@truefoundry/trueforge-core/agent-session'; +import type { ISessionStore, Sessions, TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session'; import { extractErrorLogFields } from '@truefoundry/trueforge-core/core'; import type { RedisClient, RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; import type { Context, ErrorHandler, MiddlewareHandler } from 'hono'; @@ -29,7 +24,7 @@ import { createInternalMetricsRouter } from './apis/sessionMetrics'; import { createInternalSessionsRouter, createSessionsRouter } from './apis/sessions'; import { createSettingsRouter } from './apis/settings'; import { createAvailableSkillsRouter } from './apis/skills'; -import { createTurnsRouter, type BeginTurnExecutionDeps } 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, type RequestContext } from './auth/identity'; @@ -230,68 +225,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; - /** Redis client used for redis request-reply */ - redis?: RedisClient | undefined; - /** Redis request-reply router for registering peer handlers. */ - requestReplyRouter: RequestReplyRouter; - /** Store factory for a peered rebuild. The peer body has no Hono context. */ - createPeerTurnExecutionDeps: (input: { - session: SessionHandle; - requestContext: RequestContext; - requestHeaders?: Record | undefined; - }) => Promise; - /** 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) { @@ -300,22 +287,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, - sessionStore: deps.sessionStore, - }; - + const resolveHttpAgentStore = (c: Context) => deps.resolveAgentStore(resolveRequestContext(c)); const httpMetrics = createHttpMetrics(); // Health, the scrape endpoint, and static files sit outside /api, so they never enter this middleware. app.use('/api/*', httpMetrics.middleware); @@ -326,10 +298,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({ @@ -343,7 +313,7 @@ export function createServerApp(deps: ServerDeps) { withAuth( createCapabilitiesRouter({ sandboxEnvironmentStore: deps.sandboxEnvironmentStore, - resolveWebSearchProviderStore: deps.resolveWebSearchProviderStore, + webSearchProviderStore: deps.webSearchProviderStore, withTransaction: deps.withTransaction, logger: deps.logger, resolveRequestContext, @@ -414,13 +384,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, @@ -433,8 +403,8 @@ 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, }), @@ -443,7 +413,24 @@ export function createServerApp(deps: ServerDeps) { ); app.route( '/api/internal/schedules', - withAuth(createScheduleExecutionRouter(scheduleTurnDeps), scheduleExecutionAuthMiddleware), + withAuth( + createScheduleExecutionRouter({ + 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, + }), + scheduleExecutionAuthMiddleware, + ), ); app.route( '/api/internal/sandbox-environments', @@ -461,10 +448,10 @@ export function createServerApp(deps: ServerDeps) { withAuth( createSchedulesRouter({ ...scheduleTurnDeps, - resolveAgentStore: deps.resolveAgentStore, - withTransaction: deps.withTransaction, - resolveRequestContext, authorizer: deps.authorizer, + resolveAgentStore: resolveHttpAgentStore, + resolveRequestContext, + withTransaction: deps.withTransaction, }), authMiddleware, ), @@ -477,10 +464,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, @@ -492,16 +479,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, ), @@ -512,7 +499,7 @@ export function createServerApp(deps: ServerDeps) { createInternalMetricsRouter({ sessionMetricsStore: deps.sessionMetricsStore, resolveRequestContext, - resolveAgentStore: deps.resolveAgentStore, + resolveAgentStore: resolveHttpAgentStore, authorizer: deps.authorizer, }), authMiddleware, @@ -523,7 +510,7 @@ export function createServerApp(deps: ServerDeps) { withAuth( createPermissionsRouter({ authorizer: deps.authorizer, - resolveAgentStore: deps.resolveAgentStore, + resolveAgentStore: resolveHttpAgentStore, scheduleStore: deps.scheduleStore, sessionStore: deps.sessionStore, resolveRequestContext, @@ -535,21 +522,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, ), @@ -558,23 +545,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, - resolveRequestContext, - authorizer: deps.authorizer, redis: deps.redis, requestReplyRouter: deps.requestReplyRouter, - createPeerTurnExecutionDeps: deps.createPeerTurnExecutionDeps, + resolveAgentStore: deps.resolveAgentStore, + resolveRequestContext, + resolveTurnExecutionStores: deps.resolveTurnExecutionStores, + sandboxEnvironmentStore: deps.sandboxEnvironmentStore, + sandboxProviderStore: deps.sandboxProviderStore, + sessions: deps.sessions, + sessionStore: deps.sessionStore, }), authMiddleware, ), diff --git a/packages/trueforge/src/main.ts b/packages/trueforge/src/main.ts index cf7cf39c9..1130ccf0a 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 { @@ -81,23 +80,16 @@ import { CancellationReason, Sessions, type ISessionStore, - type SessionHandle, 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'; @@ -106,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'; @@ -125,89 +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 { 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 { 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. */ @@ -231,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) { - return requestContext => - new TrueFoundrySkillStore({ - client, - context: 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. - */ -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, { @@ -373,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). */ @@ -428,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), }; } @@ -539,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({ @@ -546,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, @@ -566,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, @@ -614,25 +448,39 @@ 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({ @@ -680,189 +528,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 createPeerTurnExecutionDeps = async (input: { - session: SessionHandle; - requestContext: RequestContext; - requestHeaders?: Record | undefined; - }) => { - const { requestContext } = input; - let referencedAgent: AgentRecord | undefined; - if (input.session.record.agent.type === 'reference') { - referencedAgent = await agentStore.getAgent({ - tenant_id: input.session.tenant_id, - id: input.session.record.agent.id, - }); - } - const allowInlineResources = input.session.record.agent.type === 'inline'; - const rawInlineModelsHeader = headerValue(input.requestHeaders, X_TFG_MODELS); - const rawInlineMcpHeader = headerValue(input.requestHeaders, X_TFG_MCP); - const rawPerServerHeaders = headerValue(input.requestHeaders, X_TFG_MCP_HEADERS); - const rawInlineSkillsHeader = headerValue(input.requestHeaders, X_TFG_SKILLS); - if (rawInlineModelsHeader !== undefined) { - assertInlineResourceHeadersAllowed({ - headerNames: [X_TFG_MODELS], - isTrueFoundryMode, - allowInlineResources, - }); - } - if (rawInlineMcpHeader !== undefined) { - assertInlineResourceHeadersAllowed({ - headerNames: [X_TFG_MCP], - isTrueFoundryMode, - allowInlineResources, - }); - } - if (rawInlineSkillsHeader !== undefined) { - assertInlineResourceHeadersAllowed({ - headerNames: [X_TFG_SKILLS], - isTrueFoundryMode, - allowInlineResources, - }); - } - return { - activeTurns, - eventSubscriptions, - logger, - sandboxEnvironmentStore, - agentStore, - modelProviderStore: persistence.resolveModelProviderStore({ - rc: requestContext, - ...(referencedAgent !== undefined ? { runAsAgent: referencedAgent } : {}), - ...(rawInlineModelsHeader !== undefined - ? { inlineModelProviders: parseInlineModelProviders(rawInlineModelsHeader) } - : {}), - }), - mcpServerStore: persistence.resolveMcpServerStore({ - rc: requestContext, - ...(referencedAgent !== undefined ? { runAsAgent: referencedAgent } : {}), - ...(rawPerServerHeaders !== undefined - ? { perServerHeaders: parsePerServerMcpHeaders(rawPerServerHeaders) } - : {}), - ...(rawInlineMcpHeader !== undefined ? { inlineMcpServers: parseInlineMcpServers(rawInlineMcpHeader) } : {}), - }), - skillStore: resolveTurnSkillsStore( - rawInlineSkillsHeader === undefined ? {} : { inlineSkills: parseInlineSkills(rawInlineSkillsHeader) }, - ), - sandboxProviderStore: persistence.resolveSandboxProviderStore(requestContext), - webSearchProviderStore: persistence.resolveWebSearchProviderStore(requestContext), - sessionStore, - }; - }; + 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, - createPeerTurnExecutionDeps, - 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/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/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 411ef0ab6..f9f19d750 100644 --- a/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts +++ b/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts @@ -47,41 +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 }), - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, requestReplyRouter: new RequestReplyRouter(), - authorizer: new TrueForgeAuthorizer(), + resolveAgentStore: () => agentStore, + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + 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 c9c911df5..7ca505cae 100644 --- a/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts +++ b/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts @@ -40,26 +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), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, requestReplyRouter: new RequestReplyRouter(), - authorizer: new TrueForgeAuthorizer(), + resolveAgentStore: () => new SqliteAgentStore(db), + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + 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 6f161a6bf..e1b061e01 100644 --- a/packages/trueforge/tests/unit/apis/schedules.test.ts +++ b/packages/trueforge/tests/unit/apis/schedules.test.ts @@ -62,19 +62,19 @@ const scheduleBody = { function stubTurnExecutionDeps(agentStore: SqliteAgentStore, scheduleStore: SqliteScheduleStore) { const sessionStore = new InMemorySessionStore(); return { - scheduleStore, - sessions: new Sessions({ sessionStore }), - agentStore, activeTurns: new ActiveTurnRegistry(), + agentStore, 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 88e701189..cbba09b89 100644 --- a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts +++ b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts @@ -116,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, + 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(); @@ -139,26 +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, - resolveRequestContext: deps.resolveRequestContext, requestReplyRouter: deps.requestReplyRouter, - authorizer: deps.authorizer, + resolveAgentStore: () => agentStore, + resolveRequestContext: deps.resolveRequestContext, + 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)); diff --git a/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts b/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts index 4598260ac..511ed8700 100644 --- a/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts +++ b/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts @@ -69,26 +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), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, requestReplyRouter: new RequestReplyRouter(), - authorizer: new TrueForgeAuthorizer(), + resolveAgentStore: () => new SqliteAgentStore(db), + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + 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 e28211812..5d0b3f351 100644 --- a/packages/trueforge/tests/unit/apis/turns.test.ts +++ b/packages/trueforge/tests/unit/apis/turns.test.ts @@ -36,6 +36,17 @@ 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'); @@ -119,21 +130,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), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, requestReplyRouter: new RequestReplyRouter(), - authorizer: new TrueForgeAuthorizer(), + resolveAgentStore: () => new SqliteAgentStore(db), + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, tokenStore)), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore, }), ); @@ -216,20 +224,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, - requestReplyRouter: new RequestReplyRouter(), authorizer: { listAgentAccess: input => Promise.resolve( @@ -243,6 +238,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, }), ); @@ -319,21 +324,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), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, requestReplyRouter: new RequestReplyRouter(), - authorizer: new TrueForgeAuthorizer(), + resolveAgentStore: () => new SqliteAgentStore(db), + 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, }), ); @@ -391,21 +393,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), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, requestReplyRouter: new RequestReplyRouter(), - authorizer: new TrueForgeAuthorizer(), + resolveAgentStore: () => new SqliteAgentStore(db), + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, new SqliteOAuthTokenStore(db))), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore, }), ); @@ -500,21 +499,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), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions, - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger, - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, requestReplyRouter: new RequestReplyRouter(), - authorizer: new TrueForgeAuthorizer(), + resolveAgentStore: () => new SqliteAgentStore(db), + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, tokenStore)), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore: new SqliteSessionStore(db), }), ); @@ -610,21 +606,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), + authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger, - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, requestReplyRouter: new RequestReplyRouter(), - authorizer: new TrueForgeAuthorizer(), + resolveAgentStore: () => new SqliteAgentStore(db), + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + resolveTurnExecutionStores: () => Promise.resolve(turnExecutionStores(db, tokenStore)), + sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), + sandboxProviderStore: new SqliteSandboxProviderStore(db), + sessions, + sessionStore: new SqliteSessionStore(db), }), ); @@ -698,21 +691,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, + authorizer, eventSubscriptions: new EventSubscriptionRegistry(undefined), - resolveSandboxProviderStore: () => new SqliteSandboxProviderStore(db), - sandboxEnvironmentStore: new SqliteSandboxEnvironmentStore(db), - resolveWebSearchProviderStore: () => new SqliteWebSearchProviderStore(db), logger: createLogger({ silent: true }), - resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, requestReplyRouter: new RequestReplyRouter(), - authorizer, + resolveAgentStore: () => agentStore, + resolveRequestContext: () => STANDALONE_REQUEST_CONTEXT, + 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/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, }), From 6c355869fd06896d4b7b05cd39fbd1c52102e832 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 20 Oct 2026 20:17:37 +0530 Subject: [PATCH 11/32] Minor rename in inMemory store --- packages/trueforge/src/runtime/event-subscription/inMemory.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/trueforge/src/runtime/event-subscription/inMemory.ts b/packages/trueforge/src/runtime/event-subscription/inMemory.ts index 2dbef9ac2..328aff64d 100644 --- a/packages/trueforge/src/runtime/event-subscription/inMemory.ts +++ b/packages/trueforge/src/runtime/event-subscription/inMemory.ts @@ -52,7 +52,7 @@ export class InMemoryEventStreamStore { } /** Refreshes TTL on a live stream. Sequence numbers already follow the log length. */ - rearm(streamId: string, streamTTLSeconds?: number): void { + refreshStreamTtl(streamId: string, streamTTLSeconds?: number): void { const stream = this.getLiveStream(streamId); if (!stream || !streamTTLSeconds || streamTTLSeconds <= 0) { return; @@ -124,7 +124,7 @@ export class InMemoryEventSubscription implements EventSubscri ) {} resumeProducer(options?: EventSubscriptionPutOptions): Promise { - this.store.rearm(this.streamId, options?.streamTTLSeconds); + this.store.refreshStreamTtl(this.streamId, options?.streamTTLSeconds); return Promise.resolve(); } From 9e95d75b1444c577376342b3f6c0fcf2edf834f2 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 16:54:55 +0530 Subject: [PATCH 12/32] refactor: execute turns from stored status instead of a resume flag Live handles always carry an orchestrator, session reads return the turn row, and JSON peer bodies go through request-reply. Co-authored-by: Cursor --- .../src/agent-session/SessionHandle.ts | 29 +++++-- .../src/agent-session/TurnHandle.ts | 78 +++++-------------- .../core/runtime/AgentThreadOrchestrator.ts | 10 +-- .../trueforge-core/src/request-reply/index.ts | 2 +- .../trueforge-core/src/request-reply/types.ts | 10 +++ .../tests/agent-session/sessions.test.ts | 15 +--- packages/trueforge/src/apis/turns.ts | 39 +++++----- packages/trueforge/src/runtime/activeTurns.ts | 2 +- .../trueforge/src/runtime/turnDelivery.ts | 33 +------- 9 files changed, 83 insertions(+), 135 deletions(-) diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index 36d34061c..0132e1aaa 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -14,7 +14,7 @@ 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'; @@ -411,7 +411,6 @@ export class SessionHandle< orchestrator, resolver: input.resolver, abortController: input.abortController, - resume: true, }); await handle.replayUnconsumedInbox(); return handle; @@ -439,15 +438,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; diff --git a/packages/trueforge-core/src/agent-session/TurnHandle.ts b/packages/trueforge-core/src/agent-session/TurnHandle.ts index f3c7e72dd..6b3469de3 100644 --- a/packages/trueforge-core/src/agent-session/TurnHandle.ts +++ b/packages/trueforge-core/src/agent-session/TurnHandle.ts @@ -175,57 +175,34 @@ function resolveTerminalState(input: { }; } +/** A live turn is running or paused. Terminal rows are not executed. */ +function executableTurnStatus(state: TurnState): NonTerminalTurnState['status'] { + if (state.status === 'running' || state.status === 'paused') { + return state.status; + } + throw new Error(`Cannot execute a ${state.status} turn`); +} + 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; - /** Resume skips turn.created. Inbox replay happens before stream() is pulled. */ - private readonly resume: boolean; constructor(options: { store: ISessionStore; turn: TurnRecord; - // orchestrator and resolver are optional for store-only handles (e.g. from SessionHandle.getTurn) - orchestrator?: AgentThreadOrchestrator | undefined; - resolver?: ITurnResourceResolver | undefined; - abortController?: AbortController | undefined; - resume?: boolean | 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(); - this.resume = options.resume ?? false; - } - - /** 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 { @@ -261,11 +238,6 @@ export class TurnHandle> { } async *send(events: TurnUserEvent[]): AsyncGenerator { - if (!this.orchestrator) { - throw new Error( - `TurnHandle.send() is only available on a live turn. Perhaps this turn handle was created for read only actions?`, - ); - } yield* this.orchestrator.send(events); this.orchestrator.wake(); } @@ -304,14 +276,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() or SessionHandle.rebuildTurn()', - ); - } // 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 @@ -320,7 +285,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( @@ -331,11 +296,12 @@ export class TurnHandle> { let caughtError: Error | undefined; let executeResult: AgentThreadExecutionResult | undefined; let ownershipLost: TurnOwnershipLostError | undefined; + const turnStatus = executableTurnStatus(this.turn.state); try { - if (!this.resume) { + if (turnStatus === 'running') { yield await this.persistTurnCreated(); } - executeResult = yield* this.executeAndPersist(orchestrator, signal); + executeResult = yield* this.executeAndPersist(orchestrator, signal, turnStatus); } catch (error) { caughtError = error instanceof Error ? error : new Error(String(error)); if (caughtError instanceof TurnOwnershipLostError) { @@ -388,11 +354,6 @@ export class TurnHandle> { * crash between persist and consume does not apply the same decision twice. */ async replayUnconsumedInbox(): Promise { - if (!this.orchestrator) { - throw new Error( - `TurnHandle.replayUnconsumedInbox() is only available on a live turn. Perhaps this turn handle was created for read only actions?`, - ); - } const rows = await this.store.listUnconsumedTurnInboundEvents({ session_id: this.turn.session_id, turn_id: this.turn.turn_id, @@ -434,8 +395,9 @@ export class TurnHandle> { private async *executeAndPersist( orchestrator: AgentThreadOrchestrator, signal: AbortSignal, + turnStatus: NonTerminalTurnState['status'], ): AsyncGenerator { - const generator = orchestrator.execute({ signal, resume: this.resume }); + const generator = orchestrator.execute({ signal, turnStatus }); try { let iterResult = await generator.next(); while (!iterResult.done) { diff --git a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts index bda49c7ce..4189db2c9 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts @@ -566,14 +566,14 @@ export class AgentThreadOrchestrator { public async *execute({ signal, - resume = false, + turnStatus = 'running', }: { signal: AbortSignal; /** - * The stored turn is already paused. Start there so the first runnable batch - * persists `running`; a fresh execute would skip that transition. + * Status of the turn being executed. A rebuild passes `paused` so the first + * runnable batch persists `running`. A fresh execute starts `running`. */ - resume?: boolean; + turnStatus?: 'running' | 'paused' | undefined; }): AsyncGenerator { const { agentThreads } = this; @@ -581,7 +581,7 @@ export class AgentThreadOrchestrator { let output: ModelMessageEvent | null = null; let rootAgentError: AgentThreadExecutionResult['root_agent_error']; let rootFinished = false; - let turnPaused = resume; + let turnPaused = turnStatus === 'paused'; const mainThread = this.getMainThread(); const rootSpan = createRootAgentSpan(mainThread, this.tracing); diff --git a/packages/trueforge-core/src/request-reply/index.ts b/packages/trueforge-core/src/request-reply/index.ts index 0b20d04d5..3017ea0a7 100644 --- a/packages/trueforge-core/src/request-reply/index.ts +++ b/packages/trueforge-core/src/request-reply/index.ts @@ -19,5 +19,5 @@ 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/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index bbf51b2fd..df6b0db3e 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -32,10 +32,11 @@ import { VercelAILLM, withTimeout, } from '@truefoundry/trueforge-core/core'; -import type { - RedisClient, - RouteHandler as RequestReplyRouteHandler, - RequestReplyRouter, +import { + jsonValueFrom, + type RedisClient, + type RouteHandler as RequestReplyRouteHandler, + type RequestReplyRouter, } from '@truefoundry/trueforge-core/request-reply'; import type { Context } from 'hono'; import { HTTPException } from 'hono/http-exception'; @@ -72,12 +73,7 @@ import { resolveSandboxEnvironment, resolveSandboxProvider, } from '../runtime/sessionResources'; -import { - deliverTurnEvents, - jsonValueFrom, - SendEventsPeerBodySchema, - TURNS_SEND_EVENTS_PATH, -} from '../runtime/turnDelivery'; +import { deliverTurnEvents, SendEventsPeerBodySchema, TURNS_SEND_EVENTS_PATH } from '../runtime/turnDelivery'; import { DEFAULT_SANDBOX_ENVIRONMENT_NAME } from '../schemas/sandboxEnvironment'; import { MAX_SESSION_TITLE_LENGTH } from '../schemas/session'; import { assertGatewayMetadataRequestHeaders } from '../truefoundry/gatewayMetadata'; @@ -490,7 +486,7 @@ export async function drainTurnEvents(input: { }); if (activeTurns && event.type === EventType.TURN_UPDATE) { if (event.state.status === 'paused') { - activeTurns.armPausedKeepalive({ sessionId, turnId }); + activeTurns.startPausedKeepalive({ sessionId, turnId }); } else { activeTurns.clearPausedKeepalive({ sessionId, turnId }); } @@ -741,7 +737,7 @@ async function resumePausedTurn(input: { }); // 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.armPausedKeepalive({ + input.deps.activeTurns.startPausedKeepalive({ sessionId: input.session.session_id, turnId: turn.id, }); @@ -889,7 +885,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 => { @@ -924,7 +920,7 @@ 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); } @@ -990,7 +986,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, @@ -1291,14 +1288,14 @@ function handlePeerTurnEvents(deps: TurnsRouterDeps): RequestReplyRouteHandler { return async request => { const parsed = SendEventsPeerBodySchema.safeParse(request.body); if (!parsed.success) { - return { status: 400, body: jsonValueFrom({ message: 'Invalid turns/send-events payload' }) }; + return { status: 400, body: { message: 'Invalid turns/send-events payload' } }; } const session = await deps.sessions.get({ tenant_id: parsed.data.tenant_id, session_id: parsed.data.session_id, }); if (!session) { - return { status: 404, body: jsonValueFrom({ message: `Session not found: ${parsed.data.session_id}` }) }; + return { status: 404, body: { message: `Session not found: ${parsed.data.session_id}` } }; } const result = await deps.activeTurns.withTurnLock( { sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }, @@ -1341,14 +1338,14 @@ function handlePeerTurnEvents(deps: TurnsRouterDeps): RequestReplyRouteHandler { return { status: 200, body: jsonValueFrom({ events: result.events }) }; } if (result.type === 'invalid') { - return { status: 400, body: jsonValueFrom({ message: result.message }) }; + return { status: 400, body: { message: result.message } }; } if (result.type === 'not_owner') { - return { status: 412, body: jsonValueFrom({ message: 'Turn is not running on this executor' }) }; + return { status: 412, body: { message: 'Turn is not running on this executor' } }; } if (result.type === 'conflict') { - return { status: 409, body: jsonValueFrom({ message: result.message }) }; + return { status: 409, body: { message: result.message } }; } - return { status: 503, body: jsonValueFrom({ message: result.message }) }; + return { status: 503, body: { message: result.message } }; }; } diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index c40cd066d..79cad3168 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -166,7 +166,7 @@ export class ActiveTurnRegistry { * under the turn lock: the generator stops without a terminal write. * A later send on this owner rebuilds. `0` disables the timer. */ - armPausedKeepalive(input: { sessionId: string; turnId: string }): void { + startPausedKeepalive(input: { sessionId: string; turnId: string }): void { this.clearPausedKeepalive(input); if (this.pausedKeepaliveMs <= 0 || this.alreadyShutDownAbortReason !== undefined) { return; diff --git a/packages/trueforge/src/runtime/turnDelivery.ts b/packages/trueforge/src/runtime/turnDelivery.ts index e86a909e5..89034fa1d 100644 --- a/packages/trueforge/src/runtime/turnDelivery.ts +++ b/packages/trueforge/src/runtime/turnDelivery.ts @@ -16,11 +16,11 @@ import { } from '@truefoundry/trueforge-core/agent-session'; import { AgentHarnessError } from '@truefoundry/trueforge-core/core'; import { + jsonValueFrom, NO_RESPONDER_PUBLISH_DETAIL, NoResponderError, redisRequest, RequestTimeoutError, - type JSONValue, type RedisClient, } from '@truefoundry/trueforge-core/request-reply'; import type { Logger } from 'winston'; @@ -63,31 +63,6 @@ const PeerMessageSchema = z.object({ message: z.string(), }); -function isJsonValue(value: unknown): value is JSONValue { - if (value === null || typeof value === 'string' || typeof value === 'boolean') { - return true; - } - if (typeof value === 'number') { - return Number.isFinite(value); - } - if (Array.isArray(value)) { - return value.every(isJsonValue); - } - if (typeof value === 'object') { - return Object.values(value).every(isJsonValue); - } - return false; -} - -/** JSON.stringify drops `undefined`, which the Redis envelope cannot carry. */ -export function jsonValueFrom(value: unknown): JSONValue { - const parsed: unknown = JSON.parse(JSON.stringify(value)); - if (!isJsonValue(parsed)) { - throw new Error('Turn event payload is not JSON'); - } - return parsed; -} - export type TurnDeliveryResult = | { type: 'applied'; events: TurnUserEvent[] } | { type: 'invalid'; message: string } @@ -121,7 +96,7 @@ export async function deliverTurnEvents(input: { return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; } const status = loaded.state.status; - const owner = loaded.record.active_executor_id; + const owner = loaded.active_executor_id; const ownedByThisExecutor = owner === configuration.EXECUTOR_ID; if (!ownedByThisExecutor) { @@ -156,7 +131,7 @@ export async function deliverTurnEvents(input: { } const fresh = await input.session.getTurn(input.turnId); const freshStatus = fresh?.state.status ?? status; - const freshOwner = fresh?.record.active_executor_id ?? owner; + const freshOwner = fresh?.active_executor_id ?? owner; if (forwarded.peerOutcome === PeerOutcome.NoResponder && (freshStatus !== 'paused' || freshOwner !== owner)) { continue; } @@ -194,7 +169,7 @@ export async function deliverTurnEvents(input: { }); } finally { if (handle.state.status === 'paused') { - input.activeTurns.armPausedKeepalive({ sessionId, turnId: input.turnId }); + input.activeTurns.startPausedKeepalive({ sessionId, turnId: input.turnId }); } } } From 9b4ca95a1e719d91b05acab72284592371c0b152 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 18:29:47 +0530 Subject: [PATCH 13/32] Address some review comments --- .../src/agent-session/SessionHandle.ts | 163 ++-- packages/trueforge/src/apis/turns.ts | 753 ++++++++++++++---- packages/trueforge/src/runtime/activeTurns.ts | 5 +- packages/trueforge/src/runtime/cancelTurn.ts | 30 +- .../trueforge/src/runtime/turnDelivery.ts | 328 -------- 5 files changed, 680 insertions(+), 599 deletions(-) delete mode 100644 packages/trueforge/src/runtime/turnDelivery.ts diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index 0132e1aaa..b1f5ba034 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -191,56 +191,10 @@ 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, }); const new_context_appends = await collectContextAppends(orchestrator.send(input.input ?? [])); @@ -354,55 +308,10 @@ export class SessionHandle< throw new TurnNotRunningError(input.turn_id, turn.state); } - 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: turn.snapshot.sandbox_info ?? undefined, - previousTurn: turn, - signal: input.abortController.signal, - tracing, - }); - - const definitionsByThreadId = await this.resolveAgentDefinitions({ - previous: turn, - resolver: input.resolver, - spec, - tracing, - signal: input.abortController.signal, - }); - - 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 } = await this.assembleTurn({ previous: turn, resolver: input.resolver, - spec, - sandbox, - tracing, - definitionsByThreadId, - }); - - const createDynamicSubAgentThread = this.makeCreateDynamicSubAgentThread({ - resolver: input.resolver, - previous: turn, - spec, - sandbox, - tracing, - }); - - const orchestrator = new AgentThreadOrchestrator({ - agentThreads, - createDynamicSubAgentThread, - tracing, - logger: input.resolver.logger, + abortController: input.abortController, }); const handle = new TurnHandle({ @@ -498,6 +407,70 @@ 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; + }): 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 orchestrator = new AgentThreadOrchestrator({ + agentThreads, + createDynamicSubAgentThread, + tracing, + logger: input.resolver.logger, + }); + + return { orchestrator, agentThreads }; + } + private async resolveAgentDefinitions(input: { previous: TurnRecord | undefined; resolver: ITurnResourceResolver; diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index df6b0db3e..00839edc1 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -2,15 +2,24 @@ * DB-backed turns API (mounted at /api/v1/sessions). */ import { OpenAPIHono, type RouteHandler } from '@hono/zod-openapi'; -import type { ISessionStore, Sessions, Turn, TurnStreamingEvent } from '@truefoundry/trueforge-core/agent-session'; +import type { + ISessionStore, + Sessions, + Turn, + TurnRecord, + TurnStreamingEvent, +} from '@truefoundry/trueforge-core/agent-session'; import { CancellationReason, EventType, resolvePreviousTurnId, SessionStoreConflictError, SessionStoreNotFoundError, + TurnEventAlreadyExistsError, TurnOffloadReason, + TurnOwnershipLostError, TurnResourceResolver, + TurnUserEventSchema, type SessionHandle, type TurnHandle, type TurnRecordWithoutSnapshot, @@ -34,6 +43,10 @@ import { } from '@truefoundry/trueforge-core/core'; import { jsonValueFrom, + NO_RESPONDER_PUBLISH_DETAIL, + NoResponderError, + redisRequest, + RequestTimeoutError, type RedisClient, type RouteHandler as RequestReplyRouteHandler, type RequestReplyRouter, @@ -42,6 +55,7 @@ import type { Context } from 'hono'; import { HTTPException } from 'hono/http-exception'; import { streamSSE } from 'hono/streaming'; import type { Logger } from 'winston'; +import { z } from 'zod'; import type { Authorizer } from '../auth/authorizer'; import { requestContextFromWire, type RequestContext, type ResolveRequestContext } from '../auth/identity'; import configuration from '../config'; @@ -73,7 +87,7 @@ import { resolveSandboxEnvironment, resolveSandboxProvider, } from '../runtime/sessionResources'; -import { deliverTurnEvents, SendEventsPeerBodySchema, TURNS_SEND_EVENTS_PATH } from '../runtime/turnDelivery'; +import { PeerOutcome, resolveOwnershipAction, TurnOwnershipAction } from '../runtime/turnOwnership'; import { DEFAULT_SANDBOX_ENVIRONMENT_NAME } from '../schemas/sandboxEnvironment'; import { MAX_SESSION_TITLE_LENGTH } from '../schemas/session'; import { assertGatewayMetadataRequestHeaders } from '../truefoundry/gatewayMetadata'; @@ -81,6 +95,10 @@ import { newId } from '../utils/id'; import { resolveWebSearchProvider } from '../websearch/providers'; import { canReadAgentBoundResource, canReadSession } from './agentAccess'; +const FORBIDDEN_SESSION_ACCESS = 'Only the session creator can access this session'; +const FORBIDDEN_CREATE_TURN = 'Only the session creator can create turns'; +const MAX_OWNERSHIP_ATTEMPTS = 3; + export function toWireTurn(record: TurnRecordWithoutSnapshot): Turn { return { id: record.turn_id, @@ -461,12 +479,17 @@ 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; +} + export async function drainTurnEvents(input: { trackedStream: AsyncIterable; turnEventStream: EventSubscription; @@ -474,21 +497,22 @@ export async function drainTurnEvents(input: { turnId: string; maxExecutionTimer: NodeJS.Timeout; logger: Logger; - activeTurns?: ActiveTurnRegistry | undefined; + pausedKeepalive: TurnPausedKeepalive; onEvent?: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; }): Promise { - const { trackedStream, turnEventStream, sessionId, turnId, maxExecutionTimer, logger, activeTurns, onEvent } = input; + const { trackedStream, turnEventStream, sessionId, turnId, maxExecutionTimer, 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 (activeTurns && event.type === EventType.TURN_UPDATE) { + if (event.type === EventType.TURN_UPDATE) { if (event.state.status === 'paused') { - activeTurns.startPausedKeepalive({ sessionId, turnId }); + pausedKeepalive.start(); } else { - activeTurns.clearPausedKeepalive({ sessionId, turnId }); + pausedKeepalive.clear(); } } await onEvent?.(event, sequenceNumber); @@ -520,7 +544,63 @@ export interface TurnEventDrainInput { turnId: string; maxExecutionTimer: NodeJS.Timeout; logger: Logger; - activeTurns: ActiveTurnRegistry; + 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; + startExecutionTimeout: () => NodeJS.Timeout; +}> { + 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 }), + }, + }); + return { + abortController, + resolver, + startExecutionTimeout: () => { + const maxExecutionTimer = setTimeout(() => { + if (!abortController.signal.aborted) { + abortController.abort(CancellationReason.ServerExecutionTimeout); + } + }, configuration.SERVER_EXECUTION_TIMEOUT_SECONDS * 1000); + maxExecutionTimer.unref(); + return maxExecutionTimer; + }, + }; } /** @@ -531,36 +611,15 @@ 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 sessionAgent = session.record.agent; - - const abortController = new AbortController(); const tenant_id = session.tenant_id; - const webSearchProvider = await resolveWebSearchProvider({ - tenant_id, - 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, startExecutionTimeout } = 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 @@ -593,19 +652,13 @@ export async function beginTurnExecution( 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(); - + const maxExecutionTimer = startExecutionTimeout(); const trackedStream = deps.activeTurns.track({ abortController, stream: turn.stream(), @@ -624,7 +677,14 @@ export async function beginTurnExecution( turnId: turn.id, maxExecutionTimer, logger: deps.logger, - activeTurns: deps.activeTurns, + pausedKeepalive: { + start: () => { + deps.activeTurns.startPausedKeepalive({ sessionId, turnId: turn.id }); + }, + clear: () => { + deps.activeTurns.clearPausedKeepalive({ sessionId, turnId: turn.id }); + }, + }, }, }; } @@ -662,39 +722,14 @@ async function resumePausedTurn(input: { requestHeaders?: Record | undefined; deps: BeginTurnExecutionDeps; }): Promise<{ turn: TurnHandle; startDrain: () => void }> { - assertGatewayMetadataRequestHeaders(input.requestHeaders); - const abortController = new AbortController(); - const maxExecutionTimer = setTimeout(() => { - if (!abortController.signal.aborted) { - abortController.abort(CancellationReason.ServerExecutionTimeout); - } - }, configuration.SERVER_EXECUTION_TIMEOUT_SECONDS * 1000); - maxExecutionTimer.unref(); - - const sessionAgent = input.session.record.agent; - const webSearchProvider = await resolveWebSearchProvider({ - tenant_id: input.session.tenant_id, - store: input.deps.webSearchProviderStore, - }); - 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.requestContext.subject.id, + const { abortController, resolver, startExecutionTimeout } = await prepareTurnExecution({ 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 }), - }, + turnId: input.turnId, + userRef: input.requestContext.subject.id, + requestHeaders: input.requestHeaders, + deps: input.deps, }); + const maxExecutionTimer = startExecutionTimeout(); let turn: TurnHandle; try { @@ -748,12 +783,494 @@ async function resumePausedTurn(input: { turnId: turn.id, maxExecutionTimer, logger: input.deps.logger, - activeTurns: input.deps.activeTurns, + 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, + }); + }, + }, }); }, }; } +/** 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[] } + | { type: 'invalid'; message: string } + | { type: 'conflict'; message: string } + | { type: 'unavailable'; message: string } + | { type: 'not_owner' }; + +type ForwardedSendResult = SendEventsResult | { type: 'peer'; peerOutcome: PeerOutcome }; + +/** 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, + 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 }; +} + +/** Ask the owning replica to apply the events. A peer outcome means this replica must decide again. */ +async function forwardTurnEvents(input: { + redis: RedisClient | undefined; + owner: string; + sessionId: string; + turnId: string; + events: TurnUserEvent[]; + requestContext: RequestContext; + requestHeaders?: Record | undefined; + logger: Pick; +}): Promise { + if (!input.redis) { + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } + 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.owner, + 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, + }, + }); + if (reply.status === 200) { + const parsed = PeerEventsReplySchema.safeParse(reply.body); + if (!parsed.success) { + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } + return { type: 'applied', events: parsed.data.events }; + } + if (reply.status === 412) { + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } + if (reply.status === 409 || reply.status === 400) { + const parsed = PeerMessageSchema.safeParse(reply.body); + return { + type: reply.status === 400 ? 'invalid' : 'conflict', + message: parsed.success ? parsed.data.message : `Turn is not accepting events: ${input.turnId}`, + }; + } + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } catch (error) { + if (error instanceof NoResponderError) { + if (error.detail === NO_RESPONDER_PUBLISH_DETAIL) { + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } + return { type: 'peer', peerOutcome: PeerOutcome.NoResponder }; + } + if (error instanceof RequestTimeoutError) { + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } + input.logger.warn('Failed to reach owning executor for turn events', { + sessionId: input.sessionId, + turnId: input.turnId, + owner: input.owner, + }); + return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + } +} + +/** + * Owner path. `resolveOwnershipAction` is called as this executor with no peer + * outcome, so the result is run, rebuild, or reject. + */ +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; + const handle = input.deps.activeTurns.getTurnHandle({ sessionId, turnId }); + const action = resolveOwnershipAction({ + status: input.turn.state.status, + ownedByThisExecutor: true, + hasLocalActiveTurn: handle !== undefined, + peerOutcome: null, + }); + + 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 { + return await applyTurnEvents({ + handle, + events: input.events, + sessionStore: input.sessionStore, + sessionId, + turnId, + }); + } finally { + if (handle.state.status === 'paused') { + input.deps.activeTurns.startPausedKeepalive({ sessionId, turnId }); + } + } + } + case TurnOwnershipAction.Rebuild: { + const resumed = await resumePausedTurn({ + session: input.session, + turnId, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + deps: input.deps, + }); + try { + return await applyTurnEvents({ + handle: resumed.turn, + events: input.events, + sessionStore: input.sessionStore, + sessionId, + turnId, + }); + } finally { + resumed.startDrain(); + } + } + case TurnOwnershipAction.Reject: + return { + type: 'conflict', + message: + input.turn.state.status === 'running' + ? `Turn is running but its executor is gone: ${input.turn.turn_id}` + : `Turn is not accepting events: ${input.turn.turn_id}`, + }; + case TurnOwnershipAction.Forward: + case TurnOwnershipAction.Steal: + case TurnOwnershipAction.Retry: + throw new Error(`Unreachable ownership action for an owned turn: ${action}`); + } +} + +/** + * One ownership pass of a user send. `'retry'` means the row moved and the caller should load it again. + * A failed peer returns unavailable and does not steal. + */ +async function attemptSendTurnEvents(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) { + return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; + } + if (turn.active_executor_id === configuration.EXECUTOR_ID) { + try { + return await rebuildOptionallyAndRun({ + session: input.session, + turn, + events: input.events, + sessionStore: input.sessionStore, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + deps: input.deps, + }); + } catch (error) { + if (error instanceof TurnOwnershipLostError) { + return 'retry'; + } + throw error; + } + } + + await input.deps.activeTurns.dropLocalTurn({ sessionId, turnId: input.turnId }); + const forwarded = await forwardTurnEvents({ + redis: input.deps.redis, + owner: turn.active_executor_id, + sessionId, + turnId: input.turnId, + events: input.events, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + logger: input.deps.logger, + }); + if (forwarded.type !== 'peer') { + return forwarded; + } + const fresh = await input.session.getTurn(input.turnId); + if (fresh === undefined) { + return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; + } + // Timeout, 412, publish-to-zero, and transport errors do not prove the owner + // is gone. Fail here. Another pass could steal. + if (forwarded.peerOutcome !== PeerOutcome.NoResponder) { + if (fresh.active_executor_id === configuration.EXECUTOR_ID) { + try { + return await rebuildOptionallyAndRun({ + session: input.session, + turn: fresh, + events: input.events, + sessionStore: input.sessionStore, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + deps: input.deps, + }); + } catch (error) { + if (error instanceof TurnOwnershipLostError) { + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } + throw error; + } + } + const failedPeer = resolveOwnershipAction({ + status: fresh.state.status, + ownedByThisExecutor: false, + hasLocalActiveTurn: false, + peerOutcome: forwarded.peerOutcome, + }); + if (failedPeer === TurnOwnershipAction.Reject) { + return { + type: 'conflict', + message: + fresh.state.status === 'running' + ? `Turn is running but its executor is gone: ${fresh.turn_id}` + : `Turn is not accepting events: ${fresh.turn_id}`, + }; + } + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } + if (fresh.state.status !== 'paused' || fresh.active_executor_id !== turn.active_executor_id) { + return 'retry'; + } + const action = resolveOwnershipAction({ + status: fresh.state.status, + ownedByThisExecutor: fresh.active_executor_id === configuration.EXECUTOR_ID, + hasLocalActiveTurn: false, + peerOutcome: forwarded.peerOutcome, + }); + if (action === TurnOwnershipAction.Steal) { + const won = await input.sessionStore.claimTurnExecutor({ + session_id: sessionId, + turn_id: input.turnId, + expected_executor_id: fresh.active_executor_id, + new_executor_id: configuration.EXECUTOR_ID, + }); + if (!won) { + return 'retry'; + } + try { + return await rebuildOptionallyAndRun({ + session: input.session, + turn: fresh, + events: input.events, + sessionStore: input.sessionStore, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + deps: input.deps, + }); + } catch (error) { + if (error instanceof TurnOwnershipLostError) { + return 'retry'; + } + throw error; + } + } + if (action === TurnOwnershipAction.Reject) { + return { type: 'conflict', message: `Turn is not accepting events: ${fresh.turn_id}` }; + } + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; +} + +/** User send. Forwards to the owner and may steal a paused turn whose heartbeat is gone. */ +async function sendTurnEvents(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 attemptSendTurnEvents(input); + if (result !== 'retry') { + return result; + } + } + return { type: 'unavailable', message: `Turn ownership did not settle: ${input.turnId}` }; +} + +/** Peered send. Runs or rebuilds only when this executor already owns the turn. */ +async function acceptOwnedTurnEvents(input: { + session: SessionHandle; + sessionStore: Pick; + turnId: string; + events: TurnUserEvent[]; + requestContext: RequestContext; + requestHeaders: Record | undefined; + deps: BeginTurnExecutionDeps; +}): Promise { + const turn = await input.session.getTurn(input.turnId); + if (!turn) { + return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; + } + if (turn.active_executor_id !== configuration.EXECUTOR_ID) { + return { type: 'not_owner' }; + } + try { + return await rebuildOptionallyAndRun({ + session: input.session, + turn, + events: input.events, + sessionStore: input.sessionStore, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + deps: input.deps, + }); + } catch (error) { + if (error instanceof TurnOwnershipLostError) { + return { type: 'not_owner' }; + } + throw error; + } +} + +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' } }; + } + const session = await deps.sessions.get({ + tenant_id: parsed.data.tenant_id, + session_id: parsed.data.session_id, + }); + if (!session) { + return { status: 404, body: { message: `Session not found: ${parsed.data.session_id}` } }; + } + const result = await deps.activeTurns.withTurnLock( + { sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }, + async () => { + const requestContext = requestContextFromWire(parsed.data.request_context); + const runAsAgent = await referencedRunAsAgent({ + session, + agentStore: deps.resolveAgentStore(requestContext), + }); + const executionDeps = await executionDepsForCaller({ + deps, + session, + requestContext, + requestHeaders: parsed.data.request_headers, + ...(runAsAgent !== undefined ? { runAsAgent } : {}), + }); + return acceptOwnedTurnEvents({ + session, + sessionStore: deps.sessionStore, + turnId: parsed.data.turn_id, + events: parsed.data.events, + requestContext, + requestHeaders: parsed.data.request_headers, + deps: executionDeps, + }); + }, + ); + if (result.type === 'applied') { + return { status: 200, body: jsonValueFrom({ events: result.events }) }; + } + if (result.type === 'invalid') { + return { status: 400, body: { message: result.message } }; + } + if (result.type === 'not_owner') { + return { status: 412, body: { message: 'Turn is not running on this executor' } }; + } + if (result.type === 'conflict') { + return { status: 409, body: { message: result.message } }; + } + return { status: 503, body: { message: result.message } }; + }; +} + /** Mapped client error from turn execution; undefined means rethrow. */ export interface TurnExecutionError { status: 400 | 404 | 422; @@ -819,9 +1336,6 @@ 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'; - /** DB-backed turns (mounted at /api/v1/sessions). */ export function createTurnsRouter(deps: TurnsRouterDeps) { const listTurnsHandler: RouteHandler = async c => { @@ -1238,25 +1752,14 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { requestHeaders, ...(runAsAgent !== undefined ? { runAsAgent } : {}), }); - const result = await deliverTurnEvents({ + const result = await sendTurnEvents({ session, sessionStore: deps.sessionStore, - activeTurns: deps.activeTurns, turnId, events, - redis: deps.redis, - logger: deps.logger, - allowSteal: true, requestContext, requestHeaders, - resume: () => - resumePausedTurn({ - session, - turnId, - requestContext, - requestHeaders, - deps: executionDeps, - }), + deps: executionDeps, }); if (result.type === 'applied') { return c.json({ data: result.events }, 201); @@ -1283,69 +1786,3 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { deps.requestReplyRouter.registerRoute(TURNS_SEND_EVENTS_PATH, handlePeerTurnEvents(deps)); return router; } - -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' } }; - } - const session = await deps.sessions.get({ - tenant_id: parsed.data.tenant_id, - session_id: parsed.data.session_id, - }); - if (!session) { - return { status: 404, body: { message: `Session not found: ${parsed.data.session_id}` } }; - } - const result = await deps.activeTurns.withTurnLock( - { sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }, - async () => { - const requestContext = requestContextFromWire(parsed.data.request_context); - const runAsAgent = await referencedRunAsAgent({ - session, - agentStore: deps.resolveAgentStore(requestContext), - }); - const executionDeps = await executionDepsForCaller({ - deps, - session, - requestContext, - requestHeaders: parsed.data.request_headers, - ...(runAsAgent !== undefined ? { runAsAgent } : {}), - }); - return deliverTurnEvents({ - session, - sessionStore: deps.sessionStore, - activeTurns: deps.activeTurns, - turnId: parsed.data.turn_id, - events: parsed.data.events, - redis: undefined, - logger: deps.logger, - allowSteal: false, - requestContext, - requestHeaders: parsed.data.request_headers, - resume: () => - resumePausedTurn({ - session, - turnId: parsed.data.turn_id, - requestContext, - requestHeaders: parsed.data.request_headers, - deps: executionDeps, - }), - }); - }, - ); - if (result.type === 'applied') { - return { status: 200, body: jsonValueFrom({ events: result.events }) }; - } - if (result.type === 'invalid') { - return { status: 400, body: { message: result.message } }; - } - if (result.type === 'not_owner') { - return { status: 412, body: { message: 'Turn is not running on this executor' } }; - } - if (result.type === 'conflict') { - return { status: 409, body: { message: result.message } }; - } - return { status: 503, body: { message: result.message } }; - }; -} diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 79cad3168..3ad0be5a1 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -180,10 +180,9 @@ export class ActiveTurnRegistry { } this.keepaliveTimers.delete(key); const run = this.runs.get(key); - if (!isPausedRun(run)) { - return; + if (isPausedRun(run)) { + await this.finishRun(run, TurnOffloadReason); } - await this.finishRun(run, TurnOffloadReason); }); }, this.pausedKeepaliveMs); timer.unref(); diff --git a/packages/trueforge/src/runtime/cancelTurn.ts b/packages/trueforge/src/runtime/cancelTurn.ts index 0554d303a..65fb58a9b 100644 --- a/packages/trueforge/src/runtime/cancelTurn.ts +++ b/packages/trueforge/src/runtime/cancelTurn.ts @@ -58,6 +58,21 @@ export interface CancelTurnDeps { logger: Pick; } +/** Freeze a non-terminal turn; missing turns are a no-op (already gone). */ +async function freezeTurnIgnoringMissing( + session: Pick, + input: { turnId: string; reason: CancellationReason }, +): Promise { + try { + await session.freezeTurn({ turn_id: input.turnId, reason: input.reason }); + } catch (error) { + if (error instanceof TurnNotFoundError) { + return; + } + throw error; + } +} + /** * Cancels a `running` or `paused` turn wherever it lives: locally or on the * owning peer over Redis request-reply. Callers state the motive; default is a @@ -132,18 +147,3 @@ export async function cancelSessionTurn( await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); } } - -/** Freeze a non-terminal turn; missing turns are a no-op (already gone). */ -async function freezeTurnIgnoringMissing( - session: Pick, - input: { turnId: string; reason: CancellationReason }, -): Promise { - try { - await session.freezeTurn({ turn_id: input.turnId, reason: input.reason }); - } catch (error) { - if (error instanceof TurnNotFoundError) { - return; - } - throw error; - } -} diff --git a/packages/trueforge/src/runtime/turnDelivery.ts b/packages/trueforge/src/runtime/turnDelivery.ts deleted file mode 100644 index 89034fa1d..000000000 --- a/packages/trueforge/src/runtime/turnDelivery.ts +++ /dev/null @@ -1,328 +0,0 @@ -/** - * Sends inbound events to the turn that owns them: the local machine, a rebuild - * of a paused turn this replica owns, or a peer. Steal is a conditional claim - * after the owner's heartbeat is gone. - */ -import type { - ISessionStore, - SessionHandle, - TurnHandle, - TurnUserEvent, -} from '@truefoundry/trueforge-core/agent-session'; -import { - TurnEventAlreadyExistsError, - TurnOwnershipLostError, - TurnUserEventSchema, -} from '@truefoundry/trueforge-core/agent-session'; -import { AgentHarnessError } from '@truefoundry/trueforge-core/core'; -import { - jsonValueFrom, - NO_RESPONDER_PUBLISH_DETAIL, - NoResponderError, - redisRequest, - RequestTimeoutError, - type RedisClient, -} from '@truefoundry/trueforge-core/request-reply'; -import type { Logger } from 'winston'; -import { z } from 'zod'; -import type { RequestContext } from '../auth/identity'; -import configuration from '../config'; -import type { ActiveTurnRegistry } from './activeTurns'; -import { PeerOutcome, resolveOwnershipAction, TurnOwnershipAction } from './turnOwnership'; - -/** Request-reply path the owner serves for inbound turn events. */ -export const TURNS_SEND_EVENTS_PATH = 'turns/send-events'; - -const MAX_OWNERSHIP_ATTEMPTS = 4; - -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(), -}); - -export 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(), -}); - -export type TurnDeliveryResult = - | { type: 'applied'; events: TurnUserEvent[] } - | { type: 'invalid'; message: string } - | { type: 'conflict'; message: string } - | { type: 'unavailable'; message: string } - | { type: 'not_owner' }; - -export async function deliverTurnEvents(input: { - session: SessionHandle; - sessionStore: Pick; - activeTurns: ActiveTurnRegistry; - turnId: string; - events: TurnUserEvent[]; - redis: RedisClient | undefined; - logger: Pick; - /** - * Rebuild the paused turn this replica owns. Do not start execute yet: - * the caller queues this request's events first, then starts the drain. - */ - resume: () => Promise<{ turn: TurnHandle; startDrain: () => void }>; - /** Peer delivery is owner-local. Only the replica that received the client request may steal. */ - allowSteal: boolean; - requestContext: RequestContext; - requestHeaders?: Record | undefined; -}): Promise { - const sessionId = input.session.session_id; - - for (let attempt = 0; attempt < MAX_OWNERSHIP_ATTEMPTS; attempt++) { - const loaded = await input.session.getTurn(input.turnId); - if (!loaded) { - return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; - } - const status = loaded.state.status; - const owner = loaded.active_executor_id; - const ownedByThisExecutor = owner === configuration.EXECUTOR_ID; - - if (!ownedByThisExecutor) { - await input.activeTurns.dropLocalTurn({ sessionId, turnId: input.turnId }); - } - - if (!input.allowSteal && !ownedByThisExecutor) { - return { type: 'not_owner' }; - } - - const hasLocalActiveTurn = input.activeTurns.getTurnHandle({ sessionId, turnId: input.turnId }) !== undefined; - let action = resolveOwnershipAction({ - status, - ownedByThisExecutor, - hasLocalActiveTurn, - peerOutcome: null, - }); - - if (action === TurnOwnershipAction.Forward) { - const forwarded = await forwardTurnEvents({ - redis: input.redis, - owner, - sessionId, - turnId: input.turnId, - events: input.events, - requestContext: input.requestContext, - requestHeaders: input.requestHeaders, - logger: input.logger, - }); - if (forwarded.type !== 'peer') { - return forwarded; - } - const fresh = await input.session.getTurn(input.turnId); - const freshStatus = fresh?.state.status ?? status; - const freshOwner = fresh?.active_executor_id ?? owner; - if (forwarded.peerOutcome === PeerOutcome.NoResponder && (freshStatus !== 'paused' || freshOwner !== owner)) { - continue; - } - action = resolveOwnershipAction({ - status: freshStatus, - ownedByThisExecutor: freshOwner === configuration.EXECUTOR_ID, - hasLocalActiveTurn: false, - peerOutcome: forwarded.peerOutcome, - }); - if ( - action === TurnOwnershipAction.Forward || - action === TurnOwnershipAction.Run || - action === TurnOwnershipAction.Rebuild - ) { - continue; - } - } - - if (action === TurnOwnershipAction.Run) { - const handle = input.activeTurns.getTurnHandle({ sessionId, turnId: input.turnId }); - if (!handle) { - continue; - } - // 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.activeTurns.clearPausedKeepalive({ sessionId, turnId: input.turnId }); - try { - return await applyTurnEvents({ - handle, - events: input.events, - sessionStore: input.sessionStore, - sessionId, - turnId: input.turnId, - }); - } finally { - if (handle.state.status === 'paused') { - input.activeTurns.startPausedKeepalive({ sessionId, turnId: input.turnId }); - } - } - } - - if (action === TurnOwnershipAction.Rebuild || action === TurnOwnershipAction.Steal) { - if (action === TurnOwnershipAction.Steal) { - const won = await input.sessionStore.claimTurnExecutor({ - session_id: sessionId, - turn_id: input.turnId, - expected_executor_id: owner, - new_executor_id: configuration.EXECUTOR_ID, - }); - if (!won) { - continue; - } - } - try { - const resumed = await input.resume(); - try { - return await applyTurnEvents({ - handle: resumed.turn, - events: input.events, - sessionStore: input.sessionStore, - sessionId, - turnId: input.turnId, - }); - } finally { - resumed.startDrain(); - } - } catch (error) { - if (error instanceof TurnOwnershipLostError) { - continue; - } - throw error; - } - } - - if (action === TurnOwnershipAction.Reject) { - return { - type: 'conflict', - message: - status === 'running' - ? `Turn is running but its executor is gone: ${input.turnId}` - : `Turn is not accepting events: ${input.turnId}`, - }; - } - - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; - } - - return { type: 'unavailable', message: `Turn ownership did not settle: ${input.turnId}` }; -} - -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, - 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 }; -} - -type ForwardResult = TurnDeliveryResult | { type: 'peer'; peerOutcome: PeerOutcome }; - -async function forwardTurnEvents(input: { - redis: RedisClient | undefined; - owner: string; - sessionId: string; - turnId: string; - events: TurnUserEvent[]; - requestContext: RequestContext; - requestHeaders?: Record | undefined; - logger: Pick; -}): Promise { - if (!input.redis) { - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; - } - 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.owner, - 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, - }, - }); - if (reply.status === 200) { - const parsed = PeerEventsReplySchema.safeParse(reply.body); - if (!parsed.success) { - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; - } - return { type: 'applied', events: parsed.data.events }; - } - if (reply.status === 412) { - return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; - } - if (reply.status === 409 || reply.status === 400) { - const parsed = PeerMessageSchema.safeParse(reply.body); - return { - type: reply.status === 400 ? 'invalid' : 'conflict', - message: parsed.success ? parsed.data.message : `Turn is not accepting events: ${input.turnId}`, - }; - } - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; - } catch (error) { - if (error instanceof NoResponderError) { - if (error.detail === NO_RESPONDER_PUBLISH_DETAIL) { - return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; - } - return { type: 'peer', peerOutcome: PeerOutcome.NoResponder }; - } - if (error instanceof RequestTimeoutError) { - return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; - } - input.logger.warn('Failed to reach owning executor for turn events', { - sessionId: input.sessionId, - turnId: input.turnId, - owner: input.owner, - }); - return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; - } -} From e4e6ede4b245ac8402d15d538bbd90095d502dc4 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 19:27:59 +0530 Subject: [PATCH 14/32] Don't cancel during create turn, just freeze --- packages/trueforge/src/apis/turns.ts | 27 +++------------------------ 1 file changed, 3 insertions(+), 24 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 00839edc1..d4cd299dc 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -12,7 +12,6 @@ import type { import { CancellationReason, EventType, - resolvePreviousTurnId, SessionStoreConflictError, SessionStoreNotFoundError, TurnEventAlreadyExistsError, @@ -77,7 +76,6 @@ import { subscribeTurnRoute, } from '../routes/turnRoutes'; import type { ActiveTurnRegistry } from '../runtime/activeTurns'; -import { cancelSessionTurn } from '../runtime/cancelTurn'; import { StreamGoneError, type EventSubscription, type EventSubscriptionRegistry } from '../runtime/event-subscription'; import { validateSandboxFilePath } from '../runtime/sandboxFilePath'; import { @@ -626,28 +624,9 @@ export async function beginTurnExecution( // never overwrites an existing title. const title = session.record.last_turn_id ? undefined : deriveSessionTitle(input); - // Unpark a live previous turn so its drain emits turn.done before this turn freezes the row. - const previousId = resolvePreviousTurnId(previousTurnId, session.record.last_turn_id); - if (previousId) { - const hadLocalRun = await deps.activeTurns.abortAndWait({ - sessionId, - turnId: previousId, - abortReason: CancellationReason.CancelledForNextTurn, - }); - if (!hadLocalRun) { - await cancelSessionTurn( - { - activeTurns: deps.activeTurns, - session, - sessionStore: deps.sessionStore, - redis: deps.redis, - logger: deps.logger, - }, - { turnId: previousId, reason: CancellationReason.CancelledForNextTurn }, - ); - } - } - + // 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, From 8eaef0461c32df82fc3ace8ca266a62a7d31f557 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 19:38:19 +0530 Subject: [PATCH 15/32] Do not use send to queue up pending events --- .../src/agent-session/SessionHandle.ts | 17 +++++++----- .../src/agent-session/TurnHandle.ts | 26 +------------------ .../core/runtime/AgentThreadOrchestrator.ts | 5 +++- 3 files changed, 16 insertions(+), 32 deletions(-) diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index b1f5ba034..63c251e3a 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'; @@ -278,8 +278,8 @@ 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 queued before the handle - * is returned, so a later send() lands after them and stream() only executes. + * `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; @@ -308,21 +308,24 @@ export class SessionHandle< 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, + }); const { orchestrator } = await this.assembleTurn({ previous: turn, resolver: input.resolver, abortController: input.abortController, + pendingTurnEvents, }); - const handle = new TurnHandle({ + return new TurnHandle({ store: this.store, turn, orchestrator, resolver: input.resolver, abortController: input.abortController, }); - await handle.replayUnconsumedInbox(); - return handle; } catch (error) { await input.resolver.close().catch((closeError: unknown) => { input.resolver.logger.warn('TurnResourceResolver.close() failed after rebuildTurn() error', { @@ -415,6 +418,7 @@ export class SessionHandle< previous: TurnRecord | undefined; resolver: ITurnResourceResolver; abortController: AbortController; + pendingTurnEvents?: TurnUserEvent[] | undefined; }): Promise<{ orchestrator: AgentThreadOrchestrator; agentThreads: Map }> { const tracing = input.resolver.createTracing(); const agent = this.session.agent; @@ -466,6 +470,7 @@ export class SessionHandle< createDynamicSubAgentThread, tracing, logger: input.resolver.logger, + ...(input.pendingTurnEvents !== undefined ? { pendingTurnEvents: input.pendingTurnEvents } : {}), }); return { orchestrator, agentThreads }; diff --git a/packages/trueforge-core/src/agent-session/TurnHandle.ts b/packages/trueforge-core/src/agent-session/TurnHandle.ts index 6b3469de3..305cc4b04 100644 --- a/packages/trueforge-core/src/agent-session/TurnHandle.ts +++ b/packages/trueforge-core/src/agent-session/TurnHandle.ts @@ -1,7 +1,7 @@ /** * Durable turn handle. {@link TurnHandle.stream} is execute-once (persist-before-yield). */ -import { AgentHarnessError, InvalidAgentSendInputError } from '../core/errors'; +import { AgentHarnessError } from '../core/errors'; import type { MCPAuthRequiredEvent, MCPServerInitInfo, @@ -348,30 +348,6 @@ export class TurnHandle> { } } - /** - * Queue inbox rows onto the orchestrator before execute starts. - * A row the paused snapshot already applied fails validation; skip it so a - * crash between persist and consume does not apply the same decision twice. - */ - async replayUnconsumedInbox(): Promise { - const rows = await this.store.listUnconsumedTurnInboundEvents({ - session_id: this.turn.session_id, - turn_id: this.turn.turn_id, - }); - for (const event of rows) { - try { - for await (const batch of this.orchestrator.send([event])) { - void batch; - } - } catch (error) { - if (error instanceof InvalidAgentSendInputError) { - continue; - } - throw error; - } - } - } - private async persistTurnCreated(): Promise { const turnCreated: TurnCreatedEvent = { type: EventType.TURN_CREATED, diff --git a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts index 4189db2c9..0daf602fc 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts @@ -182,6 +182,8 @@ 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[] | undefined; } export class AgentThreadOrchestrator { @@ -192,7 +194,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 +203,7 @@ export class AgentThreadOrchestrator { this.createDynamicSubAgentThread = params.createDynamicSubAgentThread; this.tracing = params.tracing; this.logger = params.logger.child({ module: 'AgentThreadOrchestrator' }); + this.pendingTurnEvents = params.pendingTurnEvents !== undefined ? [...params.pendingTurnEvents] : []; } public wake(): void { From 86224d9fc4b313fa3635aab6f56ad12034abffea Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 20:29:12 +0530 Subject: [PATCH 16/32] Address more review comments --- .../src/agent-session/TurnHandle.ts | 109 +++++++----- .../src/agent-session/store/ISessionStore.ts | 7 +- .../store/InMemorySessionStore.ts | 1 + .../core/runtime/AgentThreadOrchestrator.ts | 7 +- .../agent-session/store/storeContractSuite.ts | 33 ++++ .../tests/agent-session/turnResume.test.ts | 1 + packages/trueforge/src/apis/sessions.ts | 162 +++++++++++++++++- packages/trueforge/src/apis/turns.ts | 29 ++-- packages/trueforge/src/app.ts | 36 ++-- .../session-store/queries/inboundEvents.ts | 12 +- .../session-store/queries/inboundEvents.ts | 12 +- packages/trueforge/src/runtime/activeTurns.ts | 43 +---- packages/trueforge/src/runtime/cancelTurn.ts | 149 ---------------- .../tests/unit/apis/sessions.test.ts | 4 +- .../trueforge/tests/unit/apis/turns.test.ts | 4 +- .../tests/unit/runtime/activeTurns.test.ts | 25 ++- .../truefoundry/TrueFoundryAgentStore.test.ts | 2 +- .../truefoundry/TrueFoundrySkillStore.test.ts | 2 +- .../truefoundry/trueFoundryNaming.test.ts | 2 +- 19 files changed, 349 insertions(+), 291 deletions(-) delete mode 100644 packages/trueforge/src/runtime/cancelTurn.ts diff --git a/packages/trueforge-core/src/agent-session/TurnHandle.ts b/packages/trueforge-core/src/agent-session/TurnHandle.ts index 305cc4b04..47a818be5 100644 --- a/packages/trueforge-core/src/agent-session/TurnHandle.ts +++ b/packages/trueforge-core/src/agent-session/TurnHandle.ts @@ -293,10 +293,9 @@ export class TurnHandle> { resolver: ITurnResourceResolver, signal: AbortSignal, ): AsyncGenerator { + const turnStatus = executableTurnStatus(this.turn.state); let caughtError: Error | undefined; let executeResult: AgentThreadExecutionResult | undefined; - let ownershipLost: TurnOwnershipLostError | undefined; - const turnStatus = executableTurnStatus(this.turn.state); try { if (turnStatus === 'running') { yield await this.persistTurnCreated(); @@ -304,50 +303,73 @@ export class TurnHandle> { 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; - // Offload ends the in-memory run only. A terminal write would cancel a turn that is still paused. - if (!ownershipLost && !isTurnOffloadReason(signal.reason)) { - 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 (ownershipLost) { + if (closeError !== undefined) { + throw closeError; + } + if (turnDone !== undefined) { + yield turnDone; + } + 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, @@ -430,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/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 0daf602fc..2be4b85f4 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts @@ -611,10 +611,9 @@ export class AgentThreadOrchestrator { return done(); } - // Take the queue, then empty it, so a send during this pass waits for the next loop. - const queued = this.pendingTurnEvents.slice(); - this.pendingTurnEvents.length = 0; - for (const event of queued) { + // 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: yield* this.applyApprovalPolicies(event); 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..88533b815 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, }); } diff --git a/packages/trueforge/src/apis/sessions.ts b/packages/trueforge/src/apis/sessions.ts index 55defe3c8..5820c4aa2 100644 --- a/packages/trueforge/src/apis/sessions.ts +++ b/packages/trueforge/src/apis/sessions.ts @@ -2,18 +2,36 @@ * DB-backed sessions APIs (mounted at /api/v1/sessions and /api/internal/sessions). */ import { OpenAPIHono, type RouteHandler } from '@hono/zod-openapi'; -import type { AgentSpec, ISessionStore, SessionRecord, Sessions } from '@truefoundry/trueforge-core/agent-session'; +import type { + AgentSpec, + ISessionStore, + SessionHandle, + SessionRecord, + Sessions, +} from '@truefoundry/trueforge-core/agent-session'; import { + CancellationReason, + isNonTerminalTurnState, SessionStoreConflictError, SessionStoreInvariantError, SessionStoreNotFoundError, + TurnNotFoundError, } from '@truefoundry/trueforge-core/agent-session'; -import { type RedisClient, type RequestReplyRouter } from '@truefoundry/trueforge-core/request-reply'; +import { extractErrorLogFields } from '@truefoundry/trueforge-core/core'; +import { + redisRequest, + RequestTimeoutError, + type RedisClient, + type RouteHandler as RequestReplyRouteHandler, + type RequestReplyRouter, +} from '@truefoundry/trueforge-core/request-reply'; import type { Context } from 'hono'; import { HTTPException } from 'hono/http-exception'; import type { Logger } from 'winston'; +import { z } from 'zod'; import type { Authorizer } from '../auth/authorizer'; import { createdBySubjectFromRequestContext, type RequestContext, type ResolveRequestContext } from '../auth/identity'; +import configuration from '../config'; import type { IAgentStore } from '../db/agentStore'; import type { IMcpServerStore } from '../db/mcpServerStore'; import type { IModelProviderStore } from '../db/modelProviderStore'; @@ -32,16 +50,12 @@ import { updateSessionRoute, } from '../routes/sessionRoutes'; import type { ActiveTurnRegistry } from '../runtime/activeTurns'; -import { cancelSessionTurn, cancelSessionTurnPeerHandler, SESSIONS_CANCEL_PATH } from '../runtime/cancelTurn'; import { validateAgentSpec } from '../runtime/sessionResources'; import { honoQueriesToRecord } from '../schemas/deepObjectQuery'; import { isSessionAgentNameRef, parseListSessionsQuery, type Session } from '../schemas/session'; import { newId } from '../utils/id'; import { agentIfAccessible, canReadAgentBoundResource, canReadSession, resolveManagedAgentIds } from './agentAccess'; -export { cancelSessionTurn, cancelSessionTurnPeerHandler, SESSIONS_CANCEL_PATH } from '../runtime/cancelTurn'; -export type { CancelTurnDeps } from '../runtime/cancelTurn'; - export function toWireSession(record: SessionRecord): Session { return { id: record.session_id, @@ -220,6 +234,142 @@ export function createInternalSessionsRouter(deps: InternalSessionsRouterDeps) { return router; } +/** Request-reply path a replica serves to cancel a turn it owns. */ +export const SESSIONS_CANCEL_PATH = 'sessions/cancel'; + +/** Wire body of a peer cancel; validated on receipt (it crosses processes via Redis). */ +const CancelPeerBodySchema = z.object({ + session_id: z.string(), + turn_id: z.string(), + reason: z.enum(CancellationReason), +}); +type CancelPeerBody = z.infer; + +/** + * 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 { + return request => { + const parsed = CancelPeerBodySchema.safeParse(request.body); + if (!parsed.success) { + return Promise.resolve({ status: 400, body: { message: 'Invalid sessions/cancel payload' } }); + } + // 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, + abortReason: parsed.data.reason, + }); + if (!found) { + return Promise.resolve({ status: 412, body: { message: 'Turn is not running on this executor' } }); + } + return Promise.resolve({ status: 200, body: {} }); + }; +} + +/** A registry to abort in, a session to freeze, durable state to read, and a way to reach peers. */ +export interface CancelTurnDeps { + activeTurns: ActiveTurnRegistry; + session: Pick; + sessionStore: Pick; + redis?: RedisClient | undefined; + logger: Pick; +} + +/** Freeze a non-terminal turn; missing turns are a no-op (already gone). */ +async function freezeTurnIgnoringMissing( + session: Pick, + input: { turnId: string; reason: CancellationReason }, +): Promise { + try { + await session.freezeTurn({ turn_id: input.turnId, reason: input.reason }); + } catch (error) { + if (error instanceof TurnNotFoundError) { + return; + } + throw error; + } +} + +/** + * Cancels a `running` or `paused` turn wherever it lives: locally or on the + * owning peer over Redis request-reply. Callers state the motive; default is a + * plain client cancel. + * + * 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. + * Later writes from that replica lose to first-terminal-write-wins. + */ +export async function cancelSessionTurn( + deps: CancelTurnDeps, + input: { turnId: string; reason?: CancellationReason }, +): Promise { + const { turnId, reason = CancellationReason.ClientCancelled } = input; + const sessionId = deps.session.session_id; + + const turn = await deps.sessionStore.getTurn({ + session_id: sessionId, + turn_id: turnId, + }); + if (!turn || !isNonTerminalTurnState(turn.state)) { + return; + } + + const owner = turn.active_executor_id; + // Without a Redis client there is no peer to ask, so a different owner falls + // through to the local lookup and freezes if the run is gone. + if (owner !== configuration.EXECUTOR_ID && deps.redis) { + try { + const reply = await redisRequest({ + redis: deps.redis, + executorId: owner, + path: SESSIONS_CANCEL_PATH, + request: { + body: { session_id: sessionId, turn_id: turnId, reason }, + }, + options: { + replyTimeoutMs: configuration.REDIS_REQUEST_REPLY_TIMEOUT_MS, + pollIntervalMs: configuration.REDIS_REQUEST_REPLY_POLL_INTERVAL_MS, + }, + }); + if (reply.status === 200) { + return; + } + } catch (error) { + const fields = { + sessionId, + turnId, + owner, + ...extractErrorLogFields(error), + }; + if (error instanceof RequestTimeoutError) { + deps.logger.warn('Timed out waiting for owning executor to cancel; freezing the turn', fields); + } else { + deps.logger.warn('Failed to reach owning executor over Redis; freezing the turn', fields); + } + } + await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); + return; + } + + const aborted = deps.activeTurns.cancelIfRunning({ + sessionId, + turnId, + abortReason: reason, + }); + if (!aborted) { + await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); + } +} + /** DB-backed sessions (mounted at /api/v1/sessions). */ export function createSessionsRouter(deps: SessionsRouterDeps) { const createSessionHandler: RouteHandler = async c => { diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index d4cd299dc..5bb3012fe 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -834,6 +834,7 @@ async function applyTurnEvents(input: { await input.sessionStore.insertTurnInboundEvents({ session_id: input.sessionId, turn_id: input.turnId, + active_executor_id: input.handle.record.active_executor_id, events: batch, }); } @@ -1057,18 +1058,18 @@ async function attemptSendTurnEvents(input: { if (forwarded.type !== 'peer') { return forwarded; } - const fresh = await input.session.getTurn(input.turnId); - if (fresh === undefined) { + const refetchedTurn = await input.session.getTurn(input.turnId); + if (refetchedTurn === undefined) { return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; } // Timeout, 412, publish-to-zero, and transport errors do not prove the owner // is gone. Fail here. Another pass could steal. if (forwarded.peerOutcome !== PeerOutcome.NoResponder) { - if (fresh.active_executor_id === configuration.EXECUTOR_ID) { + if (refetchedTurn.active_executor_id === configuration.EXECUTOR_ID) { try { return await rebuildOptionallyAndRun({ session: input.session, - turn: fresh, + turn: refetchedTurn, events: input.events, sessionStore: input.sessionStore, requestContext: input.requestContext, @@ -1083,7 +1084,7 @@ async function attemptSendTurnEvents(input: { } } const failedPeer = resolveOwnershipAction({ - status: fresh.state.status, + status: refetchedTurn.state.status, ownedByThisExecutor: false, hasLocalActiveTurn: false, peerOutcome: forwarded.peerOutcome, @@ -1092,19 +1093,19 @@ async function attemptSendTurnEvents(input: { return { type: 'conflict', message: - fresh.state.status === 'running' - ? `Turn is running but its executor is gone: ${fresh.turn_id}` - : `Turn is not accepting events: ${fresh.turn_id}`, + refetchedTurn.state.status === 'running' + ? `Turn is running but its executor is gone: ${refetchedTurn.turn_id}` + : `Turn is not accepting events: ${refetchedTurn.turn_id}`, }; } return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } - if (fresh.state.status !== 'paused' || fresh.active_executor_id !== turn.active_executor_id) { + if (refetchedTurn.state.status !== 'paused' || refetchedTurn.active_executor_id !== turn.active_executor_id) { return 'retry'; } const action = resolveOwnershipAction({ - status: fresh.state.status, - ownedByThisExecutor: fresh.active_executor_id === configuration.EXECUTOR_ID, + status: refetchedTurn.state.status, + ownedByThisExecutor: refetchedTurn.active_executor_id === configuration.EXECUTOR_ID, hasLocalActiveTurn: false, peerOutcome: forwarded.peerOutcome, }); @@ -1112,7 +1113,7 @@ async function attemptSendTurnEvents(input: { const won = await input.sessionStore.claimTurnExecutor({ session_id: sessionId, turn_id: input.turnId, - expected_executor_id: fresh.active_executor_id, + expected_executor_id: refetchedTurn.active_executor_id, new_executor_id: configuration.EXECUTOR_ID, }); if (!won) { @@ -1121,7 +1122,7 @@ async function attemptSendTurnEvents(input: { try { return await rebuildOptionallyAndRun({ session: input.session, - turn: fresh, + turn: refetchedTurn, events: input.events, sessionStore: input.sessionStore, requestContext: input.requestContext, @@ -1136,7 +1137,7 @@ async function attemptSendTurnEvents(input: { } } if (action === TurnOwnershipAction.Reject) { - return { type: 'conflict', message: `Turn is not accepting events: ${fresh.turn_id}` }; + return { type: 'conflict', message: `Turn is not accepting events: ${refetchedTurn.turn_id}` }; } return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } diff --git a/packages/trueforge/src/app.ts b/packages/trueforge/src/app.ts index a387e25e0..0313549fb 100644 --- a/packages/trueforge/src/app.ts +++ b/packages/trueforge/src/app.ts @@ -411,26 +411,24 @@ export function createServerApp(deps: ServerDeps) { 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({ - 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, - }), - scheduleExecutionAuthMiddleware, - ), + withAuth(createScheduleExecutionRouter(scheduleTurnExecutionRouterDeps), scheduleExecutionAuthMiddleware), ); app.route( '/api/internal/sandbox-environments', @@ -447,7 +445,7 @@ export function createServerApp(deps: ServerDeps) { '/api/v1/schedules', withAuth( createSchedulesRouter({ - ...scheduleTurnDeps, + ...scheduleTurnExecutionRouterDeps, authorizer: deps.authorizer, resolveAgentStore: resolveHttpAgentStore, resolveRequestContext, 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/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 3ad0be5a1..10d5ed093 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -20,15 +20,8 @@ function activeTurnKey(sessionId: string, turnId: string): string { return `${sessionId}:${turnId}`; } -/** - * `state` is a getter on TurnHandle's prototype, so it is not an own property. - * A test double may omit it; that run is not treated as paused. - */ function isPausedRun(run: ActiveTurnRun | undefined): run is ActiveTurnRun { - if (!run || !('state' in run.turn)) { - return false; - } - return run.turn.state.status === 'paused'; + return run?.turn.state.status === 'paused'; } export class ActiveTurnRegistry { @@ -69,8 +62,8 @@ 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); @@ -86,12 +79,11 @@ export class ActiveTurnRegistry { }; 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 reason = isPausedRun(run) ? TurnOffloadReason : this.alreadyShutDownAbortReason; + input.abortController.abort(reason); } const complete = () => { @@ -128,23 +120,6 @@ export class ActiveTurnRegistry { return true; } - /** - * Abort a local run and wait until its drain has finished. - * A paused generator is parked on the abort latch, so this unparks it and lets - * the drain emit `turn.done`. Returns false when this process has no run. - */ - async abortAndWait(input: { sessionId: string; turnId: string; abortReason: CancellationReason }): Promise { - const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); - if (!run) { - return false; - } - if (!run.abortController.signal.aborted) { - run.abortController.abort(input.abortReason); - } - await run.waitUntilCompleted; - return true; - } - /** * 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 @@ -214,7 +189,7 @@ export class ActiveTurnRegistry { /** * Enter shutdown mode and wait until the registry is empty. * Paused runs are offloaded so their approval state survives. Running runs abort as `abortReason`. - * Late `track()` calls abort immediately as `abortReason`. + * A later `track()` offloads a paused turn and aborts any other turn as `abortReason`. */ async shutdownAndWait(abortReason: CancellationReason): Promise { this.alreadyShutDownAbortReason = abortReason; diff --git a/packages/trueforge/src/runtime/cancelTurn.ts b/packages/trueforge/src/runtime/cancelTurn.ts deleted file mode 100644 index 65fb58a9b..000000000 --- a/packages/trueforge/src/runtime/cancelTurn.ts +++ /dev/null @@ -1,149 +0,0 @@ -import type { ISessionStore, SessionHandle } from '@truefoundry/trueforge-core/agent-session'; -import { - CancellationReason, - isNonTerminalTurnState, - TurnNotFoundError, -} from '@truefoundry/trueforge-core/agent-session'; -import { extractErrorLogFields } from '@truefoundry/trueforge-core/core'; -import { - redisRequest, - RequestTimeoutError, - type RedisClient, - type RouteHandler as RequestReplyRouteHandler, -} from '@truefoundry/trueforge-core/request-reply'; -import type { Logger } from 'winston'; -import { z } from 'zod'; -import configuration from '../config'; -import type { ActiveTurnRegistry } from './activeTurns'; - -/** Request-reply path a replica serves to cancel a turn it owns. */ -export const SESSIONS_CANCEL_PATH = 'sessions/cancel'; - -/** Wire body of a peer cancel; validated on receipt (it crosses processes via Redis). */ -const CancelPeerBodySchema = z.object({ - session_id: z.string(), - turn_id: z.string(), - reason: z.enum(CancellationReason), -}); -type CancelPeerBody = z.infer; - -/** - * Peer-facing cancel handler. 200 means the local run has finished, so its drain - * has emitted `turn.done`. 412 means it is not running here. - */ -export function cancelSessionTurnPeerHandler(activeTurns: ActiveTurnRegistry): RequestReplyRouteHandler { - return async request => { - const parsed = CancelPeerBodySchema.safeParse(request.body); - if (!parsed.success) { - return { status: 400, body: { message: 'Invalid sessions/cancel payload' } }; - } - const found = await activeTurns.abortAndWait({ - sessionId: parsed.data.session_id, - turnId: parsed.data.turn_id, - abortReason: parsed.data.reason, - }); - if (!found) { - return { status: 412, body: { message: 'Turn is not running on this executor' } }; - } - return { status: 200, body: {} }; - }; -} - -/** A registry to abort in, a session to freeze, durable state to read, and a way to reach peers. */ -export interface CancelTurnDeps { - activeTurns: ActiveTurnRegistry; - session: Pick; - sessionStore: Pick; - redis?: RedisClient | undefined; - logger: Pick; -} - -/** Freeze a non-terminal turn; missing turns are a no-op (already gone). */ -async function freezeTurnIgnoringMissing( - session: Pick, - input: { turnId: string; reason: CancellationReason }, -): Promise { - try { - await session.freezeTurn({ turn_id: input.turnId, reason: input.reason }); - } catch (error) { - if (error instanceof TurnNotFoundError) { - return; - } - throw error; - } -} - -/** - * Cancels a `running` or `paused` turn wherever it lives: locally or on the - * 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`. - * - * Redis timeout and Redis/transport failures are not a clean cancellation — - * the owning replica may still be executing — but the turn is still frozen. - * Later writes from that replica lose to first-terminal-write-wins. - */ -export async function cancelSessionTurn( - deps: CancelTurnDeps, - input: { turnId: string; reason?: CancellationReason }, -): Promise { - const { turnId, reason = CancellationReason.ClientCancelled } = input; - const sessionId = deps.session.session_id; - - const turn = await deps.sessionStore.getTurn({ - session_id: sessionId, - turn_id: turnId, - }); - if (!turn || !isNonTerminalTurnState(turn.state)) { - return; - } - - const owner = turn.active_executor_id; - // Without a Redis client there is no peer to ask, so a different owner falls - // through to the local lookup and freezes if the run is gone. - if (owner !== configuration.EXECUTOR_ID && deps.redis) { - try { - const reply = await redisRequest({ - redis: deps.redis, - executorId: owner, - path: SESSIONS_CANCEL_PATH, - request: { - body: { session_id: sessionId, turn_id: turnId, reason }, - }, - options: { - replyTimeoutMs: configuration.REDIS_REQUEST_REPLY_TIMEOUT_MS, - pollIntervalMs: configuration.REDIS_REQUEST_REPLY_POLL_INTERVAL_MS, - }, - }); - if (reply.status === 200) { - return; - } - } catch (error) { - const fields = { - sessionId, - turnId, - owner, - ...extractErrorLogFields(error), - }; - if (error instanceof RequestTimeoutError) { - deps.logger.warn('Timed out waiting for owning executor to cancel; freezing the turn', fields); - } else { - deps.logger.warn('Failed to reach owning executor over Redis; freezing the turn', fields); - } - } - await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); - return; - } - - const aborted = deps.activeTurns.cancelIfRunning({ - sessionId, - turnId, - abortReason: reason, - }); - if (!aborted) { - await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); - } -} diff --git a/packages/trueforge/tests/unit/apis/sessions.test.ts b/packages/trueforge/tests/unit/apis/sessions.test.ts index 5bd230e13..ee9707d70 100644 --- a/packages/trueforge/tests/unit/apis/sessions.test.ts +++ b/packages/trueforge/tests/unit/apis/sessions.test.ts @@ -86,7 +86,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); @@ -143,7 +143,7 @@ describe('cancelSessionTurn', () => { 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'; diff --git a/packages/trueforge/tests/unit/apis/turns.test.ts b/packages/trueforge/tests/unit/apis/turns.test.ts index 5d0b3f351..c5115e753 100644 --- a/packages/trueforge/tests/unit/apis/turns.test.ts +++ b/packages/trueforge/tests/unit/apis/turns.test.ts @@ -3,8 +3,8 @@ 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'; @@ -53,7 +53,7 @@ describe('turns', () => { }); describe('deriveSessionTitle', () => { - function userMessage(content: string): TurnInputItem[] { + function userMessage(content: string): UserMessage[] { return [{ type: 'user.message', content }]; } diff --git a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts index 6e78520ce..7b7a9123d 100644 --- a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts +++ b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts @@ -1,8 +1,8 @@ -import { CancellationReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; +import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; import { ActiveTurnRegistry } from '../../../src/runtime/activeTurns'; -function turn(sessionId: string, turnId: string): TurnHandle { - return { session_id: sessionId, id: turnId } as TurnHandle; +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 { @@ -196,4 +196,23 @@ describe('ActiveTurnRegistry', () => { void value; } }); + + it('late track after shutdownAndWait offloads a paused turn', async () => { + const registry = new ActiveTurnRegistry(); + 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/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); From ac7ccfa094746ae898f56cae7ad281f781c0f44a Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 20:45:04 +0530 Subject: [PATCH 17/32] Address more review comments --- packages/trueforge/src/apis/turns.ts | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 5bb3012fe..fa8889d50 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -12,6 +12,7 @@ import type { import { CancellationReason, EventType, + isNonTerminalTurnState, SessionStoreConflictError, SessionStoreNotFoundError, TurnEventAlreadyExistsError, @@ -1025,6 +1026,10 @@ async function attemptSendTurnEvents(input: { if (!turn) { return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; } + // done, cancelled, and error cannot accept a send. Retrying them becomes a 503. + if (!isNonTerminalTurnState(turn.state)) { + return { type: 'conflict', message: `Turn is not accepting events: ${turn.turn_id}` }; + } if (turn.active_executor_id === configuration.EXECUTOR_ID) { try { return await rebuildOptionallyAndRun({ @@ -1100,6 +1105,11 @@ async function attemptSendTurnEvents(input: { } return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } + // The owner was gone. A row that finished during the peer call is a conflict. + // A running row, or a paused row whose owner changed, is loaded again. + if (!isNonTerminalTurnState(refetchedTurn.state)) { + return { type: 'conflict', message: `Turn is not accepting events: ${refetchedTurn.turn_id}` }; + } if (refetchedTurn.state.status !== 'paused' || refetchedTurn.active_executor_id !== turn.active_executor_id) { return 'retry'; } From c6389ef959e5c96fe76d7a300e50cd7943cf46d2 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 21:06:57 +0530 Subject: [PATCH 18/32] Run the execution timer only during running phase --- packages/trueforge/src/apis/turns.ts | 65 +++++++++++++------ packages/trueforge/src/config.ts | 6 +- packages/trueforge/src/runtime/activeTurns.ts | 3 + 3 files changed, 52 insertions(+), 22 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index fa8889d50..075781064 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -489,17 +489,26 @@ interface TurnPausedKeepalive { 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; sessionId: string; turnId: string; - maxExecutionTimer: NodeJS.Timeout; + executionTimeout: TurnExecutionTimeout; logger: Logger; pausedKeepalive: TurnPausedKeepalive; onEvent?: (event: TurnStreamingEvent, sequenceNumber: number) => Promise; }): Promise { - const { trackedStream, turnEventStream, sessionId, turnId, maxExecutionTimer, logger, pausedKeepalive, onEvent } = + const { trackedStream, turnEventStream, sessionId, turnId, executionTimeout, logger, pausedKeepalive, onEvent } = input; try { for await (const event of trackedStream) { @@ -510,8 +519,10 @@ export async function drainTurnEvents(input: { if (event.type === EventType.TURN_UPDATE) { if (event.state.status === 'paused') { pausedKeepalive.start(); + executionTimeout.clear(); } else { pausedKeepalive.clear(); + executionTimeout.start(); } } await onEvent?.(event, sequenceNumber); @@ -531,7 +542,7 @@ export async function drainTurnEvents(input: { }); } } finally { - clearTimeout(maxExecutionTimer); + executionTimeout.clear(); } } @@ -541,7 +552,7 @@ export interface TurnEventDrainInput { turnEventStream: EventSubscription; sessionId: string; turnId: string; - maxExecutionTimer: NodeJS.Timeout; + executionTimeout: TurnExecutionTimeout; logger: Logger; pausedKeepalive: TurnPausedKeepalive; } @@ -559,7 +570,7 @@ async function prepareTurnExecution(input: { }): Promise<{ abortController: AbortController; resolver: TurnResourceResolver; - startExecutionTimeout: () => NodeJS.Timeout; + executionTimeout: TurnExecutionTimeout; }> { assertGatewayMetadataRequestHeaders(input.requestHeaders); const abortController = new AbortController(); @@ -587,19 +598,33 @@ async function prepareTurnExecution(input: { ...(input.requestHeaders === undefined ? {} : { requestHeaders: input.requestHeaders }), }, }); - return { - abortController, - resolver, - startExecutionTimeout: () => { - const maxExecutionTimer = setTimeout(() => { + 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); - maxExecutionTimer.unref(); - return maxExecutionTimer; + executionTimer.unref(); + }, + clear: () => { + if (executionTimer === undefined) { + return; + } + clearTimeout(executionTimer); + executionTimer = undefined; }, }; + return { + abortController, + resolver, + executionTimeout, + }; } /** @@ -613,7 +638,7 @@ export async function beginTurnExecution( const sessionId = session.session_id; const turnId = newId(); const tenant_id = session.tenant_id; - const { abortController, resolver, startExecutionTimeout } = await prepareTurnExecution({ + const { abortController, resolver, executionTimeout } = await prepareTurnExecution({ session, turnId, userRef, @@ -638,7 +663,8 @@ export async function beginTurnExecution( update_session_title_if_not_exist: title, }); - const maxExecutionTimer = startExecutionTimeout(); + // 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(), @@ -655,7 +681,7 @@ export async function beginTurnExecution( turnEventStream, sessionId, turnId: turn.id, - maxExecutionTimer, + executionTimeout, logger: deps.logger, pausedKeepalive: { start: () => { @@ -702,14 +728,13 @@ async function resumePausedTurn(input: { requestHeaders?: Record | undefined; deps: BeginTurnExecutionDeps; }): Promise<{ turn: TurnHandle; startDrain: () => void }> { - const { abortController, resolver, startExecutionTimeout } = await prepareTurnExecution({ + const { abortController, resolver, executionTimeout } = await prepareTurnExecution({ session: input.session, turnId: input.turnId, userRef: input.requestContext.subject.id, requestHeaders: input.requestHeaders, deps: input.deps, }); - const maxExecutionTimer = startExecutionTimeout(); let turn: TurnHandle; try { @@ -720,7 +745,7 @@ async function resumePausedTurn(input: { abortController, }); } catch (error) { - clearTimeout(maxExecutionTimer); + executionTimeout.clear(); throw error; } @@ -731,7 +756,7 @@ async function resumePausedTurn(input: { // 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) { - clearTimeout(maxExecutionTimer); + executionTimeout.clear(); abortController.abort(TurnOffloadReason); await turn.stream().return(undefined); throw error; @@ -761,7 +786,7 @@ async function resumePausedTurn(input: { turnEventStream, sessionId: input.session.session_id, turnId: turn.id, - maxExecutionTimer, + executionTimeout, logger: input.deps.logger, pausedKeepalive: { start: () => { diff --git a/packages/trueforge/src/config.ts b/packages/trueforge/src/config.ts index 7a439be6c..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; /** diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 10d5ed093..5ab404f4d 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -28,7 +28,10 @@ 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 alreadyShutDownAbortReason: CancellationReason | undefined; From 8863eb53e64f2b9629ba9228a3a151d856fd74cc Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 21:15:00 +0530 Subject: [PATCH 19/32] Address race conditions pointed by cursor --- packages/trueforge/src/apis/turns.ts | 72 +++++++++++- packages/trueforge/src/runtime/activeTurns.ts | 110 +++++++++++++++--- 2 files changed, 163 insertions(+), 19 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 075781064..53176705c 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -13,9 +13,11 @@ import { CancellationReason, EventType, isNonTerminalTurnState, + isTurnOffloadReason, SessionStoreConflictError, SessionStoreNotFoundError, TurnEventAlreadyExistsError, + TurnNotFoundError, TurnOffloadReason, TurnOwnershipLostError, TurnResourceResolver, @@ -717,6 +719,18 @@ 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. @@ -727,7 +741,7 @@ async function resumePausedTurn(input: { requestContext: RequestContext; requestHeaders?: Record | undefined; deps: BeginTurnExecutionDeps; -}): Promise<{ turn: TurnHandle; startDrain: () => void }> { +}): Promise<{ turn: TurnHandle; signal: AbortSignal; startDrain: () => void }> { const { abortController, resolver, executionTimeout } = await prepareTurnExecution({ session: input.session, turnId: input.turnId, @@ -735,6 +749,13 @@ async function resumePausedTurn(input: { 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 { @@ -746,6 +767,26 @@ async function resumePausedTurn(input: { }); } catch (error) { executionTimeout.clear(); + // Cancel found this reservation 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; } @@ -758,13 +799,19 @@ async function resumePausedTurn(input: { } catch (error) { executionTimeout.clear(); abortController.abort(TurnOffloadReason); - await turn.stream().return(undefined); + 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; @@ -999,6 +1046,22 @@ async function rebuildOptionallyAndRun(input: { } } case TurnOwnershipAction.Rebuild: { + // An aborted run still occupies the registry until its drain finishes. + // getTurnHandle hides it. Wait, then read the row the drain wrote. + await input.deps.activeTurns.waitForDrainingRun({ sessionId, turnId }); + const current = await input.session.getTurn(turnId); + if (!current) { + return { type: 'conflict', message: `Turn not found: ${turnId}` }; + } + if (current.active_executor_id !== configuration.EXECUTOR_ID) { + throw new TurnOwnershipLostError({ + turn_id: turnId, + active_executor_id: configuration.EXECUTOR_ID, + }); + } + if (current.state.status !== 'paused') { + return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; + } const resumed = await resumePausedTurn({ session: input.session, turnId, @@ -1007,6 +1070,11 @@ async function rebuildOptionallyAndRun(input: { deps: input.deps, }); try { + // Cancel aborted the reserved controller and will not freeze. The drain writes + // the terminal state. Skip the inbox insert so this send is not a second writer. + if (resumed.signal.aborted) { + return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; + } return await applyTurnEvents({ handle: resumed.turn, events: input.events, diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 5ab404f4d..9a7ed0498 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -10,7 +10,10 @@ import { Mutex } from 'async-mutex'; const DEFAULT_PAUSED_KEEPALIVE_MS = 300_000; interface ActiveTurnRun { - turn: TurnHandle; + /** Absent until `track()` attaches the stream. A reserved rebuild has no handle yet. */ + turn: TurnHandle | undefined; + /** A rebuild reserved before `track()`. The row is paused, so shutdown offloads it. */ + reservedPaused: boolean; abortController: AbortController; waitUntilCompleted: Promise; markCompleted: () => void; @@ -21,7 +24,13 @@ function activeTurnKey(sessionId: string, turnId: string): string { } function isPausedRun(run: ActiveTurnRun | undefined): run is ActiveTurnRun { - return run?.turn.state.status === 'paused'; + if (!run) { + return false; + } + if (run.turn) { + return run.turn.state.status === 'paused'; + } + return run.reservedPaused; } export class ActiveTurnRegistry { @@ -70,17 +79,31 @@ export class ActiveTurnRegistry { */ 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: ActiveTurnRun; + if (existing) { + // A reserved rebuild adopts this stream. Replacing the run would leave shutdown + // waiting on a promise nothing resolves. + if (existing.turn !== undefined || existing.abortController !== input.abortController) { + throw new Error(`Turn ${input.turn.id} is already registered in this process`); + } + existing.turn = input.turn; + existing.reservedPaused = false; + run = existing; + } else { + const { promise: waitUntilCompleted, resolve } = Promise.withResolvers(); + const markCompleted = (): void => { + resolve(undefined); + }; + run = { + abortController: input.abortController, + waitUntilCompleted, + markCompleted, + turn: input.turn, + reservedPaused: false, + }; + this.runs.set(key, run); + } // 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. @@ -125,12 +148,14 @@ 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 only reserved for a rebuild. + * An aborted run stays registered until its drain finishes. Call `waitForDrainingRun` before + * rebuilding so a second generator does not write with the same executor id. */ getTurnHandle(input: { sessionId: string; turnId: string }): TurnHandle | undefined { const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); - if (!run) { + if (!run?.turn) { return undefined; } if (run.abortController.signal.aborted) { @@ -139,6 +164,56 @@ export class ActiveTurnRegistry { return run.turn; } + /** + * Wait until an aborted run has left the registry. + * The caller holds the per-turn lock. The drain does not, so this does not deadlock. + * A live run is left alone. + */ + async waitForDrainingRun(input: { sessionId: string; turnId: string }): Promise { + const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); + if (!run?.abortController.signal.aborted) { + return; + } + await run.waitUntilCompleted; + } + + /** + * Register a paused rebuild before its stream exists. + * Cancel and shutdown can see the controller. `track()` attaches the stream to this same run. + * 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 { promise: waitUntilCompleted, resolve } = Promise.withResolvers(); + const markCompleted = (): void => { + resolve(undefined); + }; + this.runs.set(key, { + turn: undefined, + reservedPaused: true, + abortController: input.abortController, + waitUntilCompleted, + markCompleted, + }); + if (this.alreadyShutDownAbortReason !== undefined && !input.abortController.signal.aborted) { + input.abortController.abort(TurnOffloadReason); + } + } + + /** Drop a reserved rebuild 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 || run.turn !== undefined) { + 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. @@ -191,7 +266,8 @@ export class ActiveTurnRegistry { /** * Enter shutdown mode and wait until the registry is empty. - * Paused runs are offloaded so their approval state survives. Running runs abort as `abortReason`. + * Does not take the per-turn lock. Paused runs, including a reserved rebuild, are offloaded + * directly so their approval state survives. Running runs abort as `abortReason`. * A later `track()` offloads a paused turn and aborts any other turn as `abortReason`. */ async shutdownAndWait(abortReason: CancellationReason): Promise { From 902799ad8ab002b6ee20f7e23f9e8b4898bb161a Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 21:31:29 +0530 Subject: [PATCH 20/32] Address more race condition with heartbeat gone --- .../trueforge-core/src/request-reply/index.ts | 1 + packages/trueforge/src/apis/turns.ts | 22 +++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/packages/trueforge-core/src/request-reply/index.ts b/packages/trueforge-core/src/request-reply/index.ts index 3017ea0a7..8c90cc86e 100644 --- a/packages/trueforge-core/src/request-reply/index.ts +++ b/packages/trueforge-core/src/request-reply/index.ts @@ -21,3 +21,4 @@ export { RequestReplyRouter } from './router'; export type { RouteHandler } from './router'; export { jsonReplySchema, jsonValueFrom } from './types'; export type { JSONReply, JSONValue, RequestEnvelope, RequestHandler } from './types'; +export { heartbeatKey } from './utils'; diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 53176705c..c4b537fdc 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -44,6 +44,7 @@ import { withTimeout, } from '@truefoundry/trueforge-core/core'; import { + heartbeatKey, jsonValueFrom, NO_RESPONDER_PUBLISH_DETAIL, NoResponderError, @@ -1206,6 +1207,27 @@ async function attemptSendTurnEvents(input: { if (refetchedTurn.state.status !== 'paused' || refetchedTurn.active_executor_id !== turn.active_executor_id) { return 'retry'; } + // NoResponder covers a missing key before PUBLISH and a key that expires during the + // reply wait. The owner can refresh it before this claim. Steal only while it is still + // missing. A Redis error on this check is not proof the owner is gone. + if (input.deps.redis === undefined) { + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } + let heartbeatMissing: boolean; + try { + heartbeatMissing = (await input.deps.redis.exists(heartbeatKey(refetchedTurn.active_executor_id))) === 0; + } catch (error) { + input.deps.logger.warn('Failed to recheck owner heartbeat before steal', { + sessionId, + turnId: input.turnId, + owner: refetchedTurn.active_executor_id, + ...extractErrorLogFields(error), + }); + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + } + if (!heartbeatMissing) { + return 'retry'; + } const action = resolveOwnershipAction({ status: refetchedTurn.state.status, ownedByThisExecutor: refetchedTurn.active_executor_id === configuration.EXECUTOR_ID, From 6cda749420dfc19a5e450994c3ba950ef0a460ce Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 22:17:25 +0530 Subject: [PATCH 21/32] Add logs --- .../src/agent-session/SessionHandle.ts | 3 + .../src/agent-session/TurnHandle.ts | 5 + .../core/runtime/AgentThreadOrchestrator.ts | 3 + packages/trueforge/src/apis/turns.ts | 119 +++++++++++------- packages/trueforge/src/main.ts | 1 + packages/trueforge/src/runtime/activeTurns.ts | 100 +++++++++++++-- 6 files changed, 175 insertions(+), 56 deletions(-) diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index 63c251e3a..386101836 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -312,6 +312,9 @@ export class SessionHandle< 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, diff --git a/packages/trueforge-core/src/agent-session/TurnHandle.ts b/packages/trueforge-core/src/agent-session/TurnHandle.ts index 47a818be5..e39d9c2ab 100644 --- a/packages/trueforge-core/src/agent-session/TurnHandle.ts +++ b/packages/trueforge-core/src/agent-session/TurnHandle.ts @@ -342,6 +342,11 @@ export class TurnHandle> { 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; } diff --git a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts index 2be4b85f4..150a5d65f 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts @@ -585,6 +585,9 @@ export class AgentThreadOrchestrator { let rootAgentError: AgentThreadExecutionResult['root_agent_error']; let rootFinished = 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); diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 814436bfc..bba304170 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -70,7 +70,7 @@ import type { ISandboxProviderStore } from '../db/sandboxProviderStore'; import type { ISkillStore } from '../db/skillStore'; import type { TurnMetadata } from '../db/turnMetadata'; import type { IWebSearchProviderStore } from '../db/webSearchProviderStore'; -import { currentLogContext, updateLogContext } from '../logger'; +import { currentLogContext, runWithLogContext, updateLogContext } from '../logger'; import { createAndExecuteTurnRoute, downloadSandboxFileRoute, @@ -735,6 +735,12 @@ async function resumePausedTurn(input: { 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, @@ -1208,14 +1214,15 @@ async function attemptSendTurnEvents(input: { heartbeatMissing = (await input.deps.redis.exists(heartbeatKey(refetchedTurn.active_executor_id))) === 0; } catch (error) { input.deps.logger.warn('Failed to recheck owner heartbeat before steal', { - sessionId, - turnId: input.turnId, owner: refetchedTurn.active_executor_id, ...extractErrorLogFields(error), }); return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } if (!heartbeatMissing) { + input.deps.logger.info('Owner heartbeat is present; not stealing paused turn', { + owner: refetchedTurn.active_executor_id, + }); return 'retry'; } const action = resolveOwnershipAction({ @@ -1225,6 +1232,9 @@ async function attemptSendTurnEvents(input: { peerOutcome: forwarded.peerOutcome, }); if (action === TurnOwnershipAction.Steal) { + input.deps.logger.info('Stealing paused turn', { + fromExecutor: refetchedTurn.active_executor_id, + }); const won = await input.sessionStore.claimTurnExecutor({ session_id: sessionId, turn_id: input.turnId, @@ -1232,8 +1242,14 @@ async function attemptSendTurnEvents(input: { new_executor_id: configuration.EXECUTOR_ID, }); if (!won) { + input.deps.logger.info('Paused turn steal lost the claim', { + fromExecutor: refetchedTurn.active_executor_id, + }); return 'retry'; } + input.deps.logger.info('Paused turn steal claimed the row', { + executorId: configuration.EXECUTOR_ID, + }); try { return await rebuildOptionallyAndRun({ session: input.session, @@ -1317,55 +1333,64 @@ function handlePeerTurnEvents(deps: TurnsRouterDeps): RequestReplyRouteHandler { if (!parsed.success) { return { status: 400, body: { message: 'Invalid turns/send-events payload' } }; } - const session = await deps.sessions.get({ - tenant_id: parsed.data.tenant_id, - session_id: parsed.data.session_id, - }); - if (!session) { - return { status: 404, body: { message: `Session not found: ${parsed.data.session_id}` } }; - } - const result = await deps.activeTurns.withTurnLock( - { sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }, - async () => { - const requestContext = requestContextFromWire(parsed.data.request_context); - const runAsAgent = await referencedRunAsAgent({ - session, - agentStore: deps.resolveAgentStore(requestContext), - }); - const executionDeps = await executionDepsForCaller({ - deps, - session, - requestContext, - requestHeaders: parsed.data.request_headers, - ...(runAsAgent !== undefined ? { runAsAgent } : {}), - }); - return acceptOwnedTurnEvents({ - session, - sessionStore: deps.sessionStore, - turnId: parsed.data.turn_id, - events: parsed.data.events, - requestContext, - requestHeaders: parsed.data.request_headers, - deps: executionDeps, - }); + // 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), ); - if (result.type === 'applied') { - return { status: 200, body: jsonValueFrom({ events: result.events }) }; - } - if (result.type === 'invalid') { - return { status: 400, body: { message: result.message } }; - } - if (result.type === 'not_owner') { - return { status: 412, body: { message: 'Turn is not running on this executor' } }; - } - if (result.type === 'conflict') { - return { status: 409, body: { message: result.message } }; - } - return { status: 503, body: { message: result.message } }; }; } +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) { + return { status: 404, body: { message: `Session not found: ${body.session_id}` } }; + } + const result = await deps.activeTurns.withTurnLock({ sessionId: body.session_id, turnId: body.turn_id }, async () => { + 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 } : {}), + }); + return acceptOwnedTurnEvents({ + session, + sessionStore: deps.sessionStore, + turnId: body.turn_id, + events: body.events, + requestContext, + requestHeaders: body.request_headers, + deps: executionDeps, + }); + }); + if (result.type === 'applied') { + return { status: 200, body: jsonValueFrom({ events: result.events }) }; + } + if (result.type === 'invalid') { + return { status: 400, body: { message: result.message } }; + } + if (result.type === 'not_owner') { + return { status: 412, body: { message: 'Turn is not running on this executor' } }; + } + if (result.type === 'conflict') { + return { status: 409, body: { message: result.message } }; + } + return { status: 503, body: { message: result.message } }; +} + /** Mapped client error from turn execution; undefined means rethrow. */ export interface TurnExecutionError { status: 400 | 404 | 422; diff --git a/packages/trueforge/src/main.ts b/packages/trueforge/src/main.ts index 08df02f40..f92fb95bd 100644 --- a/packages/trueforge/src/main.ts +++ b/packages/trueforge/src/main.ts @@ -485,6 +485,7 @@ async function createServerRuntime(persistence: ServerPersistence< const activeTurns = new ActiveTurnRegistry({ pausedKeepaliveMs: configuration.TURN_PAUSED_KEEPALIVE_SECONDS * 1000, + logger, }); const requestReplyRouter = new RequestReplyRouter(); const eventSubscriptions = new EventSubscriptionRegistry(redis?.client); diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 9a7ed0498..69161b1d8 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -6,10 +6,16 @@ */ import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; import { Mutex } from 'async-mutex'; +import type { Logger } from 'winston'; +import { currentLogContext, runWithLogContext, type LogContext } from '../logger'; const DEFAULT_PAUSED_KEEPALIVE_MS = 300_000; interface ActiveTurnRun { + sessionId: string; + turnId: string; + /** Ids active when the run was registered. Timer and shutdown logs re-enter this. */ + logContext: LogContext | undefined; /** Absent until `track()` attaches the stream. A reserved rebuild has no handle yet. */ turn: TurnHandle | undefined; /** A rebuild reserved before `track()`. The row is paused, so shutdown offloads it. */ @@ -23,6 +29,14 @@ function activeTurnKey(sessionId: string, turnId: string): string { return `${sessionId}:${turnId}`; } +function snapshotLogContext(): LogContext | undefined { + const current = currentLogContext(); + if (current === undefined) { + return undefined; + } + return { ...current }; +} + function isPausedRun(run: ActiveTurnRun | undefined): run is ActiveTurnRun { if (!run) { return false; @@ -42,10 +56,12 @@ export class ActiveTurnRegistry { private readonly keepaliveTimers = new Map(); // the time to keep the turn alive for, in milliseconds private readonly pausedKeepaliveMs: number; + private readonly logger: Logger | undefined; private alreadyShutDownAbortReason: CancellationReason | undefined; - constructor(options?: { pausedKeepaliveMs?: number }) { + constructor(options?: { pausedKeepaliveMs?: number; logger?: Logger }) { this.pausedKeepaliveMs = options?.pausedKeepaliveMs ?? DEFAULT_PAUSED_KEEPALIVE_MS; + this.logger = options?.logger?.child({ module: 'ActiveTurnRegistry' }); } /** @@ -96,6 +112,9 @@ export class ActiveTurnRegistry { resolve(undefined); }; run = { + sessionId: input.turn.session_id, + turnId: input.turn.id, + logContext: snapshotLogContext(), abortController: input.abortController, waitUntilCompleted, markCompleted, @@ -108,8 +127,11 @@ export class ActiveTurnRegistry { // 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 reason = isPausedRun(run) ? TurnOffloadReason : this.alreadyShutDownAbortReason; - input.abortController.abort(reason); + const offload = isPausedRun(run); + if (offload) { + this.logOffload(run, 'shutdown'); + } + input.abortController.abort(offload ? TurnOffloadReason : this.alreadyShutDownAbortReason); } const complete = () => { @@ -191,14 +213,19 @@ export class ActiveTurnRegistry { const markCompleted = (): void => { resolve(undefined); }; - this.runs.set(key, { + const run: ActiveTurnRun = { + sessionId: input.sessionId, + turnId: input.turnId, + logContext: snapshotLogContext(), turn: undefined, reservedPaused: true, abortController: input.abortController, waitUntilCompleted, markCompleted, - }); + }; + this.runs.set(key, run); if (this.alreadyShutDownAbortReason !== undefined && !input.abortController.signal.aborted) { + this.logOffload(run, 'shutdown'); input.abortController.abort(TurnOffloadReason); } } @@ -220,8 +247,11 @@ export class ActiveTurnRegistry { * A later send on this owner rebuilds. `0` disables the timer. */ startPausedKeepalive(input: { sessionId: string; turnId: string }): void { - this.clearPausedKeepalive(input); + const replaced = this.discardPausedKeepalive(input); if (this.pausedKeepaliveMs <= 0 || this.alreadyShutDownAbortReason !== undefined) { + if (replaced) { + this.infoForTurn(input, 'Paused keepalive cleared', { pausedKeepaliveMs: this.pausedKeepaliveMs }); + } return; } const key = activeTurnKey(input.sessionId, input.turnId); @@ -234,22 +264,36 @@ export class ActiveTurnRegistry { this.keepaliveTimers.delete(key); const run = this.runs.get(key); if (isPausedRun(run)) { + this.logOffload(run, 'keepalive'); await this.finishRun(run, TurnOffloadReason); + return; } + this.infoForTurn(input, 'Paused keepalive expired; turn is no longer paused'); }); }, this.pausedKeepaliveMs); timer.unref(); this.keepaliveTimers.set(key, timer); + this.infoForTurn(input, replaced ? 'Paused keepalive restarted' : 'Paused keepalive started', { + pausedKeepaliveMs: this.pausedKeepaliveMs, + }); } clearPausedKeepalive(input: { sessionId: string; turnId: string }): void { + if (!this.discardPausedKeepalive(input)) { + return; + } + this.infoForTurn(input, '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; + return false; } clearTimeout(timer); this.keepaliveTimers.delete(key); + return true; } /** @@ -261,6 +305,9 @@ export class ActiveTurnRegistry { if (!run) { return; } + if (!run.abortController.signal.aborted) { + this.logOffload(run, 'not-owner'); + } await this.finishRun(run, TurnOffloadReason); } @@ -272,23 +319,58 @@ export class ActiveTurnRegistry { */ 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) { continue; } - const reason = isPausedRun(run) ? TurnOffloadReason : abortReason; - run.abortController.abort(reason); + const offload = isPausedRun(run); + if (offload) { + this.logOffload(run, 'shutdown'); + } + run.abortController.abort(offload ? TurnOffloadReason : abortReason); } await Promise.allSettled(pending.map(run => run.waitUntilCompleted)); } } + private logOffload(run: ActiveTurnRun, cause: 'keepalive' | 'not-owner' | 'shutdown'): void { + this.infoInContext(run.logContext, 'Offloading paused turn', { cause }); + } + + private infoForTurn( + input: { sessionId: string; turnId: string }, + message: string, + fields?: Record, + ): void { + const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); + this.infoInContext(run?.logContext, message, fields); + } + + private infoInContext(context: LogContext | undefined, message: string, fields?: Record): void { + const write = (): void => { + if (fields === undefined) { + this.logger?.info(message); + return; + } + this.logger?.info(message, fields); + }; + if (context === undefined) { + write(); + return; + } + runWithLogContext(context, write); + } + private async finishRun(run: ActiveTurnRun, reason: CancellationReason | typeof TurnOffloadReason): Promise { if (!run.abortController.signal.aborted) { run.abortController.abort(reason); From b876cc3f58e3e091a8828748b02f0aa41e40f6ce Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 23:17:17 +0530 Subject: [PATCH 22/32] Fix approval policy across turns --- .../src/agent-session/SessionHandle.ts | 4 ++ .../core/runtime/AgentThreadOrchestrator.ts | 8 +++ .../orchestrationApproval.test.ts | 56 ++++++++++++++++++- 3 files changed, 66 insertions(+), 2 deletions(-) diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index 386101836..75046a186 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -468,12 +468,16 @@ export class SessionHandle< tracing, }); + const previousMcpServers = input.previous?.snapshot.mcp_servers; const orchestrator = new AgentThreadOrchestrator({ agentThreads, createDynamicSubAgentThread, tracing, logger: input.resolver.logger, ...(input.pendingTurnEvents !== undefined ? { pendingTurnEvents: input.pendingTurnEvents } : {}), + ...(previousMcpServers !== null && previousMcpServers !== undefined + ? { mcpServers: Object.values(previousMcpServers) } + : {}), }); return { orchestrator, agentThreads }; diff --git a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts index 150a5d65f..4bb027a5e 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts @@ -184,6 +184,11 @@ export interface AgentThreadOrchestratorInput { logger: Logger; /** Inbox rows not yet in the snapshot. `execute` applies them. A rebuild passes these. */ pendingTurnEvents?: TurnUserEvent[] | undefined; + /** + * 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[] | undefined; } export class AgentThreadOrchestrator { @@ -204,6 +209,9 @@ export class AgentThreadOrchestrator { this.tracing = params.tracing; this.logger = params.logger.child({ module: 'AgentThreadOrchestrator' }); this.pendingTurnEvents = params.pendingTurnEvents !== undefined ? [...params.pendingTurnEvents] : []; + for (const server of params.mcpServers ?? []) { + this.mcpServerInitInfoById.set(server.id, server); + } } public wake(): void { diff --git a/packages/trueforge-core/tests/orchestration/orchestrationApproval.test.ts b/packages/trueforge-core/tests/orchestration/orchestrationApproval.test.ts index 05373e1b9..0b850fe69 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); @@ -791,7 +839,10 @@ function makeApprovalHarness(finalReply: string): { return { orchestrator, thread, callTool }; } -function makeApprovalHarnessWithUserToolSet(finalReply: string): { +function makeApprovalHarnessWithUserToolSet( + finalReply: string, + mcpServers?: MCPServerInitInfo[], +): { orchestrator: AgentThreadOrchestrator; thread: AgentThread; callTool: jest.Mock; @@ -835,6 +886,7 @@ function makeApprovalHarnessWithUserToolSet(finalReply: string): { createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent in approval policy test')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), + ...(mcpServers !== undefined ? { mcpServers } : {}), }); return { orchestrator, thread, callTool, toolSet }; } From 0b4bb91f60b5ce92a5b3600ab3fe0fb8a11046b6 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Wed, 21 Oct 2026 23:52:47 +0530 Subject: [PATCH 23/32] Simplify logging and reduce diff in api/sessions.ts --- .../src/useTrueForgeAgentMessages.ts | 3 +- .../src/agent-session/TurnHandle.ts | 13 +- packages/trueforge/scripts/write-openapi.ts | 2 +- packages/trueforge/src/apis/sessions.ts | 276 +++++++++--------- packages/trueforge/src/runtime/activeTurns.ts | 79 +---- .../unit/apis/deletedSessionCrud.test.ts | 2 +- .../unit/apis/sandboxFileDownload.test.ts | 2 +- .../tests/unit/apis/schedules.test.ts | 2 +- .../tests/unit/apis/sessionHttp.test.ts | 2 +- .../tests/unit/apis/sessions.test.ts | 45 +-- .../unit/apis/turnHarnessErrorStatus.test.ts | 2 +- .../trueforge/tests/unit/apis/turns.test.ts | 34 +-- .../tests/unit/runtime/activeTurns.test.ts | 21 +- 13 files changed, 217 insertions(+), 266 deletions(-) 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/trueforge-core/src/agent-session/TurnHandle.ts b/packages/trueforge-core/src/agent-session/TurnHandle.ts index e39d9c2ab..214fa9c69 100644 --- a/packages/trueforge-core/src/agent-session/TurnHandle.ts +++ b/packages/trueforge-core/src/agent-session/TurnHandle.ts @@ -175,14 +175,6 @@ function resolveTerminalState(input: { }; } -/** A live turn is running or paused. Terminal rows are not executed. */ -function executableTurnStatus(state: TurnState): NonTerminalTurnState['status'] { - if (state.status === 'running' || state.status === 'paused') { - return state.status; - } - throw new Error(`Cannot execute a ${state.status} turn`); -} - export class TurnHandle> { private readonly store: ISessionStore; private turn: TurnRecord; @@ -293,7 +285,10 @@ export class TurnHandle> { resolver: ITurnResourceResolver, signal: AbortSignal, ): AsyncGenerator { - const turnStatus = executableTurnStatus(this.turn.state); + 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; try { diff --git a/packages/trueforge/scripts/write-openapi.ts b/packages/trueforge/scripts/write-openapi.ts index 832ae5823..bb6a520bf 100644 --- a/packages/trueforge/scripts/write-openapi.ts +++ b/packages/trueforge/scripts/write-openapi.ts @@ -95,7 +95,7 @@ const app = createServerApp({ 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), diff --git a/packages/trueforge/src/apis/sessions.ts b/packages/trueforge/src/apis/sessions.ts index 463b9e71d..5316a7db2 100644 --- a/packages/trueforge/src/apis/sessions.ts +++ b/packages/trueforge/src/apis/sessions.ts @@ -57,6 +57,17 @@ import { isSessionAgentNameRef, parseListSessionsQuery, type Session } from '../ import { newId } from '../utils/id'; import { agentIfAccessible, canReadAgentBoundResource, canReadSession, resolveManagedAgentIds } from './agentAccess'; +/** Request-reply path a replica serves to cancel a turn it owns. */ +export const SESSIONS_CANCEL_PATH = 'sessions/cancel'; + +/** Wire body of a peer cancel; validated on receipt (it crosses processes via Redis). */ +const CancelPeerBodySchema = z.object({ + session_id: z.string(), + turn_id: z.string(), + reason: z.enum(CancellationReason), +}); +type CancelPeerBody = z.infer; + export function toWireSession(record: SessionRecord): Session { return { id: record.session_id, @@ -100,6 +111,133 @@ export interface SessionsRouterDeps { webSearchProviderStore: IWebSearchProviderStore; } +/** + * 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 { + return request => { + const parsed = CancelPeerBodySchema.safeParse(request.body); + if (!parsed.success) { + return Promise.resolve({ status: 400, body: { message: 'Invalid sessions/cancel payload' } }); + } + // 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, + abortReason: parsed.data.reason, + }); + if (!found) { + return Promise.resolve({ status: 412, body: { message: 'Turn is not running on this executor' } }); + } + return Promise.resolve({ status: 200, body: {} }); + }; +} + +/** A registry to abort in, a session to freeze, durable state to read, and a way to reach peers. */ +export interface CancelTurnDeps { + activeTurns: ActiveTurnRegistry; + session: Pick; + sessionStore: Pick; + redis?: RedisClient | undefined; + logger: Pick; +} + +/** + * Cancels a `running` or `paused` turn wherever it lives: locally or on the + * owning peer over Redis request-reply. Callers state the motive; default is a + * plain client cancel. + * + * 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. + * Later writes from that replica lose to first-terminal-write-wins. + */ +export async function cancelSessionTurn( + deps: CancelTurnDeps, + input: { turnId: string; reason?: CancellationReason }, +): Promise { + const { turnId, reason = CancellationReason.ClientCancelled } = input; + const sessionId = deps.session.session_id; + // Cancel is POST /sessions/{session_id}/cancel — path middleware has session_id only. + updateLogContext({ turn_id: turnId }); + + const turn = await deps.sessionStore.getTurn({ + session_id: sessionId, + turn_id: turnId, + }); + if (!turn || !isNonTerminalTurnState(turn.state)) { + return; + } + + const owner = turn.active_executor_id; + // Without a Redis client there is no peer to ask, so a different owner falls + // through to the local lookup and freezes if the run is gone. + if (owner !== configuration.EXECUTOR_ID && deps.redis) { + try { + const reply = await redisRequest({ + redis: deps.redis, + executorId: owner, + path: SESSIONS_CANCEL_PATH, + request: { + body: { session_id: sessionId, turn_id: turnId, reason }, + }, + options: { + replyTimeoutMs: configuration.REDIS_REQUEST_REPLY_TIMEOUT_MS, + pollIntervalMs: configuration.REDIS_REQUEST_REPLY_POLL_INTERVAL_MS, + }, + }); + if (reply.status === 200) { + return; + } + } catch (error) { + const fields = { + sessionId, + turnId, + owner, + ...extractErrorLogFields(error), + }; + if (error instanceof RequestTimeoutError) { + deps.logger.warn('Timed out waiting for owning executor to cancel; freezing the turn', fields); + } else { + deps.logger.warn('Failed to reach owning executor over Redis; freezing the turn', fields); + } + } + await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); + return; + } + + const aborted = deps.activeTurns.cancelIfRunning({ + sessionId, + turnId, + abortReason: reason, + }); + if (!aborted) { + await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); + } +} + +/** Freeze a non-terminal turn; missing turns are a no-op (already gone). */ +async function freezeTurnIgnoringMissing( + session: Pick, + input: { turnId: string; reason: CancellationReason }, +): Promise { + try { + await session.freezeTurn({ turn_id: input.turnId, reason: input.reason }); + } catch (error) { + if (error instanceof TurnNotFoundError) { + return; + } + throw error; + } +} + const FORBIDDEN_SESSION_ACCESS = 'Only the session creator can access this session'; function isSessionOwner({ @@ -235,144 +373,6 @@ export function createInternalSessionsRouter(deps: InternalSessionsRouterDeps) { return router; } -/** Request-reply path a replica serves to cancel a turn it owns. */ -export const SESSIONS_CANCEL_PATH = 'sessions/cancel'; - -/** Wire body of a peer cancel; validated on receipt (it crosses processes via Redis). */ -const CancelPeerBodySchema = z.object({ - session_id: z.string(), - turn_id: z.string(), - reason: z.enum(CancellationReason), -}); -type CancelPeerBody = z.infer; - -/** - * 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 { - return request => { - const parsed = CancelPeerBodySchema.safeParse(request.body); - if (!parsed.success) { - return Promise.resolve({ status: 400, body: { message: 'Invalid sessions/cancel payload' } }); - } - // 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, - abortReason: parsed.data.reason, - }); - if (!found) { - return Promise.resolve({ status: 412, body: { message: 'Turn is not running on this executor' } }); - } - return Promise.resolve({ status: 200, body: {} }); - }; -} - -/** A registry to abort in, a session to freeze, durable state to read, and a way to reach peers. */ -export interface CancelTurnDeps { - activeTurns: ActiveTurnRegistry; - session: Pick; - sessionStore: Pick; - redis?: RedisClient | undefined; - logger: Pick; -} - -/** Freeze a non-terminal turn; missing turns are a no-op (already gone). */ -async function freezeTurnIgnoringMissing( - session: Pick, - input: { turnId: string; reason: CancellationReason }, -): Promise { - try { - await session.freezeTurn({ turn_id: input.turnId, reason: input.reason }); - } catch (error) { - if (error instanceof TurnNotFoundError) { - return; - } - throw error; - } -} - -/** - * Cancels a `running` or `paused` turn wherever it lives: locally or on the - * owning peer over Redis request-reply. Callers state the motive; default is a - * plain client cancel. - * - * 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. - * Later writes from that replica lose to first-terminal-write-wins. - */ -export async function cancelSessionTurn( - deps: CancelTurnDeps, - input: { turnId: string; reason?: CancellationReason }, -): Promise { - const { turnId, reason = CancellationReason.ClientCancelled } = input; - const sessionId = deps.session.session_id; - // Cancel is POST /sessions/{session_id}/cancel — path middleware has session_id only. - updateLogContext({ turn_id: turnId }); - - const turn = await deps.sessionStore.getTurn({ - session_id: sessionId, - turn_id: turnId, - }); - if (!turn || !isNonTerminalTurnState(turn.state)) { - return; - } - - const owner = turn.active_executor_id; - // Without a Redis client there is no peer to ask, so a different owner falls - // through to the local lookup and freezes if the run is gone. - if (owner !== configuration.EXECUTOR_ID && deps.redis) { - try { - const reply = await redisRequest({ - redis: deps.redis, - executorId: owner, - path: SESSIONS_CANCEL_PATH, - request: { - body: { session_id: sessionId, turn_id: turnId, reason }, - }, - options: { - replyTimeoutMs: configuration.REDIS_REQUEST_REPLY_TIMEOUT_MS, - pollIntervalMs: configuration.REDIS_REQUEST_REPLY_POLL_INTERVAL_MS, - }, - }); - if (reply.status === 200) { - return; - } - } catch (error) { - const fields = { - sessionId, - turnId, - owner, - ...extractErrorLogFields(error), - }; - if (error instanceof RequestTimeoutError) { - deps.logger.warn('Timed out waiting for owning executor to cancel; freezing the turn', fields); - } else { - deps.logger.warn('Failed to reach owning executor over Redis; freezing the turn', fields); - } - } - await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); - return; - } - - const aborted = deps.activeTurns.cancelIfRunning({ - sessionId, - turnId, - abortReason: reason, - }); - if (!aborted) { - await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); - } -} - /** DB-backed sessions (mounted at /api/v1/sessions). */ export function createSessionsRouter(deps: SessionsRouterDeps) { const createSessionHandler: RouteHandler = async c => { diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 69161b1d8..ecafa662e 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -7,15 +7,10 @@ import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; import { Mutex } from 'async-mutex'; import type { Logger } from 'winston'; -import { currentLogContext, runWithLogContext, type LogContext } from '../logger'; const DEFAULT_PAUSED_KEEPALIVE_MS = 300_000; interface ActiveTurnRun { - sessionId: string; - turnId: string; - /** Ids active when the run was registered. Timer and shutdown logs re-enter this. */ - logContext: LogContext | undefined; /** Absent until `track()` attaches the stream. A reserved rebuild has no handle yet. */ turn: TurnHandle | undefined; /** A rebuild reserved before `track()`. The row is paused, so shutdown offloads it. */ @@ -29,14 +24,6 @@ function activeTurnKey(sessionId: string, turnId: string): string { return `${sessionId}:${turnId}`; } -function snapshotLogContext(): LogContext | undefined { - const current = currentLogContext(); - if (current === undefined) { - return undefined; - } - return { ...current }; -} - function isPausedRun(run: ActiveTurnRun | undefined): run is ActiveTurnRun { if (!run) { return false; @@ -56,12 +43,12 @@ export class ActiveTurnRegistry { private readonly keepaliveTimers = new Map(); // the time to keep the turn alive for, in milliseconds private readonly pausedKeepaliveMs: number; - private readonly logger: Logger | undefined; + private readonly logger: Logger; private alreadyShutDownAbortReason: CancellationReason | undefined; - constructor(options?: { pausedKeepaliveMs?: number; logger?: Logger }) { - this.pausedKeepaliveMs = options?.pausedKeepaliveMs ?? DEFAULT_PAUSED_KEEPALIVE_MS; - this.logger = options?.logger?.child({ module: 'ActiveTurnRegistry' }); + constructor(options: { logger: Logger; pausedKeepaliveMs?: number }) { + this.pausedKeepaliveMs = options.pausedKeepaliveMs ?? DEFAULT_PAUSED_KEEPALIVE_MS; + this.logger = options.logger.child({ module: 'ActiveTurnRegistry' }); } /** @@ -112,9 +99,6 @@ export class ActiveTurnRegistry { resolve(undefined); }; run = { - sessionId: input.turn.session_id, - turnId: input.turn.id, - logContext: snapshotLogContext(), abortController: input.abortController, waitUntilCompleted, markCompleted, @@ -129,7 +113,7 @@ export class ActiveTurnRegistry { if (this.alreadyShutDownAbortReason !== undefined && !input.abortController.signal.aborted) { const offload = isPausedRun(run); if (offload) { - this.logOffload(run, 'shutdown'); + this.logger.info('Offloading paused turn', { cause: 'shutdown' }); } input.abortController.abort(offload ? TurnOffloadReason : this.alreadyShutDownAbortReason); } @@ -214,9 +198,6 @@ export class ActiveTurnRegistry { resolve(undefined); }; const run: ActiveTurnRun = { - sessionId: input.sessionId, - turnId: input.turnId, - logContext: snapshotLogContext(), turn: undefined, reservedPaused: true, abortController: input.abortController, @@ -225,7 +206,7 @@ export class ActiveTurnRegistry { }; this.runs.set(key, run); if (this.alreadyShutDownAbortReason !== undefined && !input.abortController.signal.aborted) { - this.logOffload(run, 'shutdown'); + this.logger.info('Offloading paused turn', { cause: 'shutdown' }); input.abortController.abort(TurnOffloadReason); } } @@ -250,7 +231,7 @@ export class ActiveTurnRegistry { const replaced = this.discardPausedKeepalive(input); if (this.pausedKeepaliveMs <= 0 || this.alreadyShutDownAbortReason !== undefined) { if (replaced) { - this.infoForTurn(input, 'Paused keepalive cleared', { pausedKeepaliveMs: this.pausedKeepaliveMs }); + this.logger.info('Paused keepalive cleared', { pausedKeepaliveMs: this.pausedKeepaliveMs }); } return; } @@ -264,16 +245,16 @@ export class ActiveTurnRegistry { this.keepaliveTimers.delete(key); const run = this.runs.get(key); if (isPausedRun(run)) { - this.logOffload(run, 'keepalive'); + this.logger.info('Offloading paused turn', { cause: 'keepalive' }); await this.finishRun(run, TurnOffloadReason); return; } - this.infoForTurn(input, 'Paused keepalive expired; turn is no longer paused'); + this.logger.info('Paused keepalive expired; turn is no longer paused'); }); }, this.pausedKeepaliveMs); timer.unref(); this.keepaliveTimers.set(key, timer); - this.infoForTurn(input, replaced ? 'Paused keepalive restarted' : 'Paused keepalive started', { + this.logger.info(replaced ? 'Paused keepalive restarted' : 'Paused keepalive started', { pausedKeepaliveMs: this.pausedKeepaliveMs, }); } @@ -282,7 +263,7 @@ export class ActiveTurnRegistry { if (!this.discardPausedKeepalive(input)) { return; } - this.infoForTurn(input, 'Paused keepalive cleared'); + this.logger.info('Paused keepalive cleared'); } private discardPausedKeepalive(input: { sessionId: string; turnId: string }): boolean { @@ -297,8 +278,8 @@ export class ActiveTurnRegistry { } /** - * Abort the local run and wait until it leaves the registry. - * Caller must already hold the per-turn lock. Offload does not write a terminal state. + * 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)); @@ -306,7 +287,7 @@ export class ActiveTurnRegistry { return; } if (!run.abortController.signal.aborted) { - this.logOffload(run, 'not-owner'); + this.logger.info('Offloading paused turn', { cause: 'not-owner' }); } await this.finishRun(run, TurnOffloadReason); } @@ -325,7 +306,7 @@ export class ActiveTurnRegistry { } this.keepaliveTimers.clear(); if (timerCount > 0) { - this.logger?.info('Cleared paused keepalive timers for shutdown', { timerCount }); + this.logger.info('Cleared paused keepalive timers for shutdown', { timerCount }); } while (this.runs.size > 0) { const pending = Array.from(this.runs.values()); @@ -335,7 +316,7 @@ export class ActiveTurnRegistry { } const offload = isPausedRun(run); if (offload) { - this.logOffload(run, 'shutdown'); + this.logger.info('Offloading paused turn', { cause: 'shutdown' }); } run.abortController.abort(offload ? TurnOffloadReason : abortReason); } @@ -343,34 +324,6 @@ export class ActiveTurnRegistry { } } - private logOffload(run: ActiveTurnRun, cause: 'keepalive' | 'not-owner' | 'shutdown'): void { - this.infoInContext(run.logContext, 'Offloading paused turn', { cause }); - } - - private infoForTurn( - input: { sessionId: string; turnId: string }, - message: string, - fields?: Record, - ): void { - const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); - this.infoInContext(run?.logContext, message, fields); - } - - private infoInContext(context: LogContext | undefined, message: string, fields?: Record): void { - const write = (): void => { - if (fields === undefined) { - this.logger?.info(message); - return; - } - this.logger?.info(message, fields); - }; - if (context === undefined) { - write(); - return; - } - runWithLogContext(context, write); - } - private async finishRun(run: ActiveTurnRun, reason: CancellationReason | typeof TurnOffloadReason): Promise { if (!run.abortController.signal.aborted) { run.abortController.abort(reason); diff --git a/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts b/packages/trueforge/tests/unit/apis/deletedSessionCrud.test.ts index f9f19d750..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({ diff --git a/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts b/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts index 7ca505cae..cb3cb88e8 100644 --- a/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts +++ b/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts @@ -40,7 +40,7 @@ async function buildApp() { app.route( '/', createTurnsRouter({ - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), logger: createLogger({ silent: true }), diff --git a/packages/trueforge/tests/unit/apis/schedules.test.ts b/packages/trueforge/tests/unit/apis/schedules.test.ts index 3d221d72f..d1f0a4d6e 100644 --- a/packages/trueforge/tests/unit/apis/schedules.test.ts +++ b/packages/trueforge/tests/unit/apis/schedules.test.ts @@ -62,7 +62,7 @@ const scheduleBody = { function stubTurnExecutionDeps(agentStore: SqliteAgentStore, scheduleStore: SqliteScheduleStore) { const sessionStore = new InMemorySessionStore(); return { - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: winston.createLogger({ silent: true }) }), agentStore, eventSubscriptions: new EventSubscriptionRegistry(undefined), logger: winston.createLogger({ silent: true }), diff --git a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts index cbba09b89..45e48b3db 100644 --- a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts +++ b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts @@ -116,7 +116,7 @@ describe('sessions HTTP agent binding', () => { resolveMcpServerStore = jest.fn(() => mcpServerStore); resolveSkillStore = jest.fn(() => skillStore); const deps: SessionsRouterDeps = { - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer: new TrueForgeAuthorizer(), logger: createLogger({ silent: true }), redis: createClient(), diff --git a/packages/trueforge/tests/unit/apis/sessions.test.ts b/packages/trueforge/tests/unit/apis/sessions.test.ts index 2561be07a..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', @@ -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,7 +141,7 @@ 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(); @@ -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 511ed8700..c5b25220e 100644 --- a/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts +++ b/packages/trueforge/tests/unit/apis/turnHarnessErrorStatus.test.ts @@ -69,7 +69,7 @@ async function postTurnRejectingWith(error: AgentHarnessError): Promise { app.route( '/', createTurnsRouter({ - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), logger: createLogger({ silent: true }), @@ -225,7 +225,7 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer: { listAgentAccess: input => Promise.resolve( @@ -325,7 +325,7 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), logger: createLogger({ silent: true }), @@ -394,7 +394,7 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), logger: createLogger({ silent: true }), @@ -500,7 +500,7 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer: new TrueForgeAuthorizer(), eventSubscriptions, logger, @@ -629,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), }), ); @@ -748,7 +746,7 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer: new TrueForgeAuthorizer(), eventSubscriptions: new EventSubscriptionRegistry(undefined), logger, @@ -833,7 +831,7 @@ describe('turns', () => { app.route( '/', createTurnsRouter({ - activeTurns: new ActiveTurnRegistry(), + activeTurns: new ActiveTurnRegistry({ logger: createLogger({ silent: true }) }), authorizer, eventSubscriptions: new EventSubscriptionRegistry(undefined), logger: createLogger({ silent: true }), diff --git a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts index 7b7a9123d..dec1b9890 100644 --- a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts +++ b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts @@ -1,6 +1,9 @@ import { CancellationReason, TurnOffloadReason, type TurnHandle } from '@truefoundry/trueforge-core/agent-session'; +import { createLogger } from 'winston'; import { ActiveTurnRegistry } from '../../../src/runtime/activeTurns'; +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; } @@ -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'), @@ -120,7 +123,7 @@ describe('ActiveTurnRegistry', () => { }); it('cancelIfRunning returns false for unknown ids', () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); expect( registry.cancelIfRunning({ sessionId: 'missing', @@ -131,7 +134,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 +154,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 +182,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(); @@ -198,7 +201,7 @@ describe('ActiveTurnRegistry', () => { }); it('late track after shutdownAndWait offloads a paused turn', async () => { - const registry = new ActiveTurnRegistry(); + const registry = new ActiveTurnRegistry({ logger }); await registry.shutdownAndWait(CancellationReason.Abandoned); const abortController = new AbortController(); From 485d5520c4960a3d80b18827138f3ea928b04741 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 22 Oct 2026 00:37:38 +0530 Subject: [PATCH 24/32] refactor activeTurns.ts for rebuilding and live phase --- packages/trueforge/src/apis/turns.ts | 4 +- packages/trueforge/src/runtime/activeTurns.ts | 98 +++++++++++-------- 2 files changed, 58 insertions(+), 44 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index bba304170..3c19aeb64 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -766,7 +766,7 @@ async function resumePausedTurn(input: { }); } catch (error) { executionTimeout.clear(); - // Cancel found this reservation and did not freeze. If the stream never starts, + // 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 && @@ -1067,7 +1067,7 @@ async function rebuildOptionallyAndRun(input: { deps: input.deps, }); try { - // Cancel aborted the reserved controller and will not freeze. The drain writes + // Cancel aborted the rebuilding run and will not freeze. The drain writes // the terminal state. Skip the inbox insert so this send is not a second writer. if (resumed.signal.aborted) { return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index ecafa662e..2c16d7afd 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -10,16 +10,32 @@ import type { Logger } from 'winston'; const DEFAULT_PAUSED_KEEPALIVE_MS = 300_000; -interface ActiveTurnRun { - /** Absent until `track()` attaches the stream. A reserved rebuild has no handle yet. */ - turn: TurnHandle | undefined; - /** A rebuild reserved before `track()`. The row is paused, so shutdown offloads it. */ - reservedPaused: boolean; +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}`; } @@ -28,10 +44,10 @@ function isPausedRun(run: ActiveTurnRun | undefined): run is ActiveTurnRun { if (!run) { return false; } - if (run.turn) { - return run.turn.state.status === 'paused'; + if (run.phase === 'rebuilding') { + return true; } - return run.reservedPaused; + return run.turn.state.status === 'paused'; } export class ActiveTurnRegistry { @@ -83,27 +99,26 @@ export class ActiveTurnRegistry { track(input: { abortController: AbortController; stream: AsyncIterable; turn: TurnHandle }): AsyncGenerator { const key = activeTurnKey(input.turn.session_id, input.turn.id); const existing = this.runs.get(key); - let run: ActiveTurnRun; + let run: LiveTurnRun; if (existing) { - // A reserved rebuild adopts this stream. Replacing the run would leave shutdown - // waiting on a promise nothing resolves. - if (existing.turn !== undefined || existing.abortController !== input.abortController) { + // 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`); } - existing.turn = input.turn; - existing.reservedPaused = false; - run = existing; - } else { - const { promise: waitUntilCompleted, resolve } = Promise.withResolvers(); - const markCompleted = (): void => { - resolve(undefined); + run = { + phase: 'live', + turn: input.turn, + abortController: existing.abortController, + waitUntilCompleted: existing.waitUntilCompleted, + markCompleted: existing.markCompleted, }; + this.runs.set(key, run); + } else { run = { - abortController: input.abortController, - waitUntilCompleted, - markCompleted, + phase: 'live', turn: input.turn, - reservedPaused: false, + ...setupTurnRun(input.abortController), }; this.runs.set(key, run); } @@ -155,13 +170,13 @@ 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 still only reserved for a rebuild. + * or still in the rebuilding phase. * An aborted run stays registered until its drain finishes. Call `waitForDrainingRun` before * rebuilding so a second generator does not write with the same executor id. */ getTurnHandle(input: { sessionId: string; turnId: string }): TurnHandle | undefined { const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); - if (!run?.turn) { + if (run?.phase !== 'live') { return undefined; } if (run.abortController.signal.aborted) { @@ -184,8 +199,9 @@ export class ActiveTurnRegistry { } /** - * Register a paused rebuild before its stream exists. - * Cancel and shutdown can see the controller. `track()` attaches the stream to this same run. + * 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 { @@ -193,29 +209,24 @@ export class ActiveTurnRegistry { if (this.runs.has(key)) { throw new Error(`Turn ${input.turnId} is already registered in this process`); } - const { promise: waitUntilCompleted, resolve } = Promise.withResolvers(); - const markCompleted = (): void => { - resolve(undefined); - }; - const run: ActiveTurnRun = { - turn: undefined, - reservedPaused: true, - abortController: input.abortController, - waitUntilCompleted, - markCompleted, + 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 reserved rebuild whose stream never started, and unblock shutdown waiting on it. */ + /** 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 || run.turn !== undefined) { + if (run?.phase !== 'rebuilding') { return; } this.runs.delete(key); @@ -294,9 +305,12 @@ export class ActiveTurnRegistry { /** * Enter shutdown mode and wait until the registry is empty. - * Does not take the per-turn lock. Paused runs, including a reserved rebuild, are offloaded - * directly so their approval state survives. Running runs abort as `abortReason`. - * A later `track()` offloads a paused turn and aborts any other turn as `abortReason`. + * Does not take the per-turn lock. + * A send holds it across rebuild, `waitForDrainingRun`, 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; From b74dbb69ddfed22090e2c6e4ab5e65e7a0defeba Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 22 Oct 2026 00:50:26 +0530 Subject: [PATCH 25/32] Add comments to activeTurn.ts --- packages/trueforge/src/apis/turns.ts | 2 +- packages/trueforge/src/runtime/activeTurns.ts | 11 +++++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 3c19aeb64..7b587e30e 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -1045,7 +1045,7 @@ async function rebuildOptionallyAndRun(input: { case TurnOwnershipAction.Rebuild: { // An aborted run still occupies the registry until its drain finishes. // getTurnHandle hides it. Wait, then read the row the drain wrote. - await input.deps.activeTurns.waitForDrainingRun({ sessionId, turnId }); + await input.deps.activeTurns.waitForRunToDrainIfAborted({ sessionId, turnId }); const current = await input.session.getTurn(turnId); if (!current) { return { type: 'conflict', message: `Turn not found: ${turnId}` }; diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 2c16d7afd..859d2a70c 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -179,6 +179,8 @@ export class ActiveTurnRegistry { 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; } @@ -187,10 +189,15 @@ export class ActiveTurnRegistry { /** * Wait until an aborted run has left the registry. - * The caller holds the per-turn lock. The drain does not, so this does not deadlock. * A live run is left alone. + * + * Consider that turn was paused and then cancelled before the send event was processed, + * Cancellation just sets the abort controller and does not wait for drain. It never takes the turn lock either, + * When we start processing send events, getTurnHandle will return undefined + * The caller will see undefined and decide to rebuild the turn because it will see it as paused in store + * This is why we need to wait for the turn to drain and re-read the turn from store */ - async waitForDrainingRun(input: { sessionId: string; turnId: string }): Promise { + async waitForRunToDrainIfAborted(input: { sessionId: string; turnId: string }): Promise { const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); if (!run?.abortController.signal.aborted) { return; From 853dcd5385d55249ee3e07e203d5fe563298823d Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 22 Oct 2026 02:50:23 +0530 Subject: [PATCH 26/32] Simplify code for stealing decision --- packages/trueforge/src/apis/turns.ts | 186 +++++++++++++-------------- 1 file changed, 86 insertions(+), 100 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 7b587e30e..ecf689c2e 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -883,11 +883,11 @@ const PeerMessageSchema = z.object({ }); type SendEventsResult = - | { type: 'applied'; events: TurnUserEvent[] } - | { type: 'invalid'; message: string } - | { type: 'conflict'; message: string } - | { type: 'unavailable'; message: string } - | { type: 'not_owner' }; + | { 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: 'unavailable'; message: string } // the turn owner is unavailable + | { type: 'not_owner' }; // the turn is not owned by this executor type ForwardedSendResult = SendEventsResult | { type: 'peer'; peerOutcome: PeerOutcome }; @@ -962,10 +962,14 @@ async function forwardTurnEvents(input: { if (reply.status === 200) { const parsed = PeerEventsReplySchema.safeParse(reply.body); if (!parsed.success) { + // ideally we should never get here 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 + // We will send back 503, client will retry but will get back 400 because we already applied the approvals return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } return { type: 'applied', events: parsed.data.events }; } + // owner is not gone, but it said it does not have the turn if (reply.status === 412) { return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; } @@ -979,12 +983,17 @@ async function forwardTurnEvents(input: { return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } catch (error) { if (error instanceof NoResponderError) { + // peer is not gone, but network partition happened, publish was not received if (error.detail === NO_RESPONDER_PUBLISH_DETAIL) { return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; } + // heartbeat is gone, consider peer dead return { type: 'peer', peerOutcome: PeerOutcome.NoResponder }; } if (error instanceof RequestTimeoutError) { + // peer is not gone, but request timed out + // this can also cause a bad state for e.g. if peer accepted the batch, applied approvals, but we did not hear back in time + // We will try to peer again but will get back 400 because we already applied the approvals return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; } input.logger.warn('Failed to reach owning executor for turn events', { @@ -1067,8 +1076,10 @@ async function rebuildOptionallyAndRun(input: { deps: input.deps, }); try { - // Cancel aborted the rebuilding run and will not freeze. The drain writes - // the terminal state. Skip the inbox insert so this send is not a second writer. + // reservePausedTurn registered this controller before rebuildTurn returned. + // cancelIfRunning aborts that run and does not freeze. Freeze runs only when + // no run is registered. startDrain still consumes the stream, so skip the + // inbox insert and leave that drain as the only writer. if (resumed.signal.aborted) { return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; } @@ -1100,7 +1111,7 @@ async function rebuildOptionallyAndRun(input: { /** * One ownership pass of a user send. `'retry'` means the row moved and the caller should load it again. - * A failed peer returns unavailable and does not steal. + * `OwnerNotGone` returns unavailable and does not steal. `NoResponder` may steal a paused turn. */ async function attemptSendTurnEvents(input: { session: SessionHandle; @@ -1116,7 +1127,7 @@ async function attemptSendTurnEvents(input: { if (!turn) { return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; } - // done, cancelled, and error cannot accept a send. Retrying them becomes a 503. + // done, cancelled, and error cannot accept a send. if (!isNonTerminalTurnState(turn.state)) { return { type: 'conflict', message: `Turn is not accepting events: ${turn.turn_id}` }; } @@ -1139,7 +1150,10 @@ async function attemptSendTurnEvents(input: { } } + // 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 }); + + // Try to forward the events to the owner const forwarded = await forwardTurnEvents({ redis: input.deps.redis, owner: turn.active_executor_id, @@ -1150,63 +1164,44 @@ async function attemptSendTurnEvents(input: { 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') { 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. return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; } + // Timeout, 412, publish-to-zero, and transport errors do not prove the owner - // is gone. Fail here. Another pass could steal. - if (forwarded.peerOutcome !== PeerOutcome.NoResponder) { - if (refetchedTurn.active_executor_id === configuration.EXECUTOR_ID) { - try { - return await rebuildOptionallyAndRun({ - session: input.session, - turn: refetchedTurn, - events: input.events, - sessionStore: input.sessionStore, - requestContext: input.requestContext, - requestHeaders: input.requestHeaders, - deps: input.deps, - }); - } catch (error) { - if (error instanceof TurnOwnershipLostError) { - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; - } - throw error; - } - } - const failedPeer = resolveOwnershipAction({ - status: refetchedTurn.state.status, - ownedByThisExecutor: false, - hasLocalActiveTurn: false, - peerOutcome: forwarded.peerOutcome, - }); - if (failedPeer === TurnOwnershipAction.Reject) { + // is gone. Returning unavailable stops this request. A later request can steal. + if (forwarded.peerOutcome === PeerOutcome.OwnerNotGone) { + if (refetchedTurn.state.status !== 'running' && refetchedTurn.state.status !== 'paused') { return { type: 'conflict', - message: - refetchedTurn.state.status === 'running' - ? `Turn is running but its executor is gone: ${refetchedTurn.turn_id}` - : `Turn is not accepting events: ${refetchedTurn.turn_id}`, + message: `Turn is not accepting events: ${refetchedTurn.turn_id}`, }; } return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } + // The owner was gone. A row that finished during the peer call is a conflict. // A running row, or a paused row whose owner changed, is loaded again. if (!isNonTerminalTurnState(refetchedTurn.state)) { return { type: 'conflict', message: `Turn is not accepting events: ${refetchedTurn.turn_id}` }; } - if (refetchedTurn.state.status !== 'paused' || refetchedTurn.active_executor_id !== turn.active_executor_id) { + if (refetchedTurn.state.status === 'running' || refetchedTurn.active_executor_id !== turn.active_executor_id) { return 'retry'; } + // NoResponder covers a missing key before PUBLISH and a key that expires during the - // reply wait. The owner can refresh it before this claim. Steal only while it is still - // missing. A Redis error on this check is not proof the owner is gone. + // reply wait. Make sure the heartbeat really is gone before we steal the turn if (input.deps.redis === undefined) { + // unreachable, we know redis is set since we were able to call forwardTurnEvents and get a peer outcome return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } let heartbeatMissing: boolean; @@ -1225,52 +1220,43 @@ async function attemptSendTurnEvents(input: { }); return 'retry'; } - const action = resolveOwnershipAction({ - status: refetchedTurn.state.status, - ownedByThisExecutor: refetchedTurn.active_executor_id === configuration.EXECUTOR_ID, - hasLocalActiveTurn: false, - peerOutcome: forwarded.peerOutcome, + + // Paused, still owned by the same executor for which heartbeat is confirmed missing. + // We can Steal. + input.deps.logger.info('Stealing paused turn', { + fromExecutor: refetchedTurn.active_executor_id, }); - if (action === TurnOwnershipAction.Steal) { - input.deps.logger.info('Stealing paused turn', { + 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('Paused turn steal lost the claim', { 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, + return 'retry'; + } + input.deps.logger.info('Paused turn steal claimed the row', { + executorId: configuration.EXECUTOR_ID, + }); + try { + return await rebuildOptionallyAndRun({ + session: input.session, + turn: refetchedTurn, + events: input.events, + sessionStore: input.sessionStore, + requestContext: input.requestContext, + requestHeaders: input.requestHeaders, + deps: input.deps, }); - if (!won) { - input.deps.logger.info('Paused turn steal lost the claim', { - fromExecutor: refetchedTurn.active_executor_id, - }); + } catch (error) { + if (error instanceof TurnOwnershipLostError) { return 'retry'; } - input.deps.logger.info('Paused turn steal claimed the row', { - executorId: configuration.EXECUTOR_ID, - }); - try { - return await rebuildOptionallyAndRun({ - session: input.session, - turn: refetchedTurn, - events: input.events, - sessionStore: input.sessionStore, - requestContext: input.requestContext, - requestHeaders: input.requestHeaders, - deps: input.deps, - }); - } catch (error) { - if (error instanceof TurnOwnershipLostError) { - return 'retry'; - } - throw error; - } - } - if (action === TurnOwnershipAction.Reject) { - return { type: 'conflict', message: `Turn is not accepting events: ${refetchedTurn.turn_id}` }; + throw error; } - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } /** User send. Forwards to the owner and may steal a paused turn whose heartbeat is gone. */ @@ -1327,24 +1313,6 @@ async function acceptOwnedTurnEvents(input: { } } -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), - ); - }; -} - async function handleParsedPeerTurnEvents(deps: TurnsRouterDeps, body: z.infer) { const session = await deps.sessions.get({ tenant_id: body.tenant_id, @@ -1391,6 +1359,24 @@ async function handleParsedPeerTurnEvents(deps: TurnsRouterDeps, body: z.infer { + 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), + ); + }; +} + /** Mapped client error from turn execution; undefined means rethrow. */ export interface TurnExecutionError { status: 400 | 404 | 422; From 867ea897c385cdd025541c7ed9819626ef4abde4 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Thu, 22 Oct 2026 23:13:15 +0530 Subject: [PATCH 27/32] Simplify code in turn.ts --- .../src/request-reply/client.ts | 9 +- .../src/request-reply/errors.ts | 8 +- .../trueforge-core/src/request-reply/index.ts | 2 +- packages/trueforge/src/apis/turns.ts | 943 +++++++++--------- packages/trueforge/src/routes/turnRoutes.ts | 4 + packages/trueforge/src/runtime/activeTurns.ts | 31 +- .../trueforge/src/runtime/turnOwnership.ts | 65 -- .../tests/unit/runtime/activeTurns.test.ts | 14 + .../tests/unit/runtime/turnOwnership.test.ts | 114 --- 9 files changed, 524 insertions(+), 666 deletions(-) delete mode 100644 packages/trueforge/src/runtime/turnOwnership.ts delete mode 100644 packages/trueforge/tests/unit/runtime/turnOwnership.test.ts diff --git a/packages/trueforge-core/src/request-reply/client.ts b/packages/trueforge-core/src/request-reply/client.ts index 187208a50..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 { NO_RESPONDER_PUBLISH_DETAIL, 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'; @@ -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 4fc455c17..181cc3521 100644 --- a/packages/trueforge-core/src/request-reply/errors.ts +++ b/packages/trueforge-core/src/request-reply/errors.ts @@ -1,7 +1,9 @@ -/** Heartbeat key was absent. */ -export const NO_RESPONDER_HEARTBEAT_DETAIL = 'heartbeat not present'; +/** 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 received the published request'; +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 { diff --git a/packages/trueforge-core/src/request-reply/index.ts b/packages/trueforge-core/src/request-reply/index.ts index 8c90cc86e..27c7f70a8 100644 --- a/packages/trueforge-core/src/request-reply/index.ts +++ b/packages/trueforge-core/src/request-reply/index.ts @@ -9,6 +9,7 @@ export { redisRequest } from './client'; export type { SendRequestOptions } from './client'; export { NO_RESPONDER_HEARTBEAT_DETAIL, + NO_RESPONDER_HEARTBEAT_LOST_DETAIL, NO_RESPONDER_PUBLISH_DETAIL, NoResponderError, ReplyError, @@ -21,4 +22,3 @@ export { RequestReplyRouter } from './router'; export type { RouteHandler } from './router'; export { jsonReplySchema, jsonValueFrom } from './types'; export type { JSONReply, JSONValue, RequestEnvelope, RequestHandler } from './types'; -export { heartbeatKey } from './utils'; diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index ecf689c2e..12c8bd1e3 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -44,9 +44,8 @@ import { withTimeout, } from '@truefoundry/trueforge-core/core'; import { - heartbeatKey, jsonValueFrom, - NO_RESPONDER_PUBLISH_DETAIL, + NO_RESPONDER_HEARTBEAT_LOST_DETAIL, NoResponderError, redisRequest, RequestTimeoutError, @@ -90,7 +89,6 @@ import { resolveSandboxEnvironment, resolveSandboxProvider, } from '../runtime/sessionResources'; -import { PeerOutcome, resolveOwnershipAction, TurnOwnershipAction } from '../runtime/turnOwnership'; import { DEFAULT_SANDBOX_ENVIRONMENT_NAME } from '../schemas/sandboxEnvironment'; import { MAX_SESSION_TITLE_LENGTH } from '../schemas/session'; import { assertGatewayMetadataRequestHeaders } from '../truefoundry/gatewayMetadata'; @@ -851,6 +849,71 @@ async function resumePausedTurn(input: { }; } +/** Mapped client error from turn execution; undefined means rethrow. */ +export interface TurnExecutionError { + status: 400 | 404 | 422; + message: string; +} + +/** + * Resolve turn-start failures to HTTP status + message. + * Returns undefined when the caller should rethrow (unexpected / 5xx). + */ +export function getTurnExecutionError(error: unknown): TurnExecutionError | undefined { + if (error instanceof HTTPException) { + if (error.status === 400 || error.status === 404 || error.status === 422) { + return { status: error.status, message: error.message }; + } + return undefined; + } + if (error instanceof SessionStoreNotFoundError) { + return { status: 404, message: error.message }; + } + if (error instanceof AgentHarnessError && !(error instanceof McpConnectionError)) { + switch (error.code) { + case 'invalid_file_input': + return { status: 400, message: error.message }; + case 'invalid_send_input': + case 'agent_sandbox_required': + case 'tool_name_collision': + return { status: 422, message: error.message }; + case 'capability_state_error': + case 'mcp_connection_failed': + return undefined; + } + } + return undefined; +} + +/** + * Resume cursor: `Last-Event-ID` header wins over the body value because the + * header is updated by the SDK on every reconnect to reflect the last delivered + * event, whereas the body is the original caller-supplied cursor and never + * changes between reconnect attempts. + */ +export function resolveAfterSequenceNumber(c: Context, bodyAfterSequenceNumber?: number): number | undefined { + const lastEventId = c.req.header('last-event-id'); + if (lastEventId) { + const sequenceNumber = Number(lastEventId); + if (!Number.isInteger(sequenceNumber) || sequenceNumber < 0) { + throw new HTTPException(400, { message: 'Invalid Last-Event-Id header' }); + } + return sequenceNumber; + } + return bodyAfterSequenceNumber; +} + +/** True when the subject is the session creator (`created_by_subject.subject_id`). */ +function isSessionOwner({ + subject_id, + created_by_subject, +}: { + subject_id: string; + created_by_subject: { subject_id: string }; +}): boolean { + return created_by_subject.subject_id === subject_id; +} + /** Request-reply path the owner serves for inbound turn events. */ const TURNS_SEND_EVENTS_PATH = 'turns/send-events'; @@ -886,11 +949,30 @@ 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: 'unavailable'; message: string } // the turn owner is unavailable + | { type: 'failed'; message: string } // this executor ran the request and failed while applying the events + | { type: 'unavailable'; message: string } // the owning executor could not be reached | { type: 'not_owner' }; // the turn is not owned by this executor +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. 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; @@ -920,6 +1002,129 @@ async function applyTurnEvents(input: { 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 turn without an ActiveTurn lost its executor; it is never resumed + action = TurnOwnershipAction.Reject; + } + + 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 { + return await applyTurnEvents({ + handle, + events: input.events, + sessionStore: input.sessionStore, + sessionId, + turnId, + }); + } 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 })) { + return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; + } + 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) { + return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; + } + return await applyTurnEvents({ + handle: resumed.turn, + events: input.events, + sessionStore: input.sessionStore, + sessionId, + turnId, + }); + } finally { + resumed.startDrain(); + } + } + case TurnOwnershipAction.Reject: + return { + type: 'conflict', + message: + input.turn.state.status === 'running' + ? `Turn is running but its executor is gone: ${input.turn.turn_id}` + : `Turn is not accepting events: ${input.turn.turn_id}`, + }; + } + } catch (error) { + if (error instanceof TurnOwnershipLostError) { + return 'retry'; + } + input.deps.logger.error('Failed to apply events to an owned turn', { + sessionId, + turnId, + ...extractErrorLogFields(error), + }); + return { type: 'failed', message: `Failed to apply events to turn: ${turnId}` }; + } +} + /** Ask the owning replica to apply the events. A peer outcome means this replica must decide again. */ async function forwardTurnEvents(input: { redis: RedisClient | undefined; @@ -964,14 +1169,14 @@ async function forwardTurnEvents(input: { if (!parsed.success) { // ideally we should never get here 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 - // We will send back 503, client will retry but will get back 400 because we already applied the approvals - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + // We will send back 500, client will retry but will get back 400 because we already applied the approvals + return { type: 'failed', message: `Failed to apply events to turn: ${input.turnId}` }; } return { type: 'applied', events: parsed.data.events }; } - // owner is not gone, but it said it does not have the turn + // The peer replied that it does not own the row. The batch was not applied. if (reply.status === 412) { - return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; + return { type: 'peer', peerOutcome: PeerOutcome.NotOwner }; } if (reply.status === 409 || reply.status === 400) { const parsed = PeerMessageSchema.safeParse(reply.body); @@ -980,20 +1185,31 @@ async function forwardTurnEvents(input: { message: parsed.success ? parsed.data.message : `Turn is not accepting events: ${input.turnId}`, }; } - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + if (reply.status === 500) { + const parsed = PeerMessageSchema.safeParse(reply.body); + return { + type: 'failed', + message: parsed.success ? parsed.data.message : `Failed to apply events to turn: ${input.turnId}`, + }; + } + const parsed = PeerMessageSchema.safeParse(reply.body); + return { + type: 'unavailable', + message: parsed.success ? parsed.data.message : `Turn owner is unavailable: ${input.turnId}`, + }; } catch (error) { if (error instanceof NoResponderError) { - // peer is not gone, but network partition happened, publish was not received - if (error.detail === NO_RESPONDER_PUBLISH_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 is gone, consider peer dead + // Heartbeat was already missing, or publish reached nobody. This batch was not delivered. return { type: 'peer', peerOutcome: PeerOutcome.NoResponder }; } if (error instanceof RequestTimeoutError) { - // peer is not gone, but request timed out - // this can also cause a bad state for e.g. if peer accepted the batch, applied approvals, but we did not hear back in time - // We will try to peer again but will get back 400 because we already applied the approvals + // 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('Failed to reach owning executor for turn events', { @@ -1005,314 +1221,6 @@ async function forwardTurnEvents(input: { } } -/** - * Owner path. `resolveOwnershipAction` is called as this executor with no peer - * outcome, so the result is run, rebuild, or reject. - */ -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; - const handle = input.deps.activeTurns.getTurnHandle({ sessionId, turnId }); - const action = resolveOwnershipAction({ - status: input.turn.state.status, - ownedByThisExecutor: true, - hasLocalActiveTurn: handle !== undefined, - peerOutcome: null, - }); - - 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 { - return await applyTurnEvents({ - handle, - events: input.events, - sessionStore: input.sessionStore, - sessionId, - turnId, - }); - } finally { - if (handle.state.status === 'paused') { - input.deps.activeTurns.startPausedKeepalive({ sessionId, turnId }); - } - } - } - case TurnOwnershipAction.Rebuild: { - // An aborted run still occupies the registry until its drain finishes. - // getTurnHandle hides it. Wait, then read the row the drain wrote. - await input.deps.activeTurns.waitForRunToDrainIfAborted({ sessionId, turnId }); - const current = await input.session.getTurn(turnId); - if (!current) { - return { type: 'conflict', message: `Turn not found: ${turnId}` }; - } - if (current.active_executor_id !== configuration.EXECUTOR_ID) { - throw new TurnOwnershipLostError({ - turn_id: turnId, - active_executor_id: configuration.EXECUTOR_ID, - }); - } - if (current.state.status !== 'paused') { - return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; - } - const resumed = await resumePausedTurn({ - session: input.session, - turnId, - requestContext: input.requestContext, - requestHeaders: input.requestHeaders, - deps: input.deps, - }); - try { - // reservePausedTurn registered this controller before rebuildTurn returned. - // cancelIfRunning aborts that run and does not freeze. Freeze runs only when - // no run is registered. startDrain still consumes the stream, so skip the - // inbox insert and leave that drain as the only writer. - if (resumed.signal.aborted) { - return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; - } - return await applyTurnEvents({ - handle: resumed.turn, - events: input.events, - sessionStore: input.sessionStore, - sessionId, - turnId, - }); - } finally { - resumed.startDrain(); - } - } - case TurnOwnershipAction.Reject: - return { - type: 'conflict', - message: - input.turn.state.status === 'running' - ? `Turn is running but its executor is gone: ${input.turn.turn_id}` - : `Turn is not accepting events: ${input.turn.turn_id}`, - }; - case TurnOwnershipAction.Forward: - case TurnOwnershipAction.Steal: - case TurnOwnershipAction.Retry: - throw new Error(`Unreachable ownership action for an owned turn: ${action}`); - } -} - -/** - * One ownership pass of a user send. `'retry'` means the row moved and the caller should load it again. - * `OwnerNotGone` returns unavailable and does not steal. `NoResponder` may steal a paused turn. - */ -async function attemptSendTurnEvents(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) { - return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; - } - // done, cancelled, and error cannot accept a send. - if (!isNonTerminalTurnState(turn.state)) { - return { type: 'conflict', message: `Turn is not accepting events: ${turn.turn_id}` }; - } - if (turn.active_executor_id === configuration.EXECUTOR_ID) { - try { - return await rebuildOptionallyAndRun({ - session: input.session, - turn, - events: input.events, - sessionStore: input.sessionStore, - requestContext: input.requestContext, - requestHeaders: input.requestHeaders, - deps: input.deps, - }); - } catch (error) { - if (error instanceof TurnOwnershipLostError) { - return 'retry'; - } - throw error; - } - } - - // 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 }); - - // Try to forward the events to the owner - const forwarded = await forwardTurnEvents({ - redis: input.deps.redis, - owner: 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') { - 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. - return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; - } - - // Timeout, 412, publish-to-zero, and transport errors do not prove the owner - // is gone. Returning unavailable stops this request. A later request can steal. - if (forwarded.peerOutcome === PeerOutcome.OwnerNotGone) { - if (refetchedTurn.state.status !== 'running' && refetchedTurn.state.status !== 'paused') { - return { - type: 'conflict', - message: `Turn is not accepting events: ${refetchedTurn.turn_id}`, - }; - } - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; - } - - // The owner was gone. A row that finished during the peer call is a conflict. - // A running row, or a paused row whose owner changed, is loaded again. - if (!isNonTerminalTurnState(refetchedTurn.state)) { - return { type: 'conflict', message: `Turn is not accepting events: ${refetchedTurn.turn_id}` }; - } - if (refetchedTurn.state.status === 'running' || refetchedTurn.active_executor_id !== turn.active_executor_id) { - return 'retry'; - } - - // NoResponder covers a missing key before PUBLISH and a key that expires during the - // reply wait. Make sure the heartbeat really is gone before we steal the turn - if (input.deps.redis === undefined) { - // unreachable, we know redis is set since we were able to call forwardTurnEvents and get a peer outcome - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; - } - let heartbeatMissing: boolean; - try { - heartbeatMissing = (await input.deps.redis.exists(heartbeatKey(refetchedTurn.active_executor_id))) === 0; - } catch (error) { - input.deps.logger.warn('Failed to recheck owner heartbeat before steal', { - owner: refetchedTurn.active_executor_id, - ...extractErrorLogFields(error), - }); - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; - } - if (!heartbeatMissing) { - input.deps.logger.info('Owner heartbeat is present; not stealing paused turn', { - owner: refetchedTurn.active_executor_id, - }); - return 'retry'; - } - - // Paused, still owned by the same executor for which heartbeat is confirmed missing. - // We can Steal. - input.deps.logger.info('Stealing paused turn', { - 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('Paused turn steal lost the claim', { - fromExecutor: refetchedTurn.active_executor_id, - }); - return 'retry'; - } - input.deps.logger.info('Paused turn steal claimed the row', { - executorId: configuration.EXECUTOR_ID, - }); - try { - return await rebuildOptionallyAndRun({ - session: input.session, - turn: refetchedTurn, - events: input.events, - sessionStore: input.sessionStore, - requestContext: input.requestContext, - requestHeaders: input.requestHeaders, - deps: input.deps, - }); - } catch (error) { - if (error instanceof TurnOwnershipLostError) { - return 'retry'; - } - throw error; - } -} - -/** User send. Forwards to the owner and may steal a paused turn whose heartbeat is gone. */ -async function sendTurnEvents(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 attemptSendTurnEvents(input); - if (result !== 'retry') { - return result; - } - } - return { type: 'unavailable', message: `Turn ownership did not settle: ${input.turnId}` }; -} - -/** Peered send. Runs or rebuilds only when this executor already owns the turn. */ -async function acceptOwnedTurnEvents(input: { - session: SessionHandle; - sessionStore: Pick; - turnId: string; - events: TurnUserEvent[]; - requestContext: RequestContext; - requestHeaders: Record | undefined; - deps: BeginTurnExecutionDeps; -}): Promise { - const turn = await input.session.getTurn(input.turnId); - if (!turn) { - return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; - } - if (turn.active_executor_id !== configuration.EXECUTOR_ID) { - return { type: 'not_owner' }; - } - try { - return await rebuildOptionallyAndRun({ - session: input.session, - turn, - events: input.events, - sessionStore: input.sessionStore, - requestContext: input.requestContext, - requestHeaders: input.requestHeaders, - deps: input.deps, - }); - } catch (error) { - if (error instanceof TurnOwnershipLostError) { - return { type: 'not_owner' }; - } - throw error; - } -} - async function handleParsedPeerTurnEvents(deps: TurnsRouterDeps, body: z.infer) { const session = await deps.sessions.get({ tenant_id: body.tenant_id, @@ -1321,29 +1229,44 @@ async function handleParsedPeerTurnEvents(deps: TurnsRouterDeps, body: z.infer { - 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 } : {}), - }); - return acceptOwnedTurnEvents({ - session, - sessionStore: deps.sessionStore, - turnId: body.turn_id, - events: body.events, - requestContext, - requestHeaders: body.request_headers, - deps: executionDeps, - }); - }); + 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) { + return { type: 'conflict', message: `Turn not found: ${body.turn_id}` }; + } + if (turn.active_executor_id !== configuration.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') { + return { type: 'not_owner' }; + } + return owned; + }, + ); if (result.type === 'applied') { return { status: 200, body: jsonValueFrom({ events: result.events }) }; } @@ -1356,6 +1279,9 @@ async function handleParsedPeerTurnEvents(deps: TurnsRouterDeps, body: z.infer; + 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) { + return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; + } + // done, cancelled, and error cannot accept a send. + if (!isNonTerminalTurnState(turn.state)) { + 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 }); + + // Try to forward the events to the owner + const forwarded = await forwardTurnEvents({ + redis: input.deps.redis, + owner: 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') { + return forwarded; + } -/** - * Resolve turn-start failures to HTTP status + message. - * Returns undefined when the caller should rethrow (unexpected / 5xx). - */ -export function getTurnExecutionError(error: unknown): TurnExecutionError | undefined { - if (error instanceof HTTPException) { - if (error.status === 400 || error.status === 404 || error.status === 422) { - return { status: error.status, message: error.message }; + const refetchedTurn = await input.session.getTurn(input.turnId); + if (refetchedTurn === undefined) { + // Unreachable. Turns are not deleted, and the load above found this row. + return { type: 'conflict', message: `Turn not found: ${input.turnId}` }; + } + + // 412: the peer did not apply the batch. Load the row again and send to the current owner. + // Timeout, 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 (refetchedTurn.state.status !== 'running' && refetchedTurn.state.status !== 'paused') { + return { + type: 'conflict', + message: `Turn is not accepting events: ${refetchedTurn.turn_id}`, + }; } - return undefined; + if (forwarded.peerOutcome === PeerOutcome.NotOwner) { + return 'retry'; + } + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } - if (error instanceof SessionStoreNotFoundError) { - return { status: 404, message: error.message }; + + // This batch was not delivered. A finished row is a conflict. A running row or a new + // owner is loaded again and sent to whoever owns it now. + if (!isNonTerminalTurnState(refetchedTurn.state)) { + return { type: 'conflict', message: `Turn is not accepting events: ${refetchedTurn.turn_id}` }; } - if (error instanceof AgentHarnessError && !(error instanceof McpConnectionError)) { - switch (error.code) { - case 'invalid_file_input': - return { status: 400, message: error.message }; - case 'invalid_send_input': - case 'agent_sandbox_required': - case 'tool_name_collision': - return { status: 422, message: error.message }; - case 'capability_state_error': - case 'mcp_connection_failed': - return undefined; - } + if (refetchedTurn.state.status === 'running' || refetchedTurn.active_executor_id !== turn.active_executor_id) { + return 'retry'; } - return undefined; -} -/** - * Resume cursor: `Last-Event-ID` header wins over the body value because the - * header is updated by the SDK on every reconnect to reflect the last delivered - * event, whereas the body is the original caller-supplied cursor and never - * changes between reconnect attempts. - */ -export function resolveAfterSequenceNumber(c: Context, bodyAfterSequenceNumber?: number): number | undefined { - const lastEventId = c.req.header('last-event-id'); - if (lastEventId) { - const sequenceNumber = Number(lastEventId); - if (!Number.isInteger(sequenceNumber) || sequenceNumber < 0) { - throw new HTTPException(400, { message: 'Invalid Last-Event-Id header' }); - } - return sequenceNumber; + // Paused, still owned by the executor that did not receive this batch. + input.deps.logger.info('Stealing paused turn', { + 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('Paused turn steal lost the claim', { + fromExecutor: refetchedTurn.active_executor_id, + }); + return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; } - return bodyAfterSequenceNumber; + input.deps.logger.info('Paused turn 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, + }); } -/** True when the subject is the session creator (`created_by_subject.subject_id`). */ -function isSessionOwner({ - subject_id, - created_by_subject, -}: { - subject_id: string; - created_by_subject: { subject_id: string }; -}): boolean { - return created_by_subject.subject_id === subject_id; +/** 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; + } + } + return { type: 'unavailable', message: `Turn ownership did not settle: ${input.turnId}` }; } /** DB-backed turns (mounted at /api/v1/sessions). */ 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, + }); + if (result.type === 'applied') { + return c.json({ data: result.events }, 201); + } + if (result.type === 'invalid') { + return c.json({ error: { message: result.message } }, 400); + } + if (result.type === 'conflict') { + return c.json({ error: { message: result.message } }, 409); + } + if (result.type === 'failed') { + return c.json({ error: { message: result.message } }, 500); + } + const message = result.type === 'not_owner' ? `Turn owner is unavailable: ${turnId}` : result.message; + return c.json({ error: { message } }, 503); + }); + }; + const listTurnsHandler: RouteHandler = async c => { const { session_id: sessionId } = c.req.valid('param'); const query = c.req.valid('query'); @@ -1817,73 +1893,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); - } - - 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 sendTurnEvents({ - session, - sessionStore: deps.sessionStore, - turnId, - events, - requestContext, - requestHeaders, - deps: executionDeps, - }); - if (result.type === 'applied') { - return c.json({ data: result.events }, 201); - } - if (result.type === 'invalid') { - return c.json({ error: { message: result.message } }, 400); - } - if (result.type === 'conflict') { - return c.json({ error: { message: result.message } }, 409); - } - const message = result.type === 'not_owner' ? `Turn owner is unavailable: ${turnId}` : result.message; - return c.json({ error: { message } }, 503); - }); - }; - const router = new OpenAPIHono(); router.openapi(createAndExecuteTurnRoute, createAndExecuteTurnHandler); router.openapi(listTurnsRoute, listTurnsHandler); diff --git a/packages/trueforge/src/routes/turnRoutes.ts b/packages/trueforge/src/routes/turnRoutes.ts index e1b7612a3..2637d9062 100644 --- a/packages/trueforge/src/routes/turnRoutes.ts +++ b/packages/trueforge/src/routes/turnRoutes.ts @@ -209,6 +209,10 @@ export const postTurnEventsRoute = createRoute({ content: { 'application/json': { schema: RequestErrorResponseSchema } }, description: 'Turn is terminal, or an event id already exists.', }, + 500: { + content: { 'application/json': { schema: RequestErrorResponseSchema } }, + description: 'The executor ran the request and failed while applying the events.', + }, 503: { content: { 'application/json': { schema: RequestErrorResponseSchema } }, description: 'Turn owner is unavailable. Retry.', diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 859d2a70c..d8ff64680 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -171,8 +171,8 @@ 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 still in the rebuilding phase. - * An aborted run stays registered until its drain finishes. Call `waitForDrainingRun` before - * rebuilding so a second generator does not write with the same executor id. + * 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)); @@ -188,21 +188,24 @@ export class ActiveTurnRegistry { } /** - * Wait until an aborted run has left the registry. - * A live run is left alone. - * - * Consider that turn was paused and then cancelled before the send event was processed, - * Cancellation just sets the abort controller and does not wait for drain. It never takes the turn lock either, - * When we start processing send events, getTurnHandle will return undefined - * The caller will see undefined and decide to rebuild the turn because it will see it as paused in store - * This is why we need to wait for the turn to drain and re-read the turn from store + * 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. */ - async waitForRunToDrainIfAborted(input: { sessionId: string; turnId: string }): Promise { + isCancelled(input: { sessionId: string; turnId: string }): boolean { const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); if (!run?.abortController.signal.aborted) { - return; + return false; } - await run.waitUntilCompleted; + const { signal } = run.abortController; + return ( + signal.reason === CancellationReason.ServerExecutionTimeout || + signal.reason === CancellationReason.CancelledForNextTurn || + signal.reason === CancellationReason.Abandoned || + signal.reason === CancellationReason.ClientCancelled + ); } /** @@ -313,7 +316,7 @@ export class ActiveTurnRegistry { /** * Enter shutdown mode and wait until the registry is empty. * Does not take the per-turn lock. - * A send holds it across rebuild, `waitForDrainingRun`, and `dropLocalTurn`; + * 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. diff --git a/packages/trueforge/src/runtime/turnOwnership.ts b/packages/trueforge/src/runtime/turnOwnership.ts deleted file mode 100644 index 07d148c3f..000000000 --- a/packages/trueforge/src/runtime/turnOwnership.ts +++ /dev/null @@ -1,65 +0,0 @@ -import type { TurnState } from '@truefoundry/trueforge-core/agent-session'; - -/** - * Result of a forward that did not itself finish the request. - * A peer 200, 400, or 409 is returned directly and never becomes an outcome. - */ -export enum PeerOutcome { - /** Heartbeat missing: the owner is gone. */ - NoResponder = 'no_responder', - /** Publish-zero, timeout, 412, or transport failure. The owner is not known to be gone. */ - OwnerNotGone = 'owner_not_gone', -} - -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', - /** 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.NoResponder: - return status === 'paused' ? TurnOwnershipAction.Steal : TurnOwnershipAction.Retry; - case PeerOutcome.OwnerNotGone: - return TurnOwnershipAction.Retry; - } -} diff --git a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts index dec1b9890..c7026627e 100644 --- a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts +++ b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts @@ -120,6 +120,20 @@ 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', () => { 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 8663d02bd..000000000 --- a/packages/trueforge/tests/unit/runtime/turnOwnership.test.ts +++ /dev/null @@ -1,114 +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: '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 + forward that does not show a missing heartbeat does not steal', - input: { - status: 'paused', - ownedByThisExecutor: false, - hasLocalActiveTurn: false, - peerOutcome: PeerOutcome.OwnerNotGone, - }, - 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); - } - } - } - } - }); -}); From f9bf11fdd170cac26b40c6bb546abbfb52ab1e25 Mon Sep 17 00:00:00 2001 From: "trueforge-dev-bot[bot]" Date: Thu, 22 Oct 2026 17:46:00 +0000 Subject: [PATCH 28/32] Regenerate OpenAPI document and SDKs --- .github/fern/openapi/openapi.json | 10 +++++++ docs/openapi.json | 10 +++++++ .../api/resources/sessions/client/Client.ts | 12 ++++++++ .../trueforge-sdk/tests/wire/sessions.test.ts | 29 +++++++++++++++++++ .../src/trueforge_sdk/sessions/raw_client.py | 23 +++++++++++++++ 5 files changed, 84 insertions(+) diff --git a/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index d76dff7a3..9ebdc7e3e 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -10057,6 +10057,16 @@ }, "description": "Turn is terminal, 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." + }, "503": { "content": { "application/json": { diff --git a/docs/openapi.json b/docs/openapi.json index d76dff7a3..9ebdc7e3e 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -10057,6 +10057,16 @@ }, "description": "Turn is terminal, 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." + }, "503": { "content": { "application/json": { 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 5894fe5d8..776cd91ba 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 TrueForge.ServiceUnavailableError} * @throws {@link errors.TrueForgeError} * @throws {@link errors.TrueForgeTimeoutError} @@ -1873,6 +1874,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, + ); case 503: throw new TrueForge.ServiceUnavailableError( serializers.RequestErrorResponse.parseOrThrow(_response.error.body, { diff --git a/packages/trueforge-sdk/tests/wire/sessions.test.ts b/packages/trueforge-sdk/tests/wire/sessions.test.ts index 173a9feee..3b92f8aeb 100644 --- a/packages/trueforge-sdk/tests/wire/sessions.test.ts +++ b/packages/trueforge-sdk/tests/wire/sessions.test.ts @@ -1561,6 +1561,35 @@ describe("SessionsClient", () => { 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("create_turn_events (7)", 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") 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 8a0252b1f..674860eda 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.service_unavailable_error import ServiceUnavailableError @@ -1543,6 +1544,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(), + ), + ), + ) if _response.status_code == 503: raise ServiceUnavailableError( headers=dict(_response.headers), @@ -3201,6 +3213,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(), + ), + ), + ) if _response.status_code == 503: raise ServiceUnavailableError( headers=dict(_response.headers), From 001d09f5d9e0898a0e45c7618d09e369b32a4832 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Fri, 23 Oct 2026 02:21:54 +0530 Subject: [PATCH 29/32] update error messages --- packages/trueforge/src/apis/turns.ts | 290 ++++++++++++++------------- 1 file changed, 148 insertions(+), 142 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 12c8bd1e3..a19604b34 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -949,9 +949,7 @@ 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 - | { type: 'unavailable'; message: string } // the owning executor could not be reached - | { type: 'not_owner' }; // the turn is not owned by this executor + | { 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. */ @@ -1033,7 +1031,9 @@ async function rebuildOptionallyAndRun(input: { // a paused turn, but we don't have a local active turn action = TurnOwnershipAction.Rebuild; } else { - // a running turn without an ActiveTurn lost its executor; it is never resumed + // a running state turn, either + // - cancelled but hasn't reached turn.done + // - stuck in running state because the executor died suddenly action = TurnOwnershipAction.Reject; } @@ -1066,7 +1066,7 @@ async function rebuildOptionallyAndRun(input: { // 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 })) { - return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; + return { type: 'conflict', message: 'Turn is cancelled, cannot apply events' }; } const resumed = await resumePausedTurn({ session: input.session, @@ -1090,7 +1090,7 @@ async function rebuildOptionallyAndRun(input: { // apply? -> startDrain (without await) -> return conflict try { if (resumed.signal.aborted) { - return { type: 'conflict', message: `Turn is not accepting events: ${turnId}` }; + return { type: 'conflict', message: 'Turn is cancelled or unable to accept events' }; } return await applyTurnEvents({ handle: resumed.turn, @@ -1108,8 +1108,8 @@ async function rebuildOptionallyAndRun(input: { type: 'conflict', message: input.turn.state.status === 'running' - ? `Turn is running but its executor is gone: ${input.turn.turn_id}` - : `Turn is not accepting events: ${input.turn.turn_id}`, + ? 'Turn is running but its executor is gone' + : 'Turn is cancelled or unable to accept events', }; } } catch (error) { @@ -1121,7 +1121,8 @@ async function rebuildOptionallyAndRun(input: { turnId, ...extractErrorLogFields(error), }); - return { type: 'failed', message: `Failed to apply events to turn: ${turnId}` }; + const cause = error instanceof Error ? error.message : String(error); + return { type: 'failed', message: `Failed to apply events: ${cause}` }; } } @@ -1134,10 +1135,12 @@ async function forwardTurnEvents(input: { events: TurnUserEvent[]; requestContext: RequestContext; requestHeaders?: Record | undefined; - logger: Pick; + logger: Logger; }): Promise { if (!input.redis) { - return { type: 'unavailable', message: `Turn owner is unavailable: ${input.turnId}` }; + // 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({ @@ -1164,39 +1167,45 @@ async function forwardTurnEvents(input: { pollIntervalMs: configuration.REDIS_REQUEST_REPLY_POLL_INTERVAL_MS, }, }); - if (reply.status === 200) { - const parsed = PeerEventsReplySchema.safeParse(reply.body); - if (!parsed.success) { - // ideally we should never get here 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 - // We will send back 500, client will retry but will get back 400 because we already applied the approvals - return { type: 'failed', message: `Failed to apply events to turn: ${input.turnId}` }; + + 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 + 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`, + }; } - return { type: 'applied', events: parsed.data.events }; - } - // The peer replied that it does not own the row. The batch was not applied. - if (reply.status === 412) { - return { type: 'peer', peerOutcome: PeerOutcome.NotOwner }; - } - if (reply.status === 409 || reply.status === 400) { - const parsed = PeerMessageSchema.safeParse(reply.body); - return { - type: reply.status === 400 ? 'invalid' : 'conflict', - message: parsed.success ? parsed.data.message : `Turn is not accepting events: ${input.turnId}`, - }; - } - if (reply.status === 500) { - const parsed = PeerMessageSchema.safeParse(reply.body); - return { - type: 'failed', - message: parsed.success ? parsed.data.message : `Failed to apply events to turn: ${input.turnId}`, - }; } - const parsed = PeerMessageSchema.safeParse(reply.body); - return { - type: 'unavailable', - message: parsed.success ? parsed.data.message : `Turn owner is unavailable: ${input.turnId}`, - }; } catch (error) { if (error instanceof NoResponderError) { // The request was published, then the heartbeat disappeared while waiting for the reply. @@ -1206,8 +1215,7 @@ async function forwardTurnEvents(input: { } // Heartbeat was already missing, or publish reached nobody. This batch was not delivered. return { type: 'peer', peerOutcome: PeerOutcome.NoResponder }; - } - if (error instanceof RequestTimeoutError) { + } else if (error instanceof RequestTimeoutError) { // 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 }; @@ -1227,11 +1235,11 @@ async function handleParsedPeerTurnEvents(deps: TurnsRouterDeps, body: z.infer => { + async (): Promise => { const requestContext = requestContextFromWire(body.request_context); const runAsAgent = await referencedRunAsAgent({ session, @@ -1247,7 +1255,7 @@ async function handleParsedPeerTurnEvents(deps: TurnsRouterDeps, body: z.infer = 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, - }); - if (result.type === 'applied') { - return c.json({ data: result.events }, 201); - } - if (result.type === 'invalid') { - return c.json({ error: { message: result.message } }, 400); - } - if (result.type === 'conflict') { - return c.json({ error: { message: result.message } }, 409); - } - if (result.type === 'failed') { - return c.json({ error: { message: result.message } }, 500); - } - const message = result.type === 'not_owner' ? `Turn owner is unavailable: ${turnId}` : result.message; - return c.json({ error: { message } }, 503); - }); - }; - const listTurnsHandler: RouteHandler = async c => { const { session_id: sessionId } = c.req.valid('param'); const query = c.req.valid('query'); @@ -1812,6 +1748,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'); From dcfd21fce049418a93ff90fdc70a5947d92da484 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Fri, 23 Oct 2026 09:24:06 +0530 Subject: [PATCH 30/32] Log each turn-send decision and reject a running turn the executor did not receive. A changed owner is retried. The same running owner returns 409 instead of retrying into a 500, and the route no longer documents 503. Co-authored-by: Cursor --- packages/trueforge/src/apis/turns.ts | 152 ++++++++++++++++---- packages/trueforge/src/routes/turnRoutes.ts | 6 +- 2 files changed, 126 insertions(+), 32 deletions(-) diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index a19604b34..7ec5b1ece 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -794,6 +794,7 @@ async function resumePausedTurn(input: { // 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({ @@ -956,7 +957,7 @@ export enum PeerOutcome { 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. Do not send it again. */ + /** The request may already have been applied but we did not receive the reply. Do not send it again. */ OwnerNotGone = 'owner_not_gone', } @@ -1037,6 +1038,9 @@ async function rebuildOptionallyAndRun(input: { 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) { @@ -1047,13 +1051,15 @@ async function rebuildOptionallyAndRun(input: { // the turn paused rearms; a later `running` update clears that timer from the drain. input.deps.activeTurns.clearPausedKeepalive({ sessionId, turnId }); try { - return await applyTurnEvents({ + 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 }); @@ -1066,6 +1072,12 @@ async function rebuildOptionallyAndRun(input: { // 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({ @@ -1090,20 +1102,35 @@ async function rebuildOptionallyAndRun(input: { // 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' }; } - return await applyTurnEvents({ + 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: + 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: @@ -1111,16 +1138,18 @@ async function rebuildOptionallyAndRun(input: { ? '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', { - sessionId, - turnId, - ...extractErrorLogFields(error), - }); + 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}` }; } @@ -1129,7 +1158,7 @@ async function rebuildOptionallyAndRun(input: { /** Ask the owning replica to apply the events. A peer outcome means this replica must decide again. */ async function forwardTurnEvents(input: { redis: RedisClient | undefined; - owner: string; + activeExecutorId: string; sessionId: string; turnId: string; events: TurnUserEvent[]; @@ -1159,7 +1188,7 @@ async function forwardTurnEvents(input: { try { const reply = await redisRequest({ redis: input.redis, - executorId: input.owner, + executorId: input.activeExecutorId, path: TURNS_SEND_EVENTS_PATH, request: { body: jsonValueFrom(body) }, options: { @@ -1168,12 +1197,14 @@ async function forwardTurnEvents(input: { }, }); + 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 }; @@ -1208,6 +1239,7 @@ async function forwardTurnEvents(input: { } } 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) { @@ -1216,14 +1248,15 @@ async function forwardTurnEvents(input: { // 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('Failed to reach owning executor for turn events', { - sessionId: input.sessionId, - turnId: input.turnId, - owner: input.owner, + input.logger.warn('Turn send peer missed', { + activeExecutorId: input.activeExecutorId, + detail: 'transport', + ...extractErrorLogFields(error), }); return { type: 'peer', peerOutcome: PeerOutcome.OwnerNotGone }; } @@ -1235,6 +1268,7 @@ async function handleParsedPeerTurnEvents(deps: TurnsRouterDeps, body: z.infer Date: Fri, 23 Oct 2026 03:57:04 +0000 Subject: [PATCH 31/32] Regenerate OpenAPI document and SDKs --- .github/fern/openapi/openapi.json | 12 +------- docs/openapi.json | 12 +------- .../src/api/errors/ServiceUnavailableError.ts | 24 --------------- .../trueforge-sdk/src/api/errors/index.ts | 1 - .../api/resources/sessions/client/Client.ts | 12 -------- .../trueforge-sdk/tests/wire/sessions.test.ts | 29 ------------------- .../src/trueforge_sdk/__init__.py | 3 -- .../src/trueforge_sdk/errors/__init__.py | 3 -- .../errors/service_unavailable_error.py | 11 ------- .../src/trueforge_sdk/sessions/raw_client.py | 23 --------------- 10 files changed, 2 insertions(+), 128 deletions(-) delete mode 100644 packages/trueforge-sdk/src/api/errors/ServiceUnavailableError.ts delete mode 100644 python/trueforge_sdk/src/trueforge_sdk/errors/service_unavailable_error.py diff --git a/.github/fern/openapi/openapi.json b/.github/fern/openapi/openapi.json index 9ebdc7e3e..ec9f7375e 100644 --- a/.github/fern/openapi/openapi.json +++ b/.github/fern/openapi/openapi.json @@ -10055,7 +10055,7 @@ } } }, - "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": { @@ -10066,16 +10066,6 @@ } }, "description": "The executor ran the request and failed while applying the events." - }, - "503": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/RequestErrorResponse" - } - } - }, - "description": "Turn owner is unavailable. Retry." } }, "summary": "Create turn events", diff --git a/docs/openapi.json b/docs/openapi.json index 9ebdc7e3e..ec9f7375e 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -10055,7 +10055,7 @@ } } }, - "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": { @@ -10066,16 +10066,6 @@ } }, "description": "The executor ran the request and failed while applying the events." - }, - "503": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/RequestErrorResponse" - } - } - }, - "description": "Turn owner is unavailable. Retry." } }, "summary": "Create turn events", diff --git a/packages/trueforge-sdk/src/api/errors/ServiceUnavailableError.ts b/packages/trueforge-sdk/src/api/errors/ServiceUnavailableError.ts deleted file mode 100644 index 58e0399d1..000000000 --- a/packages/trueforge-sdk/src/api/errors/ServiceUnavailableError.ts +++ /dev/null @@ -1,24 +0,0 @@ -// This file was auto-generated by Fern from our API Definition. - -import type * as core from "../../core/index.js"; -import * as errors from "../../errors/index.js"; -import type * as TrueForge from "../index.js"; - -export class ServiceUnavailableError extends errors.TrueForgeError { - public declare readonly body: TrueForge.RequestErrorResponse; - - constructor(body: TrueForge.RequestErrorResponse, rawResponse?: core.RawResponse) { - super({ - message: "ServiceUnavailableError", - statusCode: 503, - body: body, - rawResponse: rawResponse, - }); - Object.setPrototypeOf(this, new.target.prototype); - if (Error.captureStackTrace) { - Error.captureStackTrace(this, this.constructor); - } - - this.name = "ServiceUnavailableError"; - } -} diff --git a/packages/trueforge-sdk/src/api/errors/index.ts b/packages/trueforge-sdk/src/api/errors/index.ts index e47bb54e0..5388d473c 100644 --- a/packages/trueforge-sdk/src/api/errors/index.ts +++ b/packages/trueforge-sdk/src/api/errors/index.ts @@ -8,6 +8,5 @@ export * from "./GoneError.js"; export * from "./InternalServerError.js"; export * from "./NotFoundError.js"; export * from "./PreconditionFailedError.js"; -export * from "./ServiceUnavailableError.js"; export * from "./UnauthorizedError.js"; export * from "./UnprocessableEntityError.js"; 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 776cd91ba..ba4f60323 100644 --- a/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts +++ b/packages/trueforge-sdk/src/api/resources/sessions/client/Client.ts @@ -1755,7 +1755,6 @@ export class SessionsClient { * @throws {@link TrueForge.NotFoundError} * @throws {@link TrueForge.ConflictError} * @throws {@link TrueForge.InternalServerError} - * @throws {@link TrueForge.ServiceUnavailableError} * @throws {@link errors.TrueForgeError} * @throws {@link errors.TrueForgeTimeoutError} * @@ -1885,17 +1884,6 @@ export class SessionsClient { }), _response.rawResponse, ); - case 503: - throw new TrueForge.ServiceUnavailableError( - 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 3b92f8aeb..0a7996c6e 100644 --- a/packages/trueforge-sdk/tests/wire/sessions.test.ts +++ b/packages/trueforge-sdk/tests/wire/sessions.test.ts @@ -1584,35 +1584,6 @@ describe("SessionsClient", () => { }).rejects.toThrow(TrueForgeTypes.InternalServerError); }); - test("create_turn_events (7)", 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(503) - .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.ServiceUnavailableError); - }); - test("subscribe_to_turn (1)", async () => { const server = mockServerPool.createServer(); const client = new TrueForge({ maxRetries: 0, token: "test", baseUrl: server.baseUrl }); diff --git a/python/trueforge_sdk/src/trueforge_sdk/__init__.py b/python/trueforge_sdk/src/trueforge_sdk/__init__.py index 96ec2727b..7801ca100 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/__init__.py +++ b/python/trueforge_sdk/src/trueforge_sdk/__init__.py @@ -287,7 +287,6 @@ InternalServerError, NotFoundError, PreconditionFailedError, - ServiceUnavailableError, UnauthorizedError, UnprocessableEntityError, ) @@ -507,7 +506,6 @@ "ScheduleRun": ".types", "ScheduleRunStatus": ".types", "ScheduleStatus": ".types", - "ServiceUnavailableError": ".errors", "Session": ".types", "SessionAgent": ".types", "SessionAgentInline": ".types", @@ -830,7 +828,6 @@ def __dir__(): "ScheduleRun", "ScheduleRunStatus", "ScheduleStatus", - "ServiceUnavailableError", "Session", "SessionAgent", "SessionAgentInline", diff --git a/python/trueforge_sdk/src/trueforge_sdk/errors/__init__.py b/python/trueforge_sdk/src/trueforge_sdk/errors/__init__.py index ca21b04f3..07f9aa6be 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/errors/__init__.py +++ b/python/trueforge_sdk/src/trueforge_sdk/errors/__init__.py @@ -16,7 +16,6 @@ from .internal_server_error import InternalServerError from .not_found_error import NotFoundError from .precondition_failed_error import PreconditionFailedError - from .service_unavailable_error import ServiceUnavailableError from .unauthorized_error import UnauthorizedError from .unprocessable_entity_error import UnprocessableEntityError _dynamic_imports: typing.Dict[str, str] = { @@ -30,7 +29,6 @@ "InternalServerError": ".internal_server_error", "NotFoundError": ".not_found_error", "PreconditionFailedError": ".precondition_failed_error", - "ServiceUnavailableError": ".service_unavailable_error", "UnauthorizedError": ".unauthorized_error", "UnprocessableEntityError": ".unprocessable_entity_error", } @@ -68,7 +66,6 @@ def __dir__(): "InternalServerError", "NotFoundError", "PreconditionFailedError", - "ServiceUnavailableError", "UnauthorizedError", "UnprocessableEntityError", ] diff --git a/python/trueforge_sdk/src/trueforge_sdk/errors/service_unavailable_error.py b/python/trueforge_sdk/src/trueforge_sdk/errors/service_unavailable_error.py deleted file mode 100644 index 2cb49bd7c..000000000 --- a/python/trueforge_sdk/src/trueforge_sdk/errors/service_unavailable_error.py +++ /dev/null @@ -1,11 +0,0 @@ -# This file was auto-generated by Fern from our API Definition. - -import typing - -from ..core.api_error import ApiError -from ..types.request_error_response import RequestErrorResponse - - -class ServiceUnavailableError(ApiError): - def __init__(self, body: RequestErrorResponse, headers: typing.Optional[typing.Dict[str, str]] = None): - super().__init__(status_code=503, headers=headers, body=body) 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 674860eda..8cfda1477 100644 --- a/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py +++ b/python/trueforge_sdk/src/trueforge_sdk/sessions/raw_client.py @@ -28,7 +28,6 @@ from ..errors.internal_server_error import InternalServerError from ..errors.not_found_error import NotFoundError from ..errors.precondition_failed_error import PreconditionFailedError -from ..errors.service_unavailable_error import ServiceUnavailableError from ..errors.unauthorized_error import UnauthorizedError from ..errors.unprocessable_entity_error import UnprocessableEntityError from ..types.cancel_session_response import CancelSessionResponse @@ -1555,17 +1554,6 @@ def create_turn_events( ), ), ) - if _response.status_code == 503: - raise ServiceUnavailableError( - 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) @@ -3224,17 +3212,6 @@ async def create_turn_events( ), ), ) - if _response.status_code == 503: - raise ServiceUnavailableError( - 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) From fcaa0061d38dd495ccca59ba2fc947dd6269fbe5 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Fri, 23 Oct 2026 10:00:29 +0530 Subject: [PATCH 32/32] Fix tests --- .../src/convertTurnMessages.ts | 92 +++++++++---------- .../test/convertTurnMessages.test.ts | 1 + .../src/agent-session/SessionHandle.ts | 10 +- .../core/runtime/AgentThreadOrchestrator.ts | 8 +- .../tests/agent-session/turnResume.test.ts | 8 +- .../tests/orchestration/orchestration.test.ts | 2 + .../orchestrationApproval.test.ts | 7 +- .../orchestrationMcpAuth.test.ts | 2 + .../orchestrationUserMessageReset.test.ts | 6 ++ .../orchestrationWithTools.test.ts | 2 + packages/trueforge/src/apis/sessions.ts | 4 +- .../tests/unit/apis/sessionHttp.test.ts | 6 +- 12 files changed, 84 insertions(+), 64 deletions(-) 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/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 75046a186..c2505f86c 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -195,6 +195,7 @@ export class SessionHandle< previous, resolver: input.resolver, abortController: input.abortController, + pendingTurnEvents: [], }); const new_context_appends = await collectContextAppends(orchestrator.send(input.input ?? [])); @@ -421,7 +422,7 @@ export class SessionHandle< previous: TurnRecord | undefined; resolver: ITurnResourceResolver; abortController: AbortController; - pendingTurnEvents?: TurnUserEvent[] | undefined; + pendingTurnEvents: TurnUserEvent[]; }): Promise<{ orchestrator: AgentThreadOrchestrator; agentThreads: Map }> { const tracing = input.resolver.createTracing(); const agent = this.session.agent; @@ -474,10 +475,9 @@ export class SessionHandle< createDynamicSubAgentThread, tracing, logger: input.resolver.logger, - ...(input.pendingTurnEvents !== undefined ? { pendingTurnEvents: input.pendingTurnEvents } : {}), - ...(previousMcpServers !== null && previousMcpServers !== undefined - ? { mcpServers: Object.values(previousMcpServers) } - : {}), + pendingTurnEvents: input.pendingTurnEvents, + mcpServers: + previousMcpServers !== null && previousMcpServers !== undefined ? Object.values(previousMcpServers) : [], }); return { orchestrator, agentThreads }; diff --git a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts index 4bb027a5e..bf44470a8 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThreadOrchestrator.ts @@ -183,12 +183,12 @@ export interface AgentThreadOrchestratorInput { tracing: AgentTracing; logger: Logger; /** Inbox rows not yet in the snapshot. `execute` applies them. A rebuild passes these. */ - pendingTurnEvents?: TurnUserEvent[] | undefined; + 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[] | undefined; + mcpServers: MCPServerInitInfo[]; } export class AgentThreadOrchestrator { @@ -208,8 +208,8 @@ export class AgentThreadOrchestrator { this.createDynamicSubAgentThread = params.createDynamicSubAgentThread; this.tracing = params.tracing; this.logger = params.logger.child({ module: 'AgentThreadOrchestrator' }); - this.pendingTurnEvents = params.pendingTurnEvents !== undefined ? [...params.pendingTurnEvents] : []; - for (const server of params.mcpServers ?? []) { + this.pendingTurnEvents = [...params.pendingTurnEvents]; + for (const server of params.mcpServers) { this.mcpServerInitInfoById.set(server.id, server); } } diff --git a/packages/trueforge-core/tests/agent-session/turnResume.test.ts b/packages/trueforge-core/tests/agent-session/turnResume.test.ts index 88533b815..85aeeba71 100644 --- a/packages/trueforge-core/tests/agent-session/turnResume.test.ts +++ b/packages/trueforge-core/tests/agent-session/turnResume.test.ts @@ -289,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(), @@ -298,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 0b850fe69..1d3b0d4a9 100644 --- a/packages/trueforge-core/tests/orchestration/orchestrationApproval.test.ts +++ b/packages/trueforge-core/tests/orchestration/orchestrationApproval.test.ts @@ -646,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: [], }); }); @@ -835,6 +837,8 @@ 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 }; } @@ -886,7 +890,8 @@ function makeApprovalHarnessWithUserToolSet( createDynamicSubAgentThread: () => Promise.reject(new Error('unexpected sub-agent in approval policy test')), tracing: NOOP_AGENT_TRACING, logger: makeSilentLogger(), - ...(mcpServers !== undefined ? { mcpServers } : {}), + 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/src/apis/sessions.ts b/packages/trueforge/src/apis/sessions.ts index 5316a7db2..39b8e4276 100644 --- a/packages/trueforge/src/apis/sessions.ts +++ b/packages/trueforge/src/apis/sessions.ts @@ -198,8 +198,8 @@ export async function cancelSessionTurn( } } catch (error) { const fields = { - sessionId, - turnId, + session_id: sessionId, + turn_id: turnId, owner, ...extractErrorLogFields(error), }; diff --git a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts index 45e48b3db..15673b631 100644 --- a/packages/trueforge/tests/unit/apis/sessionHttp.test.ts +++ b/packages/trueforge/tests/unit/apis/sessionHttp.test.ts @@ -498,7 +498,11 @@ describe('sessions HTTP agent binding', () => { 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`,