fix: forward custom headers in bearer token invoke paths

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
+  });
+});

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<StreamingInvokeResult> {
-  const url = buildInvokeUrl(options.region, options.runtimeArn);
+export function buildBearerInvokeHeaders(
+  options: Pick<InvokeAgentRuntimeOptions, 'bearerToken' | 'sessionId' | 'userId' | 'headers'>,
+  accept: string
+): Record<string, string> {
   const headers: Record<string, string> = {
     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<StreamingInvokeResult> {
+  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<InvokeAgentRuntimeResult> {
   const url = buildInvokeUrl(options.region, options.runtimeArn);
-  const headers: Record<string, string> = {
-    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',