@gajae-code/ai 0.18.5 → 0.18.7

package/src/providers/kiro-codewhisperer.ts

@@ -26,11 +26,11 @@ import type {
 	ToolResultMessage,
 } from "../types";
 import { AssistantMessageEventStream } from "../utils/event-stream";
-import { transportFailureFacts } from "../utils/fallback-transport";
+import { PROVIDER_PROTOCOL_MISMATCH_ERROR_CODE, transportFailureFacts } from "../utils/fallback-transport";
 import { withHttpStatus } from "../utils/http-inspector";
 import { captureUnicodeEscapeEvidence } from "../utils/json-parse";
 import { decodeEventStream } from "./aws-eventstream";
-import { isKiroApiKey, streamKiroApiKey, toKiroModelId } from "./kiro-api-key";
+import { isKiroApiKey, sanitizeKiroError, streamKiroApiKey, toKiroModelId } from "./kiro-api-key";
 
 // ─────────────────────────────────────────────────────────────────────────────
 // Provider options
@@ -74,11 +74,11 @@ interface WireUserMessage {
 		content: string;
 		modelId?: string;
 		userInputMessageContext?: {
-			tools?: { tools: WireToolSpec[] };
+			tools?: WireToolSpec[];
 			toolResults?: WireToolResult[];
 			editorStateContext?: Record<string, unknown>;
 		};
-		origin?: string;
+		origin: typeof KIRO_ORIGIN;
 	};
 }
 
