diff --git a/src/inMemory.test.ts b/src/inMemory.test.ts index f7e9e979..b6906cc8 100644 --- a/src/inMemory.test.ts +++ b/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", diff --git a/src/inMemory.ts b/src/inMemory.ts index 106a9e7e..9de100c2 100644 --- a/src/inMemory.ts +++ b/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 { // 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 { + 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 { 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 }); } } }