From b6088da01f3a8ccc393dfb7906edaa7939c3a860 Mon Sep 17 00:00:00 2001 From: Sai Guvvala Date: Fri, 2 Oct 2026 16:09:34 -0700 Subject: [PATCH 1/3] feat: remember image picker settings and save folder --- docs/image-generation-saving.md | 9 +- pkg/custom/image_picker.go | 51 ++- pkg/custom/image_picker_command.go | 3 + pkg/custom/image_picker_folder.go | 419 ++++++++++++++++++++++ pkg/custom/image_picker_folder_test.go | 405 +++++++++++++++++++++ pkg/custom/image_picker_session.go | 54 ++- pkg/custom/image_picker_session_test.go | 93 ++++- pkg/custom/image_picker_state.go | 283 +++++++++++++++ pkg/custom/image_picker_state_test.go | 356 ++++++++++++++++++ pkg/custom/image_picker_test.go | 12 + pkg/custom/image_picker_view.go | 17 +- scripts/check-image-picker-integration.py | 9 +- scripts/check-image-picker-reuse.py | 299 +++++++++++++++ 13 files changed, 1979 insertions(+), 31 deletions(-) create mode 100644 pkg/custom/image_picker_folder.go create mode 100644 pkg/custom/image_picker_folder_test.go create mode 100644 pkg/custom/image_picker_state.go create mode 100644 pkg/custom/image_picker_state_test.go create mode 100644 scripts/check-image-picker-reuse.py diff --git a/docs/image-generation-saving.md b/docs/image-generation-saving.md index d65aded3..3877a320 100644 --- a/docs/image-generation-saving.md +++ b/docs/image-generation-saving.md @@ -19,7 +19,14 @@ zsh, fish and PowerShell 7. In other or unidentified shells, generation remains available, but command printing is disabled with an explanation in the picker. After a successful generation, the picker reopens below the saved result with -the active prompt and settings. Choices last until this invocation exits. +the active prompt and settings. It remembers the last submitted prompt and +settings for next time; canceled edits are discarded. Settings stay local to +your user configuration folder and do not affect commands that supply flags. + +Choose **Save to** to use the default folder, current directory, or another +existing folder. Tab completes folder names. Saved settings are optional: if +they cannot be written, a warning appears and generation can still continue. +Unknown or unreadable saved settings are left alone. Explicit image flags, output formats, piped input and redirected output retain the direct command behavior described below. diff --git a/pkg/custom/image_picker.go b/pkg/custom/image_picker.go index ac1fae55..f06051db 100644 --- a/pkg/custom/image_picker.go +++ b/pkg/custom/image_picker.go @@ -26,9 +26,10 @@ type imagePickerOptions struct { // PromptSet lets an explicit empty prompt override the active draft. PromptSet bool // Shell selects quoting for the displayed command, never execution. - Shell string - initial *imagePickerSettings - resuming bool + Shell string + initial *imagePickerSettings + resuming bool + initialNote string } // imagePickerResult contains arguments for the existing images generate command. @@ -51,6 +52,7 @@ func runImagePicker(parent context.Context, input, output *os.File, options imag if err != nil { return imagePickerResult{}, err } + defer model.cancelFolderWork() if input == nil || output == nil || !term.IsTerminal(input.Fd()) || !term.IsTerminal(output.Fd()) { return imagePickerResult{}, errors.New("the image picker needs terminal input and output") } @@ -189,7 +191,7 @@ func (w *imagePickerOutput) Err() error { } type imagePickerSettings struct { - prompt, model, size, quality, background, format, count string + prompt, model, size, quality, background, format, count, outputDir string } type imagePicker struct { @@ -208,6 +210,7 @@ type imagePicker struct { dark bool note string result imagePickerResult + folder imagePickerFolder } type imagePickerRow struct { @@ -228,6 +231,7 @@ func newImagePicker(options imagePickerOptions) (*imagePicker, error) { m.draft, m.cursor = []rune(m.settings.prompt), len([]rune(m.settings.prompt)) } m.shell = options.Shell + m.note = options.initialNote return m, nil } @@ -243,10 +247,12 @@ func (m *imagePicker) Update(message tea.Msg) (tea.Model, tea.Cmd) { // Cancellation can still stop a submitted action before the picker exits. switch msg := message.(type) { case imagePickerStopMsg: + m.cancelFolderWork() m.result = imagePickerResult{Canceled: true, ExitCode: msg.code} return m, tea.Quit case tea.KeyPressMsg: if msg.String() == "ctrl+c" { + m.cancelFolderWork() m.result = imagePickerResult{Canceled: true} return m, tea.Quit } @@ -256,14 +262,20 @@ func (m *imagePicker) Update(message tea.Msg) (tea.Model, tea.Cmd) { return m, nil } switch msg := message.(type) { - + case imagePickerFolderMsg: + return m, m.finishFolderWork(msg) case tea.BackgroundColorMsg: m.dark = msg.IsDark() case tea.WindowSizeMsg: m.width, m.height = msg.Width, msg.Height m.clampCommandOffset() case tea.PasteMsg: - if m.focus == "prompt" { + if m.folder.busy == "submit" || m.folder.busy == "select" { + return m, nil + } + if m.focus == "path" { + m.insertFolderPath(msg.Content) + } else if m.focus == "prompt" { m.insertPrompt(msg.Content) } else { m.note = "Tab to Prompt to paste text." @@ -271,9 +283,21 @@ func (m *imagePicker) Update(message tea.Msg) (tea.Model, tea.Cmd) { case tea.KeyPressMsg: key := msg.String() if key == "esc" { + if m.page == "folder" || m.page == "path" || m.folder.busy == "submit" { + return m, m.backFolder() + } m.focusPrompt() return m, nil } + if m.folder.busy == "submit" || m.folder.busy == "select" { + return m, nil + } + if m.focus == "path" { + if m.width >= 40 && m.height >= 12 { + return m, m.editFolderPath(msg) + } + return m, nil + } // Esc immediately followed by another key can arrive as an Alt key. // Leave menus first; no Alt shortcuts are assigned // there. In the prompt itself, Alt+Enter still inserts a newline. @@ -292,6 +316,7 @@ func (m *imagePicker) Update(message tea.Msg) (tea.Model, tea.Cmd) { } } if key == "q" && m.focus != "prompt" { + m.cancelFolderWork() m.result = imagePickerResult{Canceled: true} return m, tea.Quit } @@ -363,12 +388,17 @@ func (m *imagePicker) Update(message tea.Msg) (tea.Model, tea.Cmd) { // Esc closes any menu without applying its highlighted choice. Committed // settings and prompt text remain intact. func (m *imagePicker) focusPrompt() { + m.cancelFolderWork() m.page, m.focus, m.selected, m.note = "settings", "prompt", 0, "" } func (m *imagePicker) rows() []imagePickerRow { s := m.settings switch m.page { + case "folder": + return m.folderRows() + case "path": + return m.folderCompletionRows() case "choose": return m.choices(m.field) case "more": @@ -379,7 +409,7 @@ func (m *imagePicker) rows() []imagePickerRow { {id: "size", label: "Size", value: imagePickerSizeLabel(s.size)}, {id: "quality", label: "Quality", value: imagePickerTitle(s.quality)}, } - return append(rows, imagePickerRow{id: "more", label: "More options", value: strings.ToUpper(s.format) + " / " + s.count + " image(s)"}) + return append(rows, m.folderRow(), imagePickerRow{id: "more", label: "More options", value: strings.ToUpper(s.format) + " / " + s.count + " image(s)"}) } func (m *imagePicker) cycleFocus(backwards bool) { @@ -502,6 +532,8 @@ func (m *imagePicker) apply(value string) { func (m *imagePicker) activate(row imagePickerRow) tea.Cmd { switch row.id { + case "folder", "folder-default", "folder-current", "folder-path", "folder-back": + return m.activateFolder(row.id) case "choice": m.apply(row.value) m.page, m.selected = m.returnPage, m.returnSelected @@ -517,8 +549,7 @@ func (m *imagePicker) activate(row imagePickerRow) tea.Cmd { return nil } if m.validPrompt() { - m.result = imagePickerResult{Args: m.settings.args(), PrintOnly: row.id == "print", settings: m.settings, shell: m.shell} - return tea.Quit + return m.submitWithFolder(row.id == "print") } case "back": return m.back() @@ -552,6 +583,8 @@ func (m *imagePicker) validPrompt() bool { func (m *imagePicker) back() tea.Cmd { m.note = "" switch m.page { + case "folder", "path": + return m.backFolder() case "more": m.page = "settings" m.selected = len(m.rows()) - 1 diff --git a/pkg/custom/image_picker_command.go b/pkg/custom/image_picker_command.go index 0a06cd08..054833b1 100644 --- a/pkg/custom/image_picker_command.go +++ b/pkg/custom/image_picker_command.go @@ -17,6 +17,9 @@ func (s imagePickerSettings) args() []string { // setting changes even when the prompt tail needs to be truncated. args := []string{"images", "generate", "--model", s.model, "--size", s.size, "--quality", s.quality, "--output-format", s.format, "--background", s.background, "--count", s.count} + if s.outputDir != "" { + args = append(args, "--output-dir", s.outputDir) + } return append(args, "--prompt", prompt) } diff --git a/pkg/custom/image_picker_folder.go b/pkg/custom/image_picker_folder.go new file mode 100644 index 00000000..2562160b --- /dev/null +++ b/pkg/custom/image_picker_folder.go @@ -0,0 +1,419 @@ +package custom + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "sort" + "strings" + "syscall" + + tea "charm.land/bubbletea/v2" + "github.com/openai/openai-cli/internal/imageoutput" +) + +// Folder work runs outside the input loop. Revisions make late filesystem +// replies harmless after editing, backing out, or canceling the picker. +type imagePickerFolder struct { + draft []rune + cursor int + matches []string + match int + returnPage string + returnSelected int + revision uint64 + busy string + cancel context.CancelFunc +} + +type imagePickerFolderMsg struct { + revision uint64 + kind, path string + matches []string + settings imagePickerSettings + printOnly bool + err error +} + +func (m *imagePicker) folderRow() imagePickerRow { + value := m.settings.outputDir + if value == "" { + value = "~/Downloads/gpt-images/ (Default)" + } else { + value = imagePickerFolderLabel(value) + } + return imagePickerRow{id: "folder", label: "Save to", value: value} +} + +// Abbreviate only the visible label. Stored paths and command arguments stay +// absolute, and a neighboring directory with the same prefix stays unchanged. +func imagePickerFolderLabel(path string) string { + home, err := os.UserHomeDir() + if err != nil || !filepath.IsAbs(path) { + return path + } + relative, err := filepath.Rel(home, path) + if err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) { + return path + } + if relative == "." { + return "~" + } + return "~/" + filepath.ToSlash(relative) +} + +func (m *imagePicker) folderRows() []imagePickerRow { + return []imagePickerRow{ + {id: "folder-default", label: "Default", value: "~/Downloads/gpt-images/"}, + {id: "folder-current", label: "Current folder"}, + {id: "folder-path", label: "Enter path"}, + {id: "folder-back", label: "Back to settings"}, + } +} + +func (m *imagePicker) folderCompletionRows() []imagePickerRow { + rows := make([]imagePickerRow, 0, len(m.folder.matches)) + for _, path := range m.folder.matches { + rows = append(rows, imagePickerRow{id: "folder-match", label: imagePickerLine(path)}) + } + if len(rows) == 0 { + rows = append(rows, imagePickerRow{label: "Tab completes directory names"}) + } + return rows +} + +func (m *imagePicker) activateFolder(id string) tea.Cmd { + switch id { + case "folder": + m.folder.returnPage, m.folder.returnSelected = m.page, m.selected + m.page, m.focus, m.selected, m.note = "folder", "options", 0, "" + case "folder-default": + m.commitFolder("") + case "folder-current": + return m.checkFolder("select", ".", false) + case "folder-path": + m.openFolderEditor(m.settings.outputDir) + case "folder-back": + return m.backFolder() + } + return nil +} + +func (m *imagePicker) openFolderEditor(path string) { + m.cancelFolderWork() + m.page, m.focus, m.selected = "path", "path", 0 + m.folder.draft, m.folder.cursor = []rune(path), len([]rune(path)) + m.folder.match = -1 + m.note = "" +} + +func (m *imagePicker) commitFolder(path string) { + m.cancelFolderWork() + m.settings.outputDir, m.commandOffset = path, 0 + m.page, m.focus, m.selected = m.folder.returnPage, "options", m.folder.returnSelected + if m.page == "" { + m.page = "settings" + } + m.note = "" +} + +func (m *imagePicker) backFolder() tea.Cmd { + m.cancelFolderWork() + m.note = "" + switch m.page { + case "path": + m.page, m.focus, m.selected = "folder", "options", 2 + case "folder": + m.page, m.focus, m.selected = m.folder.returnPage, "options", m.folder.returnSelected + default: + m.focus = "prompt" + } + return nil +} + +func (m *imagePicker) cancelFolderWork() { + if m.folder.cancel != nil { + m.folder.cancel() + m.folder.cancel = nil + } + m.folder.revision++ + m.folder.busy = "" + m.folder.matches, m.folder.match = nil, -1 +} + +func (m *imagePicker) folderWork(kind string) (context.Context, uint64) { + m.cancelFolderWork() + ctx, cancel := context.WithCancel(context.Background()) + m.folder.cancel, m.folder.busy = cancel, kind + return ctx, m.folder.revision +} + +func (m *imagePicker) insertFolderPath(text string) { + m.cancelFolderWork() + tail := append([]rune(nil), m.folder.draft[m.folder.cursor:]...) + m.folder.draft = append(m.folder.draft[:m.folder.cursor], []rune(text)...) + m.folder.draft = append(m.folder.draft, tail...) + m.folder.cursor += len([]rune(text)) + m.note = "" +} + +func (m *imagePicker) editFolderPath(key tea.KeyPressMsg) tea.Cmd { + f := &m.folder + switch key.String() { + case "tab", "shift+tab": + if f.busy != "" { + return nil + } + if len(f.matches) > 0 { + delta := 1 + if key.String() == "shift+tab" { + delta = -1 + } + f.match = (f.match + len(f.matches) + delta) % len(f.matches) + m.selected = f.match + return nil + } + path := string(f.draft) + ctx, revision := m.folderWork("complete") + m.note = "Finding folders…" + return func() tea.Msg { + matches, err := imagePickerCompleteFolders(ctx, path) + return imagePickerFolderMsg{revision: revision, kind: "complete", path: path, matches: matches, err: err} + } + case "enter": + if f.busy != "" { + return nil + } + if f.match >= 0 && f.match < len(f.matches) { + m.openFolderEditor(f.matches[f.match]) + return nil + } + return m.checkFolder("select", string(f.draft), false) + case "up", "down": + if len(f.matches) > 0 { + if key.String() == "up" { + f.match = max(0, f.match-1) + } else { + f.match = min(len(f.matches)-1, f.match+1) + } + m.selected = f.match + } + case "left": + f.cursor = max(0, f.cursor-1) + case "right": + f.cursor = min(len(f.draft), f.cursor+1) + case "home", "ctrl+a": + f.cursor = 0 + case "end", "ctrl+e": + f.cursor = len(f.draft) + case "ctrl+u": + m.cancelFolderWork() + f.draft, f.cursor = f.draft[f.cursor:], 0 + case "backspace": + m.cancelFolderWork() + if f.cursor > 0 { + f.draft = append(f.draft[:f.cursor-1], f.draft[f.cursor:]...) + f.cursor-- + } + case "delete": + m.cancelFolderWork() + if f.cursor < len(f.draft) { + f.draft = append(f.draft[:f.cursor], f.draft[f.cursor+1:]...) + } + default: + if key.Text != "" { + m.insertFolderPath(key.Text) + } + } + return nil +} + +func (m *imagePicker) checkFolder(kind, path string, printOnly bool) tea.Cmd { + ctx, revision := m.folderWork(kind) + snapshot := m.settings + m.note = "Checking save folder… Esc cancels" + return func() tea.Msg { + msg := imagePickerFolderMsg{revision: revision, kind: kind, path: path, settings: snapshot, printOnly: printOnly} + if err := ctx.Err(); err != nil { + msg.err = err + return msg + } + if kind == "submit" { + msg.path, msg.err = imageoutput.ResolveDirectory(path) + } else { + msg.path, msg.err = imagePickerExistingFolder(path) + } + if err := ctx.Err(); err != nil { + msg.err = err + } + return msg + } +} + +func (m *imagePicker) submitWithFolder(printOnly bool) tea.Cmd { + if printOnly && m.settings.outputDir == "" { + return m.finishSubmission(m.settings, true) + } + return m.checkFolder("submit", m.settings.outputDir, printOnly) +} + +func (m *imagePicker) finishSubmission(settings imagePickerSettings, printOnly bool) tea.Cmd { + m.result = imagePickerResult{Args: settings.args(), PrintOnly: printOnly, settings: settings, shell: m.shell} + return tea.Quit +} + +func (m *imagePicker) finishFolderWork(msg imagePickerFolderMsg) tea.Cmd { + if msg.revision != m.folder.revision || msg.kind != m.folder.busy { + return nil + } + if m.folder.cancel != nil { + m.folder.cancel() + m.folder.cancel = nil + } + m.folder.busy = "" + if errors.Is(msg.err, context.Canceled) { + return nil + } + if msg.kind == "complete" { + m.note = "" + if msg.err != nil { + m.note = "Could not read that folder." + return nil + } + m.folder.matches, m.folder.match, m.selected = msg.matches, -1, 0 + if len(msg.matches) == 1 { + m.openFolderEditor(msg.matches[0]) + } else if len(msg.matches) == 0 { + m.note = "No matching folders." + } + return nil + } + if msg.kind == "submit" && msg.settings != m.settings { + m.note = "" + return nil + } + if msg.err != nil { + path := string(m.folder.draft) + if msg.kind == "submit" { + path = msg.settings.outputDir + if path == "" { + path = "~/Downloads/gpt-images/" + } + m.folder.returnPage, m.folder.returnSelected = "settings", 0 + } else if m.page != "path" { + path = "." + } + m.openFolderEditor(path) + m.note = msg.err.Error() + return nil + } + if msg.kind == "select" { + m.commitFolder(msg.path) + return nil + } + if m.width < 40 || m.height < 12 { + m.note = "Resize before generating." + return nil + } + // Empty retains default intent; an explicit path is stable across cwd changes. + if msg.settings.outputDir != "" { + msg.settings.outputDir = msg.path + } + m.settings = msg.settings + return m.finishSubmission(msg.settings, msg.printOnly) +} + +func imagePickerAbsoluteFolder(path string) (string, error) { + if strings.ContainsRune(path, 0) { + return "", errors.New("Remove the NUL character from the folder path.") + } + if path == "~" || strings.HasPrefix(path, "~/") || strings.HasPrefix(path, `~\`) { + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + if path == "~" { + path = home + } else { + path = filepath.Join(home, path[2:]) + } + } + return filepath.Abs(path) +} + +func imagePickerExistingFolder(path string) (string, error) { + if path == "" { + return "", errors.New("Enter an existing folder, or choose Default.") + } + absolute, err := imagePickerAbsoluteFolder(path) + if err != nil { + return "", err + } + info, err := os.Stat(absolute) + if err != nil { + return "", fmt.Errorf("Choose an existing folder: %w", err) + } + if !info.IsDir() { + return "", errors.New("Choose a folder, not a file.") + } + return absolute, nil +} + +func imagePickerCompleteFolders(ctx context.Context, draft string) ([]string, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + path, err := imagePickerAbsoluteFolder(draft) + if err != nil { + return nil, err + } + parent, prefix := filepath.Dir(path), filepath.Base(path) + if draft == "" || draft == "~" || os.IsPathSeparator(draft[len(draft)-1]) { + parent, prefix = path, "" + } + // Inspect the opened handle: a path can be replaced by a FIFO between a + // separate Stat and Open. Nonblocking open keeps that replacement from hanging. + directory, err := os.OpenFile(parent, os.O_RDONLY|syscall.O_NONBLOCK, 0) + if err != nil { + return nil, err + } + defer directory.Close() + info, err := directory.Stat() + if err != nil { + return nil, err + } + if !info.IsDir() { + return nil, errors.New("Choose a folder, not a file.") + } + var matches []string + for { + if err := ctx.Err(); err != nil { + return nil, err + } + entries, err := directory.ReadDir(128) + if err != nil && !errors.Is(err, io.EOF) { + return nil, err + } + for _, entry := range entries { + if !strings.HasPrefix(entry.Name(), prefix) { + continue + } + candidate, isDir := filepath.Join(parent, entry.Name()), entry.IsDir() + if entry.Type()&os.ModeSymlink != 0 { + info, err := os.Stat(candidate) + isDir = err == nil && info.IsDir() + } + if isDir { + matches = append(matches, candidate+string(filepath.Separator)) + } + } + if errors.Is(err, io.EOF) { + break + } + } + sort.Strings(matches) + return matches, nil +} diff --git a/pkg/custom/image_picker_folder_test.go b/pkg/custom/image_picker_folder_test.go new file mode 100644 index 00000000..049f828f --- /dev/null +++ b/pkg/custom/image_picker_folder_test.go @@ -0,0 +1,405 @@ +package custom + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "syscall" + "testing" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/x/ansi" + "github.com/stretchr/testify/require" +) + +func pickerOpenFolder(t *testing.T, m *imagePicker) { + t.Helper() + m.focus = "options" + for i, row := range m.rows() { + if row.id == "folder" { + m.selected = i + pickerKey(m, tea.KeyEnter) + require.Equal(t, "folder", m.page) + return + } + } + t.Fatal("Save to row not found") +} + +func pickerOpenPath(t *testing.T, m *imagePicker) { + t.Helper() + pickerOpenFolder(t, m) + m.selected = 2 + pickerKey(m, tea.KeyEnter) + require.Equal(t, "path", m.focus) +} + +func TestImagePickerRestoredDraftAndExplicitPrompt(t *testing.T) { + original := pickerForTest(t).settings + original.prompt, original.quality, original.outputDir = "Remembered 雪", "high", t.TempDir() + for _, tc := range []struct { + prompt string + set bool + want string + }{ + {"", false, original.prompt}, {"", true, ""}, {"Override", false, "Override"}, + } { + m, err := newImagePicker(imagePickerOptions{initial: &original, initialNote: "Remembered settings", Prompt: tc.prompt, PromptSet: tc.set}) + require.NoError(t, err) + require.Equal(t, tc.want, m.settings.prompt) + require.Equal(t, tc.want, string(m.draft)) + require.Equal(t, len([]rune(tc.want)), m.cursor) + require.Equal(t, original.outputDir, m.settings.outputDir) + require.Equal(t, "high", m.settings.quality) + require.Equal(t, "Remembered settings", m.note) + } +} + +func TestImagePickerSaveFolder(t *testing.T) { + m := pickerForTest(t) + for _, size := range [][2]int{{40, 12}, {80, 24}} { + m.width, m.height = size[0], size[1] + // Short terminals scroll settings; focusing the folder keeps it visible. + m.focus, m.selected = "options", 3 + require.Contains(t, ansi.Strip(m.View().Content), "Save to") + require.Contains(t, ansi.Strip(m.View().Content), "Ctrl+C exit") + } + m.width, m.height = 90, 24 + require.Equal(t, "more", m.rows()[len(m.rows())-1].id, "End still opens More options") + pickerOpenFolder(t, m) + require.Equal(t, []string{"Default", "Current folder", "Enter path", "Back to settings"}, []string{m.rows()[0].label, m.rows()[1].label, m.rows()[2].label, m.rows()[3].label}) + m.selected = 1 + pickerKey(m, tea.KeyEnter) + cwd, err := os.Getwd() + require.NoError(t, err) + require.Equal(t, cwd, m.settings.outputDir) + require.Equal(t, "settings", m.page) + require.Contains(t, ansi.Strip(m.View().Content), "Save to") + require.Equal(t, cwd, pickerArg(t, m.settings.args(), "--output-dir")) + pickerOpenFolder(t, m) + pickerKey(m, tea.KeyEnter) // Default. + require.Empty(t, m.settings.outputDir) + require.NotContains(t, m.settings.args(), "--output-dir") +} + +func TestImagePickerFolderEditorPreservesLiteralPath(t *testing.T) { + m := pickerForTest(t) + t.Chdir(t.TempDir()) + literal := " @'$HOME' 雪 folder " + require.NoError(t, os.Mkdir(literal, 0700)) + m.settings.outputDir = t.TempDir() + pickerOpenPath(t, m) + m.Update(tea.KeyPressMsg{Code: 'u', Mod: tea.ModCtrl}) + m.Update(tea.PasteMsg{Content: literal}) + require.Equal(t, literal, string(m.folder.draft)) + require.Empty(t, m.result.Args) + pickerKey(m, tea.KeyEnter) + want, err := filepath.Abs(literal) + require.NoError(t, err) + require.Equal(t, want, m.settings.outputDir) + require.Empty(t, m.result.Args, "accepting a folder never generates") + require.Equal(t, "settings", m.page) + command := formatImagePickerCommandBash(m.settings.args()) + require.Contains(t, command, "--output-dir '") + require.Equal(t, want, pickerArg(t, m.settings.args(), "--output-dir")) +} + +func TestImagePickerFolderEscapeDiscardsDraftAndHighlight(t *testing.T) { + m := pickerForTest(t) + original := t.TempDir() + m.settings.outputDir = original + pickerOpenPath(t, m) + m.Update(tea.PasteMsg{Content: "uncommitted"}) + pickerKey(m, tea.KeyEscape) + require.Equal(t, "folder", m.page) + require.Equal(t, original, m.settings.outputDir) + m.selected = 0 // Highlight Default without accepting it. + pickerKey(m, tea.KeyEscape) + require.Equal(t, "settings", m.page) + require.Equal(t, original, m.settings.outputDir) + require.Empty(t, m.result.Args) +} + +func TestImagePickerFolderCompletionDirectoriesOnly(t *testing.T) { + root := t.TempDir() + for _, name := range []string{"雪 first", "雪 second", "$HOME"} { + require.NoError(t, os.Mkdir(filepath.Join(root, name), 0700)) + } + require.NoError(t, os.WriteFile(filepath.Join(root, "雪 file"), []byte("synthetic"), 0600)) + require.NoError(t, os.Symlink(filepath.Join(root, "雪 first"), filepath.Join(root, "雪 linked"))) + matches, err := imagePickerCompleteFolders(context.Background(), filepath.Join(root, "雪")) + require.NoError(t, err) + require.Equal(t, []string{filepath.Join(root, "雪 first") + string(filepath.Separator), filepath.Join(root, "雪 linked") + string(filepath.Separator), filepath.Join(root, "雪 second") + string(filepath.Separator)}, matches) + matches, err = imagePickerCompleteFolders(context.Background(), filepath.Join(root, "$H")) + require.NoError(t, err) + require.Equal(t, []string{filepath.Join(root, "$HOME") + string(filepath.Separator)}, matches) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err = imagePickerCompleteFolders(ctx, root+string(filepath.Separator)) + require.ErrorIs(t, err, context.Canceled) +} + +func TestImagePickerFolderCompletionRejectsFIFO(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix FIFO semantics") + } + mkfifo, err := exec.LookPath("mkfifo") + require.NoError(t, err) + for _, linked := range []bool{false, true} { + name := "direct" + if linked { + name = "symlink" + } + t.Run(name, func(t *testing.T) { + fifo := filepath.Join(t.TempDir(), "not-a-folder") + require.NoError(t, exec.Command(mkfifo, fifo).Run()) + path := fifo + if linked { + path += "-link" + require.NoError(t, os.Symlink(fifo, path)) + } + done := make(chan error, 1) + go func() { + _, err := imagePickerCompleteFolders(t.Context(), path+string(filepath.Separator)) + done <- err + }() + select { + case err := <-done: + require.Error(t, err, "a FIFO is not a completion directory") + case <-time.After(time.Second): + // Release a blocked reader so a failing regression leaves no worker. + writer, err := os.OpenFile(fifo, os.O_RDWR|syscall.O_NONBLOCK, 0) + require.NoError(t, err) + defer writer.Close() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("folder completion stayed blocked after releasing the FIFO") + } + t.Fatal("folder completion blocked opening a FIFO") + } + }) + } +} + +func TestImagePickerFolderCompletionNeedsAcceptance(t *testing.T) { + m := pickerForTest(t) + root := t.TempDir() + for _, name := range []string{"alpha", "alpine"} { + require.NoError(t, os.Mkdir(filepath.Join(root, name), 0700)) + } + pickerOpenPath(t, m) + prefix := filepath.Join(root, "al") + m.insertFolderPath(prefix) + _, cmd := m.Update(tea.KeyPressMsg{Code: tea.KeyTab}) + require.NotNil(t, cmd) + m.Update(cmd()) + require.Len(t, m.folder.matches, 2) + require.Equal(t, prefix, string(m.folder.draft)) + pickerKey(m, tea.KeyTab) + require.Zero(t, m.folder.match) + require.Equal(t, prefix, string(m.folder.draft), "highlight is not a draft edit") + pickerKey(m, tea.KeyEnter) + require.Equal(t, filepath.Join(root, "alpha")+string(filepath.Separator), string(m.folder.draft)) + require.Equal(t, "path", m.page) + require.Empty(t, m.settings.outputDir) + pickerKey(m, tea.KeyEnter) + require.Equal(t, filepath.Join(root, "alpha"), m.settings.outputDir) + require.Empty(t, m.result.Args) +} + +func TestImagePickerFolderCompletionIgnoresLateReplies(t *testing.T) { + m := pickerForTest(t) + root := t.TempDir() + require.NoError(t, os.Mkdir(filepath.Join(root, "only"), 0700)) + pickerOpenPath(t, m) + m.insertFolderPath(root + string(filepath.Separator)) + _, cmd := m.Update(tea.KeyPressMsg{Code: tea.KeyTab}) + require.NotNil(t, cmd) + reply := cmd() + m.insertFolderPath("new draft") + want := string(m.folder.draft) + m.Update(reply) + require.Equal(t, want, string(m.folder.draft)) + require.Empty(t, m.folder.matches) + m.openFolderEditor(root + string(filepath.Separator)) + _, cmd = m.Update(tea.KeyPressMsg{Code: tea.KeyTab}) + m.Update(cmd()) + require.Equal(t, filepath.Join(root, "only")+string(filepath.Separator), string(m.folder.draft)) + require.Empty(t, m.settings.outputDir, "unique completion edits draft only") +} + +func TestImagePickerInvalidFolderRetainsEditorAndEscapesErrors(t *testing.T) { + for _, invalid := range []string{"", "missing\x1b]52;c;injection\a\nnext", "nul\x00path", "not-a-folder"} { + m := pickerForTest(t) + t.Chdir(t.TempDir()) + require.NoError(t, os.WriteFile("not-a-folder", []byte("synthetic"), 0600)) + original := t.TempDir() + m.settings.outputDir = original + pickerOpenPath(t, m) + m.openFolderEditor(invalid) + pickerKey(m, tea.KeyEnter) + require.Equal(t, "path", m.page) + require.Equal(t, invalid, string(m.folder.draft)) + require.Equal(t, original, m.settings.outputDir) + require.NotEmpty(t, m.note) + require.Empty(t, m.result.Args) + view := m.View().Content + require.NotContains(t, view, "\x1b]") + require.NotContains(t, view, "\a") + } +} + +func TestImagePickerFolderValidationBeforeSubmit(t *testing.T) { + for _, printOnly := range []bool{false, true} { + m := pickerForTest(t) + m.settings.outputDir = t.TempDir() + want := m.settings + cmd := m.submitWithFolder(printOnly) + require.NotNil(t, cmd) + require.Empty(t, m.result.Args) + _, quit := m.Update(cmd()) + require.NotNil(t, quit) + require.IsType(t, tea.QuitMsg{}, quit()) + require.Equal(t, want, m.result.settings) + require.Equal(t, want.args(), m.result.Args) + require.Equal(t, printOnly, m.result.PrintOnly) + entries, err := os.ReadDir(want.outputDir) + require.NoError(t, err) + require.Empty(t, entries, "write probe must be removed") + } +} + +func TestImagePickerDefaultFolderPrintHasNoSideEffects(t *testing.T) { + m := pickerForTest(t) + path := filepath.Join(os.Getenv("HOME"), "Downloads", "gpt-images") + _, err := os.Stat(path) + require.ErrorIs(t, err, os.ErrNotExist) + _, cmd := m.Update(tea.KeyPressMsg{Code: 'p', Mod: tea.ModCtrl}) + require.NotNil(t, cmd) + require.True(t, m.result.PrintOnly) + require.Empty(t, m.result.settings.outputDir) + _, err = os.Stat(path) + require.ErrorIs(t, err, os.ErrNotExist) +} + +func TestImagePickerDefaultFolderFailureCanBeCorrected(t *testing.T) { + m := pickerForTest(t) + require.NoError(t, os.WriteFile(filepath.Join(os.Getenv("HOME"), "Downloads"), []byte("synthetic"), 0600)) + _, cmd := m.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + require.NotNil(t, cmd) + m.Update(cmd()) + require.Empty(t, m.result.Args) + require.Equal(t, "path", m.page) + require.Contains(t, m.note, "default image output") + require.Empty(t, m.settings.outputDir) + valid := t.TempDir() + m.openFolderEditor(valid) + pickerKey(m, tea.KeyEnter) + require.Equal(t, valid, m.settings.outputDir) + _, cmd = m.Update(tea.KeyPressMsg{Code: 'g', Mod: tea.ModCtrl}) + cmd = pickerFinishFolderCommand(m, cmd) + require.NotNil(t, cmd) + require.Equal(t, valid, pickerArg(t, m.result.Args, "--output-dir")) +} + +func TestImagePickerPendingFolderSubmissionCancellation(t *testing.T) { + for _, cancel := range []tea.Msg{tea.KeyPressMsg{Code: tea.KeyEscape}, tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl}, imagePickerStopMsg{code: 143}} { + m := pickerForTest(t) + m.settings.outputDir = t.TempDir() + original := m.settings + _, cmd := m.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + require.NotNil(t, cmd) + for _, queued := range []tea.Msg{tea.KeyPressMsg{Code: tea.KeyEnter}, tea.PasteMsg{Content: "do not apply"}, tea.KeyPressMsg{Code: 'x', Text: "x"}} { + _, duplicate := m.Update(queued) + require.Nil(t, duplicate) + } + require.Equal(t, original, m.settings) + reply := cmd() + m.Update(cancel) + _, quit := m.Update(reply) + require.Nil(t, quit) + require.Empty(t, m.result.Args) + require.Equal(t, original, m.settings) + } +} + +func TestImagePickerFolderSubmissionRefusesSmallResizeOrChangedState(t *testing.T) { + for _, change := range []string{"resize", "settings"} { + m := pickerForTest(t) + m.settings.outputDir = t.TempDir() + _, cmd := m.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + reply := cmd() + if change == "resize" { + m.Update(tea.WindowSizeMsg{Width: 25, Height: 7}) + } else { + m.settings.prompt = "new state" + } + _, quit := m.Update(reply) + require.Nil(t, quit) + require.Empty(t, m.result.Args) + } +} + +func TestImagePickerFolderViewBoundsAndVisibleExit(t *testing.T) { + m := pickerForTest(t) + pickerOpenPath(t, m) + m.insertFolderPath(strings.Repeat("雪 folder /", 100)) + m.folder.matches = []string{"first", "second", "third", "fourth"} + m.folder.match, m.selected = 3, 3 + m.note = "Folder error\x1b[31m\nnext" + for _, size := range [][2]int{{40, 12}, {48, 14}, {80, 24}, {120, 30}} { + m.width, m.height = size[0], size[1] + view := m.View().Content + require.Contains(t, view, "Ctrl+C exit") + require.Contains(t, ansi.Strip(view), "› fourth") + require.NotContains(t, view, "\x1b[31m") + require.LessOrEqual(t, len(strings.Split(view, "\n")), m.viewHeight()) + for _, line := range strings.Split(view, "\n") { + require.LessOrEqual(t, ansi.StringWidth(line), size[0]) + } + } +} + +func TestImagePickerReadOnlyFolderKeepsPickerOpen(t *testing.T) { + if runtime.GOOS == "windows" || os.Geteuid() == 0 { + t.Skip("requires Unix permission checks for a non-root user") + } + m := pickerForTest(t) + directory := t.TempDir() + require.NoError(t, os.Chmod(directory, 0500)) + t.Cleanup(func() { require.NoError(t, os.Chmod(directory, 0700)) }) + m.settings.outputDir = directory + _, cmd := m.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + require.NotNil(t, cmd) + m.Update(cmd()) + require.Empty(t, m.result.Args) + require.Equal(t, "path", m.page) + require.Equal(t, directory, string(m.folder.draft)) + require.Contains(t, m.note, "not writable") +} + +func TestImagePickerFolderLabelsAbbreviateOnlyHomeChildren(t *testing.T) { + m := pickerForTest(t) + home, err := os.UserHomeDir() + require.NoError(t, err) + for _, tc := range []struct{ path, want string }{ + {home, "~"}, + {filepath.Join(home, "Pictures", "雪 folder"), "~/Pictures/雪 folder"}, + {home + "-neighbor", home + "-neighbor"}, + {filepath.Dir(home), filepath.Dir(home)}, + {"relative", "relative"}, + } { + require.Equal(t, tc.want, imagePickerFolderLabel(tc.path)) + } + path := filepath.Join(home, "Pictures", "雪 folder") + m.settings.outputDir = path + require.Equal(t, "~/Pictures/雪 folder", m.folderRow().value) + require.Contains(t, ansi.Strip(m.View().Content), "~/Pictures/雪 folder") + require.Equal(t, path, m.settings.outputDir) + require.Equal(t, path, pickerArg(t, m.settings.args(), "--output-dir")) +} diff --git a/pkg/custom/image_picker_session.go b/pkg/custom/image_picker_session.go index 8900daab..938cb3cd 100644 --- a/pkg/custom/image_picker_session.go +++ b/pkg/custom/image_picker_session.go @@ -5,27 +5,56 @@ import ( "fmt" "io" "os" + "strings" + + "github.com/charmbracelet/x/term" ) -// Keep choices within the current invocation. Direct API commands never enter -// this session or acquire picker state. +// Persistence belongs to the interactive session, never direct API commands. func runImagePickerSession(ctx context.Context, input, output *os.File, options imagePickerOptions) (imagePickerResult, error) { if options.Shell == "" { options.Shell = os.Getenv("OPENAI_PICKER_SHELL") } if options.Shell == "" { + // Enter launches have no hook-provided shell. Match the immediate + // caller so the displayed command can be pasted back into that shell. options.Shell = imagePickerParentShell(ctx) } + if ctx.Err() != nil || input == nil || output == nil || !term.IsTerminal(input.Fd()) || !term.IsTerminal(output.Fd()) || strings.EqualFold(os.Getenv("TERM"), "dumb") { + return runImagePicker(ctx, input, output, options) + } + path, stateErr := imagePickerStatePath() + if stateErr == nil { + var settings imagePickerSettings + var found bool + settings, found, stateErr = loadImagePickerState(ctx, path) + if stateErr == nil && found && options.initial == nil { + options.initial = &settings + } + } + if err := ctx.Err(); err != nil { + return imagePickerResult{}, err + } + if stateErr != nil { + options.initialNote = "Could not restore saved settings. Using defaults; your saved settings are kept." + if options.initial != nil { + options.initialNote = "Could not read saved settings. Keeping this session's choices." + } + } result, err := runImagePicker(ctx, input, output, options) if err != nil || result.Canceled { return result, err } - return result, finishImagePickerSelection(ctx, output, result) + // Unknown, corrupt, or inaccessible state remains untouched this session. + if stateErr != nil { + path = "" + } + return result, finishImagePickerSelection(ctx, output, os.Stderr, path, result) } -// Echo the selected command only after restoring the terminal and before -// starting the request. A failed echo must not silently start generation. -func finishImagePickerSelection(ctx context.Context, output io.Writer, result imagePickerResult) error { +// Echo before saving history or starting a request. Recording at submission +// means a slower API response cannot overwrite a newer interactive selection. +func finishImagePickerSelection(ctx context.Context, output, diagnostics io.Writer, path string, result imagePickerResult) error { if err := ctx.Err(); err != nil { return err } @@ -42,5 +71,18 @@ func finishImagePickerSelection(ctx context.Context, output io.Writer, result im if _, err := fmt.Fprintln(output, command); err != nil { return imageSavingFailure("Could not print the command. No image request was started.", err) } + if path == "" { + return nil + } + if err := saveImagePickerState(ctx, path, result.settings); err != nil { + if canceled := ctx.Err(); canceled != nil { + return canceled + } + // Storage is optional. Do not expose stored prompts or raw filesystem + // errors, and do not prevent an otherwise valid image request. + if _, err := fmt.Fprintln(diagnostics, "Could not remember these settings."); err != nil { + return imageSavingFailure("Could not print the settings warning. No image request was started.", err) + } + } return nil } diff --git a/pkg/custom/image_picker_session_test.go b/pkg/custom/image_picker_session_test.go index 85d51809..3204ade0 100644 --- a/pkg/custom/image_picker_session_test.go +++ b/pkg/custom/image_picker_session_test.go @@ -5,29 +5,102 @@ import ( "context" "errors" "io" + "os" + "path/filepath" "testing" "github.com/stretchr/testify/require" + "github.com/urfave/cli/v3" ) +func TestImagePickerFolderFlagRestoresParsedState(t *testing.T) { + for _, explicit := range []bool{false, true} { + t.Run(map[bool]string{false: "unset", true: "explicit"}[explicit], func(t *testing.T) { + destination := "" + flag := &cli.StringFlag{Name: "output-dir", Value: "default", Destination: &destination} + app := &cli.Command{Flags: []cli.Flag{flag}} + app.Action = func(ctx context.Context, c *cli.Command) error { + before := c.String("output-dir") + beforeDestination := destination + for range 2 { + restore, err := setImagePickerFlag(c, "output-dir", "/tmp/絵 with spaces") + require.NoError(t, err) + require.Equal(t, "/tmp/絵 with spaces", c.String("output-dir")) + require.True(t, c.IsSet("output-dir")) + require.Equal(t, beforeDestination, destination) + restore() + require.Equal(t, before, c.String("output-dir")) + require.Equal(t, explicit, c.IsSet("output-dir")) + require.Equal(t, beforeDestination, destination) + } + return nil + } + args := []string{"openai"} + if explicit { + args = append(args, "--output-dir", "original") + } + require.NoError(t, app.Run(context.Background(), args)) + }) + } +} + type failedImagePickerWriter struct{} func (failedImagePickerWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe } -func TestImagePickerSelectionEchoFailurePreventsGeneration(t *testing.T) { - m, err := newImagePicker(imagePickerOptions{Prompt: "Synthetic prompt"}) +func TestImagePickerSelectionHistoryFollowsSuccessfulEcho(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "image-picker.json") + model, err := newImagePicker(imagePickerOptions{Prompt: "previous"}) require.NoError(t, err) - result := imagePickerResult{Args: m.settings.args(), settings: m.settings, shell: "bash"} - err = finishImagePickerSelection(context.Background(), failedImagePickerWriter{}, result) + require.NoError(t, saveImagePickerState(ctx, path, model.settings)) + before, err := os.ReadFile(path) + require.NoError(t, err) + model.settings.prompt = "new prompt\n'絵'" + result := imagePickerResult{Args: model.settings.args(), settings: model.settings, PrintOnly: true, shell: "bash"} + var diagnostic bytes.Buffer + err = finishImagePickerSelection(ctx, failedImagePickerWriter{}, &diagnostic, path, result) require.ErrorIs(t, err, io.ErrClosedPipe) + after, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, before, after) + require.Empty(t, diagnostic.String()) var output bytes.Buffer - require.NoError(t, finishImagePickerSelection(context.Background(), &output, result)) + require.NoError(t, finishImagePickerSelection(ctx, &output, &diagnostic, path, result)) require.Equal(t, formatImagePickerCommand(result.Args, "bash")+"\n", output.String()) + got, found, err := loadImagePickerState(ctx, path) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, model.settings, got) +} + +func TestImagePickerOptionalHistoryFailureDoesNotBlockSelection(t *testing.T) { + path := filepath.Join(t.TempDir(), "image-picker.json") + invalid := []byte(`{"version":99,"prompt":"private text"}`) + require.NoError(t, os.WriteFile(path, invalid, 0600)) + model, err := newImagePicker(imagePickerOptions{Prompt: "new"}) + require.NoError(t, err) + result := imagePickerResult{Args: model.settings.args(), settings: model.settings, shell: "bash"} + var output, diagnostics bytes.Buffer + require.NoError(t, finishImagePickerSelection(context.Background(), &output, &diagnostics, path, result)) + require.Equal(t, "Could not remember these settings.\n", diagnostics.String()) + after, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, invalid, after) + err = finishImagePickerSelection(context.Background(), &output, failedImagePickerWriter{}, path, result) + require.ErrorIs(t, err, io.ErrClosedPipe) +} + +func TestImagePickerCanceledSelectionHasNoPersistenceOrOutput(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - output.Reset() - require.True(t, errors.Is(finishImagePickerSelection(ctx, &output, result), context.Canceled)) + path := filepath.Join(t.TempDir(), "image-picker.json") + var output bytes.Buffer + err := finishImagePickerSelection(ctx, &output, &output, path, imagePickerResult{}) + require.ErrorIs(t, err, context.Canceled) require.Empty(t, output.String()) + _, err = os.Stat(path) + require.True(t, errors.Is(err, os.ErrNotExist)) } func TestImagePickerUnsupportedShellSelection(t *testing.T) { @@ -36,14 +109,14 @@ func TestImagePickerUnsupportedShellSelection(t *testing.T) { m := pickerForTest(t) result := imagePickerResult{Args: m.settings.args(), settings: m.settings, shell: shell} var output bytes.Buffer - require.NoError(t, finishImagePickerSelection(context.Background(), &output, result)) + require.NoError(t, finishImagePickerSelection(context.Background(), &output, &output, "", result)) require.Empty(t, output.String(), "generation must not echo a command for the wrong shell") result.PrintOnly = true - require.ErrorContains(t, finishImagePickerSelection(context.Background(), &output, result), imagePickerUnsupportedShell) + require.ErrorContains(t, finishImagePickerSelection(context.Background(), &output, &output, "", result), imagePickerUnsupportedShell) require.Empty(t, output.String()) ctx, cancel := context.WithCancel(context.Background()) cancel() - require.ErrorIs(t, finishImagePickerSelection(ctx, &output, result), context.Canceled) + require.ErrorIs(t, finishImagePickerSelection(ctx, &output, &output, "", result), context.Canceled) }) } } diff --git a/pkg/custom/image_picker_state.go b/pkg/custom/image_picker_state.go new file mode 100644 index 00000000..80e8622d --- /dev/null +++ b/pkg/custom/image_picker_state.go @@ -0,0 +1,283 @@ +package custom + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/json" + "errors" + "io" + "os" + "path/filepath" + "runtime" + "strings" + "syscall" + "unicode/utf8" +) + +// This bound applies only to optional local persistence, never image requests. +const imagePickerStateLimit = 1 << 20 + +var ( + errImagePickerStateInvalid = errors.New("saved image picker settings are unknown or invalid") + errImagePickerStateChanged = errors.New("saved image picker settings changed while opening") +) + +// Keep this record separate from preview preferences and root configuration. +// In particular, it cannot contain connection settings or credentials. +type imagePickerState struct { + Version int `json:"version"` + Prompt string `json:"prompt"` + Model string `json:"model"` + Size string `json:"size"` + Quality string `json:"quality"` + Background string `json:"background"` + Format string `json:"format"` + Count string `json:"count"` + OutputDir string `json:"output_dir"` +} + +func imagePickerStatePath() (string, error) { + directory, err := os.UserConfigDir() + if err != nil { + return "", err + } + return filepath.Join(directory, "openai", "image-picker.json"), nil +} + +func loadImagePickerState(ctx context.Context, path string) (imagePickerSettings, bool, error) { + if err := ctx.Err(); err != nil { + return imagePickerSettings{}, false, err + } + root, name, err := openImagePickerStateParent(path, false) + if errors.Is(err, os.ErrNotExist) { + return imagePickerSettings{}, false, nil + } + if err != nil { + return imagePickerSettings{}, false, err + } + defer root.Close() + return readImagePickerState(ctx, root, name) +} + +func readImagePickerState(ctx context.Context, root *os.Root, name string) (imagePickerSettings, bool, error) { + for range 64 { + if err := ctx.Err(); err != nil { + return imagePickerSettings{}, false, err + } + settings, found, err := readImagePickerStateSnapshot(ctx, root, name) + if !errors.Is(err, errImagePickerStateChanged) { + return settings, found, err + } + } + return imagePickerSettings{}, false, errImagePickerStateChanged +} + +func readImagePickerStateSnapshot(ctx context.Context, root *os.Root, name string) (imagePickerSettings, bool, error) { + info, err := root.Lstat(name) + if errors.Is(err, os.ErrNotExist) { + return imagePickerSettings{}, false, nil + } + if err != nil { + return imagePickerSettings{}, false, err + } + if !privateImagePickerState(info) { + return imagePickerSettings{}, false, errors.New("saved image picker settings must be a private regular file") + } + // Nonblocking open prevents a replaced FIFO from waiting for a writer. + file, err := root.OpenFile(name, os.O_RDONLY|syscall.O_NONBLOCK, 0) + if err != nil { + return imagePickerSettings{}, false, err + } + defer file.Close() + actual, err := file.Stat() + if err != nil { + return imagePickerSettings{}, false, err + } + if !privateImagePickerState(actual) || !os.SameFile(info, actual) { + return imagePickerSettings{}, false, errImagePickerStateChanged + } + current, err := root.Lstat(name) + if err != nil || !privateImagePickerState(current) || !os.SameFile(actual, current) { + return imagePickerSettings{}, false, errImagePickerStateChanged + } + if actual.Size() > imagePickerStateLimit { + return imagePickerSettings{}, false, errImagePickerStateInvalid + } + data, err := io.ReadAll(io.LimitReader(imagePickerStateReader{ctx, file}, imagePickerStateLimit+1)) + if err != nil { + return imagePickerSettings{}, false, err + } + if err := ctx.Err(); err != nil { + return imagePickerSettings{}, false, err + } + // Validate and decode the same bounded snapshot, including in-place rewrites. + settings, valid := decodeImagePickerState(data) + if err := ctx.Err(); err != nil { + return imagePickerSettings{}, false, err + } + if !valid || len(data) > imagePickerStateLimit { + return imagePickerSettings{}, false, errImagePickerStateInvalid + } + return settings, true, nil +} + +// Save replaces a whole record. Concurrent picker saves may win in either order; +// readers see only complete records, and there is no field merging. Files from +// unknown schemas or corrupt existing records are never intentionally replaced. +func saveImagePickerState(ctx context.Context, path string, settings imagePickerSettings) error { + if err := ctx.Err(); err != nil { + return err + } + if !validImagePickerState(settings) { + return errImagePickerStateInvalid + } + state := imagePickerState{1, settings.prompt, settings.model, settings.size, settings.quality, settings.background, settings.format, settings.count, settings.outputDir} + data, err := json.Marshal(state) + if err != nil { + return err + } + if len(data)+1 > imagePickerStateLimit { + return errors.New("image picker settings exceed the local persistence limit") + } + root, name, err := openImagePickerStateParent(path, true) + if err != nil { + return err + } + defer root.Close() + if _, _, err := readImagePickerState(ctx, root, name); err != nil { + return err + } + temporary := ".image-picker-" + rand.Text() + file, err := root.OpenFile(temporary, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600) + if err != nil { + return err + } + defer root.Remove(temporary) + _, writeErr := io.Copy(file, imagePickerStateReader{ctx, bytes.NewReader(append(data, '\n'))}) + if writeErr == nil { + writeErr = file.Sync() + } + closeErr := file.Close() + if err := errors.Join(writeErr, closeErr, ctx.Err()); err != nil { + return err + } + // Recheck after writing so a newly invalid record also remains untouched. + if _, _, err := readImagePickerState(ctx, root, name); err != nil { + return err + } + if err := ctx.Err(); err != nil { + return err + } + return root.Rename(temporary, name) +} + +func openImagePickerStateParent(path string, create bool) (*os.Root, string, error) { + directory, name := filepath.Dir(path), filepath.Base(path) + if create { + if err := os.MkdirAll(directory, 0700); err != nil { + return nil, "", err + } + } + info, err := os.Lstat(directory) + if err != nil { + return nil, "", err + } + if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 || (runtime.GOOS != "windows" && info.Mode().Perm()&0022 != 0) { + return nil, "", errors.New("image picker settings require a configuration directory writable only by its owner") + } + root, err := os.OpenRoot(directory) + if err != nil { + return nil, "", err + } + actual, err := root.Stat(".") + if err != nil || !os.SameFile(info, actual) { + root.Close() + if err == nil { + err = errImagePickerStateChanged + } + return nil, "", err + } + return root, name, nil +} + +func privateImagePickerState(info os.FileInfo) bool { + return info.Mode().IsRegular() && (runtime.GOOS == "windows" || info.Mode().Perm()&0077 == 0) +} + +func decodeImagePickerState(data []byte) (imagePickerSettings, bool) { + s := imagePickerSettings{} + if !utf8.Valid(data) { + return s, false + } + fields := map[string]*string{"prompt": &s.prompt, "model": &s.model, "size": &s.size, "quality": &s.quality, "background": &s.background, "format": &s.format, "count": &s.count, "output_dir": &s.outputDir} + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + token, err := decoder.Token() + if err != nil || token != json.Delim('{') { + return s, false + } + seen := map[string]bool{} + for decoder.More() { + keyToken, err := decoder.Token() + key, ok := keyToken.(string) + if err != nil || !ok || seen[key] { + return s, false + } + seen[key] = true + value, err := decoder.Token() + if err != nil { + return s, false + } + if key == "version" { + if value != json.Number("1") { + return s, false + } + } else { + text, ok := value.(string) + field, known := fields[key] + if !ok || !known { + return s, false + } + *field = text + } + } + token, err = decoder.Token() + return s, err == nil && token == json.Delim('}') && len(seen) == len(fields)+1 && decoder.Decode(new(any)) == io.EOF && validImagePickerState(s) +} + +func validImagePickerState(settings imagePickerSettings) bool { + if settings.outputDir != "" && (!filepath.IsAbs(settings.outputDir) || strings.ContainsRune(settings.outputDir, 0)) { + return false + } + total := 0 + for _, value := range []string{settings.prompt, settings.model, settings.size, settings.quality, settings.background, settings.format, settings.count, settings.outputDir} { + if !utf8.ValidString(value) || len(value) > imagePickerStateLimit-total { + return false + } + total += len(value) + } + m := imagePicker{settings: settings} + for _, field := range []string{"model", "size", "quality", "background", "format", "count"} { + found := false + for _, choice := range m.choices(field) { + found = found || choice.value == m.value(field) + } + if !found { + return false + } + } + return settings.background != "transparent" || settings.format != "jpeg" +} + +type imagePickerStateReader struct { + ctx context.Context + reader io.Reader +} + +func (r imagePickerStateReader) Read(data []byte) (int, error) { + if err := r.ctx.Err(); err != nil { + return 0, err + } + return r.reader.Read(data[:min(len(data), 32<<10)]) +} diff --git a/pkg/custom/image_picker_state_test.go b/pkg/custom/image_picker_state_test.go new file mode 100644 index 00000000..67967894 --- /dev/null +++ b/pkg/custom/image_picker_state_test.go @@ -0,0 +1,356 @@ +package custom + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func pickerStateSettings(t *testing.T) imagePickerSettings { + t.Helper() + m, err := newImagePicker(imagePickerOptions{}) + require.NoError(t, err) + m.settings.prompt = " Synthetic moon over 日本語\nsecond line 🌕\t'quoted' " + m.settings.outputDir = filepath.Join(t.TempDir(), "pictures with spaces") + return m.settings +} + +func pickerStatePath(t *testing.T) string { + t.Helper() + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("XDG_CONFIG_HOME", home) + t.Setenv("APPDATA", home) + path, err := imagePickerStatePath() + require.NoError(t, err) + require.True(t, strings.HasPrefix(path, home+string(filepath.Separator))) + require.Equal(t, "image-picker.json", filepath.Base(path)) + return path +} + +func TestImagePickerStateRoundTripAndPrivateReplacement(t *testing.T) { + path := pickerStatePath(t) + _, found, err := loadImagePickerState(t.Context(), path) + require.NoError(t, err) + require.False(t, found) + _, err = os.Stat(filepath.Dir(path)) + require.ErrorIs(t, err, os.ErrNotExist, "opening the picker must not create state") + want := pickerStateSettings(t) + for _, folder := range []string{want.outputDir, "", filepath.Join(t.TempDir(), "再利用 folder")} { + want.outputDir = folder + require.NoError(t, saveImagePickerState(t.Context(), path, want)) + got, found, err := loadImagePickerState(t.Context(), path) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, want, got) + entries, err := os.ReadDir(filepath.Dir(path)) + require.NoError(t, err) + require.Len(t, entries, 1, "completed writes must not retain temporary state") + } + if runtime.GOOS != "windows" { + info, err := os.Stat(path) + require.NoError(t, err) + require.Equal(t, os.FileMode(0600), info.Mode().Perm()) + info, err = os.Stat(filepath.Dir(path)) + require.NoError(t, err) + require.Equal(t, os.FileMode(0700), info.Mode().Perm()) + } + data, err := os.ReadFile(path) + require.NoError(t, err) + var fields map[string]any + require.NoError(t, json.Unmarshal(data, &fields)) + require.ElementsMatch(t, []string{"version", "prompt", "model", "size", "quality", "background", "format", "count", "output_dir"}, pickerStateKeys(fields)) +} + +func pickerStateKeys(fields map[string]any) []string { + keys := make([]string, 0, len(fields)) + for key := range fields { + keys = append(keys, key) + } + return keys +} + +func TestImagePickerStateMalformedAndUnknownRemainUntouched(t *testing.T) { + path := pickerStatePath(t) + settings := pickerStateSettings(t) + require.NoError(t, saveImagePickerState(t.Context(), path, settings)) + valid, err := os.ReadFile(path) + require.NoError(t, err) + for name, data := range map[string][]byte{ + "empty": {}, + "malformed": []byte("synthetic-private-prompt"), + "missing fields": []byte(`{"version":1}`), + "unknown version": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":2`), 1), + "duplicate": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":1,"version":1`), 1), + "unknown field": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":1,"future":true`), 1), + "trailing object": append(append([]byte{}, valid...), []byte(`{}`)...), + "wrong type": bytes.Replace(valid, []byte(`"count":"1"`), []byte(`"count":1`), 1), + "null": bytes.Replace(valid, []byte(`"count":"1"`), []byte(`"count":null`), 1), + "unknown choice": bytes.Replace(valid, []byte(`"count":"1"`), []byte(`"count":"11"`), 1), + "invalid UTF8": bytes.Replace(valid, []byte(`"count":"1"`), []byte{'"', 'c', 'o', 'u', 'n', 't', '"', ':', '"', 0xff, '"'}, 1), + "oversize": bytes.Repeat([]byte(" "), imagePickerStateLimit+1), + } { + t.Run(name, func(t *testing.T) { + require.NoError(t, os.WriteFile(path, data, 0600)) + _, _, err := loadImagePickerState(t.Context(), path) + require.Error(t, err) + require.NotContains(t, err.Error(), "synthetic-private-prompt") + require.Error(t, saveImagePickerState(t.Context(), path, settings)) + got, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, data, got) + }) + } +} + +func TestImagePickerStateRejectsSpecialFilesAndSymlinks(t *testing.T) { + settings := pickerStateSettings(t) + for _, kind := range []string{"directory", "symlink", "fifo"} { + t.Run(kind, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "image-picker.json") + switch kind { + case "directory": + require.NoError(t, os.Mkdir(path, 0700)) + case "symlink": + if runtime.GOOS == "windows" { + t.Skip("native Windows symlinks need additional privileges") + } + target := filepath.Join(t.TempDir(), "target") + require.NoError(t, os.WriteFile(target, []byte("unchanged"), 0600)) + require.NoError(t, os.Symlink(target, path)) + case "fifo": + if runtime.GOOS == "windows" { + t.Skip("Unix FIFO semantics") + } + mkfifo, err := exec.LookPath("mkfifo") + require.NoError(t, err) + require.NoError(t, exec.Command(mkfifo, path).Run()) + } + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + finished := make(chan error, 1) + go func() { + _, _, readErr := loadImagePickerState(ctx, path) + writeErr := saveImagePickerState(ctx, path, settings) + if readErr == nil || writeErr == nil { + finished <- errors.New("special file accepted") + } else { + finished <- nil + } + }() + select { + case err := <-finished: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("special file blocked state access") + } + }) + } + if runtime.GOOS != "windows" { + parent := filepath.Join(t.TempDir(), "linked-config") + require.NoError(t, os.Symlink(t.TempDir(), parent)) + path := filepath.Join(parent, "image-picker.json") + _, _, err := loadImagePickerState(t.Context(), path) + require.Error(t, err) + require.Error(t, saveImagePickerState(t.Context(), path, settings)) + } +} + +func TestImagePickerStateRejectsExposedPermissions(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Windows ACLs differ from Unix mode bits") + } + path := pickerStatePath(t) + settings := pickerStateSettings(t) + require.NoError(t, saveImagePickerState(t.Context(), path, settings)) + before, err := os.ReadFile(path) + require.NoError(t, err) + for _, mode := range []os.FileMode{0644, 0640, 0604, 0620, 0610} { + require.NoError(t, os.Chmod(path, mode)) + _, _, err := loadImagePickerState(t.Context(), path) + require.Error(t, err) + require.Error(t, saveImagePickerState(t.Context(), path, settings)) + } + require.NoError(t, os.Chmod(path, 0600)) + require.NoError(t, os.Chmod(filepath.Dir(path), 0770)) + _, _, err = loadImagePickerState(t.Context(), path) + require.Error(t, err) + require.Error(t, saveImagePickerState(t.Context(), path, settings)) + require.NoError(t, os.Chmod(filepath.Dir(path), 0700)) + after, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, before, after) +} + +func TestImagePickerStateWriteFailureKeepsPrevious(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Windows ACLs differ from Unix mode bits") + } + if os.Geteuid() == 0 { + t.Skip("root bypasses the write permission used by this test") + } + path := pickerStatePath(t) + settings := pickerStateSettings(t) + require.NoError(t, saveImagePickerState(t.Context(), path, settings)) + settings.prompt = "a changed synthetic prompt" + require.NoError(t, os.Chmod(filepath.Dir(path), 0500)) + defer os.Chmod(filepath.Dir(path), 0700) + require.Error(t, saveImagePickerState(t.Context(), path, settings)) + got, found, err := loadImagePickerState(t.Context(), path) + require.NoError(t, err) + require.True(t, found) + require.NotEqual(t, settings.prompt, got.prompt) +} + +type pickerCancelAfterTemporary struct { + context.Context + directory string + cancel context.CancelFunc + written bool +} + +func (c pickerCancelAfterTemporary) Err() error { + entries, _ := os.ReadDir(c.directory) + for _, entry := range entries { + if strings.HasPrefix(entry.Name(), ".image-picker-") { + info, err := entry.Info() + if err == nil && (!c.written || info.Size() > 0) { + c.cancel() + } + } + } + return c.Context.Err() +} + +func TestImagePickerStateCanceledWriteKeepsPreviousAndCleansTemporary(t *testing.T) { + path := pickerStatePath(t) + before := pickerStateSettings(t) + require.NoError(t, saveImagePickerState(t.Context(), path, before)) + after := before + after.prompt = "changed synthetic prompt" + for _, written := range []bool{false, true} { + t.Run(fmt.Sprintf("temporary written=%t", written), func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + ctx = pickerCancelAfterTemporary{ctx, filepath.Dir(path), cancel, written} + require.ErrorIs(t, saveImagePickerState(ctx, path, after), context.Canceled) + got, found, err := loadImagePickerState(t.Context(), path) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, before, got) + entries, err := os.ReadDir(filepath.Dir(path)) + require.NoError(t, err) + require.Len(t, entries, 1) + _, _, err = loadImagePickerState(ctx, path) + require.ErrorIs(t, err, context.Canceled) + }) + } +} + +func TestImagePickerStateConcurrentSavesKeepWholeRecords(t *testing.T) { + path := pickerStatePath(t) + base := pickerStateSettings(t) + require.NoError(t, saveImagePickerState(t.Context(), path, base)) + var workers sync.WaitGroup + errs := make(chan error, 25) + start := make(chan struct{}) + done, readerDone := make(chan struct{}), make(chan struct{}) + go func() { + defer close(readerDone) + <-start + for { + select { + case <-done: + return + default: + } + got, found, err := loadImagePickerState(t.Context(), path) + if err != nil || !found { + errs <- fmt.Errorf("concurrent read failed: found %t, %w", found, err) + return + } + if got != base && got.outputDir != filepath.Join(base.outputDir, strings.TrimPrefix(got.prompt, "synthetic prompt ")) { + errs <- errors.New("concurrent reader saw a mixed record") + return + } + } + }() + for i := range 24 { + workers.Add(1) + go func() { + defer workers.Done() + <-start + settings := base + settings.prompt = fmt.Sprintf("synthetic prompt %d", i) + settings.outputDir = filepath.Join(base.outputDir, fmt.Sprint(i)) + errs <- saveImagePickerState(t.Context(), path, settings) + }() + } + close(start) + workers.Wait() + close(done) + <-readerDone + close(errs) + for err := range errs { + require.NoError(t, err) + } + got, found, err := loadImagePickerState(t.Context(), path) + require.NoError(t, err) + require.True(t, found) + index := strings.TrimPrefix(got.prompt, "synthetic prompt ") + require.Equal(t, filepath.Join(base.outputDir, index), got.outputDir, "a saved record must not mix writers") + entries, err := os.ReadDir(filepath.Dir(path)) + require.NoError(t, err) + require.Len(t, entries, 1) +} + +func TestImagePickerStateLimitOnlyDisablesPersistence(t *testing.T) { + path := pickerStatePath(t) + before := pickerStateSettings(t) + require.NoError(t, saveImagePickerState(t.Context(), path, before)) + after := before + after.prompt = strings.Repeat("p", imagePickerStateLimit+1) + m := imagePicker{settings: after} + require.True(t, m.validPrompt(), "persistence size is not a generation limit") + require.Contains(t, after.args(), after.prompt) + require.Error(t, saveImagePickerState(t.Context(), path, after)) + got, _, err := loadImagePickerState(t.Context(), path) + require.NoError(t, err) + require.Equal(t, before, got) +} + +func TestImagePickerStateRejectsUnstableSaveFolders(t *testing.T) { + path := pickerStatePath(t) + settings := pickerStateSettings(t) + require.NoError(t, saveImagePickerState(t.Context(), path, settings)) + valid, err := os.ReadFile(path) + require.NoError(t, err) + for _, folder := range []string{".", "../pictures", "pictures with spaces", filepath.Join(t.TempDir(), "invalid") + "\x00tail"} { + settings.outputDir = folder + require.Error(t, saveImagePickerState(t.Context(), path, settings)) + unchanged, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, valid, unchanged) + var fields map[string]any + require.NoError(t, json.Unmarshal(valid, &fields)) + fields["output_dir"] = folder + invalid, err := json.Marshal(fields) + require.NoError(t, err) + require.NoError(t, os.WriteFile(path, invalid, 0600)) + _, _, err = loadImagePickerState(t.Context(), path) + require.Error(t, err) + require.NoError(t, os.WriteFile(path, valid, 0600)) + } +} diff --git a/pkg/custom/image_picker_test.go b/pkg/custom/image_picker_test.go index a2dc33e0..33275b7e 100644 --- a/pkg/custom/image_picker_test.go +++ b/pkg/custom/image_picker_test.go @@ -31,6 +31,15 @@ func pickerForTest(t *testing.T) *imagePicker { func pickerKey(m *imagePicker, code rune) tea.Cmd { _, cmd := m.Update(tea.KeyPressMsg{Code: code}) + return pickerFinishFolderCommand(m, cmd) +} + +// Execute the same folder command as the event loop. Race tests use Update +// directly when they need to hold or reorder the asynchronous reply. +func pickerFinishFolderCommand(m *imagePicker, cmd tea.Cmd) tea.Cmd { + if cmd != nil && (m.folder.busy == "submit" || m.folder.busy == "select") { + _, cmd = m.Update(cmd()) + } return cmd } @@ -338,6 +347,7 @@ func TestImagePickerCtrlGFromDropdown(t *testing.T) { require.Equal(t, "choose", m.page) pickerKey(m, tea.KeyDown) _, cmd := m.Update(tea.KeyPressMsg{Code: 'g', Mod: tea.ModCtrl}) + cmd = pickerFinishFolderCommand(m, cmd) require.NotNil(t, cmd) require.Equal(t, openai.ImageModelGPTImage2_5Sunburst, m.settings.model, "unconfirmed highlighted option must not apply") require.Equal(t, openai.ImageModelGPTImage2_5Sunburst, pickerArg(t, m.result.Args, "--model")) @@ -381,6 +391,7 @@ func TestImagePickerTinyWindowCannotSubmit(t *testing.T) { m.Update(tea.WindowSizeMsg{Width: 25, Height: 7}) require.Nil(t, pickerKey(m, tea.KeyEnter)) _, cmd := m.Update(tea.KeyPressMsg{Code: 'g', Mod: tea.ModCtrl}) + cmd = pickerFinishFolderCommand(m, cmd) require.Nil(t, cmd) require.Empty(t, m.result.Args) require.Contains(t, m.View().Content, "Resize") @@ -636,6 +647,7 @@ func TestImagePickerInlineFinishedView(t *testing.T) { require.False(t, view.AltScreen) require.NotEmpty(t, view.Content) _, cmd := m.Update(action) + cmd = pickerFinishFolderCommand(m, cmd) require.NotNil(t, cmd) view = m.View() require.False(t, view.AltScreen) diff --git a/pkg/custom/image_picker_view.go b/pkg/custom/image_picker_view.go index cf040b5a..3ae85134 100644 --- a/pkg/custom/image_picker_view.go +++ b/pkg/custom/image_picker_view.go @@ -39,6 +39,8 @@ func imagePickerLine(value string) string { func (m *imagePicker) heading() string { switch m.page { + case "folder", "path": + return "Save to · existing folder" case "choose": return "Choose " + m.field case "more": @@ -122,6 +124,16 @@ func (m *imagePicker) View() tea.View { lines = append(lines, "") } lines = append(lines, strong.Render(m.heading())) + if m.page == "path" { + before := imagePickerLine(string(m.folder.draft[:m.folder.cursor])) + after := imagePickerLine(string(m.folder.draft[m.folder.cursor:])) + before = ansi.TruncateLeft(before, max(0, ansi.StringWidth(before)-(width-6)), "…") + path := before + "▏" + after + if len(m.folder.draft) == 0 { + path += "Enter a folder path…" + } + lines = append(lines, highlight(path, true)) + } // Prompt, command and help remain visible. The option list scrolls // independently, keeping every selected choice accessible in a short terminal. reserved := len(lines) + 1 + len(commandLines) @@ -149,7 +161,7 @@ func (m *imagePicker) View() tea.View { } else if row.value != "" { text = fmt.Sprintf("%-14s %s", row.label, imagePickerLine(row.value)) } - active := i == m.selected && m.focus == "options" + active := i == m.selected && (m.focus == "options" || m.focus == "path" && m.folder.match >= 0) lines = append(lines, highlight(ansi.Truncate(text, width-2, "…"), active)) } if roomy { @@ -179,7 +191,8 @@ func (m *imagePicker) View() tea.View { if imagePickerShellQuoter(m.shell) != nil && ansi.StringWidth(footer+" · Ctrl+P print") <= width { footer += " · Ctrl+P print" } - + } else if m.focus == "path" { + footer = "Ctrl+C exit · Enter use · Tab complete" } if len(rows) > available && m.focus == "options" { footer = fmt.Sprintf("Ctrl+C exit · ↑↓ %d/%d · Enter select", m.selected+1, len(rows)) diff --git a/scripts/check-image-picker-integration.py b/scripts/check-image-picker-integration.py index ba92fc04..369139c9 100644 --- a/scripts/check-image-picker-integration.py +++ b/scripts/check-image-picker-integration.py @@ -190,8 +190,8 @@ def direct_pipe(name, arguments, data=b'', expected_status=1, expected_body=None assert dict(zip(arguments[2::2], arguments[3::2]))['--prompt'] == prompt assert len(server.requests) == before assert not list(home.glob('Downloads/gpt-images/*')) - assert not (home/'Library'/'Application Support'/'openai'/'image-picker.json').exists() - assert not (home/'.config'/'openai'/'image-picker.json').exists() + states = list(home.rglob('image-picker.json')) + assert len(states) == 1 and json.loads(states[0].read_text())['prompt'] == prompt passed(name) finally: try: @@ -275,7 +275,10 @@ def direct_pipe(name, arguments, data=b'', expected_status=1, expected_body=None assert len(saved) == 2 and all(f.read_bytes() == picker.PNG for f in saved) else: check_saved(home) - assert not list(home.rglob('image-picker.json')), 'picker state was persisted' + states = list(home.rglob('image-picker.json')) + assert len(states) == int(action != 'cancel'), (name, states) + if states: + assert json.loads(states[0].read_text())['prompt'] == prompt generates = action not in {'cancel', 'print'} assert len(server.requests)-before == int(generates)+int(action == 'two-generations'), name if action == 'root-connection-options': diff --git a/scripts/check-image-picker-reuse.py b/scripts/check-image-picker-reuse.py new file mode 100644 index 00000000..746b5f28 --- /dev/null +++ b/scripts/check-image-picker-reuse.py @@ -0,0 +1,299 @@ +#!/usr/bin/env python3 +"""Check remembered picker settings and save folders through the real CLI. + +Run with python3 -I -B scripts/check-image-picker-reuse.py dist/openai OUTPUT. +Uses the shared PTY/loopback fixture, temporary homes, and synthetic inputs only. +""" +import argparse +import contextlib +import importlib.util +import json +import os +import pathlib +import subprocess +import sys +import tempfile +import threading + + +spec = importlib.util.spec_from_file_location( + 'picker_harness', pathlib.Path(__file__).with_name('image_picker_harness.py')) +picker = importlib.util.module_from_spec(spec) +spec.loader.exec_module(picker) + + +def paste(terminal, text): + terminal.send(b'\x1b[200~'+text.encode()+b'\x1b[201~') + + +def open_folder_editor(terminal, from_prompt=True): + mark = len(terminal.raw) + if from_prompt: + terminal.send(picker.DOWN) + # Save to is immediately before More options in the compact settings list. + terminal.send(picker.END+picker.UP+b'\r') + terminal.wait('Current folder', after=mark) + mark = len(terminal.raw) + terminal.send(picker.HOME+picker.DOWN*2+b'\r') + terminal.wait('Tab completes directory names', after=mark) + terminal.send(b'\x15') + + +def select_folder(terminal, folder, from_prompt=True): + open_folder_editor(terminal, from_prompt) + paste(terminal, folder) + mark = len(terminal.raw) + terminal.send(b'\r') + terminal.wait('Settings', after=mark) + + +def state_path(home, env): + directory = home/'Library'/'Application Support' if sys.platform == 'darwin' else pathlib.Path(env['XDG_CONFIG_HOME']) + return directory/'openai'/'image-picker.json' + + +def check_images(directory, expected): + files = list(directory.glob('*.png')) + assert len(files) == expected, (directory, files) + assert all(path.read_bytes() == picker.PNG for path in files), 'saved bytes differ from fixture' + + +def check_body(body, prompt, model): + assert body == {'prompt': prompt, 'model': model, 'size': '1024x1024', + 'quality': 'auto', 'background': 'auto', 'output_format': 'png', 'n': 1}, body + assert 'output_dir' not in body and 'output-dir' not in body, body + + +def wait_resumed(terminal, mark): + terminal.wait('Saved image:', after=mark) + saved = terminal.raw.index(b'Saved image:', mark) + terminal.wait('Images', after=saved) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('binary') + parser.add_argument('output') + args = parser.parse_args() + binary = str(pathlib.Path(args.binary).resolve()) + output = pathlib.Path(args.output).resolve() + output.mkdir(parents=True, exist_ok=True) + results = [] + with tempfile.TemporaryDirectory(prefix='image-picker-reuse-') as temporary: + root = pathlib.Path(temporary).resolve() + server = picker.http.server.ThreadingHTTPServer(('127.0.0.1', 0), picker.Fixture) + server.requests, server.mode = [], 'ok' + server.request_started, server.release = threading.Event(), threading.Event() + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + base_env = {'PATH': '/usr/bin:/bin', 'TERM': 'xterm-256color', 'LANG': 'en_US.UTF-8', + 'OPENAI_API_KEY': 'synthetic-picker-key', 'CI': 'true', + 'OPENAI_BASE_URL': f'http://127.0.0.1:{server.server_port}/v1'} + + def environment(name): + home = root/name + home.mkdir() + env = dict(base_env, HOME=str(home), XDG_CONFIG_HOME=str(home/'config'), + APPDATA=str(home/'config')) + return home, env + + def passed(name): + results.append({'case': name, 'result': 'pass'}) + print('PASS', name, flush=True) + + @contextlib.contextmanager + def terminal_case(name, env, arguments=None, cwd=None, interactive=True): + server.request_started.clear() + previous = pathlib.Path.cwd() + try: + if cwd is not None: + os.chdir(cwd) + terminal = picker.Terminal(binary, arguments or ['images', 'generate'], env, + width=140, height=28) + finally: + # The harness has no cwd option. Only process creation changes + # cwd; cases run serially and the fixture never uses cwd. + os.chdir(previous) + try: + if interactive: + picker.ready(terminal) + terminal.wait('Images') + yield terminal + passed(name) + except BaseException: + results.append({'case': name, 'result': 'fail'}) + raise + finally: + try: + terminal.save(output/name) + finally: + terminal.close() + + try: + home, env = environment('remembered') + first_cwd, next_cwd = root/'first cwd', root/'different cwd' + first_cwd.mkdir() + next_cwd.mkdir() + folder = first_cwd/'pictures with spaces 雪' + folder.mkdir() + prompt, model = "Synthetic reuse 'moon' 雪", 'gpt-image-2.5-flare' + path = state_path(home, env) + before = len(server.requests) + with terminal_case('remember-custom-generation', env, cwd=first_cwd) as terminal: + paste(terminal, prompt) + terminal.send(picker.DOWN+b'\r') + terminal.wait('Choose model') + mark = len(terminal.raw) + terminal.send(picker.DOWN+b'\r') + terminal.wait('Settings', after=mark) + select_folder(terminal, folder.name, from_prompt=False) + mark = len(terminal.raw) + terminal.send(b'\x07') + picker.wait_for_request(terminal, server.request_started) + wait_resumed(terminal, mark) + # Capture submission output before the resumed picker's own + # teardown becomes the final bracketed-paste marker. + flags = picker.printed_flags(terminal) + terminal.send(b'\x03') + terminal.finish(130) + assert len(server.requests) == before+1 + check_body(server.requests[-1], prompt, model) + check_images(folder, 1) + assert flags['--output-dir'] == str(folder), flags + saved = json.loads(path.read_bytes()) + assert saved['prompt'] == prompt and saved['model'] == model, saved + assert saved['output_dir'] == str(folder), saved + assert not (home/'Downloads'/'gpt-images').exists() + + previous_state = path.read_bytes() + before = len(server.requests) + with terminal_case('restore-print-changed-cwd', env, cwd=next_cwd) as terminal: + terminal.wait(prompt) + terminal.wait(model) + terminal.send(picker.PRINT) + terminal.finish(0) + restored_flags = picker.printed_flags(terminal) + assert restored_flags == flags, (flags, restored_flags) + assert path.read_bytes() == previous_state + assert len(server.requests) == before + assert not list(next_cwd.iterdir()) + + with terminal_case('restore-generation-changed-cwd', env, cwd=next_cwd) as terminal: + mark = len(terminal.raw) + terminal.send(b'\r') + picker.wait_for_request(terminal, server.request_started) + wait_resumed(terminal, mark) + terminal.send(b'\x03') + terminal.finish(130) + assert len(server.requests) == before+1 + check_body(server.requests[-1], prompt, model) + check_images(folder, 2) + assert not list(next_cwd.iterdir()) + assert path.read_bytes() == previous_state + + before = len(server.requests) + with terminal_case('cancel-keeps-previous-state', env) as terminal: + terminal.send(b'\x15') + paste(terminal, 'Changed synthetic draft discarded on cancel') + terminal.wait('Changed synthetic draft') + terminal.send(b'\x03') + terminal.finish(130) + assert path.read_bytes() == previous_state + assert len(server.requests) == before + + with terminal_case('invalid-folder-keeps-previous-state', env) as terminal: + missing = root/'folder that does not exist 雪' + open_folder_editor(terminal) + paste(terminal, str(missing)) + mark = len(terminal.raw) + terminal.send(b'\r') + terminal.wait('Choose an existing folder', after=mark) + assert path.read_bytes() == previous_state + assert not missing.exists() + assert len(server.requests) == before + terminal.send(b'\x03') + terminal.finish(130) + assert path.read_bytes() == previous_state + + completion_home, completion_env = environment('completion') + completed = root/'completion parent'/'snow 雪 folder' + completed.mkdir(parents=True) + with terminal_case('complete-folder-spaces-unicode', completion_env) as terminal: + paste(terminal, 'Synthetic completion prompt') + open_folder_editor(terminal) + paste(terminal, str(completed.parent/'snow')) + mark = len(terminal.raw) + terminal.send(b'\t') + terminal.wait('snow 雪 folder', after=mark) + mark = len(terminal.raw) + terminal.send(b'\r') + terminal.wait('Settings', after=mark) + terminal.send(picker.PRINT) + terminal.finish(0) + assert picker.printed_flags(terminal)['--output-dir'] == str(completed) + assert json.loads(state_path(completion_home, completion_env).read_bytes())['output_dir'] == str(completed) + assert len(server.requests) == before + assert not list(completed.iterdir()) + assert not (completion_home/'Downloads'/'gpt-images').exists() + + default_home, default_env = environment('default-print') + with terminal_case('default-print-creates-no-image-folder', default_env) as terminal: + paste(terminal, 'Synthetic print without generation') + terminal.send(picker.PRINT) + terminal.finish(0) + assert '--output-dir' not in picker.printed_flags(terminal) + assert json.loads(state_path(default_home, default_env).read_bytes())['output_dir'] == '' + assert not (default_home/'Downloads').exists() + assert len(server.requests) == before + + poison_home, poison_env = environment('poisoned-state') + poison_path = state_path(poison_home, poison_env) + poison_path.parent.mkdir(parents=True, mode=0o700) + poison = b'{"version":999,"prompt":"SYNTHETIC_POISON_STATE"}\n' + poison_path.write_bytes(poison) + poison_path.chmod(0o600) + direct_folder = root/'direct output' + direct_folder.mkdir() + for name, arguments, generates in [ + ('poisoned-direct-prompt', ['images', 'generate', '--prompt', 'Explicit synthetic prompt', + '--output-dir', str(direct_folder), '--inline', 'off'], True), + ('poisoned-direct-help', ['images', 'generate', '--help'], False), + ('poisoned-direct-json', ['--format', 'json', 'images', 'generate', '--prompt', 'Structured synthetic prompt'], True), + ]: + before = len(server.requests) + with terminal_case(name, poison_env, arguments, interactive=False) as terminal: + terminal.finish(0, expect_picker=False) + assert len(server.requests) == before+int(generates) + assert poison_path.read_bytes() == poison + assert 'Could not restore saved settings' not in terminal.text() + assert 'Could not remember these settings' not in terminal.text() + assert 'SYNTHETIC_POISON_STATE' not in terminal.text() + if generates: + assert server.requests[-1]['prompt'] == arguments[arguments.index('--prompt')+1] + assert 'output_dir' not in server.requests[-1] + if name == 'poisoned-direct-prompt': + check_images(direct_folder, 1) + elif name == 'poisoned-direct-json': + result = json.loads(terminal.text()) + assert result['data'][0]['b64_json'] + + before = len(server.requests) + result = subprocess.run([binary, '--format', 'json', 'images', 'generate', '--prompt', 'Piped synthetic prompt'], + input=b'', capture_output=True, env=poison_env, timeout=12) + (output/'poisoned-json-pipe.stdout').write_bytes(result.stdout) + (output/'poisoned-json-pipe.stderr').write_bytes(result.stderr) + assert result.returncode == 0, (result.returncode, result.stderr) + assert json.loads(result.stdout)['data'][0]['b64_json'] + assert not result.stderr, result.stderr + assert len(server.requests) == before+1 + assert poison_path.read_bytes() == poison + passed('poisoned-json-pipe') + finally: + server.release.set() + server.shutdown() + server.server_close() + (output/'results.json').write_text(json.dumps(results, indent=2)+'\n') + + +if __name__ == '__main__': + main() From 3b88e86107b6917dbbf0a4a21ac0dbba8ca5ae43 Mon Sep 17 00:00:00 2001 From: Sai Guvvala Date: Fri, 2 Oct 2026 20:58:04 -0700 Subject: [PATCH 2/3] fix(images): remember picker settings without prompt text --- docs/image-generation-saving.md | 7 +- pkg/custom/image_picker.go | 9 +-- pkg/custom/image_picker_folder_test.go | 19 ++--- pkg/custom/image_picker_repeat_test.go | 19 +++-- pkg/custom/image_picker_session_test.go | 5 +- pkg/custom/image_picker_state.go | 16 ++-- pkg/custom/image_picker_state_test.go | 89 ++++++++++++++++------- scripts/check-image-picker-integration.py | 20 +++-- scripts/check-image-picker-reuse.py | 67 +++++++++++++++-- 9 files changed, 184 insertions(+), 67 deletions(-) diff --git a/docs/image-generation-saving.md b/docs/image-generation-saving.md index 3877a320..910b2b87 100644 --- a/docs/image-generation-saving.md +++ b/docs/image-generation-saving.md @@ -19,9 +19,10 @@ zsh, fish and PowerShell 7. In other or unidentified shells, generation remains available, but command printing is disabled with an explanation in the picker. After a successful generation, the picker reopens below the saved result with -the active prompt and settings. It remembers the last submitted prompt and -settings for next time; canceled edits are discarded. Settings stay local to -your user configuration folder and do not affect commands that supply flags. +the same image settings and an empty prompt. It remembers the last submitted +image settings and save folder for next time; the description starts empty and +is not saved in picker preferences. Canceled edits are discarded. Settings stay +local to your user configuration folder and do not affect commands that supply flags. Choose **Save to** to use the default folder, current directory, or another existing folder. Tab completes folder names. Saved settings are optional: if diff --git a/pkg/custom/image_picker.go b/pkg/custom/image_picker.go index f06051db..ecc85614 100644 --- a/pkg/custom/image_picker.go +++ b/pkg/custom/image_picker.go @@ -23,8 +23,6 @@ import ( // imagePickerOptions configures the interactive image settings. type imagePickerOptions struct { Prompt string - // PromptSet lets an explicit empty prompt override the active draft. - PromptSet bool // Shell selects quoting for the displayed command, never execution. Shell string initial *imagePickerSettings @@ -225,10 +223,9 @@ func newImagePicker(options imagePickerOptions) (*imagePicker, error) { } if options.initial != nil { m.settings = *options.initial - if options.PromptSet || options.Prompt != "" { - m.settings.prompt = options.Prompt - } - m.draft, m.cursor = []rune(m.settings.prompt), len([]rune(m.settings.prompt)) + // Carry image settings forward, but require a fresh description for + // each image. Only an explicitly supplied prompt may prefill it. + m.settings.prompt = options.Prompt } m.shell = options.Shell m.note = options.initialNote diff --git a/pkg/custom/image_picker_folder_test.go b/pkg/custom/image_picker_folder_test.go index 049f828f..03c8f667 100644 --- a/pkg/custom/image_picker_folder_test.go +++ b/pkg/custom/image_picker_folder_test.go @@ -38,24 +38,19 @@ func pickerOpenPath(t *testing.T, m *imagePicker) { require.Equal(t, "path", m.focus) } -func TestImagePickerRestoredDraftAndExplicitPrompt(t *testing.T) { +func TestImagePickerRestoredSettingsAndExplicitPrompt(t *testing.T) { original := pickerForTest(t).settings original.prompt, original.quality, original.outputDir = "Remembered 雪", "high", t.TempDir() - for _, tc := range []struct { - prompt string - set bool - want string - }{ - {"", false, original.prompt}, {"", true, ""}, {"Override", false, "Override"}, - } { - m, err := newImagePicker(imagePickerOptions{initial: &original, initialNote: "Remembered settings", Prompt: tc.prompt, PromptSet: tc.set}) + for _, prompt := range []string{"", "Override 雪"} { + m, err := newImagePicker(imagePickerOptions{initial: &original, initialNote: "Remembered settings", Prompt: prompt}) require.NoError(t, err) - require.Equal(t, tc.want, m.settings.prompt) - require.Equal(t, tc.want, string(m.draft)) - require.Equal(t, len([]rune(tc.want)), m.cursor) + require.Equal(t, prompt, m.settings.prompt) + require.Equal(t, prompt, string(m.draft)) + require.Equal(t, len([]rune(prompt)), m.cursor) require.Equal(t, original.outputDir, m.settings.outputDir) require.Equal(t, "high", m.settings.quality) require.Equal(t, "Remembered settings", m.note) + require.Equal(t, "Remembered 雪", original.prompt) } } diff --git a/pkg/custom/image_picker_repeat_test.go b/pkg/custom/image_picker_repeat_test.go index f4d29171..e90de33b 100644 --- a/pkg/custom/image_picker_repeat_test.go +++ b/pkg/custom/image_picker_repeat_test.go @@ -14,19 +14,28 @@ import ( "github.com/urfave/cli/v3" ) -func TestImagePickerResumesAtPreviousDraft(t *testing.T) { +func TestImagePickerResumesWithSettingsAndEmptyPrompt(t *testing.T) { original := pickerForTest(t).settings - original.prompt = "Keep this prompt" + original.prompt = "Previous synthetic prompt" + original.quality, original.outputDir = "high", t.TempDir() m, err := newImagePicker(imagePickerOptions{initial: &original, resuming: true}) require.NoError(t, err) m.width, m.height = 80, 24 require.Equal(t, "settings", m.page) require.Equal(t, "prompt", m.focus) - require.Equal(t, original, m.settings) - require.Equal(t, original.prompt, string(m.draft)) - require.Contains(t, ansi.Strip(m.View().Content), original.prompt) + want := original + want.prompt = "" + require.Equal(t, want, m.settings) + require.Empty(t, m.draft) + require.Zero(t, m.cursor) + require.Contains(t, ansi.Strip(m.View().Content), "Describe your image") + require.NotContains(t, ansi.Strip(m.View().Content), original.prompt) require.Contains(t, ansi.Strip(m.View().Content), "Ctrl+C") require.Empty(t, m.result.Args, "reopening alone must never submit") + pickerKey(m, tea.KeyEnter) + require.Empty(t, m.result.Args, "a new description is required") + require.Equal(t, "Add a prompt first.", m.note) + require.Equal(t, "Previous synthetic prompt", original.prompt, "reopening must not mutate the submitted selection") } func TestImagePickerResumeIgnoresQueuedAndRepeatingSubmit(t *testing.T) { diff --git a/pkg/custom/image_picker_session_test.go b/pkg/custom/image_picker_session_test.go index 3204ade0..973c719d 100644 --- a/pkg/custom/image_picker_session_test.go +++ b/pkg/custom/image_picker_session_test.go @@ -57,6 +57,7 @@ func TestImagePickerSelectionHistoryFollowsSuccessfulEcho(t *testing.T) { before, err := os.ReadFile(path) require.NoError(t, err) model.settings.prompt = "new prompt\n'絵'" + model.settings.quality = "high" result := imagePickerResult{Args: model.settings.args(), settings: model.settings, PrintOnly: true, shell: "bash"} var diagnostic bytes.Buffer err = finishImagePickerSelection(ctx, failedImagePickerWriter{}, &diagnostic, path, result) @@ -71,7 +72,9 @@ func TestImagePickerSelectionHistoryFollowsSuccessfulEcho(t *testing.T) { got, found, err := loadImagePickerState(ctx, path) require.NoError(t, err) require.True(t, found) - require.Equal(t, model.settings, got) + want := model.settings + want.prompt = "" + require.Equal(t, want, got) } func TestImagePickerOptionalHistoryFailureDoesNotBlockSelection(t *testing.T) { diff --git a/pkg/custom/image_picker_state.go b/pkg/custom/image_picker_state.go index 80e8622d..68896ddd 100644 --- a/pkg/custom/image_picker_state.go +++ b/pkg/custom/image_picker_state.go @@ -24,10 +24,9 @@ var ( ) // Keep this record separate from preview preferences and root configuration. -// In particular, it cannot contain connection settings or credentials. +// In particular, it cannot contain prompts, connection settings or credentials. type imagePickerState struct { Version int `json:"version"` - Prompt string `json:"prompt"` Model string `json:"model"` Size string `json:"size"` Quality string `json:"quality"` @@ -132,7 +131,7 @@ func saveImagePickerState(ctx context.Context, path string, settings imagePicker if !validImagePickerState(settings) { return errImagePickerStateInvalid } - state := imagePickerState{1, settings.prompt, settings.model, settings.size, settings.quality, settings.background, settings.format, settings.count, settings.outputDir} + state := imagePickerState{1, settings.model, settings.size, settings.quality, settings.background, settings.format, settings.count, settings.outputDir} data, err := json.Marshal(state) if err != nil { return err @@ -210,7 +209,7 @@ func decodeImagePickerState(data []byte) (imagePickerSettings, bool) { if !utf8.Valid(data) { return s, false } - fields := map[string]*string{"prompt": &s.prompt, "model": &s.model, "size": &s.size, "quality": &s.quality, "background": &s.background, "format": &s.format, "count": &s.count, "output_dir": &s.outputDir} + fields := map[string]*string{"model": &s.model, "size": &s.size, "quality": &s.quality, "background": &s.background, "format": &s.format, "count": &s.count, "output_dir": &s.outputDir} decoder := json.NewDecoder(bytes.NewReader(data)) decoder.UseNumber() token, err := decoder.Token() @@ -233,6 +232,12 @@ func decodeImagePickerState(data []byte) (imagePickerSettings, bool) { if value != json.Number("1") { return s, false } + } else if key == "prompt" { + // Earlier version 1 records included the submitted prompt. Accept + // its old type without restoring it or writing it on the next save. + if _, ok := value.(string); !ok { + return s, false + } } else { text, ok := value.(string) field, known := fields[key] @@ -243,6 +248,7 @@ func decodeImagePickerState(data []byte) (imagePickerSettings, bool) { } } token, err = decoder.Token() + delete(seen, "prompt") // Optional legacy field; all settings remain required. return s, err == nil && token == json.Delim('}') && len(seen) == len(fields)+1 && decoder.Decode(new(any)) == io.EOF && validImagePickerState(s) } @@ -251,7 +257,7 @@ func validImagePickerState(settings imagePickerSettings) bool { return false } total := 0 - for _, value := range []string{settings.prompt, settings.model, settings.size, settings.quality, settings.background, settings.format, settings.count, settings.outputDir} { + for _, value := range []string{settings.model, settings.size, settings.quality, settings.background, settings.format, settings.count, settings.outputDir} { if !utf8.ValidString(value) || len(value) > imagePickerStateLimit-total { return false } diff --git a/pkg/custom/image_picker_state_test.go b/pkg/custom/image_picker_state_test.go index 67967894..4fe0ee6b 100644 --- a/pkg/custom/image_picker_state_test.go +++ b/pkg/custom/image_picker_state_test.go @@ -22,7 +22,6 @@ func pickerStateSettings(t *testing.T) imagePickerSettings { t.Helper() m, err := newImagePicker(imagePickerOptions{}) require.NoError(t, err) - m.settings.prompt = " Synthetic moon over 日本語\nsecond line 🌕\t'quoted' " m.settings.outputDir = filepath.Join(t.TempDir(), "pictures with spaces") return m.settings } @@ -50,7 +49,9 @@ func TestImagePickerStateRoundTripAndPrivateReplacement(t *testing.T) { want := pickerStateSettings(t) for _, folder := range []string{want.outputDir, "", filepath.Join(t.TempDir(), "再利用 folder")} { want.outputDir = folder - require.NoError(t, saveImagePickerState(t.Context(), path, want)) + selection := want + selection.prompt = " Synthetic moon over 日本語\nsecond line 🌕\t'quoted' " + require.NoError(t, saveImagePickerState(t.Context(), path, selection)) got, found, err := loadImagePickerState(t.Context(), path) require.NoError(t, err) require.True(t, found) @@ -71,7 +72,33 @@ func TestImagePickerStateRoundTripAndPrivateReplacement(t *testing.T) { require.NoError(t, err) var fields map[string]any require.NoError(t, json.Unmarshal(data, &fields)) - require.ElementsMatch(t, []string{"version", "prompt", "model", "size", "quality", "background", "format", "count", "output_dir"}, pickerStateKeys(fields)) + require.ElementsMatch(t, []string{"version", "model", "size", "quality", "background", "format", "count", "output_dir"}, pickerStateKeys(fields)) +} + +func TestImagePickerStateIgnoresLegacyPrompt(t *testing.T) { + path := pickerStatePath(t) + want := pickerStateSettings(t) + want.quality = "high" + require.NoError(t, saveImagePickerState(t.Context(), path, want)) + data, err := os.ReadFile(path) + require.NoError(t, err) + var fields map[string]any + require.NoError(t, json.Unmarshal(data, &fields)) + fields["prompt"] = "Previous synthetic description 雪" + legacy, err := json.Marshal(fields) + require.NoError(t, err) + require.NoError(t, os.WriteFile(path, legacy, 0600)) + got, found, err := loadImagePickerState(t.Context(), path) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, want, got) + unchanged, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, legacy, unchanged, "loading must not rewrite saved settings") + require.NoError(t, saveImagePickerState(t.Context(), path, got)) + data, err = os.ReadFile(path) + require.NoError(t, err) + require.NotContains(t, string(data), "prompt") } func pickerStateKeys(fields map[string]any) []string { @@ -89,18 +116,20 @@ func TestImagePickerStateMalformedAndUnknownRemainUntouched(t *testing.T) { valid, err := os.ReadFile(path) require.NoError(t, err) for name, data := range map[string][]byte{ - "empty": {}, - "malformed": []byte("synthetic-private-prompt"), - "missing fields": []byte(`{"version":1}`), - "unknown version": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":2`), 1), - "duplicate": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":1,"version":1`), 1), - "unknown field": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":1,"future":true`), 1), - "trailing object": append(append([]byte{}, valid...), []byte(`{}`)...), - "wrong type": bytes.Replace(valid, []byte(`"count":"1"`), []byte(`"count":1`), 1), - "null": bytes.Replace(valid, []byte(`"count":"1"`), []byte(`"count":null`), 1), - "unknown choice": bytes.Replace(valid, []byte(`"count":"1"`), []byte(`"count":"11"`), 1), - "invalid UTF8": bytes.Replace(valid, []byte(`"count":"1"`), []byte{'"', 'c', 'o', 'u', 'n', 't', '"', ':', '"', 0xff, '"'}, 1), - "oversize": bytes.Repeat([]byte(" "), imagePickerStateLimit+1), + "empty": {}, + "malformed": []byte("synthetic-private-prompt"), + "missing fields": []byte(`{"version":1}`), + "unknown version": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":2`), 1), + "duplicate": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":1,"version":1`), 1), + "unknown field": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":1,"future":true`), 1), + "trailing object": append(append([]byte{}, valid...), []byte(`{}`)...), + "wrong type": bytes.Replace(valid, []byte(`"count":"1"`), []byte(`"count":1`), 1), + "null": bytes.Replace(valid, []byte(`"count":"1"`), []byte(`"count":null`), 1), + "legacy null prompt": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":1,"prompt":null`), 1), + "legacy duplicate prompt": bytes.Replace(valid, []byte(`"version":1`), []byte(`"version":1,"prompt":"one","prompt":"two"`), 1), + "unknown choice": bytes.Replace(valid, []byte(`"count":"1"`), []byte(`"count":"11"`), 1), + "invalid UTF8": bytes.Replace(valid, []byte(`"count":"1"`), []byte{'"', 'c', 'o', 'u', 'n', 't', '"', ':', '"', 0xff, '"'}, 1), + "oversize": bytes.Repeat([]byte(" "), imagePickerStateLimit+1), } { t.Run(name, func(t *testing.T) { require.NoError(t, os.WriteFile(path, data, 0600)) @@ -204,14 +233,14 @@ func TestImagePickerStateWriteFailureKeepsPrevious(t *testing.T) { path := pickerStatePath(t) settings := pickerStateSettings(t) require.NoError(t, saveImagePickerState(t.Context(), path, settings)) - settings.prompt = "a changed synthetic prompt" + settings.quality = "high" require.NoError(t, os.Chmod(filepath.Dir(path), 0500)) defer os.Chmod(filepath.Dir(path), 0700) require.Error(t, saveImagePickerState(t.Context(), path, settings)) got, found, err := loadImagePickerState(t.Context(), path) require.NoError(t, err) require.True(t, found) - require.NotEqual(t, settings.prompt, got.prompt) + require.Equal(t, "auto", got.quality) } type pickerCancelAfterTemporary struct { @@ -239,7 +268,7 @@ func TestImagePickerStateCanceledWriteKeepsPreviousAndCleansTemporary(t *testing before := pickerStateSettings(t) require.NoError(t, saveImagePickerState(t.Context(), path, before)) after := before - after.prompt = "changed synthetic prompt" + after.quality = "high" for _, written := range []bool{false, true} { t.Run(fmt.Sprintf("temporary written=%t", written), func(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) @@ -281,7 +310,7 @@ func TestImagePickerStateConcurrentSavesKeepWholeRecords(t *testing.T) { errs <- fmt.Errorf("concurrent read failed: found %t, %w", found, err) return } - if got != base && got.outputDir != filepath.Join(base.outputDir, strings.TrimPrefix(got.prompt, "synthetic prompt ")) { + if got.prompt != "" || (got != base && got.outputDir != filepath.Join(base.outputDir, got.count)) { errs <- errors.New("concurrent reader saw a mixed record") return } @@ -294,7 +323,8 @@ func TestImagePickerStateConcurrentSavesKeepWholeRecords(t *testing.T) { <-start settings := base settings.prompt = fmt.Sprintf("synthetic prompt %d", i) - settings.outputDir = filepath.Join(base.outputDir, fmt.Sprint(i)) + settings.count = fmt.Sprint(i%10 + 1) + settings.outputDir = filepath.Join(base.outputDir, settings.count) errs <- saveImagePickerState(t.Context(), path, settings) }() } @@ -309,14 +339,14 @@ func TestImagePickerStateConcurrentSavesKeepWholeRecords(t *testing.T) { got, found, err := loadImagePickerState(t.Context(), path) require.NoError(t, err) require.True(t, found) - index := strings.TrimPrefix(got.prompt, "synthetic prompt ") - require.Equal(t, filepath.Join(base.outputDir, index), got.outputDir, "a saved record must not mix writers") + require.Empty(t, got.prompt) + require.Equal(t, filepath.Join(base.outputDir, got.count), got.outputDir, "a saved record must not mix writers") entries, err := os.ReadDir(filepath.Dir(path)) require.NoError(t, err) require.Len(t, entries, 1) } -func TestImagePickerStateLimitOnlyDisablesPersistence(t *testing.T) { +func TestImagePickerStateLimitExcludesPrompt(t *testing.T) { path := pickerStatePath(t) before := pickerStateSettings(t) require.NoError(t, saveImagePickerState(t.Context(), path, before)) @@ -325,10 +355,19 @@ func TestImagePickerStateLimitOnlyDisablesPersistence(t *testing.T) { m := imagePicker{settings: after} require.True(t, m.validPrompt(), "persistence size is not a generation limit") require.Contains(t, after.args(), after.prompt) - require.Error(t, saveImagePickerState(t.Context(), path, after)) + after.quality = "high" + require.NoError(t, saveImagePickerState(t.Context(), path, after)) got, _, err := loadImagePickerState(t.Context(), path) require.NoError(t, err) - require.Equal(t, before, got) + want := after + want.prompt = "" + require.Equal(t, want, got) + // The limit still applies to settings that are actually stored. + after.outputDir = filepath.Join(t.TempDir(), strings.Repeat("p", imagePickerStateLimit+1)) + require.Error(t, saveImagePickerState(t.Context(), path, after)) + got, _, err = loadImagePickerState(t.Context(), path) + require.NoError(t, err) + require.Equal(t, want, got) } func TestImagePickerStateRejectsUnstableSaveFolders(t *testing.T) { diff --git a/scripts/check-image-picker-integration.py b/scripts/check-image-picker-integration.py index 369139c9..965e2edb 100644 --- a/scripts/check-image-picker-integration.py +++ b/scripts/check-image-picker-integration.py @@ -48,6 +48,7 @@ def wait_resumed(terminal, mark): terminal.wait('Saved image:', after=mark) saved = terminal.raw.index(b'Saved image:', mark) terminal.wait('Images', after=saved) + terminal.wait('Describe your image', after=saved) def main(): @@ -191,7 +192,7 @@ def direct_pipe(name, arguments, data=b'', expected_status=1, expected_body=None assert len(server.requests) == before assert not list(home.glob('Downloads/gpt-images/*')) states = list(home.rglob('image-picker.json')) - assert len(states) == 1 and json.loads(states[0].read_text())['prompt'] == prompt + assert len(states) == 1 and 'prompt' not in json.loads(states[0].read_text()) passed(name) finally: try: @@ -254,13 +255,19 @@ def direct_pipe(name, arguments, data=b'', expected_status=1, expected_body=None until = time.monotonic()+1.0 while time.monotonic() < until: terminal.read(0.02) + blank = len(terminal.raw) + terminal.send(b'\r') + terminal.wait('Add a prompt first.', after=blank) + assert len(server.requests) == before+1 server.request_started.clear() second = len(terminal.raw) - terminal.send(b' appended\r') + second_prompt = 'A completely new synthetic image' + terminal.send(b'\x1b[200~'+second_prompt.encode()+b'\x1b[201~\r') picker.wait_for_request(terminal, server.request_started) wait_resumed(terminal, second) - assert server.requests[-1]['prompt'] == prompt+' appended' - prompt += ' appended' + assert server.requests[-2]['prompt'] == prompt + assert server.requests[-1]['prompt'] == second_prompt + prompt = second_prompt terminal.send(b'\x03') terminal.finish(130) model = 'gpt-image-2.5-flare' if action == 'changed-model' else 'gpt-image-2.5-sunburst' @@ -278,7 +285,10 @@ def direct_pipe(name, arguments, data=b'', expected_status=1, expected_body=None states = list(home.rglob('image-picker.json')) assert len(states) == int(action != 'cancel'), (name, states) if states: - assert json.loads(states[0].read_text())['prompt'] == prompt + saved_settings = json.loads(states[0].read_text()) + assert 'prompt' not in saved_settings, saved_settings + expected_model = 'gpt-image-2.5-flare' if action == 'changed-model' else 'gpt-image-2.5-sunburst' + assert saved_settings['model'] == expected_model, saved_settings generates = action not in {'cancel', 'print'} assert len(server.requests)-before == int(generates)+int(action == 'two-generations'), name if action == 'root-connection-options': diff --git a/scripts/check-image-picker-reuse.py b/scripts/check-image-picker-reuse.py index 746b5f28..d89fe874 100644 --- a/scripts/check-image-picker-reuse.py +++ b/scripts/check-image-picker-reuse.py @@ -68,6 +68,7 @@ def wait_resumed(terminal, mark): terminal.wait('Saved image:', after=mark) saved = terminal.raw.index(b'Saved image:', mark) terminal.wait('Images', after=saved) + terminal.wait('Describe your image', after=saved) def main(): @@ -161,24 +162,36 @@ def terminal_case(name, env, arguments=None, cwd=None, interactive=True): check_images(folder, 1) assert flags['--output-dir'] == str(folder), flags saved = json.loads(path.read_bytes()) - assert saved['prompt'] == prompt and saved['model'] == model, saved + assert 'prompt' not in saved and saved['model'] == model, saved assert saved['output_dir'] == str(folder), saved assert not (home/'Downloads'/'gpt-images').exists() previous_state = path.read_bytes() before = len(server.requests) with terminal_case('restore-print-changed-cwd', env, cwd=next_cwd) as terminal: - terminal.wait(prompt) + terminal.wait('Describe your image') terminal.wait(model) + assert prompt not in terminal.text() + mark = len(terminal.raw) + terminal.send(b'\r') + terminal.wait('Add a prompt first.', after=mark) + assert len(server.requests) == before + assert path.read_bytes() == previous_state + print_prompt = 'An entirely new synthetic print prompt' + paste(terminal, print_prompt) terminal.send(picker.PRINT) terminal.finish(0) restored_flags = picker.printed_flags(terminal) - assert restored_flags == flags, (flags, restored_flags) + assert restored_flags == dict(flags, **{'--prompt': print_prompt}), (flags, restored_flags) assert path.read_bytes() == previous_state assert len(server.requests) == before assert not list(next_cwd.iterdir()) with terminal_case('restore-generation-changed-cwd', env, cwd=next_cwd) as terminal: + terminal.wait('Describe your image') + assert prompt not in terminal.text() and print_prompt not in terminal.text() + next_prompt = 'Another synthetic image with the remembered settings' + paste(terminal, next_prompt) mark = len(terminal.raw) terminal.send(b'\r') picker.wait_for_request(terminal, server.request_started) @@ -186,16 +199,20 @@ def terminal_case(name, env, arguments=None, cwd=None, interactive=True): terminal.send(b'\x03') terminal.finish(130) assert len(server.requests) == before+1 - check_body(server.requests[-1], prompt, model) + check_body(server.requests[-1], next_prompt, model) check_images(folder, 2) assert not list(next_cwd.iterdir()) assert path.read_bytes() == previous_state before = len(server.requests) with terminal_case('cancel-keeps-previous-state', env) as terminal: - terminal.send(b'\x15') paste(terminal, 'Changed synthetic draft discarded on cancel') terminal.wait('Changed synthetic draft') + terminal.send(picker.DOWN+b'\r') + terminal.wait('Choose model') + mark = len(terminal.raw) + terminal.send(picker.HOME+b'\r') + terminal.wait('Settings', after=mark) terminal.send(b'\x03') terminal.finish(130) assert path.read_bytes() == previous_state @@ -215,6 +232,46 @@ def terminal_case(name, env, arguments=None, cwd=None, interactive=True): terminal.finish(130) assert path.read_bytes() == previous_state + legacy_home, legacy_env = environment('legacy-state') + legacy_path = state_path(legacy_home, legacy_env) + legacy_path.parent.mkdir(parents=True, mode=0o700) + legacy_prompt = 'SYNTHETIC_LEGACY_PROMPT_MUST_NOT_RESTORE' + legacy = dict(saved, version=1, prompt=legacy_prompt) + legacy_bytes = (json.dumps(legacy)+'\n').encode() + legacy_path.write_bytes(legacy_bytes) + legacy_path.chmod(0o600) + with terminal_case('legacy-state-blank-and-cancel', legacy_env) as terminal: + terminal.wait('Describe your image') + terminal.wait(model) + assert legacy_prompt not in terminal.text() + assert 'Could not restore saved settings' not in terminal.text() + mark = len(terminal.raw) + terminal.send(b'\r') + terminal.wait('Add a prompt first.', after=mark) + assert legacy_path.read_bytes() == legacy_bytes + assert len(server.requests) == before + paste(terminal, 'Legacy synthetic draft discarded on cancel') + terminal.send(b'\x03') + terminal.finish(130) + assert legacy_path.read_bytes() == legacy_bytes + assert len(server.requests) == before + assert legacy_prompt not in terminal.text() + + with terminal_case('legacy-state-submit-removes-prompt', legacy_env) as terminal: + terminal.wait('Describe your image') + legacy_replacement = 'A fresh synthetic prompt with legacy settings' + paste(terminal, legacy_replacement) + terminal.wait(legacy_replacement) + assert legacy_path.read_bytes() == legacy_bytes + terminal.send(picker.PRINT) + terminal.finish(0) + assert picker.printed_flags(terminal) == dict(flags, **{'--prompt': legacy_replacement}) + assert json.loads(legacy_path.read_bytes()) == saved + assert legacy_prompt not in terminal.text() + assert 'Could not restore saved settings' not in terminal.text() + assert 'Could not remember these settings' not in terminal.text() + assert len(server.requests) == before + completion_home, completion_env = environment('completion') completed = root/'completion parent'/'snow 雪 folder' completed.mkdir(parents=True) From 5b6eb05e4b3a21392f16289cdf14b8e8e3c2d97e Mon Sep 17 00:00:00 2001 From: Sai Guvvala Date: Fri, 2 Oct 2026 21:01:02 -0700 Subject: [PATCH 3/3] fix(images): preserve settings with merged picker shell handling --- pkg/custom/image_picker_session.go | 8 ++++---- pkg/custom/image_picker_session_test.go | 23 +++++++++++++++++++---- scripts/check-image-picker-integration.py | 23 +++++++++++++++++++++++ scripts/check-image-picker-reuse.py | 2 +- 4 files changed, 47 insertions(+), 9 deletions(-) diff --git a/pkg/custom/image_picker_session.go b/pkg/custom/image_picker_session.go index 938cb3cd..6d5e687f 100644 --- a/pkg/custom/image_picker_session.go +++ b/pkg/custom/image_picker_session.go @@ -66,10 +66,10 @@ func finishImagePickerSelection(ctx context.Context, output, diagnostics io.Writ if result.PrintOnly { return imageSavingFailure(imagePickerUnsupportedShell, nil) } - return nil - } - if _, err := fmt.Fprintln(output, command); err != nil { - return imageSavingFailure("Could not print the command. No image request was started.", err) + } else { + if _, err := fmt.Fprintln(output, command); err != nil { + return imageSavingFailure("Could not print the command. No image request was started.", err) + } } if path == "" { return nil diff --git a/pkg/custom/image_picker_session_test.go b/pkg/custom/image_picker_session_test.go index 973c719d..4a0af460 100644 --- a/pkg/custom/image_picker_session_test.go +++ b/pkg/custom/image_picker_session_test.go @@ -110,16 +110,31 @@ func TestImagePickerUnsupportedShellSelection(t *testing.T) { for _, shell := range []string{"", "sh", "powershell.exe", "unknown"} { t.Run(shell, func(t *testing.T) { m := pickerForTest(t) + path := filepath.Join(t.TempDir(), "image-picker.json") result := imagePickerResult{Args: m.settings.args(), settings: m.settings, shell: shell} - var output bytes.Buffer - require.NoError(t, finishImagePickerSelection(context.Background(), &output, &output, "", result)) + var output, diagnostics bytes.Buffer + require.NoError(t, finishImagePickerSelection(context.Background(), &output, &diagnostics, path, result)) require.Empty(t, output.String(), "generation must not echo a command for the wrong shell") + require.Empty(t, diagnostics.String()) + got, found, err := loadImagePickerState(t.Context(), path) + require.NoError(t, err) + require.True(t, found, "generation still remembers settings without a printable command") + want := m.settings + want.prompt = "" + require.Equal(t, want, got) + before, err := os.ReadFile(path) + require.NoError(t, err) + result.settings.quality = "high" + result.Args = result.settings.args() result.PrintOnly = true - require.ErrorContains(t, finishImagePickerSelection(context.Background(), &output, &output, "", result), imagePickerUnsupportedShell) + require.ErrorContains(t, finishImagePickerSelection(context.Background(), &output, &diagnostics, path, result), imagePickerUnsupportedShell) require.Empty(t, output.String()) + after, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, before, after, "unsupported command printing must not save a selection") ctx, cancel := context.WithCancel(context.Background()) cancel() - require.ErrorIs(t, finishImagePickerSelection(ctx, &output, &output, "", result), context.Canceled) + require.ErrorIs(t, finishImagePickerSelection(ctx, &output, &diagnostics, path, result), context.Canceled) }) } } diff --git a/scripts/check-image-picker-integration.py b/scripts/check-image-picker-integration.py index 965e2edb..59c49535 100644 --- a/scripts/check-image-picker-integration.py +++ b/scripts/check-image-picker-integration.py @@ -142,6 +142,7 @@ def direct_pipe(name, arguments, data=b'', expected_status=1, expected_body=None assert terminal.child.poll() is None assert len(server.requests) == before assert 'openai images generate ' not in terminal.text() + assert not list(home.rglob('image-picker.json')), 'disabled printing wrote settings' mark = len(terminal.raw) terminal.send(b'\r') wait_resumed(terminal, mark) @@ -149,8 +150,30 @@ def direct_pipe(name, arguments, data=b'', expected_status=1, expected_body=None check_body(server.requests[-1], {'prompt': prompt}) check_saved(home) assert 'openai images generate ' not in terminal.text() + states = list(home.rglob('image-picker.json')) + assert len(states) == 1, states + saved_state = states[0].read_bytes() + assert json.loads(saved_state) == { + 'version': 1, 'model': 'gpt-image-2.5-sunburst', 'size': '1024x1024', + 'quality': 'auto', 'background': 'auto', 'format': 'png', + 'count': '1', 'output_dir': ''}, saved_state terminal.send(b'\x03') terminal.finish(130) + reopened = picker.Terminal(binary, ['images', 'generate'], env) + try: + reopened.wait('Describe your image') + reopened.wait('gpt-image-2.5-sunburst') + assert prompt.splitlines()[0] not in reopened.text() + assert 'Could not restore saved settings' not in reopened.text() + reopened.send(b'\x03') + reopened.finish(130) + assert states[0].read_bytes() == saved_state + assert len(server.requests) == before+1 + finally: + try: + reopened.save(output/(name+'-reopened')) + finally: + reopened.close() passed(name) finally: try: diff --git a/scripts/check-image-picker-reuse.py b/scripts/check-image-picker-reuse.py index d89fe874..621f2359 100644 --- a/scripts/check-image-picker-reuse.py +++ b/scripts/check-image-picker-reuse.py @@ -88,7 +88,7 @@ def main(): thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() base_env = {'PATH': '/usr/bin:/bin', 'TERM': 'xterm-256color', 'LANG': 'en_US.UTF-8', - 'OPENAI_API_KEY': 'synthetic-picker-key', 'CI': 'true', + 'OPENAI_API_KEY': 'synthetic-picker-key', 'OPENAI_PICKER_SHELL': 'bash', 'CI': 'true', 'OPENAI_BASE_URL': f'http://127.0.0.1:{server.server_port}/v1'} def environment(name):