responses-proxy 0.1.0

package/dist/kiro-token-store.js

@@ -0,0 +1,196 @@
+import { existsSync } from "node:fs";
+import BetterSqlite3 from "better-sqlite3";
+/**
+ * Reads (and optionally writes back) Kiro/CodeWhisperer OAuth accounts stored by
+ * the 9router app in its own SQLite database (providerConnections table).
+ *
+ * The same database is owned by a live 9router process, so write-back is scoped
+ * to only the token fields and preserves every other key 9router manages.
+ */
+export class KiroTokenStore {
+    db;
+    writeBackEnabled;
+    constructor(db, writeBackEnabled) {
+        this.db = db;
+        this.writeBackEnabled = writeBackEnabled;
+    }
+    static open(dbFile, options = {}) {
+        if (!existsSync(dbFile)) {
+            throw new KiroTokenStoreError(`9router database not found at ${dbFile}`);
+        }
+        const writeBack = options.writeBack !== false;
+        // When write-back is disabled, open read-only so we never touch a database the
+        // live 9router process owns (not even the journal_mode pragma, which is a write).
+        // SQLite can read a WAL database read-only as long as the -wal/-shm files exist.
+        const db = new BetterSqlite3(dbFile, { fileMustExist: true, readonly: !writeBack });
+        if (writeBack) {
+            // 9router runs in WAL mode; match it so our reads see committed writes and
+            // our write-backs cooperate with its connection instead of blocking it.
+            db.pragma("journal_mode = WAL");
+            db.pragma("busy_timeout = 4000");
+        }
+        return new KiroTokenStore(db, writeBack);
+    }
+    listAccounts() {
+        const rows = this.db
+            .prepare(`SELECT id, provider, authType, name, email, priority, isActive, data, createdAt, updatedAt
+         FROM providerConnections
+         WHERE provider = 'kiro'
+         ORDER BY priority IS NULL, priority, createdAt`)
+            .all();
+        return rows
+            .map((row) => mapConnectionRow(row))
+            .filter((account) => account !== undefined);
+    }
+    listAvailableAccounts() {
+        return this.listAccounts().filter((account) => account.isActive);
+    }
+    getAccount(id) {
+        const row = this.db
+            .prepare(`SELECT id, provider, authType, name, email, priority, isActive, data, createdAt, updatedAt
+         FROM providerConnections
+         WHERE provider = 'kiro' AND id = ?`)
+            .get(id);
+        return row ? mapConnectionRow(row) : undefined;
+    }
+    /**
+     * Persists refreshed tokens back into 9router's database. Re-reads the current
+     * row inside a transaction so we merge onto whatever 9router last wrote rather
+     * than clobbering concurrent updates. No-op when write-back is disabled.
+     */
+    updateTokens(id, update, now = new Date()) {
+        if (!this.writeBackEnabled) {
+            return this.getAccount(id);
+        }
+        const persist = this.db.transaction((accountId, patch) => {
+            const row = this.db
+                .prepare(`SELECT data FROM providerConnections WHERE provider = 'kiro' AND id = ?`)
+                .get(accountId);
+            if (!row) {
+                return false;
+            }
+            const data = safeParseObject(row.data);
+            data.accessToken = patch.accessToken;
+            if (patch.refreshToken) {
+                data.refreshToken = patch.refreshToken;
+            }
+            data.expiresAt = patch.expiresAt;
+            if (typeof patch.expiresIn === "number") {
+                data.expiresIn = patch.expiresIn;
+            }
+            // Clear transient error markers 9router sets so a successful refresh is reflected.
+            data.lastError = null;
+            data.lastErrorAt = null;
+            this.db
+                .prepare(`UPDATE providerConnections SET data = ?, updatedAt = ? WHERE provider = 'kiro' AND id = ?`)
+                .run(JSON.stringify(data), now.toISOString(), accountId);
+            return true;
+        });
+        const ok = persist(id, update);
+        return ok ? this.getAccount(id) : undefined;
+    }
+    /**
+     * Updates account properties (name, priority, isActive). Only works when write-back is enabled.
+     */
+    updateAccount(id, updates, now = new Date()) {
+        if (!this.writeBackEnabled) {
+            throw new KiroTokenStoreError("Cannot update account: write-back is disabled");
+        }
+        const updateTransaction = this.db.transaction((accountId, patch) => {
+            const row = this.db
+                .prepare(`SELECT name, priority, isActive, data FROM providerConnections WHERE provider = 'kiro' AND id = ?`)
+                .get(accountId);
+            if (!row) {
+                return false;
+            }
+            const updateFields = [];
+            const updateValues = [];
+            if (patch.name !== undefined) {
+                updateFields.push("name = ?");
+                updateValues.push(patch.name);
+            }
+            if (patch.priority !== undefined) {
+                updateFields.push("priority = ?");
+                updateValues.push(patch.priority);
+            }
+            if (patch.isActive !== undefined) {
+                updateFields.push("isActive = ?");
+                updateValues.push(patch.isActive ? 1 : 0);
+            }
+            if (updateFields.length === 0) {
+                return true; // No updates needed
+            }
+            updateFields.push("updatedAt = ?");
+            updateValues.push(now.toISOString());
+            updateValues.push(accountId);
+            const sql = `UPDATE providerConnections SET ${updateFields.join(", ")} WHERE provider = 'kiro' AND id = ?`;
+            this.db.prepare(sql).run(...updateValues);
+            return true;
+        });
+        const ok = updateTransaction(id, updates);
+        return ok ? this.getAccount(id) : undefined;
+    }
+    /**
+     * Deletes a Kiro account. Only works when write-back is enabled.
+     */
+    deleteAccount(id) {
+        if (!this.writeBackEnabled) {
+            throw new KiroTokenStoreError("Cannot delete account: write-back is disabled");
+        }
+        const result = this.db
+            .prepare(`DELETE FROM providerConnections WHERE provider = 'kiro' AND id = ?`)
+            .run(id);
+        return result.changes > 0;
+    }
+    close() {
+        this.db.close();
+    }
+}
+export class KiroTokenStoreError extends Error {
+}
+function mapConnectionRow(row) {
+    const data = safeParseObject(row.data);
+    const accessToken = typeof data.accessToken === "string" ? data.accessToken : "";
+    const refreshToken = typeof data.refreshToken === "string" ? data.refreshToken : "";
+    if (!accessToken && !refreshToken) {
+        return undefined;
+    }
+    const psdRaw = typeof data.providerSpecificData === "object" && data.providerSpecificData !== null
+        ? data.providerSpecificData
+        : {};
+    return {
+        id: row.id,
+        name: row.name?.trim() ? row.name.trim() : row.id,
+        priority: typeof row.priority === "number" ? row.priority : Number.MAX_SAFE_INTEGER,
+        isActive: row.isActive !== 0,
+        accessToken,
+        refreshToken,
+        expiresAt: typeof data.expiresAt === "string" ? data.expiresAt : null,
+        expiresIn: typeof data.expiresIn === "number" ? data.expiresIn : null,
+        providerSpecificData: {
+            profileArn: optionalString(psdRaw.profileArn),
+            clientId: optionalString(psdRaw.clientId),
+            clientSecret: optionalString(psdRaw.clientSecret),
+            region: optionalString(psdRaw.region),
+            authMethod: optionalString(psdRaw.authMethod),
+            startUrl: optionalString(psdRaw.startUrl),
+        },
+        raw: data,
+        createdAt: row.createdAt,
+        updatedAt: row.updatedAt,
+    };
+}
+function safeParseObject(raw) {
+    try {
+        const parsed = JSON.parse(raw);
+        return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed)
+            ? parsed
+            : {};
+    }
+    catch {
+        return {};
+    }
+}
+function optionalString(value) {
+    return typeof value === "string" && value.trim() ? value.trim() : null;
+}

