diff --git a/src/lib/__tests__/agent-api.test.ts b/src/lib/__tests__/agent-api.test.ts index 446d0a9..6243c88 100644 --- a/src/lib/__tests__/agent-api.test.ts +++ b/src/lib/__tests__/agent-api.test.ts @@ -31,20 +31,36 @@ function makeOpts(overrides: Partial = {}): StreamAgentOpts { } } -function okResponse(body: object) { +// ── SSE helpers ─────────────────────────────────────────────────────────────── + +function makeSseResponse(chunks: object[]): Promise { + const encoder = new TextEncoder() + const lines = chunks.map((c) => `data: ${JSON.stringify(c)}\n\n`).join("") + const bytes = encoder.encode(lines) + + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(bytes) + controller.close() + }, + }) + return Promise.resolve( - new Response(JSON.stringify(body), { + new Response(stream, { status: 200, - headers: { "Content-Type": "application/json" }, + headers: { "Content-Type": "text/event-stream" }, }) ) } +const defaultSseChunks = [ + { type: "text-delta", textDelta: "hello " }, + { type: "finish-message" }, +] + beforeEach(() => { vi.clearAllMocks() - mockFetch.mockImplementation(() => - okResponse({ answer: "hello", cited_ref_ids: [] }) - ) + mockFetch.mockImplementation(() => makeSseResponse(defaultSseChunks)) }) // ─── context field in POST body ────────────────────────────────────────────── @@ -91,16 +107,153 @@ describe("streamAgent – context in POST body", () => { expect(body.context?.selectedRefId).toBe("vid-1") }) - it("calls onDone with the API response answer", async () => { - mockFetch.mockImplementation(() => - okResponse({ answer: "Final answer.", cited_ref_ids: ["n1"] }) + it("includes stream: true in POST body", async () => { + const opts = makeOpts() + await streamAgent("Q?", opts) + + const [, fetchOpts] = mockFetch.mock.calls[0] as [string, RequestInit] + const body = JSON.parse(fetchOpts.body as string) + expect(body.stream).toBe(true) + }) + + it("sends Accept: text/event-stream header", async () => { + const opts = makeOpts() + await streamAgent("Q?", opts) + + const [, fetchOpts] = mockFetch.mock.calls[0] as [string, RequestInit] + const headers = fetchOpts.headers as Record + expect(headers["Accept"]).toBe("text/event-stream") + }) +}) + +// ─── SSE streaming behaviour ────────────────────────────────────────────────── + +describe("streamAgent – SSE streaming", () => { + it("text-delta chunk calls onChunk with the delta string", async () => { + const onChunk = vi.fn() + await streamAgent( + "Q?", + makeOpts({ + onChunk, + ...await makeSseResponse([ + { type: "text-delta", textDelta: "Hello!" }, + { type: "finish-message" }, + ]).then(() => ({})), + }) + ) + // Use a dedicated fetch mock for this test + const onChunk2 = vi.fn() + mockFetch.mockImplementationOnce(() => + makeSseResponse([ + { type: "text-delta", textDelta: "Hello!" }, + { type: "finish-message" }, + ]) + ) + await streamAgent("Q?", makeOpts({ onChunk: onChunk2 })) + expect(onChunk2).toHaveBeenCalledWith("Hello!") + }) + + it("tool-input-available calls onToolCall with status in-flight", async () => { + const onToolCall = vi.fn() + mockFetch.mockImplementationOnce(() => + makeSseResponse([ + { + type: "tool-input-available", + toolCallId: "tc-1", + toolName: "graph_search", + input: { q: "bitcoin" }, + }, + { type: "finish-message" }, + ]) + ) + await streamAgent("Q?", makeOpts({ onToolCall })) + expect(onToolCall).toHaveBeenCalledWith( + expect.objectContaining({ + id: "tc-1", + tool: "graph_search", + params: { q: "bitcoin" }, + status: "in-flight", + }) + ) + }) + + it("finish-step after tool-input-available calls onToolCall with status done", async () => { + const onToolCall = vi.fn() + mockFetch.mockImplementationOnce(() => + makeSseResponse([ + { + type: "tool-input-available", + toolCallId: "tc-2", + toolName: "graph_node", + input: { ref_id: "node-1" }, + }, + { type: "finish-step" }, + { type: "finish-message" }, + ]) ) + await streamAgent("Q?", makeOpts({ onToolCall })) + + const calls = onToolCall.mock.calls.map((c) => c[0]) + const inFlight = calls.find((c) => c.id === "tc-2" && c.status === "in-flight") + const done = calls.find((c) => c.id === "tc-2" && c.status === "done") + expect(inFlight).toBeDefined() + expect(done).toBeDefined() + }) + + it("finish-message calls onDone with accumulated text and empty cited_ref_ids", async () => { const onDone = vi.fn() + mockFetch.mockImplementationOnce(() => + makeSseResponse([ + { type: "text-delta", textDelta: "Foo " }, + { type: "text-delta", textDelta: "bar." }, + { type: "finish-message" }, + ]) + ) await streamAgent("Q?", makeOpts({ onDone })) + expect(onDone).toHaveBeenCalledWith({ answer: "Foo bar.", cited_ref_ids: [] }) + }) + + it("fallback onDone called when stream ends without finish-message", async () => { + const onDone = vi.fn() + mockFetch.mockImplementationOnce(() => + makeSseResponse([{ type: "text-delta", textDelta: "Partial." }]) + ) + await streamAgent("Q?", makeOpts({ onDone })) + expect(onDone).toHaveBeenCalledWith({ answer: "Partial.", cited_ref_ids: [] }) + }) +}) - expect(onDone).toHaveBeenCalledWith({ - answer: "Final answer.", - cited_ref_ids: ["n1"], +// ─── 402 retry flow ──────────────────────────────────────────────────────────── + +describe("streamAgent – 402 retry", () => { + it("pays L402 and retries on 402 response", async () => { + const { payL402 } = await import("@/lib/sphinx") + const onDone = vi.fn() + + mockFetch + .mockResolvedValueOnce(new Response(null, { status: 402 })) + .mockImplementationOnce(() => makeSseResponse(defaultSseChunks)) + + await streamAgent("Q?", makeOpts({ onDone })) + + expect(payL402).toHaveBeenCalled() + expect(onDone).toHaveBeenCalled() + }) +}) + +// ─── AbortSignal cancellation ───────────────────────────────────────────────── + +describe("streamAgent – AbortSignal", () => { + it("resolves without calling onError when aborted", async () => { + const onError = vi.fn() + const controller = new AbortController() + + mockFetch.mockImplementationOnce(() => { + controller.abort() + return Promise.reject(new DOMException("Aborted", "AbortError")) }) + + await streamAgent("Q?", makeOpts({ onError, signal: controller.signal })) + expect(onError).not.toHaveBeenCalled() }) }) diff --git a/src/lib/agent-api.ts b/src/lib/agent-api.ts index fa66200..fda0432 100644 --- a/src/lib/agent-api.ts +++ b/src/lib/agent-api.ts @@ -106,6 +106,83 @@ async function mockStreamAgent( opts.onDone({ answer, cited_ref_ids: ["mock-node-1", "mock-node-2"] }) } +async function processSSEStream(response: Response, opts: StreamAgentOpts): Promise { + if (!response.body) { + opts.onError(new Error("No response body for SSE stream")) + return + } + + const reader = response.body.getReader() + const decoder = new TextDecoder() + let buffer = "" + let accumulatedText = "" + const inFlight = new Map() + + try { + while (true) { + const { done, value } = await reader.read() + if (done) break + + buffer += decoder.decode(value, { stream: true }) + const events = buffer.split("\n\n") + // Keep the last (possibly incomplete) chunk in the buffer + buffer = events.pop() ?? "" + + for (const event of events) { + const line = event.trim() + if (!line.startsWith("data:")) continue + + const jsonStr = line.slice("data:".length).trim() + let chunk: Record + try { + chunk = JSON.parse(jsonStr) + } catch { + continue + } + + switch (chunk.type) { + case "text-delta": { + const delta = (chunk.textDelta ?? chunk.delta ?? "") as string + accumulatedText += delta + opts.onChunk(delta) + break + } + case "tool-input-available": { + const id = (chunk.toolCallId ?? `${chunk.toolName}-${Date.now()}`) as string + const event: ToolCallEvent = { + id, + tool: chunk.toolName as string, + params: (chunk.input ?? {}) as Record, + status: "in-flight", + } + inFlight.set(id, event) + opts.onToolCall(event) + break + } + case "finish-step": { + for (const stored of inFlight.values()) { + opts.onToolCall({ ...stored, status: "done" }) + } + inFlight.clear() + break + } + case "finish-message": { + opts.onDone({ answer: accumulatedText, cited_ref_ids: [] }) + return + } + } + } + } + // Fallback if stream ends without finish-message + opts.onDone({ answer: accumulatedText, cited_ref_ids: [] }) + } catch (err) { + if (err instanceof DOMException && err.name === "AbortError") return + opts.onError(err instanceof Error ? err : new Error(String(err))) + } finally { + reader.releaseLock() + } +} + export async function streamAgent( prompt: string, opts: StreamAgentOpts @@ -120,7 +197,7 @@ export async function streamAgent( const headers: Record = { "Content-Type": "application/json", - Accept: "application/json", + Accept: "text/event-stream", } if (l402) headers["Authorization"] = l402 @@ -131,6 +208,7 @@ export async function streamAgent( headers, body: JSON.stringify({ prompt, + stream: true, sessionId: opts.sessionId, ...(opts.context ? { context: opts.context } : {}), }), @@ -159,18 +237,7 @@ export async function streamAgent( return } - let data: { answer?: string; cited_ref_ids?: string[] } - try { - data = await response.json() - } catch { - opts.onError(new Error("Invalid JSON from agent")) - return - } - - opts.onDone({ - answer: data.answer ?? "", - cited_ref_ids: data.cited_ref_ids ?? [], - }) + await processSSEStream(response, opts) } return doRequest()