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_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)
}
}