authorgravatar for 71498452+StiensWout@users.noreply.github.comWout Stiens <71498452+StiensWout@users.noreply.github.com> 2026-10-08 20:30:47+02:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2026-10-08 13:30:47-05:00
logf999eb2a8e7d8db74d4f4ab7dd66f86163e1a226
tree26cece8d6945e66270c26305638f653bb23fb43c
parent62a5d3720f82f2dc3ace48fcb26c90c6e3976b50
signature Signed by PGP key B5690EEEBB952194

fix(server): Pi extension wakes get an owned continuation turn (#17214)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

2 files changed, 1130 insertions(+), 61 deletions(-)

apps/server/src/orchestration-v2/Adapters/PiAdapterV2.test.ts+831-14
......@@ -1,6 +1,7 @@
11import { assert, describe, it } from "@effect/vitest";
22import * as NodeServices from "@effect/platform-node/NodeServices";
33import {
4 CheckpointId,
45 EnvironmentId,
56 NodeId,
67 ProviderInstanceId,
......@@ -30,8 +31,10 @@ import * as Sink from "effect/Sink";
3031import * as Stream from "effect/Stream";
3132import * as TestClock from "effect/testing/TestClock";
3233import { ChildProcess, ChildProcessSpawner } from "effect/process";
34import { HostProcessEnvironment } from "@t3tools/shared/hostProcess";
3335
3436import * as ServerConfig from "../../config.ts";
37import * as ProviderContinuationRequests from "../ProviderContinuationRequests.ts";
3538import * as McpProviderSession from "../../mcp/McpProviderSession.ts";
3639import * as IdAllocator from "../IdAllocator.ts";
3740import {
......@@ -40,7 +43,12 @@ import {
4043 type ProviderAdapterV2SessionRuntime,
4144} from "../ProviderAdapter.ts";
4245import { handoffBudget } from "../ContextHandoffBudget.ts";
43import { makePiAdapterV2, PI_PROVIDER } from "./PiAdapterV2.ts";
46import {
47 makePiAdapterV2,
48 PiAdapterV2Driver,
49 PI_PROVIDER,
50 type PiAdapterV2Options,
51} from "./PiAdapterV2.ts";
4452import { makePiRpcConnection, type PiRpcRecord } from "./PiRpc.ts";
4553
4654const layerServerConfig = ServerConfig.layerTest(process.cwd(), {
......@@ -88,7 +96,9 @@ interface FakePi {
8896 readonly resolveDeferredState: (data: unknown) => Effect.Effect<void>;
8997 /** Reject the next `get_state` request. */
9098 readonly failNextState: () => void;
91 readonly deferNextLifecycle: (type: "switch_session" | "new_session") => void;
99 readonly deferNextLifecycle: (type: "switch_session" | "new_session" | "fork") => void;
100 /** Hold a model-select extension hook until its UI request is answered. */
101 readonly deferNextModelSelection: () => void;
92102 readonly queueModels: (models: ReadonlyArray<unknown>) => void;
93103 readonly vetoNextNewSession: () => void;
94104 /** Every request received by the fake process. */
......@@ -291,6 +301,9 @@ const makeFakePi: Effect.Effect<FakePi> = Effect.gen(function* () {
291301 deferNextLifecycle: (type) => {
292302 deferredLifecycle = type;
293303 },
304 deferNextModelSelection: () => {
305 deferredLifecycle = "set_model";
306 },
294307 queueModels: (value) => {
295308 models = value;
296309 },
......@@ -310,22 +323,44 @@ const makeFakePi: Effect.Effect<FakePi> = Effect.gen(function* () {
310323 } satisfies FakePi;
311324});
312325
313const makeAdapter = Effect.fnUntraced(function* (fake: FakePi, launchArgs = "", forkFake?: FakePi) {
326const makeAdapter = Effect.fnUntraced(function* (
327 fake: FakePi,
328 launchArgs = "",
329 forkFake?: FakePi,
330 continuationRequests?: PiAdapterV2Options["continuationRequests"],
331) {
314332 const idAllocator = yield* IdAllocator.IdAllocatorV2;
315333 const serverConfig = yield* ServerConfig.ServerConfig;
316334 const fileSystem = yield* FileSystem.FileSystem;
335 const spawner =
336 forkFake === undefined
337 ? fake.spawner
338 : ChildProcessSpawner.make((command) =>
339 ChildProcess.isStandardCommand(command) && command.args.includes("--fork")
340 ? forkFake.spawner.spawn(command)
341 : fake.spawner.spawn(command),
342 );
343 if (continuationRequests !== undefined) {
344 return yield* PiAdapterV2Driver.create({
345 instanceId: PI_INSTANCE_ID,
346 displayName: undefined,
347 enabled: true,
348 environment: [],
349 config: { enabled: true, binaryPath: "pi", launchArgs, customModels: [] },
350 }).pipe(
351 Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, spawner),
352 Effect.provideService(HostProcessEnvironment, {}),
353 Effect.provideService(ProviderContinuationRequests.ProviderContinuationRequests, {
354 ...continuationRequests,
355 take: Effect.never,
356 }),
357 );
358 }
317359 return makePiAdapterV2({
318360 instanceId: PI_INSTANCE_ID,
319361 settings: { enabled: true, binaryPath: "pi", launchArgs, customModels: [] },
320362 environment: {},
321 spawner:
322 forkFake === undefined
323 ? fake.spawner
324 : ChildProcessSpawner.make((command) =>
325 ChildProcess.isStandardCommand(command) && command.args.includes("--fork")
326 ? forkFake.spawner.spawn(command)
327 : fake.spawner.spawn(command),
328 ),
363 spawner,
329364 fileSystem,
330365 idAllocator,
331366 serverConfig,
......@@ -338,8 +373,9 @@ const openRuntime = Effect.fnUntraced(function* (
338373 threadId = THREAD_ID,
339374 providerSessionId = SESSION_ID,
340375 forkFake?: FakePi,
376 continuationRequests?: PiAdapterV2Options["continuationRequests"],
341377) {
342 const adapter = yield* makeAdapter(fake, "", forkFake);
378 const adapter = yield* makeAdapter(fake, "", forkFake, continuationRequests);
343379 const runtime = yield* adapter.openSession({
344380 threadId,
345381 providerSessionId,
......@@ -397,6 +433,7 @@ const startTurn = Effect.fnUntraced(function* (
397433 selection?: ModelSelection,
398434 runOrdinal = 1,
399435 threadId = THREAD_ID,
436 continuation = false,
400437) {
401438 const appThread = yield* makeAppThread(model, threadId);
402439 const runId = RunId.make(`run:${threadId}:${runOrdinal}`);
......@@ -413,8 +450,8 @@ const startTurn = Effect.fnUntraced(function* (
413450 messageId: `message:${threadId}:${runOrdinal}` as never,
414451 text,
415452 attachments,
416 createdBy: "user",
417 creationSource: "web",
453 createdBy: continuation ? "agent" : "user",
454 creationSource: continuation ? "provider" : "web",
418455 },
419456 modelSelection: selection ?? modelSelection(model),
420457 runtimePolicy,
......@@ -461,6 +498,786 @@ const expectModelFailure = (errorMessage: string) =>
461498 }).pipe(Effect.scoped, Effect.provide(layerTest));
462499
463500describe("PiAdapterV2", () => {
501 it.effect.each([false, true])(
502 "serializes rollback fork hooks with native wake ownership, wake=%s",
503 (nativeWake) =>
504 Effect.gen(function* () {
505 const fake = yield* makeFakePi;
506 const offers =
507 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
508 const { runtime, takeEvent } = yield* openRuntime(
509 fake,
510 "default",
511 THREAD_ID,
512 SESSION_ID,
513 undefined,
514 {
515 offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid),
516 },
517 );
518 const providerThread = yield* runtime.ensureThread({
519 threadId: THREAD_ID,
520 modelSelection: modelSelection("default"),
521 runtimePolicy,
522 });
523 fake.deferNextLifecycle("fork");
524 const rollingBack = yield* runtime
525 .rollbackThread({
526 providerThread,
527 target: {
528 type: "thread_start",
529 checkpointId: CheckpointId.make("checkpoint-start"),
530 appRunOrdinal: 0,
531 },
532 providerThreadTurns: [
533 {
534 id: ProviderTurnId.make("previous-turn"),
535 providerThreadId: providerThread.id,
536 nodeId: NodeId.make("previous-node"),
537 runAttemptId: null,
538 nativeTurnRef: {
539 driver: PI_PROVIDER,
540 nativeId: "user-before-wake",
541 strength: "strong",
542 },
543 ordinal: 1,
544 status: "completed",
545 startedAt: providerThread.createdAt,
546 completedAt: providerThread.createdAt,
547 },
548 ],
549 })
550 .pipe(Effect.result, Effect.forkScoped);
551 const fork = yield* fake.takeRequest("fork");
552 const overlappingTurn = yield* startTurn(runtime, providerThread).pipe(Effect.result);
553 assert.equal(overlappingTurn._tag, "Failure");
554 const overlappingRegistration = yield* runtime
555 .ensureThread({
556 threadId: THREAD_ID,
557 modelSelection: modelSelection("default"),
558 runtimePolicy,
559 })
560 .pipe(Effect.result);
561 assert.equal(overlappingRegistration._tag, "Failure");
562 yield* fake.emit({
563 type: "extension_ui_request",
564 id: "fork-hook-dialog",
565 method: "confirm",
566 title: "Restore this branch?",
567 });
568 const forkDialog = yield* takeEvent((event) => event.type === "runtime_request.updated");
569 assert.isTrue(forkDialog.type === "runtime_request.updated");
570 if (forkDialog.type !== "runtime_request.updated") return;
571 assert.isNull(forkDialog.runtimeRequest.providerTurnId);
572 yield* runtime.respondToRuntimeRequest({
573 requestId: forkDialog.runtimeRequest.id,
574 decision: "accept",
575 });
576 yield* fake.takeRequest("extension_ui_response");
577 if (nativeWake) {
578 yield* fake.emit({ type: "agent_start" });
579 yield* fake.emit({ type: "message_start", message: { role: "assistant" } });
580 yield* fake.emit({
581 type: "message_update",
582 assistantMessageEvent: {
583 type: "text_end",
584 contentIndex: 0,
585 content: "Discarded branch wake output",
586 },
587 });
588 yield* fake.emit({ type: "agent_settled" });
589 const retired = yield* takeEvent(
590 (event) =>
591 event.type === "provider_session.updated" &&
592 event.providerSession.status === "stopped",
593 ).pipe(Effect.as(true), Effect.raceFirst(Queue.take(offers).pipe(Effect.as(false))));
594 assert.isTrue(retired, "Pi offered native work while rollback was changing its binding");
595 const result = yield* Fiber.join(rollingBack);
596 assert.equal(result._tag, "Failure");
597 assert.isFalse(yield* runtime.hasPendingBackgroundWork!);
598 const staleTurn = yield* startTurn(
599 runtime,
600 providerThread,
601 "default",
602 [],
603 "Synthetic wake",
604 undefined,
605 1,
606 THREAD_ID,
607 true,
608 ).pipe(Effect.result);
609 assert.equal(staleTurn._tag, "Failure");
610 assert.isFalse(fake.allRequests().some((request) => request["type"] === "prompt"));
611 return;
612 }
613 fake.queueState({ sessionFile: "/fake/rolled-back-session.jsonl" });
614 yield* fake.emit({
615 type: "response",
616 id: fork["id"],
617 command: "fork",
618 success: true,
619 data: { cancelled: false },
620 });
621 const result = yield* Fiber.join(rollingBack);
622 assert.equal(result._tag, "Success");
623 if (result._tag !== "Success") return;
624 assert.equal(
625 result.success.providerThread.nativeThreadRef?.nativeId,
626 "/fake/rolled-back-session.jsonl",
627 );
628 yield* startTurn(runtime, result.success.providerThread);
629 const prompt = yield* fake.takeRequest("prompt");
630 assert.equal(prompt["message"], "Hello pi");
631 yield* fake.emit({ type: "response", command: "prompt", success: true });
632 yield* fake.emit({ type: "agent_start" });
633 yield* fake.emit({ type: "agent_settled" });
634 const terminal = yield* takeEvent((event) => event.type === "turn.terminal");
635 assert.isTrue(terminal.type === "turn.terminal" && terminal.status === "completed");
636 }).pipe(Effect.scoped, Effect.provide(layerTest)),
637 );
638
639 it.effect.each([false, true])(
640 "retains native wake ownership after a rejected joining prompt, compaction aborted=%s",
641 (aborted) =>
642 Effect.gen(function* () {
643 const fake = yield* makeFakePi;
644 const offers =
645 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
646 const { runtime, takeEvent } = yield* openRuntime(
647 fake,
648 "default",
649 THREAD_ID,
650 SESSION_ID,
651 undefined,
652 {
653 offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid),
654 },
655 );
656 const providerThread = yield* runtime.ensureThread({
657 threadId: THREAD_ID,
658 modelSelection: modelSelection("default"),
659 runtimePolicy,
660 });
661 yield* fake.takeRequest("get_state");
662 yield* fake.emit({ type: "agent_start" });
663 yield* Queue.take(offers);
664 yield* fake.emit({ type: "agent_settled" });
665 yield* fake.emit({ type: "compaction_start", reason: "manual" });
666 yield* startTurn(runtime, providerThread);
667 yield* fake.takeRequest("prompt");
668 const rejection = "Cannot submit a prompt while compaction is in progress.";
669 yield* fake.emit({ type: "response", command: "prompt", success: false, error: rejection });
670 yield* fake.emit({
671 type: "extension_ui_request",
672 id: "native-compaction-dialog",
673 method: "confirm",
674 title: "Continue native work?",
675 });
676 const dialog = yield* takeEvent((event) => event.type === "runtime_request.updated");
677 assert.isTrue(dialog.type === "runtime_request.updated");
678 if (dialog.type !== "runtime_request.updated") return;
679 assert.isNotNull(dialog.runtimeRequest.providerTurnId);
680 yield* runtime.respondToRuntimeRequest({
681 requestId: dialog.runtimeRequest.id,
682 decision: "accept",
683 });
684 yield* fake.takeRequest("extension_ui_response");
685 if (aborted) {
686 yield* fake.emit({ type: "compaction_end", result: null, aborted: true });
687 const terminal = yield* takeEvent((event) => event.type === "turn.terminal");
688 assert.isTrue(
689 terminal.type === "turn.terminal" &&
690 terminal.status === "failed" &&
691 terminal.failure.message === rejection,
692 );
693 return;
694 }
695 fake.queueState({ isStreaming: true });
696 yield* fake.emit({
697 type: "compaction_end",
698 result: { summary: "Native compaction finished" },
699 willRetry: true,
700 });
701 yield* fake.emit({ type: "agent_start" });
702 yield* fake.emit({ type: "message_start", message: { role: "assistant" } });
703 yield* fake.emit({
704 type: "message_update",
705 assistantMessageEvent: {
706 type: "text_end",
707 contentIndex: 0,
708 content: "Native work completed",
709 },
710 });
711 const message = yield* takeEvent((event) => event.type === "message.updated");
712 assert.isTrue(message.type === "message.updated" && message.message.runId !== null);
713 yield* fake.takeRequest("get_state");
714 yield* fake.emit({ type: "agent_settled" });
715 const terminal = yield* takeEvent((event) => event.type === "turn.terminal");
716 assert.isTrue(
717 terminal.type === "turn.terminal" &&
718 terminal.status === "failed" &&
719 terminal.failure.message === rejection,
720 );
721 }).pipe(Effect.scoped, Effect.provide(layerTest)),
722 );
723
724 it.effect("rejects rollback while native wake output is awaiting ownership", () =>
725 Effect.gen(function* () {
726 const fake = yield* makeFakePi;
727 const offers =
728 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
729 const { runtime, takeEvent } = yield* openRuntime(
730 fake,
731 "default",
732 THREAD_ID,
733 SESSION_ID,
734 undefined,
735 {
736 offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid),
737 },
738 );
739 const providerThread = yield* runtime.ensureThread({
740 threadId: THREAD_ID,
741 modelSelection: modelSelection("default"),
742 runtimePolicy,
743 });
744 yield* fake.emit({ type: "agent_start" });
745 yield* Queue.take(offers);
746 yield* fake.emit({ type: "message_start", message: { role: "assistant" } });
747 yield* fake.emit({
748 type: "message_update",
749 assistantMessageEvent: {
750 type: "text_end",
751 contentIndex: 0,
752 content: "Current branch wake output",
753 },
754 });
755 yield* fake.emit({ type: "agent_settled" });
756 yield* takeEvent(
757 (event) =>
758 event.type === "provider_session.updated" && event.providerSession.status === "ready",
759 );
760 const rolledBack = yield* runtime
761 .rollbackThread({
762 providerThread,
763 target: {
764 type: "thread_start",
765 checkpointId: CheckpointId.make("checkpoint-start"),
766 appRunOrdinal: 0,
767 },
768 providerThreadTurns: [
769 {
770 id: ProviderTurnId.make("previous-turn"),
771 providerThreadId: providerThread.id,
772 nodeId: NodeId.make("previous-node"),
773 runAttemptId: null,
774 nativeTurnRef: {
775 driver: PI_PROVIDER,
776 nativeId: "user-before-wake",
777 strength: "strong",
778 },
779 ordinal: 1,
780 status: "completed",
781 startedAt: providerThread.createdAt,
782 completedAt: providerThread.createdAt,
783 },
784 ],
785 })
786 .pipe(Effect.result);
787 assert.equal(
788 rolledBack._tag,
789 "Failure",
790 "Rollback changed the native session while its wake was still buffered",
791 );
792 assert.isFalse(fake.allRequests().some((request) => request["type"] === "fork"));
793 assert.isTrue(yield* runtime.hasPendingBackgroundWork!);
794 yield* startTurn(
795 runtime,
796 providerThread,
797 "default",
798 [],
799 "Synthetic wake",
800 undefined,
801 1,
802 THREAD_ID,
803 true,
804 );
805 const message = yield* takeEvent((event) => event.type === "message.updated");
806 assert.equal(
807 message.type === "message.updated" && message.message.text,
808 "Current branch wake output",
809 );
810 yield* takeEvent((event) => event.type === "turn.terminal");
811 }).pipe(Effect.scoped, Effect.provide(layerTest)),
812 );
813
814 it.effect("answers model-selection dialogs before a user turn claims a native wake", () =>
815 Effect.gen(function* () {
816 const fake = yield* makeFakePi;
817 const offers =
818 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
819 const { runtime, takeEvent } = yield* openRuntime(
820 fake,
821 "default",
822 THREAD_ID,
823 SESSION_ID,
824 undefined,
825 { offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid) },
826 );
827 const providerThread = yield* runtime.ensureThread({
828 threadId: THREAD_ID,
829 modelSelection: modelSelection("default"),
830 runtimePolicy,
831 });
832 yield* fake.emit({ type: "agent_start" });
833 yield* Queue.take(offers);
834 const requests: Array<Extract<ProviderAdapterV2Event, { type: "runtime_request.updated" }>> =
835 [];
836 fake.deferNextModelSelection();
837 const starting = yield* startTurn(
838 runtime,
839 providerThread,
840 "default",
841 [],
842 "Continue the review",
843 modelSelection("xai/grok-4.6"),
844 ).pipe(Effect.forkScoped);
845 const modelRequest = yield* fake.takeRequest("set_model");
846 yield* fake.emit({
847 type: "extension_ui_request",
848 id: "model-selection-dialog",
849 method: "confirm",
850 title: "Use this model?",
851 });
852 yield* fake.emit({ type: "agent_settled" });
853 yield* takeEvent((event) => {
854 if (event.type === "runtime_request.updated") requests.push(event);
855 return (
856 event.type === "provider_session.updated" && event.providerSession.status === "ready"
857 );
858 });
859 const dialog = requests[0]?.runtimeRequest;
860 assert.isDefined(dialog, "Pi buffered the dialog needed to finish model selection");
861 assert.equal(requests[0]?.threadId, THREAD_ID);
862 assert.isNull(dialog!.providerTurnId);
863 yield* runtime.respondToRuntimeRequest({ requestId: dialog!.id, decision: "accept" });
864 const response = yield* fake.takeRequest("extension_ui_response");
865 assert.equal(response["id"], "model-selection-dialog");
866 assert.isTrue(response["confirmed"]);
867 yield* fake.emit({
868 type: "response",
869 id: modelRequest["id"],
870 command: "set_model",
871 success: true,
872 data: { provider: "xai", id: "grok-4.6" },
873 });
874 yield* Fiber.join(starting);
875 const prompt = yield* fake.takeRequest("prompt");
876 assert.equal(prompt["message"], "Continue the review");
877 assert.equal(prompt["streamingBehavior"], "steer");
878 const resolved = yield* takeEvent(
879 (event) =>
880 event.type === "runtime_request.updated" && event.runtimeRequest.id === dialog!.id,
881 );
882 assert.isTrue(
883 resolved.type === "runtime_request.updated" &&
884 resolved.runtimeRequest.status === "resolved",
885 );
886 }).pipe(Effect.scoped, Effect.provide(layerTest)),
887 );
888
889 it.effect("adopts a completed native wake without sending a synthetic prompt", () =>
890 Effect.gen(function* () {
891 const fake = yield* makeFakePi;
892 const offers =
893 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
894 const { runtime, takeEvent } = yield* openRuntime(
895 fake,
896 "default",
897 THREAD_ID,
898 SESSION_ID,
899 undefined,
900 {
901 offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid),
902 },
903 );
904 const providerThread = yield* runtime.ensureThread({
905 threadId: THREAD_ID,
906 modelSelection: modelSelection("default"),
907 runtimePolicy,
908 });
909 yield* fake.emit({ type: "agent_start" });
910 const offer = yield* Queue.take(offers).pipe(
911 Effect.raceFirst(
912 takeEvent(
913 (event) =>
914 event.type === "provider_session.updated" && event.providerSession.status === "error",
915 ).pipe(
916 Effect.andThen(Effect.sync(() => assert.fail("Pi killed a routable native wake"))),
917 ),
918 ),
919 );
920 assert.equal(offer.threadId, THREAD_ID);
921 assert.equal(offer.providerThreadId, providerThread.id);
922 assert.equal(offer.delivery, "adapter_buffered");
923 assert.isTrue(yield* runtime.hasPendingBackgroundWork!);
924 let dispatchCommitted = false;
925 yield* offer.dispatchIfCurrent!(
926 Effect.sync(() => {
927 dispatchCommitted = true;
928 }),
929 );
930 assert.isTrue(dispatchCommitted);
931 assert.isTrue(yield* runtime.hasPendingBackgroundWork!);
932 yield* fake.emit({ type: "message_start", message: { role: "assistant" } });
933 yield* fake.emit({
934 type: "message_update",
935 assistantMessageEvent: {
936 type: "text_delta",
937 contentIndex: 0,
938 delta: "Background review finished",
939 },
940 });
941 yield* fake.emit({ type: "message_end", message: { role: "assistant" } });
942 yield* fake.emit({ type: "agent_settled" });
943 yield* takeEvent(
944 (event) =>
945 event.type === "provider_session.updated" && event.providerSession.status === "ready",
946 );
947 assert.isTrue(yield* runtime.hasPendingBackgroundWork!);
948 yield* startTurn(
949 runtime,
950 providerThread,
951 "default",
952 [],
953 "Synthetic wake",
954 undefined,
955 1,
956 THREAD_ID,
957 true,
958 );
959 const message = yield* takeEvent((event) => event.type === "message.updated");
960 assert.equal(
961 message.type === "message.updated" && message.message.text,
962 "Background review finished",
963 );
964 yield* takeEvent((event) => event.type === "turn.terminal");
965 assert.equal(fake.allRequests().filter((request) => request["type"] === "prompt").length, 0);
966 assert.isFalse(yield* runtime.hasPendingBackgroundWork!);
967 let dispatched = false;
968 yield* offer.dispatchIfCurrent!(
969 Effect.sync(() => {
970 dispatched = true;
971 }),
972 );
973 assert.isFalse(dispatched);
974 }).pipe(Effect.scoped, Effect.provide(layerTest)),
975 );
976
977 it.effect.each([false, true])(
978 "a user prompt takes the wake without premature settlement, native wake settled=%s",
979 (settled) =>
980 Effect.gen(function* () {
981 const fake = yield* makeFakePi;
982 const offers =
983 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
984 const { runtime, takeEvent } = yield* openRuntime(
985 fake,
986 "default",
987 THREAD_ID,
988 SESSION_ID,
989 undefined,
990 {
991 offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid),
992 },
993 );
994 const providerThread = yield* runtime.ensureThread({
995 threadId: THREAD_ID,
996 modelSelection: modelSelection("default"),
997 runtimePolicy,
998 });
999 yield* fake.emit({ type: "agent_start" });
1000 const offer = yield* Queue.take(offers);
1001 if (settled) {
1002 yield* fake.emit({ type: "agent_settled" });
1003 yield* takeEvent(
1004 (event) =>
1005 event.type === "provider_session.updated" && event.providerSession.status === "ready",
1006 );
1007 }
1008 yield* startTurn(runtime, providerThread);
1009 const prompt = yield* fake.takeRequest("prompt");
1010 assert.equal(prompt["message"], "Hello pi");
1011 assert.equal(prompt["streamingBehavior"], "steer");
1012 yield* fake.emit({ type: "response", command: "prompt", success: true });
1013 yield* fake.emit({ type: "agent_start" });
1014 yield* fake.emit({
1015 type: "extension_ui_request",
1016 id: "owned-dialog",
1017 method: "confirm",
1018 title: "Continue?",
1019 });
1020 const dialog = yield* takeEvent((event) => event.type === "runtime_request.updated");
1021 assert.isTrue(dialog.type === "runtime_request.updated");
1022 if (dialog.type !== "runtime_request.updated") return;
1023 assert.isNotNull(dialog.runtimeRequest.providerTurnId);
1024 yield* offer.clearIfCurrent!();
1025 yield* runtime.respondToRuntimeRequest({
1026 requestId: dialog.runtimeRequest.id,
1027 decision: "accept",
1028 });
1029 const response = yield* fake.takeRequest("extension_ui_response");
1030 assert.isTrue(response["confirmed"]);
1031 yield* fake.emit({ type: "message_start", message: { role: "assistant" } });
1032 yield* fake.emit({
1033 type: "message_update",
1034 assistantMessageEvent: {
1035 type: "text_end",
1036 contentIndex: 0,
1037 content: "User prompt completed",
1038 },
1039 });
1040 yield* fake.emit({ type: "agent_settled" });
1041 const message = yield* takeEvent(
1042 (event) =>
1043 event.type === "message.updated" && event.message.text === "User prompt completed",
1044 );
1045 assert.isTrue(message.type === "message.updated" && message.message.runId !== null);
1046 yield* takeEvent((event) => event.type === "turn.terminal");
1047 let dispatched = false;
1048 yield* offer.dispatchIfCurrent!(
1049 Effect.sync(() => {
1050 dispatched = true;
1051 }),
1052 );
1053 assert.isFalse(dispatched);
1054 }).pipe(Effect.scoped, Effect.provide(layerTest)),
1055 );
1056
1057 it.effect("defers a wake's late settlement until the joined user prompt is acknowledged", () =>
1058 Effect.gen(function* () {
1059 const fake = yield* makeFakePi;
1060 const offers =
1061 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
1062 const { runtime, takeEvent } = yield* openRuntime(
1063 fake,
1064 "default",
1065 THREAD_ID,
1066 SESSION_ID,
1067 undefined,
1068 {
1069 offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid),
1070 },
1071 );
1072 const providerThread = yield* runtime.ensureThread({
1073 threadId: THREAD_ID,
1074 modelSelection: modelSelection("default"),
1075 runtimePolicy,
1076 });
1077 yield* fake.emit({ type: "agent_start" });
1078 yield* Queue.take(offers);
1079 yield* startTurn(runtime, providerThread);
1080 yield* fake.takeRequest("prompt");
1081 yield* fake.emit({ type: "agent_settled" });
1082 yield* fake.emit({
1083 type: "extension_ui_request",
1084 method: "notify",
1085 message: "Ordering barrier",
1086 });
1087 yield* takeEvent(
1088 (event) =>
1089 event.type === "turn_item.updated" &&
1090 event.turnItem.type === "dynamic_tool" &&
1091 event.turnItem.toolName === "notify",
1092 );
1093 assert.equal(runtime.providerSession.status, "running");
1094 fake.queueState({ isStreaming: true });
1095 yield* fake.emit({ type: "response", command: "prompt", success: true });
1096 yield* fake.takeRequest("get_state");
1097 yield* fake.emit({ type: "agent_start" });
1098 yield* fake.emit({ type: "agent_settled" });
1099 yield* takeEvent((event) => event.type === "turn.terminal");
1100 }).pipe(Effect.scoped, Effect.provide(layerTest)),
1101 );
1102
1103 it.effect("settles a command-only user prompt after taking a completed wake", () =>
1104 Effect.gen(function* () {
1105 const fake = yield* makeFakePi;
1106 const offers =
1107 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
1108 const { runtime, takeEvent } = yield* openRuntime(
1109 fake,
1110 "default",
1111 THREAD_ID,
1112 SESSION_ID,
1113 undefined,
1114 {
1115 offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid),
1116 },
1117 );
1118 const providerThread = yield* runtime.ensureThread({
1119 threadId: THREAD_ID,
1120 modelSelection: modelSelection("default"),
1121 runtimePolicy,
1122 });
1123 yield* fake.emit({ type: "agent_start" });
1124 yield* Queue.take(offers);
1125 yield* fake.emit({ type: "agent_settled" });
1126 yield* takeEvent(
1127 (event) =>
1128 event.type === "provider_session.updated" && event.providerSession.status === "ready",
1129 );
1130 yield* startTurn(runtime, providerThread, "default", [], "/hello");
1131 yield* fake.takeRequest("prompt");
1132 yield* fake.emit({ type: "response", command: "prompt", success: true });
1133 yield* takeEvent((event) => event.type === "turn.terminal");
1134 yield* startTurn(
1135 runtime,
1136 providerThread,
1137 "default",
1138 [],
1139 "Stale continuation",
1140 undefined,
1141 2,
1142 THREAD_ID,
1143 true,
1144 );
1145 yield* takeEvent((event) => event.type === "turn.terminal");
1146 assert.equal(fake.allRequests().filter((request) => request["type"] === "prompt").length, 1);
1147 }).pipe(Effect.scoped, Effect.provide(layerTest)),
1148 );
1149
1150 it.effect("keeps a pending wake on its thread and lets Stop retire its process", () =>
1151 Effect.gen(function* () {
1152 const fake = yield* makeFakePi;
1153 const offers =
1154 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
1155 const { runtime, takeEvent } = yield* openRuntime(
1156 fake,
1157 "default",
1158 THREAD_ID,
1159 SESSION_ID,
1160 undefined,
1161 {
1162 offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid),
1163 },
1164 );
1165 const providerThread = yield* runtime.ensureThread({
1166 threadId: THREAD_ID,
1167 modelSelection: modelSelection("default"),
1168 runtimePolicy,
1169 });
1170 yield* fake.emit({ type: "agent_start" });
1171 const offer = yield* Queue.take(offers);
1172 assert.isTrue(yield* runtime.hasPendingBackgroundWorkForThread!(providerThread));
1173 const otherThread = {
1174 ...providerThread,
1175 appThreadId: ThreadId.make("other-thread"),
1176 id: ProviderThreadId.make("other-provider-thread"),
1177 };
1178 assert.isFalse(yield* runtime.hasPendingBackgroundWorkForThread!(otherThread));
1179 const failure = yield* runtime
1180 .resumeThread({ providerThread: otherThread })
1181 .pipe(Effect.flip);
1182 assert.equal(failure._tag, "ProviderAdapterResumeThreadError");
1183 yield* runtime.interruptTurn({
1184 providerThread,
1185 providerTurnId: ProviderTurnId.make("settled-turn"),
1186 requestRuntimeRestart: true,
1187 });
1188 yield* takeEvent(
1189 (event) =>
1190 event.type === "provider_session.updated" && event.providerSession.status === "stopped",
1191 );
1192 assert.isFalse(yield* runtime.hasPendingBackgroundWork!);
1193 let dispatched = false;
1194 yield* offer.dispatchIfCurrent!(
1195 Effect.sync(() => {
1196 dispatched = true;
1197 }),
1198 );
1199 assert.isFalse(dispatched);
1200 }).pipe(Effect.scoped, Effect.provide(layerTest)),
1201 );
1202
1203 it.effect.each(["archive", "stop", "transport death"] as const)(
1204 "retiring a wake by %s cancels visible dialogs and invalidates its offer",
1205 (retirement) =>
1206 Effect.gen(function* () {
1207 const fake = yield* makeFakePi;
1208 const offers =
1209 yield* Queue.unbounded<ProviderContinuationRequests.ProviderContinuationRequest>();
1210 const { runtime, takeEvent } = yield* openRuntime(
1211 fake,
1212 "default",
1213 THREAD_ID,
1214 SESSION_ID,
1215 undefined,
1216 {
1217 offer: (request) => Queue.offer(offers, request).pipe(Effect.asVoid),
1218 },
1219 );
1220 const providerThread = yield* runtime.ensureThread({
1221 threadId: THREAD_ID,
1222 modelSelection: modelSelection("default"),
1223 runtimePolicy,
1224 });
1225 yield* fake.emit({ type: "agent_start" });
1226 const offer = yield* Queue.take(offers);
1227 yield* fake.emit({
1228 type: "extension_ui_request",
1229 id: "wake-dialog",
1230 method: "confirm",
1231 title: "Continue background work?",
1232 });
1233 const dialog = yield* takeEvent((event) => event.type === "runtime_request.updated");
1234 assert.isTrue(dialog.type === "runtime_request.updated");
1235 if (dialog.type !== "runtime_request.updated") return;
1236 assert.isNull(dialog.runtimeRequest.providerTurnId);
1237 if (retirement === "archive") yield* offer.clearIfCurrent!();
1238 else if (retirement === "stop") {
1239 yield* runtime.interruptTurn({
1240 providerThread,
1241 providerTurnId: ProviderTurnId.make("settled-turn"),
1242 requestRuntimeRestart: true,
1243 });
1244 } else yield* fake.closeStdout;
1245 let requestCancelled = false;
1246 let nodeCancelled = false;
1247 let itemCancelled = false;
1248 yield* takeEvent((event) => {
1249 if (
1250 event.type === "runtime_request.updated" &&
1251 event.runtimeRequest.id === dialog.runtimeRequest.id
1252 )
1253 requestCancelled = event.runtimeRequest.status === "cancelled";
1254 if (event.type === "node.updated" && event.node.id === dialog.runtimeRequest.nodeId)
1255 nodeCancelled = event.node.status === "cancelled";
1256 if (
1257 event.type === "turn_item.updated" &&
1258 event.turnItem.nodeId === dialog.runtimeRequest.nodeId
1259 )
1260 itemCancelled = event.turnItem.status === "cancelled";
1261 return (
1262 event.type === "provider_session.updated" &&
1263 event.providerSession.status ===
1264 (retirement === "transport death" ? "error" : "stopped")
1265 );
1266 });
1267 assert.isTrue(requestCancelled, "Retired wake left a pending runtime request");
1268 assert.isTrue(nodeCancelled);
1269 assert.isTrue(itemCancelled);
1270 assert.isFalse(yield* runtime.hasPendingBackgroundWork!);
1271 let dispatched = false;
1272 yield* offer.dispatchIfCurrent!(
1273 Effect.sync(() => {
1274 dispatched = true;
1275 }),
1276 );
1277 assert.isFalse(dispatched);
1278 }).pipe(Effect.scoped, Effect.provide(layerTest)),
1279 );
1280
4641281 it.effect("stops provider-initiated work that has no T3 turn owner", () =>
4651282 Effect.gen(function* () {
4661283 const fake = yield* makeFakePi;
apps/server/src/orchestration-v2/Adapters/PiAdapterV2.ts+299-47
......@@ -48,6 +48,7 @@ import {
4848} from "@t3tools/contracts";
4949import * as Cause from "effect/Cause";
5050import * as DateTime from "effect/DateTime";
51import * as Deferred from "effect/Deferred";
5152import * as Option from "effect/Option";
5253import * as Duration from "effect/Duration";
5354import * as Effect from "effect/Effect";
......@@ -72,6 +73,7 @@ import {
7273import { mergeProviderInstanceEnvironment } from "../../provider/ProviderInstanceEnvironment.ts";
7374import * as IdAllocator from "../IdAllocator.ts";
7475import * as ProviderAdapter from "../ProviderAdapter.ts";
76import * as ProviderContinuationRequests from "../ProviderContinuationRequests.ts";
7577import {
7678 ProviderAdapterDriverCreateError,
7779 type ProviderAdapterDriver,
......@@ -228,6 +230,11 @@ export interface PiAdapterV2Options {
228230 readonly fileSystem: FileSystem.FileSystem;
229231 readonly idAllocator: IdAllocator.IdAllocatorV2["Service"];
230232 readonly serverConfig: ServerConfig.ServerConfig["Service"];
233 readonly continuationRequests?: {
234 readonly offer: (
235 request: ProviderContinuationRequests.ProviderContinuationRequest,
236 ) => Effect.Effect<void>;
237 };
231238}
232239
233240/** Concatenate the `text` fields of a Pi content-block array. */
......@@ -313,6 +320,8 @@ interface ActivePiTurn {
313320 * the turn instead.
314321 */
315322 sawAgentActivity: boolean;
323 /** A user prompt joined provider-native work whose settlement may already be queued. */
324 adoptedWake: boolean;
316325 /** Only slash-command prompts can complete without starting an agent run. */
317326 readonly promptMayBeCommandOnly: boolean;
318327 /** Pi reports context as unknown immediately after compaction; keep its estimate for the meter. */
......@@ -329,6 +338,8 @@ interface ActivePiTurn {
329338 activeCompaction: PiCompactionState | null;
330339 activeProviderRetry: PiProviderRetryState | null;
331340 failure: ReturnType<typeof makeProviderFailure> | null;
341 /** Native retry can recover its work without accepting a rejected joining prompt. */
342 rejectedPromptFailure: ReturnType<typeof makeProviderFailure> | null;
332343 /** Session-tree refs read just before Stop terminates Pi, when no read is possible later. */
333344 stopTreeRefs?: PiTurnTreeRefs | null;
334345}
......@@ -366,6 +377,11 @@ interface PiThreadState {
366377 activeTurn: ActivePiTurn | null;
367378}
368379
380interface PiWake {
381 readonly state: PiThreadState;
382 readonly events: Array<PiRpcRecord>;
383}
384
369385// ── adapter ───────────────────────────────────────────────────
370386
371387export function makePiAdapterV2(
......@@ -476,10 +492,26 @@ export function makePiAdapterV2(
476492 // Keep that intent beyond turn finalization so the later stdout close is
477493 // not mistaken for an unexpected transport failure.
478494 let stopRequested = false;
479 // Pi extensions can trigger an agent turn after the owning T3 turn has
480 // settled. Until orchestration has a first-class provider-initiated run,
481 // stop that runtime before it can execute tools without a timeline owner.
495 // Native work is held until an orchestrator-owned continuation or user
496 // turn attaches. The guard remains for sessions with no routable owner.
482497 let unsolicitedActivityDetected = false;
498 let pendingWake: PiWake | null = null;
499 let rollbackBarrier: Deferred.Deferred<
500 void,
501 ProviderAdapter.ProviderAdapterProtocolError
502 > | null = null;
503 let closed = false;
504 yield* Effect.addFinalizer(() =>
505 Effect.gen(function* () {
506 closed = true;
507 pendingWake = null;
508 if (rollbackBarrier !== null)
509 yield* Deferred.fail(
510 rollbackBarrier,
511 protocolError("Pi session closed during rollback"),
512 );
513 }),
514 );
483515 let appliedModel: string | null = null;
484516 let appliedThinking: string | null = null;
485517 /** Last thread title synced into pi's session name (`/resume` listing). */
......@@ -1437,7 +1469,7 @@ export function makePiAdapterV2(
14371469 const tokenUsage = readUsage
14381470 ? yield* readTokenUsage(turn.latestCompactionAfterTokens, completedAt)
14391471 : undefined;
1440 const failure = turn.interrupted ? null : turn.failure;
1472 const failure = turn.interrupted ? null : (turn.rejectedPromptFailure ?? turn.failure);
14411473 yield* emit({
14421474 type: "provider_turn.updated",
14431475 driver: PI_PROVIDER,
......@@ -1554,11 +1586,74 @@ export function makePiAdapterV2(
15541586 );
15551587 };
15561588
1589 const stopPendingWake = Effect.fnUntraced(function* (wake: PiWake) {
1590 if (pendingWake !== wake || threadState !== wake.state || closed) return;
1591 pendingWake = null;
1592 stopRequested = true;
1593 yield* cancelPendingPrompts(yield* DateTime.now);
1594 yield* connection.terminate;
1595 });
1596
15571597 const handleSessionEvent = Effect.fnUntraced(function* (event: PiRpcRecord) {
1598 // Model-selection hooks can await a dialog before startTurn can claim
1599 // a wake. Keep interactive requests session-scoped until a turn owns
1600 // them; buffer native output and notifications for that turn.
1601 const interactiveUiRequest =
1602 event["type"] === "extension_ui_request" &&
1603 ["select", "confirm", "input", "editor"].includes(String(event["method"]));
1604 if (
1605 threadState?.activeTurn === null &&
1606 pendingWake !== null &&
1607 !interactiveUiRequest &&
1608 !String(event["type"]).startsWith("t3.") &&
1609 event["type"] !== "response"
1610 ) {
1611 pendingWake.events.push(event);
1612 if (event["type"] === "agent_start") yield* updateProviderSession("running", null);
1613 if (event["type"] === "agent_settled") yield* updateProviderSession("ready", null);
1614 return;
1615 }
15581616 const state = threadState;
15591617 const turn = state?.activeTurn ?? null;
15601618 switch (event["type"]) {
15611619 case "agent_start": {
1620 if (
1621 turn === null &&
1622 state?.providerThread.appThreadId != null &&
1623 options.continuationRequests !== undefined &&
1624 !closed &&
1625 !stopRequested
1626 ) {
1627 const wake: PiWake = { state, events: [event] };
1628 pendingWake = wake;
1629 if (rollbackBarrier !== null) {
1630 yield* stopPendingWake(wake);
1631 return;
1632 }
1633 yield* updateProviderSession("running", null);
1634 yield* options.continuationRequests
1635 .offer({
1636 threadId: state.providerThread.appThreadId,
1637 providerThreadId: state.providerThread.id,
1638 driver: PI_PROVIDER,
1639 detail: null,
1640 delivery: "adapter_buffered",
1641 dispatchIfCurrent: (dispatch) =>
1642 sessionEventPermit.withPermits(1)(
1643 Effect.suspend(() =>
1644 pendingWake === wake && threadState === state && !closed
1645 ? dispatch.pipe(
1646 Effect.map(Option.some),
1647 Effect.tapError(() => stopPendingWake(wake)),
1648 )
1649 : Effect.succeed(Option.none()),
1650 ),
1651 ),
1652 clearIfCurrent: () => sessionEventPermit.withPermits(1)(stopPendingWake(wake)),
1653 })
1654 .pipe(Effect.tapError(() => stopPendingWake(wake)));
1655 return;
1656 }
15621657 if (turn === null) {
15631658 unsolicitedActivityDetected = true;
15641659 yield* updateProviderSession("error", PI_UNSOLICITED_ACTIVITY_ERROR);
......@@ -1787,6 +1882,17 @@ export function makePiAdapterV2(
17871882 if (turn !== null) {
17881883 turn.settleWhenIdle = true;
17891884 turn.settleProbeGeneration += 1;
1885 // A wake can settle while the newly joined user prompt is still
1886 // expanding. Its acknowledgement opens the idle-probe barrier.
1887 if (
1888 turn.adoptedWake &&
1889 pendingPromptResponses.some(
1890 (pending) =>
1891 pending.providerTurnId === turn.providerTurn.id &&
1892 pending.kind === "turn_start",
1893 )
1894 )
1895 return;
17901896 yield* scheduleSettleProbe(turn, true);
17911897 }
17921898 return;
......@@ -1845,10 +1951,11 @@ export function makePiAdapterV2(
18451951 if (
18461952 pendingPrompt?.kind === "turn_start" &&
18471953 responseTurn !== null &&
1848 responseTurn.promptMayBeCommandOnly &&
1849 !responseTurn.sawAgentActivity
1954 ((responseTurn.promptMayBeCommandOnly && !responseTurn.sawAgentActivity) ||
1955 (responseTurn.adoptedWake &&
1956 (responseTurn.settleWhenIdle || responseTurn.promptMayBeCommandOnly)))
18501957 ) {
1851 yield* scheduleSettleProbe(responseTurn);
1958 yield* scheduleSettleProbe(responseTurn, responseTurn.adoptedWake);
18521959 }
18531960 return;
18541961 }
......@@ -1869,11 +1976,18 @@ export function makePiAdapterV2(
18691976 ? turn
18701977 : null;
18711978 if (failedTurn !== null) {
1872 failedTurn.failure = makeProviderFailure({
1979 const failure = makeProviderFailure({
18731980 message: recordString(event, "error") ?? "Pi rejected the prompt.",
18741981 class: "provider_error",
18751982 });
1876 if (state !== null) yield* finalizeTurn(state);
1983 if (failedTurn.adoptedWake && command === "prompt") {
1984 failedTurn.rejectedPromptFailure = failure;
1985 failedTurn.settleWhenIdle = true;
1986 yield* scheduleSettleProbe(failedTurn, true);
1987 } else {
1988 failedTurn.failure = failure;
1989 if (state !== null) yield* finalizeTurn(state);
1990 }
18771991 }
18781992 return;
18791993 }
......@@ -1886,6 +2000,11 @@ export function makePiAdapterV2(
18862000 }
18872001 return;
18882002 }
2003 case "t3.rollback_barrier": {
2004 if (rollbackBarrier !== null && event["barrier"] === rollbackBarrier)
2005 yield* Deferred.succeed(rollbackBarrier, undefined);
2006 return;
2007 }
18892008 case "t3.settle_probe": {
18902009 // New work increments the generation before the pump can consume
18912010 // a stale idle snapshot, so only a current snapshot may settle.
......@@ -1946,6 +2065,13 @@ export function makePiAdapterV2(
19462065 // closes the provider stream cleanly; only an unexpected death
19472066 // is surfaced as an event-stream failure.
19482067 const state = threadState;
2068 if (rollbackBarrier !== null)
2069 yield* Deferred.fail(
2070 rollbackBarrier,
2071 protocolError("Pi transport closed during rollback"),
2072 );
2073 const hadPendingWake = pendingWake !== null;
2074 pendingWake = null;
19492075 const interrupted = state?.activeTurn?.interrupted === true;
19502076 if (state?.activeTurn != null) {
19512077 state.activeTurn.failure = interrupted
......@@ -1956,6 +2082,8 @@ export function makePiAdapterV2(
19562082 class: "transport_error",
19572083 });
19582084 yield* finalizeTurn(state, false);
2085 } else if (hadPendingWake) {
2086 yield* cancelPendingPrompts(yield* DateTime.now);
19592087 }
19602088 if (unsolicitedActivityDetected) {
19612089 yield* updateProviderSession("error", PI_UNSOLICITED_ACTIVITY_ERROR);
......@@ -1999,9 +2127,26 @@ export function makePiAdapterV2(
19992127 threadInput: ProviderAdapter.ProviderAdapterV2EnsureThreadInput,
20002128 publish = true,
20012129 ) {
2130 if (rollbackBarrier !== null)
2131 return yield* protocolError("Cannot register a Pi thread during rollback");
20022132 if (threadState !== null && threadState.activeTurn !== null) {
20032133 return yield* protocolError("Cannot register a Pi thread while a turn is active");
20042134 }
2135 if (pendingWake !== null) {
2136 const state = threadState;
2137 if (
2138 state !== null &&
2139 state === pendingWake.state &&
2140 state.providerThread.appThreadId === threadInput.threadId &&
2141 (threadInput.existingProviderThread === undefined ||
2142 threadInput.existingProviderThread.nativeThreadRef?.nativeId ===
2143 state.providerThread.nativeThreadRef?.nativeId)
2144 )
2145 return state.providerThread;
2146 return yield* protocolError(
2147 "Cannot switch Pi threads while native work is awaiting a turn",
2148 );
2149 }
20052150 const existing = threadInput.existingProviderThread;
20062151 const resumeId = existing?.nativeThreadRef?.nativeId;
20072152 const needsNewSession = resumeId == null && registrationAttempted;
......@@ -2219,6 +2364,11 @@ export function makePiAdapterV2(
22192364 return sessionEntity;
22202365 },
22212366 events: Stream.fromQueue(events),
2367 hasPendingBackgroundWork: Effect.sync(() => pendingWake !== null),
2368 hasPendingBackgroundWorkForThread: (providerThread) =>
2369 Effect.sync(
2370 () => pendingWake !== null && pendingWake.state.providerThread.id === providerThread.id,
2371 ),
22222372 getModelContextWindow: (selection) => {
22232373 if (selection.instanceId !== options.instanceId) return undefined;
22242374 const slug =
......@@ -2267,6 +2417,8 @@ export function makePiAdapterV2(
22672417 }),
22682418 startTurn: (turnInput) =>
22692419 Effect.gen(function* () {
2420 if (rollbackBarrier !== null)
2421 return yield* protocolError("Cannot start a Pi turn during rollback");
22702422 const state = threadState;
22712423 if (state === null) {
22722424 return yield* protocolError("Pi session has no registered thread");
......@@ -2284,8 +2436,10 @@ export function makePiAdapterV2(
22842436 }
22852437 // The orchestrator adopts a fork under its already-allocated row.
22862438 // Future session updates must retain that authoritative identity.
2287 state.providerThread = turnInput.providerThread;
2288 yield* applySelection(turnInput.modelSelection);
2439 const continuation =
2440 turnInput.message.createdBy === "agent" &&
2441 turnInput.message.creationSource === "provider";
2442 if (!continuation) yield* applySelection(turnInput.modelSelection);
22892443 // Mirror the thread title into pi's session name so the session
22902444 // stays identifiable in pi's own /resume listing. Best-effort:
22912445 // naming must never block a turn.
......@@ -2307,9 +2461,11 @@ export function makePiAdapterV2(
23072461 // extension's before_agent_start system-prompt hook, never by
23082462 // wrapping the user text: a wrapped first message would no
23092463 // longer start with "/" and slash commands would stop expanding.
2310 const compactCommand = parsePiCompactCommand(turnInput.message.text);
2464 const compactCommand = continuation
2465 ? null
2466 : parsePiCompactCommand(turnInput.message.text);
23112467 const payload =
2312 compactCommand === null
2468 !continuation && compactCommand === null
23132469 ? yield* resolvePromptPayload(turnInput.message.text, turnInput.message.attachments)
23142470 : null;
23152471 const startedAt = yield* DateTime.now;
......@@ -2340,6 +2496,7 @@ export function makePiAdapterV2(
23402496 toolStartedAt: new Map(),
23412497 interrupted: false,
23422498 sawAgentActivity: false,
2499 adoptedWake: false,
23432500 promptMayBeCommandOnly:
23442501 compactCommand !== null || (payload?.message.trimStart().startsWith("/") ?? false),
23452502 latestCompactionAfterTokens: null,
......@@ -2351,13 +2508,32 @@ export function makePiAdapterV2(
23512508 activeCompaction: null,
23522509 activeProviderRetry: null,
23532510 failure: null,
2511 rejectedPromptFailure: null,
23542512 };
23552513 // Only the install/send/start-event boundary excludes the event
23562514 // pump. Earlier correlated requests must leave the pump free so
23572515 // project trust, login, and session-switch dialogs can be shown
23582516 // and answered instead of deadlocking the caller.
23592517 yield* Effect.gen(function* () {
2518 if (
2519 threadState !== state ||
2520 state.activeTurn !== null ||
2521 closed ||
2522 stopRequested ||
2523 rollbackBarrier !== null ||
2524 state.providerThread.nativeThreadRef?.nativeId !==
2525 turnInput.providerThread.nativeThreadRef?.nativeId
2526 ) {
2527 return yield* protocolError("Pi turn ownership changed before it could start");
2528 }
2529 const wake = pendingWake?.state === state ? pendingWake : null;
2530 if (wake !== null && wake.state.providerThread.appThreadId !== turnInput.threadId) {
2531 return yield* protocolError("Pi native wake belongs to a different app thread");
2532 }
2533 pendingWake = null;
2534 state.providerThread = turnInput.providerThread;
23602535 state.activeTurn = activeTurn;
2536 activeTurn.adoptedWake = wake !== null && !continuation;
23612537 if (compactCommand !== null) {
23622538 yield* connection.send(compactRpcRecord(compactCommand));
23632539 pendingCompactResponses.push({
......@@ -2368,6 +2544,7 @@ export function makePiAdapterV2(
23682544 yield* connection.send({
23692545 type: "prompt",
23702546 message: payload.message,
2547 ...(wake === null ? {} : { streamingBehavior: "steer" }),
23712548 ...(payload.images.length === 0 ? {} : { images: payload.images }),
23722549 });
23732550 pendingPromptResponses.push({
......@@ -2387,6 +2564,18 @@ export function makePiAdapterV2(
23872564 lastRunOrdinal: turnInput.runOrdinal,
23882565 });
23892566 yield* updateProviderSession("running", null);
2567 if (wake !== null) {
2568 for (const event of wake.events) {
2569 if (!continuation && event["type"] === "agent_settled") {
2570 activeTurn.settleWhenIdle = true;
2571 continue;
2572 }
2573 yield* handleSessionEvent(event);
2574 }
2575 } else if (continuation) {
2576 // A user turn, teardown, or newer binding already took the wake.
2577 yield* finalizeTurn(state);
2578 }
23902579 if (outOfTurnExtensionErrors.length > 0) {
23912580 yield* Queue.offer(connection.events, { type: "t3.flush_extension_errors" });
23922581 }
......@@ -2478,7 +2667,11 @@ export function makePiAdapterV2(
24782667 const turn = threadState?.activeTurn ?? null;
24792668 // Stop on a settled turn: Pi runs nothing between prompts, so
24802669 // nothing of that turn is left to stop.
2481 if (turn === null && interruptInput.requestRuntimeRestart === true) return;
2670 if (turn === null && interruptInput.requestRuntimeRestart === true) {
2671 if (pendingWake !== null)
2672 yield* sessionEventPermit.withPermits(1)(stopPendingWake(pendingWake));
2673 return;
2674 }
24822675 if (turn === null || turn.providerTurn.id !== interruptInput.providerTurnId) {
24832676 return yield* protocolError(`Pi turn ${interruptInput.providerTurnId} is not active`);
24842677 }
......@@ -2662,6 +2855,9 @@ export function makePiAdapterV2(
26622855 if (state.activeTurn !== null) {
26632856 return yield* protocolError("Cannot roll back while a Pi turn is active");
26642857 }
2858 if (pendingWake !== null) {
2859 return yield* protocolError("Cannot roll back while native Pi work awaits a turn");
2860 }
26652861 // `fork(entryId)` re-roots the active branch before that user
26662862 // message, so the rollback boundary is the first user entry of
26672863 // the earliest turn being discarded.
......@@ -2673,42 +2869,94 @@ export function makePiAdapterV2(
26732869 if (forkEntryId === undefined) {
26742870 return yield* protocolError("Pi rollback target has no captured session-tree entry");
26752871 }
2676 const forkData = yield* lifecycleRequest({ type: "fork", entryId: forkEntryId });
2677 if (recordField(forkData, "cancelled") === true) {
2678 return yield* protocolError("A Pi extension cancelled the session fork");
2679 }
2680 // Pi fork replaces the session file, including for rollback. Persist
2681 // its new identity before any later request can fail or restart.
2682 // An interrupted read leaves the identity just as unknown as a failed one.
2683 const forkState = yield* request({ type: "get_state" }).pipe(
2684 Effect.onError(() =>
2685 Effect.sync(() => {
2686 threadState = null;
2872 return yield* Effect.acquireUseRelease(
2873 sessionEventPermit.withPermits(1)(
2874 Effect.gen(function* () {
2875 if (
2876 rollbackBarrier !== null ||
2877 threadState !== state ||
2878 state.activeTurn !== null ||
2879 pendingWake !== null ||
2880 closed ||
2881 stopRequested
2882 )
2883 return yield* protocolError(
2884 "Pi rollback ownership changed before its fork could start",
2885 );
2886 const barrier = yield* Deferred.make<
2887 void,
2888 ProviderAdapter.ProviderAdapterProtocolError
2889 >();
2890 rollbackBarrier = barrier;
2891 return barrier;
26872892 }),
26882893 ),
2894 (barrier) =>
2895 Effect.gen(function* () {
2896 const forkData = yield* lifecycleRequest({ type: "fork", entryId: forkEntryId });
2897 if (recordField(forkData, "cancelled") === true)
2898 return yield* protocolError("A Pi extension cancelled the session fork");
2899 const forkState = yield* request({ type: "get_state" }).pipe(
2900 Effect.onError(() =>
2901 Effect.sync(() => {
2902 threadState = null;
2903 }),
2904 ),
2905 );
2906 const forkSessionFile = recordString(forkState, "sessionFile");
2907 if (forkSessionFile === undefined) {
2908 threadState = null;
2909 return yield* protocolError("Pi fork did not return a persisted session file");
2910 }
2911 const entriesData = yield* request({ type: "get_entries" }).pipe(
2912 Effect.orElseSucceed(() => undefined),
2913 );
2914 const leafId = recordString(entriesData, "leafId") ?? null;
2915 // RPC responses bypass the pump. Drain preceding native events
2916 // before deciding whether this fork can replace their binding.
2917 const queued = yield* Queue.offer(connection.events, {
2918 type: "t3.rollback_barrier",
2919 barrier,
2920 });
2921 if (!queued) return yield* protocolError("Pi transport closed during rollback");
2922 yield* Deferred.await(barrier);
2923 return yield* sessionEventPermit.withPermits(1)(
2924 Effect.gen(function* () {
2925 if (pendingWake !== null) {
2926 yield* stopPendingWake(pendingWake);
2927 return yield* protocolError(
2928 "Native Pi work started while rollback was changing the session",
2929 );
2930 }
2931 if (
2932 threadState !== state ||
2933 state.activeTurn !== null ||
2934 closed ||
2935 stopRequested
2936 )
2937 return yield* protocolError(
2938 "Pi rollback ownership changed during its fork",
2939 );
2940 lastNativeThreadId = forkSessionFile;
2941 appliedModel = null;
2942 appliedThinking = null;
2943 appliedSessionName = null;
2944 lastKnownLeaf = leafId;
2945 leafCursorStale = entriesData === undefined;
2946 yield* updateProviderThread(state, {
2947 nativeThreadRef: providerRef(forkSessionFile),
2948 nativeConversationHeadRef: leafId === null ? null : providerRef(leafId),
2949 });
2950 rollbackBarrier = null;
2951 return piThreadSnapshot(state.providerThread);
2952 }),
2953 );
2954 }).pipe(Effect.onInterrupt(() => connection.terminate)),
2955 (barrier) =>
2956 Effect.sync(() => {
2957 if (rollbackBarrier === barrier) rollbackBarrier = null;
2958 }),
26892959 );
2690 const forkSessionFile = recordString(forkState, "sessionFile");
2691 if (forkSessionFile === undefined) {
2692 threadState = null;
2693 return yield* protocolError("Pi fork did not return a persisted session file");
2694 }
2695 lastNativeThreadId = forkSessionFile;
2696 appliedModel = null;
2697 appliedThinking = null;
2698 appliedSessionName = null;
2699 yield* updateProviderThread(state, { nativeThreadRef: providerRef(forkSessionFile) });
2700 const entriesData = yield* request({ type: "get_entries" }).pipe(
2701 Effect.orElseSucceed(() => undefined),
2702 );
2703 const leafId = recordString(entriesData, "leafId") ?? null;
2704 lastKnownLeaf = leafId;
2705 // The fork re-baselined the tree, so the cursor is trustworthy
2706 // again unless this listing itself failed.
2707 leafCursorStale = entriesData === undefined;
2708 yield* updateProviderThread(state, {
2709 nativeConversationHeadRef: leafId === null ? null : providerRef(leafId),
2710 });
2711 return piThreadSnapshot(state.providerThread);
27122960 }).pipe(
27132961 Effect.mapError(
27142962 (cause) =>
......@@ -2971,6 +3219,7 @@ export const PiAdapterV2Driver: ProviderAdapterDriver<PiSettings, PiAdapterV2Dri
29713219 const fileSystem = yield* FileSystem.FileSystem;
29723220 const idAllocator = yield* IdAllocator.IdAllocatorV2;
29733221 const serverConfig = yield* ServerConfig.ServerConfig;
3222 const continuationRequests = yield* ProviderContinuationRequests.ProviderContinuationRequests;
29743223 return makePiAdapterV2({
29753224 instanceId: input.instanceId,
29763225 settings: { ...input.config, enabled: input.enabled },
......@@ -2979,6 +3228,7 @@ export const PiAdapterV2Driver: ProviderAdapterDriver<PiSettings, PiAdapterV2Dri
29793228 fileSystem,
29803229 idAllocator,
29813230 serverConfig,
3231 continuationRequests,
29823232 });
29833233 },
29843234 (effect, input) =>
......@@ -3005,6 +3255,7 @@ const layer: Layer.Layer<ProviderAdapter.ProviderAdapterV2, never, PiAdapterV2Dr
30053255 const fileSystem = yield* FileSystem.FileSystem;
30063256 const idAllocator = yield* IdAllocator.IdAllocatorV2;
30073257 const serverConfig = yield* ServerConfig.ServerConfig;
3258 const continuationRequests = yield* ProviderContinuationRequests.ProviderContinuationRequests;
30083259 return makePiAdapterV2({
30093260 instanceId: PI_DEFAULT_INSTANCE_ID,
30103261 settings: DEFAULT_PI_SETTINGS,
......@@ -3013,6 +3264,7 @@ const layer: Layer.Layer<ProviderAdapter.ProviderAdapterV2, never, PiAdapterV2Dr
30133264 fileSystem,
30143265 idAllocator,
30153266 serverConfig,
3267 continuationRequests,
30163268 });
30173269 }),
30183270 );