Skip to content

Commit e584be4

Browse files
nicktrnTrigger.dev RepoOps
nicktrn
authored and
Trigger.dev RepoOps
committed
feat(supervisor): accept a previous workload token key
Mono-RevId: 0416a5cd5350414eed26f3acac03e6e8d9f12a27
1 parent 42fa548 commit e584be4

7 files changed

Lines changed: 812 additions & 20 deletions

File tree

‎apps/supervisor/src/env.test.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,3 +265,17 @@ describe("runner token Secret settle delay", () => {
265265
).toBe(false);
266266
});
267267
});
268+
269+
describe("previous workload token key", () => {
270+
it("defaults the cutoff mode to log", () => {
271+
expect(Env.parse(base).WORKLOAD_TOKEN_PREVIOUS_CUTOFF_MODE).toBe("log");
272+
});
273+
274+
it("refuses a run cutoff without a previous secret", () => {
275+
const cutoff = { WORKLOAD_TOKEN_PREVIOUS_RUN_CUTOFF: "2026-06-01T00:00:00.000Z" };
276+
expect(Env.safeParse({ ...base, ...cutoff }).success).toBe(false);
277+
expect(
278+
Env.safeParse({ ...base, ...cutoff, WORKLOAD_TOKEN_SECRET_PREVIOUS: "previous" }).success
279+
).toBe(true);
280+
});
281+
});

