diff --git a/README.md b/README.md index f7a24d1..fb21dc9 100644 --- a/README.md +++ b/README.md @@ -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 ` to switch providers for the current session. -`models` is optional. If present, `/model model ` is restricted to that -list. If omitted, any model name is accepted for that provider. +`models` is required. `/model model ` 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 ` 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 ` switches provider for the current session. - `/model model ` or `/model ` switches model for the current session. - `/model thinking ` changes thinking for the current session. +- `/sessions` lists saved sessions. +- `/resume [id]` resumes by ID or opens an interactive picker. +- `/rename ` 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: diff --git a/agentu.example.yaml b/agentu.example.yaml index 5549d27..e6a5a52 100644 --- a/agentu.example.yaml +++ b/agentu.example.yaml @@ -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 diff --git a/cmd/agentu/main.go b/cmd/agentu/main.go index 56e1034..6ed9766 100644 --- a/cmd/agentu/main.go +++ b/cmd/agentu/main.go @@ -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 { diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 6853499..33ee43e 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -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 = © +} + +// 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 = © + } + 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 { diff --git a/internal/agent/agent_test.go b/internal/agent/agent_test.go index 2977b29..6affae1 100644 --- a/internal/agent/agent_test.go +++ b/internal/agent/agent_test.go @@ -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) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 8678021..5dc0ea7 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -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"} } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index c58cae5..5e6d7bd 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -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) } diff --git a/internal/llm/openai.go b/internal/llm/openai.go index 47f1375..c03d85e 100644 --- a/internal/llm/openai.go +++ b/internal/llm/openai.go @@ -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, diff --git a/internal/llm/openai_test.go b/internal/llm/openai_test.go index 62915ca..b249731 100644 --- a/internal/llm/openai_test.go +++ b/internal/llm/openai_test.go @@ -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) diff --git a/internal/llm/token_estimate.go b/internal/llm/token_estimate.go new file mode 100644 index 0000000..77c99d7 --- /dev/null +++ b/internal/llm/token_estimate.go @@ -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 +} diff --git a/internal/llm/token_estimate_test.go b/internal/llm/token_estimate_test.go new file mode 100644 index 0000000..a508d0b --- /dev/null +++ b/internal/llm/token_estimate_test.go @@ -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) + } +} diff --git a/internal/llm/types.go b/internal/llm/types.go index a319f99..10c4d6f 100644 --- a/internal/llm/types.go +++ b/internal/llm/types.go @@ -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 { diff --git a/internal/session/manager.go b/internal/session/manager.go index 1c61ea0..d418a6f 100644 --- a/internal/session/manager.go +++ b/internal/session/manager.go @@ -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 | /model model | /model thinking ") 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(), }) } diff --git a/internal/session/manager_test.go b/internal/session/manager_test.go index 2bc230a..2dbc528 100644 --- a/internal/session/manager_test.go +++ b/internal/session/manager_test.go @@ -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) + } +} diff --git a/internal/session/store.go b/internal/session/store.go index 93a5086..60c9588 100644 --- a/internal/session/store.go +++ b/internal/session/store.go @@ -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. diff --git a/internal/session/store_test.go b/internal/session/store_test.go index 071d443..9b0e7e8 100644 --- a/internal/session/store_test.go +++ b/internal/session/store_test.go @@ -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") } diff --git a/internal/tools/code.go b/internal/tools/code.go new file mode 100644 index 0000000..4920223 --- /dev/null +++ b/internal/tools/code.go @@ -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() +} diff --git a/internal/tools/file.go b/internal/tools/file.go index 0c806c9..01220b4 100644 --- a/internal/tools/file.go +++ b/internal/tools/file.go @@ -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]) + } +} diff --git a/internal/tools/paths.go b/internal/tools/paths.go index 433ea6c..5436d05 100644 --- a/internal/tools/paths.go +++ b/internal/tools/paths.go @@ -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 } diff --git a/internal/tools/tools.go b/internal/tools/tools.go index 3879505..3b94f50 100644 --- a/internal/tools/tools.go +++ b/internal/tools/tools.go @@ -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 { diff --git a/internal/tools/tools_test.go b/internal/tools/tools_test.go index df8768b..b11ca74 100644 --- a/internal/tools/tools_test.go +++ b/internal/tools/tools_test.go @@ -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 ..") + } +} diff --git a/internal/tui/app.go b/internal/tui/app.go index 71e91d5..057a76a 100644 --- a/internal/tui/app.go +++ b/internal/tui/app.go @@ -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 { diff --git a/internal/tui/app_test.go b/internal/tui/app_test.go index 99b035f..4a5bb10 100644 --- a/internal/tui/app_test.go +++ b/internal/tui/app_test.go @@ -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") diff --git a/internal/tui/workflow.go b/internal/tui/workflow.go new file mode 100644 index 0000000..bc23d3b --- /dev/null +++ b/internal/tui/workflow.go @@ -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 ") +} + +func fileExists(path string) bool { + info, err := os.Stat(path) + return err == nil && !info.IsDir() +}