diff --git a/apps/webapp/app/runEngine/concerns/queues.server.ts b/apps/webapp/app/runEngine/concerns/queues.server.ts index a51ec6243..ce25696d1 100644 --- a/apps/webapp/app/runEngine/concerns/queues.server.ts +++ b/apps/webapp/app/runEngine/concerns/queues.server.ts @@ -107,19 +107,26 @@ export class DefaultQueueManager implements QueueManager { queueName = specifiedQueue.name; lockedQueueId = specifiedQueue.id; - // Only fetch task for TTL if caller didn't provide a per-trigger TTL - if (request.body.options?.ttl === undefined) { - const lockedTask = await this.replicaPrisma.backgroundWorkerTask.findFirst({ - where: { - workerId: lockedBackgroundWorker.id, - runtimeEnvironmentId: request.environment.id, - slug: request.taskId, - }, - select: { ttl: true }, - }); + // Always fetch the task so we can resolve `triggerSource` (which + // becomes `taskKind` on annotations and replicates to ClickHouse). + // Without this, AGENT/SCHEDULED runs triggered with + // `lockToVersion` + a queue override would be annotated as + // STANDARD and disappear from the run-list "Source" filter. + // `ttl` is read from the same row but only used when the caller + // didn't specify a per-trigger TTL. + const lockedTask = await this.replicaPrisma.backgroundWorkerTask.findFirst({ + where: { + workerId: lockedBackgroundWorker.id, + runtimeEnvironmentId: request.environment.id, + slug: request.taskId, + }, + select: { ttl: true, triggerSource: true }, + }); + if (request.body.options?.ttl === undefined) { taskTtl = lockedTask?.ttl; } + taskKind = lockedTask?.triggerSource; } else { // No queue override - fetch task with queue to get both default queue and TTL const lockedTask = await this.replicaPrisma.backgroundWorkerTask.findFirst({ diff --git a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts index a07f08728..d0e60242d 100644 --- a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts +++ b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts @@ -260,11 +260,15 @@ export function mockChatAgent( }); // Install the session open override so `sessions.open(id)` returns a - // SessionHandle with an in-memory `.out` that captures writes. `.in` - // stays the real SessionInputChannel — it routes through the - // `sessionStreams` global, which the mock-task-context installs as a - // TestSessionStreamManager. - __setSessionOpenImplForTests((id) => createTestSessionHandle(id, sessionOutState)); + // SessionHandle with an in-memory `.out` that captures writes. The + // `.in` channel routes record subscriptions (`on`/`once`/`peek`) + // through the `sessionStreams` global — the mock task context + // installs a `TestSessionStreamManager` there — and stubs `wait()` + // so the suspend path resolves cleanly on `runSignal.abort()` without + // touching the api client. + __setSessionOpenImplForTests((id) => + createTestSessionHandle(id, sessionOutState, () => runSignal?.signal) + ); // Install the session start override so any test path that invokes // `sessions.start()` (typically through a server action shim like diff --git a/packages/trigger-sdk/src/v3/test/test-session-handle.ts b/packages/trigger-sdk/src/v3/test/test-session-handle.ts index 6064394a5..71bc9d8d7 100644 --- a/packages/trigger-sdk/src/v3/test/test-session-handle.ts +++ b/packages/trigger-sdk/src/v3/test/test-session-handle.ts @@ -4,7 +4,7 @@ import type { StreamWriteResult, WriterStreamOptions, } from "@trigger.dev/core/v3"; -import { ensureReadableStream } from "@trigger.dev/core/v3"; +import { ensureReadableStream, ManualWaitpointPromise } from "@trigger.dev/core/v3"; import { SessionHandle, SessionInputChannel, @@ -13,6 +13,51 @@ import { SessionSubscribeOptions, } from "../sessions.js"; +/** + * Stub for `SessionInputChannel.wait` that skips the apiClient round-trip + * the production path makes via `createSessionStreamWaitpoint`. Without + * this override, every test that exercises the suspend fallback (e.g. + * the `chat.handover` idle-timeout case) throws `ApiClientMissingError` + * because `apiClientManager.clientOrThrow()` runs in a test process that + * has no `TRIGGER_SECRET_KEY`. + * + * The promise resolves with `{ ok: false, error }` when the harness + * aborts its run signal — that mimics production semantics (suspended + * until something happens, returns cleanly on abort) without making a + * network call. + */ +class TestSessionInputChannel extends SessionInputChannel { + constructor(sessionId: string, private readonly getAbortSignal: () => AbortSignal | undefined) { + super(sessionId); + } + + // Override only the `wait` path. `on` / `once` / `peek` / `send` + // continue to flow through the real `sessionStreams` global, which + // the mock task context installs as a `TestSessionStreamManager`. + wait(): ManualWaitpointPromise { + return new ManualWaitpointPromise((resolve: (value: { ok: false; error: Error }) => void) => { + const signal = this.getAbortSignal(); + if (!signal) { + // Harness hasn't wired up its run signal yet — nothing to abort + // on. Stay pending; the run loop should never reach this state + // in practice but we don't want to throw here either. + return; + } + const onAbort = () => { + resolve({ + ok: false, + error: new Error("session.in.wait() aborted by test harness"), + }); + }; + if (signal.aborted) { + onAbort(); + return; + } + signal.addEventListener("abort", onAbort, { once: true }); + }); + } +} + /** * Per-session in-memory state collected from `.out` writes during a test. * Owned by the mock-chat-agent harness; updated by {@link TestSessionOutputChannel}. @@ -201,16 +246,23 @@ export class TestSessionOutputChannel extends SessionOutputChannel { /** * Construct a {@link SessionHandle} whose `.out` channel captures writes in - * memory while `.in` reuses the real {@link SessionInputChannel} (which - * routes through the `sessionStreams` global — the mock task context - * installs a `TestSessionStreamManager` there). + * memory and whose `.in` channel routes through the `sessionStreams` + * global for record subscriptions (`on` / `once` / `peek`) but stubs + * `wait()` to skip the apiClient round-trip — see + * {@link TestSessionInputChannel}. + * + * `getAbortSignal` lets the channel observe the harness's run signal so + * `wait()` resolves cleanly on close. Pass a getter (not the signal + * directly) so the channel reads it lazily — the harness creates its + * `AbortController` after the override is installed. */ export function createTestSessionHandle( sessionId: string, - state: TestSessionOutState + state: TestSessionOutState, + getAbortSignal: () => AbortSignal | undefined = () => undefined ): SessionHandle { return new SessionHandle(sessionId, { - in: new SessionInputChannel(sessionId), + in: new TestSessionInputChannel(sessionId, getAbortSignal), out: new TestSessionOutputChannel(sessionId, state), }); }