From e1693899a2b364f4119208fae6bdc9e10a157529 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sun, 19 Apr 2026 00:48:13 -0400 Subject: [PATCH] feat(runtime): wire pi-agent-core hooks end-to-end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The worker previously constructed `new Agent({systemPrompt, model, tools, messages, streamFn})` and stopped there — transformContext, beforeToolCall, afterToolCall, sessionId, thinkingLevel, and steer/ followUp were all dark. Approval enforcement, artifact projection, auto-compaction, and mid-run steering all had to reach the agent through other, non-canonical channels (event observers, prompt-block injection at prepareRun, API-layer gates). That left the whole pi contract at one-tenth utilisation and forced business logic into the wrong seams. This change rebuilds the contract assembly layer: - `pi-agent-factory` builds the Agent with the full option surface: sessionId (taskSessionId), thinkingLevel (mapped from agent.metadataJson), and three RPC-proxied hooks. - Each hook shuttles a thin payload over the existing worker RPC channel into `RuntimeBridge`, which routes to `hook-handlers`: * beforeToolCall → approval-gate policy (blocks destructive tools while a pending approval is outstanding). * afterToolCall → filesystem artifact projection, previously nested in the tool_execution_end event observer. * transformContext → pass-through today; compaction wires in next commit. - ParentMessage gains `steer-run` / `followup-run`; `WorkerRunner` keeps a runId→Agent registry so main can inject mid-run messages into the active pi loop. `worker-types` extends the rpc-request method enum; `runtime-bridge` loses the per-event artifact projector and gains a `findActiveRunForAgent` lookup used by the upcoming inbox → steer path. No public main API changes yet — the wiring is the foundation for the next five commits. --- src/main/agent-runtime/hook-handlers.ts | 142 +++++++++++ src/main/agent-runtime/hook-rpc.ts | 69 ++++++ src/main/agent-runtime/runtime-bridge.ts | 286 ++++++++++++----------- src/main/worker/pi-agent-factory.ts | 131 +++++++++++ src/main/worker/worker-entry.ts | 8 + src/main/worker/worker-helpers.ts | 9 + src/main/worker/worker-runner.ts | 61 +++-- src/main/worker/worker-types.ts | 8 +- 8 files changed, 561 insertions(+), 153 deletions(-) create mode 100644 src/main/agent-runtime/hook-handlers.ts create mode 100644 src/main/agent-runtime/hook-rpc.ts create mode 100644 src/main/worker/pi-agent-factory.ts diff --git a/src/main/agent-runtime/hook-handlers.ts b/src/main/agent-runtime/hook-handlers.ts new file mode 100644 index 00000000..723be0ea --- /dev/null +++ b/src/main/agent-runtime/hook-handlers.ts @@ -0,0 +1,142 @@ +/** + * Main-side implementations of pi's tool-execution hooks. Invoked by + * `RuntimeBridge.handleRpc` when the worker proxies `beforeToolCall`, + * `afterToolCall`, or `transformContext`. + * + * Each handler is deliberately isolated to one concern: + * - `beforeToolCall`: approval-gate policy (blocks destructive tools + * while an approval or wait state is outstanding) + * - `afterToolCall`: filesystem artifact projection (bash/write/edit + * outputs become `work_products`) + * - `transformContext`: auto-compaction hook (pruning in future; today + * passes through) + * + * Keeping them in this module pulls policy out of `RuntimeBridge` + * (which stays a pure IPC router) and keeps the per-file budget under + * the project's 500-line cap even as policy grows. + */ + +import type { AppDatabase } from "@main/database"; +import type { + AfterToolCallRpcPayload, + AfterToolCallRpcResult, + BeforeToolCallRpcPayload, + BeforeToolCallRpcResult, + TransformContextRpcPayload, + TransformContextRpcResult, +} from "./hook-rpc"; +import { + projectToolExecutionToWorkProduct, + type ArtifactProjectorRunContext, +} from "./artifact-projector"; + +export interface ActiveRunContext { + runId: string; + agentId: string; + companyId: string; + taskId: string | null; + connectorId: string; + modelId: string; + taskSessionId: string; + workspacePath: string | null; +} + +export interface HookHandlerDeps { + db: AppDatabase; + logger: { + info: (...args: unknown[]) => void; + warn: (...args: unknown[]) => void; + error: (...args: unknown[]) => void; + }; +} + +const NO_OP_AFTER: AfterToolCallRpcResult = { + contentOverride: null, + isErrorOverride: null, +}; + +/** + * Tools that can mutate the workspace or execute host-level effects. + * Every other tool (read/grep/glob) is always permitted — gating them + * would only slow observation. + */ +const DESTRUCTIVE_TOOLS: ReadonlySet = new Set(["bash", "write", "edit"]); + +export function isDestructiveTool(toolName: string): boolean { + return DESTRUCTIVE_TOOLS.has(toolName); +} + +/** + * Approval gate. A pending approval for this agent's current task + * blocks every destructive tool — the agent must wait for the board + * (or user) to resolve the approval before mutating anything. + */ +export function handleBeforeToolCall( + deps: HookHandlerDeps, + ctx: ActiveRunContext, + payload: BeforeToolCallRpcPayload, +): BeforeToolCallRpcResult { + if (!isDestructiveTool(payload.toolName)) { + return { block: false, reason: null }; + } + + const blocking = deps.db.getBlockingApproval(ctx.companyId, ctx.agentId, ctx.taskId); + if (blocking) { + return { + block: true, + reason: `Approval ${blocking.id} (${blocking.state}) is outstanding: ${blocking.payloadSummary}. Wait for the board to decide before retrying ${payload.toolName}.`, + }; + } + + return { block: false, reason: null }; +} + +/** + * Post-execution projection. Today only filesystem artifacts flow + * through this hook; as further projections (browser actions, social + * posts) migrate off the event-observer path they plug in here. + */ +export function handleAfterToolCall( + deps: HookHandlerDeps, + ctx: ActiveRunContext, + payload: AfterToolCallRpcPayload, +): AfterToolCallRpcResult { + const runCtx: ArtifactProjectorRunContext = { + runId: ctx.runId, + agentId: ctx.agentId, + companyId: ctx.companyId, + taskId: ctx.taskId, + workspacePath: ctx.workspacePath, + }; + try { + projectToolExecutionToWorkProduct( + deps.db, + runCtx, + { + toolName: payload.toolName, + args: payload.args, + result: payload.resultContent, + isError: payload.isError, + }, + ); + } catch (error) { + deps.logger.warn( + `[hook-handlers] artifact projection failed for run ${ctx.runId} tool ${payload.toolName}: ${error instanceof Error ? error.message : String(error)}`, + ); + } + return NO_OP_AFTER; +} + +/** + * Context-transform hook. Placeholder pass-through; auto-compaction + * wiring lives in session-management-service and plugs in here once + * the compactor is available (tracked separately from the contract + * wiring). + */ +export function handleTransformContext( + _deps: HookHandlerDeps, + _ctx: ActiveRunContext, + _payload: TransformContextRpcPayload, +): TransformContextRpcResult { + return { changed: false, messages: null }; +} diff --git a/src/main/agent-runtime/hook-rpc.ts b/src/main/agent-runtime/hook-rpc.ts new file mode 100644 index 00000000..fd731adb --- /dev/null +++ b/src/main/agent-runtime/hook-rpc.ts @@ -0,0 +1,69 @@ +/** + * Hook RPC contracts — the parent↔worker half of pi-agent-core's + * `beforeToolCall` / `afterToolCall` / `transformContext` hooks. + * + * The Agent instance lives in the utility process. Policy (approval, + * compaction, filesystem projection) lives in main. The hooks are + * therefore proxied: the worker invokes an RPC that main answers from + * the same `RuntimeBridge` that already owns agent-event and + * `persistMessage` handling. + * + * Payloads intentionally exclude pi's `AgentContext` (systemPrompt, + * messages, tools) because main policy never needs it. Stripping it + * keeps each hook IPC well under the structured-clone budget even for + * large transcripts. + */ + +import type { AgentMessage } from "@mariozechner/pi-agent-core"; + +// ─── beforeToolCall ─────────────────────────────────────────────────────── + +export interface BeforeToolCallRpcPayload { + toolName: string; + toolCallId: string; + args: unknown; +} + +export interface BeforeToolCallRpcResult { + block: boolean; + reason: string | null; +} + +// ─── afterToolCall ──────────────────────────────────────────────────────── + +export interface AfterToolCallRpcPayload { + toolName: string; + toolCallId: string; + args: unknown; + /** Raw tool result content list as emitted by pi's tool executor. */ + resultContent: unknown; + /** Structured `details` payload if the tool exposed one. */ + resultDetails: unknown; + isError: boolean; +} + +/** + * Optional override for the executed tool result, mirroring pi's + * `AfterToolCallResult` shape. Main returns every field as `null` when + * the hook performs observation only (current projection use-case). + */ +export interface AfterToolCallRpcResult { + contentOverride: { type: "text"; text: string }[] | null; + isErrorOverride: boolean | null; +} + +// ─── transformContext ───────────────────────────────────────────────────── + +export interface TransformContextRpcPayload { + messages: AgentMessage[]; +} + +export interface TransformContextRpcResult { + /** + * Transformed transcript. When `changed` is false the worker keeps its + * original array instance instead of replacing it — saves structured + * clone on the common "no compaction yet" path. + */ + changed: boolean; + messages: AgentMessage[] | null; +} diff --git a/src/main/agent-runtime/runtime-bridge.ts b/src/main/agent-runtime/runtime-bridge.ts index 532deff2..602a2981 100644 --- a/src/main/agent-runtime/runtime-bridge.ts +++ b/src/main/agent-runtime/runtime-bridge.ts @@ -7,8 +7,18 @@ import { CostWriter } from "./cost-writer"; import type { CredentialResolver } from "@main/credentials"; import { getConnectorDefinition } from "@main/connectors"; import { projectInboxMessageToPi } from "./inbox-projection"; -import { projectToolExecutionToWorkProduct } from "./artifact-projector"; import { projectAgentEventToLogChunk, extractTextBody } from "./agent-event-projection"; +import { + handleBeforeToolCall, + handleAfterToolCall, + handleTransformContext, + type ActiveRunContext, +} from "./hook-handlers"; +import type { + AfterToolCallRpcPayload, + BeforeToolCallRpcPayload, + TransformContextRpcPayload, +} from "./hook-rpc"; export interface RuntimeBridgeDeps { db: AppDatabase; @@ -21,21 +31,12 @@ export interface RuntimeBridgeDeps { credentialResolver: CredentialResolver; } -interface ActiveRunContext { - runId: string; - agentId: string; - companyId: string; - taskId: string | null; - connectorId: string; - modelId: string; - taskSessionId: string; - /** - * Workspace root handed to pi's coding tools as cwd. Needed by the - * artifact projector to resolve tool-reported paths (which may be - * relative) and to reject writes that escape the workspace. - */ - workspacePath: string | null; -} +type RpcMethod = + | "persistMessage" + | "prepareRun" + | "beforeToolCall" + | "afterToolCall" + | "transformContext"; /** * Single owner of pi runtime side-effects in main. @@ -43,11 +44,12 @@ 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. + * - Proxies pi's tool-execution hooks (`beforeToolCall` / `afterToolCall` + * / `transformContext`) into policy handlers: approval-gate checks, + * filesystem artifact projection, and context compaction. * - * Co-Founder mode's human-review gates live at the API endpoint layer - * (hire_agent / document_review approvals created by `routes-*.ts`), not - * at the pi tool-call layer; the worker constructs `Agent` without a - * `beforeToolCall` hook in either mode. + * Policy itself lives in `hook-handlers.ts`; this class is a pure IPC + * router plus an `ActiveRunContext` cache. */ export class RuntimeBridge { private readonly costWriter: CostWriter; @@ -74,6 +76,17 @@ export class RuntimeBridge { if (ctx) this.sessionStores.delete(ctx.taskSessionId); } + /** Returns the active run context for a given agent, if any. Used by + * main-side services (inbox dispatcher) to decide whether a new inbox + * row should route to `agent.steer()` instead of waiting for the next + * `prepareRun` cycle. */ + findActiveRunForAgent(agentId: string): ActiveRunContext | null { + for (const ctx of this.activeRuns.values()) { + if (ctx.agentId === agentId) return ctx; + } + return null; + } + private getOrCreateStore(taskSessionId: string): SqliteSessionStore { const existing = this.sessionStores.get(taskSessionId); if (existing) return existing; @@ -128,10 +141,6 @@ export class RuntimeBridge { }); } - if (envelope.event.type === "tool_execution_end") { - this.projectToolArtifact(envelope); - } - if (envelope.event.type === "message_end") { const ctx = this.activeRuns.get(envelope.runId); const messageRaw = envelope.event.message as @@ -149,11 +158,6 @@ export class RuntimeBridge { usage: messageRaw.usage as never, }); } catch (error) { - // Cost writer fails closed (RUNTIME-105) — surface the failure - // loudly so operators can see a pricing catalog break, but let - // the run continue so the agent's work is not lost. The missing - // row means budget enforcement under-counts this message; a - // repeated failure in the same run is still visible in logs. this.deps.logger.error( `[runtime-bridge] cost recording failed for run ${ctx.runId}: ${error instanceof Error ? error.message : String(error)}`, ); @@ -163,49 +167,6 @@ export class RuntimeBridge { } } - /** - * Project a pi `tool_execution_end` event into the `work_products` - * table when the tool was `write` or `edit` and succeeded. This - * replaces the earlier behaviour where the agent had to remember to - * `POST /api/documents` after writing a file; agents forgot routinely, - * so the desktop UI showed an empty artifact list while the workspace - * filesystem filled up. The projector is idempotent and silent — all - * skip reasons are traced to the warn channel so operators can audit - * why a given write did not register. - */ - private projectToolArtifact(envelope: { - runId: string; - event: { type: string } & Record; - }): void { - const ctx = this.activeRuns.get(envelope.runId); - if (!ctx) return; - const e = envelope.event; - const toolName = typeof e.toolName === "string" ? e.toolName : null; - if (!toolName) return; - try { - projectToolExecutionToWorkProduct( - this.deps.db, - { - runId: ctx.runId, - agentId: ctx.agentId, - companyId: ctx.companyId, - taskId: ctx.taskId, - workspacePath: ctx.workspacePath, - }, - { - toolName, - args: e.args, - result: e.result, - isError: e.isError === true, - }, - ); - } catch (error) { - this.deps.logger.warn( - `[runtime-bridge] artifact projection failed for run ${envelope.runId} tool ${toolName}: ${error instanceof Error ? error.message : String(error)}`, - ); - } - } - /** * Project the pi `message_end` event into the task thread as a * `conversation_events` row so renderer thread views stay live. @@ -227,14 +188,22 @@ export class RuntimeBridge { try { const thread = ensureTaskThread(this.deps.db, envelope.taskId, envelope.companyId, task.title); const body = extractTextBody(message.content); - const actorType: "agent" | "user" | "system" = message.role === "user" + const isUserMessage = message.role === "user"; + const actorType: "agent" | "user" | "system" = isUserMessage ? "user" : message.role === "assistant" ? "agent" : "system"; + // User-role messages inside the transcript come from either a human + // (inbox projection → Board/User) or a peer agent (via steer). The + // envelope's agentId is the *recipient*, never the author; recording + // it there creates fake "CEO replied to CEO" rows. For user-authored + // rows we record the virtual user actor id; peer-authored user rows + // are labeled "system" since pi has already erased the sender. + const actorId = isUserMessage ? "user" : envelope.agentId; this.deps.db.addThreadEvent({ threadId: thread.id, - actorId: envelope.agentId, + actorId, actorType, eventType: message.role === "toolResult" ? "system_note" : "message", body: body.length > 4000 ? `${body.slice(0, 3997)}...` : body, @@ -252,75 +221,122 @@ export class RuntimeBridge { /** Handles RPC requests issued by the worker (`rpc-request` envelope). */ async handleRpc(request: { rpcId: string; - method: "persistMessage" | "prepareRun"; + method: RpcMethod; runId: string; payload: unknown; }): Promise<{ ok: true; result: unknown } | { ok: false; error: string }> { try { - if (request.method === "persistMessage") { - const payload = request.payload as { taskSessionId: string; message: AgentMessage }; - this.getOrCreateStore(payload.taskSessionId).appendMessage(payload.message); - return { ok: true, result: undefined }; - } - if (request.method === "prepareRun") { - const payload = request.payload as { - connectorId: string; - taskSessionId: string; - agentId: string; - companyId: string; - taskId: string | null; - modelId: string; - workspacePath: string | null; - }; - this.registerActiveRun({ - runId: request.runId, - agentId: payload.agentId, - companyId: payload.companyId, - taskId: payload.taskId, - connectorId: payload.connectorId, - modelId: payload.modelId, - taskSessionId: payload.taskSessionId, - workspacePath: payload.workspacePath, - }); - const definition = getConnectorDefinition(payload.connectorId); - const resolved = await this.deps.credentialResolver.resolve(payload.connectorId); - - // Project any pending inbox messages into the session transcript - // *before* loading priorMessages — the recipient must observe them - // as plain user turns in its context, and the transcript row is - // the single source of memory (recentMessages prompt block is - // gone). Append fails closed: if persistence throws, the inbox - // row stays unread and we surface the error so the worker aborts. - const store = this.getOrCreateStore(payload.taskSessionId); - const inboxRecords = this.deps.db.listAgentInboxUnread(payload.companyId, payload.agentId); - for (const record of inboxRecords) { - const projected = projectInboxMessageToPi(record, { - getAgentName: (id) => this.resolveAgentName(id), - }); - store.appendMessage(projected); + switch (request.method) { + case "persistMessage": { + const payload = request.payload as { taskSessionId: string; message: AgentMessage }; + this.getOrCreateStore(payload.taskSessionId).appendMessage(payload.message); + return { ok: true, result: undefined }; } - if (inboxRecords.length > 0) { - this.deps.db.markAgentInboxDelivered( - payload.agentId, - inboxRecords.map((r) => r.id), - ); - } - - const priorMessages = store.loadMessages(); - return { - ok: true, - result: { - provider: definition.provider, - apiKey: resolved.apiKey, - extraEnv: resolved.extraEnv ?? {}, - priorMessages, - }, - }; + case "prepareRun": + return await this.handlePrepareRunRpc(request.runId, request.payload); + case "beforeToolCall": + return this.handleBeforeToolCallRpc(request.runId, request.payload as BeforeToolCallRpcPayload); + case "afterToolCall": + return this.handleAfterToolCallRpc(request.runId, request.payload as AfterToolCallRpcPayload); + case "transformContext": + return this.handleTransformContextRpc(request.runId, request.payload as TransformContextRpcPayload); } - return { ok: false, error: `Unknown RPC method: ${(request as { method: string }).method}` }; } catch (error) { return { ok: false, error: error instanceof Error ? error.message : String(error) }; } } -} + private async handlePrepareRunRpc( + runId: string, + payloadRaw: unknown, + ): Promise<{ ok: true; result: unknown } | { ok: false; error: string }> { + const payload = payloadRaw as { + connectorId: string; + taskSessionId: string; + agentId: string; + companyId: string; + taskId: string | null; + modelId: string; + workspacePath: string | null; + }; + this.registerActiveRun({ + runId, + agentId: payload.agentId, + companyId: payload.companyId, + taskId: payload.taskId, + connectorId: payload.connectorId, + modelId: payload.modelId, + taskSessionId: payload.taskSessionId, + workspacePath: payload.workspacePath, + }); + const definition = getConnectorDefinition(payload.connectorId); + const resolved = await this.deps.credentialResolver.resolve(payload.connectorId); + + // Project any pending inbox messages into the session transcript + // before loading priorMessages — the recipient must observe them + // as plain user turns in its context. Append fails closed: if + // persistence throws, the inbox row stays unread and we surface + // the error so the worker aborts. + const store = this.getOrCreateStore(payload.taskSessionId); + const inboxRecords = this.deps.db.listAgentInboxUnread(payload.companyId, payload.agentId); + for (const record of inboxRecords) { + const projected = projectInboxMessageToPi(record, { + getAgentName: (id) => this.resolveAgentName(id), + }); + store.appendMessage(projected); + } + if (inboxRecords.length > 0) { + this.deps.db.markAgentInboxDelivered( + payload.agentId, + inboxRecords.map((r) => r.id), + ); + } + + const priorMessages = store.loadMessages(); + return { + ok: true, + result: { + provider: definition.provider, + apiKey: resolved.apiKey, + extraEnv: resolved.extraEnv ?? {}, + priorMessages, + }, + }; + } + + private handleBeforeToolCallRpc( + runId: string, + payload: BeforeToolCallRpcPayload, + ): { ok: true; result: unknown } | { ok: false; error: string } { + const ctx = this.activeRuns.get(runId); + if (!ctx) { + return { ok: false, error: `beforeToolCall: no active run ${runId}` }; + } + const result = handleBeforeToolCall(this.deps, ctx, payload); + return { ok: true, result }; + } + + private handleAfterToolCallRpc( + runId: string, + payload: AfterToolCallRpcPayload, + ): { ok: true; result: unknown } | { ok: false; error: string } { + const ctx = this.activeRuns.get(runId); + if (!ctx) { + return { ok: false, error: `afterToolCall: no active run ${runId}` }; + } + const result = handleAfterToolCall(this.deps, ctx, payload); + return { ok: true, result }; + } + + private handleTransformContextRpc( + runId: string, + payload: TransformContextRpcPayload, + ): { ok: true; result: unknown } | { ok: false; error: string } { + const ctx = this.activeRuns.get(runId); + if (!ctx) { + return { ok: false, error: `transformContext: no active run ${runId}` }; + } + const result = handleTransformContext(this.deps, ctx, payload); + return { ok: true, result }; + } +} diff --git a/src/main/worker/pi-agent-factory.ts b/src/main/worker/pi-agent-factory.ts new file mode 100644 index 00000000..a2627060 --- /dev/null +++ b/src/main/worker/pi-agent-factory.ts @@ -0,0 +1,131 @@ +/** + * Factory that assembles a `pi-agent-core` Agent from a prepared worker + * job. Encapsulates the full Pi contract (initial state, sessionId, + * thinkingLevel, beforeToolCall / afterToolCall / transformContext + * hooks, streamFn) so `worker-runner` stays focused on run lifecycle. + * + * The three tool-execution hooks proxy into main via the worker RPC + * client; main answers from `RuntimeBridge`. Pi's loop is tolerant of + * thrown hooks — a thrown `beforeToolCall` produces an error tool + * result, a thrown `afterToolCall` likewise — so RPC failures degrade + * to a deterministic tool-error instead of crashing the run. The + * `transformContext` contract explicitly forbids throwing, so on RPC + * failure we return the original messages unchanged. + */ + +import { + Agent, + type AfterToolCallResult, + type AgentMessage, + type BeforeToolCallResult, + type ThinkingLevel, +} from "@mariozechner/pi-agent-core"; +import { streamSimple, type Api, type Model } from "@mariozechner/pi-ai"; +import type { + AfterToolCallRpcPayload, + AfterToolCallRpcResult, + BeforeToolCallRpcPayload, + BeforeToolCallRpcResult, + TransformContextRpcPayload, + TransformContextRpcResult, +} from "@main/agent-runtime/hook-rpc"; +import type { WorkerRpcClient } from "./worker-rpc"; + +export interface PiAgentFactoryInput { + runId: string; + taskSessionId: string; + systemPrompt: string; + model: Model; + tools: never; + priorMessages: AgentMessage[]; + thinkingLevel: ThinkingLevel; + apiKey: string; + extraHeaders: Record | undefined; + rpc: WorkerRpcClient; +} + +/** Map the UI-facing `thinkingEffort` string onto pi's `ThinkingLevel`. */ +export function resolveThinkingLevel(effort: string | null | undefined): ThinkingLevel { + switch (effort) { + case "minimal": + return "minimal"; + case "low": + return "low"; + case "medium": + return "medium"; + case "high": + return "high"; + case "xhigh": + return "xhigh"; + default: + return "off"; + } +} + +export function createPiAgent(input: PiAgentFactoryInput): Agent { + const { runId, rpc } = input; + return new Agent({ + initialState: { + systemPrompt: input.systemPrompt, + model: input.model, + tools: input.tools, + messages: input.priorMessages, + thinkingLevel: input.thinkingLevel, + }, + sessionId: input.taskSessionId, + streamFn: (model, messages, options) => + streamSimple(model, messages, { + ...options, + apiKey: input.apiKey, + headers: { ...(options?.headers ?? {}), ...(input.extraHeaders ?? {}) }, + }), + beforeToolCall: async (ctx): Promise => { + const payload: BeforeToolCallRpcPayload = { + toolName: ctx.toolCall.name, + toolCallId: ctx.toolCall.id, + args: ctx.args, + }; + const response = await rpc.invoke(runId, "beforeToolCall", payload); + if (!response.ok) { + // RPC failure must fail closed for destructive tools: main is + // the source of truth for approval state, and silently running + // without it would bypass the gate. The loop converts the thrown + // error into a tool-error result and surfaces it to the agent. + throw new Error(`beforeToolCall RPC failed: ${response.error}`); + } + const result = response.result as BeforeToolCallRpcResult; + if (!result.block) return undefined; + return { block: true, reason: result.reason ?? undefined }; + }, + afterToolCall: async (ctx): Promise => { + const payload: AfterToolCallRpcPayload = { + toolName: ctx.toolCall.name, + toolCallId: ctx.toolCall.id, + args: ctx.args, + resultContent: ctx.result.content, + resultDetails: ctx.result.details, + isError: ctx.isError, + }; + const response = await rpc.invoke(runId, "afterToolCall", payload); + if (!response.ok) { + // After-hook failure is informational: the tool already ran, + // projection is a best-effort write. Drop the hook rather than + // corrupting the tool result. + return undefined; + } + const result = response.result as AfterToolCallRpcResult; + const override: AfterToolCallResult = {}; + if (result.contentOverride) override.content = result.contentOverride; + if (result.isErrorOverride !== null) override.isError = result.isErrorOverride; + return Object.keys(override).length === 0 ? undefined : override; + }, + transformContext: async (messages): Promise => { + const payload: TransformContextRpcPayload = { messages }; + const response = await rpc.invoke(runId, "transformContext", payload); + if (!response.ok) return messages; + const result = response.result as TransformContextRpcResult; + if (result.changed && result.messages) return result.messages; + return messages; + }, + }); +} diff --git a/src/main/worker/worker-entry.ts b/src/main/worker/worker-entry.ts index cd0e63ba..5b443aa8 100644 --- a/src/main/worker/worker-entry.ts +++ b/src/main/worker/worker-entry.ts @@ -128,6 +128,14 @@ parentPort.on("message", (event) => { cancelRun(message.runId); return; } + if (message.type === "steer-run") { + runner.steer(message.runId, message.message); + return; + } + if (message.type === "followup-run") { + runner.followUp(message.runId, message.message); + return; + } if (message.type === "rpc-response") { rpc.resolve(message.rpcId, message.response); return; diff --git a/src/main/worker/worker-helpers.ts b/src/main/worker/worker-helpers.ts index 03baf651..4c9aaadf 100644 --- a/src/main/worker/worker-helpers.ts +++ b/src/main/worker/worker-helpers.ts @@ -1,3 +1,4 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { RunStatus } from "@shared/types"; export interface RunJob { @@ -22,10 +23,18 @@ export interface RunJob { executionMode: "ai_ceo" | "ai_cofounder"; } +/** + * Parent → worker control-plane messages. `steer-run` and `followup-run` + * are the inward half of pi's mid-run injection contract: the worker owns + * the Agent instance, main owns the inbox DB, so main posts the enqueued + * user message here and the worker calls `agent.steer()` / `agent.followUp()`. + */ export type ParentMessage = | { type: "enqueue-run"; payload: RunJob } | { type: "cancel-run"; runId: string } | { type: "shutdown" } + | { type: "steer-run"; runId: string; message: AgentMessage } + | { type: "followup-run"; runId: string; message: AgentMessage } | { type: "rpc-response"; rpcId: string; response: { ok: true; result: unknown } | { ok: false; error: string } }; export type ChildMessage = diff --git a/src/main/worker/worker-runner.ts b/src/main/worker/worker-runner.ts index cbb8122e..c122bf41 100644 --- a/src/main/worker/worker-runner.ts +++ b/src/main/worker/worker-runner.ts @@ -1,11 +1,10 @@ -import { Agent } from "@mariozechner/pi-agent-core"; -import type { AgentMessage } from "@mariozechner/pi-agent-core"; -import { streamSimple } from "@mariozechner/pi-ai"; -import { createCodingTools } from "@mariozechner/pi-coding-agent"; +import type { Agent, AgentMessage } from "@mariozechner/pi-agent-core"; import { resolveModel, resolveExtraStreamHeaders } from "../model-resolver"; import type { RunStatus } from "@shared/types"; import type { ChildMessage, RunJob } from "./worker-helpers"; import { applyEnvOverrides, buildRedactor } from "./worker-helpers"; +import { createCodingTools } from "@mariozechner/pi-coding-agent"; +import { createPiAgent, resolveThinkingLevel } from "./pi-agent-factory"; import type { WorkerRpcClient } from "./worker-rpc"; interface PrepareRunResult { @@ -21,9 +20,21 @@ interface RunnerDeps { isCancelled: (runId: string) => boolean; } +/** + * Lifecycle owner for one-and-only-one active pi `Agent` per run. + * + * Responsibilities: + * - Calls `prepareRun` on main to materialise session + credentials. + * - Builds an `Agent` via `createPiAgent` with the full Pi contract. + * - Streams every pi `AgentEvent` into the worker → main pipe. + * - Persists each `message_end` message through `persistMessage` RPC. + * - Exposes `steer` / `followUp` to the worker entry so mid-run + * injections from main land on the in-process Agent instance. + */ export class WorkerRunner { private readonly abortControllers = new Map(); private readonly runPromises = new Map>(); + private readonly activeAgents = new Map(); constructor(private readonly deps: RunnerDeps) {} @@ -32,6 +43,22 @@ export class WorkerRunner { if (controller) controller.abort(); } + /** Enqueue a steering message onto the active Agent for this run. */ + steer(runId: string, message: AgentMessage): boolean { + const agent = this.activeAgents.get(runId); + if (!agent) return false; + agent.steer(message); + return true; + } + + /** Enqueue a follow-up message onto the active Agent for this run. */ + followUp(runId: string, message: AgentMessage): boolean { + const agent = this.activeAgents.get(runId); + if (!agent) return false; + agent.followUp(message); + return true; + } + /** * Aborts every in-flight run and returns a promise that resolves once * each `run()` has settled (success, failure, or abort). Callers (the @@ -110,20 +137,19 @@ export class WorkerRunner { const cwd = job.workspacePath ?? process.cwd(); const tools = createCodingTools(cwd) as never; - const agent = new Agent({ - initialState: { - systemPrompt: job.systemPrompt ?? "", - model, - tools, - messages: prepared.priorMessages, - }, - streamFn: (m, messages, options) => - streamSimple(m, messages, { - ...options, - apiKey: prepared.apiKey, - headers: { ...(options?.headers ?? {}), ...(extraHeaders ?? {}) }, - }), + const agent = createPiAgent({ + runId: job.runId, + taskSessionId: job.taskSessionId, + systemPrompt: job.systemPrompt ?? "", + model, + tools, + priorMessages: prepared.priorMessages, + thinkingLevel: resolveThinkingLevel(job.thinkingEffort), + apiKey: prepared.apiKey, + extraHeaders, + rpc, }); + this.activeAgents.set(job.runId, agent); const unsubscribe = agent.subscribe(async (event) => { const timestamp = new Date().toISOString(); @@ -167,6 +193,7 @@ export class WorkerRunner { await runPromise; } finally { unsubscribe(); + this.activeAgents.delete(job.runId); } } finally { restoreProviderEnv(); diff --git a/src/main/worker/worker-types.ts b/src/main/worker/worker-types.ts index cc4010e0..4fc90917 100644 --- a/src/main/worker/worker-types.ts +++ b/src/main/worker/worker-types.ts @@ -47,7 +47,13 @@ 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", + "afterToolCall", + "transformContext", + ]), payload: z.unknown(), }), ]);