refactor(worker): use pi-mono Agent directly
Rewrites the utility-process worker to drive pi-agent-core's Agent class against pi-ai models and pi-coding-agent tools. The entry file keeps the queue and cancellation state, delegates runs to WorkerRunner, and routes RPC responses through WorkerRpcClient so beforeToolCall/afterToolCall approvals and prepareRun credential resolution travel back to main's RuntimeBridge. Worker message schema gains `agent-event` + `rpc-request` envelopes, and run-finished no longer depends on the deleted sessionCodec.
This commit is contained in:
1 parent
892476bd36
commit
4de9cd0339
8 files changed
+404
-257
No files matched your search
@@ -20,6 +20,7 @@ export function createWorkerService({
|
||||
notify,
|
||||
budgetService,
|
||||
sessionService,
|
||||
getRuntimeBridge,
|
||||
}: WorkerServiceDependencies) {
|
||||
const MAX_RETRY_ATTEMPTS = 2;
|
||||
const RETRY_COOLDOWN_MS = 30_000; // 30 seconds
|
||||
@@ -110,7 +111,7 @@ export function createWorkerService({
|
||||
await handleWorkerMessage(
|
||||
message,
|
||||
worker,
|
||||
{ db, logger, isQuitting, emitEvent, wakeAgentIfPossible, processRunCompletion, onRunStarted, onRunCompleted, notify, budgetService, sessionService },
|
||||
{ db, logger, isQuitting, emitEvent, wakeAgentIfPossible, processRunCompletion, onRunStarted, onRunCompleted, notify, budgetService, sessionService, getRuntimeBridge },
|
||||
workerMessageSchema,
|
||||
destroyed,
|
||||
MAX_RETRY_ATTEMPTS,
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import type { z } from "zod";
|
||||
import { getConnectorDefinition } from "../../connectors";
|
||||
import { createCircuitBreaker } from "../circuit-breaker";
|
||||
import { handleFailureEscalation, handleSuccessReset } from "../failure-escalation";
|
||||
import { persistRunEvent } from "../run-event-recorder";
|
||||
@@ -53,8 +52,6 @@ export function finalizeRun(
|
||||
},
|
||||
});
|
||||
|
||||
const connector = getConnectorDefinition(run.connectorId);
|
||||
|
||||
let task: TaskRecord | null = null;
|
||||
try {
|
||||
task = db.getTask(run.taskId);
|
||||
@@ -63,17 +60,10 @@ export function finalizeRun(
|
||||
}
|
||||
|
||||
if (task) {
|
||||
if (event.sessionParams && connector.module.sessionCodec) {
|
||||
const serialized = connector.module.sessionCodec.serialize(event.sessionParams);
|
||||
db.updateTaskSession(
|
||||
task.id,
|
||||
serialized,
|
||||
connector.module.sessionCodec.getDisplayId?.(serialized ?? null) ?? event.sessionDisplayId ?? null,
|
||||
);
|
||||
}
|
||||
if (!event.sessionParams) {
|
||||
db.updateTaskSession(task.id, null, null);
|
||||
}
|
||||
// Pi-mono sessions are persisted through SqliteSessionStore via RuntimeBridge;
|
||||
// worker-reported sessionParams is an opaque JSON snapshot of whatever the
|
||||
// orchestrator wants to stamp onto the task_sessions row.
|
||||
db.updateTaskSession(task.id, event.sessionParams, event.sessionDisplayId ?? null);
|
||||
}
|
||||
|
||||
leaseRenewalMap.delete(event.runId);
|
||||
|
||||
+21
-240
@@ -1,74 +1,6 @@
|
||||
import { runningProcesses, shutdownAllRunning } from "@agentcompany/adapter-utils/server-utils";
|
||||
import { getConnectorDefinition } from "@main/connectors";
|
||||
import { buildSanitizedRuntimeEnv } from "@main/runtime-env";
|
||||
import type { ConnectorId, RunStatus } from "@shared/types";
|
||||
|
||||
function buildRedactor(env: Record<string, string>): (text: string) => string {
|
||||
const sensitivePatterns: string[] = [];
|
||||
for (const [key, value] of Object.entries(env)) {
|
||||
if (value.length < 8) continue;
|
||||
if (["PATH", "HOME", "USER", "SHELL", "LANG", "TMPDIR", "TERM"].includes(key)) continue;
|
||||
sensitivePatterns.push(value);
|
||||
}
|
||||
if (sensitivePatterns.length === 0) return (text) => text;
|
||||
sensitivePatterns.sort((a, b) => b.length - a.length);
|
||||
return (text: string) => {
|
||||
let result = text;
|
||||
for (const pattern of sensitivePatterns) {
|
||||
if (result.includes(pattern)) {
|
||||
result = result.replaceAll(pattern, "[REDACTED]");
|
||||
}
|
||||
}
|
||||
return result;
|
||||
};
|
||||
}
|
||||
|
||||
interface RunJob {
|
||||
runId: string;
|
||||
companyId: string;
|
||||
taskId: string;
|
||||
taskTitle: string;
|
||||
taskDescription: string;
|
||||
agentId: string;
|
||||
agentName: string;
|
||||
connectorId: ConnectorId;
|
||||
workspaceId: string | null;
|
||||
workspacePath: string | null;
|
||||
workspaceMode: "exclusive" | "shared" | "none";
|
||||
command: string;
|
||||
model: string | null;
|
||||
thinkingEffort: string | null;
|
||||
env: Record<string, string>;
|
||||
sessionParams: Record<string, unknown> | null;
|
||||
systemPrompt: string | null;
|
||||
trigger: string | null;
|
||||
managedSkills: Array<{ runtimeName: string; sourcePath: string }>;
|
||||
}
|
||||
|
||||
type ParentMessage =
|
||||
| { type: "enqueue-run"; payload: RunJob }
|
||||
| { type: "cancel-run"; runId: string }
|
||||
| { type: "shutdown" };
|
||||
|
||||
type ChildMessage =
|
||||
| { type: "run-log"; runId: string; stream: "stdout" | "stderr"; chunk: string }
|
||||
| { type: "run-status"; runId: string; status: RunStatus; message: string }
|
||||
| {
|
||||
type: "run-finished";
|
||||
runId: string;
|
||||
status: RunStatus;
|
||||
summary: string | null;
|
||||
errorMessage: string | null;
|
||||
exitCode: number | null;
|
||||
signal: string | null;
|
||||
model: string | null;
|
||||
sessionParams: Record<string, unknown> | null;
|
||||
sessionDisplayId: string | null;
|
||||
costUsd: number | null;
|
||||
unattributedCost: boolean;
|
||||
inputTokens: number | null;
|
||||
outputTokens: number | null;
|
||||
};
|
||||
import type { ChildMessage, ParentMessage, RunJob } from "./worker-helpers";
|
||||
import { WorkerRpcClient } from "./worker-rpc";
|
||||
import { WorkerRunner } from "./worker-runner";
|
||||
|
||||
const parentPort = process.parentPort;
|
||||
|
||||
@@ -80,74 +12,23 @@ const activeByAgent = new Map<string, string>();
|
||||
const activeByWorkspace = new Map<string, string>();
|
||||
const queue: RunJob[] = [];
|
||||
const cancelledRuns = new Set<string>();
|
||||
// S-01: Bypass is opt-in only. Packaged builds hard-refuse the dangerous CLI
|
||||
// flags unless the operator sets BOTH `AGENTCOMPANY_ALLOW_UNSAFE_CONNECTOR_BYPASS=true`
|
||||
// AND is running a dev/E2E build (AGENTCOMPANY_PACKAGED !== "true"). Main forks
|
||||
// the worker with AGENTCOMPANY_PACKAGED=<app.isPackaged>.
|
||||
const isPackaged = process.env.AGENTCOMPANY_PACKAGED === "true";
|
||||
const allowUnsafeConnectorBypass =
|
||||
!isPackaged && process.env.AGENTCOMPANY_ALLOW_UNSAFE_CONNECTOR_BYPASS === "true";
|
||||
|
||||
function post(message: ChildMessage) {
|
||||
parentPort.postMessage(message);
|
||||
}
|
||||
|
||||
const rpc = new WorkerRpcClient(post);
|
||||
const runner = new WorkerRunner({
|
||||
post,
|
||||
rpc,
|
||||
isCancelled: (runId) => cancelledRuns.has(runId),
|
||||
});
|
||||
|
||||
function workspaceKey(job: RunJob) {
|
||||
if (job.workspaceMode !== "exclusive") {
|
||||
return null;
|
||||
}
|
||||
if (job.workspaceMode !== "exclusive") return null;
|
||||
return job.workspaceId ?? job.workspacePath ?? `task:${job.taskId}`;
|
||||
}
|
||||
|
||||
function buildConnectorExecutionConfig(job: RunJob, sanitizedEnv: Record<string, string>) {
|
||||
const baseConfig: Record<string, unknown> = {
|
||||
command: job.command,
|
||||
cwd: job.workspacePath ?? process.cwd(),
|
||||
model: job.model,
|
||||
env: sanitizedEnv,
|
||||
managedSkills: job.managedSkills,
|
||||
};
|
||||
|
||||
switch (job.connectorId) {
|
||||
case "codex_local":
|
||||
return {
|
||||
...baseConfig,
|
||||
modelReasoningEffort: job.thinkingEffort ?? "",
|
||||
dangerouslyBypassApprovalsAndSandbox: allowUnsafeConnectorBypass,
|
||||
extraArgs: [
|
||||
"--skip-git-repo-check",
|
||||
"--sandbox",
|
||||
allowUnsafeConnectorBypass ? "danger-full-access" : "workspace-write",
|
||||
],
|
||||
};
|
||||
case "claude_local":
|
||||
case "claude_compatible":
|
||||
case "claude_compat_2":
|
||||
case "claude_compat_3":
|
||||
case "claude_compat_4":
|
||||
case "claude_compat_5":
|
||||
case "claude_compat_6":
|
||||
case "claude_compat_7":
|
||||
case "claude_compat_8":
|
||||
case "claude_compat_9":
|
||||
case "claude_compat_10":
|
||||
return {
|
||||
...baseConfig,
|
||||
effort: job.thinkingEffort ?? "",
|
||||
dangerouslySkipPermissions: allowUnsafeConnectorBypass,
|
||||
};
|
||||
case "gemini_local":
|
||||
return {
|
||||
...baseConfig,
|
||||
yolo: allowUnsafeConnectorBypass,
|
||||
approvalMode: allowUnsafeConnectorBypass ? "yolo" : "suggest",
|
||||
sandbox: true,
|
||||
};
|
||||
default:
|
||||
return baseConfig;
|
||||
}
|
||||
}
|
||||
|
||||
function enqueue(job: RunJob) {
|
||||
queue.push(job);
|
||||
post({
|
||||
@@ -172,115 +53,16 @@ async function maybeStartQueuedRuns() {
|
||||
}
|
||||
|
||||
async function startRun(job: RunJob) {
|
||||
const definition = getConnectorDefinition(job.connectorId);
|
||||
const lockKey = workspaceKey(job);
|
||||
activeByAgent.set(job.agentId, job.runId);
|
||||
if (lockKey) {
|
||||
activeByWorkspace.set(lockKey, job.runId);
|
||||
}
|
||||
post({
|
||||
type: "run-status",
|
||||
runId: job.runId,
|
||||
status: "running",
|
||||
message: `Running with ${definition.label}.`,
|
||||
});
|
||||
|
||||
const sanitizedEnv = buildSanitizedRuntimeEnv(
|
||||
Object.fromEntries(
|
||||
Object.entries(process.env).filter((e): e is [string, string] => typeof e[1] === "string"),
|
||||
),
|
||||
job.env,
|
||||
);
|
||||
|
||||
const redact = buildRedactor(job.env);
|
||||
if (lockKey) activeByWorkspace.set(lockKey, job.runId);
|
||||
|
||||
try {
|
||||
const result = await definition.module.execute({
|
||||
runId: job.runId,
|
||||
agent: {
|
||||
id: job.agentId,
|
||||
companyId: job.companyId,
|
||||
name: job.agentName,
|
||||
adapterType: job.connectorId,
|
||||
adapterConfig: {},
|
||||
},
|
||||
runtime: {
|
||||
sessionId:
|
||||
typeof job.sessionParams?.sessionId === "string" ? job.sessionParams.sessionId : null,
|
||||
sessionParams: job.sessionParams,
|
||||
sessionDisplayId:
|
||||
typeof job.sessionParams?.sessionId === "string" ? job.sessionParams.sessionId : null,
|
||||
taskKey: job.taskId,
|
||||
},
|
||||
config: buildConnectorExecutionConfig(job, sanitizedEnv),
|
||||
context: {
|
||||
taskId: job.taskId,
|
||||
taskTitle: job.systemPrompt ?? job.taskTitle,
|
||||
taskDescription: job.systemPrompt ? "" : job.taskDescription,
|
||||
agentCompanyWorkspace: job.workspaceId || job.workspacePath ? {
|
||||
workspaceId: job.workspaceId ?? "",
|
||||
cwd: job.workspacePath ?? "",
|
||||
source: "agent_company",
|
||||
} : undefined,
|
||||
},
|
||||
onLog: async (stream, chunk) => {
|
||||
post({
|
||||
type: "run-log",
|
||||
runId: job.runId,
|
||||
stream,
|
||||
chunk: redact(chunk),
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
const cancelled = cancelledRuns.has(job.runId);
|
||||
const finalStatus = cancelled
|
||||
? "cancelled"
|
||||
: result.timedOut
|
||||
? "timed_out"
|
||||
: result.errorMessage || result.exitCode
|
||||
? "failed"
|
||||
: "succeeded";
|
||||
|
||||
post({
|
||||
type: "run-finished",
|
||||
runId: job.runId,
|
||||
status: finalStatus,
|
||||
summary: result.summary ?? null,
|
||||
errorMessage: result.errorMessage ?? null,
|
||||
exitCode: result.exitCode,
|
||||
signal: result.signal,
|
||||
model: result.model ?? job.model,
|
||||
sessionParams: result.sessionParams ?? null,
|
||||
sessionDisplayId: result.sessionDisplayId ?? null,
|
||||
costUsd: typeof result.costUsd === "number" ? result.costUsd : null,
|
||||
unattributedCost: !result.costUsd,
|
||||
inputTokens: result.usage?.inputTokens ?? null,
|
||||
outputTokens: result.usage?.outputTokens ?? null,
|
||||
});
|
||||
} catch (error) {
|
||||
post({
|
||||
type: "run-finished",
|
||||
runId: job.runId,
|
||||
status: cancelledRuns.has(job.runId) ? "cancelled" : "failed",
|
||||
summary: null,
|
||||
errorMessage: error instanceof Error ? error.message : String(error),
|
||||
exitCode: null,
|
||||
signal: null,
|
||||
model: job.model,
|
||||
sessionParams: null,
|
||||
sessionDisplayId: null,
|
||||
costUsd: null,
|
||||
unattributedCost: true,
|
||||
inputTokens: null,
|
||||
outputTokens: null,
|
||||
});
|
||||
await runner.run(job);
|
||||
} finally {
|
||||
cancelledRuns.delete(job.runId);
|
||||
activeByAgent.delete(job.agentId);
|
||||
if (lockKey) {
|
||||
activeByWorkspace.delete(lockKey);
|
||||
}
|
||||
if (lockKey) activeByWorkspace.delete(lockKey);
|
||||
void maybeStartQueuedRuns();
|
||||
}
|
||||
}
|
||||
@@ -309,10 +91,7 @@ function cancelRun(runId: string) {
|
||||
}
|
||||
|
||||
cancelledRuns.add(runId);
|
||||
const running = runningProcesses.get(runId);
|
||||
if (running) {
|
||||
running.child.kill("SIGTERM");
|
||||
}
|
||||
runner.cancel(runId);
|
||||
}
|
||||
|
||||
parentPort.on("message", (event) => {
|
||||
@@ -325,11 +104,13 @@ parentPort.on("message", (event) => {
|
||||
cancelRun(message.runId);
|
||||
return;
|
||||
}
|
||||
if (message.type === "rpc-response") {
|
||||
rpc.resolve(message.rpcId, message.response);
|
||||
return;
|
||||
}
|
||||
if (message.type === "shutdown") {
|
||||
// Kill every running CLI process (claude, codex, gemini, etc.) and exit.
|
||||
// This is the only reliable way to ensure child processes don't survive app quit,
|
||||
// because the main process cannot kill grandchildren directly.
|
||||
shutdownAllRunning();
|
||||
runner.cancelAll();
|
||||
rpc.cancelAll("Worker shutting down.");
|
||||
process.exit(0);
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,96 @@
|
||||
import type { RunStatus } from "@shared/types";
|
||||
|
||||
export interface RunJob {
|
||||
runId: string;
|
||||
companyId: string;
|
||||
taskId: string;
|
||||
taskTitle: string;
|
||||
taskDescription: string;
|
||||
agentId: string;
|
||||
agentName: string;
|
||||
connectorId: string;
|
||||
workspaceId: string | null;
|
||||
workspacePath: string | null;
|
||||
workspaceMode: "exclusive" | "shared" | "none";
|
||||
command: string;
|
||||
model: string | null;
|
||||
thinkingEffort: string | null;
|
||||
env: Record<string, string>;
|
||||
sessionParams: Record<string, unknown> | null;
|
||||
systemPrompt: string | null;
|
||||
trigger: string | null;
|
||||
managedSkills: Array<{ runtimeName: string; sourcePath: string }>;
|
||||
taskSessionId: string;
|
||||
executionMode: "ai_ceo" | "ai_cofounder";
|
||||
}
|
||||
|
||||
export type ParentMessage =
|
||||
| { type: "enqueue-run"; payload: RunJob }
|
||||
| { type: "cancel-run"; runId: string }
|
||||
| { type: "shutdown" }
|
||||
| { type: "rpc-response"; rpcId: string; response: { ok: true; result: unknown } | { ok: false; error: string } };
|
||||
|
||||
export type ChildMessage =
|
||||
| { type: "run-log"; runId: string; stream: "stdout" | "stderr"; chunk: string }
|
||||
| { type: "run-status"; runId: string; status: RunStatus; message: string }
|
||||
| {
|
||||
type: "run-finished";
|
||||
runId: string;
|
||||
status: RunStatus;
|
||||
summary: string | null;
|
||||
errorMessage: string | null;
|
||||
exitCode: number | null;
|
||||
signal: string | null;
|
||||
model: string | null;
|
||||
sessionParams: Record<string, unknown> | null;
|
||||
sessionDisplayId: string | null;
|
||||
costUsd: number | null;
|
||||
unattributedCost: boolean;
|
||||
inputTokens: number | null;
|
||||
outputTokens: number | null;
|
||||
}
|
||||
| {
|
||||
type: "agent-event";
|
||||
runId: string;
|
||||
agentId: string;
|
||||
companyId: string;
|
||||
taskId: string | null;
|
||||
emittedAt: string;
|
||||
event: { type: string } & Record<string, unknown>;
|
||||
}
|
||||
| { type: "rpc-request"; rpcId: string; runId: string; method: string; payload: unknown };
|
||||
|
||||
export function buildRedactor(env: Record<string, string>): (text: string) => string {
|
||||
const sensitivePatterns: string[] = [];
|
||||
for (const [key, value] of Object.entries(env)) {
|
||||
if (value.length < 8) continue;
|
||||
if (["PATH", "HOME", "USER", "SHELL", "LANG", "TMPDIR", "TERM"].includes(key)) continue;
|
||||
sensitivePatterns.push(value);
|
||||
}
|
||||
if (sensitivePatterns.length === 0) return (text) => text;
|
||||
sensitivePatterns.sort((a, b) => b.length - a.length);
|
||||
return (text: string) => {
|
||||
let result = text;
|
||||
for (const pattern of sensitivePatterns) {
|
||||
if (result.includes(pattern)) {
|
||||
result = result.replaceAll(pattern, "[REDACTED]");
|
||||
}
|
||||
}
|
||||
return result;
|
||||
};
|
||||
}
|
||||
|
||||
/** Pushes extra env vars onto process.env and returns a restore function. */
|
||||
export function applyEnvOverrides(overrides: Record<string, string>): () => void {
|
||||
const saved: Record<string, string | undefined> = {};
|
||||
for (const [key, value] of Object.entries(overrides)) {
|
||||
saved[key] = process.env[key];
|
||||
process.env[key] = value;
|
||||
}
|
||||
return () => {
|
||||
for (const [key, previous] of Object.entries(saved)) {
|
||||
if (previous === undefined) delete process.env[key];
|
||||
else process.env[key] = previous;
|
||||
}
|
||||
};
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
import type { RuntimeBridge } from "../agent-runtime";
|
||||
import type { workerMessageSchema } from "./worker-types";
|
||||
import { handleRunLog } from "./handlers/run-log-handler";
|
||||
import { handleRunStatus } from "./handlers/run-status-handler";
|
||||
@@ -6,8 +7,8 @@ import type { HandlerContext, HandlerDeps } from "./handlers/handler-context";
|
||||
|
||||
export async function handleWorkerMessage(
|
||||
message: unknown,
|
||||
_worker: ReturnType<typeof import("electron").utilityProcess.fork> | null,
|
||||
deps: HandlerDeps,
|
||||
worker: ReturnType<typeof import("electron").utilityProcess.fork> | null,
|
||||
deps: HandlerDeps & { getRuntimeBridge?: () => RuntimeBridge | null },
|
||||
workerMessageSchemaRef: typeof workerMessageSchema,
|
||||
destroyed: boolean,
|
||||
MAX_RETRY_ATTEMPTS: number,
|
||||
@@ -44,5 +45,38 @@ export async function handleWorkerMessage(
|
||||
case "run-finished":
|
||||
handleRunFinished(event, deps, context);
|
||||
return;
|
||||
case "agent-event": {
|
||||
const bridge = deps.getRuntimeBridge?.();
|
||||
if (bridge) {
|
||||
bridge.handleAgentEvent({
|
||||
runId: event.runId,
|
||||
agentId: event.agentId,
|
||||
companyId: event.companyId,
|
||||
taskId: event.taskId,
|
||||
emittedAt: event.emittedAt,
|
||||
event: event.event,
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
case "rpc-request": {
|
||||
const bridge = deps.getRuntimeBridge?.();
|
||||
const response = bridge
|
||||
? await bridge.handleRpc({
|
||||
rpcId: event.rpcId,
|
||||
runId: event.runId,
|
||||
method: event.method,
|
||||
payload: event.payload,
|
||||
})
|
||||
: { ok: false as const, error: "RuntimeBridge not initialised." };
|
||||
try {
|
||||
worker?.postMessage({ type: "rpc-response", rpcId: event.rpcId, response });
|
||||
} catch (error) {
|
||||
deps.logger.error(
|
||||
`[worker] rpc-response postMessage failed: ${error instanceof Error ? error.message : String(error)}`,
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import type { ChildMessage } from "./worker-helpers";
|
||||
|
||||
interface PendingRpc {
|
||||
resolve: (value: { ok: true; result: unknown } | { ok: false; error: string }) => void;
|
||||
}
|
||||
|
||||
export class WorkerRpcClient {
|
||||
private readonly pending = new Map<string, PendingRpc>();
|
||||
|
||||
constructor(private readonly post: (message: ChildMessage) => void) {}
|
||||
|
||||
invoke(runId: string, method: string, payload: unknown): Promise<{ ok: true; result: unknown } | { ok: false; error: string }> {
|
||||
const rpcId = randomUUID();
|
||||
return new Promise((resolve) => {
|
||||
this.pending.set(rpcId, { resolve });
|
||||
this.post({ type: "rpc-request", rpcId, runId, method, payload });
|
||||
});
|
||||
}
|
||||
|
||||
resolve(rpcId: string, response: { ok: true; result: unknown } | { ok: false; error: string }): void {
|
||||
const pending = this.pending.get(rpcId);
|
||||
if (!pending) return;
|
||||
this.pending.delete(rpcId);
|
||||
pending.resolve(response);
|
||||
}
|
||||
|
||||
cancelAll(reason: string): void {
|
||||
for (const [id, pending] of this.pending) {
|
||||
pending.resolve({ ok: false, error: reason });
|
||||
this.pending.delete(id);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
import { Agent } from "@mariozechner/pi-agent-core";
|
||||
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||
import { getModel, streamSimple } from "@mariozechner/pi-ai";
|
||||
import { createCodingTools } from "@mariozechner/pi-coding-agent";
|
||||
import type { RunStatus } from "@shared/types";
|
||||
import type { ChildMessage, RunJob } from "./worker-helpers";
|
||||
import { applyEnvOverrides, buildRedactor } from "./worker-helpers";
|
||||
import type { WorkerRpcClient } from "./worker-rpc";
|
||||
|
||||
interface PrepareRunResult {
|
||||
provider: string;
|
||||
apiKey: string;
|
||||
extraEnv: Record<string, string>;
|
||||
priorMessages: AgentMessage[];
|
||||
}
|
||||
|
||||
interface RunnerDeps {
|
||||
post: (message: ChildMessage) => void;
|
||||
rpc: WorkerRpcClient;
|
||||
isCancelled: (runId: string) => boolean;
|
||||
}
|
||||
|
||||
export class WorkerRunner {
|
||||
private readonly abortControllers = new Map<string, AbortController>();
|
||||
|
||||
constructor(private readonly deps: RunnerDeps) {}
|
||||
|
||||
cancel(runId: string): void {
|
||||
const controller = this.abortControllers.get(runId);
|
||||
if (controller) controller.abort();
|
||||
}
|
||||
|
||||
cancelAll(): void {
|
||||
for (const controller of this.abortControllers.values()) {
|
||||
controller.abort();
|
||||
}
|
||||
this.abortControllers.clear();
|
||||
}
|
||||
|
||||
async run(job: RunJob): Promise<void> {
|
||||
const { post, rpc, isCancelled } = this.deps;
|
||||
const redact = buildRedactor(job.env);
|
||||
const abortController = new AbortController();
|
||||
this.abortControllers.set(job.runId, abortController);
|
||||
|
||||
post({ type: "run-status", runId: job.runId, status: "running", message: `Running agent ${job.agentName}.` });
|
||||
|
||||
let resolvedModelId: string | null = job.model;
|
||||
let finalUsage: { inputTokens: number | null; outputTokens: number | null; costUsd: number | null } = {
|
||||
inputTokens: null,
|
||||
outputTokens: null,
|
||||
costUsd: null,
|
||||
};
|
||||
let errorMessage: string | null = null;
|
||||
|
||||
try {
|
||||
if (!job.model) {
|
||||
throw new Error(`Run ${job.runId} has no model selected.`);
|
||||
}
|
||||
|
||||
const prepareResponse = await rpc.invoke(job.runId, "prepareRun", {
|
||||
connectorId: job.connectorId,
|
||||
taskSessionId: job.taskSessionId,
|
||||
agentId: job.agentId,
|
||||
companyId: job.companyId,
|
||||
taskId: job.taskId,
|
||||
modelId: job.model,
|
||||
});
|
||||
if (!prepareResponse.ok) {
|
||||
throw new Error(prepareResponse.error);
|
||||
}
|
||||
const prepared = prepareResponse.result as PrepareRunResult;
|
||||
|
||||
const restoreEnv = applyEnvOverrides(prepared.extraEnv);
|
||||
try {
|
||||
const model = getModel(prepared.provider as never, job.model as never);
|
||||
resolvedModelId = model.id;
|
||||
|
||||
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 }),
|
||||
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,
|
||||
});
|
||||
if (verdict.ok) {
|
||||
return verdict.result as { block?: boolean; reason?: string } | undefined;
|
||||
}
|
||||
return { block: true, reason: verdict.error };
|
||||
},
|
||||
afterToolCall: async (context) => {
|
||||
await rpc.invoke(job.runId, "afterToolCall", {
|
||||
toolCallId: context.toolCall.id,
|
||||
toolName: context.toolCall.name,
|
||||
isError: context.isError,
|
||||
});
|
||||
return undefined;
|
||||
},
|
||||
});
|
||||
|
||||
const unsubscribe = agent.subscribe(async (event) => {
|
||||
const timestamp = new Date().toISOString();
|
||||
post({
|
||||
type: "agent-event",
|
||||
runId: job.runId,
|
||||
agentId: job.agentId,
|
||||
companyId: job.companyId,
|
||||
taskId: job.taskId,
|
||||
emittedAt: timestamp,
|
||||
event: event as { type: string } & Record<string, unknown>,
|
||||
});
|
||||
|
||||
if (event.type === "message_end" && event.message) {
|
||||
await rpc.invoke(job.runId, "persistMessage", {
|
||||
taskSessionId: job.taskSessionId,
|
||||
message: event.message,
|
||||
});
|
||||
const maybeAssistant = event.message as { role?: string; usage?: { input?: number; output?: number; cost?: { total?: number } } };
|
||||
if (maybeAssistant.role === "assistant" && maybeAssistant.usage) {
|
||||
finalUsage = {
|
||||
inputTokens: typeof maybeAssistant.usage.input === "number" ? maybeAssistant.usage.input : finalUsage.inputTokens,
|
||||
outputTokens: typeof maybeAssistant.usage.output === "number" ? maybeAssistant.usage.output : finalUsage.outputTokens,
|
||||
costUsd: typeof maybeAssistant.usage.cost?.total === "number" ? maybeAssistant.usage.cost.total : finalUsage.costUsd,
|
||||
};
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
try {
|
||||
const userMessage = job.taskTitle
|
||||
? `${job.taskTitle}\n\n${job.taskDescription ?? ""}`.trim()
|
||||
: (job.taskDescription ?? "");
|
||||
const promptContent = userMessage.length > 0 ? userMessage : "Continue.";
|
||||
const promptMessage: AgentMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: promptContent }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
const runPromise = agent.prompt(promptMessage);
|
||||
abortController.signal.addEventListener("abort", () => agent.abort(), { once: true });
|
||||
await runPromise;
|
||||
} finally {
|
||||
unsubscribe();
|
||||
}
|
||||
} finally {
|
||||
restoreEnv();
|
||||
}
|
||||
} catch (error) {
|
||||
errorMessage = error instanceof Error ? error.message : String(error);
|
||||
post({ type: "run-log", runId: job.runId, stream: "stderr", chunk: redact(errorMessage) });
|
||||
} finally {
|
||||
this.abortControllers.delete(job.runId);
|
||||
}
|
||||
|
||||
const status: RunStatus = isCancelled(job.runId)
|
||||
? "cancelled"
|
||||
: errorMessage
|
||||
? "failed"
|
||||
: "succeeded";
|
||||
|
||||
post({
|
||||
type: "run-finished",
|
||||
runId: job.runId,
|
||||
status,
|
||||
summary: null,
|
||||
errorMessage,
|
||||
exitCode: null,
|
||||
signal: null,
|
||||
model: resolvedModelId,
|
||||
sessionParams: job.sessionParams,
|
||||
sessionDisplayId: null,
|
||||
costUsd: finalUsage.costUsd,
|
||||
unattributedCost: finalUsage.costUsd === null,
|
||||
inputTokens: finalUsage.inputTokens,
|
||||
outputTokens: finalUsage.outputTokens,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,7 @@ import type { DesktopEvent, RunRecord } from "@shared/types";
|
||||
import type { NotificationOptions } from "../notification-service";
|
||||
import type { createBudgetEnforcementService } from "../budget-enforcement-service";
|
||||
import type { createSessionManagementService } from "../session-management-service";
|
||||
import type { RuntimeBridge } from "../agent-runtime";
|
||||
|
||||
const runStatusSchema = z.enum(["queued", "running", "succeeded", "failed", "cancelled", "timed_out", "interrupted"]);
|
||||
|
||||
@@ -36,6 +37,22 @@ export const workerMessageSchema = z.union([
|
||||
inputTokens: z.number().nullable(),
|
||||
outputTokens: z.number().nullable(),
|
||||
}),
|
||||
z.object({
|
||||
type: z.literal("agent-event"),
|
||||
runId: z.string().uuid(),
|
||||
agentId: z.string(),
|
||||
companyId: z.string(),
|
||||
taskId: z.string().nullable(),
|
||||
emittedAt: z.string(),
|
||||
event: z.record(z.string(), z.unknown()).and(z.object({ type: z.string() })),
|
||||
}),
|
||||
z.object({
|
||||
type: z.literal("rpc-request"),
|
||||
rpcId: z.string(),
|
||||
runId: z.string().uuid(),
|
||||
method: z.enum(["beforeToolCall", "afterToolCall", "persistMessage", "prepareRun"]),
|
||||
payload: z.unknown(),
|
||||
}),
|
||||
]);
|
||||
|
||||
export interface WorkerServiceDependencies {
|
||||
@@ -70,4 +87,6 @@ export interface WorkerServiceDependencies {
|
||||
notify: (options: NotificationOptions) => void;
|
||||
budgetService?: ReturnType<typeof createBudgetEnforcementService>;
|
||||
sessionService?: ReturnType<typeof createSessionManagementService>;
|
||||
/** Forwards `agent-event` from worker to renderer and routes RPC to main services. */
|
||||
getRuntimeBridge?: () => RuntimeBridge | null;
|
||||
}
|
||||
Reference in new issue
Block a user