pi-ext

import { createHash } from "node:crypto";
import type {
	AssistantMessage,
	Context,
	ImageContent,
	Model,
	TextContent,
	ThinkingLevel,
	Tool,
	Usage,
} from "@earendil-works/pi-ai";
import type { KlausSessionStore } from "./cache.js";
import { dbg } from "./debug.js";

export const KLAUS_API = "klaus-agent-sdk";
export const SDK_VERSION = "0.3.281";
export const CLAUDE_VERSION = "2.1.281";

export const MODEL_SELECTORS = {
	"claude-fable-5-1": "fable",
	"claude-haiku-4-5": "haiku",
	"claude-opus-5": "opus",
	"claude-sonnet-5": "sonnet",
} as const;

/** Any Claude model id the Pi catalog can carry: non-dated, non-alias. */
export type KlausModelId = string;

/** Wire selector for the Claude Agent SDK: the short alias when klaus knows
 * one, otherwise the full catalog id (the SDK accepts both). */
export function klausSelector(modelId: KlausModelId): string {
	return Object.hasOwn(MODEL_SELECTORS, modelId)
		? MODEL_SELECTORS[modelId as keyof typeof MODEL_SELECTORS]
		: modelId;
}

export type KlausContext = Omit<Context, "messages"> & {
	messages: Exclude<Context["messages"][number], { role: "system" }>[];
};

export interface KlausTool {
	name: string;
	description: string;
	inputSchema: Record<string, unknown>;
}

export interface KlausRequest {
	modelId: KlausModelId;
	selector: string;
	systemPrompt: string;
	prompt: string;
	images?: Array<{ data: string; mimeType: string }>;
	tools: KlausTool[];
	thinking: ThinkingLevel | undefined;
	thinkingBudget?: number;
	maxTokens?: number;
	cwd: string;
	sessionId?: string;
	headers: Record<string, string | null>;
	env?: Record<string, string>;
	timeoutMs?: number;
	sessionStore?: KlausSessionStore;
	resume?: string;
	resumeAt?: string;
	forkSession?: boolean;
}

export type KlausContentEvent =
	| { type: "text-start"; index: number }
	| { type: "text-delta"; index: number; delta: string }
	| { type: "text-end"; index: number; text: string }
	| { type: "thinking-start"; index: number }
	| { type: "thinking-delta"; index: number; delta: string }
	| {
			type: "thinking-end";
			index: number;
			thinking: string;
			signature?: string;
			redacted?: boolean;
	  }
	| { type: "tool-start"; index: number; id: string; name: string }
	| { type: "tool-delta"; index: number; delta: string }
	| {
			type: "tool-end";
			index: number;
			id: string;
			name: string;
			arguments: import("@earendil-works/pi-ai").JsonObject;
	  };

export interface KlausUsage {
	input: number;
	output: number;
	cacheRead: number;
	cacheWrite: number;
}

export interface KlausQueryCallbacks {
	onReady(metadata: Record<string, string>): Promise<void>;
	onActivity(): void;
	/** Child-side conditions the human should see even though the turn continues,
	 * such as API retries and permission denials Claude decided on its own. */
	onNotice(message: string): void;
	onContent(event: KlausContentEvent): void;
	onToolBoundary(result: {
		usage: KlausUsage;
		responseId?: string;
		responseModel?: string;
	}): void;
	onResult(result: {
		usage: KlausUsage;
		responseId?: string;
		responseModel?: string;
		position?: string;
		sdkSessionId?: string;
		stopReason: "stop" | "length";
	}): void;
	onError(error: Error): void;
}

export function isKlausModelId(value: string): value is KlausModelId {
	const valid =
		/^claude-[a-z0-9]+(?:[.-][a-z0-9]+)*$/.test(value) &&
		!/-\d{8}$/.test(value) &&
		!value.endsWith("-latest");
	dbg?.("protocol.isKlausModelId");
	return valid;
}

export function toKlausTools(tools: Tool[] | undefined): KlausTool[] {
	dbg?.("protocol.toKlausTools");
	return (tools ?? []).map((tool) => ({
		name: tool.name,
		description: tool.description,
		inputSchema: tool.parameters as Record<string, unknown>,
	}));
}

export function requiredConstrainedTool(
	tools: Tool[] | undefined,
): string | undefined {
	dbg?.("protocol.requiredConstrainedTool", { toolCount: tools?.length ?? 0 });
	return tools?.find(
		(tool) =>
			tool.constrainedSampling &&
			tool.constrainedSampling.type === "json_schema" &&
			tool.constrainedSampling.strict === "require",
	)?.name;
}

export function emptyUsage(): Usage {
	dbg?.("protocol.emptyUsage");
	return {
		input: 0,
		output: 0,
		cacheRead: 0,
		cacheWrite: 0,
		totalTokens: 0,
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
	};
}

export function createAssistant(model: Model<string>): AssistantMessage {
	dbg?.("protocol.createAssistant");
	return {
		role: "assistant",
		content: [],
		api: model.api,
		provider: model.provider,
		model: model.id,
		usage: emptyUsage(),
		stopReason: "pending",
		timestamp: Date.now(),
	};
}

export function modelFacingContent(
	content: (TextContent | ImageContent)[],
): Array<
	| { type: "text"; text: string }
	| { type: "image"; data: string; mimeType: string }
> {
	dbg?.("protocol.modelFacingContent");
	return content.map((item) =>
		item.type === "text"
			? { type: "text", text: item.text }
			: { type: "image", data: item.data, mimeType: item.mimeType },
	);
}

export function contextFingerprint(context: Context, identity = ""): string {
	dbg?.("protocol.contextFingerprint.start", {
		messageCount: context.messages.length,
		toolCount: context.tools?.length ?? 0,
	});
	const messages = context.messages.flatMap<unknown>((message) => {
		if (message.role === "system") return [];
		if (message.role === "user") {
			return { role: "user", content: message.content };
		}
		if (message.role === "assistant") {
			return {
				role: "assistant",
				content: message.content.flatMap<unknown>((item) => {
					if (item.type === "thinking") return [];
					if (item.type === "text") return [{ type: "text", text: item.text }];
					return [
						{
							type: "toolCall",
							id: item.id,
							name: item.name,
							arguments: item.arguments,
						},
					];
				}),
			};
		}
		return [
			{
				role: "toolResult",
				toolCallId: message.toolCallId,
				toolName: message.toolName,
				content: message.content,
				isError: message.isError,
			},
		];
	});
	const fingerprint = createHash("sha256")
		.update(
			JSON.stringify({
				protocol: 1,
				identity,
				systemPrompt: context.systemPrompt ?? "",
				messages,
				tools: toKlausTools(context.tools),
			}),
		)
		.digest("hex");
	dbg?.("protocol.contextFingerprint.end");
	return fingerprint;
}