@synnaxlabs/client 0.49.3 → 0.52.3

package/src/arc/lsp.spec.ts

@@ -1,4 +1,4 @@
-// Copyright 2025 Synnax Labs, Inc.
+// Copyright 2026 Synnax Labs, Inc.
 //
 // Use of this software is governed by the Business Source License included in the file
 // licenses/BSL.txt.
@@ -7,52 +7,91 @@
 // License, use of this software will be governed by the Apache License, Version 2.0,
 // included in the file licenses/APL.txt.
 
-import { type jsonRPC } from "@synnaxlabs/x";
 import { describe, expect, it } from "vitest";
 
 import { createTestClient } from "@/testutil/client";
 
+interface JSONRPCRequest {
+  jsonrpc: "2.0";
+  id: number;
+  method: string;
+  params?: unknown;
+}
+
+type JSONRPCResponse =
+  | { jsonrpc: "2.0"; id: number; result: unknown }
+  | {
+      jsonrpc: "2.0";
+      id: number;
+      error: { code: number; message: string; data?: unknown };
+    };
+
+type LSPReceiver = {
+  receive: () => Promise<[{ content: string }, null] | [null, Error]>;
+};
+
+const MAX_DRAIN = 50;
+
+/** Drains messages from the stream until a JSON-RPC response with the expected id arrives. */
+const receiveResponse = async (
+  stream: LSPReceiver,
+  expectedId: number,
+): Promise<JSONRPCResponse> => {
+  for (let i = 0; i < MAX_DRAIN; i++) {
+    const [res, err] = await stream.receive();
+    if (err != null) throw err;
+    if (res == null) throw new Error("Expected response");
+    const msg = JSON.parse(res.content);
+    if (!("method" in msg) && "id" in msg && msg.id === expectedId)
+      return msg as JSONRPCResponse;
+  }
+  throw new Error(
+    `receiveResponse: drained ${MAX_DRAIN} messages without seeing id=${expectedId}`,
+  );
+};
+
+/** Drains messages from the stream until a JSON-RPC notification with the expected method arrives. */
+const receiveNotification = async (
+  stream: LSPReceiver,
+  expectedMethod: string,
+): Promise<JSONRPCRequest> => {
+  for (let i = 0; i < MAX_DRAIN; i++) {
+    const [res, err] = await stream.receive();
+    if (err != null) throw err;
+    if (res == null) throw new Error("Expected message");
+    const msg = JSON.parse(res.content);
+    if ("method" in msg && msg.method === expectedMethod) return msg as JSONRPCRequest;
+  }
+  throw new Error(
+    `receiveNotification: drained ${MAX_DRAIN} messages without seeing method=${expectedMethod}`,
+  );
+};
+
 describe("Arc LSP", () => {
   it("should open an LSP stream and handle initialize request", async () => {
     const client = createTestClient();
     const stream = await client.arcs.openLSP();
 
-    // Send LSP initialize request (raw JSON, no Content-Length headers)
-    const initializeRequest: jsonRPC.Request = {
+    const initializeRequest: JSONRPCRequest = {
       jsonrpc: "2.0",
       id: 1,
       method: "initialize",
       params: {
         processId: null,
-        clientInfo: {
-          name: "test-client",
-          version: "1.0.0",
-        },
+        clientInfo: { name: "test-client", version: "1.0.0" },
         rootUri: null,
         capabilities: {},
       },
     };
 
-    stream.send({
-      content: JSON.stringify(initializeRequest),
-    });
-
-    // Receive response (raw JSON, no Content-Length headers)
-    const [res, err] = await stream.receive();
-    expect(err).toBeNull();
-    if (!res) throw new Error("Expected response");
+    stream.send({ content: JSON.stringify(initializeRequest) });
 
-    // Parse raw JSON response directly
-    const msg = JSON.parse(res.content) as jsonRPC.Response;
+    const msg = await receiveResponse(stream, 1);
     expect(msg.jsonrpc).toBe("2.0");
     expect(msg.id).toBe(1);
-
-    // Check for error response
     if ("error" in msg) throw new Error(`LSP error: ${msg.error.message}`);
-
     expect("result" in msg).toBe(true);
 
-    // Verify capabilities are present
     if ("result" in msg) {
       const result = msg.result as Record<string, unknown>;
       expect(result).toHaveProperty("capabilities");
@@ -65,8 +104,7 @@ describe("Arc LSP", () => {
     const client = createTestClient();
     const stream = await client.arcs.openLSP();
 
-    // First initialize
-    const initializeRequest: jsonRPC.Request = {
+    const initializeRequest: JSONRPCRequest = {
       jsonrpc: "2.0",
       id: 1,
       method: "initialize",
@@ -79,62 +117,37 @@ describe("Arc LSP", () => {
     };
 
     stream.send({ content: JSON.stringify(initializeRequest) });
-    const [initResponse, initErr] = await stream.receive();
-    expect(initErr).toBeNull();
-    expect(initResponse).not.toBeNull();
-    if (!initResponse) throw new Error("Expected response");
-
-    // Parse raw JSON response
-    const initMsg = JSON.parse(initResponse.content) as jsonRPC.Response;
-    expect(initMsg.id).toBe(1);
-
-    // Check for error response
+    const initMsg = await receiveResponse(stream, 1);
     if ("error" in initMsg) throw new Error(`LSP error: ${initMsg.error.message}`);
 
-    // Send initialized notification
-    const initializedNotification: jsonRPC.Request = {
-      jsonrpc: "2.0",
-      method: "initialized",
-      params: {},
-    };
-
-    stream.send({ content: JSON.stringify(initializedNotification) });
+    stream.send({
+      content: JSON.stringify({
+        jsonrpc: "2.0",
+        method: "initialized",
+        params: {},
+      }),
+    });
 
-    // Send didOpen notification
-    const didOpenNotification: jsonRPC.Request = {
-      jsonrpc: "2.0",
-      method: "textDocument/didOpen",
-      params: {
-        textDocument: {
-          uri: "file:///test.arc",
-          languageId: "arc",
-          version: 1,
-          text: "let x = 1 + 2;",
+    stream.send({
+      content: JSON.stringify({
+        jsonrpc: "2.0",
+        method: "textDocument/didOpen",
+        params: {
+          textDocument: {
+            uri: "file:///test.arc",
+            languageId: "arc",
+            version: 1,
+            text: "let x = 1 + 2;",
+          },
         },
-      },
-    };
-
-    stream.send({ content: JSON.stringify(didOpenNotification) });
+      }),
+    });
 
-    // Notifications don't get responses, but we might get diagnostics back
-    // Wait a bit to see if we get any messages
-    const diagnosticsPromise = stream.receive();
-    const timeoutPromise = new Promise((resolve) =>
-      setTimeout(() => resolve(null), 100),
+    const diagMsg = await receiveNotification(
+      stream,
+      "textDocument/publishDiagnostics",
     );
-
-    const result = await Promise.race([diagnosticsPromise, timeoutPromise]);
-
-    if (result && Array.isArray(result)) {
-      const [diagResponse, diagErr] = result;
-      if (diagResponse && !diagErr) {
-        // Parse raw JSON notification
-        const diagMsg = JSON.parse(diagResponse.content) as jsonRPC.Message;
-        expect(diagMsg.jsonrpc).toBe("2.0");
-        if ("method" in diagMsg)
-          expect(diagMsg.method).toBe("textDocument/publishDiagnostics");
-      }
-    }
+    expect(diagMsg.jsonrpc).toBe("2.0");
 
     stream.closeSend();
     client.close();
@@ -144,8 +157,7 @@ describe("Arc LSP", () => {
     const client = createTestClient();
     const stream = await client.arcs.openLSP();
 
-    // Initialize
-    const initializeRequest: jsonRPC.Request = {
+    const initializeRequest: JSONRPCRequest = {
       jsonrpc: "2.0",
       id: 1,
       method: "initialize",
@@ -158,14 +170,9 @@ describe("Arc LSP", () => {
     };
 
     stream.send({ content: JSON.stringify(initializeRequest) });
-    const [initRes] = await stream.receive();
-    expect(initRes).not.toBeNull();
-    if (!initRes) throw new Error("Expected response");
-
-    const initMsg = JSON.parse(initRes.content) as jsonRPC.Response;
+    const initMsg = await receiveResponse(stream, 1);
     if ("error" in initMsg) throw new Error(`LSP error: ${initMsg.error.message}`);
 
-    // Send initialized notification
     stream.send({
       content: JSON.stringify({
         jsonrpc: "2.0",
@@ -174,7 +181,6 @@ describe("Arc LSP", () => {
       }),
     });
 
-    // Open a document
     stream.send({
       content: JSON.stringify({
         jsonrpc: "2.0",
@@ -190,12 +196,9 @@ describe("Arc LSP", () => {
       }),
     });
 
-    // Wait for and consume the publishDiagnostics notification
-    const [diagResponse] = await stream.receive();
-    expect(diagResponse).not.toBeNull();
+    await receiveNotification(stream, "textDocument/publishDiagnostics");
 
-    // Request hover information
-    const hoverRequest: jsonRPC.Request = {
+    const hoverRequest: JSONRPCRequest = {
       jsonrpc: "2.0",
       id: 2,
       method: "textDocument/hover",
@@ -207,17 +210,9 @@ describe("Arc LSP", () => {
 
     stream.send({ content: JSON.stringify(hoverRequest) });
 
-    const [hoverResponse, hoverErr] = await stream.receive();
-    expect(hoverErr).toBeNull();
-    expect(hoverResponse).not.toBeNull();
-    if (!hoverResponse) throw new Error("Expected response");
-
-    // Parse raw JSON response
-    const responseMsg = JSON.parse(hoverResponse.content) as jsonRPC.Response;
+    const responseMsg = await receiveResponse(stream, 2);
     expect(responseMsg.jsonrpc).toBe("2.0");
     expect(responseMsg.id).toBe(2);
-
-    // Check for error response
     if ("error" in responseMsg)
       throw new Error(`LSP error: ${responseMsg.error.message}`);
 
@@ -229,10 +224,9 @@ describe("Arc LSP", () => {
     const client = createTestClient();
     const stream = await client.arcs.openLSP();
 
-    const receivedMessages: jsonRPC.Response[] = [];
+    const receivedMessages: JSONRPCResponse[] = [];
 
-    // Initialize
-    const initializeRequest: jsonRPC.Request = {
+    const initializeRequest: JSONRPCRequest = {
       jsonrpc: "2.0",
       id: 1,
       method: "initialize",
@@ -245,14 +239,11 @@ describe("Arc LSP", () => {
     };
 
     stream.send({ content: JSON.stringify(initializeRequest) });
-    const [initResponse] = await stream.receive();
-    if (!initResponse) throw new Error("Expected response");
-    const initMsg = JSON.parse(initResponse.content) as jsonRPC.Response;
+    const initMsg = await receiveResponse(stream, 1);
     if ("error" in initMsg) throw new Error(`LSP error: ${initMsg.error.message}`);
 
     receivedMessages.push(initMsg);
 
-    // Send multiple requests with different IDs
     const requests = [
       {
         jsonrpc: "2.0",
@@ -270,20 +261,14 @@ describe("Arc LSP", () => {
 
     for (const req of requests) stream.send({ content: JSON.stringify(req) });
 
-    // Receive responses (order may vary)
-    for (let i = 0; i < requests.length; i++) {
-      const [response] = await stream.receive();
-      if (!response) throw new Error("Expected response");
-      const msg = JSON.parse(response.content) as jsonRPC.Response;
+    for (const req of requests) {
+      const msg = await receiveResponse(stream, req.id);
       if ("error" in msg) throw new Error(`LSP error: ${msg.error.message}`);
-
       receivedMessages.push(msg);
     }
 
-    // Should have init response + 2 request responses
     expect(receivedMessages.length).toBeGreaterThanOrEqual(3);
 
-    // Verify all responses have correct structure
     for (const msg of receivedMessages) {
       expect(msg.jsonrpc).toBe("2.0");
       expect(msg.id).toBeDefined();
@@ -297,31 +282,19 @@ describe("Arc LSP", () => {
     const client = createTestClient();
     const stream = await client.arcs.openLSP();
 
-    // Test that our message format is just JSON without Content-Length headers
-    const testMessage: jsonRPC.Request = {
+    const testMessage: JSONRPCRequest = {
       jsonrpc: "2.0",
       id: 999,
       method: "test/method",
       params: { data: "test" },
     };
 
-    const messageContent = JSON.stringify(testMessage);
+    stream.send({ content: JSON.stringify(testMessage) });
 
-    // Send raw JSON (no headers)
-    stream.send({ content: messageContent });
-
-    // The server should respond (even if it's an error for unknown method)
-    const [response, err] = await stream.receive();
-    expect(err).toBeNull();
-    expect(response).not.toBeNull();
-    if (!response) throw new Error("Expected response");
-
-    // Response should be parseable JSON
-    const parsed = JSON.parse(response.content) as jsonRPC.Response;
+    const parsed = await receiveResponse(stream, 999);
     expect(parsed.jsonrpc).toBe("2.0");
     expect(parsed.id).toBe(999);
 
-    // This test expects an error response for unknown method
     if ("error" in parsed) expect(parsed.error).toBeDefined();
 
     stream.closeSend();
@@ -332,8 +305,7 @@ describe("Arc LSP", () => {
     const client = createTestClient();
     const stream = await client.arcs.openLSP();
 
-    // Initialize with semantic tokens capability
-    const initializeRequest: jsonRPC.Request = {
+    const initializeRequest: JSONRPCRequest = {
       jsonrpc: "2.0",
       id: 1,
       method: "initialize",
@@ -353,21 +325,15 @@ describe("Arc LSP", () => {
     };
 
     stream.send({ content: JSON.stringify(initializeRequest) });
-    const [initRes] = await stream.receive();
-    expect(initRes).not.toBeNull();
-    if (!initRes) throw new Error("Expected response");
-
-    const initMsg = JSON.parse(initRes.content) as jsonRPC.Response;
+    const initMsg = await receiveResponse(stream, 1);
     if ("error" in initMsg) throw new Error(`LSP error: ${initMsg.error.message}`);
 
-    // Verify server advertises semantic tokens support
     if ("result" in initMsg) {
       const result = initMsg.result as Record<string, unknown>;
       const capabilities = result.capabilities as Record<string, unknown>;
       expect(capabilities).toHaveProperty("semanticTokensProvider");
     }
 
-    // Send initialized notification
     stream.send({
       content: JSON.stringify({
         jsonrpc: "2.0",
@@ -376,7 +342,6 @@ describe("Arc LSP", () => {
       }),
     });
 
-    // Open a document with Arc code
     const arcCode = "func add(x i32, y i32) i32 {\n  return x + y\n}";
     stream.send({
       content: JSON.stringify({
@@ -393,11 +358,9 @@ describe("Arc LSP", () => {
       }),
     });
 
-    // Wait for diagnostics
-    await stream.receive();
+    await receiveNotification(stream, "textDocument/publishDiagnostics");
 
-    // Request semantic tokens
-    const semanticTokensRequest: jsonRPC.Request = {
+    const semanticTokensRequest: JSONRPCRequest = {
       jsonrpc: "2.0",
       id: 2,
       method: "textDocument/semanticTokens/full",
@@ -408,26 +371,18 @@ describe("Arc LSP", () => {
 
     stream.send({ content: JSON.stringify(semanticTokensRequest) });
 
-    const [tokenResponse, tokenErr] = await stream.receive();
-    expect(tokenErr).toBeNull();
-    expect(tokenResponse).not.toBeNull();
-    if (!tokenResponse) throw new Error("Expected response");
-
-    const tokenMsg = JSON.parse(tokenResponse.content) as jsonRPC.Response;
+    const tokenMsg = await receiveResponse(stream, 2);
     expect(tokenMsg.jsonrpc).toBe("2.0");
     expect(tokenMsg.id).toBe(2);
-
     if ("error" in tokenMsg) throw new Error(`LSP error: ${tokenMsg.error.message}`);
 
-    // Verify semantic tokens are returned
     if ("result" in tokenMsg) {
       const result = tokenMsg.result as Record<string, unknown>;
       expect(result).toHaveProperty("data");
       const data = result.data as number[];
       expect(Array.isArray(data)).toBe(true);
-      // Should have tokens (encoded as [deltaLine, deltaStart, length, type, modifiers] * N)
       expect(data.length).toBeGreaterThan(0);
-      // Tokens should be in groups of 5
+      // Tokens are encoded as [deltaLine, deltaStart, length, type, modifiers] * N
       expect(data.length % 5).toBe(0);
     }
 
@@ -440,8 +395,7 @@ describe("Arc LSP", () => {
       const client = createTestClient();
       const stream = await client.arcs.openLSP();
 
-      // Initialize
-      const initializeRequest: jsonRPC.Request = {
+      const initializeRequest: JSONRPCRequest = {
         jsonrpc: "2.0",
         id: 1,
         method: "initialize",
@@ -454,14 +408,9 @@ describe("Arc LSP", () => {
       };
 
       stream.send({ content: JSON.stringify(initializeRequest) });
-      const [initRes] = await stream.receive();
-      expect(initRes).not.toBeNull();
-      if (!initRes) throw new Error("Expected response");
-
-      const initMsg = JSON.parse(initRes.content) as jsonRPC.Response;
+      const initMsg = await receiveResponse(stream, 1);
       if ("error" in initMsg) throw new Error(`LSP error: ${initMsg.error.message}`);
 
-      // Send initialized notification
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -470,13 +419,10 @@ describe("Arc LSP", () => {
         }),
       });
 
-      // Create metadata for block expression
       const metadata = { is_block: true };
       const encoded = btoa(JSON.stringify(metadata));
       const blockURI = `arc://block/test123#${encoded}`;
 
-      // Open a document with block URI containing a simple expression
-      const blockExpression = "return x * 2";
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -486,30 +432,22 @@ describe("Arc LSP", () => {
               uri: blockURI,
               languageId: "arc",
               version: 1,
-              text: blockExpression,
+              text: "return x * 2",
             },
           },
         }),
       });
 
-      // Should receive diagnostics
-      const [diagResponse, diagErr] = await stream.receive();
-      expect(diagErr).toBeNull();
-      expect(diagResponse).not.toBeNull();
-      if (!diagResponse) throw new Error("Expected diagnostics");
-
-      const diagMsg = JSON.parse(diagResponse.content) as jsonRPC.Message;
+      const diagMsg = await receiveNotification(
+        stream,
+        "textDocument/publishDiagnostics",
+      );
       expect(diagMsg.jsonrpc).toBe("2.0");
-      if ("method" in diagMsg) {
-        expect(diagMsg.method).toBe("textDocument/publishDiagnostics");
-        if ("params" in diagMsg) {
-          const params = diagMsg.params as Record<string, unknown>;
-          expect(params.uri).toBe(blockURI);
-          // Expression should be wrapped and parsed successfully
-          const diagnostics = params.diagnostics as unknown[];
-          // May have diagnostics about undefined 'x', but should parse successfully
-          expect(Array.isArray(diagnostics)).toBe(true);
-        }
+      if ("params" in diagMsg) {
+        const params = diagMsg.params as Record<string, unknown>;
+        expect(params.uri).toBe(blockURI);
+        const diagnostics = params.diagnostics as unknown[];
+        expect(Array.isArray(diagnostics)).toBe(true);
       }
 
       stream.closeSend();
@@ -520,8 +458,7 @@ describe("Arc LSP", () => {
       const client = createTestClient();
       const stream = await client.arcs.openLSP();
 
-      // Initialize
-      const initializeRequest: jsonRPC.Request = {
+      const initializeRequest: JSONRPCRequest = {
         jsonrpc: "2.0",
         id: 1,
         method: "initialize",
@@ -534,14 +471,9 @@ describe("Arc LSP", () => {
       };
 
       stream.send({ content: JSON.stringify(initializeRequest) });
-      const [initRes] = await stream.receive();
-      expect(initRes).not.toBeNull();
-      if (!initRes) throw new Error("Expected response");
-
-      const initMsg = JSON.parse(initRes.content) as jsonRPC.Response;
+      const initMsg = await receiveResponse(stream, 1);
       if ("error" in initMsg) throw new Error(`LSP error: ${initMsg.error.message}`);
 
-      // Send initialized notification
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -550,13 +482,10 @@ describe("Arc LSP", () => {
         }),
       });
 
-      // Create block URI
       const metadata = { is_block: true };
       const encoded = btoa(JSON.stringify(metadata));
       const blockURI = `arc://block/syntax-error#${encoded}`;
 
-      // Open a document with invalid syntax
-      const invalidExpression = "return x +";
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -566,21 +495,17 @@ describe("Arc LSP", () => {
               uri: blockURI,
               languageId: "arc",
               version: 1,
-              text: invalidExpression,
+              text: "return x +",
             },
           },
         }),
       });
 
-      // Should receive diagnostics with syntax error
-      const [diagResponse, diagErr] = await stream.receive();
-      expect(diagErr).toBeNull();
-      expect(diagResponse).not.toBeNull();
-      if (!diagResponse) throw new Error("Expected diagnostics");
-
-      const diagMsg = JSON.parse(diagResponse.content) as jsonRPC.Message;
-      if ("method" in diagMsg && "params" in diagMsg) {
-        expect(diagMsg.method).toBe("textDocument/publishDiagnostics");
+      const diagMsg = await receiveNotification(
+        stream,
+        "textDocument/publishDiagnostics",
+      );
+      if ("params" in diagMsg) {
         const params = diagMsg.params as Record<string, unknown>;
         const diagnostics = params.diagnostics as Array<{
           range: { start: { line: number; character: number } };
@@ -588,15 +513,9 @@ describe("Arc LSP", () => {
           severity: number;
         }>;
 
-        // Should have at least one diagnostic
         expect(diagnostics.length).toBeGreaterThan(0);
-
-        // Position should be mapped correctly (line 0 for original expression)
-        // Note: Parse errors may not be position-mapped correctly yet
-        // This is a known limitation - for now we just verify diagnostics exist
-        const firstDiag = diagnostics[0];
-        console.log("Diagnostic position:", firstDiag.range.start);
-        // expect(firstDiag.range.start.line).toBe(0);
+        // Position should be mapped to line 0 for original expression
+        expect(diagnostics[0].range.start.line).toBe(0);
       }
 
       stream.closeSend();
@@ -607,8 +526,7 @@ describe("Arc LSP", () => {
       const client = createTestClient();
       const stream = await client.arcs.openLSP();
 
-      // Initialize
-      const initializeRequest: jsonRPC.Request = {
+      const initializeRequest: JSONRPCRequest = {
         jsonrpc: "2.0",
         id: 1,
         method: "initialize",
@@ -621,14 +539,9 @@ describe("Arc LSP", () => {
       };
 
       stream.send({ content: JSON.stringify(initializeRequest) });
-      const [initRes] = await stream.receive();
-      expect(initRes).not.toBeNull();
-      if (!initRes) throw new Error("Expected response");
-
-      const initMsg = JSON.parse(initRes.content) as jsonRPC.Response;
+      const initMsg = await receiveResponse(stream, 1);
       if ("error" in initMsg) throw new Error(`LSP error: ${initMsg.error.message}`);
 
-      // Send initialized notification
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -637,14 +550,10 @@ describe("Arc LSP", () => {
         }),
       });
 
-      // Create block URI
       const metadata = { is_block: true };
       const encoded = btoa(JSON.stringify(metadata));
       const blockURI = `arc://block/multiline#${encoded}`;
 
-      // Multi-line block expression
-      const multiLineExpression =
-        "let temp = x * 2\nlet result = temp + 1\nreturn result";
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -654,25 +563,20 @@ describe("Arc LSP", () => {
               uri: blockURI,
               languageId: "arc",
               version: 1,
-              text: multiLineExpression,
+              text: "let temp = x * 2\nlet result = temp + 1\nreturn result",
             },
           },
         }),
       });
 
-      // Should receive diagnostics
-      const [diagResponse, diagErr] = await stream.receive();
-      expect(diagErr).toBeNull();
-      expect(diagResponse).not.toBeNull();
-      if (!diagResponse) throw new Error("Expected diagnostics");
-
-      const diagMsg = JSON.parse(diagResponse.content) as jsonRPC.Message;
-      if ("method" in diagMsg && "params" in diagMsg) {
-        expect(diagMsg.method).toBe("textDocument/publishDiagnostics");
+      const diagMsg = await receiveNotification(
+        stream,
+        "textDocument/publishDiagnostics",
+      );
+      if ("params" in diagMsg) {
         const params = diagMsg.params as Record<string, unknown>;
         expect(params.uri).toBe(blockURI);
         const diagnostics = params.diagnostics as unknown[];
-        // Should successfully parse multi-line expression
         expect(Array.isArray(diagnostics)).toBe(true);
       }
 
@@ -684,8 +588,7 @@ describe("Arc LSP", () => {
       const client = createTestClient();
       const stream = await client.arcs.openLSP();
 
-      // Initialize
-      const initializeRequest: jsonRPC.Request = {
+      const initializeRequest: JSONRPCRequest = {
         jsonrpc: "2.0",
         id: 1,
         method: "initialize",
@@ -698,14 +601,9 @@ describe("Arc LSP", () => {
       };
 
       stream.send({ content: JSON.stringify(initializeRequest) });
-      const [initRes] = await stream.receive();
-      expect(initRes).not.toBeNull();
-      if (!initRes) throw new Error("Expected response");
-
-      const initMsg = JSON.parse(initRes.content) as jsonRPC.Response;
+      const initMsg = await receiveResponse(stream, 1);
       if ("error" in initMsg) throw new Error(`LSP error: ${initMsg.error.message}`);
 
-      // Send initialized notification
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -714,12 +612,10 @@ describe("Arc LSP", () => {
         }),
       });
 
-      // Create block URI
       const metadata = { is_block: true };
       const encoded = btoa(JSON.stringify(metadata));
       const blockURI = `arc://block/change-test#${encoded}`;
 
-      // Open initial document
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -735,37 +631,24 @@ describe("Arc LSP", () => {
         }),
       });
 
-      // Receive initial diagnostics
-      await stream.receive();
+      await receiveNotification(stream, "textDocument/publishDiagnostics");
 
-      // Send didChange to update the expression
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
           method: "textDocument/didChange",
           params: {
-            textDocument: {
-              uri: blockURI,
-              version: 2,
-            },
-            contentChanges: [
-              {
-                text: "return x + y",
-              },
-            ],
+            textDocument: { uri: blockURI, version: 2 },
+            contentChanges: [{ text: "return x + y" }],
           },
         }),
       });
 
-      // Should receive updated diagnostics
-      const [changeDiagResponse, changeDiagErr] = await stream.receive();
-      expect(changeDiagErr).toBeNull();
-      expect(changeDiagResponse).not.toBeNull();
-      if (!changeDiagResponse) throw new Error("Expected diagnostics");
-
-      const changeDiagMsg = JSON.parse(changeDiagResponse.content) as jsonRPC.Message;
-      if ("method" in changeDiagMsg && "params" in changeDiagMsg) {
-        expect(changeDiagMsg.method).toBe("textDocument/publishDiagnostics");
+      const changeDiagMsg = await receiveNotification(
+        stream,
+        "textDocument/publishDiagnostics",
+      );
+      if ("params" in changeDiagMsg) {
         const params = changeDiagMsg.params as Record<string, unknown>;
         expect(params.uri).toBe(blockURI);
       }
@@ -778,8 +661,7 @@ describe("Arc LSP", () => {
       const client = createTestClient();
       const stream = await client.arcs.openLSP();
 
-      // Initialize
-      const initializeRequest: jsonRPC.Request = {
+      const initializeRequest: JSONRPCRequest = {
         jsonrpc: "2.0",
         id: 1,
         method: "initialize",
@@ -792,14 +674,9 @@ describe("Arc LSP", () => {
       };
 
       stream.send({ content: JSON.stringify(initializeRequest) });
-      const [initRes] = await stream.receive();
-      expect(initRes).not.toBeNull();
-      if (!initRes) throw new Error("Expected response");
-
-      const initMsg = JSON.parse(initRes.content) as jsonRPC.Response;
+      const initMsg = await receiveResponse(stream, 1);
       if ("error" in initMsg) throw new Error(`LSP error: ${initMsg.error.message}`);
 
-      // Send initialized notification
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -808,11 +685,9 @@ describe("Arc LSP", () => {
         }),
       });
 
-      // Try to open a block URI without metadata fragment
+      // Block URI without metadata fragment - expression won't be wrapped
       const invalidBlockURI = "arc://block/no-metadata";
 
-      // Open a document with invalid block expression (missing wrapping)
-      const blockExpression = "return x * 2";
       stream.send({
         content: JSON.stringify({
           jsonrpc: "2.0",
@@ -822,26 +697,20 @@ describe("Arc LSP", () => {
               uri: invalidBlockURI,
               languageId: "arc",
               version: 1,
-              text: blockExpression,
+              text: "return x * 2",
             },
           },
         }),
       });
 
-      // Should receive diagnostics with parse errors (not wrapped)
-      const [diagResponse, diagErr] = await stream.receive();
-      expect(diagErr).toBeNull();
-      expect(diagResponse).not.toBeNull();
-      if (!diagResponse) throw new Error("Expected diagnostics");
-
-      const diagMsg = JSON.parse(diagResponse.content) as jsonRPC.Message;
-      if ("method" in diagMsg && "params" in diagMsg) {
-        expect(diagMsg.method).toBe("textDocument/publishDiagnostics");
+      const diagMsg = await receiveNotification(
+        stream,
+        "textDocument/publishDiagnostics",
+      );
+      if ("params" in diagMsg) {
         const params = diagMsg.params as Record<string, unknown>;
         const diagnostics = params.diagnostics as Array<{ severity: number }>;
-
-        // Should have error diagnostics because expression isn't wrapped
-        const errors = diagnostics.filter((d) => d.severity === 1); // Error severity
+        const errors = diagnostics.filter((d) => d.severity === 1);
         expect(errors.length).toBeGreaterThan(0);
       }