diff --git a/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.test.ts b/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.test.ts index 3ec9256a29..c117e3434a 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.test.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.test.ts @@ -17,35 +17,51 @@ const end = rest.find((frame: { type: string }) => frame.type === "tool_executio function startedHost() { const host = createPiExtensionHost(); - const permission = host.onToolStart({ + host.onToolStart({ callId: start.toolCallId, toolName: start.toolName, args: start.args, status: "running", result: null, }); - if (!permission) throw new Error("Question permission missing"); - return { host, permission }; + const mapped = host.mapDialog(dialogs[0], "pi"); + if (mapped?.type !== "permission") throw new Error("Question permission missing"); + return { host, permission: mapped.request, firstDialog: dialogs[0] }; } +test("the form opens when the tool asks, and carries every question", () => { + const host = createPiExtensionHost(); + expect( + host.onToolStart({ + callId: start.toolCallId, + toolName: start.toolName, + args: start.args, + status: "running", + result: null, + }), + ).toBeUndefined(); + expect(host.mapDialog(dialogs[0], "pi")).toMatchObject({ + type: "permission", + request: { + kind: "question", + input: { + questions: [ + { + header: "Snack", + options: [{ description: "Fresh fruit" }, { description: "Sweet treat" }], + allowOther: true, + multiSelect: false, + }, + { header: "Colors", multiSelect: true }, + { header: "Project", allowOther: true }, + ], + }, + }, + }); +}); + test("real rpiv RPC dialogs use one form and preserve choices, multi-select, and custom text", () => { const { host, permission } = startedHost(); - expect(permission).toMatchObject({ - kind: "question", - input: { - questions: [ - { - header: "Snack", - options: [{ description: "Fresh fruit" }, { description: "Sweet treat" }], - allowOther: true, - multiSelect: false, - }, - { header: "Colors", multiSelect: true }, - { header: "Project", allowOther: true }, - ], - }, - }); - expect(host.mapDialog(dialogs[0], "pi")).toEqual({ type: "deferred" }); expect( host.respondToPermission(permission, { behavior: "allow", @@ -68,10 +84,12 @@ test("real rpiv RPC dialogs use one form and preserve choices, multi-select, and }); }); -test("cancellation declines each remaining real RPC dialog", () => { - const { host, permission } = startedHost(); - expect(host.respondToPermission(permission, { behavior: "deny" })).toEqual({ responses: [] }); - for (const dialog of dialogs) { +test("cancellation declines the open dialog and each remaining real RPC dialog", () => { + const { host, permission, firstDialog } = startedHost(); + expect(host.respondToPermission(permission, { behavior: "deny" })).toEqual({ + responses: [{ id: firstDialog.id, response: { cancelled: true } }], + }); + for (const dialog of dialogs.slice(1)) { expect(host.mapDialog(dialog, "pi")).toEqual({ type: "response", response: { cancelled: true }, @@ -79,6 +97,91 @@ test("cancellation declines each remaining real RPC dialog", () => { } }); +test("a questionnaire pi would reject over its own caps never opens a form", () => { + const host = createPiExtensionHost(); + const rejected = structuredClone(start.args); + rejected.questions[0].header = "Brainstorm cadence"; // 18 chars, pi caps headers at 16 + host.onToolStart({ + callId: "rejected-call", + toolName: start.toolName, + args: rejected, + status: "running", + result: null, + }); + // No orphan card, and the tool never reaches the point of asking, so the + // generic dialog fallback stays available instead. + expect( + host.mapDialog({ ...dialogs[0], id: "orphan", title: "[Brainstorm cadence] Pick one?" }, "pi") + ?.type, + ).toBeUndefined(); +}); + +test("two questionnaires keep separate forms and one form cannot cancel the other", () => { + const host = createPiExtensionHost(); + const second = { + questions: [ + { + question: "Which drink?", + header: "Drink", + options: [ + { label: "Water", description: "Plain" }, + { label: "Juice", description: "Sweet" }, + ], + }, + ], + }; + host.onToolStart({ + callId: start.toolCallId, + toolName: start.toolName, + args: start.args, + status: "running", + result: null, + }); + host.onToolStart({ + callId: "call_second", + toolName: "ask_user_question", + args: second, + status: "running", + result: null, + }); + const firstForm = host.mapDialog(dialogs[0], "pi"); + const drinkDialog = { + id: "drink-dialog", + type: "extension_ui_request", + method: "select", + title: "[Drink] Which drink?", + options: ["1. Water — Plain", "2. Juice — Sweet", "3. Type something."], + }; + const secondForm = host.mapDialog(drinkDialog, "pi"); + expect(firstForm?.type).toBe("permission"); + expect(secondForm?.type).toBe("permission"); + + // Closing the first questionnaire must not touch the second one. + expect(host.respondToPermission(firstForm.request, { behavior: "deny" })).toEqual({ + responses: [{ id: dialogs[0].id, response: { cancelled: true } }], + }); + expect( + host.respondToPermission(secondForm.request, { + behavior: "allow", + updatedInput: { answers: { Drink: "Juice" } }, + }), + ).toEqual({ + responses: [{ id: drinkDialog.id, response: { value: drinkDialog.options[1] } }], + }); +}); + +test("a form whose tool call already ended resolves nothing", () => { + const { host, permission } = startedHost(); + host.onToolEnd({ + callId: start.toolCallId, + toolName: start.toolName, + args: start.args, + status: "failed", + result: null, + }); + expect(host.respondToPermission(permission, { behavior: "deny" })).toEqual({ responses: [] }); +}); + test("captured result renders in live and history tool rows, while foreign shapes decline", () => { const call = { callId: end.toolCallId, @@ -149,9 +252,10 @@ test("the Pi agent releases a deferred captured dialog when the form is answered }); const runtime = pi.latestSession(); runtime.emit(start); + expect(session.getPendingPermissions()).toEqual([]); + runtime.emit(dialogs[0]); const [permission] = session.getPendingPermissions(); expect(permission.input?.questions).toHaveLength(3); - runtime.emit(dialogs[0]); expect(runtime.extensionUiResponses).toEqual([]); await session.respondToPermission(permission.id, { behavior: "allow", diff --git a/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.ts b/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.ts index 07906cc955..1e29afbcbb 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.ts @@ -2,17 +2,26 @@ import { z } from "zod"; import type { AgentPermissionRequest } from "../../../../agent-sdk-types.js"; import type { PiExtension, PiExtensionDialog, PiExtensionUiResponse } from "../contract.js"; +// Caps mirrored from pi's own ask_user_question schema +// (@juicesharp/rpiv-ask-user-question tool/types.ts: MAX_HEADER_LENGTH 16, +// MAX_LABEL_LENGTH 60). pi rejects an over-limit call before the tool runs, so a +// looser schema here renders a form for a questionnaire that no tool will ever +// answer: the user sees a second, identical card with nothing behind it, and any +// interaction with it lands on the live questionnaire instead. +const MAX_HEADER_LENGTH = 16; +const MAX_LABEL_LENGTH = 60; + const Params = z .object({ questions: z .array( z.object({ question: z.string(), - header: z.string(), + header: z.string().max(MAX_HEADER_LENGTH), options: z .array( z.object({ - label: z.string(), + label: z.string().max(MAX_LABEL_LENGTH), description: z.string(), preview: z.string().optional(), }), @@ -55,6 +64,16 @@ const Result = z.object({ type Question = z.infer["questions"][number]; const PERMISSION_PREFIX = "rpiv-question:"; +interface QuestionForm { + callId: string; + questions: Question[]; + answers: Record | null; + cancelled: boolean; + index: number; + customInput: boolean; + deferred: PiExtensionDialog | null; +} + function answerFor(answers: Record, question: Question): string | null { const answer = answers[question.header]; return typeof answer === "string" ? answer : null; @@ -103,34 +122,71 @@ function replyToDialog( export const rpivAskUserQuestion: PiExtension = { id: "rpiv-ask-user-question", createSession() { - let active: { - callId: string; - questions: Question[]; - answers: Record | null; - cancelled: boolean; - index: number; - customInput: boolean; - deferred: PiExtensionDialog | null; - } | null = null; + // One entry per live ask_user_question tool call. A single slot shared by + // every call lets a second questionnaire overwrite the first, and lets a + // stale form's dismissal cancel a questionnaire it never belonged to. + const forms: QuestionForm[] = []; - function answerDialog(dialog: PiExtensionDialog): PiExtensionUiResponse { - if (!active || active.cancelled) return { cancelled: true }; - const question = active.questions[active.index]; - if (!question || !active.answers) return { cancelled: true }; - const answer = answerFor(active.answers, question); + function dialogTitle(dialog: PiExtensionDialog): string { + return typeof dialog.title === "string" ? dialog.title : ""; + } + + // The provider dialog titles are `[Header] question…`, so a dialog can be + // attributed to the form whose current question it asks. When titles cannot + // decide (older provider format), the form still waiting for its form wins. + function formForDialog(dialog: PiExtensionDialog): QuestionForm | undefined { + const title = dialogTitle(dialog); + const asked = forms.find((form) => { + const question = form.questions[form.index]; + return ( + !!question && + (title.startsWith(`[${question.header}] `) || title.startsWith(question.question)) + ); + }); + if (asked) return asked; + const waiting = forms.find((form) => !form.answers && !form.cancelled); + return waiting ?? forms.find((form) => form.cancelled); + } + + function formCard(form: QuestionForm, provider: string): AgentPermissionRequest { + return { + id: `${PERMISSION_PREFIX}${form.callId}`, + provider, + name: "Pi ask_user_question", + kind: "question", + title: form.questions[0].question, + input: { + questions: form.questions.map((question) => ({ + question: question.question, + header: question.header, + options: question.options, + multiSelect: question.multiSelect ?? false, + allowOther: true, + })), + }, + }; + } + + function answerDialog(form: QuestionForm, dialog: PiExtensionDialog): PiExtensionUiResponse { + if (form.cancelled) return { cancelled: true }; + const question = form.questions[form.index]; + if (!question || !form.answers) return { cancelled: true }; + const answer = answerFor(form.answers, question); if (answer === null) return { cancelled: true }; - const mapped = replyToDialog(dialog, question, answer, active.customInput); - active.customInput = mapped.customInput; - if (!active.customInput) active.index++; + const mapped = replyToDialog(dialog, question, answer, form.customInput); + form.customInput = mapped.customInput; + if (!form.customInput) form.index++; return mapped.response; } return { - onToolStart(call, provider): AgentPermissionRequest | undefined { + onToolStart(call) { if (call.toolName !== "ask_user_question") return undefined; const parsed = Params.safeParse(call.args); if (!parsed.success) return undefined; - active = { + const existing = forms.findIndex((form) => form.callId === call.callId); + if (existing >= 0) forms.splice(existing, 1); // replayed tool start + forms.push({ callId: call.callId, questions: parsed.data.questions, answers: null, @@ -138,47 +194,42 @@ export const rpivAskUserQuestion: PiExtension = { index: 0, customInput: false, deferred: null, - }; - return { - id: `${PERMISSION_PREFIX}${call.callId}`, - provider, - name: "Pi ask_user_question", - kind: "question", - title: parsed.data.questions[0].question, - input: { - questions: parsed.data.questions.map((question) => ({ - question: question.question, - header: question.header, - options: question.options, - multiSelect: question.multiSelect ?? false, - allowOther: true, - })), - }, - }; + }); + // The card appears when the tool actually asks (its first dialog), not + // at tool start: a call pi rejects, interrupts, or finishes without + // ever asking then leaves no unanswered card behind. + return undefined; }, onToolEnd(call) { - if (active?.callId === call.callId) active = null; + const index = forms.findIndex((form) => form.callId === call.callId); + if (index >= 0) forms.splice(index, 1); }, - mapDialog(dialog) { - if (!active || (dialog.method !== "select" && dialog.method !== "input")) return undefined; - if (!active.answers && !active.cancelled) { - active.deferred = dialog; - return { type: "deferred" }; + mapDialog(dialog, provider) { + if (dialog.method !== "select" && dialog.method !== "input") return undefined; + const form = formForDialog(dialog); + if (!form) return undefined; + if (form.cancelled) return { type: "response", response: { cancelled: true } }; + if (!form.answers) { + form.deferred = dialog; + return { type: "permission", request: formCard(form, provider) }; } - return { type: "response", response: answerDialog(dialog) }; + return { type: "response", response: answerDialog(form, dialog) }; }, respondToPermission(request, response) { - if (!active || request.id !== `${PERMISSION_PREFIX}${active.callId}`) return undefined; + if (!request.id.startsWith(PERMISSION_PREFIX)) return undefined; + const callId = request.id.slice(PERMISSION_PREFIX.length); + const form = forms.find((candidate) => candidate.callId === callId); + if (!form) return { responses: [] }; // its tool already ended: nothing to corrupt const answers = response.behavior === "allow" ? response.updatedInput?.answers : undefined; - active.answers = + form.answers = answers && typeof answers === "object" && !Array.isArray(answers) ? (answers as Record) : null; - active.cancelled = response.behavior === "deny" || !active.answers; - const deferred = active.deferred; - active.deferred = null; + form.cancelled = response.behavior === "deny" || !form.answers; + const deferred = form.deferred; + form.deferred = null; return { - responses: deferred ? [{ id: deferred.id, response: answerDialog(deferred) }] : [], + responses: deferred ? [{ id: deferred.id, response: answerDialog(form, deferred) }] : [], }; }, mapToolCall(call) {