‎apps/supervisor/src/env.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,9 @@ export const Env = z
4242
// also reject invalid tokens.
4343
WORKLOAD_TOKEN_SECRET: z.string().optional(),
4444
WORKLOAD_TOKEN_ENFORCEMENT: z.enum(["disabled", "log", "enforce"]).default("disabled"),
45+
WORKLOAD_TOKEN_SECRET_PREVIOUS: z.string().optional(),
46+
WORKLOAD_TOKEN_PREVIOUS_RUN_CUTOFF: z.string().datetime().optional(),
47+
WORKLOAD_TOKEN_PREVIOUS_CUTOFF_MODE: z.enum(["log", "enforce"]).default("log"),
4548
DELETE_CHECKPOINTS_ON_COMPLETION: BoolEnv.default(false), // irreversible; enable per cluster
4649
// Absolute expiry for minted deployment tokens. Deterministic (no wall-clock issued-at) so every
4750
// pod of a deployment carries an identical token; bump before this date. Must outlive any run.
@@ -466,6 +469,14 @@ export const Env = z
466469
path: ["WORKLOAD_TOKEN_SECRET"],
467470
});
468471
}
472+
if (data.WORKLOAD_TOKEN_PREVIOUS_RUN_CUTOFF && !data.WORKLOAD_TOKEN_SECRET_PREVIOUS) {
473+
ctx.addIssue({
474+
code: z.ZodIssueCode.custom,
475+
message:
476+
"WORKLOAD_TOKEN_PREVIOUS_RUN_CUTOFF is set but WORKLOAD_TOKEN_SECRET_PREVIOUS is not, so the cutoff would never apply",
477+
path: ["WORKLOAD_TOKEN_PREVIOUS_RUN_CUTOFF"],
478+
});
479+
}
469480
if (data.DELETE_CHECKPOINTS_ON_COMPLETION && data.WORKLOAD_TOKEN_ENFORCEMENT === "disabled") {
470481
ctx.addIssue({
471482
code: z.ZodIssueCode.custom,

‎apps/supervisor/src/workloadServer/index.ts‎

Lines changed: 104 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,12 @@ import { Counter } from "prom-client";
2626
import { env } from "../env.js";
2727
import { register } from "../metrics.js";
2828
import {
29+
checkPreviousKeyRunCutoff,
30+
previousKeyCutoffEnforced,
31+
previousKeyCutoffMode,
2932
verifyDeploymentIdHeader,
33+
type WorkloadAuthTransport,
34+
type WorkloadTokenKey,
3035
workloadTokenEnforced,
3136
workloadTokensEnabled,
3237
} from "../workloadToken.js";
@@ -93,6 +98,29 @@ function legacifyCheckpointType<T extends { checkpoint?: { type: string } | null
9398
return item;
9499
}
95100

101+
type SupervisorSocketData = WorkloadClientSocketData & {
102+
workloadTokenKey?: WorkloadTokenKey;
103+
workloadTokenClaims?: WorkloadDeploymentTokenClaims;
104+
};
105+
106+
type PreviousKeyCaller = {
107+
transport: WorkloadAuthTransport;
108+
remoteAddress?: string;
109+
remotePort?: number;
110+
runnerId?: string;
111+
userAgent?: string;
112+
route: string;
113+
claims?: WorkloadDeploymentTokenClaims;
114+
};
115+
116+
function runFriendlyIdFromMessage(message: unknown): string {
117+
const friendlyId = (message as { run?: { friendlyId?: unknown } } | undefined)?.run?.friendlyId;
118+
if (typeof friendlyId !== "string") {
119+
throw new Error("Run event without a string run friendlyId");
120+
}
121+
return friendlyId;
122+
}
123+
96124
type WorkloadServerEvents = {
97125
runConnected: [
98126
{
@@ -139,7 +167,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
139167
WorkloadClientToServerEvents,
140168
WorkloadServerToClientEvents,
141169
DefaultEventsMap,
142-
WorkloadClientSocketData
170+
SupervisorSocketData
143171
>;
144172

145173
private readonly runSockets = new Map<
@@ -148,7 +176,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
148176
WorkloadClientToServerEvents,
149177
WorkloadServerToClientEvents,
150178
DefaultEventsMap,
151-
WorkloadClientSocketData
179+
SupervisorSocketData
152180
>
153181
>();
154182

@@ -224,7 +252,8 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
224252
* which is why environmentId stays gated on enforce.
225253
*/
226254
private async authorizeWorkloadRequest(
227-
req: IncomingMessage
255+
req: IncomingMessage,
256+
runFriendlyId: string
228257
): Promise<
229258
{ ok: true; environmentId?: string; claims?: WorkloadDeploymentTokenClaims } | { ok: false }
230259
> {
@@ -238,6 +267,21 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
238267
return { ok: false };
239268
}
240269

270+
if (
271+
result.key === "previous" &&
272+
!this.previousKeyRunAllowed(runFriendlyId, {
273+
transport: "http",
274+
remoteAddress: req.socket.remoteAddress,
275+
remotePort: req.socket.remotePort,
276+
runnerId: this.runnerIdFromRequest(req),
277+
userAgent: this.headerValueFromRequest(req, "user-agent"),
278+
route: `${req.method} ${req.url}`,
279+
claims: result.claims,
280+
})
281+
) {
282+
return { ok: false };
283+
}
284+
241285
return {
242286
ok: true,
243287
environmentId:
@@ -248,6 +292,31 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
248292
};
249293
}
250294

295+
private previousKeyRunAllowed(runFriendlyId: string, caller: PreviousKeyCaller): boolean {
296+
const cutoff = checkPreviousKeyRunCutoff(runFriendlyId);
297+
const { claims, ...callerFields } = caller;
298+
const logArgs = {
299+
runId: runFriendlyId,
300+
key: "previous",
301+
cutoffResult: cutoff,
302+
cutoffMode: previousKeyCutoffMode,
303+
...callerFields,
304+
deployment: claims?.deployment,
305+
environmentId: claims?.environment_id,
306+
orgId: claims?.org_id,
307+
projectId: claims?.project_id,
308+
exp: claims?.exp,
309+
};
310+
311+
if (cutoff === "pass") {
312+
this.logger.debug("Previous-key token accepted", logArgs);
313+
return true;
314+
}
315+
316+
this.logger.warn("Previous-key token used for a run created after the cutoff", logArgs);
317+
return !(workloadTokenEnforced && previousKeyCutoffEnforced);
318+
}
319+
251320
/**
252321
* reclaimCheckpoints asks the checkpoint service to delete a finished run's checkpoint storage.
253322
* Must be called after the reply is sent: it never delays the runner.
@@ -416,7 +485,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
416485
"POST",
417486
async () => {
418487
const { req, reply, params, body } = ctx;
419-
const auth = await this.authorizeWorkloadRequest(req);
488+
const auth = await this.authorizeWorkloadRequest(req, params.runFriendlyId);
420489
if (!auth.ok) {
421490
reply.empty(401);
422491
return;
@@ -462,7 +531,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
462531
"POST",
463532
async () => {
464533
const { req, reply, params, body } = ctx;
465-
const auth = await this.authorizeWorkloadRequest(req);
534+
const auth = await this.authorizeWorkloadRequest(req, params.runFriendlyId);
466535
if (!auth.ok) {
467536
reply.empty(401);
468537
return;
@@ -519,7 +588,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
519588
"POST",
520589
async () => {
521590
const { req, reply, params, body } = ctx;
522-
const auth = await this.authorizeWorkloadRequest(req);
591+
const auth = await this.authorizeWorkloadRequest(req, params.runFriendlyId);
523592
if (!auth.ok) {
524593
reply.empty(401);
525594
return;
@@ -561,7 +630,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
561630
"GET",
562631
async () => {
563632
const { reply, params, req } = ctx;
564-
const auth = await this.authorizeWorkloadRequest(req);
633+
const auth = await this.authorizeWorkloadRequest(req, params.runFriendlyId);
565634
if (!auth.ok) {
566635
reply.empty(401);
567636
return;
@@ -665,7 +734,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
665734
"GET",
666735
async () => {
667736
const { req, reply, params } = ctx;
668-
const auth = await this.authorizeWorkloadRequest(req);
737+
const auth = await this.authorizeWorkloadRequest(req, params.runFriendlyId);
669738
if (!auth.ok) {
670739
reply.empty(401);
671740
return;
@@ -728,7 +797,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
728797
"GET",
729798
async () => {
730799
const { req, reply, params } = ctx;
731-
const auth = await this.authorizeWorkloadRequest(req);
800+
const auth = await this.authorizeWorkloadRequest(req, params.runFriendlyId);
732801
if (!auth.ok) {
733802
reply.empty(401);
734803
return;
@@ -814,15 +883,15 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
814883
WorkloadClientToServerEvents,
815884
WorkloadServerToClientEvents,
816885
DefaultEventsMap,
817-
WorkloadClientSocketData
886+
SupervisorSocketData
818887
> = io.of("/workload");
819888

820889
websocketServer.on("disconnect", (socket) => {
821890
this.logger.verbose("[WS] disconnect", socket.id);
822891
});
823892
websocketServer.use(async (socket, next) => {
824893
const setSocketDataFromHeader = (
825-
dataKey: keyof typeof socket.data,
894+
dataKey: keyof WorkloadClientSocketData,
826895
headerName: string,
827896
required: boolean = true
828897
) => {
@@ -866,6 +935,9 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
866935
return;
867936
}
868937

938+
socket.data.workloadTokenKey = result.key;
939+
socket.data.workloadTokenClaims = result.claims;
940+
869941
// Re-source the deployment id from the verified claim; the raw header may be an opaque token.
870942
// A legacy bare id is itself the friendlyId, so it's safe to keep.
871943
socket.data.deploymentFriendlyId =
@@ -1000,6 +1072,17 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
10001072
emitSocketLifecycle("run_disconnected", friendlyId, reason);
10011073
};
10021074

1075+
const previousKeySocketMayBind = (friendlyId: string, eventName: "run:start" | "run:stop") =>
1076+
socket.data.workloadTokenKey !== "previous" ||
1077+
this.previousKeyRunAllowed(friendlyId, {
1078+
transport: "ws",
1079+
remoteAddress: socket.handshake.address,
1080+
runnerId: socket.data.runnerId,
1081+
userAgent: socket.handshake.headers["user-agent"],
1082+
route: eventName,
1083+
claims: socket.data.workloadTokenClaims,
1084+
});
1085+
10031086
socketLogger.debug("wsServer socket connected", { ...getSocketMetadata() });
10041087

10051088
// FIXME: where does this get set?
@@ -1044,7 +1127,11 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
10441127
log.debug("Handling run:start");
10451128

10461129
try {
1047-
runConnected(message.run.friendlyId);
1130+
const friendlyId = runFriendlyIdFromMessage(message);
1131+
if (!previousKeySocketMayBind(friendlyId, "run:start")) {
1132+
return;
1133+
}
1134+
runConnected(friendlyId);
10481135
} catch (error) {
10491136
log.error("run:start error", { error });
10501137
}
@@ -1060,7 +1147,11 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
10601147
log.debug("Handling run:stop");
10611148

10621149
try {
1063-
runDisconnected(message.run.friendlyId, "run_stop_message");
1150+
const friendlyId = runFriendlyIdFromMessage(message);
1151+
if (!previousKeySocketMayBind(friendlyId, "run:stop")) {
1152+
return;
1153+
}
1154+
runDisconnected(friendlyId, "run_stop_message");
10641155
} catch (error) {
10651156
log.error("run:stop error", { error });
10661157
}

0 commit comments

Comments
 (0)