ima2-gen 2.0.0 → 2.0.2

package/routes/nodes.js

@@ -1,6 +1,6 @@
 import { mkdir } from "fs/promises";
-import { newNodeId, saveNode, loadNodeB64, loadNodeMeta, loadAssetB64, } from "../lib/nodeStore.js";
-import { startJob, finishJob, registerJobAbortController, isJobCanceled } from "../lib/inflight.js";
+import { newNodeId, saveNode, loadNodeMeta, loadAssetB64, } from "../lib/nodeStore.js";
+import { startJob, finishJob, registerJobAbortController, isJobCanceled, isStartJobFailure, INFLIGHT_RETRY_AFTER_SECONDS } from "../lib/inflight.js";
 import { isGenerationCanceledError, makeGenerationCanceledError, throwIfJobCanceled, } from "../lib/generationCancel.js";
 import { detectImageMimeFromB64, summarizeReferencePayload, validateAndNormalizeRefs } from "../lib/refs.js";
 import { classifyUpstreamError } from "../lib/errorClassify.js";
@@ -8,76 +8,22 @@ import { normalizeOAuthParams } from "../lib/oauthNormalize.js";
 import { resolveProviderOptions } from "../lib/providerOptions.js";
 import { generateViaResponses, editViaResponses } from "../lib/responsesImageAdapter.js";
 import { generateViaGrok } from "../lib/grokImageAdapter.js";
+import { generateViaAgy } from "../lib/agyImageAdapter.js";
+import { generateViaGeminiApi } from "../lib/geminiApiImageAdapter.js";
 import { isNonRetryableGenerationError, normalizeGenerationFailure } from "../lib/generationErrors.js";
 import { logEvent, logError } from "../lib/logger.js";
 import { errInfo } from "../lib/errInfo.js";
 import { requireRuntimeContext } from "../lib/runtimeContext.js";
