@rudderhq/server 0.7.11 → 0.7.12

package/dist/services/runtime-kernel/heartbeat.js

@@ -1,7 +1,7 @@
 import { hasSessionCompactionThresholds } from "@rudderhq/agent-runtime-utils";
 import { agentRuntimeState, agents, agentTaskSessions, agentWakeupRequests, goalResultProposals, goals, heartbeatRunEvents, heartbeatRuns, issues, organizations } from "@rudderhq/db";
 import { summarizeTokenUsage, toHeartbeatRun } from "@rudderhq/shared";
-import { and, asc, desc, eq, gt, gte, inArray, lte, ne, or, sql } from "drizzle-orm";
+import { and, asc, desc, eq, gt, gte, inArray, isNull, lte, ne, or, sql } from "drizzle-orm";
 import { randomUUID } from "node:crypto";
 import { runningProcesses } from "../../agent-runtimes/index.js";
 import { killChildProcessTree, parseObject } from "../../agent-runtimes/utils.js";
@@ -24,10 +24,11 @@ import { recordProductAnalyticsEvent } from "../product-analytics.js";
 import { appendHeartbeatRunEvent } from "../run-events.js";
 import { getRunLogStore } from "../run-log-store.js";
 import { workspaceOperationService } from "../workspace-operations.js";
+import { finishLatestHeartbeatRunAttempt } from "./heartbeat-attempt-ledger.js";
 export { prioritizeProjectWorkspaceCandidatesForRun } from "../agent-run-context.js";
 import * as heartbeatCore from "./heartbeat.core.js";
 import * as heartbeatSessions from "./heartbeat.sessions.js";
