Skip to content

Commit 5a1fd1a

Browse files
felixweinbergerfelixweinberger
authored
Merge branch 'main' into fweinberger/backwards-compat-sampling
2 parents 102f4a1 + 6083600 commit 5a1fd1a

File tree

9 files changed

+476
-37
lines changed

9 files changed

+476
-37
lines changed

‎src/client/sse.test.ts‎

Lines changed: 78 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -308,14 +308,52 @@ describe('SSEClientTransport', () => {
308308

309309
await transport.start();
310310

311-
// Store original fetch
312311
const originalFetch = global.fetch;
312+
try {
313+
global.fetch = vi.fn().mockResolvedValue({ ok: true });
314+
315+
const message: JSONRPCMessage = {
316+
jsonrpc: '2.0',
317+
id: '1',
318+
method: 'test',
319+
params: {}
320+
};
321+
322+
await transport.send(message);
323+
324+
const calledHeaders = (global.fetch as Mock).mock.calls[0][1].headers;
325+
expect(calledHeaders.get('Authorization')).toBe('Bearer test-token');
326+
expect(calledHeaders.get('X-Custom-Header')).toBe('custom-value');
327+
expect(calledHeaders.get('content-type')).toBe('application/json');
328+
329+
customHeaders['X-Custom-Header'] = 'updated-value';
330+
331+
await transport.send(message);
332+
333+
const updatedHeaders = (global.fetch as Mock).mock.calls[1][1].headers;
334+
expect(updatedHeaders.get('X-Custom-Header')).toBe('updated-value');
335+
} finally {
336+
global.fetch = originalFetch;
337+
}
338+
});
339+
340+
it('passes custom headers to fetch requests (Headers class)', async () => {
341+
const customHeaders = new Headers({
342+
Authorization: 'Bearer test-token',
343+
'X-Custom-Header': 'custom-value'
344+
});
345+
346+
transport = new SSEClientTransport(resourceBaseUrl, {
347+
requestInit: {
348+
headers: customHeaders
349+
}
350+
});
351+
352+
await transport.start();
313353

354+
const originalFetch = global.fetch;
314355
try {
315-
// Mock fetch for the message sending test
316-
global.fetch = vi.fn().mockResolvedValue({
317-
ok: true
318-
});
356+
global.fetch = vi.fn().mockResolvedValue({ ok: true });
319357

320358
const message: JSONRPCMessage = {
321359
jsonrpc: '2.0',
@@ -326,20 +364,45 @@ describe('SSEClientTransport', () => {
326364

327365
await transport.send(message);
328366

329-
// Verify fetch was called with correct headers
330-
expect(global.fetch).toHaveBeenCalledWith(
331-
expect.any(URL),
332-
expect.objectContaining({
333-
headers: expect.any(Headers)
334-
})
335-
);
367+
const calledHeaders = (global.fetch as Mock).mock.calls[0][1].headers;
368+
expect(calledHeaders.get('Authorization')).toBe('Bearer test-token');
369+
expect(calledHeaders.get('X-Custom-Header')).toBe('custom-value');
370+
expect(calledHeaders.get('content-type')).toBe('application/json');
371+
372+
customHeaders.set('X-Custom-Header', 'updated-value');
373+
374+
await transport.send(message);
375+
376+
const updatedHeaders = (global.fetch as Mock).mock.calls[1][1].headers;
377+
expect(updatedHeaders.get('X-Custom-Header')).toBe('updated-value');
378+
} finally {
379+
global.fetch = originalFetch;
380+
}
381+
});
382+
383+
it('passes custom headers to fetch requests (array of tuples)', async () => {
384+
transport = new SSEClientTransport(resourceBaseUrl, {
385+
requestInit: {
386+
headers: [
387+
['Authorization', 'Bearer test-token'],
388+
['X-Custom-Header', 'custom-value']
389+
]
390+
}
391+
});
392+
393+
await transport.start();
394+
395+
const originalFetch = global.fetch;
396+
try {
397+
global.fetch = vi.fn().mockResolvedValue({ ok: true });
398+
399+
await transport.send({ jsonrpc: '2.0', id: '1', method: 'test', params: {} });
336400

337401
const calledHeaders = (global.fetch as Mock).mock.calls[0][1].headers;
338-
expect(calledHeaders.get('Authorization')).toBe(customHeaders.Authorization);
339-
expect(calledHeaders.get('X-Custom-Header')).toBe(customHeaders['X-Custom-Header']);
402+
expect(calledHeaders.get('Authorization')).toBe('Bearer test-token');
403+
expect(calledHeaders.get('X-Custom-Header')).toBe('custom-value');
340404
expect(calledHeaders.get('content-type')).toBe('application/json');
341405
} finally {
342-
// Restore original fetch
343406
global.fetch = originalFetch;
344407
}
345408
});

‎src/client/sse.ts‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { EventSource, type ErrorEvent, type EventSourceInit } from 'eventsource';
2-
import { Transport, FetchLike, createFetchWithInit } from '../shared/transport.js';
2+
import { Transport, FetchLike, createFetchWithInit, normalizeHeaders } from '../shared/transport.js';
33
import { JSONRPCMessage, JSONRPCMessageSchema } from '../types.js';
44
import { auth, AuthResult, extractWWWAuthenticateParams, OAuthClientProvider, UnauthorizedError } from './auth.js';
55

@@ -114,7 +114,7 @@ export class SSEClientTransport implements Transport {
114114
}
115115

116116
private async _commonHeaders(): Promise<Headers> {
117-
const headers: HeadersInit = {};
117+
const headers: HeadersInit & Record<string, string> = {};
118118
if (this._authProvider) {
119119
const tokens = await this._authProvider.tokens();
120120
if (tokens) {
@@ -125,7 +125,12 @@ export class SSEClientTransport implements Transport {
125125
headers['mcp-protocol-version'] = this._protocolVersion;
126126
}
127127

128-
return new Headers({ ...headers, ...this._requestInit?.headers });
128+
const extraHeaders = normalizeHeaders(this._requestInit?.headers);
129+
130+
return new Headers({
131+
...headers,
132+
...extraHeaders
133+
});
129134
}
130135

131136
private _startOrAuth(): Promise<void> {

‎src/client/streamableHttp.test.ts‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -480,6 +480,7 @@ describe('StreamableHTTPClientTransport', () => {
480480
it('should always send specified custom headers', async () => {
481481
const requestInit = {
482482
headers: {
483+
Authorization: 'Bearer test-token',
483484
'X-Custom-Header': 'CustomValue'
484485
}
485486
};
@@ -497,6 +498,7 @@ describe('StreamableHTTPClientTransport', () => {
497498
await transport.start();
498499

499500
await transport['_startOrAuthSse']({});
501+
expect((actualReqInit.headers as Headers).get('authorization')).toBe('Bearer test-token');
500502
expect((actualReqInit.headers as Headers).get('x-custom-header')).toBe('CustomValue');
501503

502504
requestInit.headers['X-Custom-Header'] = 'SecondCustomValue';
@@ -510,6 +512,7 @@ describe('StreamableHTTPClientTransport', () => {
510512
it('should always send specified custom headers (Headers class)', async () => {
511513
const requestInit = {
512514
headers: new Headers({
515+
Authorization: 'Bearer test-token',
513516
'X-Custom-Header': 'CustomValue'
514517
})
515518
};
@@ -527,6 +530,7 @@ describe('StreamableHTTPClientTransport', () => {
527530
await transport.start();
528531

529532
await transport['_startOrAuthSse']({});
533+
expect((actualReqInit.headers as Headers).get('authorization')).toBe('Bearer test-token');
530534
expect((actualReqInit.headers as Headers).get('x-custom-header')).toBe('CustomValue');
531535

532536
(requestInit.headers as Headers).set('X-Custom-Header', 'SecondCustomValue');
@@ -537,6 +541,30 @@ describe('StreamableHTTPClientTransport', () => {
537541
expect(global.fetch).toHaveBeenCalledTimes(2);
538542
});
539543

544+
it('should always send specified custom headers (array of tuples)', async () => {
545+
transport = new StreamableHTTPClientTransport(new URL('http://localhost:1234/mcp'), {
546+
requestInit: {
547+
headers: [
548+
['Authorization', 'Bearer test-token'],
549+
['X-Custom-Header', 'CustomValue']
550+
]
551+
}
552+
});
553+
554+
let actualReqInit: RequestInit = {};
555+
556+
(global.fetch as Mock).mockImplementation(async (_url, reqInit) => {
557+
actualReqInit = reqInit;
558+
return new Response(null, { status: 200, headers: { 'content-type': 'text/event-stream' } });
559+
});
560+
561+
await transport.start();
562+
563+
await transport['_startOrAuthSse']({});
564+
expect((actualReqInit.headers as Headers).get('authorization')).toBe('Bearer test-token');
565+
expect((actualReqInit.headers as Headers).get('x-custom-header')).toBe('CustomValue');
566+
});
567+
540568
it('should have exponential backoff with configurable maxRetries', () => {
541569
// This test verifies the maxRetries and backoff calculation directly
542570

‎src/server/sse.test.ts‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -547,6 +547,27 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
547547
expect(mockHandleRes.end).toHaveBeenCalledWith('Accepted');
548548
});
549549

550+
it('should accept requests without origin headers', async () => {
551+
const mockRes = createMockResponse();
552+
const transport = new SSEServerTransport('/messages', mockRes, {
553+
allowedOrigins: ['http://localhost:3000', 'https://example.com'],
554+
enableDnsRebindingProtection: true
555+
});
556+
await transport.start();
557+
558+
const mockReq = createMockRequest({
559+
headers: {
560+
'content-type': 'application/json'
561+
}
562+
});
563+
const mockHandleRes = createMockResponse();
564+
565+
await transport.handlePostMessage(mockReq, mockHandleRes, { jsonrpc: '2.0', method: 'test' });
566+
567+
expect(mockHandleRes.writeHead).toHaveBeenCalledWith(202);
568+
expect(mockHandleRes.end).toHaveBeenCalledWith('Accepted');
569+
});
570+
550571
it('should reject requests with disallowed origin headers', async () => {
551572
const mockRes = createMockResponse();
552573
const transport = new SSEServerTransport('/messages', mockRes, {

‎src/server/sse.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ export class SSEServerTransport implements Transport {
7979
// Validate Origin header if allowedOrigins is configured
8080
if (this._options.allowedOrigins && this._options.allowedOrigins.length > 0) {
8181
const originHeader = req.headers.origin;
82-
if (!originHeader || !this._options.allowedOrigins.includes(originHeader)) {
82+
if (originHeader && !this._options.allowedOrigins.includes(originHeader)) {
8383
return `Invalid Origin header: ${originHeader}`;
8484
}
8585
}

0 commit comments

Comments
 (0)