@oh-my-pi/pi-coding-agent 15.1.2 → 15.1.4

package/src/mcp/tool-bridge.ts

@@ -5,7 +5,7 @@
  */
 import type { AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
 import type { TSchema } from "@oh-my-pi/pi-ai";
-import { sanitizeSchemaForMCP } from "@oh-my-pi/pi-ai/utils/schema";
+import { normalizeSchemaForMCP } from "@oh-my-pi/pi-ai/utils/schema";
 import { untilAborted } from "@oh-my-pi/pi-utils";
 import type { SourceMeta } from "../capability/types";
 import type {
@@ -231,7 +231,7 @@ export class MCPTool implements CustomTool<TSchema, MCPToolDetails> {
 		this.name = createMCPToolName(connection.name, tool.name);
 		this.label = `${connection.name}/${tool.name}`;
 		this.description = tool.description ?? `MCP tool from ${connection.name}`;
-		this.parameters = sanitizeSchemaForMCP(tool.inputSchema) as TSchema;
+		this.parameters = normalizeSchemaForMCP(tool.inputSchema) as TSchema;
 		this.mcpToolName = tool.name;
 		this.mcpServerName = connection.name;
 	}
@@ -324,7 +324,7 @@ export class DeferredMCPTool implements CustomTool<TSchema, MCPToolDetails> {
 		this.name = createMCPToolName(serverName, tool.name);
 		this.label = `${serverName}/${tool.name}`;
 		this.description = tool.description ?? `MCP tool from ${serverName}`;
-		this.parameters = sanitizeSchemaForMCP(tool.inputSchema) as TSchema;
+		this.parameters = normalizeSchemaForMCP(tool.inputSchema) as TSchema;
 		this.mcpToolName = tool.name;
 		this.mcpServerName = serverName;
 		this.#fallbackProvider = source?.provider;

package/src/modes/acp/acp-agent.ts

@@ -66,7 +66,11 @@ import {
 import { ACP_BUILTIN_SLASH_COMMANDS, executeAcpBuiltinSlashCommand } from "../../slash-commands/acp-builtins";
 import { parseThinkingLevel } from "../../thinking";
 import { createAcpClientBridge } from "./acp-client-bridge";
-import { mapAgentSessionEventToAcpSessionUpdates, mapToolKind } from "./acp-event-mapper";
+import {
+	buildToolCallStartUpdate,
+	mapAgentSessionEventToAcpSessionUpdates,
+	normalizeReplayToolArguments,
+} from "./acp-event-mapper";
 import { ACP_TERMINAL_AUTH_FLAG } from "./terminal-auth";
 
 const ACP_DEFAULT_MODE_ID = "default";
@@ -86,6 +90,9 @@ const SESSION_PAGE_SIZE = 50;
  * wait past this guard without hard-coding the literal.
  */
 export const ACP_BOOTSTRAP_RACE_GUARD_MS = 50;
+const ACP_CANCEL_CLEANUP_TIMEOUT_MS = 5_000;
+const ACP_ASYNC_DELIVERY_DRAIN_TIMEOUT_MS = 250;
+const ACP_ASYNC_DELIVERY_DRAIN_MAX_PASSES = 3;
 
 type AgentImageContent = {
 	type: "image";
@@ -102,6 +109,13 @@ type PromptTurnState = {
 	userMessageId: string;
 	cancelRequested: boolean;
 	settled: boolean;
+	/**
+	 * `abort()` is in-flight (or its bounded-timeout race). `undefined` while the turn is
+	 * running normally and after cleanup completes. The turn occupies `record.promptTurn`
+	 * for as long as either `!settled` or `cleanup` is set — that combined window is the
+	 * "turn in flight" predicate (`isPromptTurnInFlight`) every consumer gates on.
+	 */
+	cleanup: Promise<void> | undefined;
 	usageBaseline: UsageStatistics;
 	unsubscribe: (() => void) | undefined;
 	resolve: (value: PromptResponse) => void;
@@ -109,6 +123,16 @@ type PromptTurnState = {
 	promise: Promise<PromptResponse>;
 };
 
+/**
+ * A turn is "in flight" from the moment `prompt()` reserves the slot until `settled` is
+ * true AND any cancel cleanup has completed. Fork/queue/event gating all depend on this
+ * combined window — a settled-but-still-aborting turn is not safe to fork from, queue
+ * onto, or forward late events for.
+ */
+function isPromptTurnInFlight(turn: PromptTurnState | undefined): turn is PromptTurnState {
+	return turn !== undefined && (!turn.settled || turn.cleanup !== undefined);
+}
+
 type ManagedSessionRecord = {
 	session: AgentSession;
 	mcpManager: MCPManager | undefined;
@@ -116,6 +140,7 @@ type ManagedSessionRecord = {
 	promptQueue: PromptQueueState;
 	liveMessageId: string | undefined;
 	liveMessageProgress: { textEmitted: boolean; thoughtEmitted: boolean } | undefined;
+	toolArgsById: Map<string, unknown>;
 	extensionsConfigured: boolean;
 	// Installed inside `#scheduleBootstrapUpdates` (post-race-guard); released
 	// in `#disposeSessionRecord`. Lives independent of any prompt turn.
@@ -132,6 +157,14 @@ type ReplayableMessage = {
 	isError?: boolean;
 };
 
+type ReplayableToolItem = {
+	type?: unknown;
+	id?: unknown;
+	name?: unknown;
+	arguments?: unknown;
+	input?: unknown;
+};
+
 type MCPConfigMap = {
 	[name: string]: MCPServerConfig;
 };
@@ -337,13 +370,18 @@ export class AcpAgent implements Agent {
 	#disposePromise: Promise<void> | undefined;
 	#cleanupRegistered = false;
 	#clientCapabilities: ClientCapabilities | undefined;
+	#cancelCleanupTimeoutMs = ACP_CANCEL_CLEANUP_TIMEOUT_MS;
 
-	constructor(connection: AgentSideConnection, initialSession: AgentSession, createSession: CreateAcpSession) {
+	constructor(connection: AgentSideConnection, createSession: CreateAcpSession, initialSession?: AgentSession) {
 		this.#connection = connection;
 		this.#initialSession = initialSession;
 		this.#createSession = createSession;
 	}
 
+	setCancelCleanupTimeoutForTesting(timeoutMs: number): void {
+		this.#cancelCleanupTimeoutMs = Math.max(1, timeoutMs);
+	}
+
 	async initialize(params: InitializeRequest): Promise<InitializeResponse> {
 		this.#registerConnectionCleanup();
 		this.#clientCapabilities = params.clientCapabilities;
@@ -546,9 +584,15 @@ export class AcpAgent implements Agent {
 			throw new Error("ACP prompt already in progress for this session");
 		}
 		return await this.#queuePrompt(record, async () => {
-			const queuedTurn = record.promptTurn;
-			if (queuedTurn && !queuedTurn.settled) {
-				await queuedTurn.promise.catch(() => undefined);
+			const previousTurn = record.promptTurn;
+			if (previousTurn) {
+				// Wait for any prompt that's still settling or whose cancel cleanup is
+				// still in flight. We deliberately swallow the prompt rejection (the
+				// owning caller already received it) but let cleanup rejections
+				// propagate — a timed-out cancel must fail this queued prompt instead
+				// of letting it run on a session that is about to be closed.
+				await previousTurn.promise.catch(() => undefined);
+				await previousTurn.cleanup;
 			}
 
 			const converted = this.#convertPromptBlocks(params.prompt);
@@ -557,6 +601,7 @@ export class AcpAgent implements Agent {
 				userMessageId: params.messageId ?? crypto.randomUUID(),
 				cancelRequested: false,
 				settled: false,
+				cleanup: undefined,
 				usageBaseline: this.#cloneUsageStatistics(record.session.sessionManager.getUsageStatistics()),
 				unsubscribe: undefined,
 				resolve: pendingPrompt.resolve,
@@ -604,7 +649,7 @@ export class AcpAgent implements Agent {
 		const builtinResult = await executeAcpBuiltinSlashCommand(text, {
 			session: record.session,
 			sessionManager: record.session.sessionManager,
-			settings: Settings.instance,
+			settings: record.session.settings,
 			cwd: record.session.sessionManager.getCwd(),
 			output: output => this.#emitCommandOutput(record, output),
 			refreshCommands: () => this.#emitAvailableCommandsUpdate(record),
@@ -676,16 +721,53 @@ export class AcpAgent implements Agent {
 		if (!promptTurn || promptTurn.settled) {
 			return;
 		}
-		promptTurn.cancelRequested = true;
+		const cleanup = this.#beginCancelCleanup(record, promptTurn);
 		try {
-			await record.session.abort();
-			this.#finishPrompt(record, {
-				stopReason: "cancelled",
-				usage: this.#buildTurnUsage(promptTurn.usageBaseline, record.session.sessionManager.getUsageStatistics()),
-				userMessageId: promptTurn.userMessageId,
-			});
+			await cleanup;
 		} catch (error: unknown) {
-			this.#finishPrompt(record, undefined, error);
+			logger.warn("ACP cancel cleanup timed out; closing session", { sessionId: record.session.sessionId, error });
+			await this.#closeManagedSession(record.session.sessionId, record);
+		}
+	}
+
+	/**
+	 * Transition a still-running turn into cancellation: mark intent, drop the live-event
+	 * subscription, start the bounded `abort()` race, and resolve the ACP prompt response
+	 * with `stopReason: "cancelled"` so the client sees acceptance immediately. The
+	 * returned promise is the cleanup barrier — it resolves when `abort()` completes and
+	 * rejects when the timeout fires. Idempotent: a second call returns the same barrier.
+	 */
+	#beginCancelCleanup(record: ManagedSessionRecord, promptTurn: PromptTurnState): Promise<void> {
+		if (promptTurn.cleanup) {
+			return promptTurn.cleanup;
+		}
+		promptTurn.cancelRequested = true;
+		promptTurn.unsubscribe?.();
+		const cleanup = this.#runCancelCleanup(record, promptTurn);
+		promptTurn.cleanup = cleanup;
+		this.#finishPrompt(record, {
+			stopReason: "cancelled",
+			usage: this.#buildTurnUsage(promptTurn.usageBaseline, record.session.sessionManager.getUsageStatistics()),
+			userMessageId: promptTurn.userMessageId,
+		});
+		return cleanup;
+	}
+
+	async #runCancelCleanup(record: ManagedSessionRecord, promptTurn: PromptTurnState): Promise<void> {
+		let timer: NodeJS.Timeout | undefined;
+		const timeout = new Promise<never>((_, reject) => {
+			timer = setTimeout(() => reject(new Error("ACP cancel cleanup timed out")), this.#cancelCleanupTimeoutMs);
+		});
+		try {
+			await Promise.race([record.session.abort(), timeout]);
+		} finally {
+			if (timer) clearTimeout(timer);
+			// Order matters: clear `cleanup` before evicting the slot so the slot-eviction
+			// branch matches what `#finishPrompt` saw if it ran first.
+			promptTurn.cleanup = undefined;
+			if (promptTurn.settled && record.promptTurn === promptTurn) {
+				record.promptTurn = undefined;
+			}
 		}
 	}
 
@@ -739,6 +821,9 @@ export class AcpAgent implements Agent {
 			case "_omp/usage": {
 				const [firstRecord] = this.#sessions.values();
 				const target = firstRecord?.session ?? this.#initialSession;
+				if (!target) {
+					return { reports: [] };
+				}
 				const reports = await target.fetchUsageReports();
 				return { reports: reports ?? [] };
 			}
@@ -891,6 +976,7 @@ export class AcpAgent implements Agent {
 			promptQueue: { promise: Promise.resolve(), release: undefined },
 			liveMessageId: undefined,
 			liveMessageProgress: undefined,
+			toolArgsById: new Map(),
 			extensionsConfigured: false,
 			lifetimeUnsubscribe: undefined,
 		};
@@ -929,8 +1015,7 @@ export class AcpAgent implements Agent {
 	async #resolveForkSourceSessionPath(sessionId: string): Promise<string> {
 		const loaded = this.#sessions.get(sessionId);
 		if (loaded) {
-			const promptTurn = loaded.promptTurn;
-			if (promptTurn && !promptTurn.settled) {
+			if (isPromptTurnInFlight(loaded.promptTurn)) {
 				throw new Error(`ACP session fork is unavailable while a prompt is in progress: ${sessionId}`);
 			}
 			await loaded.session.sessionManager.flush();
@@ -950,23 +1035,31 @@ export class AcpAgent implements Agent {
 
 	async #handlePromptEvent(record: ManagedSessionRecord, event: AgentSessionEvent): Promise<void> {
 		const promptTurn = record.promptTurn;
-		if (!promptTurn || promptTurn.settled) {
+		if (!promptTurn || promptTurn.settled || promptTurn.cancelRequested) {
 			return;
 		}
 
+		if (event.type === "tool_execution_start" || event.type === "tool_execution_update") {
+			record.toolArgsById.set(event.toolCallId, event.args);
+		}
+
 		this.#prepareLiveAssistantMessage(record, event);
 		for (const notification of mapAgentSessionEventToAcpSessionUpdates(event, record.session.sessionId, {
 			getMessageId: message => this.#getLiveMessageId(record, message),
 			getMessageProgress: message => this.#getLiveMessageProgress(record, message),
+			getToolArgs: toolCallId => record.toolArgsById.get(toolCallId),
 			cwd: record.session.sessionManager.getCwd(),
 		})) {
 			await this.#connection.sessionUpdate(notification);
 		}
+		if (event.type === "tool_execution_end") {
+			record.toolArgsById.delete(event.toolCallId);
+		}
 		this.#clearLiveAssistantMessageAfterEvent(record, event);
 
 		if (event.type === "agent_end") {
 			await this.#emitEndOfTurnUpdates(record);
-			await record.session.waitForIdle();
+			await this.#waitForAcpPromptIdle(record);
 			this.#finishPrompt(record, {
 				stopReason: this.#resolveStopReason(event, promptTurn.cancelRequested),
 				usage: this.#buildTurnUsage(promptTurn.usageBaseline, record.session.sessionManager.getUsageStatistics()),
@@ -975,6 +1068,20 @@ export class AcpAgent implements Agent {
 		}
 	}
 
+	async #waitForAcpPromptIdle(record: ManagedSessionRecord): Promise<void> {
+		for (let pass = 0; pass < ACP_ASYNC_DELIVERY_DRAIN_MAX_PASSES; pass++) {
+			await record.session.waitForIdle();
+			const delivered = await record.session.drainAsyncJobDeliveriesForAcp({
+				timeoutMs: ACP_ASYNC_DELIVERY_DRAIN_TIMEOUT_MS,
+			});
+			if (!delivered) {
+				return;
+			}
+		}
+
+		await record.session.waitForIdle();
+	}
+
 	#prepareLiveAssistantMessage(record: ManagedSessionRecord, event: AgentSessionEvent): void {
 		if (
 			(event.type === "message_start" || event.type === "message_update" || event.type === "message_end") &&
@@ -1019,7 +1126,11 @@ export class AcpAgent implements Agent {
 		}
 		promptTurn.settled = true;
 		promptTurn.unsubscribe?.();
-		record.promptTurn = undefined;
+		// Keep the slot occupied until cancel cleanup finishes — `#runCancelCleanup`
+		// evicts the slot in its finally block once both flags say it's safe.
+		if (!promptTurn.cleanup && record.promptTurn === promptTurn) {
+			record.promptTurn = undefined;
+		}
 		if (error !== undefined) {
 			promptTurn.reject(error);
 			return;
@@ -1227,7 +1338,7 @@ export class AcpAgent implements Agent {
 
 	#getAvailableModes(session: AgentSession): Array<{ id: string; name: string; description: string }> {
 		const modes = [{ id: ACP_DEFAULT_MODE_ID, name: "Default", description: "Standard ACP headless mode" }];
-		if (Settings.instance.get("plan.enabled")) {
+		if (session.settings.get("plan.enabled")) {
 			modes.push({
 				id: ACP_PLAN_MODE_ID,
 				name: "Plan",
@@ -1511,16 +1622,30 @@ export class AcpAgent implements Agent {
 
 	async #replaySessionHistory(record: ManagedSessionRecord): Promise<void> {
 		const cwd = record.session.sessionManager.getCwd();
+		const replayedToolCallIds = new Set<string>();
+		const replayedToolCallArgs = new Map<string, unknown>();
 		for (const message of record.session.sessionManager.buildSessionContext().messages as ReplayableMessage[]) {
-			for (const notification of this.#messageToReplayNotifications(record.session.sessionId, message, cwd)) {
+			for (const notification of this.#messageToReplayNotifications(
+				record.session.sessionId,
+				message,
+				cwd,
+				replayedToolCallIds,
+				replayedToolCallArgs,
+			)) {
 				await this.#connection.sessionUpdate(notification);
 			}
 		}
 	}
 
-	#messageToReplayNotifications(sessionId: string, message: ReplayableMessage, cwd: string): SessionNotification[] {
+	#messageToReplayNotifications(
+		sessionId: string,
+		message: ReplayableMessage,
+		cwd: string,
+		replayedToolCallIds: Set<string>,
+		replayedToolCallArgs: Map<string, unknown>,
+	): SessionNotification[] {
 		if (message.role === "assistant") {
-			return this.#replayAssistantMessage(sessionId, message);
+			return this.#replayAssistantMessage(sessionId, message, cwd, replayedToolCallIds, replayedToolCallArgs);
 		}
 		if (
 			message.role === "user" ||
@@ -1540,11 +1665,19 @@ export class AcpAgent implements Agent {
 			typeof message.toolCallId === "string" &&
 			typeof message.toolName === "string"
 		) {
-			return this.#replayToolResult(sessionId, cwd, {
-				...message,
-				toolCallId: message.toolCallId,
-				toolName: message.toolName,
-			});
+			return this.#replayToolResult(
+				sessionId,
+				cwd,
+				{
+					...message,
+					toolCallId: message.toolCallId,
+					toolName: message.toolName,
+				},
+				{
+					includeStart: !replayedToolCallIds.has(message.toolCallId),
+					toolArgs: replayedToolCallArgs.get(message.toolCallId),
+				},
+			);
 		}
 		if (
 			message.role === "bashExecution" ||
@@ -1561,7 +1694,13 @@ export class AcpAgent implements Agent {
 		return [];
 	}
 
-	#replayAssistantMessage(sessionId: string, message: ReplayableMessage): SessionNotification[] {
+	#replayAssistantMessage(
+		sessionId: string,
+		message: ReplayableMessage,
+		cwd: string,
+		replayedToolCallIds: Set<string>,
+		replayedToolCallArgs: Map<string, unknown>,
+	): SessionNotification[] {
 		const notifications: SessionNotification[] = [];
 		const messageId = crypto.randomUUID();
 		if (Array.isArray(message.content)) {
@@ -1596,24 +1735,23 @@ export class AcpAgent implements Agent {
 					});
 					continue;
 				}
+				const toolItem = item as ReplayableToolItem;
 				if (
-					(item.type === "toolCall" || item.type === "tool_use") &&
-					"id" in item &&
-					typeof item.id === "string" &&
-					"name" in item &&
-					typeof item.name === "string"
+					(toolItem.type === "toolCall" || toolItem.type === "tool_use") &&
+					typeof toolItem.id === "string" &&
+					typeof toolItem.name === "string"
 				) {
-					const update: SessionUpdate = {
-						sessionUpdate: "tool_call",
-						toolCallId: item.id,
-						title: item.name,
-						kind: mapToolKind(item.name),
+					const args = this.#buildReplayAssistantToolArgs(toolItem);
+					const update = buildToolCallStartUpdate({
+						toolCallId: toolItem.id,
+						toolName: toolItem.name,
+						args,
 						status: "completed",
-					};
-					if ("arguments" in item && typeof item.arguments === "string") {
-						update.rawInput = item.arguments;
-					}
+						cwd,
+					});
 					notifications.push({ sessionId, update });
+					replayedToolCallIds.add(toolItem.id);
+					replayedToolCallArgs.set(toolItem.id, args);
 				}
 			}
 		}
@@ -1630,10 +1768,21 @@ export class AcpAgent implements Agent {
 		return notifications;
 	}
 
+	#buildReplayAssistantToolArgs(item: ReplayableToolItem): unknown {
+		if ("arguments" in item) {
+			return normalizeReplayToolArguments(item.arguments).args;
+		}
+		if (item.type === "tool_use" && "input" in item) {
+			return item.input;
+		}
+		return {};
+	}
+
 	#replayToolResult(
 		sessionId: string,
 		cwd: string,
 		message: Required<Pick<ReplayableMessage, "toolCallId" | "toolName">> & ReplayableMessage,
+		options: { includeStart?: boolean; toolArgs?: unknown } = {},
 	): SessionNotification[] {
 		const args = this.#buildReplayToolArgs(message.details);
 		const startEvent: AgentSessionEvent = {
@@ -1653,10 +1802,14 @@ export class AcpAgent implements Agent {
 				errorMessage: message.errorMessage,
 			},
 		};
-		return [
-			...mapAgentSessionEventToAcpSessionUpdates(startEvent, sessionId, { cwd }),
-			...mapAgentSessionEventToAcpSessionUpdates(endEvent, sessionId, { cwd }),
-		];
+		const notifications = mapAgentSessionEventToAcpSessionUpdates(endEvent, sessionId, {
+			cwd,
+			getToolArgs: toolCallId => (toolCallId === message.toolCallId ? options.toolArgs : undefined),
+		});
+		if (options.includeStart === false) {
+			return notifications;
+		}
+		return [...mapAgentSessionEventToAcpSessionUpdates(startEvent, sessionId, { cwd }), ...notifications];
 	}
 
 	#buildReplayToolArgs(details: unknown): { path?: string } {
@@ -1887,22 +2040,15 @@ export class AcpAgent implements Agent {
 
 	async #cancelPromptForClose(record: ManagedSessionRecord): Promise<void> {
 		const promptTurn = record.promptTurn;
-		if (!promptTurn || promptTurn.settled) {
+		if (!isPromptTurnInFlight(promptTurn)) {
 			return;
 		}
-
-		promptTurn.cancelRequested = true;
-		promptTurn.unsubscribe?.();
+		const cleanup = promptTurn.cleanup ?? this.#beginCancelCleanup(record, promptTurn);
 		try {
-			await record.session.abort();
+			await cleanup;
 		} catch (error) {
 			logger.warn("Failed to abort ACP prompt during session close", { error });
 		}
-		this.#finishPrompt(record, {
-			stopReason: "cancelled",
-			usage: this.#buildTurnUsage(promptTurn.usageBaseline, record.session.sessionManager.getUsageStatistics()),
-			userMessageId: promptTurn.userMessageId,
-		});
 	}
 
 	async #disposeSessionRecord(record: ManagedSessionRecord): Promise<void> {

package/src/modes/acp/acp-client-bridge.ts

@@ -36,7 +36,7 @@ export function createAcpClientBridge(
 		requestPermission: true,
 	};
 
-	const bridge: ClientBridge = { capabilities };
+	const bridge: ClientBridge = { capabilities, deferAgentInitiatedTurns: true };
 
 	if (capabilities.readTextFile) {
 		bridge.readTextFile = async params => {
@@ -122,6 +122,7 @@ async function requestPermission(
 		toolCallId: toolCall.toolCallId,
 		title: toolCall.title,
 		...(toolCall.kind ? { kind: toolCall.kind as ToolCallUpdate["kind"] } : {}),
+		...(toolCall.status ? { status: toolCall.status as ToolCallUpdate["status"] } : {}),
 		...(toolCall.rawInput !== undefined ? { rawInput: toolCall.rawInput } : {}),
 		...(toolCall.locations ? { locations: toolCall.locations } : {}),
 	};