Skip to content

Commit 3cacfa4

Browse files
MariefayTrigger.dev RepoOps
Mariefay
authored and
Trigger.dev RepoOps
committed
refactor(webapp): extract the ensure-run-then-append helper from the session append route
Internal refactor of the session append endpoint. No behaviour change: appending to a session's input channel now goes through one shared server helper that other server code can reuse. Mono-RevId: c5c64aad40d4a1b19f2e8e383ff359276b97229f
1 parent 1ea559a commit 3cacfa4

5 files changed

Lines changed: 515 additions & 107 deletions

File tree

‎apps/webapp/app/routes/realtime.v1.sessions.$session.$io.append.ts‎

Lines changed: 37 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,13 @@
11
import { json } from "@remix-run/server-runtime";
2-
import { tryCatch } from "@trigger.dev/core/utils";
3-
import { nanoid } from "nanoid";
42
import { z } from "zod";
53
import { logger } from "~/services/logger.server";
6-
import { S2RealtimeStreams } from "~/services/realtime/s2realtimeStreams.server";
7-
import { ensureRunForSession } from "~/services/realtime/sessionRunManager.server";
84
import {
95
canonicalSessionAddressingKey,
106
resolveSessionWithWriterFallback,
117
} from "~/services/realtime/sessions.server";
12-
import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server";
138
import { stripClientIngressFields } from "~/services/realtime/sanitizeSessionInput.server";
14-
import {
15-
claimSessionStreamPart,
16-
releaseSessionStreamPart,
17-
} from "~/services/sessionStreamWaitpointCache.server";
18-
import { completeSessionStreamWaitpoints } from "~/services/realtime/sessionChannelAppend.server";
9+
import { ensureRunThenAppend } from "~/services/realtime/sessionChannelAppend.server";
1910
import { anyResource, createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
20-
import { ServiceValidationError } from "~/v3/services/common.server";
2111

2212
const ParamsSchema = z.object({
2313
session: z.string(),
@@ -108,116 +98,63 @@ const { action, loader } = createActionApiRoute(
10898
);
10999
}
110100

111-
const realtimeStream = getRealtimeStreamInstance(authentication.environment, "v2", {
101+
const addressingKey = canonicalSessionAddressingKey(session, params.session);
102+
103+
let part = await request.text();
104+
if (params.io === "in") {
105+
part = stripClientIngressFields(part);
106+
}
107+
108+
// Ensure a live run, then append. A failed ensure is best effort: the
109+
// record is durable in S2 and the next append retries the ensure, so the
110+
// caller still gets a 200 (the SSE tail delivers it once a run boots).
111+
const result = await ensureRunThenAppend({
112112
session,
113+
environment: authentication.environment,
114+
addressingKey,
115+
io: params.io,
116+
part,
117+
partId: request.headers.get("X-Part-Id") ?? undefined,
118+
reason: "continuation",
113119
});
114120

115-
if (!(realtimeStream instanceof S2RealtimeStreams)) {
121+
if (!result.ok && result.kind === "unsupported_backend") {
116122
return json(
117123
{ ok: false, error: "Session channels require the S2 realtime backend" },
118124
{ status: 501 }
119125
);
120126
}
121127

122-
// Probe + ensure a live run before appending. The append itself is
123-
// run-independent (S2 stream is durable, keyed on the session) but
124-
// the message is useless if no run is alive to consume it. The
125-
// probe is a single Prisma read; ensureRunForSession is no-op when
126-
// currentRunId is alive, so the steady-state cost is one extra
127-
// read in the hot path.
128-
//
129-
// Best-effort: if ensureRunForSession throws (e.g. the trigger
130-
// call fails transiently), still append to S2 — the record is
131-
// durable and the next append will retry the ensure step. Don't
132-
// surface the error to the caller; the SSE tail just won't deliver
133-
// it until a run boots.
134-
const [ensureError, ensureResult] = await tryCatch(
135-
ensureRunForSession({
136-
session,
137-
environment: authentication.environment,
138-
reason: "continuation",
139-
})
140-
);
141-
if (ensureError) {
142-
logger.error("Failed to ensureRunForSession on .in/append", {
128+
if (!result.ensure.ok) {
129+
// A close that lands mid-request is an expected race, not a failure.
130+
const level = result.ensure.sessionClosed ? "warn" : "error";
131+
logger[level]("Failed to ensureRunForSession on .in/append", {
143132
sessionId: session.id,
144133
externalId: session.externalId,
145-
error: ensureError,
134+
error: result.ensure.error,
146135
});
147136
}
148137

149-
const addressingKey = canonicalSessionAddressingKey(session, params.session);
150-
151-
let part = await request.text();
152-
if (params.io === "in") {
153-
part = stripClientIngressFields(part);
154-
}
155-
156-
const clientPartId = request.headers.get("X-Part-Id");
157-
const partId = clientPartId ?? nanoid(7);
158-
159-
// Idempotency on client-supplied part ids: atomically claim the id before
160-
// appending. A concurrent or retried POST that loses the claim skips the
161-
// append (no duplicate record) but still falls through to the drain below,
162-
// so a retry whose first attempt died before waking the waitpoint can still
163-
// recover it. The claim is released on append failure so a genuine retry
164-
// can re-claim and proceed.
165-
const wonClaim = clientPartId
166-
? await claimSessionStreamPart(
167-
authentication.environment.id,
168-
addressingKey,
169-
params.io,
170-
clientPartId
171-
)
172-
: true;
173-
174-
let appendSeq: number | undefined;
175-
if (wonClaim) {
176-
const [appendError, seq] = await tryCatch(
177-
realtimeStream.appendPartToSessionStream(part, partId, addressingKey, params.io)
178-
);
179-
appendSeq = seq ?? undefined;
180-
181-
if (appendError) {
182-
if (clientPartId) {
183-
await releaseSessionStreamPart(
184-
authentication.environment.id,
185-
addressingKey,
186-
params.io,
187-
clientPartId
188-
);
189-
}
190-
if (appendError instanceof ServiceValidationError) {
191-
return json(
192-
{ ok: false, error: appendError.message },
193-
{ status: appendError.status ?? 422 }
194-
);
195-
}
196-
logger.error("Failed to append to session stream", {
197-
sessionId: session.id,
198-
io: params.io,
199-
error: appendError,
200-
});
201-
return json(
202-
{ ok: false, error: "Something went wrong, please try again." },
203-
{ status: 500 }
204-
);
138+
if (!result.ok) {
139+
if (result.kind === "rejected") {
140+
return json({ ok: false, error: result.message }, { status: result.status });
205141
}
142+
logger.error("Failed to append to session stream", {
143+
sessionId: session.id,
144+
io: params.io,
145+
error: result.error,
146+
});
147+
return json({ ok: false, error: "Something went wrong, please try again." }, { status: 500 });
206148
}
207149

208-
await completeSessionStreamWaitpoints(
209-
authentication.environment.id,
210-
addressingKey,
211-
params.io,
212-
part
213-
);
214-
215150
// `seq` lets the client correlate this send to the turn that consumes it.
216151
return json(
217152
{
218153
ok: true,
219-
seq: appendSeq,
220-
...(ensureResult?.pendingVersion ? { pendingVersion: true } : {}),
154+
seq: result.seq,
155+
...(result.ensure.ok && result.ensure.result.pendingVersion
156+
? { pendingVersion: true }
157+
: {}),
221158
},
222159
{ status: 200 }
223160
);

‎apps/webapp/app/services/realtime/sessionChannelAppend.server.ts‎

Lines changed: 111 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,23 @@
1+
import type { Session } from "@trigger.dev/database";
12
import { tryCatch } from "@trigger.dev/core/utils";
23
import { nanoid } from "nanoid";
34
import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
45
import { logger } from "~/services/logger.server";
56
import { S2RealtimeStreams } from "~/services/realtime/s2realtimeStreams.server";
7+
import {
8+
ensureRunForSession,
9+
SessionClosedError,
10+
type EnsureRunReason,
11+
type EnsureRunResult,
12+
} from "~/services/realtime/sessionRunManager.server";
613
import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server";
7-
import { drainSessionStreamWaitpoints } from "~/services/sessionStreamWaitpointCache.server";
14+
import {
15+
claimSessionStreamPart,
16+
drainSessionStreamWaitpoints,
17+
releaseSessionStreamPart,
18+
} from "~/services/sessionStreamWaitpointCache.server";
819
import { engine } from "~/v3/runEngine.server";
20+
import { ServiceValidationError } from "~/v3/services/common.server";
921

1022
type SessionChannelIO = "out" | "in";
1123

@@ -17,7 +29,7 @@ type SessionChannelIO = "out" | "in";
1729
* key the agent registered with via `sessions.open(...).in.wait()`, so writers
1830
* and readers converge regardless of which URL form they used.
1931
*/
20-
export async function completeSessionStreamWaitpoints(
32+
async function completeSessionStreamWaitpoints(
2133
environmentId: string,
2234
addressingKey: string,
2335
io: SessionChannelIO,
@@ -109,3 +121,100 @@ export async function appendServerSessionRecord({
109121
await completeSessionStreamWaitpoints(environment.id, addressingKey, io, part);
110122
return true;
111123
}
124+
125+
type EnsureOutcome =
126+
| { ok: true; result: EnsureRunResult }
127+
/** `sessionClosed`: the run claim lost to a close, so there is no run to ensure. */
128+
| { ok: false; error: unknown; sessionClosed: boolean };
129+
130+
export type EnsureRunThenAppendResult =
131+
/** `seq` is undefined when another append already claimed `partId`. */
132+
| { ok: true; seq: number | undefined; ensure: EnsureOutcome }
133+
| { ok: false; kind: "unsupported_backend" }
134+
| { ok: false; kind: "rejected"; status: number; message: string; ensure: EnsureOutcome }
135+
| { ok: false; kind: "failed"; error: unknown; ensure: EnsureOutcome };
136+
137+
/**
138+
* Ensure the session has a live run, then append one record and wake its waiters. Outcomes are
139+
* returned, not thrown, so each caller picks its own policy.
140+
*
141+
* Ensure goes first: a revived run reads the record from the channel (continuations never carry
142+
* it in the payload), and a rejected append doesn't skip the ensure. A `partId` is claimed before
143+
* appending, so a retry appends nothing but still wakes waiters; a failed append releases it.
144+
*/
145+
export async function ensureRunThenAppend({
146+
session,
147+
environment,
148+
addressingKey,
149+
io,
150+
part,
151+
partId,
152+
reason,
153+
}: {
154+
session: Pick<
155+
Session,
156+
| "id"
157+
| "friendlyId"
158+
| "taskIdentifier"
159+
| "triggerConfig"
160+
| "currentRunId"
161+
| "currentRunVersion"
162+
| "streamBasinName"
163+
>;
164+
environment: AuthenticatedEnvironment;
165+
addressingKey: string;
166+
io: SessionChannelIO;
167+
part: string;
168+
partId?: string;
169+
reason: EnsureRunReason;
170+
}): Promise<EnsureRunThenAppendResult> {
171+
const realtimeStream = getRealtimeStreamInstance(environment, "v2", { session });
172+
173+
if (!(realtimeStream instanceof S2RealtimeStreams)) {
174+
return { ok: false, kind: "unsupported_backend" };
175+
}
176+
177+
const ensure = await ensureRun({ session, environment, reason });
178+
179+
const wonClaim = partId
180+
? await claimSessionStreamPart(environment.id, addressingKey, io, partId)
181+
: true;
182+
183+
let seq: number | undefined;
184+
if (wonClaim) {
185+
const [appendError, appendedSeq] = await tryCatch(
186+
realtimeStream.appendPartToSessionStream(part, partId ?? nanoid(7), addressingKey, io)
187+
);
188+
189+
if (appendError) {
190+
if (partId) {
191+
await releaseSessionStreamPart(environment.id, addressingKey, io, partId);
192+
}
193+
if (appendError instanceof ServiceValidationError) {
194+
return {
195+
ok: false,
196+
kind: "rejected",
197+
status: appendError.status ?? 422,
198+
message: appendError.message,
199+
ensure,
200+
};
201+
}
202+
return { ok: false, kind: "failed", error: appendError, ensure };
203+
}
204+
205+
seq = appendedSeq ?? undefined;
206+
}
207+
208+
await completeSessionStreamWaitpoints(environment.id, addressingKey, io, part);
209+
return { ok: true, seq, ensure };
210+
}
211+
212+
async function ensureRun(
213+
params: Parameters<typeof ensureRunForSession>[0]
214+
): Promise<EnsureOutcome> {
215+
const [error, result] = await tryCatch(ensureRunForSession(params));
216+
if (error) {
217+
return { ok: false, error, sessionClosed: error instanceof SessionClosedError };
218+
}
219+
return { ok: true, result };
220+
}

‎apps/webapp/app/services/realtime/sessionRunManager.server.ts‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ const SessionTriggerConfigSchema = SessionTriggerConfigZod;
2525

2626
export type SessionTriggerConfig = z.infer<typeof SessionTriggerConfigSchema>;
2727

28-
type EnsureRunReason = "initial" | "continuation" | "upgrade" | "manual";
28+
export type EnsureRunReason = "initial" | "continuation" | "upgrade" | "manual";
2929

3030
/**
3131
* Hard cap on how many times `ensureRunForSession` will recurse on the
@@ -264,7 +264,7 @@ export async function ensureRunForSession(
264264
// The claim lost to a close, not to another writer. Recursing would
265265
// trigger and cancel a run per attempt and then throw; there is nothing to
266266
// ensure for a conversation that is over.
267-
throw new SessionRunManagerError(`Session ${session.id} is closed`);
267+
throw new SessionClosedError(`Session ${session.id} is closed`);
268268
}
269269

270270
if (fresh.currentRunId) {
@@ -538,7 +538,7 @@ export async function swapSessionRun(params: SwapSessionRunParams): Promise<Swap
538538
// The claim lost to a close, not to another swap. The run triggered above
539539
// is already being cancelled; reporting a winner here would tell the caller
540540
// a closed session has a live run.
541-
throw new SessionRunManagerError(`Session ${session.id} is closed`);
541+
throw new SessionClosedError(`Session ${session.id} is closed`);
542542
}
543543

544544
// Mirror `ensureRunForSession`'s "session vanished" branch: if we
@@ -613,5 +613,10 @@ async function cancelLostRaceRun(
613613
}
614614

615615
class SessionRunManagerError extends Error {
616-
readonly name = "SessionRunManagerError";
616+
readonly name: string = "SessionRunManagerError";
617+
}
618+
619+
/** The run claim lost to a close: there is no run to ensure for this session. */
620+
export class SessionClosedError extends SessionRunManagerError {
621+
readonly name: string = "SessionClosedError";
617622
}

0 commit comments

Comments
 (0)