From 77f81a7c5a268158d6de9c5e1ff6cef08bcf715b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Fri, 17 Apr 2026 14:53:22 -0400 Subject: [PATCH] fix(runtime): revive AI Co-Founder tool-call gating via pi beforeToolCall MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An earlier cleanup deleted HumanApprovalHook wholesale on the rationale that the body always returned {}. That was the stub implementation, not the feature: 版本更新.md specifies AI Co-Founder mode as "human-reviewed execution that pauses until human review arrives". With no beforeToolCall attached, Co-Founder and CEO modes behaved identically — a silent regression of the product's core differentiator. Bring the hook back with real gating logic and split the Agent construction path by execution mode so CEO runs pay zero IPC overhead. - HumanApprovalHook.before(): create a `dangerous_command` approval row for non-whitelisted tool calls, emit approval-updated, return a promise that resolves when a human decides. - AUTO_APPROVED_TOOLS whitelist: read/ls/glob/grep/find/tree. Anything else (bash, write, edit, …) waits for human approval. - approvalHook.resolveApproval(): invoked from wire-api-server's onApprovalResolved *before* automation and thread follow-ups, so the pi Agent unblocks the instant the human decides. - approvalHook.cancelAll(): release every pending wait on shutdown so the worker can drain cleanly; wired into service-graph-pi-runtime.release. - worker-runner: attach beforeToolCall only when job.executionMode === "ai_cofounder". CEO mode constructs Agent with no hook, no RPC hops. - WorkerRpcClient.invoke: accept per-call timeoutMs override so approval RPCs can wait 30min instead of the 30s default. - worker-types schema: restore beforeToolCall in the method enum. - RuntimeBridge.handleRpc: restore beforeToolCall branch. No changes to approval DB schema, approval lifecycle, or existing automation/thread handling. CEO-mode runs retain zero-IPC-per-tool behavior from the prior cleanup. --- src/main/agent-runtime/approval-hook.ts | 136 +++++++++++++++++++++++ src/main/agent-runtime/index.ts | 1 + src/main/agent-runtime/runtime-bridge.ts | 27 ++++- src/main/bootstrap/wire-api-server.ts | 6 + src/main/service-graph-pi-runtime.ts | 4 +- src/main/worker/worker-rpc.ts | 6 +- src/main/worker/worker-runner.ts | 30 +++++ src/main/worker/worker-types.ts | 2 +- 8 files changed, 204 insertions(+), 8 deletions(-) create mode 100644 src/main/agent-runtime/approval-hook.ts diff --git a/src/main/agent-runtime/approval-hook.ts b/src/main/agent-runtime/approval-hook.ts new file mode 100644 index 00000000..d3e70925 --- /dev/null +++ b/src/main/agent-runtime/approval-hook.ts @@ -0,0 +1,136 @@ +import type { AppDatabase } from "@main/database"; +import type { DesktopEvent } from "@shared/types"; + +interface BeforeToolCallPayload { + agentId: string; + companyId: string; + taskId: string | null; + toolCallId: string; + toolName: string; + args: unknown; +} + +interface PendingApproval { + resolve: (verdict: { block?: boolean; reason?: string }) => void; +} + +export interface HumanApprovalHookDeps { + db: AppDatabase; + emitEvent: (event: DesktopEvent) => void; + logger: { + info: (...args: unknown[]) => void; + warn: (...args: unknown[]) => void; + error: (...args: unknown[]) => void; + }; +} + +/** + * Tools whose side-effects are read-only (safe to auto-approve in Co-Founder + * mode). Anything outside this set waits for a human decision. + * + * This is intentionally conservative: `bash` is gated even for obviously + * safe commands like `ls`, because parsing shell syntax for safety is its + * own class of bug. If users want per-command auto-approval they can + * approve once and the run proceeds; the alternative is a false sense of + * security. + */ +const AUTO_APPROVED_TOOLS = new Set([ + "read", + "ls", + "glob", + "grep", + "find", + "tree", +]); + +function summarizeArgs(args: unknown, maxLen = 240): string { + if (args === undefined || args === null) return ""; + try { + const json = typeof args === "string" ? args : JSON.stringify(args); + return json.length <= maxLen ? json : `${json.slice(0, maxLen - 1)}…`; + } catch { + return ""; + } +} + +/** + * Gates pi `beforeToolCall` invocations against the Human-in-the-Loop + * approval system used by AI Co-Founder mode. + * + * Contract: + * - Tools in AUTO_APPROVED_TOOLS resolve synchronously with `{}` (no-op). + * - Everything else creates a `dangerous_command` approval row, emits an + * `approval-updated` event so the renderer can surface it, and blocks + * the worker RPC until the human resolves the approval. + * - Resolutions flow in via `resolveApproval` from the approval-resolution + * pipeline in `wire-api-server.ts` (shared with automation + thread). + * - On shutdown, `cancelAll` releases every pending wait with a block + * verdict so the worker can exit cleanly. + * + * In AI CEO mode the worker never attaches `beforeToolCall` to the pi + * Agent, so this class is instantiated but sees zero traffic. + */ +export class HumanApprovalHook { + private readonly pendingByApproval = new Map(); + + constructor(private readonly deps: HumanApprovalHookDeps) {} + + async before( + runId: string, + payload: BeforeToolCallPayload, + ): Promise<{ block?: boolean; reason?: string }> { + if (AUTO_APPROVED_TOOLS.has(payload.toolName)) return {}; + + const summary = `${payload.toolName}: ${summarizeArgs(payload.args, 160)}`; + const approvalId = this.deps.db.requestApproval({ + companyId: payload.companyId, + relatedTaskId: payload.taskId ?? null, + requestedByAgentId: payload.agentId, + type: "dangerous_command", + payloadSummary: summary, + impactSummary: "Agent requested a gated tool call; Co-Founder mode requires a human decision.", + payloadJson: JSON.stringify({ + runId, + toolCallId: payload.toolCallId, + toolName: payload.toolName, + args: payload.args, + }), + }); + + this.deps.emitEvent({ type: "approval-updated", approvalId, action: "created" }); + this.deps.logger.info( + `[approval-hook] gated ${payload.toolName} for agent ${payload.agentId} (approval ${approvalId})`, + ); + + return new Promise<{ block?: boolean; reason?: string }>((resolve) => { + this.pendingByApproval.set(approvalId, { resolve }); + }); + } + + /** + * Called by the approval-resolution pipeline in `wire-api-server.ts` + * for every approval decision (automation, UI, HTTP). No-ops when the + * approvalId is not one we are waiting on. + */ + resolveApproval(approvalId: string, decision: string): void { + const pending = this.pendingByApproval.get(approvalId); + if (!pending) return; + this.pendingByApproval.delete(approvalId); + if (decision === "approved") { + pending.resolve({}); + } else { + pending.resolve({ block: true, reason: `tool call ${decision} by human reviewer` }); + } + } + + /** + * Releases every pending wait with a block verdict. Called on worker + * shutdown so in-flight RPC promises resolve and the worker can exit. + */ + cancelAll(reason: string): void { + for (const [approvalId, pending] of this.pendingByApproval) { + pending.resolve({ block: true, reason }); + this.pendingByApproval.delete(approvalId); + } + } +} diff --git a/src/main/agent-runtime/index.ts b/src/main/agent-runtime/index.ts index 667a9095..9b60f415 100644 --- a/src/main/agent-runtime/index.ts +++ b/src/main/agent-runtime/index.ts @@ -1,3 +1,4 @@ export { SqliteSessionStore } from "./sqlite-session-store"; export { CostWriter } from "./cost-writer"; +export { HumanApprovalHook } from "./approval-hook"; export { RuntimeBridge } from "./runtime-bridge"; diff --git a/src/main/agent-runtime/runtime-bridge.ts b/src/main/agent-runtime/runtime-bridge.ts index 921230fe..b2f2fa05 100644 --- a/src/main/agent-runtime/runtime-bridge.ts +++ b/src/main/agent-runtime/runtime-bridge.ts @@ -4,6 +4,7 @@ import type { DesktopEvent } from "@shared/types"; import { ensureTaskThread } from "@main/thread-integration/thread-helpers"; import { SqliteSessionStore } from "./sqlite-session-store"; import { CostWriter } from "./cost-writer"; +import { HumanApprovalHook } from "./approval-hook"; import type { CredentialResolver } from "@main/credentials"; import { getConnectorDefinition } from "@main/connectors"; import { projectInboxMessageToPi } from "./inbox-projection"; @@ -132,17 +133,23 @@ interface ActiveRunContext { * - Forwards `agent-event` from the worker to the renderer (push channel). * - Persists each `message_end` message via `SqliteSessionStore`. * - Calls `CostWriter` for assistant messages with usage. - * - * Tool-call gating runs fully in the model layer (AI CEO and AI Co-Founder - * both autonomous); no `beforeToolCall` / `afterToolCall` IPC hop exists. + * - Hosts `HumanApprovalHook` for Co-Founder mode tool gating; the worker + * only invokes the hook RPC when running in `ai_cofounder` mode, so + * CEO-mode runs pay zero IPC overhead. */ export class RuntimeBridge { + readonly approvalHook: HumanApprovalHook; private readonly costWriter: CostWriter; private readonly activeRuns = new Map(); private readonly sessionStores = new Map(); constructor(private readonly deps: RuntimeBridgeDeps) { this.costWriter = new CostWriter({ db: deps.db }); + this.approvalHook = new HumanApprovalHook({ + db: deps.db, + emitEvent: deps.emitEvent, + logger: deps.logger, + }); } private registerActiveRun(context: ActiveRunContext): void { @@ -291,7 +298,7 @@ export class RuntimeBridge { /** Handles RPC requests issued by the worker (`rpc-request` envelope). */ async handleRpc(request: { rpcId: string; - method: "persistMessage" | "prepareRun"; + method: "persistMessage" | "prepareRun" | "beforeToolCall"; runId: string; payload: unknown; }): Promise<{ ok: true; result: unknown } | { ok: false; error: string }> { @@ -301,6 +308,18 @@ export class RuntimeBridge { this.getOrCreateStore(payload.taskSessionId).appendMessage(payload.message); return { ok: true, result: undefined }; } + if (request.method === "beforeToolCall") { + const payload = request.payload as { + agentId: string; + companyId: string; + taskId: string | null; + toolCallId: string; + toolName: string; + args: unknown; + }; + const verdict = await this.approvalHook.before(request.runId, payload); + return { ok: true, result: verdict }; + } if (request.method === "prepareRun") { const payload = request.payload as { connectorId: string; diff --git a/src/main/bootstrap/wire-api-server.ts b/src/main/bootstrap/wire-api-server.ts index 2a3735e1..e8d565c7 100644 --- a/src/main/bootstrap/wire-api-server.ts +++ b/src/main/bootstrap/wire-api-server.ts @@ -79,6 +79,12 @@ export function wireApiServer(deps: WireApiServerDeps): AgentApiServer { const onApprovalResolved = (approvalId: string, companyId: string, decision: string) => { const automation = deps.serviceGraph.getAutomationService(); if (!automation) throw new Error("Automation service is not initialized."); + // Release any worker waiting on this approval *before* the business + // follow-ups fire, so the pi Agent unblocks immediately once the human + // decides. Automation state changes and thread updates are allowed to + // take longer — they do not hold back the running agent. + const piRuntime = deps.serviceGraph.getPiRuntime(); + if (piRuntime) piRuntime.runtimeBridge.approvalHook.resolveApproval(approvalId, decision); automation.handleApprovalResolved(approvalId, companyId, decision); deps.threadIntegration.onApprovalResolved(approvalId, companyId, decision); }; diff --git a/src/main/service-graph-pi-runtime.ts b/src/main/service-graph-pi-runtime.ts index c89b046c..407e0ab0 100644 --- a/src/main/service-graph-pi-runtime.ts +++ b/src/main/service-graph-pi-runtime.ts @@ -51,6 +51,8 @@ export function createPiRuntimeServices(deps: PiRuntimeServicesDeps): PiRuntimeS credentialResolver, oauthLoginService, runtimeBridge, - release: () => {}, + release: () => { + runtimeBridge.approvalHook.cancelAll("AgentCompany shutting down."); + }, }; } diff --git a/src/main/worker/worker-rpc.ts b/src/main/worker/worker-rpc.ts index 29c88beb..de7424cb 100644 --- a/src/main/worker/worker-rpc.ts +++ b/src/main/worker/worker-rpc.ts @@ -20,15 +20,17 @@ export class WorkerRpcClient { runId: string, method: string, payload: unknown, + options?: { timeoutMs?: number }, ): Promise<{ ok: true; result: unknown } | { ok: false; error: string }> { const rpcId = randomUUID(); + const effectiveTimeoutMs = options?.timeoutMs ?? this.timeoutMs; return new Promise((resolve) => { const timer = setTimeout(() => { const pending = this.pending.get(rpcId); if (!pending) return; this.pending.delete(rpcId); - pending.resolve({ ok: false, error: `RPC ${method} timed out after ${this.timeoutMs}ms` }); - }, this.timeoutMs); + pending.resolve({ ok: false, error: `RPC ${method} timed out after ${effectiveTimeoutMs}ms` }); + }, effectiveTimeoutMs); this.pending.set(rpcId, { resolve, timer }); this.post({ type: "rpc-request", rpcId, runId, method, payload }); }); diff --git a/src/main/worker/worker-runner.ts b/src/main/worker/worker-runner.ts index 5e6ab7f3..fa165429 100644 --- a/src/main/worker/worker-runner.ts +++ b/src/main/worker/worker-runner.ts @@ -110,6 +110,13 @@ export class WorkerRunner { const cwd = job.workspacePath ?? process.cwd(); const tools = createCodingTools(cwd) as never; + // Co-Founder mode gates tool calls through the HumanApprovalHook in + // main. CEO mode skips the hook entirely (zero IPC overhead per tool + // call). Two distinct Agent-construction paths instead of a runtime + // `if mode === ...` inside a shared hook body. + const gateToolCalls = job.executionMode === "ai_cofounder"; + const TOOL_APPROVAL_TIMEOUT_MS = 30 * 60 * 1000; + const agent = new Agent({ initialState: { systemPrompt: job.systemPrompt ?? "", @@ -123,6 +130,29 @@ export class WorkerRunner { apiKey: prepared.apiKey, headers: { ...(options?.headers ?? {}), ...(extraHeaders ?? {}) }, }), + ...(gateToolCalls + ? { + beforeToolCall: async (context) => { + const verdict = await rpc.invoke( + job.runId, + "beforeToolCall", + { + agentId: job.agentId, + companyId: job.companyId, + taskId: job.taskId, + toolCallId: context.toolCall.id, + toolName: context.toolCall.name, + args: context.args, + }, + { timeoutMs: TOOL_APPROVAL_TIMEOUT_MS }, + ); + if (verdict.ok) { + return verdict.result as { block?: boolean; reason?: string } | undefined; + } + return { block: true, reason: verdict.error }; + }, + } + : {}), }); const unsubscribe = agent.subscribe(async (event) => { diff --git a/src/main/worker/worker-types.ts b/src/main/worker/worker-types.ts index 0669786e..b2311d96 100644 --- a/src/main/worker/worker-types.ts +++ b/src/main/worker/worker-types.ts @@ -50,7 +50,7 @@ export const workerMessageSchema = z.union([ type: z.literal("rpc-request"), rpcId: z.string(), runId: z.string().uuid(), - method: z.enum(["persistMessage", "prepareRun"]), + method: z.enum(["persistMessage", "prepareRun", "beforeToolCall"]), payload: z.unknown(), }), ]);