Files
2026-07-13 13:39:12 +08:00

117 lines
4.4 KiB
TypeScript

/**
* Inline `<thinking>` splitter for Claude on Kiro.
*
* Background:
* When `<thinking_mode>enabled</thinking_mode>` is in the system prompt,
* Claude on Kiro emits its reasoning **inline** as `<thinking>…</thinking>`
* blocks inside `assistantResponseEvent.content`, rather than as separate
* `reasoningContentEvent` frames. To match the OpenAI streaming shape that
* downstream translators (Anthropic /thinking_blocks, Claude SSE, etc.)
* expect, we split that inline reasoning back out and route it to the
* `delta.reasoning_content` channel instead of `delta.content`.
*
* The implementation is split into pure functions so it can be unit-tested
* without dragging in the rest of the executor stack (proxy-agent, AWS
* EventStream parser, etc.). The KiroExecutor wires these helpers into its
* TransformStream by passing controller-bound emit callbacks.
*
* Ported from decolua/9router#1273 (kiroThinking.js) by Amin Fathullah.
*/
/** Mutable state carried across `splitInlineThinking` calls. */
export type KiroThinkingState = {
/** True while the cursor is inside a `<thinking>` block. */
thinkingMode: boolean;
/**
* Characters held back because they might be the start of a tag we'll
* complete on the next slice (e.g. `<thi`).
*/
pendingTag: string;
};
/**
* Stream-safe splitter. Walks one slice of upstream content at a time and
* routes characters to either the content channel or the reasoning channel
* based on the current `<thinking>` state.
*
* State is mutated on `state` so a tag split between frames (e.g. `…</think`
* followed by `ing>foo`) is still recognised.
*
* @param state Mutable state carried across calls. Initialise with
* `{ thinkingMode: false, pendingTag: "" }`.
* @param raw Next slice from `assistantResponseEvent.content`. May be empty
* or null/undefined (no-op).
* @param onContent Called with text that should land in `delta.content`.
* @param onReasoning Called with text that should land in `delta.reasoning_content`.
*/
export function splitInlineThinking(
state: KiroThinkingState,
raw: string | null | undefined,
onContent: (s: string) => void,
onReasoning: (s: string) => void
): void {
let text = (state.pendingTag || "") + (raw || "");
state.pendingTag = "";
// Maximum length of an unfinished tag we might still complete on the next
// frame: `</thinking>` is the longest at 11 chars.
const PARTIAL_MAX = 11;
while (text.length > 0) {
const target = state.thinkingMode ? "</thinking>" : "<thinking>";
const idx = text.indexOf(target);
if (idx === -1) {
// No full target tag in `text`. Look for a possible partial at the end
// so we can complete it on the next frame.
let holdFrom = text.length;
for (let i = Math.max(0, text.length - PARTIAL_MAX); i < text.length; i++) {
const tail = text.slice(i);
if (target.startsWith(tail) && tail.length > 0) {
holdFrom = i;
break;
}
}
const flushable = text.slice(0, holdFrom);
if (flushable) {
if (state.thinkingMode) onReasoning(flushable);
else onContent(flushable);
}
state.pendingTag = text.slice(holdFrom);
return;
}
// Found a complete target tag. Flush everything before it in the current
// mode, flip the mode, and keep walking the remainder.
const before = text.slice(0, idx);
if (before) {
if (state.thinkingMode) onReasoning(before);
else onContent(before);
}
state.thinkingMode = !state.thinkingMode;
text = text.slice(idx + target.length);
}
}
/**
* Drain whatever is left in `state.pendingTag` at end-of-stream. Routes the
* leftover characters to whichever channel matches the current
* `state.thinkingMode` so we don't silently lose data when the stream ends
* mid-tag (e.g. `<thi`).
*
* @param state Mutable state shared with `splitInlineThinking`.
* @param onContent Called with text that should land in `delta.content`.
* @param onReasoning Called with text that should land in `delta.reasoning_content`.
*/
export function flushPendingThinking(
state: KiroThinkingState,
onContent: (s: string) => void,
onReasoning: (s: string) => void
): void {
if (!state.pendingTag) return;
const leftover = state.pendingTag;
state.pendingTag = "";
if (state.thinkingMode) onReasoning(leftover);
else onContent(leftover);
}