fix(kiro): improve direct session cache reuse

Reshape Kiro direct requests so resumed client sessions reuse Kiro's
cache-affinity fields instead of starting unrelated CodeWhisperer
conversations.

- keep conversationState.conversationId stable when the client sends an
  explicit session id (x-session-id, session_id, conversation_id, Claude
  Code session metadata)
- add a stable conversationState.agentContinuationId per Kiro session
- send conversationState.agentTaskType: "vibe" and agentMode: "vibe",
  matching the normal Kiro CLI/KAS chat path
- move Kiro thinking instructions into Kiro-compatible systemPrompt /
  additionalModelRequestFields instead of generic top-level thinking
- keep volatile timestamp context out of the top-level systemPrompt; it
  remains only in user content fallback
- suppress additionalModelRequestFields for legacy 4.5-era Claude/Kiro
  models that reject it, while defaulting future Claude/Kiro model ids
  to supported
- preserve Kiro meteringEvent credit usage internally for accounting
  without leaking provider-specific fields into OpenAI-compatible usage
- prevent unrelated headerless Kiro requests from sharing one
  connection-wide continuation
- cap/evict continuation sessions so long-running processes do not grow
  the continuation map unbounded
- treat generated headerless Kiro sessions as one-shot so they do not
  evict real explicit-session continuations
- keep credit-only Kiro metering valid for internal persistence when
  token metrics are unavailable
This commit is contained in:
Edison42
2026-07-16 15:13:28 +07:00
committed by decolua
parent 70e8dc4974
commit 9c58ba645e
10 changed files with 786 additions and 90 deletions

View File

@@ -103,8 +103,16 @@ export function translateRequest(sourceFormat, targetFormat, model, body, stream
}
}
// Normalize thinking to the target provider-native format (config-driven, capability-aware)
applyThinking(targetFormat, model, result, provider, thinkingIntent);
// Normalize thinking to the target provider-native format (config-driven, capability-aware).
// Kiro's GenerateAssistantResponse request does not accept the generic top-level
// `thinking` field; its translators map thinking intent to KAS-compatible
// systemPrompt/additionalModelRequestFields instead.
const kiroThinkingMappedByTranslator =
targetFormat === FORMATS.KIRO &&
(sourceFormat === FORMATS.OPENAI || sourceFormat === FORMATS.CLAUDE);
if (!kiroThinkingMappedByTranslator) {
applyThinking(targetFormat, model, result, provider, thinkingIntent);
}
// Always normalize to clean OpenAI format when target is OpenAI
// This handles hybrid requests (e.g., OpenAI messages + Claude tools)

View File

