fix(studio): move AI generation behind short-lived job requests

This commit is contained in:
Clintchiz committed 2026-10-19 20:11:49 +05:30
1 parent e8c0f06d8f
commit 58687d8679
7 files changed
+248 -8

No files matched your search

+129 -1
View File
@@ -166,7 +166,7 @@ async function readBounded(source: Request | Response, limit: number) {
}
}
export function studioAI(apiKey?: string, transport: typeof fetch = fetch) {
function synchronousStudioAI(apiKey?: string, transport: typeof fetch = fetch) {
let active = 0,
windowStart = Date.now(),
count = 0;
@@ -200,6 +200,7 @@ export function studioAI(apiKey?: string, transport: typeof fetch = fetch) {
provider: "APIMart",
model: APIMART_MODEL,
templateVersion: 1,
delivery: "polling",
maxAttempts: 2,
maxOutputTokens: AI_OUTPUT_TOKENS,
timeoutMs: AI_TIMEOUT_MS,
@@ -348,3 +349,130 @@ export function studioAI(apiKey?: string, transport: typeof fetch = fetch) {
}
};
}
/** Short HTTP requests keep the deployment proxy out of the provider wait. */
export function studioAI(apiKey?: string, transport: typeof fetch = fetch) {
const run = synchronousStudioAI(apiKey, transport);
type Job = {
status: "pending" | "complete" | "failed" | "cancelled";
controller: AbortController;
expires: number;
result?: unknown;
error?: string;
};
const jobs = new Map<string, Job>();
const json = (value: unknown, status = 200) =>
Response.json(value, { status, headers: { "cache-control": "no-store" } });
return async (request: Request): Promise<Response | undefined> => {
const url = new URL(request.url);
const match =
/^\/design-studio\/api\/generate\/jobs\/([a-f0-9-]{36})$/.exec(
url.pathname,
);
const enqueue =
url.pathname === "/design-studio/api/generate" &&
request.method === "POST" &&
request.headers.get("prefer")?.includes("respond-async");
if (!apiKey || (!match && !enqueue)) return run(request);
const scheme = request.headers.get("x-forwarded-proto");
if (scheme === "http" || scheme === "https") url.protocol = scheme + ":";
if (
request.headers.get("origin") &&
request.headers.get("origin") !== url.origin
)
return json({ error: "Same-origin requests only." }, 403);
for (const [id, job] of jobs)
if (job.expires < Date.now()) {
job.controller.abort();
jobs.delete(id);
}
if (match) {
if (!["GET", "DELETE"].includes(request.method))
return json({ error: "Method not allowed." }, 405);
const job = jobs.get(match[1]!);
if (!job)
return json(
{
error:
"This generation job expired or the server restarted. Your draft is preserved; generate again.",
},
410,
);
if (request.method === "DELETE" && job.status === "pending") {
job.status = "cancelled";
job.error = "Generation cancelled. Your draft is unchanged.";
job.controller.abort();
}
return json({
status: job.status,
...(job.result ? { result: job.result } : {}),
...(job.error ? { error: job.error } : {}),
});
}
if (
!(request.headers.get("content-type") ?? "").startsWith(
"application/json",
)
)
return json({ error: "JSON required." }, 415);
if (
[...jobs.values()].filter((job) => job.status === "pending").length >= 2
)
return json({ error: "AI generation is busy. Try again shortly." }, 429);
let body: string;
try {
body = await readBounded(request, 2_000_000);
JSON.parse(body);
} catch {
return json({ error: "Invalid or oversized generation request." }, 400);
}
// Recheck after reading the body so concurrent uploads cannot bypass the limit.
if (
[...jobs.values()].filter((job) => job.status === "pending").length >= 2
)
return json({ error: "AI generation is busy. Try again shortly." }, 429);
if (jobs.size >= 20) {
const old = [...jobs].find(([, job]) => job.status !== "pending");
if (old) jobs.delete(old[0]);
}
const id = crypto.randomUUID();
const job: Job = {
status: "pending",
controller: new AbortController(),
expires: Date.now() + 5 * 60_000,
};
jobs.set(id, job);
const headers = new Headers(request.headers);
headers.delete("prefer");
const task = new Request(request.url, {
method: "POST",
headers,
body,
signal: job.controller.signal,
});
void run(task)
.then(async (response) => {
if (job.status !== "pending") return;
if (!response) throw Error("Generation is not configured.");
const result = await response.json();
if (response.ok) {
job.status = "complete";
job.result = result;
} else {
job.status = "failed";
job.error =
typeof result.error === "string"
? result.error
: "Generation did not complete.";
}
})
.catch(() => {
if (job.status === "pending") {
job.status = "failed";
job.error =
"Generation could not complete. Your draft is preserved; try again.";
}
});
return json({ job: { id, status: "pending" } }, 202);
};
}