ZacharyZhang-NY/AgentCompany
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

Enable lock-free parallel agent runs

Browse Source
This commit is contained in:
ZacharyZhang-NY committed 2026-04-22 20:23:44 -04:00
1 parent d13c9d83e8
commit 1a7eb8f1b3
18 files changed
+327 -194
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/main/worker/entry/index.ts
+6 -5
View File
@@ -15,11 +15,12 @@ interface ActiveRun {
}
/**
* Single source of truth for scheduling invariants: "one agent at a time"
* and "one exclusive-workspace lock at a time". Prior layout kept three
* parallel structures (`activeByAgent`, `activeByWorkspace`, queue) that
* had to be updated in lockstep. The Map-of-ActiveRun keyed by runId is
* the only authoritative record; agent/workspace busy checks iterate it.
* Single source of truth for scheduling invariants: one active run per
* agent and one active run per exclusive workspace lock. Prior layout
* kept three parallel structures (`activeByAgent`, `activeByWorkspace`,
* queue) that had to be updated in lockstep. The Map-of-ActiveRun keyed
* by runId is the only authoritative record; agent/workspace busy checks
* iterate it.
*/
const active = new Map<string, ActiveRun>();
const queue: RunJob[] = [];
src/main/worker/runtime/helpers.ts
-31
View File
@@ -109,34 +109,3 @@ export function buildRedactor(env: Record<string, string>): (text: string) => st
return result;
};
}
/**
* Pushes extra env vars onto `process.env` and returns a restore function
* that is **infallible** — each individual restore is guarded so a single
* malformed key cannot leave a partial restore that keeps a provider
* secret in the worker's env for the next run. Intended for short-lived
* turn-scope env (`AGENT_COMPANY_*`, connector secrets). The worker runs
* one agent at a time, so this is safe for the current scheduler; do not
* convert it to concurrent run isolation — pass per-spawn env at the
* child-process level instead.
*/
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)) {
try {
if (previous === undefined) delete process.env[key];
else process.env[key] = previous;
} catch {
// Individual restore failure is never fatal — move on so the
// remaining keys still revert. Should never happen in practice
// (the Node env-var assignment is synchronous and infallible for
// string values), but the guard lets the outer finally complete.
}
}
};
}
src/main/worker/runtime/index.ts
+104 -115
View File
@@ -2,12 +2,12 @@ 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 "./helpers";
import { applyEnvOverrides, buildRedactor } from "./helpers";
import { createCodingTools } from "@mariozechner/pi-coding-agent";
import { buildRedactor } from "./helpers";
import { createPiAgent, resolveThinkingLevel, type ErasedAgentTool } from "./pi-agent-factory";
import type { WorkerRpcClient } from "../entry/rpc";
import { createHireAgentTool, HIRE_AGENT_TOOL_NAME } from "./internal-tools/hire-agent-tool";
import { createFileApprovalTool, FILE_APPROVAL_TOOL_NAME } from "./internal-tools/file-approval-tool";
import { createRunScopedCodingTools, mergeRunEnvs } from "./run-tools";
import {
classifyRunUsage,
hasAttributableAssistantUsage,
@@ -95,7 +95,7 @@ export class WorkerRunner {
private async runImpl(job: RunJob): Promise<void> {
const { post, rpc, isCancelled } = this.deps;
const redact = buildRedactor(job.env);
let redact = buildRedactor(job.env);
const abortController = new AbortController();
this.abortControllers.set(job.runId, abortController);
@@ -110,14 +110,6 @@ export class WorkerRunner {
let terminalStopReason: TerminalAssistantStopReason | null = null;
let errorMessage: string | null = null;
// Inject the AgentCompany job env (AGENT_COMPANY_API_URL / API_KEY /
// AGENT_ID / COMPANY_ID / RUN_ID / TASK_ID / WAKE_REASON + provider
// secret bindings) into the worker's process.env BEFORE Agent
// construction, so the agent's bash-tool children can call back into
// our HTTP API via fetch(process.env.AGENT_COMPANY_API_URL). Restored
// in the outer finally regardless of success / abort / throw.
const restoreJobEnv = applyEnvOverrides(job.env);
try {
if (!job.model) {
throw new Error(`Run ${job.runId} has no model selected.`);
@@ -132,134 +124,131 @@ export class WorkerRunner {
modelId: job.model,
workspaceId: job.workspaceId ?? null,
workspacePath: job.workspacePath ?? null,
workspaceMode: job.workspaceMode,
});
if (!prepareResponse.ok) {
throw new Error(prepareResponse.error);
}
const prepared = prepareResponse.result as PrepareRunResult;
const runEnv = mergeRunEnvs(job.env, prepared.extraEnv);
redact = buildRedactor(runEnv);
const restoreProviderEnv = applyEnvOverrides(prepared.extraEnv);
try {
const model = resolveModel(prepared.provider, job.model);
if (!model) {
throw new Error(`Unknown model ${job.model} for provider ${prepared.provider}.`);
}
resolvedModelId = model.id;
const extraHeaders = resolveExtraStreamHeaders(prepared.provider, prepared.apiKey);
const model = resolveModel(prepared.provider, job.model);
if (!model) {
throw new Error(`Unknown model ${job.model} for provider ${prepared.provider}.`);
}
resolvedModelId = model.id;
const extraHeaders = resolveExtraStreamHeaders(prepared.provider, prepared.apiKey);
const cwd = job.workspacePath ?? process.cwd();
const codingTools = createCodingTools(cwd);
const internalTools: ErasedAgentTool[] = [];
for (const toolName of job.internalTools ?? []) {
switch (toolName) {
case HIRE_AGENT_TOOL_NAME:
internalTools.push(createHireAgentTool(rpc, job.runId));
break;
case FILE_APPROVAL_TOOL_NAME:
internalTools.push(createFileApprovalTool(rpc, job.runId));
break;
default: {
// Exhaustiveness guard: a new entry in the `internalTools`
// array without a matching case should be a compile-time
// regression.
const _exhaustive: never = toolName;
throw new Error(`Unknown internal tool requested: ${String(_exhaustive)}`);
}
const cwd = job.workspacePath ?? process.cwd();
const codingTools = createRunScopedCodingTools(cwd, runEnv);
const internalTools: ErasedAgentTool[] = [];
for (const toolName of job.internalTools ?? []) {
switch (toolName) {
case HIRE_AGENT_TOOL_NAME:
internalTools.push(createHireAgentTool(rpc, job.runId));
break;
case FILE_APPROVAL_TOOL_NAME:
internalTools.push(createFileApprovalTool(rpc, job.runId));
break;
default: {
// Exhaustiveness guard: a new entry in the `internalTools`
// array without a matching case should be a compile-time
// regression.
const _exhaustive: never = toolName;
throw new Error(`Unknown internal tool requested: ${String(_exhaustive)}`);
}
}
const tools: ErasedAgentTool[] = [...codingTools, ...internalTools];
}
const tools: ErasedAgentTool[] = [...codingTools, ...internalTools];
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,
onContextCompaction: (telemetry) => {
post({
type: "context-compaction",
runId: telemetry.runId,
emittedAt: telemetry.emittedAt,
estimatedTokens: telemetry.estimatedTokens,
originalMessageCount: telemetry.originalMessageCount,
compactedMessageCount: telemetry.compactedMessageCount,
compactedOlderMessageCount: telemetry.compactedOlderMessageCount,
retainedRecentMessageCount: telemetry.retainedRecentMessageCount,
});
},
});
this.activeAgents.set(job.runId, agent);
post({ type: "run-agent-state", runId: job.runId, state: "agent_active" });
const unsubscribe = agent.subscribe(async (event) => {
const timestamp = new Date().toISOString();
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,
onContextCompaction: (telemetry) => {
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>,
type: "context-compaction",
runId: telemetry.runId,
emittedAt: telemetry.emittedAt,
estimatedTokens: telemetry.estimatedTokens,
originalMessageCount: telemetry.originalMessageCount,
compactedMessageCount: telemetry.compactedMessageCount,
compactedOlderMessageCount: telemetry.compactedOlderMessageCount,
retainedRecentMessageCount: telemetry.retainedRecentMessageCount,
});
},
});
this.activeAgents.set(job.runId, agent);
post({ type: "run-agent-state", runId: job.runId, state: "agent_active" });
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;
stopReason?: string;
errorMessage?: string;
usage?: { input?: number; output?: number; cost?: { total?: number } };
};
if (maybeAssistant.role === "assistant" && isTerminalAssistantStopReason(maybeAssistant.stopReason)) {
terminalStopReason = maybeAssistant.stopReason;
errorMessage = maybeAssistant.errorMessage ?? `Agent stopped with ${maybeAssistant.stopReason}.`;
}
if (hasAttributableAssistantUsage(maybeAssistant)) {
sawAttributableUsage = true;
finalUsage = {
inputTokens: typeof maybeAssistant.usage.input === "number" ? maybeAssistant.usage.input : finalUsage.inputTokens,
outputTokens: typeof maybeAssistant.usage.output === "number" ? maybeAssistant.usage.output : finalUsage.outputTokens,
};
}
}
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>,
});
try {
const promptContent = job.initialUserTurn.trim();
if (promptContent.length === 0) {
throw new Error(`Run ${job.runId} has an empty initial user turn.`);
}
const promptMessage: AgentMessage = {
role: "user",
content: [{ type: "text", text: promptContent }],
timestamp: Date.now(),
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;
stopReason?: string;
errorMessage?: string;
usage?: { input?: number; output?: number; cost?: { total?: number } };
};
const runPromise = agent.prompt(promptMessage);
abortController.signal.addEventListener("abort", () => agent.abort(), { once: true });
await runPromise;
} finally {
post({ type: "run-agent-state", runId: job.runId, state: "closing" });
unsubscribe();
this.activeAgents.delete(job.runId);
if (maybeAssistant.role === "assistant" && isTerminalAssistantStopReason(maybeAssistant.stopReason)) {
terminalStopReason = maybeAssistant.stopReason;
errorMessage = maybeAssistant.errorMessage ?? `Agent stopped with ${maybeAssistant.stopReason}.`;
}
if (hasAttributableAssistantUsage(maybeAssistant)) {
sawAttributableUsage = true;
finalUsage = {
inputTokens: typeof maybeAssistant.usage.input === "number" ? maybeAssistant.usage.input : finalUsage.inputTokens,
outputTokens: typeof maybeAssistant.usage.output === "number" ? maybeAssistant.usage.output : finalUsage.outputTokens,
};
}
}
});
try {
const promptContent = job.initialUserTurn.trim();
if (promptContent.length === 0) {
throw new Error(`Run ${job.runId} has an empty initial user turn.`);
}
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 {
restoreProviderEnv();
post({ type: "run-agent-state", runId: job.runId, state: "closing" });
unsubscribe();
this.activeAgents.delete(job.runId);
}
} 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);
restoreJobEnv();
}
const status: RunStatus = isCancelled(job.runId)
src/main/worker/runtime/run-tools.ts
+35
View File
@@ -0,0 +1,35 @@
import { createCodingTools, type BashToolOptions } from "@mariozechner/pi-coding-agent";
import type { AgentState } from "@mariozechner/pi-agent-core";
type ErasedAgentTool = AgentState["tools"][number];
export function mergeRunEnvs(
jobEnv: Record<string, string>,
preparedEnv: Record<string, string>,
): Record<string, string> {
return { ...jobEnv, ...preparedEnv };
}
export function createRunScopedBashOptions(
runEnv: Record<string, string>,
): BashToolOptions {
const envSnapshot = { ...runEnv };
return {
spawnHook: (context) => ({
...context,
env: {
...context.env,
...envSnapshot,
},
}),
};
}
export function createRunScopedCodingTools(
cwd: string,
runEnv: Record<string, string>,
): ErasedAgentTool[] {
return createCodingTools(cwd, {
bash: createRunScopedBashOptions(runEnv),
}) as ErasedAgentTool[];
}