diff --git a/apps/server/src/mcp/DelegatedTaskProviderExit.integration.test.ts b/apps/server/src/mcp/DelegatedTaskProviderExit.integration.test.ts new file mode 100644 index 0000000000000..e303cdfcfdbc9 --- /dev/null +++ b/apps/server/src/mcp/DelegatedTaskProviderExit.integration.test.ts @@ -0,0 +1,465 @@ +import * as NodeServices from "@effect/platform-node/NodeServices"; +import { assert, it } from "@effect/vitest"; +import { + CommandId, + EnvironmentId, + MessageId, + NodeId, + ProjectId, + ProviderDriverKind, + ProviderInstanceId, + ProviderSessionId, + ProviderThreadId, + ProviderTurnId, + ThreadId, + TurnItemId, + type OrchestrationV2DomainEvent, +} from "@t3tools/contracts"; +import * as DateTime from "effect/DateTime"; +import * as Effect from "effect/Effect"; +import * as Fiber from "effect/Fiber"; +import * as Layer from "effect/Layer"; +import * as Queue from "effect/Queue"; +import * as Stream from "effect/Stream"; + +import { ClaudeProviderCapabilitiesV2 } from "../orchestration-v2/Adapters/ClaudeAdapterV2.ts"; +import * as CommandReceiptStore from "../orchestration-v2/CommandReceiptStore.ts"; +import * as EventSink from "../orchestration-v2/EventSink.ts"; +import * as EffectWorker from "../orchestration-v2/EffectWorker.ts"; +import * as Orchestrator from "../orchestration-v2/Orchestrator.ts"; +import * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; +import * as ProviderAdapterRegistry from "../orchestration-v2/ProviderAdapterRegistry.ts"; +import * as ThreadLaunchService from "../orchestration-v2/ThreadLaunchService.ts"; +import * as GitVcsDriver from "../vcs/GitVcsDriver.ts"; +import * as ThreadManagementService from "../orchestration-v2/ThreadManagementService.ts"; +import * as ServerSettings from "../serverSettings.ts"; +import * as ProviderReplayHarness from "../orchestration-v2/testkit/ProviderReplayHarness.ts"; +import { checkpointWorkspace } from "@t3tools/provider-testing/replayWorkspace"; +import * as SqlitePersistence from "../persistence/Sqlite.ts"; +import * as ProjectService from "../project/ProjectService.ts"; +import * as ProviderRegistry from "../provider/ProviderRegistry.ts"; +import * as ScheduledTaskService from "../scheduledTasks/ScheduledTaskService.ts"; +import * as RemoteDelegation from "../peer/RemoteDelegation.ts"; +import * as SecretRequests from "../secrets/SecretRequests.ts"; +import type { McpInvocationScope } from "./McpInvocationContext.ts"; +import * as OrchestratorMcpService from "./OrchestratorMcpService.ts"; + +const layerOrchestratorService = OrchestratorMcpService.layer.pipe( + Layer.provide(Layer.mock(ThreadLaunchService.ThreadLaunchService)({})), + Layer.provide(Layer.mock(GitVcsDriver.GitVcsDriver)({})), +); + +const driver = ProviderDriverKind.make("claudeAgent"); +const instanceId = ProviderInstanceId.make("claudeAgent"); +const modelSelection = { instanceId, model: "claude-test" }; + +it.effect.each([false, true])( + "a delegated child's provider exit settles native and tool work (completed root=%s)", + (completedRoot) => + Effect.scoped( + Effect.gen(function* () { + const cwd = yield* checkpointWorkspace("provider-exit-native-subagent"); + const turns: Array<{ + readonly turn: ProviderAdapter.ProviderAdapterV2TurnInput; + readonly events: Queue.Queue< + ProviderAdapter.ProviderAdapterV2Event, + ProviderAdapter.ProviderAdapterEventStreamError + >; + }> = []; + const adapter: ProviderAdapter.ProviderAdapterV2["Service"] = { + instanceId, + driver, + getCapabilities: () => Effect.succeed(ClaudeProviderCapabilitiesV2), + planSelectionTransition: () => Effect.succeed({ type: "apply_on_next_turn" }), + openSession: (input) => + Effect.gen(function* () { + const now = yield* DateTime.now; + const events = yield* Queue.unbounded< + ProviderAdapter.ProviderAdapterV2Event, + ProviderAdapter.ProviderAdapterEventStreamError + >(); + return { + instanceId, + driver, + providerSessionId: input.providerSessionId, + providerSession: { + id: input.providerSessionId, + driver, + providerInstanceId: instanceId, + status: "ready", + cwd, + model: modelSelection.model, + capabilities: ClaudeProviderCapabilitiesV2, + createdAt: now, + updatedAt: now, + lastError: null, + }, + events: Stream.fromQueue(events), + ensureThread: ({ threadId }) => + Effect.succeed({ + id: ProviderThreadId.make(`provider-thread:${threadId}`), + driver, + providerInstanceId: instanceId, + providerSessionId: input.providerSessionId, + appThreadId: threadId, + ownerNodeId: null, + nativeThreadRef: { driver, nativeId: `native:${threadId}`, strength: "strong" }, + nativeConversationHeadRef: null, + status: "idle", + firstRunOrdinal: null, + lastRunOrdinal: null, + handoffIds: [], + forkedFrom: null, + createdAt: now, + updatedAt: now, + }), + resumeThread: ({ providerThread }) => Effect.succeed(providerThread), + startTurn: (turn) => + Effect.gen(function* () { + turns.push({ turn, events }); + yield* Queue.offer(events, { + type: "provider_turn.updated", + driver, + providerTurn: { + id: ProviderTurnId.make(`provider-turn:${turn.attemptId}`), + providerThreadId: turn.providerThread.id, + nodeId: turn.rootNodeId, + runAttemptId: turn.attemptId, + nativeTurnRef: null, + ordinal: turn.providerTurnOrdinal, + status: "running", + startedAt: now, + completedAt: null, + }, + }); + }), + steerTurn: () => Effect.die("unused"), + interruptTurn: () => Effect.void, + respondToRuntimeRequest: () => Effect.die("unused"), + readThreadSnapshot: () => Effect.die("unused"), + rollbackThread: () => Effect.die("unused"), + forkThread: () => Effect.die("unused"), + }; + }), + }; + const layerAdapters = ProviderAdapterRegistry.layerSingle(adapter); + const layerOrchestrator = ProviderReplayHarness.layerWithRegistry( + { name: "provider-exit-native-subagent" }, + layerAdapters, + { runEffectWorker: false }, + ); + const layerOrchestration = Layer.merge( + layerOrchestrator, + ThreadManagementService.layer.pipe( + Layer.provide(layerOrchestrator), + Layer.provide(ServerSettings.layerTest()), + Layer.provide( + CommandReceiptStore.layer.pipe(Layer.provide(SqlitePersistence.layerMemory)), + ), + ), + ); + const layerTest = layerOrchestratorService.pipe( + Layer.provideMerge(layerOrchestration), + Layer.provide( + Layer.mergeAll( + layerAdapters, + Layer.mock(ProviderRegistry.ProviderRegistry)({ getProviders: Effect.succeed([]) }), + Layer.mock(ProjectService.ProjectService)({}), + Layer.mock(SecretRequests.SecretRequests)({}), + Layer.mock(RemoteDelegation.RemoteDelegation)({}), + ServerSettings.layerTest(), + Layer.mock(ScheduledTaskService.ScheduledTaskService)({}), + NodeServices.layer, + ), + ), + ); + + yield* Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const worker = yield* EffectWorker.OrchestrationEffectWorkerV2; + const service = yield* OrchestratorMcpService.OrchestratorMcpService; + const eventSink = yield* EventSink.EventSinkV2; + const parentThreadId = ThreadId.make("thread:provider-exit-parent"); + const watch = (predicate: (event: OrchestrationV2DomainEvent) => boolean) => + Effect.gen(function* () { + const afterSequence = yield* eventSink.latestSequence(); + return yield* eventSink.stream({ afterSequence }).pipe( + Stream.map((stored) => stored.event), + Stream.filter(predicate), + Stream.take(1), + Stream.runDrain, + Effect.forkScoped, + ); + }); + const turnRunning = () => + watch( + (event) => + event.type === "provider-turn.updated" && event.payload.status === "running", + ); + yield* orchestrator.dispatch({ + type: "thread.create", + commandId: CommandId.make("create-parent"), + threadId: parentThreadId, + projectId: ProjectId.make("project:provider-exit"), + title: "Parent", + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + branch: null, + worktreePath: cwd, + createdBy: "user", + creationSource: "web", + }); + const parentRunning = yield* turnRunning(); + yield* orchestrator.dispatch({ + type: "message.dispatch", + commandId: CommandId.make("start-parent"), + threadId: parentThreadId, + messageId: MessageId.make("message:start-parent"), + text: "Delegate a review", + attachments: [], + dispatchMode: { type: "start_immediately" }, + createdBy: "user", + creationSource: "web", + }); + yield* worker.drain(); + yield* Fiber.join(parentRunning); + + const parentRun = (yield* orchestrator.getThreadProjection(parentThreadId)).runs[0]!; + const childRunning = yield* turnRunning(); + yield* orchestrator.dispatch({ + type: "delegated_task.request", + commandId: CommandId.make("delegate"), + parentThreadId, + parentRunId: parentRun.id, + parentNodeId: parentRun.rootNodeId!, + task: "Review the change", + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + completionWake: "always", + createdBy: "agent", + creationSource: "mcp", + }); + yield* worker.drain(); + yield* Fiber.join(childRunning); + const task = (yield* orchestrator.getThreadProjection(parentThreadId)).subagents[0]!; + const child = turns.find((entry) => entry.turn.threadId === task.childThreadId)!; + const childTurn = (yield* orchestrator.getThreadProjection(task.childThreadId!)) + .providerTurns[0]!; + + const now = yield* DateTime.now; + const subagentId = NodeId.make("node:native-subagent"); + const native = { + threadId: task.childThreadId!, + runId: child.turn.runId, + origin: "provider_native", + driver, + providerInstanceId: instanceId, + providerThreadId: childTurn.providerThreadId, + childThreadId: null, + prompt: "Check the tests", + title: null, + status: "running", + result: null, + startedAt: now, + completedAt: null, + updatedAt: now, + } as const; + const toolNodeId = NodeId.make("node:ordinary-tool"); + const toolItemId = TurnItemId.make("item:ordinary-tool"); + const workPersisted = yield* watch((event) => + completedRoot + ? event.type === "run.updated" && + event.payload.id === child.turn.runId && + event.payload.status === "waiting" + : event.type === "turn-item.updated" && event.payload.id === toolItemId, + ); + for (const event of [ + { + type: "node.updated", + driver, + node: { + id: subagentId, + threadId: native.threadId, + runId: native.runId, + parentNodeId: child.turn.rootNodeId, + rootNodeId: child.turn.rootNodeId, + kind: "subagent", + status: "running", + countsForRun: false, + providerThreadId: childTurn.providerThreadId, + providerTurnId: childTurn.id, + nativeItemRef: null, + runtimeRequestId: null, + checkpointScopeId: null, + startedAt: now, + completedAt: null, + }, + }, + { + type: "subagent.updated", + driver, + subagent: { + ...native, + id: subagentId, + parentNodeId: child.turn.rootNodeId, + createdBy: "agent", + nativeTaskRef: null, + model: null, + }, + }, + { + type: "turn_item.updated", + driver, + turnItem: { + ...native, + id: TurnItemId.make("turn-item:native-subagent"), + nodeId: subagentId, + providerTurnId: childTurn.id, + nativeItemRef: null, + parentItemId: null, + ordinal: 1, + type: "subagent", + subagentId, + }, + }, + { + type: "node.updated", + driver, + node: { + id: toolNodeId, + threadId: native.threadId, + runId: native.runId, + parentNodeId: child.turn.rootNodeId, + rootNodeId: child.turn.rootNodeId, + kind: "tool_call", + status: "running", + countsForRun: false, + providerThreadId: childTurn.providerThreadId, + providerTurnId: childTurn.id, + nativeItemRef: null, + runtimeRequestId: null, + checkpointScopeId: null, + startedAt: now, + completedAt: null, + }, + }, + { + type: "turn_item.updated", + driver, + turnItem: { + id: toolItemId, + threadId: native.threadId, + runId: native.runId, + nodeId: toolNodeId, + providerThreadId: childTurn.providerThreadId, + providerTurnId: childTurn.id, + nativeItemRef: null, + parentItemId: null, + ordinal: 2, + type: "command_execution", + status: "running", + title: "Run command", + input: "echo partial", + output: "partial output", + startedAt: now, + completedAt: null, + updatedAt: now, + }, + }, + ...(completedRoot + ? ([ + { + type: "provider_turn.updated", + driver, + providerTurn: { ...childTurn, status: "completed", completedAt: now }, + }, + { + type: "turn.terminal", + driver, + providerThreadId: childTurn.providerThreadId, + providerTurnId: childTurn.id, + runOrdinal: child.turn.runOrdinal, + status: "completed", + failure: null, + threadDisposition: "reusable", + }, + ] satisfies ReadonlyArray) + : []), + ] satisfies ReadonlyArray) { + yield* Queue.offer(child.events, event); + } + yield* Fiber.join(workPersisted); + yield* worker.drain(); + + const scope: McpInvocationScope = { + environmentId: EnvironmentId.make("environment:provider-exit"), + requestNamespace: "provider-session:provider-exit", + thread: { + threadId: parentThreadId, + providerSessionId: "provider-session:provider-exit", + providerInstanceId: instanceId, + }, + client: undefined, + capabilities: new Set(["orchestration"]), + issuedAt: 1, + }; + const before = yield* service.taskStatus(scope, task.id); + assert.equal(before.workState, completedRoot ? "waiting_for_children" : "working"); + + // The parent's own task record is what wakes it; task_status alone + // derives "completed" live from the child. + const parentTaskSettled = yield* watch( + (event) => + event.type === "subagent.updated" && + event.threadId === parentThreadId && + event.payload.id === task.id && + event.payload.result !== null, + ); + const subagentFailed = yield* watch( + (event) => + event.type === "subagent.updated" && + event.payload.id === subagentId && + event.payload.status === "failed", + ); + yield* Queue.fail( + child.events, + new ProviderAdapter.ProviderAdapterEventStreamError({ + driver, + providerSessionId: ProviderSessionId.make("provider-session:provider-exit"), + cause: "provider process exited", + }), + ); + yield* Fiber.join(subagentFailed); + yield* Fiber.join(parentTaskSettled); + yield* worker.drain(); + + const childAfter = yield* orchestrator.getThreadProjection(task.childThreadId!); + assert.equal(childAfter.runs[0]?.status, completedRoot ? "completed" : "failed"); + assert.equal(childAfter.attempts.at(-1)?.status, completedRoot ? "completed" : "failed"); + assert.equal( + childAfter.subagents.find((agent) => agent.id === subagentId)?.status, + "failed", + ); + assert.equal(childAfter.nodes.find((node) => node.id === toolNodeId)?.status, "failed"); + assert.isNotNull(childAfter.nodes.find((node) => node.id === toolNodeId)?.completedAt); + const tool = childAfter.turnItems.find((item) => item.id === toolItemId); + assert.equal(tool?.status, "failed"); + assert.isNotNull(tool?.completedAt); + assert.equal(tool?.type === "command_execution" ? tool.output : null, "partial output"); + assert.equal( + childAfter.turnItems.find((item) => item.type === "subagent")?.status, + "failed", + ); + const after = yield* service.taskStatus(scope, task.id); + assert.equal(after.status, completedRoot ? "completed" : "failed"); + assert.equal(after.workState, "result_available"); + const parentTask = (yield* orchestrator.getThreadProjection( + parentThreadId, + )).subagents.find((entry) => entry.id === task.id); + assert.equal(parentTask?.status, completedRoot ? "completed" : "failed"); + }).pipe(Effect.provide(layerTest)); + }), + ), +); diff --git a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts index 90ee44f3d5b6f..a87ce5059ccbf 100644 --- a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts @@ -4168,7 +4168,7 @@ describe("CodexAdapterV2 post-settle continuation", () => { ); }); - it.effect("Stop ends a background command no Codex process tracks any more", () => + it.effect("settles a background command automatically after the Codex process exits", () => Effect.scoped( Effect.gen(function* () { const fs = yield* FileSystem.FileSystem; @@ -4244,27 +4244,19 @@ describe("CodexAdapterV2 post-settle continuation", () => { yield* worker.drain(); yield* Fiber.join(settled); yield* worker.drain(); - const before = yield* orchestrator.getThreadShell(threadId); - assert.deepEqual( - before?.pendingBackgroundTasks?.map((task) => task.kind), - ["command"], - "the thread still shows the command the gone process never finished", - ); - const run = (yield* orchestrator.getThreadProjection(threadId)).runs.at(-1)!; - yield* orchestrator.dispatch({ - type: "run.interrupt", - commandId: CommandId.make("stop-background-untracked"), - threadId, - runId: run.id, - holdQueue: true, - }); - yield* worker.drain(); const projection = yield* orchestrator.getThreadProjection(threadId); assert.deepEqual( projection.turnItems.flatMap((item) => item.type === "command_execution" ? [item.status] : [], ), - ["interrupted"], + ["failed"], + ); + const command = projection.turnItems.find((item) => item.type === "command_execution"); + assert.isDefined(command); + assert.isNotNull(command?.completedAt); + assert.notInclude( + projection.turnItems.map((item) => item.type), + "run_interrupt_request", ); assert.equal(projection.runs.at(-1)?.status, "completed"); assert.deepEqual( diff --git a/apps/server/src/orchestration-v2/EventSink.ts b/apps/server/src/orchestration-v2/EventSink.ts index b210d3c9840ae..75ea1e027414e 100644 --- a/apps/server/src/orchestration-v2/EventSink.ts +++ b/apps/server/src/orchestration-v2/EventSink.ts @@ -68,6 +68,8 @@ export type EventSinkV2Error = typeof EventSinkV2Error.Type; export interface EventSinkV2Shape { readonly write: (input: { + /** Synchronous in-memory handoff after commit, before publication or pending cancellation. */ + readonly onCommitted?: (events: ReadonlyArray) => void; /** Return no events if another thread owns this native session at commit time. */ readonly unlessNativeThreadOwned?: Pick< OrchestrationV2ProviderThread, @@ -409,6 +411,7 @@ const layerBase: Layer.Layer< const writeEffect = Effect.fn("orchestrationV2.EventSink.write")(function* ( input: Parameters[0], + onCommitted?: Parameters[0]["onCommitted"], ) { yield* Effect.annotateCurrentSpan({ "orchestration_v2.command_id": input.commandId ?? null, @@ -452,6 +455,7 @@ const layerBase: Layer.Layer< }).pipe(Effect.withTracerEnabled(false)), ({ committed, outcomes }) => Effect.gen(function* () { + yield* Effect.sync(() => onCommitted?.(committed)); if (input.effects.length > 0) { yield* effectOutbox.notifyAvailable(input.effects.length); } @@ -879,7 +883,7 @@ const layerBase: Layer.Layer< return EventSinkV2.of({ write: (input) => - writeEffect({ ...input, effects: [] }).pipe( + writeEffect({ ...input, effects: [] }, input.onCommitted).pipe( Effect.mapError( (cause) => new EventSinkWriteError({ diff --git a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts index 770da19398485..5953a51684ad2 100644 --- a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts +++ b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts @@ -763,7 +763,7 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { }), ); - it.effect("replays shared provider-session payloads across every bound thread", () => + it.effect("replays provider-session payloads independently for each bound thread", () => Effect.gen(function* () { const eventSink = yield* EventSink.EventSinkV2; const projectionStore = yield* ProjectionStore.ProjectionStoreV2; @@ -819,12 +819,20 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { ], }); - assert.equal( - (yield* projectionStore.getThreadProjection(firstThreadId)).providerSessions[0]?.cwd, - secondSession.cwd, - ); + const assertBindings = Effect.gen(function* () { + assert.deepEqual( + (yield* projectionStore.getThreadProjection(firstThreadId)).providerSessions, + [firstSession], + ); + assert.deepEqual( + (yield* projectionStore.getThreadProjection(secondThreadId)).providerSessions, + [secondSession], + ); + }); + yield* assertBindings; assert.isTrue((yield* maintenance.verify).valid); assert.isTrue((yield* maintenance.rebuild).valid); + yield* assertBindings; }), ); @@ -3666,3 +3674,88 @@ it.effect.each(["cancel transaction", "fail transaction", "cancel publication"] ); }).pipe(Effect.provide(layerDatabase)), ); + +it.effect.each(["cancel transaction", "fail transaction", "cancel publication"] as const)( + "hands off committed ownership only after a successful transaction during %s", + (scenario) => + Effect.gen(function* () { + const paused = yield* Deferred.make(); + const resume = yield* Deferred.make(); + const pause = Deferred.succeed(paused, undefined).pipe( + Effect.andThen(Deferred.await(resume)), + ); + const committed = scenario === "cancel publication"; + const stores = Layer.mergeAll( + layerStoresProvided, + Layer.effect( + EventStore.EventStoreV2, + Effect.gen(function* () { + const delegate = yield* EventStore.EventStoreV2; + return EventStore.EventStoreV2.of({ + ...delegate, + publishCommitted: (events) => + (committed ? pause : Effect.void).pipe( + Effect.andThen(delegate.publishCommitted(events)), + ), + }); + }), + ).pipe(Layer.provide(layerEventStoreProvided)), + Layer.effect( + EffectOutbox.EffectOutboxV2, + Effect.gen(function* () { + const delegate = yield* EffectOutbox.EffectOutboxV2; + return EffectOutbox.EffectOutboxV2.of({ + ...delegate, + enqueue: (effects) => + delegate + .enqueue(effects) + .pipe( + Effect.andThen(committed ? Effect.void : pause), + Effect.andThen( + scenario === "fail transaction" + ? Effect.die("Injected rollback") + : Effect.void, + ), + ), + }); + }), + ).pipe(Layer.provide(layerEffectOutboxProvided)), + layerCommandReceiptStoreProvided, + ProjectStore.layer.pipe(Layer.provide(layerDatabase)), + TurnItemPositionStore.layer.pipe(Layer.provide(layerDatabase)), + ); + yield* Effect.gen(function* () { + const sink = yield* EventSink.EventSinkV2; + const projections = yield* ProjectionStore.ProjectionStoreV2; + const eventStore = yield* EventStore.EventStoreV2; + const now = yield* DateTime.now; + const thread = makeThread(ThreadId.make("thread:commit-handoff"), now); + const events = [threadCreatedEvent({ id: "event:commit-handoff", thread, now })]; + const handedOff: Array = []; + const writer = yield* sink + .write({ + events, + onCommitted: (stored) => { + handedOff.push(...stored.map((entry) => entry.event)); + }, + }) + .pipe(Effect.forkScoped); + yield* Deferred.await(paused); + // The handoff precedes publication, and never sees transaction work that can roll back. + assert.deepEqual(handedOff, committed ? events : []); + if (scenario !== "fail transaction") writer.interruptUnsafe(); + if (scenario !== "cancel transaction") yield* Deferred.succeed(resume, undefined); + assert.equal((yield* Fiber.await(writer))._tag, "Failure"); + assert.deepEqual(handedOff, committed ? events : []); + assert.equal( + (yield* projections.getThreadShell(thread.id))?.id ?? null, + committed ? thread.id : null, + ); + const stored = yield* eventStore.read({ threadId: thread.id }).pipe(Stream.runCollect); + assert.deepEqual( + Array.from(stored, (entry) => entry.event), + committed ? events : [], + ); + }).pipe(Effect.provide(EventSink.layerFromStores.pipe(Layer.provideMerge(stores)))); + }).pipe(Effect.provide(layerDatabase)), +); diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index 0a5e56d98d639..9c1791d2eb235 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -10438,7 +10438,7 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio .map((node) => [node.id, node]), ), turnItems: new Map(), - childTurnItems: new Map( + ownedTurnItems: new Map( input.projection.turnItems .filter( (item) => diff --git a/apps/server/src/orchestration-v2/ProjectionStore.test.ts b/apps/server/src/orchestration-v2/ProjectionStore.test.ts index a29d698d76daf..8d1ef49299abd 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.test.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.test.ts @@ -4775,6 +4775,44 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { ], ); + yield* projectionStore.apply({ + id: EventId.make("event:projection-shared-provider-session:first-replaced"), + type: "provider-session.attached", + threadId: firstThreadId, + occurredAt: now, + payload: { ...session, status: "ready", lastError: null }, + }); + yield* projectionStore.apply({ + id: EventId.make("event:projection-shared-provider-session:second-released"), + type: "provider-session.updated", + threadId: secondThreadId, + occurredAt: now, + payload: { ...session, status: "stopped", lastError: "old runtime stopped" }, + }); + assert.equal( + (yield* projectionStore.getThreadProjection(firstThreadId)).providerSessions[0]?.status, + "ready", + ); + assert.equal( + (yield* projectionStore.getThreadProjection(secondThreadId)).providerSessions[0]?.status, + "stopped", + ); + assert.deepEqual( + (yield* projectionStore.getShellSnapshot()).threads + .filter((thread) => thread.id === firstThreadId || thread.id === secondThreadId) + .map((thread) => thread.lastError), + [null, "old runtime stopped"], + ); + assert.equal( + (yield* projectionStore.getRuntimeRecoveryProjection(firstThreadId)).providerSessions[0] + ?.status, + "ready", + ); + assert.lengthOf( + (yield* projectionStore.getRuntimeRecoveryProjection(secondThreadId)).providerSessions, + 0, + ); + yield* projectionStore.apply({ id: EventId.make("event:projection-shared-provider-session:first-detached"), type: "provider-session.detached", diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index 2e88fdc119081..5937186eab890 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -2221,12 +2221,18 @@ export const layer: Layer.Layer = `; if (event.type === "provider-session.attached") { yield* sql` - INSERT OR IGNORE INTO orchestration_v2_projection_provider_session_bindings ( + INSERT INTO orchestration_v2_projection_provider_session_bindings ( provider_session_id, - thread_id + thread_id, + payload_json ) - VALUES (${event.payload.id}, ${event.threadId}) + VALUES (${event.payload.id}, ${event.threadId}, ${payloadJson}) + ON CONFLICT(provider_session_id, thread_id) DO UPDATE SET payload_json = excluded.payload_json `; + } else { + yield* sql`UPDATE orchestration_v2_projection_provider_session_bindings + SET payload_json = ${payloadJson} + WHERE provider_session_id = ${event.payload.id} AND thread_id = ${event.threadId}`; } break; } @@ -3330,24 +3336,24 @@ export const layer: Layer.Layer = ? Effect.succeed([]) : window === undefined ? sql` - SELECT sessions.payload_json + SELECT COALESCE(bindings.payload_json, sessions.payload_json) AS payload_json FROM orchestration_v2_projection_provider_sessions AS sessions INNER JOIN orchestration_v2_projection_provider_session_bindings AS bindings ON bindings.provider_session_id = sessions.provider_session_id WHERE bindings.thread_id = ${threadId} - ORDER BY sessions.updated_at ASC, sessions.provider_session_id ASC + ORDER BY COALESCE(json_extract(bindings.payload_json, '$.updatedAt'), sessions.updated_at) ASC, sessions.provider_session_id ASC ` : sql` - SELECT DISTINCT sessions.payload_json + SELECT DISTINCT COALESCE(bindings.payload_json, sessions.payload_json) AS payload_json FROM orchestration_v2_projection_provider_sessions AS sessions INNER JOIN orchestration_v2_projection_provider_session_bindings AS bindings ON bindings.provider_session_id = sessions.provider_session_id LEFT JOIN orchestration_v2_projection_provider_threads AS threads ON threads.provider_session_id = sessions.provider_session_id WHERE bindings.thread_id = ${threadId} - AND (sessions.status IN ('starting','running','waiting') + AND (COALESCE(json_extract(bindings.payload_json, '$.status'), sessions.status) IN ('starting','running','waiting') OR threads.provider_thread_id IN (SELECT value FROM json_each(${cohortProviderThreadIds}))) - ORDER BY sessions.updated_at ASC, sessions.provider_session_id ASC + ORDER BY COALESCE(json_extract(bindings.payload_json, '$.updatedAt'), sessions.updated_at) ASC, sessions.provider_session_id ASC `, fields !== undefined && !fields.includes("providerThreads") ? Effect.succeed([]) @@ -3814,14 +3820,14 @@ export const layer: Layer.Layer = SELECT t.payload_json, r.run_id, r.completed_at, item.payload_json AS failure_payload_json, ( - SELECT json_extract(session.payload_json, '$.lastError') + SELECT json_extract(COALESCE(binding.payload_json, session.payload_json), '$.lastError') -- Keep bindings outermost: one provider instance can hold every session. FROM orchestration_v2_projection_provider_session_bindings binding CROSS JOIN orchestration_v2_projection_provider_sessions session ON session.provider_session_id = binding.provider_session_id WHERE binding.thread_id = t.thread_id AND session.provider_instance_id = t.provider_instance_id - ORDER BY session.updated_at DESC, session.provider_session_id DESC + ORDER BY COALESCE(json_extract(binding.payload_json, '$.updatedAt'), session.updated_at) DESC, session.provider_session_id DESC LIMIT 1 ) AS last_error FROM orchestration_v2_projection_threads t @@ -4026,7 +4032,7 @@ export const layer: Layer.Layer = FROM orchestration_v2_projection_provider_sessions AS sessions CROSS 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') + WHERE COALESCE(json_extract(bindings.payload_json, '$.status'), sessions.status) NOT IN ('stopped', 'error') UNION -- A detach drops the binding but keeps the session's status. SELECT sessions.thread_id @@ -4301,11 +4307,11 @@ export const layer: Layer.Layer = WHERE thread_id = ${threadId} AND run_id = ${runId} ORDER BY ordinal ASC, turn_item_id ASC `.pipe(Effect.flatMap(decodeRows(decodeTurnItemPayload, threadId))); const providerSessions = yield* sql` - SELECT sessions.payload_json FROM orchestration_v2_projection_provider_sessions AS sessions + SELECT COALESCE(bindings.payload_json, sessions.payload_json) AS payload_json FROM orchestration_v2_projection_provider_sessions AS sessions INNER JOIN orchestration_v2_projection_provider_session_bindings AS bindings ON bindings.provider_session_id = sessions.provider_session_id WHERE bindings.thread_id = ${threadId} - ORDER BY sessions.updated_at ASC, sessions.provider_session_id ASC + ORDER BY COALESCE(json_extract(bindings.payload_json, '$.updatedAt'), sessions.updated_at) ASC, sessions.provider_session_id ASC `.pipe(Effect.flatMap(decodeRows(decodeProviderSessionPayload, threadId))); const messages = yield* sql` SELECT payload_json FROM orchestration_v2_projection_messages @@ -4508,16 +4514,14 @@ export const layer: Layer.Layer = ORDER BY COALESCE(subagent.started_at, ''), subagent.subagent_id ASC `, sql` - SELECT session.payload_json + SELECT COALESCE(binding.payload_json, session.payload_json) AS payload_json FROM orchestration_v2_projection_provider_sessions AS session + LEFT JOIN orchestration_v2_projection_provider_session_bindings AS binding + ON binding.provider_session_id = session.provider_session_id AND binding.thread_id = ${threadId} WHERE ( - session.provider_session_id IN ( - SELECT binding.provider_session_id - FROM orchestration_v2_projection_provider_session_bindings AS binding - WHERE binding.thread_id = ${threadId} - ) + binding.thread_id IS NOT NULL AND ( - session.status NOT IN ('stopped', 'error') + COALESCE(json_extract(binding.payload_json, '$.status'), session.status) NOT IN ('stopped', 'error') OR session.provider_session_id IN ( SELECT provider_thread.provider_session_id FROM orchestration_v2_projection_provider_threads AS provider_thread @@ -4540,7 +4544,7 @@ export const layer: Layer.Layer = WHERE binding.provider_session_id = session.provider_session_id ) ) - ORDER BY session.updated_at ASC, session.provider_session_id ASC + ORDER BY COALESCE(json_extract(binding.payload_json, '$.updatedAt'), session.updated_at) ASC, session.provider_session_id ASC `, sql` SELECT provider_thread.payload_json @@ -4942,11 +4946,11 @@ export const layer: Layer.Layer = Effect.gen(function* () { const thread = yield* getThread(threadId); const sessionRows = yield* sql` - SELECT session.payload_json FROM orchestration_v2_projection_provider_sessions AS session + SELECT COALESCE(binding.payload_json, session.payload_json) AS payload_json FROM orchestration_v2_projection_provider_sessions AS session JOIN orchestration_v2_projection_provider_session_bindings AS binding ON binding.provider_session_id = session.provider_session_id WHERE binding.thread_id = ${threadId} - ORDER BY session.updated_at ASC, session.provider_session_id ASC + ORDER BY COALESCE(json_extract(binding.payload_json, '$.updatedAt'), session.updated_at) ASC, session.provider_session_id ASC `; const threadRows = targetInstanceId === undefined @@ -4997,7 +5001,7 @@ export const layer: Layer.Layer = request.responseCapability.type !== "live" ? [] : yield* sql` - SELECT session.payload_json FROM orchestration_v2_projection_provider_sessions AS session + SELECT COALESCE(binding.payload_json, session.payload_json) AS payload_json FROM orchestration_v2_projection_provider_sessions AS session JOIN orchestration_v2_projection_provider_session_bindings AS binding ON binding.provider_session_id = session.provider_session_id WHERE binding.thread_id = ${threadId} AND session.provider_session_id = ${request.responseCapability.providerSessionId}`; @@ -5191,7 +5195,7 @@ export const layer: Layer.Layer = WHERE thread_id = ${threadId} AND ordinal > ${targetOrdinal} AND status IN ('completed', 'interrupted', 'failed', 'cancelled') ORDER BY ordinal ASC`, - sql`SELECT session.payload_json + sql`SELECT COALESCE(binding.payload_json, session.payload_json) AS payload_json FROM orchestration_v2_projection_provider_sessions AS session JOIN orchestration_v2_projection_provider_session_bindings AS binding ON binding.provider_session_id = session.provider_session_id @@ -5921,14 +5925,14 @@ export const layer: Layer.Layer = LIMIT 1 ) AS activity_run_started_at, ( - SELECT json_extract(session.payload_json, '$.lastError') + SELECT json_extract(COALESCE(binding.payload_json, session.payload_json), '$.lastError') -- Keep bindings outermost: one provider instance can hold every session. FROM orchestration_v2_projection_provider_session_bindings binding CROSS JOIN orchestration_v2_projection_provider_sessions session ON session.provider_session_id = binding.provider_session_id WHERE binding.thread_id = t.thread_id AND session.provider_instance_id = t.provider_instance_id - ORDER BY session.updated_at DESC, session.provider_session_id DESC + ORDER BY COALESCE(json_extract(binding.payload_json, '$.updatedAt'), session.updated_at) DESC, session.provider_session_id DESC LIMIT 1 ) AS last_error, ( diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts index 5bbb3926e2b66..942e564bbb354 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.test.ts @@ -122,6 +122,7 @@ const layerFlakyReleaseEventSink = (flaky: FlakyReleaseWrites) => ).pipe(Layer.provide(layerTestEventSink)); const layerPausingAttachEventSink = (pause: { + readonly afterCommitResume?: Deferred.Deferred; readonly armed: Ref.Ref; readonly paused: Deferred.Deferred; }) => @@ -135,6 +136,14 @@ const layerPausingAttachEventSink = (pause: { Effect.gen(function* () { const attach = input.events.some((event) => event.type === "provider-session.attached"); if (attach && (yield* Ref.getAndSet(pause.armed, false))) { + const resume = pause.afterCommitResume; + if (resume !== undefined) { + return yield* delegate.write(input).pipe( + Effect.tap(() => Deferred.succeed(pause.paused, undefined)), + Effect.tap(() => Deferred.await(resume)), + Effect.uninterruptible, + ); + } yield* Deferred.succeed(pause.paused, undefined); return yield* Effect.never; } @@ -342,14 +351,19 @@ const makeProviderAdapter = Effect.fnUntraced(function* ( readonly beforeUnload?: Effect.Effect; readonly spawnBeforeOpen?: boolean; readonly scopeCloseReached?: Deferred.Deferred; + readonly beforeClose?: Effect.Effect; } = {}, ) { const mcpSessions = yield* McpProviderSessions.McpProviderSessions; const countClose = Effect.addFinalizer(() => - Ref.update(state, (current) => ({ - ...current, - closeCount: current.closeCount + 1, - })), + (options.beforeClose ?? Effect.void).pipe( + Effect.andThen( + Ref.update(state, (current) => ({ + ...current, + closeCount: current.closeCount + 1, + })), + ), + ), ); // Register after countClose to wedge teardown before its finalizer, matching a provider stream that never yields. const hangClose = Effect.addFinalizer(() => @@ -490,6 +504,7 @@ function layerTest(input: { readonly beforeUnload?: Effect.Effect; readonly spawnBeforeOpen?: boolean; readonly scopeCloseReached?: Deferred.Deferred; + readonly beforeClose?: Effect.Effect; readonly serverSettingsLayer?: ReturnType; readonly projectServiceLayer?: Layer.Layer; }) { @@ -506,6 +521,7 @@ function layerTest(input: { const layerRegistry = Layer.unwrap( makeProviderAdapter(input.state, { failEventStream: input.failEventStream ?? false, + ...(input.beforeClose === undefined ? {} : { beforeClose: input.beforeClose }), ...(input.capabilities === undefined ? {} : { capabilities: input.capabilities }), ...(input.mcpConfigs === undefined ? {} : { mcpConfigs: input.mcpConfigs }), ...(input.beforeOpen === undefined ? {} : { beforeOpen: input.beforeOpen }), @@ -2853,76 +2869,108 @@ it.effect("ProviderSessionManagerV2 retries release records that failed to persi ); // regression: #15048 runs got stuck after a failed session release -it.effect("ProviderSessionManagerV2 release retries leave a replacement session alone", () => - Effect.gen(function* () { - const state = yield* Ref.make(emptyState); - const flaky: FlakyReleaseWrites = { - failing: yield* Ref.make<"none" | "session" | "session-and-requests">("none"), - failures: yield* Queue.unbounded(), - }; - const effect = Effect.gen(function* () { - const eventSink = yield* EventSink.EventSinkV2; - const idAllocator = yield* IdAllocator.IdAllocatorV2; - const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; - const projectionStore = yield* ProjectionStore.ProjectionStoreV2; - const threadId = ThreadId.make("thread-provider-session-manager-release-replacement"); - const providerSessionId = yield* idAllocator.allocate.providerSession({ - providerInstanceId: modelSelection.instanceId, - threadId, - }); - const writePendingRequest = Effect.gen(function* () { - const now = yield* DateTime.now; - const request = yield* makePendingRuntimeRequestEvents({ - idAllocator, +it.effect.each([false, true])( + "ProviderSessionManagerV2 release retries fence each thread after replacement (released=%s)", + (replacementReleased) => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const flaky: FlakyReleaseWrites = { + failing: yield* Ref.make<"none" | "session" | "session-and-requests">("none"), + failures: yield* Queue.unbounded(), + }; + const effect = Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const projectionStore = yield* ProjectionStore.ProjectionStoreV2; + const threadId = ThreadId.make("thread-provider-session-manager-release-replacement"); + const providerSessionId = yield* idAllocator.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, threadId, + }); + const writePendingRequest = Effect.gen(function* () { + const now = yield* DateTime.now; + const request = yield* makePendingRuntimeRequestEvents({ + idAllocator, + threadId, + providerSessionId, + providerThread: makeProviderThread({ idAllocator, threadId, providerSessionId, now }), + now, + }); + yield* eventSink.write({ events: request.events }); + return request.requestId; + }); + yield* eventSink.write({ + events: [ + yield* makeThreadCreatedEvent({ idAllocator, threadId, now: yield* DateTime.now }), + ], + }); + const otherThreadId = ThreadId.make("thread:retry:other"); + yield* eventSink.write({ + events: [ + yield* makeThreadCreatedEvent({ + idAllocator, + threadId: otherThreadId, + now: yield* DateTime.now, + }), + ], + }); + const oldRequestId = yield* writePendingRequest; + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + yield* manager.open({ + threadId: otherThreadId, providerSessionId, - providerThread: makeProviderThread({ idAllocator, threadId, providerSessionId, now }), - now, + modelSelection, + runtimePolicy, }); - yield* eventSink.write({ events: request.events }); - return request.requestId; - }); - yield* eventSink.write({ - events: [ - yield* makeThreadCreatedEvent({ idAllocator, threadId, now: yield* DateTime.now }), - ], - }); - const oldRequestId = yield* writePendingRequest; - yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); - yield* Ref.set(flaky.failing, "session-and-requests"); - assert.isTrue(Exit.isFailure(yield* Effect.exit(manager.close(providerSessionId)))); - yield* Queue.take(flaky.failures); - yield* Queue.take(flaky.failures); - yield* Ref.set(flaky.failing, "none"); + yield* Ref.set(flaky.failing, "session-and-requests"); + assert.isTrue(Exit.isFailure(yield* Effect.exit(manager.close(providerSessionId)))); + yield* Queue.take(flaky.failures); + yield* Queue.take(flaky.failures); + yield* Ref.set(flaky.failing, "none"); - yield* TestClock.adjust("500 millis"); - yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); - const newRequestId = yield* writePendingRequest; - const replacementStatus = (yield* projectionStore.getThreadProjection( - threadId, - )).providerSessions.at(-1)?.status; - const settled = yield* eventSink - .stream({ + yield* TestClock.adjust("500 millis"); + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + const newRequestId = yield* writePendingRequest; + if (replacementReleased) + yield* manager.release({ + providerSessionId, + reason: "runtime_error", + detail: "replacement failed", + }); + const replacementStatus = (yield* projectionStore.getThreadProjection( threadId, - afterSequence: yield* eventSink.latestSequence({ threadId }), - eventType: "runtime-request.updated", - }) - .pipe(Stream.runHead, Effect.forkScoped); - yield* TestClock.adjust("500 millis"); - yield* Fiber.join(settled); + )).providerSessions.at(-1)?.status; + const settled = yield* eventSink + .stream({ + threadId: replacementReleased ? otherThreadId : threadId, + afterSequence: yield* eventSink.latestSequence({ threadId }), + eventType: replacementReleased ? "provider-session.updated" : "runtime-request.updated", + }) + .pipe(Stream.runHead, Effect.forkScoped); + yield* TestClock.adjust("500 millis"); + yield* Fiber.join(settled); - const projection = yield* projectionStore.getThreadProjection(threadId); - const request = (id: typeof oldRequestId) => - projection.runtimeRequests.find((candidate) => candidate.id === id); - assert.equal(request(oldRequestId)?.responseCapability.type, "not_resumable"); - assert.equal(request(newRequestId)?.responseCapability.type, "live"); - assert.equal(projection.providerSessions.at(-1)?.status, replacementStatus); - }); + const projection = yield* projectionStore.getThreadProjection(threadId); + const request = (id: typeof oldRequestId) => + projection.runtimeRequests.find((candidate) => candidate.id === id); + assert.equal(request(oldRequestId)?.responseCapability.type, "not_resumable"); + assert.equal( + request(newRequestId)?.responseCapability.type, + replacementReleased ? "not_resumable" : "live", + ); + assert.equal(projection.providerSessions.at(-1)?.status, replacementStatus); + assert.equal( + (yield* projectionStore.getThreadProjection(otherThreadId)).providerSessions.at(-1) + ?.status, + "stopped", + ); + }); - yield* effect.pipe( - Effect.provide(layerTest({ state, idleTimeoutMs: 60_000, flakyReleaseWrites: flaky })), - ); - }), + yield* effect.pipe( + Effect.provide(layerTest({ state, idleTimeoutMs: 60_000, flakyReleaseWrites: flaky })), + ); + }), ); // regression: #15048 runs got stuck after a failed session release @@ -3819,3 +3867,154 @@ it.effect("ProviderSessionManagerV2 gives a delegated child its project's device assert.isTrue(child?.capabilities?.has("device")); }), ); + +it.effect.each([false, true])( + "ProviderSessionManagerV2 fences a stale close per thread after replacement (released=%s)", + (replacementReleased) => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const closing = yield* Deferred.make(); + const resume = yield* Deferred.make(); + const first = yield* Ref.make(true); + yield* Effect.gen(function* () { + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const sink = yield* EventSink.EventSinkV2; + const ids = yield* IdAllocator.IdAllocatorV2; + const store = yield* ProjectionStore.ProjectionStoreV2; + const threadId = ThreadId.make("thread:release-owner:a"); + const otherThreadId = ThreadId.make("thread:release-owner:b"); + const providerSessionId = yield* ids.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + const now = yield* DateTime.now; + yield* sink.write({ + events: [ + yield* makeThreadCreatedEvent({ idAllocator: ids, threadId, now }), + yield* makeThreadCreatedEvent({ idAllocator: ids, threadId: otherThreadId, now }), + ], + }); + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + yield* manager.open({ + threadId: otherThreadId, + providerSessionId, + modelSelection, + runtimePolicy, + }); + const oldClose = yield* manager.close(providerSessionId).pipe(Effect.forkScoped); + yield* Deferred.await(closing); + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + if (replacementReleased) + yield* manager.release({ + providerSessionId, + reason: "runtime_error", + detail: "replacement failed", + }); + yield* Deferred.succeed(resume, undefined); + yield* Fiber.join(oldClose); + const a = (yield* store.getThreadProjection(threadId)).providerSessions.at(-1); + const b = (yield* store.getThreadProjection(otherThreadId)).providerSessions.at(-1); + assert.equal(a?.status, replacementReleased ? "error" : "ready"); + if (replacementReleased) assert.equal(a?.lastError, "replacement failed"); + assert.equal(b?.status, "stopped"); + }).pipe( + Effect.provide( + layerTest({ + state, + idleTimeoutMs: 60_000, + beforeClose: Ref.getAndSet(first, false).pipe( + Effect.flatMap((isFirst) => + isFirst + ? Deferred.succeed(closing, undefined).pipe( + Effect.andThen(Deferred.await(resume)), + ) + : Effect.void, + ), + ), + }), + ), + ); + }), +); + +it.effect.each([false, true])( + "ProviderSessionManagerV2 preserves release ownership when attachment cancellation follows commit=%s", + (committed) => + Effect.gen(function* () { + const state = yield* Ref.make(emptyState); + const closing = yield* Deferred.make(); + const resumeOld = yield* Deferred.make(); + const first = yield* Ref.make(true); + const armed = yield* Ref.make(false); + const paused = yield* Deferred.make(); + const resumeAttach = yield* Deferred.make(); + yield* Effect.gen(function* () { + const manager = yield* ProviderSessionManager.ProviderSessionManagerV2; + const sink = yield* EventSink.EventSinkV2; + const ids = yield* IdAllocator.IdAllocatorV2; + const store = yield* ProjectionStore.ProjectionStoreV2; + const threadId = ThreadId.make("thread:attachment-commit-interrupt"); + const providerSessionId = yield* ids.allocate.providerSession({ + providerInstanceId: modelSelection.instanceId, + threadId, + }); + yield* sink.write({ + events: [ + yield* makeThreadCreatedEvent({ + idAllocator: ids, + threadId, + now: yield* DateTime.now, + }), + ], + }); + yield* manager.open({ threadId, providerSessionId, modelSelection, runtimePolicy }); + const oldClose = yield* manager + .release({ providerSessionId, reason: "runtime_error", detail: "old runtime failure" }) + .pipe(Effect.forkScoped); + yield* Deferred.await(closing); + yield* Ref.set(armed, true); + const replacementOpen = yield* manager + .open({ + threadId, + providerSessionId, + modelSelection, + runtimePolicy, + }) + .pipe(Effect.forkScoped); + yield* Deferred.await(paused); + // Post-commit cancellation waits on masked completion; pre-commit cancellation must abort. + replacementOpen.interruptUnsafe(); + yield* Deferred.succeed(resumeAttach, undefined); + const openExit = yield* Fiber.await(replacementOpen); + assert.isTrue(Exit.isFailure(openExit)); + const replacement = (yield* store.getThreadProjection(threadId)).providerSessions[0]; + assert.equal(replacement?.status, "stopped"); + yield* Deferred.succeed(resumeOld, undefined); + yield* Fiber.join(oldClose); + const after = (yield* store.getThreadProjection(threadId)).providerSessions[0]; + assert.equal(after?.status, committed ? replacement?.status : "error"); + assert.equal(after?.lastError, committed ? replacement?.lastError : "old runtime failure"); + }).pipe( + Effect.provide( + layerTest({ + state, + idleTimeoutMs: 60_000, + pauseAttachWrite: { + armed, + paused, + ...(committed ? { afterCommitResume: resumeAttach } : {}), + }, + beforeClose: Ref.getAndSet(first, false).pipe( + Effect.flatMap((isFirst) => + isFirst + ? Deferred.succeed(closing, undefined).pipe( + Effect.andThen(Deferred.await(resumeOld)), + ) + : Effect.void, + ), + ), + }), + ), + ); + }), +); diff --git a/apps/server/src/orchestration-v2/ProviderSessionManager.ts b/apps/server/src/orchestration-v2/ProviderSessionManager.ts index 03a7278774a64..6c453165ec46c 100644 --- a/apps/server/src/orchestration-v2/ProviderSessionManager.ts +++ b/apps/server/src/orchestration-v2/ProviderSessionManager.ts @@ -406,6 +406,16 @@ export const layerWithOptions = ( const releaseRecordRetries = yield* FiberSet.make(); const nextSubscriberId = yield* Ref.make(0); const sessionOpen = yield* KeyedLock.make(); + const sessionRecords = yield* KeyedLock.make(); + // Retain ownership only while a release is closing or awaiting persistence. + const pendingReleaseThreads = new Map< + ProviderAdapter.ProviderAdapterV2SessionRuntime, + { + readonly providerSessionId: ProviderSessionId; + readonly threadIds: Set; + retrying: boolean; + } + >(); const threadAttachment = yield* KeyedLock.make(); const threadAttachmentKey = (input: { readonly providerSessionId: ProviderSessionId; @@ -582,33 +592,67 @@ export const layerWithOptions = ( const writeProviderSessionEvents = Effect.fn( "ProviderSessionManager.writeProviderSessionEvents", - )(function* (input: { - readonly runtime: ProviderAdapter.ProviderAdapterV2SessionRuntime; - readonly threadIds: Iterable; - readonly type: "provider-session.attached" | "provider-session.updated"; - readonly payload: OrchestrationV2ProviderSession; - }) { - const now = yield* DateTime.now; - const events = yield* Effect.forEach(input.threadIds, (threadId) => - Effect.gen(function* () { - return { - id: yield* idAllocator.allocate.event({ + )( + function* (input: { + readonly runtime: ProviderAdapter.ProviderAdapterV2SessionRuntime; + readonly threadIds: Iterable; + readonly type: "provider-session.attached" | "provider-session.updated"; + readonly payload: OrchestrationV2ProviderSession; + }) { + const now = yield* DateTime.now; + const events = yield* Effect.forEach(input.threadIds, (threadId) => + Effect.gen(function* () { + return { + id: yield* idAllocator.allocate.event({ + threadId, + providerSessionId: input.runtime.providerSessionId, + }), + type: input.type, threadId, - providerSessionId: input.runtime.providerSessionId, + driver: input.runtime.driver, + providerInstanceId: input.runtime.instanceId, + occurredAt: now, + payload: input.payload, + } satisfies OrchestrationV2DomainEvent; + }), + ); + if (events.length > 0) { + yield* eventSink.write({ + events, + // EventSink runs this inside its masked post-commit boundary. A pending + // interrupt must not leave the old release owning a durable attachment. + onCommitted: (committed) => { + if (input.type !== "provider-session.attached") return; + for (const [runtime, release] of pendingReleaseThreads) { + if ( + runtime !== input.runtime && + release.providerSessionId === input.runtime.providerSessionId + ) { + for (const event of committed) release.threadIds.delete(event.event.threadId); + } + } + }, + }); + } + }, + (effect, input) => + input.type !== "provider-session.attached" + ? effect + : Effect.gen(function* () { + const entry = (yield* Ref.get(sessions)).get( + sessionKey(input.runtime.providerSessionId), + ); + if (entry?.runtime !== input.runtime) return yield* Effect.interrupt; + // Pending attachment writes end with either their caller or their runtime. + yield* Effect.acquireUseRelease( + sessionRecords + .withLock(input.runtime.providerSessionId, effect) + .pipe(Effect.forkIn(entry.scope)), + Fiber.join, + Fiber.interrupt, + ); }), - type: input.type, - threadId, - driver: input.runtime.driver, - providerInstanceId: input.runtime.instanceId, - occurredAt: now, - payload: input.payload, - } satisfies OrchestrationV2DomainEvent; - }), - ); - if (events.length > 0) { - yield* eventSink.write({ events }); - } - }); + ); const writeReleasedSessionEvents = Effect.fn( "ProviderSessionManager.writeReleasedSessionEvents", @@ -627,7 +671,7 @@ export const layerWithOptions = ( }; yield* writeProviderSessionEvents({ runtime: input.entry.runtime, - threadIds: input.entry.attachedThreadIds, + threadIds: pendingReleaseThreads.get(input.entry.runtime)?.threadIds ?? [], type: "provider-session.updated", payload, }); @@ -850,13 +894,15 @@ export const layerWithOptions = ( readonly detail?: string; readonly releasedBackgroundTasks: ReadonlyArray; readonly releasedAt: DateTime.Utc; - readonly replaced: boolean; }) => Effect.all( [ - input.replaced - ? Effect.succeed(Exit.void) - : Effect.exit(writeReleasedSessionEvents(input)), + Effect.exit( + sessionRecords.withLock( + input.entry.runtime.providerSessionId, + writeReleasedSessionEvents(input), + ), + ), Effect.exit( writeReleasedRuntimeRequestEvents(input).pipe( input.entry.requestEventPermit.withPermits(1), @@ -868,20 +914,10 @@ export const layerWithOptions = ( ).pipe(Effect.flatMap(Exit.asVoidAll)); // Retry after removal from the live map; later releases cannot clear stale ready sessions and approvals. - const retryReleaseRecords = ( - input: Omit[0], "replaced">, - ) => { + const retryReleaseRecords = (input: Parameters[0]) => { const providerSessionId = input.entry.runtime.providerSessionId; const attempt = Effect.gen(function* () { - const exit = yield* Effect.exit( - sessionOpen.withLock( - providerSessionId, - Effect.gen(function* () { - const replaced = (yield* Ref.get(sessions)).has(sessionKey(providerSessionId)); - yield* writeReleaseRecords({ ...input, replaced }); - }), - ), - ); + const exit = yield* Effect.exit(writeReleaseRecords(input)); if (Exit.isSuccess(exit) || Cause.hasInterruptsOnly(exit.cause)) return yield* exit; yield* Effect.logWarning("orchestration-v2.provider-session-release-records-failed", { providerSessionId, @@ -899,7 +935,18 @@ export const layerWithOptions = ( ), }), Effect.delay("1 second"), - FiberSet.run(releaseRecordRetries), + Effect.ensuring(Effect.sync(() => pendingReleaseThreads.delete(input.entry.runtime))), + (retry) => + Effect.uninterruptibleMask((restore) => + restore(FiberSet.run(releaseRecordRetries)(retry)).pipe( + Effect.tap(() => + Effect.sync(() => { + const pending = pendingReleaseThreads.get(input.entry.runtime); + if (pending !== undefined) pending.retrying = true; + }), + ), + ), + ), ); }; @@ -953,6 +1000,11 @@ export const layerWithOptions = ( } const updated = new Map(current); updated.delete(key); + pendingReleaseThreads.set(existing.runtime, { + providerSessionId: input.providerSessionId, + threadIds: new Set(existing.attachedThreadIds), + retrying: false, + }); return ["removed", updated] as const; }), DateTime.now, @@ -1050,13 +1102,12 @@ export const layerWithOptions = ( releasedBackgroundTasks, releasedAt, }; - const recorded = yield* Effect.exit( - writeReleaseRecords({ ...records, replaced: false }), - ); + const recorded = yield* Effect.exit(writeReleaseRecords(records)); if (Exit.isFailure(recorded)) { yield* retryReleaseRecords(records); return yield* recorded; } + pendingReleaseThreads.delete(entry.runtime); if (Option.isSome(closeExit) && Exit.isFailure(closeExit.value)) { return yield* Effect.failCause(closeExit.value.cause); } @@ -1076,7 +1127,12 @@ export const layerWithOptions = ( onNone: () => Effect.void, onSome: (entry) => // Revoke recorded credentials after process death, except those held by live replacements. - Ref.get(sessions).pipe( + Effect.sync(() => { + if (pendingReleaseThreads.get(entry.runtime)?.retrying !== true) { + pendingReleaseThreads.delete(entry.runtime); + } + }).pipe( + Effect.andThen(Ref.get(sessions)), Effect.flatMap((current) => Effect.forEach( entry.mcpCredentialIdByThread, @@ -1930,7 +1986,7 @@ export const layerWithOptions = ( }); }, (effect, entry, _event) => - effect.pipe( + sessionRecords.withLock(entry.runtime.providerSessionId, effect).pipe( Effect.catchCause((cause) => Effect.logWarning("orchestration-v2.driver-session.status-persist-failed", { providerSessionId: entry.runtime.providerSessionId, diff --git a/apps/server/src/orchestration-v2/RunExecutionService.test.ts b/apps/server/src/orchestration-v2/RunExecutionService.test.ts index dee246855b815..5e814339d2359 100644 --- a/apps/server/src/orchestration-v2/RunExecutionService.test.ts +++ b/apps/server/src/orchestration-v2/RunExecutionService.test.ts @@ -3811,3 +3811,157 @@ it.effect("keeps a failed root's still-running native children live while the ru ); }), ); + +it.effect.each([false, true])( + "settles an ordinary root tool on provider exit (superseded=%s)", + (superseded) => + Effect.gen(function* () { + const key = `ordinary-tool-exit:${superseded}`; + const ids = backgroundScenarioIds(key); + const { finalEvents } = yield* captureRootRunTermination({ + key, + shouldFinalizeRun: () => Effect.succeed(true), + rejectTerminalWrite: superseded, + events: () => { + const node = { + ...makeRunOwnedSubagentNodeFixture({ ids, status: "running" }), + kind: "tool_call" as const, + }; + const item = { ...backgroundTurnItemEvent(ids, "command_execution", "running", 1) }; + if (item.type !== "turn_item.updated") throw new Error("expected tool"); + return Stream.make( + { type: "node.updated" as const, driver, node }, + { ...item, turnItem: { ...item.turnItem, nodeId: node.id, output: "partial output" } }, + { + type: "node.updated" as const, + driver, + node: { + ...node, + id: NodeId.make("foreign-tool"), + providerTurnId: ProviderTurnId.make("foreign-turn"), + }, + }, + { + ...item, + turnItem: { + ...item.turnItem, + id: TurnItemId.make("foreign-item"), + nodeId: NodeId.make("foreign-tool"), + providerTurnId: ProviderTurnId.make("foreign-turn"), + }, + }, + ).pipe( + Stream.concat( + Stream.fail( + new ProviderAdapter.ProviderAdapterEventStreamError({ + driver, + providerSessionId: ProviderSessionId.make(`session:${key}`), + cause: "process exited", + }), + ), + ), + ); + }, + }); + if (superseded) { + assert.deepEqual(finalEvents, []); + return; + } + const node = finalEvents.flatMap((event) => + event.type === "node.updated" && event.payload.id === ids.subagentNodeId + ? [event.payload] + : [], + )[0]; + const tool = finalEvents.flatMap((event) => + event.type === "turn-item.updated" && event.payload.type === "command_execution" + ? [event.payload] + : [], + )[0]; + assert.isFalse( + finalEvents.some( + (event) => + (event.type === "node.updated" && event.payload.id === NodeId.make("foreign-tool")) || + (event.type === "turn-item.updated" && + event.payload.id === TurnItemId.make("foreign-item")), + ), + ); + assert.equal(node?.status, "failed"); + assert.isNotNull(node?.completedAt); + assert.equal(tool?.status, "failed"); + assert.isNotNull(tool?.completedAt); + assert.equal(tool?.output, "partial output"); + }), +); + +it.effect.each([false, true])( + "settles native work after a completed root loses its provider (superseded=%s)", + (superseded) => + Effect.gen(function* () { + const key = `completed-root-exit:${superseded}`; + const providerInstanceId = ProviderInstanceId.make("codex"); + const ids = backgroundScenarioIds(key); + const finalized = yield* Ref.make(false); + const { finalEvents } = yield* captureRootRunTermination({ + key, + shouldFinalizeRun: () => Effect.succeed(true), + beforeTerminalWrite: () => + Ref.getAndSet(finalized, true).pipe( + Effect.map((alreadyFinalized) => !superseded || !alreadyFinalized), + ), + events: () => + Stream.make( + { + type: "subagent.updated" as const, + driver, + subagent: makeRunOwnedSubagentFixture({ + ids, + providerInstanceId, + childThreadId: ids.childThreadId, + driver, + status: "running", + }), + }, + { + type: "node.updated" as const, + driver, + node: makeRunOwnedSubagentNodeFixture({ ids, status: "running" }), + }, + { + type: "turn_item.updated" as const, + driver, + turnItem: makeRunOwnedSubagentTurnItemFixture({ + ids, + providerInstanceId, + childThreadId: ids.childThreadId, + driver, + status: "running", + }), + }, + rootTerminalEvent(ids, "completed"), + ).pipe( + Stream.concat( + Stream.fail( + new ProviderAdapter.ProviderAdapterEventStreamError({ + driver, + providerSessionId: ProviderSessionId.make(`session:${key}`), + cause: "process exited", + }), + ), + ), + ), + }); + assert.isTrue( + finalEvents.some( + (event) => event.type === "run.updated" && event.payload.status === "waiting", + ), + ); + assert.isFalse( + finalEvents.some( + (event) => event.type === "run.updated" && event.payload.status === "failed", + ), + ); + const cascaded = finalEvents.filter((event) => event.type === "subagent.updated"); + assert.equal(cascaded.length, superseded ? 0 : 1); + if (!superseded) assert.equal(cascaded[0]?.payload.status, "failed"); + }), +); diff --git a/apps/server/src/orchestration-v2/RunExecutionService.ts b/apps/server/src/orchestration-v2/RunExecutionService.ts index b9b2b09bd227e..9256e11243120 100644 --- a/apps/server/src/orchestration-v2/RunExecutionService.ts +++ b/apps/server/src/orchestration-v2/RunExecutionService.ts @@ -154,7 +154,7 @@ type SubagentTurnItem = Extract; readonly turnItems: ReadonlyMap; - readonly childTurnItems: ReadonlyMap; + readonly ownedTurnItems: ReadonlyMap; readonly nodes: ReadonlyMap; readonly linkedChildThreadIds: ReadonlySet; }; @@ -187,7 +187,7 @@ function emptyOpenRunOwnedSubagentProjection(): OpenRunOwnedSubagentProjection { return { subagents: new Map(), turnItems: new Map(), - childTurnItems: new Map(), + ownedTurnItems: new Map(), nodes: new Map(), linkedChildThreadIds: new Set(), }; @@ -296,8 +296,13 @@ export const cascadeTerminalizeRunOwnedSubagents = Effect.fn( }); } } - for (const turnItem of input.open.childTurnItems.values()) { - if (!childThreadIds.has(turnItem.threadId) || isSettledTurnItemStatus(turnItem.status)) { + for (const turnItem of input.open.ownedTurnItems.values()) { + if ( + (turnItem.threadId === input.run.threadId + ? turnItem.runId !== input.run.id + : !childThreadIds.has(turnItem.threadId)) || + isSettledTurnItemStatus(turnItem.status) + ) { continue; } events.push({ @@ -693,7 +698,7 @@ export const layer: Layer.Layer< const hasOpenSubagentProjection = open.subagents.size > 0 || open.turnItems.size > 0 || - open.childTurnItems.size > 0 || + open.ownedTurnItems.size > 0 || open.nodes.size > 0; const cascadedSubagentEvents = isRunOwnedSubagentTerminalStatus(input.terminal.status) && @@ -1150,13 +1155,18 @@ export const layer: Layer.Layer< } } if (event.type === "node.updated") { - const belongsToRootSubagent = - event.node.kind === "subagent" && event.node.runId === input.run.id; + const belongsToRootNode = + event.node.runId === input.run.id && + (event.node.kind === "subagent" || + (event.node.id !== input.rootNode.id && + event.node.threadId === input.run.threadId && + (event.node.providerTurnId === null || + event.node.providerTurnId === routing.rootProviderTurnId))); const belongsToOwnedChildThread = event.node.threadId !== input.run.threadId && routing.ownedThreadIds.has(event.node.threadId); if ( - !belongsToRootSubagent && + !belongsToRootNode && !( belongsToOwnedChildThread && (event.node.runId === null || event.node.runId === input.run.id) @@ -1198,18 +1208,22 @@ export const layer: Layer.Layer< }); } if ( - belongsToOwnedChildThread && - (event.turnItem.runId === null || belongsToRootRun) && + ((belongsToOwnedChildThread && (event.turnItem.runId === null || belongsToRootRun)) || + (belongsToRootRun && + event.turnItem.threadId === input.run.threadId && + event.turnItem.type !== "subagent" && + (event.turnItem.providerTurnId === null || + event.turnItem.providerTurnId === routing.rootProviderTurnId))) && deliverable ) { yield* Ref.update(openRunOwnedSubagents, (current) => { - const childTurnItems = new Map(current.childTurnItems); + const ownedTurnItems = new Map(current.ownedTurnItems); if (isSettledTurnItemStatus(event.turnItem.status)) { - childTurnItems.delete(event.turnItem.id); + ownedTurnItems.delete(event.turnItem.id); } else { - childTurnItems.set(event.turnItem.id, event.turnItem); + ownedTurnItems.set(event.turnItem.id, event.turnItem); } - return { ...current, childTurnItems }; + return { ...current, ownedTurnItems }; }); } if (belongsToRootRun && event.turnItem.type === "subagent") { @@ -1311,15 +1325,15 @@ export const layer: Layer.Layer< for (const subagent of current.subagents.values()) { if (subagent.threadId === moved.childThreadId) subagentIds.add(subagent.id); } - const childTurnItems = new Map( - [...current.childTurnItems].filter( + const ownedTurnItems = new Map( + [...current.ownedTurnItems].filter( ([, item]) => item.threadId !== moved.childThreadId, ), ); const turnItems = new Map(current.turnItems); turnItems.delete(moved.id); const detachedItems = new Set( - [...current.childTurnItems.values(), ...current.turnItems.values()] + [...current.ownedTurnItems.values(), ...current.turnItems.values()] .filter( (item) => item.threadId === moved.childThreadId || item.nodeId === moved.id, ) @@ -1332,7 +1346,7 @@ export const layer: Layer.Layer< [...current.subagents].filter(([id]) => !subagentIds.has(id)), ), turnItems, - childTurnItems, + ownedTurnItems, nodes: new Map( [...current.nodes].filter( ([id, node]) => id !== moved.id && node.threadId !== moved.childThreadId, @@ -1518,13 +1532,13 @@ export const layer: Layer.Layer< runId: input.run.id, cause, }); - if (yield* Ref.get(rootRunFinalized)) return; + const finalized = yield* Ref.get(rootRunFinalized); const providerThread = yield* Ref.get(latestProviderThread); const routing = yield* Ref.get(eventRouting); const rootProviderTurn = yield* Ref.get(latestRootProviderTurn); // Stop background work too; the failed consumer cannot persist terminal events. const abandonedProviderTurn = - routing.rootProviderTurnId !== null && !routing.rootTurnEnded + !finalized && routing.rootProviderTurnId !== null && !routing.rootTurnEnded ? yield* input.session .interruptTurn({ providerThread, @@ -1561,7 +1575,7 @@ export const layer: Layer.Layer< const openSubagents = yield* Ref.get(openRunOwnedSubagents); const squashed = Cause.squash(cause); let nextRetryWarningAt = completedAt; - yield* writeFinalRunEvents({ + const failRun = writeFinalRunEvents({ run: input.run, rootNode: input.rootNode, checkpointScope: input.checkpointScope, @@ -1586,7 +1600,34 @@ export const layer: Layer.Layer< ), failureItemPersisted: false, refreshAfterTurn, - }).pipe( + }); + const failRemainingWork = Effect.gen(function* () { + const events = yield* cascadeTerminalizeRunOwnedSubagents({ + run: input.run, + open: openSubagents, + status: "failed", + completedAt, + allocateEventId: () => + idAllocator.allocate.event({ threadId: input.run.threadId }), + }); + if (events.length === 0) return; + const terminal = yield* Ref.get(terminalEvent); + const statuses = + terminal?.status === "completed" + ? (["waiting", "completed"] as const) + : [terminal?.status ?? "failed"]; + for (const expectedStatus of statuses) { + const result = yield* eventSink.writeIfRunCurrent({ + threadId: input.run.threadId, + runId: input.run.id, + activeAttemptId: input.attempt.id, + expectedStatus, + events, + }); + if (result.committed) return; + } + }); + yield* (finalized ? failRemainingWork : failRun).pipe( // Storage failures can be defects; retry both failure channels. Effect.catchCause((writeCause) => Effect.gen(function* () { diff --git a/apps/server/src/orchestration-v2/testkit/fixtures/index.ts b/apps/server/src/orchestration-v2/testkit/fixtures/index.ts index 1782f6f072a1f..b5aed700ba896 100644 --- a/apps/server/src/orchestration-v2/testkit/fixtures/index.ts +++ b/apps/server/src/orchestration-v2/testkit/fixtures/index.ts @@ -1798,7 +1798,7 @@ export const ORCHESTRATOR_REPLAY_FIXTURES: ReadonlyArray thread.id === projection.thread.id); assert.deepEqual( - before?.pendingBackgroundTasks?.map((task) => task.kind), - ["command"], - "the failed turn leaves its command on the Waiting strip", + afterExit?.pendingBackgroundTasks, + [], + "provider exit settles the failed turn's background command", ); const command = projection.turnItems.find((item) => item.type === "command_execution"); - assert.equal(command?.status, "interrupted"); + assert.isDefined(command); + assert.equal(command?.status, "failed"); assert.isNotNull(command?.completedAt); const shell = result.shellSnapshot.threads.find((thread) => thread.id === projection.thread.id); assert.deepEqual(shell?.pendingBackgroundTasks ?? [], []); - assert.include( + assert.notInclude( projection.turnItems.map((item) => item.type), "run_interrupt_request", ); diff --git a/apps/server/src/persistence/Migrations.test.ts b/apps/server/src/persistence/Migrations.test.ts index 54e4a711aa8a1..957b64fdd6c60 100644 --- a/apps/server/src/persistence/Migrations.test.ts +++ b/apps/server/src/persistence/Migrations.test.ts @@ -29,6 +29,47 @@ const v2TableCount = Effect.gen(function* () { }); describe("runMigrations", () => { + it.effect.each([75, 76])( + "keeps the shared session snapshot for every existing binding after migration %s", + (throughId) => + Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + yield* runMigrations({ toMigrationInclusive: throughId }); + const payload = '{"id":"shared","status":"error","lastError":"process exited"}'; + yield* sql`INSERT INTO orchestration_v2_projection_provider_sessions + (provider_session_id, thread_id, provider, status, updated_at, payload_json) + VALUES ('shared', 'a', 'codex', 'error', '2026-10-30T00:00:00.000Z', ${payload})`; + for (const threadId of ["a", "b"]) { + yield* sql`INSERT INTO orchestration_v2_projection_provider_session_bindings + (provider_session_id, thread_id) VALUES ('shared', ${threadId})`; + } + assert.deepStrictEqual( + yield* runMigrations(), + throughId === 75 + ? [ + [76, "ProjectionProjectsPinArchive"], + [77, "ProviderSessionBindingPayloads"], + ] + : [[77, "ProviderSessionBindingPayloads"]], + ); + assert.deepStrictEqual(yield* history, migrationManifest); + assert.deepStrictEqual(yield* runMigrations(), []); + const projectColumns = yield* sql<{ + readonly name: string; + }>`PRAGMA table_info(projection_projects)`; + assert.isTrue(projectColumns.some((column) => column.name === "pinned_at")); + assert.isTrue(projectColumns.some((column) => column.name === "archived_at")); + assert.deepEqual( + yield* sql`SELECT thread_id, payload_json + FROM orchestration_v2_projection_provider_session_bindings ORDER BY thread_id`, + [ + { thread_id: "a", payload_json: payload }, + { thread_id: "b", payload_json: payload }, + ], + ); + }).pipe(Effect.provide(NodeSqliteClient.layer({ filename: ":memory:" }))), + ); + it.effect("refuses to start when another migration took a V2 id", () => Effect.gen(function* () { const sql = yield* SqlClient.SqlClient; diff --git a/apps/server/src/persistence/Migrations.ts b/apps/server/src/persistence/Migrations.ts index e13c2b449d9e7..014d9469eecd0 100644 --- a/apps/server/src/persistence/Migrations.ts +++ b/apps/server/src/persistence/Migrations.ts @@ -79,6 +79,7 @@ import Migration0073 from "./Migrations/073_ActiveTurnItemIndexes.ts"; import Migration0074 from "./Migrations/074_ExcludeAttachmentPlaceholdersFromSearch.ts"; import Migration0075 from "./Migrations/075_SubagentDeliveryState.ts"; import Migration0076 from "./Migrations/076_ProjectionProjectsPinArchive.ts"; +import Migration0077 from "./Migrations/077_ProviderSessionBindingPayloads.ts"; const migrationEntries = [ [1, "OrchestrationEvents", Migration0001], @@ -158,6 +159,7 @@ const migrationEntries = [ [74, "ExcludeAttachmentPlaceholdersFromSearch", Migration0074], [75, "SubagentDeliveryState", Migration0075], [76, "ProjectionProjectsPinArchive", Migration0076], + [77, "ProviderSessionBindingPayloads", Migration0077], ] as const; export const migrationManifest = migrationEntries.map(([id, name]) => [id, name] as const); diff --git a/apps/server/src/persistence/Migrations/077_ProviderSessionBindingPayloads.ts b/apps/server/src/persistence/Migrations/077_ProviderSessionBindingPayloads.ts new file mode 100644 index 0000000000000..b67621c39f593 --- /dev/null +++ b/apps/server/src/persistence/Migrations/077_ProviderSessionBindingPayloads.ts @@ -0,0 +1,12 @@ +import * as Effect from "effect/Effect"; +import * as SqlClient from "effect/sql/SqlClient"; + +export default Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + // Reopening one thread must not revive other bindings of the same shared session. + yield* sql`ALTER TABLE orchestration_v2_projection_provider_session_bindings ADD COLUMN payload_json TEXT`; + yield* sql`UPDATE orchestration_v2_projection_provider_session_bindings + SET payload_json = (SELECT session.payload_json + FROM orchestration_v2_projection_provider_sessions AS session + WHERE session.provider_session_id = orchestration_v2_projection_provider_session_bindings.provider_session_id)`; +}); diff --git a/apps/server/src/persistence/reconcileV2PreviewMigration.test.ts b/apps/server/src/persistence/reconcileV2PreviewMigration.test.ts index 08c5094bf49c8..ff8dfa4d6b9e2 100644 --- a/apps/server/src/persistence/reconcileV2PreviewMigration.test.ts +++ b/apps/server/src/persistence/reconcileV2PreviewMigration.test.ts @@ -56,6 +56,7 @@ describe("V2 preview upgrade", () => { [74, "ExcludeAttachmentPlaceholdersFromSearch"], [75, "SubagentDeliveryState"], [76, "ProjectionProjectsPinArchive"], + [77, "ProviderSessionBindingPayloads"], ]); assert.deepStrictEqual(yield* runMigrations(), []); assert.deepStrictEqual(yield* sql`SELECT * FROM orchestration_v2_legacy_imports`, imports); @@ -155,6 +156,7 @@ describe("V2 preview upgrade", () => { [74, "ExcludeAttachmentPlaceholdersFromSearch"], [75, "SubagentDeliveryState"], [76, "ProjectionProjectsPinArchive"], + [77, "ProviderSessionBindingPayloads"], ]); }).pipe(Effect.provide(NodeSqliteClient.layer({ filename: ":memory:" }))), ); diff --git a/apps/server/src/storageCleanup.test.ts b/apps/server/src/storageCleanup.test.ts index 9604e67cfcf60..6045003a2cca2 100644 --- a/apps/server/src/storageCleanup.test.ts +++ b/apps/server/src/storageCleanup.test.ts @@ -22,6 +22,7 @@ import { import * as HostProcess from "@t3tools/shared/HostProcess"; import * as ServerConfig from "./config.ts"; import * as StorageCleanup from "./storageCleanup.ts"; +import * as SqlitePersistence from "./persistence/Sqlite.ts"; import * as GitVcsDriver from "./vcs/GitVcsDriver.ts"; import * as GitManager from "./git/GitManager.ts"; import * as Settings from "./serverSettings.ts"; @@ -315,6 +316,7 @@ const cleanupFixture = ( worktreeAfterDays: 8, worktreeOnDelete: true, }, + sessionSql?: SqlClient.SqlClient, ) => Effect.gen(function* () { const fs = yield* FileSystem.FileSystem; @@ -420,39 +422,45 @@ const cleanupFixture = ( Effect.provideService(Orchestrator.OrchestratorV2, { streamDomainEvents: Stream.suspend(() => events), } as unknown as Orchestrator.OrchestratorV2["Service"]), - Effect.provideService(SqlClient.SqlClient, ((strings: TemplateStringsArray) => - Effect.sync(() => - openSessionCwd !== null && strings.join("").includes("projection_provider_sessions") - ? [ - { - payload_json: JSON.stringify({ - id: "session-1", - driver: "codex", - providerInstanceId: "codex", - status: "ready", - cwd: openSessionCwd, - model: null, - capabilities: CodexProviderCapabilitiesV2, - createdAt: "2026-06-01T12:00:00.000Z", - updatedAt: "2026-06-01T12:00:00.000Z", - lastError: null, - }), - }, - ] - : outboxStatus !== null && strings.join("").includes("orchestration_v2_effect_outbox") - ? [{ status: outboxStatus }] - : strings.join("").includes("deleted_at IS NOT NULL") - ? deletedThreads.map((thread) => ({ - payload_json: JSON.stringify(thread), - workspaceRoot: repo, - })) - : strings.join("").includes('AS "worktreePath"') - ? [ - ...threads.flatMap((thread) => thread.worktreePath ?? []), - ...storedWorktreePaths, - ].map((worktreePath) => ({ worktreePath })) - : [], - )) as unknown as SqlClient.SqlClient), + Effect.provideService(SqlClient.SqlClient, (( + strings: TemplateStringsArray, + ...values: unknown[] + ) => + sessionSql !== undefined && strings.join("").includes("projection_provider_sessions") + ? sessionSql(strings, ...values) + : Effect.sync(() => + openSessionCwd !== null && strings.join("").includes("projection_provider_sessions") + ? [ + { + payload_json: JSON.stringify({ + id: "session-1", + driver: "codex", + providerInstanceId: "codex", + status: "ready", + cwd: openSessionCwd, + model: null, + capabilities: CodexProviderCapabilitiesV2, + createdAt: "2026-06-01T12:00:00.000Z", + updatedAt: "2026-06-01T12:00:00.000Z", + lastError: null, + }), + }, + ] + : outboxStatus !== null && + strings.join("").includes("orchestration_v2_effect_outbox") + ? [{ status: outboxStatus }] + : strings.join("").includes("deleted_at IS NOT NULL") + ? deletedThreads.map((thread) => ({ + payload_json: JSON.stringify(thread), + workspaceRoot: repo, + })) + : strings.join("").includes('AS "worktreePath"') + ? [ + ...threads.flatMap((thread) => thread.worktreePath ?? []), + ...storedWorktreePaths, + ].map((worktreePath) => ({ worktreePath })) + : [], + )) as unknown as SqlClient.SqlClient), Effect.provideService(GitManager.GitManager, { invalidateStatus: () => Effect.void, } as unknown as GitManager.GitManager["Service"]), @@ -1986,6 +1994,53 @@ describe("storage cleanup reports and local file policies", () => { }), ), ); + it.live.each(["bindings", "legacy", "orphan"] as const)( + "protects a live provider cwd using %s session state", + (source) => + runCleanupTest( + Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + const fixture = yield* cleanupFixture(undefined, sql); + const session = { + id: "session-1", + driver: "codex", + providerInstanceId: "codex", + status: "ready", + cwd: fixture.worktree, + model: null, + capabilities: CodexProviderCapabilitiesV2, + createdAt: "2026-06-01T12:00:00.000Z", + updatedAt: "2026-06-01T12:00:00.000Z", + lastError: null, + }; + const shared = { ...session, status: source === "bindings" ? "stopped" : "ready" }; + yield* sql`INSERT INTO orchestration_v2_projection_provider_sessions + (provider_session_id, thread_id, provider, status, updated_at, payload_json) + VALUES ('session-1', 'thread-2', 'codex', ${shared.status}, ${shared.updatedAt}, ${JSON.stringify(shared)})`; + if (source !== "orphan") { + yield* sql`INSERT INTO orchestration_v2_projection_provider_session_bindings + (provider_session_id, thread_id, payload_json) + VALUES ('session-1', 'thread-1', ${source === "legacy" ? null : JSON.stringify(session)})`; + } + if (source === "bindings") { + yield* sql`INSERT INTO orchestration_v2_projection_provider_session_bindings + (provider_session_id, thread_id, payload_json) + VALUES ('session-1', 'thread-2', ${JSON.stringify(shared)})`; + } + const kept = yield* fixture.service.runNow; + expect(kept.entries[0]).toMatchObject({ + outcome: "kept", + reason: "Provider session is still open", + bytes: null, + }); + expect(yield* fixture.fs.exists(fixture.worktree)).toBe(true); + fixture.setSessionLive(false); + const removed = yield* fixture.service.runNow; + expect(removed.entries[0]).toMatchObject({ outcome: "removed" }); + expect(yield* fixture.fs.exists(fixture.worktree)).toBe(false); + }).pipe(Effect.provide(SqlitePersistence.layerMemory)), + ), + ); it.live("does not measure a worktree kept for an open provider session", () => runCleanupTest( Effect.gen(function* () { diff --git a/apps/server/src/storageCleanup.ts b/apps/server/src/storageCleanup.ts index 13ad286861b9a..94574a7a62fc9 100644 --- a/apps/server/src/storageCleanup.ts +++ b/apps/server/src/storageCleanup.ts @@ -1018,8 +1018,11 @@ const make = Effect.gen(function* () { return waitForStop("thread activity or shared worktree changed since check"); // Sessions can outlive their run and can be shared across app threads. const sessionRows = yield* sql<{ payload_json: string }>` - SELECT payload_json FROM orchestration_v2_projection_provider_sessions - WHERE status != 'stopped' + SELECT DISTINCT COALESCE(binding.payload_json, session.payload_json) AS payload_json + FROM orchestration_v2_projection_provider_sessions AS session + LEFT JOIN orchestration_v2_projection_provider_session_bindings AS binding + ON binding.provider_session_id = session.provider_session_id + WHERE COALESCE(json_extract(binding.payload_json, '$.status'), session.status) != 'stopped' `.pipe(NodeSqliteClient.readOnly); const sessions = yield* Effect.forEach(sessionRows, (row) => decodeCleanupSession(row.payload_json),