mirror of
https://github.com/mauriceboe/TREK.git
synced 2026-11-14 08:11:03 +00:00
feat(assistant): trip assistant chat with your own AI connection
Adds an in-app assistant that reads and plans the open trip. Users connect their own provider (Sign in with ChatGPT, or an Anthropic, OpenAI, Gemini or OpenRouter key) or use the one the admin set up, including a custom OpenAI-compatible server. It is an addon and only shows while a connection is usable. The chat runs on the existing MCP tools, pinned to the trip. Every change is held as a change card the user applies or discards, and applied changes can be undone. On the desktop it is a floating window that can be moved and resized; on the phone it is fullscreen. Search results show as place cards that can go onto a day, weather lookups as a strip, proposed places as ghost markers on the map with a before/after switch, and an answer can be shared into the trip chat. Booking files can be dropped in and are read into change cards.
This commit is contained in:
@@ -0,0 +1,133 @@
|
||||
/**
|
||||
* FE-AI-API-001 to FE-AI-API-007
|
||||
*
|
||||
* The two streaming calls bypass axios, so they carry by hand what the axios
|
||||
* instance would add (the cookie, the socket id) and turn a refusal before the
|
||||
* stream into a typed error the chat can name. The refusal's code comes from
|
||||
* the body when the server names one, and from the status otherwise.
|
||||
*/
|
||||
import type { AssistantChatRequest } from '@trek/shared';
|
||||
import { http, HttpResponse } from 'msw';
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest';
|
||||
import { server } from '../../tests/helpers/msw/server';
|
||||
import { assistantApi, AssistantRequestError } from './assistant';
|
||||
import { getSocketId } from './websocket';
|
||||
|
||||
const BODY: AssistantChatRequest = { locale: 'en', timeZone: 'Europe/Berlin', threadId: null, text: 'Hi' };
|
||||
|
||||
function streamResponse(text: string): Response {
|
||||
const bytes = new TextEncoder().encode(text);
|
||||
let sent = false;
|
||||
return {
|
||||
ok: true,
|
||||
status: 200,
|
||||
body: {
|
||||
getReader: () => ({
|
||||
read: async () => {
|
||||
if (sent) return { done: true, value: undefined };
|
||||
sent = true;
|
||||
return { done: false, value: bytes };
|
||||
},
|
||||
cancel: async () => {},
|
||||
}),
|
||||
},
|
||||
} as unknown as Response;
|
||||
}
|
||||
|
||||
function refusal(status: number, body: unknown): Response {
|
||||
return {
|
||||
ok: false,
|
||||
status,
|
||||
body: null,
|
||||
json: async () => {
|
||||
if (body === undefined) throw new SyntaxError('not json');
|
||||
return body;
|
||||
},
|
||||
} as unknown as Response;
|
||||
}
|
||||
|
||||
async function refusedWith(res: Response) {
|
||||
vi.spyOn(globalThis, 'fetch').mockResolvedValueOnce(res);
|
||||
try {
|
||||
await assistantApi.chat(4, BODY, new AbortController().signal);
|
||||
} catch (err) {
|
||||
return err;
|
||||
}
|
||||
throw new Error('expected a refusal');
|
||||
}
|
||||
|
||||
afterEach(() => vi.restoreAllMocks());
|
||||
|
||||
describe('assistantApi streams', () => {
|
||||
it('FE-AI-API-001: posts the message as JSON with the cookie and the socket id, and yields the chunks', async () => {
|
||||
vi.mocked(getSocketId).mockReturnValueOnce('sock-1');
|
||||
const fetchSpy = vi
|
||||
.spyOn(globalThis, 'fetch')
|
||||
.mockResolvedValueOnce(streamResponse('data: {"type":"start","messageId":"m1"}\n\ndata: [DONE]\n\n'));
|
||||
const signal = new AbortController().signal;
|
||||
|
||||
const chunks: unknown[] = [];
|
||||
for await (const chunk of await assistantApi.chat(4, BODY, signal)) chunks.push(chunk);
|
||||
|
||||
expect(chunks).toEqual([{ type: 'start', messageId: 'm1' }]);
|
||||
const [url, init] = fetchSpy.mock.calls[0] ?? [];
|
||||
expect(url).toBe('/api/trips/4/assistant/chat');
|
||||
expect(init).toMatchObject({ method: 'POST', credentials: 'include', signal, body: JSON.stringify(BODY) });
|
||||
expect(init?.headers).toMatchObject({
|
||||
'content-type': 'application/json',
|
||||
accept: 'text/event-stream',
|
||||
'X-Socket-Id': 'sock-1',
|
||||
});
|
||||
});
|
||||
|
||||
it('FE-AI-API-002: sends decisions to the thread they belong to', async () => {
|
||||
const fetchSpy = vi.spyOn(globalThis, 'fetch').mockResolvedValueOnce(streamResponse('data: [DONE]\n\n'));
|
||||
const body = { locale: 'de', timeZone: 'UTC', decisions: [{ toolCallId: 'c1', approved: true }] };
|
||||
await assistantApi.decide(4, 7, body, new AbortController().signal);
|
||||
expect(fetchSpy.mock.calls[0]?.[0]).toBe('/api/trips/4/assistant/threads/7/decisions');
|
||||
});
|
||||
|
||||
it('FE-AI-API-003: a refusal names the code the body gives', async () => {
|
||||
const err = await refusedWith(refusal(429, { error: 'Daily limit reached', code: 'daily_limit' }));
|
||||
expect(err).toBeInstanceOf(AssistantRequestError);
|
||||
expect(err).toMatchObject({ code: 'daily_limit', status: 429 });
|
||||
});
|
||||
|
||||
it('FE-AI-API-004: without a known code the status decides', async () => {
|
||||
expect(await refusedWith(refusal(429, { error: 'slow down' }))).toMatchObject({ code: 'too_many_requests' });
|
||||
expect(await refusedWith(refusal(409, { error: 'x', code: 'conflict' }))).toMatchObject({ code: 'busy' });
|
||||
expect(await refusedWith(refusal(502, undefined))).toMatchObject({ code: 'provider_unavailable' });
|
||||
expect(await refusedWith(refusal(500, { error: 'boom' }))).toMatchObject({ code: 'generic' });
|
||||
});
|
||||
|
||||
it('FE-AI-API-005: an accepted request without a body is a generic failure, not a silent empty answer', async () => {
|
||||
const err = await refusedWith({ ok: true, status: 200, body: null } as unknown as Response);
|
||||
expect(err).toMatchObject({ code: 'generic', status: 200 });
|
||||
});
|
||||
});
|
||||
|
||||
describe('assistantApi JSON calls', () => {
|
||||
it('FE-AI-API-006: undo posts an empty JSON body to the escaped tool call id', async () => {
|
||||
let seen: { path: string; body: unknown } | null = null;
|
||||
server.use(
|
||||
http.post('/api/trips/4/assistant/tool-calls/:id/undo', async ({ request, params }) => {
|
||||
seen = { path: String(params.id), body: await request.json() };
|
||||
return HttpResponse.json({ undone: true });
|
||||
})
|
||||
);
|
||||
await expect(assistantApi.undo(4, 'a/b')).resolves.toEqual({ undone: true });
|
||||
expect(seen).toEqual({ path: 'a/b', body: {} });
|
||||
});
|
||||
|
||||
it('FE-AI-API-007: stop posts an empty JSON body', async () => {
|
||||
let body: unknown = null;
|
||||
server.use(
|
||||
http.post('/api/trips/4/assistant/stop', async ({ request }) => {
|
||||
body = await request.json();
|
||||
return HttpResponse.json({ success: true });
|
||||
})
|
||||
);
|
||||
await expect(assistantApi.stop(4)).resolves.toEqual({ success: true });
|
||||
expect(body).toEqual({});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,186 @@
|
||||
import {
|
||||
ASSISTANT_ERROR_CODES,
|
||||
type AiAdminConnectionResponse,
|
||||
type AiAdminSettings,
|
||||
type AiApiKeyConnectRequest,
|
||||
type AiAuthStartResponse,
|
||||
type AiChatgptFinishRequest,
|
||||
type AiClaudeConnectorResponse,
|
||||
type AiConnectionScope,
|
||||
type AiConnectionView,
|
||||
type AiInstanceConnectionRequest,
|
||||
type AiModelListResponse,
|
||||
type AiOpenrouterFinishRequest,
|
||||
type AiPrivacyAckRequest,
|
||||
type AiStatusResponse,
|
||||
type AssistantAddFoundPlaceRequest,
|
||||
type AssistantAddFoundPlaceResponse,
|
||||
type AssistantChatRequest,
|
||||
type AssistantDecisionsRequest,
|
||||
type AssistantErrorCode,
|
||||
type AssistantShareRequest,
|
||||
type AssistantThreadListResponse,
|
||||
type AssistantThreadResponse,
|
||||
type AssistantUndoResponse,
|
||||
type CollabMessageResponse,
|
||||
type SuccessResponse,
|
||||
type UiStreamChunk,
|
||||
} from '@trek/shared';
|
||||
import apiClient from './client';
|
||||
import { parseUiMessageStream } from './uiMessageStream';
|
||||
import { getSocketId } from './websocket';
|
||||
|
||||
/**
|
||||
* The trip assistant and the AI connections behind it, in their own file
|
||||
* rather than in the `client.ts` god module.
|
||||
*
|
||||
* The two streaming calls (a new message, and the decisions on held changes)
|
||||
* go through `fetch`, not axios: the answer is a server-sent event stream that
|
||||
* runs for as long as the model works, and axios would buffer it whole behind
|
||||
* its 8 s timeout. Everything else is plain JSON through `apiClient`.
|
||||
*/
|
||||
|
||||
/** A refusal the server answered before the stream opened, named by its `AssistantErrorCode`. */
|
||||
export class AssistantRequestError extends Error {
|
||||
readonly code: AssistantErrorCode;
|
||||
readonly status: number;
|
||||
|
||||
constructor(code: AssistantErrorCode, status: number) {
|
||||
super(`Assistant request refused: ${code}`);
|
||||
this.name = 'AssistantRequestError';
|
||||
this.code = code;
|
||||
this.status = status;
|
||||
}
|
||||
}
|
||||
|
||||
const KNOWN_CODES: ReadonlySet<string> = new Set(ASSISTANT_ERROR_CODES);
|
||||
|
||||
/**
|
||||
* The code of a refusal: the body's `code` when it names one, otherwise what
|
||||
* the status alone can say.
|
||||
*/
|
||||
async function refusalCode(res: Response): Promise<AssistantErrorCode> {
|
||||
try {
|
||||
const body: unknown = await res.json();
|
||||
const code = (body as { code?: unknown } | null)?.code;
|
||||
if (typeof code === 'string' && KNOWN_CODES.has(code)) return code as AssistantErrorCode;
|
||||
} catch {
|
||||
// Not JSON (a proxy's error page); the status below still says something.
|
||||
}
|
||||
// A bare 429 is TREK's own limit or a proxy's; a provider's limit always names its code.
|
||||
if (res.status === 429) return 'too_many_requests';
|
||||
if (res.status === 409) return 'busy';
|
||||
if (res.status === 502 || res.status === 503 || res.status === 504) return 'provider_unavailable';
|
||||
return 'generic';
|
||||
}
|
||||
|
||||
async function openStream(path: string, body: unknown, signal: AbortSignal): Promise<AsyncGenerator<UiStreamChunk>> {
|
||||
// A file goes as multipart, where the browser sets the content type with its boundary.
|
||||
const form = body instanceof FormData;
|
||||
const headers: Record<string, string> = form
|
||||
? { accept: 'text/event-stream' }
|
||||
: { 'content-type': 'application/json', accept: 'text/event-stream' };
|
||||
const socketId = getSocketId();
|
||||
if (socketId) headers['X-Socket-Id'] = socketId;
|
||||
const res = await fetch(`/api${path}`, {
|
||||
method: 'POST',
|
||||
credentials: 'include',
|
||||
headers,
|
||||
body: form ? body : JSON.stringify(body),
|
||||
signal,
|
||||
});
|
||||
if (!res.ok) throw new AssistantRequestError(await refusalCode(res), res.status);
|
||||
if (!res.body) throw new AssistantRequestError('generic', res.status);
|
||||
return parseUiMessageStream(res.body);
|
||||
}
|
||||
|
||||
const tripBase = (tripId: number) => `/trips/${tripId}/assistant`;
|
||||
|
||||
/**
|
||||
* For the calls that make the server go and ask the provider (listing models,
|
||||
* testing a key, exchanging a sign-in code): the shared 8 s is TREK's own
|
||||
* budget and too short for a provider's.
|
||||
*/
|
||||
const viaProvider = { timeout: 30_000 };
|
||||
|
||||
export const assistantApi = {
|
||||
threads: (tripId: number): Promise<AssistantThreadListResponse> =>
|
||||
apiClient.get(`${tripBase(tripId)}/threads`).then((r) => r.data),
|
||||
thread: (tripId: number, threadId: number): Promise<AssistantThreadResponse> =>
|
||||
apiClient.get(`${tripBase(tripId)}/threads/${threadId}`).then((r) => r.data),
|
||||
deleteThread: (tripId: number, threadId: number): Promise<SuccessResponse> =>
|
||||
apiClient.delete(`${tripBase(tripId)}/threads/${threadId}`).then((r) => r.data),
|
||||
undo: (tripId: number, toolCallId: string): Promise<AssistantUndoResponse> =>
|
||||
apiClient.post(`${tripBase(tripId)}/tool-calls/${encodeURIComponent(toolCallId)}/undo`, {}).then((r) => r.data),
|
||||
// Ends this user's running turn, wherever it runs (another tab included).
|
||||
stop: (tripId: number): Promise<SuccessResponse> =>
|
||||
apiClient.post(`${tripBase(tripId)}/stop`, {}).then((r) => r.data),
|
||||
// A place a search found, suggested for a day: answers with the message holding the change.
|
||||
addFoundPlace: (
|
||||
tripId: number,
|
||||
threadId: number,
|
||||
body: AssistantAddFoundPlaceRequest
|
||||
): Promise<AssistantAddFoundPlaceResponse> =>
|
||||
apiClient.post(`${tripBase(tripId)}/threads/${threadId}/found-places`, body).then((r) => r.data),
|
||||
// An answer posted to the trip chat, marked as shared from the assistant; the client's socket id skips the echo.
|
||||
share: (
|
||||
tripId: number,
|
||||
threadId: number,
|
||||
messageId: string,
|
||||
body: AssistantShareRequest
|
||||
): Promise<CollabMessageResponse> =>
|
||||
apiClient
|
||||
.post(`${tripBase(tripId)}/threads/${threadId}/messages/${encodeURIComponent(messageId)}/share`, body)
|
||||
.then((r) => r.data),
|
||||
|
||||
chat: (tripId: number, body: AssistantChatRequest, signal: AbortSignal) =>
|
||||
openStream(`${tripBase(tripId)}/chat`, body, signal),
|
||||
decide: (tripId: number, threadId: number, body: AssistantDecisionsRequest, signal: AbortSignal) =>
|
||||
openStream(`${tripBase(tripId)}/threads/${threadId}/decisions`, body, signal),
|
||||
// A booking file: the server reads it and proposes each booking it finds as a change.
|
||||
attach: (tripId: number, form: FormData, signal: AbortSignal) =>
|
||||
openStream(`${tripBase(tripId)}/attachments`, form, signal),
|
||||
};
|
||||
|
||||
/** A user's own connection, and what the assistant runs on for them. */
|
||||
export const aiApi = {
|
||||
status: (): Promise<AiStatusResponse> => apiClient.get('/ai/status').then((r) => r.data),
|
||||
models: (): Promise<AiModelListResponse> => apiClient.get('/ai/models', viaProvider).then((r) => r.data),
|
||||
// Saving tests the key against the provider first.
|
||||
connectKey: (body: AiApiKeyConnectRequest): Promise<AiConnectionView> =>
|
||||
apiClient.put('/ai/connection', body, viaProvider).then((r) => r.data),
|
||||
setModel: (model: string): Promise<AiConnectionView> =>
|
||||
apiClient.patch('/ai/connection/model', { model }).then((r) => r.data),
|
||||
disconnect: (): Promise<SuccessResponse> => apiClient.delete('/ai/connection').then((r) => r.data),
|
||||
ackPrivacy: (body: AiPrivacyAckRequest): Promise<SuccessResponse> =>
|
||||
apiClient.post('/ai/privacy-ack', body).then((r) => r.data),
|
||||
|
||||
chatgptStart: (): Promise<AiAuthStartResponse> => apiClient.post('/ai/connect/chatgpt/start', {}).then((r) => r.data),
|
||||
chatgptFinish: (body: AiChatgptFinishRequest): Promise<AiConnectionView> =>
|
||||
apiClient.post('/ai/connect/chatgpt/finish', body, viaProvider).then((r) => r.data),
|
||||
chatgptSignOut: (): Promise<SuccessResponse> =>
|
||||
apiClient.post('/ai/connect/chatgpt/sign-out', {}).then((r) => r.data),
|
||||
|
||||
openrouterStart: (scope: AiConnectionScope): Promise<AiAuthStartResponse> =>
|
||||
apiClient
|
||||
.post(scope === 'instance' ? '/admin/ai/connect/openrouter/start' : '/ai/connect/openrouter/start', { scope })
|
||||
.then((r) => r.data),
|
||||
openrouterFinish: (body: AiOpenrouterFinishRequest): Promise<AiConnectionView> =>
|
||||
apiClient.post('/ai/connect/openrouter/finish', body, viaProvider).then((r) => r.data),
|
||||
};
|
||||
|
||||
/** The connection the instance provides for everyone, and the assistant's admin settings. */
|
||||
export const aiAdminApi = {
|
||||
settings: (): Promise<AiAdminSettings> => apiClient.get('/admin/ai/settings').then((r) => r.data),
|
||||
saveSettings: (body: AiAdminSettings): Promise<AiAdminSettings> =>
|
||||
apiClient.put('/admin/ai/settings', body).then((r) => r.data),
|
||||
connection: (): Promise<AiAdminConnectionResponse> => apiClient.get('/admin/ai/connection').then((r) => r.data),
|
||||
saveConnection: (body: AiInstanceConnectionRequest): Promise<AiAdminConnectionResponse> =>
|
||||
apiClient.put('/admin/ai/connection', body, viaProvider).then((r) => r.data),
|
||||
removeConnection: (): Promise<SuccessResponse> => apiClient.delete('/admin/ai/connection').then((r) => r.data),
|
||||
models: (): Promise<AiModelListResponse> => apiClient.get('/admin/ai/models', viaProvider).then((r) => r.data),
|
||||
};
|
||||
|
||||
export const claudeConnectorApi = {
|
||||
get: (): Promise<AiClaudeConnectorResponse> => apiClient.get('/integrations/claude-connector').then((r) => r.data),
|
||||
};
|
||||
@@ -0,0 +1,150 @@
|
||||
/**
|
||||
* FE-AI-STREAM-001 to FE-AI-STREAM-011
|
||||
*
|
||||
* The parser reads the server's assistant stream, so its contract is the
|
||||
* shared fixture: the same file the server's writer is tested against. What
|
||||
* the network does to that stream (cutting it at any byte, CRLF from a proxy,
|
||||
* keep-alive comments) must not change what comes out.
|
||||
*/
|
||||
import { uiStreamChunkSchema, type UiStreamChunk } from '@trek/shared';
|
||||
import { readFileSync } from 'node:fs';
|
||||
import { dirname, resolve } from 'node:path';
|
||||
import { fileURLToPath } from 'node:url';
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
import { MAX_EVENT_CHARS, parseUiMessageStream, STALL_MS } from './uiMessageStream';
|
||||
|
||||
const FIXTURE = readFileSync(
|
||||
resolve(dirname(fileURLToPath(import.meta.url)), '../../../shared/src/assistant/__fixtures__/ui-stream.v1.txt'),
|
||||
'utf8'
|
||||
);
|
||||
|
||||
/** A response body that hands out the given pieces one read at a time. */
|
||||
function bodyOf(pieces: (string | Uint8Array)[]) {
|
||||
const encoder = new TextEncoder();
|
||||
let i = 0;
|
||||
const cancel = vi.fn(async () => {});
|
||||
const body = {
|
||||
getReader: () => ({
|
||||
read: async () => {
|
||||
if (i >= pieces.length) return { done: true, value: undefined };
|
||||
const piece = pieces[i++];
|
||||
return { done: false, value: typeof piece === 'string' ? encoder.encode(piece) : piece };
|
||||
},
|
||||
cancel,
|
||||
}),
|
||||
} as unknown as ReadableStream<Uint8Array>;
|
||||
return { body, cancel };
|
||||
}
|
||||
|
||||
async function collect(body: ReadableStream<Uint8Array>): Promise<UiStreamChunk[]> {
|
||||
const chunks: UiStreamChunk[] = [];
|
||||
for await (const chunk of parseUiMessageStream(body)) chunks.push(chunk);
|
||||
return chunks;
|
||||
}
|
||||
|
||||
/** The fixture's chunks as plain JSON, the way a reader of the file sees them. */
|
||||
function fixtureChunks(): UiStreamChunk[] {
|
||||
return FIXTURE.split('\n\n')
|
||||
.map((event) => event.replace(/^data: /, '').trim())
|
||||
.filter((data) => data && data !== '[DONE]')
|
||||
.map((data) => uiStreamChunkSchema.parse(JSON.parse(data)));
|
||||
}
|
||||
|
||||
describe('parseUiMessageStream', () => {
|
||||
it('FE-AI-STREAM-001: reads the shared fixture into its chunks, in order', async () => {
|
||||
const chunks = await collect(bodyOf([FIXTURE]).body);
|
||||
expect(chunks).toEqual(fixtureChunks());
|
||||
expect(chunks.map((c) => c.type)).toContain('tool-approval-request');
|
||||
expect(chunks[chunks.length - 1]).toEqual({ type: 'finish', finishReason: 'tool-calls' });
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-002: gives the same chunks when every byte arrives in its own read', async () => {
|
||||
const bytes = new TextEncoder().encode(FIXTURE);
|
||||
const pieces = Array.from(bytes, (b) => new Uint8Array([b]));
|
||||
expect(await collect(bodyOf(pieces).body)).toEqual(fixtureChunks());
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-003: reads CRLF line ends, also when a CR and its LF arrive apart', async () => {
|
||||
const crlf = FIXTURE.replace(/\n/g, '\r\n');
|
||||
const cut = crlf.indexOf('\r\n') + 1;
|
||||
expect(await collect(bodyOf([crlf.slice(0, cut), crlf.slice(cut)]).body)).toEqual(fixtureChunks());
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-004: skips keep-alive comments and fields other than data', async () => {
|
||||
const stream = [
|
||||
': keepalive\n\n',
|
||||
'event: message\nid: 1\ndata: {"type":"start-step"}\n\n',
|
||||
': keepalive\n\n',
|
||||
'data: [DONE]\n\n',
|
||||
].join('');
|
||||
expect(await collect(bodyOf([stream]).body)).toEqual([{ type: 'start-step' }]);
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-005: skips chunk types this client does not know, and torn JSON', async () => {
|
||||
const stream = [
|
||||
'data: {"type":"reasoning-delta","delta":"hm"}\n\n',
|
||||
'data: {"type":"text-delta","id":"t"\n\n',
|
||||
'data: {"type":"text-delta","id":"t","delta":"ok"}\n\n',
|
||||
'data: [DONE]\n\n',
|
||||
].join('');
|
||||
expect(await collect(bodyOf([stream]).body)).toEqual([{ type: 'text-delta', id: 't', delta: 'ok' }]);
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-006: reads an unknown error code as generic', async () => {
|
||||
const stream = 'data: {"type":"error","errorText":"quota_exploded"}\n\ndata: [DONE]\n\n';
|
||||
expect(await collect(bodyOf([stream]).body)).toEqual([{ type: 'error', errorText: 'generic' }]);
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-007: joins the lines of a multi-line data field', async () => {
|
||||
const stream = 'data: {"type":"text-delta",\ndata: "id":"t","delta":"x"}\n\ndata: [DONE]\n\n';
|
||||
expect(await collect(bodyOf([stream]).body)).toEqual([{ type: 'text-delta', id: 't', delta: 'x' }]);
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-008: stops at [DONE] and lets go of the connection', async () => {
|
||||
const { body, cancel } = bodyOf([
|
||||
'data: {"type":"start-step"}\n\ndata: [DONE]\n\ndata: {"type":"finish-step"}\n\n',
|
||||
]);
|
||||
expect(await collect(body)).toEqual([{ type: 'start-step' }]);
|
||||
expect(cancel).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-009: a stream that closes before [DONE] broke off, after the chunks it finished', async () => {
|
||||
const { body, cancel } = bodyOf(['data: {"type":"start-step"}\n\ndata: {"type":"finish-st']);
|
||||
const seen: UiStreamChunk[] = [];
|
||||
const read = async () => {
|
||||
for await (const chunk of parseUiMessageStream(body)) seen.push(chunk);
|
||||
};
|
||||
await expect(read()).rejects.toThrow('broke off');
|
||||
expect(seen).toEqual([{ type: 'start-step' }]);
|
||||
expect(cancel).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-011: a stream that stays silent past the stall limit fails and lets go', async () => {
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
const cancel = vi.fn(async () => {});
|
||||
const body = {
|
||||
getReader: () => ({ read: () => new Promise(() => {}), cancel }),
|
||||
} as unknown as ReadableStream<Uint8Array>;
|
||||
const failed = expect(collect(body)).rejects.toThrow('stalled');
|
||||
await vi.advanceTimersByTimeAsync(STALL_MS);
|
||||
await failed;
|
||||
expect(cancel).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
|
||||
it('FE-AI-STREAM-010: refuses an event past the size limit, and a consumer leaving early cancels the read', async () => {
|
||||
const huge = bodyOf([`data: ${'x'.repeat(MAX_EVENT_CHARS + 1)}`]);
|
||||
await expect(collect(huge.body)).rejects.toThrow('too large');
|
||||
expect(huge.cancel).toHaveBeenCalledTimes(1);
|
||||
|
||||
const early = bodyOf([FIXTURE]);
|
||||
for await (const chunk of parseUiMessageStream(early.body)) {
|
||||
expect(chunk.type).toBe('start');
|
||||
break;
|
||||
}
|
||||
expect(early.cancel).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,111 @@
|
||||
import { UI_STREAM_DONE, uiStreamChunkSchema, type UiStreamChunk } from '@trek/shared';
|
||||
|
||||
/**
|
||||
* Reads the assistant's server-sent events into typed chunks.
|
||||
*
|
||||
* The wire format is the `@trek/shared` stream contract: one `data: <json>`
|
||||
* event per chunk, closed by `data: [DONE]`. Our own parser rather than the
|
||||
* `assistant-stream` one, because that one drops `tool-approval-request`, the
|
||||
* chunk a held change hangs on, and it would pull the chat runtime into every
|
||||
* chunk that only needs the data layer.
|
||||
*
|
||||
* What it tolerates, because proxies and the network do these things: events
|
||||
* split across reads at any byte, CRLF line ends, `:` keep-alive comments,
|
||||
* and fields other than `data`. A chunk that does not match the contract,
|
||||
* a type this client does not know among them, is skipped, so the server can
|
||||
* add chunk types without breaking an older tab. An event longer than
|
||||
* MAX_EVENT_CHARS ends the stream with an error instead of growing the buffer
|
||||
* without bound.
|
||||
*
|
||||
* The server always closes with `[DONE]` and sends a keep-alive every 15
|
||||
* seconds, so a stream that ends without the marker, or stays silent for
|
||||
* STALL_MS, lost its server (a restart, a dropped proxy) and fails instead of
|
||||
* leaving the chat waiting for an answer that will not come.
|
||||
*/
|
||||
|
||||
export const MAX_EVENT_CHARS = 1024 * 1024;
|
||||
export const STALL_MS = 45_000;
|
||||
|
||||
/** The next read, or a failure once nothing came for STALL_MS. */
|
||||
async function readOrStall(reader: ReadableStreamDefaultReader<Uint8Array>) {
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
const stalled = new Promise<never>((_, reject) => {
|
||||
timer = setTimeout(() => reject(new Error('Assistant stream stalled')), STALL_MS);
|
||||
});
|
||||
try {
|
||||
return await Promise.race([reader.read(), stalled]);
|
||||
} finally {
|
||||
clearTimeout(timer);
|
||||
}
|
||||
}
|
||||
|
||||
export async function* parseUiMessageStream(body: ReadableStream<Uint8Array>): AsyncGenerator<UiStreamChunk> {
|
||||
const reader = body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = '';
|
||||
let finished = false;
|
||||
try {
|
||||
for (;;) {
|
||||
const { done, value } = await readOrStall(reader);
|
||||
buffer += done ? decoder.decode() : decoder.decode(value, { stream: true });
|
||||
// A CR at the very end may be the first half of a CRLF; it waits for the next read.
|
||||
const carry = !done && buffer.endsWith('\r') ? '\r' : '';
|
||||
if (carry) buffer = buffer.slice(0, -1);
|
||||
buffer = buffer.replace(/\r\n?/g, '\n');
|
||||
|
||||
let boundary = buffer.indexOf('\n\n');
|
||||
while (boundary >= 0) {
|
||||
const data = dataOf(buffer.slice(0, boundary));
|
||||
buffer = buffer.slice(boundary + 2);
|
||||
if (data === UI_STREAM_DONE) {
|
||||
finished = true;
|
||||
return;
|
||||
}
|
||||
const chunk = data === null ? null : chunkOf(data);
|
||||
if (chunk) yield chunk;
|
||||
boundary = buffer.indexOf('\n\n');
|
||||
}
|
||||
buffer += carry;
|
||||
|
||||
if (buffer.length > MAX_EVENT_CHARS) throw new Error('Assistant stream event too large');
|
||||
// The connection closed before [DONE]: the server went away mid-answer.
|
||||
if (done) {
|
||||
finished = true;
|
||||
throw new Error('Assistant stream broke off');
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
// Leaving early (the consumer stopped, or an error) releases the connection.
|
||||
if (!finished) {
|
||||
try {
|
||||
await reader.cancel();
|
||||
} catch {
|
||||
// The stream already failed; the error that brought us here is the one that counts.
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** The joined `data` lines of one event, or null for an event without any (a comment). */
|
||||
function dataOf(event: string): string | null {
|
||||
const lines: string[] = [];
|
||||
for (const line of event.split('\n')) {
|
||||
if (line === 'data' || line.startsWith('data:')) {
|
||||
const value = line.slice(5);
|
||||
lines.push(value.startsWith(' ') ? value.slice(1) : value);
|
||||
}
|
||||
}
|
||||
return lines.length ? lines.join('\n') : null;
|
||||
}
|
||||
|
||||
function chunkOf(data: string): UiStreamChunk | null {
|
||||
let json: unknown;
|
||||
try {
|
||||
json = JSON.parse(data);
|
||||
} catch {
|
||||
// A torn or foreign event; the next one is still worth reading.
|
||||
return null;
|
||||
}
|
||||
const parsed = uiStreamChunkSchema.safeParse(json);
|
||||
return parsed.success ? parsed.data : null;
|
||||
}
|
||||
Reference in New Issue
Block a user