Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions apps/web/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,10 @@
"@ai-sdk/openai": "^3.0.41",
"@ai-sdk/openai-compatible": "^2.0.35",
"@anthropic-ai/sdk": "^0.90.0",
"@aws-crypto/sha256-js": "^5.2.0",
Comment thread
kilo-code-bot[bot] marked this conversation as resolved.
"@aws-sdk/client-s3": "^3.1009.0",
"@aws-sdk/s3-request-presigner": "^3.1009.0",
"@aws-sdk/signature-v4": "^3.374.0",
"@chat-adapter/github": "4.27.0",
"@chat-adapter/slack": "^4.27.0",
"@chat-adapter/state-memory": "^4.27.0",
Expand Down
54 changes: 54 additions & 0 deletions apps/web/src/lib/ai-gateway/providers/bedrock-signer.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
import { SignatureV4 } from '@aws-sdk/signature-v4';
import { Sha256 } from '@aws-crypto/sha256-js';
import type { CustomLlmAwsBedrock } from '@kilocode/db';
import type { SignedRequest, SignRequestArgs } from '@/lib/ai-gateway/providers/types';

export function makeBedrockSignRequest(
credentials: CustomLlmAwsBedrock,
modelId: string
): (args: SignRequestArgs) => Promise<SignedRequest> {
const signer = new SignatureV4({
credentials: {
accessKeyId: credentials.access_key_id,
secretAccessKey: credentials.secret_access_key,
},
region: credentials.region,
service: 'bedrock',
sha256: Sha256,
});

const hostname = `bedrock-runtime.${credentials.region}.amazonaws.com`;

return async ({ method, body }) => {
const isStreaming = parseStreamFlag(body);
const path = `/model/${encodeURIComponent(modelId)}/${
isStreaming ? 'invoke-with-response-stream' : 'invoke'
}`;

const signed = await signer.sign({
method,
hostname,
path,
protocol: 'https:',
headers: {
'Content-Type': 'application/json',
host: hostname,
},
body,
});

return {
url: `https://${hostname}${signed.path ?? path}`,
headers: signed.headers,
};
};
}

function parseStreamFlag(body: string): boolean {
try {
const parsed = JSON.parse(body);
return parsed !== null && typeof parsed === 'object' && parsed.stream === true;
} catch {
return false;
}
}
4 changes: 4 additions & 0 deletions apps/web/src/lib/ai-gateway/providers/get-provider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import {
addCacheBreakpoints,
injectReasoningIntoContent,
} from '@/lib/ai-gateway/providers/openrouter/request-helpers';
import { makeBedrockSignRequest } from '@/lib/ai-gateway/providers/bedrock-signer';

function inferSupportedChatApis(
aiSdkProvider: CustomLlmProvider | undefined,
Expand Down Expand Up @@ -94,6 +95,8 @@ async function checkCustomLlm(
if (!customLlm || !customLlm.organization_ids.includes(organizationId)) {
return null;
}
const bedrock = customLlm.aws_bedrock;
const signRequest = bedrock ? makeBedrockSignRequest(bedrock, customLlm.internal_id) : undefined;
return {
provider: {
id: 'custom',
Expand All @@ -120,6 +123,7 @@ async function checkCustomLlm(
injectReasoningIntoContent(context.request);
}
},
signRequest,
},
userByok: null,
bypassAccessCheck: true,
Expand Down
15 changes: 14 additions & 1 deletion apps/web/src/lib/ai-gateway/providers/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,23 @@ export type TransformRequestContext = {

export type GatewayChatApiKind = GatewayRequest['kind'];

export type SignRequestArgs = {
method: string;
url: string;
body: string;
};

export type SignedRequest = {
url?: string;
headers: Record<string, string>;
};

export type Provider = {
id: ProviderId;
apiUrl: string;
apiKey: string;
apiKey: string | undefined;
supportedChatApis: ReadonlyArray<GatewayChatApiKind>;
transformRequest(context: TransformRequestContext): void;
// When set, replaces the default `Authorization: Bearer ${apiKey}` header.
signRequest?(args: SignRequestArgs): Promise<SignedRequest>;
};
20 changes: 17 additions & 3 deletions apps/web/src/lib/ai-gateway/providers/upstream-request.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,14 +31,28 @@ export async function upstreamRequest({
for (const [key, value] of Object.entries(ATTRIBUTION_HEADERS)) {
headers.set(key, value);
}
headers.set('Authorization', `Bearer ${provider.apiKey}`);
headers.set('Content-Type', 'application/json');

Object.entries(extraHeaders).forEach(([key, value]) => {
headers.set(key, value);
});

const targetUrl = `${provider.apiUrl}${path}${search}`;
let targetUrl = `${provider.apiUrl}${path}${search}`;
const serializedBody = JSON.stringify(body);

if (provider.signRequest) {
const signed = await provider.signRequest({
method,
url: targetUrl,
body: serializedBody,
});
if (signed.url) targetUrl = signed.url;
for (const [key, value] of Object.entries(signed.headers)) {
headers.set(key, value);
}
} else if (provider.apiKey) {
headers.set('Authorization', `Bearer ${provider.apiKey}`);
Comment thread
kilo-code-bot[bot] marked this conversation as resolved.
}

const TEN_MINUTES_MS = 10 * 60 * 1000;
const timeoutSignal = AbortSignal.timeout(TEN_MINUTES_MS);
Expand All @@ -47,7 +61,7 @@ export async function upstreamRequest({
return await fetch(targetUrl, {
method,
headers,
body: JSON.stringify(body),
body: serializedBody,
// @ts-expect-error see https://github.com/node-fetch/node-fetch/issues/1769
duplex: 'half',
signal: combinedSignal,
Expand Down
13 changes: 12 additions & 1 deletion packages/db/src/schema-types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -951,13 +951,23 @@ export const CustomLlmPricingSchema = z.object({

export type CustomLlmPricing = z.infer<typeof CustomLlmPricingSchema>;

// When set, upstream requests are SigV4-signed instead of using Bearer auth;
// `internal_id` is used as the Bedrock model id in the request path.
export const CustomLlmAwsBedrockSchema = z.object({
access_key_id: z.string(),
secret_access_key: z.string(),
region: z.string(),
});

export type CustomLlmAwsBedrock = z.infer<typeof CustomLlmAwsBedrockSchema>;

export const CustomLlmDefinitionSchema = z.object({
internal_id: z.string(),
display_name: z.string(),
context_length: z.number(),
max_completion_tokens: z.number(),
base_url: z.string(),
api_key: z.string(),
api_key: z.string().optional(),
organization_ids: z.array(z.string()),
supports_image_input: z.boolean().optional(),
add_cache_breakpoints: z.boolean().optional(),
Expand All @@ -968,6 +978,7 @@ export const CustomLlmDefinitionSchema = z.object({
opencode_settings: OpenCodeSettingsSchema.optional(),
openclaw_settings: OpenClawModelSettingsSchema.optional(),
pricing: CustomLlmPricingSchema.optional(),
aws_bedrock: CustomLlmAwsBedrockSchema.optional(),
});

export type CustomLlmDefinition = z.infer<typeof CustomLlmDefinitionSchema>;
Expand Down
Loading