@selesai/code 0.13.2 → 0.13.3

package/CHANGELOG.md

@@ -2,6 +2,15 @@
 
 All notable changes to `@selesai/code` will be documented in this file.
 
+## [0.13.3] - 2026-08-31
+
+### Added
+- **Token-In multi-key auto-failover.** The `tokenin-onboarding` extension now registers a `tokenin` provider whose `streamSimple` automatically rotates to an alternative saved account when the active key fails with a 401, 429, or budget/quota error. Failed keys go on a 5-minute in-memory cooldown before being retried, and the newly active credential is persisted to both `tokenin-auth.json` and `auth.json`.
+
+### Changed
+- **Cost reconciliation replaces the finalized assistant cost.** The `cost-reconcile` fetch wrapper now finishes capturing the provider-reported cost before the provider SDK finalizes its assistant message, so `message_end` replaces `usage.cost.total` with the billed amount before the message is persisted. Providers whose payloads carry no cost keep their rate-card estimate.
+- **Default model catalog reasoning maps.** Added `thinkingLevelMap` entries (and `low`=low / `max` mappings) across bundled model defaults and enabled reasoning on Qwen3.8 27B and Gemini 3.7 Flash so thinking budgets map correctly.
+
 ## [0.13.2] - 2026-08-30
 
 ### Added

package/dist/defaults/models.json

