fix: fixed stream write callback

Author: lwin-kyaw

src/jrpc/v2/messageStream.ts

@@ -3,7 +3,7 @@ import { Duplex } from "readable-stream";
 
 import { isRequest } from "../../utils/jrpc";
 import { rpcErrors } from "../errors";
-import { JRPCRequest, Json } from "../interfaces";
+import { JRPCRequest } from "../interfaces";
 import { SafeEventEmitter } from "../safeEventEmitter";
 import { JRPCEngineV2 } from "./jrpcEngineV2";
 
@@ -21,18 +21,17 @@ export function createEngineStreamV2({ engine, notificationEmitter }: { engine:
     // noop
   }
 
-  function write(req: JRPCRequest<unknown>, _encoding: BufferEncoding, cb: (error?: Error | null) => void) {
-    engine
+  function handleRequest(req: JRPCRequest<unknown>) {
+    return engine
       .handle(req)
-      .then((res: Json | void): void => {
+      .then((res): undefined => {
         if (res !== undefined && isRequest(req)) {
           stream?.push({
             id: req.id,
             jsonrpc: "2.0",
             result: res,
           });
         }
-        cb();
         return undefined;
       })
       .catch((err: unknown) => {
@@ -45,10 +44,15 @@ export function createEngineStreamV2({ engine, notificationEmitter }: { engine:
           });
         }
         log.error(err);
-        cb(err as Error | null);
       });
   }
 
+  function write(req: JRPCRequest<unknown>, _encoding: BufferEncoding, cb: (error?: Error | null) => void) {
+    return handleRequest(req).finally(() => {
+      cb();
+    });
+  }
+
   stream = new Duplex({ objectMode: true, read: noop, write });
 
   if (notificationEmitter) {

src/jrpc/v2/providerUtils.ts

@@ -25,20 +25,24 @@ export function providerFromEngine(engine: JRPCEngineV2): SafeEventEmitterProvid
     return result as U;
   };
 
+  async function handleWithCallback<T extends JRPCParams, U>(req: JRPCRequest<T>, callback: (error: unknown, providerRes: JRPCResponse<U>) => void) {
+    try {
+      const result = await engine.handle(req as JRPCRequest);
+      callback(null, { id: req.id, jsonrpc: "2.0", result: result as U });
+    } catch (error) {
+      const serializedError = serializeJrpcError(error, {
+        shouldIncludeStack: false,
+        shouldPreserveMessage: true,
+      });
+      callback(serializedError, { id: req.id, jsonrpc: "2.0", error: serializedError });
+    }
+  }
+
   provider.send = <T extends JRPCParams, U>(req: JRPCRequest<T>, callback: (error: unknown, providerRes: JRPCResponse<U>) => void) => {
     if (typeof callback !== "function") {
       throw new Error('Must provide callback to "send" method.');
     }
-    engine
-      .handle(req as JRPCRequest)
-      .then((result) => callback(null, { id: req.id, jsonrpc: "2.0", result: result as U }))
-      .catch((error) => {
-        const serializedError = serializeJrpcError(error, {
-          shouldIncludeStack: false,
-          shouldPreserveMessage: true,
-        });
-        callback(error, { id: req.id, jsonrpc: "2.0", error: serializedError });
-      });
+    handleWithCallback(req, callback);
   };
 
   provider.request = async <T extends JRPCParams, U>(args: RequestArguments<T>) => {

test/v2/messageStream.test.ts

@@ -180,6 +180,16 @@ describe("createEngineStreamV2", () => {
       expect((pushed.error as Record<string, unknown>).message).toContain("specific failure");
     });
 
+    it("calls cb() without an error to keep the stream alive", async () => {
+      setup(createMockEngine(vi.fn().mockRejectedValue(new Error("engine failure"))));
+
+      const cb = vi.fn();
+      capturedWrite(makeRequest({ id: "err-alive" }), "utf-8", cb);
+      await flushPromises();
+
+      expect(cb).toHaveBeenCalledOnce();
+    });
+
     it("uses a fallback message for non-Error throws", async () => {
       setup(createMockEngine(vi.fn().mockRejectedValue("string error")));
 

test/v2/providerUtils.test.ts

@@ -134,10 +134,10 @@ describe("providerUtils", () => {
         const req = makeRequest({ id: "42" });
         const { error, response } = await sendPromise(provider, req);
 
-        expect(error).toBe(testError);
+        expect(error).toMatchObject({ code: -32603, message: "send error" });
+        expect(error).toBe(response.error);
         expect(response.id).toBe("42");
         expect(response.jsonrpc).toBe("2.0");
-        expect(response.error).toMatchObject({ message: "send error" });
       });
 
       it("constructs a proper JRPCResponse shape", async () => {
@@ -202,7 +202,7 @@ describe("providerUtils", () => {
 
         const { error, response } = await sendPromise(provider, makeRequest({ id: "x" }));
 
-        expect(error).toBe(thrown);
+        expect(error).toBe(response.error);
         expect(response.id).toBe("x");
         expect(response.jsonrpc).toBe("2.0");
         const rpcError = response.error as Record<string, unknown>;
@@ -323,9 +323,10 @@ describe("providerUtils", () => {
       };
       const provider = providerFromMiddleware(middleware as JRPCMiddlewareV2);
 
-      const { error } = await sendPromise(provider, makeRequest());
+      const { error, response } = await sendPromise(provider, makeRequest());
 
-      expect(error).toBe(testError);
+      expect(error).toMatchObject({ code: -32603, message: "middleware-error" });
+      expect(error).toBe(response.error);
     });
 
     it("works with a middleware that uses context", async () => {