diff --git a/apps/server/src/orchestration-v2/ProjectionRecovery.test.ts b/apps/server/src/orchestration-v2/ProjectionRecovery.test.ts index a4914fdee4509..984f92566b882 100644 --- a/apps/server/src/orchestration-v2/ProjectionRecovery.test.ts +++ b/apps/server/src/orchestration-v2/ProjectionRecovery.test.ts @@ -738,6 +738,28 @@ it.effect("includes shared sessions and provider-owned background rosters in rec new Set(yield* projections.getRecoveryThreadIds("runtime")), new Set([first, second, roster, prepared]), ); + assert.deepEqual( + new Set(yield* projections.getRecoveryThreadIds("restart-continuations")), + new Set([first, second, prepared]), + ); + assert.deepEqual( + yield* projections.getRecoveryThreadIds("restart-continuations", { include: [] }), + [], + ); + assert.deepEqual( + yield* projections.getRecoveryThreadIds("restart-continuations", { + exclude: ["project:recovery"], + }), + [], + ); + assert.deepEqual( + new Set( + yield* projections.getRecoveryThreadIds("restart-continuations", { + include: ["project:recovery"], + }), + ), + new Set([first, second, prepared]), + ); const preparedState = yield* projections.getRuntimeRecoveryProjection(prepared); assert.deepEqual( preparedState.providerSessions.map((session) => session.id), diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index bd6a61db9030b..14aa46e3a7864 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -168,6 +168,7 @@ export type ProjectionStoreV2Error = typeof ProjectionStoreV2Error.Type; export type ProjectionRecoveryKind = | "queued-runs" | "runtime" + | "restart-continuations" | "subagent-results" | "delegated-completions"; @@ -353,6 +354,7 @@ export interface ProjectionRecordFilter { readonly subagentIds?: ReadonlyArray; readonly subagentRunIds?: ReadonlyArray; readonly nodeIds?: ReadonlyArray; + readonly turnItemRequestIds?: ReadonlyArray; readonly turnItemTypes?: ReadonlyArray; readonly turnItemStatuses?: ReadonlyArray; readonly omitSettledLeafNodes?: boolean; @@ -678,6 +680,10 @@ export interface ProjectionStoreV2Shape { ) => Effect.Effect, ProjectionStoreV2Error>; readonly getRecoveryThreadIds: ( kind: ProjectionRecoveryKind, + projects?: { + readonly include?: ReadonlyArray; + readonly exclude?: ReadonlyArray; + }, ) => Effect.Effect, ProjectionStoreV2Error>; readonly getUnreadableThreadIds: () => Effect.Effect< ReadonlyArray, @@ -3121,6 +3127,7 @@ export const layer: Layer.Layer = (fields !== undefined && !fields.includes("turnItems")) || (window === undefined && matchesNothing( + filter?.turnItemRequestIds, filter?.turnItemTypes, filter?.turnItemStatuses, filter?.turnItemRunIds, @@ -3132,6 +3139,7 @@ export const layer: Layer.Layer = SELECT payload_json FROM orchestration_v2_projection_turn_items WHERE thread_id = ${threadId} + ${filter?.turnItemRequestIds === undefined ? sql`` : sql`AND json_extract(payload_json, '$.requestId') IN (SELECT value FROM json_each(${encodeIdList(filter.turnItemRequestIds)}))`} ${filter?.turnItemTypes === undefined ? sql`` : sql`AND type IN (SELECT value FROM json_each(${encodeIdList(filter.turnItemTypes)}))`} ${filter?.turnItemStatuses === undefined ? sql`` : sql`AND status IN (SELECT value FROM json_each(${encodeIdList(filter.turnItemStatuses)}))`} ${filter?.turnItemRunId === undefined ? sql`` : sql`AND run_id = ${filter.turnItemRunId}`} @@ -4108,7 +4116,13 @@ export const layer: Layer.Layer = ); const getRecoveryThreadIds = Effect.fn("ProjectionStore.getRecoveryThreadIds")( - function* (kind: ProjectionRecoveryKind) { + function* ( + kind: ProjectionRecoveryKind, + projects?: { + readonly include?: ReadonlyArray; + readonly exclude?: ReadonlyArray; + }, + ) { const candidates = (() => { switch (kind) { case "queued-runs": @@ -4159,6 +4173,27 @@ export const layer: Layer.Layer = AND type = 'subagent_result' ))) `; + case "restart-continuations": + return sql` + SELECT bindings.thread_id + FROM orchestration_v2_projection_provider_sessions AS sessions + JOIN orchestration_v2_projection_provider_session_bindings AS bindings + ON bindings.provider_session_id = sessions.provider_session_id + WHERE sessions.status NOT IN ('stopped', 'error') + UNION + SELECT thread_id FROM orchestration_v2_projection_provider_sessions AS sessions + WHERE status NOT IN ('stopped', 'error') AND thread_id IS NOT NULL + AND NOT EXISTS ( + SELECT 1 FROM orchestration_v2_projection_provider_session_bindings AS bindings + WHERE bindings.provider_session_id = sessions.provider_session_id + ) + UNION + SELECT thread_id FROM orchestration_v2_projection_runs + WHERE status = 'starting' + AND CASE WHEN json_valid(payload_json) + THEN json_extract(payload_json, '$.restartContinuationOfRunId') IS NOT NULL + ELSE 0 END + `; case "runtime": return sql` WITH pending_provider_threads AS MATERIALIZED ( @@ -4242,6 +4277,8 @@ export const layer: Layer.Layer = SELECT thread_id FROM orchestration_v2_projection_threads WHERE deleted_at IS NULL AND thread_id IN (${candidates}) + ${projects?.include === undefined ? sql`` : sql`AND project_id IN (SELECT value FROM json_each(${encodeIdList(projects.include)}))`} + ${projects?.exclude === undefined ? sql`` : sql`AND project_id NOT IN (SELECT value FROM json_each(${encodeIdList(projects.exclude)}))`} ${kind === "queued-runs" ? sql`AND archived_at IS NULL` : sql``} ORDER BY updated_at ASC, thread_id ASC `; diff --git a/apps/server/src/orchestration-v2/ProviderRuntimeRecoveryService.ts b/apps/server/src/orchestration-v2/ProviderRuntimeRecoveryService.ts index 5abaab7f1228f..96e2d890b166e 100644 --- a/apps/server/src/orchestration-v2/ProviderRuntimeRecoveryService.ts +++ b/apps/server/src/orchestration-v2/ProviderRuntimeRecoveryService.ts @@ -867,39 +867,56 @@ export const make = Effect.gen(function* () { const prepareForShutdown = Effect.gen(function* () { const enabled = yield* settings.getSettings.pipe(Effect.orElseSucceed(() => null)); if (!enabled) return; - const threadIds = yield* projections.getRecoveryThreadIds("runtime"); - for (const threadId of threadIds) { - yield* Effect.gen(function* () { - const projection = yield* projections.getRuntimeRecoveryProjection(threadId); - if ( - !resolveProjectSettings(enabled, projection.thread.projectId).settings - .continueThreadsAfterServerUpdate - ) - return; - const run = restartContinuationRun(projection, yield* DateTime.now); - if (!run) return; - const commandId = CommandId.make(`command:restart-prepare:${run.id}`); - yield* eventSink.writeWithEffects({ - commandId, - events: [], - effects: [ - { - id: `effect:restart-continuation:${run.id}`, - commandId, - threadId, - request: { type: "provider-runtime.continue", sourceRunId: run.id }, - }, - ], - }); - }).pipe( - Effect.catchCauseIf( - (cause) => !Cause.hasInterruptsOnly(cause), - (cause) => - Effect.logWarning("Failed to prepare a restart continuation", { threadId, cause }), + const enabledProjects = Object.entries(enabled.projectSettingsOverrides) + .filter(([, overrides]) => overrides.continueThreadsAfterServerUpdate === true) + .map(([id]) => id); + if (!enabled.continueThreadsAfterServerUpdate && enabledProjects.length === 0) return; + const disabledProjects = Object.entries(enabled.projectSettingsOverrides) + .filter(([, overrides]) => overrides.continueThreadsAfterServerUpdate === false) + .map(([id]) => id); + const threadIds = yield* projections.getRecoveryThreadIds( + "restart-continuations", + enabled.continueThreadsAfterServerUpdate + ? { exclude: disabledProjects } + : { include: enabledProjects }, + ); + yield* Effect.forEach( + threadIds, + (threadId) => + Effect.gen(function* () { + const projection = yield* projections.getRuntimeRecoveryProjection(threadId); + if ( + !resolveProjectSettings(enabled, projection.thread.projectId).settings + .continueThreadsAfterServerUpdate + ) + return; + const run = restartContinuationRun(projection, yield* DateTime.now); + if (!run) return; + const commandId = CommandId.make(`command:restart-prepare:${run.id}`); + yield* eventSink.writeWithEffects({ + commandId, + events: [], + effects: [ + { + id: `effect:restart-continuation:${run.id}`, + commandId, + threadId, + request: { type: "provider-runtime.continue", sourceRunId: run.id }, + }, + ], + }); + }).pipe( + Effect.catchCauseIf( + (cause) => !Cause.hasInterruptsOnly(cause), + (cause) => + Effect.logWarning("Failed to prepare a restart continuation", { threadId, cause }), + ), ), - ); - } + // Leave room in the four-reader pool for other shutdown reads. + { concurrency: 2, discard: true }, + ); }).pipe( + Effect.withSpan("ProviderRuntimeRecoveryService.prepareForShutdown"), Effect.mapError((cause) => new ProviderRuntimeRecoveryError({ operation: "reconcile", cause })), ); diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index b1e05577bfbae..02b4933221ac2 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -2859,6 +2859,27 @@ it.effect("ProviderSessionManagerV2 marks pending runtime requests non-live on r now, }); yield* eventSink.write({ events: pendingRequest.events }); + const historical = yield* makePendingRuntimeRequestEvents({ + idAllocator, + threadId, + providerSessionId, + providerThread, + now, + }); + yield* eventSink.write({ + events: historical.events.map((event) => + event.type === "runtime-request.updated" + ? { ...event, payload: { ...event.payload, status: "resolved", resolvedAt: now } } + : event, + ), + }); + assert.deepEqual( + (yield* projectionStore.getThreadRecords(threadId, ["nodes", "turnItems"], { + nodeIds: [pendingRequest.nodeId], + turnItemRequestIds: [pendingRequest.requestId], + })).turnItems.map((item) => item.id), + [pendingRequest.events.find((event) => event.type === "turn-item.updated")!.payload.id], + ); yield* manager.open({ threadId, providerSessionId, @@ -2872,7 +2893,13 @@ it.effect("ProviderSessionManagerV2 marks pending runtime requests non-live on r }); const projection = yield* projectionStore.getThreadProjection(threadId); - const request = projection.runtimeRequests.at(-1); + const request = projection.runtimeRequests.find( + (request) => request.id === pendingRequest.requestId, + ); + assert.equal( + projection.runtimeRequests.find((request) => request.id === historical.requestId)?.status, + "resolved", + ); const requestNode = projection.nodes.find((node) => node.id === request?.nodeId); const requestTurnItem = projection.turnItems.find( (item) => item.type === "approval_request" && item.requestId === request?.id, diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index f06011676c667..544602dd7f885 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -652,8 +652,8 @@ export const layerWithOptions = ( for (const threadId of input.entry.attachedThreadIds) { const projection = yield* projectionStore.getThreadRecords( threadId, - ["runtimeRequests", "nodes", "turnItems"], - { turnItemTypes: ["approval_request", "user_input_request"] }, + ["runtimeRequests"], + { runtimeRequestStatuses: ["pending"] }, ); const releasedRequests = projection.runtimeRequests.filter( (request) => @@ -665,6 +665,16 @@ export const layerWithOptions = ( !(input.reason === "server_shutdown" && request.kind === "user_input"), ); + if (releasedRequests.length === 0) continue; + const records = yield* projectionStore.getThreadRecords( + threadId, + ["nodes", "turnItems"], + { + nodeIds: releasedRequests.map((request) => request.nodeId), + turnItemRequestIds: releasedRequests.map((request) => request.id), + turnItemTypes: ["approval_request", "user_input_request"], + }, + ); for (const request of releasedRequests) { events.push({ id: yield* idAllocator.allocate.event({ @@ -687,7 +697,7 @@ export const layerWithOptions = ( }, }); - const requestNode = projection.nodes.find((node) => node.id === request.nodeId); + const requestNode = records.nodes.find((node) => node.id === request.nodeId); if (requestNode !== undefined) { events.push({ id: yield* idAllocator.allocate.event({ @@ -708,7 +718,7 @@ export const layerWithOptions = ( }); } - const turnItem = projection.turnItems.find( + const turnItem = records.turnItems.find( (item) => (item.type === "approval_request" || item.type === "user_input_request") && item.requestId === request.id, @@ -2108,7 +2118,7 @@ export const layerWithOptions = ( ), { concurrency: "unbounded", discard: true }, ); - }); + }).pipe(Effect.withSpan("ProviderSessionManager.shutdown")); yield* Effect.addFinalizer(() => shutdown.pipe( Effect.ensuring(closeScopeWithin(sessionScopes, { reason: "server_shutdown" })), diff --git a/apps/server/src/orchestration-v2/RestartContinuation.test.ts b/apps/server/src/orchestration-v2/RestartContinuation.test.ts index a318130500b3b..d2f21a772ac99 100644 --- a/apps/server/src/orchestration-v2/RestartContinuation.test.ts +++ b/apps/server/src/orchestration-v2/RestartContinuation.test.ts @@ -17,6 +17,8 @@ import { } from "@t3tools/contracts"; import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; +import * as Deferred from "effect/Deferred"; +import * as Fiber from "effect/Fiber"; import * as Layer from "effect/Layer"; import * as ServerSettings from "../serverSettings.ts"; import { restartContinuationRun, continueRestartedRun } from "./RestartContinuation.ts"; @@ -1262,3 +1264,65 @@ it.effect.each(["completed", "failed", "interrupted"] as const)( ); }), ); + +it.effect("prepares independent continuations together and waits for durable writes", () => + Effect.gen(function* () { + const bothReading = yield* Deferred.make(); + const bothWriting = yield* Deferred.make(); + const committed = yield* Deferred.make(); + let reads = 0; + let writes = 0; + const recovery = yield* ProviderRuntimeRecovery.make.pipe( + Effect.provide( + Layer.mergeAll( + ServerSettings.layerTest({ continueThreadsAfterServerUpdate: true }), + Layer.mock(ProjectionStore.ProjectionStoreV2)({ + getRecoveryThreadIds: () => Effect.succeed([threadId, ThreadId.make("thread:second")]), + getRuntimeRecoveryProjection: () => + Effect.gen(function* () { + if (++reads === 2) yield* Deferred.succeed(bothReading, undefined); + yield* Deferred.await(bothReading); + return makeProjection(); + }), + }), + Layer.mock(EventSink.EventSinkV2)({ + writeWithEffects: () => + Effect.gen(function* () { + if (++writes === 2) yield* Deferred.succeed(bothWriting, undefined); + yield* Deferred.await(committed); + return []; + }), + }), + IdAllocator.layer, + Layer.mock(EffectWorker.OrchestrationEffectWorkerV2)({}), + Layer.mock(EffectOutbox.EffectOutboxV2)({}), + ), + ), + ); + const fiber = yield* recovery.prepareForShutdown.pipe(Effect.forkChild); + yield* Deferred.await(bothWriting); + assert.equal(writes, 2); + assert.isUndefined(fiber.pollUnsafe()); + yield* Deferred.succeed(committed, undefined); + yield* Fiber.join(fiber); + }), +); + +it.effect("skips restart discovery when no project enables continuation", () => + Effect.gen(function* () { + const recovery = yield* ProviderRuntimeRecovery.make.pipe( + Effect.provide( + Layer.mergeAll( + ServerSettings.layerTest({ continueThreadsAfterServerUpdate: false }), + Layer.mock(ProjectionStore.ProjectionStoreV2)({}), + Layer.mock(EventSink.EventSinkV2)({}), + IdAllocator.layer, + Layer.mock(EffectWorker.OrchestrationEffectWorkerV2)({}), + Layer.mock(EffectOutbox.EffectOutboxV2)({}), + ), + ), + ); + yield* recovery.prepareForShutdown; + yield* recovery.prepareForShutdown; + }), +); diff --git a/apps/server/src/serverRuntimeStartup.ts b/apps/server/src/serverRuntimeStartup.ts index e37f3078548ce..26f91f59affc9 100644 --- a/apps/server/src/serverRuntimeStartup.ts +++ b/apps/server/src/serverRuntimeStartup.ts @@ -478,6 +478,7 @@ const make = Effect.fn("serverRuntimeStartup.make")(function* (options?: Startup const reconciliation = yield* providerRuntimeRecovery.reconcile("shutdown"); yield* Effect.logInfo("V2 orchestration shutdown reconciliation completed", reconciliation); }).pipe( + Effect.withSpan("server.shutdown"), Effect.catchCause((cause) => Effect.logWarning("V2 orchestration shutdown reconciliation failed", { cause: Cause.pretty(cause),