|
| 1 | +import { apiClientManager, resourceCatalog } from "@trigger.dev/core/v3"; |
| 2 | +import type { LocalsKey } from "@trigger.dev/core/v3"; |
| 3 | +import type { LanguageModel } from "ai"; |
| 4 | +import { |
| 5 | + installSessionWaitpointBackend, |
| 6 | + runInMockTaskContext, |
| 7 | + StandardSessionStreamManager, |
| 8 | +} from "@trigger.dev/core/v3/test"; |
| 9 | + |
| 10 | +export type RunRealChatAgentOptions = { |
| 11 | + agentId: string; |
| 12 | + baseUrl: string; |
| 13 | + addressingKey: string; |
| 14 | + /** |
| 15 | + * The environment secret key. The agent writes `.out` and reads `.in` as the |
| 16 | + * backend (PRIVATE auth) — the `.out` channel rejects client session tokens. |
| 17 | + */ |
| 18 | + secretKey: string; |
| 19 | + model: LanguageModel; |
| 20 | + modelLocal: LocalsKey<LanguageModel>; |
| 21 | + runId?: string; |
| 22 | + /** |
| 23 | + * Boot as a continuation of a previous run for the same session. Gates the |
| 24 | + * snapshot + `.out`/`.in` replay boot path, so the agent restores prior |
| 25 | + * history instead of treating the chat as brand new. |
| 26 | + */ |
| 27 | + continuation?: boolean; |
| 28 | + previousRunId?: string; |
| 29 | + /** |
| 30 | + * Idle window (seconds) before the turn loop falls through from the SSE |
| 31 | + * once() to the suspending `session.in.wait()`. Set this low to force the |
| 32 | + * suspend/resume path in a test. |
| 33 | + */ |
| 34 | + idleTimeoutInSeconds?: number; |
| 35 | + /** |
| 36 | + * `ctx.attempt.number`. A value greater than 1 makes the boot treat the run |
| 37 | + * as a retry (`couldHavePriorState`), restoring from the snapshot + `.in` |
| 38 | + * replay. Used to model an OOM retry re-dispatch. |
| 39 | + */ |
| 40 | + attemptNumber?: number; |
| 41 | +}; |
| 42 | + |
| 43 | +export type RunningAgent = { |
| 44 | + done: Promise<void>; |
| 45 | + close: () => Promise<void>; |
| 46 | +}; |
| 47 | + |
| 48 | +/** |
| 49 | + * Run the real `chat.agent` turn loop in-process, wired to a running webapp: |
| 50 | + * `apiClientManager` + a real `StandardSessionStreamManager` point the agent's |
| 51 | + * `.in`/`.out` at the webapp's Session streams (real S2 + SSE), the model is |
| 52 | + * injected via locals (so it survives without serialization), and turns are |
| 53 | + * driven by appending to `.in` over HTTP. Callers keep each message inside the |
| 54 | + * idle window and `close()` promptly so the run-engine suspend path is never |
| 55 | + * reached. |
| 56 | + */ |
| 57 | +export function runRealChatAgent(opts: RunRealChatAgentOptions): RunningAgent { |
| 58 | + apiClientManager.setGlobalAPIClientConfiguration({ |
| 59 | + baseURL: opts.baseUrl, |
| 60 | + accessToken: opts.secretKey, |
| 61 | + }); |
| 62 | + const apiClient = apiClientManager.clientOrThrow(); |
| 63 | + const manager = new StandardSessionStreamManager(apiClient, opts.baseUrl); |
| 64 | + const { backend, runtimeManager, restore } = installSessionWaitpointBackend(apiClient); |
| 65 | + |
| 66 | + const taskEntry = resourceCatalog.getTask(opts.agentId); |
| 67 | + if (!taskEntry) { |
| 68 | + restore(); |
| 69 | + throw new Error(`runRealChatAgent: agent "${opts.agentId}" is not registered`); |
| 70 | + } |
| 71 | + const runFn = taskEntry.fns.run as ( |
| 72 | + payload: unknown, |
| 73 | + params: { ctx: unknown; signal: AbortSignal } |
| 74 | + ) => Promise<unknown>; |
| 75 | + |
| 76 | + const runSignal = new AbortController(); |
| 77 | + runSignal.signal.addEventListener("abort", () => { |
| 78 | + try { |
| 79 | + backend.disable(); |
| 80 | + } catch {} |
| 81 | + }); |
| 82 | + const runId = opts.runId ?? `run_${opts.addressingKey}`; |
| 83 | + |
| 84 | + const idle = |
| 85 | + opts.idleTimeoutInSeconds !== undefined |
| 86 | + ? { idleTimeoutInSeconds: opts.idleTimeoutInSeconds } |
| 87 | + : {}; |
| 88 | + |
| 89 | + const done = ( |
| 90 | + runInMockTaskContext( |
| 91 | + async (drivers) => { |
| 92 | + drivers.locals.set(opts.modelLocal, opts.model); |
| 93 | + const payload = opts.continuation |
| 94 | + ? { |
| 95 | + chatId: opts.addressingKey, |
| 96 | + continuation: true, |
| 97 | + metadata: {}, |
| 98 | + ...idle, |
| 99 | + ...(opts.previousRunId ? { previousRunId: opts.previousRunId } : {}), |
| 100 | + } |
| 101 | + : { chatId: opts.addressingKey, trigger: "preload", metadata: {}, ...idle }; |
| 102 | + await runFn(payload, { ctx: drivers.ctx, signal: runSignal.signal }); |
| 103 | + }, |
| 104 | + { |
| 105 | + ctx: { |
| 106 | + run: { id: runId }, |
| 107 | + ...(opts.attemptNumber !== undefined ? { attempt: { number: opts.attemptNumber } } : {}), |
| 108 | + }, |
| 109 | + sessionStreamManager: manager, |
| 110 | + runtimeManager, |
| 111 | + } |
| 112 | + ) as Promise<void> |
| 113 | + ).finally(restore); |
| 114 | + |
| 115 | + return { |
| 116 | + done, |
| 117 | + close: async () => { |
| 118 | + try { |
| 119 | + await fetch( |
| 120 | + `${opts.baseUrl}/realtime/v1/sessions/${encodeURIComponent(opts.addressingKey)}/in/append`, |
| 121 | + { |
| 122 | + method: "POST", |
| 123 | + headers: { |
| 124 | + Authorization: `Bearer ${opts.secretKey}`, |
| 125 | + "Content-Type": "application/json", |
| 126 | + "X-Part-Id": "close", |
| 127 | + }, |
| 128 | + body: JSON.stringify({ |
| 129 | + kind: "message", |
| 130 | + payload: { chatId: opts.addressingKey, trigger: "close" }, |
| 131 | + }), |
| 132 | + } |
| 133 | + ); |
| 134 | + } catch {} |
| 135 | + runSignal.abort(); |
| 136 | + await done.catch(() => {}); |
| 137 | + }, |
| 138 | + }; |
| 139 | +} |
| 140 | + |
| 141 | +export type ChatAgentSessionOptions = Omit< |
| 142 | + RunRealChatAgentOptions, |
| 143 | + "runId" | "continuation" | "previousRunId" |
| 144 | +>; |
| 145 | + |
| 146 | +export type ChatAgentSession = { |
| 147 | + /** How many runs the session has spawned so far (1 fresh + N continuations). */ |
| 148 | + runCount: () => number; |
| 149 | + /** Close the currently-active run and stop spawning continuations. */ |
| 150 | + close: () => Promise<void>; |
| 151 | +}; |
| 152 | + |
| 153 | +/** |
| 154 | + * A session-scoped orchestrator that stands in for the run-engine's run |
| 155 | + * lifecycle: it starts a run, and whenever that run exits on its own |
| 156 | + * (`chat.endRun()` / `chat.requestUpgrade()`), spawns the next run as a |
| 157 | + * continuation (new run id, `continuation: true`, `previousRunId` threaded) |
| 158 | + * for the same session. That mirrors the server triggering a fresh run on the |
| 159 | + * next append after the previous run went terminal, and lets each continuation |
| 160 | + * restore prior history from the persisted snapshot. Runs never overlap: the |
| 161 | + * next spawn is chained on the previous run's `done` (after its manager |
| 162 | + * teardown), so the process-global managers are never installed twice at once. |
| 163 | + */ |
| 164 | +export function runChatAgentSession(opts: ChatAgentSessionOptions): ChatAgentSession { |
| 165 | + let closed = false; |
| 166 | + let index = 0; |
| 167 | + let current: RunningAgent | undefined; |
| 168 | + let previousRunId: string | undefined; |
| 169 | + |
| 170 | + const spawn = () => { |
| 171 | + index += 1; |
| 172 | + const runId = `run_${opts.addressingKey}_${index}`; |
| 173 | + current = runRealChatAgent({ |
| 174 | + ...opts, |
| 175 | + runId, |
| 176 | + continuation: index > 1, |
| 177 | + previousRunId, |
| 178 | + }); |
| 179 | + previousRunId = runId; |
| 180 | + const settle = () => { |
| 181 | + if (!closed) { |
| 182 | + spawn(); |
| 183 | + } |
| 184 | + }; |
| 185 | + current.done.then(settle, settle); |
| 186 | + }; |
| 187 | + |
| 188 | + spawn(); |
| 189 | + |
| 190 | + return { |
| 191 | + runCount: () => index, |
| 192 | + close: async () => { |
| 193 | + closed = true; |
| 194 | + await current?.close(); |
| 195 | + }, |
| 196 | + }; |
| 197 | +} |
0 commit comments