diff --git a/packages/pi-coding-agent/src/core/chat-controller-ordering.test.ts b/packages/pi-coding-agent/src/core/chat-controller-ordering.test.ts index 551ca13b6..23bba623d 100644 --- a/packages/pi-coding-agent/src/core/chat-controller-ordering.test.ts +++ b/packages/pi-coding-agent/src/core/chat-controller-ordering.test.ts @@ -150,3 +150,71 @@ test("chat-controller keeps tool output ahead of delayed assistant text for exte assert.equal(host.chatContainer.children[0]?.constructor?.name, "ToolExecutionComponent"); assert.equal(host.chatContainer.children[1]?.constructor?.name, "AssistantMessageComponent"); }); + +test("chat-controller keeps serverToolUse output ahead of assistant text when external results arrive", async () => { + (globalThis as any)[Symbol.for("@gsd/pi-coding-agent:theme")] = { + fg: (_key: string, text: string) => text, + bg: (_key: string, text: string) => text, + bold: (text: string) => text, + italic: (text: string) => text, + truncate: (text: string) => text, + }; + + const host = createHost(); + const toolId = "mcp-secure-1"; + const serverToolUse = { + type: "serverToolUse", + id: toolId, + name: "mcp__gsd-workflow__secure_env_collect", + input: { projectDir: "/tmp/project", keys: [{ key: "SECURE_PASSWORD" }], destination: "dotenv" }, + }; + + await handleAgentEvent(host, { type: "message_start", message: makeAssistant([]) } as any); + + await handleAgentEvent( + host, + { + type: "message_update", + message: makeAssistant([serverToolUse]), + assistantMessageEvent: { + type: "server_tool_use", + contentIndex: 0, + partial: makeAssistant([serverToolUse]), + }, + } as any, + ); + + assert.equal(host.streamingComponent, undefined, "assistant content should stay deferred while only tool content streams"); + assert.equal(host.chatContainer.children.length, 1, "server tool block should render immediately"); + assert.equal(host.chatContainer.children[0]?.constructor?.name, "ToolExecutionComponent"); + + host.getMarkdownThemeWithSettings = () => ({}); + const resultMessage = makeAssistant([ + { + ...serverToolUse, + externalResult: { + content: [{ type: "text", text: "secure_env_collect was cancelled by user." }], + details: {}, + isError: true, + }, + }, + { type: "text", text: "The secure password collection was cancelled." }, + ]); + + await handleAgentEvent( + host, + { + type: "message_update", + message: resultMessage, + assistantMessageEvent: { + type: "server_tool_use", + contentIndex: 0, + partial: resultMessage, + }, + } as any, + ); + + assert.equal(host.chatContainer.children.length, 2, "assistant text should render after existing server tool output"); + assert.equal(host.chatContainer.children[0]?.constructor?.name, "ToolExecutionComponent"); + assert.equal(host.chatContainer.children[1]?.constructor?.name, "AssistantMessageComponent"); +}); diff --git a/packages/pi-coding-agent/src/core/extensions/types.ts b/packages/pi-coding-agent/src/core/extensions/types.ts index a1aad0dc9..5fea6389a 100644 --- a/packages/pi-coding-agent/src/core/extensions/types.ts +++ b/packages/pi-coding-agent/src/core/extensions/types.ts @@ -88,6 +88,8 @@ export interface ExtensionUIDialogOptions { timeout?: number; /** When true, the user can select multiple options. The return type becomes `string[]`. */ allowMultiple?: boolean; + /** When true, text input dialogs should hide typed characters if supported by the client surface. */ + secure?: boolean; } /** Placement for extension widgets. */ diff --git a/packages/pi-coding-agent/src/modes/interactive/controllers/chat-controller.ts b/packages/pi-coding-agent/src/modes/interactive/controllers/chat-controller.ts index 8d72d683d..168bdaa45 100644 --- a/packages/pi-coding-agent/src/modes/interactive/controllers/chat-controller.ts +++ b/packages/pi-coding-agent/src/modes/interactive/controllers/chat-controller.ts @@ -129,19 +129,6 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { host.streamingMessage = event.message; const innerEvent = event.assistantMessageEvent; - if (!host.streamingComponent && hasVisibleAssistantContent(host.streamingMessage)) { - host.streamingComponent = new AssistantMessageComponent( - undefined, - host.hideThinkingBlock, - host.getMarkdownThemeWithSettings(), - host.settingsManager.getTimestampFormat(), - ); - host.chatContainer.addChild(host.streamingComponent); - } - if (host.streamingComponent) { - host.streamingComponent.updateContent(host.streamingMessage); - } - let externalToolResult: | { toolCallId: string; content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; details: Record; isError: boolean } | undefined; @@ -156,6 +143,18 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { isError: ext.isError ?? false, }; } + } else if (innerEvent.type === "server_tool_use") { + const idx = typeof innerEvent.contentIndex === "number" ? innerEvent.contentIndex : -1; + const block = idx >= 0 ? (host.streamingMessage.content[idx] as any) : undefined; + const ext = block?.externalResult; + if (block?.id && ext) { + externalToolResult = { + toolCallId: block.id, + content: ext.content ?? [{ type: "text", text: "" }], + details: ext.details ?? {}, + isError: ext.isError ?? false, + }; + } } const contentBlocks = host.streamingMessage.content; @@ -230,6 +229,26 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { } } + // Render assistant text/thinking after tool components so mixed + // streams keep chronological ordering in the chat container. + const hasToolBlocks = hasAssistantToolBlocks(host.streamingMessage); + if (!host.streamingComponent && hasVisibleAssistantContent(host.streamingMessage)) { + host.streamingComponent = new AssistantMessageComponent( + undefined, + host.hideThinkingBlock, + host.getMarkdownThemeWithSettings(), + host.settingsManager.getTimestampFormat(), + ); + host.chatContainer.addChild(host.streamingComponent); + } + if (host.streamingComponent) { + if (hasToolBlocks) { + host.chatContainer.removeChild(host.streamingComponent); + host.chatContainer.addChild(host.streamingComponent); + } + host.streamingComponent.updateContent(host.streamingMessage); + } + // Update index: fully processed blocks won't need re-scanning. // Keep the last block's index (it may still be accumulating data), // so we re-check it next time but skip all earlier ones. diff --git a/packages/pi-coding-agent/src/modes/rpc/rpc-mode.ts b/packages/pi-coding-agent/src/modes/rpc/rpc-mode.ts index f2f8fbe4c..7d36e563a 100644 --- a/packages/pi-coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/pi-coding-agent/src/modes/rpc/rpc-mode.ts @@ -224,7 +224,7 @@ export async function runRpcMode(session: AgentSession): Promise { ), input: (title, placeholder, opts) => - createDialogPromise(opts, undefined, { method: "input", title, placeholder, timeout: opts?.timeout }, (r) => + createDialogPromise(opts, undefined, { method: "input", title, placeholder, timeout: opts?.timeout, secure: opts?.secure }, (r) => "cancelled" in r && r.cancelled ? undefined : "value" in r ? r.value : undefined, ), diff --git a/packages/pi-coding-agent/src/modes/rpc/rpc-types.ts b/packages/pi-coding-agent/src/modes/rpc/rpc-types.ts index 20d5c2c73..d6cd25bfc 100644 --- a/packages/pi-coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/pi-coding-agent/src/modes/rpc/rpc-types.ts @@ -291,6 +291,7 @@ export type RpcExtensionUIRequest = title: string; placeholder?: string; timeout?: number; + secure?: boolean; } | { type: "extension_ui_request"; id: string; method: "editor"; title: string; prefill?: string } | { diff --git a/src/resources/extensions/claude-code-cli/stream-adapter.ts b/src/resources/extensions/claude-code-cli/stream-adapter.ts index 5c97c9a4a..9e2541221 100644 --- a/src/resources/extensions/claude-code-cli/stream-adapter.ts +++ b/src/resources/extensions/claude-code-cli/stream-adapter.ts @@ -508,15 +508,15 @@ export function extractToolResultsFromSdkUserMessage(message: SDKUserMessage): A return extracted; } -function attachExternalResultsToToolCalls( - toolCalls: AssistantMessage["content"], +function attachExternalResultsToToolBlocks( + toolBlocks: AssistantMessage["content"], toolResultsById: ReadonlyMap, ): void { - for (const block of toolCalls) { - if (block.type !== "toolCall") continue; + for (const block of toolBlocks) { + if (block.type !== "toolCall" && block.type !== "serverToolUse") continue; const externalResult = toolResultsById.get(block.id); if (!externalResult) continue; - (block as ToolCallWithExternalResult).externalResult = externalResult; + (block as ToolCallWithExternalResult & { id: string }).externalResult = externalResult; } } @@ -554,8 +554,8 @@ async function pumpSdkMessages( /** Track the last text content seen across all assistant turns for the final message. */ let lastTextContent = ""; let lastThinkingContent = ""; - /** Collect tool calls from intermediate SDK turns for tool_execution events. */ - const intermediateToolCalls: AssistantMessage["content"] = []; + /** Collect tool blocks from intermediate SDK turns for tool execution rendering. */ + const intermediateToolBlocks: AssistantMessage["content"] = []; /** Preserve real external tool results from Claude Code's synthetic user messages. */ const toolResultsById = new Map(); @@ -666,9 +666,9 @@ async function pumpSdkMessages( lastTextContent = block.text; } else if (block.type === "thinking" && block.thinking) { lastThinkingContent = block.thinking; - } else if (block.type === "toolCall") { - // Collect tool calls for externalToolExecution rendering - intermediateToolCalls.push(block); + } else if (block.type === "toolCall" || block.type === "serverToolUse") { + // Collect tool blocks for externalToolExecution rendering + intermediateToolBlocks.push(block); } } } @@ -678,24 +678,33 @@ async function pumpSdkMessages( for (const { toolUseId, result } of extractToolResultsFromSdkUserMessage(msg as SDKUserMessage)) { toolResultsById.set(toolUseId, result); } - attachExternalResultsToToolCalls(intermediateToolCalls, toolResultsById); + attachExternalResultsToToolBlocks(intermediateToolBlocks, toolResultsById); // Push a synthetic toolcall_end for each tool call from this turn // so the TUI can render tool results in real-time during the SDK // session instead of waiting until the entire session completes. if (builder) { for (const block of builder.message.content) { - if (block.type !== "toolCall") continue; const extResult = (block as ToolCallWithExternalResult).externalResult; if (!extResult) continue; - // Push a toolcall_end with result attached so the chat-controller - // can call updateResult on the pending ToolExecutionComponent. - stream.push({ - type: "toolcall_end", - contentIndex: builder.message.content.indexOf(block), - toolCall: block, - partial: builder.message, - }); + const contentIndex = builder.message.content.indexOf(block); + if (contentIndex < 0) continue; + // Push synthetic completion events with result attached so the + // chat-controller can update pending ToolExecutionComponents. + if (block.type === "toolCall") { + stream.push({ + type: "toolcall_end", + contentIndex, + toolCall: block, + partial: builder.message, + }); + } else if (block.type === "serverToolUse") { + stream.push({ + type: "server_tool_use", + contentIndex, + partial: builder.message, + }); + } } } @@ -713,8 +722,8 @@ async function pumpSdkMessages( const finalContent: AssistantMessage["content"] = []; // Add tool calls from intermediate turns first (renders above text) - attachExternalResultsToToolCalls(intermediateToolCalls, toolResultsById); - finalContent.push(...intermediateToolCalls); + attachExternalResultsToToolBlocks(intermediateToolBlocks, toolResultsById); + finalContent.push(...intermediateToolBlocks); // Add text/thinking from the last turn if (builder && builder.message.content.length > 0) { diff --git a/src/resources/extensions/get-secrets-from-user.ts b/src/resources/extensions/get-secrets-from-user.ts index a8f1cfe36..967752048 100644 --- a/src/resources/extensions/get-secrets-from-user.ts +++ b/src/resources/extensions/get-secrets-from-user.ts @@ -126,7 +126,7 @@ async function collectOneSecret( ): Promise { if (!ctx.hasUI) return null; - return ctx.ui.custom((tui: any, theme: any, _kb: any, done: (r: string | null) => void) => { + const customResult = await ctx.ui.custom((tui: any, theme: any, _kb: any, done: (r: string | null) => void) => { let value = ""; let cachedLines: string[] | undefined; @@ -223,6 +223,29 @@ async function collectOneSecret( handleInput, }; }); + + // RPC/web surfaces may not implement ctx.ui.custom(). Fall back to a + // standard input prompt so users can still provide the secret. + if (customResult !== undefined) { + return customResult; + } + + if (typeof ctx.ui?.input !== "function") { + return null; + } + + const inputTitle = `Secure value for ${keyName} (${pageIndex + 1}/${totalPages})`; + const inputPlaceholder = hint || "Enter secret value"; + const inputResult = await ctx.ui.input( + inputTitle, + inputPlaceholder, + { secure: true }, + ); + if (typeof inputResult !== "string") { + return null; + } + const trimmed = inputResult.trim(); + return trimmed.length > 0 ? trimmed : null; } /** diff --git a/src/resources/extensions/gsd/tests/secure-env-collect.test.ts b/src/resources/extensions/gsd/tests/secure-env-collect.test.ts index 18acf7dd4..3e1a5e109 100644 --- a/src/resources/extensions/gsd/tests/secure-env-collect.test.ts +++ b/src/resources/extensions/gsd/tests/secure-env-collect.test.ts @@ -317,3 +317,48 @@ test("secure_env_collect #2997: null from ctx.ui.custom() is still treated as sk "Key returning null must NOT be in applied list", ); }); + +test("secure_env_collect: falls back to secure input prompt when custom UI is unavailable", async (t) => { + const { collectSecretsFromManifest } = await loadOrchestrator(); + + const tmp = makeTempDir("sec-input-fallback-test"); + t.after(() => { + rmSync(tmp, { recursive: true, force: true }); + }); + + const manifest = makeManifest([ + { key: "SECRET_FROM_INPUT_FALLBACK", status: "pending", formatHint: "starts with sk-" }, + ]); + await writeManifestFile(tmp, manifest); + + let callIndex = 0; + const inputCalls: Array<{ title: string; placeholder?: string; opts?: { secure?: boolean } }> = []; + const mockCtx = { + cwd: tmp, + hasUI: true, + ui: { + custom: async (_factory: any) => { + callIndex++; + if (callIndex <= 1) return null; // summary screen dismiss + return undefined; // collect screen unavailable on this surface + }, + input: async (title: string, placeholder?: string, opts?: { secure?: boolean }) => { + inputCalls.push({ title, placeholder, opts }); + return " sk-test-fallback-value "; + }, + }, + }; + + const result = await collectSecretsFromManifest(tmp, "M001", mockCtx as any); + + assert.ok( + result.applied.includes("SECRET_FROM_INPUT_FALLBACK"), + "Fallback input should collect and apply the key", + ); + assert.ok( + !result.skipped.includes("SECRET_FROM_INPUT_FALLBACK"), + "Fallback input should not mark the key as skipped", + ); + assert.equal(inputCalls.length, 1, "Fallback input should be requested once"); + assert.equal(inputCalls[0]?.opts?.secure, true, "Fallback input should request secure entry when supported"); +}); diff --git a/web/components/gsd/chat-mode.tsx b/web/components/gsd/chat-mode.tsx index f298f2754..81fefd1b3 100644 --- a/web/components/gsd/chat-mode.tsx +++ b/web/components/gsd/chat-mode.tsx @@ -1801,6 +1801,7 @@ function InlineInput({ return (
setValue(e.target.value)} diff --git a/web/components/gsd/focused-panel.tsx b/web/components/gsd/focused-panel.tsx index ee5c79156..1eb3bd2a1 100644 --- a/web/components/gsd/focused-panel.tsx +++ b/web/components/gsd/focused-panel.tsx @@ -180,6 +180,7 @@ function InputRenderer({ }} > setValue(e.target.value)} placeholder={request.placeholder || "Enter a value"} diff --git a/web/lib/gsd-workspace-store.tsx b/web/lib/gsd-workspace-store.tsx index adee496d6..ae67f5f28 100644 --- a/web/lib/gsd-workspace-store.tsx +++ b/web/lib/gsd-workspace-store.tsx @@ -380,7 +380,7 @@ export interface WorkspaceLiveState { export type ExtensionUiRequestEvent = | { type: "extension_ui_request"; id: string; method: "select"; title: string; options: string[]; timeout?: number; allowMultiple?: boolean } | { type: "extension_ui_request"; id: string; method: "confirm"; title: string; message: string; timeout?: number } - | { type: "extension_ui_request"; id: string; method: "input"; title: string; placeholder?: string; timeout?: number } + | { type: "extension_ui_request"; id: string; method: "input"; title: string; placeholder?: string; timeout?: number; secure?: boolean } | { type: "extension_ui_request"; id: string; method: "editor"; title: string; prefill?: string } | { type: "extension_ui_request"; id: string; method: "notify"; message: string; notifyType?: "info" | "warning" | "error" } | { type: "extension_ui_request"; id: string; method: "setStatus"; statusKey: string; statusText: string | undefined }