WorkRoot-WRNexusJSComponent/app/lib/studio-ai.ts

import type {
  GenerationProvider,
  GenerationRequest,
} from "@wrnexus/design-studio";

export const APIMART_MODEL = "gpt-6-luna";
export const AI_TIMEOUT_MS = 60_000;
export function apimartProvider(
  apiKey: string,
  transport: typeof fetch = fetch,
): GenerationProvider {
  return {
    name: "APIMart",
    async generate(prompt, options = {}) {
      const response = await transport(
        "https://api.apimart.ai/v1/chat/completions",
        {
          method: "POST",
          signal: options.signal,
          redirect: "error",
          headers: {
            authorization: `Bearer ${apiKey}`,
            "content-type": "application/json",
          },
          body: JSON.stringify({
            model: APIMART_MODEL,
            stream: false,
            max_tokens: options.maxTokens ?? 4096,
            messages: [
              {
                role: "system",
                content:
                  options.system ?? "Return valid design proposal JSON only.",
              },
              { role: "user", content: prompt },
            ],
          }),
        },
      );
      // Never forward upstream bodies: they may echo credentials or submitted context.
      if (!response.ok)
        throw Error(`APIMart request failed (${response.status}).`);
      const raw = await readBounded(response, 200_000);
      const result = JSON.parse(raw);
      const choice = result.choices?.[0];
      if (choice?.finish_reason === "length")
        throw Error("APIMart output was truncated.");
      if (
        typeof choice?.message?.content !== "string" ||
        !choice.message.content.trim()
      )
        throw Error("APIMart returned no design proposal.");
      return {
        value: choice.message.content,
        provider: "APIMart",
        model: APIMART_MODEL,
      };
    },
  };
}

async function readBounded(source: Request | Response, limit: number) {
  const reader = source.body?.getReader();
  const decoder = new TextDecoder();
  let text = "",
    bytes = 0;
  if (!reader) return text;
  try {
    while (true) {
      const chunk = await reader.read();
      if (chunk.done) break;
      bytes += chunk.value.byteLength;
      if (bytes > limit) {
        await reader.cancel();
        throw Error("Request or response is too large.");
      }
      text += decoder.decode(chunk.value, { stream: true });
    }
    return text + decoder.decode();
  } finally {
    reader.releaseLock();
  }
}

export function studioAI(apiKey?: string, transport: typeof fetch = fetch) {
  let active = 0,
    windowStart = Date.now(),
    count = 0;
  const json = (value: unknown, status = 200) =>
    new Response(JSON.stringify(value), {
      status,
      headers: {
        "content-type": "application/json",
        "cache-control": "no-store",
      },
    });
  return async (request: Request): Promise<Response | undefined> => {
    const url = new URL(request.url);
    if (
      !apiKey ||
      ![
        "/design-studio/api/generation",
        "/design-studio/api/generate",
      ].includes(url.pathname)
    )
      return;
    const scheme = request.headers.get("x-forwarded-proto");
    if (scheme === "https" || scheme === "http") url.protocol = scheme + ":";
    if (
      request.headers.get("origin") &&
      request.headers.get("origin") !== url.origin
    )
      return json({ error: "Same-origin requests only." }, 403);
    if (url.pathname.endsWith("/generation") && request.method === "GET")
      return json({
        provider: "APIMart",
        model: APIMART_MODEL,
        templateVersion: 1,
        maxAttempts: 2,
        maxOutputTokens: 4096,
        timeoutMs: AI_TIMEOUT_MS,
      });
    if (!url.pathname.endsWith("/generate") || request.method !== "POST")
      return json({ error: "Method not allowed." }, 405);
    if (
      !(request.headers.get("content-type") ?? "").startsWith(
        "application/json",
      )
    )
      return json({ error: "JSON required." }, 415);
    if (Date.now() - windowStart >= 60000) {
      windowStart = Date.now();
      count = 0;
    }
    if (active >= 2 || count >= 10)
      return json(
        { error: "AI generation is busy. Try again in a minute." },
        429,
      );
    active++;
    try {
      const input = JSON.parse(
        await readBounded(request, 2_000_000),
      ) as GenerationRequest;
      count++;
      // Preserve the installed engine's validation, locks, revision checks and repair budget.
      const module = import.meta.resolve("@wrnexus/design-studio");
      const { generateDesigns } = (await import(
        module
      )) as typeof import("@wrnexus/design-studio");
      const result = await generateDesigns(input, {
        provider: apimartProvider(apiKey, transport),
        model: APIMART_MODEL,
        signal: request.signal,
        timeoutMs: AI_TIMEOUT_MS,
      });
      return json(result);
    } catch (error) {
      const message =
        error instanceof Error ? error.message : "Invalid generation request.";
      if (message.includes("WRN-STUDIO-CONFLICT"))
        return json(
          {
            error: "Your draft changed. Generate again from the current draft.",
          },
          409,
        );
      if (message.includes("too large"))
        return json({ error: "Request is too large." }, 413);
      if (
        request.signal.aborted ||
        (error instanceof Error && error.name === "AbortError")
      )
        return json(
          {
            error:
              "AI generation timed out or was cancelled. Your draft is unchanged.",
          },
          408,
        );
      return json(
        {
          error:
            "Could not generate a design. Check the brief and server AI configuration. Your draft is unchanged.",
        },
        400,
      );
    } finally {
      active--;
    }
  };
}