diff --git a/osprey_ui/src/components/ask_ai/AskStore.test.ts b/osprey_ui/src/components/ask_ai/AskStore.test.ts new file mode 100644 index 00000000..eb3d7ee4 --- /dev/null +++ b/osprey_ui/src/components/ask_ai/AskStore.test.ts @@ -0,0 +1,178 @@ +import { describe, expect, it } from '@rstest/core'; + +import { AskDeps, createAskStore } from './AskStore'; +import { AskStreamError } from './sseClient'; +import { AskEndpointConfig, AskEvent } from './types'; + +type StreamFn = AskDeps['stream']; + +const ENDPOINT: AskEndpointConfig = { url: 'http://test/ask' }; + +function ev(type: AskEvent['type'], payload: Record = {}, extra: Partial = {}): AskEvent { + return { version: 1, type, payload, ...extra }; +} + +function streamOf(events: AskEvent[]): StreamFn { + return (async function* () { + for (const event of events) { + yield event; + } + }) as unknown as StreamFn; +} + +function gatedStream(before: AskEvent[], gate: Promise, after: AskEvent[]): StreamFn { + return (async function* () { + for (const event of before) { + yield event; + } + await gate; + for (const event of after) { + yield event; + } + }) as unknown as StreamFn; +} + +function throwingStream(before: AskEvent[], error: unknown): StreamFn { + return (async function* () { + for (const event of before) { + yield event; + } + throw error; + }) as unknown as StreamFn; +} + +describe('AskStore', () => { + it('applies a happy-path stream and finishes idle', async () => { + const useStore = createAskStore({ + stream: streamOf([ + ev('conversation_started', {}, { conversation_id: 'c1' }), + ev('query_result', { tool_call_id: 't', name: 'lookup', is_error: false, content: 'r' }), + ev('assistant_message', { text: 'hello' }), + ev('done'), + ]), + }); + useStore.getState().configure(ENDPOINT); + await useStore.getState().sendMessage('hi'); + const s = useStore.getState(); + expect(s.status).toBe('idle'); + expect(s.conversationId).toBe('c1'); + expect(s.messages.map((m) => m.role)).toEqual(['user', 'assistant']); + const assistant = s.messages[1]; + expect(assistant.text).toBe('hello'); + expect(assistant.streaming).toBe(false); + expect(assistant.evidence).toHaveLength(1); + expect(assistant.evidence[0].content).toBe('r'); + }); + + it('prevents a stale response from overwriting a newer conversation', async () => { + let release = (): void => undefined; + const gate = new Promise((resolve) => { + release = resolve; + }); + const useStore = createAskStore({ + stream: gatedStream( + [ev('conversation_started', {}, { conversation_id: 'old' })], + gate, + [ev('assistant_message', { text: 'late answer' }), ev('done')], + ), + }); + useStore.getState().configure(ENDPOINT); + const pending = useStore.getState().sendMessage('first'); + useStore.getState().newChat(); + release(); + await pending; + const s = useStore.getState(); + expect(s.messages).toHaveLength(0); + expect(s.conversationId).toBeUndefined(); + expect(s.status).toBe('idle'); + }); + + it('abort stops streaming without surfacing an error', async () => { + let release = (): void => undefined; + const gate = new Promise((resolve) => { + release = resolve; + }); + const useStore = createAskStore({ + stream: gatedStream( + [ev('conversation_started', {}, { conversation_id: 'c1' })], + gate, + [ev('assistant_message', { text: 'late' }), ev('done')], + ), + }); + useStore.getState().configure(ENDPOINT); + const pending = useStore.getState().sendMessage('hi'); + useStore.getState().abort(); + release(); + await pending; + const s = useStore.getState(); + expect(s.status).toBe('idle'); + expect(s.error).toBeUndefined(); + const assistant = s.messages.find((m) => m.role === 'assistant'); + expect(assistant?.text).toBe(''); + }); + + it('maps an AskStreamError to a safe error state', async () => { + const useStore = createAskStore({ + stream: throwingStream([ev('conversation_started', {}, { conversation_id: 'c1' })], new AskStreamError('provider_error', 'Provider unavailable.')), + }); + useStore.getState().configure(ENDPOINT); + await useStore.getState().sendMessage('hi'); + const s = useStore.getState(); + expect(s.status).toBe('error'); + expect(s.error?.code).toBe('provider_error'); + expect(s.error?.message).toBe('Provider unavailable.'); + }); + + it('never exposes a raw thrown Error message', async () => { + const useStore = createAskStore({ + stream: throwingStream([ev('conversation_started')], new Error('secret stack xyz')), + }); + useStore.getState().configure(ENDPOINT); + await useStore.getState().sendMessage('hi'); + const s = useStore.getState(); + expect(s.status).toBe('error'); + expect(s.error?.code).toBe('internal'); + expect(JSON.stringify(s)).not.toContain('xyz'); + }); + + it('retry re-sends the last user turn without duplicating the user bubble', async () => { + let attempt = 0; + const stream: StreamFn = ((): ReturnType => { + attempt += 1; + if (attempt === 1) { + return (async function* () { + yield ev('conversation_started', {}, { conversation_id: 'c1' }); + throw new AskStreamError('provider_error', 'boom'); + })(); + } + return (async function* () { + yield ev('assistant_message', { text: 'ok now' }); + yield ev('done'); + })(); + }) as unknown as StreamFn; + const useStore = createAskStore({ stream }); + useStore.getState().configure(ENDPOINT); + await useStore.getState().sendMessage('question'); + expect(useStore.getState().status).toBe('error'); + await useStore.getState().retry(); + const s = useStore.getState(); + expect(s.status).toBe('idle'); + const users = s.messages.filter((m) => m.role === 'user'); + expect(users).toHaveLength(1); + expect(users[0].text).toBe('question'); + expect(s.messages.find((m) => m.role === 'assistant')?.text).toBe('ok now'); + }); + + it('newChat clears messages but keeps configuration', async () => { + const useStore = createAskStore({ stream: streamOf([ev('assistant_message', { text: 'a' }), ev('done')]) }); + useStore.getState().configure(ENDPOINT, 'model-x'); + await useStore.getState().sendMessage('hi'); + expect(useStore.getState().messages.length).toBeGreaterThan(0); + useStore.getState().newChat(); + const s = useStore.getState(); + expect(s.messages).toHaveLength(0); + expect(s.conversationId).toBeUndefined(); + expect(s.endpoint).toBeDefined(); + expect(s.model).toBe('model-x'); + }); +}); diff --git a/osprey_ui/src/components/ask_ai/AskStore.ts b/osprey_ui/src/components/ask_ai/AskStore.ts new file mode 100644 index 00000000..6262c77a --- /dev/null +++ b/osprey_ui/src/components/ask_ai/AskStore.ts @@ -0,0 +1,223 @@ +// Zustand conversation store for the Ask panel. +// +// createAskStore(deps) is exported so tests can inject a fake `stream`; the panel and +// app use the default useAskStore. A monotonic requestId guards against stale responses +// (a slow turn cannot mutate a newer conversation), and an AbortController cancels the +// in-flight fetch/reader. Errors are mapped to a safe {code,message} -- never a raw stack. + +import create from 'zustand'; + +import { streamAsk } from './sseClient'; +import { + AskEndpointConfig, + AskEvent, + AskEvidence, + AskMessage, + AskRequest, + AskStatus, + ErrorPayload, + QueryResultPayload, +} from './types'; + +let _counter = 0; + +function uid(): string { + _counter += 1; + return `m${_counter}-${Date.now().toString(36)}`; +} + +function toErrorPayload(e: unknown): ErrorPayload { + if (e && typeof e === 'object' && 'code' in e && 'message' in e) { + const code = (e as { code: unknown }).code; + const message = (e as { message: unknown }).message; + if (typeof code === 'string' && typeof message === 'string') { + return { code, message }; + } + } + return { code: 'internal', message: 'Something went wrong.' }; +} + +function evidenceFromPayload(payload: QueryResultPayload): AskEvidence { + return { + toolCallId: typeof payload.tool_call_id === 'string' ? payload.tool_call_id : undefined, + name: typeof payload.name === 'string' ? payload.name : undefined, + isError: payload.is_error === true, + content: typeof payload.content === 'string' ? payload.content : '', + raw: payload, + }; +} + +function dropLastTurn(messages: AskMessage[]): AskMessage[] { + // Remove the trailing assistant message and the user message before it, so retry does + // not duplicate the user bubble when it re-sends. + const next = [...messages]; + if (next.length > 0 && next[next.length - 1].role === 'assistant') { + next.pop(); + } + if (next.length > 0 && next[next.length - 1].role === 'user') { + next.pop(); + } + return next; +} + +export interface AskDeps { + stream: typeof streamAsk; +} + +export interface AskState { + status: AskStatus; + endpoint?: AskEndpointConfig; + model?: string; + conversationId?: string; + messages: AskMessage[]; + error?: ErrorPayload; + configure(endpoint: AskEndpointConfig, model?: string): void; + sendMessage(text: string, opts?: { contextRef?: string }): Promise; + retry(): Promise; + newChat(): void; + reset(): void; + abort(): void; +} + +export function createAskStore(deps: AskDeps = { stream: streamAsk }) { + let requestId = 0; + let controller: AbortController | undefined; + let lastUserText: string | undefined; + let lastContextRef: string | undefined; + + return create((set, get) => { + const applyEvent = (assistantId: string, event: AskEvent): void => { + switch (event.type) { + case 'conversation_started': + if (typeof event.conversation_id === 'string') { + set({ conversationId: event.conversation_id }); + } + break; + case 'query_result': + set((s) => ({ + messages: s.messages.map((m) => + m.id === assistantId + ? { ...m, evidence: [...m.evidence, evidenceFromPayload(event.payload as QueryResultPayload)] } + : m, + ), + })); + break; + case 'assistant_message': { + const raw = event.payload.text; + const text = typeof raw === 'string' ? raw : ''; + set((s) => ({ messages: s.messages.map((m) => (m.id === assistantId ? { ...m, text } : m)) })); + break; + } + case 'done': + set((s) => ({ + status: 'idle', + messages: s.messages.map((m) => (m.id === assistantId ? { ...m, streaming: false } : m)), + })); + break; + case 'error': { + const err = toErrorPayload(event.payload); + set((s) => ({ + status: 'error', + error: err, + messages: s.messages.map((m) => (m.id === assistantId ? { ...m, streaming: false, error: err } : m)), + })); + break; + } + case 'tool_call': + default: + break; + } + }; + + const runTurn = async (text: string, contextRef: string | undefined, appendUser: boolean): Promise => { + const endpoint = get().endpoint; + if (!endpoint) { + throw new Error('AskStore is not configured; call configure() first.'); + } + controller?.abort(); + requestId += 1; + const myId = requestId; + controller = new AbortController(); + lastUserText = text; + lastContextRef = contextRef; + const assistantId = uid(); + set((s) => { + const withUser = appendUser + ? [...s.messages, { id: uid(), role: 'user' as const, text, evidence: [], streaming: false }] + : s.messages; + return { + status: 'streaming' as const, + error: undefined, + messages: [ + ...withUser, + { id: assistantId, role: 'assistant' as const, text: '', evidence: [], streaming: true }, + ], + }; + }); + const req: AskRequest = { + message: text, + conversation_id: get().conversationId, + model: get().model, + context_ref: contextRef, + }; + try { + for await (const event of deps.stream(endpoint, req, { signal: controller.signal })) { + if (myId !== requestId) { + return; // a newer send/newChat/reset superseded this turn + } + applyEvent(assistantId, event); + } + } catch (e) { + if (myId !== requestId || controller?.signal.aborted) { + return; // stale or user-aborted: do not surface an error + } + const err = toErrorPayload(e); + set((s) => ({ + status: 'error', + error: err, + messages: s.messages.map((m) => (m.id === assistantId ? { ...m, streaming: false, error: err } : m)), + })); + } + }; + + return { + status: 'idle', + messages: [], + configure: (endpoint, model) => set({ endpoint, model }), + abort: () => { + controller?.abort(); + controller = undefined; + requestId += 1; + if (get().status === 'streaming') { + set({ status: 'idle' }); + } + }, + newChat: () => { + requestId += 1; + controller?.abort(); + controller = undefined; + set({ conversationId: undefined, messages: [], status: 'idle', error: undefined }); + }, + reset: () => { + requestId += 1; + controller?.abort(); + controller = undefined; + lastUserText = undefined; + lastContextRef = undefined; + set({ conversationId: undefined, messages: [], status: 'idle', error: undefined }); + }, + retry: async () => { + if (get().status !== 'error' || lastUserText === undefined) { + return; + } + set((s) => ({ status: 'idle', error: undefined, messages: dropLastTurn(s.messages) })); + await runTurn(lastUserText, lastContextRef, true); + }, + sendMessage: async (text, opts) => { + await runTurn(text, opts?.contextRef, true); + }, + }; + }); +} + +export const useAskStore = createAskStore(); diff --git a/osprey_ui/src/components/ask_ai/sseClient.test.ts b/osprey_ui/src/components/ask_ai/sseClient.test.ts new file mode 100644 index 00000000..c2083e2e --- /dev/null +++ b/osprey_ui/src/components/ask_ai/sseClient.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, it } from '@rstest/core'; + +import { AskStreamError, frameToEvent, parseSseFrames, streamAsk } from './sseClient'; +import { AskEndpointConfig, AskEvent, AskRequest } from './types'; + +const ENDPOINT: AskEndpointConfig = { url: 'http://test/ask' }; +const REQ: AskRequest = { message: 'hi' }; + +function frame(type: string, payload: Record = {}): string { + return `event: ${type}\ndata: ${JSON.stringify({ version: 1, type, payload })}\n\n`; +} + +function streamResponse(chunks: string[], status = 200): Response { + const encoder = new TextEncoder(); + const body = new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(encoder.encode(chunk)); + } + controller.close(); + }, + }); + return new Response(body, { status }); +} + +function setFetch(fn: () => Promise): void { + globalThis.fetch = fn as unknown as typeof globalThis.fetch; +} + +async function collect(gen: AsyncGenerator): Promise { + const out: AskEvent[] = []; + for await (const event of gen) { + out.push(event); + } + return out; +} + +describe('parseSseFrames', () => { + it('splits complete frames and keeps a partial remainder', () => { + const { frames, rest } = parseSseFrames('event: a\ndata: 1\n\nevent: b\ndata: 2\n\nevent: c\ndata:'); + expect(frames).toEqual(['event: a\ndata: 1', 'event: b\ndata: 2']); + expect(rest).toBe('event: c\ndata:'); + }); + + it('does not lose or duplicate a frame split across chunks', () => { + let buffer = ''; + const seen: string[] = []; + for (const chunk of ['event: a\nda', 'ta: 1\n\nev', 'ent: b\ndata: 2\n\n']) { + buffer += chunk; + const out = parseSseFrames(buffer); + buffer = out.rest; + seen.push(...out.frames); + } + expect(seen).toEqual(['event: a\ndata: 1', 'event: b\ndata: 2']); + expect(buffer).toBe(''); + }); +}); + +describe('frameToEvent', () => { + it('parses a valid frame', () => { + const event = frameToEvent('event: done\ndata: {"version":1,"type":"done","payload":{}}'); + expect(event.type).toBe('done'); + expect(event.version).toBe(1); + }); + + it('throws on malformed json', () => { + expect(() => frameToEvent('event: x\ndata: not-json')).toThrow(AskStreamError); + }); + + it('rejects an unknown event type (forward-compatible error)', () => { + expect(() => frameToEvent('event: nope\ndata: {"version":1,"type":"nope","payload":{}}')).toThrow(AskStreamError); + }); +}); + +describe('streamAsk', () => { + it('yields events and stops at the terminal', async () => { + setFetch(async () => + streamResponse([ + frame('conversation_started'), + frame('assistant_message', { text: 'hi' }), + frame('done'), + frame('assistant_message', { text: 'should be ignored after done' }), + ]), + ); + const events = await collect(streamAsk(ENDPOINT, REQ)); + expect(events.map((e) => e.type)).toEqual(['conversation_started', 'assistant_message', 'done']); + }); + + it('errors when the stream ends without a terminal event', async () => { + setFetch(async () => streamResponse([frame('conversation_started')])); + await expect(collect(streamAsk(ENDPOINT, REQ))).rejects.toThrow(AskStreamError); + }); + + it('maps a pre-flight HTTP error to the server code', async () => { + setFetch(async () => new Response(JSON.stringify({ error: { code: 'invalid_model', message: 'nope' } }), { status: 400 })); + await expect(collect(streamAsk(ENDPOINT, REQ))).rejects.toMatchObject({ code: 'invalid_model' }); + }); + + it('cancels the reader when iteration stops early', async () => { + let cancelled = false; + const encoder = new TextEncoder(); + const body = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(frame('conversation_started'))); + }, + cancel() { + cancelled = true; + }, + }); + setFetch(async () => new Response(body, { status: 200 })); + const gen = streamAsk(ENDPOINT, REQ); + await gen.next(); + await gen.return(undefined); + expect(cancelled).toBe(true); + }); +}); diff --git a/osprey_ui/src/components/ask_ai/sseClient.ts b/osprey_ui/src/components/ask_ai/sseClient.ts new file mode 100644 index 00000000..2ace9ea7 --- /dev/null +++ b/osprey_ui/src/components/ask_ai/sseClient.ts @@ -0,0 +1,123 @@ +// Chunk-safe POST-SSE client for the Ask endpoint. +// +// Uses fetch + a ReadableStream reader + TextDecoder to consume `event:`/`data:` frames +// separated by a blank line, tolerant of arbitrary chunk boundaries. Abortable via an +// AbortSignal; the reader is always cancelled on return/throw. A pre-flight HTTP error +// (the backend's JSON `{error:{code,message}}`) becomes an AskStreamError; a stream that +// ends without a terminal event is an error; unknown event types are rejected rather +// than silently interpreted. + +import { AskEndpointConfig, AskEvent, AskEventType, AskRequest, ErrorPayload } from './types'; + +export class AskStreamError extends Error { + code: string; + + constructor(code: string, message: string) { + super(message); + this.name = 'AskStreamError'; + this.code = code; + } +} + +const TERMINAL: ReadonlySet = new Set(['done', 'error']); +const KNOWN_TYPES: ReadonlySet = new Set([ + 'conversation_started', + 'tool_call', + 'query_result', + 'assistant_message', + 'done', + 'error', +]); + +// Pure: split a buffer into complete frames plus the trailing (possibly partial) remainder. +export function parseSseFrames(buffer: string): { frames: string[]; rest: string } { + const parts = buffer.split('\n\n'); + const rest = parts.pop() ?? ''; + return { frames: parts.filter((part) => part.trim() !== ''), rest }; +} + +// Pure: parse one raw frame into a validated AskEvent (throws on malformed/unknown type). +export function frameToEvent(frame: string): AskEvent { + let data = ''; + for (const line of frame.split('\n')) { + if (line.startsWith('data:')) { + data += line.slice(5).replace(/^ /, ''); + } + } + let parsed: unknown; + try { + parsed = JSON.parse(data); + } catch { + throw new AskStreamError('internal', 'Malformed event frame.'); + } + const type = (parsed as { type?: unknown } | null)?.type; + if (typeof parsed !== 'object' || parsed === null || typeof type !== 'string' || !KNOWN_TYPES.has(type)) { + throw new AskStreamError('internal', 'Unknown or malformed event.'); + } + return parsed as AskEvent; +} + +async function toHttpError(res: Response): Promise { + let body: { error?: Partial } | undefined; + try { + body = (await res.json()) as { error?: Partial }; + } catch { + return new AskStreamError('internal', `Request failed (${res.status}).`); + } + const err = body?.error; + if (err && typeof err.code === 'string' && typeof err.message === 'string') { + return new AskStreamError(err.code, err.message); + } + return new AskStreamError('internal', `Request failed (${res.status}).`); +} + +export interface StreamOptions { + signal?: AbortSignal; +} + +export async function* streamAsk( + endpoint: AskEndpointConfig, + req: AskRequest, + opts: StreamOptions = {}, +): AsyncGenerator { + const res = await fetch(endpoint.url, { + method: 'POST', + headers: { 'Content-Type': 'application/json', ...(endpoint.headers ?? {}) }, + body: JSON.stringify(req), + signal: opts.signal, + credentials: endpoint.withCredentials ? 'include' : 'same-origin', + }); + if (!res.ok) { + throw await toHttpError(res); + } + if (!res.body) { + throw new AskStreamError('internal', 'Response had no body to stream.'); + } + + const reader = res.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + let sawTerminal = false; + try { + for (let result = await reader.read(); !result.done; result = await reader.read()) { + buffer += decoder.decode(result.value, { stream: true }); + const { frames, rest } = parseSseFrames(buffer); + buffer = rest; + for (const frame of frames) { + const event = frameToEvent(frame); + if (TERMINAL.has(event.type)) { + sawTerminal = true; + } + yield event; + if (sawTerminal) { + return; + } + } + } + if (!sawTerminal) { + throw new AskStreamError('internal', 'Stream ended without a terminal event.'); + } + } finally { + await reader.cancel().catch(() => undefined); + } +} diff --git a/osprey_ui/src/components/ask_ai/types.ts b/osprey_ui/src/components/ask_ai/types.ts new file mode 100644 index 00000000..d61cb916 --- /dev/null +++ b/osprey_ui/src/components/ask_ai/types.ts @@ -0,0 +1,80 @@ +// Vendor-neutral Ask AI wire + domain types. +// +// Wire fields (AskRequest, AskEvent, *Payload) are snake_case to match the JSON the +// backend emits and the existing Osprey UI type conventions. Domain types the store +// owns (AskEvidence, AskMessage) are camelCase. + +export const ASK_EVENT_VERSION = 1; + +export type AskEventType = + | 'conversation_started' + | 'tool_call' + | 'query_result' + | 'assistant_message' + | 'done' + | 'error'; + +export interface AskRequest { + message: string; + conversation_id?: string; + model?: string; + context_ref?: string; +} + +export interface AskEvent { + version: number; + type: AskEventType; + conversation_id?: string; + turn_id?: string; + payload: Record; +} + +export interface ToolCallPayload { + id: string; + name: string; + arguments: Record; +} + +export interface QueryResultPayload { + tool_call_id?: string; + name?: string; + is_error?: boolean; + content?: string; + [key: string]: unknown; +} + +export interface AssistantMessagePayload { + text: string; +} + +export interface ErrorPayload { + code: string; + message: string; +} + +export interface AskEndpointConfig { + url: string; + headers?: Record; + withCredentials?: boolean; +} + +export type AskRole = 'user' | 'assistant'; + +export interface AskEvidence { + toolCallId?: string; + name?: string; + isError: boolean; + content: string; + raw: QueryResultPayload; +} + +export interface AskMessage { + id: string; + role: AskRole; + text: string; + evidence: AskEvidence[]; + streaming: boolean; + error?: ErrorPayload; +} + +export type AskStatus = 'idle' | 'streaming' | 'error';