build full example with persisting messages, adding necessary hooks, and documenting it all
This commit is contained in:
Eric Allam
@@ -379,6 +379,47 @@ async function pipeChat(
|
||||
* emits a control chunk and suspends via `messagesInput.wait()`. The frontend
|
||||
* transport resumes the same run by sending the next message via input streams.
|
||||
*/
|
||||
/**
|
||||
* Event passed to the `onChatStart` callback.
|
||||
*/
|
||||
export type ChatStartEvent = {
|
||||
/** The unique identifier for the chat session. */
|
||||
chatId: string;
|
||||
/** The initial model-ready messages for this conversation. */
|
||||
messages: ModelMessage[];
|
||||
/** Custom data from the frontend (passed via `metadata` on `sendMessage()` or the transport). */
|
||||
clientData: unknown;
|
||||
};
|
||||
|
||||
/**
|
||||
* Event passed to the `onTurnComplete` callback.
|
||||
*/
|
||||
export type TurnCompleteEvent = {
|
||||
/** The unique identifier for the chat session. */
|
||||
chatId: string;
|
||||
/** The full accumulated conversation in model format (all turns so far). */
|
||||
messages: ModelMessage[];
|
||||
/**
|
||||
* The full accumulated conversation in UI format (all turns so far).
|
||||
* This is the format expected by `useChat` — store this for persistence.
|
||||
*/
|
||||
uiMessages: UIMessage[];
|
||||
/**
|
||||
* Only the new model messages from this turn (user message(s) + assistant response).
|
||||
* Useful for appending to an existing conversation record.
|
||||
*/
|
||||
newMessages: ModelMessage[];
|
||||
/**
|
||||
* Only the new UI messages from this turn (user message(s) + assistant response).
|
||||
* Useful for inserting individual message records instead of overwriting the full history.
|
||||
*/
|
||||
newUIMessages: UIMessage[];
|
||||
/** The assistant's response for this turn (undefined if `pipeChat` was used manually). */
|
||||
responseMessage: UIMessage | undefined;
|
||||
/** The turn number (0-indexed). */
|
||||
turn: number;
|
||||
};
|
||||
|
||||
export type ChatTaskOptions<TIdentifier extends string> = Omit<
|
||||
TaskOptions<TIdentifier, ChatTaskWirePayload, unknown>,
|
||||
"run"
|
||||
@@ -394,6 +435,35 @@ export type ChatTaskOptions<TIdentifier extends string> = Omit<
|
||||
*/
|
||||
run: (payload: ChatTaskRunPayload) => Promise<unknown>;
|
||||
|
||||
/**
|
||||
* Called on the first turn (turn 0) of a new run, before the `run` function executes.
|
||||
*
|
||||
* Use this to create the chat record in your database when a new conversation starts.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* onChatStart: async ({ chatId, messages, clientData }) => {
|
||||
* await db.chat.create({ data: { id: chatId, userId: clientData.userId } });
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
onChatStart?: (event: ChatStartEvent) => Promise<void> | void;
|
||||
|
||||
/**
|
||||
* Called after each turn completes (after the response is captured, before waiting
|
||||
* for the next message). Also fires on the final turn.
|
||||
*
|
||||
* Use this to persist the conversation to your database after each assistant response.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* onTurnComplete: async ({ chatId, messages }) => {
|
||||
* await db.chat.update({ where: { id: chatId }, data: { messages } });
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
onTurnComplete?: (event: TurnCompleteEvent) => Promise<void> | void;
|
||||
|
||||
/**
|
||||
* Maximum number of conversational turns (message round-trips) a single run
|
||||
* will handle before ending. After this many turns the run completes
|
||||
@@ -456,6 +526,8 @@ function chatTask<TIdentifier extends string>(
|
||||
): Task<TIdentifier, ChatTaskWirePayload, unknown> {
|
||||
const {
|
||||
run: userRun,
|
||||
onChatStart,
|
||||
onTurnComplete,
|
||||
maxTurns = 100,
|
||||
turnTimeout = "1h",
|
||||
warmTimeoutInSeconds = 30,
|
||||
@@ -478,6 +550,10 @@ function chatTask<TIdentifier extends string>(
|
||||
// user message(s) and the captured assistant response.
|
||||
let accumulatedMessages: ModelMessage[] = [];
|
||||
|
||||
// Accumulated UI messages for persistence. Mirrors the model accumulator
|
||||
// but in frontend-friendly UIMessage format (with parts, id, etc.).
|
||||
let accumulatedUIMessages: UIMessage[] = [];
|
||||
|
||||
// Mutable reference to the current turn's stop controller so the
|
||||
// stop input stream listener (registered once) can abort the right turn.
|
||||
let currentStopController: AbortController | undefined;
|
||||
@@ -549,15 +625,52 @@ function chatTask<TIdentifier extends string>(
|
||||
// Turn 2+: only the new message(s) → appended to the accumulator.
|
||||
const incomingModelMessages = await convertToModelMessages(uiMessages);
|
||||
|
||||
// Track new messages for this turn (user input + assistant response).
|
||||
const turnNewModelMessages: ModelMessage[] = [];
|
||||
const turnNewUIMessages: UIMessage[] = [];
|
||||
|
||||
if (turn === 0) {
|
||||
accumulatedMessages = incomingModelMessages;
|
||||
accumulatedUIMessages = [...uiMessages];
|
||||
// On first turn, the "new" messages are just the last user message
|
||||
// (the rest is history). We'll add the response after streaming.
|
||||
if (uiMessages.length > 0) {
|
||||
turnNewUIMessages.push(uiMessages[uiMessages.length - 1]!);
|
||||
const lastModel = incomingModelMessages[incomingModelMessages.length - 1];
|
||||
if (lastModel) turnNewModelMessages.push(lastModel);
|
||||
}
|
||||
} else if (currentWirePayload.trigger === "regenerate-message") {
|
||||
// Regenerate: frontend sent full history with last assistant message
|
||||
// removed. Reset the accumulator to match.
|
||||
accumulatedMessages = incomingModelMessages;
|
||||
accumulatedUIMessages = [...uiMessages];
|
||||
// No new user messages for regenerate — just the response (added below)
|
||||
} else {
|
||||
// Submit: frontend sent only the new user message(s). Append to accumulator.
|
||||
accumulatedMessages.push(...incomingModelMessages);
|
||||
accumulatedUIMessages.push(...uiMessages);
|
||||
turnNewModelMessages.push(...incomingModelMessages);
|
||||
turnNewUIMessages.push(...uiMessages);
|
||||
}
|
||||
|
||||
// Fire onChatStart on the first turn
|
||||
if (turn === 0 && onChatStart) {
|
||||
await tracer.startActiveSpan(
|
||||
"onChatStart()",
|
||||
async () => {
|
||||
await onChatStart({
|
||||
chatId: currentWirePayload.chatId,
|
||||
messages: accumulatedMessages,
|
||||
clientData: wireMetadata,
|
||||
});
|
||||
},
|
||||
{
|
||||
attributes: {
|
||||
[SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart",
|
||||
[SemanticInternalAttributes.COLLAPSED]: true,
|
||||
},
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
// Captured by the onFinish callback below — works even on abort/stop.
|
||||
@@ -602,11 +715,14 @@ function chatTask<TIdentifier extends string>(
|
||||
// The onFinish callback fires even on abort/stop, so partial responses
|
||||
// from stopped generation are captured correctly.
|
||||
if (capturedResponseMessage) {
|
||||
accumulatedUIMessages.push(capturedResponseMessage);
|
||||
turnNewUIMessages.push(capturedResponseMessage);
|
||||
try {
|
||||
const responseModelMessages = await convertToModelMessages([
|
||||
stripProviderMetadata(capturedResponseMessage),
|
||||
]);
|
||||
accumulatedMessages.push(...responseModelMessages);
|
||||
turnNewModelMessages.push(...responseModelMessages);
|
||||
} catch {
|
||||
// Conversion failed — skip accumulation for this turn
|
||||
}
|
||||
@@ -618,6 +734,30 @@ function chatTask<TIdentifier extends string>(
|
||||
|
||||
if (runSignal.aborted) return "exit";
|
||||
|
||||
// Fire onTurnComplete after response capture
|
||||
if (onTurnComplete) {
|
||||
await tracer.startActiveSpan(
|
||||
"onTurnComplete()",
|
||||
async () => {
|
||||
await onTurnComplete({
|
||||
chatId: currentWirePayload.chatId,
|
||||
messages: accumulatedMessages,
|
||||
uiMessages: accumulatedUIMessages,
|
||||
newMessages: turnNewModelMessages,
|
||||
newUIMessages: turnNewUIMessages,
|
||||
responseMessage: capturedResponseMessage,
|
||||
turn,
|
||||
});
|
||||
},
|
||||
{
|
||||
attributes: {
|
||||
[SemanticInternalAttributes.STYLE_ICON]: "task-hook-onComplete",
|
||||
[SemanticInternalAttributes.COLLAPSED]: true,
|
||||
},
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
// Write turn-complete control chunk so frontend closes its stream
|
||||
await writeTurnCompleteChunk(currentWirePayload.chatId);
|
||||
|
||||
@@ -629,9 +769,12 @@ function chatTask<TIdentifier extends string>(
|
||||
|
||||
// Phase 1: Keep the run warm for quick response to the next message.
|
||||
// The run stays active (using compute) during this window.
|
||||
if (warmTimeoutInSeconds > 0) {
|
||||
const effectiveWarmTimeout =
|
||||
(metadata.get(WARM_TIMEOUT_METADATA_KEY) as number | undefined) ?? warmTimeoutInSeconds;
|
||||
|
||||
if (effectiveWarmTimeout > 0) {
|
||||
const warm = await messagesInput.once({
|
||||
timeoutMs: warmTimeoutInSeconds * 1000,
|
||||
timeoutMs: effectiveWarmTimeout * 1000,
|
||||
spanName: "waiting (warm)",
|
||||
});
|
||||
|
||||
@@ -643,8 +786,11 @@ function chatTask<TIdentifier extends string>(
|
||||
}
|
||||
|
||||
// Phase 2: Suspend the task (frees compute) until the next message arrives
|
||||
const effectiveTurnTimeout =
|
||||
(metadata.get(TURN_TIMEOUT_METADATA_KEY) as string | undefined) ?? turnTimeout;
|
||||
|
||||
const next = await messagesInput.wait({
|
||||
timeout: turnTimeout,
|
||||
timeout: effectiveTurnTimeout,
|
||||
spanName: "waiting (suspended)",
|
||||
});
|
||||
|
||||
@@ -693,6 +839,74 @@ function chatTask<TIdentifier extends string>(
|
||||
* const token = await chat.createAccessToken("my-chat");
|
||||
* ```
|
||||
*/
|
||||
// ---------------------------------------------------------------------------
|
||||
// Runtime configuration helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const TURN_TIMEOUT_METADATA_KEY = "chat.turnTimeout";
|
||||
const WARM_TIMEOUT_METADATA_KEY = "chat.warmTimeout";
|
||||
|
||||
/**
|
||||
* Override the turn timeout for subsequent turns in the current run.
|
||||
*
|
||||
* The turn timeout controls how long the run stays suspended (freeing compute)
|
||||
* waiting for the next user message. When it expires, the run completes
|
||||
* gracefully and the next message starts a fresh run.
|
||||
*
|
||||
* Call from inside a `chatTask` run function to adjust based on context.
|
||||
*
|
||||
* @param duration - A duration string (e.g. `"5m"`, `"1h"`, `"30s"`)
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* run: async ({ messages, signal }) => {
|
||||
* chat.setTurnTimeout("2h");
|
||||
* return streamText({ model, messages, abortSignal: signal });
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
function setTurnTimeout(duration: string): void {
|
||||
metadata.set(TURN_TIMEOUT_METADATA_KEY, duration);
|
||||
}
|
||||
|
||||
/**
|
||||
* Override the turn timeout in seconds for subsequent turns in the current run.
|
||||
*
|
||||
* @param seconds - Number of seconds to wait for the next message before ending the run
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* run: async ({ messages, signal }) => {
|
||||
* chat.setTurnTimeoutInSeconds(3600); // 1 hour
|
||||
* return streamText({ model, messages, abortSignal: signal });
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
function setTurnTimeoutInSeconds(seconds: number): void {
|
||||
metadata.set(TURN_TIMEOUT_METADATA_KEY, `${seconds}s`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Override the warm timeout for subsequent turns in the current run.
|
||||
*
|
||||
* The warm timeout controls how long the run stays active (using compute)
|
||||
* after each turn, waiting for the next message. During this window,
|
||||
* responses are instant. After it expires, the run suspends.
|
||||
*
|
||||
* @param seconds - Number of seconds to stay warm (0 to suspend immediately)
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* run: async ({ messages, signal }) => {
|
||||
* chat.setWarmTimeoutInSeconds(60);
|
||||
* return streamText({ model, messages, abortSignal: signal });
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
function setWarmTimeoutInSeconds(seconds: number): void {
|
||||
metadata.set(WARM_TIMEOUT_METADATA_KEY, seconds);
|
||||
}
|
||||
|
||||
export const chat = {
|
||||
/** Create a chat task. See {@link chatTask}. */
|
||||
task: chatTask,
|
||||
@@ -700,6 +914,12 @@ export const chat = {
|
||||
pipe: pipeChat,
|
||||
/** Create a public access token for a chat task. See {@link createChatAccessToken}. */
|
||||
createAccessToken: createChatAccessToken,
|
||||
/** Override the turn timeout at runtime (duration string). See {@link setTurnTimeout}. */
|
||||
setTurnTimeout,
|
||||
/** Override the turn timeout at runtime (seconds). See {@link setTurnTimeoutInSeconds}. */
|
||||
setTurnTimeoutInSeconds,
|
||||
/** Override the warm timeout at runtime. See {@link setWarmTimeoutInSeconds}. */
|
||||
setWarmTimeoutInSeconds,
|
||||
};
|
||||
|
||||
/**
|
||||
|
||||
@@ -23,7 +23,7 @@
|
||||
* ```
|
||||
*/
|
||||
|
||||
import { useRef } from "react";
|
||||
import { useEffect, useRef } from "react";
|
||||
import {
|
||||
TriggerChatTransport,
|
||||
type TriggerChatTransportOptions,
|
||||
@@ -57,6 +57,9 @@ export type UseTriggerChatTransportOptions<TTask extends AnyTask = AnyTask> = Om
|
||||
* For dynamic access tokens, pass a function — it will be called on each
|
||||
* request without needing to recreate the transport.
|
||||
*
|
||||
* The `onSessionChange` callback is kept in a ref so the transport always
|
||||
* calls the latest version without needing to be recreated.
|
||||
*
|
||||
* @example
|
||||
* ```tsx
|
||||
* import { useChat } from "@ai-sdk/react";
|
||||
@@ -80,5 +83,12 @@ export function useTriggerChatTransport<TTask extends AnyTask = AnyTask>(
|
||||
if (ref.current === null) {
|
||||
ref.current = new TriggerChatTransport(options);
|
||||
}
|
||||
|
||||
// Keep onSessionChange up to date without recreating the transport
|
||||
const { onSessionChange } = options;
|
||||
useEffect(() => {
|
||||
ref.current?.setOnSessionChange(onSessionChange);
|
||||
}, [onSessionChange]);
|
||||
|
||||
return ref.current;
|
||||
}
|
||||
|
||||
@@ -1727,4 +1727,187 @@ describe("TriggerChatTransport", () => {
|
||||
expect(triggerCallCount).toBe(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("onSessionChange", () => {
|
||||
it("should fire when a new session is created", async () => {
|
||||
const onSessionChange = vi.fn();
|
||||
const triggerRunId = "run_session_new";
|
||||
const publicToken = "pub_session_new";
|
||||
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
|
||||
if (urlStr.includes("/trigger")) {
|
||||
return new Response(JSON.stringify({ id: triggerRunId }), {
|
||||
status: 200,
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"x-trigger-jwt": publicToken,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (urlStr.includes("/realtime/v1/streams/")) {
|
||||
const chunks = [
|
||||
...sampleChunks,
|
||||
{ type: "__trigger_turn_complete" },
|
||||
];
|
||||
return new Response(createSSEStream(sseEncode(chunks)), {
|
||||
status: 200,
|
||||
headers: {
|
||||
"content-type": "text/event-stream",
|
||||
"X-Stream-Version": "v1",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
throw new Error(`Unexpected fetch URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-task",
|
||||
accessToken: "token",
|
||||
baseURL: "https://api.test.trigger.dev",
|
||||
onSessionChange,
|
||||
});
|
||||
|
||||
const stream = await transport.sendMessages({
|
||||
trigger: "submit-message",
|
||||
chatId: "chat-1",
|
||||
messageId: undefined,
|
||||
messages: [createUserMessage("Hello")],
|
||||
abortSignal: undefined,
|
||||
});
|
||||
|
||||
// Session created notification should have fired
|
||||
expect(onSessionChange).toHaveBeenCalledWith("chat-1", {
|
||||
runId: triggerRunId,
|
||||
publicAccessToken: publicToken,
|
||||
lastEventId: undefined,
|
||||
});
|
||||
|
||||
// Consume stream
|
||||
const reader = stream.getReader();
|
||||
while (!(await reader.read()).done) {}
|
||||
|
||||
// Should also fire with updated lastEventId on turn complete
|
||||
const lastCall = onSessionChange.mock.calls[onSessionChange.mock.calls.length - 1]!;
|
||||
expect(lastCall![0]).toBe("chat-1");
|
||||
expect(lastCall![1]).not.toBeNull();
|
||||
expect(lastCall![1].lastEventId).toBeDefined();
|
||||
});
|
||||
|
||||
it("should fire with null when session is deleted (stream ends naturally)", async () => {
|
||||
const onSessionChange = vi.fn();
|
||||
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
|
||||
if (urlStr.includes("/trigger")) {
|
||||
return new Response(JSON.stringify({ id: "run_end" }), {
|
||||
status: 200,
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"x-trigger-jwt": "pub_end",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (urlStr.includes("/realtime/v1/streams/")) {
|
||||
// No turn-complete chunk — stream ends naturally (run completed)
|
||||
return new Response(createSSEStream(sseEncode(sampleChunks)), {
|
||||
status: 200,
|
||||
headers: {
|
||||
"content-type": "text/event-stream",
|
||||
"X-Stream-Version": "v1",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
throw new Error(`Unexpected fetch URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-task",
|
||||
accessToken: "token",
|
||||
baseURL: "https://api.test.trigger.dev",
|
||||
onSessionChange,
|
||||
});
|
||||
|
||||
const stream = await transport.sendMessages({
|
||||
trigger: "submit-message",
|
||||
chatId: "chat-end",
|
||||
messageId: undefined,
|
||||
messages: [createUserMessage("Hello")],
|
||||
abortSignal: undefined,
|
||||
});
|
||||
|
||||
// Consume the stream fully
|
||||
const reader = stream.getReader();
|
||||
while (!(await reader.read()).done) {}
|
||||
|
||||
// Session should have been created then deleted
|
||||
expect(onSessionChange).toHaveBeenCalledWith("chat-end", expect.objectContaining({
|
||||
runId: "run_end",
|
||||
}));
|
||||
expect(onSessionChange).toHaveBeenCalledWith("chat-end", null);
|
||||
});
|
||||
|
||||
it("should be updatable via setOnSessionChange", async () => {
|
||||
const onSessionChange1 = vi.fn();
|
||||
const onSessionChange2 = vi.fn();
|
||||
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
|
||||
if (urlStr.includes("/trigger")) {
|
||||
return new Response(JSON.stringify({ id: "run_update" }), {
|
||||
status: 200,
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"x-trigger-jwt": "pub_update",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (urlStr.includes("/realtime/v1/streams/")) {
|
||||
const chunks = [
|
||||
...sampleChunks,
|
||||
{ type: "__trigger_turn_complete" },
|
||||
];
|
||||
return new Response(createSSEStream(sseEncode(chunks)), {
|
||||
status: 200,
|
||||
headers: {
|
||||
"content-type": "text/event-stream",
|
||||
"X-Stream-Version": "v1",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
throw new Error(`Unexpected fetch URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-task",
|
||||
accessToken: "token",
|
||||
baseURL: "https://api.test.trigger.dev",
|
||||
onSessionChange: onSessionChange1,
|
||||
});
|
||||
|
||||
// Update the callback before sending
|
||||
transport.setOnSessionChange(onSessionChange2);
|
||||
|
||||
await transport.sendMessages({
|
||||
trigger: "submit-message",
|
||||
chatId: "chat-update",
|
||||
messageId: undefined,
|
||||
messages: [createUserMessage("Hello")],
|
||||
abortSignal: undefined,
|
||||
});
|
||||
|
||||
// Only onSessionChange2 should have been called
|
||||
expect(onSessionChange1).not.toHaveBeenCalled();
|
||||
expect(onSessionChange2).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -100,6 +100,57 @@ export type TriggerChatTransportOptions = {
|
||||
* ```
|
||||
*/
|
||||
metadata?: Record<string, unknown>;
|
||||
|
||||
/**
|
||||
* Restore active chat sessions from external storage (e.g. localStorage).
|
||||
*
|
||||
* After a page refresh, pass previously persisted sessions here so the
|
||||
* transport can reconnect to existing runs instead of starting new ones.
|
||||
* Use `getSession()` to retrieve session state for persistence.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* new TriggerChatTransport({
|
||||
* task: "my-chat",
|
||||
* accessToken,
|
||||
* sessions: {
|
||||
* "chat-abc": { runId: "run_123", publicAccessToken: "...", lastEventId: "42" },
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
*/
|
||||
sessions?: Record<string, { runId: string; publicAccessToken: string; lastEventId?: string }>;
|
||||
|
||||
/**
|
||||
* Called whenever a chat session's state changes.
|
||||
*
|
||||
* Fires when:
|
||||
* - A new session is created (after triggering a task)
|
||||
* - A turn completes (lastEventId updated)
|
||||
* - A session is removed (run ended or input stream send failed) — `session` will be `null`
|
||||
*
|
||||
* Use this to persist session state for reconnection after page refreshes,
|
||||
* without needing to call `getSession()` manually.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* new TriggerChatTransport({
|
||||
* task: "my-chat",
|
||||
* accessToken,
|
||||
* onSessionChange: (chatId, session) => {
|
||||
* if (session) {
|
||||
* localStorage.setItem(`session:${chatId}`, JSON.stringify(session));
|
||||
* } else {
|
||||
* localStorage.removeItem(`session:${chatId}`);
|
||||
* }
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
*/
|
||||
onSessionChange?: (
|
||||
chatId: string,
|
||||
session: { runId: string; publicAccessToken: string; lastEventId?: string } | null
|
||||
) => void;
|
||||
};
|
||||
|
||||
/**
|
||||
@@ -151,6 +202,12 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
private readonly extraHeaders: Record<string, string>;
|
||||
private readonly streamTimeoutSeconds: number;
|
||||
private readonly defaultMetadata: Record<string, unknown> | undefined;
|
||||
private _onSessionChange:
|
||||
| ((
|
||||
chatId: string,
|
||||
session: { runId: string; publicAccessToken: string; lastEventId?: string } | null
|
||||
) => void)
|
||||
| undefined;
|
||||
|
||||
private sessions: Map<string, ChatSessionState> = new Map();
|
||||
|
||||
@@ -165,6 +222,18 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
this.extraHeaders = options.headers ?? {};
|
||||
this.streamTimeoutSeconds = options.streamTimeoutSeconds ?? DEFAULT_STREAM_TIMEOUT_SECONDS;
|
||||
this.defaultMetadata = options.metadata;
|
||||
this._onSessionChange = options.onSessionChange;
|
||||
|
||||
// Restore sessions from external storage
|
||||
if (options.sessions) {
|
||||
for (const [chatId, session] of Object.entries(options.sessions)) {
|
||||
this.sessions.set(chatId, {
|
||||
runId: session.runId,
|
||||
publicAccessToken: session.publicAccessToken,
|
||||
lastEventId: session.lastEventId,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sendMessages = async (
|
||||
@@ -216,6 +285,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
} catch {
|
||||
// If sending fails (run died, etc.), fall through to trigger a new run.
|
||||
this.sessions.delete(chatId);
|
||||
this.notifySessionChange(chatId, null);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -236,10 +306,12 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
? (triggerResponse as { publicAccessToken?: string }).publicAccessToken
|
||||
: undefined;
|
||||
|
||||
this.sessions.set(chatId, {
|
||||
const newSession: ChatSessionState = {
|
||||
runId,
|
||||
publicAccessToken: publicAccessToken ?? currentToken,
|
||||
});
|
||||
};
|
||||
this.sessions.set(chatId, newSession);
|
||||
this.notifySessionChange(chatId, newSession);
|
||||
return this.subscribeToStream(
|
||||
runId,
|
||||
publicAccessToken ?? currentToken,
|
||||
@@ -261,6 +333,62 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
return this.subscribeToStream(session.runId, session.publicAccessToken, undefined, options.chatId);
|
||||
};
|
||||
|
||||
/**
|
||||
* Get the current session state for a chat, suitable for external persistence.
|
||||
*
|
||||
* Returns `undefined` if no active session exists for this chatId.
|
||||
* Persist the returned value to localStorage so it can be restored
|
||||
* after a page refresh via `restoreSession()`.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* const session = transport.getSession(chatId);
|
||||
* if (session) {
|
||||
* localStorage.setItem(`session:${chatId}`, JSON.stringify(session));
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
getSession = (chatId: string): { runId: string; publicAccessToken: string; lastEventId?: string } | undefined => {
|
||||
const session = this.sessions.get(chatId);
|
||||
if (!session) return undefined;
|
||||
return {
|
||||
runId: session.runId,
|
||||
publicAccessToken: session.publicAccessToken,
|
||||
lastEventId: session.lastEventId,
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Update the `onSessionChange` callback.
|
||||
* Useful for React hooks that need to update the callback without recreating the transport.
|
||||
*/
|
||||
setOnSessionChange(
|
||||
callback:
|
||||
| ((
|
||||
chatId: string,
|
||||
session: { runId: string; publicAccessToken: string; lastEventId?: string } | null
|
||||
) => void)
|
||||
| undefined
|
||||
): void {
|
||||
this._onSessionChange = callback;
|
||||
}
|
||||
|
||||
private notifySessionChange(
|
||||
chatId: string,
|
||||
session: ChatSessionState | null
|
||||
): void {
|
||||
if (!this._onSessionChange) return;
|
||||
if (session) {
|
||||
this._onSessionChange(chatId, {
|
||||
runId: session.runId,
|
||||
publicAccessToken: session.publicAccessToken,
|
||||
lastEventId: session.lastEventId,
|
||||
});
|
||||
} else {
|
||||
this._onSessionChange(chatId, null);
|
||||
}
|
||||
}
|
||||
|
||||
private subscribeToStream(
|
||||
runId: string,
|
||||
accessToken: string,
|
||||
@@ -331,6 +459,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
// the next message via input streams.
|
||||
if (chatId && !combinedSignal.aborted) {
|
||||
this.sessions.delete(chatId);
|
||||
this.notifySessionChange(chatId, null);
|
||||
}
|
||||
controller.close();
|
||||
return;
|
||||
@@ -348,6 +477,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
session.lastEventId = value.id;
|
||||
}
|
||||
|
||||
|
||||
// Guard against heartbeat or malformed SSE events
|
||||
if (value.chunk != null && typeof value.chunk === "object") {
|
||||
const chunk = value.chunk as Record<string, unknown>;
|
||||
@@ -363,6 +493,10 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
}
|
||||
|
||||
if (chunk.type === "__trigger_turn_complete" && chatId) {
|
||||
// Notify with updated lastEventId before closing
|
||||
if (session) {
|
||||
this.notifySessionChange(chatId, session);
|
||||
}
|
||||
internalAbort.abort();
|
||||
try {
|
||||
controller.close();
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
lib/generated/
|
||||
@@ -6,15 +6,22 @@
|
||||
"dev": "next dev --turbopack",
|
||||
"build": "next build",
|
||||
"start": "next start",
|
||||
"dev:trigger": "trigger dev"
|
||||
"dev:trigger": "trigger dev",
|
||||
"db:migrate": "prisma migrate dev",
|
||||
"db:push": "prisma db push",
|
||||
"db:generate": "prisma generate",
|
||||
"postinstall": "prisma generate"
|
||||
},
|
||||
"dependencies": {
|
||||
"@ai-sdk/anthropic": "^3.0.0",
|
||||
"@ai-sdk/openai": "^3.0.0",
|
||||
"@ai-sdk/react": "^3.0.0",
|
||||
"@prisma/adapter-pg": "^7.4.2",
|
||||
"@prisma/client": "^7.4.2",
|
||||
"@trigger.dev/sdk": "workspace:*",
|
||||
"ai": "^6.0.0",
|
||||
"next": "15.3.3",
|
||||
"pg": "^8.16.3",
|
||||
"react": "^19.0.0",
|
||||
"react-dom": "^19.0.0",
|
||||
"streamdown": "^2.3.0",
|
||||
@@ -27,6 +34,7 @@
|
||||
"@types/react": "^19",
|
||||
"@types/react-dom": "^19",
|
||||
"tailwindcss": "^4",
|
||||
"prisma": "^7.4.2",
|
||||
"trigger.dev": "workspace:*",
|
||||
"typescript": "^5"
|
||||
}
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
import "dotenv/config";
|
||||
import { defineConfig, env } from "prisma/config";
|
||||
|
||||
export default defineConfig({
|
||||
schema: "prisma/schema.prisma",
|
||||
migrations: {
|
||||
path: "prisma/migrations",
|
||||
},
|
||||
datasource: {
|
||||
url: env("DATABASE_URL"),
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,20 @@
|
||||
-- CreateTable
|
||||
CREATE TABLE "Chat" (
|
||||
"id" TEXT NOT NULL,
|
||||
"title" TEXT NOT NULL,
|
||||
"messages" JSONB NOT NULL DEFAULT '[]',
|
||||
"createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updatedAt" TIMESTAMP(3) NOT NULL,
|
||||
|
||||
CONSTRAINT "Chat_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "ChatSession" (
|
||||
"id" TEXT NOT NULL,
|
||||
"runId" TEXT NOT NULL,
|
||||
"publicAccessToken" TEXT NOT NULL,
|
||||
"lastEventId" TEXT,
|
||||
|
||||
CONSTRAINT "ChatSession_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
@@ -0,0 +1,3 @@
|
||||
# Please do not edit this file manually
|
||||
# It should be added in your version-control system (e.g., Git)
|
||||
provider = "postgresql"
|
||||
@@ -0,0 +1,23 @@
|
||||
generator client {
|
||||
provider = "prisma-client"
|
||||
output = "../lib/generated/prisma"
|
||||
}
|
||||
|
||||
datasource db {
|
||||
provider = "postgresql"
|
||||
}
|
||||
|
||||
model Chat {
|
||||
id String @id
|
||||
title String
|
||||
messages Json @default("[]")
|
||||
createdAt DateTime @default(now())
|
||||
updatedAt DateTime @updatedAt
|
||||
}
|
||||
|
||||
model ChatSession {
|
||||
id String @id // chatId
|
||||
runId String
|
||||
publicAccessToken String
|
||||
lastEventId String?
|
||||
}
|
||||
@@ -2,5 +2,79 @@
|
||||
|
||||
import { chat } from "@trigger.dev/sdk/ai";
|
||||
import type { aiChat } from "@/trigger/chat";
|
||||
import { prisma } from "@/lib/prisma";
|
||||
|
||||
export const getChatToken = async () => chat.createAccessToken<typeof aiChat>("ai-chat");
|
||||
|
||||
export async function getChatList() {
|
||||
const chats = await prisma.chat.findMany({
|
||||
select: { id: true, title: true, createdAt: true, updatedAt: true },
|
||||
orderBy: { updatedAt: "desc" },
|
||||
});
|
||||
return chats.map((c) => ({
|
||||
id: c.id,
|
||||
title: c.title,
|
||||
createdAt: c.createdAt.getTime(),
|
||||
updatedAt: c.updatedAt.getTime(),
|
||||
}));
|
||||
}
|
||||
|
||||
export async function getChatMessages(chatId: string) {
|
||||
const found = await prisma.chat.findUnique({ where: { id: chatId } });
|
||||
if (!found) return [];
|
||||
return found.messages as any[];
|
||||
}
|
||||
|
||||
export async function saveChatMessages(chatId: string, messages: unknown[]) {
|
||||
await prisma.chat.update({
|
||||
where: { id: chatId },
|
||||
data: { messages: messages as any },
|
||||
}).catch(() => {});
|
||||
}
|
||||
|
||||
export async function deleteChat(chatId: string) {
|
||||
await prisma.chat.delete({ where: { id: chatId } }).catch(() => {});
|
||||
await prisma.chatSession.delete({ where: { id: chatId } }).catch(() => {});
|
||||
}
|
||||
|
||||
export async function updateChatTitle(chatId: string, title: string) {
|
||||
await prisma.chat.update({ where: { id: chatId }, data: { title } }).catch(() => {});
|
||||
}
|
||||
|
||||
export async function saveSessionAction(
|
||||
chatId: string,
|
||||
session: { runId: string; publicAccessToken: string; lastEventId?: string }
|
||||
) {
|
||||
await prisma.chatSession.upsert({
|
||||
where: { id: chatId },
|
||||
create: {
|
||||
id: chatId,
|
||||
runId: session.runId,
|
||||
publicAccessToken: session.publicAccessToken,
|
||||
lastEventId: session.lastEventId,
|
||||
},
|
||||
update: {
|
||||
runId: session.runId,
|
||||
publicAccessToken: session.publicAccessToken,
|
||||
lastEventId: session.lastEventId,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
export async function deleteSessionAction(chatId: string) {
|
||||
await prisma.chatSession.delete({ where: { id: chatId } }).catch(() => {});
|
||||
}
|
||||
|
||||
export async function getAllSessions() {
|
||||
const sessions = await prisma.chatSession.findMany();
|
||||
const result: Record<string, { runId: string; publicAccessToken: string; lastEventId?: string }> =
|
||||
{};
|
||||
for (const s of sessions) {
|
||||
result[s.id] = {
|
||||
runId: s.runId,
|
||||
publicAccessToken: s.publicAccessToken,
|
||||
lastEventId: s.lastEventId ?? undefined,
|
||||
};
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
@@ -1,9 +1,58 @@
|
||||
import { Chat } from "@/components/chat";
|
||||
"use client";
|
||||
|
||||
import type { UIMessage } from "ai";
|
||||
import { useEffect, useState } from "react";
|
||||
import { ChatApp } from "@/components/chat-app";
|
||||
import {
|
||||
getChatList,
|
||||
getChatMessages,
|
||||
getAllSessions,
|
||||
} from "@/app/actions";
|
||||
|
||||
type ChatMeta = {
|
||||
id: string;
|
||||
title: string;
|
||||
createdAt: number;
|
||||
updatedAt: number;
|
||||
};
|
||||
|
||||
export default function Home() {
|
||||
const [chatList, setChatList] = useState<ChatMeta[]>([]);
|
||||
const [activeChatId, setActiveChatId] = useState<string | null>(null);
|
||||
const [initialMessages, setInitialMessages] = useState<UIMessage[]>([]);
|
||||
const [initialSessions, setInitialSessions] = useState<
|
||||
Record<string, { runId: string; publicAccessToken: string; lastEventId?: string }>
|
||||
>({});
|
||||
const [loaded, setLoaded] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
async function load() {
|
||||
const [list, sessions] = await Promise.all([getChatList(), getAllSessions()]);
|
||||
setChatList(list);
|
||||
setInitialSessions(sessions);
|
||||
|
||||
let firstChatId: string | null = null;
|
||||
let firstMessages: UIMessage[] = [];
|
||||
if (list.length > 0) {
|
||||
firstChatId = list[0]!.id;
|
||||
firstMessages = await getChatMessages(firstChatId);
|
||||
}
|
||||
|
||||
setActiveChatId(firstChatId);
|
||||
setInitialMessages(firstMessages);
|
||||
setLoaded(true);
|
||||
}
|
||||
load();
|
||||
}, []);
|
||||
|
||||
if (!loaded) return null;
|
||||
|
||||
return (
|
||||
<main className="flex h-screen flex-col">
|
||||
<Chat />
|
||||
</main>
|
||||
<ChatApp
|
||||
initialChatList={chatList}
|
||||
initialActiveChatId={activeChatId}
|
||||
initialMessages={initialMessages}
|
||||
initialSessions={initialSessions}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
"use client";
|
||||
|
||||
import type { UIMessage } from "ai";
|
||||
import { generateId } from "ai";
|
||||
import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react";
|
||||
import { useCallback, useEffect, useState } from "react";
|
||||
import { Chat } from "@/components/chat";
|
||||
import { ChatSidebar } from "@/components/chat-sidebar";
|
||||
import {
|
||||
getChatToken,
|
||||
getChatList,
|
||||
getChatMessages,
|
||||
deleteChat as deleteChatAction,
|
||||
updateChatTitle,
|
||||
saveSessionAction,
|
||||
deleteSessionAction,
|
||||
saveChatMessages,
|
||||
} from "@/app/actions";
|
||||
|
||||
type ChatMeta = {
|
||||
id: string;
|
||||
title: string;
|
||||
createdAt: number;
|
||||
updatedAt: number;
|
||||
};
|
||||
|
||||
type ChatAppProps = {
|
||||
initialChatList: ChatMeta[];
|
||||
initialActiveChatId: string | null;
|
||||
initialMessages: UIMessage[];
|
||||
initialSessions: Record<
|
||||
string,
|
||||
{ runId: string; publicAccessToken: string; lastEventId?: string }
|
||||
>;
|
||||
};
|
||||
|
||||
export function ChatApp({
|
||||
initialChatList,
|
||||
initialActiveChatId,
|
||||
initialMessages,
|
||||
initialSessions,
|
||||
}: ChatAppProps) {
|
||||
const [chatList, setChatList] = useState<ChatMeta[]>(initialChatList);
|
||||
const [activeChatId, setActiveChatId] = useState<string | null>(initialActiveChatId);
|
||||
const [messages, setMessages] = useState<UIMessage[]>(initialMessages);
|
||||
|
||||
const handleSessionChange = useCallback(
|
||||
(
|
||||
chatId: string,
|
||||
session: { runId: string; publicAccessToken: string; lastEventId?: string } | null
|
||||
) => {
|
||||
if (session) {
|
||||
saveSessionAction(chatId, session);
|
||||
} else {
|
||||
deleteSessionAction(chatId);
|
||||
}
|
||||
},
|
||||
[]
|
||||
);
|
||||
|
||||
const transport = useTriggerChatTransport({
|
||||
task: "ai-chat",
|
||||
accessToken: getChatToken,
|
||||
baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL,
|
||||
sessions: initialSessions,
|
||||
onSessionChange: handleSessionChange,
|
||||
});
|
||||
|
||||
// Load messages when active chat changes
|
||||
useEffect(() => {
|
||||
if (!activeChatId) {
|
||||
setMessages([]);
|
||||
return;
|
||||
}
|
||||
// Don't reload if we already have the initial messages for the initial chat
|
||||
if (activeChatId === initialActiveChatId && messages === initialMessages) {
|
||||
return;
|
||||
}
|
||||
getChatMessages(activeChatId).then(setMessages);
|
||||
}, [activeChatId]);
|
||||
|
||||
function handleNewChat() {
|
||||
const id = generateId();
|
||||
setActiveChatId(id);
|
||||
setMessages([]);
|
||||
}
|
||||
|
||||
function handleSelectChat(id: string) {
|
||||
setActiveChatId(id);
|
||||
}
|
||||
|
||||
async function handleDeleteChat(id: string) {
|
||||
await deleteChatAction(id);
|
||||
const list = await getChatList();
|
||||
setChatList(list);
|
||||
if (activeChatId === id) {
|
||||
if (list.length > 0) {
|
||||
setActiveChatId(list[0]!.id);
|
||||
} else {
|
||||
setActiveChatId(null);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const handleFirstMessage = useCallback(async (chatId: string, text: string) => {
|
||||
const title = text.slice(0, 40).trim() || "New chat";
|
||||
await updateChatTitle(chatId, title);
|
||||
const list = await getChatList();
|
||||
setChatList(list);
|
||||
}, []);
|
||||
|
||||
const handleMessagesChange = useCallback(async (_chatId: string, _messages: UIMessage[]) => {
|
||||
// Messages are persisted server-side via onTurnComplete.
|
||||
// Refresh the chat list to update timestamps.
|
||||
const list = await getChatList();
|
||||
setChatList(list);
|
||||
}, []);
|
||||
|
||||
return (
|
||||
<main className="flex h-screen">
|
||||
<ChatSidebar
|
||||
chats={chatList}
|
||||
activeChatId={activeChatId}
|
||||
onSelectChat={handleSelectChat}
|
||||
onNewChat={handleNewChat}
|
||||
onDeleteChat={handleDeleteChat}
|
||||
/>
|
||||
<div className="flex-1">
|
||||
{activeChatId ? (
|
||||
<Chat
|
||||
key={activeChatId}
|
||||
chatId={activeChatId}
|
||||
initialMessages={messages}
|
||||
transport={transport}
|
||||
onFirstMessage={handleFirstMessage}
|
||||
onMessagesChange={handleMessagesChange}
|
||||
/>
|
||||
) : (
|
||||
<div className="flex h-full items-center justify-center">
|
||||
<div className="text-center">
|
||||
<p className="text-sm text-gray-400">No conversation selected</p>
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleNewChat}
|
||||
className="mt-3 rounded-lg bg-blue-600 px-4 py-2 text-sm font-medium text-white hover:bg-blue-700"
|
||||
>
|
||||
Start a new chat
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</main>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
"use client";
|
||||
|
||||
type ChatMeta = {
|
||||
id: string;
|
||||
title: string;
|
||||
createdAt: number;
|
||||
updatedAt: number;
|
||||
};
|
||||
|
||||
function timeAgo(ts: number): string {
|
||||
const seconds = Math.floor((Date.now() - ts) / 1000);
|
||||
if (seconds < 60) return "just now";
|
||||
const minutes = Math.floor(seconds / 60);
|
||||
if (minutes < 60) return `${minutes}m ago`;
|
||||
const hours = Math.floor(minutes / 60);
|
||||
if (hours < 24) return `${hours}h ago`;
|
||||
const days = Math.floor(hours / 24);
|
||||
return `${days}d ago`;
|
||||
}
|
||||
|
||||
type ChatSidebarProps = {
|
||||
chats: ChatMeta[];
|
||||
activeChatId: string | null;
|
||||
onSelectChat: (id: string) => void;
|
||||
onNewChat: () => void;
|
||||
onDeleteChat: (id: string) => void;
|
||||
};
|
||||
|
||||
export function ChatSidebar({
|
||||
chats,
|
||||
activeChatId,
|
||||
onSelectChat,
|
||||
onNewChat,
|
||||
onDeleteChat,
|
||||
}: ChatSidebarProps) {
|
||||
const sorted = [...chats].sort((a, b) => b.updatedAt - a.updatedAt);
|
||||
|
||||
return (
|
||||
<div className="flex h-full w-64 shrink-0 flex-col border-r border-gray-200 bg-gray-50">
|
||||
<div className="p-3">
|
||||
<button
|
||||
type="button"
|
||||
onClick={onNewChat}
|
||||
className="w-full rounded-lg border border-gray-300 bg-white px-3 py-2 text-sm font-medium text-gray-700 hover:bg-gray-100"
|
||||
>
|
||||
+ New Chat
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<div className="flex-1 overflow-y-auto">
|
||||
{sorted.length === 0 && (
|
||||
<p className="px-3 py-8 text-center text-xs text-gray-400">No conversations yet</p>
|
||||
)}
|
||||
|
||||
{sorted.map((chat) => (
|
||||
<button
|
||||
key={chat.id}
|
||||
type="button"
|
||||
onClick={() => onSelectChat(chat.id)}
|
||||
className={`group flex w-full items-start gap-2 px-3 py-2.5 text-left text-sm hover:bg-gray-100 ${
|
||||
activeChatId === chat.id ? "bg-white" : ""
|
||||
}`}
|
||||
>
|
||||
<div className="min-w-0 flex-1">
|
||||
<div className="truncate font-medium text-gray-800">{chat.title}</div>
|
||||
<div className="text-[10px] text-gray-400">{timeAgo(chat.updatedAt)}</div>
|
||||
</div>
|
||||
<span
|
||||
onClick={(e) => {
|
||||
e.stopPropagation();
|
||||
onDeleteChat(chat.id);
|
||||
}}
|
||||
className="mt-0.5 hidden shrink-0 rounded p-0.5 text-xs text-gray-400 hover:bg-red-100 hover:text-red-600 group-hover:inline-block"
|
||||
>
|
||||
×
|
||||
</span>
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -1,12 +1,11 @@
|
||||
"use client";
|
||||
|
||||
import type { UIMessage } from "ai";
|
||||
import { useChat } from "@ai-sdk/react";
|
||||
import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react";
|
||||
import type { TriggerChatTransport } from "@trigger.dev/sdk/chat";
|
||||
import { useEffect, useRef, useState } from "react";
|
||||
import { Streamdown } from "streamdown";
|
||||
import { getChatToken } from "@/app/actions";
|
||||
import { MODEL_OPTIONS, DEFAULT_MODEL } from "@/trigger/chat";
|
||||
import type { aiChat } from "@/trigger/chat";
|
||||
import { MODEL_OPTIONS, DEFAULT_MODEL } from "@/lib/models";
|
||||
|
||||
function ToolInvocation({ part }: { part: any }) {
|
||||
const [expanded, setExpanded] = useState(false);
|
||||
@@ -71,41 +70,73 @@ function ToolInvocation({ part }: { part: any }) {
|
||||
);
|
||||
}
|
||||
|
||||
export function Chat() {
|
||||
type ChatProps = {
|
||||
chatId: string;
|
||||
initialMessages: UIMessage[];
|
||||
transport: TriggerChatTransport;
|
||||
onFirstMessage?: (chatId: string, text: string) => void;
|
||||
onMessagesChange?: (chatId: string, messages: UIMessage[]) => void;
|
||||
};
|
||||
|
||||
export function Chat({
|
||||
chatId,
|
||||
initialMessages,
|
||||
transport,
|
||||
onFirstMessage,
|
||||
onMessagesChange,
|
||||
}: ChatProps) {
|
||||
const [input, setInput] = useState("");
|
||||
const [model, setModel] = useState(DEFAULT_MODEL);
|
||||
// Track which model was used for each assistant message (keyed by the preceding user message ID)
|
||||
const modelByUserMsgId = useRef<Map<string, string>>(new Map());
|
||||
|
||||
const transport = useTriggerChatTransport<typeof aiChat>({
|
||||
task: "ai-chat",
|
||||
accessToken: getChatToken,
|
||||
baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL,
|
||||
});
|
||||
const hasCalledFirstMessage = useRef(false);
|
||||
|
||||
const { messages, sendMessage, stop, status, error } = useChat({
|
||||
id: chatId,
|
||||
messages: initialMessages,
|
||||
transport,
|
||||
});
|
||||
|
||||
// Notify parent of first user message (for chat metadata creation)
|
||||
useEffect(() => {
|
||||
if (hasCalledFirstMessage.current) return;
|
||||
const firstUser = messages.find((m) => m.role === "user");
|
||||
if (firstUser) {
|
||||
hasCalledFirstMessage.current = true;
|
||||
const text = firstUser.parts
|
||||
.filter((p: any) => p.type === "text")
|
||||
.map((p: any) => p.text)
|
||||
.join(" ");
|
||||
onFirstMessage?.(chatId, text);
|
||||
}
|
||||
}, [messages, chatId, onFirstMessage]);
|
||||
|
||||
// Pending message to send after the current turn completes
|
||||
const [pendingMessage, setPendingMessage] = useState<{ text: string; model: string } | null>(null);
|
||||
|
||||
// Auto-send the pending message when the turn completes
|
||||
// Handle turn completion: persist messages and auto-send pending message
|
||||
const prevStatus = useRef(status);
|
||||
useEffect(() => {
|
||||
if (prevStatus.current === "streaming" && status === "ready" && pendingMessage) {
|
||||
const turnCompleted = prevStatus.current === "streaming" && status === "ready";
|
||||
prevStatus.current = status;
|
||||
|
||||
if (!turnCompleted) return;
|
||||
|
||||
// Persist messages when a turn completes — this ensures the final assistant
|
||||
// message content is saved (not the empty placeholder from mid-stream).
|
||||
if (messages.length > 0) {
|
||||
onMessagesChange?.(chatId, messages);
|
||||
}
|
||||
|
||||
// Auto-send the pending message
|
||||
if (pendingMessage) {
|
||||
const { text, model: pendingMsgModel } = pendingMessage;
|
||||
setPendingMessage(null);
|
||||
pendingModel.current = pendingMsgModel;
|
||||
sendMessage({ text }, { metadata: { model: pendingMsgModel } });
|
||||
}
|
||||
prevStatus.current = status;
|
||||
}, [status, sendMessage, pendingMessage]);
|
||||
}, [status, messages, chatId, onMessagesChange, sendMessage, pendingMessage]);
|
||||
|
||||
// Build a map of assistant message index -> model used
|
||||
// Each assistant message follows a user message, so we track by position
|
||||
function getModelForAssistantAt(index: number): string | undefined {
|
||||
// Walk backwards to find the preceding user message
|
||||
for (let i = index - 1; i >= 0; i--) {
|
||||
if (messages[i]?.role === "user") {
|
||||
return modelByUserMsgId.current.get(messages[i].id);
|
||||
@@ -114,16 +145,13 @@ export function Chat() {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
// When sending, record which model is selected for this user message
|
||||
const originalSendMessage = sendMessage;
|
||||
function trackedSendMessage(msg: Parameters<typeof sendMessage>[0], opts?: Parameters<typeof sendMessage>[1]) {
|
||||
// We'll track it after the message appears — use a ref to store the pending model
|
||||
pendingModel.current = model;
|
||||
originalSendMessage(msg, opts);
|
||||
}
|
||||
const pendingModel = useRef<string>(model);
|
||||
|
||||
// Track model for new user messages as they appear
|
||||
const trackedUserIds = useRef<Set<string>>(new Set());
|
||||
for (const msg of messages) {
|
||||
if (msg.role === "user" && !trackedUserIds.current.has(msg.id)) {
|
||||
@@ -146,7 +174,6 @@ export function Chat() {
|
||||
className={`flex ${message.role === "user" ? "justify-end" : "justify-start"}`}
|
||||
>
|
||||
<div className={`max-w-[80%] ${message.role === "user" ? "" : "w-full"}`}>
|
||||
{/* Model badge for assistant messages */}
|
||||
{message.role === "assistant" && (
|
||||
<div className="mb-1 flex items-center gap-2 text-[10px] text-gray-400">
|
||||
<span className="rounded bg-gray-200 px-1.5 py-0.5 font-medium text-gray-500">
|
||||
@@ -199,7 +226,6 @@ export function Chat() {
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Queued message indicator */}
|
||||
{pendingMessage && (
|
||||
<div className="flex justify-end">
|
||||
<div className="max-w-[80%]">
|
||||
@@ -214,20 +240,17 @@ export function Chat() {
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Error */}
|
||||
{error && (
|
||||
<div className="shrink-0 border-t border-red-100 bg-red-50 px-4 py-2 text-sm text-red-600">
|
||||
{error.message}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Input */}
|
||||
<form
|
||||
onSubmit={(e) => {
|
||||
e.preventDefault();
|
||||
if (!input.trim()) return;
|
||||
if (status === "streaming") {
|
||||
// Buffer the message — it will be sent when the current turn completes
|
||||
setPendingMessage({ text: input, model });
|
||||
} else {
|
||||
trackedSendMessage({ text: input }, { metadata: { model } });
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
export const MODEL_OPTIONS = [
|
||||
"gpt-4o-mini",
|
||||
"gpt-4o",
|
||||
"claude-sonnet-4-6",
|
||||
"claude-opus-4-6",
|
||||
];
|
||||
|
||||
export const DEFAULT_MODEL = "gpt-4o-mini";
|
||||
@@ -0,0 +1,15 @@
|
||||
import { PrismaPg } from "@prisma/adapter-pg";
|
||||
import { PrismaClient } from "../../lib/generated/prisma/client";
|
||||
|
||||
const globalForPrisma = globalThis as unknown as { prisma: PrismaClient | undefined };
|
||||
|
||||
function createClient() {
|
||||
const adapter = new PrismaPg({ connectionString: process.env.DATABASE_URL! });
|
||||
return new PrismaClient({ adapter });
|
||||
}
|
||||
|
||||
export const prisma = globalForPrisma.prisma ?? createClient();
|
||||
|
||||
if (process.env.NODE_ENV !== "production") {
|
||||
globalForPrisma.prisma = prisma;
|
||||
}
|
||||
@@ -5,6 +5,13 @@ import { openai } from "@ai-sdk/openai";
|
||||
import { anthropic } from "@ai-sdk/anthropic";
|
||||
import { z } from "zod";
|
||||
import os from "node:os";
|
||||
import { PrismaPg } from "@prisma/adapter-pg";
|
||||
import { PrismaClient } from "../../lib/generated/prisma/client";
|
||||
|
||||
const adapter = new PrismaPg({ connectionString: process.env.DATABASE_URL! });
|
||||
const prisma = new PrismaClient({ adapter });
|
||||
|
||||
import { DEFAULT_MODEL } from "@/lib/models";
|
||||
|
||||
const MODELS: Record<string, () => LanguageModel> = {
|
||||
"gpt-4o-mini": () => openai("gpt-4o-mini"),
|
||||
@@ -13,9 +20,6 @@ const MODELS: Record<string, () => LanguageModel> = {
|
||||
"claude-opus-4-6": () => anthropic("claude-opus-4-6"),
|
||||
};
|
||||
|
||||
export const MODEL_OPTIONS = Object.keys(MODELS);
|
||||
export const DEFAULT_MODEL = "gpt-4o-mini";
|
||||
|
||||
function getModel(modelId?: string): LanguageModel {
|
||||
const factory = MODELS[modelId ?? DEFAULT_MODEL];
|
||||
if (!factory) return MODELS[DEFAULT_MODEL]!();
|
||||
@@ -83,6 +87,19 @@ declare const Deno: unknown;
|
||||
export const aiChat = chat.task({
|
||||
id: "ai-chat",
|
||||
warmTimeoutInSeconds: 10,
|
||||
onChatStart: async ({ chatId }) => {
|
||||
await prisma.chat.upsert({
|
||||
where: { id: chatId },
|
||||
create: { id: chatId, title: "New chat" },
|
||||
update: {},
|
||||
});
|
||||
},
|
||||
onTurnComplete: async ({ chatId, uiMessages }) => {
|
||||
await prisma.chat.update({
|
||||
where: { id: chatId },
|
||||
data: { messages: uiMessages as any },
|
||||
});
|
||||
},
|
||||
run: async ({ messages, clientData, stopSignal }) => {
|
||||
const { model: modelId } = z
|
||||
.object({ model: z.string().optional() })
|
||||
|
||||
@@ -1,7 +1,15 @@
|
||||
import { defineConfig } from "@trigger.dev/sdk";
|
||||
import { prismaExtension } from "@trigger.dev/build/extensions/prisma";
|
||||
|
||||
export default defineConfig({
|
||||
project: process.env.TRIGGER_PROJECT_REF!,
|
||||
dirs: ["./src/trigger"],
|
||||
maxDuration: 300,
|
||||
build: {
|
||||
extensions: [
|
||||
prismaExtension({
|
||||
mode: "modern",
|
||||
}),
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user