Pass extra object instead of authInfo directly

This commit is contained in:
Allen Zhou 2025-04-12 16:38:55 -07:00
commit c4381b857b
No known key found for this signature in database
GPG key ID: 5817C145DE5C5556
4 changed files with 21 additions and 24 deletions

View file

@ -52,12 +52,12 @@ describe("InMemoryTransport", () => {
let receivedMessage: JSONRPCMessage | undefined;
let receivedAuthInfo: AuthInfo | undefined;
serverTransport.onmessage = (msg, auth) => {
serverTransport.onmessage = (msg, extra) => {
receivedMessage = msg;
receivedAuthInfo = auth;
receivedAuthInfo = extra?.authInfo;
};
await clientTransport.sendWithAuth(message, authInfo);
await clientTransport.send(message, { authInfo });
expect(receivedMessage).toEqual(message);
expect(receivedAuthInfo).toEqual(authInfo);
});

View file

@ -4,7 +4,7 @@ import { AuthInfo } from "./server/auth/types.js";
interface QueuedMessage {
message: JSONRPCMessage;
authInfo?: AuthInfo;
extra?: { authInfo?: AuthInfo };
}
/**
@ -16,7 +16,7 @@ export class InMemoryTransport implements Transport {
onclose?: () => void;
onerror?: (error: Error) => void;
onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void;
onmessage?: (message: JSONRPCMessage, extra?: { authInfo?: AuthInfo }) => void;
sessionId?: string;
/**
@ -34,7 +34,7 @@ export class InMemoryTransport implements Transport {
// Process any messages that were queued before start was called
while (this._messageQueue.length > 0) {
const queuedMessage = this._messageQueue.shift()!;
this.onmessage?.(queuedMessage.message, queuedMessage.authInfo);
this.onmessage?.(queuedMessage.message, queuedMessage.extra);
}
}
@ -45,23 +45,19 @@ export class InMemoryTransport implements Transport {
this.onclose?.();
}
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 send(message: JSONRPCMessage, options?: { authInfo?: AuthInfo }): Promise<void>
async send(message: JSONRPCMessage, options?: { authInfo?: AuthInfo }): Promise<void> {
if (!this._otherTransport) {
throw new Error("Not connected");
}
if (this._otherTransport.onmessage) {
this._otherTransport.onmessage(message, authInfo);
this._otherTransport.onmessage(message, { authInfo: options?.authInfo });
} else {
this._otherTransport._messageQueue.push({ message, authInfo });
this._otherTransport._messageQueue.push({ message, extra: { authInfo: options?.authInfo } });
}
}
}

View file

@ -19,7 +19,7 @@ export class SSEServerTransport implements Transport {
onclose?: () => void;
onerror?: (error: Error) => void;
onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void;
onmessage?: (message: JSONRPCMessage, extra?: { authInfo?: AuthInfo }) => void;
/**
* Creates a new SSE server transport, which will direct the client to POST messages to the relative or absolute URL identified by `_endpoint`.
@ -96,7 +96,7 @@ export class SSEServerTransport implements Transport {
}
try {
await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body, authInfo);
await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body, { authInfo });
} catch {
res.writeHead(400).end(`Invalid message: ${body}`);
return;
@ -108,7 +108,7 @@ export class SSEServerTransport implements Transport {
/**
* Handle a client message, regardless of how it arrived. This can be used to inform the server of messages that arrive via a means different than HTTP POST.
*/
async handleMessage(message: unknown, authInfo?: AuthInfo): Promise<void> {
async handleMessage(message: unknown, extra?: { authInfo?: AuthInfo }): Promise<void> {
let parsedMessage: JSONRPCMessage;
try {
parsedMessage = JSONRPCMessageSchema.parse(message);
@ -117,7 +117,7 @@ export class SSEServerTransport implements Transport {
throw error;
}
this.onmessage?.(parsedMessage, authInfo);
this.onmessage?.(parsedMessage, extra);
}
async close(): Promise<void> {

View file

@ -246,11 +246,11 @@ export abstract class Protocol<
this._onerror(error);
};
this._transport.onmessage = (message, authInfo) => {
this._transport.onmessage = (message, extra) => {
if (!("method" in message)) {
this._onresponse(message);
} else if ("id" in message) {
this._onrequest(message, authInfo);
this._onrequest(message, extra);
} else {
this._onnotification(message);
}
@ -296,7 +296,7 @@ export abstract class Protocol<
);
}
private _onrequest(request: JSONRPCRequest, authInfo?: AuthInfo): void {
private _onrequest(request: JSONRPCRequest, extra?: { authInfo?: AuthInfo }): void {
const handler =
this._requestHandlers.get(request.method) ?? this.fallbackRequestHandler;
@ -321,16 +321,17 @@ export abstract class Protocol<
const abortController = new AbortController();
this._requestHandlerAbortControllers.set(request.id, abortController);
// Create extra object with both abort signal and sessionId from transport
const extra: RequestHandlerExtra = {
// Create fullExtra object with both abort signal and sessionId from transport
// in addition to any authInfo provided
const fullExtra: RequestHandlerExtra = {
signal: abortController.signal,
sessionId: this._transport?.sessionId,
authInfo,
authInfo: extra?.authInfo,
};
// Starting with Promise.resolve() puts any synchronous errors into the monad as well.
Promise.resolve()
.then(() => handler(request, extra))
.then(() => handler(request, fullExtra))
.then(
(result) => {
if (abortController.signal.aborted) {