Minor fixes around reconnecting streams

This commit is contained in:
Eric Allam
2026-03-06 10:47:30 +00:00
parent a59e37c07e
commit be405bbe28
2 changed files with 88 additions and 22 deletions
+26 -22
View File
@@ -210,7 +210,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
| undefined;
private sessions: Map<string, ChatSessionState> = new Map();
private activeReconnects: Map<string, AbortController> = new Map();
private activeStreams: Map<string, AbortController> = new Map();
constructor(options: TriggerChatTransportOptions) {
this.taskId = options.task;
@@ -277,6 +277,15 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
const apiClient = new ApiClient(this.baseURL, session.publicAccessToken);
await apiClient.sendInputStream(session.runId, CHAT_MESSAGES_STREAM_ID, minimalPayload);
// Cancel any active reconnect stream for this chatId before
// opening a new subscription for the new turn.
const activeStream = this.activeStreams.get(chatId);
if (activeStream) {
activeStream.abort();
this.activeStreams.delete(chatId);
}
return this.subscribeToStream(
session.runId,
session.publicAccessToken,
@@ -331,20 +340,21 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
return null;
}
// Abort any previous reconnect for this chatId (e.g. React strict mode
// double-firing the effect) to avoid duplicate SSE connections.
const prev = this.activeReconnects.get(options.chatId);
if (prev) {
prev.abort();
// Deduplicate: if there's already an active stream for this chatId,
// return null so the second caller no-ops.
if (this.activeStreams.has(options.chatId)) {
return null;
}
const reconnectAbort = new AbortController();
this.activeReconnects.set(options.chatId, reconnectAbort);
const abortController = new AbortController();
this.activeStreams.set(options.chatId, abortController);
return this.subscribeToStream(
session.runId,
session.publicAccessToken,
reconnectAbort.signal,
options.chatId
abortController.signal,
options.chatId,
{ sendStopOnAbort: false }
);
};
@@ -408,7 +418,8 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
runId: string,
accessToken: string,
abortSignal: AbortSignal | undefined,
chatId?: string
chatId?: string,
options?: { sendStopOnAbort?: boolean }
): ReadableStream<UIMessageChunk> {
const headers: Record<string, string> = {
Authorization: `Bearer ${accessToken}`,
@@ -427,13 +438,14 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
? AbortSignal.any([abortSignal, internalAbort.signal])
: internalAbort.signal;
// When the caller aborts (user calls stop()), send a stop signal to the
// running task via input streams, then close the SSE connection.
// When the caller aborts (user calls stop()), close the SSE connection.
// Only send a stop signal to the task if this is a user-initiated stop
// (sendStopOnAbort), not an internal stream management abort.
if (abortSignal) {
abortSignal.addEventListener(
"abort",
() => {
if (session) {
if (options?.sendStopOnAbort !== false && session) {
session.skipToTurnComplete = true;
const api = new ApiClient(this.baseURL, session.publicAccessToken);
api
@@ -468,14 +480,6 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
const { done, value } = await reader.read();
if (done) {
// Only delete session if the stream ended naturally (not aborted by stop).
// When the user clicks stop, the abort closes the SSE reader which
// returns done=true, but the run is still alive and waiting for
// the next message via input streams.
if (chatId && !combinedSignal.aborted) {
this.sessions.delete(chatId);
this.notifySessionChange(chatId, null);
}
controller.close();
return;
}
+62
View File
@@ -0,0 +1,62 @@
# AI Chat Reference App
A multi-turn chat app built with the AI SDK's `useChat` hook and Trigger.dev's `chat.task`. Conversations run as durable Trigger.dev tasks with realtime streaming, automatic message accumulation, and persistence across page refreshes.
## Data Models
### Chat
The conversation itself — your application data.
| Column | Description |
| ---------- | ---------------------------------------- |
| `id` | Unique chat ID (generated on the client) |
| `title` | Display title for the sidebar |
| `messages` | Full `UIMessage[]` history (JSON) |
A Chat lives forever (until the user deletes it). It is independent of any particular Trigger.dev run.
### ChatSession
The transport's connection state for a chat — what the frontend needs to reconnect to the same Trigger.dev run after a page refresh.
| Column | Description |
| ------------------- | --------------------------------------------------------------------------- |
| `id` | Same as the chat ID (1:1 relationship) |
| `runId` | The Trigger.dev run handling this conversation |
| `publicAccessToken` | Scoped token for reading the run's stream and sending input stream messages |
| `lastEventId` | Stream position — used to resume without replaying old events |
A Chat can outlive many ChatSessions. When the run ends (turn timeout, max turns reached, crash), the ChatSession is gone but the Chat and its messages remain. The next message from the user starts a fresh run and creates a new ChatSession for the same Chat.
**Think of it as: Chat = the conversation, ChatSession = the live connection to the run handling it.**
## Lifecycle Hooks
Persistence is handled server-side in the Trigger.dev task via three hooks:
- **`onChatStart`** — Creates the Chat and ChatSession records when a new conversation starts (turn 0).
- **`onTurnStart`** — Saves messages and updates the session _before_ streaming begins, so a mid-stream page refresh still shows the user's message.
- **`onTurnComplete`** — Saves the assistant's response and the `lastEventId` for stream resumption.
## Setup
```bash
# From the repo root
pnpm run docker # Start PostgreSQL, Redis, Electric
pnpm run db:migrate # Run webapp migrations
pnpm run db:seed # Seed the database
# Set up the reference app's database
cd references/ai-chat
cp .env.example .env # Edit DATABASE_URL if needed
npx prisma migrate deploy
# Build and run
pnpm run build --filter trigger.dev --filter @trigger.dev/sdk
pnpm run dev --filter webapp # In one terminal
cd references/ai-chat && pnpm exec trigger dev # In another
cd references/ai-chat && pnpm run dev # In another
```
Open http://localhost:3000 to use the chat app.