@@ -19,7 +19,7 @@
           "maxTokens": 64000,
           "thinkingLevelMap": {
             "minimal": null,
-            "low": null,
+            "low": "low",
             "medium": null,
             "high": "high",
             "xhigh": null,
@@ -38,7 +38,7 @@
           "maxTokens": 64000,
           "thinkingLevelMap": {
             "minimal": null,
-            "low": null,
+            "low": "low",
             "medium": null,
             "high": "high",
             "xhigh": null,
@@ -57,7 +57,7 @@
           "maxTokens": 64000,
           "thinkingLevelMap": {
             "minimal": null,
-            "low": null,
+            "low": "low",
             "medium": null,
             "high": "high",
             "xhigh": null,
@@ -114,7 +114,7 @@
           "maxTokens": 64000,
           "thinkingLevelMap": {
             "minimal": null,
-            "low": null,
+            "low": "low",
             "medium": null,
             "high": "high",
             "xhigh": null,
@@ -130,7 +130,15 @@
           "reasoning": true,
           "input": ["text", "image"],
           "contextWindow": 512000,
-          "maxTokens": 128000
+          "maxTokens": 128000,
+          "thinkingLevelMap": {
+            "minimal": null,
+            "low": "low",
+            "medium": null,
+            "high": "high",
+            "xhigh": null,
+            "max": "max"
+          }
         },
         {
           "id": "gpt-5.6-terra",
@@ -138,7 +146,15 @@
           "reasoning": true,
           "input": ["text", "image"],
           "contextWindow": 512000,
-          "maxTokens": 128000
+          "maxTokens": 128000,
+          "thinkingLevelMap": {
+            "minimal": null,
+            "low": "low",
+            "medium": null,
+            "high": "high",
+            "xhigh": null,
+            "max": null
+          }
         },
         {
           "id": "gpt-5.6-sol",
@@ -146,7 +162,15 @@
           "reasoning": true,
           "input": ["text", "image"],
           "contextWindow": 512000,
-          "maxTokens": 128000
+          "maxTokens": 128000,
+          "thinkingLevelMap": {
+            "minimal": null,
+            "low": "low",
+            "medium": null,
+            "high": "high",
+            "xhigh": null,
+            "max": null
+          }
         },
         {
           "id": "kimi-k3",
@@ -209,16 +233,24 @@
             "medium": null,
             "high": "high",
             "xhigh": null,
-            "max": null
+            "max": "max"
           }
         },
         {
           "id": "qwen3.8-27b",
           "name": "Qwen3.8 27B",
-          "reasoning": false,
+          "reasoning": true,
           "input": ["text", "image"],
           "contextWindow": 256000,
           "maxTokens": 64000,
+          "thinkingLevelMap": {
+            "minimal": null,
+            "low": "low",
+            "medium": null,
+            "high": "high",
+            "xhigh": null,
+            "max": "max"
+          },
           "compat": {
             "thinkingFormat": "qwen"
           }
@@ -226,10 +258,18 @@
         {
           "id": "gemini-3.7-flash",
           "name": "Gemini 3.7 Flash",
-          "reasoning": false,
-          "input": ["text", "image" ],
+          "reasoning": true,
+          "input": ["text", "image"],
           "contextWindow": 512000,
-          "maxTokens": 64000
+          "maxTokens": 64000,
+          "thinkingLevelMap": {
+            "minimal": null,
+            "low": "low",
+            "medium": null,
+            "high": "high",
+            "xhigh": null,
+            "max": null
+          }
         },
         {
           "id": "Qwen3-VL",
@@ -238,6 +278,14 @@
           "input": ["text", "image"],
           "contextWindow": 256000,
           "maxTokens": 8192,
+          "thinkingLevelMap": {
+            "minimal": null,
+            "low": "low",
+            "medium": null,
+            "high": "high",
+            "xhigh": null,
+            "max": null
+          },
           "compat": {
             "supportsDeveloperRole": false,
             "supportsReasoningEffort": false,
@@ -251,6 +299,14 @@
           "input": ["text", "image"],
           "contextWindow": 256000,
           "maxTokens": 8192,
+          "thinkingLevelMap": {
+            "minimal": null,
+            "low": "low",
+            "medium": null,
+            "high": "high",
+            "xhigh": null,
+            "max": null
+          },
           "compat": {
             "supportsDeveloperRole": false,
             "supportsReasoningEffort": false,

package/dist/extensions/cost-reconcile.test.ts

@@ -1,5 +1,63 @@
-import { describe, expect, it } from "vitest";
-import { extractCosts, parseCostHeader } from "./cost-reconcile.ts";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import costReconcileExtension, { extractCosts, parseCostHeader } from "./cost-reconcile.ts";
+
+type Handler = (event: any, ctx: any) => any;
+
+const FETCH_PATCHED = Symbol.for("selesai.cost-reconcile.fetch-patched");
+let originalFetch: typeof globalThis.fetch;
+
+function createHarness() {
+	const handlers = new Map<string, Handler>();
+	const appendEntry = vi.fn();
+	const pi = {
+		on: vi.fn((event: string, handler: Handler) => handlers.set(event, handler)),
+		appendEntry,
+	};
+	costReconcileExtension(pi as any);
+	return { appendEntry, handlers };
+}
+
+async function startSession(handlers: Map<string, Handler>): Promise<void> {
+	await handlers.get("session_start")!({}, { sessionManager: { getEntries: () => [] } });
+}
+
+function assistantMessage(responseId: string, total = 0.75) {
+	return {
+		role: "assistant",
+		provider: "openrouter",
+		model: "openai/gpt-4.1-mini",
+		responseModel: "openai/gpt-4.1-mini",
+		responseId,
+		stopReason: "stop",
+		content: [],
+		usage: {
+			input: 10,
+			output: 5,
+			cacheRead: 2,
+			cacheWrite: 1,
+			cost: { input: 0.2, output: 0.4, cacheRead: 0.1, cacheWrite: 0.05, total },
+		},
+	};
+}
+
+function installFetch(body: string, headers: Record<string, string> = {}) {
+	const upstreamFetch = vi.fn(async () =>
+		new Response(body, { headers: { "content-type": "text/event-stream", ...headers } }),
+	);
+	globalThis.fetch = upstreamFetch as typeof globalThis.fetch;
+	return upstreamFetch;
+}
+
+beforeEach(() => {
+	originalFetch = globalThis.fetch;
+	delete (globalThis as typeof globalThis & { [FETCH_PATCHED]?: boolean })[FETCH_PATCHED];
+});
+
+afterEach(() => {
+	globalThis.fetch = originalFetch;
+	delete (globalThis as typeof globalThis & { [FETCH_PATCHED]?: boolean })[FETCH_PATCHED];
+	vi.restoreAllMocks();
+});
 
 describe("parseCostHeader", () => {
 	it("parses LiteLLM scientific-notation cost", () => {
@@ -10,11 +68,17 @@ describe("parseCostHeader", () => {
 		expect(parseCostHeader("0.00123")).toBe(0.00123);
 	});
 
-	it("rejects null, garbage, and negatives", () => {
+	it("rejects null, blank, garbage, and negative values", () => {
 		expect(parseCostHeader(null)).toBeUndefined();
+		expect(parseCostHeader("")).toBeUndefined();
+		expect(parseCostHeader(" \t ")).toBeUndefined();
 		expect(parseCostHeader("abc")).toBeUndefined();
 		expect(parseCostHeader("-1")).toBeUndefined();
 	});
+
+	it("trims valid zero-valued headers", () => {
+		expect(parseCostHeader(" 0 ")).toBe(0);
+	});
 });
 
 describe("extractCosts", () => {
@@ -73,4 +137,117 @@ describe("extractCosts", () => {
 		const { ids } = extractCosts(body);
 		expect(ids.length).toBe(50);
 	});
-});
+});
+
+describe("live assistant usage reconciliation", () => {
+	it("replaces finalized assistant usage only after the streamed body completes and keeps the custom entry", async () => {
+		const responseId = "live-reconcile-1";
+		const body = `data: {"id":"${responseId}","usage":{"cost":0.0123}}\n\ndata: [DONE]\n\n`;
+		installFetch(body);
+		const { appendEntry, handlers } = createHarness();
+		await startSession(handlers);
+
+		const response = await globalThis.fetch("https://openrouter.ai/api/v1/chat/completions");
+		const messageEnd = handlers.get("message_end")!;
+		expect(await messageEnd({ message: assistantMessage(responseId) }, {})).toBeUndefined();
+
+		expect(await response.text()).toBe(body);
+		const result = await messageEnd({ message: assistantMessage(responseId) }, {});
+		expect(result.message.usage.cost).toEqual({
+			input: 0.2,
+			output: 0.4,
+			cacheRead: 0.1,
+			cacheWrite: 0.05,
+			total: 0.0123,
+		});
+		expect(appendEntry).toHaveBeenCalledWith(
+			"cost-reconcile",
+			expect.objectContaining({ provider: "openrouter", responseId, cost: 0.0123, source: "payload" }),
+		);
+	});
+
+	it("uses a valid zero-valued billed header over the payload cost", async () => {
+		const responseId = "live-header-zero-2";
+		installFetch(`data: {"id":"${responseId}","usage":{"cost":0.0123}}\n\n`, {
+			"x-litellm-response-cost": "0",
+		});
+		const { handlers } = createHarness();
+		await startSession(handlers);
+
+		const response = await globalThis.fetch("https://gateway.example/v1/chat/completions");
+		await response.text();
+		const result = await handlers.get("message_end")!({ message: assistantMessage(responseId) }, {});
+		expect(result.message.usage.cost.total).toBe(0);
+	});
+
+	it("falls back when a captured body has multiple candidate response ids", async () => {
+		const responseId = "live-ambiguous-response-3";
+		installFetch(
+			`data: {"id":"${responseId}","usage":{"cost":0.0123},"tool_calls":[{"id":"tool-ambiguous-3"}]}\n\n`,
+		);
+		const { appendEntry, handlers } = createHarness();
+		await startSession(handlers);
+
+		const response = await globalThis.fetch("https://gateway.example/v1/chat/completions");
+		await response.text();
+		const message = assistantMessage(responseId);
+		expect(await handlers.get("message_end")!({ message }, {})).toBeUndefined();
+		expect(message.usage.cost.total).toBe(0.75);
+		expect(appendEntry).not.toHaveBeenCalled();
+	});
+
+	it("falls back when the process cache sees a duplicate response id", async () => {
+		const responseId = "live-duplicate-response-4";
+		installFetch(`data: {"id":"${responseId}","usage":{"cost":0.0123}}\n\n`);
+		const { appendEntry, handlers } = createHarness();
+		await startSession(handlers);
+
+		const first = await globalThis.fetch("https://gateway.example/v1/chat/completions");
+		const second = await globalThis.fetch("https://gateway.example/v1/chat/completions");
+		await Promise.all([first.text(), second.text()]);
+		const message = assistantMessage(responseId);
+		expect(await handlers.get("message_end")!({ message }, {})).toBeUndefined();
+		expect(message.usage.cost.total).toBe(0.75);
+		expect(appendEntry).not.toHaveBeenCalled();
+	});
+
+	it("preserves the original message when no valid billed cost was captured", async () => {
+		const responseId = "live-no-cost-3";
+		installFetch(`data: {"id":"${responseId}","usage":{"cost":-1}}\n\n`);
+		const { appendEntry, handlers } = createHarness();
+		await startSession(handlers);
+
+		const response = await globalThis.fetch("https://gateway.example/v1/chat/completions");
+		await response.text();
+		const message = assistantMessage(responseId);
+		expect(await handlers.get("message_end")!({ message }, {})).toBeUndefined();
+		expect(message.usage.cost.total).toBe(0.75);
+		expect(appendEntry).not.toHaveBeenCalled();
+	});
+
+	it("preserves terminal error messages even when a billed cost was captured", async () => {
+		const responseId = "live-terminal-4";
+		installFetch(`data: {"id":"${responseId}","usage":{"cost":0.0123}}\n\n`);
+		const { appendEntry, handlers } = createHarness();
+		await startSession(handlers);
+
+		const response = await globalThis.fetch("https://gateway.example/v1/chat/completions");
+		await response.text();
+		const message = { ...assistantMessage(responseId), stopReason: "error" };
+		expect(await handlers.get("message_end")!({ message }, {})).toBeUndefined();
+		expect(message.usage.cost.total).toBe(0.75);
+		expect(appendEntry).not.toHaveBeenCalled();
+	});
+
+	it("does not capture unrelated fetches", async () => {
+		const responseId = "unrelated-fetch-5";
+		const body = `data: {"id":"${responseId}","usage":{"cost":0.0123}}\n\n`;
+		installFetch(body);
+		const { handlers } = createHarness();
+		await startSession(handlers);
+
+		const response = await globalThis.fetch("https://example.com/data.json");
+		expect(await response.text()).toBe(body);
+		expect(await handlers.get("message_end")!({ message: assistantMessage(responseId) }, {})).toBeUndefined();
+	});
+});

package/dist/extensions/cost-reconcile.ts

@@ -9,19 +9,16 @@
  *
  * Extensions cannot hook the response body through the event API
  * (`after_provider_response` exposes only status/headers), but provider SDKs
- * fall back to `globalThis.fetch`. This extension installs a tee-wrapping
- * fetch once per process: LLM API responses are scanned for a
- * provider-reported cost (LiteLLM `x-litellm-response-cost` header,
- * OpenRouter `usage.cost`, `total_cost`, plus object-shaped `cost.total`)
- * and its response id, then recorded as a custom session entry on
- * `message_end`.
+ * fall back to `globalThis.fetch`. This extension installs a stream-wrapping
+ * fetch once per process: LLM API responses are scanned for a provider-reported
+ * cost (LiteLLM `x-litellm-response-cost` header, OpenRouter `usage.cost`,
+ * `total_cost`, plus object-shaped `cost.total`) and response id. The wrapper
+ * finishes capture before the provider SDK finalizes its assistant message, so
+ * `message_end` can replace `usage.cost.total` before session persistence.
  *
- * Providers whose payloads carry no cost (first-party OpenAI/Anthropic due to
- * their usage APIs) simply record nothing — their rate-card estimate matches
- * their bill anyway when catalog prices are configured.
- *
- * The zentui footer prefers reconciled entries over the rate-card total when
- * both exist for the same response id.
+ * Providers whose payloads carry no cost retain their rate-card estimate.
+ * Reconciled entries remain persisted for zentui and existing-session
+ * compatibility.
  */
 
 import type { ExtensionAPI } from "@selesai/code";
@@ -31,11 +28,10 @@ const ENTRY_VERSION = 1 as const;
 
 // Guard against pathological bodies; real usage chunks arrive well under 1MB.
 const MAX_CAPTURE_BYTES = 4 * 1024 * 1024;
-// Only LLM API paths are tee'd and scanned; unrelated fetches (web fetch tool,
-// model catalog refreshes) pass through untouched.
+// Only LLM API paths are scanned; unrelated fetches (web fetch tool, model
+// catalog refreshes) pass through untouched.
 const LLM_PATH_RE = /(\/chat\/completions|\/responses|\/messages\b|:streamGenerateContent)/;
 const MAX_CAPTURED_IDS = 50;
-const MAX_PENDING_LOOKUPS = 200;
 const MAX_SESSION_ENTRIES = 2_000;
 
 type ReconcileEntry = {
@@ -49,46 +45,34 @@ type ReconcileEntry = {
 	source: "payload";
 };
 
-/**
- * Scan a captured SSE/JSON body for provider-reported cost and response ids.
- * Exported for tests.
- *
- * In streams, usage lands in the final chunk, so the LAST cost match wins.
- */
 type MessageInfo = {
 	provider: string;
 	model: string;
 	responseId: string;
 };
 
-// Process-wide capture state. The fetch patch is installed once per process
-// and must not capture session-bound `pi`; per-session message_end handlers
-// read these maps and write entries through their own session's appendEntry.
-const processCaptured = new Map<string, number>();
-const processWaiting = new Map<
-	string,
-	{ info: MessageInfo; write: (info: MessageInfo, cost: number) => void }
->();
+// Process-wide capture state. The fetch patch is installed once per process;
+// per-session message_end handlers consume the matching response id. A duplicate
+// capture for the same id is marked ambiguous rather than assigning either bill.
+const processCaptured = new Map<string, number | null>();
 
-/**
- * Scan a captured SSE/JSON body for provider-reported cost and response ids.
- * Pure function; exported for tests.
- *
- * In streams, usage lands in the final chunk, so the LAST cost match wins.
- */
 /**
  * Parse a LiteLLM `x-litellm-response-cost` header value (a USD float,
  * possibly in scientific notation). Exported for tests.
  */
 export function parseCostHeader(value: string | null): number | undefined {
-	if (value === null) return undefined;
-	const cost = Number(value);
+	const trimmed = value?.trim();
+	if (!trimmed) return undefined;
+	const cost = Number(trimmed);
 	return Number.isFinite(cost) && cost >= 0 ? cost : undefined;
 }
 
-export function extractCosts(
-	text: string,
-): { ids: string[]; cost?: number } {
+/**
+ * Scan a captured SSE/JSON body for provider-reported cost and response ids.
+ * Exported for tests. In streams, usage lands in the final chunk, so the LAST
+ * valid cost match wins.
+ */
+export function extractCosts(text: string): { ids: string[]; cost?: number } {
 	const ids = new Set<string>();
 	for (const match of text.matchAll(/"id"\s*:\s*"([^"]{1,200})"/g)) {
 		ids.add(match[1]);
@@ -151,26 +135,22 @@ export default function costReconcileExtension(pi: ExtensionAPI): void {
 		try {
 			pi.appendEntry(ENTRY_TYPE, entry);
 		} catch {
-			// transcript persistence failure must not break the session
+			// Transcript persistence failure must not break the session.
 		}
 	};
 
 	const finishCapture = (rawBody: string, headerCost?: number): void => {
-		// ponytail: one naive regex pass over the whole body; a streaming
-		// parser only matters if bodies grow past the 4MB cap.
+		// Ponytail: one naive regex pass over the whole body; a streaming parser
+		// only matters if bodies grow past the 4MB cap.
 		const { ids, cost } = extractCosts(rawBody);
-		// LiteLLM's header is the amount the gateway billed; prefer it over
-		// any cost the upstream payload reports.
+		// LiteLLM's header is the amount the gateway billed; prefer it over any
+		// cost the upstream payload reports.
 		const effectiveCost = headerCost ?? cost;
-		if (effectiveCost === undefined) return;
-		for (const responseId of ids) {
-			processCaptured.set(responseId, effectiveCost);
-			const waiting = processWaiting.get(responseId);
-			if (waiting) {
-				processWaiting.delete(responseId);
-				waiting.write(waiting.info, effectiveCost);
-			}
-		}
+		// A cost must identify exactly one response. Applying it to every `id` in
+		// a body can mistake nested tool/message ids for the assistant response.
+		if (effectiveCost === undefined || ids.length !== 1) return;
+		const responseId = ids[0]!;
+		processCaptured.set(responseId, processCaptured.has(responseId) ? null : effectiveCost);
 		if (processCaptured.size > 500) {
 			for (const key of [...processCaptured.keys()].slice(0, processCaptured.size - 500))
 				processCaptured.delete(key);
@@ -191,7 +171,7 @@ export default function costReconcileExtension(pi: ExtensionAPI): void {
 				if (typeof responseId === "string") seen.add(responseId);
 			}
 		} catch {
-			// best effort only
+			// Best effort only.
 		}
 
 		const globalFetch = globalThis as typeof globalThis & { [FETCH_PATCHED]?: boolean };
@@ -205,26 +185,36 @@ export default function costReconcileExtension(pi: ExtensionAPI): void {
 				const contentType = response.headers.get("content-type") ?? "";
 				if (!contentType.includes("json") && !contentType.includes("event-stream")) return response;
 				if (!response.body) return response;
+
 				const headerCost = parseCostHeader(response.headers.get("x-litellm-response-cost"));
-				const [main, tee] = response.body.tee();
-				void (async () => {
-					let raw = "";
-					try {
-						const reader = tee.getReader();
-						const decoder = new TextDecoder();
-						for (;;) {
-							const { done, value } = await reader.read();
-							if (done) break;
-							raw += decoder.decode(value, { stream: true });
-							if (raw.length > MAX_CAPTURE_BYTES) break;
+				let raw = "";
+				let captureFailed = false;
+				const decoder = new TextDecoder();
+				const capture = new TransformStream<Uint8Array, Uint8Array>({
+					transform(chunk, controller) {
+						controller.enqueue(chunk);
+						if (captureFailed) return;
+						try {
+							raw += decoder.decode(chunk, { stream: true });
+							if (raw.length > MAX_CAPTURE_BYTES) captureFailed = true;
+						} catch {
+							captureFailed = true;
 						}
-					} catch {
-						// capture is best-effort; the main branch is unaffected
-					}
-					finishCapture(raw, headerCost);
-				})();
-				// Preserve response identity as seen by the SDK: body swapped, rest identical.
-				return new Response(main, {
+					},
+					flush() {
+						if (captureFailed) return;
+						try {
+							raw += decoder.decode();
+							finishCapture(raw, headerCost);
+						} catch {
+							// Capture is best-effort; the response stream is unaffected.
+						}
+					},
+				});
+				// flush() runs before the SDK observes EOF, so message_end can safely
+				// replace the finalized message usage without waiting on another fetch.
+				const body = response.body.pipeThrough(capture);
+				return new Response(body, {
 					status: response.status,
 					statusText: response.statusText,
 					headers: response.headers,
@@ -236,43 +226,32 @@ export default function costReconcileExtension(pi: ExtensionAPI): void {
 	});
 
 	pi.on("message_end", (event) => {
-		const message = event.message as
-			| {
-					role?: string;
-					provider?: string;
-					model?: string;
-					responseId?: string;
-					stopReason?: string;
-				}
-			| undefined;
-		if (!message || message.role !== "assistant") return;
+		const message = event.message;
+		if (message.role !== "assistant") return;
 		if (message.stopReason === "error" || message.stopReason === "aborted") return;
 		const responseId = message.responseId;
 		if (typeof responseId !== "string" || responseId.length === 0) return;
 
-		const info: MessageInfo = {
-			provider: message.provider ?? "",
-			model: message.responseModel ?? message.model ?? "",
-			responseId,
-		};
 		const cost = processCaptured.get(responseId);
-		if (cost !== undefined) {
-			writeEntry(info, cost);
+		if (cost === undefined || cost === null || !message.usage?.cost) {
+			processCaptured.delete(responseId);
 			return;
 		}
-		// The tee branch may finish after message_end; remember the message so
-		// the capture can match it on completion and write via this session.
-		if (!seen.has(responseId) && !processWaiting.has(responseId)) {
-			waitingLimitGuard();
-			processWaiting.set(responseId, { info, write: writeEntry });
-		}
+		processCaptured.delete(responseId);
+
+		writeEntry(
+			{
+				provider: message.provider ?? "",
+				model: message.responseModel ?? message.model ?? "",
+				responseId,
+			},
+			cost,
+		);
+		return {
+			message: {
+				...message,
+				usage: { ...message.usage, cost: { ...message.usage.cost, total: cost } },
+			},
+		};
 	});
 }
-
-// Trim oldest waiting entries so unclaimed captures cannot grow unbounded.
-function waitingLimitGuard(): void {
-	if (processWaiting.size >= MAX_PENDING_LOOKUPS) {
-		for (const key of [...processWaiting.keys()].slice(0, processWaiting.size - MAX_PENDING_LOOKUPS + 1))
-			processWaiting.delete(key);
-	}
-}

package/dist/extensions/tokenin-onboarding.ts

@@ -13,6 +13,16 @@ import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "n
 import { dirname, join } from "node:path";
 import { getAgentDir, getModelsPath } from "@selesai/code";
 import type { AuthStorage, ExtensionAPI, ExtensionCommandContext, ExtensionContext, SessionStartEvent } from "@selesai/code";
+import {
+	createAssistantMessageEventStream,
+	lazyStream,
+	type AssistantMessageEvent,
+	type AssistantMessageEventStream,
+	type Context,
+	type Model,
+	type SimpleStreamOptions,
+} from "@earendil-works/pi-ai";
+import { getApiProvider } from "@earendil-works/pi-ai/compat";
 
 export const TOKEN_IN_PROVIDER = "tokenin";
 export const TOKEN_IN_DASHBOARD_URL = "https://token.selesai.in/dashboard/tokens";
@@ -216,6 +226,291 @@ function isOnboardingComplete(agentDir: string = getAgentDir()): boolean {
 	return existsSync(getOnboardingMarkerPath(agentDir));
 }
 
+// ---------------------------------------------------------------------------
+// Multi-key load balancing / auto-failover
+// ---------------------------------------------------------------------------
+
+export const DEFAULT_TOKENIN_COOLDOWN_MS = 5 * 60 * 1000; // 5 minutes
+
+/** In-memory cooldown tracker mapping token/accountId -> cooldown expires timestamp (ms). */
+const tokenCooldowns = new Map<string, number>();
+
+/** Clear all in-memory cooldowns. Exported for tests. */
+export function clearTokenInCooldowns(): void {
+	tokenCooldowns.clear();
+}
+
+/** Check if a token/account is currently on cooldown. */
+export function isTokenInOnCooldown(tokenId: string, now: number = Date.now()): boolean {
+	const expires = tokenCooldowns.get(tokenId);
+	if (expires === undefined) return false;
+	if (now >= expires) {
+		tokenCooldowns.delete(tokenId);
+		return false;
+	}
+	return true;
+}
+
+/** Put a token/account on cooldown for a given duration. */
+export function setTokenInCooldown(tokenId: string, durationMs: number = DEFAULT_TOKENIN_COOLDOWN_MS, now: number = Date.now()): void {
+	tokenCooldowns.set(tokenId, now + durationMs);
+}
+
+/** Detect if an error message or AssistantMessage stop indicates an auth/quota failure eligible for failover. */
+export function isRotateableTokenInError(error: unknown): boolean {
+	if (!error) return false;
+	let text = "";
+	if (typeof error === "string") {
+		text = error;
+	} else if (error instanceof Error) {
+		text = `${error.name}: ${error.message}`;
+	} else if (typeof error === "object" && error !== null) {
+		const obj = error as Record<string, unknown>;
+		if (typeof obj.errorMessage === "string") {
+			text = obj.errorMessage;
+		} else if (typeof obj.message === "string") {
+			text = obj.message;
+		} else if (Array.isArray(obj.content)) {
+			// AssistantMessage with content blocks containing error text or errorMessage
+			for (const block of obj.content) {
+				if (typeof block === "object" && block !== null) {
+					const b = block as Record<string, unknown>;
+					if (typeof b.errorMessage === "string") text += ` ${b.errorMessage}`;
+					if (typeof b.text === "string") text += ` ${b.text}`;
+				}
+			}
+		}
+		if (!text) text = JSON.stringify(error);
+	} else {
+		text = JSON.stringify(error);
+	}
+
+	const lower = text.toLowerCase();
+	// HTTP 401 / Invalid Key
+	if (lower.includes("401") || lower.includes("unauthorized") || lower.includes("invalid_api_key") || lower.includes("invalid api key") || lower.includes("authentication_error")) {
+		return true;
+	}
+	// HTTP 429 / Rate limit
+	if (lower.includes("429") || lower.includes("rate_limit") || lower.includes("rate limit") || lower.includes("too many requests")) {
+		return true;
+	}
+	// Budget / Quota / Credits
+	if (
+		lower.includes("budget has been exceeded") ||
+		lower.includes("budget_exceeded") ||
+		lower.includes("max_budget") ||
+		lower.includes("insufficient_quota") ||
+		lower.includes("quota exceeded") ||
+		lower.includes("exceeds max_budget") ||
+		lower.includes("credit limit") ||
+		lower.includes("out of credits")
+	) {
+		return true;
+	}
+	return false;
+}
+
+/** Helper to persist active credential to auth.json and tokenin-auth.json without needing AuthStorage instance. */
+function persistActiveTokenInAccount(account: TokenInAccount, authPath: string = getTokenInAuthPath()): void {
+	try {
+		// Update tokenin-auth.json
+		const auth = readTokenInAuth(authPath);
+		auth.activeId = account.id;
+		writeTokenInAuth(auth, authPath);
+
+		// Update auth.json
+		const agentDir = dirname(authPath);
+		const globalAuthPath = join(agentDir, "auth.json");
+		const current = existsSync(globalAuthPath) ? JSON.parse(readFileSync(globalAuthPath, "utf-8")) as Record<string, unknown> : {};
+		current[TOKEN_IN_PROVIDER] = { type: "api_key", key: account.apiKey };
+		mkdirSync(dirname(globalAuthPath), { recursive: true, mode: 0o700 });
+		writeFileSync(globalAuthPath, `${JSON.stringify(current, null, 2)}\n`, { encoding: "utf-8", mode: 0o600 });
+		try {
+			chmodSync(globalAuthPath, 0o600);
+		} catch {
+			// best-effort
+		}
+	} catch {
+		// best-effort sync
+	}
+}
+
+/**
+ * Custom streamSimple implementation for tokenin provider that wraps OpenAI completions
+ * and automatically rotates to alternative saved keys on 401/429/budget exceeded errors.
+ */
+export function createTokenInStreamSimple(options?: {
+	authPath?: string;
+	getAuthStorage?: () => AuthStorage | undefined;
+	onRotate?: (failedAccount: TokenInAccount, nextAccount: TokenInAccount, reason: string) => void;
+	streamSimple?: (model: Model<any>, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream;
+}) {
+	const authPath = options?.authPath ?? getTokenInAuthPath();
+	const baseApi = options?.streamSimple ? undefined : getApiProvider("openai-completions");
+	const streamSimple = options?.streamSimple ?? baseApi!.streamSimple.bind(baseApi!);
+
+	const activate = (account: TokenInAccount): void => {
+		const authStorage = options?.getAuthStorage?.();
+		if (authStorage) {
+			applyTokenInAccountToAuth(account, authStorage, authPath);
+		} else {
+			persistActiveTokenInAccount(account, authPath);
+		}
+	};
+
+	const withAccountAuthorization = (streamOptions: SimpleStreamOptions | undefined, account: TokenInAccount): SimpleStreamOptions => {
+		const headers = Object.fromEntries(
+			Object.entries(streamOptions?.headers ?? {}).filter(([name]) => name.toLowerCase() !== "authorization"),
+		);
+		return {
+			...streamOptions,
+			apiKey: account.apiKey,
+			headers: { ...headers, Authorization: `Bearer ${account.apiKey}` },
+		};
+	};
+
+	return function tokenInStreamSimple(
+		model: Model<any>,
+		context: Context,
+		streamOptions?: SimpleStreamOptions,
+	): AssistantMessageEventStream {
+		return lazyStream(model, async () => {
+			const auth = readTokenInAuth(authPath);
+			const accounts = auth.accounts;
+
+			// If no accounts saved, fall back directly to base implementation
+			if (accounts.length === 0) {
+				return streamSimple(model, context, streamOptions);
+			}
+
+			// Order accounts: active account first, followed by remaining accounts
+			const activeIndex = accounts.findIndex((a) => a.id === auth.activeId || a.apiKey === streamOptions?.apiKey);
+			const orderedAccounts: TokenInAccount[] = [];
+			if (activeIndex >= 0) {
+				orderedAccounts.push(accounts[activeIndex]!);
+				for (let i = 0; i < accounts.length; i++) {
+					if (i !== activeIndex) orderedAccounts.push(accounts[i]!);
+				}
+			} else {
+				orderedAccounts.push(...accounts);
+			}
+
+			// Filter/order candidates: prioritize those not currently on cooldown
+			const availableAccounts = orderedAccounts.filter((a) => !isTokenInOnCooldown(a.id));
+			const candidateAccounts = (availableAccounts.length > 0 ? availableAccounts : orderedAccounts).map((account) => ({ ...account }));
+
+			let lastErrorEvent: AssistantMessageEvent | undefined;
+			let lastThrownError: unknown;
+
+			for (let i = 0; i < candidateAccounts.length; i++) {
+				const currentAccount = candidateAccounts[i]!;
+
+				const requestOptions = withAccountAuthorization(streamOptions, currentAccount);
+
+				let underlyingStream: AssistantMessageEventStream;
+				try {
+					underlyingStream = streamSimple(model, context, requestOptions);
+				} catch (err) {
+					lastThrownError = err;
+					if (isRotateableTokenInError(err) && i + 1 < candidateAccounts.length) {
+						setTokenInCooldown(currentAccount.id);
+						const nextAccount = candidateAccounts[i + 1]!;
+						activate(nextAccount);
+						options?.onRotate?.(currentAccount, nextAccount, String(err));
+						continue;
+					}
+					throw err;
+				}
+
+				// Hold a leading start event until the request proves it produced output. This
+				// permits an immediate upstream error to be retried without exposing two starts.
+				const iterator = underlyingStream[Symbol.asyncIterator]();
+				let firstResult: IteratorResult<AssistantMessageEvent>;
+				try {
+					firstResult = await iterator.next();
+				} catch (err) {
+					lastThrownError = err;
+					if (isRotateableTokenInError(err) && i + 1 < candidateAccounts.length) {
+						setTokenInCooldown(currentAccount.id);
+						const nextAccount = candidateAccounts[i + 1]!;
+						activate(nextAccount);
+						options?.onRotate?.(currentAccount, nextAccount, String(err));
+						continue;
+					}
+					throw err;
+				}
+
+				if (firstResult.done) {
+					const passthrough = createAssistantMessageEventStream();
+					passthrough.end();
+					return passthrough;
+				}
+
+				const bufferedEvents = [firstResult.value];
+				if (bufferedEvents[0].type === "start") {
+					try {
+						const next = await iterator.next();
+						if (!next.done) bufferedEvents.push(next.value);
+					} catch (err) {
+						lastThrownError = err;
+						if (isRotateableTokenInError(err) && i + 1 < candidateAccounts.length) {
+							setTokenInCooldown(currentAccount.id);
+							const nextAccount = candidateAccounts[i + 1]!;
+							activate(nextAccount);
+							options?.onRotate?.(currentAccount, nextAccount, String(err));
+							continue;
+						}
+						throw err;
+					}
+				}
+
+				const immediateError = bufferedEvents.find((event) => event.type === "error");
+				if (immediateError?.type === "error" && isRotateableTokenInError(immediateError.error)) {
+					lastErrorEvent = immediateError;
+					setTokenInCooldown(currentAccount.id);
+					if (i + 1 < candidateAccounts.length) {
+						const nextAccount = candidateAccounts[i + 1]!;
+						activate(nextAccount);
+						options?.onRotate?.(currentAccount, nextAccount, immediateError.error.errorMessage ?? "Error");
+						continue;
+					}
+				}
+
+				const outputStream = createAssistantMessageEventStream();
+				(async () => {
+					try {
+						for (const event of bufferedEvents) outputStream.push(event);
+						while (true) {
+							const next = await iterator.next();
+							if (next.done) break;
+							outputStream.push(next.value);
+						}
+						outputStream.end();
+					} catch {
+						outputStream.end();
+					}
+				})();
+
+				return outputStream;
+			}
+
+			// If loop exhausted with a rotateable error event, return a stream with that error
+			if (lastErrorEvent) {
+				const errorStream = createAssistantMessageEventStream();
+				errorStream.push(lastErrorEvent);
+				errorStream.end();
+				return errorStream;
+			}
+
+			if (lastThrownError) {
+				throw lastThrownError;
+			}
+
+			return streamSimple(model, context, streamOptions);
+		});
+	};
+}
+
 async function openDashboard(pi: ExtensionAPI): Promise<void> {
 	const url = TOKEN_IN_DASHBOARD_URL;
 	let command: string;
@@ -537,7 +832,14 @@ async function runOnboarding(pi: ExtensionAPI, ctx: ExtensionContext): Promise<v
 }
 
 export default function tokenInOnboardingExtension(pi: ExtensionAPI): void {
+	let authStorage: AuthStorage | undefined;
+	pi.registerProvider("tokenin", {
+		api: "openai-completions",
+		streamSimple: createTokenInStreamSimple({ getAuthStorage: () => authStorage }),
+	});
+
 	pi.on("session_start", async (event: SessionStartEvent, ctx: ExtensionContext) => {
+		authStorage = ctx.modelRegistry.authStorage;
 		if (event.reason !== "startup") return;
 		if (isOnboardingComplete()) return;
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@selesai/code",
-  "version": "0.13.2",
+  "version": "0.13.3",
   "description": "Maintained, extension-first Pi coding agent with built-in workflows, subagents, web research, questions, skills, and an enhanced terminal UI.",
   "type": "module",
   "repository": {