fix(orchestration): drain queued recovery blockers

This commit is contained in:
ZacharyZhang-NY committed 2026-04-28 23:47:02 -04:00
1 parent 7dbaaf0022
commit e34503df68
2 files changed
+249 -3

No files matched your search

+64 -3
View File
@@ -214,6 +214,31 @@ export function createTaskRunner(
}
const snapshot = db.projection.snapshot.listSnapshot();
const companiesById = new Map(snapshot.companies.map((company) => [company.id, company]));
const reservedAgentIds = new Set<string>();
const reservedTaskIds = new Set<string>();
function finishQueuedRecoveryRun(
run: ReturnType<typeof db.core.runs.listQueuedRuns>[number],
status: "cancelled" | "interrupted",
summary: string,
errorMessage: string | null,
): void {
db.core.runs.finishRun({
runId: run.id,
status,
summary,
errorMessage,
exitCode: null,
signal: null,
model: null,
unattributedCost: true,
recoverTaskForRetry: run.taskId !== null,
});
if (run.workspaceId) {
db.core.workspaces.releaseWorkspaceLock(run.workspaceId, run.id);
}
emitEvent({ type: "run-updated", runId: run.id, action: "completed" });
}
for (const run of db.core.runs.listQueuedRuns()) {
const company = companiesById.get(run.companyId);
@@ -221,13 +246,23 @@ export function createTaskRunner(
continue;
}
const hasBlockingSiblingRun = snapshot.runs.some((entry) =>
const hasRunningSiblingRun = snapshot.runs.some((entry) =>
entry.id !== run.id &&
entry.companyId === run.companyId &&
["queued", "running"].includes(entry.status) &&
entry.status === "running" &&
(entry.taskId === run.taskId || entry.agentId === run.agentId),
);
if (hasBlockingSiblingRun) {
if (hasRunningSiblingRun) {
continue;
}
if (reservedAgentIds.has(run.agentId) || (run.taskId !== null && reservedTaskIds.has(run.taskId))) {
finishQueuedRecoveryRun(
run,
"cancelled",
"Queued run superseded during recovery.",
null,
);
continue;
}
@@ -257,6 +292,21 @@ export function createTaskRunner(
detail: `Recovered queued run blocked by connector/model policy for task "${task?.title ?? "(heartbeat)"}"`,
});
emitEvent({ type: "run-updated", runId: run.id, action: "completed" });
} else {
finishQueuedRecoveryRun(
run,
"cancelled",
"Queued run recovery deferred by eligibility gate.",
eligibilityFailure,
);
db.core.activity.addActivity({
companyId: run.companyId,
actor: "system",
action: "run.recovery_deferred",
entityType: "run",
entityId: run.id,
detail: `Recovered queued run deferred by ${eligibilityFailure} for task "${task?.title ?? "(heartbeat)"}"`,
});
}
continue;
}
@@ -275,8 +325,19 @@ export function createTaskRunner(
snapshot,
trigger: "recovery",
});
const currentRun = db.core.runs.getRun(run.id);
if (currentRun.status === "queued" || currentRun.status === "running") {
reservedAgentIds.add(run.agentId);
if (run.taskId !== null) reservedTaskIds.add(run.taskId);
}
} catch (error) {
logger.warn(`[run] Skipped queued run recovery for ${run.id}: ${error instanceof Error ? error.message : String(error)}`);
finishQueuedRecoveryRun(
run,
"interrupted",
"Queued run recovery failed.",
error instanceof Error ? error.message : String(error),
);
}
}
}
+185
View File
@@ -0,0 +1,185 @@
import { mkdtemp, rm } from "node:fs/promises";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { afterEach, describe, expect, it, vi } from "vitest";
import { AppDatabase } from "../../src/main/database";
import { createBudgetEnforcementService } from "../../src/main/budget-enforcement-service";
import { createRunDispatch } from "../../src/main/orchestration/run-dispatch";
import { createTaskRunner } from "../../src/main/orchestration/task-runner";
vi.mock("electron", () => ({
safeStorage: {
isEncryptionAvailable: () => false,
getSelectedStorageBackend: () => "basic_text",
encryptString: (value: string) => Buffer.from(value, "utf8"),
decryptString: (value: Buffer) => value.toString("utf8"),
},
}));
const tempDirs: string[] = [];
async function createTempDir(prefix: string) {
const dir = await mkdtemp(join(tmpdir(), prefix));
tempDirs.push(dir);
return dir;
}
afterEach(async () => {
await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
});
function createLogger() {
return {
debug: vi.fn(),
info: vi.fn(),
warn: vi.fn(),
error: vi.fn(),
};
}
function createHarness(db: AppDatabase) {
const logger = createLogger();
const enqueueWorkerRun = vi.fn();
const budgetService = createBudgetEnforcementService({
db,
logger,
notify: vi.fn(),
publishDomainChanged: vi.fn(),
});
const deps = {
db,
logger,
publishDomainChanged: vi.fn(),
emitEvent: vi.fn(),
notify: vi.fn(),
createRunToken: () => "token",
getApiPort: () => 3210,
dispatchAutomation: vi.fn(),
hasWorkerReady: () => true,
enqueueWorkerRun,
isQuitting: () => false,
listRuntimeSkillEntries: () => [],
budgetService,
checkpointEngine: { shouldPauseHeartbeat: () => false },
};
const runDispatch = createRunDispatch(deps);
const taskRunner = createTaskRunner(
deps,
runDispatch.resolveWorkspace,
runDispatch.getRunEligibilityFailure,
runDispatch.dispatchRunToWorker,
(filter) => db.core.runs.hasActiveRun(filter),
);
return { enqueueWorkerRun, taskRunner };
}
function createCompanyGraph(db: AppDatabase, profileDir: string) {
db.core.connectors.updateConnectorHealth({
id: "openai-codex",
status: "ready",
authState: "authenticated",
version: "test",
lastError: null,
});
const companyId = db.core.company.saveCompany({
name: "Queued Recovery Co",
description: "Recovers queued runs deterministically",
status: "active",
});
const workspaceId = db.core.workspaces.saveWorkspace({
companyId,
name: "Primary",
localPath: profileDir,
repoUrl: "",
repoRef: "",
isPrimary: true,
});
const agentId = db.core.agents.saveAgent({
companyId,
name: "Engineer",
role: "Engineer",
title: "Engineer",
status: "idle",
reportsTo: null,
connectorId: "openai-codex",
workspaceId,
model: "gpt-5.4",
capabilities: "Runs queued recovery tasks",
budgetMonthlyUsd: 100,
});
return { companyId, workspaceId, agentId };
}
function createQueuedTaskRun(db: AppDatabase, input: {
companyId: string;
workspaceId: string;
agentId: string;
title: string;
}) {
const taskId = db.core.tasks.saveTask({
companyId: input.companyId,
title: input.title,
description: "Queued before worker startup recovery.",
assigneeAgentId: input.agentId,
workspaceId: input.workspaceId,
priority: "high",
status: "todo",
});
const run = db.core.runs.createTaskRun({
companyId: input.companyId,
taskId,
agentId: input.agentId,
workspaceId: input.workspaceId,
connectorId: "openai-codex",
});
if (!run) throw new Error("Expected queued run to be created.");
return { taskId, runId: run.id };
}
describe("queued run recovery", () => {
it("dispatches the first queued run and releases same-agent queued siblings", async () => {
const profileDir = await createTempDir("agentcompany-queued-recovery-");
const db = new AppDatabase(profileDir);
db.init();
const { companyId, workspaceId, agentId } = createCompanyGraph(db, profileDir);
const first = createQueuedTaskRun(db, { companyId, workspaceId, agentId, title: "First queued task" });
const second = createQueuedTaskRun(db, { companyId, workspaceId, agentId, title: "Second queued task" });
const { enqueueWorkerRun, taskRunner } = createHarness(db);
taskRunner.resumeQueuedRuns();
expect(enqueueWorkerRun).toHaveBeenCalledOnce();
expect(enqueueWorkerRun.mock.calls[0]?.[0]).toEqual(expect.objectContaining({
runId: first.runId,
taskId: first.taskId,
agentId,
trigger: "recovery",
}));
expect(db.core.runs.getRun(second.runId).status).toBe("cancelled");
expect(db.core.entity.getTask(second.taskId).status).toBe("todo");
db.close();
});
it("releases queued runs when recovery is blocked by connector readiness", async () => {
const profileDir = await createTempDir("agentcompany-queued-connector-block-");
const db = new AppDatabase(profileDir);
db.init();
const { companyId, workspaceId, agentId } = createCompanyGraph(db, profileDir);
const queued = createQueuedTaskRun(db, { companyId, workspaceId, agentId, title: "Connector blocked task" });
db.core.connectors.updateConnectorHealth({
id: "openai-codex",
status: "auth_required",
authState: "missing",
version: "test",
lastError: "Login required",
});
const { enqueueWorkerRun, taskRunner } = createHarness(db);
taskRunner.resumeQueuedRuns();
expect(enqueueWorkerRun).not.toHaveBeenCalled();
expect(db.core.runs.getRun(queued.runId).status).toBe("cancelled");
expect(db.core.entity.getTask(queued.taskId).status).toBe("todo");
db.close();
});
});