3 Commits
38 changed files with 5491 additions and 211 deletions
+1
View File
@@ -3,3 +3,4 @@ agentu*.yaml
/coverage.out /coverage.out
*.swp *.swp
chat-history.txt chat-history.txt
dist/
+59 -18
View File
@@ -1,19 +1,24 @@
# agentu # 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 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. - TUI-first chat experience with a light theme by default.
- OpenAI-compatible `/v1/chat/completions` provider support. - OpenAI-compatible `/v1/chat/completions` provider support.
- OpenAI Responses API (`/v1/responses`) provider support. - OpenAI Responses API (`/v1/responses`) provider support.
- Multiple providers from `~/.agentu/config.yaml`. - Multiple providers from `~/.agentu/config.yaml`.
- Session-only provider, model, and thinking-level switching. - Session-only provider, model, and thinking-level switching.
- Read-only project tools enabled by default. - 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. - Built-in web search and URL fetching.
- `--yolo` mode for file writes and shell execution. - `--yolo` mode for automatic file writes and shell execution.
- Temporary `/btw` side questions that can run while the main task continues.
## Quick Start ## Quick Start
@@ -30,7 +35,9 @@ Use dark mode when your terminal is dark:
go run . --theme dark go run . --theme dark
``` ```
Use `--yolo` only when you want agentu to write files or run shell commands: 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 ```sh
go run . --yolo go run . --yolo
@@ -74,9 +81,10 @@ responses OpenAI Responses API /v1/responses
``` ```
The Responses API client converts the conversation history to Responses input The Responses API client converts the conversation history to Responses input
items (`user`/`assistant`/`system` messages, `function_call`, and items (`user`/`assistant` messages, `function_call`, and
`function_call_output`), streams `response.output_text.delta` events, and `function_call_output`), sends the system prompt as the top-level
supports function calling through `response.output_item.added`, `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`. `response.function_call_arguments.delta`, and `response.output_item.done`.
`models` is required. `/model model <name>` is restricted to that list. Each `models` is required. `/model model <name>` is restricted to that list. Each
@@ -98,10 +106,11 @@ none, middle, high, xhigh, max
``` ```
When configured or changed with `/model thinking ...`, the value is sent as a When configured or changed with `/model thinking ...`, the value is sent as a
top-level request field named `thinking`. For OpenAI reasoning models on the top-level request field. The default field depends on `api_type`: `chat` uses
Responses API, set `thinking_param: reasoning` so the level is sent as the `thinking`, while `responses` uses the official `reasoning` field, sent as
official `reasoning: {"effort": ...}` field; agentu maps its `middle` level to `{"effort": ...}` with agentu's `middle` level mapped to the official `medium`
the official `medium` value and passes `none`, `high`, `xhigh`, `max` through. value (`none`, `high`, `xhigh`, `max` pass through). Override with
`thinking_param` if a provider expects a different field name.
## CLI Flags ## CLI Flags
@@ -112,7 +121,7 @@ go run . [flags]
- `--config <path>` loads a YAML config file. Default: `~/.agentu/config.yaml`. - `--config <path>` loads a YAML config file. Default: `~/.agentu/config.yaml`.
- `--theme light|dark` selects the TUI theme. Default: `light`. - `--theme light|dark` selects the TUI theme. Default: `light`.
- `--plain` uses the simple line-based REPL instead of the TUI. - `--plain` uses the simple line-based REPL instead of the TUI.
- `--yolo` enables automatic file writes and shell execution. - `--yolo` automatically approves model-requested file writes and shell execution.
- `--resume <id>` resumes a saved session from `~/.agentu/sessions`. - `--resume <id>` resumes a saved session from `~/.agentu/sessions`.
## TUI Controls ## TUI Controls
@@ -123,16 +132,23 @@ go run . [flags]
- The input starts at one line and grows up to six lines. - The input starts at one line and grows up to six lines.
- `page up` / `page down` scrolls the transcript. - `page up` / `page down` scrolls the transcript.
- `esc` cancels a running response. - `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. - Type `/` to show available commands.
- In `--yolo` mode, type `! <command>` to run a shell command from the configured working directory. - In `--yolo` mode, type `! <command>` to run a shell command from the configured working directory.
Slash commands: 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. - `/clear` resets conversation history.
- `/compact` summarizes older conversation history and keeps the newest ~20% of context, with a small recent-message floor. - `/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`. - `/status` shows changed files with `git status --short`.
- `/diff [--stat] [path...]` shows the unstaged git diff. - `/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. - `/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` shows current provider, model, thinking, and switch usage.
- `/model provider <name>` switches provider for the current session. - `/model provider <name>` switches provider for the current session.
- `/model model <name>` or `/model <name>` switches model for the current session. - `/model model <name>` or `/model <name>` switches model for the current session.
@@ -146,26 +162,36 @@ Slash commands:
## Local Tools ## Local Tools
Read-only tools are available by default: These tools are always available:
- `file_read` - `file_read`
- `file_list` - `file_list`
- `file_search` with optional `context_lines` and `max_results` - `file_search` with optional `context_lines` and `max_results`
- `code_symbols` for Go packages, imports, types, functions, and methods - `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_status` shows changed files
- `git_diff` shows unstaged or staged diff output - `git_diff` shows unstaged or staged diff output
- `git_log` shows recent commit history - `git_log` shows recent commit history
Web tools are also available in read-only mode: Web tools are also available:
- `web_search` - `web_search`
- `web_fetch` - `web_fetch`
`--yolo` additionally enables: In the TUI, these tools require approval by default and run automatically with
`--yolo`:
- `file_write` for creating or replacing whole files - `file_write` for creating or replacing whole files
- `file_edit` for surgical line-range or pattern edits to existing files - `file_edit` for surgical line-range or pattern edits to existing files
- `shell_run` for non-interactive shell commands - `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 Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
example: example:
@@ -174,7 +200,20 @@ example:
请帮我运行 `go test ./...` 并总结结果 请帮我运行 `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 ## Development
@@ -197,3 +236,5 @@ Local config and build outputs are intentionally ignored by git.
API type. API type.
- The config schema is intentionally strict during MVP development. Deprecated - The config schema is intentionally strict during MVP development. Deprecated
fields such as `provider` or `active_provider` are rejected. 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.
+308 -36
View File
@@ -24,6 +24,8 @@ const (
toolLogArgumentLimit = 400 toolLogArgumentLimit = 400
toolLogOutputPreviewBytes = 4096 toolLogOutputPreviewBytes = 4096
toolLogOutputMarker = "[output]" 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" + 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" + "- Key facts, decisions, and conclusions\n" +
@@ -40,6 +42,7 @@ type toolExecutionResult struct {
} }
type Agent struct { type Agent struct {
mu sync.RWMutex
provider llm.Provider provider llm.Provider
providerName string providerName string
model string model string
@@ -48,6 +51,8 @@ type Agent struct {
thinkingEnabled bool thinkingEnabled bool
systemPrompt string systemPrompt string
toolRegistry *tools.Registry toolRegistry *tools.Registry
authorizer ToolAuthorizer
rollbackStore *tools.RollbackStore
toolTimeout time.Duration toolTimeout time.Duration
maxToolRounds int maxToolRounds int
maxContextTokens int maxContextTokens int
@@ -64,6 +69,8 @@ type Options struct {
ThinkingEnabled bool ThinkingEnabled bool
SystemPrompt string SystemPrompt string
ToolRegistry *tools.Registry ToolRegistry *tools.Registry
Authorizer ToolAuthorizer
RollbackStore *tools.RollbackStore
ToolTimeout time.Duration ToolTimeout time.Duration
MaxToolIterations int MaxToolIterations int
MaxContextTokens int MaxContextTokens int
@@ -87,6 +94,8 @@ func New(opts Options) *Agent {
thinkingEnabled: opts.ThinkingEnabled, thinkingEnabled: opts.ThinkingEnabled,
systemPrompt: opts.SystemPrompt, systemPrompt: opts.SystemPrompt,
toolRegistry: opts.ToolRegistry, toolRegistry: opts.ToolRegistry,
authorizer: opts.Authorizer,
rollbackStore: opts.RollbackStore,
toolTimeout: toolTimeout, toolTimeout: toolTimeout,
maxToolRounds: maxToolRounds, maxToolRounds: maxToolRounds,
maxContextTokens: opts.MaxContextTokens, maxContextTokens: opts.MaxContextTokens,
@@ -115,6 +124,8 @@ type RuntimeInfo struct {
} }
func (a *Agent) SetRuntime(opts RuntimeOptions) { func (a *Agent) SetRuntime(opts RuntimeOptions) {
a.mu.Lock()
defer a.mu.Unlock()
if opts.Provider != nil { if opts.Provider != nil {
a.provider = opts.Provider a.provider = opts.Provider
} }
@@ -131,6 +142,8 @@ func (a *Agent) SetRuntime(opts RuntimeOptions) {
} }
func (a *Agent) RuntimeInfo() RuntimeInfo { func (a *Agent) RuntimeInfo() RuntimeInfo {
a.mu.RLock()
defer a.mu.RUnlock()
return RuntimeInfo{ return RuntimeInfo{
ProviderName: a.providerName, ProviderName: a.providerName,
Model: a.model, Model: a.model,
@@ -142,6 +155,12 @@ func (a *Agent) RuntimeInfo() RuntimeInfo {
} }
func (a *Agent) Clear() { func (a *Agent) Clear() {
a.mu.Lock()
defer a.mu.Unlock()
a.clear()
}
func (a *Agent) clear() {
a.lastUsage = nil a.lastUsage = nil
a.messages = nil a.messages = nil
if strings.TrimSpace(a.systemPrompt) != "" { if strings.TrimSpace(a.systemPrompt) != "" {
@@ -152,8 +171,27 @@ 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. // Messages returns a copy of the current conversation messages.
func (a *Agent) Messages() []llm.Message { func (a *Agent) Messages() []llm.Message {
a.mu.RLock()
defer a.mu.RUnlock()
out := make([]llm.Message, len(a.messages)) out := make([]llm.Message, len(a.messages))
copy(out, a.messages) copy(out, a.messages)
return out return out
@@ -161,11 +199,23 @@ func (a *Agent) Messages() []llm.Message {
// SetMessages replaces the conversation messages (used for session restore). // SetMessages replaces the conversation messages (used for session restore).
func (a *Agent) SetMessages(messages []llm.Message) { func (a *Agent) SetMessages(messages []llm.Message) {
a.mu.Lock()
defer a.mu.Unlock()
a.messages = append([]llm.Message(nil), messages...) 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. // LastUsage returns a copy of the most recent provider token usage data.
func (a *Agent) LastUsage() *llm.Usage { func (a *Agent) LastUsage() *llm.Usage {
a.mu.RLock()
defer a.mu.RUnlock()
if a.lastUsage == nil { if a.lastUsage == nil {
return nil return nil
} }
@@ -175,6 +225,12 @@ func (a *Agent) LastUsage() *llm.Usage {
// SetLastUsage restores the most recent provider token usage data. // SetLastUsage restores the most recent provider token usage data.
func (a *Agent) SetLastUsage(usage *llm.Usage) { 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 { if usage == nil {
a.lastUsage = nil a.lastUsage = nil
return return
@@ -182,6 +238,7 @@ func (a *Agent) SetLastUsage(usage *llm.Usage) {
copy := *usage copy := *usage
a.lastUsage = &copy a.lastUsage = &copy
} }
// compactRetainedTokenBudget returns the number of estimated tokens to retain // compactRetainedTokenBudget returns the number of estimated tokens to retain
// based on compactRecentRetentionPercent, using integer ceiling division. // based on compactRecentRetentionPercent, using integer ceiling division.
func compactRetainedTokenBudget(totalTokens int) int { func compactRetainedTokenBudget(totalTokens int) int {
@@ -237,31 +294,35 @@ func compactRecentStart(messages []llm.Message, prefixEnd int) int {
return compactTurnBoundary(messages, prefixEnd, start) return compactTurnBoundary(messages, prefixEnd, start)
} }
// Compact summarizes older conversation messages while preserving the system // Compact summarizes older conversation messages while preserving the system
// prompt and recent messages intact. // prompt and recent messages intact.
func (a *Agent) Compact(ctx context.Context) error { func (a *Agent) Compact(ctx context.Context) error {
if len(a.messages) < 4 { 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 return nil
} }
if a.provider == nil { if provider == nil {
return fmt.Errorf("provider is not configured") return fmt.Errorf("provider is not configured")
} }
prefixEnd := 0 prefixEnd := 0
if a.messages[0].Role == llm.RoleSystem { if messages[0].Role == llm.RoleSystem {
prefixEnd = 1 prefixEnd = 1
} }
recentStart := compactRecentStart(a.messages, prefixEnd) recentStart := compactRecentStart(messages, prefixEnd)
if recentStart <= prefixEnd || recentStart >= len(a.messages) { if recentStart <= prefixEnd || recentStart >= len(messages) {
return nil return nil
} }
middle := append([]llm.Message(nil), a.messages[prefixEnd:recentStart]...) middle := append([]llm.Message(nil), messages[prefixEnd:recentStart]...)
if len(middle) == 0 { if len(middle) == 0 {
return nil return nil
} }
summary, usage, err := a.summarizeMessages(ctx, middle) summary, usage, err := summarizeMessages(ctx, provider, model, middle)
if err != nil { if err != nil {
return fmt.Errorf("compact context: %w", err) return fmt.Errorf("compact context: %w", err)
} }
@@ -270,13 +331,15 @@ func (a *Agent) Compact(ctx context.Context) error {
return fmt.Errorf("compact context: summary was empty") return fmt.Errorf("compact context: summary was empty")
} }
compressed := make([]llm.Message, 0, prefixEnd+1+len(a.messages[recentStart:])) compressed := make([]llm.Message, 0, prefixEnd+1+len(messages[recentStart:]))
compressed = append(compressed, a.messages[:prefixEnd]...) compressed = append(compressed, messages[:prefixEnd]...)
compressed = append(compressed, llm.Message{ compressed = append(compressed, llm.Message{
Role: llm.RoleAssistant, Role: llm.RoleAssistant,
Content: contextSummaryPrefix + "\n" + summary, Content: contextSummaryPrefix + "\n" + summary,
}) })
compressed = append(compressed, a.messages[recentStart:]...) compressed = append(compressed, messages[recentStart:]...)
a.mu.Lock()
defer a.mu.Unlock()
a.messages = compressed a.messages = compressed
if usage != nil { if usage != nil {
a.lastUsage = usage a.lastUsage = usage
@@ -285,6 +348,8 @@ func (a *Agent) Compact(ctx context.Context) error {
} }
func (a *Agent) shouldCompact() bool { func (a *Agent) shouldCompact() bool {
a.mu.RLock()
defer a.mu.RUnlock()
if a.maxContextTokens <= 0 { if a.maxContextTokens <= 0 {
return false return false
} }
@@ -294,9 +359,9 @@ func (a *Agent) shouldCompact() bool {
return llm.EstimateTokens(a.messages, a.toolDefinitions()) >= a.maxContextTokens return llm.EstimateTokens(a.messages, a.toolDefinitions()) >= a.maxContextTokens
} }
func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (string, *llm.Usage, error) { func summarizeMessages(ctx context.Context, provider llm.Provider, model string, messages []llm.Message) (string, *llm.Usage, error) {
req := llm.ChatRequest{ req := llm.ChatRequest{
Model: a.model, Model: model,
Messages: []llm.Message{ Messages: []llm.Message{
{Role: llm.RoleSystem, Content: summarizerSystemPrompt}, {Role: llm.RoleSystem, Content: summarizerSystemPrompt},
{Role: llm.RoleUser, Content: formatMessagesForSummary(messages)}, {Role: llm.RoleUser, Content: formatMessagesForSummary(messages)},
@@ -305,7 +370,7 @@ func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (
var summary strings.Builder var summary strings.Builder
var usage *llm.Usage var usage *llm.Usage
if err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error { if err := provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
if event.Content != "" { if event.Content != "" {
summary.WriteString(event.Content) summary.WriteString(event.Content)
} }
@@ -353,12 +418,93 @@ func roundSignature(calls []llm.ToolCall) string {
return b.String() 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 { func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
input = strings.TrimSpace(input) input = strings.TrimSpace(input)
if input == "" { if input == "" {
return nil return nil
} }
a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input}) a.appendMessage(llm.Message{Role: llm.RoleUser, Content: input})
if a.shouldCompact() { if a.shouldCompact() {
if err := a.Compact(ctx); err != nil { if err := a.Compact(ctx); err != nil {
if logs != nil { if logs != nil {
@@ -375,14 +521,14 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
return err return err
} }
if len(calls) == 0 { if len(calls) == 0 {
a.messages = append(a.messages, llm.Message{ a.appendMessage(llm.Message{
Role: llm.RoleAssistant, Role: llm.RoleAssistant,
Content: content, Content: content,
}) })
return nil return nil
} }
a.messages = append(a.messages, llm.Message{ a.appendMessage(llm.Message{
Role: llm.RoleAssistant, Role: llm.RoleAssistant,
Content: content, Content: content,
ToolCalls: calls, ToolCalls: calls,
@@ -417,7 +563,7 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
strings.Join(toolNames, ", "), strings.Join(toolNames, ", "),
doomLoopThreshold, doomLoopThreshold,
) )
a.messages = append(a.messages, llm.Message{ a.appendMessage(llm.Message{
Role: llm.RoleAssistant, Role: llm.RoleAssistant,
Content: msg, Content: msg,
}) })
@@ -427,7 +573,7 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
results := a.executeTools(ctx, calls, logs) results := a.executeTools(ctx, calls, logs)
for _, result := range results { for _, result := range results {
a.messages = append(a.messages, llm.Message{ a.appendMessage(llm.Message{
Role: llm.RoleTool, Role: llm.RoleTool,
ToolCallID: result.call.ID, ToolCallID: result.call.ID,
Content: result.output, Content: result.output,
@@ -443,23 +589,33 @@ func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []l
var content strings.Builder var content strings.Builder
calls := make(map[int]*toolCallBuilder) 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{ req := llm.ChatRequest{
Model: a.model, Model: model,
Messages: append([]llm.Message(nil), a.messages...), Messages: expandContinuationMessages(messages),
Tools: a.toolDefinitions(), Tools: a.toolDefinitions(),
} }
if a.thinkingEnabled { if thinkingEnabled {
thinkingParam := a.thinkingParam
if thinkingParam == "" { if thinkingParam == "" {
thinkingParam = "thinking" thinkingParam = "thinking"
} }
req.Extra = map[string]any{thinkingParam: a.thinking} req.Extra = map[string]any{thinkingParam: thinking}
} }
if len(req.Tools) > 0 { if len(req.Tools) > 0 {
req.ToolChoice = "auto" 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 { if event.Usage != nil {
a.SetLastUsage(event.Usage) a.SetLastUsage(event.Usage)
} }
@@ -495,32 +651,144 @@ func (a *Agent) toolDefinitions() []llm.Tool {
return a.toolRegistry.Definitions() return a.toolRegistry.Definitions()
} }
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall) string { func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) (string, bool) {
name := call.Function.Name name := call.Function.Name
if a.toolRegistry == nil { 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) tool, ok := a.toolRegistry.Get(name)
if !ok { 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) args := strings.TrimSpace(call.Function.Arguments)
if args == "" { if args == "" {
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) toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
defer cancel() 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 err != nil {
if output != "" { 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
} }
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 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 { func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
@@ -529,7 +797,11 @@ func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
} }
for _, call := range calls { for _, call := range calls {
if tool, ok := a.toolRegistry.Get(call.Function.Name); ok { if tool, ok := a.toolRegistry.Get(call.Function.Name); ok {
if tools.IsMutating(tool) { args := strings.TrimSpace(call.Function.Arguments)
if args == "" {
args = "{}"
}
if tools.IsMutatingCall(tool, json.RawMessage(args)) {
return true return true
} }
} }
@@ -550,9 +822,9 @@ func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.
if a.hasMutatingCall(calls) { if a.hasMutatingCall(calls) {
// Sequential: avoid racing mutating tools. // Sequential: avoid racing mutating tools.
for i, call := range calls { for i, call := range calls {
output := a.executeTool(ctx, call) output, streamed := a.executeTool(ctx, call, logs)
results[i] = toolExecutionResult{call: call, output: output} results[i] = toolExecutionResult{call: call, output: output}
if logs != nil { if logs != nil && !streamed {
logToolOutput(logs, call, output) logToolOutput(logs, call, output)
} }
} }
@@ -566,9 +838,9 @@ func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.
wg.Add(1) wg.Add(1)
go func(idx int, c llm.ToolCall) { go func(idx int, c llm.ToolCall) {
defer wg.Done() defer wg.Done()
output := a.executeTool(ctx, c) output, streamed := a.executeTool(ctx, c, logs)
results[idx] = toolExecutionResult{call: c, output: output} results[idx] = toolExecutionResult{call: c, output: output}
if logs != nil { if logs != nil && !streamed {
mu.Lock() mu.Lock()
logToolOutput(logs, c, output) logToolOutput(logs, c, output)
mu.Unlock() mu.Unlock()
+366
View File
@@ -3,8 +3,13 @@ package agent
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"os"
"path/filepath"
"reflect"
"strings" "strings"
"sync"
"testing" "testing"
"agentu/internal/tools" "agentu/internal/tools"
@@ -87,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 { type contentProvider struct {
called bool called bool
text string text string
@@ -97,6 +182,173 @@ func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, e
return emit(llm.StreamEvent{Content: p.text}) return emit(llm.StreamEvent{Content: p.text})
} }
type 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 { type compactProvider struct {
summaryCalls int summaryCalls int
answerCalls int answerCalls int
@@ -549,6 +801,120 @@ type parallelProvider struct {
calls int 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 { func (p *parallelProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.calls++ p.calls++
if p.calls <= len(p.responses) { if p.calls <= len(p.responses) {
+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)
}
+32 -11
View File
@@ -13,6 +13,7 @@ import (
"agentu/internal/agent" "agentu/internal/agent"
"agentu/internal/session" "agentu/internal/session"
"agentu/internal/skills"
"agentu/internal/tools" "agentu/internal/tools"
"agentu/internal/tui" "agentu/internal/tui"
"agentu/pkg/config" "agentu/pkg/config"
@@ -40,8 +41,20 @@ func Run() error {
return err 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) providerName, providerConfig, provider := defaultProvider(cfg, http.DefaultClient)
assistant := newAssistant(cfg, providerName, providerConfig, provider, opts.yolo, http.DefaultClient) assistant := newAssistant(cfg, providerName, providerConfig, provider, opts, globalSkills, authorizer, rollbackStore, http.DefaultClient)
modelManager, err := newSessionManager(cfg, assistant, http.DefaultClient) modelManager, err := newSessionManager(cfg, assistant, http.DefaultClient)
if err != nil { if err != nil {
return err return err
@@ -54,7 +67,7 @@ func Run() error {
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
defer stop() defer stop()
return runInterface(ctx, assistant, modelManager, providerConfig, cfg.Agent.WorkingDir, opts, themeMode) return runInterface(ctx, assistant, modelManager, providerConfig, cfg.Agent.WorkingDir, opts, themeMode, approvalBroker, rollbackStore)
} }
func parseCLIOptions() cliOptions { func parseCLIOptions() cliOptions {
@@ -102,7 +115,8 @@ func newProviderClient(providerConfig config.ProviderConfig, httpClient *http.Cl
return llm.NewOpenAICompatibleClient(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, yolo bool, httpClient *http.Client) *agent.Agent { 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{ return agent.New(agent.Options{
Provider: provider, Provider: provider,
ProviderName: providerName, ProviderName: providerName,
@@ -110,16 +124,21 @@ func newAssistant(cfg *config.Config, providerName string, providerConfig config
Thinking: providerConfig.Thinking, Thinking: providerConfig.Thinking,
ThinkingParam: providerConfig.ThinkingParam, ThinkingParam: providerConfig.ThinkingParam,
ThinkingEnabled: providerConfig.ThinkingConfigured, ThinkingEnabled: providerConfig.ThinkingConfigured,
SystemPrompt: systemPrompt(cfg, yolo), SystemPrompt: systemPrompt(cfg, writesEnabled, opts.yolo, globalSkills),
ToolRegistry: newToolRegistry(cfg.Agent.WorkingDir, yolo, httpClient), ToolRegistry: newToolRegistry(cfg.Agent.WorkingDir, writesEnabled, httpClient),
Authorizer: authorizer,
RollbackStore: rollbackStore,
ToolTimeout: cfg.Agent.ToolTimeout, ToolTimeout: cfg.Agent.ToolTimeout,
MaxContextTokens: maxContextTokens(cfg, providerConfig), MaxContextTokens: maxContextTokens(cfg, providerConfig),
}) })
} }
func systemPrompt(cfg *config.Config, yolo bool) string { func systemPrompt(cfg *config.Config, writesEnabled bool, yolo bool, globalSkills []skills.Skill) string {
instructions := tools.AgentInstructions(cfg.Agent.WorkingDir, yolo) parts := []string{cfg.Agent.SystemPrompt, tools.AgentInstructionsForMode(cfg.Agent.WorkingDir, writesEnabled, yolo)}
return strings.TrimSpace(cfg.Agent.SystemPrompt + "\n\n" + instructions) 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 { func maxContextTokens(cfg *config.Config, providerConfig config.ProviderConfig) int {
@@ -129,9 +148,9 @@ func maxContextTokens(cfg *config.Config, providerConfig config.ProviderConfig)
return providerConfig.ContextTokens return providerConfig.ContextTokens
} }
func newToolRegistry(workingDir string, yolo bool, httpClient *http.Client) *tools.Registry { func newToolRegistry(workingDir string, writesEnabled bool, httpClient *http.Client) *tools.Registry {
var registry *tools.Registry var registry *tools.Registry
if yolo { if writesEnabled {
registry = tools.Builtins(workingDir) registry = tools.Builtins(workingDir)
} else { } else {
registry = tools.ReadOnlyBuiltins(workingDir) registry = tools.ReadOnlyBuiltins(workingDir)
@@ -154,7 +173,7 @@ func newSessionManager(cfg *config.Config, assistant *agent.Agent, httpClient *h
return session.NewManager(cfg, assistant, httpClient, store) 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) error { 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 { if opts.plain {
return repl(ctx, assistant, modelManager) return repl(ctx, assistant, modelManager)
} }
@@ -165,6 +184,8 @@ func runInterface(ctx context.Context, assistant *agent.Agent, modelManager *ses
ThemeMode: themeMode, ThemeMode: themeMode,
ModelManager: modelManager, ModelManager: modelManager,
WorkingDir: workingDir, WorkingDir: workingDir,
Approval: approvalBroker,
Rollback: rollbackStore,
}) })
printExitSession(modelManager) printExitSession(modelManager)
return err return err
+12
View File
@@ -3,10 +3,12 @@ package app
import ( import (
"context" "context"
"path/filepath" "path/filepath"
"strings"
"testing" "testing"
"agentu/internal/agent" "agentu/internal/agent"
"agentu/internal/session" "agentu/internal/session"
"agentu/internal/skills"
"agentu/pkg/config" "agentu/pkg/config"
"agentu/pkg/llm" "agentu/pkg/llm"
) )
@@ -99,3 +101,13 @@ func TestInitializeSessionResumesExplicitID(t *testing.T) {
t.Fatal("explicit resume did not restore saved context") 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)
}
}
}
+36 -1
View File
@@ -9,6 +9,7 @@ import (
"agentu/internal/agent" "agentu/internal/agent"
"agentu/internal/session" "agentu/internal/session"
"agentu/internal/skills"
) )
func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error { func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error {
@@ -16,7 +17,7 @@ func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Man
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
fmt.Println("agentu") fmt.Println("agentu")
fmt.Println("Type /exit to quit, /clear to reset context, /compact to compress context.") fmt.Println("Type /exit to quit, /clear to reset context, /compact to compress context, /btw <question> for a side question.")
for { for {
fmt.Print("> ") fmt.Print("> ")
if !scanner.Scan() { if !scanner.Scan() {
@@ -29,6 +30,7 @@ func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Man
} }
input := strings.TrimSpace(scanner.Text()) input := strings.TrimSpace(scanner.Text())
fields := strings.Fields(input)
switch input { switch input {
case "": case "":
continue continue
@@ -50,6 +52,39 @@ func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Man
} }
continue 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 { if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
fmt.Fprintln(os.Stderr, "\nerror:", err) fmt.Fprintln(os.Stderr, "\nerror:", 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
}
+23 -1
View File
@@ -68,11 +68,29 @@ func (m *Manager) Resume(id string) error {
return err return err
} }
m.currentSession = sess m.currentSession = sess
m.restoreRuntime(sess)
m.agent.SetMessages(sess.Messages) m.agent.SetMessages(sess.Messages)
m.agent.SetLastUsage(sess.LastUsage) m.agent.SetLastUsage(sess.LastUsage)
return nil 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. // Save persists the current session to disk.
func (m *Manager) Save() error { func (m *Manager) Save() error {
if m.store == nil || m.currentSession == nil { if m.store == nil || m.currentSession == nil {
@@ -82,6 +100,8 @@ func (m *Manager) Save() error {
m.currentSession.LastUsage = m.agent.LastUsage() m.currentSession.LastUsage = m.agent.LastUsage()
m.currentSession.Provider = m.providerName m.currentSession.Provider = m.providerName
m.currentSession.Model = m.provider.Model m.currentSession.Model = m.provider.Model
m.currentSession.APIType = m.provider.APIType
m.currentSession.Thinking = m.provider.Thinking
// Auto-name if empty // Auto-name if empty
if m.currentSession.Name == "" { if m.currentSession.Name == "" {
m.currentSession.Name = DefaultName(m.currentSession.Messages) m.currentSession.Name = DefaultName(m.currentSession.Messages)
@@ -150,6 +170,8 @@ func (m *Manager) NewSession() (string, error) {
m.currentSession.Messages = msgs m.currentSession.Messages = msgs
m.currentSession.Provider = m.providerName m.currentSession.Provider = m.providerName
m.currentSession.Model = m.provider.Model m.currentSession.Model = m.provider.Model
m.currentSession.APIType = m.provider.APIType
m.currentSession.Thinking = m.provider.Thinking
if m.currentSession.Name == "" { if m.currentSession.Name == "" {
m.currentSession.Name = DefaultName(msgs) m.currentSession.Name = DefaultName(msgs)
} }
@@ -257,7 +279,7 @@ func (m *Manager) SetThinking(level string) (string, error) {
m.provider.Thinking = level m.provider.Thinking = level
m.provider.ThinkingConfigured = true m.provider.ThinkingConfigured = true
if m.provider.ThinkingParam == "" { if m.provider.ThinkingParam == "" {
m.provider.ThinkingParam = "thinking" m.provider.ThinkingParam = config.DefaultThinkingParam(m.provider.APIType)
} }
m.cfg.Providers[m.providerName] = m.provider m.cfg.Providers[m.providerName] = m.provider
m.syncAgent(nil) m.syncAgent(nil)
+129
View File
@@ -403,3 +403,132 @@ func TestManagerUsesModelContextForAgentCompaction(t *testing.T) {
t.Fatalf("answerCalls = %d, want 1", provider.answerCalls) 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)
}
}
+83 -6
View File
@@ -11,6 +11,7 @@ import (
"strings" "strings"
"time" "time"
"agentu/internal/safeio"
"agentu/pkg/llm" "agentu/pkg/llm"
) )
@@ -22,6 +23,8 @@ type Session struct {
UpdatedAt time.Time `json:"updated_at"` UpdatedAt time.Time `json:"updated_at"`
Provider string `json:"provider"` Provider string `json:"provider"`
Model string `json:"model"` Model string `json:"model"`
APIType string `json:"api_type,omitempty"`
Thinking string `json:"thinking,omitempty"`
Messages []llm.Message `json:"messages"` Messages []llm.Message `json:"messages"`
LastUsage *llm.Usage `json:"last_usage,omitempty"` LastUsage *llm.Usage `json:"last_usage,omitempty"`
} }
@@ -52,8 +55,11 @@ func NewStore(dir string) (*Store, error) {
// Save writes a session to disk. // Save writes a session to disk.
func (s *Store) Save(sess *Session) error { func (s *Store) Save(sess *Session) error {
if sess.ID == "" { if sess == nil {
return fmt.Errorf("session ID is required") return fmt.Errorf("session is required")
}
if err := validateSessionID(sess.ID); err != nil {
return err
} }
sess.UpdatedAt = time.Now().UTC() sess.UpdatedAt = time.Now().UTC()
if sess.CreatedAt.IsZero() { if sess.CreatedAt.IsZero() {
@@ -64,14 +70,17 @@ func (s *Store) Save(sess *Session) error {
return fmt.Errorf("marshal session: %w", err) return fmt.Errorf("marshal session: %w", err)
} }
path := s.filePath(sess.ID) path := s.filePath(sess.ID)
if err := os.WriteFile(path, data, 0o644); err != nil { if err := safeio.AtomicWriteFile(path, data, 0o644); err != nil {
return fmt.Errorf("write session: %w", err) return fmt.Errorf("write session: %w", err)
} }
return nil return s.updateIndex(metaFromSession(sess))
} }
// Load reads a session by ID. // Load reads a session by ID.
func (s *Store) Load(id string) (*Session, error) { func (s *Store) Load(id string) (*Session, error) {
if err := validateSessionID(id); err != nil {
return nil, err
}
path := s.filePath(id) path := s.filePath(id)
data, err := os.ReadFile(path) data, err := os.ReadFile(path)
if err != nil { if err != nil {
@@ -89,13 +98,27 @@ func (s *Store) Load(id string) (*Session, error) {
// List returns metadata for all sessions, sorted by updated_at descending. // List returns metadata for all sessions, sorted by updated_at descending.
func (s *Store) List() ([]SessionMeta, error) { 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) entries, err := os.ReadDir(s.dir)
if err != nil { if err != nil {
return nil, fmt.Errorf("read session dir: %w", err) return nil, fmt.Errorf("read session dir: %w", err)
} }
var metas []SessionMeta var metas []SessionMeta
for _, entry := range entries { for _, entry := range entries {
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") { if entry.IsDir() || entry.Name() == sessionIndexFile || !strings.HasSuffix(entry.Name(), ".json") {
continue continue
} }
path := filepath.Join(s.dir, entry.Name()) path := filepath.Join(s.dir, entry.Name())
@@ -125,6 +148,9 @@ func (s *Store) List() ([]SessionMeta, error) {
// Delete removes a session file. // Delete removes a session file.
func (s *Store) Delete(id string) error { func (s *Store) Delete(id string) error {
if err := validateSessionID(id); err != nil {
return err
}
path := s.filePath(id) path := s.filePath(id)
if err := os.Remove(path); err != nil { if err := os.Remove(path); err != nil {
if os.IsNotExist(err) { if os.IsNotExist(err) {
@@ -132,13 +158,64 @@ func (s *Store) Delete(id string) error {
} }
return fmt.Errorf("delete session: %w", err) return fmt.Errorf("delete session: %w", err)
} }
return nil 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 { func (s *Store) filePath(id string) string {
return filepath.Join(s.dir, id+".json") 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. // GenerateID creates a random 8-character hex session ID.
func GenerateID() string { func GenerateID() string {
b := make([]byte, 4) b := make([]byte, 4)
+40
View File
@@ -119,6 +119,9 @@ func TestDeleteNotFound(t *testing.T) {
func TestSaveRequiresID(t *testing.T) { func TestSaveRequiresID(t *testing.T) {
store := tempStore(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"}) err := store.Save(&Session{Name: "no id"})
if err == nil { if err == nil {
t.Fatal("expected error for empty ID") t.Fatal("expected error for empty ID")
@@ -172,3 +175,40 @@ func TestNewStoreCreatesDir(t *testing.T) {
t.Error("expected directory") 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...)
}
+44 -16
View File
@@ -13,6 +13,7 @@ import (
"sort" "sort"
"strings" "strings"
"agentu/internal/safeio"
"agentu/pkg/llm" "agentu/pkg/llm"
) )
@@ -83,6 +84,33 @@ func NewFileWriteTool(workingDir string) *FileWriteTool {
} }
func (*FileWriteTool) Mutates() bool { return true } 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 { func (t *FileWriteTool) Definition() llm.Tool {
return llm.Tool{ return llm.Tool{
Type: "function", Type: "function",
@@ -117,28 +145,28 @@ func (t *FileWriteTool) Execute(ctx context.Context, raw json.RawMessage) (strin
if strings.TrimSpace(args.Path) == "" { if strings.TrimSpace(args.Path) == "" {
return "", errors.New("path is required") return "", errors.New("path is required")
} }
path, err := resolvePath(t.workingDir, args.Path) path, err := resolvePathForWrite(t.workingDir, args.Path)
if err != nil { if err != nil {
return "", err 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 { if err := ctx.Err(); err != nil {
return "", err 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" action := "wrote"
if args.Append { if args.Append {
action = "appended" action = "appended"
+32 -1
View File
@@ -8,6 +8,7 @@ import (
"os" "os"
"strings" "strings"
"agentu/internal/safeio"
"agentu/pkg/llm" "agentu/pkg/llm"
) )
@@ -21,6 +22,36 @@ func NewFileEditTool(workingDir string) *FileEditTool {
func (*FileEditTool) Mutates() bool { return true } 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 { func (t *FileEditTool) Definition() llm.Tool {
return llm.Tool{ return llm.Tool{
Type: "function", Type: "function",
@@ -111,7 +142,7 @@ func (t *FileEditTool) Execute(ctx context.Context, raw json.RawMessage) (string
if err := ctx.Err(); err != nil { if err := ctx.Err(); err != nil {
return "", err return "", err
} }
if err := os.WriteFile(path, []byte(result), perm); err != nil { if err := safeio.AtomicWriteFile(path, []byte(result), perm); err != nil {
return "", err return "", 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
}
+68 -5
View File
@@ -2,22 +2,85 @@ package tools
import ( import (
"fmt" "fmt"
"os"
"path/filepath" "path/filepath"
"strings" "strings"
) )
func resolvePath(baseDir, path string) (string, error) { 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 == "" { if path == "" {
path = "." path = "."
} }
if filepath.IsAbs(path) { if filepath.IsAbs(path) {
return "", fmt.Errorf("absolute paths are not allowed; use a path relative to the project directory") return "", "", fmt.Errorf("absolute paths are not allowed; use a path relative to the project directory")
} }
resolved := filepath.Clean(filepath.Join(baseDir, path)) absBase, err := filepath.Abs(baseDir)
if !isWithinDir(resolved, baseDir) { if err != nil {
return "", fmt.Errorf("path escapes the project directory: %s", path) return "", "", err
} }
return resolved, nil 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 { func isWithinDir(path, dir string) bool {
+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
}
+51 -5
View File
@@ -1,6 +1,7 @@
package tools package tools
import ( import (
"bytes"
"context" "context"
"encoding/json" "encoding/json"
"errors" "errors"
@@ -8,6 +9,7 @@ import (
"os" "os"
"os/exec" "os/exec"
"strings" "strings"
"sync"
"agentu/pkg/llm" "agentu/pkg/llm"
) )
@@ -21,6 +23,17 @@ func NewShellRunTool(workingDir string) *ShellRunTool {
} }
func (*ShellRunTool) Mutates() bool { return true } 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 { func (t *ShellRunTool) Definition() llm.Tool {
return llm.Tool{ return llm.Tool{
Type: "function", Type: "function",
@@ -46,6 +59,10 @@ type shellRunArgs struct {
} }
func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) { 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 var args shellRunArgs
if err := json.Unmarshal(raw, &args); err != nil { if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err) return "", fmt.Errorf("parse arguments: %w", err)
@@ -66,16 +83,45 @@ func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string
cmd := exec.CommandContext(ctx, shell, "-lc", args.Command) cmd := exec.CommandContext(ctx, shell, "-lc", args.Command)
cmd.Dir = workdir cmd.Dir = workdir
output, err := cmd.CombinedOutput() var mu sync.Mutex
text := FormatResult(string(output)) 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) { if errors.Is(ctx.Err(), context.DeadlineExceeded) {
return text, fmt.Errorf("command timed out: %w", ctx.Err()) return text, fmt.Errorf("command timed out: %w", ctx.Err())
} }
if err != nil { if runErr != nil {
if text != "" { 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 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
}
+55 -5
View File
@@ -19,18 +19,54 @@ type Tool interface {
Execute(ctx context.Context, args json.RawMessage) (string, error) 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. // 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. // It is used by the agent to determine whether tool calls can run in parallel.
type MutatingTool interface { type MutatingTool interface {
Mutates() bool 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. // IsMutating reports whether tool implements MutatingTool and returns true.
func IsMutating(tool Tool) bool { func IsMutating(tool Tool) bool {
mutating, ok := tool.(MutatingTool) mutating, ok := tool.(MutatingTool)
return ok && mutating.Mutates() 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 { type Registry struct {
tools map[string]Tool tools map[string]Tool
primary []string primary []string
@@ -81,8 +117,10 @@ func ReadOnlyBuiltins(workingDir string) *Registry {
gitStatus := NewGitStatusTool(workingDir) gitStatus := NewGitStatusTool(workingDir)
gitDiff := NewGitDiffTool(workingDir) gitDiff := NewGitDiffTool(workingDir)
gitLog := NewGitLogTool(workingDir) gitLog := NewGitLogTool(workingDir)
projectOutline := NewProjectOutlineTool(workingDir)
todo := NewTodoTool()
registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols, gitStatus, gitDiff, gitLog) registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols, projectOutline, todo, gitStatus, gitDiff, gitLog)
registry.RegisterAlias("file.read", fileRead) registry.RegisterAlias("file.read", fileRead)
registry.RegisterAlias("file.list", fileList) registry.RegisterAlias("file.list", fileList)
registry.RegisterAlias("file.search", fileSearch) registry.RegisterAlias("file.search", fileSearch)
@@ -97,36 +135,48 @@ func Builtins(workingDir string) *Registry {
registry := ReadOnlyBuiltins(workingDir) registry := ReadOnlyBuiltins(workingDir)
fileWrite := NewFileWriteTool(workingDir) fileWrite := NewFileWriteTool(workingDir)
fileEdit := NewFileEditTool(workingDir) fileEdit := NewFileEditTool(workingDir)
applyPatch := NewApplyPatchTool(workingDir)
shellRun := NewShellRunTool(workingDir) shellRun := NewShellRunTool(workingDir)
registry.Register(fileWrite) registry.Register(fileWrite)
registry.Register(fileEdit) registry.Register(fileEdit)
registry.Register(applyPatch)
registry.Register(shellRun) registry.Register(shellRun)
registry.RegisterAlias("file.write", fileWrite) registry.RegisterAlias("file.write", fileWrite)
registry.RegisterAlias("file.edit", fileEdit) registry.RegisterAlias("file.edit", fileEdit)
registry.RegisterAlias("apply.patch", applyPatch)
registry.RegisterAlias("shell.run", shellRun) registry.RegisterAlias("shell.run", shellRun)
return registry return registry
} }
func AgentInstructions(workingDir string, yolo bool) string { func AgentInstructions(workingDir string, yolo bool) string {
return AgentInstructionsForMode(workingDir, yolo, yolo)
}
func AgentInstructionsForMode(workingDir string, writesEnabled bool, yolo bool) string {
searchInstructions := webSearchInstructions() searchInstructions := webSearchInstructions()
readOnlyTools := "file_list, file_read, file_search, code_symbols, git_status, git_diff, git_log, web_search, web_fetch" readOnlyTools := "file_list, file_read, file_search, code_symbols, project_outline, todo, git_status, git_diff, git_log, web_search, web_fetch"
yoloTools := "file_list, file_read, file_search, code_symbols, file_write, file_edit, shell_run, 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 { if yolo {
approval = "Mutating tools execute automatically because --yolo is enabled."
}
return fmt.Sprintf(`Local project context: return fmt.Sprintf(`Local project context:
- The project working directory is %q. - The project working directory is %q.
- Tools are available, but use them only when needed to satisfy the user's request. - 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 call tools for greetings, small talk, simple Q&A, or general questions that do not require local project or command context.
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files. - Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files. - When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files. - When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
- When the user asks to run, execute, test, check, debug, or inspect a shell command, use shell_run. This includes local commands written in backticks, such as go, make, git, pwd, ls, ps, and similar terminal commands. - When the user asks to run, execute, test, check, debug, or inspect a shell command, use shell_run. This includes local commands written in backticks, such as go, make, git, pwd, ls, ps, and similar terminal commands.
- When editing existing files, prefer file_edit for small or localized changes; use file_write only for creating files, replacing whole files intentionally, or appending. - 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. - 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. - 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. - Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions.
- %s - %s
- Tool paths should normally be relative to the project working directory. - Tool paths should normally be relative to the project working directory.
- Available tools: %s.`, workingDir, searchInstructions, yoloTools) - Available tools: %s.`, workingDir, approval, searchInstructions, writeTools)
} }
return fmt.Sprintf(`Local project context: return fmt.Sprintf(`Local project context:
- The project working directory is %q. - The project working directory is %q.
+3 -3
View File
@@ -104,7 +104,7 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
} }
} }
for _, want := range []string{"file_list", "file_read", "file_search", "code_symbols", "git_status", "git_diff", "git_log"} { 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) { if !slices.Contains(names, want) {
t.Fatalf("definitions missing %s: %#v", want, names) t.Fatalf("definitions missing %s: %#v", want, names)
} }
@@ -173,7 +173,7 @@ func TestRegistryCanExposeWebSearch(t *testing.T) {
func TestBuiltinsExposeYoloTools(t *testing.T) { func TestBuiltinsExposeYoloTools(t *testing.T) {
registry := Builtins(t.TempDir()) registry := Builtins(t.TempDir())
for _, name := range []string{"file_write", "file_edit", "shell_run", "file.write", "file.edit", "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 { if _, ok := registry.Get(name); !ok {
t.Fatalf("tool missing: %s", name) t.Fatalf("tool missing: %s", name)
} }
@@ -198,7 +198,7 @@ func TestAgentInstructionsExplainCommandModes(t *testing.T) {
func TestAgentInstructionsExplainWebSearch(t *testing.T) { func TestAgentInstructionsExplainWebSearch(t *testing.T) {
withSearch := AgentInstructions("/tmp/project", false) withSearch := AgentInstructions("/tmp/project", false)
for _, want := range []string{"use web_search", "Use web_fetch", "latest information", "source URLs", "Available tools: file_list, file_read, file_search, code_symbols, git_status, git_diff, git_log, web_search, web_fetch"} { for _, want := range []string{"use web_search", "Use web_fetch", "latest information", "source URLs", "Available tools: file_list, file_read, file_search, code_symbols, project_outline, todo, git_status, git_diff, git_log, web_search, web_fetch"} {
if !strings.Contains(withSearch, want) { if !strings.Contains(withSearch, want) {
t.Fatalf("web search instructions missing %q:\n%s", want, withSearch) t.Fatalf("web search instructions missing %q:\n%s", want, withSearch)
} }
+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)
}
}
+414 -44
View File
@@ -2,6 +2,7 @@ package tui
import ( import (
"context" "context"
"encoding/json"
"errors" "errors"
"fmt" "fmt"
"io" "io"
@@ -9,8 +10,11 @@ import (
"sort" "sort"
"strings" "strings"
"unicode" "unicode"
"unicode/utf8"
"agentu/internal/agent" "agentu/internal/agent"
"agentu/internal/skills"
"agentu/internal/tools"
"agentu/pkg/llm" "agentu/pkg/llm"
"github.com/charmbracelet/bubbles/spinner" "github.com/charmbracelet/bubbles/spinner"
@@ -33,6 +37,8 @@ type Options struct {
ThemeMode ThemeMode ThemeMode ThemeMode
ModelManager ModelManager ModelManager ModelManager
WorkingDir string WorkingDir string
Approval *ApprovalBroker
Rollback *tools.RollbackStore
} }
type ModelManager interface { type ModelManager interface {
@@ -67,6 +73,8 @@ const (
roleTool role = "tool" roleTool role = "tool"
roleError role = "error" roleError role = "error"
roleSystem role = "system" roleSystem role = "system"
roleBTWUser role = "btw-user"
roleBTW role = "btw"
) )
type message struct { type message struct {
@@ -81,6 +89,8 @@ type model struct {
yolo bool yolo bool
models ModelManager models ModelManager
workingDir string workingDir string
approvalBroker *ApprovalBroker
rollback *tools.RollbackStore
theme theme theme theme
styles styles styles styles
@@ -101,7 +111,10 @@ type model struct {
running bool running bool
cancel context.CancelFunc cancel context.CancelFunc
events <-chan tea.Msg btwRunning bool
btwCancel context.CancelFunc
events chan tea.Msg
approval *toolApprovalMsg
// Session picker state // Session picker state
picking bool picking bool
@@ -117,11 +130,15 @@ type pickerItem struct {
} }
type assistantChunkMsg string type assistantChunkMsg string
type btwChunkMsg string
type toolLogMsg string type toolLogMsg string
type turnStartedMsg struct{} type turnStartedMsg struct{}
type turnDoneMsg struct { type turnDoneMsg struct {
err error err error
} }
type btwDoneMsg struct {
err error
}
var errOpenPicker = fmt.Errorf("open picker") var errOpenPicker = fmt.Errorf("open picker")
@@ -142,7 +159,10 @@ type commandSuggestion struct {
executable bool executable bool
} }
var btwSlashCommand = slashCommand{Name: "/btw", Description: "ask a temporary side question"}
var slashCommands = []slashCommand{ var slashCommands = []slashCommand{
btwSlashCommand,
{Name: "/clear", Description: "reset context"}, {Name: "/clear", Description: "reset context"},
{Name: "/compact", Description: "compress context"}, {Name: "/compact", Description: "compress context"},
{Name: "/exit", Description: "quit"}, {Name: "/exit", Description: "quit"},
@@ -154,6 +174,9 @@ var slashCommands = []slashCommand{
{Name: "/new", Description: "new session"}, {Name: "/new", Description: "new session"},
{Name: "/status", Description: "show changed files"}, {Name: "/status", Description: "show changed files"},
{Name: "/diff", Description: "show unstaged diff"}, {Name: "/diff", Description: "show unstaged diff"},
{Name: "/test", Description: "run project tests"},
{Name: "/skills", Description: "list or install global skills"},
{Name: "/rollback", Description: "restore the latest file snapshot"},
{Name: "/theme", Description: "switch theme light|dark"}, {Name: "/theme", Description: "switch theme light|dark"},
} }
@@ -190,6 +213,8 @@ func newModel(ctx context.Context, assistant *agent.Agent, opts Options) model {
yolo: opts.Yolo, yolo: opts.Yolo,
models: opts.ModelManager, models: opts.ModelManager,
workingDir: opts.WorkingDir, workingDir: opts.WorkingDir,
approvalBroker: opts.Approval,
rollback: opts.Rollback,
theme: th, theme: th,
styles: st, styles: st,
viewport: vp, viewport: vp,
@@ -242,7 +267,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, nil return m, nil
case spinner.TickMsg: case spinner.TickMsg:
if !m.running { if !m.active() {
return m, nil return m, nil
} }
var cmd tea.Cmd var cmd tea.Cmd
@@ -250,6 +275,9 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, cmd return m, cmd
case tea.KeyMsg: case tea.KeyMsg:
if m.approval != nil {
return m.updateApproval(msg)
}
if m.picking { if m.picking {
return m.updatePicker(msg) return m.updatePicker(msg)
} }
@@ -261,31 +289,30 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
if m.cancel != nil { if m.cancel != nil {
m.cancel() m.cancel()
} }
if m.btwCancel != nil {
m.btwCancel()
}
return m, tea.Quit return m, tea.Quit
case "esc": case "esc":
if m.running && m.cancel != nil { if m.cancel != nil {
m.cancel() m.cancel()
}
if m.btwCancel != nil {
m.btwCancel()
}
if m.active() {
m.status = "canceling..." m.status = "canceling..."
return m, nil return m, nil
} }
case "enter": case "enter":
if m.running {
return m, nil
}
return m.submit() return m.submit()
case "up": case "up":
if m.running {
return m, nil
}
if m.input.Line() == 0 { if m.input.Line() == 0 {
m.previousInputHistory() m.previousInputHistory()
m.afterInputChanged() m.afterInputChanged()
return m, nil return m, nil
} }
case "down": case "down":
if m.running {
return m, nil
}
value := m.input.Value() value := m.input.Value()
lineCount := 1 lineCount := 1
if value != "" { if value != "" {
@@ -297,16 +324,10 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, nil return m, nil
} }
case "ctrl+p": case "ctrl+p":
if m.running {
return m, nil
}
m.previousInputHistory() m.previousInputHistory()
m.afterInputChanged() m.afterInputChanged()
return m, nil return m, nil
case "ctrl+n": case "ctrl+n":
if m.running {
return m, nil
}
m.nextInputHistory() m.nextInputHistory()
m.afterInputChanged() m.afterInputChanged()
return m, nil return m, nil
@@ -320,6 +341,14 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.refreshViewport(true) m.refreshViewport(true)
return m, m.waitForEvent() return m, m.waitForEvent()
case btwChunkMsg:
m.appendBTWChunk(string(msg))
if !m.running {
m.status = "answering btw..."
}
m.refreshViewport(true)
return m, m.waitForEvent()
case toolLogMsg: case toolLogMsg:
text := cleanToolLog(string(msg)) text := cleanToolLog(string(msg))
if text != "" { if text != "" {
@@ -331,13 +360,25 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
} }
return m, m.waitForEvent() return m, m.waitForEvent()
case toolApprovalMsg:
m.approval = &msg
m.status = "approval required"
m.resize()
m.refreshViewport(true)
return m, m.waitForEvent()
case turnDoneMsg: case turnDoneMsg:
m.running = false m.running = false
m.events = nil
m.cancel = nil m.cancel = nil
m.approval = nil
if m.approvalBroker != nil {
m.approvalBroker.ClearEvents()
}
if msg.err != nil { if msg.err != nil {
m.messages = append(m.messages, message{role: roleError, content: msg.err.Error()}) m.messages = append(m.messages, message{role: roleError, content: msg.err.Error()})
m.status = "error" m.status = "error"
} else if m.btwRunning {
m.status = "answering btw..."
} else { } else {
m.status = "ready" m.status = "ready"
} }
@@ -349,6 +390,29 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.resize() m.resize()
m.refreshViewport(true) m.refreshViewport(true)
cmds = append(cmds, textarea.Blink) cmds = append(cmds, textarea.Blink)
if m.active() {
cmds = append(cmds, m.waitForEvent())
} else {
m.events = nil
}
case btwDoneMsg:
m.btwRunning = false
m.btwCancel = nil
if msg.err != nil && !errors.Is(msg.err, context.Canceled) {
m.messages = append(m.messages, message{role: roleBTW, content: "Error: " + msg.err.Error()})
}
if !m.running {
m.status = "ready"
}
m.resize()
m.refreshViewport(true)
cmds = append(cmds, textarea.Blink)
if m.active() {
cmds = append(cmds, m.waitForEvent())
} else {
m.events = nil
}
} }
if keyMsg, ok := msg.(tea.KeyMsg); ok { if keyMsg, ok := msg.(tea.KeyMsg); ok {
@@ -357,7 +421,6 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.viewport = nextViewport m.viewport = nextViewport
return m, cmd return m, cmd
} }
if !m.running {
nextInput, cmd := m.input.Update(msg) nextInput, cmd := m.input.Update(msg)
m.input = nextInput m.input = nextInput
m.historyIndex = -1 m.historyIndex = -1
@@ -365,10 +428,8 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.afterInputChanged() m.afterInputChanged()
return m, cmd return m, cmd
} }
return m, nil
}
if !m.running { if !m.active() {
nextInput, cmd := m.input.Update(msg) nextInput, cmd := m.input.Update(msg)
m.input = nextInput m.input = nextInput
cmds = append(cmds, cmd) cmds = append(cmds, cmd)
@@ -392,6 +453,15 @@ func (m model) View() string {
body := lipgloss.JoinVertical(lipgloss.Left, parts...) body := lipgloss.JoinVertical(lipgloss.Left, parts...)
return m.styles.App.Width(m.width).Height(m.height).Render(body) return m.styles.App.Width(m.width).Height(m.height).Render(body)
} }
if m.approval != nil {
parts := []string{m.viewport.View()}
if activity := m.activityView(); activity != "" {
parts = append(parts, activity)
}
parts = append(parts, m.approvalView(), m.footerView())
body := lipgloss.JoinVertical(lipgloss.Left, parts...)
return m.styles.App.Width(m.width).Height(m.height).Render(body)
}
parts := []string{m.viewport.View()} parts := []string{m.viewport.View()}
if activity := m.activityView(); activity != "" { if activity := m.activityView(); activity != "" {
@@ -407,6 +477,14 @@ func (m *model) submit() (tea.Model, tea.Cmd) {
if input == "" { if input == "" {
return *m, nil return *m, nil
} }
fields := strings.Fields(input)
if len(fields) > 0 && fields[0] == "/btw" {
m.rememberInput(input)
return m.submitBTW(strings.TrimSpace(strings.TrimPrefix(input, "/btw")))
}
if m.active() {
return *m, nil
}
m.rememberInput(input) m.rememberInput(input)
switch input { switch input {
@@ -446,20 +524,19 @@ func (m *model) submit() (tea.Model, tea.Cmd) {
events := make(chan tea.Msg, 256) events := make(chan tea.Msg, 256)
m.events = events m.events = events
if m.approvalBroker != nil {
m.approvalBroker.SetEvents(events)
}
start := func() tea.Msg { start := func() tea.Msg {
go func() { go func() {
defer close(events)
out := eventWriter{ctx: turnCtx, events: events, assistant: true} out := eventWriter{ctx: turnCtx, events: events, assistant: true}
logs := eventWriter{ctx: turnCtx, events: events} logs := eventWriter{ctx: turnCtx, events: events}
err := m.agent.RunTurn(turnCtx, input, out, logs) err := m.agent.RunTurn(turnCtx, input, out, logs)
if errors.Is(err, context.Canceled) { if errors.Is(err, context.Canceled) {
err = nil err = nil
} }
select { events <- turnDoneMsg{err: err}
case events <- turnDoneMsg{err: err}:
case <-turnCtx.Done():
}
}() }()
return turnStartedMsg{} return turnStartedMsg{}
} }
@@ -467,6 +544,54 @@ func (m *model) submit() (tea.Model, tea.Cmd) {
return *m, tea.Batch(start, m.waitForEvent(), m.spinner.Tick) return *m, tea.Batch(start, m.waitForEvent(), m.spinner.Tick)
} }
func (m *model) submitBTW(question string) (tea.Model, tea.Cmd) {
if question == "" {
m.messages = append(m.messages, message{role: roleBTW, content: "Usage: /btw <question>"})
m.input.SetValue("")
m.afterInputChanged()
m.refreshViewport(true)
return *m, nil
}
if m.agent == nil {
m.messages = append(m.messages, message{role: roleBTW, content: "Error: agent is not configured"})
return *m, nil
}
if m.btwRunning {
m.messages = append(m.messages, message{role: roleBTW, content: "A /btw question is already running."})
return *m, nil
}
hadActive := m.active()
m.messages = append(m.messages, message{role: roleBTWUser, content: question})
m.input.SetValue("")
m.afterInputChanged()
m.btwRunning = true
if !m.running {
m.status = "answering btw..."
}
m.resize()
m.refreshViewport(true)
btwCtx, cancel := context.WithCancel(m.ctx)
m.btwCancel = cancel
if m.events == nil {
m.events = make(chan tea.Msg, 256)
}
events := m.events
start := func() tea.Msg {
go func() {
out := eventWriter{ctx: btwCtx, events: events, btw: true}
err := m.agent.BTW(btwCtx, question, out)
events <- btwDoneMsg{err: err}
}()
return turnStartedMsg{}
}
if hadActive {
return *m, start
}
return *m, tea.Batch(start, m.waitForEvent(), m.spinner.Tick)
}
func (m *model) insertInputNewline() (tea.Model, tea.Cmd) { func (m *model) insertInputNewline() (tea.Model, tea.Cmd) {
enter := tea.KeyMsg{Type: tea.KeyEnter} enter := tea.KeyMsg{Type: tea.KeyEnter}
nextInput, cmd := m.input.Update(enter) nextInput, cmd := m.input.Update(enter)
@@ -505,10 +630,38 @@ func (m *model) appendAssistantChunk(chunk string) {
if chunk == "" { if chunk == "" {
return return
} }
if len(m.messages) == 0 || m.messages[len(m.messages)-1].role != roleAssistant { for i := len(m.messages) - 1; i >= 0; i-- {
m.messages = append(m.messages, message{role: roleAssistant}) switch m.messages[i].role {
case roleAssistant:
m.messages[i].content += chunk
return
case roleBTW, roleBTWUser:
continue
default:
i = -1
} }
m.messages[len(m.messages)-1].content += chunk }
m.messages = append(m.messages, message{role: roleAssistant, content: chunk})
}
func (m *model) appendBTWChunk(chunk string) {
if chunk == "" {
return
}
for i := len(m.messages) - 1; i >= 0; i-- {
if m.messages[i].role == roleBTW {
m.messages[i].content += chunk
return
}
if m.messages[i].role == roleBTWUser {
break
}
}
m.messages = append(m.messages, message{role: roleBTW, content: chunk})
}
func (m model) active() bool {
return m.running || m.btwRunning
} }
func (m *model) resize() { func (m *model) resize() {
@@ -518,10 +671,13 @@ func (m *model) resize() {
m.syncInputHeight() m.syncInputHeight()
activityHeight := 0 activityHeight := 0
if m.running { if m.active() {
activityHeight = 1 activityHeight = 1
} }
inputHeight := m.inputBlockHeight() inputHeight := m.inputBlockHeight()
if m.approval != nil {
inputHeight = lipgloss.Height(m.approvalView())
}
footerHeight := 2 footerHeight := 2
m.viewport.Width = width m.viewport.Width = width
m.viewport.Height = max(4, m.height-activityHeight-inputHeight-footerHeight) m.viewport.Height = max(4, m.height-activityHeight-inputHeight-footerHeight)
@@ -564,6 +720,12 @@ func (m model) messageView(r role, content string, width int) string {
switch r { switch r {
case roleUser: case roleUser:
style = m.styles.UserMessage style = m.styles.UserMessage
case roleBTWUser:
style = m.styles.UserMessage
labelStyle = m.styles.SystemLabel
case roleBTW:
style = m.styles.AssistantMsg
labelStyle = m.styles.SystemLabel
case roleTool: case roleTool:
style = m.styles.ToolMessage style = m.styles.ToolMessage
labelStyle = m.styles.ToolLabel labelStyle = m.styles.ToolLabel
@@ -576,7 +738,7 @@ func (m model) messageView(r role, content string, width int) string {
} }
bodyContent := content bodyContent := content
if r == roleAssistant { if r == roleAssistant || r == roleBTW {
bodyWidth := max(minMarkdownWidth, width-style.GetHorizontalFrameSize()) bodyWidth := max(minMarkdownWidth, width-style.GetHorizontalFrameSize())
bodyContent = renderMarkdown(content, bodyWidth, m.theme) bodyContent = renderMarkdown(content, bodyWidth, m.theme)
} }
@@ -614,6 +776,10 @@ func roleLabel(r role) string {
return "Error" return "Error"
case roleSystem: case roleSystem:
return "Note" return "Note"
case roleBTWUser:
return "BTW · You"
case roleBTW:
return "BTW"
default: default:
return string(r) return string(r)
} }
@@ -649,6 +815,11 @@ func parseToolLog(content string) toolLogParts {
func (m model) toolMessageView(content string, width int) string { func (m model) toolMessageView(content string, width int) string {
parts := parseToolLog(content) parts := parseToolLog(content)
if isShellRun(parts.name) {
if args, ok := parseShellRunDisplayArgs(parts.args); ok {
return m.shellRunMessageView(args, parts.output, width)
}
}
header := m.styles.ToolName.Render(strings.TrimSpace(iconTool + " " + parts.name)) header := m.styles.ToolName.Render(strings.TrimSpace(iconTool + " " + parts.name))
lines := []string{header} lines := []string{header}
if parts.args != "" { if parts.args != "" {
@@ -669,18 +840,132 @@ func (m model) toolMessageView(content string, width int) string {
const maxToolOutputPreviewLines = 8 const maxToolOutputPreviewLines = 8
type shellRunDisplayArgs struct {
Command string `json:"command"`
Workdir string `json:"workdir"`
}
func isShellRun(name string) bool {
return name == "shell_run" || name == "shell.run"
}
func parseShellRunDisplayArgs(raw string) (shellRunDisplayArgs, bool) {
var args shellRunDisplayArgs
if json.Unmarshal([]byte(raw), &args) != nil || strings.TrimSpace(args.Command) == "" {
return shellRunDisplayArgs{}, false
}
args.Command = strings.TrimSpace(args.Command)
args.Workdir = strings.TrimSpace(args.Workdir)
return args, true
}
func (m model) shellRunMessageView(args shellRunDisplayArgs, output string, width int) string {
innerWidth := max(1, width-m.styles.ToolMessage.GetHorizontalFrameSize())
lines := []string{m.styles.ToolName.Render(fitLine("$ "+args.Command, innerWidth))}
if args.Workdir != "" {
lines = append(lines, m.styles.ToolKey.Render("cwd ")+m.styles.ToolValue.Render(fitLine(args.Workdir, max(1, innerWidth-4))))
}
if output != "" {
lines = append(lines, m.styles.ToolKey.Render("output"))
for _, line := range shellOutputPreviewLines(output, innerWidth) {
lines = append(lines, m.styles.ToolValue.Render(line))
}
}
bodyContent := strings.Join(lines, "\n")
bodyWidth := messageBodyWidth(bodyContent, width, m.styles.ToolMessage)
return m.styles.ToolMessage.Width(bodyWidth).Render(bodyContent)
}
func shellOutputPreviewLines(output string, width int) []string {
output = sanitizeTerminalOutput(output)
lines := strings.Split(strings.TrimRight(output, "\n"), "\n")
if len(lines) > maxToolOutputPreviewLines {
head := maxToolOutputPreviewLines / 2
tail := maxToolOutputPreviewLines - head - 1
omitted := len(lines) - head - tail
lines = append(append(append([]string(nil), lines[:head]...), fmt.Sprintf("… %d lines omitted …", omitted)), lines[len(lines)-tail:]...)
}
result := make([]string, 0, len(lines))
for _, line := range lines {
result = append(result, fitLine(line, width))
}
return result
}
func sanitizeTerminalOutput(output string) string {
output = strings.ReplaceAll(output, "\r\n", "\n")
var clean strings.Builder
for i := 0; i < len(output); {
if end, ok := terminalEscapeEnd(output, i); ok {
i = end
continue
}
r, size := utf8.DecodeRuneInString(output[i:])
if r == utf8.RuneError && size == 0 {
break
}
if r == '\n' || r == '\r' || r == '\t' || r >= ' ' {
clean.WriteRune(r)
}
i += size
}
lines := strings.Split(clean.String(), "\n")
for i, line := range lines {
segments := strings.Split(line, "\r")
for j := len(segments) - 1; j >= 0; j-- {
if segments[j] != "" || j == 0 {
lines[i] = segments[j]
break
}
}
}
return strings.Join(lines, "\n")
}
func terminalEscapeEnd(value string, start int) (int, bool) {
if start+1 >= len(value) || value[start] != '\x1b' {
return 0, false
}
if value[start+1] == '[' {
if end, ok := ansiSequenceEnd(value, start); ok {
return end, true
}
return len(value), true
}
if value[start+1] == ']' {
for i := start + 2; i < len(value); i++ {
if value[i] == '\a' {
return i + 1, true
}
if value[i] == '\x1b' && i+1 < len(value) && value[i+1] == '\\' {
return i + 2, true
}
}
return len(value), true
}
return min(start+2, len(value)), true
}
func toolStatusText(content string) string { func toolStatusText(content string) string {
content = strings.TrimSpace(content) content = strings.TrimSpace(content)
if content == "" { if content == "" {
return "" return ""
} }
lines := strings.SplitN(content, "\n", 2) lines := strings.SplitN(content, "\n", 2)
parts := parseToolLog(lines[0])
if isShellRun(parts.name) {
if args, ok := parseShellRunDisplayArgs(parts.args); ok {
return "$ " + args.Command
}
}
return lines[0] return lines[0]
} }
func (m *model) mergeToolLog(text string) bool { func (m *model) mergeToolLog(text string) bool {
incoming := parseToolLog(text) incoming := parseToolLog(text)
if incoming.output == "" { incomingHeader, _, ok := strings.Cut(text, "\n[output]\n")
if !ok {
return false return false
} }
for i := len(m.messages) - 1; i >= 0; i-- { for i := len(m.messages) - 1; i >= 0; i-- {
@@ -688,10 +973,18 @@ func (m *model) mergeToolLog(text string) bool {
continue continue
} }
existing := parseToolLog(m.messages[i].content) existing := parseToolLog(m.messages[i].content)
if existing.name == incoming.name && existing.args == incoming.args && existing.output == "" { if existing.name != incoming.name || existing.args != incoming.args {
m.messages[i].content = text continue
return true
} }
if !strings.Contains(m.messages[i].content, "\n[output]\n") {
// Placeholder card from the tool-start log; first chunk fills it.
m.messages[i].content = text
} else {
// Streaming tools send incremental chunks; append the raw delta
// so line boundaries are preserved inside the same card.
m.messages[i].content += strings.TrimPrefix(text, incomingHeader+"\n[output]\n")
}
return true
} }
return false return false
} }
@@ -740,20 +1033,28 @@ func (m model) inputView() string {
} }
func (m model) activityView() string { func (m model) activityView() string {
if !m.running { if !m.active() {
return "" return ""
} }
status := strings.TrimSpace(m.status) status := strings.TrimSpace(m.status)
if status == "" || status == "ready" { if status == "" || status == "ready" {
status = "working..." status = "working..."
} }
line := m.styles.StatusInfo.Render(iconWorking) + " Working " + m.spinner.View() + " " + status + " · esc cancels" label := "Working"
if !m.running {
label = "BTW"
} else if m.btwRunning {
label = "Working + BTW"
}
line := m.styles.StatusInfo.Render(iconWorking) + " " + label + " " + m.spinner.View() + " " + status + " · esc cancels"
return m.styles.Activity.Render(fitLine(line, max(1, m.width-2))) return m.styles.Activity.Render(fitLine(line, max(1, m.width-2)))
} }
func (m model) footerView() string { func (m model) footerView() string {
hint := "" hint := ""
if m.running { if m.btwRunning {
hint = "esc cancel · ctrl+c quit" hint = "btw running · esc cancel · ctrl+c quit"
} else if m.running {
hint = "type /btw <question> · esc cancel · ctrl+c quit"
} }
if m.showCommandPalette() { if m.showCommandPalette() {
if sug, ok := m.selectedCommandSuggestion(); ok && sug.executable { if sug, ok := m.selectedCommandSuggestion(); ok && sug.executable {
@@ -835,6 +1136,7 @@ type eventWriter struct {
ctx context.Context ctx context.Context
events chan<- tea.Msg events chan<- tea.Msg
assistant bool assistant bool
btw bool
} }
func (w eventWriter) Write(p []byte) (int, error) { func (w eventWriter) Write(p []byte) (int, error) {
@@ -843,7 +1145,9 @@ func (w eventWriter) Write(p []byte) (int, error) {
} }
text := string(append([]byte(nil), p...)) text := string(append([]byte(nil), p...))
var msg tea.Msg var msg tea.Msg
if w.assistant { if w.btw {
msg = btwChunkMsg(text)
} else if w.assistant {
msg = assistantChunkMsg(text) msg = assistantChunkMsg(text)
} else { } else {
msg = toolLogMsg(text) msg = toolLogMsg(text)
@@ -857,8 +1161,12 @@ func (w eventWriter) Write(p []byte) (int, error) {
} }
func cleanToolLog(text string) string { func cleanToolLog(text string) string {
text = strings.TrimSpace(text) text = strings.TrimLeft(text, "\r\n")
text = strings.TrimPrefix(text, "[tool]") text = strings.TrimPrefix(text, "[tool]")
text = strings.TrimLeft(text, " \t")
if strings.Contains(text, "\n[output]\n") {
return text
}
return strings.TrimSpace(text) return strings.TrimSpace(text)
} }
@@ -1108,6 +1416,17 @@ func subsequenceScore(candidate string, query string, base int) (int, bool) {
func (m model) commandSuggestions() []commandSuggestion { func (m model) commandSuggestions() []commandSuggestion {
input := m.input.Value() input := m.input.Value()
if m.active() {
value := strings.TrimSpace(input)
if m.running && strings.HasPrefix(value, "/") && !strings.Contains(value, " ") {
query := commandPaletteQuery(input)
command := btwSlashCommand
if _, ok := scoreCommand(command, query); ok {
return []commandSuggestion{{label: command.Name, description: command.Description, value: command.Name}}
}
}
return nil
}
if sugs := m.modelCommandSuggestions(input); len(sugs) > 0 { if sugs := m.modelCommandSuggestions(input); len(sugs) > 0 {
return sugs return sugs
} }
@@ -1461,11 +1780,62 @@ func (m model) handleLocalCommand(input string) (string, error) {
return m.handleDiffCommand(fields[1:]) return m.handleDiffCommand(fields[1:])
case "/test": case "/test":
return m.handleTestCommand(fields[1:]) return m.handleTestCommand(fields[1:])
case "/skills":
return m.handleSkillsCommand(fields[1:])
case "/rollback":
return m.handleRollbackCommand(fields[1:])
default: default:
return "", fmt.Errorf("unknown command: %s", fields[0]) return "", fmt.Errorf("unknown command: %s", fields[0])
} }
} }
func (m model) handleSkillsCommand(args []string) (string, error) {
if len(args) == 0 || args[0] == "list" {
if len(args) > 1 {
return "", fmt.Errorf("usage: /skills [list|install <source>]")
}
items, err := skills.LoadGlobal()
if err != nil {
return "", err
}
return skills.FormatList(items), nil
}
if args[0] != "install" || len(args) < 2 {
return "", fmt.Errorf("usage: /skills [list|install <source>]")
}
result, err := skills.InstallGlobal(m.ctx, strings.Join(args[1:], " "))
if err != nil {
return "", err
}
if m.agent != nil {
m.agent.AddSystemInstructions(skills.FormatForSystemPrompt([]skills.Skill{result.Skill}))
}
return fmt.Sprintf("Installed skill %q from %s. It is active in this conversation.", result.Skill.Name, result.Skill.Path), nil
}
func (m model) handleRollbackCommand(args []string) (string, error) {
if m.rollback == nil {
return "", fmt.Errorf("rollback is not configured")
}
if len(args) > 1 {
return "", fmt.Errorf("usage: /rollback [id]")
}
id := ""
if len(args) == 1 {
id = strings.TrimSpace(args[0])
} else {
id = m.rollback.LastID()
}
if id == "" {
return "", fmt.Errorf("no rollback snapshots available")
}
snapshot, err := m.rollback.Restore(id)
if err != nil {
return "", err
}
return fmt.Sprintf("Rolled back %s (%s).", snapshot.Path, snapshot.ID), nil
}
func (m model) handleCompactCommand() (string, error) { func (m model) handleCompactCommand() (string, error) {
if m.agent == nil { if m.agent == nil {
return "", fmt.Errorf("agent is not configured") return "", fmt.Errorf("agent is not configured")
@@ -1634,7 +2004,7 @@ func (m *model) updatePicker(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
} }
func (m *model) updateCommandSuggestions(msg tea.KeyMsg) (bool, tea.Model, tea.Cmd) { func (m *model) updateCommandSuggestions(msg tea.KeyMsg) (bool, tea.Model, tea.Cmd) {
if m.running || len(m.commandSuggestions()) == 0 { if len(m.commandSuggestions()) == 0 {
return false, *m, nil return false, *m, nil
} }
m.clampCommandCursor() m.clampCommandCursor()
+318
View File
@@ -2,13 +2,17 @@ package tui
import ( import (
"context" "context"
"fmt"
"os" "os"
"os/exec" "os/exec"
"path/filepath" "path/filepath"
"reflect"
"strings" "strings"
"testing" "testing"
"time"
"agentu/internal/agent" "agentu/internal/agent"
"agentu/internal/tools"
"agentu/pkg/llm" "agentu/pkg/llm"
tea "github.com/charmbracelet/bubbletea" tea "github.com/charmbracelet/bubbletea"
@@ -126,6 +130,116 @@ func TestRunningStateRendersActivityIndicator(t *testing.T) {
} }
} }
type tuiBTWProvider struct{}
func (tuiBTWProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
return emit(llm.StreamEvent{Content: "side answer"})
}
func TestBTWCanBeSubmittedWhileMainTurnIsRunning(t *testing.T) {
assistant := agent.New(agent.Options{Provider: tuiBTWProvider{}, Model: "test-model"})
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
m.running = true
m.events = make(chan tea.Msg, 8)
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("/btw why?")})
m = next.(model)
if m.input.Value() != "/btw why?" {
t.Fatalf("running input = %q", m.input.Value())
}
next, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = next.(model)
if !m.btwRunning || m.input.Value() != "" {
t.Fatalf("btw running/input = %v/%q", m.btwRunning, m.input.Value())
}
if got := m.messages[len(m.messages)-1]; got.role != roleBTWUser || got.content != "why?" {
t.Fatalf("btw question message = %+v", got)
}
if cmd == nil {
t.Fatal("btw did not return a start command")
}
_ = cmd()
chunk := <-m.events
done := <-m.events
if _, ok := chunk.(btwChunkMsg); !ok {
t.Fatalf("first btw event = %T", chunk)
}
if _, ok := done.(btwDoneMsg); !ok {
t.Fatalf("final btw event = %T", done)
}
if len(assistant.Messages()) != 0 {
t.Fatalf("btw polluted agent history: %+v", assistant.Messages())
}
}
func TestBTWRequiresQuestion(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.input.SetValue("/btw")
next, _ := m.submit()
m = next.(model)
if m.btwRunning {
t.Fatal("empty btw question started a request")
}
last := m.messages[len(m.messages)-1]
if last.role != roleBTW || !strings.Contains(last.content, "Usage: /btw") {
t.Fatalf("empty btw message = %+v", last)
}
}
func TestBTWMessagesHaveDistinctLabels(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
question := stripANSI(m.messageView(roleBTWUser, "question", 80))
answer := stripANSI(m.messageView(roleBTW, "answer", 80))
if !strings.Contains(question, "BTW · You") || !strings.Contains(answer, "BTW") {
t.Fatalf("btw labels missing:\n%s\n%s", question, answer)
}
}
func TestBTWAndMainStreamChunksStaySeparate(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleUser, content: "main"})
m.appendAssistantChunk("main one ")
m.messages = append(m.messages, message{role: roleBTWUser, content: "side"})
m.appendBTWChunk("side one ")
m.appendAssistantChunk("main two")
m.appendBTWChunk("side two")
var mainAnswers, sideAnswers []string
for _, msg := range m.messages {
switch msg.role {
case roleAssistant:
mainAnswers = append(mainAnswers, msg.content)
case roleBTW:
sideAnswers = append(sideAnswers, msg.content)
}
}
if !reflect.DeepEqual(mainAnswers, []string{"main one main two"}) {
t.Fatalf("main answers = %#v", mainAnswers)
}
if !reflect.DeepEqual(sideAnswers, []string{"side one side two"}) {
t.Fatalf("side answers = %#v", sideAnswers)
}
}
func TestMainCompletionKeepsBTWEventStreamAlive(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
events := make(chan tea.Msg, 1)
m.events = events
m.running = true
m.btwRunning = true
next, _ := m.Update(turnDoneMsg{})
m = next.(model)
if m.running || !m.btwRunning || m.events != events {
t.Fatalf("state after main completion: running=%v btw=%v events=%p", m.running, m.btwRunning, m.events)
}
next, _ = m.Update(btwDoneMsg{})
m = next.(model)
if m.active() || m.events != nil {
t.Fatalf("state after btw completion: active=%v events=%p", m.active(), m.events)
}
}
func TestUserMessageRendersFullWidthBubble(t *testing.T) { func TestUserMessageRendersFullWidthBubble(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"}) m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30}) next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
@@ -395,6 +509,98 @@ func TestToolLogOutputUpdatesExistingToolCard(t *testing.T) {
} }
} }
func TestToolLogStreamingChunksAppendToSameCard(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleTool, content: `shell_run echo hi`})
if !m.mergeToolLog("shell_run echo hi\n[output]\none\n") {
t.Fatal("expected first mergeToolLog to return true")
}
if !m.mergeToolLog("shell_run echo hi\n[output]\ntwo\n") {
t.Fatal("expected second mergeToolLog to return true")
}
if len(m.messages) != 1 {
t.Fatalf("messages count = %d, want 1", len(m.messages))
}
plain := stripANSI(m.renderMessages())
if !strings.Contains(plain, "one") || !strings.Contains(plain, "two") {
t.Fatalf("rendered output missing streamed lines:\n%s", plain)
}
if strings.Contains(plain, "onetwo") {
t.Fatalf("streamed lines were concatenated:\n%s", plain)
}
}
func TestShellRunCardShowsCommandWorkdirAndOutput(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{
role: roleTool,
content: "shell_run {\"command\":\"go test ./...\",\"workdir\":\"internal/tui\"}\n[output]\nok agentu/internal/tui",
})
plain := stripANSI(m.renderMessages())
for _, want := range []string{"$ go test ./...", "cwd internal/tui", "output", "ok agentu/internal/tui"} {
if !strings.Contains(plain, want) {
t.Fatalf("shell card missing %q:\n%s", want, plain)
}
}
if strings.Contains(plain, `{"command"`) || strings.Contains(plain, iconTool+" shell_run") {
t.Fatalf("shell card still exposes raw tool arguments:\n%s", plain)
}
if status := toolStatusText(`shell_run {"command":"go test ./...","workdir":"internal/tui"}`); status != "$ go test ./..." {
t.Fatalf("shell status = %q", status)
}
}
func TestShellRunStreamingEventsPreserveLineBreaks(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
args := `{"command":"printf lines"}`
m.messages = append(m.messages, message{role: roleTool, content: "shell_run " + args})
next, _ := m.Update(toolLogMsg("\n[tool] shell_run " + args + "\n[output]\none\n"))
m = next.(model)
next, _ = m.Update(toolLogMsg("\n[tool] shell_run " + args + "\n[output]\n\n"))
m = next.(model)
next, _ = m.Update(toolLogMsg("\n[tool] shell_run " + args + "\n[output]\ntwo\n"))
m = next.(model)
parts := parseToolLog(m.messages[0].content)
if parts.output != "one\n\ntwo" {
t.Fatalf("streamed output = %q", parts.output)
}
if len(m.messages) != 1 {
t.Fatalf("streaming created %d tool cards, want 1", len(m.messages))
}
plain := stripANSI(m.renderMessages())
if strings.Contains(plain, "onetwo") {
t.Fatalf("streamed lines were concatenated:\n%s", plain)
}
}
func TestShellOutputPreviewKeepsHeadAndLatestLines(t *testing.T) {
var lines []string
for i := 1; i <= 12; i++ {
lines = append(lines, fmt.Sprintf("line %d", i))
}
preview := strings.Join(shellOutputPreviewLines(strings.Join(lines, "\n"), 80), "\n")
for _, want := range []string{"line 1", "line 4", "5 lines omitted", "line 10", "line 12"} {
if !strings.Contains(preview, want) {
t.Fatalf("preview missing %q:\n%s", want, preview)
}
}
if strings.Contains(preview, "line 5\n") || strings.Contains(preview, "line 9\n") {
t.Fatalf("preview contains a line that should be omitted:\n%s", preview)
}
}
func TestSanitizeTerminalOutputRemovesEscapesAndFoldsCarriageReturns(t *testing.T) {
got := sanitizeTerminalOutput("\x1b[31mred\x1b[0m\nprogress 10%\rprogress 20%\r\nclean\x00text")
want := "red\nprogress 20%\ncleantext"
if got != want {
t.Fatalf("sanitized output = %q, want %q", got, want)
}
}
type tuiCompactProvider struct{} type tuiCompactProvider struct{}
func (tuiCompactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error { func (tuiCompactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
@@ -912,3 +1118,115 @@ func TestCleanToolLog(t *testing.T) {
t.Fatalf("got %q want %q", got, want) t.Fatalf("got %q want %q", got, want)
} }
} }
func TestApprovalPromptAllowsAndDeniesToolCalls(t *testing.T) {
for _, tc := range []struct {
key tea.KeyMsg
want agent.ToolAuthorizationDecision
}{
{key: tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'y'}}, want: agent.ToolAuthorizationAllow},
{key: tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'n'}}, want: agent.ToolAuthorizationDeny},
} {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
responses := make(chan agent.ToolAuthorizationDecision, 1)
next, _ = m.Update(toolApprovalMsg{
request: agent.ToolAuthorizationRequest{ToolName: "file_write", Preview: "write 5 bytes", RollbackAvailable: true},
respond: responses,
})
m = next.(model)
view := stripANSI(m.View())
for _, want := range []string{"Approve file_write?", "write 5 bytes", "rollback snapshot", "y approve"} {
if !strings.Contains(view, want) {
t.Fatalf("approval view missing %q:\n%s", want, view)
}
}
next, _ = m.Update(tc.key)
m = next.(model)
if m.approval != nil {
t.Fatal("approval state was not cleared")
}
if got := <-responses; got != tc.want {
t.Fatalf("decision = %q, want %q", got, tc.want)
}
}
}
func TestApprovalBrokerBridgesAgentAndTUI(t *testing.T) {
broker := NewApprovalBroker()
events := make(chan tea.Msg, 1)
broker.SetEvents(events)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
result := make(chan agent.ToolAuthorizationDecision, 1)
go func() {
decision, _ := broker.AuthorizeToolCall(ctx, agent.ToolAuthorizationRequest{ToolName: "file_edit"})
result <- decision
}()
msg, ok := (<-events).(toolApprovalMsg)
if !ok {
t.Fatal("approval broker emitted unexpected message")
}
msg.respond <- agent.ToolAuthorizationAllow
if decision := <-result; decision != agent.ToolAuthorizationAllow {
t.Fatalf("decision = %q", decision)
}
}
func TestRollbackCommandRestoresLatestSnapshot(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "state.txt")
if err := os.WriteFile(path, []byte("before"), 0o644); err != nil {
t.Fatal(err)
}
rollback := tools.NewRollbackStore()
if _, err := rollback.Capture("file_edit", "state.txt", path); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("after"), 0o644); err != nil {
t.Fatal(err)
}
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Rollback: rollback})
out, err := m.handleRollbackCommand(nil)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "Rolled back state.txt") {
t.Fatalf("output = %q", out)
}
data, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(data) != "before" {
t.Fatalf("content = %q", data)
}
}
func TestSkillsInstallCommandActivatesSkill(t *testing.T) {
home := t.TempDir()
t.Setenv("HOME", home)
sourceDir := t.TempDir()
source := filepath.Join(sourceDir, "review skill.md")
content := "---\nname: review\ndescription: Review changes\n---\n\nCheck every diff.\n"
if err := os.WriteFile(source, []byte(content), 0o644); err != nil {
t.Fatal(err)
}
assistant := agent.New(agent.Options{SystemPrompt: "base"})
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
out, err := m.handleLocalCommand("/skills install " + source)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, `Installed skill "review"`) {
t.Fatalf("output = %q", out)
}
if _, err := os.Stat(filepath.Join(home, ".agentu", "skills", "review", "SKILL.md")); err != nil {
t.Fatalf("installed skill missing: %v", err)
}
messages := assistant.Messages()
if len(messages) == 0 || !strings.Contains(messages[0].Content, "Check every diff.") {
t.Fatalf("skill not activated in system prompt: %+v", messages)
}
}
+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"))
}
+14 -1
View File
@@ -66,6 +66,7 @@ type rawProviderConfig struct {
BaseURL string `yaml:"base_url"` BaseURL string `yaml:"base_url"`
Models []rawModelConfig `yaml:"models"` Models []rawModelConfig `yaml:"models"`
APIKey string `yaml:"api_key"` APIKey string `yaml:"api_key"`
ThinkingParam string `yaml:"thinking_param"`
} }
type rawModelConfig struct { type rawModelConfig struct {
@@ -235,13 +236,14 @@ func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
Models: models, Models: models,
ModelConfigs: modelConfigs, ModelConfigs: modelConfigs,
APIKey: raw.APIKey, APIKey: raw.APIKey,
ThinkingParam: strings.TrimSpace(raw.ThinkingParam),
} }
if len(provider.Models) > 0 { if len(provider.Models) > 0 {
provider.Model = provider.Models[0] provider.Model = provider.Models[0]
} }
applyModelDefaults(&provider) applyModelDefaults(&provider)
if provider.ThinkingConfigured && provider.ThinkingParam == "" { if provider.ThinkingConfigured && provider.ThinkingParam == "" {
provider.ThinkingParam = "thinking" provider.ThinkingParam = DefaultThinkingParam(provider.APIType)
} }
return provider return provider
} }
@@ -268,6 +270,17 @@ func IsAPIType(value string) bool {
return false 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) { func normalizeModelConfigs(rawModels []rawModelConfig) ([]string, map[string]ModelConfig) {
models := make([]string, 0, len(rawModels)) models := make([]string, 0, len(rawModels))
configs := make(map[string]ModelConfig, len(rawModels)) configs := make(map[string]ModelConfig, len(rawModels))
+61
View File
@@ -441,3 +441,64 @@ providers:
t.Fatalf("err = %v", err) 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)
}
}
+22 -1
View File
@@ -124,6 +124,9 @@ func buildResponsesRequest(req ChatRequest) map[string]any {
"input": buildResponsesInput(req.Messages), "input": buildResponsesInput(req.Messages),
"stream": true, "stream": true,
} }
if instructions := responsesInstructions(req.Messages); instructions != "" {
payload["instructions"] = instructions
}
if len(req.Tools) > 0 { if len(req.Tools) > 0 {
payload["tools"] = buildResponsesTools(req.Tools) payload["tools"] = buildResponsesTools(req.Tools)
} }
@@ -153,7 +156,9 @@ func buildResponsesInput(messages []Message) []any {
for _, msg := range messages { for _, msg := range messages {
switch msg.Role { switch msg.Role {
case RoleSystem: case RoleSystem:
items = append(items, responsesMessageInput{Role: "system", Content: msg.Content}) // System instructions are sent as the top-level "instructions"
// field and must not be duplicated inside the input array.
continue
case RoleUser: case RoleUser:
items = append(items, responsesMessageInput{Role: "user", Content: msg.Content}) items = append(items, responsesMessageInput{Role: "user", Content: msg.Content})
case RoleAssistant: case RoleAssistant:
@@ -179,6 +184,22 @@ func buildResponsesInput(messages []Message) []any {
return items 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 { func buildResponsesTools(tools []Tool) []responsesTool {
out := make([]responsesTool, 0, len(tools)) out := make([]responsesTool, 0, len(tools))
for _, tool := range tools { for _, tool := range tools {
+41 -9
View File
@@ -59,12 +59,15 @@ func TestBuildResponsesRequest(t *testing.T) {
if raw["model"] != "gpt-test" || raw["stream"] != true { if raw["model"] != "gpt-test" || raw["stream"] != true {
t.Fatalf("model/stream = %#v / %#v", raw["model"], raw["stream"]) 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" { if raw["tool_choice"] != "auto" {
t.Fatalf("tool_choice = %#v", raw["tool_choice"]) t.Fatalf("tool_choice = %#v", raw["tool_choice"])
} }
input, ok := raw["input"].([]any) input, ok := raw["input"].([]any)
if !ok || len(input) != 6 { if !ok || len(input) != 5 {
t.Fatalf("input = %#v", raw["input"]) t.Fatalf("input = %#v", raw["input"])
} }
checkInputItem := func(index int, wantType, wantKey, wantValue string) { checkInputItem := func(index int, wantType, wantKey, wantValue string) {
@@ -80,16 +83,15 @@ func TestBuildResponsesRequest(t *testing.T) {
t.Fatalf("input[%d].type = %#v, want %q", index, item["type"], wantType) t.Fatalf("input[%d].type = %#v, want %q", index, item["type"], wantType)
} }
} }
checkInputItem(0, "", "role", "system") checkInputItem(0, "", "role", "user")
checkInputItem(1, "", "role", "user") checkInputItem(1, "", "role", "assistant")
checkInputItem(2, "", "role", "assistant") checkInputItem(2, "function_call", "call_id", "call_1")
checkInputItem(3, "function_call", "call_id", "call_1") checkInputItem(3, "function_call_output", "call_id", "call_1")
checkInputItem(4, "function_call_output", "call_id", "call_1") checkInputItem(4, "", "role", "assistant")
checkInputItem(5, "", "role", "assistant") if item := input[2].(map[string]any); item["name"] != "file_read" || item["arguments"] != `{"path":"README.md"}` {
if item := input[3].(map[string]any); item["name"] != "file_read" || item["arguments"] != `{"path":"README.md"}` {
t.Fatalf("function_call item = %#v", item) t.Fatalf("function_call item = %#v", item)
} }
if item := input[4].(map[string]any); item["output"] != "contents" { if item := input[3].(map[string]any); item["output"] != "contents" {
t.Fatalf("function_call_output item = %#v", item) t.Fatalf("function_call_output item = %#v", item)
} }
@@ -117,6 +119,36 @@ func TestBuildResponsesRequest(t *testing.T) {
} }
} }
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) { func TestOpenAIResponsesClientStreamsContentAndToolCalls(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/responses" { if r.URL.Path != "/v1/responses" {