4 Commits

Author SHA1 Message Date
loveuer 7b8bad5e62 feat: crush-inspired TUI improvements, activity bar polish, startup fresh session
TUI improvements:
- semantic icon constants (✓ ◆ ▣ ✕ ⌘ ◌ ◔ ◉)
- pill-style metadata bar with model/provider/session/context
- fuzzy command palette with deterministic subsequence scoring
- input history navigation (ctrl+p / ctrl+n)
- structured tool log cards (▣ tool_name + args preview)
- no-background activity indicator and meta bar for cleaner rendering

Startup behavior:
- default startup creates a fresh session
- explicit --resume <id> restores saved context
- added main_test.go for initializeSession coverage
2026-06-24 18:16:53 -07:00
loveuer f235b8ce90 refactor: restructure project layout — main.go to root, llm+config to pkg/
- Move cmd/agentu/main.go → main.go (project root)
- Move internal/llm/ → pkg/llm/ (leaf package, no internal deps)
- Move internal/config/ → pkg/config/ (leaf package, no internal deps)
- Update all import paths: agentu/internal/llm → agentu/pkg/llm (14 files)
- Update all import paths: agentu/internal/config → agentu/pkg/config (3 files)
- Update README: build/run commands, CLI flags section
- agent/session/tools/tui stay in internal/ (app-specific business logic)

