diff --git a/.gitignore b/.gitignore index e8c9945..5f91704 100644 --- a/.gitignore +++ b/.gitignore @@ -34,3 +34,6 @@ Thumbs.db # Superpowers specs and plans (local design docs, not shipped) docs/superpowers/ + +# Local agent evidence and workflow state +.omo/ diff --git a/AGENTS.md b/AGENTS.md index b108bf5..4ed9c9e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,6 +1,6 @@ # Repository Guidelines -Potaco is a Go CLI for image generation and editing via multi-provider adapters (OpenAI, fal, Vercel AI Gateway, and custom OpenAI-compatible endpoints) with encrypted credential storage and interactive TUI flows. +Potaco is a Go CLI for image generation and editing via multi-provider adapters (OpenAI, fal, Replicate, and custom OpenAI-compatible endpoints) with encrypted credential storage and interactive TUI flows. ## Project Structure @@ -10,8 +10,9 @@ models` or `potaco config set model `. Generation is assumed available fo the selected model. Edit capability is user-configured per model with `potaco config set model.edit true` or `potaco config set providers..models..edit true`; do not infer it -from provider discovery. The `custom` provider has no preset and requires a -user-supplied base URL. +from provider discovery. Generation/edit parameters are user-configured per model with `potaco config set model.parameters. ` or `potaco config set providers..models..parameters. `. +Generation/edit may fall back across configured providers with the same or similar model ID only when provider/API key/base URL are not explicitly overridden by flags or environment. +The `custom` provider has no preset and requires a user-supplied base URL. Layered monolith dependency graph: @@ -19,7 +20,7 @@ Layered monolith dependency graph: cli --> adapter, auth, config, credential, tui, image tui --> adapter, auth auth --> config, credential -adapter/openai|fal|vercel|custom --> adapter (parent), config +adapter/openai|fal|replicate|custom --> adapter (parent), config config, credential, image --> (no internal deps) ``` @@ -68,21 +69,23 @@ Config file lives at `~/.potaco/config.yaml`. Update check cache lives at `~/.po POTACO_BASE_URL=https://api.openai.com POTACO_API_KEY=sk-test \ ./potaco gen --prompt "a cat" --dry-run ``` +`potaco config show` prints the raw `config.yaml` contents without formatting or migration. Provider credentials are managed via `auth add`: ``` ./potaco auth add openai --api-key sk-... # Connect a provider ./potaco auth list # List connected providers ./potaco use openai # Switch active provider -./potaco models # Discover available models (interactive) +./potaco models # Discover available models (interactive, includes Custom) ./potaco status # Show current provider/model status +./potaco stats --non-interactive # Show local command analytics ``` ## Coding Style & Naming Conventions - Go 1.26, pure Go only (no CGO). - Standard `gofmt` formatting. Run `gofmt -l .` before committing; fix any flagged files. -- Key dependencies: `spf13/cobra` (CLI), `charmbracelet/huh` (TUI forms), `charm.land/lipgloss/v2` (styling), `golang.org/x/image` (WebP decode), `golang.org/x/term` (TTY detection), `gopkg.in/yaml.v3` (config). +- Key dependencies: `spf13/cobra` (CLI), `github.com/charmbracelet/bubbletea` (hand-crafted TUI programs), `charm.land/lipgloss/v2` (styling), `golang.org/x/image` (WebP decode), `golang.org/x/term` (TTY detection), `gopkg.in/yaml.v3` (config). - Internal image package is imported as `img` in CLI files to avoid collision with stdlib `image`. - No panics in library code. Use `fmt.Errorf` with `%w` for error wrapping. - Keep files focused: one responsibility per file, one subcommand per file in `cli/`. @@ -98,7 +101,7 @@ Provider credentials are managed via `auth add`: - Adapter tests use `httptest.Server` mocks and override `Adapter.backoff` and `Adapter.sleep` (via `SetBackoff`/`SetSleep`) to 1ms for fast retry tests. - Credential tests verify encrypt/decrypt roundtrips with test keys and temp directories. - Image tests use `t.TempDir()` for temp files and `bytes.Buffer` for in-memory roundtrips. -- TUI tests are minimal (smoke tests) since interactive forms require a TTY. +- TUI tests are minimal (smoke tests) since interactive prompts require a TTY. ## Commit & Pull Request Guidelines @@ -137,6 +140,7 @@ Subject line is lowercase, no period. Use `feat(scope):` for new features, `fix( ### Input Validation - Normalize user-supplied URLs (trim trailing slashes) before joining with path segments to avoid double-slash endpoints. +- Validate provider base URLs before model discovery or verification. Unencrypted `http://` base URLs require an explicit interactive trust confirmation; non-interactive commands should reject them. - Validate file size and dimensions before decoding user-supplied image files to prevent OOM. Use shared budget checks (`maxImageFileBytes`, `validateImageDimensionsFromBytes`) for every image input path. - Do not write multiple binary blobs to stdout in sequence. Reject combinations of flags that would produce undecodable output early with a `UserError`. - When an interactive picker returns a zero-value result on cancel, the caller must check for it and return early. Do not let a cancelled selection flow through to downstream operations. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 30f0f13..e68977a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -48,16 +48,20 @@ Provider credentials are managed via `auth add`: ```sh ./potaco auth add openai --api-key sk-... # Connect a provider +./potaco auth add replicate --api-key r8_... # Connect Replicate ./potaco auth add custom --api-key sk-... --base-url https://api.example.com/v1 # Connect custom endpoint -./potaco auth add openrouter --type openai-compatible --api-key sk-... --base-url https://openrouter.ai/api/v1 # Connect named custom endpoint +./potaco auth add openrouter --type openai --api-key sk-... --base-url https://openrouter.ai/api/v1 # Connect named OpenAI-style endpoint ./potaco auth list # List connected providers ./potaco use openai # Switch active provider -./potaco models # Pick a model interactively -./potaco models list # List available models without changing selection +./potaco models # Pick a model interactively, including Custom +./potaco models list # List available models by display name ./potaco status # Show current provider/model status +./potaco stats --non-interactive # Show local command analytics ./potaco config set model gpt-image-2 # Set active provider model ./potaco config set model.edit true # Mark active provider model edit capable +./potaco config set model.parameters.size 1536x1024 # Set active model size ./potaco config set auto_update false # Disable automatic update prompts +./potaco config show # Print raw config.yaml ``` ## Coding Style @@ -118,7 +122,7 @@ The following checks run on every PR: - Adapter tests use `httptest.Server` mocks and override `Adapter.backoff` and `Adapter.sleep` (via `SetBackoff`/`SetSleep`) to 1ms for fast retry tests - Credential tests verify encrypt/decrypt roundtrips with test keys and temp directories - Image tests use `t.TempDir()` for temp files and `bytes.Buffer` for in-memory roundtrips -- TUI tests are minimal (smoke tests) since interactive forms require a TTY +- TUI tests are minimal (smoke tests) since interactive prompts require a TTY - Use `--dry-run` for local testing without a real provider ## Releasing (Maintainers) diff --git a/DESIGN.md b/DESIGN.md new file mode 100644 index 0000000..2d50e55 --- /dev/null +++ b/DESIGN.md @@ -0,0 +1,54 @@ +# Potaco Design System + +## 1. Atmosphere & Identity + +Potaco feels like a compact creative console: direct, precise, and image-native without becoming flashy. The signature is a quiet terminal surface with clear replacement prompts, muted helper explanations, and decisive success or failure messages. + +## 2. Color + +| Role | Token | Terminal | Usage | +|------|-------|----------|-------| +| Surface/primary | `potaco-surface` | terminal default | Terminal background | +| Text/primary | `potaco-text` | terminal default | Titles, body text, options, keystrokes, entered values | +| Text/secondary | `potaco-muted` | ANSI 240 | Helper action explanations, placeholders, and captions | +| Status/success | `potaco-success` | ANSI 42 | Completed action messages only | +| Status/error | `potaco-error` | ANSI 203 | Failed, empty, destructive, and cancelled states | + +## 3. Typography + +Terminal typography uses the user's monospace font. Prompt titles are bold, options and helper keystrokes are regular white/default text, and helper action explanations are muted. No negative tracking or decorative glyphs. + +## 4. Spacing & Layout + +TUI flows follow the compact redraw pattern used by Vercel's `skills` CLI: one active prompt surface, stable row replacement, explicit completion/cancel states, and no scrollback pile-up while a multi-step flow is running. All prompts use one leading blank line, single-line titles aligned to the option/input text column, one blank line before controls, and one blank line before help. Prompt steps render in the alternate screen so completing one step replaces it with the next instead of appending below it. Lists use a two-cell cursor gutter with a Nerd Font chevron so selected and unselected rows never shift. Layout decisions must use display-cell width, not byte count, because ANSI styling and Nerd Font glyphs should not move columns. + +## 5. Components + +### Prompt Select +- **Structure**: title, option list, help row. +- **Variants**: standard provider/model selection and destructive provider removal selection. +- **States**: default option, focused option, empty option list, cancelled. +- **Accessibility**: arrow keys and vim-style `ctrl+n`/`ctrl+p`; `esc` and `ctrl+c` cancel. +- **Motion**: none. + +### Prompt Input +- **Structure**: title, input line, help row. +- **Variants**: plain text and password mask. +- **States**: empty, filled, submitted, cancelled. +- **Accessibility**: keyboard-only; typed CJK text is preserved as runes. +- **Motion**: none. + +### Prompt Confirm +- **Structure**: inline title plus `[Y/n]` typed answer input, help row. +- **Variants**: neutral and destructive copy. +- **States**: empty default yes, typed yes, typed no, submitted, cancelled. +- **Accessibility**: `y`, `n`, backspace, enter, `esc`, and `ctrl+c`. +- **Motion**: none. + +## 6. Motion & Interaction + +Terminal prompts do not animate. State changes are immediate and visible through replacement, bold titles, cursor position, and status lines. + +## 7. Depth & Surface + +Depth strategy is tonal-shift only. The TUI uses text hierarchy and muted captions instead of boxes, shadows, accent-heavy styling, or decorative borders. diff --git a/README.md b/README.md index ae43fc9..ba022c0 100644 --- a/README.md +++ b/README.md @@ -9,14 +9,14 @@ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) [![Latest Release](https://img.shields.io/github/v/release/ncxton/potaco)](https://github.com/ncxton/potaco/releases/latest) -Potaco is a Go CLI for image generation and editing via multiple AI image providers. Connect to OpenAI, fal, Vercel AI Gateway, or any OpenAI-compatible API via the custom provider. +Potaco is a Go CLI for image generation and editing via multiple AI image providers. Connect to OpenAI, fal, Replicate, or any OpenAI-compatible API via the custom provider. > [!WARNING] > This project is still in **beta**. Therefore, critical breakages or bugs are to be expected. Use at your own risk, and please report any issues you encounter. ## Features -- **Multi-provider** support for OpenAI, fal, Vercel AI Gateway, and custom OpenAI-compatible endpoints +- **Multi-provider** support for OpenAI, fal, Replicate, and custom OpenAI-compatible endpoints - **Encrypted credentials** stored locally - **Auth management** with `auth add/remove/list` commands - **Provider switching** with `potaco use ` @@ -24,8 +24,8 @@ Potaco is a Go CLI for image generation and editing via multiple AI image provid - **Generate** images from text prompts with size, quality, seed, and guidance control - **Edit** existing images with inpainting (mask-based) and outpainting (canvas extension) - **Status** and **models** commands to inspect current state and available models -- **Interactive TUI** with interactive forms for auth, model selection, and provider switching -- **Info** inspect image metadata (dimensions, format, file size, color model) +- **Stats** dashboard for local command analytics and estimated spend +- **Interactive TUI** with hand-crafted prompts for auth, model selection, and provider switching - **Retry** with exponential backoff on rate-limit and server errors - **JSON output** for scripting and piping - **Dry-run mode** to validate requests without calling the API @@ -62,7 +62,7 @@ Supported platforms: ## Quick Start ```sh -# Connect a provider (interactive TUI prompts for key, verification, and model) +# Connect a provider (hand-crafted TUI prompts for key, verification, and model) potaco auth add openai # Or non-interactive: @@ -77,8 +77,8 @@ potaco use fal # List or pick available models from the active provider potaco models # interactive picker (persists the selected model) -potaco models list # static list of models -potaco models list openai # static list for a specific provider +potaco models list # discovered model list +potaco models list openai # discovered list for a specific provider # Check current status potaco status @@ -94,13 +94,13 @@ Generate new images from a text prompt. |------|------|----------|---------|-------------| | `--prompt`, `-p` | string | yes | | Text description of the desired image(s) | | `--model` | string | no | from config | Model to use (e.g., `gpt-image-2`) | -| `--size` | string | no | `1024x1024` | Image dimensions (WxH) | -| `--quality` | string | no | `auto` | Image quality (`low`, `medium`, `high`, or `auto`) | -| `--n` | int | no | `1` | Number of images to generate | -| `--seed` | int | no | `0` | Reproducibility seed | -| `--guidance-scale` | float | no | `0` | Guidance scale | -| `--negative-prompt` | string | no | | Negative prompt | -| `--response-format` | string | no | `b64_json` | Response format (`url` or `b64_json`) | +| `--size` | string | no | from model config | Override image dimensions (WxH) | +| `--quality` | string | no | from model config | Override image quality (`low`, `medium`, `high`, or `auto`) | +| `--n` | int | no | from model config | Override number of images to generate | +| `--seed` | int | no | from model config | Override reproducibility seed | +| `--guidance-scale` | float | no | from model config | Override guidance scale | +| `--negative-prompt` | string | no | from model config | Override negative prompt | +| `--response-format` | string | no | from model config | Override response format (`url` or `b64_json`) | | `--output`, `-o` | string | no | auto | Output file path (auto: `potaco-YYYYMMDD-HHMMSS.png`) | | `--output-format` | string | no | `png` | Output format (`png` or `jpeg`) | | `--stdout` | bool | no | `false` | Pipe raw image bytes to stdout | @@ -116,12 +116,18 @@ Generate new images from a text prompt. ```sh potaco gen --prompt "a red fox in a forest" -potaco gen --prompt "a cityscape at night" --size 1536x1024 --quality high --n 2 -potaco gen --prompt "portrait of a woman" --seed 42 --json +potaco config set model.parameters.size 1536x1024 +potaco config set model.parameters.quality high +potaco gen --prompt "a cityscape at night" +potaco config set model.parameters.seed 42 +potaco gen --prompt "portrait of a woman" --json +potaco gen --prompt "portrait of a woman" --size 1024x1792 --quality medium potaco gen --prompt "test" --dry-run potaco gen --prompt "test" --provider fal --api-key --dry-run ``` +When no provider, API key, or base URL override is supplied, generation can fall back to another configured provider with the same or similar model ID after a transient provider failure such as a rate limit, timeout, or server error. Fallback providers are tried after the active provider in A-Z order. + ### `potaco edit` -- Edit existing images Edit an existing image with inpainting (mask-based editing) or outpainting (canvas extension). @@ -135,9 +141,9 @@ Edit an existing image with inpainting (mask-based editing) or outpainting (canv | `--mask-circle` | string | no | | Circular mask: `x,y,r` in pixels | | `--extend` | string | no | | Outpaint: `top=N,bottom=N,left=N,right=N` or `all=N` | | `--model` | string | no | from config | Model to use | -| `--size` | string | no | `1024x1024` | Image dimensions (WxH) | -| `--n` | int | no | `1` | Number of images to generate | -| `--response-format` | string | no | `b64_json` | Response format (`url` or `b64_json`) | +| `--size` | string | no | from model config | Override image dimensions (WxH) | +| `--n` | int | no | from model config | Override number of images to generate | +| `--response-format` | string | no | from model config | Override response format (`url` or `b64_json`) | | `--output`, `-o` | string | no | auto | Output file path | | `--output-format` | string | no | `png` | Output format (`png` or `jpeg`) | | `--stdout` | bool | no | `false` | Pipe raw image bytes to stdout | @@ -164,6 +170,8 @@ potaco config set model.edit true potaco config set providers.openai.models.gpt-image-2.edit true ``` +Edit requests use the same transient fallback behavior as generation, but only to configured providers whose matching model is marked edit capable. + ### `potaco auth` -- Manage provider credentials Connect providers, add API keys, and manage credentials. Credentials are encrypted at rest. @@ -172,9 +180,9 @@ Connect providers, add API keys, and manage credentials. Credentials are encrypt potaco auth add openai # interactive TUI flow potaco auth add openai --api-key sk-xxx # non-interactive potaco auth add fal --api-key # connect fal -potaco auth add vercel --api-key # connect Vercel AI Gateway +potaco auth add replicate --api-key # connect Replicate potaco auth add custom --api-key --base-url # connect any OpenAI-compatible endpoint -potaco auth add openrouter --type openai-compatible --api-key --base-url # named custom endpoint +potaco auth add openrouter --type openai --api-key --base-url # named OpenAI-style endpoint potaco auth add staging-openai --type openai --api-key --base-url # named provider using a built-in adapter type potaco auth remove openai # disconnect a provider potaco auth list # list connected providers @@ -202,20 +210,33 @@ potaco status # text output potaco status --json # JSON output ``` +### `potaco stats` -- Command analytics + +Show a local analytics dashboard for command count, duration, errors, providers, recent runs, and estimated cost. Interactive terminals get a btop-style dashboard; non-interactive shells get a plain text report. + +```sh +potaco stats # interactive dashboard +potaco stats --non-interactive # plain text fallback +potaco --json stats --non-interactive # JSON summary +potaco stats --clear # clear local stats history +``` + +Dashboard keys: `tab` switches focused panel, `r` refreshes, `c` clears local stats, and `q` or `ctrl+c` closes. Cost is estimated from OpenRouter model pricing when available. + ### `potaco models` -- Pick or list available models -Discover models from the active provider or a specified provider. By default, the command launches an interactive picker that persists the selected model to config and asks whether it can edit images. In non-interactive mode it falls back to a static list. +Discover models from the active provider or a specified provider. By default, the command launches an interactive picker that persists the selected model to config and asks whether it can edit images. In non-interactive mode it prints the discovered model list without changing selection. ```sh -potaco models # interactive picker for the active provider -potaco models openai # interactive picker for a specific provider -potaco models --non-interactive # static list for the active provider -potaco models list # static list for the active provider -potaco models list openai # static list for a specific provider -potaco models --json --non-interactive # JSON output (static list) +potaco models # interactive picker with Custom option for the active provider +potaco models openai # interactive picker with Custom option for a specific provider +potaco models --non-interactive # discovered list by display name for the active provider +potaco models list # discovered list by display name for the active provider +potaco models list openai # discovered list by display name for a specific provider +potaco models --json --non-interactive # JSON output ``` -Use `potaco models list` when you only want to see available model IDs without changing the active model. Model discovery does not configure or display generation/edit capability; generation is assumed available, and edit capability is user-configured. +Use `potaco models list` when you only want to see available models without changing the active model. Interactive pickers include a bottom `Custom` option for manually entering a model ID and display name. Custom display names are saved with a `[Custom]` suffix. Model discovery does not configure or display generation/edit capability; generation is assumed available, and edit capability is user-configured. ### `potaco config` -- Provider settings @@ -223,24 +244,24 @@ Manage per-provider settings stored in `~/.potaco/config.yaml`. API keys are sto ```sh potaco config set model gpt-image-2 # change active provider model -potaco config set providers.vercel.model openai/gpt-image-2 # set a provider model +potaco config set providers.replicate.model black-forest-labs/flux-schnell # set a provider model potaco config set model.edit true # mark active provider model edit capable potaco config set providers.openai.models.gpt-image-2.edit true # mark a specific model edit capable +potaco config set model.parameters.size 1536x1024 # set active model image size +potaco config set model.parameters.quality high # set active model quality +potaco config set model.parameters.n 2 # set active model output count +potaco config set model.parameters.response_format url # set active model response format +potaco config set model.parameters.seed 42 # set active model seed +potaco config set model.parameters.guidance_scale 7.5 # set active model guidance +potaco config set model.parameters.negative_prompt blurry # set active model negative prompt potaco config set base_url https://api.example.com/v1 # set active provider base URL potaco config set retries 3 # set active provider retries potaco config set timeout 120 # set active provider timeout in seconds potaco config set auto_update false # disable automatic update prompts -potaco config show # display current config +potaco config show # print raw config.yaml ``` -### `potaco info` -- Image metadata - -Print metadata about an image file. - -```sh -potaco info output.png -potaco info output.png --json -``` +Generation and edit parameters are stored per model under `providers..models..parameters`. Command flags such as `--size`, `--quality`, and `--response-format` still work and override the model configuration for that run. `potaco edit` warns when generation-only parameters such as `quality`, `seed`, `guidance_scale`, or `negative_prompt` are configured but not applied to edit requests. ### `potaco version` -- Print version @@ -302,7 +323,7 @@ providers: timeout: 120 ``` -The `base_url` field is optional for built-in provider names (`openai`, `fal`, `vercel`) and overrides their preset URL. It is required for any other provider name, including `custom` and aliases that use a built-in adapter type. +The `base_url` field is optional for built-in provider names (`openai`, `fal`, `replicate`) and overrides their preset URL. It is required for any other provider name, including `custom` and aliases that use a built-in adapter type. Base URLs must be full `http://` or `https://` URLs; interactive flows warn before continuing with unencrypted `http://` endpoints. API keys are stored separately in `~/.potaco/credentials.enc`, encrypted with a machine-derived key. Update check metadata is stored separately in `~/.potaco/.potaco.json`. @@ -311,7 +332,7 @@ Update check metadata is stored separately in `~/.potaco/.potaco.json`. | Variable | Description | |----------|-------------| -| `POTACO_PROVIDER` | Active provider name (e.g., `openai`, `fal`, `vercel`, `custom`) | +| `POTACO_PROVIDER` | Active provider name (e.g., `openai`, `fal`, `replicate`, `custom`) | | `POTACO_API_KEY` | API key for the active provider | | `POTACO_MODEL` | Default model for the active provider | | `POTACO_BASE_URL` | Override the provider's base URL (required for `custom` and provider aliases) | @@ -327,7 +348,9 @@ The known provider names below ship with preset base URLs. Any other auth name, |----------|-----------|--------------| | `openai` | Bearer | Yes | | `fal` | Key | Yes | -| `vercel` | Bearer | No | +| `replicate` | Bearer | No | + +Replicate model discovery is not automatic; configure an image model ID manually in `owner/model` format, for example `potaco config set model black-forest-labs/flux-schnell`. ## Contributing diff --git a/go.mod b/go.mod index be10274..5d915ed 100644 --- a/go.mod +++ b/go.mod @@ -3,10 +3,9 @@ module github.com/ncxton/potaco go 1.26 require ( + charm.land/bubbletea/v2 v2.0.7 charm.land/lipgloss/v2 v2.0.4 - github.com/charmbracelet/bubbles v1.0.0 - github.com/charmbracelet/bubbletea v1.3.10 - github.com/charmbracelet/huh v1.0.0 + github.com/charmbracelet/colorprofile v0.4.3 github.com/spf13/cobra v1.10.2 github.com/spf13/pflag v1.0.10 golang.org/x/image v0.43.0 @@ -15,34 +14,19 @@ require ( ) require ( - github.com/atotto/clipboard v0.1.4 // indirect - github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect - github.com/catppuccin/go v0.3.0 // indirect - github.com/charmbracelet/colorprofile v0.4.3 // indirect - github.com/charmbracelet/lipgloss v1.1.0 // indirect - github.com/charmbracelet/ultraviolet v0.0.0-20251205161215-1948445e3318 // indirect + github.com/charmbracelet/ultraviolet v0.0.0-20260525132238-948f4557a654 // indirect github.com/charmbracelet/x/ansi v0.11.7 // indirect - github.com/charmbracelet/x/cellbuf v0.0.15 // indirect - github.com/charmbracelet/x/exp/strings v0.0.0-20240722160745-212f7b056ed0 // indirect github.com/charmbracelet/x/term v0.2.2 // indirect github.com/charmbracelet/x/termios v0.1.1 // indirect github.com/charmbracelet/x/windows v0.2.2 // indirect github.com/clipperhouse/displaywidth v0.11.0 // indirect github.com/clipperhouse/uax29/v2 v2.7.0 // indirect - github.com/dustin/go-humanize v1.0.1 // indirect - github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/lucasb-eyer/go-colorful v1.4.0 // indirect - github.com/mattn/go-isatty v0.0.20 // indirect - github.com/mattn/go-localereader v0.0.1 // indirect github.com/mattn/go-runewidth v0.0.23 // indirect - github.com/mitchellh/hashstructure/v2 v2.0.2 // indirect - github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect github.com/muesli/cancelreader v0.2.2 // indirect - github.com/muesli/termenv v0.16.0 // indirect github.com/rivo/uniseg v0.4.7 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect golang.org/x/sync v0.21.0 // indirect golang.org/x/sys v0.46.0 // indirect - golang.org/x/text v0.38.0 // indirect ) diff --git a/go.sum b/go.sum index ab805ef..bbd37d1 100644 --- a/go.sum +++ b/go.sum @@ -1,76 +1,36 @@ +charm.land/bubbletea/v2 v2.0.7 h1:7qw2tTAVar7m7klOPBYfTB0mniv/RuexsYwMRNxSeL0= +charm.land/bubbletea/v2 v2.0.7/go.mod h1:DGW2q8gvzHnOpMpZTORs0aySVHCox5C+2Svk0fci1qs= charm.land/lipgloss/v2 v2.0.4 h1:lcPeVtcp23SNra7lHy8iYE4UC2aIipVQ47sbGyyxR5Q= charm.land/lipgloss/v2 v2.0.4/go.mod h1:0653x8epbZSzdDfO/XPS1a/uYPOBeSsCssOpJOqDzik= -github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ= -github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE= -github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= -github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= -github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= -github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/aymanbagabas/go-udiff v0.4.1 h1:OEIrQ8maEeDBXQDoGCbbTTXYJMYRCRO1fnodZ12Gv5o= github.com/aymanbagabas/go-udiff v0.4.1/go.mod h1:0L9PGwj20lrtmEMeyw4WKJ/TMyDtvAoK9bf2u/mNo3w= -github.com/catppuccin/go v0.3.0 h1:d+0/YicIq+hSTo5oPuRi5kOpqkVA5tAsU6dNhvRu+aY= -github.com/catppuccin/go v0.3.0/go.mod h1:8IHJuMGaUUjQM82qBrGNBv7LFq6JI3NnQCF6MOlZjpc= -github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc= -github.com/charmbracelet/bubbles v1.0.0/go.mod h1:9d/Zd5GdnauMI5ivUIVisuEm3ave1XwXtD1ckyV6r3E= -github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw= -github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4= github.com/charmbracelet/colorprofile v0.4.3 h1:QPa1IWkYI+AOB+fE+mg/5/4HRMZcaXex9t5KX76i20Q= github.com/charmbracelet/colorprofile v0.4.3/go.mod h1:/zT4BhpD5aGFpqQQqw7a+VtHCzu+zrQtt1zhMt9mR4Q= -github.com/charmbracelet/huh v1.0.0 h1:wOnedH8G4qzJbmhftTqrpppyqHakl/zbbNdXIWJyIxw= -github.com/charmbracelet/huh v1.0.0/go.mod h1:5YVc+SlZ1IhQALxRPpkGwwEKftN/+OlJlnJYlDRFqN4= -github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY= -github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30= -github.com/charmbracelet/ultraviolet v0.0.0-20251205161215-1948445e3318 h1:OqDqxQZliC7C8adA7KjelW3OjtAxREfeHkNcd66wpeI= -github.com/charmbracelet/ultraviolet v0.0.0-20251205161215-1948445e3318/go.mod h1:Y6kE2GzHfkyQQVCSL9r2hwokSrIlHGzZG+71+wDYSZI= +github.com/charmbracelet/ultraviolet v0.0.0-20260525132238-948f4557a654 h1:FpSYhY28ucg9ZRr+2wj67FAQ0Ey5yiK0072PmRDJNek= +github.com/charmbracelet/ultraviolet v0.0.0-20260525132238-948f4557a654/go.mod h1:hFpumms29Smx3LStRfku8vcCTBe1Kq8aCXtHUJa3mjY= github.com/charmbracelet/x/ansi v0.11.7 h1:kzv1kJvjg2S3r9KHo8hDdHFQLEqn4RBCb39dAYC84jI= github.com/charmbracelet/x/ansi v0.11.7/go.mod h1:9qGpnAVYz+8ACONkZBUWPtL7lulP9No6p1epAihUZwQ= -github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI= -github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q= -github.com/charmbracelet/x/conpty v0.1.0 h1:4zc8KaIcbiL4mghEON8D72agYtSeIgq8FSThSPQIb+U= -github.com/charmbracelet/x/conpty v0.1.0/go.mod h1:rMFsDJoDwVmiYM10aD4bH2XiRgwI7NYJtQgl5yskjEQ= -github.com/charmbracelet/x/errors v0.0.0-20240508181413-e8d8b6e2de86 h1:JSt3B+U9iqk37QUU2Rvb6DSBYRLtWqFqfxf8l5hOZUA= -github.com/charmbracelet/x/errors v0.0.0-20240508181413-e8d8b6e2de86/go.mod h1:2P0UgXMEa6TsToMSuFqKFQR+fZTO9CNGUNokkPatT/0= github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f h1:pk6gmGpCE7F3FcjaOEKYriCvpmIN4+6OS/RD0vm4uIA= github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f/go.mod h1:IfZAMTHB6XkZSeXUqriemErjAWCCzT0LwjKFYCZyw0I= -github.com/charmbracelet/x/exp/strings v0.0.0-20240722160745-212f7b056ed0 h1:qko3AQ4gK1MTS/de7F5hPGx6/k1u0w4TeYmBFwzYVP4= -github.com/charmbracelet/x/exp/strings v0.0.0-20240722160745-212f7b056ed0/go.mod h1:pBhA0ybfXv6hDjQUZ7hk1lVxBiUbupdw5R31yPUViVQ= github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk= github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI= github.com/charmbracelet/x/termios v0.1.1 h1:o3Q2bT8eqzGnGPOYheoYS8eEleT5ZVNYNy8JawjaNZY= github.com/charmbracelet/x/termios v0.1.1/go.mod h1:rB7fnv1TgOPOyyKRJ9o+AsTU/vK5WHJ2ivHeut/Pcwo= github.com/charmbracelet/x/windows v0.2.2 h1:IofanmuvaxnKHuV04sC0eBy/smG6kIKrWG2/jYn2GuM= github.com/charmbracelet/x/windows v0.2.2/go.mod h1:/8XtdKZzedat74NQFn0NGlGL4soHB0YQZrETF96h75k= -github.com/charmbracelet/x/xpty v0.1.2 h1:Pqmu4TEJ8KeA9uSkISKMU3f+C1F6OGBn8ABuGlqCbtI= -github.com/charmbracelet/x/xpty v0.1.2/go.mod h1:XK2Z0id5rtLWcpeNiMYBccNNBrP2IJnzHI0Lq13Xzq4= github.com/clipperhouse/displaywidth v0.11.0 h1:lBc6kY44VFw+TDx4I8opi/EtL9m20WSEFgwIwO+UVM8= github.com/clipperhouse/displaywidth v0.11.0/go.mod h1:bkrFNkf81G8HyVqmKGxsPufD3JhNl3dSqnGhOoSD/o0= github.com/clipperhouse/uax29/v2 v2.7.0 h1:+gs4oBZ2gPfVrKPthwbMzWZDaAFPGYK72F0NJv2v7Vk= github.com/clipperhouse/uax29/v2 v2.7.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= -github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s= -github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= -github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= -github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= -github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= -github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4= github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= -github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= -github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= -github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4= -github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= github.com/mattn/go-runewidth v0.0.23 h1:7ykA0T0jkPpzSvMS5i9uoNn2Xy3R383f9HDx3RybWcw= github.com/mattn/go-runewidth v0.0.23/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= -github.com/mitchellh/hashstructure/v2 v2.0.2 h1:vGKWl0YJqUNxE8d+h8f6NJLcCJrgbhC4NcD46KavDd4= -github.com/mitchellh/hashstructure/v2 v2.0.2/go.mod h1:MG3aRVU/N29oo/V/IhBX8GR/zz4kQkprJgF2EVszyDE= -github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= -github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= -github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= -github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= @@ -88,14 +48,10 @@ golang.org/x/image v0.43.0 h1:FLxcP4ec2350nTfOC8ysKtqYSIFbk/QGjw1ZHNP4tsY= golang.org/x/image v0.43.0/go.mod h1:rrpelvGFt+kLPAjPM4HeWPgrl0FtafueU//e5N0qk/Q= golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/term v0.44.0 h1:0rLvDRCtNj0gZkyIXhCyOb2OAzEhLVqc4B+hrsBhrmc= golang.org/x/term v0.44.0/go.mod h1:7ze4MdzUzLXpSAoFP1H0bOI9aXDqveSvatT5vKcFh2Y= -golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= -golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= diff --git a/internal/adapter/custom/discover.go b/internal/adapter/custom/discover.go index ec07647..82ba3ee 100644 --- a/internal/adapter/custom/discover.go +++ b/internal/adapter/custom/discover.go @@ -9,9 +9,8 @@ import ( "github.com/ncxton/potaco/internal/adapter" ) -// DiscoverModels calls GET /v1/models and returns every model from the -// endpoint. Unlike openai, we do not filter for known image model IDs -// because a custom provider may expose any model identifiers. +// DiscoverModels calls GET /v1/models and filters the endpoint through +// OpenRouter's image model catalog. // On any API failure it returns an error instead of falling back to a // hardcoded list. func (a *Adapter) DiscoverModels(ctx context.Context) ([]adapter.Model, error) { @@ -42,11 +41,20 @@ func (a *Adapter) DiscoverModels(ctx context.Context) ([]adapter.Model, error) { return nil, fmt.Errorf("decode models response: %w", err) } + imageModels, err := adapter.OpenRouterImageModelNames(ctx, a.http) + if err != nil { + return nil, err + } + var models []adapter.Model for _, m := range result.Data { + displayName, ok := imageModels[m.ID] + if !ok { + continue + } models = append(models, adapter.Model{ ID: m.ID, - DisplayName: m.ID, + DisplayName: displayName, SupportsGen: true, SupportsEdit: true, }) diff --git a/internal/adapter/custom/discover_test.go b/internal/adapter/custom/discover_test.go index 7fd8b5a..c38781e 100644 --- a/internal/adapter/custom/discover_test.go +++ b/internal/adapter/custom/discover_test.go @@ -3,6 +3,7 @@ package custom import ( "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "strings" @@ -11,7 +12,41 @@ import ( "github.com/ncxton/potaco/internal/adapter" ) -func TestCustomDiscoverModelsReturnsAllModels(t *testing.T) { +func setOpenRouterImageCatalog(t *testing.T, ids ...string) { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + data := make([]map[string]any, 0, len(ids)) + for _, id := range ids { + data = append(data, map[string]any{ + "id": id, + "name": "Display " + id, + "architecture": map[string]any{ + "output_modalities": []string{"image"}, + }, + }) + } + json.NewEncoder(w).Encode(map[string]any{"data": data}) + })) + t.Cleanup(srv.Close) + t.Setenv("POTACO_OPENROUTER_MODELS_URL", srv.URL+"/models?output_modalities=image") +} + +func setOpenRouterCatalogTrap(t *testing.T) *int { + t.Helper() + calls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + t.Setenv("POTACO_OPENROUTER_MODELS_URL", srv.URL+"/models?output_modalities=image") + return &calls +} + +func TestCustomDiscoverModelsFiltersWithOpenRouterImageCatalog(t *testing.T) { + setOpenRouterImageCatalog(t, "custom/custom-image-model", "custom/another-image-model") + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/models" && r.URL.Path != "/models" { t.Errorf("path = %q, want /v1/models", r.URL.Path) @@ -42,19 +77,47 @@ func TestCustomDiscoverModelsReturnsAllModels(t *testing.T) { if !m.SupportsEdit { t.Errorf("model %s should have SupportsEdit=true", m.ID) } + if !strings.HasPrefix(m.DisplayName, "Display custom/") { + t.Errorf("model %s display name = %q, want OpenRouter name", m.ID, m.DisplayName) + } } if !ids["custom-image-model"] { t.Error("should include custom-image-model") } - if !ids["text-model-1"] { - t.Error("should include text-model-1 (no filtering)") + if ids["text-model-1"] { + t.Error("should not include text-model-1") } if !ids["another-image-model"] { t.Error("should include another-image-model") } } +func TestCustomDiscoverModelsReturnsNoFallbackWhenCatalogHasNoMatch(t *testing.T) { + setOpenRouterImageCatalog(t, "custom/image-model") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "text-only-model", "object": "model", "owned_by": "custom"}, + }, + }) + })) + defer server.Close() + + a := New("sk-test", adapter.AdapterOpts{BaseURL: server.URL}) + + models, err := a.DiscoverModels(context.Background()) + if !errors.Is(err, adapter.ErrDiscoveryFailed) { + t.Fatalf("DiscoverModels error = %v, want ErrDiscoveryFailed", err) + } + if len(models) != 0 { + t.Fatalf("models len = %d, want 0", len(models)) + } +} + func TestCustomDiscoverModelsError(t *testing.T) { + openRouterCalls := setOpenRouterCatalogTrap(t) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusUnauthorized) })) @@ -66,9 +129,14 @@ func TestCustomDiscoverModelsError(t *testing.T) { if err == nil { t.Fatal("DiscoverModels should return error on API failure") } + if *openRouterCalls != 0 { + t.Fatalf("OpenRouter catalog calls = %d, want 0 after provider failure", *openRouterCalls) + } } func TestCustomDiscoverModelsMalformedJSON(t *testing.T) { + openRouterCalls := setOpenRouterCatalogTrap(t) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) w.Write([]byte("not-json")) @@ -81,6 +149,9 @@ func TestCustomDiscoverModelsMalformedJSON(t *testing.T) { if err == nil { t.Fatal("DiscoverModels should return error on malformed JSON") } + if *openRouterCalls != 0 { + t.Fatalf("OpenRouter catalog calls = %d, want 0 after provider decode failure", *openRouterCalls) + } } func TestCustomVerifySuccess(t *testing.T) { @@ -147,10 +218,3 @@ func TestCustomVerifyServerError(t *testing.T) { t.Errorf("error should mention verification failed or 500, got: %v", err) } } - -func TestCustomNoFallbackModels(t *testing.T) { - // This is a compile-time/static check: ensure no fallbackModels or - // hardcodedModelParams exist in the custom package. The custom - // adapter must not contain hardcoded model lists. - t.Log("custom adapter intentionally has no fallbackModels or hardcodedModelParams") -} diff --git a/internal/adapter/fal/discover.go b/internal/adapter/fal/discover.go index ee626fd..24a886d 100644 --- a/internal/adapter/fal/discover.go +++ b/internal/adapter/fal/discover.go @@ -43,7 +43,15 @@ func (a *Adapter) DiscoverModels(ctx context.Context) (models []adapter.Model, e return nil, fmt.Errorf("decode models response: %w", err) } + imageModels, err := adapter.OpenRouterImageModelIDs(ctx, a.http) + if err != nil { + return nil, err + } + for _, model := range result.Models { + if !imageModels[model.ID] { + continue + } displayName := model.Metadata.DisplayName if displayName == "" { displayName = model.ID diff --git a/internal/adapter/fal/fal_test.go b/internal/adapter/fal/fal_test.go index 1b46cf1..8f9d32c 100644 --- a/internal/adapter/fal/fal_test.go +++ b/internal/adapter/fal/fal_test.go @@ -11,6 +11,37 @@ import ( "github.com/ncxton/potaco/internal/adapter" ) +func setOpenRouterImageCatalog(t *testing.T, ids ...string) { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + data := make([]map[string]any, 0, len(ids)) + for _, id := range ids { + data = append(data, map[string]any{ + "id": id, + "architecture": map[string]any{ + "output_modalities": []string{"image"}, + }, + }) + } + json.NewEncoder(w).Encode(map[string]any{"data": data}) + })) + t.Cleanup(srv.Close) + t.Setenv("POTACO_OPENROUTER_MODELS_URL", srv.URL+"/models?output_modalities=image") +} + +func setOpenRouterCatalogTrap(t *testing.T) *int { + t.Helper() + calls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + t.Setenv("POTACO_OPENROUTER_MODELS_URL", srv.URL+"/models?output_modalities=image") + return &calls +} + func TestFalAdapterName(t *testing.T) { ad := New("test-key", adapter.AdapterOpts{}) if ad.Name() != "fal" { @@ -77,6 +108,8 @@ func TestFalRegisteredInRegistry(t *testing.T) { } func TestFalDiscoverModels(t *testing.T) { + setOpenRouterImageCatalog(t, "fal-ai/flux/dev", "fal-ai/nano-banana") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/models" { t.Errorf("path = %q, want /v1/models", r.URL.Path) @@ -123,12 +156,15 @@ func TestFalDiscoverModels(t *testing.T) { if err != nil { t.Fatalf("DiscoverModels: %v", err) } - if len(models) != 3 { - t.Fatalf("models len = %d, want 3", len(models)) + if len(models) != 2 { + t.Fatalf("models len = %d, want 2", len(models)) } if models[0].ID != "fal-ai/flux/dev" { t.Errorf("model[0] ID = %q, want 'fal-ai/flux/dev'", models[0].ID) } + if models[1].ID != "fal-ai/nano-banana" { + t.Errorf("model[1] ID = %q, want 'fal-ai/nano-banana'", models[1].ID) + } if models[0].DisplayName != "Flux Dev" { t.Errorf("model[0] DisplayName = %q, want 'Flux Dev'", models[0].DisplayName) } @@ -138,18 +174,13 @@ func TestFalDiscoverModels(t *testing.T) { if !models[0].SupportsEdit { t.Error("model[0] should support edit") } - if models[1].ID != "fal-ai/flux/dev/image-to-image" { - t.Errorf("model[1] ID = %q", models[1].ID) - } if !models[1].SupportsEdit { t.Error("model[1] should support edit") } - if !models[2].SupportsEdit { - t.Error("model[2] should support edit") - } } func TestFalDiscoverModelsError(t *testing.T) { + openRouterCalls := setOpenRouterCatalogTrap(t) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusInternalServerError) })) @@ -163,6 +194,9 @@ func TestFalDiscoverModelsError(t *testing.T) { if err == nil { t.Fatal("DiscoverModels should return error on API failure") } + if *openRouterCalls != 0 { + t.Fatalf("OpenRouter catalog calls = %d, want 0 after provider failure", *openRouterCalls) + } } func TestFalVerify(t *testing.T) { diff --git a/internal/adapter/openai/discover.go b/internal/adapter/openai/discover.go index 5239e37..b92421a 100644 --- a/internal/adapter/openai/discover.go +++ b/internal/adapter/openai/discover.go @@ -40,16 +40,22 @@ func (a *Adapter) DiscoverModels(ctx context.Context) ([]adapter.Model, error) { return nil, fmt.Errorf("decode models response: %w", err) } + imageModels, err := adapter.OpenRouterImageModelNames(ctx, a.http) + if err != nil { + return nil, err + } + var models []adapter.Model for _, m := range result.Data { - if !imageModelIDs[m.ID] { + displayName, ok := imageModels[m.ID] + if !ok { continue } models = append(models, adapter.Model{ ID: m.ID, - DisplayName: m.ID, + DisplayName: displayName, SupportsGen: true, - SupportsEdit: editCapableModels[m.ID], + SupportsEdit: editCapableModels[m.ID] || editCapableModels[adapter.SuffixModelID(m.ID)], }) } if len(models) == 0 { diff --git a/internal/adapter/openai/models.go b/internal/adapter/openai/models.go index 90222b9..fd93f40 100644 --- a/internal/adapter/openai/models.go +++ b/internal/adapter/openai/models.go @@ -1,11 +1,5 @@ package openai -// imageModelIDs is the set of OpenAI model IDs that support image generation. -// Only gpt-image-2 is current; the rest are deprecated. -var imageModelIDs = map[string]bool{ - "gpt-image-2": true, -} - // editCapableModels is the set of image model IDs that support image editing. var editCapableModels = map[string]bool{ "gpt-image-2": true, diff --git a/internal/adapter/openai/openrouter_discover_test.go b/internal/adapter/openai/openrouter_discover_test.go new file mode 100644 index 0000000..ed0201a --- /dev/null +++ b/internal/adapter/openai/openrouter_discover_test.go @@ -0,0 +1,142 @@ +package openai + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "testing" + + "github.com/ncxton/potaco/internal/adapter" +) + +var openRouterCatalogServer *httptest.Server + +func TestMain(m *testing.M) { + openRouterCatalogServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + openRouterCatalogModel("openai/gpt-image-2"), + }, + }) + })) + os.Setenv("POTACO_OPENROUTER_MODELS_URL", openRouterCatalogServer.URL+"/models?output_modalities=image") + code := m.Run() + openRouterCatalogServer.Close() + os.Exit(code) +} + +func TestDiscoverModelsUsesOpenRouterImageCatalogForAnyBaseURL(t *testing.T) { + openRouterSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/v1/models" { + t.Errorf("openrouter path = %q, want /api/v1/models", r.URL.Path) + } + if r.URL.Query().Get("output_modalities") != "image" { + t.Errorf("output_modalities = %q, want image", r.URL.Query().Get("output_modalities")) + } + if auth := r.Header.Get("Authorization"); auth != "" { + t.Errorf("OpenRouter Authorization header = %q, want empty", auth) + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + { + "id": "test-provider/image-model", + "name": "Image Model", + "architecture": map[string]any{ + "output_modalities": []string{"image"}, + }, + }, + }, + }) + })) + defer openRouterSrv.Close() + t.Setenv("POTACO_OPENROUTER_MODELS_URL", openRouterSrv.URL+"/api/v1/models?output_modalities=image") + + providerSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Errorf("provider path = %q, want /v1/models", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "image-model", "owned_by": "test-provider"}, + {"id": "text-model", "owned_by": "test-provider"}, + }, + }) + })) + defer providerSrv.Close() + + a := New("sk-test", adapter.AdapterOpts{BaseURL: providerSrv.URL}) + models, err := a.DiscoverModels(context.Background()) + if err != nil { + t.Fatalf("DiscoverModels error: %v", err) + } + if len(models) != 1 { + t.Fatalf("models len = %d, want 1", len(models)) + } + if models[0].ID != "image-model" { + t.Fatalf("model ID = %q, want image-model", models[0].ID) + } + if models[0].DisplayName != "Image Model" { + t.Fatalf("model display name = %q, want Image Model", models[0].DisplayName) + } +} + +func TestDiscoverModelsDoesNotCallOpenRouterAfterProviderError(t *testing.T) { + openRouterCalls := setOpenRouterCatalogTrap(t) + providerSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer providerSrv.Close() + + a := New("sk-test", adapter.AdapterOpts{BaseURL: providerSrv.URL}) + _, err := a.DiscoverModels(context.Background()) + if err == nil { + t.Fatal("DiscoverModels should return error on provider failure") + } + if *openRouterCalls != 0 { + t.Fatalf("OpenRouter catalog calls = %d, want 0 after provider failure", *openRouterCalls) + } +} + +func TestDiscoverModelsDoesNotCallOpenRouterAfterProviderMalformedJSON(t *testing.T) { + openRouterCalls := setOpenRouterCatalogTrap(t) + providerSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("not-json")) + })) + defer providerSrv.Close() + + a := New("sk-test", adapter.AdapterOpts{BaseURL: providerSrv.URL}) + _, err := a.DiscoverModels(context.Background()) + if err == nil { + t.Fatal("DiscoverModels should return error on malformed provider JSON") + } + if *openRouterCalls != 0 { + t.Fatalf("OpenRouter catalog calls = %d, want 0 after provider decode failure", *openRouterCalls) + } +} + +func setOpenRouterCatalogTrap(t *testing.T) *int { + t.Helper() + calls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + t.Setenv("POTACO_OPENROUTER_MODELS_URL", srv.URL+"/models?output_modalities=image") + return &calls +} + +func openRouterCatalogModel(id string) map[string]any { + return map[string]any{ + "id": id, + "architecture": map[string]any{ + "output_modalities": []string{"image"}, + }, + } +} diff --git a/internal/adapter/openrouter.go b/internal/adapter/openrouter.go new file mode 100644 index 0000000..31aad0d --- /dev/null +++ b/internal/adapter/openrouter.go @@ -0,0 +1,177 @@ +package adapter + +import ( + "context" + "encoding/json" + "flag" + "fmt" + "net/http" + "os" + "strings" +) + +const defaultOpenRouterModelsURL = "https://openrouter.ai/api/v1/models?output_modalities=image" + +type OpenRouterModel struct { + ID string + DisplayName string +} + +func OpenRouterImageModelIDs(ctx context.Context, client *http.Client) (map[string]bool, error) { + names, err := OpenRouterImageModelNames(ctx, client) + if err != nil { + return nil, err + } + models := make(map[string]bool, len(names)) + for id := range names { + models[id] = true + } + return models, nil +} + +func OpenRouterImageModelNames(ctx context.Context, client *http.Client) (map[string]string, error) { + catalog, err := OpenRouterImageModels(ctx, client) + if err != nil { + return nil, err + } + models := make(map[string]string, len(catalog)*2) + for _, m := range catalog { + displayName := m.DisplayName + if displayName == "" { + displayName = m.ID + } + models[m.ID] = displayName + models[SuffixModelID(m.ID)] = displayName + } + if len(models) == 0 { + return nil, ErrDiscoveryFailed + } + return models, nil +} + +func OpenRouterImageModels(ctx context.Context, client *http.Client) ([]OpenRouterModel, error) { + url := defaultOpenRouterModelsURL + if override := testOpenRouterModelsURL(); override != "" { + url = override + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, fmt.Errorf("create OpenRouter models request: %w", err) + } + + resp, err := client.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("discover OpenRouter image models: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode >= 400 { + return nil, fmt.Errorf("discover OpenRouter image models failed (HTTP %d)", resp.StatusCode) + } + + var result struct { + Data []struct { + ID string `json:"id"` + Name string `json:"name"` + Architecture struct { + OutputModalities []string `json:"output_modalities"` + } `json:"architecture"` + } `json:"data"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("decode OpenRouter models response: %w", err) + } + + models := make([]OpenRouterModel, 0, len(result.Data)) + for _, m := range result.Data { + if m.ID == "" || !openRouterModelOutputsImage(m.Architecture.OutputModalities) { + continue + } + models = append(models, OpenRouterModel{ID: m.ID, DisplayName: m.Name}) + } + if len(models) == 0 { + return nil, ErrDiscoveryFailed + } + return models, nil +} + +func BestOpenRouterImageModelMatch(modelID string, models []OpenRouterModel) (OpenRouterModel, bool) { + query := strings.ToLower(modelID) + querySuffix := strings.ToLower(SuffixModelID(modelID)) + best := OpenRouterModel{} + bestScore := 0 + for _, model := range models { + id := strings.ToLower(model.ID) + suffix := strings.ToLower(SuffixModelID(model.ID)) + score := modelMatchScore(query, querySuffix, id, suffix) + if score > bestScore { + best = model + bestScore = score + } + } + return best, bestScore >= 60 +} + +func modelMatchScore(query, querySuffix, id, suffix string) int { + switch { + case query == id: + return 100 + case query == suffix || querySuffix == suffix: + return 95 + case strings.Contains(id, query) || strings.Contains(id, querySuffix): + return 80 + case editDistance(querySuffix, suffix) <= 2: + return 70 + default: + return 0 + } +} + +func editDistance(a, b string) int { + if a == "" { + return len(b) + } + if b == "" { + return len(a) + } + prev := make([]int, len(b)+1) + for j := range prev { + prev[j] = j + } + for i := 1; i <= len(a); i++ { + curr := make([]int, len(b)+1) + curr[0] = i + for j := 1; j <= len(b); j++ { + cost := 0 + if a[i-1] != b[j-1] { + cost = 1 + } + curr[j] = min(prev[j]+1, curr[j-1]+1, prev[j-1]+cost) + } + prev = curr + } + return prev[len(b)] +} + +func openRouterModelOutputsImage(modalities []string) bool { + for _, modality := range modalities { + if modality == "image" { + return true + } + } + return false +} + +func SuffixModelID(modelID string) string { + if idx := strings.LastIndex(modelID, "/"); idx >= 0 { + return modelID[idx+1:] + } + return modelID +} + +func testOpenRouterModelsURL() string { + if flag.Lookup("test.v") == nil { + return "" + } + return os.Getenv("POTACO_OPENROUTER_MODELS_URL") +} diff --git a/internal/adapter/openrouter_test.go b/internal/adapter/openrouter_test.go new file mode 100644 index 0000000..4f57c10 --- /dev/null +++ b/internal/adapter/openrouter_test.go @@ -0,0 +1,26 @@ +package adapter + +import "testing" + +func TestBestOpenRouterImageModelMatch_acceptsSuffixOrCloseMatch(t *testing.T) { + models := []OpenRouterModel{ + {ID: "openai/gpt-image-2", DisplayName: "GPT Image 2"}, + {ID: "openrouter/image-model", DisplayName: "Image Model"}, + } + + got, ok := BestOpenRouterImageModelMatch("image-modl", models) + if !ok { + t.Fatal("expected best match") + } + if got.ID != "openrouter/image-model" { + t.Fatalf("match = %q, want openrouter/image-model", got.ID) + } +} + +func TestBestOpenRouterImageModelMatch_rejectsDistantMatch(t *testing.T) { + models := []OpenRouterModel{{ID: "openai/gpt-image-2"}} + + if got, ok := BestOpenRouterImageModelMatch("totally-different", models); ok { + t.Fatalf("match = %+v, want no match", got) + } +} diff --git a/internal/adapter/replicate/discover.go b/internal/adapter/replicate/discover.go new file mode 100644 index 0000000..09b4c2e --- /dev/null +++ b/internal/adapter/replicate/discover.go @@ -0,0 +1,39 @@ +package replicate + +import ( + "context" + "fmt" + "net/http" + + "github.com/ncxton/potaco/internal/adapter" +) + +func (a *Adapter) DiscoverModels(ctx context.Context) (models []adapter.Model, err error) { + return nil, adapter.ErrDiscoveryFailed +} + +func (a *Adapter) Verify(ctx context.Context) (err error) { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, a.modelsURL(), nil) + if err != nil { + return fmt.Errorf("create request: %w", err) + } + httpReq.Header.Set("Authorization", a.AuthHeader(a.apiKey)) + + resp, err := a.http.Do(httpReq) + if err != nil { + return fmt.Errorf("verify: %w", err) + } + defer func() { + if closeErr := resp.Body.Close(); closeErr != nil && err == nil { + err = fmt.Errorf("close response body: %w", closeErr) + } + }() + + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) + } + if resp.StatusCode >= http.StatusBadRequest { + return fmt.Errorf("verification failed (HTTP %d)", resp.StatusCode) + } + return nil +} diff --git a/internal/adapter/vercel/edit.go b/internal/adapter/replicate/edit.go similarity index 65% rename from internal/adapter/vercel/edit.go rename to internal/adapter/replicate/edit.go index b7b159d..1b99b31 100644 --- a/internal/adapter/vercel/edit.go +++ b/internal/adapter/replicate/edit.go @@ -1,4 +1,4 @@ -package vercel +package replicate import ( "context" @@ -6,7 +6,7 @@ import ( "github.com/ncxton/potaco/internal/adapter" ) -// Edit returns ErrEditNotSupported because Vercel AI Gateway does not support image editing. +// Edit returns ErrEditNotSupported because Replicate edit inputs are model-specific. func (a *Adapter) Edit(context.Context, adapter.EditRequest) (*adapter.GenerateResponse, error) { return nil, adapter.ErrEditNotSupported } diff --git a/internal/adapter/replicate/generate.go b/internal/adapter/replicate/generate.go new file mode 100644 index 0000000..a46bef0 --- /dev/null +++ b/internal/adapter/replicate/generate.go @@ -0,0 +1,135 @@ +package replicate + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + + "github.com/ncxton/potaco/internal/adapter" +) + +func (a *Adapter) Generate(ctx context.Context, req adapter.GenerateRequest) (result *adapter.GenerateResponse, err error) { + if err := validateModelID(req.Model); err != nil { + return nil, err + } + body := map[string]any{ + "input": a.buildInput(req), + } + bodyJSON, err := json.Marshal(body) + if err != nil { + return nil, fmt.Errorf("marshal request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.generateURL(req.Model), bytes.NewReader(bodyJSON)) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Authorization", a.AuthHeader(a.apiKey)) + httpReq.Header.Set("Prefer", "wait") + + resp, err := a.doWithRetry(ctx, httpReq) + if err != nil { + return nil, err + } + defer func() { + if closeErr := resp.Body.Close(); closeErr != nil && err == nil { + err = fmt.Errorf("close response body: %w", closeErr) + } + }() + + pred, err := parsePrediction(resp) + if err != nil { + return nil, err + } + if isIncompletePrediction(pred) { + return a.pollPrediction(ctx, pred.URLs.Get) + } + return predictionToGenerateResponse(pred) +} + +func validateModelID(modelID string) error { + owner, name, ok := strings.Cut(modelID, "/") + if !ok || owner == "" || name == "" || strings.Contains(name, "/") { + return fmt.Errorf("replicate model must use owner/model format, got %q", modelID) + } + return nil +} + +func (a *Adapter) buildInput(req adapter.GenerateRequest) map[string]any { + input := map[string]any{ + "prompt": req.Prompt, + } + if req.N > 0 { + input["num_outputs"] = req.N + } + if req.Size != "" { + if aspectRatio := sizeToAspectRatio(req.Size); aspectRatio != "" { + input["aspect_ratio"] = aspectRatio + } else { + input["image_size"] = req.Size + } + } + if req.Seed != 0 { + input["seed"] = req.Seed + } + if req.GuidanceScale != 0 { + input["guidance_scale"] = req.GuidanceScale + } + if req.NegativePrompt != "" { + input["negative_prompt"] = req.NegativePrompt + } + for key, value := range req.ExtraParams { + input[key] = value + } + return input +} + +func (a *Adapter) pollPrediction(ctx context.Context, getURL string) (*adapter.GenerateResponse, error) { + for attempt := 0; attempt <= a.retries; attempt++ { + if attempt > 0 { + if err := a.sleepBeforeRetry(ctx, a.backoffOrDefault(attempt-1)); err != nil { + return nil, err + } + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, getURL, nil) + if err != nil { + return nil, fmt.Errorf("create prediction poll request: %w", err) + } + httpReq.Header.Set("Authorization", a.AuthHeader(a.apiKey)) + + resp, err := a.http.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("poll prediction: %w", err) + } + pred, parseErr := parsePrediction(resp) + closeErr := resp.Body.Close() + if closeErr != nil { + return nil, fmt.Errorf("close response body: %w", closeErr) + } + if parseErr != nil { + return nil, parseErr + } + if isIncompletePrediction(pred) { + continue + } + return predictionToGenerateResponse(pred) + } + return nil, fmt.Errorf("prediction did not complete after %d polls", a.retries+1) +} + +func sizeToAspectRatio(size string) string { + switch size { + case "1024x1024": + return "1:1" + case "1024x1792": + return "9:16" + case "1792x1024": + return "16:9" + default: + return "" + } +} diff --git a/internal/adapter/vercel/vercel.go b/internal/adapter/replicate/replicate.go similarity index 69% rename from internal/adapter/vercel/vercel.go rename to internal/adapter/replicate/replicate.go index 45f02c3..279b741 100644 --- a/internal/adapter/vercel/vercel.go +++ b/internal/adapter/replicate/replicate.go @@ -1,18 +1,18 @@ -// Package vercel implements the adapter.Adapter interface for the -// Vercel AI Gateway. -package vercel +// Package replicate implements the adapter.Adapter interface for Replicate. +package replicate import ( "context" "net/http" + "strings" "time" "github.com/ncxton/potaco/internal/adapter" ) -const defaultBaseURL = "https://ai-gateway.vercel.sh/v1" +const defaultBaseURL = "https://api.replicate.com/v1" -// Adapter implements adapter.Adapter for the Vercel AI Gateway. +// Adapter implements adapter.Adapter for the Replicate API. type Adapter struct { apiKey string baseURL string @@ -23,11 +23,11 @@ type Adapter struct { sleep func(ctx context.Context, d time.Duration) } -// New creates a Vercel adapter with the given API key and options. +// New creates a Replicate adapter with the given API key and options. func New(apiKey string, opts adapter.AdapterOpts) adapter.Adapter { baseURL := defaultBaseURL if opts.BaseURL != "" { - baseURL = opts.BaseURL + baseURL = strings.TrimRight(opts.BaseURL, "/") } timeout := 120 * time.Second if opts.Timeout > 0 { @@ -59,7 +59,7 @@ func (a *Adapter) SetSleep(fn func(context.Context, time.Duration)) { } // Name returns the provider name. -func (a *Adapter) Name() string { return "vercel" } +func (a *Adapter) Name() string { return "replicate" } // SupportsGenerate reports whether this provider supports image generation. func (a *Adapter) SupportsGenerate() bool { return true } @@ -72,20 +72,16 @@ func (a *Adapter) AuthHeader(apiKey string) string { return "Bearer " + apiKey } -func (a *Adapter) generateURL() string { - return a.baseURL + "/images/generations" +func (a *Adapter) generateURL(modelID string) string { + return a.baseURL + "/models/" + modelID + "/predictions" } func (a *Adapter) modelsURL() string { return a.baseURL + "/models" } -func (a *Adapter) endpointsURL(modelID string) string { - return a.baseURL + "/models/" + modelID + "/endpoints" -} - func init() { - adapter.Register("vercel", func(apiKey string, opts adapter.AdapterOpts) (adapter.Adapter, error) { + adapter.Register("replicate", func(apiKey string, opts adapter.AdapterOpts) (adapter.Adapter, error) { return New(apiKey, opts), nil }) } diff --git a/internal/adapter/replicate/replicate_test.go b/internal/adapter/replicate/replicate_test.go new file mode 100644 index 0000000..8931706 --- /dev/null +++ b/internal/adapter/replicate/replicate_test.go @@ -0,0 +1,177 @@ +package replicate + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/ncxton/potaco/internal/adapter" +) + +func TestReplicateAdapterName(t *testing.T) { + ad := New("test-key", adapter.AdapterOpts{}) + if ad.Name() != "replicate" { + t.Fatalf("Name() = %q, want replicate", ad.Name()) + } +} + +func TestReplicateAuthHeader(t *testing.T) { + ad := New("test-key", adapter.AdapterOpts{}) + if got := ad.AuthHeader("test-key"); got != "Bearer test-key" { + t.Fatalf("AuthHeader() = %q, want Bearer test-key", got) + } +} + +func TestReplicateRegisteredInRegistry(t *testing.T) { + _, err := adapter.Get("replicate", "test-key", adapter.AdapterOpts{}) + if err != nil { + t.Fatalf("adapter.Get(\"replicate\") failed: %v", err) + } +} + +func TestReplicateGenerate(t *testing.T) { + var gotPath string + var gotBody struct { + Input map[string]any `json:"input"` + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + if r.Method != http.MethodPost { + t.Fatalf("method = %s, want POST", r.Method) + } + if auth := r.Header.Get("Authorization"); auth != "Bearer test-key" { + t.Fatalf("auth = %q, want Bearer test-key", auth) + } + if prefer := r.Header.Get("Prefer"); prefer != "wait" { + t.Fatalf("Prefer = %q, want wait", prefer) + } + if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil { + t.Fatalf("decode request body: %v", err) + } + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(map[string]any{ + "output": []string{"https://replicate.delivery/result.png"}, + }); err != nil { + t.Fatalf("write response: %v", err) + } + })) + defer srv.Close() + + ad := New("test-key", adapter.AdapterOpts{BaseURL: srv.URL + "/v1"}) + resp, err := ad.Generate(context.Background(), adapter.GenerateRequest{ + Prompt: "a cat", + Model: "black-forest-labs/flux-schnell", + N: 2, + Size: "1024x1024", + Seed: 42, + GuidanceScale: 3.5, + NegativePrompt: "low quality", + ExtraParams: map[string]any{ + "output_format": "png", + }, + }) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if gotPath != "/v1/models/black-forest-labs/flux-schnell/predictions" { + t.Fatalf("path = %q, want model predictions path", gotPath) + } + if gotBody.Input["prompt"] != "a cat" { + t.Fatalf("input prompt = %v, want a cat", gotBody.Input["prompt"]) + } + if gotBody.Input["num_outputs"] != float64(2) { + t.Fatalf("input num_outputs = %v, want 2", gotBody.Input["num_outputs"]) + } + if gotBody.Input["aspect_ratio"] != "1:1" { + t.Fatalf("input aspect_ratio = %v, want 1:1", gotBody.Input["aspect_ratio"]) + } + if gotBody.Input["output_format"] != "png" { + t.Fatalf("input output_format = %v, want png", gotBody.Input["output_format"]) + } + if len(resp.Data) != 1 || resp.Data[0].URL != "https://replicate.delivery/result.png" { + t.Fatalf("response data = %#v, want Replicate output URL", resp.Data) + } +} + +func TestReplicateGeneratePollsIncompletePrediction(t *testing.T) { + getCalls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/v1/models/black-forest-labs/flux-schnell/predictions": + if err := json.NewEncoder(w).Encode(map[string]any{ + "id": "pred-1", + "status": "processing", + "urls": map[string]any{ + "get": "http://" + r.Host + "/v1/predictions/pred-1", + }, + }); err != nil { + t.Fatalf("write create response: %v", err) + } + case "/v1/predictions/pred-1": + getCalls++ + if err := json.NewEncoder(w).Encode(map[string]any{ + "id": "pred-1", + "status": "succeeded", + "output": "https://replicate.delivery/polled.png", + }); err != nil { + t.Fatalf("write get response: %v", err) + } + default: + t.Fatalf("unexpected path %q", r.URL.Path) + } + })) + defer srv.Close() + + ad := New("test-key", adapter.AdapterOpts{BaseURL: srv.URL + "/v1"}).(*Adapter) + ad.SetBackoff(func(int) time.Duration { return time.Millisecond }) + ad.SetSleep(func(context.Context, time.Duration) {}) + + resp, err := ad.Generate(context.Background(), adapter.GenerateRequest{ + Prompt: "a cat", + Model: "black-forest-labs/flux-schnell", + }) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if getCalls != 1 { + t.Fatalf("poll calls = %d, want 1", getCalls) + } + if len(resp.Data) != 1 || resp.Data[0].URL != "https://replicate.delivery/polled.png" { + t.Fatalf("response data = %#v, want polled output", resp.Data) + } +} + +func TestReplicateGenerateRequiresOwnerModel(t *testing.T) { + ad := New("test-key", adapter.AdapterOpts{}) + _, err := ad.Generate(context.Background(), adapter.GenerateRequest{ + Prompt: "a cat", + Model: "flux-schnell", + }) + if err == nil { + t.Fatal("expected owner/model error") + } + if !strings.Contains(err.Error(), "owner/model") { + t.Fatalf("error = %v, want owner/model", err) + } +} + +func TestReplicateDiscoverModelsUnsupported(t *testing.T) { + ad := New("test-key", adapter.AdapterOpts{}) + _, err := ad.DiscoverModels(context.Background()) + if err != adapter.ErrDiscoveryFailed { + t.Fatalf("DiscoverModels error = %v, want ErrDiscoveryFailed", err) + } +} + +func TestReplicateEditNotSupported(t *testing.T) { + ad := New("test-key", adapter.AdapterOpts{}) + _, err := ad.Edit(context.Background(), adapter.EditRequest{Prompt: "test"}) + if err != adapter.ErrEditNotSupported { + t.Fatalf("Edit error = %v, want ErrEditNotSupported", err) + } +} diff --git a/internal/adapter/replicate/response.go b/internal/adapter/replicate/response.go new file mode 100644 index 0000000..f76aa7c --- /dev/null +++ b/internal/adapter/replicate/response.go @@ -0,0 +1,112 @@ +package replicate + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + + "github.com/ncxton/potaco/internal/adapter" +) + +var maxResponseBytes int64 = 128 << 20 + +type predictionResponse struct { + Status string `json:"status"` + Error any `json:"error"` + Output any `json:"output"` + URLs struct { + Get string `json:"get"` + } `json:"urls"` +} + +func parsePrediction(resp *http.Response) (*predictionResponse, error) { + respBody, err := readLimitedBody(resp.Body, maxResponseBytes, "provider response") + if err != nil { + return nil, err + } + + if resp.StatusCode >= http.StatusBadRequest { + message := parseErrorMessage(respBody) + if message != "" { + return nil, fmt.Errorf("API error (HTTP %d): %s", resp.StatusCode, message) + } + return nil, fmt.Errorf("API error (HTTP %d): %s", resp.StatusCode, string(respBody)) + } + + var pred predictionResponse + if err := json.Unmarshal(respBody, &pred); err != nil { + return nil, fmt.Errorf("parse response: %w", err) + } + return &pred, nil +} + +func predictionToGenerateResponse(pred *predictionResponse) (*adapter.GenerateResponse, error) { + if pred.Error != nil { + return nil, fmt.Errorf("prediction failed: %v", pred.Error) + } + switch pred.Status { + case "failed", "canceled": + return nil, fmt.Errorf("prediction %s", pred.Status) + } + + var urls []string + collectOutputStrings(pred.Output, &urls) + imgResp := &adapter.GenerateResponse{} + for _, url := range urls { + imgResp.Data = append(imgResp.Data, adapter.ImageData{URL: url}) + } + if len(imgResp.Data) == 0 { + return nil, fmt.Errorf("prediction returned no image output") + } + return imgResp, nil +} + +func isIncompletePrediction(pred *predictionResponse) bool { + switch pred.Status { + case "starting", "processing": + return pred.URLs.Get != "" + default: + return false + } +} + +func parseErrorMessage(respBody []byte) string { + var errResp struct { + Detail string `json:"detail"` + Error any `json:"error"` + } + if err := json.Unmarshal(respBody, &errResp); err != nil { + return "" + } + if errResp.Detail != "" { + return errResp.Detail + } + if errResp.Error != nil { + return fmt.Sprint(errResp.Error) + } + return "" +} + +func collectOutputStrings(value any, out *[]string) { + switch v := value.(type) { + case string: + *out = append(*out, v) + case []any: + for _, item := range v { + collectOutputStrings(item, out) + } + } +} + +func readLimitedBody(r io.Reader, limit int64, label string) ([]byte, error) { + limited := io.LimitReader(r, limit+1) + data, err := io.ReadAll(limited) + if err != nil { + return nil, fmt.Errorf("read %s: %w", label, err) + } + if int64(len(data)) > limit { + return nil, fmt.Errorf("%s too large: limit is %d bytes", label, limit) + } + return data, nil +} diff --git a/internal/adapter/vercel/retry.go b/internal/adapter/replicate/retry.go similarity index 99% rename from internal/adapter/vercel/retry.go rename to internal/adapter/replicate/retry.go index 6c4ea6c..083e58d 100644 --- a/internal/adapter/vercel/retry.go +++ b/internal/adapter/replicate/retry.go @@ -1,4 +1,4 @@ -package vercel +package replicate import ( "context" diff --git a/internal/adapter/vercel/discover.go b/internal/adapter/vercel/discover.go deleted file mode 100644 index 549f1e9..0000000 --- a/internal/adapter/vercel/discover.go +++ /dev/null @@ -1,114 +0,0 @@ -package vercel - -import ( - "context" - "encoding/json" - "fmt" - "net/http" - - "github.com/ncxton/potaco/internal/adapter" -) - -func (a *Adapter) DiscoverModels(ctx context.Context) (models []adapter.Model, err error) { - httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, a.modelsURL(), nil) - if err != nil { - return nil, fmt.Errorf("create request: %w", err) - } - httpReq.Header.Set("Authorization", a.AuthHeader(a.apiKey)) - - resp, err := a.http.Do(httpReq) - if err != nil { - return nil, fmt.Errorf("discover models: %w", err) - } - defer func() { - if closeErr := resp.Body.Close(); closeErr != nil && err == nil { - err = fmt.Errorf("close response body: %w", closeErr) - } - }() - - if resp.StatusCode >= http.StatusBadRequest { - return nil, fmt.Errorf("discover models failed (HTTP %d)", resp.StatusCode) - } - - var result struct { - Data []struct { - ID string `json:"id"` - Type string `json:"type"` - } `json:"data"` - } - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { - return nil, fmt.Errorf("decode models response: %w", err) - } - - for _, model := range result.Data { - if model.Type != "image" { - continue - } - models = append(models, adapter.Model{ - ID: model.ID, - DisplayName: stripProviderPrefix(model.ID), - SupportsGen: true, - SupportsEdit: false, - }) - } - if len(models) == 0 { - return nil, adapter.ErrDiscoveryFailed - } - return models, nil -} - -func (a *Adapter) Verify(ctx context.Context) error { - if err := a.verifyReachable(ctx); err != nil { - return err - } - if err := a.verifyKey(ctx); err != nil { - return err - } - return nil -} - -func (a *Adapter) verifyReachable(ctx context.Context) (err error) { - httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, a.modelsURL(), nil) - if err != nil { - return fmt.Errorf("create request: %w", err) - } - httpReq.Header.Set("Authorization", a.AuthHeader(a.apiKey)) - - resp, err := a.http.Do(httpReq) - if err != nil { - return fmt.Errorf("verify reachability: %w", err) - } - defer func() { - if closeErr := resp.Body.Close(); closeErr != nil && err == nil { - err = fmt.Errorf("close response body: %w", closeErr) - } - }() - - if resp.StatusCode >= http.StatusInternalServerError { - return fmt.Errorf("verification failed (HTTP %d)", resp.StatusCode) - } - return nil -} - -func (a *Adapter) verifyKey(ctx context.Context) (err error) { - httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, a.endpointsURL("openai/gpt-image-2"), nil) - if err != nil { - return fmt.Errorf("create key validation request: %w", err) - } - httpReq.Header.Set("Authorization", a.AuthHeader(a.apiKey)) - - resp, err := a.http.Do(httpReq) - if err != nil { - return fmt.Errorf("verify key: %w", err) - } - defer func() { - if closeErr := resp.Body.Close(); closeErr != nil && err == nil { - err = fmt.Errorf("close response body: %w", closeErr) - } - }() - - if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { - return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) - } - return nil -} diff --git a/internal/adapter/vercel/generate.go b/internal/adapter/vercel/generate.go deleted file mode 100644 index ebe69ff..0000000 --- a/internal/adapter/vercel/generate.go +++ /dev/null @@ -1,70 +0,0 @@ -package vercel - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "net/http" - - "github.com/ncxton/potaco/internal/adapter" -) - -// Generate calls POST /v1/images/generations with an OpenAI-compatible body. -func (a *Adapter) Generate(ctx context.Context, req adapter.GenerateRequest) (result *adapter.GenerateResponse, err error) { - body := a.buildGenerateBody(req) - bodyJSON, err := json.Marshal(body) - if err != nil { - return nil, fmt.Errorf("marshal request: %w", err) - } - - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.generateURL(), bytes.NewReader(bodyJSON)) - if err != nil { - return nil, fmt.Errorf("create request: %w", err) - } - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Authorization", a.AuthHeader(a.apiKey)) - - resp, err := a.doWithRetry(ctx, httpReq) - if err != nil { - return nil, err - } - defer func() { - if closeErr := resp.Body.Close(); closeErr != nil && err == nil { - err = fmt.Errorf("close response body: %w", closeErr) - } - }() - - return parseResponse(resp) -} - -func (a *Adapter) buildGenerateBody(req adapter.GenerateRequest) map[string]any { - body := map[string]any{ - "prompt": req.Prompt, - } - if req.Model != "" { - body["model"] = req.Model - } - if req.N > 0 { - body["n"] = req.N - } - if req.Size != "" { - body["size"] = req.Size - } - if req.Quality != "" { - body["quality"] = req.Quality - } - if req.Style != "" { - body["style"] = req.Style - } - if req.ResponseFormat != "" { - body["response_format"] = req.ResponseFormat - } - if req.Seed != 0 { - body["seed"] = req.Seed - } - for key, value := range req.ExtraParams { - body[key] = value - } - return body -} diff --git a/internal/adapter/vercel/generate_test.go b/internal/adapter/vercel/generate_test.go deleted file mode 100644 index 25c0e66..0000000 --- a/internal/adapter/vercel/generate_test.go +++ /dev/null @@ -1,160 +0,0 @@ -package vercel - -import ( - "context" - "encoding/json" - "io" - "net/http" - "net/http/httptest" - "strings" - "testing" - "time" - - "github.com/ncxton/potaco/internal/adapter" -) - -func TestVercelGenerate(t *testing.T) { - var gotBody map[string]any - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method != "POST" { - t.Errorf("method = %s, want POST", r.Method) - } - if r.URL.Path != "/v1/images/generations" { - t.Errorf("path = %q, want /v1/images/generations", r.URL.Path) - } - if auth := r.Header.Get("Authorization"); auth != "Bearer test-key" { - t.Errorf("auth = %q, want 'Bearer test-key'", auth) - } - body, err := io.ReadAll(r.Body) - if err != nil { - t.Errorf("read request body: %v", err) - return - } - if err := json.Unmarshal(body, &gotBody); err != nil { - t.Errorf("decode request body: %v", err) - return - } - - w.Header().Set("Content-Type", "application/json") - if err := json.NewEncoder(w).Encode(map[string]any{ - "created": 1700000000, - "data": []map[string]any{ - {"b64_json": "iVBORw0KGgo="}, - }, - }); err != nil { - t.Errorf("write response: %v", err) - } - })) - defer srv.Close() - - ad := New("test-key", adapter.AdapterOpts{BaseURL: srv.URL + "/v1"}) - va := ad.(*Adapter) - va.SetBackoff(func(int) time.Duration { return 1 * time.Millisecond }) - va.SetSleep(func(context.Context, time.Duration) {}) - - resp, err := ad.Generate(context.Background(), adapter.GenerateRequest{ - Prompt: "a cat", - Model: "openai/gpt-image-2", - N: 1, - Size: "1024x1024", - }) - if err != nil { - t.Fatalf("Generate: %v", err) - } - if len(resp.Data) != 1 { - t.Fatalf("Data len = %d, want 1", len(resp.Data)) - } - if resp.Data[0].B64JSON != "iVBORw0KGgo=" { - t.Errorf("b64_json = %q, want 'iVBORw0KGgo='", resp.Data[0].B64JSON) - } - if gotBody["model"] != "openai/gpt-image-2" { - t.Errorf("body model = %v, want 'openai/gpt-image-2'", gotBody["model"]) - } -} - -func TestVercelGenerateWithProviderOptions(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - body, err := io.ReadAll(r.Body) - if err != nil { - t.Errorf("read request body: %v", err) - return - } - var gotBody map[string]any - if err := json.Unmarshal(body, &gotBody); err != nil { - t.Errorf("decode request body: %v", err) - return - } - - po, ok := gotBody["providerOptions"] - if !ok { - t.Fatal("providerOptions not found in body") - } - poMap, ok := po.(map[string]any) - if !ok { - t.Fatalf("providerOptions is %T, want map", po) - } - if _, ok := poMap["blackForestLabs"]; !ok { - t.Error("providerOptions.blackForestLabs not found") - } - - w.Header().Set("Content-Type", "application/json") - if err := json.NewEncoder(w).Encode(map[string]any{ - "data": []map[string]any{ - {"url": "https://example.com/result.png"}, - }, - }); err != nil { - t.Errorf("write response: %v", err) - } - })) - defer srv.Close() - - ad := New("test-key", adapter.AdapterOpts{BaseURL: srv.URL + "/v1"}) - va := ad.(*Adapter) - va.SetBackoff(func(int) time.Duration { return 1 * time.Millisecond }) - va.SetSleep(func(context.Context, time.Duration) {}) - - _, err := ad.Generate(context.Background(), adapter.GenerateRequest{ - Prompt: "a cat", - Model: "bfl/flux-2-pro", - ExtraParams: map[string]any{ - "providerOptions": map[string]any{ - "blackForestLabs": map[string]any{ - "outputFormat": "png", - }, - }, - }, - }) - if err != nil { - t.Fatalf("Generate: %v", err) - } -} - -func TestVercelGenerateAPIError(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusBadRequest) - if err := json.NewEncoder(w).Encode(map[string]any{ - "error": map[string]any{ - "message": "invalid model id", - }, - }); err != nil { - t.Errorf("write response: %v", err) - } - })) - defer srv.Close() - - ad := New("test-key", adapter.AdapterOpts{BaseURL: srv.URL + "/v1"}) - va := ad.(*Adapter) - va.SetBackoff(func(int) time.Duration { return 1 * time.Millisecond }) - va.SetSleep(func(context.Context, time.Duration) {}) - - _, err := ad.Generate(context.Background(), adapter.GenerateRequest{ - Prompt: "test", - Model: "invalid/model", - }) - if err == nil { - t.Fatal("expected error, got nil") - } - if !strings.Contains(err.Error(), "invalid model id") { - t.Errorf("error should contain 'invalid model id', got: %v", err) - } -} diff --git a/internal/adapter/vercel/models.go b/internal/adapter/vercel/models.go deleted file mode 100644 index 42652ec..0000000 --- a/internal/adapter/vercel/models.go +++ /dev/null @@ -1,10 +0,0 @@ -package vercel - -import "strings" - -func stripProviderPrefix(modelID string) string { - if idx := strings.Index(modelID, "/"); idx >= 0 { - return modelID[idx+1:] - } - return modelID -} diff --git a/internal/adapter/vercel/response.go b/internal/adapter/vercel/response.go deleted file mode 100644 index 1dce9e9..0000000 --- a/internal/adapter/vercel/response.go +++ /dev/null @@ -1,49 +0,0 @@ -package vercel - -import ( - "encoding/json" - "fmt" - "io" - "net/http" - - "github.com/ncxton/potaco/internal/adapter" -) - -var maxResponseBytes int64 = 128 << 20 - -func parseResponse(resp *http.Response) (*adapter.GenerateResponse, error) { - respBody, err := readLimitedBody(resp.Body, maxResponseBytes, "provider response") - if err != nil { - return nil, err - } - - if resp.StatusCode >= http.StatusBadRequest { - var errResp struct { - Error struct { - Message string `json:"message"` - } `json:"error"` - } - if jsonErr := json.Unmarshal(respBody, &errResp); jsonErr == nil && errResp.Error.Message != "" { - return nil, fmt.Errorf("API error (HTTP %d): %s", resp.StatusCode, errResp.Error.Message) - } - return nil, fmt.Errorf("API error (HTTP %d): %s", resp.StatusCode, string(respBody)) - } - - var imgResp adapter.GenerateResponse - if err := json.Unmarshal(respBody, &imgResp); err != nil { - return nil, fmt.Errorf("parse response: %w", err) - } - return &imgResp, nil -} - -func readLimitedBody(r io.Reader, limit int64, label string) ([]byte, error) { - limited := io.LimitReader(r, limit+1) - data, err := io.ReadAll(limited) - if err != nil { - return nil, fmt.Errorf("read %s: %w", label, err) - } - if int64(len(data)) > limit { - return nil, fmt.Errorf("%s too large: limit is %d bytes", label, limit) - } - return data, nil -} diff --git a/internal/adapter/vercel/retry_test.go b/internal/adapter/vercel/retry_test.go deleted file mode 100644 index 6056192..0000000 --- a/internal/adapter/vercel/retry_test.go +++ /dev/null @@ -1,125 +0,0 @@ -package vercel - -import ( - "context" - "net/http" - "net/http/httptest" - "sync/atomic" - "testing" - "time" - - "github.com/ncxton/potaco/internal/adapter" -) - -// TestVercelRetryOn5xxThenSuccess verifies that a 500 response is retried -// and the subsequent successful response is returned. -func TestVercelRetryOn5xxThenSuccess(t *testing.T) { - var attempts int32 - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - n := atomic.AddInt32(&attempts, 1) - if n == 1 { - w.WriteHeader(http.StatusInternalServerError) - return - } - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"created":123,"data":[{"b64_json":"aGVsbG8="}]}`)) - })) - defer server.Close() - - ad := New("test-key", adapter.AdapterOpts{BaseURL: server.URL + "/v1"}) - va := ad.(*Adapter) - va.SetBackoff(func(int) time.Duration { return 1 * time.Millisecond }) - va.SetSleep(func(context.Context, time.Duration) {}) - - resp, err := ad.Generate(context.Background(), adapter.GenerateRequest{ - Prompt: "test", - Model: "openai/gpt-image-2", - }) - if err != nil { - t.Fatalf("Generate error: %v", err) - } - if len(resp.Data) != 1 { - t.Fatalf("Data len = %d, want 1", len(resp.Data)) - } - if resp.Data[0].B64JSON != "aGVsbG8=" { - t.Errorf("B64JSON = %q, want 'aGVsbG8='", resp.Data[0].B64JSON) - } - if attempts != 2 { - t.Errorf("attempts = %d, want 2", attempts) - } -} - -// TestVercelRetryExhaustion verifies that when the server always returns 429, -// the adapter gives up after maxRetries attempts and returns the last -// error response. -func TestVercelRetryExhaustion(t *testing.T) { - var attempts int32 - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - atomic.AddInt32(&attempts, 1) - w.WriteHeader(http.StatusTooManyRequests) - _, _ = w.Write([]byte(`{"error":{"message":"rate limited"}}`)) - })) - defer server.Close() - - ad := New("test-key", adapter.AdapterOpts{BaseURL: server.URL + "/v1", Retries: 2}) - va := ad.(*Adapter) - va.SetBackoff(func(int) time.Duration { return 1 * time.Millisecond }) - va.SetSleep(func(context.Context, time.Duration) {}) - - _, err := ad.Generate(context.Background(), adapter.GenerateRequest{ - Prompt: "test", - Model: "openai/gpt-image-2", - }) - if err == nil { - t.Fatal("Generate should return error after retry exhaustion") - } - // With retries=2: 1 initial attempt + 2 retries = 3 total attempts - if attempts != 3 { - t.Errorf("attempts = %d, want 3", attempts) - } -} - -// TestVercelRetryAfterHeaderHonored verifies that when the server sends a -// Retry-After header, the delay from that header is used instead of the -// default backoff. We capture the actual sleep duration via SetSleep. -func TestVercelRetryAfterHeaderHonored(t *testing.T) { - var attempts int32 - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - n := atomic.AddInt32(&attempts, 1) - if n == 1 { - w.Header().Set("Retry-After", "1") - w.WriteHeader(http.StatusTooManyRequests) - return - } - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"created":123,"data":[{"b64_json":"aGVsbG8="}]}`)) - })) - defer server.Close() - - ad := New("test-key", adapter.AdapterOpts{BaseURL: server.URL + "/v1"}) - va := ad.(*Adapter) - - var capturedDelay time.Duration - // Backoff returns a huge value that should NOT be used when Retry-After is present. - va.SetBackoff(func(int) time.Duration { return 100 * time.Hour }) - va.SetSleep(func(ctx context.Context, d time.Duration) { - capturedDelay = d - }) - - resp, err := ad.Generate(context.Background(), adapter.GenerateRequest{ - Prompt: "test", - Model: "openai/gpt-image-2", - }) - if err != nil { - t.Fatalf("Generate error: %v", err) - } - if len(resp.Data) != 1 { - t.Fatalf("Data len = %d, want 1", len(resp.Data)) - } - if attempts != 2 { - t.Errorf("attempts = %d, want 2", attempts) - } - if capturedDelay != 1*time.Second { - t.Errorf("sleep duration = %v, want 1s (from Retry-After header)", capturedDelay) - } -} diff --git a/internal/adapter/vercel/vercel_test.go b/internal/adapter/vercel/vercel_test.go deleted file mode 100644 index 9ac3e5d..0000000 --- a/internal/adapter/vercel/vercel_test.go +++ /dev/null @@ -1,200 +0,0 @@ -package vercel - -import ( - "context" - "encoding/json" - "net/http" - "net/http/httptest" - "strings" - "testing" - - "github.com/ncxton/potaco/internal/adapter" -) - -func TestVercelAdapterName(t *testing.T) { - ad := New("test-key", adapter.AdapterOpts{}) - if ad.Name() != "vercel" { - t.Errorf("Name() = %q, want %q", ad.Name(), "vercel") - } -} - -func TestVercelAuthHeader(t *testing.T) { - ad := New("my-key", adapter.AdapterOpts{}) - got := ad.AuthHeader("my-key") - want := "Bearer my-key" - if got != want { - t.Errorf("AuthHeader() = %q, want %q", got, want) - } -} - -func TestVercelSupportsGenerate(t *testing.T) { - ad := New("key", adapter.AdapterOpts{}) - if !ad.SupportsGenerate() { - t.Error("SupportsGenerate should be true") - } -} - -func TestVercelSupportsEdit(t *testing.T) { - ad := New("key", adapter.AdapterOpts{}) - if ad.SupportsEdit() { - t.Error("SupportsEdit should be false") - } -} - -func TestVercelNewWithDefaults(t *testing.T) { - ad := New("key", adapter.AdapterOpts{}) - va, ok := ad.(*Adapter) - if !ok { - t.Fatalf("expected *vercel.Adapter, got %T", ad) - } - if va.baseURL != "https://ai-gateway.vercel.sh/v1" { - t.Errorf("baseURL = %q, want %q", va.baseURL, "https://ai-gateway.vercel.sh/v1") - } -} - -func TestVercelModelsURL(t *testing.T) { - ad := New("key", adapter.AdapterOpts{BaseURL: "https://example.com/v1"}) - va, ok := ad.(*Adapter) - if !ok { - t.Fatalf("expected *vercel.Adapter, got %T", ad) - } - - got := va.modelsURL() - want := "https://example.com/v1/models" - if got != want { - t.Errorf("modelsURL() = %q, want %q", got, want) - } -} - -func TestVercelEndpointsURL(t *testing.T) { - ad := New("key", adapter.AdapterOpts{BaseURL: "https://example.com/v1"}) - va, ok := ad.(*Adapter) - if !ok { - t.Fatalf("expected *vercel.Adapter, got %T", ad) - } - - got := va.endpointsURL("openai/gpt-image-2") - want := "https://example.com/v1/models/openai/gpt-image-2/endpoints" - if got != want { - t.Errorf("endpointsURL() = %q, want %q", got, want) - } -} - -func TestVercelRegisteredInRegistry(t *testing.T) { - _, err := adapter.Get("vercel", "key", adapter.AdapterOpts{}) - if err != nil { - t.Fatalf("adapter.Get(\"vercel\") failed: %v", err) - } -} - -func TestVercelEditNotSupported(t *testing.T) { - ad := New("key", adapter.AdapterOpts{}) - _, err := ad.Edit(context.Background(), adapter.EditRequest{ - Prompt: "test", - Model: "openai/gpt-image-2", - ImagePath: "/tmp/test.png", - }) - if err != adapter.ErrEditNotSupported { - t.Errorf("Edit error = %v, want ErrEditNotSupported", err) - } -} - -func TestVercelDiscoverModels(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/v1/models" { - t.Errorf("path = %q, want /v1/models", r.URL.Path) - } - w.Header().Set("Content-Type", "application/json") - if err := json.NewEncoder(w).Encode(map[string]any{ - "data": []map[string]any{ - {"id": "openai/gpt-image-2", "type": "image"}, - {"id": "openai/text-embedding-3", "type": "embedding"}, - {"id": "bfl/flux-2-pro", "type": "image"}, - {"id": "meta/llama-3", "type": "chat"}, - }, - }); err != nil { - t.Errorf("write response: %v", err) - } - })) - defer srv.Close() - - ad := New("test-key", adapter.AdapterOpts{BaseURL: srv.URL + "/v1"}) - models, err := ad.DiscoverModels(context.Background()) - if err != nil { - t.Fatalf("DiscoverModels: %v", err) - } - if len(models) != 2 { - t.Fatalf("models len = %d, want 2 (only image type)", len(models)) - } - if models[0].ID != "openai/gpt-image-2" { - t.Errorf("model[0] ID = %q, want 'openai/gpt-image-2'", models[0].ID) - } - if models[1].ID != "bfl/flux-2-pro" { - t.Errorf("model[1] ID = %q, want 'bfl/flux-2-pro'", models[1].ID) - } - if models[0].DisplayName != "gpt-image-2" { - t.Errorf("model[0] DisplayName = %q, want 'gpt-image-2' (prefix stripped)", models[0].DisplayName) - } -} - -func TestVercelDiscoverModelsError(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusInternalServerError) - })) - defer srv.Close() - - ad := New("test-key", adapter.AdapterOpts{BaseURL: srv.URL + "/v1"}) - _, err := ad.DiscoverModels(context.Background()) - if err == nil { - t.Fatal("DiscoverModels should return error on API failure") - } -} - -func TestVercelVerify(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // The verify flow makes two requests: GET /v1/models (reachability) - // and GET /v1/models/openai/gpt-image-2/endpoints (key validation). - if r.URL.Path == "/v1/models" { - w.WriteHeader(http.StatusOK) - return - } - if r.URL.Path == "/v1/models/openai/gpt-image-2/endpoints" { - if auth := r.Header.Get("Authorization"); auth != "Bearer test-key" { - t.Errorf("endpoints auth = %q, want 'Bearer test-key'", auth) - } - w.WriteHeader(http.StatusOK) - return - } - w.WriteHeader(http.StatusNotFound) - })) - defer srv.Close() - - ad := New("test-key", adapter.AdapterOpts{BaseURL: srv.URL + "/v1"}) - if err := ad.Verify(context.Background()); err != nil { - t.Errorf("Verify: %v", err) - } -} - -func TestVercelVerifyInvalidKey(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/v1/models" { - w.WriteHeader(http.StatusOK) - return - } - if r.URL.Path == "/v1/models/openai/gpt-image-2/endpoints" { - w.WriteHeader(http.StatusUnauthorized) - return - } - w.WriteHeader(http.StatusNotFound) - })) - defer srv.Close() - - ad := New("bad-key", adapter.AdapterOpts{BaseURL: srv.URL + "/v1"}) - err := ad.Verify(context.Background()) - if err == nil { - t.Fatal("expected error for invalid key") - } - if !strings.Contains(err.Error(), "invalid API key") { - t.Errorf("error should mention invalid key, got: %v", err) - } -} diff --git a/internal/auth/auth.go b/internal/auth/auth.go index 19b67fc..a85ee76 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -60,36 +60,11 @@ func (m *AuthManager) Add(provider, apiKey string) error { // AddProvider stores the API key for a provider key using the given provider // type for adapter resolution, creates a config entry, and sets it active. func (m *AuthManager) AddProvider(provider, providerType, apiKey string) error { - if err := m.store.Set(provider, apiKey); err != nil { - return fmt.Errorf("store credential: %w", err) - } - - cfg, err := m.LoadConfig() - if err != nil { - return fmt.Errorf("load config: %w", err) - } - - if cfg.Providers == nil { - cfg.Providers = make(map[string]config.ProviderConfig) - } - - pc := cfg.Providers[provider] - if providerType == "" { - providerType = config.ResolveProviderType(provider, pc) - } - pc.Type = providerType - if pc.Retries == 0 { - pc.Retries = 2 - } - if pc.Timeout == 0 { - pc.Timeout = 120 - } - cfg.Providers[provider] = pc - - cfg.ActiveProvider = provider - cfg.ActiveModel = cfg.Providers[provider].Model - - return m.saveConfig(cfg) + return m.AddProviderSetup(ProviderSetup{ + ProviderName: provider, + ProviderType: providerType, + APIKey: apiKey, + }) } // Remove deletes the provider's credential and config entry. @@ -223,6 +198,49 @@ func (m *AuthManager) SetModelEdit(provider, model string, edit bool) error { return m.saveConfig(cfg) } +func (m *AuthManager) SetModelMetadata(provider, model, displayName string, custom bool) error { + cfg, err := m.LoadConfig() + if err != nil { + return fmt.Errorf("load config: %w", err) + } + + pc, ok := cfg.Providers[provider] + if !ok { + return fmt.Errorf("provider %q is not configured. Use 'potaco auth add %s' first", provider, provider) + } + if pc.Models == nil { + pc.Models = make(map[string]config.ModelConfig) + } + mc := pc.Models[model] + mc.DisplayName = displayName + mc.Custom = custom + pc.Models[model] = mc + cfg.Providers[provider] = pc + + return m.saveConfig(cfg) +} + +func (m *AuthManager) SetModelParameters(provider, model string, params config.ModelParameters) error { + cfg, err := m.LoadConfig() + if err != nil { + return fmt.Errorf("load config: %w", err) + } + + pc, ok := cfg.Providers[provider] + if !ok { + return fmt.Errorf("provider %q is not configured. Use 'potaco auth add %s' first", provider, provider) + } + if pc.Models == nil { + pc.Models = make(map[string]config.ModelConfig) + } + mc := pc.Models[model] + mc.Parameters = params + pc.Models[model] = mc + cfg.Providers[provider] = pc + + return m.saveConfig(cfg) +} + // GetActiveAPIKey returns the API key for the active provider. func (m *AuthManager) GetActiveAPIKey() (string, error) { cfg, err := m.LoadConfig() diff --git a/internal/auth/auth_test.go b/internal/auth/auth_test.go index b310a71..e51a0f3 100644 --- a/internal/auth/auth_test.go +++ b/internal/auth/auth_test.go @@ -2,10 +2,12 @@ package auth import ( + "os" "path/filepath" "strings" "testing" + "github.com/ncxton/potaco/internal/config" "github.com/ncxton/potaco/internal/credential" ) @@ -87,7 +89,7 @@ func TestAuthAddDoesNotSetDefaultModel(t *testing.T) { func TestAuthManagerAddProviderStoresType(t *testing.T) { auth := newTestAuth(t) - if err := auth.AddProvider("openrouter", "openai-compatible", "sk-test"); err != nil { + if err := auth.AddProvider("openrouter", "openai", "sk-test"); err != nil { t.Fatalf("AddProvider: %v", err) } @@ -96,14 +98,54 @@ func TestAuthManagerAddProviderStoresType(t *testing.T) { t.Fatalf("LoadConfig: %v", err) } pc := cfg.Providers["openrouter"] - if pc.Type != "openai-compatible" { - t.Fatalf("Type = %q, want openai-compatible", pc.Type) + if pc.Type != "openai" { + t.Fatalf("Type = %q, want openai", pc.Type) } if cfg.ActiveProvider != "openrouter" { t.Fatalf("ActiveProvider = %q, want openrouter", cfg.ActiveProvider) } } +func TestAddProviderSetupRestoresExistingCredentialOnConfigFailure(t *testing.T) { + auth := newTestAuth(t) + if err := auth.AddProvider("openai", "openai", "sk-old"); err != nil { + t.Fatalf("AddProvider: %v", err) + } + previous, err := auth.store.GetCredential("openai") + if err != nil { + t.Fatalf("get previous credential: %v", err) + } + + cfg, err := auth.LoadConfig() + if err != nil { + t.Fatalf("LoadConfig: %v", err) + } + if err := config.SaveMultiProvider(auth.configPath, cfg); err != nil { + t.Fatalf("SaveMultiProvider: %v", err) + } + configDir := filepath.Dir(auth.configPath) + if err := os.Chmod(configDir, 0500); err != nil { + t.Fatalf("chmod config dir: %v", err) + } + t.Cleanup(func() { _ = os.Chmod(configDir, 0700) }) + + err = auth.AddProviderSetup(ProviderSetup{ + ProviderName: "openai", + ProviderType: "openai", + APIKey: "sk-new", + }) + if err == nil { + t.Fatal("AddProviderSetup should fail when config cannot be saved") + } + restored, err := auth.store.GetCredential("openai") + if err != nil { + t.Fatalf("get restored credential: %v", err) + } + if restored.APIKey != previous.APIKey || !restored.AddedAt.Equal(previous.AddedAt) { + t.Fatalf("credential = %+v, want restored %+v", restored, previous) + } +} + func TestAuthAddSetsActiveProvider(t *testing.T) { auth := newTestAuth(t) diff --git a/internal/auth/setup.go b/internal/auth/setup.go new file mode 100644 index 0000000..1ae2fc4 --- /dev/null +++ b/internal/auth/setup.go @@ -0,0 +1,75 @@ +package auth + +import ( + "fmt" + + "github.com/ncxton/potaco/internal/config" + "github.com/ncxton/potaco/internal/credential" +) + +type ProviderSetup struct { + ProviderName string + ProviderType string + APIKey string + BaseURL string + ModelID string + ModelConfig config.ModelConfig +} + +func (m *AuthManager) AddProviderSetup(setup ProviderSetup) error { + provider := setup.ProviderName + providerType := setup.ProviderType + previous, previousErr := m.store.GetCredential(provider) + hadPrevious := previousErr == nil + if err := m.store.Set(provider, setup.APIKey); err != nil { + return fmt.Errorf("store credential: %w", err) + } + + cfg, err := m.LoadConfig() + if err != nil { + m.rollbackCredential(provider, previous, hadPrevious) + return fmt.Errorf("load config: %w", err) + } + if cfg.Providers == nil { + cfg.Providers = make(map[string]config.ProviderConfig) + } + + pc := cfg.Providers[provider] + if providerType == "" { + providerType = config.ResolveProviderType(provider, pc) + } + pc.Type = providerType + if pc.Retries == 0 { + pc.Retries = 2 + } + if pc.Timeout == 0 { + pc.Timeout = 120 + } + if setup.BaseURL != "" { + pc.BaseURL = setup.BaseURL + } + if setup.ModelID != "" { + pc.Model = setup.ModelID + if pc.Models == nil { + pc.Models = make(map[string]config.ModelConfig) + } + pc.Models[setup.ModelID] = setup.ModelConfig + } + cfg.Providers[provider] = pc + cfg.ActiveProvider = provider + cfg.ActiveModel = pc.Model + + if err := m.saveConfig(cfg); err != nil { + m.rollbackCredential(provider, previous, hadPrevious) + return err + } + return nil +} + +func (m *AuthManager) rollbackCredential(provider string, previous credential.ProviderCredential, hadPrevious bool) { + if hadPrevious { + _ = m.store.SetCredential(provider, previous) + return + } + _ = m.store.Remove(provider) +} diff --git a/internal/cli/auth_add_base_url.go b/internal/cli/auth_add_base_url.go new file mode 100644 index 0000000..cc1a6c4 --- /dev/null +++ b/internal/cli/auth_add_base_url.go @@ -0,0 +1,34 @@ +package cli + +import ( + "fmt" + + "github.com/ncxton/potaco/internal/config" + "github.com/ncxton/potaco/internal/tui" +) + +var ( + authAddIsInteractive = tui.IsInteractive + authAddConfirmAction = tui.ConfirmAction +) + +func confirmAuthAddBaseURLSafety(baseURL string) error { + if !config.BaseURLIsInsecureHTTP(baseURL) { + return nil + } + if !authAddIsInteractive() { + return configUserErr( + "Refusing unencrypted HTTP base URL.", + "Run interactively and confirm you trust the URL, or use https://.", + fmt.Errorf("insecure base URL: %s", baseURL), + ) + } + trusted, err := authAddConfirmAction(fmt.Sprintf(" Danger: %s uses unencrypted HTTP. Continue only if you trust this URL?", baseURL)) + if err != nil { + return configError(err) + } + if !trusted { + return configError(fmt.Errorf("cancelled by user")) + } + return nil +} diff --git a/internal/cli/auth_add_custom_cmd_test.go b/internal/cli/auth_add_custom_cmd_test.go index dd3c6e7..31c8cab 100644 --- a/internal/cli/auth_add_custom_cmd_test.go +++ b/internal/cli/auth_add_custom_cmd_test.go @@ -24,6 +24,38 @@ func TestAuthAddCustomRequiresBaseURL(t *testing.T) { } } +func TestAuthAddCustomRejectsInvalidBaseURL(t *testing.T) { + tmpHome, _ := newAuthTest(t) + rootCmd.SetArgs([]string{"auth", "add", "custom", "--api-key", "sk-test", "--base-url", "not-a-url", "--force"}) + + err := rootCmd.Execute() + if err == nil { + t.Fatal("auth add custom with invalid base-url should error") + } + if !strings.Contains(err.Error(), "base URL") { + t.Fatalf("error = %v, want base URL", err) + } + if _, statErr := os.Stat(filepath.Join(tmpHome, ".potaco", "config.yaml")); !os.IsNotExist(statErr) { + t.Fatalf("config should not be written, stat err = %v", statErr) + } +} + +func TestAuthAddCustomRejectsHTTPBaseURLInNonInteractiveMode(t *testing.T) { + tmpHome, _ := newAuthTest(t) + rootCmd.SetArgs([]string{"auth", "add", "custom", "--api-key", "sk-test", "--base-url", "http://localhost:8080/v1", "--force", "--non-interactive"}) + + err := rootCmd.Execute() + if err == nil { + t.Fatal("auth add custom with http base-url should error in non-interactive mode") + } + if !strings.Contains(err.Error(), "http://") { + t.Fatalf("error = %v, want http warning", err) + } + if _, statErr := os.Stat(filepath.Join(tmpHome, ".potaco", "config.yaml")); !os.IsNotExist(statErr) { + t.Fatalf("config should not be written, stat err = %v", statErr) + } +} + func TestAuthAddCustomWithBaseURLForce(t *testing.T) { tmpHome, buf := newAuthTest(t) rootCmd.SetArgs([]string{"auth", "add", "custom", "--api-key", "sk-test", "--base-url", "https://example.com/v1", "--force"}) @@ -50,8 +82,8 @@ func TestAuthAddCustomWithBaseURLForce(t *testing.T) { if err != nil { t.Fatalf("load config: %v", err) } - if got := cfg.Providers["custom"].Type; got != "openai-compatible" { - t.Fatalf("custom type = %q, want openai-compatible", got) + if got := cfg.Providers["custom"].Type; got != "openai" { + t.Fatalf("custom type = %q, want openai", got) } } @@ -115,6 +147,15 @@ func TestAuthAddCustomVerifyUsesBaseURL(t *testing.T) { defer srv.Close() tmpHome, _ := newAuthTest(t) + oldIsInteractive := authAddIsInteractive + oldConfirmAction := authAddConfirmAction + t.Cleanup(func() { + authAddIsInteractive = oldIsInteractive + authAddConfirmAction = oldConfirmAction + }) + authAddIsInteractive = func() bool { return true } + authAddConfirmAction = func(prompt string) (bool, error) { return true, nil } + rootCmd.SetArgs([]string{"auth", "add", "custom", "--api-key", "sk-test", "--base-url", srv.URL}) err := rootCmd.Execute() diff --git a/internal/cli/auth_add_provider_type_cmd_test.go b/internal/cli/auth_add_provider_type_cmd_test.go index 131f2d8..cb61e2d 100644 --- a/internal/cli/auth_add_provider_type_cmd_test.go +++ b/internal/cli/auth_add_provider_type_cmd_test.go @@ -26,7 +26,7 @@ func TestAuthAddCustomNamedProviderWithTypeAndBaseURL(t *testing.T) { newAuthTest(t) rootCmd.SetArgs([]string{ "auth", "add", "openrouter", - "--type", "openai-compatible", + "--type", "openai", "--api-key", "sk-test", "--base-url", "https://openrouter.ai/api/v1", "--force", @@ -50,8 +50,8 @@ func TestAuthAddCustomNamedProviderWithTypeAndBaseURL(t *testing.T) { if !ok { t.Fatal("openrouter provider should be configured") } - if pc.Type != "openai-compatible" { - t.Fatalf("provider type = %q, want openai-compatible", pc.Type) + if pc.Type != "openai" { + t.Fatalf("provider type = %q, want openai", pc.Type) } if pc.BaseURL != "https://openrouter.ai/api/v1" { t.Fatalf("base URL = %q, want https://openrouter.ai/api/v1", pc.BaseURL) @@ -59,7 +59,7 @@ func TestAuthAddCustomNamedProviderWithTypeAndBaseURL(t *testing.T) { } func TestAuthAddAliasWithBuiltInTypeRequiresBaseURL(t *testing.T) { - tests := []string{"openai", "fal", "vercel"} + tests := []string{"openai", "fal", "replicate"} for _, providerType := range tests { t.Run(providerType, func(t *testing.T) { newAuthTest(t) @@ -188,9 +188,9 @@ func TestShouldRunInteractiveAuthAdd(t *testing.T) { want: false, }, { - name: "openai-compatible provider without base URL prompts in interactive mode", + name: "openai alias without base URL prompts in interactive mode", providerName: "openrouter", - providerTypeFlag: "openai-compatible", + providerTypeFlag: "openai", apiKey: "sk-test", interactive: true, want: true, @@ -212,9 +212,9 @@ func TestShouldRunInteractiveAuthAdd(t *testing.T) { want: true, }, { - name: "vercel alias without base URL prompts in interactive mode", - providerName: "staging-vercel", - providerTypeFlag: "vercel", + name: "replicate alias without base URL prompts in interactive mode", + providerName: "staging-replicate", + providerTypeFlag: "replicate", apiKey: "sk-test", interactive: true, want: true, @@ -226,9 +226,9 @@ func TestShouldRunInteractiveAuthAdd(t *testing.T) { want: true, }, { - name: "complete openai-compatible provider does not prompt", + name: "complete openai alias does not prompt", providerName: "openrouter", - providerTypeFlag: "openai-compatible", + providerTypeFlag: "openai", apiKey: "sk-test", baseURL: "https://openrouter.ai/api/v1", interactive: true, diff --git a/internal/cli/auth_add_provider_type_reject_test.go b/internal/cli/auth_add_provider_type_reject_test.go new file mode 100644 index 0000000..1edf2ee --- /dev/null +++ b/internal/cli/auth_add_provider_type_reject_test.go @@ -0,0 +1,44 @@ +package cli + +import ( + "strings" + "testing" + + "github.com/ncxton/potaco/internal/config" +) + +func TestAuthAddCustomNamedProviderRejectsOpenAICompatibleType(t *testing.T) { + newAuthTest(t) + rootCmd.SetArgs([]string{ + "auth", "add", "openrouter", + "--type", "openai-compatible", + "--api-key", "sk-test", + "--base-url", "https://openrouter.ai/api/v1", + "--force", + "--non-interactive", + }) + + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for openai-compatible provider type") + } + if !strings.Contains(err.Error(), "unknown provider type: openai-compatible") { + t.Fatalf("error = %v, want openai-compatible type rejection", err) + } +} + +func TestResolveAuthProviderTypeNormalizesLegacyOpenAICompatible(t *testing.T) { + cfg := &config.MultiProviderConfig{ + Providers: map[string]config.ProviderConfig{ + "openrouter": {Type: "openai-compatible"}, + }, + } + + providerType, err := resolveAuthProviderType("openrouter", "", cfg) + if err != nil { + t.Fatalf("resolve provider type: %v", err) + } + if providerType != "openai" { + t.Fatalf("provider type = %q, want openai", providerType) + } +} diff --git a/internal/cli/auth_cmd.go b/internal/cli/auth_cmd.go index 52df891..93d90fa 100644 --- a/internal/cli/auth_cmd.go +++ b/internal/cli/auth_cmd.go @@ -47,7 +47,7 @@ func init() { authAddCmd.Flags().Bool("force", false, "skip provider verification") authAddCmd.Flags().String("model", "", "override the default model for this provider") authAddCmd.Flags().String("base-url", "", "override API base URL for verification") - authAddCmd.Flags().String("type", "", "provider adapter type (openai, fal, vercel, openai-compatible)") + authAddCmd.Flags().String("type", "", "provider adapter type (openai, fal, replicate)") authCmd.AddCommand(authAddCmd) authCmd.AddCommand(authRemoveCmd) @@ -98,10 +98,19 @@ func runAuthAdd(cmd *cobra.Command, args []string) error { if err != nil { return configError(err) } - adapterType := config.AdapterType(providerType) + adapterType := config.AdapterTypeForProvider(providerName, providerType) if baseURL == "" && !authAddRequiresBaseURL(providerName, providerType) { baseURL = resolveBaseURL(cmd, providerName, cfg) } + if baseURL != "" { + baseURL, err = config.NormalizeBaseURL(baseURL) + if err != nil { + return configUserErr("Invalid base URL.", "Use a full http:// or https:// URL.", err) + } + if err := confirmAuthAddBaseURLSafety(baseURL); err != nil { + return err + } + } if apiKey == "" { return configError(fmt.Errorf("API key required: use --api-key or set POTACO_API_KEY")) @@ -121,29 +130,26 @@ func runAuthAdd(cmd *cobra.Command, args []string) error { return configError(fmt.Errorf("create adapter: %w", err)) } if err := ad.Verify(context.Background()); err != nil { - return apiError(fmt.Errorf("verification failed: %w\nUse --force to add anyway", err)) - } - } - - if err := mgr.AddProvider(providerName, providerType, apiKey); err != nil { - return configError(fmt.Errorf("add provider: %w", err)) - } - - if shouldPersistAuthBaseURL(cmd, providerType, baseURL) { - if err := mgr.SetBaseURL(providerName, baseURL); err != nil { - return configError(fmt.Errorf("set base URL: %w", err)) + return apiUserErr("Verification failed.", "Use --force to add anyway.", fmt.Errorf("verification failed: %w", err)) } } model, _ := cmd.Flags().GetString("model") - if model != "" { - if err := mgr.SetActiveProvider(providerName, model); err != nil { - return configError(fmt.Errorf("set model: %w", err)) - } + persistedBaseURL := "" + if authAddRequiresBaseURL(providerName, providerType) || cmd.Flags().Changed("base-url") || os.Getenv("POTACO_BASE_URL") != "" { + persistedBaseURL = baseURL + } + if err := mgr.AddProviderSetup(auth.ProviderSetup{ + ProviderName: providerName, + ProviderType: providerType, + APIKey: apiKey, + BaseURL: persistedBaseURL, + ModelID: model, + }); err != nil { + return configError(fmt.Errorf("add provider: %w", err)) } - fmt.Fprintf(cmd.OutOrStdout(), "Provider '%s' added successfully.\n", providerName) - fmt.Fprintf(cmd.OutOrStdout(), "Use 'potaco use %s' to switch to it.\n", providerName) + renderSuccessf(cmd.OutOrStdout(), "Provider '%s' added successfully. Use 'potaco use %s' to switch to it.", displayProviderName(providerName), providerName) return nil } @@ -180,7 +186,7 @@ func runAuthRemove(cmd *cobra.Command, args []string) error { if err := mgr.Remove(providerName); err != nil { return configError(fmt.Errorf("remove provider: %w", err)) } - fmt.Fprintf(cmd.OutOrStdout(), "Provider '%s' removed.\n", providerName) + renderSuccessf(cmd.OutOrStdout(), "Provider '%s' removed.", displayProviderName(providerName)) return nil } @@ -237,17 +243,17 @@ func runAuthList(cmd *cobra.Command, args []string) error { fmt.Fprintln(out, titleStyle.Render("Connected providers:")) fmt.Fprintln(out) for _, p := range providers { - name := p.Name + name := displayProviderName(p.Name) if p.IsActive { - name = activeStyle.Render(p.Name + " (active)") + name = activeStyle.Render(name + " (active)") } keyStatus := keyOkStyle.Render("configured") if !p.HasKey { keyStatus = keyMissingStyle.Render("missing") } - fmt.Fprintf(out, " %s %s key: %s\n", name, p.Model, keyStatus) + fmt.Fprintf(out, " %s %s key: %s\n", name, safeTerminalText(p.Model), keyStatus) if p.BaseURL != "" { - fmt.Fprintf(out, " base_url: %s\n", p.BaseURL) + fmt.Fprintf(out, " base_url: %s\n", safeTerminalText(p.BaseURL)) } } return nil @@ -264,9 +270,9 @@ func runAuthList(cmd *cobra.Command, args []string) error { if p.HasKey { keyStatus = "configured" } - fmt.Fprintf(out, " %s\t%s\tkey: %s%s\n", p.Name, p.Model, keyStatus, active) + fmt.Fprintf(out, " %s\t%s\tkey: %s%s\n", displayProviderName(p.Name), safeTerminalText(p.Model), keyStatus, active) if p.BaseURL != "" { - fmt.Fprintf(out, " base_url: %s\n", p.BaseURL) + fmt.Fprintf(out, " base_url: %s\n", safeTerminalText(p.BaseURL)) } } return nil diff --git a/internal/cli/auth_cmd_test_helpers_test.go b/internal/cli/auth_cmd_test_helpers_test.go index 50c5a10..ca401d7 100644 --- a/internal/cli/auth_cmd_test_helpers_test.go +++ b/internal/cli/auth_cmd_test_helpers_test.go @@ -32,3 +32,15 @@ func resetAuthAddFlags(t *testing.T) { flag.Changed = false } } + +func trustHTTPBaseURLs(t *testing.T) { + t.Helper() + oldIsInteractive := authAddIsInteractive + oldConfirmAction := authAddConfirmAction + t.Cleanup(func() { + authAddIsInteractive = oldIsInteractive + authAddConfirmAction = oldConfirmAction + }) + authAddIsInteractive = func() bool { return true } + authAddConfirmAction = func(string) (bool, error) { return true, nil } +} diff --git a/internal/cli/auth_list_cmd_test.go b/internal/cli/auth_list_cmd_test.go index 2a523e7..2f5bde9 100644 --- a/internal/cli/auth_list_cmd_test.go +++ b/internal/cli/auth_list_cmd_test.go @@ -33,8 +33,11 @@ func TestAuthListCommand(t *testing.T) { if !strings.Contains(output, "openai") { t.Errorf("list should include openai, got: %q", output) } - if !strings.Contains(output, "fal") { - t.Errorf("list should include fal, got: %q", output) + if !strings.Contains(output, "Fal") { + t.Errorf("list should display Fal, got: %q", output) + } + if strings.Contains(output, " fal\t") { + t.Errorf("list should not display lowercase fal row, got: %q", output) } } diff --git a/internal/cli/auth_provider_type.go b/internal/cli/auth_provider_type.go index 6e2803a..8370d29 100644 --- a/internal/cli/auth_provider_type.go +++ b/internal/cli/auth_provider_type.go @@ -2,10 +2,8 @@ package cli import ( "fmt" - "os" "github.com/ncxton/potaco/internal/config" - "github.com/spf13/cobra" ) func resolveAuthProviderType(providerName, providerType string, cfg *config.MultiProviderConfig) (string, error) { @@ -17,30 +15,23 @@ func resolveAuthProviderType(providerName, providerType string, cfg *config.Mult } if cfg != nil { if pc, ok := cfg.Providers[providerName]; ok && pc.Type != "" { - return pc.Type, nil + return config.ResolveProviderType(providerName, pc), nil } } if providerName == "custom" { - return "openai-compatible", nil + return "openai", nil } if isKnownProvider(providerName, nil) { return providerName, nil } - return "", fmt.Errorf("provider type required for %q: use --type openai-compatible, openai, fal, or vercel", providerName) + return "", fmt.Errorf("provider type required for %q: use --type openai, fal, or replicate", providerName) } func isAllowedAuthProviderType(providerType string) bool { switch providerType { - case "openai", "fal", "vercel", "openai-compatible": + case "openai", "fal", "replicate": return true default: return false } } - -func shouldPersistAuthBaseURL(cmd *cobra.Command, providerType, baseURL string) bool { - if baseURL == "" { - return false - } - return providerType == "openai-compatible" || cmd.Flags().Changed("base-url") || os.Getenv("POTACO_BASE_URL") != "" -} diff --git a/internal/cli/base_url.go b/internal/cli/base_url.go new file mode 100644 index 0000000..8b739e3 --- /dev/null +++ b/internal/cli/base_url.go @@ -0,0 +1,29 @@ +package cli + +import "github.com/ncxton/potaco/internal/config" + +func validateCommandBaseURL(baseURL string) (string, error) { + return validateBaseURL(baseURL) +} + +func validateDiscoveryBaseURL(baseURL string) (string, error) { + return validateBaseURL(baseURL) +} + +func validateRuntimeBaseURL(baseURL string) (string, error) { + return validateBaseURL(baseURL) +} + +func validateBaseURL(baseURL string) (string, error) { + if baseURL == "" { + return "", nil + } + normalized, err := config.NormalizeBaseURL(baseURL) + if err != nil { + return "", configUserErr("Invalid base URL.", "Use a full http:// or https:// URL.", err) + } + if err := confirmAuthAddBaseURLSafety(normalized); err != nil { + return "", err + } + return normalized, nil +} diff --git a/internal/cli/config_cmd.go b/internal/cli/config_cmd.go index d6b65ec..58e0391 100644 --- a/internal/cli/config_cmd.go +++ b/internal/cli/config_cmd.go @@ -2,7 +2,7 @@ package cli import ( "fmt" - "sort" + "os" "strconv" "strings" @@ -60,7 +60,7 @@ func runConfigSet(cmd *cobra.Command, args []string) error { if err := config.SaveMultiProvider(path, cfg); err != nil { return configError(err) } - fmt.Fprintf(cmd.OutOrStdout(), "Configuration saved to %s\n", path) + renderSuccessf(cmd.OutOrStdout(), "Configuration saved to %s.", path) return nil } @@ -82,7 +82,11 @@ func runConfigSet(cmd *cobra.Command, args []string) error { } if cmd.Flags().Changed("base-url") { baseURL, _ := cmd.Flags().GetString("base-url") - pc.BaseURL = strings.TrimRight(baseURL, "/") + normalized, err := validateCommandBaseURL(baseURL) + if err != nil { + return err + } + pc.BaseURL = normalized changed = true } if cmd.Flags().Changed("retries") { @@ -110,7 +114,7 @@ func runConfigSet(cmd *cobra.Command, args []string) error { return configError(err) } - fmt.Fprintf(cmd.OutOrStdout(), "Configuration saved to %s\n", path) + renderSuccessf(cmd.OutOrStdout(), "Configuration saved to %s.", path) return nil } @@ -129,6 +133,13 @@ func setConfigKeyValue(cfg *config.MultiProviderConfig, key, value string) error if len(parts) == 3 && parts[1] != "" { return setProviderConfigValue(cfg, parts[1], parts[2], value) } + if len(parts) >= 6 && parts[1] != "" && parts[2] == "models" && parts[len(parts)-2] == "parameters" { + modelID := strings.Join(parts[3:len(parts)-2], ".") + if modelID == "" { + return fmt.Errorf("unknown config key %q", key) + } + return setProviderModelParameterValue(cfg, parts[1], modelID, parts[len(parts)-1], value) + } if len(parts) >= 5 && parts[1] != "" && parts[2] == "models" && parts[len(parts)-1] == "edit" { modelID := strings.Join(parts[3:len(parts)-1], ".") if modelID == "" { @@ -157,6 +168,17 @@ func setConfigKeyValue(cfg *config.MultiProviderConfig, key, value string) error } return setProviderModelEditValue(cfg, cfg.ActiveProvider, pc.Model, value) default: + if strings.HasPrefix(key, "model.parameters.") { + if cfg.ActiveProvider == "" { + return fmt.Errorf("no active provider. Use 'potaco auth add ' to connect one") + } + pc := cfg.Providers[cfg.ActiveProvider] + if pc.Model == "" { + return fmt.Errorf("active provider %q has no model configured", cfg.ActiveProvider) + } + field := strings.TrimPrefix(key, "model.parameters.") + return setProviderModelParameterValue(cfg, cfg.ActiveProvider, pc.Model, field, value) + } return fmt.Errorf("unknown config key %q", key) } } @@ -176,7 +198,11 @@ func setProviderConfigValue(cfg *config.MultiProviderConfig, providerName, field cfg.ActiveModel = value } case "base_url": - pc.BaseURL = strings.TrimRight(value, "/") + normalized, err := validateCommandBaseURL(value) + if err != nil { + return err + } + pc.BaseURL = normalized case "retries": retries, err := strconv.Atoi(value) if err != nil { @@ -221,75 +247,20 @@ func setProviderModelEditValue(cfg *config.MultiProviderConfig, providerName, mo func runConfigShow(cmd *cobra.Command, args []string) error { path := config.DefaultConfigPath() - cfg, err := config.LoadMultiProvider(path) + raw, err := os.ReadFile(path) if err != nil { - return configError(err) + if os.IsNotExist(err) { + // Keep this as normal command output: no config file is an empty state, + // not a failed command. + fmt.Fprintln(cmd.OutOrStdout(), "No configuration found. Use 'potaco auth add ' to connect.") + return nil + } + return configError(fmt.Errorf("read config: %w", err)) } - - // LoadMultiProvider returns an empty config (not an error) when the - // file does not exist. Treat an empty config as "not configured". - if cfg.ActiveProvider == "" && len(cfg.Providers) == 0 { + if len(raw) == 0 { fmt.Fprintln(cmd.OutOrStdout(), "No configuration found. Use 'potaco auth add ' to connect.") return nil } - - out := cmd.OutOrStdout() - fmt.Fprintf(out, "Config file: %s\n\n", path) - fmt.Fprintf(out, "Active provider: %s\n", cfg.ActiveProvider) - fmt.Fprintf(out, "Active model: %s\n", cfg.ActiveModel) - fmt.Fprintln(out) - - if len(cfg.Providers) == 0 { - return nil - } - - names := make([]string, 0, len(cfg.Providers)) - for name := range cfg.Providers { - names = append(names, name) - } - sort.Strings(names) - - fmt.Fprintln(out, "Providers:") - for _, name := range names { - pc := cfg.Providers[name] - active := "" - if name == cfg.ActiveProvider { - active = " (active)" - } - fmt.Fprintf(out, " %s%s\n", name, active) - fmt.Fprintf(out, " model: %s\n", pc.Model) - fmt.Fprintf(out, " base_url: %s\n", formatBaseURL(pc.BaseURL)) - fmt.Fprintf(out, " retries: %d\n", pc.Retries) - fmt.Fprintf(out, " timeout: %s\n", formatTimeout(pc.Timeout)) - if len(pc.Models) > 0 { - models := make([]string, 0, len(pc.Models)) - for model := range pc.Models { - models = append(models, model) - } - sort.Strings(models) - fmt.Fprintln(out, " models:") - for _, model := range models { - fmt.Fprintf(out, " %s: edit=%t\n", model, pc.Models[model].Edit) - } - } - } + fmt.Fprint(cmd.OutOrStdout(), string(raw)) return nil } - -// formatTimeout renders the timeout (in seconds) for display, returning -// "default" for a zero value so the output is informative. -func formatTimeout(secs int) string { - if secs <= 0 { - return "default" - } - return fmt.Sprintf("%ds", secs) -} - -// formatBaseURL renders the base URL for display, returning "default" -// for a zero value so providers without an override are easy to spot. -func formatBaseURL(baseURL string) string { - if baseURL == "" { - return "default" - } - return baseURL -} diff --git a/internal/cli/config_cmd_test.go b/internal/cli/config_cmd_test.go index d2fdfcd..5321084 100644 --- a/internal/cli/config_cmd_test.go +++ b/internal/cli/config_cmd_test.go @@ -316,45 +316,35 @@ func TestConfigSetWritesPrivateFileMode(t *testing.T) { } } -func TestConfigShowDisplaysActiveProvider(t *testing.T) { +func TestConfigShowPrintsRawConfigFile(t *testing.T) { path, buf := newConfigTest(t) - writeMultiProviderConfig(t, path, &config.MultiProviderConfig{ - ActiveProvider: "openai", - ActiveModel: "gpt-image-2", - Providers: map[string]config.ProviderConfig{ - "openai": {Model: "gpt-image-2", Retries: 3, Timeout: 90}, - }, - }) + raw := []byte("active_provider: openai\nactive_model: gpt-image-2\nproviders:\n openai:\n model: gpt-image-2\n retries: 3\n timeout: 90\n") + if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil { + t.Fatalf("create config dir: %v", err) + } + if err := os.WriteFile(path, raw, 0600); err != nil { + t.Fatalf("write config: %v", err) + } rootCmd.SetArgs([]string{"config", "show"}) if err := rootCmd.Execute(); err != nil { t.Fatalf("config show error: %v", err) } - output := buf.String() - for _, want := range []string{"openai", "gpt-image-2", "Active provider:", "Active model:"} { - if !strings.Contains(output, want) { - t.Errorf("config show should contain %q, got: %q", want, output) - } - } - if !strings.Contains(output, "retries: 3") { - t.Errorf("config show should display retries: 3, got: %q", output) - } - if !strings.Contains(output, "timeout: 90s") { - t.Errorf("config show should display timeout: 90s, got: %q", output) + if got := buf.String(); got != string(raw) { + t.Fatalf("config show output = %q, want raw config %q", got, string(raw)) } } -func TestConfigShowMarksActiveProvider(t *testing.T) { +func TestConfigShowPrintsRawConfigWithoutFormattingActiveProvider(t *testing.T) { path, buf := newConfigTest(t) - writeMultiProviderConfig(t, path, &config.MultiProviderConfig{ - ActiveProvider: "fal", - ActiveModel: "fal-ai/flux/dev", - Providers: map[string]config.ProviderConfig{ - "openai": {Model: "gpt-image-2", Retries: 2, Timeout: 120}, - "fal": {Model: "fal-ai/flux/dev", Retries: 3, Timeout: 60}, - }, - }) + raw := []byte("active_provider: fal\nactive_model: fal-ai/flux/dev\nproviders:\n fal:\n model: fal-ai/flux/dev\n") + if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil { + t.Fatalf("create config dir: %v", err) + } + if err := os.WriteFile(path, raw, 0600); err != nil { + t.Fatalf("write config: %v", err) + } rootCmd.SetArgs([]string{"config", "show"}) if err := rootCmd.Execute(); err != nil { @@ -362,15 +352,11 @@ func TestConfigShowMarksActiveProvider(t *testing.T) { } output := buf.String() - if !strings.Contains(output, "fal (active)") { - t.Errorf("config show should mark 'fal (active)', got: %q", output) + if output != string(raw) { + t.Fatalf("config show output = %q, want raw config %q", output, string(raw)) } - if !strings.Contains(output, "openai") { - t.Errorf("config show should list openai provider, got: %q", output) - } - // openai should NOT be marked active. - if strings.Contains(output, "openai (active)") { - t.Errorf("config show should NOT mark openai as active, got: %q", output) + if strings.Contains(output, "(active)") || strings.Contains(output, "Config file:") { + t.Fatalf("config show should not add formatted labels, got: %q", output) } } @@ -391,19 +377,12 @@ func TestConfigShowNoConfig(t *testing.T) { } } -func TestConfigShowDoesNotPrintAPIKey(t *testing.T) { +func TestConfigShowPrintsRawFileContent(t *testing.T) { path, buf := newConfigTest(t) - writeMultiProviderConfig(t, path, &config.MultiProviderConfig{ - ActiveProvider: "openai", - ActiveModel: "gpt-image-2", - Providers: map[string]config.ProviderConfig{ - "openai": {Model: "gpt-image-2", Retries: 2, Timeout: 120}, - }, - }) - - // Drop a fake-looking secret into the config file to ensure show - // never reads or prints credential-store contents. raw := []byte("active_provider: openai\nactive_model: gpt-image-2\nproviders:\n openai:\n model: gpt-image-2\n retries: 2\n timeout: 120\n") + if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil { + t.Fatalf("create config dir: %v", err) + } if err := os.WriteFile(path, raw, 0600); err != nil { t.Fatalf("write config: %v", err) } @@ -413,14 +392,8 @@ func TestConfigShowDoesNotPrintAPIKey(t *testing.T) { t.Fatalf("config show error: %v", err) } - output := buf.String() - // API keys live in the credential store, not the config file, so - // there should be no api_key field in the output at all. - if strings.Contains(output, "api_key") { - t.Errorf("config show should not display api_key, got: %q", output) - } - if strings.Contains(output, "sk-") { - t.Errorf("config show should not print any API key, got: %q", output) + if got := buf.String(); got != string(raw) { + t.Fatalf("config show output = %q, want raw config %q", got, string(raw)) } } @@ -459,6 +432,24 @@ func TestConfigSetBaseURL(t *testing.T) { } } +func TestConfigSetRejectsUntrustedHTTPBaseURL(t *testing.T) { + newConfigTest(t) + setupAuthProvider(t, "sk-test") + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"config", "set", "--base-url", "http://example.com/v1", "--non-interactive"}) + + err := rootCmd.Execute() + if err == nil { + t.Fatal("config set should reject untrusted http base_url") + } + if !strings.Contains(err.Error(), "insecure base URL") { + t.Fatalf("error = %v, want insecure base URL", err) + } +} + func TestConfigSetBaseURLTrimsTrailingSlash(t *testing.T) { path, _ := newConfigTest(t) writeMultiProviderConfig(t, path, &config.MultiProviderConfig{ @@ -569,27 +560,27 @@ func TestConfigSetKeyValueExplicitProviderFields(t *testing.T) { check func(t *testing.T, pc config.ProviderConfig) }{ { - key: "providers.vercel.model", - value: "openai/gpt-image-2", + key: "providers.replicate.model", + value: "black-forest-labs/flux-schnell", check: func(t *testing.T, pc config.ProviderConfig) { t.Helper() - if pc.Model != "openai/gpt-image-2" { - t.Fatalf("model = %q, want openai/gpt-image-2", pc.Model) + if pc.Model != "black-forest-labs/flux-schnell" { + t.Fatalf("model = %q, want black-forest-labs/flux-schnell", pc.Model) } }, }, { - key: "providers.vercel.base_url", - value: "https://gateway.ai/v1/", + key: "providers.replicate.base_url", + value: "https://api.replicate.com/v1/", check: func(t *testing.T, pc config.ProviderConfig) { t.Helper() - if pc.BaseURL != "https://gateway.ai/v1" { + if pc.BaseURL != "https://api.replicate.com/v1" { t.Fatalf("base_url = %q, want trimmed URL", pc.BaseURL) } }, }, { - key: "providers.vercel.retries", + key: "providers.replicate.retries", value: "5", check: func(t *testing.T, pc config.ProviderConfig) { t.Helper() @@ -599,7 +590,7 @@ func TestConfigSetKeyValueExplicitProviderFields(t *testing.T) { }, }, { - key: "providers.vercel.timeout", + key: "providers.replicate.timeout", value: "75", check: func(t *testing.T, pc config.ProviderConfig) { t.Helper() @@ -615,8 +606,8 @@ func TestConfigSetKeyValueExplicitProviderFields(t *testing.T) { ActiveProvider: "openai", ActiveModel: "gpt-image-2", Providers: map[string]config.ProviderConfig{ - "openai": {Model: "gpt-image-2", Retries: 2, Timeout: 120}, - "vercel": {Model: "old", Retries: 1, Timeout: 30}, + "openai": {Model: "gpt-image-2", Retries: 2, Timeout: 120}, + "replicate": {Model: "old", Retries: 1, Timeout: 30}, }, }) @@ -629,7 +620,7 @@ func TestConfigSetKeyValueExplicitProviderFields(t *testing.T) { if err != nil { t.Fatalf("load config: %v", err) } - tc.check(t, cfg.Providers["vercel"]) + tc.check(t, cfg.Providers["replicate"]) }) } } @@ -687,6 +678,71 @@ func TestConfigSetActiveModelEditCapability(t *testing.T) { } } +func TestConfigSetActiveModelParameters(t *testing.T) { + path, _ := newConfigTest(t) + writeMultiProviderConfig(t, path, &config.MultiProviderConfig{ + ActiveProvider: "openai", + ActiveModel: "gpt-image-2", + Providers: map[string]config.ProviderConfig{ + "openai": {Model: "gpt-image-2"}, + }, + }) + + tests := []struct { + key string + value string + }{ + {"model.parameters.size", "1536x1024"}, + {"model.parameters.quality", "high"}, + {"model.parameters.n", "2"}, + {"model.parameters.response_format", "url"}, + {"model.parameters.seed", "42"}, + {"model.parameters.guidance_scale", "7.5"}, + {"model.parameters.negative_prompt", "blurry"}, + } + for _, tt := range tests { + rootCmd.SetArgs([]string{"config", "set", tt.key, tt.value}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("config set %s error: %v", tt.key, err) + } + } + + cfg, err := config.LoadMultiProvider(path) + if err != nil { + t.Fatalf("load config: %v", err) + } + params := cfg.Providers["openai"].Models["gpt-image-2"].Parameters + if params.Size != "1536x1024" || params.Quality != "high" || params.N != 2 || + params.ResponseFormat != "url" || params.Seed != 42 || + params.GuidanceScale != 7.5 || params.NegativePrompt != "blurry" { + t.Fatalf("parameters = %+v, want configured values", params) + } +} + +func TestConfigSetProviderModelParameterWithSlashID(t *testing.T) { + path, _ := newConfigTest(t) + writeMultiProviderConfig(t, path, &config.MultiProviderConfig{ + ActiveProvider: "replicate", + ActiveModel: "owner/model", + Providers: map[string]config.ProviderConfig{ + "replicate": {Model: "owner/model"}, + }, + }) + + rootCmd.SetArgs([]string{"config", "set", "providers.replicate.models.owner/model.parameters.guidance_scale", "3.5"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("config set provider model parameter error: %v", err) + } + + cfg, err := config.LoadMultiProvider(path) + if err != nil { + t.Fatalf("load config: %v", err) + } + if got := cfg.Providers["replicate"].Models["owner/model"].Parameters.GuidanceScale; got != 3.5 { + t.Fatalf("guidance_scale = %v, want 3.5", got) + } +} + func TestConfigSetAutoUpdate(t *testing.T) { path, _ := newConfigTest(t) writeMultiProviderConfig(t, path, &config.MultiProviderConfig{ @@ -800,7 +856,7 @@ func TestConfigSetNoFlagsIncludesBaseURL(t *testing.T) { _ = path } -func TestConfigShowDisplaysBaseURL(t *testing.T) { +func TestConfigShowRawOutputIncludesBaseURL(t *testing.T) { path, buf := newConfigTest(t) writeMultiProviderConfig(t, path, &config.MultiProviderConfig{ ActiveProvider: "openai", @@ -817,11 +873,11 @@ func TestConfigShowDisplaysBaseURL(t *testing.T) { output := buf.String() if !strings.Contains(output, "base_url: https://api.example.com/v1") { - t.Errorf("config show should display base_url, got: %q", output) + t.Errorf("config show raw output should include base_url, got: %q", output) } } -func TestConfigShowDisplaysModelEditCapability(t *testing.T) { +func TestConfigShowRawOutputIncludesModelEditCapability(t *testing.T) { path, buf := newConfigTest(t) writeMultiProviderConfig(t, path, &config.MultiProviderConfig{ ActiveProvider: "openai", @@ -842,12 +898,12 @@ func TestConfigShowDisplaysModelEditCapability(t *testing.T) { } output := buf.String() - if !strings.Contains(output, "models:") || !strings.Contains(output, "gpt-image-2: edit=true") { - t.Errorf("config show should include model edit capability, got: %s", output) + if !strings.Contains(output, "models:") || !strings.Contains(output, "edit: true") { + t.Errorf("config show raw output should include model edit capability, got: %s", output) } } -func TestConfigShowOldConfigWithoutBaseURL(t *testing.T) { +func TestConfigShowRawOldConfigWithoutBaseURL(t *testing.T) { path, buf := newConfigTest(t) // Old-format config without base_url. if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil { @@ -865,9 +921,9 @@ func TestConfigShowOldConfigWithoutBaseURL(t *testing.T) { output := buf.String() if !strings.Contains(output, "openai") { - t.Errorf("config show should display provider, got: %q", output) + t.Errorf("config show raw output should include provider, got: %q", output) } - if !strings.Contains(output, "base_url: default") { - t.Errorf("config show should display default base_url for old config, got: %q", output) + if strings.Contains(output, "base_url: default") { + t.Errorf("config show raw output should not synthesize default base_url, got: %q", output) } } diff --git a/internal/cli/config_model_parameters.go b/internal/cli/config_model_parameters.go new file mode 100644 index 0000000..85010ed --- /dev/null +++ b/internal/cli/config_model_parameters.go @@ -0,0 +1,47 @@ +package cli + +import ( + "fmt" + "strconv" + + "github.com/ncxton/potaco/internal/config" +) + +func setProviderModelParameterValue(cfg *config.MultiProviderConfig, providerName, modelID, field, value string) error { + if cfg.Providers == nil { + cfg.Providers = make(map[string]config.ProviderConfig) + } + pc, ok := cfg.Providers[providerName] + if !ok { + return fmt.Errorf("provider %q is not configured. Use 'potaco auth add %s' first", providerName, providerName) + } + if pc.Models == nil { + pc.Models = make(map[string]config.ModelConfig) + } + mc := pc.Models[modelID] + var err error + switch field { + case "size": + mc.Parameters.Size = value + case "quality": + mc.Parameters.Quality = value + case "n": + mc.Parameters.N, err = strconv.Atoi(value) + case "response_format": + mc.Parameters.ResponseFormat = value + case "seed": + mc.Parameters.Seed, err = strconv.Atoi(value) + case "guidance_scale": + mc.Parameters.GuidanceScale, err = strconv.ParseFloat(value, 64) + case "negative_prompt": + mc.Parameters.NegativePrompt = value + default: + return fmt.Errorf("unknown model parameter %q", field) + } + if err != nil { + return fmt.Errorf("model parameter %s has invalid value %q: %w", field, value, err) + } + pc.Models[modelID] = mc + cfg.Providers[providerName] = pc + return nil +} diff --git a/internal/cli/edit.go b/internal/cli/edit.go index aebe6dc..e0aad03 100644 --- a/internal/cli/edit.go +++ b/internal/cli/edit.go @@ -1,13 +1,12 @@ package cli import ( - "context" "fmt" "os" "time" - "github.com/ncxton/potaco/internal/adapter" img "github.com/ncxton/potaco/internal/image" + "github.com/ncxton/potaco/internal/stats" "github.com/spf13/cobra" ) @@ -32,9 +31,9 @@ func init() { // Shared flags from gen editCmd.Flags().String("model", "", "model to use") - editCmd.Flags().String("size", "1024x1024", "image dimensions (WxH)") - editCmd.Flags().Int("n", 1, "number of images to generate") - editCmd.Flags().String("response-format", "b64_json", "response format (url or b64_json)") + editCmd.Flags().String("size", "", "override configured image dimensions (WxH)") + editCmd.Flags().Int("n", 0, "override configured number of images") + editCmd.Flags().String("response-format", "", "override configured response format (url or b64_json)") // Output flags editCmd.Flags().StringP("output", "o", "", "output file path") @@ -42,7 +41,7 @@ func init() { editCmd.Flags().Bool("stdout", false, "pipe raw image bytes to stdout") // Provider override flags - editCmd.Flags().String("provider", "", "provider preset (openai, fal, vercel)") + editCmd.Flags().String("provider", "", "provider preset (openai, fal, replicate)") editCmd.Flags().String("base-url", "", "override API base URL") editCmd.Flags().String("api-key", "", "override API key") editCmd.Flags().Int("retries", 0, "max retry attempts") @@ -54,7 +53,18 @@ func init() { rootCmd.AddCommand(editCmd) } -func runEdit(cmd *cobra.Command, args []string) error { +func runEdit(cmd *cobra.Command, args []string) (err error) { + start := time.Now() + event := stats.Event{Command: cmd.CommandPath(), Operation: "edit"} + defer func() { + event.DurationMS = time.Since(start).Milliseconds() + event.Success = err == nil + if err != nil { + event.Error = err.Error() + } + recordStatsEvent(event) + }() + prompt := flagString(cmd, "prompt") if prompt == "" { return configUserErr("A prompt is required.", "Use 'potaco edit --prompt \"your description\"'.", fmt.Errorf("prompt cannot be empty")) @@ -85,6 +95,10 @@ func runEdit(cmd *cobra.Command, args []string) error { return err } model := resolved.Model + event.Provider = resolved.ProviderName + event.Model = model + event.DryRun = flagBool(cmd, "dry-run") + params := editParametersFromFlags(cmd, resolved.Parameters) if !resolved.Adapter.SupportsEdit() { return apiUserErr( @@ -106,8 +120,9 @@ func runEdit(cmd *cobra.Command, args []string) error { maskPath := "" if dryRun { + warnUnsupportedEditParameters(cmd, resolved.Adapter.Name(), params) authHeader := resolved.Adapter.AuthHeader("[REDACTED]") - return printEditDryRun(cmd, resolved.BaseURL, resolved.Adapter.Name(), authHeader, prompt, model, imagePath, cmd.Flags()) + return printEditDryRun(cmd, resolved.BaseURL, resolved.Adapter.Name(), authHeader, prompt, model, imagePath, params, cmd.Flags()) } cleanup := noopCleanup @@ -121,7 +136,7 @@ func runEdit(cmd *cobra.Command, args []string) error { // the source image. Use the actual canvas size so the API produces an // output matching the extended area, rather than the --size default // which would crop or resize away the new pixels. - size := flagString(cmd, "size") + extendedSize := "" extendFlag, _ := cmd.Flags().GetString("extend") if extendFlag != "" { canvas, _, err := img.ReadImage(editImagePath) @@ -133,25 +148,17 @@ func runEdit(cmd *cobra.Command, args []string) error { ) } bounds := canvas.Bounds() - size = fmt.Sprintf("%dx%d", bounds.Dx(), bounds.Dy()) + extendedSize = fmt.Sprintf("%dx%d", bounds.Dx(), bounds.Dy()) } - req := adapter.EditRequest{ - Prompt: prompt, - Model: model, - N: flagInt(cmd, "n"), - Size: size, - ResponseFormat: flagString(cmd, "response-format"), - ImagePath: editImagePath, - MaskPath: maskPath, - } - - sp := startSpinner(cmd, "Editing image...") - start := time.Now() - ctx := context.Background() - resp, err := resolved.Adapter.Edit(ctx, req) - sp.stop() - elapsed := time.Since(start) + requestStart := time.Now() + result, err := editWithFallback(cmd, resolved, editInput{ + prompt: prompt, + imagePath: editImagePath, + maskPath: maskPath, + extendedSize: extendedSize, + }) + elapsed := time.Since(requestStart) if err != nil { return apiUserErr( "Image editing failed.", @@ -159,6 +166,12 @@ func runEdit(cmd *cobra.Command, args []string) error { fmt.Errorf("edit: %w", err), ) } + resp := result.response + model = result.provider.Model + req := result.request + event.Provider = result.provider.ProviderName + event.Model = model + event.Outputs = len(resp.Data) if err := processAndOutput(cmd, outputContext{ resp: resp, diff --git a/internal/cli/edit_mask.go b/internal/cli/edit_mask.go index d0a2d50..2145ce4 100644 --- a/internal/cli/edit_mask.go +++ b/internal/cli/edit_mask.go @@ -8,12 +8,13 @@ import ( "strconv" "strings" + "github.com/ncxton/potaco/internal/config" img "github.com/ncxton/potaco/internal/image" "github.com/spf13/cobra" "github.com/spf13/pflag" ) -func printEditDryRun(cmd *cobra.Command, baseURL, providerName, authHeader, prompt, model, imagePath string, flags *pflag.FlagSet) error { +func printEditDryRun(cmd *cobra.Command, baseURL, providerName, authHeader, prompt, model, imagePath string, params config.ModelParameters, flags *pflag.FlagSet) error { extendFlag, _ := flags.GetString("extend") maskFlag, _ := flags.GetString("mask") maskRectFlag, _ := flags.GetString("mask-rect") @@ -33,9 +34,12 @@ func printEditDryRun(cmd *cobra.Command, baseURL, providerName, authHeader, prom } body := map[string]any{ - "prompt": prompt, - "model": model, - "mode": mode, + "prompt": prompt, + "model": model, + "mode": mode, + "size": params.Size, + "n": params.N, + "response_format": params.ResponseFormat, } editURL := baseURL + "/v1/images/edits" diff --git a/internal/cli/edit_test.go b/internal/cli/edit_test.go index 1e2f229..24a391b 100644 --- a/internal/cli/edit_test.go +++ b/internal/cli/edit_test.go @@ -337,6 +337,7 @@ func TestPrepareEditImageCleanupRemovesOutpaintDir(t *testing.T) { } func TestEditOutpaintUsesExpandedCanvasSize(t *testing.T) { + trustHTTPBaseURLs(t) resetRootCmdFlags(t) resetAuthAddFlags(t) resetEditCmdFlags(t) @@ -409,6 +410,7 @@ func TestEditOutpaintUsesExpandedCanvasSize(t *testing.T) { } func TestEditCleansGeneratedMaskDirAfterUpload(t *testing.T) { + trustHTTPBaseURLs(t) resetRootCmdFlags(t) resetAuthAddFlags(t) resetEditCmdFlags(t) @@ -656,11 +658,11 @@ func TestEditDryRunFalProvider(t *testing.T) { } } -func TestEditDryRunVercelNotSupported(t *testing.T) { +func TestEditDryRunReplicateNotSupported(t *testing.T) { resetRootCmdFlags(t) resetEditCmdFlags(t) t.Cleanup(func() { resetEditCmdFlags(t) }) - setupAuthProviderForProvider(t, "vercel", "vkey", "openai/gpt-image-2") + setupAuthProviderForProvider(t, "replicate", "r8_key", "black-forest-labs/flux-schnell") tmpDir := t.TempDir() imgPath := tmpDir + "/test.png" @@ -669,24 +671,24 @@ func TestEditDryRunVercelNotSupported(t *testing.T) { var buf bytes.Buffer rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) - rootCmd.SetArgs([]string{"edit", "--prompt", "test", "--image", imgPath, "--dry-run", "--provider", "vercel"}) + rootCmd.SetArgs([]string{"edit", "--prompt", "test", "--image", imgPath, "--dry-run", "--provider", "replicate"}) err := rootCmd.Execute() if err == nil { - t.Fatal("expected error for vercel edit, got nil") + t.Fatal("expected error for replicate edit, got nil") } if !strings.Contains(err.Error(), "not supported") { t.Errorf("error should mention 'not supported', got: %v", err) } - if !strings.Contains(err.Error(), "vercel") { - t.Errorf("error should mention provider 'vercel', got: %v", err) + if !strings.Contains(err.Error(), "replicate") { + t.Errorf("error should mention provider 'replicate', got: %v", err) } } -func TestEditVercelNotSupportedNonDryRun(t *testing.T) { +func TestEditReplicateNotSupportedNonDryRun(t *testing.T) { resetRootCmdFlags(t) resetEditCmdFlags(t) t.Cleanup(func() { resetEditCmdFlags(t) }) - setupAuthProviderForProvider(t, "vercel", "vkey", "openai/gpt-image-2") + setupAuthProviderForProvider(t, "replicate", "r8_key", "black-forest-labs/flux-schnell") tmpDir := t.TempDir() imgPath := tmpDir + "/test.png" createTestPNG(t, imgPath, 4, 4) @@ -695,16 +697,16 @@ func TestEditVercelNotSupportedNonDryRun(t *testing.T) { rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) // No --dry-run: should still be rejected before touching the image. - rootCmd.SetArgs([]string{"edit", "--prompt", "test", "--image", imgPath, "--provider", "vercel"}) + rootCmd.SetArgs([]string{"edit", "--prompt", "test", "--image", imgPath, "--provider", "replicate"}) err := rootCmd.Execute() if err == nil { - t.Fatal("expected error for vercel edit (non-dry-run), got nil") + t.Fatal("expected error for replicate edit (non-dry-run), got nil") } if !strings.Contains(err.Error(), "not supported") { t.Errorf("error should mention 'not supported', got: %v", err) } - if !strings.Contains(err.Error(), "vercel") { - t.Errorf("error should mention provider 'vercel', got: %v", err) + if !strings.Contains(err.Error(), "replicate") { + t.Errorf("error should mention provider 'replicate', got: %v", err) } } @@ -790,6 +792,93 @@ func TestEditDryRunOpenAIProviderMultipartContentType(t *testing.T) { } } +func TestEditWarnsForConfiguredGenerationOnlyParameters(t *testing.T) { + resetRootCmdFlags(t) + resetEditCmdFlags(t) + t.Cleanup(func() { resetEditCmdFlags(t) }) + setupAuthProvider(t, "sk-test") + for _, args := range [][]string{ + {"config", "set", "model.parameters.quality", "high"}, + {"config", "set", "model.parameters.seed", "42"}, + {"config", "set", "model.parameters.guidance_scale", "7.5"}, + {"config", "set", "model.parameters.negative_prompt", "blurry"}, + } { + rootCmd.SetArgs(args) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("%v: %v", args, err) + } + } + + tmpDir := t.TempDir() + imgPath := filepath.Join(tmpDir, "test.png") + createTestPNG(t, imgPath, 4, 4) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"edit", "--prompt", "test", "--image", imgPath, "--dry-run"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + output := buf.String() + for _, want := range []string{" Parameter quality", " Parameter seed", " Parameter guidance_scale", " Parameter negative_prompt"} { + if !strings.Contains(output, want) { + t.Fatalf("output missing warning %q:\n%s", want, output) + } + } +} + +func TestEditDryRunFlagsOverrideConfiguredModelParameters(t *testing.T) { + resetRootCmdFlags(t) + resetEditCmdFlags(t) + t.Cleanup(func() { resetEditCmdFlags(t) }) + setupAuthProvider(t, "sk-test") + + for _, args := range [][]string{ + {"config", "set", "model.parameters.size", "1536x1024"}, + {"config", "set", "model.parameters.n", "2"}, + {"config", "set", "model.parameters.response_format", "url"}, + } { + rootCmd.SetArgs(args) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("%v: %v", args, err) + } + } + + tmpDir := t.TempDir() + imgPath := filepath.Join(tmpDir, "test.png") + createTestPNG(t, imgPath, 4, 4) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{ + "edit", + "--prompt", "test", + "--image", imgPath, + "--size", "1024x1792", + "--n", "1", + "--response-format", "b64_json", + "--dry-run", + }) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + output := buf.String() + for _, want := range []string{`"size": "1024x1792"`, `"n": 1`, `"response_format": "b64_json"`} { + if !strings.Contains(output, want) { + t.Fatalf("dry-run output missing override %s:\n%s", want, output) + } + } + for _, stale := range []string{`"size": "1536x1024"`, `"n": 2`, `"response_format": "url"`} { + if strings.Contains(output, stale) { + t.Fatalf("dry-run output contains configured value despite flag override %s:\n%s", stale, output) + } + } +} + // resetEditCmdFlags restores edit subcommand flags to their defaults so that // values set by earlier tests (e.g. --dry-run, --extend, --provider) do not // leak in when tests run in shuffled order. diff --git a/internal/cli/errors.go b/internal/cli/errors.go index 18368de..4d447b7 100644 --- a/internal/cli/errors.go +++ b/internal/cli/errors.go @@ -41,11 +41,6 @@ func configError(err error) error { return &UserError{Category: CatConfig, Message: err.Error(), Raw: err} } -// apiError wraps an error so Execute() exits with the API exit code (3). -func apiError(err error) error { - return &UserError{Category: CatAPI, Message: err.Error(), Raw: err} -} - // imageError wraps an error so Execute() exits with the image exit code (4). func imageError(err error) error { return &UserError{Category: CatImage, Message: err.Error(), Raw: err} diff --git a/internal/cli/fallback.go b/internal/cli/fallback.go new file mode 100644 index 0000000..5e1fa52 --- /dev/null +++ b/internal/cli/fallback.go @@ -0,0 +1,264 @@ +package cli + +import ( + "context" + "errors" + "fmt" + "net" + "os" + "regexp" + "sort" + "strings" + + "github.com/ncxton/potaco/internal/adapter" + "github.com/ncxton/potaco/internal/auth" + "github.com/ncxton/potaco/internal/config" + "github.com/spf13/cobra" +) + +var httpStatusPattern = regexp.MustCompile(`HTTP ([0-9]{3})`) + +func resolveFallbackProviders(cmd *cobra.Command, mgr *auth.AuthManager, cfg *config.MultiProviderConfig, primary *resolvedConfig) []*resolvedConfig { + if !fallbackAllowed(cmd) || cfg == nil || primary == nil || primary.Model == "" { + return nil + } + + names := make([]string, 0, len(cfg.Providers)) + for name, pc := range cfg.Providers { + if name == primary.ProviderName || !similarModelID(primary.Model, pc.Model) { + continue + } + names = append(names, name) + } + sort.Strings(names) + + fallbacks := make([]*resolvedConfig, 0, len(names)) + for _, name := range names { + resolved, err := resolveFallbackProvider(cmd, mgr, cfg, name) + if err == nil { + fallbacks = append(fallbacks, resolved) + } + } + return fallbacks +} + +func fallbackAllowed(cmd *cobra.Command) bool { + for _, name := range []string{"provider", "base-url", "api-key"} { + if cmd.Flags().Changed(name) { + return false + } + } + for _, name := range []string{"POTACO_PROVIDER", "POTACO_BASE_URL", "POTACO_API_KEY"} { + if os.Getenv(name) != "" { + return false + } + } + return true +} + +func resolveFallbackProvider(cmd *cobra.Command, mgr *auth.AuthManager, cfg *config.MultiProviderConfig, providerName string) (*resolvedConfig, error) { + pc, ok := cfg.Providers[providerName] + if !ok || pc.Model == "" { + return nil, fmt.Errorf("provider %q has no model", providerName) + } + apiKey, err := mgr.GetAPIKey(providerName) + if err != nil { + return nil, err + } + providerType := config.ResolveProviderType(providerName, pc) + baseURL := resolveBaseURL(cmd, providerName, cfg) + baseURL, err = validateRuntimeBaseURL(baseURL) + if err != nil { + return nil, err + } + if config.ProviderRequiresBaseURL(providerName, providerType) && baseURL == "" { + return nil, fmt.Errorf("provider %q requires base URL", providerName) + } + retries, timeout, err := resolveRetriesTimeout(cmd, cfg, providerName) + if err != nil { + return nil, err + } + opts := adapter.AdapterOpts{BaseURL: baseURL, Retries: retries} + if timeout > 0 { + opts.Timeout = timeout + } + ad, err := adapter.Get(config.AdapterTypeForProvider(providerName, providerType), apiKey, opts) + if err != nil { + return nil, err + } + return &resolvedConfig{ + Adapter: ad, + ProviderName: providerName, + Model: pc.Model, + BaseURL: baseURL, + ModelCanEdit: modelEditCapable(cfg, providerName, pc.Model), + Parameters: modelParameters(cfg, providerName, pc.Model), + }, nil +} + +func similarModelID(a, b string) bool { + a = normalizeModelID(a) + b = normalizeModelID(b) + if a == "" || b == "" { + return false + } + if a == b { + return true + } + return modelScoped(a) != modelScoped(b) && modelLeaf(a) == modelLeaf(b) +} + +func normalizeModelID(model string) string { + return strings.Trim(strings.ToLower(model), " /") +} + +func modelLeaf(model string) string { + parts := strings.Split(model, "/") + return parts[len(parts)-1] +} + +func modelScoped(model string) bool { + return strings.Contains(model, "/") +} + +func providerAttempts(primary *resolvedConfig) []*resolvedConfig { + if primary == nil { + return nil + } + attempts := make([]*resolvedConfig, 0, 1+len(primary.Fallbacks)) + attempts = append(attempts, primary) + attempts = append(attempts, primary.Fallbacks...) + return attempts +} + +func transientProviderError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + return true + } + var netErr net.Error + if errors.As(err, &netErr) { + return true + } + msg := strings.ToLower(err.Error()) + for _, token := range []string{"timeout", "temporary", "connection refused", "connection reset", "network is unreachable", "no such host", "http request:"} { + if strings.Contains(msg, token) { + return true + } + } + matches := httpStatusPattern.FindStringSubmatch(err.Error()) + if len(matches) == 2 { + switch matches[1] { + case "408", "409", "425", "429", "500", "502", "503", "504": + return true + } + if strings.HasPrefix(matches[1], "5") { + return true + } + } + return false +} + +func warnProviderFallback(cmd *cobra.Command, from, to string) { + renderWarningf(cmd.ErrOrStderr(), "Provider '%s' failed; trying '%s'.", displayProviderName(from), displayProviderName(to)) +} + +type generateResult struct { + response *adapter.GenerateResponse + provider *resolvedConfig + request adapter.GenerateRequest +} + +func generateWithFallback(cmd *cobra.Command, primary *resolvedConfig, prompt string) (generateResult, error) { + attempts := providerAttempts(primary) + var lastErr error + for i, attempt := range attempts { + params := generationParametersFromFlags(cmd, attempt.Parameters) + warnUnsupportedGenerateParameters(cmd, attempt.Adapter.Name(), params) + req := adapter.GenerateRequest{ + Prompt: prompt, + Model: attempt.Model, + Size: params.Size, + Quality: params.Quality, + N: params.N, + ResponseFormat: params.ResponseFormat, + Seed: params.Seed, + GuidanceScale: params.GuidanceScale, + NegativePrompt: params.NegativePrompt, + } + + sp := startSpinner(cmd, fmt.Sprintf("Generating image with %s...", displayProviderName(attempt.ProviderName))) + resp, err := attempt.Adapter.Generate(context.Background(), req) + sp.stop() + if err == nil { + return generateResult{response: resp, provider: attempt, request: req}, nil + } + lastErr = err + if i == len(attempts)-1 || !transientProviderError(err) { + return generateResult{}, err + } + warnProviderFallback(cmd, attempt.ProviderName, attempts[i+1].ProviderName) + } + return generateResult{}, lastErr +} + +type editInput struct { + prompt string + imagePath string + maskPath string + extendedSize string +} + +type editResult struct { + response *adapter.GenerateResponse + provider *resolvedConfig + request adapter.EditRequest +} + +func editWithFallback(cmd *cobra.Command, primary *resolvedConfig, input editInput) (editResult, error) { + attempts := editAttempts(primary) + var lastErr error + for i, attempt := range attempts { + params := editParametersFromFlags(cmd, attempt.Parameters) + warnUnsupportedEditParameters(cmd, attempt.Adapter.Name(), params) + size := params.Size + if input.extendedSize != "" { + size = input.extendedSize + } + req := adapter.EditRequest{ + Prompt: input.prompt, + Model: attempt.Model, + N: params.N, + Size: size, + ResponseFormat: params.ResponseFormat, + ImagePath: input.imagePath, + MaskPath: input.maskPath, + } + + sp := startSpinner(cmd, fmt.Sprintf("Editing image with %s...", displayProviderName(attempt.ProviderName))) + resp, err := attempt.Adapter.Edit(context.Background(), req) + sp.stop() + if err == nil { + return editResult{response: resp, provider: attempt, request: req}, nil + } + lastErr = err + if i == len(attempts)-1 || !transientProviderError(err) { + return editResult{}, err + } + warnProviderFallback(cmd, attempt.ProviderName, attempts[i+1].ProviderName) + } + return editResult{}, lastErr +} + +func editAttempts(primary *resolvedConfig) []*resolvedConfig { + attempts := providerAttempts(primary) + filtered := make([]*resolvedConfig, 0, len(attempts)) + for i, attempt := range attempts { + if i == 0 || attempt.Adapter.SupportsEdit() && attempt.ModelCanEdit { + filtered = append(filtered, attempt) + } + } + return filtered +} diff --git a/internal/cli/fallback_test.go b/internal/cli/fallback_test.go new file mode 100644 index 0000000..a75328b --- /dev/null +++ b/internal/cli/fallback_test.go @@ -0,0 +1,15 @@ +package cli + +import "testing" + +func TestSimilarModelIDRejectsDifferentScopedProviders(t *testing.T) { + if similarModelID("vendor-a/foo", "vendor-b/foo") { + t.Fatal("same leaf under different scoped providers should not match") + } +} + +func TestSimilarModelIDMatchesScopedAndUnscopedAlias(t *testing.T) { + if !similarModelID("foo", "vendor/foo") { + t.Fatal("scoped and unscoped aliases with the same leaf should match") + } +} diff --git a/internal/cli/gen.go b/internal/cli/gen.go index 35d5496..39c24bf 100644 --- a/internal/cli/gen.go +++ b/internal/cli/gen.go @@ -1,12 +1,12 @@ package cli import ( - "context" "fmt" "strings" "time" "github.com/ncxton/potaco/internal/adapter" + "github.com/ncxton/potaco/internal/stats" "github.com/spf13/cobra" ) @@ -21,20 +21,19 @@ func init() { _ = genCmd.MarkFlagRequired("prompt") genCmd.Flags().String("model", "", "model to use (e.g., gpt-image-2)") - genCmd.Flags().String("size", "1024x1024", "image dimensions (WxH)") - genCmd.Flags().String("quality", "auto", "image quality (low, medium, high, or auto)") - genCmd.Flags().Int("n", 1, "number of images to generate") - - genCmd.Flags().Int("seed", 0, "reproducibility seed") - genCmd.Flags().Float64("guidance-scale", 0, "guidance scale") - genCmd.Flags().String("negative-prompt", "", "negative prompt") - genCmd.Flags().String("response-format", "b64_json", "response format (url or b64_json)") + genCmd.Flags().String("size", "", "override configured image dimensions (WxH)") + genCmd.Flags().String("quality", "", "override configured image quality") + genCmd.Flags().Int("n", 0, "override configured number of images") + genCmd.Flags().Int("seed", 0, "override configured reproducibility seed") + genCmd.Flags().Float64("guidance-scale", 0, "override configured guidance scale") + genCmd.Flags().String("negative-prompt", "", "override configured negative prompt") + genCmd.Flags().String("response-format", "", "override configured response format (url or b64_json)") genCmd.Flags().StringP("output", "o", "", "output file path") genCmd.Flags().String("output-format", "png", "output format (png or jpeg)") genCmd.Flags().Bool("stdout", false, "pipe raw image bytes to stdout") - genCmd.Flags().String("provider", "", "provider preset (openai, fal, vercel)") + genCmd.Flags().String("provider", "", "provider preset (openai, fal, replicate)") genCmd.Flags().String("base-url", "", "override API base URL") genCmd.Flags().String("api-key", "", "override API key") genCmd.Flags().Int("retries", 0, "max retry attempts") @@ -45,7 +44,18 @@ func init() { rootCmd.AddCommand(genCmd) } -func runGen(cmd *cobra.Command, args []string) error { +func runGen(cmd *cobra.Command, args []string) (err error) { + start := time.Now() + event := stats.Event{Command: cmd.CommandPath(), Operation: "generate"} + defer func() { + event.DurationMS = time.Since(start).Milliseconds() + event.Success = err == nil + if err != nil { + event.Error = err.Error() + } + recordStatsEvent(event) + }() + prompt := flagString(cmd, "prompt") if prompt == "" { return configUserErr("A prompt is required.", "Use 'potaco gen --prompt \"your description\"'.", fmt.Errorf("prompt cannot be empty")) @@ -63,21 +73,26 @@ func runGen(cmd *cobra.Command, args []string) error { return err } model := resolved.Model + event.Provider = resolved.ProviderName + event.Model = model + event.DryRun = flagBool(cmd, "dry-run") + params := generationParametersFromFlags(cmd, resolved.Parameters) req := adapter.GenerateRequest{ Prompt: prompt, Model: model, - Size: flagString(cmd, "size"), - Quality: flagString(cmd, "quality"), - N: flagInt(cmd, "n"), - ResponseFormat: flagString(cmd, "response-format"), - Seed: flagInt(cmd, "seed"), - GuidanceScale: flagFloat64(cmd, "guidance-scale"), - NegativePrompt: flagString(cmd, "negative-prompt"), + Size: params.Size, + Quality: params.Quality, + N: params.N, + ResponseFormat: params.ResponseFormat, + Seed: params.Seed, + GuidanceScale: params.GuidanceScale, + NegativePrompt: params.NegativePrompt, } dryRun := flagBool(cmd, "dry-run") if dryRun { + warnUnsupportedGenerateParameters(cmd, resolved.Adapter.Name(), params) var dryRunURL string if resolved.Adapter.Name() == "fal" { if model != "" { @@ -85,6 +100,15 @@ func runGen(cmd *cobra.Command, args []string) error { } else { dryRunURL = resolved.BaseURL + "/" } + } else if resolved.Adapter.Name() == "replicate" { + if err := validateReplicateDryRunModel(model); err != nil { + return err + } + if model != "" { + dryRunURL = resolved.BaseURL + "/models/" + model + "/predictions" + } else { + dryRunURL = resolved.BaseURL + "/models///predictions" + } } else if strings.HasSuffix(resolved.BaseURL, "/v1") { dryRunURL = resolved.BaseURL + "/images/generations" } else { @@ -94,12 +118,9 @@ func runGen(cmd *cobra.Command, args []string) error { return printDryRun(cmd, "POST", dryRunURL, "application/json", authHeader, req) } - sp := startSpinner(cmd, "Generating image...") - start := time.Now() - ctx := context.Background() - resp, err := resolved.Adapter.Generate(ctx, req) - sp.stop() - elapsed := time.Since(start) + requestStart := time.Now() + result, err := generateWithFallback(cmd, resolved, prompt) + elapsed := time.Since(requestStart) if err != nil { return apiUserErr( "Image generation failed.", @@ -107,6 +128,12 @@ func runGen(cmd *cobra.Command, args []string) error { fmt.Errorf("generate: %w", err), ) } + resp := result.response + model = result.provider.Model + req = result.request + event.Provider = result.provider.ProviderName + event.Model = model + event.Outputs = len(resp.Data) if err := processAndOutput(cmd, outputContext{ resp: resp, @@ -123,3 +150,15 @@ func runGen(cmd *cobra.Command, args []string) error { } return nil } + +func validateReplicateDryRunModel(model string) error { + owner, name, ok := strings.Cut(model, "/") + if !ok || owner == "" || name == "" || strings.Contains(name, "/") { + return configUserErr( + "Replicate models must use owner/model format.", + "Set a model such as 'black-forest-labs/flux-schnell'.", + fmt.Errorf("replicate model must use owner/model format, got %q", model), + ) + } + return nil +} diff --git a/internal/cli/gen_test.go b/internal/cli/gen_test.go index 59123ac..18875cc 100644 --- a/internal/cli/gen_test.go +++ b/internal/cli/gen_test.go @@ -114,6 +114,131 @@ func TestGenDryRunBodyUsesLowercaseJSONKeys(t *testing.T) { } } +func TestGenDryRunUsesConfiguredModelParameters(t *testing.T) { + setupAuthProvider(t, "sk-test") + resetGenCmdFlags(t) + t.Cleanup(func() { resetGenCmdFlags(t) }) + + for _, args := range [][]string{ + {"config", "set", "model.parameters.size", "1536x1024"}, + {"config", "set", "model.parameters.quality", "high"}, + {"config", "set", "model.parameters.n", "2"}, + {"config", "set", "model.parameters.response_format", "url"}, + {"config", "set", "model.parameters.seed", "42"}, + {"config", "set", "model.parameters.guidance_scale", "7.5"}, + {"config", "set", "model.parameters.negative_prompt", "blurry"}, + } { + rootCmd.SetArgs(args) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("%v: %v", args, err) + } + } + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"gen", "--prompt", "a cat", "--dry-run"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("gen --dry-run: %v", err) + } + + output := buf.String() + for _, want := range []string{`"size": "1536x1024"`, `"quality": "high"`, `"n": 2`, `"response_format": "url"`, `"seed": 42`, `"guidance_scale": 7.5`, `"negative_prompt": "blurry"`} { + if !strings.Contains(output, want) { + t.Fatalf("dry-run output missing %s:\n%s", want, output) + } + } +} + +func TestGenRejectsUntrustedHTTPBaseURLBeforeRequest(t *testing.T) { + setupAuthProvider(t, "sk-test") + resetGenCmdFlags(t) + t.Cleanup(func() { resetGenCmdFlags(t) }) + + calls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + mgr, err := auth.New() + if err != nil { + t.Fatalf("auth.New: %v", err) + } + if err := mgr.SetBaseURL("openai", srv.URL); err != nil { + t.Fatalf("set base URL: %v", err) + } + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"gen", "--prompt", "a cat"}) + + err = rootCmd.Execute() + if err == nil { + t.Fatal("gen should reject untrusted HTTP base URL") + } + if !strings.Contains(err.Error(), "insecure base URL") { + t.Fatalf("error = %v, want insecure base URL", err) + } + if calls != 0 { + t.Fatalf("server calls = %d, want 0", calls) + } +} + +func TestGenDryRunFlagsOverrideConfiguredModelParameters(t *testing.T) { + setupAuthProvider(t, "sk-test") + resetGenCmdFlags(t) + t.Cleanup(func() { resetGenCmdFlags(t) }) + + for _, args := range [][]string{ + {"config", "set", "model.parameters.size", "1536x1024"}, + {"config", "set", "model.parameters.quality", "high"}, + {"config", "set", "model.parameters.n", "2"}, + {"config", "set", "model.parameters.response_format", "url"}, + {"config", "set", "model.parameters.seed", "42"}, + {"config", "set", "model.parameters.guidance_scale", "7.5"}, + {"config", "set", "model.parameters.negative_prompt", "blurry"}, + } { + rootCmd.SetArgs(args) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("%v: %v", args, err) + } + } + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{ + "gen", + "--prompt", "a cat", + "--size", "1024x1792", + "--quality", "medium", + "--n", "1", + "--response-format", "b64_json", + "--seed", "99", + "--guidance-scale", "3.25", + "--negative-prompt", "rain", + "--dry-run", + }) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("gen --dry-run: %v", err) + } + + output := buf.String() + for _, want := range []string{`"size": "1024x1792"`, `"quality": "medium"`, `"n": 1`, `"response_format": "b64_json"`, `"seed": 99`, `"guidance_scale": 3.25`, `"negative_prompt": "rain"`} { + if !strings.Contains(output, want) { + t.Fatalf("dry-run output missing override %s:\n%s", want, output) + } + } + for _, stale := range []string{`"size": "1536x1024"`, `"quality": "high"`, `"n": 2`, `"response_format": "url"`, `"seed": 42`, `"guidance_scale": 7.5`, `"negative_prompt": "blurry"`} { + if strings.Contains(output, stale) { + t.Fatalf("dry-run output contains configured value despite flag override %s:\n%s", stale, output) + } + } +} + func TestGenCommandMissingConfigError(t *testing.T) { resetRootCmdFlags(t) resetAuthAddFlags(t) @@ -172,7 +297,8 @@ func resetGenCmdFlags(t *testing.T) { t.Helper() flags := genCmd.Flags() for _, name := range []string{"dry-run", "prompt", "model", "output", "stdout", - "provider", "base-url", "api-key", "retries", "timeout"} { + "provider", "base-url", "api-key", "retries", "timeout", + "size", "quality", "n", "seed", "guidance-scale", "negative-prompt", "response-format"} { if fl := flags.Lookup(name); fl != nil { if err := flags.Set(name, fl.DefValue); err != nil { t.Fatalf("reset gen flag %s: %v", name, err) @@ -269,48 +395,62 @@ func TestGenDryRunFalProvider(t *testing.T) { } } -func TestGenDryRunVercelProvider(t *testing.T) { - setupAuthProviderForProvider(t, "vercel", "vkey", "openai/gpt-image-2") +func TestGenDryRunReplicateProvider(t *testing.T) { + setupAuthProviderForProvider(t, "replicate", "r8_key", "black-forest-labs/flux-schnell") resetGenCmdFlags(t) t.Cleanup(func() { resetGenCmdFlags(t) }) var buf bytes.Buffer rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) - rootCmd.SetArgs([]string{"gen", "--prompt", "test", "--dry-run", "--provider", "vercel", "--base-url", "https://ai-gateway.vercel.sh/v1"}) + rootCmd.SetArgs([]string{"gen", "--prompt", "test", "--dry-run", "--provider", "replicate", "--base-url", "https://api.replicate.com/v1"}) if err := rootCmd.Execute(); err != nil { t.Fatalf("Execute: %v", err) } output := buf.String() if !strings.Contains(output, "Bearer [REDACTED]") { - t.Errorf("dry-run should show 'Bearer [REDACTED]' for vercel, got: %s", output) + t.Errorf("dry-run should show 'Bearer [REDACTED]' for replicate, got: %s", output) } - if !strings.Contains(output, "images/generations") { - t.Errorf("dry-run should show images/generations URL, got: %s", output) + if !strings.Contains(output, "models/black-forest-labs/flux-schnell/predictions") { + t.Errorf("dry-run should show Replicate predictions URL, got: %s", output) } } -// TestGenDryRunVercelProviderPresetNoBaseURL verifies that the Vercel provider -// preset includes /v1 in the base URL so that the generate endpoint URL is -// constructed correctly even without an explicit --base-url override. -// This is a regression test for a 404 bug where the preset was missing /v1. -func TestGenDryRunVercelProviderPresetNoBaseURL(t *testing.T) { - setupAuthProviderForProvider(t, "vercel", "vkey", "openai/gpt-image-2") +func TestGenDryRunReplicateProviderPresetNoBaseURL(t *testing.T) { + setupAuthProviderForProvider(t, "replicate", "r8_key", "black-forest-labs/flux-schnell") resetGenCmdFlags(t) t.Cleanup(func() { resetGenCmdFlags(t) }) var buf bytes.Buffer rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) - rootCmd.SetArgs([]string{"gen", "--prompt", "test", "--dry-run", "--provider", "vercel"}) + rootCmd.SetArgs([]string{"gen", "--prompt", "test", "--dry-run", "--provider", "replicate"}) if err := rootCmd.Execute(); err != nil { t.Fatalf("Execute: %v", err) } output := buf.String() - if !strings.Contains(output, "ai-gateway.vercel.sh/v1/images/generations") { - t.Errorf("dry-run should show full URL with /v1/images/generations from preset (no --base-url), got: %s", output) + if !strings.Contains(output, "api.replicate.com/v1/models/black-forest-labs/flux-schnell/predictions") { + t.Errorf("dry-run should show full URL with Replicate preset, got: %s", output) + } +} + +func TestGenDryRunReplicateRejectsInvalidModelID(t *testing.T) { + setupAuthProviderForProvider(t, "replicate", "r8_key", "flux-schnell") + resetGenCmdFlags(t) + t.Cleanup(func() { resetGenCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"gen", "--prompt", "test", "--dry-run", "--provider", "replicate"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected invalid model ID error") + } + if !strings.Contains(err.Error(), "owner/model") { + t.Fatalf("error = %v, want owner/model", err) } } @@ -388,6 +528,7 @@ func TestGenWithApiKeyOverride(t *testing.T) { // path (not the legacy provider client) to call the API and save output. // A mock server mimics an OpenAI-compatible generations endpoint. func TestGenCommandAdapterEndToEnd(t *testing.T) { + trustHTTPBaseURLs(t) resetRootCmdFlags(t) resetAuthAddFlags(t) resetGenCmdFlags(t) @@ -456,6 +597,165 @@ func TestGenCommandAdapterEndToEnd(t *testing.T) { } } +func TestGenFallsBackToSimilarModelProviderOnTransientError(t *testing.T) { + trustHTTPBaseURLs(t) + resetRootCmdFlags(t) + resetAuthAddFlags(t) + resetGenCmdFlags(t) + t.Cleanup(func() { resetGenCmdFlags(t) }) + tmpHome := t.TempDir() + t.Setenv("HOME", tmpHome) + t.Setenv("POTACO_API_KEY", "") + t.Setenv("POTACO_BASE_URL", "") + t.Setenv("POTACO_PROVIDER", "") + t.Setenv("POTACO_MODEL", "") + + openaiCalls := 0 + openaiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + openaiCalls++ + http.Error(w, `{"error":{"message":"rate limited"}}`, http.StatusTooManyRequests) + })) + defer openaiServer.Close() + + tile := image.NewRGBA(image.Rect(0, 0, 2, 2)) + var tileBuf bytes.Buffer + if err := png.Encode(&tileBuf, tile); err != nil { + t.Fatalf("encode png: %v", err) + } + + replicateCalls := 0 + replicateServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1/models/owner/gpt-image-2/predictions": + replicateCalls++ + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(map[string]any{ + "output": "http://" + r.Host + "/result.png", + }); err != nil { + t.Fatalf("write prediction: %v", err) + } + case "/result.png": + w.Header().Set("Content-Type", "image/png") + if _, err := w.Write(tileBuf.Bytes()); err != nil { + t.Fatalf("write image: %v", err) + } + default: + t.Fatalf("unexpected replicate path %q", r.URL.Path) + } + })) + defer replicateServer.Close() + + mgr, err := auth.New() + if err != nil { + t.Fatalf("auth.New: %v", err) + } + if err := mgr.AddProvider("openai", "openai", "sk-openai"); err != nil { + t.Fatalf("add openai: %v", err) + } + if err := mgr.SetBaseURL("openai", openaiServer.URL); err != nil { + t.Fatalf("set openai base URL: %v", err) + } + if err := mgr.SetActiveProvider("openai", "gpt-image-2"); err != nil { + t.Fatalf("set openai model: %v", err) + } + if err := mgr.AddProvider("replicate", "replicate", "r8-replicate"); err != nil { + t.Fatalf("add replicate: %v", err) + } + if err := mgr.SetBaseURL("replicate", replicateServer.URL+"/v1"); err != nil { + t.Fatalf("set replicate base URL: %v", err) + } + if err := mgr.SetActiveProvider("replicate", "owner/gpt-image-2"); err != nil { + t.Fatalf("set replicate model: %v", err) + } + if err := mgr.SetActiveProvider("openai", "gpt-image-2"); err != nil { + t.Fatalf("restore active provider: %v", err) + } + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"gen", "--prompt", "fallback", "--retries", "-1"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("gen fallback: %v\n%s", err, buf.String()) + } + + if openaiCalls != 1 { + t.Fatalf("openai calls = %d, want 1", openaiCalls) + } + if replicateCalls != 1 { + t.Fatalf("replicate calls = %d, want 1", replicateCalls) + } + if !strings.Contains(buf.String(), "replicate") { + t.Fatalf("output should mention fallback provider, got:\n%s", buf.String()) + } + if !strings.Contains(buf.String(), "result.png") { + t.Fatalf("output should include fallback result URL, got:\n%s", buf.String()) + } +} + +func TestGenExplicitProviderDisablesFallback(t *testing.T) { + resetRootCmdFlags(t) + resetAuthAddFlags(t) + resetGenCmdFlags(t) + t.Cleanup(func() { resetGenCmdFlags(t) }) + tmpHome := t.TempDir() + t.Setenv("HOME", tmpHome) + t.Setenv("POTACO_API_KEY", "") + t.Setenv("POTACO_BASE_URL", "") + t.Setenv("POTACO_PROVIDER", "") + t.Setenv("POTACO_MODEL", "") + + openaiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"error":{"message":"rate limited"}}`, http.StatusTooManyRequests) + })) + defer openaiServer.Close() + + replicateCalls := 0 + replicateServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + replicateCalls++ + w.WriteHeader(http.StatusOK) + })) + defer replicateServer.Close() + + mgr, err := auth.New() + if err != nil { + t.Fatalf("auth.New: %v", err) + } + if err := mgr.AddProvider("openai", "openai", "sk-openai"); err != nil { + t.Fatalf("add openai: %v", err) + } + if err := mgr.SetBaseURL("openai", openaiServer.URL); err != nil { + t.Fatalf("set openai base URL: %v", err) + } + if err := mgr.SetActiveProvider("openai", "gpt-image-2"); err != nil { + t.Fatalf("set openai model: %v", err) + } + if err := mgr.AddProvider("replicate", "replicate", "r8-replicate"); err != nil { + t.Fatalf("add replicate: %v", err) + } + if err := mgr.SetBaseURL("replicate", replicateServer.URL+"/v1"); err != nil { + t.Fatalf("set replicate base URL: %v", err) + } + if err := mgr.SetActiveProvider("replicate", "owner/gpt-image-2"); err != nil { + t.Fatalf("set replicate model: %v", err) + } + if err := mgr.SetActiveProvider("openai", "gpt-image-2"); err != nil { + t.Fatalf("restore active provider: %v", err) + } + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"gen", "--prompt", "fallback", "--provider", "openai", "--retries", "-1"}) + err = rootCmd.Execute() + if err == nil { + t.Fatal("gen should fail without fallback when --provider is explicit") + } + if replicateCalls != 0 { + t.Fatalf("replicate calls = %d, want 0", replicateCalls) + } +} + func TestGenUsesConfigBaseURL(t *testing.T) { resetRootCmdFlags(t) resetAuthAddFlags(t) diff --git a/internal/cli/helpers.go b/internal/cli/helpers.go index 2d158e4..bc9178f 100644 --- a/internal/cli/helpers.go +++ b/internal/cli/helpers.go @@ -11,10 +11,10 @@ import ( "time" "github.com/ncxton/potaco/internal/adapter" - _ "github.com/ncxton/potaco/internal/adapter/custom" // register custom adapter - _ "github.com/ncxton/potaco/internal/adapter/fal" // register fal adapter - _ "github.com/ncxton/potaco/internal/adapter/openai" // register openai adapter - _ "github.com/ncxton/potaco/internal/adapter/vercel" // register vercel adapter + _ "github.com/ncxton/potaco/internal/adapter/custom" // register custom adapter + _ "github.com/ncxton/potaco/internal/adapter/fal" // register fal adapter + _ "github.com/ncxton/potaco/internal/adapter/openai" // register openai adapter + _ "github.com/ncxton/potaco/internal/adapter/replicate" // register replicate adapter img "github.com/ncxton/potaco/internal/image" "github.com/spf13/cobra" ) @@ -25,9 +25,9 @@ type providerPreset struct { } var providerPresets = map[string]providerPreset{ - "openai": {BaseURL: "https://api.openai.com/v1"}, - "fal": {BaseURL: "https://fal.run"}, - "vercel": {BaseURL: "https://ai-gateway.vercel.sh/v1"}, + "openai": {BaseURL: "https://api.openai.com/v1"}, + "fal": {BaseURL: "https://fal.run"}, + "replicate": {BaseURL: "https://api.replicate.com/v1"}, // "custom" has no preset: users must supply a base URL. } @@ -36,6 +36,13 @@ func getProviderPreset(name string) (providerPreset, bool) { return p, ok } +func displayProviderName(name string) string { + if name == "fal" { + return "Fal" + } + return safeTerminalText(name) +} + func flagString(cmd *cobra.Command, name string) string { v, _ := cmd.Flags().GetString(name) return v diff --git a/internal/cli/helpers_test.go b/internal/cli/helpers_test.go index e2e9b69..36664bb 100644 --- a/internal/cli/helpers_test.go +++ b/internal/cli/helpers_test.go @@ -9,7 +9,7 @@ import ( func TestAllProvidersRegistered(t *testing.T) { names := adapter.List() - want := []string{"custom", "fal", "openai", "vercel"} + want := []string{"custom", "fal", "openai", "replicate"} if len(names) != len(want) { t.Fatalf("registered providers = %v, want %v", names, want) } @@ -27,13 +27,13 @@ func TestAllProvidersRegistered(t *testing.T) { } } -func TestVercelPresetExists(t *testing.T) { - preset, ok := getProviderPreset("vercel") +func TestReplicatePresetExists(t *testing.T) { + preset, ok := getProviderPreset("replicate") if !ok { - t.Fatal("vercel preset not found") + t.Fatal("replicate preset not found") } if preset.BaseURL == "" { - t.Error("vercel preset BaseURL should not be empty") + t.Error("replicate preset BaseURL should not be empty") } } @@ -55,9 +55,9 @@ func TestProviderPresetsNoCustomEntry(t *testing.T) { func TestProviderPresetsOnlyKnownProviders(t *testing.T) { want := map[string]bool{ - "openai": true, - "fal": true, - "vercel": true, + "openai": true, + "fal": true, + "replicate": true, } for name := range providerPresets { if !want[name] { diff --git a/internal/cli/info.go b/internal/cli/info.go deleted file mode 100644 index c813d1d..0000000 --- a/internal/cli/info.go +++ /dev/null @@ -1,67 +0,0 @@ -package cli - -import ( - "encoding/json" - "fmt" - "os" - - img "github.com/ncxton/potaco/internal/image" - "github.com/spf13/cobra" -) - -var infoCmd = &cobra.Command{ - Use: "info [path]", - Short: "Print metadata about an image file", - Args: cobra.ExactArgs(1), - RunE: runInfo, -} - -func init() { - rootCmd.AddCommand(infoCmd) -} - -func runInfo(cmd *cobra.Command, args []string) error { - path := args[0] - - stat, err := os.Stat(path) - if err != nil { - return imageError(fmt.Errorf("file: %w", err)) - } - - image, format, err := img.ReadImage(path) - if err != nil { - return imageError(fmt.Errorf("read image: %w", err)) - } - - bounds := image.Bounds() - width := bounds.Dx() - height := bounds.Dy() - fileSize := stat.Size() - colorModel := fmt.Sprintf("%v", image.ColorModel()) - - jsonMode, _ := cmd.Root().PersistentFlags().GetBool("json") - - if jsonMode { - output := map[string]any{ - "path": path, - "format": format, - "width": width, - "height": height, - "file_size": fileSize, - "color_model": colorModel, - } - b, err := json.MarshalIndent(output, "", " ") - if err != nil { - return fmt.Errorf("marshal JSON: %w", err) - } - fmt.Fprintln(cmd.OutOrStdout(), string(b)) - } else { - fmt.Fprintf(cmd.OutOrStdout(), "File: %s\n", path) - fmt.Fprintf(cmd.OutOrStdout(), "Format: %s\n", format) - fmt.Fprintf(cmd.OutOrStdout(), "Dimensions: %dx%d\n", width, height) - fmt.Fprintf(cmd.OutOrStdout(), "File size: %d bytes\n", fileSize) - fmt.Fprintf(cmd.OutOrStdout(), "Color: %s\n", colorModel) - } - - return nil -} diff --git a/internal/cli/info_test.go b/internal/cli/info_test.go deleted file mode 100644 index 710d9b0..0000000 --- a/internal/cli/info_test.go +++ /dev/null @@ -1,120 +0,0 @@ -package cli - -import ( - "bytes" - "encoding/json" - "image" - "image/color" - "image/png" - "os" - "path/filepath" - "strings" - "testing" -) - -func TestInfoCommandExists(t *testing.T) { - found := false - for _, cmd := range rootCmd.Commands() { - if cmd.Use == "info" || strings.HasPrefix(cmd.Use, "info ") { - found = true - break - } - } - if !found { - t.Fatal("root command should have 'info' subcommand") - } -} - -func TestInfoCommandOutput(t *testing.T) { - resetRootCmdFlags(t) - dir := t.TempDir() - imgPath := filepath.Join(dir, "test.png") - createInfoTestPNG(t, imgPath, 100, 200) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetErr(&buf) - rootCmd.SetArgs([]string{"info", imgPath}) - - err := rootCmd.Execute() - if err != nil { - t.Fatalf("info command error: %v", err) - } - - output := buf.String() - if !strings.Contains(output, "png") { - t.Errorf("output should mention format 'png', got: %q", output) - } - if !strings.Contains(output, "100x200") { - t.Errorf("output should contain dimensions 100x200, got: %q", output) - } - if !strings.Contains(output, imgPath) { - t.Errorf("output should contain file path, got: %q", output) - } -} - -func TestInfoCommandJSONOutput(t *testing.T) { - resetRootCmdFlags(t) - dir := t.TempDir() - imgPath := filepath.Join(dir, "test.png") - createInfoTestPNG(t, imgPath, 64, 32) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetErr(&buf) - rootCmd.SetArgs([]string{"info", "--json", imgPath}) - - err := rootCmd.Execute() - if err != nil { - t.Fatalf("info --json command error: %v", err) - } - - output := buf.String() - var parsed map[string]any - if err := json.Unmarshal([]byte(output), &parsed); err != nil { - t.Fatalf("info --json output should be valid JSON, got parse error: %v, output: %q", err, output) - } - if parsed["format"] != "png" { - t.Errorf("JSON output format should be 'png', got: %v", parsed["format"]) - } - if parsed["width"] != float64(64) { - t.Errorf("JSON output width should be 64, got: %v", parsed["width"]) - } - if parsed["height"] != float64(32) { - t.Errorf("JSON output height should be 32, got: %v", parsed["height"]) - } - if parsed["path"] != imgPath { - t.Errorf("JSON output path should be %q, got: %v", imgPath, parsed["path"]) - } -} - -func TestInfoCommandMissingFile(t *testing.T) { - resetRootCmdFlags(t) - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetErr(&buf) - rootCmd.SetArgs([]string{"info", "/nonexistent/file.png"}) - - err := rootCmd.Execute() - if err == nil { - t.Fatal("info should error on missing file") - } -} - -func createInfoTestPNG(t *testing.T, path string, w, h int) { - t.Helper() - f, err := os.Create(path) - if err != nil { - t.Fatalf("create file: %v", err) - } - defer f.Close() - img := image.NewRGBA(image.Rect(0, 0, w, h)) - for y := 0; y < h; y++ { - for x := 0; x < w; x++ { - img.Set(x, y, color.RGBA{R: 255, G: 0, B: 0, A: 255}) - } - } - if err := png.Encode(f, img); err != nil { - t.Fatalf("encode png: %v", err) - } -} diff --git a/internal/cli/model_parameters.go b/internal/cli/model_parameters.go new file mode 100644 index 0000000..a9534eb --- /dev/null +++ b/internal/cli/model_parameters.go @@ -0,0 +1,80 @@ +package cli + +import ( + "github.com/ncxton/potaco/internal/config" + "github.com/spf13/cobra" +) + +func generationParametersFromFlags(cmd *cobra.Command, params config.ModelParameters) config.ModelParameters { + if cmd.Flags().Changed("size") { + params.Size = flagString(cmd, "size") + } + if cmd.Flags().Changed("quality") { + params.Quality = flagString(cmd, "quality") + } + if cmd.Flags().Changed("n") { + params.N = flagInt(cmd, "n") + } + if cmd.Flags().Changed("response-format") { + params.ResponseFormat = flagString(cmd, "response-format") + } + if cmd.Flags().Changed("seed") { + params.Seed = flagInt(cmd, "seed") + } + if cmd.Flags().Changed("guidance-scale") { + params.GuidanceScale = flagFloat64(cmd, "guidance-scale") + } + if cmd.Flags().Changed("negative-prompt") { + params.NegativePrompt = flagString(cmd, "negative-prompt") + } + return params.WithDefaults() +} + +func editParametersFromFlags(cmd *cobra.Command, params config.ModelParameters) config.ModelParameters { + if cmd.Flags().Changed("size") { + params.Size = flagString(cmd, "size") + } + if cmd.Flags().Changed("n") { + params.N = flagInt(cmd, "n") + } + if cmd.Flags().Changed("response-format") { + params.ResponseFormat = flagString(cmd, "response-format") + } + return params.WithDefaults() +} + +func warnUnsupportedGenerateParameters(cmd *cobra.Command, providerName string, params config.ModelParameters) { + switch providerName { + case "openai", "custom": + return + } + defaults := config.DefaultModelParameters() + if cmd.Flags().Changed("quality") || params.Quality != defaults.Quality { + renderWarning(cmd.ErrOrStderr(), "Parameter quality is configured but not applied to generation requests for this provider.") + } + if cmd.Flags().Changed("response-format") || params.ResponseFormat != defaults.ResponseFormat { + renderWarning(cmd.ErrOrStderr(), "Parameter response_format is configured but not applied to generation requests for this provider.") + } +} + +func warnUnsupportedEditParameters(cmd *cobra.Command, providerName string, params config.ModelParameters) { + var ignored []string + if params.Quality != "" && params.Quality != config.DefaultModelParameters().Quality { + ignored = append(ignored, "quality") + } + if providerName == "fal" && params.ResponseFormat != "" && params.ResponseFormat != config.DefaultModelParameters().ResponseFormat { + ignored = append(ignored, "response_format") + } + if params.Seed != 0 { + ignored = append(ignored, "seed") + } + if params.GuidanceScale != 0 { + ignored = append(ignored, "guidance_scale") + } + if params.NegativePrompt != "" { + ignored = append(ignored, "negative_prompt") + } + for _, name := range ignored { + renderWarningf(cmd.ErrOrStderr(), "Parameter %s is configured but not applied to edit requests.", name) + } +} diff --git a/internal/cli/model_parameters_test.go b/internal/cli/model_parameters_test.go new file mode 100644 index 0000000..4334819 --- /dev/null +++ b/internal/cli/model_parameters_test.go @@ -0,0 +1,61 @@ +package cli + +import ( + "bytes" + "strings" + "testing" + + "github.com/ncxton/potaco/internal/config" + "github.com/spf13/cobra" +) + +func TestWarnUnsupportedGenerateParametersForNonOpenAIProvider(t *testing.T) { + cmd := &cobra.Command{} + cmd.Flags().String("quality", "", "") + cmd.Flags().String("response-format", "", "") + var buf bytes.Buffer + cmd.SetErr(&buf) + + warnUnsupportedGenerateParameters(cmd, "replicate", config.ModelParameters{ + Quality: "high", + ResponseFormat: "url", + }) + + output := buf.String() + for _, want := range []string{" Parameter quality", " Parameter response_format"} { + if !strings.Contains(output, want) { + t.Fatalf("warning output missing %q:\n%s", want, output) + } + } +} + +func TestWarnUnsupportedGenerateParametersAllowsCustomProvider(t *testing.T) { + cmd := &cobra.Command{} + cmd.Flags().String("quality", "", "") + cmd.Flags().String("response-format", "", "") + var buf bytes.Buffer + cmd.SetErr(&buf) + + warnUnsupportedGenerateParameters(cmd, "custom", config.ModelParameters{ + Quality: "high", + ResponseFormat: "url", + }) + + if buf.Len() != 0 { + t.Fatalf("custom provider applies generation parameters, got warning:\n%s", buf.String()) + } +} + +func TestWarnUnsupportedEditParametersForFalResponseFormat(t *testing.T) { + cmd := &cobra.Command{} + var buf bytes.Buffer + cmd.SetErr(&buf) + + warnUnsupportedEditParameters(cmd, "fal", config.ModelParameters{ + ResponseFormat: "url", + }) + + if !strings.Contains(buf.String(), " Parameter response_format") { + t.Fatalf("warning output missing response_format:\n%s", buf.String()) + } +} diff --git a/internal/cli/models_cmd.go b/internal/cli/models_cmd.go index 919fe5a..6de25f4 100644 --- a/internal/cli/models_cmd.go +++ b/internal/cli/models_cmd.go @@ -116,6 +116,10 @@ func resolveModelsProvider(cmd *cobra.Command, args []string) (resolvedModelsPro } providerType := config.ResolveProviderType(providerName, pc) baseURL := resolveBaseURL(cmd, providerName, cfg) + baseURL, err = validateDiscoveryBaseURL(baseURL) + if err != nil { + return resolvedModelsProvider{}, err + } if config.ProviderRequiresBaseURL(providerName, providerType) && baseURL == "" { return resolvedModelsProvider{}, configUserErr( "A base URL is required for this provider.", @@ -126,7 +130,7 @@ func resolveModelsProvider(cmd *cobra.Command, args []string) (resolvedModelsPro return resolvedModelsProvider{ ProviderName: providerName, - AdapterType: config.AdapterType(providerType), + AdapterType: config.AdapterTypeForProvider(providerName, providerType), APIKey: apiKey, BaseURL: baseURL, }, nil @@ -200,13 +204,20 @@ func printModelsText(out io.Writer, models []adapter.Model) error { fmt.Fprintln(out, "No models found.") return nil } - fmt.Fprintf(out, "%-40s %s\n", "MODEL ID", "DISPLAY NAME") + fmt.Fprintf(out, "%-40s %s\n", "DISPLAY NAME", "MODEL ID") for _, m := range models { - fmt.Fprintf(out, "%-40s %s\n", m.ID, m.DisplayName) + fmt.Fprintf(out, "%-40s %s\n", safeTerminalText(displayNameOrID(m)), safeTerminalText(m.ID)) } return nil } +func displayNameOrID(model adapter.Model) string { + if model.DisplayName != "" { + return model.DisplayName + } + return model.ID +} + func printModelsJSON(out io.Writer, models []adapter.Model) error { type modelJSON struct { ID string `json:"id"` diff --git a/internal/cli/models_cmd_errors_test.go b/internal/cli/models_cmd_errors_test.go new file mode 100644 index 0000000..d28b633 --- /dev/null +++ b/internal/cli/models_cmd_errors_test.go @@ -0,0 +1,200 @@ +package cli + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/spf13/cobra" +) + +func TestModelsListSpecificProviderNotConnected(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + resetRootCmdFlags(t) + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "list", "fal"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for unconnected provider") + } + if !strings.Contains(err.Error(), "not connected") { + t.Errorf("error should mention not connected, got: %v", err) + } +} + +func TestModelsUnknownProvider(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + resetRootCmdFlags(t) + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "unknown-provider"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for unknown provider") + } + if !strings.Contains(err.Error(), "unknown provider") { + t.Errorf("error should mention unknown provider, got: %v", err) + } +} + +func TestModelsListUnknownProvider(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + resetRootCmdFlags(t) + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "list", "unknown-provider"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for unknown provider") + } + if !strings.Contains(err.Error(), "unknown provider") { + t.Errorf("error should mention unknown provider, got: %v", err) + } +} + +func TestModelsDiscoveryError(t *testing.T) { + trustHTTPBaseURLs(t) + setupAuthProviderForProvider(t, "openai", "sk-test", "gpt-image-2") + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer srv.Close() + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "list", "--base-url", srv.URL}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error when discovery fails") + } + if !strings.Contains(err.Error(), "discover models") { + t.Errorf("error should mention discover models, got: %v", err) + } +} + +func TestModelsHelp(t *testing.T) { + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "--help"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + output := buf.String() + if !strings.Contains(output, "models [provider]") { + t.Errorf("help should show usage 'models [provider]', got: %s", output) + } + if !strings.Contains(output, "list") { + t.Errorf("help should list subcommand 'list', got: %s", output) + } + if strings.Contains(output, "--params") { + t.Errorf("help should not list --params flag, got: %s", output) + } +} + +func TestModelsListHelp(t *testing.T) { + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "list", "--help"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + output := buf.String() + if !strings.Contains(output, "list [provider]") { + t.Errorf("help should show usage 'list [provider]', got: %s", output) + } +} + +func TestModelsMaximumNArgs(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + resetRootCmdFlags(t) + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "openai", "extra"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for too many args") + } + if !strings.Contains(err.Error(), "accepts at most 1 arg") { + t.Errorf("error should mention arg limit, got: %v", err) + } +} + +func TestModelsListMaximumNArgs(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + resetRootCmdFlags(t) + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "list", "openai", "extra"}) + err := rootCmd.Execute() + if err == nil { + t.Fatal("expected error for too many args") + } + if !strings.Contains(err.Error(), "accepts at most 1 arg") { + t.Errorf("error should mention arg limit, got: %v", err) + } +} + +func resetModelsCmdFlags(t *testing.T) { + t.Helper() + for _, name := range []string{"base-url", "api-key"} { + flag := modelsCmd.PersistentFlags().Lookup(name) + if flag == nil { + continue + } + if err := flag.Value.Set(flag.DefValue); err != nil { + t.Fatalf("reset %s flag: %v", name, err) + } + flag.Changed = false + } + for _, name := range []string{"base-url", "api-key"} { + flag := modelsListCmd.Flags().Lookup(name) + if flag == nil { + continue + } + if err := flag.Value.Set(flag.DefValue); err != nil { + t.Fatalf("reset %s child flag: %v", name, err) + } + flag.Changed = false + } + for _, cmd := range []*cobra.Command{modelsCmd, modelsListCmd} { + for _, name := range []string{"help", "version"} { + flag := cmd.Flags().Lookup(name) + if flag == nil { + continue + } + if err := flag.Value.Set("false"); err != nil { + t.Fatalf("reset %s flag on %s: %v", name, cmd.Name(), err) + } + flag.Changed = false + } + } +} diff --git a/internal/cli/models_cmd_provider_test.go b/internal/cli/models_cmd_provider_test.go new file mode 100644 index 0000000..938e0d8 --- /dev/null +++ b/internal/cli/models_cmd_provider_test.go @@ -0,0 +1,244 @@ +package cli + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/ncxton/potaco/internal/auth" + "github.com/ncxton/potaco/internal/config" +) + +func TestModelsSpecificProvider(t *testing.T) { + trustHTTPBaseURLs(t) + resetRootCmdFlags(t) + resetAuthAddFlags(t) + t.Setenv("HOME", t.TempDir()) + t.Setenv("POTACO_API_KEY", "") + t.Setenv("POTACO_BASE_URL", "") + t.Setenv("POTACO_PROVIDER", "") + t.Setenv("POTACO_MODEL", "") + + var setupBuf bytes.Buffer + rootCmd.SetOut(&setupBuf) + rootCmd.SetErr(&setupBuf) + + rootCmd.SetArgs([]string{"auth", "add", "openai", "--api-key", "sk-openai", "--force", "--model", "gpt-image-2"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("setup auth add openai: %v", err) + } + resetAuthAddFlags(t) + resetRootCmdFlags(t) + + rootCmd.SetArgs([]string{"auth", "add", "fal", "--api-key", "fal-key", "--force", "--model", "fal-ai/flux/dev"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("setup auth add fal: %v", err) + } + resetAuthAddFlags(t) + resetRootCmdFlags(t) + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Errorf("path = %q, want /v1/models", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "gpt-image-2", "owned_by": "openai"}, + }, + }) + })) + defer srv.Close() + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "openai", "--base-url", srv.URL}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + output := buf.String() + if !strings.Contains(output, "gpt-image-2") { + t.Errorf("models openai should list openai models, got: %s", output) + } +} + +func TestModelsListSpecificProvider(t *testing.T) { + trustHTTPBaseURLs(t) + resetRootCmdFlags(t) + resetAuthAddFlags(t) + t.Setenv("HOME", t.TempDir()) + t.Setenv("POTACO_API_KEY", "") + t.Setenv("POTACO_BASE_URL", "") + t.Setenv("POTACO_PROVIDER", "") + t.Setenv("POTACO_MODEL", "") + + var setupBuf bytes.Buffer + rootCmd.SetOut(&setupBuf) + rootCmd.SetErr(&setupBuf) + + rootCmd.SetArgs([]string{"auth", "add", "openai", "--api-key", "sk-openai", "--force", "--model", "gpt-image-2"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("setup auth add openai: %v", err) + } + resetAuthAddFlags(t) + resetRootCmdFlags(t) + + rootCmd.SetArgs([]string{"auth", "add", "fal", "--api-key", "fal-key", "--force", "--model", "fal-ai/flux/dev"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("setup auth add fal: %v", err) + } + resetAuthAddFlags(t) + resetRootCmdFlags(t) + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Errorf("path = %q, want /v1/models", r.URL.Path) + } + authHeader := r.Header.Get("Authorization") + if authHeader != "Bearer sk-openai" { + t.Errorf("Authorization = %q, want Bearer sk-openai", authHeader) + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "gpt-image-2", "owned_by": "openai"}, + }, + }) + })) + defer srv.Close() + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "list", "openai", "--base-url", srv.URL}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + output := buf.String() + if !strings.Contains(output, "gpt-image-2") { + t.Errorf("models list openai should list openai models, got: %s", output) + } +} + +func TestModelsListWithApiKeyOverride(t *testing.T) { + trustHTTPBaseURLs(t) + setupAuthProviderForProvider(t, "openai", "sk-stored", "gpt-image-2") + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + var gotAuth string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "gpt-image-2", "owned_by": "openai"}, + }, + }) + })) + defer srv.Close() + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "list", "--api-key", "sk-override", "--base-url", srv.URL}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + if gotAuth != "Bearer sk-override" { + t.Errorf("Authorization header = %q, want Bearer sk-override", gotAuth) + } +} + +func TestModelsListUsesBaseURLFromConfig(t *testing.T) { + trustHTTPBaseURLs(t) + t.Setenv("HOME", t.TempDir()) + resetRootCmdFlags(t) + resetAuthAddFlags(t) + + mgr, err := auth.New() + if err != nil { + t.Fatalf("init auth: %v", err) + } + if err := mgr.Add("openai", "sk-test"); err != nil { + t.Fatalf("add provider: %v", err) + } + if err := mgr.SetBaseURL("openai", "https://config.example.com/v1"); err != nil { + t.Fatalf("set base url: %v", err) + } + + var gotURL string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotURL = r.URL.Path + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "gpt-image-2", "owned_by": "openai"}, + }, + }) + })) + defer srv.Close() + + cfg, err := mgr.LoadConfig() + if err != nil { + t.Fatalf("load config: %v", err) + } + pc := cfg.Providers["openai"] + pc.BaseURL = srv.URL + cfg.Providers["openai"] = pc + if err := config.SaveMultiProvider(config.DefaultConfigPath(), cfg); err != nil { + t.Fatalf("save config: %v", err) + } + + resetRootCmdFlags(t) + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "list"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + if gotURL != "/v1/models" { + t.Errorf("request path = %q, want /v1/models", gotURL) + } + if !strings.Contains(buf.String(), "gpt-image-2") { + t.Errorf("output should list gpt-image-2, got: %s", buf.String()) + } +} + +func TestModelsListAliasWithBuiltInTypeRequiresBaseURL(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + resetRootCmdFlags(t) + resetModelsCmdFlags(t) + t.Cleanup(func() { resetModelsCmdFlags(t) }) + + mgr, err := auth.New() + if err != nil { + t.Fatalf("init auth: %v", err) + } + if err := mgr.AddProvider("openrouter", "openai", "sk-test"); err != nil { + t.Fatalf("add provider: %v", err) + } + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetArgs([]string{"models", "list", "openrouter"}) + err = rootCmd.Execute() + if err == nil { + t.Fatal("expected error for missing alias base URL") + } + if !strings.Contains(err.Error(), "base URL") { + t.Fatalf("error = %v, want base URL", err) + } +} diff --git a/internal/cli/models_cmd_test.go b/internal/cli/models_cmd_test.go index 07ff555..dd8e53c 100644 --- a/internal/cli/models_cmd_test.go +++ b/internal/cli/models_cmd_test.go @@ -7,13 +7,10 @@ import ( "net/http/httptest" "strings" "testing" - - "github.com/ncxton/potaco/internal/auth" - "github.com/ncxton/potaco/internal/config" - "github.com/spf13/cobra" ) func TestModelsListActiveProvider(t *testing.T) { + trustHTTPBaseURLs(t) resetModelsCmdFlags(t) t.Cleanup(func() { resetModelsCmdFlags(t) }) setupAuthProviderForProvider(t, "openai", "sk-test", "gpt-image-2") @@ -50,6 +47,7 @@ func TestModelsListActiveProvider(t *testing.T) { } func TestModelsListJSON(t *testing.T) { + trustHTTPBaseURLs(t) resetModelsCmdFlags(t) t.Cleanup(func() { resetModelsCmdFlags(t) }) setupAuthProviderForProvider(t, "openai", "sk-test", "gpt-image-2") @@ -89,6 +87,7 @@ func TestModelsListJSON(t *testing.T) { } func TestModelsNonInteractiveFallback(t *testing.T) { + trustHTTPBaseURLs(t) resetModelsCmdFlags(t) t.Cleanup(func() { resetModelsCmdFlags(t) }) setupAuthProviderForProvider(t, "openai", "sk-test", "gpt-image-2") @@ -112,9 +111,9 @@ func TestModelsNonInteractiveFallback(t *testing.T) { output := buf.String() if !strings.Contains(output, "gpt-image-2") { - t.Errorf("models should fall back to static list, got: %s", output) + t.Errorf("models should print discovered image models, got: %s", output) } - if !strings.Contains(output, "MODEL ID") { + if !strings.Contains(output, "DISPLAY NAME") { t.Errorf("models should print text table, got: %s", output) } if strings.Contains(output, "CAPABILITIES") || strings.Contains(output, "[edit]") { @@ -123,6 +122,7 @@ func TestModelsNonInteractiveFallback(t *testing.T) { } func TestModelsJSONNonInteractive(t *testing.T) { + trustHTTPBaseURLs(t) resetModelsCmdFlags(t) t.Cleanup(func() { resetModelsCmdFlags(t) }) setupAuthProviderForProvider(t, "openai", "sk-test", "gpt-image-2") @@ -188,422 +188,3 @@ func TestModelsListNoActiveProvider(t *testing.T) { t.Errorf("error should mention no active provider, got: %v", err) } } - -func TestModelsSpecificProvider(t *testing.T) { - resetRootCmdFlags(t) - resetAuthAddFlags(t) - t.Setenv("HOME", t.TempDir()) - t.Setenv("POTACO_API_KEY", "") - t.Setenv("POTACO_BASE_URL", "") - t.Setenv("POTACO_PROVIDER", "") - t.Setenv("POTACO_MODEL", "") - - var setupBuf bytes.Buffer - rootCmd.SetOut(&setupBuf) - rootCmd.SetErr(&setupBuf) - - rootCmd.SetArgs([]string{"auth", "add", "openai", "--api-key", "sk-openai", "--force", "--model", "gpt-image-2"}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("setup auth add openai: %v", err) - } - resetAuthAddFlags(t) - resetRootCmdFlags(t) - - rootCmd.SetArgs([]string{"auth", "add", "fal", "--api-key", "fal-key", "--force", "--model", "fal-ai/flux/dev"}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("setup auth add fal: %v", err) - } - resetAuthAddFlags(t) - resetRootCmdFlags(t) - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/v1/models" { - t.Errorf("path = %q, want /v1/models", r.URL.Path) - } - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]any{ - "data": []map[string]any{ - {"id": "gpt-image-2", "owned_by": "openai"}, - }, - }) - })) - defer srv.Close() - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "openai", "--base-url", srv.URL}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("Execute: %v", err) - } - - output := buf.String() - if !strings.Contains(output, "gpt-image-2") { - t.Errorf("models openai should list openai models, got: %s", output) - } -} - -func TestModelsListSpecificProvider(t *testing.T) { - resetRootCmdFlags(t) - resetAuthAddFlags(t) - t.Setenv("HOME", t.TempDir()) - t.Setenv("POTACO_API_KEY", "") - t.Setenv("POTACO_BASE_URL", "") - t.Setenv("POTACO_PROVIDER", "") - t.Setenv("POTACO_MODEL", "") - - var setupBuf bytes.Buffer - rootCmd.SetOut(&setupBuf) - rootCmd.SetErr(&setupBuf) - - rootCmd.SetArgs([]string{"auth", "add", "openai", "--api-key", "sk-openai", "--force", "--model", "gpt-image-2"}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("setup auth add openai: %v", err) - } - resetAuthAddFlags(t) - resetRootCmdFlags(t) - - rootCmd.SetArgs([]string{"auth", "add", "fal", "--api-key", "fal-key", "--force", "--model", "fal-ai/flux/dev"}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("setup auth add fal: %v", err) - } - resetAuthAddFlags(t) - resetRootCmdFlags(t) - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/v1/models" { - t.Errorf("path = %q, want /v1/models", r.URL.Path) - } - authHeader := r.Header.Get("Authorization") - if authHeader != "Bearer sk-openai" { - t.Errorf("Authorization = %q, want Bearer sk-openai", authHeader) - } - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]any{ - "data": []map[string]any{ - {"id": "gpt-image-2", "owned_by": "openai"}, - }, - }) - })) - defer srv.Close() - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "list", "openai", "--base-url", srv.URL}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("Execute: %v", err) - } - - output := buf.String() - if !strings.Contains(output, "gpt-image-2") { - t.Errorf("models list openai should list openai models, got: %s", output) - } -} - -func TestModelsListSpecificProviderNotConnected(t *testing.T) { - t.Setenv("HOME", t.TempDir()) - resetRootCmdFlags(t) - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "list", "fal"}) - err := rootCmd.Execute() - if err == nil { - t.Fatal("expected error for unconnected provider") - } - if !strings.Contains(err.Error(), "not connected") { - t.Errorf("error should mention not connected, got: %v", err) - } -} - -func TestModelsUnknownProvider(t *testing.T) { - t.Setenv("HOME", t.TempDir()) - resetRootCmdFlags(t) - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "unknown-provider"}) - err := rootCmd.Execute() - if err == nil { - t.Fatal("expected error for unknown provider") - } - if !strings.Contains(err.Error(), "unknown provider") { - t.Errorf("error should mention unknown provider, got: %v", err) - } -} - -func TestModelsListUnknownProvider(t *testing.T) { - t.Setenv("HOME", t.TempDir()) - resetRootCmdFlags(t) - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "list", "unknown-provider"}) - err := rootCmd.Execute() - if err == nil { - t.Fatal("expected error for unknown provider") - } - if !strings.Contains(err.Error(), "unknown provider") { - t.Errorf("error should mention unknown provider, got: %v", err) - } -} - -func TestModelsListWithApiKeyOverride(t *testing.T) { - setupAuthProviderForProvider(t, "openai", "sk-stored", "gpt-image-2") - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - var gotAuth string - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotAuth = r.Header.Get("Authorization") - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]any{ - "data": []map[string]any{ - {"id": "gpt-image-2", "owned_by": "openai"}, - }, - }) - })) - defer srv.Close() - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "list", "--api-key", "sk-override", "--base-url", srv.URL}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("Execute: %v", err) - } - - if gotAuth != "Bearer sk-override" { - t.Errorf("Authorization header = %q, want Bearer sk-override", gotAuth) - } -} - -func TestModelsListUsesBaseURLFromConfig(t *testing.T) { - t.Setenv("HOME", t.TempDir()) - resetRootCmdFlags(t) - resetAuthAddFlags(t) - - mgr, err := auth.New() - if err != nil { - t.Fatalf("init auth: %v", err) - } - if err := mgr.Add("openai", "sk-test"); err != nil { - t.Fatalf("add provider: %v", err) - } - if err := mgr.SetBaseURL("openai", "https://config.example.com/v1"); err != nil { - t.Fatalf("set base url: %v", err) - } - - var gotURL string - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotURL = r.URL.Path - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]any{ - "data": []map[string]any{ - {"id": "gpt-image-2", "owned_by": "openai"}, - }, - }) - })) - defer srv.Close() - - // Override the stored base URL with the test server URL so the request lands - // on the mock server, but keep the resolution path. - cfg, err := mgr.LoadConfig() - if err != nil { - t.Fatalf("load config: %v", err) - } - pc := cfg.Providers["openai"] - pc.BaseURL = srv.URL - cfg.Providers["openai"] = pc - if err := config.SaveMultiProvider(config.DefaultConfigPath(), cfg); err != nil { - t.Fatalf("save config: %v", err) - } - - resetRootCmdFlags(t) - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "list"}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("Execute: %v", err) - } - - if gotURL != "/v1/models" { - t.Errorf("request path = %q, want /v1/models", gotURL) - } - if !strings.Contains(buf.String(), "gpt-image-2") { - t.Errorf("output should list gpt-image-2, got: %s", buf.String()) - } -} - -func TestModelsListAliasWithBuiltInTypeRequiresBaseURL(t *testing.T) { - t.Setenv("HOME", t.TempDir()) - resetRootCmdFlags(t) - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - mgr, err := auth.New() - if err != nil { - t.Fatalf("init auth: %v", err) - } - if err := mgr.AddProvider("openrouter", "openai", "sk-test"); err != nil { - t.Fatalf("add provider: %v", err) - } - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "list", "openrouter"}) - err = rootCmd.Execute() - if err == nil { - t.Fatal("expected error for missing alias base URL") - } - if !strings.Contains(err.Error(), "base URL") { - t.Fatalf("error = %v, want base URL", err) - } -} - -func TestModelsDiscoveryError(t *testing.T) { - setupAuthProviderForProvider(t, "openai", "sk-test", "gpt-image-2") - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusInternalServerError) - })) - defer srv.Close() - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "list", "--base-url", srv.URL}) - err := rootCmd.Execute() - if err == nil { - t.Fatal("expected error when discovery fails") - } - if !strings.Contains(err.Error(), "discover models") { - t.Errorf("error should mention discover models, got: %v", err) - } -} - -func TestModelsHelp(t *testing.T) { - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "--help"}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("Execute: %v", err) - } - - output := buf.String() - if !strings.Contains(output, "models [provider]") { - t.Errorf("help should show usage 'models [provider]', got: %s", output) - } - if !strings.Contains(output, "list") { - t.Errorf("help should list subcommand 'list', got: %s", output) - } - if strings.Contains(output, "--params") { - t.Errorf("help should not list --params flag, got: %s", output) - } -} - -func TestModelsListHelp(t *testing.T) { - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "list", "--help"}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("Execute: %v", err) - } - - output := buf.String() - if !strings.Contains(output, "list [provider]") { - t.Errorf("help should show usage 'list [provider]', got: %s", output) - } -} - -func TestModelsMaximumNArgs(t *testing.T) { - t.Setenv("HOME", t.TempDir()) - resetRootCmdFlags(t) - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "openai", "extra"}) - err := rootCmd.Execute() - if err == nil { - t.Fatal("expected error for too many args") - } - if !strings.Contains(err.Error(), "accepts at most 1 arg") { - t.Errorf("error should mention arg limit, got: %v", err) - } -} - -func TestModelsListMaximumNArgs(t *testing.T) { - t.Setenv("HOME", t.TempDir()) - resetRootCmdFlags(t) - resetModelsCmdFlags(t) - t.Cleanup(func() { resetModelsCmdFlags(t) }) - - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetArgs([]string{"models", "list", "openai", "extra"}) - err := rootCmd.Execute() - if err == nil { - t.Fatal("expected error for too many args") - } - if !strings.Contains(err.Error(), "accepts at most 1 arg") { - t.Errorf("error should mention arg limit, got: %v", err) - } -} - -func resetModelsCmdFlags(t *testing.T) { - t.Helper() - for _, name := range []string{"base-url", "api-key"} { - flag := modelsCmd.PersistentFlags().Lookup(name) - if flag == nil { - continue - } - if err := flag.Value.Set(flag.DefValue); err != nil { - t.Fatalf("reset %s flag: %v", name, err) - } - flag.Changed = false - } - for _, name := range []string{"base-url", "api-key"} { - flag := modelsListCmd.Flags().Lookup(name) - if flag == nil { - continue - } - if err := flag.Value.Set(flag.DefValue); err != nil { - t.Fatalf("reset %s child flag: %v", name, err) - } - flag.Changed = false - } - // The help flag is added by Cobra to each command and may be left set after - // a --help invocation, which would suppress arg validation on the next run. - for _, cmd := range []*cobra.Command{modelsCmd, modelsListCmd} { - for _, name := range []string{"help", "version"} { - flag := cmd.Flags().Lookup(name) - if flag == nil { - continue - } - if err := flag.Value.Set("false"); err != nil { - t.Fatalf("reset %s flag on %s: %v", name, cmd.Name(), err) - } - flag.Changed = false - } - } -} diff --git a/internal/cli/openrouter_test.go b/internal/cli/openrouter_test.go new file mode 100644 index 0000000..c4f04b9 --- /dev/null +++ b/internal/cli/openrouter_test.go @@ -0,0 +1,36 @@ +package cli + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "testing" +) + +var openRouterCatalogServer *httptest.Server + +func TestMain(m *testing.M) { + openRouterCatalogServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + openRouterCatalogModel("openai/gpt-image-2"), + openRouterCatalogModel("openrouter/image-model"), + }, + }) + })) + os.Setenv("POTACO_OPENROUTER_MODELS_URL", openRouterCatalogServer.URL+"/models?output_modalities=image") + code := m.Run() + openRouterCatalogServer.Close() + os.Exit(code) +} + +func openRouterCatalogModel(id string) map[string]any { + return map[string]any{ + "id": id, + "architecture": map[string]any{ + "output_modalities": []string{"image"}, + }, + } +} diff --git a/internal/cli/output.go b/internal/cli/output.go index 5c484ae..42ce2c9 100644 --- a/internal/cli/output.go +++ b/internal/cli/output.go @@ -38,7 +38,7 @@ func FormatResult(result OutputResult, opts OutputOptions) (string, error) { var lines []string for _, path := range result.Paths { - lines = append(lines, fmt.Sprintf("Saved to: %s", path)) + lines = append(lines, fmt.Sprintf("Saved to: %s", safeTerminalText(path))) } return strings.Join(lines, "\n"), nil } diff --git a/internal/cli/output_test.go b/internal/cli/output_test.go index 45643c3..365d17f 100644 --- a/internal/cli/output_test.go +++ b/internal/cli/output_test.go @@ -184,6 +184,29 @@ func TestFormatResultDefault(t *testing.T) { } } +func TestFormatResultSanitizesPlainPaths(t *testing.T) { + result := OutputResult{ + Paths: []string{"https://example.test/image.png\x1b[31m\r\nspoofed"}, + Format: "png", + Widths: []int{1024}, + Heights: []int{1024}, + Model: "gpt-image-2", + Params: map[string]any{"size": "1024x1024"}, + } + opts := OutputOptions{JSON: false, Stdout: false} + + output, err := FormatResult(result, opts) + if err != nil { + t.Fatalf("FormatResult error: %v", err) + } + if strings.ContainsAny(output, "\x1b\r\n") { + t.Fatalf("plain output should be terminal-safe, got %q", output) + } + if !strings.Contains(output, "https://example.test/image.png spoofed") { + t.Fatalf("plain output should preserve sanitized path text, got %q", output) + } +} + func TestFormatResultJSON(t *testing.T) { result := OutputResult{ Paths: []string{"output.png"}, diff --git a/internal/cli/provider_type_test.go b/internal/cli/provider_type_test.go index 4854040..c2f6926 100644 --- a/internal/cli/provider_type_test.go +++ b/internal/cli/provider_type_test.go @@ -20,6 +20,7 @@ import ( ) func TestGenUsesConfiguredProviderTypeAndProviderKey(t *testing.T) { + trustHTTPBaseURLs(t) resetRootCmdFlags(t) resetAuthAddFlags(t) resetGenCmdFlags(t) @@ -90,6 +91,7 @@ func TestGenUsesConfiguredProviderTypeAndProviderKey(t *testing.T) { } func TestGenUsesSelectedProviderModelWhenProviderOverrideHasNoModelOverride(t *testing.T) { + trustHTTPBaseURLs(t) resetRootCmdFlags(t) resetGenCmdFlags(t) t.Cleanup(func() { resetGenCmdFlags(t) }) @@ -148,7 +150,8 @@ func TestGenUsesSelectedProviderModelWhenProviderOverrideHasNoModelOverride(t *t } } -func TestModelsListConfiguredOpenAICompatibleProvider(t *testing.T) { +func TestModelsListConfiguredOpenAIProviderAlias(t *testing.T) { + trustHTTPBaseURLs(t) t.Setenv("HOME", t.TempDir()) t.Setenv("POTACO_API_KEY", "") t.Setenv("POTACO_BASE_URL", "") @@ -160,11 +163,19 @@ func TestModelsListConfiguredOpenAICompatibleProvider(t *testing.T) { var gotAuth string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotAuth = r.Header.Get("Authorization") + if r.URL.Path == "/v1/models" { + gotAuth = r.Header.Get("Authorization") + } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]any{ "data": []map[string]any{ - {"id": "openrouter/image-model", "owned_by": "openrouter"}, + { + "id": "openrouter/image-model", + "owned_by": "openrouter", + "architecture": map[string]any{ + "output_modalities": []string{"image"}, + }, + }, }, }) })) @@ -208,7 +219,7 @@ func addOpenAICompatibleProvider(t *testing.T, baseURL string) *auth.AuthManager openaiConfig.Model = "gpt-image-2" cfg.Providers["openai"] = openaiConfig openrouterConfig := cfg.Providers["openrouter"] - openrouterConfig.Type = "openai-compatible" + openrouterConfig.Type = "openai" openrouterConfig.BaseURL = baseURL cfg.Providers["openrouter"] = openrouterConfig cfg.ActiveProvider = "openai" diff --git a/internal/cli/resolve.go b/internal/cli/resolve.go index 5dd6676..af657de 100644 --- a/internal/cli/resolve.go +++ b/internal/cli/resolve.go @@ -23,6 +23,8 @@ type resolvedConfig struct { Model string BaseURL string ModelCanEdit bool + Parameters config.ModelParameters + Fallbacks []*resolvedConfig } // resolveAdapterForCommand resolves the active provider, credential, and @@ -62,8 +64,12 @@ func resolveAdapterForCommand(cmd *cobra.Command) (*resolvedConfig, error) { } model := resolveModel(cmd, cfg, providerName) providerType := config.ResolveProviderType(providerName, pc) - adapterType := config.AdapterType(providerType) + adapterType := config.AdapterTypeForProvider(providerName, providerType) baseURL := resolveBaseURL(cmd, providerName, cfg) + baseURL, err = validateRuntimeBaseURL(baseURL) + if err != nil { + return nil, err + } retries, timeout, err := resolveRetriesTimeout(cmd, cfg, providerName) if err != nil { @@ -95,13 +101,16 @@ func resolveAdapterForCommand(cmd *cobra.Command) (*resolvedConfig, error) { ) } - return &resolvedConfig{ + resolved := &resolvedConfig{ Adapter: ad, ProviderName: providerName, Model: model, BaseURL: baseURL, ModelCanEdit: modelEditCapable(cfg, providerName, model), - }, nil + Parameters: modelParameters(cfg, providerName, model), + } + resolved.Fallbacks = resolveFallbackProviders(cmd, mgr, cfg, resolved) + return resolved, nil } func resolveProvider(cmd *cobra.Command, mgr *auth.AuthManager) (string, error) { @@ -195,6 +204,21 @@ func modelEditCapable(cfg *config.MultiProviderConfig, providerName, model strin return ok && mc.Edit } +func modelParameters(cfg *config.MultiProviderConfig, providerName, model string) config.ModelParameters { + if cfg == nil || model == "" { + return config.DefaultModelParameters() + } + pc, ok := cfg.Providers[providerName] + if !ok { + return config.DefaultModelParameters() + } + mc, ok := pc.Models[model] + if !ok { + return config.DefaultModelParameters() + } + return mc.Parameters.WithDefaults() +} + func resolveRetriesTimeout(cmd *cobra.Command, cfg *config.MultiProviderConfig, providerName string) (int, time.Duration, error) { retries := 2 timeout := 120 * time.Second diff --git a/internal/cli/root.go b/internal/cli/root.go index 9a88f5e..34f2fba 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -128,7 +128,7 @@ func shouldSkipAutoUpdate(cmd *cobra.Command) bool { } for c := cmd; c != nil; c = c.Parent() { switch c.Name() { - case "update", "version", "uninstall", "help", "completion": + case "update", "version", "uninstall", "stats", "help", "completion": return true } } @@ -200,14 +200,24 @@ func init() { rootCmd.PersistentFlags().Bool("non-interactive", false, "force non-interactive mode (skip TUI)") rootCmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error { + startTelemetry(cmd) ni, err := cmd.Flags().GetBool("non-interactive") if err != nil { return fmt.Errorf("read non-interactive flag: %w", err) } tui.SetNonInteractive(ni) - if err := runConfigMigrations(); err != nil { - return err + if !shouldSkipConfigMigration(cmd) { + if err := runConfigMigrations(); err != nil { + return err + } } return runAutoUpdateCheck(cmd) } + rootCmd.PersistentPostRun = func(cmd *cobra.Command, args []string) { + recordSuccessfulCommand(cmd) + } +} + +func shouldSkipConfigMigration(cmd *cobra.Command) bool { + return cmd.Name() == "show" && cmd.Parent() != nil && cmd.Parent().Name() == "config" } diff --git a/internal/cli/root_autoupdate_skip_test.go b/internal/cli/root_autoupdate_skip_test.go new file mode 100644 index 0000000..4339087 --- /dev/null +++ b/internal/cli/root_autoupdate_skip_test.go @@ -0,0 +1,36 @@ +package cli + +import ( + "testing" + + "github.com/spf13/cobra" +) + +func TestAutoUpdateSkipsCommandPaths(t *testing.T) { + for _, cmd := range []*cobra.Command{ + updateCmd, + versionCmd, + uninstallCmd, + {Use: "help"}, + findRootSubcommand(t, "completion"), + } { + t.Run(cmd.Name(), func(t *testing.T) { + resetAutoUpdateTest(t) + autoUpdateIsInteractive = func() bool { return true } + if !shouldSkipAutoUpdate(cmd) { + t.Fatalf("%s path should skip auto-update", cmd.Name()) + } + }) + } +} + +func findRootSubcommand(t *testing.T, name string) *cobra.Command { + t.Helper() + for _, cmd := range rootCmd.Commands() { + if cmd.Name() == name { + return cmd + } + } + t.Fatalf("root command %q not found", name) + return nil +} diff --git a/internal/cli/root_autoupdate_test.go b/internal/cli/root_autoupdate_test.go new file mode 100644 index 0000000..55b142d --- /dev/null +++ b/internal/cli/root_autoupdate_test.go @@ -0,0 +1,253 @@ +package cli + +import ( + "bytes" + "errors" + "path/filepath" + "testing" + "time" + + "github.com/ncxton/potaco/internal/config" + "github.com/ncxton/potaco/internal/tui" + "github.com/spf13/cobra" +) + +func resetAutoUpdateTest(t *testing.T) { + t.Helper() + resetRootCmdFlags(t) + resetVersionCache() + resetConfigSetFlags(t) + resetUpdateCmdFlags(t) + + origVersion := Version + origCheck := autoUpdateCheckLatest + origInstall := autoUpdateInstall + origPrompt := autoUpdatePrompt + origInteractive := autoUpdateIsInteractive + origNow := autoUpdateNow + origTTL := autoUpdateCacheTTL + + Version = "v1.0.0" + autoUpdateCheckLatest = checkLatestVersion + autoUpdateInstall = installUpdate + autoUpdatePrompt = promptAutoUpdate + autoUpdateIsInteractive = tui.IsInteractive + autoUpdateNow = time.Now + autoUpdateCacheTTL = 24 * time.Hour + + t.Cleanup(func() { + Version = origVersion + autoUpdateCheckLatest = origCheck + autoUpdateInstall = origInstall + autoUpdatePrompt = origPrompt + autoUpdateIsInteractive = origInteractive + autoUpdateNow = origNow + autoUpdateCacheTTL = origTTL + resetRootCmdFlags(t) + resetConfigSetFlags(t) + resetUpdateCmdFlags(t) + }) +} + +func seedAutoUpdateConfig(t *testing.T, home string, autoUpdate *bool) { + t.Helper() + cfg := &config.MultiProviderConfig{ + AutoUpdate: autoUpdate, + ActiveProvider: "openai", + Providers: map[string]config.ProviderConfig{ + "openai": {Model: "gpt-image-2"}, + }, + } + if err := config.SaveMultiProvider(filepath.Join(home, ".potaco", "config.yaml"), cfg); err != nil { + t.Fatalf("seed config: %v", err) + } +} + +func runAutoUpdateRoot(t *testing.T, args ...string) (*bytes.Buffer, error) { + t.Helper() + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs(args) + return &buf, rootCmd.Execute() +} + +func TestAutoUpdateFreshCacheAvoidsNetwork(t *testing.T) { + resetAutoUpdateTest(t) + home := t.TempDir() + t.Setenv("HOME", home) + seedAutoUpdateConfig(t, home, nil) + now := time.Date(2026, 6, 28, 12, 0, 0, 0, time.UTC) + autoUpdateNow = func() time.Time { return now } + if err := config.SaveUpdateCache(config.DefaultUpdateCachePath(), &config.UpdateCache{ + LastUpdateCheck: now.Add(-time.Hour), + LatestVersion: "v1.0.0", + }); err != nil { + t.Fatalf("seed update cache: %v", err) + } + + autoUpdateIsInteractive = func() bool { return true } + autoUpdateCheckLatest = func() (string, error) { + t.Fatal("fresh cache should not call latest-version network check") + return "", nil + } + + if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { + t.Fatalf("Execute: %v", err) + } +} + +func TestAutoUpdateStaleCacheChecksLatestAndWritesCache(t *testing.T) { + resetAutoUpdateTest(t) + home := t.TempDir() + t.Setenv("HOME", home) + seedAutoUpdateConfig(t, home, nil) + now := time.Date(2026, 6, 28, 12, 0, 0, 0, time.UTC) + autoUpdateNow = func() time.Time { return now } + if err := config.SaveUpdateCache(config.DefaultUpdateCachePath(), &config.UpdateCache{ + LastUpdateCheck: now.Add(-48 * time.Hour), + LatestVersion: "v1.0.0", + }); err != nil { + t.Fatalf("seed update cache: %v", err) + } + + calls := 0 + autoUpdateIsInteractive = func() bool { return true } + autoUpdateCheckLatest = func() (string, error) { + calls++ + return "v1.0.0", nil + } + + if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { + t.Fatalf("Execute: %v", err) + } + if calls != 1 { + t.Fatalf("latest check calls = %d, want 1", calls) + } + cache, err := config.LoadUpdateCache(config.DefaultUpdateCachePath()) + if err != nil { + t.Fatalf("LoadUpdateCache: %v", err) + } + if !cache.LastUpdateCheck.Equal(now) || cache.LatestVersion != "v1.0.0" { + t.Fatalf("cache = %+v, want refreshed at %v with v1.0.0", cache, now) + } +} + +func TestAutoUpdateDismissedVersionSuppressesPrompt(t *testing.T) { + resetAutoUpdateTest(t) + home := t.TempDir() + t.Setenv("HOME", home) + seedAutoUpdateConfig(t, home, nil) + now := time.Date(2026, 6, 28, 12, 0, 0, 0, time.UTC) + autoUpdateNow = func() time.Time { return now } + if err := config.SaveUpdateCache(config.DefaultUpdateCachePath(), &config.UpdateCache{ + LastUpdateCheck: now.Add(-time.Hour), + LatestVersion: "v1.1.0", + DismissedVersion: "v1.1.0", + }); err != nil { + t.Fatalf("seed update cache: %v", err) + } + + autoUpdateIsInteractive = func() bool { return true } + autoUpdatePrompt = func(string) (bool, error) { + t.Fatal("dismissed version should not prompt") + return false, nil + } + + if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { + t.Fatalf("Execute: %v", err) + } +} + +func TestAutoUpdateNoStoresDismissedVersion(t *testing.T) { + resetAutoUpdateTest(t) + home := t.TempDir() + t.Setenv("HOME", home) + seedAutoUpdateConfig(t, home, nil) + now := time.Date(2026, 6, 28, 12, 0, 0, 0, time.UTC) + autoUpdateNow = func() time.Time { return now } + + autoUpdateIsInteractive = func() bool { return true } + autoUpdateCheckLatest = func() (string, error) { return "v1.1.0", nil } + autoUpdatePrompt = func(string) (bool, error) { return false, nil } + autoUpdateInstall = func(*cobra.Command, string) error { + t.Fatal("declined update should not run installer") + return nil + } + + if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { + t.Fatalf("Execute: %v", err) + } + cache, err := config.LoadUpdateCache(config.DefaultUpdateCachePath()) + if err != nil { + t.Fatalf("LoadUpdateCache: %v", err) + } + if cache.DismissedVersion != "v1.1.0" { + t.Fatalf("DismissedVersion = %q, want v1.1.0", cache.DismissedVersion) + } +} + +func TestAutoUpdateInstallerFailureReturnsError(t *testing.T) { + resetAutoUpdateTest(t) + home := t.TempDir() + t.Setenv("HOME", home) + seedAutoUpdateConfig(t, home, nil) + + autoUpdateIsInteractive = func() bool { return true } + autoUpdateCheckLatest = func() (string, error) { return "v1.1.0", nil } + autoUpdatePrompt = func(string) (bool, error) { return true, nil } + autoUpdateInstall = func(*cobra.Command, string) error { return errors.New("installer failed") } + + _, err := runAutoUpdateRoot(t, "config", "show") + if err == nil { + t.Fatal("accepted update should return installer failure") + } +} + +func TestAutoUpdateSkips(t *testing.T) { + for _, tc := range []struct { + name string + args []string + prep func() + }{ + {name: "unknown version", args: []string{"config", "show"}, prep: func() { Version = "unknown" }}, + {name: "non interactive flag", args: []string{"--non-interactive", "config", "show"}, prep: func() { autoUpdateIsInteractive = tui.IsInteractive }}, + {name: "json", args: []string{"--json", "config", "show"}}, + {name: "json after subcommand", args: []string{"config", "show", "--json"}}, + {name: "disable auto update", args: []string{"config", "set", "auto_update", "false"}}, + } { + t.Run(tc.name, func(t *testing.T) { + resetAutoUpdateTest(t) + home := t.TempDir() + t.Setenv("HOME", home) + seedAutoUpdateConfig(t, home, nil) + autoUpdateIsInteractive = func() bool { return true } + autoUpdateCheckLatest = func() (string, error) { + t.Fatal("skip path should not check latest version") + return "", nil + } + if tc.prep != nil { + tc.prep() + } + + _, _ = runAutoUpdateRoot(t, tc.args...) + }) + } +} + +func TestAutoUpdateDisabledInConfigSkips(t *testing.T) { + resetAutoUpdateTest(t) + home := t.TempDir() + t.Setenv("HOME", home) + disabled := false + seedAutoUpdateConfig(t, home, &disabled) + autoUpdateIsInteractive = func() bool { return true } + autoUpdateCheckLatest = func() (string, error) { + t.Fatal("auto_update false should skip latest version check") + return "", nil + } + + if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { + t.Fatalf("Execute: %v", err) + } +} diff --git a/internal/cli/root_migrate_test.go b/internal/cli/root_migrate_test.go new file mode 100644 index 0000000..e3b1b70 --- /dev/null +++ b/internal/cli/root_migrate_test.go @@ -0,0 +1,55 @@ +package cli + +import ( + "bytes" + "os" + "path/filepath" + "testing" + + "github.com/ncxton/potaco/internal/config" + "github.com/ncxton/potaco/internal/tui" +) + +const legacyCustomProviderConfigYAML = ` +active_provider: custom +providers: + custom: + base_url: https://api.example.com/v1 + model: gpt-image-1 + retries: 2 + timeout: 120 +` + +func TestRootCommandMigratesConfigOnStartup(t *testing.T) { + resetRootCmdFlags(t) + tmpHome := t.TempDir() + t.Setenv("HOME", tmpHome) + t.Cleanup(func() { tui.SetNonInteractive(false) }) + + configPath := filepath.Join(tmpHome, ".potaco", "config.yaml") + if err := os.MkdirAll(filepath.Dir(configPath), 0700); err != nil { + t.Fatalf("mkdir config dir: %v", err) + } + if err := os.WriteFile(configPath, []byte(legacyCustomProviderConfigYAML), 0600); err != nil { + t.Fatalf("write config: %v", err) + } + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + loaded, err := config.LoadMultiProvider(configPath) + if err != nil { + t.Fatalf("LoadMultiProvider: %v", err) + } + if loaded.SchemaVersion != config.CurrentSchemaVersion { + t.Fatalf("SchemaVersion = %d, want %d", loaded.SchemaVersion, config.CurrentSchemaVersion) + } + if got := loaded.Providers["custom"].Type; got != "openai" { + t.Fatalf("custom Type = %q, want openai", got) + } +} diff --git a/internal/cli/root_test.go b/internal/cli/root_test.go index 5b6e1a2..e0b8dac 100644 --- a/internal/cli/root_test.go +++ b/internal/cli/root_test.go @@ -2,28 +2,12 @@ package cli import ( "bytes" - "errors" - "os" - "path/filepath" "strings" "testing" - "time" - "github.com/ncxton/potaco/internal/config" "github.com/ncxton/potaco/internal/tui" - "github.com/spf13/cobra" ) -const legacyCustomProviderConfigYAML = ` -active_provider: custom -providers: - custom: - base_url: https://api.example.com/v1 - model: gpt-image-1 - retries: 2 - timeout: 120 -` - // resetRootCmdFlags restores persistent flags on the shared global rootCmd to // their default values. Tests that dispatch rootCmd.SetArgs must call this // first so that flags left set by earlier tests (e.g. --json) do not leak in @@ -95,24 +79,24 @@ func TestRootCommandHasNonInteractiveFlag(t *testing.T) { } } -func TestNonInteractiveFlagWiresToTUI(t *testing.T) { +func TestInfoCommandIsUnavailable(t *testing.T) { resetRootCmdFlags(t) t.Setenv("HOME", t.TempDir()) - t.Cleanup(func() { tui.SetNonInteractive(false) }) - var buf bytes.Buffer rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) - rootCmd.SetArgs([]string{"--non-interactive"}) - if err := rootCmd.Execute(); err != nil { - t.Fatalf("Execute: %v", err) + rootCmd.SetArgs([]string{"info"}) + + err := rootCmd.Execute() + if err == nil { + t.Fatal("potaco info should be unavailable") } - if tui.IsInteractive() { - t.Error("IsInteractive() should return false after --non-interactive is set") + if !strings.Contains(err.Error(), "unknown command") { + t.Fatalf("error = %v, want unknown command", err) } } -func TestNonInteractiveFlagDefaultsToInteractive(t *testing.T) { +func TestNonInteractiveFlagWiresToTUI(t *testing.T) { resetRootCmdFlags(t) t.Setenv("HOME", t.TempDir()) t.Cleanup(func() { tui.SetNonInteractive(false) }) @@ -120,31 +104,20 @@ func TestNonInteractiveFlagDefaultsToInteractive(t *testing.T) { var buf bytes.Buffer rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) - rootCmd.SetArgs([]string{}) + rootCmd.SetArgs([]string{"--non-interactive"}) if err := rootCmd.Execute(); err != nil { t.Fatalf("Execute: %v", err) } - - // In a test environment (no TTY), IsInteractive should return false if tui.IsInteractive() { - t.Error("IsInteractive() should return false in test environment (no TTY)") + t.Error("IsInteractive() should return false after --non-interactive is set") } } -func TestRootCommandMigratesConfigOnStartup(t *testing.T) { +func TestNonInteractiveFlagDefaultsToInteractive(t *testing.T) { resetRootCmdFlags(t) - tmpHome := t.TempDir() - t.Setenv("HOME", tmpHome) + t.Setenv("HOME", t.TempDir()) t.Cleanup(func() { tui.SetNonInteractive(false) }) - configPath := filepath.Join(tmpHome, ".potaco", "config.yaml") - if err := os.MkdirAll(filepath.Dir(configPath), 0700); err != nil { - t.Fatalf("mkdir config dir: %v", err) - } - if err := os.WriteFile(configPath, []byte(legacyCustomProviderConfigYAML), 0600); err != nil { - t.Fatalf("write config: %v", err) - } - var buf bytes.Buffer rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) @@ -153,285 +126,9 @@ func TestRootCommandMigratesConfigOnStartup(t *testing.T) { t.Fatalf("Execute: %v", err) } - loaded, err := config.LoadMultiProvider(configPath) - if err != nil { - t.Fatalf("LoadMultiProvider: %v", err) - } - if loaded.SchemaVersion != config.CurrentSchemaVersion { - t.Fatalf("SchemaVersion = %d, want %d", loaded.SchemaVersion, config.CurrentSchemaVersion) - } - if got := loaded.Providers["custom"].Type; got != "openai-compatible" { - t.Fatalf("custom Type = %q, want openai-compatible", got) - } -} - -func resetAutoUpdateTest(t *testing.T) { - t.Helper() - resetRootCmdFlags(t) - resetVersionCache() - resetConfigSetFlags(t) - resetUpdateCmdFlags(t) - - origVersion := Version - origCheck := autoUpdateCheckLatest - origInstall := autoUpdateInstall - origPrompt := autoUpdatePrompt - origInteractive := autoUpdateIsInteractive - origNow := autoUpdateNow - origTTL := autoUpdateCacheTTL - origURL := githubReleaseURL - - Version = "v1.0.0" - autoUpdateCheckLatest = checkLatestVersion - autoUpdateInstall = installUpdate - autoUpdatePrompt = promptAutoUpdate - autoUpdateIsInteractive = tui.IsInteractive - autoUpdateNow = time.Now - autoUpdateCacheTTL = 24 * time.Hour - tui.SetNonInteractive(false) - - t.Cleanup(func() { - Version = origVersion - autoUpdateCheckLatest = origCheck - autoUpdateInstall = origInstall - autoUpdatePrompt = origPrompt - autoUpdateIsInteractive = origInteractive - autoUpdateNow = origNow - autoUpdateCacheTTL = origTTL - githubReleaseURL = origURL - tui.SetNonInteractive(false) - }) -} - -func seedAutoUpdateConfig(t *testing.T, home string, autoUpdate *bool) { - t.Helper() - cfg := &config.MultiProviderConfig{ - AutoUpdate: autoUpdate, - ActiveProvider: "openai", - Providers: map[string]config.ProviderConfig{ - "openai": {Model: "gpt-image-2"}, - }, - } - if err := config.SaveMultiProvider(filepath.Join(home, ".potaco", "config.yaml"), cfg); err != nil { - t.Fatalf("seed config: %v", err) - } -} - -func runAutoUpdateRoot(t *testing.T, args ...string) (*bytes.Buffer, error) { - t.Helper() - var buf bytes.Buffer - rootCmd.SetOut(&buf) - rootCmd.SetErr(&buf) - rootCmd.SetArgs(args) - return &buf, rootCmd.Execute() -} - -func TestAutoUpdateFreshCacheAvoidsNetwork(t *testing.T) { - resetAutoUpdateTest(t) - home := t.TempDir() - t.Setenv("HOME", home) - seedAutoUpdateConfig(t, home, nil) - now := time.Date(2026, 6, 28, 12, 0, 0, 0, time.UTC) - autoUpdateNow = func() time.Time { return now } - if err := config.SaveUpdateCache(config.DefaultUpdateCachePath(), &config.UpdateCache{ - LastUpdateCheck: now.Add(-time.Hour), - LatestVersion: "v1.0.0", - }); err != nil { - t.Fatalf("seed update cache: %v", err) - } - - autoUpdateIsInteractive = func() bool { return true } - autoUpdateCheckLatest = func() (string, error) { - t.Fatal("fresh cache should not call latest-version network check") - return "", nil - } - - if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { - t.Fatalf("Execute: %v", err) - } -} - -func TestAutoUpdateStaleCacheChecksLatestAndWritesCache(t *testing.T) { - resetAutoUpdateTest(t) - home := t.TempDir() - t.Setenv("HOME", home) - seedAutoUpdateConfig(t, home, nil) - now := time.Date(2026, 6, 28, 12, 0, 0, 0, time.UTC) - autoUpdateNow = func() time.Time { return now } - if err := config.SaveUpdateCache(config.DefaultUpdateCachePath(), &config.UpdateCache{ - LastUpdateCheck: now.Add(-48 * time.Hour), - LatestVersion: "v1.0.0", - }); err != nil { - t.Fatalf("seed update cache: %v", err) - } - - calls := 0 - autoUpdateIsInteractive = func() bool { return true } - autoUpdateCheckLatest = func() (string, error) { - calls++ - return "v1.0.0", nil - } - - if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { - t.Fatalf("Execute: %v", err) - } - if calls != 1 { - t.Fatalf("latest check calls = %d, want 1", calls) - } - cache, err := config.LoadUpdateCache(config.DefaultUpdateCachePath()) - if err != nil { - t.Fatalf("LoadUpdateCache: %v", err) - } - if !cache.LastUpdateCheck.Equal(now) || cache.LatestVersion != "v1.0.0" { - t.Fatalf("cache = %+v, want refreshed at %v with v1.0.0", cache, now) - } -} - -func TestAutoUpdateDismissedVersionSuppressesPrompt(t *testing.T) { - resetAutoUpdateTest(t) - home := t.TempDir() - t.Setenv("HOME", home) - seedAutoUpdateConfig(t, home, nil) - now := time.Date(2026, 6, 28, 12, 0, 0, 0, time.UTC) - autoUpdateNow = func() time.Time { return now } - if err := config.SaveUpdateCache(config.DefaultUpdateCachePath(), &config.UpdateCache{ - LastUpdateCheck: now.Add(-time.Hour), - LatestVersion: "v1.1.0", - DismissedVersion: "v1.1.0", - }); err != nil { - t.Fatalf("seed update cache: %v", err) - } - - autoUpdateIsInteractive = func() bool { return true } - autoUpdatePrompt = func(string) (bool, error) { - t.Fatal("dismissed version should not prompt") - return false, nil - } - - if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { - t.Fatalf("Execute: %v", err) - } -} - -func TestAutoUpdateNoStoresDismissedVersion(t *testing.T) { - resetAutoUpdateTest(t) - home := t.TempDir() - t.Setenv("HOME", home) - seedAutoUpdateConfig(t, home, nil) - now := time.Date(2026, 6, 28, 12, 0, 0, 0, time.UTC) - autoUpdateNow = func() time.Time { return now } - - autoUpdateIsInteractive = func() bool { return true } - autoUpdateCheckLatest = func() (string, error) { return "v1.1.0", nil } - autoUpdatePrompt = func(string) (bool, error) { return false, nil } - autoUpdateInstall = func(*cobra.Command, string) error { - t.Fatal("declined update should not run installer") - return nil - } - - if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { - t.Fatalf("Execute: %v", err) - } - cache, err := config.LoadUpdateCache(config.DefaultUpdateCachePath()) - if err != nil { - t.Fatalf("LoadUpdateCache: %v", err) - } - if cache.DismissedVersion != "v1.1.0" { - t.Fatalf("DismissedVersion = %q, want v1.1.0", cache.DismissedVersion) - } -} - -func TestAutoUpdateInstallerFailureReturnsError(t *testing.T) { - resetAutoUpdateTest(t) - home := t.TempDir() - t.Setenv("HOME", home) - seedAutoUpdateConfig(t, home, nil) - - autoUpdateIsInteractive = func() bool { return true } - autoUpdateCheckLatest = func() (string, error) { return "v1.1.0", nil } - autoUpdatePrompt = func(string) (bool, error) { return true, nil } - autoUpdateInstall = func(*cobra.Command, string) error { return errors.New("installer failed") } - - _, err := runAutoUpdateRoot(t, "config", "show") - if err == nil { - t.Fatal("accepted update should return installer failure") - } -} - -func TestAutoUpdateSkips(t *testing.T) { - for _, tc := range []struct { - name string - args []string - prep func() - }{ - {name: "unknown version", args: []string{"config", "show"}, prep: func() { Version = "unknown" }}, - {name: "non interactive flag", args: []string{"--non-interactive", "config", "show"}, prep: func() { autoUpdateIsInteractive = tui.IsInteractive }}, - {name: "json", args: []string{"--json", "config", "show"}}, - {name: "json after subcommand", args: []string{"config", "show", "--json"}}, - {name: "disable auto update", args: []string{"config", "set", "auto_update", "false"}}, - } { - t.Run(tc.name, func(t *testing.T) { - resetAutoUpdateTest(t) - home := t.TempDir() - t.Setenv("HOME", home) - seedAutoUpdateConfig(t, home, nil) - autoUpdateIsInteractive = func() bool { return true } - autoUpdateCheckLatest = func() (string, error) { - t.Fatal("skip path should not check latest version") - return "", nil - } - if tc.prep != nil { - tc.prep() - } - - _, _ = runAutoUpdateRoot(t, tc.args...) - }) - } -} - -func TestAutoUpdateSkipsCommandPaths(t *testing.T) { - for _, cmd := range []*cobra.Command{ - updateCmd, - versionCmd, - uninstallCmd, - {Use: "help"}, - findRootSubcommand(t, "completion"), - } { - t.Run(cmd.Name(), func(t *testing.T) { - resetAutoUpdateTest(t) - autoUpdateIsInteractive = func() bool { return true } - if !shouldSkipAutoUpdate(cmd) { - t.Fatalf("%s path should skip auto-update", cmd.Name()) - } - }) - } -} - -func findRootSubcommand(t *testing.T, name string) *cobra.Command { - t.Helper() - for _, cmd := range rootCmd.Commands() { - if cmd.Name() == name { - return cmd - } - } - t.Fatalf("root command %q not found", name) - return nil -} - -func TestAutoUpdateDisabledInConfigSkips(t *testing.T) { - resetAutoUpdateTest(t) - home := t.TempDir() - t.Setenv("HOME", home) - disabled := false - seedAutoUpdateConfig(t, home, &disabled) - autoUpdateIsInteractive = func() bool { return true } - autoUpdateCheckLatest = func() (string, error) { - t.Fatal("auto_update false should skip latest version check") - return "", nil - } - - if _, err := runAutoUpdateRoot(t, "config", "show"); err != nil { - t.Fatalf("Execute: %v", err) + // In a test environment (no TTY), IsInteractive should return false + if tui.IsInteractive() { + t.Error("IsInteractive() should return false in test environment (no TTY)") } } diff --git a/internal/cli/stats_cmd.go b/internal/cli/stats_cmd.go new file mode 100644 index 0000000..055b567 --- /dev/null +++ b/internal/cli/stats_cmd.go @@ -0,0 +1,146 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/colorprofile" + "github.com/ncxton/potaco/internal/stats" + "github.com/ncxton/potaco/internal/tui" + "github.com/spf13/cobra" +) + +var statsCmd = &cobra.Command{ + Use: "stats", + Short: "Show local command analytics", + RunE: runStats, +} + +func init() { + statsCmd.Flags().Bool("clear", false, "clear local stats history") + statsCmd.Flags().Bool("no-prices", false, "skip OpenRouter pricing fetch") + rootCmd.AddCommand(statsCmd) +} + +func runStats(cmd *cobra.Command, args []string) error { + if flagBool(cmd, "clear") { + if err := stats.Clear(stats.DefaultPath()); err != nil { + return configUserErr("Could not clear stats.", "Check that ~/.potaco/ is writable.", err) + } + renderSuccess(cmd.OutOrStdout(), "Stats cleared.") + return nil + } + + events, err := stats.Load(stats.DefaultPath()) + if err != nil { + return configUserErr("Could not load stats.", "Check that ~/.potaco/stats.jsonl is readable.", err) + } + prices := map[string]stats.Price{} + priceStatus := "OpenRouter pricing disabled" + if !flagBool(cmd, "no-prices") { + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + if fetched, err := stats.FetchOpenRouterPrices(ctx, &http.Client{Timeout: 3 * time.Second}); err == nil { + prices = fetched + priceStatus = "OpenRouter pricing active" + } else { + priceStatus = "OpenRouter pricing unavailable" + } + } + + summary := stats.Summarize(events, prices) + jsonMode, _ := cmd.Root().PersistentFlags().GetBool("json") + if jsonMode { + return printStatsJSON(cmd.OutOrStdout(), summary, priceStatus) + } + if !tui.IsInteractive() { + printStatsText(cmd.OutOrStdout(), summary, priceStatus) + return nil + } + program := tea.NewProgram(newStatsDashboard(summary, prices, priceStatus), tea.WithOutput(cmd.OutOrStdout()), tea.WithColorProfile(colorprofile.ANSI256)) + _, err = program.Run() + return err +} + +func printStatsJSON(out io.Writer, summary stats.Summary, priceStatus string) error { + body := map[string]any{ + "total": summary.Total, + "successes": summary.Successes, + "errors": summary.Errors, + "avg_duration_ms": summary.AvgDuration.Milliseconds(), + "estimated_cost": summary.CostUSD, + "pricing_status": priceStatus, + "by_command": statsBucketJSON(summary.ByCommand), + "by_provider": statsBucketJSON(summary.ByProvider), + "recent_commands": summary.Recent, + "stats_file": stats.DefaultPath(), + "cost_is_estimate": true, + } + data, err := json.MarshalIndent(body, "", " ") + if err != nil { + return fmt.Errorf("marshal stats JSON: %w", err) + } + fmt.Fprintln(out, string(data)) + return nil +} + +func statsBucketJSON(buckets []stats.Bucket) []map[string]any { + out := make([]map[string]any, 0, len(buckets)) + for _, bucket := range buckets { + out = append(out, map[string]any{ + "name": bucket.Name, + "count": bucket.Count, + "errors": bucket.Errors, + "duration_ms": bucket.DurationSum.Milliseconds(), + "cost_usd": bucket.CostUSD, + }) + } + return out +} + +func printStatsText(out io.Writer, summary stats.Summary, priceStatus string) { + fmt.Fprintln(out, "Potaco Stats") + fmt.Fprintf(out, "Commands: %d Success: %d Errors: %d Avg: %s Estimated cost: $%.4f\n", summary.Total, summary.Successes, summary.Errors, formatDuration(summary.AvgDuration), summary.CostUSD) + fmt.Fprintf(out, "Pricing: %s\n", priceStatus) + fmt.Fprintln(out) + fmt.Fprintln(out, "Top commands") + for _, bucket := range firstBuckets(summary.ByCommand, 6) { + fmt.Fprintf(out, " %-18s %4d errors:%d cost:$%.4f\n", safeTerminalText(bucket.Name), bucket.Count, bucket.Errors, bucket.CostUSD) + } + fmt.Fprintln(out) + fmt.Fprintln(out, "Providers") + for _, bucket := range firstBuckets(summary.ByProvider, 6) { + fmt.Fprintf(out, " %-18s %4d errors:%d cost:$%.4f\n", displayProviderName(bucket.Name), bucket.Count, bucket.Errors, bucket.CostUSD) + } + fmt.Fprintln(out) + fmt.Fprintln(out, "Recent") + for _, event := range summary.Recent { + status := "ok" + if !event.Success { + status = "err" + } + fmt.Fprintf(out, " %-8s %-18s %-12s %-24s %s\n", status, safeTerminalText(event.Command), displayProviderName(event.Provider), safeTerminalText(event.Model), formatDuration(time.Duration(event.DurationMS)*time.Millisecond)) + } +} + +func firstBuckets(buckets []stats.Bucket, limit int) []stats.Bucket { + if len(buckets) <= limit { + return buckets + } + return buckets[:limit] +} + +func formatDuration(d time.Duration) string { + if d <= 0 { + return "0ms" + } + if d < time.Second { + return fmt.Sprintf("%dms", d.Milliseconds()) + } + return d.Round(100 * time.Millisecond).String() +} diff --git a/internal/cli/stats_cmd_test.go b/internal/cli/stats_cmd_test.go new file mode 100644 index 0000000..bb4227f --- /dev/null +++ b/internal/cli/stats_cmd_test.go @@ -0,0 +1,193 @@ +package cli + +import ( + "bytes" + "encoding/json" + "strings" + "testing" + "time" + + "github.com/ncxton/potaco/internal/stats" +) + +func TestStatsCommandExists(t *testing.T) { + found := false + for _, cmd := range rootCmd.Commands() { + if cmd.Name() == "stats" { + found = true + break + } + } + if !found { + t.Fatal("root command should have stats subcommand") + } +} + +func TestStatsNonInteractiveShowsSummary(t *testing.T) { + resetRootCmdFlags(t) + resetStatsCmdFlags(t) + t.Cleanup(func() { resetStatsCmdFlags(t) }) + t.Setenv("HOME", t.TempDir()) + seedStatsEvent(t, stats.Event{Timestamp: time.Unix(1, 0), Command: "potaco gen", Operation: "generate", Provider: "openai", Model: "gpt-image-2", DurationMS: 1200, Success: true, Outputs: 1}) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"stats", "--non-interactive", "--no-prices"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("stats: %v", err) + } + + output := buf.String() + for _, want := range []string{"Potaco Stats", "Commands: 1", "Estimated cost:", "potaco gen", "openai"} { + if !strings.Contains(output, want) { + t.Fatalf("stats output missing %q:\n%s", want, output) + } + } +} + +func TestStatsTextSanitizesTerminalControls(t *testing.T) { + var buf bytes.Buffer + summary := stats.Summary{ + Total: 1, + ByCommand: []stats.Bucket{ + {Name: "potaco\x1b[0m gen", Count: 1}, + }, + ByProvider: []stats.Bucket{ + {Name: "evil\x1b]0;title\aprovider", Count: 1}, + }, + Recent: []stats.Event{ + {Command: "potaco\nedit", Provider: "evil\x1b[31mprovider", Model: "model\rname", Success: true}, + }, + } + + printStatsText(&buf, summary, "pricing") + + output := buf.String() + if strings.ContainsAny(output, "\x1b\a\r") { + t.Fatalf("stats output contains terminal controls: %q", output) + } + for _, want := range []string{"potaco gen", "evilprovider", "potaco edit", "model name"} { + if !strings.Contains(output, want) { + t.Fatalf("stats output missing sanitized text %q:\n%s", want, output) + } + } +} + +func TestStatsJSON(t *testing.T) { + resetRootCmdFlags(t) + resetStatsCmdFlags(t) + t.Cleanup(func() { resetStatsCmdFlags(t) }) + t.Setenv("HOME", t.TempDir()) + seedStatsEvent(t, stats.Event{Timestamp: time.Unix(1, 0), Command: "potaco gen", Operation: "generate", Provider: "openai", Model: "gpt-image-2", DurationMS: 1200, Success: true, Outputs: 1}) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"--json", "stats", "--non-interactive", "--no-prices"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("stats --json: %v", err) + } + + var body map[string]any + if err := json.Unmarshal(buf.Bytes(), &body); err != nil { + t.Fatalf("decode json: %v\n%s", err, buf.String()) + } + if body["total"] != float64(1) { + t.Fatalf("total = %v, want 1", body["total"]) + } +} + +func TestStatsClear(t *testing.T) { + resetRootCmdFlags(t) + resetStatsCmdFlags(t) + t.Cleanup(func() { resetStatsCmdFlags(t) }) + t.Setenv("HOME", t.TempDir()) + seedStatsEvent(t, stats.Event{Timestamp: time.Unix(1, 0), Command: "potaco gen", Success: true}) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"stats", "--clear", "--non-interactive"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("stats --clear: %v", err) + } + if !strings.Contains(buf.String(), iconSuccess+" Stats cleared.") { + t.Fatalf("clear output = %q", buf.String()) + } + events, err := stats.Load(stats.DefaultPath()) + if err != nil { + t.Fatalf("load stats: %v", err) + } + if len(events) != 0 { + t.Fatalf("events after clear = %d, want 0", len(events)) + } +} + +func TestGenDryRunRecordsStatsEvent(t *testing.T) { + setupAuthProvider(t, "sk-test") + resetGenCmdFlags(t) + t.Cleanup(func() { resetGenCmdFlags(t) }) + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"gen", "--prompt", "stats", "--dry-run"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("gen --dry-run: %v", err) + } + + events, err := stats.Load(stats.DefaultPath()) + if err != nil { + t.Fatalf("load stats: %v", err) + } + for _, event := range events { + if event.Operation == "generate" && event.Provider == "openai" && event.Model == "gpt-image-2" && event.Success && event.DryRun { + return + } + } + t.Fatalf("missing generate stats event: %+v", events) +} + +func TestRootRecordsSuccessfulCommandStatsEvent(t *testing.T) { + resetRootCmdFlags(t) + resetAuthAddFlags(t) + t.Setenv("HOME", t.TempDir()) + t.Setenv("POTACO_API_KEY", "") + t.Setenv("POTACO_BASE_URL", "") + + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + rootCmd.SetArgs([]string{"auth", "list"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("auth list: %v", err) + } + + events, err := stats.Load(stats.DefaultPath()) + if err != nil { + t.Fatalf("load stats: %v", err) + } + if len(events) != 1 || events[0].Command != "potaco auth list" || !events[0].Success { + t.Fatalf("events = %+v, want auth list success", events) + } +} + +func seedStatsEvent(t *testing.T, event stats.Event) { + t.Helper() + if err := stats.Append(stats.DefaultPath(), event); err != nil { + t.Fatalf("seed stats: %v", err) + } +} + +func resetStatsCmdFlags(t *testing.T) { + t.Helper() + for _, name := range []string{"clear", "no-prices"} { + flag := statsCmd.Flags().Lookup(name) + if flag == nil { + continue + } + _ = flag.Value.Set(flag.DefValue) + flag.Changed = false + } +} diff --git a/internal/cli/stats_dashboard.go b/internal/cli/stats_dashboard.go new file mode 100644 index 0000000..c3b620f --- /dev/null +++ b/internal/cli/stats_dashboard.go @@ -0,0 +1,184 @@ +package cli + +import ( + "fmt" + "strings" + "time" + + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + "github.com/ncxton/potaco/internal/stats" +) + +type statsDashboard struct { + summary stats.Summary + prices map[string]stats.Price + priceStatus string + activePanel int + width int + height int + status string +} + +var ( + statsTitleStyle = lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("42")) + statsMutedStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + statsPanelStyle = lipgloss.NewStyle().Border(lipgloss.RoundedBorder()).Padding(0, 1) + statsActiveStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("42")).Bold(true) +) + +func newStatsDashboard(summary stats.Summary, prices map[string]stats.Price, priceStatus string) statsDashboard { + return statsDashboard{summary: summary, prices: prices, priceStatus: priceStatus, width: 100, height: 32} +} + +func (m statsDashboard) Init() tea.Cmd { + return nil +} + +func (m statsDashboard) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.width = msg.Width + m.height = msg.Height + case tea.KeyPressMsg: + switch msg.Keystroke() { + case "q", "esc", "ctrl+c": + return m, tea.Quit + case "tab": + m.activePanel = (m.activePanel + 1) % 3 + case "r": + events, err := stats.Load(stats.DefaultPath()) + if err != nil { + m.status = iconWarning + " could not refresh stats" + return m, nil + } + m.summary = stats.Summarize(events, m.prices) + m.status = iconSuccess + " refreshed" + case "c": + if err := stats.Clear(stats.DefaultPath()); err != nil { + m.status = iconWarning + " could not clear stats" + return m, nil + } + m.summary = stats.Summary{} + m.status = iconSuccess + " cleared" + } + } + return m, nil +} + +func (m statsDashboard) View() tea.View { + width := m.width + if width <= 0 { + width = 100 + } + panelBodyWidth := max(1, width-4) + panelWidth := max(1, (width-12)/3) + header := statsTitleStyle.Render("Potaco Stats") + statsMutedStyle.Render(" "+m.priceStatus) + if m.status != "" { + header += " " + m.renderStatus() + } + summary := m.panel(0, panelWidth, m.summaryPanel()) + commands := m.panel(1, panelWidth, m.commandPanel()) + providers := m.panel(2, panelWidth, m.providerPanel()) + topBody := lipgloss.JoinHorizontal(lipgloss.Top, summary, " ", commands, " ", providers) + if width < 72 { + summary = m.panel(0, panelBodyWidth, m.summaryPanel()) + commands = m.panel(1, panelBodyWidth, m.commandPanel()) + providers = m.panel(2, panelBodyWidth, m.providerPanel()) + topBody = strings.Join([]string{summary, commands, providers}, "\n") + } + top := statsPanelStyle.Width(panelBodyWidth).Render(topBody) + recent := statsPanelStyle.Width(panelBodyWidth).Render(m.recentPanel(max(1, width-8))) + help := statsMutedStyle.Render("tab focus r refresh c clear q/ctrl+c close") + view := tea.NewView(strings.Join([]string{header, top, recent, help}, "\n")) + view.AltScreen = true + return view +} + +func (m statsDashboard) renderStatus() string { + if strings.HasPrefix(m.status, iconSuccess) { + return successStyle.Render(m.status) + } + return warnStyle.Render(m.status) +} + +func (m statsDashboard) panel(index, width int, body string) string { + body = padLines(body, 6) + style := lipgloss.NewStyle().Width(width) + if m.activePanel == index { + style = statsActiveStyle.Width(width) + } + return style.Render(body) +} + +func (m statsDashboard) summaryPanel() string { + return strings.Join([]string{ + statsTitleStyle.Render("Overview"), + fmt.Sprintf("Commands %d", m.summary.Total), + fmt.Sprintf("Success %d", m.summary.Successes), + fmt.Sprintf("Errors %d", m.summary.Errors), + fmt.Sprintf("Avg duration %s", formatDuration(m.summary.AvgDuration)), + fmt.Sprintf("Est. cost $%.4f", m.summary.CostUSD), + }, "\n") +} + +func (m statsDashboard) commandPanel() string { + lines := []string{statsTitleStyle.Render("Commands")} + for _, bucket := range firstBuckets(m.summary.ByCommand, 6) { + lines = append(lines, fmt.Sprintf("%-14s %4d err %d", shortText(safeTerminalText(bucket.Name), 14), bucket.Count, bucket.Errors)) + } + if len(lines) == 1 { + lines = append(lines, statsMutedStyle.Render("No command data")) + } + return strings.Join(lines, "\n") +} + +func (m statsDashboard) providerPanel() string { + lines := []string{statsTitleStyle.Render("Providers")} + for _, bucket := range firstBuckets(m.summary.ByProvider, 6) { + lines = append(lines, fmt.Sprintf("%-14s %4d $%.3f", shortText(displayProviderName(bucket.Name), 14), bucket.Count, bucket.CostUSD)) + } + if len(lines) == 1 { + lines = append(lines, statsMutedStyle.Render("No provider data")) + } + return strings.Join(lines, "\n") +} + +func (m statsDashboard) recentPanel(width int) string { + lines := []string{statsTitleStyle.Render("Recent commands")} + if len(m.summary.Recent) == 0 { + return strings.Join(append(lines, statsMutedStyle.Render("No command history yet")), "\n") + } + lines = append(lines, fmt.Sprintf("%-5s %-16s %-12s %-24s %8s", "state", "command", "provider", "model", "time")) + for _, event := range m.summary.Recent { + state := "ok" + if !event.Success { + state = "err" + } + line := fmt.Sprintf("%-5s %-16s %-12s %-24s %8s", state, shortText(safeTerminalText(event.Command), 16), shortText(displayProviderName(event.Provider), 12), shortText(safeTerminalText(event.Model), 24), formatDuration(time.Duration(event.DurationMS)*time.Millisecond)) + lines = append(lines, shortText(line, width)) + } + return strings.Join(lines, "\n") +} + +func shortText(value string, limit int) string { + if limit <= 0 || lipgloss.Width(value) <= limit { + return value + } + if limit <= 1 { + return "." + } + runes := []rune(value) + for len(runes) > 0 && lipgloss.Width(string(runes))+3 > limit { + runes = runes[:len(runes)-1] + } + return string(runes) + "..." +} + +func padLines(value string, count int) string { + lines := strings.Split(value, "\n") + for len(lines) < count { + lines = append(lines, "") + } + return strings.Join(lines, "\n") +} diff --git a/internal/cli/status_cmd.go b/internal/cli/status_cmd.go index 320b4b4..c5d004e 100644 --- a/internal/cli/status_cmd.go +++ b/internal/cli/status_cmd.go @@ -55,17 +55,17 @@ func runStatus(cmd *cobra.Command, args []string) error { fmt.Fprintln(out, "No active provider configured.") fmt.Fprintln(out, "Use 'potaco auth add ' to connect one.") } else if tui.IsTTY() { - fmt.Fprintf(out, "%s %s\n", statusLabelStyle.Render("Active provider:"), statusActiveStyle.Render(provider)) - fmt.Fprintf(out, "%s %s\n", statusLabelStyle.Render("Active model:"), model) + fmt.Fprintf(out, "%s %s\n", statusLabelStyle.Render("Active provider:"), statusActiveStyle.Render(displayProviderName(provider))) + fmt.Fprintf(out, "%s %s\n", statusLabelStyle.Render("Active model:"), safeTerminalText(model)) fmt.Fprintf(out, "%s %s\n", statusLabelStyle.Render("Base URL:"), formatStatusBaseURL(activeBaseURL)) } else { - fmt.Fprintf(out, "Active provider: %s\n", provider) - fmt.Fprintf(out, "Active model: %s\n", model) + fmt.Fprintf(out, "Active provider: %s\n", displayProviderName(provider)) + fmt.Fprintf(out, "Active model: %s\n", safeTerminalText(model)) fmt.Fprintf(out, "Base URL: %s\n", formatStatusBaseURL(activeBaseURL)) } fmt.Fprintln(out) - fmt.Fprintf(out, "Config file: %s\n", configPath) - fmt.Fprintf(out, "Credentials: %s\n", credPath) + fmt.Fprintf(out, "Config file: %s\n", safeTerminalText(configPath)) + fmt.Fprintf(out, "Credentials: %s\n", safeTerminalText(credPath)) fmt.Fprintln(out) if len(providers) == 0 { @@ -88,7 +88,7 @@ func runStatus(cmd *cobra.Command, args []string) error { if p.AddedAt != "" { added = " added: " + p.AddedAt } - fmt.Fprintf(out, " %s\t%s\tbase_url: %s\tkey: %s%s%s\n", p.Name, p.Model, baseURL, keyStatus, active, added) + fmt.Fprintf(out, " %s\t%s\tbase_url: %s\tkey: %s%s%s\n", displayProviderName(p.Name), safeTerminalText(p.Model), baseURL, keyStatus, active, added) } } @@ -99,7 +99,7 @@ func formatStatusBaseURL(baseURL string) string { if baseURL == "" { return "default" } - return baseURL + return safeTerminalText(baseURL) } func printStatusJSON(out io.Writer, provider, model, activeBaseURL, configPath, credPath string, providers []auth.ProviderInfo) error { diff --git a/internal/cli/telemetry.go b/internal/cli/telemetry.go new file mode 100644 index 0000000..9a193a1 --- /dev/null +++ b/internal/cli/telemetry.go @@ -0,0 +1,60 @@ +package cli + +import ( + "context" + "time" + + "github.com/ncxton/potaco/internal/stats" + "github.com/spf13/cobra" +) + +type telemetryKey struct{} + +func startTelemetry(cmd *cobra.Command) { + cmd.SetContext(context.WithValue(cmd.Context(), telemetryKey{}, time.Now())) +} + +func recordSuccessfulCommand(cmd *cobra.Command) { + if skipTelemetry(cmd) { + return + } + start := telemetryStart(cmd) + if start.IsZero() { + start = time.Now() + } + recordStatsEvent(stats.Event{ + Command: cmd.CommandPath(), + Operation: cmd.Name(), + DurationMS: time.Since(start).Milliseconds(), + Success: true, + }) +} + +func recordStatsEvent(event stats.Event) { + if event.Command == "" { + event.Command = "potaco" + } + if event.Timestamp.IsZero() { + event.Timestamp = time.Now() + } + _ = stats.Append(stats.DefaultPath(), event) +} + +func telemetryStart(cmd *cobra.Command) time.Time { + if start, ok := cmd.Context().Value(telemetryKey{}).(time.Time); ok { + return start + } + return time.Time{} +} + +func skipTelemetry(cmd *cobra.Command) bool { + for c := cmd; c != nil; c = c.Parent() { + switch c.Name() { + case "stats", "help", "completion": + return true + case "gen", "edit": + return true + } + } + return false +} diff --git a/internal/cli/terminal_text.go b/internal/cli/terminal_text.go new file mode 100644 index 0000000..f548885 --- /dev/null +++ b/internal/cli/terminal_text.go @@ -0,0 +1,68 @@ +package cli + +import ( + "strings" + "unicode/utf8" +) + +func safeTerminalText(s string) string { + var b strings.Builder + b.Grow(len(s)) + for i := 0; i < len(s); { + r, size := utf8.DecodeRuneInString(s[i:]) + switch { + case r == '\x1b': + i = skipANSI(s, i+size) + case r == '\x9b': + i = skipCSI(s, i+size) + case r == '\n' || r == '\r' || r == '\t': + b.WriteByte(' ') + i += size + case r < 0x20 || (r >= 0x7f && r <= 0x9f): + i += size + default: + b.WriteRune(r) + i += size + } + } + return strings.Join(strings.Fields(b.String()), " ") +} + +func skipANSI(s string, i int) int { + if i >= len(s) { + return i + } + switch s[i] { + case '[': + return skipCSI(s, i+1) + case ']': + return skipOSC(s, i+1) + default: + _, size := utf8.DecodeRuneInString(s[i:]) + return i + size + } +} + +func skipCSI(s string, i int) int { + for i < len(s) { + c := s[i] + i++ + if c >= 0x40 && c <= 0x7e { + return i + } + } + return i +} + +func skipOSC(s string, i int) int { + for i < len(s) { + if s[i] == '\a' { + return i + 1 + } + if s[i] == '\x1b' && i+1 < len(s) && s[i+1] == '\\' { + return i + 2 + } + i++ + } + return i +} diff --git a/internal/cli/uninstall_cmd.go b/internal/cli/uninstall_cmd.go index 2ab7532..7bd149b 100644 --- a/internal/cli/uninstall_cmd.go +++ b/internal/cli/uninstall_cmd.go @@ -20,6 +20,7 @@ var uninstallCmd = &cobra.Command{ // symlinks. It is a package-level variable so tests can override it with // a fake path instead of resolving the test binary itself. var findPotacoBinaryFn = findPotacoBinary +var confirmActionFn = tui.ConfirmAction func init() { uninstallCmd.Flags().BoolP("yes", "y", false, "skip confirmation prompts") @@ -33,7 +34,7 @@ func runUninstall(cmd *cobra.Command, args []string) error { // Locate the binary binaryPath, findErr := findPotacoBinaryFn() if findErr != nil && nonInteractive { - fmt.Fprintf(cmd.OutOrStdout(), "Warning: potaco binary not found: %v\n", findErr) + renderWarningf(cmd.OutOrStdout(), "potaco binary not found: %v", findErr) } // Determine home directory @@ -61,7 +62,7 @@ func runUninstallNonInteractive(cmd *cobra.Command, binaryPath, configDir string if binaryPath != "" { if err := os.Remove(binaryPath); err != nil { if os.IsNotExist(err) { - fmt.Fprintf(out, "Binary not found at %s (may have been removed already).\n", binaryPath) + fmt.Fprintf(out, "Binary not found at %s (may have been removed already).\n", safeTerminalText(binaryPath)) } else { return configUserErr( "Cannot remove the binary.", @@ -70,11 +71,11 @@ func runUninstallNonInteractive(cmd *cobra.Command, binaryPath, configDir string ) } } else { - fmt.Fprintf(out, "Binary removed: %s\n", binaryPath) + fmt.Fprintf(out, "Binary removed: %s\n", safeTerminalText(binaryPath)) } } - fmt.Fprintf(out, "Uninstall complete.\n") + renderSuccess(out, "Uninstall complete.") return nil } @@ -83,13 +84,13 @@ func runUninstallInteractive(cmd *cobra.Command, binaryPath, configDir string, y // Step 1: Confirm binary removal if !yes { - confirmed, err := tui.ConfirmAction(fmt.Sprintf("This will remove the potaco binary at '%s'. Continue?", binaryPath)) + confirmed, err := confirmActionFn(fmt.Sprintf("This will remove the potaco binary at '%s'. Continue?", binaryPath)) if err != nil { - fmt.Fprintln(out, "Cancelled.") + renderError(out, "Cancelled.") return nil } if !confirmed { - fmt.Fprintln(out, "Cancelled.") + renderError(out, "Cancelled.") return nil } } @@ -97,7 +98,12 @@ func runUninstallInteractive(cmd *cobra.Command, binaryPath, configDir string, y // Step 2: Ask about config removal removeConfig := false if !yes { - removeConfig, _ = tui.ConfirmAction(fmt.Sprintf("Also remove configuration and credentials at '%s'?", configDir)) + var err error + removeConfig, err = confirmActionFn(fmt.Sprintf("Also remove configuration and credentials at '%s'?", configDir)) + if err != nil { + renderError(out, "Cancelled.") + return nil + } } // Step 3: Final confirmation @@ -107,13 +113,13 @@ func runUninstallInteractive(cmd *cobra.Command, binaryPath, configDir string, y summary += " and config directory" } summary += "?" - finalConfirm, err := tui.ConfirmAction(summary) + finalConfirm, err := confirmActionFn(summary) if err != nil { - fmt.Fprintln(out, "Cancelled.") + renderError(out, "Cancelled.") return nil } if !finalConfirm { - fmt.Fprintln(out, "Cancelled.") + renderError(out, "Cancelled.") return nil } } @@ -122,7 +128,7 @@ func runUninstallInteractive(cmd *cobra.Command, binaryPath, configDir string, y if binaryPath != "" { if err := os.Remove(binaryPath); err != nil { if os.IsNotExist(err) { - fmt.Fprintf(out, "Binary not found at %s (may have been removed already).\n", binaryPath) + fmt.Fprintf(out, "Binary not found at %s (may have been removed already).\n", safeTerminalText(binaryPath)) } else { return configUserErr( "Cannot remove the binary.", @@ -131,19 +137,19 @@ func runUninstallInteractive(cmd *cobra.Command, binaryPath, configDir string, y ) } } else { - fmt.Fprintf(out, "Binary removed: %s\n", binaryPath) + fmt.Fprintf(out, "Binary removed: %s\n", safeTerminalText(binaryPath)) } } if removeConfig { if err := os.RemoveAll(configDir); err != nil { - fmt.Fprintf(out, "Warning: could not remove config directory: %v\n", err) + renderWarningf(out, "could not remove config directory: %v", err) } else { - fmt.Fprintf(out, "Config directory removed: %s\n", configDir) + fmt.Fprintf(out, "Config directory removed: %s\n", safeTerminalText(configDir)) } } - fmt.Fprintf(out, "Uninstall complete.\n") + renderSuccess(out, "Uninstall complete.") return nil } diff --git a/internal/cli/uninstall_cmd_test.go b/internal/cli/uninstall_cmd_test.go index b333e25..aeb2212 100644 --- a/internal/cli/uninstall_cmd_test.go +++ b/internal/cli/uninstall_cmd_test.go @@ -2,10 +2,13 @@ package cli import ( "bytes" + "errors" "os" "path/filepath" "strings" "testing" + + "github.com/spf13/cobra" ) func resetUninstallCmdFlags(t *testing.T) { @@ -34,6 +37,13 @@ func withBinaryFinder(t *testing.T, path string) { t.Cleanup(func() { findPotacoBinaryFn = orig }) } +func withConfirmAction(t *testing.T, fn func(string) (bool, error)) { + t.Helper() + orig := confirmActionFn + confirmActionFn = fn + t.Cleanup(func() { confirmActionFn = orig }) +} + func TestUninstallCommandExists(t *testing.T) { found := false for _, cmd := range rootCmd.Commands() { @@ -169,3 +179,44 @@ func TestUninstallBinaryNotFoundWarns(t *testing.T) { t.Errorf("output should mention binary not found, got: %q", output) } } + +func TestUninstallInteractiveCancelConfigPromptDoesNotRemove(t *testing.T) { + tmpHome := t.TempDir() + binPath := filepath.Join(tmpHome, "potaco") + if err := os.WriteFile(binPath, []byte("fake binary"), 0755); err != nil { + t.Fatalf("write binary: %v", err) + } + configDir := filepath.Join(tmpHome, ".potaco") + if err := os.MkdirAll(configDir, 0755); err != nil { + t.Fatalf("mkdir config: %v", err) + } + + calls := 0 + withConfirmAction(t, func(string) (bool, error) { + calls++ + if calls == 1 { + return true, nil + } + return false, errors.New("cancelled") + }) + + var buf bytes.Buffer + cmd := &cobra.Command{} + cmd.SetOut(&buf) + + if err := runUninstallInteractive(cmd, binPath, configDir, false); err != nil { + t.Fatalf("runUninstallInteractive: %v", err) + } + if calls != 2 { + t.Fatalf("confirm calls = %d, want 2", calls) + } + if _, err := os.Stat(binPath); err != nil { + t.Fatalf("binary should remain after cancel: %v", err) + } + if _, err := os.Stat(configDir); err != nil { + t.Fatalf("config dir should remain after cancel: %v", err) + } + if !strings.Contains(buf.String(), "Cancelled.") { + t.Fatalf("output = %q, want Cancelled.", buf.String()) + } +} diff --git a/internal/cli/update_cmd.go b/internal/cli/update_cmd.go index 09b13e5..04445b3 100644 --- a/internal/cli/update_cmd.go +++ b/internal/cli/update_cmd.go @@ -44,7 +44,7 @@ func runUpdate(cmd *cobra.Command, args []string) error { } if !force && Version != "unknown" && Version == latest { - fmt.Fprintf(cmd.OutOrStdout(), "Already up to date (%s).\n", Version) + renderSuccessf(cmd.OutOrStdout(), "Already up to date (%s).", Version) return nil } @@ -56,7 +56,7 @@ func runUpdate(cmd *cobra.Command, args []string) error { return err } - fmt.Fprintf(cmd.OutOrStdout(), "Update complete.\n") + renderSuccess(cmd.OutOrStdout(), "Update complete.") return nil } diff --git a/internal/cli/update_cmd_test.go b/internal/cli/update_cmd_test.go index dc6577f..c7c30c7 100644 --- a/internal/cli/update_cmd_test.go +++ b/internal/cli/update_cmd_test.go @@ -213,8 +213,8 @@ YAML if loaded.SchemaVersion != config.CurrentSchemaVersion { t.Fatalf("SchemaVersion = %d, want %d", loaded.SchemaVersion, config.CurrentSchemaVersion) } - if got := loaded.Providers["custom"].Type; got != "openai-compatible" { - t.Fatalf("custom Type = %q, want openai-compatible", got) + if got := loaded.Providers["custom"].Type; got != "openai" { + t.Fatalf("custom Type = %q, want openai", got) } backups, err := filepath.Glob(configPath + ".bak-*") if err != nil { diff --git a/internal/cli/use_cmd.go b/internal/cli/use_cmd.go index 622f235..49b2f63 100644 --- a/internal/cli/use_cmd.go +++ b/internal/cli/use_cmd.go @@ -1,9 +1,13 @@ package cli import ( + "context" + "errors" "fmt" + "github.com/ncxton/potaco/internal/adapter" "github.com/ncxton/potaco/internal/auth" + "github.com/ncxton/potaco/internal/config" "github.com/ncxton/potaco/internal/tui" "github.com/spf13/cobra" ) @@ -36,14 +40,76 @@ func runUse(cmd *cobra.Command, args []string) error { return configError(fmt.Errorf("init auth: %w", err)) } + if model != "" { + if err := validateUseModel(cmd, mgr, providerName, model); err != nil { + return err + } + } + if err := mgr.SetActiveProvider(providerName, model); err != nil { return configError(fmt.Errorf("switch provider: %w", err)) } if model != "" { - fmt.Fprintf(cmd.OutOrStdout(), "Switched to provider '%s' with model '%s'.\n", providerName, model) + renderSuccessf(cmd.OutOrStdout(), "Switched to provider '%s' with model '%s'.", displayProviderName(providerName), model) } else { - fmt.Fprintf(cmd.OutOrStdout(), "Switched to provider '%s'.\n", providerName) + renderSuccessf(cmd.OutOrStdout(), "Switched to provider '%s'.", displayProviderName(providerName)) } return nil } + +func validateUseModel(cmd *cobra.Command, mgr *auth.AuthManager, providerName, model string) error { + apiKey, err := resolveAPIKey(cmd, mgr, providerName) + if err != nil { + return err + } + cfg, _ := mgr.LoadConfig() + pc := config.ProviderConfig{} + if cfg != nil { + if configured, ok := cfg.Providers[providerName]; ok { + pc = configured + } + } + providerType := config.ResolveProviderType(providerName, pc) + baseURL := resolveBaseURL(cmd, providerName, cfg) + baseURL, err = validateDiscoveryBaseURL(baseURL) + if err != nil { + return err + } + if config.ProviderRequiresBaseURL(providerName, providerType) && baseURL == "" { + return configUserErr( + "A base URL is required for this provider.", + fmt.Sprintf("Use POTACO_BASE_URL or run 'potaco config set providers.%s.base_url '.", providerName), + fmt.Errorf("base URL required for provider %s", providerName), + ) + } + ad, err := adapter.Get(config.AdapterTypeForProvider(providerName, providerType), apiKey, adapter.AdapterOpts{BaseURL: baseURL}) + if err != nil { + return configUserErr( + fmt.Sprintf("Could not connect to provider '%s'.", providerName), + "Check that the provider name is correct. Use 'potaco auth list' to see connected providers.", + fmt.Errorf("create adapter: %w", err), + ) + } + models, err := ad.DiscoverModels(context.Background()) + if err != nil { + if errors.Is(err, adapter.ErrDiscoveryFailed) { + return nil + } + return apiUserErr( + "Could not discover models.", + "Check your API key, base URL, and network connection.", + fmt.Errorf("discover models: %w", err), + ) + } + for _, discovered := range models { + if discovered.ID == model { + return nil + } + } + return configUserErr( + fmt.Sprintf("Model '%s' was not found for provider '%s'.", model, providerName), + "Run 'potaco models list ' to see available image generation models.", + fmt.Errorf("model %q not found for provider %s", model, providerName), + ) +} diff --git a/internal/cli/use_cmd_test.go b/internal/cli/use_cmd_test.go index b53a55a..7163e20 100644 --- a/internal/cli/use_cmd_test.go +++ b/internal/cli/use_cmd_test.go @@ -2,8 +2,14 @@ package cli import ( "bytes" + "encoding/json" + "net/http" + "net/http/httptest" "strings" "testing" + + "github.com/ncxton/potaco/internal/auth" + "github.com/ncxton/potaco/internal/config" ) func TestUseCommandExists(t *testing.T) { @@ -22,6 +28,8 @@ func TestUseCommandExists(t *testing.T) { func TestUseSwitchesProvider(t *testing.T) { resetRootCmdFlags(t) resetAuthAddFlags(t) + resetUseCmdFlags(t) + t.Cleanup(func() { resetUseCmdFlags(t) }) tmpHome := t.TempDir() t.Setenv("HOME", tmpHome) var buf bytes.Buffer @@ -33,6 +41,7 @@ func TestUseSwitchesProvider(t *testing.T) { rootCmd.Execute() resetAuthAddFlags(t) resetRootCmdFlags(t) + resetUseCmdFlags(t) rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) @@ -40,6 +49,7 @@ func TestUseSwitchesProvider(t *testing.T) { rootCmd.Execute() resetAuthAddFlags(t) resetRootCmdFlags(t) + resetUseCmdFlags(t) rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) @@ -57,8 +67,11 @@ func TestUseSwitchesProvider(t *testing.T) { } func TestUseWithModel(t *testing.T) { + trustHTTPBaseURLs(t) resetRootCmdFlags(t) resetAuthAddFlags(t) + resetUseCmdFlags(t) + t.Cleanup(func() { resetUseCmdFlags(t) }) tmpHome := t.TempDir() t.Setenv("HOME", tmpHome) var buf bytes.Buffer @@ -69,11 +82,36 @@ func TestUseWithModel(t *testing.T) { rootCmd.Execute() resetAuthAddFlags(t) resetRootCmdFlags(t) + resetUseCmdFlags(t) rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) + providerSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "gpt-image-2", "owned_by": "openai"}, + }, + }) + })) + defer providerSrv.Close() + mgr, err := auth.New() + if err != nil { + t.Fatalf("init auth: %v", err) + } + cfg, err := mgr.LoadConfig() + if err != nil { + t.Fatalf("load config: %v", err) + } + pc := cfg.Providers["openai"] + pc.BaseURL = providerSrv.URL + cfg.Providers["openai"] = pc + if err := config.SaveMultiProvider(config.DefaultConfigPath(), cfg); err != nil { + t.Fatalf("save config: %v", err) + } + rootCmd.SetArgs([]string{"use", "openai", "--model", "gpt-image-2"}) - err := rootCmd.Execute() + err = rootCmd.Execute() if err != nil { t.Fatalf("use with model: %v", err) } @@ -84,9 +122,49 @@ func TestUseWithModel(t *testing.T) { } } +func TestUseReplicateWithModelSkipsDiscoveryFailure(t *testing.T) { + resetRootCmdFlags(t) + resetAuthAddFlags(t) + resetUseCmdFlags(t) + t.Cleanup(func() { resetUseCmdFlags(t) }) + t.Setenv("HOME", t.TempDir()) + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + + rootCmd.SetArgs([]string{"auth", "add", "replicate", "--api-key", "r8-test", "--force"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("auth add replicate: %v", err) + } + resetAuthAddFlags(t) + resetRootCmdFlags(t) + resetUseCmdFlags(t) + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + + rootCmd.SetArgs([]string{"use", "replicate", "--model", "owner/model"}) + if err := rootCmd.Execute(); err != nil { + t.Fatalf("use replicate with model: %v", err) + } + + mgr, err := auth.New() + if err != nil { + t.Fatalf("init auth: %v", err) + } + cfg, err := mgr.LoadConfig() + if err != nil { + t.Fatalf("load config: %v", err) + } + if cfg.ActiveProvider != "replicate" || cfg.ActiveModel != "owner/model" { + t.Fatalf("active = %q/%q, want replicate/owner/model", cfg.ActiveProvider, cfg.ActiveModel) + } +} + func TestUseNoArgs(t *testing.T) { resetRootCmdFlags(t) resetAuthAddFlags(t) + resetUseCmdFlags(t) + t.Cleanup(func() { resetUseCmdFlags(t) }) tmpHome := t.TempDir() t.Setenv("HOME", tmpHome) var buf bytes.Buffer @@ -103,6 +181,8 @@ func TestUseNoArgs(t *testing.T) { func TestUseUnknownProvider(t *testing.T) { resetRootCmdFlags(t) resetAuthAddFlags(t) + resetUseCmdFlags(t) + t.Cleanup(func() { resetUseCmdFlags(t) }) tmpHome := t.TempDir() t.Setenv("HOME", tmpHome) var buf bytes.Buffer @@ -115,3 +195,74 @@ func TestUseUnknownProvider(t *testing.T) { t.Fatal("use with unknown provider should error") } } + +func TestUseWithModelRejectsUnknownModelNonInteractive(t *testing.T) { + trustHTTPBaseURLs(t) + resetRootCmdFlags(t) + resetAuthAddFlags(t) + resetUseCmdFlags(t) + t.Cleanup(func() { resetUseCmdFlags(t) }) + tmpHome := t.TempDir() + t.Setenv("HOME", tmpHome) + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + + mgr, err := auth.New() + if err != nil { + t.Fatalf("init auth: %v", err) + } + if err := mgr.AddProvider("openai", "openai", "sk-test"); err != nil { + t.Fatalf("add provider: %v", err) + } + + providerSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + {"id": "gpt-image-2", "owned_by": "openai"}, + }, + }) + })) + defer providerSrv.Close() + cfg, err := mgr.LoadConfig() + if err != nil { + t.Fatalf("load config: %v", err) + } + pc := cfg.Providers["openai"] + pc.BaseURL = providerSrv.URL + cfg.Providers["openai"] = pc + if err := config.SaveMultiProvider(config.DefaultConfigPath(), cfg); err != nil { + t.Fatalf("save config: %v", err) + } + + rootCmd.SetArgs([]string{"use", "openai", "--model", "missing-model", "--non-interactive"}) + err = rootCmd.Execute() + if err == nil { + t.Fatal("expected error for missing model") + } + if !strings.Contains(err.Error(), "missing-model") { + t.Fatalf("error = %v, want missing model name", err) + } + _, model, err := mgr.GetActiveProvider() + if err != nil { + t.Fatalf("get active provider: %v", err) + } + if model == "missing-model" { + t.Fatal("missing model should not be persisted") + } +} + +func resetUseCmdFlags(t *testing.T) { + t.Helper() + for _, name := range []string{"model"} { + flag := useCmd.Flags().Lookup(name) + if flag == nil { + continue + } + if err := flag.Value.Set(flag.DefValue); err != nil { + t.Fatalf("reset %s flag: %v", name, err) + } + flag.Changed = false + } +} diff --git a/internal/cli/usererr.go b/internal/cli/usererr.go index 1661457..1474e5e 100644 --- a/internal/cli/usererr.go +++ b/internal/cli/usererr.go @@ -77,24 +77,63 @@ func renderUserError(w io.Writer, ue *UserError) { renderUserErrorColored(w, ue) return } - fmt.Fprintf(w, "Error: %s\n", ue.Message) + fmt.Fprintf(w, "%s %s\n", iconError, safeTerminalText(userErrorLine(ue))) +} + +const ( + iconWarning = "" + iconError = "" + iconSuccess = "" +) + +func userErrorLine(ue *UserError) string { if ue.Hint != "" { - fmt.Fprintf(w, "Hint: %s\n", ue.Hint) + return ue.Message + " " + ue.Hint } + return ue.Message } var ( - errLabelStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("1")).Bold(true) - errMsgStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("15")) - errHintLabelStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("3")).Bold(true) - errHintMsgStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("3")) + errStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("203")) + warnStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("3")) + successStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("42")) ) func renderUserErrorColored(w io.Writer, ue *UserError) { - fmt.Fprintf(w, "%s %s\n", errLabelStyle.Render("Error:"), errMsgStyle.Render(ue.Message)) - if ue.Hint != "" { - fmt.Fprintf(w, "%s %s\n", errHintLabelStyle.Render("Hint:"), errHintMsgStyle.Render(ue.Hint)) + renderColoredStatusLine(w, errStyle, iconError, userErrorLine(ue)) +} + +func renderError(w io.Writer, message string) { + renderStatusLine(w, errStyle, iconError, message) +} + +func renderWarning(w io.Writer, message string) { + renderStatusLine(w, warnStyle, iconWarning, message) +} + +func renderWarningf(w io.Writer, format string, args ...any) { + renderWarning(w, fmt.Sprintf(format, args...)) +} + +func renderSuccess(w io.Writer, message string) { + renderStatusLine(w, successStyle, iconSuccess, message) +} + +func renderSuccessf(w io.Writer, format string, args ...any) { + renderSuccess(w, fmt.Sprintf(format, args...)) +} + +func renderStatusLine(w io.Writer, style lipgloss.Style, icon, message string) { + message = safeTerminalText(message) + if isWriterTTY(w) { + renderColoredStatusLine(w, style, icon, message) + return } + fmt.Fprintf(w, "%s %s\n", icon, message) +} + +func renderColoredStatusLine(w io.Writer, style lipgloss.Style, icon, message string) { + fmt.Fprintln(w, style.Render(icon+" "+safeTerminalText(message))) } func isWriterTTY(w io.Writer) bool { @@ -154,9 +193,9 @@ func renderAnyError(w io.Writer, err error) { return } if isWriterTTY(w) { - fmt.Fprintf(w, "%s %s\n", errLabelStyle.Render("Error:"), errMsgStyle.Render(err.Error())) + renderColoredStatusLine(w, errStyle, iconError, err.Error()) } else { - fmt.Fprintf(w, "Error: %s\n", err.Error()) + fmt.Fprintf(w, "%s %s\n", iconError, safeTerminalText(err.Error())) } } diff --git a/internal/cli/usererr_test.go b/internal/cli/usererr_test.go index a3996ae..33b24f1 100644 --- a/internal/cli/usererr_test.go +++ b/internal/cli/usererr_test.go @@ -55,14 +55,20 @@ func TestRenderUserErrorNoHint(t *testing.T) { var buf bytes.Buffer renderUserError(&buf, ue) output := buf.String() - if !strings.Contains(output, "Error:") { - t.Errorf("should contain 'Error:', got: %q", output) + if strings.Contains(output, "Error:") { + t.Errorf("should not contain type prefix, got: %q", output) + } + if strings.Contains(output, "Hint:") { + t.Errorf("should not render a separate hint line, got: %q", output) + } + if !strings.Contains(output, "") { + t.Errorf("should contain cancel/error icon, got: %q", output) } if !strings.Contains(output, "Something went wrong.") { t.Errorf("should contain message, got: %q", output) } - if strings.Contains(output, "Hint:") { - t.Errorf("should not contain Hint when empty, got: %q", output) + if got := strings.Count(strings.TrimSpace(output), "\n"); got != 0 { + t.Errorf("should render exactly one message line, got: %q", output) } } @@ -71,18 +77,71 @@ func TestRenderUserErrorWithHint(t *testing.T) { var buf bytes.Buffer renderUserError(&buf, ue) output := buf.String() - if !strings.Contains(output, "Error:") { - t.Errorf("should contain 'Error:', got: %q", output) + if strings.Contains(output, "Error:") || strings.Contains(output, "Hint:") { + t.Errorf("should not contain type prefixes, got: %q", output) } if !strings.Contains(output, "API failed.") { t.Errorf("should contain message, got: %q", output) } - if !strings.Contains(output, "Hint:") { - t.Errorf("should contain 'Hint:', got: %q", output) - } if !strings.Contains(output, "Check your key.") { t.Errorf("should contain hint text, got: %q", output) } + if got := strings.Count(strings.TrimSpace(output), "\n"); got != 0 { + t.Errorf("should render exactly one message line, got: %q", output) + } +} + +func TestRenderUserErrorColoredColorsWholeLine(t *testing.T) { + ue := &UserError{Category: CatConfig, Message: "base URL is required"} + var buf bytes.Buffer + renderUserErrorColored(&buf, ue) + + want := "\x1b[38;5;203m base URL is required\x1b[m\n" + if buf.String() != want { + t.Fatalf("colored error = %q, want %q", buf.String(), want) + } +} + +func TestRenderColoredStatusLinesUseTypeColors(t *testing.T) { + tests := []struct { + name string + want string + render func(*bytes.Buffer) + }{ + {name: "error", want: "\x1b[38;5;203m" + iconError + " message", render: func(buf *bytes.Buffer) { renderColoredStatusLine(buf, errStyle, iconError, "message") }}, + {name: "warning", want: "\x1b[33m" + iconWarning + " message", render: func(buf *bytes.Buffer) { renderColoredStatusLine(buf, warnStyle, iconWarning, "message") }}, + {name: "success", want: "\x1b[38;5;42m" + iconSuccess + " message", render: func(buf *bytes.Buffer) { renderColoredStatusLine(buf, successStyle, iconSuccess, "message") }}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var buf bytes.Buffer + tc.render(&buf) + if !strings.Contains(buf.String(), tc.want) { + t.Fatalf("colored status = %q, want %q", buf.String(), tc.want) + } + }) + } +} + +func TestRenderStatusLineStripsTerminalControls(t *testing.T) { + var buf bytes.Buffer + renderColoredStatusLine(&buf, errStyle, iconError, "bad\x1b[0m\n\x1b]0;title\aok") + + want := "\x1b[38;5;203m bad ok\x1b[m\n" + if buf.String() != want { + t.Fatalf("colored status = %q, want %q", buf.String(), want) + } +} + +func TestRenderUserErrorStripsTerminalControlsWithoutTTY(t *testing.T) { + ue := &UserError{Category: CatConfig, Message: "bad\x1b[0m\nok"} + var buf bytes.Buffer + + renderUserError(&buf, ue) + + if got := buf.String(); got != " bad ok\n" { + t.Fatalf("plain error = %q, want sanitized one-line error", got) + } } func TestValidateOutputPathEmpty(t *testing.T) { diff --git a/internal/config/base_url.go b/internal/config/base_url.go new file mode 100644 index 0000000..71710d1 --- /dev/null +++ b/internal/config/base_url.go @@ -0,0 +1,32 @@ +package config + +import ( + "fmt" + "net/url" + "strings" +) + +func NormalizeBaseURL(raw string) (string, error) { + baseURL := strings.TrimRight(strings.TrimSpace(raw), "/") + if baseURL == "" { + return "", fmt.Errorf("base URL is required") + } + parsed, err := url.Parse(baseURL) + if err != nil { + return "", fmt.Errorf("invalid base URL %q: %w", raw, err) + } + switch parsed.Scheme { + case "http", "https": + default: + return "", fmt.Errorf("base URL must use http or https, got %q", parsed.Scheme) + } + if parsed.Host == "" { + return "", fmt.Errorf("base URL must include a host") + } + return baseURL, nil +} + +func BaseURLIsInsecureHTTP(baseURL string) bool { + parsed, err := url.Parse(baseURL) + return err == nil && parsed.Scheme == "http" +} diff --git a/internal/config/base_url_test.go b/internal/config/base_url_test.go new file mode 100644 index 0000000..93f1f1c --- /dev/null +++ b/internal/config/base_url_test.go @@ -0,0 +1,46 @@ +package config + +import "testing" + +func TestNormalizeBaseURL_acceptsHTTPAndHTTPSSchemes(t *testing.T) { + tests := []struct { + name string + raw string + want string + }{ + {name: "https", raw: " https://api.example.com/v1/ ", want: "https://api.example.com/v1"}, + {name: "http", raw: "http://localhost:8080/", want: "http://localhost:8080"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := NormalizeBaseURL(tt.raw) + if err != nil { + t.Fatalf("NormalizeBaseURL(%q): %v", tt.raw, err) + } + if got != tt.want { + t.Fatalf("NormalizeBaseURL(%q) = %q, want %q", tt.raw, got, tt.want) + } + }) + } +} + +func TestNormalizeBaseURL_rejectsInvalidURLs(t *testing.T) { + tests := []string{"", "not-a-url", "ftp://api.example.com", "https://", "https:///v1"} + for _, raw := range tests { + t.Run(raw, func(t *testing.T) { + if got, err := NormalizeBaseURL(raw); err == nil { + t.Fatalf("NormalizeBaseURL(%q) = %q, want error", raw, got) + } + }) + } +} + +func TestBaseURLIsInsecureHTTP(t *testing.T) { + if !BaseURLIsInsecureHTTP("http://localhost:8080/v1") { + t.Fatal("http base URL should be insecure") + } + if BaseURLIsInsecureHTTP("https://api.example.com/v1") { + t.Fatal("https base URL should not be insecure") + } +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go index ac29ea1..fffe7e7 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -105,7 +105,7 @@ func TestLoadMultiProviderTypeRoundTrip(t *testing.T) { ActiveModel: "openai/gpt-image-1", Providers: map[string]ProviderConfig{ "openrouter": { - Type: "openai-compatible", + Type: "openai", Model: "openai/gpt-image-1", BaseURL: "https://openrouter.ai/api/v1", Retries: 2, @@ -125,8 +125,8 @@ func TestLoadMultiProviderTypeRoundTrip(t *testing.T) { if loaded.SchemaVersion != CurrentSchemaVersion { t.Fatalf("SchemaVersion = %d, want %d", loaded.SchemaVersion, CurrentSchemaVersion) } - if got := loaded.Providers["openrouter"].Type; got != "openai-compatible" { - t.Fatalf("Type = %q, want openai-compatible", got) + if got := loaded.Providers["openrouter"].Type; got != "openai" { + t.Fatalf("Type = %q, want openai", got) } } @@ -176,8 +176,8 @@ providers: if got := loaded.Providers["openai"].Type; got != "openai" { t.Fatalf("openai Type = %q, want openai", got) } - if got := loaded.Providers["custom"].Type; got != "openai-compatible" { - t.Fatalf("custom Type = %q, want openai-compatible", got) + if got := loaded.Providers["custom"].Type; got != "openai" { + t.Fatalf("custom Type = %q, want openai", got) } } diff --git a/internal/config/migrate.go b/internal/config/migrate.go index 28ab655..fc26afe 100644 --- a/internal/config/migrate.go +++ b/internal/config/migrate.go @@ -8,10 +8,13 @@ import ( "time" ) +const VercelGatewayBaseURL = "https://ai-gateway.vercel.sh/v1" + type migrationFunc func(*MultiProviderConfig) bool var migrations = map[int]migrationFunc{ 2: migrateProviderTypes, + 3: migrateProviderTypes, } func MigrateConfigFile(path string, now func() time.Time) (bool, string, error) { @@ -56,16 +59,30 @@ func migrateConfig(cfg *MultiProviderConfig) bool { changed = true } } + if migrateProviderTypes(cfg) { + changed = true + } return changed } func migrateProviderTypes(cfg *MultiProviderConfig) bool { changed := false for name, pc := range cfg.Providers { - if pc.Type != "" { + providerType := ResolveProviderType(name, pc) + if name == "vercel" || providerType == "vercel" { + providerType = "openai" + if pc.BaseURL == "" { + pc.BaseURL = VercelGatewayBaseURL + } + } + if pc.Type == providerType { + if cfg.Providers[name].BaseURL != pc.BaseURL { + cfg.Providers[name] = pc + changed = true + } continue } - pc.Type = ResolveProviderType(name, pc) + pc.Type = providerType cfg.Providers[name] = pc changed = true } @@ -74,10 +91,13 @@ func migrateProviderTypes(cfg *MultiProviderConfig) bool { func ResolveProviderType(providerName string, pc ProviderConfig) string { if pc.Type != "" { + if pc.Type == "openai-compatible" { + return "openai" + } return pc.Type } if providerName == "custom" { - return "openai-compatible" + return "openai" } return providerName } @@ -87,7 +107,7 @@ func ProviderRequiresBaseURL(providerName, providerType string) bool { return true } switch providerName { - case "openai", "fal", "vercel": + case "openai", "fal", "replicate": return false default: return true @@ -96,11 +116,18 @@ func ProviderRequiresBaseURL(providerName, providerType string) bool { func AdapterType(providerType string) string { if providerType == "openai-compatible" { - return "custom" + return "openai" } return providerType } +func AdapterTypeForProvider(providerName, providerType string) string { + if providerName == "custom" { + return "custom" + } + return AdapterType(providerType) +} + func backupConfig(path string, now func() time.Time) (string, error) { data, err := os.ReadFile(path) if err != nil { diff --git a/internal/config/provider_type_migrate_test.go b/internal/config/provider_type_migrate_test.go new file mode 100644 index 0000000..ea72bcd --- /dev/null +++ b/internal/config/provider_type_migrate_test.go @@ -0,0 +1,117 @@ +package config + +import ( + "fmt" + "os" + "path/filepath" + "testing" + "time" +) + +func TestMigrateConfigFileRewritesOpenAICompatibleTypes(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + content := []byte(` +schema_version: 1 +active_provider: openrouter +providers: + openrouter: + type: openai-compatible + base_url: https://openrouter.ai/api/v1 + model: openai/gpt-image-2 +`) + if err := os.WriteFile(path, content, 0600); err != nil { + t.Fatalf("write config: %v", err) + } + now := func() time.Time { return time.Date(2026, 7, 2, 23, 55, 0, 0, time.UTC) } + + changed, _, err := MigrateConfigFile(path, now) + if err != nil { + t.Fatalf("MigrateConfigFile: %v", err) + } + if !changed { + t.Fatal("changed = false, want true") + } + + loaded, err := LoadMultiProvider(path) + if err != nil { + t.Fatalf("LoadMultiProvider: %v", err) + } + if got := loaded.Providers["openrouter"].Type; got != "openai" { + t.Fatalf("openrouter Type = %q, want openai", got) + } +} + +func TestMigrateConfigFileRewritesCurrentSchemaOpenAICompatibleTypes(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + content := []byte(fmt.Sprintf(` +schema_version: %d +active_provider: openrouter +providers: + openrouter: + type: openai-compatible + base_url: https://openrouter.ai/api/v1 + model: openai/gpt-image-2 +`, CurrentSchemaVersion)) + if err := os.WriteFile(path, content, 0600); err != nil { + t.Fatalf("write config: %v", err) + } + now := func() time.Time { return time.Date(2026, 7, 2, 23, 58, 0, 0, time.UTC) } + + changed, _, err := MigrateConfigFile(path, now) + if err != nil { + t.Fatalf("MigrateConfigFile: %v", err) + } + if !changed { + t.Fatal("changed = false, want true") + } + + loaded, err := LoadMultiProvider(path) + if err != nil { + t.Fatalf("LoadMultiProvider: %v", err) + } + if loaded.SchemaVersion != CurrentSchemaVersion { + t.Fatalf("SchemaVersion = %d, want %d", loaded.SchemaVersion, CurrentSchemaVersion) + } + if got := loaded.Providers["openrouter"].Type; got != "openai" { + t.Fatalf("openrouter Type = %q, want openai", got) + } +} + +func TestMigrateConfigFileRewritesVercelToOpenAICompatibleAlias(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + content := []byte(fmt.Sprintf(` +schema_version: %d +active_provider: vercel +providers: + vercel: + type: vercel + model: openai/gpt-image-2 +`, CurrentSchemaVersion)) + if err := os.WriteFile(path, content, 0600); err != nil { + t.Fatalf("write config: %v", err) + } + now := func() time.Time { return time.Date(2026, 7, 3, 10, 0, 0, 0, time.UTC) } + + changed, _, err := MigrateConfigFile(path, now) + if err != nil { + t.Fatalf("MigrateConfigFile: %v", err) + } + if !changed { + t.Fatal("changed = false, want true") + } + + loaded, err := LoadMultiProvider(path) + if err != nil { + t.Fatalf("LoadMultiProvider: %v", err) + } + pc := loaded.Providers["vercel"] + if pc.Type != "openai" { + t.Fatalf("vercel Type = %q, want openai", pc.Type) + } + if pc.BaseURL != "https://ai-gateway.vercel.sh/v1" { + t.Fatalf("vercel BaseURL = %q, want Vercel AI Gateway URL", pc.BaseURL) + } +} diff --git a/internal/config/types.go b/internal/config/types.go index 201fab4..ffacee5 100644 --- a/internal/config/types.go +++ b/internal/config/types.go @@ -1,6 +1,6 @@ package config -const CurrentSchemaVersion = 2 +const CurrentSchemaVersion = 3 // ProviderConfig holds per-provider settings in the multi-provider config format. // Timeout is stored in seconds (plain integer in YAML) so users do not need @@ -15,7 +15,46 @@ type ProviderConfig struct { } type ModelConfig struct { - Edit bool `yaml:"edit"` + Edit bool `yaml:"edit"` + DisplayName string `yaml:"display_name,omitempty"` + Custom bool `yaml:"custom,omitempty"` + Parameters ModelParameters `yaml:"parameters,omitempty"` +} + +type ModelParameters struct { + Size string `yaml:"size,omitempty"` + Quality string `yaml:"quality,omitempty"` + N int `yaml:"n,omitempty"` + ResponseFormat string `yaml:"response_format,omitempty"` + Seed int `yaml:"seed,omitempty"` + GuidanceScale float64 `yaml:"guidance_scale,omitempty"` + NegativePrompt string `yaml:"negative_prompt,omitempty"` +} + +func DefaultModelParameters() ModelParameters { + return ModelParameters{ + Size: "1024x1024", + Quality: "auto", + N: 1, + ResponseFormat: "b64_json", + } +} + +func (p ModelParameters) WithDefaults() ModelParameters { + defaults := DefaultModelParameters() + if p.Size == "" { + p.Size = defaults.Size + } + if p.Quality == "" { + p.Quality = defaults.Quality + } + if p.N == 0 { + p.N = defaults.N + } + if p.ResponseFormat == "" { + p.ResponseFormat = defaults.ResponseFormat + } + return p } // MultiProviderConfig is the v2 config format supporting multiple providers diff --git a/internal/credential/store.go b/internal/credential/store.go index 1abd52f..6351006 100644 --- a/internal/credential/store.go +++ b/internal/credential/store.go @@ -111,6 +111,12 @@ func (s *CredentialStore) Set(provider, apiKey string) error { return s.save() } +// SetCredential restores a complete credential record for the named provider. +func (s *CredentialStore) SetCredential(provider string, cred ProviderCredential) error { + s.data.Providers[provider] = cred + return s.save() +} + // Remove deletes the named provider from the credential store and persists. func (s *CredentialStore) Remove(provider string) error { if _, ok := s.data.Providers[provider]; !ok { diff --git a/internal/credential/store_test.go b/internal/credential/store_test.go index cd37b82..363411d 100644 --- a/internal/credential/store_test.go +++ b/internal/credential/store_test.go @@ -87,7 +87,7 @@ func TestStoreList(t *testing.T) { store.Set("openai", "sk-1") store.Set("fal", "sk-2") - store.Set("vercel", "sk-3") + store.Set("replicate", "r8-3") providers := store.List() if len(providers) != 3 { @@ -98,7 +98,7 @@ func TestStoreList(t *testing.T) { for _, p := range providers { found[p] = true } - if !found["openai"] || !found["fal"] || !found["vercel"] { + if !found["openai"] || !found["fal"] || !found["replicate"] { t.Errorf("List missing providers, got: %v", providers) } } diff --git a/internal/stats/openrouter.go b/internal/stats/openrouter.go new file mode 100644 index 0000000..d24503f --- /dev/null +++ b/internal/stats/openrouter.go @@ -0,0 +1,70 @@ +package stats + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "strconv" + "time" +) + +const defaultOpenRouterModelsURL = "https://openrouter.ai/api/v1/models?output_modalities=image" + +func FetchOpenRouterPrices(ctx context.Context, client *http.Client) (map[string]Price, error) { + if client == nil { + client = &http.Client{Timeout: 3 * time.Second} + } + url := defaultOpenRouterModelsURL + if override := os.Getenv("POTACO_OPENROUTER_MODELS_URL"); override != "" { + url = override + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, fmt.Errorf("create OpenRouter pricing request: %w", err) + } + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("fetch OpenRouter pricing: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode >= 400 { + return nil, fmt.Errorf("fetch OpenRouter pricing failed (HTTP %d)", resp.StatusCode) + } + + var result struct { + Data []struct { + ID string `json:"id"` + Name string `json:"name"` + Pricing struct { + Request string `json:"request"` + Image string `json:"image"` + } `json:"pricing"` + } `json:"data"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("decode OpenRouter pricing: %w", err) + } + prices := make(map[string]Price, len(result.Data)) + for _, model := range result.Data { + if model.ID == "" { + continue + } + prices[model.ID] = Price{ + ID: model.ID, + Name: model.Name, + RequestUSD: parsePrice(model.Pricing.Request), + ImageUSD: parsePrice(model.Pricing.Image), + } + } + return prices, nil +} + +func parsePrice(value string) float64 { + n, err := strconv.ParseFloat(value, 64) + if err != nil { + return 0 + } + return n +} diff --git a/internal/stats/stats.go b/internal/stats/stats.go new file mode 100644 index 0000000..aea663b --- /dev/null +++ b/internal/stats/stats.go @@ -0,0 +1,234 @@ +package stats + +import ( + "bufio" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "time" +) + +type Event struct { + Timestamp time.Time `json:"timestamp"` + Command string `json:"command"` + Operation string `json:"operation,omitempty"` + Provider string `json:"provider,omitempty"` + Model string `json:"model,omitempty"` + DurationMS int64 `json:"duration_ms"` + Success bool `json:"success"` + Error string `json:"error,omitempty"` + Outputs int `json:"outputs,omitempty"` + DryRun bool `json:"dry_run,omitempty"` +} + +type Price struct { + ID string + Name string + RequestUSD float64 + ImageUSD float64 +} + +type Summary struct { + Total int + Successes int + Errors int + AvgDuration time.Duration + CostUSD float64 + ByCommand []Bucket + ByProvider []Bucket + Recent []Event +} + +type Bucket struct { + Name string + Count int + Errors int + DurationSum time.Duration + CostUSD float64 +} + +func DefaultPath() string { + home, err := os.UserHomeDir() + if err != nil { + return filepath.Join(".potaco", "stats.jsonl") + } + return filepath.Join(home, ".potaco", "stats.jsonl") +} + +func Append(path string, event Event) error { + if event.Timestamp.IsZero() { + event.Timestamp = time.Now() + } + if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil { + return fmt.Errorf("create stats dir: %w", err) + } + f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600) + if err != nil { + return fmt.Errorf("open stats: %w", err) + } + defer f.Close() + if err := json.NewEncoder(f).Encode(event); err != nil { + return fmt.Errorf("write stats: %w", err) + } + return nil +} + +func Load(path string) ([]Event, error) { + f, err := os.Open(path) + if os.IsNotExist(err) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("open stats: %w", err) + } + defer f.Close() + + var events []Event + scanner := bufio.NewScanner(f) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + var event Event + if err := json.Unmarshal([]byte(line), &event); err != nil { + return nil, fmt.Errorf("decode stats event: %w", err) + } + events = append(events, event) + } + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("read stats: %w", err) + } + return events, nil +} + +func Clear(path string) error { + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("clear stats: %w", err) + } + return nil +} + +func Summarize(events []Event, prices map[string]Price) Summary { + summary := Summary{Total: len(events)} + commandBuckets := map[string]*Bucket{} + providerBuckets := map[string]*Bucket{} + var durationSum time.Duration + for _, event := range events { + if event.Success { + summary.Successes++ + } else { + summary.Errors++ + } + duration := time.Duration(event.DurationMS) * time.Millisecond + durationSum += duration + cost := EstimateCost(event, prices) + summary.CostUSD += cost + addBucket(commandBuckets, event.Command, event, duration, cost) + addBucket(providerBuckets, providerBucketName(event.Provider), event, duration, cost) + } + if summary.Total > 0 { + summary.AvgDuration = durationSum / time.Duration(summary.Total) + } + summary.ByCommand = sortedBuckets(commandBuckets) + summary.ByProvider = sortedBuckets(providerBuckets) + summary.Recent = recentEvents(events, 10) + return summary +} + +func EstimateCost(event Event, prices map[string]Price) float64 { + if event.DryRun || !event.Success || event.Model == "" { + return 0 + } + price, ok := matchPrice(event.Model, prices) + if !ok { + return 0 + } + outputs := event.Outputs + if outputs <= 0 { + outputs = 1 + } + return price.RequestUSD + price.ImageUSD*float64(outputs) +} + +func matchPrice(model string, prices map[string]Price) (Price, bool) { + if prices == nil { + return Price{}, false + } + if price, ok := prices[model]; ok { + return price, true + } + suffix := modelSuffix(model) + for id, price := range prices { + if modelSuffix(id) == suffix { + return price, true + } + } + return Price{}, false +} + +func addBucket(buckets map[string]*Bucket, name string, event Event, duration time.Duration, cost float64) { + if name == "" { + name = "unknown" + } + bucket := buckets[name] + if bucket == nil { + bucket = &Bucket{Name: name} + buckets[name] = bucket + } + bucket.Count++ + if !event.Success { + bucket.Errors++ + } + bucket.DurationSum += duration + bucket.CostUSD += cost +} + +func sortedBuckets(buckets map[string]*Bucket) []Bucket { + out := make([]Bucket, 0, len(buckets)) + for _, bucket := range buckets { + out = append(out, *bucket) + } + sort.Slice(out, func(i, j int) bool { + if out[i].Count == out[j].Count { + return out[i].Name < out[j].Name + } + return out[i].Count > out[j].Count + }) + return out +} + +func recentEvents(events []Event, limit int) []Event { + if len(events) <= limit { + out := append([]Event(nil), events...) + reverseEvents(out) + return out + } + out := append([]Event(nil), events[len(events)-limit:]...) + reverseEvents(out) + return out +} + +func reverseEvents(events []Event) { + for i, j := 0, len(events)-1; i < j; i, j = i+1, j-1 { + events[i], events[j] = events[j], events[i] + } +} + +func providerBucketName(provider string) string { + if provider == "" { + return "unknown" + } + return provider +} + +func modelSuffix(model string) string { + model = strings.ToLower(strings.TrimSpace(model)) + if idx := strings.LastIndex(model, "/"); idx >= 0 { + return model[idx+1:] + } + return model +} diff --git a/internal/stats/stats_test.go b/internal/stats/stats_test.go new file mode 100644 index 0000000..fc11b42 --- /dev/null +++ b/internal/stats/stats_test.go @@ -0,0 +1,44 @@ +package stats + +import ( + "path/filepath" + "testing" + "time" +) + +func TestAppendLoadAndSummarize(t *testing.T) { + path := filepath.Join(t.TempDir(), "stats.jsonl") + events := []Event{ + {Timestamp: time.Unix(1, 0), Command: "potaco gen", Operation: "generate", Provider: "openai", Model: "gpt-image-2", DurationMS: 1000, Success: true, Outputs: 2}, + {Timestamp: time.Unix(2, 0), Command: "potaco edit", Operation: "edit", Provider: "openai", Model: "gpt-image-2", DurationMS: 3000, Success: false, Error: "api"}, + } + for _, event := range events { + if err := Append(path, event); err != nil { + t.Fatalf("append: %v", err) + } + } + + loaded, err := Load(path) + if err != nil { + t.Fatalf("load: %v", err) + } + if len(loaded) != 2 { + t.Fatalf("loaded events = %d, want 2", len(loaded)) + } + + summary := Summarize(loaded, map[string]Price{ + "openai/gpt-image-2": {RequestUSD: 0.01, ImageUSD: 0.02}, + }) + if summary.Total != 2 || summary.Successes != 1 || summary.Errors != 1 { + t.Fatalf("summary counts = %+v, want total=2 success=1 errors=1", summary) + } + if summary.AvgDuration != 2*time.Second { + t.Fatalf("avg duration = %v, want 2s", summary.AvgDuration) + } + if summary.CostUSD != 0.05 { + t.Fatalf("cost = %v, want 0.05", summary.CostUSD) + } + if len(summary.Recent) != 2 || summary.Recent[0].Command != "potaco edit" { + t.Fatalf("recent = %+v, want newest first", summary.Recent) + } +} diff --git a/internal/tui/auth_add.go b/internal/tui/auth_add.go index c1a374e..444a3cb 100644 --- a/internal/tui/auth_add.go +++ b/internal/tui/auth_add.go @@ -4,20 +4,15 @@ import ( "context" "errors" "fmt" + "slices" "strings" - "github.com/charmbracelet/huh" - "github.com/ncxton/potaco/internal/adapter" "github.com/ncxton/potaco/internal/auth" "github.com/ncxton/potaco/internal/config" ) -// errCancelled is returned by form helpers when the user aborts a TUI flow. -// It is normalized to a nil error by the callers so the command exits cleanly. -var errCancelled = errors.New("cancelled") - -// RunAuthAdd launches the interactive auth add flow using huh forms. +// RunAuthAdd launches the interactive auth add flow. // It prompts for an API key, verifies the provider, discovers models, // lets the user pick a model, and stores the credential. // If providerName is empty, shows a provider picker first. @@ -44,32 +39,36 @@ func RunAuthAdd(providerName string) error { return normalizeCancel(err) } - ad, err := adapter.Get(config.AdapterType(providerType), apiKey, adapter.AdapterOpts{BaseURL: baseURL}) + ad, err := adapter.Get(config.AdapterTypeForProvider(providerName, providerType), apiKey, adapter.AdapterOpts{BaseURL: baseURL}) if err != nil { return fmt.Errorf("create adapter: %w", err) } + var selected modelSelection verifyErr := ad.Verify(context.Background()) if verifyErr != nil { - proceed, err := confirmVerification(verifyErr) + printUnableToFindModels() + selected, err = promptCustomModel() if err != nil { return normalizeCancel(err) } - if !proceed { - return fmt.Errorf("cancelled by user") + } else { + models, discoverErr := ad.DiscoverModels(context.Background()) + if discoverErr != nil { + printUnableToFindModels() + selected, err = promptCustomModel() + if err != nil { + return normalizeCancel(err) + } + } else { + selected, err = promptModel(providerName, models) + if err != nil { + return normalizeCancel(err) + } } } - selected := modelSelection{} - models, discoverErr := ad.DiscoverModels(context.Background()) - if discoverErr == nil { - selected, err = promptModel(providerName, models) - if err != nil { - return normalizeCancel(err) - } - } - - return addProvider(providerName, providerType, apiKey, baseURL, selected.ID, selected.Edit) + return addProviderSelection(providerName, providerType, apiKey, baseURL, selected) } // normalizeCancel converts errCancelled into a nil error so the CLI exits @@ -100,18 +99,19 @@ func promptProvider() (string, error) { if len(available) == 0 { return "", fmt.Errorf("no providers available") } - options := make([]huh.Option[string], 0, len(available)) + options := make([]promptOption, 0, len(available)) for _, name := range available { - options = append(options, huh.NewOption(name, name)) - } - var providerName string - form := newForm(huh.NewGroup( - huh.NewSelect[string](). - Title("Select a provider to connect:"). - Options(options...). - Value(&providerName), - )) - if err := runForm(form, "provider select"); err != nil { + options = append(options, promptOption{Label: name, Value: name}) + } + providerName, err := runSelectPrompt("Select a provider to connect:", options) + if err != nil { + if isCancelled(err) { + printCancelled() + return "", errCancelled + } + return "", fmt.Errorf("provider select: %w", err) + } + if providerName == "" { return "", err } return providerName, nil @@ -119,35 +119,30 @@ func promptProvider() (string, error) { // isKnownProvider reports whether name is a registered provider. func isKnownProvider(name string) bool { - for _, n := range adapter.List() { - if n == name { - return true - } - } - return false + return slices.Contains(adapter.List(), name) } func ensureProviderType(providerName string) (string, error) { if providerName == "custom" { - return "openai-compatible", nil + return "openai", nil } if isKnownProvider(providerName) { return providerName, nil } - var providerType string - options := []huh.Option[string]{ - huh.NewOption("OpenAI-compatible", "openai-compatible"), - huh.NewOption("OpenAI", "openai"), - huh.NewOption("fal", "fal"), - huh.NewOption("Vercel AI Gateway", "vercel"), - } - form := newForm(huh.NewGroup( - huh.NewSelect[string](). - Title("Select provider type:"). - Options(options...). - Value(&providerType), - )) - if err := runForm(form, "provider type select"); err != nil { + options := []promptOption{ + {Label: "OpenAI", Value: "openai"}, + {Label: "Fal", Value: "fal"}, + {Label: "Replicate", Value: "replicate"}, + } + providerType, err := runSelectPrompt("Select provider type:", options) + if err != nil { + if isCancelled(err) { + printCancelled() + return "", errCancelled + } + return "", fmt.Errorf("provider type select: %w", err) + } + if providerType == "" { return "", err } return providerType, nil @@ -163,82 +158,91 @@ func promptBaseURL(providerName, providerType string) (string, error) { if !authAddRequiresBaseURL(providerName, providerType) { return "", nil } - var baseURL string - form := newForm(huh.NewGroup( - huh.NewInput(). - Title("Enter base URL for the provider:"). - Value(&baseURL), - )) - if err := runForm(form, "base URL input"); err != nil { + baseURL, err := runInputPrompt("Enter base URL for the provider:", "", false) + if err != nil { + if isCancelled(err) { + printCancelled() + return "", errCancelled + } + return "", fmt.Errorf("base URL input: %w", err) + } + baseURL, err = config.NormalizeBaseURL(baseURL) + if err != nil { return "", err } - baseURL = strings.TrimRight(baseURL, "/") - if baseURL == "" { - return "", fmt.Errorf("base URL is required for this provider") + if config.BaseURLIsInsecureHTTP(baseURL) { + trusted, err := confirmInsecureBaseURL(baseURL) + if err != nil { + return "", err + } + if !trusted { + printCancelled() + return "", errCancelled + } } return baseURL, nil } -// promptAPIKey prompts for the provider API key. -func promptAPIKey(providerName string) (string, error) { - var apiKey string - form := newForm(huh.NewGroup( - huh.NewInput(). - Title(fmt.Sprintf("Enter API key for %s:", providerName)). - EchoMode(huh.EchoModePassword). - Value(&apiKey), - )) - if err := runForm(form, "key input"); err != nil { - return "", err - } - if apiKey == "" { - return "", fmt.Errorf("API key is required") +func confirmInsecureBaseURL(baseURL string) (bool, error) { + trusted, err := runConfirmPrompt(fmt.Sprintf(" Danger: %s uses unencrypted HTTP. Continue only if you trust this URL?", baseURL)) + if err != nil { + if isCancelled(err) { + printCancelled() + return false, errCancelled + } + return false, err } - return apiKey, nil + return trusted, nil } -// runForm runs a huh form and maps user aborts to errCancelled. -func runForm(form *huh.Form, label string) error { - if err := form.Run(); err != nil { +// promptAPIKey prompts for the provider API key. +func promptAPIKey(providerName string) (string, error) { + apiKey, err := runInputPrompt(fmt.Sprintf("Enter API key for %s:", providerName), "", true) + if err != nil { if isCancelled(err) { - fmt.Println("Cancelled.") - return errCancelled + printCancelled() + return "", errCancelled } - return fmt.Errorf("%s: %w", label, err) + return "", fmt.Errorf("key input: %w", err) } - return nil -} - -// confirmVerification asks whether to add a provider after verification failed. -func confirmVerification(verifyErr error) (bool, error) { - var proceed bool - form := newForm(huh.NewGroup( - huh.NewConfirm(). - Title(fmt.Sprintf("Verification failed: %s\nAdd anyway?", verifyErr)). - Value(&proceed), - )) - if err := runForm(form, "confirm"); err != nil { - return false, err + if apiKey == "" { + return "", fmt.Errorf("API key is required") } - return proceed, nil + return apiKey, nil } // promptModel shows a model picker when discovery succeeds. func promptModel(providerName string, models []adapter.Model) (modelSelection, error) { if len(models) == 0 { - return modelSelection{}, nil + printUnableToFindModels() + return modelSelection{}, errCancelled } return PickModel(providerName, models) } +func promptManualModel() (string, error) { + modelID, err := runInputPrompt("Enter model ID:", "", false) + if err != nil { + if isCancelled(err) { + printCancelled() + return "", errCancelled + } + return "", fmt.Errorf("model input: %w", err) + } + modelID = strings.TrimSpace(modelID) + if modelID == "" { + return "", fmt.Errorf("model ID is required") + } + return modelID, nil +} + func promptModelEditCapable(modelID string) (bool, error) { - var edit bool - form := newForm(huh.NewGroup( - huh.NewConfirm(). - Title(fmt.Sprintf("Can %s edit images?", modelID)). - Value(&edit), - )) - if err := runForm(form, "model edit capability"); err != nil { + edit, err := runConfirmPrompt(fmt.Sprintf("Can %s edit images?", modelID)) + if err != nil { + if isCancelled(err) { + printCancelled() + return false, errCancelled + } return false, err } return edit, nil @@ -246,26 +250,32 @@ func promptModelEditCapable(modelID string) (bool, error) { // addProvider stores the credential and config for the connected provider. func addProvider(providerName, providerType, apiKey, baseURL, modelID string, edit bool) error { + return addProviderSelection(providerName, providerType, apiKey, baseURL, modelSelection{ID: modelID, Edit: edit}) +} + +func addProviderSelection(providerName, providerType, apiKey, baseURL string, selected modelSelection) error { mgr, err := auth.New() if err != nil { return fmt.Errorf("init auth: %w", err) } - if err := mgr.AddProvider(providerName, providerType, apiKey); err != nil { - return fmt.Errorf("add provider: %w", err) - } - if baseURL != "" { - if err := mgr.SetBaseURL(providerName, baseURL); err != nil { - return fmt.Errorf("set base URL: %w", err) + setup := auth.ProviderSetup{ + ProviderName: providerName, + ProviderType: providerType, + APIKey: apiKey, + BaseURL: baseURL, + ModelID: selected.ID, + } + if selected.ID != "" { + setup.ModelConfig = config.ModelConfig{ + Edit: selected.Edit, + DisplayName: selected.DisplayName, + Custom: selected.Custom, + Parameters: selected.Parameters, } } - if modelID != "" { - if err := mgr.SetActiveProvider(providerName, modelID); err != nil { - return fmt.Errorf("set model: %w", err) - } - if err := mgr.SetModelEdit(providerName, modelID, edit); err != nil { - return fmt.Errorf("set model edit capability: %w", err) - } + if err := mgr.AddProviderSetup(setup); err != nil { + return fmt.Errorf("add provider: %w", err) } - fmt.Printf("Provider '%s' added successfully.\n", providerName) + printSuccess("Provider '%s' added successfully.", providerName) return nil } diff --git a/internal/tui/auth_add_base_url_test.go b/internal/tui/auth_add_base_url_test.go new file mode 100644 index 0000000..9272ea9 --- /dev/null +++ b/internal/tui/auth_add_base_url_test.go @@ -0,0 +1,114 @@ +package tui + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func trustHTTPBaseURLs(t *testing.T) { + t.Helper() + oldConfirmPrompt := runConfirmPrompt + t.Cleanup(func() { runConfirmPrompt = oldConfirmPrompt }) + runConfirmPrompt = func(string) (bool, error) { return true, nil } +} + +func TestPromptBaseURLRejectsInvalidURL(t *testing.T) { + oldInputPrompt := runInputPrompt + t.Cleanup(func() { runInputPrompt = oldInputPrompt }) + + runInputPrompt = func(label, placeholder string, secret bool) (string, error) { + return "not-a-url", nil + } + + if got, err := promptBaseURL("custom", "openai"); err == nil { + t.Fatalf("promptBaseURL = %q, want invalid URL error", got) + } +} + +func TestPromptBaseURLConfirmsInsecureHTTP(t *testing.T) { + oldInputPrompt := runInputPrompt + oldConfirmPrompt := runConfirmPrompt + t.Cleanup(func() { + runInputPrompt = oldInputPrompt + runConfirmPrompt = oldConfirmPrompt + }) + + runInputPrompt = func(label, placeholder string, secret bool) (string, error) { + return "http://localhost:8080/v1/", nil + } + var confirmLabel string + runConfirmPrompt = func(label string) (bool, error) { + confirmLabel = label + return true, nil + } + + got, err := promptBaseURL("custom", "openai") + if err != nil { + t.Fatalf("promptBaseURL: %v", err) + } + if got != "http://localhost:8080/v1" { + t.Fatalf("base URL = %q, want normalized HTTP URL", got) + } + if !strings.Contains(confirmLabel, "") || !strings.Contains(confirmLabel, "Danger") || !strings.Contains(confirmLabel, "http://localhost:8080/v1") { + t.Fatalf("confirm prompt = %q, want warning icon, danger text, and URL", confirmLabel) + } +} + +func TestPromptBaseURLCancelsInsecureHTTPWhenNotTrusted(t *testing.T) { + oldInputPrompt := runInputPrompt + oldConfirmPrompt := runConfirmPrompt + t.Cleanup(func() { + runInputPrompt = oldInputPrompt + runConfirmPrompt = oldConfirmPrompt + }) + + runInputPrompt = func(label, placeholder string, secret bool) (string, error) { + return "http://localhost:8080/v1", nil + } + runConfirmPrompt = func(label string) (bool, error) { + return false, nil + } + + if got, err := promptBaseURL("custom", "openai"); !isCancelled(err) { + t.Fatalf("promptBaseURL = %q, %v; want cancelled", got, err) + } +} + +func TestRunAuthAddCancelsManualModelWithoutWritingConfigOrCredentials(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + oldInputPrompt := runInputPrompt + oldConfirmPrompt := runConfirmPrompt + defer func() { + runInputPrompt = oldInputPrompt + runConfirmPrompt = oldConfirmPrompt + }() + + runInputPrompt = func(label, placeholder string, secret bool) (string, error) { + switch { + case strings.Contains(label, "API key"): + return "r8_test", nil + case strings.Contains(label, "model"): + return "", nil + default: + return "", nil + } + } + runConfirmPrompt = func(label string) (bool, error) { + t.Fatalf("RunAuthAdd should not confirm failed verification; prompt = %q", label) + return false, nil + } + + err := RunAuthAdd("replicate") + if err == nil { + t.Fatal("RunAuthAdd should error when manual model ID is empty") + } + if _, statErr := os.Stat(filepath.Join(home, ".potaco", "config.yaml")); !os.IsNotExist(statErr) { + t.Fatalf("config should not be written, stat err = %v", statErr) + } + if _, statErr := os.Stat(filepath.Join(home, ".potaco", "credentials.enc")); !os.IsNotExist(statErr) { + t.Fatalf("credentials should not be written, stat err = %v", statErr) + } +} diff --git a/internal/tui/auth_add_test.go b/internal/tui/auth_add_test.go index a05200d..1b93e6f 100644 --- a/internal/tui/auth_add_test.go +++ b/internal/tui/auth_add_test.go @@ -1,6 +1,9 @@ package tui import ( + "io" + "os" + "strings" "testing" "github.com/ncxton/potaco/internal/adapter" @@ -35,13 +38,13 @@ func TestAuthAddRequiresBaseURLForAliases(t *testing.T) { providerType string want bool }{ - {name: "custom", providerName: "custom", providerType: "openai-compatible", want: true}, + {name: "custom", providerName: "custom", providerType: "openai", want: true}, {name: "openai alias", providerName: "staging-openai", providerType: "openai", want: true}, {name: "fal alias", providerName: "staging-fal", providerType: "fal", want: true}, - {name: "vercel alias", providerName: "staging-vercel", providerType: "vercel", want: true}, + {name: "replicate alias", providerName: "staging-replicate", providerType: "replicate", want: true}, {name: "openai built-in", providerName: "openai", providerType: "openai", want: false}, {name: "fal built-in", providerName: "fal", providerType: "fal", want: false}, - {name: "vercel built-in", providerName: "vercel", providerType: "vercel", want: false}, + {name: "replicate built-in", providerName: "replicate", providerType: "replicate", want: false}, } for _, tt := range tests { @@ -53,24 +56,98 @@ func TestAuthAddRequiresBaseURLForAliases(t *testing.T) { } } -func TestEnsureProviderTypeReturnsOpenAICompatibleForCustomProvider(t *testing.T) { +func TestEnsureProviderTypeReturnsOpenAIForCustomProvider(t *testing.T) { providerType, err := ensureProviderType("custom") if err != nil { t.Fatalf("ensureProviderType: %v", err) } - if providerType != "openai-compatible" { - t.Fatalf("provider type = %q, want openai-compatible", providerType) + if providerType != "openai" { + t.Fatalf("provider type = %q, want openai", providerType) } } -func TestPromptModelReturnsEmptyWhenNoModels(t *testing.T) { - selected, err := promptModel("openai", []adapter.Model{}) +func TestPromptModelCancelsWhenNoModels(t *testing.T) { + oldStdout := os.Stdout + r, w, err := os.Pipe() if err != nil { - t.Fatalf("promptModel: %v", err) + t.Fatalf("pipe: %v", err) + } + os.Stdout = w + t.Cleanup(func() { os.Stdout = oldStdout }) + + selected, err := promptModel("openai", []adapter.Model{}) + if closeErr := w.Close(); closeErr != nil { + t.Fatalf("close writer: %v", closeErr) + } + out, readErr := io.ReadAll(r) + if readErr != nil { + t.Fatalf("read output: %v", readErr) + } + if err == nil { + t.Fatal("promptModel should cancel when no image models are available") + } + if !isCancelled(err) { + t.Fatalf("promptModel error = %v, want cancelled", err) } if selected.ID != "" { t.Fatalf("model ID = %q, want empty", selected.ID) } + if !strings.Contains(string(out), " Unable to find any image generation models.") { + t.Fatalf("empty discovery output = %q, want unable message", string(out)) + } +} + +func TestRunAuthAddFallsBackToManualModelWhenDiscoveryFails(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + oldInputPrompt := runInputPrompt + oldConfirmPrompt := runConfirmPrompt + defer func() { + runInputPrompt = oldInputPrompt + runConfirmPrompt = oldConfirmPrompt + }() + + runInputPrompt = func(label, placeholder string, secret bool) (string, error) { + switch { + case strings.Contains(label, "base URL"): + return "", nil + case strings.Contains(label, "API key"): + return "r8_test", nil + case strings.Contains(label, "model"): + return "black-forest-labs/flux-schnell", nil + default: + return "", nil + } + } + runConfirmPrompt = func(label string) (bool, error) { + if strings.Contains(label, "add anyway") { + t.Fatalf("RunAuthAdd should not ask whether to add after verification failure; prompt = %q", label) + } + if strings.Contains(label, "edit images") { + return true, nil + } + return false, nil + } + + err := RunAuthAdd("replicate") + if err != nil { + t.Fatalf("RunAuthAdd replicate: %v", err) + } + + mgr, err := auth.New() + if err != nil { + t.Fatalf("init auth: %v", err) + } + cfg, err := mgr.LoadConfig() + if err != nil { + t.Fatalf("load config: %v", err) + } + pc, ok := cfg.Providers["replicate"] + if !ok { + t.Fatal("replicate provider should be configured") + } + if pc.Model != "black-forest-labs/flux-schnell" { + t.Fatalf("model = %q, want black-forest-labs/flux-schnell", pc.Model) + } } func TestAddProviderStoresProviderTypeAndBaseURL(t *testing.T) { @@ -78,7 +155,7 @@ func TestAddProviderStoresProviderTypeAndBaseURL(t *testing.T) { err := addProvider( "openrouter", - "openai-compatible", + "openai", "sk-test", "https://openrouter.ai/api/v1", "openrouter/image-model", @@ -100,8 +177,8 @@ func TestAddProviderStoresProviderTypeAndBaseURL(t *testing.T) { if !ok { t.Fatal("openrouter provider should be configured") } - if pc.Type != "openai-compatible" { - t.Fatalf("provider type = %q, want openai-compatible", pc.Type) + if pc.Type != "openai" { + t.Fatalf("provider type = %q, want openai", pc.Type) } if pc.BaseURL != "https://openrouter.ai/api/v1" { t.Fatalf("base URL = %q, want https://openrouter.ai/api/v1", pc.BaseURL) diff --git a/internal/tui/auth_remove.go b/internal/tui/auth_remove.go index 5e7026a..9e83dbe 100644 --- a/internal/tui/auth_remove.go +++ b/internal/tui/auth_remove.go @@ -3,8 +3,6 @@ package tui import ( "fmt" - "github.com/charmbracelet/huh" - "github.com/ncxton/potaco/internal/auth" ) @@ -25,41 +23,37 @@ func RunAuthRemove(providerName string) error { return nil } - options := make([]huh.Option[string], 0, len(providers)) + options := make([]promptOption, 0, len(providers)) for _, p := range providers { label := p.Name if p.IsActive { label += " (active)" } label += " - " + p.Model - options = append(options, huh.NewOption(label, p.Name)) + options = append(options, promptOption{Label: label, Value: p.Name}) } - selectForm := newForm(huh.NewGroup( - huh.NewSelect[string](). - Title("Select a provider to remove:"). - Options(options...). - Value(&providerName), - )) - if err := selectForm.Run(); err != nil { + selected, err := runSelectPrompt("Select a provider to remove:", options) + if err != nil { if isCancelled(err) { - fmt.Println("Cancelled.") + printCancelled() return nil } return fmt.Errorf("provider select: %w", err) } + providerName = selected } confirmed, err := ConfirmAction(fmt.Sprintf("Remove provider '%s' and its credentials?", providerName)) if err != nil { if isCancelled(err) { - fmt.Println("Cancelled.") + printCancelled() return nil } return fmt.Errorf("confirm: %w", err) } if !confirmed { - fmt.Println("Cancelled.") + printCancelled() return nil } @@ -67,6 +61,6 @@ func RunAuthRemove(providerName string) error { return fmt.Errorf("remove provider: %w", err) } - fmt.Printf("Provider '%s' removed.\n", providerName) + printSuccess("Provider '%s' removed.", providerName) return nil } diff --git a/internal/tui/handcrafted_test.go b/internal/tui/handcrafted_test.go new file mode 100644 index 0000000..cd8830a --- /dev/null +++ b/internal/tui/handcrafted_test.go @@ -0,0 +1,34 @@ +package tui + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestTUIDoesNotUsePrebuiltPromptComponents(t *testing.T) { + huhImport := `"github.com/charmbracelet/` + `huh"` + bubblesImport := `"github.com/charmbracelet/` + `bubbles/` + entries, err := os.ReadDir(".") + if err != nil { + t.Fatalf("read tui package: %v", err) + } + + for _, entry := range entries { + if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".go") { + continue + } + data, err := os.ReadFile(filepath.Clean(entry.Name())) + if err != nil { + t.Fatalf("read %s: %v", entry.Name(), err) + } + source := string(data) + if strings.Contains(source, huhImport) { + t.Fatalf("%s imports huh prebuilt prompts", entry.Name()) + } + if strings.Contains(source, bubblesImport) { + t.Fatalf("%s imports bubbles prebuilt components", entry.Name()) + } + } +} diff --git a/internal/tui/model_list.go b/internal/tui/model_list.go index 05ce6ae..9765890 100644 --- a/internal/tui/model_list.go +++ b/internal/tui/model_list.go @@ -3,10 +3,8 @@ package tui import ( "context" "fmt" - "os" - - tea "github.com/charmbracelet/bubbletea" - "github.com/charmbracelet/huh" + "net/http" + "strings" "github.com/ncxton/potaco/internal/adapter" "github.com/ncxton/potaco/internal/auth" @@ -14,8 +12,11 @@ import ( ) type modelSelection struct { - ID string - Edit bool + ID string + DisplayName string + Edit bool + Custom bool + Parameters config.ModelParameters } // modelPicker selects a model from the discovered list. It is abstracted so @@ -38,36 +39,19 @@ func runModelListWithPicker(providerName, apiKey, baseURL string, picker modelPi return fmt.Errorf("init auth: %w", err) } - if providerName == "" { - providerName, _, err = mgr.GetActiveProvider() - if err != nil || providerName == "" { - return fmt.Errorf("no active provider. Use 'potaco auth add ' to connect one") - } - } - - if apiKey == "" { - apiKey, err = mgr.GetAPIKey(providerName) - if err != nil { - return fmt.Errorf("provider %q is not connected. Use 'potaco auth add %s' first", providerName, providerName) - } - } - - cfg, _ := mgr.LoadConfig() - pc := config.ProviderConfig{} - if cfg != nil { - if configured, ok := cfg.Providers[providerName]; ok { - pc = configured - if baseURL == "" && configured.BaseURL != "" { - baseURL = configured.BaseURL - } + resolved, err := resolveModelListProvider(mgr, modelListRequest{ + ProviderName: providerName, + APIKey: apiKey, + BaseURL: baseURL, + }) + if err != nil { + if isCancelled(err) { + return nil } - } - providerType := config.ResolveProviderType(providerName, pc) - if config.ProviderRequiresBaseURL(providerName, providerType) && baseURL == "" { - return fmt.Errorf("base URL required for provider %s", providerName) + return err } - ad, err := adapter.Get(config.AdapterType(providerType), apiKey, adapter.AdapterOpts{BaseURL: baseURL}) + ad, err := adapter.Get(config.AdapterTypeForProvider(resolved.Name, resolved.ProviderType), resolved.APIKey, adapter.AdapterOpts{BaseURL: resolved.BaseURL}) if err != nil { return fmt.Errorf("create adapter: %w", err) } @@ -77,13 +61,13 @@ func runModelListWithPicker(providerName, apiKey, baseURL string, picker modelPi return fmt.Errorf("discover models: %w", err) } if len(models) == 0 { - return fmt.Errorf("no models found for %s", providerName) + return fmt.Errorf("no models found for %s", resolved.Name) } - selected, err := picker(providerName, models) + selected, err := picker(resolved.Name, models) if err != nil { if isCancelled(err) { - fmt.Println("Cancelled.") + printCancelled() return nil } return err @@ -92,41 +76,143 @@ func runModelListWithPicker(providerName, apiKey, baseURL string, picker modelPi return nil } - if err := mgr.SetActiveProvider(providerName, selected.ID); err != nil { + if err := mgr.SetActiveProvider(resolved.Name, selected.ID); err != nil { return fmt.Errorf("set active provider: %w", err) } - if err := mgr.SetModelEdit(providerName, selected.ID, selected.Edit); err != nil { + if err := mgr.SetModelEdit(resolved.Name, selected.ID, selected.Edit); err != nil { return fmt.Errorf("set model edit capability: %w", err) } - fmt.Printf("Switched to model '%s'.\n", selected.ID) + if selected.DisplayName != "" || selected.Custom { + if err := mgr.SetModelMetadata(resolved.Name, selected.ID, selected.DisplayName, selected.Custom); err != nil { + return fmt.Errorf("set model metadata: %w", err) + } + } + if selected.Parameters != (config.ModelParameters{}) { + if err := mgr.SetModelParameters(resolved.Name, selected.ID, selected.Parameters); err != nil { + return fmt.Errorf("set model parameters: %w", err) + } + } + printSuccess("Switched to model '%s'.", selected.ID) return nil } // PickModel renders a Bubble Tea search program listing the // discovered models and returns the selected model ID. The user can // type to filter the list in real-time and navigate with arrow keys. -// Returns huh.ErrUserAborted if the user cancels with Esc or Ctrl-C. +// Returns errCancelled if the user cancels with Esc or Ctrl-C. // Returns an empty string when the filter matches no models and the user // presses Enter. func PickModel(providerName string, models []adapter.Model) (modelSelection, error) { m := newSearchModel(providerName, models) - p := tea.NewProgram(m, tea.WithOutput(os.Stderr)) + p := newPromptProgram(m) result, err := p.Run() if err != nil { return modelSelection{}, fmt.Errorf("model search: %w", err) } if sm, ok := result.(*searchModel); ok { if sm.quitted { - return modelSelection{}, huh.ErrUserAborted + return modelSelection{}, errCancelled } if sm.selected == "" { return modelSelection{}, nil } + if sm.selected == customModelSentinel { + return promptCustomModel() + } edit, err := promptModelEditCapable(sm.selected) if err != nil { return modelSelection{}, err } - return modelSelection{ID: sm.selected, Edit: edit}, nil + selected := sm.selectedModel() + return modelSelection{ + ID: sm.selected, + DisplayName: selected.DisplayName, + Edit: edit, + }, nil } return modelSelection{}, fmt.Errorf("unexpected model type") } + +func promptCustomModel() (modelSelection, error) { + modelID, err := promptManualModel() + if err != nil { + return modelSelection{}, err + } + modelID, err = confirmOpenRouterModelSuggestion(modelID) + if err != nil { + return modelSelection{}, err + } + displayName, err := runInputPrompt("Enter display name:", modelID, false) + if err != nil { + if isCancelled(err) { + printCancelled() + return modelSelection{}, errCancelled + } + return modelSelection{}, fmt.Errorf("display name input: %w", err) + } + displayName = strings.TrimSpace(displayName) + if displayName == "" { + displayName = modelID + } + displayName = ensureCustomDisplayName(displayName) + + edit, err := promptModelEditCapable(modelID) + if err != nil { + return modelSelection{}, err + } + tweak, err := runConfirmPrompt("Tweak advanced settings?") + if err != nil { + if isCancelled(err) { + printCancelled() + return modelSelection{}, errCancelled + } + return modelSelection{}, err + } + if tweak { + params, err := promptModelParameters(config.DefaultModelParameters()) + if err != nil { + return modelSelection{}, err + } + return modelSelection{ID: modelID, DisplayName: displayName, Edit: edit, Custom: true, Parameters: params}, nil + } + + return modelSelection{ID: modelID, DisplayName: displayName, Edit: edit, Custom: true}, nil +} + +func confirmOpenRouterModelSuggestion(modelID string) (string, error) { + models, err := adapter.OpenRouterImageModels(context.Background(), http.DefaultClient) + if err != nil { + return modelID, nil + } + ids := make(map[string]bool, len(models)*2) + for _, model := range models { + ids[model.ID] = true + ids[adapter.SuffixModelID(model.ID)] = true + } + if ids[modelID] { + return modelID, nil + } + suggestion, ok := adapter.BestOpenRouterImageModelMatch(modelID, models) + if !ok { + return modelID, nil + } + replace, err := runConfirmPrompt(fmt.Sprintf(" Model %q was not found on OpenRouter. Use %q instead?", modelID, suggestion.ID)) + if err != nil { + if isCancelled(err) { + printCancelled() + return "", errCancelled + } + return "", err + } + if replace { + return suggestion.ID, nil + } + return modelID, nil +} + +func ensureCustomDisplayName(displayName string) string { + if strings.HasSuffix(displayName, " [Custom]") { + return displayName + } + return displayName + " [Custom]" +} diff --git a/internal/tui/model_list_custom_test.go b/internal/tui/model_list_custom_test.go new file mode 100644 index 0000000..14af49b --- /dev/null +++ b/internal/tui/model_list_custom_test.go @@ -0,0 +1,121 @@ +package tui + +import "testing" + +func TestPromptCustomModelStoresCustomSuffixAndDefaultsEdit(t *testing.T) { + oldInputPrompt := runInputPrompt + oldConfirmPrompt := runConfirmPrompt + t.Cleanup(func() { + runInputPrompt = oldInputPrompt + runConfirmPrompt = oldConfirmPrompt + }) + + inputs := []string{"custom/image-model", "Friendly Image"} + runInputPrompt = func(label, value string, secret bool) (string, error) { + if len(inputs) == 0 { + t.Fatalf("unexpected input prompt %q", label) + } + out := inputs[0] + inputs = inputs[1:] + return out, nil + } + runConfirmPrompt = func(label string) (bool, error) { + if label == "Can custom/image-model edit images?" { + return true, nil + } + return false, nil + } + + selected, err := promptCustomModel() + if err != nil { + t.Fatalf("promptCustomModel: %v", err) + } + if selected.ID != "custom/image-model" { + t.Fatalf("id = %q, want custom/image-model", selected.ID) + } + if selected.DisplayName != "Friendly Image [Custom]" { + t.Fatalf("display name = %q, want Friendly Image [Custom]", selected.DisplayName) + } + if !selected.Custom { + t.Fatal("custom flag = false, want true") + } + if !selected.Edit { + t.Fatal("edit = false, want default true") + } +} + +func TestPromptCustomModelCanTweakAdvancedParameters(t *testing.T) { + oldInputPrompt := runInputPrompt + oldConfirmPrompt := runConfirmPrompt + t.Cleanup(func() { + runInputPrompt = oldInputPrompt + runConfirmPrompt = oldConfirmPrompt + }) + + inputs := []string{"custom/private-art-model", "", "1536x1024", "high", "2", "url", "42", "7.5", "blurry"} + runInputPrompt = func(label, value string, secret bool) (string, error) { + if len(inputs) == 0 { + t.Fatalf("unexpected input prompt %q", label) + } + out := inputs[0] + inputs = inputs[1:] + return out, nil + } + confirmCount := 0 + runConfirmPrompt = func(label string) (bool, error) { + confirmCount++ + return true, nil + } + + selected, err := promptCustomModel() + if err != nil { + t.Fatalf("promptCustomModel: %v", err) + } + if confirmCount < 2 { + t.Fatalf("confirm prompts = %d, want advanced and edit prompts", confirmCount) + } + if selected.DisplayName != "custom/private-art-model [Custom]" { + t.Fatalf("display name = %q, want default id suffix", selected.DisplayName) + } + params := selected.Parameters + if params.Size != "1536x1024" || params.Quality != "high" || params.N != 2 || + params.ResponseFormat != "url" || params.Seed != 42 || + params.GuidanceScale != 7.5 || params.NegativePrompt != "blurry" { + t.Fatalf("parameters = %+v, want configured values", params) + } +} + +func TestPromptCustomModelCanUseOpenRouterSuggestion(t *testing.T) { + oldInputPrompt := runInputPrompt + oldConfirmPrompt := runConfirmPrompt + t.Cleanup(func() { + runInputPrompt = oldInputPrompt + runConfirmPrompt = oldConfirmPrompt + }) + + inputs := []string{"image-modl", ""} + runInputPrompt = func(label, value string, secret bool) (string, error) { + if len(inputs) == 0 { + t.Fatalf("unexpected input prompt %q", label) + } + out := inputs[0] + inputs = inputs[1:] + return out, nil + } + confirmCount := 0 + runConfirmPrompt = func(label string) (bool, error) { + confirmCount++ + return confirmCount == 1, nil + } + + selected, err := promptCustomModel() + if err != nil { + t.Fatalf("promptCustomModel: %v", err) + } + if selected.ID != "openrouter/image-model" { + t.Fatalf("id = %q, want OpenRouter suggestion", selected.ID) + } + if selected.DisplayName != "openrouter/image-model [Custom]" { + t.Fatalf("display name = %q, want default suggested id suffix", selected.DisplayName) + } +} diff --git a/internal/tui/model_list_provider.go b/internal/tui/model_list_provider.go new file mode 100644 index 0000000..5bd4138 --- /dev/null +++ b/internal/tui/model_list_provider.go @@ -0,0 +1,91 @@ +package tui + +import ( + "fmt" + + "github.com/ncxton/potaco/internal/auth" + "github.com/ncxton/potaco/internal/config" +) + +type modelListRequest struct { + ProviderName string + APIKey string + BaseURL string +} + +type resolvedModelListProvider struct { + Name string + APIKey string + BaseURL string + ProviderType string +} + +func resolveModelListProvider(mgr *auth.AuthManager, req modelListRequest) (resolvedModelListProvider, error) { + providerName := req.ProviderName + if providerName == "" { + activeProvider, _, err := mgr.GetActiveProvider() + if err != nil || activeProvider == "" { + return resolvedModelListProvider{}, fmt.Errorf("no active provider. Use 'potaco auth add ' to connect one") + } + providerName = activeProvider + } + + apiKey := req.APIKey + if apiKey == "" { + key, err := mgr.GetAPIKey(providerName) + if err != nil { + return resolvedModelListProvider{}, fmt.Errorf("provider %q is not connected. Use 'potaco auth add %s' first", providerName, providerName) + } + apiKey = key + } + + baseURL := req.BaseURL + pc := config.ProviderConfig{} + cfg, _ := mgr.LoadConfig() + if cfg != nil { + if configured, ok := cfg.Providers[providerName]; ok { + pc = configured + if baseURL == "" && configured.BaseURL != "" { + baseURL = configured.BaseURL + } + } + } + + providerType := config.ResolveProviderType(providerName, pc) + baseURL, err := normalizeTrustedModelListBaseURL(baseURL) + if err != nil { + return resolvedModelListProvider{}, err + } + if config.ProviderRequiresBaseURL(providerName, providerType) && baseURL == "" { + return resolvedModelListProvider{}, fmt.Errorf("base URL required for provider %s", providerName) + } + + return resolvedModelListProvider{ + Name: providerName, + APIKey: apiKey, + BaseURL: baseURL, + ProviderType: providerType, + }, nil +} + +func normalizeTrustedModelListBaseURL(baseURL string) (string, error) { + if baseURL == "" { + return "", nil + } + normalized, err := config.NormalizeBaseURL(baseURL) + if err != nil { + return "", err + } + if !config.BaseURLIsInsecureHTTP(normalized) { + return normalized, nil + } + trusted, err := confirmInsecureBaseURL(normalized) + if err != nil { + return "", err + } + if !trusted { + printCancelled() + return "", errCancelled + } + return normalized, nil +} diff --git a/internal/tui/model_list_test.go b/internal/tui/model_list_test.go index 51bc996..3386749 100644 --- a/internal/tui/model_list_test.go +++ b/internal/tui/model_list_test.go @@ -7,12 +7,11 @@ import ( "strings" "testing" - "github.com/charmbracelet/huh" "github.com/ncxton/potaco/internal/adapter" _ "github.com/ncxton/potaco/internal/adapter/custom" _ "github.com/ncxton/potaco/internal/adapter/fal" _ "github.com/ncxton/potaco/internal/adapter/openai" - _ "github.com/ncxton/potaco/internal/adapter/vercel" + _ "github.com/ncxton/potaco/internal/adapter/replicate" "github.com/ncxton/potaco/internal/auth" ) @@ -39,6 +38,7 @@ func TestRunModelListNotConnectedReturnsError(t *testing.T) { } func TestRunModelListPersistsSelection(t *testing.T) { + trustHTTPBaseURLs(t) t.Setenv("HOME", t.TempDir()) mgr, err := auth.New() if err != nil { @@ -68,7 +68,7 @@ func TestRunModelListPersistsSelection(t *testing.T) { if len(models) != 1 { t.Errorf("models = %d, want 1", len(models)) } - return modelSelection{ID: "gpt-image-2", Edit: true}, nil + return modelSelection{ID: "gpt-image-2", DisplayName: "GPT Image 2", Edit: true}, nil } err = runModelListWithPicker("openai", "sk-test", srv.URL, mockPicker) @@ -97,9 +97,13 @@ func TestRunModelListPersistsSelection(t *testing.T) { if !cfg.Providers["openai"].Models["gpt-image-2"].Edit { t.Fatal("model edit capability = false, want true") } + if cfg.Providers["openai"].Models["gpt-image-2"].DisplayName != "GPT Image 2" { + t.Fatalf("model display name = %q, want GPT Image 2", cfg.Providers["openai"].Models["gpt-image-2"].DisplayName) + } } func TestRunModelListCancelDoesNotPersist(t *testing.T) { + trustHTTPBaseURLs(t) t.Setenv("HOME", t.TempDir()) mgr, err := auth.New() if err != nil { @@ -125,7 +129,7 @@ func TestRunModelListCancelDoesNotPersist(t *testing.T) { } mockPicker := func(providerName string, models []adapter.Model) (modelSelection, error) { - return modelSelection{}, huh.ErrUserAborted + return modelSelection{}, errCancelled } err = runModelListWithPicker("openai", "sk-test", srv.URL, mockPicker) @@ -143,6 +147,7 @@ func TestRunModelListCancelDoesNotPersist(t *testing.T) { } func TestRunModelListEmptySelectionDoesNotPersist(t *testing.T) { + trustHTTPBaseURLs(t) t.Setenv("HOME", t.TempDir()) mgr, err := auth.New() if err != nil { @@ -186,6 +191,7 @@ func TestRunModelListEmptySelectionDoesNotPersist(t *testing.T) { } func TestRunModelListUsesBaseURLFromConfig(t *testing.T) { + trustHTTPBaseURLs(t) t.Setenv("HOME", t.TempDir()) mgr, err := auth.New() if err != nil { @@ -250,6 +256,7 @@ func TestRunModelListAliasWithBuiltInTypeRequiresBaseURL(t *testing.T) { } func TestRunModelListDiscoveryError(t *testing.T) { + trustHTTPBaseURLs(t) t.Setenv("HOME", t.TempDir()) mgr, err := auth.New() if err != nil { diff --git a/internal/tui/model_parameters.go b/internal/tui/model_parameters.go new file mode 100644 index 0000000..9ecff25 --- /dev/null +++ b/internal/tui/model_parameters.go @@ -0,0 +1,75 @@ +package tui + +import ( + "fmt" + "strings" + + "github.com/ncxton/potaco/internal/config" +) + +func promptModelParameters(defaults config.ModelParameters) (config.ModelParameters, error) { + params := defaults + var err error + if params.Size, err = promptParameterString("Image size:", defaults.Size); err != nil { + return config.ModelParameters{}, err + } + if params.Quality, err = promptParameterString("Quality:", defaults.Quality); err != nil { + return config.ModelParameters{}, err + } + if params.N, err = promptParameterInt("Number of images:", defaults.N); err != nil { + return config.ModelParameters{}, err + } + if params.ResponseFormat, err = promptParameterString("Response format:", defaults.ResponseFormat); err != nil { + return config.ModelParameters{}, err + } + if params.Seed, err = promptParameterInt("Seed:", defaults.Seed); err != nil { + return config.ModelParameters{}, err + } + if params.GuidanceScale, err = promptParameterFloat("Guidance scale:", defaults.GuidanceScale); err != nil { + return config.ModelParameters{}, err + } + if params.NegativePrompt, err = promptParameterString("Negative prompt:", defaults.NegativePrompt); err != nil { + return config.ModelParameters{}, err + } + return params, nil +} + +func promptParameterString(label, value string) (string, error) { + got, err := runInputPrompt(label, value, false) + if err != nil { + if isCancelled(err) { + printCancelled() + return "", errCancelled + } + return "", err + } + got = strings.TrimSpace(got) + if got == "" { + return value, nil + } + return got, nil +} + +func promptParameterInt(label string, value int) (int, error) { + raw, err := promptParameterString(label, fmt.Sprintf("%d", value)) + if err != nil { + return 0, err + } + var out int + if _, err := fmt.Sscanf(raw, "%d", &out); err != nil { + return 0, fmt.Errorf("%s must be a number", label) + } + return out, nil +} + +func promptParameterFloat(label string, value float64) (float64, error) { + raw, err := promptParameterString(label, fmt.Sprintf("%g", value)) + if err != nil { + return 0, err + } + var out float64 + if _, err := fmt.Sscanf(raw, "%f", &out); err != nil { + return 0, fmt.Errorf("%s must be a number", label) + } + return out, nil +} diff --git a/internal/tui/model_search.go b/internal/tui/model_search.go index bdf8696..95739dc 100644 --- a/internal/tui/model_search.go +++ b/internal/tui/model_search.go @@ -3,19 +3,12 @@ package tui import ( "strings" - "charm.land/lipgloss/v2" - "github.com/charmbracelet/bubbles/textinput" - tea "github.com/charmbracelet/bubbletea" + tea "charm.land/bubbletea/v2" "github.com/ncxton/potaco/internal/adapter" ) -var ( - modelTitleStyle = lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("42")) - modelFocusStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("6")) - modelMutedStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("8")) - modelEmptyStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("203")) -) +const customModelSentinel = "__potaco_custom_model__" // searchModel is a custom Bubble Tea model for searching and selecting // from a list of models by typing. The filter updates in real-time as @@ -25,7 +18,7 @@ type searchModel struct { models []adapter.Model filtered []adapter.Model cursor int - query textinput.Model + query string selected string quitted bool } @@ -33,28 +26,22 @@ type searchModel struct { // newSearchModel creates a new searchModel initialized with the given // models. All models are shown initially (no filter applied). func newSearchModel(providerName string, models []adapter.Model) *searchModel { - ti := textinput.New() - ti.Prompt = "> " - ti.Placeholder = "Type to search..." - ti.Focus() - return &searchModel{ providerName: providerName, models: models, filtered: models, - query: ti, } } // Init implements tea.Model. func (m *searchModel) Init() tea.Cmd { - return textinput.Blink + return nil } // applyFilter updates m.filtered based on the current query value. // Filtering is case-insensitive and matches against both ID and DisplayName. func (m *searchModel) applyFilter() { - q := strings.ToLower(m.query.Value()) + q := strings.ToLower(m.query) if q == "" { m.filtered = m.models m.clampCursor() @@ -75,69 +62,77 @@ func (m *searchModel) clampCursor() { if m.cursor < 0 { m.cursor = 0 } - if m.cursor > len(m.filtered)-1 { - m.cursor = max(0, len(m.filtered)-1) + if m.cursor > m.optionCount()-1 { + m.cursor = max(0, m.optionCount()-1) } } +func (m *searchModel) optionCount() int { + return len(m.filtered) + 1 +} + // Update implements tea.Model. func (m *searchModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch msg := msg.(type) { - case tea.KeyMsg: - switch msg.Type { - case tea.KeyEsc, tea.KeyCtrlC: + case tea.PasteMsg: + m.query += msg.Content + m.applyFilter() + return m, nil + case tea.KeyPressMsg: + switch msg.Keystroke() { + case "esc", "ctrl+c": m.quitted = true return m, tea.Quit - case tea.KeyEnter: - if len(m.filtered) > 0 { + case "enter": + if m.cursor == len(m.filtered) { + m.selected = customModelSentinel + } else if len(m.filtered) > 0 { m.selected = m.filtered[m.cursor].ID } return m, tea.Quit - case tea.KeyDown, tea.KeyCtrlJ, tea.KeyCtrlN: - if len(m.filtered) > 0 { - m.cursor = (m.cursor + 1) % len(m.filtered) - } - case tea.KeyUp, tea.KeyCtrlK, tea.KeyCtrlP: - if len(m.filtered) > 0 { - m.cursor = (m.cursor - 1 + len(m.filtered)) % len(m.filtered) - } - default: - // Any other key: let the text input handle it, then re-filter. - var cmd tea.Cmd - m.query, cmd = m.query.Update(msg) + case "down", "ctrl+j", "ctrl+n": + m.cursor = (m.cursor + 1) % m.optionCount() + case "up", "ctrl+k", "ctrl+p": + m.cursor = (m.cursor - 1 + m.optionCount()) % m.optionCount() + case "backspace": + m.query = trimLastRune(m.query) m.applyFilter() - return m, cmd + default: + if msg.Key().Text != "" { + m.query += msg.Key().Text + m.applyFilter() + } } return m, nil default: - // Forward non-key messages to the text input (e.g. blink). - var cmd tea.Cmd - m.query, cmd = m.query.Update(msg) - return m, cmd + return m, nil } } // View implements tea.Model. -func (m *searchModel) View() string { +func (m *searchModel) View() tea.View { var b strings.Builder - b.WriteString(modelTitleStyle.Render("Select a model for " + m.providerName)) + b.WriteString("\n") + b.WriteString(renderPromptTitle("Select a model for " + m.providerName)) b.WriteString("\n\n") - b.WriteString(m.query.View()) - b.WriteString("\n\n") + b.WriteString(promptTextInset) + if m.query == "" { + b.WriteString(promptMutedStyle.Render("Type to search...")) + } else { + b.WriteString(m.query) + } + b.WriteString("\n") for i, model := range m.filtered { cursor := " " if i == m.cursor { - cursor = "> " - } - label := model.ID - if model.DisplayName != "" && model.DisplayName != model.ID { - label += " " + model.DisplayName + cursor = " " } + label := modelLabel(model) if i == m.cursor { - label = modelFocusStyle.Render(label) + label = promptPickerStyle.Render(label) } b.WriteString(cursor) b.WriteString(label) @@ -145,13 +140,43 @@ func (m *searchModel) View() string { } if len(m.filtered) == 0 { - b.WriteString(modelEmptyStyle.Render(" No matching models.")) + b.WriteString(promptErrorStyle.Render(" No matching models.")) b.WriteString("\n") } + customCursor := " " + if m.cursor == len(m.filtered) { + customCursor = " " + } + customLabel := promptMutedStyle.Render("Custom") + if m.cursor == len(m.filtered) { + customLabel = promptPickerStyle.Render("Custom") + } + b.WriteString(customCursor) + b.WriteString(customLabel) + b.WriteString("\n") b.WriteString("\n") - b.WriteString(modelMutedStyle.Render(" \u2191\u2193 navigate enter select esc cancel")) + b.WriteString(renderHelp("", "↑↓", " navigate ", "enter", " select ", "esc", " cancel")) b.WriteString("\n") - return b.String() + view := tea.NewView(b.String()) + view.AltScreen = true + view.Cursor = tea.NewCursor(inputCursorX(m.query), inputLineY) + return view +} + +func modelLabel(model adapter.Model) string { + if model.DisplayName == "" || model.DisplayName == model.ID { + return model.ID + } + return model.DisplayName + " " + model.ID +} + +func (m *searchModel) selectedModel() adapter.Model { + for _, model := range m.models { + if model.ID == m.selected { + return model + } + } + return adapter.Model{ID: m.selected} } diff --git a/internal/tui/model_search_custom_test.go b/internal/tui/model_search_custom_test.go new file mode 100644 index 0000000..3dcc955 --- /dev/null +++ b/internal/tui/model_search_custom_test.go @@ -0,0 +1,44 @@ +package tui + +import ( + "strings" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/ncxton/potaco/internal/adapter" +) + +func TestSearchModelViewShowsDisplayNameBeforeIDAndCustomBottom(t *testing.T) { + models := []adapter.Model{ + {ID: "gpt-image-2", DisplayName: "GPT Image 2"}, + } + m := newSearchModel("openai", models) + + view := m.View().Content + if !strings.Contains(view, "GPT Image 2 gpt-image-2") { + t.Fatalf("view should prefer display name before id:\n%s", view) + } + customIdx := strings.LastIndex(view, "Custom") + modelIdx := strings.Index(view, "GPT Image 2") + if customIdx < modelIdx { + t.Fatalf("custom option should render after discovered models:\n%s", view) + } +} + +func TestSearchEnterSelectsCustomOption(t *testing.T) { + models := []adapter.Model{ + {ID: "gpt-image-2", DisplayName: "GPT Image 2"}, + } + m := newSearchModel("openai", models) + m.cursor = 1 + + _, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyEnter})) + + if cmd == nil { + t.Error("expected non-nil cmd (tea.Quit)") + } + if m.selected != customModelSentinel { + t.Errorf("selected = %q, want custom sentinel", m.selected) + } +} diff --git a/internal/tui/model_search_test.go b/internal/tui/model_search_test.go index 7c05a76..693f696 100644 --- a/internal/tui/model_search_test.go +++ b/internal/tui/model_search_test.go @@ -4,7 +4,7 @@ import ( "strings" "testing" - tea "github.com/charmbracelet/bubbletea" + tea "charm.land/bubbletea/v2" "github.com/ncxton/potaco/internal/adapter" ) @@ -15,10 +15,16 @@ func TestSearchModelViewIncludesStyledTitleAndBadgeText(t *testing.T) { } m := newSearchModel("openai", models) - view := m.View() + view := m.View().Content if !strings.Contains(view, "Select a model for openai") { t.Fatalf("view missing title:\n%s", view) } + if strings.Contains(view, "> Type to search") { + t.Fatalf("search filter should not show picker marker:\n%s", view) + } + if !strings.Contains(view, " ") || strings.Contains(view, "> gpt-image-1") { + t.Fatalf("model list should use Nerd Font picker cursor:\n%s", view) + } if strings.Contains(view, "[edit]") { t.Fatalf("view should not show discovered edit badge:\n%s", view) } @@ -50,7 +56,7 @@ func TestSearchFilterByID(t *testing.T) { {ID: "flux-schnell", DisplayName: "Flux Schnell"}, } m := newSearchModel("openai", models) - m.query.SetValue("gpt") + m.query = "gpt" m.applyFilter() @@ -69,7 +75,7 @@ func TestSearchFilterByDisplayName(t *testing.T) { {ID: "flux-schnell", DisplayName: "Flux Schnell"}, } m := newSearchModel("openai", models) - m.query.SetValue("flux") + m.query = "flux" m.applyFilter() @@ -87,7 +93,7 @@ func TestSearchFilterCaseInsensitive(t *testing.T) { {ID: "flux-pro", DisplayName: "Flux Pro"}, } m := newSearchModel("openai", models) - m.query.SetValue("gpt") + m.query = "gpt" m.applyFilter() @@ -102,7 +108,7 @@ func TestSearchFilterEmptyShowsAll(t *testing.T) { {ID: "flux-pro", DisplayName: "Flux Pro"}, } m := newSearchModel("openai", models) - m.query.SetValue("") + m.query = "" m.applyFilter() @@ -116,13 +122,16 @@ func TestSearchFilterNoMatch(t *testing.T) { {ID: "gpt-image-2", DisplayName: "GPT Image 2"}, } m := newSearchModel("openai", models) - m.query.SetValue("xyz123") + m.query = "xyz123" m.applyFilter() if len(m.filtered) != 0 { t.Errorf("expected 0 filtered models, got %d", len(m.filtered)) } + if !strings.Contains(m.View().Content, "Custom") { + t.Fatal("custom option should remain visible when no discovered model matches") + } } func TestSearchCursorClampsToFiltered(t *testing.T) { @@ -135,12 +144,12 @@ func TestSearchCursorClampsToFiltered(t *testing.T) { m.cursor = 2 // at last item // Filter to only 1 item - m.query.SetValue("schnell") + m.query = "schnell" m.applyFilter() - // applyFilter calls clampCursor internally, so cursor should be 0 - if m.cursor != 0 { - t.Errorf("expected cursor clamped to 0, got %d", m.cursor) + // applyFilter clamps to the last visible option, which is Custom. + if m.cursor != 1 { + t.Errorf("expected cursor clamped to Custom at 1, got %d", m.cursor) } } @@ -150,7 +159,7 @@ func TestSearchEscQuits(t *testing.T) { } m := newSearchModel("openai", models) - _, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEsc}) + _, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyEsc})) if !m.quitted { t.Error("expected quitted=true after Esc") @@ -167,7 +176,7 @@ func TestSearchEnterSelects(t *testing.T) { } m := newSearchModel("openai", models) - _, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter}) + _, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyEnter})) if cmd == nil { t.Error("expected non-nil cmd (tea.Quit)") @@ -187,17 +196,81 @@ func TestSearchArrowDownMovesCursor(t *testing.T) { {ID: "flux-schnell", DisplayName: "Flux Schnell"}, } m := newSearchModel("openai", models) - m.Update(tea.KeyMsg{Type: tea.KeyDown}) + m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyDown})) if m.cursor != 1 { t.Errorf("expected cursor at 1, got %d", m.cursor) } - m.Update(tea.KeyMsg{Type: tea.KeyDown}) + m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyDown})) if m.cursor != 2 { t.Errorf("expected cursor at 2, got %d", m.cursor) } - // Wrap at bottom - m.Update(tea.KeyMsg{Type: tea.KeyDown}) + m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyDown})) + if m.cursor != 3 { + t.Errorf("expected cursor at Custom row 3, got %d", m.cursor) + } + // Wrap after Custom. + m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyDown})) if m.cursor != 0 { t.Errorf("expected cursor to wrap to 0, got %d", m.cursor) } } + +func TestSearchKeyboardFilteringAndNavigationWithBubbleTeaV2(t *testing.T) { + models := []adapter.Model{ + {ID: "gpt-image-2", DisplayName: "GPT Image 2"}, + {ID: "flux-pro", DisplayName: "Flux Pro"}, + {ID: "flux-schnell", DisplayName: "Flux Schnell"}, + } + m := newSearchModel("openai", models) + + m.Update(tea.KeyPressMsg(tea.Key{Text: "flux", Code: tea.KeyExtended})) + if m.query != "flux" { + t.Fatalf("query = %q, want flux", m.query) + } + if len(m.filtered) != 2 { + t.Fatalf("filtered length = %d, want 2", len(m.filtered)) + } + + m.Update(tea.KeyPressMsg(tea.Key{Code: 'n', Mod: tea.ModCtrl})) + if m.cursor != 1 { + t.Fatalf("ctrl+n cursor = %d, want 1", m.cursor) + } + + m.Update(tea.KeyPressMsg(tea.Key{Code: 'j', Mod: tea.ModCtrl})) + if m.cursor != 2 { + t.Fatalf("ctrl+j should move cursor to custom row 2, got %d", m.cursor) + } + + m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyUp})) + if m.cursor != 1 { + t.Fatalf("up should wrap cursor to 1, got %d", m.cursor) + } + + m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyBackspace})) + if m.query != "flu" { + t.Fatalf("backspace query = %q, want flu", m.query) + } +} + +func TestSearchHandlesPasteAndShowsCursor(t *testing.T) { + models := []adapter.Model{ + {ID: "gpt-image-2", DisplayName: "GPT Image 2"}, + {ID: "flux-pro", DisplayName: "Flux Pro"}, + } + m := newSearchModel("openai", models) + + m.Update(tea.PasteMsg{Content: "flux"}) + + if m.query != "flux" { + t.Fatalf("query = %q, want pasted filter", m.query) + } + if len(m.filtered) != 1 || m.filtered[0].ID != "flux-pro" { + t.Fatalf("filtered = %#v, want only flux-pro", m.filtered) + } + if m.View().Cursor == nil { + t.Fatal("search view should expose a cursor") + } + if !m.View().AltScreen { + t.Fatal("search view should use alternate screen") + } +} diff --git a/internal/tui/openrouter_test.go b/internal/tui/openrouter_test.go new file mode 100644 index 0000000..189b6d0 --- /dev/null +++ b/internal/tui/openrouter_test.go @@ -0,0 +1,36 @@ +package tui + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "testing" +) + +var openRouterCatalogServer *httptest.Server + +func TestMain(m *testing.M) { + openRouterCatalogServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "data": []map[string]any{ + openRouterCatalogModel("openai/gpt-image-2"), + openRouterCatalogModel("openrouter/image-model"), + }, + }) + })) + os.Setenv("POTACO_OPENROUTER_MODELS_URL", openRouterCatalogServer.URL+"/models?output_modalities=image") + code := m.Run() + openRouterCatalogServer.Close() + os.Exit(code) +} + +func openRouterCatalogModel(id string) map[string]any { + return map[string]any{ + "id": id, + "architecture": map[string]any{ + "output_modalities": []string{"image"}, + }, + } +} diff --git a/internal/tui/prompt.go b/internal/tui/prompt.go new file mode 100644 index 0000000..296f3fe --- /dev/null +++ b/internal/tui/prompt.go @@ -0,0 +1,212 @@ +package tui + +import ( + "errors" + "fmt" + "os" + "strings" + "unicode/utf8" + + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/colorprofile" +) + +var ( + errCancelled = errors.New("cancelled") + + runSelectPrompt = runSelect + runInputPrompt = runInput + runConfirmPrompt = ConfirmAction + + promptTitleStyle = lipgloss.NewStyle().Bold(true) + promptFocusStyle = lipgloss.NewStyle() + promptPickerStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("42")) + promptMutedStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + promptKeyStyle = lipgloss.NewStyle() + promptSuccessStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("42")) + promptWarningStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("3")) + promptErrorStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("203")) +) + +const ( + promptTextInset = " " + inputLineY = 3 +) + +type promptOption struct { + Label string + Value string +} + +type promptKind int + +const ( + promptKindSelect promptKind = iota + promptKindInput + promptKindConfirm +) + +type promptState struct { + kind promptKind + title string + options []promptOption + cursor int + value string + secret bool + confirm bool + done bool + cancel bool + selected string +} + +func runSelect(title string, options []promptOption) (string, error) { + if len(options) == 0 { + return "", nil + } + m := &promptState{kind: promptKindSelect, title: title, options: options} + result, err := runPrompt(m) + if err != nil { + return "", err + } + if result.cancel { + return "", errCancelled + } + return result.selected, nil +} + +func runInput(title, value string, secret bool) (string, error) { + m := &promptState{kind: promptKindInput, title: title, value: value, secret: secret} + result, err := runPrompt(m) + if err != nil { + return "", err + } + if result.cancel { + return "", errCancelled + } + return result.value, nil +} + +func ConfirmAction(prompt string) (bool, error) { + m := &promptState{kind: promptKindConfirm, title: prompt} + result, err := runPrompt(m) + if err != nil { + return false, err + } + if result.cancel { + return false, errCancelled + } + return result.confirm, nil +} + +func runPrompt(m *promptState) (*promptState, error) { + p := newPromptProgram(m) + result, err := p.Run() + if err != nil { + return nil, err + } + pm, ok := result.(*promptState) + if !ok { + return nil, fmt.Errorf("unexpected prompt model") + } + return pm, nil +} + +func newPromptProgram(m tea.Model) *tea.Program { + return tea.NewProgram(m, tea.WithOutput(os.Stderr), tea.WithColorProfile(colorprofile.ANSI256)) +} + +func (m *promptState) Init() tea.Cmd { + return nil +} + +func (m *promptState) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.PasteMsg: + if m.kind == promptKindInput { + m.value += msg.Content + } + if m.kind == promptKindConfirm { + m.updateConfirmValue(msg.Content) + } + return m, nil + default: + key, ok := msg.(tea.KeyPressMsg) + if !ok { + return m, nil + } + k := key.Key() + if k.Code == tea.KeyEsc || k.Code == 'c' && k.Mod == tea.ModCtrl { + m.cancel = true + return m, tea.Quit + } + switch m.kind { + case promptKindSelect: + return m.updateSelect(key) + case promptKindInput: + return m.updateInput(key) + case promptKindConfirm: + return m.updateConfirm(key) + default: + return m, nil + } + } +} + +func (m *promptState) updateSelect(key tea.KeyPressMsg) (tea.Model, tea.Cmd) { + switch key.Keystroke() { + case "enter": + m.selected = m.options[m.cursor].Value + return m, tea.Quit + case "down", "ctrl+n", "ctrl+j": + m.cursor = (m.cursor + 1) % len(m.options) + case "up", "ctrl+p", "ctrl+k": + m.cursor = (m.cursor - 1 + len(m.options)) % len(m.options) + } + return m, nil +} + +func (m *promptState) updateInput(key tea.KeyPressMsg) (tea.Model, tea.Cmd) { + switch key.Keystroke() { + case "enter": + m.done = true + return m, tea.Quit + case "backspace": + m.value = trimLastRune(m.value) + default: + m.value += key.Key().Text + } + return m, nil +} + +func (m *promptState) updateConfirm(key tea.KeyPressMsg) (tea.Model, tea.Cmd) { + switch key.Keystroke() { + case "enter": + m.confirm = m.value != "n" + m.done = true + return m, tea.Quit + case "backspace": + m.value = trimLastRune(m.value) + default: + m.updateConfirmValue(key.Key().Text) + } + return m, nil +} + +func (m *promptState) updateConfirmValue(text string) { + for _, r := range strings.ToLower(text) { + switch r { + case 'y', 'n': + m.value = string(r) + return + } + } +} + +func trimLastRune(s string) string { + if s == "" { + return "" + } + _, size := utf8.DecodeLastRuneInString(s) + return s[:len(s)-size] +} diff --git a/internal/tui/prompt_view.go b/internal/tui/prompt_view.go new file mode 100644 index 0000000..3ce1d6d --- /dev/null +++ b/internal/tui/prompt_view.go @@ -0,0 +1,137 @@ +package tui + +import ( + "strings" + "unicode/utf8" + + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" +) + +func (m *promptState) View() tea.View { + var b strings.Builder + b.WriteString("\n") + if m.kind == promptKindConfirm { + b.WriteString(m.confirmLines()) + b.WriteString("\n\n") + b.WriteString(m.help()) + b.WriteString("\n") + view := tea.NewView(b.String()) + view.AltScreen = true + view.Cursor = tea.NewCursor(confirmCursorX(m.title, m.value), confirmCursorY(m.title)) + return view + } + + b.WriteString(renderPromptTitle(m.title)) + b.WriteString("\n\n") + switch m.kind { + case promptKindSelect: + m.viewSelect(&b) + case promptKindInput: + m.viewInput(&b) + } + b.WriteString("\n") + b.WriteString(m.help()) + b.WriteString("\n") + view := tea.NewView(b.String()) + view.AltScreen = true + if m.kind == promptKindInput { + view.Cursor = tea.NewCursor(inputCursorX(m.inputDisplayValue()), inputLineY) + } + return view +} + +func (m *promptState) viewSelect(b *strings.Builder) { + if len(m.options) == 0 { + b.WriteString(promptErrorStyle.Render(" No options available.")) + b.WriteString("\n") + return + } + for i, option := range m.options { + cursor := " " + label := option.Label + if i == m.cursor { + cursor = " " + label = promptPickerStyle.Render(label) + } + b.WriteString(cursor) + b.WriteString(label) + b.WriteString("\n") + } +} + +func (m *promptState) viewInput(b *strings.Builder) { + b.WriteString(promptFocusStyle.Render("> ")) + b.WriteString(m.inputDisplayValue()) + b.WriteString("\n") +} + +func (m *promptState) confirmLines() string { + lines := strings.Split(m.title, "\n") + last := len(lines) - 1 + for i, line := range lines { + lines[i] = promptTextInset + promptTitleStyle.Render(line) + if i == last { + lines[i] += " " + promptFocusStyle.Render("[Y/n] ") + m.value + } + } + return strings.Join(lines, "\n") +} + +func (m *promptState) help() string { + switch m.kind { + case promptKindSelect: + return renderHelp("", "up/down", " navigate ", "enter", " select ", "esc", " cancel") + case promptKindInput: + return renderHelp("", "enter", " submit ", "esc", " cancel") + case promptKindConfirm: + return renderHelp("", "y/n", " answer ", "enter", " submit ", "esc", " cancel") + default: + return renderHelp("", "esc", " cancel") + } +} + +func (m *promptState) inputDisplayValue() string { + if m.secret && m.value != "" { + return strings.Repeat("*", utf8.RuneCountInString(m.value)) + } + return m.value +} + +func inputCursorX(value string) int { + return 2 + lipgloss.Width(value) +} + +func confirmCursorX(title, value string) int { + lines := strings.Split(title, "\n") + lastLine := lines[len(lines)-1] + return lipgloss.Width(promptTextInset) + lipgloss.Width(lastLine) + 1 + lipgloss.Width("[Y/n] ") + lipgloss.Width(value) +} + +func confirmCursorY(title string) int { + return len(strings.Split(title, "\n")) +} + +func renderPromptTitle(title string) string { + lines := strings.Split(title, "\n") + for i, line := range lines { + lines[i] = promptTextInset + promptTitleStyle.Render(line) + } + return strings.Join(lines, "\n") +} + +func renderHelp(parts ...string) string { + var b strings.Builder + b.WriteString(promptTextInset) + for i, part := range parts { + if part == "" { + continue + } + if i%2 == 1 { + b.WriteString(promptKeyStyle.Render(part)) + continue + } + b.WriteString(promptMutedStyle.Render(part)) + } + return b.String() +} diff --git a/internal/tui/provider_type_test.go b/internal/tui/provider_type_test.go index c93dce6..cb4f7ba 100644 --- a/internal/tui/provider_type_test.go +++ b/internal/tui/provider_type_test.go @@ -7,28 +7,36 @@ import ( "testing" "github.com/ncxton/potaco/internal/adapter" - _ "github.com/ncxton/potaco/internal/adapter/custom" "github.com/ncxton/potaco/internal/auth" "github.com/ncxton/potaco/internal/config" ) func TestRunModelListUsesConfiguredProviderType(t *testing.T) { + trustHTTPBaseURLs(t) t.Setenv("HOME", t.TempDir()) mgr, err := auth.New() if err != nil { t.Fatalf("init auth: %v", err) } - if err := mgr.AddProvider("openrouter", "openai-compatible", "sk-openrouter"); err != nil { + if err := mgr.AddProvider("openrouter", "openai", "sk-openrouter"); err != nil { t.Fatalf("add provider: %v", err) } var gotAuth string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotAuth = r.Header.Get("Authorization") + if r.URL.Path == "/v1/models" { + gotAuth = r.Header.Get("Authorization") + } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]any{ "data": []map[string]any{ - {"id": "openrouter/image-model", "owned_by": "openrouter"}, + { + "id": "openrouter/image-model", + "owned_by": "openrouter", + "architecture": map[string]any{ + "output_modalities": []string{"image"}, + }, + }, }, }) })) diff --git a/internal/tui/tui.go b/internal/tui/tui.go index 14d88b5..b1458fe 100644 --- a/internal/tui/tui.go +++ b/internal/tui/tui.go @@ -3,16 +3,27 @@ package tui import ( "errors" + "fmt" "os" - "github.com/charmbracelet/bubbles/key" - "github.com/charmbracelet/huh" "golang.org/x/term" ) // nonInteractive is set by the CLI when --non-interactive is passed. var nonInteractive bool +var printCancelled = func() { + fmt.Println(promptErrorStyle.Render(" Cancelled.")) +} + +var printUnableToFindModels = func() { + fmt.Println(promptWarningStyle.Render(" Unable to find any image generation models.")) +} + +var printSuccess = func(format string, args ...any) { + fmt.Printf(promptSuccessStyle.Render(" "+format)+"\n", args...) +} + // SetNonInteractive enables or disables non-interactive mode. When true, // IsInteractive returns false regardless of TTY status. func SetNonInteractive(v bool) { @@ -37,37 +48,7 @@ func IsTTY() bool { return term.IsTerminal(int(os.Stdin.Fd())) } -// newForm creates a huh.Form with the Esc key bound to quit alongside -// the default ctrl+c. This allows users to cancel any interactive TUI -// flow by pressing Esc. -func newForm(groups ...*huh.Group) *huh.Form { - form := huh.NewForm(groups...) - keymap := huh.NewDefaultKeyMap() - keymap.Quit = key.NewBinding( - key.WithKeys("ctrl+c", "esc"), - key.WithHelp("ctrl+c / esc", "quit"), - ) - form.WithKeyMap(keymap) - return form -} - -// isCancelled returns true when the error from form.Run() indicates the -// user aborted (pressed Esc or Ctrl+C). +// isCancelled returns true when the user aborts an interactive TUI flow. func isCancelled(err error) bool { - return errors.Is(err, huh.ErrUserAborted) -} - -// ConfirmAction shows a yes/no confirmation prompt and returns the result. -// Returns (false, error) when the user presses Esc/Ctrl+C. -func ConfirmAction(prompt string) (bool, error) { - var result bool - form := newForm(huh.NewGroup( - huh.NewConfirm(). - Title(prompt). - Value(&result), - )) - if err := form.Run(); err != nil { - return false, err - } - return result, nil + return errors.Is(err, errCancelled) } diff --git a/internal/tui/tui_test.go b/internal/tui/tui_test.go index b3b265d..d2c6a77 100644 --- a/internal/tui/tui_test.go +++ b/internal/tui/tui_test.go @@ -1,9 +1,12 @@ package tui import ( + "io" + "os" + "strings" "testing" - "github.com/charmbracelet/huh" + tea "charm.land/bubbletea/v2" ) func TestIsInteractiveReturnsFalseInTestEnv(t *testing.T) { @@ -34,20 +37,239 @@ func TestIsTTYReturnsFalseInTestEnv(t *testing.T) { } } -func TestNewFormReturnsFormWithEscQuit(t *testing.T) { - form := newForm(huh.NewGroup( - huh.NewConfirm().Title("test"), - )) - if form == nil { - t.Fatal("newForm should return a non-nil form") +func TestIsCancelledRecognizesPromptCancel(t *testing.T) { + if !isCancelled(errCancelled) { + t.Fatal("isCancelled should recognize prompt cancellation") } } -func TestNewFormNotNil(t *testing.T) { - form := newForm(huh.NewGroup( - huh.NewInput().Title("test"), - )) - if form == nil { - t.Fatal("newForm should return non-nil form") +func TestPrintCancelledRendersRedIcon(t *testing.T) { + oldStdout := os.Stdout + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("pipe: %v", err) + } + os.Stdout = w + t.Cleanup(func() { os.Stdout = oldStdout }) + + printCancelled() + + if err := w.Close(); err != nil { + t.Fatalf("close writer: %v", err) + } + out, err := io.ReadAll(r) + if err != nil { + t.Fatalf("read output: %v", err) + } + if !strings.Contains(string(out), "\x1b[38;5;203m Cancelled.") { + t.Fatalf("cancel output should be red with Nerd Font icon prefix, got %q", string(out)) + } +} + +func TestPrintSuccessRendersGreenMessage(t *testing.T) { + oldStdout := os.Stdout + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("pipe: %v", err) + } + os.Stdout = w + t.Cleanup(func() { os.Stdout = oldStdout }) + + printSuccess("Provider '%s' added successfully.", "openai") + + if err := w.Close(); err != nil { + t.Fatalf("close writer: %v", err) + } + out, err := io.ReadAll(r) + if err != nil { + t.Fatalf("read output: %v", err) + } + if !strings.Contains(string(out), "\x1b[38;5;42m Provider 'openai' added successfully.") { + t.Fatalf("success output should be green, got %q", string(out)) + } +} + +func TestPromptInputMasksSecretAndPreservesValue(t *testing.T) { + m := &promptState{kind: promptKindInput, title: "Enter key", secret: true} + + m.Update(tea.KeyPressMsg(tea.Key{Text: "sk-test"})) + + if m.value != "sk-test" { + t.Fatalf("value = %q, want sk-test", m.value) + } + view := m.View().Content + if strings.Contains(view, "sk-test") { + t.Fatalf("secret view leaked value:\n%s", view) + } + if !strings.Contains(view, "*******") { + t.Fatalf("secret view missing mask:\n%s", view) + } +} + +func TestPromptInputHandlesPasteAndShowsCursor(t *testing.T) { + m := &promptState{kind: promptKindInput, title: "Enter URL"} + + m.Update(tea.PasteMsg{Content: "https://api.example.com/v1"}) + + if m.value != "https://api.example.com/v1" { + t.Fatalf("value = %q, want pasted URL", m.value) + } + view := m.View() + if view.Cursor == nil { + t.Fatal("input view should expose a cursor") + } + if !view.AltScreen { + t.Fatal("input view should use alternate screen") + } +} + +func TestPromptHelpRendersPlainKeysAndMutedActions(t *testing.T) { + m := &promptState{kind: promptKindInput, title: "Enter"} + + view := m.View().Content + if strings.Contains(view, "🥔") { + t.Fatalf("prompt title should not include emoji:\n%q", view) + } + if !strings.Contains(view, " \x1b[1mEnter") { + t.Fatalf("title should align with prompt text column:\n%q", view) + } + if !strings.Contains(view, " \x1b[1mEnter\x1b[m\n\n") { + t.Fatalf("title should have one blank line before content:\n%q", view) + } + if !strings.Contains(view, "\n enter") { + t.Fatalf("help should align with prompt text column:\n%q", view) + } + if strings.Contains(view, "\x1b[1menter") || strings.Contains(view, "\x1b[37menter") { + t.Fatalf("help key should not be bold:\n%q", view) + } + if strings.Contains(view, "\x1b[38;5;240menter") || strings.Contains(view, "\x1b[38;5;240mesc") { + t.Fatalf("help key should not be muted:\n%q", view) + } + if !strings.Contains(view, "\x1b[38;5;240m submit") { + t.Fatalf("help description should stay muted:\n%q", view) + } + if !strings.Contains(view, "\x1b[mesc\x1b[38;5;240m cancel") { + t.Fatalf("help cancel action should be muted while esc stays plain:\n%q", view) + } +} + +func TestPromptSelectMovesAndSelectsValue(t *testing.T) { + m := &promptState{ + kind: promptKindSelect, + title: "Select provider", + options: []promptOption{ + {Label: "OpenAI", Value: "openai"}, + {Label: "fal", Value: "fal"}, + }, + } + + if view := m.View().Content; !strings.Contains(view, " ") || strings.Contains(view, "> OpenAI") { + t.Fatalf("select view should use Nerd Font picker cursor:\n%q", view) + } + if view := m.View().Content; strings.Contains(view, "\x1b[7m") || !strings.Contains(view, "\x1b[38;5;42mOpenAI") { + t.Fatalf("focused select option should use color-only highlight:\n%q", view) + } + + m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyDown})) + _, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyEnter})) + + if cmd == nil { + t.Fatal("enter should quit the prompt") + } + if m.selected != "fal" { + t.Fatalf("selected = %q, want fal", m.selected) + } +} + +func TestPromptConfirmRendersTypedYesNoInput(t *testing.T) { + m := &promptState{kind: promptKindConfirm, title: "Confirm?"} + + m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeySpace})) + + if m.value != "" { + t.Fatalf("space should not change confirm input, got %q", m.value) + } + view := m.View() + if view.Cursor == nil { + t.Fatal("confirm view should expose a cursor") + } + if view.Cursor.Y != 1 { + t.Fatalf("confirm cursor Y = %d, want inline title row 1", view.Cursor.Y) + } + if strings.Contains(view.Content, "> Yes") || strings.Contains(view.Content, "> No") { + t.Fatalf("confirm view should not show picker marker:\n%q", view.Content) + } + if !strings.Contains(view.Content, "Confirm?\x1b[m [Y/n] ") { + t.Fatalf("confirm view should show inline typed default:\n%q", view.Content) + } + if strings.Contains(view.Content, "Confirm?\x1b[m\n\n") { + t.Fatalf("confirm input should stay on title line:\n%q", view.Content) + } +} + +func TestPromptConfirmKeepsInputOnLastQuestionLine(t *testing.T) { + m := &promptState{kind: promptKindConfirm, title: "Verification failed: unsupported protocol\nAdd anyway?"} + + view := m.View() + + if view.Cursor == nil { + t.Fatal("confirm view should expose a cursor") + } + if view.Cursor.Y != 2 { + t.Fatalf("confirm cursor Y = %d, want last title row 2", view.Cursor.Y) + } + if !strings.Contains(view.Content, "\n \x1b[1mAdd anyway?\x1b[m [Y/n] ") { + t.Fatalf("confirm input should stay on last question line:\n%q", view.Content) + } + if strings.Contains(view.Content, "unsupported protocol\x1b[m [Y/n]") { + t.Fatalf("confirm input should not attach to the error line:\n%q", view.Content) + } +} + +func TestPromptConfirmChooseYesNoAndCtrlCCancels(t *testing.T) { + m := &promptState{kind: promptKindConfirm, title: "Confirm?"} + + _, cmd := m.Update(tea.KeyPressMsg(tea.Key{Text: "y", Code: 'y'})) + if cmd != nil { + t.Fatal("typing y should not quit until enter") + } + _, cmd = m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyEnter})) + + if cmd == nil { + t.Fatal("enter should quit prompt") + } + if !m.confirm { + t.Fatal("y should confirm") + } + + m = &promptState{kind: promptKindConfirm, title: "Confirm?", confirm: true} + _, cmd = m.Update(tea.KeyPressMsg(tea.Key{Text: "n", Code: 'n'})) + if cmd != nil { + t.Fatal("typing n should not quit until enter") + } + _, cmd = m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyEnter})) + if cmd == nil { + t.Fatal("enter should quit prompt") + } + if m.confirm { + t.Fatal("n should reject") + } + + m = &promptState{kind: promptKindConfirm, title: "Confirm?"} + _, cmd = m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeyEnter})) + if cmd == nil { + t.Fatal("enter should quit prompt") + } + if !m.confirm { + t.Fatal("empty answer should default to yes") + } + + m = &promptState{kind: promptKindInput, title: "Enter"} + _, cmd = m.Update(tea.KeyPressMsg(tea.Key{Code: 'c', Mod: tea.ModCtrl})) + if cmd == nil { + t.Fatal("ctrl+c should quit prompt") + } + if !m.cancel { + t.Fatal("ctrl+c should cancel") } } diff --git a/internal/tui/use_picker.go b/internal/tui/use_picker.go index 84efc22..ebaa389 100644 --- a/internal/tui/use_picker.go +++ b/internal/tui/use_picker.go @@ -2,8 +2,7 @@ package tui import ( "fmt" - - "github.com/charmbracelet/huh" + "strings" "github.com/ncxton/potaco/internal/auth" ) @@ -24,6 +23,9 @@ func RunUsePicker() error { providerName, err := pickProvider(providers) if err != nil { + if isCancelled(err) { + return nil + } return err } if providerName == "" { @@ -32,6 +34,9 @@ func RunUsePicker() error { modelID, err := pickModel(providers, providerName) if err != nil { + if isCancelled(err) { + return nil + } return err } @@ -39,34 +44,28 @@ func RunUsePicker() error { return fmt.Errorf("set active provider: %w", err) } - fmt.Printf("Switched to provider '%s'.\n", providerName) + printSuccess("Switched to provider '%s'.", displayProviderName(providerName)) return nil } -// pickProvider renders a huh form listing connected providers and returns the +// pickProvider renders a provider prompt and returns the // chosen provider name. func pickProvider(providers []auth.ProviderInfo) (string, error) { - options := make([]huh.Option[string], 0, len(providers)) + options := make([]promptOption, 0, len(providers)) for _, p := range providers { - label := p.Name + label := displayProviderName(p.Name) if p.IsActive { label += " (active)" } label += " - " + p.Model - options = append(options, huh.NewOption(label, p.Name)) + options = append(options, promptOption{Label: label, Value: p.Name}) } - var providerName string - form := newForm(huh.NewGroup( - huh.NewSelect[string](). - Title("Select a provider:"). - Options(options...). - Value(&providerName), - )) - if err := form.Run(); err != nil { + providerName, err := runSelectPrompt("Select a provider:", options) + if err != nil { if isCancelled(err) { - fmt.Println("Cancelled.") - return "", nil + printCancelled() + return "", errCancelled } return "", fmt.Errorf("provider select: %w", err) } @@ -77,16 +76,11 @@ func pickProvider(providers []auth.ProviderInfo) (string, error) { // if so, prompts for the new model ID. It returns the model to set, or an // empty string to keep the existing configured model. func pickModel(providers []auth.ProviderInfo, providerName string) (string, error) { - var changeModel bool - confirmForm := newForm(huh.NewGroup( - huh.NewConfirm(). - Title("Change the model for this provider?"). - Value(&changeModel), - )) - if err := confirmForm.Run(); err != nil { + changeModel, err := runConfirmPrompt("Change the model for this provider?") + if err != nil { if isCancelled(err) { - fmt.Println("Cancelled.") - return "", nil + printCancelled() + return "", errCancelled } return "", fmt.Errorf("confirm model change: %w", err) } @@ -96,15 +90,11 @@ func pickModel(providers []auth.ProviderInfo, providerName string) (string, erro } modelID := currentModel(providers, providerName) - modelForm := newForm(huh.NewGroup( - huh.NewInput(). - Title("Enter model ID:"). - Value(&modelID), - )) - if err := modelForm.Run(); err != nil { + modelID, err = runInputPrompt("Enter model ID:", modelID, false) + if err != nil { if isCancelled(err) { - fmt.Println("Cancelled.") - return "", nil + printCancelled() + return "", errCancelled } return "", fmt.Errorf("model input: %w", err) } @@ -121,3 +111,10 @@ func currentModel(providers []auth.ProviderInfo, providerName string) string { } return "" } + +func displayProviderName(name string) string { + if strings.EqualFold(name, "fal") { + return "Fal" + } + return name +} diff --git a/internal/tui/use_picker_test.go b/internal/tui/use_picker_test.go index 5879a9f..392b386 100644 --- a/internal/tui/use_picker_test.go +++ b/internal/tui/use_picker_test.go @@ -2,6 +2,8 @@ package tui import ( "testing" + + "github.com/ncxton/potaco/internal/auth" ) func TestRunUsePickerNoProvidersReturnsError(t *testing.T) { @@ -11,3 +13,44 @@ func TestRunUsePickerNoProvidersReturnsError(t *testing.T) { t.Fatal("expected error when no providers connected") } } + +func TestRunUsePickerCancelModelPromptDoesNotSwitchProvider(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + mgr, err := auth.New() + if err != nil { + t.Fatalf("init auth: %v", err) + } + if err := mgr.AddProvider("openai", "openai", "sk-openai"); err != nil { + t.Fatalf("add openai: %v", err) + } + if err := mgr.AddProvider("fal", "fal", "fal-test"); err != nil { + t.Fatalf("add fal: %v", err) + } + if err := mgr.SetActiveProvider("openai", "gpt-image-1"); err != nil { + t.Fatalf("set active: %v", err) + } + + origSelect := runSelectPrompt + origConfirm := runConfirmPrompt + runSelectPrompt = func(string, []promptOption) (string, error) { + return "fal", nil + } + runConfirmPrompt = func(string) (bool, error) { + return false, errCancelled + } + t.Cleanup(func() { + runSelectPrompt = origSelect + runConfirmPrompt = origConfirm + }) + + if err := RunUsePicker(); err != nil { + t.Fatalf("RunUsePicker: %v", err) + } + provider, model, err := mgr.GetActiveProvider() + if err != nil { + t.Fatalf("get active provider: %v", err) + } + if provider != "openai" || model != "gpt-image-1" { + t.Fatalf("active provider changed after cancel: provider=%q model=%q", provider, model) + } +} diff --git a/skills/potaco/SKILL.md b/skills/potaco/SKILL.md index a49c31d..3afa9cb 100644 --- a/skills/potaco/SKILL.md +++ b/skills/potaco/SKILL.md @@ -2,12 +2,12 @@ name: potaco version: 1.5.0 description: | - Use when you want to generate images from text prompts, edit existing images, perform inpainting or outpainting, set up provider credentials, discover models, or troubleshoot potaco CLI failures. Covers OpenAI, fal, Vercel AI Gateway, and custom OpenAI-compatible providers. + Use when you want to generate images from text prompts, edit existing images, perform inpainting or outpainting, set up provider credentials, discover models, or troubleshoot potaco CLI failures. Covers OpenAI, fal, Replicate, and custom OpenAI-compatible providers. --- # Potaco CLI -Potaco generates and edits images through OpenAI, fal, Vercel AI Gateway, and custom OpenAI-compatible providers. +Potaco generates and edits images through OpenAI, fal, Replicate, and custom OpenAI-compatible providers. ## Non-negotiable security rules @@ -49,7 +49,7 @@ Preferred user-run credential setup: potaco auth add ``` -Available built-in providers: `openai`, `fal`, `vercel`, `custom`. `custom` and non-built-in provider names also prompt for base URL. If no provider is configured, launch auth only when the user explicitly asks you to start that flow; otherwise instruct them to run it themselves. +Available built-in providers: `openai`, `fal`, `replicate`, `custom`. `custom` and non-built-in provider names also prompt for base URL. Vercel AI Gateway can be configured as an OpenAI-compatible alias with `--type openai --base-url https://ai-gateway.vercel.sh/v1`. If no provider is configured, launch auth only when the user explicitly asks you to start that flow; otherwise instruct them to run it themselves. ## On-demand references diff --git a/skills/potaco/references/commands.md b/skills/potaco/references/commands.md index 75f9176..f67d5ca 100644 --- a/skills/potaco/references/commands.md +++ b/skills/potaco/references/commands.md @@ -46,7 +46,7 @@ Do not enter credentials yourself unless the user explicitly asks you to start a ```sh potaco auth add openai potaco auth add custom -potaco auth add openrouter --type openai-compatible --base-url https://openrouter.ai/api/v1 +potaco auth add openrouter --type openai --base-url https://openrouter.ai/api/v1 potaco auth add staging-openai --type openai --base-url https://proxy.example/v1 potaco auth list potaco auth list --json @@ -57,7 +57,7 @@ potaco status potaco status --json ``` -Built-in names (`openai`, `fal`, `vercel`) use preset base URLs. Any other auth name, including aliases using `--type openai`, `--type fal`, or `--type vercel`, must include `--base-url` or have `base_url` in config. +Built-in names (`openai`, `fal`, `replicate`) use preset base URLs. Any other auth name, including aliases using `--type openai`, `--type fal`, or `--type replicate`, must include `--base-url` or have `base_url` in config. ## Generate @@ -80,9 +80,9 @@ potaco edit --prompt "extend the landscape" --image photo.png --extend right=256 potaco edit --prompt "test edit" --image input.png --dry-run ``` -Vercel does not support editing. +Replicate does not support editing through Potaco because edit inputs are model-specific. -## Models and info +## Models ```sh potaco models @@ -91,8 +91,6 @@ potaco models list openai potaco models openai potaco models --json potaco models list --json -potaco info image.png -potaco info image.png --json ``` Interactive `models` persists selection. `models list` only displays. In non-interactive mode, `models` behaves like `models list`. diff --git a/skills/potaco/references/configuration.md b/skills/potaco/references/configuration.md index deef86a..212d838 100644 --- a/skills/potaco/references/configuration.md +++ b/skills/potaco/references/configuration.md @@ -32,13 +32,14 @@ Security: There is no provider preset model. Use `potaco models` or `potaco config set model `. Generation capability is assumed available. Edit capability is not auto-detected; configure it per model with `potaco config set model.edit true` or `potaco config set providers..models..edit true`. +For Replicate, set the model manually in `owner/model` format. ### Base URL 1. `--base-url` CLI flag 2. `POTACO_BASE_URL` env var 3. `base_url` from the provider config in config file -4. Provider preset default (openai: `https://api.openai.com/v1`, fal: `https://fal.run`, vercel: `https://ai-gateway.vercel.sh/v1`) +4. Provider preset default (openai: `https://api.openai.com/v1`, fal: `https://fal.run`, replicate: `https://api.replicate.com/v1`) 5. Empty (custom and non-built-in provider aliases have no usable setup preset; error if not configured) For `custom` and any non-built-in provider name, the base URL is required during setup, then persisted. Confirm remote custom or alias endpoints are trusted before saving. @@ -109,11 +110,11 @@ providers: potaco config show ``` -Shows config path, active provider/model, and per-provider model/base URL/retries/timeout. +Prints the raw `config.yaml` contents. ```sh potaco config set model gpt-image-2 -potaco config set providers.vercel.model openai/gpt-image-2 +potaco config set providers.replicate.model black-forest-labs/flux-schnell potaco config set model.edit true potaco config set providers.openai.models.gpt-image-2.edit true potaco config set base_url https://api.example.com/v1 @@ -159,7 +160,7 @@ Presets store only a base URL: |----------|----------| | openai | `https://api.openai.com/v1` | | fal | `https://fal.run` | -| vercel | `https://ai-gateway.vercel.sh/v1` | +| replicate | `https://api.replicate.com/v1` | | custom | (no preset; user-supplied) | No preset default model exists. diff --git a/skills/potaco/references/editing.md b/skills/potaco/references/editing.md index 1034ec8..330f363 100644 --- a/skills/potaco/references/editing.md +++ b/skills/potaco/references/editing.md @@ -80,13 +80,12 @@ The expanded canvas is checked against a pixel budget. If the expanded dimension **Dry run output**: mode is `"outpaint"` with the parsed extend config shown. -## Edit with Vercel provider +## Edit with Replicate provider -The Vercel AI Gateway does not support image editing. If the active provider is `vercel`, the edit command returns: +Replicate edit inputs are model-specific, so Potaco does not expose editing through the `replicate` provider. If the active provider is `replicate`, the edit command returns: ``` -Error: Image editing is not supported by the Vercel AI Gateway provider. -Hint: Use 'potaco use' to switch to a provider that supports editing. + Image editing is not supported by the 'replicate' provider. Use 'potaco use' to switch to a provider that supports editing. ``` Switch to a provider that supports editing: diff --git a/skills/potaco/references/providers.md b/skills/potaco/references/providers.md index e6a5399..e71ab07 100644 --- a/skills/potaco/references/providers.md +++ b/skills/potaco/references/providers.md @@ -1,14 +1,14 @@ # Providers -Potaco supports built-in `openai`, `fal`, `vercel`, plus aliases and `custom` endpoints. +Potaco supports built-in `openai`, `fal`, `replicate`, plus aliases and `custom` endpoints. ## Provider Support Matrix -| Feature | openai | fal | vercel | custom | +| Feature | openai | fal | replicate | custom | |---------|--------|-----|--------|--------| | Generation | Yes | Yes | Yes | Yes | | Editing | Yes | Yes | No | Yes | -| Model discovery | GET `/v1/models` | GET `api.fal.ai/v1/models` | GET `/v1/models` | GET `/v1/models` | +| Model discovery | GET `/v1/models` | GET `api.fal.ai/v1/models` | Manual `owner/model` config | GET `/v1/models` | | Auth header | `Bearer ` | `Key ` | `Bearer ` | `Bearer ` | | Requires `--base-url` | No | No | No | Yes | @@ -18,12 +18,14 @@ Built-in provider names use preset base URLs. Any other provider name, including - `openai`: base `https://api.openai.com/v1`, bearer auth, gen/edit via Images API, discovery filters to known image IDs (`gpt-image-2`). - `fal`: base `https://fal.run`, discovery `https://api.fal.ai/v1/models?category=image`, auth `Key `, edits use `/image-to-image`. -- `vercel`: base `https://ai-gateway.vercel.sh/v1`, bearer auth, generation only, discovery filters `type == "image"`, model IDs use `provider/model`. +- `replicate`: base `https://api.replicate.com/v1`, bearer auth, generation only, model IDs use `owner/model` and generation posts to `/models/{owner}/{model}/predictions`. Configure the model manually with `potaco config set model `. - `custom`: no default base URL, bearer auth, assumes OpenAI-compatible `/v1/images/generations`, `/v1/images/edits`, and `/v1/models`. -Vercel verification: -1. `GET /v1/models` to confirm the endpoint is reachable (5xx = fail). -2. `GET /v1/models/openai/gpt-image-2/endpoints` to validate the API key (401/403 = invalid). +Vercel AI Gateway is OpenAI-compatible. Configure it as an alias: + +```sh +potaco auth add vercel --type openai --api-key --base-url https://ai-gateway.vercel.sh/v1 +``` ## Custom provider security diff --git a/skills/potaco/references/troubleshooting.md b/skills/potaco/references/troubleshooting.md index 5d1d5ad..d30a947 100644 --- a/skills/potaco/references/troubleshooting.md +++ b/skills/potaco/references/troubleshooting.md @@ -27,16 +27,15 @@ Friendly errors print to stderr/TTY. Raw `UserError` details append to `~/.potac #### "No active provider configured" ``` -Error: No active provider configured. -Hint: Run 'potaco auth add ' to connect one. + No active provider configured. Run 'potaco auth add ' to connect one. ``` -Fix: tell the user to run `potaco auth add openai` (or `fal`/`vercel`) interactively unless they explicitly asked you to configure credentials. +Fix: tell the user to run `potaco auth add openai` (or `fal`/`replicate`) interactively unless they explicitly asked you to configure credentials. #### "No API key found for the active provider" ``` -Error: No API key found for the active provider. + No API key found for the active provider. ``` Fix: the CLI hint may mention `--api-key`, but the safer instruction is for the user to re-add the provider interactively: `potaco auth add openai`. @@ -44,8 +43,7 @@ Fix: the CLI hint may mention `--api-key`, but the safer instruction is for the #### "Could not load configuration" ``` -Error: Could not load configuration. -Hint: Check that ~/.potaco/ is readable. + Could not load configuration. Check that ~/.potaco/ is readable. ``` Fix: check permissions on `~/.potaco/` and `~/.potaco/config.yaml`. If corrupted, regenerate via interactive auth. @@ -53,8 +51,7 @@ Fix: check permissions on `~/.potaco/` and `~/.potaco/config.yaml`. If corrupted #### "Image generation failed" ``` -Error: Image generation failed. -Hint: Check your API key, network connection, and model name. + Image generation failed. Check your API key, network connection, and model name. ``` Fix: if safe, check `~/.potaco/debug.log` for the raw error. Common causes: @@ -67,12 +64,10 @@ Fix: if safe, check `~/.potaco/debug.log` for the raw error. Common causes: Same as generation but for the edit endpoint. Also check that the source image file exists and is a valid PNG/JPEG/WebP. -#### "Image editing is not supported by the Vercel AI Gateway provider" +#### "Image editing is not supported by the Replicate provider" ``` -Error: Image editing is not supported by the Vercel AI Gateway provider. -Hint: Use 'potaco use openai' or 'potaco use fal' to switch to a provider -that supports editing. + Image editing is not supported by the Replicate provider. Use 'potaco use openai' or 'potaco use fal' to switch to a provider that supports editing. ``` Fix: Switch to a provider that supports editing: `potaco use openai`, `potaco use fal`, or `potaco use custom`. @@ -80,8 +75,7 @@ Fix: Switch to a provider that supports editing: `potaco use openai`, `potaco us #### "Cannot write to ''" ``` -Error: Cannot write to '/tmp/output.png'. -Hint: Check that the path is valid and you have write permissions. + Cannot write to '/tmp/output.png'. Check that the path is valid and you have write permissions. ``` Fix: output path must be a writable file path and parent directory must exist. @@ -89,8 +83,7 @@ Fix: output path must be a writable file path and parent directory must exist. #### "'' is a directory, not a file." ``` -Error: '/output/' is a directory, not a file. -Hint: Specify a filename ending in .png or .jpeg, or omit -o to auto-generate one. + '/output/' is a directory, not a file. Specify a filename ending in .png or .jpeg, or omit -o to auto-generate one. ``` Fix: provide a filename, not a directory path. Example: `-o /output/result.png`. @@ -98,8 +91,7 @@ Fix: provide a filename, not a directory path. Example: `-o /output/result.png`. #### "Cannot write multiple images to stdout" ``` -Error: Cannot write multiple images to stdout. -Hint: Use --output to write to files, or request a single image with --n 1. + Cannot write multiple images to stdout. Use --output to write to files, or request a single image with --n 1. ``` Fix: use `--output` or set `--n 1`. @@ -107,9 +99,7 @@ Fix: use `--output` or set `--n 1`. #### "Could not decode the image returned by the provider" ``` -Error: Could not decode the image returned by the provider. -Hint: The provider may have returned an unsupported image format. Try a -different model or check '~/.potaco/debug.log' for details. + Could not decode the image returned by the provider. The provider may have returned an unsupported image format. Try a different model or check '~/.potaco/debug.log' for details. ``` Fix: try a different model or inspect the debug log if safe. @@ -117,19 +107,18 @@ Fix: try a different model or inspect the debug log if safe. #### "Unknown provider: " ``` -Error: unknown provider: foo (available: [custom fal openai vercel]) + unknown provider: foo (available: [custom fal openai replicate]) ``` -Fix: Use one of the registered provider names: `openai`, `fal`, `vercel`, or `custom`. +Fix: Use one of the registered provider names: `openai`, `fal`, `replicate`, or `custom`. -#### "A base URL is required for OpenAI-compatible providers." +#### "A base URL is required for custom provider aliases." ``` -Error: A base URL is required for OpenAI-compatible providers. -Hint: Use --base-url, set POTACO_BASE_URL, or run 'potaco config set providers.openrouter.base_url '. + A base URL is required for this provider. Use --base-url, set POTACO_BASE_URL, or run 'potaco config set providers.openrouter.base_url '. ``` -Fix: prefer `potaco auth add custom` or `potaco auth add --type openai-compatible` interactively with a trusted base URL. Use `--base-url`, `POTACO_BASE_URL`, or `potaco config set providers..base_url ` only in explicitly approved non-interactive workflows. +Fix: prefer `potaco auth add custom` or `potaco auth add --type openai` interactively with a trusted base URL. Use `--base-url`, `POTACO_BASE_URL`, or `potaco config set providers..base_url ` only in explicitly approved non-interactive workflows. ## Debugging Workflow @@ -162,7 +151,6 @@ Common JSON outputs: - `potaco auth list --json`: array of providers with name, model, base_url, has_key, is_active. - `potaco models --json` / `potaco models list --json`: array of models with id and display_name. - `potaco version --json`: current, latest, update_available fields. -- `potaco info --json`: path, format, width, height, file_size, color_model. - `potaco gen --json` / `potaco edit --json`: path, format, dimensions, model, params, latency_ms. Use `--json` when an agent needs to parse output programmatically. JSON output goes to stdout, error output goes to stderr.