13 Commits
Author SHA1 Message Date
loveuer 065b30a84c feat: complete v0.3.1 agent workflow 2026-09-15 03:39:22 -07:00
loveuer 47671f3098 feat(tools): stream shell_run output live into the TUI 2026-08-13 13:56:24 +08:00
loveuer 28195519d1 feat(session,llm): restore session runtime on resume; align Responses API details 2026-08-13 13:46:26 +08:00
loveuer 425f6ee5ee feat(llm): support OpenAI Responses API via per-provider api_type 2026-08-13 11:46:23 +08:00
loveuer c8d8329dbb feat(tui): codex-style input, history, bang commands, single-line status bar
- Arrow-key input history (up/down with multiline awareness)
- Bang shell commands via ! prefix (requires --yolo)
- Codex-style input bar with mode indicator (R/Y/⌘)
- Command palette virtual scrolling
- Multiline input viewport fix
- Merged model pills + footer into single status line
- Removed shift+enter (wait for TUI library support)
- Removed default footer hints
2026-06-25 20:40:50 -07:00
loveuer dabf5bfecc compact: retain newest ~20% of context; refactor startup into internal/app
- Replace fixed 6-message compact retention with ratio-based logic:
  retain ~20% of compactable estimated tokens, with a 6-message floor
  and user-turn boundary alignment
- Add compactRetainedTokenBudget, compactRecentStart, compactTurnBoundary
  helpers in internal/agent/agent.go
- Add TestAgentCompactKeepsTwentyPercentRecentContext
- Move startup/session/repl/bootstrap logic from root into internal/app
  so main.go is a thin binary entry point (15 lines)
- Update /compact docs in README.md
2026-06-24 23:59:21 -07:00
loveuer 9abd5c3b5f feat(tui): implement dark mode with runtime /theme switching
- Paint dark background at app shell level via appStyle() helper
- Add /theme slash command for runtime light/dark switching
- Add comprehensive dark-mode style tests
- Document --theme dark in README Quick Start
2026-06-24 20:45:06 -07:00
loveuer 3d403bd685 feat: Phase 1 code-agent capabilities — file_edit, git tools, parallel execution, tool output preview
New tools:
- file_edit: surgical line-range or pattern edits (yolo only)
- git_status: show changed files (read-only)
- git_diff: show unstaged/staged diff (read-only)
- git_log: show recent commits (read-only)

Agent improvements:
- parallel execution for independent read-only tool calls
- mutating tools (file_write, file_edit, shell_run) run sequentially
- tool output preview in logs with [output] marker

TUI improvements:
- tool cards now show truncated output preview (up to 8 lines)
- mergeToolLog updates placeholder cards with results
- toolStatusText keeps activity bar single-line

