fix(runtime): revive AI Co-Founder tool-call gating via pi beforeToolCall

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.
This commit is contained in:
ZacharyZhang-NY committed 2026-04-17 14:53:22 -04:00
1 parent b50716d4e8
commit 77f81a7c5a
8 files changed
+204 -8

No files matched your search

+136
View File
@@ -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<string>([
"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 "<unserialisable>";
}
}
/**
* 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<string, PendingApproval>();
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);
}
}
}
+1
View File
@@ -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";
+23 -4
View File
@@ -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<string, ActiveRunContext>();
private readonly sessionStores = new Map<string, SqliteSessionStore>();
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;
+6
View File
@@ -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);
};
+3 -1
View File
@@ -51,6 +51,8 @@ export function createPiRuntimeServices(deps: PiRuntimeServicesDeps): PiRuntimeS
credentialResolver,
oauthLoginService,
runtimeBridge,
release: () => {},
release: () => {
runtimeBridge.approvalHook.cancelAll("AgentCompany shutting down.");
},
};
}
+4 -2
View File
@@ -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 });
});
+30
View File
@@ -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) => {
+1 -1
View File
@@ -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(),
}),
]);