1273 lines
40 KiB
Go
1273 lines
40 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// ---------- helpers ----------
|
|
|
|
func testHome(t *testing.T) string {
|
|
t.Helper()
|
|
h := t.TempDir()
|
|
t.Setenv("HOME", h)
|
|
return h
|
|
}
|
|
|
|
func writeCfg(t *testing.T, content string) string {
|
|
t.Helper()
|
|
p := filepath.Join(t.TempDir(), "model.cfg")
|
|
if err := os.WriteFile(p, []byte(content), 0644); err != nil {
|
|
t.Fatalf("writeCfg: %v", err)
|
|
}
|
|
return p
|
|
}
|
|
|
|
func writeSession(t *testing.T, dir, id string, msgs []Message) {
|
|
t.Helper()
|
|
s := Session{ID: id, Created: "2026-01-01 00:00:00", Summary: "s", Messages: msgs}
|
|
b, err := json.Marshal(s)
|
|
if err != nil {
|
|
t.Fatalf("writeSession marshal: %v", err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(dir, id+".json"), b, 0644); err != nil {
|
|
t.Fatalf("writeSession: %v", err)
|
|
}
|
|
}
|
|
|
|
// ---------- parseStream ----------
|
|
|
|
func TestParseStreamReasoningNoDuplication(t *testing.T) {
|
|
// Simulate SSE stream where chunk 1 has reasoning_content, chunk 2 has reasoning_content,
|
|
// chunk 3 has content (and NO reasoning_content), chunk 4 has tool_calls (and NO reasoning_content)
|
|
sseData := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"reasoning_content":"Thinking step 1. "}}]}`,
|
|
`data: {"choices":[{"delta":{"reasoning_content":"Thinking step 2."}}]}`,
|
|
`data: {"choices":[{"delta":{"content":"Hello user"}}]}`,
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"shell_exec","arguments":"{\"command\":\"ls\"}"}}]}}]}`,
|
|
`data: [DONE]`,
|
|
}, "\n")
|
|
|
|
msg, err := parseStream(bytes.NewBufferString(sseData))
|
|
if err != nil {
|
|
t.Fatalf("unexpected parseStream error: %v", err)
|
|
}
|
|
|
|
expectedReasoning := "Thinking step 1. Thinking step 2."
|
|
if msg.ReasoningContent != expectedReasoning {
|
|
t.Errorf("expected ReasoningContent %q, got %q", expectedReasoning, msg.ReasoningContent)
|
|
}
|
|
|
|
expectedContent := "Hello user"
|
|
if msg.Content == nil || *msg.Content != expectedContent {
|
|
t.Errorf("expected Content %q, got %v", expectedContent, msg.Content)
|
|
}
|
|
|
|
if len(msg.ToolCalls) != 1 {
|
|
t.Fatalf("expected 1 tool call, got %d", len(msg.ToolCalls))
|
|
}
|
|
if msg.ToolCalls[0].Function.Name != "shell_exec" {
|
|
t.Errorf("expected tool call function name shell_exec, got %q", msg.ToolCalls[0].Function.Name)
|
|
}
|
|
}
|
|
|
|
func TestParseStreamReasoningAlias(t *testing.T) {
|
|
// Test reasoning field alias
|
|
sseData := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"reasoning":"Thought A. "}}]}`,
|
|
`data: {"choices":[{"delta":{"reasoning":"Thought B."}}]}`,
|
|
`data: {"choices":[{"delta":{"content":"Result"}}]}`,
|
|
`data: [DONE]`,
|
|
}, "\n")
|
|
|
|
msg, err := parseStream(bytes.NewBufferString(sseData))
|
|
if err != nil {
|
|
t.Fatalf("unexpected parseStream error: %v", err)
|
|
}
|
|
|
|
expectedReasoning := "Thought A. Thought B."
|
|
if msg.ReasoningContent != expectedReasoning {
|
|
t.Errorf("expected ReasoningContent %q, got %q", expectedReasoning, msg.ReasoningContent)
|
|
}
|
|
}
|
|
|
|
func TestParseStreamContentOnly(t *testing.T) {
|
|
sseData := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"content":"Hello"}}]}`,
|
|
`data: {"choices":[{"delta":{"content":" world"}}]}`,
|
|
`data: [DONE]`,
|
|
}, "\n")
|
|
|
|
msg, err := parseStream(bytes.NewBufferString(sseData))
|
|
if err != nil {
|
|
t.Fatalf("unexpected parseStream error: %v", err)
|
|
}
|
|
if msg.Content == nil || *msg.Content != "Hello world" {
|
|
t.Errorf("expected Content %q, got %v", "Hello world", msg.Content)
|
|
}
|
|
if msg.ReasoningContent != "" {
|
|
t.Errorf("expected empty ReasoningContent, got %q", msg.ReasoningContent)
|
|
}
|
|
if len(msg.ToolCalls) != 0 {
|
|
t.Errorf("expected no tool calls, got %d", len(msg.ToolCalls))
|
|
}
|
|
}
|
|
|
|
func TestParseStreamReasoningOnly(t *testing.T) {
|
|
sseData := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"reasoning_content":"Just thinking."}}]}`,
|
|
`data: [DONE]`,
|
|
}, "\n")
|
|
|
|
msg, err := parseStream(bytes.NewBufferString(sseData))
|
|
if err != nil {
|
|
t.Fatalf("unexpected parseStream error: %v", err)
|
|
}
|
|
if msg.ReasoningContent != "Just thinking." {
|
|
t.Errorf("expected ReasoningContent %q, got %q", "Just thinking.", msg.ReasoningContent)
|
|
}
|
|
if msg.Content != nil {
|
|
t.Errorf("expected nil Content, got %q", *msg.Content)
|
|
}
|
|
}
|
|
|
|
func TestParseStreamEmpty(t *testing.T) {
|
|
for _, in := range []string{"", "\n\n", "event: message\n\n"} {
|
|
msg, err := parseStream(strings.NewReader(in))
|
|
if err != nil {
|
|
t.Fatalf("unexpected parseStream error for input %q: %v", in, err)
|
|
}
|
|
if msg.Content != nil || msg.ReasoningContent != "" || len(msg.ToolCalls) != 0 {
|
|
t.Errorf("expected empty Message for input %q, got %+v", in, msg)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestParseStreamToolCallSplitAcrossChunks(t *testing.T) {
|
|
sseData := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"shell_","arguments":""}}]}}]}`,
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"exec","arguments":"{\"com"}}]}}]}`,
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"mand\":\"ls\"}"}}]}}]}`,
|
|
`data: [DONE]`,
|
|
}, "\n")
|
|
|
|
msg, err := parseStream(bytes.NewBufferString(sseData))
|
|
if err != nil {
|
|
t.Fatalf("unexpected parseStream error: %v", err)
|
|
}
|
|
if len(msg.ToolCalls) != 1 {
|
|
t.Fatalf("expected 1 tool call, got %d", len(msg.ToolCalls))
|
|
}
|
|
tc := msg.ToolCalls[0]
|
|
if tc.ID != "call_1" {
|
|
t.Errorf("expected id call_1, got %q", tc.ID)
|
|
}
|
|
if tc.Function.Name != "shell_exec" {
|
|
t.Errorf("expected name shell_exec, got %q", tc.Function.Name)
|
|
}
|
|
if tc.Function.Arguments != `{"command":"ls"}` {
|
|
t.Errorf("expected args %q, got %q", `{"command":"ls"}`, tc.Function.Arguments)
|
|
}
|
|
}
|
|
|
|
func TestParseStreamMultipleToolCallsKeepFirstAppearanceOrder(t *testing.T) {
|
|
sseData := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":1,"id":"c2","function":{"name":"run_subagent","arguments":"{\"prompt\":\"p\"}"}}]}}]}`,
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"shell_exec","arguments":"{\"command\":\"ls\"}"}}]}}]}`,
|
|
`data: [DONE]`,
|
|
}, "\n")
|
|
|
|
msg, err := parseStream(bytes.NewBufferString(sseData))
|
|
if err != nil {
|
|
t.Fatalf("unexpected parseStream error: %v", err)
|
|
}
|
|
if len(msg.ToolCalls) != 2 {
|
|
t.Fatalf("expected 2 tool calls, got %d", len(msg.ToolCalls))
|
|
}
|
|
// order follows first appearance: index 1 before index 0
|
|
if msg.ToolCalls[0].ID != "c2" || msg.ToolCalls[1].ID != "c1" {
|
|
t.Errorf("expected order [c2 c1], got [%s %s]", msg.ToolCalls[0].ID, msg.ToolCalls[1].ID)
|
|
}
|
|
}
|
|
|
|
func TestParseStreamJunkAndNoChoicesIgnored(t *testing.T) {
|
|
sseData := strings.Join([]string{
|
|
`event: message`,
|
|
`data: {"foo":"bar"}`,
|
|
`data: {"choices":[]}`,
|
|
`data: {"choices":[{"delta":{"content":"x"}}]}`,
|
|
`data: [DONE]`,
|
|
}, "\n")
|
|
|
|
msg, err := parseStream(bytes.NewBufferString(sseData))
|
|
if err != nil {
|
|
t.Fatalf("unexpected parseStream error: %v", err)
|
|
}
|
|
if msg.Content == nil || *msg.Content != "x" {
|
|
t.Errorf("expected Content %q, got %v", "x", msg.Content)
|
|
}
|
|
}
|
|
|
|
func TestParseStreamReasoningAfterContent(t *testing.T) {
|
|
sseData := strings.Join([]string{
|
|
`data: {"choices":[{"delta":{"content":"A"}}]}`,
|
|
`data: {"choices":[{"delta":{"reasoning":"B"}}]}`,
|
|
`data: [DONE]`,
|
|
}, "\n")
|
|
|
|
msg, err := parseStream(bytes.NewBufferString(sseData))
|
|
if err != nil {
|
|
t.Fatalf("unexpected parseStream error: %v", err)
|
|
}
|
|
if msg.Content == nil || *msg.Content != "A" {
|
|
t.Errorf("expected Content %q, got %v", "A", msg.Content)
|
|
}
|
|
if msg.ReasoningContent != "B" {
|
|
t.Errorf("expected ReasoningContent %q, got %q", "B", msg.ReasoningContent)
|
|
}
|
|
}
|
|
|
|
// ---------- prompt / getCfg / atoiD ----------
|
|
|
|
func TestDefaultSystemPromptFallback(t *testing.T) {
|
|
p := prompt("non_existent_file.txt")
|
|
if !strings.Contains(p, "You are Bantam, a tiny, powerful AI agent.") {
|
|
t.Errorf("expected prompt to contain base description, got: %q", p)
|
|
}
|
|
if !strings.Contains(p, "shell_exec") || !strings.Contains(p, "run_subagent") {
|
|
t.Errorf("expected prompt to list shell_exec and run_subagent tools, got: %q", p)
|
|
}
|
|
}
|
|
|
|
func TestPromptReadsAndTrimsFile(t *testing.T) {
|
|
p := filepath.Join(t.TempDir(), "sys.txt")
|
|
if err := os.WriteFile(p, []byte(" hello\nworld \n"), 0644); err != nil {
|
|
t.Fatalf("write: %v", err)
|
|
}
|
|
got := prompt(p)
|
|
if got != "hello\nworld" {
|
|
t.Errorf("expected trimmed content, got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestGetCfgDefaults(t *testing.T) {
|
|
t.Setenv("OPENAI_API_KEY", "")
|
|
cfg := getCfg(filepath.Join(t.TempDir(), "missing.cfg"))
|
|
if cfg.Endpoint != defCfg.Endpoint || cfg.Model != defCfg.Model || cfg.APIKey != defCfg.APIKey {
|
|
t.Errorf("defaults mismatch: %+v", cfg)
|
|
}
|
|
if cfg.Temperature != 0.7 || cfg.Timeout != 300 || cfg.ShellTimeout != 120 || cfg.MaxALIterations != 1000 {
|
|
t.Errorf("default numeric values mismatch: %+v", cfg)
|
|
}
|
|
if !cfg.Stream || cfg.Color != "auto" {
|
|
t.Errorf("default stream/color mismatch: %+v", cfg)
|
|
}
|
|
}
|
|
|
|
func TestGetCfgParsesFile(t *testing.T) {
|
|
p := writeCfg(t, strings.Join([]string{
|
|
"endpoint=http://localhost:9999/v1",
|
|
"model=test-model",
|
|
"temperature=0.5",
|
|
"api_key=secret",
|
|
"stream=false",
|
|
"color=never",
|
|
"timeout=42",
|
|
"shell_timeout=7",
|
|
"max_al_iterations=9",
|
|
"",
|
|
}, "\n"))
|
|
cfg := getCfg(p)
|
|
if cfg.Endpoint != "http://localhost:9999/v1" {
|
|
t.Errorf("endpoint: got %q", cfg.Endpoint)
|
|
}
|
|
if cfg.Model != "test-model" {
|
|
t.Errorf("model: got %q", cfg.Model)
|
|
}
|
|
if cfg.Temperature != 0.5 {
|
|
t.Errorf("temperature: got %v", cfg.Temperature)
|
|
}
|
|
if cfg.APIKey != "secret" {
|
|
t.Errorf("api_key: got %q", cfg.APIKey)
|
|
}
|
|
if cfg.Stream {
|
|
t.Errorf("stream: expected false")
|
|
}
|
|
if cfg.Color != "never" {
|
|
t.Errorf("color: got %q", cfg.Color)
|
|
}
|
|
if cfg.Timeout != 42 || cfg.ShellTimeout != 7 || cfg.MaxALIterations != 9 {
|
|
t.Errorf("timeouts: got %+v", cfg)
|
|
}
|
|
}
|
|
|
|
func TestGetCfgIgnoresCommentsBlankAndInvalid(t *testing.T) {
|
|
p := writeCfg(t, strings.Join([]string{
|
|
"# comment",
|
|
"",
|
|
"no-equals-line",
|
|
"temperature=abc",
|
|
"timeout=xyz",
|
|
"shell_timeout=",
|
|
"max_al_iterations=1.5",
|
|
"stream=banana",
|
|
"color=",
|
|
"",
|
|
}, "\n"))
|
|
cfg := getCfg(p)
|
|
if cfg.Temperature != 0.7 {
|
|
t.Errorf("invalid temperature should keep default, got %v", cfg.Temperature)
|
|
}
|
|
if cfg.Timeout != 300 || cfg.ShellTimeout != 120 || cfg.MaxALIterations != 1000 {
|
|
t.Errorf("invalid timeouts should keep defaults, got %+v", cfg)
|
|
}
|
|
if cfg.Stream {
|
|
t.Errorf("stream=banana should parse as false (matches Python/Perl)")
|
|
}
|
|
if cfg.Color != "" {
|
|
t.Errorf("empty color should stay empty (matches other ports), got %q", cfg.Color)
|
|
}
|
|
}
|
|
|
|
func TestGetCfgStreamTruthyVariants(t *testing.T) {
|
|
for _, tc := range []struct{ v, want string }{
|
|
{"true", "true"}, {"1", "true"}, {"yes", "true"},
|
|
{"false", "false"}, {"TRUE", "false"}, {"0", "false"},
|
|
} {
|
|
p := writeCfg(t, "stream="+tc.v+"\n")
|
|
got := getCfg(p).Stream
|
|
want := tc.want == "true"
|
|
if got != want {
|
|
t.Errorf("stream=%s: got %v, want %v", tc.v, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestGetCfgAPIKeyEnvFallback(t *testing.T) {
|
|
t.Setenv("OPENAI_API_KEY", "sk-env")
|
|
t.Setenv("HOME", t.TempDir())
|
|
|
|
// no api_key line at all -> env fallback (documented behavior)
|
|
if got := getCfg(writeCfg(t, "model=m\n")).APIKey; got != "sk-env" {
|
|
t.Errorf("no api_key + env: got %q, want sk-env", got)
|
|
}
|
|
// api_key=- -> env fallback
|
|
if got := getCfg(writeCfg(t, "api_key=-\n")).APIKey; got != "sk-env" {
|
|
t.Errorf("api_key=- + env: got %q, want sk-env", got)
|
|
}
|
|
// api_key= (empty) -> env fallback
|
|
if got := getCfg(writeCfg(t, "api_key=\n")).APIKey; got != "sk-env" {
|
|
t.Errorf("api_key= + env: got %q, want sk-env", got)
|
|
}
|
|
// explicit key wins over env
|
|
if got := getCfg(writeCfg(t, "api_key=real\n")).APIKey; got != "real" {
|
|
t.Errorf("explicit api_key: got %q, want real", got)
|
|
}
|
|
// no env and no key -> stays "-"
|
|
t.Setenv("OPENAI_API_KEY", "")
|
|
if got := getCfg(writeCfg(t, "api_key=-\n")).APIKey; got != "-" {
|
|
t.Errorf("api_key=- without env: got %q, want -", got)
|
|
}
|
|
if got := getCfg(writeCfg(t, "model=m\n")).APIKey; got != "-" {
|
|
t.Errorf("no api_key without env: got %q, want -", got)
|
|
}
|
|
}
|
|
|
|
func TestAtoiD(t *testing.T) {
|
|
cases := []struct {
|
|
s string
|
|
d, w int
|
|
}{
|
|
{"42", 0, 42}, {"-3", 0, -3}, {" 7 ", 0, 7},
|
|
{"abc", 5, 5}, {"", 5, 5}, {"1.5", 5, 5},
|
|
}
|
|
for _, c := range cases {
|
|
if got := atoiD(c.s, c.d); got != c.w {
|
|
t.Errorf("atoiD(%q, %d) = %d, want %d", c.s, c.d, got, c.w)
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------- sanitizeMessages / isInvalidAssistantErr ----------
|
|
|
|
func TestSanitizeMessages(t *testing.T) {
|
|
msgs := []Message{
|
|
{
|
|
Role: "assistant",
|
|
ToolCalls: []ToolCall{
|
|
{
|
|
ID: "tc1",
|
|
Type: "function",
|
|
Function: struct {
|
|
Name string `json:"name"`
|
|
Arguments string `json:"arguments"`
|
|
}{
|
|
Name: "shell_exec",
|
|
Arguments: `{"command": "ls"`, // invalid JSON: missing closing brace
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
sanitizeMessages(msgs)
|
|
if msgs[0].ToolCalls[0].Function.Arguments == `{"command": "ls"` {
|
|
t.Errorf("expected arguments to be sanitized to valid JSON, but remained raw")
|
|
}
|
|
if !strings.Contains(msgs[0].ToolCalls[0].Function.Arguments, "invalid_raw") {
|
|
t.Errorf("expected sanitized arguments to contain invalid_raw, got: %q", msgs[0].ToolCalls[0].Function.Arguments)
|
|
}
|
|
}
|
|
|
|
func TestSanitizeMessagesLeavesValidAndOtherRolesAlone(t *testing.T) {
|
|
validArgs := `{"command":"ls"}`
|
|
msgs := []Message{
|
|
{Role: "user", Content: strp("hi")},
|
|
{Role: "assistant", Content: strp("ok"), ToolCalls: []ToolCall{{ID: "a", Type: "function", Function: struct {
|
|
Name string `json:"name"`
|
|
Arguments string `json:"arguments"`
|
|
}{Name: "shell_exec", Arguments: validArgs}}}},
|
|
{Role: "assistant", Content: strp("no calls")},
|
|
{Role: "tool", ToolCallID: "a", Content: strp("out")},
|
|
}
|
|
sanitizeMessages(msgs)
|
|
if msgs[1].ToolCalls[0].Function.Arguments != validArgs {
|
|
t.Errorf("valid args were altered: %q", msgs[1].ToolCalls[0].Function.Arguments)
|
|
}
|
|
if msgs[0].Content == nil || *msgs[0].Content != "hi" {
|
|
t.Errorf("user message altered: %+v", msgs[0])
|
|
}
|
|
if len(msgs[2].ToolCalls) != 0 || msgs[3].Content == nil {
|
|
t.Errorf("unexpected alteration: %+v", msgs)
|
|
}
|
|
}
|
|
|
|
func TestIsInvalidAssistantErr(t *testing.T) {
|
|
cases := []struct {
|
|
msg string
|
|
want bool
|
|
}{
|
|
{"HTTP 400: [invalid_request_error] Invalid assistant message: content or tool_calls must be set", true},
|
|
{"HTTP 400: [invalid_request_error] Invalid assistant message: content or tool_calls must be set (HTTP 400)", true},
|
|
{"HTTP 400: [invalid_request_error] Invalid assistant message", true},
|
|
{"HTTP 400: content or tool_calls must be set", true},
|
|
{"HTTP 500: internal server error", false},
|
|
{"HTTP 429: rate limited", false},
|
|
{"network error: connection refused", false},
|
|
}
|
|
for _, c := range cases {
|
|
if got := isInvalidAssistantErr(errors.New(c.msg)); got != c.want {
|
|
t.Errorf("isInvalidAssistantErr(%q) = %v, want %v", c.msg, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------- shell ----------
|
|
|
|
func TestShellBasic(t *testing.T) {
|
|
res := shell("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)
|
|
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)
|
|
if !strings.Contains(res, "exit: 127") {
|
|
t.Errorf("expected exit 127, got %q", res)
|
|
}
|
|
}
|
|
|
|
func TestShellTimeout(t *testing.T) {
|
|
res := shell("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)
|
|
}
|
|
}
|
|
|
|
// ---------- last / summary ----------
|
|
|
|
func TestLast(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
msgs []Message
|
|
want string
|
|
}{
|
|
{"empty", nil, ""},
|
|
{"no assistant", []Message{{Role: "user", Content: strp("u")}}, ""},
|
|
{"last assistant wins", []Message{
|
|
{Role: "assistant", Content: strp("first")},
|
|
{Role: "tool", ToolCallID: "x", Content: strp("r")},
|
|
{Role: "assistant", Content: strp("second")},
|
|
}, "second"},
|
|
{"nil and empty content skipped", []Message{
|
|
{Role: "assistant"},
|
|
{Role: "assistant", Content: strp("")},
|
|
{Role: "assistant", Content: strp("real")},
|
|
}, "real"},
|
|
}
|
|
for _, c := range cases {
|
|
if got := last(c.msgs); got != c.want {
|
|
t.Errorf("%s: last() = %q, want %q", c.name, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestSummary(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
msgs []Message
|
|
want string
|
|
}{
|
|
{"empty", nil, "(empty session)"},
|
|
{"system only", []Message{{Role: "system", Content: strp("sys")}}, "(empty session)"},
|
|
{"whitespace user skipped", []Message{{Role: "user", Content: strp(" ")}}, "(empty session)"},
|
|
{"first user wins", []Message{
|
|
{Role: "user", Content: strp("hello world")},
|
|
{Role: "user", Content: strp("second")},
|
|
}, "hello world"},
|
|
{"long truncated", []Message{{Role: "user", Content: strp(strings.Repeat("a", 100))}}, strings.Repeat("a", 80) + "..."},
|
|
}
|
|
for _, c := range cases {
|
|
if got := summary(c.msgs); got != c.want {
|
|
t.Errorf("%s: summary() = %q, want %q", c.name, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------- homeDir / sdir / fileExists ----------
|
|
|
|
func TestHomeDir(t *testing.T) {
|
|
t.Setenv("HOME", "/tmp/bantam-test-home")
|
|
if got := homeDir(); got != "/tmp/bantam-test-home" {
|
|
t.Errorf("homeDir() = %q", got)
|
|
}
|
|
t.Setenv("HOME", "")
|
|
if got := homeDir(); got != "." {
|
|
t.Errorf("homeDir() with empty HOME = %q, want .", got)
|
|
}
|
|
}
|
|
|
|
func TestSdirCreatesDir(t *testing.T) {
|
|
h := testHome(t)
|
|
d := sdir()
|
|
want := filepath.Join(h, ".bantam", "sessions")
|
|
if d != want {
|
|
t.Errorf("sdir() = %q, want %q", d, want)
|
|
}
|
|
if fi, err := os.Stat(d); err != nil || !fi.IsDir() {
|
|
t.Errorf("sdir() did not create directory: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestFileExists(t *testing.T) {
|
|
p := filepath.Join(t.TempDir(), "f")
|
|
if fileExists(p) {
|
|
t.Errorf("fileExists(%q) = true before creation", p)
|
|
}
|
|
if err := os.WriteFile(p, []byte("x"), 0644); err != nil {
|
|
t.Fatalf("write: %v", err)
|
|
}
|
|
if !fileExists(p) {
|
|
t.Errorf("fileExists(%q) = false after creation", p)
|
|
}
|
|
}
|
|
|
|
// ---------- sessions ----------
|
|
|
|
func TestSaveSessionAndLoad(t *testing.T) {
|
|
h := testHome(t)
|
|
msgs := []Message{
|
|
{Role: "system", Content: strp("sys")},
|
|
{Role: "user", Content: strp("hello")},
|
|
}
|
|
sid, sm := saveSession(msgs)
|
|
if sm != "hello" {
|
|
t.Errorf("summary = %q, want hello", sm)
|
|
}
|
|
if sid == "" {
|
|
t.Fatalf("empty session id")
|
|
}
|
|
if !fileExists(filepath.Join(h, ".bantam", "sessions", sid+".json")) {
|
|
t.Errorf("session file not written")
|
|
}
|
|
loaded, err := loadSession(sid)
|
|
if err != nil {
|
|
t.Fatalf("loadSession: %v", err)
|
|
}
|
|
if len(loaded) != 2 || loaded[1].Role != "user" || *loaded[1].Content != "hello" {
|
|
t.Errorf("loaded messages mismatch: %+v", loaded)
|
|
}
|
|
}
|
|
|
|
func TestSaveSessionCollisionSuffix(t *testing.T) {
|
|
testHome(t)
|
|
base := time.Now().Format("20060102-150405")
|
|
if err := os.WriteFile(filepath.Join(sdir(), base+".json"), []byte("{}"), 0644); err != nil {
|
|
t.Fatalf("write: %v", err)
|
|
}
|
|
sid, _ := saveSession([]Message{{Role: "user", Content: strp("x")}})
|
|
if sid != base+"-1" {
|
|
t.Errorf("expected collision suffix %q, got %q", base+"-1", sid)
|
|
}
|
|
}
|
|
|
|
func TestSessionsSortAndFilter(t *testing.T) {
|
|
h := testHome(t)
|
|
d := filepath.Join(h, ".bantam", "sessions")
|
|
if err := os.MkdirAll(d, 0755); err != nil {
|
|
t.Fatalf("mkdir: %v", err)
|
|
}
|
|
writeSession(t, d, "b", []Message{{Role: "user", Content: strp("u")}})
|
|
writeSession(t, d, "a", []Message{{Role: "user", Content: strp("u")}})
|
|
os.WriteFile(filepath.Join(d, "junk.txt"), []byte("nope"), 0644)
|
|
os.WriteFile(filepath.Join(d, "corrupt.json"), []byte("not json"), 0644)
|
|
os.Mkdir(filepath.Join(d, "subdir"), 0755)
|
|
|
|
ss := sessions()
|
|
if len(ss) != 2 {
|
|
t.Fatalf("expected 2 sessions, got %d", len(ss))
|
|
}
|
|
if ss[0].ID != "b" || ss[1].ID != "a" {
|
|
t.Errorf("expected descending order [b a], got [%s %s]", ss[0].ID, ss[1].ID)
|
|
}
|
|
}
|
|
|
|
func TestLoadSessionExactPrefixAmbiguousNotFound(t *testing.T) {
|
|
h := testHome(t)
|
|
d := filepath.Join(h, ".bantam", "sessions")
|
|
if err := os.MkdirAll(d, 0755); err != nil {
|
|
t.Fatalf("mkdir: %v", err)
|
|
}
|
|
writeSession(t, d, "aaa", []Message{{Role: "user", Content: strp("one")}})
|
|
writeSession(t, d, "aab", []Message{{Role: "user", Content: strp("two")}})
|
|
writeSession(t, d, "zzz", []Message{{Role: "user", Content: strp("three")}})
|
|
|
|
if _, err := loadSession("aaa"); err != nil {
|
|
t.Errorf("exact match failed: %v", err)
|
|
}
|
|
if _, err := loadSession("zz"); err != nil {
|
|
t.Errorf("unique prefix failed: %v", err)
|
|
}
|
|
if _, err := loadSession("aa"); err == nil || !strings.Contains(err.Error(), "ambiguous") {
|
|
t.Errorf("expected ambiguous error, got %v", err)
|
|
}
|
|
if _, err := loadSession("qq"); err == nil || !strings.Contains(err.Error(), "not found") {
|
|
t.Errorf("expected not found error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestAutosave(t *testing.T) {
|
|
h := testHome(t)
|
|
msgs := []Message{{Role: "user", Content: strp("turn")}}
|
|
autosave(msgs)
|
|
p := filepath.Join(h, ".bantam", "sessions", "autosave.json")
|
|
if !fileExists(p) {
|
|
t.Fatalf("autosave.json not written")
|
|
}
|
|
loaded, err := loadSession("autosave")
|
|
if err != nil {
|
|
t.Fatalf("loadSession(autosave): %v", err)
|
|
}
|
|
if len(loaded) != 1 || *loaded[0].Content != "turn" {
|
|
t.Errorf("autosave messages mismatch: %+v", loaded)
|
|
}
|
|
}
|
|
|
|
// ---------- 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") {
|
|
t.Errorf("expected no-conversation error, got %v", err)
|
|
}
|
|
if _, err := summarize(&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)
|
|
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")}})
|
|
if err2 == nil {
|
|
t.Errorf("expected error when first message is not system")
|
|
}
|
|
if len(msgs2) != 1 {
|
|
t.Errorf("expected original messages returned, got %d", len(msgs2))
|
|
}
|
|
}
|
|
|
|
// ---------- history ----------
|
|
|
|
func TestHistoryRoundTrip(t *testing.T) {
|
|
oldHist, oldHistF := hist, histF
|
|
t.Cleanup(func() { hist, histF = oldHist, oldHistF })
|
|
hist, histF = nil, ""
|
|
|
|
h := testHome(t)
|
|
loadHistory()
|
|
if histF != filepath.Join(h, ".bantam_history") {
|
|
t.Errorf("histF = %q", histF)
|
|
}
|
|
if len(hist) != 0 {
|
|
t.Errorf("expected empty history, got %v", hist)
|
|
}
|
|
addHistory("one")
|
|
addHistory("one") // duplicate ignored
|
|
addHistory("two")
|
|
addHistory("") // empty ignored
|
|
if len(hist) != 2 || hist[0] != "one" || hist[1] != "two" {
|
|
t.Errorf("hist = %v", hist)
|
|
}
|
|
saveHistory()
|
|
|
|
hist = nil
|
|
loadHistory()
|
|
if len(hist) != 2 || hist[0] != "one" || hist[1] != "two" {
|
|
t.Errorf("history not reloaded: %v", hist)
|
|
}
|
|
}
|
|
|
|
// ---------- visibleLen / textPos ----------
|
|
|
|
func TestVisibleLen(t *testing.T) {
|
|
cases := []struct {
|
|
s string
|
|
want int
|
|
}{
|
|
{"", 0},
|
|
{"hello", 5},
|
|
{"héllo", 5},
|
|
{"🙂x", 2},
|
|
{"\033[31mred\033[0m", 3},
|
|
{"a\001\033[1m\002b\001\033[0m\002c", 3},
|
|
}
|
|
for _, c := range cases {
|
|
if got := visibleLen(c.s); got != c.want {
|
|
t.Errorf("visibleLen(%q) = %d, want %d", c.s, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestTextPos(t *testing.T) {
|
|
cases := []struct {
|
|
pl, W int
|
|
s string
|
|
pos int
|
|
r, col int
|
|
}{
|
|
{0, 80, "hello", 3, 0, 3},
|
|
{5, 80, "hello", 0, 0, 5},
|
|
{0, 80, "ab\ncd", 4, 1, 1},
|
|
{0, 80, "ab\ncd", 5, 1, 2},
|
|
{0, 5, "abcde", 5, 0, 4}, // 5th char sits at last column, next would wrap
|
|
{0, 5, "abcdef", 6, 1, 1}, // wrap to next row
|
|
{0, 80, "", 0, 0, 0},
|
|
}
|
|
for _, c := range cases {
|
|
r, col := textPos(c.pl, c.W, c.s, c.pos)
|
|
if r != c.r || col != c.col {
|
|
t.Errorf("textPos(%d,%d,%q,%d) = (%d,%d), want (%d,%d)", c.pl, c.W, c.s, c.pos, r, col, c.r, c.col)
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------- editor.histNav ----------
|
|
|
|
func TestHistNav(t *testing.T) {
|
|
e := &editor{hpos: -1, hist: []string{"first", "second"}, buf: []rune("draft"), pos: 5}
|
|
e.histNav(true) // up: newest entry
|
|
if e.hpos != 1 || string(e.buf) != "second" || e.draft != "draft" || e.pos != len(e.buf) {
|
|
t.Errorf("after first up: hpos=%d buf=%q draft=%q pos=%d", e.hpos, string(e.buf), e.draft, e.pos)
|
|
}
|
|
e.histNav(true) // up again
|
|
if e.hpos != 0 || string(e.buf) != "first" {
|
|
t.Errorf("after second up: hpos=%d buf=%q", e.hpos, string(e.buf))
|
|
}
|
|
e.histNav(true) // at oldest, stays
|
|
if e.hpos != 0 || string(e.buf) != "first" {
|
|
t.Errorf("after third up: hpos=%d buf=%q", e.hpos, string(e.buf))
|
|
}
|
|
e.histNav(false) // down
|
|
if e.hpos != 1 || string(e.buf) != "second" {
|
|
t.Errorf("after down: hpos=%d buf=%q", e.hpos, string(e.buf))
|
|
}
|
|
e.histNav(false) // down past end -> restore draft
|
|
if e.hpos != -1 || string(e.buf) != "draft" {
|
|
t.Errorf("after down to draft: hpos=%d buf=%q", e.hpos, string(e.buf))
|
|
}
|
|
e.histNav(false) // no-op when not navigating
|
|
if e.hpos != -1 || string(e.buf) != "draft" {
|
|
t.Errorf("after extra down: hpos=%d buf=%q", e.hpos, string(e.buf))
|
|
}
|
|
|
|
empty := &editor{hpos: -1}
|
|
empty.histNav(true)
|
|
if empty.hpos != -1 {
|
|
t.Errorf("histNav with empty history changed hpos to %d", empty.hpos)
|
|
}
|
|
}
|
|
|
|
// ---------- readPlain / readLine ----------
|
|
|
|
func TestReadPlain(t *testing.T) {
|
|
oldStdin := stdin
|
|
t.Cleanup(func() { stdin = oldStdin })
|
|
|
|
stdin = bufio.NewReader(strings.NewReader("hello\n"))
|
|
got, ok := readPlain("> ")
|
|
if !ok || got != "hello" {
|
|
t.Errorf("readPlain = (%q, %v), want (hello, true)", got, ok)
|
|
}
|
|
|
|
stdin = bufio.NewReader(strings.NewReader("no-newline"))
|
|
got, ok = readPlain("> ")
|
|
if !ok || got != "no-newline" {
|
|
t.Errorf("readPlain no-newline = (%q, %v)", got, ok)
|
|
}
|
|
|
|
stdin = bufio.NewReader(strings.NewReader(""))
|
|
got, ok = readPlain("> ")
|
|
if ok || got != "" {
|
|
t.Errorf("readPlain EOF = (%q, %v), want (\"\", false)", got, ok)
|
|
}
|
|
}
|
|
|
|
func TestReadLineFallsBackToPlainWhenNotTTY(t *testing.T) {
|
|
if isTerminal(int(os.Stdin.Fd())) {
|
|
t.Skip("stdin is a terminal; readLine would enter raw mode")
|
|
}
|
|
oldStdin := stdin
|
|
t.Cleanup(func() { stdin = oldStdin })
|
|
|
|
stdin = bufio.NewReader(strings.NewReader("line\n"))
|
|
got, ok := readLine("> ")
|
|
if !ok || got != "line" {
|
|
t.Errorf("readLine = (%q, %v), want (line, true)", got, ok)
|
|
}
|
|
}
|
|
|
|
// ---------- colors ----------
|
|
|
|
func TestColorHelperC(t *testing.T) {
|
|
old := COL
|
|
t.Cleanup(func() { COL = old })
|
|
|
|
COL = false
|
|
if got := c("x", 31); got != "x" {
|
|
t.Errorf("c with COL=false = %q", got)
|
|
}
|
|
COL = true
|
|
if got := c("x", 31); got != "\033[31mx\033[0m" {
|
|
t.Errorf("c(x,31) = %q", got)
|
|
}
|
|
if got := c("x", 1, 32); got != "\033[1;32mx\033[0m" {
|
|
t.Errorf("c(x,1,32) = %q", got)
|
|
}
|
|
if got := c("x"); got != "x" {
|
|
t.Errorf("c(x) with no codes = %q", got)
|
|
}
|
|
}
|
|
|
|
func TestCol(t *testing.T) {
|
|
t.Setenv("NO_COLOR", "")
|
|
t.Setenv("BANTAM_NO_COLOR", "")
|
|
|
|
if !col(Cfg{Color: "always"}) {
|
|
t.Errorf("color=always should be true")
|
|
}
|
|
if col(Cfg{Color: "never"}) {
|
|
t.Errorf("color=never should be false")
|
|
}
|
|
if col(Cfg{Color: "auto"}) {
|
|
t.Errorf("color=auto should be false when stdout is not a TTY")
|
|
}
|
|
if col(Cfg{Color: "garbage"}) {
|
|
t.Errorf("unknown color value should fall back to TTY detection")
|
|
}
|
|
|
|
t.Setenv("NO_COLOR", "1")
|
|
if col(Cfg{Color: "always"}) {
|
|
t.Errorf("NO_COLOR should override color=always")
|
|
}
|
|
}
|
|
|
|
// ---------- llm / AL / summarize / compact via httptest (no real network) ----------
|
|
|
|
func TestLLMNonStreamingAndHeaders(t *testing.T) {
|
|
var gotPath, gotAuth, gotUA string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
gotAuth = r.Header.Get("Authorization")
|
|
gotUA = r.Header.Get("User-Agent")
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"hi","reasoning_content":"think"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
cfg.Endpoint = srv.URL
|
|
cfg.Stream = false
|
|
cfg.APIKey = "secret"
|
|
m, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, TOOLS)
|
|
if err != nil {
|
|
t.Fatalf("llm: %v", err)
|
|
}
|
|
if gotPath != "/chat/completions" {
|
|
t.Errorf("path = %q", gotPath)
|
|
}
|
|
if gotAuth != "Bearer secret" {
|
|
t.Errorf("auth = %q", gotAuth)
|
|
}
|
|
if !strings.Contains(gotUA, "Bantam/1.0") {
|
|
t.Errorf("user-agent = %q", gotUA)
|
|
}
|
|
if m.Content == nil || *m.Content != "hi" || m.ReasoningContent != "think" {
|
|
t.Errorf("message = %+v", m)
|
|
}
|
|
}
|
|
|
|
func TestLLMNoAuthHeaderWhenNoKey(t *testing.T) {
|
|
var gotAuth string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotAuth = r.Header.Get("Authorization")
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"x"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
cfg.Endpoint = srv.URL
|
|
cfg.Stream = false
|
|
cfg.APIKey = "-"
|
|
if _, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, nil); err != nil {
|
|
t.Fatalf("llm: %v", err)
|
|
}
|
|
if gotAuth != "" {
|
|
t.Errorf("expected no Authorization header, got %q", gotAuth)
|
|
}
|
|
}
|
|
|
|
func TestLLMStreaming(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"R\"}}]}\n\n"))
|
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"C\"}}]}\n\n"))
|
|
w.Write([]byte("data: [DONE]\n\n"))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
cfg.Endpoint = srv.URL
|
|
cfg.Stream = true
|
|
cfg.APIKey = "-"
|
|
m, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
|
|
if err != nil {
|
|
t.Fatalf("llm: %v", err)
|
|
}
|
|
if m.Content == nil || *m.Content != "C" {
|
|
t.Errorf("content = %v", m.Content)
|
|
}
|
|
if m.ReasoningContent != "R" {
|
|
t.Errorf("reasoning = %q", m.ReasoningContent)
|
|
}
|
|
}
|
|
|
|
func TestLLM4xxReturnsImmediately(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(400)
|
|
w.Write([]byte(`{"error":"Invalid assistant message: content or tool_calls must be set"}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
cfg.Endpoint = srv.URL
|
|
cfg.Stream = false
|
|
cfg.APIKey = "-"
|
|
start := time.Now()
|
|
_, err := llm(&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)
|
|
}
|
|
if time.Since(start) > time.Second {
|
|
t.Errorf("4xx should not be retried, took %v", time.Since(start))
|
|
}
|
|
}
|
|
|
|
func TestLLMRetriesOn5xxThenSucceeds(t *testing.T) {
|
|
var calls int
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
calls++
|
|
if calls < 3 {
|
|
w.WriteHeader(500)
|
|
w.Write([]byte("boom"))
|
|
return
|
|
}
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"ok"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
cfg.Endpoint = srv.URL
|
|
cfg.Stream = false
|
|
cfg.APIKey = "-"
|
|
m, err := llm(&cfg, []Message{{Role: "user", Content: strp("hi")}}, nil)
|
|
if err != nil {
|
|
t.Fatalf("llm after retries: %v", err)
|
|
}
|
|
if calls != 3 {
|
|
t.Errorf("expected 3 calls, got %d", calls)
|
|
}
|
|
if m.Content == nil || *m.Content != "ok" {
|
|
t.Errorf("content = %v", m.Content)
|
|
}
|
|
}
|
|
|
|
func TestLLMEmptyChoices(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Write([]byte(`{"choices":[]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
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") {
|
|
t.Errorf("expected empty choices error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestALToolLoop(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var req struct {
|
|
Messages []Message `json:"messages"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
t.Errorf("decode request: %v", err)
|
|
}
|
|
for _, m := range req.Messages {
|
|
if m.Role == "tool" {
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"done"}}]}`))
|
|
return
|
|
}
|
|
}
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"shell_exec","arguments":"{\"command\":\"echo hello\"}"}}]}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
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)
|
|
if err != nil {
|
|
t.Fatalf("AL: %v", err)
|
|
}
|
|
if got := last(msgs); got != "done" {
|
|
t.Errorf("last = %q, want done", got)
|
|
}
|
|
var toolMsgs int
|
|
for _, m := range msgs {
|
|
if m.Role == "tool" {
|
|
toolMsgs++
|
|
if m.ToolCallID != "c1" {
|
|
t.Errorf("tool msg tool_call_id = %q", m.ToolCallID)
|
|
}
|
|
if m.Content == nil || !strings.Contains(*m.Content, "hello") || !strings.Contains(*m.Content, "exit: 0") {
|
|
t.Errorf("tool result = %v", m.Content)
|
|
}
|
|
}
|
|
}
|
|
if toolMsgs != 1 {
|
|
t.Errorf("expected 1 tool message, got %d", toolMsgs)
|
|
}
|
|
}
|
|
|
|
func TestALRunSubagent(t *testing.T) {
|
|
var n int
|
|
var mu sync.Mutex
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
n++
|
|
cur := n
|
|
mu.Unlock()
|
|
var req struct {
|
|
Messages []Message `json:"messages"`
|
|
}
|
|
json.NewDecoder(r.Body).Decode(&req)
|
|
switch cur {
|
|
case 1: // parent's first call -> delegate to subagent
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"run_subagent","arguments":"{\"prompt\":\"inner task\"}"}}]}}]}`))
|
|
case 2: // child's call -> verify it got the child system prompt, then finish
|
|
isChild := false
|
|
for _, m := range req.Messages {
|
|
if m.Role == "system" && m.Content != nil && strings.Contains(*m.Content, "Important: this is a child agent") {
|
|
isChild = true
|
|
}
|
|
}
|
|
if !isChild {
|
|
t.Errorf("request #2 did not carry the child system prompt")
|
|
}
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"child done"}}]}`))
|
|
case 3: // parent's second call -> finish
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"parent done"}}]}`))
|
|
default:
|
|
t.Errorf("unexpected request #%d", cur)
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
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)
|
|
if err != nil {
|
|
t.Fatalf("AL: %v", err)
|
|
}
|
|
if got := last(msgs); got != "parent done" {
|
|
t.Errorf("last = %q, want parent done", got)
|
|
}
|
|
found := false
|
|
for _, m := range msgs {
|
|
if m.Role == "tool" && m.Content != nil && strings.Contains(*m.Content, "child done") {
|
|
found = true
|
|
}
|
|
}
|
|
if !found {
|
|
t.Errorf("expected subagent result 'child done' in tool messages")
|
|
}
|
|
}
|
|
|
|
func TestALSubagentDepthLimit(t *testing.T) {
|
|
// depth >= MAX_DEPTH must not spawn an HTTP request at all
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"run_subagent","arguments":"{\"prompt\":\"deep\"}"}}]}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
cfg.Endpoint = srv.URL
|
|
cfg.Stream = false
|
|
cfg.APIKey = "-"
|
|
msgs, err := AL(&cfg, []Message{{Role: "user", Content: strp("go")}}, "sys", MAX_DEPTH)
|
|
if err != nil {
|
|
t.Fatalf("AL: %v", err)
|
|
}
|
|
found := false
|
|
for _, m := range msgs {
|
|
if m.Role == "tool" && m.Content != nil && strings.Contains(*m.Content, "depth limit") {
|
|
found = true
|
|
}
|
|
}
|
|
if !found {
|
|
t.Errorf("expected depth-limit tool result")
|
|
}
|
|
}
|
|
|
|
func TestALStripsInvalidAssistantAndRetries(t *testing.T) {
|
|
var n int
|
|
var mu sync.Mutex
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
n++
|
|
cur := n
|
|
mu.Unlock()
|
|
switch cur {
|
|
case 1:
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"shell_exec","arguments":"{\"command\":\"echo x\"}"}}]}}]}`))
|
|
case 2:
|
|
w.WriteHeader(400)
|
|
w.Write([]byte(`{"error":"Invalid assistant message: content or tool_calls must be set"}`))
|
|
case 3:
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"recovered"}}]}`))
|
|
default:
|
|
t.Errorf("unexpected request #%d", cur)
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
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)
|
|
if err != nil {
|
|
t.Fatalf("AL: %v", err)
|
|
}
|
|
if got := last(msgs); got != "recovered" {
|
|
t.Errorf("last = %q, want recovered", got)
|
|
}
|
|
for _, m := range msgs {
|
|
if len(m.ToolCalls) > 0 {
|
|
t.Errorf("expected malformed assistant tool-call message to be stripped, found %+v", m)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestSummarizeHappyPath(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"the summary"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
cfg.Endpoint = srv.URL
|
|
cfg.Stream = false
|
|
cfg.APIKey = "-"
|
|
s, err := summarize(&cfg, []Message{{Role: "user", Content: strp("hello world")}})
|
|
if err != nil {
|
|
t.Fatalf("summarize: %v", err)
|
|
}
|
|
if s != "the summary" {
|
|
t.Errorf("summary = %q", s)
|
|
}
|
|
}
|
|
|
|
func TestCompactHappyPath(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"the summary"}}]}`))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := defCfg
|
|
cfg.Endpoint = srv.URL
|
|
cfg.Stream = false
|
|
cfg.APIKey = "-"
|
|
orig := []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("hello world")}}
|
|
msgs, sm, err := compact(&cfg, orig)
|
|
if err != nil {
|
|
t.Fatalf("compact: %v", err)
|
|
}
|
|
if sm != "the summary" {
|
|
t.Errorf("summary = %q", sm)
|
|
}
|
|
if len(msgs) != 2 || msgs[0].Role != "system" || *msgs[0].Content != "sys" {
|
|
t.Errorf("compacted messages = %+v", msgs)
|
|
}
|
|
if msgs[1].Role != "user" || msgs[1].Content == nil || !strings.Contains(*msgs[1].Content, "the summary") {
|
|
t.Errorf("continuation message = %+v", msgs[1])
|
|
}
|
|
}
|