Skip to content

Commit 3dccd97

Browse files
authored
fix: forward custom headers in bearer token invoke paths (#1065)
Custom headers passed via -H/--header were silently dropped when using CUSTOM_JWT auth because invokeWithBearerToken and invokeWithBearerTokenStreaming did not merge options.headers into the request headers. Extract buildBearerInvokeHeaders() (paralleling the existing buildMcpBearerHeaders) and use it in both invoke paths. Closes #1052
1 parent a4c37a2 commit 3dccd97

2 files changed

Lines changed: 64 additions & 15 deletions

File tree

src/cli/aws/__tests__/agentcore.test.ts

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { extractResult, parseA2AResponse, parseSSE, parseSSELine } from '../agentcore.js';
1+
import { buildBearerInvokeHeaders, extractResult, parseA2AResponse, parseSSE, parseSSELine } from '../agentcore.js';
22
import { describe, expect, it } from 'vitest';
33

44
describe('parseSSELine', () => {
@@ -176,3 +176,43 @@ describe('parseA2AResponse', () => {
176176
expect(parseA2AResponse('not json')).toBe('not json');
177177
});
178178
});
179+
180+
describe('buildBearerInvokeHeaders', () => {
181+
it('includes custom headers from options.headers', () => {
182+
const headers = buildBearerInvokeHeaders(
183+
{
184+
bearerToken: 'tok',
185+
headers: {
186+
'x-amzn-bedrock-agentcore-runtime-custom-foo': 'bar',
187+
'x-amzn-bedrock-agentcore-runtime-custom-baz': 'qux',
188+
},
189+
},
190+
'application/json'
191+
);
192+
expect(headers['x-amzn-bedrock-agentcore-runtime-custom-foo']).toBe('bar');
193+
expect(headers['x-amzn-bedrock-agentcore-runtime-custom-baz']).toBe('qux');
194+
});
195+
196+
it('sets Authorization, Content-Type, Accept, and default user ID', () => {
197+
const headers = buildBearerInvokeHeaders({ bearerToken: 'tok' }, 'application/json');
198+
expect(headers.Authorization).toBe('Bearer tok');
199+
expect(headers['Content-Type']).toBe('application/json');
200+
expect(headers.Accept).toBe('application/json');
201+
expect(headers['X-Amzn-Bedrock-AgentCore-Runtime-User-Id']).toBe('default-user');
202+
});
203+
204+
it('sets session ID header when provided', () => {
205+
const headers = buildBearerInvokeHeaders({ bearerToken: 'tok', sessionId: 's1' }, 'application/json');
206+
expect(headers['X-Amzn-Bedrock-AgentCore-Runtime-Session-Id']).toBe('s1');
207+
});
208+
209+
it('omits session ID header when not provided', () => {
210+
const headers = buildBearerInvokeHeaders({ bearerToken: 'tok' }, 'application/json');
211+
expect(headers).not.toHaveProperty('X-Amzn-Bedrock-AgentCore-Runtime-Session-Id');
212+
});
213+
214+
it('returns correct headers when options.headers is undefined', () => {
215+
const headers = buildBearerInvokeHeaders({ bearerToken: 'tok' }, 'application/json');
216+
expect(Object.keys(headers)).toHaveLength(4); // Authorization, Content-Type, Accept, User-Id
217+
});
218+
});

src/cli/aws/agentcore.ts

Lines changed: 23 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -151,20 +151,37 @@ function buildInvokeUrl(region: string, runtimeArn: string): string {
151151
}
152152

153153
/**
154-
* Invoke an AgentCore Runtime using bearer token auth (raw HTTP, no SigV4).
155-
* Used when the runtime has CUSTOM_JWT authorizer configured.
154+
* Build headers for bearer-token invoke requests.
155+
* Shared by both streaming and non-streaming invoke paths.
156156
*/
157-
async function invokeWithBearerTokenStreaming(options: InvokeAgentRuntimeOptions): Promise<StreamingInvokeResult> {
158-
const url = buildInvokeUrl(options.region, options.runtimeArn);
157+
export function buildBearerInvokeHeaders(
158+
options: Pick<InvokeAgentRuntimeOptions, 'bearerToken' | 'sessionId' | 'userId' | 'headers'>,
159+
accept: string
160+
): Record<string, string> {
159161
const headers: Record<string, string> = {
160162
Authorization: `Bearer ${options.bearerToken}`,
161163
'Content-Type': 'application/json',
162-
Accept: 'application/json, text/event-stream',
164+
Accept: accept,
163165
};
164166
if (options.sessionId) {
165167
headers['X-Amzn-Bedrock-AgentCore-Runtime-Session-Id'] = options.sessionId;
166168
}
167169
headers['X-Amzn-Bedrock-AgentCore-Runtime-User-Id'] = options.userId ?? DEFAULT_RUNTIME_USER_ID;
170+
if (options.headers) {
171+
for (const [name, value] of Object.entries(options.headers)) {
172+
headers[name] = value;
173+
}
174+
}
175+
return headers;
176+
}
177+
178+
/**
179+
* Invoke an AgentCore Runtime using bearer token auth (raw HTTP, no SigV4).
180+
* Used when the runtime has CUSTOM_JWT authorizer configured.
181+
*/
182+
async function invokeWithBearerTokenStreaming(options: InvokeAgentRuntimeOptions): Promise<StreamingInvokeResult> {
183+
const url = buildInvokeUrl(options.region, options.runtimeArn);
184+
const headers = buildBearerInvokeHeaders(options, 'application/json, text/event-stream');
168185

169186
const res = await fetch(url, {
170187
method: 'POST',
@@ -250,15 +267,7 @@ async function invokeWithBearerTokenStreaming(options: InvokeAgentRuntimeOptions
250267
*/
251268
async function invokeWithBearerToken(options: InvokeAgentRuntimeOptions): Promise<InvokeAgentRuntimeResult> {
252269
const url = buildInvokeUrl(options.region, options.runtimeArn);
253-
const headers: Record<string, string> = {
254-
Authorization: `Bearer ${options.bearerToken}`,
255-
'Content-Type': 'application/json',
256-
Accept: 'application/json',
257-
};
258-
if (options.sessionId) {
259-
headers['X-Amzn-Bedrock-AgentCore-Runtime-Session-Id'] = options.sessionId;
260-
}
261-
headers['X-Amzn-Bedrock-AgentCore-Runtime-User-Id'] = options.userId ?? DEFAULT_RUNTIME_USER_ID;
270+
const headers = buildBearerInvokeHeaders(options, 'application/json');
262271

263272
const res = await fetch(url, {
264273
method: 'POST',

0 commit comments

Comments
 (0)