fix(imagine,krea,leonardo): poll for the full genCtx budget instead of a short hardcoded cap

imagine/krea capped polling at 4min and leonardo at 5min while the caller's genCtx budget is 8min, so slow-but-valid generations were killed early with 'generation timed out'. Now the poll deadline tracks ctx.Deadline() minus 60s download headroom (fallback to the old fixed cap if ctx has no deadline). chatgpt already ties its budget to the deadline (pollBudget) and runway/grok honor ctx, so no change needed there.
2026-07-09 16:52:16 +08:00
parent 5198b08bdf
commit fba6772584
3 changed files with 16 additions and 1 deletions
@@ -181,7 +181,12 @@ func (c *Client) GenerateImage(ctx context.Context, cookie, prompt string, width
 func (c *Client) pollImage(ctx context.Context, cookie, jobID string) (string, error) {
 	ticker := time.NewTicker(3 * time.Second)
 	defer ticker.Stop()
+	// Poll for the full generation budget (caller's genCtx), leaving headroom for
+	// the download, instead of a shorter hardcoded cap that killed slow jobs early.
 	deadline := time.Now().Add(4 * time.Minute)
+	if dl, ok := ctx.Deadline(); ok {
+		deadline = dl.Add(-60 * time.Second)
+	}

 	for {
 		body, status, err := c.apiGetP(ctx, cookie, "/api/job-status?id="+jobID, false)