feat(codex): full image-generation lifecycle (started/generating/keepalive/partial/completed)

The diagnostic trace from a real image-gen run revealed five SSE event
types we were silently ignoring during the 30-50s render wait:

  +1.1s   response.image_generation_call.in_progress
  +31s    response.image_generation_call.generating
  +47s    keepalive
  +49s    response.image_generation_call.partial_image
  +49s    response.output_item.done (final image)

That's why the previous UI sat on an opaque typing indicator the whole
time. Plumb every stage through:

- New CodexImageEvent enum with .started / .generating / .keepalive /
  .partial(CodexImage) / .completed(CodexImage). CodexStreamPart now
  carries .imageEvent(CodexImageEvent) instead of a single
  .generatedImage. (Pre-1.0 breaking — only consumer is the demo,
  updated in the same change.)
- CodexImage gains an `isPartial` flag.
- Parser handles the four new event types and decodes
  partial_image_b64 (fallback to b64_json / partial_image field
  names defensively).
- CodexTool.imageGeneration now takes partialImages: Int. Default
  CodexTool.imageGenerationPNG asks for 2 preview frames. Wire JSON
  serialises as
  `{"type": "image_generation", "output_format": "png", "partial_images": 2}`.

Strip the diagnostic prints — root cause is found, lifecycle is now
plumbed cleanly. 46 kit tests still pass.

Author: atom2ueki

Sources/CodingPlanCodex/CodexStreamPart.swift

@@ -7,38 +7,67 @@
 import Foundation
 
 /// One slice of a streaming Codex response. Either a text delta to be
-/// appended to the running reply, or a fully-generated image artifact.
+/// appended to the running reply, or an image-lifecycle event from the
+/// `image_generation` tool.
 public enum CodexStreamPart: Sendable, Equatable {
     /// One streaming text delta. Concatenate every `.textDelta` to
     /// reconstruct the full assistant reply.
     case textDelta(String)
-    /// A fully-generated image emitted by the `image_generation` tool.
-    case generatedImage(CodexImage)
+    /// An image-generation lifecycle event — start, progress, partial
+    /// preview, or the final completed image. See ``CodexImageEvent``.
+    case imageEvent(CodexImageEvent)
 }
 
-/// A finished image generated mid-turn by the `image_generation` tool.
+/// Stages an image generation passes through during a streaming turn.
 ///
-/// `pngData` carries the decoded PNG bytes ready for
+/// On the wire the Codex backend emits `response.image_generation_call.*`
+/// events at different points. This enum collapses them into a small
+/// surface a chat UI can drive a placeholder + image swap from.
+public enum CodexImageEvent: Sendable, Equatable {
+    /// Backend acknowledged the image-generation call and queued it.
+    /// Show "Starting image generation…" or similar.
+    case started(callId: String)
+    /// Backend is actively rendering. Drives a shimmering placeholder.
+    case generating(callId: String)
+    /// A connection keep-alive pulse — no progress, but the server is
+    /// still working. Useful for resetting client-side stalled timers.
+    case keepalive
+    /// A low-fidelity preview of the image arrived. Swap any placeholder
+    /// for this PNG; expect a final ``completed(_:)`` shortly after.
+    case partial(CodexImage)
+    /// The final, fully-rendered image. Replace any partials with this.
+    case completed(CodexImage)
+}
+
+/// A finished (or partial) image emitted by the `image_generation` tool.
+///
+/// `pngData` carries decoded PNG bytes ready for
 /// `UIImage(data:)` / `NSImage(data:)`.
 public struct CodexImage: Sendable, Equatable {
     /// Backend identifier for this generation call (e.g. `"ig_123"`).
     public let id: String
-    /// Status reported by the backend. Typically `"completed"`.
+    /// Status reported by the backend. `"completed"` for the final image,
+    /// or a transient label like `"generating"` for partial previews.
     public let status: String
     /// The prompt the model used after revision, when the backend reports it.
     public let revisedPrompt: String?
     /// Decoded PNG bytes.
     public let pngData: Data
+    /// `true` when this is a partial preview; more frames may follow.
+    /// `false` when this is the final image.
+    public let isPartial: Bool
 
     public init(
         id: String,
         status: String,
         revisedPrompt: String? = nil,
-        pngData: Data
+        pngData: Data,
+        isPartial: Bool = false
     ) {
         self.id = id
         self.status = status
         self.revisedPrompt = revisedPrompt
         self.pngData = pngData
+        self.isPartial = isPartial
     }
 }

Sources/CodingPlanCodex/CodexTool.swift

