added C-c handling

main.go

@@ -15,6 +15,7 @@ import (
   "net/http"
   "os"
   "os/exec"
+  "os/signal"
   "path/filepath"
   "regexp"
   "sort"
@@ -628,7 +629,8 @@ func isInvalidAssistantErr(err error) bool {
   return strings.Contains(s, "Invalid assistant message") || strings.Contains(s, "content or tool_calls must be set")
 }
 
-func llm(cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
+func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
+  if err := ctx.Err(); err != nil { return Message{}, err }
   sanitizeMessages(msgs)
   p := map[string]any{"model": cfg.Model, "temperature": cfg.Temperature, "messages": msgs, "stream": cfg.Stream}
   if tools != nil { p["tools"] = tools }
@@ -650,13 +652,18 @@ func llm(cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
     DialContext:           (&net.Dialer{Timeout: time.Duration(cfg.Timeout) * time.Second}).DialContext,
     ResponseHeaderTimeout: time.Duration(cfg.Timeout) * time.Second,
   }}
+  defer client.CloseIdleConnections()
   fib := []int{1, 1, 2, 3, 5, 8, 13, 21, 34}
   pend := c("...requesting...", 1, 2)
   var resp *http.Response
   var err error
   for i := 0; i <= len(fib); i++ {
+    if err := ctx.Err(); err != nil {
+      if COL { fmt.Print("\r\033[K") }
+      return Message{}, err
+    }
     if COL { fmt.Print("\r" + pend) } else { fmt.Println(pend) }
-    req, _ := http.NewRequest("POST", strings.TrimRight(cfg.Endpoint, "/")+"/chat/completions", bytes.NewReader(body))
+    req, _ := http.NewRequestWithContext(ctx, "POST", strings.TrimRight(cfg.Endpoint, "/")+"/chat/completions", bytes.NewReader(body))
     req.Header.Set("Content-Type", "application/json")
     req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; Bantam/1.0)")
     if cfg.APIKey != "" && cfg.APIKey != "-" { req.Header.Set("Authorization", "Bearer "+cfg.APIKey) }
@@ -673,12 +680,19 @@ func llm(cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
     }
     if err == nil { break }
     if COL { fmt.Print("\r\033[K") }
+    if errors.Is(err, context.Canceled) || ctx.Err() != nil {
+      return Message{}, ctx.Err()
+    }
     if is4xxClientErr {
       return Message{}, err
     }
     if i < len(fib) {
       fmt.Println(c(fmt.Sprintf("[network error: %v, retrying in %ds...]", err, fib[i]), 31))
-      time.Sleep(time.Duration(fib[i]) * time.Second)
+      select {
+      case <-ctx.Done():
+        return Message{}, ctx.Err()
+      case <-time.After(time.Duration(fib[i]) * time.Second):
+      }
     }
   }
   if err != nil {
@@ -696,16 +710,19 @@ func llm(cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
         } `json:"message"`
       } `json:"choices"`
     }
-    if err := json.NewDecoder(resp.Body).Decode(&cr); err != nil { return Message{}, err }
+    if err := json.NewDecoder(resp.Body).Decode(&cr); err != nil {
+      if errors.Is(err, context.Canceled) || ctx.Err() != nil { return Message{}, ctx.Err() }
+      return Message{}, err
+    }
     if len(cr.Choices) == 0 { return Message{}, errors.New("empty choices in LLM response") }
     m := cr.Choices[0].Message.Message
     if m.ReasoningContent == "" { m.ReasoningContent = cr.Choices[0].Message.Reasoning }
     return m, nil
   }
-  return parseStream(resp.Body)
+  return parseStream(ctx, resp.Body)
 }
 