package/dist/kiro-token-store.test.js

@@ -0,0 +1,207 @@
+import assert from "node:assert/strict";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import path from "node:path";
+import test from "node:test";
+import BetterSqlite3 from "better-sqlite3";
+import { KiroTokenStore, KiroTokenStoreError } from "./kiro-token-store.js";
+function createDbFile(rows) {
+    const dir = mkdtempSync(path.join(tmpdir(), "responses-proxy-kiro-store-"));
+    const file = path.join(dir, "data.sqlite");
+    const db = new BetterSqlite3(file);
+    db.exec(`
+    CREATE TABLE providerConnections (
+      id TEXT PRIMARY KEY,
+      provider TEXT NOT NULL,
+      authType TEXT,
+      name TEXT,
+      email TEXT,
+      priority INTEGER,
+      isActive INTEGER,
+      data TEXT NOT NULL,
+      createdAt TEXT NOT NULL,
+      updatedAt TEXT NOT NULL
+    );
+  `);
+    const insert = db.prepare(`INSERT INTO providerConnections
+       (id, provider, authType, name, email, priority, isActive, data, createdAt, updatedAt)
+     VALUES (@id, @provider, @authType, @name, @email, @priority, @isActive, @data, @createdAt, @updatedAt)`);
+    for (const row of rows) {
+        insert.run({
+            id: row.id,
+            provider: row.provider ?? "kiro",
+            authType: row.authType ?? "oauth",
+            name: row.name ?? null,
+            email: row.email ?? null,
+            priority: row.priority ?? null,
+            isActive: row.isActive ?? 1,
+            data: JSON.stringify(row.data),
+            createdAt: row.createdAt ?? "2026-05-01T00:00:00.000Z",
+            updatedAt: row.updatedAt ?? "2026-05-01T00:00:00.000Z",
+        });
+    }
+    db.close();
+    return { dir, file };
+}
+test("open throws when the database file does not exist", () => {
+    const dir = mkdtempSync(path.join(tmpdir(), "responses-proxy-kiro-missing-"));
+    try {
+        assert.throws(() => KiroTokenStore.open(path.join(dir, "nope.sqlite")), (error) => error instanceof KiroTokenStoreError);
+    }
+    finally {
+        rmSync(dir, { recursive: true, force: true });
+    }
+});
+test("listAccounts maps rows, filters non-kiro and tokenless rows, and orders by priority", () => {
+    const { dir, file } = createDbFile([
+        {
+            id: "acct-b",
+            priority: 5,
+            data: {
+                accessToken: "access-b",
+                refreshToken: "refresh-b",
+                expiresAt: "2026-06-01T00:00:00.000Z",
+                expiresIn: 3600,
+                providerSpecificData: {
+                    profileArn: "arn:aws:codewhisperer:profile/b",
+                    region: "us-west-2",
+                    clientId: "client-b",
+                    clientSecret: "secret-b",
+                },
+            },
+        },
+        {
+            id: "acct-a",
+            priority: 1,
+            data: { accessToken: "access-a", refreshToken: "refresh-a" },
+        },
+        {
+            id: "no-tokens",
+            data: { providerSpecificData: { region: "us-east-1" } },
+        },
+        {
+            id: "other-provider",
+            provider: "chatgpt",
+            data: { accessToken: "x", refreshToken: "y" },
+        },
+    ]);
+    try {
+        const store = KiroTokenStore.open(file);
+        const accounts = store.listAccounts();
+        assert.deepEqual(accounts.map((a) => a.id), ["acct-a", "acct-b"]);
+        const b = accounts.find((a) => a.id === "acct-b");
+        assert.equal(b?.accessToken, "access-b");
+        assert.equal(b?.providerSpecificData.profileArn, "arn:aws:codewhisperer:profile/b");
+        assert.equal(b?.providerSpecificData.region, "us-west-2");
+        assert.equal(b?.providerSpecificData.clientId, "client-b");
+        store.close();
+    }
+    finally {
+        rmSync(dir, { recursive: true, force: true });
+    }
+});
+test("listAvailableAccounts excludes inactive accounts", () => {
+    const { dir, file } = createDbFile([
+        { id: "active", isActive: 1, data: { accessToken: "a", refreshToken: "r" } },
+        { id: "inactive", isActive: 0, data: { accessToken: "a2", refreshToken: "r2" } },
+    ]);
+    try {
+        const store = KiroTokenStore.open(file);
+        assert.deepEqual(store.listAvailableAccounts().map((a) => a.id), ["active"]);
+        store.close();
+    }
+    finally {
+        rmSync(dir, { recursive: true, force: true });
+    }
+});
+test("getAccount returns a single mapped account or undefined", () => {
+    const { dir, file } = createDbFile([
+        { id: "acct-a", data: { accessToken: "a", refreshToken: "r" } },
+    ]);
+    try {
+        const store = KiroTokenStore.open(file);
+        assert.equal(store.getAccount("acct-a")?.id, "acct-a");
+        assert.equal(store.getAccount("missing"), undefined);
+        store.close();
+    }
+    finally {
+        rmSync(dir, { recursive: true, force: true });
+    }
+});
+test("updateTokens writes back token fields, clears error markers, and preserves other keys", () => {
+    const { dir, file } = createDbFile([
+        {
+            id: "acct-a",
+            data: {
+                accessToken: "old-access",
+                refreshToken: "old-refresh",
+                expiresAt: "2026-05-01T00:00:00.000Z",
+                expiresIn: 3600,
+                lastError: "boom",
+                lastErrorAt: "2026-05-01T00:00:00.000Z",
+                providerSpecificData: { region: "us-east-1", profileArn: "arn:keep" },
+                someOtherKey: "keep-me",
+            },
+        },
+    ]);
+    try {
+        const store = KiroTokenStore.open(file);
+        const now = new Date("2026-05-15T12:00:00.000Z");
+        const updated = store.updateTokens("acct-a", {
+            accessToken: "new-access",
+            refreshToken: "new-refresh",
+            expiresAt: "2026-05-15T13:00:00.000Z",
+            expiresIn: 3600,
+        }, now);
+        assert.equal(updated?.accessToken, "new-access");
+        assert.equal(updated?.refreshToken, "new-refresh");
+        assert.equal(updated?.expiresAt, "2026-05-15T13:00:00.000Z");
+        // Re-open from a fresh handle to confirm persistence and preserved keys.
+        const verifyDb = new BetterSqlite3(file, { readonly: true });
+        const row = verifyDb
+            .prepare(`SELECT data, updatedAt FROM providerConnections WHERE id = 'acct-a'`)
+            .get();
+        const data = JSON.parse(row.data);
+        assert.equal(data.accessToken, "new-access");
+        assert.equal(data.lastError, null);
+        assert.equal(data.lastErrorAt, null);
+        assert.equal(data.someOtherKey, "keep-me");
+        assert.deepEqual(data.providerSpecificData, {
+            region: "us-east-1",
+            profileArn: "arn:keep",
+        });
+        assert.equal(row.updatedAt, now.toISOString());
+        verifyDb.close();
+        store.close();
+    }
+    finally {
+        rmSync(dir, { recursive: true, force: true });
+    }
+});
+test("updateTokens is a no-op when write-back is disabled", () => {
+    const { dir, file } = createDbFile([
+        {
+            id: "acct-a",
+            data: { accessToken: "old-access", refreshToken: "old-refresh" },
+        },
+    ]);
+    try {
+        const store = KiroTokenStore.open(file, { writeBack: false });
+        const result = store.updateTokens("acct-a", {
+            accessToken: "new-access",
+            expiresAt: "2026-05-15T13:00:00.000Z",
+        });
+        // Returns the current (unchanged) account.
+        assert.equal(result?.accessToken, "old-access");
+        const verifyDb = new BetterSqlite3(file, { readonly: true });
+        const row = verifyDb
+            .prepare(`SELECT data FROM providerConnections WHERE id = 'acct-a'`)
+            .get();
+        assert.equal(JSON.parse(row.data).accessToken, "old-access");
+        verifyDb.close();
+        store.close();
+    }
+    finally {
+        rmSync(dir, { recursive: true, force: true });
+    }
+});