@@ -108,25 +108,19 @@ interface GenerateAssistantResponseRequest {
 // ─────────────────────────────────────────────────────────────────────────────
 
 interface AssistantResponseEvent {
-	assistantResponseEvent?: {
-		content?: string;
-	};
+	content?: string;
 }
 
 interface ToolUseEventPayload {
-	toolUseEvent?: {
-		toolUseId?: string;
-		name?: string;
-		input?: unknown;
-		stop?: { stopReason?: string };
-	};
+	toolUseId?: string;
+	name?: string;
+	input?: string;
+	stop?: boolean;
 }
 
 interface MessageMetadataEvent {
-	messageMetadataEvent?: {
-		conversationId?: string;
-		utteranceId?: string;
-	};
+	conversationId?: string;
+	utteranceId?: string;
 }
 
 interface ErrorPayload {
@@ -142,6 +136,12 @@ interface ErrorPayload {
 
 const DEFAULT_REGION = "us-east-1";
 
+/**
+ * Without `origin` the service ignores `modelId` and routes every request to
+ * `auto`; with it, response frames report the requested model.
+ */
+const KIRO_ORIGIN = "AI_EDITOR";
+
 type Block = (TextContent | ToolCall) & { index?: number; partialJson?: string };
 
 // ─────────────────────────────────────────────────────────────────────────────
@@ -163,6 +163,8 @@ export const streamKiroCodeWhisperer: StreamFunction<"kiro-codewhisperer-stream"
 	(async () => {
 		const startTime = Date.now();
 		let firstTokenTime: number | undefined;
+		// Accumulator for streaming tool input fragments, keyed by toolUseId
+		const toolInputAccumulator = new Map<string, { name: string; input: string }>();
 
 		const output: AssistantMessage = {
 			role: "assistant",
@@ -184,6 +186,7 @@ export const streamKiroCodeWhisperer: StreamFunction<"kiro-codewhisperer-stream"
 
 		const blocks = output.content as Block[];
 		const region = options.region ?? $env.KIRO_REGION ?? $env.AWS_REGION ?? $env.AWS_DEFAULT_REGION ?? DEFAULT_REGION;
+		let started = false; // Track whether start event has been emitted
 
 		try {
 			assertAwsRegionLabel(region);
@@ -216,38 +219,53 @@ export const streamKiroCodeWhisperer: StreamFunction<"kiro-codewhisperer-stream"
 			const bodyText = JSON.stringify(requestBody);
 			const body = new TextEncoder().encode(bodyText);
 			const requestHeaders: Record<string, string> = {
-				"content-type": "application/json",
-				accept: "application/vnd.amazon.eventstream",
-				authorization: `Bearer ${bearerToken}`,
-				"amzn-X-amz-target": "AmazonCodeWhispererService.GenerateAssistantResponse",
+				"Content-Type": "application/x-amz-json-1.0",
+				Accept: "application/vnd.amazon.eventstream",
+				Authorization: `Bearer ${bearerToken}`,
+				"X-Amz-Target": "AmazonCodeWhispererStreamingService.GenerateAssistantResponse",
 			};
 
-			if (options.profileArn) {
-				requestHeaders["x-amzn-codewhisperer-proflearn"] = options.profileArn;
-			}
-
-			// Merge user-provided headers
+			// Merge user-provided headers (case-insensitively)
+			const headersList = new Headers(requestHeaders);
 			if (options.headers) {
-				Object.assign(requestHeaders, options.headers);
+				const userHeaders = new Headers(options.headers);
+				userHeaders.forEach((value, key) => {
+					headersList.set(key, value);
+				});
 			}
 
 			const response = await fetch(url, {
 				method: "POST",
-				headers: requestHeaders,
+				headers: headersList,
 				body,
 				redirect: "error",
 				signal: options.signal,
 			});
 
 			if (!response.ok) {
-				const errBody = await response.text().catch(() => "");
+				const errBody = await readBodyPrefix(response);
 				throw withHttpStatus(
-					new Error(`Kiro CodeWhisperer HTTP ${response.status}: ${errBody.slice(0, 1000)}`),
+					new Error(sanitizeKiroError(`Kiro CodeWhisperer HTTP ${response.status}: ${errBody}`, bearerToken)),
 					response.status,
 				);
 			}
 
-			if (!response.body) throw new Error("Kiro CodeWhisperer response has no body");
+			// Verify content-type is eventstream; if not, read and report the actual error
+			const contentType = response.headers.get("content-type") ?? "";
+			const mediaType = contentType.split(";", 1)[0]?.trim().toLowerCase() ?? "";
+			if (!response.body || mediaType !== "application/vnd.amazon.eventstream") {
+				const errBody = await readBodyPrefix(response);
+				// Deliberately no HTTP status metadata and no "HTTP <code>" wording: a 2xx
+				// protocol mismatch is deterministic, and a status would materialize
+				// transport-failure facts that session fallback retries.
+				// Sanitize each untrusted part on its own so the fixed prefix does not eat
+				// into the body's 1000-character diagnostic budget.
+				const safeContentType = contentType ? sanitizeKiroError(contentType, bearerToken) : "";
+				const safeBody = errBody ? sanitizeKiroError(errBody, bearerToken) : "";
+				throw new KiroNonEventStreamError(
+					`Kiro CodeWhisperer returned a non-eventstream ${response.status} response (${safeContentType}): ${safeBody}`,
+				);
+			}
 
 			// Decode eventstream
 			for await (const message of decodeEventStream(response.body)) {
@@ -279,11 +297,12 @@ export const streamKiroCodeWhisperer: StreamFunction<"kiro-codewhisperer-stream"
 				switch (eventType) {
 					case "assistantResponseEvent": {
 						const ev = payload as AssistantResponseEvent;
-						const content = ev.assistantResponseEvent?.content;
+						const content = ev.content;
 						if (content) {
 							if (!firstTokenTime) firstTokenTime = Date.now();
-							if (blocks.length === 0) {
+							if (!started) {
 								stream.push({ type: "start", partial: output });
+								started = true;
 							}
 							handleTextDelta(content, blocks, output, stream);
 						}
@@ -292,16 +311,21 @@ export const streamKiroCodeWhisperer: StreamFunction<"kiro-codewhisperer-stream"
 					case "toolUseEvent": {
 						const ev = payload as ToolUseEventPayload;
 						if (!firstTokenTime) firstTokenTime = Date.now();
-						if (blocks.length === 0) {
+						if (!started) {
 							stream.push({ type: "start", partial: output });
+							started = true;
+						}
+						handleToolUseEvent(ev, blocks, output, stream, toolInputAccumulator);
+						// Clear accumulator for completed tool
+						if (ev.stop) {
+							toolInputAccumulator.delete(ev.toolUseId ?? "");
 						}
-						handleToolUseEvent(ev, blocks, output, stream);
 						break;
 					}
 					case "messageMetadataEvent": {
 						const ev = payload as MessageMetadataEvent;
-						if (ev.messageMetadataEvent?.conversationId) {
-							output.responseId = ev.messageMetadataEvent.conversationId;
+						if (ev.conversationId) {
+							output.responseId = ev.conversationId;
 						}
 						break;
 					}
@@ -328,6 +352,12 @@ export const streamKiroCodeWhisperer: StreamFunction<"kiro-codewhisperer-stream"
 
 			if (options.signal?.aborted) throw new Error("Request was aborted");
 
+			// Reject EOF while tool input is still being accumulated
+			if (toolInputAccumulator.size > 0) {
+				const unfinishedIds = Array.from(toolInputAccumulator.keys()).join(", ");
+				throw new Error(`Kiro CodeWhisperer stream ended with incomplete tool calls: ${unfinishedIds}`);
+			}
+
 			// Finalize blocks
 			for (const block of blocks) {
 				delete block.index;
@@ -348,8 +378,10 @@ export const streamKiroCodeWhisperer: StreamFunction<"kiro-codewhisperer-stream"
 				delete (block as Block).partialJson;
 			}
 			output.stopReason = options.signal?.aborted ? "aborted" : "error";
-			output.errorStatus = extractHttpStatusFromError(error);
+			// The non-eventstream diagnostic embeds untrusted body text; never parse a status out of it.
+			output.errorStatus = error instanceof KiroNonEventStreamError ? undefined : extractHttpStatusFromError(error);
 			output.transportFailure = transportFailureFacts(error);
+			if (error instanceof KiroNonEventStreamError) output.errorCode = PROVIDER_PROTOCOL_MISMATCH_ERROR_CODE;
 			const baseMessage = error instanceof Error ? error.message : JSON.stringify(error);
 			output.errorMessage = baseMessage;
 			output.duration = Date.now() - startTime;
@@ -386,7 +418,20 @@ function buildConversationState(
 	const history: WireHistoryMessage[] = [];
 	const systemPrompt = context.systemPrompt?.join("\n") ?? "";
 
-	for (let i = 0; i < messages.length - 1; i++) {
+	// Parallel tool calls end the conversation with several consecutive toolResult
+	// messages. They all answer the last assistant turn, so they must travel
+	// together in currentMessage; leaving the earlier ones as a separate history
+	// entry is rejected upstream with TOOL_USE_RESULT_MISMATCH.
+	let currentStart = messages.length - 1;
+	while (
+		currentStart > 0 &&
+		messages[currentStart].role === "toolResult" &&
+		messages[currentStart - 1].role === "toolResult"
+	) {
+		currentStart--;
+	}
+
+	for (let i = 0; i < currentStart; i++) {
 		const msg = messages[i];
 		const previous = history[history.length - 1];
 		// Parallel tool calls produce consecutive toolResult messages; they answer one
@@ -406,18 +451,21 @@ function buildConversationState(
 		history.push(convertToWireMessage(msg, modelId, i === 0 ? systemPrompt : undefined));
 	}
 
-	// Convert the last message as currentMessage
+	// Convert the last message as currentMessage, carrying every trailing parallel result.
 	const lastMsg = messages[messages.length - 1];
 	const currentMessage = convertToWireUserMessage(lastMsg, modelId, systemPrompt);
+	if (currentStart < messages.length - 1) {
+		currentMessage.userInputMessage.userInputMessageContext = {
+			toolResults: messages.slice(currentStart).map(msg => convertToolResult(msg as ToolResultMessage)),
+		};
+	}
 
 	// Add tools to the current message context
 	if (context.tools && context.tools.length > 0) {
 		if (!currentMessage.userInputMessage.userInputMessageContext) {
 			currentMessage.userInputMessage.userInputMessageContext = {};
 		}
-		currentMessage.userInputMessage.userInputMessageContext.tools = {
-			tools: convertTools(context.tools),
-		};
+		currentMessage.userInputMessage.userInputMessageContext.tools = convertTools(context.tools);
 	}
 
 	return {
@@ -480,6 +528,7 @@ function convertToWireUserMessage(
 		userInputMessage: {
 			content,
 			modelId,
+			origin: KIRO_ORIGIN,
 		},
 	};
 
@@ -559,50 +608,97 @@ function handleToolUseEvent(
 	blocks: Block[],
 	output: AssistantMessage,
 	stream: AssistantMessageEventStream,
+	accumulator: Map<string, { name: string; input: string }>,
 ): void {
-	const toolEvent = ev.toolUseEvent;
-	if (!toolEvent) return;
+	const name = ev.name ?? "";
+	const input = ev.input ?? "";
+
+	// Determine the tool ID: use provided ID, or find the currently active tool call
+	let toolUseId = ev.toolUseId;
+	if (!toolUseId) {
+		// ID-less fragment: associate with the currently active tool call
+		// Find the last tool call being accumulated (the most recently started tool)
+		const keys = Array.from(accumulator.keys());
+		if (keys.length > 0) {
+			toolUseId = keys[keys.length - 1];
+		} else {
+			// No active tool calls; create a placeholder entry (should rarely happen)
+			toolUseId = "";
+		}
+	}
 
-	const toolUseId = toolEvent.toolUseId ?? "";
-	const name = toolEvent.name ?? "";
+	// Accumulate input fragments per toolUseId
+	let accumulated = accumulator.get(toolUseId);
+	if (!accumulated) {
+		accumulated = { name, input: "" };
+		accumulator.set(toolUseId, accumulated);
+	}
+	accumulated.input += input;
+
+	// Update name if provided and not yet set
+	if (name && !accumulated.name) {
+		accumulated.name = name;
+	}
 
-	// If input is provided as a complete object, emit toolcall_end
-	if (toolEvent.input !== undefined && toolEvent.input !== null) {
-		const inputStr = typeof toolEvent.input === "string" ? toolEvent.input : JSON.stringify(toolEvent.input);
+	// Emit toolcall_end only when we have a stop signal
+	if (ev.stop) {
+		const inputStr = accumulated.input;
 		const toolCall: ToolCall = {
 			type: "toolCall",
 			id: toolUseId,
-			name,
+			name: accumulated.name,
 			arguments: safeParseJson(inputStr) as Record<string, any>,
 		};
-		if (typeof toolEvent.input === "string") captureUnicodeEscapeEvidence(toolCall, inputStr);
+		captureUnicodeEscapeEvidence(toolCall, inputStr);
 
 		const newBlock: Block = { ...toolCall, index: blocks.length };
-		if (typeof toolEvent.input === "string") captureUnicodeEscapeEvidence(newBlock, inputStr);
+		captureUnicodeEscapeEvidence(newBlock, inputStr);
 		blocks.push(newBlock);
 		stream.push({ type: "toolcall_end", contentIndex: newBlock.index!, toolCall, partial: output });
-		return;
+		accumulator.delete(toolUseId);
 	}
-
-	// Otherwise, accumulate partial input (if the service streams it in chunks)
-	// The published model does not document chunked tool input for CodeWhisperer,
-	// so this path handles the case defensively but expects complete input per event.
-	const toolCall: ToolCall = {
-		type: "toolCall",
-		id: toolUseId,
-		name,
-		arguments: {},
-	};
-
-	const newBlock: Block = { ...toolCall, index: blocks.length };
-	blocks.push(newBlock);
-	stream.push({ type: "toolcall_end", contentIndex: newBlock.index!, toolCall, partial: output });
 }
 
 // ─────────────────────────────────────────────────────────────────────────────
 // Helpers
 // ─────────────────────────────────────────────────────────────────────────────
 
+/** A 2xx response that is not an AWS event stream: a deterministic protocol error, never a transport failure. */
+class KiroNonEventStreamError extends Error {
+	override name = "KiroNonEventStreamError";
+}
+
+/** Byte cap for diagnostic error bodies; the rest of the stream is cancelled unread. */
+const ERROR_BODY_PREFIX_BYTES = 4096;
+
+/** Read at most {@link ERROR_BODY_PREFIX_BYTES} of a response body for diagnostics, then cancel the remainder. */
+async function readBodyPrefix(response: Response): Promise<string> {
+	if (!response.body) return "";
+	const reader = response.body.getReader();
+	const chunks: Uint8Array[] = [];
+	let total = 0;
+	try {
+		while (total < ERROR_BODY_PREFIX_BYTES) {
+			const { done, value } = await reader.read();
+			if (done) break;
+			const take = value.subarray(0, ERROR_BODY_PREFIX_BYTES - total);
+			chunks.push(take);
+			total += take.length;
+		}
+	} catch {
+		// Diagnostic read only; report whatever prefix arrived.
+	} finally {
+		reader.cancel().catch(() => {});
+	}
+	const bytes = new Uint8Array(total);
+	let offset = 0;
+	for (const chunk of chunks) {
+		bytes.set(chunk, offset);
+		offset += chunk.length;
+	}
+	return new TextDecoder().decode(bytes);
+}
+
 function resolveBearerToken(apiKey: string | undefined): string | undefined {
 	if (!apiKey) {
 		return $credentialEnv("KIRO_API_KEY") ?? $credentialEnv("AWS_BEARER_TOKEN_KIRO") ?? undefined;

package/src/providers/openai-codex-responses.d.ts

@@ -11,6 +11,8 @@ export interface OpenAICodexResponsesOptions extends StreamOptions {
     preferWebsockets?: boolean;
     serviceTier?: ServiceTier;
 }
+/** @internal Exported for tests. */
+export declare function createCodexStreamProgressClassifier(): (event: unknown) => boolean;
 type CodexTransport = "sse" | "websocket";
 export interface OpenAICodexWebSocketDebugStats {
     fullContextRequests: number;
@@ -75,3 +77,4 @@ type CodexToolPayload = {
 /** @internal Exported for tests. */
 export declare function convertOpenAICodexResponsesTools(tools: Tool[], model: Model<"openai-codex-responses">): CodexToolPayload[];
 declare function isRetryableCodexFailureEvent(rawEvent: Record<string, unknown>): boolean;
+export declare function isExplicitCodexTerminalVeto(code?: string, providerMessage?: string): boolean;