-func parseStream(r io.Reader) (Message, error) {
+func parseStream(ctx context.Context, r io.Reader) (Message, error) {
   var content, reas string
   var rh, ch bool
   var lineBuf string
@@ -724,6 +741,7 @@ func parseStream(r io.Reader) (Message, error) {
   sc := bufio.NewScanner(r)
   sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
   for sc.Scan() {
+    if err := ctx.Err(); err != nil { return Message{}, err }
     ln := strings.TrimSpace(sc.Text())
     if !strings.HasPrefix(ln, "data:") { continue }
     data := strings.TrimSpace(ln[5:])
@@ -774,6 +792,7 @@ func parseStream(r io.Reader) (Message, error) {
       if tc.Function.Arguments != "" { t.Function.Arguments += tc.Function.Arguments }
     }
   }
+  if err := ctx.Err(); err != nil { return Message{}, err }
   flushTable()
   if lineBuf != "" {
     if !mdSt.inCode && isTableLine(lineBuf) {
@@ -796,12 +815,17 @@ func parseStream(r io.Reader) (Message, error) {
   return m, sc.Err()
 }
 
-func shell(cmd string, timeout int) string {
+func shell(ctx context.Context, cmd string, timeout int) string {
-  ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
+  cmdCtx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
   defer cancel()
-  out, err := exec.CommandContext(ctx, "sh", "-c", cmd).CombinedOutput()
+  c := exec.CommandContext(cmdCtx, "sh", "-c", cmd)
+  c.WaitDelay = 100 * time.Millisecond
+  out, err := c.CombinedOutput()
   res := strings.TrimSpace(string(out))
-  if ctx.Err() == context.DeadlineExceeded {
+  if ctx.Err() != nil {
+    return "[interrupted]\n\nexit: -1"
+  }
+  if cmdCtx.Err() == context.DeadlineExceeded {
     return fmt.Sprintf("%s\n\n[shell timeout after %ds]\nexit: -1", res, timeout)
   }
   code := 0
@@ -824,11 +848,15 @@ func last(msgs []Message) string {
   return ""
 }
 
-func AL(cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, error) {
+func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, error) {
   done := false
   for i := 0; i < cfg.MaxALIterations && !done; i++ {
-    m, err := llm(cfg, msgs, TOOLS)
+    if err := ctx.Err(); err != nil { return msgs, err }
+    m, err := llm(ctx, cfg, msgs, TOOLS)
     if err != nil {
+      if errors.Is(err, context.Canceled) || ctx.Err() != nil {
+        return msgs, err
+      }
       if isInvalidAssistantErr(err) {
         stripped := false
         for j := len(msgs) - 1; j >= 0; j-- {
@@ -861,6 +889,7 @@ func AL(cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, error) {
     }
     if len(m.ToolCalls) == 0 { done = true; break }
     for _, tc := range m.ToolCalls {
+      if err := ctx.Err(); err != nil { return msgs, err }
       fn, astr := tc.Function.Name, tc.Function.Arguments
       fmt.Println(c(fmt.Sprintf("[tool call: %s(%s)]", fn, astr), 33))
       res, sty := "", 2
@@ -871,18 +900,28 @@ func AL(cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, error) {
         switch fn {
         case "shell_exec":
           cmd, _ := a["command"].(string)
-          res = shell(cmd, cfg.ShellTimeout)
+          res = shell(ctx, cmd, cfg.ShellTimeout)
+          if err := ctx.Err(); err != nil { return msgs, err }
         case "run_subagent":
           pr, _ := a["prompt"].(string)
           if depth >= MAX_DEPTH {
             res, sty = fmt.Sprintf("[subagent depth limit (%d) reached, child not spawned]", MAX_DEPTH), 31
-          } else if subr, err := AL(cfg, []Message{{Role: "system", Content: strp(sp + "\n\nImportant: this is a child agent")}, {Role: "user", Content: strp(pr)}}, sp, depth+1); err != nil {
-            res, sty = "[subagent error: "+err.Error()+"]", 31
           } else {
-            res, sty = last(subr), 2
+            subMsgs := []Message{
+              {Role: "system", Content: strp(sp + "\n\nImportant: this is a child agent")},
+              {Role: "user", Content: strp(pr)},
+            }
+            if subr, err := AL(ctx, cfg, subMsgs, sp, depth+1); err != nil {
+              if errors.Is(err, context.Canceled) || ctx.Err() != nil {
+                return msgs, err
+              }
+              res, sty = "[subagent error: "+err.Error()+"]", 31
+            } else {
+              res, sty = last(subr), 2
+            }
           }
         default:
-          res, sty = "Unknown tool: "+fn, 31
+          res, sty = "Unknown tool: " + fn, 31
         }
       }
       fmt.Println(c("[tool result: "+fn+"]", 32) + "\n" + c(res, sty) + "\n")
@@ -989,7 +1028,7 @@ func autosave(msgs []Message) {
   os.WriteFile(filepath.Join(sdir(), "autosave.json"), b, 0644)
 }
 
-func summarize(cfg *Cfg, msgs []Message) (string, error) {
+func summarize(ctx context.Context, cfg *Cfg, msgs []Message) (string, error) {
   var sb strings.Builder
   for _, m := range msgs {
     if m.Role == "system" { continue }
@@ -1013,7 +1052,7 @@ func summarize(cfg *Cfg, msgs []Message) (string, error) {
   sys := "You are a conversation summarizer for an AI agent's context window. Summarize concisely but completely, preserving all important facts, decisions, code, errors, and the current task state, so the agent can continue the work without the original messages. Output only the summary."
   cc := *cfg
   cc.Stream = false
-  m, err := llm(&cc, []Message{{Role: "system", Content: strp(sys)}, {Role: "user", Content: strp("Summarize this conversation:\n\n" + joined)}}, nil)
+  m, err := llm(ctx, &cc, []Message{{Role: "system", Content: strp(sys)}, {Role: "user", Content: strp("Summarize this conversation:\n\n" + joined)}}, nil)
   if err != nil { return "", err }
   s := ""
   if m.Content != nil { s = *m.Content }
@@ -1022,11 +1061,11 @@ func summarize(cfg *Cfg, msgs []Message) (string, error) {
   return strings.TrimSpace(s), nil
 }
 
-func compact(cfg *Cfg, msgs []Message) ([]Message, string, error) {
+func compact(ctx context.Context, cfg *Cfg, msgs []Message) ([]Message, string, error) {
   if len(msgs) == 0 || msgs[0].Role != "system" {
     return msgs, "", errors.New("session has no system message")
   }
-  s, err := summarize(cfg, msgs)
+  s, err := summarize(ctx, cfg, msgs)
   if err != nil { return msgs, "", err }
   return []Message{{Role: "system", Content: msgs[0].Content}, {Role: "user", Content: strp("Summary of the previous conversation:\n" + s + "\n\nPlease continue from here.")}}, s, nil
 }
@@ -1189,8 +1228,8 @@ func readLine(prompt string) (string, bool) {
       e.buf[e.pos] = '\n'
       e.pos++
     case 0x03:
-      fmt.Print("\r\n")
+      fmt.Print("^C\r\n")
-      return "", false
+      return "", true
     case 0x04:
       if len(e.buf) == 0 {
         fmt.Print("\r\n")
@@ -1247,14 +1286,23 @@ func main() {
       if cmd != "" {
         astr, _ := json.Marshal(map[string]string{"command": cmd})
         fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33))
-        res := shell(cmd, cfg.ShellTimeout)
+        sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
+        res := shell(sigCtx, cmd, cfg.ShellTimeout)
+        cancel()
         fmt.Println(c("[tool result: shell_exec]", 32) + "\n" + c(res, 2))
       }
       return
     }
     msgs = append(msgs, Message{Role: "user", Content: strp(u)})
-    if msgs, err = AL(&cfg, msgs, sp, 0); err != nil {
+    sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
-      fmt.Println(c("[error: "+err.Error()+"]", 31))
+    msgs, err = AL(sigCtx, &cfg, msgs, sp, 0)
+    cancel()
+    if err != nil {
+      if errors.Is(err, context.Canceled) || sigCtx.Err() != nil {
+        fmt.Println(c("\n[interrupted]", 33))
+      } else {
+        fmt.Println(c("[error: "+err.Error()+"]", 31))
+      }
       os.Exit(1)
     }
     autosave(msgs)
@@ -1317,9 +1365,15 @@ func main() {
         continue
       }
       fmt.Println(c("[compacting conversation...]", 33))
-      nm, sm, err := compact(&cfg, msgs)
+      sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
+      nm, sm, err := compact(sigCtx, &cfg, msgs)
+      cancel()
       if err != nil {
-        fmt.Println(c("[compact failed: "+err.Error()+"]", 31))
+        if errors.Is(err, context.Canceled) || sigCtx.Err() != nil {
+          fmt.Println(c("\n[interrupted]", 33))
+        } else {
+          fmt.Println(c("[compact failed: "+err.Error()+"]", 31))
+        }
         continue
       }
       msgs = nm
@@ -1354,7 +1408,9 @@ func main() {
       if cmd != "" {
         astr, _ := json.Marshal(map[string]string{"command": cmd})
         fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33))
-        res := shell(cmd, cfg.ShellTimeout)
+        sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
+        res := shell(sigCtx, cmd, cfg.ShellTimeout)
+        cancel()
         fmt.Println(c("[tool result: shell_exec]", 32) + "\n" + c(res, 2))
       }
       continue
@@ -1365,11 +1421,20 @@ func main() {
       }
       continue
     }
-    msgs = append(msgs, Message{Role: "user", Content: strp(u)})
+    turnMsgs := append([]Message{}, msgs...)
-    var err error
+    turnMsgs = append(turnMsgs, Message{Role: "user", Content: strp(u)})
-    if msgs, err = AL(&cfg, msgs, sp, 0); err != nil {
+    sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
-      fmt.Println(c("[error: "+err.Error()+"]", 31))
+    resMsgs, err := AL(sigCtx, &cfg, turnMsgs, sp, 0)
+    cancel()
+    if err != nil {
+      if errors.Is(err, context.Canceled) || sigCtx.Err() != nil {
+        fmt.Println(c("\n[interrupted]", 33))
+      } else {
+        fmt.Println(c("[error: "+err.Error()+"]", 31))
+      }
+      continue
     }
+    msgs = resMsgs
     autosave(msgs)
   }
 done:

main_test.go

@@ -3,6 +3,7 @@ package main
 import (
   "bufio"
   "bytes"
+  "context"
   "encoding/json"
   "errors"
   "net/http"
@@ -58,7 +59,7 @@ func TestParseStreamReasoningNoDuplication(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(bytes.NewBufferString(sseData))
+  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -90,7 +91,7 @@ func TestParseStreamReasoningAlias(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(bytes.NewBufferString(sseData))
+  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -108,7 +109,7 @@ func TestParseStreamContentOnly(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(bytes.NewBufferString(sseData))
+  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -129,7 +130,7 @@ func TestParseStreamReasoningOnly(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(bytes.NewBufferString(sseData))
+  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -143,7 +144,7 @@ func TestParseStreamReasoningOnly(t *testing.T) {
 
 func TestParseStreamEmpty(t *testing.T) {
   for _, in := range []string{"", "\n\n", "event: message\n\n"} {
-    msg, err := parseStream(strings.NewReader(in))
+    msg, err := parseStream(context.Background(), strings.NewReader(in))
     if err != nil {
       t.Fatalf("unexpected parseStream error for input %q: %v", in, err)
     }
@@ -161,7 +162,7 @@ func TestParseStreamToolCallSplitAcrossChunks(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(bytes.NewBufferString(sseData))
+  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -187,7 +188,7 @@ func TestParseStreamMultipleToolCallsKeepFirstAppearanceOrder(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(bytes.NewBufferString(sseData))
+  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -209,7 +210,7 @@ func TestParseStreamJunkAndNoChoicesIgnored(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(bytes.NewBufferString(sseData))
+  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -225,7 +226,7 @@ func TestParseStreamReasoningAfterContent(t *testing.T) {
     `data: [DONE]`,
   }, "\n")
 
-  msg, err := parseStream(bytes.NewBufferString(sseData))
+  msg, err := parseStream(context.Background(), bytes.NewBufferString(sseData))
   if err != nil {
     t.Fatalf("unexpected parseStream error: %v", err)
   }
@@ -475,33 +476,45 @@ func TestIsInvalidAssistantErr(t *testing.T) {
 // ---------- shell ----------
 
 func TestShellBasic(t *testing.T) {
-  res := shell("echo hi", 10)
+  res := shell(context.Background(), "echo hi", 10)
   if !strings.Contains(res, "hi") || !strings.HasSuffix(res, "exit: 0") {
     t.Errorf("shell(echo hi) = %q", res)
   }
 }
 
 func TestShellExitCodeAndStderr(t *testing.T) {
-  res := shell("echo out; echo err >&2; exit 7", 10)
+  res := shell(context.Background(), "echo out; echo err >&2; exit 7", 10)
   if !strings.Contains(res, "out") || !strings.Contains(res, "err") || !strings.HasSuffix(res, "exit: 7") {
     t.Errorf("shell multi = %q", res)
   }
 }
 
 func TestShellUnknownCommand(t *testing.T) {
-  res := shell("definitely_not_a_command_xyz", 10)
+  res := shell(context.Background(), "definitely_not_a_command_xyz", 10)
   if !strings.Contains(res, "exit: 127") {
     t.Errorf("expected exit 127, got %q", res)
   }
 }
 
 func TestShellTimeout(t *testing.T) {
-  res := shell("sleep 5", 1)
+  res := shell(context.Background(), "sleep 5", 1)
   if !strings.Contains(res, "[shell timeout after 1s]") || !strings.HasSuffix(res, "exit: -1") {
     t.Errorf("expected timeout marker and exit -1, got %q", res)
   }
 }
 
+func TestShellContextCancellation(t *testing.T) {
+  ctx, cancel := context.WithCancel(context.Background())
+  go func() {
+    time.Sleep(50 * time.Millisecond)
+    cancel()
+  }()
+  res := shell(ctx, "sleep 5", 10)
+  if !strings.Contains(res, "[interrupted]") || !strings.HasSuffix(res, "exit: -1") {
+    t.Errorf("expected interrupted result on cancellation, got %q", res)
+  }
+}
+
 // ---------- last / summary ----------
 
 func TestLast(t *testing.T) {
@@ -694,23 +707,23 @@ func TestAutosave(t *testing.T) {
 // ---------- summarize / compact (error paths only, no network) ----------
 
 func TestSummarizeEmpty(t *testing.T) {
-  if _, err := summarize(&Cfg{}, nil); err == nil || !strings.Contains(err.Error(), "no conversation") {
+  if _, err := summarize(context.Background(), &Cfg{}, nil); err == nil || !strings.Contains(err.Error(), "no conversation") {
     t.Errorf("expected no-conversation error, got %v", err)
   }
-  if _, err := summarize(&Cfg{}, []Message{{Role: "system", Content: strp("sys")}}); err == nil {
+  if _, err := summarize(context.Background(), &Cfg{}, []Message{{Role: "system", Content: strp("sys")}}); err == nil {
     t.Errorf("expected error for system-only conversation")
   }
 }
 
 func TestCompactNoSystem(t *testing.T) {
-  msgs, _, err := compact(&Cfg{}, nil)
+  msgs, _, err := compact(context.Background(), &Cfg{}, nil)
   if err == nil || !strings.Contains(err.Error(), "no system message") {
     t.Errorf("expected no-system error, got %v", err)
   }
   if msgs != nil {
     t.Errorf("expected original messages on error")
   }
-  msgs2, _, err2 := compact(&Cfg{}, []Message{{Role: "user", Content: strp("x")}})
+  msgs2, _, err2 := compact(context.Background(), &Cfg{}, []Message{{Role: "user", Content: strp("x")}})
   if err2 == nil {
     t.Errorf("expected error when first message is not system")
   }
@@ -930,7 +943,7 @@ func TestLLMNonStreamingAndHeaders(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "secret"
-  m, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, TOOLS)
+  m, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, TOOLS)
   if err != nil {
     t.Fatalf("llm: %v", err)
   }
@@ -960,7 +973,7 @@ func TestLLMNoAuthHeaderWhenNoKey(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  if _, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, nil); err != nil {
+  if _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil); err != nil {
     t.Fatalf("llm: %v", err)
   }
   if gotAuth != "" {
@@ -981,7 +994,7 @@ func TestLLMStreaming(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = true
   cfg.APIKey = "-"
-  m, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  m, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
   if err != nil {
     t.Fatalf("llm: %v", err)
   }
@@ -1005,7 +1018,7 @@ func TestLLM4xxReturnsImmediately(t *testing.T) {
   cfg.Stream = false
   cfg.APIKey = "-"
   start := time.Now()
-  _, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
   if err == nil || !strings.Contains(err.Error(), "Invalid assistant message") {
     t.Fatalf("expected 400 error, got %v", err)
   }
@@ -1031,7 +1044,7 @@ func TestLLMRetriesOn5xxThenSucceeds(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  m, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  m, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
   if err != nil {
     t.Fatalf("llm after retries: %v", err)
   }
@@ -1053,7 +1066,7 @@ func TestLLMEmptyChoices(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  if _, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, nil); err == nil || !strings.Contains(err.Error(), "empty choices") {
+  if _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil); err == nil || !strings.Contains(err.Error(), "empty choices") {
     t.Errorf("expected empty choices error, got %v", err)
   }
 }
@@ -1080,7 +1093,7 @@ func TestALToolLoop(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  msgs, err := AL(&cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("run")}}, "sys", 0)
+  msgs, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("run")}}, "sys", 0)
   if err != nil {
     t.Fatalf("AL: %v", err)
   }
@@ -1142,7 +1155,7 @@ func TestALRunSubagent(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  msgs, err := AL(&cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("parent task")}}, "sys", 0)
+  msgs, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("parent task")}}, "sys", 0)
   if err != nil {
     t.Fatalf("AL: %v", err)
   }
@@ -1185,7 +1198,7 @@ func TestALSubagentDepthLimit(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  msgs, err := AL(&cfg, []Message{{Role: "user", Content: strp("go")}}, "sys", MAX_DEPTH)
+  msgs, err := AL(context.Background(), &cfg, []Message{{Role: "user", Content: strp("go")}}, "sys", MAX_DEPTH)
   if err != nil {
     t.Fatalf("AL: %v", err)
   }
@@ -1232,7 +1245,7 @@ func TestALStripsInvalidAssistantAndRetries(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  msgs, err := AL(&cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("go")}}, "sys", 0)
+  msgs, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("go")}}, "sys", 0)
   if err != nil {
     t.Fatalf("AL: %v", err)
   }
@@ -1256,7 +1269,7 @@ func TestSummarizeHappyPath(t *testing.T) {
   cfg.Endpoint = srv.URL
   cfg.Stream = false
   cfg.APIKey = "-"
-  s, err := summarize(&cfg, []Message{{Role: "user", Content: strp("hello world")}})
+  s, err := summarize(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hello world")}})
   if err != nil {
     t.Fatalf("summarize: %v", err)
   }
@@ -1276,7 +1289,7 @@ func TestCompactHappyPath(t *testing.T) {
   cfg.Stream = false
   cfg.APIKey = "-"
   orig := []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("hello world")}}
-  msgs, sm, err := compact(&cfg, orig)
+  msgs, sm, err := compact(context.Background(), &cfg, orig)
   if err != nil {
     t.Fatalf("compact: %v", err)
   }
@@ -1339,7 +1352,7 @@ func TestLLMForwardsRelevantParameters(t *testing.T) {
   }, "\n"))
 
   cfg := getCfg(cfgFile)
-  _, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  _, err := llm(context.Background(), &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
   if err != nil {
     t.Fatalf("llm: %v", err)
   }
@@ -1591,7 +1604,7 @@ func TestRenderMDNoColorFallback(t *testing.T) {
 
 func TestDirectShellExecution(t *testing.T) {
   cmd := "echo direct_exec_test"
-  res := shell(cmd, 10)
+  res := shell(context.Background(), cmd, 10)
   if !strings.HasPrefix(res, "direct_exec_test") || !strings.Contains(res, "exit: 0") {
     t.Errorf("direct shell exec failed, got %q", res)
   }
@@ -1604,6 +1617,80 @@ func TestDirectShellExecution(t *testing.T) {
   }
 }
 
+func TestALContextCancellation(t *testing.T) {
+  done := make(chan struct{})
+  srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    select {
+    case <-r.Context().Done():
+    case <-done:
+    }
+  }))
+  defer func() {
+    close(done)
+    srv.CloseClientConnections()
+    srv.Close()
+  }()
+
+  cfg := defCfg
+  cfg.Endpoint = srv.URL
+  cfg.Stream = false
+  cfg.APIKey = "-"
+
+  ctx, cancel := context.WithCancel(context.Background())
+  go func() {
+    time.Sleep(30 * time.Millisecond)
+    cancel()
+  }()
+
+  origMsgs := []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("hello")}}
+  inputMsgs := append([]Message{}, origMsgs...)
+  msgs, err := AL(ctx, &cfg, inputMsgs, "sys", 0)
+  if err == nil {
+    t.Fatalf("expected context cancellation error, got nil")
+  }
+  if !errors.Is(err, context.Canceled) && ctx.Err() == nil {
+    t.Errorf("expected context.Canceled, got %v", err)
+  }
+  // Verify that input messages slice was not mutated with partial assistant messages
+  if len(msgs) != len(origMsgs) {
+    t.Errorf("expected %d messages after cancellation, got %d: %+v", len(origMsgs), len(msgs), msgs)
+  }
+}
+
+func TestLLMContextCancellation(t *testing.T) {
+  done := make(chan struct{})
+  srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    select {
+    case <-r.Context().Done():
+    case <-done:
+    }
+  }))
+  defer func() {
+    close(done)
+    srv.CloseClientConnections()
+    srv.Close()
+  }()
+
+  cfg := defCfg
+  cfg.Endpoint = srv.URL
+  cfg.Stream = false
+  cfg.APIKey = "-"
+
+  ctx, cancel := context.WithCancel(context.Background())
+  go func() {
+    time.Sleep(30 * time.Millisecond)
+    cancel()
+  }()
+
+  _, err := llm(ctx, &cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
+  if err == nil {
+    t.Fatalf("expected context cancellation error, got nil")
+  }
+  if !errors.Is(err, context.Canceled) && ctx.Err() == nil {
+    t.Errorf("expected context.Canceled, got %v", err)
+  }
+}
+