From 011381a801ffac5f096d3645234faf74fa49b33b Mon Sep 17 00:00:00 2001
From: shrimbly <falloon.willie@gmail.com>
Date: Sun, 19 Apr 2026 19:57:29 +1200
Subject: [PATCH] feat: return polling response for Kie tasks instead of
 blocking

The /api/generate route now submits Kie tasks and returns immediately
with polling metadata (taskId, pollProvider, etc.) instead of holding
the HTTP connection open for 2-10 minutes. Export buildMediaResponse
for reuse by the poll endpoint.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/app/api/generate/route.ts | 33 ++++++++++++++++-----------------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/src/app/api/generate/route.ts b/src/app/api/generate/route.ts
index 87d1f59e..a4b2049a 100644
--- a/src/app/api/generate/route.ts
+++ b/src/app/api/generate/route.ts
@@ -16,7 +16,7 @@ import { GenerationInput, ModelCapability } from "@/lib/providers/types";
 import { generateWithGemini, generateWithGeminiVideo } from "./providers/gemini";
 import { generateWithReplicate } from "./providers/replicate";
 import { clearFalInputMappingCache as _clearFalInputMappingCache, generateWithFalQueue } from "./providers/fal";
-import { generateWithKie } from "./providers/kie";
+import { submitKieTask } from "./providers/kie";
 import { generateWithWaveSpeed } from "./providers/wavespeed";
 
 // Re-export for backward compatibility (test file imports from route)
@@ -37,7 +37,7 @@ interface MultiProviderGenerateRequest extends GenerateRequest {
 }
 
 
-function buildMediaResponse(output: { type: string; data: string; url?: string }): NextResponse {
+export function buildMediaResponse(output: { type: string; data: string; url?: string }): NextResponse {
   if (output.type === "3d") {
     return NextResponse.json<GenerateResponse>({
       success: true,
@@ -341,28 +341,27 @@ export async function POST(request: NextRequest) {
         dynamicInputs: processedDynamicInputs,
       };
 
-      const result = await generateWithKie(requestId, kieApiKey, genInput);
-
-      if (!result.success) {
+      // Submit task and return immediately — client polls for completion
+      try {
+        const { taskId, isVeo } = await submitKieTask(requestId, kieApiKey, genInput);
+        return NextResponse.json<GenerateResponse>({
+          success: true,
+          polling: true,
+          taskId,
+          pollProvider: 'kie',
+          pollModelId: selectedModel.modelId,
+          pollModelName: selectedModel.displayName,
+          pollMediaType: mediaType || 'image',
+        });
+      } catch (error) {
         return NextResponse.json<GenerateResponse>(
           {
             success: false,
-            error: result.error || "Generation failed",
+            error: error instanceof Error ? error.message : "Task submission failed",
           },
           { status: 500 }
         );
       }
-
-      // Return first output
-      const output = result.outputs?.[0];
-      if (!output?.data && !output?.url) {
-        return NextResponse.json<GenerateResponse>(
-          { success: false, error: "No output in generation result" },
-          { status: 500 }
-        );
-      }
-
-      return buildMediaResponse(output);
     }
 
     if (provider === "wavespeed") {