authorgravatar for 71498452+StiensWout@users.noreply.github.comWout Stiens <71498452+StiensWout@users.noreply.github.com> 2026-10-11 00:49:04+02:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2026-10-10 17:49:04-05:00
log6e9c5a851c92c52e2f4a8a80c655e60db133f9ac
tree704f8a986c17fc684e4e061da490ff3bd8205a5c
parent62da1cd48bb43def11a79cd6aa033dbc478e4dfd
signature Signed by PGP key B5690EEEBB952194

fix(pi): expire and cancel extension approvals (#17840)


4 files changed, 172 insertions(+), 6 deletions(-)

packages/provider-pi/src/server/adapter.test.ts+104
......@@ -2912,6 +2912,110 @@ describe("PiAdapterV2", () => {
29122912 }).pipe(Effect.scoped, Effect.provide(layerTest)),
29132913 );
29142914
2915 it.effect("expires native dialogs and rejects stale answers without caching approval", () =>
2916 Effect.gen(function* () {
2917 const fake = yield* makeFakePi;
2918 const { runtime, takeEvent } = yield* openRuntime(fake);
2919 yield* runtime.ensureThread({
2920 threadId: THREAD_ID,
2921 modelSelection: modelSelection("default"),
2922 runtimePolicy,
2923 });
2924 const dialog = {
2925 type: "extension_ui_request",
2926 method: "confirm",
2927 title: "Allow extension?",
2928 message: "Approve this operation",
2929 };
2930 yield* fake.emit({ ...dialog, id: "timed-dialog", timeout: 1000 });
2931 const pending = yield* takeEvent(
2932 (event) => event.type === "turn_item.updated" && event.turnItem.type === "approval_request",
2933 );
2934 assert.isTrue(
2935 pending.type === "turn_item.updated" && pending.turnItem.type === "approval_request",
2936 );
2937 if (pending.type !== "turn_item.updated" || pending.turnItem.type !== "approval_request")
2938 return;
2939 const requestId = pending.turnItem.requestId;
2940 yield* TestClock.adjust(Duration.seconds(1));
2941 yield* takeEvent(
2942 (event) =>
2943 event.type === "runtime_request.updated" &&
2944 event.runtimeRequest.id === requestId &&
2945 event.runtimeRequest.status === "cancelled",
2946 );
2947 assert.equal((yield* fake.takeRequest("extension_ui_response"))["cancelled"], true);
2948 const late = yield* Effect.result(
2949 runtime.respondToRuntimeRequest({
2950 requestId: pending.turnItem.requestId,
2951 decision: "acceptForSession",
2952 }),
2953 );
2954 assert.equal(late._tag, "Failure");
2955 yield* fake.emit({ ...dialog, id: "fresh-dialog", timeout: 1000 });
2956 const fresh = yield* takeEvent(
2957 (event) =>
2958 event.type === "runtime_request.updated" &&
2959 event.runtimeRequest.status === "pending" &&
2960 event.runtimeRequest.nativeRequestRef?.nativeId === "fresh-dialog",
2961 );
2962 if (fresh.type !== "runtime_request.updated") return;
2963 yield* runtime.respondToRuntimeRequest({
2964 requestId: fresh.runtimeRequest.id,
2965 decision: "accept",
2966 });
2967 assert.equal((yield* fake.takeRequest("extension_ui_response"))["id"], "fresh-dialog");
2968 // The old expiry must not cancel an already answered request, and zero disables expiry.
2969 yield* fake.emit({ ...dialog, id: "zero-timeout", timeout: 0 });
2970 const zero = yield* takeEvent(
2971 (event) =>
2972 event.type === "runtime_request.updated" &&
2973 event.runtimeRequest.status === "pending" &&
2974 event.runtimeRequest.nativeRequestRef?.nativeId === "zero-timeout",
2975 );
2976 yield* TestClock.adjust(Duration.seconds(60));
2977 if (zero.type !== "runtime_request.updated") return;
2978 yield* runtime.respondToRuntimeRequest({
2979 requestId: zero.runtimeRequest.id,
2980 decision: "decline",
2981 });
2982 const response = yield* fake.takeRequest("extension_ui_response");
2983 assert.equal(response["id"], "zero-timeout");
2984 assert.equal(response["confirmed"], false);
2985 }).pipe(Effect.scoped, Effect.provide(layerTest)),
2986 );
2987
2988 it.effect("cancels startup dialogs when Pi exits before a turn exists", () =>
2989 Effect.gen(function* () {
2990 const fake = yield* makeFakePi;
2991 const { runtime, takeEvent } = yield* openRuntime(fake);
2992 yield* runtime.ensureThread({
2993 threadId: THREAD_ID,
2994 modelSelection: modelSelection("default"),
2995 runtimePolicy,
2996 });
2997 yield* fake.emit({
2998 type: "extension_ui_request",
2999 id: "startup",
3000 method: "confirm",
3001 title: "Trust project?",
3002 });
3003 const pending = yield* takeEvent(
3004 (event) =>
3005 event.type === "runtime_request.updated" && event.runtimeRequest.status === "pending",
3006 );
3007 if (pending.type !== "runtime_request.updated") return;
3008 yield* fake.closeStdout;
3009 const cancelled = yield* takeEvent(
3010 (event) =>
3011 event.type === "runtime_request.updated" &&
3012 event.runtimeRequest.id === pending.runtimeRequest.id &&
3013 event.runtimeRequest.status === "cancelled",
3014 );
3015 assert.isTrue(cancelled.type === "runtime_request.updated");
3016 }).pipe(Effect.scoped, Effect.provide(layerTest)),
3017 );
3018
29153019 it.effect("remembers session approvals only for identical confirmation content", () =>
29163020 Effect.gen(function* () {
29173021 const fake = yield* makeFakePi;
packages/provider-pi/src/server/adapter.ts+53-4
......@@ -56,6 +56,7 @@ import * as Option from "effect/Option";
5656import * as Duration from "effect/Duration";
5757import * as Effect from "effect/Effect";
5858import * as FileSystem from "effect/FileSystem";
59import * as Fiber from "effect/Fiber";
5960import * as Layer from "effect/Layer";
6061import * as Queue from "effect/Queue";
6162import * as Schema from "effect/Schema";
......@@ -456,6 +457,8 @@ interface PendingPiPrompt {
456457 readonly method: "select" | "confirm" | "input" | "editor";
457458 readonly questionId: string;
458459 readonly approvalKey: string;
460 readonly expiresAt: number | null;
461 expirationFiber: Fiber.Fiber<void> | null;
459462 runtimeRequest: OrchestrationV2RuntimeRequest;
460463 readonly node: OrchestrationV2ExecutionNode;
461464 readonly turnItem: OrchestrationV2TurnItem;
......@@ -1270,6 +1273,7 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
12701273
12711274 const cancelPrompt = (pending: PendingPiPrompt, resolvedAt: DateTime.Utc) =>
12721275 Effect.gen(function* () {
1276 if (pending.expirationFiber !== null) yield* Fiber.interrupt(pending.expirationFiber);
12731277 yield* connection
12741278 .send({
12751279 type: "extension_ui_response",
......@@ -1368,6 +1372,12 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
13681372 const state = threadState;
13691373 const turn = state?.activeTurn ?? null;
13701374 const createdAt = yield* DateTime.now;
1375 const timeout = recordNumber(event, "timeout");
1376 // Pi treats zero as no timeout and negative durations as immediate expiry.
1377 const expiresAt =
1378 timeout === undefined || timeout === 0
1379 ? null
1380 : DateTime.toEpochMillis(createdAt) + Math.max(0, timeout);
13711381 const requestId = yield* idAllocator.allocate.runtimeRequest({
13721382 driver: PI_PROVIDER,
13731383 ...(turn === null ? {} : { providerTurnId: turn.providerTurn.id }),
......@@ -1440,15 +1450,18 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
14401450 requestId,
14411451 questions: [piQuestion(nativeRequestId, method, title, event)],
14421452 };
1443 pendingPrompts.set(String(requestId), {
1453 const pending: PendingPiPrompt = {
14441454 nativeRequestId,
14451455 method,
14461456 questionId: nativeRequestId,
14471457 approvalKey,
1458 expiresAt,
1459 expirationFiber: null,
14481460 runtimeRequest,
14491461 node,
14501462 turnItem,
1451 });
1463 };
1464 pendingPrompts.set(String(requestId), pending);
14521465 yield* emit({
14531466 type: "runtime_request.updated",
14541467 driver: PI_PROVIDER,
......@@ -1457,6 +1470,22 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
14571470 });
14581471 yield* emit({ type: "node.updated", driver: PI_PROVIDER, node });
14591472 yield* emit({ type: "turn_item.updated", driver: PI_PROVIDER, turnItem });
1473 if (expiresAt !== null) {
1474 const now = yield* DateTime.now;
1475 pending.expirationFiber = yield* Effect.sleep(
1476 Duration.millis(Math.max(0, expiresAt - DateTime.toEpochMillis(now))),
1477 ).pipe(
1478 Effect.andThen(
1479 Queue.offer(connection.events, {
1480 type: "t3.ui_request_expired",
1481 requestId: String(requestId),
1482 pending,
1483 }),
1484 ),
1485 Effect.asVoid,
1486 Effect.forkIn(scope),
1487 );
1488 }
14601489 });
14611490
14621491 const emitExtensionError = Effect.fnUntraced(function* (event: PiRpcRecord) {
......@@ -2132,6 +2161,15 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
21322161 }
21332162 return;
21342163 }
2164 case "t3.ui_request_expired": {
2165 const requestId = recordString(event, "requestId");
2166 if (requestId === undefined) return;
2167 const pending = pendingPrompts.get(requestId);
2168 if (pending === undefined || pending !== event["pending"]) return;
2169 pendingPrompts.delete(requestId);
2170 yield* cancelPrompt(pending, yield* DateTime.now);
2171 return;
2172 }
21352173 case "t3.flush_extension_errors": {
21362174 // Startup extension failures are informational and do not block
21372175 // Pi, so attach them to the next real turn instead of creating a
......@@ -2211,7 +2249,6 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
22112249 rollbackBarrier,
22122250 protocolError("Pi transport closed during rollback"),
22132251 );
2214 const hadPendingWake = pendingWake !== null;
22152252 pendingWake = null;
22162253 const interrupted = state?.activeTurn?.interrupted === true;
22172254 if (state?.activeTurn != null) {
......@@ -2223,7 +2260,7 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
22232260 class: "transport_error",
22242261 });
22252262 yield* finalizeTurn(state, false);
2226 } else if (hadPendingWake) {
2263 } else {
22272264 yield* cancelPendingPrompts(yield* DateTime.now);
22282265 }
22292266 if (unsolicitedActivityDetected) {
......@@ -2844,6 +2881,17 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
28442881 `No pending Pi extension request ${requestInput.requestId}`,
28452882 );
28462883 }
2884 const answeredAt = yield* DateTime.now;
2885 if (
2886 pending.expiresAt !== null &&
2887 DateTime.toEpochMillis(answeredAt) >= pending.expiresAt
2888 ) {
2889 pendingPrompts.delete(String(requestInput.requestId));
2890 yield* cancelPrompt(pending, answeredAt);
2891 return yield* protocolError(
2892 `Pi extension request ${requestInput.requestId} has expired`,
2893 );
2894 }
28472895 const response = piUiResponse(pending, requestInput.decision, requestInput.answers);
28482896 yield* connection.send({
28492897 type: "extension_ui_response",
......@@ -2853,6 +2901,7 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
28532901 // Dropped only once Pi has the answer, so a failed send leaves the
28542902 // request retryable and still cancellable during teardown.
28552903 pendingPrompts.delete(String(requestInput.requestId));
2904 if (pending.expirationFiber !== null) yield* Fiber.interrupt(pending.expirationFiber);
28562905 if (pending.method === "confirm" && requestInput.decision === "acceptForSession") {
28572906 sessionApprovals.add(pending.approvalKey);
28582907 }
packages/provider-pi/src/server/mcpExtensionSource.test.ts+14-2
......@@ -249,7 +249,16 @@ describe("Pi tool discovery permissions", () => {
249249 it("allows discovery without confirmation and still gates the discovered tool", async () => {
250250 type ToolCallHook = (
251251 event: { toolName: string; input: unknown },
252 ctx: { ui: { confirm: (title: string, detail: string) => Promise<boolean> } },
252 ctx: {
253 signal: AbortSignal;
254 ui: {
255 confirm: (
256 title: string,
257 detail: string,
258 options: { signal: AbortSignal },
259 ) => Promise<boolean>;
260 };
261 },
253262 ) => Promise<{ block: true; reason: string } | undefined>;
254263 let toolCall: ToolCallHook | undefined;
255264 let searchPath = "builtin:tool-search";
......@@ -270,9 +279,12 @@ describe("Pi tool discovery permissions", () => {
270279 });
271280 assert.isDefined(toolCall);
272281 const confirmations: string[] = [];
282 const controller = new AbortController();
273283 const ctx = {
284 signal: controller.signal,
274285 ui: {
275 confirm: async (title: string) => {
286 confirm: async (title: string, _detail: string, options: { signal: AbortSignal }) => {
287 assert.strictEqual(options.signal, controller.signal);
276288 confirmations.push(title);
277289 return false;
278290 },
packages/provider-pi/src/server/mcpExtensionSource.ts+1
......@@ -318,6 +318,7 @@ export default async function t3McpExtension(pi: ExtensionAPI) {
318318 const approved = await ctx.ui.confirm(
319319 \`Allow \${event.toolName}?\`,
320320 toolInputSummary(event.input),
321 { signal: ctx.signal },
321322 );
322323 if (!approved) {
323324 return { block: true, reason: \`\${event.toolName} was declined in T3 Code.\` };