diff --git a/src/cli/aws/__tests__/agentcore.test.ts b/src/cli/aws/__tests__/agentcore.test.ts index e26e4324e..f23ff5c83 100644 --- a/src/cli/aws/__tests__/agentcore.test.ts +++ b/src/cli/aws/__tests__/agentcore.test.ts @@ -1,4 +1,4 @@ -import { extractResult, parseA2AResponse, parseSSE, parseSSELine } from '../agentcore.js'; +import { buildBearerInvokeHeaders, extractResult, parseA2AResponse, parseSSE, parseSSELine } from '../agentcore.js'; import { describe, expect, it } from 'vitest'; describe('parseSSELine', () => { @@ -176,3 +176,43 @@ describe('parseA2AResponse', () => { expect(parseA2AResponse('not json')).toBe('not json'); }); }); + +describe('buildBearerInvokeHeaders', () => { + it('includes custom headers from options.headers', () => { + const headers = buildBearerInvokeHeaders( + { + bearerToken: 'tok', + headers: { + 'x-amzn-bedrock-agentcore-runtime-custom-foo': 'bar', + 'x-amzn-bedrock-agentcore-runtime-custom-baz': 'qux', + }, + }, + 'application/json' + ); + expect(headers['x-amzn-bedrock-agentcore-runtime-custom-foo']).toBe('bar'); + expect(headers['x-amzn-bedrock-agentcore-runtime-custom-baz']).toBe('qux'); + }); + + it('sets Authorization, Content-Type, Accept, and default user ID', () => { + const headers = buildBearerInvokeHeaders({ bearerToken: 'tok' }, 'application/json'); + expect(headers.Authorization).toBe('Bearer tok'); + expect(headers['Content-Type']).toBe('application/json'); + expect(headers.Accept).toBe('application/json'); + expect(headers['X-Amzn-Bedrock-AgentCore-Runtime-User-Id']).toBe('default-user'); + }); + + it('sets session ID header when provided', () => { + const headers = buildBearerInvokeHeaders({ bearerToken: 'tok', sessionId: 's1' }, 'application/json'); + expect(headers['X-Amzn-Bedrock-AgentCore-Runtime-Session-Id']).toBe('s1'); + }); + + it('omits session ID header when not provided', () => { + const headers = buildBearerInvokeHeaders({ bearerToken: 'tok' }, 'application/json'); + expect(headers).not.toHaveProperty('X-Amzn-Bedrock-AgentCore-Runtime-Session-Id'); + }); + + it('returns correct headers when options.headers is undefined', () => { + const headers = buildBearerInvokeHeaders({ bearerToken: 'tok' }, 'application/json'); + expect(Object.keys(headers)).toHaveLength(4); // Authorization, Content-Type, Accept, User-Id + }); +}); diff --git a/src/cli/aws/agentcore.ts b/src/cli/aws/agentcore.ts index 55c19d2d0..2cf54ac22 100644 --- a/src/cli/aws/agentcore.ts +++ b/src/cli/aws/agentcore.ts @@ -151,20 +151,37 @@ function buildInvokeUrl(region: string, runtimeArn: string): string { } /** - * Invoke an AgentCore Runtime using bearer token auth (raw HTTP, no SigV4). - * Used when the runtime has CUSTOM_JWT authorizer configured. + * Build headers for bearer-token invoke requests. + * Shared by both streaming and non-streaming invoke paths. */ -async function invokeWithBearerTokenStreaming(options: InvokeAgentRuntimeOptions): Promise { - const url = buildInvokeUrl(options.region, options.runtimeArn); +export function buildBearerInvokeHeaders( + options: Pick, + accept: string +): Record { const headers: Record = { Authorization: `Bearer ${options.bearerToken}`, 'Content-Type': 'application/json', - Accept: 'application/json, text/event-stream', + Accept: accept, }; if (options.sessionId) { headers['X-Amzn-Bedrock-AgentCore-Runtime-Session-Id'] = options.sessionId; } headers['X-Amzn-Bedrock-AgentCore-Runtime-User-Id'] = options.userId ?? DEFAULT_RUNTIME_USER_ID; + if (options.headers) { + for (const [name, value] of Object.entries(options.headers)) { + headers[name] = value; + } + } + return headers; +} + +/** + * Invoke an AgentCore Runtime using bearer token auth (raw HTTP, no SigV4). + * Used when the runtime has CUSTOM_JWT authorizer configured. + */ +async function invokeWithBearerTokenStreaming(options: InvokeAgentRuntimeOptions): Promise { + const url = buildInvokeUrl(options.region, options.runtimeArn); + const headers = buildBearerInvokeHeaders(options, 'application/json, text/event-stream'); const res = await fetch(url, { method: 'POST', @@ -250,15 +267,7 @@ async function invokeWithBearerTokenStreaming(options: InvokeAgentRuntimeOptions */ async function invokeWithBearerToken(options: InvokeAgentRuntimeOptions): Promise { const url = buildInvokeUrl(options.region, options.runtimeArn); - const headers: Record = { - Authorization: `Bearer ${options.bearerToken}`, - 'Content-Type': 'application/json', - Accept: 'application/json', - }; - if (options.sessionId) { - headers['X-Amzn-Bedrock-AgentCore-Runtime-Session-Id'] = options.sessionId; - } - headers['X-Amzn-Bedrock-AgentCore-Runtime-User-Id'] = options.userId ?? DEFAULT_RUNTIME_USER_ID; + const headers = buildBearerInvokeHeaders(options, 'application/json'); const res = await fetch(url, { method: 'POST',