Skip to content
129 changes: 112 additions & 17 deletions src/adapters/kiro.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import type {
import type { ProviderAdapter } from "./base";
import type { AdapterFetchContext, AdapterRequest } from "./base";
import { extractKiroImages, normalizeKiroImages, type KiroImage } from "./kiro-images";
import { sniffImageDimensions } from "./anthropic-image-guard";
import { fetchKiroWithRetry } from "./kiro-retry";
import { convertKiroToolContext } from "./kiro-tools";
import { neutralizeIdentity } from "./identity";
Expand Down Expand Up @@ -134,6 +135,48 @@ function messageLogText(msg: OcxMessage): string {
}).filter(Boolean).join("\n");
}

function estimateKiroImageTokens(image: KiroImage): number {
const dimensions = sniffImageDimensions(image.source.bytes);
if (dimensions) {
return Math.max(256, Math.ceil(dimensions.width * dimensions.height / 750));
}
const decodedBytes = Math.floor(image.source.bytes.length * 3 / 4);
return Math.max(256, Math.ceil(decodedBytes / 512));
}

function estimateKiroPayloadInputTokens(payload: Record<string, unknown>, modelId: string): number {
const conversationState = (payload as {
conversationState?: {
history?: KiroHistoryEntry[];
currentMessage?: KiroHistoryEntry;
};
}).conversationState;
if (!conversationState) return 0;

const parts: string[] = [];
let imageTokens = 0;
const entries = [
...(conversationState.history ?? []),
...(conversationState.currentMessage ? [conversationState.currentMessage] : []),
];
for (const entry of entries) {
const user = entry.userInputMessage;
if (user) {
if (user.content) parts.push(user.content);
for (const image of user.images ?? []) imageTokens += estimateKiroImageTokens(image);
const context = user.userInputMessageContext;
if (context?.tools?.length) parts.push(serializeForUsage(context.tools));
if (context?.toolResults?.length) parts.push(serializeForUsage(context.toolResults));
}
const assistant = entry.assistantResponseMessage;
if (assistant) {
if (assistant.content) parts.push(assistant.content);
if (assistant.toolUses?.length) parts.push(serializeForUsage(assistant.toolUses));
}
}
return estimateTokens(parts.join("\n"), modelId) + imageTokens;
}

function shouldCountStablePromptOverhead(parsed: OcxParsedRequest): boolean {
return !parsed.previousResponseId && !parsed.context.messages.some(m => m.role === "assistant");
}
Expand All @@ -158,15 +201,11 @@ function estimateKiroLogInputTokens(parsed: OcxParsedRequest): number {
return Math.max(estimateKiroInputTokens(parsed), estimateTokens(parts.join("\n"), parsed.modelId));
}

