diff --git a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts index eecf0bf2..9c0f2732 100644 --- a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts +++ b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts @@ -5,6 +5,7 @@ import { agents, agentWakeupRequests, companies, + costEvents, createDb, documentRevisions, documents, @@ -96,6 +97,7 @@ async function cleanupHeartbeatInvalidationFixture(db: ReturnType; }; type SeedResult = { @@ -201,6 +204,7 @@ describeEmbeddedPostgres("heartbeat stale queued-run invalidation", () => { heartbeat: { wakeOnDemand: true, maxConcurrentRuns: opts.maxConcurrentRuns ?? 1, + ...(opts.heartbeatConfig ?? {}), }, }, permissions: {}, @@ -288,6 +292,701 @@ describeEmbeddedPostgres("heartbeat stale queued-run invalidation", () => { }); } + it("skips generic timer wakes with no actionable assigned work before adapter execution", async () => { + const { agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + enabled: true, + skipTimerWhenNoActionableWork: true, + }, + }); + + const run = await heartbeat.wakeup(agentId, { + source: "timer", + triggerDetail: "schedule", + }); + + expect(run).toBeNull(); + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const [wakeup] = await db + .select({ + status: agentWakeupRequests.status, + reason: agentWakeupRequests.reason, + payload: agentWakeupRequests.payload, + }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.agentId, agentId)); + const runRows = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns); + + expect(wakeup).toMatchObject({ + status: "skipped", + reason: "heartbeat.timer.no_actionable_work", + }); + expect(wakeup?.payload).toMatchObject({ + heartbeatSkip: { + reason: expect.stringContaining("No assigned todo or in_progress issue"), + }, + }); + expect(runRows).toHaveLength(0); + }); + + it("rate-limits skipped generic timer wakes by advancing the timer baseline", async () => { + const { agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + enabled: true, + intervalSec: 60, + skipTimerWhenNoActionableWork: true, + }, + }); + const now = new Date(); + await db + .update(agents) + .set({ lastHeartbeatAt: new Date(now.getTime() - 120_000) }) + .where(eq(agents.id, agentId)); + + const firstTick = await heartbeat.tickTimers(now); + const secondTick = await heartbeat.tickTimers(now); + + expect(firstTick.skipped).toBe(1); + expect(secondTick.skipped).toBe(0); + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const wakeups = await db + .select({ reason: agentWakeupRequests.reason }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.agentId, agentId)); + const [agent] = await db + .select({ lastHeartbeatAt: agents.lastHeartbeatAt }) + .from(agents) + .where(eq(agents.id, agentId)); + + expect(wakeups).toHaveLength(1); + expect(wakeups[0]?.reason).toBe("heartbeat.timer.no_actionable_work"); + expect(agent?.lastHeartbeatAt).toBeInstanceOf(Date); + expect(agent?.lastHeartbeatAt?.getTime()).toBeGreaterThan(now.getTime() - 120_000); + }); + + it("allows generic timer wakes when the agent has assigned todo work", async () => { + const { companyId, agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + enabled: true, + skipTimerWhenNoActionableWork: true, + }, + }); + await db.insert(issues).values({ + id: randomUUID(), + companyId, + title: "Assigned work", + status: "todo", + priority: "high", + assigneeAgentId: agentId, + }); + + const run = await heartbeat.wakeup(agentId, { + source: "timer", + triggerDetail: "schedule", + }); + + expect(run).not.toBeNull(); + await waitForCondition(async () => countExecuteCallsForRun(run!.id) > 0); + + expect(countExecuteCallsForRun(run!.id)).toBe(1); + }); + + it("runs generic timer wakes by default for proactive agents without assigned issue work", async () => { + const { agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + enabled: true, + }, + }); + + const run = await heartbeat.wakeup(agentId, { + source: "timer", + triggerDetail: "schedule", + }); + + expect(run).not.toBeNull(); + await waitForCondition(async () => countExecuteCallsForRun(run!.id) > 0); + + expect(countExecuteCallsForRun(run!.id)).toBe(1); + }); + + it("skips wakes before queueing when per-agent daily run cap is reached", async () => { + const { companyId, agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + maxDailyRuns: 1, + }, + }); + await db.insert(heartbeatRuns).values({ + id: randomUUID(), + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "succeeded", + createdAt: new Date(), + startedAt: new Date(), + finishedAt: new Date(), + contextSnapshot: {}, + }); + + const run = await heartbeat.wakeup(agentId, { + source: "on_demand", + triggerDetail: "manual", + }); + + expect(run).toBeNull(); + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const [wakeup] = await db + .select({ + status: agentWakeupRequests.status, + reason: agentWakeupRequests.reason, + payload: agentWakeupRequests.payload, + }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.agentId, agentId)); + + expect(wakeup).toMatchObject({ + status: "skipped", + reason: "heartbeat.daily_run_limit", + }); + expect(wakeup?.payload).toMatchObject({ + heartbeatSkip: { + observed: 1, + limit: 1, + }, + }); + }); + + it("treats zero daily run cap as a hard stop", async () => { + const { agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + maxDailyRuns: 0, + }, + }); + + const run = await heartbeat.wakeup(agentId, { + source: "on_demand", + triggerDetail: "manual", + }); + + expect(run).toBeNull(); + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const [wakeup] = await db + .select({ + status: agentWakeupRequests.status, + reason: agentWakeupRequests.reason, + payload: agentWakeupRequests.payload, + }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.agentId, agentId)); + + expect(wakeup).toMatchObject({ + status: "skipped", + reason: "heartbeat.daily_run_limit", + }); + expect(wakeup?.payload).toMatchObject({ + heartbeatSkip: { + observed: 0, + limit: 0, + }, + }); + }); + + it("counts started cancelled runs toward the per-agent daily run cap", async () => { + const { companyId, agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + maxDailyRuns: 1, + }, + }); + await db.insert(heartbeatRuns).values({ + id: randomUUID(), + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "cancelled", + createdAt: new Date(), + startedAt: new Date(), + finishedAt: new Date(), + contextSnapshot: {}, + }); + + const run = await heartbeat.wakeup(agentId, { + source: "on_demand", + triggerDetail: "manual", + }); + + expect(run).toBeNull(); + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const [wakeup] = await db + .select({ + status: agentWakeupRequests.status, + reason: agentWakeupRequests.reason, + payload: agentWakeupRequests.payload, + }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.agentId, agentId)); + + expect(wakeup).toMatchObject({ + status: "skipped", + reason: "heartbeat.daily_run_limit", + }); + expect(wakeup?.payload).toMatchObject({ + heartbeatSkip: { + observed: 1, + limit: 1, + }, + }); + }); + + it("coalesces same-issue wakes before enforcing the daily run cap", async () => { + const { companyId, agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + maxDailyRuns: 1, + }, + }); + const issueId = randomUUID(); + const wakeupRequestId = randomUUID(); + const queuedRunId = randomUUID(); + await db.insert(heartbeatRuns).values({ + id: randomUUID(), + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "succeeded", + createdAt: new Date(), + startedAt: new Date(), + finishedAt: new Date(), + contextSnapshot: {}, + }); + await db.insert(agentWakeupRequests).values({ + id: wakeupRequestId, + companyId, + agentId, + source: "on_demand", + triggerDetail: "manual", + reason: "manual", + payload: { issueId }, + status: "queued", + }); + await db.insert(heartbeatRuns).values({ + id: queuedRunId, + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "queued", + wakeupRequestId, + contextSnapshot: { issueId }, + }); + await db.insert(issues).values({ + id: issueId, + companyId, + title: "Queued issue work", + status: "in_progress", + priority: "high", + assigneeAgentId: agentId, + executionRunId: queuedRunId, + }); + await db + .update(agentWakeupRequests) + .set({ runId: queuedRunId }) + .where(eq(agentWakeupRequests.id, wakeupRequestId)); + + const run = await heartbeat.wakeup(agentId, { + source: "on_demand", + triggerDetail: "manual", + payload: { issueId }, + }); + + expect(run?.id).toBe(queuedRunId); + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const wakeups = await db + .select({ + status: agentWakeupRequests.status, + reason: agentWakeupRequests.reason, + runId: agentWakeupRequests.runId, + }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.agentId, agentId)); + + expect(wakeups).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + status: "coalesced", + reason: "issue_execution_same_name", + runId: queuedRunId, + }), + ]), + ); + }); + + it("skips wakes before queueing when per-agent daily cost cap is reached", async () => { + const { companyId, agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + maxDailyCostCents: 75, + }, + }); + await db.insert(costEvents).values({ + companyId, + agentId, + provider: "test", + biller: "test", + billingType: "metered_api", + model: "test-model", + inputTokens: 100, + outputTokens: 50, + costCents: 75, + occurredAt: new Date(), + }); + + const run = await heartbeat.wakeup(agentId, { + source: "on_demand", + triggerDetail: "manual", + }); + + expect(run).toBeNull(); + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const [wakeup] = await db + .select({ + status: agentWakeupRequests.status, + reason: agentWakeupRequests.reason, + payload: agentWakeupRequests.payload, + }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.agentId, agentId)); + + expect(wakeup).toMatchObject({ + status: "skipped", + reason: "heartbeat.daily_cost_limit", + }); + expect(wakeup?.payload).toMatchObject({ + heartbeatSkip: { + observed: 75, + limit: 75, + }, + }); + }); + + it("treats zero daily cost cap as a hard stop", async () => { + const { agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + maxDailyCostCents: 0, + }, + }); + + const run = await heartbeat.wakeup(agentId, { + source: "on_demand", + triggerDetail: "manual", + }); + + expect(run).toBeNull(); + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const [wakeup] = await db + .select({ + status: agentWakeupRequests.status, + reason: agentWakeupRequests.reason, + payload: agentWakeupRequests.payload, + }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.agentId, agentId)); + + expect(wakeup).toMatchObject({ + status: "skipped", + reason: "heartbeat.daily_cost_limit", + }); + expect(wakeup?.payload).toMatchObject({ + heartbeatSkip: { + observed: 0, + limit: 0, + }, + }); + }); + + it("skips already queued runs before adapter execution when the daily cost cap is reached", async () => { + const { companyId, agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + maxDailyCostCents: 75, + }, + }); + const wakeupRequestId = randomUUID(); + const runId = randomUUID(); + await db.insert(agentWakeupRequests).values({ + id: wakeupRequestId, + companyId, + agentId, + source: "on_demand", + triggerDetail: "manual", + reason: "manual", + payload: {}, + status: "queued", + }); + await db.insert(heartbeatRuns).values({ + id: runId, + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "queued", + wakeupRequestId, + contextSnapshot: {}, + }); + await db + .update(agentWakeupRequests) + .set({ runId }) + .where(eq(agentWakeupRequests.id, wakeupRequestId)); + await db.insert(costEvents).values({ + companyId, + agentId, + provider: "test", + biller: "test", + billingType: "metered_api", + model: "test-model", + inputTokens: 100, + outputTokens: 50, + costCents: 75, + occurredAt: new Date(), + }); + + await heartbeat.resumeQueuedRuns(); + + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const [run] = await db + .select({ + status: heartbeatRuns.status, + errorCode: heartbeatRuns.errorCode, + resultJson: heartbeatRuns.resultJson, + }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)); + const [wakeup] = await db + .select({ + status: agentWakeupRequests.status, + error: agentWakeupRequests.error, + }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequestId)); + + expect(run).toMatchObject({ + status: "cancelled", + errorCode: "heartbeat.daily_cost_limit", + }); + expect(run?.resultJson).toMatchObject({ + stopReason: "heartbeat.daily_cost_limit", + observed: 75, + limit: 75, + }); + expect(wakeup).toMatchObject({ + status: "skipped", + error: expect.stringContaining("per-day heartbeat budget cap"), + }); + }); + + it("skips already queued issue runs at the daily run cap and releases the execution lock", async () => { + const { companyId, agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + maxDailyRuns: 1, + }, + }); + const issueId = randomUUID(); + const wakeupRequestId = randomUUID(); + const queuedRunId = randomUUID(); + await db.insert(heartbeatRuns).values({ + id: randomUUID(), + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "succeeded", + createdAt: new Date(), + startedAt: new Date(), + finishedAt: new Date(), + contextSnapshot: {}, + }); + await db.insert(agentWakeupRequests).values({ + id: wakeupRequestId, + companyId, + agentId, + source: "on_demand", + triggerDetail: "manual", + reason: "manual", + payload: { issueId }, + status: "queued", + }); + await db.insert(heartbeatRuns).values({ + id: queuedRunId, + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "queued", + wakeupRequestId, + contextSnapshot: { issueId }, + }); + await db.insert(issues).values({ + id: issueId, + companyId, + title: "Queued issue work", + status: "in_progress", + priority: "high", + assigneeAgentId: agentId, + executionRunId: queuedRunId, + }); + await db + .update(agentWakeupRequests) + .set({ runId: queuedRunId }) + .where(eq(agentWakeupRequests.id, wakeupRequestId)); + + await heartbeat.resumeQueuedRuns(); + + expect(mockAdapterExecute).not.toHaveBeenCalled(); + + const [run] = await db + .select({ + status: heartbeatRuns.status, + errorCode: heartbeatRuns.errorCode, + }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, queuedRunId)); + const [wakeup] = await db + .select({ status: agentWakeupRequests.status }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequestId)); + const [issue] = await db + .select({ executionRunId: issues.executionRunId }) + .from(issues) + .where(eq(issues.id, issueId)); + + expect(run).toMatchObject({ + status: "cancelled", + errorCode: "heartbeat.daily_run_limit", + }); + expect(wakeup).toMatchObject({ status: "skipped" }); + expect(issue?.executionRunId).toBeNull(); + }); + + it("promotes deferred issue wakes when a queued holder is cancelled by the daily run cap", async () => { + const { companyId, agentId } = await seedCompanyAndAgent({ + heartbeatConfig: { + maxDailyRuns: 1, + }, + }); + const peerAgentId = randomUUID(); + const issueId = randomUUID(); + const wakeupRequestId = randomUUID(); + const queuedRunId = randomUUID(); + const deferredWakeupId = randomUUID(); + await db.insert(agents).values({ + id: peerAgentId, + companyId, + name: "PeerAgent", + role: "engineer", + status: "active", + adapterType: "codex_local", + adapterConfig: {}, + runtimeConfig: { + heartbeat: { + wakeOnDemand: true, + maxConcurrentRuns: 1, + }, + }, + permissions: {}, + }); + await db.insert(heartbeatRuns).values({ + id: randomUUID(), + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "succeeded", + createdAt: new Date(), + startedAt: new Date(), + finishedAt: new Date(), + contextSnapshot: {}, + }); + await db.insert(agentWakeupRequests).values({ + id: wakeupRequestId, + companyId, + agentId, + source: "on_demand", + triggerDetail: "manual", + reason: "manual", + payload: { issueId }, + status: "queued", + }); + await db.insert(heartbeatRuns).values({ + id: queuedRunId, + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "queued", + wakeupRequestId, + contextSnapshot: { issueId }, + }); + await db.insert(issues).values({ + id: issueId, + companyId, + title: "Queued issue work", + status: "in_progress", + priority: "high", + assigneeAgentId: agentId, + executionRunId: queuedRunId, + }); + await db.insert(agentWakeupRequests).values({ + id: deferredWakeupId, + companyId, + agentId: peerAgentId, + source: "comment", + triggerDetail: "mention", + reason: "issue_execution_deferred", + payload: { + issueId, + _paperclipWakeContext: { + issueId, + wakeReason: "issue_mention", + }, + }, + status: "deferred_issue_execution", + }); + await db + .update(agentWakeupRequests) + .set({ runId: queuedRunId }) + .where(eq(agentWakeupRequests.id, wakeupRequestId)); + + await heartbeat.resumeQueuedRuns(); + await waitForCondition(async () => { + const [deferred] = await db + .select({ status: agentWakeupRequests.status, runId: agentWakeupRequests.runId }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, deferredWakeupId)); + return Boolean(deferred?.runId) && deferred?.status !== "deferred_issue_execution"; + }); + + const [deferred] = await db + .select({ status: agentWakeupRequests.status, runId: agentWakeupRequests.runId }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, deferredWakeupId)); + const [promotedRun] = deferred?.runId + ? await db + .select({ agentId: heartbeatRuns.agentId }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, deferred.runId)) + : []; + + expect(deferred?.status).not.toBe("deferred_issue_execution"); + expect(promotedRun?.agentId).toBe(peerAgentId); + }); + it("cancels queued runs when the issue assignee changes before the run starts", async () => { const { companyId, agentId } = await seedCompanyAndAgent({ agentName: "OriginalCoder" }); const replacementAgentId = randomUUID(); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 466328d0..72423127 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -3,7 +3,7 @@ import path from "node:path"; import { execFile as execFileCallback } from "node:child_process"; import { promisify } from "node:util"; import { randomUUID } from "node:crypto"; -import { and, asc, desc, eq, getTableColumns, gt, inArray, isNull, lt, lte, notInArray, or, sql } from "drizzle-orm"; +import { and, asc, desc, eq, getTableColumns, gt, gte, inArray, isNull, lt, lte, notInArray, or, sql } from "drizzle-orm"; import type { Db } from "@paperclipai/db"; import { AGENT_DEFAULT_MAX_CONCURRENT_RUNS, @@ -32,6 +32,7 @@ import { approvals, companySkills as companySkillsTable, companies, + costEvents, documentAnnotationComments, documentAnnotationThreads, documentRevisions, @@ -238,6 +239,7 @@ const EXECUTION_PATH_HEARTBEAT_RUN_STATUSES = ["queued", "running", "scheduled_r const CANCELLABLE_HEARTBEAT_RUN_STATUSES = ["queued", "running", "scheduled_retry"] as const; const HEARTBEAT_RUN_TERMINAL_STATUSES = ["succeeded", "failed", "cancelled", "timed_out"] as const; const UNSUCCESSFUL_HEARTBEAT_RUN_TERMINAL_STATUSES = ["failed", "cancelled", "timed_out"] as const; +const TIMER_ACTIONABLE_ISSUE_STATUSES = ["todo", "in_progress"] as const; export { ACTIVE_RUN_OUTPUT_CONTINUE_REARM_MS, ACTIVE_RUN_OUTPUT_CRITICAL_THRESHOLD_MS, @@ -6835,9 +6837,169 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) intervalSec: Math.max(0, asNumber(heartbeat.intervalSec, 0)), wakeOnDemand: asBoolean(heartbeat.wakeOnDemand ?? heartbeat.wakeOnAssignment ?? heartbeat.wakeOnOnDemand ?? heartbeat.wakeOnAutomation, true), maxConcurrentRuns: normalizeMaxConcurrentRuns(heartbeat.maxConcurrentRuns), + skipTimerWhenNoActionableWork: asBoolean( + heartbeat.skipTimerWhenNoActionableWork ?? + heartbeat.requireActionableTimerWork ?? + heartbeat.issueOnlyTimer, + false, + ), + maxDailyRuns: normalizeOptionalNonNegativeInteger( + heartbeat.maxDailyRuns ?? heartbeat.dailyRunLimit ?? heartbeat.dailyRunCap ?? heartbeat.maxRunsPerDay, + ), + maxDailyCostCents: normalizeOptionalNonNegativeInteger( + heartbeat.maxDailyCostCents ?? + heartbeat.dailyCostCentsLimit ?? + heartbeat.dailySpendCentsLimit ?? + heartbeat.dailyBudgetCents, + ), }; } + function normalizeOptionalNonNegativeInteger(value: unknown) { + if (value === null || value === undefined || value === "") return null; + const normalized = Math.floor(asNumber(value, 0)); + return normalized >= 0 ? normalized : null; + } + + function currentUtcDayWindow(now = new Date()) { + const start = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), now.getUTCDate(), 0, 0, 0, 0)); + const end = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), now.getUTCDate() + 1, 0, 0, 0, 0)); + return { start, end }; + } + + async function getHeartbeatDailyCapBlock( + agent: typeof agents.$inferSelect, + policy: ReturnType, + options: { checkRunCap?: boolean; checkCostCap?: boolean; excludeRunId?: string | null } = {}, + client: Pick = db, + ) { + const checkRunCap = options.checkRunCap ?? true; + const checkCostCap = options.checkCostCap ?? true; + const { start, end } = currentUtcDayWindow(); + if (checkRunCap && policy.maxDailyRuns !== null) { + const conditions = [ + eq(heartbeatRuns.companyId, agent.companyId), + eq(heartbeatRuns.agentId, agent.id), + gte(heartbeatRuns.startedAt, start), + lt(heartbeatRuns.startedAt, end), + notInArray(heartbeatRuns.status, ["queued", "scheduled_retry"]), + ]; + if (options.excludeRunId) { + conditions.push(sql`${heartbeatRuns.id} <> ${options.excludeRunId}`); + } + const [row] = await client + .select({ total: sql`count(*)::integer` }) + .from(heartbeatRuns) + .where(and(...conditions)); + const observed = Number(row?.total ?? 0); + if (observed >= policy.maxDailyRuns) { + return { + reason: "heartbeat.daily_run_limit", + observed, + limit: policy.maxDailyRuns, + }; + } + } + + if (checkCostCap && policy.maxDailyCostCents !== null) { + const [row] = await client + .select({ total: sql`coalesce(sum(${costEvents.costCents})::bigint, 0)` }) + .from(costEvents) + .where( + and( + eq(costEvents.companyId, agent.companyId), + eq(costEvents.agentId, agent.id), + gte(costEvents.occurredAt, start), + lt(costEvents.occurredAt, end), + ), + ); + const observed = Number(row?.total ?? 0); + if (observed >= policy.maxDailyCostCents) { + return { + reason: "heartbeat.daily_cost_limit", + observed, + limit: policy.maxDailyCostCents, + }; + } + } + + return null; + } + + async function cancelQueuedRunForHeartbeatDailyCap( + run: typeof heartbeatRuns.$inferSelect, + dailyCapBlock: NonNullable>>, + ) { + const now = new Date(); + const reason = "Cancelled because the agent reached a per-day heartbeat budget cap before adapter invocation"; + const cancelled = await setRunStatus(run.id, "cancelled", { + finishedAt: now, + error: reason, + errorCode: dailyCapBlock.reason, + resultJson: { + ...parseObject(run.resultJson), + stopReason: dailyCapBlock.reason, + observed: dailyCapBlock.observed, + limit: dailyCapBlock.limit, + effectiveTimeoutSec: 0, + timeoutConfigured: false, + timeoutSource: "heartbeat_daily_cap_gate", + timeoutFired: false, + }, + }); + if (!cancelled) return null; + + await setWakeupStatus(run.wakeupRequestId, "skipped", { + finishedAt: now, + error: reason, + }); + + await appendRunEvent(cancelled, await nextRunEventSeq(cancelled.id), { + eventType: "lifecycle", + stream: "system", + level: "warn", + message: reason, + payload: { + reason: dailyCapBlock.reason, + observed: dailyCapBlock.observed, + limit: dailyCapBlock.limit, + }, + }); + + await releaseIssueExecutionAndPromote(cancelled, { suppressImmediateRecovery: true }); + + return cancelled; + } + + async function hasActionableTimerWork(agent: typeof agents.$inferSelect) { + const row = await db + .select({ id: issues.id }) + .from(issues) + .where( + and( + eq(issues.companyId, agent.companyId), + eq(issues.assigneeAgentId, agent.id), + isNull(issues.assigneeUserId), + isNull(issues.hiddenAt), + inArray(issues.status, [...TIMER_ACTIONABLE_ISSUE_STATUSES]), + ), + ) + .limit(1) + .then((rows) => rows[0] ?? null); + return Boolean(row); + } + + async function markTimerHeartbeatChecked(agentId: string, source: WakeupOptions["source"]) { + if (source !== "timer") return; + await db + .update(agents) + .set({ + lastHeartbeatAt: new Date(), + updatedAt: new Date(), + }) + .where(eq(agents.id, agentId)); + } + function parseMaxTurnContinuationPolicy(agent: typeof agents.$inferSelect): MaxTurnContinuationPolicy { const runtimeConfig = parseObject(agent.runtimeConfig); const heartbeat = parseObject(runtimeConfig.heartbeat); @@ -6915,6 +7077,16 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) return null; } + const dailyCapBlock = await getHeartbeatDailyCapBlock(agent, parseHeartbeatPolicy(agent), { + excludeRunId: run.id, + checkRunCap: true, + checkCostCap: true, + }); + if (dailyCapBlock) { + await cancelQueuedRunForHeartbeatDailyCap(run, dailyCapBlock); + return null; + } + const issueId = readNonEmptyString(context.issueId); if (issueId) { const activePauseHold = await treeControlSvc.getActivePauseHoldGate(run.companyId, issueId); @@ -9801,7 +9973,10 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) ); } - async function releaseIssueExecutionAndPromote(run: typeof heartbeatRuns.$inferSelect) { + async function releaseIssueExecutionAndPromote( + run: typeof heartbeatRuns.$inferSelect, + options: { suppressImmediateRecovery?: boolean } = {}, + ) { const runContext = parseObject(run.contextSnapshot); const contextIssueId = readNonEmptyString(runContext.issueId); const taskKey = deriveTaskKeyWithHeartbeatFallback(runContext, null); @@ -10188,6 +10363,9 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) if (!issueNeedsImmediateRecovery) { return { kind: "released" as const }; } + if (options.suppressImmediateRecovery) { + return { kind: "released" as const }; + } const existingExecutionPath = await tx .select({ id: heartbeatRuns.id }) @@ -10410,6 +10588,14 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) ...patch, }); }; + const writeSkippedHeartbeatRequest = async (skipReason: string, details: Record) => { + await writeSkippedRequest(skipReason, { + payload: { + ...(payload ?? {}), + heartbeatSkip: details, + }, + }); + }; const company = await db .select({ status: companies.status }) @@ -10523,6 +10709,20 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) return null; } + const genericTimerWake = + source === "timer" && + !issueId && + !wakeCommentId && + !readNonEmptyString(enrichedContextSnapshot.taskId) && + !readNonEmptyString(enrichedContextSnapshot.taskKey); + if (policy.skipTimerWhenNoActionableWork && genericTimerWake && !(await hasActionableTimerWork(agent))) { + await writeSkippedHeartbeatRequest("heartbeat.timer.no_actionable_work", { + reason: "No assigned todo or in_progress issue requires this agent before timer adapter invocation.", + }); + await markTimerHeartbeatChecked(agentId, source); + return null; + } + if (issueId) { const activePauseHold = await treeControlSvc.getActivePauseHoldGate(agent.companyId, issueId); if (activePauseHold) { @@ -10996,6 +11196,41 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) } } + const dailyCapBlock = await getHeartbeatDailyCapBlock(agent, policy, {}, tx); + if (dailyCapBlock) { + const now = new Date(); + await tx.insert(agentWakeupRequests).values({ + companyId: agent.companyId, + agentId, + source, + triggerDetail, + reason: dailyCapBlock.reason, + payload: { + ...(payload ?? {}), + heartbeatSkip: { + reason: "Per-agent heartbeat daily cap reached before adapter invocation.", + observed: dailyCapBlock.observed, + limit: dailyCapBlock.limit, + }, + }, + status: "skipped", + requestedByActorType: opts.requestedByActorType ?? null, + requestedByActorId: opts.requestedByActorId ?? null, + idempotencyKey: opts.idempotencyKey ?? null, + finishedAt: now, + }); + if (source === "timer") { + await tx + .update(agents) + .set({ + lastHeartbeatAt: now, + updatedAt: now, + }) + .where(eq(agents.id, agentId)); + } + return { kind: "skipped" as const }; + } + const wakeupRequest = await tx .insert(agentWakeupRequests) .values({ @@ -11130,46 +11365,92 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) return mergedRun; } - const wakeupRequest = await db - .insert(agentWakeupRequests) - .values({ - companyId: agent.companyId, - agentId, - source, - triggerDetail, - reason, - payload, - status: "queued", - requestedByActorType: opts.requestedByActorType ?? null, - requestedByActorId: opts.requestedByActorId ?? null, - idempotencyKey: opts.idempotencyKey ?? null, - }) - .returning() - .then((rows) => rows[0]); + const queueOutcome = await db.transaction(async (tx) => { + await tx.execute( + sql`select id from agents where id = ${agentId} and company_id = ${agent.companyId} for update`, + ); - const newRun = await db - .insert(heartbeatRuns) - .values({ - companyId: agent.companyId, - agentId, - invocationSource: source, - triggerDetail, - status: "queued", - wakeupRequestId: wakeupRequest.id, - contextSnapshot: enrichedContextSnapshot, - sessionIdBefore: sessionBefore, - continuationAttempt, - }) - .returning() - .then((rows) => rows[0]); + const dailyCapBlock = await getHeartbeatDailyCapBlock(agent, policy, {}, tx); + if (dailyCapBlock) { + const now = new Date(); + await tx.insert(agentWakeupRequests).values({ + companyId: agent.companyId, + agentId, + source, + triggerDetail, + reason: dailyCapBlock.reason, + payload: { + ...(payload ?? {}), + heartbeatSkip: { + reason: "Per-agent heartbeat daily cap reached before adapter invocation.", + observed: dailyCapBlock.observed, + limit: dailyCapBlock.limit, + }, + }, + status: "skipped", + requestedByActorType: opts.requestedByActorType ?? null, + requestedByActorId: opts.requestedByActorId ?? null, + idempotencyKey: opts.idempotencyKey ?? null, + finishedAt: now, + }); + if (source === "timer") { + await tx + .update(agents) + .set({ + lastHeartbeatAt: now, + updatedAt: now, + }) + .where(eq(agents.id, agentId)); + } + return { kind: "skipped" as const }; + } - await db - .update(agentWakeupRequests) - .set({ - runId: newRun.id, - updatedAt: new Date(), - }) - .where(eq(agentWakeupRequests.id, wakeupRequest.id)); + const wakeupRequest = await tx + .insert(agentWakeupRequests) + .values({ + companyId: agent.companyId, + agentId, + source, + triggerDetail, + reason, + payload, + status: "queued", + requestedByActorType: opts.requestedByActorType ?? null, + requestedByActorId: opts.requestedByActorId ?? null, + idempotencyKey: opts.idempotencyKey ?? null, + }) + .returning() + .then((rows) => rows[0]); + + const newRun = await tx + .insert(heartbeatRuns) + .values({ + companyId: agent.companyId, + agentId, + invocationSource: source, + triggerDetail, + status: "queued", + wakeupRequestId: wakeupRequest.id, + contextSnapshot: enrichedContextSnapshot, + sessionIdBefore: sessionBefore, + continuationAttempt, + }) + .returning() + .then((rows) => rows[0]); + + await tx + .update(agentWakeupRequests) + .set({ + runId: newRun.id, + updatedAt: new Date(), + }) + .where(eq(agentWakeupRequests.id, wakeupRequest.id)); + + return { kind: "queued" as const, run: newRun }; + }); + + if (queueOutcome.kind === "skipped") return null; + const newRun = queueOutcome.run; publishLiveEvent({ companyId: newRun.companyId,