Documentation:
- README updated with git tools and file_edit descriptions
2026-06-24 19:47:14 -07:00
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
loveuer 8dbcb1147f v0.0.3: add built-in web tools
Add no-key web_search and web_fetch tools.
Fetch pages with net/http and parse HTML content.
Document built-in web capabilities.
2026-06-23 09:59:24 +08:00
60 changed files with 13261 additions and 861 deletions
+2
View File
@@ -2,3 +2,5 @@ agentu*.yaml
/agentu
/coverage.out
*.swp
chat-history.txt
dist/
+129 -33
View File
@@ -1,17 +1,24 @@
# agentu
agentu is a small local agent written in Go. The current MVP is a terminal chat
agentu is a small local coding agent written in Go. It provides a terminal chat
UI with OpenAI-compatible providers, streaming responses, session-level
model/provider switching, and local tools.
model/provider switching, skills, approvals, and local tools.
## Current MVP
## v0.3.1
- TUI-first chat experience with a light theme by default.
- OpenAI-compatible `/v1/chat/completions` provider support.
- OpenAI Responses API (`/v1/responses`) provider support.
- Multiple providers from `~/.agentu/config.yaml`.
- Session-only provider, model, and thinking-level switching.
- Read-only project tools enabled by default.
- `--yolo` mode for file writes and shell execution.
- Interactive approval for model-requested file writes and shell commands.
- Atomic file writes, per-file rollback snapshots, and multi-file patches.
- Project outlines with Go, TypeScript/JavaScript, and Rust symbols.
- Global Codex-style skills loaded from `~/.agentu/skills`.
- Persistent session index and in-memory task tracking.
- Built-in web search and URL fetching.
- `--yolo` mode for automatic file writes and shell execution.
- Temporary `/btw` side questions that can run while the main task continues.
## Quick Start
@@ -19,13 +26,21 @@ 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:
Use dark mode when your terminal is dark:
```sh
go run ./cmd/agentu --yolo
go run . --theme dark
```
By default, the TUI asks for approval before every model-requested file change
or shell command. Use `--yolo` only when those operations should run without
confirmation:
```sh
go run . --yolo
```
## Config
@@ -41,81 +56,142 @@ Use `--config <path>` to load another file.
```yaml
providers:
loveuer:
api_type: chat
base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY}
model: deepseek-v4-flash
models:
- deepseek-v4-flash
- id: 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
context: 256000
```
Web tools are enabled by default with a built-in no-key backend. agentu exposes
`web_search` for current or external web information and `web_fetch` for reading
a known URL.
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.
`api_type` selects which API protocol the provider speaks:
`thinking` is optional. Supported values are:
```text
chat OpenAI-compatible /v1/chat/completions (default)
responses OpenAI Responses API /v1/responses
```
The Responses API client converts the conversation history to Responses input
items (`user`/`assistant` messages, `function_call`, and
`function_call_output`), sends the system prompt as the top-level
`instructions` field, streams `response.output_text.delta` events, and supports
function calling through `response.output_item.added`,
`response.function_call_arguments.delta`, and `response.output_item.done`.
`models` is required. `/model model <name>` is restricted to that list. Each
entry is an object with:
- `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 request field. The default field depends on `api_type`: `chat` uses
`thinking`, while `responses` uses the official `reasoning` field, sent as
`{"effort": ...}` with agentu's `middle` level mapped to the official `medium`
value (`none`, `high`, `xhigh`, `max` pass through). Override with
`thinking_param` if a provider expects a different field name.
## 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.
- `--yolo` automatically approves model-requested file writes and shell execution.
- `--resume <id>` resumes a saved session from `~/.agentu/sessions`.
## TUI Controls
- `enter` sends the message.
- `ctrl+j` or `alt+enter` inserts a newline.
- `ctrl+j` inserts a newline.
- `up` recalls the previous input; `ctrl+p` / `ctrl+n` navigate input history.
- The input starts at one line and grows up to six lines.
- `page up` / `page down` scrolls the transcript.
- `esc` cancels a running response.
- While a task is running, `/btw <question>` asks a concurrent side question.
- Entering only `.` means “continue with the current task.”
- At a tool approval prompt, `y` approves and `n` or `esc` denies the operation.
- Type `/` to show available commands.
- In `--yolo` mode, type `! <command>` to run a shell command from the configured working directory.
Slash commands:
- `/btw <question>` answers from a snapshot of the current context without adding the question or answer to conversation history. It never exposes tools and can run alongside the main task.
- `/clear` resets conversation history.
- `/compact` summarizes older conversation history and keeps the newest ~20% of context, with a small recent-message floor.
- `/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.
- `/skills` lists global skills.
- `/skills install <path-or-url>` installs a skill and activates it in the current conversation.
- `/rollback [id]` restores a file snapshot; without an ID it restores the newest snapshot.
- `/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.
In the TUI, type `/model ` or `/model model ` to choose from configured model IDs in the candidate palette.
- `/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
Read-only tools are available by default:
These tools are always available:
- `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
- `project_outline` for a gitignore-aware tree with Go, TypeScript/JavaScript, and Rust symbols
- `todo` for conversation-local task tracking
- `git_status` shows changed files
- `git_diff` shows unstaged or staged diff output
- `git_log` shows recent commit history
`--yolo` additionally enables:
Web tools are also available:
- `file_write`
- `shell_run`
- `web_search`
- `web_fetch`
In the TUI, these tools require approval by default and run automatically with
`--yolo`:
- `file_write` for creating or replacing whole files
- `file_edit` for surgical line-range or pattern edits to existing files
- `apply_patch` for atomic-style multi-file V4 patches, including dry runs
- `shell_run` for non-interactive shell commands; output streams live into the
TUI while the command runs instead of appearing only after completion
Before a file mutation, agentu stores an in-memory snapshot and includes its
`rb_...` ID in the tool result. Use `/rollback` or `/rollback <id>` to restore
it. Snapshots last for the lifetime of the process. The plain REPL remains
read-only unless `--yolo` is supplied because it has no interactive approval UI.
Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
example:
@@ -124,6 +200,21 @@ example:
请帮我运行 `go test ./...` 并总结结果
```
## Global Skills
At startup, agentu loads both Codex-style directories containing `SKILL.md`
and legacy flat Markdown files from `~/.agentu/skills`. Skill instructions are
added to the system prompt. Install a local directory, Markdown file, or direct
HTTP(S) Markdown URL from the TUI:
```text
/skills install ./my-skill
/skills install https://example.com/SKILL.md
```
Directory installs preserve supporting files such as `references/`; symlinks
inside a source directory are rejected.
## Development
Run the verification suite:
@@ -131,7 +222,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.
@@ -140,5 +231,10 @@ Local config and build outputs are intentionally ignored by git.
- Provider/model/thinking changes are session-only; agentu does not write them
back to `~/.agentu/config.yaml`.
- `api_type` is per provider and cannot be changed at runtime; use
`/model provider <name>` to switch to a provider configured with a different
API type.
- The config schema is intentionally strict during MVP development. Deprecated
fields such as `provider` or `active_provider` are rejected.
- Session metadata is indexed in `~/.agentu/sessions/index.json`; complete
conversations remain stored as one JSON file per session.
+5 -11
View File
@@ -1,17 +1,11 @@
providers:
loveuer:
# API protocol: "chat" (OpenAI-compatible /v1/chat/completions, default)
# or "responses" (OpenAI Responses API /v1/responses).
api_type: chat
base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY}
model: deepseek-v4-flash
models:
- deepseek-v4-flash
- id: 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
context: 256000
-129
View File
@@ -1,129 +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)
}
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()
}
}
+1 -1
View File
@@ -7,6 +7,7 @@ require (
github.com/charmbracelet/bubbletea v1.3.10
github.com/charmbracelet/glamour v1.0.0
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834
golang.org/x/net v0.38.0
gopkg.in/yaml.v3 v3.0.1
)
@@ -39,7 +40,6 @@ require (
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
github.com/yuin/goldmark v1.7.13 // indirect
github.com/yuin/goldmark-emoji v1.0.6 // indirect
golang.org/x/net v0.38.0 // indirect
golang.org/x/sys v0.38.0 // indirect
golang.org/x/term v0.36.0 // indirect
golang.org/x/text v0.30.0 // indirect
+659 -35
View File
@@ -7,15 +7,42 @@ import (
"io"
"sort"
"strings"
"sync"
"time"
"agentu/internal/llm"
"agentu/internal/tools"
"agentu/pkg/llm"
)
const defaultMaxToolIterations = 8
const (
defaultMaxToolRounds = 50
doomLoopThreshold = 3
compactMinRecentMessages = 6
compactRecentRetentionPercent = 20
summaryMessageLimit = 8000
contextSummaryPrefix = "[context summary]"
toolLogArgumentLimit = 400
toolLogOutputPreviewBytes = 4096
toolLogOutputMarker = "[output]"
toolStreamLogLimit = tools.MaxToolOutputBytes
continuationPrompt = "Continue with the current task."
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 toolExecutionResult struct {
call llm.ToolCall
output string
}
type Agent struct {
mu sync.RWMutex
provider llm.Provider
providerName string
model string
@@ -24,8 +51,12 @@ type Agent struct {
thinkingEnabled bool
systemPrompt string
toolRegistry *tools.Registry
authorizer ToolAuthorizer
rollbackStore *tools.RollbackStore
toolTimeout time.Duration
maxToolIterations int
maxToolRounds int
maxContextTokens int
lastUsage *llm.Usage
messages []llm.Message
}
@@ -38,14 +69,17 @@ type Options struct {
ThinkingEnabled bool
SystemPrompt string
ToolRegistry *tools.Registry
Authorizer ToolAuthorizer
RollbackStore *tools.RollbackStore
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 {
@@ -60,8 +94,11 @@ func New(opts Options) *Agent {
thinkingEnabled: opts.ThinkingEnabled,
systemPrompt: opts.SystemPrompt,
toolRegistry: opts.ToolRegistry,
authorizer: opts.Authorizer,
rollbackStore: opts.RollbackStore,
toolTimeout: toolTimeout,
maxToolIterations: maxToolIterations,
maxToolRounds: maxToolRounds,
maxContextTokens: opts.MaxContextTokens,
}
a.Clear()
return a
@@ -74,6 +111,7 @@ type RuntimeOptions struct {
Thinking string
ThinkingParam string
ThinkingEnabled bool
MaxContextTokens int
}
type RuntimeInfo struct {
@@ -81,9 +119,13 @@ type RuntimeInfo struct {
Model string
Thinking string
ThinkingEnabled bool
ContextTokens int
MaxContextTokens int
}
func (a *Agent) SetRuntime(opts RuntimeOptions) {
a.mu.Lock()
defer a.mu.Unlock()
if opts.Provider != nil {
a.provider = opts.Provider
}
@@ -96,18 +138,30 @@ 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 {
a.mu.RLock()
defer a.mu.RUnlock()
return RuntimeInfo{
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.mu.Lock()
defer a.mu.Unlock()
a.clear()
}
func (a *Agent) clear() {
a.lastUsage = nil
a.messages = nil
if strings.TrimSpace(a.systemPrompt) != "" {
a.messages = append(a.messages, llm.Message{
@@ -117,27 +171,364 @@ func (a *Agent) Clear() {
}
}
// AddSystemInstructions makes newly installed instructions effective for the
// current conversation and for future Clear calls.
func (a *Agent) AddSystemInstructions(instructions string) {
instructions = strings.TrimSpace(instructions)
if instructions == "" {
return
}
a.mu.Lock()
defer a.mu.Unlock()
a.systemPrompt = strings.TrimSpace(a.systemPrompt + "\n\n" + instructions)
if len(a.messages) > 0 && a.messages[0].Role == llm.RoleSystem {
a.messages[0].Content = a.systemPrompt
return
}
a.messages = append([]llm.Message{{Role: llm.RoleSystem, Content: a.systemPrompt}}, a.messages...)
}
// Messages returns a copy of the current conversation messages.
func (a *Agent) Messages() []llm.Message {
a.mu.RLock()
defer a.mu.RUnlock()
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.mu.Lock()
defer a.mu.Unlock()
a.messages = append([]llm.Message(nil), messages...)
if strings.TrimSpace(a.systemPrompt) == "" {
return
}
if len(a.messages) > 0 && a.messages[0].Role == llm.RoleSystem {
a.messages[0].Content = a.systemPrompt
return
}
a.messages = append([]llm.Message{{Role: llm.RoleSystem, Content: a.systemPrompt}}, a.messages...)
}
// LastUsage returns a copy of the most recent provider token usage data.
func (a *Agent) LastUsage() *llm.Usage {
a.mu.RLock()
defer a.mu.RUnlock()
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) {
a.mu.Lock()
defer a.mu.Unlock()
a.setLastUsage(usage)
}
func (a *Agent) setLastUsage(usage *llm.Usage) {
if usage == nil {
a.lastUsage = nil
return
}
copy := *usage
a.lastUsage = &copy
}
// compactRetainedTokenBudget returns the number of estimated tokens to retain
// based on compactRecentRetentionPercent, using integer ceiling division.
func compactRetainedTokenBudget(totalTokens int) int {
if totalTokens <= 0 {
return 1
}
return (totalTokens*compactRecentRetentionPercent + 99) / 100
}
// compactTurnBoundary walks backward from start until it finds a user message
// or reaches prefixEnd, ensuring the retained suffix starts at a user turn.
func compactTurnBoundary(messages []llm.Message, prefixEnd int, start int) int {
for start > prefixEnd && messages[start].Role != llm.RoleUser {
start--
}
return start
}
// compactRecentStart returns the index where the retained recent suffix begins.
// It balances the 20% token-budget ratio with the compactMinRecentMessages floor
// and aligns to a user-turn boundary.
func compactRecentStart(messages []llm.Message, prefixEnd int) int {
compactableCount := len(messages) - prefixEnd
if compactableCount <= compactMinRecentMessages {
return len(messages)
}
compactable := messages[prefixEnd:]
totalTokens := llm.EstimateTokens(compactable, nil)
retainTokens := compactRetainedTokenBudget(totalTokens)
// Walk backward from the end to find the earliest suffix fitting the budget.
ratioStart := len(messages) - 1
for i := len(messages) - 1; i >= prefixEnd; i-- {
if llm.EstimateTokens(messages[i:], nil) <= retainTokens {
ratioStart = i
} else {
break
}
}
// Enforce the minimum recent-message floor.
minStart := len(messages) - compactMinRecentMessages
if minStart < prefixEnd {
minStart = prefixEnd
}
// Choose the earlier start so the floor is never weakened.
start := minStart
if ratioStart < start {
start = ratioStart
}
return compactTurnBoundary(messages, prefixEnd, start)
}
// Compact summarizes older conversation messages while preserving the system
// prompt and recent messages intact.
func (a *Agent) Compact(ctx context.Context) error {
a.mu.RLock()
messages := append([]llm.Message(nil), a.messages...)
provider := a.provider
model := a.model
a.mu.RUnlock()
if len(messages) < 4 {
return nil
}
if provider == nil {
return fmt.Errorf("provider is not configured")
}
prefixEnd := 0
if messages[0].Role == llm.RoleSystem {
prefixEnd = 1
}
recentStart := compactRecentStart(messages, prefixEnd)
if recentStart <= prefixEnd || recentStart >= len(messages) {
return nil
}
middle := append([]llm.Message(nil), messages[prefixEnd:recentStart]...)
if len(middle) == 0 {
return nil
}
summary, usage, err := summarizeMessages(ctx, provider, model, 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(messages[recentStart:]))
compressed = append(compressed, messages[:prefixEnd]...)
compressed = append(compressed, llm.Message{
Role: llm.RoleAssistant,
Content: contextSummaryPrefix + "\n" + summary,
})
compressed = append(compressed, messages[recentStart:]...)
a.mu.Lock()
defer a.mu.Unlock()
a.messages = compressed
if usage != nil {
a.lastUsage = usage
}
return nil
}
func (a *Agent) shouldCompact() bool {
a.mu.RLock()
defer a.mu.RUnlock()
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 summarizeMessages(ctx context.Context, provider llm.Provider, model string, messages []llm.Message) (string, *llm.Usage, error) {
req := llm.ChatRequest{
Model: 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 := 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()
}
// BTW answers a temporary side question using a snapshot of the current
// conversation. It intentionally exposes no tools and does not mutate the
// conversation, usage, or compaction state.
func (a *Agent) BTW(ctx context.Context, input string, out io.Writer) error {
input = strings.TrimSpace(input)
if input == "" {
return fmt.Errorf("btw question is required")
}
a.mu.RLock()
provider := a.provider
model := a.model
messages := append([]llm.Message(nil), a.messages...)
thinking := a.thinking
thinkingParam := a.thinkingParam
thinkingEnabled := a.thinkingEnabled
a.mu.RUnlock()
if provider == nil {
return fmt.Errorf("provider is not configured")
}
// A snapshot taken while tools are running may end with unanswered tool
// calls, which providers reject. Keep only the latest complete boundary.
messages = completeConversationPrefix(messages)
messages = append(messages, llm.Message{Role: llm.RoleUser, Content: input})
messages = expandContinuationMessages(messages)
req := llm.ChatRequest{Model: model, Messages: messages}
if thinkingEnabled {
if thinkingParam == "" {
thinkingParam = "thinking"
}
req.Extra = map[string]any{thinkingParam: thinking}
}
return provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
if event.Content == "" || out == nil {
return nil
}
_, err := io.WriteString(out, event.Content)
return err
})
}
func completeConversationPrefix(messages []llm.Message) []llm.Message {
safeEnd := 0
pending := make(map[string]struct{})
for i, msg := range messages {
if len(pending) > 0 && msg.Role != llm.RoleTool {
break
}
switch {
case msg.Role == llm.RoleAssistant && len(msg.ToolCalls) > 0:
for _, call := range msg.ToolCalls {
pending[call.ID] = struct{}{}
}
case msg.Role == llm.RoleTool:
delete(pending, msg.ToolCallID)
}
if len(pending) == 0 {
safeEnd = i + 1
}
}
return append([]llm.Message(nil), messages[:safeEnd]...)
}
func expandContinuationMessages(messages []llm.Message) []llm.Message {
expanded := append([]llm.Message(nil), messages...)
for i := range expanded {
if expanded[i].Role == llm.RoleUser && strings.TrimSpace(expanded[i].Content) == "." {
expanded[i].Content = continuationPrompt
}
}
return expanded
}
func (a *Agent) appendMessage(message llm.Message) {
a.mu.Lock()
a.messages = append(a.messages, message)
a.mu.Unlock()
}
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})
a.appendMessage(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
}
if len(calls) == 0 {
a.messages = append(a.messages, llm.Message{
a.appendMessage(llm.Message{
Role: llm.RoleAssistant,
Content: content,
})
return nil
}
a.messages = append(a.messages, llm.Message{
a.appendMessage(llm.Message{
Role: llm.RoleAssistant,
Content: content,
ToolCalls: calls,
@@ -146,40 +537,88 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
fmt.Fprintln(out)
}
for _, call := range calls {
result := a.executeTool(ctx, call, logs)
a.messages = append(a.messages, llm.Message{
Role: llm.RoleTool,
ToolCallID: call.ID,
Content: result,
// 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.appendMessage(llm.Message{
Role: llm.RoleAssistant,
Content: msg,
})
return fmt.Errorf("%s", msg)
}
}
return fmt.Errorf("tool iteration limit reached: %d", a.maxToolIterations)
results := a.executeTools(ctx, calls, logs)
for _, result := range results {
a.appendMessage(llm.Message{
Role: llm.RoleTool,
ToolCallID: result.call.ID,
Content: result.output,
})
}
}
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) {
var content strings.Builder
calls := make(map[int]*toolCallBuilder)
a.mu.RLock()
provider := a.provider
model := a.model
messages := append([]llm.Message(nil), a.messages...)
thinking := a.thinking
thinkingParam := a.thinkingParam
thinkingEnabled := a.thinkingEnabled
a.mu.RUnlock()
if provider == nil {
return "", nil, fmt.Errorf("provider is not configured")
}
req := llm.ChatRequest{
Model: a.model,
Messages: append([]llm.Message(nil), a.messages...),
Model: model,
Messages: expandContinuationMessages(messages),
Tools: a.toolDefinitions(),
}
if a.thinkingEnabled {
thinkingParam := a.thinkingParam
if thinkingEnabled {
if thinkingParam == "" {
thinkingParam = "thinking"
}
req.Extra = map[string]any{thinkingParam: a.thinking}
req.Extra = map[string]any{thinkingParam: thinking}
}
if len(req.Tools) > 0 {
req.ToolChoice = "auto"
}
err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
err := 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 {
@@ -212,39 +651,224 @@ func (a *Agent) toolDefinitions() []llm.Tool {
return a.toolRegistry.Definitions()
}
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) string {
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) (string, bool) {
name := call.Function.Name
if logs != nil {
fmt.Fprintf(logs, "\n[tool] %s %s\n", name, compact(call.Function.Arguments, 400))
}
if a.toolRegistry == nil {
return tools.FormatError(fmt.Errorf("tool registry is disabled"))
return tools.FormatError(fmt.Errorf("tool registry is disabled")), false
}
tool, ok := a.toolRegistry.Get(name)
if !ok {
return tools.FormatError(fmt.Errorf("unknown tool: %s", name))
return tools.FormatError(fmt.Errorf("unknown tool: %s", name)), false
}
args := strings.TrimSpace(call.Function.Arguments)
if args == "" {
args = "{}"
}
raw := json.RawMessage(args)
rollbackNotes, err := a.prepareMutation(ctx, tool, call, raw, args)
if err != nil {
return tools.FormatError(err), false
}
toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
defer cancel()
output, err := tool.Execute(toolCtx, json.RawMessage(args))
if streaming, ok := tool.(tools.StreamingTool); ok {
output, streamed := a.executeStreamingTool(streaming, toolCtx, raw, name, args, logs)
return appendRollbackNotes(output, rollbackNotes), streamed
}
output, err := tool.Execute(toolCtx, raw)
if err != nil {
if output != "" {
return output + "\n" + tools.FormatError(err)
return appendRollbackNotes(output+"\n"+tools.FormatError(err), rollbackNotes), false
}
return tools.FormatError(err)
return appendRollbackNotes(tools.FormatError(err), rollbackNotes), false
}
if logs != nil {
fmt.Fprintf(logs, "[tool] %s done\n", name)
return appendRollbackNotes(output, rollbackNotes), false
}
func (a *Agent) prepareMutation(ctx context.Context, tool tools.Tool, call llm.ToolCall, raw json.RawMessage, args string) ([]string, error) {
if !tools.IsMutatingCall(tool, raw) {
return nil, nil
}
var targets []string
if targetTool, ok := tool.(tools.TargetPathTool); ok {
var err error
targets, err = targetTool.TargetPaths(raw)
if err != nil {
return nil, fmt.Errorf("resolve mutation targets: %w", err)
}
}
preview := compact(args, toolLogArgumentLimit)
if previewTool, ok := tool.(tools.PreviewableTool); ok {
value, err := previewTool.Preview(ctx, raw)
if err != nil {
return nil, fmt.Errorf("preview mutation: %w", err)
}
preview = value
}
if a.authorizer != nil {
decision, err := a.authorizer.AuthorizeToolCall(ctx, ToolAuthorizationRequest{
ID: call.ID,
ToolName: call.Function.Name,
Arguments: args,
Preview: preview,
RollbackAvailable: a.rollbackStore != nil && len(targets) > 0,
})
if err != nil {
return nil, fmt.Errorf("authorize %s: %w", call.Function.Name, err)
}
if decision != ToolAuthorizationAllow {
return nil, fmt.Errorf("tool call denied by user: %s", call.Function.Name)
}
}
if a.rollbackStore == nil {
return nil, nil
}
notes := make([]string, 0, len(targets))
for _, path := range targets {
id, err := a.rollbackStore.Capture(call.Function.Name, path, path)
if err != nil {
return nil, fmt.Errorf("capture rollback snapshot for %s: %w", path, err)
}
notes = append(notes, fmt.Sprintf("[rollback] %s %s", id, path))
}
return notes, nil
}
func appendRollbackNotes(output string, notes []string) string {
if len(notes) == 0 {
return output
}
if strings.TrimSpace(output) == "" {
return strings.Join(notes, "\n")
}
return output + "\n" + strings.Join(notes, "\n")
}
func (a *Agent) executeStreamingTool(streaming tools.StreamingTool, ctx context.Context, raw json.RawMessage, name, args string, logs io.Writer) (string, bool) {
var streamed int
emit := func(chunk string) error {
if chunk == "" || logs == nil {
return nil
}
if streamed >= toolStreamLogLimit {
return nil
}
if len(chunk) > toolStreamLogLimit-streamed {
chunk = chunk[:toolStreamLogLimit-streamed]
streamed = toolStreamLogLimit
_, err := fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s\n[streaming output truncated after %d bytes]", name, compact(args, toolLogArgumentLimit), toolLogOutputMarker, chunk, toolStreamLogLimit)
return err
}
streamed += len(chunk)
_, err := fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s", name, compact(args, toolLogArgumentLimit), toolLogOutputMarker, chunk)
return err
}
output, err := streaming.ExecuteStream(ctx, raw, emit)
if logs != nil {
if err != nil {
_ = emit("\n" + tools.FormatError(err))
} else if streamed == 0 {
_ = emit("(no output)")
}
}
if err != nil {
if output != "" {
return output + "\n" + tools.FormatError(err), true
}
return tools.FormatError(err), true
}
return output, true
}
func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
if a.toolRegistry == nil {
return false
}
for _, call := range calls {
if tool, ok := a.toolRegistry.Get(call.Function.Name); ok {
args := strings.TrimSpace(call.Function.Arguments)
if args == "" {
args = "{}"
}
if tools.IsMutatingCall(tool, json.RawMessage(args)) {
return true
}
}
}
return false
}
func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.Writer) []toolExecutionResult {
results := make([]toolExecutionResult, len(calls))
// Log all tool starts sequentially so TUI gets ordered placeholders.
if logs != nil {
for _, call := range calls {
logToolStart(logs, call)
}
}
if a.hasMutatingCall(calls) {
// Sequential: avoid racing mutating tools.
for i, call := range calls {
output, streamed := a.executeTool(ctx, call, logs)
results[i] = toolExecutionResult{call: call, output: output}
if logs != nil && !streamed {
logToolOutput(logs, call, output)
}
}
return results
}
// Concurrent: independent read-only tools.
var mu sync.Mutex
var wg sync.WaitGroup
for i, call := range calls {
wg.Add(1)
go func(idx int, c llm.ToolCall) {
defer wg.Done()
output, streamed := a.executeTool(ctx, c, logs)
results[idx] = toolExecutionResult{call: c, output: output}
if logs != nil && !streamed {
mu.Lock()
logToolOutput(logs, c, output)
mu.Unlock()
}
}(i, call)
}
wg.Wait()
return results
}
func logToolStart(logs io.Writer, call llm.ToolCall) {
fmt.Fprintf(logs, "\n[tool] %s %s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit))
}
func logToolOutput(logs io.Writer, call llm.ToolCall, output string) {
fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit), toolLogOutputMarker, toolOutputPreview(output))
}
func toolOutputPreview(output string) string {
output = strings.TrimSpace(output)
if output == "" {
return "(no output)"
}
if len(output) <= toolLogOutputPreviewBytes {
return output
}
return output[:toolLogOutputPreviewBytes] + fmt.Sprintf("\n[truncated: showing first %d bytes; omitted %d bytes]", toolLogOutputPreviewBytes, len(output)-toolLogOutputPreviewBytes)
}
type toolCallBuilder struct {
index int
+786 -1
View File
@@ -3,11 +3,17 @@ package agent
import (
"context"
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"reflect"
"strings"
"sync"
"testing"
"agentu/internal/llm"
"agentu/internal/tools"
"agentu/pkg/llm"
)
type fakeProvider struct {
@@ -86,6 +92,86 @@ func TestAgentRunsToolLoop(t *testing.T) {
}
}
type streamToolProvider struct {
calls int
t *testing.T
}
func (p *streamToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.calls++
switch p.calls {
case 1:
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
{Index: 0, ID: "call_1", Type: "function", Name: "test_stream", Arguments: `{}`},
}})
case 2:
last := req.Messages[len(req.Messages)-1]
if last.Role != llm.RoleTool || last.Content != "full result" {
p.t.Fatalf("last message = %#v", last)
}
return emit(llm.StreamEvent{Content: "done"})
default:
p.t.Fatalf("unexpected call count %d", p.calls)
return nil
}
}
type streamEchoTool struct{}
func (streamEchoTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "test_stream",
Description: "streaming test",
Parameters: json.RawMessage(`{"type":"object"}`),
},
}
}
func (streamEchoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
return "full result", nil
}
func (streamEchoTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
if err := emit("one\n"); err != nil {
return "", err
}
if err := emit("two\n"); err != nil {
return "", err
}
return "full result", nil
}
func TestAgentStreamsToolOutputToLogs(t *testing.T) {
provider := &streamToolProvider{t: t}
a := New(Options{
Provider: provider,
Model: "test",
SystemPrompt: "system",
ToolRegistry: tools.NewRegistry(streamEchoTool{}),
})
var out strings.Builder
var logs strings.Builder
if err := a.RunTurn(context.Background(), "stream", &out, &logs); err != nil {
t.Fatal(err)
}
if out.String() != "done" {
t.Fatalf("out = %q", out.String())
}
logText := logs.String()
if strings.Count(logText, "[output]") != 2 {
t.Fatalf("logs = %q", logText)
}
if !strings.Contains(logText, "one\n") || !strings.Contains(logText, "two\n") {
t.Fatalf("logs missing streamed chunks: %q", logText)
}
if strings.Contains(logText, "[output]\nfull result") {
t.Fatalf("logs contain duplicate full output: %q", logText)
}
}
type contentProvider struct {
called bool
text string
@@ -96,6 +182,346 @@ func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, e
return emit(llm.StreamEvent{Content: p.text})
}
type btwProvider struct {
req llm.ChatRequest
err error
}
func (p *btwProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.req = req
if p.err != nil {
return p.err
}
return emit(llm.StreamEvent{Content: "side answer", Usage: &llm.Usage{TotalTokens: 999}})
}
func TestBTWUsesContextWithoutMutatingConversation(t *testing.T) {
provider := &btwProvider{}
a := New(Options{Provider: provider, Model: "test", Thinking: "high", ThinkingEnabled: true})
a.SetMessages([]llm.Message{
{Role: llm.RoleUser, Content: "main question"},
{Role: llm.RoleAssistant, Content: "main answer"},
})
a.SetLastUsage(&llm.Usage{TotalTokens: 7})
before := a.Messages()
var out strings.Builder
if err := a.BTW(context.Background(), "what does that mean?", &out); err != nil {
t.Fatal(err)
}
if out.String() != "side answer" {
t.Fatalf("output = %q", out.String())
}
if !reflect.DeepEqual(a.Messages(), before) {
t.Fatalf("conversation changed:\nbefore: %+v\nafter: %+v", before, a.Messages())
}
if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 7 {
t.Fatalf("usage changed: %+v", usage)
}
if len(provider.req.Tools) != 0 || provider.req.ToolChoice != "" {
t.Fatalf("btw exposed tools: %+v / %q", provider.req.Tools, provider.req.ToolChoice)
}
last := provider.req.Messages[len(provider.req.Messages)-1]
if last.Role != llm.RoleUser || last.Content != "what does that mean?" {
t.Fatalf("last request message = %+v", last)
}
if provider.req.Extra["thinking"] != "high" {
t.Fatalf("thinking extra = %+v", provider.req.Extra)
}
}
func TestBTWRejectsEmptyQuestionAndReturnsProviderError(t *testing.T) {
provider := &btwProvider{err: errors.New("boom")}
a := New(Options{Provider: provider, Model: "test"})
if err := a.BTW(context.Background(), " ", nil); err == nil {
t.Fatal("empty question did not fail")
}
if err := a.BTW(context.Background(), "question", nil); err == nil || !strings.Contains(err.Error(), "boom") {
t.Fatalf("provider error = %v", err)
}
}
func TestDotContinuesCurrentTaskWithoutRewritingHistory(t *testing.T) {
provider := &btwProvider{}
a := New(Options{Provider: provider, Model: "test"})
if err := a.RunTurn(context.Background(), " . ", nil, nil); err != nil {
t.Fatal(err)
}
requestMessage := provider.req.Messages[len(provider.req.Messages)-1]
if requestMessage.Content != continuationPrompt {
t.Fatalf("request content = %q, want %q", requestMessage.Content, continuationPrompt)
}
history := a.Messages()
if history[0].Content != "." {
t.Fatalf("stored input = %q, want literal dot", history[0].Content)
}
}
func TestExpandContinuationMessagesOnlyExpandsSingleDot(t *testing.T) {
messages := []llm.Message{
{Role: llm.RoleUser, Content: "."},
{Role: llm.RoleUser, Content: " . "},
{Role: llm.RoleUser, Content: ".."},
{Role: llm.RoleUser, Content: ". continue"},
{Role: llm.RoleAssistant, Content: "."},
}
expanded := expandContinuationMessages(messages)
if expanded[0].Content != continuationPrompt || expanded[1].Content != continuationPrompt {
t.Fatalf("single dots were not expanded: %+v", expanded)
}
for _, index := range []int{2, 3, 4} {
if expanded[index].Content != messages[index].Content {
t.Fatalf("message %d changed from %q to %q", index, messages[index].Content, expanded[index].Content)
}
}
if messages[0].Content != "." || messages[1].Content != " . " {
t.Fatalf("source messages were mutated: %+v", messages)
}
}
func TestCompleteConversationPrefixDropsPendingToolRound(t *testing.T) {
messages := []llm.Message{
{Role: llm.RoleUser, Content: "first"},
{Role: llm.RoleAssistant, Content: "answer"},
{Role: llm.RoleUser, Content: "main"},
{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "call_1"}, {ID: "call_2"}}},
{Role: llm.RoleTool, ToolCallID: "call_1", Content: "one"},
}
got := completeConversationPrefix(messages)
if len(got) != 3 || got[len(got)-1].Content != "main" {
t.Fatalf("complete prefix = %+v", got)
}
}
type concurrentBTWProvider struct {
mainStarted chan struct{}
releaseMain chan struct{}
once sync.Once
mu sync.Mutex
btwRequest llm.ChatRequest
}
func (p *concurrentBTWProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
last := req.Messages[len(req.Messages)-1].Content
if last == "main task" {
p.once.Do(func() { close(p.mainStarted) })
select {
case <-p.releaseMain:
return emit(llm.StreamEvent{Content: "main answer"})
case <-ctx.Done():
return ctx.Err()
}
}
p.mu.Lock()
p.btwRequest = req
p.mu.Unlock()
return emit(llm.StreamEvent{Content: "side answer"})
}
func TestBTWCanRunDuringMainTurn(t *testing.T) {
provider := &concurrentBTWProvider{mainStarted: make(chan struct{}), releaseMain: make(chan struct{})}
a := New(Options{Provider: provider, Model: "test"})
mainDone := make(chan error, 1)
go func() { mainDone <- a.RunTurn(context.Background(), "main task", nil, nil) }()
<-provider.mainStarted
var side strings.Builder
if err := a.BTW(context.Background(), "side question", &side); err != nil {
t.Fatal(err)
}
if side.String() != "side answer" {
t.Fatalf("side output = %q", side.String())
}
provider.mu.Lock()
req := provider.btwRequest
provider.mu.Unlock()
if len(req.Messages) < 2 || req.Messages[len(req.Messages)-2].Content != "main task" {
t.Fatalf("btw request did not include active main turn: %+v", req.Messages)
}
close(provider.releaseMain)
if err := <-mainDone; err != nil {
t.Fatal(err)
}
if len(a.Messages()) != 2 {
t.Fatalf("main conversation contains btw messages: %+v", a.Messages())
}
}
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 messages[2].Content != "recent-user-0" {
t.Fatalf("first retained recent message = %#v", messages[2])
}
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 manyTurnConversation(turns int) []llm.Message {
messages := []llm.Message{{Role: llm.RoleSystem, Content: "system"}}
for i := 0; i < turns; i++ {
messages = append(messages,
llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("user-%02d", i)},
llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("assistant-%02d", i)},
)
}
return messages
}
func TestAgentCompactKeepsTwentyPercentRecentContext(t *testing.T) {
provider := &compactProvider{}
a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
a.SetMessages(manyTurnConversation(25))
if err := a.Compact(context.Background()); err != nil {
t.Fatal(err)
}
messages := a.Messages()
// 1 system + 1 summary + 10 retained recent messages (5 user/assistant turns)
if len(messages) != 12 {
t.Fatalf("messages = %d, want 12; contents: %v", len(messages), joinMessageContents(messages))
}
if provider.summaryCalls != 1 {
t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
}
if messages[2].Content != "user-20" {
t.Fatalf("first retained message = %#v, want user-20", messages[2])
}
if messages[len(messages)-1].Content != "assistant-24" {
t.Fatalf("last retained message = %#v, want assistant-24", messages[len(messages)-1])
}
if !strings.Contains(provider.lastSummary, "user-19") {
t.Fatalf("summary should contain user-19 (part of summarized 80%%): %q", provider.lastSummary)
}
if strings.Contains(provider.lastSummary, "user-20") {
t.Fatalf("summary should not contain user-20 (part of retained 20%%): %q", provider.lastSummary)
}
}
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 +563,362 @@ 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)
}
}
// slowTool blocks on a channel before returning.
type slowTool struct {
name string
started chan<- string
release <-chan struct{}
}
func (t *slowTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: t.name,
Description: "slow test tool",
Parameters: json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}}}`),
},
}
}
func (t *slowTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args struct {
Text string `json:"text"`
}
if err := json.Unmarshal(raw, &args); err != nil {
return "", err
}
t.started <- t.name
<-t.release
return args.Text, nil
}
// mutatingSlowTool is like slowTool but implements MutatingTool.
type mutatingSlowTool struct {
slowTool
}
func (*mutatingSlowTool) Mutates() bool { return true }
func TestAgentExecutesIndependentToolCallsInParallel(t *testing.T) {
started := make(chan string, 2)
release := make(chan struct{})
tool1 := &slowTool{name: "slow_one", started: started, release: release}
tool2 := &slowTool{name: "slow_two", started: started, release: release}
provider := &parallelProvider{
responses: [][]llm.ToolCallDelta{
{
{Index: 0, ID: "call_1", Type: "function", Name: "slow_one", Arguments: `{"text":"one"}`},
{Index: 1, ID: "call_2", Type: "function", Name: "slow_two", Arguments: `{"text":"two"}`},
},
},
finalContent: "done",
}
a := New(Options{
Provider: provider,
Model: "test",
SystemPrompt: "system",
ToolRegistry: tools.NewRegistry(tool1, tool2),
})
var out strings.Builder
done := make(chan error, 1)
go func() {
done <- a.RunTurn(context.Background(), "run both", &out, nil)
}()
// Both tools must start before we release them.
names := make(map[string]bool)
for i := 0; i < 2; i++ {
select {
case name := <-started:
names[name] = true
case <-make(chan struct{}):
// Use a timer in real code, but for test simplicity we rely on deadlock detection.
t.Fatal("timed out waiting for tool to start")
}
}
if !names["slow_one"] || !names["slow_two"] {
t.Fatalf("expected both tools to start, got %v", names)
}
close(release)
if err := <-done; err != nil {
t.Fatal(err)
}
if out.String() != "done" {
t.Fatalf("out = %q", out.String())
}
// Verify tool results were in order.
msgs := a.Messages()
var toolMsgs []string
for _, m := range msgs {
if m.Role == llm.RoleTool {
toolMsgs = append(toolMsgs, m.Content)
}
}
if len(toolMsgs) != 2 || toolMsgs[0] != "one" || toolMsgs[1] != "two" {
t.Fatalf("tool messages = %v, want [one two]", toolMsgs)
}
}
func TestAgentRunsMutatingToolCallsSequentially(t *testing.T) {
started := make(chan string, 2)
release := make(chan struct{})
tool1 := &mutatingSlowTool{slowTool{name: "mut_one", started: started, release: release}}
tool2 := &mutatingSlowTool{slowTool{name: "mut_two", started: started, release: release}}
provider := &parallelProvider{
responses: [][]llm.ToolCallDelta{
{
{Index: 0, ID: "call_1", Type: "function", Name: "mut_one", Arguments: `{"text":"one"}`},
{Index: 1, ID: "call_2", Type: "function", Name: "mut_two", Arguments: `{"text":"two"}`},
},
},
finalContent: "done",
}
a := New(Options{
Provider: provider,
Model: "test",
SystemPrompt: "system",
ToolRegistry: tools.NewRegistry(tool1, tool2),
})
var out strings.Builder
done := make(chan error, 1)
go func() {
done <- a.RunTurn(context.Background(), "run both", &out, nil)
}()
// Only the first tool should start.
select {
case name := <-started:
if name != "mut_one" {
t.Fatalf("expected mut_one to start first, got %q", name)
}
case <-make(chan struct{}):
t.Fatal("timed out waiting for first tool to start")
}
// Second tool should NOT have started yet.
select {
case name := <-started:
t.Fatalf("second tool started too early: %q", name)
default:
// Expected: second tool hasn't started.
}
close(release)
if err := <-done; err != nil {
t.Fatal(err)
}
}
// parallelProvider emits a batch of tool calls on the first request, then returns finalContent.
type parallelProvider struct {
responses [][]llm.ToolCallDelta
finalContent string
calls int
}
type authorizationProvider struct {
call llm.ToolCallDelta
calls int
toolOutput string
}
func (p *authorizationProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.calls++
if p.calls == 1 {
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{p.call}})
}
p.toolOutput = req.Messages[len(req.Messages)-1].Content
return emit(llm.StreamEvent{Content: "done"})
}
type recordingAuthorizer struct {
decision ToolAuthorizationDecision
requests []ToolAuthorizationRequest
}
func (a *recordingAuthorizer) AuthorizeToolCall(ctx context.Context, req ToolAuthorizationRequest) (ToolAuthorizationDecision, error) {
a.requests = append(a.requests, req)
return a.decision, ctx.Err()
}
func TestAgentAuthorizesMutationAndCreatesRollbackSnapshot(t *testing.T) {
dir := t.TempDir()
provider := &authorizationProvider{call: llm.ToolCallDelta{
Index: 0, ID: "write_1", Type: "function", Name: "file_write",
Arguments: `{"path":"new.txt","content":"hello"}`,
}}
authorizer := &recordingAuthorizer{decision: ToolAuthorizationAllow}
rollback := tools.NewRollbackStore()
a := New(Options{
Provider: provider,
Model: "test",
ToolRegistry: tools.Builtins(dir),
Authorizer: authorizer,
RollbackStore: rollback,
})
if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
t.Fatal(err)
}
if len(authorizer.requests) != 1 || !authorizer.requests[0].RollbackAvailable {
t.Fatalf("authorization requests = %+v", authorizer.requests)
}
if !strings.Contains(authorizer.requests[0].Preview, "write 5 bytes to new.txt") {
t.Fatalf("preview = %q", authorizer.requests[0].Preview)
}
rollbackID := rollback.LastID()
if rollbackID == "" || !strings.Contains(provider.toolOutput, rollbackID) {
t.Fatalf("tool output %q missing rollback ID %q", provider.toolOutput, rollbackID)
}
if _, err := rollback.Restore(rollbackID); err != nil {
t.Fatal(err)
}
if _, err := os.Stat(filepath.Join(dir, "new.txt")); !os.IsNotExist(err) {
t.Fatalf("new file still exists after rollback: %v", err)
}
}
func TestAgentDeniedMutationDoesNotExecute(t *testing.T) {
dir := t.TempDir()
provider := &authorizationProvider{call: llm.ToolCallDelta{
Index: 0, ID: "write_1", Type: "function", Name: "file_write",
Arguments: `{"path":"denied.txt","content":"no"}`,
}}
authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})
if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
t.Fatal(err)
}
if _, err := os.Stat(filepath.Join(dir, "denied.txt")); !os.IsNotExist(err) {
t.Fatalf("denied file exists: %v", err)
}
if !strings.Contains(provider.toolOutput, "denied by user") {
t.Fatalf("tool output = %q", provider.toolOutput)
}
}
func TestAgentDoesNotAuthorizeApplyPatchDryRun(t *testing.T) {
dir := t.TempDir()
provider := &authorizationProvider{call: llm.ToolCallDelta{
Index: 0, ID: "patch_1", Type: "function", Name: "apply_patch",
Arguments: `{"patch":"*** Begin Patch\n*** Add File: dry.txt\n+dry\n*** End Patch","dry_run":true}`,
}}
authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})
if err := a.RunTurn(context.Background(), "preview", nil, nil); err != nil {
t.Fatal(err)
}
if len(authorizer.requests) != 0 {
t.Fatalf("dry run requested authorization: %+v", authorizer.requests)
}
if _, err := os.Stat(filepath.Join(dir, "dry.txt")); !os.IsNotExist(err) {
t.Fatalf("dry run created a file: %v", err)
}
}
func TestSetMessagesRefreshesCurrentSystemInstructions(t *testing.T) {
a := New(Options{SystemPrompt: "current instructions"})
a.SetMessages([]llm.Message{
{Role: llm.RoleSystem, Content: "old instructions"},
{Role: llm.RoleUser, Content: "hello"},
})
messages := a.Messages()
if messages[0].Content != "current instructions" || messages[1].Content != "hello" {
t.Fatalf("messages = %+v", messages)
}
}
func (p *parallelProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.calls++
if p.calls <= len(p.responses) {
return emit(llm.StreamEvent{ToolCalls: p.responses[p.calls-1]})
}
return emit(llm.StreamEvent{Content: p.finalContent})
}
+26
View File
@@ -0,0 +1,26 @@
package agent
import "context"
// ToolAuthorizationDecision is the result of an authorization check.
type ToolAuthorizationDecision string
const (
ToolAuthorizationAllow ToolAuthorizationDecision = "allow"
ToolAuthorizationDeny ToolAuthorizationDecision = "deny"
)
// ToolAuthorizationRequest is sent to an authorizer before executing a
// mutating tool call.
type ToolAuthorizationRequest struct {
ID string
ToolName string
Arguments string
Preview string
RollbackAvailable bool
}
// ToolAuthorizer decides whether a mutating tool call should proceed.
type ToolAuthorizer interface {
AuthorizeToolCall(ctx context.Context, req ToolAuthorizationRequest) (ToolAuthorizationDecision, error)
}
+192
View File
@@ -0,0 +1,192 @@
package app
import (
"context"
"errors"
"flag"
"fmt"
"net/http"
"os"
"os/signal"
"path/filepath"
"strings"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/internal/skills"
"agentu/internal/tools"
"agentu/internal/tui"
"agentu/pkg/config"
"agentu/pkg/llm"
)
type cliOptions struct {
configPath string
yolo bool
plain bool
themeName string
resumeID string
}
func Run() error {
opts := parseCLIOptions()
themeMode, err := tui.ParseThemeMode(opts.themeName)
if err != nil {
return err
}
cfg, err := loadConfig(opts.configPath)
if err != nil {
return err
}
globalSkills, err := skills.LoadGlobal()
if err != nil {
return fmt.Errorf("load global skills: %w", err)
}
rollbackStore := tools.NewRollbackStore()
var approvalBroker *tui.ApprovalBroker
var authorizer agent.ToolAuthorizer
if !opts.yolo && !opts.plain {
approvalBroker = tui.NewApprovalBroker()
authorizer = approvalBroker
}
providerName, providerConfig, provider := defaultProvider(cfg, http.DefaultClient)
assistant := newAssistant(cfg, providerName, providerConfig, provider, opts, globalSkills, authorizer, rollbackStore, http.DefaultClient)
modelManager, err := newSessionManager(cfg, assistant, http.DefaultClient)
if err != nil {
return err
}
if err := initializeSession(modelManager, opts.resumeID); err != nil {
return err
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
defer stop()
return runInterface(ctx, assistant, modelManager, providerConfig, cfg.Agent.WorkingDir, opts, themeMode, approvalBroker, rollbackStore)
}
func parseCLIOptions() cliOptions {
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()
return cliOptions{
configPath: *configPath,
yolo: *yolo,
plain: *plain,
themeName: *themeName,
resumeID: *resumeID,
}
}
func loadConfig(configPath string) (*config.Config, error) {
cfg, err := config.Load(configPath)
if err == nil {
return cfg, nil
}
if !errors.Is(err, os.ErrNotExist) {
return nil, err
}
if bootstrapErr := bootstrapConfig(configPath); bootstrapErr != nil {
return nil, bootstrapErr
}
return nil, fmt.Errorf("created default config at %s; edit it to set your API key and provider, then restart agentu", configPath)
}
func defaultProvider(cfg *config.Config, httpClient *http.Client) (string, config.ProviderConfig, llm.Provider) {
providerName := cfg.DefaultProviderName()
providerConfig := cfg.Providers[providerName]
provider := newProviderClient(providerConfig, httpClient)
return providerName, providerConfig, provider
}
func newProviderClient(providerConfig config.ProviderConfig, httpClient *http.Client) llm.Provider {
if providerConfig.APIType == config.APITypeResponses {
return llm.NewOpenAIResponsesClient(providerConfig.BaseURL, providerConfig.APIKey, httpClient)
}
return llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, httpClient)
}
func newAssistant(cfg *config.Config, providerName string, providerConfig config.ProviderConfig, provider llm.Provider, opts cliOptions, globalSkills []skills.Skill, authorizer agent.ToolAuthorizer, rollbackStore *tools.RollbackStore, httpClient *http.Client) *agent.Agent {
writesEnabled := opts.yolo || !opts.plain
return agent.New(agent.Options{
Provider: provider,
ProviderName: providerName,
Model: providerConfig.Model,
Thinking: providerConfig.Thinking,
ThinkingParam: providerConfig.ThinkingParam,
ThinkingEnabled: providerConfig.ThinkingConfigured,
SystemPrompt: systemPrompt(cfg, writesEnabled, opts.yolo, globalSkills),
ToolRegistry: newToolRegistry(cfg.Agent.WorkingDir, writesEnabled, httpClient),
Authorizer: authorizer,
RollbackStore: rollbackStore,
ToolTimeout: cfg.Agent.ToolTimeout,
MaxContextTokens: maxContextTokens(cfg, providerConfig),
})
}
func systemPrompt(cfg *config.Config, writesEnabled bool, yolo bool, globalSkills []skills.Skill) string {
parts := []string{cfg.Agent.SystemPrompt, tools.AgentInstructionsForMode(cfg.Agent.WorkingDir, writesEnabled, yolo)}
if formatted := skills.FormatForSystemPrompt(globalSkills); formatted != "" {
parts = append(parts, formatted)
}
return strings.TrimSpace(strings.Join(parts, "\n\n"))
}
func maxContextTokens(cfg *config.Config, providerConfig config.ProviderConfig) int {
if cfg.Agent.MaxContextTokens != 0 {
return cfg.Agent.MaxContextTokens
}
return providerConfig.ContextTokens
}
func newToolRegistry(workingDir string, writesEnabled bool, httpClient *http.Client) *tools.Registry {
var registry *tools.Registry
if writesEnabled {
registry = tools.Builtins(workingDir)
} else {
registry = tools.ReadOnlyBuiltins(workingDir)
}
registry.Register(tools.NewBuiltinSearchTool(httpClient))
registry.Register(tools.NewBuiltinFetchTool(httpClient))
return registry
}
func newSessionManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Client) (*session.Manager, error) {
home, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("resolve home dir: %w", err)
}
sessionsDir := filepath.Join(home, ".agentu", "sessions")
store, err := session.NewStore(sessionsDir)
if err != nil {
return nil, fmt.Errorf("init session store: %w", err)
}
return session.NewManager(cfg, assistant, httpClient, store)
}
func runInterface(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager, providerConfig config.ProviderConfig, workingDir string, opts cliOptions, themeMode tui.ThemeMode, approvalBroker *tui.ApprovalBroker, rollbackStore *tools.RollbackStore) error {
if opts.plain {
return repl(ctx, assistant, modelManager)
}
err := tui.Run(ctx, assistant, tui.Options{
ModelName: providerConfig.Model,
Yolo: opts.yolo,
ThemeMode: themeMode,
ModelManager: modelManager,
WorkingDir: workingDir,
Approval: approvalBroker,
Rollback: rollbackStore,
})
printExitSession(modelManager)
return err
}
+38
View File
@@ -0,0 +1,38 @@
package app
import (
"fmt"
"os"
"path/filepath"
"agentu/pkg/config"
)
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 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
}
+113
View File
@@ -0,0 +1,113 @@
package app
import (
"context"
"path/filepath"
"strings"
"testing"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/internal/skills"
"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")
}
}
func TestSystemPromptIncludesGlobalSkillsAndApprovalMode(t *testing.T) {
cfg := &config.Config{Agent: config.AgentConfig{SystemPrompt: "base", WorkingDir: "/tmp/project"}}
prompt := systemPrompt(cfg, true, false, []skills.Skill{{Name: "review", Path: "/tmp/review/SKILL.md", Content: "Review carefully."}})
for _, want := range []string{"base", "interactive user approval", "### review", "Review carefully."} {
if !strings.Contains(prompt, want) {
t.Fatalf("prompt missing %q:\n%s", want, prompt)
}
}
}
+95
View File
@@ -0,0 +1,95 @@
package app
import (
"bufio"
"context"
"fmt"
"os"
"strings"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/internal/skills"
)
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, /btw <question> for a side question.")
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())
fields := strings.Fields(input)
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 len(fields) > 0 && fields[0] == "/btw" {
question := strings.TrimSpace(strings.TrimPrefix(input, "/btw"))
if question == "" {
fmt.Fprintln(os.Stderr, "usage: /btw <question>")
continue
}
fmt.Print("[btw] ")
if err := assistant.BTW(ctx, question, os.Stdout); err != nil {
fmt.Fprintln(os.Stderr, "\nbtw error:", err)
} else {
fmt.Println()
}
continue
}
if input == "/skills" || input == "/skills list" {
items, err := skills.LoadGlobal()
if err != nil {
fmt.Fprintln(os.Stderr, "skills error:", err)
} else {
fmt.Println(skills.FormatList(items))
}
continue
}
if len(fields) >= 3 && fields[0] == "/skills" && fields[1] == "install" {
result, err := skills.InstallGlobal(ctx, strings.Join(fields[2:], " "))
if err != nil {
fmt.Fprintln(os.Stderr, "skills error:", err)
} else {
assistant.AddSystemInstructions(skills.FormatForSystemPrompt([]skills.Skill{result.Skill}))
fmt.Printf("installed skill %q; active in this conversation\n", result.Skill.Name)
}
continue
}
if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
fmt.Fprintln(os.Stderr, "\nerror:", err)
}
_ = modelManager.Save()
fmt.Println()
}
}
+38
View File
@@ -0,0 +1,38 @@
package app
import (
"fmt"
"os"
"agentu/internal/session"
)
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 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)
}
}
-257
View File
@@ -1,257 +0,0 @@
package config
import (
"os"
"path/filepath"
"strings"
"testing"
"time"
)
func TestLoadExpandsEnvAndDefaults(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
model: test-model
api_key: ${AGENTU_TEST_KEY}
agent:
working_dir: .
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
provider := cfg.Providers[cfg.DefaultProviderName()]
if provider.APIKey != "sk-test" {
t.Fatalf("api key = %q", provider.APIKey)
}
if providerName := cfg.DefaultProviderName(); providerName != "test" {
t.Fatalf("default provider = %q", providerName)
}
if cfg.Agent.ToolTimeout != 30*time.Second {
t.Fatalf("tool timeout = %s", cfg.Agent.ToolTimeout)
}
if !filepath.IsAbs(cfg.Agent.WorkingDir) {
t.Fatalf("working dir is not absolute: %s", cfg.Agent.WorkingDir)
}
}
func TestLoadUsesDefaultConfigPath(t *testing.T) {
dir := t.TempDir()
t.Setenv("HOME", dir)
t.Setenv("AGENTU_TEST_KEY", "sk-test")
configDir := filepath.Join(dir, ".agentu")
if err := os.MkdirAll(configDir, 0o755); err != nil {
t.Fatal(err)
}
path := filepath.Join(configDir, "config.yaml")
if err := os.WriteFile(path, []byte(`
providers:
test:
base_url: https://example.com
model: test-model
api_key: ${AGENTU_TEST_KEY}
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load("")
if err != nil {
t.Fatal(err)
}
if got := cfg.DefaultProviderName(); got != "test" {
t.Fatalf("default provider = %q", got)
}
}
func TestResolvePathExpandsHome(t *testing.T) {
dir := t.TempDir()
t.Setenv("HOME", dir)
got, err := ResolvePath("~/.agentu/config.yaml")
if err != nil {
t.Fatal(err)
}
want := filepath.Join(dir, ".agentu", "config.yaml")
if got != want {
t.Fatalf("path = %q, want %q", got, want)
}
}
func TestLoadDefaultSystemPromptDiscouragesEagerTools(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
model: test-model
api_key: ${AGENTU_TEST_KEY}
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
for _, want := range []string{"Answer simple conversation directly", "Use tools only when"} {
if !strings.Contains(cfg.Agent.SystemPrompt, want) {
t.Fatalf("system prompt missing %q: %s", want, cfg.Agent.SystemPrompt)
}
}
}
func TestLoadMultiProviderConfig(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:
loveuer:
base_url: https://example.com
api_key: ${AGENTU_TEST_KEY}
model: model-a
models:
- model-a
- model-b
thinking: xhigh
thinking_param: thinking
backup:
base_url: https://backup.example.com
api_key: ${AGENTU_TEST_KEY}
models:
- model-c
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.DefaultProviderName(); got != "backup" {
t.Fatalf("default provider = %q", got)
}
provider := cfg.Providers["loveuer"]
if got := provider.Model; got != "model-a" {
t.Fatalf("active model = %q", got)
}
if got := cfg.Providers["backup"].Model; got != "model-c" {
t.Fatalf("backup default model = %q", got)
}
if !provider.ThinkingConfigured || provider.Thinking != "xhigh" || provider.ThinkingParam != "thinking" {
t.Fatalf("thinking config = %#v", provider)
}
}
func TestLoadDefaultProviderNameIsDeterministic(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:
two:
base_url: https://two.example.com
api_key: ${AGENTU_TEST_KEY}
model: model-b
one:
base_url: https://one.example.com
api_key: ${AGENTU_TEST_KEY}
model: model-a
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.DefaultProviderName(); got != "one" {
t.Fatalf("default provider = %q", got)
}
}
func TestLoadRejectsInvalidThinking(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
model: test-model
api_key: ${AGENTU_TEST_KEY}
thinking: huge
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "thinking") {
t.Fatalf("err = %v", err)
}
}
func TestLoadRequiresProviderFields(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
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil {
t.Fatal("expected error")
}
}
func TestLoadRejectsLegacyProviderConfig(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
provider:
base_url: https://example.com
model: test-model
api_key: sk-test
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "field provider not found") {
t.Fatalf("err = %v", err)
}
}
func TestLoadRejectsActiveProviderConfig(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
active_provider: loveuer
providers:
loveuer:
base_url: https://example.com
model: test-model
api_key: sk-test
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "field active_provider not found") {
t.Fatalf("err = %v", err)
}
}
+56
View File
@@ -0,0 +1,56 @@
package safeio
import (
"fmt"
"io/fs"
"os"
"path/filepath"
)
// AtomicWriteFile writes data to path atomically by writing to a temp file in
// the same directory and renaming it into place. On any error before rename,
// the temp file is removed. The parent directory is created if needed.
func AtomicWriteFile(path string, data []byte, perm fs.FileMode) error {
dir := filepath.Dir(path)
if err := os.MkdirAll(dir, 0o755); err != nil {
return fmt.Errorf("create parent directory: %w", err)
}
tmp, err := os.CreateTemp(dir, ".atomic-*")
if err != nil {
return fmt.Errorf("create temp file: %w", err)
}
tmpPath := tmp.Name()
cleanup := true
defer func() {
if cleanup {
os.Remove(tmpPath)
}
}()
if _, err := tmp.Write(data); err != nil {
tmp.Close()
return fmt.Errorf("write temp file: %w", err)
}
if err := tmp.Sync(); err != nil {
tmp.Close()
return fmt.Errorf("sync temp file: %w", err)
}
if err := tmp.Close(); err != nil {
return fmt.Errorf("close temp file: %w", err)
}
if err := os.Chmod(tmpPath, perm); err != nil {
return fmt.Errorf("chmod temp file: %w", err)
}
if err := os.Rename(tmpPath, path); err != nil {
return fmt.Errorf("rename temp file: %w", err)
}
cleanup = false
// Best-effort sync of the parent directory.
if f, err := os.Open(dir); err == nil {
_ = f.Sync()
f.Close()
}
return nil
}
+58
View File
@@ -0,0 +1,58 @@
package safeio
import (
"os"
"path/filepath"
"testing"
)
func TestAtomicWriteFileWritesContentAndReplacesExisting(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "data.txt")
// Write initial content.
if err := os.WriteFile(path, []byte("old"), 0o644); err != nil {
t.Fatal(err)
}
// Atomic replace.
if err := AtomicWriteFile(path, []byte("new"), 0o644); err != nil {
t.Fatal(err)
}
got, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(got) != "new" {
t.Fatalf("content = %q, want %q", string(got), "new")
}
// No temp files left behind.
entries, err := os.ReadDir(dir)
if err != nil {
t.Fatal(err)
}
for _, e := range entries {
if e.Name() != "data.txt" {
t.Fatalf("unexpected file in directory: %s", e.Name())
}
}
}
func TestAtomicWriteFileCreatesParentDirs(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "sub", "dir", "file.txt")
if err := AtomicWriteFile(path, []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
got, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(got) != "hello" {
t.Fatalf("content = %q, want %q", string(got), "hello")
}
}
+58
View File
@@ -0,0 +1,58 @@
package session
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"sort"
"agentu/internal/safeio"
)
const sessionIndexFile = "index.json"
type sessionIndex struct {
Sessions []SessionMeta `json:"sessions"`
}
func metaFromSession(sess *Session) SessionMeta {
return SessionMeta{
ID: sess.ID,
Name: sess.Name,
CreatedAt: sess.CreatedAt,
UpdatedAt: sess.UpdatedAt,
Provider: sess.Provider,
Model: sess.Model,
MessageCount: len(sess.Messages),
}
}
func (s *Store) saveIndexFromSessions(metas []SessionMeta) error {
sort.Slice(metas, func(i, j int) bool {
return metas[i].UpdatedAt.After(metas[j].UpdatedAt)
})
idx := sessionIndex{Sessions: metas}
data, err := json.MarshalIndent(idx, "", " ")
if err != nil {
return fmt.Errorf("marshal session index: %w", err)
}
path := filepath.Join(s.dir, sessionIndexFile)
return safeio.AtomicWriteFile(path, data, 0o644)
}
func (s *Store) loadIndex() ([]SessionMeta, error) {
path := filepath.Join(s.dir, sessionIndexFile)
data, err := os.ReadFile(path)
if err != nil {
return nil, err
}
var idx sessionIndex
if err := json.Unmarshal(data, &idx); err != nil {
return nil, err
}
sort.Slice(idx.Sessions, func(i, j int) bool {
return idx.Sessions[i].UpdatedAt.After(idx.Sessions[j].UpdatedAt)
})
return idx.Sessions, nil
}
+220 -16
View File
@@ -5,10 +5,11 @@ import (
"net/http"
"sort"
"strings"
"time"
"agentu/internal/agent"
"agentu/internal/config"
"agentu/internal/llm"
"agentu/pkg/config"
"agentu/pkg/llm"
)
type Manager struct {
@@ -17,9 +18,11 @@ type Manager struct {
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,184 @@ 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.restoreRuntime(sess)
m.agent.SetMessages(sess.Messages)
m.agent.SetLastUsage(sess.LastUsage)
return nil
}
// restoreRuntime switches the manager and agent back to the provider, model,
// and thinking level recorded on a resumed session. Values that are no longer
// valid for the current config are skipped and the defaults remain active.
func (m *Manager) restoreRuntime(sess *Session) {
if sess.Provider != "" && sess.Provider != m.providerName {
if _, ok := m.cfg.Providers[sess.Provider]; ok {
_, _ = m.SetProvider(sess.Provider)
}
}
if sess.Model != "" && sess.Model != m.provider.Model {
_, _ = m.SetModel(sess.Model)
}
if sess.Thinking != "" && sess.Thinking != m.provider.Thinking {
_, _ = m.SetThinking(sess.Thinking)
}
}
// 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
m.currentSession.APIType = m.provider.APIType
m.currentSession.Thinking = m.provider.Thinking
// 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
m.currentSession.APIType = m.provider.APIType
m.currentSession.Thinking = m.provider.Thinking
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 {
@@ -52,17 +226,21 @@ func (m *Manager) CurrentThinking() string {
return m.provider.Thinking
}
func (m *Manager) CurrentAPIType() string {
if m.provider.APIType == "" {
return config.APITypeChat
}
return m.provider.APIType
}
func (m *Manager) Info() string {
var b strings.Builder
fmt.Fprintf(&b, "provider: %s\n", m.CurrentProvider())
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
fmt.Fprintf(&b, "api: %s\n", m.CurrentAPIType())
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, "thinking levels: %s\n", strings.Join(config.ThinkingLevels(), ", "))
b.WriteString("usage: /model provider <name> | /model model <name> | /model thinking <level>")
return b.String()
@@ -75,7 +253,7 @@ func (m *Manager) SetProvider(name string) (string, error) {
}
m.providerName = name
m.provider = provider
m.syncAgent(llm.NewOpenAICompatibleClient(provider.BaseURL, provider.APIKey, m.httpClient))
m.syncAgent(providerClient(provider, m.httpClient))
return "Switched provider.\n" + m.summary(), nil
}
@@ -84,10 +262,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
@@ -101,7 +279,7 @@ func (m *Manager) SetThinking(level string) (string, error) {
m.provider.Thinking = level
m.provider.ThinkingConfigured = true
if m.provider.ThinkingParam == "" {
m.provider.ThinkingParam = "thinking"
m.provider.ThinkingParam = config.DefaultThinkingParam(m.provider.APIType)
}
m.cfg.Providers[m.providerName] = m.provider
m.syncAgent(nil)
@@ -109,7 +287,21 @@ func (m *Manager) SetThinking(level string) (string, error) {
}
func (m *Manager) summary() string {
return fmt.Sprintf("provider: %s\nmodel: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentThinking())
return fmt.Sprintf("provider: %s\nmodel: %s\napi: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentAPIType(), m.CurrentThinking())
}
func providerClient(provider config.ProviderConfig, httpClient *http.Client) llm.Provider {
if provider.APIType == config.APITypeResponses {
return llm.NewOpenAIResponsesClient(provider.BaseURL, provider.APIKey, httpClient)
}
return llm.NewOpenAICompatibleClient(provider.BaseURL, provider.APIKey, httpClient)
}
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) {
@@ -120,6 +312,7 @@ func (m *Manager) syncAgent(provider llm.Provider) {
Thinking: m.provider.Thinking,
ThinkingParam: m.provider.ThinkingParam,
ThinkingEnabled: m.provider.ThinkingConfigured,
MaxContextTokens: m.maxContextTokens(),
})
}
@@ -132,11 +325,22 @@ func (m *Manager) providerNames() []string {
return names
}
func (m *Manager) AvailableModels() []string {
return append([]string(nil), m.provider.Models...)
}
func contains(values []string, value string) bool {
for _, item := range values {
if item == value {
for _, v := range values {
if v == value {
return true
}
}
return false
}
func fitName(name string, width int) string {
if len(name) <= width {
return name
}
return name[:width-3] + "..."
}
+446 -6
View File
@@ -2,12 +2,15 @@ package session
import (
"context"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"testing"
"agentu/internal/agent"
"agentu/internal/config"
"agentu/internal/llm"
"agentu/pkg/config"
"agentu/pkg/llm"
)
type recordingProvider struct {
@@ -19,8 +22,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 +37,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 +51,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 +97,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 +110,425 @@ func TestManagerRejectsUnknownModelAndThinking(t *testing.T) {
t.Fatal("expected invalid thinking error")
}
}
func TestManagerSwitchesToResponsesProvider(t *testing.T) {
var hitPath string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hitPath = r.URL.Path
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"ok"}` + "\n\n"))
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed"}}` + "\n\n"))
}))
defer server.Close()
cfg := &config.Config{
Providers: map[string]config.ProviderConfig{
"chat": {
Name: "chat",
APIType: config.APITypeChat,
BaseURL: "https://chat.example.com",
APIKey: "sk-test",
Model: "model-a",
Models: []string{"model-a"},
},
"responses": {
Name: "responses",
APIType: config.APITypeResponses,
BaseURL: server.URL,
APIKey: "sk-test",
Model: "model-b",
Models: []string{"model-b"},
},
},
}
assistant := agent.New(agent.Options{Provider: &recordingProvider{}, Model: "model-a"})
store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
if err != nil {
t.Fatal(err)
}
manager, err := NewManager(cfg, assistant, server.Client(), store)
if err != nil {
t.Fatal(err)
}
if _, err := manager.SetProvider("responses"); err != nil {
t.Fatal(err)
}
if manager.CurrentAPIType() != config.APITypeResponses {
t.Fatalf("api type = %q", manager.CurrentAPIType())
}
var out strings.Builder
if err := assistant.RunTurn(context.Background(), "hi", &out, nil); err != nil {
t.Fatal(err)
}
if hitPath != "/v1/responses" {
t.Fatalf("path = %q, want /v1/responses", hitPath)
}
if out.String() != "ok" {
t.Fatalf("output = %q", out.String())
}
}
func TestManagerExposesModelCommandCandidates(t *testing.T) {
manager, _, _ := testManager(t)
models := manager.AvailableModels()
if len(models) != 2 || models[0] != "model-a" || models[1] != "model-b" {
t.Fatalf("AvailableModels() = %v", models)
}
models[0] = "mutated"
if got := manager.AvailableModels()[0]; got != "model-a" {
t.Fatalf("mutation leaked: %q", got)
}
}
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)
}
}
func TestManagerResumeRestoresRuntime(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"ok"}` + "\n\n"))
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed"}}` + "\n\n"))
}))
defer server.Close()
cfg := &config.Config{
Providers: map[string]config.ProviderConfig{
"chat": {
Name: "chat",
APIType: config.APITypeChat,
BaseURL: "https://chat.example.com",
APIKey: "sk-test",
Model: "model-a",
Models: []string{"model-a"},
Thinking: "none",
ThinkingParam: "thinking",
ThinkingConfigured: true,
},
"resp": {
Name: "resp",
APIType: config.APITypeResponses,
BaseURL: server.URL,
APIKey: "sk-test",
Model: "m1",
Models: []string{"m1", "m2"},
Thinking: "none",
ThinkingParam: "reasoning",
ThinkingConfigured: true,
},
},
}
assistant := agent.New(agent.Options{
Provider: &recordingProvider{},
Model: "model-a",
Thinking: "none",
ThinkingParam: "thinking",
ThinkingEnabled: true,
})
store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
if err != nil {
t.Fatal(err)
}
manager, err := NewManager(cfg, assistant, server.Client(), store)
if err != nil {
t.Fatal(err)
}
// Build a session on the responses provider with a custom model and thinking.
manager.newSession()
if _, err := manager.SetProvider("resp"); err != nil {
t.Fatal(err)
}
if _, err := manager.SetModel("m2"); err != nil {
t.Fatal(err)
}
if _, err := manager.SetThinking("xhigh"); err != nil {
t.Fatal(err)
}
if err := manager.Save(); err != nil {
t.Fatal(err)
}
id := manager.CurrentSessionID()
// Drift away from the saved runtime.
if _, err := manager.SetProvider("chat"); err != nil {
t.Fatal(err)
}
if _, err := manager.SetModel("model-a"); err != nil {
t.Fatal(err)
}
if _, err := manager.SetThinking("none"); err != nil {
t.Fatal(err)
}
if err := manager.Resume(id); err != nil {
t.Fatal(err)
}
if got := manager.CurrentProvider(); got != "resp" {
t.Fatalf("provider = %q, want resp", got)
}
if got := manager.CurrentModel(); got != "m2" {
t.Fatalf("model = %q, want m2", got)
}
if got := manager.CurrentThinking(); got != "xhigh" {
t.Fatalf("thinking = %q, want xhigh", got)
}
if got := manager.CurrentAPIType(); got != config.APITypeResponses {
t.Fatalf("api type = %q, want responses", got)
}
runtime := assistant.RuntimeInfo()
if runtime.ProviderName != "resp" || runtime.Model != "m2" || runtime.Thinking != "xhigh" {
t.Fatalf("agent runtime = %#v", runtime)
}
}
func TestManagerResumeFallsBackWhenRuntimeMissing(t *testing.T) {
manager, assistant, _ := testManager(t)
sess := &Session{
ID: GenerateID(),
Provider: "ghost",
Model: "ghost-model",
Thinking: "xhigh",
Messages: []llm.Message{{Role: llm.RoleUser, Content: "hi"}},
}
if err := manager.store.Save(sess); err != nil {
t.Fatal(err)
}
if err := manager.Resume(sess.ID); err != nil {
t.Fatal(err)
}
if got := manager.CurrentProvider(); got != "one" {
t.Fatalf("provider = %q, want default one", got)
}
if got := manager.CurrentModel(); got != "model-a" {
t.Fatalf("model = %q, want default model-a", got)
}
if got := manager.CurrentThinking(); got != "xhigh" {
t.Fatalf("thinking = %q, want restored xhigh", got)
}
runtime := assistant.RuntimeInfo()
if runtime.ProviderName != "one" || runtime.Model != "model-a" {
t.Fatalf("agent runtime = %#v", runtime)
}
}
+251
View File
@@ -0,0 +1,251 @@
package session
import (
"crypto/rand"
"encoding/hex"
"encoding/json"
"fmt"
"os"
"path/filepath"
"sort"
"strings"
"time"
"agentu/internal/safeio"
"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"`
APIType string `json:"api_type,omitempty"`
Thinking string `json:"thinking,omitempty"`
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 == nil {
return fmt.Errorf("session is required")
}
if err := validateSessionID(sess.ID); err != nil {
return err
}
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 := safeio.AtomicWriteFile(path, data, 0o644); err != nil {
return fmt.Errorf("write session: %w", err)
}
return s.updateIndex(metaFromSession(sess))
}
// Load reads a session by ID.
func (s *Store) Load(id string) (*Session, error) {
if err := validateSessionID(id); err != nil {
return nil, err
}
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) {
if metas, err := s.loadIndex(); err == nil {
return metas, nil
}
metas, err := s.scanSessions()
if err != nil {
return nil, err
}
if err := s.saveIndexFromSessions(metas); err != nil {
return nil, err
}
return metas, nil
}
func (s *Store) scanSessions() ([]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() || entry.Name() == sessionIndexFile || !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 {
if err := validateSessionID(id); err != nil {
return err
}
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)
}
metas, err := s.loadIndex()
if err != nil {
metas, err = s.scanSessions()
if err != nil {
return err
}
}
filtered := metas[:0]
for _, meta := range metas {
if meta.ID != id {
filtered = append(filtered, meta)
}
}
return s.saveIndexFromSessions(filtered)
}
func (s *Store) filePath(id string) string {
return filepath.Join(s.dir, id+".json")
}
func validateSessionID(id string) error {
if id == "" {
return fmt.Errorf("session ID is required")
}
if id == "." || id == ".." {
return fmt.Errorf("invalid session ID %q", id)
}
for _, r := range id {
if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-' || r == '_' || r == '.' {
continue
}
return fmt.Errorf("invalid session ID %q", id)
}
return nil
}
func (s *Store) updateIndex(meta SessionMeta) error {
metas, err := s.loadIndex()
if err != nil {
metas, err = s.scanSessions()
if err != nil {
return err
}
}
found := false
for i := range metas {
if metas[i].ID == meta.ID {
metas[i] = meta
found = true
break
}
}
if !found {
metas = append(metas, meta)
}
return s.saveIndexFromSessions(metas)
}
// 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 ""
}
+214
View File
@@ -0,0 +1,214 @@
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)
if err := store.Save(nil); err == nil {
t.Fatal("expected error for nil session")
}
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")
}
}
func TestStoreMaintainsSessionIndex(t *testing.T) {
store := tempStore(t)
if err := store.Save(&Session{ID: "indexed", Name: "Indexed"}); err != nil {
t.Fatal(err)
}
if _, err := os.Stat(filepath.Join(store.dir, sessionIndexFile)); err != nil {
t.Fatalf("index not created: %v", err)
}
metas, err := store.loadIndex()
if err != nil {
t.Fatal(err)
}
if len(metas) != 1 || metas[0].ID != "indexed" {
t.Fatalf("index metadata = %+v", metas)
}
if err := store.Delete("indexed"); err != nil {
t.Fatal(err)
}
metas, err = store.loadIndex()
if err != nil {
t.Fatal(err)
}
if len(metas) != 0 {
t.Fatalf("index after delete = %+v", metas)
}
}
func TestStoreRejectsUnsafeSessionID(t *testing.T) {
store := tempStore(t)
if _, err := store.Load("../outside"); err == nil {
t.Fatal("expected unsafe ID rejection")
}
if err := store.Save(&Session{ID: "a/b"}); err == nil {
t.Fatal("expected unsafe ID rejection")
}
}
+583
View File
@@ -0,0 +1,583 @@
package skills
import (
"context"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"regexp"
"sort"
"strings"
"time"
"agentu/internal/safeio"
"agentu/pkg/config"
"gopkg.in/yaml.v3"
)
// DefaultDir is the default directory for global skills.
const DefaultDir = "~/.agentu/skills"
// SkillFileName is the required instruction file inside a Codex-style skill directory.
const SkillFileName = "SKILL.md"
// maxInstallDownloadBytes caps HTTP downloads at 1 MiB.
const maxInstallDownloadBytes = 1 << 20
// Skill is a single global Markdown skill.
type Skill struct {
Name string
Description string
Path string
Content string
}
// InstallResult describes a successful skill installation.
type InstallResult struct {
Skill Skill
TargetDir string
}
// validNameRe allows ASCII letters, digits, dots, underscores, and hyphens.
var validNameRe = regexp.MustCompile(`^[A-Za-z0-9._-]+$`)
// validateSkillName rejects empty, path-like, and otherwise unsafe names.
func validateSkillName(name string) error {
if name == "" || name == "." || name == ".." {
return fmt.Errorf("invalid skill name %q: use letters, numbers, dots, underscores, or hyphens", name)
}
if strings.ContainsAny(name, `/\`) {
return fmt.Errorf("invalid skill name %q: use letters, numbers, dots, underscores, or hyphens", name)
}
if strings.ContainsAny(name, " \t\n\r") {
return fmt.Errorf("invalid skill name %q: use letters, numbers, dots, underscores, or hyphens", name)
}
if !validNameRe.MatchString(name) {
return fmt.Errorf("invalid skill name %q: use letters, numbers, dots, underscores, or hyphens", name)
}
return nil
}
// skillFrontmatter is the decoded YAML frontmatter.
type skillFrontmatter struct {
Name string `yaml:"name"`
Description string `yaml:"description"`
}
// parseSkillFile reads raw bytes and returns the skill body, frontmatter name,
// and frontmatter description. hasFrontmatter indicates whether the source
// bytes started with a --- delimiter pair.
func parseSkillFile(path string, data []byte) (body string, fmName string, fmDesc string, hasFrontmatter bool, err error) {
normalized := strings.ReplaceAll(string(data), "\r\n", "\n")
lines := strings.Split(normalized, "\n")
if len(lines) >= 2 && lines[0] == "---" {
for i := 1; i < len(lines); i++ {
if lines[i] == "---" {
yamlBlock := strings.Join(lines[1:i], "\n")
var fm skillFrontmatter
if err := yaml.Unmarshal([]byte(yamlBlock), &fm); err != nil {
return "", "", "", false, fmt.Errorf("parse skill %s: %w", path, err)
}
bodyText := strings.TrimSpace(strings.Join(lines[i+1:], "\n"))
return bodyText, fm.Name, fm.Description, true, nil
}
}
return "", "", "", false, fmt.Errorf("parse skill %s: missing closing ---", path)
}
return strings.TrimSpace(normalized), "", "", false, nil
}
// LoadGlobal loads global skills from DefaultDir.
func LoadGlobal() ([]Skill, error) {
return LoadDir(DefaultDir)
}
// LoadDir loads skills from a directory. Missing directories return nil, nil.
func LoadDir(dir string) ([]Skill, error) {
resolved, err := config.ResolvePath(dir)
if err != nil {
return nil, err
}
resolved, err = filepath.Abs(resolved)
if err != nil {
return nil, err
}
info, err := os.Stat(resolved)
if err != nil {
if os.IsNotExist(err) {
return nil, nil
}
return nil, fmt.Errorf("read skills directory: %w", err)
}
if !info.IsDir() {
return nil, fmt.Errorf("skills path is not a directory: %s", resolved)
}
entries, err := os.ReadDir(resolved)
if err != nil {
return nil, fmt.Errorf("read skills directory: %w", err)
}
var skills []Skill
for _, entry := range entries {
name := entry.Name()
if strings.HasPrefix(name, ".") {
continue
}
entryPath := filepath.Join(resolved, name)
// Shape 1: Codex-style directory with SKILL.md.
isDir, err := isDirectory(entryPath)
if err != nil {
return nil, err
}
if isDir {
skillFile := filepath.Join(entryPath, SkillFileName)
fi, err := os.Stat(skillFile)
if err != nil {
if os.IsNotExist(err) {
continue
}
return nil, fmt.Errorf("read skill %s: %w", skillFile, err)
}
if !fi.Mode().IsRegular() {
continue
}
skill, err := loadSkillFile(name, skillFile, true)
if err != nil {
return nil, err
}
if skill != nil {
skills = append(skills, *skill)
}
continue
}
// Shape 2: Legacy flat .md file.
if !entry.Type().IsRegular() {
continue
}
if !strings.EqualFold(filepath.Ext(name), ".md") {
continue
}
skill, err := loadSkillFile(name[:len(name)-len(filepath.Ext(name))], entryPath, false)
if err != nil {
return nil, err
}
if skill != nil {
skills = append(skills, *skill)
}
}
sort.Slice(skills, func(i, j int) bool {
if skills[i].Name != skills[j].Name {
return skills[i].Name < skills[j].Name
}
return skills[i].Path < skills[j].Path
})
return skills, nil
}
// isDirectory reports whether path is a directory. For symlinks it resolves
// the target and checks the resolved type.
func isDirectory(path string) (bool, error) {
fi, err := os.Lstat(path)
if err != nil {
if os.IsNotExist(err) {
return false, nil
}
return false, fmt.Errorf("read skill %s: %w", path, err)
}
if fi.IsDir() {
return true, nil
}
if fi.Mode()&os.ModeSymlink != 0 {
resolved, err := filepath.EvalSymlinks(path)
if err != nil {
return false, nil
}
target, err := os.Stat(resolved)
if err != nil {
return false, nil
}
return target.IsDir(), nil
}
return false, nil
}
// loadSkillFile reads one skill markdown file and returns a Skill or nil when
// the instruction body is empty. legacyBaseName is used as the default name
// when frontmatter does not supply one and isDir is false.
func loadSkillFile(legacyBaseName string, path string, isDirSkill bool) (*Skill, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("read skill %s: %w", path, err)
}
body, fmName, fmDesc, _, err := parseSkillFile(path, data)
if err != nil {
return nil, err
}
if body == "" {
return nil, nil
}
name := fmName
if name == "" {
if isDirSkill {
name = filepath.Base(filepath.Dir(path))
} else {
name = legacyBaseName
}
}
if err := validateSkillName(name); err != nil {
return nil, err
}
return &Skill{
Name: name,
Description: strings.TrimSpace(fmDesc),
Path: path,
Content: body,
}, nil
}
// FormatForSystemPrompt formats skills for injection into the system prompt.
// Returns an empty string if items is empty.
func FormatForSystemPrompt(items []Skill) string {
if len(items) == 0 {
return ""
}
var b strings.Builder
b.WriteString("Global skills:\n")
b.WriteString("The following global skills are available for every conversation. Apply them when relevant.")
for _, skill := range items {
b.WriteString("\n\n### ")
b.WriteString(skill.Name)
b.WriteString("\nSource: ")
b.WriteString(skill.Path)
if skill.Description != "" {
b.WriteString("\nDescription: ")
b.WriteString(skill.Description)
}
b.WriteString("\n\n")
b.WriteString(skill.Content)
}
return b.String()
}
// FormatList formats skills for slash-command output.
func FormatList(items []Skill) string {
if len(items) == 0 {
return "No global skills installed."
}
var b strings.Builder
b.WriteString("Global skills:\n")
for _, skill := range items {
b.WriteString("- ")
b.WriteString(skill.Name)
if skill.Description != "" {
b.WriteString(" — ")
b.WriteString(skill.Description)
}
b.WriteString("\n Source: ")
b.WriteString(skill.Path)
b.WriteString("\n")
}
return b.String()
}
// InstallGlobal installs a skill into DefaultDir.
func InstallGlobal(ctx context.Context, source string) (InstallResult, error) {
return InstallDir(ctx, DefaultDir, source)
}
// InstallDir installs a Codex-style skill directory under the target root.
func InstallDir(ctx context.Context, dir string, source string) (InstallResult, error) {
source = strings.TrimSpace(source)
if source == "" {
return InstallResult{}, fmt.Errorf("skill source is required")
}
resolvedRoot, err := config.ResolvePath(dir)
if err != nil {
return InstallResult{}, err
}
resolvedRoot, err = filepath.Abs(resolvedRoot)
if err != nil {
return InstallResult{}, err
}
// HTTP(S) source.
if strings.HasPrefix(source, "http://") || strings.HasPrefix(source, "https://") {
return installFromHTTP(ctx, resolvedRoot, source)
}
// Local source.
resolvedSource, err := config.ResolvePath(source)
if err != nil {
return InstallResult{}, err
}
resolvedSource, err = filepath.Abs(resolvedSource)
if err != nil {
return InstallResult{}, err
}
fi, err := os.Stat(resolvedSource)
if err != nil {
return InstallResult{}, fmt.Errorf("read skill source: %w", err)
}
if fi.IsDir() {
return installFromDirectory(ctx, resolvedRoot, resolvedSource)
}
if fi.Mode().IsRegular() && strings.EqualFold(filepath.Ext(resolvedSource), ".md") {
return installFromFile(ctx, resolvedRoot, resolvedSource)
}
return InstallResult{}, fmt.Errorf("skill source must be a directory, .md file, or http(s) URL")
}
func installFromHTTP(ctx context.Context, root string, url string) (InstallResult, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return InstallResult{}, fmt.Errorf("download skill: %w", err)
}
client := &http.Client{Timeout: 30 * time.Second}
resp, err := client.Do(req)
if err != nil {
return InstallResult{}, fmt.Errorf("download skill: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return InstallResult{}, fmt.Errorf("download skill: HTTP %d", resp.StatusCode)
}
limited := io.LimitReader(resp.Body, int64(maxInstallDownloadBytes)+1)
data, err := io.ReadAll(limited)
if err != nil {
return InstallResult{}, fmt.Errorf("download skill: %w", err)
}
if len(data) > maxInstallDownloadBytes {
return InstallResult{}, fmt.Errorf("skill file is too large: max %d bytes", maxInstallDownloadBytes)
}
return installMarkdownBytes(root, data, deriveNameFromURL(url))
}
func installFromDirectory(ctx context.Context, root string, sourceDir string) (InstallResult, error) {
skillFile := filepath.Join(sourceDir, SkillFileName)
fi, err := os.Stat(skillFile)
if err != nil {
return InstallResult{}, fmt.Errorf("read skill source: %w", err)
}
if !fi.Mode().IsRegular() {
return InstallResult{}, fmt.Errorf("skill source must contain a regular %s file", SkillFileName)
}
data, err := os.ReadFile(skillFile)
if err != nil {
return InstallResult{}, fmt.Errorf("read skill source: %w", err)
}
body, fmName, _, _, err := parseSkillFile(skillFile, data)
if err != nil {
return InstallResult{}, err
}
_ = body
name := fmName
if name == "" {
name = filepath.Base(sourceDir)
}
if err := validateSkillName(name); err != nil {
return InstallResult{}, err
}
targetDir := filepath.Join(root, name)
if err := createSkillTarget(root, targetDir); err != nil {
return InstallResult{}, err
}
// Copy tree.
if err := copyTree(ctx, sourceDir, targetDir); err != nil {
os.RemoveAll(targetDir)
return InstallResult{}, err
}
installedPath := filepath.Join(targetDir, SkillFileName)
installedData, _ := os.ReadFile(installedPath)
instBody, instName, instDesc, _, _ := parseSkillFile(installedPath, installedData)
if instName == "" {
instName = name
}
return InstallResult{
Skill: Skill{
Name: instName,
Description: strings.TrimSpace(instDesc),
Path: installedPath,
Content: instBody,
},
TargetDir: targetDir,
}, nil
}
func installFromFile(ctx context.Context, root string, sourcePath string) (InstallResult, error) {
data, err := os.ReadFile(sourcePath)
if err != nil {
return InstallResult{}, fmt.Errorf("read skill source: %w", err)
}
baseName := filepath.Base(sourcePath)
baseNameNoExt := baseName[:len(baseName)-len(filepath.Ext(baseName))]
// When the source is literally SKILL.md, use the parent directory name.
if strings.EqualFold(baseName, SkillFileName) {
baseNameNoExt = filepath.Base(filepath.Dir(sourcePath))
}
return installMarkdownBytes(root, data, baseNameNoExt)
}
// installMarkdownBytes parses markdown bytes and writes a Codex-style
// SKILL.md under <root>/<name>/. If the source had frontmatter, the original
// bytes are preserved. Otherwise a normalized header is prepended.
func installMarkdownBytes(root string, data []byte, fallbackName string) (InstallResult, error) {
normalized := strings.ReplaceAll(string(data), "\r\n", "\n")
lines := strings.Split(normalized, "\n")
hasFrontmatter := false
if len(lines) >= 2 && lines[0] == "---" {
for i := 1; i < len(lines); i++ {
if lines[i] == "---" {
hasFrontmatter = true
break
}
}
}
var body string
var fmName string
var fmDesc string
var err error
if hasFrontmatter {
body, fmName, fmDesc, _, err = parseSkillFile("<source>", data)
if err != nil {
return InstallResult{}, err
}
} else {
body = strings.TrimSpace(normalized)
}
name := fmName
if name == "" {
name = fallbackName
}
if err := validateSkillName(name); err != nil {
return InstallResult{}, err
}
if body == "" {
return InstallResult{}, fmt.Errorf("skill body is empty")
}
targetDir := filepath.Join(root, name)
if err := createSkillTarget(root, targetDir); err != nil {
return InstallResult{}, err
}
var writeData []byte
if hasFrontmatter {
writeData = data
} else {
var b strings.Builder
b.WriteString("---\n")
fmt.Fprintf(&b, "name: %s\n", name)
b.WriteString("description: \"\"\n")
b.WriteString("---\n\n")
b.WriteString(body)
b.WriteString("\n")
writeData = []byte(b.String())
}
targetFile := filepath.Join(targetDir, SkillFileName)
if err := safeio.AtomicWriteFile(targetFile, writeData, 0o644); err != nil {
os.RemoveAll(targetDir)
return InstallResult{}, err
}
return InstallResult{
Skill: Skill{
Name: name,
Description: strings.TrimSpace(fmDesc),
Path: targetFile,
Content: body,
},
TargetDir: targetDir,
}, nil
}
func createSkillTarget(root, targetDir string) error {
if err := os.MkdirAll(root, 0o755); err != nil {
return fmt.Errorf("create skills directory: %w", err)
}
if err := os.Mkdir(targetDir, 0o755); err != nil {
if os.IsExist(err) {
return fmt.Errorf("skill already exists: %s", targetDir)
}
return fmt.Errorf("create skill directory: %w", err)
}
return nil
}
// deriveNameFromURL extracts a filename stem from a URL for use as a fallback name.
func deriveNameFromURL(rawURL string) string {
// Strip query and fragment.
u := rawURL
if idx := strings.IndexAny(u, "?#"); idx >= 0 {
u = u[:idx]
}
base := filepath.Base(u)
ext := filepath.Ext(base)
if ext != "" {
base = base[:len(base)-len(ext)]
}
if base == "" || strings.EqualFold(base, "SKILL") {
return "downloaded-skill"
}
return base
}
// copyTree copies srcDir to dstDir, rejecting symlinks inside the source tree.
func copyTree(ctx context.Context, srcDir string, dstDir string) error {
return filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
if err != nil {
return err
}
if err := ctx.Err(); err != nil {
return err
}
rel, err := filepath.Rel(srcDir, path)
if err != nil {
return err
}
target := filepath.Join(dstDir, rel)
// Reject any symlink inside the source tree.
lfi, lerr := os.Lstat(path)
if lerr != nil {
return lerr
}
if lfi.Mode()&os.ModeSymlink != 0 {
return fmt.Errorf("skill source contains symlink: %s", path)
}
if lfi.IsDir() {
return os.MkdirAll(target, 0o755)
}
if !lfi.Mode().IsRegular() {
return nil
}
data, err := os.ReadFile(path)
if err != nil {
return err
}
return safeio.AtomicWriteFile(target, data, 0o644)
})
}
+404
View File
@@ -0,0 +1,404 @@
package skills
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
)
func TestLoadDirMissingReturnsEmpty(t *testing.T) {
dir := filepath.Join(t.TempDir(), "missing")
skills, err := LoadDir(dir)
if err != nil {
t.Fatalf("LoadDir(missing) returned error: %v", err)
}
if len(skills) != 0 {
t.Fatalf("expected 0 skills, got %d", len(skills))
}
}
func TestLoadDirRejectsFilePath(t *testing.T) {
file := filepath.Join(t.TempDir(), "not-a-dir")
if err := os.WriteFile(file, []byte("x"), 0o644); err != nil {
t.Fatal(err)
}
_, err := LoadDir(file)
if err == nil {
t.Fatal("expected error for file path")
}
if !strings.Contains(err.Error(), "skills path is not a directory") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestLoadDirLoadsDirectorySkillsSorted(t *testing.T) {
dir := t.TempDir()
// b-skill loads second by name.
bDir := filepath.Join(dir, "b-skill")
if err := os.MkdirAll(bDir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(bDir, SkillFileName), []byte("---\nname: b-skill\ndescription: Beta skill\n---\n\nBeta content.\n"), 0o644); err != nil {
t.Fatal(err)
}
// a-skill loads first by name.
aDir := filepath.Join(dir, "a-skill")
if err := os.MkdirAll(aDir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(aDir, SkillFileName), []byte("---\nname: a-skill\ndescription: Alpha skill\n---\n\nAlpha content.\n"), 0o644); err != nil {
t.Fatal(err)
}
// Visible directory without SKILL.md — ignored.
noSkillDir := filepath.Join(dir, "no-skill-here")
if err := os.MkdirAll(noSkillDir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(noSkillDir, "README.md"), []byte("not a skill"), 0o644); err != nil {
t.Fatal(err)
}
skills, err := LoadDir(dir)
if err != nil {
t.Fatalf("LoadDir: %v", err)
}
if len(skills) != 2 {
t.Fatalf("expected 2 skills, got %d: %+v", len(skills), skills)
}
if skills[0].Name != "a-skill" {
t.Fatalf("skills[0].Name = %q", skills[0].Name)
}
if skills[0].Description != "Alpha skill" {
t.Fatalf("skills[0].Description = %q", skills[0].Description)
}
if skills[0].Content != "Alpha content." {
t.Fatalf("skills[0].Content = %q", skills[0].Content)
}
if !strings.HasSuffix(skills[0].Path, SkillFileName) {
t.Fatalf("skills[0].Path = %q, want suffix %s", skills[0].Path, SkillFileName)
}
if skills[1].Name != "b-skill" {
t.Fatalf("skills[1].Name = %q", skills[1].Name)
}
if skills[1].Description != "Beta skill" {
t.Fatalf("skills[1].Description = %q", skills[1].Description)
}
}
func TestLoadDirKeepsLegacyMarkdownSkills(t *testing.T) {
dir := t.TempDir()
write := func(name, content string) {
if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644); err != nil {
t.Fatal(err)
}
}
write("legacy.md", " legacy content ")
write(".hidden.md", "hidden")
write("notes.txt", "not a skill")
write("empty.md", " \n ")
skills, err := LoadDir(dir)
if err != nil {
t.Fatalf("LoadDir: %v", err)
}
if len(skills) != 1 {
t.Fatalf("expected 1 skill, got %d: %+v", len(skills), skills)
}
if skills[0].Name != "legacy" {
t.Fatalf("skill.Name = %q", skills[0].Name)
}
if skills[0].Content != "legacy content" {
t.Fatalf("skill.Content = %q", skills[0].Content)
}
}
func TestInstallDirInstallsMarkdownFileAsSkillDirectory(t *testing.T) {
sourceDir := t.TempDir()
targetRoot := filepath.Join(t.TempDir(), "skills")
sourcePath := filepath.Join(sourceDir, "source.md")
sourceContent := "---\nname: uiux-pro-max\ndescription: UI workflow\n---\n\nReview the UI carefully.\n"
if err := os.WriteFile(sourcePath, []byte(sourceContent), 0o644); err != nil {
t.Fatal(err)
}
result, err := InstallDir(context.Background(), targetRoot, sourcePath)
if err != nil {
t.Fatalf("InstallDir: %v", err)
}
absRoot, _ := filepath.Abs(targetRoot)
wantDir := filepath.Join(absRoot, "uiux-pro-max")
if result.TargetDir != wantDir {
t.Fatalf("TargetDir = %q, want %q", result.TargetDir, wantDir)
}
skillFile := filepath.Join(wantDir, SkillFileName)
if _, err := os.Stat(skillFile); err != nil {
t.Fatalf("SKILL.md not created: %v", err)
}
loaded, err := LoadDir(targetRoot)
if err != nil {
t.Fatalf("LoadDir: %v", err)
}
if len(loaded) != 1 {
t.Fatalf("expected 1 loaded skill, got %d", len(loaded))
}
if loaded[0].Name != "uiux-pro-max" {
t.Fatalf("loaded name = %q", loaded[0].Name)
}
if loaded[0].Content != "Review the UI carefully." {
t.Fatalf("loaded content = %q", loaded[0].Content)
}
}
func TestInstallDirCopiesDirectorySkill(t *testing.T) {
sourceDir := t.TempDir()
targetRoot := filepath.Join(t.TempDir(), "skills")
skillSrc := filepath.Join(sourceDir, "grill-me")
if err := os.MkdirAll(filepath.Join(skillSrc, "references"), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(skillSrc, SkillFileName), []byte("---\nname: grill-me\ndescription: Grill skill\n---\n\nGrill content.\n"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(skillSrc, "references", "guide.md"), []byte("# Guide\n"), 0o644); err != nil {
t.Fatal(err)
}
result, err := InstallDir(context.Background(), targetRoot, skillSrc)
if err != nil {
t.Fatalf("InstallDir: %v", err)
}
skillFile := filepath.Join(result.TargetDir, SkillFileName)
if _, err := os.Stat(skillFile); err != nil {
t.Fatalf("SKILL.md not copied: %v", err)
}
guideFile := filepath.Join(result.TargetDir, "references", "guide.md")
if _, err := os.Stat(guideFile); err != nil {
t.Fatalf("references/guide.md not copied: %v", err)
}
loaded, err := LoadDir(targetRoot)
if err != nil {
t.Fatalf("LoadDir: %v", err)
}
if len(loaded) != 1 || loaded[0].Name != "grill-me" {
t.Fatalf("loaded = %+v", loaded)
}
}
func TestInstallDirRejectsExistingSkill(t *testing.T) {
targetRoot := filepath.Join(t.TempDir(), "skills")
existingDir := filepath.Join(targetRoot, "grill-me")
if err := os.MkdirAll(existingDir, 0o755); err != nil {
t.Fatal(err)
}
sourceDir := t.TempDir()
sourcePath := filepath.Join(sourceDir, "source.md")
if err := os.WriteFile(sourcePath, []byte("---\nname: grill-me\ndescription: d\n---\n\nbody\n"), 0o644); err != nil {
t.Fatal(err)
}
_, err := InstallDir(context.Background(), targetRoot, sourcePath)
if err == nil {
t.Fatal("expected error for existing skill")
}
if !strings.Contains(err.Error(), "skill already exists") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestInstallDirRejectsInvalidName(t *testing.T) {
targetRoot := filepath.Join(t.TempDir(), "skills")
sourceDir := t.TempDir()
sourcePath := filepath.Join(sourceDir, "source.md")
if err := os.WriteFile(sourcePath, []byte("---\nname: ../bad\ndescription: d\n---\n\nbody\n"), 0o644); err != nil {
t.Fatal(err)
}
_, err := InstallDir(context.Background(), targetRoot, sourcePath)
if err == nil {
t.Fatal("expected error for invalid name")
}
if !strings.Contains(err.Error(), "invalid skill name") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestInstallDirRequiresSource(t *testing.T) {
targetRoot := filepath.Join(t.TempDir(), "skills")
_, err := InstallDir(context.Background(), targetRoot, "")
if err == nil {
t.Fatal("expected error for empty source")
}
if !strings.Contains(err.Error(), "skill source is required") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestInstallDirRejectsBadSourceType(t *testing.T) {
targetRoot := filepath.Join(t.TempDir(), "skills")
sourceDir := t.TempDir()
sourcePath := filepath.Join(sourceDir, "something.yaml")
if err := os.WriteFile(sourcePath, []byte("key: value\n"), 0o644); err != nil {
t.Fatal(err)
}
_, err := InstallDir(context.Background(), targetRoot, sourcePath)
if err == nil {
t.Fatal("expected error for non-.md non-dir source")
}
if !strings.Contains(err.Error(), "skill source must be a directory") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestInstallDirWithoutFrontmatterCreatesNormalizedHeader(t *testing.T) {
targetRoot := filepath.Join(t.TempDir(), "skills")
sourceDir := t.TempDir()
sourcePath := filepath.Join(sourceDir, "my-skill.md")
if err := os.WriteFile(sourcePath, []byte("Just body content.\n"), 0o644); err != nil {
t.Fatal(err)
}
result, err := InstallDir(context.Background(), targetRoot, sourcePath)
if err != nil {
t.Fatalf("InstallDir: %v", err)
}
if result.Skill.Name != "my-skill" {
t.Fatalf("name = %q", result.Skill.Name)
}
installed, err := os.ReadFile(filepath.Join(result.TargetDir, SkillFileName))
if err != nil {
t.Fatal(err)
}
content := string(installed)
if !strings.HasPrefix(content, "---\nname: my-skill\n") {
t.Fatalf("installed SKILL.md missing normalized header: %q", content)
}
if !strings.Contains(content, "Just body content.") {
t.Fatalf("installed SKILL.md missing body: %q", content)
}
}
func TestFormatForSystemPromptIncludesDescription(t *testing.T) {
items := []Skill{
{Name: "first", Path: "/tmp/first.md", Content: "skill one"},
{Name: "second", Description: "Described", Path: "/tmp/second.md", Content: "skill two"},
}
out := FormatForSystemPrompt(items)
if out == "" {
t.Fatal("FormatForSystemPrompt returned empty string")
}
if !strings.Contains(out, "Global skills:") {
t.Fatal("missing 'Global skills:' header")
}
if !strings.Contains(out, "### first") {
t.Fatal("missing '### first'")
}
if !strings.Contains(out, "### second") {
t.Fatal("missing '### second'")
}
if !strings.Contains(out, "Source: /tmp/first.md") {
t.Fatal("missing source for first")
}
if !strings.Contains(out, "Source: /tmp/second.md") {
t.Fatal("missing source for second")
}
if !strings.Contains(out, "skill one") {
t.Fatal("missing content for first")
}
if !strings.Contains(out, "skill two") {
t.Fatal("missing content for second")
}
// First skill has no description — the Description: line must be absent.
firstBlock := out[:strings.Index(out, "### second")]
if strings.Contains(firstBlock, "Description:") {
t.Fatal("first skill should not have Description line")
}
// Second skill has description — the line must appear.
if !strings.Contains(out, "Description: Described") {
t.Fatal("missing Description line for second skill")
}
}
func TestFormatForSystemPromptEmptyReturnsEmpty(t *testing.T) {
if got := FormatForSystemPrompt(nil); got != "" {
t.Fatalf("expected empty string, got %q", got)
}
}
func TestFormatListNonEmpty(t *testing.T) {
items := []Skill{
{Name: "alpha", Description: "Alpha desc", Path: "/tmp/alpha/SKILL.md", Content: "..."},
{Name: "beta", Path: "/tmp/beta/SKILL.md", Content: "..."},
}
out := FormatList(items)
if !strings.Contains(out, "Global skills:") {
t.Fatal("missing header")
}
if !strings.Contains(out, "- alpha — Alpha desc") {
t.Fatal("missing alpha with description")
}
if !strings.Contains(out, " Source: /tmp/alpha/SKILL.md") {
t.Fatal("missing alpha source")
}
if !strings.Contains(out, "- beta\n") {
t.Fatal("missing beta without description")
}
// Ensure beta has no em-dash line.
if strings.Contains(out, "- beta —") {
t.Fatal("beta should not have description separator")
}
}
func TestFormatListEmpty(t *testing.T) {
out := FormatList(nil)
if out != "No global skills installed." {
t.Fatalf("unexpected: %q", out)
}
}
func TestLoadDirRejectsInvalidFrontmatterName(t *testing.T) {
dir := t.TempDir()
skillDir := filepath.Join(dir, "bad-skill")
if err := os.MkdirAll(skillDir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(skillDir, SkillFileName), []byte("---\nname: ../bad\n---\n\nbody\n"), 0o644); err != nil {
t.Fatal(err)
}
_, err := LoadDir(dir)
if err == nil {
t.Fatal("expected error for invalid frontmatter name")
}
if !strings.Contains(err.Error(), "invalid skill name") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestParseSkillFileMissingClosingDelimiter(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "bad.md")
_, _, _, _, err := parseSkillFile(path, []byte("---\nname: foo\nbody\n"))
if err == nil {
t.Fatal("expected error for missing closing ---")
}
if !strings.Contains(err.Error(), "missing closing ---") {
t.Fatalf("unexpected error: %v", err)
}
}
+153
View File
@@ -0,0 +1,153 @@
package tools
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"agentu/pkg/llm"
)
// ────────────────────────────────────────────────────────────────────────────
// ApplyPatchTool — multi-file patch using V4 format
// ────────────────────────────────────────────────────────────────────────────
type ApplyPatchTool struct {
workingDir string
}
func NewApplyPatchTool(workingDir string) *ApplyPatchTool {
return &ApplyPatchTool{workingDir: workingDir}
}
// MutatingTool interface.
func (*ApplyPatchTool) Mutates() bool { return true }
// CallMutatingTool — dry_run is non-mutating.
func (t *ApplyPatchTool) MutatesCall(args json.RawMessage) bool {
var a applyPatchArgs
if err := json.Unmarshal(args, &a); err != nil {
return true // conservative
}
return !a.DryRun
}
// PreviewableTool interface.
func (t *ApplyPatchTool) Preview(ctx context.Context, args json.RawMessage) (string, error) {
var a applyPatchArgs
if err := json.Unmarshal(args, &a); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
return summarizePatch(a.Patch)
}
// TargetPathTool interface.
func (t *ApplyPatchTool) TargetPaths(args json.RawMessage) ([]string, error) {
var a applyPatchArgs
if err := json.Unmarshal(args, &a); err != nil {
return nil, fmt.Errorf("parse arguments: %w", err)
}
return extractTargetPaths(t.workingDir, a.Patch)
}
// Definition.
func (t *ApplyPatchTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "apply_patch",
Description: "Apply a multi-file patch to create, update, or delete files. Uses a V4 diff format with *** Begin Patch / *** Add File / *** Update File / *** Delete File markers. Prefer this over multiple file_edit calls when changing 2+ files.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"patch": {"type": "string", "description": "V4 format patch text. Must start with *** Begin Patch and end with *** End Patch."},
"dry_run": {"type": "boolean", "description": "Preview the patch without writing any files."}
},
"required": ["patch"],
"additionalProperties": false
}`),
},
}
}
type applyPatchArgs struct {
Patch string `json:"patch"`
DryRun bool `json:"dry_run"`
}
// Execute.
func (t *ApplyPatchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args applyPatchArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
if strings.TrimSpace(args.Patch) == "" {
return "", fmt.Errorf("patch is empty")
}
pp, err := parsePatch(args.Patch)
if err != nil {
return "", err
}
if args.DryRun {
summary, err := summarizePatch(args.Patch)
if err != nil {
return "", err
}
return "dry run — " + summary, nil
}
rollback := NewRollbackStore()
var snapshotIDs []string
for _, pf := range pp.files {
path, err := resolvePathForWrite(t.workingDir, pf.path)
if err != nil {
return "", err
}
id, err := rollback.Capture("apply_patch", pf.path, path)
if err != nil {
return "", err
}
snapshotIDs = append(snapshotIDs, id)
}
var results []string
for _, pf := range pp.files {
if err := ctx.Err(); err != nil {
if rollbackErr := restoreSnapshots(rollback, snapshotIDs); rollbackErr != nil {
return "", fmt.Errorf("%w; rollback failed: %v", err, rollbackErr)
}
return "", err
}
result, err := applyPatchFile(t.workingDir, pf)
if err != nil {
if rollbackErr := restoreSnapshots(rollback, snapshotIDs); rollbackErr != nil {
return "", fmt.Errorf("apply patch: %w; rollback failed: %v", err, rollbackErr)
}
return "", fmt.Errorf("apply patch: %w", err)
}
results = append(results, result)
}
summary := fmt.Sprintf("applied patch: %d file(s) changed\n%s", len(results), strings.Join(results, "\n"))
return summary, nil
}
func restoreSnapshots(store *RollbackStore, ids []string) error {
var errs []error
for i := len(ids) - 1; i >= 0; i-- {
if _, err := store.Restore(ids[i]); err != nil {
errs = append(errs, err)
}
}
return errors.Join(errs...)
}
+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()
}
+103 -23
View File
@@ -13,7 +13,8 @@ import (
"sort"
"strings"
"agentu/internal/llm"
"agentu/internal/safeio"
"agentu/pkg/llm"
)
type FileReadTool struct {
@@ -81,6 +82,34 @@ type FileWriteTool struct {
func NewFileWriteTool(workingDir string) *FileWriteTool {
return &FileWriteTool{workingDir: workingDir}
}
func (*FileWriteTool) Mutates() bool { return true }
func (t *FileWriteTool) Preview(ctx context.Context, raw json.RawMessage) (string, error) {
var args fileWriteArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
action := "write"
if args.Append {
action = "append"
}
return fmt.Sprintf("%s %d bytes to %s", action, len(args.Content), args.Path), ctx.Err()
}
func (t *FileWriteTool) TargetPaths(raw json.RawMessage) ([]string, error) {
var args fileWriteArgs
if err := json.Unmarshal(raw, &args); err != nil {
return nil, fmt.Errorf("parse arguments: %w", err)
}
if strings.TrimSpace(args.Path) == "" {
return nil, errors.New("path is required")
}
path, err := resolvePathForWrite(t.workingDir, args.Path)
if err != nil {
return nil, err
}
return []string{path}, nil
}
func (t *FileWriteTool) Definition() llm.Tool {
return llm.Tool{
@@ -116,28 +145,28 @@ func (t *FileWriteTool) Execute(ctx context.Context, raw json.RawMessage) (strin
if strings.TrimSpace(args.Path) == "" {
return "", errors.New("path is required")
}
path, err := resolvePath(t.workingDir, args.Path)
path, err := resolvePathForWrite(t.workingDir, args.Path)
if err != nil {
return "", err
}
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
return "", err
}
flag := os.O_CREATE | os.O_WRONLY | os.O_TRUNC
if args.Append {
flag = os.O_CREATE | os.O_WRONLY | os.O_APPEND
}
file, err := os.OpenFile(path, flag, 0o644)
if err != nil {
return "", err
}
defer file.Close()
if _, err := file.WriteString(args.Content); err != nil {
return "", err
}
if err := ctx.Err(); err != nil {
return "", err
}
content := []byte(args.Content)
perm := os.FileMode(0o644)
if info, statErr := os.Stat(path); statErr == nil {
perm = info.Mode().Perm()
}
if args.Append {
existing, readErr := os.ReadFile(path)
if readErr != nil && !os.IsNotExist(readErr) {
return "", readErr
}
content = append(existing, content...)
}
if err := safeio.AtomicWriteFile(path, content, perm); err != nil {
return "", err
}
action := "wrote"
if args.Append {
action = "appended"
@@ -223,7 +252,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
@@ -236,6 +267,8 @@ type fileSearchArgs struct {
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 +289,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 +338,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 +371,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 +394,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])
}
}
+209
View File
@@ -0,0 +1,209 @@
package tools
import (
"context"
"encoding/json"
"errors"
"fmt"
"os"
"strings"
"agentu/internal/safeio"
"agentu/pkg/llm"
)
type FileEditTool struct {
workingDir string
}
func NewFileEditTool(workingDir string) *FileEditTool {
return &FileEditTool{workingDir: workingDir}
}
func (*FileEditTool) Mutates() bool { return true }
func (t *FileEditTool) Preview(ctx context.Context, raw json.RawMessage) (string, error) {
var args fileEditArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
if args.StartLine > 0 {
end := args.EndLine
if end == 0 {
end = args.StartLine
}
return fmt.Sprintf("edit %s lines %d-%d", args.Path, args.StartLine, end), ctx.Err()
}
return fmt.Sprintf("replace text in %s", args.Path), ctx.Err()
}
func (t *FileEditTool) TargetPaths(raw json.RawMessage) ([]string, error) {
var args fileEditArgs
if err := json.Unmarshal(raw, &args); err != nil {
return nil, fmt.Errorf("parse arguments: %w", err)
}
if strings.TrimSpace(args.Path) == "" {
return nil, errors.New("path is required")
}
path, err := resolvePath(t.workingDir, args.Path)
if err != nil {
return nil, err
}
return []string{path}, nil
}
func (t *FileEditTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "file_edit",
Description: "Edit an existing UTF-8 text file by replacing an inclusive 1-based line range or an exact old text pattern.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Existing file path relative to the configured agent working directory."},
"start_line": {"type": "integer", "description": "1-based first line for a line-range edit."},
"end_line": {"type": "integer", "description": "Optional 1-based final line; defaults to start_line when omitted."},
"old": {"type": "string", "description": "Exact text to replace for a pattern edit; not regex."},
"new": {"type": "string", "description": "Replacement text; may be empty string for deletion."},
"replace_all": {"type": "boolean", "description": "When true, replace every exact old occurrence; otherwise require exactly one occurrence."}
},
"required": ["path", "new"],
"additionalProperties": false
}`),
},
}
}
type fileEditArgs struct {
Path string `json:"path"`
StartLine int `json:"start_line"`
EndLine int `json:"end_line"`
Old string `json:"old"`
New *string `json:"new"`
ReplaceAll bool `json:"replace_all"`
}
func (t *FileEditTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args fileEditArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
if strings.TrimSpace(args.Path) == "" {
return "", errors.New("path is required")
}
if args.New == nil {
return "", errors.New("new is required")
}
path, err := resolvePath(t.workingDir, args.Path)
if err != nil {
return "", err
}
info, err := os.Stat(path)
if err != nil {
return "", fmt.Errorf("cannot access %s: %w", args.Path, err)
}
if info.IsDir() {
return "", fmt.Errorf("%s is a directory, not a file", args.Path)
}
perm := info.Mode().Perm()
hasLine := args.StartLine > 0 || args.EndLine > 0
hasOld := strings.TrimSpace(args.Old) != ""
if hasLine && hasOld {
return "", errors.New("provide either start_line/end_line or old, not both")
}
if !hasLine && !hasOld {
return "", errors.New("provide either start_line/end_line or old")
}
if err := ctx.Err(); err != nil {
return "", err
}
data, err := os.ReadFile(path)
if err != nil {
return "", err
}
content := string(data)
var result string
if hasOld {
result, err = editByPattern(content, args.Path, args.Old, *args.New, args.ReplaceAll)
} else {
result, err = editByLineRange(content, args.Path, args.StartLine, args.EndLine, *args.New)
}
if err != nil {
return "", err
}
if err := ctx.Err(); err != nil {
return "", err
}
if err := safeio.AtomicWriteFile(path, []byte(result), perm); err != nil {
return "", err
}
switch {
case hasOld:
count := strings.Count(content, args.Old)
if args.ReplaceAll {
return fmt.Sprintf("edited %s: replaced %d occurrences", args.Path, count), nil
}
return fmt.Sprintf("edited %s: replaced 1 occurrence", args.Path), nil
default:
endLine := args.EndLine
if endLine == 0 {
endLine = args.StartLine
}
return fmt.Sprintf("edited %s: replaced lines %d-%d", args.Path, args.StartLine, endLine), nil
}
}
func editByLineRange(content, displayPath string, startLine, endLine int, replacement string) (string, error) {
if startLine < 1 {
return "", errors.New("start_line must be >= 1")
}
if endLine == 0 {
endLine = startLine
}
if endLine < startLine {
return "", fmt.Errorf("end_line (%d) must be >= start_line (%d)", endLine, startLine)
}
lines := strings.Split(content, "\n")
if startLine > len(lines) {
return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
}
if endLine > len(lines) {
return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
}
var b strings.Builder
b.WriteString(strings.Join(lines[:startLine-1], "\n"))
if startLine > 1 {
b.WriteString("\n")
}
b.WriteString(replacement)
if endLine < len(lines) {
if !strings.HasSuffix(replacement, "\n") && replacement != "" {
b.WriteString("\n")
}
b.WriteString(strings.Join(lines[endLine:], "\n"))
}
return b.String(), nil
}
func editByPattern(content, displayPath, old, replacement string, replaceAll bool) (string, error) {
count := strings.Count(content, old)
if count == 0 {
return "", errors.New("old text not found")
}
if !replaceAll && count != 1 {
return "", fmt.Errorf("old text matched %d occurrences; set replace_all=true or provide a more specific old value", count)
}
return strings.ReplaceAll(content, old, replacement), nil
}
+136
View File
@@ -0,0 +1,136 @@
package tools
import (
"context"
"encoding/json"
"os"
"path/filepath"
"strings"
"testing"
)
func TestFileEditReplacesLineRange(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("one\ntwo\nthree\n"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","start_line":2,"end_line":2,"new":"TWO\n"}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "replaced lines 2-2") {
t.Fatalf("output = %q", out)
}
got, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(got) != "one\nTWO\nthree\n" {
t.Fatalf("file content = %q", string(got))
}
}
func TestFileEditReplacesExactPattern(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("alpha beta alpha"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"beta","new":"BETA"}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "replaced 1 occurrence") {
t.Fatalf("output = %q", out)
}
got, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(got) != "alpha BETA alpha" {
t.Fatalf("file content = %q", string(got))
}
}
func TestFileEditRejectsAmbiguousPattern(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("x x"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"x","new":"y"}`))
if err == nil {
t.Fatal("expected error for ambiguous pattern")
}
if !strings.Contains(err.Error(), "matched 2 occurrences") {
t.Fatalf("error = %v", err)
}
}
func TestFileEditReplaceAll(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("x x"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"x","new":"y","replace_all":true}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "replaced 2 occurrences") {
t.Fatalf("output = %q", out)
}
got, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(got) != "y y" {
t.Fatalf("file content = %q", string(got))
}
}
func TestFileEditRejectsEscapingPath(t *testing.T) {
dir := t.TempDir()
edit := NewFileEditTool(dir)
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"../outside.txt","old":"x","new":"y"}`))
if err == nil {
t.Fatal("expected error for escaping path")
}
if !strings.Contains(err.Error(), "escapes") {
t.Fatalf("error = %v", err)
}
}
func TestFileEditRejectsMissingNew(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"hello"}`))
if err == nil {
t.Fatal("expected error for missing new")
}
if !strings.Contains(err.Error(), "new is required") {
t.Fatalf("error = %v", err)
}
}
func TestFileEditRejectsNonExistentFile(t *testing.T) {
dir := t.TempDir()
edit := NewFileEditTool(dir)
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"missing.txt","old":"x","new":"y"}`))
if err == nil {
t.Fatal("expected error for non-existent file")
}
}
+246
View File
@@ -0,0 +1,246 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"os/exec"
"path/filepath"
"strings"
"agentu/pkg/llm"
)
const (
defaultGitLogMaxCount = 20
maxGitLogMaxCount = 100
)
func runGit(ctx context.Context, workingDir string, args ...string) (string, error) {
gitArgs := append([]string{"-C", workingDir}, args...)
cmd := exec.CommandContext(ctx, "git", gitArgs...)
output, err := cmd.CombinedOutput()
text := FormatResult(string(output))
if err != nil {
if strings.Contains(err.Error(), "context deadline exceeded") || strings.Contains(err.Error(), "signal: killed") {
if text != "" {
return text, fmt.Errorf("git command timed out: %w", err)
}
return "", fmt.Errorf("git command timed out: %w", err)
}
if text != "" {
return text, fmt.Errorf("git command failed: %w", err)
}
return "", fmt.Errorf("git command failed: %w", err)
}
return text, nil
}
func gitPathspec(baseDir, path string) (string, error) {
if strings.TrimSpace(path) == "" {
return "", nil
}
resolved, err := resolvePath(baseDir, path)
if err != nil {
return "", err
}
rel, err := filepath.Rel(baseDir, resolved)
if err != nil {
return "", err
}
return filepath.ToSlash(filepath.Clean(rel)), nil
}
func normalizeGitLogMaxCount(value int) int {
if value <= 0 {
return defaultGitLogMaxCount
}
if value > maxGitLogMaxCount {
return maxGitLogMaxCount
}
return value
}
// git_status
type GitStatusTool struct {
workingDir string
}
func NewGitStatusTool(workingDir string) *GitStatusTool {
return &GitStatusTool{workingDir: workingDir}
}
func (t *GitStatusTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "git_status",
Description: "Show changed files using git status --short.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Optional path to limit status to. Relative paths are resolved from the configured agent working directory."}
},
"additionalProperties": false
}`),
},
}
}
type gitStatusArgs struct {
Path string `json:"path"`
}
func (t *GitStatusTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args gitStatusArgs
if len(raw) > 0 {
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
}
gitArgs := []string{"status", "--short"}
if args.Path != "" {
pathspec, err := gitPathspec(t.workingDir, args.Path)
if err != nil {
return "", err
}
gitArgs = append(gitArgs, "--", pathspec)
}
text, err := runGit(ctx, t.workingDir, gitArgs...)
if err != nil {
return text, err
}
if strings.TrimSpace(text) == "" {
return "Working tree clean.", nil
}
return "Changed files:\n" + text, nil
}
// git_diff
type GitDiffTool struct {
workingDir string
}
func NewGitDiffTool(workingDir string) *GitDiffTool {
return &GitDiffTool{workingDir: workingDir}
}
func (t *GitDiffTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "git_diff",
Description: "Show unstaged or staged git diff output, optionally as --stat.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Optional path to limit diff to. Relative paths are resolved from the configured agent working directory."},
"staged": {"type": "boolean", "description": "Show staged diff instead of unstaged."},
"stat": {"type": "boolean", "description": "Show diffstat summary instead of full diff."}
},
"additionalProperties": false
}`),
},
}
}
type gitDiffArgs struct {
Path string `json:"path"`
Staged bool `json:"staged"`
Stat bool `json:"stat"`
}
func (t *GitDiffTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args gitDiffArgs
if len(raw) > 0 {
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
}
gitArgs := []string{"diff"}
if args.Staged {
gitArgs = append(gitArgs, "--cached")
}
if args.Stat {
gitArgs = append(gitArgs, "--stat")
}
if args.Path != "" {
pathspec, err := gitPathspec(t.workingDir, args.Path)
if err != nil {
return "", err
}
gitArgs = append(gitArgs, "--", pathspec)
}
text, err := runGit(ctx, t.workingDir, gitArgs...)
if err != nil {
return text, err
}
if strings.TrimSpace(text) == "" {
if args.Staged {
return "No staged diff.", nil
}
return "No unstaged diff.", nil
}
return text, nil
}
// git_log
type GitLogTool struct {
workingDir string
}
func NewGitLogTool(workingDir string) *GitLogTool {
return &GitLogTool{workingDir: workingDir}
}
func (t *GitLogTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "git_log",
Description: "Show recent git commits with git log --oneline --decorate.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Optional path to limit log to. Relative paths are resolved from the configured agent working directory."},
"max_count": {"type": "integer", "description": "Maximum number of commits to show. Defaults to 20, capped at 100."}
},
"additionalProperties": false
}`),
},
}
}
type gitLogArgs struct {
Path string `json:"path"`
MaxCount int `json:"max_count"`
}
func (t *GitLogTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args gitLogArgs
if len(raw) > 0 {
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
}
maxCount := normalizeGitLogMaxCount(args.MaxCount)
gitArgs := []string{"log", "--oneline", "--decorate", fmt.Sprintf("-n%d", maxCount)}
if args.Path != "" {
pathspec, err := gitPathspec(t.workingDir, args.Path)
if err != nil {
return "", err
}
gitArgs = append(gitArgs, "--", pathspec)
}
text, err := runGit(ctx, t.workingDir, gitArgs...)
if err != nil {
return text, err
}
if strings.TrimSpace(text) == "" {
return "No git history.", nil
}
return text, nil
}
+142
View File
@@ -0,0 +1,142 @@
package tools
import (
"context"
"encoding/json"
"os"
"os/exec"
"path/filepath"
"strings"
"testing"
)
func requireGit(t *testing.T) {
t.Helper()
if _, err := exec.LookPath("git"); err != nil {
t.Skip("git not installed")
}
}
func runGitForTest(t *testing.T, dir string, args ...string) {
t.Helper()
gitArgs := append([]string{"-C", dir}, args...)
cmd := exec.Command("git", gitArgs...)
out, err := cmd.CombinedOutput()
if err != nil {
t.Fatalf("git %v failed: %v\n%s", args, err, out)
}
}
func initTestRepo(t *testing.T) string {
t.Helper()
dir := t.TempDir()
runGitForTest(t, dir, "init")
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "--allow-empty", "-m", "initial")
return dir
}
func TestGitStatusToolShowsChangedFiles(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
path := filepath.Join(dir, "tracked.txt")
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
runGitForTest(t, dir, "add", "tracked.txt")
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
t.Fatal(err)
}
tool := NewGitStatusTool(dir)
out, err := tool.Execute(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "Changed files:") {
t.Fatalf("output missing 'Changed files:': %q", out)
}
if !strings.Contains(out, "M tracked.txt") {
t.Fatalf("output missing 'M tracked.txt': %q", out)
}
}
func TestGitDiffToolShowsUnstagedDiff(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
path := filepath.Join(dir, "tracked.txt")
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
runGitForTest(t, dir, "add", "tracked.txt")
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
t.Fatal(err)
}
tool := NewGitDiffTool(dir)
out, err := tool.Execute(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "diff --git") {
t.Fatalf("output missing 'diff --git': %q", out)
}
}
func TestGitDiffToolShowsStat(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
path := filepath.Join(dir, "tracked.txt")
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
runGitForTest(t, dir, "add", "tracked.txt")
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
t.Fatal(err)
}
tool := NewGitDiffTool(dir)
out, err := tool.Execute(context.Background(), json.RawMessage(`{"stat":true}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "tracked.txt") {
t.Fatalf("output missing 'tracked.txt': %q", out)
}
}
func TestGitLogToolShowsCommits(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
tool := NewGitLogTool(dir)
out, err := tool.Execute(context.Background(), json.RawMessage(`{"max_count":1}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "initial") {
t.Fatalf("output missing 'initial': %q", out)
}
}
func TestGitToolsRejectEscapingPath(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
tool := NewGitStatusTool(dir)
_, err := tool.Execute(context.Background(), json.RawMessage(`{"path":"../outside"}`))
if err == nil {
t.Fatal("expected error for escaping path")
}
if !strings.Contains(err.Error(), "escapes") {
t.Fatalf("error = %v", err)
}
}
+244
View File
@@ -0,0 +1,244 @@
package tools
import (
"context"
"encoding/json"
"errors"
"fmt"
"io/fs"
"os"
"path/filepath"
"strings"
"agentu/pkg/llm"
)
const defaultOutlineMaxFiles = 200
// ────────────────────────────────────────────────────────────────────────────
// ProjectOutlineTool — gitignore-aware file tree + symbol outline
// ────────────────────────────────────────────────────────────────────────────
type ProjectOutlineTool struct {
workingDir string
}
func NewProjectOutlineTool(workingDir string) *ProjectOutlineTool {
return &ProjectOutlineTool{workingDir: workingDir}
}
func (t *ProjectOutlineTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "project_outline",
Description: "Get the project file tree and code symbol outline in one call. Respects .gitignore. Use this first to understand project structure before reading individual files.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Subdirectory to outline. Defaults to the project root."},
"max_files": {"type": "integer", "description": "Maximum files to include in the outline. Defaults to 200."}
},
"additionalProperties": false
}`),
},
}
}
type outlineArgs struct {
Path string `json:"path"`
MaxFiles int `json:"max_files"`
}
func (t *ProjectOutlineTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args outlineArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
root, err := resolvePath(t.workingDir, args.Path)
if err != nil {
return "", err
}
limit := args.MaxFiles
if limit <= 0 {
limit = defaultOutlineMaxFiles
}
// Parse .gitignore files from the root upward.
ignorer := newGitignoreMatcher(root)
var out strings.Builder
out.WriteString(root)
out.WriteString("\n")
count := 0
truncated := false
err = filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
if err != nil {
return nil
}
if ctx.Err() != nil {
return ctx.Err()
}
// Skip always-ignored directories.
if d.IsDir() {
name := d.Name()
// Always skip VCS and common build-artifact directories.
switch name {
case ".git", ".hg", ".svn", "vendor", "node_modules",
"dist", "build", "__pycache__", "target":
return filepath.SkipDir
}
// Skip hidden directories other than .github and similar.
if strings.HasPrefix(name, ".") && name != ".github" && name != ".vscode" && name != ".devcontainer" {
return filepath.SkipDir
}
}
rel, _ := filepath.Rel(root, path)
if rel == "." {
return nil
}
// Check gitignore.
if ignorer.shouldIgnore(rel, d.IsDir()) {
if d.IsDir() {
return filepath.SkipDir
}
return nil
}
if count >= limit {
truncated = true
return filepath.SkipAll
}
// Build indented tree entry.
depth := strings.Count(rel, string(filepath.Separator))
indent := strings.Repeat(" ", depth)
name := d.Name()
if d.IsDir() {
out.WriteString(fmt.Sprintf("%s%s/\n", indent, name))
return nil
}
count++
out.WriteString(indent + name)
// Extract symbols for source files.
ext := filepath.Ext(name)
if sourceFileExtensions[ext] {
data, readErr := os.ReadFile(path)
if readErr == nil && len(data) <= 256*1024 {
symbols, symErr := extractSymbols(path, data)
if symErr == nil && len(symbols) > 0 {
out.WriteString(" — ")
out.WriteString(compactSymbolSummary(symbols))
}
}
}
out.WriteString("\n")
return nil
})
if err != nil && !errors.Is(err, filepath.SkipAll) {
return "", err
}
if truncated {
out.WriteString(fmt.Sprintf("\n... (%d files shown, more available; use file_list or code_symbols for details)\n", count))
}
return FormatResult(out.String()), nil
}
// ────────────────────────────────────────────────────────────────────────────
// Compact symbol summary for inline display in the file tree
// ────────────────────────────────────────────────────────────────────────────
func compactSymbolSummary(symbols []symbolInfo) string {
var parts []string
// Group by kind.
funcs := []string{}
methods := []string{}
types := []string{}
for _, s := range symbols {
switch s.Kind {
case "function":
funcs = append(funcs, s.Name)
case "method":
methods = append(methods, s.Name)
case "class", "interface", "struct", "enum", "trait", "alias", "func", "type", "impl":
types = append(types, s.Name)
}
// Skip imports and package for compact view.
}
if len(funcs) > 0 {
parts = append(parts, fmt.Sprintf("fn: %s", strings.Join(funcs, ", ")))
}
if len(methods) > 0 {
parts = append(parts, fmt.Sprintf("method: %s", strings.Join(methods, ", ")))
}
if len(types) > 0 {
parts = append(parts, fmt.Sprintf("type: %s", strings.Join(types, ", ")))
}
return strings.Join(parts, "; ")
}
// ────────────────────────────────────────────────────────────────────────────
// Simple .gitignore matcher
// ────────────────────────────────────────────────────────────────────────────
type gitignoreMatcher struct {
patterns []string
}
func newGitignoreMatcher(root string) *gitignoreMatcher {
m := &gitignoreMatcher{}
// Walk up from root, collecting .gitignore patterns.
for dir := root; dir != "" && dir != string(filepath.Separator); dir = filepath.Dir(dir) {
giPath := filepath.Join(dir, ".gitignore")
data, err := os.ReadFile(giPath)
if err != nil {
continue
}
for _, line := range strings.Split(string(data), "\n") {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
continue
}
// Remove leading / for anchored patterns; we handle all as relative.
line = strings.TrimPrefix(line, "/")
// Remove trailing / for directory-only patterns.
line = strings.TrimSuffix(line, "/")
m.patterns = append(m.patterns, line)
}
}
return m
}
func (m *gitignoreMatcher) shouldIgnore(path string, isDir bool) bool {
name := filepath.Base(path)
for _, pat := range m.patterns {
// Simple match against full relative path or just the basename.
matched, _ := filepath.Match(pat, path)
if matched {
return true
}
matched, _ = filepath.Match(pat, name)
if matched {
return true
}
}
return false
}
+430
View File
@@ -0,0 +1,430 @@
package tools
import (
"fmt"
"os"
"path/filepath"
"strings"
"agentu/internal/safeio"
)
// ────────────────────────────────────────────────────────────────────────────
// V4 patch format data structures
// ────────────────────────────────────────────────────────────────────────────
type patchAction string
const (
actionAdd patchAction = "add"
actionUpdate patchAction = "update"
actionDelete patchAction = "delete"
)
type patchFile struct {
action patchAction
path string
newContent string // for "add" — full file content
lines []string // for "update" — lines to match/replace; for "delete" — unused
}
type parsedPatch struct {
files []patchFile
}
// ────────────────────────────────────────────────────────────────────────────
// Parser
// ────────────────────────────────────────────────────────────────────────────
// parsePatch parses a V4-format patch string.
//
// Format:
//
// *** Begin Patch
// *** Add File: <path>
// <file content>
// *** Update File: <path>
// @@
// <context line>
// @@
// -<old line>
// +<new line>
// *** Delete File: <path>
// *** End Patch
func parsePatch(patch string) (*parsedPatch, error) {
patch = strings.TrimSpace(patch)
if patch == "" {
return nil, fmt.Errorf("patch is empty")
}
lines := strings.Split(patch, "\n")
if len(lines) == 0 {
return nil, fmt.Errorf("patch is empty")
}
if strings.TrimSpace(lines[0]) != "*** Begin Patch" {
return nil, fmt.Errorf("patch must start with *** Begin Patch")
}
var files []patchFile
seenPaths := make(map[string]struct{})
seenEnd := false
i := 1
for i < len(lines) {
line := strings.TrimSpace(lines[i])
switch {
case line == "*** End Patch":
seenEnd = true
i = len(lines)
case strings.HasPrefix(line, "*** Add File:"):
path := strings.TrimSpace(strings.TrimPrefix(line, "*** Add File:"))
if err := validatePatchPath(path, seenPaths); err != nil {
return nil, err
}
i++
var contentLines []string
for i < len(lines) {
next := strings.TrimSpace(lines[i])
if strings.HasPrefix(next, "***") {
break
}
content := lines[i]
if strings.HasPrefix(content, "+") {
content = strings.TrimPrefix(content, "+")
}
contentLines = append(contentLines, content)
i++
}
files = append(files, patchFile{
action: actionAdd,
path: path,
newContent: strings.Join(contentLines, "\n") + "\n",
})
case strings.HasPrefix(line, "*** Delete File:"):
path := strings.TrimSpace(strings.TrimPrefix(line, "*** Delete File:"))
if err := validatePatchPath(path, seenPaths); err != nil {
return nil, err
}
files = append(files, patchFile{
action: actionDelete,
path: path,
})
i++
case strings.HasPrefix(line, "*** Update File:"):
path := strings.TrimSpace(strings.TrimPrefix(line, "*** Update File:"))
if err := validatePatchPath(path, seenPaths); err != nil {
return nil, err
}
i++
var hunkLines []string
for i < len(lines) {
next := strings.TrimSpace(lines[i])
if strings.HasPrefix(next, "***") {
break
}
hunkLines = append(hunkLines, lines[i])
i++
}
files = append(files, patchFile{
action: actionUpdate,
path: path,
lines: hunkLines,
})
default:
i++
}
}
if len(files) == 0 {
return nil, fmt.Errorf("patch contains no file operations")
}
if !seenEnd {
return nil, fmt.Errorf("patch must end with *** End Patch")
}
return &parsedPatch{files: files}, nil
}
func validatePatchPath(path string, seen map[string]struct{}) error {
if path == "" {
return fmt.Errorf("patch file path is required")
}
if _, ok := seen[path]; ok {
return fmt.Errorf("patch contains duplicate file path: %s", path)
}
seen[path] = struct{}{}
return nil
}
// ────────────────────────────────────────────────────────────────────────────
// Applier
// ────────────────────────────────────────────────────────────────────────────
// applyPatchFile applies a single patch file operation.
func applyPatchFile(baseDir string, pf patchFile) (string, error) {
switch pf.action {
case actionAdd:
return applyAdd(baseDir, pf)
case actionUpdate:
return applyUpdate(baseDir, pf)
case actionDelete:
return applyDelete(baseDir, pf)
default:
return "", fmt.Errorf("unknown patch action: %s", pf.action)
}
}
func applyAdd(baseDir string, pf patchFile) (string, error) {
path, err := resolvePathForWrite(baseDir, pf.path)
if err != nil {
return "", err
}
if _, err := os.Stat(path); err == nil {
return "", fmt.Errorf("cannot add existing file: %s", pf.path)
} else if !os.IsNotExist(err) {
return "", fmt.Errorf("stat new file %s: %w", pf.path, err)
}
if err := safeio.AtomicWriteFile(path, []byte(pf.newContent), 0o644); err != nil {
return "", fmt.Errorf("write new file %s: %w", pf.path, err)
}
return fmt.Sprintf("+ %s", pf.path), nil
}
func applyDelete(baseDir string, pf patchFile) (string, error) {
path, err := resolvePath(baseDir, pf.path)
if err != nil {
return "", err
}
if info, err := os.Stat(path); err != nil {
if os.IsNotExist(err) {
return "", fmt.Errorf("cannot delete non-existent file: %s", pf.path)
}
return "", err
} else if info.IsDir() {
return "", fmt.Errorf("cannot delete directory: %s", pf.path)
}
if err := os.Remove(path); err != nil {
return "", fmt.Errorf("delete file %s: %w", pf.path, err)
}
// Try to remove empty parent directories.
dir := filepath.Dir(path)
for dir != baseDir && dir != "." {
if entries, err := os.ReadDir(dir); err == nil && len(entries) == 0 {
os.Remove(dir)
dir = filepath.Dir(dir)
} else {
break
}
}
return fmt.Sprintf("- %s", pf.path), nil
}
func applyUpdate(baseDir string, pf patchFile) (string, error) {
path, err := resolvePath(baseDir, pf.path)
if err != nil {
return "", fmt.Errorf("cannot update non-existent file: %s: %w", pf.path, err)
}
existing, err := os.ReadFile(path)
if err != nil {
return "", fmt.Errorf("cannot update non-existent file: %s", pf.path)
}
info, err := os.Stat(path)
if err != nil {
return "", fmt.Errorf("stat file %s: %w", pf.path, err)
}
fileLines := strings.Split(string(existing), "\n")
// Parse the hunk lines into: context blocks + old lines + new lines.
result, err := applyHunkLines(fileLines, pf.lines, pf.path)
if err != nil {
return "", err
}
newContent := strings.Join(result, "\n")
if err := safeio.AtomicWriteFile(path, []byte(newContent), info.Mode().Perm()); err != nil {
return "", fmt.Errorf("write updated file %s: %w", pf.path, err)
}
return fmt.Sprintf("~ %s", pf.path), nil
}
// applyHunkLines applies a sequence of hunk lines to the file content.
// The hunk lines are a mix of:
//
// @@ — hunk separator (a context line block starts)
// <plain line> — context line (must appear in both old and new)
// -<line> — line to remove
// +<line> — line to add
//
// Hunks are matched greedily: for each group of contiguous context + - + + lines,
// find the first position in the file where context + - lines match,
// replace with context + + lines.
func applyHunkLines(fileLines []string, hunkLines []string, path string) ([]string, error) {
if len(hunkLines) == 0 {
return fileLines, nil
}
// Split the hunk lines into separate hunks at @@ markers.
type hunk struct {
lines []string
}
var hunks []hunk
var current hunk
for _, hl := range hunkLines {
trimmed := strings.TrimSpace(hl)
if strings.HasPrefix(trimmed, "@@") {
if len(current.lines) > 0 {
hunks = append(hunks, current)
current = hunk{}
}
continue
}
current.lines = append(current.lines, hl)
}
if len(current.lines) > 0 {
hunks = append(hunks, current)
}
if len(hunks) == 0 {
return fileLines, nil
}
result := make([]string, len(fileLines))
copy(result, fileLines)
offset := 0 // tracks net line count change from previous hunks
for hi, h := range hunks {
// Separate the hunk lines into a search pattern and replacement.
searchLines, replaceLines := buildHunkSearchReplace(h.lines)
if len(searchLines) == 0 && len(replaceLines) == 0 {
continue
}
matchStart := findHunkMatch(result, searchLines, offset)
if matchStart < 0 {
return nil, fmt.Errorf("hunk context not found in %s", path)
}
// Replace: remove searchLines, insert replaceLines.
matchEnd := matchStart + len(searchLines)
newResult := make([]string, 0, len(result)-len(searchLines)+len(replaceLines))
newResult = append(newResult, result[:matchStart]...)
newResult = append(newResult, replaceLines...)
newResult = append(newResult, result[matchEnd:]...)
result = newResult
offset = matchStart + len(replaceLines)
// Multiple matches check for this hunk only.
secondMatch := findHunkMatch(result, searchLines, offset)
if secondMatch >= 0 {
return nil, fmt.Errorf("hunk context matched multiple times in %s", path)
}
_ = hi
}
return result, nil
}
// buildHunkSearchReplace splits a hunk's lines into the search pattern
// (context + old lines) and replacement (context + new lines).
func buildHunkSearchReplace(lines []string) (search, replace []string) {
for _, l := range lines {
if strings.HasPrefix(l, "-") {
search = append(search, l[1:])
} else if strings.HasPrefix(l, "+") {
replace = append(replace, l[1:])
} else {
contextLine := strings.TrimPrefix(l, " ")
search = append(search, contextLine)
replace = append(replace, contextLine)
}
}
return
}
// findHunkMatch finds the first position >= offset in fileLines where
// searchLines matches exactly (line-by-line comparison).
func findHunkMatch(fileLines []string, searchLines []string, offset int) int {
for i := offset; i <= len(fileLines)-len(searchLines); i++ {
match := true
for j, sl := range searchLines {
if fileLines[i+j] != sl {
match = false
break
}
}
if match {
return i
}
}
return -1
}
// ────────────────────────────────────────────────────────────────────────────
// Preview
// ────────────────────────────────────────────────────────────────────────────
// summarizePatch returns a human-readable summary of the patch.
func summarizePatch(patch string) (string, error) {
pp, err := parsePatch(patch)
if err != nil {
return "", err
}
var b strings.Builder
added, updated, deleted := 0, 0, 0
for _, pf := range pp.files {
switch pf.action {
case actionAdd:
added++
fmt.Fprintf(&b, "+ add %s\n", pf.path)
case actionUpdate:
updated++
fmt.Fprintf(&b, "~ update %s\n", pf.path)
case actionDelete:
deleted++
fmt.Fprintf(&b, "- delete %s\n", pf.path)
}
}
summary := fmt.Sprintf("patch: %d file(s) changed", len(pp.files))
if added > 0 {
summary += fmt.Sprintf(" (%d added)", added)
}
if updated > 0 {
summary += fmt.Sprintf(" (%d updated)", updated)
}
if deleted > 0 {
summary += fmt.Sprintf(" (%d deleted)", deleted)
}
return summary + "\n" + b.String(), nil
}
// extractTargetPaths returns all file paths referenced in the patch.
func extractTargetPaths(baseDir, patch string) ([]string, error) {
pp, err := parsePatch(patch)
if err != nil {
return nil, err
}
var paths []string
for _, pf := range pp.files {
// Use resolvePathForWrite for all files — TargetPaths needs paths
// even for files that don't exist yet (rollback snapshots).
resolved, err := resolvePathForWrite(baseDir, pf.path)
if err != nil {
return nil, fmt.Errorf("file %s: %w", pf.path, err)
}
paths = append(paths, resolved)
}
return paths, nil
}
+82 -2
View File
@@ -1,15 +1,95 @@
package tools
import (
"fmt"
"os"
"path/filepath"
"strings"
)
func resolvePath(baseDir, path string) (string, error) {
resolved, canonicalBase, err := resolveLexicalPath(baseDir, path)
if err != nil {
return "", err
}
canonical, err := filepath.EvalSymlinks(resolved)
if err != nil {
return "", err
}
if !isWithinDir(canonical, canonicalBase) {
return "", fmt.Errorf("path escapes the project directory through a symbolic link: %s", path)
}
return canonical, nil
}
// resolvePathForWrite resolves a path that may not exist yet while preventing
// existing parent-directory symlinks from escaping the project directory.
func resolvePathForWrite(baseDir, path string) (string, error) {
resolved, canonicalBase, err := resolveLexicalPath(baseDir, path)
if err != nil {
return "", err
}
existing := resolved
var suffix []string
for {
_, err := os.Lstat(existing)
if err == nil {
break
}
if !os.IsNotExist(err) {
return "", err
}
parent := filepath.Dir(existing)
if parent == existing {
return "", err
}
suffix = append(suffix, filepath.Base(existing))
existing = parent
}
canonical, err := filepath.EvalSymlinks(existing)
if err != nil {
return "", err
}
for i := len(suffix) - 1; i >= 0; i-- {
canonical = filepath.Join(canonical, suffix[i])
}
if !isWithinDir(canonical, canonicalBase) {
return "", fmt.Errorf("path escapes the project directory through a symbolic link: %s", path)
}
return canonical, nil
}
func resolveLexicalPath(baseDir, path string) (string, string, error) {
if path == "" {
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))
absBase, err := filepath.Abs(baseDir)
if err != nil {
return "", "", err
}
canonicalBase, err := filepath.EvalSymlinks(absBase)
if err != nil {
return "", "", err
}
resolved := filepath.Clean(filepath.Join(absBase, path))
if !isWithinDir(resolved, absBase) {
return "", "", fmt.Errorf("path escapes the project directory: %s", path)
}
return resolved, canonicalBase, 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
}
+117
View File
@@ -0,0 +1,117 @@
package tools
import (
"crypto/rand"
"encoding/hex"
"fmt"
"io/fs"
"os"
"sync"
"time"
"agentu/internal/safeio"
)
// RollbackSnapshot captures the state of a file before a mutation.
type RollbackSnapshot struct {
ID string
Tool string
Path string
ResolvedPath string
Existed bool
Data []byte
Mode fs.FileMode
CreatedAt time.Time
}
// RollbackStore manages in-memory rollback snapshots for the current process.
type RollbackStore struct {
mu sync.Mutex
snapshots []RollbackSnapshot
}
// NewRollbackStore creates an empty rollback store.
func NewRollbackStore() *RollbackStore {
return &RollbackStore{}
}
// generateSnapshotID creates a rollback ID: "rb_" + 12 lowercase hex chars.
func generateSnapshotID() (string, error) {
b := make([]byte, 6)
if _, err := rand.Read(b); err != nil {
return "", fmt.Errorf("generate snapshot id: %w", err)
}
return "rb_" + hex.EncodeToString(b), nil
}
// Capture stores the current state of the file at resolvedPath before a tool
// modifies it. Returns the snapshot ID.
func (s *RollbackStore) Capture(tool string, displayPath string, resolvedPath string) (string, error) {
id, err := generateSnapshotID()
if err != nil {
return "", err
}
snap := RollbackSnapshot{
ID: id,
Tool: tool,
Path: displayPath,
ResolvedPath: resolvedPath,
CreatedAt: time.Now().UTC(),
}
info, err := os.Stat(resolvedPath)
if err != nil {
if os.IsNotExist(err) {
snap.Existed = false
} else {
return "", fmt.Errorf("stat file: %w", err)
}
} else {
snap.Existed = true
snap.Mode = info.Mode().Perm()
data, err := os.ReadFile(resolvedPath)
if err != nil {
return "", fmt.Errorf("read file for snapshot: %w", err)
}
snap.Data = data
}
s.mu.Lock()
s.snapshots = append(s.snapshots, snap)
s.mu.Unlock()
return id, nil
}
// Restore rolls back a file to the state captured in the given snapshot.
func (s *RollbackStore) Restore(id string) (RollbackSnapshot, error) {
s.mu.Lock()
defer s.mu.Unlock()
for i := len(s.snapshots) - 1; i >= 0; i-- {
if s.snapshots[i].ID == id {
snap := s.snapshots[i]
if snap.Existed {
if err := safeio.AtomicWriteFile(snap.ResolvedPath, snap.Data, snap.Mode); err != nil {
return RollbackSnapshot{}, fmt.Errorf("restore file: %w", err)
}
} else {
if err := os.Remove(snap.ResolvedPath); err != nil && !os.IsNotExist(err) {
return RollbackSnapshot{}, fmt.Errorf("remove created file: %w", err)
}
}
s.snapshots = append(s.snapshots[:i], s.snapshots[i+1:]...)
return snap, nil
}
}
return RollbackSnapshot{}, fmt.Errorf("rollback snapshot %s not found", id)
}
// LastID returns the most recent snapshot ID, or "" if empty.
func (s *RollbackStore) LastID() string {
s.mu.Lock()
defer s.mu.Unlock()
if len(s.snapshots) == 0 {
return ""
}
return s.snapshots[len(s.snapshots)-1].ID
}
+53 -6
View File
@@ -1,6 +1,7 @@
package tools
import (
"bytes"
"context"
"encoding/json"
"errors"
@@ -8,8 +9,9 @@ import (
"os"
"os/exec"
"strings"
"sync"
"agentu/internal/llm"
"agentu/pkg/llm"
)
type ShellRunTool struct {
@@ -19,6 +21,18 @@ type ShellRunTool struct {
func NewShellRunTool(workingDir string) *ShellRunTool {
return &ShellRunTool{workingDir: workingDir}
}
func (*ShellRunTool) Mutates() bool { return true }
func (t *ShellRunTool) Preview(ctx context.Context, raw json.RawMessage) (string, error) {
var args shellRunArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
if strings.TrimSpace(args.Command) == "" {
return "", errors.New("command is required")
}
return "$ " + strings.TrimSpace(args.Command), ctx.Err()
}
func (t *ShellRunTool) Definition() llm.Tool {
return llm.Tool{
@@ -45,6 +59,10 @@ type shellRunArgs struct {
}
func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
return t.ExecuteStream(ctx, raw, nil)
}
func (t *ShellRunTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
var args shellRunArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
@@ -65,16 +83,45 @@ func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string
cmd := exec.CommandContext(ctx, shell, "-lc", args.Command)
cmd.Dir = workdir
output, err := cmd.CombinedOutput()
text := FormatResult(string(output))
var mu sync.Mutex
var combined bytes.Buffer
writer := &shellStreamWriter{mu: &mu, buf: &combined, emit: emit}
cmd.Stdout = writer
cmd.Stderr = writer
runErr := cmd.Run()
text := FormatResult(combined.String())
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
return text, fmt.Errorf("command timed out: %w", ctx.Err())
}
if err != nil {
if runErr != nil {
if text != "" {
return text, fmt.Errorf("command failed: %w", err)
return text, fmt.Errorf("command failed: %w", runErr)
}
return "", fmt.Errorf("command failed: %w", err)
return "", fmt.Errorf("command failed: %w", runErr)
}
return text, nil
}
// shellStreamWriter captures command output for the final result while
// forwarding every write to the streaming callback as it happens.
type shellStreamWriter struct {
mu *sync.Mutex
buf *bytes.Buffer
emit func(string) error
}
func (w *shellStreamWriter) Write(p []byte) (int, error) {
w.mu.Lock()
n, writeErr := w.buf.Write(p)
w.mu.Unlock()
if writeErr != nil {
return n, writeErr
}
if w.emit != nil {
if err := w.emit(string(p)); err != nil {
return 0, err
}
}
return len(p), nil
}
+71
View File
@@ -0,0 +1,71 @@
package tools
import (
"context"
"encoding/json"
"strings"
"sync"
"testing"
)
func TestShellRunStreamsOutputWhileRunning(t *testing.T) {
tool := NewShellRunTool(t.TempDir())
var mu sync.Mutex
var chunks []string
output, err := tool.ExecuteStream(context.Background(), json.RawMessage(`{"command":"printf 'a\\n'; sleep 0.05; printf 'b\\n'"}`), func(chunk string) error {
mu.Lock()
chunks = append(chunks, chunk)
mu.Unlock()
return nil
})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(output, "a") || !strings.Contains(output, "b") {
t.Fatalf("output = %q", output)
}
mu.Lock()
joined := strings.Join(chunks, "")
mu.Unlock()
if len(chunks) == 0 {
t.Fatal("expected at least one streamed chunk")
}
if !strings.Contains(joined, "a") || !strings.Contains(joined, "b") {
t.Fatalf("streamed chunks = %q", joined)
}
}
func TestShellRunExecuteCollectsFullOutput(t *testing.T) {
tool := NewShellRunTool(t.TempDir())
output, err := tool.Execute(context.Background(), json.RawMessage(`{"command":"printf 'hello world'"}`))
if err != nil {
t.Fatal(err)
}
if strings.TrimSpace(output) != "hello world" {
t.Fatalf("output = %q", output)
}
}
func TestShellRunStreamsStderrAndStdout(t *testing.T) {
tool := NewShellRunTool(t.TempDir())
var mu sync.Mutex
var chunks []string
output, err := tool.ExecuteStream(context.Background(), json.RawMessage(`{"command":"echo out; echo err >&2"}`), func(chunk string) error {
mu.Lock()
chunks = append(chunks, chunk)
mu.Unlock()
return nil
})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(output, "out") || !strings.Contains(output, "err") {
t.Fatalf("output = %q", output)
}
mu.Lock()
joined := strings.Join(chunks, "")
mu.Unlock()
if !strings.Contains(joined, "out") || !strings.Contains(joined, "err") {
t.Fatalf("streamed chunks = %q", joined)
}
}
+530
View File
@@ -0,0 +1,530 @@
package tools
import (
"bufio"
"bytes"
"fmt"
"go/ast"
"go/parser"
"go/printer"
"go/token"
"regexp"
"sort"
"strings"
)
// symbolInfo holds extracted symbol metadata.
type symbolInfo struct {
Name string
Kind string // "import", "function", "method", "class", "type", "interface", "trait", "enum", "struct"
StartLine int
}
// sourceFileExtensions lists extensions that we can extract symbols from.
var sourceFileExtensions = map[string]bool{
".go": true,
".ts": true,
".tsx": true,
".js": true,
".jsx": true,
".mjs": true,
".py": true,
".rs": true,
}
// extractSymbols extracts symbols from a source file using language-specific
// parsing (go/ast for Go, regex for other languages).
func extractSymbols(path string, data []byte) ([]symbolInfo, error) {
ext := ""
if idx := strings.LastIndex(path, "."); idx >= 0 {
ext = path[idx:]
}
switch ext {
case ".go":
return extractGoSymbols(path, data)
case ".py":
return extractPythonsymbols(data)
case ".ts", ".tsx", ".js", ".jsx", ".mjs":
return extractTSSymbols(data)
case ".rs":
return extractRustSymbols(data)
default:
return nil, nil
}
}
// ────────────────────────────────────────────────────────────────────────────
// Go symbol extraction (using go/ast)
// ────────────────────────────────────────────────────────────────────────────
func extractGoSymbols(path string, data []byte) ([]symbolInfo, error) {
fset := token.NewFileSet()
// Parse imports only first for the import list.
importFile, err := parser.ParseFile(fset, path, data, parser.ImportsOnly|parser.ParseComments)
if err != nil {
return nil, err
}
full, err := parser.ParseFile(fset, path, data, 0)
if err != nil {
return nil, err
}
var symbols []symbolInfo
// Package
symbols = append(symbols, symbolInfo{
Name: full.Name.Name,
Kind: "package",
StartLine: fset.Position(full.Name.Pos()).Line,
})
// Imports
for _, imp := range importFile.Imports {
symbols = append(symbols, symbolInfo{
Name: strings.Trim(imp.Path.Value, "\""),
Kind: "import",
StartLine: fset.Position(imp.Pos()).Line,
})
}
// Declarations
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 := goTypeKind(ts.Type)
symbols = append(symbols, symbolInfo{
Name: ts.Name.Name,
Kind: kind,
StartLine: fset.Position(ts.Pos()).Line,
})
}
}
case *ast.FuncDecl:
line := fset.Position(d.Pos()).Line
if d.Recv != nil {
symbols = append(symbols, symbolInfo{
Name: fmt.Sprintf("%s.%s", goReceiverName(d.Recv), d.Name.Name),
Kind: "method",
StartLine: line,
})
} else {
symbols = append(symbols, symbolInfo{
Name: d.Name.Name,
Kind: "function",
StartLine: line,
})
}
}
}
sortSymbols(symbols)
return symbols, nil
}
func goTypeKind(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 goReceiverName(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()
}
// ────────────────────────────────────────────────────────────────────────────
// Python symbol extraction (regex-based)
// ────────────────────────────────────────────────────────────────────────────
var (
pyFuncRe = regexp.MustCompile(`^\s*def\s+(\w+)\s*\(`)
pyClassRe = regexp.MustCompile(`^\s*class\s+(\w+)\s*[(:]`)
pyImportRe = regexp.MustCompile(`^\s*(?:from\s+\S+\s+)?import\s+`)
pyDecoratorRe = regexp.MustCompile(`^\s*@`)
)
func extractPythonsymbols(data []byte) ([]symbolInfo, error) {
var symbols []symbolInfo
scanner := bufio.NewScanner(bytes.NewReader(data))
lineNum := 0
inDecorator := false
for scanner.Scan() {
lineNum++
line := scanner.Text()
// Track decorator stacks
if pyDecoratorRe.MatchString(line) {
inDecorator = true
continue
}
// Import lines
if pyImportRe.MatchString(line) {
// Extract the module name for a simple "import foo" or "from foo import bar"
importPath := strings.TrimSpace(line)
symbols = append(symbols, symbolInfo{
Name: importPath,
Kind: "import",
StartLine: lineNum,
})
inDecorator = false
continue
}
matches := pyClassRe.FindStringSubmatch(line)
if matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "class",
StartLine: lineNum,
})
inDecorator = false
continue
}
matches = pyFuncRe.FindStringSubmatch(line)
if matches != nil {
kind := "function"
if inDecorator {
// Could be a decorated function or method — still mark as function
// (class-method detection requires indentation tracking)
}
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: kind,
StartLine: lineNum,
})
inDecorator = false
continue
}
inDecorator = false
}
sortSymbols(symbols)
return symbols, scanner.Err()
}
// ────────────────────────────────────────────────────────────────────────────
// TypeScript / JavaScript symbol extraction (regex-based)
// ────────────────────────────────────────────────────────────────────────────
var (
tsFuncRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?(?:async\s+)?function\s+(\w+)\s*[<(]`)
tsArrowRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s*)?\s*\(`)
tsClassRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?class\s+(\w+)`)
tsInterfaceRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?interface\s+(\w+)`)
tsTypeRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?type\s+(\w+)\s*=`)
tsEnumRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?enum\s+(\w+)`)
tsImportRe = regexp.MustCompile(`^\s*import\s+`)
tsMethodRe = regexp.MustCompile(`^\s*(?:public\s+|private\s+|protected\s+|static\s+|async\s+)*(\w+)\s*[<(][^)]*\)\s*(?::\s*\S+\s*)?\{`)
)
func extractTSSymbols(data []byte) ([]symbolInfo, error) {
var symbols []symbolInfo
scanner := bufio.NewScanner(bytes.NewReader(data))
lineNum := 0
// Track brace depth to detect class body methods.
braceDepth := 0
inClass := false
for scanner.Scan() {
lineNum++
line := scanner.Text()
trimmed := strings.TrimSpace(line)
// Count braces on this line (applied at end of iteration).
opens := strings.Count(line, "{")
closes := strings.Count(line, "}")
// Import lines — skip everything else, just track braces.
if tsImportRe.MatchString(line) {
importPath := trimmed
symbols = append(symbols, symbolInfo{
Name: importPath,
Kind: "import",
StartLine: lineNum,
})
goto applyBraces
}
// Skip empty and comment-only lines.
if trimmed == "" || strings.HasPrefix(trimmed, "//") || strings.HasPrefix(trimmed, "/*") {
goto applyBraces
}
// Interface
if matches := tsInterfaceRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "interface",
StartLine: lineNum,
})
goto applyBraces
}
// Type alias
if matches := tsTypeRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "type",
StartLine: lineNum,
})
goto applyBraces
}
// Enum
if matches := tsEnumRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "enum",
StartLine: lineNum,
})
goto applyBraces
}
// Class
if matches := tsClassRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "class",
StartLine: lineNum,
})
// Set inClass BEFORE applying braces for this line, so that
// methods on the same line (unlikely) or next line are detected.
if opens > 0 {
inClass = true
}
goto applyBraces
}
// Method inside class body (at depth 1 brace level).
if inClass && braceDepth == 1 {
if matches := tsMethodRe.FindStringSubmatch(line); matches != nil {
name := matches[1]
if !isKeyword(name) {
symbols = append(symbols, symbolInfo{
Name: name,
Kind: "method",
StartLine: lineNum,
})
}
}
}
// Function declaration (top-level only, or nested function).
if !inClass || braceDepth > 1 {
if matches := tsFuncRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "function",
StartLine: lineNum,
})
}
}
// Arrow function assigned to const/let/var.
if !inClass || braceDepth > 1 {
if matches := tsArrowRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "function",
StartLine: lineNum,
})
}
}
applyBraces:
braceDepth += opens - closes
if braceDepth <= 0 {
braceDepth = 0
inClass = false
}
}
sortSymbols(symbols)
return symbols, scanner.Err()
}
func isKeyword(s string) bool {
switch s {
case "if", "else", "for", "while", "switch", "case", "return", "throw",
"try", "catch", "finally", "new", "delete", "typeof", "instanceof",
"constructor", "get", "set":
return true
}
return false
}
// ────────────────────────────────────────────────────────────────────────────
// Rust symbol extraction (regex-based)
// ────────────────────────────────────────────────────────────────────────────
var (
rsFnRe = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?(?:async\s+)?(?:unsafe\s+)?fn\s+(\w+)\s*[<(]`)
rsStructRe = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?struct\s+(\w+)`)
rsEnumRe = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?enum\s+(\w+)`)
rsTraitRe = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?(?:unsafe\s+)?trait\s+(\w+)`)
rsImplRe = regexp.MustCompile(`^\s*impl\s+`)
rsUseRe = regexp.MustCompile(`^\s*(?:pub\s+)?use\s+`)
)
func extractRustSymbols(data []byte) ([]symbolInfo, error) {
var symbols []symbolInfo
scanner := bufio.NewScanner(bytes.NewReader(data))
lineNum := 0
for scanner.Scan() {
lineNum++
line := scanner.Text()
trimmed := strings.TrimSpace(line)
// Skip empty and comments
if trimmed == "" || strings.HasPrefix(trimmed, "//") || strings.HasPrefix(trimmed, "/*") {
continue
}
// Use
if rsUseRe.MatchString(line) {
symbols = append(symbols, symbolInfo{
Name: trimmed,
Kind: "import",
StartLine: lineNum,
})
continue
}
// Struct
if matches := rsStructRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "struct",
StartLine: lineNum,
})
continue
}
// Enum
if matches := rsEnumRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "enum",
StartLine: lineNum,
})
continue
}
// Trait
if matches := rsTraitRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "trait",
StartLine: lineNum,
})
continue
}
// Impl
if rsImplRe.MatchString(line) {
// Extract type name from "impl Foo" or "impl<T> Foo" or "impl Trait for Foo"
rest := strings.TrimPrefix(trimmed, "impl")
rest = strings.TrimSpace(rest)
// Remove generics
if idx := strings.Index(rest, "<"); idx >= 0 {
depth := 1
for i := idx + 1; i < len(rest); i++ {
if rest[i] == '<' {
depth++
} else if rest[i] == '>' {
depth--
if depth == 0 {
rest = strings.TrimSpace(rest[:idx] + rest[i+1:])
break
}
}
}
}
// Handle "for Type"
if idx := strings.Index(rest, " for "); idx >= 0 {
rest = strings.TrimSpace(rest[idx+5:])
}
name := strings.Fields(rest)[0]
symbols = append(symbols, symbolInfo{
Name: name,
Kind: "impl",
StartLine: lineNum,
})
continue
}
// Function
if matches := rsFnRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "function",
StartLine: lineNum,
})
}
}
sortSymbols(symbols)
return symbols, scanner.Err()
}
// ────────────────────────────────────────────────────────────────────────────
// Helpers
// ────────────────────────────────────────────────────────────────────────────
func sortSymbols(symbols []symbolInfo) {
sort.Slice(symbols, func(i, j int) bool {
return symbols[i].StartLine < symbols[j].StartLine
})
}
// groupSymbols groups symbols by kind and returns formatted strings.
func groupSymbols(symbols []symbolInfo) map[string][]string {
groups := make(map[string][]string)
for _, s := range symbols {
kind := s.Kind
switch kind {
case "package":
groups["package"] = append(groups["package"], s.Name)
case "import":
groups["imports"] = append(groups["imports"], s.Name)
case "function":
groups["functions"] = append(groups["functions"], s.Name)
case "method":
groups["methods"] = append(groups["methods"], s.Name)
case "class", "interface", "struct", "enum", "trait", "alias", "func":
// All structural types go under "types" for consistency.
groups["types"] = append(groups["types"], fmt.Sprintf("%s %s", s.Name, kind))
case "type":
// TS type alias
groups["types"] = append(groups["types"], fmt.Sprintf("%s type", s.Name))
case "impl":
groups["types"] = append(groups["types"], fmt.Sprintf("impl %s", s.Name))
}
}
return groups
}
+173
View File
@@ -0,0 +1,173 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"strings"
"sync"
"agentu/pkg/llm"
)
// ────────────────────────────────────────────────────────────────────────────
// TodoTool — in-memory task tracking
// ────────────────────────────────────────────────────────────────────────────
type todoItem struct {
Content string `json:"content"`
Status string `json:"status"` // "pending", "in_progress", "completed"
}
type TodoTool struct {
mu sync.Mutex
items []todoItem
}
func NewTodoTool() *TodoTool {
return &TodoTool{}
}
func (t *TodoTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "todo",
Description: "Manage a task list for the current conversation. Create, update, and track multi-step work. Use this to plan complex tasks and show progress.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"action": {"type": "string", "enum": ["list", "create", "update", "start", "complete", "clear"], "description": "Operation to perform."},
"items": {"type": "array", "items": {"type": "string"}, "description": "For 'create': list of task descriptions to add."},
"task": {"type": "string", "description": "For 'update'/'start'/'complete': the exact task content string to target."},
"status": {"type": "string", "enum": ["pending", "in_progress", "completed"], "description": "For 'update': new status for the task."}
},
"required": ["action"],
"additionalProperties": false
}`),
},
}
}
type todoArgs struct {
Action string `json:"action"`
Items []string `json:"items"`
Task string `json:"task"`
Status string `json:"status"`
}
func (t *TodoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args todoArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
t.mu.Lock()
defer t.mu.Unlock()
if err := ctx.Err(); err != nil {
return "", err
}
switch args.Action {
case "list":
return t.listTasks(), nil
case "create":
return t.createTasks(args.Items)
case "start":
return t.startTask(args.Task)
case "complete":
return t.completeTask(args.Task)
case "update":
return t.updateTask(args.Task, args.Status)
case "clear":
t.items = nil
return "Tasks cleared.", nil
default:
return "", fmt.Errorf("unknown todo action: %s", args.Action)
}
}
func (t *TodoTool) listTasks() string {
if len(t.items) == 0 {
return "No tasks."
}
var b strings.Builder
b.WriteString("Tasks:\n")
for _, item := range t.items {
fmt.Fprintf(&b, "[%s] %s\n", item.Status, item.Content)
}
return b.String()
}
func (t *TodoTool) createTasks(items []string) (string, error) {
if len(items) == 0 {
return "", fmt.Errorf("items is required for create")
}
for _, content := range items {
content = strings.TrimSpace(content)
if content == "" {
return "", fmt.Errorf("task content must not be empty")
}
t.items = append(t.items, todoItem{
Content: content,
Status: "pending",
})
}
return t.listTasks(), nil
}
func (t *TodoTool) startTask(task string) (string, error) {
idx := -1
for i, item := range t.items {
if item.Content == task {
idx = i
break
}
}
if idx < 0 {
return "", fmt.Errorf("task not found: %s", task)
}
// Set all in_progress items to pending, then mark the target as in_progress.
for i := range t.items {
if t.items[i].Status == "in_progress" {
t.items[i].Status = "pending"
}
}
t.items[idx].Status = "in_progress"
return t.listTasks(), nil
}
func (t *TodoTool) completeTask(task string) (string, error) {
idx := -1
for i, item := range t.items {
if item.Content == task {
idx = i
break
}
}
if idx < 0 {
return "", fmt.Errorf("task not found: %s", task)
}
t.items[idx].Status = "completed"
return t.listTasks(), nil
}
func (t *TodoTool) updateTask(task, status string) (string, error) {
if status != "pending" && status != "in_progress" && status != "completed" {
return "", fmt.Errorf("status must be one of: pending, in_progress, completed")
}
idx := -1
for i, item := range t.items {
if item.Content == task {
idx = i
break
}
}
if idx < 0 {
return "", fmt.Errorf("task not found: %s", task)
}
t.items[idx].Status = status
return t.listTasks(), nil
}
+126 -6
View File
@@ -5,17 +5,68 @@ 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)
}
// StreamingTool is optionally implemented by tools that can report partial
// output while they execute. emit receives chunks of the tool's output as
// they are produced; the returned string is still the complete result that
// gets sent back to the model.
type StreamingTool interface {
Tool
ExecuteStream(ctx context.Context, args json.RawMessage, emit func(string) error) (string, error)
}
// MutatingTool is optionally implemented by tools that modify external state.
// It is used by the agent to determine whether tool calls can run in parallel.
type MutatingTool interface {
Mutates() bool
}
// CallMutatingTool can refine a tool's mutation classification per call.
// This is useful for tools such as apply_patch whose dry-run mode is read-only.
type CallMutatingTool interface {
MutatesCall(args json.RawMessage) bool
}
// PreviewableTool describes a pending mutation for interactive approval.
type PreviewableTool interface {
Preview(ctx context.Context, args json.RawMessage) (string, error)
}
// TargetPathTool reports files that should be snapshotted before a mutation.
type TargetPathTool interface {
TargetPaths(args json.RawMessage) ([]string, error)
}
// IsMutating reports whether tool implements MutatingTool and returns true.
func IsMutating(tool Tool) bool {
mutating, ok := tool.(MutatingTool)
return ok && mutating.Mutates()
}
// IsMutatingCall reports whether this invocation modifies external state.
func IsMutatingCall(tool Tool, args json.RawMessage) bool {
if !IsMutating(tool) {
return false
}
if call, ok := tool.(CallMutatingTool); ok {
return call.MutatesCall(args)
}
return true
}
type Registry struct {
tools map[string]Tool
primary []string
@@ -62,38 +113,70 @@ func ReadOnlyBuiltins(workingDir string) *Registry {
fileRead := NewFileReadTool(workingDir)
fileList := NewFileListTool(workingDir)
fileSearch := NewFileSearchTool(workingDir)
codeSymbols := NewCodeSymbolsTool(workingDir)
gitStatus := NewGitStatusTool(workingDir)
gitDiff := NewGitDiffTool(workingDir)
gitLog := NewGitLogTool(workingDir)
projectOutline := NewProjectOutlineTool(workingDir)
todo := NewTodoTool()
registry := NewRegistry(fileRead, fileList, fileSearch)
registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols, projectOutline, todo, gitStatus, gitDiff, gitLog)
registry.RegisterAlias("file.read", fileRead)
registry.RegisterAlias("file.list", fileList)
registry.RegisterAlias("file.search", fileSearch)
registry.RegisterAlias("code.symbols", codeSymbols)
registry.RegisterAlias("git.status", gitStatus)
registry.RegisterAlias("git.diff", gitDiff)
registry.RegisterAlias("git.log", gitLog)
return registry
}
func Builtins(workingDir string) *Registry {
registry := ReadOnlyBuiltins(workingDir)
fileWrite := NewFileWriteTool(workingDir)
fileEdit := NewFileEditTool(workingDir)
applyPatch := NewApplyPatchTool(workingDir)
shellRun := NewShellRunTool(workingDir)
registry.Register(fileWrite)
registry.Register(fileEdit)
registry.Register(applyPatch)
registry.Register(shellRun)
registry.RegisterAlias("file.write", fileWrite)
registry.RegisterAlias("file.edit", fileEdit)
registry.RegisterAlias("apply.patch", applyPatch)
registry.RegisterAlias("shell.run", shellRun)
return registry
}
func AgentInstructions(workingDir string, yolo bool) string {
return AgentInstructionsForMode(workingDir, yolo, yolo)
}
func AgentInstructionsForMode(workingDir string, writesEnabled bool, yolo bool) string {
searchInstructions := webSearchInstructions()
readOnlyTools := "file_list, file_read, file_search, code_symbols, project_outline, todo, git_status, git_diff, git_log, web_search, web_fetch"
writeTools := "file_list, file_read, file_search, code_symbols, project_outline, todo, file_write, file_edit, apply_patch, shell_run, git_status, git_diff, git_log, web_search, web_fetch"
if writesEnabled {
approval := "Mutating tools require interactive user approval before execution. Explain the intended change through the tool call and wait for the decision."
if yolo {
approval = "Mutating tools execute automatically because --yolo is enabled."
}
return fmt.Sprintf(`Local project context:
- The project working directory is %q.
- Tools are available, but use them only when needed to satisfy the user's request.
- %s
- 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.
- Use project_outline first when the user asks for a broad understanding of a project; prefer file_edit for localized changes, apply_patch for related multi-file changes, and file_write for whole-file creation or replacement.
- Do not use file tools as a substitute for a command execution request.
- When the user asks about changed files, diffs, or git history, use git_status, git_diff, or git_log instead of shell_run.
- Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions.
- %s
- Tool paths should normally be relative to the project working directory.
- Available tools: file_list, file_read, file_search, file_write, shell_run.`, workingDir)
- Available tools: %s.`, workingDir, approval, searchInstructions, writeTools)
}
return fmt.Sprintf(`Local project context:
- The project working directory is %q.
@@ -101,9 +184,16 @@ 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.
- When the user asks about changed files, diffs, or git history, use git_status, git_diff, or git_log instead of shell_run.
- %s
- Tool paths should normally be relative to the project working directory.
- Available tools: file_list, file_read, file_search.`, workingDir)
- Available tools: %s.`, workingDir, searchInstructions, readOnlyTools)
}
func webSearchInstructions() string {
return "When the user's request depends on current or external web information, use web_search. This includes news, prices, laws, schedules, software versions, product recommendations, or requests for the latest information. Use web_fetch to read a source URL when search snippets are not enough, when the user provides a URL, or when the user names a specific website to inspect. Base answers on returned sources and include source URLs when useful."
}
func JSONSchema(schema string) json.RawMessage {
@@ -111,10 +201,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 {
+170 -5
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", "project_outline", "todo", "git_status", "git_diff", "git_log"} {
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", "git.status", "git.diff", "git.log"} {
if _, ok := registry.Get(alias); !ok {
t.Fatalf("alias missing: %s", alias)
}
@@ -82,11 +117,63 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
if _, ok := registry.Get("file_write"); ok {
t.Fatal("read-only registry should not expose file_write")
}
if _, ok := registry.Get("file_edit"); ok {
t.Fatal("read-only registry should not expose file_edit")
}
}
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))
registry.Register(NewBuiltinFetchTool(nil))
var names []string
for _, def := range registry.Definitions() {
names = append(names, def.Function.Name)
}
for _, want := range []string{"web_search", "web_fetch"} {
if !slices.Contains(names, want) {
t.Fatalf("definitions missing %s: %#v", want, names)
}
}
for _, want := range []string{"web_search", "web_fetch"} {
if _, ok := registry.Get(want); !ok {
t.Fatalf("%s missing", want)
}
}
}
func TestBuiltinsExposeYoloTools(t *testing.T) {
registry := Builtins(t.TempDir())
for _, name := range []string{"file_write", "shell_run", "file.write", "shell.run"} {
for _, name := range []string{"file_write", "file_edit", "apply_patch", "shell_run", "file.write", "file.edit", "apply.patch", "shell.run"} {
if _, ok := registry.Get(name); !ok {
t.Fatalf("tool missing: %s", name)
}
@@ -95,16 +182,94 @@ func TestBuiltinsExposeYoloTools(t *testing.T) {
func TestAgentInstructionsExplainCommandModes(t *testing.T) {
readOnly := AgentInstructions("/tmp/project", false)
for _, want := range []string{"Shell command execution is disabled", "--yolo", "Do not call tools for greetings", "Do not explore the project preemptively", "Do not call file_list just to discover context"} {
for _, want := range []string{"Shell command execution is disabled", "--yolo", "Do not call tools for greetings", "Do not explore the project preemptively", "Do not call file_list just to discover context", "git_status, git_diff, git_log"} {
if !strings.Contains(readOnly, want) {
t.Fatalf("read-only instructions missing %q:\n%s", want, readOnly)
}
}
yolo := AgentInstructions("/tmp/project", true)
for _, want := range []string{"use shell_run", "commands written in backticks", "Do not start interactive", "non-interactive local commands", "Do not call tools for greetings"} {
for _, want := range []string{"use shell_run", "commands written in backticks", "Do not start interactive", "non-interactive local commands", "Do not call tools for greetings", "prefer file_edit"} {
if !strings.Contains(yolo, want) {
t.Fatalf("yolo instructions missing %q:\n%s", want, yolo)
}
}
}
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, code_symbols, project_outline, todo, git_status, git_diff, git_log, 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 ..")
}
}
func TestIsMutating(t *testing.T) {
dir := t.TempDir()
if IsMutating(NewFileReadTool(dir)) {
t.Fatal("FileReadTool should not be mutating")
}
if !IsMutating(NewFileWriteTool(dir)) {
t.Fatal("FileWriteTool should be mutating")
}
if !IsMutating(NewShellRunTool(dir)) {
t.Fatal("ShellRunTool should be mutating")
}
}
+200
View File
@@ -0,0 +1,200 @@
package tools
import (
"context"
"encoding/json"
"os"
"path/filepath"
"strings"
"testing"
)
func TestApplyPatchChangesMultipleFiles(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "old.txt"), []byte("before\n"), 0o600); err != nil {
t.Fatal(err)
}
patch := `*** Begin Patch
*** Update File: old.txt
@@
-before
+after
*** Add File: new.txt
+created
*** End Patch`
out, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch}))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "2 file(s) changed") {
t.Fatalf("output = %q", out)
}
assertFileContent(t, filepath.Join(dir, "old.txt"), "after\n")
assertFileContent(t, filepath.Join(dir, "new.txt"), "created\n")
if info, err := os.Stat(filepath.Join(dir, "old.txt")); err != nil || info.Mode().Perm() != 0o600 {
t.Fatalf("updated mode = %v, err = %v", info.Mode().Perm(), err)
}
}
func TestApplyPatchRollsBackEarlierFilesOnFailure(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "first.txt")
if err := os.WriteFile(path, []byte("original\n"), 0o644); err != nil {
t.Fatal(err)
}
patch := `*** Begin Patch
*** Update File: first.txt
@@
-original
+changed
*** Update File: missing.txt
@@
-missing
+changed
*** End Patch`
_, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch}))
if err == nil {
t.Fatal("expected patch failure")
}
assertFileContent(t, path, "original\n")
}
func TestApplyPatchDryRunDoesNotWrite(t *testing.T) {
dir := t.TempDir()
patch := "*** Begin Patch\n*** Add File: dry.txt\n+dry\n*** End Patch"
out, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch, "dry_run": true}))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "dry run") {
t.Fatalf("output = %q", out)
}
if _, err := os.Stat(filepath.Join(dir, "dry.txt")); !os.IsNotExist(err) {
t.Fatalf("dry run created file: %v", err)
}
}
func TestApplyPatchRejectsMalformedEnvelopeAndDuplicatePaths(t *testing.T) {
for _, patch := range []string{
"*** Begin Patch\n*** Add File: x.txt\n+x",
"prefix\n*** Begin Patch\n*** Add File: x.txt\n+x\n*** End Patch",
"*** Begin Patch\n*** Add File: x.txt\n+x\n*** Delete File: x.txt\n*** End Patch",
} {
if _, err := parsePatch(patch); err == nil {
t.Fatalf("expected parse failure for %q", patch)
}
}
}
func TestResolvePathRejectsSymlinkEscape(t *testing.T) {
parent := t.TempDir()
base := filepath.Join(parent, "project")
outside := filepath.Join(parent, "outside")
if err := os.MkdirAll(base, 0o755); err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(outside, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(outside, "secret.txt"), []byte("secret"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.Symlink(outside, filepath.Join(base, "link")); err != nil {
t.Skipf("symlinks unavailable: %v", err)
}
if _, err := resolvePath(base, "link/secret.txt"); err == nil {
t.Fatal("expected read symlink escape rejection")
}
if _, err := resolvePathForWrite(base, "link/new.txt"); err == nil {
t.Fatal("expected write symlink escape rejection")
}
}
func TestTodoRejectsInvalidData(t *testing.T) {
tool := NewTodoTool()
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"create","items":[""]}`)); err == nil {
t.Fatal("expected empty task rejection")
}
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"create","items":["one"]}`)); err != nil {
t.Fatal(err)
}
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"update","task":"one","status":"invalid"}`)); err == nil {
t.Fatal("expected invalid status rejection")
}
}
func TestProjectOutlineIncludesGoTypeScriptAndRustSymbols(t *testing.T) {
dir := t.TempDir()
files := map[string]string{
"main.go": "package demo\nfunc Run() {}\n",
"ui.ts": "export class Panel {}\nexport function render() {}\n",
"lib.rs": "pub struct Worker;\npub fn execute() {}\n",
}
for name, content := range files {
if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644); err != nil {
t.Fatal(err)
}
}
out, err := NewProjectOutlineTool(dir).Execute(context.Background(), json.RawMessage(`{}`))
if err != nil {
t.Fatal(err)
}
for _, want := range []string{"main.go", "Run", "ui.ts", "Panel", "render", "lib.rs", "Worker", "execute"} {
if !strings.Contains(out, want) {
t.Fatalf("outline missing %q:\n%s", want, out)
}
}
}
func TestRollbackStoreRestoresInReverseOrder(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "state.txt")
if err := os.WriteFile(path, []byte("one"), 0o644); err != nil {
t.Fatal(err)
}
store := NewRollbackStore()
first, err := store.Capture("edit", "state.txt", path)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("two"), 0o644); err != nil {
t.Fatal(err)
}
second, err := store.Capture("edit", "state.txt", path)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("three"), 0o644); err != nil {
t.Fatal(err)
}
if store.LastID() != second {
t.Fatalf("last ID = %q", store.LastID())
}
if _, err := store.Restore(second); err != nil {
t.Fatal(err)
}
assertFileContent(t, path, "two")
if store.LastID() != first {
t.Fatalf("last ID after restore = %q, want %q", store.LastID(), first)
}
}
func mustJSON(t *testing.T, value any) json.RawMessage {
t.Helper()
data, err := json.Marshal(value)
if err != nil {
t.Fatal(err)
}
return data
}
func assertFileContent(t *testing.T, path, want string) {
t.Helper()
data, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(data) != want {
t.Fatalf("%s content = %q, want %q", path, string(data), want)
}
}
+597
View File
@@ -0,0 +1,597 @@
package tools
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"agentu/pkg/llm"
"golang.org/x/net/html"
)
const (
defaultSearchBaseURL = "https://html.duckduckgo.com/html/"
defaultUserAgent = "agentu/0.0.3"
maxSearchResults = 20
maxWebReadBytes = 2 * 1024 * 1024
)
type BuiltinSearchTool struct {
baseURL string
maxResults int
client *http.Client
}
func NewBuiltinSearchTool(client *http.Client) *BuiltinSearchTool {
return NewBuiltinSearchToolWithBaseURL(defaultSearchBaseURL, 5, client)
}
func NewBuiltinSearchToolWithBaseURL(baseURL string, maxResults int, client *http.Client) *BuiltinSearchTool {
if strings.TrimSpace(baseURL) == "" {
baseURL = defaultSearchBaseURL
}
if maxResults <= 0 {
maxResults = 5
}
if client == nil {
client = http.DefaultClient
}
return &BuiltinSearchTool{
baseURL: baseURL,
maxResults: maxResults,
client: client,
}
}
func (t *BuiltinSearchTool) Definition() llm.Tool {
return webSearchDefinition()
}
func webSearchDefinition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "web_search",
Description: "Search the public web for current or external information. Use for recent facts, news, prices, schedules, versions, policies, recommendations, or other questions that need online sources. Results include source URLs.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"query": {"type": "string", "description": "Search query."},
"max_results": {"type": "integer", "description": "Optional result count from 1 to 20. Defaults to 5.", "minimum": 1, "maximum": 20},
"topic": {"type": "string", "description": "Optional topic hint. Accepted for compatibility but the built-in backend may ignore it.", "enum": ["general", "news", "finance"]}
},
"required": ["query"],
"additionalProperties": false
}`),
},
}
}
func (t *BuiltinSearchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
args, err := parseWebSearchArgs(raw, t.maxResults)
if err != nil {
return "", err
}
endpoint, err := url.Parse(t.baseURL)
if err != nil {
return "", fmt.Errorf("parse search endpoint: %w", err)
}
query := endpoint.Query()
query.Set("q", args.Query)
endpoint.RawQuery = query.Encode()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint.String(), nil)
if err != nil {
return "", fmt.Errorf("create search request: %w", err)
}
req.Header.Set("User-Agent", defaultUserAgent)
req.Header.Set("Accept", "text/html,application/xhtml+xml")
resp, err := t.client.Do(req)
if err != nil {
return "", fmt.Errorf("search request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 8*1024))
return "", fmt.Errorf("web search failed: status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
}
results, err := parseSearchResults(resp.Body, args.MaxResults)
if err != nil {
return "", err
}
return FormatResult(formatWebSearchResults(args.Query, "", results)), nil
}
type BuiltinFetchTool struct {
client *http.Client
}
type webSearchArgs struct {
Query string `json:"query"`
MaxResults int `json:"max_results"`
Topic string `json:"topic"`
}
type webSearchResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
Score float64 `json:"score"`
}
type webFetchArgs struct {
URL string `json:"url"`
Query string `json:"query"`
}
type webFetchResponse struct {
Results []webFetchResult `json:"results"`
FailedResults []webFetchFailed `json:"failed_results"`
}
type webFetchResult struct {
URL string `json:"url"`
RawContent string `json:"raw_content"`
}
type webFetchFailed struct {
URL string `json:"url"`
Error string `json:"error"`
}
func NewBuiltinFetchTool(client *http.Client) *BuiltinFetchTool {
if client == nil {
client = http.DefaultClient
}
return &BuiltinFetchTool{client: client}
}
func (t *BuiltinFetchTool) Definition() llm.Tool {
return webFetchDefinition()
}
func webFetchDefinition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "web_fetch",
Description: "Fetch and extract readable Markdown content from a public web URL. Use after web_search when snippets are not enough, or when the user gives a URL or website to inspect. Results include the source URL.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"url": {"type": "string", "description": "Public http or https URL to fetch. URLs without a scheme are treated as https URLs."},
"query": {"type": "string", "description": "Optional extraction focus. Accepted for compatibility but the built-in backend may ignore it."}
},
"required": ["url"],
"additionalProperties": false
}`),
},
}
}
func (t *BuiltinFetchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
args, err := parseWebFetchArgs(raw)
if err != nil {
return "", err
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, args.URL, nil)
if err != nil {
return "", fmt.Errorf("create fetch request: %w", err)
}
req.Header.Set("User-Agent", defaultUserAgent)
req.Header.Set("Accept", "text/html,text/plain,application/xhtml+xml")
resp, err := t.client.Do(req)
if err != nil {
return "", fmt.Errorf("fetch request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 8*1024))
return "", fmt.Errorf("web fetch failed: status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
}
data, err := io.ReadAll(io.LimitReader(resp.Body, maxWebReadBytes))
if err != nil {
return "", fmt.Errorf("read response: %w", err)
}
contentType := strings.ToLower(resp.Header.Get("Content-Type"))
var content string
if strings.Contains(contentType, "html") || contentType == "" {
content, err = htmlToMarkdown(string(data), args.URL)
if err != nil {
return "", err
}
} else if strings.Contains(contentType, "text/plain") || strings.Contains(contentType, "markdown") {
content = strings.TrimSpace(string(data))
} else {
return "", fmt.Errorf("unsupported content type: %s", contentType)
}
result := webFetchResponse{
Results: []webFetchResult{{
URL: args.URL,
RawContent: content,
}},
}
return FormatResult(formatWebFetchResults(args.URL, result)), nil
}
func parseWebSearchArgs(raw json.RawMessage, defaultMaxResults int) (webSearchArgs, error) {
var args webSearchArgs
if err := json.Unmarshal(raw, &args); err != nil {
return args, fmt.Errorf("parse arguments: %w", err)
}
args.Query = strings.TrimSpace(args.Query)
if args.Query == "" {
return args, errors.New("query is required")
}
if args.MaxResults == 0 {
args.MaxResults = defaultMaxResults
}
if args.MaxResults < 1 || args.MaxResults > maxSearchResults {
return args, fmt.Errorf("max_results must be between 1 and %d", maxSearchResults)
}
args.Topic = strings.ToLower(strings.TrimSpace(args.Topic))
if args.Topic == "" {
args.Topic = "general"
}
switch args.Topic {
case "general", "news", "finance":
default:
return args, fmt.Errorf("topic must be one of: general, news, finance")
}
return args, nil
}
func parseWebFetchArgs(raw json.RawMessage) (webFetchArgs, error) {
var args webFetchArgs
if err := json.Unmarshal(raw, &args); err != nil {
return args, fmt.Errorf("parse arguments: %w", err)
}
args.URL = strings.TrimSpace(args.URL)
if args.URL == "" {
return args, errors.New("url is required")
}
if !strings.Contains(args.URL, "://") {
args.URL = "https://" + args.URL
}
parsed, err := url.Parse(args.URL)
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" {
return args, errors.New("url must be a valid http or https URL")
}
args.Query = strings.TrimSpace(args.Query)
return args, nil
}
func formatWebSearchResults(query string, answer string, results []webSearchResult) string {
if len(results) == 0 {
return fmt.Sprintf("No web search results for %q.", query)
}
var b strings.Builder
fmt.Fprintf(&b, "Web search results for %q:\n", query)
if strings.TrimSpace(answer) != "" {
fmt.Fprintf(&b, "\nSummary: %s\n", strings.TrimSpace(answer))
}
for i, result := range results {
title := strings.TrimSpace(result.Title)
if title == "" {
title = "Untitled"
}
fmt.Fprintf(&b, "\n%d. %s\n", i+1, title)
fmt.Fprintf(&b, " Source: %s\n", strings.TrimSpace(result.URL))
if content := strings.TrimSpace(result.Content); content != "" {
fmt.Fprintf(&b, " Snippet: %s\n", content)
}
if result.Score > 0 {
fmt.Fprintf(&b, " Score: %.3f\n", result.Score)
}
}
return strings.TrimRight(b.String(), "\n")
}
func formatWebFetchResults(inputURL string, response webFetchResponse) string {
var b strings.Builder
if len(response.Results) == 0 {
fmt.Fprintf(&b, "No readable web content extracted from %s.", inputURL)
} else {
fmt.Fprintf(&b, "Fetched web content:\n")
for i, result := range response.Results {
sourceURL := strings.TrimSpace(result.URL)
if sourceURL == "" {
sourceURL = inputURL
}
fmt.Fprintf(&b, "\n%d. Source: %s\n", i+1, sourceURL)
if content := strings.TrimSpace(result.RawContent); content != "" {
fmt.Fprintf(&b, "\n%s\n", content)
}
}
}
for _, failed := range response.FailedResults {
sourceURL := strings.TrimSpace(failed.URL)
if sourceURL == "" {
sourceURL = inputURL
}
if strings.TrimSpace(failed.Error) != "" {
fmt.Fprintf(&b, "\nFetch failed for %s: %s", sourceURL, strings.TrimSpace(failed.Error))
}
}
return strings.TrimRight(b.String(), "\n")
}
func parseSearchResults(r io.Reader, maxResults int) ([]webSearchResult, error) {
doc, err := html.Parse(r)
if err != nil {
return nil, fmt.Errorf("parse search results: %w", err)
}
results := make([]webSearchResult, 0, maxResults)
var walk func(*html.Node)
walk = func(n *html.Node) {
if len(results) >= maxResults || n == nil {
return
}
if n.Type == html.ElementNode && n.Data == "div" && hasClass(n, "result") {
if result, ok := parseResultNode(n); ok {
results = append(results, result)
}
return
}
for child := n.FirstChild; child != nil; child = child.NextSibling {
walk(child)
}
}
walk(doc)
return results, nil
}
func parseResultNode(n *html.Node) (webSearchResult, bool) {
var result webSearchResult
var walk func(*html.Node)
walk = func(node *html.Node) {
if node == nil {
return
}
if node.Type == html.ElementNode && node.Data == "a" {
if hasClass(node, "result__a") && result.URL == "" {
result.Title = cleanWebText(textContent(node))
result.URL = decodeSearchURL(attr(node, "href"))
}
if hasClass(node, "result__snippet") && result.Content == "" {
result.Content = cleanWebText(textContent(node))
}
}
if node.Type == html.ElementNode && (node.Data == "div" || node.Data == "span") && hasClass(node, "result__snippet") && result.Content == "" {
result.Content = cleanWebText(textContent(node))
}
for child := node.FirstChild; child != nil; child = child.NextSibling {
walk(child)
}
}
walk(n)
return result, result.Title != "" && result.URL != ""
}
func decodeSearchURL(raw string) string {
raw = strings.TrimSpace(raw)
if raw == "" {
return ""
}
parsed, err := url.Parse(raw)
if err == nil {
if target := parsed.Query().Get("uddg"); target != "" {
if decoded, err := url.QueryUnescape(target); err == nil {
return decoded
}
return target
}
}
return raw
}
func htmlToMarkdown(input string, sourceURL string) (string, error) {
doc, err := html.Parse(strings.NewReader(input))
if err != nil {
return "", fmt.Errorf("parse html: %w", err)
}
var b strings.Builder
var title string
var walk func(*html.Node, int)
walk = func(n *html.Node, listDepth int) {
if n == nil || skipHTMLNode(n) {
return
}
if n.Type == html.ElementNode && n.Data == "title" && title == "" {
title = cleanWebText(textContent(n))
return
}
if n.Type == html.TextNode {
writeInlineText(&b, n.Data)
return
}
if n.Type != html.ElementNode {
for child := n.FirstChild; child != nil; child = child.NextSibling {
walk(child, listDepth)
}
return
}
switch n.Data {
case "h1", "h2", "h3", "h4", "h5", "h6":
writeBlockBreak(&b)
level := int(n.Data[1] - '0')
b.WriteString(strings.Repeat("#", level))
b.WriteString(" ")
writeChildren(&b, n, listDepth, walk)
writeBlockBreak(&b)
case "p", "div", "section", "article", "header", "footer", "main", "tr":
writeBlockBreak(&b)
writeChildren(&b, n, listDepth, walk)
writeBlockBreak(&b)
case "br":
b.WriteString("\n")
case "li":
writeBlockBreak(&b)
b.WriteString(strings.Repeat(" ", listDepth))
b.WriteString("- ")
writeChildren(&b, n, listDepth+1, walk)
writeBlockBreak(&b)
case "a":
label := cleanWebText(textContent(n))
href := strings.TrimSpace(attr(n, "href"))
if href != "" && label != "" {
href = resolveURL(sourceURL, href)
b.WriteString(label)
b.WriteString(" (")
b.WriteString(href)
b.WriteString(")")
return
}
writeChildren(&b, n, listDepth, walk)
default:
writeChildren(&b, n, listDepth, walk)
}
}
walk(doc, 0)
content := normalizeMarkdownText(b.String())
if title != "" && !strings.Contains(content, title) {
content = "# " + title + "\n\n" + content
}
if strings.TrimSpace(content) == "" {
return "", errors.New("no readable content found")
}
return content, nil
}
func writeChildren(b *strings.Builder, n *html.Node, listDepth int, walk func(*html.Node, int)) {
for child := n.FirstChild; child != nil; child = child.NextSibling {
walk(child, listDepth)
}
}
func writeInlineText(b *strings.Builder, text string) {
text = cleanWebText(text)
if text == "" {
return
}
if b.Len() > 0 {
last := b.String()[b.Len()-1]
if last != '\n' && last != ' ' {
b.WriteByte(' ')
}
}
b.WriteString(text)
}
func writeBlockBreak(b *strings.Builder) {
if b.Len() == 0 {
return
}
text := b.String()
if strings.HasSuffix(text, "\n\n") {
return
}
if strings.HasSuffix(text, "\n") {
b.WriteByte('\n')
return
}
b.WriteString("\n\n")
}
func skipHTMLNode(n *html.Node) bool {
if n.Type != html.ElementNode {
return false
}
switch n.Data {
case "script", "style", "noscript", "svg", "canvas", "template":
return true
default:
return false
}
}
func textContent(n *html.Node) string {
var b strings.Builder
var walk func(*html.Node)
walk = func(node *html.Node) {
if node == nil || skipHTMLNode(node) {
return
}
if node.Type == html.TextNode {
b.WriteString(node.Data)
b.WriteByte(' ')
return
}
for child := node.FirstChild; child != nil; child = child.NextSibling {
walk(child)
}
}
walk(n)
return cleanWebText(b.String())
}
func hasClass(n *html.Node, class string) bool {
for _, item := range strings.Fields(attr(n, "class")) {
if item == class {
return true
}
}
return false
}
func attr(n *html.Node, name string) string {
for _, item := range n.Attr {
if item.Key == name {
return item.Val
}
}
return ""
}
func resolveURL(baseURL, raw string) string {
ref, err := url.Parse(raw)
if err != nil {
return raw
}
base, err := url.Parse(baseURL)
if err != nil {
return raw
}
return base.ResolveReference(ref).String()
}
func cleanWebText(text string) string {
return strings.Join(strings.Fields(text), " ")
}
func normalizeMarkdownText(text string) string {
lines := strings.Split(text, "\n")
out := make([]string, 0, len(lines))
blank := false
for _, line := range lines {
line = strings.TrimSpace(line)
if line == "" {
if !blank && len(out) > 0 {
out = append(out, "")
}
blank = true
continue
}
out = append(out, line)
blank = false
}
return strings.TrimSpace(strings.Join(out, "\n"))
}
+96
View File
@@ -0,0 +1,96 @@
package tools
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestBuiltinSearchTool(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Fatalf("method = %s", r.Method)
}
if got := r.URL.Query().Get("q"); got != "agentu" {
t.Fatalf("query = %q", got)
}
w.Header().Set("Content-Type", "text/html")
_, _ = w.Write([]byte(`
<html><body>
<div class="result">
<a class="result__a" href="/l/?uddg=https%3A%2F%2Fexample.com%2Fagentu">Agentu release</a>
<a class="result__snippet">Agentu adds built-in web tools.</a>
</div>
</body></html>`))
}))
defer server.Close()
tool := NewBuiltinSearchToolWithBaseURL(server.URL, 5, server.Client())
out, err := tool.Execute(context.Background(), json.RawMessage(`{"query":"agentu"}`))
if err != nil {
t.Fatal(err)
}
for _, want := range []string{"Web search results", "Agentu release", "https://example.com/agentu", "Agentu adds built-in web tools"} {
if !strings.Contains(out, want) {
t.Fatalf("output missing %q:\n%s", want, out)
}
}
}
func TestBuiltinFetchTool(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Fatalf("method = %s", r.Method)
}
w.Header().Set("Content-Type", "text/html; charset=utf-8")
_, _ = w.Write([]byte(`
<html>
<head><title>Agentu</title><script>hidden()</script></head>
<body>
<h1>Agentu</h1>
<p>Built-in fetch reads HTML.</p>
<ul><li>Search</li><li>Fetch</li></ul>
</body>
</html>`))
}))
defer server.Close()
tool := NewBuiltinFetchTool(server.Client())
out, err := tool.Execute(context.Background(), json.RawMessage(`{"url":"`+server.URL+`"}`))
if err != nil {
t.Fatal(err)
}
for _, want := range []string{"Fetched web content", server.URL, "# Agentu", "Built-in fetch reads HTML", "- Search", "- Fetch"} {
if !strings.Contains(out, want) {
t.Fatalf("output missing %q:\n%s", want, out)
}
}
if strings.Contains(out, "hidden") {
t.Fatalf("script content should be ignored:\n%s", out)
}
}
func TestBuiltinFetchToolAcceptsURLsWithoutScheme(t *testing.T) {
args, err := parseWebFetchArgs(json.RawMessage(`{"url":"example.com/page"}`))
if err != nil {
t.Fatal(err)
}
if args.URL != "https://example.com/page" {
t.Fatalf("url = %q", args.URL)
}
}
func TestBuiltinFetchToolFormatsFailedResults(t *testing.T) {
response := webFetchResponse{
FailedResults: []webFetchFailed{{URL: "https://example.com", Error: "blocked"}},
}
out := formatWebFetchResults("https://example.com", response)
for _, want := range []string{"No readable web content", "Fetch failed", "blocked"} {
if !strings.Contains(out, want) {
t.Fatalf("output missing %q:\n%s", want, out)
}
}
}
+1454 -90
View File
File diff suppressed because it is too large Load Diff
+996 -30
View File
File diff suppressed because it is too large Load Diff
+128
View File
@@ -0,0 +1,128 @@
package tui
import (
"context"
"fmt"
"strings"
"sync"
tea "github.com/charmbracelet/bubbletea"
"agentu/internal/agent"
)
// toolApprovalMsg is sent to the TUI when a mutating tool call needs approval.
type toolApprovalMsg struct {
request agent.ToolAuthorizationRequest
respond chan agent.ToolAuthorizationDecision
}
// ApprovalBroker bridges the agent's ToolAuthorizer interface with the TUI's
// event loop, allowing interactive approval of mutating tool calls.
type ApprovalBroker struct {
mu sync.Mutex
events chan<- tea.Msg
}
// NewApprovalBroker creates a new approval broker.
func NewApprovalBroker() *ApprovalBroker {
return &ApprovalBroker{}
}
// SetEvents sets the current event channel for sending approval requests.
func (b *ApprovalBroker) SetEvents(events chan<- tea.Msg) {
b.mu.Lock()
defer b.mu.Unlock()
b.events = events
}
// ClearEvents removes the current event channel.
func (b *ApprovalBroker) ClearEvents() {
b.mu.Lock()
defer b.mu.Unlock()
b.events = nil
}
// AuthorizeToolCall implements agent.ToolAuthorizer by sending an approval
// request to the TUI and waiting for the user's decision.
func (b *ApprovalBroker) AuthorizeToolCall(ctx context.Context, req agent.ToolAuthorizationRequest) (agent.ToolAuthorizationDecision, error) {
b.mu.Lock()
events := b.events
b.mu.Unlock()
if events == nil {
return "", fmt.Errorf("approval UI is not available")
}
respond := make(chan agent.ToolAuthorizationDecision, 1)
msg := toolApprovalMsg{
request: req,
respond: respond,
}
select {
case events <- msg:
case <-ctx.Done():
return "", ctx.Err()
}
select {
case decision := <-respond:
return decision, nil
case <-ctx.Done():
return "", ctx.Err()
}
}
func (m *model) updateApproval(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
if m.approval == nil {
return *m, nil
}
switch msg.String() {
case "y", "Y":
m.respondToApproval(agent.ToolAuthorizationAllow)
return *m, nil
case "n", "N", "esc":
m.respondToApproval(agent.ToolAuthorizationDeny)
return *m, nil
case "ctrl+c":
m.respondToApproval(agent.ToolAuthorizationDeny)
if m.cancel != nil {
m.cancel()
}
return *m, tea.Quit
default:
return *m, nil
}
}
func (m *model) respondToApproval(decision agent.ToolAuthorizationDecision) {
request := m.approval
m.approval = nil
if request != nil {
request.respond <- decision
}
m.status = "working..."
m.resize()
m.refreshViewport(true)
}
func (m model) approvalView() string {
if m.approval == nil {
return ""
}
req := m.approval.request
preview := strings.TrimSpace(req.Preview)
if preview == "" {
preview = strings.TrimSpace(req.Arguments)
}
lines := []string{
m.styles.ErrorLabel.Render("Approve " + req.ToolName + "?"),
m.styles.ToolValue.Render(preview),
}
if req.RollbackAvailable {
lines = append(lines, m.styles.Muted.Render("A rollback snapshot will be created."))
}
lines = append(lines, m.styles.Muted.Render("y approve · n/esc deny"))
return m.styles.InputBox.Width(max(20, m.width)).Render(strings.Join(lines, "\n"))
}
+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 {
+128 -45
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:
@@ -32,6 +43,7 @@ type theme struct {
Background string
Surface string
SurfaceAlt string
InputBg string
Border string
Text string
Muted string
@@ -46,17 +58,28 @@ type theme struct {
type styles struct {
App lipgloss.Style
Header lipgloss.Style
Title lipgloss.Style
Muted lipgloss.Style
Activity lipgloss.Style
ModelMeta lipgloss.Style
Header lipgloss.Style
Footer lipgloss.Style
Viewport lipgloss.Style
InputBox lipgloss.Style
CommandHint lipgloss.Style
InputModeReadonly lipgloss.Style
InputModeYolo lipgloss.Style
InputModeCommand 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
AssistantLabel lipgloss.Style
ToolLabel lipgloss.Style
ErrorLabel lipgloss.Style
SystemLabel lipgloss.Style
@@ -78,8 +101,8 @@ func lightTheme() theme {
return theme{
Mode: ThemeLight,
Background: "#F8FAFC",
Surface: "#FFFFFF",
SurfaceAlt: "#EEF2F7",
InputBg: "#E4E4E4",
Border: "#CBD5E1",
Text: "#111827",
Muted: "#64748B",
@@ -96,8 +119,8 @@ func darkTheme() theme {
return theme{
Mode: ThemeDark,
Background: "#111827",
Surface: "#1F2937",
SurfaceAlt: "#0F172A",
InputBg: "#383838",
Border: "#475569",
Text: "#E5E7EB",
Muted: "#94A3B8",
@@ -110,67 +133,116 @@ func darkTheme() theme {
}
}
func appStyle(t theme) lipgloss.Style {
style := lipgloss.NewStyle().Foreground(lipgloss.Color(t.Text))
if t.Mode == ThemeDark {
style = style.Background(lipgloss.Color(t.Background))
}
return style
}
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)),
App: appStyle(t),
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),
ModelMeta: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.SurfaceAlt)).
Padding(0, 1),
Header: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Title)),
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()).
BorderForeground(lipgloss.Color(t.Border)).
Padding(0, 1),
Background(lipgloss.Color(t.InputBg)).
Padding(1, 2, 1, 1),
CommandHint: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1),
InputModeReadonly: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Muted)).
Background(lipgloss.Color(t.InputBg)).
Padding(0, 1),
InputModeYolo: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Error)).
Background(lipgloss.Color(t.InputBg)).
Padding(0, 1),
InputModeCommand: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.User)).
Background(lipgloss.Color(t.InputBg)).
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),
PillAccent: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Title)).
Background(lipgloss.Color(t.SurfaceAlt)).
Padding(0, 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),
AssistantLabel: labelStyle(t.Assistant),
ToolLabel: labelStyle(t.Tool),
ErrorLabel: labelStyle(t.Error),
SystemLabel: labelStyle(t.System),
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),
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 +252,32 @@ 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)),
Background(lipgloss.Color(t.InputBg)),
CursorLine: lipgloss.NewStyle().
Background(lipgloss.Color(t.Surface)),
Background(lipgloss.Color(t.InputBg)),
Placeholder: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)),
Prompt: lipgloss.NewStyle().
@@ -205,7 +288,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().Background(lipgloss.Color(t.InputBg))
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()
}
+15
View File
@@ -0,0 +1,15 @@
package main
import (
"fmt"
"os"
"agentu/internal/app"
)
func main() {
if err := app.Run(); err != nil {
fmt.Fprintln(os.Stderr, "agentu:", err)
os.Exit(1)
}
}
@@ -14,6 +14,11 @@ import (
const DefaultPath = "~/.agentu/config.yaml"
const (
APITypeChat = "chat"
APITypeResponses = "responses"
)
type Config struct {
Providers map[string]ProviderConfig `yaml:"providers"`
Agent AgentConfig `yaml:"agent"`
@@ -21,19 +26,29 @@ 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"`
APIType string
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"`
MaxContextTokens int `yaml:"max_context_tokens"`
}
type rawConfig struct {
@@ -42,18 +57,24 @@ type rawConfig struct {
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 {
APIType string `yaml:"api_type"`
BaseURL string `yaml:"base_url"`
Model string `yaml:"model"`
Models []string `yaml:"models"`
Models []rawModelConfig `yaml:"models"`
APIKey string `yaml:"api_key"`
Thinking *string `yaml:"thinking"`
ThinkingParam string `yaml:"thinking_param"`
}
type rawModelConfig struct {
ID string `yaml:"id"`
Thinking string `yaml:"thinking"`
ContextTokens int `yaml:"context"`
}
func Load(path string) (*Config, error) {
path, err := ResolvePath(path)
if err != nil {
@@ -78,6 +99,7 @@ func Load(path string) (*Config, error) {
Agent: AgentConfig{
SystemPrompt: raw.Agent.SystemPrompt,
WorkingDir: raw.Agent.WorkingDir,
MaxContextTokens: raw.Agent.MaxContextTokens,
},
}
@@ -139,11 +161,14 @@ func (c *Config) Validate() error {
for name, provider := range c.Providers {
var missing []string
prefix := "providers." + name
if !IsAPIType(provider.APIType) {
return fmt.Errorf("%s.api_type must be one of: %s", prefix, strings.Join(APITypes(), ", "))
}
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 +176,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 +196,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 +228,94 @@ func normalizeProviders(raw rawConfig) map[string]ProviderConfig {
}
func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
models, modelConfigs := normalizeModelConfigs(raw.Models)
provider := ProviderConfig{
Name: name,
APIType: normalizeAPIType(raw.APIType),
BaseURL: raw.BaseURL,
Model: raw.Model,
Models: append([]string(nil), raw.Models...),
Models: models,
ModelConfigs: modelConfigs,
APIKey: raw.APIKey,
ThinkingParam: raw.ThinkingParam,
ThinkingParam: strings.TrimSpace(raw.ThinkingParam),
}
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 = DefaultThinkingParam(provider.APIType)
}
return provider
}
func normalizeAPIType(value string) string {
value = strings.ToLower(strings.TrimSpace(value))
if value == "" {
return APITypeChat
}
return value
}
func APITypes() []string {
return []string{APITypeChat, APITypeResponses}
}
func IsAPIType(value string) bool {
value = strings.ToLower(strings.TrimSpace(value))
for _, apiType := range APITypes() {
if value == apiType {
return true
}
}
return false
}
// DefaultThinkingParam returns the default request field used for the
// thinking level. Chat-completions providers historically use a top-level
// "thinking" field; the OpenAI Responses API uses "reasoning" with an
// {"effort": ...} value.
func DefaultThinkingParam(apiType string) string {
if apiType == APITypeResponses {
return "reasoning"
}
return "thinking"
}
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"}
}
+504
View File
@@ -0,0 +1,504 @@
package config
import (
"os"
"path/filepath"
"strings"
"testing"
"time"
)
func TestLoadExpandsEnvAndDefaults(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: test-model
agent:
working_dir: .
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
provider := cfg.Providers[cfg.DefaultProviderName()]
if provider.APIKey != "sk-test" {
t.Fatalf("api key = %q", provider.APIKey)
}
if providerName := cfg.DefaultProviderName(); providerName != "test" {
t.Fatalf("default provider = %q", providerName)
}
if cfg.Agent.ToolTimeout != 30*time.Second {
t.Fatalf("tool timeout = %s", cfg.Agent.ToolTimeout)
}
if !filepath.IsAbs(cfg.Agent.WorkingDir) {
t.Fatalf("working dir is not absolute: %s", cfg.Agent.WorkingDir)
}
}
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)
t.Setenv("AGENTU_TEST_KEY", "sk-test")
configDir := filepath.Join(dir, ".agentu")
if err := os.MkdirAll(configDir, 0o755); err != nil {
t.Fatal(err)
}
path := filepath.Join(configDir, "config.yaml")
if err := os.WriteFile(path, []byte(`
providers:
test:
base_url: https://example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load("")
if err != nil {
t.Fatal(err)
}
if got := cfg.DefaultProviderName(); got != "test" {
t.Fatalf("default provider = %q", got)
}
}
func TestResolvePathExpandsHome(t *testing.T) {
dir := t.TempDir()
t.Setenv("HOME", dir)
got, err := ResolvePath("~/.agentu/config.yaml")
if err != nil {
t.Fatal(err)
}
want := filepath.Join(dir, ".agentu", "config.yaml")
if got != want {
t.Fatalf("path = %q, want %q", got, want)
}
}
func TestLoadDefaultSystemPromptDiscouragesEagerTools(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: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
for _, want := range []string{"Answer simple conversation directly", "Use tools only when"} {
if !strings.Contains(cfg.Agent.SystemPrompt, want) {
t.Fatalf("system prompt missing %q: %s", want, cfg.Agent.SystemPrompt)
}
}
}
func TestLoadMultiProviderConfig(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:
loveuer:
base_url: https://example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: model-a
thinking: xhigh
- id: model-b
backup:
base_url: https://backup.example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: model-c
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.DefaultProviderName(); got != "backup" {
t.Fatalf("default provider = %q", got)
}
provider := cfg.Providers["loveuer"]
if got := provider.Model; got != "model-a" {
t.Fatalf("active model = %q", got)
}
if got := cfg.Providers["backup"].Model; got != "model-c" {
t.Fatalf("backup default model = %q", got)
}
if !provider.ThinkingConfigured || provider.Thinking != "xhigh" || provider.ThinkingParam != "thinking" {
t.Fatalf("thinking config = %#v", provider)
}
}
func TestLoadDefaultProviderNameIsDeterministic(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:
two:
base_url: https://two.example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: model-b
one:
base_url: https://one.example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: model-a
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.DefaultProviderName(); got != "one" {
t.Fatalf("default provider = %q", got)
}
}
func TestLoadRejectsInvalidThinking(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: test-model
thinking: huge
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "thinking") {
t.Fatalf("err = %v", err)
}
}
func TestLoadRequiresProviderFields(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
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil {
t.Fatal("expected error")
}
}
func TestLoadRejectsLegacyProviderConfig(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
provider:
base_url: https://example.com
api_key: sk-test
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "field provider not found") {
t.Fatalf("err = %v", err)
}
}
func TestLoadRejectsActiveProviderConfig(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
active_provider: loveuer
providers:
loveuer:
base_url: https://example.com
api_key: sk-test
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "field active_provider not found") {
t.Fatalf("err = %v", err)
}
}
func TestLoadDefaultsAPITypeToChat(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: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.Providers["test"].APIType; got != APITypeChat {
t.Fatalf("api type = %q, want %q", got, APITypeChat)
}
}
func TestLoadParsesResponsesAPIType(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:
openai:
api_type: RESPONSES
base_url: https://api.openai.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: gpt-test
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.Providers["openai"].APIType; got != APITypeResponses {
t.Fatalf("api type = %q, want %q", got, APITypeResponses)
}
}
func TestLoadRejectsInvalidAPIType(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:
api_type: completions
base_url: https://example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "api_type") {
t.Fatalf("err = %v", err)
}
}
func TestLoadDefaultsThinkingParamByAPIType(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:
chat:
api_type: chat
base_url: https://chat.example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: chat-model
thinking: none
responses:
api_type: responses
base_url: https://responses.example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: responses-model
thinking: high
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.Providers["chat"].ThinkingParam; got != "thinking" {
t.Fatalf("chat thinking param = %q, want %q", got, "thinking")
}
if got := cfg.Providers["responses"].ThinkingParam; got != "reasoning" {
t.Fatalf("responses thinking param = %q, want %q", got, "reasoning")
}
}
func TestLoadCustomThinkingParam(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
custom:
api_type: chat
base_url: https://example.com
api_key: sk-test
thinking_param: reasoning_effort
models:
- id: custom-model
thinking: high
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.Providers["custom"].ThinkingParam; got != "reasoning_effort" {
t.Fatalf("thinking param = %q", got)
}
}
+62 -9
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,11 +36,21 @@ 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)
}
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)
@@ -48,18 +63,56 @@ func (c *OpenAICompatibleClient) ChatStream(ctx context.Context, req ChatRequest
resp, err := c.httpClient.Do(httpReq)
if err != nil {
return fmt.Errorf("send chat request: %w", err)
lastErr = fmt.Errorf("send chat request: %w", err)
if !shouldRetryRequest(ctx, attempt, 0) {
return lastErr
}
if waitErr := waitBeforeRetry(ctx); waitErr != nil {
return waitErr
}
continue
}
defer resp.Body.Close()
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)))
_ = 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
}
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 {
reader := bufio.NewReader(r)
@@ -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)
+481
View File
@@ -0,0 +1,481 @@
package llm
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"sort"
"strings"
)
const responsesEndpoint = "/v1/responses"
// OpenAIResponsesClient implements llm.Provider against the OpenAI Responses
// API (/v1/responses). It maps the agent's chat-style conversation history to
// Responses input items (messages, function_call, function_call_output) and
// converts Responses streaming events back into the shared StreamEvent shape.
type OpenAIResponsesClient struct {
baseURL string
apiKey string
httpClient *http.Client
}
func NewOpenAIResponsesClient(baseURL, apiKey string, httpClient *http.Client) *OpenAIResponsesClient {
if httpClient == nil {
httpClient = http.DefaultClient
}
return &OpenAIResponsesClient{
baseURL: strings.TrimRight(baseURL, "/"),
apiKey: apiKey,
httpClient: httpClient,
}
}
func (c *OpenAIResponsesClient) ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error {
body, err := json.Marshal(buildResponsesRequest(req))
if err != nil {
return fmt.Errorf("marshal responses request: %w", err)
}
state := &responsesStreamState{}
return postResponsesSSE(ctx, c.baseURL, c.apiKey, c.httpClient, body, func(payload string) error {
return state.handle(payload, emit)
})
}
func postResponsesSSE(ctx context.Context, baseURL, apiKey string, httpClient *http.Client, body []byte, handle func(payload string) error) error {
var lastErr error
for attempt := 1; attempt <= maxChatAttempts; attempt++ {
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+responsesEndpoint, bytes.NewReader(body))
if err != nil {
return fmt.Errorf("create responses request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
if apiKey != "" {
httpReq.Header.Set("Authorization", "Bearer "+apiKey)
}
resp, err := httpClient.Do(httpReq)
if err != nil {
lastErr = fmt.Errorf("send responses 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)
_ = resp.Body.Close()
lastErr = fmt.Errorf("responses 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
}
defer resp.Body.Close()
return readResponsesSSE(resp.Body, handle)
}
return lastErr
}
func readResponsesSSE(r io.Reader, handle func(payload string) error) error {
reader := bufio.NewReader(r)
for {
line, err := reader.ReadString('\n')
if err != nil && !errors.Is(err, io.EOF) {
return fmt.Errorf("read responses stream: %w", err)
}
line = strings.TrimSpace(line)
if strings.HasPrefix(line, "data:") {
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if payload == "[DONE]" {
return nil
}
if handleErr := handle(payload); handleErr != nil {
return handleErr
}
}
if errors.Is(err, io.EOF) {
return nil
}
}
}
// buildResponsesRequest converts a ChatRequest into the /v1/responses body.
func buildResponsesRequest(req ChatRequest) map[string]any {
payload := map[string]any{
"model": req.Model,
"input": buildResponsesInput(req.Messages),
"stream": true,
}
if instructions := responsesInstructions(req.Messages); instructions != "" {
payload["instructions"] = instructions
}
if len(req.Tools) > 0 {
payload["tools"] = buildResponsesTools(req.Tools)
}
if req.ToolChoice != "" {
payload["tool_choice"] = req.ToolChoice
}
for key, value := range req.Extra {
if _, exists := payload[key]; exists {
continue
}
if key == "reasoning" {
if level, ok := value.(string); ok {
payload[key] = map[string]any{"effort": mapReasoningEffort(level)}
continue
}
}
payload[key] = value
}
return payload
}
// buildResponsesInput maps chat-style messages to Responses input items.
// Assistant tool calls become function_call items and tool results become
// function_call_output items, preserving their position in the conversation.
func buildResponsesInput(messages []Message) []any {
items := make([]any, 0, len(messages)+4)
for _, msg := range messages {
switch msg.Role {
case RoleSystem:
// System instructions are sent as the top-level "instructions"
// field and must not be duplicated inside the input array.
continue
case RoleUser:
items = append(items, responsesMessageInput{Role: "user", Content: msg.Content})
case RoleAssistant:
if msg.Content != "" {
items = append(items, responsesMessageInput{Role: "assistant", Content: msg.Content})
}
for _, call := range msg.ToolCalls {
items = append(items, responsesFunctionCallInput{
Type: "function_call",
CallID: call.ID,
Name: call.Function.Name,
Arguments: call.Function.Arguments,
})
}
case RoleTool:
items = append(items, responsesFunctionCallOutput{
Type: "function_call_output",
CallID: msg.ToolCallID,
Output: msg.Content,
})
}
}
return items
}
// responsesInstructions extracts system messages into a single instructions
// string, which is the OpenAI Responses API way to pass developer/system
// guidance (and the most prompt-cache friendly).
func responsesInstructions(messages []Message) string {
var instructions []string
for _, msg := range messages {
if msg.Role != RoleSystem {
continue
}
if text := strings.TrimSpace(msg.Content); text != "" {
instructions = append(instructions, text)
}
}
return strings.Join(instructions, "\n\n")
}
func buildResponsesTools(tools []Tool) []responsesTool {
out := make([]responsesTool, 0, len(tools))
for _, tool := range tools {
out = append(out, responsesTool{
Type: "function",
Name: tool.Function.Name,
Description: tool.Function.Description,
Parameters: tool.Function.Parameters,
})
}
return out
}
func mapReasoningEffort(level string) string {
level = strings.ToLower(strings.TrimSpace(level))
if level == "middle" {
return "medium"
}
return level
}
type responsesMessageInput struct {
Role string `json:"role"`
Content string `json:"content"`
}
type responsesFunctionCallInput struct {
Type string `json:"type"`
CallID string `json:"call_id"`
Name string `json:"name"`
Arguments string `json:"arguments"`
}
type responsesFunctionCallOutput struct {
Type string `json:"type"`
CallID string `json:"call_id"`
Output string `json:"output"`
}
type responsesTool struct {
Type string `json:"type"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters json.RawMessage `json:"parameters,omitempty"`
}
// responsesStreamState tracks in-flight output items while parsing SSE events.
type responsesStreamState struct {
calls map[int]*responsesCallBuilder
done bool
}
type responsesCallBuilder struct {
index int
id string
callType string
name string
arguments strings.Builder
finalized bool
}
func (s *responsesStreamState) handle(payload string, emit func(StreamEvent) error) error {
if s.done {
return nil
}
var event responsesStreamEvent
if err := json.Unmarshal([]byte(payload), &event); err != nil {
return fmt.Errorf("parse responses stream payload: %w", err)
}
switch event.Type {
case "error":
s.done = true
return event.err()
case "response.failed":
s.done = true
if event.Response != nil && event.Response.Error != nil {
return fmt.Errorf("provider error: %s: %s", event.Response.Error.Code, event.Response.Error.Message)
}
return fmt.Errorf("provider error: response failed")
case "response.incomplete":
s.done = true
return s.finish(emit, "length", nil)
case "response.completed":
s.done = true
var usage *Usage
if event.Response != nil && event.Response.Usage != nil {
usage = &Usage{
PromptTokens: event.Response.Usage.InputTokens,
CompletionTokens: event.Response.Usage.OutputTokens,
TotalTokens: event.Response.Usage.TotalTokens,
}
}
return s.finish(emit, "stop", usage)
case "response.output_text.delta":
if event.Delta != "" {
return emit(StreamEvent{Content: event.Delta})
}
case "response.output_item.added":
if event.Item != nil && event.Item.Type == "function_call" {
s.builder(event.OutputIndex).applyAdded(*event.Item)
}
case "response.function_call_arguments.delta":
if event.Delta != "" {
s.builder(event.OutputIndex).arguments.WriteString(event.Delta)
}
case "response.function_call_arguments.done":
builder := s.builder(event.OutputIndex)
if event.Name != "" {
builder.name = event.Name
}
if event.Arguments != "" {
builder.arguments.Reset()
builder.arguments.WriteString(event.Arguments)
}
case "response.output_item.done":
if event.Item != nil && event.Item.Type == "function_call" {
builder := s.builder(event.OutputIndex)
builder.applyDone(*event.Item)
return s.emitCall(builder, emit)
}
}
return nil
}
// finish flushes any not-yet-finalized function calls and reports the
// terminal finish reason.
func (s *responsesStreamState) finish(emit func(StreamEvent) error, finishReason string, usage *Usage) error {
if len(s.calls) == 0 {
return emit(StreamEvent{FinishReason: finishReason, Usage: usage})
}
indexes := make([]int, 0, len(s.calls))
for index := range s.calls {
indexes = append(indexes, index)
}
sort.Ints(indexes)
var deltas []ToolCallDelta
for _, index := range indexes {
builder := s.calls[index]
if builder.finalized {
continue
}
builder.finalized = true
deltas = append(deltas, builder.delta())
}
if len(deltas) > 0 {
if err := emit(StreamEvent{ToolCalls: deltas}); err != nil {
return err
}
}
return emit(StreamEvent{FinishReason: finishReason, Usage: usage})
}
func (s *responsesStreamState) emitCall(builder *responsesCallBuilder, emit func(StreamEvent) error) error {
if builder.finalized {
return nil
}
builder.finalized = true
return emit(StreamEvent{ToolCalls: []ToolCallDelta{builder.delta()}})
}
func (s *responsesStreamState) builder(index int) *responsesCallBuilder {
if s.calls == nil {
s.calls = make(map[int]*responsesCallBuilder)
}
builder := s.calls[index]
if builder == nil {
builder = &responsesCallBuilder{index: index}
s.calls[index] = builder
}
return builder
}
func (b *responsesCallBuilder) applyAdded(item responsesStreamItem) {
if b.id == "" {
b.id = item.CallID
if b.id == "" {
b.id = item.ID
}
}
if b.name == "" {
b.name = item.Name
}
if item.Arguments != "" {
b.arguments.Reset()
b.arguments.WriteString(item.Arguments)
}
}
func (b *responsesCallBuilder) applyDone(item responsesStreamItem) {
if item.CallID != "" {
b.id = item.CallID
} else if item.ID != "" {
b.id = item.ID
}
if item.Name != "" {
b.name = item.Name
}
if item.Arguments != "" {
b.arguments.Reset()
b.arguments.WriteString(item.Arguments)
}
}
func (b *responsesCallBuilder) delta() ToolCallDelta {
id := b.id
if id == "" {
id = fmt.Sprintf("call_%d", b.index)
}
callType := b.callType
if callType == "" {
callType = "function"
}
return ToolCallDelta{
Index: b.index,
ID: id,
Type: callType,
Name: b.name,
Arguments: b.arguments.String(),
}
}
type responsesStreamEvent struct {
Type string `json:"type"`
OutputIndex int `json:"output_index"`
Delta string `json:"delta"`
Arguments string `json:"arguments"`
Name string `json:"name"`
Code string `json:"code"`
Message string `json:"message"`
Param string `json:"param"`
Item *responsesStreamItem `json:"item"`
Response *responsesStreamResponse `json:"response"`
Error *responsesStreamError `json:"error"`
}
type responsesStreamItem struct {
Type string `json:"type"`
ID string `json:"id"`
CallID string `json:"call_id"`
Name string `json:"name"`
Arguments string `json:"arguments"`
}
type responsesStreamResponse struct {
Status string `json:"status"`
Error *responsesStreamError `json:"error"`
Usage *responsesStreamUsage `json:"usage"`
}
type responsesStreamUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
}
type responsesStreamError struct {
Type string `json:"type"`
Code string `json:"code"`
Message string `json:"message"`
Param string `json:"param"`
}
func (e responsesStreamEvent) err() error {
code, message, param := e.Code, e.Message, e.Param
if code == "" && message == "" && e.Error != nil {
code, message, param = e.Error.Code, e.Error.Message, e.Error.Param
}
if message == "" {
if param != "" {
return fmt.Errorf("provider error: %s", param)
}
return fmt.Errorf("provider error")
}
if code != "" {
return fmt.Errorf("provider error: %s: %s", code, message)
}
return fmt.Errorf("provider error: %s", message)
}
+337
View File
@@ -0,0 +1,337 @@
package llm
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestBuildResponsesRequest(t *testing.T) {
req := ChatRequest{
Model: "gpt-test",
Messages: []Message{
{Role: RoleSystem, Content: "be concise"},
{Role: RoleUser, Content: "hi"},
{
Role: RoleAssistant,
Content: "let me check",
ToolCalls: []ToolCall{{
ID: "call_1",
Type: "function",
Function: FunctionCall{
Name: "file_read",
Arguments: `{"path":"README.md"}`,
},
}},
},
{Role: RoleTool, ToolCallID: "call_1", Content: "contents"},
{Role: RoleAssistant, Content: "done"},
},
Tools: []Tool{{
Type: "function",
Function: ToolFunction{
Name: "file_read",
Description: "Read a local file",
Parameters: json.RawMessage(`{"type":"object","properties":{}}`),
},
}},
ToolChoice: "auto",
Extra: map[string]any{
"reasoning": "middle",
"thinking": "high",
"temperature": 0.2,
},
}
payload := buildResponsesRequest(req)
data, err := json.Marshal(payload)
if err != nil {
t.Fatal(err)
}
var raw map[string]any
if err := json.Unmarshal(data, &raw); err != nil {
t.Fatal(err)
}
if raw["model"] != "gpt-test" || raw["stream"] != true {
t.Fatalf("model/stream = %#v / %#v", raw["model"], raw["stream"])
}
if raw["instructions"] != "be concise" {
t.Fatalf("instructions = %#v", raw["instructions"])
}
if raw["tool_choice"] != "auto" {
t.Fatalf("tool_choice = %#v", raw["tool_choice"])
}
input, ok := raw["input"].([]any)
if !ok || len(input) != 5 {
t.Fatalf("input = %#v", raw["input"])
}
checkInputItem := func(index int, wantType, wantKey, wantValue string) {
t.Helper()
item, ok := input[index].(map[string]any)
if !ok {
t.Fatalf("input[%d] = %#v", index, input[index])
}
if item[wantKey] != wantValue {
t.Fatalf("input[%d].%s = %#v, want %q", index, wantKey, item[wantKey], wantValue)
}
if wantType != "" && item["type"] != wantType {
t.Fatalf("input[%d].type = %#v, want %q", index, item["type"], wantType)
}
}
checkInputItem(0, "", "role", "user")
checkInputItem(1, "", "role", "assistant")
checkInputItem(2, "function_call", "call_id", "call_1")
checkInputItem(3, "function_call_output", "call_id", "call_1")
checkInputItem(4, "", "role", "assistant")
if item := input[2].(map[string]any); item["name"] != "file_read" || item["arguments"] != `{"path":"README.md"}` {
t.Fatalf("function_call item = %#v", item)
}
if item := input[3].(map[string]any); item["output"] != "contents" {
t.Fatalf("function_call_output item = %#v", item)
}
tools, ok := raw["tools"].([]any)
if !ok || len(tools) != 1 {
t.Fatalf("tools = %#v", raw["tools"])
}
tool := tools[0].(map[string]any)
if tool["type"] != "function" || tool["name"] != "file_read" || tool["description"] != "Read a local file" {
t.Fatalf("tool = %#v", tool)
}
if _, ok := tool["parameters"].(map[string]any); !ok {
t.Fatalf("tool parameters = %#v", tool["parameters"])
}
reasoning, ok := raw["reasoning"].(map[string]any)
if !ok || reasoning["effort"] != "medium" {
t.Fatalf("reasoning = %#v", raw["reasoning"])
}
if raw["thinking"] != "high" {
t.Fatalf("thinking = %#v", raw["thinking"])
}
if raw["temperature"] != 0.2 {
t.Fatalf("temperature = %#v", raw["temperature"])
}
}
func TestResponsesInstructionsCombinesSystemMessages(t *testing.T) {
payload := buildResponsesRequest(ChatRequest{
Model: "gpt-test",
Messages: []Message{
{Role: RoleSystem, Content: "first instruction"},
{Role: RoleUser, Content: "hi"},
{Role: RoleSystem, Content: "second instruction"},
},
})
data, err := json.Marshal(payload)
if err != nil {
t.Fatal(err)
}
var raw map[string]any
if err := json.Unmarshal(data, &raw); err != nil {
t.Fatal(err)
}
if raw["instructions"] != "first instruction\n\nsecond instruction" {
t.Fatalf("instructions = %#v", raw["instructions"])
}
input, ok := raw["input"].([]any)
if !ok || len(input) != 1 {
t.Fatalf("input = %#v", raw["input"])
}
item, ok := input[0].(map[string]any)
if !ok || item["role"] != "user" {
t.Fatalf("input[0] = %#v", input[0])
}
}
func TestOpenAIResponsesClientStreamsContentAndToolCalls(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/responses" {
t.Fatalf("path = %s", r.URL.Path)
}
if got := r.Header.Get("Authorization"); got != "Bearer sk-test" {
t.Fatalf("authorization = %q", got)
}
body := mustReadBody(t, r)
var raw map[string]any
if err := json.NewDecoder(strings.NewReader(body)).Decode(&raw); err != nil {
t.Fatal(err)
}
if raw["stream"] != true {
t.Fatal("request did not enable stream")
}
if raw["tool_choice"] != "auto" {
t.Fatalf("tool_choice = %#v", raw["tool_choice"])
}
input, ok := raw["input"].([]any)
if !ok || len(input) == 0 {
t.Fatalf("input = %#v", raw["input"])
}
w.Header().Set("Content-Type", "text/event-stream")
events := []string{
`{"type":"response.output_item.added","output_index":0,"item":{"type":"message","id":"msg_1","role":"assistant","content":[]}}`,
`{"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hel"}`,
`{"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"lo"}`,
`{"type":"response.output_item.added","output_index":1,"item":{"type":"function_call","id":"fc_1","call_id":"call_1","name":"file_read","arguments":""}}`,
`{"type":"response.function_call_arguments.delta","item_id":"fc_1","output_index":1,"delta":"{\"path\""}`,
`{"type":"response.function_call_arguments.delta","item_id":"fc_1","output_index":1,"delta":":\"README.md\"}"}`,
`{"type":"response.function_call_arguments.done","item_id":"fc_1","output_index":1,"name":"file_read","arguments":"{\"path\":\"README.md\"}"}`,
`{"type":"response.output_item.done","output_index":1,"item":{"type":"function_call","id":"fc_1","call_id":"call_1","name":"file_read","arguments":"{\"path\":\"README.md\"}"}}`,
`{"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_1","role":"assistant","content":[{"type":"output_text","text":"Hello"}]}}`,
`{"type":"response.completed","response":{"id":"resp_1","status":"completed","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}}`,
}
for _, event := range events {
_, _ = w.Write([]byte("data: " + event + "\n\n"))
}
}))
defer server.Close()
client := NewOpenAIResponsesClient(server.URL, "sk-test", server.Client())
var content strings.Builder
var calls []ToolCallDelta
var finishes []string
var usage *Usage
err := client.ChatStream(context.Background(), ChatRequest{
Model: "gpt-test",
Messages: []Message{{Role: RoleUser, Content: "hi"}},
ToolChoice: "auto",
Tools: []Tool{{
Type: "function",
Function: ToolFunction{
Name: "file_read",
Description: "Read a local file",
Parameters: json.RawMessage(`{"type":"object"}`),
},
}},
}, func(event StreamEvent) error {
content.WriteString(event.Content)
calls = append(calls, event.ToolCalls...)
if event.Usage != nil {
usage = event.Usage
}
if event.FinishReason != "" {
finishes = append(finishes, event.FinishReason)
}
return nil
})
if err != nil {
t.Fatal(err)
}
if content.String() != "Hello" {
t.Fatalf("content = %q", content.String())
}
if len(calls) != 1 {
t.Fatalf("calls = %#v", calls)
}
call := calls[0]
if call.Index != 1 || call.ID != "call_1" || call.Name != "file_read" {
t.Fatalf("call = %#v", call)
}
if call.Arguments != `{"path":"README.md"}` {
t.Fatalf("arguments = %q", call.Arguments)
}
if len(finishes) != 1 || finishes[0] != "stop" {
t.Fatalf("finishes = %#v", finishes)
}
if usage == nil || usage.PromptTokens != 10 || usage.CompletionTokens != 20 || usage.TotalTokens != 30 {
t.Fatalf("usage = %#v", usage)
}
}
func TestOpenAIResponsesClientErrorEvent(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"error","code":"invalid_request_error","message":"bad request"}` + "\n\n"))
}))
defer server.Close()
client := NewOpenAIResponsesClient(server.URL, "sk-test", server.Client())
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(StreamEvent) error {
return nil
})
if err == nil || !strings.Contains(err.Error(), "bad request") {
t.Fatalf("err = %v", err)
}
}
func TestOpenAIResponsesClientFailedEvent(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"server_error","message":"boom"}}}` + "\n\n"))
}))
defer server.Close()
client := NewOpenAIResponsesClient(server.URL, "sk-test", server.Client())
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(StreamEvent) error {
return nil
})
if err == nil || !strings.Contains(err.Error(), "boom") {
t.Fatalf("err = %v", err)
}
}
func TestOpenAIResponsesClientIncomplete(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"partial"}` + "\n\n"))
_, _ = w.Write([]byte(`data: {"type":"response.incomplete","response":{"id":"resp_1","status":"incomplete","incomplete_details":{"reason":"max_tokens"}}}` + "\n\n"))
}))
defer server.Close()
client := NewOpenAIResponsesClient(server.URL, "sk-test", server.Client())
var content strings.Builder
var finishes []string
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
content.WriteString(event.Content)
if event.FinishReason != "" {
finishes = append(finishes, event.FinishReason)
}
return nil
})
if err != nil {
t.Fatal(err)
}
if content.String() != "partial" {
t.Fatalf("content = %q", content.String())
}
if len(finishes) != 1 || finishes[0] != "length" {
t.Fatalf("finishes = %#v", finishes)
}
}
func TestOpenAIResponsesClientRetriesTransientHTTPError(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: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"ok"}` + "\n\n"))
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed"}}` + "\n\n"))
}))
defer server.Close()
client := NewOpenAIResponsesClient(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())
}
}
+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 {