Skip to content

Commit 420c96e

Browse files
wirjovincentkoc
wirjo
and
vincentkoc
authored
fix(amazon-bedrock-mantle): refresh IAM bearer token via resolveConfigApiKey cache lookup (#68903)
* fix(amazon-bedrock-mantle): refresh IAM bearer token via resolveConfigApiKey cache lookup The Mantle plugin generates a bearer token from IAM credentials at discovery time and bakes it as a static string into the provider config. After the token's cache TTL expires (~1hr), requests fail because resolveConfigApiKey only handled the explicit AWS_BEARER_TOKEN_BEDROCK env var case. Fix: expose getCachedIamToken() as a sync read from the existing iamTokenCache, and wire it into resolveConfigApiKey as a fallback when no explicit env var is set. The catalog.run still generates/refreshes the token on discovery; this change ensures the cached token is served at auth resolution time. Fixes #68900 * fix(amazon-bedrock-mantle): refresh runtime IAM bearer auth * docs(changelog): note Mantle IAM refresh * fix(agents): apply runtime auth in simple completion --------- Co-authored-by: Vincent Koc <[email protected]>
1 parent 2321d67 commit 420c96e

7 files changed

Lines changed: 243 additions & 10 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ Docs: https://docs.openclaw.ai
2323

2424
### Fixes
2525

26+
- Providers/Amazon Bedrock Mantle: refresh IAM-backed bearer tokens at runtime instead of baking discovery-time tokens into provider config, so long-lived Mantle sessions keep working after the initial token ages out. Thanks @wirjo.
2627
- Codex harness: rotate the shared app-server websocket client when the configured bearer token changes, so auth-token refreshes reconnect with the new `Authorization` header instead of reusing a stale socket. (#70328) Thanks @Lucenx9.
2728
- Telegram/sandbox: keep Telegram bot DMs on per-account sender session keys even when `session.dmScope=main`, so sandbox/tool policy can distinguish Telegram-originated direct chats from the agent main session.
2829
- Config/models: merge provider-scoped model allowlist updates and protect model/provider map writes from accidental full replacement, adding `config set --merge` for additive updates and `--replace` for intentional clobbers. Fixes #65920, #68392, and #68653.

‎extensions/amazon-bedrock-mantle/api.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
export {
22
discoverMantleModels,
33
generateBearerTokenFromIam,
4+
getCachedIamToken,
5+
MANTLE_IAM_TOKEN_MARKER,
46
mergeImplicitMantleProvider,
57
resetIamTokenCacheForTest,
68
resetMantleDiscoveryCacheForTest,
79
resolveImplicitMantleProvider,
810
resolveMantleBearerToken,
11+
resolveMantleRuntimeBearerToken,
912
} from "./discovery.js";

‎extensions/amazon-bedrock-mantle/discovery.test.ts‎

Lines changed: 75 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,14 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
22
import {
33
discoverMantleModels,
44
generateBearerTokenFromIam,
5+
getCachedIamToken,
6+
MANTLE_IAM_TOKEN_MARKER,
57
mergeImplicitMantleProvider,
68
resetIamTokenCacheForTest,
79
resetMantleDiscoveryCacheForTest,
810
resolveMantleBearerToken,
911
resolveImplicitMantleProvider,
12+
resolveMantleRuntimeBearerToken,
1013
} from "./api.js";
1114

1215
const mocks = vi.hoisted(() => ({
@@ -80,7 +83,7 @@ describe("bedrock mantle discovery", () => {
8083
let now = 1000;
8184

8285
const t1 = await generateBearerTokenFromIam({ region: "us-east-1", now: () => now });
83-
now += 1800_000; // 30 min — within 1hr cache TTL
86+
now += 1800_000; // 30 min — within 2hr cache TTL
8487
const t2 = await generateBearerTokenFromIam({ region: "us-east-1", now: () => now });
8588

8689
expect(t1).toEqual(t2);
@@ -118,6 +121,33 @@ describe("bedrock mantle discovery", () => {
118121
await expect(generateBearerTokenFromIam({ region: "us-east-1" })).resolves.toBeUndefined();
119122
});
120123

124+
it("getCachedIamToken returns cached token when valid", async () => {
125+
const tokenProvider = vi.fn(async () => "bedrock-cached-token"); // pragma: allowlist secret
126+
mocks.getTokenProvider.mockReturnValue(tokenProvider);
127+
128+
// Generate a token to populate the cache
129+
await generateBearerTokenFromIam({ region: "us-east-1" });
130+
131+
// Sync read should return the cached token
132+
expect(getCachedIamToken("us-east-1")).toBe("bedrock-cached-token");
133+
});
134+
135+
it("getCachedIamToken returns undefined when cache is empty", () => {
136+
expect(getCachedIamToken("us-east-1")).toBeUndefined();
137+
});
138+
139+
it("getCachedIamToken returns undefined when cache is expired", async () => {
140+
const tokenProvider = vi.fn(async () => "bedrock-expired-token"); // pragma: allowlist secret
141+
mocks.getTokenProvider.mockReturnValue(tokenProvider);
142+
143+
// Generate with a time far in the past so it's already expired
144+
await generateBearerTokenFromIam({ region: "us-east-1", now: () => 1000 });
145+
146+
// The cache entry exists but expiresAt is 1000 + 3600000 = 3601000
147+
// Current Date.now() is way past that, so it should be expired
148+
expect(getCachedIamToken("us-east-1")).toBeUndefined();
149+
});
150+
121151
// ---------------------------------------------------------------------------
122152
// Model discovery
123153
// ---------------------------------------------------------------------------
@@ -383,7 +413,7 @@ describe("bedrock mantle discovery", () => {
383413
});
384414

385415
expect(provider).not.toBeNull();
386-
expect(provider?.apiKey).toBe("bedrock-api-key-iam");
416+
expect(provider?.apiKey).toBe(MANTLE_IAM_TOKEN_MARKER);
387417
expect(tokenProvider).toHaveBeenCalledTimes(1);
388418
expect(mockFetch).toHaveBeenCalledWith(
389419
"https://bedrock-mantle.us-east-1.api.aws/v1/models",
@@ -395,6 +425,49 @@ describe("bedrock mantle discovery", () => {
395425
);
396426
});
397427

428+
it("resolves Mantle runtime auth from the cached IAM token marker", async () => {
429+
const tokenProvider = vi.fn(async () => "bedrock-api-key-runtime"); // pragma: allowlist secret
430+
mocks.getTokenProvider.mockReturnValue(tokenProvider);
431+
432+
await generateBearerTokenFromIam({
433+
region: "us-east-1",
434+
now: () => 1000,
435+
});
436+
437+
await expect(
438+
resolveMantleRuntimeBearerToken({
439+
apiKey: MANTLE_IAM_TOKEN_MARKER,
440+
env: {
441+
AWS_REGION: "us-east-1",
442+
} as NodeJS.ProcessEnv,
443+
now: () => 2000,
444+
}),
445+
).resolves.toMatchObject({
446+
apiKey: "bedrock-api-key-runtime",
447+
expiresAt: 1000 + 7200_000,
448+
});
449+
expect(tokenProvider).toHaveBeenCalledTimes(1);
450+
});
451+
452+
it("generates a fresh Mantle runtime IAM token when the cache is cold", async () => {
453+
const tokenProvider = vi.fn(async () => "bedrock-api-key-fresh"); // pragma: allowlist secret
454+
mocks.getTokenProvider.mockReturnValue(tokenProvider);
455+
456+
await expect(
457+
resolveMantleRuntimeBearerToken({
458+
apiKey: MANTLE_IAM_TOKEN_MARKER,
459+
env: {
460+
AWS_REGION: "us-east-1",
461+
} as NodeJS.ProcessEnv,
462+
now: () => 5000,
463+
}),
464+
).resolves.toMatchObject({
465+
apiKey: "bedrock-api-key-fresh",
466+
expiresAt: 5000 + 7200_000,
467+
});
468+
expect(tokenProvider).toHaveBeenCalledTimes(1);
469+
});
470+
398471
it("returns null for unsupported regions", async () => {
399472
const provider = await resolveImplicitMantleProvider({
400473
env: {

‎extensions/amazon-bedrock-mantle/discovery.ts‎

Lines changed: 62 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ const DEFAULT_COST = {
1818
const DEFAULT_CONTEXT_WINDOW = 32000;
1919
const DEFAULT_MAX_TOKENS = 4096;
2020
const DEFAULT_REFRESH_INTERVAL_SECONDS = 3600; // 1 hour
21+
export const MANTLE_IAM_TOKEN_MARKER = "__amazon_bedrock_mantle_iam__";
2122

2223
// ---------------------------------------------------------------------------
2324
// Mantle region & endpoint helpers
@@ -69,7 +70,22 @@ export function resolveMantleBearerToken(env: NodeJS.ProcessEnv = process.env):
6970

7071
/** Token cache for IAM-derived bearer tokens, keyed by region. */
7172
const iamTokenCache = new Map<string, { token: string; expiresAt: number }>();
72-
const IAM_TOKEN_TTL_MS = 3600_000; // Refresh every 1 hour (tokens valid up to 12h)
73+
const IAM_TOKEN_TTL_MS = 7200_000; // Matches the 2h token lifetime we request below.
74+
75+
function resolveMantleRegion(env: NodeJS.ProcessEnv): string {
76+
return env.AWS_REGION ?? env.AWS_DEFAULT_REGION ?? "us-east-1";
77+
}
78+
79+
function getCachedIamTokenEntry(
80+
region: string,
81+
now: number = Date.now(),
82+
): { token: string; expiresAt: number } | undefined {
83+
const cached = iamTokenCache.get(region);
84+
if (cached && cached.expiresAt > now) {
85+
return cached;
86+
}
87+
return undefined;
88+
}
7389

7490
/**
7591
* Generate a bearer token from IAM credentials using `@aws/bedrock-token-generator`.
@@ -82,9 +98,9 @@ export async function generateBearerTokenFromIam(params: {
8298
now?: () => number;
8399
}): Promise<string | undefined> {
84100
const now = params.now?.() ?? Date.now();
85-
const cached = iamTokenCache.get(params.region);
101+
const cached = getCachedIamTokenEntry(params.region, now);
86102

87-
if (cached && cached.expiresAt > now) {
103+
if (cached) {
88104
return cached.token;
89105
}
90106

@@ -110,6 +126,47 @@ export async function generateBearerTokenFromIam(params: {
110126
}
111127
}
112128

129+
/**
130+
* Read a cached IAM bearer token for the given region (sync, no generation).
131+
*
132+
* Returns the token if it exists and has not expired, undefined otherwise.
133+
* Used by Mantle runtime auth and tests to inspect the current cache.
134+
*/
135+
export function getCachedIamToken(region: string): string | undefined {
136+
return getCachedIamTokenEntry(region)?.token;
137+
}
138+
139+
export async function resolveMantleRuntimeBearerToken(params: {
140+
apiKey: string;
141+
env?: NodeJS.ProcessEnv;
142+
now?: () => number;
143+
}): Promise<{ apiKey: string; expiresAt?: number } | undefined> {
144+
if (params.apiKey !== MANTLE_IAM_TOKEN_MARKER) {
145+
return { apiKey: params.apiKey };
146+
}
147+
const now = params.now?.() ?? Date.now();
148+
const region = resolveMantleRegion(params.env ?? process.env);
149+
const cached = getCachedIamTokenEntry(region, now);
150+
if (cached) {
151+
return {
152+
apiKey: cached.token,
153+
expiresAt: cached.expiresAt,
154+
};
155+
}
156+
const token = await generateBearerTokenFromIam({
157+
region,
158+
now: params.now,
159+
});
160+
if (!token) {
161+
return undefined;
162+
}
163+
const refreshed = getCachedIamTokenEntry(region, now);
164+
return {
165+
apiKey: refreshed?.token ?? token,
166+
expiresAt: refreshed?.expiresAt ?? now + IAM_TOKEN_TTL_MS,
167+
};
168+
}
169+
113170
/** Reset the IAM token cache (for testing). */
114171
export function resetIamTokenCacheForTest(): void {
115172
iamTokenCache.clear();
@@ -259,7 +316,7 @@ export async function resolveImplicitMantleProvider(params: {
259316
fetchFn?: typeof fetch;
260317
}): Promise<ModelProviderConfig | null> {
261318
const env = params.env ?? process.env;
262-
const region = env.AWS_REGION ?? env.AWS_DEFAULT_REGION ?? "us-east-1";
319+
const region = resolveMantleRegion(env);
263320
const explicitBearerToken = resolveMantleBearerToken(env);
264321

265322
if (!isSupportedRegion(region)) {
@@ -290,7 +347,7 @@ export async function resolveImplicitMantleProvider(params: {
290347
baseUrl: `${mantleEndpoint(region)}/v1`,
291348
api: "openai-completions",
292349
auth: "api-key",
293-
apiKey: explicitBearerToken ? "env:AWS_BEARER_TOKEN_BEDROCK" : bearerToken,
350+
apiKey: explicitBearerToken ? "env:AWS_BEARER_TOKEN_BEDROCK" : MANTLE_IAM_TOKEN_MARKER,
294351
models,
295352
};
296353
}

‎extensions/amazon-bedrock-mantle/register.sync.runtime.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import type { OpenClawPluginApi } from "openclaw/plugin-sdk/plugin-entry";
22
import {
33
mergeImplicitMantleProvider,
44
resolveImplicitMantleProvider,
5+
resolveMantleRuntimeBearerToken,
56
resolveMantleBearerToken,
67
} from "./discovery.js";
78

@@ -31,7 +32,12 @@ export function registerBedrockMantlePlugin(api: OpenClawPluginApi): void {
3132
},
3233
},
3334
resolveConfigApiKey: ({ env }) =>
34-
resolveMantleBearerToken(env) ? "AWS_BEARER_TOKEN_BEDROCK" : undefined,
35+
resolveMantleBearerToken(env) ? "env:AWS_BEARER_TOKEN_BEDROCK" : undefined,
36+
prepareRuntimeAuth: async ({ apiKey, env }) =>
37+
await resolveMantleRuntimeBearerToken({
38+
apiKey,
39+
env,
40+
}),
3541
matchesContextOverflowError: ({ errorMessage }) =>
3642
/context_length_exceeded|max.*tokens.*exceeded/i.test(errorMessage),
3743
classifyFailoverReason: ({ errorMessage }) => {

‎src/agents/simple-completion-runtime.test.ts‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ const hoisted = vi.hoisted(() => ({
66
applyLocalNoAuthHeaderOverrideMock: vi.fn(),
77
setRuntimeApiKeyMock: vi.fn(),
88
resolveCopilotApiTokenMock: vi.fn(),
9+
prepareProviderRuntimeAuthMock: vi.fn(),
910
}));
1011

1112
vi.mock("./pi-embedded-runner/model.js", () => ({
@@ -21,6 +22,10 @@ vi.mock("./github-copilot-token.js", () => ({
2122
resolveCopilotApiToken: hoisted.resolveCopilotApiTokenMock,
2223
}));
2324

25+
vi.mock("../plugins/provider-runtime.runtime.js", () => ({
26+
prepareProviderRuntimeAuth: hoisted.prepareProviderRuntimeAuthMock,
27+
}));
28+
2429
let prepareSimpleCompletionModel: typeof import("./simple-completion-runtime.js").prepareSimpleCompletionModel;
2530

2631
beforeAll(async () => {
@@ -33,6 +38,7 @@ beforeEach(() => {
3338
hoisted.applyLocalNoAuthHeaderOverrideMock.mockReset();
3439
hoisted.setRuntimeApiKeyMock.mockReset();
3540
hoisted.resolveCopilotApiTokenMock.mockReset();
41+
hoisted.prepareProviderRuntimeAuthMock.mockReset();
3642

3743
hoisted.applyLocalNoAuthHeaderOverrideMock.mockImplementation((model: unknown) => model);
3844

@@ -57,6 +63,7 @@ beforeEach(() => {
5763
source: "cache:/tmp/copilot-token.json",
5864
baseUrl: "https://api.individual.githubcopilot.com",
5965
});
66+
hoisted.prepareProviderRuntimeAuthMock.mockResolvedValue(undefined);
6067
});
6168

6269
describe("prepareSimpleCompletionModel", () => {
@@ -340,4 +347,62 @@ describe("prepareSimpleCompletionModel", () => {
340347
}),
341348
);
342349
});
350+
351+
it("applies provider runtime auth before storing simple-completion credentials", async () => {
352+
hoisted.resolveModelMock.mockReturnValueOnce({
353+
model: {
354+
provider: "amazon-bedrock-mantle",
355+
id: "anthropic.claude-opus-4-7",
356+
baseUrl: "https://bedrock-mantle.us-east-1.api.aws/anthropic",
357+
},
358+
authStorage: {
359+
setRuntimeApiKey: hoisted.setRuntimeApiKeyMock,
360+
},
361+
modelRegistry: {},
362+
});
363+
hoisted.getApiKeyForModelMock.mockResolvedValueOnce({
364+
apiKey: "__amazon_bedrock_mantle_iam__",
365+
source: "models.providers.amazon-bedrock-mantle.apiKey",
366+
mode: "api-key",
367+
profileId: "mantle",
368+
});
369+
hoisted.prepareProviderRuntimeAuthMock.mockResolvedValueOnce({
370+
apiKey: "bedrock-runtime-token",
371+
baseUrl: "https://bedrock-mantle.us-east-1.api.aws/anthropic",
372+
});
373+
374+
const result = await prepareSimpleCompletionModel({
375+
cfg: undefined,
376+
provider: "amazon-bedrock-mantle",
377+
modelId: "anthropic.claude-opus-4-7",
378+
agentDir: "/tmp/openclaw-agent",
379+
});
380+
381+
expect(hoisted.prepareProviderRuntimeAuthMock).toHaveBeenCalledWith(
382+
expect.objectContaining({
383+
provider: "amazon-bedrock-mantle",
384+
workspaceDir: "/tmp/openclaw-agent",
385+
context: expect.objectContaining({
386+
apiKey: "__amazon_bedrock_mantle_iam__",
387+
authMode: "api-key",
388+
modelId: "anthropic.claude-opus-4-7",
389+
profileId: "mantle",
390+
}),
391+
}),
392+
);
393+
expect(hoisted.setRuntimeApiKeyMock).toHaveBeenCalledWith(
394+
"amazon-bedrock-mantle",
395+
"bedrock-runtime-token",
396+
);
397+
expect(result).toEqual(
398+
expect.objectContaining({
399+
model: expect.objectContaining({
400+
baseUrl: "https://bedrock-mantle.us-east-1.api.aws/anthropic",
401+
}),
402+
auth: expect.objectContaining({
403+
apiKey: "bedrock-runtime-token",
404+
}),
405+
}),
406+
);
407+
});
343408
});

0 commit comments

Comments
 (0)