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

refactor(db): group database domain modules

Browse Source
This commit is contained in:
ZacharyZhang-NY committed 2026-04-22 19:24:25 -04:00
1 parent e43a89096d
commit 6da0a58769
160 files changed
+2124 -1999
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/main/worker/entry/crash-recovery.ts
+14 -14
View File
@@ -22,7 +22,7 @@ export function handleWorkerExit(
logger.warn(`[worker] Worker exited with code ${code}, recovering...`);
const snapshot = db.snap.listSnapshot();
const snapshot = db.projection.snapshot.listSnapshot();
const affectedAgentIds = new Set<string>();
const affectedCompanyIds = new Map<string, string>();
const requeueCandidates: Array<{
@@ -36,7 +36,7 @@ export function handleWorkerExit(
for (const run of snapshot.runs) {
if (run.status === "running") {
db.runs.finishRun({
db.core.runs.finishRun({
runId: run.id,
status: "interrupted",
summary: "Worker process crashed.",
@@ -47,16 +47,16 @@ export function handleWorkerExit(
unattributedCost: true,
});
if (run.workspaceId) {
db.workspacesMod.releaseWorkspaceLock(run.workspaceId, run.id);
db.core.workspaces.releaseWorkspaceLock(run.workspaceId, run.id);
}
// EXEC-02: Release execution lock on crash recovery
if (run.taskId) {
db.tasks.releaseTaskLock(run.taskId, run.id);
db.core.tasks.releaseTaskLock(run.taskId, run.id);
}
affectedAgentIds.add(run.agentId);
affectedCompanyIds.set(run.agentId, run.companyId);
const chain = db.agents.getChainOfCommand(run.agentId);
const chain = db.core.agents.getChainOfCommand(run.agentId);
if (chain.length > 0) {
wakeAgentIfPossible(chain[0].id, run.companyId, "report_failed");
}
@@ -75,11 +75,11 @@ export function handleWorkerExit(
for (const agentId of affectedAgentIds) {
// EXEC-05: Promote deferred wakes after crash recovery
const cid = affectedCompanyIds.get(agentId) ?? "";
const promoted = db.wakeup.promoteWakeupRequests(agentId, cid);
const promoted = db.automation.wakeup.promoteWakeupRequests(agentId, cid);
if (promoted > 0) {
logger.info(`[worker] Promoted ${promoted} deferred wake(s) for agent ${agentId} after crash recovery`);
}
const pendingWakeTrigger = db.wakeup.getPreferredWakeTrigger(agentId, cid);
const pendingWakeTrigger = db.automation.wakeup.getPreferredWakeTrigger(agentId, cid);
if (pendingWakeTrigger) {
wakeAgentIfPossible(agentId, cid, pendingWakeTrigger);
}
@@ -92,18 +92,18 @@ export function handleWorkerExit(
const worker = getWorker();
if (requeueCandidates.length > 0 && worker) {
logger.info(`[worker] Re-enqueuing ${requeueCandidates.length} queued run(s) after worker restart.`);
const recoverySnapshot = db.snap.listSnapshot();
const recoverySnapshot = db.projection.snapshot.listSnapshot();
for (const candidate of requeueCandidates) {
try {
// The crashed run may still own a workspace lock acquired during the
// original dispatch; release it before re-dispatch so the replay
// does not self-collide with "Workspace is locked by another run".
if (candidate.workspaceId) {
db.workspacesMod.releaseWorkspaceLock(candidate.workspaceId, candidate.id);
db.core.workspaces.releaseWorkspaceLock(candidate.workspaceId, candidate.id);
}
const task = candidate.taskId ? db.entityCrud.getTask(candidate.taskId) : null;
const agent = db.entityCrud.getAgent(candidate.agentId);
const connector = db.connectorsCrud.getConnector(agent.connectorId);
const task = candidate.taskId ? db.core.entity.getTask(candidate.taskId) : null;
const agent = db.core.entity.getAgent(candidate.agentId);
const connector = db.core.connectors.getConnector(agent.connectorId);
const { workspace, lockMode } = resolveWorkspaceWithLockMode(db, task, agent, candidate.companyId);
const recoveryCompany = recoverySnapshot.companies.find((company) => company.id === candidate.companyId);
@@ -122,7 +122,7 @@ export function handleWorkerExit(
});
} catch (error) {
logger.warn(`[worker] Failed to re-enqueue run ${candidate.id}: ${error instanceof Error ? error.message : String(error)}`);
db.runs.finishRun({
db.core.runs.finishRun({
runId: candidate.id,
status: "interrupted",
summary: "Could not re-enqueue after worker restart.",
@@ -133,7 +133,7 @@ export function handleWorkerExit(
unattributedCost: true,
});
if (candidate.workspaceId) {
db.workspacesMod.releaseWorkspaceLock(candidate.workspaceId, candidate.id);
db.core.workspaces.releaseWorkspaceLock(candidate.workspaceId, candidate.id);
}
}
}
src/main/worker/entry/types.ts
+5 -5
View File
@@ -111,14 +111,14 @@ export interface WorkerServiceDependencies {
dispatchRunToWorker: (params: {
runId: string;
companyId: string;
agent: ReturnType<AppDatabase["entityCrud"]["getAgent"]>;
task: ReturnType<AppDatabase["entityCrud"]["getTask"]> | null;
connector: ReturnType<AppDatabase["connectorsCrud"]["getConnector"]>;
workspace: ReturnType<AppDatabase["workspacesMod"]["getWorkspace"]>;
agent: ReturnType<AppDatabase["core"]["entity"]["getAgent"]>;
task: ReturnType<AppDatabase["core"]["entity"]["getTask"]> | null;
connector: ReturnType<AppDatabase["core"]["connectors"]["getConnector"]>;
workspace: ReturnType<AppDatabase["core"]["workspaces"]["getWorkspace"]>;
workspaceMode: "exclusive" | "shared" | "none";
companyName: string;
companyDescription: string;
snapshot: ReturnType<AppDatabase["snap"]["listSnapshot"]>;
snapshot: ReturnType<AppDatabase["projection"]["snapshot"]["listSnapshot"]>;
trigger: string;
}) => void;
processRunCompletion: (run: RunRecord, event: z.infer<typeof workerMessageSchema> & { type: "run-finished" }) => void;
src/main/worker/handlers/context-compaction-handler.ts
+1 -1
View File
@@ -6,7 +6,7 @@ export function handleContextCompaction(
deps: HandlerDeps,
_context: HandlerContext,
): void {
const runState = deps.db.runs.getRunAuthState(event.runId);
const runState = deps.db.core.runs.getRunAuthState(event.runId);
if (!runState) {
deps.logger.warn(`[context-compaction] run ${event.runId} disappeared before telemetry could be recorded.`);
return;
src/main/worker/handlers/run-finished-helpers.ts
+23 -23
View File
@@ -7,13 +7,13 @@ import { isUsageStatusUnattributed } from "../../run-usage-status";
import { decideRunAttempt, runAttemptKey } from "../../run-attempt-policy";
type RunFinishedEvent = Extract<WorkerMessageEvent, { type: "run-finished" }>;
type RunRecord = ReturnType<HandlerDeps["db"]["runs"]["getRun"]>;
type TaskRecord = ReturnType<HandlerDeps["db"]["entityCrud"]["getTask"]>;
type RunRecord = ReturnType<HandlerDeps["db"]["core"]["runs"]["getRun"]>;
type TaskRecord = ReturnType<HandlerDeps["db"]["core"]["entity"]["getTask"]>;
export interface RunFinishedContext {
run: RunRecord;
task: TaskRecord | null;
runOutcome: ReturnType<HandlerDeps["db"]["runs"]["finishRun"]>;
runOutcome: ReturnType<HandlerDeps["db"]["core"]["runs"]["finishRun"]>;
}
export function finalizeRun(
@@ -26,7 +26,7 @@ export function finalizeRun(
let run: RunRecord;
try {
run = db.runs.getRun(event.runId);
run = db.core.runs.getRun(event.runId);
} catch (error) {
logger.warn(
`[finalizeRun] getRun(${event.runId}) failed; run may have been deleted: ${error instanceof Error ? error.message : String(error)}`,
@@ -60,7 +60,7 @@ export function finalizeRun(
let task: TaskRecord | null = null;
if (run.taskId) {
try {
task = db.entityCrud.getTask(run.taskId);
task = db.core.entity.getTask(run.taskId);
} catch (error) {
logger.warn(
`[finalizeRun] getTask(${run.taskId}) failed; task may have been deleted: ${error instanceof Error ? error.message : String(error)}`,
@@ -70,11 +70,11 @@ export function finalizeRun(
}
leaseRenewalMap.delete(event.runId);
const runOutcome = db.runs.finishRun({ ...event, unattributedCost });
const runOutcome = db.core.runs.finishRun({ ...event, unattributedCost });
if (run.workspaceId) {
db.workspacesMod.releaseWorkspaceLock(run.workspaceId, run.id);
db.core.workspaces.releaseWorkspaceLock(run.workspaceId, run.id);
}
db.activityMod.addActivity({
db.core.activity.addActivity({
companyId: run.companyId,
actor: run.agentId,
action: `run.${event.status}`,
@@ -101,9 +101,9 @@ export function accumulateAndRotateSession(
// heartbeat run — the exact failure mode that lets a taskless transcript
// grow past the model's context window. Resolve the session through the
// agent, never through the task.
let agentRecord: ReturnType<HandlerDeps["db"]["entityCrud"]["getAgent"]>;
let agentRecord: ReturnType<HandlerDeps["db"]["core"]["entity"]["getAgent"]>;
try {
agentRecord = db.entityCrud.getAgent(run.agentId);
agentRecord = db.core.entity.getAgent(run.agentId);
} catch (error) {
logger.warn(
`[session] accumulateAndRotateSession: getAgent(${run.agentId}) failed; skipping: ${error instanceof Error ? error.message : String(error)}`,
@@ -111,7 +111,7 @@ export function accumulateAndRotateSession(
return;
}
if (!agentRecord.sessionId) return;
const session = db.tasks.getSessionById(agentRecord.sessionId);
const session = db.core.tasks.getSessionById(agentRecord.sessionId);
if (!session) return;
sessionService.accumulateTokenUsage(
@@ -124,7 +124,7 @@ export function accumulateAndRotateSession(
// that include this run, otherwise a run that just crossed the
// threshold waits a full additional cycle to rotate — the exact
// failure mode that let transcripts balloon past the context window.
const updatedSession = db.tasks.getSessionById(agentRecord.sessionId);
const updatedSession = db.core.tasks.getSessionById(agentRecord.sessionId);
if (!updatedSession) return;
const rotationResult = sessionService.evaluateRotation(updatedSession);
if (rotationResult.shouldRotate) {
@@ -158,9 +158,9 @@ export function emitAndNotifyCompletion(
});
}
let updatedAgent: ReturnType<HandlerDeps["db"]["entityCrud"]["getAgent"]> | null = null;
let updatedAgent: ReturnType<HandlerDeps["db"]["core"]["entity"]["getAgent"]> | null = null;
try {
updatedAgent = db.entityCrud.getAgent(run.agentId);
updatedAgent = db.core.entity.getAgent(run.agentId);
} catch (error) {
logger.warn(
`[emitAndNotifyCompletion] getAgent(${run.agentId}) failed; skipping budget-paused notification: ${error instanceof Error ? error.message : String(error)}`,
@@ -169,7 +169,7 @@ export function emitAndNotifyCompletion(
}
if (updatedAgent && runOutcome.budgetHardStopped) {
db.activityMod.addActivity({
db.core.activity.addActivity({
companyId: updatedAgent.companyId,
actor: "system",
action: "agent.auto_paused",
@@ -233,8 +233,8 @@ export function runPostCompletionAutomation(
// EXEC-05: Promote deferred wakes after run completion
let hadDeferredWakes = false;
if (!isQuitting() && !destroyed) {
const promoted = db.wakeup.promoteWakeupRequests(run.agentId, run.companyId);
const pendingWakeTrigger = db.wakeup.getPreferredWakeTrigger(run.agentId, run.companyId);
const promoted = db.automation.wakeup.promoteWakeupRequests(run.agentId, run.companyId);
const pendingWakeTrigger = db.automation.wakeup.getPreferredWakeTrigger(run.agentId, run.companyId);
if (promoted > 0) {
logger.info(`[worker] Promoted ${promoted} deferred wake(s) for agent ${run.agentId}`);
}
@@ -297,7 +297,7 @@ export function triggerSelfWake(
if (hadDeferredWakes || isQuitting() || destroyed) return;
try {
const actionableTasks = db.tasks.listAgentTasks(run.agentId, run.companyId, ["todo", "in_progress"]);
const actionableTasks = db.core.tasks.listAgentTasks(run.agentId, run.companyId, ["todo", "in_progress"]);
const hasWorkToContinue = actionableTasks.length > 0 || hasUnfinishedTeamBuilding(db, run.agentId, run.companyId);
const consecutiveSelfWakes = selfWakeAttemptMap.get(attemptKey) ?? 0;
const decision = decideRunAttempt({
@@ -316,7 +316,7 @@ export function triggerSelfWake(
);
// Escalate to manager since agent is stuck
try {
const chain = db.agents.getChainOfCommand(run.agentId);
const chain = db.core.agents.getChainOfCommand(run.agentId);
if (chain.length > 0) {
wakeAgentIfPossible(chain[0].id, run.companyId, "report_stuck");
}
@@ -354,15 +354,15 @@ function hasUnfinishedTeamBuilding(
companyId: string,
): boolean {
try {
const agent = db.entityCrud.getAgent(agentId);
const agent = db.core.entity.getAgent(agentId);
if (agent.reportsTo !== null) return false;
const activeReports = db.agents
const activeReports = db.core.agents
.getDirectReports(agentId, companyId)
.filter((r) => r.status !== "pending_approval" && r.status !== "terminated");
if (activeReports.length > 0) return false;
const goalCount = db.entityCrud.countActiveGoalsForCompany(companyId);
const goalCount = db.core.entity.countActiveGoalsForCompany(companyId);
if (goalCount === 0) return false;
const pendingHires = db.approvalsMod
const pendingHires = db.core.approvals
.listCompanyApprovalsByType(companyId, "hire_agent")
.filter((a) => a.state === "pending");
return pendingHires.length === 0;
src/main/worker/handlers/run-log-handler.ts
+4 -4
View File
@@ -16,8 +16,8 @@ export async function handleRunLog(
const filtered = filterRunLogChunk(event.stream, event.chunk);
if (filtered === null) return;
await db.runs.appendRunLog(event.runId, event.stream, filtered);
const runState = db.runs.getRunAuthState(event.runId);
await db.core.runs.appendRunLog(event.runId, event.stream, filtered);
const runState = db.core.runs.getRunAuthState(event.runId);
if (runState) {
persistRunEvent(db, logger, {
companyId: runState.companyId,
@@ -36,9 +36,9 @@ export async function handleRunLog(
const lastRenewal = leaseRenewalMap.get(event.runId) ?? 0;
if (now - lastRenewal >= LEASE_RENEWAL_INTERVAL_MS) {
try {
const run = db.runs.getRun(event.runId);
const run = db.core.runs.getRun(event.runId);
if (run.taskId) {
db.tasks.renewTaskLease(run.taskId, run.id);
db.core.tasks.renewTaskLease(run.taskId, run.id);
}
leaseRenewalMap.set(event.runId, now);
} catch (err) {
src/main/worker/handlers/run-status-handler.ts
+4 -4
View File
@@ -9,8 +9,8 @@ export function handleRunStatus(
): void {
const { db, logger, emitEvent, onRunStarted } = deps;
const runState = db.runs.getRunAuthState(event.runId);
db.runs.markRunStatus(event.runId, event.status, event.message);
const runState = db.core.runs.getRunAuthState(event.runId);
db.core.runs.markRunStatus(event.runId, event.status, event.message);
if (runState) {
persistRunEvent(db, logger, {
companyId: runState.companyId,
@@ -30,9 +30,9 @@ export function handleRunStatus(
// dispatchRunToWorker before enqueue. Only lease renewal and
// run-started housekeeping run here.
try {
const run = db.runs.getRun(event.runId);
const run = db.core.runs.getRun(event.runId);
if (run.taskId) {
db.tasks.renewTaskLease(run.taskId, run.id);
db.core.tasks.renewTaskLease(run.taskId, run.id);
}
onRunStarted?.(run.id, run.taskId, run.agentId, run.companyId);
} catch (err) {
src/main/worker/observability/failure-escalation.ts
+18 -16
View File
@@ -13,15 +13,17 @@ import { decideRunAttempt } from "../../run-attempt-policy";
export interface EscalationDeps {
db: {
tasks: {
getTaskRetryCount: (taskId: string) => number;
incrementTaskRetryCount: (taskId: string) => void;
resetTaskRetryCount: (taskId: string) => void;
advanceTaskStatus: (taskId: string, status: string, note: string) => unknown;
};
agents: {
getChainOfCommand: (agentId: string) => Array<{ id: string; name: string }>;
getCeoForCompany: (companyId: string) => { id: string; name: string } | null;
core: {
tasks: {
getTaskRetryCount: (taskId: string) => number;
incrementTaskRetryCount: (taskId: string) => void;
resetTaskRetryCount: (taskId: string) => void;
advanceTaskStatus: (taskId: string, status: string, note: string) => unknown;
};
agents: {
getChainOfCommand: (agentId: string) => Array<{ id: string; name: string }>;
getCeoForCompany: (companyId: string) => { id: string; name: string } | null;
};
};
};
wakeAgentIfPossible: (agentId: string, companyId: string, trigger: string) => unknown;
@@ -48,7 +50,7 @@ export function handleFailureEscalation(
scheduleRetry: (fn: () => void, delayMs: number) => void,
): EscalationResult {
const { db, wakeAgentIfPossible, notify, logger } = deps;
const retryCount = db.tasks.getTaskRetryCount(run.taskId);
const retryCount = db.core.tasks.getTaskRetryCount(run.taskId);
const retryDecision = decideRunAttempt({
kind: "task_failure",
retryCount,
@@ -56,8 +58,8 @@ export function handleFailureEscalation(
});
if (retryDecision.action === "retry") {
db.tasks.incrementTaskRetryCount(run.taskId);
db.tasks.advanceTaskStatus(
db.core.tasks.incrementTaskRetryCount(run.taskId);
db.core.tasks.advanceTaskStatus(
run.taskId,
"todo",
`Run ${eventStatus} (attempt ${retryDecision.attemptNumber}/${maxRetryAttempts}). Auto-retrying after ${retryCooldownMs / 1000}s cooldown.`,
@@ -72,14 +74,14 @@ export function handleFailureEscalation(
}
// Retry limit reached — escalate up the chain
db.tasks.resetTaskRetryCount(run.taskId);
db.core.tasks.resetTaskRetryCount(run.taskId);
logger.warn(
`[escalation] Task ${run.taskId} failed ${maxRetryAttempts} times. Escalating to manager.`,
);
let escalated = false;
try {
const chain = db.agents.getChainOfCommand(run.agentId);
const chain = db.core.agents.getChainOfCommand(run.agentId);
if (chain.length > 0) {
wakeAgentIfPossible(chain[0].id, run.companyId, "report_failed");
escalated = true;
@@ -95,7 +97,7 @@ export function handleFailureEscalation(
if (!escalated) {
try {
const ceo = db.agents.getCeoForCompany(run.companyId);
const ceo = db.core.agents.getCeoForCompany(run.companyId);
if (ceo) {
wakeAgentIfPossible(ceo.id, run.companyId, "report_failed");
escalated = true;
@@ -131,7 +133,7 @@ export function handleSuccessReset(
logger?: EscalationDeps["logger"],
): void {
try {
db.tasks.resetTaskRetryCount(taskId);
db.core.tasks.resetTaskRetryCount(taskId);
} catch (error) {
logger?.warn(
`[escalation] resetTaskRetryCount failed for task ${taskId}: ${error instanceof Error ? error.message : String(error)}`,
src/main/worker/observability/run-event-recorder.ts
+1 -1
View File
@@ -22,7 +22,7 @@ export function persistRunEvent(
input: PersistInput,
) {
try {
db.runs.appendRunEvent(input);
db.core.runs.appendRunEvent(input);
} catch (error) {
logger.warn(
`[run-events] Failed to append ${input.eventType} for ${input.runId}: ${
src/main/worker/runtime/internal-tools/file-approval-tool.ts
+1 -1
View File
@@ -4,7 +4,7 @@
* Collapses the `bash + node -e + fetch POST /api/companies/:id/approvals`
* chain into a single tool-call the CEO (or any agent with request
* authority) can invoke directly. The handler in main reuses the same
* `db.approvalsMod.requestApproval` primitive and auto-approval policy
* `db.core.approvals.requestApproval` primitive and auto-approval policy
* the HTTP route uses, so both paths produce identical rows and the
* `approve_ceo_strategy` ordering gate fires identically.
*