Skip to content

Commit c3639f7

Browse files
aidandaly24aidandaly24
committed
refactor(runtime): split invoke transports
1 parent fa77bd0 commit c3639f7

1 file changed

Lines changed: 103 additions & 84 deletions

File tree

‎src/core/runtime.tsx‎

Lines changed: 103 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ export class RuntimeClient implements CoreRuntimeClient {
3535
options: CoreOptions,
3636
signal?: AbortSignal,
3737
): Promise<RuntimeInvokeResponse> {
38-
const { runtimeId, applicationHeaders, bearerToken, ...input } = request;
38+
const { runtimeId, bearerToken } = request;
3939
const logger = this.logger.child({
4040
operation: "invokeRuntime",
4141
authMode: bearerToken === undefined ? "IAM" : "CUSTOM_JWT",
@@ -44,92 +44,111 @@ export class RuntimeClient implements CoreRuntimeClient {
4444
region: options.region,
4545
});
4646
if (bearerToken !== undefined) {
47-
const client = this.clients.data(toClientConfig(options));
48-
const endpoint = client.config.endpointProvider({
49-
Region: options.region,
50-
Endpoint: options.endpointUrl,
51-
});
52-
const url = new URL(endpoint.url);
53-
if (url.protocol !== "https:") {
54-
throw new TypeError("CUSTOM_JWT requires an HTTPS endpoint");
55-
}
56-
url.pathname = `${url.pathname.replace(/\/?$/, "/")}runtimes/${encodeURIComponent(runtimeId)}/invocations`;
57-
url.search = new URLSearchParams({
58-
accountId: request.accountId,
59-
qualifier: request.qualifier,
60-
}).toString();
61-
const headers = new Headers(applicationHeaders);
62-
try {
63-
headers.set("Authorization", `Bearer ${bearerToken}`);
64-
} catch {
65-
throw new TypeError("Invalid bearer token");
66-
}
67-
try {
68-
for (const [name, value] of [
69-
["Content-Type", request.contentType],
70-
["Accept", request.accept],
71-
["Mcp-Session-Id", request.mcpSessionId],
72-
["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id", request.runtimeSessionId],
73-
["Mcp-Protocol-Version", request.mcpProtocolVersion],
74-
["Mcp-Method", request.mcpMethod],
75-
["Mcp-Name", request.mcpName],
76-
["X-Amzn-Bedrock-AgentCore-Runtime-User-Id", request.runtimeUserId],
77-
["X-Amzn-Trace-Id", request.traceId],
78-
["traceparent", request.traceParent],
79-
["tracestate", request.traceState],
80-
["baggage", request.baggage],
81-
] as const) {
82-
if (value !== undefined) headers.set(name, value);
83-
}
84-
} catch {
85-
throw new TypeError("Invalid Runtime request header");
86-
}
87-
let response: Response;
88-
try {
89-
response = await this.fetch(url, {
90-
method: "POST",
91-
redirect: "error",
92-
headers,
93-
body: request.payload as RequestInit["body"],
94-
signal,
95-
});
96-
} catch (error) {
97-
if (signal?.aborted) throw signal.reason ?? error;
98-
logger
99-
.child({
100-
errorName:
101-
error instanceof TypeError
102-
? "TypeError"
103-
: error instanceof Error
104-
? "Error"
105-
: typeof error,
106-
})
107-
.debug("Runtime invocation transport failed");
108-
throw new Error("Runtime invocation failed");
109-
}
110-
if (!response.ok) {
111-
logger
112-
.child({ httpStatusCode: response.status })
113-
.debug("Runtime invocation returned a non-success response");
114-
await response.body?.cancel().catch(() => undefined);
115-
throw new Error(`HTTP ${response.status}`);
47+
return this.invokeRuntimeWithCustomJwt(request, bearerToken, options, logger, signal);
48+
}
49+
return this.invokeRuntimeWithIam(request, options, logger, signal);
50+
}
51+
52+
private async invokeRuntimeWithCustomJwt(
53+
request: RuntimeInvokeRequest,
54+
bearerToken: string,
55+
options: CoreOptions,
56+
logger: Logger,
57+
signal?: AbortSignal,
58+
): Promise<RuntimeInvokeResponse> {
59+
const client = this.clients.data(toClientConfig(options));
60+
const endpoint = client.config.endpointProvider({
61+
Region: options.region,
62+
Endpoint: options.endpointUrl,
63+
});
64+
const url = new URL(endpoint.url);
65+
if (url.protocol !== "https:") {
66+
throw new TypeError("CUSTOM_JWT requires an HTTPS endpoint");
67+
}
68+
url.pathname = `${url.pathname.replace(/\/?$/, "/")}runtimes/${encodeURIComponent(request.runtimeId)}/invocations`;
69+
url.search = new URLSearchParams({
70+
accountId: request.accountId,
71+
qualifier: request.qualifier,
72+
}).toString();
73+
const headers = new Headers(request.applicationHeaders);
74+
try {
75+
headers.set("Authorization", `Bearer ${bearerToken}`);
76+
} catch {
77+
throw new TypeError("Invalid bearer token");
78+
}
79+
try {
80+
for (const [name, value] of [
81+
["Content-Type", request.contentType],
82+
["Accept", request.accept],
83+
["Mcp-Session-Id", request.mcpSessionId],
84+
["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id", request.runtimeSessionId],
85+
["Mcp-Protocol-Version", request.mcpProtocolVersion],
86+
["Mcp-Method", request.mcpMethod],
87+
["Mcp-Name", request.mcpName],
88+
["X-Amzn-Bedrock-AgentCore-Runtime-User-Id", request.runtimeUserId],
89+
["X-Amzn-Trace-Id", request.traceId],
90+
["traceparent", request.traceParent],
91+
["tracestate", request.traceState],
92+
["baggage", request.baggage],
93+
] as const) {
94+
if (value !== undefined) headers.set(name, value);
11695
}
117-
const body = (response.body as AsyncIterable<Uint8Array> | null) ?? emptyBody();
118-
return {
119-
statusCode: response.status,
120-
contentType: response.headers.get("content-type") ?? "",
121-
runtimeSessionId:
122-
response.headers.get("x-amzn-bedrock-agentcore-runtime-session-id") ?? undefined,
123-
mcpSessionId: response.headers.get("mcp-session-id") ?? undefined,
124-
mcpProtocolVersion: response.headers.get("mcp-protocol-version") ?? undefined,
125-
traceId: response.headers.get("x-amzn-trace-id") ?? undefined,
126-
traceParent: response.headers.get("traceparent") ?? undefined,
127-
traceState: response.headers.get("tracestate") ?? undefined,
128-
baggage: response.headers.get("baggage") ?? undefined,
129-
body: signal ? abortable(body, signal) : body,
130-
};
96+
} catch {
97+
throw new TypeError("Invalid Runtime request header");
13198
}
99+
let response: Response;
100+
try {
101+
response = await this.fetch(url, {
102+
method: "POST",
103+
redirect: "error",
104+
headers,
105+
body: request.payload as RequestInit["body"],
106+
signal,
107+
});
108+
} catch (error) {
109+
if (signal?.aborted) throw signal.reason ?? error;
110+
logger
111+
.child({
112+
errorName:
113+
error instanceof TypeError
114+
? "TypeError"
115+
: error instanceof Error
116+
? "Error"
117+
: typeof error,
118+
})
119+
.debug("Runtime invocation transport failed");
120+
throw new Error("Runtime invocation failed");
121+
}
122+
if (!response.ok) {
123+
logger
124+
.child({ httpStatusCode: response.status })
125+
.debug("Runtime invocation returned a non-success response");
126+
await response.body?.cancel().catch(() => undefined);
127+
throw new Error(`HTTP ${response.status}`);
128+
}
129+
const body = (response.body as AsyncIterable<Uint8Array> | null) ?? emptyBody();
130+
return {
131+
statusCode: response.status,
132+
contentType: response.headers.get("content-type") ?? "",
133+
runtimeSessionId:
134+
response.headers.get("x-amzn-bedrock-agentcore-runtime-session-id") ?? undefined,
135+
mcpSessionId: response.headers.get("mcp-session-id") ?? undefined,
136+
mcpProtocolVersion: response.headers.get("mcp-protocol-version") ?? undefined,
137+
traceId: response.headers.get("x-amzn-trace-id") ?? undefined,
138+
traceParent: response.headers.get("traceparent") ?? undefined,
139+
traceState: response.headers.get("tracestate") ?? undefined,
140+
baggage: response.headers.get("baggage") ?? undefined,
141+
body: signal ? abortable(body, signal) : body,
142+
};
143+
}
132144

145+
private async invokeRuntimeWithIam(
146+
request: RuntimeInvokeRequest,
147+
options: CoreOptions,
148+
logger: Logger,
149+
signal?: AbortSignal,
150+
): Promise<RuntimeInvokeResponse> {
151+
const { runtimeId, applicationHeaders, bearerToken: _bearerToken, ...input } = request;
133152
const command = new InvokeAgentRuntimeCommand({ ...input, agentRuntimeArn: runtimeId });
134153
if (applicationHeaders?.length) {
135154
command.middlewareStack.add(

0 commit comments

Comments
 (0)