feat(05-02): replace flat budget check with policy-based gate, wire incident checks into run completion and heartbeat
- Replace flat budgetMonthlyUsd check with evaluateBudgetGate in getRunEligibilityFailure - Add projectId parameter to getRunEligibilityFailure for project-scope policy evaluation - Update all 3 call sites to pass projectId where task context is available - Wire checkAndCreateIncidents into worker-service run-finished handler for post-run incident creation - Wire autoCloseExpiredIncidents into heartbeat for window rollover cleanup (D-06) - Pass budgetService dependency to orchestration-service and worker-service via service-graph - Mark autoPauseBudgetExceededAgent as deprecated (superseded by BudgetEnforcementService)
This commit is contained in:
1 parent
feea5c1fe9
commit
4e5429929a
3 files changed
+36
-9
No files matched your search
@@ -5,6 +5,7 @@ import { getConnectorExecutionReadinessIssue } from "@shared/connector-policy";
|
||||
import type { DesktopResult } from "@shared/contracts";
|
||||
import type { ConnectorRecord } from "@shared/types";
|
||||
import type { NotificationOptions } from "./notification-service";
|
||||
import type { createBudgetEnforcementService } from "./budget-enforcement-service";
|
||||
|
||||
type WorkspaceLockMode = "exclusive" | "shared" | "none";
|
||||
|
||||
@@ -64,6 +65,7 @@ interface OrchestrationDependencies {
|
||||
enqueueWorkerRun: (payload: EnqueueRunPayload) => void;
|
||||
isQuitting: () => boolean;
|
||||
listRuntimeSkillEntries: (companyId: string, agentId: string, connectorId: string) => Array<{ runtimeName: string; sourcePath: string }>;
|
||||
budgetService: ReturnType<typeof createBudgetEnforcementService>;
|
||||
}
|
||||
|
||||
export function createOrchestrationService({
|
||||
@@ -78,6 +80,7 @@ export function createOrchestrationService({
|
||||
enqueueWorkerRun,
|
||||
isQuitting,
|
||||
listRuntimeSkillEntries,
|
||||
budgetService,
|
||||
}: OrchestrationDependencies) {
|
||||
let heartbeatTimer: ReturnType<typeof setInterval> | null = null;
|
||||
const scheduleLastFiredMap = new Map<string, number>();
|
||||
@@ -122,6 +125,7 @@ export function createOrchestrationService({
|
||||
return { workspace: null, lockMode: "none", source: "none" };
|
||||
}
|
||||
|
||||
// DEPRECATED: Superseded by BudgetEnforcementService.evaluateBudgetGate (Phase 5)
|
||||
function autoPauseBudgetExceededAgent(agent: { id: string; companyId: string; name: string; spentMonthlyUsd: number; budgetMonthlyUsd: number }) {
|
||||
if (!(agent.budgetMonthlyUsd > 0) || agent.spentMonthlyUsd < agent.budgetMonthlyUsd) {
|
||||
return false;
|
||||
@@ -156,14 +160,16 @@ export function createOrchestrationService({
|
||||
function getRunEligibilityFailure(
|
||||
agent: { status: string; budgetMonthlyUsd: number; spentMonthlyUsd: number; name: string; connectorId: string; id: string; companyId: string },
|
||||
connector: Pick<ConnectorRecord, "id" | "label" | "status" | "capabilityMatrix">,
|
||||
projectId: string | null = null,
|
||||
) {
|
||||
if (["terminated", "paused", "pending_approval"].includes(agent.status)) {
|
||||
return "AGENT_NOT_RUNNABLE";
|
||||
}
|
||||
if (agent.budgetMonthlyUsd > 0 && agent.spentMonthlyUsd >= agent.budgetMonthlyUsd) {
|
||||
autoPauseBudgetExceededAgent(agent);
|
||||
logger.info(`[run] Agent ${agent.name} budget exhausted, skipping`);
|
||||
return "AGENT_BUDGET_EXHAUSTED";
|
||||
// Policy-based budget check replaces flat field check (D-03, D-04, D-05)
|
||||
const gate = budgetService.evaluateBudgetGate(agent.companyId, agent.id, projectId);
|
||||
if (!gate.allowed) {
|
||||
logger.info(`[run] Agent ${agent.name} budget policy blocked: ${gate.blockedBy.map(b => `${b.scopeType}(${b.policyId})`).join(", ")}`);
|
||||
return "BUDGET_POLICY_BLOCKED";
|
||||
}
|
||||
const connectorIssue = getConnectorExecutionReadinessIssue(connector);
|
||||
if (connectorIssue) {
|
||||
@@ -342,7 +348,7 @@ export function createOrchestrationService({
|
||||
return null;
|
||||
}
|
||||
const connector = db.getConnector(agent.connectorId);
|
||||
if (getRunEligibilityFailure(agent, connector)) {
|
||||
if (getRunEligibilityFailure(agent, connector, null)) {
|
||||
return null;
|
||||
}
|
||||
if (hasActiveRun({ agentId: agent.id, companyId })) {
|
||||
@@ -530,12 +536,12 @@ export function createOrchestrationService({
|
||||
return fail("AGENT_COMPANY_MISMATCH", "The assigned agent does not belong to the task company.");
|
||||
}
|
||||
const connector = db.getConnector(agent.connectorId);
|
||||
const eligibilityFailure = getRunEligibilityFailure(agent, connector);
|
||||
const eligibilityFailure = getRunEligibilityFailure(agent, connector, task.projectId ?? null);
|
||||
if (eligibilityFailure === "AGENT_NOT_RUNNABLE") {
|
||||
return fail("AGENT_NOT_RUNNABLE", "The assigned agent is paused, pending approval, or terminated.");
|
||||
}
|
||||
if (eligibilityFailure === "AGENT_BUDGET_EXHAUSTED") {
|
||||
return fail("AGENT_BUDGET_EXHAUSTED", "The assigned agent has exhausted its monthly budget.");
|
||||
if (eligibilityFailure === "BUDGET_POLICY_BLOCKED") {
|
||||
return fail("BUDGET_POLICY_BLOCKED", "A budget policy blocks this run. Review budget settings.");
|
||||
}
|
||||
if (eligibilityFailure === "CONNECTOR_NOT_READY") {
|
||||
return fail("CONNECTOR_NOT_READY", "The assigned agent connector is not ready.");
|
||||
@@ -650,7 +656,7 @@ export function createOrchestrationService({
|
||||
const task = db.getTask(run.taskId);
|
||||
const agent = db.getAgent(run.agentId);
|
||||
const connector = db.getConnector(run.connectorId);
|
||||
if (getRunEligibilityFailure(agent, connector)) {
|
||||
if (getRunEligibilityFailure(agent, connector, task.projectId ?? null)) {
|
||||
continue;
|
||||
}
|
||||
const workspaceResolution = resolveWorkspace(task, agent, run.companyId);
|
||||
@@ -757,6 +763,12 @@ export function createOrchestrationService({
|
||||
db.setLastBudgetResetMonth(company.id, currentMonth);
|
||||
logger.info(`[heartbeat] Reset monthly budgets for company ${company.name} (${currentMonth})`);
|
||||
}
|
||||
|
||||
// D-06: Auto-close incidents from previous budget windows
|
||||
const closedIncidents = budgetService.autoCloseExpiredIncidents(company.id);
|
||||
if (closedIncidents > 0) {
|
||||
logger.info(`[heartbeat] Auto-closed ${closedIncidents} expired budget incident(s) for company ${company.name}`);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`[heartbeat] Error processing company ${company.id}: ${error instanceof Error ? error.message : String(error)}`);
|
||||
}
|
||||
|
||||
@@ -191,6 +191,7 @@ export function createServiceGraph({
|
||||
dispatchRunToWorker,
|
||||
processRunCompletion: (run, event) => processRunCompletion(run, event),
|
||||
notify,
|
||||
budgetService: getOrCreateBudgetEnforcementService(),
|
||||
});
|
||||
|
||||
orchestrationService = createOrchestrationService({
|
||||
@@ -219,6 +220,7 @@ export function createServiceGraph({
|
||||
}
|
||||
};
|
||||
})(),
|
||||
budgetService: getOrCreateBudgetEnforcementService(),
|
||||
});
|
||||
workerService.createWorker();
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ import { getConnectorDefinition } from "./connectors";
|
||||
import type { AppDatabase } from "./database";
|
||||
import type { RunRecord } from "@shared/types";
|
||||
import type { NotificationOptions } from "./notification-service";
|
||||
import type { createBudgetEnforcementService } from "./budget-enforcement-service";
|
||||
|
||||
const runStatusSchema = z.enum(["queued", "running", "succeeded", "failed", "cancelled", "timed_out", "interrupted"]);
|
||||
export const workerMessageSchema = z.union([
|
||||
@@ -63,6 +64,7 @@ interface WorkerServiceDependencies {
|
||||
}) => void;
|
||||
processRunCompletion: (run: RunRecord, event: z.infer<typeof workerMessageSchema> & { type: "run-finished" }) => void;
|
||||
notify: (options: NotificationOptions) => void;
|
||||
budgetService?: ReturnType<typeof createBudgetEnforcementService>;
|
||||
}
|
||||
|
||||
export function createWorkerService({
|
||||
@@ -76,6 +78,7 @@ export function createWorkerService({
|
||||
dispatchRunToWorker,
|
||||
processRunCompletion,
|
||||
notify,
|
||||
budgetService,
|
||||
}: WorkerServiceDependencies) {
|
||||
const MAX_RETRY_ATTEMPTS = 2;
|
||||
const RETRY_COOLDOWN_MS = 30_000; // 30 seconds
|
||||
@@ -330,6 +333,16 @@ export function createWorkerService({
|
||||
publishDomainChanged();
|
||||
}
|
||||
|
||||
// Policy-based incident check after run completion (Phase 5)
|
||||
if (budgetService && typeof event.costUsd === "number" && Number.isFinite(event.costUsd)) {
|
||||
try {
|
||||
const task = db.getTask(run.taskId);
|
||||
budgetService.checkAndCreateIncidents(run.companyId, run.agentId, task?.projectId ?? null);
|
||||
} catch (err) {
|
||||
logger.error(`[budget] Incident check failed: ${err instanceof Error ? err.message : String(err)}`);
|
||||
}
|
||||
}
|
||||
|
||||
if ((event.status === "failed" || event.status === "timed_out") && !isQuitting() && !destroyed) {
|
||||
const retryCount = db.getTaskRetryCount(run.taskId);
|
||||
|
||||
|
||||
Reference in new issue
Block a user