pi-plus 0.1.8 → 0.1.11

package/src/usage.ts

@@ -1,201 +0,0 @@
-/**
- * Ports of Claude Code's usage-based context measurement and rough tail
- * estimation.
- *
- * Sources (claude-code repo):
- *  - src/utils/tokens.ts                — tokenCountWithEstimation, getTokenCountFromUsage
- *  - src/services/tokenEstimation.ts    — roughTokenCountEstimationForBlock
- *  - src/services/compact/microCompact.ts — IMAGE_MAX_TOKEN_SIZE = 2000
- *
- * pi's message model differs from CC's: tool results are separate
- * `toolResult` messages (CC: `user` messages containing tool_result blocks)
- * and assistant tool calls use the `toolCall` block type (CC: `tool_use`).
- * The estimation below maps those 1:1 so the arithmetic matches CC.
- */
-
-/** Structural shape of pi's Usage we rely on (pi-ai Usage). */
-export interface UsageSnapshot {
-	input: number;
-	output: number;
-	cacheRead: number;
-	cacheWrite: number;
-}
-
-/** Structural minimal message shape — accepts pi AgentMessage without a
- * hard runtime dependency on pi-ai. */
-export interface RoughMessage {
-	role?: string;
-	content?: unknown;
-	usage?: Partial<UsageSnapshot> | null;
-	stopReason?: string;
-	timestamp?: number;
-	toolsAdded?: unknown[];
-	toolsRemoved?: unknown[];
-}
-
-/** CC: images and PDFs are billed ~2000 tokens after client-side resizing;
- * a 1MB base64 PDF must NOT fall through to the stringify catch-all. */
-const IMAGE_DOCUMENT_TOKEN_ESTIMATE = 2000;
-const CHARS_PER_TOKEN = 4;
-
-/**
- * CC: getTokenCountFromUsage — the full context size at the time of that API
- * call: input + cache creation (5m and 1h) + cache read + output.
- * pi's `cacheWrite` already includes the 1h subset, and `reasoning` is a
- * subset of `output`, so summing the four fields is exactly CC's formula.
- */
-export function getTokenCountFromUsage(usage: UsageSnapshot): number {
-	return usage.input + usage.cacheWrite + usage.cacheRead + usage.output;
-}
-
-function hasPositiveTokenCount(usage: Partial<UsageSnapshot> | null | undefined): usage is UsageSnapshot {
-	if (!usage) return false;
-	const total = getTokenCountFromUsage({
-		input: usage.input ?? 0,
-		output: usage.output ?? 0,
-		cacheRead: usage.cacheRead ?? 0,
-		cacheWrite: usage.cacheWrite ?? 0,
-	});
-	return total > 0;
-}
-
-function safeJsonStringify(value: unknown): string {
-	try {
-		return JSON.stringify(value) ?? "undefined";
-	} catch {
-		return "[unserializable]";
-	}
-}
-
-/** CC: roughTokenCountEstimation — characters / 4, rounded. */
-export function roughTokenCountEstimation(content: string, bytesPerToken: number = CHARS_PER_TOKEN): number {
-	return Math.round(content.length / bytesPerToken);
-}
-
-/**
- * Find the last assistant message whose usage validly describes the current
- * context prefix. Port of pi's getLastAssistantUsageInfo (estimate.ts), which
- * matches CC's getTokenUsage behavior: skip aborted/error responses, skip
- * responses older than a newer prefix message (e.g. a compaction summary
- * inserted afterwards), skip zero-usage records.
- */
-export function findLastUsageInfo(
-	messages: readonly RoughMessage[],
-): { usage: UsageSnapshot; index: number } | undefined {
-	let latestPrefixTimestamp = Number.NEGATIVE_INFINITY;
-	let usageInfo: { usage: UsageSnapshot; index: number } | undefined;
-
-	for (let i = 0; i < messages.length; i++) {
-		const message = messages[i];
-		if (message.role === "assistant" && hasPositiveTokenCount(message.usage)) {
-			const usageAppliesToPrefix = (message.timestamp ?? 0) >= latestPrefixTimestamp;
-			if (usageAppliesToPrefix && message.stopReason !== "aborted" && message.stopReason !== "error") {
-				usageInfo = {
-					usage: {
-						input: message.usage.input ?? 0,
-						output: message.usage.output ?? 0,
-						cacheRead: message.usage.cacheRead ?? 0,
-						cacheWrite: message.usage.cacheWrite ?? 0,
-					},
-					index: i,
-				};
-			}
-		}
-		if (message.timestamp !== undefined) {
-			latestPrefixTimestamp = Math.max(latestPrefixTimestamp, message.timestamp);
-		}
-	}
-	return usageInfo;
-}
-
-/** CC: roughTokenCountEstimationForBlock — per content block. */
-function estimateBlockTokens(block: unknown): number {
-	if (typeof block === "string") return roughTokenCountEstimation(block);
-	if (!block || typeof block !== "object") return 0;
-	const b = block as Record<string, unknown>;
-	switch (b.type) {
-		case "text":
-			return roughTokenCountEstimation(String(b.text ?? ""));
-		case "thinking":
-			return roughTokenCountEstimation(String(b.thinking ?? ""));
-		case "image":
-		case "document":
-			return IMAGE_DOCUMENT_TOKEN_ESTIMATE;
-		// pi toolResult message content (CC: tool_result block inside a user message)
-		case "toolResult":
-			return estimateContentTokens(b.content);
-		// pi assistant tool call (CC: tool_use block — name + serialized input)
-		case "toolCall":
-			return roughTokenCountEstimation(String(b.name ?? "") + safeJsonStringify(b.arguments ?? {}));
-		default:
-			// server_tool_use, web_search_tool_result, mcp_tool_use, etc. — text-like
-			// payloads where the stringify length tracks the serialized form.
-			return roughTokenCountEstimation(safeJsonStringify(block));
-	}
-}
-
-function estimateContentTokens(content: unknown): number {
-	if (!content) return 0;
-	if (typeof content === "string") return roughTokenCountEstimation(content);
-	if (!Array.isArray(content)) return roughTokenCountEstimation(safeJsonStringify(content));
-	let total = 0;
-	for (const block of content) total += estimateBlockTokens(block);
-	return total;
-}
-
-/**
- * CC: roughTokenCountEstimationForMessage — CC counts only assistant and user
- * message content (tool results ride inside user messages there). pi carries
- * tool results as their own `toolResult` messages, estimated via their content
- * blocks. pi also keeps `system` messages in the branch (CC has none in the
- * tail), so those are counted as chars/4 plus their tool definitions.
- */
-export function estimateMessageTokens(message: RoughMessage): number {
-	if (message.role === "system") {
-		let chars = 0;
-		const content = message.content;
-		if (typeof content === "string") {
-			chars += content.length;
-		} else if (Array.isArray(content)) {
-			for (const block of content) {
-				if (block && typeof block === "object" && "text" in block) {
-					chars += String((block as { text: unknown }).text ?? "").length;
-				}
-			}
-		}
-		const toolsAdded = message.toolsAdded?.length ? safeJsonStringify(message.toolsAdded).length : 0;
-		const toolsRemoved = message.toolsRemoved?.length ? safeJsonStringify(message.toolsRemoved).length : 0;
-		chars += toolsAdded + toolsRemoved;
-		return Math.round(chars / CHARS_PER_TOKEN);
-	}
-	if (
-		(message.role === "assistant" || message.role === "user" || message.role === "toolResult") &&
-		message.content != null
-	) {
-		return estimateContentTokens(message.content);
-	}
-	return 0;
-}
-
-/**
- * CC: tokenCountWithEstimation — THE canonical context measurement for
- * threshold checks. Last valid API usage (exact tokens for everything up to
- * that response) plus a rough estimate of everything appended since.
- *
- * Note: CC walks back past sibling assistant records split from the same API
- * response (shared message.id, interleaved tool_results). pi does not split
- * assistant records per content block, so no walk-back is needed.
- */
-export function tokenCountWithEstimation(messages: readonly RoughMessage[]): number {
-	const usageInfo = findLastUsageInfo(messages);
-	if (usageInfo) {
-		let trailingTokens = 0;
-		for (let i = usageInfo.index + 1; i < messages.length; i++) {
-			trailingTokens += estimateMessageTokens(messages[i]);
-		}
-		return getTokenCountFromUsage(usageInfo.usage) + trailingTokens;
-	}
-	let tokens = 0;
-	for (const message of messages) tokens += estimateMessageTokens(message);
-	return tokens;
-}