-const { MAX_LIVE_LOG_CHUNK_BYTES, HEARTBEAT_MAX_CONCURRENT_RUNS_DEFAULT, HEARTBEAT_MAX_CONCURRENT_RUNS_MIN, HEARTBEAT_MAX_CONCURRENT_RUNS_MAX, DEFERRED_WAKE_CONTEXT_KEY, DETACHED_PROCESS_ERROR_CODE, ORPHANED_PROCESS_TERMINATION_GRACE_MS, ORPHANED_PROCESS_KILL_WAIT_MS, ORPHANED_PROCESS_POLL_INTERVAL_MS, startLocksByAgent, MAX_RECOVERY_CHAIN_DEPTH, ISSUE_PASSIVE_FOLLOWUP_REASON, ISSUE_PASSIVE_FOLLOWUP_WAKE_SOURCE, ISSUE_PASSIVE_FOLLOWUP_FAILURE_REASON, ISSUE_PASSIVE_FOLLOWUP_MAX_ATTEMPTS, ISSUE_REVIEW_CLOSEOUT_REASON, ISSUE_REVIEW_CLOSEOUT_FAILURE_REASON, ISSUE_REVIEW_CLOSEOUT_MAX_ATTEMPTS, ISSUE_PASSIVE_FOLLOWUP_COOLDOWN_MS_BY_ATTEMPT, ISSUE_PASSIVE_FOLLOWUP_TIMER_CONTINUITY_MAX_WINDOW_MS, SESSIONED_LOCAL_ADAPTERS, heartbeatRunListColumns, appendExcerpt, appendTranscriptEntriesFromChunk, normalizeMaxConcurrentRuns, withAgentStartLock, readNonEmptyString, isIssueCommentMentionWake, buildHeartbeatAdapterInvokePayload, buildRecentDateKeys, buildDateKeysBetween, fallbackSkillLabel, normalizeLoadedSkill, normalizeLoadedSkillForPayload, emptySkillEvidenceCounts, incrementSkillEvidenceCount, strongestSkillEvidence, resolveSkillEvidence, readSkillEvidenceFromPayload, extractSkillSlugFromPath, collectSkillPathsFromText, collectStringValues, normalizeSkillUseFromPath, dedupeSkillUses, collectSkillUsesFromText, readToolCommandInput, isCommandTranscriptTool, isReadTranscriptTool, inferUsedSkillsFromTranscript, normalizeSkillCandidate, addSkillCandidate, readSkillReferenceSlug, collectSkillReferences, inferUsedSkillsFromPrompt, normalizeLedgerBillingType, resolveLedgerBiller, normalizeBilledCostCents, resolveLedgerScopeForRun } = heartbeatCore;
+const { MAX_LIVE_LOG_CHUNK_BYTES, HEARTBEAT_MAX_CONCURRENT_RUNS_DEFAULT, HEARTBEAT_MAX_CONCURRENT_RUNS_MIN, HEARTBEAT_MAX_CONCURRENT_RUNS_MAX, DEFERRED_WAKE_CONTEXT_KEY, DETACHED_PROCESS_ERROR_CODE, ORPHANED_PROCESS_TERMINATION_GRACE_MS, ORPHANED_PROCESS_KILL_WAIT_MS, ORPHANED_PROCESS_POLL_INTERVAL_MS, startLocksByAgent, MAX_RECOVERY_CHAIN_DEPTH, ISSUE_PASSIVE_FOLLOWUP_REASON, ISSUE_PASSIVE_FOLLOWUP_WAKE_SOURCE, ISSUE_PASSIVE_FOLLOWUP_FAILURE_REASON, ISSUE_PASSIVE_FOLLOWUP_MAX_ATTEMPTS, ISSUE_REVIEW_CLOSEOUT_REASON, ISSUE_REVIEW_CLOSEOUT_FAILURE_REASON, ISSUE_REVIEW_CLOSEOUT_MAX_ATTEMPTS, ISSUE_PASSIVE_FOLLOWUP_COOLDOWN_MS_BY_ATTEMPT, ISSUE_PASSIVE_FOLLOWUP_TIMER_CONTINUITY_MAX_WINDOW_MS, NETWORK_WAIT_MAX_ATTEMPTS, NETWORK_WAIT_EXHAUSTED_ERROR_CODE, NETWORK_WAIT_EXHAUSTED_ERROR, SESSIONED_LOCAL_ADAPTERS, heartbeatRunListColumns, appendExcerpt, appendTranscriptEntriesFromChunk, normalizeMaxConcurrentRuns, withAgentStartLock, readNonEmptyString, isIssueCommentMentionWake, buildHeartbeatAdapterInvokePayload, buildRecentDateKeys, buildDateKeysBetween, fallbackSkillLabel, normalizeLoadedSkill, normalizeLoadedSkillForPayload, emptySkillEvidenceCounts, incrementSkillEvidenceCount, strongestSkillEvidence, resolveSkillEvidence, readSkillEvidenceFromPayload, extractSkillSlugFromPath, collectSkillPathsFromText, collectStringValues, normalizeSkillUseFromPath, dedupeSkillUses, collectSkillUsesFromText, readToolCommandInput, isCommandTranscriptTool, isReadTranscriptTool, inferUsedSkillsFromTranscript, normalizeSkillCandidate, addSkillCandidate, readSkillReferenceSlug, collectSkillReferences, inferUsedSkillsFromPrompt, normalizeLedgerBillingType, resolveLedgerBiller, normalizeBilledCostCents, resolveLedgerScopeForRun } = heartbeatCore;
 const { buildExplicitResumeSessionOverride, normalizeUsageTotals, readRawUsageTotals, deriveNormalizedUsageDelta, formatCount, parseSessionCompactionPolicy, resolveRuntimeSessionParamsForWorkspace, parseIssueAssigneeAgentRuntimeOverrides, deriveTaskKey, shouldResetTaskSessionForWake, formatRuntimeWorkspaceWarningLog, describeSessionResetReason, deriveCommentId, enrichWakeContextSnapshot, mergeCoalescedContextSnapshot, issueCommentAuthorKind, issueCommentAuthorLabel, buildDeferredWakePayload, readDeferredWakeContext, readDeferredWakePayload, deriveDeferredWakeTaskKey, hydrateWakeContextSnapshot, firstNonEmptyLine, deriveRecoveryFailureKind, deriveRecoveryFailureSummary, mergeMissingRecoveryContextFields, hydrateRecoveryBaseContextSnapshot, buildRecoveryContextSnapshot, normalizePassiveFollowupContext, normalizeReviewCloseoutContext, passiveFollowupCooldownMs, issueHasReviewer, isAgentEligibleForTimerContinuation, hasCredibleTimerContinuation, buildPassiveFollowupContextSnapshot, runTaskKey, isSameTaskScope, isTrackedLocalChildProcessAdapter, isProcessAlive, waitForProcessExit, terminateOrphanedProcess, truncateDisplayId, normalizeAgentNameKey, defaultSessionCodec, getAgentRuntimeSessionCodec, normalizeSessionParams, resolveNextSessionState } = heartbeatSessions;
 import { createHeartbeatExecuteHandlers } from "./heartbeat.execute.js";
 import { createHeartbeatMiscHandlers } from "./heartbeat.misc.js";