@@ -10,23 +10,33 @@ import Foundation
 /// A tool the Codex backend can invoke as part of a response.
 public enum CodexTool: Sendable, Equatable, Hashable {
     /// Image generation. The model decides when to invoke it; the
-    /// resulting image is yielded as ``CodexStreamPart/generatedImage(_:)``.
-    /// `outputFormat` is the file format the backend should return
-    /// (currently only `"png"` is documented upstream).
-    case imageGeneration(outputFormat: String)
+    /// resulting image is yielded as ``CodexImageEvent``-wrapped
+    /// ``CodexStreamPart/imageEvent(_:)`` events.
+    ///
+    /// - Parameters:
+    ///   - outputFormat: File format (currently only `"png"` is documented).
+    ///   - partialImages: How many low-fidelity preview frames the backend
+    ///     should stream before the final image. `0` disables previews.
+    ///     The Codex backend supports `0...3`; defaults to `2` for a
+    ///     visible build-up effect.
+    case imageGeneration(outputFormat: String, partialImages: Int)
 
-    /// Convenience: PNG image generation, the upstream default.
-    public static let imageGenerationPNG: CodexTool = .imageGeneration(outputFormat: "png")
+    /// Convenience: PNG image generation with two preview frames.
+    public static let imageGenerationPNG: CodexTool = .imageGeneration(
+        outputFormat: "png",
+        partialImages: 2
+    )
 }
 
 extension CodexTool {
     /// Wire JSON object suitable for the Codex `/codex/responses` request body.
     var jsonObject: [String: Any] {
         switch self {
-        case .imageGeneration(let outputFormat):
+        case .imageGeneration(let outputFormat, let partialImages):
             return [
                 "type": "image_generation",
                 "output_format": outputFormat,
+                "partial_images": partialImages,
             ]
         }
     }

Sources/CodingPlanCodex/OpenAICodexClient.swift

@@ -321,40 +321,43 @@ public struct OpenAICodexClient: Sendable {
                     var current: [String] = []
                     var anyDeltaYielded = false
                     var pendingItemText: String?
-                    let traceImageEvents = !tools.isEmpty
-                    let startedAt = Date()
-                    var eventCount = 0
 
                     func flushEvent() throws {
                         defer { current.removeAll(keepingCapacity: true) }
                         guard let event = Self.decodeSSEEvent(dataLines: current) else { return }
-                        eventCount += 1
-                        let type = event["type"] as? String ?? "<no-type>"
-                        if traceImageEvents {
-                            let elapsed = String(format: "%.1f", Date().timeIntervalSince(startedAt))
-                            // For image_generation_call sub-events, include the
-                            // status so we can see in_progress / completed.
-                            let status = (event["item"] as? [String: Any])?["status"] as? String
-                                ?? event["status"] as? String
-                                ?? "—"
-                            print("[CodexStream +\(elapsed)s] event #\(eventCount) type=\(type) status=\(status)")
-                        }
                         switch event["type"] as? String {
                         case "response.output_text.delta":
                             if let delta = event["delta"] as? String, !delta.isEmpty {
                                 continuation.yield(.textDelta(delta))
                                 anyDeltaYielded = true
                             }
+                        case "response.output_item.added":
+                            // The backend adds an image_generation_call item when
+                            // the tool starts. Surface that as `.started`.
+                            if let item = event["item"] as? [String: Any],
+                               (item["type"] as? String) == "image_generation_call",
+                               let id = item["id"] as? String {
+                                continuation.yield(.imageEvent(.started(callId: id)))
+                            }
+                        case "response.image_generation_call.in_progress":
+                            if let id = Self.itemId(in: event) {
+                                continuation.yield(.imageEvent(.started(callId: id)))
+                            }
+                        case "response.image_generation_call.generating":
+                            if let id = Self.itemId(in: event) {
+                                continuation.yield(.imageEvent(.generating(callId: id)))
+                            }
+                        case "keepalive":
+                            continuation.yield(.imageEvent(.keepalive))
+                        case "response.image_generation_call.partial_image":
+                            if let image = Self.parsePartialImage(from: event) {
+                                continuation.yield(.imageEvent(.partial(image)))
+                            }
                         case "response.output_item.done":
                             guard let item = event["item"] as? [String: Any] else { return }
-                            // image_generation_call items carry a generated image
-                            // regardless of whether text deltas were also present.
+                            // image_generation_call items carry the final image.
                             if let image = Self.parseImageGenerationCall(from: item) {
-                                if traceImageEvents {
-                                    let elapsed = String(format: "%.1f", Date().timeIntervalSince(startedAt))
-                                    print("[CodexStream +\(elapsed)s] yielding image, pngBytes=\(image.pngData.count)")
-                                }
-                                continuation.yield(.generatedImage(image))
+                                continuation.yield(.imageEvent(.completed(image)))
                                 return
                             }
                             // Assistant message item: hold its text until response.completed
@@ -431,10 +434,40 @@ public struct OpenAICodexClient: Sendable {
             id: id,
             status: status,
             revisedPrompt: item["revised_prompt"] as? String,
-            pngData: pngData
+            pngData: pngData,
+            isPartial: false
         )
     }
 
+    /// Parse a `response.image_generation_call.partial_image` event payload.
+    /// Wire shape: `{ "type": "...partial_image", "item_id": "ig_…",
+    /// "partial_image_b64": "…", "partial_image_index": 0 }`. Falls back
+    /// to alternative field names defensively in case the upstream renames.
+    private static func parsePartialImage(from event: [String: Any]) -> CodexImage? {
+        let id = (event["item_id"] as? String) ?? (event["id"] as? String) ?? ""
+        let base64 = (event["partial_image_b64"] as? String)
+            ?? (event["b64_json"] as? String)
+            ?? (event["partial_image"] as? String)
+        guard let base64, let pngData = Data(base64Encoded: base64) else {
+            return nil
+        }
+        return CodexImage(
+            id: id,
+            status: "generating",
+            revisedPrompt: nil,
+            pngData: pngData,
+            isPartial: true
+        )
+    }
+
+    /// Pull the `item_id` from a `response.image_generation_call.*` event.
+    private static func itemId(in event: [String: Any]) -> String? {
+        if let id = event["item_id"] as? String { return id }
+        if let item = event["item"] as? [String: Any],
+           let id = item["id"] as? String { return id }
+        return nil
+    }
+
     private static func decodeSSEEvent(dataLines: [String]) -> [String: Any]? {
         guard !dataLines.isEmpty else { return nil }
         let payload = dataLines.joined(separator: "\n")