diff --git a/.github/scripts/test_codeql_source.py b/.github/scripts/test_codeql_source.py new file mode 100644 index 00000000..08e01e07 --- /dev/null +++ b/.github/scripts/test_codeql_source.py @@ -0,0 +1,162 @@ +"""Run the reviewed source-preparation validator against real shallow Git data.""" +import copy +import json +import os +from pathlib import Path +import subprocess +import sys +import tarfile +import tempfile +import textwrap +import unittest + +ROOT = Path(__file__).resolve().parents[2] +WORKFLOW = ROOT / ".github/workflows/prepare-ci-source.yml" +CHECKOUT_REF = "${{ inputs.current-pr-merge && (github.event_name == 'pull_request' && github.ref || github.sha) || '' }}" + + +def validator_command(): + workflow = WORKFLOW.read_text() + step = workflow.split(" - name: Verify selected pull request merge\n", 1)[1].split("\n - name:", 1)[0] + return textwrap.dedent(step.split(" run: |\n", 1)[1]) + + +class SourceSelectionTest(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.origin = self.root / "origin.git" + self.candidate = self.root / "candidate" + self.env = dict(os.environ, GIT_AUTHOR_NAME="Fixture", GIT_AUTHOR_EMAIL="fixture@example.test", + GIT_COMMITTER_NAME="Fixture", GIT_COMMITTER_EMAIL="fixture@example.test", + GIT_AUTHOR_DATE="2026-01-01T00:00:00Z", GIT_COMMITTER_DATE="2026-01-01T00:00:00Z", + GITHUB_EVENT_PATH=str(self.root / "event.json"), GITHUB_REF="refs/pull/3/merge", + GITHUB_REPOSITORY="test/cli", RUNNER_TEMP=str(self.root)) + self.git(self.root, "init", "--bare", str(self.origin)) + self.tree = self.git(self.origin, "mktree", data="") + self.base = self.commit("base") + self.head = self.commit("head", self.base) + self.old = self.commit("old merge", self.base, self.head) + self.current = self.commit("regenerated merge", self.base, self.head) + self.env["GITHUB_SHA"] = self.old + self.git(self.origin, "update-ref", "refs/pull/3/merge", self.current) + self.git(self.root, "init", str(self.candidate)) + self.git(self.candidate, "remote", "add", "origin", str(self.origin)) + self.event = {"number": 3, "repository": {"id": 1, "full_name": "test/cli"}, + "pull_request": {"number": 3, "state": "open", "base": { + "sha": self.base, "repo": {"id": 1, "full_name": "test/cli"}}, + "head": {"sha": self.head}}} + + def git(self, cwd, *args, data=None): + return subprocess.check_output(["git", "-C", str(cwd), *args], input=data, + text=True, stderr=subprocess.PIPE, env=self.env).strip() + + def commit(self, message, *parents): + args = ["commit-tree", self.tree] + for parent in parents: + args.extend(["-p", parent]) + return self.git(self.origin, *args, data=message + "\n") + + def checkout(self, ref): + self.git(self.candidate, "fetch", "--depth=1", "origin", ref) + self.git(self.candidate, "checkout", "--detach", "FETCH_HEAD") + + def validate(self, event=None, **env): + (self.root / "event.json").write_text(json.dumps(event if event is not None else self.event)) + return subprocess.run(["bash", "-e", "-c", validator_command()], cwd=self.root, + env=dict(self.env, **env), capture_output=True, text=True) + + def test_current_merge_replaces_stale_event_sha_and_is_packaged_exactly(self): + self.checkout("refs/pull/3/merge") + self.assertNotEqual(self.current, self.old) + self.assertEqual(self.git(self.candidate, "rev-parse", "HEAD"), self.current) + # Ordinary %P output hides shallow parents; the validator must read the + # actual commit object rather than treating this as a parentless commit. + self.assertEqual(self.git(self.candidate, "rev-parse", "--is-shallow-repository"), "true") + result = self.validate() + self.assertEqual(result.returncode, 0, result.stderr) + command = WORKFLOW.read_text().split(" - name: Package candidate checkout\n", 1)[1].splitlines()[0] + subprocess.run(["bash", "-e", "-c", command.removeprefix(" run: ")], + cwd=self.root, env=self.env, check=True, capture_output=True) + with tarfile.open(self.root / "ci-source.tar.gz") as archive: + self.assertEqual(archive.extractfile("./.git/HEAD").read().decode().strip(), self.current) + # The existing recorder must bind actual Git HEAD, not stale GITHUB_SHA. + for language in ["actions", "go"]: + env = dict(self.env, CODEQL_LANGUAGE=language, SARIF_DIRECTORY=str(self.root), + GITHUB_RUN_ID="42", GITHUB_RUN_ATTEMPT="1") + subprocess.run([sys.executable, "-I", str(ROOT / ".github/scripts/codeql_upload.py"), "record"], + cwd=self.candidate, env=env, check=True, capture_output=True) + self.assertEqual(json.loads((self.root / "context.json").read_text()), { + "sha": self.current, "ref": "refs/pull/3/merge", "language": language, + "run_id": 42, "run_attempt": 1}) + + def test_changed_or_unexpected_merge_parents_fail_before_packaging(self): + new_base, new_head = self.commit("new base"), self.commit("new head", self.base) + for parents in [(new_base, self.head), (self.base, new_head), (self.head, self.base), + (self.head,), (self.base, self.head, new_head)]: + with self.subTest(parents=parents): + changed = self.commit("changed merge", *parents) + self.git(self.origin, "update-ref", "refs/pull/3/merge", changed) + self.checkout("refs/pull/3/merge") + result = self.validate() + self.assertNotEqual(result.returncode, 0) + self.assertIn("head or base changed", result.stderr) + self.assertFalse((self.root / "ci-source.tar.gz").exists()) + + def test_wrong_pr_repository_state_or_ref_is_rejected(self): + self.checkout("refs/pull/3/merge") + mutations = [lambda e: e.update(number=True), + lambda e: e["pull_request"].update(number=4), + lambda e: e["pull_request"].update(state="closed"), + lambda e: e["repository"].update(full_name="other/cli"), + lambda e: e["pull_request"]["base"]["repo"].update(id=2), + lambda e: e["pull_request"]["base"]["repo"].update(full_name="other/cli"), + lambda e: e["pull_request"]["head"].update(sha="--not-a-sha")] + for mutation in mutations: + event = copy.deepcopy(self.event) + mutation(event) + with self.subTest(event=event): + self.assertNotEqual(self.validate(event).returncode, 0) + for ref in ["refs/pull/4/merge", "refs/pull/3/head", "refs/heads/main"]: + with self.subTest(ref=ref): + self.assertNotEqual(self.validate(GITHUB_REF=ref).returncode, 0) + + def test_candidate_code_is_data_not_executed(self): + self.checkout("refs/pull/3/merge") + payload = "from pathlib import Path\nPath('executed').write_text('bad')\n" + (self.root / "sitecustomize.py").write_text(payload) + (self.candidate / "sitecustomize.py").write_text(payload) + scripts = self.candidate / ".github/scripts" + scripts.mkdir(parents=True) + (scripts / "codeql_upload.py").write_text(payload) + result = self.validate(PYTHONPATH=str(self.candidate)) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertFalse((self.root / "executed").exists()) + self.assertFalse((self.candidate / "executed").exists()) + + def test_non_pr_events_remain_pinned_when_branch_advances(self): + # The workflow explicitly selects github.sha for opted-in non-PRs. + # Check that exact-SHA checkout stays fixed while the branch advances. + for ref in ["refs/heads/main", "refs/heads/gh-readonly-queue/main/pr-3-test"]: + with self.subTest(ref=ref): + self.git(self.origin, "update-ref", ref, self.current) + self.checkout(self.old) + self.assertEqual(self.git(self.candidate, "rev-parse", "HEAD"), self.old) + + +class WorkflowBoundaryTest(unittest.TestCase): + def test_only_codeql_opts_in_and_other_events_use_exact_event_sha(self): + source = WORKFLOW.read_text() + self.assertIn("type: boolean\n default: false", source) + self.assertIn("ref: " + CHECKOUT_REF, source) + self.assertIn("if: inputs.current-pr-merge && github.event_name == 'pull_request'", source) + self.assertIn("persist-credentials: false", source) + self.assertIn("python3 -I - <<'PYTHON'", validator_command()) + self.assertIn("current-pr-merge: true", (ROOT / ".github/workflows/codeql.yml").read_text()) + for name in ["ci.yml", "help-compatibility.yml"]: + self.assertNotIn("current-pr-merge", (ROOT / ".github/workflows" / name).read_text()) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/test_codeql_upload.py b/.github/scripts/test_codeql_upload.py index e75b1922..55ab3d17 100644 --- a/.github/scripts/test_codeql_upload.py +++ b/.github/scripts/test_codeql_upload.py @@ -154,6 +154,39 @@ def test_required_statuses_wait_for_successful_publication(self): publisher.report(self.get, "test/cli", 1, dict(id=42, run_attempt=2)) self.assertEqual([p["state"] for p in self.writes], ["pending", "pending", "failure", "failure"]) + def test_regenerated_merge_still_requires_exact_analyzed_sha(self): + for index, language in enumerate(publisher.CATEGORIES): + self.responses[f"actions/artifacts/{index}/zip"] = self.archive( + language, mutate=lambda context, report: context.update(sha="d" * 40)) + self.responses["actions/runs/42/artifacts?per_page=100&page=1"] = {"artifacts": [ + dict(id=i, name=f"codeql-sarif-42-2-{lang}", expired=False) + for i, lang in enumerate(publisher.CATEGORIES)]} + # Even if a regenerated merge has the same parents/tree, old reports + # cannot be attributed to its different commit object. + with self.assertRaises(publisher.StaleRun): + publisher.report(self.get, "test/cli", 1, dict(id=42, run_attempt=2)) + self.assertEqual([p["state"] for p in self.writes], ["pending", "pending"]) + + def test_merge_ref_moving_during_artifact_download_cannot_publish(self): + for index, language in enumerate(publisher.CATEGORIES): + self.responses[f"actions/artifacts/{index}/zip"] = self.archive(language) + self.responses["actions/runs/42/artifacts?per_page=100&page=1"] = {"artifacts": [ + dict(id=i, name=f"codeql-sarif-42-2-{lang}", expired=False) + for i, lang in enumerate(publisher.CATEGORIES)]} + + def get(route, **kwargs): + result = self.get(route, **kwargs) + if route.endswith("actions/artifacts/1/zip"): + replacement = "d" * 40 + self.pr["merge_commit_sha"] = replacement + self.responses["git/ref/pull/3/merge"]["object"]["sha"] = replacement + self.responses["git/commits/" + replacement] = {"parents": [{"sha": BASE}, {"sha": HEAD}]} + return result + + with self.assertRaises(publisher.StaleRun): + publisher.report(get, "test/cli", 1, dict(id=42, run_attempt=2)) + self.assertEqual([p["state"] for p in self.writes], ["pending", "pending"]) + def test_pending_and_failed_analysis_cannot_pass_required_status(self): for status, conclusion in [("in_progress", None), ("completed", "failure"), ("completed", "cancelled")]: with self.subTest(status=status, conclusion=conclusion): diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index b977e2d0..f051c6a8 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -16,6 +16,8 @@ permissions: {} jobs: source: uses: ./.github/workflows/prepare-ci-source.yml + with: + current-pr-merge: true permissions: contents: read actions: read @@ -50,6 +52,10 @@ jobs: shell: bash run: tar -xzf - < "$RUNNER_TEMP/ci-inputs/ci-source.tar.gz" + - name: Test CodeQL source and publication boundaries + if: matrix.language == 'actions' + run: python3 -B -I -m unittest discover -s .github/scripts -p 'test_codeql*.py' + - name: Set up Go for internal cross-linking if: matrix.language == 'go' && github.repository == 'openai/openai-cli-internal' && github.event_name == 'pull_request' && github.event.pull_request.head.repo.id == github.event.repository.id && github.event.pull_request.base.ref == 'main' uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 diff --git a/.github/workflows/prepare-ci-source.yml b/.github/workflows/prepare-ci-source.yml index bb316664..9d2821ca 100644 --- a/.github/workflows/prepare-ci-source.yml +++ b/.github/workflows/prepare-ci-source.yml @@ -2,6 +2,11 @@ name: Prepare CI source on: workflow_call: + inputs: + current-pr-merge: + description: Fetch the current PR merge for exact-revision CodeQL publication + type: boolean + default: false outputs: artifact-id: value: ${{ jobs.prepare.outputs.artifact-id }} @@ -25,6 +30,36 @@ jobs: with: path: candidate persist-credentials: false + # An event SHA can describe a previous incarnation of the PR merge. + # Opted-in PRs fetch the merge ref; other callers keep checkout's default. + ref: ${{ inputs.current-pr-merge && (github.event_name == 'pull_request' && github.ref || github.sha) || '' }} + + - name: Verify selected pull request merge + if: inputs.current-pr-merge && github.event_name == 'pull_request' + run: | + python3 -I - <<'PYTHON' + import json, os, re, subprocess + from pathlib import Path + event = json.loads(Path(os.environ["GITHUB_EVENT_PATH"]).read_text()) + pr, repository, number = event["pull_request"], event["repository"], event["number"] + if (type(number) is not int or number < 1 or pr["number"] != number or + pr["state"] != "open" or + os.environ["GITHUB_REF"] != f"refs/pull/{number}/merge" or + repository["full_name"] != os.environ["GITHUB_REPOSITORY"] or + pr["base"]["repo"]["id"] != repository["id"] or + pr["base"]["repo"]["full_name"] != repository["full_name"]): + raise RuntimeError("Unexpected pull request source identity") + expected = [pr["base"]["sha"], pr["head"]["sha"]] + if not all(isinstance(sha, str) and re.fullmatch(r"[0-9a-f]{40}", sha) for sha in expected): + raise RuntimeError("Invalid pull request parent revision") + # cat-file retains parents even in checkout's shallow repository. + # Read Git objects as data; never execute code from candidate/ here. + commit = subprocess.check_output(["git", "-C", "candidate", "cat-file", "commit", "HEAD"], text=True) + parents = [line.removeprefix("parent ") for line in commit.split("\n\n", 1)[0].splitlines() + if line.startswith("parent ")] + if parents != expected: + raise RuntimeError("Pull request head or base changed before source preparation") + PYTHON - name: Package candidate checkout run: tar -czf "$RUNNER_TEMP/ci-source.tar.gz" -C candidate . diff --git a/cmd/openai/image_output_compatibility_unix_test.go b/cmd/openai/image_output_compatibility_unix_test.go new file mode 100644 index 00000000..e60f4c22 --- /dev/null +++ b/cmd/openai/image_output_compatibility_unix_test.go @@ -0,0 +1,154 @@ +//go:build darwin || linux + +package main + +import ( + "bytes" + "context" + "encoding/base64" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/charmbracelet/x/term" +) + +// Run under a private sized PTY. The capture checks one native frame for API +// and stdin cases, and three for progress, between COMPAT-BEGIN/END markers. +func TestMainImageOutputSurvivesExecutableReplacementTerminal(t *testing.T) { + if !term.IsTerminal(os.Stdout.Fd()) { + t.Skip("requires actual terminal stdout") + } + sourcePath, err := os.Executable() + if err != nil { + t.Fatal(err) + } + payload := imageGenerationPNG(t) + for _, phase := range []string{"api", "stdin", "progress"} { + for _, replace := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/replace-%v", phase, replace), func(t *testing.T) { + dir := t.TempDir() + live := filepath.Join(dir, "running cli") + source, err := os.Open(sourcePath) + if err != nil { + t.Fatal(err) + } + copy, err := os.OpenFile(live, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0700) + if err != nil { + source.Close() + t.Fatal(err) + } + _, copyErr := io.Copy(copy, source) + sourceErr := source.Close() + closeErr := copy.Close() + if copyErr != nil || sourceErr != nil || closeErr != nil { + t.Fatalf("copy executable: %v, %v, %v", copyErr, sourceErr, closeErr) + } + change := func() error { + if err := os.Rename(live, filepath.Join(dir, "original executable")); err != nil { + return err + } + if replace { + return os.WriteFile(live, []byte("#!/bin/sh\nexit 23\n"), 0700) + } + return nil + } + requestErr := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, err := io.Copy(io.Discard, r.Body) + if err == nil && phase != "stdin" { + err = change() + } + requestErr <- err + if err != nil { + http.Error(w, "synthetic failure", 500) + return + } + if phase == "progress" { + w.Header().Set("Content-Type", "text/event-stream") + encoded := base64.StdEncoding.EncodeToString(payload) + for index := 0; index < 2; index++ { + fmt.Fprintf(w, "data: {\"type\":\"image_generation.partial_image\",\"partial_image_index\":%d,\"b64_json\":%q}\n\n", index, encoded) + w.(http.Flusher).Flush() + } + fmt.Fprintf(w, "data: {\"type\":\"image_generation.completed\",\"b64_json\":%q}\n\n", encoded) + } else { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, imageGenerationResponse(payload)) + } + })) + defer server.Close() + ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second) + defer cancel() + args := []string{"-test.run=^TestMainDispatchProcess$", "--", "openai", "images", "generate", "--inline", "on", "--name", "result", "--output-dir", dir} + if phase != "stdin" { + args = append(args, "--prompt", "synthetic") + } + if phase == "progress" { + args = append(args, "--partial-images", "2") + } + command := exec.CommandContext(ctx, live, args...) + command.Env = append(imageGenerationEnv(server, dir), "PATH=/usr/bin:/bin", "TERM=xterm-kitty", "TERM_PROGRAM=kitty", "CI=false", "OPENAI_CLI_MAIN_DISPATCH_PROCESS=1") + command.Stdout = os.Stdout + var diagnostic bytes.Buffer + command.Stderr = &diagnostic + var input io.WriteCloser + if phase == "stdin" { + input, err = command.StdinPipe() + if err != nil { + t.Fatal(err) + } + defer input.Close() + } + fmt.Fprintln(os.Stdout, "COMPAT-BEGIN", t.Name()) + if err := command.Start(); err != nil { + t.Fatal(err) + } + if phase == "stdin" { + // Legal JSON whitespace exceeds ordinary Unix pipe capacity. + // Completion proves the CLI is consuming it; withhold EOF and + // the JSON body until its executable has changed. + _, err = io.WriteString(input, strings.Repeat(" ", 4<<20)) + if err == nil { + err = change() + } + if err == nil { + _, err = io.WriteString(input, `{"prompt":"synthetic"}`) + } + closeErr := input.Close() + if err != nil || closeErr != nil { + command.Process.Kill() + command.Wait() + t.Fatalf("prepare stdin replacement: %v, %v", err, closeErr) + } + } + if err := command.Wait(); err != nil { + t.Fatalf("preview after replacement failed: %v; stderr=%q", err, diagnostic.String()) + } + if diagnostic.Len() != 0 { + t.Fatalf("unexpected diagnostic: %q", diagnostic.String()) + } + select { + case err := <-requestErr: + if err != nil { + t.Fatal(err) + } + default: + t.Fatal("synthetic API request missing") + } + original, err := os.ReadFile(filepath.Join(dir, "result.png")) + if err != nil || !bytes.Equal(original, payload) { + t.Fatalf("saved original changed: %v", err) + } + fmt.Fprintln(os.Stdout, "COMPAT-END", t.Name()) + }) + } + } +} diff --git a/cmd/openai/kitty_output_helper_test.go b/cmd/openai/kitty_output_helper_test.go new file mode 100644 index 00000000..9e897ffa --- /dev/null +++ b/cmd/openai/kitty_output_helper_test.go @@ -0,0 +1,19 @@ +package main + +import ( + "os" + "testing" + + "github.com/openai/openai-cli/internal/terminalimage" +) + +// Match the executable dispatch when a terminal test launches its writer helper. +func TestMain(m *testing.M) { + if handled, err := terminalimage.RunKittyOutputHelper(os.Args); handled { + if err != nil { + os.Exit(1) + } + os.Exit(0) + } + os.Exit(m.Run()) +} diff --git a/cmd/openai/main.go b/cmd/openai/main.go index de1e6b8a..10e8c287 100644 --- a/cmd/openai/main.go +++ b/cmd/openai/main.go @@ -6,12 +6,20 @@ import ( "fmt" "os" + "github.com/openai/openai-cli/internal/terminalimage" "github.com/openai/openai-cli/pkg/cmd" "github.com/openai/openai-cli/pkg/custom" "github.com/urfave/cli/v3" ) func main() { + if handled, err := terminalimage.RunKittyOutputHelper(os.Args); handled { + if err != nil { + fmt.Fprintln(os.Stderr, "Image output helper failed.") + os.Exit(1) + } + return + } app := cmd.Command app.Flags = append(app.Flags, cmd.NewRequestHeaderFlag()) diff --git a/internal/terminalimage/image.go b/internal/terminalimage/image.go index a89ba248..2e114e21 100644 --- a/internal/terminalimage/image.go +++ b/internal/terminalimage/image.go @@ -120,13 +120,31 @@ func (w contextWriter) Write(data []byte) (int, error) { // Encode into a checked writer: kitty.EncodeGraphics buffers the entire PNG // without checking cancellation. Reuse Charm's options/framing after encoding. -func writeKittyImage(ctx context.Context, out io.Writer, img image.Image, columns int) error { +func writeKittyImage(ctx context.Context, out io.Writer, img image.Image, columns int) (err error) { var encoded bytes.Buffer if err := imagefont.EncodePNG(ctx, &png.Encoder{}, &encoded, img); err != nil { return err } + return writeKittyOutput(ctx, out, func(destination io.Writer) error { + return writeKittyFrames(ctx, destination, &encoded, columns) + }) +} + +func writeKittyFrames(ctx context.Context, out io.Writer, encoded *bytes.Buffer, columns int) (err error) { options := (&kitty.Options{Action: kitty.TransmitAndPut, Transmission: kitty.Direct, Format: kitty.PNG, Quite: 2, Columns: columns}).Options() destination := contextWriter{ctx, out} + incomplete := false + defer func() { + if incomplete { + // ST closes an APC, but not a chunked upload. Finish that upload + // quietly so a later image cannot be appended to its partial PNG. + // Do not delete images: earlier placements belong to scrollback. + // Keep cancellation: an arbitrary writer may block in cleanup. + // Native terminal workers have a separate bounded transport reset. + _, cleanupErr := io.WriteString(destination, "\x1b_Gq=2,m=0;\x1b\\") + err = errors.Join(err, cleanupErr) + } + }() const rawChunk = kitty.MaxChunkSize / 4 * 3 first := true for encoded.Len() > 0 { @@ -146,14 +164,24 @@ func writeKittyImage(ctx context.Context, out io.Writer, img image.Image, column opts = append(opts, "m=0") } frame := ansi.KittyGraphics(payload, opts...) - if n, err := io.WriteString(destination, frame); err != nil { + n, writeErr := io.WriteString(destination, frame) + if n > 0 { + incomplete = n < len(frame) || encoded.Len() > 0 + } + if writeErr != nil { if n > 0 && n < len(frame) { // Each chunk owns one APC. A partial write can leave it open; - // attempt only ST after cancellation, retaining both errors. - _, cleanupErr := io.WriteString(contextWriter{context.WithoutCancel(ctx), out}, "\x1b\\") - return errors.Join(err, cleanupErr) + // close it before the deferred final chunk, retaining all errors. + terminator := "\x1b\\" + if frame[n-1] == '\x1b' { + // Complete an accepted ESC rather than doubling it, which + // can leave a literal backslash in the terminal's text. + terminator = "\\" + } + _, cleanupErr := io.WriteString(destination, terminator) + return errors.Join(writeErr, cleanupErr) } - return err + return writeErr } first = false } diff --git a/internal/terminalimage/image_test.go b/internal/terminalimage/image_test.go index 7f812baf..e3c73415 100644 --- a/internal/terminalimage/image_test.go +++ b/internal/terminalimage/image_test.go @@ -97,7 +97,7 @@ func TestWriteErrorsAndCancellation(t *testing.T) { require.ErrorIs(t, Write(ctx, writer, img, protocol, 50), context.Canceled) wantWrites := 1 if protocol == "iterm" { - wantWrites++ // Close the streamed OSC after header cancellation. + wantWrites++ // Close the OSC. } require.Equal(t, wantWrites, writes) require.ErrorIs(t, Write(ctx, writer, img, protocol, 50), context.Canceled) diff --git a/internal/terminalimage/kitty_interrupt_darwin.go b/internal/terminalimage/kitty_interrupt_darwin.go new file mode 100644 index 00000000..361ab1db --- /dev/null +++ b/internal/terminalimage/kitty_interrupt_darwin.go @@ -0,0 +1,30 @@ +//go:build darwin + +package terminalimage + +import ( + "syscall" + "unsafe" +) + +// Used only inside the private worker. Never combine these constant dispositions +// with os/signal SIGINT handling or restore a Go handler: its runtime bookkeeping +// and Darwin's user trampoline are deliberately outside this isolated process. +func kittyInterruptDisposition(ignore bool) error { + // Darwin's kernel input is __sigaction, which includes a trampoline slot + // absent from libc's public sigaction structure. Both supported ports are + // 64-bit. Constant SIG_DFL/SIG_IGN need no trampoline, mask or flags. + var action struct { + handler, trampoline uintptr + mask uint32 + flags int32 + } + if ignore { + action.handler = 1 // SIG_IGN; zero is SIG_DFL. + } + _, _, errno := syscall.RawSyscall(syscall.SYS_SIGACTION, uintptr(syscall.SIGINT), uintptr(unsafe.Pointer(&action)), 0) + if errno != 0 { + return errno + } + return nil +} diff --git a/internal/terminalimage/kitty_interrupt_linux.go b/internal/terminalimage/kitty_interrupt_linux.go new file mode 100644 index 00000000..30e112bf --- /dev/null +++ b/internal/terminalimage/kitty_interrupt_linux.go @@ -0,0 +1,32 @@ +//go:build linux + +package terminalimage + +import ( + "runtime" + "syscall" + "unsafe" +) + +// Used only inside the private worker, without os/signal SIGINT registration, +// Go handler restoration, or child processes. The kernel owns active SIGINT. +func kittyInterruptDisposition(ignore bool) error { + // Only the handler is nonzero; every ABI's flags, restorer and mask stay + // zero. Reserve more than the largest kernel action instead of assuming + // identical field order on every Go Linux port. MIPS puts flags first and + // uses a 128-bit mask; other ports put the handler first and use 64 bits. + var action [16]uintptr + maskSize, handler := uintptr(8), 0 + switch runtime.GOARCH { + case "mips", "mipsle", "mips64", "mips64le": + maskSize, handler = 16, 1 + } + if ignore { + action[handler] = 1 // SIG_IGN; zero is SIG_DFL. + } + _, _, errno := syscall.RawSyscall6(syscall.SYS_RT_SIGACTION, uintptr(syscall.SIGINT), uintptr(unsafe.Pointer(&action[0])), 0, maskSize, 0, 0) + if errno != 0 { + return errno + } + return nil +} diff --git a/internal/terminalimage/kitty_output_helper_test.go b/internal/terminalimage/kitty_output_helper_test.go new file mode 100644 index 00000000..83dd65d9 --- /dev/null +++ b/internal/terminalimage/kitty_output_helper_test.go @@ -0,0 +1,17 @@ +package terminalimage + +import ( + "os" + "testing" +) + +// Match the executable dispatch when a terminal test launches its writer helper. +func TestMain(m *testing.M) { + if handled, err := RunKittyOutputHelper(os.Args); handled { + if err != nil { + os.Exit(1) + } + os.Exit(0) + } + os.Exit(m.Run()) +} diff --git a/internal/terminalimage/kitty_output_helper_unix.go b/internal/terminalimage/kitty_output_helper_unix.go new file mode 100644 index 00000000..9bdcdee4 --- /dev/null +++ b/internal/terminalimage/kitty_output_helper_unix.go @@ -0,0 +1,127 @@ +//go:build darwin || linux + +package terminalimage + +import ( + "errors" + "io" + "net" + "os" + + "golang.org/x/sys/unix" +) + +const kittyOutputHelperArgument = "__image-output" + +func kittyReadPipe(fd int) error { + var info unix.Stat_t + if err := unix.Fstat(fd, &info); err != nil || info.Mode&unix.S_IFMT != unix.S_IFIFO { + return errors.New("image output helper requires pipe descriptors") + } + flags, err := unix.FcntlInt(uintptr(fd), unix.F_GETFL, 0) + if err != nil || flags&unix.O_ACCMODE != unix.O_RDONLY { + return errors.New("image output helper requires read-only pipes") + } + return nil +} + +// RunKittyOutputHelper handles only the private resident writer invocation. +// The process itself owns output. Parent death exits it even during a blocked +// write, and the parent can always kill and reap its direct child on failure. +func RunKittyOutputHelper(args []string) (bool, error) { + if len(args) < 2 || args[1] != kittyOutputHelperArgument { + return false, nil + } + if len(args) != 2 { + return true, errors.New("invalid image output helper invocation") + } + if err := kittyReadPipe(3); err != nil { + return true, err + } + socketType, err := unix.GetsockoptInt(4, unix.SOL_SOCKET, unix.SO_TYPE) + if err != nil || socketType != unix.SOCK_STREAM { + return true, errors.New("image output helper requires a control socket") + } + address, err := unix.Getsockname(4) + if _, ok := address.(*unix.SockaddrUnix); err != nil || !ok { + return true, errors.New("image output helper requires a Unix socket") + } + if _, err := unix.Getpeername(4); err != nil { + return true, errors.New("image output helper requires a connected control socket") + } + unix.CloseOnExec(3) + unix.CloseOnExec(4) + if err := unix.SetNonblock(3, true); err != nil { + return true, err + } + life := os.NewFile(3, "image-output-lifeline") + defer life.Close() + control, err := kittyControl(os.NewFile(4, "image-output-control")) + if err != nil { + return true, err + } + defer control.Close() + if err := kittyInterruptDisposition(true); err != nil { + return true, err + } + // Only the helper uses process exit: it has no children or user state to + // unwind. This stops a blocked terminal syscall when the parent disappears. + finished := make(chan struct{}) + defer close(finished) + go func() { + var value [1]byte + _, _ = life.Read(value[:]) + select { + case <-finished: + return + default: + os.Exit(0) + } + }() + if _, _, err := control.WriteMsgUnix([]byte{'R'}, nil, nil); err != nil { + return true, err + } + for { + command, descriptors, err := readKittyControl(control) + if err != nil { + if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) { + return true, nil + } + return true, err + } + if command != 'J' || len(descriptors) != 1 { + closeKittyDescriptors(descriptors) + return true, errors.New("invalid image output job") + } + jobErr := runKittyJob(control, descriptors[0]) + status := byte('0') + if jobErr != nil { + status = '1' + } + if _, _, err := control.WriteMsgUnix([]byte{status}, nil, nil); err != nil { + return true, err + } + } +} + +func runKittyJob(control *net.UnixConn, fd int) error { + input := os.NewFile(uintptr(fd), "image-output-data") + defer input.Close() + if err := kittyReadPipe(fd); err != nil { + return err + } + // Go's SIGINT handler can resume writes after the terminal flushes their + // headers. Use the kernel's default disposition only inside this worker. + // No data is sent until the parent receives this readiness acknowledgement. + if err := kittyInterruptDisposition(false); err != nil { + return err + } + _, _, readyErr := control.WriteMsgUnix([]byte{'A'}, nil, nil) + var writeErr error + if readyErr == nil { + _, writeErr = io.Copy(os.Stdout, input) + } + // Ignore only while idle, including the spare reserved for protocol reset. + // Restoring a Go handler/trampoline is unnecessary and deliberately avoided. + return errors.Join(readyErr, writeErr, kittyInterruptDisposition(true)) +} diff --git a/internal/terminalimage/kitty_output_session_unix.go b/internal/terminalimage/kitty_output_session_unix.go new file mode 100644 index 00000000..e4272e76 --- /dev/null +++ b/internal/terminalimage/kitty_output_session_unix.go @@ -0,0 +1,314 @@ +//go:build darwin || linux + +package terminalimage + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "os" + "os/exec" + "sync" + "sync/atomic" + "syscall" + "time" + + "golang.org/x/sys/unix" +) + +type kittyOutputWorker struct { + command *exec.Cmd + control *net.UnixConn + life *os.File + used atomic.Bool + retired atomic.Bool + closeOnce sync.Once + closeErr error +} + +// A spare prepared before input survives SIGINT while idle and can terminate an +// interrupted frame even if the executable is removed or replaced later. +// A failed worker is retired; normal jobs reuse the current live worker. +type kittyOutputSession struct { + workers []*kittyOutputWorker + gate chan struct{} + used atomic.Bool +} + +func startKittySession(ctx context.Context, path string, out *os.File) (*kittyOutputSession, error) { + s := &kittyOutputSession{gate: make(chan struct{}, 1)} + s.gate <- struct{}{} + for range 2 { + worker, err := startKittyWorker(ctx, path, out) + if err != nil { + return nil, errors.Join(err, s.Close()) + } + s.workers = append(s.workers, worker) + } + return s, nil +} + +func (s *kittyOutputSession) Close() error { + errs := make([]error, len(s.workers)) + var closing sync.WaitGroup + for i, worker := range s.workers { + closing.Go(func() { errs[i] = worker.Close() }) + } + closing.Wait() + var err error + for i, worker := range s.workers { + // Retired workers were already reaped and reported by write, including + // intentional cancellation. Do not poison a successful reset with it. + // An unused reserve cannot affect output that completed elsewhere. + if worker.used.Load() && !worker.retired.Load() { + err = errors.Join(err, errs[i]) + } + } + return err +} + +func (s *kittyOutputSession) write(ctx context.Context, write func(io.Writer) error) (bool, error) { + for _, worker := range s.workers { + if !worker.retired.Load() { + return worker.write(ctx, write) + } + } + return false, errors.New("native image output workers have stopped") +} + +func kittyControl(file *os.File) (*net.UnixConn, error) { + connection, err := net.FileConn(file) + _ = file.Close() + if err != nil { + return nil, err + } + control, ok := connection.(*net.UnixConn) + if !ok { + _ = connection.Close() + return nil, errors.New("image output requires a Unix socket") + } + return control, nil +} + +func startKittyWorker(ctx context.Context, path string, out *os.File) (*kittyOutputWorker, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + lifeRead, lifeWrite, err := os.Pipe() + if err != nil { + return nil, err + } + defer lifeRead.Close() + pair, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM, 0) + if err != nil { + lifeWrite.Close() + return nil, err + } + unix.CloseOnExec(pair[0]) + unix.CloseOnExec(pair[1]) + parent := os.NewFile(uintptr(pair[0]), "image-control-parent") + child := os.NewFile(uintptr(pair[1]), "image-control-child") + defer child.Close() + control, err := kittyControl(parent) + if err != nil { + lifeWrite.Close() + return nil, err + } + command := exec.Command(path, kittyOutputHelperArgument) + command.Stdout = out + command.ExtraFiles = []*os.File{lifeRead, child} + if err := command.Start(); err != nil { + control.Close() + lifeWrite.Close() + return nil, fmt.Errorf("start native image output: %w", err) + } + lifeRead.Close() + child.Close() + session := &kittyOutputWorker{command: command, control: control, life: lifeWrite} + startup, cancelStartup := context.WithTimeout(ctx, 2*time.Second) + defer cancelStartup() + stopped := make(chan struct{}) + stop := context.AfterFunc(startup, func() { defer close(stopped); lifeWrite.Close(); control.Close() }) + ready, descriptors, err := readKittyControl(control) + closeKittyDescriptors(descriptors) + if !stop() { + <-stopped + } + if err != nil || ready != 'R' || len(descriptors) != 0 || startup.Err() != nil { + // No job was sent. An incompatible or stopped executable cannot be + // trusted to honor the lifeline, and has no native writer to orphan. + _ = command.Process.Kill() + return nil, errors.Join(errors.New("native image output helper did not become ready"), err, startup.Err(), session.Close()) + } + return session, nil +} + +func (s *kittyOutputWorker) Close() error { + s.closeOnce.Do(func() { + _ = s.life.Close() + _ = s.control.Close() + // A stopped or unresponsive worker cannot read its lifeline. Bound + // shutdown only; ordinary image output has no timeout. + stopped := make(chan struct{}) + timer := time.AfterFunc(kittyOutputCleanupTimeout, func() { + defer close(stopped) + _ = s.command.Process.Kill() + }) + s.closeErr = kittyWaitError(s.command.Wait()) + if !timer.Stop() { + <-stopped + } + }) + return s.closeErr +} + +func kittyWaitError(err error) error { + if err == nil { + return nil + } + if exit, exited := err.(*exec.ExitError); exited { + // The worker can observe foreground SIGINT before the parent's Go + // cancellation goroutine. Preserve cancellation classification then. + if status, ok := exit.Sys().(syscall.WaitStatus); ok && status.Signaled() && status.Signal() == syscall.SIGINT { + return context.Canceled + } + return fmt.Errorf("write image to terminal: %v", err) + } + return fmt.Errorf("write image to terminal: %w", err) +} + +func (s *kittyOutputWorker) retire() { + s.retired.Store(true) + _ = s.command.Process.Kill() + _ = s.control.Close() +} + +func (s *kittyOutputWorker) write(ctx context.Context, write func(io.Writer) error) (written bool, err error) { + if err := ctx.Err(); err != nil { + return false, err + } + dataRead, dataWrite, err := os.Pipe() + if err != nil { + return false, err + } + defer dataRead.Close() + defer dataWrite.Close() + canceled := make(chan struct{}) + stop := context.AfterFunc(ctx, func() { + defer close(canceled) + s.retire() + _ = dataWrite.Close() + }) + var done chan error + acknowledged := false + completed := false + defer func() { + // A callback that has started must finish before reusing any worker. + if !stop() { + <-canceled + } + err = errors.Join(err, ctx.Err()) + if !completed || err != nil { + s.retire() + _ = dataWrite.Close() + if done != nil && !acknowledged { + <-done + } + // Kill and Wait refer to this directly owned child. No descendant + // can remain writing after we return to text or the reset worker. + err = errors.Join(err, s.Close()) + } + }() + s.used.Store(true) + _, _, err = s.control.WriteMsgUnix([]byte{'J'}, unix.UnixRights(int(dataRead.Fd())), nil) + if err != nil { + return false, err + } + dataRead.Close() + ready, fds, readyErr := readKittyControl(s.control) + closeKittyDescriptors(fds) + if readyErr != nil || ready != 'A' || len(fds) != 0 { + return false, errors.Join(errors.New("native image writer did not become ready"), readyErr) + } + done = make(chan error, 1) + go func() { + status, fds, readErr := readKittyControl(s.control) + closeKittyDescriptors(fds) + if status != '0' || len(fds) != 0 { + readErr = errors.Join(readErr, errors.New("write image to terminal: native writer failed")) + } + if readErr != nil { + s.retire() + _ = dataWrite.Close() + } + done <- readErr + }() + writer := &kittyPipeWriter{Writer: dataWrite} + writeErr := write(writer) + if writeErr != nil { + s.retire() + } + closeErr := dataWrite.Close() + ack := <-done + acknowledged = true + completed = true + return writer.written, errors.Join(writeErr, closeErr, ack) +} + +func closeKittyDescriptors(fds []int) { + for _, fd := range fds { + _ = unix.Close(fd) + } +} + +// Only control bytes and at most one owned pipe descriptor belong here. Keep +// ancillary storage bounded and close all delivered descriptors on rejection. +func readKittyControl(control *net.UnixConn) (command byte, fds []int, err error) { + var data [1]byte + // Reserve above Darwin's 512-FD and Linux's 253-FD SCM_RIGHTS limits. + // On Darwin, truncation can install FDs whose numbers are not returned. + ancillary := make([]byte, 4096) + n, oobn, flags, _, readErr := control.ReadMsgUnix(data[:], ancillary) + remaining := ancillary[:oobn] + for len(remaining) >= unix.CmsgLen(0) { + header, payload, rest, parseErr := unix.ParseOneSocketControlMessage(remaining) + if parseErr != nil { + err = errors.Join(err, parseErr) + break + } + remaining = rest + if header.Level != unix.SOL_SOCKET || header.Type != unix.SCM_RIGHTS { + err = errors.Join(err, errors.New("unexpected image control metadata")) + continue + } + if len(payload)%4 != 0 { + err = errors.Join(err, errors.New("invalid image control descriptors")) + payload = payload[:len(payload)/4*4] + } + rights, parseErr := unix.ParseUnixRights(&unix.SocketControlMessage{Header: header, Data: payload}) + err = errors.Join(err, parseErr) + for _, fd := range rights { + unix.CloseOnExec(fd) + } + fds = append(fds, rights...) + } + if len(fds) > 1 { + err = errors.Join(err, errors.New("too many image control descriptors")) + } + if n == 0 { + err = errors.Join(err, io.EOF) + } else if n != 1 { + err = errors.Join(err, io.ErrUnexpectedEOF) + } + if flags&unix.MSG_CTRUNC != 0 { + err = errors.Join(err, io.ErrUnexpectedEOF) + } + err = errors.Join(err, readErr) + if err != nil { + closeKittyDescriptors(fds) + return 0, nil, err + } + return data[0], fds, nil +} diff --git a/internal/terminalimage/kitty_output_session_unix_test.go b/internal/terminalimage/kitty_output_session_unix_test.go new file mode 100644 index 00000000..6ab2ab40 --- /dev/null +++ b/internal/terminalimage/kitty_output_session_unix_test.go @@ -0,0 +1,379 @@ +//go:build darwin || linux + +package terminalimage + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "syscall" + "testing" + "time" + + "golang.org/x/sys/unix" +) + +func TestKittySessionReusesResidentAfterJobCancellation(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + out, err := os.Create(filepath.Join(t.TempDir(), "output")) + if err != nil { + t.Fatal(err) + } + defer out.Close() + session, err := startKittySession(t.Context(), path, out) + if err != nil { + t.Fatal(err) + } + defer session.Close() + ctx, cancel := context.WithCancel(t.Context()) + _, err = session.write(ctx, func(w io.Writer) error { + _, err := io.WriteString(w, "interrupted") + cancel() + return err + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("cancellation lost: %v", err) + } + _, err = session.write(t.Context(), func(w io.Writer) error { _, err := io.WriteString(w, "cleanup"); return err }) + if err != nil { + t.Fatalf("resident helper could not run cleanup: %v", err) + } + if err := session.Close(); err != nil { + t.Fatal(err) + } + data, err := os.ReadFile(out.Name()) + if err != nil || len(data) < len("cleanup") || string(data[len(data)-len("cleanup"):]) != "cleanup" { + t.Fatalf("cleanup was not complete before acknowledgement: %q, %v", data, err) + } +} + +func TestKittySessionReportsReserveFailureOnlyWhenNeeded(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + for _, needed := range []bool{false, true} { + t.Run(fmt.Sprintf("needed=%t", needed), func(t *testing.T) { + out, err := os.Create(filepath.Join(t.TempDir(), "output")) + if err != nil { + t.Fatal(err) + } + defer out.Close() + session, err := startKittySession(t.Context(), path, out) + if err != nil { + t.Fatal(err) + } + defer session.Close() + if err := session.workers[1].command.Process.Kill(); err != nil { + t.Fatal(err) + } + const image = "complete image" + _, err = session.write(t.Context(), func(w io.Writer) error { + _, err := io.WriteString(w, image) + return err + }) + if err != nil { + t.Fatalf("primary output failed: %v", err) + } + if needed { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + _, err := session.write(ctx, func(io.Writer) error { cancel(); return nil }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("primary cancellation lost: %v", err) + } + _, err = session.write(t.Context(), func(io.Writer) error { + t.Error("failed reserve accepted output") + return nil + }) + if err == nil { + t.Fatal("required reserve failure was hidden") + } + } + if err := session.Close(); err != nil { + t.Fatalf("cleanup reported an unused or already reported failure: %v", err) + } + data, err := os.ReadFile(out.Name()) + if err != nil || string(data) != image { + t.Fatalf("primary output changed: %q, %v", data, err) + } + for _, worker := range session.workers { + if worker.command.ProcessState == nil { + t.Fatal("worker was not reaped") + } + } + }) + } +} + +func TestKittyWorkerFailureReapsBeforeReturning(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + for _, mode := range []string{"killed", "stopped-then-canceled", "interrupt"} { + t.Run(mode, func(t *testing.T) { + outRead, outWrite, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer outRead.Close() + defer outWrite.Close() + session, err := startKittySession(t.Context(), path, outWrite) + if err != nil { + t.Fatal(err) + } + defer session.Close() + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + done := make(chan error, 1) + go func() { + _, err := session.write(ctx, func(out io.Writer) error { + _, err := io.Copy(out, strings.NewReader(strings.Repeat("image", 1<<20))) + return err + }) + done <- err + }() + if err := outRead.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + var first [1]byte + if _, err := outRead.Read(first[:]); err != nil { + t.Fatal(err) + } + worker := session.workers[0] + sig := syscall.SIGKILL + if mode == "stopped-then-canceled" { + sig = syscall.SIGSTOP + } else if mode == "interrupt" { + sig = syscall.SIGINT + } + if err := worker.command.Process.Signal(sig); err != nil { + t.Fatal(err) + } + if mode == "stopped-then-canceled" { + cancel() + } + // Never drain output to make cancellation or helper failure finish. + select { + case err := <-done: + if err == nil || (mode != "killed" && !errors.Is(err, context.Canceled)) { + t.Fatalf("worker failure lost: %v", err) + } + case <-time.After(3 * time.Second): + outRead.Close() + t.Fatal("worker failure left producer blocked") + } + if worker.command.ProcessState == nil { + t.Fatal("worker was not reaped before write returned") + } + if !worker.retired.Load() || session.workers[1].retired.Load() { + t.Fatal("failure consumed the idle reset worker") + } + }) + } +} + +func TestKittyCloseReapsStoppedIdleWorkers(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + out, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0) + if err != nil { + t.Fatal(err) + } + defer out.Close() + session, err := startKittySession(t.Context(), path, out) + if err != nil { + t.Fatal(err) + } + defer session.Close() + if _, err := session.write(t.Context(), func(w io.Writer) error { + _, err := io.WriteString(w, "image") + return err + }); err != nil { + t.Fatal(err) + } + for _, worker := range session.workers { + if err := worker.command.Process.Signal(syscall.SIGSTOP); err != nil { + t.Fatal(err) + } + } + started := time.Now() + if err := session.Close(); err == nil { + t.Fatal("forced shutdown failure was not reported") + } + if elapsed := time.Since(started); elapsed > 4*time.Second { + t.Fatalf("stopped idle workers outlived shutdown: %v", elapsed) + } + for _, worker := range session.workers { + if worker.command.ProcessState == nil { + t.Fatal("idle worker was not reaped") + } + } +} + +func TestKittyJobRejectsWritePipeBeforeReadiness(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + out, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0) + if err != nil { + t.Fatal(err) + } + defer out.Close() + worker, err := startKittyWorker(t.Context(), path, out) + if err != nil { + t.Fatal(err) + } + defer worker.Close() + r, w, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer r.Close() + defer w.Close() + flags, err := unix.FcntlInt(w.Fd(), unix.F_GETFL, 0) + if err != nil { + t.Fatal(err) + } + if _, _, err := worker.control.WriteMsgUnix([]byte{'J'}, unix.UnixRights(int(w.Fd())), nil); err != nil { + t.Fatal(err) + } + status, fds, err := readKittyControl(worker.control) + closeKittyDescriptors(fds) + if err != nil || status != '1' || len(fds) != 0 { + t.Fatalf("invalid job entered output phase: status=%q, err=%v", status, err) + } + after, err := unix.FcntlInt(w.Fd(), unix.F_GETFL, 0) + if err != nil || flags != after { + t.Fatalf("rejected pipe flags changed: %v, %v, %v", flags, after, err) + } +} + +func TestKittyReadinessCancellationReapsUnresponsiveExecutable(t *testing.T) { + path := filepath.Join(t.TempDir(), "unresponsive helper") + if err := os.WriteFile(path, []byte("#!/bin/sh\nexec /bin/sleep 10\n"), 0700); err != nil { + t.Fatal(err) + } + out, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0) + if err != nil { + t.Fatal(err) + } + defer out.Close() + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + started := time.Now() + _, err = startKittySession(ctx, path, out) + if !errors.Is(err, context.DeadlineExceeded) || time.Since(started) > 3*time.Second { + t.Fatalf("unready helper outlived cancellation: %v (%v)", err, time.Since(started)) + } + var status interface{ ExitCode() int } + if errors.As(err, &status) { + t.Fatalf("helper ExitCoder escaped: %v", err) + } +} + +func TestKittyControlRejectsExcessDescriptorsWithoutLeaks(t *testing.T) { + pair, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM, 0) + if err != nil { + t.Fatal(err) + } + left, err := kittyControl(os.NewFile(uintptr(pair[0]), "left")) + if err != nil { + t.Fatal(err) + } + defer left.Close() + right, err := kittyControl(os.NewFile(uintptr(pair[1]), "right")) + if err != nil { + t.Fatal(err) + } + defer right.Close() + r, w, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer r.Close() + defer w.Close() + fds := make([]int, 64) + for i := range fds { + fds[i] = int(w.Fd()) + } + if _, _, err := left.WriteMsgUnix([]byte{'J'}, unix.UnixRights(fds...), nil); err != nil { + t.Fatal(err) + } + _, received, err := readKittyControl(right) + if err == nil || len(received) != 0 { + closeKittyDescriptors(received) + t.Fatalf("excess control descriptors accepted: %v", err) + } + w.Close() + if err := r.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatal(err) + } + var b [1]byte + if _, err := r.Read(b[:]); !errors.Is(err, io.EOF) { + t.Fatalf("transferred writer leaked after rejection: %v", err) + } +} + +func TestKittyHelperRejectsUnconnectedControlBeforeChangingFlags(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + for _, listening := range []bool{false, true} { + t.Run(fmt.Sprintf("listening-%v", listening), func(t *testing.T) { + life, write, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer life.Close() + defer write.Close() + fd, err := unix.Socket(unix.AF_UNIX, unix.SOCK_STREAM, 0) + if err != nil { + t.Fatal(err) + } + socket := os.NewFile(uintptr(fd), "unconnected-control") + defer socket.Close() + if listening { + directory, err := os.MkdirTemp("", "kitty-control-") + if err != nil { + t.Fatal(err) + } + defer os.RemoveAll(directory) + if err := unix.Bind(fd, &unix.SockaddrUnix{Name: filepath.Join(directory, "control")}); err != nil { + t.Fatal(err) + } + if err := unix.Listen(fd, 1); err != nil { + t.Fatal(err) + } + } + flags, err := unix.FcntlInt(life.Fd(), unix.F_GETFL, 0) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + command := exec.CommandContext(ctx, path, kittyOutputHelperArgument) + command.ExtraFiles = []*os.File{life, socket} + if err := command.Run(); err == nil || ctx.Err() != nil { + t.Fatalf("unconnected control accepted or hung: %v, %v", err, ctx.Err()) + } + after, err := unix.FcntlInt(life.Fd(), unix.F_GETFL, 0) + if err != nil || flags != after { + t.Fatalf("rejected invocation changed pipe flags: %v, %v, %v", flags, after, err) + } + }) + } +} diff --git a/internal/terminalimage/kitty_retry_test.go b/internal/terminalimage/kitty_retry_test.go new file mode 100644 index 00000000..a6c6d656 --- /dev/null +++ b/internal/terminalimage/kitty_retry_test.go @@ -0,0 +1,90 @@ +package terminalimage + +import ( + "bytes" + "encoding/base64" + "errors" + "image" + "image/color" + "image/png" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +// kittyTransfers models chunk assembly, not a terminal's visible rendering. +// A new a=T does not replace a pending transfer: subsequent payloads belong to +// that transfer until m=0, including across CLI invocations. See: +// https://sw.kovidgoyal.net/kitty/graphics-protocol/#remote-client +type kittyTransfers struct { + pending []byte + completed [][]byte +} + +func (s *kittyTransfers) accept(t *testing.T, wire []byte) { + t.Helper() + for _, frame := range strings.Split(strings.TrimSuffix(string(wire), "\x1b\\"), "\x1b\\") { + header, payload, ok := strings.Cut(frame, ";") + require.True(t, ok) + require.True(t, strings.HasPrefix(header, "\x1b_G")) + require.Contains(t, header, "q=2", "recovery must not inject a reply into shell input") + require.NotContains(t, header, "a=d", "recovery must preserve earlier images") + decoded, err := base64.StdEncoding.DecodeString(payload) + require.NoError(t, err) + s.pending = append(s.pending, decoded...) + if !strings.Contains(header, "m=1") { + s.completed = append(s.completed, bytes.Clone(s.pending)) + s.pending = nil + } + } +} + +func TestKittyFailedTransferPreservesNextImage(t *testing.T) { + for _, failure := range []string{"error after chunk", "error before next chunk"} { + t.Run(failure, func(t *testing.T) { + var state kittyTransfers + var previous bytes.Buffer + require.NoError(t, Write(t.Context(), &previous, testImage(t), "kitty", 50)) + state.accept(t, previous.Bytes()) + require.Len(t, state.completed, 1) + prior := bytes.Clone(state.completed[0]) + + var interrupted bytes.Buffer + writes := 0 + sinkErr := errors.New("synthetic chunk failure") + out := writerFunc(func(data []byte) (int, error) { + writes++ + if failure == "error before next chunk" && writes == 2 { + return 0, sinkErr + } + n, err := interrupted.Write(data) + if failure == "error after chunk" && writes == 1 { + err = sinkErr + } + return n, err + }) + err := Write(t.Context(), out, testImage(t), "kitty", 50) + require.ErrorIs(t, err, sinkErr) + state.accept(t, interrupted.Bytes()) + assertedPending := len(state.pending) + + img := image.NewNRGBA(image.Rect(0, 0, 5, 3)) + img.SetNRGBA(0, 0, color.NRGBA{R: 200, G: 5, B: 61, A: 128}) + var retry bytes.Buffer + require.NoError(t, Write(t.Context(), &retry, img, "kitty", 20)) + state.accept(t, retry.Bytes()) + rendered, decodeErr := png.Decode(bytes.NewReader(state.completed[len(state.completed)-1])) + require.NoError(t, decodeErr, "the next image must not contain the failed upload") + require.Zero(t, assertedPending, "a closed APC must also complete the chunked transfer") + require.Empty(t, state.pending) + require.Equal(t, img.Bounds(), rendered.Bounds()) + for y := 0; y < img.Bounds().Dy(); y++ { + for x := 0; x < img.Bounds().Dx(); x++ { + require.Equal(t, img.At(x, y), rendered.At(x, y)) + } + } + require.Equal(t, prior, state.completed[0], "earlier image data must remain unchanged") + }) + } +} diff --git a/internal/terminalimage/kitty_test.go b/internal/terminalimage/kitty_test.go index a04f65a5..d39fb2ef 100644 --- a/internal/terminalimage/kitty_test.go +++ b/internal/terminalimage/kitty_test.go @@ -6,17 +6,78 @@ import ( "errors" "fmt" "io" + "strings" "testing" + "time" "github.com/charmbracelet/x/ansi/kitty" "github.com/stretchr/testify/require" ) +func TestKittyCancellationDoesNotEnterBlockingCleanup(t *testing.T) { + for _, at := range []string{"chunk", "partial chunk", "terminator"} { + t.Run(at, func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + cleanupEntered := make(chan struct{}, 2) + release := make(chan struct{}) + writes := 0 + out := writerFunc(func(data []byte) (int, error) { + writes++ + if writes == 1 { + if at == "terminator" { + return 3, io.ErrShortWrite + } + cancel() + if at == "partial chunk" { + return 3, nil + } + return len(data), nil + } + if at == "terminator" && writes == 2 { + cancel() + return len(data), nil + } + cleanupEntered <- struct{}{} + <-release + return len(data), nil + }) + img := testImage(t) + done := make(chan error, 1) + go func() { done <- Write(ctx, out, img, "kitty", 50) }() + select { + case err := <-done: + close(release) + require.ErrorIs(t, err, context.Canceled) + wantWrites := 1 + if at == "terminator" { + wantWrites = 2 + require.ErrorIs(t, err, io.ErrShortWrite) + } + require.Equal(t, wantWrites, writes) + case <-cleanupEntered: + // Release and join the writer even on the old implementation. + close(release) + <-done + t.Fatal("canceled output entered an uninterruptible cleanup write") + case <-time.After(5 * time.Second): + close(release) + <-done + t.Fatal("canceled output did not return") + } + }) + } +} + func TestKittyPartialFrameCleanup(t *testing.T) { for _, frameNumber := range []int{1, 2} { for _, location := range []string{"escape", "header", "payload", "terminator"} { for _, failure := range []string{"error", "short write", "cancellation"} { t.Run(fmt.Sprintf("frame%d/%s/%s", frameNumber, location, failure), func(t *testing.T) { + terminator := "\x1b\\" + if location == "escape" || location == "terminator" { + terminator = "\\" + } ctx, cancel := context.WithCancel(t.Context()) defer cancel() sentinel := errors.New("synthetic partial Kitty write") @@ -29,7 +90,9 @@ func TestKittyPartialFrameCleanup(t *testing.T) { return output.Write(data) } if writes > frameNumber { - require.Equal(t, "\x1b\\", string(data), "cleanup emits only ST") + cleanup := []string{terminator, "\x1b_Gq=2,m=0;\x1b\\"} + require.LessOrEqual(t, writes-frameNumber, len(cleanup)) + require.Equal(t, cleanup[writes-frameNumber-1], string(data), "close APC before finishing the upload") return output.Write(data) } accepted := 1 @@ -42,7 +105,9 @@ func TestKittyPartialFrameCleanup(t *testing.T) { accepted = len(data) - 1 } expected.Write(data[:accepted]) - expected.WriteString("\x1b\\") + if failure != "cancellation" { + expected.WriteString(terminator + "\x1b_Gq=2,m=0;\x1b\\") + } n, err := output.Write(data[:accepted]) if failure == "error" { err = sentinel @@ -60,7 +125,11 @@ func TestKittyPartialFrameCleanup(t *testing.T) { case "cancellation": require.ErrorIs(t, err, context.Canceled) } - require.Equal(t, frameNumber+1, writes) + wantWrites := frameNumber + 2 + if failure == "cancellation" { + wantWrites = frameNumber + } + require.Equal(t, wantWrites, writes) require.Equal(t, expected.Bytes(), output.Bytes()) }) } @@ -68,11 +137,12 @@ func TestKittyPartialFrameCleanup(t *testing.T) { } } -func TestKittyCleanupRetainsBothErrors(t *testing.T) { +func TestKittyCleanupRetainsErrors(t *testing.T) { for _, canceled := range []bool{false, true} { ctx, cancel := context.WithCancel(t.Context()) defer cancel() payloadErr, cleanupErr := errors.New("synthetic Kitty write failure"), errors.New("synthetic ST failure") + finishErr := errors.New("synthetic final-chunk failure") writes := 0 out := writerFunc(func(data []byte) (int, error) { writes++ @@ -82,21 +152,32 @@ func TestKittyCleanupRetainsBothErrors(t *testing.T) { } return 3, payloadErr } - require.Equal(t, "\x1b\\", string(data)) - return 0, cleanupErr + if writes == 2 { + require.Equal(t, "\x1b\\", string(data)) + return 0, cleanupErr + } + require.Equal(t, "\x1b_Gq=2,m=0;\x1b\\", string(data)) + return 0, finishErr }) err := Write(ctx, out, testImage(t), "kitty", 50) require.ErrorIs(t, err, payloadErr) - require.ErrorIs(t, err, cleanupErr) if canceled { require.ErrorIs(t, err, context.Canceled) + require.Equal(t, 1, writes, "cancellation must not enter arbitrary cleanup writes") + } else { + require.ErrorIs(t, err, cleanupErr) + require.ErrorIs(t, err, finishErr) + require.Equal(t, 3, writes, "attempt each cleanup write only once even if it fails") } - require.Equal(t, 2, writes, "attempt cleanup once even if it fails") } } -func TestKittyClosedFramesNeedNoCleanup(t *testing.T) { - for _, frameNumber := range []int{1, 2} { +func TestKittyClosedFramesFinishPendingTransfer(t *testing.T) { + var success bytes.Buffer + require.NoError(t, Write(t.Context(), &success, testImage(t), "kitty", 50)) + lastFrame := strings.Count(success.String(), "\x1b_G") + require.Greater(t, lastFrame, 2) + for _, frameNumber := range []int{1, 2, lastFrame} { for _, accepted := range []bool{false, true} { for _, canceled := range []bool{false, true} { ctx, cancel := context.WithCancel(t.Context()) @@ -104,6 +185,10 @@ func TestKittyClosedFramesNeedNoCleanup(t *testing.T) { writes := 0 out := writerFunc(func(data []byte) (int, error) { writes++ + if writes > frameNumber { + require.Equal(t, "\x1b_Gq=2,m=0;\x1b\\", string(data), "closed APC needs only the final chunk") + return len(data), nil + } n := len(data) if writes == frameNumber { if !accepted { @@ -124,7 +209,11 @@ func TestKittyClosedFramesNeedNoCleanup(t *testing.T) { } else { require.ErrorIs(t, err, sentinel) } - require.Equal(t, frameNumber, writes, "zero acceptance or a complete ST cannot leave a partial frame") + wantWrites := frameNumber + if !canceled && (frameNumber > 1 || accepted) && !(frameNumber == lastFrame && accepted) { + wantWrites++ + } + require.Equal(t, wantWrites, writes, "finish only an upload that may still be pending") } } } diff --git a/internal/terminalimage/kitty_transport_other.go b/internal/terminalimage/kitty_transport_other.go new file mode 100644 index 00000000..5214f632 --- /dev/null +++ b/internal/terminalimage/kitty_transport_other.go @@ -0,0 +1,18 @@ +//go:build !darwin && !linux + +package terminalimage + +import ( + "context" + "io" +) + +func writeKittyOutput(_ context.Context, out io.Writer, write func(io.Writer) error) error { + return write(out) +} + +func RunKittyOutputHelper(_ []string) (bool, error) { return false, nil } + +func PrepareKittyOutput(ctx context.Context, _ io.Writer) (context.Context, func() error) { + return ctx, func() error { return nil } +} diff --git a/internal/terminalimage/kitty_transport_unix.go b/internal/terminalimage/kitty_transport_unix.go new file mode 100644 index 00000000..6b68c508 --- /dev/null +++ b/internal/terminalimage/kitty_transport_unix.go @@ -0,0 +1,114 @@ +//go:build darwin || linux + +package terminalimage + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "time" + + "github.com/charmbracelet/x/term" +) + +const kittyOutputCleanupTimeout = 2 * time.Second + +type kittyOutputContextKey struct{} +type kittyOutputBinding struct { + output *os.File + session *kittyOutputSession + err error +} + +// PrepareKittyOutput binds the writer before a command reads input or waits for +// the API. Preparation errors are deferred until a preview actually uses it. +// Close reports helper failures only after use, preserving non-preview paths. +func PrepareKittyOutput(ctx context.Context, out io.Writer) (context.Context, func() error) { + file, ok := out.(*os.File) + if !ok || !term.IsTerminal(file.Fd()) { + return ctx, func() error { return nil } + } + // A picker owns this binding across its requests. Nested saving workflows + // reuse it without taking ownership of closure or reopening the executable. + // writeKittyOutput still rejects a different output descriptor. + if binding, _ := ctx.Value(kittyOutputContextKey{}).(*kittyOutputBinding); binding != nil { + return ctx, func() error { return nil } + } + binding := &kittyOutputBinding{output: file} + path, err := os.Executable() + if err == nil { + binding.session, err = startKittySession(ctx, path, file) + } + binding.err = err + return context.WithValue(ctx, kittyOutputContextKey{}, binding), func() error { + if binding.session == nil { + return nil + } + err := binding.session.Close() + if !binding.session.used.Load() { + return nil + } + return err + } +} + +// Go's interrupted tty writes can resume after VINTR flushes their header. +// Only a directly owned worker with kernel-default SIGINT writes graphics. +// Terminal modes, user descriptor flags and parent signal handlers are unchanged. +func writeKittyOutput(ctx context.Context, out io.Writer, write func(io.Writer) error) (err error) { + file, ok := out.(*os.File) + if !ok || !term.IsTerminal(file.Fd()) { + return write(out) + } + binding, _ := ctx.Value(kittyOutputContextKey{}).(*kittyOutputBinding) + if binding == nil { + // Direct internal-library callers may have no command preparation phase. + var closeOutput func() error + ctx, closeOutput = PrepareKittyOutput(ctx, out) + defer func() { err = errors.Join(err, closeOutput()) }() + binding, _ = ctx.Value(kittyOutputContextKey{}).(*kittyOutputBinding) + } + if binding == nil || binding.output != file { + return errors.New("native image output changed after preparation") + } + if binding.err != nil { + return fmt.Errorf("prepare native image output: %w", binding.err) + } + session := binding.session + session.used.Store(true) + select { + case <-ctx.Done(): + return ctx.Err() + case <-session.gate: + } + defer func() { session.gate <- struct{}{} }() + written, err := session.write(ctx, write) + if err == nil || !written { + return err + } + // Pipe acceptance does not identify the terminal's cut. Reap failed workers before + // NUL + ST + quiet final chunk. Keep the lease through bounded cleanup. + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), kittyOutputCleanupTimeout) + defer cancel() + _, cleanupErr := session.write(cleanupCtx, func(destination io.Writer) error { + _, err := io.WriteString(destination, "\x00\x1b\\\x1b_Gq=2,m=0;\x1b\\") + return err + }) + if cleanupErr != nil { + cleanupErr = fmt.Errorf("reset interrupted image output: %w", cleanupErr) + } + return errors.Join(err, cleanupErr) +} + +type kittyPipeWriter struct { + io.Writer + written bool +} + +func (w *kittyPipeWriter) Write(data []byte) (int, error) { + n, err := w.Writer.Write(data) + w.written = w.written || n > 0 + return n, err +} diff --git a/internal/terminalimage/kitty_transport_unix_test.go b/internal/terminalimage/kitty_transport_unix_test.go new file mode 100644 index 00000000..8d17308e --- /dev/null +++ b/internal/terminalimage/kitty_transport_unix_test.go @@ -0,0 +1,281 @@ +//go:build darwin || linux + +package terminalimage + +import ( + "bytes" + "context" + "errors" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + + "golang.org/x/sys/unix" +) + +// One-shot session wrapper for focused native transport tests. +func runKittyWriter(ctx context.Context, path string, out *os.File, write func(io.Writer) error) (written bool, err error) { + session, err := startKittySession(ctx, path, out) + if err != nil { + return false, err + } + defer func() { err = errors.Join(err, session.Close()) }() + return session.write(ctx, write) +} + +func TestKittyCatPreservesBytesAndWaits(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + out, err := os.Create(filepath.Join(t.TempDir(), "output")) + if err != nil { + t.Fatal(err) + } + defer out.Close() + payload := strings.Repeat("\x1b_Gq=2,m=1;YWJj\x1b\\", 8192) + "\x1b_Gq=2,m=0;\x1b\\" + written, err := runKittyWriter(context.Background(), path, out, func(destination io.Writer) error { + _, err := io.WriteString(destination, payload) + return err + }) + if err != nil || !written { + t.Fatalf("native output failed: written=%v, err=%v", written, err) + } + got, err := os.ReadFile(out.Name()) + if err != nil { + t.Fatal(err) + } + if string(got) != payload { + t.Fatalf("native output was incomplete when call returned: got %d, want %d bytes", len(got), len(payload)) + } +} + +func TestKittyCatRetainsFailures(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + out, err := os.Create(filepath.Join(t.TempDir(), "output")) + if err != nil { + t.Fatal(err) + } + defer out.Close() + t.Run("spawn", func(t *testing.T) { + called := false + written, err := runKittyWriter(context.Background(), filepath.Join(t.TempDir(), "missing-cat"), out, func(io.Writer) error { + called = true + return nil + }) + if called || written || !errors.Is(err, os.ErrNotExist) { + t.Fatalf("startup failure wrote data or lost cause: called=%v, written=%v, err=%v", called, written, err) + } + }) + t.Run("writer", func(t *testing.T) { + cause := errors.New("synthetic encoder failure") + written, err := runKittyWriter(context.Background(), path, out, func(destination io.Writer) error { + _, err := io.WriteString(destination, "partial") + return errors.Join(cause, err) + }) + if !written || !errors.Is(err, cause) { + t.Fatalf("writer failure lost cause: written=%v, err=%v", written, err) + } + }) + t.Run("cancellation", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + written, err := runKittyWriter(ctx, path, out, func(destination io.Writer) error { + _, err := io.WriteString(destination, "partial") + cancel() + return err + }) + if !written || !errors.Is(err, context.Canceled) { + t.Fatalf("cancellation lost after process wait: written=%v, err=%v", written, err) + } + var exit interface{ ExitCode() int } + if errors.As(err, &exit) { + t.Fatalf("helper exit status escaped into CLI error: %v", err) + } + }) +} + +func TestKittyOutputNonTerminalPreservesWriter(t *testing.T) { + var out bytes.Buffer + cause := errors.New("synthetic output failure") + err := writeKittyOutput(context.Background(), &out, func(destination io.Writer) error { + if destination != &out { + t.Fatal("non-terminal output was wrapped") + } + _, err := io.WriteString(destination, "unchanged") + return errors.Join(cause, err) + }) + if out.String() != "unchanged" || !errors.Is(err, cause) { + t.Fatalf("non-terminal output changed: %q, %v", out.String(), err) + } +} + +func TestKittyCatOutputFailureDoesNotSetCLIExitCode(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + // A read-only descriptor makes the native helper fail its own output write. + out, err := os.Open(os.DevNull) + if err != nil { + t.Fatal(err) + } + defer out.Close() + _, err = runKittyWriter(context.Background(), path, out, func(destination io.Writer) error { + _, err := io.WriteString(destination, "payload") + return err + }) + var exit interface{ ExitCode() int } + if err == nil || errors.As(err, &exit) || !strings.Contains(err.Error(), "write image to terminal") { + t.Fatalf("helper output failure lost or exposed as CLI exit code: %v", err) + } +} + +func TestKittyWriterPanicReapsSupervisor(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + out, err := os.Create(filepath.Join(t.TempDir(), "output")) + if err != nil { + t.Fatal(err) + } + defer out.Close() + cause := errors.New("synthetic writer panic") + defer func() { + if got := recover(); got != cause { + t.Errorf("writer panic was replaced: %v", got) + } + }() + _, _ = runKittyWriter(t.Context(), path, out, func(io.Writer) error { + panic(cause) + }) +} + +func TestKittySupervisorLifelineStopsBlockedWriter(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + outRead, outWrite, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer outRead.Close() + defer outWrite.Close() + session, err := startKittySession(t.Context(), path, outWrite) + if err != nil { + t.Fatal(err) + } + defer session.Close() + writes := make(chan error, 1) + go func() { + _, err := session.write(t.Context(), func(out io.Writer) error { + _, err := io.Copy(out, strings.NewReader(strings.Repeat("image bytes", 1<<18))) + return err + }) + writes <- err + }() + if err := outRead.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + var first [1]byte + if _, err := outRead.Read(first[:]); err != nil { + t.Fatal(err) + } + closed := make(chan error, 1) + go func() { closed <- session.Close() }() + select { + case err := <-closed: + if err != nil { + t.Fatalf("session closure failed: %v", err) + } + case <-time.After(5 * time.Second): + outRead.Close() + t.Fatal("supervisor did not reap blocked native writer") + } + select { + case err := <-writes: + if err == nil { + t.Fatal("blocked producer unexpectedly succeeded") + } + case <-time.After(5 * time.Second): + t.Fatal("native writer retained data pipe after closure") + } +} + +func TestKittyOutputHelperRejectsMissingPipes(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + command := exec.Command(path, kittyOutputHelperArgument) + if err := command.Run(); err == nil { + t.Fatal("private helper accepted an invocation without its pipe descriptors") + } + if handled, err := RunKittyOutputHelper([]string{"openai", kittyOutputHelperArgument, "extra"}); !handled || err == nil { + t.Fatalf("private helper accepted extra arguments: handled=%v, err=%v", handled, err) + } +} + +func TestKittyOutputHelperRejectsInvalidLifeline(t *testing.T) { + path, err := os.Executable() + if err != nil { + t.Fatal(err) + } + for _, scenario := range []string{"regular-file", "write-end", "closed-parent"} { + t.Run(scenario, func(t *testing.T) { + dataRead, dataWrite, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer dataRead.Close() + defer dataWrite.Close() + lifeRead, lifeWrite, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer lifeRead.Close() + defer lifeWrite.Close() + life := lifeRead + switch scenario { + case "regular-file": + life, err = os.Create(filepath.Join(t.TempDir(), "not-a-lifeline")) + if err != nil { + t.Fatal(err) + } + defer life.Close() + case "write-end": + life = lifeWrite + case "closed-parent": + _ = lifeWrite.Close() + } + flags, err := unix.FcntlInt(life.Fd(), unix.F_GETFL, 0) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + command := exec.CommandContext(ctx, path, kittyOutputHelperArgument) + command.Stdin = dataRead + command.ExtraFiles = []*os.File{life} + out, err := command.Output() + if err == nil || ctx.Err() != nil || len(out) != 0 { + t.Fatalf("invalid lifeline emitted data or hung: output=%q, err=%v, context=%v", out, err, ctx.Err()) + } + if scenario != "closed-parent" { + after, err := unix.FcntlInt(life.Fd(), unix.F_GETFL, 0) + if err != nil || after != flags { + t.Fatalf("invalid lifeline flags changed: before=%v, after=%v, err=%v", flags, after, err) + } + } + }) + } +} diff --git a/pkg/custom/image_picker_generate.go b/pkg/custom/image_picker_generate.go index aef852ce..d13107bf 100644 --- a/pkg/custom/image_picker_generate.go +++ b/pkg/custom/image_picker_generate.go @@ -2,6 +2,7 @@ package custom import ( "context" + "errors" "fmt" "os" "slices" @@ -15,13 +16,19 @@ import ( // Only a bare interactive command enters the picker. Explicit flags and piped // bodies retain their existing validation, defaults and output behavior. func imagePickerWorkflow(next cli.ActionFunc) cli.ActionFunc { - return func(ctx context.Context, command *cli.Command) error { + return func(ctx context.Context, command *cli.Command) (resultErr error) { out, ok := command.Root().Writer.(*os.File) if !ok || out != os.Stdout || !term.IsTerminal(os.Stdin.Fd()) || !term.IsTerminal(out.Fd()) || strings.EqualFold(os.Getenv("TERM"), "dumb") || !imagePickerRequested(command) { return next(ctx, command) } options := imagePickerOptions{} + // Keep the executable binding through initial input and every request. + // Each native writer still finishes before the picker reads input again. + ctx, closeOutput := prepareNativeImageOutput(ctx, command) + defer func() { + resultErr = errors.Join(resultErr, closeOutput()) + }() for { result, err := runImagePickerSession(ctx, os.Stdin, out, options) if err != nil { diff --git a/pkg/custom/image_saving_commands.go b/pkg/custom/image_saving_commands.go index 33c7cdd4..40541d7d 100644 --- a/pkg/custom/image_saving_commands.go +++ b/pkg/custom/image_saving_commands.go @@ -99,11 +99,13 @@ type imagePresentation struct { // Prepare once so stdin and @file values are not consumed twice. The generated // handler still owns SDK dispatch and response/stream lifetime. func imageSavingWorkflow(next cli.ActionFunc) cli.ActionFunc { - return func(ctx context.Context, command *cli.Command) error { + return func(ctx context.Context, command *cli.Command) (resultErr error) { inline := command.String("inline") if inline != "auto" && inline != "on" && inline != "off" { return imageSavingFailure("--inline must be auto, on or off.", nil) } + ctx, closeOutput := prepareNativeImageOutput(ctx, command) + defer func() { resultErr = errors.Join(resultErr, closeOutput()) }() if command.Metadata == nil { command.Metadata = make(map[string]any) } diff --git a/pkg/custom/kitty_output_helper_test.go b/pkg/custom/kitty_output_helper_test.go new file mode 100644 index 00000000..f7947312 --- /dev/null +++ b/pkg/custom/kitty_output_helper_test.go @@ -0,0 +1,19 @@ +package custom + +import ( + "os" + "testing" + + "github.com/openai/openai-cli/internal/terminalimage" +) + +// Match the executable dispatch when a terminal test launches its writer helper. +func TestMain(m *testing.M) { + if handled, err := terminalimage.RunKittyOutputHelper(os.Args); handled { + if err != nil { + os.Exit(1) + } + os.Exit(0) + } + os.Exit(m.Run()) +} diff --git a/pkg/custom/prepare_native_image_output.go b/pkg/custom/prepare_native_image_output.go new file mode 100644 index 00000000..019cb82f --- /dev/null +++ b/pkg/custom/prepare_native_image_output.go @@ -0,0 +1,29 @@ +package custom + +import ( + "context" + "os" + "strings" + + "github.com/openai/openai-cli/internal/terminalimage" + "github.com/urfave/cli/v3" +) + +// Bind before stdin or uploads can block. Body-selected URL output may make the +// session unnecessary; its unused startup/close errors must not affect API data. +func prepareNativeImageOutput(ctx context.Context, command *cli.Command) (context.Context, func() error) { + skip := func() (context.Context, func() error) { return ctx, func() error { return nil } } + root := command.Root() + format := strings.ToLower(root.String("format")) + if (format != "" && format != "auto" && format != "text") || root.String("transform") != "" || root.Bool("raw-output") { + return skip() + } + if response, ok := command.Value("response-format").(*string); ok && command.IsSet("response-format") && response != nil && *response == "url" { + return skip() + } + out := imageCommandWriter(command) + if savedImageProtocol(command.String("inline"), isTerminal(out), os.Getenv) != "kitty" { + return skip() + } + return terminalimage.PrepareKittyOutput(ctx, out) +} diff --git a/pkg/custom/preview_saved_image.go b/pkg/custom/preview_saved_image.go index f10e1404..8ff6b3db 100644 --- a/pkg/custom/preview_saved_image.go +++ b/pkg/custom/preview_saved_image.go @@ -49,7 +49,7 @@ func registerImagePreviewCommands(root *cli.Command) { } } -func handleImagesPreview(ctx context.Context, command *cli.Command) error { +func handleImagesPreview(ctx context.Context, command *cli.Command) (resultErr error) { if command.Args().Len() != 1 { return imageSavingFailure("Provide one saved image: openai images preview \"path/to/image.png\". Quote filenames containing spaces.", nil) } @@ -74,6 +74,11 @@ func handleImagesPreview(ctx context.Context, command *cli.Command) error { } ctx, stop := signal.NotifyContext(ctx, os.Interrupt) defer stop() + if protocol == "kitty" { + var closeOutput func() error + ctx, closeOutput = terminalimage.PrepareKittyOutput(ctx, out) + defer func() { resultErr = errors.Join(resultErr, closeOutput()) }() + } img, err := terminalimage.ReadSaved(ctx, path) if err != nil { return imageSavingFailure("Could not preview that file. Choose a readable PNG, JPEG or WebP, at most 64 MiB and 16 megapixels. The original is unchanged.", err) diff --git a/scripts/check-image-output-compatibility.py b/scripts/check-image-output-compatibility.py new file mode 100644 index 00000000..b14350b8 --- /dev/null +++ b/scripts/check-image-output-compatibility.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""Check resident image output after executable replacement, offline. + + go test -c -o /tmp/openai-image.test ./cmd/openai + python3 -I scripts/check-image-output-compatibility.py /tmp/openai-image.test OUTPUT_DIR + +Uses the existing PTY runner; neither a graphical terminal nor image libraries +are required. This verifies protocol pixels, not native visual presentation. +""" +import base64 +import binascii +import json +import os +import pathlib +import re +import struct +import sys +import zlib + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +from terminal_test import run_terminal_test + +TEST = 'TestMainImageOutputSurvivesExecutableReplacementTerminal' +PIXELS = bytes([230, 0, 0, 255, 0, 0, 200, 127]) # imageGenerationPNG fixture + + +def check(condition, message): + if not condition: + raise RuntimeError(message) + + +def check_fixture_png(png): + """Decode only this test's 2x1 RGBA8 PNG, including any PNG row filter.""" + check(png.startswith(b'\x89PNG\r\n\x1a\n'), 'image payload is not PNG') + offset, compressed, header, ended = 8, bytearray(), False, False + while offset < len(png): + check(offset + 12 <= len(png), 'truncated PNG chunk') + length = struct.unpack('>I', png[offset:offset + 4])[0] + kind = png[offset + 4:offset + 8] + end = offset + 8 + length + check(end + 4 <= len(png), 'truncated PNG payload') + data = png[offset + 8:end] + crc = struct.unpack('>I', png[end:end + 4])[0] + check(binascii.crc32(kind + data) & 0xffffffff == crc, 'PNG checksum mismatch') + if kind == b'IHDR': + check(not header and offset == 8, 'unexpected PNG header') + check(data == struct.pack('>IIBBBBB', 2, 1, 8, 6, 0, 0, 0), 'fixture PNG shape changed') + header = True + elif kind == b'IDAT': + compressed.extend(data) + elif kind == b'IEND': + check(length == 0 and end + 4 == len(png), 'invalid PNG end') + ended = True + offset = end + 4 + check(header and ended, 'incomplete PNG') + decoder = zlib.decompressobj() + row = decoder.decompress(compressed, 10) + check(len(row) == 9 and decoder.eof and not decoder.unused_data, 'invalid PNG row data') + filter_type, pixels = row[0], bytearray(row[1:]) + check(filter_type <= 4, 'unknown PNG filter') + for index in range(len(pixels)): + left = pixels[index - 4] if index >= 4 else 0 + # On the first and only row, Up is zero and Paeth selects Left. + prediction = left if filter_type in (1, 4) else left // 2 if filter_type == 3 else 0 + pixels[index] = (pixels[index] + prediction) & 255 + check(pixels == PIXELS, 'rendered fixture pixels changed') + + +def validate(code, captured): + check(code == 0 and b'--- SKIP' not in captured, 'replacement tests failed or skipped') + expected = {TEST + '/' + phase + '/replace-' + replace: 3 if phase == 'progress' else 1 + for phase in ('api', 'stdin', 'progress') for replace in ('false', 'true')} + sections = re.findall(rb'COMPAT-BEGIN ([^\r\n]+)\r?\n(.*?)COMPAT-END \1\r?\n', captured, re.S) + check(len(sections) == 6, 'expected six completed replacement cases') + results = {} + total = 0 + for raw_name, section in sections: + name = raw_name.decode('ascii') + check(name in expected and name not in results, 'unexpected or duplicate replacement case') + frames = re.findall(rb'\x1b_G([^;]*);(.*?)\x1b\\', section, re.S) + check(len(frames) == expected[name], 'wrong native image count: ' + name) + for header, payload in frames: + fields = dict(field.split(b'=', 1) for field in header.split(b',')) + check(all(fields.get(k) == v for k, v in {b'a': b'T', b'f': b'100', b'm': b'0', b'q': b'2'}.items()), + 'unexpected graphics frame: ' + name) + check_fixture_png(base64.b64decode(payload, validate=True)) + total += len(frames) + results[name] = {'images': len(frames), 'exact_rgba': True, 'passed': True} + check(captured.count(b'\x1b_G') == total, 'incomplete or unexpected graphics output') + return results + + +def main(): + if len(sys.argv) != 3: + raise SystemExit('usage: check-image-output-compatibility.py MAIN_TEST OUTPUT_DIR') + binary = str(pathlib.Path(sys.argv[1]).resolve()) + output = pathlib.Path(sys.argv[2]).resolve() + output.mkdir(parents=True, exist_ok=True) + result_path = output / 'image-output-compatibility.json' + # A failed rerun must not leave a previous successful result beside new logs. + result_path.unlink(missing_ok=True) + env = {k: v for k, v in os.environ.items() if not k.startswith('OPENAI_')} + code, captured = run_terminal_test(binary, TEST, env, timeout=100) + (output / 'image-output-compatibility.raw').write_bytes(captured) + results = validate(code, captured) + result_path.write_text(json.dumps(results, indent=2) + '\n') + print(json.dumps(results, indent=2)) + + +if __name__ == '__main__': + main() diff --git a/scripts/check-image-output-lifecycle.py b/scripts/check-image-output-lifecycle.py new file mode 100644 index 00000000..d45825a2 --- /dev/null +++ b/scripts/check-image-output-lifecycle.py @@ -0,0 +1,263 @@ +#!/usr/bin/env python3 +"""Exercise native image helper death and cancellation through the real CLI. + +Run with python3 -I -B scripts/check-image-output-lifecycle.py BINARY OUTPUT. +Uses private PTYs and synthetic loopback responses. Requires ps and, on macOS, +lsof permission to inspect this test's own children. No real terminal settings +or production APIs are used; native graphical appearance is not verified. +""" +import argparse +import base64 +import importlib.util +import json +import os +import pathlib +import random +import re +import signal +import struct +import subprocess +import sys +import tempfile +import threading +import time +import zlib + + +def load_script(name): + spec = importlib.util.spec_from_file_location( + name, pathlib.Path(__file__).with_name(name + '.py')) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +picker = load_script('image_picker_harness') +compatibility = load_script('check-image-output-compatibility') + + +def fixture_png(large=False): + width, height, channels = (1024, 1024, 3) if large else (2, 1, 4) + pixels = random.Random(342).randbytes(width * height * channels) if large else compatibility.PIXELS + stride = width * channels + rows = b''.join(b'\x00' + pixels[i:i + stride] for i in range(0, len(pixels), stride)) + + def chunk(kind, data): + return (struct.pack('>I', len(data)) + kind + data + + struct.pack('>I', zlib.crc32(kind + data) & 0xffffffff)) + + return (b'\x89PNG\r\n\x1a\n' + + chunk(b'IHDR', struct.pack('>IIBBBBB', width, height, 8, 2 if large else 6, 0, 0, 0)) + + chunk(b'IDAT', zlib.compress(rows)) + chunk(b'IEND', b'')) + + +def process_group(group): + # Read executable names, never unrelated process arguments or environments. + output = subprocess.check_output( + ['/bin/ps', '-axo', 'pid=,ppid=,pgid=,stat=,comm='], text=True) + result = {} + for line in output.splitlines(): + fields = line.split(None, 4) + if len(fields) == 5 and int(fields[2]) == group and not fields[3].startswith('Z'): + result[int(fields[0])] = {'parent': int(fields[1]), 'state': fields[3], 'executable': fields[4]} + return result + + +def pipe_count(pid): + if sys.platform == 'linux': + count = 0 + for path in pathlib.Path('/proc', str(pid), 'fd').iterdir(): + try: + count += os.readlink(path).startswith('pipe:') + except FileNotFoundError: + pass # An unrelated descriptor in this owned process closed. + return count + output = subprocess.check_output( + ['/usr/sbin/lsof', '-nP', '-a', '-p', str(pid), '-Fft'], text=True) + return output.splitlines().count('tPIPE') + + +def helpers(terminal): + return {pid: info for pid, info in process_group(terminal.child.pid).items() + if pid != terminal.child.pid} + + +def signal_helper(terminal, pid, value): + # Every process in this private session's group descends from our CLI. + # Only select a live, observed helper; never signal a name-matched process. + assert pid in helpers(terminal), 'owned helper disappeared before signal' + os.kill(pid, value) + + +def stopped_helpers(terminal, timeout=4): + deadline = time.monotonic() + timeout + while True: + remaining = helpers(terminal) + if not remaining or time.monotonic() >= deadline: + return remaining + time.sleep(0.04) # Deliberately do not release terminal backpressure. + + +def close_terminal(terminal): + # The CLI is our session leader. While it is alive, its group cannot be + # reused. If it died, kill only still-observed members of that owned group. + if terminal.child.poll() is None: + try: + os.killpg(terminal.child.pid, signal.SIGKILL) + except ProcessLookupError: + # The process group exited between poll() and killpg(). + pass + else: + for pid in helpers(terminal): + try: + signal_helper(terminal, pid, signal.SIGKILL) + except (ProcessLookupError, AssertionError): + # The owned helper exited between observation and signaling. + pass + terminal.close() + + +class Fixture(picker.http.server.BaseHTTPRequestHandler): + def log_message(self, *_): + pass + + def do_POST(self): + self.server.requests.append(json.loads(self.rfile.read(int(self.headers['Content-Length'])))) + self.server.request_started.set() + self.server.release.wait(15) + data = json.dumps({'data': [{'b64_json': base64.b64encode(self.server.png).decode()}]}).encode() + self.send_response(200) + self.send_header('Content-Type', 'application/json') + self.send_header('Content-Length', str(len(data))) + self.end_headers() + try: + self.wfile.write(data) + except (BrokenPipeError, ConnectionResetError): + pass # Failure cases deliberately disconnect. + + +def run_case(binary, output, name, png): + result = {'case': name, 'passed': False} + with tempfile.TemporaryDirectory(prefix='image-output-lifecycle-') as temporary: + home = pathlib.Path(temporary) + server = picker.http.server.ThreadingHTTPServer(('127.0.0.1', 0), Fixture) + server.png, server.requests = png, [] + server.request_started, server.release = threading.Event(), threading.Event() + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + env = {'PATH': '/usr/bin:/bin', 'HOME': str(home), 'TERM': 'xterm-kitty', + 'TERM_PROGRAM': 'kitty', 'LANG': 'en_US.UTF-8', 'CI': 'false', + 'OPENAI_API_KEY': 'synthetic-lifecycle-key', + 'OPENAI_BASE_URL': f'http://127.0.0.1:{server.server_port}/v1'} + terminal = None + try: + terminal = picker.Terminal(binary, ['images', 'generate', '--prompt', 'synthetic lifecycle', + '--inline', 'on', '--output-dir', str(home), '--name', 'result'], env) + picker.wait_for_request(terminal, server.request_started) + idle = {pid: pipe_count(pid) for pid in helpers(terminal)} + assert idle, 'no prepared native helper observed' + result['idle_pipe_counts'] = idle + server.release.set() + terminal.wait('\x1b_G') + if name == 'normal': + terminal.finish(0, expect_picker=False) + frames = re.findall(rb'\x1b_G([^;]*);(.*?)\x1b\\', terminal.raw, re.S) + assert len(frames) == terminal.raw.count(b'\x1b_G') == 1, 'incomplete image output' + compatibility.check_fixture_png(base64.b64decode(frames[0][1], validate=True)) + result['exact_pixel_previews'] = 1 + else: + # A blocked job retains its data and cancellation pipes. Observe + # the increase instead of depending on helper spawn order: this + # selects the old cat supervisor or a redesigned direct worker. + topology = helpers(terminal) + supervisors = {info['parent'] for info in topology.values() + if pathlib.Path(info['executable']).name == 'cat' and info['parent'] in idle} + counts = {pid: pipe_count(pid) for pid in idle if pid in topology} + active = list(supervisors) if supervisors else [pid for pid, count in counts.items() + if count > idle[pid]] + result['topology_before'] = topology + result['active_pipe_counts'] = counts + assert len(active) == 1, ('cannot identify active output owner', active) + result['active_owner'] = active[0] + started = time.monotonic() + if name == 'owner-death': + signal_helper(terminal, active[0], signal.SIGKILL) + terminal.child.send_signal(signal.SIGINT) + elif name == 'cancel': + terminal.child.send_signal(signal.SIGINT) + else: + terminal.child.send_signal(signal.SIGTERM if name == 'parent-term' else signal.SIGKILL) + remaining = stopped_helpers(terminal) + result['helpers_after_deadline'] = remaining + result['helper_stop_seconds'] = time.monotonic() - started + # Preserve failing evidence, then release/kill only owned work + # in finally. Draining here would hide a blocked-output leak. + assert not remaining, 'native output survived cancellation or owner death' + code = {'owner-death': 130, 'cancel': 130, 'parent-term': -signal.SIGTERM, + 'parent-kill': -signal.SIGKILL}[name] + terminal.finish(code, expect_picker=False) + result['exit'] = code + assert not helpers(terminal), 'helper survived CLI completion' + assert (home / 'result.png').read_bytes() == png, 'saved original changed' + assert len(server.requests) == 1, 'unexpected API request count' + # Exit and joined helpers establish the writer boundary. Linux can + # retain more queued PTY bytes than Terminal.finish's short drain; + # they are old output, not proof of a write after process exit. + drain_deadline = time.monotonic() + 3 + quiet_until = time.monotonic() + 0.1 + while time.monotonic() < quiet_until: + before = len(terminal.raw) + terminal.read(0.02) + if len(terminal.raw) != before: + quiet_until = time.monotonic() + 0.1 + assert time.monotonic() < drain_deadline, 'terminal queue did not become quiet' + after = len(terminal.raw) + until = time.monotonic() + 0.2 + while time.monotonic() < until: + terminal.read(0.02) + assert len(terminal.raw) == after, 'output continued after command completion' + result.update(passed=True, saved_originals=1, requests=1, exit=terminal.child.returncode) + except Exception as error: + result['error'] = str(error) + result['exit_before_cleanup'] = terminal.child.poll() if terminal is not None else None + saved = home / 'result.png' + result['saved_original_preserved'] = saved.is_file() and saved.read_bytes() == png + finally: + server.release.set() + try: + if terminal is not None: + try: + terminal.save(output / name) + finally: + close_terminal(terminal) + finally: + server.shutdown() + server.server_close() + thread.join() + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('binary') + parser.add_argument('output') + args = parser.parse_args() + if sys.platform not in ('darwin', 'linux'): + parser.error('native helper lifecycle checks require macOS or Linux') + binary = str(pathlib.Path(args.binary).resolve()) + output = pathlib.Path(args.output).resolve() + output.mkdir(parents=True, exist_ok=True) + summary = output / 'results.json' + summary.unlink(missing_ok=True) + small, large = fixture_png(), fixture_png(large=True) + results = [] + for name in ('normal', 'cancel', 'owner-death', 'parent-term', 'parent-kill'): + result = run_case(binary, output, name, small if name == 'normal' else large) + results.append(result) + summary.write_text(json.dumps(results, indent=2) + '\n') + print(('PASS' if result['passed'] else 'FAIL'), name, result.get('error', ''), flush=True) + return 0 if all(result['passed'] for result in results) else 1 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/scripts/check-image-picker-native.py b/scripts/check-image-picker-native.py new file mode 100644 index 00000000..b7af19a0 --- /dev/null +++ b/scripts/check-image-picker-native.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Check native previews across picker requests and executable replacement. + +Run with python3 -I -B scripts/check-image-picker-native.py BINARY OUTPUT. +Uses private PTYs, synthetic loopback responses and temporary homes. This checks +protocol pixels and terminal restoration, not graphical terminal appearance. +""" +import argparse +import base64 +import importlib.util +import json +import pathlib +import re +import shutil +import struct +import tempfile +import threading +import time +import zlib + + +def load_script(name): + spec = importlib.util.spec_from_file_location( + name, pathlib.Path(__file__).with_name(name + '.py')) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +picker = load_script('image_picker_harness') +compatibility = load_script('check-image-output-compatibility') + + +def fixture_png(): + def chunk(kind, data): + return (struct.pack('>I', len(data)) + kind + data + + struct.pack('>I', zlib.crc32(kind + data) & 0xffffffff)) + return (b'\x89PNG\r\n\x1a\n' + + chunk(b'IHDR', struct.pack('>IIBBBBB', 2, 1, 8, 6, 0, 0, 0)) + + chunk(b'IDAT', zlib.compress(b'\x00' + compatibility.PIXELS)) + + chunk(b'IEND', b'')) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('binary') + parser.add_argument('output') + args = parser.parse_args() + binary = pathlib.Path(args.binary).resolve() + output = pathlib.Path(args.output).resolve() + output.mkdir(parents=True, exist_ok=True) + summary = output / 'results.json' + summary.unlink(missing_ok=True) + results = [] + png = fixture_png() + cases = [('api', 'unchanged'), ('api', 'removed'), ('api', 'replaced'), + ('picker', 'removed'), ('picker', 'replaced')] + for phase, mode in cases: + name = phase + '-' + mode + with tempfile.TemporaryDirectory(prefix='image-picker-native-') as temporary: + home = pathlib.Path(temporary) + live = home / 'running cli' + shutil.copyfile(binary, live) + live.chmod(0o700) + + def replace_executable(): + live.rename(home / 'original cli') + if mode == 'replaced': + live.write_text('#!/bin/sh\nexit 23\n') + live.chmod(0o700) + + class Fixture(picker.http.server.BaseHTTPRequestHandler): + def log_message(self, *_): + pass + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers['Content-Length']))) + self.server.requests.append(body) + if phase == 'api' and len(self.server.requests) == 1 and mode != 'unchanged': + replace_executable() + data = json.dumps({'data': [{'b64_json': base64.b64encode(png).decode()}]}).encode() + self.send_response(200) + self.send_header('Content-Type', 'application/json') + self.send_header('Content-Length', str(len(data))) + self.end_headers() + self.wfile.write(data) + + server = picker.http.server.ThreadingHTTPServer(('127.0.0.1', 0), Fixture) + server.requests = [] + worker = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + env = {'PATH': '/usr/bin:/bin', 'HOME': str(home), 'TERM': 'xterm-kitty', + 'TERM_PROGRAM': 'kitty', 'LANG': 'en_US.UTF-8', 'CI': 'false', + 'OPENAI_PICKER_SHELL': 'bash', 'OPENAI_API_KEY': 'synthetic-picker-key', + 'OPENAI_BASE_URL': f'http://127.0.0.1:{server.server_port}/v1'} + terminal = None + try: + terminal = picker.Terminal(str(live), ['images', 'generate'], env) + picker.ready(terminal) + if phase == 'picker': + replace_executable() + for number in range(2): + mark = len(terminal.raw) + terminal.send(b'\x1b[200~synthetic image ' + str(number).encode() + b'\x1b[201~\r') + terminal.wait('Saved image:', after=mark) + saved = terminal.raw.index(b'Saved image:', mark) + terminal.wait('Describe your image', after=saved) + # Respect the resumed picker's existing submit-key quiet guard. + until = time.monotonic() + 1 + while time.monotonic() < until: + terminal.read(0.02) + terminal.send(b'\x03') + terminal.finish(130) + assert len(server.requests) == 2, 'unexpected request count' + files = list((home / 'Downloads' / 'gpt-images').glob('*.png')) + assert len(files) == 2 and all(path.read_bytes() == png for path in files) + frames = re.findall(rb'\x1b_G([^;]*);(.*?)\x1b\\', terminal.raw, re.S) + assert len(frames) == terminal.raw.count(b'\x1b_G') == 2, 'missing or partial preview' + for header, payload in frames: + fields = dict(field.split(b'=', 1) for field in header.split(b',')) + assert all(fields.get(k) == v for k, v in + {b'a': b'T', b'f': b'100', b'm': b'0', b'q': b'2'}.items()) + compatibility.check_fixture_png(base64.b64decode(payload, validate=True)) + results.append({'case': name, 'requests': 2, 'exact_pixel_previews': 2, + 'saved_originals': 2, 'picker_returns': 2, 'exit': 130, 'passed': True}) + print('PASS', name, flush=True) + finally: + try: + if terminal is not None: + try: + terminal.save(output / name) + finally: + terminal.close() + finally: + server.shutdown() + server.server_close() + worker.join() + summary.write_text(json.dumps(results, indent=2) + '\n') + + +if __name__ == '__main__': + main()