package/dist/krouter-usage.js

@@ -0,0 +1,243 @@
+export class KRouterUsageLimitError extends Error {
+    statusCode = 429;
+    code = "KROUTER_TOKEN_LIMIT_REACHED";
+    usage;
+    constructor(message, usage) {
+        super(message);
+        this.name = "KRouterUsageLimitError";
+        this.usage = usage;
+    }
+}
+export async function fetchKRouterUsage({ apiKey, requestId, logger, timeoutMs, url, onEvent, }) {
+    if (!apiKey) {
+        return undefined;
+    }
+    const controller = new AbortController();
+    const timeout = setTimeout(() => controller.abort(), timeoutMs);
+    const startedAt = Date.now();
+    try {
+        const response = await fetch(url, {
+            method: "POST",
+            headers: {
+                "Content-Type": "application/json",
+            },
+            body: JSON.stringify({ apiKey }),
+            signal: controller.signal,
+        });
+        const rawText = await response.text();
+        const rawPayload = parseJsonSafely(rawText);
+        if (!response.ok) {
+            logger.warn({
+                requestId,
+                usageCheckStatus: response.status,
+                usageCheckBody: rawText,
+            }, "krouter usage check request failed");
+            await onEvent?.({
+                event: "krouter_usage_check_failed",
+                requestId,
+                usageCheckStatus: response.status,
+                usageCheckBody: rawText,
+            });
+            return undefined;
+        }
+        const usage = extractUsageSnapshot(rawPayload);
+        logger.info({
+            requestId,
+            usageCheckStatus: response.status,
+            usageCheckMs: Date.now() - startedAt,
+            usageAllowed: usage.allowed,
+            usageRemaining: usage.remaining,
+            usageLimit: usage.limit,
+            usageUsed: usage.used,
+        }, "krouter usage check completed");
+        await onEvent?.({
+            event: "krouter_usage_checked",
+            requestId,
+            usageCheckStatus: response.status,
+            usageCheckMs: Date.now() - startedAt,
+            usageAllowed: usage.allowed,
+            usageRemaining: usage.remaining,
+            usageLimit: usage.limit,
+            usageUsed: usage.used,
+        });
+        return usage;
+    }
+    catch (error) {
+        logger.warn({
+            err: error,
+            requestId,
+        }, "krouter usage check skipped after request error");
+        await onEvent?.({
+            event: "krouter_usage_check_error",
+            requestId,
+            errorMessage: error instanceof Error ? error.message : "Unknown usage check error",
+        });
+        return undefined;
+    }
+    finally {
+        clearTimeout(timeout);
+    }
+}
+export async function ensureKRouterUsageAvailable(args) {
+    const usage = await fetchKRouterUsage(args);
+    if (!usage) {
+        return undefined;
+    }
+    if (usage.allowed === false) {
+        throw new KRouterUsageLimitError("KRouter API key is not allowed to make more requests", usage);
+    }
+    if (usage.remaining !== undefined && usage.remaining <= 0) {
+        throw new KRouterUsageLimitError("KRouter token limit has been reached", usage);
+    }
+    return usage;
+}
+function extractUsageSnapshot(raw) {
+    const limit = firstFiniteNumber(raw, [
+        ["dailyTokenLimit"],
+        ["limit"],
+        ["token_limit"],
+        ["tokens_limit"],
+        ["quota"],
+        ["quota_limit"],
+        ["data", "dailyTokenLimit"],
+        ["data", "limit"],
+        ["data", "token_limit"],
+        ["data", "tokens_limit"],
+        ["data", "quota"],
+        ["result", "dailyTokenLimit"],
+        ["result", "limit"],
+        ["result", "token_limit"],
+    ]);
+    const used = firstFiniteNumber(raw, [
+        ["dailyTokensUsed"],
+        ["used"],
+        ["used_tokens"],
+        ["tokens_used"],
+        ["consumed_tokens"],
+        ["data", "dailyTokensUsed"],
+        ["data", "used"],
+        ["data", "used_tokens"],
+        ["data", "tokens_used"],
+        ["result", "dailyTokensUsed"],
+        ["result", "used"],
+        ["result", "used_tokens"],
+    ]);
+    const remaining = firstFiniteNumber(raw, [
+        ["remaining_tokens"],
+        ["remaining_token"],
+        ["tokens_remaining"],
+        ["token_remaining"],
+        ["remaining"],
+        ["quota_remaining"],
+        ["remaining_quota"],
+        ["dailyTokensRemaining"],
+        ["data", "remaining_tokens"],
+        ["data", "tokens_remaining"],
+        ["data", "remaining"],
+        ["data", "quota_remaining"],
+        ["data", "dailyTokensRemaining"],
+        ["result", "remaining_tokens"],
+        ["result", "tokens_remaining"],
+        ["result", "remaining"],
+        ["result", "dailyTokensRemaining"],
+    ]);
+    const allowed = firstBoolean(raw, [
+        ["allowed"],
+        ["isActive"],
+        ["isExpired"],
+        ["active"],
+        ["valid"],
+        ["enabled"],
+        ["can_use"],
+        ["data", "allowed"],
+        ["data", "isActive"],
+        ["data", "isExpired"],
+        ["data", "active"],
+        ["data", "valid"],
+        ["result", "allowed"],
+        ["result", "isActive"],
+        ["result", "isExpired"],
+        ["result", "active"],
+    ]);
+    const normalizedAllowed = deriveAllowed(raw, allowed);
+    const normalizedRemaining = remaining ?? (limit !== undefined && used !== undefined ? Math.max(limit - used, 0) : undefined);
+    return {
+        allowed: normalizedAllowed,
+        remaining: normalizedRemaining,
+        limit,
+        used,
+        raw,
+    };
+}
+function deriveAllowed(raw, allowed) {
+    const isExpired = firstBoolean(raw, [["isExpired"], ["data", "isExpired"], ["result", "isExpired"]]);
+    const isActive = firstBoolean(raw, [["isActive"], ["data", "isActive"], ["result", "isActive"]]);
+    if (isExpired === true) {
+        return false;
+    }
+    if (isActive !== undefined) {
+        return isActive;
+    }
+    return allowed;
+}
+function parseJsonSafely(value) {
+    if (!value.trim()) {
+        return undefined;
+    }
+    try {
+        return JSON.parse(value);
+    }
+    catch {
+        return value;
+    }
+}
+function firstFiniteNumber(payload, paths) {
+    for (const path of paths) {
+        const value = readPath(payload, path);
+        const parsed = toFiniteNumber(value);
+        if (parsed !== undefined) {
+            return parsed;
+        }
+    }
+    return undefined;
+}
+function firstBoolean(payload, paths) {
+    for (const path of paths) {
+        const value = readPath(payload, path);
+        if (typeof value === "boolean") {
+            return value;
+        }
+        if (typeof value === "string") {
+            const normalized = value.trim().toLowerCase();
+            if (normalized === "true") {
+                return true;
+            }
+            if (normalized === "false") {
+                return false;
+            }
+        }
+    }
+    return undefined;
+}
+function readPath(payload, path) {
+    let current = payload;
+    for (const key of path) {
+        if (typeof current !== "object" || current === null || Array.isArray(current)) {
+            return undefined;
+        }
+        current = current[key];
+    }
+    return current;
+}
+function toFiniteNumber(value) {
+    if (typeof value === "number" && Number.isFinite(value)) {
+        return value;
+    }
+    if (typeof value === "string" && value.trim()) {
+        const parsed = Number(value);
+        if (Number.isFinite(parsed)) {
+            return parsed;
+        }
+    }
+    return undefined;
+}