fix: emit MCP progress for perplexity_reason SSE (follow-up to #115)

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<Uint8Array>({
+        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<Uint8Array> | undefined;
+        const stream = new ReadableStream<Uint8Array>({
+          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");
+    });
+  });
 });

src/server.ts

@@ -118,7 +118,55 @@ async function makeApiRequest(
   return response;
 }
 
-export async function consumeSSEStream(response: Response): Promise<ChatCompletionResponse> {
+export type SSEProgressStats = {
+  chunkCount: number;
+  totalChars: number;
+};
+
+export type SSEProgressCallback = (stats: SSEProgressStats) => void | Promise<void>;
+
+/** 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<void>;
+};
+
+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<ChatCompletionResponse> {
   const body = response.body;
   if (!body) {
     throw new Error("Response body is null");
@@ -134,41 +182,65 @@ export async function consumeSSEStream(response: Response): Promise<ChatCompleti
   let model: string | undefined;
   let created: number | undefined;
   let buffer = "";
+  let chunkCount = 0;
+  let totalChars = 0;
 
-  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);
+  const emitProgress = async () => {
+    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<string> {
   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 },