diff --git a/src/server.test.ts b/src/server.test.ts index d1efb3f..7061956 100644 --- a/src/server.test.ts +++ b/src/server.test.ts @@ -1,5 +1,13 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { stripThinkingTokens, getProxyUrl, proxyAwareFetch, validateMessages } from "./server.js"; +import { + stripThinkingTokens, + getProxyUrl, + proxyAwareFetch, + validateMessages, + consumeSSEStream, + createMcpProgressReporter, + SSE_HEARTBEAT_MS, +} from "./server.js"; describe("Server Utility Functions", () => { describe("stripThinkingTokens", () => { @@ -253,4 +261,126 @@ describe("Server Utility Functions", () => { ], "test_tool")).toThrow("Invalid message at index 2: 'content' must be a string"); }); }); + + describe("consumeSSEStream", () => { + function sseResponse(chunks: string[]): Response { + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(encoder.encode(chunk)); + } + controller.close(); + }, + }); + return new Response(stream, { status: 200 }); + } + + it("should assemble streamed content deltas", async () => { + const response = sseResponse([ + 'data: {"choices":[{"delta":{"content":"Hello"}}]}\n\n', + 'data: {"choices":[{"delta":{"content":" world"}}]}\n\n', + "data: [DONE]\n\n", + ]); + + const result = await consumeSSEStream(response); + expect(result.choices[0].message.content).toBe("Hello world"); + }); + + it("should invoke onProgress for each content delta", async () => { + const response = sseResponse([ + 'data: {"choices":[{"delta":{"content":"a"}}]}\n\n', + 'data: {"choices":[{"delta":{"content":"b"}}]}\n\n', + ]); + const progressCalls: Array<{ chunkCount: number; totalChars: number }> = []; + + await consumeSSEStream(response, async (stats) => { + progressCalls.push(stats); + }); + + expect(progressCalls).toEqual([ + { chunkCount: 1, totalChars: 1 }, + { chunkCount: 2, totalChars: 2 }, + ]); + }); + + it("should emit heartbeat progress during silent SSE gaps", async () => { + vi.useFakeTimers(); + try { + const encoder = new TextEncoder(); + let controllerRef: ReadableStreamDefaultController | undefined; + const stream = new ReadableStream({ + start(controller) { + controllerRef = controller; + }, + }); + const response = new Response(stream, { status: 200 }); + const progressCalls: number[] = []; + + const pending = consumeSSEStream(response, async () => { + progressCalls.push(Date.now()); + }); + + await vi.advanceTimersByTimeAsync(SSE_HEARTBEAT_MS); + expect(progressCalls.length).toBeGreaterThanOrEqual(1); + + controllerRef?.enqueue( + encoder.encode('data: {"choices":[{"delta":{"content":"done"}}]}\n\n'), + ); + controllerRef?.close(); + await pending; + + expect(progressCalls.length).toBeGreaterThanOrEqual(2); + } finally { + vi.useRealTimers(); + } + }); + }); + + describe("createMcpProgressReporter", () => { + it("should return undefined when progressToken is missing", () => { + const reporter = createMcpProgressReporter({ + sendNotification: vi.fn().mockResolvedValue(undefined), + }); + expect(reporter).toBeUndefined(); + }); + + it("should emit notifications/progress with incrementing counter", async () => { + const sendNotification = vi.fn().mockResolvedValue(undefined); + const reporter = createMcpProgressReporter({ + _meta: { progressToken: "token-1" }, + sendNotification, + }); + expect(reporter).toBeDefined(); + + await reporter!({ chunkCount: 2, totalChars: 40 }); + await reporter!({ chunkCount: 5, totalChars: 120 }); + + expect(sendNotification).toHaveBeenCalledTimes(2); + expect(sendNotification.mock.calls[0][0]).toEqual({ + method: "notifications/progress", + params: { + progressToken: "token-1", + progress: 1, + message: "Streaming response… 2 chunks, 40 chars streamed", + }, + }); + expect(sendNotification.mock.calls[1][0].params.progress).toBe(2); + }); + + it("should use a custom progress label", async () => { + const sendNotification = vi.fn().mockResolvedValue(undefined); + const reporter = createMcpProgressReporter( + { + _meta: { progressToken: "token-2" }, + sendNotification, + }, + "Reasoning", + ); + + await reporter!({ chunkCount: 1, totalChars: 10 }); + + expect(sendNotification.mock.calls[0][0].params.message).toContain("Reasoning"); + }); + }); }); diff --git a/src/server.ts b/src/server.ts index d649ccb..69a2cb2 100644 --- a/src/server.ts +++ b/src/server.ts @@ -118,7 +118,55 @@ async function makeApiRequest( return response; } -export async function consumeSSEStream(response: Response): Promise { +export type SSEProgressStats = { + chunkCount: number; + totalChars: number; +}; + +export type SSEProgressCallback = (stats: SSEProgressStats) => void | Promise; + +/** Heartbeat interval while the API is in a silent search/reasoning phase (no SSE deltas). */ +export const SSE_HEARTBEAT_MS = 10_000; + +type McpProgressExtra = { + _meta?: { progressToken?: string | number }; + sendNotification: (notification: { + method: "notifications/progress"; + params: { + progressToken: string | number; + progress: number; + message?: string; + }; + }) => Promise; +}; + +export function createMcpProgressReporter( + extra: McpProgressExtra, + label = "Streaming response", +): SSEProgressCallback | undefined { + const progressToken = extra._meta?.progressToken; + if (progressToken === undefined) { + return undefined; + } + + let progress = 0; + return async ({ chunkCount, totalChars }) => { + progress += 1; + await extra.sendNotification({ + method: "notifications/progress", + params: { + progressToken, + progress, + message: `${label}… ${chunkCount} chunks, ${totalChars} chars streamed`, + }, + }); + }; +} + +export async function consumeSSEStream( + response: Response, + onProgress?: SSEProgressCallback, +): Promise { const body = response.body; if (!body) { throw new Error("Response body is null"); @@ -134,41 +182,65 @@ export async function consumeSSEStream(response: Response): Promise { + if (!onProgress) { + return; + } + await onProgress({ chunkCount, totalChars }); + }; - if (parsed.id) id = parsed.id; - if (parsed.model) model = parsed.model; - if (parsed.created) created = parsed.created; - if (parsed.citations) citations = parsed.citations; - if (parsed.usage) usage = parsed.usage; + const heartbeatTimer = onProgress + ? setInterval(() => { + void emitProgress(); + }, SSE_HEARTBEAT_MS) + : undefined; - const delta = parsed.choices?.[0]?.delta; - if (delta?.content) { - contentParts.push(delta.content); + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + buffer += decoder.decode(value, { stream: true }); + + const lines = buffer.split("\n"); + // Keep the last potentially incomplete line in the buffer + buffer = lines.pop() || ""; + + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed || !trimmed.startsWith("data:")) continue; + + const data = trimmed.slice("data:".length).trim(); + if (data === "[DONE]") continue; + + try { + const parsed = JSON.parse(data); + + if (parsed.id) id = parsed.id; + if (parsed.model) model = parsed.model; + if (parsed.created) created = parsed.created; + if (parsed.citations) citations = parsed.citations; + if (parsed.usage) usage = parsed.usage; + + const delta = parsed.choices?.[0]?.delta; + if (delta?.content) { + contentParts.push(delta.content); + chunkCount += 1; + totalChars += delta.content.length; + await emitProgress(); + } + } catch { + // Skip malformed JSON chunks (e.g. keep-alive pings) } - } catch { - // Skip malformed JSON chunks (e.g. keep-alive pings) } } + } finally { + if (heartbeatTimer) { + clearInterval(heartbeatTimer); + } } const assembled: ChatCompletionResponse = { @@ -194,7 +266,8 @@ export async function performChatCompletion( model: string = "sonar-pro", stripThinking: boolean = false, serviceOrigin?: string, - options?: ChatCompletionOptions + options?: ChatCompletionOptions, + onProgress?: SSEProgressCallback, ): Promise { const useStreaming = model === "sonar-deep-research"; @@ -213,7 +286,7 @@ export async function performChatCompletion( let data: ChatCompletionResponse; try { if (useStreaming) { - data = await consumeSSEStream(response); + data = await consumeSSEStream(response, onProgress); } else { const json = await response.json(); data = ChatCompletionResponseSchema.parse(json); @@ -420,7 +493,7 @@ export function createPerplexityServer(serviceOrigin?: string) { destructiveHint: false, }, }, - async (args: any) => { + async (args: any, extra: McpProgressExtra) => { const { messages, strip_thinking, reasoning_effort } = args as { messages: Message[]; strip_thinking?: boolean; @@ -431,7 +504,15 @@ export function createPerplexityServer(serviceOrigin?: string) { const options = { ...(reasoning_effort && { reasoning_effort }), }; - const result = await performChatCompletion(messages, "sonar-deep-research", stripThinking, serviceOrigin, Object.keys(options).length > 0 ? options : undefined); + const onProgress = createMcpProgressReporter(extra, "Deep research"); + const result = await performChatCompletion( + messages, + "sonar-deep-research", + stripThinking, + serviceOrigin, + Object.keys(options).length > 0 ? options : undefined, + onProgress, + ); return { content: [{ type: "text" as const, text: result }], structuredContent: { response: result }, @@ -473,7 +554,15 @@ export function createPerplexityServer(serviceOrigin?: string) { ...(search_domain_filter && { search_domain_filter }), ...(search_context_size && { search_context_size }), }; - const result = await performChatCompletion(messages, "sonar-reasoning-pro", stripThinking, serviceOrigin, Object.keys(options).length > 0 ? options : undefined); + const onProgress = createMcpProgressReporter(extra, "Reasoning"); + const result = await performChatCompletion( + messages, + "sonar-reasoning-pro", + stripThinking, + serviceOrigin, + Object.keys(options).length > 0 ? options : undefined, + onProgress, + ); return { content: [{ type: "text" as const, text: result }], structuredContent: { response: result },