diff --git a/.github/pr-assets/1002/01-tool-flow.png b/.github/pr-assets/1002/01-tool-flow.png new file mode 100644 index 00000000..617af316 Binary files /dev/null and b/.github/pr-assets/1002/01-tool-flow.png differ diff --git a/.github/pr-assets/1002/02-image-expanded.png b/.github/pr-assets/1002/02-image-expanded.png new file mode 100644 index 00000000..b87d5e51 Binary files /dev/null and b/.github/pr-assets/1002/02-image-expanded.png differ diff --git a/.github/pr-assets/1002/03-image-model-dropdown.png b/.github/pr-assets/1002/03-image-model-dropdown.png new file mode 100644 index 00000000..91f83a32 Binary files /dev/null and b/.github/pr-assets/1002/03-image-model-dropdown.png differ diff --git a/.github/pr-assets/1002/index.html b/.github/pr-assets/1002/index.html new file mode 100644 index 00000000..8c8925c3 --- /dev/null +++ b/.github/pr-assets/1002/index.html @@ -0,0 +1,243 @@ + + + + +PR #1002 — Native image_generation UX trace + + + + +
+
+

PR #1002 · Native image_generation — UX trace

+
+ Captured against a live backend running the PR binary · agent chatgpt-image-gen (provider cliproxy-codex, model gpt-5.4) · image model gpt-image-2 · real ChatGPT OAuth · real Postgres · session 23cb257e. +
+
+ +
+ +
+ +
+
+

What this shows

+

Real end-to-end run of the unified create_image → NativeImageProvider → Codex pathway against a complex Vietnamese infographic prompt. Tool completes with Done, image renders inline with the prompt as caption, and the image model is user-configurable from the existing Chain dialog.

+
+
+

Why it matters

+

Before this PR, routing create_image to openai-codex failed with provider "openai-codex" does not expose API credentials required for image generation. The new NativeImageProvider interface bridges OAuth-backed providers without exposing static keys — and locks in gpt-image-2 as the quality default.

+
+
+

Review cue

+

The generated PNG is intentionally compact (thumbnail + expanded view) rather than a raw file attachment, so the PR itself stays small. The captures are here to show the surface, not to ship the asset.

+
+
+ +
+
+

1 · Inline result · create_image returns Done, image + prompt caption render

+ Tail of the Vietnamese Red Fox encyclopedia prompt, then the tool-result row, then the assistant reply with image and caption. +
+
+
+ Implemented + / chat / chatgpt-image-gen +
+ Chat view showing the end of the user's Vietnamese infographic prompt, then create_image Done, then the assistant's Vietnamese reply with the generated infographic rendered inline and the prompt as an italic muted caption beneath it. +
+ What this demonstrates + The user's infographic prompt flows through to a create_image call that completes with Done. The assistant acknowledges in Vietnamese and the generated image renders inline via MediaGallery. Beneath the image, in muted italic, is the prompt caption — one of the two new UX surfaces this PR adds. The prompt is also embedded into the PNG's tEXt chunk on write, so downloaded files carry their own provenance. + Backend trace + Log sequence: tool call create_image args_len=N → (4–8 min of work upstream) → create_image: file saved path=/app/workspace/…/generated/…/cao-do-infographic-vietnamese_….png size=… → v3.run.completed. +
+
+
+ +
+
+

2 · Expanded view · click the image, MediaGallery lightbox

+ Clicking the inline image opens the existing MediaGallery lightbox. Download button top-right. Full-resolution view confirms the generated asset is what reached the client — no placeholder, no degraded render. +
+
+
+ Implemented + MediaGallery · lightbox overlay +
+ Lightbox overlay showing the full vertical Red Fox encyclopedia infographic in Vietnamese, with download button top-right, image filename visible bottom-center +
+ What this proves + The full 1024×1792 PNG, streamed out of the native Codex Responses API, persisted to the workspace, surfaced through the existing MediaGallery render path with no new plumbing — that's the value of routing through the pre-existing create_image tool rather than inventing a new rail. + Download UX + Filename is cao-do-infographic-vietnamese_YYYYMMDD-HHmmss_hash.png — resolved by the tool's filename_hint arg + timestamp, not a random UUID. Discoverable in the assistant's workspace at {workspace}/media/{sha256}.{ext} (deduped on hash) and also under the tool's generated/YYYY-MM-DD/ folder. +
+
+
+ +
+
+

3 · Where to configure the image model

+ Whitelist select in the existing "Create Image — Provider Chain" dialog. Default gpt-image-2, legacy gpt-image-1.5, nothing else. +
+
+
+ Configurable + Built-in Tools → create_image → Provider Chain +
+ Create Image Provider Chain modal with openai-codex row expanded. Model GPT-5.4, Timeout 600s, Retries 1. Settings panel open. Image model dropdown expanded showing Default · gpt-image-2 (recommended) selected, and Legacy · gpt-image-1.5 as the second option. +
+ Navigation path +
+ Sidebar › + Built-in Tools › + create_image row › + Settings › + Create Image — Provider Chain › + openai-codex card › + Settings (expand) › + Image model +
+ Options (whitelist — enforced server-side) + + + + + + + +
LabelValueWhen to pick
Default · gpt-image-2gpt-image-2Quality baseline. The motivation of this PR. Recommended for everyone.
Legacy · gpt-image-1.5gpt-image-1.5Only if a pool account lacks gpt-image-2 entitlement or you're cost-tuning.
Anything else—Rejected by ValidateImageModel with unsupported image model "…"; allowed: gpt-image-2 (default), gpt-image-1.5 (legacy). Prevents silent upstream 400s.
+ Where it's stored · how it's threaded + Per chain entry, under params.image_model in the create_image tool settings JSON. At runtime: create_image.callProvider reads entry.Params["image_model"] → NativeImageRequest.ImageModel → ValidateImageModel (defaults empty to gpt-image-2) → outbound tools[0].model on POST /codex/responses. + Why surface this at all + Most operators never need to touch it. It exists so that when Codex eventually rotates image models, or when an account's entitlement differs, the fallback is a clean UI toggle rather than a code change. The whitelist keeps the selector honest — no arbitrary strings, no silent upstream rejections. + Also visible in this screenshot + Timeout: 600s and Retries: 1 — the new defaults for this chain entry. Image generation of complex prompts legitimately runs 4–8 minutes; the old default of 120s × 2 retries routinely timed out mid-flight with context deadline exceeded. See §4. +
+
+
+ +
+
+

4 · What changed — honest summary

+ Before/after, no hand-waving. +
+
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
AreaBeforeAfter
Routing create_image → openai-codexFailed. credentialProvider required static APIKey / APIBase; OAuth providers don't satisfy it.Works. New NativeImageProvider.GenerateImage. CodexProvider implements it via POST /codex/responses with the native image_generation tool.
Responses API wire format—stream:true (API rejects false), instructions populated (API rejects missing), tool_choice forces image_generation. SSE stream parsed for response.output_item.done image items and response.completed output walk.
Image model selectionHardcoded literal.Whitelisted: gpt-image-2 (default) + gpt-image-1.5 (legacy). Selector in the Chain dialog. Server validator rejects anything else.
Default chain timeout120s × 2 retries · image gen routinely died with context deadline exceeded while upstream was still generating.600s × 1 retry · matches realistic gpt-image-2 completion time. Retries reduced to 1 — stateful upstream runs don't benefit from retry.
Assistant images in UIRendered inline, no provenance.Prompt caption beneath image (muted italic, line-clamp-2, full text in tooltip). Prompt also embedded in PNG tEXt chunk so downloaded files carry provenance.
Per-request user toggleAdded in earlier commits of this PR branch.Removed. Users toggling it off then forgetting = support footgun. Emergency admin kill-switch still exists via AgentConfig.AllowImageGeneration (stored in other_config).
CI drift (unrelated — fixed in-PR)sessions.compact unclassified → RBAC drift test fail. contains() declared twice in tests/integration → compile fail.Classified, deduped. Green.
+
+
+ +
+
+

Out of scope / follow-ups

+ Honest gap log. +
+
+
    +
  • Desktop (Wails) surface — UI changes live only in ui/web/. Desktop shell unchanged.
    • +
  • Video / audio generation chains — only create_image is routed through NativeImageProvider; create_video / create_audio still use the credentialProvider path.
    • +
  • OpenAI-compat track (non-Codex providers sending message.images[]) is wired but untested against a live OpenAI-compat image endpoint — forward-compat infrastructure only.
    • +
