Add AuthInfo test

Author: allenzhou101

src/inMemory.test.ts

@@ -1,5 +1,6 @@
 import { InMemoryTransport } from "./inMemory.js";
 import { JSONRPCMessage } from "./types.js";
+import { AuthInfo } from "./server/auth/types.js";
 
 describe("InMemoryTransport", () => {
   let clientTransport: InMemoryTransport;
@@ -35,6 +36,32 @@ describe("InMemoryTransport", () => {
     expect(receivedMessage).toEqual(message);
   });
 
+  test("should send message with auth info from client to server", async () => {
+    const message: JSONRPCMessage = {
+      jsonrpc: "2.0",
+      method: "test",
+      id: 1,
+    };
+
+    const authInfo: AuthInfo = {
+      token: "test-token",
+      clientId: "test-client",
+      scopes: ["read", "write"],
+      expiresAt: Date.now() / 1000 + 3600,
+    };
+
+    let receivedMessage: JSONRPCMessage | undefined;
+    let receivedAuthInfo: AuthInfo | undefined;
+    serverTransport.onmessage = (msg, auth) => {
+      receivedMessage = msg;
+      receivedAuthInfo = auth;
+    };
+
+    await clientTransport.sendWithAuth(message, authInfo);
+    expect(receivedMessage).toEqual(message);
+    expect(receivedAuthInfo).toEqual(authInfo);
+  });
+
   test("should send message from server to client", async () => {
     const message: JSONRPCMessage = {
       jsonrpc: "2.0",

src/inMemory.ts

@@ -1,16 +1,22 @@
 import { Transport } from "./shared/transport.js";
 import { JSONRPCMessage } from "./types.js";
+import { AuthInfo } from "./server/auth/types.js";
+
+interface QueuedMessage {
+  message: JSONRPCMessage;
+  authInfo?: AuthInfo;
+}
 
 /**
  * In-memory transport for creating clients and servers that talk to each other within the same process.
  */
 export class InMemoryTransport implements Transport {
   private _otherTransport?: InMemoryTransport;
-  private _messageQueue: JSONRPCMessage[] = [];
+  private _messageQueue: QueuedMessage[] = [];
 
   onclose?: () => void;
   onerror?: (error: Error) => void;
-  onmessage?: (message: JSONRPCMessage) => void;
+  onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void;
   sessionId?: string;
 
   /**
@@ -27,10 +33,8 @@ export class InMemoryTransport implements Transport {
   async start(): Promise<void> {
     // Process any messages that were queued before start was called
     while (this._messageQueue.length > 0) {
-      const message = this._messageQueue.shift();
-      if (message) {
-        this.onmessage?.(message);
-      }
+      const queuedMessage = this._messageQueue.shift()!;
+      this.onmessage?.(queuedMessage.message, queuedMessage.authInfo);
     }
   }
 
@@ -42,14 +46,22 @@ export class InMemoryTransport implements Transport {
   }
 
   async send(message: JSONRPCMessage): Promise<void> {
+    return this.sendWithAuth(message);
+  }
+
+  /**
+   * Sends a message with optional auth info.
+   * This is useful for testing authentication scenarios.
+   */
+  async sendWithAuth(message: JSONRPCMessage, authInfo?: AuthInfo): Promise<void> {
     if (!this._otherTransport) {
       throw new Error("Not connected");
     }
 
     if (this._otherTransport.onmessage) {
-      this._otherTransport.onmessage(message);
+      this._otherTransport.onmessage(message, authInfo);
     } else {
-      this._otherTransport._messageQueue.push(message);
+      this._otherTransport._messageQueue.push({ message, authInfo });
     }
   }
 }