git.iko.soy/crush
on main help sign in
fix(shell): ctx-aware jq builtin
af002dcChristian Rocha
.agents/skills/shell-builtins/SKILL.md+34−13
compare with main
unifiedsplit
@@ -11,38 +11,59 @@ emulation. Commands can be intercepted before they reach the OS by adding
11
11
12## How Builtins Work
12## How Builtins Work
13
13
14−Builtins live in `Shell.builtinHandler()` in `internal/shell/shell.go`.
14+Builtins live in `builtinHandler()` in `internal/shell/run.go`. This is an
15−This is an `interp.ExecHandlerFunc` middleware registered in
15+`interp.ExecHandlerFunc` middleware registered in `standardHandlers()`
16−`execHandlers()` **before** the block handler, so builtins run even for
16+**before** the block handler, so builtins run even for commands that would
17−commands that would otherwise be blocked.
17+otherwise be blocked. The same handler chain is shared by the stateful
18+`Shell` type and the stateless `Run` entrypoint used by the hook runner,
19+so builtins are available identically in the `bash` tool and in hooks.
18
20
19The handler is a switch on `args[0]`. Each case either handles the command
21The handler is a switch on `args[0]`. Each case either handles the command
20inline or delegates to a helper function.
22inline or delegates to a helper function.
21
23
22## Adding a New Builtin
24## Adding a New Builtin
23
25
24−1. **Add the case** to the switch in `builtinHandler()` in `shell.go`.
26+1. **Add the case** to the switch in `builtinHandler()` in `run.go`.
252. **Get I/O from the handler context**, not from `os.Stdin`/`os.Stdout`.
272. **Get I/O from the handler context**, not from `os.Stdin`/`os.Stdout`.
26 This ensures the builtin works with pipes and redirections:
28 This ensures the builtin works with pipes and redirections:
27 ```go
29 ```go
28 case "mycommand":
30 case "mycommand":
29 hc := interp.HandlerCtx(ctx)
31 hc := interp.HandlerCtx(ctx)
30− return handleMyCommand(args, hc.Stdin, hc.Stdout, hc.Stderr)
32+ return handleMyCommand(ctx, args, hc.Stdin, hc.Stdout, hc.Stderr)
31 ```
33 ```
323. **Implement the handler** in its own file (e.g.,
343. **Implement the handler** in its own file (e.g.,
33− `internal/shell/mycommand.go`). The function signature should accept
35+ `internal/shell/mycommand.go`). The function signature must accept a
34− args, stdin, stdout, and stderr:
36+ `context.Context` as the first parameter, plus args, stdin, stdout, and
37+ stderr:
35 ```go
38 ```go
36− func handleMyCommand(args []string, stdin io.Reader, stdout, stderr io.Writer) error {
39+ func handleMyCommand(ctx context.Context, args []string, stdin io.Reader, stdout, stderr io.Writer) error {
37 // args[0] is the command name ("mycommand"), args[1:] are arguments.
40 // args[0] is the command name ("mycommand"), args[1:] are arguments.
38 // Write output to stdout, errors to stderr.
41 // Write output to stdout, errors to stderr.
39 // Return nil on success, or interp.ExitStatus(n) for non-zero exit codes.
42 // Return nil on success, or interp.ExitStatus(n) for non-zero exit codes.
40 }
43 }
41 ```
44 ```
45+4. **Poll `ctx` in every unbounded loop.** Builtins that iterate over
46+ input, emit values in a generator-style loop, or do any other work
47+ that can exceed a few milliseconds MUST check `ctx.Err()` on each
48+ iteration and return it verbatim when non-nil. Hook timeouts rely on
49+ this: an unbounded builtin that never polls ctx cannot be interrupted
50+ by a hook's `timeout_sec`, and the hook runner will have to abandon
51+ the goroutine (see `internal/hooks/runner.go`). Returning `ctx.Err()`
52+ (not `interp.ExitStatus(n)`) lets callers distinguish "command exited
53+ non-zero" from "we ran out of time".
54+ ```go
55+ for _, item := range items {
56+ if err := ctx.Err(); err != nil {
57+ return err
58+ }
59+ // ... process item
60+ }
61+ ```
42−4. **Return values**: return `nil` for success, `interp.ExitStatus(n)` for
62+5. **Return values**: return `nil` for success, `interp.ExitStatus(n)` for
63+ non-zero exit codes, or `ctx.Err()` on cancellation. Write error
43− non-zero exit codes. Write error messages to `stderr` before returning.
64+ messages to `stderr` before returning.
44−5. **No extra wiring needed** — `builtinHandler()` is already registered
65+6. **No extra wiring needed** — `builtinHandler()` is already registered
45− in `execHandlers()`.
66+ in `standardHandlers()`.
46
67
47## Existing Builtins
68## Existing Builtins
48
69
internal/shell/jq.go+69−4
@@ -1,6 +1,7 @@
1package shell
1package shell
2
2
3import (
3import (
4+ "context"
4 "encoding/json"
5 "encoding/json"
5 "fmt"
6 "fmt"
6 "io"
7 "io"
@@ -36,10 +37,16 @@ Options:
36// flags: -r (raw output), -c (compact output), -s (slurp), -n (null input),
37// flags: -r (raw output), -c (compact output), -s (slurp), -n (null input),
37// -e (exit status), -R (raw input), and --arg name value.
38// -e (exit status), -R (raw input), and --arg name value.
38//
39//
40+// ctx is polled at each iteration of the output loop and at each reader in
41+// [readInputs] so that hook timeouts or other cancellations can interrupt
42+// long-running queries. A cancelled context surfaces as ctx.Err(), not an
43+// [interp.ExitStatus], so callers (e.g. the hook runner) can distinguish
44+// "filter exited non-zero" from "we ran out of time".
45+//
39// Note that this is somewhat of a reimplmentation of the CLI of the glorious
46// Note that this is somewhat of a reimplmentation of the CLI of the glorious
40// github.com/itchyny/gojq, and we'd ideally get the CLI exposed upstream to
47// github.com/itchyny/gojq, and we'd ideally get the CLI exposed upstream to
41// avoid this falling out of sync.
48// avoid this falling out of sync.
42−func handleJQ(args []string, stdin io.Reader, stdout, stderr io.Writer) error {
49+func handleJQ(ctx context.Context, args []string, stdin io.Reader, stdout, stderr io.Writer) error {
43 var (
50 var (
44 rawOutput bool
51 rawOutput bool
45 compact bool
52 compact bool
@@ -143,8 +150,13 @@ func handleJQ(args []string, stdin io.Reader, stdout, stderr io.Writer) error {
143 }
150 }
144
151
145 // Build input values.
152 // Build input values.
146− inputs, err := readInputs(stdin, fileArgs, nullInput, rawInput, slurp)
153+ inputs, err := readInputs(ctx, stdin, fileArgs, nullInput, rawInput, slurp)
147 if err != nil {
154 if err != nil {
155+ // Prefer surfacing ctx cancellation verbatim so timeouts are
156+ // distinguishable from user input errors.
157+ if ctxErr := ctx.Err(); ctxErr != nil {
158+ return ctxErr
159+ }
148 fmt.Fprintf(stderr, "jq: %s\n", err)
160 fmt.Fprintf(stderr, "jq: %s\n", err)
149 return interp.ExitStatus(2)
161 return interp.ExitStatus(2)
150 }
162 }
@@ -153,6 +165,12 @@ func handleJQ(args []string, stdin io.Reader, stdout, stderr io.Writer) error {
153 for _, input := range inputs {
165 for _, input := range inputs {
154 iter := code.Run(input, argValues...)
166 iter := code.Run(input, argValues...)
155 for {
167 for {
168+ // Poll ctx on every value so a long-running filter (e.g. a
169+ // generator over a slurped array) can be interrupted by hook
170+ // timeouts without waiting for iter.Next to yield.
171+ if err := ctx.Err(); err != nil {
172+ return err
173+ }
156 v, ok := iter.Next()
174 v, ok := iter.Next()
157 if !ok {
175 if !ok {
158 break
176 break
@@ -177,7 +195,22 @@ func handleJQ(args []string, stdin io.Reader, stdout, stderr io.Writer) error {
177}
195}
178
196
179// readInputs reads JSON (or raw) input values from stdin or files.
197// readInputs reads JSON (or raw) input values from stdin or files.
198+//
199+// ctx is polled in three places so that a cancellation observed mid-read
200+// short-circuits promptly:
201+// - between readers (before opening the next file / consuming stdin);
202+// - on every io.Read call via ctxReader, so io.ReadAll on a large but
203+// non-blocking source (e.g. the bytes.NewReader payload the hook
204+// runner supplies) returns ctx.Err() on the next chunk boundary;
205+// - inside the post-read value accumulation loops (raw-input line
206+// split and JSON stream decode), which are otherwise unbounded in
207+// the size of the input.
208+//
209+// A reader that blocks forever in Read (e.g. an unterminated pipe) can
210+// still outlast ctx; the outer abandon-goroutine path in the hook
211+// runner (internal/hooks/runner.go) is the authoritative enforcer for
212+// that case.
180−func readInputs(stdin io.Reader, files []string, nullInput, rawInput, slurp bool) ([]any, error) {
213+func readInputs(ctx context.Context, stdin io.Reader, files []string, nullInput, rawInput, slurp bool) ([]any, error) {
181 if nullInput {
214 if nullInput {
182 return []any{nil}, nil
215 return []any{nil}, nil
183 }
216 }
@@ -198,8 +231,16 @@ func readInputs(stdin io.Reader, files []string, nullInput, rawInput, slurp bool
198
231
199 var vals []any
232 var vals []any
200 for _, r := range readers {
233 for _, r := range readers {
234+ if err := ctx.Err(); err != nil {
235+ return nil, err
236+ }
201− data, err := io.ReadAll(r)
237+ data, err := io.ReadAll(ctxReader{ctx: ctx, r: r})
202 if err != nil {
238 if err != nil {
239+ // ctxReader surfaces ctx.Err() verbatim; preserve it so the
240+ // caller can distinguish cancellation from a parse error.
241+ if ctxErr := ctx.Err(); ctxErr != nil {
242+ return nil, ctxErr
243+ }
203 return nil, err
244 return nil, err
204 }
245 }
205
246
@@ -209,6 +250,9 @@ func readInputs(stdin io.Reader, files []string, nullInput, rawInput, slurp bool
209 vals = append(vals, strings.Join(lines, "\n"))
250 vals = append(vals, strings.Join(lines, "\n"))
210 } else {
251 } else {
211 for _, line := range lines {
252 for _, line := range lines {
253+ if err := ctx.Err(); err != nil {
254+ return nil, err
255+ }
212 if line != "" || !slurp {
256 if line != "" || !slurp {
213 vals = append(vals, line)
257 vals = append(vals, line)
214 }
258 }
@@ -221,6 +265,9 @@ func readInputs(stdin io.Reader, files []string, nullInput, rawInput, slurp bool
221 dec := json.NewDecoder(strings.NewReader(string(data)))
265 dec := json.NewDecoder(strings.NewReader(string(data)))
222 var streamVals []any
266 var streamVals []any
223 for {
267 for {
268+ if err := ctx.Err(); err != nil {
269+ return nil, err
270+ }
224 var v any
271 var v any
225 if err := dec.Decode(&v); err != nil {
272 if err := dec.Decode(&v); err != nil {
226 if err == io.EOF {
273 if err == io.EOF {
@@ -244,6 +291,24 @@ func readInputs(stdin io.Reader, files []string, nullInput, rawInput, slurp bool
244 return vals, nil
291 return vals, nil
245}
292}
246
293
294+// ctxReader wraps an io.Reader so that each Read call checks ctx first.
295+// This makes io.ReadAll over a large but non-blocking source (e.g. a
296+// bytes.Reader of the hook stdin payload) cancellable on the next chunk
297+// boundary. A reader that itself blocks in Read will still outlast ctx —
298+// the hook runner's abandon-goroutine path is the enforcer of last resort
299+// for that case.
300+type ctxReader struct {
301+ ctx context.Context
302+ r io.Reader
303+}
304+
305+func (cr ctxReader) Read(p []byte) (int, error) {
306+ if err := cr.ctx.Err(); err != nil {
307+ return 0, err
308+ }
309+ return cr.r.Read(p)
310+}
311+
247// writeValue writes a single jq output value.
312// writeValue writes a single jq output value.
248func writeValue(w io.Writer, v any, raw, compact, join bool) error {
313func writeValue(w io.Writer, v any, raw, compact, join bool) error {
249 if raw {
314 if raw {
internal/shell/jq_test.gonew file+199
@@ -0,0 +1,199 @@
1+package shell
2+
3+import (
4+ "bytes"
5+ "context"
6+ "errors"
7+ "io"
8+ "strings"
9+ "testing"
10+ "time"
11+)
12+
13+// TestJQ_CtxCancel verifies that handleJQ polls ctx during iteration and
14+// returns ctx.Err() (not an interp.ExitStatus) when the context is
15+// cancelled. This is what lets hook timeouts interrupt long-running jq
16+// filters rather than waiting for the iterator to terminate naturally.
17+func TestJQ_CtxCancel(t *testing.T) {
18+ t.Parallel()
19+
20+ // `range(N)` generates a large stream of values. With a slurped input
21+ // the filter produces all N values in sequence; ctx cancellation
22+ // between values should short-circuit the loop.
23+ const filter = "range(10000000)"
24+ stdin := strings.NewReader("null\n")
25+
26+ ctx, cancel := context.WithCancel(t.Context())
27+ // Cancel almost immediately so we catch the next iteration check.
28+ cancel()
29+
30+ err := handleJQ(ctx, []string{"jq", filter}, stdin, io.Discard, io.Discard)
31+ if err == nil {
32+ t.Fatal("expected ctx cancel error, got nil")
33+ }
34+ if !errors.Is(err, context.Canceled) {
35+ t.Fatalf("expected context.Canceled, got %v", err)
36+ }
37+}
38+
39+// TestJQ_CtxCancel_DuringFilter verifies cancellation mid-stream: ctx is
40+// cancelled after jq has started producing output, and the loop must
41+// observe the cancel on the next iteration rather than running to
42+// completion.
43+func TestJQ_CtxCancel_DuringFilter(t *testing.T) {
44+ t.Parallel()
45+
46+ ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
47+ defer cancel()
48+
49+ // 100M values; without ctx polling this would take many seconds to
50+ // fully emit. With ctx polling the loop exits shortly after the
51+ // deadline.
52+ stdin := strings.NewReader("null\n")
53+ var stdout, stderr bytes.Buffer
54+
55+ start := time.Now()
56+ err := handleJQ(ctx, []string{"jq", "-c", "range(100000000)"}, stdin, &stdout, &stderr)
57+ elapsed := time.Since(start)
58+
59+ if err == nil {
60+ t.Fatal("expected ctx timeout error, got nil")
61+ }
62+ if !errors.Is(err, context.DeadlineExceeded) {
63+ t.Fatalf("expected context.DeadlineExceeded, got %v", err)
64+ }
65+ // Allow generous slack for slow CI; the important invariant is that we
66+ // don't run all 100M iterations (which would take orders of magnitude
67+ // longer than 1s).
68+ if elapsed > time.Second {
69+ t.Fatalf("handleJQ took %v after 50ms timeout; ctx polling is not tight enough", elapsed)
70+ }
71+}
72+
73+// slowReader serves bytes in small chunks with a fixed delay between
74+// Read calls. It never blocks indefinitely — each Read returns after
75+// chunkDelay — so cancellation must be observed via ctxReader's ctx
76+// check, not by the underlying reader itself. That isolates the
77+// behavior we want to test: the wrapper polling ctx between chunks.
78+type slowReader struct {
79+ remaining []byte
80+ chunk int
81+ chunkDelay time.Duration
82+}
83+
84+func (s *slowReader) Read(p []byte) (int, error) {
85+ if len(s.remaining) == 0 {
86+ return 0, io.EOF
87+ }
88+ time.Sleep(s.chunkDelay)
89+ n := min(len(p), min(s.chunk, len(s.remaining)))
90+ copy(p, s.remaining[:n])
91+ s.remaining = s.remaining[n:]
92+ return n, nil
93+}
94+
95+// TestJQ_CtxCancel_MidReadAll verifies that ctx cancellation observed
96+// *during* io.ReadAll — after several chunks have already been consumed
97+// — short-circuits the read via ctxReader, rather than draining the
98+// whole source. This is the guarantee the hook runner relies on when
99+// it feeds a large bytes.Reader payload.
100+//
101+// The reader serves bytes in 512-byte chunks with a 5ms gap between
102+// reads. ctx is cancelled after ~50ms, so several chunks have already
103+// been read when ctxReader first observes the cancellation. The test
104+// asserts that (a) we got a context.Canceled error and (b) the call
105+// returned well before the reader would have been fully drained.
106+func TestJQ_CtxCancel_MidReadAll(t *testing.T) {
107+ t.Parallel()
108+
109+ const (
110+ size = 64 * 1024 * 1024 // 64 MiB
111+ chunk = 512
112+ chunkDelay = 5 * time.Millisecond
113+ )
114+ // At 512 bytes / 5ms, draining 64 MiB would take ~11 minutes. Any
115+ // return within a second proves cancel was observed mid-stream, not
116+ // after EOF.
117+ reader := &slowReader{
118+ remaining: bytes.Repeat([]byte("a"), size),
119+ chunk: chunk,
120+ chunkDelay: chunkDelay,
121+ }
122+
123+ ctx, cancel := context.WithCancel(t.Context())
124+ defer cancel()
125+
126+ // Cancel after enough time that several Read calls have completed
127+ // and io.ReadAll is actively consuming the source.
128+ go func() {
129+ time.Sleep(50 * time.Millisecond)
130+ cancel()
131+ }()
132+
133+ start := time.Now()
134+ err := handleJQ(ctx, []string{"jq", "-R", "."}, reader, io.Discard, io.Discard)
135+ elapsed := time.Since(start)
136+
137+ if !errors.Is(err, context.Canceled) {
138+ t.Fatalf("expected context.Canceled, got %v", err)
139+ }
140+ // Generous slack for slow CI; the invariant is orders-of-magnitude
141+ // faster than draining the full source.
142+ if elapsed > time.Second {
143+ t.Fatalf("mid-ReadAll cancel took %v; ctxReader is not polling between chunks", elapsed)
144+ }
145+ // Sanity check: we should have been cancelled mid-stream, not
146+ // before any reads happened. If remaining == size, cancel fired so
147+ // early nothing was consumed — that's a fast-fail path, not the
148+ // mid-read guarantee we want to verify.
149+ consumed := size - len(reader.remaining)
150+ if consumed == 0 {
151+ t.Fatal("reader was never read from; test did not exercise mid-ReadAll cancel")
152+ }
153+ if consumed >= size {
154+ t.Fatal("reader was fully drained; cancel was not observed mid-read")
155+ }
156+}
157+
158+// TestJQ_CtxCancel_PreCancel verifies the fast-fail path: a ctx already
159+// cancelled before handleJQ is called returns context.Canceled
160+// immediately via the outer-loop guard, never entering io.ReadAll.
161+// Complements TestJQ_CtxCancel_MidReadAll.
162+func TestJQ_CtxCancel_PreCancel(t *testing.T) {
163+ t.Parallel()
164+
165+ ctx, cancel := context.WithCancel(t.Context())
166+ cancel()
167+
168+ start := time.Now()
169+ err := handleJQ(ctx, []string{"jq", "-R", "."},
170+ bytes.NewReader(bytes.Repeat([]byte("a"), 1024)),
171+ io.Discard, io.Discard)
172+ elapsed := time.Since(start)
173+
174+ if !errors.Is(err, context.Canceled) {
175+ t.Fatalf("expected context.Canceled, got %v", err)
176+ }
177+ if elapsed > 100*time.Millisecond {
178+ t.Fatalf("pre-cancel fast-fail took %v; outer guard is not firing", elapsed)
179+ }
180+}
181+
182+// TestJQ_Success confirms the ctx-aware refactor did not regress the
183+// success path.
184+func TestJQ_Success(t *testing.T) {
185+ t.Parallel()
186+
187+ var stdout bytes.Buffer
188+ err := handleJQ(t.Context(),
189+ []string{"jq", "-c", ".a"},
190+ strings.NewReader(`{"a":1}`),
191+ &stdout, io.Discard,
192+ )
193+ if err != nil {
194+ t.Fatalf("handleJQ returned error: %v", err)
195+ }
196+ if got := stdout.String(); got != "1\n" {
197+ t.Fatalf("stdout = %q, want %q", got, "1\n")
198+ }
199+}
internal/shell/run.go+1−1
@@ -130,7 +130,7 @@ func builtinHandler() func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {
130 switch args[0] {
130 switch args[0] {
131 case "jq":
131 case "jq":
132 hc := interp.HandlerCtx(ctx)
132 hc := interp.HandlerCtx(ctx)
133− return handleJQ(args, hc.Stdin, hc.Stdout, hc.Stderr)
133+ return handleJQ(ctx, args, hc.Stdin, hc.Stdout, hc.Stderr)
134 default:
134 default:
135 return next(ctx, args)
135 return next(ctx, args)
136 }
136 }