feat(pi): bridge extension UI dialogs (#1134)

2026-05-22 00:03:42 +08:00
parent db44a3e0e1
commit af10e64f82
7 changed files with 307 additions and 7 deletions
@@ -102,6 +102,20 @@ class SessionEvents {
    );
  }

+  nextPermissionRequest(): Promise<Extract<AgentStreamEvent, { type: "permission_requested" }>> {
+    return this.nextEvent(
+      (event): event is Extract<AgentStreamEvent, { type: "permission_requested" }> =>
+        event.type === "permission_requested",
+    );
+  }
+
+  nextPermissionResolution(): Promise<Extract<AgentStreamEvent, { type: "permission_resolved" }>> {
+    return this.nextEvent(
+      (event): event is Extract<AgentStreamEvent, { type: "permission_resolved" }> =>
+        event.type === "permission_resolved",
+    );
+  }
+
  private nextEvent<T extends AgentStreamEvent>(
    predicate: (event: AgentStreamEvent) => event is T,
  ): Promise<T> {
@@ -119,6 +133,125 @@ class SessionEvents {
 }

 describe("PiRpcAgentSession", () => {
+  test("bridges Pi RPC select extension UI requests through question permissions", async () => {
+    const { pi, session, events } = await createSession();
+    const fakeSession = pi.latestSession();
+
+    await session.startTurn("ask");
+    fakeSession.emit({
+      type: "extension_ui_request",
+      id: "ui-1",
+      method: "select",
+      title: "Pick one",
+      options: ["A", "B"],
+    });
+
+    const permission = await events.nextPermissionRequest();
+    expect(permission.request).toMatchObject({
+      id: "ui-1",
+      provider: "pi",
+      kind: "question",
+      title: "Pick one",
+      input: {
+        questions: [
+          {
+            question: "Pick one",
+            header: "Response",
+            options: [{ label: "A" }, { label: "B" }],
+            multiSelect: false,
+          },
+        ],
+      },
+      metadata: { extensionUiMethod: "select" },
+    });
+    expect(session.getPendingPermissions()).toHaveLength(1);
+
+    await session.respondToPermission("ui-1", {
+      behavior: "allow",
+      updatedInput: { answers: { Response: "B" } },
+    });
+
+    expect(fakeSession.extensionUiResponses).toEqual([{ id: "ui-1", response: { value: "B" } }]);
+    expect(session.getPendingPermissions()).toEqual([]);
+    await expect(events.nextPermissionResolution()).resolves.toMatchObject({
+      requestId: "ui-1",
+      resolution: { behavior: "allow" },
+    });
+  });
+
+  test("bridges Pi RPC input and confirm extension UI responses", async () => {
+    const { pi, session, events } = await createSession();
+    const fakeSession = pi.latestSession();
+
+    fakeSession.emit({
+      type: "extension_ui_request",
+      id: "input-1",
+      method: "input",
+      title: "Your name",
+      placeholder: "name",
+    });
+    await events.nextPermissionRequest();
+    await session.respondToPermission("input-1", {
+      behavior: "allow",
+      updatedInput: { answers: { Response: "Ada" } },
+    });
+
+    fakeSession.emit({
+      type: "extension_ui_request",
+      id: "confirm-1",
+      method: "confirm",
+      title: "Proceed?",
+    });
+    await events.nextPermissionRequest();
+    await session.respondToPermission("confirm-1", {
+      behavior: "allow",
+      updatedInput: { answers: { Response: "No" } },
+    });
+
+    expect(fakeSession.extensionUiResponses).toEqual([
+      { id: "input-1", response: { value: "Ada" } },
+      { id: "confirm-1", response: { confirmed: false } },
+    ]);
+  });
+
+  test("cancels Pi RPC extension UI dialogs when question permission is denied", async () => {
+    const { pi, session, events } = await createSession();
+    const fakeSession = pi.latestSession();
+
+    fakeSession.emit({
+      type: "extension_ui_request",
+      id: "ui-cancel",
+      method: "select",
+      title: "Pick one",
+      options: ["A", "B"],
+    });
+    await events.nextPermissionRequest();
+
+    await session.respondToPermission("ui-cancel", {
+      behavior: "deny",
+      message: "Dismissed by user",
+    });
+
+    expect(fakeSession.extensionUiResponses).toEqual([
+      { id: "ui-cancel", response: { cancelled: true } },
+    ]);
+  });
+
+  test("ignores Pi RPC fire-and-forget extension UI requests", async () => {
+    const { pi } = await createSession();
+    const fakeSession = pi.latestSession();
+
+    fakeSession.emit({
+      type: "extension_ui_request",
+      id: "notify-1",
+      method: "notify",
+      message: "hello",
+    });
+
+    expect(fakeSession.extensionUiResponses).toEqual([]);
+    expect(fakeSession.canceledExtensionUiRequests).toEqual([]);
+  });
+
  test("streams assistant text, reasoning, and tool calls from Pi events", async () => {
    const { pi, session, events } = await createSession();
    const fakeSession = pi.latestSession();
@@ -403,6 +403,108 @@ function latestPiErrorMessage(messages: PiAgentMessage[]): string | null {
    : null;
 }

+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function optionalString(value: unknown): string | undefined {
+  return typeof value === "string" ? value : undefined;
+}
+
+function mapExtensionUiRequestToPermission(
+  event: Extract<PiRuntimeEvent, { type: "extension_ui_request" }>,
+): AgentPermissionRequest | null {
+  switch (event.method) {
+    case "select":
+      return buildExtensionUiQuestionPermission(event, {
+        question: optionalString(event.title) ?? "Select an option",
+        options: Array.isArray(event.options)
+          ? event.options.filter((option): option is string => typeof option === "string")
+          : [],
+        multiSelect: false,
+      });
+    case "input":
+      return buildExtensionUiQuestionPermission(event, {
+        question: optionalString(event.title) ?? "Enter a value",
+        options: [],
+        multiSelect: false,
+      });
+    case "editor":
+      return buildExtensionUiQuestionPermission(event, {
+        question: optionalString(event.title) ?? "Edit text",
+        options: [],
+        multiSelect: false,
+      });
+    case "confirm":
+      return buildExtensionUiQuestionPermission(event, {
+        question: [optionalString(event.title), optionalString(event.message)]
+          .filter(Boolean)
+          .join("\n\n"),
+        options: ["Yes", "No"],
+        multiSelect: false,
+      });
+    default:
+      return null;
+  }
+}
+
+function buildExtensionUiQuestionPermission(
+  event: Extract<PiRuntimeEvent, { type: "extension_ui_request" }>,
+  input: { question: string; options: string[]; multiSelect: boolean },
+): AgentPermissionRequest {
+  const header = "Response";
+  return {
+    id: event.id,
+    provider: PI_PROVIDER,
+    name: `Pi ${event.method}`,
+    kind: "question",
+    title: input.question,
+    input: {
+      questions: [
+        {
+          question: input.question,
+          header,
+          options: input.options.map((label) => ({ label })),
+          multiSelect: input.multiSelect,
+        },
+      ],
+    },
+    metadata: {
+      extensionUiMethod: event.method,
+      answerHeader: header,
+    },
+  };
+}
+
+function firstPermissionAnswer(input: AgentMetadata | undefined): string | null {
+  const answers = isRecord(input?.answers) ? input.answers : null;
+  if (!answers) {
+    return null;
+  }
+  const first = Object.values(answers).find((value) => typeof value === "string");
+  return typeof first === "string" ? first : null;
+}
+
+function buildExtensionUiResponse(
+  request: AgentPermissionRequest,
+  response: AgentPermissionResponse,
+): { value?: string; confirmed?: boolean; cancelled?: boolean } {
+  if (response.behavior === "deny") {
+    return { cancelled: true };
+  }
+
+  const method = optionalString(request.metadata?.extensionUiMethod);
+  const answer = firstPermissionAnswer(response.updatedInput);
+  if (answer === null) {
+    return { cancelled: true };
+  }
+
+  if (method === "confirm") {
+    return { confirmed: /^yes$/i.test(answer.trim()) };
+  }
+  return { value: answer };
+}
+
 function mapPiModel(model: PiModel): AgentModelDefinition {
  return {
    provider: PI_PROVIDER,
@@ -428,6 +530,7 @@ export class PiRpcAgentSession implements AgentSession {

  private readonly subscribers = new Set<(event: AgentStreamEvent) => void>();
  private readonly activeToolCalls = new Map<string, PiTrackedToolCall>();
+  private readonly pendingExtensionUiRequests = new Map<string, AgentPermissionRequest>();
  private activeTurnId: string | null = null;
  private lastKnownThinkingOptionId: string | null;
  private state: PiSessionState;
@@ -540,12 +643,27 @@ export class PiRpcAgentSession implements AgentSession {
  }

  getPendingPermissions(): AgentPermissionRequest[] {
-    return [];
+    return [...this.pendingExtensionUiRequests.values()];
  }

  async respondToPermission(requestId: string, response: AgentPermissionResponse): Promise<void> {
-    void requestId;
-    void response;
+    const request = this.pendingExtensionUiRequests.get(requestId);
+    if (!request) {
+      throw new Error(`No pending permission request with id '${requestId}'`);
+    }
+    this.pendingExtensionUiRequests.delete(requestId);
+
+    this.runtimeSession.respondToExtensionUiRequest(
+      requestId,
+      buildExtensionUiResponse(request, response),
+    );
+    this.emit({
+      type: "permission_resolved",
+      provider: PI_PROVIDER,
+      requestId,
+      resolution: response,
+      turnId: this.currentTurnIdForEvent(),
+    });
  }

  describePersistence(): AgentPersistenceHandle | null {
@@ -624,9 +742,26 @@ export class PiRpcAgentSession implements AgentSession {
    return this.activeTurnId ?? undefined;
  }

+  private handleExtensionUiRequest(
+    event: Extract<PiRuntimeEvent, { type: "extension_ui_request" }>,
+  ): void {
+    const request = mapExtensionUiRequestToPermission(event);
+    if (!request) {
+      return;
+    }
+
+    this.pendingExtensionUiRequests.set(request.id, request);
+    this.emit({
+      type: "permission_requested",
+      provider: PI_PROVIDER,
+      request,
+      turnId: this.currentTurnIdForEvent(),
+    });
+  }
+
  private handleRuntimeEvent(event: PiRuntimeEvent): void {
    if (event.type === "extension_ui_request") {
-      this.runtimeSession.cancelExtensionUiRequest(event.id);
+      this.handleExtensionUiRequest(event);
      return;
    }
    if (event.type === "process_exit") {
@@ -166,8 +166,15 @@ class PiCliRuntimeSession implements PiRuntimeSession {
    return data.commands ?? [];
  }

+  respondToExtensionUiRequest(
+    id: string,
+    response: { value?: string; confirmed?: boolean; cancelled?: boolean },
+  ): void {
+    this.writeJsonLine({ type: "extension_ui_response", id, ...response });
+  }
+
  cancelExtensionUiRequest(id: string): void {
-    this.writeJsonLine({ type: "extension_ui_response", id, cancelled: true });
+    this.respondToExtensionUiRequest(id, { cancelled: true });
  }

  async close(): Promise<void> {
@@ -43,6 +43,10 @@ export interface PiRuntimeSession {
  setThinkingLevel(level: string): Promise<void>;
  getSessionStats(): Promise<PiSessionStats>;
  getCommands(): Promise<PiRpcSlashCommand[]>;
+  respondToExtensionUiRequest(
+    id: string,
+    response: { value?: string; confirmed?: boolean; cancelled?: boolean },
+  ): void;
  cancelExtensionUiRequest(id: string): void;
  close(): Promise<void>;
 }
@@ -55,6 +55,10 @@ export class FakePiSession implements PiRuntimeSession {
  readonly setThinkingLevelRequests: string[] = [];
  abortRequested = false;
  readonly canceledExtensionUiRequests: string[] = [];
+  readonly extensionUiResponses: Array<{
+    id: string;
+    response: { value?: string; confirmed?: boolean; cancelled?: boolean };
+  }> = [];
  setModelResult: PiModel | null = null;
  models: PiModel[] = [];
  messages: PiAgentMessage[] = [];
@@ -130,8 +134,16 @@ export class FakePiSession implements PiRuntimeSession {
    return this.commands;
  }

+  respondToExtensionUiRequest(
+    id: string,
+    response: { value?: string; confirmed?: boolean; cancelled?: boolean },
+  ): void {
+    this.extensionUiResponses.push({ id, response });
+  }
+
  cancelExtensionUiRequest(id: string): void {
    this.canceledExtensionUiRequests.push(id);
+    this.respondToExtensionUiRequest(id, { cancelled: true });
  }

  async close(): Promise<void> {}