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:
@@ -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 = {};
|
||||
|
||||
@@ -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 = {};
|
||||
|
||||
Reference in New Issue
Block a user