diff --git a/packages/sdk/src/proxy/client.ts b/packages/sdk/src/proxy/client.ts index 70ab0ff..b66d98d 100644 --- a/packages/sdk/src/proxy/client.ts +++ b/packages/sdk/src/proxy/client.ts @@ -23,6 +23,20 @@ export interface ProxyContext { remoteTools: Tool[]; } +/** + * Build the appropriate auth headers based on config. + */ +function buildAuthHeaders(config: StitchProxyConfig): Record { + if (config.accessToken) { + const headers: Record = { Authorization: `Bearer ${config.accessToken}` }; + if (config.quotaProjectId) { + headers['X-Goog-User-Project'] = config.quotaProjectId; + } + return headers; + } + return { 'X-Goog-Api-Key': config.apiKey! }; +} + /** * Forward a JSON-RPC request to Stitch. */ @@ -45,7 +59,7 @@ export async function forwardToStitch( headers: { 'Content-Type': 'application/json', Accept: 'application/json', - 'X-Goog-Api-Key': config.apiKey!, + ...buildAuthHeaders(config), }, body: JSON.stringify(request), }); @@ -92,7 +106,7 @@ export async function initializeStitchConnection( headers: { 'Content-Type': 'application/json', Accept: 'application/json', - 'X-Goog-Api-Key': ctx.config.apiKey!, + ...buildAuthHeaders(ctx.config), }, body: JSON.stringify({ jsonrpc: '2.0', diff --git a/packages/sdk/src/proxy/core.ts b/packages/sdk/src/proxy/core.ts index 72633be..ba45892 100644 --- a/packages/sdk/src/proxy/core.ts +++ b/packages/sdk/src/proxy/core.ts @@ -32,6 +32,8 @@ export class StitchProxy implements StitchProxySpec { constructor(inputConfig?: Partial) { const rawConfig = { apiKey: inputConfig?.apiKey || process.env.STITCH_API_KEY, + accessToken: inputConfig?.accessToken || process.env.STITCH_ACCESS_TOKEN, + quotaProjectId: inputConfig?.quotaProjectId || process.env.STITCH_PROJECT_ID || process.env.GOOGLE_CLOUD_PROJECT, url: inputConfig?.url || process.env.STITCH_MCP_URL, name: inputConfig?.name, version: inputConfig?.version, @@ -40,8 +42,8 @@ export class StitchProxy implements StitchProxySpec { // Validate config this.config = StitchProxyConfigSchema.parse(rawConfig); - if (!this.config.apiKey) { - throw new Error('StitchProxy requires an API key (STITCH_API_KEY)'); + if (!this.config.apiKey && !this.config.accessToken) { + throw new Error('StitchProxy requires an API key (STITCH_API_KEY) or access token (STITCH_ACCESS_TOKEN)'); } this.server = new McpServer( diff --git a/packages/sdk/src/spec/proxy.ts b/packages/sdk/src/spec/proxy.ts index 7c1dbe3..ed88dfb 100644 --- a/packages/sdk/src/spec/proxy.ts +++ b/packages/sdk/src/spec/proxy.ts @@ -23,6 +23,12 @@ export const StitchProxyConfigSchema = z.object({ /** API key for Stitch authentication. Falls back to STITCH_API_KEY. */ apiKey: z.string().optional(), + /** Access token for Stitch authentication (Bearer). Falls back to STITCH_ACCESS_TOKEN. */ + accessToken: z.string().optional(), + + /** Quota project ID for billing. Required with accessToken auth. Falls back to STITCH_PROJECT_ID or GOOGLE_CLOUD_PROJECT. */ + quotaProjectId: z.string().optional(), + /** Target Stitch MCP URL. Default: https://stitch.googleapis.com/mcp */ url: z.string().default(DEFAULT_STITCH_API_URL), diff --git a/packages/sdk/test/proxy.test.ts b/packages/sdk/test/proxy.test.ts index 7114ec3..2e7a623 100644 --- a/packages/sdk/test/proxy.test.ts +++ b/packages/sdk/test/proxy.test.ts @@ -43,9 +43,25 @@ describe('StitchProxy', () => { expect(proxy).toBeDefined(); }); - it('should throw if no API key is provided', () => { + it('should throw if neither API key nor access token is provided', () => { delete process.env.STITCH_API_KEY; - expect(() => new StitchProxy({})).toThrow("StitchProxy requires an API key"); + delete process.env.STITCH_ACCESS_TOKEN; + expect(() => new StitchProxy({})).toThrow("StitchProxy requires an API key (STITCH_API_KEY) or access token (STITCH_ACCESS_TOKEN)"); + }); + + it('should initialize with accessToken instead of apiKey', () => { + delete process.env.STITCH_API_KEY; + delete process.env.STITCH_ACCESS_TOKEN; + const proxy = new StitchProxy({ accessToken: 'test-token' }); + expect(proxy).toBeDefined(); + }); + + it('should initialize with STITCH_ACCESS_TOKEN env var', () => { + delete process.env.STITCH_API_KEY; + process.env.STITCH_ACCESS_TOKEN = 'env-token'; + const proxy = new StitchProxy({}); + expect(proxy).toBeDefined(); + delete process.env.STITCH_ACCESS_TOKEN; }); it('should connect to stitch and fetch tools on start', async () => { @@ -99,6 +115,37 @@ describe('Proxy Client Error Handling', () => { vi.clearAllMocks(); }); + it('forwardToStitch should send Authorization: Bearer header when accessToken is configured', async () => { + mockFetch.mockResolvedValue({ + ok: true, + json: async () => ({ result: 'ok' }) + } as Response); + + await forwardToStitch({ url: 'http://test', accessToken: 'my-token', quotaProjectId: 'my-project' } as any, 'testMethod'); + + const fetchCall = mockFetch.mock.calls[0]; + expect(fetchCall[1].headers).toEqual(expect.objectContaining({ + Authorization: 'Bearer my-token', + 'X-Goog-User-Project': 'my-project', + })); + expect(fetchCall[1].headers).not.toHaveProperty('X-Goog-Api-Key'); + }); + + it('forwardToStitch should send X-Goog-Api-Key header when only apiKey is configured', async () => { + mockFetch.mockResolvedValue({ + ok: true, + json: async () => ({ result: 'ok' }) + } as Response); + + await forwardToStitch({ url: 'http://test', apiKey: 'my-api-key' } as any, 'testMethod'); + + const fetchCall = mockFetch.mock.calls[0]; + expect(fetchCall[1].headers).toEqual(expect.objectContaining({ + 'X-Goog-Api-Key': 'my-api-key', + })); + expect(fetchCall[1].headers).not.toHaveProperty('Authorization'); + }); + it('forwardToStitch should throw Stitch API error on non-ok response', async () => { mockFetch.mockResolvedValue({ ok: false,