luxferre/bantam
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added C-c handling

Browse Source
This commit is contained in:
Luxferre
2026-08-15 17:09:26 +03:00
parent 4509851243
commit 3c55ce14d7
2 changed files with 216 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.go
+92 -27
View File
@@ -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 {
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
func shell(ctx context.Context, cmd string, timeout int) string {
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,16 +900,26 @@ 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 {
} else {
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
}
@@ -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")
return "", false
fmt.Print("^C\r\n")
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)
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 {
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)})
var err error
if msgs, err = AL(&cfg, msgs, sp, 0); err != nil {
turnMsgs := append([]Message{}, msgs...)
turnMsgs = append(turnMsgs, Message{Role: "user", Content: strp(u)})
sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
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
+118 -31
View File
@@ -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)
}
}