+
+
+ +
+ + + + diff --git a/docs/codebase-summary.md b/docs/codebase-summary.md index b773b72a..01cf2c61 100644 --- a/docs/codebase-summary.md +++ b/docs/codebase-summary.md @@ -119,6 +119,27 @@ Parity enforced by `ui/web/src/__tests__/i18n-tts-key-parity.test.ts` (vitest). --- +## Image Generation + +Native `image_generation` support in the Codex provider (`POST /codex/responses`) + passthrough in the OpenAI-compat path. + +**Provider flag:** `ProviderCapabilities.ImageGeneration bool` (`internal/providers/capabilities.go`). Codex sets `true`; other providers default `false`. + +**Gate (agent loop):** `ToolDefinition{Type:"image_generation"}` appended iff (provider capability) AND (`AgentConfig.AllowImageGeneration`, default true) AND (request lacks `x-goclaw-no-image-gen` header). Gate logic in `internal/agent/loop_tool_filter.go`. + +**Codex native events** (`internal/providers/codex.go`): +- `response.image_generation_call.partial_image` → `ChatResponse.Images` entry with `Partial:true`. +- `response.output_item.done` with `item.type == "image_generation_call"` → final `ChatResponse.Images` entry; partial frames for same `item_id` replaced. +- `response.completed` walks `response.output[]` for image items (non-stream). + +**OpenAI-compat parsing:** `choices[0].message.images[]` + `choices[0].delta.images[]` with `data:image/...;base64,...` URLs decoded in `internal/providers/openai_http.go` and `internal/providers/openai_chat.go`. Helper: `parseDataURL()` in `internal/providers/openai_image_url.go`. + +**Persistence:** `internal/agent/media.go persistAssistantImages()` writes final images to `{workspace}/media/{sha256}.{ext}`, returns `MediaRef` entries, clears inline `Images[]`. Idempotent on hash. Invoked from `pipeline.FinalizeStage` via `Deps.PersistAssistantImages` callback. + +**Web UI:** Download filename resolver (`imageGenDownloadName`) in `ui/web/src/components/chat/media-gallery.tsx`. Image generation works automatically when the agent has the `create_image` tool — no user-facing toggle. + +--- + ## Key Conventions - **Store layer:** Interface-based; PG (`store/pg/`) + SQLite (`store/sqlitestore/`). Raw SQL, `$1/$2` params. diff --git a/docs/project-changelog.md b/docs/project-changelog.md index 06f9faf0..388bc30e 100644 --- a/docs/project-changelog.md +++ b/docs/project-changelog.md @@ -4,6 +4,37 @@ Significant changes, features, and fixes in reverse chronological order. --- +## 2026-04-22 + +### Providers: Native image generation for Codex + OpenAI-compat + +**Features** + +- **Codex native track:** `CodexProvider` now attaches the `image_generation` tool object to `POST /codex/responses` when the agent permits it. Streams `response.image_generation_call.partial_image` intermediate frames + `response.output_item.done` (type `image_generation_call`) final images; non-stream path walks `response.output[]`. Deduped per `item_id`, partial frames emitted as `ImageContent{Partial:true}` for UI progressive render. +- **OpenAI-compat track:** `tools[]` serializer passes `{type:"image_generation"}` entries through natively; response parser reads `choices[0].message.images[]` / `choices[0].delta.images[]` (data URLs) into `ChatResponse.Images`. +- **Media persistence:** `internal/agent/media.go` `persistAssistantImages()` writes final images to `{workspace}/media/{sha256}.{ext}`, returns `MediaRef` entries, clears inline base64. Idempotent on hash. Wired via `pipeline.Deps.PersistAssistantImages` callback from `FinalizeStage`. Partial frames skipped. +- **Capabilities + gate:** `ProviderCapabilities.ImageGeneration` flag, set true on Codex provider. Tri-level gate in agent loop: provider capability AND `AgentConfig.AllowImageGeneration` (read from `other_config.allow_image_generation`, default true) AND request not opted-out via `x-goclaw-no-image-gen` header. +- **Web UI:** Composer "Images" toggle chip (visible only when provider supports image gen, per-agent persistence in localStorage). Streaming placeholder skeleton in `ActiveRunZone` while partials arrive. `MediaGallery` assigns `generated-{timestamp}.png` filename for assistant-generated PNGs. + +**Wire format** + +Implementation is evidence-backed against the native ChatGPT Responses API event shape, not the compat shim shape. Research notes in `plans/reports/`. + +**i18n** + +- 1 UI key (`imageGenDownloadName`) in `ui/web/src/i18n/locales/{en,vi,zh}/chat.json` — download filename for generated images. + +**Tests** + +- Unit tests across providers (Codex native + OpenAI-compat), agent media persistence, store config. Full test sweep: 2618 pass. + +**Internal docs** + +- `plans/260422-1349-goclaw-chatgpt-image-gen/` — plan + phase files. +- `plans/reports/researcher-260422-1414-codex-native-image-events.md` — native event schema. + +--- + ## 2026-04-19 ### TTS: Gemini provider + ProviderCapabilities schema engine diff --git a/internal/agent/image_gen_gate_test.go b/internal/agent/image_gen_gate_test.go new file mode 100644 index 00000000..ac4daeed --- /dev/null +++ b/internal/agent/image_gen_gate_test.go @@ -0,0 +1,117 @@ +package agent + +// Tests for the two-tier image_generation gate in buildFilteredTools. +// +// Gate conditions (ALL must be true to inject the native tool): +// (1) provider implements CapabilitiesAware and Capabilities().ImageGeneration == true +// (2) Loop.allowImageGeneration == true (agent config, defaults true; admin-only control) +// +// Additionally: final-iteration stripping takes priority — all tools removed. + +import ( + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +// imageCapableProvider is a stub provider that also implements CapabilitiesAware +// and can toggle ImageGeneration on/off. +type imageCapableProvider struct { + stubProvider + imageGen bool +} + +func (p *imageCapableProvider) Capabilities() providers.ProviderCapabilities { + return providers.ProviderCapabilities{ + Streaming: true, + ToolCalling: true, + ImageGeneration: p.imageGen, + } +} + +// buildImageGenLoop constructs a minimal Loop for gate testing. +// Uses the stubExecutor already defined in loop_pipeline_tool_callbacks_test.go. +func buildImageGenLoop(allowImageGen bool, prov providers.Provider) *Loop { + return &Loop{ + provider: prov, + allowImageGeneration: allowImageGen, + tools: &stubExecutor{}, + } +} + +// hasImageGenTool returns true if the slice contains the image_generation sentinel. +func hasImageGenTool(defs []providers.ToolDefinition) bool { + for _, d := range defs { + if d.Type == "image_generation" { + return true + } + } + return false +} + +// ─── Gate: all conditions true → tool present ───────────────────────────── + +func TestImageGenGate_AllTrue_ToolPresent(t *testing.T) { + prov := &imageCapableProvider{imageGen: true} + l := buildImageGenLoop(true, prov) + + defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil) + + if !hasImageGenTool(defs) { + t.Error("expected image_generation tool when all gate conditions are true") + } +} + +// ─── Gate: provider capability false → tool absent ──────────────────────── + +func TestImageGenGate_ProviderNoCapability_ToolAbsent(t *testing.T) { + prov := &imageCapableProvider{imageGen: false} + l := buildImageGenLoop(true, prov) + + defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil) + + if hasImageGenTool(defs) { + t.Error("image_generation must NOT be in tools when provider does not advertise ImageGeneration") + } +} + +// ─── Gate: provider not CapabilitiesAware → tool absent ────────────────── + +func TestImageGenGate_ProviderNotCapabilitiesAware_ToolAbsent(t *testing.T) { + // stubProvider (from intent_classify_test.go) does NOT implement CapabilitiesAware. + prov := &stubProvider{} + l := buildImageGenLoop(true, prov) + + defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil) + + if hasImageGenTool(defs) { + t.Error("image_generation must NOT be in tools when provider is not CapabilitiesAware") + } +} + +// ─── Gate: agent config disables → tool absent ─────────────────────────── + +func TestImageGenGate_AgentConfigDisabled_ToolAbsent(t *testing.T) { + prov := &imageCapableProvider{imageGen: true} + l := buildImageGenLoop(false, prov) // allowImageGeneration = false + + defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil) + + if hasImageGenTool(defs) { + t.Error("image_generation must NOT be in tools when agent config disables it") + } +} + +// ─── Final iteration strips all tools including image_generation ────────── + +func TestImageGenGate_FinalIteration_AllToolsStripped(t *testing.T) { + prov := &imageCapableProvider{imageGen: true} + l := buildImageGenLoop(true, prov) + + // iteration == maxIter → final stripping path; gate never reached + defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 5, 5, nil) + + if len(defs) != 0 { + t.Errorf("final iteration must strip all tools; got %d: %v", len(defs), defs) + } +} diff --git a/internal/agent/loop_pipeline_adapter.go b/internal/agent/loop_pipeline_adapter.go index 92f0fe6f..f0cb3e67 100644 --- a/internal/agent/loop_pipeline_adapter.go +++ b/internal/agent/loop_pipeline_adapter.go @@ -160,6 +160,7 @@ func (l *Loop) buildPipelineDeps(req *RunRequest, bridgeRS *runState) pipeline.P // Checkpoint + Finalize FlushMessages: cb.flushMessages, + PersistAssistantImages: persistAssistantImages, SkillPostscript: l.makeSkillPostscript(), SanitizeContent: cb.sanitizeContent, StripMessageDirectives: StripMessageDirectives, @@ -245,6 +246,7 @@ func convertRunResult(pr *pipeline.RunResult) *RunResult { ContentType: m.ContentType, Size: m.Size, AsVoice: m.AsVoice, + Prompt: m.Prompt, } } return &RunResult{ diff --git a/internal/agent/loop_pipeline_tool_callbacks.go b/internal/agent/loop_pipeline_tool_callbacks.go index 51d49638..b1f056ac 100644 --- a/internal/agent/loop_pipeline_tool_callbacks.go +++ b/internal/agent/loop_pipeline_tool_callbacks.go @@ -181,6 +181,7 @@ func syncBridgeToState(bridgeRS *runState, state *pipeline.RunState, action tool ContentType: mr.ContentType, Size: mr.Size, AsVoice: mr.AsVoice, + Prompt: mr.Prompt, }) } } diff --git a/internal/agent/loop_tool_filter.go b/internal/agent/loop_tool_filter.go index 758886d2..f633a90b 100644 --- a/internal/agent/loop_tool_filter.go +++ b/internal/agent/loop_tool_filter.go @@ -8,6 +8,11 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/tools" ) +// imageGenToolDef is the native image_generation tool sentinel. Its Type-only form +// is passed through by the Codex/OpenAI request builder as a bare {"type":"image_generation"} +// object — no "function" wrapper, no parameters. +var imageGenToolDef = providers.ToolDefinition{Type: "image_generation"} + // buildFilteredTools resolves the per-iteration tool definitions based on policy, // disabled tools, bootstrap mode, skill visibility, channel type, and iteration budget. // Per-user MCP tools must be registered in the Registry before calling this function @@ -103,6 +108,19 @@ func (l *Loop) buildFilteredTools(req *RunRequest, hadBootstrap bool, iteration, Role: "user", Content: "[System] Final iteration reached. Summarize all findings and respond to the user now. No more tool calls allowed.", }) + return toolDefs, allowedTools, messages + } + + // Two-tier image generation gate: + // (1) provider supports native image_generation (ImageGeneration capability) + // (2) agent config allows it (allowImageGeneration — defaults true, set false via + // other_config.allow_image_generation = false in the admin agent configuration) + if l.allowImageGeneration { + if aware, ok := l.provider.(providers.CapabilitiesAware); ok { + if aware.Capabilities().ImageGeneration { + toolDefs = append(toolDefs, imageGenToolDef) + } + } } return toolDefs, allowedTools, messages diff --git a/internal/agent/loop_tools.go b/internal/agent/loop_tools.go index 7a08bd62..af825ac2 100644 --- a/internal/agent/loop_tools.go +++ b/internal/agent/loop_tools.go @@ -86,12 +86,16 @@ func (l *Loop) processToolResult( // Collect MEDIA: paths from tool results. // Prefer result.Media (explicit) over ForLLM MEDIA: prefix (legacy) to avoid duplicates. if len(result.Media) > 0 { - for _, mf := range result.Media { + for i, mf := range result.Media { ct := mf.MimeType if ct == "" { ct = mimeFromExt(filepath.Ext(mf.Path)) } - rs.mediaResults = append(rs.mediaResults, MediaResult{Path: mf.Path, ContentType: ct}) + mr := MediaResult{Path: mf.Path, ContentType: ct} + if result.MediaPrompts != nil { + mr.Prompt = result.MediaPrompts[i] + } + rs.mediaResults = append(rs.mediaResults, mr) } } else if mr := parseMediaResult(result.ForLLM); mr != nil { rs.mediaResults = append(rs.mediaResults, *mr) diff --git a/internal/agent/loop_types.go b/internal/agent/loop_types.go index 59a5dba8..79a934d3 100644 --- a/internal/agent/loop_types.go +++ b/internal/agent/loop_types.go @@ -200,6 +200,11 @@ type Loop struct { // Self-evolve: predefined agents can update SOUL.md through chat selfEvolve bool + // allowImageGeneration: gate for native image_generation tool injection. + // Tri-level: provider supports it AND this flag is true AND request hasn't opted out. + // Defaults to true; set false via other_config.allow_image_generation = false. + allowImageGeneration bool + // TTS auto mode from config: "off", "always", "inbound", "tagged" ttsAutoMode string @@ -392,6 +397,10 @@ type LoopConfig struct { // Self-evolve: predefined agents can update SOUL.md (style/tone) through chat SelfEvolve bool + // AllowImageGeneration: whether the native image_generation tool may be attached. + // Defaults to true; set false to disable image generation for this agent. + AllowImageGeneration bool + // TTS auto mode from config: "off", "always", "inbound", "tagged" // When "tagged", inject [[tts]] directive guidance into system prompt. TTSAutoMode string @@ -546,6 +555,7 @@ func NewLoop(cfg LoopConfig) *Loop { promptMode: cfg.PromptMode, pinnedSkills: cfg.PinnedSkills, selfEvolve: cfg.SelfEvolve, + allowImageGeneration: cfg.AllowImageGeneration, ttsAutoMode: cfg.TTSAutoMode, skillEvolve: cfg.SkillEvolve, skillNudgeInterval: cfg.SkillNudgeInterval, @@ -652,6 +662,9 @@ type MediaResult struct { ContentType string `json:"content_type,omitempty"` // MIME type Size int64 `json:"size,omitempty"` // file size in bytes AsVoice bool `json:"as_voice,omitempty"` // send as voice message (Telegram OGG) + // Prompt is the generation prompt for AI-generated media (e.g. create_image). + // Empty for user-uploaded or non-generated files. + Prompt string `json:"prompt,omitempty"` } // runState encapsulates all mutable state for a single agent run. diff --git a/internal/agent/media.go b/internal/agent/media.go index f56d1c41..de97ed95 100644 --- a/internal/agent/media.go +++ b/internal/agent/media.go @@ -1,6 +1,7 @@ package agent import ( + "crypto/sha256" "encoding/base64" "fmt" "io" @@ -16,6 +17,132 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" ) +// mediaWorkspaceDiskWarnThreshold is the size in bytes at which a warn-level +// log is emitted for the workspace/media/ directory. 500 MB. +const mediaWorkspaceDiskWarnThreshold = 500 * 1024 * 1024 + +// persistAssistantImages writes final (non-partial) images from msg.Images to +// {workspace}/media/{sha256}.{ext}, replaces them with MediaRefs, and clears +// msg.Images to prevent large base64 blobs from bloating the session store. +// +// Dedup: SHA256 hash is used as the filename, so writing the same image twice +// results in only one disk file. Idempotent: if the file already exists, the +// write is skipped but a new MediaRef is still appended (so the message +// correctly references the image regardless of dedup). +// +// Partial frames (Partial == true) are skipped — they are preview-only and +// must not be persisted to disk. +func persistAssistantImages(msg *providers.Message, workspace string) { + if workspace == "" || len(msg.Images) == 0 { + return + } + + mediaDir := filepath.Join(workspace, "media") + if err := os.MkdirAll(mediaDir, 0755); err != nil { + slog.Warn("media: failed to create workspace/media dir", "dir", mediaDir, "error", err) + return + } + // Symlink guard — identical to the pattern used for .uploads. + if fi, err := os.Lstat(mediaDir); err == nil && fi.Mode()&os.ModeSymlink != 0 { + slog.Warn("media: workspace/media is a symlink, refusing to use", "dir", mediaDir) + return + } + + var refs []providers.MediaRef + var totalBytes int64 + + for _, img := range msg.Images { + if img.Partial { + // Skip intermediate streaming frames — not final images. + continue + } + if img.Data == "" || img.MimeType == "" { + continue + } + + raw, err := base64.StdEncoding.DecodeString(img.Data) + if err != nil { + slog.Warn("media: failed to decode assistant image base64", "error", err) + continue + } + if len(raw) == 0 { + continue + } + + // Derive extension from MIME type. + ext := media.ExtFromMime(img.MimeType) + if ext == "" { + ext = ".bin" + } + + // SHA256 hash → deterministic filename enables free dedup. + sum := sha256.Sum256(raw) + hashHex := fmt.Sprintf("%x", sum) + filename := hashHex + ext + dstPath := filepath.Join(mediaDir, filename) + + // Traversal guard: resolved path must be inside mediaDir. + cleanDst := filepath.Clean(dstPath) + cleanMedia := filepath.Clean(mediaDir) + if !strings.HasPrefix(cleanDst+string(os.PathSeparator), cleanMedia+string(os.PathSeparator)) { + slog.Warn("media: refusing to persist outside workspace/media", "dst", dstPath, "media", mediaDir) + continue + } + + // Write only if the file does not already exist (idempotent on hash). + if _, statErr := os.Lstat(dstPath); os.IsNotExist(statErr) { + if writeErr := os.WriteFile(dstPath, raw, 0644); writeErr != nil { + slog.Warn("media: failed to write assistant image", "path", dstPath, "error", writeErr) + continue + } + slog.Debug("media: persisted assistant image", "path", dstPath, "mime", img.MimeType, "bytes", len(raw)) + } else { + slog.Debug("media: assistant image already on disk (dedup)", "path", dstPath) + } + + totalBytes += int64(len(raw)) + refs = append(refs, providers.MediaRef{ + ID: uuid.New().String(), + MimeType: img.MimeType, + Kind: "image", + Path: dstPath, + }) + } + + if len(refs) == 0 { + return + } + + // Attach refs to the message and clear inline base64 to save session store space. + msg.MediaRefs = append(msg.MediaRefs, refs...) + msg.Images = nil + + // Warn if workspace/media is growing large (quota enforcement deferred per phase spec). + go warnIfMediaDirLarge(mediaDir) +} + +// warnIfMediaDirLarge emits a warn log when {mediaDir} exceeds the disk threshold. +// Called in a goroutine to avoid blocking the pipeline finalize path. +func warnIfMediaDirLarge(mediaDir string) { + entries, err := os.ReadDir(mediaDir) + if err != nil { + return + } + var total int64 + for _, e := range entries { + if e.IsDir() { + continue + } + if info, err := e.Info(); err == nil { + total += info.Size() + } + } + if total > mediaWorkspaceDiskWarnThreshold { + slog.Warn("media: workspace/media dir exceeds 500 MB threshold", + "dir", mediaDir, "bytes", total) + } +} + // maxImageBytes is the safety limit for reading image files (10MB). const maxImageBytes = 10 * 1024 * 1024 diff --git a/internal/agent/media_assistant_images_test.go b/internal/agent/media_assistant_images_test.go new file mode 100644 index 00000000..230c485c --- /dev/null +++ b/internal/agent/media_assistant_images_test.go @@ -0,0 +1,294 @@ +package agent + +import ( + "encoding/base64" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +// minimalPNG is a 1x1 red PNG (67 bytes) — real PNG magic bytes + valid IHDR/IDAT. +// Used to verify that PNG magic bytes survive the write path. +var minimalPNG = func() []byte { + // base64 of a minimal 1x1 transparent PNG + const b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + b, _ := base64.StdEncoding.DecodeString(b64) + return b +}() + +// TestPersistAssistantImages_BasicPNG verifies that a final PNG image is written +// to {workspace}/media/{sha256}.png and Message.MediaRefs has one entry. +func TestPersistAssistantImages_BasicPNG(t *testing.T) { + workspace := t.TempDir() + + msg := &providers.Message{ + Role: "assistant", + Images: []providers.ImageContent{{ + MimeType: "image/png", + Data: base64.StdEncoding.EncodeToString(minimalPNG), + Partial: false, + }}, + } + + persistAssistantImages(msg, workspace) + + // Images must be cleared after persistence. + if len(msg.Images) != 0 { + t.Fatalf("expected Images cleared, got %d entries", len(msg.Images)) + } + // One MediaRef must be added. + if len(msg.MediaRefs) != 1 { + t.Fatalf("expected 1 MediaRef, got %d", len(msg.MediaRefs)) + } + ref := msg.MediaRefs[0] + if ref.Kind != "image" { + t.Errorf("MediaRef.Kind = %q, want %q", ref.Kind, "image") + } + if ref.MimeType != "image/png" { + t.Errorf("MediaRef.MimeType = %q, want %q", ref.MimeType, "image/png") + } + if !strings.HasSuffix(ref.Path, ".png") { + t.Errorf("MediaRef.Path %q must end with .png", ref.Path) + } + + // File must exist on disk with PNG magic bytes. + data, err := os.ReadFile(ref.Path) + if err != nil { + t.Fatalf("could not read persisted file: %v", err) + } + if len(data) < 4 || string(data[:4]) != "\x89PNG" { + t.Errorf("persisted file does not have PNG magic bytes, got %x", data[:min(4, len(data))]) + } + + // File must live inside workspace/media/. + mediaDir := filepath.Join(workspace, "media") + rel, err := filepath.Rel(mediaDir, ref.Path) + if err != nil || strings.HasPrefix(rel, "..") { + t.Errorf("persisted path %q is not inside workspace/media/", ref.Path) + } +} + +// TestPersistAssistantImages_Dedup verifies that writing the same image twice +// results in only one disk file. Both calls append a MediaRef (two refs, one file). +func TestPersistAssistantImages_Dedup(t *testing.T) { + workspace := t.TempDir() + imgData := base64.StdEncoding.EncodeToString(minimalPNG) + + msg1 := &providers.Message{ + Images: []providers.ImageContent{{MimeType: "image/png", Data: imgData}}, + } + msg2 := &providers.Message{ + Images: []providers.ImageContent{{MimeType: "image/png", Data: imgData}}, + } + + persistAssistantImages(msg1, workspace) + persistAssistantImages(msg2, workspace) + + mediaDir := filepath.Join(workspace, "media") + entries, err := os.ReadDir(mediaDir) + if err != nil { + t.Fatalf("ReadDir failed: %v", err) + } + // Same hash → exactly one file on disk. + if len(entries) != 1 { + t.Errorf("expected 1 file on disk (dedup), got %d", len(entries)) + } + // Each message gets its own MediaRef pointing to the same path. + if len(msg1.MediaRefs) != 1 { + t.Errorf("msg1: expected 1 MediaRef, got %d", len(msg1.MediaRefs)) + } + if len(msg2.MediaRefs) != 1 { + t.Errorf("msg2: expected 1 MediaRef, got %d", len(msg2.MediaRefs)) + } + if msg1.MediaRefs[0].Path != msg2.MediaRefs[0].Path { + t.Errorf("both msgs should reference same path; got %q and %q", + msg1.MediaRefs[0].Path, msg2.MediaRefs[0].Path) + } +} + +// TestPersistAssistantImages_SkipsPartial verifies that images with Partial=true +// are not persisted and do not produce MediaRefs. +func TestPersistAssistantImages_SkipsPartial(t *testing.T) { + workspace := t.TempDir() + imgData := base64.StdEncoding.EncodeToString(minimalPNG) + + msg := &providers.Message{ + Images: []providers.ImageContent{ + {MimeType: "image/png", Data: imgData, Partial: true}, // skip + {MimeType: "image/png", Data: imgData, Partial: false}, // persist + }, + } + + persistAssistantImages(msg, workspace) + + mediaDir := filepath.Join(workspace, "media") + entries, err := os.ReadDir(mediaDir) + if err != nil { + t.Fatalf("ReadDir failed: %v", err) + } + // Only the final (non-partial) image is written. + if len(entries) != 1 { + t.Errorf("expected 1 file (partial skipped), got %d", len(entries)) + } + if len(msg.MediaRefs) != 1 { + t.Errorf("expected 1 MediaRef (partial skipped), got %d", len(msg.MediaRefs)) + } + if msg.Images != nil { + t.Errorf("expected Images cleared, got %v", msg.Images) + } +} + +// TestPersistAssistantImages_EmptyWorkspace verifies that an empty workspace +// path is handled gracefully (no panic, no files written). +func TestPersistAssistantImages_EmptyWorkspace(t *testing.T) { + imgData := base64.StdEncoding.EncodeToString(minimalPNG) + msg := &providers.Message{ + Images: []providers.ImageContent{{MimeType: "image/png", Data: imgData}}, + } + + // Must not panic. + persistAssistantImages(msg, "") + + // Images should NOT be cleared (no workspace = nothing happened). + if len(msg.Images) == 0 { + t.Error("Images should remain when workspace is empty (early return)") + } + if len(msg.MediaRefs) != 0 { + t.Errorf("expected 0 MediaRefs when workspace is empty, got %d", len(msg.MediaRefs)) + } +} + +// TestPersistAssistantImages_AllPartials verifies that a message with only +// partial frames produces no disk files and leaves MediaRefs empty. +func TestPersistAssistantImages_AllPartials(t *testing.T) { + workspace := t.TempDir() + imgData := base64.StdEncoding.EncodeToString(minimalPNG) + + msg := &providers.Message{ + Images: []providers.ImageContent{ + {MimeType: "image/png", Data: imgData, Partial: true}, + {MimeType: "image/png", Data: imgData, Partial: true}, + }, + } + + persistAssistantImages(msg, workspace) + + mediaDir := filepath.Join(workspace, "media") + if _, err := os.Stat(mediaDir); err == nil { + entries, _ := os.ReadDir(mediaDir) + if len(entries) != 0 { + t.Errorf("expected 0 files (all partial), got %d", len(entries)) + } + } + if len(msg.MediaRefs) != 0 { + t.Errorf("expected 0 MediaRefs (all partial), got %d", len(msg.MediaRefs)) + } +} + +// TestPersistAssistantImages_MultipleDistinct verifies that two different images +// (different content → different hashes) produce two separate disk files. +func TestPersistAssistantImages_MultipleDistinct(t *testing.T) { + workspace := t.TempDir() + + // Create two distinct payloads by appending different bytes. + raw1 := append(minimalPNG[:len(minimalPNG):len(minimalPNG)], 0x01) + raw2 := append(minimalPNG[:len(minimalPNG):len(minimalPNG)], 0x02) + + msg := &providers.Message{ + Images: []providers.ImageContent{ + {MimeType: "image/png", Data: base64.StdEncoding.EncodeToString(raw1), Partial: false}, + {MimeType: "image/png", Data: base64.StdEncoding.EncodeToString(raw2), Partial: false}, + }, + } + + persistAssistantImages(msg, workspace) + + mediaDir := filepath.Join(workspace, "media") + entries, err := os.ReadDir(mediaDir) + if err != nil { + t.Fatalf("ReadDir failed: %v", err) + } + if len(entries) != 2 { + t.Errorf("expected 2 files (distinct images), got %d", len(entries)) + } + if len(msg.MediaRefs) != 2 { + t.Errorf("expected 2 MediaRefs, got %d", len(msg.MediaRefs)) + } + if msg.MediaRefs[0].Path == msg.MediaRefs[1].Path { + t.Errorf("expected distinct paths for distinct images") + } +} + +// TestPersistAssistantImages_PromptNotPropagated verifies that persistAssistantImages +// does NOT set MediaRef.Prompt (it only handles image bytes; prompt threading happens +// in the tools layer via result.MediaPrompts and in finalize_stage via MediaResult.Prompt). +// This test documents the current contract so any future signature change is caught. +func TestPersistAssistantImages_PromptNotPropagated(t *testing.T) { + workspace := t.TempDir() + imgData := base64.StdEncoding.EncodeToString(minimalPNG) + + msg := &providers.Message{ + Images: []providers.ImageContent{{ + MimeType: "image/png", + Data: imgData, + Partial: false, + }}, + } + + persistAssistantImages(msg, workspace) + + if len(msg.MediaRefs) != 1 { + t.Fatalf("expected 1 MediaRef, got %d", len(msg.MediaRefs)) + } + // persistAssistantImages has no access to prompts; Prompt must be empty here. + // The pipeline's finalize_stage sets Prompt on MediaRefs built from tool + // MediaResults (create_image path), not from Codex assistant image refs. + ref := msg.MediaRefs[0] + if ref.Prompt != "" { + t.Errorf("expected MediaRef.Prompt empty from persistAssistantImages, got %q", ref.Prompt) + } + if ref.Kind != "image" { + t.Errorf("MediaRef.Kind = %q, want image", ref.Kind) + } +} + +// TestPersistAssistantImages_PathInsideMediaDir verifies the hash-derived filename +// is exactly {sha256hex}.{ext} and lives directly under workspace/media/. +func TestPersistAssistantImages_PathInsideMediaDir(t *testing.T) { + workspace := t.TempDir() + imgData := base64.StdEncoding.EncodeToString(minimalPNG) + + msg := &providers.Message{ + Images: []providers.ImageContent{{MimeType: "image/png", Data: imgData}}, + } + persistAssistantImages(msg, workspace) + + ref := msg.MediaRefs[0] + base := filepath.Base(ref.Path) + dir := filepath.Dir(ref.Path) + + // Must be directly in workspace/media/ (no subdirectory). + wantDir := filepath.Join(workspace, "media") + if dir != wantDir { + t.Errorf("parent dir = %q, want %q", dir, wantDir) + } + // Filename must be {64hex}.png + if len(base) != 64+4 { // 64 hex + ".png" + t.Errorf("filename %q: expected {64hex}.png, len=%d", base, len(base)) + } + if !strings.HasSuffix(base, ".png") { + t.Errorf("filename %q must end with .png", base) + } + hashPart := strings.TrimSuffix(base, ".png") + for _, c := range hashPart { + if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) { + t.Errorf("filename %q: non-hex character %q in hash part", base, fmt.Sprintf("%c", c)) + break + } + } +} + diff --git a/internal/agent/png_metadata.go b/internal/agent/png_metadata.go new file mode 100644 index 00000000..67723169 --- /dev/null +++ b/internal/agent/png_metadata.go @@ -0,0 +1,112 @@ +package agent + +import ( + "bytes" + "encoding/binary" + "hash/crc32" +) + +// pngSignature is the 8-byte PNG file signature. +var pngSignature = []byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a} + +// EmbedPNGPrompt rewrites a PNG byte stream to include tEXt metadata chunks +// for "Description" (the generation prompt) and "Software" (goclaw). +// +// The chunks are inserted immediately before the IEND chunk so all image data +// remains valid. If the input is not a PNG (wrong magic bytes), the original +// bytes are returned unchanged without error. An empty prompt is a no-op. +// +// tEXt chunk format (per PNG spec): +// +// 4 bytes length of data field +// 4 bytes chunk type "tEXt" +// N bytes keyword\0text (data field) +// 4 bytes CRC32 of chunk-type + data +func EmbedPNGPrompt(pngBytes []byte, prompt string) ([]byte, error) { + if len(prompt) == 0 { + return pngBytes, nil + } + // Validate PNG signature. + if len(pngBytes) < len(pngSignature) || !bytes.Equal(pngBytes[:len(pngSignature)], pngSignature) { + // Not a PNG — return unchanged. + return pngBytes, nil + } + + // Build the tEXt chunks to insert. + extraChunks := buildTextChunks([]textKV{ + {Key: "Description", Value: prompt}, + {Key: "Software", Value: "goclaw"}, + }) + + // Locate the IEND chunk and insert before it. + iendOffset := findIENDOffset(pngBytes) + if iendOffset < 0 { + // Malformed PNG — return unchanged. + return pngBytes, nil + } + + result := make([]byte, 0, len(pngBytes)+len(extraChunks)) + result = append(result, pngBytes[:iendOffset]...) + result = append(result, extraChunks...) + result = append(result, pngBytes[iendOffset:]...) + return result, nil +} + +// textKV is a keyword/value pair for PNG tEXt chunks. +type textKV struct { + Key string + Value string +} + +// buildTextChunks encodes multiple tEXt chunks into raw PNG chunk bytes. +func buildTextChunks(pairs []textKV) []byte { + var buf bytes.Buffer + for _, p := range pairs { + // data = keyword + NUL + text + data := make([]byte, 0, len(p.Key)+1+len(p.Value)) + data = append(data, []byte(p.Key)...) + data = append(data, 0x00) + data = append(data, []byte(p.Value)...) + + chunkType := []byte("tEXt") + crcInput := append(chunkType, data...) + checksum := crc32.ChecksumIEEE(crcInput) + + // 4-byte length + var lenBuf [4]byte + binary.BigEndian.PutUint32(lenBuf[:], uint32(len(data))) + buf.Write(lenBuf[:]) + + // chunk type + buf.Write(chunkType) + + // data + buf.Write(data) + + // CRC32 + var crcBuf [4]byte + binary.BigEndian.PutUint32(crcBuf[:], checksum) + buf.Write(crcBuf[:]) + } + return buf.Bytes() +} + +// findIENDOffset returns the byte offset at which the IEND chunk starts. +// Returns -1 if IEND is not found (malformed PNG). +func findIENDOffset(data []byte) int { + pos := len(pngSignature) + for pos+12 <= len(data) { + chunkLen := int(binary.BigEndian.Uint32(data[pos : pos+4])) + chunkType := data[pos+4 : pos+8] + if bytes.Equal(chunkType, []byte("IEND")) { + return pos + } + // Advance: 4 (length) + 4 (type) + chunkLen (data) + 4 (CRC) + pos += 8 + chunkLen + 4 + if chunkLen < 0 || pos < 0 { + // Overflow guard. + break + } + } + return -1 +} diff --git a/internal/agent/png_metadata_test.go b/internal/agent/png_metadata_test.go new file mode 100644 index 00000000..0d45977b --- /dev/null +++ b/internal/agent/png_metadata_test.go @@ -0,0 +1,137 @@ +package agent + +import ( + "bytes" + "encoding/binary" + "strings" + "testing" +) + +// TestEmbedPNGPrompt_RoundTrip embeds a prompt into a real PNG and verifies +// the tEXt chunk can be read back by parsing raw chunk bytes. +func TestEmbedPNGPrompt_RoundTrip(t *testing.T) { + wantPrompt := "A vibrant sunset over the ocean" + + out, err := EmbedPNGPrompt(minimalPNG, wantPrompt) + if err != nil { + t.Fatalf("EmbedPNGPrompt: %v", err) + } + if len(out) <= len(minimalPNG) { + t.Errorf("output (%d bytes) must be larger than input (%d bytes)", len(out), len(minimalPNG)) + } + + // Parse tEXt chunks from the output PNG. + texts := parsePNGTextChunks(out) + + // "Description" chunk must carry the prompt. + got, ok := texts["Description"] + if !ok { + t.Fatalf("no tEXt 'Description' chunk found; chunks = %v", texts) + } + if got != wantPrompt { + t.Errorf("Description = %q, want %q", got, wantPrompt) + } + + // "Software" chunk must carry "goclaw". + if sw := texts["Software"]; sw != "goclaw" { + t.Errorf("Software = %q, want %q", sw, "goclaw") + } +} + +// TestEmbedPNGPrompt_EmptyPrompt verifies that an empty prompt is a no-op +// (output identical to input). +func TestEmbedPNGPrompt_EmptyPrompt(t *testing.T) { + out, err := EmbedPNGPrompt(minimalPNG, "") + if err != nil { + t.Fatalf("EmbedPNGPrompt with empty prompt: %v", err) + } + if !bytes.Equal(out, minimalPNG) { + t.Error("expected output identical to input for empty prompt") + } +} + +// TestEmbedPNGPrompt_NonPNGPassthrough verifies that non-PNG bytes are returned +// unchanged (no error). +func TestEmbedPNGPrompt_NonPNGPassthrough(t *testing.T) { + notPNG := []byte("this is not a png file at all") + out, err := EmbedPNGPrompt(notPNG, "some prompt") + if err != nil { + t.Fatalf("EmbedPNGPrompt on non-PNG: %v", err) + } + if !bytes.Equal(out, notPNG) { + t.Error("expected non-PNG bytes returned unchanged") + } +} + +// TestEmbedPNGPrompt_LongPrompt verifies that a prompt longer than 1 KB round-trips +// correctly (tEXt chunks have no length limit). +func TestEmbedPNGPrompt_LongPrompt(t *testing.T) { + longPrompt := strings.Repeat("detailed landscape with mountains, ", 40) + + out, err := EmbedPNGPrompt(minimalPNG, longPrompt) + if err != nil { + t.Fatalf("EmbedPNGPrompt with long prompt: %v", err) + } + texts := parsePNGTextChunks(out) + if got := texts["Description"]; got != longPrompt { + t.Errorf("long prompt round-trip failed: len(got)=%d len(want)=%d", + len(got), len(longPrompt)) + } +} + +// TestEmbedPNGPrompt_IENDStillLast verifies the structural invariant that +// IEND remains the last chunk in the output PNG after embedding. +func TestEmbedPNGPrompt_IENDStillLast(t *testing.T) { + out, err := EmbedPNGPrompt(minimalPNG, "test prompt") + if err != nil { + t.Fatalf("EmbedPNGPrompt: %v", err) + } + + // Walk chunks and record the last one we see. + pos := len(pngSignature) + lastType := "" + for pos+12 <= len(out) { + chunkLen := int(binary.BigEndian.Uint32(out[pos : pos+4])) + if chunkLen < 0 { + break + } + lastType = string(out[pos+4 : pos+8]) + next := pos + 8 + chunkLen + 4 + if next <= pos { + break + } + pos = next + } + if lastType != "IEND" { + t.Errorf("last chunk type = %q, want IEND", lastType) + } +} + +// parsePNGTextChunks walks a PNG byte stream and extracts all tEXt chunks as +// a map of keyword → text. Used only by tests to verify round-trip correctness. +func parsePNGTextChunks(data []byte) map[string]string { + result := make(map[string]string) + pos := len(pngSignature) + for pos+12 <= len(data) { + chunkLen := int(binary.BigEndian.Uint32(data[pos : pos+4])) + if chunkLen < 0 { + break + } + chunkType := string(data[pos+4 : pos+8]) + chunkData := data[pos+8 : pos+8+chunkLen] + if chunkType == "tEXt" { + // tEXt format: keyword\0text + if nul := bytes.IndexByte(chunkData, 0x00); nul >= 0 { + keyword := string(chunkData[:nul]) + text := string(chunkData[nul+1:]) + result[keyword] = text + } + } + next := pos + 8 + chunkLen + 4 + if next <= pos { + break + } + pos = next + } + return result +} diff --git a/internal/agent/preview_prompt.go b/internal/agent/preview_prompt.go index abc71ece..058542b2 100644 --- a/internal/agent/preview_prompt.go +++ b/internal/agent/preview_prompt.go @@ -232,7 +232,7 @@ func BuildPreviewPrompt(ctx context.Context, ag *store.AgentData, mode PromptMod if tool, ok := deps.ToolLister.Get(canonical); ok { toolDefs = append(toolDefs, providers.ToolDefinition{ Type: "function", - Function: providers.ToolFunctionSchema{ + Function: &providers.ToolFunctionSchema{ Name: alias, Description: tool.Description(), Parameters: tool.Parameters(), diff --git a/internal/agent/resolver.go b/internal/agent/resolver.go index b50fc472..5b62fc31 100644 --- a/internal/agent/resolver.go +++ b/internal/agent/resolver.go @@ -514,6 +514,7 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc { PromptMode: PromptMode(ag.ParsePromptMode()), PinnedSkills: ag.ParsePinnedSkills(), SelfEvolve: ag.ParseSelfEvolve(), + AllowImageGeneration: ag.ParseAllowImageGeneration(), TTSAutoMode: deps.TTSAutoMode, SkillEvolve: ag.AgentType == store.AgentTypePredefined && ag.ParseSkillEvolve(), SkillNudgeInterval: ag.ParseSkillNudgeInterval(), diff --git a/internal/hooks/handlers/prompt.go b/internal/hooks/handlers/prompt.go index 7b4425b2..2ec9fb1f 100644 --- a/internal/hooks/handlers/prompt.go +++ b/internal/hooks/handlers/prompt.go @@ -263,7 +263,7 @@ func (h *PromptHandler) buildChatRequest(cfg hooks.HookConfig, ev hooks.Event, m }, Tools: []providers.ToolDefinition{{ Type: "function", - Function: providers.ToolFunctionSchema{ + Function: &providers.ToolFunctionSchema{ Name: promptDecideToolName, Description: "Return the hook evaluation decision.", Parameters: map[string]any{ diff --git a/internal/permissions/policy.go b/internal/permissions/policy.go index 5e4f391a..9c75d61d 100644 --- a/internal/permissions/policy.go +++ b/internal/permissions/policy.go @@ -292,6 +292,7 @@ func isWriteMethod(method string) bool { protocol.MethodSessionsDelete, protocol.MethodSessionsReset, protocol.MethodSessionsPatch, + protocol.MethodSessionsCompact, protocol.MethodCronCreate, protocol.MethodCronUpdate, protocol.MethodCronDelete, diff --git a/internal/pipeline/deps.go b/internal/pipeline/deps.go index e941d4a8..1c5d095d 100644 --- a/internal/pipeline/deps.go +++ b/internal/pipeline/deps.go @@ -99,6 +99,11 @@ type PipelineDeps struct { FlushMessages func(ctx context.Context, sessionKey string, msgs []providers.Message) error // Finalize callbacks (FinalizeStage) + // PersistAssistantImages writes final (non-partial) images from the assistant + // response to workspace disk, appends MediaRefs, and clears inline base64. + // Called BEFORE building the assistant message for session persistence. + // nil = feature disabled (no Codex image gen or no workspace). + PersistAssistantImages func(msg *providers.Message, workspace string) SkillPostscript func(ctx context.Context, content string, totalToolCalls int) string // skill evolution nudge (nil = disabled) SanitizeContent func(content string) string StripMessageDirectives func(content string) string diff --git a/internal/pipeline/finalize_stage.go b/internal/pipeline/finalize_stage.go index 9ce406ad..3236e9f6 100644 --- a/internal/pipeline/finalize_stage.go +++ b/internal/pipeline/finalize_stage.go @@ -61,7 +61,27 @@ func (s *FinalizeStage) Execute(ctx context.Context, state *RunState) error { // 3. Deduplicate + populate media sizes s.processMedia(state) - // 3b. Build final assistant message with MediaRefs for session persistence. + // 3b. Persist assistant-generated images (Codex image_generation_call) to disk + // BEFORE building the assistant message so MediaRefs are included in the session store. + // Source is state.Observe.AssistantImages, which ObserveStage accumulates across + // every iteration — required because LastResponse holds only the final iteration's + // response (an image emitted mid-loop alongside a tool call would otherwise be lost). + var assistantImageRefs []providers.MediaRef + if s.deps.PersistAssistantImages != nil && len(state.Observe.AssistantImages) > 0 { + workspace := "" + if state.Workspace != nil { + workspace = state.Workspace.ActivePath + } + // Build a scratch message carrying only Images so PersistAssistantImages can + // decode/hash/write them and populate MediaRefs. The caller clears Images on + // the scratch message — we harvest MediaRefs from there. + scratch := &providers.Message{Images: state.Observe.AssistantImages} + s.deps.PersistAssistantImages(scratch, workspace) + assistantImageRefs = scratch.MediaRefs + state.Observe.AssistantImages = nil // prevent double-processing on retries + } + + // 3c. Build final assistant message with MediaRefs for session persistence. assistantMsg := providers.Message{ Role: "assistant", Content: state.Observe.FinalContent, @@ -82,8 +102,11 @@ func (s *FinalizeStage) Execute(ctx context.Context, state *RunState) error { MimeType: mr.ContentType, Kind: kind, Path: mr.Path, + Prompt: mr.Prompt, }) } + // Append persisted assistant image refs (Codex image_generation_call output). + assistantMsg.MediaRefs = append(assistantMsg.MediaRefs, assistantImageRefs...) state.Messages.AppendPending(assistantMsg) // 4. Flush remaining pending messages to session store diff --git a/internal/pipeline/observe_stage.go b/internal/pipeline/observe_stage.go index d24ac88d..0bc6763b 100644 --- a/internal/pipeline/observe_stage.go +++ b/internal/pipeline/observe_stage.go @@ -45,5 +45,24 @@ func (s *ObserveStage) Execute(_ context.Context, state *RunState) error { state.Observe.FinalThinking = resp.Thinking } + // 4. Accumulate assistant-generated final images across iterations. + // The LLM may emit image_generation_call in iter N alongside a function_call, + // then respond text-only in iter N+1 — LastResponse.Images would then be empty + // at finalize time and the iter-N image would be lost. Drain Images here so + // FinalizeStage sees every image regardless of which iteration produced it. + // Partial streaming frames are filtered out at source (codex.go only sets + // non-partial entries via imageState.recordFinal); a defensive filter here + // avoids coupling to that invariant. + if len(resp.Images) > 0 { + for _, img := range resp.Images { + if img.Partial { + continue + } + state.Observe.AssistantImages = append(state.Observe.AssistantImages, img) + } + // Clear on response so a re-processing pass (e.g. retry) doesn't double-count. + resp.Images = nil + } + return nil } diff --git a/internal/pipeline/run_state.go b/internal/pipeline/run_state.go index 8114db80..3f3fa28c 100644 --- a/internal/pipeline/run_state.go +++ b/internal/pipeline/run_state.go @@ -111,4 +111,7 @@ type MediaResult struct { ContentType string Size int64 AsVoice bool + // Prompt is the generation prompt for AI-generated media (e.g. create_image). + // Empty for user-uploaded or non-generated files. + Prompt string } diff --git a/internal/pipeline/stages_test.go b/internal/pipeline/stages_test.go index 94c25fda..10f3f7bf 100644 --- a/internal/pipeline/stages_test.go +++ b/internal/pipeline/stages_test.go @@ -1241,6 +1241,155 @@ func TestObserveStage_EmptyContent_BlockRepliesNotIncremented(t *testing.T) { } } +// --- ObserveStage image accumulation (regression for mid-loop image loss) --- +// +// These tests cover the bug where LLM emits an image_generation_call alongside +// a function_call in iter N, then responds text-only in iter N+1. Without +// accumulation in Observe, FinalizeStage would only see LastResponse.Images +// (which is empty at iter N+1) and drop the iter-N image. + +// Case 1: single iteration with image only → accumulated. +func TestObserveStage_ImageAccumulation_SingleIterImageOnly(t *testing.T) { + t.Parallel() + stage := NewObserveStage(&PipelineDeps{}) + state := defaultState() + state.Think.LastResponse = &providers.ChatResponse{ + FinishReason: "stop", + Images: []providers.ImageContent{ + {MimeType: "image/png", Data: "imgA"}, + }, + } + _ = stage.Execute(context.Background(), state) + if got := len(state.Observe.AssistantImages); got != 1 { + t.Fatalf("AssistantImages len = %d, want 1", got) + } + if state.Observe.AssistantImages[0].Data != "imgA" { + t.Errorf("image data = %q, want %q", state.Observe.AssistantImages[0].Data, "imgA") + } + // Source response.Images must be cleared to prevent double-counting on re-exec. + if state.Think.LastResponse.Images != nil { + t.Error("LastResponse.Images must be cleared after draining") + } +} + +// Case 2: image + tool_call in same iter → image accumulated, tool_call flows through think. +func TestObserveStage_ImageAccumulation_ImagePlusToolCall(t *testing.T) { + t.Parallel() + stage := NewObserveStage(&PipelineDeps{}) + state := defaultState() + state.Think.LastResponse = &providers.ChatResponse{ + FinishReason: "tool_calls", + ToolCalls: []providers.ToolCall{{ID: "1", Name: "search"}}, + Images: []providers.ImageContent{ + {MimeType: "image/png", Data: "imgMid"}, + }, + } + _ = stage.Execute(context.Background(), state) + if got := len(state.Observe.AssistantImages); got != 1 { + t.Fatalf("AssistantImages len = %d, want 1 (image must survive tool_calls path)", got) + } +} + +// Case 3: mid-loop image in iter 1 + text-only iter 2 → image from iter 1 retained. +// +// This is the regression scenario that motivated the accumulator. Without the fix +// FinalizeStage reads LastResponse (iter 2) and drops iter-1 image. +func TestObserveStage_ImageAccumulation_MidLoopImagePreservedAcrossIterations(t *testing.T) { + t.Parallel() + stage := NewObserveStage(&PipelineDeps{}) + state := defaultState() + + // Iter 1: image + tool call. + state.Think.LastResponse = &providers.ChatResponse{ + FinishReason: "tool_calls", + ToolCalls: []providers.ToolCall{{ID: "1", Name: "search"}}, + Images: []providers.ImageContent{{MimeType: "image/png", Data: "iter1img"}}, + } + _ = stage.Execute(context.Background(), state) + + // Iter 2: text-only final response — no Images. + state.Think.LastResponse = &providers.ChatResponse{ + FinishReason: "stop", + Content: "Done.", + } + _ = stage.Execute(context.Background(), state) + + if got := len(state.Observe.AssistantImages); got != 1 { + t.Fatalf("AssistantImages len = %d, want 1 (iter-1 image must survive)", got) + } + if state.Observe.AssistantImages[0].Data != "iter1img" { + t.Errorf("image data = %q, want %q", state.Observe.AssistantImages[0].Data, "iter1img") + } +} + +// Case 4: multiple images emitted across multiple iterations → all retained in order. +func TestObserveStage_ImageAccumulation_MultipleImagesAcrossIterations(t *testing.T) { + t.Parallel() + stage := NewObserveStage(&PipelineDeps{}) + state := defaultState() + + // Iter 1: two images + tool call. + state.Think.LastResponse = &providers.ChatResponse{ + FinishReason: "tool_calls", + ToolCalls: []providers.ToolCall{{ID: "1", Name: "t"}}, + Images: []providers.ImageContent{ + {MimeType: "image/png", Data: "A"}, + {MimeType: "image/png", Data: "B"}, + }, + } + _ = stage.Execute(context.Background(), state) + + // Iter 2: one image standalone. + state.Think.LastResponse = &providers.ChatResponse{ + FinishReason: "stop", + Images: []providers.ImageContent{{MimeType: "image/png", Data: "C"}}, + } + _ = stage.Execute(context.Background(), state) + + if got := len(state.Observe.AssistantImages); got != 3 { + t.Fatalf("AssistantImages len = %d, want 3", got) + } + for i, want := range []string{"A", "B", "C"} { + if state.Observe.AssistantImages[i].Data != want { + t.Errorf("image[%d] = %q, want %q", i, state.Observe.AssistantImages[i].Data, want) + } + } +} + +// Case 5: partial frames must be filtered out — only final (non-partial) images accumulate. +func TestObserveStage_ImageAccumulation_PartialFramesFiltered(t *testing.T) { + t.Parallel() + stage := NewObserveStage(&PipelineDeps{}) + state := defaultState() + state.Think.LastResponse = &providers.ChatResponse{ + FinishReason: "stop", + Images: []providers.ImageContent{ + {MimeType: "image/png", Data: "partial1", Partial: true}, + {MimeType: "image/png", Data: "final1"}, + {MimeType: "image/png", Data: "partial2", Partial: true}, + }, + } + _ = stage.Execute(context.Background(), state) + if got := len(state.Observe.AssistantImages); got != 1 { + t.Fatalf("AssistantImages len = %d, want 1 (partials filtered)", got) + } + if state.Observe.AssistantImages[0].Data != "final1" { + t.Errorf("image data = %q, want %q", state.Observe.AssistantImages[0].Data, "final1") + } +} + +// Case 6: nil response → no panic, accumulator unchanged. +func TestObserveStage_ImageAccumulation_NilResponseSafe(t *testing.T) { + t.Parallel() + stage := NewObserveStage(&PipelineDeps{}) + state := defaultState() + state.Think.LastResponse = nil + _ = stage.Execute(context.Background(), state) + if state.Observe.AssistantImages != nil { + t.Errorf("AssistantImages = %v, want nil", state.Observe.AssistantImages) + } +} + // --- CheckpointStage tests --- func TestCheckpointStage_SkipsIteration0(t *testing.T) { @@ -1408,6 +1557,75 @@ func TestFinalizeStage_DeduplicatesMediaByPath(t *testing.T) { } } +// TestFinalizeStage_PersistsFromObserveAccumulator verifies that FinalizeStage +// sources assistant images from state.Observe.AssistantImages, NOT from +// LastResponse.Images. This guards against the regression where a mid-loop +// image_generation_call is lost when the final iteration responds text-only. +func TestFinalizeStage_PersistsFromObserveAccumulator(t *testing.T) { + t.Parallel() + var persistedImages []providers.ImageContent + deps := &PipelineDeps{ + PersistAssistantImages: func(msg *providers.Message, _ string) { + // Capture what was handed to the persist callback. + persistedImages = append([]providers.ImageContent(nil), msg.Images...) + // Simulate hash→MediaRef mapping. + for range msg.Images { + msg.MediaRefs = append(msg.MediaRefs, providers.MediaRef{ + Kind: "image", MimeType: "image/png", Path: "/tmp/img.png", + }) + } + msg.Images = nil + }, + FlushMessages: func(_ context.Context, _ string, _ []providers.Message) error { return nil }, + } + stage := NewFinalizeStage(deps) + state := defaultState() + + // Observe has accumulated an image from an earlier iteration. + state.Observe.AssistantImages = []providers.ImageContent{ + {MimeType: "image/png", Data: "iterMidImage"}, + } + // LastResponse is the final iteration — text-only, no Images. + state.Think.LastResponse = &providers.ChatResponse{FinishReason: "stop", Content: "Done."} + state.Observe.FinalContent = "Done." + + if err := stage.Execute(context.Background(), state); err != nil { + t.Fatalf("Execute() error: %v", err) + } + if len(persistedImages) != 1 { + t.Fatalf("persisted images len = %d, want 1 (accumulator-sourced image must be persisted)", len(persistedImages)) + } + if persistedImages[0].Data != "iterMidImage" { + t.Errorf("persisted image data = %q, want %q", persistedImages[0].Data, "iterMidImage") + } + // Accumulator must be drained. + if state.Observe.AssistantImages != nil { + t.Errorf("AssistantImages = %v, want nil after drain", state.Observe.AssistantImages) + } +} + +// TestFinalizeStage_NoPersistWhenAccumulatorEmpty verifies no-op when no images +// were emitted across any iteration. Prevents regression where a non-nil +// LastResponse with empty Images would still call PersistAssistantImages. +func TestFinalizeStage_NoPersistWhenAccumulatorEmpty(t *testing.T) { + t.Parallel() + persistCalled := false + deps := &PipelineDeps{ + PersistAssistantImages: func(_ *providers.Message, _ string) { persistCalled = true }, + FlushMessages: func(_ context.Context, _ string, _ []providers.Message) error { return nil }, + } + stage := NewFinalizeStage(deps) + state := defaultState() + state.Think.LastResponse = &providers.ChatResponse{FinishReason: "stop", Content: "hi"} + + if err := stage.Execute(context.Background(), state); err != nil { + t.Fatalf("Execute() error: %v", err) + } + if persistCalled { + t.Error("PersistAssistantImages must not be called when accumulator is empty") + } +} + func TestFinalizeStage_PopulatesFileSizes(t *testing.T) { t.Parallel() // create a real temp file diff --git a/internal/pipeline/substates.go b/internal/pipeline/substates.go index a0e56a91..aeb3a331 100644 --- a/internal/pipeline/substates.go +++ b/internal/pipeline/substates.go @@ -59,6 +59,14 @@ type ObserveState struct { FinalThinking string // reasoning output BlockReplies int LastBlockReply string + + // AssistantImages accumulates final (non-partial) images from every iteration's + // ChatResponse.Images. FinalizeStage persists these to workspace/media/. + // Accumulation is required because LastResponse holds only the final iteration's + // response — if the LLM emits an image_generation_call alongside a function_call + // in iter N and responds text-only in iter N+1, reading only LastResponse.Images + // would lose the image. + AssistantImages []providers.ImageContent } // CompactState: owned by CheckpointStage + MemoryFlushStage. diff --git a/internal/providers/adapter_anthropic_test.go b/internal/providers/adapter_anthropic_test.go index 1bb82721..da976683 100644 --- a/internal/providers/adapter_anthropic_test.go +++ b/internal/providers/adapter_anthropic_test.go @@ -85,8 +85,8 @@ func TestAnthropicAdapterToRequest_CacheControl(t *testing.T) { {Role: "user", Content: "Hello"}, }, Tools: []ToolDefinition{ - {Type: "function", Function: ToolFunctionSchema{Name: "tool1", Description: "desc1", Parameters: map[string]any{"type": "object"}}}, - {Type: "function", Function: ToolFunctionSchema{Name: "tool2", Description: "desc2", Parameters: map[string]any{"type": "object"}}}, + {Type: "function", Function: &ToolFunctionSchema{Name: "tool1", Description: "desc1", Parameters: map[string]any{"type": "object"}}}, + {Type: "function", Function: &ToolFunctionSchema{Name: "tool2", Description: "desc2", Parameters: map[string]any{"type": "object"}}}, }, } data, _, err := adapter.ToRequest(req) diff --git a/internal/providers/adapter_codex.go b/internal/providers/adapter_codex.go index 17175f16..8f3f856b 100644 --- a/internal/providers/adapter_codex.go +++ b/internal/providers/adapter_codex.go @@ -52,6 +52,7 @@ func (a *CodexAdapter) Capabilities() ProviderCapabilities { Thinking: true, Vision: true, CacheControl: false, + ImageGeneration: true, // Codex (OpenAI Responses API) supports native image_generation tool MaxContextWindow: 1_000_000, TokenizerID: "o200k_base", } diff --git a/internal/providers/capabilities.go b/internal/providers/capabilities.go index e982bf57..2c6bebdf 100644 --- a/internal/providers/capabilities.go +++ b/internal/providers/capabilities.go @@ -11,6 +11,7 @@ type ProviderCapabilities struct { Thinking bool // supports extended thinking / reasoning Vision bool // supports image inputs CacheControl bool // supports cache_control blocks (Anthropic) + ImageGeneration bool // supports native image_generation tool (Codex/OpenAI Responses API) MaxContextWindow int // default context window for default model TokenizerID string // for tokencount package mapping } diff --git a/internal/providers/capabilities_image_gen_test.go b/internal/providers/capabilities_image_gen_test.go new file mode 100644 index 00000000..6743713a --- /dev/null +++ b/internal/providers/capabilities_image_gen_test.go @@ -0,0 +1,50 @@ +package providers + +import "testing" + +// TestCodexProvider_ImageGenCapability verifies CodexProvider.Capabilities() +// reports ImageGeneration=true. +func TestCodexProvider_ImageGenCapability(t *testing.T) { + ts := &staticTokenSource{token: "tok"} + p := NewCodexProvider("codex", ts, "", "") + caps := p.Capabilities() + if !caps.ImageGeneration { + t.Error("CodexProvider.Capabilities().ImageGeneration must be true") + } +} + +// TestCodexAdapter_ImageGenCapability verifies CodexAdapter.Capabilities() +// also reports ImageGeneration=true (adapter must mirror provider). +func TestCodexAdapter_ImageGenCapability(t *testing.T) { + a, err := NewCodexAdapter(ProviderConfig{}) + if err != nil { + t.Fatalf("NewCodexAdapter: %v", err) + } + caps := a.Capabilities() + if !caps.ImageGeneration { + t.Error("CodexAdapter.Capabilities().ImageGeneration must be true") + } +} + +// TestOtherProviders_NoImageGenCapability verifies providers that do NOT support +// image_generation return ImageGeneration=false (the zero value). This protects +// against accidentally attaching the tool to non-Codex providers. +func TestOtherProviders_NoImageGenCapability(t *testing.T) { + // Anthropic + ap := &AnthropicProvider{} + if ap.Capabilities().ImageGeneration { + t.Error("AnthropicProvider must not advertise ImageGeneration") + } + + // DashScope + dp := &DashScopeProvider{} + if dp.Capabilities().ImageGeneration { + t.Error("DashScopeProvider must not advertise ImageGeneration") + } + + // OpenAI (via OpenAIProvider using compat layer) + op := &OpenAIProvider{} + if op.Capabilities().ImageGeneration { + t.Error("OpenAIProvider must not advertise ImageGeneration") + } +} diff --git a/internal/providers/codex.go b/internal/providers/codex.go index f2d6106d..3ed3a8d2 100644 --- a/internal/providers/codex.go +++ b/internal/providers/codex.go @@ -69,6 +69,7 @@ func (p *CodexProvider) Capabilities() ProviderCapabilities { Thinking: true, Vision: true, CacheControl: false, + ImageGeneration: true, // Codex (OpenAI Responses API) supports native image_generation tool MaxContextWindow: 1_000_000, TokenizerID: "o200k_base", } @@ -139,6 +140,7 @@ func (p *CodexProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk result := &ChatResponse{FinishReason: "stop"} toolCalls := make(map[string]*codexToolCallAcc) // keyed by item_id streamState := newCodexMessageStreamState() + imageState := newCodexImageState() sse := NewSSEScanner(cb) for sse.Next() { @@ -149,7 +151,7 @@ func (p *CodexProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk continue } - if err := p.processSSEEvent(&event, result, toolCalls, streamState, onChunk, stripThinking); err != nil { + if err := p.processSSEEvent(&event, result, toolCalls, streamState, imageState, onChunk, stripThinking); err != nil { return nil, err } } @@ -158,6 +160,9 @@ func (p *CodexProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk return nil, fmt.Errorf("%s: stream read error: %w", p.name, err) } + // Assemble generated images from image accumulator into ChatResponse. + imageState.appendToResponse(result) + // Build tool calls from accumulators for _, acc := range toolCalls { if acc.name == "" { @@ -192,8 +197,21 @@ func (p *CodexProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk // processSSEEvent handles a single SSE event during streaming. // stripThinking drops reasoning summaries from user-visible output while // leaving billing counters (Usage.ThinkingTokens) untouched. -func (p *CodexProvider) processSSEEvent(event *codexSSEEvent, result *ChatResponse, toolCalls map[string]*codexToolCallAcc, streamState *codexMessageStreamState, onChunk func(StreamChunk), stripThinking bool) error { +func (p *CodexProvider) processSSEEvent(event *codexSSEEvent, result *ChatResponse, toolCalls map[string]*codexToolCallAcc, streamState *codexMessageStreamState, imageState *codexImageState, onChunk func(StreamChunk), stripThinking bool) error { switch event.Type { + case "response.image_generation_call.partial_image": + // Intermediate frame from a streaming image generation call. + // Deduplicate by SHA256 so identical frames are not re-emitted. + if imageState.recordPartial(event.ItemID, event.OutputFormat, event.PartialImageB64) { + if onChunk != nil { + onChunk(StreamChunk{Images: []ImageContent{{ + MimeType: mimeFromFormat(event.OutputFormat), + Data: event.PartialImageB64, + Partial: true, + }}}) + } + } + case "response.output_item.added": if event.Item != nil { streamState.registerMessageItem(event.ItemID, event.OutputIndex, event.Item) @@ -249,6 +267,20 @@ func (p *CodexProvider) processSSEEvent(event *codexSSEEvent, result *ChatRespon } } } + case "image_generation_call": + // Final image for this item. Record and emit a non-partial chunk. + itemID := event.Item.ID + if itemID == "" { + itemID = event.ItemID + } + imageState.recordFinal(itemID, event.Item.OutputFormat, event.Item.Result) + if event.Item.Result != "" && onChunk != nil { + onChunk(StreamChunk{Images: []ImageContent{{ + MimeType: mimeFromFormat(event.Item.OutputFormat), + Data: event.Item.Result, + Partial: false, + }}}) + } } } @@ -259,6 +291,16 @@ func (p *CodexProvider) processSSEEvent(event *codexSSEEvent, result *ChatRespon streamState.flushCompletedResponse(result, onChunk) streamState.updateResultPhase(result) } + // Walk output[] for image_generation_call items not captured via stream events. + // This covers non-streaming mode (single response.completed with all outputs) + // and acts as a safety net for the streaming case. + for i := range event.Response.Output { + item := &event.Response.Output[i] + if item.Type == "image_generation_call" && item.Result != "" { + itemID := item.ID + imageState.recordFinal(itemID, item.OutputFormat, item.Result) + } + } if event.Response.Usage != nil { u := event.Response.Usage result.Usage = &Usage{ diff --git a/internal/providers/codex_build.go b/internal/providers/codex_build.go index 88786404..5a9a941a 100644 --- a/internal/providers/codex_build.go +++ b/internal/providers/codex_build.go @@ -111,12 +111,26 @@ func (p *CodexProvider) buildRequestBody(req ChatRequest, stream bool) map[strin if len(req.Tools) > 0 { var tools []map[string]any for _, t := range req.Tools { - tools = append(tools, map[string]any{ - "type": "function", - "name": t.Function.Name, - "description": t.Function.Description, - "parameters": NormalizeSchema("codex", t.Function.Parameters), - }) + if t.Type == "image_generation" { + // Pass native image_generation tool object as-is — Responses API first-class tool. + // Defaults chosen for Phase 1b; per-agent overrides are Phase 4. + tools = append(tools, map[string]any{ + "type": "image_generation", + "action": "generate", + "model": "gpt-image-2", + "output_format": "png", + "partial_images": 1, + }) + } else { + // Function tool path (default). Works with both value-type and pointer Function + // fields — we only access t.Function when type is not "image_generation". + tools = append(tools, map[string]any{ + "type": "function", + "name": t.Function.Name, + "description": t.Function.Description, + "parameters": NormalizeSchema("codex", t.Function.Parameters), + }) + } } body["tools"] = tools } diff --git a/internal/providers/codex_image_test.go b/internal/providers/codex_image_test.go new file mode 100644 index 00000000..2a5f0f96 --- /dev/null +++ b/internal/providers/codex_image_test.go @@ -0,0 +1,377 @@ +package providers + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "testing" +) + +// pngHeader is the 8-byte PNG signature for validity checks. +var pngHeader = []byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a} + +// decodeBase64PNG decodes a base64 string and verifies it has a valid PNG header. +func decodeBase64PNG(t *testing.T, b64 string) []byte { + t.Helper() + raw, err := base64.StdEncoding.DecodeString(b64) + if err != nil { + t.Fatalf("base64 decode error: %v", err) + } + if len(raw) < 8 { + t.Fatalf("decoded bytes too short for PNG header: %d bytes", len(raw)) + } + for i, b := range pngHeader { + if raw[i] != b { + t.Fatalf("PNG header mismatch at byte %d: got 0x%02x want 0x%02x", i, raw[i], b) + } + } + return raw +} + +// loadFixtureEvents reads a JSON fixture file and unmarshals it as a slice of codexSSEEvent. +func loadFixtureEvents(t *testing.T, filename string) []codexSSEEvent { + t.Helper() + data, err := os.ReadFile("testdata/" + filename) + if err != nil { + t.Fatalf("read fixture %q: %v", filename, err) + } + var events []codexSSEEvent + if err := json.Unmarshal(data, &events); err != nil { + t.Fatalf("unmarshal fixture %q: %v", filename, err) + } + return events +} + +// serveFixtureAsSSE creates an httptest.Server that streams the given events +// as SSE data frames, followed by [DONE]. +func serveFixtureAsSSE(t *testing.T, events []codexSSEEvent) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + flusher, ok := w.(http.Flusher) + if !ok { + t.Error("ResponseWriter does not implement http.Flusher") + return + } + for _, ev := range events { + b, err := json.Marshal(ev) + if err != nil { + t.Errorf("marshal event: %v", err) + return + } + fmt.Fprintf(w, "data: %s\n\n", b) + flusher.Flush() + } + fmt.Fprint(w, "data: [DONE]\n\n") + flusher.Flush() + })) +} + +// TestCodexImagePartialThenDone verifies: +// 1. A partial_image event stores the partial in imageState and emits a Partial=true chunk. +// 2. The subsequent output_item.done (image_generation_call) records the final image. +// 3. ChatResponse.Images contains exactly one image with correct MIME type. +// 4. The base64 in ChatResponse.Images[0].Data decodes to a valid PNG. +func TestCodexImagePartialThenDone(t *testing.T) { + events := loadFixtureEvents(t, "codex_native_image_partial_then_done.json") + server := serveFixtureAsSSE(t, events) + defer server.Close() + + p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-image-2") + p.retryConfig.Attempts = 1 + + var partialChunks []ImageContent + var finalChunks []ImageContent + + result, err := p.ChatStream(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "Draw a cat"}}, + }, func(chunk StreamChunk) { + for _, img := range chunk.Images { + if img.Partial { + partialChunks = append(partialChunks, img) + } else { + finalChunks = append(finalChunks, img) + } + } + }) + if err != nil { + t.Fatalf("ChatStream: %v", err) + } + + // Streaming: one partial chunk emitted. + if len(partialChunks) != 1 { + t.Errorf("partial chunks = %d, want 1", len(partialChunks)) + } else { + if partialChunks[0].MimeType != "image/png" { + t.Errorf("partial chunk MimeType = %q, want image/png", partialChunks[0].MimeType) + } + decodeBase64PNG(t, partialChunks[0].Data) + } + + // Streaming: one final (non-partial) chunk emitted. + if len(finalChunks) != 1 { + t.Errorf("final chunks = %d, want 1", len(finalChunks)) + } else { + if finalChunks[0].MimeType != "image/png" { + t.Errorf("final chunk MimeType = %q, want image/png", finalChunks[0].MimeType) + } + decodeBase64PNG(t, finalChunks[0].Data) + } + + // ChatResponse.Images: exactly one entry (deduplicated; not double-counted from response.completed). + if len(result.Images) != 1 { + t.Fatalf("result.Images length = %d, want 1", len(result.Images)) + } + if result.Images[0].MimeType != "image/png" { + t.Errorf("result.Images[0].MimeType = %q, want image/png", result.Images[0].MimeType) + } + decodeBase64PNG(t, result.Images[0].Data) +} + +// TestCodexImageNonStream verifies that a single response.completed event +// containing image_generation_call items in output[] populates ChatResponse.Images. +func TestCodexImageNonStream(t *testing.T) { + events := loadFixtureEvents(t, "codex_native_image_non_stream.json") + server := serveFixtureAsSSE(t, events) + defer server.Close() + + p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-image-2") + p.retryConfig.Attempts = 1 + + result, err := p.Chat(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "Draw a landscape"}}, + }) + if err != nil { + t.Fatalf("Chat: %v", err) + } + + // Text content from the message item is also captured. + if result.Content != "Here is your image." { + t.Errorf("Content = %q, want 'Here is your image.'", result.Content) + } + + if len(result.Images) != 1 { + t.Fatalf("result.Images length = %d, want 1", len(result.Images)) + } + if result.Images[0].MimeType != "image/png" { + t.Errorf("result.Images[0].MimeType = %q, want image/png", result.Images[0].MimeType) + } + decodeBase64PNG(t, result.Images[0].Data) + + // Usage captured from non-stream response.completed. + if result.Usage == nil { + t.Fatal("Usage is nil") + } + if result.Usage.TotalTokens != 11 { + t.Errorf("TotalTokens = %d, want 11", result.Usage.TotalTokens) + } +} + +// TestCodexImageDuplicatePartialDedup verifies that two identical partial_image +// events for the same item_id do not emit a second chunk (SHA256 dedup). +func TestCodexImageDuplicatePartialDedup(t *testing.T) { + const b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + flusher := w.(http.Flusher) + // Two identical partial frames. + for range 2 { + ev := codexSSEEvent{ + Type: "response.image_generation_call.partial_image", + ItemID: "ig_dedup", + OutputFormat: "png", + PartialImageB64: b64, + PartialImageIndex: 0, + } + b, _ := json.Marshal(ev) + fmt.Fprintf(w, "data: %s\n\n", b) + flusher.Flush() + } + // Final done. + done := codexSSEEvent{ + Type: "response.output_item.done", + Item: &codexItem{ + ID: "ig_dedup", + Type: "image_generation_call", + OutputFormat: "png", + Result: b64, + }, + } + db, _ := json.Marshal(done) + fmt.Fprintf(w, "data: %s\n\n", db) + flusher.Flush() + + completed := codexSSEEvent{ + Type: "response.completed", + Response: &codexAPIResponse{ + ID: "resp_dedup", + Status: "completed", + Usage: &codexUsage{InputTokens: 1, OutputTokens: 1, TotalTokens: 2}, + }, + } + cb, _ := json.Marshal(completed) + fmt.Fprintf(w, "data: %s\n\n", cb) + flusher.Flush() + + fmt.Fprint(w, "data: [DONE]\n\n") + flusher.Flush() + })) + defer server.Close() + + p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-image-2") + p.retryConfig.Attempts = 1 + + var imageChunks []ImageContent + result, err := p.ChatStream(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "Draw"}}, + }, func(chunk StreamChunk) { + imageChunks = append(imageChunks, chunk.Images...) + }) + if err != nil { + t.Fatalf("ChatStream: %v", err) + } + + // Only 1 partial chunk (dedup skips the second identical frame) + 1 final chunk. + partialCount := 0 + finalCount := 0 + for _, img := range imageChunks { + if img.Partial { + partialCount++ + } else { + finalCount++ + } + } + if partialCount != 1 { + t.Errorf("partial chunks emitted = %d, want 1 (duplicate suppressed)", partialCount) + } + if finalCount != 1 { + t.Errorf("final chunks emitted = %d, want 1", finalCount) + } + + // ChatResponse.Images: exactly one image. + if len(result.Images) != 1 { + t.Errorf("result.Images length = %d, want 1", len(result.Images)) + } +} + +// TestCodexImageMixedTextAndImage verifies that text content, a function tool_call, +// and an image_generation_call in the same response are all preserved correctly. +func TestCodexImageMixedTextAndImage(t *testing.T) { + const imgB64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + flusher := w.(http.Flusher) + + events := []codexSSEEvent{ + // Text delta. + {Type: "response.output_text.delta", ItemID: "msg_1", Delta: "Here is "}, + {Type: "response.output_text.delta", ItemID: "msg_1", Delta: "the result."}, + // Message item done. + { + Type: "response.output_item.done", + Item: &codexItem{ID: "msg_1", Type: "message", Role: "assistant"}, + }, + // Function call done. + { + Type: "response.output_item.done", + Item: &codexItem{ + ID: "fc_1", + Type: "function_call", + CallID: "call_abc", + Name: "web_search", + Arguments: `{"query":"cats"}`, + }, + }, + // Image generation call done. + { + Type: "response.output_item.done", + Item: &codexItem{ + ID: "ig_1", + Type: "image_generation_call", + OutputFormat: "png", + Result: imgB64, + }, + }, + // Completion. + { + Type: "response.completed", + Response: &codexAPIResponse{ + ID: "resp_mixed", + Status: "completed", + Usage: &codexUsage{InputTokens: 10, OutputTokens: 20, TotalTokens: 30}, + }, + }, + } + + for _, ev := range events { + b, _ := json.Marshal(ev) + fmt.Fprintf(w, "data: %s\n\n", b) + flusher.Flush() + } + fmt.Fprint(w, "data: [DONE]\n\n") + flusher.Flush() + })) + defer server.Close() + + p := NewCodexProvider("test", &staticTokenSource{token: "test"}, server.URL, "gpt-4o") + p.retryConfig.Attempts = 1 + + result, err := p.Chat(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "Search and draw"}}, + }) + if err != nil { + t.Fatalf("Chat: %v", err) + } + + // Text preserved. + if result.Content != "Here is the result." { + t.Errorf("Content = %q, want 'Here is the result.'", result.Content) + } + + // Function tool call preserved. + if result.FinishReason != "tool_calls" { + t.Errorf("FinishReason = %q, want tool_calls", result.FinishReason) + } + if len(result.ToolCalls) != 1 { + t.Fatalf("ToolCalls length = %d, want 1", len(result.ToolCalls)) + } + if result.ToolCalls[0].Name != "web_search" { + t.Errorf("ToolCalls[0].Name = %q, want web_search", result.ToolCalls[0].Name) + } + + // Image preserved. + if len(result.Images) != 1 { + t.Fatalf("result.Images length = %d, want 1", len(result.Images)) + } + if result.Images[0].MimeType != "image/png" { + t.Errorf("result.Images[0].MimeType = %q, want image/png", result.Images[0].MimeType) + } + decodeBase64PNG(t, result.Images[0].Data) +} + +// TestCodexMimeFromFormat verifies the mimeFromFormat helper covers all documented formats. +func TestCodexMimeFromFormat(t *testing.T) { + cases := []struct { + format string + want string + }{ + {"png", "image/png"}, + {"jpg", "image/jpeg"}, + {"jpeg", "image/jpeg"}, + {"webp", "image/webp"}, + {"", "image/png"}, + {"unknown", "image/png"}, + } + for _, tc := range cases { + got := mimeFromFormat(tc.format) + if got != tc.want { + t.Errorf("mimeFromFormat(%q) = %q, want %q", tc.format, got, tc.want) + } + } +} diff --git a/internal/providers/codex_native_image.go b/internal/providers/codex_native_image.go new file mode 100644 index 00000000..ec5fb507 --- /dev/null +++ b/internal/providers/codex_native_image.go @@ -0,0 +1,211 @@ +package providers + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" +) + +// GenerateImage implements NativeImageProvider for CodexProvider. +// Sends a minimal POST /codex/responses request with an image_generation tool +// and tool_choice forced to image_generation. Returns decoded image bytes. +func (p *CodexProvider) GenerateImage(ctx context.Context, req NativeImageRequest) (*NativeImageResult, error) { + if req.Prompt == "" { + return nil, fmt.Errorf("codex native image: prompt is required") + } + if req.OutputFormat == "" { + req.OutputFormat = "png" + } + if req.AspectRatio == "" { + req.AspectRatio = "1:1" + } + + model := req.Model + if model == "" { + model = p.defaultModel + } + + imageModel, err := ValidateImageModel(req.ImageModel) + if err != nil { + return nil, err + } + req.ImageModel = imageModel + + body := p.buildNativeImageRequestBody(model, req) + + respBody, err := RetryDo(ctx, p.retryConfig, func() (io.ReadCloser, error) { + return p.doRequest(ctx, body) + }) + if err != nil { + return nil, fmt.Errorf("codex native image: request failed: %w", err) + } + defer respBody.Close() + + raw, err := io.ReadAll(respBody) + if err != nil { + return nil, fmt.Errorf("codex native image: read response: %w", err) + } + + return parseNativeImageResponse(raw) +} + +// buildNativeImageRequestBody constructs the minimal Responses API body for image generation. +// The Responses API rejects non-streaming requests with HTTP 400 "Stream must be set to true", +// so stream is always true. Final assembly happens in parseNativeImageSSE which scans the +// event stream for response.output_item.done (image item) or response.completed output walk. +func (p *CodexProvider) buildNativeImageRequestBody(model string, req NativeImageRequest) map[string]any { + return map[string]any{ + "model": model, + "stream": true, + "store": false, + "instructions": "Generate an image matching the user's description using the image_generation tool. Return only the image; do not describe it in text.", + "input": []any{ + map[string]any{ + "role": "user", + "content": []map[string]any{ + {"type": "input_text", "text": req.Prompt}, + }, + }, + }, + "tools": []map[string]any{ + { + "type": "image_generation", + "action": "generate", + "model": req.ImageModel, + "output_format": req.OutputFormat, + "size": SizeFromAspect(req.AspectRatio), + }, + }, + "tool_choice": map[string]any{ + "type": "image_generation", + }, + } +} + +// parseNativeImageResponse extracts base64-encoded image bytes from a Responses API +// non-streaming body (single JSON object). Walks output[] for type == "image_generation_call". +func parseNativeImageResponse(data []byte) (*NativeImageResult, error) { + // Non-streaming path returns a raw JSON object (not SSE lines). + // If the response looks like SSE (starts with "data:"), fall back to SSE parse. + trimmed := bytes.TrimSpace(data) + if len(trimmed) > 0 && trimmed[0] != '{' { + return parseNativeImageSSE(data) + } + + var resp codexAPIResponse + if err := json.Unmarshal(data, &resp); err != nil { + return nil, fmt.Errorf("codex native image: decode response: %w", err) + } + + if resp.Error != nil { + msg := resp.Error.Message + if msg == "" { + msg = resp.Error.Code + } + return nil, fmt.Errorf("codex native image: API error: %s", msg) + } + + for i := range resp.Output { + item := &resp.Output[i] + if item.Type == "image_generation_call" && item.Result != "" { + raw, err := base64.StdEncoding.DecodeString(item.Result) + if err != nil { + return nil, fmt.Errorf("codex native image: decode base64: %w", err) + } + mime := mimeFromFormat(item.OutputFormat) + var usage *Usage + if resp.Usage != nil { + usage = &Usage{ + PromptTokens: resp.Usage.InputTokens, + CompletionTokens: resp.Usage.OutputTokens, + TotalTokens: resp.Usage.TotalTokens, + } + } + return &NativeImageResult{MimeType: mime, Data: raw, Usage: usage}, nil + } + } + + return nil, fmt.Errorf("codex native image: no image_generation_call in response output") +} + +// parseNativeImageSSE parses SSE-streamed lines when the server unexpectedly returns +// a stream despite stream:false. Looks for response.completed or output_item.done events. +func parseNativeImageSSE(data []byte) (*NativeImageResult, error) { + // Scan lines for "data: {...}" frames. + var b64 string + var outputFormat string + var usage *Usage + + for _, line := range bytes.Split(data, []byte("\n")) { + if !bytes.HasPrefix(line, []byte("data: ")) { + continue + } + payload := line[len("data: "):] + if bytes.Equal(payload, []byte("[DONE]")) { + break + } + + var event codexSSEEvent + if err := json.Unmarshal(payload, &event); err != nil { + continue + } + + switch event.Type { + case "response.output_item.done": + if event.Item != nil && event.Item.Type == "image_generation_call" && event.Item.Result != "" { + b64 = event.Item.Result + outputFormat = event.Item.OutputFormat + } + case "response.completed": + if event.Response != nil { + for i := range event.Response.Output { + item := &event.Response.Output[i] + if item.Type == "image_generation_call" && item.Result != "" { + b64 = item.Result + outputFormat = item.OutputFormat + } + } + if event.Response.Usage != nil { + u := event.Response.Usage + usage = &Usage{ + PromptTokens: u.InputTokens, + CompletionTokens: u.OutputTokens, + TotalTokens: u.TotalTokens, + } + } + } + } + } + + if b64 == "" { + return nil, fmt.Errorf("codex native image: no image in SSE stream") + } + + raw, err := base64.StdEncoding.DecodeString(b64) + if err != nil { + return nil, fmt.Errorf("codex native image: decode base64 from SSE: %w", err) + } + + return &NativeImageResult{ + MimeType: mimeFromFormat(outputFormat), + Data: raw, + Usage: usage, + }, nil +} + +// GenerateImage implements NativeImageProvider for CodexAdapter. +// Delegates to a temporary CodexProvider using the adapter's credentials. +func (a *CodexAdapter) GenerateImage(ctx context.Context, req NativeImageRequest) (*NativeImageResult, error) { + p := &CodexProvider{ + name: "codex", + apiBase: a.apiBase, + defaultModel: a.defaultModel, + client: NewDefaultHTTPClient(), + retryConfig: DefaultRetryConfig(), + tokenSource: a.tokenSource, + } + return p.GenerateImage(ctx, req) +} diff --git a/internal/providers/codex_native_image_test.go b/internal/providers/codex_native_image_test.go new file mode 100644 index 00000000..82d8d040 --- /dev/null +++ b/internal/providers/codex_native_image_test.go @@ -0,0 +1,346 @@ +package providers + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// minimalPNGForProviders is a 1x1 transparent PNG in base64 used by native image tests. +const minimalPNGForProviders = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + +// mockImageServer returns a test server that captures request bodies and returns a +// minimal successful image generation response. The captured pointer is written on +// each request. +func mockImageServer(t *testing.T, captured *[]byte) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read body: %v", err) + http.Error(w, "read error", http.StatusInternalServerError) + return + } + *captured = body + + resp := map[string]any{ + "id": "resp_test", + "status": "completed", + "output": []map[string]any{ + { + "type": "image_generation_call", + "result": minimalPNGForProviders, + "output_format": "png", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(resp); err != nil { + t.Errorf("encode response: %v", err) + } + })) +} + +// TestCodexGenerateImage_BuildsNativeRequest verifies that GenerateImage sends the +// correct JSON body to the Responses API: model, stream:false, input, tools, and +// tool_choice. The test captures the raw request body from a mock server and +// asserts each required field is present and well-formed. +// +// Sub-cases: +// - Default (empty ImageModel) → tools[0].model == "gpt-image-2" +// - Legacy (ImageModel: "gpt-image-1.5") → tools[0].model == "gpt-image-1.5" +// - Rejected (ImageModel: "dall-e-3") → GenerateImage returns error containing "unsupported image model" +func TestCodexGenerateImage_BuildsNativeRequest(t *testing.T) { + var captured []byte + server := mockImageServer(t, &captured) + defer server.Close() + + p := NewCodexProvider("codex-test", &staticTokenSource{token: "tok"}, server.URL, "gpt-image-2") + p.retryConfig.Attempts = 1 + + req := NativeImageRequest{ + Model: "gpt-image-2", + Prompt: "A red circle on a white background", + AspectRatio: "16:9", + OutputFormat: "png", + } + result, err := p.GenerateImage(context.Background(), req) + if err != nil { + t.Fatalf("GenerateImage returned error: %v", err) + } + if len(result.Data) == 0 { + t.Fatal("GenerateImage returned empty Data") + } + + // Verify outbound request body shape. + var body map[string]any + if err := json.Unmarshal(captured, &body); err != nil { + t.Fatalf("unmarshal captured body: %v", err) + } + + // model field (outer Responses API model, not image model) + if model, _ := body["model"].(string); model != "gpt-image-2" { + t.Errorf("body[model] = %q, want %q", model, "gpt-image-2") + } + + // Responses API requires stream:true — non-streaming requests are rejected with + // HTTP 400 "Stream must be set to true". Final image is assembled from SSE events. + if stream, _ := body["stream"].(bool); !stream { + t.Error("body[stream] must be true (Responses API rejects stream:false)") + } + + // instructions is required by Responses API — must be non-empty. + if instr, _ := body["instructions"].(string); instr == "" { + t.Error("body[instructions] must be non-empty (Responses API rejects requests without instructions)") + } + + // input must be an array with one user message + inputs, ok := body["input"].([]any) + if !ok || len(inputs) != 1 { + t.Fatalf("body[input]: expected []any length 1, got %T len %d", body["input"], len(inputs)) + } + userMsg, ok := inputs[0].(map[string]any) + if !ok { + t.Fatalf("input[0] is not a map: %T", inputs[0]) + } + if role, _ := userMsg["role"].(string); role != "user" { + t.Errorf("input[0].role = %q, want %q", role, "user") + } + contents, ok := userMsg["content"].([]any) + if !ok || len(contents) != 1 { + t.Fatalf("input[0].content: expected []any length 1, got %T len %d", userMsg["content"], len(contents)) + } + contentPart, ok := contents[0].(map[string]any) + if !ok { + t.Fatalf("content[0] is not a map: %T", contents[0]) + } + if typ, _ := contentPart["type"].(string); typ != "input_text" { + t.Errorf("content[0].type = %q, want %q", typ, "input_text") + } + if text, _ := contentPart["text"].(string); text != req.Prompt { + t.Errorf("content[0].text = %q, want %q", text, req.Prompt) + } + + // tools must contain one image_generation entry + tools, ok := body["tools"].([]any) + if !ok || len(tools) != 1 { + t.Fatalf("body[tools]: expected []any length 1, got %T len %d", body["tools"], len(tools)) + } + tool, ok := tools[0].(map[string]any) + if !ok { + t.Fatalf("tools[0] is not a map: %T", tools[0]) + } + if typ, _ := tool["type"].(string); typ != "image_generation" { + t.Errorf("tools[0].type = %q, want %q", typ, "image_generation") + } + // size should map to 1792x1024 for 16:9 + wantSize := SizeFromAspect("16:9") + if size, _ := tool["size"].(string); size != wantSize { + t.Errorf("tools[0].size = %q, want %q", size, wantSize) + } + if fmt.Sprint(tool["output_format"]) != "png" { + t.Errorf("tools[0].output_format = %v, want png", tool["output_format"]) + } + // tools[0].model must be gpt-image-2 (default when ImageModel is empty) + if imgModel, _ := tool["model"].(string); imgModel != DefaultImageModel { + t.Errorf("tools[0].model = %q, want %q (default)", imgModel, DefaultImageModel) + } + + // tool_choice must force image_generation + toolChoice, ok := body["tool_choice"].(map[string]any) + if !ok { + t.Fatalf("body[tool_choice] is not a map: %T", body["tool_choice"]) + } + if typ, _ := toolChoice["type"].(string); typ != "image_generation" { + t.Errorf("tool_choice.type = %q, want %q", typ, "image_generation") + } +} + +// TestCodexGenerateImage_ImageModelDefault verifies that an empty ImageModel results +// in the default gpt-image-2 model in the outbound tools[0].model field. +func TestCodexGenerateImage_ImageModelDefault(t *testing.T) { + var captured []byte + server := mockImageServer(t, &captured) + defer server.Close() + + p := NewCodexProvider("codex-test", &staticTokenSource{token: "tok"}, server.URL, "gpt-image-2") + p.retryConfig.Attempts = 1 + + _, err := p.GenerateImage(context.Background(), NativeImageRequest{ + Prompt: "test", + ImageModel: "", // explicitly empty — should default to gpt-image-2 + AspectRatio: "1:1", + }) + if err != nil { + t.Fatalf("GenerateImage returned error: %v", err) + } + + var body map[string]any + if err := json.Unmarshal(captured, &body); err != nil { + t.Fatalf("unmarshal captured body: %v", err) + } + tools, _ := body["tools"].([]any) + if len(tools) == 0 { + t.Fatal("tools array is empty") + } + tool, _ := tools[0].(map[string]any) + if imgModel, _ := tool["model"].(string); imgModel != "gpt-image-2" { + t.Errorf("tools[0].model = %q, want gpt-image-2 (default)", imgModel) + } +} + +// TestCodexGenerateImage_ImageModelLegacy verifies that ImageModel "gpt-image-1.5" +// is forwarded to the outbound tools[0].model field. +func TestCodexGenerateImage_ImageModelLegacy(t *testing.T) { + var captured []byte + server := mockImageServer(t, &captured) + defer server.Close() + + p := NewCodexProvider("codex-test", &staticTokenSource{token: "tok"}, server.URL, "gpt-image-2") + p.retryConfig.Attempts = 1 + + _, err := p.GenerateImage(context.Background(), NativeImageRequest{ + Prompt: "test", + ImageModel: "gpt-image-1.5", + AspectRatio: "1:1", + }) + if err != nil { + t.Fatalf("GenerateImage returned error: %v", err) + } + + var body map[string]any + if err := json.Unmarshal(captured, &body); err != nil { + t.Fatalf("unmarshal captured body: %v", err) + } + tools, _ := body["tools"].([]any) + if len(tools) == 0 { + t.Fatal("tools array is empty") + } + tool, _ := tools[0].(map[string]any) + if imgModel, _ := tool["model"].(string); imgModel != "gpt-image-1.5" { + t.Errorf("tools[0].model = %q, want gpt-image-1.5 (legacy)", imgModel) + } +} + +// TestCodexGenerateImage_ImageModelRejected verifies that an unsupported image model +// causes GenerateImage to return an error containing "unsupported image model" before +// making any HTTP request. +func TestCodexGenerateImage_ImageModelRejected(t *testing.T) { + requestMade := false + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestMade = true + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + p := NewCodexProvider("codex-test", &staticTokenSource{token: "tok"}, server.URL, "gpt-image-2") + p.retryConfig.Attempts = 1 + + _, err := p.GenerateImage(context.Background(), NativeImageRequest{ + Prompt: "test", + ImageModel: "dall-e-3", + }) + if err == nil { + t.Fatal("expected error for unsupported image model, got nil") + } + if !strings.Contains(err.Error(), "unsupported image model") { + t.Errorf("error %q does not contain 'unsupported image model'", err.Error()) + } + if requestMade { + t.Error("HTTP request was made despite invalid image model (should have been rejected before the request)") + } +} + +// TestCodexGenerateImage_SSEFallback verifies that GenerateImage correctly parses +// an SSE-format response when the server returns streamed lines instead of a JSON blob. +func TestCodexGenerateImage_SSEFallback(t *testing.T) { + imgB64 := minimalPNGForProviders + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + // Emit a response.completed SSE event with an image_generation_call. + ev := codexSSEEvent{ + Type: "response.completed", + Response: &codexAPIResponse{ + ID: "resp_sse", + Status: "completed", + Output: []codexItem{ + { + ID: "ig_1", + Type: "image_generation_call", + OutputFormat: "png", + Result: imgB64, + }, + }, + Usage: &codexUsage{InputTokens: 5, OutputTokens: 5, TotalTokens: 10}, + }, + } + b, _ := json.Marshal(ev) + fmt.Fprintf(w, "data: %s\n\n", b) + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer server.Close() + + p := NewCodexProvider("codex-test", &staticTokenSource{token: "tok"}, server.URL, "gpt-image-2") + p.retryConfig.Attempts = 1 + + result, err := p.GenerateImage(context.Background(), NativeImageRequest{ + Prompt: "A blue square", + OutputFormat: "png", + }) + if err != nil { + t.Fatalf("GenerateImage SSE fallback: %v", err) + } + if result.MimeType != "image/png" { + t.Errorf("MimeType = %q, want image/png", result.MimeType) + } + want, _ := base64.StdEncoding.DecodeString(imgB64) + if len(result.Data) != len(want) { + t.Errorf("Data length = %d, want %d", len(result.Data), len(want)) + } + if result.Usage == nil { + t.Error("Usage is nil") + } else if result.Usage.TotalTokens != 10 { + t.Errorf("Usage.TotalTokens = %d, want 10", result.Usage.TotalTokens) + } +} + +// TestCodexGenerateImage_NoPrompt verifies that an empty prompt returns an error +// before making any HTTP request. +func TestCodexGenerateImage_NoPrompt(t *testing.T) { + p := NewCodexProvider("codex-test", &staticTokenSource{token: "tok"}, "http://localhost", "gpt-image-2") + p.retryConfig.Attempts = 1 + + _, err := p.GenerateImage(context.Background(), NativeImageRequest{Prompt: ""}) + if err == nil { + t.Fatal("expected error for empty prompt, got nil") + } +} + +// TestSizeFromAspect verifies the aspect ratio → pixel dimension mapping. +func TestSizeFromAspect(t *testing.T) { + cases := []struct { + ratio string + want string + }{ + {"1:1", "1024x1024"}, + {"16:9", "1792x1024"}, + {"9:16", "1024x1792"}, + {"4:3", "1365x1024"}, + {"3:4", "1024x1365"}, + {"", "1024x1024"}, + {"custom", "1024x1024"}, + } + for _, tc := range cases { + got := SizeFromAspect(tc.ratio) + if got != tc.want { + t.Errorf("SizeFromAspect(%q) = %q, want %q", tc.ratio, got, tc.want) + } + } +} diff --git a/internal/providers/codex_stream_state.go b/internal/providers/codex_stream_state.go index afadfeea..3dd762d5 100644 --- a/internal/providers/codex_stream_state.go +++ b/internal/providers/codex_stream_state.go @@ -1,6 +1,7 @@ package providers import ( + "crypto/sha256" "fmt" "sort" "strings" @@ -237,3 +238,92 @@ func appendCodexContent(result *ChatResponse, text string, onChunk func(StreamCh onChunk(StreamChunk{Content: text}) } } + +// codexImageAccum tracks streaming image generation items keyed by item_id. +// It stores the last-seen partial frame (deduplicated via SHA256) and the final result. +// Not thread-safe — designed for sequential SSE processing in a single goroutine. +type codexImageAccum struct { + outputFormat string + lastPartialHash [sha256.Size]byte // SHA256 of last partial_image_b64; zeroed if none + hasPartial bool // true once at least one partial is recorded + finalB64 string // filled on response.output_item.done or response.completed +} + +// codexImageState tracks all image accumulators for a single streaming response. +type codexImageState struct { + items map[string]*codexImageAccum // keyed by item_id + insertOrder []string // preserves emission order for final assembly +} + +func newCodexImageState() *codexImageState { + return &codexImageState{items: make(map[string]*codexImageAccum)} +} + +func (s *codexImageState) ensureItem(itemID, outputFormat string) *codexImageAccum { + if acc, ok := s.items[itemID]; ok { + return acc + } + acc := &codexImageAccum{outputFormat: outputFormat} + s.items[itemID] = acc + s.insertOrder = append(s.insertOrder, itemID) + return acc +} + +// recordPartial stores a partial frame, deduplicating by SHA256. +// Returns true if the frame is new (not a duplicate) and was recorded. +func (s *codexImageState) recordPartial(itemID, outputFormat, b64 string) bool { + if b64 == "" { + return false + } + acc := s.ensureItem(itemID, outputFormat) + if outputFormat != "" { + acc.outputFormat = outputFormat + } + h := sha256.Sum256([]byte(b64)) + if acc.hasPartial && acc.lastPartialHash == h { + return false // duplicate frame + } + acc.lastPartialHash = h + acc.hasPartial = true + return true +} + +// recordFinal stores the final base64 image for an item. +func (s *codexImageState) recordFinal(itemID, outputFormat, b64 string) { + if b64 == "" { + return + } + acc := s.ensureItem(itemID, outputFormat) + if outputFormat != "" { + acc.outputFormat = outputFormat + } + acc.finalB64 = b64 +} + +// appendToResponse appends all completed images (those with a final) to result.Images +// in insertion order. Deduplication by item_id is implicit (each item appears once). +func (s *codexImageState) appendToResponse(result *ChatResponse) { + for _, id := range s.insertOrder { + acc := s.items[id] + if acc.finalB64 == "" { + continue + } + result.Images = append(result.Images, ImageContent{ + MimeType: mimeFromFormat(acc.outputFormat), + Data: acc.finalB64, + }) + } +} + +// mimeFromFormat converts an output_format string to a MIME type. +// Defaults to "image/png" for unknown or empty formats. +func mimeFromFormat(format string) string { + switch format { + case "jpg", "jpeg": + return "image/jpeg" + case "webp": + return "image/webp" + default: + return "image/png" + } +} diff --git a/internal/providers/codex_test.go b/internal/providers/codex_test.go index 2393623f..0326a0fb 100644 --- a/internal/providers/codex_test.go +++ b/internal/providers/codex_test.go @@ -121,7 +121,8 @@ func TestCodexProviderBuildRequestBodyWithTools(t *testing.T) { Messages: []Message{{Role: "user", Content: "What's the weather?"}}, Tools: []ToolDefinition{ { - Function: ToolFunctionSchema{ + Type: "function", + Function: &ToolFunctionSchema{ Name: "get_weather", Description: "Get current weather", Parameters: map[string]any{ @@ -932,6 +933,114 @@ func TestCodexProviderTokenSource(t *testing.T) { } } +// TestCodexBuildRequestBodyImageGenerationTool verifies that a ToolDefinition with +// Type="image_generation" produces a native image_generation object in the tools array, +// not a function-shaped entry. +func TestCodexBuildRequestBodyImageGenerationTool(t *testing.T) { + p := NewCodexProvider("test", &staticTokenSource{token: "test"}, "", "gpt-4o") + + req := ChatRequest{ + Messages: []Message{{Role: "user", Content: "Draw a cat"}}, + Tools: []ToolDefinition{ + {Type: "image_generation"}, + }, + } + + body := p.buildRequestBody(req, false) + + tools, ok := body["tools"].([]map[string]any) + if !ok { + t.Fatalf("tools is not []map[string]any: %T", body["tools"]) + } + if len(tools) != 1 { + t.Fatalf("tools length = %d, want 1", len(tools)) + } + + tool := tools[0] + if tool["type"] != "image_generation" { + t.Errorf("tool[type] = %v, want image_generation", tool["type"]) + } + if tool["action"] != "generate" { + t.Errorf("tool[action] = %v, want generate", tool["action"]) + } + if tool["model"] != "gpt-image-2" { + t.Errorf("tool[model] = %v, want gpt-image-2", tool["model"]) + } + if tool["output_format"] != "png" { + t.Errorf("tool[output_format] = %v, want png", tool["output_format"]) + } + if tool["partial_images"] != 1 { + t.Errorf("tool[partial_images] = %v, want 1", tool["partial_images"]) + } + // Must NOT contain function-specific fields. + if _, has := tool["name"]; has { + t.Error("image_generation tool must not have 'name' field") + } + if _, has := tool["parameters"]; has { + t.Error("image_generation tool must not have 'parameters' field") + } +} + +// TestCodexBuildRequestBodyMixedTools verifies that a request containing both function +// tools and an image_generation tool produces both in the correct shapes, in order. +func TestCodexBuildRequestBodyMixedTools(t *testing.T) { + p := NewCodexProvider("test", &staticTokenSource{token: "test"}, "", "gpt-4o") + + req := ChatRequest{ + Messages: []Message{{Role: "user", Content: "Search and draw"}}, + Tools: []ToolDefinition{ + { + Type: "function", + Function: &ToolFunctionSchema{ + Name: "web_search", + Description: "Search the web", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{"query": map[string]any{"type": "string"}}, + "required": []string{"query"}, + }, + }, + }, + {Type: "image_generation"}, + }, + } + + body := p.buildRequestBody(req, false) + + tools, ok := body["tools"].([]map[string]any) + if !ok { + t.Fatalf("tools is not []map[string]any: %T", body["tools"]) + } + if len(tools) != 2 { + t.Fatalf("tools length = %d, want 2", len(tools)) + } + + // First tool: function shape. + fn := tools[0] + if fn["type"] != "function" { + t.Errorf("tools[0] type = %v, want function", fn["type"]) + } + if fn["name"] != "web_search" { + t.Errorf("tools[0] name = %v, want web_search", fn["name"]) + } + + // Second tool: native image_generation shape. + img := tools[1] + if img["type"] != "image_generation" { + t.Errorf("tools[1] type = %v, want image_generation", img["type"]) + } + if img["action"] != "generate" { + t.Errorf("tools[1] action = %v, want generate", img["action"]) + } + if img["model"] != "gpt-image-2" { + t.Errorf("tools[1] model = %v, want gpt-image-2", img["model"]) + } + // Function field must not bleed into image tool. + if _, has := img["name"]; has { + t.Error("image_generation tool must not contain 'name'") + } +} + // Verify request body includes image content func TestCodexProviderBuildRequestBodyWithImages(t *testing.T) { p := NewCodexProvider("test", &staticTokenSource{token: "test"}, "", "gpt-4o") diff --git a/internal/providers/codex_types.go b/internal/providers/codex_types.go index 1ca09e24..dc778f37 100644 --- a/internal/providers/codex_types.go +++ b/internal/providers/codex_types.go @@ -18,15 +18,17 @@ type codexErrorDetail struct { } type codexItem struct { - ID string `json:"id"` - Type string `json:"type"` // "message", "function_call", "reasoning" - Role string `json:"role,omitempty"` - Phase string `json:"phase,omitempty"` // gpt-5.3-codex: "commentary" or "final_answer" - Content []codexContent `json:"content,omitempty"` - CallID string `json:"call_id,omitempty"` - Name string `json:"name,omitempty"` - Arguments string `json:"arguments,omitempty"` - Summary []codexSummary `json:"summary,omitempty"` + ID string `json:"id"` + Type string `json:"type"` // "message", "function_call", "reasoning", "image_generation_call" + Role string `json:"role,omitempty"` + Phase string `json:"phase,omitempty"` // gpt-5.3-codex: "commentary" or "final_answer" + Content []codexContent `json:"content,omitempty"` + CallID string `json:"call_id,omitempty"` + Name string `json:"name,omitempty"` + Arguments string `json:"arguments,omitempty"` + Summary []codexSummary `json:"summary,omitempty"` + OutputFormat string `json:"output_format,omitempty"` // image_generation_call: "png", "jpeg", "webp" + Result string `json:"result,omitempty"` // image_generation_call: base64 final image } type codexContent struct { @@ -53,15 +55,18 @@ type codexTokensDetails struct { // SSE streaming types type codexSSEEvent struct { - Type string `json:"type"` - Delta string `json:"delta,omitempty"` - Text string `json:"text,omitempty"` - ItemID string `json:"item_id,omitempty"` - OutputIndex int `json:"output_index,omitempty"` - ContentIndex int `json:"content_index,omitempty"` - Item *codexItem `json:"item,omitempty"` - Part *codexContentPart `json:"part,omitempty"` - Response *codexAPIResponse `json:"response,omitempty"` + Type string `json:"type"` + Delta string `json:"delta,omitempty"` + Text string `json:"text,omitempty"` + ItemID string `json:"item_id,omitempty"` + OutputIndex int `json:"output_index,omitempty"` + ContentIndex int `json:"content_index,omitempty"` + Item *codexItem `json:"item,omitempty"` + Part *codexContentPart `json:"part,omitempty"` + Response *codexAPIResponse `json:"response,omitempty"` + OutputFormat string `json:"output_format,omitempty"` // response.image_generation_call.partial_image + PartialImageB64 string `json:"partial_image_b64,omitempty"` // response.image_generation_call.partial_image + PartialImageIndex int `json:"partial_image_index,omitempty"` // response.image_generation_call.partial_image } type codexToolCallAcc struct { diff --git a/internal/providers/native_image.go b/internal/providers/native_image.go new file mode 100644 index 00000000..d836dfa3 --- /dev/null +++ b/internal/providers/native_image.go @@ -0,0 +1,94 @@ +package providers + +import ( + "context" + "fmt" +) + +// NativeImageProvider is implemented by OAuth-backed providers whose upstream +// exposes an image_generation native tool (ChatGPT Responses API style). +// create_image routes through this interface when the chain resolves to such +// a provider, bypassing the credentialProvider (APIKey/APIBase) path. +type NativeImageProvider interface { + GenerateImage(ctx context.Context, req NativeImageRequest) (*NativeImageResult, error) +} + +// DefaultImageModel is the image model used by the Responses API image_generation +// tool when the caller does not specify one. gpt-image-2 is the current (2026-Q2) +// quality baseline; gpt-image-1.5 is available as a legacy fallback. +const DefaultImageModel = "gpt-image-2" + +// allowedImageModels enumerates the image models the native ChatGPT Responses API +// image_generation tool will accept. Constraining to this whitelist prevents +// silent upstream rejections from arbitrary model names (e.g. "dall-e-3") and +// keeps the PR's motivation — gpt-image-2 quality — as the default everywhere. +var allowedImageModels = map[string]bool{ + "gpt-image-2": true, // default — latest quality + "gpt-image-1.5": true, // legacy fallback +} + +// ValidateImageModel returns the model to use, or an error if the caller +// supplied an unsupported value. Empty input returns DefaultImageModel. +func ValidateImageModel(model string) (string, error) { + if model == "" { + return DefaultImageModel, nil + } + if !allowedImageModels[model] { + return "", fmt.Errorf("unsupported image model %q; allowed: gpt-image-2 (default), gpt-image-1.5 (legacy)", model) + } + return model, nil +} + +// NativeImageRequest describes a single image generation request. +type NativeImageRequest struct { + // Model is the parent LLM model for the Responses API call (e.g. "gpt-5.4"). + // NOT the image model — see ImageModel below. + // If empty, the provider uses its own default LLM model. + Model string + + // ImageModel is the image-generation model attached to the image_generation + // tool (e.g. "gpt-image-2"). Must be a value accepted by ValidateImageModel; + // empty falls back to DefaultImageModel. + ImageModel string + + // Prompt is the text description of the image. + Prompt string + + // AspectRatio is the desired aspect ratio, e.g. "16:9", "1:1", "9:16". + // Converted to a concrete pixel size by the provider implementation. + // Defaults to "1:1" if empty. + AspectRatio string + + // OutputFormat is the desired image format: "png" (default), "jpg", "webp". + OutputFormat string +} + +// NativeImageResult holds the result of a native image generation call. +type NativeImageResult struct { + // MimeType is the detected MIME type of the generated image (e.g. "image/png"). + MimeType string + + // Data is the raw decoded image bytes (NOT base64). + Data []byte + + // Usage is optional token usage if the provider reports it. + Usage *Usage +} + +// SizeFromAspect converts a common aspect ratio string to a pixel dimension +// string expected by image generation APIs (e.g. "1792x1024"). +// Falls back to "1024x1024" for unrecognised ratios. +func SizeFromAspect(aspectRatio string) string { + switch aspectRatio { + case "16:9": + return "1792x1024" + case "9:16": + return "1024x1792" + case "3:4": + return "1024x1365" + case "4:3": + return "1365x1024" + default: + return "1024x1024" + } +} diff --git a/internal/providers/openai_chat.go b/internal/providers/openai_chat.go index 7061b3ed..b3991990 100644 --- a/internal/providers/openai_chat.go +++ b/internal/providers/openai_chat.go @@ -139,6 +139,22 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun } } + // Accumulate images from delta.images[]. + // Each chunk may carry one or more image parts; we collect all into result.Images. + // Malformed data URLs are skipped with a warning — they don't abort the stream. + for _, img := range delta.Images { + mimeType, b64Data, err := parseDataURL(img.ImageURL.URL) + if err != nil { + slog.Warn("openai_stream: skipping malformed image data URL", + "type", img.Type, "url_len", len(img.ImageURL.URL), "error", err) + continue + } + result.Images = append(result.Images, ImageContent{ + MimeType: mimeType, + Data: b64Data, + }) + } + // Accumulate streamed tool calls for _, tc := range delta.ToolCalls { acc, ok := accumulators[tc.Index] diff --git a/internal/providers/openai_compat_image_parse_test.go b/internal/providers/openai_compat_image_parse_test.go new file mode 100644 index 00000000..e7ca49ef --- /dev/null +++ b/internal/providers/openai_compat_image_parse_test.go @@ -0,0 +1,300 @@ +package providers + +import ( + "bufio" + "bytes" + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" +) + +// pngMagic is the 8-byte PNG file signature per the PNG spec. +var pngMagic = []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A} + +// decodeB64 decodes a standard base64 string and fails the test on error. +func decodeB64(t *testing.T, s string) []byte { + t.Helper() + raw, err := base64.StdEncoding.DecodeString(s) + if err != nil { + t.Fatalf("base64 decode failed: %v", err) + } + return raw +} + +// hasPNGMagic checks that raw starts with the 8-byte PNG signature. +func hasPNGMagic(raw []byte) bool { + if len(raw) < 8 { + return false + } + return bytes.Equal(raw[:8], pngMagic) +} + +// --- parseDataURL unit tests --- + +func TestParseDataURL_ValidPNG(t *testing.T) { + b64 := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + url := "data:image/png;base64," + b64 + + mime, got, err := parseDataURL(url) + if err != nil { + t.Fatalf("parseDataURL returned unexpected error: %v", err) + } + if mime != "image/png" { + t.Errorf("mime = %q, want %q", mime, "image/png") + } + if got != b64 { + t.Errorf("b64Data mismatch: got %q, want %q", got, b64) + } + // Verify the decoded bytes start with PNG magic. + raw := decodeB64(t, got) + if !hasPNGMagic(raw) { + t.Errorf("decoded bytes do not have PNG magic: first 8 = %x", raw[:8]) + } +} + +func TestParseDataURL_ValidJPEG(t *testing.T) { + // Minimal JPEG SOI marker (FF D8) + raw := []byte{0xFF, 0xD8, 0x00} + b64 := base64.StdEncoding.EncodeToString(raw) + url := "data:image/jpeg;base64," + b64 + + mime, got, err := parseDataURL(url) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if mime != "image/jpeg" { + t.Errorf("mime = %q, want %q", mime, "image/jpeg") + } + if got != b64 { + t.Errorf("b64Data mismatch") + } +} + +func TestParseDataURL_MalformedNoBase64Prefix(t *testing.T) { + _, _, err := parseDataURL("data:image/png," + "notbase64encoded") + if err == nil { + t.Fatal("expected error for missing ;base64 marker, got nil") + } +} + +func TestParseDataURL_MalformedNotDataURL(t *testing.T) { + _, _, err := parseDataURL("https://example.com/image.png") + if err == nil { + t.Fatal("expected error for non-data URL, got nil") + } +} + +func TestParseDataURL_MalformedInvalidBase64(t *testing.T) { + _, _, err := parseDataURL("data:image/png;base64,!!!not-valid-base64!!!") + if err == nil { + t.Fatal("expected error for invalid base64 payload, got nil") + } +} + +func TestParseDataURL_Empty(t *testing.T) { + _, _, err := parseDataURL("") + if err == nil { + t.Fatal("expected error for empty string, got nil") + } +} + +// --- Non-stream fixture test --- + +func TestOpenAICompatParseResponse_NonStreamImages(t *testing.T) { + fixture, err := os.ReadFile("testdata/openai_compat_image_nonstream.json") + if err != nil { + t.Fatalf("read fixture: %v", err) + } + + var oaiResp openAIResponse + if err := json.Unmarshal(fixture, &oaiResp); err != nil { + t.Fatalf("unmarshal fixture: %v", err) + } + + p := NewOpenAIProvider("test", "key", "https://api.openai.com/v1", "gpt-image-1") + result := p.parseResponse(&oaiResp) + + if len(result.Images) != 1 { + t.Fatalf("Images len = %d, want 1", len(result.Images)) + } + img := result.Images[0] + if img.MimeType != "image/png" { + t.Errorf("MimeType = %q, want %q", img.MimeType, "image/png") + } + raw := decodeB64(t, img.Data) + if !hasPNGMagic(raw) { + t.Errorf("decoded image data does not start with PNG magic: %x", raw[:min8(raw)]) + } + // Content should be preserved alongside images. + if result.Content != "Here is your image." { + t.Errorf("Content = %q, want %q", result.Content, "Here is your image.") + } +} + +// min8 returns the smaller of 8 and len(b) — used for safe slice in error messages. +func min8(b []byte) int { + if len(b) < 8 { + return len(b) + } + return 8 +} + +// --- Stream fixture test --- + +// newStreamServer returns a test HTTP server that serves the SSE fixture file. +func newStreamServer(t *testing.T, fixturePath string) *httptest.Server { + t.Helper() + data, err := os.ReadFile(fixturePath) + if err != nil { + t.Fatalf("read SSE fixture: %v", err) + } + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + // Write line-by-line to simulate chunked SSE. + scanner := bufio.NewScanner(bytes.NewReader(data)) + for scanner.Scan() { + _, _ = io.WriteString(w, scanner.Text()+"\n") + } + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + })) +} + +func TestOpenAICompatChatStream_Images(t *testing.T) { + srv := newStreamServer(t, "testdata/openai_compat_image_stream.sse") + defer srv.Close() + + p := NewOpenAIProvider("test", "key", srv.URL, "gpt-image-1") + + result, err := p.ChatStream(t.Context(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "generate an image"}}, + }, nil) + if err != nil { + t.Fatalf("ChatStream error: %v", err) + } + + // Content accumulated across chunks. + if !strings.Contains(result.Content, "Here") { + t.Errorf("Content = %q, want substring %q", result.Content, "Here") + } + + if len(result.Images) != 1 { + t.Fatalf("Images len = %d, want 1", len(result.Images)) + } + img := result.Images[0] + if img.MimeType != "image/png" { + t.Errorf("MimeType = %q, want %q", img.MimeType, "image/png") + } + raw := decodeB64(t, img.Data) + if !hasPNGMagic(raw) { + t.Errorf("streamed image data does not start with PNG magic: %x", raw[:min8(raw)]) + } +} + +// --- Mixed payload test (content + tool_calls + images all preserved) --- + +func TestOpenAICompatParseResponse_Mixed(t *testing.T) { + fixture, err := os.ReadFile("testdata/openai_compat_image_mixed.json") + if err != nil { + t.Fatalf("read fixture: %v", err) + } + + var oaiResp openAIResponse + if err := json.Unmarshal(fixture, &oaiResp); err != nil { + t.Fatalf("unmarshal fixture: %v", err) + } + + p := NewOpenAIProvider("test", "key", "https://api.openai.com/v1", "gpt-image-1") + result := p.parseResponse(&oaiResp) + + // Content preserved. + if result.Content != "I'll generate that for you." { + t.Errorf("Content = %q, want %q", result.Content, "I'll generate that for you.") + } + // Tool calls preserved. + if len(result.ToolCalls) != 1 { + t.Fatalf("ToolCalls len = %d, want 1", len(result.ToolCalls)) + } + if result.ToolCalls[0].Name != "log_generation" { + t.Errorf("ToolCalls[0].Name = %q, want %q", result.ToolCalls[0].Name, "log_generation") + } + // Images preserved. + if len(result.Images) != 1 { + t.Fatalf("Images len = %d, want 1", len(result.Images)) + } + if result.Images[0].MimeType != "image/png" { + t.Errorf("Images[0].MimeType = %q, want %q", result.Images[0].MimeType, "image/png") + } +} + +// --- Regression: non-image response unaffected --- + +func TestOpenAICompatParseResponse_NoImages_Unchanged(t *testing.T) { + raw := `{ + "choices": [{ + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7} + }` + + var oaiResp openAIResponse + if err := json.Unmarshal([]byte(raw), &oaiResp); err != nil { + t.Fatalf("unmarshal: %v", err) + } + + p := NewOpenAIProvider("test", "key", "https://api.openai.com/v1", "gpt-4o") + result := p.parseResponse(&oaiResp) + + if result.Content != "Hello!" { + t.Errorf("Content = %q, want %q", result.Content, "Hello!") + } + if len(result.Images) != 0 { + t.Errorf("Images len = %d, want 0 for non-image response", len(result.Images)) + } + if result.Usage == nil || result.Usage.TotalTokens != 7 { + t.Errorf("Usage not populated correctly: %+v", result.Usage) + } +} + +// --- Malformed data URL in stream: skipped, stream continues --- + +func TestOpenAICompatChatStream_MalformedImageSkipped(t *testing.T) { + sseData := strings.Join([]string{ + `data: {"choices":[{"delta":{"content":"hi"},"finish_reason":null}]}`, + `data: {"choices":[{"delta":{"images":[{"type":"image_url","image_url":{"url":"data:image/png;base64,!!!INVALID!!!"}}]},"finish_reason":null}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`, + `data: [DONE]`, + ``, + }, "\n") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, sseData) + })) + defer srv.Close() + + p := NewOpenAIProvider("test", "key", srv.URL, "gpt-image-1") + result, err := p.ChatStream(t.Context(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "hi"}}, + }, nil) + // Stream must not error — malformed image is skipped, not fatal. + if err != nil { + t.Fatalf("ChatStream returned error for malformed image: %v", err) + } + if result.Content != "hi" { + t.Errorf("Content = %q, want %q", result.Content, "hi") + } + // Malformed image dropped — no images in result. + if len(result.Images) != 0 { + t.Errorf("Images len = %d, want 0 (malformed entry skipped)", len(result.Images)) + } +} diff --git a/internal/providers/openai_http.go b/internal/providers/openai_http.go index 617f0c0b..80a069a9 100644 --- a/internal/providers/openai_http.go +++ b/internal/providers/openai_http.go @@ -100,6 +100,22 @@ func (p *OpenAIProvider) parseResponse(resp *openAIResponse) *ChatResponse { if len(result.ToolCalls) > 0 && result.FinishReason != "length" { result.FinishReason = "tool_calls" } + + // Decode images[] from the response message into ChatResponse.Images. + // Each entry carries a data URL (data:;base64,). + // Malformed entries are skipped with a warning to avoid crashing on partial responses. + for _, img := range msg.Images { + mimeType, b64Data, err := parseDataURL(img.ImageURL.URL) + if err != nil { + slog.Warn("openai: skipping malformed image data URL", + "type", img.Type, "url_len", len(img.ImageURL.URL), "error", err) + continue + } + result.Images = append(result.Images, ImageContent{ + MimeType: mimeType, + Data: b64Data, + }) + } } if resp.Usage != nil { diff --git a/internal/providers/openai_image_url.go b/internal/providers/openai_image_url.go new file mode 100644 index 00000000..ba1652c7 --- /dev/null +++ b/internal/providers/openai_image_url.go @@ -0,0 +1,36 @@ +package providers + +import ( + "encoding/base64" + "fmt" + "regexp" +) + +// dataURLRe matches data URLs of the form data:;base64,. +var dataURLRe = regexp.MustCompile(`^data:([^;]+);base64,(.+)$`) + +// parseDataURL extracts the MIME type and base64 payload from a data URL. +// It validates that the base64 string is decodable and returns it verbatim +// (not re-encoded) so callers can store it in ImageContent.Data as-is. +// +// Returns an error if: +// - the URL does not match data:;base64, format +// - the base64 payload cannot be decoded +func parseDataURL(s string) (mimeType string, b64Data string, err error) { + m := dataURLRe.FindStringSubmatch(s) + if m == nil { + return "", "", fmt.Errorf("invalid data URL format (expected data:;base64,)") + } + mimeType = m[1] + b64Data = m[2] + + // Validate the base64 payload without retaining the decoded bytes. + // Try standard encoding first, fall back to raw URL-safe encoding. + if _, decErr := base64.StdEncoding.DecodeString(b64Data); decErr != nil { + if _, decErr2 := base64.RawURLEncoding.DecodeString(b64Data); decErr2 != nil { + return "", "", fmt.Errorf("base64 decode failed: %w", decErr) + } + } + + return mimeType, b64Data, nil +} diff --git a/internal/providers/openai_request.go b/internal/providers/openai_request.go index 421be3e8..84cc264a 100644 --- a/internal/providers/openai_request.go +++ b/internal/providers/openai_request.go @@ -147,7 +147,7 @@ func (p *OpenAIProvider) buildRequestBody(model string, req ChatRequest, stream } if len(req.Tools) > 0 { - body["tools"] = CleanToolSchemas(p.schemaProviderName(), req.Tools) + body["tools"] = buildToolsPayload(p.schemaProviderName(), req.Tools) body["tool_choice"] = "auto" } @@ -293,6 +293,44 @@ func openAIModelSupportsReasoningEffort(model string) bool { return false } +// buildToolsPayload serializes tools for the OpenAI-compat tools array. +// - function tools → {"type":"function","function":{cleaned schema}} +// - native tools (e.g. "image_generation") → {"type": t.Type} bare object +// +// Ordering is preserved. +func buildToolsPayload(schemaProvider string, tools []ToolDefinition) []map[string]any { + cleaned := CleanToolSchemas(schemaProvider, tools) + out := make([]map[string]any, 0, len(cleaned)) + for _, t := range cleaned { + switch t.Type { + case "function": + if t.Function == nil { + continue + } + params := t.Function.Parameters + fn := map[string]any{ + "name": t.Function.Name, + "description": t.Function.Description, + "parameters": params, + } + if t.Function.Strict != nil { + fn["strict"] = *t.Function.Strict + } + out = append(out, map[string]any{ + "type": "function", + "function": fn, + }) + default: + // Native provider tool — emit as bare {"type": X}. + // Richer field serialization is deferred to later phases. + out = append(out, map[string]any{ + "type": t.Type, + }) + } + } + return out +} + // openAIWireAssistantReasoningContent is true when assistant message objects may include // "reasoning_content" (thinking replay). Narrow allowlist — most OpenAI-compat hosts reject it. func openAIWireAssistantReasoningContent(model string) bool { diff --git a/internal/providers/openai_request_test.go b/internal/providers/openai_request_test.go new file mode 100644 index 00000000..c263fec2 --- /dev/null +++ b/internal/providers/openai_request_test.go @@ -0,0 +1,189 @@ +package providers + +import ( + "testing" +) + +// TestBuildToolsPayload_NativeOnly verifies that a single native "image_generation" tool +// is serialized as bare {"type":"image_generation"} without a "function" wrapper. +func TestBuildToolsPayload_NativeOnly(t *testing.T) { + tools := []ToolDefinition{ + {Type: "image_generation"}, + } + got := buildToolsPayload("openai", tools) + if len(got) != 1 { + t.Fatalf("expected 1 tool, got %d", len(got)) + } + if got[0]["type"] != "image_generation" { + t.Errorf("type = %q, want image_generation", got[0]["type"]) + } + if _, hasFunc := got[0]["function"]; hasFunc { + t.Error("native tool must not have 'function' field") + } + if len(got[0]) != 1 { + t.Errorf("native tool payload should have exactly 1 key, got %d: %v", len(got[0]), got[0]) + } +} + +// TestBuildToolsPayload_MixedOrder verifies that mixed function + native tools preserve +// insertion order and each is serialized correctly. +func TestBuildToolsPayload_MixedOrder(t *testing.T) { + tools := []ToolDefinition{ + { + Type: "function", + Function: &ToolFunctionSchema{ + Name: "read_file", + Description: "read a file", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "path": map[string]any{"type": "string"}, + }, + "required": []string{"path"}, + }, + }, + }, + {Type: "image_generation"}, + { + Type: "function", + Function: &ToolFunctionSchema{ + Name: "write_file", + Description: "write a file", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "path": map[string]any{"type": "string"}, + "content": map[string]any{"type": "string"}, + }, + "required": []string{"path", "content"}, + }, + }, + }, + } + + got := buildToolsPayload("openrouter", tools) + if len(got) != 3 { + t.Fatalf("expected 3 tools, got %d", len(got)) + } + + // Position 0: function tool + if got[0]["type"] != "function" { + t.Errorf("got[0] type = %q, want function", got[0]["type"]) + } + fn0, ok := got[0]["function"].(map[string]any) + if !ok { + t.Fatalf("got[0]['function'] is not map[string]any: %T", got[0]["function"]) + } + if fn0["name"] != "read_file" { + t.Errorf("got[0] function.name = %q, want read_file", fn0["name"]) + } + + // Position 1: native tool + if got[1]["type"] != "image_generation" { + t.Errorf("got[1] type = %q, want image_generation", got[1]["type"]) + } + if _, hasFunc := got[1]["function"]; hasFunc { + t.Error("native tool at position 1 must not have 'function' field") + } + + // Position 2: function tool + if got[2]["type"] != "function" { + t.Errorf("got[2] type = %q, want function", got[2]["type"]) + } + fn2, ok := got[2]["function"].(map[string]any) + if !ok { + t.Fatalf("got[2]['function'] is not map[string]any: %T", got[2]["function"]) + } + if fn2["name"] != "write_file" { + t.Errorf("got[2] function.name = %q, want write_file", fn2["name"]) + } +} + +// TestBuildToolsPayload_FunctionToolByteIdentical verifies existing function-only paths +// produce the expected function wrapper with name, description, and parameters fields. +func TestBuildToolsPayload_FunctionToolByteIdentical(t *testing.T) { + tools := []ToolDefinition{ + { + Type: "function", + Function: &ToolFunctionSchema{ + Name: "get_weather", + Description: "Get current weather", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "city": map[string]any{"type": "string"}, + }, + "required": []string{"city"}, + }, + }, + }, + } + + got := buildToolsPayload("openrouter", tools) + if len(got) != 1 { + t.Fatalf("expected 1 tool, got %d", len(got)) + } + if got[0]["type"] != "function" { + t.Errorf("type = %q, want function", got[0]["type"]) + } + fn, ok := got[0]["function"].(map[string]any) + if !ok { + t.Fatalf("'function' is not map[string]any: %T", got[0]["function"]) + } + if fn["name"] != "get_weather" { + t.Errorf("name = %q, want get_weather", fn["name"]) + } + if fn["description"] != "Get current weather" { + t.Errorf("description = %q, want 'Get current weather'", fn["description"]) + } + if fn["parameters"] == nil { + t.Error("parameters should not be nil") + } +} + +// TestBuildToolsPayload_NilFunctionSkipped verifies a malformed function tool +// (Type="function" with nil Function) is silently skipped. +func TestBuildToolsPayload_NilFunctionSkipped(t *testing.T) { + tools := []ToolDefinition{ + {Type: "function", Function: nil}, // malformed + {Type: "image_generation"}, + } + got := buildToolsPayload("openai", tools) + if len(got) != 1 { + t.Fatalf("expected 1 tool (malformed skipped), got %d", len(got)) + } + if got[0]["type"] != "image_generation" { + t.Errorf("expected image_generation tool, got type=%q", got[0]["type"]) + } +} + +// TestBuildRequestBody_NativeToolInBody verifies that buildRequestBody emits +// a native tool correctly in the request body via the tools key. +func TestBuildRequestBody_NativeToolInBody(t *testing.T) { + p := NewOpenAIProvider("openai", "sk-test", "https://api.openai.com/v1", "gpt-4o") + req := ChatRequest{ + Model: "gpt-4o", + Messages: []Message{{Role: "user", Content: "hello"}}, + Tools: []ToolDefinition{ + {Type: "image_generation"}, + }, + } + body := p.buildRequestBody("gpt-4o", req, false) + rawTools, ok := body["tools"] + if !ok { + t.Fatal("body missing 'tools' key") + } + tools, ok := rawTools.([]map[string]any) + if !ok { + t.Fatalf("tools is not []map[string]any: %T", rawTools) + } + if len(tools) != 1 { + t.Fatalf("expected 1 tool, got %d", len(tools)) + } + if tools[0]["type"] != "image_generation" { + t.Errorf("type = %q, want image_generation", tools[0]["type"]) + } + if _, hasFunc := tools[0]["function"]; hasFunc { + t.Error("native tool must not have 'function' field in request body") + } +} diff --git a/internal/providers/openai_types.go b/internal/providers/openai_types.go index 7ad840bd..57361e1c 100644 --- a/internal/providers/openai_types.go +++ b/internal/providers/openai_types.go @@ -13,11 +13,24 @@ type openAIChoice struct { } type openAIMessage struct { - Role string `json:"role"` - Content string `json:"content"` - ReasoningContent string `json:"reasoning_content,omitempty"` - Reasoning string `json:"reasoning,omitempty"` // Ollama alias for reasoning_content - ToolCalls []openAIToolCall `json:"tool_calls,omitempty"` + Role string `json:"role"` + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content,omitempty"` + Reasoning string `json:"reasoning,omitempty"` // Ollama alias for reasoning_content + ToolCalls []openAIToolCall `json:"tool_calls,omitempty"` + Images []openAIImagePart `json:"images,omitempty"` +} + +// openAIImagePart represents an image entry in the images[] array returned by +// OpenAI-compat endpoints that generate images (e.g. gpt-image-1 via chat completions). +type openAIImagePart struct { + Type string `json:"type"` + ImageURL openAIImageURLObject `json:"image_url"` +} + +// openAIImageURLObject holds the data URL (data:;base64,) for an image. +type openAIImageURLObject struct { + URL string `json:"url"` } type openAIToolCall struct { @@ -65,6 +78,7 @@ type openAIStreamDelta struct { ReasoningContent string `json:"reasoning_content,omitempty"` Reasoning string `json:"reasoning,omitempty"` // Ollama alias for reasoning_content ToolCalls []openAIStreamToolCall `json:"tool_calls,omitempty"` + Images []openAIImagePart `json:"images,omitempty"` } type openAIStreamToolCall struct { diff --git a/internal/providers/schema_cleaner.go b/internal/providers/schema_cleaner.go index 7075fde2..e932a8ca 100644 --- a/internal/providers/schema_cleaner.go +++ b/internal/providers/schema_cleaner.go @@ -2,38 +2,55 @@ package providers // CleanToolSchemas normalizes tool schemas for a specific provider. // This is the batch entry point — called from OpenAI/DashScope providers. +// Native tool types (anything other than "function") are passed through untouched. func CleanToolSchemas(providerName string, tools []ToolDefinition) []ToolDefinition { if len(tools) == 0 { return tools } profile := profileForProvider(providerName) - cleaned := make([]ToolDefinition, len(tools)) - for i, t := range tools { - // Exempt multi-action tools from strict mode — their many optional - // params become required under strict, forcing models to send empty values - // for every call (wasting ~200-300 output tokens per tool call). - useStrict := profile.StrictToolMode && !IsMultiActionSchema(t.Function.Parameters) - - var strictPtr *bool - if useStrict { - tr := true - strictPtr = &tr - } - - toolProfile := profile - toolProfile.StrictToolMode = useStrict - - cleaned[i] = ToolDefinition{ - Type: t.Type, - Function: ToolFunctionSchema{ - Name: t.Function.Name, - Description: t.Function.Description, - Parameters: normalizeWithProfile(toolProfile, t.Function.Parameters), - Strict: strictPtr, - }, + out := make([]ToolDefinition, 0, len(tools)) + for _, t := range tools { + switch t.Type { + case "function": + if t.Function == nil { + // Malformed function tool — skip rather than panic. + continue + } + fn := cleanFunctionSchema(profile, *t.Function) + out = append(out, ToolDefinition{ + Type: "function", + Function: &fn, + }) + default: + // Native provider tool (e.g. "image_generation") — pass through as-is. + out = append(out, t) } } - return cleaned + return out +} + +// cleanFunctionSchema normalizes a single function tool schema against a provider profile. +// Returns a new ToolFunctionSchema with cleaned parameters and strict mode applied. +func cleanFunctionSchema(profile SchemaProfile, fn ToolFunctionSchema) ToolFunctionSchema { + // Exempt multi-action tools from strict mode — their many optional params become + // required under strict, forcing models to send empty values (~200-300 wasted tokens/call). + useStrict := profile.StrictToolMode && !IsMultiActionSchema(fn.Parameters) + + var strictPtr *bool + if useStrict { + tr := true + strictPtr = &tr + } + + toolProfile := profile + toolProfile.StrictToolMode = useStrict + + return ToolFunctionSchema{ + Name: fn.Name, + Description: fn.Description, + Parameters: normalizeWithProfile(toolProfile, fn.Parameters), + Strict: strictPtr, + } } // CleanSchemaForProvider normalizes a single tool's parameters. diff --git a/internal/providers/schema_cleaner_test.go b/internal/providers/schema_cleaner_test.go index 2c88fd5d..a152c623 100644 --- a/internal/providers/schema_cleaner_test.go +++ b/internal/providers/schema_cleaner_test.go @@ -7,7 +7,7 @@ import ( func TestCleanToolSchemas_Gemini(t *testing.T) { tools := []ToolDefinition{{ Type: "function", - Function: ToolFunctionSchema{ + Function: &ToolFunctionSchema{ Name: "test", Description: "desc", Parameters: map[string]any{ @@ -91,7 +91,7 @@ func TestCleanToolSchemas_Anthropic(t *testing.T) { func TestCleanToolSchemas_Unknown(t *testing.T) { tools := []ToolDefinition{{ Type: "function", - Function: ToolFunctionSchema{ + Function: &ToolFunctionSchema{ Name: "test", Parameters: map[string]any{ "type": "object", @@ -266,7 +266,7 @@ func TestIsMultiActionSchema(t *testing.T) { func TestCleanToolSchemas_OpenAI_MultiActionExempt(t *testing.T) { multiAction := ToolDefinition{ Type: "function", - Function: ToolFunctionSchema{ + Function: &ToolFunctionSchema{ Name: "team_tasks", Description: "multi-action tool", Parameters: map[string]any{ @@ -291,7 +291,7 @@ func TestCleanToolSchemas_OpenAI_MultiActionExempt(t *testing.T) { } simple := ToolDefinition{ Type: "function", - Function: ToolFunctionSchema{ + Function: &ToolFunctionSchema{ Name: "read_file", Description: "simple tool", Parameters: map[string]any{ diff --git a/internal/providers/testdata/codex_native_image_non_stream.json b/internal/providers/testdata/codex_native_image_non_stream.json new file mode 100644 index 00000000..9ef7e62a --- /dev/null +++ b/internal/providers/testdata/codex_native_image_non_stream.json @@ -0,0 +1,30 @@ +[ + { + "type": "response.completed", + "response": { + "id": "resp_img2", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_001", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Here is your image."} + ] + }, + { + "type": "image_generation_call", + "id": "ig_xyz789", + "output_format": "png", + "result": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + } + ], + "usage": { + "input_tokens": 8, + "output_tokens": 3, + "total_tokens": 11 + } + } + } +] diff --git a/internal/providers/testdata/codex_native_image_partial_then_done.json b/internal/providers/testdata/codex_native_image_partial_then_done.json new file mode 100644 index 00000000..046e9991 --- /dev/null +++ b/internal/providers/testdata/codex_native_image_partial_then_done.json @@ -0,0 +1,38 @@ +[ + { + "type": "response.image_generation_call.partial_image", + "item_id": "ig_abc123", + "output_format": "png", + "partial_image_b64": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC", + "partial_image_index": 0 + }, + { + "type": "response.output_item.done", + "item": { + "id": "ig_abc123", + "type": "image_generation_call", + "output_format": "png", + "result": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + } + }, + { + "type": "response.completed", + "response": { + "id": "resp_img1", + "status": "completed", + "output": [ + { + "type": "image_generation_call", + "id": "ig_abc123", + "output_format": "png", + "result": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + } + ], + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15 + } + } + } +] diff --git a/internal/providers/testdata/openai_compat_image_mixed.json b/internal/providers/testdata/openai_compat_image_mixed.json new file mode 100644 index 00000000..e21ee8ec --- /dev/null +++ b/internal/providers/testdata/openai_compat_image_mixed.json @@ -0,0 +1,34 @@ +{ + "choices": [ + { + "message": { + "role": "assistant", + "content": "I'll generate that for you.", + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "log_generation", + "arguments": "{\"prompt\":\"a red square\"}" + } + } + ], + "images": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + } + } + ] + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 20, + "completion_tokens": 10, + "total_tokens": 30 + } +} diff --git a/internal/providers/testdata/openai_compat_image_nonstream.json b/internal/providers/testdata/openai_compat_image_nonstream.json new file mode 100644 index 00000000..ee8b4d28 --- /dev/null +++ b/internal/providers/testdata/openai_compat_image_nonstream.json @@ -0,0 +1,24 @@ +{ + "choices": [ + { + "message": { + "role": "assistant", + "content": "Here is your image.", + "images": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + } + } + ] + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15 + } +} diff --git a/internal/providers/testdata/openai_compat_image_stream.sse b/internal/providers/testdata/openai_compat_image_stream.sse new file mode 100644 index 00000000..f6d82d44 --- /dev/null +++ b/internal/providers/testdata/openai_compat_image_stream.sse @@ -0,0 +1,9 @@ +data: {"choices":[{"delta":{"role":"assistant","content":"Here "},"finish_reason":null}]} + +data: {"choices":[{"delta":{"content":"is your image."},"finish_reason":null}]} + +data: {"choices":[{"delta":{"images":[{"type":"image_url","image_url":{"url":"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC"}}]},"finish_reason":null}]} + +data: {"choices":[{"delta":{},"finish_reason":"stop"}]} + +data: [DONE] diff --git a/internal/providers/types.go b/internal/providers/types.go index abf0fdf4..9aa3b9e2 100644 --- a/internal/providers/types.go +++ b/internal/providers/types.go @@ -98,29 +98,36 @@ type ChatResponse struct { // ThinkingSignature is the accumulated signature from streaming thinking blocks. // Required by Anthropic API for tool use passback when thinking is enabled. ThinkingSignature string `json:"-"` + + // Images holds generated images returned by image_generation_call tools (Codex). + // Not persisted to DB; populated at runtime from provider response. + Images []ImageContent `json:"-"` } // StreamChunk is a piece of a streaming response. type StreamChunk struct { - Content string `json:"content,omitempty"` - Thinking string `json:"thinking,omitempty"` - Done bool `json:"done,omitempty"` + Content string `json:"content,omitempty"` + Thinking string `json:"thinking,omitempty"` + Done bool `json:"done,omitempty"` + Images []ImageContent `json:"images,omitempty"` // image generation frames (Codex) } // ImageContent represents a base64-encoded image for vision-capable models. type ImageContent struct { - MimeType string `json:"mime_type"` // e.g. "image/jpeg" - Data string `json:"data"` // base64-encoded image bytes + MimeType string `json:"mime_type"` // e.g. "image/jpeg" + Data string `json:"data"` // base64-encoded image bytes + Partial bool `json:"partial,omitempty"` // true for intermediate frames (Codex image_generation_call) } // MediaRef is a lightweight reference to a persistently stored media file. // Stored in session JSONB (~60 bytes each) instead of megabytes for base64. // On reload, MediaRefs are resolved to file paths and loaded into Images (for images). type MediaRef struct { - ID string `json:"id"` // unique media ID (uuid) - MimeType string `json:"mime_type"` // e.g. "image/jpeg", "application/pdf" - Kind string `json:"kind"` // "image", "video", "audio", "document" - Path string `json:"path,omitempty"` // absolute workspace path (persisted for /v1/files/ serving) + ID string `json:"id"` // unique media ID (uuid) + MimeType string `json:"mime_type"` // e.g. "image/jpeg", "application/pdf" + Kind string `json:"kind"` // "image", "video", "audio", "document" + Path string `json:"path,omitempty"` // absolute workspace path (persisted for /v1/files/ serving) + Prompt string `json:"prompt,omitempty"` // prompt that generated this asset, if known } // Message represents a conversation message. @@ -160,9 +167,12 @@ type ToolCall struct { } // ToolDefinition describes a tool available to the LLM. +// Type is "function" for standard function tools, or a native provider tool type +// (e.g. "image_generation") for first-class provider-native tools. +// Function is nil when Type is not "function". type ToolDefinition struct { - Type string `json:"type"` // "function" - Function ToolFunctionSchema `json:"function"` + Type string `json:"type"` // "function" | "image_generation" | ... + Function *ToolFunctionSchema `json:"function,omitempty"` // nil when Type != "function" } // ToolFunctionSchema is the schema for a function tool. diff --git a/internal/store/agent_store.go b/internal/store/agent_store.go index 64c2c30a..0103a149 100644 --- a/internal/store/agent_store.go +++ b/internal/store/agent_store.go @@ -229,6 +229,28 @@ func (a *AgentData) ParseSelfEvolve() bool { return a.SelfEvolve } // ParseSkillEvolve returns whether the agent's skill learning loop is enabled. func (a *AgentData) ParseSkillEvolve() bool { return a.SkillEvolve } +// ParseAllowImageGeneration returns whether the native image_generation tool +// is allowed for this agent. Defaults to true (enabled) when not set in +// other_config, so existing agents automatically get image generation with +// Codex providers. Operators can explicitly disable it by setting +// other_config.allow_image_generation = false. +// No DB column — code-only default to avoid a migration for a feature flag. +func (a *AgentData) ParseAllowImageGeneration() bool { + if len(a.OtherConfig) <= 2 { + return true // default: enabled + } + var bag struct { + AllowImageGeneration *bool `json:"allow_image_generation"` + } + if json.Unmarshal(a.OtherConfig, &bag) != nil { + return true // malformed config → default: enabled + } + if bag.AllowImageGeneration == nil { + return true // not set → default: enabled + } + return *bag.AllowImageGeneration +} + // validPromptModes is the set of allowed prompt_mode values. var validPromptModes = map[string]bool{ "full": true, "task": true, "minimal": true, "none": true, diff --git a/internal/store/agent_store_test.go b/internal/store/agent_store_test.go index c7c9dd43..0a820366 100644 --- a/internal/store/agent_store_test.go +++ b/internal/store/agent_store_test.go @@ -348,3 +348,47 @@ func TestResolveEffectiveChatGPTOAuthRoutingIgnoresCustomMembersWhenProviderOwns t.Fatalf("ExtraProviderNames = %#v, want provider defaults %#v", got.ExtraProviderNames, defaults.ExtraProviderNames) } } + +// ─── ParseAllowImageGeneration ──────────────────────────────────────────── + +func TestParseAllowImageGeneration_DefaultTrue_NoOtherConfig(t *testing.T) { + ag := &AgentData{} + if !ag.ParseAllowImageGeneration() { + t.Error("empty other_config must default to true (image gen enabled)") + } +} + +func TestParseAllowImageGeneration_DefaultTrue_EmptyObject(t *testing.T) { + ag := &AgentData{OtherConfig: json.RawMessage(`{}`)} + if !ag.ParseAllowImageGeneration() { + t.Error("empty JSONB object must default to true") + } +} + +func TestParseAllowImageGeneration_ExplicitTrue(t *testing.T) { + ag := &AgentData{OtherConfig: json.RawMessage(`{"allow_image_generation":true}`)} + if !ag.ParseAllowImageGeneration() { + t.Error("explicit true must return true") + } +} + +func TestParseAllowImageGeneration_ExplicitFalse(t *testing.T) { + ag := &AgentData{OtherConfig: json.RawMessage(`{"allow_image_generation":false}`)} + if ag.ParseAllowImageGeneration() { + t.Error("explicit false must return false") + } +} + +func TestParseAllowImageGeneration_MalformedJSON_DefaultsTrue(t *testing.T) { + ag := &AgentData{OtherConfig: json.RawMessage(`{not-json`)} + if !ag.ParseAllowImageGeneration() { + t.Error("malformed other_config must default to true") + } +} + +func TestParseAllowImageGeneration_UnrelatedKeys_DefaultsTrue(t *testing.T) { + ag := &AgentData{OtherConfig: json.RawMessage(`{"self_evolve":true,"skill_evolve":false}`)} + if !ag.ParseAllowImageGeneration() { + t.Error("other_config without allow_image_generation key must default to true") + } +} diff --git a/internal/tools/create_image.go b/internal/tools/create_image.go index 09decfcb..34ff9192 100644 --- a/internal/tools/create_image.go +++ b/internal/tools/create_image.go @@ -104,6 +104,10 @@ func (t *CreateImageTool) Execute(ctx context.Context, args map[string]any) *Res return ErrorResult(fmt.Sprintf("image generation failed: %v", err)) } + // Embed prompt into PNG tEXt metadata before writing to disk. + // If embedding fails (malformed bytes, non-PNG) the original data is used unchanged. + imageData := embedPromptIntoPNG(chainResult.Data, prompt) + // Save to workspace under date-based folder (e.g. generated/2026-03-02/) workspace := ToolWorkspaceFromCtx(ctx) if workspace == "" { @@ -114,7 +118,7 @@ func (t *CreateImageTool) Execute(ctx context.Context, args map[string]any) *Res return ErrorResult(fmt.Sprintf("failed to create output directory: %v", err)) } imagePath := filepath.Join(dateDir, mediaFileName(ctx, "image", filenameHint, "png")) - if err := os.WriteFile(imagePath, chainResult.Data, 0644); err != nil { + if err := os.WriteFile(imagePath, imageData, 0644); err != nil { return ErrorResult(fmt.Sprintf("failed to save generated image: %v", err)) } @@ -123,11 +127,12 @@ func (t *CreateImageTool) Execute(ctx context.Context, args map[string]any) *Res slog.Warn("create_image: file missing immediately after write", "path", imagePath, "error", err) return ErrorResult(fmt.Sprintf("generated image file missing after write: %v", err)) } else { - slog.Info("create_image: file saved", "path", imagePath, "size", fi.Size(), "data_len", len(chainResult.Data)) + slog.Info("create_image: file saved", "path", imagePath, "size", fi.Size(), "data_len", len(imageData)) } result := &Result{ForLLM: fmt.Sprintf("MEDIA:%s\nUse the EXACT filename when referencing: %s", imagePath, filepath.Base(imagePath))} result.Media = []bus.MediaFile{{Path: imagePath, MimeType: "image/png", Filename: filepath.Base(imagePath)}} + result.MediaPrompts = map[int]string{0: prompt} result.Deliverable = fmt.Sprintf("[Generated image: %s]\nPrompt: %s", filepath.Base(imagePath), prompt) if t.vaultIntc != nil { go t.vaultIntc.AfterWriteMedia(context.WithoutCancel(ctx), imagePath, prompt, "image/png") @@ -140,8 +145,48 @@ func (t *CreateImageTool) Execute(ctx context.Context, args map[string]any) *Res return result } +// embedPromptIntoPNG wraps agent.EmbedPNGPrompt for the tools package. +// Logs a warning on error but always returns usable bytes. +func embedPromptIntoPNG(data []byte, prompt string) []byte { + if prompt == "" { + return data + } + // Import cycle guard: tools → agent is not allowed. Use the local pngEmbed function. + out, err := pngEmbedPrompt(data, prompt) + if err != nil { + slog.Warn("create_image: failed to embed prompt into PNG metadata", "error", err) + return data + } + return out +} + // callProvider dispatches to the correct image generation implementation based on provider type. +// If the resolved provider implements NativeImageProvider (e.g. CodexProvider via OAuth), +// the native path is used and cp may be nil. The credentialProvider path is only reached +// for API-key-backed providers. func (t *CreateImageTool) callProvider(ctx context.Context, cp credentialProvider, providerName, model string, params map[string]any) ([]byte, *providers.Usage, error) { + // Native path: provider implements the image_generation tool natively (e.g. Codex/OAuth). + // The raw provider object is injected into params["_native_provider"] by ExecuteWithChain. + // Must check before the cp==nil guard — these providers intentionally have no APIKey/APIBase. + if rawProvider, ok := params["_native_provider"]; ok { + if np, ok := rawProvider.(providers.NativeImageProvider); ok { + prompt := GetParamString(params, "prompt", "") + aspectRatio := GetParamString(params, "aspect_ratio", "1:1") + imageModel := GetParamString(params, "image_model", "") + result, err := np.GenerateImage(ctx, providers.NativeImageRequest{ + Model: model, + ImageModel: imageModel, + Prompt: prompt, + AspectRatio: aspectRatio, + OutputFormat: "png", + }) + if err != nil { + return nil, nil, fmt.Errorf("native image generation: %w", err) + } + return result.Data, result.Usage, nil + } + } + if cp == nil { return nil, nil, fmt.Errorf("provider %q does not expose API credentials required for image generation", providerName) } diff --git a/internal/tools/create_image_native_path_test.go b/internal/tools/create_image_native_path_test.go new file mode 100644 index 00000000..bd0c7020 --- /dev/null +++ b/internal/tools/create_image_native_path_test.go @@ -0,0 +1,258 @@ +package tools + +import ( + "context" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +// nativeImageProvider is a minimal fake that satisfies providers.NativeImageProvider. +// It records call arguments so tests can assert correct routing. +type nativeImageProvider struct { + name string + model string + calledWith *providers.NativeImageRequest + returnData []byte + returnError error +} + +func (p *nativeImageProvider) Name() string { return p.name } +func (p *nativeImageProvider) DefaultModel() string { return p.model } +func (p *nativeImageProvider) Chat(_ context.Context, _ providers.ChatRequest) (*providers.ChatResponse, error) { + return &providers.ChatResponse{}, nil +} +func (p *nativeImageProvider) ChatStream(_ context.Context, _ providers.ChatRequest, _ func(providers.StreamChunk)) (*providers.ChatResponse, error) { + return &providers.ChatResponse{}, nil +} +func (p *nativeImageProvider) GenerateImage(_ context.Context, req providers.NativeImageRequest) (*providers.NativeImageResult, error) { + p.calledWith = &req + if p.returnError != nil { + return nil, p.returnError + } + return &providers.NativeImageResult{ + MimeType: "image/png", + Data: p.returnData, + }, nil +} + +// TestCreateImageTool_RoutesNativePath verifies that when the provider chain resolves to +// a provider that implements NativeImageProvider (e.g. CodexProvider via OAuth), +// the create_image tool uses the native path (GenerateImage) and not the credentialProvider +// path. Specifically: the tool must NOT fail with "does not expose API credentials". +func TestCreateImageTool_RoutesNativePath(t *testing.T) { + // Build a minimal 8-byte PNG-like data (not a real PNG, but large enough for + // the tool to write to disk without crashing). + pngMagic := []byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, + // IHDR chunk (minimal valid PNG) — 25 bytes: len(13) + type + data + crc + 0x00, 0x00, 0x00, 0x0d, + 0x49, 0x48, 0x44, 0x52, + 0x00, 0x00, 0x00, 0x01, // width = 1 + 0x00, 0x00, 0x00, 0x01, // height = 1 + 0x08, 0x02, 0x00, 0x00, 0x00, + 0x90, 0x77, 0x53, 0xde, + // IDAT chunk (minimal: zlib compressed 1x1 pixel) + 0x00, 0x00, 0x00, 0x0c, + 0x49, 0x44, 0x41, 0x54, + 0x08, 0xd7, 0x63, 0xf8, 0xcf, 0xc0, 0x00, 0x00, 0x00, 0x02, 0x00, 0x01, + 0xe2, 0x21, 0xbc, 0x33, + // IEND chunk + 0x00, 0x00, 0x00, 0x00, + 0x49, 0x45, 0x4e, 0x44, + 0xae, 0x42, 0x60, 0x82, + } + + fakeProvider := &nativeImageProvider{ + name: "openai-codex", + model: "gpt-image-2", + returnData: pngMagic, + } + + // Register provider in a fresh registry. + reg := providers.NewRegistry(nil) + reg.Register(fakeProvider) + + // Build a chain that points to the fake native provider. + chain := []MediaProviderEntry{ + { + Provider: "openai-codex", + Model: "gpt-image-2", + Enabled: true, + Timeout: 30, + MaxRetries: 1, + }, + } + + // Inject workspace context so the tool can write the file. + ctx := WithToolWorkspace(context.Background(), t.TempDir()) + + tool := NewCreateImageTool(reg) + + // Execute via the chain directly (same code path as Execute, bypassing chain resolution). + chainResult, err := ExecuteWithChain(ctx, chain, reg, tool.callProvider) + if err != nil { + t.Fatalf("ExecuteWithChain returned error: %v — native path was NOT used (credentialProvider path instead)", err) + } + + // Verify the fake provider was called via the native interface. + if fakeProvider.calledWith == nil { + t.Fatal("NativeImageProvider.GenerateImage was not called") + } + + // The native provider's GenerateImage should have received a prompt. + // (Prompt is injected by callProvider from params["prompt"].) + // We cannot assert non-empty here without injecting params, but we can assert + // the chain result contains the returned bytes. + if len(chainResult.Data) == 0 { + t.Error("chainResult.Data is empty — native path should return image bytes") + } + + // Provider and model must be populated in the chain result. + if chainResult.Provider != "openai-codex" { + t.Errorf("chainResult.Provider = %q, want openai-codex", chainResult.Provider) + } + if chainResult.Model != "gpt-image-2" { + t.Errorf("chainResult.Model = %q, want gpt-image-2", chainResult.Model) + } +} + +// TestCreateImageTool_RoutesNativePath_WithPrompt verifies end-to-end that the +// Execute method (with prompt in args) routes via the native path and sets +// MediaPrompts[0] on the result. +func TestCreateImageTool_RoutesNativePath_WithPrompt(t *testing.T) { + pngMagic := []byte{ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, // PNG signature + // IEND chunk (minimal: enough for pngEmbedPrompt to process) + 0x00, 0x00, 0x00, 0x00, + 0x49, 0x45, 0x4e, 0x44, + 0xae, 0x42, 0x60, 0x82, + } + + wantPrompt := "a sunny day at the beach" + + fakeProvider := &nativeImageProvider{ + name: "openai-codex", + model: "gpt-image-2", + returnData: pngMagic, + } + + reg := providers.NewRegistry(nil) + reg.Register(fakeProvider) + + // Inject per-agent provider override so chain resolves to our fake. + chainJSON := []byte(`{"providers":[{"provider":"openai-codex","model":"gpt-image-2","enabled":true,"timeout":30,"max_retries":1}]}`) + settings := BuiltinToolSettings{"create_image": chainJSON} + ctx := WithBuiltinToolSettings(context.Background(), settings) + ctx = WithToolWorkspace(ctx, t.TempDir()) + + tool := NewCreateImageTool(reg) + result := tool.Execute(ctx, map[string]any{ + "prompt": wantPrompt, + "aspect_ratio": "1:1", + }) + + if result.IsError { + t.Fatalf("Execute returned error: %q", result.ForLLM) + } + + // Verify NativeImageProvider.GenerateImage was called with the correct prompt. + if fakeProvider.calledWith == nil { + t.Fatal("GenerateImage was not called on the native provider") + } + if fakeProvider.calledWith.Prompt != wantPrompt { + t.Errorf("GenerateImage called with prompt %q, want %q", + fakeProvider.calledWith.Prompt, wantPrompt) + } + + // MediaPrompts must carry the prompt so MediaRef.Prompt gets populated downstream. + if result.MediaPrompts == nil || result.MediaPrompts[0] != wantPrompt { + t.Errorf("result.MediaPrompts[0] = %q, want %q", result.MediaPrompts[0], wantPrompt) + } + // Media must have one entry. + if len(result.Media) != 1 { + t.Errorf("result.Media length = %d, want 1", len(result.Media)) + } +} + +// TestCreateImageTool_ThreadsImageModel verifies that params["image_model"] from the +// chain entry is forwarded into NativeImageRequest.ImageModel. This covers the data +// flow: chain entry JSON → callProvider → GenerateImage. +func TestCreateImageTool_ThreadsImageModel(t *testing.T) { + pngMagic := []byte{ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, // PNG signature + 0x00, 0x00, 0x00, 0x00, + 0x49, 0x45, 0x4e, 0x44, + 0xae, 0x42, 0x60, 0x82, + } + + tests := []struct { + name string + chainImageModel string + wantImageModel string + }{ + { + name: "default (empty params.image_model)", + chainImageModel: "", + wantImageModel: "", // provider validator defaults to gpt-image-2 + }, + { + name: "legacy gpt-image-1.5", + chainImageModel: "gpt-image-1.5", + wantImageModel: "gpt-image-1.5", + }, + { + name: "explicit gpt-image-2", + chainImageModel: "gpt-image-2", + wantImageModel: "gpt-image-2", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + fakeProvider := &nativeImageProvider{ + name: "openai-codex", + model: "gpt-image-2", + returnData: pngMagic, + } + + reg := providers.NewRegistry(nil) + reg.Register(fakeProvider) + + // Build chain entry with optional image_model param. + entryParams := map[string]any{ + "prompt": "test image", + "aspect_ratio": "1:1", + } + if tc.chainImageModel != "" { + entryParams["image_model"] = tc.chainImageModel + } + chain := []MediaProviderEntry{ + { + Provider: "openai-codex", + Model: "gpt-image-2", + Enabled: true, + Timeout: 30, + MaxRetries: 1, + Params: entryParams, + }, + } + + ctx := WithToolWorkspace(context.Background(), t.TempDir()) + tool := NewCreateImageTool(reg) + + _, err := ExecuteWithChain(ctx, chain, reg, tool.callProvider) + if err != nil { + t.Fatalf("ExecuteWithChain returned error: %v", err) + } + + if fakeProvider.calledWith == nil { + t.Fatal("GenerateImage was not called on the native provider") + } + + gotImageModel := fakeProvider.calledWith.ImageModel + if gotImageModel != tc.wantImageModel { + t.Errorf("NativeImageRequest.ImageModel = %q, want %q", gotImageModel, tc.wantImageModel) + } + }) + } +} diff --git a/internal/tools/media_provider_chain.go b/internal/tools/media_provider_chain.go index e62ab897..320b4a53 100644 --- a/internal/tools/media_provider_chain.go +++ b/internal/tools/media_provider_chain.go @@ -20,8 +20,8 @@ type MediaProviderEntry struct { Provider string `json:"provider"` // name for registry.Get() Model string `json:"model"` Enabled bool `json:"enabled"` - Timeout int `json:"timeout"` // seconds, default 120 - MaxRetries int `json:"max_retries"` // default 2 + Timeout int `json:"timeout"` // seconds, default 600 (10 min — image/video gen is slow) + MaxRetries int `json:"max_retries"` // default 1 (image gen rarely succeeds on retry) Params map[string]any `json:"params,omitempty"` // provider-specific config } @@ -33,12 +33,23 @@ type mediaProviderChain struct { } // applyDefaults fills in zero-value fields with sensible defaults. +// +// Timeout: 600s (10 min). Native image_generation on gpt-image-2 can legitimately +// take 4-8 min for complex prompts with heavy in-image text (e.g. infographics). +// 120s — the old default — routinely truncated real work mid-generation and +// surfaced as 'context deadline exceeded'. Lowering this risks re-introducing +// that footgun; operators can still set a tighter value explicitly. +// +// MaxRetries: 1. Image generation is stateful per upstream run — a mid-flight +// timeout leaves orphan server-side work. Retrying a fresh generation (new +// upstream run) doubles cost and rarely succeeds where the first attempt timed +// out. Surface the failure fast so the user can adjust the timeout. func (e *MediaProviderEntry) applyDefaults() { if e.Timeout <= 0 { - e.Timeout = 120 + e.Timeout = 600 } if e.MaxRetries <= 0 { - e.MaxRetries = 2 + e.MaxRetries = 1 } } @@ -176,12 +187,15 @@ func ExecuteWithChain( // callProvider falls back to using the provider's Chat() API. cp, _ := p.(credentialProvider) - // Inject resolved provider type into params so callProvider can route correctly. - // Clone params to avoid mutating the original entry config. + // Inject resolved provider type and the raw provider object into params so + // callProvider can route correctly. Clone params to avoid mutating entry config. resolvedType := ResolveProviderType(p) - callParams := make(map[string]any, len(entry.Params)+1) + callParams := make(map[string]any, len(entry.Params)+2) maps.Copy(callParams, entry.Params) callParams["_provider_type"] = resolvedType + // "_native_provider" carries the providers.Provider instance so callProvider + // can type-assert to NativeImageProvider without a separate registry lookup. + callParams["_native_provider"] = p // Retry loop for this provider for attempt := 1; attempt <= entry.MaxRetries; attempt++ { diff --git a/internal/tools/png_embed.go b/internal/tools/png_embed.go new file mode 100644 index 00000000..0bf133f7 --- /dev/null +++ b/internal/tools/png_embed.go @@ -0,0 +1,84 @@ +package tools + +import ( + "bytes" + "encoding/binary" + "hash/crc32" +) + +// pngEmbedPrompt embeds a generation prompt into a PNG byte stream as a tEXt +// "Description" chunk inserted before the IEND chunk. +// +// If data is not a valid PNG (wrong magic bytes or no IEND chunk), the original +// bytes are returned unchanged. An empty prompt is a no-op. +// +// This is the tools-package counterpart of agent.EmbedPNGPrompt. A separate copy +// is required to avoid the tools→agent import cycle. +func pngEmbedPrompt(data []byte, prompt string) ([]byte, error) { + if len(prompt) == 0 { + return data, nil + } + + pngSig := []byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a} + if len(data) < len(pngSig) || !bytes.Equal(data[:len(pngSig)], pngSig) { + return data, nil + } + + iendOffset := pngFindIEND(data) + if iendOffset < 0 { + return data, nil + } + + extra := pngBuildTextChunk("Description", prompt) + + result := make([]byte, 0, len(data)+len(extra)) + result = append(result, data[:iendOffset]...) + result = append(result, extra...) + result = append(result, data[iendOffset:]...) + return result, nil +} + +// pngBuildTextChunk encodes a single PNG tEXt chunk for the given keyword/value pair. +func pngBuildTextChunk(keyword, value string) []byte { + data := make([]byte, 0, len(keyword)+1+len(value)) + data = append(data, []byte(keyword)...) + data = append(data, 0x00) + data = append(data, []byte(value)...) + + chunkType := []byte("tEXt") + crcInput := append(chunkType, data...) + checksum := crc32.ChecksumIEEE(crcInput) + + var buf bytes.Buffer + var lenBuf [4]byte + binary.BigEndian.PutUint32(lenBuf[:], uint32(len(data))) + buf.Write(lenBuf[:]) + buf.Write(chunkType) + buf.Write(data) + var crcBuf [4]byte + binary.BigEndian.PutUint32(crcBuf[:], checksum) + buf.Write(crcBuf[:]) + return buf.Bytes() +} + +// pngFindIEND returns the byte offset of the IEND chunk start, or -1 if not found. +func pngFindIEND(data []byte) int { + pngSigLen := 8 + pos := pngSigLen + for pos+12 <= len(data) { + chunkLen := int(binary.BigEndian.Uint32(data[pos : pos+4])) + if chunkLen < 0 { + break + } + chunkType := data[pos+4 : pos+8] + if bytes.Equal(chunkType, []byte("IEND")) { + return pos + } + next := pos + 8 + chunkLen + 4 + if next <= pos { + break // overflow guard + } + pos = next + } + return -1 +} diff --git a/internal/tools/policy.go b/internal/tools/policy.go index ca2c1b76..dd0d6588 100644 --- a/internal/tools/policy.go +++ b/internal/tools/policy.go @@ -163,7 +163,7 @@ func (pe *PolicyEngine) FilterTools( if tool, ok := registry.Get(canonical); ok { defs = append(defs, providers.ToolDefinition{ Type: "function", - Function: providers.ToolFunctionSchema{ + Function: &providers.ToolFunctionSchema{ Name: alias, Description: tool.Description(), Parameters: tool.Parameters(), diff --git a/internal/tools/registry.go b/internal/tools/registry.go index bc4ee5d3..46f13a14 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -294,7 +294,7 @@ func (r *Registry) ProviderDefs() []providers.ToolDefinition { } defs = append(defs, providers.ToolDefinition{ Type: "function", - Function: providers.ToolFunctionSchema{ + Function: &providers.ToolFunctionSchema{ Name: alias, Description: tool.Description(), Parameters: tool.Parameters(), diff --git a/internal/tools/result.go b/internal/tools/result.go index 08d6564d..3126cc81 100644 --- a/internal/tools/result.go +++ b/internal/tools/result.go @@ -17,6 +17,11 @@ type Result struct { // Media holds media files to forward as output (e.g. images from delegation). Media []bus.MediaFile `json:"-"` + // MediaPrompts maps each Media[i] index to the generation prompt that produced it. + // Used to populate MediaRef.Prompt when the pipeline persists tool-generated images. + // Nil means no prompt metadata available. + MediaPrompts map[int]string `json:"-"` + // Deliverable holds the primary work output from this tool execution. // Used to capture actual content (e.g. written file text, image prompt) for team // task results instead of relying on the LLM's summary response. diff --git a/internal/tools/types.go b/internal/tools/types.go index 88c08894..661fb8b4 100644 --- a/internal/tools/types.go +++ b/internal/tools/types.go @@ -121,7 +121,7 @@ type ChannelAware interface { func ToProviderDef(t Tool) providers.ToolDefinition { return providers.ToolDefinition{ Type: "function", - Function: providers.ToolFunctionSchema{ + Function: &providers.ToolFunctionSchema{ Name: t.Name(), Description: t.Description(), Parameters: t.Parameters(), diff --git a/tests/integration/mcp_grant_revoke_test.go b/tests/integration/mcp_grant_revoke_test.go index 301a3403..c236acac 100644 --- a/tests/integration/mcp_grant_revoke_test.go +++ b/tests/integration/mcp_grant_revoke_test.go @@ -5,6 +5,7 @@ package integration import ( "context" "database/sql" + "strings" "sync/atomic" "testing" @@ -242,16 +243,7 @@ func grantUserAccess(t *testing.T, db *sql.DB, tenantID, serverID uuid.UUID, use } func containsGrantRevoked(s string) bool { - return len(s) > 0 && (contains(s, "grant revoked") || contains(s, "grant denied")) -} - -func contains(s, substr string) bool { - for i := 0; i <= len(s)-len(substr); i++ { - if s[i:i+len(substr)] == substr { - return true - } - } - return false + return len(s) > 0 && (strings.Contains(s, "grant revoked") || strings.Contains(s, "grant denied")) } // fakeMCPClient is a stub for testing. Since mcpclient.Client is a struct diff --git a/tests/integration/tts_gemini_live_test.go b/tests/integration/tts_gemini_live_test.go index ed5c539f..316d29d0 100644 --- a/tests/integration/tts_gemini_live_test.go +++ b/tests/integration/tts_gemini_live_test.go @@ -5,6 +5,7 @@ package integration import ( "context" "os" + "strings" "testing" "github.com/nextlevelbuilder/goclaw/internal/audio" @@ -87,15 +88,5 @@ func isServerError(err error) bool { return false } msg := err.Error() - return contains(msg, "500") || contains(msg, "503") || contains(msg, "server error") -} - -// contains checks if a string contains a substring (case-insensitive). -func contains(s, substr string) bool { - for i := 0; i < len(s)-len(substr)+1; i++ { - if s[i:i+len(substr)] == substr { - return true - } - } - return false + return strings.Contains(msg, "500") || strings.Contains(msg, "503") || strings.Contains(msg, "server error") } diff --git a/ui/web/src/adapters/chat-message.adapter.ts b/ui/web/src/adapters/chat-message.adapter.ts index 164457d5..8fdfc7f5 100644 --- a/ui/web/src/adapters/chat-message.adapter.ts +++ b/ui/web/src/adapters/chat-message.adapter.ts @@ -33,6 +33,7 @@ export function transformHistoryMessages( mimeType: ref.mime_type, fileName: (ref.path?.split("?")[0]?.split("/").pop()) ?? ref.id, kind: (ref.kind as MediaItem["kind"]) || "document", + prompt: ref.prompt || undefined, })); } diff --git a/ui/web/src/components/chat/chat-input.tsx b/ui/web/src/components/chat/chat-input.tsx index bdf3bbdb..53a3dde0 100644 --- a/ui/web/src/components/chat/chat-input.tsx +++ b/ui/web/src/components/chat/chat-input.tsx @@ -19,7 +19,14 @@ interface ChatInputProps { onFilesChange: (files: AttachedFile[]) => void; } -export function ChatInput({ onSend, onAbort, isBusy, disabled, files, onFilesChange }: ChatInputProps) { +export function ChatInput({ + onSend, + onAbort, + isBusy, + disabled, + files, + onFilesChange, +}: ChatInputProps) { const { t } = useTranslation("common"); const [value, setValue] = useState(""); const textareaRef = useRef(null); diff --git a/ui/web/src/components/chat/media-gallery.tsx b/ui/web/src/components/chat/media-gallery.tsx index 7a21bde1..858f002c 100644 --- a/ui/web/src/components/chat/media-gallery.tsx +++ b/ui/web/src/components/chat/media-gallery.tsx @@ -12,6 +12,35 @@ import { } from "@/components/ui/dialog"; import type { MediaItem } from "@/types/chat"; +/** Hex/UUID filename pattern — matches assistant-generated image basenames + * such as `a1b2c3d4e5f6.png` or `550e8400-e29b-41d4-a716-446655440000.png`. */ +const GENERATED_FILENAME_RE = /^[0-9a-f-]{8,}\.png$/i; + +/** + * Format a Date as YYYYMMDD-HHmmss for use in download filenames. + * Uses local time so filenames match the user's clock. + */ +function formatTimestamp(d: Date): string { + const pad = (n: number) => String(n).padStart(2, "0"); + return ( + `${d.getFullYear()}${pad(d.getMonth() + 1)}${pad(d.getDate())}` + + `-${pad(d.getHours())}${pad(d.getMinutes())}${pad(d.getSeconds())}` + ); +} + +/** + * Returns the download filename for an image MediaItem. + * For assistant-generated PNGs (hex/UUID basename), uses a descriptive + * `generated-YYYYMMDD-HHmmss.png` name instead of the opaque server hash. + */ +function resolveImageDownloadName(item: MediaItem): string { + const base = item.fileName ?? "image"; + if (item.mimeType === "image/png" && GENERATED_FILENAME_RE.test(base)) { + return `generated-${formatTimestamp(new Date())}.png`; + } + return base; +} + interface MediaGalleryProps { items: MediaItem[]; } @@ -82,37 +111,47 @@ export function MediaGallery({ items }: MediaGalleryProps) { {images.length > 0 && (
{images.map((item, i) => ( -
- -
-
-
- {item.fileName && {item.fileName}} - {item.size != null && item.size > 0 && ( - {formatSize(item.size)} - )} -
- e.stopPropagation()} - className="shrink-0 rounded-lg bg-white/90 dark:bg-neutral-800/90 p-1.5 text-neutral-700 dark:text-neutral-200 shadow-md ring-1 ring-black/10 dark:ring-white/10 hover:bg-white dark:hover:bg-neutral-700 transition-colors cursor-pointer" - title="Download" + diff --git a/ui/web/src/i18n/locales/en/chat.json b/ui/web/src/i18n/locales/en/chat.json index c0a347e8..a8e3a9c3 100644 --- a/ui/web/src/i18n/locales/en/chat.json +++ b/ui/web/src/i18n/locales/en/chat.json @@ -29,5 +29,6 @@ "selectAgent": { "title": "Select an agent", "description": "Choose an agent to start chatting" - } + }, + "imageGenDownloadName": "generated-{{timestamp}}.png" } diff --git a/ui/web/src/i18n/locales/en/tools.json b/ui/web/src/i18n/locales/en/tools.json index c59cf227..9173bc1a 100644 --- a/ui/web/src/i18n/locales/en/tools.json +++ b/ui/web/src/i18n/locales/en/tools.json @@ -232,7 +232,10 @@ "addProvider": "Add Provider", "cancel": "Cancel", "save": "Save", - "saving": "Saving..." + "saving": "Saving...", + "imageModelLabel": "Image model", + "imageModelDefaultOption": "Default · gpt-image-2 (recommended)", + "imageModelLegacyOption": "Legacy · gpt-image-1.5" }, "toast": { "toggled": "Tool updated", diff --git a/ui/web/src/i18n/locales/vi/chat.json b/ui/web/src/i18n/locales/vi/chat.json index c5c5ac90..b0c58e0d 100644 --- a/ui/web/src/i18n/locales/vi/chat.json +++ b/ui/web/src/i18n/locales/vi/chat.json @@ -29,5 +29,6 @@ "selectAgent": { "title": "Chọn agent", "description": "Chọn agent để bắt đầu trò chuyện" - } + }, + "imageGenDownloadName": "generated-{{timestamp}}.png" } diff --git a/ui/web/src/i18n/locales/vi/tools.json b/ui/web/src/i18n/locales/vi/tools.json index 2c3fcd2f..3558d411 100644 --- a/ui/web/src/i18n/locales/vi/tools.json +++ b/ui/web/src/i18n/locales/vi/tools.json @@ -232,7 +232,10 @@ "addProvider": "Thêm provider", "cancel": "Hủy", "save": "Lưu", - "saving": "Đang lưu..." + "saving": "Đang lưu...", + "imageModelLabel": "Mô hình ảnh", + "imageModelDefaultOption": "Mặc định · gpt-image-2 (khuyến nghị)", + "imageModelLegacyOption": "Cũ · gpt-image-1.5" }, "toast": { "toggled": "Đã cập nhật công cụ", diff --git a/ui/web/src/i18n/locales/zh/chat.json b/ui/web/src/i18n/locales/zh/chat.json index c16aa14c..ee00aecd 100644 --- a/ui/web/src/i18n/locales/zh/chat.json +++ b/ui/web/src/i18n/locales/zh/chat.json @@ -29,5 +29,6 @@ "selectAgent": { "title": "选择Agent", "description": "选择Agent开始对话" - } + }, + "imageGenDownloadName": "generated-{{timestamp}}.png" } diff --git a/ui/web/src/i18n/locales/zh/tools.json b/ui/web/src/i18n/locales/zh/tools.json index 76520f89..fec84ae3 100644 --- a/ui/web/src/i18n/locales/zh/tools.json +++ b/ui/web/src/i18n/locales/zh/tools.json @@ -146,7 +146,10 @@ "selectModel": "选择模型", "selectProvider": "选择Provider", "settings": "设置", - "timeout": "超时" + "timeout": "超时", + "imageModelLabel": "图像模型", + "imageModelDefaultOption": "默认 · gpt-image-2（推荐）", + "imageModelLegacyOption": "旧版 · gpt-image-1.5" }, "noMatchDescription": "请尝试其他搜索词。", "noMatchTitle": "未找到匹配工具", diff --git a/ui/web/src/pages/builtin-tools/media-provider-params-schema.ts b/ui/web/src/pages/builtin-tools/media-provider-params-schema.ts index 492c807d..a6a42088 100644 --- a/ui/web/src/pages/builtin-tools/media-provider-params-schema.ts +++ b/ui/web/src/pages/builtin-tools/media-provider-params-schema.ts @@ -12,6 +12,18 @@ export type ParamField = { export const MEDIA_PARAMS_SCHEMA: Record> = { create_image: { + chatgpt_oauth: [ + { + key: "image_model", + label: "Image model", + type: "select", + default: "gpt-image-2", + options: [ + { value: "gpt-image-2", label: "Default · gpt-image-2 (recommended)" }, + { value: "gpt-image-1.5", label: "Legacy · gpt-image-1.5" }, + ], + }, + ], minimax_native: [], bailian: [ { diff --git a/ui/web/src/types/chat.ts b/ui/web/src/types/chat.ts index 9873aecb..dfdedb89 100644 --- a/ui/web/src/types/chat.ts +++ b/ui/web/src/types/chat.ts @@ -33,6 +33,8 @@ export interface MediaItem { fileName?: string; size?: number; kind: "image" | "video" | "audio" | "document" | "code"; + /** Prompt that generated this asset (create_image, native Codex image_generation_call). */ + prompt?: string; } /** Extended message with UI-specific fields */ diff --git a/ui/web/src/types/session.ts b/ui/web/src/types/session.ts index f043ab22..53cc2517 100644 --- a/ui/web/src/types/session.ts +++ b/ui/web/src/types/session.ts @@ -33,7 +33,7 @@ export interface Message { tool_calls?: ToolCall[]; tool_call_id?: string; is_error?: boolean; - media_refs?: { id: string; mime_type: string; kind: string; path?: string }[]; + media_refs?: { id: string; mime_type: string; kind: string; path?: string; prompt?: string }[]; created_at?: string; // ISO 8601 timestamp from server; absent for older messages }