From fba6772584269e20fef5fb9a62b4a0886cc5a2be Mon Sep 17 00:00:00 2001 From: chiyi Date: Thu, 9 Jul 2026 16:52:16 +0800 Subject: [PATCH] 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. --- backend/internal/provider/imagine/image.go | 5 +++++ backend/internal/provider/krea/image.go | 5 +++++ backend/internal/provider/leonardo/image.go | 7 ++++++- 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/backend/internal/provider/imagine/image.go b/backend/internal/provider/imagine/image.go index 125d616..07a126b 100644 --- a/backend/internal/provider/imagine/image.go +++ b/backend/internal/provider/imagine/image.go @@ -97,7 +97,12 @@ func (c *Client) GenerateImage(ctx context.Context, cred string, styleID int, re func (c *Client) pollImage(ctx context.Context, token, userID, batchID 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) + } url := teamsBase + "/v1/org/" + userID + "/objects?batch=true&limit=50&service=image,chat-image" for { diff --git a/backend/internal/provider/krea/image.go b/backend/internal/provider/krea/image.go index d877e90..86199ed 100644 --- a/backend/internal/provider/krea/image.go +++ b/backend/internal/provider/krea/image.go @@ -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) diff --git a/backend/internal/provider/leonardo/image.go b/backend/internal/provider/leonardo/image.go index 33b8e91..fb0d26d 100644 --- a/backend/internal/provider/leonardo/image.go +++ b/backend/internal/provider/leonardo/image.go @@ -252,8 +252,13 @@ func (c *Client) pollImage(ctx context.Context, accessToken, genID string) (stri ticker := time.NewTicker(3 * time.Second) defer ticker.Stop() - // Cap the wait independent of the parent deadline so a stuck job can't hang. + // 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. + // ctx already bounds the wait, so a stuck job still can't hang indefinitely. deadline := time.Now().Add(5 * time.Minute) + if dl, ok := ctx.Deadline(); ok { + deadline = dl.Add(-60 * time.Second) + } for { body, status, err := c.graphqlP(ctx, accessToken, payload, false)