superoc 0.1.0

package/dist/index.js

@@ -0,0 +1,809 @@
+import { loadStore, saveStore, addKey, getNextKey, getActiveKeys, getDefaultStore, recordRateLimit, resetRateLimit, recordModelRateLimit, } from "./storage.js";
+import { describeError, is429Error, isStatusMessageRateLimited, shouldRetryForError, } from "./errors.js";
+import { detectProviderForRequest, getProviderHeaders } from "./provider.js";
+import { authorizeAntigravity, exchangeAntigravity, getOrRefreshAntigravityAccessToken, getAntigravityHeaders, } from "./antigravity.js";
+import { getNormalizedQuota } from "./quota.js";
+const PROVIDERS = ["nvidia", "google", "antigravity"];
+const NIM_BASE_URL = "https://integrate.api.nvidia.com";
+const GEMINI_BASE_URL = "https://generativelanguage.googleapis.com";
+const VALID_STRATEGIES = ["round-robin", "least-failures"];
+function isValidStrategy(val) {
+    return val === "round-robin" || val === "least-failures";
+}
+function modelKey(model) {
+    return `${model.providerID}/${model.modelID}`;
+}
+function getEnvKeyName(provider) {
+    if (provider === "nvidia")
+        return "NVIDIA_API_KEY";
+    if (provider === "google")
+        return "GOOGLE_API_KEY";
+    return "ANTIGRAVITY_API_KEY";
+}
+function isProviderRequest(provider, modelApiStr, providerId, modelProviderId) {
+    const detected = detectProviderForRequest({
+        provider: { info: { id: providerId } },
+        model: { providerID: modelProviderId, api: modelApiStr },
+    });
+    return detected === provider;
+}
+async function isSubagentSession(client, sessionID) {
+    try {
+        const res = await client.session.get({ path: { id: sessionID } });
+        const data = res && typeof res === "object" && "data" in res
+            ? res.data
+            : res;
+        if (!data || typeof data !== "object")
+            return false;
+        return data?.parentID !== undefined;
+    }
+    catch (err) {
+        console.debug(`[nimsuper] isSubagentSession failed for ${sessionID}:`, err);
+        return false;
+    }
+}
+const SUBAGENT_CACHE_MAX_SIZE = 1000;
+const SUBAGENT_CACHE_TTL_MS = 60_000;
+const ERROR_DEDUP_WINDOW_MS = 500;
+const SESSIONS_MAX_SIZE = 500;
+const SESSIONS_MAX_AGE_MS = 10 * 60 * 1000;
+const subAgentCache = new Map();
+async function isSubagentSessionCached(client, sessionID) {
+    const cached = subAgentCache.get(sessionID);
+    if (cached !== undefined) {
+        if (cached > Date.now())
+            return true;
+        subAgentCache.delete(sessionID);
+    }
+    const result = await isSubagentSession(client, sessionID);
+    if (result) {
+        if (subAgentCache.size >= SUBAGENT_CACHE_MAX_SIZE) {
+            const firstKey = subAgentCache.keys().next().value;
+            if (firstKey !== undefined)
+                subAgentCache.delete(firstKey);
+        }
+        subAgentCache.set(sessionID, Date.now() + SUBAGENT_CACHE_TTL_MS);
+    }
+    return result;
+}
+function findChainIndex(chain, model) {
+    if (!model)
+        return -1;
+    return chain.findIndex((entry) => entry.id === model.modelID);
+}
+export const SuperocPlugin = async (input, options) => {
+    const client = input.client;
+    const config = {
+        storePath: options?.storePath,
+        rotationStrategy: isValidStrategy(options?.rotationStrategy)
+            ? options.rotationStrategy
+            : "round-robin",
+    };
+    const store = loadStore(config) ?? getDefaultStore();
+    if (!store.fallbackChains)
+        store.fallbackChains = { nvidia: [], google: [], antigravity: [] };
+    const sessions = new Map();
+    const reloadFromDisk = () => {
+        let fresh = null;
+        try {
+            fresh = loadStore(config);
+        }
+        catch (err) {
+            console.debug("[nimsuper] Failed to reload store from disk:", err);
+            return;
+        }
+        if (fresh === null)
+            return;
+        try {
+            store.keys = fresh.keys;
+            store.currentIndex = fresh.currentIndex;
+            store.rotationStrategy = fresh.rotationStrategy;
+            store.updatedAt = fresh.updatedAt;
+            store.lastUsedKeyId = fresh.lastUsedKeyId;
+            store.fallbackChains = {
+                nvidia: Array.isArray(fresh.fallbackChains?.nvidia) ? fresh.fallbackChains.nvidia : [],
+                google: Array.isArray(fresh.fallbackChains?.google) ? fresh.fallbackChains.google : [],
+                antigravity: Array.isArray(fresh.fallbackChains?.antigravity) ? fresh.fallbackChains.antigravity : [],
+            };
+            store.maxRateLimitFailures =
+                typeof fresh.maxRateLimitFailures === "number" &&
+                    Number.isFinite(fresh.maxRateLimitFailures) &&
+                    fresh.maxRateLimitFailures >= 1
+                    ? fresh.maxRateLimitFailures
+                    : getDefaultStore().maxRateLimitFailures;
+        }
+        catch (err) {
+            console.debug("[nimsuper] Failed to apply reloaded store:", err);
+        }
+    };
+    const safeSaveStore = () => {
+        try {
+            saveStore(store, config);
+        }
+        catch (err) {
+            console.error("[nimsuper] Failed to save store:", err);
+        }
+    };
+    for (const provider of PROVIDERS) {
+        const activeKeys = getActiveKeys(store, provider);
+        if (activeKeys.length === 0) {
+            const envKey = process.env[getEnvKeyName(provider)];
+            if (envKey) {
+                const existing = store.keys.find((k) => k.name === "env-default" && k.provider === provider);
+                if (!existing) {
+                    addKey(store, "env-default", envKey, provider);
+                    safeSaveStore();
+                }
+            }
+        }
+    }
+    const showToast = async (variant, message) => {
+        try {
+            await client.tui?.showToast?.({ body: { title: "Model Fallback", message, variant } });
+        }
+        catch (err) {
+            console.debug("[nimsuper] showToast failed:", err);
+        }
+    };
+    const getState = (sessionID) => {
+        const existing = sessions.get(sessionID);
+        if (existing)
+            return existing;
+        if (sessions.size >= SESSIONS_MAX_SIZE) {
+            const now = Date.now();
+            let oldestId;
+            let oldestTime = Infinity;
+            for (const [id, s] of sessions) {
+                if (s.createdAt < oldestTime) {
+                    oldestTime = s.createdAt;
+                    oldestId = id;
+                }
+            }
+            if (oldestId)
+                sessions.delete(oldestId);
+            for (const [id, s] of sessions) {
+                if (now - s.createdAt > SESSIONS_MAX_AGE_MS) {
+                    sessions.delete(id);
+                }
+            }
+        }
+        const next = {
+            attemptIndex: 0,
+            inRetry: false,
+            aborting: false,
+            pendingRetryIndex: undefined,
+            lastUserMessageID: undefined,
+            activeChainKey: undefined,
+            activeChainModelId: undefined,
+            rateLimitCount: 0,
+            currentModelId: undefined,
+            lastFailedModelId: undefined,
+            lastErrorHandledAt: 0,
+            createdAt: Date.now(),
+            sessionProviderId: undefined,
+            lastUsedKeyId: undefined,
+        };
+        sessions.set(sessionID, next);
+        return next;
+    };
+    const cleanupSession = (sessionID) => {
+        sessions.delete(sessionID);
+    };
+    const waitForSessionIdle = async (sessionID, timeoutMs = 2000) => {
+        const start = Date.now();
+        while (Date.now() - start < timeoutMs) {
+            try {
+                const res = await client.session.status({});
+                const data = res && typeof res === "object" && "data" in res
+                    ? res.data
+                    : res;
+                if (data && typeof data === "object") {
+                    const statusMap = data;
+                    const status = statusMap[sessionID];
+                    if (status?.type === "idle")
+                        return true;
+                    if (!status)
+                        return true;
+                }
+            }
+            catch {
+                // status endpoint might not be available, keep polling
+            }
+            await new Promise((resolve) => setTimeout(resolve, 50));
+        }
+        console.debug(`[nimsuper] waitForSessionIdle timed out for ${sessionID}`);
+        return false;
+    };
+    const getChainForProvider = (provider) => store.fallbackChains[provider];
+    const triggerRetry = async (sessionID, state, reason) => {
+        const provider = state.sessionProviderId ?? "nvidia";
+        const chain = getChainForProvider(provider);
+        if (chain.length < 2)
+            return false;
+        let nextIndex = (state.attemptIndex + 1) % chain.length;
+        if (state.lastFailedModelId &&
+            chain[nextIndex]?.id === state.lastFailedModelId &&
+            chain.length > 2) {
+            nextIndex = (nextIndex + 1) % chain.length;
+        }
+        state.inRetry = true;
+        state.pendingRetryIndex = nextIndex;
+        try {
+            const source = chain[state.attemptIndex];
+            const target = chain[nextIndex];
+            if (!source || !target)
+                return false;
+            await showToast("warning", `${source.name} → ${target.name}${reason ? `: ${reason}` : ""}`);
+            const messagesResult = await client.session.messages({ path: { id: sessionID } });
+            const entries = messagesResult && "data" in messagesResult ? messagesResult.data : messagesResult;
+            if (!Array.isArray(entries))
+                return false;
+            const userMessages = entries.filter((entry) => entry?.info?.role === "user");
+            if (userMessages.length === 0)
+                return false;
+            const lastUser = userMessages[userMessages.length - 1];
+            const lastUserInfo = lastUser.info;
+            const lastUserParts = lastUser.parts;
+            if (state.lastUserMessageID &&
+                lastUserInfo?.id !== state.lastUserMessageID) {
+                return false;
+            }
+            const promptParts = [];
+            if (Array.isArray(lastUserParts)) {
+                for (const part of lastUserParts) {
+                    if (part?.type === "text") {
+                        promptParts.push({
+                            type: "text",
+                            id: part.id,
+                            text: part.text,
+                            synthetic: part.synthetic,
+                            ignored: part.ignored,
+                        });
+                    }
+                }
+            }
+            state.aborting = true;
+            try {
+                await client.session.abort({ path: { id: sessionID } });
+            }
+            catch (abortErr) {
+                console.debug(`[nimsuper] abort failed for ${sessionID}:`, abortErr);
+            }
+            const idle = await waitForSessionIdle(sessionID);
+            if (!idle) {
+                console.debug(`[nimsuper] session ${sessionID} did not go idle after abort`);
+                state.pendingRetryIndex = undefined;
+                return false;
+            }
+            await client.session.prompt({
+                path: { id: sessionID },
+                body: {
+                    messageID: lastUserInfo?.id,
+                    agent: lastUserInfo?.agent,
+                    model: {
+                        providerID: state.sessionProviderId ?? provider,
+                        modelID: target.id,
+                    },
+                    parts: promptParts,
+                },
+            });
+            return true;
+        }
+        catch (err) {
+            console.debug(`[nimsuper] triggerRetry failed for ${sessionID}:`, err);
+            state.pendingRetryIndex = undefined;
+            return false;
+        }
+        finally {
+            state.inRetry = false;
+        }
+    };
+    const handleSessionError = async (event) => {
+        const props = event.properties;
+        const error = props?.error;
+        const sessionID = props?.sessionID;
+        if (is429Error(error)) {
+            const stateForBlacklist = sessionID ? sessions.get(sessionID) : undefined;
+            const errorKeyId = stateForBlacklist?.lastUsedKeyId ?? store.lastUsedKeyId;
+            reloadFromDisk();
+            if (errorKeyId) {
+                recordRateLimit(store, errorKeyId);
+                const modelForBlacklist = stateForBlacklist?.currentModelId ?? stateForBlacklist?.activeChainModelId;
+                if (modelForBlacklist) {
+                    recordModelRateLimit(store, errorKeyId, modelForBlacklist);
+                }
+                if (stateForBlacklist) {
+                    stateForBlacklist.lastFailedModelId = modelForBlacklist;
+                }
+            }
+            safeSaveStore();
+        }
+        if (!sessionID)
+            return;
+        const state = sessions.get(sessionID);
+        if (!state)
+            return;
+        if (state.aborting) {
+            state.aborting = false;
+            return;
+        }
+        if (state.inRetry)
+            return;
+        const now = Date.now();
+        if (now - state.lastErrorHandledAt < ERROR_DEDUP_WINDOW_MS)
+            return;
+        state.lastErrorHandledAt = now;
+        if (!shouldRetryForError(error, state)) {
+            if (!is429Error(error))
+                state.rateLimitCount = 0;
+            return;
+        }
+        if (await isSubagentSessionCached(client, sessionID)) {
+            if (is429Error(error)) {
+                await showToast("warning", "Subagent rate limited — model switch skipped to preserve parent task");
+            }
+            return;
+        }
+        if (is429Error(error)) {
+            state.rateLimitCount++;
+            if (state.rateLimitCount < store.maxRateLimitFailures)
+                return;
+        }
+        else {
+            state.rateLimitCount = 0;
+            return;
+        }
+        const reason = describeError(error, state, store.maxRateLimitFailures);
+        await triggerRetry(sessionID, state, reason);
+    };
+    const handleSessionStatusRetry = async (sessionID, status) => {
+        const message = status.message;
+        const is429 = isStatusMessageRateLimited(message);
+        if (!is429)
+            return;
+        const state = sessions.get(sessionID);
+        if (!state)
+            return;
+        if (state.inRetry)
+            return;
+        if (await isSubagentSessionCached(client, sessionID))
+            return;
+        reloadFromDisk();
+        const errorKeyId = state?.lastUsedKeyId ?? store.lastUsedKeyId;
+        if (errorKeyId) {
+            recordRateLimit(store, errorKeyId);
+            const modelForBlacklist = state.currentModelId ?? state.activeChainModelId;
+            if (modelForBlacklist) {
+                recordModelRateLimit(store, errorKeyId, modelForBlacklist);
+            }
+            state.lastFailedModelId = modelForBlacklist;
+        }
+        safeSaveStore();
+        state.rateLimitCount++;
+        if (state.rateLimitCount < store.maxRateLimitFailures)
+            return;
+        const reason = `Rate limited (429) — ${state.rateLimitCount}/${store.maxRateLimitFailures} consecutive`;
+        await triggerRetry(sessionID, state, reason);
+    };
+    const handleSessionStepFailed = async (event) => {
+        const props = event.properties;
+        const sessionID = props?.sessionID;
+        if (!sessionID)
+            return;
+        const error = props?.error;
+        const errorMessage = typeof error?.message === "string" ? error.message : undefined;
+        if (!isStatusMessageRateLimited(errorMessage) && !is429Error(error))
+            return;
+        const state = sessions.get(sessionID);
+        if (!state)
+            return;
+        if (state.inRetry)
+            return;
+        const now = Date.now();
+        if (now - state.lastErrorHandledAt < ERROR_DEDUP_WINDOW_MS)
+            return;
+        state.lastErrorHandledAt = now;
+        const errorKeyId = state?.lastUsedKeyId ?? store.lastUsedKeyId;
+        reloadFromDisk();
+        if (errorKeyId) {
+            recordRateLimit(store, errorKeyId);
+            const modelForBlacklist = state.currentModelId ?? state.activeChainModelId;
+            if (modelForBlacklist) {
+                recordModelRateLimit(store, errorKeyId, modelForBlacklist);
+            }
+            state.lastFailedModelId = modelForBlacklist;
+        }
+        safeSaveStore();
+        if (await isSubagentSessionCached(client, sessionID)) {
+            await showToast("warning", "Subagent rate limited — model switch skipped to preserve parent task");
+            return;
+        }
+        state.rateLimitCount++;
+        if (state.rateLimitCount < store.maxRateLimitFailures)
+            return;
+        const reason = `Rate limited (429) — ${state.rateLimitCount}/${store.maxRateLimitFailures} consecutive`;
+        await triggerRetry(sessionID, state, reason);
+    };
+    function createSseUnwrapTransform() {
+        const decoder = new TextDecoder();
+        const encoder = new TextEncoder();
+        let buffer = "";
+        return new TransformStream({
+            transform(chunk, controller) {
+                buffer += decoder.decode(chunk, { stream: true });
+                const lines = buffer.split("\n");
+                buffer = lines.pop() || "";
+                for (const line of lines) {
+                    if (line.startsWith("data:")) {
+                        const jsonStr = line.slice(5).trim();
+                        if (!jsonStr) {
+                            controller.enqueue(encoder.encode(line + "\n"));
+                            continue;
+                        }
+                        try {
+                            const parsed = JSON.parse(jsonStr);
+                            if (parsed.response !== undefined) {
+                                controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed.response)}\n`));
+                                continue;
+                            }
+                        }
+                        catch { }
+                    }
+                    controller.enqueue(encoder.encode(line + "\n"));
+                }
+            },
+            flush(controller) {
+                if (buffer.length > 0) {
+                    if (buffer.startsWith("data:")) {
+                        const jsonStr = buffer.slice(5).trim();
+                        try {
+                            const parsed = JSON.parse(jsonStr);
+                            if (parsed.response !== undefined) {
+                                controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed.response)}\n`));
+                                return;
+                            }
+                        }
+                        catch { }
+                    }
+                    controller.enqueue(encoder.encode(buffer));
+                }
+            },
+        });
+    }
+    const hooks = {
+        auth: {
+            provider: "google",
+            loader: async (_getAuth, providerContext) => {
+                if (providerContext?.models) {
+                    try {
+                        reloadFromDisk();
+                        for (const [id, model] of Object.entries(providerContext.models)) {
+                            if (id.startsWith("antigravity-") && model) {
+                                const quota = await getNormalizedQuota(store.keys, id);
+                                const cleanBase = (model.name ?? id).replace(/\s+5h:.*$/, "");
+                                model.name = `${cleanBase} 5h: ${quota.fiveHourPercent}% W: ${quota.weeklyPercent}%`;
+                            }
+                        }
+                    }
+                    catch { }
+                }
+                return {
+                    apiKey: "",
+                    async fetch(input, init) {
+                        const urlString = typeof input === "string"
+                            ? input
+                            : input instanceof URL
+                                ? input.toString()
+                                : input.url;
+                        if (urlString.includes("generativelanguage.googleapis.com")) {
+                            const match = urlString.match(/\/models\/([^:]+):(\w+)/);
+                            const rawModel = match ? match[1] : "";
+                            const action = match ? match[2] : "streamGenerateContent";
+                            const isStreaming = action === "streamGenerateContent" || urlString.includes("alt=sse");
+                            const isAntigravityModel = rawModel.startsWith("antigravity-") ||
+                                /claude|gpt-oss|gemini-3|gemini-pro-agent/i.test(rawModel);
+                            reloadFromDisk();
+                            const activeAntigravityKeys = getActiveKeys(store, "antigravity");
+                            if (isAntigravityModel || activeAntigravityKeys.length > 0) {
+                                let attempts = 0;
+                                const maxAttempts = Math.max(1, activeAntigravityKeys.length);
+                                while (attempts < maxAttempts) {
+                                    attempts++;
+                                    const next = getNextKey(store, config, rawModel, "antigravity");
+                                    if (!next)
+                                        break;
+                                    const authRes = await getOrRefreshAntigravityAccessToken(next.key.key);
+                                    if (!authRes) {
+                                        recordRateLimit(store, next.key.id);
+                                        safeSaveStore();
+                                        continue;
+                                    }
+                                    const effectiveModel = rawModel.replace(/^antigravity-/, "");
+                                    let bodyStr = init?.body;
+                                    let parsedBody = typeof bodyStr === "string" ? JSON.parse(bodyStr) : bodyStr;
+                                    const wrappedBody = JSON.stringify({
+                                        project: authRes.projectId || "rising-fact-p41fc",
+                                        model: effectiveModel,
+                                        request: parsedBody,
+                                        requestType: "agent",
+                                        userAgent: "antigravity",
+                                    });
+                                    const headers = new Headers(init?.headers ?? {});
+                                    headers.set("Authorization", `Bearer ${authRes.accessToken}`);
+                                    headers.set("User-Agent", `Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Antigravity/1.18.3 Chrome/138.0.7204.235 Electron/37.3.1 Safari/537.36`);
+                                    headers.set("X-Goog-Api-Client", "google-cloud-sdk vscode_cloudshelleditor/0.1");
+                                    headers.set("Client-Metadata", `{"ideType":"ANTIGRAVITY","platform":"WINDOWS","pluginType":"GEMINI"}`);
+                                    headers.delete("x-goog-api-key");
+                                    headers.delete("x-api-key");
+                                    headers.delete("x-goog-user-project");
+                                    if (isStreaming)
+                                        headers.set("Accept", "text/event-stream");
+                                    const endpoints = [
+                                        "https://daily-cloudcode-pa.sandbox.googleapis.com",
+                                        "https://cloudcode-pa.googleapis.com",
+                                    ];
+                                    let gotRes = null;
+                                    for (const ep of endpoints) {
+                                        const transformedUrl = `${ep}/v1internal:${action}${isStreaming ? "?alt=sse" : ""}`;
+                                        const r = await fetch(transformedUrl, {
+                                            ...init,
+                                            headers,
+                                            body: wrappedBody,
+                                        });
+                                        if (r.ok) {
+                                            gotRes = r;
+                                            break;
+                                        }
+                                        if (r.status === 429) {
+                                            gotRes = r;
+                                        }
+                                    }
+                                    if (gotRes && gotRes.ok) {
+                                        // Update in-memory quota labels in background after successful response
+                                        if (providerContext?.models) {
+                                            setTimeout(async () => {
+                                                try {
+                                                    reloadFromDisk();
+                                                    const freshQuota = await getNormalizedQuota(store.keys, rawModel, true);
+                                                    const targetModel = providerContext.models[rawModel] ??
+                                                        providerContext.models[`antigravity-${rawModel}`];
+                                                    if (targetModel) {
+                                                        const cleanBase = (targetModel.name ?? rawModel).replace(/\s+5h:.*$/, "");
+                                                        targetModel.name = `${cleanBase} 5h: ${freshQuota.fiveHourPercent}% W: ${freshQuota.weeklyPercent}%`;
+                                                    }
+                                                }
+                                                catch { }
+                                            }, 2000);
+                                        }
+                                        if (isStreaming && gotRes.body) {
+                                            const transformedStream = gotRes.body.pipeThrough(createSseUnwrapTransform());
+                                            return new Response(transformedStream, {
+                                                status: gotRes.status,
+                                                statusText: gotRes.statusText,
+                                                headers: gotRes.headers,
+                                            });
+                                        }
+                                        return gotRes;
+                                    }
+                                    if (gotRes && gotRes.status === 429) {
+                                        recordRateLimit(store, next.key.id);
+                                        safeSaveStore();
+                                        continue;
+                                    }
+                                    if (gotRes)
+                                        return gotRes;
+                                }
+                            }
+                        }
+                        return fetch(input, init);
+                    },
+                };
+            },
+            methods: [
+                {
+                    type: "api",
+                    label: "Enter NVIDIA NIM API Key",
+                    async authorize(inputs) {
+                        const key = inputs?.["apiKey"];
+                        if (!key)
+                            return { type: "failed" };
+                        try {
+                            const res = await fetch(`${NIM_BASE_URL}/v1/models`, { headers: { Authorization: `Bearer ${key}` } });
+                            if (!res.ok)
+                                return { type: "failed" };
+                        }
+                        catch (err) {
+                            console.debug("[nimsuper] authorize fetch failed:", err);
+                            return { type: "failed" };
+                        }
+                        return { type: "success", key, provider: "nvidia" };
+                    },
+                },
+                {
+                    type: "oauth",
+                    label: "OAuth with Google (Antigravity)",
+                    async authorize() {
+                        const auth = authorizeAntigravity();
+                        return {
+                            url: auth.url,
+                            instructions: "Log in with your Google account in your browser",
+                            method: "code",
+                            async callback(code) {
+                                const res = await exchangeAntigravity(code, auth.state);
+                                if (res.type === "success") {
+                                    return {
+                                        type: "success",
+                                        refresh: res.refresh,
+                                        access: res.access,
+                                        expires: res.expires,
+                                        provider: "antigravity",
+                                    };
+                                }
+                                return { type: "failed" };
+                            },
+                        };
+                    },
+                },
+            ],
+        },
+        "chat.headers": async (_input, _output) => {
+            const provider = detectProviderForRequest(_input);
+            if (!provider)
+                return;
+            reloadFromDisk();
+            const prevKeyId = store.lastUsedKeyId;
+            const modelIdForRotation = _input.model?.id;
+            const next = getNextKey(store, config, modelIdForRotation, provider);
+            if (next) {
+                if (provider === "antigravity") {
+                    const authRes = await getOrRefreshAntigravityAccessToken(next.key.key);
+                    if (authRes) {
+                        const headers = getAntigravityHeaders(authRes.accessToken, authRes.projectId);
+                        Object.assign(_output.headers, headers);
+                    }
+                }
+                else {
+                    const headers = getProviderHeaders(provider, next.key.key);
+                    Object.assign(_output.headers, headers);
+                }
+                if (prevKeyId && prevKeyId !== next.key.id) {
+                    resetRateLimit(store, prevKeyId);
+                }
+                safeSaveStore();
+            }
+            if (modelIdForRotation && _input.sessionID) {
+                const state = getState(_input.sessionID);
+                state.currentModelId = modelIdForRotation;
+                if (next) {
+                    state.lastUsedKeyId = next.key.id;
+                }
+            }
+        },
+        "chat.message": async (input, output) => {
+            const reqModel = output.message.model ?? input.model;
+            const reqProviderId = reqModel?.providerID;
+            let matchedProvider = null;
+            for (const provider of PROVIDERS) {
+                const chain = getChainForProvider(provider);
+                if (reqModel && findChainIndex(chain, reqModel) >= 0) {
+                    matchedProvider = provider;
+                    break;
+                }
+                const isProvider = typeof reqProviderId === "string" && reqProviderId.toLowerCase().includes(provider);
+                if (isProvider) {
+                    matchedProvider = provider;
+                }
+            }
+            if (!matchedProvider)
+                return;
+            const chain = getChainForProvider(matchedProvider);
+            if (chain.length === 0)
+                return;
+            const sessionID = input.sessionID;
+            const state = getState(sessionID);
+            const requestedModel = output.message.model ?? input.model;
+            let activeChainKey = state.activeChainKey;
+            let activeChainKeyStr = activeChainKey;
+            if (!activeChainKeyStr || state.pendingRetryIndex === undefined) {
+                if (!requestedModel) {
+                    cleanupSession(sessionID);
+                    return;
+                }
+                activeChainKeyStr = modelKey(requestedModel);
+            }
+            const chainIndex = findChainIndex(chain, requestedModel);
+            if (chainIndex < 0 && state.pendingRetryIndex === undefined) {
+                cleanupSession(sessionID);
+                return;
+            }
+            let desiredIndex;
+            if (state.pendingRetryIndex !== undefined) {
+                desiredIndex = state.pendingRetryIndex;
+            }
+            else {
+                desiredIndex = chainIndex >= 0 ? chainIndex : 0;
+            }
+            const target = chain[desiredIndex];
+            if (!target) {
+                cleanupSession(sessionID);
+                return;
+            }
+            output.message.model = {
+                providerID: requestedModel?.providerID ?? matchedProvider,
+                modelID: target.id,
+            };
+            state.activeChainKey = activeChainKeyStr;
+            state.activeChainModelId = target.id;
+            state.sessionProviderId = matchedProvider;
+            state.attemptIndex = desiredIndex;
+            state.lastUserMessageID = output.message.id;
+        },
+        "shell.env": async (_input, output) => {
+            reloadFromDisk();
+            for (const provider of PROVIDERS) {
+                const envKeyName = getEnvKeyName(provider);
+                if (output.env[envKeyName] !== undefined || getActiveKeys(store, provider).length > 0) {
+                    const next = getNextKey(store, config, undefined, provider);
+                    if (next) {
+                        output.env[envKeyName] = next.key.key;
+                        safeSaveStore();
+                    }
+                }
+            }
+        },
+        event: async ({ event }) => {
+            if (event.type === "session.error") {
+                await handleSessionError(event);
+                return;
+            }
+            if (event.type === "session.next.step.failed") {
+                await handleSessionStepFailed(event);
+                return;
+            }
+            if (event.type === "session.status") {
+                const props = event.properties;
+                const sessionID = props?.sessionID;
+                const status = props?.status;
+                const statusType = status?.type;
+                if (statusType === "retry" && sessionID && status) {
+                    await handleSessionStatusRetry(sessionID, status);
+                    return;
+                }
+                if (statusType === "idle" && sessionID) {
+                    const state = sessions.get(sessionID);
+                    if (!state)
+                        return;
+                    state.rateLimitCount = 0;
+                    state.pendingRetryIndex = undefined;
+                    state.lastFailedModelId = undefined;
+                    if (state.inRetry)
+                        return;
+                    cleanupSession(sessionID);
+                    return;
+                }
+            }
+            if (event.type === "session.idle") {
+                const sessionID = event.properties?.sessionID;
+                if (!sessionID)
+                    return;
+                const state = sessions.get(sessionID);
+                if (!state)
+                    return;
+                if (state.inRetry)
+                    return;
+                state.pendingRetryIndex = undefined;
+                state.lastFailedModelId = undefined;
+                cleanupSession(sessionID);
+            }
+            if (event.type === "session.deleted") {
+                const sessionID = event.properties?.info?.id;
+                if (sessionID)
+                    cleanupSession(sessionID);
+            }
+        },
+    };
+    return hooks;
+};
+export const NimSuperPlugin = SuperocPlugin;
+export default SuperocPlugin;
+//# sourceMappingURL=index.js.map