271 lines
8.3 KiB
TypeScript
271 lines
8.3 KiB
TypeScript
/**
|
|
* Content-block mapping helpers and streaming state tracker.
|
|
*
|
|
* Translates the Claude Agent SDK's `BetaRawMessageStreamEvent` sequence
|
|
* into GSD's `AssistantMessageEvent` deltas for incremental TUI rendering.
|
|
*/
|
|
|
|
import type {
|
|
AssistantMessage,
|
|
AssistantMessageEvent,
|
|
ServerToolUseContent,
|
|
StopReason,
|
|
TextContent,
|
|
ThinkingContent,
|
|
ToolCall,
|
|
Usage,
|
|
WebSearchResultContent,
|
|
} from "@gsd/pi-ai";
|
|
import { hasXmlParameterTags, repairToolJson } from "@gsd/pi-ai";
|
|
import type { BetaContentBlock, BetaRawMessageStreamEvent, NonNullableUsage } from "./sdk-types.js";
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Content-block mapping helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Convert a single BetaContentBlock to the corresponding GSD content type.
|
|
*/
|
|
export function mapContentBlock(
|
|
block: BetaContentBlock,
|
|
): TextContent | ThinkingContent | ToolCall | ServerToolUseContent | WebSearchResultContent {
|
|
switch (block.type) {
|
|
case "text":
|
|
return { type: "text", text: block.text } satisfies TextContent;
|
|
|
|
case "thinking":
|
|
return {
|
|
type: "thinking",
|
|
thinking: block.thinking,
|
|
...(block.signature ? { thinkingSignature: block.signature } : {}),
|
|
} satisfies ThinkingContent;
|
|
|
|
case "tool_use":
|
|
return {
|
|
type: "toolCall",
|
|
id: block.id,
|
|
name: block.name,
|
|
arguments: block.input,
|
|
} satisfies ToolCall;
|
|
|
|
case "server_tool_use":
|
|
return {
|
|
type: "serverToolUse",
|
|
id: block.id,
|
|
name: block.name,
|
|
input: block.input,
|
|
} satisfies ServerToolUseContent;
|
|
|
|
case "web_search_tool_result":
|
|
return {
|
|
type: "webSearchResult",
|
|
toolUseId: block.tool_use_id,
|
|
content: block.content,
|
|
} satisfies WebSearchResultContent;
|
|
|
|
default: {
|
|
const unknown = block as Record<string, unknown>;
|
|
return { type: "text", text: `[unknown content block: ${JSON.stringify(unknown)}]` };
|
|
}
|
|
}
|
|
}
|
|
|
|
export function mapStopReason(reason: string | null): StopReason {
|
|
switch (reason) {
|
|
case "end_turn":
|
|
case "stop_sequence":
|
|
return "stop";
|
|
case "max_tokens":
|
|
return "length";
|
|
case "tool_use":
|
|
return "toolUse";
|
|
default:
|
|
return "stop";
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Convert SDK usage + total_cost_usd into GSD's Usage shape.
|
|
*
|
|
* The SDK does not break cost down per-bucket, so all cost is
|
|
* attributed to `cost.total`.
|
|
*/
|
|
export function mapUsage(sdkUsage: NonNullableUsage, totalCostUsd: number): Usage {
|
|
return {
|
|
input: sdkUsage.input_tokens,
|
|
output: sdkUsage.output_tokens,
|
|
cacheRead: sdkUsage.cache_read_input_tokens,
|
|
cacheWrite: sdkUsage.cache_creation_input_tokens,
|
|
totalTokens:
|
|
sdkUsage.input_tokens +
|
|
sdkUsage.output_tokens +
|
|
sdkUsage.cache_read_input_tokens +
|
|
sdkUsage.cache_creation_input_tokens,
|
|
cost: {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
total: totalCostUsd,
|
|
},
|
|
};
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Zero-cost usage constant
|
|
// ---------------------------------------------------------------------------
|
|
|
|
export const ZERO_USAGE: Usage = {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
};
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Streaming partial-message state tracker
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Mutable accumulator that tracks the partial AssistantMessage being built
|
|
* from a sequence of stream_event messages. Produces AssistantMessageEvent
|
|
* deltas that the TUI can render incrementally.
|
|
*/
|
|
export class PartialMessageBuilder {
|
|
private partial: AssistantMessage;
|
|
/** Map from stream-event `index` to our content array index. */
|
|
private indexMap = new Map<number, number>();
|
|
/** Accumulated JSON input string per tool_use block (keyed by stream index). */
|
|
private toolJsonAccum = new Map<number, string>();
|
|
|
|
constructor(model: string) {
|
|
this.partial = {
|
|
role: "assistant",
|
|
content: [],
|
|
api: "anthropic-messages",
|
|
provider: "claude-code",
|
|
model,
|
|
usage: { ...ZERO_USAGE },
|
|
stopReason: "stop",
|
|
timestamp: Date.now(),
|
|
};
|
|
}
|
|
|
|
get message(): AssistantMessage {
|
|
return this.partial;
|
|
}
|
|
|
|
/**
|
|
* Feed a BetaRawMessageStreamEvent and return the corresponding
|
|
* AssistantMessageEvent (or null if the event is not mapped).
|
|
*/
|
|
handleEvent(event: BetaRawMessageStreamEvent): AssistantMessageEvent | null {
|
|
const streamIndex = event.index ?? 0;
|
|
|
|
switch (event.type) {
|
|
// ---- Block start ----
|
|
case "content_block_start": {
|
|
const block = event.content_block;
|
|
if (!block) return null;
|
|
|
|
const contentIndex = this.partial.content.length;
|
|
this.indexMap.set(streamIndex, contentIndex);
|
|
|
|
if (block.type === "text") {
|
|
this.partial.content.push({ type: "text", text: "" });
|
|
return { type: "text_start", contentIndex, partial: this.partial };
|
|
}
|
|
if (block.type === "thinking") {
|
|
this.partial.content.push({ type: "thinking", thinking: "" });
|
|
return { type: "thinking_start", contentIndex, partial: this.partial };
|
|
}
|
|
if (block.type === "tool_use") {
|
|
this.toolJsonAccum.set(streamIndex, "");
|
|
this.partial.content.push({
|
|
type: "toolCall",
|
|
id: block.id,
|
|
name: block.name,
|
|
arguments: {},
|
|
});
|
|
return { type: "toolcall_start", contentIndex, partial: this.partial };
|
|
}
|
|
if (block.type === "server_tool_use") {
|
|
this.partial.content.push({
|
|
type: "serverToolUse",
|
|
id: block.id,
|
|
name: block.name,
|
|
input: block.input,
|
|
});
|
|
return { type: "server_tool_use", contentIndex, partial: this.partial };
|
|
}
|
|
return null;
|
|
}
|
|
|
|
// ---- Block delta ----
|
|
case "content_block_delta": {
|
|
const contentIndex = this.indexMap.get(streamIndex);
|
|
if (contentIndex === undefined) return null;
|
|
const delta = event.delta;
|
|
if (!delta) return null;
|
|
|
|
if (delta.type === "text_delta" && typeof delta.text === "string") {
|
|
const existing = this.partial.content[contentIndex] as TextContent;
|
|
existing.text += delta.text;
|
|
return { type: "text_delta", contentIndex, delta: delta.text, partial: this.partial };
|
|
}
|
|
if (delta.type === "thinking_delta" && typeof delta.thinking === "string") {
|
|
const existing = this.partial.content[contentIndex] as ThinkingContent;
|
|
existing.thinking += delta.thinking;
|
|
return { type: "thinking_delta", contentIndex, delta: delta.thinking, partial: this.partial };
|
|
}
|
|
if (delta.type === "input_json_delta" && typeof delta.partial_json === "string") {
|
|
const accum = (this.toolJsonAccum.get(streamIndex) ?? "") + delta.partial_json;
|
|
this.toolJsonAccum.set(streamIndex, accum);
|
|
return { type: "toolcall_delta", contentIndex, delta: delta.partial_json, partial: this.partial };
|
|
}
|
|
return null;
|
|
}
|
|
|
|
// ---- Block stop ----
|
|
case "content_block_stop": {
|
|
const contentIndex = this.indexMap.get(streamIndex);
|
|
if (contentIndex === undefined) return null;
|
|
const block = this.partial.content[contentIndex];
|
|
|
|
if (block.type === "text") {
|
|
return { type: "text_end", contentIndex, content: block.text, partial: this.partial };
|
|
}
|
|
if (block.type === "thinking") {
|
|
return { type: "thinking_end", contentIndex, content: block.thinking, partial: this.partial };
|
|
}
|
|
if (block.type === "toolCall") {
|
|
const jsonStr = this.toolJsonAccum.get(streamIndex) ?? "{}";
|
|
const jsonForParse = hasXmlParameterTags(jsonStr) ? repairToolJson(jsonStr) : jsonStr;
|
|
try {
|
|
block.arguments = JSON.parse(jsonForParse);
|
|
} catch {
|
|
// JSON.parse failed — attempt repair for YAML-style bullet
|
|
// lists that LLMs copy from template formatting (#2660).
|
|
try {
|
|
block.arguments = JSON.parse(repairToolJson(jsonForParse));
|
|
} catch {
|
|
// Repair also failed — stream was truncated or garbage.
|
|
// Preserve the raw string for diagnostics but signal the
|
|
// malformation explicitly so downstream consumers can
|
|
// distinguish this from a healthy tool completion (#2574).
|
|
block.arguments = { _raw: jsonStr };
|
|
return { type: "toolcall_end", contentIndex, toolCall: block, partial: this.partial, malformedArguments: true };
|
|
}
|
|
}
|
|
return { type: "toolcall_end", contentIndex, toolCall: block, partial: this.partial };
|
|
}
|
|
return null;
|
|
}
|
|
|
|
default:
|
|
return null;
|
|
}
|
|
}
|
|
}
|