Skip to content

Include Authorization Info in Tool Calls (and request handlers generally) #166

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions src/server/sse.ts
Original file line number Diff line number Diff line change
@@ -4,6 +4,7 @@ import { Transport } from "../shared/transport.js";
import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
import getRawBody from "raw-body";
import contentType from "content-type";
import { AuthInfo } from "./auth/types.js";

const MAXIMUM_MESSAGE_SIZE = "4mb";

@@ -18,7 +19,7 @@ export class SSEServerTransport implements Transport {

onclose?: () => void;
onerror?: (error: Error) => void;
onmessage?: (message: JSONRPCMessage) => void;
onmessage?: (message: JSONRPCMessage, 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`.
@@ -66,7 +67,7 @@ export class SSEServerTransport implements Transport {
* This should be called when a POST request is made to send a message to the server.
*/
async handlePostMessage(
req: IncomingMessage,
req: IncomingMessage & { auth?: AuthInfo },
res: ServerResponse,
parsedBody?: unknown,
): Promise<void> {
@@ -75,6 +76,7 @@ export class SSEServerTransport implements Transport {
res.writeHead(500).end(message);
throw new Error(message);
}
const authInfo: AuthInfo | undefined = req.auth;

let body: string | unknown;
try {
@@ -94,7 +96,7 @@ export class SSEServerTransport implements Transport {
}

try {
await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body);
await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body, authInfo);
} catch {
res.writeHead(400).end(`Invalid message: ${body}`);
return;
@@ -106,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): Promise<void> {
async handleMessage(message: unknown, authInfo?: AuthInfo): Promise<void> {
let parsedMessage: JSONRPCMessage;
try {
parsedMessage = JSONRPCMessageSchema.parse(message);
@@ -115,7 +117,7 @@ export class SSEServerTransport implements Transport {
throw error;
}

this.onmessage?.(parsedMessage);
this.onmessage?.(parsedMessage, authInfo);
}

async close(): Promise<void> {
13 changes: 10 additions & 3 deletions src/shared/protocol.ts
Original file line number Diff line number Diff line change
@@ -19,6 +19,7 @@ import {
ServerCapabilities,
} from "../types.js";
import { Transport } from "./transport.js";
import { AuthInfo } from "../server/auth/types.js";

/**
* Callback for progress notifications.
@@ -89,6 +90,11 @@ export type RequestHandlerExtra = {
*/
signal: AbortSignal;

/**
* Information about a validated access token, provided to request handlers.
*/
authInfo?: AuthInfo;

/**
* The session ID from the transport, if available.
*/
@@ -240,11 +246,11 @@ export abstract class Protocol<
this._onerror(error);
};

this._transport.onmessage = (message) => {
this._transport.onmessage = (message, authInfo) => {
if (!("method" in message)) {
this._onresponse(message);
} else if ("id" in message) {
this._onrequest(message);
this._onrequest(message, authInfo);
} else {
this._onnotification(message);
}
@@ -290,7 +296,7 @@ export abstract class Protocol<
);
}

private _onrequest(request: JSONRPCRequest): void {
private _onrequest(request: JSONRPCRequest, authInfo?: AuthInfo): void {
const handler =
this._requestHandlers.get(request.method) ?? this.fallbackRequestHandler;

@@ -319,6 +325,7 @@ export abstract class Protocol<
const extra: RequestHandlerExtra = {
signal: abortController.signal,
sessionId: this._transport?.sessionId,
authInfo,
};

// Starting with Promise.resolve() puts any synchronous errors into the monad as well.
6 changes: 5 additions & 1 deletion src/shared/transport.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import { AuthInfo } from "../server/auth/types.js";
import { JSONRPCMessage } from "../types.js";

/**
@@ -39,8 +40,11 @@ export interface Transport {

/**
* Callback for when a message (request or response) is received over the connection.
*
* Includes the authInfo if the transport is authenticated.
*
*/
onmessage?: (message: JSONRPCMessage) => void;
onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void;

/**
* The session ID generated for this connection.