feat(chat): add custom agent mailbox helpers

Author: gtremper

.changeset/tidy-mailboxes-wait.md

@@ -0,0 +1,6 @@
+---
+"@trigger.dev/core": patch
+"@trigger.dev/sdk": patch
+---
+
+Custom agent loops can now inspect pending chat input without consuming it and consume one mailbox record at a time with `chat.messages.hasPending()` and `chat.messages.next()`. Mailbox records include stable identifiers for tracing and redelivery.

docs/ai-chat/custom-agents.mdx

@@ -213,14 +213,54 @@ For full control, skip `createSession` and compose the primitives directly:
 
 | Primitive                       | Description                                                                                  |
 | ------------------------------- | -------------------------------------------------------------------------------------------- |
-| `chat.messages`                 | Input stream for incoming messages — use `.waitWithIdleTimeout()` to wait for the next turn  |
+| `chat.messages`                 | Mailbox for incoming messages — inspect buffered input, consume one record, or suspend until the next turn |
 | `chat.createStopSignal()`       | Create a managed stop signal wired to the stop input stream                                  |
 | `chat.pipeAndCapture(result)`   | Pipe a stream and capture the response; returns `{ message, status, error }`                 |
 | `chat.writeTurnComplete()`      | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors             |
 | `chat.MessageAccumulator`       | Accumulates conversation messages across turns                                               |
 | `chat.pipe(stream)`             | Pipe a stream to the frontend (no response capture)                                          |
 | `chat.cleanupAbortedParts(msg)` | Clean up incomplete parts from a stopped response                                            |
 
