From c6881c331abb837c7509fdf52fcc2025e01faa41 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Tue, 1 Sep 2026 01:35:25 -0600 Subject: [PATCH 1/4] feat(prime): adopt active turns after restart Closes #84 --- .../orchestration/Layers/CheckpointReactor.ts | 19 + .../Layers/ProviderRuntimeIngestion.ts | 22 + apps/server/src/persistence/Migrations.ts | 2 + .../050_PrimeAgentRecoveryLedger.ts | 48 ++ .../src/provider/Drivers/PrimeAgentDriver.ts | 38 +- .../provider/Layers/ProviderService.test.ts | 58 +- .../src/provider/Layers/ProviderService.ts | 152 +++- .../src/provider/Services/ProviderAdapter.ts | 20 + .../src/provider/Services/ProviderService.ts | 3 + .../prime/PrimeAgentBackendSelection.ts | 4 + .../provider/prime/PrimeAgentDaemonAdapter.ts | 786 ++++++++++++++++-- .../provider/prime/PrimeAgentDaemonBridge.ts | 141 +++- .../prime/PrimeAgentDaemonManager.test.ts | 61 +- .../provider/prime/PrimeAgentDaemonManager.ts | 81 +- .../PrimeAgentDaemonSessionRuntime.test.ts | 213 +++++ .../prime/PrimeAgentDaemonSessionRuntime.ts | 605 +++++++++++--- .../prime/PrimeAgentRecoveryLedger.test.ts | 139 ++++ .../prime/PrimeAgentRecoveryLedger.ts | 478 +++++++++++ .../PrimeAgentRestartAdoption.real.test.mjs | 194 +++++ .../prime/PrimeAgentRestartReplay.test.ts | 52 ++ apps/server/src/server.ts | 5 +- .../serverRuntimeStartup.reconcile.test.ts | 48 +- apps/server/src/serverRuntimeStartup.ts | 5 + docs/internals/prime-agent-daemon-parity.md | 36 +- docs/internals/prime-agent-native-parity.md | 2 +- docs/user/providers-prime-agent.md | 12 +- 26 files changed, 3014 insertions(+), 210 deletions(-) create mode 100644 apps/server/src/persistence/Migrations/050_PrimeAgentRecoveryLedger.ts create mode 100644 apps/server/src/provider/prime/PrimeAgentRecoveryLedger.test.ts create mode 100644 apps/server/src/provider/prime/PrimeAgentRecoveryLedger.ts create mode 100644 apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs create mode 100644 apps/server/src/provider/prime/PrimeAgentRestartReplay.test.ts diff --git a/apps/server/src/orchestration/Layers/CheckpointReactor.ts b/apps/server/src/orchestration/Layers/CheckpointReactor.ts index 0f40c2286..6a1916f51 100644 --- a/apps/server/src/orchestration/Layers/CheckpointReactor.ts +++ b/apps/server/src/orchestration/Layers/CheckpointReactor.ts @@ -29,6 +29,7 @@ import { } from "../../checkpointing/Utils.ts"; import * as CheckpointStore from "../../checkpointing/CheckpointStore.ts"; import { ProviderService } from "../../provider/Services/ProviderService.ts"; +import { PrimeAgentRecoveryLedger } from "../../provider/prime/PrimeAgentRecoveryLedger.ts"; import { CheckpointReactor, type CheckpointReactorShape } from "../Services/CheckpointReactor.ts"; import { forkParked } from "../../serverActivation.ts"; import { OrchestrationEngineService } from "../Services/OrchestrationEngine.ts"; @@ -85,6 +86,9 @@ const make = Effect.gen(function* () { const orchestrationEngine = yield* OrchestrationEngineService; const projectionSnapshotQuery = yield* ProjectionSnapshotQuery; const providerService = yield* ProviderService; + const recoveryLedger = Option.getOrUndefined( + yield* Effect.serviceOption(PrimeAgentRecoveryLedger), + ); const checkpointStore = yield* CheckpointStore.CheckpointStore; const receiptBus = yield* RuntimeReceiptBus; const workspaceEntries = yield* WorkspaceEntries.WorkspaceEntries; @@ -823,6 +827,21 @@ const make = Effect.gen(function* () { ), ), ); + if (recoveryLedger !== undefined) { + yield* Effect.gen(function* () { + yield* recoveryLedger.markCheckpointQuiesced({ + threadId: event.threadId, + updatedAt: yield* nowIso, + }); + yield* recoveryLedger.deleteIfSettled(event.threadId); + }).pipe( + Effect.catchCause(() => + Effect.logWarning("failed to settle Prime Agent recovery checkpoint proof", { + threadId: event.threadId, + }), + ), + ); + } return; } }); diff --git a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts index 9b3d425ce..c6cb285c5 100644 --- a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts +++ b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts @@ -37,6 +37,7 @@ import { makeDrainableWorker } from "@t3tools/shared/DrainableWorker"; import { ProviderRegistry } from "../../provider/Services/ProviderRegistry.ts"; import { ProviderService } from "../../provider/Services/ProviderService.ts"; +import { PrimeAgentRecoveryLedger } from "../../provider/prime/PrimeAgentRecoveryLedger.ts"; import { rateLimitFromRuntimeEventPayload, usageWindowsFromRuntimeEventPayload, @@ -1507,8 +1508,23 @@ const make = Effect.gen(function* () { const orchestrationEngine = yield* OrchestrationEngineService; const projectionSnapshotQuery = yield* ProjectionSnapshotQuery; const providerService = yield* ProviderService; + const recoveryLedger = Option.getOrUndefined( + yield* Effect.serviceOption(PrimeAgentRecoveryLedger), + ); const providerRegistry = yield* ProviderRegistry; const projectionTurnRepository = yield* ProjectionTurnRepository; + const settleRecoveryTerminalProjection = (threadId: ThreadId, updatedAt: string) => + recoveryLedger === undefined + ? Effect.void + : recoveryLedger.markTerminalProjected({ threadId, updatedAt }).pipe( + Effect.andThen(recoveryLedger.deleteIfSettled(threadId)), + Effect.catchCause(() => + Effect.logWarning("failed to settle Prime Agent terminal projection proof", { + threadId, + }), + ), + Effect.asVoid, + ); const serverSettingsService = yield* ServerSettingsService; const providerCommandId = (event: ProviderRuntimeEvent, tag: string) => crypto.randomUUIDv4.pipe( @@ -2228,6 +2244,9 @@ const make = Effect.gen(function* () { if ((cleared.eventCount ?? 0) > 0 && stoppedLineageStillProjected) { yield* clearTurnStateForSession(thread.id); } + if ((cleared.eventCount ?? 0) > 0) { + yield* settleRecoveryTerminalProjection(thread.id, event.createdAt); + } return; } const eventMatchesStoppedSession = @@ -2566,6 +2585,9 @@ const make = Effect.gen(function* () { "thread-session-set", ); if ((applied.eventCount ?? 0) === 0) return; + if (event.type === "session.exited") { + yield* settleRecoveryTerminalProjection(thread.id, event.createdAt); + } } if (event.type === "turn.started" && acceptedTurnStartedSourcePlan !== null) { diff --git a/apps/server/src/persistence/Migrations.ts b/apps/server/src/persistence/Migrations.ts index 54cecd0fe..341e53616 100644 --- a/apps/server/src/persistence/Migrations.ts +++ b/apps/server/src/persistence/Migrations.ts @@ -61,6 +61,7 @@ import Migration0046 from "./Migrations/046_ProjectionThreadLinkedPullRequest.ts import Migration0047 from "./Migrations/047_ProjectionThreadsUnsettledAt.ts"; import Migration0048 from "./Migrations/048_ProjectionThreadSessionPendingTurnRequest.ts"; import Migration0049 from "./Migrations/049_ProjectionThreadSessionPendingStop.ts"; +import Migration0050 from "./Migrations/050_PrimeAgentRecoveryLedger.ts"; /** * Migration loader with all migrations defined inline. * @@ -144,6 +145,7 @@ export const migrationEntries = [ [47, "ProjectionThreadsUnsettledAt", Migration0047], [48, "ProjectionThreadSessionPendingTurnRequest", Migration0048], [49, "ProjectionThreadSessionPendingStop", Migration0049], + [50, "PrimeAgentRecoveryLedger", Migration0050], ] as const; export const migrationManifest = migrationEntries.map(([id, name]) => [id, name] as const); diff --git a/apps/server/src/persistence/Migrations/050_PrimeAgentRecoveryLedger.ts b/apps/server/src/persistence/Migrations/050_PrimeAgentRecoveryLedger.ts new file mode 100644 index 000000000..5b2d7098d --- /dev/null +++ b/apps/server/src/persistence/Migrations/050_PrimeAgentRecoveryLedger.ts @@ -0,0 +1,48 @@ +import * as Effect from "effect/Effect"; +import * as SqlClient from "effect/unstable/sql/SqlClient"; + +export default Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + + yield* sql` + CREATE TABLE IF NOT EXISTS prime_agent_recovery_ledger ( + thread_id TEXT PRIMARY KEY, + provider_instance_id TEXT NOT NULL, + session_incarnation_id TEXT NOT NULL, + admission_request_id TEXT NOT NULL, + turn_id TEXT, + package_root TEXT NOT NULL, + package_version TEXT NOT NULL, + managed_build_id TEXT NOT NULL, + sdk_features_json TEXT NOT NULL, + daemon_capabilities_json TEXT NOT NULL, + protocol_name TEXT NOT NULL, + protocol_version INTEGER NOT NULL, + schema_revision INTEGER NOT NULL, + active_session_id TEXT NOT NULL, + native_session_id TEXT NOT NULL, + recovery_handle TEXT NOT NULL, + supervisor_generation TEXT NOT NULL, + ownership_generation INTEGER NOT NULL, + cursor_generation TEXT NOT NULL, + cursor_sequence INTEGER NOT NULL, + correlation_id TEXT NOT NULL, + mcp_owner_id TEXT NOT NULL, + recovery_config_json TEXT NOT NULL, + launch_environment_json TEXT NOT NULL, + transcript_message_count INTEGER NOT NULL DEFAULT 0, + transcript_fingerprints_json TEXT NOT NULL DEFAULT '[]', + owner_token TEXT NOT NULL, + state TEXT NOT NULL, + native_cleanup_proven INTEGER NOT NULL DEFAULT 0, + terminal_projected INTEGER NOT NULL DEFAULT 0, + checkpoint_quiesced INTEGER NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL + ) + `; + + yield* sql` + CREATE INDEX IF NOT EXISTS idx_prime_agent_recovery_ledger_active + ON prime_agent_recovery_ledger(state, provider_instance_id, updated_at) + `; +}); diff --git a/apps/server/src/provider/Drivers/PrimeAgentDriver.ts b/apps/server/src/provider/Drivers/PrimeAgentDriver.ts index 41d42e4d9..19b3d33bb 100644 --- a/apps/server/src/provider/Drivers/PrimeAgentDriver.ts +++ b/apps/server/src/provider/Drivers/PrimeAgentDriver.ts @@ -4,12 +4,13 @@ import { type ServerProvider, type ServerProviderDistribution, } from "@t3tools/contracts"; -import { HostProcessPlatform } from "@t3tools/shared/hostProcess"; +import { HostProcessArchitecture, HostProcessPlatform } from "@t3tools/shared/hostProcess"; import { resolveCommandPath } from "@t3tools/shared/shell"; import * as Crypto from "effect/Crypto"; import * as Effect from "effect/Effect"; import * as FileSystem from "effect/FileSystem"; import * as Path from "effect/Path"; +import * as Result from "effect/Result"; import * as Schema from "effect/Schema"; import { HttpClient } from "effect/unstable/http"; import { ChildProcessSpawner } from "effect/unstable/process"; @@ -123,6 +124,7 @@ export const PrimeAgentDriver: ProviderDriver Effect.gen(function* () { const hostPlatform = yield* HostProcessPlatform; + const hostArchitecture = yield* HostProcessArchitecture; if (!isPrimeAgentProviderPlatformSupported(hostPlatform)) { return yield* new ProviderDriverError({ driver: DRIVER_KIND, @@ -219,6 +221,37 @@ export const PrimeAgentDriver: ProviderDriver + inspectPrimeAgentDistribution( + { + stateDir: serverConfig.stateDir, + instanceId, + packageRoot: publicPackage.packageRoot, + platform: hostPlatform, + checkedAt: "1970-01-01T00:00:00.000Z", + enableUpdateChecks: false, + }, + { loadLatestVerifiedPublication }, + ), + ); + return { publicPackage, distribution }; + }), + ); + const recoveryManagedBuildId = + Result.isSuccess(recoveryDistribution) && + recoveryDistribution.success.distribution.classification === "pylon-managed" && + recoveryDistribution.success.distribution.buildId !== null + ? recoveryDistribution.success.distribution.buildId + : undefined; + const backend = yield* negotiatePrimeAgentBackend( { enabled: effectiveConfig.enabled, @@ -228,6 +261,8 @@ export const PrimeAgentDriver: ProviderDriver @@ -383,6 +418,7 @@ export const PrimeAgentDriver: ProviderDriver { const directoryLayer = ProviderSessionDirectoryLive.pipe(Layer.provide(runtimeRepositoryLayer)); return makeProviderServiceLive(options).pipe( Layer.provide(providerAdapterLayer), - Layer.provide(directoryLayer), + Layer.provideMerge(directoryLayer), Layer.provide(ServerSettings.ServerSettingsService.layerTest({ enableAgentBrowserAccess })), Layer.provide(serverConfigTestLayer), Layer.provide(AnalyticsService.layerTest), @@ -3602,6 +3602,62 @@ describe("agent browser access", () => { }).pipe(Effect.provide(NodeServices.layer)), ); + it.effect("restores MCP and the directory binding before exposing recovered activity", () => + Effect.gen(function* () { + const threadId = asThreadId("thread-restart-adoption"); + const codex = makeFakeCodexAdapter(); + const order: string[] = []; + const recoveryAdapter: ProviderAdapterShape = { + ...codex.adapter, + recoverSession: (input) => + Effect.gen(function* () { + assert.isDefined(McpProviderSession.readMcpProviderSession(threadId)); + order.push("recover"); + return yield* codex.startSession(input); + }), + activateRecoveredSession: () => + Effect.sync(() => { + order.push("activate"); + }), + }; + const providerLayer = makeAgentBrowserProviderLayer( + true, + { ...codex, adapter: recoveryAdapter }, + { + issueMcpCredential: (request) => + Effect.sync(() => { + order.push("mcp"); + return issuedBrowserCredential(request.threadId); + }), + revokeMcpCredential: () => Effect.void, + }, + ); + + yield* Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const directory = yield* ProviderSessionDirectory.ProviderSessionDirectory; + yield* provider.startSession(threadId, { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + cwd: "/tmp/project", + runtimeMode: "full-access", + }); + codex.removeSession(threadId); + order.length = 0; + + yield* provider.recoverRestartSessions!(); + + assert.deepEqual(order, ["mcp", "recover", "activate"]); + const binding = Option.getOrThrow(yield* directory.getBinding(threadId)); + assert.equal( + (binding.runtimePayload as { readonly lastRuntimeEvent?: string }).lastRuntimeEvent, + "provider.restart-adopted", + ); + }).pipe(Effect.provide(providerLayer)); + }).pipe(Effect.provide(NodeServices.layer)), + ); + it.effect("revokes the MCP credential even when explicit adapter stop fails", () => Effect.gen(function* () { const threadId = asThreadId("thread-browser-stop-failure"); diff --git a/apps/server/src/provider/Layers/ProviderService.ts b/apps/server/src/provider/Layers/ProviderService.ts index fc4b08888..57acd7048 100644 --- a/apps/server/src/provider/Layers/ProviderService.ts +++ b/apps/server/src/provider/Layers/ProviderService.ts @@ -1278,6 +1278,9 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( }, }); } + if (routed.adapter.prepareTurnRecovery !== undefined) { + yield* routed.adapter.prepareTurnRecovery(input); + } const turn = yield* routed.adapter.sendTurn(input); yield* directory.upsert({ threadId: input.threadId, @@ -1326,6 +1329,97 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( ); }); + const recoverRestartSessions: ProviderServiceMethod<"recoverRestartSessions"> = Effect.fn( + "recoverRestartSessions", + )(function* () { + const bindings = yield* directory.listBindings(); + for (const binding of bindings) { + let adoptedAdapter: ProviderAdapterShape | undefined; + let mcpPrepared = false; + yield* Effect.gen(function* () { + const instanceId = yield* requireBindingInstanceId( + "ProviderService.recoverRestartSessions", + binding, + ); + const adapter = yield* registry.getByInstance(instanceId); + if ( + adapter.recoverSession === undefined || + adapter.activateRecoveredSession === undefined + ) { + return; + } + const rawIncarnation = readRuntimePayloadString( + binding.runtimePayload, + "sessionIncarnationId", + ); + const cwd = readPersistedCwd(binding.runtimePayload); + if ( + rawIncarnation === undefined || + cwd === undefined || + binding.resumeCursor === undefined || + binding.resumeCursor === null + ) { + return; + } + yield* prepareMcpSession(binding.threadId, instanceId); + mcpPrepared = true; + const modelSelection = readPersistedModelSelection(binding.runtimePayload); + const recovered = yield* adapter.recoverSession({ + threadId: binding.threadId, + providerInstanceId: instanceId, + sessionIncarnationId: RuntimeSessionId.make(rawIncarnation), + runtimeMode: binding.runtimeMode ?? "full-access", + cwd, + ...(modelSelection === undefined ? {} : { modelSelection }), + resumeCursor: binding.resumeCursor, + }); + if (recovered === null) { + yield* clearMcpSession(binding.threadId); + return; + } + adoptedAdapter = adapter; + currentSessionIncarnations.set(binding.threadId, { + id: RuntimeSessionId.make(rawIncarnation), + instanceId, + adapter, + }); + yield* upsertSessionBinding( + { ...recovered, providerInstanceId: instanceId }, + binding.threadId, + { + lastRuntimeEvent: "provider.restart-adopted", + lastRuntimeEventAt: yield* nowIso, + }, + ); + yield* adapter.activateRecoveredSession(binding.threadId); + }).pipe( + Effect.catchCause((cause) => + Effect.gen(function* () { + if (adoptedAdapter !== undefined) { + yield* adoptedAdapter.stopSession(binding.threadId).pipe( + Effect.catchCause((cleanupCause) => + Effect.logWarning("failed to clean up adopted provider session", { + threadId: binding.threadId, + errorTag: causeErrorTag(cleanupCause), + }), + ), + ); + const current = currentSessionIncarnations.get(binding.threadId); + if (current?.adapter === adoptedAdapter) { + currentSessionIncarnations.delete(binding.threadId); + } + } + if (mcpPrepared) yield* clearMcpSession(binding.threadId); + yield* Effect.logWarning("failed to adopt recoverable provider session", { + threadId: binding.threadId, + errorTag: causeErrorTag(cause), + }); + }), + ), + ); + } + }); + const interruptTurn: ProviderServiceMethod<"interruptTurn"> = Effect.fn("interruptTurn")( function* (rawInput) { const input = yield* decodeInputOrValidationError({ @@ -2612,8 +2706,63 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( yield* analytics.flush; }); + const runShutdown = Effect.fn("runShutdown")(function* () { + const currentAdapters = yield* getAdapterEntries; + if (currentAdapters.every(([, adapter]) => adapter.shutdown === undefined)) { + return yield* runStopAll(); + } + const bindings = yield* directory.listBindings().pipe(Effect.orElseSucceed(() => [])); + yield* Effect.forEach( + currentAdapters, + ([instanceId, adapter]) => + adapter.shutdown !== undefined + ? adapter.shutdown() + : Effect.gen(function* () { + const activeSessions = yield* adapter.listSessions(); + yield* Effect.forEach(activeSessions, (session) => + Effect.flatMap(nowIso, (lastRuntimeEventAt) => + upsertSessionBinding( + { ...session, providerInstanceId: instanceId }, + session.threadId, + { + lastRuntimeEvent: "provider.stopAll", + lastRuntimeEventAt, + }, + ), + ), + ); + yield* adapter.stopAll().pipe( + Effect.ensuring( + Effect.forEach(activeSessions, (session) => clearMcpSession(session.threadId), { + discard: true, + }), + ), + ); + yield* Effect.forEach( + bindings.filter((binding) => binding.providerInstanceId === instanceId), + (binding) => + Effect.flatMap(nowIso, (lastRuntimeEventAt) => + directory.upsert({ + threadId: binding.threadId, + provider: binding.provider, + providerInstanceId: instanceId, + status: "stopped", + runtimePayload: { + activeTurnId: null, + lastRuntimeEvent: "provider.stopAll", + lastRuntimeEventAt, + }, + }), + ), + { discard: true }, + ); + }), + { concurrency: "unbounded", discard: true }, + ); + }); + yield* Effect.addFinalizer(() => - runStopAll().pipe( + runShutdown().pipe( Effect.catchCause((cause) => Effect.logWarning("failed to stop provider service", { errorTag: causeErrorTag(cause), @@ -2625,6 +2774,7 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( return { startSession, sendTurn, + recoverRestartSessions, interruptTurn, respondToRequest, respondToUserInput, diff --git a/apps/server/src/provider/Services/ProviderAdapter.ts b/apps/server/src/provider/Services/ProviderAdapter.ts index 7ef90105f..087725b5a 100644 --- a/apps/server/src/provider/Services/ProviderAdapter.ts +++ b/apps/server/src/provider/Services/ProviderAdapter.ts @@ -99,6 +99,26 @@ export interface ProviderAdapterShape { input: ProviderSendTurnInput, ) => Effect.Effect; + /** Server-private hook. Prime uses it to replace an idle ordinary owner with a recoverable one. */ + readonly prepareTurnRecovery?: (input: ProviderSendTurnInput) => Effect.Effect; + + /** Server-private startup hook. It must fail closed and never submit provider input. */ + readonly recoverSession?: (input: { + readonly threadId: ThreadId; + readonly providerInstanceId: import("@t3tools/contracts").ProviderInstanceId; + readonly sessionIncarnationId: import("@t3tools/contracts").RuntimeSessionId; + readonly runtimeMode: ProviderSessionStartInput["runtimeMode"]; + readonly cwd: string; + readonly modelSelection?: import("@t3tools/contracts").ModelSelection; + readonly resumeCursor: unknown; + }) => Effect.Effect; + + /** Releases retained frames only after ProviderService installs exact incarnation fencing. */ + readonly activateRecoveredSession?: (threadId: ThreadId) => Effect.Effect; + + /** Process shutdown can detach recoverable ownership instead of implementing explicit Stop. */ + readonly shutdown?: () => Effect.Effect; + /** * Interrupt an active turn. */ diff --git a/apps/server/src/provider/Services/ProviderService.ts b/apps/server/src/provider/Services/ProviderService.ts index c324e9d82..f74f6cc2b 100644 --- a/apps/server/src/provider/Services/ProviderService.ts +++ b/apps/server/src/provider/Services/ProviderService.ts @@ -90,6 +90,9 @@ export interface ProviderServiceShape { input: ProviderSendTurnInput, ) => Effect.Effect; + /** Adopt eligible surviving Prime executions before startup orphan reconciliation. */ + readonly recoverRestartSessions?: () => Effect.Effect; + /** * Interrupt a running provider turn. */ diff --git a/apps/server/src/provider/prime/PrimeAgentBackendSelection.ts b/apps/server/src/provider/prime/PrimeAgentBackendSelection.ts index 2d27a8e9c..60dc920c3 100644 --- a/apps/server/src/provider/prime/PrimeAgentBackendSelection.ts +++ b/apps/server/src/provider/prime/PrimeAgentBackendSelection.ts @@ -33,6 +33,8 @@ export interface PrimeAgentBackendNegotiationInput { readonly environment: NodeJS.ProcessEnv; readonly stateDir: string; readonly providerInstanceId: ProviderInstanceId; + readonly recoveryEnabled?: boolean; + readonly architecture?: string; } export interface PrimeAgentBackendNegotiationDependencies< @@ -105,6 +107,8 @@ export function negotiatePrimeAgentBackend< environment: input.environment, stateDir: input.stateDir, providerInstanceId: input.providerInstanceId, + ...(input.recoveryEnabled === undefined ? {} : { recoveryEnabled: input.recoveryEnabled }), + ...(input.architecture === undefined ? {} : { architecture: input.architecture }), }), ); if (Result.isFailure(manager)) { diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts b/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts index 043c34870..15add700a 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts @@ -2,6 +2,7 @@ import * as NodeCrypto from "node:crypto"; import { ApprovalRequestId, + CommandId, EventId, ProviderAskSessionSideQuestionInput, PROVIDER_SESSION_AGENT_DEPTH_MAX_SETTABLE, @@ -16,7 +17,9 @@ import { type ProviderSessionSideQuestionRequestId, type ProviderRefineSessionHarnessResult, type ProviderRuntimeEvent, + type ProviderSendTurnInput, type ProviderSession, + type ModelSelection, ProviderDriverKind, ProviderInstanceId, RuntimeItemId, @@ -89,6 +92,11 @@ import { type PrimeDaemonUsage, } from "./PrimeAgentDaemonEvents.ts"; import type { PrimeAgentDaemonManager } from "./PrimeAgentDaemonManager.ts"; +import { + PrimeAgentRecoveryLedger, + type PrimeAgentRecoveryAuthority, + type PrimeAgentRecoveryLedgerShape, +} from "./PrimeAgentRecoveryLedger.ts"; import { makePrimeAgentEventPubSub, shutdownPrimeAgentEventPubSub, @@ -158,6 +166,9 @@ export interface PrimeAgentDaemonAdapterLiveOptions { readonly nativeEventLogPath?: string; readonly nativeEventLogger?: EventNdjsonLogger; readonly instanceId?: ProviderInstanceId; + /** Present only after the exact selected package passed Pylon managed-distribution proof. */ + readonly recoveryManagedBuildId?: string; + readonly recoveryLedger?: PrimeAgentRecoveryLedgerShape; readonly runtimeFactory?: ( input: PrimeAgentDaemonSessionRuntimeInput, ) => Effect.Effect< @@ -352,6 +363,8 @@ interface PrimeAgentDaemonSessionContext { nativeTranscript: Array; nativeTranscriptMessageCount: number; readonly nativeTranscriptFingerprints: Set; + recoveryTranscriptMessageCount: number; + recoveryTranscriptFingerprints: Array; readonly pendingInteractions: Map< SessionInteractionRequestId, PrimeAgentDaemonPendingInteraction @@ -393,6 +406,9 @@ interface PrimeAgentDaemonSessionContext { teardownStarted: boolean; readonly teardownCompletion: Deferred.Deferred; readonly teardownResourcesStarted: Deferred.Deferred; + readonly recoveryOwnerToken?: string; + readonly recoveryBacklog: ReadonlyArray; + recoveryPendingActivation: boolean; } function observeNativeRunStarted( @@ -513,10 +529,70 @@ function primeAgentRunCompletedNeedsHandoff(event: PrimeAgentRunCompletedEvent): ); } +function sameStrings(left: ReadonlyArray, right: ReadonlyArray): boolean { + if (left.length !== right.length) return false; + const rightSet = new Set(right); + return rightSet.size === right.length && left.every((value) => rightSet.has(value)); +} + +function sameStringRecord( + left: Readonly>, + right: Readonly>, +): boolean { + const leftKeys = Object.keys(left); + return ( + leftKeys.length === Object.keys(right).length && + leftKeys.every((key) => left[key] === right[key]) + ); +} + function primeDaemonMessageFingerprint(message: PrimeDaemonMessage): string { return NodeCrypto.createHash("sha256").update(JSON.stringify(message), "utf8").digest("hex"); } +export function planPrimeAgentRestartReplay(input: { + readonly authorityMessageCount: number; + readonly authorityFingerprints: ReadonlyArray; + readonly snapshotMessageCount: number; + readonly snapshotMessages: ReadonlyArray; +}): + | { readonly valid: true; readonly backlog: ReadonlyArray } + | { + readonly valid: false; + } { + const expectedFingerprintCount = Math.min( + input.authorityMessageCount, + PRIME_AGENT_DAEMON_TRANSCRIPT_MAX_MESSAGES, + ); + const snapshotStart = input.snapshotMessageCount - input.snapshotMessages.length; + const authorityStart = input.authorityMessageCount - input.authorityFingerprints.length; + if ( + input.authorityFingerprints.length !== expectedFingerprintCount || + input.snapshotMessages.length !== + Math.min(input.snapshotMessageCount, PRIME_AGENT_DAEMON_TRANSCRIPT_MAX_MESSAGES) || + input.snapshotMessageCount < input.authorityMessageCount || + snapshotStart > input.authorityMessageCount + ) { + return { valid: false }; + } + const overlapStart = Math.max(authorityStart, snapshotStart); + for ( + let absoluteIndex = overlapStart; + absoluteIndex < input.authorityMessageCount; + absoluteIndex += 1 + ) { + const expected = input.authorityFingerprints[absoluteIndex - authorityStart]; + const observed = input.snapshotMessages[absoluteIndex - snapshotStart]; + if (observed === undefined || primeDaemonMessageFingerprint(observed) !== expected) { + return { valid: false }; + } + } + return { + valid: true, + backlog: input.snapshotMessages.slice(input.authorityMessageCount - snapshotStart), + }; +} + // Reconnect snapshots keep only a bounded completed-message tail. Absolute // message counts make a shifted tail exact without retaining the full history. function reconcileTranscriptTail(input: { @@ -652,6 +728,39 @@ export function makePrimeAgentDaemonAdapter( const fileSystem = yield* FileSystem.FileSystem; const path = yield* Path.Path; const serverConfig = yield* ServerConfig; + const ledgerService = yield* Effect.serviceOption(PrimeAgentRecoveryLedger); + const rawRecoveryLedger = options?.recoveryLedger ?? Option.getOrUndefined(ledgerService); + const recoveryLedger = + rawRecoveryLedger === undefined + ? undefined + : { + putPrepared: (input: Parameters[0]) => + rawRecoveryLedger.putPrepared(input).pipe(Effect.orDie), + get: (threadId: string) => rawRecoveryLedger.get(threadId).pipe(Effect.orDie), + discardPrepared: ( + input: Parameters[0], + ) => rawRecoveryLedger.discardPrepared(input).pipe(Effect.orDie), + markAdmitted: (input: Parameters[0]) => + rawRecoveryLedger.markAdmitted(input).pipe(Effect.orDie), + updateTranscriptProgress: ( + input: Parameters[0], + ) => rawRecoveryLedger.updateTranscriptProgress(input).pipe(Effect.orDie), + claim: (input: Parameters[0]) => + rawRecoveryLedger.claim(input).pipe(Effect.orDie), + releaseClaim: (input: Parameters[0]) => + rawRecoveryLedger.releaseClaim(input).pipe(Effect.orDie), + commitAdoption: ( + input: Parameters[0], + ) => rawRecoveryLedger.commitAdoption(input).pipe(Effect.orDie), + markNativeCleanup: ( + input: Parameters[0], + ) => rawRecoveryLedger.markNativeCleanup(input).pipe(Effect.orDie), + markTerminalProjected: ( + input: Parameters[0], + ) => rawRecoveryLedger.markTerminalProjected(input).pipe(Effect.orDie), + deleteIfSettled: (threadId: string) => + rawRecoveryLedger.deleteIfSettled(threadId).pipe(Effect.orDie), + }; const crypto = yield* Crypto.Crypto; const runtimeContext = yield* Effect.context(); const runPromise = Effect.runPromiseWith(runtimeContext); @@ -668,6 +777,26 @@ export function makePrimeAgentDaemonAdapter( void options?.environment; const sessions = new Map(); + type PendingRecoveryStart = + | { + readonly kind: "create"; + readonly admissionRequestId: string; + readonly correlationId: string; + readonly mcpOwnerId: string; + readonly ownerToken: string; + readonly transcriptMessageCount: number; + readonly transcriptFingerprints: ReadonlyArray; + } + | { + readonly kind: "adopt"; + readonly authority: PrimeAgentRecoveryAuthority; + readonly previousOwnerToken: string; + readonly ownerToken: string; + readonly requestId: string; + readonly mcpOwnerId: string; + readonly sessionFile: string; + }; + const pendingRecoveryStarts = new Map(); const activeTeardowns = new Map< ThreadId, { @@ -1706,6 +1835,14 @@ export function makePrimeAgentDaemonAdapter( updatedAt: yield* nowIso, }; yield* Deferred.succeed(turn.completed, undefined).pipe(Effect.ignore); + if (context.recoveryOwnerToken !== undefined && !context.stopRequested) { + context.stopRequested = true; + yield* Effect.forkDetach( + Effect.yieldNow.pipe( + Effect.andThen(withThreadLock(context.threadId, stopSessionInternal(context))), + ), + ); + } return true; }); @@ -3100,6 +3237,42 @@ export function makePrimeAgentDaemonAdapter( yield* refreshContextUsage(context).pipe(Effect.forkDetach); } }).pipe( + Effect.tap(() => { + if (event._tag === "MessageCompleted") { + context.recoveryTranscriptMessageCount += 1; + context.recoveryTranscriptFingerprints.push( + primeDaemonMessageFingerprint(event.message), + ); + if ( + context.recoveryTranscriptFingerprints.length > + PRIME_AGENT_DAEMON_TRANSCRIPT_MAX_MESSAGES + ) { + context.recoveryTranscriptFingerprints.splice( + 0, + context.recoveryTranscriptFingerprints.length - + PRIME_AGENT_DAEMON_TRANSCRIPT_MAX_MESSAGES, + ); + } + } else if (event._tag === "SessionResynced") { + context.recoveryTranscriptMessageCount = event.state.messageCount; + context.recoveryTranscriptFingerprints = event.messages.map( + primeDaemonMessageFingerprint, + ); + } + const ownerToken = context.recoveryOwnerToken; + const cursor = context.runtime.recoveryCursorForEvent?.(event); + if (ownerToken === undefined || cursor === undefined) return Effect.void; + return Effect.gen(function* () { + yield* recoveryLedger!.updateTranscriptProgress({ + threadId: context.threadId, + ownerToken, + cursor, + messageCount: context.recoveryTranscriptMessageCount, + fingerprints: [...context.recoveryTranscriptFingerprints], + updatedAt: yield* nowIso, + }); + }).pipe(Effect.asVoid); + }), Effect.ensuring( withThreadLock( context.threadId, @@ -3286,7 +3459,24 @@ export function makePrimeAgentDaemonAdapter( [ // Dispose owns its own lane so a full batch of hung native // cancellation calls cannot delay process/resource teardown. - runCleanupStep({ label: "runtime-dispose", effect: context.runtime.dispose }), + runCleanupStep({ + label: "runtime-dispose", + effect: + context.recoveryOwnerToken === undefined + ? context.runtime.dispose + : context.runtime.dispose.pipe( + Effect.tap(() => + Effect.gen(function* () { + yield* recoveryLedger!.markNativeCleanup({ + threadId: context.threadId, + ownerToken: context.recoveryOwnerToken!, + updatedAt: yield* nowIso, + }); + yield* recoveryLedger!.deleteIfSettled(context.threadId); + }), + ), + ), + }), Effect.forEach(cancellationSteps, runCleanupStep, { concurrency: PRIME_AGENT_SESSION_CLEANUP_CONCURRENCY, discard: true, @@ -3484,6 +3674,14 @@ export function makePrimeAgentDaemonAdapter( }); } const approvalRequired = input.runtimeMode === "approval-required"; + const recoveryStart = pendingRecoveryStarts.get(input.threadId); + if (recoveryStart !== undefined && approvalRequired) { + return yield* new ProviderAdapterValidationError({ + provider: PROVIDER, + operation: "startSession", + issue: "Recoverable Prime Agent execution requires full-access mode.", + }); + } const existing = sessions.get(input.threadId); if (existing !== undefined && !existing.stopped) { @@ -3669,6 +3867,102 @@ export function makePrimeAgentDaemonAdapter( }, } : {}), + ...(recoveryStart === undefined + ? {} + : recoveryStart.kind === "create" + ? { + recovery: { + kind: "create" as const, + requestId: yield* randomUUIDv4, + correlationId: recoveryStart.correlationId, + mcpOwnerId: recoveryStart.mcpOwnerId, + onAuthorityReady: (authority) => + runPromise( + Effect.gen(function* () { + const sessionIncarnationId = input.sessionIncarnationId; + if (sessionIncarnationId === undefined) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: + "Recoverable Prime Agent execution is missing its session incarnation.", + }); + } + yield* recoveryLedger!.putPrepared({ + threadId: input.threadId, + providerInstanceId: boundInstanceId, + sessionIncarnationId, + admissionRequestId: recoveryStart.admissionRequestId, + turnId: null, + packageRoot: manager.bridge.packageRoot, + packageVersion: manager.bridge.version, + managedBuildId: options?.recoveryManagedBuildId ?? "", + sdkFeatures: [...(manager.bridge.sdkFeatures ?? [])], + daemonCapabilities: [...authority.daemonCapabilities], + protocolName: manager.bridge.protocolName, + protocolVersion: manager.bridge.protocolVersion, + schemaRevision: authority.schemaRevision, + activeSessionId: authority.activeSessionId, + nativeSessionId: authority.sessionId, + recoveryHandle: authority.recoveryHandle, + supervisorGeneration: authority.supervisorGeneration, + ownershipGeneration: authority.ownershipGeneration, + cursor: authority.cursor, + correlationId: recoveryStart.correlationId, + mcpOwnerId: recoveryStart.mcpOwnerId, + recoveryConfig: authority.recoveryConfig, + launchEnvironment: authority.launchEnvironment, + transcriptMessageCount: recoveryStart.transcriptMessageCount, + transcriptFingerprints: [...recoveryStart.transcriptFingerprints], + ownerToken: recoveryStart.ownerToken, + state: "prepared", + nativeCleanupProven: false, + terminalProjected: false, + checkpointQuiesced: false, + updatedAt: yield* nowIso, + }); + }), + ), + }, + } + : { + recovery: { + kind: "adopt" as const, + requestId: recoveryStart.requestId, + recoveryHandle: recoveryStart.authority.recoveryHandle, + expectedSupervisorGeneration: recoveryStart.authority.supervisorGeneration, + activeSessionId: recoveryStart.authority.activeSessionId, + sessionId: recoveryStart.authority.nativeSessionId, + sessionFile: recoveryStart.sessionFile, + correlationId: recoveryStart.authority.correlationId, + cursor: recoveryStart.authority.cursor, + previousMcpOwnerId: recoveryStart.authority.mcpOwnerId, + mcpOwnerId: recoveryStart.mcpOwnerId, + recoveryConfig: recoveryStart.authority.recoveryConfig, + launchEnvironment: recoveryStart.authority.launchEnvironment, + onAdoptionCommitted: ({ recoveryHandle, proof }) => + runPromise( + Effect.gen(function* () { + const committed = yield* recoveryLedger!.commitAdoption({ + threadId: input.threadId, + ownerToken: recoveryStart.ownerToken, + recoveryHandle, + ownershipGeneration: proof.ownershipGeneration, + cursor: proof.cursor, + mcpOwnerId: proof.mcpOwnerId, + updatedAt: yield* nowIso, + }); + if (!committed) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Recoverable Prime Agent ownership was superseded.", + }); + } + }), + ), + }, + }), ...(input.resumeCursor === undefined ? {} : { resumeCursor: input.resumeCursor }), ...(resumeSessionId === undefined ? {} : { resumeSessionId }), }).pipe( @@ -3797,19 +4091,44 @@ export function makePrimeAgentDaemonAdapter( ), ); + let recoveryBacklog: ReadonlyArray = []; + if (recoveryStart?.kind === "adopt") { + const authority = recoveryStart.authority; + const replay = planPrimeAgentRestartReplay({ + authorityMessageCount: authority.transcriptMessageCount, + authorityFingerprints: authority.transcriptFingerprints, + snapshotMessageCount: runtime.initialSnapshot.state.messageCount, + snapshotMessages: runtime.initialSnapshot.messages, + }); + if (authority.turnId === null || !replay.valid) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Prime Agent restart recovery could not prove complete event continuity.", + }); + } + recoveryBacklog = replay.backlog; + } + const now = yield* nowIso; const sessionIncarnationId = input.sessionIncarnationId ?? RuntimeSessionId.make(yield* randomUUIDv4); const session: ProviderSession = { provider: PROVIDER, providerInstanceId: boundInstanceId, - status: "ready", + status: recoveryStart?.kind === "adopt" ? "running" : "ready", runtimeMode: input.runtimeMode, cwd, model, threadId: input.threadId, resumeCursor: runtime.resumeCursor, ...(input.resumeCursor !== undefined ? { restored: true } : {}), + ...(recoveryStart?.kind === "adopt" && recoveryStart.authority.turnId !== null + ? { + activeTurnId: TurnId.make(recoveryStart.authority.turnId), + activeTurnRequestId: CommandId.make(recoveryStart.authority.admissionRequestId), + } + : {}), sessionIncarnationId, createdAt: now, updatedAt: now, @@ -3859,11 +4178,45 @@ export function makePrimeAgentDaemonAdapter( nativeTranscriptFingerprints: new Set( runtime.initialSnapshot.messages.map(primeDaemonMessageFingerprint), ), + recoveryTranscriptMessageCount: runtime.initialSnapshot.state.messageCount, + recoveryTranscriptFingerprints: runtime.initialSnapshot.messages.map( + primeDaemonMessageFingerprint, + ), pendingInteractions: new Map(), pendingApprovals: new Map(), permissionToken, approvalsAcceptedForSession: false, - activeTurn: undefined, + activeTurn: + recoveryStart?.kind === "adopt" && recoveryStart.authority.turnId !== null + ? { + id: TurnId.make(recoveryStart.authority.turnId), + controller: new AbortController(), + completed: yield* Deferred.make(), + correlationId: recoveryStart.authority.correlationId, + cancellationRequested: false, + assistantTextStreamed: false, + assistantTextEmitted: "", + assistantTextRecoveryComparable: true, + nextAssistantMessageSequence: 0, + activeAssistantItemId: undefined, + lastAssistantHadRenderableText: false, + runCompletionHandoffSequence: 0, + terminalQuiescenceGeneration: 0, + terminalQuiescenceToken: undefined, + pendingRunCompletionHandoff: undefined, + queuedInputCount: 0, + awaitingQueuedRun: false, + queuedActionObserved: false, + completedRunMessages: [], + nativeTranscriptBaselineMessageCount: + recoveryStart.authority.transcriptMessageCount, + observedToolStarts: new Set(), + observedToolCompletions: new Set(), + durableToolCallNames: new Map(), + completedToolCallNames: new Map(), + projectedPlanToolCallIds: new Set(), + } + : undefined, nativeRunActive: runtime.initialSnapshot.state.isStreaming, backgroundQuiescenceGeneration: 0, backgroundQuiescencePending: false, @@ -3896,6 +4249,11 @@ export function makePrimeAgentDaemonAdapter( teardownStarted: false, teardownCompletion: yield* Deferred.make(), teardownResourcesStarted: yield* Deferred.make(), + ...(recoveryStart === undefined + ? {} + : { recoveryOwnerToken: recoveryStart.ownerToken }), + recoveryBacklog, + recoveryPendingActivation: recoveryStart !== undefined, }; context.agentDepth = { ...context.agentDepth, @@ -3907,52 +4265,57 @@ export function makePrimeAgentDaemonAdapter( }; sessions.set(input.threadId, context); scopeTransferred = true; - yield* publishRuntimeEvent(context, { - type: "session.started", - ...(yield* makeEventStamp()), - provider: PROVIDER, - providerInstanceId: boundInstanceId, - threadId: input.threadId, - payload: { resume: input.resumeCursor !== undefined }, - }); - yield* publishSessionResources(context, runtime.initialResources); - yield* publishSessionAgentDepth(context, context.agentDepth); - yield* publishSessionCompaction(context, context.compaction); - yield* publishSessionGoal(context, context.goal); - yield* publishSessionInputQueue(context, context.inputQueue); - yield* publishRuntimeEvent(context, { - type: "session.state.changed", - ...(yield* makeEventStamp()), - provider: PROVIDER, - providerInstanceId: boundInstanceId, - threadId: input.threadId, - payload: { state: "ready", reason: "Prime Agent daemon session ready" }, - }); - yield* publishRuntimeEvent(context, { - type: "thread.started", - ...(yield* makeEventStamp()), - provider: PROVIDER, - providerInstanceId: boundInstanceId, - threadId: input.threadId, - payload: {}, - }); - if (!context.agentRosterProjected) { - for (const child of runtime.initialSnapshot.children) { - if (child.status === "queued" || child.status === "running") { - yield* publishDrafts(context, { _tag: "ChildUpdated", child }, undefined); + if (recoveryStart === undefined) { + yield* publishRuntimeEvent(context, { + type: "session.started", + ...(yield* makeEventStamp()), + provider: PROVIDER, + providerInstanceId: boundInstanceId, + threadId: input.threadId, + payload: { resume: input.resumeCursor !== undefined }, + }); + yield* publishSessionResources(context, runtime.initialResources); + yield* publishSessionAgentDepth(context, context.agentDepth); + yield* publishSessionCompaction(context, context.compaction); + yield* publishSessionGoal(context, context.goal); + yield* publishSessionInputQueue(context, context.inputQueue); + yield* publishRuntimeEvent(context, { + type: "session.state.changed", + ...(yield* makeEventStamp()), + provider: PROVIDER, + providerInstanceId: boundInstanceId, + threadId: input.threadId, + payload: { state: "ready", reason: "Prime Agent daemon session ready" }, + }); + yield* publishRuntimeEvent(context, { + type: "thread.started", + ...(yield* makeEventStamp()), + provider: PROVIDER, + providerInstanceId: boundInstanceId, + threadId: input.threadId, + payload: {}, + }); + if (!context.agentRosterProjected) { + for (const child of runtime.initialSnapshot.children) { + if (child.status === "queued" || child.status === "running") { + yield* publishDrafts(context, { _tag: "ChildUpdated", child }, undefined); + } } + context.agentRosterProjected = true; } - context.agentRosterProjected = true; } - context.eventFiber = yield* runtime.events.pipe( - Stream.runForEach((event) => consumeEvent(context, event)), - Effect.forkChild, - ); - if (runtime.inputAdmissionBusy) { + if (recoveryStart === undefined) { + context.eventFiber = yield* runtime.events.pipe( + Stream.runForEach((event) => consumeEvent(context, event)), + Effect.forkChild, + ); + } + if (recoveryStart === undefined && runtime.inputAdmissionBusy) { yield* startBackgroundQuiescenceWatchLocked(context); } context.lifecycleStarted = true; + pendingRecoveryStarts.delete(input.threadId); yield* refreshContextUsage(context).pipe(Effect.forkDetach); yield* refreshDiscoveredModels(context); return { _tag: "Started" as const, session }; @@ -4027,6 +4390,235 @@ export function makePrimeAgentDaemonAdapter( } }); + const recoveryPlatformEligible = + recoveryLedger !== undefined && + manager.recoveryEnabled && + options?.recoveryManagedBuildId !== undefined && + (manager.platform === "darwin" || manager.platform === "linux") && + (manager.architecture === "arm64" || manager.architecture === "x64"); + + const silentlyCloseSessionForRecovery = (context: PrimeAgentDaemonSessionContext) => + Effect.gen(function* () { + context.stopped = true; + sessions.delete(context.threadId); + if (context.eventFiber !== undefined) yield* Fiber.interrupt(context.eventFiber); + context.backgroundQuiescenceController?.abort(); + context.backgroundQuiescenceController = undefined; + yield* Scope.close(context.scope, Exit.void).pipe(Effect.ignore); + }); + + const prepareTurnRecovery = Effect.fn("PrimeAgentDaemonAdapter.prepareTurnRecovery")(function* ( + input: ProviderSendTurnInput, + ) { + if (!recoveryPlatformEligible) return; + const plan = yield* withThreadMutationLock( + input.threadId, + Effect.gen(function* () { + const context = sessions.get(input.threadId); + if ( + context === undefined || + context.stopped || + context.session.status !== "ready" || + context.activeTurn !== undefined || + context.session.runtimeMode !== "full-access" || + context.sessionIncarnationId === undefined || + context.recoveryOwnerToken !== undefined + ) { + return undefined; + } + const admissionRequestId = input.admissionRequestId?.trim(); + if (admissionRequestId === undefined || admissionRequestId.length === 0) { + return undefined; + } + const ownerToken = yield* randomUUIDv4; + const mcpSession = McpProviderSession.readMcpProviderSession(input.threadId); + const recoveryStart: PendingRecoveryStart = { + kind: "create", + admissionRequestId, + correlationId: yield* randomUUIDv4, + mcpOwnerId: + mcpSession === undefined + ? `pylon:none:${yield* randomUUIDv4}` + : `pylon:${mcpSession.providerSessionId}`, + ownerToken, + transcriptMessageCount: context.recoveryTranscriptMessageCount, + transcriptFingerprints: [...context.recoveryTranscriptFingerprints], + }; + const restartInput = { + threadId: context.threadId, + provider: PROVIDER, + providerInstanceId: boundInstanceId, + runtimeMode: context.session.runtimeMode, + ...(context.session.cwd === undefined ? {} : { cwd: context.session.cwd }), + ...(context.session.model === undefined + ? {} + : { + modelSelection: { + instanceId: boundInstanceId, + model: context.session.model, + }, + }), + resumeCursor: context.session.resumeCursor, + sessionIncarnationId: context.sessionIncarnationId, + } as const; + yield* silentlyCloseSessionForRecovery(context); + pendingRecoveryStarts.set(input.threadId, recoveryStart); + return { restartInput, ownerToken } as const; + }), + ); + if (plan === undefined) return; + const recoveryResult = yield* Effect.result(startSession(plan.restartInput)); + if (Result.isSuccess(recoveryResult)) return; + + pendingRecoveryStarts.delete(input.threadId); + yield* recoveryLedger!.discardPrepared({ + threadId: input.threadId, + ownerToken: plan.ownerToken, + }); + const fallback = yield* Effect.result(startSession(plan.restartInput)); + if (Result.isFailure(fallback)) return yield* fallback.failure; + }); + + const recoverSession = Effect.fn("PrimeAgentDaemonAdapter.recoverSession")(function* (input: { + readonly threadId: ThreadId; + readonly providerInstanceId: ProviderInstanceId; + readonly sessionIncarnationId: RuntimeSessionId; + readonly runtimeMode: Parameters[0]["runtimeMode"]; + readonly cwd: string; + readonly modelSelection?: ModelSelection; + readonly resumeCursor: unknown; + }) { + if (!recoveryPlatformEligible || input.runtimeMode !== "full-access") return null; + const authorityOption = yield* recoveryLedger!.get(input.threadId); + const authority = Option.getOrUndefined(authorityOption); + if ( + authority === undefined || + authority.state !== "active" || + authority.turnId === null || + authority.providerInstanceId !== input.providerInstanceId || + authority.sessionIncarnationId !== input.sessionIncarnationId || + authority.packageRoot !== manager.bridge.packageRoot || + authority.packageVersion !== manager.bridge.version || + authority.managedBuildId !== options?.recoveryManagedBuildId || + authority.protocolName !== manager.bridge.protocolName || + authority.protocolVersion !== manager.bridge.protocolVersion || + !sameStrings(authority.sdkFeatures, manager.bridge.sdkFeatures ?? []) + ) { + return null; + } + yield* manager.prepare().pipe( + Effect.mapError( + (cause) => + new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Could not validate the surviving Prime Agent daemon.", + cause, + }), + ), + ); + const readinessClient = yield* manager.openClient().pipe( + Effect.mapError( + (cause) => + new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Could not inspect the surviving Prime Agent daemon.", + cause, + }), + ), + ); + const hello = readinessClient.hello; + readinessClient.close(); + if ( + hello?.supervisorGeneration !== authority.supervisorGeneration || + hello?.schemaRevision !== authority.schemaRevision || + !sameStrings(hello?.serverCapabilities ?? [], authority.daemonCapabilities) || + !sameStringRecord(manager.launchEnvironment ?? {}, authority.launchEnvironment) + ) { + return null; + } + const ownerToken = yield* randomUUIDv4; + const claimedAt = yield* nowIso; + const claimed = yield* recoveryLedger!.claim({ + threadId: input.threadId, + expectedOwnerToken: authority.ownerToken, + nextOwnerToken: ownerToken, + updatedAt: claimedAt, + }); + if (Option.isNone(claimed)) return null; + const requestId = yield* randomUUIDv4; + const mcpSession = McpProviderSession.readMcpProviderSession(input.threadId); + const mcpOwnerId = + mcpSession === undefined + ? `pylon:none:${yield* randomUUIDv4}` + : `pylon:${mcpSession.providerSessionId}`; + pendingRecoveryStarts.set(input.threadId, { + kind: "adopt", + authority, + previousOwnerToken: authority.ownerToken, + ownerToken, + requestId, + mcpOwnerId, + sessionFile: `${authority.nativeSessionId}.jsonl`, + }); + const started = yield* Effect.result( + startSession({ + threadId: input.threadId, + provider: PROVIDER, + providerInstanceId: input.providerInstanceId, + runtimeMode: input.runtimeMode, + cwd: input.cwd, + ...(input.modelSelection === undefined ? {} : { modelSelection: input.modelSelection }), + resumeCursor: input.resumeCursor, + sessionIncarnationId: input.sessionIncarnationId, + }), + ); + if (Result.isFailure(started)) { + pendingRecoveryStarts.delete(input.threadId); + yield* recoveryLedger!.releaseClaim({ + threadId: input.threadId, + ownerToken, + previousOwnerToken: authority.ownerToken, + updatedAt: yield* nowIso, + }); + return null; + } + return started.success; + }); + + const activateRecoveredSession = Effect.fn("PrimeAgentDaemonAdapter.activateRecoveredSession")( + function* (threadId: ThreadId) { + yield* withThreadLock( + threadId, + Effect.gen(function* () { + const context = sessions.get(threadId); + if (context === undefined || context.stopped || !context.recoveryPendingActivation) + return; + const turn = context.activeTurn; + if (turn === undefined) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId, + detail: "Recovered Prime Agent execution lost its admitted turn.", + }); + } + for (const message of context.recoveryBacklog) { + yield* publishDrafts(context, { _tag: "MessageCompleted", message }, turn); + } + context.recoveryPendingActivation = false; + context.eventFiber = yield* context.runtime.events.pipe( + Stream.runForEach((event) => consumeEvent(context, event)), + Effect.forkChild, + ); + if (context.runtime.inputAdmissionBusy) { + yield* startBackgroundQuiescenceWatchLocked(context); + } + }), + ); + }, + ); + const sendTurn: PrimeAgentAdapterShape["sendTurn"] = (input) => Effect.uninterruptibleMask((restore) => Effect.gen(function* () { @@ -4237,7 +4829,7 @@ export function makePrimeAgentDaemonAdapter( } const turnId = TurnId.make(yield* randomUUIDv4); const correlationId = context.runtime.correlatedPromptLifecycleAvailable - ? yield* randomUUIDv4 + ? (context.runtime.recoveryCorrelationId ?? (yield* randomUUIDv4)) : undefined; const turn: PrimeAgentDaemonActiveTurn = { id: turnId, @@ -4312,18 +4904,6 @@ export function makePrimeAgentDaemonAdapter( const initialRlmQuiescenceToken = turn.terminalQuiescenceToken; const runPrompt = Effect.gen(function* () { const turnModel = requestedModel || context.session.model || "default"; - yield* publishRuntimeEvent(context, { - type: "turn.started", - ...(yield* makeEventStamp()), - provider: PROVIDER, - providerInstanceId: boundInstanceId, - threadId: input.threadId, - turnId: turn.id, - ...(input.admissionRequestId !== undefined - ? { admissionRequestId: input.admissionRequestId } - : {}), - payload: { model: turnModel }, - }); if (turn.correlationId !== undefined) { const lifecycle = yield* context.runtime .submitCorrelatedPrompt({ @@ -4358,6 +4938,51 @@ export function makePrimeAgentDaemonAdapter( ), ); } + if (context.recoveryOwnerToken !== undefined) { + const admissionRequestId = input.admissionRequestId?.trim(); + if (admissionRequestId === undefined || admissionRequestId.length === 0) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Recoverable Prime Agent admission lost its durable request identity.", + }); + } + const admitted = yield* recoveryLedger!.markAdmitted({ + threadId: input.threadId, + ownerToken: context.recoveryOwnerToken, + turnId: turn.id, + updatedAt: yield* nowIso, + }); + if (!admitted) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Recoverable Prime Agent admission lost its durable owner.", + }); + } + } + yield* publishRuntimeEvent(context, { + type: "turn.started", + ...(yield* makeEventStamp()), + provider: PROVIDER, + providerInstanceId: boundInstanceId, + threadId: input.threadId, + turnId: turn.id, + ...(input.admissionRequestId !== undefined + ? { admissionRequestId: input.admissionRequestId } + : {}), + payload: { model: turnModel }, + }); + if (context.recoveryPendingActivation) { + context.recoveryPendingActivation = false; + context.eventFiber = yield* context.runtime.events.pipe( + Stream.runForEach((event) => consumeEvent(context, event)), + Effect.forkChild, + ); + if (context.runtime.inputAdmissionBusy) { + yield* startBackgroundQuiescenceWatchLocked(context); + } + } if (initialRlmQuiescenceToken !== undefined) { yield* awaitRlmQuiescence(context, turn, initialRlmQuiescenceToken).pipe( Effect.catch((error) => @@ -4392,8 +5017,15 @@ export function makePrimeAgentDaemonAdapter( }); return yield* restore(runPrompt).pipe( - Effect.catch(() => + Effect.catch((error) => Effect.gen(function* () { + if (context.recoveryPendingActivation && context.recoveryOwnerToken !== undefined) { + yield* stopSessionInternal( + context, + "Prime Agent recoverable prompt admission could not be proven.", + ).pipe(Effect.ignore); + return yield* error; + } if (turn.correlationId !== undefined && turn.cancellationRequested) { yield* Deferred.await(turn.completed); return result; @@ -6298,11 +6930,47 @@ export function makePrimeAgentDaemonAdapter( } }); + const shutdown: NonNullable = () => + Effect.gen(function* () { + const contexts = Array.from(sessions.values()); + const ordinaryCompletions = yield* Effect.forEach( + contexts, + (context) => + withThreadMutationLock( + context.threadId, + Effect.gen(function* () { + if (sessions.get(context.threadId) !== context || context.stopped) return undefined; + if (context.recoveryOwnerToken !== undefined) { + context.stopped = true; + sessions.delete(context.threadId); + if (context.eventFiber !== undefined) yield* Fiber.interrupt(context.eventFiber); + context.backgroundQuiescenceController?.abort(); + context.backgroundQuiescenceController = undefined; + yield* (context.runtime.detach ?? context.runtime.dispose).pipe( + Effect.mapError((error) => + runtimeOperationError(context.threadId, "shutdown", error), + ), + ); + yield* Scope.close(context.scope, Exit.void).pipe(Effect.ignore); + return undefined; + } + return yield* stopSessionInternal(context); + }), + ), + { concurrency: "unbounded" }, + ); + yield* Effect.forEach( + ordinaryCompletions.filter((completion) => completion !== undefined), + Deferred.await, + { concurrency: "unbounded", discard: true }, + ); + }); + yield* Effect.addFinalizer(() => shutdownPrimeAgentEventPubSub({ component: "daemon", pubSub: runtimeEventPubSub, - drain: stopAll().pipe( + drain: shutdown().pipe( Effect.andThen( Effect.suspend(() => Effect.forEach(Array.from(pendingTerminalDeliveries), Deferred.await, { @@ -6325,6 +6993,10 @@ export function makePrimeAgentDaemonAdapter( conversationRollback: BUILT_IN_ADAPTER_CONVERSATION_ROLLBACK_MODES.primeDaemon, }, startSession, + prepareTurnRecovery, + recoverSession, + activateRecoveredSession, + shutdown, sendTurn, interruptTurn, respondToRequest, diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonBridge.ts b/apps/server/src/provider/prime/PrimeAgentDaemonBridge.ts index 936056de8..4f2f0c2b2 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonBridge.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonBridge.ts @@ -41,8 +41,24 @@ export class PrimeAgentDaemonBridgeError extends Schema.TaggedErrorClass; +} + +export interface PrimeAgentDaemonEventCursor { + readonly generation: string; + readonly sequence: number; +} + export interface PrimeAgentDaemonClient { readonly isConnected: boolean; + readonly hello?: PrimeAgentDaemonHello; readonly connect: (timeoutMs?: number) => Promise; readonly waitForHello: (timeoutMs?: number) => Promise; readonly request: ( @@ -50,9 +66,9 @@ export interface PrimeAgentDaemonClient { timeoutMs?: number, ) => Promise; readonly enableRequestRecovery?: () => void; - readonly supportsServerCapability?: ( - capability: "queue_message_mutation" | "correlated_prompt_lifecycle_v1", - ) => boolean; + readonly supportsServerCapability?: { + bivarianceHack(capability: string): boolean; + }["bivarianceHack"]; readonly enableAutoReconnect?: (options: { readonly recoverDaemon: () => Promise; readonly timeoutMs?: number; @@ -200,6 +216,8 @@ export interface PrimeAgentDaemonAgentConnection { readonly getSessionStats: () => Promise; readonly getRlmMaxDepthStatus?: () => Promise; readonly setRlmMaxDepth?: (maxDepth: number) => Promise; + readonly getOwnedSessionContractProof?: () => unknown; + readonly disposeOwnedSession?: (options?: { readonly timeoutMs?: number }) => Promise; readonly dispose: () => Promise; } @@ -225,6 +243,59 @@ export interface PrimeAgentDaemonAgentConnectionConstructor { ) => Promise; } +export interface PrimeAgentRecoverableOwnedSessionCreation { + readonly connection: PrimeAgentDaemonAgentConnection; + readonly state: Readonly>; + readonly recoveryHandle: string; + readonly supervisorGeneration: string; + readonly ownershipGeneration: number; +} + +export interface PrimeAgentRecoverableOwnedSessionAdoptionProof { + readonly feature: "recoverable_owned_session_adoption_v1"; + readonly status: "adopted"; + readonly supervisorGeneration: string; + readonly ownershipGeneration: number; + readonly activeSessionId: string; + readonly sessionId: string; + readonly correlationId: string; + readonly lifecycle: unknown; + readonly cursor: PrimeAgentDaemonEventCursor; + readonly mcpOwnerId: string; +} + +export interface PrimeAgentRecoverableOwnedSessionAdoption { + readonly connection: PrimeAgentDaemonAgentConnection; + readonly recoveryHandle: string; + readonly proof: PrimeAgentRecoverableOwnedSessionAdoptionProof; +} + +export interface PrimeAgentRecoverableOwnedSessionCreateOptions { + readonly requestId: string; + readonly correlationId: string; + readonly mcpOwnerId: string; + readonly config: Readonly>; + readonly sessionPath?: string; + readonly continueRecent?: boolean; + readonly launchEnv: Readonly>; + readonly connectionOptions?: Readonly>; +} + +export interface PrimeAgentRecoverableOwnedSessionAdoptionOptions { + readonly requestId: string; + readonly recoveryHandle: string; + readonly expectedSupervisorGeneration: string; + readonly activeSessionId: string; + readonly sessionId: string; + readonly correlationId: string; + readonly cursor: PrimeAgentDaemonEventCursor; + readonly previousMcpOwnerId: string; + readonly mcpOwnerId: string; + readonly config: Readonly>; + readonly launchEnv: Readonly>; + readonly connectionOptions?: Readonly>; +} + export interface PrimeAgentPublicPackage { readonly packageRoot: string; readonly moduleEntryPath: string; @@ -236,6 +307,25 @@ export interface PrimeAgentDaemonBridge extends PrimeAgentPublicPackage { readonly protocolVersion: number; /** True only for the frozen Prime SDK feature contract, never from method presence. */ readonly negotiatedDaemonSessionCapabilitiesAvailable: boolean; + /** Exact frozen client capability evidence used with the post-connect daemon gates. */ + readonly sdkFeatures?: ReadonlyArray; + readonly recoverableOwnedSessionAdoptionAvailable?: boolean; + readonly createRecoverableOwnedSession?: ( + client: PrimeAgentDaemonClient, + options: PrimeAgentRecoverableOwnedSessionCreateOptions, + ) => Promise; + readonly adoptRecoverableOwnedSession?: ( + client: PrimeAgentDaemonClient, + options: PrimeAgentRecoverableOwnedSessionAdoptionOptions, + ) => Promise; + readonly confirmRecoverableOwnedSessionAdoption?: ( + client: PrimeAgentDaemonClient, + confirmation: { + readonly requestId: string; + readonly recoveryHandle: string; + readonly proof: PrimeAgentRecoverableOwnedSessionAdoptionProof; + }, + ) => Promise; readonly DaemonClient: PrimeAgentDaemonClientConstructor; readonly DaemonAgentConnection: PrimeAgentDaemonAgentConnectionConstructor; readonly defaultDaemonSocketPath: () => string; @@ -465,15 +555,12 @@ export const locatePrimeAgentPublicPackage = Effect.fn("locatePrimeAgentPublicPa }); }); -function hasFrozenNegotiatedDaemonSessionCapabilitiesFeature(loadedModule: unknown): boolean { - if (!Predicate.isObject(loadedModule)) return false; +function frozenSdkFeatures(loadedModule: unknown): ReadonlyArray { + if (!Predicate.isObject(loadedModule)) return []; const features = loadedModule.PRIME_AGENT_SDK_FEATURES; - return ( - Array.isArray(features) && - Object.isFrozen(features) && - features.every(Predicate.isString) && - features.includes(PRIME_AGENT_NEGOTIATED_DAEMON_SESSION_CAPABILITIES_FEATURE) - ); + return Array.isArray(features) && Object.isFrozen(features) && features.every(Predicate.isString) + ? [...features] + : []; } function requireDaemonExports(input: { @@ -520,8 +607,20 @@ function requireDaemonExports(input: { const daemonClient = input.loadedModule.DaemonClient; const daemonAgentConnection = input.loadedModule.DaemonAgentConnection; const defaultDaemonSocketPath = input.loadedModule.defaultDaemonSocketPath; - const negotiatedDaemonSessionCapabilitiesAvailable = - hasFrozenNegotiatedDaemonSessionCapabilitiesFeature(input.loadedModule); + const sdkFeatures = frozenSdkFeatures(input.loadedModule); + const negotiatedDaemonSessionCapabilitiesAvailable = sdkFeatures.includes( + PRIME_AGENT_NEGOTIATED_DAEMON_SESSION_CAPABILITIES_FEATURE, + ); + const createRecoverableOwnedSession = input.loadedModule.createRecoverableOwnedSession; + const adoptRecoverableOwnedSession = input.loadedModule.adoptRecoverableOwnedSession; + const confirmRecoverableOwnedSessionAdoption = + input.loadedModule.confirmRecoverableOwnedSessionAdoption; + const recoverableOwnedSessionAdoptionAvailable = + sdkFeatures.includes("recoverable_owned_session_adoption_v1") && + sdkFeatures.includes("caller_owned_session_environment_cleanup_v1") && + Predicate.isFunction(createRecoverableOwnedSession) && + Predicate.isFunction(adoptRecoverableOwnedSession) && + Predicate.isFunction(confirmRecoverableOwnedSessionAdoption); if ( !Predicate.isFunction(daemonClient) || !Predicate.isObject(daemonClient.prototype) || @@ -577,6 +676,22 @@ function requireDaemonExports(input: { protocolName: PRIME_AGENT_DAEMON_PROTOCOL_NAME, protocolVersion: metadata.value.DAEMON_PROTOCOL_VERSION, negotiatedDaemonSessionCapabilitiesAvailable, + sdkFeatures, + recoverableOwnedSessionAdoptionAvailable, + ...(recoverableOwnedSessionAdoptionAvailable + ? { + createRecoverableOwnedSession: createRecoverableOwnedSession as NonNullable< + PrimeAgentDaemonBridge["createRecoverableOwnedSession"] + >, + adoptRecoverableOwnedSession: adoptRecoverableOwnedSession as NonNullable< + PrimeAgentDaemonBridge["adoptRecoverableOwnedSession"] + >, + confirmRecoverableOwnedSessionAdoption: + confirmRecoverableOwnedSessionAdoption as NonNullable< + PrimeAgentDaemonBridge["confirmRecoverableOwnedSessionAdoption"] + >, + } + : {}), DaemonClient: daemonClient as PrimeAgentDaemonClientConstructor, DaemonAgentConnection: daemonAgentConnection as unknown as PrimeAgentDaemonAgentConnectionConstructor, diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonManager.test.ts b/apps/server/src/provider/prime/PrimeAgentDaemonManager.test.ts index ef14f378d..2c890117a 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonManager.test.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonManager.test.ts @@ -103,6 +103,7 @@ function fakeBridge(input: { class FakeClient implements PrimeAgentDaemonClient { isConnected = false; + hello = hello as NonNullable; readonly socketPath: string; constructor(socketPath: string) { @@ -188,6 +189,8 @@ function fakeBridge(input: { protocolName: "prime-agent.daemon", protocolVersion: 7, negotiatedDaemonSessionCapabilitiesAvailable: false, + sdkFeatures: [], + recoverableOwnedSessionAdoptionAvailable: false, DaemonClient: FakeClient, DaemonAgentConnection: FakeAgentConnection, defaultDaemonSocketPath: () => "/tmp/user-prime-agent.sock", @@ -203,6 +206,7 @@ function managerFixture(options?: { readonly tempDir?: string; readonly platform?: NodeJS.Platform; readonly injectBridge?: boolean; + readonly recoverable?: boolean; }) { const commands: CapturedCommand[] = []; const processes: FakeProcess[] = []; @@ -218,6 +222,21 @@ function managerFixture(options?: { platform: options?.platform ?? "linux", tempDir: options?.tempDir ?? "/tmp", }); + const recoveryHello = options?.recoverable + ? { + type: "daemon_hello", + socketPath: paths.socket, + protocol: { name: "prime-agent.daemon", version: 7 }, + schemaRevision: 30, + supervisorGeneration: "supervisor-1", + serverCapabilities: [ + ...PRIME_AGENT_REQUIRED_DAEMON_CAPABILITIES, + "daemon_recoverable_owned_session_adoption_v1", + "caller_owned_session_environment_cleanup_v1", + "authoritative_owned_session_cleanup_v1", + ], + } + : undefined; const bridge = fakeBridge({ socket: paths.socket, processes, @@ -227,9 +246,27 @@ function managerFixture(options?: { readinessFailures, connectionAvailable, calls, - ...(options?.hello === undefined ? {} : { hello: options.hello }), + ...(options?.hello === undefined && recoveryHello === undefined + ? {} + : { hello: options?.hello ?? recoveryHello }), ...(options?.failConnect === undefined ? {} : { failConnect: options.failConnect }), }); + if (options?.recoverable) { + Object.assign(bridge, { + sdkFeatures: [ + "recoverable_owned_session_adoption_v1", + "caller_owned_session_environment_cleanup_v1", + ], + recoverableOwnedSessionAdoptionAvailable: true, + createRecoverableOwnedSession: async () => { + throw new Error("not used by manager test"); + }, + adoptRecoverableOwnedSession: async () => { + throw new Error("not used by manager test"); + }, + confirmRecoverableOwnedSessionAdoption: async () => undefined, + }); + } const spawner = Layer.succeed( ChildProcessSpawner.ChildProcessSpawner, ChildProcessSpawner.make((command) => @@ -261,6 +298,8 @@ function managerFixture(options?: { readinessRetryDelay: Duration.zero, readinessRetries: 4, shutdownTimeout: Duration.zero, + recoveryEnabled: options?.recoverable === true, + architecture: "arm64", ...(options?.injectBridge === false ? {} : { bridge }), }).pipe(Effect.provide(Layer.merge(NodeServices.layer, spawner))); return { @@ -586,4 +625,24 @@ describe("PrimeAgentDaemonManager lifecycle", () => { ); }, ); + + it.effect("never shuts down a compatible recovery supervisor retained by another process", () => { + const fixture = managerFixture({ existingLive: true, recoverable: true }); + return Effect.scoped( + Effect.gen(function* () { + const manager = yield* fixture.make; + const client = yield* manager.openClient(); + client.close(); + expect(fixture.shutdownRequests).toEqual([]); + expect(fixture.commands).toHaveLength(0); + }), + ).pipe( + Effect.andThen( + Effect.sync(() => { + expect(fixture.shutdownRequests).toEqual([]); + expect(fixture.commands).toHaveLength(0); + }), + ), + ); + }); }); diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonManager.ts b/apps/server/src/provider/prime/PrimeAgentDaemonManager.ts index 2d9eee8b0..bc476d353 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonManager.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonManager.ts @@ -78,6 +78,13 @@ export interface PrimeAgentDaemonManager { PrimeAgentDaemonClient, PrimeAgentDaemonManagerOpenError >; + /** Exact caller-owned worker environment captured before any Prime worker launch. */ + readonly launchEnvironment?: Readonly>; + readonly recoveryEnabled?: boolean; + readonly platform?: NodeJS.Platform; + readonly architecture?: string; + /** Keeps the compatible supervisor alive while at least one ledger authority can be adopted. */ + readonly retainForRecovery?: () => () => void; /** Directly accepted by DaemonClient.enableAutoReconnect({ recoverDaemon }). */ readonly recover: () => Promise; } @@ -96,6 +103,9 @@ export interface PrimeAgentDaemonManagerInput { readonly platform?: NodeJS.Platform; /** Test-only temp-directory injection. */ readonly tempDir?: string; + /** Enables retention only after the selected package passed Pylon managed-distribution proof. */ + readonly recoveryEnabled?: boolean; + readonly architecture?: string; /** Tests may supply the already validated public bridge without importing a real installation. */ readonly bridge?: PrimeAgentDaemonBridge; } @@ -112,6 +122,10 @@ const daemonHelloSchema = Schema.Struct({ name: Schema.String, version: Schema.Int, }), + schemaRevision: Schema.optional(Schema.Int), + appVersion: Schema.optional(Schema.String), + buildId: Schema.optional(Schema.String), + supervisorGeneration: Schema.optional(Schema.String), serverCapabilities: Schema.Array(Schema.String), }); const decodeDaemonHello = Schema.decodeUnknownOption(daemonHelloSchema); @@ -282,6 +296,26 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage ); } const bridge = input.bridge ?? (yield* loadPrimeAgentDaemonBridge(input.executablePath)); + const recoveryEnabled = + input.recoveryEnabled === true && bridge.recoverableOwnedSessionAdoptionAvailable === true; + const launchEnvironment = Object.fromEntries( + Object.entries( + makePrimeAgentDaemonEnvironment({ + settings: input.settings, + environment: input.environment ?? hostEnvironment, + }), + ).filter((entry): entry is [string, string] => typeof entry[1] === "string"), + ); + let recoveryRetainers = 0; + const retainForRecovery = () => { + recoveryRetainers += 1; + let released = false; + return () => { + if (released) return; + released = true; + recoveryRetainers = Math.max(0, recoveryRetainers - 1); + }; + }; const defaultSocket = bridge.defaultDaemonSocketPath(); const socket = paths.socket === defaultSocket @@ -296,6 +330,7 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage const shutdownTimeout = input.shutdownTimeout ?? Duration.seconds(5); const semaphore = yield* Semaphore.make(1); let running: RunningDaemon | undefined; + let retainedExistingDaemon = false; let closing = false; const removeSocket = () => @@ -666,19 +701,32 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage ); const existing = yield* probeExistingDaemon(); if (Option.isSome(existing)) { + const hello = decodeDaemonHello(existing.value.hello); + const recoverable = + recoveryEnabled && + Option.isSome(hello) && + (hello.value.schemaRevision ?? 0) >= 30 && + typeof hello.value.supervisorGeneration === "string" && + hello.value.supervisorGeneration.length > 0 && + [ + "daemon_recoverable_owned_session_adoption_v1", + "caller_owned_session_environment_cleanup_v1", + "authoritative_owned_session_cleanup_v1", + ].every((capability) => hello.value.serverCapabilities.includes(capability)); + if (recoverable) { + retainedExistingDaemon = true; + return existing.value; + } yield* retireExistingDaemon(existing.value); } + retainedExistingDaemon = false; yield* removeSocket(); const processScope = yield* Scope.make("sequential"); - const environment = makePrimeAgentDaemonEnvironment({ - settings: input.settings, - environment: input.environment ?? hostEnvironment, - }); const command = ChildProcess.make( input.executablePath, ["--mode", "daemon", "--daemon-socket", socket, "--offline", "--session-dir", sessionDir], - { env: environment, extendEnv: false }, + { env: launchEnvironment, extendEnv: false }, ); const handle = yield* spawner.spawn(command).pipe( Effect.provideService(Scope.Scope, processScope), @@ -699,6 +747,7 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage Effect.onError(() => stopCapturedDaemon(state, false)), ); running = state; + retainedExistingDaemon = false; return readinessClient; }); @@ -719,10 +768,23 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage closing = true; const captured = running; running = undefined; - if (captured) yield* stopCapturedDaemon(captured); + if (recoveryRetainers > 0) { + // The standalone process scope is deliberately left open. The daemon is detached + // from this Pylon process and retains the exact same supervisor generation. + return; + } + if (captured) { + yield* stopCapturedDaemon(captured); + return; + } + if (retainedExistingDaemon) { + // This process did not spawn the compatible supervisor. A competing replacement + // may already own it, so shutdown must never revoke that process's authority. + return; + } }), ); - yield* Effect.addFinalizer(() => shutdown); + yield* Effect.addFinalizer(() => shutdown.pipe(Effect.ignore)); return { bridge, @@ -730,6 +792,11 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage sessionDir, prepare, openClient, + launchEnvironment, + recoveryEnabled, + platform, + architecture: input.architecture ?? "unsupported", + retainForRecovery, recover: () => runPromise(prepare()), } satisfies PrimeAgentDaemonManager; }); diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts index a03d07d68..59dce0a54 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts @@ -301,6 +301,7 @@ function fixture(options?: { readonly verifyManagedSourceImpl?: () => Promise; readonly unsubscribeImpl?: () => void; readonly disposeImpl?: () => Promise; + readonly recoveryMode?: "create" | "adopt"; }) { const captures: Captures = { order: [], @@ -334,6 +335,19 @@ function fixture(options?: { class FakeClient implements PrimeAgentDaemonClient { isConnected = true; + hello = { + type: "daemon_hello" as const, + protocol: { name: "prime-agent.daemon" as const, version: 7 }, + socketPath: "/tmp/pylon-prime.sock", + appVersion: "0.7.1", + schemaRevision: 30, + supervisorGeneration: "supervisor-1", + serverCapabilities: [ + "daemon_recoverable_owned_session_adoption_v1", + "caller_owned_session_environment_cleanup_v1", + "authoritative_owned_session_cleanup_v1", + ], + }; connect(): Promise { return Promise.resolve(); } @@ -848,6 +862,10 @@ function fixture(options?: { Promise.resolve({ maxDepth, source: "chat", globalSaved: false }) ); } + disposeOwnedSession(): Promise { + captures.order.push("dispose-owned"); + return Promise.resolve({ status: "completed" }); + } dispose(): Promise { captures.order.push("dispose"); captures.disposeCount += 1; @@ -865,6 +883,51 @@ function fixture(options?: { options?.correlatedPromptLifecycleSdkFeature ?? options?.correlatedPromptLifecycleCapability ?? false, + sdkFeatures: + options?.recoveryMode === undefined + ? [] + : ["recoverable_owned_session_adoption_v1", "caller_owned_session_environment_cleanup_v1"], + recoverableOwnedSessionAdoptionAvailable: options?.recoveryMode !== undefined, + ...(options?.recoveryMode === undefined + ? {} + : { + createRecoverableOwnedSession: async () => { + captures.order.push("create-recoverable"); + return { + connection: new FakeConnection(), + recoveryHandle: "handle-1", + supervisorGeneration: "supervisor-1", + ownershipGeneration: 0, + state: { + activeSessionId: "active-secret-1", + sessionId: "session-1", + sessionFile: "/state/provider-sessions/thread-safe/session.jsonl", + }, + }; + }, + adoptRecoverableOwnedSession: async () => { + captures.order.push("adopt-recoverable"); + return { + connection: new FakeConnection(), + recoveryHandle: "handle-2", + proof: { + feature: "recoverable_owned_session_adoption_v1" as const, + status: "adopted" as const, + lifecycle: { phase: "owned" }, + supervisorGeneration: "supervisor-1", + activeSessionId: "active-secret-1", + sessionId: "session-1", + correlationId: "correlation-1", + mcpOwnerId: "pylon:mcp-2", + ownershipGeneration: 1, + cursor: { generation: "events-1", sequence: 9 }, + }, + }; + }, + confirmRecoverableOwnedSessionAdoption: async () => { + captures.order.push("confirm-adoption"); + }, + }), DaemonClient: FakeClient, DaemonAgentConnection: FakeConnection, defaultDaemonSocketPath: () => "/tmp/prime-agent.sock", @@ -873,6 +936,14 @@ function fixture(options?: { bridge, socket: "/tmp/pylon-prime.sock", sessionDir: "/state/shared-daemon-sessions", + launchEnvironment: { HOME: "/private/home" }, + recoveryEnabled: options?.recoveryMode !== undefined, + platform: "darwin", + architecture: "arm64", + retainForRecovery: () => { + captures.order.push("retain-daemon"); + return () => captures.order.push("release-daemon"); + }, prepare: () => Effect.void, openClient: () => Effect.sync(() => { @@ -890,6 +961,7 @@ function fixture(options?: { resumeSessionId?: string, mcpServer?: PrimeAgentDaemonSessionRuntimeInput["mcpServer"], expectedExtension?: { readonly path: string; readonly markerCommand: string }, + recovery?: PrimeAgentDaemonSessionRuntimeInput["recovery"], ) => makePrimeAgentDaemonSessionRuntime({ manager, @@ -917,6 +989,7 @@ function fixture(options?: { ...(resumeCursor === undefined ? {} : { resumeCursor }), ...(resumeSessionId === undefined ? {} : { resumeSessionId }), ...(mcpServer === undefined ? {} : { mcpServer }), + ...(recovery === undefined ? {} : { recovery }), }); const emit = (event: unknown) => Promise.resolve(listener?.(event)); const emitWatch = (event: unknown) => Promise.resolve(watcherListener?.(event)); @@ -12923,4 +12996,144 @@ describe("Prime Agent live activity privacy boundary", () => { }), ), ); + it.effect("persists create authority before exposing a recoverable runtime", () => + Effect.scoped( + Effect.gen(function* () { + const side = fixture({ + recoveryMode: "create", + rawSnapshot: { + ...snapshot(), + lastEventCursor: { generation: "events-1", sequence: 4 }, + }, + }); + const runtime = yield* side.make( + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { + kind: "create", + requestId: "request-create-1", + correlationId: "correlation-1", + mcpOwnerId: "pylon:none:1", + onAuthorityReady: async (authority) => { + expect(authority.cursor).toEqual({ generation: "events-1", sequence: 4 }); + side.captures.order.push("authority-durable"); + }, + }, + ); + + expect(side.captures.order.indexOf("create-recoverable")).toBeLessThan( + side.captures.order.indexOf("authority-durable"), + ); + expect(runtime.recoveryCorrelationId).toBe("correlation-1"); + yield* runtime.detach!; + }), + ), + ); + + it.effect("commits adoption, restores MCP, then confirms before replay is exposed", () => + Effect.scoped( + Effect.gen(function* () { + const side = fixture({ + recoveryMode: "adopt", + rawSnapshot: { + ...snapshot(), + lastEventCursor: { generation: "events-1", sequence: 9 }, + }, + }); + const runtime = yield* side.make( + PRIME_AGENT_DAEMON_RESUME_CURSOR, + undefined, + undefined, + "session-1", + { + ownerId: "pylon:mcp-2", + server: { + name: "t3-code", + type: "http", + url: "http://127.0.0.1/mcp", + headers: {}, + }, + }, + undefined, + { + kind: "adopt", + requestId: "request-adopt-1", + recoveryHandle: "handle-1", + expectedSupervisorGeneration: "supervisor-1", + activeSessionId: "active-secret-1", + sessionId: "session-1", + sessionFile: "/state/provider-sessions/thread-safe/session.jsonl", + correlationId: "correlation-1", + cursor: { generation: "events-1", sequence: 4 }, + previousMcpOwnerId: "pylon:mcp-1", + mcpOwnerId: "pylon:mcp-2", + recoveryConfig: { cwd: "/work/project" }, + launchEnvironment: { HOME: "/private/home" }, + onAdoptionCommitted: async () => { + side.captures.order.push("ledger-committed"); + }, + }, + ); + const initial = yield* Stream.runHead(runtime.events); + + expect(initial._tag).toBe("Some"); + expect(side.captures.order).toEqual( + expect.arrayContaining([ + "adopt-recoverable", + "retain-daemon", + "ledger-committed", + "replace-mcp", + "confirm-adoption", + ]), + ); + expect(side.captures.order.indexOf("ledger-committed")).toBeLessThan( + side.captures.order.indexOf("replace-mcp"), + ); + expect(side.captures.order.indexOf("replace-mcp")).toBeLessThan( + side.captures.order.indexOf("confirm-adoption"), + ); + yield* runtime.detach!; + }), + ), + ); + + it.effect("uses authoritative owned cleanup and releases daemon retention only after proof", () => + Effect.scoped( + Effect.gen(function* () { + const side = fixture({ + recoveryMode: "create", + rawSnapshot: { + ...snapshot(), + lastEventCursor: { generation: "events-1", sequence: 4 }, + }, + }); + const runtime = yield* side.make( + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { + kind: "create", + requestId: "request-create-cleanup", + correlationId: "correlation-cleanup", + mcpOwnerId: "pylon:none:cleanup", + onAuthorityReady: async () => undefined, + }, + ); + yield* runtime.dispose; + + expect(side.captures.order).toContain("dispose-owned"); + expect(side.captures.order.indexOf("dispose-owned")).toBeLessThan( + side.captures.order.indexOf("release-daemon"), + ); + expect(side.captures.order).not.toContain("dispose"); + }), + ), + ); }); diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts index e9d2213b3..bad48409c 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts @@ -37,6 +37,8 @@ import * as Stream from "effect/Stream"; import { type PrimeAgentDaemonAcpMcpServer, type PrimeAgentDaemonAgentConnection, + type PrimeAgentDaemonEventCursor, + type PrimeAgentRecoverableOwnedSessionAdoptionProof, type PrimeAgentDaemonExtensionUiResponse, type PrimeAgentDaemonImage, type PrimeAgentDaemonQueueMode, @@ -883,6 +885,45 @@ export interface PrimeAgentDaemonSessionRuntimeInput { readonly resumeCursor?: unknown; /** Private stable native id selected from the server-owned identity sidecar. */ readonly resumeSessionId?: string; + readonly recovery?: + | { + readonly kind: "create"; + readonly requestId: string; + readonly correlationId: string; + readonly mcpOwnerId: string; + readonly onAuthorityReady: (authority: { + readonly recoveryHandle: string; + readonly supervisorGeneration: string; + readonly ownershipGeneration: number; + readonly activeSessionId: string; + readonly sessionId: string; + readonly sessionFile: string; + readonly cursor: PrimeAgentDaemonEventCursor; + readonly recoveryConfig: Readonly>; + readonly launchEnvironment: Readonly>; + readonly daemonCapabilities: ReadonlyArray; + readonly schemaRevision: number; + }) => Promise; + } + | { + readonly kind: "adopt"; + readonly requestId: string; + readonly recoveryHandle: string; + readonly expectedSupervisorGeneration: string; + readonly activeSessionId: string; + readonly sessionId: string; + readonly sessionFile: string; + readonly correlationId: string; + readonly cursor: PrimeAgentDaemonEventCursor; + readonly previousMcpOwnerId: string; + readonly mcpOwnerId: string; + readonly recoveryConfig: Readonly>; + readonly launchEnvironment: Readonly>; + readonly onAdoptionCommitted: (authority: { + readonly recoveryHandle: string; + readonly proof: PrimeAgentRecoverableOwnedSessionAdoptionProof; + }) => Promise; + }; } export interface PrimeAgentDaemonPromptInput { @@ -1139,6 +1180,14 @@ export interface PrimeAgentDaemonSessionRuntime { PrimeAgentDaemonSessionStats, PrimeAgentDaemonSessionRuntimeError >; + readonly recoveryCorrelationId?: string; + readonly recoveryCursor?: PrimeAgentDaemonEventCursor; + readonly recoveryCursorForEvent?: ( + event: PrimeDaemonEvent, + ) => PrimeAgentDaemonEventCursor | undefined; + /** Close only this Pylon owner. The recoverable worker and MCP authority stay with the daemon. */ + readonly detach?: Effect.Effect; + /** Explicit cleanup requires Prime's authoritative owned-session result. */ readonly dispose: Effect.Effect; } @@ -1298,6 +1347,19 @@ function safeCatalogModels( return Option.some(safeModels); } +function recoveryCursorFromSnapshot(value: unknown): PrimeAgentDaemonEventCursor | undefined { + if (!Predicate.isObject(value) || !Predicate.isObject(value.lastEventCursor)) return undefined; + const generation = value.lastEventCursor.generation; + const sequence = value.lastEventCursor.sequence; + return typeof generation === "string" && + generation.length > 0 && + Number.isSafeInteger(sequence) && + typeof sequence === "number" && + sequence >= 0 + ? { generation, sequence } + : undefined; +} + export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemonSessionRuntime")( function* ( input: PrimeAgentDaemonSessionRuntimeInput, @@ -1351,6 +1413,7 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo let connection: PrimeAgentDaemonAgentConnection | undefined; let unsubscribe: (() => void) | undefined; let disposed = false; + let detached = false; let disposeStarted = false; const disposeCompletion = yield* Deferred.make(); let needsResumeAfterAbort = false; @@ -1770,100 +1833,262 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo ...(configuredModel && configuredModel !== "default" ? { model: configuredModel } : {}), ...(input.thinkingLevel === undefined ? {} : { thinking: input.thinkingLevel }), }; - const createCommand = { - type: "create", - lifecycle: "client_owned", - ...(resumeSessionId === undefined - ? { continueRecent: shouldContinue } - : { sessionPath: resumeSessionId, continueRecent: false }), - config: sessionRuntimeConfig, - } as const; - const requestCreate = Effect.tryPromise({ - try: () => client.request(createCommand, COMMAND_TIMEOUT_MS), - catch: () => - runtimeError( + let activeSessionId: string; + let sessionId: string; + let sessionFile: string; + let createdRecovery: + | { + readonly recoveryHandle: string; + readonly supervisorGeneration: string; + readonly ownershipGeneration: number; + } + | undefined; + let adoptedRecovery: + | { + readonly recoveryHandle: string; + readonly proof: PrimeAgentRecoverableOwnedSessionAdoptionProof; + } + | undefined; + let releaseManagerRecoveryRetention: (() => void) | undefined; + let confirmAdoption: (() => Promise) | undefined; + const recoveryConnectionOptions = { + closeClientOnDispose: false, + supportsExtensionUi: true, + ...(input.disableAutoReconnect === true ? {} : { recoverDaemon: input.manager.recover }), + }; + + if (input.recovery?.kind === "create") { + const recovery = input.recovery; + const createRecoverableOwnedSession = input.manager.bridge.createRecoverableOwnedSession; + if ( + !input.manager.recoveryEnabled || + !input.manager.bridge.recoverableOwnedSessionAdoptionAvailable || + !Predicate.isFunction(createRecoverableOwnedSession) + ) { + yield* closeClient; + return yield* runtimeError( "create-session", - "request-failed", - "The daemon did not complete the create command.", + "incompatible-api", + "Recoverable Prime Agent execution is unavailable.", + ); + } + const created = yield* Effect.tryPromise({ + try: () => + createRecoverableOwnedSession(client, { + requestId: recovery.requestId, + correlationId: recovery.correlationId, + mcpOwnerId: recovery.mcpOwnerId, + config: sessionRuntimeConfig, + ...(resumeSessionId === undefined + ? { continueRecent: shouldContinue } + : { sessionPath: resumeSessionId, continueRecent: false }), + launchEnv: input.manager.launchEnvironment ?? {}, + connectionOptions: recoveryConnectionOptions, + }), + catch: () => + runtimeError( + "create-session", + "request-failed", + "Recoverable Prime Agent execution could not be created.", + ), + }).pipe(Effect.onError(() => closeClient)); + connection = created.connection; + const state = created.state; + activeSessionId = + typeof state.activeSessionId === "string" ? state.activeSessionId.trim() : ""; + sessionId = typeof state.sessionId === "string" ? state.sessionId.trim() : ""; + sessionFile = typeof state.sessionFile === "string" ? state.sessionFile.trim() : ""; + createdRecovery = { + recoveryHandle: created.recoveryHandle, + supervisorGeneration: created.supervisorGeneration, + ownershipGeneration: created.ownershipGeneration, + }; + releaseManagerRecoveryRetention = input.manager.retainForRecovery?.(); + } else if (input.recovery?.kind === "adopt") { + const recovery = input.recovery; + const adoptRecoverableOwnedSession = input.manager.bridge.adoptRecoverableOwnedSession; + const confirmRecoverableOwnedSessionAdoption = + input.manager.bridge.confirmRecoverableOwnedSessionAdoption; + if ( + !input.manager.recoveryEnabled || + !input.manager.bridge.recoverableOwnedSessionAdoptionAvailable || + !Predicate.isFunction(adoptRecoverableOwnedSession) || + !Predicate.isFunction(confirmRecoverableOwnedSessionAdoption) + ) { + yield* closeClient; + return yield* runtimeError( + "attach-session", + "incompatible-api", + "Recoverable Prime Agent execution is unavailable.", + ); + } + const adopted = yield* Effect.tryPromise({ + try: () => + adoptRecoverableOwnedSession(client, { + requestId: recovery.requestId, + recoveryHandle: recovery.recoveryHandle, + expectedSupervisorGeneration: recovery.expectedSupervisorGeneration, + activeSessionId: recovery.activeSessionId, + sessionId: recovery.sessionId, + correlationId: recovery.correlationId, + cursor: recovery.cursor, + previousMcpOwnerId: recovery.previousMcpOwnerId, + mcpOwnerId: recovery.mcpOwnerId, + config: recovery.recoveryConfig, + launchEnv: recovery.launchEnvironment, + connectionOptions: recoveryConnectionOptions, + }), + catch: () => + runtimeError( + "attach-session", + "request-failed", + "Recoverable Prime Agent execution could not be adopted.", + ), + }).pipe(Effect.onError(() => closeClient)); + releaseManagerRecoveryRetention = input.manager.retainForRecovery?.(); + yield* Effect.tryPromise({ + try: () => + recovery.onAdoptionCommitted({ + recoveryHandle: adopted.recoveryHandle, + proof: adopted.proof, + }), + catch: () => + runtimeError( + "attach-session", + "request-failed", + "Recoverable Prime Agent ownership could not be durably recorded.", + ), + }).pipe( + Effect.onError(() => + Effect.promise(async () => { + await adopted.connection.dispose().catch(() => undefined); + client.close(); + }), ), - }).pipe(Effect.onError(() => closeClient)); - const createResponse = yield* Effect.gen(function* () { - let response = yield* requestCreate; - for (const delay of OWNED_SESSION_RELEASE_RETRY_DELAYS_MS) { - if (!shouldContinue || Option.isNone(decodeCreateSessionAlreadyActiveFailure(response))) - break; - yield* Effect.sleep(delay); - response = yield* requestCreate; - } - return response; - }).pipe(Effect.onInterrupt(() => closeClient)); - const created = decodeCreateSuccess(createResponse); - if (Option.isNone(created)) { - yield* closeClient; - const alreadyActive = decodeCreateSessionAlreadyActiveFailure(createResponse); - if (Option.isSome(alreadyActive)) { + ); + confirmAdoption = () => + confirmRecoverableOwnedSessionAdoption(client, { + requestId: recovery.requestId, + recoveryHandle: adopted.recoveryHandle, + proof: adopted.proof, + }); + connection = adopted.connection; + activeSessionId = recovery.activeSessionId; + sessionId = recovery.sessionId; + sessionFile = recovery.sessionFile; + adoptedRecovery = { recoveryHandle: adopted.recoveryHandle, proof: adopted.proof }; + } else { + const createCommand = { + type: "create", + lifecycle: "client_owned", + ...(resumeSessionId === undefined + ? { continueRecent: shouldContinue } + : { sessionPath: resumeSessionId, continueRecent: false }), + config: sessionRuntimeConfig, + } as const; + const requestCreate = Effect.tryPromise({ + try: () => client.request(createCommand, COMMAND_TIMEOUT_MS), + catch: () => + runtimeError( + "create-session", + "request-failed", + "The daemon did not complete the create command.", + ), + }).pipe(Effect.onError(() => closeClient)); + const createResponse = yield* Effect.gen(function* () { + let response = yield* requestCreate; + for (const delay of OWNED_SESSION_RELEASE_RETRY_DELAYS_MS) { + if (!shouldContinue || Option.isNone(decodeCreateSessionAlreadyActiveFailure(response))) + break; + yield* Effect.sleep(delay); + response = yield* requestCreate; + } + return response; + }).pipe(Effect.onInterrupt(() => closeClient)); + const created = decodeCreateSuccess(createResponse); + if (Option.isNone(created)) { + yield* closeClient; + const alreadyActive = decodeCreateSessionAlreadyActiveFailure(createResponse); + if (Option.isSome(alreadyActive)) { + return yield* runtimeError( + "create-session", + "session-already-active", + "SessionAlreadyActiveError: Prime Agent session is already active in another client.", + ); + } + const failed = decodeCreateFailure(createResponse); return yield* runtimeError( "create-session", - "session-already-active", - "SessionAlreadyActiveError: Prime Agent session is already active in another client.", + Option.isSome(failed) ? "request-failed" : "invalid-response", + Option.isSome(failed) + ? "The daemon rejected the create command." + : "The daemon returned an invalid create response.", ); } - const failed = decodeCreateFailure(createResponse); - return yield* runtimeError( - "create-session", - Option.isSome(failed) ? "request-failed" : "invalid-response", - Option.isSome(failed) - ? "The daemon rejected the create command." - : "The daemon returned an invalid create response.", - ); - } - const activeSessionId = created.value.data.activeSessionId.trim(); - if (activeSessionId.length === 0) { - client.close(); - return yield* runtimeError( - "create-session", - "invalid-response", - "The daemon create response omitted its active session identifier.", - ); + activeSessionId = created.value.data.activeSessionId.trim(); + sessionId = created.value.data.sessionId.trim(); + sessionFile = created.value.data.sessionFile.trim(); + const completeUnattachedOwnedSession = Effect.tryPromise({ + try: () => + client.request({ type: "complete_owned_session", activeSessionId }, COMMAND_TIMEOUT_MS), + catch: () => undefined, + }).pipe(Effect.ignore, Effect.ensuring(closeClient)); + if (activeSessionId.length === 0) { + yield* completeUnattachedOwnedSession; + return yield* runtimeError( + "create-session", + "invalid-response", + "The daemon create response omitted its active session identifier.", + ); + } + if ( + !/^[A-Za-z0-9_-]{1,256}$/.test(sessionId) || + primeAgentSessionFileName(sessionDir, sessionFile) === undefined || + (resumeSessionId !== undefined && sessionId !== resumeSessionId) + ) { + yield* completeUnattachedOwnedSession; + return yield* runtimeError( + "create-session", + "invalid-response", + "The daemon create response did not match the isolated durable session identity.", + ); + } + connection = yield* Effect.tryPromise({ + try: () => + input.manager.bridge.DaemonAgentConnection.attach(client, activeSessionId, { + closeClientOnDispose: false, + supportsExtensionUi: true, + ownedSession: true, + ownedSessionRecoveryConfig: sessionRuntimeConfig, + ...(input.disableAutoReconnect === true + ? {} + : { recoverDaemon: input.manager.recover }), + }), + catch: () => + runtimeError( + "attach-session", + "request-failed", + "Could not attach to the created daemon session.", + ), + }).pipe(Effect.onError(() => completeUnattachedOwnedSession)); } - const completeUnattachedOwnedSession = Effect.tryPromise({ - try: () => - client.request({ type: "complete_owned_session", activeSessionId }, COMMAND_TIMEOUT_MS), - catch: () => undefined, - }).pipe(Effect.ignore, Effect.ensuring(closeClient)); - - const sessionId = created.value.data.sessionId.trim(); - const sessionFile = created.value.data.sessionFile.trim(); + if ( + activeSessionId.length === 0 || !/^[A-Za-z0-9_-]{1,256}$/.test(sessionId) || primeAgentSessionFileName(sessionDir, sessionFile) === undefined || (resumeSessionId !== undefined && sessionId !== resumeSessionId) ) { - yield* completeUnattachedOwnedSession; + yield* Effect.promise(() => connection?.dispose().catch(() => undefined)); + yield* closeClient; + releaseManagerRecoveryRetention?.(); return yield* runtimeError( "create-session", "invalid-response", - "The daemon create response did not match the isolated durable session identity.", + "The daemon session authority did not match the isolated durable session identity.", ); } - connection = yield* Effect.tryPromise({ - try: () => - input.manager.bridge.DaemonAgentConnection.attach(client, activeSessionId, { - closeClientOnDispose: false, - supportsExtensionUi: true, - ownedSession: true, - ownedSessionRecoveryConfig: sessionRuntimeConfig, - ...(input.disableAutoReconnect === true ? {} : { recoverDaemon: input.manager.recover }), - }), - catch: () => - runtimeError( - "attach-session", - "request-failed", - "Could not attach to the created daemon session.", - ), - }).pipe(Effect.onError(() => completeUnattachedOwnedSession)); - const closeUnusableAttachedConnection = Effect.promise(async () => { await connection?.dispose().catch(() => undefined); client.close(); @@ -1957,7 +2182,9 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo interface QueuedRuntimeEvent { readonly event: PrimeDaemonEvent; readonly weight: number; + readonly recoveryCursor?: PrimeAgentDaemonEventCursor; } + const consumedRecoveryCursors = new WeakMap(); let queuedRuntimeEventWeight = 0; let runtimeEventIngressFailed = false; let runtimeEventCapacityFailed = false; @@ -2097,6 +2324,7 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo onCommit: () => void = () => undefined, ordinaryIngressFence?: OrdinaryIngressFence, providerRouteRetirement?: ProviderRouteRetirement, + recoveryCursor?: PrimeAgentDaemonEventCursor, ) => { const routeIsCurrent = () => ordinaryIngressFenceIsCurrent(ordinaryIngressFence) && @@ -2107,7 +2335,11 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo return proofEpoch === undefined ? Effect.void : Effect.fail(CORRELATED_PROOF_FENCE_RETIRED); } const weight = boundedCorrelatedProofRouteWeight(event); - const queued = { event, weight } satisfies QueuedRuntimeEvent; + const queued = { + event, + weight, + ...(recoveryCursor === undefined ? {} : { recoveryCursor }), + } satisfies QueuedRuntimeEvent; if (event._tag === "SessionClosed") { runtimeEventIngressFailed = true; settleReconnectResolution(connectionGeneration, false); @@ -2781,6 +3013,19 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo ) { return Effect.void; } + const rawRecoveryCursor = + Predicate.isObject(raw) && + Predicate.isObject(raw.meta) && + Predicate.isObject(raw.meta.cursor) && + typeof raw.meta.cursor.generation === "string" && + typeof raw.meta.cursor.sequence === "number" && + Number.isSafeInteger(raw.meta.cursor.sequence) && + raw.meta.cursor.sequence >= 0 + ? { + generation: raw.meta.cursor.generation, + sequence: raw.meta.cursor.sequence, + } + : undefined; let decoded = safeEvent( decodePrimeAgentDaemonEvent(raw, { correlatedPromptLifecycle: correlatedPromptLifecycleAvailable, @@ -3069,6 +3314,7 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo }, ordinaryIngressFence, providerRouteRetirement, + rawRecoveryCursor, ); }; const routeRawEvent = ( @@ -7437,6 +7683,71 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo } satisfies PrimeAgentDaemonSessionStats; }); + const recoveryCursor = adoptedRecovery?.proof.cursor ?? recoveryCursorFromSnapshot(rawSnapshot); + if (createdRecovery !== undefined) { + const creationRecovery = input.recovery; + if (recoveryCursor === undefined || creationRecovery?.kind !== "create") { + unsubscribe(); + yield* Effect.promise(() => connection!.dispose().catch(() => undefined)); + client.close(); + releaseManagerRecoveryRetention?.(); + return yield* runtimeError( + "initial-snapshot", + "invalid-response", + "Recoverable Prime Agent execution omitted its authoritative event cursor.", + ); + } + yield* Effect.tryPromise({ + try: () => + creationRecovery.onAuthorityReady({ + ...createdRecovery, + activeSessionId, + sessionId, + sessionFile, + cursor: recoveryCursor, + recoveryConfig: sessionRuntimeConfig, + launchEnvironment: input.manager.launchEnvironment ?? {}, + daemonCapabilities: [...(client.hello?.serverCapabilities ?? [])], + schemaRevision: client.hello?.schemaRevision ?? 0, + }), + catch: () => + runtimeError( + "create-session", + "request-failed", + "Recoverable Prime Agent authority could not be durably recorded.", + ), + }).pipe( + Effect.onError(() => + Effect.promise(async () => { + await connection + ?.disposeOwnedSession?.({ timeoutMs: COMMAND_TIMEOUT_MS }) + .catch(() => undefined); + client.close(); + releaseManagerRecoveryRetention?.(); + }), + ), + ); + } + + if (confirmAdoption !== undefined) { + yield* Effect.tryPromise({ + try: confirmAdoption, + catch: () => + runtimeError( + "attach-session", + "request-failed", + "Recoverable Prime Agent ownership could not be durably confirmed.", + ), + }).pipe( + Effect.onError(() => + Effect.promise(async () => { + await connection?.dispose().catch(() => undefined); + client.close(); + }), + ), + ); + } + // The initial proof cannot survive any overlapping attachment generation, even // if a newer proof has already appeared by the end of these asynchronous reads. if ( @@ -7461,6 +7772,7 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo yield* offerBackpressuredRuntimeEvent({ event: initialRuntimeEvent, weight: boundedCorrelatedProofRouteWeight(initialRuntimeEvent), + ...(recoveryCursor === undefined ? {} : { recoveryCursor }), }).pipe( Effect.asVoid, Effect.catch(() => @@ -7498,28 +7810,91 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo // synchronous assignment. Later events enter the normal bounded queue path. initializing = false; + const retireLocalOwner = Effect.sync(() => { + retireOrdinaryIngressFence(); + retireCurrentOrdinaryWorkerCloseRoute(); + retireProviderRoute(correlatedProviderRouteRetirement); + settleReconnectResolution(connectionGeneration, false); + settleManagedRecovery(managedRecoveryResolution, false); + mcpRecoveryPending = false; + mcpRecoveryFailed = true; + settleQuiescenceMcpRecovery(quiescenceMcpRecovery, false); + const workerRecovery = activeWorkerRecovery; + if (workerRecovery !== undefined) { + workerRecovery.provisionalSnapshot = undefined; + settleReconnectResolution(workerRecovery.resolution.generation, false); + if (activeWorkerRecovery === workerRecovery) activeWorkerRecovery = undefined; + } + unsubscribe?.(); + }); + + const detach = Effect.uninterruptible( + Effect.gen(function* () { + if (detached || disposed) return; + detached = true; + disposeStarted = true; + yield* retireLocalOwner; + yield* closeClient; + yield* Queue.shutdown(eventQueue); + yield* Queue.shutdown(runtimeEventWeightCapacityAvailable); + yield* Deferred.succeed(disposeCompletion, undefined).pipe(Effect.ignore); + }), + ); + const dispose = Effect.uninterruptibleMask((restore) => { // Scope cleanup and explicit teardown can race. Every caller joins the // first bounded native disposal instead of treating "started" as done. if (disposeStarted) return restore(Deferred.await(disposeCompletion)); disposeStarted = true; - const nativeDispose = Effect.tryPromise({ - try: () => connection!.dispose(), - catch: () => - runtimeError("dispose", "request-failed", "Could not dispose the daemon session."), - }).pipe( - Effect.flatMap((output) => - output === undefined - ? Effect.void - : Effect.fail( - runtimeError( - "dispose", - "invalid-response", - "The daemon dispose operation returned an invalid response.", - ), + const recoveryOwned = input.recovery !== undefined; + const nativeDispose = recoveryOwned + ? Effect.tryPromise({ + try: async () => { + const cleanup = connection?.disposeOwnedSession; + if (!Predicate.isFunction(cleanup)) { + throw new Error("authoritative cleanup is unavailable"); + } + const result = await cleanup.call(connection, { timeoutMs: COMMAND_TIMEOUT_MS }); + if ( + !Predicate.isObject(result) || + (result.status !== "completed" && result.status !== "already_completed") + ) { + throw new Error("authoritative cleanup was not proven"); + } + }, + catch: () => + runtimeError( + "dispose", + "request-failed", + "Prime Agent could not prove authoritative native cleanup.", ), - ), + }).pipe( + Effect.tap(() => + Effect.sync(() => { + releaseManagerRecoveryRetention?.(); + releaseManagerRecoveryRetention = undefined; + }), + ), + ) + : Effect.tryPromise({ + try: () => connection!.dispose(), + catch: () => + runtimeError("dispose", "request-failed", "Could not dispose the daemon session."), + }).pipe( + Effect.flatMap((output) => + output === undefined + ? Effect.void + : Effect.fail( + runtimeError( + "dispose", + "invalid-response", + "The daemon dispose operation returned an invalid response.", + ), + ), + ), + ); + const boundedNativeDispose = nativeDispose.pipe( Effect.timeoutOrElse({ duration: COMMAND_TIMEOUT_MS, orElse: () => @@ -7530,37 +7905,21 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo ), }), ); - const beginDisposeOwner = Effect.sync(() => { - retireOrdinaryIngressFence(); - retireCurrentOrdinaryWorkerCloseRoute(); - retireProviderRoute(correlatedProviderRouteRetirement); - settleReconnectResolution(connectionGeneration, false); - settleManagedRecovery(managedRecoveryResolution, false); - mcpRecoveryPending = false; - mcpRecoveryFailed = true; - settleQuiescenceMcpRecovery(quiescenceMcpRecovery, false); - const workerRecovery = activeWorkerRecovery; - if (workerRecovery !== undefined) { - workerRecovery.provisionalSnapshot = undefined; - settleReconnectResolution(workerRecovery.resolution.generation, false); - if (activeWorkerRecovery === workerRecovery) activeWorkerRecovery = undefined; - } - unsubscribe?.(); - return [...activePrivateSideQuestions.entries()]; - }); + const beginDisposeOwner = retireLocalOwner.pipe( + Effect.map(() => [...activePrivateSideQuestions.entries()] as const), + ); const disposeOwnerBody = ( nativeSideQuestions: ReadonlyArray, ) => Effect.forEach(nativeSideQuestions, ([nativeId, active]) => bestEffortAbortSideQuestion(nativeId, active), - ).pipe(Effect.andThen(failActivePrivateSideQuestions()), Effect.andThen(releaseMcpServer)); + ).pipe( + Effect.andThen(failActivePrivateSideQuestions()), + Effect.andThen(recoveryOwned ? Effect.void : releaseMcpServer), + ); return beginDisposeOwner.pipe( - // Election stays masked through synchronous route retirement and unsubscribe. - // Only the remaining owner body observes an interrupt pending since election. Effect.flatMap((nativeSideQuestions) => restore(disposeOwnerBody(nativeSideQuestions))), - // These finalizers must be installed outside restore: a pending interrupt may - // prevent the restored body from starting, but it cannot skip native teardown. - Effect.onExit(() => nativeDispose), + Effect.onExit(() => boundedNativeDispose), Effect.ensuring( Effect.gen(function* () { disposed = true; @@ -7572,12 +7931,11 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo yield* Queue.shutdown(runtimeEventWeightCapacityAvailable); }), ), - // Publish only the final Exit after every cleanup and combined cleanup defect. Effect.onExit((exit) => Deferred.done(disposeCompletion, exit).pipe(Effect.ignore)), ); }); - yield* Effect.addFinalizer(() => dispose.pipe(Effect.ignore)); + yield* Effect.addFinalizer(() => (detached ? Effect.void : dispose.pipe(Effect.ignore))); return { resumeCursor: PRIME_AGENT_DAEMON_RESUME_CURSOR, @@ -7615,6 +7973,9 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo events: Stream.fromQueue(eventQueue).pipe( Stream.map((queued) => { releaseQueuedRuntimeEventWeight(queued.weight); + if (queued.recoveryCursor !== undefined) { + consumedRecoveryCursors.set(queued.event, queued.recoveryCursor); + } return queued.event; }), ), @@ -7680,6 +8041,14 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo setServiceTier, respondToExtensionUiRequest, getSessionStats, + ...(input.recovery === undefined + ? {} + : { + recoveryCorrelationId: input.recovery.correlationId, + ...(recoveryCursor === undefined ? {} : { recoveryCursor }), + }), + recoveryCursorForEvent: (event) => consumedRecoveryCursors.get(event), + detach, dispose, } satisfies PrimeAgentDaemonSessionRuntime; }, diff --git a/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.test.ts b/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.test.ts new file mode 100644 index 000000000..3cebfc290 --- /dev/null +++ b/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.test.ts @@ -0,0 +1,139 @@ +import { assert, it } from "@effect/vitest"; +import * as Effect from "effect/Effect"; +import * as Exit from "effect/Exit"; +import * as Layer from "effect/Layer"; +import * as Option from "effect/Option"; + +import migration050 from "../../persistence/Migrations/050_PrimeAgentRecoveryLedger.ts"; +import * as NodeSqliteClient from "../../persistence/NodeSqliteClient.ts"; +import { make, type PrimeAgentRecoveryAuthority } from "./PrimeAgentRecoveryLedger.ts"; + +const layer = it.layer(Layer.mergeAll(NodeSqliteClient.layerMemory())); + +const authority = { + threadId: "thread-recovery", + providerInstanceId: "primeAgent", + sessionIncarnationId: "incarnation-1", + admissionRequestId: "admission-1", + turnId: null, + packageRoot: "/private/prime-package", + packageVersion: "0.5.1", + managedBuildId: "managed-build-1", + sdkFeatures: [ + "recoverable_owned_session_adoption_v1", + "caller_owned_session_environment_cleanup_v1", + ], + daemonCapabilities: [ + "daemon_recoverable_owned_session_adoption_v1", + "caller_owned_session_environment_cleanup_v1", + "authoritative_owned_session_cleanup_v1", + ], + protocolName: "prime-agent.daemon", + protocolVersion: 4, + schemaRevision: 30, + activeSessionId: "active-1", + nativeSessionId: "native-1", + recoveryHandle: "private-handle-1", + supervisorGeneration: "supervisor-1", + ownershipGeneration: 0, + cursor: { generation: "events-1", sequence: 7 }, + correlationId: "correlation-1", + mcpOwnerId: "pylon:mcp-1", + recoveryConfig: { cwd: "/private/worktree", model: "anthropic/claude" }, + launchEnvironment: { HOME: "/private/home", PRIME_TOKEN: "private-token" }, + transcriptMessageCount: 2, + transcriptFingerprints: ["fingerprint-1", "fingerprint-2"], + ownerToken: "owner-1", + state: "prepared", + nativeCleanupProven: false, + terminalProjected: false, + checkpointQuiesced: false, + updatedAt: "2026-01-01T00:00:00.000Z", +} satisfies PrimeAgentRecoveryAuthority; + +layer("PrimeAgentRecoveryLedger", (it) => { + it.effect("CAS-claims one owner and deletes only after all three cleanup proofs", () => + Effect.gen(function* () { + yield* migration050; + const ledger = yield* make; + yield* ledger.putPrepared(authority); + const replacement = yield* Effect.exit( + ledger.putPrepared({ ...authority, ownerToken: "owner-replacement" }), + ); + assert.isTrue(Exit.isFailure(replacement)); + assert.equal( + Option.getOrThrow(yield* ledger.get(authority.threadId)).ownerToken, + authority.ownerToken, + ); + assert.isTrue( + yield* ledger.markAdmitted({ + threadId: authority.threadId, + ownerToken: authority.ownerToken, + turnId: "turn-1", + updatedAt: "2026-01-01T00:00:01.000Z", + }), + ); + + const firstClaim = yield* ledger.claim({ + threadId: authority.threadId, + expectedOwnerToken: authority.ownerToken, + nextOwnerToken: "owner-2", + updatedAt: "2026-01-01T00:00:02.000Z", + }); + const competingClaim = yield* ledger.claim({ + threadId: authority.threadId, + expectedOwnerToken: authority.ownerToken, + nextOwnerToken: "owner-3", + updatedAt: "2026-01-01T00:00:02.000Z", + }); + assert.isTrue(Option.isSome(firstClaim)); + assert.isTrue(Option.isNone(competingClaim)); + + assert.isTrue( + yield* ledger.commitAdoption({ + threadId: authority.threadId, + ownerToken: "owner-2", + recoveryHandle: "private-handle-2", + ownershipGeneration: 1, + cursor: { generation: "events-1", sequence: 11 }, + mcpOwnerId: "pylon:mcp-2", + updatedAt: "2026-01-01T00:00:03.000Z", + }), + ); + assert.isTrue( + yield* ledger.updateTranscriptProgress({ + threadId: authority.threadId, + ownerToken: "owner-2", + cursor: { generation: "events-1", sequence: 19 }, + messageCount: 3, + fingerprints: ["fingerprint-1", "fingerprint-2", "fingerprint-3"], + updatedAt: "2026-01-01T00:00:04.000Z", + }), + ); + const adopted = Option.getOrThrow(yield* ledger.get(authority.threadId)); + assert.equal(adopted.recoveryHandle, "private-handle-2"); + assert.deepEqual(adopted.cursor, { generation: "events-1", sequence: 19 }); + assert.equal(adopted.transcriptMessageCount, 3); + + assert.isFalse(yield* ledger.deleteIfSettled(authority.threadId)); + assert.isTrue( + yield* ledger.markNativeCleanup({ + threadId: authority.threadId, + ownerToken: "owner-2", + updatedAt: "2026-01-01T00:00:05.000Z", + }), + ); + yield* ledger.markTerminalProjected({ + threadId: authority.threadId, + updatedAt: "2026-01-01T00:00:06.000Z", + }); + assert.isFalse(yield* ledger.deleteIfSettled(authority.threadId)); + yield* ledger.markCheckpointQuiesced({ + threadId: authority.threadId, + updatedAt: "2026-01-01T00:00:07.000Z", + }); + assert.isTrue(yield* ledger.deleteIfSettled(authority.threadId)); + assert.isTrue(Option.isNone(yield* ledger.get(authority.threadId))); + }), + ); +}); diff --git a/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.ts b/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.ts new file mode 100644 index 000000000..7c133f4a5 --- /dev/null +++ b/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.ts @@ -0,0 +1,478 @@ +import * as Context from "effect/Context"; +import * as Effect from "effect/Effect"; +import * as Layer from "effect/Layer"; +import * as Option from "effect/Option"; +import * as Schema from "effect/Schema"; +import * as SqlClient from "effect/unstable/sql/SqlClient"; + +const NonNegativeInt = Schema.Number.pipe( + Schema.check(Schema.isInt(), Schema.isGreaterThanOrEqualTo(0)), +); + +const RecoveryCursor = Schema.Struct({ + generation: Schema.String, + sequence: NonNegativeInt, +}); + +export const PrimeAgentRecoveryAuthority = Schema.Struct({ + threadId: Schema.String, + providerInstanceId: Schema.String, + sessionIncarnationId: Schema.String, + admissionRequestId: Schema.String, + turnId: Schema.NullOr(Schema.String), + packageRoot: Schema.String, + packageVersion: Schema.String, + managedBuildId: Schema.String, + sdkFeatures: Schema.Array(Schema.String), + daemonCapabilities: Schema.Array(Schema.String), + protocolName: Schema.String, + protocolVersion: Schema.Int, + schemaRevision: Schema.Int, + activeSessionId: Schema.String, + nativeSessionId: Schema.String, + recoveryHandle: Schema.String, + supervisorGeneration: Schema.String, + ownershipGeneration: NonNegativeInt, + cursor: RecoveryCursor, + correlationId: Schema.String, + mcpOwnerId: Schema.String, + recoveryConfig: Schema.Record(Schema.String, Schema.Unknown), + launchEnvironment: Schema.Record(Schema.String, Schema.String), + transcriptMessageCount: NonNegativeInt, + transcriptFingerprints: Schema.Array(Schema.String), + ownerToken: Schema.String, + state: Schema.Literals(["prepared", "active", "adopting", "terminal"]), + nativeCleanupProven: Schema.Boolean, + terminalProjected: Schema.Boolean, + checkpointQuiesced: Schema.Boolean, + updatedAt: Schema.String, +}); +export type PrimeAgentRecoveryAuthority = typeof PrimeAgentRecoveryAuthority.Type; + +export class PrimeAgentRecoveryLedgerError extends Schema.TaggedErrorClass()( + "PrimeAgentRecoveryLedgerError", + { + operation: Schema.String, + cause: Schema.optional(Schema.Defect()), + }, +) { + override get message(): string { + return `Prime Agent recovery ledger failed during ${this.operation}.`; + } +} + +export interface PrimeAgentRecoveryLedgerShape { + readonly putPrepared: ( + authority: PrimeAgentRecoveryAuthority, + ) => Effect.Effect; + readonly get: ( + threadId: string, + ) => Effect.Effect, PrimeAgentRecoveryLedgerError>; + readonly listActive: () => Effect.Effect< + ReadonlyArray, + PrimeAgentRecoveryLedgerError + >; + readonly markAdmitted: (input: { + readonly threadId: string; + readonly ownerToken: string; + readonly turnId: string; + readonly updatedAt: string; + }) => Effect.Effect; + readonly discardPrepared: (input: { + readonly threadId: string; + readonly ownerToken: string; + }) => Effect.Effect; + readonly updateTranscriptProgress: (input: { + readonly threadId: string; + readonly ownerToken: string; + readonly cursor: typeof RecoveryCursor.Type; + readonly messageCount: number; + readonly fingerprints: ReadonlyArray; + readonly updatedAt: string; + }) => Effect.Effect; + /** Compare-and-swap the last durable owner. Exactly one restarted Pylon process can win. */ + readonly claim: (input: { + readonly threadId: string; + readonly expectedOwnerToken: string; + readonly nextOwnerToken: string; + readonly updatedAt: string; + }) => Effect.Effect, PrimeAgentRecoveryLedgerError>; + readonly releaseClaim: (input: { + readonly threadId: string; + readonly ownerToken: string; + readonly previousOwnerToken: string; + readonly updatedAt: string; + }) => Effect.Effect; + /** Persist the rotated bearer authority before the SDK confirmation step. */ + readonly commitAdoption: (input: { + readonly threadId: string; + readonly ownerToken: string; + readonly recoveryHandle: string; + readonly ownershipGeneration: number; + readonly cursor: typeof RecoveryCursor.Type; + readonly mcpOwnerId: string; + readonly updatedAt: string; + }) => Effect.Effect; + readonly markNativeCleanup: (input: { + readonly threadId: string; + readonly ownerToken: string; + readonly updatedAt: string; + }) => Effect.Effect; + readonly markTerminalProjected: (input: { + readonly threadId: string; + readonly updatedAt: string; + }) => Effect.Effect; + readonly markCheckpointQuiesced: (input: { + readonly threadId: string; + readonly updatedAt: string; + }) => Effect.Effect; + /** Deletes only after native cleanup, terminal projection, and checkpoint quiescence all hold. */ + readonly deleteIfSettled: ( + threadId: string, + ) => Effect.Effect; +} + +export class PrimeAgentRecoveryLedger extends Context.Service< + PrimeAgentRecoveryLedger, + PrimeAgentRecoveryLedgerShape +>()("t3/provider/prime/PrimeAgentRecoveryLedger") {} + +const RawRow = Schema.Struct({ + threadId: Schema.String, + providerInstanceId: Schema.String, + sessionIncarnationId: Schema.String, + admissionRequestId: Schema.String, + turnId: Schema.NullOr(Schema.String), + packageRoot: Schema.String, + packageVersion: Schema.String, + managedBuildId: Schema.String, + sdkFeaturesJson: Schema.String, + daemonCapabilitiesJson: Schema.String, + protocolName: Schema.String, + protocolVersion: Schema.Int, + schemaRevision: Schema.Int, + activeSessionId: Schema.String, + nativeSessionId: Schema.String, + recoveryHandle: Schema.String, + supervisorGeneration: Schema.String, + ownershipGeneration: Schema.Int, + cursorGeneration: Schema.String, + cursorSequence: Schema.Int, + correlationId: Schema.String, + mcpOwnerId: Schema.String, + recoveryConfigJson: Schema.String, + launchEnvironmentJson: Schema.String, + transcriptMessageCount: Schema.Int, + transcriptFingerprintsJson: Schema.String, + ownerToken: Schema.String, + state: Schema.String, + nativeCleanupProven: Schema.Int, + terminalProjected: Schema.Int, + checkpointQuiesced: Schema.Int, + updatedAt: Schema.String, +}); + +const decodeRawRows = Schema.decodeUnknownSync(Schema.Array(RawRow)); +const decodeAuthority = Schema.decodeUnknownSync(PrimeAgentRecoveryAuthority); +const selectColumns = ` + thread_id AS threadId, + provider_instance_id AS providerInstanceId, + session_incarnation_id AS sessionIncarnationId, + admission_request_id AS admissionRequestId, + turn_id AS turnId, + package_root AS packageRoot, + package_version AS packageVersion, + managed_build_id AS managedBuildId, + sdk_features_json AS sdkFeaturesJson, + daemon_capabilities_json AS daemonCapabilitiesJson, + protocol_name AS protocolName, + protocol_version AS protocolVersion, + schema_revision AS schemaRevision, + active_session_id AS activeSessionId, + native_session_id AS nativeSessionId, + recovery_handle AS recoveryHandle, + supervisor_generation AS supervisorGeneration, + ownership_generation AS ownershipGeneration, + cursor_generation AS cursorGeneration, + cursor_sequence AS cursorSequence, + correlation_id AS correlationId, + mcp_owner_id AS mcpOwnerId, + recovery_config_json AS recoveryConfigJson, + launch_environment_json AS launchEnvironmentJson, + transcript_message_count AS transcriptMessageCount, + transcript_fingerprints_json AS transcriptFingerprintsJson, + owner_token AS ownerToken, + state, + native_cleanup_proven AS nativeCleanupProven, + terminal_projected AS terminalProjected, + checkpoint_quiesced AS checkpointQuiesced, + updated_at AS updatedAt +`; + +function ledgerError(operation: string, cause?: unknown): PrimeAgentRecoveryLedgerError { + return new PrimeAgentRecoveryLedgerError({ + operation, + ...(cause === undefined ? {} : { cause }), + }); +} + +function decodeRows(rows: unknown, operation: string): ReadonlyArray { + try { + return decodeRawRows(rows).map((row) => + decodeAuthority({ + threadId: row.threadId, + providerInstanceId: row.providerInstanceId, + sessionIncarnationId: row.sessionIncarnationId, + admissionRequestId: row.admissionRequestId, + turnId: row.turnId, + packageRoot: row.packageRoot, + packageVersion: row.packageVersion, + managedBuildId: row.managedBuildId, + sdkFeatures: JSON.parse(row.sdkFeaturesJson), + daemonCapabilities: JSON.parse(row.daemonCapabilitiesJson), + protocolName: row.protocolName, + protocolVersion: row.protocolVersion, + schemaRevision: row.schemaRevision, + activeSessionId: row.activeSessionId, + nativeSessionId: row.nativeSessionId, + recoveryHandle: row.recoveryHandle, + supervisorGeneration: row.supervisorGeneration, + ownershipGeneration: row.ownershipGeneration, + cursor: { generation: row.cursorGeneration, sequence: row.cursorSequence }, + correlationId: row.correlationId, + mcpOwnerId: row.mcpOwnerId, + recoveryConfig: JSON.parse(row.recoveryConfigJson), + launchEnvironment: JSON.parse(row.launchEnvironmentJson), + transcriptMessageCount: row.transcriptMessageCount, + transcriptFingerprints: JSON.parse(row.transcriptFingerprintsJson), + ownerToken: row.ownerToken, + state: row.state, + nativeCleanupProven: row.nativeCleanupProven === 1, + terminalProjected: row.terminalProjected === 1, + checkpointQuiesced: row.checkpointQuiesced === 1, + updatedAt: row.updatedAt, + }), + ); + } catch (cause) { + throw ledgerError(operation, cause); + } +} + +export const make = Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + const mapSqlError = (operation: string) => (cause: unknown) => ledgerError(operation, cause); + const decodeEffect = (rows: unknown, operation: string) => + Effect.try({ + try: () => decodeRows(rows, operation), + catch: (cause) => + Schema.is(PrimeAgentRecoveryLedgerError)(cause) ? cause : ledgerError(operation, cause), + }); + + const queryByThread = (threadId: string) => + sql + .unsafe(`SELECT ${selectColumns} FROM prime_agent_recovery_ledger WHERE thread_id = ?`, [ + threadId, + ]) + .pipe( + Effect.mapError(mapSqlError("get")), + Effect.flatMap((rows) => decodeEffect(rows, "get")), + ); + + const get: PrimeAgentRecoveryLedgerShape["get"] = (threadId) => + queryByThread(threadId).pipe(Effect.map((rows) => Option.fromNullishOr(rows[0]))); + + const putPrepared: PrimeAgentRecoveryLedgerShape["putPrepared"] = (authority) => + sql + .unsafe( + `INSERT INTO prime_agent_recovery_ledger ( + thread_id, provider_instance_id, session_incarnation_id, admission_request_id, turn_id, + package_root, package_version, managed_build_id, sdk_features_json, daemon_capabilities_json, + protocol_name, protocol_version, schema_revision, active_session_id, native_session_id, + recovery_handle, supervisor_generation, ownership_generation, cursor_generation, cursor_sequence, + correlation_id, mcp_owner_id, recovery_config_json, launch_environment_json, + transcript_message_count, transcript_fingerprints_json, owner_token, state, + native_cleanup_proven, terminal_projected, checkpoint_quiesced, updated_at + ) VALUES (${Array.from({ length: 32 }, () => "?").join(",")}) +`, + [ + authority.threadId, + authority.providerInstanceId, + authority.sessionIncarnationId, + authority.admissionRequestId, + authority.turnId, + authority.packageRoot, + authority.packageVersion, + authority.managedBuildId, + JSON.stringify(authority.sdkFeatures), + JSON.stringify(authority.daemonCapabilities), + authority.protocolName, + authority.protocolVersion, + authority.schemaRevision, + authority.activeSessionId, + authority.nativeSessionId, + authority.recoveryHandle, + authority.supervisorGeneration, + authority.ownershipGeneration, + authority.cursor.generation, + authority.cursor.sequence, + authority.correlationId, + authority.mcpOwnerId, + JSON.stringify(authority.recoveryConfig), + JSON.stringify(authority.launchEnvironment), + authority.transcriptMessageCount, + JSON.stringify(authority.transcriptFingerprints), + authority.ownerToken, + authority.state, + authority.nativeCleanupProven ? 1 : 0, + authority.terminalProjected ? 1 : 0, + authority.checkpointQuiesced ? 1 : 0, + authority.updatedAt, + ], + ) + .pipe(Effect.mapError(mapSqlError("putPrepared")), Effect.asVoid); + + const listActive: PrimeAgentRecoveryLedgerShape["listActive"] = () => + sql + .unsafe( + `SELECT ${selectColumns} FROM prime_agent_recovery_ledger WHERE state IN ('active','adopting') ORDER BY updated_at, thread_id`, + ) + .pipe( + Effect.mapError(mapSqlError("listActive")), + Effect.flatMap((rows) => decodeEffect(rows, "listActive")), + ); + + const conditionalUpdate = ( + operation: string, + statement: string, + parameters: ReadonlyArray, + ) => + sql.unsafe(statement, parameters).pipe( + Effect.mapError(mapSqlError(operation)), + Effect.map((rows) => Array.isArray(rows) && rows.length === 1), + ); + + const markAdmitted: PrimeAgentRecoveryLedgerShape["markAdmitted"] = (input) => + conditionalUpdate( + "markAdmitted", + `UPDATE prime_agent_recovery_ledger SET turn_id=?, state='active', updated_at=? + WHERE thread_id=? AND owner_token=? AND state='prepared' RETURNING thread_id`, + [input.turnId, input.updatedAt, input.threadId, input.ownerToken], + ); + + const discardPrepared: PrimeAgentRecoveryLedgerShape["discardPrepared"] = (input) => + conditionalUpdate( + "discardPrepared", + `DELETE FROM prime_agent_recovery_ledger + WHERE thread_id=? AND owner_token=? AND state='prepared' RETURNING thread_id`, + [input.threadId, input.ownerToken], + ); + + const updateTranscriptProgress: PrimeAgentRecoveryLedgerShape["updateTranscriptProgress"] = ( + input, + ) => + conditionalUpdate( + "updateTranscriptProgress", + `UPDATE prime_agent_recovery_ledger + SET cursor_generation=?, cursor_sequence=?, transcript_message_count=?, + transcript_fingerprints_json=?, updated_at=? + WHERE thread_id=? AND owner_token=? AND state IN ('active','adopting') RETURNING thread_id`, + [ + input.cursor.generation, + input.cursor.sequence, + input.messageCount, + JSON.stringify(input.fingerprints), + input.updatedAt, + input.threadId, + input.ownerToken, + ], + ); + + const claim: PrimeAgentRecoveryLedgerShape["claim"] = (input) => + conditionalUpdate( + "claim", + `UPDATE prime_agent_recovery_ledger SET owner_token=?, state='adopting', updated_at=? + WHERE thread_id=? AND owner_token=? AND state='active' RETURNING thread_id`, + [input.nextOwnerToken, input.updatedAt, input.threadId, input.expectedOwnerToken], + ).pipe( + Effect.flatMap((claimed) => (claimed ? get(input.threadId) : Effect.succeed(Option.none()))), + ); + + const releaseClaim: PrimeAgentRecoveryLedgerShape["releaseClaim"] = (input) => + conditionalUpdate( + "releaseClaim", + `UPDATE prime_agent_recovery_ledger SET owner_token=?, state='active', updated_at=? + WHERE thread_id=? AND owner_token=? AND state='adopting' RETURNING thread_id`, + [input.previousOwnerToken, input.updatedAt, input.threadId, input.ownerToken], + ); + + const commitAdoption: PrimeAgentRecoveryLedgerShape["commitAdoption"] = (input) => + conditionalUpdate( + "commitAdoption", + `UPDATE prime_agent_recovery_ledger + SET recovery_handle=?, ownership_generation=?, cursor_generation=?, cursor_sequence=?, + mcp_owner_id=?, state='active', updated_at=? + WHERE thread_id=? AND owner_token=? AND state='adopting' RETURNING thread_id`, + [ + input.recoveryHandle, + input.ownershipGeneration, + input.cursor.generation, + input.cursor.sequence, + input.mcpOwnerId, + input.updatedAt, + input.threadId, + input.ownerToken, + ], + ); + + const markNativeCleanup: PrimeAgentRecoveryLedgerShape["markNativeCleanup"] = (input) => + conditionalUpdate( + "markNativeCleanup", + `UPDATE prime_agent_recovery_ledger + SET native_cleanup_proven=1, state='terminal', updated_at=? + WHERE thread_id=? AND owner_token=? RETURNING thread_id`, + [input.updatedAt, input.threadId, input.ownerToken], + ); + + const markTerminalProjected: PrimeAgentRecoveryLedgerShape["markTerminalProjected"] = (input) => + sql + .unsafe( + `UPDATE prime_agent_recovery_ledger SET terminal_projected=1, updated_at=? WHERE thread_id=?`, + [input.updatedAt, input.threadId], + ) + .pipe(Effect.mapError(mapSqlError("markTerminalProjected")), Effect.asVoid); + + const markCheckpointQuiesced: PrimeAgentRecoveryLedgerShape["markCheckpointQuiesced"] = (input) => + sql + .unsafe( + `UPDATE prime_agent_recovery_ledger SET checkpoint_quiesced=1, updated_at=? WHERE thread_id=?`, + [input.updatedAt, input.threadId], + ) + .pipe(Effect.mapError(mapSqlError("markCheckpointQuiesced")), Effect.asVoid); + + const deleteIfSettled: PrimeAgentRecoveryLedgerShape["deleteIfSettled"] = (threadId) => + conditionalUpdate( + "deleteIfSettled", + `DELETE FROM prime_agent_recovery_ledger + WHERE thread_id=? AND native_cleanup_proven=1 AND terminal_projected=1 AND checkpoint_quiesced=1 + RETURNING thread_id`, + [threadId], + ); + + return { + putPrepared, + get, + listActive, + markAdmitted, + discardPrepared, + updateTranscriptProgress, + claim, + releaseClaim, + commitAdoption, + markNativeCleanup, + markTerminalProjected, + markCheckpointQuiesced, + deleteIfSettled, + } satisfies PrimeAgentRecoveryLedgerShape; +}); + +export const layer = Layer.effect(PrimeAgentRecoveryLedger, make); diff --git a/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs b/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs new file mode 100644 index 000000000..9a34d0ff0 --- /dev/null +++ b/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs @@ -0,0 +1,194 @@ +import * as ChildProcess from "node:child_process"; +import { randomUUID } from "node:crypto"; +import * as FSP from "node:fs/promises"; +import * as OS from "node:os"; +import * as Path from "node:path"; +import { pathToFileURL } from "node:url"; + +import { describe, expect, it } from "vite-plus/test"; + +const packageRoot = process.env.PRIME_AGENT_RECOVERY_REAL_PACKAGE_ROOT; +const exactHead = "507a52239d3ace7bb2b2965ade7779988fdb6344"; + +const waitForExit = (child, timeoutMs) => + new Promise((resolve, reject) => { + const timer = setTimeout(() => reject(new Error("subprocess exit timed out")), timeoutMs); + child.once("exit", (code, signal) => { + clearTimeout(timer); + resolve({ code, signal }); + }); + }); + +const runCaptured = (command, args, options, timeoutMs) => + new Promise((resolve, reject) => { + const child = ChildProcess.spawn(command, args, options); + const stdout = []; + const stderr = []; + let stdoutBytes = 0; + let stderrBytes = 0; + const maximumBytes = 1024 * 1024; + const timer = setTimeout(() => { + child.kill("SIGTERM"); + reject(new Error("owner subprocess timed out")); + }, timeoutMs); + child.stdout.on("data", (chunk) => { + stdoutBytes += chunk.length; + if (stdoutBytes > maximumBytes) child.kill("SIGTERM"); + else stdout.push(chunk); + }); + child.stderr.on("data", (chunk) => { + stderrBytes += chunk.length; + if (stderrBytes > maximumBytes) child.kill("SIGTERM"); + else stderr.push(chunk); + }); + child.once("error", (error) => { + clearTimeout(timer); + reject(error); + }); + child.once("exit", (code) => { + clearTimeout(timer); + const output = Buffer.concat(stdout).toString("utf8"); + const errorOutput = Buffer.concat(stderr).toString("utf8"); + if (code === 0) resolve(output); + else reject(new Error(`owner subprocess failed (${code}): ${errorOutput}`)); + }); + }); + +describe.skipIf(!packageRoot)("Prime Agent exact-head restart adoption subprocess", () => { + it("adopts after the creating owner process exits and proves authoritative cleanup", async () => { + const gitHead = ChildProcess.execFileSync("git", ["-C", packageRoot, "rev-parse", "HEAD"], { + encoding: "utf8", + timeout: 5_000, + }).trim(); + expect(gitHead).toBe(exactHead); + + const codingAgentRoot = Path.join(packageRoot, "packages", "coding-agent"); + const sdkEntry = Path.join(codingAgentRoot, "dist", "index.js"); + const cliEntry = Path.join(codingAgentRoot, "dist", "bundle", "cli.js"); + const temp = await FSP.mkdtemp(Path.join(OS.tmpdir(), "pylon-prime-restart-")); + const socket = Path.join(temp, "daemon.sock"); + const sessionDir = Path.join(temp, "sessions"); + const agentDir = Path.join(temp, "agent-home"); + await FSP.mkdir(sessionDir, { recursive: true }); + await FSP.mkdir(agentDir, { recursive: true }); + const launchEnv = { + HOME: process.env.HOME ?? temp, + PATH: process.env.PATH ?? "/usr/bin:/bin", + PRIME_AGENT_CODING_AGENT_DIR: agentDir, + }; + const daemon = ChildProcess.spawn( + process.execPath, + [ + cliEntry, + "--mode", + "daemon", + "--daemon-socket", + socket, + "--offline", + "--session-dir", + sessionDir, + ], + { env: launchEnv, stdio: ["ignore", "ignore", "pipe"] }, + ); + const daemonErrors = []; + daemon.stderr.on("data", (chunk) => daemonErrors.push(chunk)); + + try { + const helper = Path.join(temp, "create-owner.mjs"); + await FSP.writeFile( + helper, + `import { randomUUID } from "node:crypto"; +import { DaemonClient, createRecoverableOwnedSession } from ${JSON.stringify(pathToFileURL(sdkEntry).href)}; +const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); +const client = new DaemonClient(${JSON.stringify(socket)}); +let connected = false; +for (let attempt = 0; attempt < 80; attempt += 1) { + try { await client.connect(); connected = true; break; } catch { await sleep(25); } +} +if (!connected) throw new Error("daemon readiness timed out"); +await client.waitForHello(); +const config = ${JSON.stringify({ cwd: temp, sessionDir, noBuiltinTools: true, noExtensions: true, noSkills: true, noContextFiles: true })}; +const created = await createRecoverableOwnedSession(client, { + requestId: randomUUID(), correlationId: "correlation-real-1", mcpOwnerId: "pylon:mcp-real-1", + config, continueRecent: false, launchEnv: ${JSON.stringify(launchEnv)}, + connectionOptions: { closeClientOnDispose: false, supportsExtensionUi: true }, +}); +await created.connection.submitCorrelatedPrompt("/help", { + correlationId: "correlation-real-1", + queueIfBusy: true, +}); +for (let attempt = 0; attempt < 80; attempt += 1) { + const lifecycles = await created.connection.getPromptLifecycles(); + if (lifecycles.records?.some((entry) => entry.correlationId === "correlation-real-1") || + lifecycles.expired?.some((entry) => entry.correlationId === "correlation-real-1")) break; + await sleep(25); +} +const snapshot = await created.connection.getInitialSnapshot(); +process.stdout.write(JSON.stringify({ + recoveryHandle: created.recoveryHandle, + supervisorGeneration: created.supervisorGeneration, + activeSessionId: created.state.activeSessionId, + sessionId: created.state.sessionId, + cursor: snapshot.lastEventCursor, + config, +})); +client.close(); +`, + "utf8", + ); + + const authority = JSON.parse( + await runCaptured( + process.execPath, + [helper], + { env: launchEnv, stdio: ["ignore", "pipe", "pipe"] }, + 20_000, + ), + ); + expect(authority.cursor).toEqual( + expect.objectContaining({ generation: expect.any(String), sequence: expect.any(Number) }), + ); + + const sdk = await import(pathToFileURL(sdkEntry).href); + const client = new sdk.DaemonClient(socket); + await client.connect(); + const hello = await client.waitForHello(); + expect(hello.supervisorGeneration).toBe(authority.supervisorGeneration); + const adoptionRequestId = randomUUID(); + const adopted = await sdk.adoptRecoverableOwnedSession(client, { + requestId: adoptionRequestId, + recoveryHandle: authority.recoveryHandle, + expectedSupervisorGeneration: authority.supervisorGeneration, + activeSessionId: authority.activeSessionId, + sessionId: authority.sessionId, + correlationId: "correlation-real-1", + cursor: authority.cursor, + previousMcpOwnerId: "pylon:mcp-real-1", + mcpOwnerId: "pylon:mcp-real-2", + config: authority.config, + launchEnv, + connectionOptions: { closeClientOnDispose: false, supportsExtensionUi: true }, + }); + expect(adopted.recoveryHandle).not.toBe(authority.recoveryHandle); + expect(adopted.proof.ownershipGeneration).toBeGreaterThan(0); + await sdk.confirmRecoverableOwnedSessionAdoption(client, { + requestId: adoptionRequestId, + recoveryHandle: adopted.recoveryHandle, + proof: adopted.proof, + }); + const cleanup = await adopted.connection.disposeOwnedSession({ timeoutMs: 10_000 }); + expect(["completed", "already_completed"]).toContain(cleanup.status); + await client.request({ type: "shutdown" }, 5_000); + client.close(); + const exit = await waitForExit(daemon, 10_000); + expect(exit.code).toBe(0); + } finally { + if (daemon.exitCode === null && daemon.signalCode === null) daemon.kill("SIGTERM"); + await waitForExit(daemon, 5_000).catch(() => undefined); + await FSP.rm(temp, { recursive: true, force: true }); + if (daemon.exitCode && daemon.exitCode !== 0) { + throw new Error(Buffer.concat(daemonErrors).toString("utf8")); + } + } + }, 40_000); +}); diff --git a/apps/server/src/provider/prime/PrimeAgentRestartReplay.test.ts b/apps/server/src/provider/prime/PrimeAgentRestartReplay.test.ts new file mode 100644 index 000000000..a6fe41ac6 --- /dev/null +++ b/apps/server/src/provider/prime/PrimeAgentRestartReplay.test.ts @@ -0,0 +1,52 @@ +// @effect-diagnostics nodeBuiltinImport:off +import * as NodeCrypto from "node:crypto"; + +import { describe, expect, it } from "@effect/vitest"; + +import type { PrimeDaemonMessage } from "./PrimeAgentDaemonEvents.ts"; +import { planPrimeAgentRestartReplay } from "./PrimeAgentDaemonAdapter.ts"; + +const message = (index: number): PrimeDaemonMessage => ({ + role: "user", + timestamp: index, + text: `message-${index}`, + imageMimeTypes: [], + imageDigests: [], +}); +const fingerprint = (value: PrimeDaemonMessage) => + NodeCrypto.createHash("sha256").update(JSON.stringify(value), "utf8").digest("hex"); + +describe("planPrimeAgentRestartReplay", () => { + it("replays only the exact suffix after proving the overlap", () => { + const messages = [1, 2, 3, 4, 5].map(message); + const replay = planPrimeAgentRestartReplay({ + authorityMessageCount: 3, + authorityFingerprints: messages.slice(0, 3).map(fingerprint), + snapshotMessageCount: 5, + snapshotMessages: messages, + }); + expect(replay).toEqual({ valid: true, backlog: messages.slice(3) }); + }); + + it("fails closed on changed overlap or a transcript retention gap", () => { + const messages = [1, 2, 3, 4, 5].map(message); + expect( + planPrimeAgentRestartReplay({ + authorityMessageCount: 3, + authorityFingerprints: [messages[0]!, messages[1]!, message(30)].map(fingerprint), + snapshotMessageCount: 5, + snapshotMessages: messages, + }), + ).toEqual({ valid: false }); + + const longTranscript = Array.from({ length: 1_026 }, (_, index) => message(index)); + expect( + planPrimeAgentRestartReplay({ + authorityMessageCount: 1, + authorityFingerprints: [fingerprint(longTranscript[0]!)], + snapshotMessageCount: longTranscript.length, + snapshotMessages: longTranscript.slice(-1_024), + }), + ).toEqual({ valid: false }); + }); +}); diff --git a/apps/server/src/server.ts b/apps/server/src/server.ts index 234b50750..835d4eda9 100644 --- a/apps/server/src/server.ts +++ b/apps/server/src/server.ts @@ -31,6 +31,7 @@ import * as ServerLifecycleEvents from "./serverLifecycleEvents.ts"; import * as AnalyticsService from "./telemetry/AnalyticsService.ts"; import { ProviderSessionDirectoryLive } from "./provider/Layers/ProviderSessionDirectory.ts"; import * as ProviderSessionRuntime from "./persistence/ProviderSessionRuntime.ts"; +import * as PrimeAgentRecoveryLedger from "./provider/prime/PrimeAgentRecoveryLedger.ts"; import { ProviderAdapterRegistryLive } from "./provider/Layers/ProviderAdapterRegistry.ts"; import * as ModelManifest from "./provider/ModelManifest.ts"; import * as ProviderEventLoggers from "./provider/Layers/ProviderEventLoggers.ts"; @@ -276,7 +277,9 @@ const ProviderLayerLive = ProviderServiceLive.pipe( Layer.provideMerge(ProviderSessionDirectoryLayerLive), ); -const PersistenceLayerLive = Layer.empty.pipe(Layer.provideMerge(SqlitePersistenceLayerLive)); +const PersistenceLayerLive = PrimeAgentRecoveryLedger.layer.pipe( + Layer.provideMerge(SqlitePersistenceLayerLive), +); const VcsDriverRegistryLayerLive = VcsDriverRegistry.layer.pipe( Layer.provide(VcsProjectConfig.layer), diff --git a/apps/server/src/serverRuntimeStartup.reconcile.test.ts b/apps/server/src/serverRuntimeStartup.reconcile.test.ts index 0b5974384..4e8a2530b 100644 --- a/apps/server/src/serverRuntimeStartup.reconcile.test.ts +++ b/apps/server/src/serverRuntimeStartup.reconcile.test.ts @@ -43,10 +43,14 @@ const makeThread = ( }, }); -const makeProviderService = (liveThreadIds: ReadonlyArray = []) => +const makeProviderService = ( + liveThreadIds: ReadonlyArray = [], + recoverRestartSessions: () => Effect.Effect = () => Effect.void, +) => ({ startSession: () => Effect.die("unused"), sendTurn: () => Effect.die("unused"), + recoverRestartSessions, interruptTurn: () => Effect.die("unused"), respondToRequest: () => Effect.die("unused"), respondToUserInput: () => Effect.die("unused"), @@ -286,6 +290,48 @@ it.effect("retries failed projections and continues after a persistent failure", ); }); +it.effect("runs restart adoption before taking the orphan inventory", () => { + const recovered = makeThread("thread-recovered", "running", TurnId.make("turn-recovered")); + const liveThreadIds: ThreadId[] = []; + const order: string[] = []; + const provider = { + ...makeProviderService(liveThreadIds, () => + Effect.sync(() => { + order.push("recover"); + liveThreadIds.push(recovered.id); + }), + ), + listSessions: () => + Effect.sync(() => { + order.push("inventory"); + return liveThreadIds.map((threadId) => ({ threadId }) as never); + }), + }; + return ServerRuntimeStartup.reconcileProviderSessions.pipe( + Effect.provideService( + ProjectionSnapshotQuery.ProjectionSnapshotQuery, + queryWithThreads([recovered]), + ), + Effect.provideService(ProviderService.ProviderService, provider), + Effect.provideService(ProviderSessionDirectory.ProviderSessionDirectory, { + getBinding: () => Effect.die("recovered thread must not be orphaned"), + upsert: () => Effect.die("recovered thread must not be orphaned"), + removeExact: () => Effect.die("unused"), + getProvider: () => Effect.die("unused"), + listThreadIds: () => Effect.die("unused"), + listBindings: () => Effect.die("unused"), + }), + Effect.provideService(OrchestrationEngine.OrchestrationEngineService, { + readEvents: () => Stream.empty, + dispatch: () => Effect.die("recovered thread must not be orphaned"), + streamDomainEvents: Stream.empty, + latestSequence: Effect.succeed(0), + }), + Effect.provide(NodeServices.layer), + Effect.tap(() => Effect.sync(() => assert.deepStrictEqual(order, ["recover", "inventory"]))), + ); +}); + it.effect("does not fail startup when the live provider session inventory cannot be read", () => { let queried = false; return ServerRuntimeStartup.reconcileProviderSessions.pipe( diff --git a/apps/server/src/serverRuntimeStartup.ts b/apps/server/src/serverRuntimeStartup.ts index 3912dabc0..bbcc66466 100644 --- a/apps/server/src/serverRuntimeStartup.ts +++ b/apps/server/src/serverRuntimeStartup.ts @@ -303,6 +303,11 @@ export const reconcileProviderSessions = Effect.gen(function* () { const providerService = yield* ProviderService.ProviderService; const query = yield* ProjectionSnapshotQuery.ProjectionSnapshotQuery; + // Prime restart adoption must install exact incarnation fencing and release retained + // replay before generic orphan settlement can observe the thread as dead. + if (providerService.recoverRestartSessions !== undefined) { + yield* providerService.recoverRestartSessions(); + } const liveThreadIds = new Set( (yield* providerService.listSessions()).map((session) => session.threadId), ); diff --git a/docs/internals/prime-agent-daemon-parity.md b/docs/internals/prime-agent-daemon-parity.md index 1ebffe0dd..59fffa415 100644 --- a/docs/internals/prime-agent-daemon-parity.md +++ b/docs/internals/prime-agent-daemon-parity.md @@ -2,7 +2,7 @@ This ledger records Pylon's treatment of the public `DaemonAgentConnection` surface shipped by Prime Agent 0.8.1 (daemon protocol 7, schema 22) and Pylon's optional fork extension at protocol 7, -schema 24. Parity here means that every useful public outcome is either integrated through a typed +schema 30. Parity here means that every useful public outcome is either integrated through a typed provider-neutral contract or has an explicit product and safety decision. It does not mean exposing a raw method tunnel. @@ -207,11 +207,35 @@ modes because bounded plan progress reaches Pylon. It does not advertise `propos hidden instead of synthesizing one with a hidden prompt. Any transcript mismatch, incomplete streaming snapshot, MCP reattachment failure, or unvalidated barrier -fails the canonical turn once and disposes the uncertain native session. This recovery is in-memory only; -a Pylon server restart does not adopt native execution that outlives the server process. Instead, a -replacement waits beyond Prime's bounded client-owned disconnect grace for the daemon to release the old -worker, then recreates the exact saved session. It never steals ownership from a live client; if ownership -remains active after that bounded wait, Pylon preserves the structured `SessionAlreadyActiveError`. When +fails the canonical turn once and disposes the uncertain native session. + +### Pylon process restart recovery + +An admitted Full access turn can outlive the Pylon server process only when a Pylon-managed Prime +publication, its exact package root and build, the supported macOS/Linux host architecture, and the same +retained supervisor generation all match the private recovery ledger. The ledger also binds the provider +instance and Pylon session incarnation to Prime's protocol, schema, capabilities, active/native session, +correlation, cursor, recovery configuration, exact launch environment, MCP owner, and bounded transcript +progress. It is server-private: recovery handles, native identifiers, correlations, cursors, snapshots, +paths, prompts, tool data, and transport errors never enter public contracts, events, receipts, or logs. + +Preparation writes native ownership before prompt submission and records the exact turn only after Prime +admits that prompt. At startup, a replacement claims ownership with a SQLite compare-and-swap before SDK +adoption. It rotates the ledger authority, restores scoped MCP ownership, confirms adoption, installs the +existing Pylon incarnation fence, and only then releases exact retained replay and live events. Startup +performs this adoption pass before generic orphan reconciliation. A terminal reached while Pylon was down +therefore settles the original turn and checkpoint once without another prompt or `turn.started`. + +Graceful process shutdown detaches an eligible owned worker and leaves its compatible supervisor alive; +explicit Stop and normal terminal cleanup still require Prime's authoritative owned-session cleanup proof. +The private row is deleted only after that proof, terminal projection delivery, and checkpoint quiescence. +A competing Pylon process cannot win the same ledger generation, and a process that did not spawn a +compatible supervisor never shuts it down. + +Supervised and other approval-required sessions, ACP mode, stock/manual or unverified distributions, +native Windows, copied state, a replaced supervisor, an unsupported host, unresolved interaction state, +partial replay, transcript mismatch, or any other unproven identity/continuity retain the existing precise +orphan result. Pylon does not steal, re-submit, disclose, or synthesize native work in those cases. When both bounded stats reads succeed, the usage delta includes child billing that Prime attributes after the original message event. The native autonomous-status result is discarded at the provider boundary. Older connections without the barrier retain response-boundary behavior. An error- or tool-terminated native run completion without a diff --git a/docs/internals/prime-agent-native-parity.md b/docs/internals/prime-agent-native-parity.md index 075d67143..8046e2949 100644 --- a/docs/internals/prime-agent-native-parity.md +++ b/docs/internals/prime-agent-native-parity.md @@ -19,7 +19,7 @@ Use the independently installed Prime Agent package and its public detached-daem - use a short, stable Pylon-owned socket name so the user's normal Prime daemon is untouched; contain it in an owner-only directory below a trusted private or sticky temporary root; - strip every inherited `PRIME_AGENT_INTERNAL_*` variable before launch, because Pylon may itself be running inside a Prime worker; - create one client-owned Prime daemon session per live Pylon thread; -- persist exact Prime session identity in a server-private thread sidecar while keeping the client-visible provider resume cursor opaque, then rehydrate after server restart; +- persist exact Prime session identity in server-private state while keeping the client-visible provider resume cursor opaque; exact Pylon-managed Full access sessions can adopt one proven active turn after restart, while Supervised/approval-required sessions, unproven continuity, and native Windows retain the existing orphan result without exposing a recovery handle; - keep ACP as an explicit compatibility fallback on supported hosts when daemon setup cannot be used; - support provider execution on macOS, Linux, and WSL2 only; reject native `win32` at the driver boundary before any Prime process or probe and direct users to WSL2. diff --git a/docs/user/providers-prime-agent.md b/docs/user/providers-prime-agent.md index aa2f22eda..b76a7ba37 100644 --- a/docs/user/providers-prime-agent.md +++ b/docs/user/providers-prime-agent.md @@ -128,8 +128,16 @@ negotiated, recovery is stricter: Pylon never copies missing prompt output from provide complete event continuity, and the completed-message transcript must exactly match messages Pylon already received through attributed live events. Any extra snapshot message closes the uncertain session rather than guessing whether it was your answer or unrelated background output. -Restarting the Pylon server is a separate boundary and does not yet adopt Prime work that is still running -in another process. +An active Full access turn can also survive a Pylon server restart when the exact Prime installation is +Pylon managed and the replacement server can prove the same retained native execution and complete +event history. Pylon restores the turn's scoped browser/MCP access before showing recovered activity and +never sends your prompt again. The recovery identity and handle remain private to the server and are not +sent to clients or written to public thread history. + +Supervised or other approval-required sessions do not use restart adoption. Neither do ACP sessions, +stock or manually installed Prime distributions, native Windows, a replaced Prime supervisor, or any turn +whose identity or complete event continuity cannot be proven. Those cases keep the existing orphaned +session result rather than guessing, replaying the prompt, or exposing partial native work. Native Windows is not a Prime Agent provider runtime. Pylon does not fall back to ACP there. Use WSL2, where the server runs as Linux, or connect this client to another supported environment. From d4fd9e01ece708ce5ec00b915a474970101ed978 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Tue, 1 Sep 2026 09:24:22 -0600 Subject: [PATCH 2/4] fix(prime): preserve turns across server restarts Refs #84 --- .../provider/Layers/ProviderService.test.ts | 110 +- .../src/provider/Layers/ProviderService.ts | 44 +- .../prime/PrimeAgentDaemonAdapter.test.ts | 66 + .../provider/prime/PrimeAgentDaemonAdapter.ts | 81 +- .../prime/PrimeAgentDaemonManager.test.ts | 49 + .../provider/prime/PrimeAgentDaemonManager.ts | 30 +- .../PrimeAgentDaemonSessionRuntime.test.ts | 13 +- .../prime/PrimeAgentDaemonSessionRuntime.ts | 20 +- .../PrimeAgentRestartAdoption.real.test.mjs | 1404 +++++++++++++++-- 9 files changed, 1609 insertions(+), 208 deletions(-) diff --git a/apps/server/src/provider/Layers/ProviderService.test.ts b/apps/server/src/provider/Layers/ProviderService.test.ts index f6f131594..eeeaa7730 100644 --- a/apps/server/src/provider/Layers/ProviderService.test.ts +++ b/apps/server/src/provider/Layers/ProviderService.test.ts @@ -2188,6 +2188,83 @@ routing.layer("ProviderServiceLive routing", (it) => { }), ); + it.effect("persists turn.started before a pending provider send completes", () => + Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const runtimeRepository = yield* ProviderSessionRuntime.ProviderSessionRuntimeRepository; + const sendStarted = yield* Deferred.make(); + const releaseSend = yield* Deferred.make(); + routing.codex.sendTurn.mockImplementationOnce((input) => + Effect.gen(function* () { + yield* Deferred.succeed(sendStarted, undefined); + yield* Deferred.await(releaseSend); + return { + threadId: input.threadId, + turnId: TurnId.make(`turn-${String(input.threadId)}`), + }; + }), + ); + + const threadId = asThreadId("thread-started-event-directory"); + const session = yield* provider.startSession(threadId, { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + }); + const admissionRequestId = CommandId.make("request-started-event-directory"); + const turnId = TurnId.make("turn-started-event-directory"); + const eventId = asEventId("evt-started-event-directory"); + const published = yield* provider.streamEvents.pipe( + Stream.filter((event) => event.eventId === eventId), + Stream.take(1), + Stream.runDrain, + Effect.forkChild, + ); + yield* Effect.yieldNow; + const send = yield* provider + .sendTurn({ + threadId, + input: "hold until durable", + attachments: [], + admissionRequestId, + sessionIncarnationId: session.sessionIncarnationId, + }) + .pipe(Effect.forkChild); + yield* Deferred.await(sendStarted); + + routing.codex.emit({ + type: "turn.started", + eventId, + provider: CODEX_DRIVER, + createdAt: "2026-01-01T00:00:00.000Z", + threadId, + turnId, + admissionRequestId, + sessionIncarnationId: session.sessionIncarnationId, + payload: {}, + }); + yield* Fiber.join(published); + + const persisted = Option.getOrThrow(yield* runtimeRepository.getByThreadId({ threadId })); + assert.equal(persisted.status, "running"); + assert.deepEqual(persisted.runtimePayload, { + activeTurnId: turnId, + activeTurnRequestId: admissionRequestId, + admissionRequestId, + cwd: process.cwd(), + lastError: null, + lastRuntimeEvent: "turn.started", + lastRuntimeEventAt: "2026-01-01T00:00:00.000Z", + model: null, + sessionIncarnationId: session.sessionIncarnationId, + }); + + yield* Deferred.succeed(releaseSend, undefined); + yield* Fiber.join(send); + }), + ); + it.effect("does not persist running after a concurrent send is interrupted", () => Effect.gen(function* () { const provider = yield* ProviderService.ProviderService; @@ -3607,17 +3684,37 @@ describe("agent browser access", () => { const threadId = asThreadId("thread-restart-adoption"); const codex = makeFakeCodexAdapter(); const order: string[] = []; + const recoveredTurnId = TurnId.make("turn-restart-adoption"); + const recoveredRequestId = CommandId.make("request-restart-adoption"); + let recoveredSession: ProviderSession | undefined; const recoveryAdapter: ProviderAdapterShape = { ...codex.adapter, recoverSession: (input) => Effect.gen(function* () { assert.isDefined(McpProviderSession.readMcpProviderSession(threadId)); order.push("recover"); - return yield* codex.startSession(input); + const session = yield* codex.startSession(input); + recoveredSession = { + ...session, + status: "running", + activeTurnId: recoveredTurnId, + activeTurnRequestId: recoveredRequestId, + }; + return recoveredSession; }), activateRecoveredSession: () => Effect.sync(() => { order.push("activate"); + codex.emit({ + type: "content.delta", + eventId: asEventId("evt-restart-adopted-output"), + provider: CODEX_DRIVER, + threadId, + turnId: recoveredTurnId, + sessionIncarnationId: recoveredSession!.sessionIncarnationId, + createdAt: "2026-01-01T00:00:00.000Z", + delta: "recovered", + }); }), }; const providerLayer = makeAgentBrowserProviderLayer( @@ -3645,10 +3742,21 @@ describe("agent browser access", () => { }); codex.removeSession(threadId); order.length = 0; + const recoveredEvents = yield* Ref.make>([]); + const consumer = yield* Stream.take(provider.streamEvents, 1).pipe( + Stream.runForEach((event) => Ref.update(recoveredEvents, (events) => [...events, event])), + Effect.forkChild, + ); + yield* Effect.yieldNow; yield* provider.recoverRestartSessions!(); + yield* Fiber.join(consumer); assert.deepEqual(order, ["mcp", "recover", "activate"]); + const [recoveredEvent] = yield* Ref.get(recoveredEvents); + assert.equal(recoveredEvent?.eventId, asEventId("evt-restart-adopted-output")); + assert.equal(recoveredEvent?.turnId, recoveredTurnId); + assert.equal(recoveredEvent?.admissionRequestId, recoveredRequestId); const binding = Option.getOrThrow(yield* directory.getBinding(threadId)); assert.equal( (binding.runtimePayload as { readonly lastRuntimeEvent?: string }).lastRuntimeEvent, diff --git a/apps/server/src/provider/Layers/ProviderService.ts b/apps/server/src/provider/Layers/ProviderService.ts index 57acd7048..ada472642 100644 --- a/apps/server/src/provider/Layers/ProviderService.ts +++ b/apps/server/src/provider/Layers/ProviderService.ts @@ -656,15 +656,31 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( }; } - if ( - canonicalEvent.type === "turn.started" && - canonicalEvent.admissionRequestId !== undefined - ) { - activeTurnAdmissions.set(canonicalEvent.threadId, { - requestId: canonicalEvent.admissionRequestId, - sessionIncarnationId: currentIncarnation.id, - turnId: canonicalEvent.turnId, - }); + if (canonicalEvent.type === "turn.started") { + if (canonicalEvent.admissionRequestId !== undefined) { + activeTurnAdmissions.set(canonicalEvent.threadId, { + requestId: canonicalEvent.admissionRequestId, + sessionIncarnationId: currentIncarnation.id, + turnId: canonicalEvent.turnId, + }); + } + yield* directory + .upsert({ + threadId: canonicalEvent.threadId, + provider: canonicalEvent.provider, + providerInstanceId: source.instanceId, + status: "running", + runtimePayload: { + activeTurnId: canonicalEvent.turnId, + ...(canonicalEvent.admissionRequestId === undefined + ? {} + : { activeTurnRequestId: canonicalEvent.admissionRequestId }), + sessionIncarnationId: currentIncarnation.id, + lastRuntimeEvent: "turn.started", + lastRuntimeEventAt: canonicalEvent.createdAt, + }, + }) + .pipe(Effect.orDie); } yield* increment(providerRuntimeEventsTotal, { @@ -1378,11 +1394,19 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( return; } adoptedAdapter = adapter; + const sessionIncarnationId = RuntimeSessionId.make(rawIncarnation); currentSessionIncarnations.set(binding.threadId, { - id: RuntimeSessionId.make(rawIncarnation), + id: sessionIncarnationId, instanceId, adapter, }); + if (recovered.activeTurnId !== undefined && recovered.activeTurnRequestId !== undefined) { + activeTurnAdmissions.set(binding.threadId, { + requestId: recovered.activeTurnRequestId, + sessionIncarnationId, + turnId: recovered.activeTurnId, + }); + } yield* upsertSessionBinding( { ...recovered, providerInstanceId: instanceId }, binding.threadId, diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.test.ts b/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.test.ts index 969811e98..6044249fa 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.test.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.test.ts @@ -1375,6 +1375,10 @@ describe("PrimeAgentDaemonAdapter", () => { _tag: "PromptLifecycleUpdated", lifecycle: lifecycleSnapshot(correlationId, "delivered", 2), }); + yield* offer(captures, { + _tag: "PromptLifecycleUpdated", + lifecycle: lifecycleSnapshot(correlationId, "owned", 1), + }); yield* offer(captures, { _tag: "MessageCompleted", message: assistantMessage("wrong owner contamination"), @@ -2104,6 +2108,68 @@ describe("PrimeAgentDaemonAdapter", () => { ).pipe(Effect.provide(testLayer)), ); + it.effect("replays one missed terminal response after the admitted user boundary", () => + Effect.scoped( + Effect.gen(function* () { + const captures = makeCaptures(); + captures.correlatedPromptLifecycleAvailable = true; + captures.correlatedPromptObserved = yield* Queue.unbounded(); + captures.rlmConnectionGeneration = 1; + captures.rlmContinuityValid = false; + const adapter = yield* makePrimeAgentDaemonAdapter(decodeSettings({}), manager, { + instanceId, + runtimeFactory: fakeRuntimeFactory(captures), + }); + const subscription = yield* subscribe(adapter); + yield* adapter.startSession({ threadId, cwd: process.cwd(), runtimeMode: "full-access" }); + const turnFiber = yield* adapter + .sendTurn({ threadId, input: "response missed during reconnect" }) + .pipe(Effect.forkChild); + const correlationId = yield* Queue.take(captures.correlatedPromptObserved); + yield* offer(captures, { + _tag: "PromptLifecycleUpdated", + lifecycle: lifecycleSnapshot(correlationId, "delivered", 2), + }); + const prompt = { + role: "user", + timestamp: 1, + text: "response missed during reconnect", + imageMimeTypes: [], + imageDigests: [], + } satisfies PrimeDaemonMessage; + const answer = { ...assistantMessage("replayed terminal answer"), timestamp: 2 }; + yield* offer(captures, { + _tag: "MessageCompleted", + message: prompt, + attribution: { scope: "prompt", correlationId }, + }); + + yield* offer(captures, { + ...initialSnapshot(), + state: { ...initialSnapshot().state, messageCount: 2 }, + messages: [prompt, answer], + replayContinuity: "complete", + connectionGeneration: 1, + promptLifecycles: { + records: [lifecycleSnapshot(correlationId, "completed", 3, { usage })], + expired: [], + }, + }); + const result = yield* Fiber.join(turnFiber); + const turnEvents = subscription.events.filter((event) => event.turnId === result.turnId); + expect(captures.reconnectResolutions).toContainEqual({ + generation: 1, + reconciled: true, + terminalResponseObserved: false, + }); + expect(encodeUnknownJson(turnEvents)).toContain("replayed terminal answer"); + expect(turnEvents.findLast((event) => event.type === "turn.completed")).toMatchObject({ + payload: { state: "completed", usage: { totalTokens: usage.totalTokens } }, + }); + }), + ).pipe(Effect.provide(testLayer)), + ); + it.effect("does not apply a terminal correlated snapshot when proof settlement is rejected", () => Effect.scoped( Effect.gen(function* () { diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts b/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts index 15add700a..32810e785 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts @@ -1412,20 +1412,6 @@ export function makePrimeAgentDaemonAdapter( } }); - const correlatedTranscriptSnapshotIsExact = ( - context: PrimeAgentDaemonSessionContext, - event: Extract, - ): boolean => { - if (event.replayContinuity !== "complete") return false; - const reconciliation = reconcileTranscriptTail({ - observed: context.nativeTranscript, - observedCount: context.nativeTranscriptMessageCount, - snapshot: event.messages, - snapshotCount: event.state.messageCount, - }); - return reconciliation !== undefined && reconciliation.missingMessages.length === 0; - }; - const reconcileTranscriptSnapshotLocked = ( context: PrimeAgentDaemonSessionContext, event: Extract, @@ -1875,6 +1861,9 @@ export function makePrimeAgentDaemonAdapter( if (turn === undefined || turn.correlationId !== lifecycle.correlationId) return false; const current = turn.correlatedLifecycle; if (current !== undefined) { + // The submit response and its lifecycle notification travel on separate + // channels, so the older notification can arrive after the response. + if (lifecycle.revision < current.revision) return false; if (lifecycle.revision === current.revision) { if (primeAgentPromptLifecycleIsSame(lifecycle, current)) return false; return yield* failCorrelatedProtocolLocked(context, turn); @@ -2124,7 +2113,13 @@ export function makePrimeAgentDaemonAdapter( /** Must be called with the thread lock held. */ const startBackgroundQuiescenceWatchLocked = (context: PrimeAgentDaemonSessionContext) => { - if (!context.runtime.rlmQuiescenceAvailable || context.stopped) return Effect.void; + if ( + !context.runtime.rlmQuiescenceAvailable || + context.stopped || + context.activeTurn?.correlationId !== undefined + ) { + return Effect.void; + } if (context.backgroundQuiescencePending) { // The native call may finish later, but its aborted signal prevents that older watch // from clearing activity observed by the replacement generation. @@ -2282,8 +2277,37 @@ export function makePrimeAgentDaemonAdapter( context.managedPlanProjectionEnabled = true; const activeTurn = context.activeTurn; if (context.runtime.correlatedPromptLifecycleAvailable) { - if (event.initialSnapshot !== true) { - if (!correlatedTranscriptSnapshotIsExact(context, event)) { + if (event.initialSnapshot === true) { + const lifecycle = + activeTurn?.correlationId === undefined + ? undefined + : event.promptLifecycles?.records.find( + (candidate) => candidate.correlationId === activeTurn.correlationId, + ); + if (lifecycle !== undefined) { + yield* applyCorrelatedPromptLifecycleLocked(context, lifecycle, { + authoritativeSnapshot: true, + }); + } + } else { + const transcriptPlan = reconcileTranscriptTail({ + observed: context.nativeTranscript, + observedCount: context.nativeTranscriptMessageCount, + snapshot: event.messages, + snapshotCount: event.state.messageCount, + }); + const missingMessages = transcriptPlan?.missingMessages ?? []; + const snapshotIsExactOrCurrentTerminal = + missingMessages.length === 0 || + (missingMessages.length === 1 && + missingMessages[0]?.role === "assistant" && + context.nativeTranscript.at(-1)?.role === "user"); + const transcriptReconciled = + event.replayContinuity === "complete" && + transcriptPlan !== undefined && + snapshotIsExactOrCurrentTerminal && + (yield* reconcileTranscriptSnapshotLocked(context, event)); + if (!transcriptReconciled) { if (reconnectGeneration !== undefined) { context.runtime.resolveReconnectSnapshot(reconnectGeneration, false, false); } @@ -4397,13 +4421,17 @@ export function makePrimeAgentDaemonAdapter( (manager.platform === "darwin" || manager.platform === "linux") && (manager.architecture === "arm64" || manager.architecture === "x64"); - const silentlyCloseSessionForRecovery = (context: PrimeAgentDaemonSessionContext) => - Effect.gen(function* () { + const reserveSessionForRecoveryRestart = (context: PrimeAgentDaemonSessionContext) => + Effect.sync(() => { context.stopped = true; sessions.delete(context.threadId); - if (context.eventFiber !== undefined) yield* Fiber.interrupt(context.eventFiber); context.backgroundQuiescenceController?.abort(); context.backgroundQuiescenceController = undefined; + }); + + const closeReservedRecoverySession = (context: PrimeAgentDaemonSessionContext) => + Effect.gen(function* () { + if (context.eventFiber !== undefined) yield* Fiber.interrupt(context.eventFiber); yield* Scope.close(context.scope, Exit.void).pipe(Effect.ignore); }); @@ -4461,12 +4489,13 @@ export function makePrimeAgentDaemonAdapter( resumeCursor: context.session.resumeCursor, sessionIncarnationId: context.sessionIncarnationId, } as const; - yield* silentlyCloseSessionForRecovery(context); + yield* reserveSessionForRecoveryRestart(context); pendingRecoveryStarts.set(input.threadId, recoveryStart); - return { restartInput, ownerToken } as const; + return { restartInput, ownerToken, context } as const; }), ); if (plan === undefined) return; + yield* closeReservedRecoverySession(plan.context); const recoveryResult = yield* Effect.result(startSession(plan.restartInput)); if (Result.isSuccess(recoveryResult)) return; @@ -4538,6 +4567,10 @@ export function makePrimeAgentDaemonAdapter( ) { return null; } + const recoverySessionDir = authority.recoveryConfig.sessionDir; + if (typeof recoverySessionDir !== "string") { + return null; + } const ownerToken = yield* randomUUIDv4; const claimedAt = yield* nowIso; const claimed = yield* recoveryLedger!.claim({ @@ -4560,7 +4593,7 @@ export function makePrimeAgentDaemonAdapter( ownerToken, requestId, mcpOwnerId, - sessionFile: `${authority.nativeSessionId}.jsonl`, + sessionFile: path.join(recoverySessionDir, `${authority.nativeSessionId}.jsonl`), }); const started = yield* Effect.result( startSession({ @@ -4609,7 +4642,7 @@ export function makePrimeAgentDaemonAdapter( context.recoveryPendingActivation = false; context.eventFiber = yield* context.runtime.events.pipe( Stream.runForEach((event) => consumeEvent(context, event)), - Effect.forkChild, + Effect.forkIn(context.scope), ); if (context.runtime.inputAdmissionBusy) { yield* startBackgroundQuiescenceWatchLocked(context); diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonManager.test.ts b/apps/server/src/provider/prime/PrimeAgentDaemonManager.test.ts index 2c890117a..156422c28 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonManager.test.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonManager.test.ts @@ -31,6 +31,10 @@ interface CapturedCommand { readonly options: { readonly env?: NodeJS.ProcessEnv; readonly extendEnv?: boolean; + readonly detached?: boolean; + readonly stdin?: "ignore"; + readonly stdout?: "ignore"; + readonly stderr?: "ignore"; }; } @@ -423,6 +427,10 @@ describe("PrimeAgentDaemonManager lifecycle", () => { manager.sessionDir, ]); expect(command.options.extendEnv).toBe(false); + expect(command.options).not.toHaveProperty("detached"); + expect(command.options).not.toHaveProperty("stdin"); + expect(command.options).not.toHaveProperty("stdout"); + expect(command.options).not.toHaveProperty("stderr"); expect(command.options.env).toMatchObject({ PATH: "/usr/bin", KEEP_ME: "yes", @@ -485,6 +493,26 @@ describe("PrimeAgentDaemonManager lifecycle", () => { ), ); + it.effect("detaches a recovery supervisor from Pylon stdio while retaining its handle", () => { + const fixture = managerFixture({ recoverable: true }); + return Effect.gen(function* () { + const manager = yield* fixture.make; + const client = yield* manager.openClient(); + client.close(); + + expect(fixture.commands).toHaveLength(1); + expect(fixture.processes).toHaveLength(1); + expect(fixture.commands[0]!.options).toMatchObject({ + detached: true, + stdin: "ignore", + stdout: "ignore", + stderr: "ignore", + extendEnv: false, + }); + expect(fixture.processes[0]!.handle.pid).toBe(1); + }).pipe(Effect.scoped); + }); + it.effect( "fails readiness without publishing a daemon and interrupts only its captured handle", () => { @@ -645,4 +673,25 @@ describe("PrimeAgentDaemonManager lifecycle", () => { ), ); }); + + it.effect("retires an adopted recovery supervisor on clean shutdown", () => { + const fixture = managerFixture({ existingLive: true, recoverable: true }); + return Effect.scoped( + Effect.gen(function* () { + const manager = yield* fixture.make; + const client = yield* manager.openClient(); + client.close(); + const release = manager.retainForRecovery!(); + release(); + expect(fixture.shutdownRequests).toEqual([]); + }), + ).pipe( + Effect.andThen( + Effect.sync(() => { + expect(fixture.shutdownRequests).toEqual([fixture.paths.socket]); + expect(fixture.commands).toHaveLength(0); + }), + ), + ); + }); }); diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonManager.ts b/apps/server/src/provider/prime/PrimeAgentDaemonManager.ts index bc476d353..d22fa01f5 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonManager.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonManager.ts @@ -307,8 +307,11 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage ).filter((entry): entry is [string, string] => typeof entry[1] === "string"), ); let recoveryRetainers = 0; + let retainedExistingDaemon = false; + let adoptedExistingDaemon = false; const retainForRecovery = () => { recoveryRetainers += 1; + if (retainedExistingDaemon) adoptedExistingDaemon = true; let released = false; return () => { if (released) return; @@ -330,7 +333,6 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage const shutdownTimeout = input.shutdownTimeout ?? Duration.seconds(5); const semaphore = yield* Semaphore.make(1); let running: RunningDaemon | undefined; - let retainedExistingDaemon = false; let closing = false; const removeSocket = () => @@ -726,7 +728,16 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage const command = ChildProcess.make( input.executablePath, ["--mode", "daemon", "--daemon-socket", socket, "--offline", "--session-dir", sessionDir], - { env: launchEnvironment, extendEnv: false }, + recoveryEnabled + ? { + env: launchEnvironment, + extendEnv: false, + detached: true, + stdin: "ignore", + stdout: "ignore", + stderr: "ignore", + } + : { env: launchEnvironment, extendEnv: false }, ); const handle = yield* spawner.spawn(command).pipe( Effect.provideService(Scope.Scope, processScope), @@ -736,8 +747,10 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage Effect.onError(() => Scope.close(processScope, Exit.void).pipe(Effect.ignore)), ); const state = { handle, scope: processScope } satisfies RunningDaemon; - yield* Effect.forkIn(drainProcessOutput(handle.stdout, "stdout"), processScope); - yield* Effect.forkIn(drainProcessOutput(handle.stderr, "stderr"), processScope); + if (!recoveryEnabled) { + yield* Effect.forkIn(drainProcessOutput(handle.stdout, "stdout"), processScope); + yield* Effect.forkIn(drainProcessOutput(handle.stderr, "stderr"), processScope); + } const readinessClient = yield* connectClient({ bridge, socket, timeoutMs }).pipe( Effect.retry({ @@ -778,9 +791,12 @@ export const makePrimeAgentDaemonManager = Effect.fn("makePrimeAgentDaemonManage return; } if (retainedExistingDaemon) { - // This process did not spawn the compatible supervisor. A competing replacement - // may already own it, so shutdown must never revoke that process's authority. - return; + // Do not revoke an untouched compatible supervisor owned by another + // process. Once this manager adopts one of its recoverable sessions, + // it owns the surviving supervisor and must retire it on clean shutdown. + if (!adoptedExistingDaemon) return; + const client = yield* connectClient({ bridge, socket, timeoutMs }); + yield* retireExistingDaemon(client); } }), ); diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts index 59dce0a54..e86b8cd7b 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts @@ -13034,7 +13034,7 @@ describe("Prime Agent live activity privacy boundary", () => { ), ); - it.effect("commits adoption, restores MCP, then confirms before replay is exposed", () => + it.effect("commits adoption before replay and defers MCP replacement to the next prompt", () => Effect.scoped( Effect.gen(function* () { const side = fixture({ @@ -13086,16 +13086,19 @@ describe("Prime Agent live activity privacy boundary", () => { "adopt-recoverable", "retain-daemon", "ledger-committed", - "replace-mcp", "confirm-adoption", ]), ); + expect(side.captures.order).not.toContain("replace-mcp"); expect(side.captures.order.indexOf("ledger-committed")).toBeLessThan( - side.captures.order.indexOf("replace-mcp"), - ); - expect(side.captures.order.indexOf("replace-mcp")).toBeLessThan( side.captures.order.indexOf("confirm-adoption"), ); + + yield* runtime.prompt({ text: "next prompt after adoption" }); + expect(side.captures.order).toContain("replace-mcp"); + expect(side.captures.order.indexOf("confirm-adoption")).toBeLessThan( + side.captures.order.indexOf("replace-mcp"), + ); yield* runtime.detach!; }), ), diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts index bad48409c..8cc6f3fd1 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts @@ -2525,7 +2525,23 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo }); if (mayCommit()) mcpAttached = true; }); - yield* configureMcpServer().pipe(Effect.onError(() => closeAttachedSession)); + let adoptionMcpRefreshPending = + input.recovery?.kind === "adopt" && input.mcpServer !== undefined; + if (adoptionMcpRefreshPending) { + // Recoverable adoption retags the surviving worker's existing MCP servers + // to the new owner. The admitted prompt is still running, so Prime cannot + // replace those servers until the next prompt admission boundary. + mcpAttached = true; + } else { + yield* configureMcpServer().pipe(Effect.onError(() => closeAttachedSession)); + } + const refreshMcpAfterAdoption = Effect.fn( + "PrimeAgentDaemonSessionRuntime.refreshMcpAfterAdoption", + )(function* () { + if (!adoptionMcpRefreshPending) return; + yield* configureMcpServer(); + adoptionMcpRefreshPending = false; + }); let verifiedInventory: | readonly [typeof resourceSnapshotSchema.Type, typeof commandsSchema.Type] | undefined; @@ -6749,6 +6765,7 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo yield* requireCorrelatedPromptLifecycleAdmission("prompt"); yield* requireCurrentCorrelatedPromptLifecycleProof("prompt"); yield* awaitProviderRecovery; + yield* refreshMcpAfterAdoption(); yield* requireCorrelatedPromptLifecycleAdmission("prompt"); const proofEpoch = yield* requireCurrentCorrelatedPromptLifecycleProof("prompt"); const images = yield* validateImages("prompt", promptInput.images); @@ -6878,6 +6895,7 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo yield* requireCorrelatedPromptLifecycleAdmission("prompt"); } yield* awaitProviderRecovery; + yield* refreshMcpAfterAdoption(); const resumedAfterAbort = yield* resumeAfterAbort(); const images = yield* validateImages("prompt", promptInput.images); yield* validatePromptContent("prompt", promptInput.text, images); diff --git a/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs b/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs index 9a34d0ff0..3a6e44340 100644 --- a/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs +++ b/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs @@ -1,194 +1,1278 @@ -import * as ChildProcess from "node:child_process"; -import { randomUUID } from "node:crypto"; -import * as FSP from "node:fs/promises"; -import * as OS from "node:os"; -import * as Path from "node:path"; -import { pathToFileURL } from "node:url"; +/* eslint-disable t3code/no-manual-effect-runtime-in-tests -- This opt-in POSIX proof drives two external server processes through their public RPC boundary. */ +import * as NodeChildProcess from "node:child_process"; +import * as NodeCrypto from "node:crypto"; +import * as NodeFSP from "node:fs/promises"; +import * as NodeHttp from "node:http"; +import * as NodeNet from "node:net"; +import * as NodeOS from "node:os"; +import * as NodePath from "node:path"; +import * as NodeProcess from "node:process"; +import * as NodeSqlite from "node:sqlite"; +import * as NodeURL from "node:url"; +import * as NodeSocket from "@effect/platform-node/NodeSocket"; +import { ORCHESTRATION_WS_METHODS, WS_METHODS, WsRpcGroup } from "@t3tools/contracts"; +import * as Deferred from "effect/Deferred"; +import * as Effect from "effect/Effect"; +import * as Layer from "effect/Layer"; +import * as Socket from "effect/unstable/socket/Socket"; +import * as Stream from "effect/Stream"; +import { RpcClient, RpcSerialization } from "effect/unstable/rpc"; import { describe, expect, it } from "vite-plus/test"; -const packageRoot = process.env.PRIME_AGENT_RECOVERY_REAL_PACKAGE_ROOT; +import { persistPrimeManagedReceipt } from "./PrimeAgentDistributionVerifier.ts"; + +const packageRoot = NodeProcess.env.PRIME_AGENT_REAL_PACKAGE_ROOT?.trim(); const exactHead = "507a52239d3ace7bb2b2965ade7779988fdb6344"; +const skipReason = + NodeProcess.platform === "win32" + ? "native Windows is unsupported; run the POSIX proof in WSL2 with a Linux PRIME_AGENT_REAL_PACKAGE_ROOT" + : "set PRIME_AGENT_REAL_PACKAGE_ROOT to the built exact Prime checkout at 507a52239d3ace7bb2b2965ade7779988fdb6344"; +const enabled = NodeProcess.platform !== "win32" && Boolean(packageRoot); +const outerSafetyMs = 120_000; +const maximumOutputBytes = 2 * 1024 * 1024; +const providerInstanceId = "primeAgent"; +const modelSelection = { + instanceId: providerInstanceId, + model: "faux-adoption/faux-adoption", +}; -const waitForExit = (child, timeoutMs) => +const withSafetyCeiling = (promise, timeoutMs, label) => new Promise((resolve, reject) => { - const timer = setTimeout(() => reject(new Error("subprocess exit timed out")), timeoutMs); - child.once("exit", (code, signal) => { - clearTimeout(timer); - resolve({ code, signal }); - }); + const timer = setTimeout( + () => reject(new Error(`${label} exceeded its safety ceiling`)), + timeoutMs, + ); + promise.then( + (value) => { + clearTimeout(timer); + resolve(value); + }, + (error) => { + clearTimeout(timer); + reject(error); + }, + ); }); -const runCaptured = (command, args, options, timeoutMs) => +const waitForExit = (child, timeoutMs, label) => { + if (child.exitCode !== null || child.signalCode !== null) { + return Promise.resolve({ code: child.exitCode, signal: child.signalCode }); + } + return withSafetyCeiling( + new Promise((resolve, reject) => { + child.once("error", reject); + child.once("exit", (code, signal) => resolve({ code, signal })); + }), + timeoutMs, + label, + ); +}; + +const runCaptured = (command, args, options, timeoutMs, label) => new Promise((resolve, reject) => { - const child = ChildProcess.spawn(command, args, options); + const child = NodeChildProcess.spawn(command, args, options); const stdout = []; const stderr = []; let stdoutBytes = 0; let stderrBytes = 0; - const maximumBytes = 1024 * 1024; + let settled = false; + const finish = (effect) => { + if (settled) return; + settled = true; + clearTimeout(timer); + effect(); + }; const timer = setTimeout(() => { child.kill("SIGTERM"); - reject(new Error("owner subprocess timed out")); + finish(() => reject(new Error(`${label} exceeded its safety ceiling`))); }, timeoutMs); child.stdout.on("data", (chunk) => { stdoutBytes += chunk.length; - if (stdoutBytes > maximumBytes) child.kill("SIGTERM"); - else stdout.push(chunk); + if (stdoutBytes > maximumOutputBytes) { + child.kill("SIGTERM"); + finish(() => reject(new Error(`${label} exceeded its stdout budget`))); + } else { + stdout.push(chunk); + } }); child.stderr.on("data", (chunk) => { stderrBytes += chunk.length; - if (stderrBytes > maximumBytes) child.kill("SIGTERM"); - else stderr.push(chunk); - }); - child.once("error", (error) => { - clearTimeout(timer); - reject(error); + if (stderrBytes > maximumOutputBytes) { + child.kill("SIGTERM"); + finish(() => reject(new Error(`${label} exceeded its stderr budget`))); + } else { + stderr.push(chunk); + } }); - child.once("exit", (code) => { - clearTimeout(timer); + child.once("error", (error) => finish(() => reject(error))); + child.once("exit", (code, signal) => { const output = Buffer.concat(stdout).toString("utf8"); const errorOutput = Buffer.concat(stderr).toString("utf8"); - if (code === 0) resolve(output); - else reject(new Error(`owner subprocess failed (${code}): ${errorOutput}`)); + finish(() => { + if (code === 0) resolve(output); + else reject(new Error(`${label} failed (${code ?? signal}): ${errorOutput || output}`)); + }); }); }); -describe.skipIf(!packageRoot)("Prime Agent exact-head restart adoption subprocess", () => { - it("adopts after the creating owner process exits and proves authoritative cleanup", async () => { - const gitHead = ChildProcess.execFileSync("git", ["-C", packageRoot, "rev-parse", "HEAD"], { - encoding: "utf8", - timeout: 5_000, - }).trim(); - expect(gitHead).toBe(exactHead); - - const codingAgentRoot = Path.join(packageRoot, "packages", "coding-agent"); - const sdkEntry = Path.join(codingAgentRoot, "dist", "index.js"); - const cliEntry = Path.join(codingAgentRoot, "dist", "bundle", "cli.js"); - const temp = await FSP.mkdtemp(Path.join(OS.tmpdir(), "pylon-prime-restart-")); - const socket = Path.join(temp, "daemon.sock"); - const sessionDir = Path.join(temp, "sessions"); - const agentDir = Path.join(temp, "agent-home"); - await FSP.mkdir(sessionDir, { recursive: true }); - await FSP.mkdir(agentDir, { recursive: true }); - const launchEnv = { - HOME: process.env.HOME ?? temp, - PATH: process.env.PATH ?? "/usr/bin:/bin", - PRIME_AGENT_CODING_AGENT_DIR: agentDir, - }; - const daemon = ChildProcess.spawn( - process.execPath, - [ - cliEntry, - "--mode", - "daemon", - "--daemon-socket", - socket, - "--offline", - "--session-dir", - sessionDir, - ], - { env: launchEnv, stdio: ["ignore", "ignore", "pipe"] }, - ); - const daemonErrors = []; - daemon.stderr.on("data", (chunk) => daemonErrors.push(chunk)); +const reserveEphemeralPort = () => + new Promise((resolve, reject) => { + const server = NodeNet.createServer(); + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + const address = server.address(); + if (address === null || typeof address === "string") { + server.close(); + reject(new Error("ephemeral port reservation returned no TCP address")); + return; + } + server.close((error) => (error ? reject(error) : resolve(address.port))); + }); + }); - try { - const helper = Path.join(temp, "create-owner.mjs"); - await FSP.writeFile( - helper, - `import { randomUUID } from "node:crypto"; -import { DaemonClient, createRecoverableOwnedSession } from ${JSON.stringify(pathToFileURL(sdkEntry).href)}; -const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); -const client = new DaemonClient(${JSON.stringify(socket)}); -let connected = false; -for (let attempt = 0; attempt < 80; attempt += 1) { - try { await client.connect(); connected = true; break; } catch { await sleep(25); } -} -if (!connected) throw new Error("daemon readiness timed out"); -await client.waitForHello(); -const config = ${JSON.stringify({ cwd: temp, sessionDir, noBuiltinTools: true, noExtensions: true, noSkills: true, noContextFiles: true })}; -const created = await createRecoverableOwnedSession(client, { - requestId: randomUUID(), correlationId: "correlation-real-1", mcpOwnerId: "pylon:mcp-real-1", - config, continueRecent: false, launchEnv: ${JSON.stringify(launchEnv)}, - connectionOptions: { closeClientOnDispose: false, supportsExtensionUi: true }, -}); -await created.connection.submitCorrelatedPrompt("/help", { - correlationId: "correlation-real-1", - queueIfBusy: true, -}); -for (let attempt = 0; attempt < 80; attempt += 1) { - const lifecycles = await created.connection.getPromptLifecycles(); - if (lifecycles.records?.some((entry) => entry.correlationId === "correlation-real-1") || - lifecycles.expired?.some((entry) => entry.correlationId === "correlation-real-1")) break; - await sleep(25); -} -const snapshot = await created.connection.getInitialSnapshot(); -process.stdout.write(JSON.stringify({ - recoveryHandle: created.recoveryHandle, - supervisorGeneration: created.supervisorGeneration, - activeSessionId: created.state.activeSessionId, - sessionId: created.state.sessionId, - cursor: snapshot.lastEventCursor, - config, -})); -client.close(); -`, - "utf8", - ); +const createGate = () => { + let resolveGate; + let settled = false; + const promise = new Promise((resolve) => { + resolveGate = resolve; + }); + return { + promise, + settle: (value) => { + if (settled) return false; + settled = true; + resolveGate(value); + return true; + }, + get settled() { + return settled; + }, + }; +}; + +const startFixtureBackend = async () => { + const firstAdmission = createGate(); + const secondAdmission = createGate(); + const firstConnectionClosed = createGate(); + const records = []; + const sockets = new Set(); + let heldFirstResponse; - const authority = JSON.parse( - await runCaptured( - process.execPath, - [helper], - { env: launchEnv, stdio: ["ignore", "pipe", "pipe"] }, - 20_000, + const contentChunk = (content) => + `data: ${JSON.stringify({ + id: `fixture-${records.length}`, + object: "chat.completion.chunk", + created: 0, + model: modelSelection.model, + choices: [{ index: 0, delta: { role: "assistant", content }, finish_reason: null }], + })}\n\n`; + const terminalChunks = () => + [ + `data: ${JSON.stringify({ + id: `fixture-${records.length}`, + object: "chat.completion.chunk", + created: 0, + model: modelSelection.model, + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + })}\n\n`, + "data: [DONE]\n\n", + ].join(""); + const start = (response, content) => { + response.writeHead(200, { + "Content-Type": "text/event-stream", + Connection: "close", + "Cache-Control": "no-cache", + }); + response.write(contentChunk(content)); + }; + const finish = (response, content) => { + start(response, content); + response.end(terminalChunks()); + }; + + const messageText = (message) => + typeof message.content === "string" + ? message.content + : (message.content ?? []) + .map((part) => (typeof part === "string" ? part : (part.text ?? ""))) + .join(""); + + const server = NodeHttp.createServer((request, response) => { + if (request.method === "GET" && request.url?.endsWith("/models")) { + response.writeHead(200, { "Content-Type": "application/json", Connection: "close" }); + response.end(JSON.stringify({ object: "list", data: [] })); + return; + } + if (request.method !== "POST" || !request.url?.endsWith("/chat/completions")) { + response.writeHead(404, { Connection: "close" }); + response.end(); + return; + } + let buffered = ""; + request.setEncoding("utf8"); + request.on("data", (chunk) => { + buffered += chunk; + }); + request.once("end", () => { + const payload = JSON.parse(buffered); + const authorization = request.headers.authorization ?? ""; + const workerPidMatch = /^Bearer (\d+)$/.exec(authorization); + if (workerPidMatch === null) { + response.destroy( + new Error(`fixture request omitted its worker identity: ${authorization}`), + ); + return; + } + const record = { + callCount: records.length + 1, + workerPid: Number(workerPidMatch[1]), + messages: (payload.messages ?? []).map((message) => ({ + role: message.role, + text: messageText(message), + })), + }; + records.push(record); + if (record.callCount === 1) { + heldFirstResponse = response; + request.socket.once("close", () => firstConnectionClosed.settle()); + firstAdmission.settle(record); + return; + } + if (record.callCount === 2) { + secondAdmission.settle(record); + finish(response, "SECOND_TURN_COMPLETE"); + return; + } + response.destroy(new Error(`unexpected fixture call ${record.callCount}`)); + }); + }); + server.on("connection", (socket) => { + sockets.add(socket); + socket.once("close", () => sockets.delete(socket)); + }); + const port = await withSafetyCeiling( + new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + const address = server.address(); + if (address === null || typeof address === "string") { + reject(new Error("fixture backend returned no HTTP address")); + return; + } + resolve(address.port); + }); + }), + 5_000, + "fixture backend listen", + ); + return { + port, + records, + firstAdmission: firstAdmission.promise, + secondAdmission: secondAdmission.promise, + firstConnectionClosed: firstConnectionClosed.promise, + releaseFirst: () => { + if (heldFirstResponse === undefined || heldFirstResponse.destroyed) { + throw new Error("the first native provider request is not held"); + } + start(heldFirstResponse, "FIRST_TURN_COMPLETE"); + }, + finishFirst: () => { + if (heldFirstResponse === undefined || heldFirstResponse.destroyed) { + throw new Error("the first native provider response cannot be finished"); + } + heldFirstResponse.end(terminalChunks()); + }, + close: () => + withSafetyCeiling( + new Promise((resolve, reject) => { + for (const socket of sockets) socket.destroy(); + server.close((error) => (error ? reject(error) : resolve())); + }), + 5_000, + "fixture backend close", + ), + }; +}; + +const sanitizeServerEnvironment = (home) => { + const environment = { ...NodeProcess.env }; + for (const name of Object.keys(environment)) { + if ( + name.startsWith("PRIME_AGENT_INTERNAL_") || + name.startsWith("RLM_") || + name === "PRIME_AGENT_CODING_AGENT_DIR" || + name === "PI_CODING_AGENT_DIR" || + name === "PRIME_AGENT_REAL_PACKAGE_ROOT" || + name === "FORCE_COLOR" || + name === "VITEST" || + name.startsWith("VITEST_") || + name === "JEST_WORKER_ID" || + name === "NODE_CHANNEL_FD" || + name === "NODE_UNIQUE_ID" + ) { + delete environment[name]; + } + } + return { + ...environment, + HOME: home, + SHELL: "/bin/sh", + NO_COLOR: "1", + T3CODE_LOG_LEVEL: "Debug", + T3CODE_TRACE_TIMING_ENABLED: "false", + }; +}; + +const createPrimeFacade = async (temp, sourceRoot, sourceCommit, sourceTree) => { + const codingAgentRoot = NodePath.join(sourceRoot, "packages", "coding-agent"); + const sdkEntry = NodePath.join(codingAgentRoot, "dist", "index.js"); + const cliEntry = NodePath.join(codingAgentRoot, "dist", "bundle", "cli.js"); + const aiEntry = NodePath.join(sourceRoot, "packages", "ai", "dist", "index.js"); + for (const required of [sdkEntry, cliEntry, aiEntry]) { + await NodeFSP.access(required); + } + + const facadeRoot = NodePath.join(temp, "prime-package"); + await NodeFSP.mkdir(facadeRoot, { recursive: true, mode: 0o700 }); + const executable = NodePath.join(facadeRoot, "prime-agent"); + const moduleEntry = NodePath.join(facadeRoot, "index.mjs"); + const buildId = `pylon-build-g${sourceCommit.slice(0, 12)}-r1`; + await NodeFSP.writeFile( + executable, + `#!/usr/bin/env node\nprocess.argv[1] = ${JSON.stringify(cliEntry)};\nawait import(${JSON.stringify(NodeURL.pathToFileURL(cliEntry).href)});\n`, + { mode: 0o700 }, + ); + await NodeFSP.writeFile( + moduleEntry, + `export * from ${JSON.stringify(NodeURL.pathToFileURL(sdkEntry).href)};\n`, + "utf8", + ); + await NodeFSP.writeFile( + NodePath.join(facadeRoot, "package.json"), + `${JSON.stringify( + { + name: "prime-agent", + version: "0.8.1", + type: "module", + exports: "./index.mjs", + bin: { "prime-agent": "./prime-agent" }, + pylonDistribution: { + schemaVersion: 1, + repository: "https://github.com/pylon-code/prime-agent", + sourceCommit, + sourceTree, + buildId, + recipeRevision: 1, + node: "22.23.2", + npm: "11.10.1", + packageLockSha256: "0".repeat(64), + }, + }, + null, + 2, + )}\n`, + "utf8", + ); + return { facadeRoot, executable, sdkEntry, aiEntry, buildId }; +}; + +const writeFixtureModelConfig = async (agentHome, port) => { + await NodeFSP.mkdir(agentHome, { recursive: true, mode: 0o700 }); + await NodeFSP.writeFile( + NodePath.join(agentHome, "models.json"), + `${JSON.stringify( + { + providers: { + "faux-adoption": { + baseUrl: `http://127.0.0.1:${port}/v1`, + api: "openai-completions", + apiKey: '!printf "$PPID"', + authHeader: true, + compat: { + supportsDeveloperRole: false, + supportsReasoningEffort: false, + supportsUsageInStreaming: false, + maxTokensField: "max_tokens", + }, + models: [ + { + id: "faux-adoption", + name: "Faux Adoption", + reasoning: false, + input: ["text"], + contextWindow: 128_000, + maxTokens: 4_096, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + }, + ], + }, + }, + }, + null, + 2, + )}\n`, + { mode: 0o600 }, + ); +}; + +const preparePylonState = async ( + baseDir, + executable, + agentHome, + facadeRoot, + buildId, + sourceCommit, + sourceTree, +) => { + const stateDir = NodePath.join(baseDir, "userdata"); + await NodeFSP.mkdir(stateDir, { recursive: true, mode: 0o700 }); + await NodeFSP.writeFile( + NodePath.join(stateDir, "settings.json"), + `${JSON.stringify( + { + enableProviderUpdateChecks: false, + enableLegacyTokenStreaming: true, + providers: { + primeAgent: { + enabled: true, + binaryPath: executable, + agentHomePath: agentHome, + launchArgs: "", + customModels: [modelSelection.model], + }, + }, + }, + null, + 2, + )}\n`, + { mode: 0o600 }, + ); + await persistPrimeManagedReceipt({ + stateDir, + instanceId: providerInstanceId, + packageRoot: facadeRoot, + platform: NodeProcess.platform, + publication: { + channel: "preview", + sequenceEpoch: 1, + sequence: 1, + buildId, + sourceCommit, + sourceTree, + recipeRevision: 1, + rootAsset: "pylon-prime-agent-0.8.1.tgz", + rootSha256: "1".repeat(64), + }, + }); + return stateDir; +}; + +const spawnPylonServer = async ({ repoRoot, baseDir, projectDir, home, port, label }) => { + const output = []; + let outputBytes = 0; + let pairingBuffer = ""; + const pairing = createGate(); + const child = NodeChildProcess.spawn( + NodeProcess.execPath, + [ + NodePath.join(repoRoot, "apps", "server", "src", "bin.ts"), + "serve", + "--host", + "127.0.0.1", + "--port", + String(port), + "--base-dir", + baseDir, + projectDir, + ], + { + cwd: repoRoot, + env: sanitizeServerEnvironment(home), + stdio: ["ignore", "pipe", "pipe"], + }, + ); + const capture = (chunk) => { + outputBytes += chunk.length; + if (outputBytes > maximumOutputBytes) { + child.kill("SIGTERM"); + pairing.settle(Promise.reject(new Error(`${label} exceeded its output budget`))); + return; + } + output.push(chunk); + pairingBuffer += chunk.toString("utf8"); + const match = /Pairing URL:\s+(https?:\/\/[^\s]+#token=([^\s]+))/u.exec(pairingBuffer); + if (match) pairing.settle({ pairingUrl: match[1], credential: match[2] }); + if (pairingBuffer.length > 128 * 1024) pairingBuffer = pairingBuffer.slice(-64 * 1024); + }; + child.stdout.on("data", capture); + child.stderr.on("data", capture); + child.once("error", (error) => { + if (!pairing.settled) pairing.settle(Promise.reject(error)); + }); + child.once("exit", (code, signal) => { + if (!pairing.settled) { + pairing.settle( + Promise.reject( + new Error( + `${label} exited before publishing pairing readiness (${code ?? signal}): ${Buffer.concat(output).toString("utf8")}`, + ), ), ); - expect(authority.cursor).toEqual( - expect.objectContaining({ generation: expect.any(String), sequence: expect.any(Number) }), - ); + } + }); + const access = await withSafetyCeiling(pairing.promise, 30_000, `${label} pairing readiness`); + return { + child, + access, + output: () => Buffer.concat(output).toString("utf8"), + baseUrl: `http://127.0.0.1:${port}`, + }; +}; - const sdk = await import(pathToFileURL(sdkEntry).href); - const client = new sdk.DaemonClient(socket); - await client.connect(); - const hello = await client.waitForHello(); - expect(hello.supervisorGeneration).toBe(authority.supervisorGeneration); - const adoptionRequestId = randomUUID(); - const adopted = await sdk.adoptRecoverableOwnedSession(client, { - requestId: adoptionRequestId, - recoveryHandle: authority.recoveryHandle, - expectedSupervisorGeneration: authority.supervisorGeneration, - activeSessionId: authority.activeSessionId, - sessionId: authority.sessionId, - correlationId: "correlation-real-1", - cursor: authority.cursor, - previousMcpOwnerId: "pylon:mcp-real-1", - mcpOwnerId: "pylon:mcp-real-2", - config: authority.config, - launchEnv, - connectionOptions: { closeClientOnDispose: false, supportsExtensionUi: true }, - }); - expect(adopted.recoveryHandle).not.toBe(authority.recoveryHandle); - expect(adopted.proof.ownershipGeneration).toBeGreaterThan(0); - await sdk.confirmRecoverableOwnedSessionAdoption(client, { - requestId: adoptionRequestId, - recoveryHandle: adopted.recoveryHandle, - proof: adopted.proof, - }); - const cleanup = await adopted.connection.disposeOwnedSession({ timeoutMs: 10_000 }); - expect(["completed", "already_completed"]).toContain(cleanup.status); - await client.request({ type: "shutdown" }, 5_000); - client.close(); - const exit = await waitForExit(daemon, 10_000); - expect(exit.code).toBe(0); +const fetchJson = async (url, options, label) => { + const response = await withSafetyCeiling(fetch(url, options), 10_000, label); + const text = await withSafetyCeiling(response.text(), 10_000, `${label} body`); + if (!response.ok) throw new Error(`${label} returned ${response.status}: ${text}`); + return text.length === 0 ? undefined : JSON.parse(text); +}; + +const exchangePairingCredential = (baseUrl, credential) => + fetchJson( + `${baseUrl}/oauth/token`, + { + method: "POST", + headers: { "content-type": "application/x-www-form-urlencoded" }, + body: new URLSearchParams({ + grant_type: "urn:ietf:params:oauth:grant-type:token-exchange", + subject_token: credential, + subject_token_type: "urn:t3:params:oauth:token-type:environment-bootstrap", + requested_token_type: "urn:ietf:params:oauth:token-type:access_token", + client_label: "Prime restart adoption proof", + client_device_type: "desktop", + client_os: NodeProcess.platform, + }), + }, + "pairing-token exchange", + ); + +const issueWebSocketUrl = async (baseUrl, bearerToken) => { + const issued = await fetchJson( + `${baseUrl}/api/auth/websocket-ticket`, + { method: "POST", headers: { authorization: `Bearer ${bearerToken}` } }, + "websocket ticket issuance", + ); + const url = new URL(baseUrl.replace(/^http/u, "ws")); + url.pathname = "/ws"; + url.searchParams.set("wsTicket", issued.ticket); + return url.toString(); +}; + +const wsRpcProtocolLayer = (url) => + RpcClient.layerProtocolSocket().pipe( + Layer.provide( + Socket.layerWebSocket(url).pipe( + Layer.provide( + Layer.succeed( + Socket.WebSocketConstructor, + (socketUrl, protocols) => new NodeSocket.NodeWS.WebSocket(socketUrl, protocols), + ), + ), + ), + ), + Layer.provide(RpcSerialization.layerJson), + ); + +const makeWsRpcClient = RpcClient.make(WsRpcGroup); +const runRpc = (wsUrl, operation, label) => + withSafetyCeiling( + Effect.runPromise( + Effect.scoped( + makeWsRpcClient.pipe(Effect.flatMap(operation), Effect.provide(wsRpcProtocolLayer(wsUrl))), + ), + ), + 15_000, + label, + ); + +const dispatch = (wsUrl, command, label) => + runRpc(wsUrl, (client) => client[ORCHESTRATION_WS_METHODS.dispatchCommand](command), label); + +const readThreadSnapshot = (baseUrl, bearerToken, threadId) => + fetchJson( + `${baseUrl}/api/orchestration/threads/${encodeURIComponent(threadId)}`, + { headers: { authorization: `Bearer ${bearerToken}` } }, + "thread snapshot", + ); + +const readDaemonState = async (sdkEntry, socketPath) => { + const sdk = await import(NodeURL.pathToFileURL(sdkEntry).href); + const client = new sdk.DaemonClient(socketPath); + try { + await client.connect(2_000); + await client.waitForHello(2_000); + const response = await client.request({ type: "list", includeClientOwned: true }, 3_000); + if (response.success !== true || !Array.isArray(response.data?.sessions)) { + throw new Error(`invalid Prime daemon list response: ${JSON.stringify(response)}`); + } + return response.data; + } finally { + client.close(); + } +}; + +const readLedger = (databasePath, threadId) => { + const database = new NodeSqlite.DatabaseSync(databasePath, { readOnly: true }); + try { + return database + .prepare( + `SELECT thread_id, provider_instance_id, session_incarnation_id, admission_request_id, + turn_id, package_root, active_session_id, native_session_id, recovery_handle, + supervisor_generation, ownership_generation, cursor_generation, cursor_sequence, + correlation_id, mcp_owner_id, owner_token, state + FROM prime_agent_recovery_ledger WHERE thread_id = ?`, + ) + .get(threadId); + } finally { + database.close(); + } +}; + +const readProviderSessionRuntime = (databasePath, threadId) => { + const database = new NodeSqlite.DatabaseSync(databasePath, { readOnly: true }); + try { + const row = database + .prepare( + `SELECT provider_instance_id, status, runtime_payload_json + FROM provider_session_runtime WHERE thread_id = ?`, + ) + .get(threadId); + if (row === undefined) return undefined; + return { + providerInstanceId: row.provider_instance_id, + status: row.status, + runtimePayload: + typeof row.runtime_payload_json === "string" ? JSON.parse(row.runtime_payload_json) : {}, + }; + } finally { + database.close(); + } +}; + +const waitForDurableActiveRuntime = async (databasePath, threadId, expected, timeoutMs) => { + const deadline = Date.now() + timeoutMs; + let observed; + while (Date.now() < deadline) { + observed = readProviderSessionRuntime(databasePath, threadId); + if ( + observed?.status === "running" && + observed.providerInstanceId === providerInstanceId && + observed.runtimePayload.activeTurnId === expected.turnId && + observed.runtimePayload.activeTurnRequestId === expected.admissionRequestId && + observed.runtimePayload.sessionIncarnationId === expected.sessionIncarnationId + ) { + return observed; + } + await new Promise((resolve) => setTimeout(resolve, 20)); + } + throw new Error( + `provider runtime did not become durably active before restart (status=${String(observed?.status)})`, + ); +}; + +const readRawThreadEvents = (databasePath, threadId) => { + const database = new NodeSqlite.DatabaseSync(databasePath, { readOnly: true }); + try { + return database + .prepare( + "SELECT sequence, event_type, payload_json FROM orchestration_events WHERE stream_id = ? ORDER BY sequence", + ) + .all(threadId) + .map((row) => ({ ...row, payload: JSON.parse(row.payload_json) })); + } finally { + database.close(); + } +}; + +const shutdownCapturedDaemon = async (sdkEntry, socketPath) => { + if (sdkEntry === undefined || socketPath === undefined) return; + try { + const sdk = await import(NodeURL.pathToFileURL(sdkEntry).href); + const client = new sdk.DaemonClient(socketPath); + try { + await client.connect(2_000); + await client.waitForHello(2_000); + const response = await client.request({ type: "shutdown" }, 3_000); + if (response.success !== true) throw new Error("Prime daemon rejected shutdown"); } finally { - if (daemon.exitCode === null && daemon.signalCode === null) daemon.kill("SIGTERM"); - await waitForExit(daemon, 5_000).catch(() => undefined); - await FSP.rm(temp, { recursive: true, force: true }); - if (daemon.exitCode && daemon.exitCode !== 0) { - throw new Error(Buffer.concat(daemonErrors).toString("utf8")); - } + client.close(); } - }, 40_000); -}); + } catch (error) { + if ( + error && + typeof error === "object" && + (error.code === "ENOENT" || + (error instanceof Error && error.message.includes(`ENOENT ${socketPath}`))) + ) { + return; + } + throw error; + } +}; + +const stopCaptured = async (server, signal = "SIGTERM") => { + if (server === undefined) return; + const { child } = server; + if (child.exitCode === null && child.signalCode === null) child.kill(signal); + try { + await waitForExit(child, 10_000, `captured Pylon ${signal} exit`); + } catch (error) { + if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL"); + await waitForExit(child, 5_000, "captured Pylon forced exit").catch(() => undefined); + throw error; + } +}; + +const processExists = (pid) => { + try { + NodeProcess.kill(pid, 0); + return true; + } catch (error) { + if (error && typeof error === "object" && error.code === "ESRCH") return false; + throw error; + } +}; + +const waitForProcessExit = async (pid, timeoutMs, label) => { + const deadline = Date.now() + timeoutMs; + while (processExists(pid)) { + if (Date.now() >= deadline) throw new Error(`${label} exceeded its safety ceiling`); + await new Promise((resolve) => setTimeout(resolve, 20)); + } +}; + +const listSessionJsonlFiles = async (stateDir) => { + const root = NodePath.join(stateDir, "provider-sessions", "prime-agent"); + const found = []; + const visit = async (directory) => { + for (const entry of await NodeFSP.readdir(directory, { withFileTypes: true })) { + const path = NodePath.join(directory, entry.name); + if (entry.isDirectory()) await visit(path); + else if (entry.isFile() && entry.name.endsWith(".jsonl")) found.push(path); + } + }; + await visit(root); + return found; +}; + +const runRestartedTurn = ({ wsUrl, threadId, fixture, onRecoveredActivity }) => + Effect.runPromise( + Effect.scoped( + Effect.gen(function* () { + const client = yield* makeWsRpcClient; + const synchronized = yield* Deferred.make(); + const firstAssistantMessage = yield* Deferred.make(); + const firstCheckpoint = yield* Deferred.make(); + const secondCheckpoint = yield* Deferred.make(); + const stopped = yield* Deferred.make(); + const items = []; + let checkpointCount = 0; + let stopRequested = false; + + yield* client[ORCHESTRATION_WS_METHODS.subscribeThread]({ + threadId, + afterSequence: 0, + requestCompletionMarker: true, + }).pipe( + Stream.runForEach((item) => + Effect.gen(function* () { + items.push(item); + if (item.kind === "synchronized") yield* Deferred.succeed(synchronized, undefined); + if (item.kind !== "event") return; + if ( + item.event.type === "thread.message-sent" && + item.event.payload.role === "assistant" && + item.event.payload.text === "FIRST_TURN_COMPLETE" + ) { + yield* Deferred.succeed(firstAssistantMessage, item.event); + } + if (item.event.type === "thread.turn-diff-completed") { + checkpointCount += 1; + if (checkpointCount === 1) yield* Deferred.succeed(firstCheckpoint, item.event); + if (checkpointCount === 2) yield* Deferred.succeed(secondCheckpoint, item.event); + } + if ( + stopRequested && + item.event.type === "thread.session-set" && + item.event.payload.session.status === "stopped" + ) { + yield* Deferred.succeed(stopped, item.event); + } + }), + ), + Effect.forkScoped, + ); + + yield* Deferred.await(synchronized); + expect(fixture.records).toHaveLength(1); + yield* Effect.sync(() => fixture.releaseFirst()); + const firstAssistantMessageEvent = yield* Deferred.await(firstAssistantMessage); + yield* Effect.tryPromise(() => onRecoveredActivity(firstAssistantMessageEvent)); + const firstCheckpointEvent = yield* Deferred.await(firstCheckpoint); + + const secondCommandId = `cmd-${NodeCrypto.randomUUID()}`; + yield* client[ORCHESTRATION_WS_METHODS.dispatchCommand]({ + type: "thread.turn.start", + commandId: secondCommandId, + threadId, + message: { + messageId: `message-${NodeCrypto.randomUUID()}`, + role: "user", + text: "second prompt proves exact native transcript continuity", + attachments: [], + }, + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + createdAt: new Date().toISOString(), + }); + const secondCheckpointEvent = yield* Deferred.await(secondCheckpoint); + + stopRequested = true; + yield* client[ORCHESTRATION_WS_METHODS.dispatchCommand]({ + type: "thread.session.stop", + commandId: `stop-${NodeCrypto.randomUUID()}`, + threadId, + createdAt: new Date().toISOString(), + }); + const stoppedEvent = yield* Deferred.await(stopped); + return { items, firstCheckpointEvent, secondCheckpointEvent, stoppedEvent }; + }).pipe(Effect.provide(wsRpcProtocolLayer(wsUrl))), + ), + ); + +describe.skipIf(!enabled)( + `Prime Agent two-Pylon-server restart adoption (${enabled ? "enabled" : skipReason})`, + () => { + it( + "adopts one live owned worker across the real server boundary and cleans it authoritatively", + async () => { + const repoRoot = NodePath.resolve(import.meta.dirname, "../../../../.."); + const sourceRoot = NodePath.resolve(packageRoot); + const sourceHead = ( + await runCaptured( + "git", + ["-C", sourceRoot, "rev-parse", "HEAD"], + { stdio: ["ignore", "pipe", "pipe"] }, + 5_000, + "Prime source HEAD", + ) + ).trim(); + const sourceTree = ( + await runCaptured( + "git", + ["-C", sourceRoot, "rev-parse", "HEAD^{tree}"], + { stdio: ["ignore", "pipe", "pipe"] }, + 5_000, + "Prime source tree", + ) + ).trim(); + expect(sourceHead).toBe(exactHead); + + const temp = await NodeFSP.realpath( + await NodeFSP.mkdtemp(NodePath.join(NodeOS.tmpdir(), "pylon-two-server-adoption-")), + ); + const home = NodePath.join(temp, "home"); + const baseDir = NodePath.join(temp, "server-home"); + const agentHome = NodePath.join(home, ".prime", "agent"); + const projectDir = NodePath.join(temp, "project"); + let fixture; + let serverA; + let serverB; + let bearerToken; + let daemonSocket; + let primeSdkEntry; + let workerPid; + let secondWorkerPid; + try { + await Promise.all([ + NodeFSP.mkdir(home, { recursive: true, mode: 0o700 }), + NodeFSP.mkdir(agentHome, { recursive: true, mode: 0o700 }), + NodeFSP.mkdir(projectDir, { recursive: true, mode: 0o700 }), + ]); + await runCaptured( + "git", + ["init", "--initial-branch=main", projectDir], + { stdio: ["ignore", "pipe", "pipe"], env: sanitizeServerEnvironment(home) }, + 10_000, + "fixture repository initialization", + ); + await runCaptured( + "git", + [ + "-C", + projectDir, + "-c", + "user.name=Pylon Test", + "-c", + "user.email=pylon@test.invalid", + "commit", + "--allow-empty", + "-m", + "fixture", + ], + { stdio: ["ignore", "pipe", "pipe"], env: sanitizeServerEnvironment(home) }, + 10_000, + "fixture repository seed commit", + ); + + fixture = await startFixtureBackend(); + const primeFacade = await createPrimeFacade(temp, sourceRoot, sourceHead, sourceTree); + primeSdkEntry = primeFacade.sdkEntry; + await writeFixtureModelConfig(agentHome, fixture.port); + const stateDir = await preparePylonState( + baseDir, + primeFacade.executable, + agentHome, + primeFacade.facadeRoot, + primeFacade.buildId, + sourceHead, + sourceTree, + ); + const databasePath = NodePath.join(stateDir, "state.sqlite"); + daemonSocket = NodePath.join( + NodeOS.tmpdir(), + `pylon-prime-agent-${NodeCrypto.createHash("sha256") + .update(`${NodePath.resolve(stateDir)}\0${providerInstanceId}`) + .digest("hex") + .slice(0, 20)}`, + "daemon.sock", + ); + const portA = await reserveEphemeralPort(); + serverA = await spawnPylonServer({ + repoRoot, + baseDir, + projectDir, + home, + port: portA, + label: "server A", + }); + const exchanged = await exchangePairingCredential( + serverA.baseUrl, + serverA.access.credential, + ); + bearerToken = exchanged.access_token; + expect(typeof bearerToken).toBe("string"); + const wsA = await issueWebSocketUrl(serverA.baseUrl, bearerToken); + await runRpc( + wsA, + (client) => client[WS_METHODS.serverProbe]({}), + "server A command readiness", + ); + + const projectId = `project-${NodeCrypto.randomUUID()}`; + const threadId = `thread-${NodeCrypto.randomUUID()}`; + const createdAt = new Date().toISOString(); + await dispatch( + wsA, + { + type: "project.create", + commandId: `project-command-${NodeCrypto.randomUUID()}`, + projectId, + title: "Two-server adoption proof", + workspaceRoot: projectDir, + defaultModelSelection: modelSelection, + createdAt, + }, + "project creation", + ); + await dispatch( + wsA, + { + type: "thread.create", + commandId: `thread-command-${NodeCrypto.randomUUID()}`, + threadId, + projectId, + title: "Restart adoption", + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + branch: "main", + worktreePath: null, + createdAt, + }, + "thread creation", + ); + const firstPrompt = "first prompt must be admitted once before owner A exits"; + await dispatch( + wsA, + { + type: "thread.turn.start", + commandId: `turn-command-${NodeCrypto.randomUUID()}`, + threadId, + message: { + messageId: `message-${NodeCrypto.randomUUID()}`, + role: "user", + text: firstPrompt, + attachments: [], + }, + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + createdAt: new Date().toISOString(), + }, + "first prompt admission", + ); + + const firstNativeRecord = await withSafetyCeiling( + fixture.firstAdmission, + 5_000, + "first native provider activity", + ); + expect(firstNativeRecord.callCount).toBe(1); + expect( + firstNativeRecord.messages.filter((message) => message.text.includes(firstPrompt)), + ).toHaveLength(1); + const daemonStateA = await readDaemonState(primeFacade.sdkEntry, daemonSocket); + expect(daemonStateA).toMatchObject({ sessions: [], busyClientOwnedSessionCount: 1 }); + workerPid = firstNativeRecord.workerPid; + expect(workerPid).toEqual(expect.any(Number)); + const activeSnapshot = await readThreadSnapshot(serverA.baseUrl, bearerToken, threadId); + expect(activeSnapshot.thread.id).toBe(threadId); + if (activeSnapshot.thread.session.status !== "running") { + throw new Error( + `${JSON.stringify(activeSnapshot.thread.session)}\n${serverA.output()}`, + ); + } + expect(activeSnapshot.thread.session).toMatchObject({ + status: "running", + providerInstanceId, + activeTurnId: expect.any(String), + }); + expect( + activeSnapshot.thread.messages.filter((message) => message.text === firstPrompt), + ).toHaveLength(1); + + const ledgerA = readLedger(databasePath, threadId); + expect(ledgerA).toMatchObject({ + thread_id: threadId, + provider_instance_id: providerInstanceId, + turn_id: activeSnapshot.thread.session.activeTurnId, + state: "active", + }); + const sessionFilesBefore = await listSessionJsonlFiles(stateDir); + expect(sessionFilesBefore).toHaveLength(1); + expect(await NodeFSP.lstat(daemonSocket)).toMatchObject({}); + await waitForDurableActiveRuntime( + databasePath, + threadId, + { + turnId: ledgerA.turn_id, + admissionRequestId: ledgerA.admission_request_id, + sessionIncarnationId: ledgerA.session_incarnation_id, + }, + 5_000, + ); + + serverA.child.kill("SIGKILL"); + const ownerAExit = await waitForExit(serverA.child, 10_000, "server A abrupt exit"); + expect(ownerAExit.signal).toBe("SIGKILL"); + expect(processExists(workerPid)).toBe(true); + expect(fixture.records).toHaveLength(1); + expect(await readDaemonState(primeFacade.sdkEntry, daemonSocket)).toMatchObject({ + sessions: [], + busyClientOwnedSessionCount: 1, + }); + + const portB = await reserveEphemeralPort(); + serverB = await spawnPylonServer({ + repoRoot, + baseDir, + projectDir, + home, + port: portB, + label: "server B", + }); + const wsB = await issueWebSocketUrl(serverB.baseUrl, bearerToken); + await runRpc( + wsB, + (client) => client[WS_METHODS.serverProbe]({}), + "server B command readiness after adoption", + ); + expect(fixture.records).toHaveLength(1); + if (!processExists(workerPid)) { + throw new Error( + `captured Prime worker ${workerPid} exited before recovered activity\n${serverB.output()}`, + ); + } + + let ledgerB; + const restarted = await withSafetyCeiling( + runRestartedTurn({ + wsUrl: wsB, + threadId, + fixture, + onRecoveredActivity: async () => { + ledgerB = readLedger(databasePath, threadId); + expect(ledgerB).toMatchObject({ + thread_id: ledgerA.thread_id, + session_incarnation_id: ledgerA.session_incarnation_id, + admission_request_id: ledgerA.admission_request_id, + turn_id: ledgerA.turn_id, + active_session_id: ledgerA.active_session_id, + native_session_id: ledgerA.native_session_id, + supervisor_generation: ledgerA.supervisor_generation, + correlation_id: ledgerA.correlation_id, + state: "active", + }); + expect(ledgerB.recovery_handle).not.toBe(ledgerA.recovery_handle); + expect(ledgerB.owner_token).not.toBe(ledgerA.owner_token); + expect(ledgerB.mcp_owner_id).not.toBe(ledgerA.mcp_owner_id); + expect(ledgerB.ownership_generation).toBeGreaterThan(ledgerA.ownership_generation); + expect(ledgerB.cursor_sequence).toBeGreaterThanOrEqual(ledgerA.cursor_sequence); + const daemonStateB = await readDaemonState(primeFacade.sdkEntry, daemonSocket); + expect(daemonStateB).toMatchObject({ + sessions: [], + busyClientOwnedSessionCount: 1, + }); + fixture.finishFirst(); + }, + }), + 60_000, + "restarted turn completion and explicit stop", + ); + const secondNativeRecord = await withSafetyCeiling( + fixture.secondAdmission, + 10_000, + "second native provider activity", + ); + expect(fixture.records).toHaveLength(2); + secondWorkerPid = secondNativeRecord.workerPid; + expect(secondWorkerPid).not.toBe(workerPid); + expect( + secondNativeRecord.messages.filter((message) => message.text.includes(firstPrompt)), + ).toHaveLength(1); + expect( + secondNativeRecord.messages.filter((message) => + message.text.includes("FIRST_TURN_COMPLETE"), + ), + ).toHaveLength(1); + expect( + secondNativeRecord.messages.filter((message) => + message.text.includes("second prompt proves exact native transcript continuity"), + ), + ).toHaveLength(1); + + const publicEvents = restarted.items.filter((item) => item.kind === "event"); + const persistedThreadEvents = readRawThreadEvents(databasePath, threadId); + expect( + persistedThreadEvents.filter( + (event) => event.event_type === "thread.turn-start-requested", + ), + ).toHaveLength(2); + const checkpointEvents = persistedThreadEvents.filter( + (event) => event.event_type === "thread.turn-diff-completed", + ); + expect(checkpointEvents).toHaveLength(2); + expect(new Set(checkpointEvents.map((event) => event.payload.turnId)).size).toBe(2); + expect(restarted.firstCheckpointEvent.payload.turnId).toBe(ledgerA.turn_id); + expect(restarted.secondCheckpointEvent.payload.turnId).not.toBe(ledgerA.turn_id); + expect(checkpointEvents.map((event) => event.payload.status)).toEqual(["ready", "ready"]); + + const finalSnapshot = await readThreadSnapshot(serverB.baseUrl, bearerToken, threadId); + expect(finalSnapshot.thread.id).toBe(threadId); + expect(finalSnapshot.thread.session).toMatchObject({ + status: "stopped", + activeTurnId: null, + }); + expect( + finalSnapshot.thread.messages.filter((message) => message.text === firstPrompt), + ).toHaveLength(1); + expect( + finalSnapshot.thread.messages.filter( + (message) => message.text === "FIRST_TURN_COMPLETE", + ), + ).toHaveLength(1); + expect( + finalSnapshot.thread.messages.filter( + (message) => message.text === "SECOND_TURN_COMPLETE", + ), + ).toHaveLength(1); + expect(readLedger(databasePath, threadId)).toBeUndefined(); + expect(await listSessionJsonlFiles(stateDir)).toEqual(sessionFilesBefore); + await waitForProcessExit(workerPid, 5_000, "adopted Prime worker exit after its turn"); + await waitForProcessExit(secondWorkerPid, 5_000, "second Prime worker exit after Stop"); + + const publicSurface = JSON.stringify({ snapshot: finalSnapshot, events: publicEvents }); + const privateValues = [ + ledgerA.recovery_handle, + ledgerB.recovery_handle, + ledgerA.owner_token, + ledgerB.owner_token, + ledgerA.active_session_id, + ledgerA.native_session_id, + ledgerA.cursor_generation, + ledgerA.correlation_id, + ledgerA.mcp_owner_id, + ledgerB.mcp_owner_id, + primeFacade.facadeRoot, + sourceRoot, + home, + agentHome, + daemonSocket, + serverA.access.credential, + bearerToken, + String(serverA.child.pid), + String(serverB.child.pid), + String(workerPid), + ]; + for (const privateValue of privateValues) { + expect(publicSurface).not.toContain(privateValue); + } + const logSafeResult = [serverA.output(), serverB.output()] + .join("\n") + .replace(/^.*(?:Pairing URL|Connection string):.*$/gmu, "[startup access redacted]"); + for (const privateValue of [ + ledgerA.recovery_handle, + ledgerB.recovery_handle, + ledgerA.owner_token, + ledgerB.owner_token, + ledgerA.active_session_id, + ledgerA.native_session_id, + ledgerA.cursor_generation, + ledgerA.correlation_id, + ledgerA.mcp_owner_id, + ledgerB.mcp_owner_id, + daemonSocket, + ]) { + expect(logSafeResult).not.toContain(privateValue); + } + + await waitForProcessExit(workerPid, 5_000, "adopted Prime worker exit after its turn"); + await waitForProcessExit(secondWorkerPid, 5_000, "second Prime worker exit after Stop"); + await stopCaptured(serverB, "SIGTERM"); + serverB = undefined; + await expect(NodeFSP.access(daemonSocket)).rejects.toMatchObject({ code: "ENOENT" }); + } finally { + await stopCaptured(serverB, "SIGTERM").catch(() => undefined); + await stopCaptured(serverA, "SIGKILL").catch(() => undefined); + if (fixture !== undefined) { + await fixture.close().catch(() => undefined); + } + if (primeSdkEntry !== undefined && daemonSocket !== undefined) { + await shutdownCapturedDaemon(primeSdkEntry, daemonSocket); + } + for (const [pid, label] of [ + [workerPid, "adopted Prime worker test cleanup"], + [secondWorkerPid, "second Prime worker test cleanup"], + ]) { + if (pid !== undefined) await waitForProcessExit(pid, 5_000, label); + } + await NodeFSP.rm(temp, { recursive: true, force: true, maxRetries: 5, retryDelay: 50 }); + } + }, + outerSafetyMs, + ); + }, +); From a8b8cde2cab242c28965743c1d84d222aca053aa Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Tue, 1 Sep 2026 11:21:38 -0600 Subject: [PATCH 3/4] fix(prime): recover interrupted session adoption Refs #84 --- .../050_PrimeAgentRecoveryLedger.ts | 8 + .../provider/prime/PrimeAgentDaemonAdapter.ts | 487 ++++++++++++++++-- .../PrimeAgentDaemonSessionRuntime.test.ts | 21 + .../prime/PrimeAgentDaemonSessionRuntime.ts | 50 +- .../prime/PrimeAgentRecoveryLedger.test.ts | 368 +++++++++++-- .../prime/PrimeAgentRecoveryLedger.ts | 287 ++++++++++- .../PrimeAgentRestartAdoption.real.test.mjs | 177 ++++++- 7 files changed, 1283 insertions(+), 115 deletions(-) diff --git a/apps/server/src/persistence/Migrations/050_PrimeAgentRecoveryLedger.ts b/apps/server/src/persistence/Migrations/050_PrimeAgentRecoveryLedger.ts index 5b2d7098d..729c862f2 100644 --- a/apps/server/src/persistence/Migrations/050_PrimeAgentRecoveryLedger.ts +++ b/apps/server/src/persistence/Migrations/050_PrimeAgentRecoveryLedger.ts @@ -34,6 +34,14 @@ export default Effect.gen(function* () { transcript_fingerprints_json TEXT NOT NULL DEFAULT '[]', owner_token TEXT NOT NULL, state TEXT NOT NULL, + adoption_previous_owner_token TEXT, + adoption_owner_token TEXT, + adoption_request_id TEXT, + adoption_mcp_owner_id TEXT, + adoption_phase TEXT, + adoption_attempt INTEGER NOT NULL DEFAULT 0, + adoption_recovery_handle TEXT, + adoption_proof_json TEXT, native_cleanup_proven INTEGER NOT NULL DEFAULT 0, terminal_projected INTEGER NOT NULL DEFAULT 0, checkpoint_quiesced INTEGER NOT NULL DEFAULT 0, diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts b/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts index 32810e785..a98a9d990 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonAdapter.ts @@ -1,4 +1,5 @@ import * as NodeCrypto from "node:crypto"; +import * as NodeUtil from "node:util"; import { ApprovalRequestId, @@ -93,7 +94,9 @@ import { } from "./PrimeAgentDaemonEvents.ts"; import type { PrimeAgentDaemonManager } from "./PrimeAgentDaemonManager.ts"; import { + PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS, PrimeAgentRecoveryLedger, + type PrimeAgentRecoveryAdoptionProof, type PrimeAgentRecoveryAuthority, type PrimeAgentRecoveryLedgerShape, } from "./PrimeAgentRecoveryLedger.ts"; @@ -149,6 +152,17 @@ export const PRIME_AGENT_FAILED_RUN_SETTLEMENT_GRACE_MS = 3_000; export const PRIME_AGENT_SESSION_TEARDOWN_TIMEOUT_MS = 5_000; const PRIME_AGENT_SESSION_CLEANUP_CONCURRENCY = 4; const PRIME_AGENT_TERMINAL_EVENT_TIMEOUT_MS = 100; +const PRIME_AGENT_RECOVERY_TEST_CRASH_STAGE = "PRIME_AGENT_INTERNAL_PYLON_RECOVERY_CRASH_STAGE"; +type PrimeAgentRecoveryTestCrashStage = + | "after-claim-persisted" + | "after-native-response-before-commit" + | "after-commit-before-confirm"; + +function crashAtPrimeAgentRecoveryTestBarrier(stage: PrimeAgentRecoveryTestCrashStage): void { + if (process.env[PRIME_AGENT_RECOVERY_TEST_CRASH_STAGE] === stage) { + process.kill(process.pid, "SIGKILL"); + } +} export const PRIME_AGENT_SIDE_QUESTION_TIMEOUT_MS = 2 * 60_000; const PRIME_AGENT_SIDE_QUESTION_MAX_ACTIVE = 4; const unavailableSessionGoal: SessionGoalUpdatedPayload = { @@ -546,6 +560,61 @@ function sameStringRecord( ); } +function isExactPrimeAgentAdoptionProof(input: { + readonly authority: PrimeAgentRecoveryAuthority; + readonly proof: PrimeAgentRecoveryAdoptionProof; + readonly mcpOwnerId: string; +}): boolean { + const { authority, proof } = input; + return ( + proof.feature === "recoverable_owned_session_adoption_v1" && + proof.status === "adopted" && + proof.supervisorGeneration === authority.supervisorGeneration && + proof.ownershipGeneration > authority.ownershipGeneration && + proof.activeSessionId === authority.activeSessionId && + proof.sessionId === authority.nativeSessionId && + proof.correlationId === authority.correlationId && + proof.mcpOwnerId === input.mcpOwnerId && + proof.cursor.generation === authority.cursor.generation && + proof.cursor.sequence >= authority.cursor.sequence + ); +} + +function samePrimeAgentAdoptionProof( + left: PrimeAgentRecoveryAdoptionProof | null, + right: PrimeAgentRecoveryAdoptionProof, +): boolean { + return left !== null && NodeUtil.isDeepStrictEqual(left, right); +} + +function completePrimeAgentAdoptionRoute(authority: PrimeAgentRecoveryAuthority): + | { + readonly previousOwnerToken: string; + readonly ownerToken: string; + readonly requestId: string; + readonly mcpOwnerId: string; + } + | undefined { + if ( + authority.state !== "adopting" || + authority.adoptionPreviousOwnerToken !== authority.ownerToken || + authority.adoptionOwnerToken === null || + authority.adoptionRequestId === null || + !/^[0-9a-f]{48}$/u.test(authority.adoptionRequestId) || + authority.adoptionMcpOwnerId === null || + authority.adoptionPhase === null || + authority.adoptionPhase === "quarantined" + ) { + return undefined; + } + return { + previousOwnerToken: authority.adoptionPreviousOwnerToken, + ownerToken: authority.adoptionOwnerToken, + requestId: authority.adoptionRequestId, + mcpOwnerId: authority.adoptionMcpOwnerId, + }; +} + function primeDaemonMessageFingerprint(message: PrimeDaemonMessage): string { return NodeCrypto.createHash("sha256").update(JSON.stringify(message), "utf8").digest("hex"); } @@ -747,11 +816,23 @@ export function makePrimeAgentDaemonAdapter( ) => rawRecoveryLedger.updateTranscriptProgress(input).pipe(Effect.orDie), claim: (input: Parameters[0]) => rawRecoveryLedger.claim(input).pipe(Effect.orDie), + beginAdoptionAttempt: ( + input: Parameters[0], + ) => rawRecoveryLedger.beginAdoptionAttempt(input).pipe(Effect.orDie), releaseClaim: (input: Parameters[0]) => rawRecoveryLedger.releaseClaim(input).pipe(Effect.orDie), commitAdoption: ( input: Parameters[0], ) => rawRecoveryLedger.commitAdoption(input).pipe(Effect.orDie), + beginAdoptionConfirmation: ( + input: Parameters[0], + ) => rawRecoveryLedger.beginAdoptionConfirmation(input).pipe(Effect.orDie), + finalizeAdoption: ( + input: Parameters[0], + ) => rawRecoveryLedger.finalizeAdoption(input).pipe(Effect.orDie), + quarantineAdoption: ( + input: Parameters[0], + ) => rawRecoveryLedger.quarantineAdoption(input).pipe(Effect.orDie), markNativeCleanup: ( input: Parameters[0], ) => rawRecoveryLedger.markNativeCleanup(input).pipe(Effect.orDie), @@ -904,6 +985,18 @@ export function makePrimeAgentDaemonAdapter( }), ), ); + const randomAdoptionRequestId = crypto.randomBytes(24).pipe( + Effect.map((bytes) => Buffer.from(bytes).toString("hex")), + Effect.mapError( + (cause) => + new ProviderAdapterRequestError({ + provider: PROVIDER, + method: "crypto/randomBytes", + detail: "Failed to generate Prime Agent adoption request authority.", + cause, + }), + ), + ); const makeEventStamp = () => Effect.all({ eventId: Effect.map(randomUUIDv4, EventId.make), createdAt: nowIso }); const offerRuntimeEvent = (event: ProviderRuntimeEvent) => @@ -3858,7 +3951,10 @@ export function makePrimeAgentDaemonAdapter( ? {} : { mcpServer: { - ownerId: `pylon:${mcpSession.providerSessionId}`, + ownerId: + recoveryStart?.kind === "adopt" + ? recoveryStart.mcpOwnerId + : `pylon:${mcpSession.providerSessionId}`, server: { name: "t3-code", type: "http" as const, @@ -3940,6 +4036,14 @@ export function makePrimeAgentDaemonAdapter( transcriptFingerprints: [...recoveryStart.transcriptFingerprints], ownerToken: recoveryStart.ownerToken, state: "prepared", + adoptionPreviousOwnerToken: null, + adoptionOwnerToken: null, + adoptionRequestId: null, + adoptionMcpOwnerId: null, + adoptionPhase: null, + adoptionAttempt: 0, + adoptionRecoveryHandle: null, + adoptionProof: null, nativeCleanupProven: false, terminalProjected: false, checkpointQuiesced: false, @@ -3964,16 +4068,65 @@ export function makePrimeAgentDaemonAdapter( mcpOwnerId: recoveryStart.mcpOwnerId, recoveryConfig: recoveryStart.authority.recoveryConfig, launchEnvironment: recoveryStart.authority.launchEnvironment, + onAdoptionAttemptStarted: () => + runPromise( + Effect.gen(function* () { + const started = yield* recoveryLedger!.beginAdoptionAttempt({ + threadId: input.threadId, + ownerToken: recoveryStart.ownerToken, + requestId: recoveryStart.requestId, + updatedAt: yield* nowIso, + }); + if (Option.isSome(started)) return; + const current = Option.getOrUndefined( + yield* recoveryLedger!.get(input.threadId), + ); + if ( + current?.state === "adopting" && + current.adoptionOwnerToken === recoveryStart.ownerToken && + current.adoptionRequestId === recoveryStart.requestId && + current.adoptionAttempt >= PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS + ) { + yield* recoveryLedger!.quarantineAdoption({ + threadId: input.threadId, + ownerToken: recoveryStart.ownerToken, + requestId: recoveryStart.requestId, + updatedAt: yield* nowIso, + }); + } + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Recoverable Prime Agent adoption retry was superseded.", + }); + }), + ), onAdoptionCommitted: ({ recoveryHandle, proof }) => runPromise( Effect.gen(function* () { + crashAtPrimeAgentRecoveryTestBarrier( + "after-native-response-before-commit", + ); + if ( + !isExactPrimeAgentAdoptionProof({ + authority: recoveryStart.authority, + proof, + mcpOwnerId: recoveryStart.mcpOwnerId, + }) + ) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: + "Recoverable Prime Agent returned mismatched ownership proof.", + }); + } const committed = yield* recoveryLedger!.commitAdoption({ threadId: input.threadId, ownerToken: recoveryStart.ownerToken, + requestId: recoveryStart.requestId, recoveryHandle, - ownershipGeneration: proof.ownershipGeneration, - cursor: proof.cursor, - mcpOwnerId: proof.mcpOwnerId, + proof, updatedAt: yield* nowIso, }); if (!committed) { @@ -3983,6 +4136,68 @@ export function makePrimeAgentDaemonAdapter( detail: "Recoverable Prime Agent ownership was superseded.", }); } + crashAtPrimeAgentRecoveryTestBarrier("after-commit-before-confirm"); + }), + ), + onAdoptionConfirming: ({ recoveryHandle, proof }) => + runPromise( + Effect.gen(function* () { + if ( + !isExactPrimeAgentAdoptionProof({ + authority: recoveryStart.authority, + proof, + mcpOwnerId: recoveryStart.mcpOwnerId, + }) + ) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Recoverable Prime Agent confirmation proof mismatched.", + }); + } + const confirming = yield* recoveryLedger!.beginAdoptionConfirmation({ + threadId: input.threadId, + ownerToken: recoveryStart.ownerToken, + requestId: recoveryStart.requestId, + updatedAt: yield* nowIso, + }); + if (Option.isNone(confirming)) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Recoverable Prime Agent confirmation was superseded.", + }); + } + if ( + confirming.value.adoptionRecoveryHandle !== recoveryHandle || + !samePrimeAgentAdoptionProof(confirming.value.adoptionProof, proof) + ) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Recoverable Prime Agent confirmation receipt mismatched.", + }); + } + }), + ), + onAdoptionConfirmed: ({ recoveryHandle, proof }) => + runPromise( + Effect.gen(function* () { + const finalized = yield* recoveryLedger!.finalizeAdoption({ + threadId: input.threadId, + ownerToken: recoveryStart.ownerToken, + requestId: recoveryStart.requestId, + recoveryHandle, + proof, + updatedAt: yield* nowIso, + }); + if (!finalized) { + return yield* new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Confirmed Prime Agent ownership was superseded.", + }); + } }), ), }, @@ -4518,23 +4733,57 @@ export function makePrimeAgentDaemonAdapter( readonly resumeCursor: unknown; }) { if (!recoveryPlatformEligible || input.runtimeMode !== "full-access") return null; - const authorityOption = yield* recoveryLedger!.get(input.threadId); - const authority = Option.getOrUndefined(authorityOption); + let authority = Option.getOrUndefined(yield* recoveryLedger!.get(input.threadId)); + const authorityMatches = (candidate: PrimeAgentRecoveryAuthority) => + candidate.threadId === input.threadId && + candidate.turnId !== null && + candidate.providerInstanceId === input.providerInstanceId && + candidate.sessionIncarnationId === input.sessionIncarnationId && + candidate.packageRoot === manager.bridge.packageRoot && + candidate.packageVersion === manager.bridge.version && + candidate.managedBuildId === options?.recoveryManagedBuildId && + candidate.protocolName === manager.bridge.protocolName && + candidate.protocolVersion === manager.bridge.protocolVersion && + sameStrings(candidate.sdkFeatures, manager.bridge.sdkFeatures ?? []); + const activeAuthorityHasAdoptionResidue = + authority?.state === "active" && + (authority.adoptionPreviousOwnerToken !== null || + authority.adoptionOwnerToken !== null || + authority.adoptionRequestId !== null || + authority.adoptionMcpOwnerId !== null || + authority.adoptionPhase !== null || + authority.adoptionAttempt !== 0 || + authority.adoptionRecoveryHandle !== null || + authority.adoptionProof !== null); if ( authority === undefined || - authority.state !== "active" || - authority.turnId === null || - authority.providerInstanceId !== input.providerInstanceId || - authority.sessionIncarnationId !== input.sessionIncarnationId || - authority.packageRoot !== manager.bridge.packageRoot || - authority.packageVersion !== manager.bridge.version || - authority.managedBuildId !== options?.recoveryManagedBuildId || - authority.protocolName !== manager.bridge.protocolName || - authority.protocolVersion !== manager.bridge.protocolVersion || - !sameStrings(authority.sdkFeatures, manager.bridge.sdkFeatures ?? []) + (authority.state !== "active" && authority.state !== "adopting") || + activeAuthorityHasAdoptionResidue || + !authorityMatches(authority) ) { return null; } + + const recoverySessionDir = authority.recoveryConfig.sessionDir; + const recoveryCwd = authority.recoveryConfig.cwd; + const expectedSessionDir = primeAgentSessionDirectory({ + stateDir: serverConfig.stateDir, + instanceId: boundInstanceId, + threadId: input.threadId, + join: path.join, + }); + if ( + typeof recoverySessionDir !== "string" || + path.resolve(recoverySessionDir) !== path.resolve(expectedSessionDir) || + typeof recoveryCwd !== "string" || + path.resolve(recoveryCwd) !== path.resolve(input.cwd) || + authority.nativeSessionId === "." || + authority.nativeSessionId === ".." || + path.basename(authority.nativeSessionId) !== authority.nativeSessionId + ) { + return null; + } + yield* manager.prepare().pipe( Effect.mapError( (cause) => @@ -4567,32 +4816,162 @@ export function makePrimeAgentDaemonAdapter( ) { return null; } - const recoverySessionDir = authority.recoveryConfig.sessionDir; - if (typeof recoverySessionDir !== "string") { + + if ( + authority.state === "adopting" && + (authority.adoptionPhase === "committed" || authority.adoptionPhase === "confirming") + ) { + const route = completePrimeAgentAdoptionRoute(authority); + const proof = authority.adoptionProof; + const recoveryHandle = authority.adoptionRecoveryHandle; + const confirmRecoverableOwnedSessionAdoption = + manager.bridge.confirmRecoverableOwnedSessionAdoption; + if ( + route === undefined || + proof === null || + recoveryHandle === null || + confirmRecoverableOwnedSessionAdoption === undefined || + !isExactPrimeAgentAdoptionProof({ + authority, + proof, + mcpOwnerId: route.mcpOwnerId, + }) + ) { + return null; + } + const confirming = yield* recoveryLedger!.beginAdoptionConfirmation({ + threadId: input.threadId, + ownerToken: route.ownerToken, + requestId: route.requestId, + updatedAt: yield* nowIso, + }); + if (Option.isNone(confirming)) { + if (authority.adoptionAttempt >= PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS) { + yield* recoveryLedger!.quarantineAdoption({ + threadId: input.threadId, + ownerToken: route.ownerToken, + requestId: route.requestId, + updatedAt: yield* nowIso, + }); + } + return null; + } + const confirmationClient = yield* manager.openClient().pipe( + Effect.mapError( + (cause) => + new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Could not confirm surviving Prime Agent ownership.", + cause, + }), + ), + ); + yield* Effect.tryPromise({ + try: () => + confirmRecoverableOwnedSessionAdoption(confirmationClient, { + requestId: route.requestId, + recoveryHandle, + proof, + }), + catch: (cause) => + new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: input.threadId, + detail: "Could not confirm surviving Prime Agent ownership.", + cause, + }), + }).pipe(Effect.ensuring(Effect.sync(() => confirmationClient.close()))); + const finalized = yield* recoveryLedger!.finalizeAdoption({ + threadId: input.threadId, + ownerToken: route.ownerToken, + requestId: route.requestId, + recoveryHandle, + proof, + updatedAt: yield* nowIso, + }); + if (!finalized) return null; + authority = Option.getOrUndefined(yield* recoveryLedger!.get(input.threadId)); + if ( + authority === undefined || + authority.state !== "active" || + !authorityMatches(authority) + ) { + return null; + } + } + + let route: + | { + readonly previousOwnerToken: string; + readonly ownerToken: string; + readonly requestId: string; + readonly mcpOwnerId: string; + } + | undefined; + if (authority.state === "active") { + const ownerToken = yield* randomUUIDv4; + const requestId = yield* randomAdoptionRequestId; + const mcpSession = McpProviderSession.readMcpProviderSession(input.threadId); + const mcpOwnerId = + mcpSession === undefined + ? `pylon:none:${yield* randomUUIDv4}` + : `pylon:${mcpSession.providerSessionId}`; + const claimed = yield* recoveryLedger!.claim({ + threadId: input.threadId, + expectedOwnerToken: authority.ownerToken, + nextOwnerToken: ownerToken, + requestId, + mcpOwnerId, + updatedAt: yield* nowIso, + }); + if (Option.isNone(claimed)) return null; + authority = claimed.value; + route = completePrimeAgentAdoptionRoute(authority); + crashAtPrimeAgentRecoveryTestBarrier("after-claim-persisted"); + } else { + route = completePrimeAgentAdoptionRoute(authority); + } + if ( + route === undefined || + authority.adoptionAttempt >= PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS + ) { + if (route !== undefined) { + yield* recoveryLedger!.quarantineAdoption({ + threadId: input.threadId, + ownerToken: route.ownerToken, + requestId: route.requestId, + updatedAt: yield* nowIso, + }); + } return null; } - const ownerToken = yield* randomUUIDv4; - const claimedAt = yield* nowIso; - const claimed = yield* recoveryLedger!.claim({ - threadId: input.threadId, - expectedOwnerToken: authority.ownerToken, - nextOwnerToken: ownerToken, - updatedAt: claimedAt, - }); - if (Option.isNone(claimed)) return null; - const requestId = yield* randomUUIDv4; - const mcpSession = McpProviderSession.readMcpProviderSession(input.threadId); - const mcpOwnerId = - mcpSession === undefined - ? `pylon:none:${yield* randomUUIDv4}` - : `pylon:${mcpSession.providerSessionId}`; + const stagedReceiptExpected = + authority.adoptionPhase === "committed" || authority.adoptionPhase === "confirming"; + const stagedReceiptComplete = + authority.adoptionRecoveryHandle !== null && authority.adoptionProof !== null; + const stagedReceiptPresent = + authority.adoptionRecoveryHandle !== null || authority.adoptionProof !== null; + if ( + stagedReceiptExpected !== stagedReceiptComplete || + (!stagedReceiptExpected && stagedReceiptPresent) || + (authority.adoptionProof !== null && + !isExactPrimeAgentAdoptionProof({ + authority, + proof: authority.adoptionProof, + mcpOwnerId: route.mcpOwnerId, + })) + ) { + return null; + } + pendingRecoveryStarts.set(input.threadId, { kind: "adopt", authority, - previousOwnerToken: authority.ownerToken, - ownerToken, - requestId, - mcpOwnerId, + previousOwnerToken: route.previousOwnerToken, + ownerToken: route.ownerToken, + requestId: route.requestId, + mcpOwnerId: route.mcpOwnerId, sessionFile: path.join(recoverySessionDir, `${authority.nativeSessionId}.jsonl`), }); const started = yield* Effect.result( @@ -4609,12 +4988,34 @@ export function makePrimeAgentDaemonAdapter( ); if (Result.isFailure(started)) { pendingRecoveryStarts.delete(input.threadId); - yield* recoveryLedger!.releaseClaim({ - threadId: input.threadId, - ownerToken, - previousOwnerToken: authority.ownerToken, - updatedAt: yield* nowIso, - }); + const current = Option.getOrUndefined(yield* recoveryLedger!.get(input.threadId)); + if ( + current?.state === "adopting" && + current.adoptionOwnerToken === route.ownerToken && + current.adoptionRequestId === route.requestId && + current.adoptionPhase === "claimed" && + current.adoptionAttempt === 0 + ) { + yield* recoveryLedger!.releaseClaim({ + threadId: input.threadId, + ownerToken: route.ownerToken, + previousOwnerToken: route.previousOwnerToken, + requestId: route.requestId, + updatedAt: yield* nowIso, + }); + } else if ( + current?.state === "adopting" && + current.adoptionOwnerToken === route.ownerToken && + current.adoptionRequestId === route.requestId && + current.adoptionAttempt >= PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS + ) { + yield* recoveryLedger!.quarantineAdoption({ + threadId: input.threadId, + ownerToken: route.ownerToken, + requestId: route.requestId, + updatedAt: yield* nowIso, + }); + } return null; } return started.success; diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts index e86b8cd7b..6d7eb6b49 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.test.ts @@ -13073,9 +13073,18 @@ describe("Prime Agent live activity privacy boundary", () => { mcpOwnerId: "pylon:mcp-2", recoveryConfig: { cwd: "/work/project" }, launchEnvironment: { HOME: "/private/home" }, + onAdoptionAttemptStarted: async () => { + side.captures.order.push("adoption-attempt-durable"); + }, onAdoptionCommitted: async () => { side.captures.order.push("ledger-committed"); }, + onAdoptionConfirming: async () => { + side.captures.order.push("confirmation-attempt-durable"); + }, + onAdoptionConfirmed: async () => { + side.captures.order.push("ledger-finalized"); + }, }, ); const initial = yield* Stream.runHead(runtime.events); @@ -13083,16 +13092,28 @@ describe("Prime Agent live activity privacy boundary", () => { expect(initial._tag).toBe("Some"); expect(side.captures.order).toEqual( expect.arrayContaining([ + "adoption-attempt-durable", "adopt-recoverable", "retain-daemon", "ledger-committed", + "confirmation-attempt-durable", "confirm-adoption", + "ledger-finalized", ]), ); expect(side.captures.order).not.toContain("replace-mcp"); + expect(side.captures.order.indexOf("adoption-attempt-durable")).toBeLessThan( + side.captures.order.indexOf("adopt-recoverable"), + ); expect(side.captures.order.indexOf("ledger-committed")).toBeLessThan( + side.captures.order.indexOf("confirmation-attempt-durable"), + ); + expect(side.captures.order.indexOf("confirmation-attempt-durable")).toBeLessThan( side.captures.order.indexOf("confirm-adoption"), ); + expect(side.captures.order.indexOf("confirm-adoption")).toBeLessThan( + side.captures.order.indexOf("ledger-finalized"), + ); yield* runtime.prompt({ text: "next prompt after adoption" }); expect(side.captures.order).toContain("replace-mcp"); diff --git a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts index 8cc6f3fd1..af862da3c 100644 --- a/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts +++ b/apps/server/src/provider/prime/PrimeAgentDaemonSessionRuntime.ts @@ -919,10 +919,19 @@ export interface PrimeAgentDaemonSessionRuntimeInput { readonly mcpOwnerId: string; readonly recoveryConfig: Readonly>; readonly launchEnvironment: Readonly>; + readonly onAdoptionAttemptStarted: () => Promise; readonly onAdoptionCommitted: (authority: { readonly recoveryHandle: string; readonly proof: PrimeAgentRecoverableOwnedSessionAdoptionProof; }) => Promise; + readonly onAdoptionConfirming: (authority: { + readonly recoveryHandle: string; + readonly proof: PrimeAgentRecoverableOwnedSessionAdoptionProof; + }) => Promise; + readonly onAdoptionConfirmed: (authority: { + readonly recoveryHandle: string; + readonly proof: PrimeAgentRecoverableOwnedSessionAdoptionProof; + }) => Promise; }; } @@ -1922,6 +1931,15 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo "Recoverable Prime Agent execution is unavailable.", ); } + yield* Effect.tryPromise({ + try: recovery.onAdoptionAttemptStarted, + catch: () => + runtimeError( + "attach-session", + "request-failed", + "Recoverable Prime Agent adoption attempt could not be durably recorded.", + ), + }).pipe(Effect.onError(() => closeClient)); const adopted = yield* Effect.tryPromise({ try: () => adoptRecoverableOwnedSession(client, { @@ -7747,7 +7765,28 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo ); } - if (confirmAdoption !== undefined) { + if (confirmAdoption !== undefined && adoptedRecovery !== undefined) { + const adoptionRecovery = input.recovery; + if (adoptionRecovery?.kind !== "adopt") { + return yield* runtimeError( + "attach-session", + "invalid-response", + "Recoverable Prime Agent adoption state was lost before confirmation.", + ); + } + const adoptionAuthority = { + recoveryHandle: adoptedRecovery.recoveryHandle, + proof: adoptedRecovery.proof, + }; + yield* Effect.tryPromise({ + try: () => adoptionRecovery.onAdoptionConfirming(adoptionAuthority), + catch: () => + runtimeError( + "attach-session", + "request-failed", + "Recoverable Prime Agent confirmation attempt could not be durably recorded.", + ), + }); yield* Effect.tryPromise({ try: confirmAdoption, catch: () => @@ -7764,6 +7803,15 @@ export const makePrimeAgentDaemonSessionRuntime = Effect.fn("makePrimeAgentDaemo }), ), ); + yield* Effect.tryPromise({ + try: () => adoptionRecovery.onAdoptionConfirmed(adoptionAuthority), + catch: () => + runtimeError( + "attach-session", + "request-failed", + "Confirmed Prime Agent ownership could not be durably finalized.", + ), + }); } // The initial proof cannot survive any overlapping attachment generation, even diff --git a/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.test.ts b/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.test.ts index 3cebfc290..3393101cc 100644 --- a/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.test.ts +++ b/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.test.ts @@ -3,10 +3,17 @@ import * as Effect from "effect/Effect"; import * as Exit from "effect/Exit"; import * as Layer from "effect/Layer"; import * as Option from "effect/Option"; +import * as SqlClient from "effect/unstable/sql/SqlClient"; import migration050 from "../../persistence/Migrations/050_PrimeAgentRecoveryLedger.ts"; import * as NodeSqliteClient from "../../persistence/NodeSqliteClient.ts"; -import { make, type PrimeAgentRecoveryAuthority } from "./PrimeAgentRecoveryLedger.ts"; +import { + make, + PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS, + type PrimeAgentRecoveryAdoptionProof, + type PrimeAgentRecoveryAuthority, + type PrimeAgentRecoveryLedgerShape, +} from "./PrimeAgentRecoveryLedger.ts"; const layer = it.layer(Layer.mergeAll(NodeSqliteClient.layerMemory())); @@ -45,81 +52,362 @@ const authority = { transcriptFingerprints: ["fingerprint-1", "fingerprint-2"], ownerToken: "owner-1", state: "prepared", + adoptionPreviousOwnerToken: null, + adoptionOwnerToken: null, + adoptionRequestId: null, + adoptionMcpOwnerId: null, + adoptionPhase: null, + adoptionAttempt: 0, + adoptionRecoveryHandle: null, + adoptionProof: null, nativeCleanupProven: false, terminalProjected: false, checkpointQuiesced: false, updatedAt: "2026-01-01T00:00:00.000Z", } satisfies PrimeAgentRecoveryAuthority; +const proof = { + feature: "recoverable_owned_session_adoption_v1", + status: "adopted", + supervisorGeneration: authority.supervisorGeneration, + ownershipGeneration: 1, + activeSessionId: authority.activeSessionId, + sessionId: authority.nativeSessionId, + correlationId: authority.correlationId, + lifecycle: { phase: "owned" }, + cursor: { generation: authority.cursor.generation, sequence: 11 }, + mcpOwnerId: "pylon:mcp-2", +} satisfies PrimeAgentRecoveryAdoptionProof; + +const resetLedger = Effect.gen(function* () { + yield* migration050; + const sql = yield* SqlClient.SqlClient; + yield* sql.unsafe("DELETE FROM prime_agent_recovery_ledger"); +}); + +const admit = (ledger: PrimeAgentRecoveryLedgerShape) => + ledger.markAdmitted({ + threadId: authority.threadId, + ownerToken: authority.ownerToken, + turnId: "turn-1", + updatedAt: "2026-01-01T00:00:01.000Z", + }); + layer("PrimeAgentRecoveryLedger", (it) => { - it.effect("CAS-claims one owner and deletes only after all three cleanup proofs", () => + it.effect( + "keeps prior authority while one stable adoption route advances through every phase", + () => + Effect.gen(function* () { + yield* resetLedger; + const ledger = yield* make; + yield* ledger.putPrepared(authority); + assert.isTrue(yield* admit(ledger)); + + const claimed = Option.getOrThrow( + yield* ledger.claim({ + threadId: authority.threadId, + expectedOwnerToken: authority.ownerToken, + nextOwnerToken: "owner-2", + requestId: "a".repeat(48), + mcpOwnerId: proof.mcpOwnerId, + updatedAt: "2026-01-01T00:00:02.000Z", + }), + ); + assert.equal(claimed.state, "adopting"); + assert.equal(claimed.ownerToken, authority.ownerToken); + assert.equal(claimed.recoveryHandle, authority.recoveryHandle); + assert.equal(claimed.adoptionPreviousOwnerToken, authority.ownerToken); + assert.equal(claimed.adoptionOwnerToken, "owner-2"); + assert.equal(claimed.adoptionRequestId, "a".repeat(48)); + assert.equal(claimed.adoptionPhase, "claimed"); + assert.equal(claimed.adoptionAttempt, 0); + assert.isTrue( + Option.isNone( + yield* ledger.claim({ + threadId: authority.threadId, + expectedOwnerToken: authority.ownerToken, + nextOwnerToken: "owner-3", + requestId: "b".repeat(48), + mcpOwnerId: "pylon:mcp-3", + updatedAt: "2026-01-01T00:00:02.000Z", + }), + ), + ); + + const requested = Option.getOrThrow( + yield* ledger.beginAdoptionAttempt({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:03.000Z", + }), + ); + assert.equal(requested.adoptionPhase, "requested"); + assert.equal(requested.adoptionAttempt, 1); + assert.isTrue( + yield* ledger.commitAdoption({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + recoveryHandle: "private-handle-2", + proof, + updatedAt: "2026-01-01T00:00:04.000Z", + }), + ); + const committed = Option.getOrThrow(yield* ledger.get(authority.threadId)); + assert.equal(committed.state, "adopting"); + assert.equal(committed.adoptionPhase, "committed"); + assert.equal(committed.recoveryHandle, authority.recoveryHandle); + assert.equal(committed.ownerToken, authority.ownerToken); + assert.equal(committed.adoptionRecoveryHandle, "private-handle-2"); + assert.deepEqual(committed.adoptionProof, proof); + + const retried = Option.getOrThrow( + yield* ledger.beginAdoptionAttempt({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:05.000Z", + }), + ); + assert.equal(retried.adoptionPhase, "committed"); + assert.equal(retried.adoptionAttempt, 2); + assert.isTrue( + yield* ledger.commitAdoption({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + recoveryHandle: "private-handle-2", + proof, + updatedAt: "2026-01-01T00:00:06.000Z", + }), + ); + assert.isFalse( + yield* ledger.commitAdoption({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + recoveryHandle: "wrong-handle", + proof, + updatedAt: "2026-01-01T00:00:06.000Z", + }), + ); + + const confirming = Option.getOrThrow( + yield* ledger.beginAdoptionConfirmation({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:07.000Z", + }), + ); + assert.equal(confirming.adoptionPhase, "confirming"); + assert.equal(confirming.adoptionAttempt, 3); + const confirmationRetried = Option.getOrThrow( + yield* ledger.beginAdoptionConfirmation({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:07.500Z", + }), + ); + assert.equal(confirmationRetried.adoptionPhase, "confirming"); + assert.equal(confirmationRetried.adoptionAttempt, 4); + assert.isTrue( + yield* ledger.finalizeAdoption({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + recoveryHandle: "private-handle-2", + proof, + updatedAt: "2026-01-01T00:00:08.000Z", + }), + ); + const adopted = Option.getOrThrow(yield* ledger.get(authority.threadId)); + assert.equal(adopted.state, "active"); + assert.equal(adopted.ownerToken, "owner-2"); + assert.equal(adopted.recoveryHandle, "private-handle-2"); + assert.equal(adopted.ownershipGeneration, proof.ownershipGeneration); + assert.deepEqual(adopted.cursor, proof.cursor); + assert.equal(adopted.mcpOwnerId, proof.mcpOwnerId); + assert.equal(adopted.adoptionRequestId, null); + assert.equal(adopted.adoptionProof, null); + }), + ); + + it.effect("releases only a claim whose native attempt provably never started", () => Effect.gen(function* () { - yield* migration050; + yield* resetLedger; const ledger = yield* make; yield* ledger.putPrepared(authority); - const replacement = yield* Effect.exit( - ledger.putPrepared({ ...authority, ownerToken: "owner-replacement" }), - ); - assert.isTrue(Exit.isFailure(replacement)); - assert.equal( - Option.getOrThrow(yield* ledger.get(authority.threadId)).ownerToken, - authority.ownerToken, + assert.isTrue(yield* admit(ledger)); + yield* ledger.claim({ + threadId: authority.threadId, + expectedOwnerToken: authority.ownerToken, + nextOwnerToken: "owner-2", + requestId: "a".repeat(48), + mcpOwnerId: proof.mcpOwnerId, + updatedAt: "2026-01-01T00:00:02.000Z", + }); + assert.isFalse( + yield* ledger.releaseClaim({ + threadId: authority.threadId, + ownerToken: "wrong-owner", + previousOwnerToken: authority.ownerToken, + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:03.000Z", + }), ); assert.isTrue( - yield* ledger.markAdmitted({ + yield* ledger.releaseClaim({ threadId: authority.threadId, - ownerToken: authority.ownerToken, - turnId: "turn-1", - updatedAt: "2026-01-01T00:00:01.000Z", + ownerToken: "owner-2", + previousOwnerToken: authority.ownerToken, + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:03.000Z", }), ); + const released = Option.getOrThrow(yield* ledger.get(authority.threadId)); + assert.equal(released.state, "active"); + assert.equal(released.ownerToken, authority.ownerToken); - const firstClaim = yield* ledger.claim({ + yield* ledger.claim({ threadId: authority.threadId, expectedOwnerToken: authority.ownerToken, nextOwnerToken: "owner-2", - updatedAt: "2026-01-01T00:00:02.000Z", + requestId: "a".repeat(48), + mcpOwnerId: proof.mcpOwnerId, + updatedAt: "2026-01-01T00:00:04.000Z", }); - const competingClaim = yield* ledger.claim({ + yield* ledger.beginAdoptionAttempt({ threadId: authority.threadId, - expectedOwnerToken: authority.ownerToken, - nextOwnerToken: "owner-3", - updatedAt: "2026-01-01T00:00:02.000Z", + ownerToken: "owner-2", + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:05.000Z", }); - assert.isTrue(Option.isSome(firstClaim)); - assert.isTrue(Option.isNone(competingClaim)); - - assert.isTrue( - yield* ledger.commitAdoption({ + assert.isFalse( + yield* ledger.releaseClaim({ threadId: authority.threadId, ownerToken: "owner-2", - recoveryHandle: "private-handle-2", - ownershipGeneration: 1, - cursor: { generation: "events-1", sequence: 11 }, - mcpOwnerId: "pylon:mcp-2", - updatedAt: "2026-01-01T00:00:03.000Z", + previousOwnerToken: authority.ownerToken, + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:06.000Z", }), ); + assert.equal( + Option.getOrThrow(yield* ledger.get(authority.threadId)).adoptionPhase, + "requested", + ); + }), + ); + + it.effect("bounds ambiguous retries, quarantines authority, and rejects wrong routes", () => + Effect.gen(function* () { + yield* resetLedger; + const ledger = yield* make; + yield* ledger.putPrepared(authority); + assert.isTrue(yield* admit(ledger)); + yield* ledger.claim({ + threadId: authority.threadId, + expectedOwnerToken: authority.ownerToken, + nextOwnerToken: "owner-2", + requestId: "a".repeat(48), + mcpOwnerId: proof.mcpOwnerId, + updatedAt: "2026-01-01T00:00:02.000Z", + }); assert.isTrue( - yield* ledger.updateTranscriptProgress({ + Option.isNone( + yield* ledger.beginAdoptionAttempt({ + threadId: "thread-other", + ownerToken: "owner-2", + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:03.000Z", + }), + ), + ); + assert.isTrue( + Option.isNone( + yield* ledger.beginAdoptionAttempt({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "wrong-request", + updatedAt: "2026-01-01T00:00:03.000Z", + }), + ), + ); + for (let attempt = 0; attempt < PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS; attempt += 1) { + assert.isTrue( + Option.isSome( + yield* ledger.beginAdoptionAttempt({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + updatedAt: `2026-01-01T00:00:${String(attempt + 10).padStart(2, "0")}.000Z`, + }), + ), + ); + } + assert.isTrue( + Option.isNone( + yield* ledger.beginAdoptionAttempt({ + threadId: authority.threadId, + ownerToken: "owner-2", + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:30.000Z", + }), + ), + ); + assert.isTrue( + yield* ledger.quarantineAdoption({ threadId: authority.threadId, ownerToken: "owner-2", - cursor: { generation: "events-1", sequence: 19 }, - messageCount: 3, - fingerprints: ["fingerprint-1", "fingerprint-2", "fingerprint-3"], - updatedAt: "2026-01-01T00:00:04.000Z", + requestId: "a".repeat(48), + updatedAt: "2026-01-01T00:00:31.000Z", }), ); - const adopted = Option.getOrThrow(yield* ledger.get(authority.threadId)); - assert.equal(adopted.recoveryHandle, "private-handle-2"); - assert.deepEqual(adopted.cursor, { generation: "events-1", sequence: 19 }); - assert.equal(adopted.transcriptMessageCount, 3); + assert.equal(Option.getOrThrow(yield* ledger.get(authority.threadId)).state, "quarantined"); + assert.deepEqual(yield* ledger.listActive(), []); + }), + ); + it.effect("fails closed when a staged exact proof is corrupt", () => + Effect.gen(function* () { + yield* resetLedger; + const ledger = yield* make; + const sql = yield* SqlClient.SqlClient; + yield* ledger.putPrepared(authority); + assert.isTrue(yield* admit(ledger)); + yield* ledger.claim({ + threadId: authority.threadId, + expectedOwnerToken: authority.ownerToken, + nextOwnerToken: "owner-2", + requestId: "a".repeat(48), + mcpOwnerId: proof.mcpOwnerId, + updatedAt: "2026-01-01T00:00:02.000Z", + }); + yield* sql.unsafe( + `UPDATE prime_agent_recovery_ledger + SET adoption_phase='committed', adoption_recovery_handle='private-handle-2', + adoption_proof_json='{not-json' + WHERE thread_id=?`, + [authority.threadId], + ); + const loaded = yield* Effect.exit(ledger.get(authority.threadId)); + assert.isTrue(Exit.isFailure(loaded)); + }), + ); + + it.effect("deletes only after all three cleanup proofs", () => + Effect.gen(function* () { + yield* resetLedger; + const ledger = yield* make; + yield* ledger.putPrepared(authority); + assert.isTrue(yield* admit(ledger)); assert.isFalse(yield* ledger.deleteIfSettled(authority.threadId)); assert.isTrue( yield* ledger.markNativeCleanup({ threadId: authority.threadId, - ownerToken: "owner-2", + ownerToken: authority.ownerToken, updatedAt: "2026-01-01T00:00:05.000Z", }), ); diff --git a/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.ts b/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.ts index 7c133f4a5..bc27622c4 100644 --- a/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.ts +++ b/apps/server/src/provider/prime/PrimeAgentRecoveryLedger.ts @@ -14,6 +14,22 @@ const RecoveryCursor = Schema.Struct({ sequence: NonNegativeInt, }); +export const PrimeAgentRecoveryAdoptionProof = Schema.Struct({ + feature: Schema.Literal("recoverable_owned_session_adoption_v1"), + status: Schema.Literal("adopted"), + supervisorGeneration: Schema.String, + ownershipGeneration: NonNegativeInt, + activeSessionId: Schema.String, + sessionId: Schema.String, + correlationId: Schema.String, + lifecycle: Schema.Unknown, + cursor: RecoveryCursor, + mcpOwnerId: Schema.String, +}); +export type PrimeAgentRecoveryAdoptionProof = typeof PrimeAgentRecoveryAdoptionProof.Type; + +export const PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS = 8; + export const PrimeAgentRecoveryAuthority = Schema.Struct({ threadId: Schema.String, providerInstanceId: Schema.String, @@ -41,7 +57,17 @@ export const PrimeAgentRecoveryAuthority = Schema.Struct({ transcriptMessageCount: NonNegativeInt, transcriptFingerprints: Schema.Array(Schema.String), ownerToken: Schema.String, - state: Schema.Literals(["prepared", "active", "adopting", "terminal"]), + state: Schema.Literals(["prepared", "active", "adopting", "quarantined", "terminal"]), + adoptionPreviousOwnerToken: Schema.NullOr(Schema.String), + adoptionOwnerToken: Schema.NullOr(Schema.String), + adoptionRequestId: Schema.NullOr(Schema.String), + adoptionMcpOwnerId: Schema.NullOr(Schema.String), + adoptionPhase: Schema.NullOr( + Schema.Literals(["claimed", "requested", "committed", "confirming", "quarantined"]), + ), + adoptionAttempt: NonNegativeInt, + adoptionRecoveryHandle: Schema.NullOr(Schema.String), + adoptionProof: Schema.NullOr(PrimeAgentRecoveryAdoptionProof), nativeCleanupProven: Schema.Boolean, terminalProjected: Schema.Boolean, checkpointQuiesced: Schema.Boolean, @@ -61,6 +87,8 @@ export class PrimeAgentRecoveryLedgerError extends Schema.TaggedErrorClass; readonly updatedAt: string; }) => Effect.Effect; - /** Compare-and-swap the last durable owner. Exactly one restarted Pylon process can win. */ + /** + * Compare-and-swap one stable adoption route without replacing the current native authority. + * The old owner and bearer handle remain current until native confirmation succeeds. + */ readonly claim: (input: { readonly threadId: string; readonly expectedOwnerToken: string; readonly nextOwnerToken: string; + readonly requestId: string; + readonly mcpOwnerId: string; readonly updatedAt: string; }) => Effect.Effect, PrimeAgentRecoveryLedgerError>; + /** Records an attempt before the native prepare request can begin. */ + readonly beginAdoptionAttempt: (input: { + readonly threadId: string; + readonly ownerToken: string; + readonly requestId: string; + readonly updatedAt: string; + }) => Effect.Effect, PrimeAgentRecoveryLedgerError>; + /** Releases only a never-started claim. Any native ambiguity keeps the durable route. */ readonly releaseClaim: (input: { readonly threadId: string; readonly ownerToken: string; readonly previousOwnerToken: string; + readonly requestId: string; readonly updatedAt: string; }) => Effect.Effect; - /** Persist the rotated bearer authority before the SDK confirmation step. */ + /** Persist the complete rotated receipt separately before the SDK confirmation step. */ readonly commitAdoption: (input: { readonly threadId: string; readonly ownerToken: string; + readonly requestId: string; readonly recoveryHandle: string; - readonly ownershipGeneration: number; - readonly cursor: typeof RecoveryCursor.Type; - readonly mcpOwnerId: string; + readonly proof: PrimeAgentRecoveryAdoptionProof; + readonly updatedAt: string; + }) => Effect.Effect; + /** Records the confirmation attempt before it can close the old-handle retry window. */ + readonly beginAdoptionConfirmation: (input: { + readonly threadId: string; + readonly ownerToken: string; + readonly requestId: string; + readonly updatedAt: string; + }) => Effect.Effect, PrimeAgentRecoveryLedgerError>; + /** Promotes the staged receipt only after exact native confirmation. */ + readonly finalizeAdoption: (input: { + readonly threadId: string; + readonly ownerToken: string; + readonly requestId: string; + readonly recoveryHandle: string; + readonly proof: PrimeAgentRecoveryAdoptionProof; + readonly updatedAt: string; + }) => Effect.Effect; + readonly quarantineAdoption: (input: { + readonly threadId: string; + readonly ownerToken: string; + readonly requestId: string; readonly updatedAt: string; }) => Effect.Effect; readonly markNativeCleanup: (input: { @@ -166,6 +229,14 @@ const RawRow = Schema.Struct({ transcriptFingerprintsJson: Schema.String, ownerToken: Schema.String, state: Schema.String, + adoptionPreviousOwnerToken: Schema.NullOr(Schema.String), + adoptionOwnerToken: Schema.NullOr(Schema.String), + adoptionRequestId: Schema.NullOr(Schema.String), + adoptionMcpOwnerId: Schema.NullOr(Schema.String), + adoptionPhase: Schema.NullOr(Schema.String), + adoptionAttempt: Schema.Int, + adoptionRecoveryHandle: Schema.NullOr(Schema.String), + adoptionProofJson: Schema.NullOr(Schema.String), nativeCleanupProven: Schema.Int, terminalProjected: Schema.Int, checkpointQuiesced: Schema.Int, @@ -203,12 +274,35 @@ const selectColumns = ` transcript_fingerprints_json AS transcriptFingerprintsJson, owner_token AS ownerToken, state, + adoption_previous_owner_token AS adoptionPreviousOwnerToken, + adoption_owner_token AS adoptionOwnerToken, + adoption_request_id AS adoptionRequestId, + adoption_mcp_owner_id AS adoptionMcpOwnerId, + adoption_phase AS adoptionPhase, + adoption_attempt AS adoptionAttempt, + adoption_recovery_handle AS adoptionRecoveryHandle, + adoption_proof_json AS adoptionProofJson, native_cleanup_proven AS nativeCleanupProven, terminal_projected AS terminalProjected, checkpoint_quiesced AS checkpointQuiesced, updated_at AS updatedAt `; +function encodeAdoptionProof(proof: PrimeAgentRecoveryAdoptionProof): string { + return JSON.stringify({ + feature: proof.feature, + status: proof.status, + supervisorGeneration: proof.supervisorGeneration, + ownershipGeneration: proof.ownershipGeneration, + activeSessionId: proof.activeSessionId, + sessionId: proof.sessionId, + correlationId: proof.correlationId, + lifecycle: proof.lifecycle, + cursor: proof.cursor, + mcpOwnerId: proof.mcpOwnerId, + }); +} + function ledgerError(operation: string, cause?: unknown): PrimeAgentRecoveryLedgerError { return new PrimeAgentRecoveryLedgerError({ operation, @@ -247,6 +341,14 @@ function decodeRows(rows: unknown, operation: string): ReadonlyArray decodeRows(rows, operation), catch: (cause) => - Schema.is(PrimeAgentRecoveryLedgerError)(cause) ? cause : ledgerError(operation, cause), + isPrimeAgentRecoveryLedgerError(cause) ? cause : ledgerError(operation, cause), }); const queryByThread = (threadId: string) => @@ -291,8 +393,11 @@ export const make = Effect.gen(function* () { recovery_handle, supervisor_generation, ownership_generation, cursor_generation, cursor_sequence, correlation_id, mcp_owner_id, recovery_config_json, launch_environment_json, transcript_message_count, transcript_fingerprints_json, owner_token, state, - native_cleanup_proven, terminal_projected, checkpoint_quiesced, updated_at - ) VALUES (${Array.from({ length: 32 }, () => "?").join(",")}) + adoption_previous_owner_token, adoption_owner_token, adoption_request_id, + adoption_mcp_owner_id, adoption_phase, adoption_attempt, adoption_recovery_handle, + adoption_proof_json, native_cleanup_proven, terminal_projected, checkpoint_quiesced, + updated_at + ) VALUES (${Array.from({ length: 40 }, () => "?").join(",")}) `, [ authority.threadId, @@ -323,6 +428,14 @@ export const make = Effect.gen(function* () { JSON.stringify(authority.transcriptFingerprints), authority.ownerToken, authority.state, + authority.adoptionPreviousOwnerToken, + authority.adoptionOwnerToken, + authority.adoptionRequestId, + authority.adoptionMcpOwnerId, + authority.adoptionPhase, + authority.adoptionAttempt, + authority.adoptionRecoveryHandle, + authority.adoptionProof === null ? null : JSON.stringify(authority.adoptionProof), authority.nativeCleanupProven ? 1 : 0, authority.terminalProjected ? 1 : 0, authority.checkpointQuiesced ? 1 : 0, @@ -390,37 +503,161 @@ export const make = Effect.gen(function* () { const claim: PrimeAgentRecoveryLedgerShape["claim"] = (input) => conditionalUpdate( "claim", - `UPDATE prime_agent_recovery_ledger SET owner_token=?, state='adopting', updated_at=? - WHERE thread_id=? AND owner_token=? AND state='active' RETURNING thread_id`, - [input.nextOwnerToken, input.updatedAt, input.threadId, input.expectedOwnerToken], + `UPDATE prime_agent_recovery_ledger + SET state='adopting', adoption_previous_owner_token=owner_token, + adoption_owner_token=?, adoption_request_id=?, adoption_mcp_owner_id=?, + adoption_phase='claimed', adoption_attempt=0, adoption_recovery_handle=NULL, + adoption_proof_json=NULL, updated_at=? + WHERE thread_id=? AND owner_token=? AND state='active' + RETURNING thread_id`, + [ + input.nextOwnerToken, + input.requestId, + input.mcpOwnerId, + input.updatedAt, + input.threadId, + input.expectedOwnerToken, + ], ).pipe( Effect.flatMap((claimed) => (claimed ? get(input.threadId) : Effect.succeed(Option.none()))), ); + const beginAdoptionAttempt: PrimeAgentRecoveryLedgerShape["beginAdoptionAttempt"] = (input) => + conditionalUpdate( + "beginAdoptionAttempt", + `UPDATE prime_agent_recovery_ledger + SET adoption_phase=CASE WHEN adoption_phase='claimed' THEN 'requested' ELSE adoption_phase END, + adoption_attempt=adoption_attempt+1, updated_at=? + WHERE thread_id=? AND state='adopting' AND adoption_owner_token=? + AND adoption_request_id=? AND adoption_phase IN ('claimed','requested','committed') + AND adoption_attempt < ? + RETURNING thread_id`, + [ + input.updatedAt, + input.threadId, + input.ownerToken, + input.requestId, + PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS, + ], + ).pipe( + Effect.flatMap((started) => (started ? get(input.threadId) : Effect.succeed(Option.none()))), + ); + const releaseClaim: PrimeAgentRecoveryLedgerShape["releaseClaim"] = (input) => conditionalUpdate( "releaseClaim", - `UPDATE prime_agent_recovery_ledger SET owner_token=?, state='active', updated_at=? - WHERE thread_id=? AND owner_token=? AND state='adopting' RETURNING thread_id`, - [input.previousOwnerToken, input.updatedAt, input.threadId, input.ownerToken], + `UPDATE prime_agent_recovery_ledger + SET state='active', adoption_previous_owner_token=NULL, adoption_owner_token=NULL, + adoption_request_id=NULL, adoption_mcp_owner_id=NULL, adoption_phase=NULL, + adoption_attempt=0, adoption_recovery_handle=NULL, adoption_proof_json=NULL, + updated_at=? + WHERE thread_id=? AND owner_token=? AND state='adopting' + AND adoption_previous_owner_token=? AND adoption_owner_token=? + AND adoption_request_id=? AND adoption_phase='claimed' AND adoption_attempt=0 + RETURNING thread_id`, + [ + input.updatedAt, + input.threadId, + input.previousOwnerToken, + input.previousOwnerToken, + input.ownerToken, + input.requestId, + ], ); - const commitAdoption: PrimeAgentRecoveryLedgerShape["commitAdoption"] = (input) => - conditionalUpdate( + const commitAdoption: PrimeAgentRecoveryLedgerShape["commitAdoption"] = (input) => { + const proofJson = encodeAdoptionProof(input.proof); + return conditionalUpdate( "commitAdoption", + `UPDATE prime_agent_recovery_ledger + SET adoption_recovery_handle=?, adoption_proof_json=?, adoption_phase='committed', + updated_at=? + WHERE thread_id=? AND state='adopting' AND adoption_owner_token=? + AND adoption_request_id=? + AND (adoption_phase='requested' OR + (adoption_phase='committed' AND adoption_recovery_handle=? AND adoption_proof_json=?)) + RETURNING thread_id`, + [ + input.recoveryHandle, + proofJson, + input.updatedAt, + input.threadId, + input.ownerToken, + input.requestId, + input.recoveryHandle, + proofJson, + ], + ); + }; + + const beginAdoptionConfirmation: PrimeAgentRecoveryLedgerShape["beginAdoptionConfirmation"] = ( + input, + ) => + conditionalUpdate( + "beginAdoptionConfirmation", + `UPDATE prime_agent_recovery_ledger + SET adoption_phase='confirming', adoption_attempt=adoption_attempt+1, updated_at=? + WHERE thread_id=? AND state='adopting' AND adoption_owner_token=? + AND adoption_request_id=? AND adoption_phase IN ('committed','confirming') + AND adoption_recovery_handle IS NOT NULL AND adoption_proof_json IS NOT NULL + AND adoption_attempt < ? + RETURNING thread_id`, + [ + input.updatedAt, + input.threadId, + input.ownerToken, + input.requestId, + PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS, + ], + ).pipe( + Effect.flatMap((started) => (started ? get(input.threadId) : Effect.succeed(Option.none()))), + ); + + const finalizeAdoption: PrimeAgentRecoveryLedgerShape["finalizeAdoption"] = (input) => { + const proofJson = encodeAdoptionProof(input.proof); + return conditionalUpdate( + "finalizeAdoption", `UPDATE prime_agent_recovery_ledger SET recovery_handle=?, ownership_generation=?, cursor_generation=?, cursor_sequence=?, - mcp_owner_id=?, state='active', updated_at=? - WHERE thread_id=? AND owner_token=? AND state='adopting' RETURNING thread_id`, + mcp_owner_id=?, owner_token=?, state='active', adoption_previous_owner_token=NULL, + adoption_owner_token=NULL, adoption_request_id=NULL, adoption_mcp_owner_id=NULL, + adoption_phase=NULL, adoption_attempt=0, adoption_recovery_handle=NULL, + adoption_proof_json=NULL, updated_at=? + WHERE thread_id=? AND state='adopting' AND adoption_owner_token=? + AND adoption_request_id=? AND adoption_phase='confirming' + AND adoption_recovery_handle=? AND adoption_proof_json=? + RETURNING thread_id`, [ input.recoveryHandle, - input.ownershipGeneration, - input.cursor.generation, - input.cursor.sequence, - input.mcpOwnerId, + input.proof.ownershipGeneration, + input.proof.cursor.generation, + input.proof.cursor.sequence, + input.proof.mcpOwnerId, + input.ownerToken, + input.updatedAt, + input.threadId, + input.ownerToken, + input.requestId, + input.recoveryHandle, + proofJson, + ], + ); + }; + + const quarantineAdoption: PrimeAgentRecoveryLedgerShape["quarantineAdoption"] = (input) => + conditionalUpdate( + "quarantineAdoption", + `UPDATE prime_agent_recovery_ledger + SET state='quarantined', adoption_phase='quarantined', updated_at=? + WHERE thread_id=? AND state='adopting' AND adoption_owner_token=? + AND adoption_request_id=? AND adoption_attempt >= ? + RETURNING thread_id`, + [ input.updatedAt, input.threadId, input.ownerToken, + input.requestId, + PRIME_AGENT_RECOVERY_ADOPTION_MAX_ATTEMPTS, ], ); @@ -466,8 +703,12 @@ export const make = Effect.gen(function* () { discardPrepared, updateTranscriptProgress, claim, + beginAdoptionAttempt, releaseClaim, commitAdoption, + beginAdoptionConfirmation, + finalizeAdoption, + quarantineAdoption, markNativeCleanup, markTerminalProjected, markCheckpointQuiesced, diff --git a/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs b/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs index 3a6e44340..34bdf1fce 100644 --- a/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs +++ b/apps/server/src/provider/prime/PrimeAgentRestartAdoption.real.test.mjs @@ -29,7 +29,7 @@ const skipReason = ? "native Windows is unsupported; run the POSIX proof in WSL2 with a Linux PRIME_AGENT_REAL_PACKAGE_ROOT" : "set PRIME_AGENT_REAL_PACKAGE_ROOT to the built exact Prime checkout at 507a52239d3ace7bb2b2965ade7779988fdb6344"; const enabled = NodeProcess.platform !== "win32" && Boolean(packageRoot); -const outerSafetyMs = 120_000; +const outerSafetyMs = 180_000; const maximumOutputBytes = 2 * 1024 * 1024; const providerInstanceId = "primeAgent"; const modelSelection = { @@ -469,7 +469,16 @@ const preparePylonState = async ( return stateDir; }; -const spawnPylonServer = async ({ repoRoot, baseDir, projectDir, home, port, label }) => { +const spawnPylonServer = async ({ + repoRoot, + baseDir, + projectDir, + home, + port, + label, + environment = {}, + waitForPairing = true, +}) => { const output = []; let outputBytes = 0; let pairingBuffer = ""; @@ -489,7 +498,7 @@ const spawnPylonServer = async ({ repoRoot, baseDir, projectDir, home, port, lab ], { cwd: repoRoot, - env: sanitizeServerEnvironment(home), + env: { ...sanitizeServerEnvironment(home), ...environment }, stdio: ["ignore", "pipe", "pipe"], }, ); @@ -522,7 +531,10 @@ const spawnPylonServer = async ({ repoRoot, baseDir, projectDir, home, port, lab ); } }); - const access = await withSafetyCeiling(pairing.promise, 30_000, `${label} pairing readiness`); + const access = waitForPairing + ? await withSafetyCeiling(pairing.promise, 30_000, `${label} pairing readiness`) + : undefined; + if (!waitForPairing) void pairing.promise.catch(() => undefined); return { child, access, @@ -630,7 +642,10 @@ const readLedger = (databasePath, threadId) => { `SELECT thread_id, provider_instance_id, session_incarnation_id, admission_request_id, turn_id, package_root, active_session_id, native_session_id, recovery_handle, supervisor_generation, ownership_generation, cursor_generation, cursor_sequence, - correlation_id, mcp_owner_id, owner_token, state + correlation_id, mcp_owner_id, owner_token, state, + adoption_previous_owner_token, adoption_owner_token, adoption_request_id, + adoption_mcp_owner_id, adoption_phase, adoption_attempt, + adoption_recovery_handle, adoption_proof_json FROM prime_agent_recovery_ledger WHERE thread_id = ?`, ) .get(threadId); @@ -660,6 +675,17 @@ const readProviderSessionRuntime = (databasePath, threadId) => { } }; +const waitForLedger = async (databasePath, threadId, predicate, timeoutMs) => { + const deadline = Date.now() + timeoutMs; + let observed; + while (Date.now() < deadline) { + observed = readLedger(databasePath, threadId); + if (observed !== undefined && predicate(observed)) return observed; + await new Promise((resolve) => setTimeout(resolve, 20)); + } + return observed; +}; + const waitForDurableActiveRuntime = async (databasePath, threadId, expected, timeoutMs) => { const deadline = Date.now() + timeoutMs; let observed; @@ -734,6 +760,17 @@ const stopCaptured = async (server, signal = "SIGTERM") => { } }; +const spawnRecoveryCrashServer = async (input, stage) => { + const server = await spawnPylonServer({ + ...input, + waitForPairing: false, + environment: { PRIME_AGENT_INTERNAL_PYLON_RECOVERY_CRASH_STAGE: stage }, + }); + const exited = await waitForExit(server.child, 30_000, `${input.label} deterministic crash`); + expect(exited).toMatchObject({ code: null, signal: "SIGKILL" }); + return server; +}; + const processExists = (pid) => { try { NodeProcess.kill(pid, 0); @@ -890,6 +927,8 @@ describe.skipIf(!enabled)( let fixture; let serverA; let serverB; + const crashedServers = []; + const crashLedgers = []; let bearerToken; let daemonSocket; let primeSdkEntry; @@ -1085,6 +1124,81 @@ describe.skipIf(!enabled)( busyClientOwnedSessionCount: 1, }); + const crashCases = [ + { + label: "server B after durable claim", + stage: "after-claim-persisted", + phase: "claimed", + attempt: 0, + }, + { + label: "server C after native response", + stage: "after-native-response-before-commit", + phase: "requested", + attempt: 1, + }, + { + label: "server D after durable rotated receipt", + stage: "after-commit-before-confirm", + phase: "committed", + attempt: 2, + }, + ]; + let adoptionRequestId; + let adoptionOwnerToken; + for (const crashCase of crashCases) { + const crashed = await spawnRecoveryCrashServer( + { + repoRoot, + baseDir, + projectDir, + home, + port: await reserveEphemeralPort(), + label: crashCase.label, + }, + crashCase.stage, + ); + crashedServers.push(crashed); + const crashLedger = readLedger(databasePath, threadId); + crashLedgers.push(crashLedger); + expect(crashLedger).toMatchObject({ + thread_id: ledgerA.thread_id, + session_incarnation_id: ledgerA.session_incarnation_id, + admission_request_id: ledgerA.admission_request_id, + turn_id: ledgerA.turn_id, + recovery_handle: ledgerA.recovery_handle, + owner_token: ledgerA.owner_token, + state: "adopting", + adoption_previous_owner_token: ledgerA.owner_token, + adoption_phase: crashCase.phase, + adoption_attempt: crashCase.attempt, + }); + expect(crashLedger.adoption_request_id).toMatch(/^[0-9a-f]{48}$/u); + if (adoptionRequestId === undefined) { + adoptionRequestId = crashLedger.adoption_request_id; + adoptionOwnerToken = crashLedger.adoption_owner_token; + } else { + expect(crashLedger.adoption_request_id).toBe(adoptionRequestId); + expect(crashLedger.adoption_owner_token).toBe(adoptionOwnerToken); + } + if (crashCase.phase === "committed") { + expect(crashLedger.adoption_recovery_handle).toEqual(expect.any(String)); + expect(crashLedger.adoption_recovery_handle).not.toBe(ledgerA.recovery_handle); + expect(JSON.parse(crashLedger.adoption_proof_json)).toMatchObject({ + feature: "recoverable_owned_session_adoption_v1", + status: "adopted", + activeSessionId: ledgerA.active_session_id, + sessionId: ledgerA.native_session_id, + correlationId: ledgerA.correlation_id, + }); + } else { + expect(crashLedger.adoption_recovery_handle).toBeNull(); + expect(crashLedger.adoption_proof_json).toBeNull(); + } + expect(fixture.records).toHaveLength(1); + expect(processExists(workerPid)).toBe(true); + } + const portB = await reserveEphemeralPort(); serverB = await spawnPylonServer({ repoRoot, @@ -1092,15 +1206,39 @@ describe.skipIf(!enabled)( projectDir, home, port: portB, - label: "server B", + label: "server E", }); const wsB = await issueWebSocketUrl(serverB.baseUrl, bearerToken); await runRpc( wsB, (client) => client[WS_METHODS.serverProbe]({}), - "server B command readiness after adoption", + "server E command readiness after repeated adoption recovery", ); expect(fixture.records).toHaveLength(1); + const ledgerAfterRepeatedCrash = await waitForLedger( + databasePath, + threadId, + (ledger) => ledger.state === "active", + 10_000, + ); + if (ledgerAfterRepeatedCrash?.state !== "active") { + throw new Error( + `recovery did not finalize (phase=${String(ledgerAfterRepeatedCrash?.adoption_phase)}, attempt=${String(ledgerAfterRepeatedCrash?.adoption_attempt)})`, + ); + } + expect({ + state: ledgerAfterRepeatedCrash?.state, + adoptionPhase: ledgerAfterRepeatedCrash?.adoption_phase, + adoptionAttempt: ledgerAfterRepeatedCrash?.adoption_attempt, + ownerRotated: ledgerAfterRepeatedCrash?.owner_token !== ledgerA.owner_token, + handleRotated: ledgerAfterRepeatedCrash?.recovery_handle !== ledgerA.recovery_handle, + }).toEqual({ + state: "active", + adoptionPhase: null, + adoptionAttempt: 0, + ownerRotated: true, + handleRotated: true, + }); if (!processExists(workerPid)) { throw new Error( `captured Prime worker ${workerPid} exited before recovered activity\n${serverB.output()}`, @@ -1125,6 +1263,11 @@ describe.skipIf(!enabled)( supervisor_generation: ledgerA.supervisor_generation, correlation_id: ledgerA.correlation_id, state: "active", + adoption_request_id: null, + adoption_phase: null, + adoption_attempt: 0, + adoption_recovery_handle: null, + adoption_proof_json: null, }); expect(ledgerB.recovery_handle).not.toBe(ledgerA.recovery_handle); expect(ledgerB.owner_token).not.toBe(ledgerA.owner_token); @@ -1216,6 +1359,13 @@ describe.skipIf(!enabled)( ledgerA.correlation_id, ledgerA.mcp_owner_id, ledgerB.mcp_owner_id, + adoptionRequestId, + adoptionOwnerToken, + ...crashLedgers.flatMap((ledger) => + [ledger.adoption_recovery_handle, ledger.adoption_mcp_owner_id].filter( + (value) => typeof value === "string", + ), + ), primeFacade.facadeRoot, sourceRoot, home, @@ -1230,7 +1380,11 @@ describe.skipIf(!enabled)( for (const privateValue of privateValues) { expect(publicSurface).not.toContain(privateValue); } - const logSafeResult = [serverA.output(), serverB.output()] + const logSafeResult = [ + serverA.output(), + ...crashedServers.map((server) => server.output()), + serverB.output(), + ] .join("\n") .replace(/^.*(?:Pairing URL|Connection string):.*$/gmu, "[startup access redacted]"); for (const privateValue of [ @@ -1244,6 +1398,13 @@ describe.skipIf(!enabled)( ledgerA.correlation_id, ledgerA.mcp_owner_id, ledgerB.mcp_owner_id, + adoptionRequestId, + adoptionOwnerToken, + ...crashLedgers.flatMap((ledger) => + [ledger.adoption_recovery_handle, ledger.adoption_mcp_owner_id].filter( + (value) => typeof value === "string", + ), + ), daemonSocket, ]) { expect(logSafeResult).not.toContain(privateValue); From 3a8d20fb0b8f85fd2efdf339f43e3dc78e3b39dc Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Tue, 1 Sep 2026 13:04:39 -0600 Subject: [PATCH 4/4] fix(server): preserve recovered Prime turn events --- .../Layers/ProviderRuntimeIngestion.ts | 2 +- .../provider/Layers/ProviderService.test.ts | 138 ++++++++++++++++++ .../src/provider/Layers/ProviderService.ts | 57 +++++++- 3 files changed, 193 insertions(+), 4 deletions(-) diff --git a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts index c6cb285c5..1b714b5a4 100644 --- a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts +++ b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts @@ -2585,7 +2585,7 @@ const make = Effect.gen(function* () { "thread-session-set", ); if ((applied.eventCount ?? 0) === 0) return; - if (event.type === "session.exited") { + if (event.type === "turn.completed" || event.type === "session.exited") { yield* settleRecoveryTerminalProjection(thread.id, event.createdAt); } } diff --git a/apps/server/src/provider/Layers/ProviderService.test.ts b/apps/server/src/provider/Layers/ProviderService.test.ts index eeeaa7730..1f082cd47 100644 --- a/apps/server/src/provider/Layers/ProviderService.test.ts +++ b/apps/server/src/provider/Layers/ProviderService.test.ts @@ -450,11 +450,26 @@ function makeFakeCodexAdapter( sessions.set(threadId, update(existing)); }; + const setPrepareTurnRecovery = ( + prepareTurnRecovery: + | NonNullable["prepareTurnRecovery"]> + | undefined, + ): void => { + const mutable = adapter as { + prepareTurnRecovery?: NonNullable< + ProviderAdapterShape["prepareTurnRecovery"] + >; + }; + if (prepareTurnRecovery === undefined) delete mutable.prepareTurnRecovery; + else mutable.prepareTurnRecovery = prepareTurnRecovery; + }; + return { adapter, emit, removeSession, updateSession, + setPrepareTurnRecovery, startSession, sendTurn, followUp, @@ -2707,6 +2722,129 @@ fanout.layer("ProviderServiceLive fanout", (it) => { }), ); + it.effect("retains a live same-incarnation replacement after an old session exit", () => + Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const threadId = asThreadId("thread-live-same-incarnation-exit"); + const session = yield* provider.startSession(threadId, { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + }); + const received = yield* Ref.make>([]); + const consumer = yield* Stream.take(provider.streamEvents, 2).pipe( + Stream.runForEach((event) => Ref.update(received, (events) => [...events, event])), + Effect.forkChild, + ); + yield* Effect.yieldNow; + + fanout.codex.emit({ + type: "session.exited", + eventId: asEventId("evt-live-same-incarnation-old-exit"), + provider: CODEX_DRIVER, + threadId, + sessionIncarnationId: session.sessionIncarnationId, + createdAt: "2026-01-01T00:00:00.000Z", + payload: { exitKind: "graceful" }, + }); + fanout.codex.emit({ + type: "content.delta", + eventId: asEventId("evt-live-same-incarnation-current-output"), + provider: CODEX_DRIVER, + threadId, + sessionIncarnationId: session.sessionIncarnationId, + createdAt: "2026-01-01T00:00:01.000Z", + delta: "current", + }); + yield* Fiber.join(consumer); + + assert.deepEqual( + (yield* Ref.get(received)).map((event) => event.eventId), + [ + asEventId("evt-live-same-incarnation-old-exit"), + asEventId("evt-live-same-incarnation-current-output"), + ], + ); + }), + ); + + it.effect("keeps a logical incarnation current across a recovery attachment swap", () => + Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + const threadId = asThreadId("thread-recovery-attachment-swap"); + const session = yield* provider.startSession(threadId, { + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: "full-access", + }); + const admissionRequestId = CommandId.make("cmd-recovery-attachment-swap"); + fanout.codex.setPrepareTurnRecovery((input) => + Effect.gen(function* () { + fanout.codex.removeSession(threadId); + fanout.codex.emit({ + type: "session.exited", + eventId: asEventId("evt-recovery-attachment-old-exit"), + provider: CODEX_DRIVER, + threadId, + sessionIncarnationId: session.sessionIncarnationId, + createdAt: "2026-01-01T00:00:00.000Z", + payload: { exitKind: "graceful" }, + }); + yield* Effect.yieldNow; + yield* fanout.codex.startSession({ + provider: CODEX_DRIVER, + providerInstanceId: codexInstanceId, + threadId, + runtimeMode: session.runtimeMode, + ...(session.cwd === undefined ? {} : { cwd: session.cwd }), + sessionIncarnationId: session.sessionIncarnationId, + }); + fanout.codex.emit({ + type: "turn.started", + eventId: asEventId("evt-recovery-attachment-current-turn"), + provider: CODEX_DRIVER, + threadId, + turnId: asTurnId("turn-recovery-attachment-swap"), + admissionRequestId: input.admissionRequestId, + sessionIncarnationId: session.sessionIncarnationId, + createdAt: "2026-01-01T00:00:01.000Z", + payload: {}, + }); + }), + ); + const received = yield* Ref.make>([]); + const consumer = yield* Stream.take(provider.streamEvents, 2).pipe( + Stream.runForEach((event) => Ref.update(received, (events) => [...events, event])), + Effect.forkChild, + ); + yield* Effect.yieldNow; + + yield* provider.sendTurn({ + threadId, + input: "continue after the exact attachment swap", + admissionRequestId, + sessionIncarnationId: session.sessionIncarnationId, + }); + yield* Fiber.join(consumer); + + assert.deepEqual( + (yield* Ref.get(received)).map((event) => event.eventId), + [ + asEventId("evt-recovery-attachment-old-exit"), + asEventId("evt-recovery-attachment-current-turn"), + ], + ); + }).pipe( + Effect.ensuring( + Effect.sync(() => { + fanout.codex.setPrepareTurnRecovery(undefined); + }), + ), + ), + ); + it.effect("retains a stopping incarnation until its delayed exit is ingested", () => Effect.gen(function* () { const provider = yield* ProviderService.ProviderService; diff --git a/apps/server/src/provider/Layers/ProviderService.ts b/apps/server/src/provider/Layers/ProviderService.ts index ada472642..c70b6b7b0 100644 --- a/apps/server/src/provider/Layers/ProviderService.ts +++ b/apps/server/src/provider/Layers/ProviderService.ts @@ -332,6 +332,9 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( readonly turnId: ProviderRuntimeEvent["turnId"]; } >(); + // Some adapters replace their native attachment inside prepareTurnRecovery while preserving + // the logical Pylon incarnation. Keep exit events from tearing down that incarnation mid-swap. + const turnRecoveryPreparations = new Set(); type StartReservationToken = string; type StartReservationEntry = { readonly currentToken: StartReservationToken; @@ -704,13 +707,15 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( } } if (canonicalEvent.type === "session.exited") { + const stillActive = + turnRecoveryPreparations.has(canonicalEvent.threadId) || + (yield* source.adapter.hasSession(canonicalEvent.threadId)); const retained = currentSessionIncarnations.get(canonicalEvent.threadId); - if (retained?.id === currentIncarnation.id) { + if (!stillActive && retained?.id === currentIncarnation.id) { currentSessionIncarnations.delete(canonicalEvent.threadId); } const mcpSession = McpProviderSession.readMcpProviderSession(canonicalEvent.threadId); if (mcpSession?.providerInstanceId !== source.instanceId) return; - const stillActive = yield* source.adapter.hasSession(canonicalEvent.threadId); if (!stillActive) yield* clearMcpSession(canonicalEvent.threadId); } }); @@ -1295,7 +1300,53 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( }); } if (routed.adapter.prepareTurnRecovery !== undefined) { - yield* routed.adapter.prepareTurnRecovery(input); + turnRecoveryPreparations.add(input.threadId); + yield* routed.adapter.prepareTurnRecovery(input).pipe( + Effect.onError(() => + routed.adapter.hasSession(input.threadId).pipe( + Effect.flatMap((stillLive) => { + if (stillLive) return Effect.void; + const retained = currentSessionIncarnations.get(input.threadId); + if ( + retained?.instanceId === routed.instanceId && + retained.adapter === routed.adapter + ) { + currentSessionIncarnations.delete(input.threadId); + } + return clearMcpSession(input.threadId); + }), + Effect.ignore, + ), + ), + Effect.ensuring( + Effect.sync(() => { + turnRecoveryPreparations.delete(input.threadId); + }), + ), + ); + } + if ( + routed.adapter.prepareTurnRecovery !== undefined && + input.sessionIncarnationId !== undefined + ) { + const preparedSessionIsLive = yield* routed.adapter + .listSessions() + .pipe( + Effect.map((sessions) => + sessions.some( + (session) => + session.threadId === input.threadId && + session.sessionIncarnationId === input.sessionIncarnationId, + ), + ), + ); + if (preparedSessionIsLive) { + currentSessionIncarnations.set(input.threadId, { + id: input.sessionIncarnationId, + instanceId: routed.instanceId, + adapter: routed.adapter, + }); + } } const turn = yield* routed.adapter.sendTurn(input); yield* directory.upsert({