function configuredKiroContextWindow(provider: OcxProviderConfig, modelId: string | undefined): number | undefined {
function kiroUpstreamContextWindow(modelId: string | undefined): number | undefined {
if (!modelId) return undefined;
const normalizedModelId = normalizeKiroModelId(modelId);
if (normalizedModelId === "auto") return undefined;
Comment thread
coseung2 marked this conversation as resolved.
const window =
modelRecordValue(provider.modelContextWindows, modelId)
?? modelRecordValue(provider.modelContextWindows, normalizedModelId)
?? provider.contextWindow
?? modelRecordValue(KIRO_MODEL_CONTEXT_WINDOWS, modelId)
const window = modelRecordValue(KIRO_MODEL_CONTEXT_WINDOWS, modelId)
?? modelRecordValue(KIRO_MODEL_CONTEXT_WINDOWS, normalizedModelId);
return typeof window === "number" && Number.isFinite(window) && window > 0 ? window : undefined;
}
Expand Down Expand Up @@ -465,17 +504,26 @@ interface KiroAttemptResult {
interface KiroFallbackAttempt {
response: Response;
inputTokens: number;
contextInputEstimate: number;
nameMap: Map<string, string>;
conversationId: string;
}

interface KiroContextWindowState {
value?: number;
}

type KiroFallbackFactory = (
conversationId: string | undefined,
assistantText: string,
sawReasoning: boolean,
) => Promise<KiroFallbackAttempt>;

function mergeKiroUsage(first: OcxUsage | undefined, second: OcxUsage | undefined): OcxUsage | undefined {
function mergeKiroUsage(
first: OcxUsage | undefined,
second: OcxUsage | undefined,
preserveFirstContextGrowth = false,
): OcxUsage | undefined {
if (!first) return second;
if (!second) return first;
const sumOptional = (key: keyof OcxUsage): number | undefined => {
Expand All @@ -488,9 +536,21 @@ function mergeKiroUsage(first: OcxUsage | undefined, second: OcxUsage | undefine
const totalTokens = typeof first.totalTokens === "number" && typeof second.totalTokens === "number"
? first.totalTokens + second.totalTokens
: undefined;
const carriedContextTotal = preserveFirstContextGrowth && typeof first.contextTotalTokens === "number"
? first.contextTotalTokens + second.outputTokens
: undefined;
return {
inputTokens: first.inputTokens + second.inputTokens,
outputTokens: first.outputTokens + second.outputTokens,
...(typeof first.contextTotalTokens === "number" || typeof second.contextTotalTokens === "number"
? {
contextTotalTokens: Math.max(
first.contextTotalTokens ?? 0,
second.contextTotalTokens ?? 0,
carriedContextTotal ?? 0,
Comment thread
coseung2 marked this conversation as resolved.
),
}
: {}),
Comment thread
coseung2 marked this conversation as resolved.
...(totalTokens !== undefined ? { totalTokens } : {}),
...(sumOptional("cachedInputTokens") !== undefined ? { cachedInputTokens: sumOptional("cachedInputTokens") } : {}),
...(sumOptional("cacheReadInputTokens") !== undefined ? { cacheReadInputTokens: sumOptional("cacheReadInputTokens") } : {}),
Expand Down Expand Up @@ -530,10 +590,11 @@ async function* parseKiroAttempt(
mode: KiroCompletionMode,
modelId: string | undefined,
inputTokens: number,
contextWindow: number | undefined,
contextWindowState: KiroContextWindowState,
nameMap: Map<string, string> | undefined,
conversationId: string | undefined,
previousAssistantText?: string,
contextInputEstimate?: number,
): AsyncGenerator<AdapterEvent, KiroAttemptResult> {
const emptyResult = (): KiroAttemptResult => ({ assistantText: "", sawReasoning: false });
if (!response.body) {
Expand All @@ -560,11 +621,28 @@ async function* parseKiroAttempt(
const providerState = (): { kiro: { conversationId: string } } | undefined =>
returnedConversationId ? { kiro: { conversationId: returnedConversationId } } : undefined;

const usage = (): OcxUsage => authoritativeUsage ?? ({
const contextUsageTotalFloor = (): number | undefined => {
if (contextUsagePercentage === undefined || !contextWindowState.value) return undefined;
const floor = Math.ceil(contextWindowState.value * Math.min(contextUsagePercentage, 100) / 100);
return Number.isFinite(floor) && floor > 0 ? floor : undefined;
};
const usage = (): OcxUsage => {
const base = authoritativeUsage ?? {
inputTokens,
outputTokens: estimateTokens(outputChars, modelId),
estimated: true,
});
};
const estimatedContextTotal = contextInputEstimate !== undefined
? contextInputEstimate + base.outputTokens
: undefined;
const authoritativeTurnTotal = base.inputTokens + base.outputTokens;
const contextTotal = Math.max(
estimatedContextTotal ?? 0,
contextUsageTotalFloor() ?? 0,
authoritativeTurnTotal,
);
return contextTotal > 0 ? { ...base, contextTotalTokens: contextTotal } : base;
};

const classifiedTerminal = (failure: KiroErrorClassification): AdapterEvent => ({
type: "error",
Expand Down Expand Up @@ -743,6 +821,9 @@ async function* parseKiroAttempt(
if (isValidKiroConversationId(ev.conversationId)) returnedConversationId = ev.conversationId;
break;
case "content":
if (ev.modelId) {
contextWindowState.value = kiroUpstreamContextWindow(ev.modelId) ?? contextWindowState.value;
}
if (open) {
open = null;
return { assistantText, sawReasoning, terminal: protocolTerminal(kiroTruncationErrorMessage("content arrived before tool stop")) };
Expand Down Expand Up @@ -843,7 +924,7 @@ async function* parseKiroAttempt(
if (contextUsagePercentage !== undefined) {
debugProviderDiagnostic("kiro", "context_usage", {
contextUsagePercentage,
...(contextWindow ? { configuredContextWindow: contextWindow } : {}),
...(contextWindowState.value ? { upstreamContextWindow: contextWindowState.value } : {}),
});
}
debugProviderDiagnostic("kiro", "attempt_complete", {
Expand Down Expand Up @@ -970,15 +1051,19 @@ export async function* parseKiroStream(
conversationId?: string,
completionMode: KiroCompletionMode = "disabled",
fallbackFactory?: KiroFallbackFactory,
contextInputEstimate?: number,
): AsyncGenerator<AdapterEvent> {
const contextWindowState: KiroContextWindowState = { value: contextWindow };
const first = parseKiroAttempt(
response,
completionMode,
modelId,
inputTokens,
contextWindow,
contextWindowState,
nameMap,
conversationId,
undefined,
contextInputEstimate,
);
let firstNext = await first.next();
while (!firstNext.done) {
Expand Down Expand Up @@ -1039,10 +1124,11 @@ export async function* parseKiroStream(
"text_fallback",
modelId,
fallback.inputTokens,
contextWindow,
contextWindowState,
fallback.nameMap,
fallback.conversationId,
firstResult.assistantText,
fallback.contextInputEstimate,
);
let secondNext = await second.next();
while (!secondNext.done) {
Expand All @@ -1054,23 +1140,24 @@ export async function* parseKiroStream(
yield retryableKiroIncomplete(
"empty_kiro_fallback",
"Kiro's bounded completion retry ended without a terminal result",
mergeKiroUsage(firstResult.usage, secondResult.usage) ?? { inputTokens, outputTokens: 0, estimated: true },
mergeKiroUsage(firstResult.usage, secondResult.usage, Boolean(firstResult.assistantText))
?? { inputTokens, outputTokens: 0, estimated: true },
secondResult.providerState ?? firstResult.providerState,
);
return;
}
if (secondResult.terminal.type === "done" || secondResult.terminal.type === "incomplete") {
yield {
...secondResult.terminal,
usage: mergeKiroUsage(firstResult.usage, secondResult.terminal.usage),
usage: mergeKiroUsage(firstResult.usage, secondResult.terminal.usage, Boolean(firstResult.assistantText)),
providerState: secondResult.terminal.providerState ?? firstResult.providerState,
};
return;
}
yield {
...secondResult.terminal,
...(secondResult.terminal.type === "error"
? { usage: mergeKiroUsage(firstResult.usage, secondResult.terminal.usage) }
? { usage: mergeKiroUsage(firstResult.usage, secondResult.terminal.usage, Boolean(firstResult.assistantText)) }
Comment thread
coderabbitai[bot] marked this conversation as resolved.
: {}),
};
}
Expand All @@ -1080,6 +1167,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
// Per-request closure (resolveAdapter builds a fresh adapter per request — server.ts:440 — so this
// is race-free) carrying the heuristic input-token estimate from buildRequest into the stream.
let inputTokens = 0;
let contextInputEstimate = 0;
let modelId: string | undefined;
let contextWindow: number | undefined;
let toolNameMap: Map<string, string> | undefined;
Expand All @@ -1097,6 +1185,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
conversationId: string;
completionMode: KiroCompletionMode;
inputTokens: number;
contextInputEstimate: number;
}> => {
if (typeof provider.apiKey !== "string" || provider.apiKey.trim() === "") {
throw new Error("kiro token missing — run ocx login kiro");
Expand All @@ -1118,6 +1207,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
if (profileArn) headers["x-amzn-kiro-profile-arn"] = profileArn;
const built = buildKiroPayload(parsed, profileArn, forcedCompletionMode);
await normalizeKiroImages(built.payload);
const contextInputEstimate = estimateKiroPayloadInputTokens(built.payload, parsed.modelId);
Comment thread
coseung2 marked this conversation as resolved.
const body = JSON.stringify(built.payload);
debugProviderDiagnostic("kiro", "request", {
region,
Expand All @@ -1141,6 +1231,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
conversationId: built.conversationId,
completionMode: built.completionMode,
inputTokens: estimateKiroInputTokens(parsed),
contextInputEstimate,
};
};

Expand Down Expand Up @@ -1179,6 +1270,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
return {
response,
inputTokens: retry.inputTokens,
contextInputEstimate: retry.contextInputEstimate,
nameMap: retry.nameMap,
conversationId: retry.conversationId,
};
Expand All @@ -1189,8 +1281,9 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
async buildRequest(parsed: OcxParsedRequest, incoming) {
const built = await build(parsed);
modelId = parsed.modelId;
contextWindow = configuredKiroContextWindow(provider, parsed.modelId);
contextWindow = kiroUpstreamContextWindow(parsed.modelId);
inputTokens = built.inputTokens;
contextInputEstimate = built.contextInputEstimate;
toolNameMap = built.nameMap;
conversationId = built.conversationId;
completionMode = built.completionMode;
Expand All @@ -1209,6 +1302,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
conversationId,
completionMode,
completionMode === "required" ? fallbackFactory : undefined,
contextInputEstimate,
);
},

Expand Down Expand Up @@ -1239,6 +1333,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
conversationId,
completionMode,
completionMode === "required" ? fallbackFactory : undefined,
contextInputEstimate,
)) events.push(e);
return events;
},
Expand Down
19 changes: 14 additions & 5 deletions src/bridge.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,20 +15,29 @@ function sseEvent(name: string, data: Record<string, unknown>): string {

function responsesUsage(usage: OcxUsage | undefined): Record<string, unknown> {
if (!usage) return { input_tokens: 0, output_tokens: 0, total_tokens: 0 };
// inputTokens is already inclusive of cache read/write (types.ts convention).
const inputTokens = usage.inputTokens;
// Stateful providers may report an absolute active-context checkpoint separately from their
// per-attempt usage. Split that checkpoint into input + output without adding output twice.
const inputTokens = usage.contextTotalTokens !== undefined
? Math.max(0, usage.contextTotalTokens - usage.outputTokens)
: usage.inputTokens;
Comment thread
coseung2 marked this conversation as resolved.
const out: Record<string, unknown> = {
input_tokens: inputTokens,
Comment thread
coseung2 marked this conversation as resolved.
output_tokens: usage.outputTokens,
total_tokens: usageDisplayTotalTokens(usage) ?? inputTokens + usage.outputTokens,
total_tokens: usage.contextTotalTokens !== undefined
? usage.contextTotalTokens
: usageDisplayTotalTokens(usage) ?? inputTokens + usage.outputTokens,
};
const inputDetails: Record<string, number> = {};
if (usage.cachedInputTokens !== undefined) {
// cached_tokens carries cache READS only, matching OpenAI semantics.
inputDetails.cached_tokens = usage.cachedInputTokens;
inputDetails.cached_tokens = Math.min(usage.cachedInputTokens, inputTokens);
}
if (usage.cacheCreationInputTokens !== undefined) {
inputDetails.cache_write_tokens = usage.cacheCreationInputTokens;
const cacheRead = inputDetails.cached_tokens ?? 0;
inputDetails.cache_write_tokens = Math.min(
usage.cacheCreationInputTokens,
Math.max(0, inputTokens - cacheRead),
);
}
if (Object.keys(inputDetails).length > 0) {
out.input_tokens_details = inputDetails;
Expand Down
6 changes: 6 additions & 0 deletions src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -308,6 +308,12 @@ export interface OcxUrlCitation {
export interface OcxUsage {
inputTokens: number;
outputTokens: number;
/**
* Absolute active-context size after the response. Stateful providers can expose this separately
* from their per-attempt usage. Responses serialization derives the input side from
* `contextTotalTokens - outputTokens` so output is never added to an absolute checkpoint twice.
*/
contextTotalTokens?: number;
totalTokens?: number;
cachedInputTokens?: number;
cacheReadInputTokens?: number;
Expand Down
Loading
Loading