+### `chat.messages` mailbox
+
+`chat.messages` exposes the incoming message mailbox for hand-rolled loops:
+
+| Method | Behavior |
+| --- | --- |
+| `peek()` | Return the buffer head when it is a message, without consuming it; otherwise return `undefined` |
+| `hasPending()` | Resolve `true` when any message is buffered; does not consume it |
+| `next({ timeoutInSeconds? })` | Consume exactly one message record in channel order, or resolve `undefined` when the optional timeout elapses |
+| `on(handler)` | Consume messages as they arrive and invoke the handler |
+| `waitWithIdleTimeout(options)` | Wait warm, then suspend the run until the next message arrives |
+
+`hasPending()` checks the local, already-delivered buffer. It does not query the
+remote Session channel or start a subscription. Use `waitWithIdleTimeout()` when
+the loop needs to idle until future input arrives.
+
+`next()` returns a readonly record envelope:
+
+```ts
+const record = await chat.messages.next({ timeoutInSeconds: 5 });
+if (record) {
+  console.log(record.id, record.seqNum);
+  currentPayload = record.payload;
+}
+```
+
+- `id` is the append's stable idempotency key.
+- `seqNum` is the monotonic sequence on this Session's `.in` channel.
+- `payload` is the existing `ChatTaskWirePayload` delivered by the other mailbox methods.
+
+Both identifiers remain the same if the record is delivered again after a
+reconnect. Each `next()` call commits only the record it returns, so a loop that
+owns its own turn sequencing never advances past input it has not taken. By
+contrast, `on()` commits a record as soon as it dispatches the handler; avoid
+mixing `on()` and `next()` when a single loop owns mailbox consumption.
+
+The Session `.in` channel also carries control records such as handovers. If one
+comes before a message, `next()` leaves it for its own consumer and waits until
+that record has been handled.
+
 A complete loop:
 
 ```ts trigger/my-chat-raw.ts

docs/ai-chat/reference.mdx

@@ -506,7 +506,7 @@ All methods available on the `chat` object from `@trigger.dev/sdk/ai`.
 | `chat.pipeAndCapture(source, options?)`     | Pipe and capture the response; returns `{ message, status, error }`                                                         |
 | `chat.writeTurnComplete(options?)`          | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors                                            |
 | `chat.createStopSignal()`                   | Create a managed stop signal wired to the stop input stream                                                                  |
-| `chat.messages`                             | Input stream for incoming messages — use `.waitWithIdleTimeout()`                                                            |
+| `chat.messages`                             | Incoming message mailbox; supports non-consuming `.peek()` / `.hasPending()`, single-record `.next()`, `.on()`, and suspend-aware `.waitWithIdleTimeout()` |
 | `chat.local<T>({ id })`                     | Create a per-run typed local (see [`chat.local`](/ai-chat/chat-local))                             |
 | `chat.createStartSessionAction(taskId, options?)` | Returns a server action that creates a chat Session + triggers the first run + returns a session-scoped PAT. Idempotent on `(env, externalId)`.   |
 | `chat.waitForHandover(options)`             | Wait for a [`chat.headStart`](/ai-chat/fast-starts#handover-with-custom-agents) handover signal in a custom loop. Returns the signal or `null`. `chat.MessageAccumulator` wraps this as `consumeHandover()` / `applyHandover()` |

packages/core/src/v3/apiClient/runStream.test.ts

@@ -492,6 +492,7 @@ describe("SSEStreamSubscription v2 batch parsing — record kinds", () => {
   });
 
   type ParsedPart = {
+    recordId?: string;
     id: string;
     chunk: unknown;
     headers?: ReadonlyArray<readonly [string, string]>;
@@ -548,6 +549,7 @@ describe("SSEStreamSubscription v2 batch parsing — record kinds", () => {
     const parts = await sub.subscribe().then(drain);
 
     expect(parts).toHaveLength(1);
+    expect(parts[0]!.recordId).toBe("p1");
     expect(parts[0]!.id).toBe("5");
     expect(parts[0]!.chunk).toEqual({ type: "text-delta", delta: "hi" });
     expect(parts[0]!.headers).toEqual([]);

packages/core/src/v3/apiClient/runStream.ts

@@ -170,6 +170,9 @@ export interface StreamSubscriptionFactory {
 }
 
 export type SSEStreamPart<TChunk = unknown> = {
+  /** Stable logical record id from the S2 data envelope (`X-Part-Id` on append). */
+  recordId?: string;
+  /** S2 sequence number in decimal-string form. */
   id: string;
   chunk: TChunk;
   timestamp: number;
@@ -502,6 +505,7 @@ export class SSEStreamSubscription implements StreamSubscription {
                     chunkController.enqueue({
                       type: "part",
                       part: {
+                        recordId: parsedBody?.id,
                         id: record.seq_num.toString(),
                         chunk: parsedBody?.data,
                         timestamp: record.timestamp,

packages/core/src/v3/sessionStreams/index.ts

@@ -1,6 +1,12 @@
 import { getGlobal, registerGlobal } from "../utils/globals.js";
 import { NoopSessionStreamManager } from "./noopManager.js";
-import type { InputStreamOncePromise, SessionChannelIO, SessionStreamManager } from "./types.js";
+import type {
+  InputStreamOncePromise,
+  SessionChannelIO,
+  SessionStreamManager,
+  SessionStreamRecord,
+  SessionStreamRecordPredicate,
+} from "./types.js";
 import type { InputStreamOnceOptions } from "../realtimeStreams/types.js";
 
 const API_NAME = "session-streams";
@@ -43,10 +49,51 @@ export class SessionStreamsAPI implements SessionStreamManager {
     return this.#getManager().once(sessionId, io, options);
   }
 
+  public onceRecord(
+    sessionId: string,
+    io: SessionChannelIO,
+    options?: InputStreamOnceOptions
+  ): InputStreamOncePromise<SessionStreamRecord> {
+    const manager = this.#getManager();
+    if (!manager.onceRecord) {
+      throw new Error("The configured Session stream manager does not support record metadata");
+    }
+    return manager.onceRecord(sessionId, io, options);
+  }
+
+  public onceRecordWhere(
+    sessionId: string,
+    io: SessionChannelIO,
+    predicate: SessionStreamRecordPredicate,
+    options?: InputStreamOnceOptions
+  ): InputStreamOncePromise<SessionStreamRecord> {
+    const manager = this.#getManager();
+    if (!manager.onceRecordWhere) {
+      throw new Error("The configured Session stream manager does not support selective records");
+    }
+    return manager.onceRecordWhere(sessionId, io, predicate, options);
+  }
+
   public peek(sessionId: string, io: SessionChannelIO): unknown | undefined {
     return this.#getManager().peek(sessionId, io);
   }
 
+  public peekRecord(sessionId: string, io: SessionChannelIO): SessionStreamRecord | undefined {
+    return this.#getManager().peekRecord?.(sessionId, io);
+  }
+
+  public peekRecordWhere(
+    sessionId: string,
+    io: SessionChannelIO,
+    predicate: SessionStreamRecordPredicate
+  ): SessionStreamRecord | undefined {
+    const manager = this.#getManager();
+    if (!manager.peekRecordWhere) {
+      throw new Error("The configured Session stream manager does not support selective records");
+    }
+    return manager.peekRecordWhere(sessionId, io, predicate);
+  }
+
   public lastSeqNum(sessionId: string, io: SessionChannelIO): number | undefined {
     return this.#getManager().lastSeqNum(sessionId, io);
   }

packages/core/src/v3/sessionStreams/manager.test.ts

@@ -11,7 +11,7 @@ import type { SSEStreamPart } from "../apiClient/runStream.js";
 // an empty stream synchronously triggers a tight reconnect loop, so the
 // mock parks indefinitely instead.
 function singleShotApiClient(
-  records: Array<{ id: string; chunk: unknown; timestamp: number }>
+  records: Array<{ id: string; recordId?: string; chunk: unknown; timestamp: number }>
 ): ApiClient {
   let delivered = false;
   return {
@@ -160,3 +160,125 @@ describe("StandardSessionStreamManager — minTimestamp filter", () => {
     manager.disconnect();
   });
 });
+
+describe("StandardSessionStreamManager — record metadata", () => {
+  const sessionId = "session-records";
+  const io = "in" as const;
+  const records = [
+    {
+      id: "41",
+      recordId: "part-stable-1",
+      chunk: { kind: "message", payload: { id: "u1" } },
+      timestamp: 1000,
+    },
+    {
+      id: "42",
+      recordId: "part-stable-2",
+      chunk: { kind: "message", payload: { id: "u2" } },
+      timestamp: 2000,
+    },
+  ];
+
+  it("consumes one record at a time with stable id and sequence metadata", async () => {
+    const manager = new StandardSessionStreamManager(
+      singleShotApiClient(records),
+      "http://localhost"
+    );
+
+    const first = await manager.onceRecord(sessionId, io);
+    expect(first).toEqual({
+      ok: true,
+      output: {
+        id: "part-stable-1",
+        seqNum: 41,
+        data: { kind: "message", payload: { id: "u1" } },
+      },
+    });
+    expect(manager.peekRecord(sessionId, io)).toEqual({
+      id: "part-stable-2",
+      seqNum: 42,
+      data: { kind: "message", payload: { id: "u2" } },
+    });
+    expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(41);
+
+    const second = await manager.onceRecord(sessionId, io);
+    expect(second.ok && second.output.id).toBe("part-stable-2");
+    expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(42);
+
+    manager.disconnectStream(sessionId, io);
+    manager.disconnect();
+  });
+
+  it("returns the same envelope when a record is redelivered", async () => {
+    const firstDelivery = new StandardSessionStreamManager(
+      singleShotApiClient([records[0]!]),
+      "http://localhost"
+    );
+    const redelivery = new StandardSessionStreamManager(
+      singleShotApiClient([records[0]!]),
+      "http://localhost"
+    );
+
+    const first = await firstDelivery.onceRecord(sessionId, io);
+    const replayed = await redelivery.onceRecord(sessionId, io);
+
+    expect(first).toEqual(replayed);
+
+    firstDelivery.disconnectStream(sessionId, io);
+    firstDelivery.disconnect();
+    redelivery.disconnectStream(sessionId, io);
+    redelivery.disconnect();
+  });
+
+  it("does not consume a matching record past an earlier unmatched record", async () => {
+    const manager = new StandardSessionStreamManager(
+      singleShotApiClient([
+        {
+          id: "50",
+          recordId: "handover-1",
+          chunk: { kind: "handover" },
+          timestamp: 1000,
+        },
+        {
+          id: "51",
+          recordId: "message-1",
+          chunk: { kind: "message", payload: { id: "u1" } },
+          timestamp: 2000,
+        },
+      ]),
+      "http://localhost"
+    );
+
+    const pendingMessage = manager.onceRecordWhere(
+      sessionId,
+      io,
+      (record) => (record.data as { kind?: string }).kind === "message",
+      { timeoutMs: 200 }
+    );
+
+    expect(manager.peekRecordWhere(sessionId, io, (record) => record.id === "message-1")).toEqual({
+      id: "message-1",
+      seqNum: 51,
+      data: { kind: "message", payload: { id: "u1" } },
+    });
+    expect(manager.lastDispatchedSeqNum(sessionId, io)).toBeUndefined();
+
+    const handover = await manager.onceRecord(sessionId, io);
+    expect(handover).toEqual({
+      ok: true,
+      output: { id: "handover-1", seqNum: 50, data: { kind: "handover" } },
+    });
+    await expect(pendingMessage).resolves.toEqual({
+      ok: true,
+      output: {
+        id: "message-1",
+        seqNum: 51,
+        data: { kind: "message", payload: { id: "u1" } },
+      },
+    });
+    expect(manager.lastDispatchedSeqNum(sessionId, io)).toBe(51);
+
+    manager.disconnectStream(sessionId, io);
+    manager.disconnect();
+  });
+});