|
| 1 | +import { getMeter } from "@internal/tracing"; |
| 2 | +import type { ExecutionResult, TaskRunExecutionStatus } from "@trigger.dev/core/v3"; |
| 3 | +import { |
| 4 | + RUN_TOKEN_INSTALL_FAILURE_REASONS, |
| 5 | + RUN_TOKEN_REFRESH_BEFORE_EXPIRY_SECONDS, |
| 6 | + RUN_TOKEN_REFRESH_TRIGGERS, |
| 7 | + type RunTokenHeartbeatReport, |
| 8 | + type RunTokenRefreshTrigger, |
| 9 | +} from "@trigger.dev/core/v3/workers"; |
| 10 | +import type { PrismaClientOrTransaction } from "~/db.server"; |
| 11 | +import { env } from "~/env.server"; |
| 12 | +import { logger } from "~/services/logger.server"; |
| 13 | +import { mintRunToken, runTokenInjectionEnabled } from "~/services/oidc/runToken.server"; |
| 14 | +import { singleton } from "~/utils/singleton"; |
| 15 | + |
| 16 | +// Statuses in which the heartbeating runner still owns a live attempt, including the restore |
| 17 | +// window before it continues. |
| 18 | +const LIVE_ATTEMPT_STATUSES: ReadonlySet<TaskRunExecutionStatus> = new Set<TaskRunExecutionStatus>([ |
| 19 | + "PENDING_EXECUTING", |
| 20 | + "EXECUTING", |
| 21 | + "EXECUTING_WITH_WAITPOINTS", |
| 22 | + "QUEUED_EXECUTING", |
| 23 | +]); |
| 24 | + |
| 25 | +const REMAINING_LIFETIME_BUCKETS_SECONDS = [0, 60, 300, 600, 900, 1800, 2700, 3600]; |
| 26 | + |
| 27 | +const ISSUE_ERROR_REASONS = ["lookup_error", "sign_error"] as const; |
| 28 | + |
| 29 | +type IssueOutcome = |
| 30 | + | { result: "success"; reason: "none"; token: string } |
| 31 | + | { |
| 32 | + result: "skipped"; |
| 33 | + reason: "injection_disabled" | "stale_snapshot" | "not_executing" | "identity_mismatch"; |
| 34 | + } |
| 35 | + | { result: "error"; reason: (typeof ISSUE_ERROR_REASONS)[number] }; |
| 36 | + |
| 37 | +// Names and label values are relied on by internal dashboards and monitors. Labels stay bounded. |
| 38 | +const telemetry = singleton("runTokenRefreshTelemetry", () => { |
| 39 | + const meter = getMeter("run-token-refresh"); |
| 40 | + const issued = meter.createCounter("run_token.refresh.issued", { |
| 41 | + description: "Due run-token refresh requests handled by the webapp. Not delivery success", |
| 42 | + }); |
| 43 | + const installations = meter.createCounter("run_token.refresh.installations", { |
| 44 | + description: "Runner-reported installs of refreshed run tokens into the task process", |
| 45 | + }); |
| 46 | + |
| 47 | + // Export the rare failure series from zero, so a rate query sees a process's first failure. |
| 48 | + for (const trigger of RUN_TOKEN_REFRESH_TRIGGERS) { |
| 49 | + for (const reason of ISSUE_ERROR_REASONS) issued.add(0, { trigger, result: "error", reason }); |
| 50 | + for (const reason of RUN_TOKEN_INSTALL_FAILURE_REASONS) { |
| 51 | + installations.add(0, { trigger, result: "failure", reason }); |
| 52 | + } |
| 53 | + } |
| 54 | + |
| 55 | + return { |
| 56 | + issued, |
| 57 | + installations, |
| 58 | + remainingLifetime: meter.createHistogram("run_token.refresh.remaining_lifetime_seconds", { |
| 59 | + description: |
| 60 | + "Remaining lifetime of the run token installed in the task process, per heartbeat", |
| 61 | + unit: "s", |
| 62 | + advice: { explicitBucketBoundaries: REMAINING_LIFETIME_BUCKETS_SECONDS }, |
| 63 | + }), |
| 64 | + }; |
| 65 | +}); |
| 66 | + |
| 67 | +export type FindRunForRefresh = (runId: string) => Promise<{ |
| 68 | + taskIdentifier: string; |
| 69 | + runtimeEnvironmentId: string; |
| 70 | + lockedById: string | null; |
| 71 | +} | null>; |
| 72 | + |
| 73 | +export type RunTokenRefreshRequest = { |
| 74 | + report: RunTokenHeartbeatReport | undefined; |
| 75 | + /** The engine's heartbeat result, i.e. the run's latest snapshot. */ |
| 76 | + heartbeat: ExecutionResult; |
| 77 | + snapshotFriendlyId: string; |
| 78 | + /** Verified workload environment, when the supervisor forwarded one. */ |
| 79 | + environmentId?: string; |
| 80 | + findRun: FindRunForRefresh; |
| 81 | + replica: PrismaClientOrTransaction; |
| 82 | + now?: Date; |
| 83 | +}; |
| 84 | + |
| 85 | +/** |
| 86 | + * Records the runner's installed-token telemetry and, only when its token is due, mints a |
| 87 | + * replacement. Identity comes from the run's server-side state, never from the runner's token. |
| 88 | + */ |
| 89 | +export async function refreshRunTokenOnHeartbeat( |
| 90 | + request: RunTokenRefreshRequest |
| 91 | +): Promise<string | undefined> { |
| 92 | + const { report } = request; |
| 93 | + if (!report) return undefined; |
| 94 | + |
| 95 | + const { trigger } = report; |
| 96 | + const remainingSeconds = report.expiresAt - (request.now ?? new Date()).getTime() / 1_000; |
| 97 | + // OTel JS drops negative histogram values; zero stands for expired. |
| 98 | + telemetry.remainingLifetime.record(Math.max(0, remainingSeconds), { trigger }); |
| 99 | + |
| 100 | + if (report.installation) { |
| 101 | + const { installation } = report; |
| 102 | + telemetry.installations.add(1, { |
| 103 | + trigger: installation.trigger, |
| 104 | + result: installation.result, |
| 105 | + reason: installation.result === "success" ? "none" : installation.reason, |
| 106 | + }); |
| 107 | + } |
| 108 | + |
| 109 | + if (remainingSeconds > RUN_TOKEN_REFRESH_BEFORE_EXPIRY_SECONDS) return undefined; |
| 110 | + |
| 111 | + const outcome = await issue(request, trigger); |
| 112 | + telemetry.issued.add(1, { trigger, result: outcome.result, reason: outcome.reason }); |
| 113 | + return outcome.result === "success" ? outcome.token : undefined; |
| 114 | +} |
| 115 | + |
| 116 | +async function issue( |
| 117 | + { heartbeat, snapshotFriendlyId, environmentId, findRun, replica }: RunTokenRefreshRequest, |
| 118 | + trigger: RunTokenRefreshTrigger |
| 119 | +): Promise<IssueOutcome> { |
| 120 | + if (!env.RUN_TOKEN_ISSUER_URL) return { result: "skipped", reason: "injection_disabled" }; |
| 121 | + // Legacy workload calls carry no verified tenant scope, even in enforcement mode. |
| 122 | + if (!environmentId) return { result: "skipped", reason: "identity_mismatch" }; |
| 123 | + |
| 124 | + const { snapshot, run } = heartbeat; |
| 125 | + if (snapshot.friendlyId !== snapshotFriendlyId) { |
| 126 | + return { result: "skipped", reason: "stale_snapshot" }; |
| 127 | + } |
| 128 | + if (!LIVE_ATTEMPT_STATUSES.has(snapshot.executionStatus)) { |
| 129 | + return { result: "skipped", reason: "not_executing" }; |
| 130 | + } |
| 131 | + if (!run.attemptNumber) return { result: "skipped", reason: "identity_mismatch" }; |
| 132 | + |
| 133 | + let identity; |
| 134 | + try { |
| 135 | + identity = await loadIdentity(run.id, findRun, replica); |
| 136 | + } catch (error) { |
| 137 | + logger.warn("Run token refresh lookup failed", { |
| 138 | + runId: run.friendlyId, |
| 139 | + trigger, |
| 140 | + error: error instanceof Error ? error.message : String(error), |
| 141 | + }); |
| 142 | + return { result: "error", reason: "lookup_error" }; |
| 143 | + } |
| 144 | + |
| 145 | + if (!identity || identity.environment.id !== environmentId) { |
| 146 | + return { result: "skipped", reason: "identity_mismatch" }; |
| 147 | + } |
| 148 | + |
| 149 | + const { environment, taskIdentifier, deploymentFriendlyId } = identity; |
| 150 | + if (!runTokenInjectionEnabled(environment.organization.featureFlags)) { |
| 151 | + return { result: "skipped", reason: "injection_disabled" }; |
| 152 | + } |
| 153 | + |
| 154 | + try { |
| 155 | + const token = await mintRunToken({ |
| 156 | + organizationId: environment.organizationId, |
| 157 | + projectId: environment.projectId, |
| 158 | + environmentId: environment.id, |
| 159 | + environmentType: environment.type, |
| 160 | + runFriendlyId: run.friendlyId, |
| 161 | + attemptNumber: run.attemptNumber, |
| 162 | + taskIdentifier, |
| 163 | + deploymentFriendlyId, |
| 164 | + }); |
| 165 | + return { result: "success", reason: "none", token }; |
| 166 | + } catch (error) { |
| 167 | + logger.warn("Run token refresh signing failed", { |
| 168 | + runId: run.friendlyId, |
| 169 | + trigger, |
| 170 | + error: error instanceof Error ? error.message : String(error), |
| 171 | + }); |
| 172 | + return { result: "error", reason: "sign_error" }; |
| 173 | + } |
| 174 | +} |
| 175 | + |
| 176 | +// Runs once per due refresh (roughly once per token lifetime per live attempt), not per heartbeat. |
| 177 | +async function loadIdentity( |
| 178 | + runId: string, |
| 179 | + findRun: FindRunForRefresh, |
| 180 | + replica: PrismaClientOrTransaction |
| 181 | +) { |
| 182 | + const run = await findRun(runId); |
| 183 | + if (!run) return undefined; |
| 184 | + |
| 185 | + const [environment, lockedTask] = await Promise.all([ |
| 186 | + replica.runtimeEnvironment.findFirst({ |
| 187 | + where: { id: run.runtimeEnvironmentId }, |
| 188 | + select: { |
| 189 | + id: true, |
| 190 | + type: true, |
| 191 | + organizationId: true, |
| 192 | + projectId: true, |
| 193 | + organization: { select: { featureFlags: true } }, |
| 194 | + }, |
| 195 | + }), |
| 196 | + run.lockedById |
| 197 | + ? replica.backgroundWorkerTask.findFirst({ |
| 198 | + where: { id: run.lockedById }, |
| 199 | + select: { worker: { select: { deployment: { select: { friendlyId: true } } } } }, |
| 200 | + }) |
| 201 | + : undefined, |
| 202 | + ]); |
| 203 | + if (!environment) return undefined; |
| 204 | + |
| 205 | + return { |
| 206 | + environment, |
| 207 | + taskIdentifier: run.taskIdentifier, |
| 208 | + deploymentFriendlyId: lockedTask?.worker.deployment?.friendlyId, |
| 209 | + }; |
| 210 | +} |
0 commit comments