Build: go build . | Test: go test ./... | Vet: go vet ./...
2026-06-24 07:06:26 -07:00
loveuer 950c63adaa v0.0.5: context compaction, minimal config, retry/search/code symbols/diff/status/test commands 2026-06-24 02:20:19 -07:00
loveuer e4c75c7c0e v0.0.4: add sessions and refine tui 2026-06-23 17:49:36 +08:00
32 changed files with 3907 additions and 495 deletions
+1
View File
@@ -2,3 +2,4 @@ agentu*.yaml
/agentu
/coverage.out
*.swp
chat-history.txt
+31 -22
View File
@@ -20,13 +20,13 @@ model/provider switching, and local tools.
mkdir -p ~/.agentu
cp agentu.example.yaml ~/.agentu/config.yaml
export AGENTU_API_KEY='your-api-key'
go run ./cmd/agentu
go run .
```
Use `--yolo` only when you want agentu to write files or run shell commands:
```sh
go run ./cmd/agentu --yolo
go run . --yolo
```
## Config
@@ -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,29 +58,38 @@ 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
```sh
go run ./cmd/agentu [flags]
go run . [flags]
```
- `--config <path>` loads a YAML config file. Default: `~/.agentu/config.yaml`.
- `--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:
@@ -141,7 +150,7 @@ Run the verification suite:
```sh
go test ./...
go vet ./...
go build ./cmd/agentu
go build .
```
Local config and build outputs are intentionally ignored by git.
+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
-131
View File
@@ -1,131 +0,0 @@
package main
import (
"bufio"
"context"
"errors"
"flag"
"fmt"
"net/http"
"os"
"os/signal"
"strings"
"agentu/internal/agent"
"agentu/internal/config"
"agentu/internal/llm"
"agentu/internal/session"
"agentu/internal/tools"
"agentu/internal/tui"
)
func main() {
if err := run(); err != nil {
fmt.Fprintln(os.Stderr, "agentu:", err)
os.Exit(1)
}
}
func run() error {
configPath := flag.String("config", config.DefaultPath, "path to YAML config")
yolo := flag.Bool("yolo", false, "also enable automatic file writes and shell execution")
plain := flag.Bool("plain", false, "use the simple line-based REPL instead of the TUI")
themeName := flag.String("theme", "light", "TUI theme: light or dark")
flag.Parse()
themeMode, err := tui.ParseThemeMode(*themeName)
if err != nil {
return err
}
cfg, err := config.Load(*configPath)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
resolvedPath, resolveErr := config.ResolvePath(*configPath)
if resolveErr != nil {
return resolveErr
}
return fmt.Errorf("config not found at %s; create ~/.agentu/config.yaml from agentu.example.yaml and set AGENTU_API_KEY", resolvedPath)
}
return err
}
providerName := cfg.DefaultProviderName()
providerConfig := cfg.Providers[providerName]
provider := llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, http.DefaultClient)
var registry *tools.Registry
if *yolo {
registry = tools.Builtins(cfg.Agent.WorkingDir)
} else {
registry = tools.ReadOnlyBuiltins(cfg.Agent.WorkingDir)
}
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))
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,
})
modelManager, err := session.NewManager(cfg, assistant, http.DefaultClient)
if err != nil {
return err
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
defer stop()
if !*plain {
return tui.Run(ctx, assistant, tui.Options{
ModelName: providerConfig.Model,
Yolo: *yolo,
ThemeMode: themeMode,
ModelManager: modelManager,
})
}
return repl(ctx, assistant)
}
func repl(ctx context.Context, assistant *agent.Agent) 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.")
for {
fmt.Print("> ")
if !scanner.Scan() {
if err := scanner.Err(); err != nil {
return err
}
fmt.Println()
return nil
}
input := strings.TrimSpace(scanner.Text())
switch input {
case "":
continue
case "/exit", "/quit":
return nil
case "/clear":
assistant.Clear()
fmt.Println("context cleared")
continue
}
if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
fmt.Fprintln(os.Stderr, "\nerror:", err)
}
fmt.Println()
}
}
+263 -42
View File
@@ -9,24 +9,40 @@ import (
"strings"
"time"
"agentu/internal/llm"
"agentu/pkg/llm"
"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{
@@ -117,14 +142,173 @@ func (a *Agent) Clear() {
}
}
// Messages returns a copy of the current conversation messages.
func (a *Agent) Messages() []llm.Message {
out := make([]llm.Message, len(a.messages))
copy(out, a.messages)
return out
}
// SetMessages replaces the conversation messages (used for session restore).
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
@@ -146,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{
@@ -156,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) {
@@ -180,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 {
+204 -1
View File
@@ -3,10 +3,11 @@ package agent
import (
"context"
"encoding/json"
"fmt"
"strings"
"testing"
"agentu/internal/llm"
"agentu/pkg/llm"
"agentu/internal/tools"
)
@@ -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)
}
}
+186 -23
View File
@@ -5,21 +5,24 @@ import (
"net/http"
"sort"
"strings"
"time"
"agentu/internal/agent"
"agentu/internal/config"
"agentu/internal/llm"
"agentu/pkg/config"
"agentu/pkg/llm"
)
type Manager struct {
cfg *config.Config
agent *agent.Agent
httpClient *http.Client
providerName string
provider config.ProviderConfig
cfg *config.Config
agent *agent.Agent
httpClient *http.Client
providerName string
provider config.ProviderConfig
store *Store
currentSession *Session
}
func NewManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Client) (*Manager, error) {
func NewManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Client, store *Store) (*Manager, error) {
if httpClient == nil {
httpClient = http.DefaultClient
}
@@ -28,13 +31,162 @@ func NewManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Cli
if !ok {
return nil, fmt.Errorf("no provider is configured")
}
return &Manager{
m := &Manager{
cfg: cfg,
agent: assistant,
httpClient: httpClient,
providerName: providerName,
provider: provider,
}, nil
store: store,
}
return m, nil
}
// ResumeLatest loads the most recent session, or creates a new one.
func (m *Manager) ResumeLatest() error {
if m.store == nil {
return nil
}
metas, err := m.store.List()
if err != nil {
return err
}
if len(metas) > 0 {
return m.Resume(metas[0].ID)
}
m.newSession()
return nil
}
// Resume loads a session by ID and restores it into the agent.
func (m *Manager) Resume(id string) error {
if m.store == nil {
return fmt.Errorf("session storage is not configured")
}
sess, err := m.store.Load(id)
if err != nil {
return err
}
m.currentSession = sess
m.agent.SetMessages(sess.Messages)
m.agent.SetLastUsage(sess.LastUsage)
return nil
}
// Save persists the current session to disk.
func (m *Manager) Save() error {
if m.store == nil || m.currentSession == nil {
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
if m.currentSession.Name == "" {
m.currentSession.Name = DefaultName(m.currentSession.Messages)
}
return m.store.Save(m.currentSession)
}
// Rename changes the current session name and saves.
func (m *Manager) Rename(name string) (string, error) {
name = strings.TrimSpace(name)
if name == "" {
return "", fmt.Errorf("name is required")
}
if m.currentSession == nil {
return "", fmt.Errorf("no active session")
}
m.currentSession.Name = name
if err := m.Save(); err != nil {
return "", err
}
return fmt.Sprintf("Renamed session to %q.", name), nil
}
// Sessions returns a formatted list of all sessions.
func (m *Manager) Sessions() (string, error) {
if m.store == nil {
return "", fmt.Errorf("session storage is not configured")
}
metas, err := m.store.List()
if err != nil {
return "", err
}
if len(metas) == 0 {
return "No saved sessions.", nil
}
var b strings.Builder
b.WriteString("Sessions:\n")
for _, meta := range metas {
name := meta.Name
if name == "" {
name = "(unnamed)"
}
current := ""
if m.currentSession != nil && meta.ID == m.currentSession.ID {
current = " (current)"
}
ts := meta.UpdatedAt.Local().Format("2006-01-02 15:04")
fmt.Fprintf(&b, " %s %-30s %s %d msgs%s\n", meta.ID, fitName(name, 30), ts, meta.MessageCount, current)
}
return b.String(), nil
}
// NewSession starts a fresh session, saving the current one first.
func (m *Manager) NewSession() (string, error) {
// Save current if it has user messages
if m.currentSession != nil {
msgs := m.agent.Messages()
hasUser := false
for _, msg := range msgs {
if msg.Role == llm.RoleUser {
hasUser = true
break
}
}
if hasUser {
m.currentSession.Messages = msgs
m.currentSession.Provider = m.providerName
m.currentSession.Model = m.provider.Model
if m.currentSession.Name == "" {
m.currentSession.Name = DefaultName(msgs)
}
if m.store != nil {
_ = m.store.Save(m.currentSession)
}
}
}
m.newSession()
return "Session saved. Starting new session.", nil
}
func (m *Manager) newSession() {
m.currentSession = &Session{
ID: GenerateID(),
CreatedAt: time.Now().UTC(),
Provider: m.providerName,
Model: m.provider.Model,
}
m.agent.Clear()
}
func (m *Manager) CurrentSessionID() string {
if m.currentSession == nil {
return ""
}
return m.currentSession.ID
}
func (m *Manager) CurrentSessionName() string {
if m.currentSession == nil {
return ""
}
if m.currentSession.Name != "" {
return m.currentSession.Name
}
return m.currentSession.ID
}
func (m *Manager) CurrentProvider() string {
@@ -58,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()
@@ -84,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
@@ -112,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(),
})
}
@@ -140,3 +296,10 @@ func contains(values []string, value string) bool {
}
return false
}
func fitName(name string, width int) string {
if len(name) <= width {
return name
}
return name[:width-3] + "..."
}
+245 -6
View File
@@ -2,12 +2,13 @@ package session
import (
"context"
"path/filepath"
"strings"
"testing"
"agentu/internal/agent"
"agentu/internal/config"
"agentu/internal/llm"
"agentu/pkg/config"
"agentu/pkg/llm"
)
type recordingProvider struct {
@@ -19,8 +20,8 @@ func (p *recordingProvider) ChatStream(ctx context.Context, req llm.ChatRequest,
return emit(llm.StreamEvent{Content: "ok"})
}
func TestManagerSwitchesModelAndThinking(t *testing.T) {
cfg := &config.Config{
func testConfig() *config.Config {
return &config.Config{
Providers: map[string]config.ProviderConfig{
"one": {
Name: "one",
@@ -34,6 +35,11 @@ func TestManagerSwitchesModelAndThinking(t *testing.T) {
},
},
}
}
func testManager(t *testing.T) (*Manager, *agent.Agent, *recordingProvider) {
t.Helper()
cfg := testConfig()
provider := &recordingProvider{}
assistant := agent.New(agent.Options{
Provider: provider,
@@ -43,10 +49,19 @@ func TestManagerSwitchesModelAndThinking(t *testing.T) {
ThinkingParam: "thinking",
ThinkingEnabled: true,
})
manager, err := NewManager(cfg, assistant, nil)
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)
}
return manager, assistant, provider
}
func TestManagerSwitchesModelAndThinking(t *testing.T) {
manager, assistant, provider := testManager(t)
if _, err := manager.SetModel("model-b"); err != nil {
t.Fatal(err)
@@ -80,7 +95,8 @@ func TestManagerRejectsUnknownModelAndThinking(t *testing.T) {
},
}
assistant := agent.New(agent.Options{Provider: &recordingProvider{}, Model: "model-a"})
manager, err := NewManager(cfg, assistant, nil)
store, _ := NewStore(filepath.Join(t.TempDir(), "sessions"))
manager, err := NewManager(cfg, assistant, nil, store)
if err != nil {
t.Fatal(err)
}
@@ -92,3 +108,226 @@ func TestManagerRejectsUnknownModelAndThinking(t *testing.T) {
t.Fatal("expected invalid thinking error")
}
}
func TestManagerSaveAndResume(t *testing.T) {
manager, assistant, _ := testManager(t)
// Start new session and have a conversation
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 {
t.Fatalf("save: %v", err)
}
savedID := manager.CurrentSessionID()
if savedID == "" {
t.Fatal("expected session ID after save")
}
// Name should be auto-derived from first user message
if manager.CurrentSessionName() != "hello world" {
t.Fatalf("session name = %q, want %q", manager.CurrentSessionName(), "hello world")
}
// Start a new session
manager.newSession()
if manager.CurrentSessionID() == savedID {
t.Fatal("new session should have different ID")
}
// Resume
if err := manager.Resume(savedID); err != nil {
t.Fatalf("resume: %v", err)
}
if manager.CurrentSessionID() != savedID {
t.Fatalf("after resume, ID = %q, want %q", manager.CurrentSessionID(), savedID)
}
msgs := assistant.Messages()
hasUser := false
for _, m := range msgs {
if m.Role == llm.RoleUser && m.Content == "hello world" {
hasUser = true
}
}
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) {
manager, _, _ := testManager(t)
manager.newSession()
msg, err := manager.Rename("my cool session")
if err != nil {
t.Fatalf("rename: %v", err)
}
if !strings.Contains(msg, "my cool session") {
t.Fatalf("rename message = %q", msg)
}
if manager.CurrentSessionName() != "my cool session" {
t.Fatalf("name = %q", manager.CurrentSessionName())
}
// Verify persisted
_, err = manager.store.Load(manager.CurrentSessionID())
if err != nil {
t.Fatalf("load after rename: %v", err)
}
}
func TestManagerRenameEmpty(t *testing.T) {
manager, _, _ := testManager(t)
manager.newSession()
_, err := manager.Rename(" ")
if err == nil {
t.Fatal("expected error for empty name")
}
}
func TestManagerSessions(t *testing.T) {
manager, _, _ := testManager(t)
// No sessions yet
out, err := manager.Sessions()
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "No saved sessions") {
t.Fatalf("expected 'no saved sessions', got: %s", out)
}
// Create and save a session
manager.newSession()
manager.Save()
out, err = manager.Sessions()
if err != nil {
t.Fatal(err)
}
if strings.Contains(out, "No saved sessions") {
t.Fatal("expected session list")
}
}
func TestManagerNewSession(t *testing.T) {
manager, assistant, _ := testManager(t)
manager.newSession()
var out strings.Builder
assistant.RunTurn(context.Background(), "first message", &out, nil)
msg, err := manager.NewSession()
if err != nil {
t.Fatalf("new session: %v", err)
}
if !strings.Contains(msg, "Starting new session") {
t.Fatalf("message = %q", msg)
}
// Old session should be saved
metas, _ := manager.store.List()
if len(metas) != 1 {
t.Fatalf("expected 1 saved session, got %d", len(metas))
}
}
func TestManagerResumeLatest(t *testing.T) {
manager, assistant, _ := testManager(t)
// Create two sessions
manager.newSession()
var out strings.Builder
assistant.RunTurn(context.Background(), "session one", &out, nil)
manager.Save()
id1 := manager.CurrentSessionID()
manager.NewSession()
assistant.RunTurn(context.Background(), "session two", &out, nil)
manager.Save()
id2 := manager.CurrentSessionID()
// Clear and resume latest
manager.newSession()
if err := manager.ResumeLatest(); err != nil {
t.Fatalf("resume latest: %v", err)
}
// Should resume the most recently updated (id2)
if manager.CurrentSessionID() != id2 {
t.Fatalf("expected to resume %s, got %s", id2, manager.CurrentSessionID())
}
// Also verify id1 is resumable
if err := manager.Resume(id1); err != nil {
t.Fatalf("resume id1: %v", err)
}
if manager.CurrentSessionID() != id1 {
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)
}
}
+174
View File
@@ -0,0 +1,174 @@
package session
import (
"crypto/rand"
"encoding/hex"
"encoding/json"
"fmt"
"os"
"path/filepath"
"sort"
"strings"
"time"
"agentu/pkg/llm"
)
// Session represents a persisted conversation.
type Session struct {
ID string `json:"id"`
Name string `json:"name"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
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.
type SessionMeta struct {
ID string `json:"id"`
Name string `json:"name"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
Provider string `json:"provider"`
Model string `json:"model"`
MessageCount int `json:"message_count"`
}
// Store manages session files under a directory.
type Store struct {
dir string
}
// NewStore creates a store, ensuring the directory exists.
func NewStore(dir string) (*Store, error) {
if err := os.MkdirAll(dir, 0o755); err != nil {
return nil, fmt.Errorf("create session dir: %w", err)
}
return &Store{dir: dir}, nil
}
// Save writes a session to disk.
func (s *Store) Save(sess *Session) error {
if sess.ID == "" {
return fmt.Errorf("session ID is required")
}
sess.UpdatedAt = time.Now().UTC()
if sess.CreatedAt.IsZero() {
sess.CreatedAt = sess.UpdatedAt
}
data, err := json.MarshalIndent(sess, "", " ")
if err != nil {
return fmt.Errorf("marshal session: %w", err)
}
path := s.filePath(sess.ID)
if err := os.WriteFile(path, data, 0o644); err != nil {
return fmt.Errorf("write session: %w", err)
}
return nil
}
// Load reads a session by ID.
func (s *Store) Load(id string) (*Session, error) {
path := s.filePath(id)
data, err := os.ReadFile(path)
if err != nil {
if os.IsNotExist(err) {
return nil, fmt.Errorf("session %q not found", id)
}
return nil, fmt.Errorf("read session: %w", err)
}
var sess Session
if err := json.Unmarshal(data, &sess); err != nil {
return nil, fmt.Errorf("parse session: %w", err)
}
return &sess, nil
}
// List returns metadata for all sessions, sorted by updated_at descending.
func (s *Store) List() ([]SessionMeta, error) {
entries, err := os.ReadDir(s.dir)
if err != nil {
return nil, fmt.Errorf("read session dir: %w", err)
}
var metas []SessionMeta
for _, entry := range entries {
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
continue
}
path := filepath.Join(s.dir, entry.Name())
data, err := os.ReadFile(path)
if err != nil {
continue
}
var sess Session
if err := json.Unmarshal(data, &sess); err != nil {
continue
}
metas = append(metas, SessionMeta{
ID: sess.ID,
Name: sess.Name,
CreatedAt: sess.CreatedAt,
UpdatedAt: sess.UpdatedAt,
Provider: sess.Provider,
Model: sess.Model,
MessageCount: len(sess.Messages),
})
}
sort.Slice(metas, func(i, j int) bool {
return metas[i].UpdatedAt.After(metas[j].UpdatedAt)
})
return metas, nil
}
// Delete removes a session file.
func (s *Store) Delete(id string) error {
path := s.filePath(id)
if err := os.Remove(path); err != nil {
if os.IsNotExist(err) {
return fmt.Errorf("session %q not found", id)
}
return fmt.Errorf("delete session: %w", err)
}
return nil
}
func (s *Store) filePath(id string) string {
return filepath.Join(s.dir, id+".json")
}
// GenerateID creates a random 8-character hex session ID.
func GenerateID() string {
b := make([]byte, 4)
_, _ = rand.Read(b)
return hex.EncodeToString(b)
}
// DefaultName derives a session name from the first user message.
func DefaultName(messages []llm.Message) string {
for _, msg := range messages {
if msg.Role != llm.RoleUser {
continue
}
text := strings.TrimSpace(msg.Content)
if text == "" {
continue
}
// Take first line, truncate at 40 chars at word boundary
if idx := strings.IndexByte(text, '\n'); idx >= 0 {
text = text[:idx]
}
text = strings.TrimSpace(text)
if len(text) > 40 {
text = text[:40]
if idx := strings.LastIndexByte(text, ' '); idx > 20 {
text = text[:idx]
}
text += "..."
}
return text
}
return ""
}
+174
View File
@@ -0,0 +1,174 @@
package session
import (
"os"
"path/filepath"
"testing"
"agentu/pkg/llm"
)
func tempStore(t *testing.T) *Store {
t.Helper()
dir := t.TempDir()
store, err := NewStore(dir)
if err != nil {
t.Fatal(err)
}
return store
}
func TestSaveAndLoad(t *testing.T) {
store := tempStore(t)
sess := &Session{
ID: GenerateID(),
Name: "test session",
Provider: "loveuer",
Model: "deepseek-v4-flash",
Messages: []llm.Message{
{Role: llm.RoleSystem, Content: "you are helpful"},
{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)
}
loaded, err := store.Load(sess.ID)
if err != nil {
t.Fatalf("load: %v", err)
}
if loaded.ID != sess.ID {
t.Errorf("ID = %q, want %q", loaded.ID, sess.ID)
}
if loaded.Name != sess.Name {
t.Errorf("Name = %q, want %q", loaded.Name, sess.Name)
}
if len(loaded.Messages) != 3 {
t.Errorf("Messages len = %d, want 3", len(loaded.Messages))
}
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")
}
}
func TestLoadNotFound(t *testing.T) {
store := tempStore(t)
_, err := store.Load("nonexistent")
if err == nil {
t.Fatal("expected error for nonexistent session")
}
}
func TestList(t *testing.T) {
store := tempStore(t)
// Empty list
metas, err := store.List()
if err != nil {
t.Fatalf("list: %v", err)
}
if len(metas) != 0 {
t.Errorf("expected empty list, got %d", len(metas))
}
// Save two sessions
store.Save(&Session{ID: "aaa", Name: "first", Messages: []llm.Message{{Role: "user", Content: "a"}}})
store.Save(&Session{ID: "bbb", Name: "second", Messages: []llm.Message{{Role: "user", Content: "b"}, {Role: "assistant", Content: "c"}}})
metas, err = store.List()
if err != nil {
t.Fatalf("list: %v", err)
}
if len(metas) != 2 {
t.Fatalf("expected 2 sessions, got %d", len(metas))
}
// Sorted by updated_at desc
if metas[0].ID != "bbb" && metas[0].ID != "aaa" {
t.Errorf("unexpected first session ID: %q", metas[0].ID)
}
}
func TestDelete(t *testing.T) {
store := tempStore(t)
store.Save(&Session{ID: "todelete", Name: "temp"})
if err := store.Delete("todelete"); err != nil {
t.Fatalf("delete: %v", err)
}
_, err := store.Load("todelete")
if err == nil {
t.Fatal("expected error after delete")
}
}
func TestDeleteNotFound(t *testing.T) {
store := tempStore(t)
err := store.Delete("nonexistent")
if err == nil {
t.Fatal("expected error for nonexistent delete")
}
}
func TestSaveRequiresID(t *testing.T) {
store := tempStore(t)
err := store.Save(&Session{Name: "no id"})
if err == nil {
t.Fatal("expected error for empty ID")
}
}
func TestGenerateID(t *testing.T) {
id := GenerateID()
if len(id) != 8 {
t.Errorf("ID length = %d, want 8", len(id))
}
id2 := GenerateID()
if id == id2 {
t.Error("expected unique IDs")
}
}
func TestDefaultName(t *testing.T) {
tests := []struct {
name string
messages []llm.Message
want string
}{
{"empty", nil, ""},
{"no user msg", []llm.Message{{Role: "system", Content: "hi"}}, ""},
{"short", []llm.Message{{Role: "user", Content: "hello world"}}, "hello world"},
{"long", []llm.Message{{Role: "user", Content: "this is a really long message that should be truncated at some point in the middle"}}, "this is a really long message that..."},
{"multiline", []llm.Message{{Role: "user", Content: "first line\nsecond line"}}, "first line"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := DefaultName(tt.messages)
if got != tt.want {
t.Errorf("DefaultName = %q, want %q", got, tt.want)
}
})
}
}
func TestNewStoreCreatesDir(t *testing.T) {
dir := filepath.Join(t.TempDir(), "nested", "sessions")
store, err := NewStore(dir)
if err != nil {
t.Fatalf("NewStore: %v", err)
}
info, err := os.Stat(store.dir)
if err != nil {
t.Fatalf("stat: %v", err)
}
if !info.IsDir() {
t.Error("expected directory")
}
}
+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/pkg/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()
}
+62 -11
View File
@@ -13,7 +13,7 @@ import (
"sort"
"strings"
"agentu/internal/llm"
"agentu/pkg/llm"
)
type FileReadTool struct {
@@ -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])
}
}
+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
}
+1 -1
View File
@@ -9,7 +9,7 @@ import (
"os/exec"
"strings"
"agentu/internal/llm"
"agentu/pkg/llm"
)
type ShellRunTool struct {
+43 -6
View File
@@ -5,12 +5,15 @@ import (
"encoding/json"
"fmt"
"sort"
"strings"
"agentu/internal/llm"
"agentu/pkg/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 {
+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 ..")
}
}
+1 -1
View File
@@ -10,7 +10,7 @@ import (
"net/url"
"strings"
"agentu/internal/llm"
"agentu/pkg/llm"
"golang.org/x/net/html"
)
+649 -72
View File
@@ -6,9 +6,12 @@ import (
"fmt"
"io"
"math"
"sort"
"strings"
"unicode"
"agentu/internal/agent"
"agentu/pkg/llm"
"github.com/charmbracelet/bubbles/spinner"
"github.com/charmbracelet/bubbles/textarea"
@@ -18,8 +21,10 @@ import (
)
const (
minInputLines = 1
maxInputLines = 6
minInputLines = 1
maxInputLines = 6
maxCommandPaletteItems = 8
maxInputHistory = 100
)
type Options struct {
@@ -27,16 +32,24 @@ type Options struct {
Yolo bool
ThemeMode ThemeMode
ModelManager ModelManager
WorkingDir string
}
type ModelManager interface {
CurrentProvider() string
CurrentModel() string
CurrentThinking() string
CurrentSessionID() string
CurrentSessionName() string
Info() string
SetProvider(name string) (string, error)
SetModel(model string) (string, error)
SetThinking(level string) (string, error)
Save() error
Resume(id string) error
Rename(name string) (string, error)
Sessions() (string, error)
NewSession() (string, error)
}
func Run(ctx context.Context, assistant *agent.Agent, opts Options) error {
@@ -61,13 +74,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
@@ -76,11 +90,28 @@ 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
inputHistory []string
historyIndex int
draftInput string
running bool
cancel context.CancelFunc
events <-chan tea.Msg
// Session picker state
picking bool
pickerItems []pickerItem
pickerCursor int
}
type pickerItem struct {
id string
name string
detail string
isCurrent bool
}
type assistantChunkMsg string
@@ -90,16 +121,31 @@ type turnDoneMsg struct {
err error
}
var errOpenPicker = fmt.Errorf("open picker")
type slashCommand struct {
Name string
Description string
}
type commandMatch struct {
command slashCommand
score int
}
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"},
{Name: "/sessions", Description: "list sessions"},
{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 {
@@ -128,25 +174,55 @@ 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",
historyIndex: -1,
}
m.refreshContextUsage()
return m
}
func (m model) Init() tea.Cmd {
m.syncMessages()
return textarea.Blink
}
// syncMessages rebuilds the TUI display messages from the agent's conversation history.
func (m *model) syncMessages() {
if m.agent == nil {
return
}
msgs := m.agent.Messages()
m.messages = m.messages[:0]
for _, msg := range msgs {
switch msg.Role {
case llm.RoleUser:
m.messages = append(m.messages, message{role: roleUser, content: msg.Content})
case llm.RoleAssistant:
if msg.Content != "" {
m.messages = append(m.messages, message{role: roleAssistant, content: msg.Content})
}
case llm.RoleTool:
text := cleanToolLog(msg.Content)
if text != "" {
m.messages = append(m.messages, message{role: roleTool, content: text})
}
}
}
m.refreshViewport(true)
}
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
var cmds []tea.Cmd
@@ -167,6 +243,9 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, cmd
case tea.KeyMsg:
if m.picking {
return m.updatePicker(msg)
}
switch msg.String() {
case "ctrl+c":
if m.cancel != nil {
@@ -184,6 +263,20 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, nil
}
return m.submit()
case "ctrl+p":
if m.running {
return m, nil
}
m.previousInputHistory()
m.afterInputChanged()
return m, nil
case "ctrl+n":
if m.running {
return m, nil
}
m.nextInputHistory()
m.afterInputChanged()
return m, nil
case "alt+enter", "ctrl+j":
return m.insertInputNewline()
}
@@ -213,6 +306,11 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
} else {
m.status = "ready"
}
// Auto-save session after turn completes
if m.models != nil {
_ = m.models.Save()
}
m.refreshContextUsage()
m.resize()
m.refreshViewport(true)
cmds = append(cmds, textarea.Blink)
@@ -227,6 +325,8 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
if !m.running {
nextInput, cmd := m.input.Update(msg)
m.input = nextInput
m.historyIndex = -1
m.draftInput = ""
m.afterInputChanged()
return m, cmd
}
@@ -252,7 +352,14 @@ func (m model) View() string {
return "agentu"
}
parts := []string{m.headerView(), m.viewport.View()}
if m.picking {
parts := []string{m.viewport.View()}
parts = append(parts, m.pickerView(), m.modelMetaView(), m.footerView())
body := lipgloss.JoinVertical(lipgloss.Left, parts...)
return m.styles.App.Width(m.width).Height(m.height).Render(body)
}
parts := []string{m.viewport.View()}
if activity := m.activityView(); activity != "" {
parts = append(parts, activity)
}
@@ -266,19 +373,21 @@ func (m *model) submit() (tea.Model, tea.Cmd) {
if input == "" {
return *m, nil
}
m.rememberInput(input)
switch input {
case "/exit", "/quit":
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)
}
@@ -324,6 +433,8 @@ func (m *model) insertInputNewline() (tea.Model, tea.Cmd) {
enter := tea.KeyMsg{Type: tea.KeyEnter}
nextInput, cmd := m.input.Update(enter)
m.input = nextInput
m.historyIndex = -1
m.draftInput = ""
m.afterInputChanged()
return *m, cmd
}
@@ -354,11 +465,10 @@ func (m *model) appendAssistantChunk(chunk string) {
func (m *model) resize() {
width := max(40, m.width)
inputWidth := max(20, width-4)
inputWidth := max(20, width-m.styles.InputBox.GetHorizontalFrameSize())
m.input.SetWidth(inputWidth)
m.syncInputHeight()
headerHeight := 1
activityHeight := 0
if m.running {
activityHeight = 1
@@ -366,7 +476,7 @@ func (m *model) resize() {
inputHeight := m.inputBlockHeight()
footerHeight := 2
m.viewport.Width = width
m.viewport.Height = max(4, m.height-headerHeight-activityHeight-inputHeight-footerHeight)
m.viewport.Height = max(4, m.height-activityHeight-inputHeight-footerHeight)
m.viewport.Style = m.styles.Viewport
}
@@ -379,7 +489,7 @@ func (m *model) refreshViewport(bottom bool) {
func (m model) renderMessages() string {
width := max(40, m.width)
contentWidth := max(20, width-6)
contentWidth := width
parts := make([]string, 0, len(m.messages))
for _, msg := range m.messages {
content := strings.TrimRight(msg.content, "\n")
@@ -395,14 +505,17 @@ func (m model) renderMessages() string {
}
func (m model) messageView(r role, content string, width int) string {
if r == roleTool {
return m.toolMessageView(content, width)
}
label := roleLabel(r)
style := m.styles.AssistantMsg
labelStyle := m.styles.AssistantLabel
labelStyle := m.styles.SystemLabel
switch r {
case roleUser:
style = m.styles.UserMessage
labelStyle = m.styles.UserLabel
case roleTool:
style = m.styles.ToolMessage
labelStyle = m.styles.ToolLabel
@@ -420,17 +533,33 @@ func (m model) messageView(r role, content string, width int) string {
bodyContent = renderMarkdown(content, bodyWidth, m.theme)
}
bodyWidth := messageBodyWidth(bodyContent, width, style)
if r == roleUser {
bodyWidth = width
}
body := style.Width(bodyWidth).Render(bodyContent)
if label == "" {
return body
}
labelLine := labelStyle.Render(label)
body := style.Width(width).Render(bodyContent)
return lipgloss.JoinVertical(lipgloss.Left, labelLine, body)
}
func messageBodyWidth(content string, maxWidth int, style lipgloss.Style) int {
maxWidth = max(1, maxWidth)
contentWidth := 1
for _, line := range strings.Split(content, "\n") {
contentWidth = max(contentWidth, lipgloss.Width(line))
}
return min(maxWidth, contentWidth+style.GetHorizontalFrameSize())
}
func roleLabel(r role) string {
switch r {
case roleUser:
return "You"
return ""
case roleAssistant:
return "Agentu"
return ""
case roleTool:
return "Tool"
case roleError:
@@ -442,11 +571,72 @@ func roleLabel(r role) string {
}
}
func (m model) headerView() string {
title := m.styles.Title.Render("agentu")
meta := m.styles.Muted.Render("local agent")
line := lipgloss.JoinHorizontal(lipgloss.Center, title, " ", meta)
return m.styles.Header.Width(max(1, m.width)).Render(line)
type toolLogParts struct {
name string
args string
}
func parseToolLog(content string) toolLogParts {
content = strings.TrimSpace(content)
if content == "" {
return toolLogParts{}
}
idx := strings.IndexFunc(content, unicode.IsSpace)
if idx < 0 {
return toolLogParts{name: content}
}
return toolLogParts{
name: strings.TrimSpace(content[:idx]),
args: strings.TrimSpace(content[idx+1:]),
}
}
func (m model) toolMessageView(content string, width int) string {
parts := parseToolLog(content)
header := m.styles.ToolName.Render(strings.TrimSpace(iconTool + " " + parts.name))
lines := []string{header}
if parts.args != "" {
args := fitLine(parts.args, max(1, width-4))
lines = append(lines, m.styles.ToolValue.Render(args))
}
bodyContent := strings.Join(lines, "\n")
bodyWidth := messageBodyWidth(bodyContent, width, m.styles.ToolMessage)
return m.styles.ToolMessage.Width(bodyWidth).Render(bodyContent)
}
type metaPill struct {
text string
accent bool
}
func (m model) renderMetaPills(parts []metaPill) string {
if len(parts) == 0 {
return ""
}
plain := make([]string, len(parts))
for i, part := range parts {
plain[i] = part.text
}
width := max(1, m.width-2)
visible := len(plain)
for visible > 1 && lipgloss.Width(strings.Join(plain[:visible], " ")) > width {
visible--
}
if lipgloss.Width(strings.Join(plain[:visible], " ")) > width {
return m.styles.PillAccent.Render(fitLine(plain[0], width))
}
rendered := make([]string, 0, visible)
for i := 0; i < visible; i++ {
style := m.styles.Pill
if parts[i].accent {
style = m.styles.PillAccent
}
rendered = append(rendered, style.Render(parts[i].text))
}
return strings.Join(rendered, " ")
}
func (m model) modelMetaView() string {
@@ -454,28 +644,39 @@ func (m model) modelMetaView() string {
if m.yolo {
mode = "yolo tools"
}
modelName := m.modelName
parts := []string{}
if m.models != nil {
if provider := m.models.CurrentProvider(); provider != "" {
parts = append(parts, "provider "+provider)
}
modelName = m.models.CurrentModel()
if currentThinking := m.models.CurrentThinking(); currentThinking != "" && currentThinking != "unset" {
parts = append(parts, "thinking "+currentThinking)
if currentModel := m.models.CurrentModel(); currentModel != "" {
modelName = currentModel
}
}
parts = append([]string{modelName, mode}, parts...)
parts = append(parts, "theme "+string(m.theme.Mode))
line := fitLine(strings.Join(parts, " · "), max(1, m.width-2))
return m.styles.ModelMeta.Width(max(1, m.width)).Render(line)
parts := []metaPill{{text: iconModel + " " + modelName, accent: true}}
if m.models != nil {
if provider := m.models.CurrentProvider(); provider != "" {
parts = append(parts, metaPill{text: "provider " + provider})
}
if currentThinking := m.models.CurrentThinking(); currentThinking != "" && currentThinking != "unset" {
parts = append(parts, metaPill{text: "thinking " + currentThinking})
}
if sessionName := m.models.CurrentSessionName(); sessionName != "" {
parts = append(parts, metaPill{text: iconSession + " " + sessionName})
}
}
if m.contextUsage != "" {
parts = append(parts, metaPill{text: iconContext + " " + m.contextUsage})
}
parts = append(parts, metaPill{text: mode}, metaPill{text: "theme " + string(m.theme.Mode)})
return m.styles.ModelMeta.Width(max(1, m.width)).Render(m.renderMetaPills(parts))
}
func (m model) inputView() string {
suggestions := m.slashSuggestionsView()
palette := m.commandPaletteView()
box := m.styles.InputBox.Width(max(20, m.width-2)).Render(m.input.View())
if suggestions != "" {
return lipgloss.JoinVertical(lipgloss.Left, suggestions, box)
if palette != "" {
return lipgloss.JoinVertical(lipgloss.Left, palette, box)
}
return box
}
@@ -488,12 +689,12 @@ func (m model) activityView() string {
if status == "" || status == "ready" {
status = "working..."
}
line := fitLine(fmt.Sprintf("%s Working · %s · esc cancels", m.spinner.View(), status), max(1, m.width-2))
return m.styles.Activity.Width(max(1, m.width)).Render(line)
line := m.styles.StatusInfo.Render(iconWorking) + " Working " + m.spinner.View() + " " + status + " · esc cancels"
return m.styles.Activity.Render(fitLine(line, max(1, m.width-2)))
}
func (m model) footerView() string {
hint := "enter send · ctrl+j newline · alt+enter newline · / commands"
hint := "enter send · ctrl+j newline · alt+enter newline · ctrl+p/ctrl+n history · / commands"
if m.running {
hint = "esc cancel · ctrl+c quit"
}
@@ -548,6 +749,51 @@ func (m *model) afterInputChanged() {
}
}
func (m *model) rememberInput(input string) {
input = strings.TrimSpace(input)
if input == "" {
return
}
if len(m.inputHistory) > 0 && m.inputHistory[len(m.inputHistory)-1] == input {
m.historyIndex = -1
m.draftInput = ""
return
}
m.inputHistory = append(m.inputHistory, input)
if len(m.inputHistory) > maxInputHistory {
m.inputHistory = m.inputHistory[len(m.inputHistory)-maxInputHistory:]
}
m.historyIndex = -1
m.draftInput = ""
}
func (m *model) previousInputHistory() {
if len(m.inputHistory) == 0 {
return
}
if m.historyIndex == -1 {
m.draftInput = m.input.Value()
m.historyIndex = len(m.inputHistory) - 1
} else if m.historyIndex > 0 {
m.historyIndex--
}
m.input.SetValue(m.inputHistory[m.historyIndex])
}
func (m *model) nextInputHistory() {
if len(m.inputHistory) == 0 || m.historyIndex == -1 {
return
}
if m.historyIndex < len(m.inputHistory)-1 {
m.historyIndex++
m.input.SetValue(m.inputHistory[m.historyIndex])
return
}
m.input.SetValue(m.draftInput)
m.historyIndex = -1
m.draftInput = ""
}
func (m *model) syncInputHeight() {
m.input.SetHeight(m.desiredInputHeight())
}
@@ -570,33 +816,132 @@ func (m model) desiredInputHeight() int {
func (m model) inputBlockHeight() int {
height := m.input.Height() + m.styles.InputBox.GetVerticalFrameSize()
if m.showSlashSuggestions() {
height++
if m.showCommandPalette() {
height += m.commandPaletteHeight()
}
return height
}
func (m model) showSlashSuggestions() bool {
func (m model) commandPaletteHeight() int {
matches := commandMatches(commandPaletteQuery(m.input.Value()))
count := len(matches)
if count > maxCommandPaletteItems {
count = maxCommandPaletteItems
}
if count == 0 {
count = 1
}
return 1 + count + m.styles.Palette.GetVerticalFrameSize()
}
func (m model) showCommandPalette() bool {
value := strings.TrimSpace(m.input.Value())
return strings.HasPrefix(value, "/") && !strings.Contains(value, " ")
}
func (m model) slashSuggestionsView() string {
if !m.showSlashSuggestions() {
func (m model) commandPaletteView() string {
if !m.showCommandPalette() {
return ""
}
prefix := strings.TrimSpace(m.input.Value())
var parts []string
matches := commandMatches(commandPaletteQuery(m.input.Value()))
lines := []string{m.styles.Header.Render(iconCommand + " Commands")}
if len(matches) == 0 {
lines = append(lines, m.styles.Muted.Render("No matching commands"))
} else {
for i, match := range matches {
if i >= maxCommandPaletteItems {
break
}
lines = append(lines, fmt.Sprintf("%s %s", m.styles.PaletteMatch.Render(match.command.Name), m.styles.Muted.Render(match.command.Description)))
}
}
return m.styles.Palette.Width(max(20, m.width-2)).Render(strings.Join(lines, "\n"))
}
func commandPaletteQuery(value string) string {
value = strings.TrimSpace(value)
if !strings.HasPrefix(value, "/") {
return ""
}
return strings.TrimPrefix(value, "/")
}
func commandMatches(query string) []commandMatch {
matches := make([]commandMatch, 0, len(slashCommands))
for _, command := range slashCommands {
if !strings.HasPrefix(command.Name, prefix) {
score, ok := scoreCommand(command, query)
if !ok {
continue
}
parts = append(parts, fmt.Sprintf("%s %s", command.Name, command.Description))
matches = append(matches, commandMatch{command: command, score: score})
}
if len(parts) == 0 {
parts = append(parts, "No matching commands")
sort.Slice(matches, func(i, j int) bool {
if matches[i].score != matches[j].score {
return matches[i].score < matches[j].score
}
return matches[i].command.Name < matches[j].command.Name
})
return matches
}
func scoreCommand(command slashCommand, query string) (int, bool) {
query = strings.ToLower(strings.TrimSpace(query))
name := strings.TrimPrefix(strings.ToLower(command.Name), "/")
if query == "" || query == name {
return 0, true
}
return m.styles.CommandHint.Width(max(20, m.width-2)).Render(strings.Join(parts, " "))
if strings.HasPrefix(name, query) {
return 100 + len(name) - len(query), true
}
if score, ok := subsequenceScore(name, query, 200); ok {
return score, true
}
candidate := name + " " + strings.ToLower(command.Description)
if score, ok := subsequenceScore(candidate, query, 300); ok {
return score, true
}
return 0, false
}
func subsequenceScore(candidate string, query string, base int) (int, bool) {
candidate = strings.ToLower(candidate)
query = strings.ToLower(query)
if query == "" {
return base, true
}
start := 0
last := -1
gaps := 0
for _, r := range query {
found := -1
for i, c := range candidate[start:] {
if c == r {
found = start + i
break
}
}
if found < 0 {
return 0, false
}
if last >= 0 {
gaps += found - last - 1
}
last = found
start = found + len(string(r))
}
return base + gaps, true
}
func (m model) slashSuggestionCount() int {
return m.commandPaletteHeight()
}
func (m model) showSlashSuggestions() bool {
return m.showCommandPalette()
}
func (m model) slashSuggestionsView() string {
return m.commandPaletteView()
}
func (m model) modelInfo() string {
@@ -610,6 +955,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 {
@@ -635,6 +1032,10 @@ func fitLine(text string, width int) string {
func (m *model) runLocalCommand(input string) (tea.Model, tea.Cmd) {
output, err := m.handleLocalCommand(input)
if err == errOpenPicker {
m.openSessionPicker()
return *m, nil
}
role := roleSystem
status := "command"
if err != nil {
@@ -642,10 +1043,14 @@ func (m *model) runLocalCommand(input string) (tea.Model, tea.Cmd) {
role = roleError
status = "command error"
}
m.messages = append(m.messages, message{role: role, content: output})
m.input.SetValue("")
m.afterInputChanged()
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)
return *m, nil
}
@@ -656,13 +1061,39 @@ 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":
return m.handleSessionsCommand()
case "/resume":
return m.handleResumeCommand(fields[1:])
case "/rename":
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 {
@@ -697,4 +1128,150 @@ func (m model) handleModelCommand(args []string) (string, error) {
}
}
func (m model) handleSessionsCommand() (string, error) {
if m.models == nil {
return "", fmt.Errorf("session management is not configured")
}
return m.models.Sessions()
}
func (m model) handleResumeCommand(args []string) (string, error) {
if m.models == nil {
return "", fmt.Errorf("session management is not configured")
}
if len(args) > 0 {
id := strings.TrimSpace(args[0])
if err := m.models.Resume(id); err != nil {
return "", err
}
return fmt.Sprintf("Resumed session %q.", m.models.CurrentSessionName()), nil
}
// No args — open interactive picker (handled in runLocalCommand)
return "", errOpenPicker
}
func (m model) handleRenameCommand(args []string) (string, error) {
if m.models == nil {
return "", fmt.Errorf("session management is not configured")
}
if len(args) == 0 {
return "", fmt.Errorf("usage: /rename <name>")
}
name := strings.Join(args, " ")
return m.models.Rename(name)
}
func (m model) handleNewCommand() (string, error) {
if m.models == nil {
return "", fmt.Errorf("session management is not configured")
}
return m.models.NewSession()
}
func (m *model) openSessionPicker() {
if m.models == nil {
return
}
output, err := m.models.Sessions()
if err != nil {
m.messages = append(m.messages, message{role: roleError, content: err.Error()})
return
}
// Parse session list into picker items
m.pickerItems = m.pickerItems[:0]
m.pickerCursor = 0
for _, line := range strings.Split(output, "\n") {
line = strings.TrimSpace(line)
if line == "" || line == "Sessions:" || strings.HasPrefix(line, "No saved") {
continue
}
// Format: " id name date msgs (current)"
fields := strings.Fields(line)
if len(fields) < 2 {
continue
}
id := fields[0]
isCurrent := strings.Contains(line, "(current)")
// Reconstruct display
m.pickerItems = append(m.pickerItems, pickerItem{
id: id,
name: line,
isCurrent: isCurrent,
})
}
if len(m.pickerItems) == 0 {
m.messages = append(m.messages, message{role: roleSystem, content: "No saved sessions."})
m.refreshViewport(true)
return
}
// Skip current session in picker, pre-select next one
for i, item := range m.pickerItems {
if !item.isCurrent {
m.pickerCursor = i
break
}
}
m.picking = true
m.input.SetValue("")
m.input.Blur()
m.refreshViewport(true)
}
func (m *model) updatePicker(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
switch msg.String() {
case "ctrl+c", "esc":
m.picking = false
m.input.Focus()
return *m, textarea.Blink
case "up", "k":
if m.pickerCursor > 0 {
m.pickerCursor--
}
return *m, nil
case "down", "j":
if m.pickerCursor < len(m.pickerItems)-1 {
m.pickerCursor++
}
return *m, nil
case "enter":
if m.pickerCursor >= 0 && m.pickerCursor < len(m.pickerItems) {
item := m.pickerItems[m.pickerCursor]
m.picking = false
m.input.Focus()
if err := m.models.Resume(item.id); err != nil {
m.messages = append(m.messages, message{role: roleError, content: err.Error()})
} else {
m.syncMessages()
m.messages = append(m.messages, message{role: roleSystem, content: fmt.Sprintf("Resumed session %q.", m.models.CurrentSessionName())})
}
m.refreshViewport(true)
}
return *m, textarea.Blink
}
return *m, nil
}
func (m model) pickerView() string {
if !m.picking || len(m.pickerItems) == 0 {
return ""
}
width := max(40, m.width-4)
var b strings.Builder
b.WriteString(m.styles.Muted.Render("Select a session (↑/↓ navigate, enter select, esc cancel)"))
b.WriteString("\n\n")
for i, item := range m.pickerItems {
line := item.name
if item.isCurrent {
line += " (current)"
}
if i == m.pickerCursor {
b.WriteString(m.styles.UserLabel.Render("▸ " + line))
} else {
b.WriteString(m.styles.Muted.Render(" " + line))
}
b.WriteString("\n")
}
return m.styles.InputBox.Width(width).Render(b.String())
}
var _ io.Writer = eventWriter{}
+333 -23
View File
@@ -2,22 +2,75 @@ package tui
import (
"context"
"os"
"os/exec"
"path/filepath"
"strings"
"testing"
"agentu/internal/agent"
"agentu/pkg/llm"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
)
func TestModelRendersChatShell(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
rendered := next.(model).View()
m = next.(model)
rendered := m.View()
for _, want := range []string{"agentu", "test-model", "yolo tools", "light", "ctrl+j", "alt+enter"} {
for _, want := range []string{"test-model", "yolo tools", "light", "ctrl+j", "alt+enter", "ctrl+p/ctrl+n history"} {
if !strings.Contains(rendered, want) {
t.Fatalf("rendered view missing %q:\n%s", want, rendered)
}
}
if strings.Contains(rendered, "local agent") {
t.Fatalf("top header should not render:\n%s", rendered)
}
if _, ok := m.styles.App.GetBackground().(lipgloss.NoColor); !ok {
t.Fatalf("app should not paint a global background")
}
if _, ok := m.styles.Viewport.GetBackground().(lipgloss.NoColor); !ok {
t.Fatalf("viewport should not paint an assistant background")
}
}
func TestInputTextDoesNotPaintBackground(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
for name, color := range map[string]any{
"input box": m.styles.InputBox.GetBackground(),
"focused base": m.input.FocusedStyle.Base.GetBackground(),
"focused cursor line": m.input.FocusedStyle.CursorLine.GetBackground(),
"blurred base": m.input.BlurredStyle.Base.GetBackground(),
"blurred cursor line": m.input.BlurredStyle.CursorLine.GetBackground(),
} {
if _, ok := color.(lipgloss.NoColor); !ok {
t.Fatalf("%s should not paint input text background", name)
}
}
}
func TestInputBoxIsRoomier(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
if got := m.styles.InputBox.GetVerticalFrameSize(); got != 2 {
t.Fatalf("input box vertical frame should stay compact, got %d", got)
}
if got := m.styles.InputBox.GetHorizontalFrameSize(); got < 6 {
t.Fatalf("input box horizontal frame should include roomier padding, got %d", got)
}
if got, want := m.input.Width(), 100-m.styles.InputBox.GetHorizontalFrameSize(); got != want {
t.Fatalf("input content width = %d, want %d", got, want)
}
rendered := m.inputView()
line := strings.Split(rendered, "\n")[0]
if got := lipgloss.Width(line); got != 100 {
t.Fatalf("input box rendered width = %d, want 100:\n%s", got, rendered)
}
}
func TestModelMetaRendersBelowInput(t *testing.T) {
@@ -35,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})
@@ -56,43 +123,97 @@ func TestRunningStateRendersActivityIndicator(t *testing.T) {
m.resize()
rendered := m.View()
for _, want := range []string{"Working", "thinking...", "esc cancels"} {
for _, want := range []string{iconWorking, "Working", "thinking...", "esc cancels"} {
if !strings.Contains(rendered, want) {
t.Fatalf("running view missing %q:\n%s", want, rendered)
}
}
if backgroundPattern.MatchString(m.activityView()) {
t.Fatalf("activity view should not paint a background:\n%q", m.activityView())
}
}
func TestUserMessageRendersOnLeft(t *testing.T) {
func TestUserMessageRendersFullWidthBubble(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
rendered := m.messageView(roleUser, "hello", 80)
if strings.HasPrefix(rendered, " ") {
t.Fatalf("user message should not be right-aligned:\n%q", rendered)
}
if !strings.Contains(rendered, "You") || !strings.Contains(rendered, "hello") {
rendered := m.messageView(roleUser, "hi", 80)
if !strings.Contains(rendered, "hi") {
t.Fatalf("user message missing content:\n%s", rendered)
}
if strings.Contains(rendered, "You") {
t.Fatalf("user message should not render label:\n%s", rendered)
}
if !strings.Contains(rendered, " hi") || !strings.Contains(rendered, "hi ") {
t.Fatalf("user message should have horizontal padding:\n%q", rendered)
}
lines := strings.Split(rendered, "\n")
if len(lines) < 3 {
t.Fatalf("user message background should include vertical padding:\n%q", rendered)
}
for _, line := range lines {
if got := lipgloss.Width(line); got != 80 {
t.Fatalf("user message background should span full width, line width = %d:\n%q", got, rendered)
}
}
m.width = 80
m.messages = []message{{role: roleUser, content: "hi"}}
rendered = m.renderMessages()
for _, line := range strings.Split(rendered, "\n") {
if got := lipgloss.Width(line); got != 80 {
t.Fatalf("rendered user transcript should span full width, line width = %d:\n%q", got, rendered)
}
}
}
func TestAssistantMessagesRenderMarkdown(t *testing.T) {
func TestAssistantMessagesRenderMarkdownWithoutChrome(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
rendered := m.messageView(roleAssistant, "# Title\n\n- item", 80)
for _, want := range []string{"Agentu", "Title", "item"} {
for _, want := range []string{"Title", "item"} {
if !strings.Contains(rendered, want) {
t.Fatalf("assistant markdown missing %q:\n%s", want, rendered)
}
}
for _, unwanted := range []string{"Agentu", "│", "┃", "|"} {
if strings.Contains(rendered, unwanted) {
t.Fatalf("assistant message should not render chrome %q:\n%s", unwanted, rendered)
}
}
if strings.Contains(rendered, "# Title") {
t.Fatalf("assistant markdown heading was not rendered:\n%s", rendered)
}
}
func TestAssistantMessageRendersPaddedBodyWithoutBackground(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
if _, ok := m.styles.AssistantMsg.GetBackground().(lipgloss.NoColor); !ok {
t.Fatalf("assistant message should not set a background color")
}
rendered := m.messageView(roleAssistant, "hello", 80)
if !strings.Contains(rendered, "hello") {
t.Fatalf("assistant message missing content:\n%s", rendered)
}
lines := strings.Split(rendered, "\n")
if len(lines) < 3 {
t.Fatalf("assistant message should include vertical padding:\n%q", rendered)
}
if got := lipgloss.Width(rendered); got <= len("hello") {
t.Fatalf("assistant message should be wider than text, width = %d:\n%q", got, rendered)
}
for _, line := range lines {
if got := lipgloss.Width(line); got >= 80 {
t.Fatalf("assistant message should not fill the row, line width = %d:\n%q", got, rendered)
}
}
}
func TestUserMessagesDoNotRenderMarkdown(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
@@ -106,7 +227,7 @@ func TestUserMessagesDoNotRenderMarkdown(t *testing.T) {
}
}
func TestMessageRoleLabels(t *testing.T) {
func TestMessageRoleLabelsDoNotLabelUserOrAssistant(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
@@ -116,9 +237,14 @@ func TestMessageRoleLabels(t *testing.T) {
}
rendered := m.renderMessages()
for _, want := range []string{"You", "Agentu"} {
for _, want := range []string{"hello", "hi"} {
if !strings.Contains(rendered, want) {
t.Fatalf("messages missing %q:\n%s", want, rendered)
t.Fatalf("messages missing content %q:\n%s", want, rendered)
}
}
for _, unwanted := range []string{"You", "Agentu"} {
if strings.Contains(rendered, unwanted) {
t.Fatalf("message should not render label %q:\n%s", unwanted, rendered)
}
}
}
@@ -153,19 +279,168 @@ func TestInputStartsAtOneLineAndGrows(t *testing.T) {
}
}
func TestSlashCommandSuggestions(t *testing.T) {
func TestCommandPaletteFuzzyMatches(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/")
m.input.SetValue("/mdl")
m.afterInputChanged()
rendered := m.View()
for _, want := range []string{"/clear reset context", "/exit quit", "/model model/provider/thinking"} {
if !strings.Contains(rendered, want) {
t.Fatalf("rendered suggestions missing %q:\n%s", want, rendered)
plain := stripANSI(m.View())
for _, want := range []string{iconCommand + " Commands", "/model model/provider/thinking"} {
if !strings.Contains(plain, want) {
t.Fatalf("rendered palette missing %q:\n%s", want, plain)
}
}
matches := commandMatches("m")
modelIndex := -1
renameIndex := -1
for i, match := range matches {
switch match.command.Name {
case "/model":
modelIndex = i
case "/rename":
renameIndex = i
}
}
if modelIndex < 0 || renameIndex < 0 || modelIndex > renameIndex {
t.Fatalf("/model should sort before /rename for /m query: %#v", matches)
}
}
func TestInputHistoryNavigation(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.rememberInput("first")
m.rememberInput("second")
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlP})
m = next.(model)
if got := m.input.Value(); got != "second" {
t.Fatalf("after ctrl+p input = %q", got)
}
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlP})
m = next.(model)
if got := m.input.Value(); got != "first" {
t.Fatalf("after second ctrl+p input = %q", got)
}
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlN})
m = next.(model)
if got := m.input.Value(); got != "second" {
t.Fatalf("after ctrl+n input = %q", got)
}
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlN})
m = next.(model)
if got := m.input.Value(); got != "" {
t.Fatalf("after second ctrl+n input = %q", got)
}
}
func TestToolMessageRendersStructuredCard(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`})
plain := stripANSI(m.renderMessages())
for _, want := range []string{iconTool + " file_read", "README.md"} {
if !strings.Contains(plain, want) {
t.Fatalf("tool card missing %q:\n%s", want, plain)
}
}
}
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) {
@@ -197,9 +472,11 @@ func TestModelThinkingCommand(t *testing.T) {
}
type fakeModelManager struct {
provider string
model string
thinking string
provider string
model string
thinking string
sessionName string
saved bool
}
func (f *fakeModelManager) CurrentProvider() string {
@@ -217,6 +494,17 @@ func (f *fakeModelManager) CurrentThinking() string {
return f.thinking
}
func (f *fakeModelManager) CurrentSessionID() string {
return "test-id"
}
func (f *fakeModelManager) CurrentSessionName() string {
if f.sessionName == "" {
return "test-session"
}
return f.sessionName
}
func (f *fakeModelManager) Info() string {
return "provider: " + f.CurrentProvider() + "\nmodel: " + f.model + "\nthinking: " + f.thinking
}
@@ -236,6 +524,28 @@ func (f *fakeModelManager) SetThinking(level string) (string, error) {
return f.Info(), nil
}
func (f *fakeModelManager) Save() error {
f.saved = true
return nil
}
func (f *fakeModelManager) Resume(id string) error {
return nil
}
func (f *fakeModelManager) Rename(name string) (string, error) {
f.sessionName = name
return "Renamed.", nil
}
func (f *fakeModelManager) Sessions() (string, error) {
return "No sessions.", nil
}
func (f *fakeModelManager) NewSession() (string, error) {
return "New session.", nil
}
func TestParseThemeMode(t *testing.T) {
for input, want := range map[string]ThemeMode{
"": ThemeLight,
+92 -2
View File
@@ -2,6 +2,8 @@ package tui
import (
"strings"
"unicode"
"unicode/utf8"
"github.com/charmbracelet/glamour"
"github.com/charmbracelet/glamour/ansi"
@@ -24,7 +26,7 @@ func renderMarkdown(content string, width int, th theme) string {
if err != nil {
return content
}
return strings.TrimRight(rendered, "\n")
return trimRenderedMarkdown(rendered)
}
func markdownStyle(th theme) ansi.StyleConfig {
@@ -32,6 +34,7 @@ func markdownStyle(th theme) ansi.StyleConfig {
if th.Mode == ThemeDark {
cfg = glamstyles.DarkStyleConfig
}
clearMarkdownBackgrounds(&cfg)
zero := uint(0)
cfg.Document.Margin = &zero
@@ -51,7 +54,7 @@ func markdownStyle(th theme) ansi.StyleConfig {
cfg.LinkText.Color = stringPtr(th.User)
cfg.Code.Color = stringPtr(th.Tool)
cfg.Code.BackgroundColor = stringPtr(th.SurfaceAlt)
cfg.Code.BackgroundColor = nil
cfg.CodeBlock.Color = stringPtr(th.Text)
cfg.CodeBlock.BackgroundColor = nil
cfg.CodeBlock.Margin = &zero
@@ -59,6 +62,93 @@ func markdownStyle(th theme) ansi.StyleConfig {
return cfg
}
func clearMarkdownBackgrounds(cfg *ansi.StyleConfig) {
cfg.Document.BackgroundColor = nil
cfg.BlockQuote.BackgroundColor = nil
cfg.Paragraph.BackgroundColor = nil
cfg.Heading.BackgroundColor = nil
cfg.H1.BackgroundColor = nil
cfg.H2.BackgroundColor = nil
cfg.H3.BackgroundColor = nil
cfg.H4.BackgroundColor = nil
cfg.H5.BackgroundColor = nil
cfg.H6.BackgroundColor = nil
cfg.Text.BackgroundColor = nil
cfg.Strikethrough.BackgroundColor = nil
cfg.Emph.BackgroundColor = nil
cfg.Strong.BackgroundColor = nil
cfg.HorizontalRule.BackgroundColor = nil
cfg.Item.BackgroundColor = nil
cfg.Enumeration.BackgroundColor = nil
cfg.Task.BackgroundColor = nil
cfg.Link.BackgroundColor = nil
cfg.LinkText.BackgroundColor = nil
cfg.Image.BackgroundColor = nil
cfg.ImageText.BackgroundColor = nil
cfg.Code.BackgroundColor = nil
cfg.CodeBlock.BackgroundColor = nil
cfg.Table.BackgroundColor = nil
cfg.DefinitionList.BackgroundColor = nil
cfg.DefinitionTerm.BackgroundColor = nil
cfg.DefinitionDescription.BackgroundColor = nil
cfg.HTMLBlock.BackgroundColor = nil
cfg.HTMLSpan.BackgroundColor = nil
}
func trimRenderedMarkdown(value string) string {
lines := strings.Split(strings.TrimRight(value, "\n"), "\n")
for i, line := range lines {
lines[i] = trimStyledTrailingSpaces(line)
}
return strings.Join(lines, "\n")
}
func trimStyledTrailingSpaces(line string) string {
cut := 0
seenNonSpace := false
seenWhitespaceAfterLastNonSpace := false
for i := 0; i < len(line); {
if end, ok := ansiSequenceEnd(line, i); ok {
if seenNonSpace && !seenWhitespaceAfterLastNonSpace {
cut = end
}
i = end
continue
}
r, size := utf8.DecodeRuneInString(line[i:])
if r == utf8.RuneError && size == 0 {
break
}
end := i + size
if unicode.IsSpace(r) {
if seenNonSpace {
seenWhitespaceAfterLastNonSpace = true
}
} else {
seenNonSpace = true
seenWhitespaceAfterLastNonSpace = false
cut = end
}
i = end
}
if cut == 0 {
return ""
}
return line[:cut]
}
func ansiSequenceEnd(value string, start int) (int, bool) {
if start+2 >= len(value) || value[start] != '\x1b' || value[start+1] != '[' {
return 0, false
}
for i := start + 2; i < len(value); i++ {
if value[i] >= 0x40 && value[i] <= 0x7e {
return i + 1, true
}
}
return 0, false
}
func stringPtr(value string) *string {
return &value
}
+20
View File
@@ -37,6 +37,26 @@ func TestRenderMarkdownSupportsThemesAndNarrowWidth(t *testing.T) {
}
}
func TestRenderMarkdownDoesNotEmitBackgroundColors(t *testing.T) {
content := "# Title\n\n`inline`\n\n```go\nfmt.Println(\"x\")\n```"
for _, mode := range []ThemeMode{ThemeLight, ThemeDark} {
rendered := renderMarkdown(content, 80, themeForMode(mode))
if backgroundPattern.MatchString(rendered) {
t.Fatalf("markdown output for %s should not emit background colors:\n%q", mode, rendered)
}
}
}
func TestRenderMarkdownTrimsGlamourTrailingSpaces(t *testing.T) {
rendered := renderMarkdown("hello", 80, themeForMode(ThemeLight))
plain := stripANSI(rendered)
if plain != "hello" {
t.Fatalf("markdown should not pad a plain assistant line:\n%q", plain)
}
}
var backgroundPattern = regexp.MustCompile(`\x1b\[[0-9;:]*48[;:][0-9;:]*m`)
var ansiPattern = regexp.MustCompile(`\x1b\[[0-9;:]*[A-Za-z]`)
func stripANSI(value string) string {
+119 -61
View File
@@ -15,6 +15,17 @@ const (
ThemeDark ThemeMode = "dark"
)
const (
iconReady = "✓"
iconWorking = "◆"
iconTool = "▣"
iconError = "✕"
iconCommand = "⌘"
iconSession = "◌"
iconContext = "◔"
iconModel = "◉"
)
func ParseThemeMode(value string) (ThemeMode, error) {
switch ThemeMode(strings.ToLower(strings.TrimSpace(value))) {
case "", ThemeLight:
@@ -45,26 +56,35 @@ type theme struct {
}
type styles struct {
App lipgloss.Style
Header lipgloss.Style
Title lipgloss.Style
Muted lipgloss.Style
Activity lipgloss.Style
ModelMeta lipgloss.Style
Footer lipgloss.Style
Viewport lipgloss.Style
InputBox lipgloss.Style
CommandHint lipgloss.Style
UserLabel lipgloss.Style
AssistantLabel lipgloss.Style
ToolLabel lipgloss.Style
ErrorLabel lipgloss.Style
SystemLabel lipgloss.Style
UserMessage lipgloss.Style
AssistantMsg lipgloss.Style
ToolMessage lipgloss.Style
ErrorMessage lipgloss.Style
SystemMessage lipgloss.Style
App lipgloss.Style
Muted lipgloss.Style
Activity lipgloss.Style
Header lipgloss.Style
ModelMeta lipgloss.Style
Footer lipgloss.Style
Viewport lipgloss.Style
InputBox lipgloss.Style
CommandHint lipgloss.Style
StatusReady lipgloss.Style
StatusWarn lipgloss.Style
StatusError lipgloss.Style
StatusInfo lipgloss.Style
Pill lipgloss.Style
PillAccent lipgloss.Style
Palette lipgloss.Style
PaletteMatch lipgloss.Style
ToolName lipgloss.Style
ToolKey lipgloss.Style
ToolValue lipgloss.Style
UserLabel lipgloss.Style
ToolLabel lipgloss.Style
ErrorLabel lipgloss.Style
SystemLabel lipgloss.Style
UserMessage lipgloss.Style
AssistantMsg lipgloss.Style
ToolMessage lipgloss.Style
ErrorMessage lipgloss.Style
SystemMessage lipgloss.Style
}
func themeForMode(mode ThemeMode) theme {
@@ -113,64 +133,93 @@ func darkTheme() theme {
func (t theme) styles() styles {
return styles{
App: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.Background)),
Header: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.SurfaceAlt)).
Padding(0, 1),
Title: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Title)),
Foreground(lipgloss.Color(t.Text)),
Muted: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)),
Activity: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Background)).
Background(lipgloss.Color(t.Title)).
Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1),
Header: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Title)),
ModelMeta: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.SurfaceAlt)).
Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1),
Footer: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)).
Background(lipgloss.Color(t.Background)).
Padding(0, 1),
Viewport: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.Background)),
Foreground(lipgloss.Color(t.Text)),
InputBox: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.Surface)).
Border(lipgloss.NormalBorder()).
Border(lipgloss.ThickBorder()).
BorderForeground(lipgloss.Color(t.Border)).
Padding(0, 1),
Padding(0, 2),
CommandHint: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1),
StatusReady: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Assistant)),
StatusWarn: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Tool)),
StatusError: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Error)),
StatusInfo: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.User)),
Pill: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.SurfaceAlt)).
Padding(0, 1),
UserLabel: labelStyle(t.User),
AssistantLabel: labelStyle(t.Assistant),
ToolLabel: labelStyle(t.Tool),
ErrorLabel: labelStyle(t.Error),
SystemLabel: labelStyle(t.System),
PillAccent: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Title)).
Background(lipgloss.Color(t.SurfaceAlt)).
Padding(0, 1),
UserMessage: messageStyle(t.Text, t.User),
AssistantMsg: messageStyle(t.Text, t.Assistant),
ToolMessage: messageStyle(t.Text, t.Tool),
ErrorMessage: messageStyle(t.Error, t.Error),
SystemMessage: lipgloss.NewStyle().Foreground(lipgloss.Color(t.System)).PaddingLeft(1),
Palette: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Border(lipgloss.ThickBorder()).
BorderForeground(lipgloss.Color(t.Border)).
Padding(0, 2),
PaletteMatch: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.User)),
ToolName: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Tool)),
ToolKey: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)),
ToolValue: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)),
UserLabel: labelStyle(t.User),
ToolLabel: labelStyle(t.Tool),
ErrorLabel: labelStyle(t.Error),
SystemLabel: labelStyle(t.System),
UserMessage: bubbleMessageStyle(t.Text, t.SurfaceAlt),
AssistantMsg: assistantMessageStyle(t.Text),
ToolMessage: messageStyle(t.Text),
ErrorMessage: messageStyle(t.Error),
SystemMessage: lipgloss.NewStyle().Foreground(lipgloss.Color(t.System)).Padding(0, 1),
}
}
@@ -180,21 +229,30 @@ func labelStyle(color string) lipgloss.Style {
Foreground(lipgloss.Color(color))
}
func messageStyle(textColor, accentColor string) lipgloss.Style {
func messageStyle(textColor string) lipgloss.Style {
return lipgloss.NewStyle().
Foreground(lipgloss.Color(textColor)).
Border(lipgloss.ThickBorder(), false, false, false, true).
BorderForeground(lipgloss.Color(accentColor)).
PaddingLeft(1)
Padding(0, 1)
}
func assistantMessageStyle(textColor string) lipgloss.Style {
return lipgloss.NewStyle().
Foreground(lipgloss.Color(textColor)).
Padding(1, 2)
}
func bubbleMessageStyle(textColor, backgroundColor string) lipgloss.Style {
return lipgloss.NewStyle().
Foreground(lipgloss.Color(textColor)).
Background(lipgloss.Color(backgroundColor)).
Padding(1, 2)
}
func applyTextareaTheme(input *textarea.Model, t theme) {
focused := textarea.Style{
Base: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.Surface)),
CursorLine: lipgloss.NewStyle().
Background(lipgloss.Color(t.Surface)),
Foreground(lipgloss.Color(t.Text)),
CursorLine: lipgloss.NewStyle(),
Placeholder: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)),
Prompt: lipgloss.NewStyle().
@@ -205,7 +263,7 @@ func applyTextareaTheme(input *textarea.Model, t theme) {
Foreground(lipgloss.Color(t.Surface)),
}
blurred := focused
blurred.CursorLine = lipgloss.NewStyle().Background(lipgloss.Color(t.Surface))
blurred.CursorLine = lipgloss.NewStyle()
input.FocusedStyle = focused
input.BlurredStyle = blurred
+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()
}
+228
View File
@@ -0,0 +1,228 @@
package main
import (
"bufio"
"context"
"errors"
"flag"
"fmt"
"net/http"
"os"
"os/signal"
"path/filepath"
"strings"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/internal/tools"
"agentu/internal/tui"
"agentu/pkg/config"
"agentu/pkg/llm"
)
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)
os.Exit(1)
}
}
func run() error {
configPath := flag.String("config", config.DefaultPath, "path to YAML config")
yolo := flag.Bool("yolo", false, "also enable automatic file writes and shell execution")
plain := flag.Bool("plain", false, "use the simple line-based REPL instead of the TUI")
themeName := flag.String("theme", "light", "TUI theme: light or dark")
resumeID := flag.String("resume", "", "resume a specific session by ID")
flag.Parse()
themeMode, err := tui.ParseThemeMode(*themeName)
if err != nil {
return err
}
cfg, err := config.Load(*configPath)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
if bootstrapErr := bootstrapConfig(*configPath); bootstrapErr != nil {
return bootstrapErr
}
return fmt.Errorf("created default config at %s; edit it to set your API key and provider, then restart agentu", *configPath)
}
return err
}
providerName := cfg.DefaultProviderName()
providerConfig := cfg.Providers[providerName]
provider := llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, http.DefaultClient)
var registry *tools.Registry
if *yolo {
registry = tools.Builtins(cfg.Agent.WorkingDir)
} else {
registry = tools.ReadOnlyBuiltins(cfg.Agent.WorkingDir)
}
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,
MaxContextTokens: maxContextTokens,
})
// Initialize session store
home, err := os.UserHomeDir()
if err != nil {
return fmt.Errorf("resolve home dir: %w", err)
}
sessionsDir := filepath.Join(home, ".agentu", "sessions")
store, err := session.NewStore(sessionsDir)
if err != nil {
return fmt.Errorf("init session store: %w", err)
}
modelManager, err := session.NewManager(cfg, assistant, http.DefaultClient, store)
if err != nil {
return err
}
if err := initializeSession(modelManager, *resumeID); err != nil {
return err
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
defer stop()
if !*plain {
err := tui.Run(ctx, assistant, tui.Options{
ModelName: providerConfig.Model,
Yolo: *yolo,
ThemeMode: themeMode,
ModelManager: modelManager,
WorkingDir: cfg.Agent.WorkingDir,
})
printExitSession(modelManager)
return err
}
return repl(ctx, assistant, modelManager)
}
func initializeSession(modelManager *session.Manager, resumeID string) error {
if resumeID != "" {
if err := modelManager.Resume(resumeID); err != nil {
return fmt.Errorf("resume session: %w", err)
}
return nil
}
if _, err := modelManager.NewSession(); err != nil {
return fmt.Errorf("start new session: %w", err)
}
return nil
}
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, /compact to compress context.")
for {
fmt.Print("> ")
if !scanner.Scan() {
if err := scanner.Err(); err != nil {
return err
}
fmt.Println()
printExitSession(modelManager)
return nil
}
input := strings.TrimSpace(scanner.Text())
switch input {
case "":
continue
case "/exit", "/quit":
printExitSession(modelManager)
return nil
case "/clear":
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 {
fmt.Fprintln(os.Stderr, "\nerror:", err)
}
_ = modelManager.Save()
fmt.Println()
}
}
func printExitSession(m *session.Manager) {
if m == nil {
return
}
_ = m.Save()
id := m.CurrentSessionID()
name := m.CurrentSessionName()
if id != "" {
if name != "" && name != id {
fmt.Fprintf(os.Stderr, "session: %s (%s)\n", id, name)
} else {
fmt.Fprintf(os.Stderr, "session: %s\n", id)
}
fmt.Fprintf(os.Stderr, "resume: agentu --resume %s\n", id)
}
}
+101
View File
@@ -0,0 +1,101 @@
package main
import (
"context"
"path/filepath"
"testing"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/pkg/config"
"agentu/pkg/llm"
)
type startupProvider struct{}
func (p *startupProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
return emit(llm.StreamEvent{Content: "ok"})
}
func startupTestManager(t *testing.T) (*session.Manager, *agent.Agent) {
t.Helper()
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"},
Thinking: "none",
ThinkingParam: "thinking",
},
},
}
assistant := agent.New(agent.Options{
Provider: &startupProvider{},
ProviderName: "one",
Model: "model-a",
})
store, err := session.NewStore(filepath.Join(t.TempDir(), "sessions"))
if err != nil {
t.Fatal(err)
}
manager, err := session.NewManager(cfg, assistant, nil, store)
if err != nil {
t.Fatal(err)
}
return manager, assistant
}
func TestInitializeSessionStartsNewByDefault(t *testing.T) {
manager, assistant := startupTestManager(t)
if _, err := manager.NewSession(); err != nil {
t.Fatalf("new session: %v", err)
}
savedID := manager.CurrentSessionID()
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: "old context"}})
if err := manager.Save(); err != nil {
t.Fatalf("save: %v", err)
}
if err := initializeSession(manager, ""); err != nil {
t.Fatalf("initialize: %v", err)
}
if manager.CurrentSessionID() == savedID {
t.Fatalf("default startup resumed latest session %s", savedID)
}
for _, msg := range assistant.Messages() {
if msg.Content == "old context" {
t.Fatal("default startup carried old context into new session")
}
}
}
func TestInitializeSessionResumesExplicitID(t *testing.T) {
manager, assistant := startupTestManager(t)
if _, err := manager.NewSession(); err != nil {
t.Fatalf("new session: %v", err)
}
savedID := manager.CurrentSessionID()
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: "old context"}})
if err := manager.Save(); err != nil {
t.Fatalf("save: %v", err)
}
if err := initializeSession(manager, savedID); err != nil {
t.Fatalf("initialize: %v", err)
}
if manager.CurrentSessionID() != savedID {
t.Fatalf("expected resumed session %s, got %s", savedID, manager.CurrentSessionID())
}
found := false
for _, msg := range assistant.Messages() {
if msg.Content == "old context" {
found = true
}
}
if !found {
t.Fatal("explicit resume did not restore saved context")
}
}
@@ -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"}
}
@@ -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)
}
+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,
@@ -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)
+43
View File
@@ -0,0 +1,43 @@
package llm
import "encoding/json"
const tokenEstimateCharRatio = 4
// EstimateTokens estimates the token cost of messages and tool definitions in a
// chat completion request. It intentionally uses a simple JSON-size heuristic:
// roughly four UTF-8 bytes per token plus a small per-message framing overhead.
func EstimateTokens(messages []Message, tools []Tool) int {
if len(messages) == 0 && len(tools) == 0 {
return 0
}
payload := struct {
Messages []Message `json:"messages,omitempty"`
Tools []Tool `json:"tools,omitempty"`
}{
Messages: messages,
Tools: tools,
}
data, err := json.Marshal(payload)
if err != nil {
return 0
}
return estimateBytesAsTokens(len(data)) + len(messages)*4
}
func estimateMessageTokens(msg Message) int {
data, err := json.Marshal(msg)
if err != nil {
return 0
}
return estimateBytesAsTokens(len(data)) + 4
}
func estimateBytesAsTokens(n int) int {
if n <= 0 {
return 0
}
return (n + tokenEstimateCharRatio - 1) / tokenEstimateCharRatio
}
+45
View File
@@ -0,0 +1,45 @@
package llm
import "testing"
func TestEstimateTokensEmpty(t *testing.T) {
if got := EstimateTokens(nil, nil); got != 0 {
t.Fatalf("EstimateTokens(nil, nil) = %d, want 0", got)
}
}
func TestEstimateTokens(t *testing.T) {
shortMessages := []Message{{Role: RoleUser, Content: "hello"}}
short := EstimateTokens(shortMessages, nil)
if short <= 4 {
t.Fatalf("short estimate = %d, want message framing plus content", short)
}
if short > 50 {
t.Fatalf("short estimate = %d, looks too high for one short message", short)
}
longMessages := []Message{
{Role: RoleUser, Content: "hello"},
{Role: RoleAssistant, Content: "This is a longer response with enough content to exceed the short estimate."},
}
long := EstimateTokens(longMessages, nil)
if long <= short {
t.Fatalf("long estimate = %d, want > short estimate %d", long, short)
}
withTools := EstimateTokens(shortMessages, []Tool{{
Type: "function",
Function: ToolFunction{
Name: "file_read",
Description: "Read a file",
Parameters: []byte(`{"type":"object"}`),
},
}})
if withTools <= short {
t.Fatalf("withTools estimate = %d, want > short estimate %d", withTools, short)
}
if msgTokens := estimateMessageTokens(shortMessages[0]); msgTokens <= 4 {
t.Fatalf("estimateMessageTokens = %d, want message framing plus content", msgTokens)
}
}
@@ -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 {