loveuer/agentu
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+27 -18
View File
@@ -44,19 +44,10 @@ providers:
loveuer:
base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY}
model: deepseek-v4-flash
models:
- deepseek-v4-flash
thinking: none
thinking_param: thinking
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
- id: deepseek-v4-flash
thinking: none
context: 256000
```
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
current session.
`models` is optional. If present, `/model model <name>` is restricted to that
list. If omitted, any model name is accepted for that provider.
`models` is required. `/model model <name>` is restricted to that list. Each
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
none, middle, high, xhigh, max
```
When configured or changed with `/model thinking ...`, the value is sent as a
top-level OpenAI-compatible request field. `thinking_param` controls that field
name and defaults to `thinking`.
top-level OpenAI-compatible request field named `thinking`.
## CLI Flags
@@ -90,6 +89,7 @@ go run ./cmd/agentu [flags]
- `--theme light|dark` selects the TUI theme. Default: `light`.
- `--plain` uses the simple line-based REPL instead of the TUI.
- `--yolo` enables automatic file writes and shell execution.
- `--resume <id>` resumes a saved session from `~/.agentu/sessions`.
## TUI Controls
@@ -103,10 +103,18 @@ go run ./cmd/agentu [flags]
Slash commands:
- `/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 provider <name>` switches provider 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.
- `/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.
## Local Tools
@@ -115,7 +123,8 @@ Read-only tools are available by default:
- `file_read`
- `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:
agentu.example.yaml
+3 -12
View File
@@ -2,16 +2,7 @@ providers:
loveuer:
base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY}
model: deepseek-v4-flash
models:
- deepseek-v4-flash
thinking: none
thinking_param: thinking
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
- id: deepseek-v4-flash
thinking: none
context: 256000
cmd/agentu/main.go
+58 -14
View File
@@ -20,6 +20,20 @@ import (
"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() {
if err := run(); err != nil {
fmt.Fprintln(os.Stderr, "agentu:", err)
@@ -43,11 +57,10 @@ func run() error {
cfg, err := config.Load(*configPath)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
resolvedPath, resolveErr := config.ResolvePath(*configPath)
if resolveErr != nil {
return resolveErr
if bootstrapErr := bootstrapConfig(*configPath); bootstrapErr != nil {
return bootstrapErr
}
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
}
@@ -65,17 +78,22 @@ func run() error {
registry.Register(tools.NewBuiltinSearchTool(http.DefaultClient))
registry.Register(tools.NewBuiltinFetchTool(http.DefaultClient))
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{
Provider: provider,
ProviderName: providerName,
Model: providerConfig.Model,
Thinking: providerConfig.Thinking,
ThinkingParam: providerConfig.ThinkingParam,
ThinkingEnabled: providerConfig.ThinkingConfigured,
SystemPrompt: systemPrompt,
ToolRegistry: registry,
ToolTimeout: cfg.Agent.ToolTimeout,
Provider: provider,
ProviderName: providerName,
Model: providerConfig.Model,
Thinking: providerConfig.Thinking,
ThinkingParam: providerConfig.ThinkingParam,
ThinkingEnabled: providerConfig.ThinkingConfigured,
SystemPrompt: systemPrompt,
ToolRegistry: registry,
ToolTimeout: cfg.Agent.ToolTimeout,
MaxContextTokens: maxContextTokens,
})
// Initialize session store
@@ -114,6 +132,7 @@ func run() error {
Yolo: *yolo,
ThemeMode: themeMode,
ModelManager: modelManager,
WorkingDir: cfg.Agent.WorkingDir,
})
printExitSession(modelManager)
return err
@@ -122,12 +141,27 @@ func run() error {
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 {
scanner := bufio.NewScanner(os.Stdin)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
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 {
fmt.Print("> ")
if !scanner.Scan() {
@@ -150,6 +184,16 @@ func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Man
assistant.Clear()
fmt.Println("context cleared")
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 {
internal/agent/agent.go
+250 -41
View File
@@ -13,20 +13,36 @@ import (
"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 {
provider llm.Provider
providerName string
model string
thinking string
thinkingParam string
thinkingEnabled bool
systemPrompt string
toolRegistry *tools.Registry
toolTimeout time.Duration
maxToolIterations int
messages []llm.Message
provider llm.Provider
providerName string
model string
thinking string
thinkingParam string
thinkingEnabled bool
systemPrompt string
toolRegistry *tools.Registry
toolTimeout time.Duration
maxToolRounds int
maxContextTokens int
lastUsage *llm.Usage
messages []llm.Message
}
type Options struct {
@@ -40,47 +56,52 @@ type Options struct {
ToolRegistry *tools.Registry
ToolTimeout time.Duration
MaxToolIterations int
MaxContextTokens int
}
func New(opts Options) *Agent {
maxToolIterations := opts.MaxToolIterations
if maxToolIterations <= 0 {
maxToolIterations = defaultMaxToolIterations
maxToolRounds := opts.MaxToolIterations
if maxToolRounds <= 0 {
maxToolRounds = defaultMaxToolRounds
}
toolTimeout := opts.ToolTimeout
if toolTimeout <= 0 {
toolTimeout = 30 * time.Second
}
a := &Agent{
provider: opts.Provider,
providerName: opts.ProviderName,
model: opts.Model,
thinking: opts.Thinking,
thinkingParam: opts.ThinkingParam,
thinkingEnabled: opts.ThinkingEnabled,
systemPrompt: opts.SystemPrompt,
toolRegistry: opts.ToolRegistry,
toolTimeout: toolTimeout,
maxToolIterations: maxToolIterations,
provider: opts.Provider,
providerName: opts.ProviderName,
model: opts.Model,
thinking: opts.Thinking,
thinkingParam: opts.ThinkingParam,
thinkingEnabled: opts.ThinkingEnabled,
systemPrompt: opts.SystemPrompt,
toolRegistry: opts.ToolRegistry,
toolTimeout: toolTimeout,
maxToolRounds: maxToolRounds,
maxContextTokens: opts.MaxContextTokens,
}
a.Clear()
return a
}
type RuntimeOptions struct {
Provider llm.Provider
ProviderName string
Model string
Thinking string
ThinkingParam string
ThinkingEnabled bool
Provider llm.Provider
ProviderName string
Model string
Thinking string
ThinkingParam string
ThinkingEnabled bool
MaxContextTokens int
}
type RuntimeInfo struct {
ProviderName string
Model string
Thinking string
ThinkingEnabled bool
ProviderName string
Model string
Thinking string
ThinkingEnabled bool
ContextTokens int
MaxContextTokens int
}
func (a *Agent) SetRuntime(opts RuntimeOptions) {
@@ -96,18 +117,22 @@ func (a *Agent) SetRuntime(opts RuntimeOptions) {
a.thinking = opts.Thinking
a.thinkingParam = opts.ThinkingParam
a.thinkingEnabled = opts.ThinkingEnabled
a.maxContextTokens = opts.MaxContextTokens
}
func (a *Agent) RuntimeInfo() RuntimeInfo {
return RuntimeInfo{
ProviderName: a.providerName,
Model: a.model,
Thinking: a.thinking,
ThinkingEnabled: a.thinkingEnabled,
ProviderName: a.providerName,
Model: a.model,
Thinking: a.thinking,
ThinkingEnabled: a.thinkingEnabled,
ContextTokens: llm.EstimateTokens(a.messages, a.toolDefinitions()),
MaxContextTokens: a.maxContextTokens,
}
}
func (a *Agent) Clear() {
a.lastUsage = nil
a.messages = nil
if strings.TrimSpace(a.systemPrompt) != "" {
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...)
}
// 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 {
input = strings.TrimSpace(input)
if input == "" {
return nil
}
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)
if err != nil {
return err
@@ -158,6 +330,40 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
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 {
result := a.executeTool(ctx, call, logs)
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) {
@@ -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 {
if event.Usage != nil {
a.SetLastUsage(event.Usage)
}
if event.Content != "" {
content.WriteString(event.Content)
if out != nil {
internal/agent/agent_test.go
+203
View File
@@ -3,6 +3,7 @@ package agent
import (
"context"
"encoding/json"
"fmt"
"strings"
"testing"
@@ -96,6 +97,134 @@ func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, e
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) {
provider := &contentProvider{text: "shell is disabled"}
a := New(Options{
@@ -137,3 +266,77 @@ func TestAgentDoesNotBlockFileRequest(t *testing.T) {
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)
}
}
internal/config/config.go
+97 -35
View File
@@ -21,37 +21,50 @@ type Config struct {
type ProviderConfig struct {
Name string
BaseURL string `yaml:"base_url"`
Model string `yaml:"model"`
Models []string `yaml:"models"`
APIKey string `yaml:"api_key"`
Thinking string `yaml:"thinking"`
ThinkingParam string `yaml:"thinking_param"`
BaseURL string
Model string
Models []string
ModelConfigs map[string]ModelConfig
ContextTokens int
APIKey string
Thinking string
ThinkingParam string
ThinkingConfigured bool
}
type ModelConfig struct {
ID string `yaml:"id"`
Thinking string `yaml:"thinking"`
ContextTokens int `yaml:"context"`
}
type AgentConfig struct {
SystemPrompt string `yaml:"system_prompt"`
WorkingDir string `yaml:"working_dir"`
ToolTimeout time.Duration `yaml:"tool_timeout"`
SystemPrompt string `yaml:"system_prompt"`
WorkingDir string `yaml:"working_dir"`
ToolTimeout time.Duration `yaml:"tool_timeout"`
MaxContextTokens int `yaml:"max_context_tokens"`
}
type rawConfig struct {
Providers map[string]rawProviderConfig `yaml:"providers"`
Agent struct {
SystemPrompt string `yaml:"system_prompt"`
WorkingDir string `yaml:"working_dir"`
ToolTimeout string `yaml:"tool_timeout"`
SystemPrompt string `yaml:"system_prompt"`
WorkingDir string `yaml:"working_dir"`
ToolTimeout string `yaml:"tool_timeout"`
MaxContextTokens int `yaml:"max_context_tokens"`
} `yaml:"agent"`
}
type rawProviderConfig struct {
BaseURL string `yaml:"base_url"`
Model string `yaml:"model"`
Models []string `yaml:"models"`
APIKey string `yaml:"api_key"`
Thinking *string `yaml:"thinking"`
ThinkingParam string `yaml:"thinking_param"`
BaseURL string `yaml:"base_url"`
Models []rawModelConfig `yaml:"models"`
APIKey string `yaml:"api_key"`
}
type rawModelConfig struct {
ID string `yaml:"id"`
Thinking string `yaml:"thinking"`
ContextTokens int `yaml:"context"`
}
func Load(path string) (*Config, error) {
@@ -76,8 +89,9 @@ func Load(path string) (*Config, error) {
cfg := &Config{
Providers: normalizeProviders(raw),
Agent: AgentConfig{
SystemPrompt: raw.Agent.SystemPrompt,
WorkingDir: raw.Agent.WorkingDir,
SystemPrompt: raw.Agent.SystemPrompt,
WorkingDir: raw.Agent.WorkingDir,
MaxContextTokens: raw.Agent.MaxContextTokens,
},
}
@@ -142,8 +156,8 @@ func (c *Config) Validate() error {
if strings.TrimSpace(provider.BaseURL) == "" {
missing = append(missing, prefix+".base_url")
}
if strings.TrimSpace(provider.Model) == "" {
missing = append(missing, prefix+".model")
if len(provider.Models) == 0 {
missing = append(missing, prefix+".models")
}
if strings.TrimSpace(provider.APIKey) == "" {
missing = append(missing, prefix+".api_key")
@@ -151,6 +165,19 @@ func (c *Config) Validate() error {
if len(missing) > 0 {
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) {
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 {
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
}
@@ -187,27 +217,59 @@ func normalizeProviders(raw rawConfig) map[string]ProviderConfig {
}
func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
models, modelConfigs := normalizeModelConfigs(raw.Models)
provider := ProviderConfig{
Name: name,
BaseURL: raw.BaseURL,
Model: raw.Model,
Models: append([]string(nil), raw.Models...),
APIKey: raw.APIKey,
ThinkingParam: raw.ThinkingParam,
Name: name,
BaseURL: raw.BaseURL,
Models: models,
ModelConfigs: modelConfigs,
APIKey: raw.APIKey,
}
if provider.Model == "" && len(provider.Models) > 0 {
if len(provider.Models) > 0 {
provider.Model = provider.Models[0]
}
if raw.Thinking != nil {
provider.Thinking = strings.ToLower(strings.TrimSpace(*raw.Thinking))
provider.ThinkingConfigured = true
if provider.ThinkingParam == "" {
provider.ThinkingParam = "thinking"
}
applyModelDefaults(&provider)
if provider.ThinkingConfigured && provider.ThinkingParam == "" {
provider.ThinkingParam = "thinking"
}
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 {
return []string{"none", "middle", "high", "xhigh", "max"}
}
internal/config/config_test.go
+130 -15
View File
@@ -16,8 +16,9 @@ func TestLoadExpandsEnvAndDefaults(t *testing.T) {
providers:
test:
base_url: https://example.com
model: test-model
api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
agent:
working_dir: .
`), 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) {
dir := t.TempDir()
t.Setenv("HOME", dir)
@@ -56,8 +166,9 @@ func TestLoadUsesDefaultConfigPath(t *testing.T) {
providers:
test:
base_url: https://example.com
model: test-model
api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
@@ -93,8 +204,9 @@ func TestLoadDefaultSystemPromptDiscouragesEagerTools(t *testing.T) {
providers:
test:
base_url: https://example.com
model: test-model
api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
@@ -119,17 +231,15 @@ providers:
loveuer:
base_url: https://example.com
api_key: ${AGENTU_TEST_KEY}
model: model-a
models:
- model-a
- model-b
thinking: xhigh
thinking_param: thinking
- id: model-a
thinking: xhigh
- id: model-b
backup:
base_url: https://backup.example.com
api_key: ${AGENTU_TEST_KEY}
models:
- model-c
- id: model-c
`), 0o644); err != nil {
t.Fatal(err)
}
@@ -162,11 +272,13 @@ providers:
two:
base_url: https://two.example.com
api_key: ${AGENTU_TEST_KEY}
model: model-b
models:
- id: model-b
one:
base_url: https://one.example.com
api_key: ${AGENTU_TEST_KEY}
model: model-a
models:
- id: model-a
`), 0o644); err != nil {
t.Fatal(err)
}
@@ -188,9 +300,10 @@ func TestLoadRejectsInvalidThinking(t *testing.T) {
providers:
test:
base_url: https://example.com
model: test-model
api_key: ${AGENTU_TEST_KEY}
thinking: huge
models:
- id: test-model
thinking: huge
`), 0o644); err != nil {
t.Fatal(err)
}
@@ -224,8 +337,9 @@ func TestLoadRejectsLegacyProviderConfig(t *testing.T) {
if err := os.WriteFile(path, []byte(`
provider:
base_url: https://example.com
model: test-model
api_key: sk-test
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
@@ -244,8 +358,9 @@ active_provider: loveuer
providers:
loveuer:
base_url: https://example.com
model: test-model
api_key: sk-test
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
internal/llm/openai.go
+79 -26
View File
@@ -10,8 +10,13 @@ import (
"io"
"net/http"
"strings"
"time"
)
const maxChatAttempts = 3
const retryDelay = 100 * time.Millisecond
type OpenAICompatibleClient struct {
baseURL 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 {
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)
if err != nil {
return fmt.Errorf("marshal chat request: %w", err)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
if err != nil {
return fmt.Errorf("create chat request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
if c.apiKey != "" {
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
}
var lastErr error
for attempt := 1; attempt <= maxChatAttempts; attempt++ {
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
if err != nil {
return fmt.Errorf("create chat request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
if c.apiKey != "" {
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
}
resp, err := c.httpClient.Do(httpReq)
if err != nil {
return fmt.Errorf("send chat request: %w", err)
}
defer resp.Body.Close()
resp, err := c.httpClient.Do(httpReq)
if err != nil {
lastErr = fmt.Errorf("send chat request: %w", err)
if !shouldRetryRequest(ctx, attempt, 0) {
return lastErr
}
if waitErr := waitBeforeRetry(ctx); waitErr != nil {
return waitErr
}
continue
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
limit := io.LimitReader(resp.Body, 4096)
data, _ := io.ReadAll(limit)
return fmt.Errorf("chat request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
limit := io.LimitReader(resp.Body, 4096)
data, _ := io.ReadAll(limit)
_ = 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 {
@@ -79,7 +132,7 @@ func readSSE(r io.Reader, emit func(StreamEvent) error) error {
if parseErr != nil {
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 {
return emitErr
}
@@ -100,6 +153,7 @@ type streamPayload struct {
} `json:"delta"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
Usage *Usage `json:"usage,omitempty"`
Error *struct {
Message string `json:"message"`
Type string `json:"type"`
@@ -127,15 +181,14 @@ func parseStreamPayload(payload string) (StreamEvent, error) {
}
return StreamEvent{}, fmt.Errorf("provider error: %s", decoded.Error.Message)
}
event := StreamEvent{Usage: decoded.Usage}
if len(decoded.Choices) == 0 {
return StreamEvent{}, nil
return event, nil
}
choice := decoded.Choices[0]
event := StreamEvent{
Content: choice.Delta.Content,
FinishReason: choice.FinishReason,
}
event.Content = choice.Delta.Content
event.FinishReason = choice.FinishReason
for _, tc := range choice.Delta.ToolCalls {
event.ToolCalls = append(event.ToolCalls, ToolCallDelta{
Index: tc.Index,
internal/llm/openai_test.go
+73
View File
@@ -30,6 +30,10 @@ func TestOpenAICompatibleClientStreamsContentAndToolCalls(t *testing.T) {
if raw["thinking"] != "high" {
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 {
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 {
t.Helper()
data, err := io.ReadAll(r.Body)
internal/llm/token_estimate.go
+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
}
internal/llm/token_estimate_test.go
+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)
}
}
internal/llm/types.go
+7
View File
@@ -79,6 +79,13 @@ type StreamEvent struct {
Content string
ToolCalls []ToolCallDelta
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 {
internal/session/manager.go
+19 -13
View File
@@ -69,6 +69,7 @@ func (m *Manager) Resume(id string) error {
}
m.currentSession = sess
m.agent.SetMessages(sess.Messages)
m.agent.SetLastUsage(sess.LastUsage)
return nil
}
@@ -78,6 +79,7 @@ func (m *Manager) Save() error {
return nil
}
m.currentSession.Messages = m.agent.Messages()
m.currentSession.LastUsage = m.agent.LastUsage()
m.currentSession.Provider = m.providerName
m.currentSession.Model = m.provider.Model
// Auto-name if empty
@@ -208,11 +210,7 @@ func (m *Manager) Info() string {
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking())
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, ", "))
} else {
b.WriteString("models: any model name is accepted for this provider\n")
}
fmt.Fprintf(&b, "models: %s\n", strings.Join(m.provider.Models, ", "))
fmt.Fprintf(&b, "thinking levels: %s\n", strings.Join(config.ThinkingLevels(), ", "))
b.WriteString("usage: /model provider <name> | /model model <name> | /model thinking <level>")
return b.String()
@@ -234,10 +232,10 @@ func (m *Manager) SetModel(model string) (string, error) {
if model == "" {
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, ", "))
}
m.provider.Model = model
m.provider = m.provider.WithModel(model)
m.cfg.Providers[m.providerName] = m.provider
m.syncAgent(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())
}
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) {
m.agent.SetRuntime(agent.RuntimeOptions{
Provider: provider,
ProviderName: m.providerName,
Model: m.provider.Model,
Thinking: m.provider.Thinking,
ThinkingParam: m.provider.ThinkingParam,
ThinkingEnabled: m.provider.ThinkingConfigured,
Provider: provider,
ProviderName: m.providerName,
Model: m.provider.Model,
Thinking: m.provider.Thinking,
ThinkingParam: m.provider.ThinkingParam,
ThinkingEnabled: m.provider.ThinkingConfigured,
MaxContextTokens: m.maxContextTokens(),
})
}
internal/session/manager_test.go
+67
View File
@@ -116,6 +116,7 @@ func TestManagerSaveAndResume(t *testing.T) {
manager.newSession()
var out strings.Builder
assistant.RunTurn(context.Background(), "hello world", &out, nil)
assistant.SetLastUsage(&llm.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15})
// Save
if err := manager.Save(); err != nil {
@@ -153,6 +154,9 @@ func TestManagerSaveAndResume(t *testing.T) {
if !hasUser {
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) {
@@ -264,3 +268,66 @@ func TestManagerResumeLatest(t *testing.T) {
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)
}
}
internal/session/store.go
+1
View File
@@ -23,6 +23,7 @@ type Session struct {
Provider string `json:"provider"`
Model string `json:"model"`
Messages []llm.Message `json:"messages"`
LastUsage *llm.Usage `json:"last_usage,omitempty"`
}
// SessionMeta is a lightweight summary for listing.
internal/session/store_test.go
+4
View File
@@ -30,6 +30,7 @@ func TestSaveAndLoad(t *testing.T) {
{Role: llm.RoleUser, Content: "hello"},
{Role: llm.RoleAssistant, Content: "hi there"},
},
LastUsage: &llm.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15},
}
if err := store.Save(sess); err != nil {
t.Fatalf("save: %v", err)
@@ -51,6 +52,9 @@ func TestSaveAndLoad(t *testing.T) {
if 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() {
t.Error("UpdatedAt should be set")
}
internal/tools/code.go
+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()
}
internal/tools/file.go
+61 -10
View File
@@ -223,7 +223,9 @@ func (t *FileSearchTool) Definition() llm.Tool {
"properties": {
"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."},
"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"],
"additionalProperties": false
@@ -233,9 +235,11 @@ func (t *FileSearchTool) Definition() llm.Tool {
}
type fileSearchArgs struct {
Query string `json:"query"`
Path string `json:"path"`
Glob string `json:"glob"`
Query string `json:"query"`
Path string `json:"path"`
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) {
@@ -256,8 +260,31 @@ func (t *FileSearchTool) Execute(ctx context.Context, raw json.RawMessage) (stri
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) {
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 != "" {
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) {
contextLines, maxResults := normalizedSearchOptions(args)
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 {
if err != nil {
return nil
@@ -312,16 +342,22 @@ func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, ro
}
lines := strings.Split(string(data), "\n")
for i, line := range lines {
if strings.Contains(line, args.Query) {
fmt.Fprintf(&out, "%s:%d:%s\n", path, i+1, line)
if out.Len() > MaxToolOutputBytes {
return errors.New("output limit reached")
}
if !strings.Contains(line, args.Query) {
continue
}
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
})
if err != nil && err.Error() != "output limit reached" {
if err != nil && !errors.Is(err, stop) {
return "", err
}
if out.Len() == 0 {
@@ -329,3 +365,18 @@ func (t *FileSearchTool) walkSearch(ctx context.Context, args fileSearchArgs, ro
}
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])
}
}
internal/tools/paths.go
+19 -2
View File
@@ -1,7 +1,9 @@
package tools
import (
"fmt"
"path/filepath"
"strings"
)
func resolvePath(baseDir, path string) (string, error) {
@@ -9,7 +11,22 @@ func resolvePath(baseDir, path string) (string, error) {
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
}
internal/tools/tools.go
+42 -5
View File
@@ -5,12 +5,15 @@ import (
"encoding/json"
"fmt"
"sort"
"strings"
"agentu/internal/llm"
)
const MaxToolOutputBytes = 64 * 1024
const truncationNoticeBudget = 160
type Tool interface {
Definition() llm.Tool
Execute(ctx context.Context, args json.RawMessage) (string, error)
@@ -62,11 +65,13 @@ func ReadOnlyBuiltins(workingDir string) *Registry {
fileRead := NewFileReadTool(workingDir)
fileList := NewFileListTool(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.list", fileList)
registry.RegisterAlias("file.search", fileSearch)
registry.RegisterAlias("code.symbols", codeSymbols)
return registry
}
@@ -83,8 +88,8 @@ func Builtins(workingDir string) *Registry {
func AgentInstructions(workingDir string, yolo bool) string {
searchInstructions := webSearchInstructions()
readOnlyTools := "file_list, file_read, file_search, web_search, web_fetch"
yoloTools := "file_list, file_read, file_search, file_write, shell_run, web_search, web_fetch"
readOnlyTools := "file_list, file_read, file_search, code_symbols, web_search, web_fetch"
yoloTools := "file_list, file_read, file_search, code_symbols, file_write, shell_run, web_search, web_fetch"
if yolo {
return fmt.Sprintf(`Local project context:
- 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 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 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.
- 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.
@@ -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 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 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.
- %s
- 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 {
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[: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 {
internal/tools/tools_test.go
+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) {
dir := t.TempDir()
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) {
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 {
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) {
registry := ReadOnlyBuiltins(t.TempDir())
registry.Register(NewBuiltinSearchTool(nil))
@@ -132,9 +195,65 @@ func TestAgentInstructionsExplainCommandModes(t *testing.T) {
func TestAgentInstructionsExplainWebSearch(t *testing.T) {
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) {
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 ..")
}
}
internal/tui/app.go
+111 -25
View File
@@ -28,6 +28,7 @@ type Options struct {
Yolo bool
ThemeMode ThemeMode
ModelManager ModelManager
WorkingDir string
}
type ModelManager interface {
@@ -69,13 +70,14 @@ type message struct {
}
type model struct {
ctx context.Context
agent *agent.Agent
modelName string
yolo bool
models ModelManager
theme theme
styles styles
ctx context.Context
agent *agent.Agent
modelName string
yolo bool
models ModelManager
workingDir string
theme theme
styles styles
viewport viewport.Model
input textarea.Model
@@ -84,11 +86,12 @@ type model struct {
width int
height int
messages []message
status string
running bool
cancel context.CancelFunc
events <-chan tea.Msg
messages []message
status string
contextUsage string
running bool
cancel context.CancelFunc
events <-chan tea.Msg
// Session picker state
picking bool
@@ -119,6 +122,7 @@ type slashCommand struct {
var slashCommands = []slashCommand{
{Name: "/clear", Description: "reset context"},
{Name: "/compact", Description: "compress context"},
{Name: "/exit", Description: "quit"},
{Name: "/quit", Description: "quit"},
{Name: "/model", Description: "model/provider/thinking"},
@@ -126,6 +130,9 @@ var slashCommands = []slashCommand{
{Name: "/resume", Description: "resume session"},
{Name: "/rename", Description: "rename 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 {
@@ -154,19 +161,22 @@ func newModel(ctx context.Context, assistant *agent.Agent, opts Options) model {
vp.Style = st.Viewport
vp.SetContent("")
return model{
ctx: ctx,
agent: assistant,
modelName: opts.ModelName,
yolo: opts.Yolo,
models: opts.ModelManager,
theme: th,
styles: st,
viewport: vp,
input: input,
spinner: spin,
status: "ready",
m := model{
ctx: ctx,
agent: assistant,
modelName: opts.ModelName,
yolo: opts.Yolo,
models: opts.ModelManager,
workingDir: opts.WorkingDir,
theme: th,
styles: st,
viewport: vp,
input: input,
spinner: spin,
status: "ready",
}
m.refreshContextUsage()
return m
}
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 {
_ = m.models.Save()
}
m.refreshContextUsage()
m.resize()
m.refreshViewport(true)
cmds = append(cmds, textarea.Blink)
@@ -338,13 +349,14 @@ func (m *model) submit() (tea.Model, tea.Cmd) {
return *m, tea.Quit
case "/clear":
m.agent.Clear()
m.refreshContextUsage()
m.messages = []message{{role: roleSystem, content: "Context cleared."}}
m.input.SetValue("")
m.afterInputChanged()
m.status = "ready"
m.refreshViewport(true)
return *m, nil
case "/model":
case "/model", "/compact":
return m.runLocalCommand(input)
}
@@ -541,6 +553,9 @@ func (m model) modelMetaView() string {
parts = append(parts, "thinking "+currentThinking)
}
}
if m.contextUsage != "" {
parts = append(parts, m.contextUsage)
}
parts = append([]string{modelName, mode}, parts...)
parts = append(parts, "theme "+string(m.theme.Mode))
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)
}
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 {
width = max(1, width)
if lipgloss.Width(text) <= width {
@@ -741,6 +808,7 @@ func (m *model) runLocalCommand(input string) (tea.Model, tea.Cmd) {
m.status = status
// Rebuild display from agent history (matters after /resume and /new)
m.syncMessages()
m.refreshContextUsage()
// Append command output after synced history
m.messages = append(m.messages, message{role: role, content: output})
m.refreshViewport(true)
@@ -753,6 +821,8 @@ func (m model) handleLocalCommand(input string) (string, error) {
return "", fmt.Errorf("empty command")
}
switch fields[0] {
case "/compact":
return m.handleCompactCommand()
case "/model":
return m.handleModelCommand(fields[1:])
case "/sessions":
@@ -763,11 +833,27 @@ func (m model) handleLocalCommand(input string) (string, error) {
return m.handleRenameCommand(fields[1:])
case "/new":
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:
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) {
if m.models == nil {
if len(args) > 0 {
internal/tui/app_test.go
+116 -1
View File
@@ -2,9 +2,15 @@ package tui
import (
"context"
"os"
"os/exec"
"path/filepath"
"strings"
"testing"
"agentu/internal/agent"
"agentu/internal/llm"
tea "github.com/charmbracelet/bubbletea"
"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) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
@@ -264,13 +284,108 @@ func TestSlashCommandSuggestions(t *testing.T) {
m.afterInputChanged()
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) {
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) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
m.input.SetValue("/model")
internal/tui/workflow.go
+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()
}