@@ -546,6 +547,19 @@ export function heartbeatService(db, testHooks) {
         });
         if (updated) {
             publishRunStatus(updated);
+            await finishLatestHeartbeatRunAttempt(db, updated.id, {
+                status,
+                usageDeltaJson: updated.usageJson,
+                costUsd: updated.usageJson?.costUsd,
+                sessionDisplayId: updated.sessionIdAfter,
+                sessionParamsJson: updated.sessionParamsAfterJson,
+                errorCode: updated.errorCode,
+                error: updated.error,
+                finishedAt: updated.finishedAt ?? new Date(),
+            }).catch((error) => {
+                logger.warn({ err: error, runId: updated.id }, "failed to persist heartbeat attempt terminal state");
+                return null;
+            });
             const context = parseObject(updated.contextSnapshot);
             const requestId = agentIssueCreationSettlement?.requestId
                 ?? getAgentIssueCreationRequestIdFromRunContext(context);
@@ -1002,6 +1016,7 @@ export function heartbeatService(db, testHooks) {
                 .update(heartbeatRuns)
                 .set({
                 status: "running",
+                runningSubstate: "executing",
                 startedAt: currentRun.startedAt ?? claimedAt,
                 sessionIdBefore: sessionSelection.sessionDisplayId ?? readNonEmptyString(sessionSelection.sessionParams?.sessionId),
                 sessionParamsBeforeJson: sessionSelection.sessionParams,
@@ -1130,7 +1145,7 @@ export function heartbeatService(db, testHooks) {
             })
                 .from(heartbeatRuns)
                 .innerJoin(agents, eq(heartbeatRuns.agentId, agents.id))
-                .where(or(eq(heartbeatRuns.status, "running"), eq(heartbeatRuns.terminalEffectsPending, true)));
+                .where(or(and(eq(heartbeatRuns.status, "running"), or(isNull(heartbeatRuns.runningSubstate), ne(heartbeatRuns.runningSubstate, "waiting_for_network"))), eq(heartbeatRuns.terminalEffectsPending, true)));
             for (const { run, agentRuntimeType } of activeRuns) {
                 if (run.terminalEffectsPending) {
                     if (!run.processExitedAt) {
@@ -1276,7 +1291,7 @@ export function heartbeatService(db, testHooks) {
             .from(heartbeatRuns)
             .innerJoin(agents, eq(heartbeatRuns.agentId, agents.id))
             .leftJoin(heartbeatRunEvents, eq(heartbeatRunEvents.runId, heartbeatRuns.id))
-            .where(eq(heartbeatRuns.status, "running"))
+            .where(and(eq(heartbeatRuns.status, "running"), or(isNull(heartbeatRuns.runningSubstate), ne(heartbeatRuns.runningSubstate, "waiting_for_network"))))
             .groupBy(heartbeatRuns.id, agents.agentRuntimeType);
         const timedOut = [];
         for (const { run, agentRuntimeType, lastEventAt, eventCount } of activeRuns) {
@@ -1362,7 +1377,7 @@ export function heartbeatService(db, testHooks) {
         })
             .from(heartbeatRuns)
             .innerJoin(agents, eq(heartbeatRuns.agentId, agents.id))
-            .where(eq(heartbeatRuns.status, "running"));
+            .where(and(eq(heartbeatRuns.status, "running"), or(isNull(heartbeatRuns.runningSubstate), ne(heartbeatRuns.runningSubstate, "waiting_for_network"))));
         const timedOut = [];
         for (const { run, agentRuntimeType } of activeRuns) {
             if (opts?.recoveryCutoff && new Date(run.createdAt).getTime() >= opts.recoveryCutoff.getTime())
@@ -1372,7 +1387,7 @@ export function heartbeatService(db, testHooks) {
             const startedAt = run.startedAt ? new Date(run.startedAt).getTime() : null;
             if (!startedAt || !Number.isFinite(startedAt))
                 continue;
-            const runtimeMs = now.getTime() - startedAt;
+            const runtimeMs = Math.max(0, now.getTime() - startedAt - (run.networkWaitDurationMs ?? 0));
             if (runtimeMs < maxRuntimeMs)
                 continue;
             const message = `Run exceeded maximum duration of ${formatDurationMs(maxRuntimeMs)}`;
@@ -1797,6 +1812,145 @@ export function heartbeatService(db, testHooks) {
     const { releaseIssueExecutionAndPromote } = releaseHandlers;
     const { executeRun } = executeHandlers;
     const { resumeDeferredWakeupsForAgent, listProjectScopedRunIds, listProjectScopedWakeupIds, cancelPendingWakeupsForBudgetScope, cancelRunInternal, cancelActiveForAgentInternal, cancelBudgetScopeWork, retryRunInternal, buildSkillAnalytics } = miscHandlers;
+    async function recoverNetworkWaitingRunsLocked(opts) {
+        const now = opts?.now ?? new Date();
+        const candidates = await db
+            .select()
+            .from(heartbeatRuns)
+            .where(and(eq(heartbeatRuns.status, "running"), eq(heartbeatRuns.runningSubstate, "waiting_for_network"), or(isNull(heartbeatRuns.networkWaitNextRetryAt), lte(heartbeatRuns.networkWaitNextRetryAt, now))))
+            .orderBy(asc(heartbeatRuns.networkWaitNextRetryAt), asc(heartbeatRuns.updatedAt))
+            .limit(100);
+        const resumed = [];
+        for (const candidate of candidates) {
+            if (activeRunExecutions.has(candidate.id))
+                continue;
+            const agent = await getAgent(candidate.agentId);
+            if (!agent || ["paused", "terminated", "pending_approval"].includes(agent.status))
+                continue;
+            const checkpoint = candidate.recoveryCheckpoint ?? {};
+            const continuation = checkpoint.continuation;
+            const submissionPhase = checkpoint.submissionPhase;
+            const pristineReplayAllowed = continuation === "fresh_if_pristine"
+                && submissionPhase === "pre_submission"
+                && checkpoint.modelOutputObserved !== true
+                && checkpoint.toolActivityObserved !== true
+                && checkpoint.sideEffectRisk === "none";
+            const sameSessionAllowed = continuation === "resume_same_session"
+                && (typeof checkpoint.sessionId === "string"
+                    || Boolean(candidate.sessionIdBefore)
+                    || Boolean(candidate.sessionParamsBeforeJson));
+            if (!pristineReplayAllowed && !sameSessionAllowed) {
+                const unsafe = await transitionRunToTerminal(candidate.id, "failed", {
+                    finishedAt: now,
+                    error: "Network recovery cannot safely resume this provider attempt",
+                    errorCode: "network_resume_unsafe",
+                }, {
+                    expectedStatuses: ["running"],
+                    processExitedAt: now,
+                });
+                if (unsafe) {
+                    await setWakeupStatus(candidate.wakeupRequestId, "failed", {
+                        finishedAt: now,
+                        error: "Network recovery cannot safely resume this provider attempt",
+                    });
+                    await appendRunEvent(unsafe, {
+                        eventType: "network.resume_unsafe",
+                        stream: "system",
+                        level: "error",
+                        message: "network recovery failed closed because continuation was unsafe",
+                        payload: {
+                            continuation: continuation ?? null,
+                            submissionPhase: submissionPhase ?? null,
+                        },
+                    });
+                }
+                continue;
+            }
+            if ((candidate.networkWaitAttemptCount ?? 0) >= NETWORK_WAIT_MAX_ATTEMPTS) {
+                const claim = await claimExpiredHeartbeatRunExecution(db, candidate.id, { now });
+                if (!claim)
+                    continue;
+                if (candidate.chatConversationId && testHooks?.onNetworkWaitingRun) {
+                    await Promise.resolve(testHooks.onNetworkWaitingRun({
+                        ...claim.run,
+                        networkRecoveryExhausted: true,
+                    })).catch((error) => {
+                        logger.warn({ err: error, runId: candidate.id }, "network retry exhaustion handler failed");
+                    });
+                }
+                const exhausted = await transitionRunToTerminal(candidate.id, "failed", {
+                    finishedAt: now,
+                    error: NETWORK_WAIT_EXHAUSTED_ERROR,
+                    errorCode: NETWORK_WAIT_EXHAUSTED_ERROR_CODE,
+                }, {
+                    expectedStatuses: ["running"],
+                    processExitedAt: now,
+                    expectedExecutionOwnerToken: claim.ownerToken,
+                });
+                if (exhausted) {
+                    await setWakeupStatus(candidate.wakeupRequestId, "failed", {
+                        finishedAt: now,
+                        error: NETWORK_WAIT_EXHAUSTED_ERROR,
+                    });
+                    await appendRunEvent(exhausted, {
+                        eventType: "network.retry_exhausted",
+                        stream: "system",
+                        level: "error",
+                        message: "network recovery retry policy exhausted",
+                        payload: {
+                            attempt: candidate.networkWaitAttemptCount,
+                            maxAttempts: NETWORK_WAIT_MAX_ATTEMPTS,
+                            transport: checkpoint.transport ?? null,
+                            provider: checkpoint.provider ?? null,
+                            model: checkpoint.model ?? null,
+                        },
+                    });
+                }
+                continue;
+            }
+            const claim = await claimExpiredHeartbeatRunExecution(db, candidate.id, { now });
+            if (!claim)
+                continue;
+            const waitStartedAt = candidate.networkWaitStartedAt ? new Date(candidate.networkWaitStartedAt).getTime() : now.getTime();
+            const waitedMs = Math.max(0, now.getTime() - waitStartedAt);
+            const resumedRun = await db
+                .update(heartbeatRuns)
+                .set({
+                runningSubstate: "executing",
+                networkWaitStartedAt: null,
+                networkWaitNextRetryAt: null,
+                networkWaitDurationMs: sql `${heartbeatRuns.networkWaitDurationMs} + ${waitedMs}`,
+                executionOwnerToken: claim.ownerToken,
+                executionLeaseExpiresAt: new Date(now.getTime() + RUN_EXECUTION_LEASE_MS),
+                updatedAt: now,
+            })
+                .where(and(eq(heartbeatRuns.id, candidate.id), eq(heartbeatRuns.status, "running"), eq(heartbeatRuns.runningSubstate, "waiting_for_network"), eq(heartbeatRuns.executionOwnerToken, claim.ownerToken)))
+                .returning()
+                .then((rows) => rows[0] ?? null);
+            if (!resumedRun)
+                continue;
+            await appendRunEvent(resumedRun, {
+                eventType: "network.resumed",
+                stream: "system",
+                level: "info",
+                message: "run resumed after network recovery",
+                payload: { waitedMs, attempt: resumedRun.networkWaitAttemptCount },
+            });
+            resumed.push(resumedRun.id);
+            if (resumedRun.chatConversationId && testHooks?.onNetworkWaitingRun) {
+                const handled = await testHooks.onNetworkWaitingRun(resumedRun);
+                if (handled)
+                    continue;
+            }
+            void executeRun(resumedRun.id).catch((error) => {
+                logger.error({ err: error, runId: resumedRun.id }, "network-wait heartbeat recovery failed");
+            });
+        }
+        return { resumed: resumed.length, runIds: resumed };
+    }
+    async function recoverNetworkWaitingRuns(opts) {
+        return withHeartbeatRecoveryLock(() => recoverNetworkWaitingRunsLocked(opts));
+    }
     async function resumePendingWakeupRequests(opts = {}) {
         const pendingWakeups = await db.select().from(agentWakeupRequests).where(and(opts.orgId ? eq(agentWakeupRequests.orgId, opts.orgId) : undefined, inArray(agentWakeupRequests.status, [
             "queued",
@@ -2079,6 +2233,7 @@ export function heartbeatService(db, testHooks) {
         reapOrphanedRuns,
         reapInactiveRuns,
         reapTimedOutRuns,
+        recoverNetworkWaitingRuns,
         startNextQueuedRunForAgent,
         resumePendingWakeupRequests,
         resumeQueuedRuns,