@@ -24,13 +24,15 @@
*/
import { register } from "../index.js";
import { FORMATS } from "../formats.js";
import { v4 as uuidv4 } from "uuid";
import { applyKiroSessionReplay } from "../../utils/kiroSessionReplay.js";
import { resolveContinuationId, resolveSessionIdentity } from "../../utils/sessionManager.js";
import {
resolveKiroModel,
resolveKiroThinkingBudget,
buildThinkingSystemPrefix,
KIRO_AGENTIC_SYSTEM_PROMPT,
resolveDefaultProfileArn,
buildKiroAdditionalModelRequestFieldsForModel,
} from "../../config/kiroConstants.js";
import { DEFAULT_IMAGE_MIME } from "../schema/index.js";
import { ROLE, CLAUDE_BLOCK } from "../schema/index.js";
@@ -363,6 +365,18 @@ function reconcileOrphanedToolResults(history, currentMessage) {
}
}
function extractClaudeSystemText(system) {
if (!system) return "";
if (typeof system === "string") return system;
if (Array.isArray(system)) {
return system.map((s) => {
if (typeof s === "string") return s;
return s?.text || "";
}).filter(Boolean).join("\n");
}
return "";
}
/**
* Build a Kiro payload directly from a Claude Messages API request body.
*/
@@ -402,62 +416,75 @@ export function claudeToKiroRequest(model, body, stream, credentials) {
? (credentials?.providerSpecificData?.profileArn || "")
: (credentials?.providerSpecificData?.profileArn || resolveDefaultProfileArn(authMethod));
let finalContent = currentMessage?.userInputMessage?.content || "";
// System prompt: pass via native systemInstruction field (Kiro/Q API supports it)
// and also prepend as <instructions> in user content as fallback for upstreams
// that don't support the native field.
let systemInstruction = undefined;
if (body.system) {
let systemText = "";
if (typeof body.system === "string") {
systemText = body.system;
} else if (Array.isArray(body.system)) {
systemText = body.system.map((s) => s.text || "").join("\n");
}
if (systemText) {
systemInstruction = systemText;
finalContent = `<instructions>\n${systemText}\n</instructions>\n\n${finalContent}`;
}
}
// Prefix order: thinking_mode tag, timestamp marker, then agentic prompt.
// Kiro CLI/KAS sends system prompt as top-level `systemPrompt`. Keep a
// content fallback too because the CodeWhisperer surface does not always
// enforce top-level systemPrompt for direct calls.
const timestamp = new Date().toISOString();
const prefixParts = [];
if (thinkingBudget !== null) prefixParts.push(buildThinkingSystemPrefix(thinkingBudget));
prefixParts.push(`[Context: Current time is ${timestamp}]`);
if (agentic) prefixParts.push(KIRO_AGENTIC_SYSTEM_PROMPT);
finalContent = `${prefixParts.join("\n\n")}\n\n${finalContent}`;
const systemPromptParts = [];
if (thinkingBudget !== null) systemPromptParts.push(buildThinkingSystemPrefix(thinkingBudget));
if (agentic) systemPromptParts.push(KIRO_AGENTIC_SYSTEM_PROMPT);
const systemInstruction = extractClaudeSystemText(body.system);
if (systemInstruction) systemPromptParts.push(systemInstruction);
const systemPrompt = systemPromptParts.filter(Boolean).join("\n\n");
const currentTimeContext = `[Context: Current time is ${timestamp}]`;
const contentPrefix = [systemPrompt, currentTimeContext].filter(Boolean).join("\n\n");
const sessionIdentity = resolveSessionIdentity({
headers: credentials?.rawHeaders,
body,
connectionId: credentials?.connectionId,
scope: "kiro",
});
const conversationId = sessionIdentity.sessionId;
const continuationId = resolveContinuationId({
sessionId: conversationId,
connectionId: credentials?.connectionId,
scope: "kiro",
ephemeral: sessionIdentity.ephemeral,
});
const replay = applyKiroSessionReplay({
conversationId,
connectionId: credentials?.connectionId,
modelId: upstreamModel,
systemPrompt,
contentPrefix,
currentContentPrefix: currentTimeContext,
history,
currentMessage,
});
const replayCurrent = replay.currentMessage?.userInputMessage || {};
const userInputMessage = {
content: finalContent,
content: replayCurrent.content || "",
modelId: upstreamModel,
origin: "AI_EDITOR",
...(currentMessage?.userInputMessage?.userInputMessageContext && {
userInputMessageContext:
currentMessage.userInputMessage.userInputMessageContext,
...(replayCurrent.userInputMessageContext && {
userInputMessageContext: replayCurrent.userInputMessageContext,
}),
...(currentMessage?.userInputMessage?.images && {
images: currentMessage.userInputMessage.images,
...(replayCurrent.images && {
images: replayCurrent.images,
}),
};
if (systemInstruction) {
userInputMessage.systemInstruction = systemInstruction;
}
const payload = {
conversationState: {
chatTriggerType: "MANUAL",
conversationId: uuidv4(),
conversationId,
agentContinuationId: continuationId,
agentTaskType: "vibe",
currentMessage: {
userInputMessage,
},
history,
history: replay.history,
},
agentMode: "vibe",
};
if (profileArn) payload.profileArn = profileArn;
if (systemPrompt) payload.systemPrompt = systemPrompt;
const additionalModelRequestFields = buildKiroAdditionalModelRequestFieldsForModel(body, upstreamModel);
if (additionalModelRequestFields) {
payload.additionalModelRequestFields = additionalModelRequestFields;
}
if (maxTokens || temperature !== undefined || topP !== undefined) {
payload.inferenceConfig = {};

View File

@@ -5,13 +5,15 @@
import { register } from "../index.js";
import { FORMATS } from "../formats.js";
import { v4 as uuidv4 } from "uuid";
import { resolveSessionId } from "../../utils/sessionManager.js";
import { applyKiroSessionReplay } from "../../utils/kiroSessionReplay.js";
import { resolveContinuationId, resolveSessionIdentity } from "../../utils/sessionManager.js";
import {
resolveKiroModel,
resolveKiroThinkingBudget,
buildThinkingSystemPrefix,
KIRO_AGENTIC_SYSTEM_PROMPT,
resolveDefaultProfileArn
resolveDefaultProfileArn,
buildKiroAdditionalModelRequestFieldsForModel
} from "../../config/kiroConstants.js";
import { parseDataUri } from "../concerns/image.js";
import { DEFAULT_IMAGE_MIME } from "../schema/index.js";
@@ -546,47 +548,74 @@ export function openaiToKiroRequest(model, body, stream, credentials) {
? (credentials?.providerSpecificData?.profileArn || "")
: (credentials?.providerSpecificData?.profileArn || resolveDefaultProfileArn(authMethod));
let finalContent = currentMessage?.userInputMessage?.content || "";
const timestamp = new Date().toISOString();
// Build the system-prompt prefix that goes ABOVE the user message body.
// Order: thinking_mode tag first (so Kiro sees it before any user text),
// then context/timestamp marker, then optional agentic chunked-write prompt.
const prefixParts = [];
// Kiro CLI/KAS sends these as top-level systemPrompt. Keep a content fallback
// too because the CodeWhisperer surface does not always enforce top-level
// systemPrompt for direct calls.
const systemPromptParts = [];
if (thinkingBudget !== null) {
prefixParts.push(buildThinkingSystemPrefix(thinkingBudget));
systemPromptParts.push(buildThinkingSystemPrefix(thinkingBudget));
}
prefixParts.push(`[Context: Current time is ${timestamp}]`);
if (agentic) {
prefixParts.push(KIRO_AGENTIC_SYSTEM_PROMPT);
systemPromptParts.push(KIRO_AGENTIC_SYSTEM_PROMPT);
}
finalContent = `${prefixParts.join("\n\n")}\n\n${finalContent}`;
const systemPrompt = systemPromptParts.filter(Boolean).join("\n\n");
const currentTimeContext = `[Context: Current time is ${timestamp}]`;
const contentPrefix = [systemPrompt, currentTimeContext].filter(Boolean).join("\n\n");
const sessionIdentity = resolveSessionIdentity({ headers: credentials?.rawHeaders, body, connectionId: credentials?.connectionId, scope: "kiro" });
const conversationId = sessionIdentity.sessionId;
const continuationId = resolveContinuationId({
sessionId: conversationId,
connectionId: credentials?.connectionId,
scope: "kiro",
ephemeral: sessionIdentity.ephemeral,
});
const replay = applyKiroSessionReplay({
conversationId,
connectionId: credentials?.connectionId,
modelId: upstreamModel,
systemPrompt,
contentPrefix,
currentContentPrefix: currentTimeContext,
history,
currentMessage,
});
const replayCurrent = replay.currentMessage?.userInputMessage || {};
const payload = {
conversationState: {
chatTriggerType: "MANUAL",
conversationId: resolveSessionId({ headers: credentials?.rawHeaders, body, connectionId: credentials?.connectionId, scope: "kiro" }),
conversationId,
agentContinuationId: continuationId,
agentTaskType: "vibe",
currentMessage: {
userInputMessage: {
content: finalContent,
content: replayCurrent.content || "",
modelId: upstreamModel,
origin: "AI_EDITOR",
...(currentMessage?.userInputMessage?.images?.length > 0 && {
images: currentMessage.userInputMessage.images
...(replayCurrent.images?.length > 0 && {
images: replayCurrent.images
}),
...(currentMessage?.userInputMessage?.userInputMessageContext && {
userInputMessageContext: currentMessage.userInputMessage.userInputMessageContext
...(replayCurrent.userInputMessageContext && {
userInputMessageContext: replayCurrent.userInputMessageContext
})
}
},
history: history
}
history: replay.history
},
agentMode: "vibe",
};
if (profileArn) {
payload.profileArn = profileArn;
}
if (systemPrompt) payload.systemPrompt = systemPrompt;
const additionalModelRequestFields = buildKiroAdditionalModelRequestFieldsForModel(body, upstreamModel);
if (additionalModelRequestFields) {
payload.additionalModelRequestFields = additionalModelRequestFields;
}
if (maxTokens || temperature !== undefined || topP !== undefined) {
payload.inferenceConfig = {};