feat(image): stream text-to-image generation with SSE and fallback

Send stream=true for text-to-image requests and consume the SSE
response, falling back to a non-streamed request when the stream errors
in a way that warrants it. Raises the scanner line limit to handle large
base64 image payloads.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

server/internal/infrastructure/piagent/image_generation.go

@@ -1,6 +1,7 @@
 package piagent
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/base64"
@@ -28,6 +29,7 @@ const (
 	defaultImageGenerationTimeout               = 10 * time.Minute
 	defaultImageGenerationResponseHeaderTimeout = 4 * time.Minute
 	maxImageGenerationTimeout                   = 10 * time.Minute
+	imageGenerationStreamMaxLineBytes           = 64 << 20
 )
 
 type ImageGenerator interface {
@@ -132,70 +134,95 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq
 	images := normalizeImageInputs(opts.Images)
 
 	idempotencyKey := strings.TrimSpace(opts.IdempotencyKey)
-	body, contentType, err := c.imageRequestBody(ctx, model, prompt, size, count, images)
+	streamPlans := []bool{false}
-	if err != nil {
+	if len(images) == 0 {
-		return GeneratedImage{}, err
+		streamPlans = []bool{true, false}
 	}
 
 	var lastErr error
 	startedAt := time.Now()
 	promptRunes := utf8.RuneCountInString(prompt)
-	for attempt := 0; attempt < 3; attempt++ {
+	for _, stream := range streamPlans {
-		image, err := c.generateOnce(ctx, body, contentType, idempotencyKey, len(images) > 0)
+		body, contentType, err := c.imageRequestBody(ctx, model, prompt, size, count, images, stream)
-		if err == nil {
+		if err != nil {
-			logx.Infof(
+			return GeneratedImage{}, err
-				"gpt image generation succeeded model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s timeout=%s images=%d idempotency_key=%s",
+		}
+		for attempt := 0; attempt < 3; attempt++ {
+			image, err := c.generateOnce(ctx, body, contentType, idempotencyKey, len(images) > 0, stream, count)
+			if err == nil {
+				logx.Infof(
+					"gpt image generation succeeded model=%s size=%s count=%d prompt_runes=%d attempt=%d stream=%t elapsed=%s timeout=%s images=%d idempotency_key=%s",
+					model,
+					size,
+					count,
+					promptRunes,
+					attempt+1,
+					stream,
+					time.Since(startedAt).Round(time.Millisecond),
+					c.client.Timeout,
+					len(image.URLs),
+					idempotencyKey,
+				)
+				return image, nil
+			}
+			if ctx.Err() != nil {
+				return GeneratedImage{}, ctx.Err()
+			}
+			lastErr = err
+			if stream && isImageGenerationStreamFallbackError(err) {
+				logx.Errorf(
+					"gpt image generation stream fallback model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s idempotency_key=%s err=%v",
+					model,
+					size,
+					count,
+					promptRunes,
+					attempt+1,
+					time.Since(startedAt).Round(time.Millisecond),
+					idempotencyKey,
+					err,
+				)
+				break
+			}
+			timeoutRetryable := idempotencyKey != "" && isImageGenerationTimeoutError(err)
+			retryable := isRetryableImageGenerationError(err) || timeoutRetryable
+			terminal := isImageGenerationTerminalError(err) && !timeoutRetryable
+			logx.Errorf(
+				"gpt image generation failed model=%s size=%s count=%d prompt_runes=%d attempt=%d stream=%t elapsed=%s timeout=%s retryable=%t terminal=%t idempotency_key=%s err=%v",
 				model,
 				size,
 				count,
 				promptRunes,
 				attempt+1,
+				stream,
 				time.Since(startedAt).Round(time.Millisecond),
 				c.client.Timeout,
-				len(image.URLs),
+				retryable,
+				terminal,
 				idempotencyKey,
+				err,
 			)
-			return image, nil
+			if !retryable || attempt == 2 {
+				break
+			}
+			select {
+			case <-ctx.Done():
+				return GeneratedImage{}, ctx.Err()
+			case <-time.After(time.Duration(attempt+1) * 1500 * time.Millisecond):
+			}
 		}
-		if ctx.Err() != nil {
+		if !(stream && shouldFallbackImageGenerationStream(lastErr, idempotencyKey)) {
-			return GeneratedImage{}, ctx.Err()
-		}
-		lastErr = err
-		timeoutRetryable := idempotencyKey != "" && isImageGenerationTimeoutError(err)
-		retryable := isRetryableImageGenerationError(err) || timeoutRetryable
-		terminal := isImageGenerationTerminalError(err) && !timeoutRetryable
-		logx.Errorf(
-			"gpt image generation failed model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s timeout=%s retryable=%t terminal=%t idempotency_key=%s err=%v",
-			model,
-			size,
-			count,
-			promptRunes,
-			attempt+1,
-			time.Since(startedAt).Round(time.Millisecond),
-			c.client.Timeout,
-			retryable,
-			terminal,
-			idempotencyKey,
-			err,
-		)
-		if !retryable || attempt == 2 {
 			break
 		}
-		select {
-		case <-ctx.Done():
-			return GeneratedImage{}, ctx.Err()
-		case <-time.After(time.Duration(attempt+1) * 1500 * time.Millisecond):
-		}
 	}
 	return GeneratedImage{}, lastErr
 }
 
-func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt string, size string, count int, images []string) ([]byte, string, error) {
+func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) {
 	if len(images) > 0 {
 		if c.inputImageTransport == "url" {
-			return jsonImageEditBody(model, prompt, size, count, images)
+			return jsonImageEditBody(model, prompt, size, count, images, stream)
 		}
-		return c.multipartImageEditBody(ctx, model, prompt, size, count, images)
+		return c.multipartImageEditBody(ctx, model, prompt, size, count, images, stream)
 	}
 
 	requestPayload := map[string]any{
@@ -208,11 +235,14 @@ func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt
 	if size != "" {
 		requestPayload["size"] = size
 	}
+	if stream {
+		requestPayload["stream"] = true
+	}
 	body, err := json.Marshal(requestPayload)
 	return body, "application/json", err
 }
 
-func jsonImageEditBody(model string, prompt string, size string, count int, images []string) ([]byte, string, error) {
+func jsonImageEditBody(model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) {
 	requestPayload := map[string]any{
 		"model":  model,
 		"prompt": prompt,
@@ -224,11 +254,14 @@ func jsonImageEditBody(model string, prompt string, size string, count int, imag
 	if size != "" {
 		requestPayload["size"] = size
 	}
+	if stream {
+		requestPayload["stream"] = true
+	}
 	body, err := json.Marshal(requestPayload)
 	return body, "application/json", err
 }
 
-func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string, prompt string, size string, count int, images []string) ([]byte, string, error) {
+func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) {
 	var body bytes.Buffer
 	writer := multipart.NewWriter(&body)
 	fields := map[string]string{
@@ -241,6 +274,9 @@ func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string,
 	if count > 1 {
 		fields["n"] = strconv.Itoa(count)
 	}
+	if stream {
+		fields["stream"] = "true"
+	}
 	for key, value := range fields {
 		if err := writer.WriteField(key, value); err != nil {
 			return nil, "", err
@@ -413,7 +449,7 @@ func imageExtension(contentType string, fallback string) string {
 	}
 }
 
-func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType string, idempotencyKey string, edit bool) (GeneratedImage, error) {
+func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType string, idempotencyKey string, edit bool, stream bool, expectedCount int) (GeneratedImage, error) {
 	endpoint := "/v1/images/generations"
 	if edit {
 		endpoint = "/v1/images/edits"
@@ -425,6 +461,9 @@ func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType
 	req.Header.Set("Authorization", "Bearer "+c.apiKey)
 	req.Header.Set("Content-Type", contentType)
 	req.Header.Set("Connection", "close")
+	if stream {
+		req.Header.Set("Accept", "text/event-stream")
+	}
 	if idempotencyKey != "" {
 		req.Header.Set("Idempotency-Key", idempotencyKey)
 	}
@@ -437,16 +476,21 @@ func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType
 	defer resp.Body.Close()
 
 	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		var detail struct {
+		respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
-			Error any `json:"error"`
+		if readErr != nil {
+			return GeneratedImage{}, readErr
 		}
-		_ = json.NewDecoder(resp.Body).Decode(&detail)
+		detail := imageErrorBody(respBody)
-		if detail.Error != nil {
+		if detail != "" {
-			return GeneratedImage{}, fmt.Errorf("image generation failed: %s: %v", resp.Status, detail.Error)
+			return GeneratedImage{}, fmt.Errorf("image generation failed: %s: %s", resp.Status, detail)
 		}
 		return GeneratedImage{}, fmt.Errorf("image generation failed: %s", resp.Status)
 	}
 
+	if stream {
+		return decodeImageGenerationStreamOrJSON(resp.Body, expectedCount)
+	}
+
 	var responsePayload any
 	if err := json.NewDecoder(resp.Body).Decode(&responsePayload); err != nil {
 		return GeneratedImage{}, err
@@ -459,6 +503,180 @@ func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType
 	return GeneratedImage{URL: urls[0], URLs: urls}, nil
 }
 
+type imageStreamEvent struct {
+	name      string
+	dataLines []string
+	hasData   bool
+}
+
+func decodeImageGenerationStreamOrJSON(body io.Reader, expectedCount int) (GeneratedImage, error) {
+	if expectedCount <= 0 {
+		expectedCount = 1
+	}
+	scanner := bufio.NewScanner(body)
+	scanner.Buffer(make([]byte, 0, 64*1024), imageGenerationStreamMaxLineBytes)
+	var raw strings.Builder
+	event := imageStreamEvent{name: "message"}
+	sawEvents := false
+	urls := make([]string, 0, expectedCount)
+	seen := make(map[string]struct{})
+	appendURLs := func(values []string) {
+		for _, value := range values {
+			if value == "" {
+				continue
+			}
+			if _, ok := seen[value]; ok {
+				continue
+			}
+			seen[value] = struct{}{}
+			urls = append(urls, value)
+		}
+	}
+	flush := func() (bool, error) {
+		if !event.hasData {
+			event = imageStreamEvent{name: "message"}
+			return false, nil
+		}
+		sawEvents = true
+		name := event.name
+		dataText := strings.TrimSpace(strings.Join(event.dataLines, "\n"))
+		event = imageStreamEvent{name: "message"}
+		if dataText == "" {
+			return false, nil
+		}
+		if dataText == "[DONE]" {
+			return true, nil
+		}
+		if strings.Contains(strings.ToLower(name), "partial_image") {
+			return false, nil
+		}
+		var payload any
+		if err := json.Unmarshal([]byte(dataText), &payload); err != nil {
+			return false, nil
+		}
+		if isPartialImagePayload(payload) {
+			return false, nil
+		}
+		if message := imageStreamFailureMessage(payload); message != "" {
+			return true, fmt.Errorf("image generation stream failed: %s", message)
+		}
+		appendURLs(collectGeneratedImageURLs(payload))
+		return len(urls) >= expectedCount, nil
+	}
+
+	for scanner.Scan() {
+		line := strings.TrimSuffix(scanner.Text(), "\r")
+		trimmed := strings.TrimSpace(line)
+		switch {
+		case line == "":
+			done, err := flush()
+			if err != nil {
+				return GeneratedImage{}, err
+			}
+			if done && len(urls) > 0 {
+				return GeneratedImage{URL: urls[0], URLs: urls}, nil
+			}
+		case strings.HasPrefix(line, ":"):
+			continue
+		case strings.HasPrefix(line, "event:"):
+			event.name = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
+		case strings.HasPrefix(line, "data:"):
+			event.hasData = true
+			event.dataLines = append(event.dataLines, strings.TrimSpace(strings.TrimPrefix(line, "data:")))
+		default:
+			if !sawEvents && trimmed != "" {
+				raw.WriteString(line)
+				raw.WriteByte('\n')
+			}
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return GeneratedImage{}, err
+	}
+	done, err := flush()
+	if err != nil {
+		return GeneratedImage{}, err
+	}
+	if len(urls) > 0 {
+		return GeneratedImage{URL: urls[0], URLs: urls}, nil
+	}
+	if sawEvents || done {
+		return GeneratedImage{}, fmt.Errorf("image generation stream returned no complete image")
+	}
+
+	var responsePayload any
+	if err := json.Unmarshal([]byte(raw.String()), &responsePayload); err != nil {
+		return GeneratedImage{}, err
+	}
+	urls = collectGeneratedImageURLs(responsePayload)
+	if len(urls) == 0 {
+		return GeneratedImage{}, fmt.Errorf("image generation returned empty image")
+	}
+	return GeneratedImage{URL: urls[0], URLs: urls}, nil
+}
+
+func imageErrorBody(body []byte) string {
+	body = bytes.TrimSpace(body)
+	if len(body) == 0 {
+		return ""
+	}
+	var detail struct {
+		Error any `json:"error"`
+	}
+	if err := json.Unmarshal(body, &detail); err == nil && detail.Error != nil {
+		return fmt.Sprint(detail.Error)
+	}
+	return trimErrorBody(body)
+}
+
+func isPartialImagePayload(payload any) bool {
+	item, ok := payload.(map[string]any)
+	if !ok {
+		return false
+	}
+	if value, ok := item["type"].(string); ok && strings.Contains(strings.ToLower(value), "partial_image") {
+		return true
+	}
+	return false
+}
+
+func imageStreamFailureMessage(payload any) string {
+	item, ok := payload.(map[string]any)
+	if !ok {
+		return ""
+	}
+	if message := imageErrorMessage(item["error"]); message != "" {
+		return message
+	}
+	typeValue, _ := item["type"].(string)
+	statusValue, _ := item["status"].(string)
+	if strings.Contains(strings.ToLower(typeValue), "failed") || strings.EqualFold(statusValue, "failed") {
+		if message, ok := item["message"].(string); ok && strings.TrimSpace(message) != "" {
+			return strings.TrimSpace(message)
+		}
+		return "upstream reported failed status"
+	}
+	return ""
+}
+
+func imageErrorMessage(value any) string {
+	switch typed := value.(type) {
+	case nil:
+		return ""
+	case string:
+		return strings.TrimSpace(typed)
+	case map[string]any:
+		for _, key := range []string{"message", "detail", "type", "code"} {
+			if message, ok := typed[key].(string); ok && strings.TrimSpace(message) != "" {
+				return strings.TrimSpace(message)
+			}
+		}
+		return fmt.Sprint(typed)
+	default:
+		return fmt.Sprint(typed)
+	}
+}
+
 func collectGeneratedImageURLs(payload any) []string {
 	urls := make([]string, 0)
 	seen := make(map[string]struct{})
@@ -516,13 +734,36 @@ func isGeneratedImageURLKey(key string) bool {
 
 func isGeneratedImageBase64Key(key string) bool {
 	switch key {
-	case "b64_json", "b64json", "base64", "image_base64", "imagebase64":
+	case "b64_json", "b64json", "base64", "image_base64", "imagebase64", "result":
 		return true
 	default:
 		return false
 	}
 }
 
+func isImageGenerationStreamFallbackError(err error) bool {
+	if err == nil {
+		return false
+	}
+	message := strings.ToLower(err.Error())
+	return strings.Contains(message, "unknown parameter") && strings.Contains(message, "stream") ||
+		strings.Contains(message, "unsupported") && strings.Contains(message, "stream") ||
+		strings.Contains(message, "not support") && strings.Contains(message, "stream") ||
+		strings.Contains(message, "stream returned no complete image") ||
+		strings.Contains(message, "incomplete stream") ||
+		strings.Contains(message, "text/event-stream")
+}
+
+func shouldFallbackImageGenerationStream(err error, idempotencyKey string) bool {
+	if isImageGenerationStreamFallbackError(err) {
+		return true
+	}
+	if strings.TrimSpace(idempotencyKey) == "" {
+		return false
+	}
+	return isRetryableImageGenerationError(err) || isImageGenerationTimeoutError(err)
+}
+
 func isRetryableImageGenerationError(err error) bool {
 	if err == nil {
 		return false

server/internal/infrastructure/piagent/image_generation_test.go

@@ -97,6 +97,9 @@ func TestImageClientSendsTransportDefaultSizeWhenUnset(t *testing.T) {
 	if requestBody["size"] != "1024x1024" {
 		t.Fatalf("expected transport default size, got %#v", requestBody["size"])
 	}
+	if requestBody["stream"] != true {
+		t.Fatalf("expected image generations to use stream transport, got %#v", requestBody["stream"])
+	}
 	if _, ok := requestBody["n"]; ok {
 		t.Fatalf("expected single image request to omit n for compatible transports, got %#v", requestBody["n"])
 	}
@@ -306,6 +309,112 @@ func TestImageClientParsesAlternateImageResponseShapes(t *testing.T) {
 	}
 }
 
+func TestImageClientParsesStreamingImageResponse(t *testing.T) {
+	var requestBody map[string]any
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Fatal(err)
+		}
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = w.Write([]byte("event: image_generation.completed\n"))
+		_, _ = w.Write([]byte(`data: {"type":"image_generation.completed","result":"Z2VuZXJhdGVkIGJ5dGVz"}` + "\n\n"))
+		_, _ = w.Write([]byte("data: [DONE]\n\n"))
+	}))
+	defer server.Close()
+
+	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
+	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if requestBody["stream"] != true {
+		t.Fatalf("expected stream request body, got %#v", requestBody)
+	}
+	if image.URL != "data:image/png;base64,Z2VuZXJhdGVkIGJ5dGVz" {
+		t.Fatalf("expected streamed base64 image, got %#v", image)
+	}
+}
+
+func TestImageClientFallsBackWhenStreamIsUnsupported(t *testing.T) {
+	calls := 0
+	streamValues := make([]any, 0, 2)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		calls++
+		var requestBody map[string]any
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Fatal(err)
+		}
+		streamValues = append(streamValues, requestBody["stream"])
+		if calls == 1 {
+			http.Error(w, `{"error":{"message":"unknown parameter: stream"}}`, http.StatusBadRequest)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
+	}))
+	defer server.Close()
+
+	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
+	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if image.URL != "https://example.com/image.png" {
+		t.Fatalf("expected fallback image, got %#v", image)
+	}
+	if calls != 2 {
+		t.Fatalf("expected stream then json fallback, got %d calls", calls)
+	}
+	if len(streamValues) != 2 || streamValues[0] != true || streamValues[1] != nil {
+		t.Fatalf("expected first request stream=true and fallback without stream, got %#v", streamValues)
+	}
+}
+
+func TestImageClientFallsBackAfterIdempotentStreamFailures(t *testing.T) {
+	calls := 0
+	streamValues := make([]any, 0, 4)
+	keys := make([]string, 0, 4)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		calls++
+		keys = append(keys, r.Header.Get("Idempotency-Key"))
+		var requestBody map[string]any
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Fatal(err)
+		}
+		streamValues = append(streamValues, requestBody["stream"])
+		if calls <= 3 {
+			http.Error(w, `{"error":{"message":"temporary stream failure"}}`, http.StatusServiceUnavailable)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/recovered.png"}]}`))
+	}))
+	defer server.Close()
+
+	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
+	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{
+		Model:          "gpt-image-2",
+		IdempotencyKey: "same-generation",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if image.URL != "https://example.com/recovered.png" {
+		t.Fatalf("expected fallback image, got %#v", image)
+	}
+	if calls != 4 {
+		t.Fatalf("expected three stream attempts then json fallback, got %d calls", calls)
+	}
+	if len(streamValues) != 4 || streamValues[0] != true || streamValues[1] != true || streamValues[2] != true || streamValues[3] != nil {
+		t.Fatalf("expected stream retries then json fallback, got %#v", streamValues)
+	}
+	for _, key := range keys {
+		if key != "same-generation" {
+			t.Fatalf("expected idempotency key to be preserved, got %#v", keys)
+		}
+	}
+}
+
 func TestImageClientRetriesTransientFailures(t *testing.T) {
 	calls := 0
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {