v0.0.5: context compaction, minimal config, retry/search/code symbols/diff/status/test commands

This commit is contained in:
loveuer
2026-06-24 02:20:19 -07:00
parent e4c75c7c0e
commit 950c63adaa
24 changed files with 1939 additions and 220 deletions
+27 -18
View File
@@ -44,19 +44,10 @@ providers:
loveuer: loveuer:
base_url: https://ai.loveuer.com base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY} api_key: ${AGENTU_API_KEY}
model: deepseek-v4-flash
models: models:
- deepseek-v4-flash - id: deepseek-v4-flash
thinking: none thinking: none
thinking_param: thinking context: 256000
agent:
system_prompt: |
You are agentu, a concise local agent. Answer simple conversation directly.
Use tools only when the user's request needs local project context or local
command execution.
working_dir: .
tool_timeout: 30s
``` ```
Web tools are enabled by default with a built-in no-key backend. agentu exposes Web tools are enabled by default with a built-in no-key backend. agentu exposes
@@ -67,18 +58,26 @@ Provider names are the keys under `providers`. At startup, agentu selects the
first provider by name. Use `/model provider <name>` to switch providers for the first provider by name. Use `/model provider <name>` to switch providers for the
current session. current session.
`models` is optional. If present, `/model model <name>` is restricted to that `models` is required. `/model model <name>` is restricted to that list. Each
list. If omitted, any model name is accepted for that provider. entry is an object with:
`thinking` is optional. Supported values are: - `id`: model name sent to the provider.
- `thinking`: optional default thinking level for that model.
- `context`: optional context-window token count used for auto-compaction.
If `model` is omitted, agentu uses the first `models` entry. The `agent` block
is optional: `working_dir`, `tool_timeout`, `system_prompt`, and
`max_context_tokens` all have defaults. `max_context_tokens` overrides the
selected model's `context` when set.
Supported thinking values are:
```text ```text
none, middle, high, xhigh, max none, middle, high, xhigh, max
``` ```
When configured or changed with `/model thinking ...`, the value is sent as a When configured or changed with `/model thinking ...`, the value is sent as a
top-level OpenAI-compatible request field. `thinking_param` controls that field top-level OpenAI-compatible request field named `thinking`.
name and defaults to `thinking`.
## CLI Flags ## CLI Flags
@@ -90,6 +89,7 @@ go run ./cmd/agentu [flags]
- `--theme light|dark` selects the TUI theme. Default: `light`. - `--theme light|dark` selects the TUI theme. Default: `light`.
- `--plain` uses the simple line-based REPL instead of the TUI. - `--plain` uses the simple line-based REPL instead of the TUI.
- `--yolo` enables automatic file writes and shell execution. - `--yolo` enables automatic file writes and shell execution.
- `--resume <id>` resumes a saved session from `~/.agentu/sessions`.
## TUI Controls ## TUI Controls
@@ -103,10 +103,18 @@ go run ./cmd/agentu [flags]
Slash commands: Slash commands:
- `/clear` resets conversation history. - `/clear` resets conversation history.
- `/compact` summarizes older conversation history and keeps recent context.
- `/status` shows changed files with `git status --short`.
- `/diff [--stat] [path...]` shows the unstaged git diff.
- `/test [command...]` runs a project test command; defaults to `go test ./...` for Go modules or `npm test` for Node projects.
- `/model` shows current provider, model, thinking, and switch usage. - `/model` shows current provider, model, thinking, and switch usage.
- `/model provider <name>` switches provider for the current session. - `/model provider <name>` switches provider for the current session.
- `/model model <name>` or `/model <name>` switches model for the current session. - `/model model <name>` or `/model <name>` switches model for the current session.
- `/model thinking <none|middle|high|xhigh|max>` changes thinking for the current session. - `/model thinking <none|middle|high|xhigh|max>` changes thinking for the current session.
- `/sessions` lists saved sessions.
- `/resume [id]` resumes by ID or opens an interactive picker.
- `/rename <name>` renames the current session.
- `/new` saves the current session and starts a fresh one.
- `/exit` or `/quit` exits. - `/exit` or `/quit` exits.
## Local Tools ## Local Tools
@@ -115,7 +123,8 @@ Read-only tools are available by default:
- `file_read` - `file_read`
- `file_list` - `file_list`
- `file_search` - `file_search` with optional `context_lines` and `max_results`
- `code_symbols` for Go packages, imports, types, functions, and methods
Web tools are also available in read-only mode: Web tools are also available in read-only mode:
+3 -12
View File
@@ -2,16 +2,7 @@ providers:
loveuer: loveuer:
base_url: https://ai.loveuer.com base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY} api_key: ${AGENTU_API_KEY}
model: deepseek-v4-flash
models: models:
- deepseek-v4-flash - id: deepseek-v4-flash
thinking: none thinking: none
thinking_param: thinking context: 256000
agent:
system_prompt: |
You are agentu, a concise local agent. Answer simple conversation directly.
Use tools only when the user's request needs local project context or local
command execution.
working_dir: .
tool_timeout: 30s
+58 -14
View File
@@ -20,6 +20,20 @@ import (
"agentu/internal/tui" "agentu/internal/tui"
) )
const defaultConfigContent = `# agentu configuration
# Replace the placeholder values below with your actual provider settings.
# See agentu.example.yaml for the full reference.
providers:
default:
base_url: https://api.openai.com
api_key: YOUR_API_KEY_HERE
models:
- id: gpt-4o
context: 128000
- id: gpt-4o-mini
context: 128000
`
func main() { func main() {
if err := run(); err != nil { if err := run(); err != nil {
fmt.Fprintln(os.Stderr, "agentu:", err) fmt.Fprintln(os.Stderr, "agentu:", err)
@@ -43,11 +57,10 @@ func run() error {
cfg, err := config.Load(*configPath) cfg, err := config.Load(*configPath)
if err != nil { if err != nil {
if errors.Is(err, os.ErrNotExist) { if errors.Is(err, os.ErrNotExist) {
resolvedPath, resolveErr := config.ResolvePath(*configPath) if bootstrapErr := bootstrapConfig(*configPath); bootstrapErr != nil {
if resolveErr != nil { return bootstrapErr
return resolveErr
} }
return fmt.Errorf("config not found at %s; create ~/.agentu/config.yaml from agentu.example.yaml and set AGENTU_API_KEY", resolvedPath) return fmt.Errorf("created default config at %s; edit it to set your API key and provider, then restart agentu", *configPath)
} }
return err return err
} }
@@ -65,17 +78,22 @@ func run() error {
registry.Register(tools.NewBuiltinSearchTool(http.DefaultClient)) registry.Register(tools.NewBuiltinSearchTool(http.DefaultClient))
registry.Register(tools.NewBuiltinFetchTool(http.DefaultClient)) registry.Register(tools.NewBuiltinFetchTool(http.DefaultClient))
systemPrompt := strings.TrimSpace(cfg.Agent.SystemPrompt + "\n\n" + tools.AgentInstructions(cfg.Agent.WorkingDir, *yolo)) systemPrompt := strings.TrimSpace(cfg.Agent.SystemPrompt + "\n\n" + tools.AgentInstructions(cfg.Agent.WorkingDir, *yolo))
maxContextTokens := cfg.Agent.MaxContextTokens
if maxContextTokens == 0 {
maxContextTokens = providerConfig.ContextTokens
}
assistant := agent.New(agent.Options{ assistant := agent.New(agent.Options{
Provider: provider, Provider: provider,
ProviderName: providerName, ProviderName: providerName,
Model: providerConfig.Model, Model: providerConfig.Model,
Thinking: providerConfig.Thinking, Thinking: providerConfig.Thinking,
ThinkingParam: providerConfig.ThinkingParam, ThinkingParam: providerConfig.ThinkingParam,
ThinkingEnabled: providerConfig.ThinkingConfigured, ThinkingEnabled: providerConfig.ThinkingConfigured,
SystemPrompt: systemPrompt, SystemPrompt: systemPrompt,
ToolRegistry: registry, ToolRegistry: registry,
ToolTimeout: cfg.Agent.ToolTimeout, ToolTimeout: cfg.Agent.ToolTimeout,
MaxContextTokens: maxContextTokens,
}) })
// Initialize session store // Initialize session store
@@ -114,6 +132,7 @@ func run() error {
Yolo: *yolo, Yolo: *yolo,
ThemeMode: themeMode, ThemeMode: themeMode,
ModelManager: modelManager, ModelManager: modelManager,
WorkingDir: cfg.Agent.WorkingDir,
}) })
printExitSession(modelManager) printExitSession(modelManager)
return err return err
@@ -122,12 +141,27 @@ func run() error {
return repl(ctx, assistant, modelManager) return repl(ctx, assistant, modelManager)
} }
func bootstrapConfig(configPath string) error {
resolved, err := config.ResolvePath(configPath)
if err != nil {
return fmt.Errorf("resolve config path: %w", err)
}
dir := filepath.Dir(resolved)
if err := os.MkdirAll(dir, 0o755); err != nil {
return fmt.Errorf("create config directory: %w", err)
}
if err := os.WriteFile(resolved, []byte(defaultConfigContent), 0o644); err != nil {
return fmt.Errorf("write default config: %w", err)
}
return nil
}
func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error { func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error {
scanner := bufio.NewScanner(os.Stdin) scanner := bufio.NewScanner(os.Stdin)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
fmt.Println("agentu") fmt.Println("agentu")
fmt.Println("Type /exit to quit, /clear to reset context.") fmt.Println("Type /exit to quit, /clear to reset context, /compact to compress context.")
for { for {
fmt.Print("> ") fmt.Print("> ")
if !scanner.Scan() { if !scanner.Scan() {
@@ -150,6 +184,16 @@ func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Man
assistant.Clear() assistant.Clear()
fmt.Println("context cleared") fmt.Println("context cleared")
continue continue
case "/compact":
if err := assistant.Compact(ctx); err != nil {
fmt.Fprintln(os.Stderr, "compact error:", err)
} else {
fmt.Println("context compacted")
if modelManager != nil {
_ = modelManager.Save()
}
}
continue
} }
if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil { if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
+250 -41
View File
@@ -13,20 +13,36 @@ import (
"agentu/internal/tools" "agentu/internal/tools"
) )
const defaultMaxToolIterations = 8 const (
defaultMaxToolRounds = 50
doomLoopThreshold = 3
compactRecentMessages = 6
summaryMessageLimit = 8000
contextSummaryPrefix = "[context summary]"
summarizerSystemPrompt = "You are a conversation summarizer. Summarize the following conversation history into a concise but detailed summary. Preserve:\n" +
"- Key facts, decisions, and conclusions\n" +
"- File paths and code changes made\n" +
"- Tool call results (especially errors)\n" +
"- Current task state and next steps\n\n" +
"Do NOT include: tool call syntax, raw JSON, or verbatim file contents.\n" +
"Output only the summary, no preamble."
)
type Agent struct { type Agent struct {
provider llm.Provider provider llm.Provider
providerName string providerName string
model string model string
thinking string thinking string
thinkingParam string thinkingParam string
thinkingEnabled bool thinkingEnabled bool
systemPrompt string systemPrompt string
toolRegistry *tools.Registry toolRegistry *tools.Registry
toolTimeout time.Duration toolTimeout time.Duration
maxToolIterations int maxToolRounds int
messages []llm.Message maxContextTokens int
lastUsage *llm.Usage
messages []llm.Message
} }
type Options struct { type Options struct {
@@ -40,47 +56,52 @@ type Options struct {
ToolRegistry *tools.Registry ToolRegistry *tools.Registry
ToolTimeout time.Duration ToolTimeout time.Duration
MaxToolIterations int MaxToolIterations int
MaxContextTokens int
} }
func New(opts Options) *Agent { func New(opts Options) *Agent {
maxToolIterations := opts.MaxToolIterations maxToolRounds := opts.MaxToolIterations
if maxToolIterations <= 0 { if maxToolRounds <= 0 {
maxToolIterations = defaultMaxToolIterations maxToolRounds = defaultMaxToolRounds
} }
toolTimeout := opts.ToolTimeout toolTimeout := opts.ToolTimeout
if toolTimeout <= 0 { if toolTimeout <= 0 {
toolTimeout = 30 * time.Second toolTimeout = 30 * time.Second
} }
a := &Agent{ a := &Agent{
provider: opts.Provider, provider: opts.Provider,
providerName: opts.ProviderName, providerName: opts.ProviderName,
model: opts.Model, model: opts.Model,
thinking: opts.Thinking, thinking: opts.Thinking,
thinkingParam: opts.ThinkingParam, thinkingParam: opts.ThinkingParam,
thinkingEnabled: opts.ThinkingEnabled, thinkingEnabled: opts.ThinkingEnabled,
systemPrompt: opts.SystemPrompt, systemPrompt: opts.SystemPrompt,
toolRegistry: opts.ToolRegistry, toolRegistry: opts.ToolRegistry,
toolTimeout: toolTimeout, toolTimeout: toolTimeout,
maxToolIterations: maxToolIterations, maxToolRounds: maxToolRounds,
maxContextTokens: opts.MaxContextTokens,
} }
a.Clear() a.Clear()
return a return a
} }
type RuntimeOptions struct { type RuntimeOptions struct {
Provider llm.Provider Provider llm.Provider
ProviderName string ProviderName string
Model string Model string
Thinking string Thinking string
ThinkingParam string ThinkingParam string
ThinkingEnabled bool ThinkingEnabled bool
MaxContextTokens int
} }
type RuntimeInfo struct { type RuntimeInfo struct {
ProviderName string ProviderName string
Model string Model string
Thinking string Thinking string
ThinkingEnabled bool ThinkingEnabled bool
ContextTokens int
MaxContextTokens int
} }
func (a *Agent) SetRuntime(opts RuntimeOptions) { func (a *Agent) SetRuntime(opts RuntimeOptions) {
@@ -96,18 +117,22 @@ func (a *Agent) SetRuntime(opts RuntimeOptions) {
a.thinking = opts.Thinking a.thinking = opts.Thinking
a.thinkingParam = opts.ThinkingParam a.thinkingParam = opts.ThinkingParam
a.thinkingEnabled = opts.ThinkingEnabled a.thinkingEnabled = opts.ThinkingEnabled
a.maxContextTokens = opts.MaxContextTokens
} }
func (a *Agent) RuntimeInfo() RuntimeInfo { func (a *Agent) RuntimeInfo() RuntimeInfo {
return RuntimeInfo{ return RuntimeInfo{
ProviderName: a.providerName, ProviderName: a.providerName,
Model: a.model, Model: a.model,
Thinking: a.thinking, Thinking: a.thinking,
ThinkingEnabled: a.thinkingEnabled, ThinkingEnabled: a.thinkingEnabled,
ContextTokens: llm.EstimateTokens(a.messages, a.toolDefinitions()),
MaxContextTokens: a.maxContextTokens,
} }
} }
func (a *Agent) Clear() { func (a *Agent) Clear() {
a.lastUsage = nil
a.messages = nil a.messages = nil
if strings.TrimSpace(a.systemPrompt) != "" { if strings.TrimSpace(a.systemPrompt) != "" {
a.messages = append(a.messages, llm.Message{ a.messages = append(a.messages, llm.Message{
@@ -129,14 +154,161 @@ func (a *Agent) SetMessages(messages []llm.Message) {
a.messages = append([]llm.Message(nil), messages...) a.messages = append([]llm.Message(nil), messages...)
} }
// LastUsage returns a copy of the most recent provider token usage data.
func (a *Agent) LastUsage() *llm.Usage {
if a.lastUsage == nil {
return nil
}
usage := *a.lastUsage
return &usage
}
// SetLastUsage restores the most recent provider token usage data.
func (a *Agent) SetLastUsage(usage *llm.Usage) {
if usage == nil {
a.lastUsage = nil
return
}
copy := *usage
a.lastUsage = &copy
}
// Compact summarizes older conversation messages while preserving the system
// prompt and recent messages intact.
func (a *Agent) Compact(ctx context.Context) error {
if len(a.messages) < 4 {
return nil
}
if a.provider == nil {
return fmt.Errorf("provider is not configured")
}
prefixEnd := 0
if a.messages[0].Role == llm.RoleSystem {
prefixEnd = 1
}
if len(a.messages)-prefixEnd <= compactRecentMessages {
return nil
}
recentStart := len(a.messages) - compactRecentMessages
if recentStart <= prefixEnd {
return nil
}
middle := append([]llm.Message(nil), a.messages[prefixEnd:recentStart]...)
if len(middle) == 0 {
return nil
}
summary, usage, err := a.summarizeMessages(ctx, middle)
if err != nil {
return fmt.Errorf("compact context: %w", err)
}
summary = strings.TrimSpace(summary)
if summary == "" {
return fmt.Errorf("compact context: summary was empty")
}
compressed := make([]llm.Message, 0, prefixEnd+1+len(a.messages[recentStart:]))
compressed = append(compressed, a.messages[:prefixEnd]...)
compressed = append(compressed, llm.Message{
Role: llm.RoleAssistant,
Content: contextSummaryPrefix + "\n" + summary,
})
compressed = append(compressed, a.messages[recentStart:]...)
a.messages = compressed
if usage != nil {
a.lastUsage = usage
}
return nil
}
func (a *Agent) shouldCompact() bool {
if a.maxContextTokens <= 0 {
return false
}
if a.lastUsage != nil && a.lastUsage.TotalTokens >= a.maxContextTokens {
return true
}
return llm.EstimateTokens(a.messages, a.toolDefinitions()) >= a.maxContextTokens
}
func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (string, *llm.Usage, error) {
req := llm.ChatRequest{
Model: a.model,
Messages: []llm.Message{
{Role: llm.RoleSystem, Content: summarizerSystemPrompt},
{Role: llm.RoleUser, Content: formatMessagesForSummary(messages)},
},
}
var summary strings.Builder
var usage *llm.Usage
if err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
if event.Content != "" {
summary.WriteString(event.Content)
}
if event.Usage != nil {
copy := *event.Usage
usage = &copy
}
return nil
}); err != nil {
return "", nil, err
}
return strings.TrimSpace(summary.String()), usage, nil
}
func formatMessagesForSummary(messages []llm.Message) string {
var b strings.Builder
b.WriteString("Conversation history to summarize:\n\n")
for index, msg := range messages {
fmt.Fprintf(&b, "Message %d (%s):\n", index+1, msg.Role)
if msg.ToolCallID != "" {
fmt.Fprintf(&b, "tool_call_id: %s\n", msg.ToolCallID)
}
if msg.Content != "" {
fmt.Fprintf(&b, "%s\n", compact(msg.Content, summaryMessageLimit))
}
for _, call := range msg.ToolCalls {
fmt.Fprintf(&b, "tool_call: %s %s\n", call.Function.Name, compact(call.Function.Arguments, summaryMessageLimit))
}
b.WriteByte('\n')
}
return b.String()
}
// roundSignature returns a deterministic signature for a set of tool calls
// in a single round, used for doom loop detection.
func roundSignature(calls []llm.ToolCall) string {
var b strings.Builder
for _, c := range calls {
b.WriteString(c.Function.Name)
b.WriteByte('\x00')
b.WriteString(strings.TrimSpace(c.Function.Arguments))
b.WriteByte('\x00')
}
return b.String()
}
func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error { func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
input = strings.TrimSpace(input) input = strings.TrimSpace(input)
if input == "" { if input == "" {
return nil return nil
} }
a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input}) a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input})
if a.shouldCompact() {
if err := a.Compact(ctx); err != nil {
if logs != nil {
fmt.Fprintf(logs, "[compact] warning: %v\n", err)
}
}
}
for iteration := 0; iteration <= a.maxToolIterations; iteration++ { var recentSignatures []string
for round := 0; round < a.maxToolRounds; round++ {
content, calls, err := a.streamAssistant(ctx, out) content, calls, err := a.streamAssistant(ctx, out)
if err != nil { if err != nil {
return err return err
@@ -158,6 +330,40 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
fmt.Fprintln(out) fmt.Fprintln(out)
} }
// Doom loop detection: check if the last N rounds all called
// the exact same tools with the exact same arguments.
sig := roundSignature(calls)
recentSignatures = append(recentSignatures, sig)
if len(recentSignatures) > doomLoopThreshold {
recentSignatures = recentSignatures[len(recentSignatures)-doomLoopThreshold:]
}
if len(recentSignatures) == doomLoopThreshold {
allSame := true
for i := 1; i < len(recentSignatures); i++ {
if recentSignatures[i] != recentSignatures[0] {
allSame = false
break
}
}
if allSame {
toolNames := make([]string, 0, len(calls))
for _, c := range calls {
toolNames = append(toolNames, c.Function.Name)
}
msg := fmt.Sprintf(
"error: doom loop detected — the same tool call (%s) with identical arguments has been repeated %d times. Stopping. Try a different approach or break the task into smaller steps.",
strings.Join(toolNames, ", "),
doomLoopThreshold,
)
a.messages = append(a.messages, llm.Message{
Role: llm.RoleAssistant,
Content: msg,
})
return fmt.Errorf("%s", msg)
}
}
for _, call := range calls { for _, call := range calls {
result := a.executeTool(ctx, call, logs) result := a.executeTool(ctx, call, logs)
a.messages = append(a.messages, llm.Message{ a.messages = append(a.messages, llm.Message{
@@ -168,7 +374,7 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
} }
} }
return fmt.Errorf("tool iteration limit reached: %d", a.maxToolIterations) return fmt.Errorf("tool round limit reached after %d rounds; the task may be too complex for a single turn", a.maxToolRounds)
} }
func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []llm.ToolCall, error) { func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []llm.ToolCall, error) {
@@ -192,6 +398,9 @@ func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []l
} }
err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error { err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
if event.Usage != nil {
a.SetLastUsage(event.Usage)
}
if event.Content != "" { if event.Content != "" {
content.WriteString(event.Content) content.WriteString(event.Content)
if out != nil { if out != nil {
+203
View File
@@ -3,6 +3,7 @@ package agent
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"fmt"
"strings" "strings"
"testing" "testing"
@@ -96,6 +97,134 @@ func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, e
return emit(llm.StreamEvent{Content: p.text}) return emit(llm.StreamEvent{Content: p.text})
} }
type compactProvider struct {
summaryCalls int
answerCalls int
lastSummary string
}
func (p *compactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
if len(req.Messages) == 2 && req.Messages[0].Role == llm.RoleSystem && req.Messages[0].Content == summarizerSystemPrompt {
p.summaryCalls++
p.lastSummary = req.Messages[1].Content
if err := emit(llm.StreamEvent{Content: "old work summarized"}); err != nil {
return err
}
return emit(llm.StreamEvent{Usage: &llm.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}})
}
p.answerCalls++
return emit(llm.StreamEvent{Content: "done"})
}
func TestAgentCompact(t *testing.T) {
provider := &compactProvider{}
a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
a.SetMessages(longConversation())
if err := a.Compact(context.Background()); err != nil {
t.Fatal(err)
}
messages := a.Messages()
if len(messages) != 8 {
t.Fatalf("messages = %d, want 8", len(messages))
}
if messages[0].Role != llm.RoleSystem || messages[0].Content != "system" {
t.Fatalf("system message = %#v", messages[0])
}
if messages[1].Role != llm.RoleAssistant || !strings.HasPrefix(messages[1].Content, contextSummaryPrefix) {
t.Fatalf("summary message = %#v", messages[1])
}
if !strings.Contains(messages[1].Content, "old work summarized") {
t.Fatalf("summary content = %q", messages[1].Content)
}
if !strings.Contains(provider.lastSummary, "old-user-0") {
t.Fatalf("summary input did not include old history: %q", provider.lastSummary)
}
if messages[len(messages)-1].Content != "recent-answer-2" {
t.Fatalf("last recent message = %#v", messages[len(messages)-1])
}
if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 12 {
t.Fatalf("usage = %#v", usage)
}
}
func TestAgentAutoCompact(t *testing.T) {
provider := &compactProvider{}
a := New(Options{
Provider: provider,
Model: "test",
SystemPrompt: "system",
MaxContextTokens: 1,
ToolRegistry: tools.NewRegistry(),
MaxToolIterations: 1,
})
a.SetMessages(longConversation())
var out strings.Builder
if err := a.RunTurn(context.Background(), "new question", &out, nil); err != nil {
t.Fatal(err)
}
if provider.summaryCalls != 1 {
t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
}
if provider.answerCalls != 1 {
t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
}
if out.String() != "done" {
t.Fatalf("out = %q", out.String())
}
if !strings.Contains(joinMessageContents(a.Messages()), contextSummaryPrefix) {
t.Fatalf("messages missing summary: %#v", a.Messages())
}
}
func TestAgentCompactSkipsShortHistory(t *testing.T) {
provider := &contentProvider{text: "unused"}
a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
a.SetMessages([]llm.Message{
{Role: llm.RoleSystem, Content: "system"},
{Role: llm.RoleUser, Content: "question"},
{Role: llm.RoleAssistant, Content: "answer"},
})
if err := a.Compact(context.Background()); err != nil {
t.Fatal(err)
}
if provider.called {
t.Fatal("provider should not be called for short history")
}
if len(a.Messages()) != 3 {
t.Fatalf("messages = %d, want 3", len(a.Messages()))
}
}
func longConversation() []llm.Message {
messages := []llm.Message{{Role: llm.RoleSystem, Content: "system"}}
for i := 0; i < 3; i++ {
messages = append(messages,
llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("old-user-%d", i)},
llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("old-answer-%d", i)},
)
}
for i := 0; i < 3; i++ {
messages = append(messages,
llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("recent-user-%d", i)},
llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("recent-answer-%d", i)},
)
}
return messages
}
func joinMessageContents(messages []llm.Message) string {
var b strings.Builder
for _, msg := range messages {
b.WriteString(msg.Content)
b.WriteByte('\n')
}
return b.String()
}
func TestAgentLetsProviderDecideShellRequestWhenShellDisabled(t *testing.T) { func TestAgentLetsProviderDecideShellRequestWhenShellDisabled(t *testing.T) {
provider := &contentProvider{text: "shell is disabled"} provider := &contentProvider{text: "shell is disabled"}
a := New(Options{ a := New(Options{
@@ -137,3 +266,77 @@ func TestAgentDoesNotBlockFileRequest(t *testing.T) {
t.Fatalf("out = %q", out.String()) t.Fatalf("out = %q", out.String())
} }
} }
// alwaysToolProvider always returns a tool call with the same args on every request.
type alwaysToolProvider struct {
calls int
}
func (p *alwaysToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.calls++
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo", Arguments: `{"text":"loop"}`},
}})
}
func TestAgentStopsOnDoomLoop(t *testing.T) {
provider := &alwaysToolProvider{}
a := New(Options{
Provider: provider,
Model: "test",
SystemPrompt: "system",
ToolRegistry: tools.NewRegistry(echoTool{}),
MaxToolIterations: 50, // High hard limit — doom loop should trigger first
})
var out strings.Builder
err := a.RunTurn(context.Background(), "loop forever", &out, nil)
if err == nil {
t.Fatal("expected error when doom loop is detected")
}
if !strings.Contains(err.Error(), "doom loop detected") {
t.Fatalf("unexpected error: %v", err)
}
// Should have stopped at doomLoopThreshold (3) rounds, not reached hard limit.
if provider.calls != doomLoopThreshold {
t.Fatalf("calls = %d, want %d (doom loop threshold)", provider.calls, doomLoopThreshold)
}
}
// varyingToolProvider alternates between different tool calls to avoid doom loop detection.
type varyingToolProvider struct {
calls int
}
func (p *varyingToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.calls++
// Vary the arguments each round to avoid doom loop detection.
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo",
Arguments: fmt.Sprintf(`{"text":"round-%d"}`, p.calls)},
}})
}
func TestAgentHardLimitStopsVaryingCalls(t *testing.T) {
provider := &varyingToolProvider{}
const limit = 5
a := New(Options{
Provider: provider,
Model: "test",
SystemPrompt: "system",
ToolRegistry: tools.NewRegistry(echoTool{}),
MaxToolIterations: limit,
})
var out strings.Builder
err := a.RunTurn(context.Background(), "vary forever", &out, nil)
if err == nil {
t.Fatal("expected error when hard limit is reached")
}
if !strings.Contains(err.Error(), "tool round limit reached") {
t.Fatalf("unexpected error: %v", err)
}
if provider.calls != limit {
t.Fatalf("calls = %d, want %d", provider.calls, limit)
}
}
+97 -35
View File
@@ -21,37 +21,50 @@ type Config struct {
type ProviderConfig struct { type ProviderConfig struct {
Name string Name string
BaseURL string `yaml:"base_url"` BaseURL string
Model string `yaml:"model"` Model string
Models []string `yaml:"models"` Models []string
APIKey string `yaml:"api_key"` ModelConfigs map[string]ModelConfig
Thinking string `yaml:"thinking"` ContextTokens int
ThinkingParam string `yaml:"thinking_param"` APIKey string
Thinking string
ThinkingParam string
ThinkingConfigured bool ThinkingConfigured bool
} }
type ModelConfig struct {
ID string `yaml:"id"`
Thinking string `yaml:"thinking"`
ContextTokens int `yaml:"context"`
}
type AgentConfig struct { type AgentConfig struct {
SystemPrompt string `yaml:"system_prompt"` SystemPrompt string `yaml:"system_prompt"`
WorkingDir string `yaml:"working_dir"` WorkingDir string `yaml:"working_dir"`
ToolTimeout time.Duration `yaml:"tool_timeout"` ToolTimeout time.Duration `yaml:"tool_timeout"`
MaxContextTokens int `yaml:"max_context_tokens"`
} }
type rawConfig struct { type rawConfig struct {
Providers map[string]rawProviderConfig `yaml:"providers"` Providers map[string]rawProviderConfig `yaml:"providers"`
Agent struct { Agent struct {
SystemPrompt string `yaml:"system_prompt"` SystemPrompt string `yaml:"system_prompt"`
WorkingDir string `yaml:"working_dir"` WorkingDir string `yaml:"working_dir"`
ToolTimeout string `yaml:"tool_timeout"` ToolTimeout string `yaml:"tool_timeout"`
MaxContextTokens int `yaml:"max_context_tokens"`
} `yaml:"agent"` } `yaml:"agent"`
} }
type rawProviderConfig struct { type rawProviderConfig struct {
BaseURL string `yaml:"base_url"` BaseURL string `yaml:"base_url"`
Model string `yaml:"model"` Models []rawModelConfig `yaml:"models"`
Models []string `yaml:"models"` APIKey string `yaml:"api_key"`
APIKey string `yaml:"api_key"` }
Thinking *string `yaml:"thinking"`
ThinkingParam string `yaml:"thinking_param"` type rawModelConfig struct {
ID string `yaml:"id"`
Thinking string `yaml:"thinking"`
ContextTokens int `yaml:"context"`
} }
func Load(path string) (*Config, error) { func Load(path string) (*Config, error) {
@@ -76,8 +89,9 @@ func Load(path string) (*Config, error) {
cfg := &Config{ cfg := &Config{
Providers: normalizeProviders(raw), Providers: normalizeProviders(raw),
Agent: AgentConfig{ Agent: AgentConfig{
SystemPrompt: raw.Agent.SystemPrompt, SystemPrompt: raw.Agent.SystemPrompt,
WorkingDir: raw.Agent.WorkingDir, WorkingDir: raw.Agent.WorkingDir,
MaxContextTokens: raw.Agent.MaxContextTokens,
}, },
} }
@@ -142,8 +156,8 @@ func (c *Config) Validate() error {
if strings.TrimSpace(provider.BaseURL) == "" { if strings.TrimSpace(provider.BaseURL) == "" {
missing = append(missing, prefix+".base_url") missing = append(missing, prefix+".base_url")
} }
if strings.TrimSpace(provider.Model) == "" { if len(provider.Models) == 0 {
missing = append(missing, prefix+".model") missing = append(missing, prefix+".models")
} }
if strings.TrimSpace(provider.APIKey) == "" { if strings.TrimSpace(provider.APIKey) == "" {
missing = append(missing, prefix+".api_key") missing = append(missing, prefix+".api_key")
@@ -151,6 +165,19 @@ func (c *Config) Validate() error {
if len(missing) > 0 { if len(missing) > 0 {
return fmt.Errorf("missing required config: %s", strings.Join(missing, ", ")) return fmt.Errorf("missing required config: %s", strings.Join(missing, ", "))
} }
for _, model := range provider.Models {
if strings.TrimSpace(model) == "" {
return fmt.Errorf("%s.models.id is required", prefix)
}
}
for _, model := range provider.ModelConfigs {
if model.ContextTokens < 0 {
return fmt.Errorf("%s.models.context must be greater than or equal to zero", prefix)
}
if model.Thinking != "" && !IsThinkingLevel(model.Thinking) {
return fmt.Errorf("%s.models.thinking must be one of: %s", prefix, strings.Join(ThinkingLevels(), ", "))
}
}
if provider.ThinkingConfigured && !IsThinkingLevel(provider.Thinking) { if provider.ThinkingConfigured && !IsThinkingLevel(provider.Thinking) {
return fmt.Errorf("%s.thinking must be one of: %s", prefix, strings.Join(ThinkingLevels(), ", ")) return fmt.Errorf("%s.thinking must be one of: %s", prefix, strings.Join(ThinkingLevels(), ", "))
} }
@@ -158,6 +185,9 @@ func (c *Config) Validate() error {
if c.Agent.ToolTimeout <= 0 { if c.Agent.ToolTimeout <= 0 {
return errors.New("agent.tool_timeout must be greater than zero") return errors.New("agent.tool_timeout must be greater than zero")
} }
if c.Agent.MaxContextTokens < 0 {
return errors.New("agent.max_context_tokens must be greater than or equal to zero")
}
return nil return nil
} }
@@ -187,27 +217,59 @@ func normalizeProviders(raw rawConfig) map[string]ProviderConfig {
} }
func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig { func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
models, modelConfigs := normalizeModelConfigs(raw.Models)
provider := ProviderConfig{ provider := ProviderConfig{
Name: name, Name: name,
BaseURL: raw.BaseURL, BaseURL: raw.BaseURL,
Model: raw.Model, Models: models,
Models: append([]string(nil), raw.Models...), ModelConfigs: modelConfigs,
APIKey: raw.APIKey, APIKey: raw.APIKey,
ThinkingParam: raw.ThinkingParam,
} }
if provider.Model == "" && len(provider.Models) > 0 { if len(provider.Models) > 0 {
provider.Model = provider.Models[0] provider.Model = provider.Models[0]
} }
if raw.Thinking != nil { applyModelDefaults(&provider)
provider.Thinking = strings.ToLower(strings.TrimSpace(*raw.Thinking)) if provider.ThinkingConfigured && provider.ThinkingParam == "" {
provider.ThinkingConfigured = true provider.ThinkingParam = "thinking"
if provider.ThinkingParam == "" {
provider.ThinkingParam = "thinking"
}
} }
return provider return provider
} }
func normalizeModelConfigs(rawModels []rawModelConfig) ([]string, map[string]ModelConfig) {
models := make([]string, 0, len(rawModels))
configs := make(map[string]ModelConfig, len(rawModels))
for _, raw := range rawModels {
model := ModelConfig{
ID: strings.TrimSpace(raw.ID),
Thinking: strings.ToLower(strings.TrimSpace(raw.Thinking)),
ContextTokens: raw.ContextTokens,
}
models = append(models, model.ID)
if model.ID != "" {
configs[model.ID] = model
}
}
return models, configs
}
func applyModelDefaults(provider *ProviderConfig) {
model, ok := provider.ModelConfigs[provider.Model]
if !ok {
return
}
provider.ContextTokens = model.ContextTokens
if model.Thinking != "" {
provider.Thinking = model.Thinking
provider.ThinkingConfigured = true
}
}
func (p ProviderConfig) WithModel(model string) ProviderConfig {
p.Model = strings.TrimSpace(model)
applyModelDefaults(&p)
return p
}
func ThinkingLevels() []string { func ThinkingLevels() []string {
return []string{"none", "middle", "high", "xhigh", "max"} return []string{"none", "middle", "high", "xhigh", "max"}
} }
+130 -15
View File
@@ -16,8 +16,9 @@ func TestLoadExpandsEnvAndDefaults(t *testing.T) {
providers: providers:
test: test:
base_url: https://example.com base_url: https://example.com
model: test-model
api_key: ${AGENTU_TEST_KEY} api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
agent: agent:
working_dir: . working_dir: .
`), 0o644); err != nil { `), 0o644); err != nil {
@@ -43,6 +44,115 @@ agent:
} }
} }
func TestLoadMinimalModelObjectConfig(t *testing.T) {
dir := t.TempDir()
t.Setenv("AGENTU_TEST_KEY", "sk-test")
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
test:
base_url: https://example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: model-a
thinking: none
context: 256000
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
provider := cfg.Providers["test"]
if provider.Model != "model-a" {
t.Fatalf("model = %q", provider.Model)
}
if len(provider.Models) != 1 || provider.Models[0] != "model-a" {
t.Fatalf("models = %#v", provider.Models)
}
if provider.ContextTokens != 256000 {
t.Fatalf("context tokens = %d", provider.ContextTokens)
}
if !provider.ThinkingConfigured || provider.Thinking != "none" || provider.ThinkingParam != "thinking" {
t.Fatalf("thinking config = %#v", provider)
}
if cfg.Agent.ToolTimeout != 30*time.Second {
t.Fatalf("tool timeout = %s", cfg.Agent.ToolTimeout)
}
if cfg.Agent.SystemPrompt == "" {
t.Fatal("system prompt default was not applied")
}
}
func TestLoadRejectsStringModelEntries(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
test:
base_url: https://example.com
api_key: sk-test
models:
- test-model
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "cannot unmarshal") {
t.Fatalf("err = %v", err)
}
}
func TestLoadMaxContextTokens(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
test:
base_url: https://example.com
api_key: sk-test
models:
- id: test-model
agent:
max_context_tokens: 120000
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if cfg.Agent.MaxContextTokens != 120000 {
t.Fatalf("max context tokens = %d", cfg.Agent.MaxContextTokens)
}
}
func TestLoadRejectsNegativeMaxContextTokens(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
test:
base_url: https://example.com
api_key: sk-test
models:
- id: test-model
agent:
max_context_tokens: -1
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "agent.max_context_tokens") {
t.Fatalf("err = %v", err)
}
}
func TestLoadUsesDefaultConfigPath(t *testing.T) { func TestLoadUsesDefaultConfigPath(t *testing.T) {
dir := t.TempDir() dir := t.TempDir()
t.Setenv("HOME", dir) t.Setenv("HOME", dir)
@@ -56,8 +166,9 @@ func TestLoadUsesDefaultConfigPath(t *testing.T) {
providers: providers:
test: test:
base_url: https://example.com base_url: https://example.com
model: test-model
api_key: ${AGENTU_TEST_KEY} api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
`), 0o644); err != nil { `), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -93,8 +204,9 @@ func TestLoadDefaultSystemPromptDiscouragesEagerTools(t *testing.T) {
providers: providers:
test: test:
base_url: https://example.com base_url: https://example.com
model: test-model
api_key: ${AGENTU_TEST_KEY} api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
`), 0o644); err != nil { `), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -119,17 +231,15 @@ providers:
loveuer: loveuer:
base_url: https://example.com base_url: https://example.com
api_key: ${AGENTU_TEST_KEY} api_key: ${AGENTU_TEST_KEY}
model: model-a
models: models:
- model-a - id: model-a
- model-b thinking: xhigh
thinking: xhigh - id: model-b
thinking_param: thinking
backup: backup:
base_url: https://backup.example.com base_url: https://backup.example.com
api_key: ${AGENTU_TEST_KEY} api_key: ${AGENTU_TEST_KEY}
models: models:
- model-c - id: model-c
`), 0o644); err != nil { `), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -162,11 +272,13 @@ providers:
two: two:
base_url: https://two.example.com base_url: https://two.example.com
api_key: ${AGENTU_TEST_KEY} api_key: ${AGENTU_TEST_KEY}
model: model-b models:
- id: model-b
one: one:
base_url: https://one.example.com base_url: https://one.example.com
api_key: ${AGENTU_TEST_KEY} api_key: ${AGENTU_TEST_KEY}
model: model-a models:
- id: model-a
`), 0o644); err != nil { `), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -188,9 +300,10 @@ func TestLoadRejectsInvalidThinking(t *testing.T) {
providers: providers:
test: test:
base_url: https://example.com base_url: https://example.com
model: test-model
api_key: ${AGENTU_TEST_KEY} api_key: ${AGENTU_TEST_KEY}
thinking: huge models:
- id: test-model
thinking: huge
`), 0o644); err != nil { `), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -224,8 +337,9 @@ func TestLoadRejectsLegacyProviderConfig(t *testing.T) {
if err := os.WriteFile(path, []byte(` if err := os.WriteFile(path, []byte(`
provider: provider:
base_url: https://example.com base_url: https://example.com
model: test-model
api_key: sk-test api_key: sk-test
models:
- id: test-model
`), 0o644); err != nil { `), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
@@ -244,8 +358,9 @@ active_provider: loveuer
providers: providers:
loveuer: loveuer:
base_url: https://example.com base_url: https://example.com
model: test-model
api_key: sk-test api_key: sk-test
models:
- id: test-model
`), 0o644); err != nil { `), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
+79 -26
View File
@@ -10,8 +10,13 @@ import (
"io" "io"
"net/http" "net/http"
"strings" "strings"
"time"
) )
const maxChatAttempts = 3
const retryDelay = 100 * time.Millisecond
type OpenAICompatibleClient struct { type OpenAICompatibleClient struct {
baseURL string baseURL string
apiKey string apiKey string
@@ -31,34 +36,82 @@ func NewOpenAICompatibleClient(baseURL, apiKey string, httpClient *http.Client)
func (c *OpenAICompatibleClient) ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error { func (c *OpenAICompatibleClient) ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error {
req.Stream = true req.Stream = true
extra := make(map[string]any, len(req.Extra)+1)
for key, value := range req.Extra {
extra[key] = value
}
if _, exists := extra["stream_options"]; !exists {
extra["stream_options"] = map[string]any{"include_usage": true}
}
req.Extra = extra
body, err := json.Marshal(req) body, err := json.Marshal(req)
if err != nil { if err != nil {
return fmt.Errorf("marshal chat request: %w", err) return fmt.Errorf("marshal chat request: %w", err)
} }
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body)) var lastErr error
if err != nil { for attempt := 1; attempt <= maxChatAttempts; attempt++ {
return fmt.Errorf("create chat request: %w", err) httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
} if err != nil {
httpReq.Header.Set("Content-Type", "application/json") return fmt.Errorf("create chat request: %w", err)
httpReq.Header.Set("Accept", "text/event-stream") }
if c.apiKey != "" { httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) httpReq.Header.Set("Accept", "text/event-stream")
} if c.apiKey != "" {
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
}
resp, err := c.httpClient.Do(httpReq) resp, err := c.httpClient.Do(httpReq)
if err != nil { if err != nil {
return fmt.Errorf("send chat request: %w", err) lastErr = fmt.Errorf("send chat request: %w", err)
} if !shouldRetryRequest(ctx, attempt, 0) {
defer resp.Body.Close() return lastErr
}
if waitErr := waitBeforeRetry(ctx); waitErr != nil {
return waitErr
}
continue
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 { if resp.StatusCode < 200 || resp.StatusCode >= 300 {
limit := io.LimitReader(resp.Body, 4096) limit := io.LimitReader(resp.Body, 4096)
data, _ := io.ReadAll(limit) data, _ := io.ReadAll(limit)
return fmt.Errorf("chat request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data))) _ = resp.Body.Close()
} lastErr = fmt.Errorf("chat request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
if !shouldRetryRequest(ctx, attempt, resp.StatusCode) {
return lastErr
}
if waitErr := waitBeforeRetry(ctx); waitErr != nil {
return waitErr
}
continue
}
return readSSE(resp.Body, emit) defer resp.Body.Close()
return readSSE(resp.Body, emit)
}
return lastErr
}
func shouldRetryRequest(ctx context.Context, attempt int, status int) bool {
if ctx.Err() != nil || attempt >= maxChatAttempts {
return false
}
if status == 0 {
return true
}
return status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500
}
func waitBeforeRetry(ctx context.Context) error {
timer := time.NewTimer(retryDelay)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
} }
func readSSE(r io.Reader, emit func(StreamEvent) error) error { func readSSE(r io.Reader, emit func(StreamEvent) error) error {
@@ -79,7 +132,7 @@ func readSSE(r io.Reader, emit func(StreamEvent) error) error {
if parseErr != nil { if parseErr != nil {
return parseErr return parseErr
} }
if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" { if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" || event.Usage != nil {
if emitErr := emit(event); emitErr != nil { if emitErr := emit(event); emitErr != nil {
return emitErr return emitErr
} }
@@ -100,6 +153,7 @@ type streamPayload struct {
} `json:"delta"` } `json:"delta"`
FinishReason string `json:"finish_reason"` FinishReason string `json:"finish_reason"`
} `json:"choices"` } `json:"choices"`
Usage *Usage `json:"usage,omitempty"`
Error *struct { Error *struct {
Message string `json:"message"` Message string `json:"message"`
Type string `json:"type"` Type string `json:"type"`
@@ -127,15 +181,14 @@ func parseStreamPayload(payload string) (StreamEvent, error) {
} }
return StreamEvent{}, fmt.Errorf("provider error: %s", decoded.Error.Message) return StreamEvent{}, fmt.Errorf("provider error: %s", decoded.Error.Message)
} }
event := StreamEvent{Usage: decoded.Usage}
if len(decoded.Choices) == 0 { if len(decoded.Choices) == 0 {
return StreamEvent{}, nil return event, nil
} }
choice := decoded.Choices[0] choice := decoded.Choices[0]
event := StreamEvent{ event.Content = choice.Delta.Content
Content: choice.Delta.Content, event.FinishReason = choice.FinishReason
FinishReason: choice.FinishReason,
}
for _, tc := range choice.Delta.ToolCalls { for _, tc := range choice.Delta.ToolCalls {
event.ToolCalls = append(event.ToolCalls, ToolCallDelta{ event.ToolCalls = append(event.ToolCalls, ToolCallDelta{
Index: tc.Index, Index: tc.Index,
+73
View File
@@ -30,6 +30,10 @@ func TestOpenAICompatibleClientStreamsContentAndToolCalls(t *testing.T) {
if raw["thinking"] != "high" { if raw["thinking"] != "high" {
t.Fatalf("thinking = %#v", raw["thinking"]) t.Fatalf("thinking = %#v", raw["thinking"])
} }
streamOptions, ok := raw["stream_options"].(map[string]any)
if !ok || streamOptions["include_usage"] != true {
t.Fatalf("stream_options = %#v", raw["stream_options"])
}
if !req.Stream { if !req.Stream {
t.Fatal("request did not enable stream") t.Fatal("request did not enable stream")
} }
@@ -66,6 +70,75 @@ func TestOpenAICompatibleClientStreamsContentAndToolCalls(t *testing.T) {
} }
} }
func TestChatStreamCapturesUsage(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body := mustReadBody(t, r)
var raw map[string]any
if err := json.NewDecoder(strings.NewReader(body)).Decode(&raw); err != nil {
t.Fatal(err)
}
streamOptions, ok := raw["stream_options"].(map[string]any)
if !ok || streamOptions["include_usage"] != true {
t.Fatalf("stream_options = %#v", raw["stream_options"])
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"ok"},"finish_reason":""}]}` + "\n\n"))
_, _ = w.Write([]byte(`data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15}}` + "\n\n"))
_, _ = w.Write([]byte("data: [DONE]\n\n"))
}))
defer server.Close()
client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
var usage *Usage
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
if event.Usage != nil {
usage = event.Usage
}
return nil
})
if err != nil {
t.Fatal(err)
}
if usage == nil {
t.Fatal("usage was not emitted")
}
if usage.PromptTokens != 12 || usage.CompletionTokens != 3 || usage.TotalTokens != 15 {
t.Fatalf("usage = %#v", usage)
}
}
func TestOpenAICompatibleClientRetriesTransientHTTPError(t *testing.T) {
attempts := 0
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
attempts++
if attempts == 1 {
http.Error(w, "temporary", http.StatusBadGateway)
return
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"ok"},"finish_reason":""}]}` + "\n\n"))
_, _ = w.Write([]byte("data: [DONE]\n\n"))
}))
defer server.Close()
client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
var content strings.Builder
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
content.WriteString(event.Content)
return nil
})
if err != nil {
t.Fatal(err)
}
if attempts != 2 {
t.Fatalf("attempts = %d, want 2", attempts)
}
if content.String() != "ok" {
t.Fatalf("content = %q", content.String())
}
}
func mustReadBody(t *testing.T, r *http.Request) string { func mustReadBody(t *testing.T, r *http.Request) string {
t.Helper() t.Helper()
data, err := io.ReadAll(r.Body) data, err := io.ReadAll(r.Body)
+43
View File
@@ -0,0 +1,43 @@
package llm
import "encoding/json"
const tokenEstimateCharRatio = 4
// EstimateTokens estimates the token cost of messages and tool definitions in a
// chat completion request. It intentionally uses a simple JSON-size heuristic:
// roughly four UTF-8 bytes per token plus a small per-message framing overhead.
func EstimateTokens(messages []Message, tools []Tool) int {
if len(messages) == 0 && len(tools) == 0 {
return 0
}
payload := struct {
Messages []Message `json:"messages,omitempty"`
Tools []Tool `json:"tools,omitempty"`
}{
Messages: messages,
Tools: tools,
}
data, err := json.Marshal(payload)
if err != nil {
return 0
}
return estimateBytesAsTokens(len(data)) + len(messages)*4
}
func estimateMessageTokens(msg Message) int {
data, err := json.Marshal(msg)
if err != nil {
return 0
}
return estimateBytesAsTokens(len(data)) + 4
}
func estimateBytesAsTokens(n int) int {
if n <= 0 {
return 0
}
return (n + tokenEstimateCharRatio - 1) / tokenEstimateCharRatio
}
+45
View File
@@ -0,0 +1,45 @@
package llm
import "testing"
func TestEstimateTokensEmpty(t *testing.T) {
if got := EstimateTokens(nil, nil); got != 0 {
t.Fatalf("EstimateTokens(nil, nil) = %d, want 0", got)
}
}
func TestEstimateTokens(t *testing.T) {
shortMessages := []Message{{Role: RoleUser, Content: "hello"}}
short := EstimateTokens(shortMessages, nil)
if short <= 4 {
t.Fatalf("short estimate = %d, want message framing plus content", short)
}
if short > 50 {
t.Fatalf("short estimate = %d, looks too high for one short message", short)
}
longMessages := []Message{
{Role: RoleUser, Content: "hello"},
{Role: RoleAssistant, Content: "This is a longer response with enough content to exceed the short estimate."},
}
long := EstimateTokens(longMessages, nil)
if long <= short {
t.Fatalf("long estimate = %d, want > short estimate %d", long, short)
}
withTools := EstimateTokens(shortMessages, []Tool{{
Type: "function",
Function: ToolFunction{
Name: "file_read",
Description: "Read a file",
Parameters: []byte(`{"type":"object"}`),
},
}})
if withTools <= short {
t.Fatalf("withTools estimate = %d, want > short estimate %d", withTools, short)
}
if msgTokens := estimateMessageTokens(shortMessages[0]); msgTokens <= 4 {
t.Fatalf("estimateMessageTokens = %d, want message framing plus content", msgTokens)
}
}
+7
View File
@@ -79,6 +79,13 @@ type StreamEvent struct {
Content string Content string
ToolCalls []ToolCallDelta ToolCalls []ToolCallDelta
FinishReason string FinishReason string
Usage *Usage
}
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
} }
type ToolCallDelta struct { type ToolCallDelta struct {
+19 -13
View File
@@ -69,6 +69,7 @@ func (m *Manager) Resume(id string) error {
} }
m.currentSession = sess m.currentSession = sess
m.agent.SetMessages(sess.Messages) m.agent.SetMessages(sess.Messages)
m.agent.SetLastUsage(sess.LastUsage)
return nil return nil
} }
@@ -78,6 +79,7 @@ func (m *Manager) Save() error {
return nil return nil
} }
m.currentSession.Messages = m.agent.Messages() m.currentSession.Messages = m.agent.Messages()
m.currentSession.LastUsage = m.agent.LastUsage()
m.currentSession.Provider = m.providerName m.currentSession.Provider = m.providerName
m.currentSession.Model = m.provider.Model m.currentSession.Model = m.provider.Model
// Auto-name if empty // Auto-name if empty
@@ -208,11 +210,7 @@ func (m *Manager) Info() string {
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel()) fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking()) fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking())
fmt.Fprintf(&b, "providers: %s\n", strings.Join(m.providerNames(), ", ")) fmt.Fprintf(&b, "providers: %s\n", strings.Join(m.providerNames(), ", "))
if len(m.provider.Models) > 0 { fmt.Fprintf(&b, "models: %s\n", strings.Join(m.provider.Models, ", "))
fmt.Fprintf(&b, "models: %s\n", strings.Join(m.provider.Models, ", "))
} else {
b.WriteString("models: any model name is accepted for this provider\n")
}
fmt.Fprintf(&b, "thinking levels: %s\n", strings.Join(config.ThinkingLevels(), ", ")) fmt.Fprintf(&b, "thinking levels: %s\n", strings.Join(config.ThinkingLevels(), ", "))
b.WriteString("usage: /model provider <name> | /model model <name> | /model thinking <level>") b.WriteString("usage: /model provider <name> | /model model <name> | /model thinking <level>")
return b.String() return b.String()
@@ -234,10 +232,10 @@ func (m *Manager) SetModel(model string) (string, error) {
if model == "" { if model == "" {
return "", fmt.Errorf("model is required") return "", fmt.Errorf("model is required")
} }
if len(m.provider.Models) > 0 && !contains(m.provider.Models, model) { if !contains(m.provider.Models, model) {
return "", fmt.Errorf("model %q is not configured for provider %q; available: %s", model, m.providerName, strings.Join(m.provider.Models, ", ")) return "", fmt.Errorf("model %q is not configured for provider %q; available: %s", model, m.providerName, strings.Join(m.provider.Models, ", "))
} }
m.provider.Model = model m.provider = m.provider.WithModel(model)
m.cfg.Providers[m.providerName] = m.provider m.cfg.Providers[m.providerName] = m.provider
m.syncAgent(nil) m.syncAgent(nil)
return "Switched model.\n" + m.summary(), nil return "Switched model.\n" + m.summary(), nil
@@ -262,14 +260,22 @@ func (m *Manager) summary() string {
return fmt.Sprintf("provider: %s\nmodel: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentThinking()) return fmt.Sprintf("provider: %s\nmodel: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentThinking())
} }
func (m *Manager) maxContextTokens() int {
if m.cfg != nil && m.cfg.Agent.MaxContextTokens > 0 {
return m.cfg.Agent.MaxContextTokens
}
return m.provider.ContextTokens
}
func (m *Manager) syncAgent(provider llm.Provider) { func (m *Manager) syncAgent(provider llm.Provider) {
m.agent.SetRuntime(agent.RuntimeOptions{ m.agent.SetRuntime(agent.RuntimeOptions{
Provider: provider, Provider: provider,
ProviderName: m.providerName, ProviderName: m.providerName,
Model: m.provider.Model, Model: m.provider.Model,
Thinking: m.provider.Thinking, Thinking: m.provider.Thinking,
ThinkingParam: m.provider.ThinkingParam, ThinkingParam: m.provider.ThinkingParam,
ThinkingEnabled: m.provider.ThinkingConfigured, ThinkingEnabled: m.provider.ThinkingConfigured,
MaxContextTokens: m.maxContextTokens(),
}) })
} }
+67
View File
@@ -116,6 +116,7 @@ func TestManagerSaveAndResume(t *testing.T) {
manager.newSession() manager.newSession()
var out strings.Builder var out strings.Builder
assistant.RunTurn(context.Background(), "hello world", &out, nil) assistant.RunTurn(context.Background(), "hello world", &out, nil)
assistant.SetLastUsage(&llm.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15})
// Save // Save
if err := manager.Save(); err != nil { if err := manager.Save(); err != nil {
@@ -153,6 +154,9 @@ func TestManagerSaveAndResume(t *testing.T) {
if !hasUser { if !hasUser {
t.Fatal("resumed session should contain original user message") t.Fatal("resumed session should contain original user message")
} }
if usage := assistant.LastUsage(); usage == nil || usage.TotalTokens != 15 {
t.Fatalf("resumed usage = %#v", usage)
}
} }
func TestManagerRename(t *testing.T) { func TestManagerRename(t *testing.T) {
@@ -264,3 +268,66 @@ func TestManagerResumeLatest(t *testing.T) {
t.Fatalf("expected %s, got %s", id1, manager.CurrentSessionID()) t.Fatalf("expected %s, got %s", id1, manager.CurrentSessionID())
} }
} }
type compactingProvider struct {
summaryCalls int
answerCalls int
}
func (p *compactingProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
if len(req.Messages) == 2 && req.Messages[0].Role == llm.RoleSystem {
p.summaryCalls++
return emit(llm.StreamEvent{Content: "summary"})
}
p.answerCalls++
return emit(llm.StreamEvent{Content: "answer"})
}
func TestManagerUsesModelContextForAgentCompaction(t *testing.T) {
provider := &compactingProvider{}
cfg := &config.Config{Providers: map[string]config.ProviderConfig{
"one": {
Name: "one",
BaseURL: "https://one.example.com",
APIKey: "sk-test",
Model: "model-a",
Models: []string{"model-a"},
ModelConfigs: map[string]config.ModelConfig{"model-a": {ID: "model-a", Thinking: "none", ContextTokens: 1}},
ContextTokens: 1,
},
}}
assistant := agent.New(agent.Options{Provider: provider, ProviderName: "one", Model: "model-a"})
store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
if err != nil {
t.Fatal(err)
}
manager, err := NewManager(cfg, assistant, nil, store)
if err != nil {
t.Fatal(err)
}
if _, err := manager.SetModel("model-a"); err != nil {
t.Fatal(err)
}
assistant.SetMessages([]llm.Message{
{Role: llm.RoleSystem, Content: "system"},
{Role: llm.RoleUser, Content: "old question 1"},
{Role: llm.RoleAssistant, Content: "old answer 1"},
{Role: llm.RoleUser, Content: "old question 2"},
{Role: llm.RoleAssistant, Content: "old answer 2"},
{Role: llm.RoleUser, Content: "recent question 1"},
{Role: llm.RoleAssistant, Content: "recent answer 1"},
{Role: llm.RoleUser, Content: "recent question 2"},
{Role: llm.RoleAssistant, Content: "recent answer 2"},
})
var out strings.Builder
if err := assistant.RunTurn(context.Background(), "new question", &out, nil); err != nil {
t.Fatal(err)
}
if provider.summaryCalls != 1 {
t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
}
if provider.answerCalls != 1 {
t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
}
}
+1
View File
@@ -23,6 +23,7 @@ type Session struct {
Provider string `json:"provider"` Provider string `json:"provider"`
Model string `json:"model"` Model string `json:"model"`
Messages []llm.Message `json:"messages"` Messages []llm.Message `json:"messages"`
LastUsage *llm.Usage `json:"last_usage,omitempty"`
} }
// SessionMeta is a lightweight summary for listing. // SessionMeta is a lightweight summary for listing.
+4
View File
@@ -30,6 +30,7 @@ func TestSaveAndLoad(t *testing.T) {
{Role: llm.RoleUser, Content: "hello"}, {Role: llm.RoleUser, Content: "hello"},
{Role: llm.RoleAssistant, Content: "hi there"}, {Role: llm.RoleAssistant, Content: "hi there"},
}, },
LastUsage: &llm.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15},
} }
if err := store.Save(sess); err != nil { if err := store.Save(sess); err != nil {
t.Fatalf("save: %v", err) t.Fatalf("save: %v", err)
@@ -51,6 +52,9 @@ func TestSaveAndLoad(t *testing.T) {
if loaded.Messages[1].Content != "hello" { if loaded.Messages[1].Content != "hello" {
t.Errorf("Messages[1].Content = %q, want %q", loaded.Messages[1].Content, "hello") t.Errorf("Messages[1].Content = %q, want %q", loaded.Messages[1].Content, "hello")
} }
if loaded.LastUsage == nil || loaded.LastUsage.TotalTokens != 15 {
t.Errorf("LastUsage = %#v", loaded.LastUsage)
}
if loaded.UpdatedAt.IsZero() { if loaded.UpdatedAt.IsZero() {
t.Error("UpdatedAt should be set") t.Error("UpdatedAt should be set")
} }
+219
View File
@@ -0,0 +1,219 @@
package tools
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"go/ast"
"go/parser"
"go/printer"
"go/token"
"io/fs"
"os"
"path/filepath"
"sort"
"strings"
"agentu/internal/llm"
)
const defaultCodeSymbolFileLimit = 50
type CodeSymbolsTool struct {
workingDir string
}
func NewCodeSymbolsTool(workingDir string) *CodeSymbolsTool {
return &CodeSymbolsTool{workingDir: workingDir}
}
func (t *CodeSymbolsTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "code_symbols",
Description: "List Go code symbols (packages, imports, types, functions, methods) for a file or directory without reading full file contents.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Go file or directory to inspect. Defaults to the configured agent working directory."},
"max_files": {"type": "integer", "description": "Maximum number of Go files to inspect when path is a directory. Defaults to 50."}
},
"additionalProperties": false
}`),
},
}
}
type codeSymbolsArgs struct {
Path string `json:"path"`
MaxFiles int `json:"max_files"`
}
func (t *CodeSymbolsTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args codeSymbolsArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
path, err := resolvePath(t.workingDir, args.Path)
if err != nil {
return "", err
}
limit := args.MaxFiles
if limit <= 0 {
limit = defaultCodeSymbolFileLimit
}
files, err := goFiles(ctx, path, limit)
if err != nil {
return "", err
}
if len(files) == 0 {
return "no Go files found", nil
}
fset := token.NewFileSet()
var out strings.Builder
for i, file := range files {
if ctxErr := ctx.Err(); ctxErr != nil {
return FormatResult(out.String()), ctxErr
}
if i > 0 {
out.WriteString("\n")
}
if err := appendGoSymbols(&out, fset, file); err != nil {
fmt.Fprintf(&out, "%s\nerror: %v\n", file, err)
}
if out.Len() > MaxToolOutputBytes {
break
}
}
return FormatResult(out.String()), nil
}
func goFiles(ctx context.Context, path string, limit int) ([]string, error) {
info, err := os.Stat(path)
if err != nil {
return nil, err
}
if !info.IsDir() {
if strings.HasSuffix(path, ".go") {
return []string{path}, nil
}
return nil, fmt.Errorf("path is not a Go file: %s", path)
}
var files []string
stop := errors.New("file limit reached")
err = filepath.WalkDir(path, func(current string, d fs.DirEntry, err error) error {
if err != nil {
return nil
}
if ctxErr := ctx.Err(); ctxErr != nil {
return ctxErr
}
if d.IsDir() {
switch d.Name() {
case ".git", "vendor", "node_modules":
return filepath.SkipDir
}
return nil
}
if strings.HasSuffix(current, ".go") {
files = append(files, current)
if len(files) >= limit {
return stop
}
}
return nil
})
if err != nil && !errors.Is(err, stop) {
return nil, err
}
sort.Strings(files)
return files, nil
}
func appendGoSymbols(out *strings.Builder, fset *token.FileSet, path string) error {
file, err := parser.ParseFile(fset, path, nil, parser.ImportsOnly|parser.ParseComments)
if err != nil {
return err
}
full, err := parser.ParseFile(fset, path, nil, 0)
if err != nil {
return err
}
fmt.Fprintf(out, "%s\n", path)
fmt.Fprintf(out, " package %s\n", full.Name.Name)
if len(file.Imports) > 0 {
imports := make([]string, 0, len(file.Imports))
for _, imp := range file.Imports {
imports = append(imports, strings.Trim(imp.Path.Value, "\""))
}
fmt.Fprintf(out, " imports: %s\n", strings.Join(imports, ", "))
}
var types, funcs, methods []string
for _, decl := range full.Decls {
switch d := decl.(type) {
case *ast.GenDecl:
if d.Tok != token.TYPE {
continue
}
for _, spec := range d.Specs {
if ts, ok := spec.(*ast.TypeSpec); ok {
kind := typeKind(ts.Type)
types = append(types, fmt.Sprintf("%d: type %s %s", fset.Position(ts.Pos()).Line, ts.Name.Name, kind))
}
}
case *ast.FuncDecl:
line := fset.Position(d.Pos()).Line
if d.Recv == nil {
funcs = append(funcs, fmt.Sprintf("%d: func %s", line, d.Name.Name))
continue
}
methods = append(methods, fmt.Sprintf("%d: method %s.%s", line, receiverName(d.Recv), d.Name.Name))
}
}
appendSymbolGroup(out, "types", types)
appendSymbolGroup(out, "funcs", funcs)
appendSymbolGroup(out, "methods", methods)
return nil
}
func appendSymbolGroup(out *strings.Builder, title string, values []string) {
if len(values) == 0 {
return
}
fmt.Fprintf(out, " %s:\n", title)
for _, value := range values {
fmt.Fprintf(out, " - %s\n", value)
}
}
func typeKind(expr ast.Expr) string {
switch expr.(type) {
case *ast.StructType:
return "struct"
case *ast.InterfaceType:
return "interface"
case *ast.FuncType:
return "func"
default:
return "alias"
}
}
func receiverName(recv *ast.FieldList) string {
if recv == nil || len(recv.List) == 0 {
return "?"
}
var b bytes.Buffer
if err := printer.Fprint(&b, token.NewFileSet(), recv.List[0].Type); err != nil {
return "?"
}
return b.String()
}
+61 -10
View File
@@ -223,7 +223,9 @@ func (t *FileSearchTool) Definition() llm.Tool {
"properties": { "properties": {
"query": {"type": "string", "description": "Text or regular expression to search for."}, "query": {"type": "string", "description": "Text or regular expression to search for."},
"path": {"type": "string", "description": "Optional directory or file to search. Defaults to the configured agent working directory."}, "path": {"type": "string", "description": "Optional directory or file to search. Defaults to the configured agent working directory."},
"glob": {"type": "string", "description": "Optional file glob, for example '*.go'."} "glob": {"type": "string", "description": "Optional file glob, for example '*.go'."},
"context_lines": {"type": "integer", "description": "Optional number of context lines around each match. Defaults to 0."},
"max_results": {"type": "integer", "description": "Optional maximum matches to return. Defaults to 100."}
}, },
"required": ["query"], "required": ["query"],
"additionalProperties": false "additionalProperties": false
@@ -233,9 +235,11 @@ func (t *FileSearchTool) Definition() llm.Tool {
} }
type fileSearchArgs struct { type fileSearchArgs struct {
Query string `json:"query"` Query string `json:"query"`
Path string `json:"path"` Path string `json:"path"`
Glob string `json:"glob"` Glob string `json:"glob"`
ContextLines int `json:"context_lines"`
MaxResults int `json:"max_results"`
} }
func (t *FileSearchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) { func (t *FileSearchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
@@ -256,8 +260,31 @@ func (t *FileSearchTool) Execute(ctx context.Context, raw json.RawMessage) (stri
return t.walkSearch(ctx, args, path) return t.walkSearch(ctx, args, path)
} }
func normalizedSearchOptions(args fileSearchArgs) (contextLines int, maxResults int) {
contextLines = args.ContextLines
if contextLines < 0 {
contextLines = 0
}
if contextLines > 5 {
contextLines = 5
}
maxResults = args.MaxResults
if maxResults <= 0 {
maxResults = 100
}
if maxResults > 1000 {
maxResults = 1000
}
return contextLines, maxResults
}
func (t *FileSearchTool) ripgrep(ctx context.Context, args fileSearchArgs, path string) (string, error) { func (t *FileSearchTool) ripgrep(ctx context.Context, args fileSearchArgs, path string) (string, error) {
rgArgs := []string{"--line-number", "--color", "never"} rgArgs := []string{"--line-number", "--color", "never"}
contextLines, maxResults := normalizedSearchOptions(args)
if contextLines > 0 {
rgArgs = append(rgArgs, "--context", fmt.Sprint(contextLines))
}
rgArgs = append(rgArgs, "--max-count", fmt.Sprint(maxResults))
if args.Glob != "" { if args.Glob != "" {
rgArgs = append(rgArgs, "--glob", args.Glob) rgArgs = append(rgArgs, "--glob", args.Glob)
} }
@@ -282,7 +309,10 @@ func (t *FileSearchTool) ripgrep(ctx context.Context, args fileSearchArgs, path
} }
func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, root string) (string, error) { func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, root string) (string, error) {
contextLines, maxResults := normalizedSearchOptions(args)
var out bytes.Buffer var out bytes.Buffer
matchCount := 0
stop := errors.New("result limit reached")
err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error { err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
if err != nil { if err != nil {
return nil return nil
@@ -312,16 +342,22 @@ func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, ro
} }
lines := strings.Split(string(data), "\n") lines := strings.Split(string(data), "\n")
for i, line := range lines { for i, line := range lines {
if strings.Contains(line, args.Query) { if !strings.Contains(line, args.Query) {
fmt.Fprintf(&out, "%s:%d:%s\n", path, i+1, line) continue
if out.Len() > MaxToolOutputBytes { }
return errors.New("output limit reached") matchCount++
} appendSearchMatch(&out, path, lines, i, contextLines)
if matchCount >= maxResults {
fmt.Fprintf(&out, "[stopped after %d matches; narrow path/query or raise max_results]\n", maxResults)
return stop
}
if out.Len() > MaxToolOutputBytes {
return stop
} }
} }
return nil return nil
}) })
if err != nil && err.Error() != "output limit reached" { if err != nil && !errors.Is(err, stop) {
return "", err return "", err
} }
if out.Len() == 0 { if out.Len() == 0 {
@@ -329,3 +365,18 @@ func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, ro
} }
return FormatResult(out.String()), nil return FormatResult(out.String()), nil
} }
func appendSearchMatch(out *bytes.Buffer, path string, lines []string, matchIndex int, contextLines int) {
start := max(0, matchIndex-contextLines)
end := min(len(lines)-1, matchIndex+contextLines)
if contextLines > 0 {
fmt.Fprintf(out, "-- %s:%d --\n", path, matchIndex+1)
}
for i := start; i <= end; i++ {
marker := ":"
if i == matchIndex {
marker = "*"
}
fmt.Fprintf(out, "%s%s%d:%s\n", path, marker, i+1, lines[i])
}
}
+19 -2
View File
@@ -1,7 +1,9 @@
package tools package tools
import ( import (
"fmt"
"path/filepath" "path/filepath"
"strings"
) )
func resolvePath(baseDir, path string) (string, error) { func resolvePath(baseDir, path string) (string, error) {
@@ -9,7 +11,22 @@ func resolvePath(baseDir, path string) (string, error) {
path = "." path = "."
} }
if filepath.IsAbs(path) { if filepath.IsAbs(path) {
return filepath.Clean(path), nil return "", fmt.Errorf("absolute paths are not allowed; use a path relative to the project directory")
} }
return filepath.Abs(filepath.Join(baseDir, path)) resolved := filepath.Clean(filepath.Join(baseDir, path))
if !isWithinDir(resolved, baseDir) {
return "", fmt.Errorf("path escapes the project directory: %s", path)
}
return resolved, nil
}
func isWithinDir(path, dir string) bool {
rel, err := filepath.Rel(dir, path)
if err != nil {
return false
}
if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
return false
}
return true
} }
+42 -5
View File
@@ -5,12 +5,15 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"sort" "sort"
"strings"
"agentu/internal/llm" "agentu/internal/llm"
) )
const MaxToolOutputBytes = 64 * 1024 const MaxToolOutputBytes = 64 * 1024
const truncationNoticeBudget = 160
type Tool interface { type Tool interface {
Definition() llm.Tool Definition() llm.Tool
Execute(ctx context.Context, args json.RawMessage) (string, error) Execute(ctx context.Context, args json.RawMessage) (string, error)
@@ -62,11 +65,13 @@ func ReadOnlyBuiltins(workingDir string) *Registry {
fileRead := NewFileReadTool(workingDir) fileRead := NewFileReadTool(workingDir)
fileList := NewFileListTool(workingDir) fileList := NewFileListTool(workingDir)
fileSearch := NewFileSearchTool(workingDir) fileSearch := NewFileSearchTool(workingDir)
codeSymbols := NewCodeSymbolsTool(workingDir)
registry := NewRegistry(fileRead, fileList, fileSearch) registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols)
registry.RegisterAlias("file.read", fileRead) registry.RegisterAlias("file.read", fileRead)
registry.RegisterAlias("file.list", fileList) registry.RegisterAlias("file.list", fileList)
registry.RegisterAlias("file.search", fileSearch) registry.RegisterAlias("file.search", fileSearch)
registry.RegisterAlias("code.symbols", codeSymbols)
return registry return registry
} }
@@ -83,8 +88,8 @@ func Builtins(workingDir string) *Registry {
func AgentInstructions(workingDir string, yolo bool) string { func AgentInstructions(workingDir string, yolo bool) string {
searchInstructions := webSearchInstructions() searchInstructions := webSearchInstructions()
readOnlyTools := "file_list, file_read, file_search, web_search, web_fetch" readOnlyTools := "file_list, file_read, file_search, code_symbols, web_search, web_fetch"
yoloTools := "file_list, file_read, file_search, file_write, shell_run, web_search, web_fetch" yoloTools := "file_list, file_read, file_search, code_symbols, file_write, shell_run, web_search, web_fetch"
if yolo { if yolo {
return fmt.Sprintf(`Local project context: return fmt.Sprintf(`Local project context:
- The project working directory is %q. - The project working directory is %q.
@@ -92,6 +97,7 @@ func AgentInstructions(workingDir string, yolo bool) string {
- Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project or command context. - Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project or command context.
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files. - Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files. - When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
- When the user asks to run, execute, test, check, debug, or inspect a shell command, use shell_run. This includes local commands written in backticks, such as go, make, git, pwd, ls, ps, and similar terminal commands. - When the user asks to run, execute, test, check, debug, or inspect a shell command, use shell_run. This includes local commands written in backticks, such as go, make, git, pwd, ls, ps, and similar terminal commands.
- Do not use file tools as a substitute for a command execution request. - Do not use file tools as a substitute for a command execution request.
- Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions. - Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions.
@@ -105,6 +111,7 @@ func AgentInstructions(workingDir string, yolo bool) string {
- Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project context. - Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project context.
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files. - Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files. - When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
- Shell command execution is disabled in this mode. When the user asks to run, execute, test, check, debug, or inspect a shell command, do not use file tools as a substitute; explain that they need to restart agentu with --yolo to enable shell_run. - Shell command execution is disabled in this mode. When the user asks to run, execute, test, check, debug, or inspect a shell command, do not use file tools as a substitute; explain that they need to restart agentu with --yolo to enable shell_run.
- %s - %s
- Tool paths should normally be relative to the project working directory. - Tool paths should normally be relative to the project working directory.
@@ -120,10 +127,40 @@ func JSONSchema(schema string) json.RawMessage {
} }
func FormatResult(output string) string { func FormatResult(output string) string {
if len(output) <= MaxToolOutputBytes { return FormatResultLimit(output, MaxToolOutputBytes)
}
func FormatResultLimit(output string, limit int) string {
if limit <= 0 || len(output) <= limit {
return output return output
} }
return output[:MaxToolOutputBytes] + fmt.Sprintf("\n\n[truncated after %d bytes]", MaxToolOutputBytes) if limit <= truncationNoticeBudget {
return output[:limit]
}
notice := fmt.Sprintf("\n\n[truncated: showing head and tail; omitted %d bytes]\n\n", len(output)-limit)
keep := limit - len(notice)
if keep <= 0 {
return output[:limit]
}
headLen := keep / 2
tailLen := keep - headLen
head := trimHeadAtLine(output[:headLen])
tail := trimTailAtLine(output[len(output)-tailLen:])
return head + notice + tail
}
func trimHeadAtLine(value string) string {
if idx := strings.LastIndexByte(value, '\n'); idx > 0 {
return value[:idx+1]
}
return value
}
func trimTailAtLine(value string) string {
if idx := strings.IndexByte(value, '\n'); idx >= 0 && idx+1 < len(value) {
return value[idx+1:]
}
return value
} }
func FormatError(err error) string { func FormatError(err error) string {
+122 -3
View File
@@ -43,6 +43,41 @@ func TestFileTools(t *testing.T) {
} }
} }
func TestFileSearchSupportsContextAndMaxResults(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
content := strings.Join([]string{"before", "needle one", "middle", "needle two", "after"}, "\n")
if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
t.Fatal(err)
}
search := NewFileSearchTool(dir)
out, err := search.walkSearch(context.Background(), fileSearchArgs{Query: "needle", ContextLines: 1, MaxResults: 1}, dir)
if err != nil {
t.Fatal(err)
}
for _, want := range []string{"-- ", "before", "needle one", "middle", "stopped after 1 matches"} {
if !strings.Contains(out, want) {
t.Fatalf("search output missing %q:\n%s", want, out)
}
}
if strings.Contains(out, "needle two") {
t.Fatalf("search should stop after max_results:\n%s", out)
}
}
func TestFormatResultKeepsHeadAndTail(t *testing.T) {
input := strings.Join([]string{"head-1", "head-2", strings.Repeat("x", 220), "tail-1", "tail-2"}, "\n")
out := FormatResultLimit(input, 220)
for _, want := range []string{"head-1", "truncated: showing head and tail", "tail-2"} {
if !strings.Contains(out, want) {
t.Fatalf("formatted output missing %q:\n%s", want, out)
}
}
if len(out) > 220 {
t.Fatalf("formatted output len = %d, want <= 220", len(out))
}
}
func TestShellRunTool(t *testing.T) { func TestShellRunTool(t *testing.T) {
dir := t.TempDir() dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "x.txt"), []byte("x"), 0o644); err != nil { if err := os.WriteFile(filepath.Join(dir, "x.txt"), []byte("x"), 0o644); err != nil {
@@ -69,12 +104,12 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
} }
} }
for _, want := range []string{"file_list", "file_read", "file_search"} { for _, want := range []string{"file_list", "file_read", "file_search", "code_symbols"} {
if !slices.Contains(names, want) { if !slices.Contains(names, want) {
t.Fatalf("definitions missing %s: %#v", want, names) t.Fatalf("definitions missing %s: %#v", want, names)
} }
} }
for _, alias := range []string{"file.list", "file.read", "file.search"} { for _, alias := range []string{"file.list", "file.read", "file.search", "code.symbols"} {
if _, ok := registry.Get(alias); !ok { if _, ok := registry.Get(alias); !ok {
t.Fatalf("alias missing: %s", alias) t.Fatalf("alias missing: %s", alias)
} }
@@ -84,6 +119,34 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
} }
} }
func TestCodeSymbolsToolListsGoSymbols(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "sample.go")
if err := os.WriteFile(path, []byte(`package sample
import "context"
type Worker struct{}
type Runner interface{ Run(context.Context) error }
func NewWorker() *Worker { return &Worker{} }
func (w *Worker) Run(ctx context.Context) error { return nil }
`), 0o644); err != nil {
t.Fatal(err)
}
tool := NewCodeSymbolsTool(dir)
out, err := tool.Execute(context.Background(), json.RawMessage(`{"path":"sample.go"}`))
if err != nil {
t.Fatal(err)
}
for _, want := range []string{"package sample", "imports: context", "type Worker struct", "type Runner interface", "func NewWorker", "method *Worker.Run"} {
if !strings.Contains(out, want) {
t.Fatalf("symbols output missing %q:\n%s", want, out)
}
}
}
func TestRegistryCanExposeWebSearch(t *testing.T) { func TestRegistryCanExposeWebSearch(t *testing.T) {
registry := ReadOnlyBuiltins(t.TempDir()) registry := ReadOnlyBuiltins(t.TempDir())
registry.Register(NewBuiltinSearchTool(nil)) registry.Register(NewBuiltinSearchTool(nil))
@@ -132,9 +195,65 @@ func TestAgentInstructionsExplainCommandModes(t *testing.T) {
func TestAgentInstructionsExplainWebSearch(t *testing.T) { func TestAgentInstructionsExplainWebSearch(t *testing.T) {
withSearch := AgentInstructions("/tmp/project", false) withSearch := AgentInstructions("/tmp/project", false)
for _, want := range []string{"use web_search", "Use web_fetch", "latest information", "source URLs", "Available tools: file_list, file_read, file_search, web_search, web_fetch"} { for _, want := range []string{"use web_search", "Use web_fetch", "latest information", "source URLs", "Available tools: file_list, file_read, file_search, code_symbols, web_search, web_fetch"} {
if !strings.Contains(withSearch, want) { if !strings.Contains(withSearch, want) {
t.Fatalf("web search instructions missing %q:\n%s", want, withSearch) t.Fatalf("web search instructions missing %q:\n%s", want, withSearch)
} }
} }
} }
func TestResolvePathRejectsAbsolutePaths(t *testing.T) {
base := t.TempDir()
_, err := resolvePath(base, "/etc/passwd")
if err == nil {
t.Fatal("expected error for absolute path")
}
if !strings.Contains(err.Error(), "absolute paths are not allowed") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestResolvePathRejectsEscapingBaseDir(t *testing.T) {
base := filepath.Join(t.TempDir(), "project")
if err := os.MkdirAll(base, 0o755); err != nil {
t.Fatal(err)
}
_, err := resolvePath(base, "../../../etc/passwd")
if err == nil {
t.Fatal("expected error for path escaping base dir")
}
if !strings.Contains(err.Error(), "escapes the project directory") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestResolvePathAllowsSubdirTraversal(t *testing.T) {
base := t.TempDir()
if err := os.MkdirAll(filepath.Join(base, "a", "b"), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(base, "a", "target.txt"), []byte("ok"), 0o644); err != nil {
t.Fatal(err)
}
// "a/b/../target.txt" resolves to "a/target.txt" which is within base.
resolved, err := resolvePath(base, "a/b/../target.txt")
if err != nil {
t.Fatal(err)
}
want := filepath.Join(base, "a", "target.txt")
if resolved != want {
t.Fatalf("resolved = %q, want %q", resolved, want)
}
}
func TestResolvePathRejectsDotDotThatEscapes(t *testing.T) {
base := t.TempDir()
sub := filepath.Join(base, "a")
if err := os.MkdirAll(sub, 0o755); err != nil {
t.Fatal(err)
}
_, err := resolvePath(sub, "..")
if err == nil {
t.Fatal("expected error for path escaping via ..")
}
}
+111 -25
View File
@@ -28,6 +28,7 @@ type Options struct {
Yolo bool Yolo bool
ThemeMode ThemeMode ThemeMode ThemeMode
ModelManager ModelManager ModelManager ModelManager
WorkingDir string
} }
type ModelManager interface { type ModelManager interface {
@@ -69,13 +70,14 @@ type message struct {
} }
type model struct { type model struct {
ctx context.Context ctx context.Context
agent *agent.Agent agent *agent.Agent
modelName string modelName string
yolo bool yolo bool
models ModelManager models ModelManager
theme theme workingDir string
styles styles theme theme
styles styles
viewport viewport.Model viewport viewport.Model
input textarea.Model input textarea.Model
@@ -84,11 +86,12 @@ type model struct {
width int width int
height int height int
messages []message messages []message
status string status string
running bool contextUsage string
cancel context.CancelFunc running bool
events <-chan tea.Msg cancel context.CancelFunc
events <-chan tea.Msg
// Session picker state // Session picker state
picking bool picking bool
@@ -119,6 +122,7 @@ type slashCommand struct {
var slashCommands = []slashCommand{ var slashCommands = []slashCommand{
{Name: "/clear", Description: "reset context"}, {Name: "/clear", Description: "reset context"},
{Name: "/compact", Description: "compress context"},
{Name: "/exit", Description: "quit"}, {Name: "/exit", Description: "quit"},
{Name: "/quit", Description: "quit"}, {Name: "/quit", Description: "quit"},
{Name: "/model", Description: "model/provider/thinking"}, {Name: "/model", Description: "model/provider/thinking"},
@@ -126,6 +130,9 @@ var slashCommands = []slashCommand{
{Name: "/resume", Description: "resume session"}, {Name: "/resume", Description: "resume session"},
{Name: "/rename", Description: "rename session"}, {Name: "/rename", Description: "rename session"},
{Name: "/new", Description: "new session"}, {Name: "/new", Description: "new session"},
{Name: "/status", Description: "show changed files"},
{Name: "/diff", Description: "show unstaged diff"},
{Name: "/test", Description: "run project tests"},
} }
func newModel(ctx context.Context, assistant *agent.Agent, opts Options) model { func newModel(ctx context.Context, assistant *agent.Agent, opts Options) model {
@@ -154,19 +161,22 @@ func newModel(ctx context.Context, assistant *agent.Agent, opts Options) model {
vp.Style = st.Viewport vp.Style = st.Viewport
vp.SetContent("") vp.SetContent("")
return model{ m := model{
ctx: ctx, ctx: ctx,
agent: assistant, agent: assistant,
modelName: opts.ModelName, modelName: opts.ModelName,
yolo: opts.Yolo, yolo: opts.Yolo,
models: opts.ModelManager, models: opts.ModelManager,
theme: th, workingDir: opts.WorkingDir,
styles: st, theme: th,
viewport: vp, styles: st,
input: input, viewport: vp,
spinner: spin, input: input,
status: "ready", spinner: spin,
status: "ready",
} }
m.refreshContextUsage()
return m
} }
func (m model) Init() tea.Cmd { func (m model) Init() tea.Cmd {
@@ -272,6 +282,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
if m.models != nil { if m.models != nil {
_ = m.models.Save() _ = m.models.Save()
} }
m.refreshContextUsage()
m.resize() m.resize()
m.refreshViewport(true) m.refreshViewport(true)
cmds = append(cmds, textarea.Blink) cmds = append(cmds, textarea.Blink)
@@ -338,13 +349,14 @@ func (m *model) submit() (tea.Model, tea.Cmd) {
return *m, tea.Quit return *m, tea.Quit
case "/clear": case "/clear":
m.agent.Clear() m.agent.Clear()
m.refreshContextUsage()
m.messages = []message{{role: roleSystem, content: "Context cleared."}} m.messages = []message{{role: roleSystem, content: "Context cleared."}}
m.input.SetValue("") m.input.SetValue("")
m.afterInputChanged() m.afterInputChanged()
m.status = "ready" m.status = "ready"
m.refreshViewport(true) m.refreshViewport(true)
return *m, nil return *m, nil
case "/model": case "/model", "/compact":
return m.runLocalCommand(input) return m.runLocalCommand(input)
} }
@@ -541,6 +553,9 @@ func (m model) modelMetaView() string {
parts = append(parts, "thinking "+currentThinking) parts = append(parts, "thinking "+currentThinking)
} }
} }
if m.contextUsage != "" {
parts = append(parts, m.contextUsage)
}
parts = append([]string{modelName, mode}, parts...) parts = append([]string{modelName, mode}, parts...)
parts = append(parts, "theme "+string(m.theme.Mode)) parts = append(parts, "theme "+string(m.theme.Mode))
line := fitLine(strings.Join(parts, " · "), max(1, m.width-2)) line := fitLine(strings.Join(parts, " · "), max(1, m.width-2))
@@ -700,6 +715,58 @@ func (m model) modelInfo() string {
return fmt.Sprintf("model: %s\nmode: %s\ntheme: %s", m.modelName, mode, m.theme.Mode) return fmt.Sprintf("model: %s\nmode: %s\ntheme: %s", m.modelName, mode, m.theme.Mode)
} }
func (m *model) refreshContextUsage() {
m.contextUsage = contextUsageMeta(m.agent)
}
func contextUsageMeta(assistant *agent.Agent) string {
if assistant == nil {
return ""
}
info := assistant.RuntimeInfo()
if info.ContextTokens <= 0 && info.MaxContextTokens <= 0 {
return ""
}
used := formatTokenCount(info.ContextTokens)
if info.MaxContextTokens <= 0 {
return "ctx ~" + used
}
return fmt.Sprintf("ctx ~%s/%s (%s)", used, formatTokenCount(info.MaxContextTokens), contextPercent(info.ContextTokens, info.MaxContextTokens))
}
func contextPercent(used int, limit int) string {
if limit <= 0 {
return ""
}
if used <= 0 {
return "0%"
}
percent := used * 100 / limit
if percent == 0 {
return "<1%"
}
return fmt.Sprintf("%d%%", percent)
}
func formatTokenCount(tokens int) string {
if tokens < 0 {
tokens = 0
}
if tokens < 1000 {
return fmt.Sprintf("%d", tokens)
}
if tokens < 1_000_000 {
if tokens%1000 == 0 {
return fmt.Sprintf("%dk", tokens/1000)
}
return fmt.Sprintf("%.1fk", float64(tokens)/1000)
}
if tokens%1_000_000 == 0 {
return fmt.Sprintf("%dm", tokens/1_000_000)
}
return fmt.Sprintf("%.1fm", float64(tokens)/1_000_000)
}
func fitLine(text string, width int) string { func fitLine(text string, width int) string {
width = max(1, width) width = max(1, width)
if lipgloss.Width(text) <= width { if lipgloss.Width(text) <= width {
@@ -741,6 +808,7 @@ func (m *model) runLocalCommand(input string) (tea.Model, tea.Cmd) {
m.status = status m.status = status
// Rebuild display from agent history (matters after /resume and /new) // Rebuild display from agent history (matters after /resume and /new)
m.syncMessages() m.syncMessages()
m.refreshContextUsage()
// Append command output after synced history // Append command output after synced history
m.messages = append(m.messages, message{role: role, content: output}) m.messages = append(m.messages, message{role: role, content: output})
m.refreshViewport(true) m.refreshViewport(true)
@@ -753,6 +821,8 @@ func (m model) handleLocalCommand(input string) (string, error) {
return "", fmt.Errorf("empty command") return "", fmt.Errorf("empty command")
} }
switch fields[0] { switch fields[0] {
case "/compact":
return m.handleCompactCommand()
case "/model": case "/model":
return m.handleModelCommand(fields[1:]) return m.handleModelCommand(fields[1:])
case "/sessions": case "/sessions":
@@ -763,11 +833,27 @@ func (m model) handleLocalCommand(input string) (string, error) {
return m.handleRenameCommand(fields[1:]) return m.handleRenameCommand(fields[1:])
case "/new": case "/new":
return m.handleNewCommand() return m.handleNewCommand()
case "/status":
return m.handleStatusCommand(fields[1:])
case "/diff":
return m.handleDiffCommand(fields[1:])
case "/test":
return m.handleTestCommand(fields[1:])
default: default:
return "", fmt.Errorf("unknown command: %s", fields[0]) return "", fmt.Errorf("unknown command: %s", fields[0])
} }
} }
func (m model) handleCompactCommand() (string, error) {
if m.agent == nil {
return "", fmt.Errorf("agent is not configured")
}
if err := m.agent.Compact(m.ctx); err != nil {
return "", err
}
return "Context compacted.", nil
}
func (m model) handleModelCommand(args []string) (string, error) { func (m model) handleModelCommand(args []string) (string, error) {
if m.models == nil { if m.models == nil {
if len(args) > 0 { if len(args) > 0 {
+116 -1
View File
@@ -2,9 +2,15 @@ package tui
import ( import (
"context" "context"
"os"
"os/exec"
"path/filepath"
"strings" "strings"
"testing" "testing"
"agentu/internal/agent"
"agentu/internal/llm"
tea "github.com/charmbracelet/bubbletea" tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss" "github.com/charmbracelet/lipgloss"
) )
@@ -82,6 +88,20 @@ func TestModelMetaRendersBelowInput(t *testing.T) {
} }
} }
func TestModelMetaRendersContextUsage(t *testing.T) {
assistant := agent.New(agent.Options{Model: "test-model", MaxContextTokens: 1000})
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("x", 800)}})
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 120, Height: 30})
rendered := next.(model).modelMetaView()
for _, want := range []string{"ctx ~", "/1k", "%"} {
if !strings.Contains(rendered, want) {
t.Fatalf("model meta missing %q:\n%s", want, rendered)
}
}
}
func TestInitialViewDoesNotRenderTopNote(t *testing.T) { func TestInitialViewDoesNotRenderTopNote(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"}) m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30}) next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
@@ -264,13 +284,108 @@ func TestSlashCommandSuggestions(t *testing.T) {
m.afterInputChanged() m.afterInputChanged()
rendered := m.View() rendered := m.View()
for _, want := range []string{"/clear reset context", "/exit quit", "/model model/provider/thinking"} { for _, want := range []string{"/clear reset context", "/compact compress context", "/exit quit", "/model model/provider/thinking"} {
if !strings.Contains(rendered, want) { if !strings.Contains(rendered, want) {
t.Fatalf("rendered suggestions missing %q:\n%s", want, rendered) t.Fatalf("rendered suggestions missing %q:\n%s", want, rendered)
} }
} }
} }
type tuiCompactProvider struct{}
func (tuiCompactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
return emit(llm.StreamEvent{Content: "summary from tui"})
}
func TestCompactCommand(t *testing.T) {
assistant := agent.New(agent.Options{Provider: tuiCompactProvider{}, Model: "test", SystemPrompt: "system"})
assistant.SetMessages([]llm.Message{
{Role: llm.RoleSystem, Content: "system"},
{Role: llm.RoleUser, Content: "old question 1"},
{Role: llm.RoleAssistant, Content: "old answer 1"},
{Role: llm.RoleUser, Content: "old question 2"},
{Role: llm.RoleAssistant, Content: "old answer 2"},
{Role: llm.RoleUser, Content: "recent question 1"},
{Role: llm.RoleAssistant, Content: "recent answer 1"},
{Role: llm.RoleUser, Content: "recent question 2"},
{Role: llm.RoleAssistant, Content: "recent answer 2"},
})
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
out, err := m.handleLocalCommand("/compact")
if err != nil {
t.Fatal(err)
}
if out != "Context compacted." {
t.Fatalf("output = %q", out)
}
if !strings.Contains(messageContents(assistant.Messages()), "[context summary]") {
t.Fatalf("messages missing summary: %#v", assistant.Messages())
}
}
func messageContents(messages []llm.Message) string {
var b strings.Builder
for _, msg := range messages {
b.WriteString(msg.Content)
b.WriteByte('\n')
}
return b.String()
}
func TestWorkflowStatusDiffAndTestCommands(t *testing.T) {
if _, err := exec.LookPath("git"); err != nil {
t.Skip("git not installed")
}
dir := t.TempDir()
runInDir(t, dir, "git", "init")
runInDir(t, dir, "git", "config", "user.email", "agentu@example.test")
runInDir(t, dir, "git", "config", "user.name", "agentu")
if err := os.WriteFile(filepath.Join(dir, "go.mod"), []byte("module example.com/p1\n\ngo 1.26\n"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(dir, "main_test.go"), []byte("package p1\n\nimport \"testing\"\n\nfunc TestOK(t *testing.T) {}\n"), 0o644); err != nil {
t.Fatal(err)
}
runInDir(t, dir, "git", "add", ".")
runInDir(t, dir, "git", "commit", "-m", "init")
if err := os.WriteFile(filepath.Join(dir, "main_test.go"), []byte("package p1\n\nimport \"testing\"\n\nfunc TestOK(t *testing.T) { t.Log(\"changed\") }\n"), 0o644); err != nil {
t.Fatal(err)
}
m := newModel(context.Background(), nil, Options{ModelName: "test-model", WorkingDir: dir})
status, err := m.handleStatusCommand(nil)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(status, "main_test.go") {
t.Fatalf("status missing changed file:\n%s", status)
}
diff, err := m.handleDiffCommand(nil)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(diff, "changed") {
t.Fatalf("diff missing changed content:\n%s", diff)
}
out, err := m.handleTestCommand(nil)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "$ go test ./...") {
t.Fatalf("test output missing default command:\n%s", out)
}
}
func runInDir(t *testing.T, dir string, name string, args ...string) {
t.Helper()
cmd := exec.Command(name, args...)
cmd.Dir = dir
if output, err := cmd.CombinedOutput(); err != nil {
t.Fatalf("%s %v failed: %v\n%s", name, args, err, output)
}
}
func TestModelCommand(t *testing.T) { func TestModelCommand(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true}) m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
m.input.SetValue("/model") m.input.SetValue("/model")
+143
View File
@@ -0,0 +1,143 @@
package tui
import (
"bytes"
"context"
"errors"
"fmt"
"os"
"os/exec"
"path/filepath"
"strings"
"time"
"agentu/internal/tools"
)
const workflowCommandTimeout = 2 * time.Minute
func (m model) handleStatusCommand(args []string) (string, error) {
if len(args) != 0 {
return "", fmt.Errorf("usage: /status")
}
out, err := m.runGitCommand("status", "--short")
if err != nil {
return out, err
}
if strings.TrimSpace(out) == "" {
return "Working tree clean.", nil
}
return "Changed files:\n" + out, nil
}
func (m model) handleDiffCommand(args []string) (string, error) {
if len(args) > 0 && args[0] == "--stat" {
out, err := m.runGitCommand(append([]string{"diff", "--stat", "--"}, args[1:]...)...)
if err != nil {
return out, err
}
if strings.TrimSpace(out) == "" {
return "No unstaged diff.", nil
}
return out, nil
}
out, err := m.runGitCommand(append([]string{"diff", "--"}, args...)...)
if err != nil {
return out, err
}
if strings.TrimSpace(out) == "" {
return "No unstaged diff.", nil
}
return out, nil
}
func (m model) handleTestCommand(args []string) (string, error) {
command := strings.TrimSpace(strings.Join(args, " "))
if command == "" {
var err error
command, err = defaultTestCommand(m.workflowDir())
if err != nil {
return "", err
}
}
out, err := m.runShellCommand(command)
if err != nil {
if strings.TrimSpace(out) != "" {
return out, err
}
return "", err
}
if strings.TrimSpace(out) == "" {
return fmt.Sprintf("$ %s\n(no output)", command), nil
}
return out, nil
}
func (m model) runGitCommand(args ...string) (string, error) {
ctx, cancel := context.WithTimeout(m.ctx, workflowCommandTimeout)
defer cancel()
cmd := exec.CommandContext(ctx, "git", append([]string{"-C", m.workflowDir()}, args...)...)
output, err := cmd.CombinedOutput()
text := tools.FormatResult(string(output))
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
return text, fmt.Errorf("git command timed out: %w", ctx.Err())
}
if err != nil {
if strings.TrimSpace(text) != "" {
return text, fmt.Errorf("git command failed: %w", err)
}
return "", fmt.Errorf("git command failed: %w", err)
}
return text, nil
}
func (m model) runShellCommand(command string) (string, error) {
ctx, cancel := context.WithTimeout(m.ctx, workflowCommandTimeout)
defer cancel()
shell := os.Getenv("SHELL")
if shell == "" {
shell = "/bin/sh"
}
cmd := exec.CommandContext(ctx, shell, "-lc", command)
cmd.Dir = m.workflowDir()
var combined bytes.Buffer
cmd.Stdout = &combined
cmd.Stderr = &combined
err := cmd.Run()
text := strings.TrimRight(tools.FormatResult(combined.String()), "\n")
if text != "" {
text = fmt.Sprintf("$ %s\n%s", command, text)
}
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
return text, fmt.Errorf("command timed out: %w", ctx.Err())
}
if err != nil {
return text, fmt.Errorf("command failed: %w", err)
}
return text, nil
}
func (m model) workflowDir() string {
if strings.TrimSpace(m.workingDir) == "" {
return "."
}
return m.workingDir
}
func defaultTestCommand(dir string) (string, error) {
if fileExists(filepath.Join(dir, "go.mod")) {
return "go test ./...", nil
}
if fileExists(filepath.Join(dir, "package.json")) {
return "npm test", nil
}
return "", fmt.Errorf("no default test command found; use /test <command>")
}
func fileExists(path string) bool {
info, err := os.Stat(path)
return err == nil && !info.IsDir()
}