feat(runtime): wire pi-agent-core hooks end-to-end

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.
This commit is contained in:
ZacharyZhang-NY committed 2026-04-19 00:48:13 -04:00
1 parent f041336721
commit e1693899a2
8 files changed
+561 -153

No files matched your search

+44 -17
View File
@@ -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<string, AbortController>();
private readonly runPromises = new Map<string, Promise<void>>();
private readonly activeAgents = new Map<string, Agent>();
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();