diff --git a/apps/server/src/orchestration-v2/Adapters/PiAdapterV2.test.ts b/apps/server/src/orchestration-v2/Adapters/PiAdapterV2.test.ts index 06728c0410664c756d303c8a28107017c570c08c..6e385f1a1489113c98c52f992cda8199f0a98ceb 100644 --- a/apps/server/src/orchestration-v2/Adapters/PiAdapterV2.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/PiAdapterV2.test.ts @@ -1,6 +1,7 @@ import { assert, describe, it } from "@effect/vitest"; import * as NodeServices from "@effect/platform-node/NodeServices"; import { + CheckpointId, EnvironmentId, NodeId, ProviderInstanceId, @@ -30,8 +31,10 @@ import * as Sink from "effect/Sink"; import * as Stream from "effect/Stream"; import * as TestClock from "effect/testing/TestClock"; import { ChildProcess, ChildProcessSpawner } from "effect/process"; +import { HostProcessEnvironment } from "@t3tools/shared/hostProcess"; import * as ServerConfig from "../../config.ts"; +import * as ProviderContinuationRequests from "../ProviderContinuationRequests.ts"; import * as McpProviderSession from "../../mcp/McpProviderSession.ts"; import * as IdAllocator from "../IdAllocator.ts"; import { @@ -40,7 +43,12 @@ import { type ProviderAdapterV2SessionRuntime, } from "../ProviderAdapter.ts"; import { handoffBudget } from "../ContextHandoffBudget.ts"; -import { makePiAdapterV2, PI_PROVIDER } from "./PiAdapterV2.ts"; +import { + makePiAdapterV2, + PiAdapterV2Driver, + PI_PROVIDER, + type PiAdapterV2Options, +} from "./PiAdapterV2.ts"; import { makePiRpcConnection, type PiRpcRecord } from "./PiRpc.ts"; const layerServerConfig = ServerConfig.layerTest(process.cwd(), { @@ -88,7 +96,9 @@ interface FakePi { readonly resolveDeferredState: (data: unknown) => Effect.Effect; /** Reject the next `get_state` request. */ readonly failNextState: () => void; - readonly deferNextLifecycle: (type: "switch_session" | "new_session") => void; + readonly deferNextLifecycle: (type: "switch_session" | "new_session" | "fork") => void; + /** Hold a model-select extension hook until its UI request is answered. */ + readonly deferNextModelSelection: () => void; readonly queueModels: (models: ReadonlyArray) => void; readonly vetoNextNewSession: () => void; /** Every request received by the fake process. */ @@ -291,6 +301,9 @@ const makeFakePi: Effect.Effect = Effect.gen(function* () { deferNextLifecycle: (type) => { deferredLifecycle = type; }, + deferNextModelSelection: () => { + deferredLifecycle = "set_model"; + }, queueModels: (value) => { models = value; }, @@ -310,22 +323,44 @@ const makeFakePi: Effect.Effect = Effect.gen(function* () { } satisfies FakePi; }); -const makeAdapter = Effect.fnUntraced(function* (fake: FakePi, launchArgs = "", forkFake?: FakePi) { +const makeAdapter = Effect.fnUntraced(function* ( + fake: FakePi, + launchArgs = "", + forkFake?: FakePi, + continuationRequests?: PiAdapterV2Options["continuationRequests"], +) { const idAllocator = yield* IdAllocator.IdAllocatorV2; const serverConfig = yield* ServerConfig.ServerConfig; const fileSystem = yield* FileSystem.FileSystem; + const spawner = + forkFake === undefined + ? fake.spawner + : ChildProcessSpawner.make((command) => + ChildProcess.isStandardCommand(command) && command.args.includes("--fork") + ? forkFake.spawner.spawn(command) + : fake.spawner.spawn(command), + ); + if (continuationRequests !== undefined) { + return yield* PiAdapterV2Driver.create({ + instanceId: PI_INSTANCE_ID, + displayName: undefined, + enabled: true, + environment: [], + config: { enabled: true, binaryPath: "pi", launchArgs, customModels: [] }, + }).pipe( + Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, spawner), + Effect.provideService(HostProcessEnvironment, {}), + Effect.provideService(ProviderContinuationRequests.ProviderContinuationRequests, { + ...continuationRequests, + take: Effect.never, + }), + ); + } return makePiAdapterV2({ instanceId: PI_INSTANCE_ID, settings: { enabled: true, binaryPath: "pi", launchArgs, customModels: [] }, environment: {}, - spawner: - forkFake === undefined - ? fake.spawner - : ChildProcessSpawner.make((command) => - ChildProcess.isStandardCommand(command) && command.args.includes("--fork") - ? forkFake.spawner.spawn(command) - : fake.spawner.spawn(command), - ), + spawner, fileSystem, idAllocator, serverConfig, @@ -338,8 +373,9 @@ const openRuntime = Effect.fnUntraced(function* ( threadId = THREAD_ID, providerSessionId = SESSION_ID, forkFake?: FakePi, + continuationRequests?: PiAdapterV2Options["continuationRequests"], ) { - const adapter = yield* makeAdapter(fake, "", forkFake); + const adapter = yield* makeAdapter(fake, "", forkFake, continuationRequests); const runtime = yield* adapter.openSession({ threadId, providerSessionId, @@ -397,6 +433,7 @@ const startTurn = Effect.fnUntraced(function* ( selection?: ModelSelection, runOrdinal = 1, threadId = THREAD_ID, + continuation = false, ) { const appThread = yield* makeAppThread(model, threadId); const runId = RunId.make(`run:${threadId}:${runOrdinal}`); @@ -413,8 +450,8 @@ const startTurn = Effect.fnUntraced(function* ( messageId: `message:${threadId}:${runOrdinal}` as never, text, attachments, - createdBy: "user", - creationSource: "web", + createdBy: continuation ? "agent" : "user", + creationSource: continuation ? "provider" : "web", }, modelSelection: selection ?? modelSelection(model), runtimePolicy, @@ -461,6 +498,786 @@ const expectModelFailure = (errorMessage: string) => }).pipe(Effect.scoped, Effect.provide(layerTest)); describe("PiAdapterV2", () => { + it.effect.each([false, true])( + "serializes rollback fork hooks with native wake ownership, wake=%s", + (nativeWake) => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { + offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid), + }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + fake.deferNextLifecycle("fork"); + const rollingBack = yield* runtime + .rollbackThread({ + providerThread, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("checkpoint-start"), + appRunOrdinal: 0, + }, + providerThreadTurns: [ + { + id: ProviderTurnId.make("previous-turn"), + providerThreadId: providerThread.id, + nodeId: NodeId.make("previous-node"), + runAttemptId: null, + nativeTurnRef: { + driver: PI_PROVIDER, + nativeId: "user-before-wake", + strength: "strong", + }, + ordinal: 1, + status: "completed", + startedAt: providerThread.createdAt, + completedAt: providerThread.createdAt, + }, + ], + }) + .pipe(Effect.result, Effect.forkScoped); + const fork = yield* fake.takeRequest("fork"); + const overlappingTurn = yield* startTurn(runtime, providerThread).pipe(Effect.result); + assert.equal(overlappingTurn._tag, "Failure"); + const overlappingRegistration = yield* runtime + .ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }) + .pipe(Effect.result); + assert.equal(overlappingRegistration._tag, "Failure"); + yield* fake.emit({ + type: "extension_ui_request", + id: "fork-hook-dialog", + method: "confirm", + title: "Restore this branch?", + }); + const forkDialog = yield* takeEvent((event) => event.type === "runtime_request.updated"); + assert.isTrue(forkDialog.type === "runtime_request.updated"); + if (forkDialog.type !== "runtime_request.updated") return; + assert.isNull(forkDialog.runtimeRequest.providerTurnId); + yield* runtime.respondToRuntimeRequest({ + requestId: forkDialog.runtimeRequest.id, + decision: "accept", + }); + yield* fake.takeRequest("extension_ui_response"); + if (nativeWake) { + yield* fake.emit({ type: "agent_start" }); + yield* fake.emit({ type: "message_start", message: { role: "assistant" } }); + yield* fake.emit({ + type: "message_update", + assistantMessageEvent: { + type: "text_end", + contentIndex: 0, + content: "Discarded branch wake output", + }, + }); + yield* fake.emit({ type: "agent_settled" }); + const retired = yield* takeEvent( + (event) => + event.type === "provider_session.updated" && + event.providerSession.status === "stopped", + ).pipe(Effect.as(true), Effect.raceFirst(Queue.take(offers).pipe(Effect.as(false)))); + assert.isTrue(retired, "Pi offered native work while rollback was changing its binding"); + const result = yield* Fiber.join(rollingBack); + assert.equal(result._tag, "Failure"); + assert.isFalse(yield* runtime.hasPendingBackgroundWork!); + const staleTurn = yield* startTurn( + runtime, + providerThread, + "default", + [], + "Synthetic wake", + undefined, + 1, + THREAD_ID, + true, + ).pipe(Effect.result); + assert.equal(staleTurn._tag, "Failure"); + assert.isFalse(fake.allRequests().some((request) => request["type"] === "prompt")); + return; + } + fake.queueState({ sessionFile: "/fake/rolled-back-session.jsonl" }); + yield* fake.emit({ + type: "response", + id: fork["id"], + command: "fork", + success: true, + data: { cancelled: false }, + }); + const result = yield* Fiber.join(rollingBack); + assert.equal(result._tag, "Success"); + if (result._tag !== "Success") return; + assert.equal( + result.success.providerThread.nativeThreadRef?.nativeId, + "/fake/rolled-back-session.jsonl", + ); + yield* startTurn(runtime, result.success.providerThread); + const prompt = yield* fake.takeRequest("prompt"); + assert.equal(prompt["message"], "Hello pi"); + yield* fake.emit({ type: "response", command: "prompt", success: true }); + yield* fake.emit({ type: "agent_start" }); + yield* fake.emit({ type: "agent_settled" }); + const terminal = yield* takeEvent((event) => event.type === "turn.terminal"); + assert.isTrue(terminal.type === "turn.terminal" && terminal.status === "completed"); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect.each([false, true])( + "retains native wake ownership after a rejected joining prompt, compaction aborted=%s", + (aborted) => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { + offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid), + }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* fake.takeRequest("get_state"); + yield* fake.emit({ type: "agent_start" }); + yield* Queue.take(offers); + yield* fake.emit({ type: "agent_settled" }); + yield* fake.emit({ type: "compaction_start", reason: "manual" }); + yield* startTurn(runtime, providerThread); + yield* fake.takeRequest("prompt"); + const rejection = "Cannot submit a prompt while compaction is in progress."; + yield* fake.emit({ type: "response", command: "prompt", success: false, error: rejection }); + yield* fake.emit({ + type: "extension_ui_request", + id: "native-compaction-dialog", + method: "confirm", + title: "Continue native work?", + }); + const dialog = yield* takeEvent((event) => event.type === "runtime_request.updated"); + assert.isTrue(dialog.type === "runtime_request.updated"); + if (dialog.type !== "runtime_request.updated") return; + assert.isNotNull(dialog.runtimeRequest.providerTurnId); + yield* runtime.respondToRuntimeRequest({ + requestId: dialog.runtimeRequest.id, + decision: "accept", + }); + yield* fake.takeRequest("extension_ui_response"); + if (aborted) { + yield* fake.emit({ type: "compaction_end", result: null, aborted: true }); + const terminal = yield* takeEvent((event) => event.type === "turn.terminal"); + assert.isTrue( + terminal.type === "turn.terminal" && + terminal.status === "failed" && + terminal.failure.message === rejection, + ); + return; + } + fake.queueState({ isStreaming: true }); + yield* fake.emit({ + type: "compaction_end", + result: { summary: "Native compaction finished" }, + willRetry: true, + }); + yield* fake.emit({ type: "agent_start" }); + yield* fake.emit({ type: "message_start", message: { role: "assistant" } }); + yield* fake.emit({ + type: "message_update", + assistantMessageEvent: { + type: "text_end", + contentIndex: 0, + content: "Native work completed", + }, + }); + const message = yield* takeEvent((event) => event.type === "message.updated"); + assert.isTrue(message.type === "message.updated" && message.message.runId !== null); + yield* fake.takeRequest("get_state"); + yield* fake.emit({ type: "agent_settled" }); + const terminal = yield* takeEvent((event) => event.type === "turn.terminal"); + assert.isTrue( + terminal.type === "turn.terminal" && + terminal.status === "failed" && + terminal.failure.message === rejection, + ); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect("rejects rollback while native wake output is awaiting ownership", () => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { + offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid), + }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* fake.emit({ type: "agent_start" }); + yield* Queue.take(offers); + yield* fake.emit({ type: "message_start", message: { role: "assistant" } }); + yield* fake.emit({ + type: "message_update", + assistantMessageEvent: { + type: "text_end", + contentIndex: 0, + content: "Current branch wake output", + }, + }); + yield* fake.emit({ type: "agent_settled" }); + yield* takeEvent( + (event) => + event.type === "provider_session.updated" && event.providerSession.status === "ready", + ); + const rolledBack = yield* runtime + .rollbackThread({ + providerThread, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("checkpoint-start"), + appRunOrdinal: 0, + }, + providerThreadTurns: [ + { + id: ProviderTurnId.make("previous-turn"), + providerThreadId: providerThread.id, + nodeId: NodeId.make("previous-node"), + runAttemptId: null, + nativeTurnRef: { + driver: PI_PROVIDER, + nativeId: "user-before-wake", + strength: "strong", + }, + ordinal: 1, + status: "completed", + startedAt: providerThread.createdAt, + completedAt: providerThread.createdAt, + }, + ], + }) + .pipe(Effect.result); + assert.equal( + rolledBack._tag, + "Failure", + "Rollback changed the native session while its wake was still buffered", + ); + assert.isFalse(fake.allRequests().some((request) => request["type"] === "fork")); + assert.isTrue(yield* runtime.hasPendingBackgroundWork!); + yield* startTurn( + runtime, + providerThread, + "default", + [], + "Synthetic wake", + undefined, + 1, + THREAD_ID, + true, + ); + const message = yield* takeEvent((event) => event.type === "message.updated"); + assert.equal( + message.type === "message.updated" && message.message.text, + "Current branch wake output", + ); + yield* takeEvent((event) => event.type === "turn.terminal"); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect("answers model-selection dialogs before a user turn claims a native wake", () => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid) }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* fake.emit({ type: "agent_start" }); + yield* Queue.take(offers); + const requests: Array> = + []; + fake.deferNextModelSelection(); + const starting = yield* startTurn( + runtime, + providerThread, + "default", + [], + "Continue the review", + modelSelection("xai/grok-4.6"), + ).pipe(Effect.forkScoped); + const modelRequest = yield* fake.takeRequest("set_model"); + yield* fake.emit({ + type: "extension_ui_request", + id: "model-selection-dialog", + method: "confirm", + title: "Use this model?", + }); + yield* fake.emit({ type: "agent_settled" }); + yield* takeEvent((event) => { + if (event.type === "runtime_request.updated") requests.push(event); + return ( + event.type === "provider_session.updated" && event.providerSession.status === "ready" + ); + }); + const dialog = requests[0]?.runtimeRequest; + assert.isDefined(dialog, "Pi buffered the dialog needed to finish model selection"); + assert.equal(requests[0]?.threadId, THREAD_ID); + assert.isNull(dialog!.providerTurnId); + yield* runtime.respondToRuntimeRequest({ requestId: dialog!.id, decision: "accept" }); + const response = yield* fake.takeRequest("extension_ui_response"); + assert.equal(response["id"], "model-selection-dialog"); + assert.isTrue(response["confirmed"]); + yield* fake.emit({ + type: "response", + id: modelRequest["id"], + command: "set_model", + success: true, + data: { provider: "xai", id: "grok-4.6" }, + }); + yield* Fiber.join(starting); + const prompt = yield* fake.takeRequest("prompt"); + assert.equal(prompt["message"], "Continue the review"); + assert.equal(prompt["streamingBehavior"], "steer"); + const resolved = yield* takeEvent( + (event) => + event.type === "runtime_request.updated" && event.runtimeRequest.id === dialog!.id, + ); + assert.isTrue( + resolved.type === "runtime_request.updated" && + resolved.runtimeRequest.status === "resolved", + ); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect("adopts a completed native wake without sending a synthetic prompt", () => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { + offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid), + }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* fake.emit({ type: "agent_start" }); + const offer = yield* Queue.take(offers).pipe( + Effect.raceFirst( + takeEvent( + (event) => + event.type === "provider_session.updated" && event.providerSession.status === "error", + ).pipe( + Effect.andThen(Effect.sync(() => assert.fail("Pi killed a routable native wake"))), + ), + ), + ); + assert.equal(offer.threadId, THREAD_ID); + assert.equal(offer.providerThreadId, providerThread.id); + assert.equal(offer.delivery, "adapter_buffered"); + assert.isTrue(yield* runtime.hasPendingBackgroundWork!); + let dispatchCommitted = false; + yield* offer.dispatchIfCurrent!( + Effect.sync(() => { + dispatchCommitted = true; + }), + ); + assert.isTrue(dispatchCommitted); + assert.isTrue(yield* runtime.hasPendingBackgroundWork!); + yield* fake.emit({ type: "message_start", message: { role: "assistant" } }); + yield* fake.emit({ + type: "message_update", + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "Background review finished", + }, + }); + yield* fake.emit({ type: "message_end", message: { role: "assistant" } }); + yield* fake.emit({ type: "agent_settled" }); + yield* takeEvent( + (event) => + event.type === "provider_session.updated" && event.providerSession.status === "ready", + ); + assert.isTrue(yield* runtime.hasPendingBackgroundWork!); + yield* startTurn( + runtime, + providerThread, + "default", + [], + "Synthetic wake", + undefined, + 1, + THREAD_ID, + true, + ); + const message = yield* takeEvent((event) => event.type === "message.updated"); + assert.equal( + message.type === "message.updated" && message.message.text, + "Background review finished", + ); + yield* takeEvent((event) => event.type === "turn.terminal"); + assert.equal(fake.allRequests().filter((request) => request["type"] === "prompt").length, 0); + assert.isFalse(yield* runtime.hasPendingBackgroundWork!); + let dispatched = false; + yield* offer.dispatchIfCurrent!( + Effect.sync(() => { + dispatched = true; + }), + ); + assert.isFalse(dispatched); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect.each([false, true])( + "a user prompt takes the wake without premature settlement, native wake settled=%s", + (settled) => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { + offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid), + }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* fake.emit({ type: "agent_start" }); + const offer = yield* Queue.take(offers); + if (settled) { + yield* fake.emit({ type: "agent_settled" }); + yield* takeEvent( + (event) => + event.type === "provider_session.updated" && event.providerSession.status === "ready", + ); + } + yield* startTurn(runtime, providerThread); + const prompt = yield* fake.takeRequest("prompt"); + assert.equal(prompt["message"], "Hello pi"); + assert.equal(prompt["streamingBehavior"], "steer"); + yield* fake.emit({ type: "response", command: "prompt", success: true }); + yield* fake.emit({ type: "agent_start" }); + yield* fake.emit({ + type: "extension_ui_request", + id: "owned-dialog", + method: "confirm", + title: "Continue?", + }); + const dialog = yield* takeEvent((event) => event.type === "runtime_request.updated"); + assert.isTrue(dialog.type === "runtime_request.updated"); + if (dialog.type !== "runtime_request.updated") return; + assert.isNotNull(dialog.runtimeRequest.providerTurnId); + yield* offer.clearIfCurrent!(); + yield* runtime.respondToRuntimeRequest({ + requestId: dialog.runtimeRequest.id, + decision: "accept", + }); + const response = yield* fake.takeRequest("extension_ui_response"); + assert.isTrue(response["confirmed"]); + yield* fake.emit({ type: "message_start", message: { role: "assistant" } }); + yield* fake.emit({ + type: "message_update", + assistantMessageEvent: { + type: "text_end", + contentIndex: 0, + content: "User prompt completed", + }, + }); + yield* fake.emit({ type: "agent_settled" }); + const message = yield* takeEvent( + (event) => + event.type === "message.updated" && event.message.text === "User prompt completed", + ); + assert.isTrue(message.type === "message.updated" && message.message.runId !== null); + yield* takeEvent((event) => event.type === "turn.terminal"); + let dispatched = false; + yield* offer.dispatchIfCurrent!( + Effect.sync(() => { + dispatched = true; + }), + ); + assert.isFalse(dispatched); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect("defers a wake's late settlement until the joined user prompt is acknowledged", () => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { + offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid), + }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* fake.emit({ type: "agent_start" }); + yield* Queue.take(offers); + yield* startTurn(runtime, providerThread); + yield* fake.takeRequest("prompt"); + yield* fake.emit({ type: "agent_settled" }); + yield* fake.emit({ + type: "extension_ui_request", + method: "notify", + message: "Ordering barrier", + }); + yield* takeEvent( + (event) => + event.type === "turn_item.updated" && + event.turnItem.type === "dynamic_tool" && + event.turnItem.toolName === "notify", + ); + assert.equal(runtime.providerSession.status, "running"); + fake.queueState({ isStreaming: true }); + yield* fake.emit({ type: "response", command: "prompt", success: true }); + yield* fake.takeRequest("get_state"); + yield* fake.emit({ type: "agent_start" }); + yield* fake.emit({ type: "agent_settled" }); + yield* takeEvent((event) => event.type === "turn.terminal"); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect("settles a command-only user prompt after taking a completed wake", () => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { + offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid), + }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* fake.emit({ type: "agent_start" }); + yield* Queue.take(offers); + yield* fake.emit({ type: "agent_settled" }); + yield* takeEvent( + (event) => + event.type === "provider_session.updated" && event.providerSession.status === "ready", + ); + yield* startTurn(runtime, providerThread, "default", [], "/hello"); + yield* fake.takeRequest("prompt"); + yield* fake.emit({ type: "response", command: "prompt", success: true }); + yield* takeEvent((event) => event.type === "turn.terminal"); + yield* startTurn( + runtime, + providerThread, + "default", + [], + "Stale continuation", + undefined, + 2, + THREAD_ID, + true, + ); + yield* takeEvent((event) => event.type === "turn.terminal"); + assert.equal(fake.allRequests().filter((request) => request["type"] === "prompt").length, 1); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect("keeps a pending wake on its thread and lets Stop retire its process", () => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { + offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid), + }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* fake.emit({ type: "agent_start" }); + const offer = yield* Queue.take(offers); + assert.isTrue(yield* runtime.hasPendingBackgroundWorkForThread!(providerThread)); + const otherThread = { + ...providerThread, + appThreadId: ThreadId.make("other-thread"), + id: ProviderThreadId.make("other-provider-thread"), + }; + assert.isFalse(yield* runtime.hasPendingBackgroundWorkForThread!(otherThread)); + const failure = yield* runtime + .resumeThread({ providerThread: otherThread }) + .pipe(Effect.flip); + assert.equal(failure._tag, "ProviderAdapterResumeThreadError"); + yield* runtime.interruptTurn({ + providerThread, + providerTurnId: ProviderTurnId.make("settled-turn"), + requestRuntimeRestart: true, + }); + yield* takeEvent( + (event) => + event.type === "provider_session.updated" && event.providerSession.status === "stopped", + ); + assert.isFalse(yield* runtime.hasPendingBackgroundWork!); + let dispatched = false; + yield* offer.dispatchIfCurrent!( + Effect.sync(() => { + dispatched = true; + }), + ); + assert.isFalse(dispatched); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect.each(["archive", "stop", "transport death"] as const)( + "retiring a wake by %s cancels visible dialogs and invalidates its offer", + (retirement) => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const offers = + yield* Queue.unbounded(); + const { runtime, takeEvent } = yield* openRuntime( + fake, + "default", + THREAD_ID, + SESSION_ID, + undefined, + { + offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid), + }, + ); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* fake.emit({ type: "agent_start" }); + const offer = yield* Queue.take(offers); + yield* fake.emit({ + type: "extension_ui_request", + id: "wake-dialog", + method: "confirm", + title: "Continue background work?", + }); + const dialog = yield* takeEvent((event) => event.type === "runtime_request.updated"); + assert.isTrue(dialog.type === "runtime_request.updated"); + if (dialog.type !== "runtime_request.updated") return; + assert.isNull(dialog.runtimeRequest.providerTurnId); + if (retirement === "archive") yield* offer.clearIfCurrent!(); + else if (retirement === "stop") { + yield* runtime.interruptTurn({ + providerThread, + providerTurnId: ProviderTurnId.make("settled-turn"), + requestRuntimeRestart: true, + }); + } else yield* fake.closeStdout; + let requestCancelled = false; + let nodeCancelled = false; + let itemCancelled = false; + yield* takeEvent((event) => { + if ( + event.type === "runtime_request.updated" && + event.runtimeRequest.id === dialog.runtimeRequest.id + ) + requestCancelled = event.runtimeRequest.status === "cancelled"; + if (event.type === "node.updated" && event.node.id === dialog.runtimeRequest.nodeId) + nodeCancelled = event.node.status === "cancelled"; + if ( + event.type === "turn_item.updated" && + event.turnItem.nodeId === dialog.runtimeRequest.nodeId + ) + itemCancelled = event.turnItem.status === "cancelled"; + return ( + event.type === "provider_session.updated" && + event.providerSession.status === + (retirement === "transport death" ? "error" : "stopped") + ); + }); + assert.isTrue(requestCancelled, "Retired wake left a pending runtime request"); + assert.isTrue(nodeCancelled); + assert.isTrue(itemCancelled); + assert.isFalse(yield* runtime.hasPendingBackgroundWork!); + let dispatched = false; + yield* offer.dispatchIfCurrent!( + Effect.sync(() => { + dispatched = true; + }), + ); + assert.isFalse(dispatched); + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + it.effect("stops provider-initiated work that has no T3 turn owner", () => Effect.gen(function* () { const fake = yield* makeFakePi; diff --git a/apps/server/src/orchestration-v2/Adapters/PiAdapterV2.ts b/apps/server/src/orchestration-v2/Adapters/PiAdapterV2.ts index 9270253d5c74bc01a0f59ba42d1b9290bd0c8959..a319b4b2f0b7e8681809aec6d6ed0d50fe2c1108 100644 --- a/apps/server/src/orchestration-v2/Adapters/PiAdapterV2.ts +++ b/apps/server/src/orchestration-v2/Adapters/PiAdapterV2.ts @@ -48,6 +48,7 @@ import { } from "@t3tools/contracts"; import * as Cause from "effect/Cause"; import * as DateTime from "effect/DateTime"; +import * as Deferred from "effect/Deferred"; import * as Option from "effect/Option"; import * as Duration from "effect/Duration"; import * as Effect from "effect/Effect"; @@ -72,6 +73,7 @@ import { import { mergeProviderInstanceEnvironment } from "../../provider/ProviderInstanceEnvironment.ts"; import * as IdAllocator from "../IdAllocator.ts"; import * as ProviderAdapter from "../ProviderAdapter.ts"; +import * as ProviderContinuationRequests from "../ProviderContinuationRequests.ts"; import { ProviderAdapterDriverCreateError, type ProviderAdapterDriver, @@ -228,6 +230,11 @@ export interface PiAdapterV2Options { readonly fileSystem: FileSystem.FileSystem; readonly idAllocator: IdAllocator.IdAllocatorV2["Service"]; readonly serverConfig: ServerConfig.ServerConfig["Service"]; + readonly continuationRequests?: { + readonly offer: ( + request: ProviderContinuationRequests.ProviderContinuationRequest, + ) => Effect.Effect; + }; } /** Concatenate the `text` fields of a Pi content-block array. */ @@ -313,6 +320,8 @@ interface ActivePiTurn { * the turn instead. */ sawAgentActivity: boolean; + /** A user prompt joined provider-native work whose settlement may already be queued. */ + adoptedWake: boolean; /** Only slash-command prompts can complete without starting an agent run. */ readonly promptMayBeCommandOnly: boolean; /** Pi reports context as unknown immediately after compaction; keep its estimate for the meter. */ @@ -329,6 +338,8 @@ interface ActivePiTurn { activeCompaction: PiCompactionState | null; activeProviderRetry: PiProviderRetryState | null; failure: ReturnType | null; + /** Native retry can recover its work without accepting a rejected joining prompt. */ + rejectedPromptFailure: ReturnType | null; /** Session-tree refs read just before Stop terminates Pi, when no read is possible later. */ stopTreeRefs?: PiTurnTreeRefs | null; } @@ -366,6 +377,11 @@ interface PiThreadState { activeTurn: ActivePiTurn | null; } +interface PiWake { + readonly state: PiThreadState; + readonly events: Array; +} + // ── adapter ─────────────────────────────────────────────────── export function makePiAdapterV2( @@ -476,10 +492,26 @@ export function makePiAdapterV2( // Keep that intent beyond turn finalization so the later stdout close is // not mistaken for an unexpected transport failure. let stopRequested = false; - // Pi extensions can trigger an agent turn after the owning T3 turn has - // settled. Until orchestration has a first-class provider-initiated run, - // stop that runtime before it can execute tools without a timeline owner. + // Native work is held until an orchestrator-owned continuation or user + // turn attaches. The guard remains for sessions with no routable owner. let unsolicitedActivityDetected = false; + let pendingWake: PiWake | null = null; + let rollbackBarrier: Deferred.Deferred< + void, + ProviderAdapter.ProviderAdapterProtocolError + > | null = null; + let closed = false; + yield* Effect.addFinalizer(() => + Effect.gen(function* () { + closed = true; + pendingWake = null; + if (rollbackBarrier !== null) + yield* Deferred.fail( + rollbackBarrier, + protocolError("Pi session closed during rollback"), + ); + }), + ); let appliedModel: string | null = null; let appliedThinking: string | null = null; /** Last thread title synced into pi's session name (`/resume` listing). */ @@ -1437,7 +1469,7 @@ export function makePiAdapterV2( const tokenUsage = readUsage ? yield* readTokenUsage(turn.latestCompactionAfterTokens, completedAt) : undefined; - const failure = turn.interrupted ? null : turn.failure; + const failure = turn.interrupted ? null : (turn.rejectedPromptFailure ?? turn.failure); yield* emit({ type: "provider_turn.updated", driver: PI_PROVIDER, @@ -1554,11 +1586,74 @@ export function makePiAdapterV2( ); }; + const stopPendingWake = Effect.fnUntraced(function* (wake: PiWake) { + if (pendingWake !== wake || threadState !== wake.state || closed) return; + pendingWake = null; + stopRequested = true; + yield* cancelPendingPrompts(yield* DateTime.now); + yield* connection.terminate; + }); + const handleSessionEvent = Effect.fnUntraced(function* (event: PiRpcRecord) { + // Model-selection hooks can await a dialog before startTurn can claim + // a wake. Keep interactive requests session-scoped until a turn owns + // them; buffer native output and notifications for that turn. + const interactiveUiRequest = + event["type"] === "extension_ui_request" && + ["select", "confirm", "input", "editor"].includes(String(event["method"])); + if ( + threadState?.activeTurn === null && + pendingWake !== null && + !interactiveUiRequest && + !String(event["type"]).startsWith("t3.") && + event["type"] !== "response" + ) { + pendingWake.events.push(event); + if (event["type"] === "agent_start") yield* updateProviderSession("running", null); + if (event["type"] === "agent_settled") yield* updateProviderSession("ready", null); + return; + } const state = threadState; const turn = state?.activeTurn ?? null; switch (event["type"]) { case "agent_start": { + if ( + turn === null && + state?.providerThread.appThreadId != null && + options.continuationRequests !== undefined && + !closed && + !stopRequested + ) { + const wake: PiWake = { state, events: [event] }; + pendingWake = wake; + if (rollbackBarrier !== null) { + yield* stopPendingWake(wake); + return; + } + yield* updateProviderSession("running", null); + yield* options.continuationRequests + .offer({ + threadId: state.providerThread.appThreadId, + providerThreadId: state.providerThread.id, + driver: PI_PROVIDER, + detail: null, + delivery: "adapter_buffered", + dispatchIfCurrent: (dispatch) => + sessionEventPermit.withPermits(1)( + Effect.suspend(() => + pendingWake === wake && threadState === state && !closed + ? dispatch.pipe( + Effect.map(Option.some), + Effect.tapError(() => stopPendingWake(wake)), + ) + : Effect.succeed(Option.none()), + ), + ), + clearIfCurrent: () => sessionEventPermit.withPermits(1)(stopPendingWake(wake)), + }) + .pipe(Effect.tapError(() => stopPendingWake(wake))); + return; + } if (turn === null) { unsolicitedActivityDetected = true; yield* updateProviderSession("error", PI_UNSOLICITED_ACTIVITY_ERROR); @@ -1787,6 +1882,17 @@ export function makePiAdapterV2( if (turn !== null) { turn.settleWhenIdle = true; turn.settleProbeGeneration += 1; + // A wake can settle while the newly joined user prompt is still + // expanding. Its acknowledgement opens the idle-probe barrier. + if ( + turn.adoptedWake && + pendingPromptResponses.some( + (pending) => + pending.providerTurnId === turn.providerTurn.id && + pending.kind === "turn_start", + ) + ) + return; yield* scheduleSettleProbe(turn, true); } return; @@ -1845,10 +1951,11 @@ export function makePiAdapterV2( if ( pendingPrompt?.kind === "turn_start" && responseTurn !== null && - responseTurn.promptMayBeCommandOnly && - !responseTurn.sawAgentActivity + ((responseTurn.promptMayBeCommandOnly && !responseTurn.sawAgentActivity) || + (responseTurn.adoptedWake && + (responseTurn.settleWhenIdle || responseTurn.promptMayBeCommandOnly))) ) { - yield* scheduleSettleProbe(responseTurn); + yield* scheduleSettleProbe(responseTurn, responseTurn.adoptedWake); } return; } @@ -1869,11 +1976,18 @@ export function makePiAdapterV2( ? turn : null; if (failedTurn !== null) { - failedTurn.failure = makeProviderFailure({ + const failure = makeProviderFailure({ message: recordString(event, "error") ?? "Pi rejected the prompt.", class: "provider_error", }); - if (state !== null) yield* finalizeTurn(state); + if (failedTurn.adoptedWake && command === "prompt") { + failedTurn.rejectedPromptFailure = failure; + failedTurn.settleWhenIdle = true; + yield* scheduleSettleProbe(failedTurn, true); + } else { + failedTurn.failure = failure; + if (state !== null) yield* finalizeTurn(state); + } } return; } @@ -1886,6 +2000,11 @@ export function makePiAdapterV2( } return; } + case "t3.rollback_barrier": { + if (rollbackBarrier !== null && event["barrier"] === rollbackBarrier) + yield* Deferred.succeed(rollbackBarrier, undefined); + return; + } case "t3.settle_probe": { // New work increments the generation before the pump can consume // a stale idle snapshot, so only a current snapshot may settle. @@ -1946,6 +2065,13 @@ export function makePiAdapterV2( // closes the provider stream cleanly; only an unexpected death // is surfaced as an event-stream failure. const state = threadState; + if (rollbackBarrier !== null) + yield* Deferred.fail( + rollbackBarrier, + protocolError("Pi transport closed during rollback"), + ); + const hadPendingWake = pendingWake !== null; + pendingWake = null; const interrupted = state?.activeTurn?.interrupted === true; if (state?.activeTurn != null) { state.activeTurn.failure = interrupted @@ -1956,6 +2082,8 @@ export function makePiAdapterV2( class: "transport_error", }); yield* finalizeTurn(state, false); + } else if (hadPendingWake) { + yield* cancelPendingPrompts(yield* DateTime.now); } if (unsolicitedActivityDetected) { yield* updateProviderSession("error", PI_UNSOLICITED_ACTIVITY_ERROR); @@ -1999,9 +2127,26 @@ export function makePiAdapterV2( threadInput: ProviderAdapter.ProviderAdapterV2EnsureThreadInput, publish = true, ) { + if (rollbackBarrier !== null) + return yield* protocolError("Cannot register a Pi thread during rollback"); if (threadState !== null && threadState.activeTurn !== null) { return yield* protocolError("Cannot register a Pi thread while a turn is active"); } + if (pendingWake !== null) { + const state = threadState; + if ( + state !== null && + state === pendingWake.state && + state.providerThread.appThreadId === threadInput.threadId && + (threadInput.existingProviderThread === undefined || + threadInput.existingProviderThread.nativeThreadRef?.nativeId === + state.providerThread.nativeThreadRef?.nativeId) + ) + return state.providerThread; + return yield* protocolError( + "Cannot switch Pi threads while native work is awaiting a turn", + ); + } const existing = threadInput.existingProviderThread; const resumeId = existing?.nativeThreadRef?.nativeId; const needsNewSession = resumeId == null && registrationAttempted; @@ -2219,6 +2364,11 @@ export function makePiAdapterV2( return sessionEntity; }, events: Stream.fromQueue(events), + hasPendingBackgroundWork: Effect.sync(() => pendingWake !== null), + hasPendingBackgroundWorkForThread: (providerThread) => + Effect.sync( + () => pendingWake !== null && pendingWake.state.providerThread.id === providerThread.id, + ), getModelContextWindow: (selection) => { if (selection.instanceId !== options.instanceId) return undefined; const slug = @@ -2267,6 +2417,8 @@ export function makePiAdapterV2( }), startTurn: (turnInput) => Effect.gen(function* () { + if (rollbackBarrier !== null) + return yield* protocolError("Cannot start a Pi turn during rollback"); const state = threadState; if (state === null) { return yield* protocolError("Pi session has no registered thread"); @@ -2284,8 +2436,10 @@ export function makePiAdapterV2( } // The orchestrator adopts a fork under its already-allocated row. // Future session updates must retain that authoritative identity. - state.providerThread = turnInput.providerThread; - yield* applySelection(turnInput.modelSelection); + const continuation = + turnInput.message.createdBy === "agent" && + turnInput.message.creationSource === "provider"; + if (!continuation) yield* applySelection(turnInput.modelSelection); // Mirror the thread title into pi's session name so the session // stays identifiable in pi's own /resume listing. Best-effort: // naming must never block a turn. @@ -2307,9 +2461,11 @@ export function makePiAdapterV2( // extension's before_agent_start system-prompt hook, never by // wrapping the user text: a wrapped first message would no // longer start with "/" and slash commands would stop expanding. - const compactCommand = parsePiCompactCommand(turnInput.message.text); + const compactCommand = continuation + ? null + : parsePiCompactCommand(turnInput.message.text); const payload = - compactCommand === null + !continuation && compactCommand === null ? yield* resolvePromptPayload(turnInput.message.text, turnInput.message.attachments) : null; const startedAt = yield* DateTime.now; @@ -2340,6 +2496,7 @@ export function makePiAdapterV2( toolStartedAt: new Map(), interrupted: false, sawAgentActivity: false, + adoptedWake: false, promptMayBeCommandOnly: compactCommand !== null || (payload?.message.trimStart().startsWith("/") ?? false), latestCompactionAfterTokens: null, @@ -2351,13 +2508,32 @@ export function makePiAdapterV2( activeCompaction: null, activeProviderRetry: null, failure: null, + rejectedPromptFailure: null, }; // Only the install/send/start-event boundary excludes the event // pump. Earlier correlated requests must leave the pump free so // project trust, login, and session-switch dialogs can be shown // and answered instead of deadlocking the caller. yield* Effect.gen(function* () { + if ( + threadState !== state || + state.activeTurn !== null || + closed || + stopRequested || + rollbackBarrier !== null || + state.providerThread.nativeThreadRef?.nativeId !== + turnInput.providerThread.nativeThreadRef?.nativeId + ) { + return yield* protocolError("Pi turn ownership changed before it could start"); + } + const wake = pendingWake?.state === state ? pendingWake : null; + if (wake !== null && wake.state.providerThread.appThreadId !== turnInput.threadId) { + return yield* protocolError("Pi native wake belongs to a different app thread"); + } + pendingWake = null; + state.providerThread = turnInput.providerThread; state.activeTurn = activeTurn; + activeTurn.adoptedWake = wake !== null && !continuation; if (compactCommand !== null) { yield* connection.send(compactRpcRecord(compactCommand)); pendingCompactResponses.push({ @@ -2368,6 +2544,7 @@ export function makePiAdapterV2( yield* connection.send({ type: "prompt", message: payload.message, + ...(wake === null ? {} : { streamingBehavior: "steer" }), ...(payload.images.length === 0 ? {} : { images: payload.images }), }); pendingPromptResponses.push({ @@ -2387,6 +2564,18 @@ export function makePiAdapterV2( lastRunOrdinal: turnInput.runOrdinal, }); yield* updateProviderSession("running", null); + if (wake !== null) { + for (const event of wake.events) { + if (!continuation && event["type"] === "agent_settled") { + activeTurn.settleWhenIdle = true; + continue; + } + yield* handleSessionEvent(event); + } + } else if (continuation) { + // A user turn, teardown, or newer binding already took the wake. + yield* finalizeTurn(state); + } if (outOfTurnExtensionErrors.length > 0) { yield* Queue.offer(connection.events, { type: "t3.flush_extension_errors" }); } @@ -2478,7 +2667,11 @@ export function makePiAdapterV2( const turn = threadState?.activeTurn ?? null; // Stop on a settled turn: Pi runs nothing between prompts, so // nothing of that turn is left to stop. - if (turn === null && interruptInput.requestRuntimeRestart === true) return; + if (turn === null && interruptInput.requestRuntimeRestart === true) { + if (pendingWake !== null) + yield* sessionEventPermit.withPermits(1)(stopPendingWake(pendingWake)); + return; + } if (turn === null || turn.providerTurn.id !== interruptInput.providerTurnId) { return yield* protocolError(`Pi turn ${interruptInput.providerTurnId} is not active`); } @@ -2662,6 +2855,9 @@ export function makePiAdapterV2( if (state.activeTurn !== null) { return yield* protocolError("Cannot roll back while a Pi turn is active"); } + if (pendingWake !== null) { + return yield* protocolError("Cannot roll back while native Pi work awaits a turn"); + } // `fork(entryId)` re-roots the active branch before that user // message, so the rollback boundary is the first user entry of // the earliest turn being discarded. @@ -2673,42 +2869,94 @@ export function makePiAdapterV2( if (forkEntryId === undefined) { return yield* protocolError("Pi rollback target has no captured session-tree entry"); } - const forkData = yield* lifecycleRequest({ type: "fork", entryId: forkEntryId }); - if (recordField(forkData, "cancelled") === true) { - return yield* protocolError("A Pi extension cancelled the session fork"); - } - // Pi fork replaces the session file, including for rollback. Persist - // its new identity before any later request can fail or restart. - // An interrupted read leaves the identity just as unknown as a failed one. - const forkState = yield* request({ type: "get_state" }).pipe( - Effect.onError(() => - Effect.sync(() => { - threadState = null; + return yield* Effect.acquireUseRelease( + sessionEventPermit.withPermits(1)( + Effect.gen(function* () { + if ( + rollbackBarrier !== null || + threadState !== state || + state.activeTurn !== null || + pendingWake !== null || + closed || + stopRequested + ) + return yield* protocolError( + "Pi rollback ownership changed before its fork could start", + ); + const barrier = yield* Deferred.make< + void, + ProviderAdapter.ProviderAdapterProtocolError + >(); + rollbackBarrier = barrier; + return barrier; }), ), + (barrier) => + Effect.gen(function* () { + const forkData = yield* lifecycleRequest({ type: "fork", entryId: forkEntryId }); + if (recordField(forkData, "cancelled") === true) + return yield* protocolError("A Pi extension cancelled the session fork"); + const forkState = yield* request({ type: "get_state" }).pipe( + Effect.onError(() => + Effect.sync(() => { + threadState = null; + }), + ), + ); + const forkSessionFile = recordString(forkState, "sessionFile"); + if (forkSessionFile === undefined) { + threadState = null; + return yield* protocolError("Pi fork did not return a persisted session file"); + } + const entriesData = yield* request({ type: "get_entries" }).pipe( + Effect.orElseSucceed(() => undefined), + ); + const leafId = recordString(entriesData, "leafId") ?? null; + // RPC responses bypass the pump. Drain preceding native events + // before deciding whether this fork can replace their binding. + const queued = yield* Queue.offer(connection.events, { + type: "t3.rollback_barrier", + barrier, + }); + if (!queued) return yield* protocolError("Pi transport closed during rollback"); + yield* Deferred.await(barrier); + return yield* sessionEventPermit.withPermits(1)( + Effect.gen(function* () { + if (pendingWake !== null) { + yield* stopPendingWake(pendingWake); + return yield* protocolError( + "Native Pi work started while rollback was changing the session", + ); + } + if ( + threadState !== state || + state.activeTurn !== null || + closed || + stopRequested + ) + return yield* protocolError( + "Pi rollback ownership changed during its fork", + ); + lastNativeThreadId = forkSessionFile; + appliedModel = null; + appliedThinking = null; + appliedSessionName = null; + lastKnownLeaf = leafId; + leafCursorStale = entriesData === undefined; + yield* updateProviderThread(state, { + nativeThreadRef: providerRef(forkSessionFile), + nativeConversationHeadRef: leafId === null ? null : providerRef(leafId), + }); + rollbackBarrier = null; + return piThreadSnapshot(state.providerThread); + }), + ); + }).pipe(Effect.onInterrupt(() => connection.terminate)), + (barrier) => + Effect.sync(() => { + if (rollbackBarrier === barrier) rollbackBarrier = null; + }), ); - const forkSessionFile = recordString(forkState, "sessionFile"); - if (forkSessionFile === undefined) { - threadState = null; - return yield* protocolError("Pi fork did not return a persisted session file"); - } - lastNativeThreadId = forkSessionFile; - appliedModel = null; - appliedThinking = null; - appliedSessionName = null; - yield* updateProviderThread(state, { nativeThreadRef: providerRef(forkSessionFile) }); - const entriesData = yield* request({ type: "get_entries" }).pipe( - Effect.orElseSucceed(() => undefined), - ); - const leafId = recordString(entriesData, "leafId") ?? null; - lastKnownLeaf = leafId; - // The fork re-baselined the tree, so the cursor is trustworthy - // again unless this listing itself failed. - leafCursorStale = entriesData === undefined; - yield* updateProviderThread(state, { - nativeConversationHeadRef: leafId === null ? null : providerRef(leafId), - }); - return piThreadSnapshot(state.providerThread); }).pipe( Effect.mapError( (cause) => @@ -2971,6 +3219,7 @@ export const PiAdapterV2Driver: ProviderAdapterDriver @@ -3005,6 +3255,7 @@ const layer: Layer.Layer