Skip to content

Commit 354f8fb

Browse files
committed
feat(sdk,core,webapp): the transcript storage option, read API and conformance suite
chat.agent takes a storage option, with the platform snapshot as the default and the TranscriptStorage types, defaultStorage, snapshotTranscriptStorage, memoryTranscriptStorage and reduceTranscriptChanges exported from @trigger.dev/sdk/ai. chat.createLoadTranscriptAction(storage) reads a conversation on the server for any storage, and useLoadTranscript renders it in the browser and seeds the transport's resume cursor. runTranscriptStorageTests from @trigger.dev/sdk/ai/test is the contract a storage implementation has to meet. A new secret-key endpoint, GET /api/v1/sessions/:id/transcript, pages the platform snapshot server-side so rendering the last page of a long conversation does not download all of it; the default storage uses it for paged reads and falls back to the whole blob.
1 parent 10a1d68 commit 354f8fb

13 files changed

Lines changed: 948 additions & 23 deletions

File tree

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
import { json } from "@remix-run/server-runtime";
2+
import { pageTranscriptEntries, parseTranscriptSnapshot } from "@trigger.dev/core/v3";
3+
import { z } from "zod";
4+
import { $replica } from "~/db.server";
5+
import { chatSnapshotStorageKey } from "~/services/realtime/chatSnapshot.server";
6+
import { resolveSessionByIdOrExternalId } from "~/services/realtime/sessions.server";
7+
import { anyResource, createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server";
8+
import { downloadPacketFromObjectStore } from "~/v3/objectStore.server";
9+
10+
const ParamsSchema = z.object({
11+
sessionId: z.string(),
12+
});
13+
14+
const SearchParamsSchema = z.object({
15+
limit: z.coerce.number().int().min(1).max(1000).optional(),
16+
before: z.string().optional(),
17+
});
18+
19+
function sessionResource(
20+
paramId: string,
21+
session: { friendlyId: string; externalId: string | null } | null | undefined
22+
) {
23+
const ids = new Set<string>([paramId]);
24+
if (session) {
25+
ids.add(session.friendlyId);
26+
if (session.externalId) ids.add(session.externalId);
27+
}
28+
return anyResource([...ids].map((id) => ({ type: "sessions" as const, id })));
29+
}
30+
31+
export const loader = createLoaderApiRoute(
32+
{
33+
params: ParamsSchema,
34+
searchParams: SearchParamsSchema,
35+
corsStrategy: "none",
36+
findResource: async (params, auth) =>
37+
resolveSessionByIdOrExternalId($replica, auth.environment.id, params.sessionId),
38+
authorization: {
39+
action: "read",
40+
resource: (session, params) => sessionResource(params.sessionId, session),
41+
},
42+
},
43+
async ({ authentication, resource: session, searchParams }) => {
44+
if (!session) {
45+
return json({ error: "Session not found" }, { status: 404 });
46+
}
47+
48+
let body: unknown;
49+
try {
50+
const packet = await downloadPacketFromObjectStore(
51+
{ dataType: "application/store", data: chatSnapshotStorageKey(session) },
52+
authentication.environment
53+
);
54+
body = typeof packet.data === "string" ? JSON.parse(packet.data) : undefined;
55+
} catch {
56+
body = undefined;
57+
}
58+
59+
const snapshot = parseTranscriptSnapshot(body);
60+
if (!snapshot) {
61+
return json({ messages: [], state: null });
62+
}
63+
64+
const page = pageTranscriptEntries(snapshot.messages, searchParams);
65+
return json({
66+
messages: page.entries.map((entry) => entry.message),
67+
state: snapshot.state,
68+
cursors: {
69+
lastOutEventId: snapshot.lastOutEventId,
70+
lastInEventId: snapshot.lastInEventId,
71+
},
72+
nextCursor: page.nextCursor,
73+
});
74+
}
75+
);

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ import {
5151
CreateSessionStreamWaitpointResponseBody,
5252
CreateStreamResponseBody,
5353
CreateUploadPayloadUrlResponseBody,
54+
SessionTranscriptResponseBody,
5455
CreateWaitpointTokenResponseBody,
5556
CreatedSessionResponseBody,
5657
DeletedScheduleObject,
@@ -709,6 +710,30 @@ export class ApiClient {
709710
);
710711
}
711712

713+
/**
714+
* One page of a `chat.agent` session's persisted transcript, most recent
715+
* messages first when `limit` is set. Secret key only.
716+
*/
717+
getSessionTranscript(
718+
sessionId: string,
719+
options?: { limit?: number; before?: string },
720+
requestOptions?: ZodFetchOptions
721+
) {
722+
const query = new URLSearchParams();
723+
if (options?.limit !== undefined) query.set("limit", String(options.limit));
724+
if (options?.before !== undefined) query.set("before", options.before);
725+
const suffix = query.size > 0 ? `?${query.toString()}` : "";
726+
return zodfetch(
727+
SessionTranscriptResponseBody,
728+
`${this.baseUrl}/api/v1/sessions/${encodeURIComponent(sessionId)}/transcript${suffix}`,
729+
{
730+
method: "GET",
731+
headers: this.#getHeaders(false),
732+
},
733+
mergeRequestOptions(this.defaultRequestOptions, requestOptions)
734+
);
735+
}
736+
712737
retrieveRun(runId: string, requestOptions?: ZodFetchOptions) {
713738
return zodfetch(
714739
RetrieveRunResponse,

packages/core/src/v3/schemas/api.ts

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1005,6 +1005,21 @@ export const CreateUploadPayloadUrlResponseBody = z.object({
10051005
storagePath: z.string().optional(),
10061006
});
10071007

1008+
/** One page of a `chat.agent` session transcript, from `GET /api/v1/sessions/{id}/transcript`. */
1009+
export const SessionTranscriptResponseBody = z.object({
1010+
messages: z.array(z.unknown()),
1011+
state: z.unknown().nullable(),
1012+
cursors: z
1013+
.object({
1014+
lastOutEventId: z.string().optional(),
1015+
lastInEventId: z.string().optional(),
1016+
})
1017+
.optional(),
1018+
nextCursor: z.string().optional(),
1019+
});
1020+
1021+
export type SessionTranscriptResponseBody = z.infer<typeof SessionTranscriptResponseBody>;
1022+
10081023
export const WorkersListResponseBody = z
10091024
.object({
10101025
type: z.string(),

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,29 @@ export function parseTranscriptSnapshot<TUIMessage extends UIMessage = UIMessage
142142
return undefined;
143143
}
144144

145+
/**
146+
* Select one page of transcript entries, newest last. `before` keeps only the
147+
* entries ordered before that id; `limit` keeps the last that many. The
148+
* returned `nextCursor` is the id to pass as `before` for the previous page,
149+
* absent when there is no earlier page.
150+
*/
151+
export function pageTranscriptEntries<TUIMessage extends UIMessage = UIMessage>(
152+
all: TranscriptSnapshotEntry<TUIMessage>[],
153+
opts: { limit?: number; before?: string } | undefined
154+
): { entries: TranscriptSnapshotEntry<TUIMessage>[]; nextCursor: string | undefined } {
155+
let entries = all;
156+
if (opts?.before !== undefined) {
157+
const idx = entries.findIndex((e) => e.id === opts.before);
158+
if (idx !== -1) entries = entries.slice(0, idx);
159+
}
160+
let nextCursor: string | undefined;
161+
if (opts?.limit !== undefined && opts.limit >= 0 && entries.length > opts.limit) {
162+
entries = entries.slice(entries.length - opts.limit);
163+
nextCursor = entries[0]?.id;
164+
}
165+
return { entries, nextCursor };
166+
}
167+
145168
/**
146169
* S3 key suffix for a session's snapshot blob. The webapp's presigned
147170
* URL routes prefix this with `packets/{projectRef}/{envSlug}/`.

packages/core/test/chatSnapshot.test.ts

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,43 @@
11
import { describe, expect, it } from "vitest";
22
import {
3+
pageTranscriptEntries,
34
parseTranscriptSnapshot,
45
type ChatSnapshotV1,
6+
type TranscriptSnapshotEntry,
57
type TranscriptSnapshotV2,
68
} from "../src/v3/sessionStreams/chatSnapshot.js";
79

10+
describe("pageTranscriptEntries", () => {
11+
const entries: TranscriptSnapshotEntry[] = ["a", "b", "c", "d", "e"].map((id) => ({
12+
id,
13+
final: true,
14+
message: { id, role: "user", parts: [] },
15+
}));
16+
17+
it("returns everything with no options and no cursor", () => {
18+
expect(pageTranscriptEntries(entries, undefined)).toEqual({ entries, nextCursor: undefined });
19+
});
20+
21+
it("returns the newest `limit` entries with the cursor for the page before", () => {
22+
const page = pageTranscriptEntries(entries, { limit: 2 });
23+
expect(page.entries.map((e) => e.id)).toEqual(["d", "e"]);
24+
expect(page.nextCursor).toBe("d");
25+
});
26+
27+
it("pages backwards with `before` until the cursor runs out", () => {
28+
const page = pageTranscriptEntries(entries, { limit: 2, before: "d" });
29+
expect(page.entries.map((e) => e.id)).toEqual(["b", "c"]);
30+
expect(page.nextCursor).toBe("b");
31+
const last = pageTranscriptEntries(entries, { limit: 2, before: "b" });
32+
expect(last.entries.map((e) => e.id)).toEqual(["a"]);
33+
expect(last.nextCursor).toBeUndefined();
34+
});
35+
36+
it("ignores an unknown `before` id", () => {
37+
expect(pageTranscriptEntries(entries, { before: "zz" }).entries).toHaveLength(5);
38+
});
39+
});
40+
841
const user = { id: "u-1", role: "user" as const, parts: [{ type: "text" as const, text: "hi" }] };
942
const assistant = {
1043
id: "a-1",

packages/trigger-sdk/src/v3/ai.ts

Lines changed: 125 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ import {
8181
restoreModelLane,
8282
type TranscriptChange,
8383
type TranscriptChangeReason,
84+
type TranscriptLoadResult,
8485
type TranscriptRuntimeState,
8586
type TranscriptShadow,
8687
type TranscriptStorage,
@@ -277,6 +278,25 @@ export {
277278
__writeChatSnapshotProductionPathForTests,
278279
} from "./chatSnapshotIo.js";
279280

281+
export {
282+
defaultStorage,
283+
memoryTranscriptStorage,
284+
reduceTranscriptChanges,
285+
snapshotTranscriptStorage,
286+
type LoadContextEvent,
287+
type MemoryTranscriptStorage,
288+
type TranscriptChange,
289+
type TranscriptChangeReason,
290+
type TranscriptChangeset,
291+
type TranscriptCursors,
292+
type TranscriptLoadOptions,
293+
type TranscriptLoadResult,
294+
type TranscriptScope,
295+
type TranscriptState,
296+
type TranscriptStorage,
297+
type TranscriptStorageContext,
298+
} from "./transcriptStorage.js";
299+
280300
/**
281301
* Merge two `UIMessage[]` lists by `id`, with the second list winning on
282302
* collision. Used at run boot to combine the snapshot's persisted history
@@ -6140,6 +6160,33 @@ export type ChatAgentOptions<
61406160
event: HydrateMessagesEvent<inferSchemaOut<TClientDataSchema>, TUIMessage>
61416161
) => TUIMessage[] | Promise<TUIMessage[]>;
61426162

6163+
/**
6164+
* Where the conversation is persisted. The runtime calls `save` after
6165+
* every turn, failed turn and history-changing action with the changes
6166+
* since the last save, and `load` once when a new run boots to continue
6167+
* the conversation.
6168+
*
6169+
* Defaults to `defaultStorage`, the platform's snapshot in object storage
6170+
* that the Sessions dashboard renders. Bring your own to write each change
6171+
* to your database; `memoryTranscriptStorage()` is the reference
6172+
* implementation and `runTranscriptStorageTests` from
6173+
* `@trigger.dev/sdk/ai/test` checks yours against the contract.
6174+
*
6175+
* A storage with `loadContext` also owns the model's context on every
6176+
* turn, which is what `hydrateMessages` did. The two cannot be combined.
6177+
*
6178+
* @example
6179+
* ```ts
6180+
* chat.agent({
6181+
* id: "my-chat",
6182+
* storage: myPostgresTranscriptStorage,
6183+
* run: async ({ messages, signal, streamText }) =>
6184+
* streamText({ model, messages, abortSignal: signal }),
6185+
* });
6186+
* ```
6187+
*/
6188+
storage?: TranscriptStorage<inferSchemaOut<TClientDataSchema>>;
6189+
61436190
/**
61446191
* Called at the start of every turn, after message accumulation and `onChatStart` (turn 0),
61456192
* but before the `run` function executes.
@@ -6726,6 +6773,7 @@ function chatAgent<
67266773
onChatStart,
67276774
onValidateMessages,
67286775
hydrateMessages,
6776+
storage,
67296777
actionSchema,
67306778
onAction,
67316779
onTurnStart,
@@ -6755,12 +6803,17 @@ function chatAgent<
67556803
} = options;
67566804

67576805
if (hydrateMessages) {
6758-
const storageAtDefinition = transcriptStorageOverride ?? defaultStorage;
6759-
if (typeof storageAtDefinition.loadContext === "function") {
6806+
if (storage) {
6807+
throw new Error(
6808+
`chat.agent: "${options.id}" sets both \`hydrateMessages\` and \`storage\`. ` +
6809+
"`hydrateMessages` is deprecated and replaced by the storage: `save` receives every " +
6810+
"change and `loadContext` on the storage owns the model's context. Remove `hydrateMessages`."
6811+
);
6812+
}
6813+
if (typeof (transcriptStorageOverride ?? defaultStorage).loadContext === "function") {
67606814
throw new Error(
67616815
`chat.agent: "${options.id}" sets \`hydrateMessages\` and uses a transcript storage with ` +
6762-
"`loadContext`. Both would own the model's context; keep one. `hydrateMessages` is " +
6763-
"deprecated, so prefer `loadContext` on the storage."
6816+
"`loadContext`. Both would own the model's context; keep one."
67646817
);
67656818
}
67666819
warnHydrateMessagesDeprecatedOnce(options.id);
@@ -6915,7 +6968,10 @@ function chatAgent<
69156968
// collectively cost ~600ms on every first-message TTFC. Both reads
69166969
// swallow errors internally; the agent stays available either way.
69176970
const sessionIdForSnapshot = payload.sessionId ?? payload.chatId;
6918-
const transcriptStorage = transcriptStorageOverride ?? defaultStorage;
6971+
const transcriptStorage: TranscriptStorage<unknown> =
6972+
(storage as TranscriptStorage<unknown> | undefined) ??
6973+
transcriptStorageOverride ??
6974+
defaultStorage;
69196975
const storageLoadContext = transcriptStorage.loadContext?.bind(transcriptStorage);
69206976
/**
69216977
* Who supplies the model's context each turn: the deprecated
@@ -12604,6 +12660,68 @@ async function mintPublicTokenWithOverride(args: {
1260412660
});
1260512661
}
1260612662

12663+
export type CreateChatLoadTranscriptActionOptions = {
12664+
/**
12665+
* Scope the action to a specific API client configuration (secret key,
12666+
* base URL) instead of the process-wide one. The default storage reads
12667+
* through this client.
12668+
*/
12669+
apiClient?: ApiClientConfiguration;
12670+
/** Page size when the caller passes none. */
12671+
limit?: number;
12672+
};
12673+
12674+
export type ChatLoadTranscriptParams<TClientData = unknown> = {
12675+
chatId: string;
12676+
clientData?: TClientData;
12677+
limit?: number;
12678+
before?: string;
12679+
};
12680+
12681+
/**
12682+
* Creates a server-side helper that reads a conversation from a transcript
12683+
* storage, for rendering history before the chat connects. Works the same
12684+
* for every storage, the platform default included, so the browser never
12685+
* reads a store directly and the secret key stays on the server.
12686+
*
12687+
* Wrap it in a Next.js server action (or any server-side handler), scope it
12688+
* to the authenticated user through `clientData`, and pass the result to
12689+
* `useLoadTranscript` in the browser.
12690+
*
12691+
* @example
12692+
* ```ts
12693+
* // actions.ts
12694+
* "use server";
12695+
* import { chat, defaultStorage } from "@trigger.dev/sdk/ai";
12696+
*
12697+
* export const loadTranscript = chat.createLoadTranscriptAction(defaultStorage, { limit: 50 });
12698+
* ```
12699+
*/
12700+
function createChatLoadTranscriptAction<TClientData = unknown>(
12701+
storage: TranscriptStorage<TClientData>,
12702+
options?: CreateChatLoadTranscriptActionOptions
12703+
): (params: ChatLoadTranscriptParams<TClientData>) => Promise<TranscriptLoadResult> {
12704+
return async (params) => {
12705+
if (!params.chatId) {
12706+
throw new Error("chat.createLoadTranscriptAction: params.chatId is required.");
12707+
}
12708+
if (options?.apiClient) {
12709+
const { apiClient, ...rest } = options;
12710+
return apiClientManager.runWithConfig(apiClient, () =>
12711+
createChatLoadTranscriptAction(storage, rest)(params)
12712+
);
12713+
}
12714+
const limit = params.limit ?? options?.limit;
12715+
return storage.load(
12716+
{ chatId: params.chatId, clientData: params.clientData as TClientData },
12717+
{
12718+
...(limit !== undefined ? { limit } : {}),
12719+
...(params.before !== undefined ? { before: params.before } : {}),
12720+
}
12721+
);
12722+
};
12723+
}
12724+
1260712725
export const chat = {
1260812726
/** Create a chat agent. See {@link chatAgent}. */
1260912727
agent: chatAgent,
@@ -12615,6 +12733,8 @@ export const chat = {
1261512733
withClientData,
1261612734
/** Create a server-side helper for starting (or resuming) a Session for a chatId. See {@link createChatStartSessionAction}. */
1261712735
createStartSessionAction: createChatStartSessionAction,
12736+
/** Returns a server-side helper that reads a conversation from a transcript storage. */
12737+
createLoadTranscriptAction: createChatLoadTranscriptAction,
1261812738
/** Pipe a stream to the chat transport. See {@link pipeChat}. */
1261912739
pipe: pipeChat,
1262012740
/** Return from `onAction` to run a turn on the edited history. See {@link chatTurn}. */

0 commit comments

Comments
 (0)