-function asUpstream(e) {
-    return (e && typeof e === "object" ? e : {});
-}
-function validateModeration(ctx, moderation) {
-    if (typeof moderation !== "string" || !ctx.config.oauth.validModeration.has(moderation)) {
-        return { error: "moderation must be one of: auto, low" };
-    }
-    return { moderation };
-}
-function wantsSse(req) {
-    const accept = typeof req.headers.accept === "string" ? req.headers.accept : "";
-    return accept.includes("text/event-stream");
-}
-function writeSse(res, event, data) {
-    res.write(`event: ${event}\n`);
-    res.write(`data: ${JSON.stringify(data)}\n\n`);
-}
-function writeNodeError(res, status, code, message, parentNodeId, details = {}) {
-    if (res.headersSent) {
-        writeSse(res, "error", {
-            error: { code, message },
-            parentNodeId,
-            status,
-            ...details,
-        });
-        res.end();
-        return;
-    }
-    res.status(status).json({
-        error: { code, message },
-        parentNodeId,
-        status,
-        ...details,
-    });
-}
-function dataUrlFromB64(format, b64) {
-    return `data:image/${format === "jpeg" ? "jpeg" : format};base64,${b64}`;
-}
-function imageFormatFromMime(mime) {
-    if (mime === "image/jpeg")
-        return "jpeg";
-    if (mime === "image/webp")
-        return "webp";
-    return "png";
-}
-async function loadParentNodeB64(ctx, nodeId) {
-    for (const ext of ["png", "jpeg", "webp"]) {
-        const meta = await loadNodeMeta(ctx.rootDir, nodeId, ext, ctx.config.storage.generatedDir);
-        if (meta)
-            return loadNodeB64(ctx.rootDir, `${nodeId}.${ext}`, ctx.config.storage.generatedDir);
-    }
-    return loadNodeB64(ctx.rootDir, `${nodeId}.png`, ctx.config.storage.generatedDir);
-}
-function toGrokReferences(parentB64, refs) {
-    const parentMime = parentB64 ? detectImageMimeFromB64(parentB64) : null;
-    const parentRefs = parentB64
-        ? [{ b64: parentB64, declaredMime: parentMime, detectedMime: parentMime }]
-        : [];
-    const normalizedRefs = refs.map((ref) => typeof ref === "string" ? { b64: ref } : ref);
-    return [...parentRefs, ...normalizedRefs];
-}
+import { validateModeration, imageFormatFromMime, writeSse, dataUrlFromB64 } from "../lib/routeHelpers.js";
+import { publish } from "../lib/eventBus.js";
+import { publishJobEvent } from "../lib/ssePublish.js";
+import { asUpstream, wantsSse, writeNodeError, loadParentNodeB64, toGrokReferences, nodeErrorDetails, } from "../lib/nodeHelpers.js";
 export function registerNodeRoutes(app, ctxRaw) {
     const ctx = requireRuntimeContext(ctxRaw);
     app.post("/api/node/generate", async (req, res) => {
         const body = (req.body ?? {});
-        const streamResponse = wantsSse(req);
+        const asyncMode = body.async === true;
+        const streamResponse = !asyncMode && wantsSse(req);
         const parentNodeId = (typeof body.parentNodeId === "string" ? body.parentNodeId : null);
         const requestId = typeof body.requestId === "string" ? body.requestId : (req.id ?? "");
         const sessionId = typeof body.sessionId === "string" ? body.sessionId : null;
@@ -89,21 +35,6 @@ export function registerNodeRoutes(app, ctxRaw) {
         let finishCanceled = false;
         const cancelController = new AbortController();
         const referencePayload = summarizeReferencePayload(body.references);
-        startJob({
-            requestId,
-            kind: "node",
-            prompt: body.prompt,
-            meta: {
-                kind: "node",
-                sessionId,
-                parentNodeId,
-                clientNodeId,
-                refsCount: referencePayload.refsCount,
-                referenceBytes: referencePayload.referenceBytes,
-                referenceB64Chars: referencePayload.referenceB64Chars,
-            },
-        });
-        registerJobAbortController(requestId, cancelController);
         try {
             const { prompt, quality: rawQuality = "medium", size = "1024x1024", format = "png", moderation = "low", references = [], externalSrc = null, mode: promptMode = "auto", contextMode: rawContextMode = "parent-plus-refs", searchMode: rawSearchMode = "on", model: rawModel, reasoningEffort: rawReasoningEffort, } = body;
             const { provider = "oauth" } = body;
@@ -135,7 +66,7 @@ export function registerNodeRoutes(app, ctxRaw) {
             const effectiveSize = providerOptions.size;
             const webSearchEnabled = providerOptions.webSearchEnabled;
             const activeProvider = providerOptions.provider;
-            const effectiveImageModel = activeProvider === "grok" && quality === "high"
+            const effectiveImageModel = (activeProvider === "grok" || activeProvider === "grok-api") && quality === "high"
                 ? "grok-imagine-image-quality"
                 : imageModel;
             if (contextMode === "ancestry") {
@@ -193,19 +124,47 @@ export function registerNodeRoutes(app, ctxRaw) {
             const refsForRequest = contextMode === "parent-only" ? [] : (refCheck.refDetails || refCheck.refs);
             const parentImagePresent = !!parentB64;
             const inputImageCount = (parentImagePresent ? 1 : 0) + refsForRequest.length;
-            if (activeProvider === "grok" && inputImageCount > 3) {
+            if ((activeProvider === "grok" || activeProvider === "agy" || activeProvider === "grok-api" || activeProvider === "gemini-api") && inputImageCount > 3) {
                 finishStatus = "error";
                 finishHttpStatus = 400;
-                finishErrorCode = "GROK_REF_TOO_MANY";
+                const code = activeProvider === "agy" ? "AGY_REF_TOO_MANY" : "GROK_REF_TOO_MANY";
                 return res.status(400).json({
                     error: {
-                        code: "GROK_REF_TOO_MANY",
-                        message: "Grok image editing supports up to 3 reference images.",
+                        code,
+                        message: `${activeProvider === "agy" ? "Agy" : "Grok"} image editing supports up to 3 reference images.`,
                     },
-                    code: "GROK_REF_TOO_MANY",
+                    code,
                     parentNodeId,
                 });
             }
+            const started = startJob({
+                requestId,
+                kind: "node",
+                prompt: body.prompt,
+                meta: {
+                    kind: "node",
+                    sessionId,
+                    parentNodeId,
+                    clientNodeId,
+                    refsCount: referencePayload.refsCount,
+                    referenceBytes: referencePayload.referenceBytes,
+                    referenceB64Chars: referencePayload.referenceB64Chars,
+                },
+            });
+            if (started && isStartJobFailure(started)) {
+                finishStatus = "error";
+                finishHttpStatus = started.code === "TOO_MANY_JOBS" ? 429 : 409;
+                finishErrorCode = started.code;
+                if (started.code === "TOO_MANY_JOBS") {
+                    res.setHeader("Retry-After", String(INFLIGHT_RETRY_AFTER_SECONDS));
+                }
+                return writeNodeError(res, finishHttpStatus, started.code, started.code === "TOO_MANY_JOBS"
+                    ? "Too many concurrent generation jobs"
+                    : "Request ID already in use", parentNodeId, {}, requestId);
+            }
+            registerJobAbortController(requestId, cancelController);
+            if (asyncMode)
+                res.status(202).json({ requestId });
             logEvent("node", "request", {
                 requestId,
                 operation,
@@ -229,6 +188,7 @@ export function registerNodeRoutes(app, ctxRaw) {
                 promptChars: prompt.length,
                 promptMode: normalizedPromptMode,
             });
+            const emitProgress = streamResponse || asyncMode;
             if (streamResponse) {
                 res.writeHead(200, {
                     "Content-Type": "text/event-stream; charset=utf-8",
@@ -236,9 +196,14 @@ export function registerNodeRoutes(app, ctxRaw) {
                     Connection: "keep-alive",
                 });
                 writeSse(res, "phase", { requestId, phase: "streaming" });
+                publish(requestId, "phase", { requestId, phase: "streaming" });
+            }
+            else if (asyncMode) {
+                publish(requestId, "phase", { requestId, phase: "streaming" });
             }
             let b64, usage, webSearchCalls = 0, revisedPrompt = null;
-            let resultFormat = activeProvider === "grok" ? "jpeg" : format;
+            const grokDirectApiKey = activeProvider === "grok-api" ? ctx.xaiApiKey : undefined;
+            let resultFormat = activeProvider === "grok" || activeProvider === "agy" || activeProvider === "grok-api" || activeProvider === "gemini-api" ? "jpeg" : format;
             const MAX_RETRIES = 1;
             let lastErr = null;
             for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
@@ -261,46 +226,66 @@ export function registerNodeRoutes(app, ctxRaw) {
                         searchMode,
                         webSearchEnabled,
                     });
-                    const r = activeProvider === "grok"
-                        ? await generateViaGrok(prompt, ctx, {
+                    const r = activeProvider === "gemini-api"
+                        ? await generateViaGeminiApi(parentB64 ? `Edit this image: ${prompt}` : prompt, requireRuntimeContext(ctx), {
                             model: effectiveImageModel,
                             size: effectiveSize,
-                            requestId,
                             signal: cancelController.signal,
-                            references: toGrokReferences(parentB64, refsForRequest),
+                            requestId,
+                            references: parentB64
+                                ? [{ b64: parentB64, declaredMime: null, detectedMime: null }, ...(refCheck.refDetails || [])]
+                                : refCheck.refDetails,
                         })
-                        : parentB64
-                            ? await editViaResponses(activeProvider, prompt, parentB64, quality, effectiveSize, moderation, normalizedPromptMode, ctx, requestId, {
-                                model: effectiveImageModel,
-                                references: refsForRequest,
-                                searchMode,
-                                reasoningEffort,
-                                webSearchEnabled,
+                        : activeProvider === "agy"
+                            ? await generateViaAgy(parentB64 ? `Edit this image: ${prompt}` : prompt, {
+                                references: parentB64
+                                    ? [{ b64: parentB64, declaredMime: null, detectedMime: null }]
+                                    : undefined,
                                 signal: cancelController.signal,
+                                requestId,
                             })
-                            : await generateViaResponses(activeProvider, prompt, quality, effectiveSize, moderation, refsForRequest, requestId, normalizedPromptMode, ctx, {
-                                model: effectiveImageModel,
-                                reasoningEffort,
-                                webSearchEnabled,
-                                signal: cancelController.signal,
-                                partialImages: streamResponse ? 2 : 0,
-                                onPartialImage: streamResponse
-                                    ? (partial) => isJobCanceled(requestId)
-                                        ? undefined
-                                        : writeSse(res, "partial", {
-                                            requestId,
-                                            image: dataUrlFromB64(format, partial.b64),
-                                            index: partial.index,
-                                        })
-                                    : null,
-                            });
+                            : activeProvider === "grok" || activeProvider === "grok-api"
+                                ? await generateViaGrok(prompt, ctx, {
+                                    model: effectiveImageModel,
+                                    size: effectiveSize,
+                                    requestId,
+                                    signal: cancelController.signal,
+                                    references: toGrokReferences(parentB64, refsForRequest),
+                                    directApiKey: grokDirectApiKey,
+                                })
+                                : parentB64
+                                    ? await editViaResponses(activeProvider, prompt, parentB64, quality, effectiveSize, moderation, normalizedPromptMode, ctx, requestId, {
+                                        model: effectiveImageModel,
+                                        references: refsForRequest,
+                                        searchMode,
+                                        reasoningEffort,
+                                        webSearchEnabled,
+                                        signal: cancelController.signal,
+                                    })
+                                    : await generateViaResponses(activeProvider, prompt, quality, effectiveSize, moderation, refsForRequest, requestId, normalizedPromptMode, ctx, {
+                                        model: effectiveImageModel,
+                                        reasoningEffort,
+                                        webSearchEnabled,
+                                        signal: cancelController.signal,
+                                        partialImages: emitProgress ? 2 : 0,
+                                        onPartialImage: emitProgress
+                                            ? (partial) => {
+                                                if (isJobCanceled(requestId))
+                                                    return;
+                                                const pd = { requestId, image: dataUrlFromB64(format, partial.b64), index: partial.index };
+                                                if (streamResponse)
+                                                    writeSse(res, "partial", pd);
+                                                publish(requestId, "partial", pd);
+                                            }
+                                            : null,
+                                    });
                     throwIfJobCanceled(requestId);
                     if (r.b64) {
                         b64 = r.b64;
                         usage = r.usage;
                         webSearchCalls = r.webSearchCalls || 0;
                         revisedPrompt = r.revisedPrompt || null;
-                        if (activeProvider === "grok") {
+                        if (activeProvider === "grok" || activeProvider === "grok-api" || activeProvider === "gemini-api") {
                             resultFormat = imageFormatFromMime(("mime" in r ? r.mime : undefined) || detectImageMimeFromB64(r.b64) || "image/jpeg");
                         }
                         break;
@@ -346,18 +331,7 @@ export function registerNodeRoutes(app, ctxRaw) {
                     outerHttpAlreadyCommitted: res.headersSent,
                     sseErrorSent: streamResponse,
                 });
-                return writeNodeError(res, finishHttpStatus ?? 500, finishErrorCode ?? "NODE_GEN_FAILED", finalErr.message, parentNodeId, {
-                    upstreamCode: lastErr?.upstreamCode || lastErr?.code || null,
-                    upstreamType: lastErr?.upstreamType || null,
-                    upstreamParam: lastErr?.upstreamParam || null,
-                    errorEventType: lastErr?.eventType || null,
-                    errorEventCount: lastErr?.eventCount ?? null,
-                    diagnosticReason: finalErr.diagnosticReason || lastErr?.diagnosticReason || null,
-                    retryKind: finalErr.retryKind || lastErr?.retryKind || null,
-                    referencesDroppedOnRetry: finalErr.referencesDroppedOnRetry ?? lastErr?.referencesDroppedOnRetry ?? null,
-                    refsCount: finalErr.refsCount ?? lastErr?.refsCount ?? null,
-                    inputImageCount: finalErr.inputImageCount ?? lastErr?.inputImageCount ?? null,
-                });
+                return writeNodeError(res, finishHttpStatus ?? 500, finishErrorCode ?? "NODE_GEN_FAILED", finalErr.message, parentNodeId, nodeErrorDetails(finalErr, lastErr), requestId);
             }
             const nodeId = newNodeId();
             throwIfJobCanceled(requestId);
@@ -433,7 +407,11 @@ export function registerNodeRoutes(app, ctxRaw) {
                 revisedPrompt,
                 promptMode: normalizedPromptMode,
             };
-            if (streamResponse) {
+            publishJobEvent(requestId, "done", payload);
+            if (res.writableEnded) {
+                // async mode — response already sent
+            }
+            else if (streamResponse) {
                 writeSse(res, "done", payload);
                 res.end();
             }
@@ -450,7 +428,7 @@ export function registerNodeRoutes(app, ctxRaw) {
                 finishCanceled = true;
                 finishHttpStatus = canceled.status;
                 finishErrorCode = canceled.code;
-                return writeNodeError(res, canceled.status, canceled.code, canceled.message, parentNodeId);
+                return writeNodeError(res, canceled.status, canceled.code, canceled.message, parentNodeId, {}, requestId);
             }
             finishStatus = "error";
             finishHttpStatus = err.status || 500;
@@ -460,7 +438,7 @@ export function registerNodeRoutes(app, ctxRaw) {
                 upstreamCode: ext.upstreamCode || null,
                 upstreamType: ext.upstreamType || null,
                 upstreamParam: ext.upstreamParam || null,
-            });
+            }, requestId);
         }
         finally {
             finishJob(requestId, {

package/routes/quota.js

@@ -1,4 +1,4 @@
-import { readFileSync } from "node:fs";
+import { existsSync, readFileSync } from "node:fs";
 import { homedir } from "node:os";
 import { join } from "node:path";
 function readCodexTokens() {
@@ -53,14 +53,65 @@ async function fetchCodexUsage(tokens) {
         return { provider: "codex", error: true, windows: [] };
     }
 }
+function grokTierFromLimit(val) {
+    if (val >= 150_000)
+        return "SuperGrok Heavy";
+    if (val >= 15_000)
+        return "SuperGrok";
+    return `SuperGrok (${val} val)`;
+}
+async function fetchGrokBilling() {
+    try {
+        const authPath = join(homedir(), ".progrok", "auth.json");
+        if (!existsSync(authPath))
+            return { provider: "grok", authenticated: false, windows: [] };
+        const auth = JSON.parse(readFileSync(authPath, "utf8"));
+        if (!auth.accessToken)
+            return { provider: "grok", authenticated: false, windows: [] };
+        const headers = { Authorization: `Bearer ${auth.accessToken}` };
+        const [billingRes, userRes] = await Promise.allSettled([
+            fetch("https://cli-chat-proxy.grok.com/v1/billing", { headers, signal: AbortSignal.timeout(8000) }),
+            fetch("https://cli-chat-proxy.grok.com/v1/user", { headers, signal: AbortSignal.timeout(5000) }),
+        ]);
+        if (billingRes.status !== "fulfilled" || !billingRes.value.ok) {
+            return { provider: "grok", authenticated: true, windows: [] };
+        }
+        const billing = (await billingRes.value.json()).config;
+        const limit = billing.monthlyLimit.val;
+        const used = billing.used.val;
+        let email = null;
+        if (userRes.status === "fulfilled" && userRes.value.ok) {
+            const user = await userRes.value.json();
+            email = user.email ?? null;
+        }
+        const tier = grokTierFromLimit(limit);
+        return {
+            provider: "grok",
+            account: { email, plan: tier },
+            windows: [{
+                    label: "monthly",
+                    percent: limit > 0 ? Math.round((used / limit) * 100) : 0,
+                    resetsAt: billing.billingPeriodEnd,
+                }],
+            billing: { usedUsd: used / 100, limitUsd: limit / 100 },
+        };
+    }
+    catch {
+        return { provider: "grok", error: true, windows: [] };
+    }
+}
 export function registerQuotaRoutes(app, _ctx) {
     app.get("/api/quota", async (_req, res) => {
-        const tokens = readCodexTokens();
-        if (!tokens) {
-            res.json({ codex: { provider: "codex", authenticated: false, windows: [] } });
-            return;
+        try {
+            const tokens = readCodexTokens();
+            const [codex, grok] = await Promise.all([
+                tokens ? fetchCodexUsage(tokens) : Promise.resolve({ provider: "codex", authenticated: false, windows: [] }),
+                fetchGrokBilling(),
+            ]);
+            res.json({ codex, grok });
+        }
+        catch (e) {
+            res.status(500).json({ error: "Failed to fetch quota" });
         }
-        const codex = await fetchCodexUsage(tokens);
-        res.json({ codex });
     });
 }