From 717bb37d94e82fd7c8235416bc2f0f7d272f12f1 Mon Sep 17 00:00:00 2001 From: Heer Ambavi Date: Wed, 23 Sep 2026 16:36:02 +0530 Subject: [PATCH 1/5] add turn peering route and logic --- packages/trueforge/src/apis/peering.ts | 197 ++++++++++++ packages/trueforge/src/apis/sessions.ts | 131 ++------ packages/trueforge/src/runtime/activeTurns.ts | 21 +- .../trueforge/tests/unit/apis/peering.test.ts | 302 ++++++++++++++++++ .../tests/unit/apis/sessions.test.ts | 35 +- .../tests/unit/runtime/activeTurns.test.ts | 85 +++-- 6 files changed, 605 insertions(+), 166 deletions(-) create mode 100644 packages/trueforge/src/apis/peering.ts create mode 100644 packages/trueforge/tests/unit/apis/peering.test.ts diff --git a/packages/trueforge/src/apis/peering.ts b/packages/trueforge/src/apis/peering.ts new file mode 100644 index 000000000..1f33e33f3 --- /dev/null +++ b/packages/trueforge/src/apis/peering.ts @@ -0,0 +1,197 @@ +/** + * Redis request-reply peering: locate/cancel handlers, peer reachability, + * and turn-ownership resolution used by send/subscribe/cancel. + */ +import type { ISessionStore, TurnState } from '@truefoundry/trueforge-core/agent-session'; +import { CancellationReason, TurnNotFoundError } from '@truefoundry/trueforge-core/agent-session'; +import { + NoResponderError, + redisRequest, + type JSONValue, + type RedisClient, + type RouteHandler as RequestReplyRouteHandler, + type RequestReplyRouter, +} from '@truefoundry/trueforge-core/request-reply'; +import type { Logger } from 'winston'; +import { z } from 'zod'; +import configuration from '../config'; +import type { ActiveTurnRegistry } from '../runtime/activeTurns'; + +/** Request-reply path a replica serves to cancel a turn it owns. */ +export const SESSIONS_CANCEL_PATH = 'sessions/cancel'; + +/** Peer probe: does this executor have an ActiveTurn for the id? */ +export const TURNS_LOCATE_PATH = 'turns/locate'; + +/** 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), +}); + +const LocatePeerBodySchema = z.object({ + session_id: z.string(), + turn_id: z.string(), +}); + +/** + * Outcome of a Redis request-reply to another executor (locate, cancel, …). + * HTTP 200 → `ok`; heartbeat gone → `no_responder`; timeout / 412 / Redis is down → `failed`. + */ +export type PeerResult = 'ok' | 'no_responder' | 'failed'; + +/** + * What this replica should do: + * + * - `run` — execute here (we own it and have an ActiveTurn). + * - `rebuild` — rebuild ActiveTurn here (we own a paused tip with no run). + * - `forward` — send the work to the remote owner. + * - `steal` — claim ownership, then rebuild. + * - `reject` — do not continue (tip terminal, or local running with no ActiveTurn). + * - `retry` — do not continue now (remote owner not usable; caller may try again). + */ +export type OwnershipAction = 'run' | 'rebuild' | 'forward' | 'steal' | 'reject' | 'retry'; + +export interface ResolveTurnOwnershipDeps { + activeTurns: Pick; + sessionStore: Pick; + redis?: RedisClient | undefined; + logger: Pick; +} + +export function resolveOwnershipAction(input: { + status: TurnState['status']; + ownerIsLocal: boolean; + hasActiveTurn: boolean; + peerResult?: PeerResult; +}): OwnershipAction { + if (input.status === 'cancelled' || input.status === 'done' || input.status === 'error') { + return 'reject'; + } + + if (input.ownerIsLocal) { + if (input.peerResult === 'ok') { + throw new Error('peerResult "ok" is only valid for a remote owner'); + } + if (input.hasActiveTurn) { + return 'run'; + } + if (input.status === 'paused') { + return 'rebuild'; + } + return 'reject'; + } + + if (input.peerResult === 'ok') { + return 'forward'; + } + if (input.peerResult === 'no_responder' && input.status === 'paused') { + return 'steal'; + } + return 'retry'; +} + +/** Send a request-reply to `executorId`. Maps transport to {@link PeerResult}. */ +export async function callPeer(input: { + redis: RedisClient; + executorId: string; + path: string; + body: JSONValue; +}): Promise { + try { + const reply = await redisRequest({ + redis: input.redis, + executorId: input.executorId, + path: input.path, + request: { body: input.body }, + options: { + replyTimeoutMs: configuration.REDIS_REQUEST_REPLY_TIMEOUT_MS, + pollIntervalMs: configuration.REDIS_REQUEST_REPLY_POLL_INTERVAL_MS, + }, + }); + return reply.status === 200 ? 'ok' : 'failed'; + } catch (error) { + return error instanceof NoResponderError ? 'no_responder' : 'failed'; + } +} + +/** + * Load the turn, peer if another replica owns it, then {@link resolveOwnershipAction}. + * Steal is returned, not performed. + */ +export async function resolveTurnOwnership( + deps: ResolveTurnOwnershipDeps, + input: { sessionId: string; turnId: string }, +): Promise { + const turn = await deps.sessionStore.getTurn({ + session_id: input.sessionId, + turn_id: input.turnId, + }); + if (!turn) { + throw new TurnNotFoundError(input.turnId); + } + + const owner = turn.active_executor_id; + const ownerIsLocal = owner === configuration.EXECUTOR_ID; + const peerResult = + !ownerIsLocal && deps.redis + ? await callPeer({ + redis: deps.redis, + executorId: owner, + path: TURNS_LOCATE_PATH, + body: { session_id: input.sessionId, turn_id: input.turnId }, + }) + : undefined; + + return resolveOwnershipAction({ + status: turn.state.status, + ownerIsLocal, + hasActiveTurn: deps.activeTurns.has({ sessionId: input.sessionId, turnId: input.turnId }), + ...(peerResult === undefined ? {} : { peerResult }), + }); +} + +/** + * Peer-facing locate: 200 if this process has an ActiveTurn, 412 if not. + */ +export function locateTurnPeerHandler(activeTurns: ActiveTurnRegistry): RequestReplyRouteHandler { + return request => { + const parsed = LocatePeerBodySchema.safeParse(request.body); + if (!parsed.success) { + return Promise.resolve({ status: 400, body: { message: 'Invalid turns/locate payload' } }); + } + const found = activeTurns.has({ sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }); + return Promise.resolve( + found ? { status: 200, body: {} } : { status: 412, body: { message: 'Turn is not on this executor' } }, + ); + }; +} + +/** Peer-facing cancel: 200 if we had the run and aborted, 412 if not here. */ +export function cancelSessionTurnPeerHandler(activeTurns: ActiveTurnRegistry): RequestReplyRouteHandler { + // Synchronous by nature; the transport expects a Promise and require-await + // forbids an async fn without awaits. + return request => { + const parsed = CancelPeerBodySchema.safeParse(request.body); + if (!parsed.success) { + return Promise.resolve({ status: 400, body: { message: 'Invalid sessions/cancel payload' } }); + } + const found = activeTurns.has({ sessionId: parsed.data.session_id, turnId: parsed.data.turn_id }); + if (found) { + activeTurns.cancel({ + sessionId: parsed.data.session_id, + turnId: parsed.data.turn_id, + abortReason: parsed.data.reason, + }); + } + return Promise.resolve( + found ? { status: 200, body: {} } : { status: 412, body: { message: 'Turn is not running on this executor' } }, + ); + }; +} + +export function registerPeerRoutes(requestReplyRouter: RequestReplyRouter, activeTurns: ActiveTurnRegistry): void { + requestReplyRouter.registerRoute(SESSIONS_CANCEL_PATH, cancelSessionTurnPeerHandler(activeTurns)); + requestReplyRouter.registerRoute(TURNS_LOCATE_PATH, locateTurnPeerHandler(activeTurns)); +} diff --git a/packages/trueforge/src/apis/sessions.ts b/packages/trueforge/src/apis/sessions.ts index cbdf0a944..5a4633232 100644 --- a/packages/trueforge/src/apis/sessions.ts +++ b/packages/trueforge/src/apis/sessions.ts @@ -10,17 +10,9 @@ import { 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, 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'; @@ -45,19 +37,9 @@ import { honoQueriesToRecord } from '../schemas/deepObjectQuery'; import { isSessionAgentNameRef, parseListSessionsQuery, type Session } from '../schemas/session'; import { newId } from '../utils/id'; import { agentIfAccessible, canReadAgentBoundResource, resolveManagedAgentIds } from './agentAccess'; +import { callPeer, registerPeerRoutes, SESSIONS_CANCEL_PATH } from './peering'; 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 function toWireSession(record: SessionRecord): Session { return { id: record.session_id, @@ -89,40 +71,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, - }); - return Promise.resolve( - found ? { status: 200, body: {} } : { status: 412, body: { message: 'Turn is not running on this executor' } }, - ); - }; -} - /** A registry to abort in, a session to freeze, durable state to read, and a way to reach peers. */ export interface CancelTurnDeps { activeTurns: ActiveTurnRegistry; @@ -140,9 +88,10 @@ export interface CancelTurnDeps { * the terminal state. If abort cannot be confirmed, this replica freezes the * turn in the store so the session is not stuck `running`. * - * 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. + * Peer timeout, 412, and Redis/transport failures classify as `unavailable` + * They are not a clean cancellation — the owning replica may still be executing, + * but the turn is still frozen. + */ export async function cancelSessionTurn( deps: CancelTurnDeps, @@ -161,55 +110,35 @@ export async function cancelSessionTurn( } 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 running turn', fields); - } else { - deps.logger.warn('Failed to reach owning executor over Redis; freezing the running turn', fields); - } - } - await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); + const ownerIsLocal = owner === configuration.EXECUTOR_ID; + + if (ownerIsLocal && deps.activeTurns.has({ sessionId, turnId })) { + deps.activeTurns.cancel({ sessionId, turnId, abortReason: reason }); return; } - const aborted = cancelTurnOnThisExecutor(deps.activeTurns, { sessionId, turnId, reason }); - if (!aborted) { - await freezeTurnIgnoringMissing(deps.session, { turnId, reason }); + if (!ownerIsLocal && deps.redis) { + const peerResult = await callPeer({ + redis: deps.redis, + executorId: owner, + path: SESSIONS_CANCEL_PATH, + body: { session_id: sessionId, turn_id: turnId, reason }, + }); + // 200: peer already aborted the run. Do not freeze on top of that. + if (peerResult === 'ok') { + return; + } + deps.logger.warn('Owning executor not usable; freezing the running turn', { + sessionId, + turnId, + owner, + peerResult, + }); } -} -/** Freeze a running turn; missing turns are a no-op (already gone). */ -async function freezeTurnIgnoringMissing( - session: Pick, - input: { turnId: string; reason: CancellationReason }, -): Promise { + // Without Redis there is no peer to ask; freeze if the run is gone. try { - await session.freezeTurn({ turn_id: input.turnId, reason: input.reason }); + await deps.session.freezeTurn({ turn_id: turnId, reason }); } catch (error) { if (error instanceof TurnNotFoundError) { return; @@ -598,6 +527,6 @@ export function createSessionsRouter(deps: SessionsRouterDeps) { router.openapi(listSessionsRoute, listSessionsHandler); router.openapi(cancelSessionRoute, cancelSessionHandler); router.openapi(listSessionEventsRoute, listSessionEventsHandler); - deps.requestReplyRouter.registerRoute(SESSIONS_CANCEL_PATH, cancelSessionTurnPeerHandler(deps.activeTurns)); + registerPeerRoutes(deps.requestReplyRouter, deps.activeTurns); return router; } diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index a8fa87d6e..5ee5f8cbc 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -68,21 +68,20 @@ export class ActiveTurnRegistry { return tracked(); } + has(input: { sessionId: string; turnId: string }): boolean { + return this.runs.has(activeTurnKey(input.sessionId, input.turnId)); + } + /** - * Aborts the given turn if it is running in this process. Returns true when - * the run was found (already-aborted runs are not re-aborted). Cancelling a - * turn that is not running is a no-op, mirroring the store's - * first-terminal-write-wins rule. + * Aborts the turn if it is tracked here. Missing and already-aborted runs + * are a no-op (first abort wins). */ - cancelIfRunning(input: { sessionId: string; turnId: string; abortReason: CancellationReason }): boolean { + cancel(input: { sessionId: string; turnId: string; abortReason: CancellationReason }): void { const run = this.runs.get(activeTurnKey(input.sessionId, input.turnId)); - if (!run) { - return false; - } - if (!run.abortController.signal.aborted) { - run.abortController.abort(input.abortReason); + if (!run || run.abortController.signal.aborted) { + return; } - return true; + run.abortController.abort(input.abortReason); } /** diff --git a/packages/trueforge/tests/unit/apis/peering.test.ts b/packages/trueforge/tests/unit/apis/peering.test.ts new file mode 100644 index 000000000..2995bbe72 --- /dev/null +++ b/packages/trueforge/tests/unit/apis/peering.test.ts @@ -0,0 +1,302 @@ +import type { ISessionStore, TurnRecord, TurnState } from '@truefoundry/trueforge-core/agent-session'; +import { TurnNotFoundError } from '@truefoundry/trueforge-core/agent-session'; +import { NoResponderError, redisRequest, RequestTimeoutError } from '@truefoundry/trueforge-core/request-reply'; +import type { RedisClientType } from 'redis'; +import { + resolveOwnershipAction, + resolveTurnOwnership, + TURNS_LOCATE_PATH, + type OwnershipAction, + type PeerResult, +} from '../../../src/apis/peering'; +import configuration from '../../../src/config'; +import { ActiveTurnRegistry } from '../../../src/runtime/activeTurns'; + +jest.mock('@truefoundry/trueforge-core/request-reply', () => { + const actual = jest.requireActual( + '@truefoundry/trueforge-core/request-reply', + ); + return { + ...actual, + redisRequest: jest.fn(), + }; +}); + +const redisRequestMock = jest.mocked(redisRequest); + +const SESSION_ID = 's1'; +const REDIS = {} as RedisClientType; +const REMOTE_EXECUTOR = 'other1'; +const pausedState = { status: 'paused' as const, action_required_on_events: [] }; + +function silentLogger(): { warn: jest.Mock } { + return { warn: jest.fn() }; +} + +function turnRecord(input: { turnId: string; state: TurnState; activeExecutorId?: string }): TurnRecord { + return { + turn_id: input.turnId, + session_id: SESSION_ID, + first_turn_id: input.turnId, + ancestor_ids: [], + previous_turn_id: null, + active_executor_id: input.activeExecutorId ?? configuration.EXECUTOR_ID, + state: input.state, + input: [], + snapshot: { threads: {}, mcp_servers: null, sandbox_info: null }, + created_at: new Date('2026-07-31T00:00:00.000Z'), + updated_at: new Date('2026-07-31T00:00:00.000Z'), + custom: null, + }; +} + +function storeReturning(turn: TurnRecord | undefined): Pick { + return { getTurn: () => Promise.resolve(turn) }; +} + +function ownershipDeps(input: { + activeTurns: ActiveTurnRegistry; + turn: TurnRecord | undefined; + redis?: RedisClientType; + logger?: { warn: jest.Mock }; +}): { + activeTurns: ActiveTurnRegistry; + sessionStore: Pick; + redis?: RedisClientType; + logger: { warn: jest.Mock }; +} { + return { + activeTurns: input.activeTurns, + sessionStore: storeReturning(input.turn), + ...(input.redis === undefined ? {} : { redis: input.redis }), + logger: input.logger ?? silentLogger(), + }; +} + +/** Registers a live run without consuming its stream, mirroring a turn mid-execution. */ +function trackRun(registry: ActiveTurnRegistry, turnId: string): AbortController { + const abortController = new AbortController(); + registry.track({ + sessionId: SESSION_ID, + turnId, + abortController, + stream: (async function* () { + await new Promise(() => undefined); + yield 'never'; + })(), + }); + return abortController; +} + +describe('resolveTurnOwnership', () => { + beforeEach(() => { + redisRequestMock.mockReset(); + }); + + it('throws when the turn is missing', async () => { + await expect( + resolveTurnOwnership(ownershipDeps({ activeTurns: new ActiveTurnRegistry(), turn: undefined }), { + sessionId: SESSION_ID, + turnId: 'missing', + }), + ).rejects.toBeInstanceOf(TurnNotFoundError); + }); + + it('runs locally when this executor owns a live ActiveTurn', async () => { + const activeTurns = new ActiveTurnRegistry(); + const turnId = 'turn-local'; + trackRun(activeTurns, turnId); + + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns, + turn: turnRecord({ turnId, state: { status: 'running' } }), + }), + { sessionId: SESSION_ID, turnId }, + ), + ).resolves.toBe('run'); + }); + + it('rebuilds when this executor owns a paused turn with no ActiveTurn', async () => { + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns: new ActiveTurnRegistry(), + turn: turnRecord({ turnId: 'paused-local', state: pausedState }), + }), + { sessionId: SESSION_ID, turnId: 'paused-local' }, + ), + ).resolves.toBe('rebuild'); + }); + + it('returns forward when the owning replica still has the turn', async () => { + redisRequestMock.mockResolvedValue({ status: 200, body: {} }); + + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns: new ActiveTurnRegistry(), + turn: turnRecord({ turnId: 'remote-ok', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), + redis: REDIS, + }), + { sessionId: SESSION_ID, turnId: 'remote-ok' }, + ), + ).resolves.toBe('forward'); + expect(redisRequestMock).toHaveBeenCalledWith( + expect.objectContaining({ executorId: REMOTE_EXECUTOR, path: TURNS_LOCATE_PATH }), + ); + }); + + it('does not steal on 412 / timeout (unavailable)', async () => { + redisRequestMock.mockResolvedValue({ status: 412, body: { message: 'Turn is not on this executor' } }); + + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns: new ActiveTurnRegistry(), + turn: turnRecord({ turnId: 'remote-412', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), + redis: REDIS, + }), + { sessionId: SESSION_ID, turnId: 'remote-412' }, + ), + ).resolves.toBe('retry'); + }); + + it('steals a paused turn when there is no responder', async () => { + redisRequestMock.mockRejectedValue(new NoResponderError(REMOTE_EXECUTOR)); + + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns: new ActiveTurnRegistry(), + turn: turnRecord({ turnId: 'remote-steal', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), + redis: REDIS, + }), + { sessionId: SESSION_ID, turnId: 'remote-steal' }, + ), + ).resolves.toBe('steal'); + }); + + it('does not steal a running turn when there is no responder', async () => { + redisRequestMock.mockRejectedValue(new NoResponderError(REMOTE_EXECUTOR)); + + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns: new ActiveTurnRegistry(), + turn: turnRecord({ + turnId: 'remote-running', + state: { status: 'running' }, + activeExecutorId: REMOTE_EXECUTOR, + }), + redis: REDIS, + }), + { sessionId: SESSION_ID, turnId: 'remote-running' }, + ), + ).resolves.toBe('retry'); + }); + + it('does not steal on peer timeout', async () => { + redisRequestMock.mockRejectedValue(new RequestTimeoutError(60_000)); + + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns: new ActiveTurnRegistry(), + turn: turnRecord({ turnId: 'remote-timeout', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), + redis: REDIS, + }), + { sessionId: SESSION_ID, turnId: 'remote-timeout' }, + ), + ).resolves.toBe('retry'); + }); +}); + +describe('resolveOwnershipAction', () => { + it.each([ + { + name: 'local paused with ActiveTurn → run here', + input: { status: 'paused' as const, ownerIsLocal: true, hasActiveTurn: true }, + expect: 'run' satisfies OwnershipAction, + }, + { + name: 'local paused without ActiveTurn → run here (rebuild)', + input: { status: 'paused' as const, ownerIsLocal: true, hasActiveTurn: false }, + expect: 'rebuild' satisfies OwnershipAction, + }, + { + name: 'local running with ActiveTurn → run here', + input: { status: 'running' as const, ownerIsLocal: true, hasActiveTurn: true }, + expect: 'run' satisfies OwnershipAction, + }, + { + name: 'local running without ActiveTurn → reject', + input: { status: 'running' as const, ownerIsLocal: true, hasActiveTurn: false }, + expect: 'reject' satisfies OwnershipAction, + }, + { + name: 'remote + ok → forward', + input: { + status: 'paused' as const, + ownerIsLocal: false, + hasActiveTurn: false, + peerResult: 'ok' as const, + }, + expect: 'forward' satisfies OwnershipAction, + }, + { + name: 'remote paused + no_responder → steal', + input: { + status: 'paused' as const, + ownerIsLocal: false, + hasActiveTurn: false, + peerResult: 'no_responder' as PeerResult, + }, + expect: 'steal' satisfies OwnershipAction, + }, + { + name: 'remote paused + failed → retry (no steal)', + input: { + status: 'paused' as const, + ownerIsLocal: false, + hasActiveTurn: false, + peerResult: 'failed' as PeerResult, + }, + expect: 'retry' satisfies OwnershipAction, + }, + { + name: 'remote running + no_responder → retry (no steal)', + input: { + status: 'running' as const, + ownerIsLocal: false, + hasActiveTurn: false, + peerResult: 'no_responder' as PeerResult, + }, + expect: 'retry' satisfies OwnershipAction, + }, + { + name: 'cancelled → reject', + input: { status: 'cancelled' as const, ownerIsLocal: true, hasActiveTurn: false }, + expect: 'reject' satisfies OwnershipAction, + }, + { + name: 'done → reject', + input: { status: 'done' as const, ownerIsLocal: true, hasActiveTurn: false }, + expect: 'reject' satisfies OwnershipAction, + }, + ])('$name', ({ input, expect: expected }) => { + expect(resolveOwnershipAction(input)).toEqual(expected); + }); + + it('throws if ok is passed for a local owner', () => { + expect(() => + resolveOwnershipAction({ + status: 'running', + ownerIsLocal: true, + hasActiveTurn: true, + peerResult: 'ok', + }), + ).toThrow('peerResult "ok" is only valid for a remote owner'); + }); +}); diff --git a/packages/trueforge/tests/unit/apis/sessions.test.ts b/packages/trueforge/tests/unit/apis/sessions.test.ts index ddc8ef3db..1c22994f2 100644 --- a/packages/trueforge/tests/unit/apis/sessions.test.ts +++ b/packages/trueforge/tests/unit/apis/sessions.test.ts @@ -2,6 +2,7 @@ import type { ISessionStore, SessionHandle, TurnRecord, TurnState } from '@truef 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 { SESSIONS_CANCEL_PATH } from '../../../src/apis/peering'; import { cancelSessionTurn } from '../../../src/apis/sessions'; import configuration from '../../../src/config'; import { ActiveTurnRegistry } from '../../../src/runtime/activeTurns'; @@ -171,6 +172,7 @@ describe('cancelSessionTurn', () => { expect(redisRequestMock).toHaveBeenCalledWith( expect.objectContaining({ executorId: REMOTE_EXECUTOR, + path: SESSIONS_CANCEL_PATH, }), ); expect(session.freezeTurn).not.toHaveBeenCalled(); @@ -180,6 +182,7 @@ describe('cancelSessionTurn', () => { const activeTurns = new ActiveTurnRegistry(); const turnId = 'turn-remote-412'; const session = sessionHandle(); + const logger = silentLogger(); redisRequestMock.mockResolvedValue({ status: 412, body: { message: 'Turn is not running on this executor' } }); await cancelSessionTurn( @@ -188,11 +191,21 @@ describe('cancelSessionTurn', () => { turn: turnRecord({ turnId: turnId, state: { status: 'running' }, activeExecutorId: REMOTE_EXECUTOR }), session, redis: REDIS, + logger, }), { turnId }, ); expect(session.freezeTurn).toHaveBeenCalled(); + expect(logger.warn).toHaveBeenCalledWith( + 'Owning executor not usable; freezing the running turn', + expect.objectContaining({ + sessionId: SESSION_ID, + turnId, + owner: REMOTE_EXECUTOR, + peerResult: 'failed', + }), + ); }); it('freezes when the owning executor is unreachable', async () => { @@ -216,8 +229,13 @@ describe('cancelSessionTurn', () => { ).resolves.toBeUndefined(); expect(session.freezeTurn).toHaveBeenCalled(); expect(logger.warn).toHaveBeenCalledWith( - 'Failed to reach owning executor over Redis; freezing the running turn', - expect.objectContaining({ sessionId: SESSION_ID, turnId, owner: REMOTE_EXECUTOR }), + 'Owning executor not usable; freezing the running turn', + expect.objectContaining({ + sessionId: SESSION_ID, + turnId, + owner: REMOTE_EXECUTOR, + peerResult: 'no_responder', + }), ); }); @@ -242,8 +260,13 @@ describe('cancelSessionTurn', () => { ).resolves.toBeUndefined(); expect(session.freezeTurn).toHaveBeenCalled(); expect(logger.warn).toHaveBeenCalledWith( - 'Timed out waiting for owning executor to cancel; freezing the running turn', - expect.objectContaining({ sessionId: SESSION_ID, turnId, owner: REMOTE_EXECUTOR }), + 'Owning executor not usable; freezing the running turn', + expect.objectContaining({ + sessionId: SESSION_ID, + turnId, + owner: REMOTE_EXECUTOR, + peerResult: 'failed', + }), ); }); @@ -268,12 +291,12 @@ describe('cancelSessionTurn', () => { ).resolves.toBeUndefined(); expect(session.freezeTurn).toHaveBeenCalled(); expect(logger.warn).toHaveBeenCalledWith( - 'Failed to reach owning executor over Redis; freezing the running turn', + 'Owning executor not usable; freezing the running turn', expect.objectContaining({ sessionId: SESSION_ID, turnId, owner: REMOTE_EXECUTOR, - error: 'Redis connection closed', + peerResult: 'failed', }), ); }); diff --git a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts index 6754ea907..7e4e678d7 100644 --- a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts +++ b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts @@ -42,13 +42,7 @@ describe('ActiveTurnRegistry', () => { seen.push(value); } expect(seen).toEqual([1, 2, 3]); - expect( - registry.cancelIfRunning({ - sessionId: 's1', - turnId: 't1', - abortReason: CancellationReason.ClientCancelled, - }), - ).toBe(false); + expect(registry.has({ sessionId: 's1', turnId: 't1' })).toBe(false); }); it('track cleans up when the consumer breaks early', async () => { @@ -65,13 +59,7 @@ describe('ActiveTurnRegistry', () => { expect(value).toBe(1); break; } - expect( - registry.cancelIfRunning({ - sessionId: 's1', - turnId: 't1', - abortReason: CancellationReason.ClientCancelled, - }), - ).toBe(false); + expect(registry.has({ sessionId: 's1', turnId: 't1' })).toBe(false); }); it('track cleans up when the stream throws', async () => { @@ -89,16 +77,27 @@ describe('ActiveTurnRegistry', () => { void value; } }).rejects.toThrow(/stream boom/); - expect( - registry.cancelIfRunning({ - sessionId: 's1', - turnId: 't1', - abortReason: CancellationReason.ClientCancelled, - }), - ).toBe(false); + expect(registry.has({ sessionId: 's1', turnId: 't1' })).toBe(false); + }); + + it('has is true only while the run is tracked', async () => { + const registry = new ActiveTurnRegistry(); + const abortController = new AbortController(); + expect(registry.has({ sessionId: 's1', turnId: 't1' })).toBe(false); + const tracked = registry.track({ + sessionId: 's1', + turnId: 't1', + abortController, + stream: values([1]), + }); + expect(registry.has({ sessionId: 's1', turnId: 't1' })).toBe(true); + for await (const value of tracked) { + void value; + } + expect(registry.has({ sessionId: 's1', turnId: 't1' })).toBe(false); }); - it('cancelIfRunning aborts with the given reason and returns true', () => { + it('cancel aborts with the given reason', () => { const registry = new ActiveTurnRegistry(); const abortController = new AbortController(); void registry.track({ @@ -108,29 +107,27 @@ describe('ActiveTurnRegistry', () => { stream: values([1]), }); - expect( - registry.cancelIfRunning({ - sessionId: 's1', - turnId: 't1', - abortReason: CancellationReason.ClientCancelled, - }), - ).toBe(true); + registry.cancel({ + sessionId: 's1', + turnId: 't1', + abortReason: CancellationReason.ClientCancelled, + }); expect(abortController.signal.aborted).toBe(true); expect(abortController.signal.reason).toBe(CancellationReason.ClientCancelled); }); - it('cancelIfRunning returns false for unknown ids', () => { + it('cancel is a no-op for unknown ids', () => { const registry = new ActiveTurnRegistry(); - expect( - registry.cancelIfRunning({ + expect(() => + registry.cancel({ sessionId: 'missing', turnId: 'missing', abortReason: CancellationReason.ClientCancelled, }), - ).toBe(false); + ).not.toThrow(); }); - it('cancelIfRunning does not re-abort an already-aborted controller', () => { + it('cancel does not re-abort an already-aborted controller', () => { const registry = new ActiveTurnRegistry(); const abortController = new AbortController(); void registry.track({ @@ -141,13 +138,11 @@ describe('ActiveTurnRegistry', () => { }); abortController.abort(CancellationReason.ClientCancelled); - expect( - registry.cancelIfRunning({ - sessionId: 's1', - turnId: 't1', - abortReason: CancellationReason.Abandoned, - }), - ).toBe(true); + registry.cancel({ + sessionId: 's1', + turnId: 't1', + abortReason: CancellationReason.Abandoned, + }); expect(abortController.signal.reason).toBe(CancellationReason.ClientCancelled); }); @@ -171,13 +166,7 @@ describe('ActiveTurnRegistry', () => { expect(abortController.signal.aborted).toBe(true); expect(abortController.signal.reason).toBe(CancellationReason.Abandoned); await drain; - expect( - registry.cancelIfRunning({ - sessionId: 's1', - turnId: 't1', - abortReason: CancellationReason.ClientCancelled, - }), - ).toBe(false); + expect(registry.has({ sessionId: 's1', turnId: 't1' })).toBe(false); }); it('late track after shutdownAndWait aborts immediately with the shutdown reason', async () => { From a775f43b69f9d917ef597ac3bed8871deeb095f9 Mon Sep 17 00:00:00 2001 From: Heer Ambavi Date: Wed, 23 Sep 2026 16:49:47 +0530 Subject: [PATCH 2/5] add turn ownership update - handle steal --- .../trueforge-core/src/agent-session/index.ts | 1 + .../src/agent-session/store/ISessionStore.ts | 15 ++++ .../store/InMemorySessionStore.ts | 11 +++ .../agent-session/store/storeContractSuite.ts | 62 +++++++++++++++ packages/trueforge/src/apis/peering.ts | 76 +++++++++++-------- .../session-store/PostgresSessionStore.ts | 6 ++ .../postgres/session-store/queries/turns.ts | 36 +++++++++ .../session-store/SqliteSessionStore.ts | 6 ++ .../db/sqlite/session-store/queries/turns.ts | 32 ++++++++ packages/trueforge/src/runtime/activeTurns.ts | 25 ++++++ .../trueforge/tests/unit/apis/peering.test.ts | 43 +++++++++-- .../tests/unit/runtime/activeTurns.test.ts | 63 +++++++++++++++ 12 files changed, 339 insertions(+), 37 deletions(-) diff --git a/packages/trueforge-core/src/agent-session/index.ts b/packages/trueforge-core/src/agent-session/index.ts index f380dd585..8cdc0a9d3 100644 --- a/packages/trueforge-core/src/agent-session/index.ts +++ b/packages/trueforge-core/src/agent-session/index.ts @@ -81,6 +81,7 @@ export type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, + ClaimTurnOwnershipInput, CreateSessionInput, CreateTurnInput, DeleteSessionInput, diff --git a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts index e423cfff6..fe070f8a2 100644 --- a/packages/trueforge-core/src/agent-session/store/ISessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/ISessionStore.ts @@ -164,6 +164,14 @@ export interface UpdateTurnStateInput { turn_done_event: PersistedTurnEvent; } +/** Steal: one winner when two replicas claim a paused turn. */ +export interface ClaimTurnOwnershipInput { + session_id: string; + turn_id: string; + expected_active_executor_id: string; + new_active_executor_id: string; +} + export interface AppendToEventsInput { session_id: string; turn_id: string; @@ -364,6 +372,13 @@ export interface ISessionStore< */ updateTurnState(input: UpdateTurnStateInput): Promise; + /** + * Claim `active_executor_id` when the turn is owned by + * `expected_active_executor_id`. True if this caller won. False if another + * replica already claimed or the tip is not stealable. Missing → {@link TurnNotFoundError}. + */ + claimTurnOwnership(input: ClaimTurnOwnershipInput): Promise; + /** * Durable event log for the turn. MUST include lifecycle rows: a * `TurnCreatedEvent` at the start of the stream and a terminal `TurnDoneEvent` diff --git a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts index a0a19615a..0e4a23761 100644 --- a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts @@ -10,6 +10,7 @@ import type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, + ClaimTurnOwnershipInput, CreateSessionInput, CreateTurnInput, DeleteSessionInput, @@ -490,6 +491,16 @@ export class InMemorySessionStore< this.addTerminalSessionMetrics(input.session_id, turn.created_at, input.state); } + async claimTurnOwnership(input: ClaimTurnOwnershipInput): Promise { + const turn = this.requireTurn(input.session_id, input.turn_id); + if (turn.state.status !== 'paused' || turn.active_executor_id !== input.expected_active_executor_id) { + return false; + } + turn.active_executor_id = input.new_active_executor_id; + turn.updated_at = new Date(); + return true; + } + async appendToEvents(input: AppendToEventsInput): Promise { this.requireRunningTurn(input.session_id, input.turn_id); const tKey = turnKey(input); diff --git a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts index 36a8f0eba..6d77fab73 100644 --- a/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts +++ b/packages/trueforge-core/tests/agent-session/store/storeContractSuite.ts @@ -2199,6 +2199,68 @@ export function runStoreContractSuite(createStore: () => ISessionStore) { }); }); + describe('claimTurnOwnership', () => { + it('missing turn → not found', async () => { + const store = createStore(); + await seedSession(store); + await expect( + store.claimTurnOwnership({ + session_id: sessionId, + turn_id: missingTurnId, + expected_active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + new_active_executor_id: 'stealer', + }), + ).rejects.toBeInstanceOf(TurnNotFoundError); + }); + + it('running turn → false and owner unchanged', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + await expect( + store.claimTurnOwnership({ + session_id: sessionId, + turn_id: 'turn-1', + expected_active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + new_active_executor_id: 'stealer', + }), + ).resolves.toBe(false); + const turn = await store.getTurn({ session_id: sessionId, turn_id: 'turn-1' }); + expect(mustGet(turn).active_executor_id).toBe(TEST_ACTIVE_EXECUTOR_ID); + }); + + it('wrong expected owner → false', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + await expect( + store.claimTurnOwnership({ + session_id: sessionId, + turn_id: 'turn-1', + expected_active_executor_id: 'someone-else', + new_active_executor_id: 'stealer', + }), + ).resolves.toBe(false); + const turn = await store.getTurn({ session_id: sessionId, turn_id: 'turn-1' }); + expect(mustGet(turn).active_executor_id).toBe(TEST_ACTIVE_EXECUTOR_ID); + }); + + it('terminal turn → false', async () => { + const store = createStore(); + await seedSession(store); + await store.createTurn(makeCreateTurnInput({ sessionId, turnId: 'turn-1' })); + await finishTurn(store, 'turn-1'); + await expect( + store.claimTurnOwnership({ + session_id: sessionId, + turn_id: 'turn-1', + expected_active_executor_id: TEST_ACTIVE_EXECUTOR_ID, + new_active_executor_id: 'stealer', + }), + ).resolves.toBe(false); + }); + }); + describe('events + threads + capability_state', () => { it('appendToEvents orders by monotonic event id, not append call order', async () => { const store = createStore(); diff --git a/packages/trueforge/src/apis/peering.ts b/packages/trueforge/src/apis/peering.ts index 1f33e33f3..f3b15cb29 100644 --- a/packages/trueforge/src/apis/peering.ts +++ b/packages/trueforge/src/apis/peering.ts @@ -45,17 +45,17 @@ export type PeerResult = 'ok' | 'no_responder' | 'failed'; * What this replica should do: * * - `run` — execute here (we own it and have an ActiveTurn). - * - `rebuild` — rebuild ActiveTurn here (we own a paused tip with no run). + * - `rebuild` — rebuild ActiveTurn here (we own a paused turn with no run). * - `forward` — send the work to the remote owner. - * - `steal` — claim ownership, then rebuild. - * - `reject` — do not continue (tip terminal, or local running with no ActiveTurn). + * - `steal` — table-only: claim the paused turn (resolveTurnOwnership then rebuilds or retries). + * - `reject` — do not continue (turn terminal, or local running with no ActiveTurn). * - `retry` — do not continue now (remote owner not usable; caller may try again). */ export type OwnershipAction = 'run' | 'rebuild' | 'forward' | 'steal' | 'reject' | 'retry'; export interface ResolveTurnOwnershipDeps { - activeTurns: Pick; - sessionStore: Pick; + activeTurns: Pick; + sessionStore: Pick; redis?: RedisClient | undefined; logger: Pick; } @@ -117,38 +117,52 @@ export async function callPeer(input: { } /** - * Load the turn, peer if another replica owns it, then {@link resolveOwnershipAction}. - * Steal is returned, not performed. + * Load the turn under the per-turn lock, peer if another replica owns it, then + * {@link resolveOwnershipAction}. A `steal` is claimed here (R5): winner → + * `rebuild`, loser → `retry`. */ export async function resolveTurnOwnership( deps: ResolveTurnOwnershipDeps, input: { sessionId: string; turnId: string }, ): Promise { - const turn = await deps.sessionStore.getTurn({ - session_id: input.sessionId, - turn_id: input.turnId, - }); - if (!turn) { - throw new TurnNotFoundError(input.turnId); - } + return deps.activeTurns.withTurnLock({ sessionId: input.sessionId, turnId: input.turnId }, async () => { + const turn = await deps.sessionStore.getTurn({ + session_id: input.sessionId, + turn_id: input.turnId, + }); + if (!turn) { + throw new TurnNotFoundError(input.turnId); + } - const owner = turn.active_executor_id; - const ownerIsLocal = owner === configuration.EXECUTOR_ID; - const peerResult = - !ownerIsLocal && deps.redis - ? await callPeer({ - redis: deps.redis, - executorId: owner, - path: TURNS_LOCATE_PATH, - body: { session_id: input.sessionId, turn_id: input.turnId }, - }) - : undefined; - - return resolveOwnershipAction({ - status: turn.state.status, - ownerIsLocal, - hasActiveTurn: deps.activeTurns.has({ sessionId: input.sessionId, turnId: input.turnId }), - ...(peerResult === undefined ? {} : { peerResult }), + const owner = turn.active_executor_id; + const ownerIsLocal = owner === configuration.EXECUTOR_ID; + const peerResult = + !ownerIsLocal && deps.redis + ? await callPeer({ + redis: deps.redis, + executorId: owner, + path: TURNS_LOCATE_PATH, + body: { session_id: input.sessionId, turn_id: input.turnId }, + }) + : undefined; + + const action = resolveOwnershipAction({ + status: turn.state.status, + ownerIsLocal, + hasActiveTurn: deps.activeTurns.has({ sessionId: input.sessionId, turnId: input.turnId }), + ...(peerResult === undefined ? {} : { peerResult }), + }); + if (action !== 'steal') { + return action; + } + + const won = await deps.sessionStore.claimTurnOwnership({ + session_id: input.sessionId, + turn_id: input.turnId, + expected_active_executor_id: owner, + new_active_executor_id: configuration.EXECUTOR_ID, + }); + return won ? 'rebuild' : 'retry'; }); } diff --git a/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts b/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts index fff1f5ca6..994984bde 100644 --- a/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts +++ b/packages/trueforge/src/db/postgres/session-store/PostgresSessionStore.ts @@ -13,6 +13,7 @@ import type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, + ClaimTurnOwnershipInput, CreateSessionInput, CreateTurnInput, DeleteSessionInput, @@ -73,6 +74,7 @@ import { } from './queries/threads'; import type { NewThreadRegistration } from './queries/turns'; import { + claimTurnOwnership as claimTurnOwnershipQuery, createTurn as createTurnQuery, freezeAndGetTurn as freezeAndGetTurnQuery, getTurn as getTurnQuery, @@ -214,6 +216,10 @@ export class PostgresSessionStore implements ISessionStore { + return claimTurnOwnershipQuery(this.db, input); + } + appendToEvents(input: AppendToEventsInput): Promise { return appendToEventsQuery(this.db, input); } diff --git a/packages/trueforge/src/db/postgres/session-store/queries/turns.ts b/packages/trueforge/src/db/postgres/session-store/queries/turns.ts index 4180945f8..0233a3d7e 100644 --- a/packages/trueforge/src/db/postgres/session-store/queries/turns.ts +++ b/packages/trueforge/src/db/postgres/session-store/queries/turns.ts @@ -7,6 +7,7 @@ import { } from '@truefoundry/trueforge-core/agent-session/schemas/turn'; import { assertCreateTurnThreadDelta } from '@truefoundry/trueforge-core/agent-session/store/assertCreateTurnThreadDelta'; import type { + ClaimTurnOwnershipInput, FreezeAndGetTurnInput, TurnRecordWithoutSnapshot, UpdateTurnStateInput, @@ -796,3 +797,38 @@ export async function updateTurnState(db: Kysely, input: UpdateTurnSta .execute(); }); } + +/** + * Steal CAS: one row updated iff still paused and owned by `expected`. + * 0 rows → false (or TurnNotFoundError if the turn is missing). + */ +export async function claimTurnOwnership(db: Kysely, input: ClaimTurnOwnershipInput): Promise { + const result = await db + .updateTable('turn') + .set({ + active_executor_id: input.new_active_executor_id, + updated_at: sql`now()`, + }) + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .where('active_executor_id', '=', input.expected_active_executor_id) + .where(sql`state->>'status' = 'paused'`) + .returning(['turn_id']) + .executeTakeFirst(); + + if (result !== undefined) { + return true; + } + + const existing = await db + .selectFrom('turn') + .select('turn_id') + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .executeTakeFirst(); + + if (!existing) { + throw new TurnNotFoundError(input.turn_id); + } + return false; +} diff --git a/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts b/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts index 57b4fe3b5..2e8e1c660 100644 --- a/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts +++ b/packages/trueforge/src/db/sqlite/session-store/SqliteSessionStore.ts @@ -6,6 +6,7 @@ import type { AddThreadsInput, AppendToEventsInput, AppendToThreadContextInput, + ClaimTurnOwnershipInput, CreateSessionInput, CreateTurnInput, DeleteSessionInput, @@ -61,6 +62,7 @@ import { } from './queries/threads'; import type { NewThreadRegistration } from './queries/turns'; import { + claimTurnOwnership as claimTurnOwnershipQuery, createTurn as createTurnQuery, freezeAndGetTurn as freezeAndGetTurnQuery, getTurn as getTurnQuery, @@ -189,6 +191,10 @@ export class SqliteSessionStore implements ISessionStore { + return claimTurnOwnershipQuery(this.db, input); + } + appendToEvents(input: AppendToEventsInput): Promise { return appendToEventsQuery(this.db, input); } diff --git a/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts b/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts index 160558471..443d558be 100644 --- a/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts +++ b/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts @@ -7,6 +7,7 @@ import { } from '@truefoundry/trueforge-core/agent-session/schemas/turn'; import { assertCreateTurnThreadDelta } from '@truefoundry/trueforge-core/agent-session/store/assertCreateTurnThreadDelta'; import type { + ClaimTurnOwnershipInput, FreezeAndGetTurnInput, TurnRecordWithoutSnapshot, UpdateTurnStateInput, @@ -848,3 +849,34 @@ export async function updateTurnState(db: Kysely, input: UpdateTurnSta .execute(); }); } + +export async function claimTurnOwnership(db: Kysely, input: ClaimTurnOwnershipInput): Promise { + const result = await db + .updateTable('turn') + .set({ + active_executor_id: input.new_active_executor_id, + updated_at: nowIso(), + }) + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .where('active_executor_id', '=', input.expected_active_executor_id) + .where(sql`state->>'status' = 'paused'`) + .returning(['turn_id']) + .executeTakeFirst(); + + if (result !== undefined) { + return true; + } + + const existing = await db + .selectFrom('turn') + .select('turn_id') + .where('session_id', '=', input.session_id) + .where('turn_id', '=', input.turn_id) + .executeTakeFirst(); + + if (!existing) { + throw new TurnNotFoundError(input.turn_id); + } + return false; +} diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index 5ee5f8cbc..c0be08344 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -18,8 +18,33 @@ function activeTurnKey(sessionId: string, turnId: string): string { export class ActiveTurnRegistry { private readonly runs = new Map(); + private readonly locks = new Map>(); private alreadyShutDownAbortReason: CancellationReason | undefined; + /** + * Serialize work for one turn in this process. Waiters queue; different keys + * run in parallel. Re-read registry / DB inside `fn` — do not trust values + * from before the lock. + */ + async withTurnLock(input: { sessionId: string; turnId: string }, fn: () => Promise): Promise { + const key = activeTurnKey(input.sessionId, input.turnId); + const previous = this.locks.get(key) ?? Promise.resolve(); + let release!: () => void; + const held = new Promise(resolve => { + release = resolve; + }); + this.locks.set(key, held); + await previous; + try { + return await fn(); + } finally { + release(); + if (this.locks.get(key) === held) { + this.locks.delete(key); + } + } + } + /** * Registers the run immediately, then returns a generator that forwards * `stream` and removes the run when the stream completes (or the consumer diff --git a/packages/trueforge/tests/unit/apis/peering.test.ts b/packages/trueforge/tests/unit/apis/peering.test.ts index 2995bbe72..4699d189c 100644 --- a/packages/trueforge/tests/unit/apis/peering.test.ts +++ b/packages/trueforge/tests/unit/apis/peering.test.ts @@ -50,8 +50,14 @@ function turnRecord(input: { turnId: string; state: TurnState; activeExecutorId? }; } -function storeReturning(turn: TurnRecord | undefined): Pick { - return { getTurn: () => Promise.resolve(turn) }; +function storeReturning( + turn: TurnRecord | undefined, + claimTurnOwnership: ISessionStore['claimTurnOwnership'] = () => Promise.resolve(true), +): Pick { + return { + getTurn: () => Promise.resolve(turn), + claimTurnOwnership, + }; } function ownershipDeps(input: { @@ -59,15 +65,16 @@ function ownershipDeps(input: { turn: TurnRecord | undefined; redis?: RedisClientType; logger?: { warn: jest.Mock }; + claimTurnOwnership?: ISessionStore['claimTurnOwnership']; }): { activeTurns: ActiveTurnRegistry; - sessionStore: Pick; + sessionStore: Pick; redis?: RedisClientType; logger: { warn: jest.Mock }; } { return { activeTurns: input.activeTurns, - sessionStore: storeReturning(input.turn), + sessionStore: storeReturning(input.turn, input.claimTurnOwnership), ...(input.redis === undefined ? {} : { redis: input.redis }), logger: input.logger ?? silentLogger(), }; @@ -163,8 +170,9 @@ describe('resolveTurnOwnership', () => { ).resolves.toBe('retry'); }); - it('steals a paused turn when there is no responder', async () => { + it('claims a paused turn with no responder and rebuilds on a winning CAS', async () => { redisRequestMock.mockRejectedValue(new NoResponderError(REMOTE_EXECUTOR)); + const claimTurnOwnership = jest.fn().mockResolvedValue(true); await expect( resolveTurnOwnership( @@ -172,10 +180,33 @@ describe('resolveTurnOwnership', () => { activeTurns: new ActiveTurnRegistry(), turn: turnRecord({ turnId: 'remote-steal', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), redis: REDIS, + claimTurnOwnership, }), { sessionId: SESSION_ID, turnId: 'remote-steal' }, ), - ).resolves.toBe('steal'); + ).resolves.toBe('rebuild'); + expect(claimTurnOwnership).toHaveBeenCalledWith({ + session_id: SESSION_ID, + turn_id: 'remote-steal', + expected_active_executor_id: REMOTE_EXECUTOR, + new_active_executor_id: configuration.EXECUTOR_ID, + }); + }); + + it('retries when the steal CAS loses', async () => { + redisRequestMock.mockRejectedValue(new NoResponderError(REMOTE_EXECUTOR)); + + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns: new ActiveTurnRegistry(), + turn: turnRecord({ turnId: 'remote-steal-lose', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), + redis: REDIS, + claimTurnOwnership: () => Promise.resolve(false), + }), + { sessionId: SESSION_ID, turnId: 'remote-steal-lose' }, + ), + ).resolves.toBe('retry'); }); it('does not steal a running turn when there is no responder', async () => { diff --git a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts index 7e4e678d7..4560d7e92 100644 --- a/packages/trueforge/tests/unit/runtime/activeTurns.test.ts +++ b/packages/trueforge/tests/unit/runtime/activeTurns.test.ts @@ -188,4 +188,67 @@ describe('ActiveTurnRegistry', () => { void value; } }); + + it('withTurnLock serializes the same turn and runs different turns in parallel', async () => { + const registry = new ActiveTurnRegistry(); + const order: string[] = []; + let releaseFirst!: () => void; + const firstHold = new Promise(resolve => { + releaseFirst = resolve; + }); + let firstEntered!: () => void; + const firstInside = new Promise(resolve => { + firstEntered = resolve; + }); + + const first = registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, async () => { + order.push('t1-a'); + firstEntered(); + await firstHold; + order.push('t1-a-done'); + return 'a'; + }); + await firstInside; + + let secondStarted = false; + const second = registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, async () => { + secondStarted = true; + order.push('t1-b'); + return 'b'; + }); + + let otherEntered!: () => void; + const otherInside = new Promise(resolve => { + otherEntered = resolve; + }); + let releaseOther!: () => void; + const otherHold = new Promise(resolve => { + releaseOther = resolve; + }); + const other = registry.withTurnLock({ sessionId: 's1', turnId: 't2' }, async () => { + otherEntered(); + await otherHold; + order.push('t2'); + return 'other'; + }); + + await otherInside; + expect(secondStarted).toBe(false); + + releaseFirst(); + releaseOther(); + await expect(Promise.all([first, second, other])).resolves.toEqual(['a', 'b', 'other']); + expect(order.filter(step => step.startsWith('t1'))).toEqual(['t1-a', 't1-a-done', 't1-b']); + expect(order).toContain('t2'); + }); + + it('withTurnLock releases after a thrown fn so the next waiter runs', async () => { + const registry = new ActiveTurnRegistry(); + await expect( + registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, async () => { + throw new Error('lock boom'); + }), + ).rejects.toThrow(/lock boom/); + await expect(registry.withTurnLock({ sessionId: 's1', turnId: 't1' }, async () => 'ok')).resolves.toBe('ok'); + }); }); From 8ba401c0e847974d6ca15cbd840637b6d348bdff Mon Sep 17 00:00:00 2001 From: Heer Ambavi Date: Wed, 23 Sep 2026 17:06:51 +0530 Subject: [PATCH 3/5] use mutex, use generator --- .../src/agent-session/SessionHandle.ts | 3 +- .../src/agent-session/activeExecutorId.ts | 31 +++++++++++++++++++ .../trueforge-core/src/agent-session/index.ts | 1 + .../store/InMemorySessionStore.ts | 3 +- .../agent-session/activeExecutorId.test.ts | 23 ++++++++++++++ packages/trueforge/package.json | 1 + packages/trueforge/src/apis/peering.ts | 16 +++++++--- packages/trueforge/src/apis/sessions.ts | 6 ++-- .../postgres/session-store/queries/turns.ts | 4 +-- .../db/sqlite/session-store/queries/turns.ts | 4 +-- packages/trueforge/src/runtime/activeTurns.ts | 20 +++++------- pnpm-lock.yaml | 10 ++++++ 12 files changed, 96 insertions(+), 26 deletions(-) create mode 100644 packages/trueforge-core/src/agent-session/activeExecutorId.ts create mode 100644 packages/trueforge-core/tests/agent-session/activeExecutorId.test.ts diff --git a/packages/trueforge-core/src/agent-session/SessionHandle.ts b/packages/trueforge-core/src/agent-session/SessionHandle.ts index ecaa2aeaa..099cfa4f1 100644 --- a/packages/trueforge-core/src/agent-session/SessionHandle.ts +++ b/packages/trueforge-core/src/agent-session/SessionHandle.ts @@ -18,6 +18,7 @@ import { import type { CreateDynamicSubAgentThread } from '../core/runtime/CreateDynamicSubAgentThread'; import type { Sandbox } from '../core/sandbox/Sandbox'; import type { AgentTracing } from '../core/tracing/AgentTracing'; +import { mintActiveExecutorId } from './activeExecutorId'; import { builtinsFromSpec } from './builtinsFromSpec'; import type { ITurnResourceResolver, ResolvedAgentDefinition } from './ITurnResourceResolver'; import type { SessionRecord } from './models/SessionRecord'; @@ -285,7 +286,7 @@ export class SessionHandle< first_turn_id: previous?.first_turn_id ?? turnId, ancestor_ids: previous ? [...previous.ancestor_ids, previous.turn_id].slice(-MAX_TURN_ANCESTORS) : [], previous_turn_id: previousTurnId, - active_executor_id: input.active_executor_id, + active_executor_id: mintActiveExecutorId(input.active_executor_id), state: { status: 'running' }, input: input.input ?? [], created_at: now, diff --git a/packages/trueforge-core/src/agent-session/activeExecutorId.ts b/packages/trueforge-core/src/agent-session/activeExecutorId.ts new file mode 100644 index 000000000..d86fedadc --- /dev/null +++ b/packages/trueforge-core/src/agent-session/activeExecutorId.ts @@ -0,0 +1,31 @@ +/** + * `active_executor_id` is `{executorId}.{generation}`. Generation is 4 hex + * chars from `randomBytes` so a stale in-memory run cannot persist after + * steal/rebuild. Rows minted before this grammar parse as executor-only. + */ +import { randomBytes } from 'node:crypto'; + +const GENERATION = /^[0-9a-f]{4}$/; + +export function parseActiveExecutorId(value: string): { executorId: string; generation: string | undefined } { + const lastDot = value.lastIndexOf('.'); + if (lastDot <= 0) { + return { executorId: value, generation: undefined }; + } + const generation = value.slice(lastDot + 1); + if (!GENERATION.test(generation)) { + return { executorId: value, generation: undefined }; + } + return { executorId: value.slice(0, lastDot), generation }; +} + +/** `{executorId}.{4 hex chars}`, never equal to `previous`. */ +export function mintActiveExecutorId(executorId: string, previous?: string): string { + const id = parseActiveExecutorId(executorId).executorId; + for (;;) { + const next = `${id}.${randomBytes(2).toString('hex')}`; + if (next !== previous) { + return next; + } + } +} diff --git a/packages/trueforge-core/src/agent-session/index.ts b/packages/trueforge-core/src/agent-session/index.ts index 8cdc0a9d3..cde2e57fd 100644 --- a/packages/trueforge-core/src/agent-session/index.ts +++ b/packages/trueforge-core/src/agent-session/index.ts @@ -71,6 +71,7 @@ export type { export { TokenPaginationSchema } from './schemas/pagination'; export type { TokenPagination } from './schemas/pagination'; +export { mintActiveExecutorId, parseActiveExecutorId } from './activeExecutorId'; export type { SessionRecord } from './models/SessionRecord'; export { MAIN_THREAD_ID } from './models/TurnRecord'; export type { TurnRecord, TurnSnapshot } from './models/TurnRecord'; diff --git a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts index 0e4a23761..9188b0031 100644 --- a/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts +++ b/packages/trueforge-core/src/agent-session/store/InMemorySessionStore.ts @@ -1,5 +1,6 @@ import type { AgentThreadSnapshot } from '../../core/runtime/AgentThread.types'; import { getEmptyCurrentContextUsage } from '../../core/runtime/contextUsage'; +import { mintActiveExecutorId } from '../activeExecutorId'; import type { SessionRecord } from '../models/SessionRecord'; import type { TurnRecord, TurnSnapshot } from '../models/TurnRecord'; import type { PersistedTurnEvent, SessionEventItem } from '../schemas/events'; @@ -496,7 +497,7 @@ export class InMemorySessionStore< if (turn.state.status !== 'paused' || turn.active_executor_id !== input.expected_active_executor_id) { return false; } - turn.active_executor_id = input.new_active_executor_id; + turn.active_executor_id = mintActiveExecutorId(input.new_active_executor_id, turn.active_executor_id); turn.updated_at = new Date(); return true; } diff --git a/packages/trueforge-core/tests/agent-session/activeExecutorId.test.ts b/packages/trueforge-core/tests/agent-session/activeExecutorId.test.ts new file mode 100644 index 000000000..bfc750415 --- /dev/null +++ b/packages/trueforge-core/tests/agent-session/activeExecutorId.test.ts @@ -0,0 +1,23 @@ +import { mintActiveExecutorId, parseActiveExecutorId } from '../../src/agent-session/activeExecutorId'; + +describe('parseActiveExecutorId', () => { + it('splits executorId.generation', () => { + expect(parseActiveExecutorId('abc123.a1f0')).toEqual({ executorId: 'abc123', generation: 'a1f0' }); + }); + + it('treats a bare id as executor-only', () => { + expect(parseActiveExecutorId('abc123')).toEqual({ executorId: 'abc123', generation: undefined }); + }); + + it('does not treat a short suffix as generation', () => { + expect(parseActiveExecutorId('abc123.12')).toEqual({ executorId: 'abc123.12', generation: undefined }); + }); +}); + +describe('mintActiveExecutorId', () => { + it('appends 4 hex chars and strips an existing suffix', () => { + const minted = mintActiveExecutorId('abc123.a1f0', 'abc123.a1f0'); + expect(minted).toMatch(/^abc123\.[0-9a-f]{4}$/); + expect(minted).not.toBe('abc123.a1f0'); + }); +}); diff --git a/packages/trueforge/package.json b/packages/trueforge/package.json index a2c464801..217e56db4 100644 --- a/packages/trueforge/package.json +++ b/packages/trueforge/package.json @@ -69,6 +69,7 @@ "@sentry/node": "^10.74.0", "@truefoundry/trueforge-core": "workspace:*", "@truefoundry/trueforge-sdk": "workspace:*", + "async-mutex": "^0.5.0", "better-sqlite3": "^13.0.3", "cron-parser": "^5.4.0", "env-paths": "^4.0.0", diff --git a/packages/trueforge/src/apis/peering.ts b/packages/trueforge/src/apis/peering.ts index f3b15cb29..d2c8b10df 100644 --- a/packages/trueforge/src/apis/peering.ts +++ b/packages/trueforge/src/apis/peering.ts @@ -3,7 +3,11 @@ * and turn-ownership resolution used by send/subscribe/cancel. */ import type { ISessionStore, TurnState } from '@truefoundry/trueforge-core/agent-session'; -import { CancellationReason, TurnNotFoundError } from '@truefoundry/trueforge-core/agent-session'; +import { + CancellationReason, + parseActiveExecutorId, + TurnNotFoundError, +} from '@truefoundry/trueforge-core/agent-session'; import { NoResponderError, redisRequest, @@ -118,13 +122,14 @@ export async function callPeer(input: { /** * Load the turn under the per-turn lock, peer if another replica owns it, then - * {@link resolveOwnershipAction}. A `steal` is claimed here (R5): winner → - * `rebuild`, loser → `retry`. + * {@link resolveOwnershipAction}. A `steal` is claimed here. */ export async function resolveTurnOwnership( deps: ResolveTurnOwnershipDeps, input: { sessionId: string; turnId: string }, ): Promise { + // One request at a time per turn so two handlers don't both steal or rebuild + // from a stale read. Re-read the row inside the lock. return deps.activeTurns.withTurnLock({ sessionId: input.sessionId, turnId: input.turnId }, async () => { const turn = await deps.sessionStore.getTurn({ session_id: input.sessionId, @@ -135,12 +140,13 @@ export async function resolveTurnOwnership( } const owner = turn.active_executor_id; - const ownerIsLocal = owner === configuration.EXECUTOR_ID; + const ownerExecutorId = parseActiveExecutorId(owner).executorId; + const ownerIsLocal = ownerExecutorId === configuration.EXECUTOR_ID; const peerResult = !ownerIsLocal && deps.redis ? await callPeer({ redis: deps.redis, - executorId: owner, + executorId: ownerExecutorId, path: TURNS_LOCATE_PATH, body: { session_id: input.sessionId, turn_id: input.turnId }, }) diff --git a/packages/trueforge/src/apis/sessions.ts b/packages/trueforge/src/apis/sessions.ts index 5a4633232..cc41aa434 100644 --- a/packages/trueforge/src/apis/sessions.ts +++ b/packages/trueforge/src/apis/sessions.ts @@ -5,6 +5,7 @@ import { OpenAPIHono, type RouteHandler } from '@hono/zod-openapi'; import type { ISessionStore, SessionHandle, SessionRecord, Sessions } from '@truefoundry/trueforge-core/agent-session'; import { CancellationReason, + parseActiveExecutorId, SessionStoreConflictError, SessionStoreInvariantError, SessionStoreNotFoundError, @@ -110,7 +111,8 @@ export async function cancelSessionTurn( } const owner = turn.active_executor_id; - const ownerIsLocal = owner === configuration.EXECUTOR_ID; + const ownerExecutorId = parseActiveExecutorId(owner).executorId; + const ownerIsLocal = ownerExecutorId === configuration.EXECUTOR_ID; if (ownerIsLocal && deps.activeTurns.has({ sessionId, turnId })) { deps.activeTurns.cancel({ sessionId, turnId, abortReason: reason }); @@ -120,7 +122,7 @@ export async function cancelSessionTurn( if (!ownerIsLocal && deps.redis) { const peerResult = await callPeer({ redis: deps.redis, - executorId: owner, + executorId: ownerExecutorId, path: SESSIONS_CANCEL_PATH, body: { session_id: sessionId, turn_id: turnId, reason }, }); diff --git a/packages/trueforge/src/db/postgres/session-store/queries/turns.ts b/packages/trueforge/src/db/postgres/session-store/queries/turns.ts index 0233a3d7e..80f599e72 100644 --- a/packages/trueforge/src/db/postgres/session-store/queries/turns.ts +++ b/packages/trueforge/src/db/postgres/session-store/queries/turns.ts @@ -1,4 +1,4 @@ -import type { SessionMetrics } from '@truefoundry/trueforge-core/agent-session'; +import { mintActiveExecutorId, type SessionMetrics } from '@truefoundry/trueforge-core/agent-session'; import type { TurnRecord, TurnSnapshot } from '@truefoundry/trueforge-core/agent-session/models/TurnRecord'; import { type TerminalTurnState, @@ -806,7 +806,7 @@ export async function claimTurnOwnership(db: Kysely, input: ClaimTurnO const result = await db .updateTable('turn') .set({ - active_executor_id: input.new_active_executor_id, + active_executor_id: mintActiveExecutorId(input.new_active_executor_id, input.expected_active_executor_id), updated_at: sql`now()`, }) .where('session_id', '=', input.session_id) diff --git a/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts b/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts index 443d558be..f28344850 100644 --- a/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts +++ b/packages/trueforge/src/db/sqlite/session-store/queries/turns.ts @@ -1,4 +1,4 @@ -import type { SessionMetrics } from '@truefoundry/trueforge-core/agent-session'; +import { mintActiveExecutorId, type SessionMetrics } from '@truefoundry/trueforge-core/agent-session'; import type { TurnRecord, TurnSnapshot } from '@truefoundry/trueforge-core/agent-session/models/TurnRecord'; import { type TerminalTurnState, @@ -854,7 +854,7 @@ export async function claimTurnOwnership(db: Kysely, input: ClaimTurnO const result = await db .updateTable('turn') .set({ - active_executor_id: input.new_active_executor_id, + active_executor_id: mintActiveExecutorId(input.new_active_executor_id, input.expected_active_executor_id), updated_at: nowIso(), }) .where('session_id', '=', input.session_id) diff --git a/packages/trueforge/src/runtime/activeTurns.ts b/packages/trueforge/src/runtime/activeTurns.ts index c0be08344..c72a1b5de 100644 --- a/packages/trueforge/src/runtime/activeTurns.ts +++ b/packages/trueforge/src/runtime/activeTurns.ts @@ -5,6 +5,7 @@ * `track()` owns registration and cleanup around the stream lifecycle. */ import { CancellationReason } from '@truefoundry/trueforge-core/agent-session'; +import { Mutex } from 'async-mutex'; interface ActiveTurnRun { abortController: AbortController; @@ -18,28 +19,21 @@ function activeTurnKey(sessionId: string, turnId: string): string { export class ActiveTurnRegistry { private readonly runs = new Map(); - private readonly locks = new Map>(); + private readonly locks = new Map(); private alreadyShutDownAbortReason: CancellationReason | undefined; /** * Serialize work for one turn in this process. Waiters queue; different keys - * run in parallel. Re-read registry / DB inside `fn` — do not trust values - * from before the lock. + * run in parallel. */ async withTurnLock(input: { sessionId: string; turnId: string }, fn: () => Promise): Promise { const key = activeTurnKey(input.sessionId, input.turnId); - const previous = this.locks.get(key) ?? Promise.resolve(); - let release!: () => void; - const held = new Promise(resolve => { - release = resolve; - }); - this.locks.set(key, held); - await previous; + const mutex = this.locks.get(key) ?? new Mutex(); + this.locks.set(key, mutex); try { - return await fn(); + return await mutex.runExclusive(fn); } finally { - release(); - if (this.locks.get(key) === held) { + if (!mutex.isLocked() && this.locks.get(key) === mutex) { this.locks.delete(key); } } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b3b0a135a..59f3dd2d1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -187,6 +187,9 @@ importers: '@truefoundry/trueforge-sdk': specifier: workspace:* version: link:../trueforge-sdk + async-mutex: + specifier: ^0.5.0 + version: 0.5.0 better-sqlite3: specifier: ^13.0.3 version: 13.0.3 @@ -4212,6 +4215,9 @@ packages: ast-v8-to-istanbul@1.0.5: resolution: {integrity: sha512-UPAgKJFSEGMWSDr3LX4tqnAb4f7KGT8O40Tyx8wbYmmZ/yn58lNCm8h3svs3eXgiGd5AXxz8NDOvXWvicq+rJA==} + async-mutex@0.5.0: + resolution: {integrity: sha512-1A94B18jkJ3DYq284ohPxoXbfTA5HsQ7/Mf4DEhcyLx3Bz27Rh59iScbB6EPiP+B+joue6YCxcMXSbFC1tZKwA==} + async@3.2.6: resolution: {integrity: sha512-htCUDlxyyCLMgaM3xXg0C0LW2xqfuQ6p05pCEIsXuyQ+a1koYKTuBMzRNwmybfLgvJDMd0r1LTn4+E0Ti6C2AA==} @@ -11618,6 +11624,10 @@ snapshots: estree-walker: 3.0.3 js-tokens: 10.0.0 + async-mutex@0.5.0: + dependencies: + tslib: 2.8.1 + async@3.2.6: {} asynckit@0.4.0: {} From cbee34a7e81fb9d4ea00c427ffb5b07af0faf3ac Mon Sep 17 00:00:00 2001 From: Heer Ambavi Date: Wed, 23 Sep 2026 17:30:27 +0530 Subject: [PATCH 4/5] wip --- packages/trueforge/src/apis/peering.ts | 139 ++++++++++++------ .../trueforge/tests/unit/apis/peering.test.ts | 109 ++++++++++---- 2 files changed, 178 insertions(+), 70 deletions(-) diff --git a/packages/trueforge/src/apis/peering.ts b/packages/trueforge/src/apis/peering.ts index d2c8b10df..bbefac4e1 100644 --- a/packages/trueforge/src/apis/peering.ts +++ b/packages/trueforge/src/apis/peering.ts @@ -51,12 +51,28 @@ export type PeerResult = 'ok' | 'no_responder' | 'failed'; * - `run` — execute here (we own it and have an ActiveTurn). * - `rebuild` — rebuild ActiveTurn here (we own a paused turn with no run). * - `forward` — send the work to the remote owner. - * - `steal` — table-only: claim the paused turn (resolveTurnOwnership then rebuilds or retries). + * - `steal` — claim the paused turn (CAS). Winner must rebuild under this lock (not yet). * - `reject` — do not continue (turn terminal, or local running with no ActiveTurn). * - `retry` — do not continue now (remote owner not usable; caller may try again). */ export type OwnershipAction = 'run' | 'rebuild' | 'forward' | 'steal' | 'reject' | 'retry'; +export class OwnershipRejectedError extends Error { + readonly action = 'reject' as const; + constructor(message = 'Turn cannot continue on this replica') { + super(message); + this.name = 'OwnershipRejectedError'; + } +} + +export class OwnershipRetryError extends Error { + readonly action = 'retry' as const; + constructor(message = 'This turn is temporarily unavailable. Please try again.') { + super(message); + this.name = 'OwnershipRetryError'; + } +} + export interface ResolveTurnOwnershipDeps { activeTurns: Pick; sessionStore: Pick; @@ -121,55 +137,90 @@ export async function callPeer(input: { } /** - * Load the turn under the per-turn lock, peer if another replica owns it, then - * {@link resolveOwnershipAction}. A `steal` is claimed here. + * Decide ownership under the per-turn lock (re-read + locate). Steal CAS + * runs here. `true` — run here. `false` — request was forwarded. Throws + * {@link OwnershipRejectedError} / {@link OwnershipRetryError}. + * `rebuild` is not handled yet (must run under this same lock later). */ export async function resolveTurnOwnership( deps: ResolveTurnOwnershipDeps, - input: { sessionId: string; turnId: string }, -): Promise { - // One request at a time per turn so two handlers don't both steal or rebuild - // from a stale read. Re-read the row inside the lock. - return deps.activeTurns.withTurnLock({ sessionId: input.sessionId, turnId: input.turnId }, async () => { - const turn = await deps.sessionStore.getTurn({ - session_id: input.sessionId, - turn_id: input.turnId, - }); - if (!turn) { - throw new TurnNotFoundError(input.turnId); - } + input: { + sessionId: string; + turnId: string; + /** Sent to the owning replica when the table says `forward`. */ + forward: { path: string; body: JSONValue }; + }, +): Promise { + const decided = await deps.activeTurns.withTurnLock( + { sessionId: input.sessionId, turnId: input.turnId }, + async () => { + const turn = await deps.sessionStore.getTurn({ + session_id: input.sessionId, + turn_id: input.turnId, + }); + if (!turn) { + throw new TurnNotFoundError(input.turnId); + } - const owner = turn.active_executor_id; - const ownerExecutorId = parseActiveExecutorId(owner).executorId; - const ownerIsLocal = ownerExecutorId === configuration.EXECUTOR_ID; - const peerResult = - !ownerIsLocal && deps.redis - ? await callPeer({ - redis: deps.redis, - executorId: ownerExecutorId, - path: TURNS_LOCATE_PATH, - body: { session_id: input.sessionId, turn_id: input.turnId }, - }) - : undefined; - - const action = resolveOwnershipAction({ - status: turn.state.status, - ownerIsLocal, - hasActiveTurn: deps.activeTurns.has({ sessionId: input.sessionId, turnId: input.turnId }), - ...(peerResult === undefined ? {} : { peerResult }), - }); - if (action !== 'steal') { - return action; - } + const owner = turn.active_executor_id; + const ownerExecutorId = parseActiveExecutorId(owner).executorId; + const ownerIsLocal = ownerExecutorId === configuration.EXECUTOR_ID; + const peerResult = + !ownerIsLocal && deps.redis + ? await callPeer({ + redis: deps.redis, + executorId: ownerExecutorId, + path: TURNS_LOCATE_PATH, + body: { session_id: input.sessionId, turn_id: input.turnId }, + }) + : undefined; - const won = await deps.sessionStore.claimTurnOwnership({ - session_id: input.sessionId, - turn_id: input.turnId, - expected_active_executor_id: owner, - new_active_executor_id: configuration.EXECUTOR_ID, - }); - return won ? 'rebuild' : 'retry'; + const action = resolveOwnershipAction({ + status: turn.state.status, + ownerIsLocal, + hasActiveTurn: deps.activeTurns.has({ sessionId: input.sessionId, turnId: input.turnId }), + ...(peerResult === undefined ? {} : { peerResult }), + }); + switch (action) { + case 'run': + case 'rebuild': + // TODO: Implement run and rebuild. + return undefined; + case 'forward': + return ownerExecutorId; + case 'steal': + await deps.sessionStore.claimTurnOwnership({ + session_id: input.sessionId, + turn_id: input.turnId, + expected_active_executor_id: owner, + new_active_executor_id: configuration.EXECUTOR_ID, + }); + // Winner still needs rebuild under this lock; not implemented yet. + throw new OwnershipRetryError(); + case 'retry': + throw new OwnershipRetryError(); + case 'reject': + throw new OwnershipRejectedError(); + } + }, + ); + + if (decided === undefined) { + return true; + } + if (!deps.redis) { + throw new OwnershipRetryError(); + } + const forwarded = await callPeer({ + redis: deps.redis, + executorId: decided, + path: input.forward.path, + body: input.forward.body, }); + if (forwarded !== 'ok') { + throw new OwnershipRetryError(); + } + return false; } /** diff --git a/packages/trueforge/tests/unit/apis/peering.test.ts b/packages/trueforge/tests/unit/apis/peering.test.ts index 4699d189c..38da2ef6e 100644 --- a/packages/trueforge/tests/unit/apis/peering.test.ts +++ b/packages/trueforge/tests/unit/apis/peering.test.ts @@ -3,6 +3,8 @@ import { TurnNotFoundError } from '@truefoundry/trueforge-core/agent-session'; import { NoResponderError, redisRequest, RequestTimeoutError } from '@truefoundry/trueforge-core/request-reply'; import type { RedisClientType } from 'redis'; import { + OwnershipRejectedError, + OwnershipRetryError, resolveOwnershipAction, resolveTurnOwnership, TURNS_LOCATE_PATH, @@ -27,8 +29,21 @@ const redisRequestMock = jest.mocked(redisRequest); const SESSION_ID = 's1'; const REDIS = {} as RedisClientType; const REMOTE_EXECUTOR = 'other1'; +const SEND_PATH = 'turns/send'; const pausedState = { status: 'paused' as const, action_required_on_events: [] }; +function ownershipInput(turnId: string): { + sessionId: string; + turnId: string; + forward: { path: string; body: { session_id: string; turn_id: string } }; +} { + return { + sessionId: SESSION_ID, + turnId, + forward: { path: SEND_PATH, body: { session_id: SESSION_ID, turn_id: turnId } }, + }; +} + function silentLogger(): { warn: jest.Mock } { return { warn: jest.fn() }; } @@ -102,10 +117,10 @@ describe('resolveTurnOwnership', () => { it('throws when the turn is missing', async () => { await expect( - resolveTurnOwnership(ownershipDeps({ activeTurns: new ActiveTurnRegistry(), turn: undefined }), { - sessionId: SESSION_ID, - turnId: 'missing', - }), + resolveTurnOwnership( + ownershipDeps({ activeTurns: new ActiveTurnRegistry(), turn: undefined }), + ownershipInput('missing'), + ), ).rejects.toBeInstanceOf(TurnNotFoundError); }); @@ -120,24 +135,36 @@ describe('resolveTurnOwnership', () => { activeTurns, turn: turnRecord({ turnId, state: { status: 'running' } }), }), - { sessionId: SESSION_ID, turnId }, + ownershipInput(turnId), + ), + ).resolves.toBe(true); + }); + + it('rejects a local running turn with no ActiveTurn', async () => { + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns: new ActiveTurnRegistry(), + turn: turnRecord({ turnId: 'gone-local', state: { status: 'running' } }), + }), + ownershipInput('gone-local'), ), - ).resolves.toBe('run'); + ).rejects.toBeInstanceOf(OwnershipRejectedError); }); - it('rebuilds when this executor owns a paused turn with no ActiveTurn', async () => { + it('retries when a local rebuild would be needed (not handled yet)', async () => { await expect( resolveTurnOwnership( ownershipDeps({ activeTurns: new ActiveTurnRegistry(), turn: turnRecord({ turnId: 'paused-local', state: pausedState }), }), - { sessionId: SESSION_ID, turnId: 'paused-local' }, + ownershipInput('paused-local'), ), - ).resolves.toBe('rebuild'); + ).rejects.toBeInstanceOf(OwnershipRetryError); }); - it('returns forward when the owning replica still has the turn', async () => { + it('forwards the request when the owning replica still has the turn', async () => { redisRequestMock.mockResolvedValue({ status: 200, body: {} }); await expect( @@ -147,15 +174,37 @@ describe('resolveTurnOwnership', () => { turn: turnRecord({ turnId: 'remote-ok', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), redis: REDIS, }), - { sessionId: SESSION_ID, turnId: 'remote-ok' }, + ownershipInput('remote-ok'), ), - ).resolves.toBe('forward'); - expect(redisRequestMock).toHaveBeenCalledWith( + ).resolves.toBe(false); + expect(redisRequestMock).toHaveBeenNthCalledWith( + 1, expect.objectContaining({ executorId: REMOTE_EXECUTOR, path: TURNS_LOCATE_PATH }), ); + expect(redisRequestMock).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ executorId: REMOTE_EXECUTOR, path: SEND_PATH }), + ); + }); + + it('retries when the forwarded request fails', async () => { + redisRequestMock + .mockResolvedValueOnce({ status: 200, body: {} }) + .mockResolvedValueOnce({ status: 412, body: { message: 'Turn is not on this executor' } }); + + await expect( + resolveTurnOwnership( + ownershipDeps({ + activeTurns: new ActiveTurnRegistry(), + turn: turnRecord({ turnId: 'remote-forward-fail', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), + redis: REDIS, + }), + ownershipInput('remote-forward-fail'), + ), + ).rejects.toBeInstanceOf(OwnershipRetryError); }); - it('does not steal on 412 / timeout (unavailable)', async () => { + it('retries when locate is 412 / timeout (unavailable)', async () => { redisRequestMock.mockResolvedValue({ status: 412, body: { message: 'Turn is not on this executor' } }); await expect( @@ -165,12 +214,12 @@ describe('resolveTurnOwnership', () => { turn: turnRecord({ turnId: 'remote-412', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), redis: REDIS, }), - { sessionId: SESSION_ID, turnId: 'remote-412' }, + ownershipInput('remote-412'), ), - ).resolves.toBe('retry'); + ).rejects.toBeInstanceOf(OwnershipRetryError); }); - it('claims a paused turn with no responder and rebuilds on a winning CAS', async () => { + it('claims a paused turn with no responder, then retries (rebuild not handled yet)', async () => { redisRequestMock.mockRejectedValue(new NoResponderError(REMOTE_EXECUTOR)); const claimTurnOwnership = jest.fn().mockResolvedValue(true); @@ -182,9 +231,9 @@ describe('resolveTurnOwnership', () => { redis: REDIS, claimTurnOwnership, }), - { sessionId: SESSION_ID, turnId: 'remote-steal' }, + ownershipInput('remote-steal'), ), - ).resolves.toBe('rebuild'); + ).rejects.toBeInstanceOf(OwnershipRetryError); expect(claimTurnOwnership).toHaveBeenCalledWith({ session_id: SESSION_ID, turn_id: 'remote-steal', @@ -195,6 +244,7 @@ describe('resolveTurnOwnership', () => { it('retries when the steal CAS loses', async () => { redisRequestMock.mockRejectedValue(new NoResponderError(REMOTE_EXECUTOR)); + const claimTurnOwnership = jest.fn().mockResolvedValue(false); await expect( resolveTurnOwnership( @@ -202,15 +252,17 @@ describe('resolveTurnOwnership', () => { activeTurns: new ActiveTurnRegistry(), turn: turnRecord({ turnId: 'remote-steal-lose', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), redis: REDIS, - claimTurnOwnership: () => Promise.resolve(false), + claimTurnOwnership, }), - { sessionId: SESSION_ID, turnId: 'remote-steal-lose' }, + ownershipInput('remote-steal-lose'), ), - ).resolves.toBe('retry'); + ).rejects.toBeInstanceOf(OwnershipRetryError); + expect(claimTurnOwnership).toHaveBeenCalled(); }); it('does not steal a running turn when there is no responder', async () => { redisRequestMock.mockRejectedValue(new NoResponderError(REMOTE_EXECUTOR)); + const claimTurnOwnership = jest.fn(); await expect( resolveTurnOwnership( @@ -222,14 +274,17 @@ describe('resolveTurnOwnership', () => { activeExecutorId: REMOTE_EXECUTOR, }), redis: REDIS, + claimTurnOwnership, }), - { sessionId: SESSION_ID, turnId: 'remote-running' }, + ownershipInput('remote-running'), ), - ).resolves.toBe('retry'); + ).rejects.toBeInstanceOf(OwnershipRetryError); + expect(claimTurnOwnership).not.toHaveBeenCalled(); }); it('does not steal on peer timeout', async () => { redisRequestMock.mockRejectedValue(new RequestTimeoutError(60_000)); + const claimTurnOwnership = jest.fn(); await expect( resolveTurnOwnership( @@ -237,10 +292,12 @@ describe('resolveTurnOwnership', () => { activeTurns: new ActiveTurnRegistry(), turn: turnRecord({ turnId: 'remote-timeout', state: pausedState, activeExecutorId: REMOTE_EXECUTOR }), redis: REDIS, + claimTurnOwnership, }), - { sessionId: SESSION_ID, turnId: 'remote-timeout' }, + ownershipInput('remote-timeout'), ), - ).resolves.toBe('retry'); + ).rejects.toBeInstanceOf(OwnershipRetryError); + expect(claimTurnOwnership).not.toHaveBeenCalled(); }); }); From 0a1628475f2f7517953d69c1fb55be350b324afc Mon Sep 17 00:00:00 2001 From: Heer Ambavi Date: Thu, 24 Sep 2026 10:39:21 +0530 Subject: [PATCH 5/5] wip --- packages/trueforge/src/apis/peering.ts | 3 --- 1 file changed, 3 deletions(-) diff --git a/packages/trueforge/src/apis/peering.ts b/packages/trueforge/src/apis/peering.ts index bbefac4e1..880408e47 100644 --- a/packages/trueforge/src/apis/peering.ts +++ b/packages/trueforge/src/apis/peering.ts @@ -91,9 +91,6 @@ export function resolveOwnershipAction(input: { } if (input.ownerIsLocal) { - if (input.peerResult === 'ok') { - throw new Error('peerResult "ok" is only valid for a remote owner'); - } if (input.hasActiveTurn) { return 'run'; }