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:
ZacharyZhang-NY committed
1 parent
892476bd36
commit
4de9cd0339
8 files changed
+404
-257
No files matched your search
+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);
|
||||
}
|
||||
});
|
||||
Reference in new issue
Block a user