Skip to content

optimization: When using custom SSE request,Authorization header can still be automatically attached to the SSE request. #478

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

Open
wants to merge 10 commits into
base: main
Choose a base branch
from
Open
4 changes: 4 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -128,5 +128,9 @@ out
.yarn/install-state.gz
.pnp.*

# IDE
.vscode/
.idea/

.DS_Store
dist/
31 changes: 29 additions & 2 deletions src/client/sse.test.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { createServer, type IncomingMessage, type Server } from "http";
import { createServer, IncomingMessage, Server, ServerResponse } from "http";
import { AddressInfo } from "net";
import { JSONRPCMessage } from "../types.js";
import { SSEClientTransport } from "./sse.js";
Expand All @@ -12,8 +12,21 @@ describe("SSEClientTransport", () => {
let resourceBaseUrl: URL;
let authBaseUrl: URL;
let lastServerRequest: IncomingMessage;
const serverRequests: Record<string, IncomingMessage[]> = {};
let sendServerMessage: ((message: string) => void) | null = null;

const recordServerRequest = (req: IncomingMessage, res: ServerResponse) => {
lastServerRequest = req;

const key = `${req.method} ${req.url}`;
serverRequests[key] = serverRequests[key] || [];
serverRequests[key].push(req);

res.on('finish', () => {
console.log(`[server] ${req.method} ${req.url} -> ${res.statusCode} ${res.statusMessage}`);
});
};

beforeEach((done) => {
// Reset state
lastServerRequest = null as unknown as IncomingMessage;
Expand Down Expand Up @@ -606,6 +619,8 @@ describe("SSEClientTransport", () => {
authServer.close();

authServer = createServer((req, res) => {
recordServerRequest(req, res);

if (req.url === "/.well-known/oauth-authorization-server") {
res.writeHead(404).end();
return;
Expand All @@ -618,7 +633,7 @@ describe("SSEClientTransport", () => {
req.on("end", () => {
const params = new URLSearchParams(body);
if (params.get("grant_type") === "refresh_token" &&
params.get("refresh_token") === "refresh-token" &&
params.get("refresh_token")?.includes("refresh-token") &&
params.get("client_id") === "test-client-id" &&
params.get("client_secret") === "test-client-secret") {
res.writeHead(200, { "Content-Type": "application/json" });
Expand Down Expand Up @@ -649,6 +664,7 @@ describe("SSEClientTransport", () => {

let connectionAttempts = 0;
resourceServer = createServer((req, res) => {
recordServerRequest(req, res);
lastServerRequest = req;

if (req.url === "/.well-known/oauth-protected-resource") {
Expand Down Expand Up @@ -698,6 +714,14 @@ describe("SSEClientTransport", () => {

transport = new SSEClientTransport(resourceBaseUrl, {
authProvider: mockAuthProvider,
eventSourceInit: {
fetch: (url, init) => {
return fetch(url, { ...init, headers: {
...init?.headers,
'X-Custom-Header': 'custom-value'
} });
}
},
});

await transport.start();
Expand All @@ -709,6 +733,9 @@ describe("SSEClientTransport", () => {
});
expect(connectionAttempts).toBe(1);
expect(lastServerRequest.headers.authorization).toBe("Bearer new-token");
expect(serverRequests["GET /"]).toHaveLength(2);
expect(serverRequests["GET /"]
.every(req => req.headers["x-custom-header"] === "custom-value")).toBe(true);
});

it("refreshes expired token during POST request", async () => {
Expand Down
16 changes: 14 additions & 2 deletions src/client/sse.ts
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ export class SSEClientTransport implements Transport {
private _abortController?: AbortController;
private _url: URL;
private _resourceMetadataUrl?: URL;
private _eventSourceInit?: EventSourceInit;
private _eventSourceInit: EventSourceInit;
private _requestInit?: RequestInit;
private _authProvider?: OAuthClientProvider;
private _fetch?: FetchLike;
Expand All @@ -80,7 +80,19 @@ export class SSEClientTransport implements Transport {
) {
this._url = url;
this._resourceMetadataUrl = undefined;
this._eventSourceInit = opts?.eventSourceInit;

const actualFetch = opts?.eventSourceInit?.fetch ?? opts?.fetch ?? fetch;
this._eventSourceInit = {
...(opts?.eventSourceInit ?? {}),
fetch: (url, init) => this._commonHeaders().then((headers) => actualFetch(url, {
...init,
headers: {
...Object.fromEntries(headers.entries()),
Accept: "text/event-stream"
}
})),
};

this._requestInit = opts?.requestInit;
this._authProvider = opts?.authProvider;
this._fetch = opts?.fetch;
Expand Down