5 Commits
Author SHA1 Message Date
loveuer 065b30a84c feat: complete v0.3.1 agent workflow 2026-09-15 03:39:22 -07:00
loveuer 47671f3098 feat(tools): stream shell_run output live into the TUI 2026-08-13 13:56:24 +08:00
loveuer 28195519d1 feat(session,llm): restore session runtime on resume; align Responses API details 2026-08-13 13:46:26 +08:00
loveuer 425f6ee5ee feat(llm): support OpenAI Responses API via per-provider api_type 2026-08-13 11:46:23 +08:00
loveuer c8d8329dbb feat(tui): codex-style input, history, bang commands, single-line status bar
- Arrow-key input history (up/down with multiline awareness)
- Bang shell commands via ! prefix (requires --yolo)
- Codex-style input bar with mode indicator (R/Y/⌘)
- Command palette virtual scrolling
- Multiline input viewport fix
- Merged model pills + footer into single status line
- Removed shift+enter (wait for TUI library support)
- Removed default footer hints
2026-06-25 20:40:50 -07:00
40 changed files with 7174 additions and 364 deletions
+1
View File
@@ -3,3 +3,4 @@ agentu*.yaml
/coverage.out /coverage.out
*.swp *.swp
chat-history.txt chat-history.txt
dist/
+78 -13
View File
@@ -1,18 +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.
- 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
@@ -29,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
@@ -48,6 +56,7 @@ Use `--config <path>` to load another file.
```yaml ```yaml
providers: providers:
loveuer: loveuer:
api_type: chat
base_url: https://ai.loveuer.com base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY} api_key: ${AGENTU_API_KEY}
models: models:
@@ -64,6 +73,20 @@ Provider names are the keys under `providers`. At startup, agentu selects the
first provider by name. Use `/model provider <name>` to switch providers for the first provider by name. Use `/model provider <name>` to switch providers for the
current session. current session.
`api_type` selects which API protocol the provider speaks:
```text
chat OpenAI-compatible /v1/chat/completions (default)
responses OpenAI Responses API /v1/responses
```
The Responses API client converts the conversation history to Responses input
items (`user`/`assistant` messages, `function_call`, and
`function_call_output`), sends the system prompt as the top-level
`instructions` field, streams `response.output_text.delta` events, and supports
function calling through `response.output_item.added`,
`response.function_call_arguments.delta`, and `response.output_item.done`.
`models` is required. `/model model <name>` is restricted to that list. Each `models` is required. `/model model <name>` is restricted to that list. Each
entry is an object with: entry is an object with:
@@ -83,7 +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 OpenAI-compatible request field named `thinking`. top-level request field. The default field depends on `api_type`: `chat` uses
`thinking`, while `responses` uses the official `reasoning` field, sent as
`{"effort": ...}` with agentu's `middle` level mapped to the official `medium`
value (`none`, `high`, `xhigh`, `max` pass through). Override with
`thinking_param` if a provider expects a different field name.
## CLI Flags ## CLI Flags
@@ -94,28 +121,38 @@ 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
- `enter` sends the message. - `enter` sends the message.
- `ctrl+j` or `alt+enter` inserts a newline. - `ctrl+j` inserts a newline.
- `up` recalls the previous input; `ctrl+p` / `ctrl+n` navigate input history.
- The input starts at one line and grows up to six lines. - 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.
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.
In the TUI, type `/model ` or `/model model ` to choose from configured model IDs in the candidate palette.
- `/model thinking <none|middle|high|xhigh|max>` changes thinking for the current session. - `/model thinking <none|middle|high|xhigh|max>` changes thinking for the current session.
- `/sessions` lists saved sessions. - `/sessions` lists saved sessions.
- `/resume [id]` resumes by ID or opens an interactive picker. - `/resume [id]` resumes by ID or opens an interactive picker.
@@ -125,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:
@@ -153,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
@@ -171,5 +231,10 @@ Local config and build outputs are intentionally ignored by git.
- Provider/model/thinking changes are session-only; agentu does not write them - Provider/model/thinking changes are session-only; agentu does not write them
back to `~/.agentu/config.yaml`. back to `~/.agentu/config.yaml`.
- `api_type` is per provider and cannot be changed at runtime; use
`/model provider <name>` to switch to a provider configured with a different
API type.
- The config schema is intentionally strict during MVP development. Deprecated - 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.
+3
View File
@@ -1,5 +1,8 @@
providers: providers:
loveuer: loveuer:
# API protocol: "chat" (OpenAI-compatible /v1/chat/completions, default)
# or "responses" (OpenAI Responses API /v1/responses).
api_type: chat
base_url: https://ai.loveuer.com base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY} api_key: ${AGENTU_API_KEY}
models: models:
+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,33 +651,145 @@ 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 {
if a.toolRegistry == nil { if a.toolRegistry == nil {
@@ -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)
}
+40 -12
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 {
@@ -91,11 +104,19 @@ func loadConfig(configPath string) (*config.Config, error) {
func defaultProvider(cfg *config.Config, httpClient *http.Client) (string, config.ProviderConfig, llm.Provider) { func defaultProvider(cfg *config.Config, httpClient *http.Client) (string, config.ProviderConfig, llm.Provider) {
providerName := cfg.DefaultProviderName() providerName := cfg.DefaultProviderName()
providerConfig := cfg.Providers[providerName] providerConfig := cfg.Providers[providerName]
provider := llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, httpClient) provider := newProviderClient(providerConfig, httpClient)
return providerName, providerConfig, provider return providerName, providerConfig, provider
} }
func newAssistant(cfg *config.Config, providerName string, providerConfig config.ProviderConfig, provider llm.Provider, yolo bool, httpClient *http.Client) *agent.Agent { func newProviderClient(providerConfig config.ProviderConfig, httpClient *http.Client) llm.Provider {
if providerConfig.APIType == config.APITypeResponses {
return llm.NewOpenAIResponsesClient(providerConfig.BaseURL, providerConfig.APIKey, httpClient)
}
return llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, httpClient)
}
func newAssistant(cfg *config.Config, providerName string, providerConfig config.ProviderConfig, provider llm.Provider, opts cliOptions, globalSkills []skills.Skill, authorizer agent.ToolAuthorizer, rollbackStore *tools.RollbackStore, httpClient *http.Client) *agent.Agent {
writesEnabled := opts.yolo || !opts.plain
return agent.New(agent.Options{ return agent.New(agent.Options{
Provider: provider, Provider: provider,
ProviderName: providerName, ProviderName: providerName,
@@ -103,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 {
@@ -122,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)
@@ -147,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)
} }
@@ -158,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
}
+46 -5
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)
} }
@@ -204,10 +226,18 @@ func (m *Manager) CurrentThinking() string {
return m.provider.Thinking return m.provider.Thinking
} }
func (m *Manager) CurrentAPIType() string {
if m.provider.APIType == "" {
return config.APITypeChat
}
return m.provider.APIType
}
func (m *Manager) Info() string { func (m *Manager) Info() string {
var b strings.Builder var b strings.Builder
fmt.Fprintf(&b, "provider: %s\n", m.CurrentProvider()) fmt.Fprintf(&b, "provider: %s\n", m.CurrentProvider())
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel()) fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
fmt.Fprintf(&b, "api: %s\n", m.CurrentAPIType())
fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking()) fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking())
fmt.Fprintf(&b, "providers: %s\n", strings.Join(m.providerNames(), ", ")) fmt.Fprintf(&b, "providers: %s\n", strings.Join(m.providerNames(), ", "))
fmt.Fprintf(&b, "models: %s\n", strings.Join(m.provider.Models, ", ")) fmt.Fprintf(&b, "models: %s\n", strings.Join(m.provider.Models, ", "))
@@ -223,7 +253,7 @@ func (m *Manager) SetProvider(name string) (string, error) {
} }
m.providerName = name m.providerName = name
m.provider = provider m.provider = provider
m.syncAgent(llm.NewOpenAICompatibleClient(provider.BaseURL, provider.APIKey, m.httpClient)) m.syncAgent(providerClient(provider, m.httpClient))
return "Switched provider.\n" + m.summary(), nil return "Switched provider.\n" + m.summary(), nil
} }
@@ -249,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)
@@ -257,7 +287,14 @@ func (m *Manager) SetThinking(level string) (string, error) {
} }
func (m *Manager) summary() string { func (m *Manager) summary() string {
return fmt.Sprintf("provider: %s\nmodel: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentThinking()) return fmt.Sprintf("provider: %s\nmodel: %s\napi: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentAPIType(), m.CurrentThinking())
}
func providerClient(provider config.ProviderConfig, httpClient *http.Client) llm.Provider {
if provider.APIType == config.APITypeResponses {
return llm.NewOpenAIResponsesClient(provider.BaseURL, provider.APIKey, httpClient)
}
return llm.NewOpenAICompatibleClient(provider.BaseURL, provider.APIKey, httpClient)
} }
func (m *Manager) maxContextTokens() int { func (m *Manager) maxContextTokens() int {
@@ -288,9 +325,13 @@ func (m *Manager) providerNames() []string {
return names return names
} }
func (m *Manager) AvailableModels() []string {
return append([]string(nil), m.provider.Models...)
}
func contains(values []string, value string) bool { func contains(values []string, value string) bool {
for _, item := range values { for _, v := range values {
if item == value { if v == value {
return true return true
} }
} }
+201
View File
@@ -2,6 +2,8 @@ package session
import ( import (
"context" "context"
"net/http"
"net/http/httptest"
"path/filepath" "path/filepath"
"strings" "strings"
"testing" "testing"
@@ -109,6 +111,76 @@ func TestManagerRejectsUnknownModelAndThinking(t *testing.T) {
} }
} }
func TestManagerSwitchesToResponsesProvider(t *testing.T) {
var hitPath string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hitPath = r.URL.Path
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"ok"}` + "\n\n"))
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed"}}` + "\n\n"))
}))
defer server.Close()
cfg := &config.Config{
Providers: map[string]config.ProviderConfig{
"chat": {
Name: "chat",
APIType: config.APITypeChat,
BaseURL: "https://chat.example.com",
APIKey: "sk-test",
Model: "model-a",
Models: []string{"model-a"},
},
"responses": {
Name: "responses",
APIType: config.APITypeResponses,
BaseURL: server.URL,
APIKey: "sk-test",
Model: "model-b",
Models: []string{"model-b"},
},
},
}
assistant := agent.New(agent.Options{Provider: &recordingProvider{}, Model: "model-a"})
store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
if err != nil {
t.Fatal(err)
}
manager, err := NewManager(cfg, assistant, server.Client(), store)
if err != nil {
t.Fatal(err)
}
if _, err := manager.SetProvider("responses"); err != nil {
t.Fatal(err)
}
if manager.CurrentAPIType() != config.APITypeResponses {
t.Fatalf("api type = %q", manager.CurrentAPIType())
}
var out strings.Builder
if err := assistant.RunTurn(context.Background(), "hi", &out, nil); err != nil {
t.Fatal(err)
}
if hitPath != "/v1/responses" {
t.Fatalf("path = %q, want /v1/responses", hitPath)
}
if out.String() != "ok" {
t.Fatalf("output = %q", out.String())
}
}
func TestManagerExposesModelCommandCandidates(t *testing.T) {
manager, _, _ := testManager(t)
models := manager.AvailableModels()
if len(models) != 2 || models[0] != "model-a" || models[1] != "model-b" {
t.Fatalf("AvailableModels() = %v", models)
}
models[0] = "mutated"
if got := manager.AvailableModels()[0]; got != "model-a" {
t.Fatalf("mutation leaked: %q", got)
}
}
func TestManagerSaveAndResume(t *testing.T) { func TestManagerSaveAndResume(t *testing.T) {
manager, assistant, _ := testManager(t) manager, assistant, _ := testManager(t)
@@ -331,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)
}
}
+822 -126
View File
File diff suppressed because it is too large Load Diff
+545 -35
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"
@@ -21,7 +25,7 @@ func TestModelRendersChatShell(t *testing.T) {
m = next.(model) m = next.(model)
rendered := m.View() rendered := m.View()
for _, want := range []string{"test-model", "yolo tools", "light", "ctrl+j", "alt+enter", "ctrl+p/ctrl+n history"} { for _, want := range []string{"test-model"} {
if !strings.Contains(rendered, want) { if !strings.Contains(rendered, want) {
t.Fatalf("rendered view missing %q:\n%s", want, rendered) t.Fatalf("rendered view missing %q:\n%s", want, rendered)
} }
@@ -37,18 +41,16 @@ func TestModelRendersChatShell(t *testing.T) {
} }
} }
func TestInputTextDoesNotPaintBackground(t *testing.T) { func TestInputTextBackgroundMatchesInputBox(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"}) m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
for name, color := range map[string]any{ want := lightTheme().InputBg
"input box": m.styles.InputBox.GetBackground(), for name, color := range map[string]lipgloss.TerminalColor{
"focused base": m.input.FocusedStyle.Base.GetBackground(), "focused base": m.input.FocusedStyle.Base.GetBackground(),
"focused cursor line": m.input.FocusedStyle.CursorLine.GetBackground(), "focused cursor line": m.input.FocusedStyle.CursorLine.GetBackground(),
"blurred base": m.input.BlurredStyle.Base.GetBackground(), "blurred base": m.input.BlurredStyle.Base.GetBackground(),
"blurred cursor line": m.input.BlurredStyle.CursorLine.GetBackground(), "blurred cursor line": m.input.BlurredStyle.CursorLine.GetBackground(),
} { } {
if _, ok := color.(lipgloss.NoColor); !ok { requireStyleColor(t, name, color, want)
t.Fatalf("%s should not paint input text background", name)
}
} }
} }
@@ -58,12 +60,12 @@ func TestInputBoxIsRoomier(t *testing.T) {
m = next.(model) m = next.(model)
if got := m.styles.InputBox.GetVerticalFrameSize(); got != 2 { if got := m.styles.InputBox.GetVerticalFrameSize(); got != 2 {
t.Fatalf("input box vertical frame should stay compact, got %d", got) t.Fatalf("input box vertical frame should be 2 with vertical padding, got %d", got)
} }
if got := m.styles.InputBox.GetHorizontalFrameSize(); got < 6 { if got := m.styles.InputBox.GetHorizontalFrameSize(); got != 3 {
t.Fatalf("input box horizontal frame should include roomier padding, got %d", got) t.Fatalf("input box horizontal frame should be 3 (asymmetric padding), got %d", got)
} }
if got, want := m.input.Width(), 100-m.styles.InputBox.GetHorizontalFrameSize(); got != want { if got, want := m.input.Width(), 100-m.styles.InputBox.GetHorizontalFrameSize()-m.inputModeWidth()-1; got != want {
t.Fatalf("input content width = %d, want %d", got, want) t.Fatalf("input content width = %d, want %d", got, want)
} }
rendered := m.inputView() rendered := m.inputView()
@@ -78,13 +80,8 @@ func TestModelMetaRendersBelowInput(t *testing.T) {
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30}) next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
rendered := next.(model).View() rendered := next.(model).View()
inputIndex := strings.Index(rendered, "Message agentu") if !strings.Contains(rendered, "test-model") {
modelIndex := strings.Index(rendered, "test-model") t.Fatalf("rendered view missing %q:\n%s", "test-model", rendered)
if inputIndex < 0 || modelIndex < 0 {
t.Fatalf("rendered view missing input or model meta:\n%s", rendered)
}
if modelIndex < inputIndex {
t.Fatalf("model meta should render below input:\n%s", rendered)
} }
} }
@@ -93,7 +90,7 @@ func TestModelMetaRendersContextUsage(t *testing.T) {
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("x", 800)}}) assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("x", 800)}})
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"}) m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 120, Height: 30}) next, _ := m.Update(tea.WindowSizeMsg{Width: 120, Height: 30})
rendered := next.(model).modelMetaView() rendered := next.(model).footerView()
for _, want := range []string{"ctx ~", "/1k", "%"} { for _, want := range []string{"ctx ~", "/1k", "%"} {
if !strings.Contains(rendered, want) { if !strings.Contains(rendered, want) {
@@ -133,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})
@@ -336,6 +443,23 @@ func TestInputHistoryNavigation(t *testing.T) {
} }
} }
func TestUpArrowRecallsPreviousInput(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.rememberInput("first")
m.rememberInput("second")
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyUp})
m = next.(model)
if got := m.input.Value(); got != "second" {
t.Fatalf("after first up input = %q, want %q", got, "second")
}
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyUp})
m = next.(model)
if got := m.input.Value(); got != "first" {
t.Fatalf("after second up input = %q, want %q", got, "first")
}
}
func TestToolMessageRendersStructuredCard(t *testing.T) { func TestToolMessageRendersStructuredCard(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"}) m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`}) m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`})
@@ -385,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 {
@@ -480,16 +696,60 @@ func runInDir(t *testing.T, dir string, name string, args ...string) {
} }
} }
func TestBangShellCommandRunsFromInput(t *testing.T) {
dir := t.TempDir()
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true, WorkingDir: dir})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("! printf hello")
next, _ = m.submit()
m = next.(model)
if m.input.Value() != "" {
t.Fatalf("input not cleared after bang command: %q", m.input.Value())
}
rendered := stripANSI(m.renderMessages())
if !strings.Contains(rendered, "$ printf hello") {
t.Fatalf("rendered output missing shell echo:\n%s", rendered)
}
if !strings.Contains(rendered, "hello") {
t.Fatalf("rendered output missing 'hello':\n%s", rendered)
}
}
func TestBangShellCommandRequiresYolo(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: false})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("! printf hello")
next, _ = m.submit()
m = next.(model)
rendered := stripANSI(m.renderMessages())
if !strings.Contains(rendered, "shell command input requires --yolo") {
t.Fatalf("rendered output missing yolo requirement:\n%s", rendered)
}
}
func TestBangShellCommandRejectsEmptyCommand(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("!")
next, _ = m.submit()
m = next.(model)
rendered := stripANSI(m.renderMessages())
if !strings.Contains(rendered, "usage: ! <command>") {
t.Fatalf("rendered output missing usage hint:\n%s", rendered)
}
}
func TestModelCommand(t *testing.T) { func TestModelCommand(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true}) m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
m.input.SetValue("/model") m.input.SetValue("/model")
next, _ := m.submit() next, _ := m.submit()
rendered := next.(model).renderMessages() rendered := next.(model).renderMessages()
for _, want := range []string{"model: test-model", "mode: yolo tools", "theme: light"} { if !strings.Contains(rendered, "test-model") {
if !strings.Contains(rendered, want) { t.Fatalf("model command output missing %q:\n%s", "test-model", rendered)
t.Fatalf("model command output missing %q:\n%s", want, rendered)
}
} }
} }
@@ -508,11 +768,124 @@ func TestModelThinkingCommand(t *testing.T) {
} }
} }
func TestModelCommandPaletteShowsConfiguredModels(t *testing.T) {
manager := &fakeModelManager{model: "model-a", models: []string{"model-a", "model-b"}}
m := newModel(context.Background(), nil, Options{ModelName: "model-a", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model ")
m.afterInputChanged()
plain := stripANSI(m.View())
for _, want := range []string{iconCommand + " Model candidates", "model-a", "current model", "model-b", "switch model"} {
if !strings.Contains(plain, want) {
t.Fatalf("rendered palette missing %q:\n%s", want, plain)
}
}
}
func TestModelCommandPaletteSelectsConfiguredModel(t *testing.T) {
manager := &fakeModelManager{model: "model-a", models: []string{"model-a", "model-b"}}
m := newModel(context.Background(), nil, Options{ModelName: "model-a", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model ")
m.afterInputChanged()
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyDown})
m = next.(model)
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = next.(model)
if manager.model != "model-b" {
t.Fatalf("model = %q", manager.model)
}
if got := m.input.Value(); got != "" {
t.Fatalf("input = %q", got)
}
rendered := stripANSI(m.renderMessages())
if !strings.Contains(rendered, "model: model-b") {
t.Fatalf("rendered output missing model switch:\n%s", rendered)
}
}
func TestModelCommandPaletteFiltersBareModelPartial(t *testing.T) {
manager := &fakeModelManager{model: "claude-sonnet", models: []string{"claude-sonnet", "deepseek-v4-flash"}}
m := newModel(context.Background(), nil, Options{ModelName: "claude-sonnet", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model deep")
m.afterInputChanged()
palette := stripANSI(m.commandPaletteView())
if !strings.Contains(palette, "deepseek-v4-flash") {
t.Fatalf("expected deepseek-v4-flash in palette:\n%s", palette)
}
if strings.Contains(palette, "claude-sonnet") {
t.Fatalf("unexpected claude-sonnet in palette:\n%s", palette)
}
}
func TestModelCommandPaletteFiltersExplicitModelSubcommand(t *testing.T) {
manager := &fakeModelManager{model: "claude-sonnet", models: []string{"claude-sonnet", "deepseek-v4-flash"}}
m := newModel(context.Background(), nil, Options{ModelName: "claude-sonnet", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model model deep")
m.afterInputChanged()
palette := stripANSI(m.commandPaletteView())
if !strings.Contains(palette, "deepseek-v4-flash") {
t.Fatalf("expected deepseek-v4-flash in palette:\n%s", palette)
}
if strings.Contains(palette, "claude-sonnet") {
t.Fatalf("unexpected claude-sonnet in palette:\n%s", palette)
}
}
func TestModelCommandPaletteTabFillsConfiguredModel(t *testing.T) {
manager := &fakeModelManager{model: "claude-sonnet", models: []string{"claude-sonnet", "deepseek-v4-flash"}}
m := newModel(context.Background(), nil, Options{ModelName: "claude-sonnet", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model deep")
m.afterInputChanged()
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyTab})
m = next.(model)
if got := m.input.Value(); got != "/model model deepseek-v4-flash" {
t.Fatalf("input = %q", got)
}
if manager.model != "claude-sonnet" {
t.Fatalf("model should not change on tab: %q", manager.model)
}
}
func TestModelCommandPaletteIgnoresProviderAndThinkingSubcommands(t *testing.T) {
manager := &fakeModelManager{model: "model-a", models: []string{"model-a", "model-b"}}
m := newModel(context.Background(), nil, Options{ModelName: "model-a", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model provider ")
m.afterInputChanged()
if sugs := m.commandSuggestions(); len(sugs) != 0 {
t.Fatalf("provider suggestions should be empty, got %d", len(sugs))
}
if view := m.commandPaletteView(); view != "" {
t.Fatalf("provider palette should be empty:\n%s", view)
}
m.input.SetValue("/model thinking ")
m.afterInputChanged()
if sugs := m.commandSuggestions(); len(sugs) != 0 {
t.Fatalf("thinking suggestions should be empty, got %d", len(sugs))
}
if view := m.commandPaletteView(); view != "" {
t.Fatalf("thinking palette should be empty:\n%s", view)
}
}
type fakeModelManager struct { type fakeModelManager struct {
provider string provider string
model string model string
thinking string thinking string
sessionName string sessionName string
models []string
saved bool saved bool
} }
@@ -583,6 +956,16 @@ func (f *fakeModelManager) NewSession() (string, error) {
return "New session.", nil return "New session.", nil
} }
func (f *fakeModelManager) AvailableModels() []string {
if len(f.models) > 0 {
return append([]string(nil), f.models...)
}
if f.model != "" {
return []string{f.model}
}
return nil
}
func TestParseThemeMode(t *testing.T) { func TestParseThemeMode(t *testing.T) {
for input, want := range map[string]ThemeMode{ for input, want := range map[string]ThemeMode{
"": ThemeLight, "": ThemeLight,
@@ -627,9 +1010,7 @@ func TestDarkThemePaintsAppBackground(t *testing.T) {
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30}) next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model) m = next.(model)
rendered := stripANSI(m.View()) rendered := stripANSI(m.View())
if !strings.Contains(rendered, "theme dark") { _ = rendered
t.Fatalf("rendered view missing 'theme dark':\n%s", rendered)
}
if m.theme.Mode != ThemeDark { if m.theme.Mode != ThemeDark {
t.Fatalf("theme.Mode = %q, want %q", m.theme.Mode, ThemeDark) t.Fatalf("theme.Mode = %q, want %q", m.theme.Mode, ThemeDark)
} }
@@ -643,8 +1024,7 @@ func TestDarkThemeKeepsComponentBackgroundsLightweight(t *testing.T) {
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30}) next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model) m = next.(model)
requireNoStyleColor(t, "Viewport.Background", m.styles.Viewport.GetBackground()) requireNoStyleColor(t, "Viewport.Background", m.styles.Viewport.GetBackground())
requireNoStyleColor(t, "InputBox.Background", m.styles.InputBox.GetBackground()) requireStyleColor(t, "InputBox.Background", m.styles.InputBox.GetBackground(), darkTheme().InputBg)
requireNoStyleColor(t, "AssistantMsg.Background", m.styles.AssistantMsg.GetBackground())
requireNoStyleColor(t, "Activity.Background", m.styles.Activity.GetBackground()) requireNoStyleColor(t, "Activity.Background", m.styles.Activity.GetBackground())
m.running = true m.running = true
m.status = "thinking..." m.status = "thinking..."
@@ -662,16 +1042,34 @@ func TestDarkThemeStylesUseDarkPalette(t *testing.T) {
requireStyleColor(t, "Pill.Background", m.styles.Pill.GetBackground(), want.SurfaceAlt) requireStyleColor(t, "Pill.Background", m.styles.Pill.GetBackground(), want.SurfaceAlt)
requireStyleColor(t, "input.FocusedStyle.Base.Foreground", m.input.FocusedStyle.Base.GetForeground(), want.Text) requireStyleColor(t, "input.FocusedStyle.Base.Foreground", m.input.FocusedStyle.Base.GetForeground(), want.Text)
requireStyleColor(t, "input.FocusedStyle.Placeholder.Foreground", m.input.FocusedStyle.Placeholder.GetForeground(), want.Muted) requireStyleColor(t, "input.FocusedStyle.Placeholder.Foreground", m.input.FocusedStyle.Placeholder.GetForeground(), want.Muted)
requireStyleColor(t, "InputModeYolo.Foreground", m.styles.InputModeYolo.GetForeground(), want.Error)
requireStyleColor(t, "InputModeCommand.Foreground", m.styles.InputModeCommand.GetForeground(), want.User)
} }
func TestModelCommandShowsDarkTheme(t *testing.T) { func TestInputModeSymbolReflectsYoloAndBangCommand(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark}) m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
m.input.SetValue("/model") next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
next, _ := m.submit() m = next.(model)
rendered := next.(model).renderMessages()
if !strings.Contains(rendered, "theme: dark") { rendered := stripANSI(m.inputView())
t.Fatalf("model command output missing 'theme: dark':\n%s", rendered) if !strings.Contains(rendered, "Y") {
t.Fatalf("YOLO input view should contain Y symbol:\n%s", rendered)
} }
m.input.SetValue("! pwd")
m.afterInputChanged()
rendered = stripANSI(m.inputView())
if !strings.Contains(rendered, iconCommand) {
t.Fatalf("bang input view should contain command icon:\n%s", rendered)
}
if strings.Contains(rendered, "Y") {
t.Fatalf("bang input view should not contain Y symbol:\n%s", rendered)
}
}
func TestInputModeYoloUsesWarningColor(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
requireStyleColor(t, "InputModeYolo.Foreground", m.styles.InputModeYolo.GetForeground(), lightTheme().Error)
} }
func TestThemeSlashCommandSwitchesAtRuntime(t *testing.T) { func TestThemeSlashCommandSwitchesAtRuntime(t *testing.T) {
@@ -720,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"))
}
+31 -13
View File
@@ -43,6 +43,7 @@ type theme struct {
Background string Background string
Surface string Surface string
SurfaceAlt string SurfaceAlt string
InputBg string
Border string Border string
Text string Text string
Muted string Muted string
@@ -60,11 +61,13 @@ type styles struct {
Muted lipgloss.Style Muted lipgloss.Style
Activity lipgloss.Style Activity lipgloss.Style
Header lipgloss.Style Header lipgloss.Style
ModelMeta lipgloss.Style
Footer lipgloss.Style Footer lipgloss.Style
Viewport lipgloss.Style Viewport lipgloss.Style
InputBox lipgloss.Style InputBox lipgloss.Style
CommandHint lipgloss.Style CommandHint lipgloss.Style
InputModeReadonly lipgloss.Style
InputModeYolo lipgloss.Style
InputModeCommand lipgloss.Style
StatusReady lipgloss.Style StatusReady lipgloss.Style
StatusWarn lipgloss.Style StatusWarn lipgloss.Style
StatusError lipgloss.Style StatusError lipgloss.Style
@@ -98,8 +101,8 @@ func lightTheme() theme {
return theme{ return theme{
Mode: ThemeLight, Mode: ThemeLight,
Background: "#F8FAFC", Background: "#F8FAFC",
Surface: "#FFFFFF",
SurfaceAlt: "#EEF2F7", SurfaceAlt: "#EEF2F7",
InputBg: "#E4E4E4",
Border: "#CBD5E1", Border: "#CBD5E1",
Text: "#111827", Text: "#111827",
Muted: "#64748B", Muted: "#64748B",
@@ -116,8 +119,8 @@ func darkTheme() theme {
return theme{ return theme{
Mode: ThemeDark, Mode: ThemeDark,
Background: "#111827", Background: "#111827",
Surface: "#1F2937",
SurfaceAlt: "#0F172A", SurfaceAlt: "#0F172A",
InputBg: "#383838",
Border: "#475569", Border: "#475569",
Text: "#E5E7EB", Text: "#E5E7EB",
Muted: "#94A3B8", Muted: "#94A3B8",
@@ -153,10 +156,6 @@ func (t theme) styles() styles {
Bold(true). Bold(true).
Foreground(lipgloss.Color(t.Title)), Foreground(lipgloss.Color(t.Title)),
ModelMeta: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1),
Footer: lipgloss.NewStyle(). Footer: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)). Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1), Padding(0, 1),
@@ -166,14 +165,31 @@ func (t theme) styles() styles {
InputBox: lipgloss.NewStyle(). InputBox: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)). Foreground(lipgloss.Color(t.Text)).
Border(lipgloss.ThickBorder()). Background(lipgloss.Color(t.InputBg)).
BorderForeground(lipgloss.Color(t.Border)). Padding(1, 2, 1, 1),
Padding(0, 2),
CommandHint: lipgloss.NewStyle(). CommandHint: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)). Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1), Padding(0, 1),
InputModeReadonly: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Muted)).
Background(lipgloss.Color(t.InputBg)).
Padding(0, 1),
InputModeYolo: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Error)).
Background(lipgloss.Color(t.InputBg)).
Padding(0, 1),
InputModeCommand: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.User)).
Background(lipgloss.Color(t.InputBg)).
Padding(0, 1),
StatusReady: lipgloss.NewStyle(). StatusReady: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Assistant)), Foreground(lipgloss.Color(t.Assistant)),
@@ -258,8 +274,10 @@ func bubbleMessageStyle(textColor, backgroundColor string) lipgloss.Style {
func applyTextareaTheme(input *textarea.Model, t theme) { func applyTextareaTheme(input *textarea.Model, t theme) {
focused := textarea.Style{ focused := textarea.Style{
Base: lipgloss.NewStyle(). Base: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)), Foreground(lipgloss.Color(t.Text)).
CursorLine: lipgloss.NewStyle(), Background(lipgloss.Color(t.InputBg)),
CursorLine: lipgloss.NewStyle().
Background(lipgloss.Color(t.InputBg)),
Placeholder: lipgloss.NewStyle(). Placeholder: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)), Foreground(lipgloss.Color(t.Muted)),
Prompt: lipgloss.NewStyle(). Prompt: lipgloss.NewStyle().
@@ -270,7 +288,7 @@ func applyTextareaTheme(input *textarea.Model, t theme) {
Foreground(lipgloss.Color(t.Surface)), Foreground(lipgloss.Color(t.Surface)),
} }
blurred := focused blurred := focused
blurred.CursorLine = lipgloss.NewStyle() blurred.CursorLine = lipgloss.NewStyle().Background(lipgloss.Color(t.InputBg))
input.FocusedStyle = focused input.FocusedStyle = focused
input.BlurredStyle = blurred input.BlurredStyle = blurred
+47 -1
View File
@@ -14,6 +14,11 @@ import (
const DefaultPath = "~/.agentu/config.yaml" const DefaultPath = "~/.agentu/config.yaml"
const (
APITypeChat = "chat"
APITypeResponses = "responses"
)
type Config struct { type Config struct {
Providers map[string]ProviderConfig `yaml:"providers"` Providers map[string]ProviderConfig `yaml:"providers"`
Agent AgentConfig `yaml:"agent"` Agent AgentConfig `yaml:"agent"`
@@ -21,6 +26,7 @@ type Config struct {
type ProviderConfig struct { type ProviderConfig struct {
Name string Name string
APIType string
BaseURL string BaseURL string
Model string Model string
Models []string Models []string
@@ -56,9 +62,11 @@ type rawConfig struct {
} }
type rawProviderConfig struct { type rawProviderConfig struct {
APIType string `yaml:"api_type"`
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 {
@@ -153,6 +161,9 @@ func (c *Config) Validate() error {
for name, provider := range c.Providers { for name, provider := range c.Providers {
var missing []string var missing []string
prefix := "providers." + name prefix := "providers." + name
if !IsAPIType(provider.APIType) {
return fmt.Errorf("%s.api_type must be one of: %s", prefix, strings.Join(APITypes(), ", "))
}
if strings.TrimSpace(provider.BaseURL) == "" { if strings.TrimSpace(provider.BaseURL) == "" {
missing = append(missing, prefix+".base_url") missing = append(missing, prefix+".base_url")
} }
@@ -220,21 +231,56 @@ func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
models, modelConfigs := normalizeModelConfigs(raw.Models) models, modelConfigs := normalizeModelConfigs(raw.Models)
provider := ProviderConfig{ provider := ProviderConfig{
Name: name, Name: name,
APIType: normalizeAPIType(raw.APIType),
BaseURL: raw.BaseURL, BaseURL: raw.BaseURL,
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
} }
func normalizeAPIType(value string) string {
value = strings.ToLower(strings.TrimSpace(value))
if value == "" {
return APITypeChat
}
return value
}
func APITypes() []string {
return []string{APITypeChat, APITypeResponses}
}
func IsAPIType(value string) bool {
value = strings.ToLower(strings.TrimSpace(value))
for _, apiType := range APITypes() {
if value == apiType {
return true
}
}
return false
}
// DefaultThinkingParam returns the default request field used for the
// thinking level. Chat-completions providers historically use a top-level
// "thinking" field; the OpenAI Responses API uses "reasoning" with an
// {"effort": ...} value.
func DefaultThinkingParam(apiType string) string {
if apiType == APITypeResponses {
return "reasoning"
}
return "thinking"
}
func normalizeModelConfigs(rawModels []rawModelConfig) ([]string, map[string]ModelConfig) { 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))
+132
View File
@@ -370,3 +370,135 @@ providers:
t.Fatalf("err = %v", err) t.Fatalf("err = %v", err)
} }
} }
func TestLoadDefaultsAPITypeToChat(t *testing.T) {
dir := t.TempDir()
t.Setenv("AGENTU_TEST_KEY", "sk-test")
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
test:
base_url: https://example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.Providers["test"].APIType; got != APITypeChat {
t.Fatalf("api type = %q, want %q", got, APITypeChat)
}
}
func TestLoadParsesResponsesAPIType(t *testing.T) {
dir := t.TempDir()
t.Setenv("AGENTU_TEST_KEY", "sk-test")
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
openai:
api_type: RESPONSES
base_url: https://api.openai.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: gpt-test
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.Providers["openai"].APIType; got != APITypeResponses {
t.Fatalf("api type = %q, want %q", got, APITypeResponses)
}
}
func TestLoadRejectsInvalidAPIType(t *testing.T) {
dir := t.TempDir()
t.Setenv("AGENTU_TEST_KEY", "sk-test")
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
test:
api_type: completions
base_url: https://example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: test-model
`), 0o644); err != nil {
t.Fatal(err)
}
_, err := Load(path)
if err == nil || !strings.Contains(err.Error(), "api_type") {
t.Fatalf("err = %v", err)
}
}
func TestLoadDefaultsThinkingParamByAPIType(t *testing.T) {
dir := t.TempDir()
t.Setenv("AGENTU_TEST_KEY", "sk-test")
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
chat:
api_type: chat
base_url: https://chat.example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: chat-model
thinking: none
responses:
api_type: responses
base_url: https://responses.example.com
api_key: ${AGENTU_TEST_KEY}
models:
- id: responses-model
thinking: high
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.Providers["chat"].ThinkingParam; got != "thinking" {
t.Fatalf("chat thinking param = %q, want %q", got, "thinking")
}
if got := cfg.Providers["responses"].ThinkingParam; got != "reasoning" {
t.Fatalf("responses thinking param = %q, want %q", got, "reasoning")
}
}
func TestLoadCustomThinkingParam(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "agentu.yaml")
if err := os.WriteFile(path, []byte(`
providers:
custom:
api_type: chat
base_url: https://example.com
api_key: sk-test
thinking_param: reasoning_effort
models:
- id: custom-model
thinking: high
`), 0o644); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if err != nil {
t.Fatal(err)
}
if got := cfg.Providers["custom"].ThinkingParam; got != "reasoning_effort" {
t.Fatalf("thinking param = %q", got)
}
}
+481
View File
@@ -0,0 +1,481 @@
package llm
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"sort"
"strings"
)
const responsesEndpoint = "/v1/responses"
// OpenAIResponsesClient implements llm.Provider against the OpenAI Responses
// API (/v1/responses). It maps the agent's chat-style conversation history to
// Responses input items (messages, function_call, function_call_output) and
// converts Responses streaming events back into the shared StreamEvent shape.
type OpenAIResponsesClient struct {
baseURL string
apiKey string
httpClient *http.Client
}
func NewOpenAIResponsesClient(baseURL, apiKey string, httpClient *http.Client) *OpenAIResponsesClient {
if httpClient == nil {
httpClient = http.DefaultClient
}
return &OpenAIResponsesClient{
baseURL: strings.TrimRight(baseURL, "/"),
apiKey: apiKey,
httpClient: httpClient,
}
}
func (c *OpenAIResponsesClient) ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error {
body, err := json.Marshal(buildResponsesRequest(req))
if err != nil {
return fmt.Errorf("marshal responses request: %w", err)
}
state := &responsesStreamState{}
return postResponsesSSE(ctx, c.baseURL, c.apiKey, c.httpClient, body, func(payload string) error {
return state.handle(payload, emit)
})
}
func postResponsesSSE(ctx context.Context, baseURL, apiKey string, httpClient *http.Client, body []byte, handle func(payload string) error) error {
var lastErr error
for attempt := 1; attempt <= maxChatAttempts; attempt++ {
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+responsesEndpoint, bytes.NewReader(body))
if err != nil {
return fmt.Errorf("create responses request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
if apiKey != "" {
httpReq.Header.Set("Authorization", "Bearer "+apiKey)
}
resp, err := httpClient.Do(httpReq)
if err != nil {
lastErr = fmt.Errorf("send responses request: %w", err)
if !shouldRetryRequest(ctx, attempt, 0) {
return lastErr
}
if waitErr := waitBeforeRetry(ctx); waitErr != nil {
return waitErr
}
continue
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
limit := io.LimitReader(resp.Body, 4096)
data, _ := io.ReadAll(limit)
_ = resp.Body.Close()
lastErr = fmt.Errorf("responses request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
if !shouldRetryRequest(ctx, attempt, resp.StatusCode) {
return lastErr
}
if waitErr := waitBeforeRetry(ctx); waitErr != nil {
return waitErr
}
continue
}
defer resp.Body.Close()
return readResponsesSSE(resp.Body, handle)
}
return lastErr
}
func readResponsesSSE(r io.Reader, handle func(payload string) error) error {
reader := bufio.NewReader(r)
for {
line, err := reader.ReadString('\n')
if err != nil && !errors.Is(err, io.EOF) {
return fmt.Errorf("read responses stream: %w", err)
}
line = strings.TrimSpace(line)
if strings.HasPrefix(line, "data:") {
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if payload == "[DONE]" {
return nil
}
if handleErr := handle(payload); handleErr != nil {
return handleErr
}
}
if errors.Is(err, io.EOF) {
return nil
}
}
}
// buildResponsesRequest converts a ChatRequest into the /v1/responses body.
func buildResponsesRequest(req ChatRequest) map[string]any {
payload := map[string]any{
"model": req.Model,
"input": buildResponsesInput(req.Messages),
"stream": true,
}
if instructions := responsesInstructions(req.Messages); instructions != "" {
payload["instructions"] = instructions
}
if len(req.Tools) > 0 {
payload["tools"] = buildResponsesTools(req.Tools)
}
if req.ToolChoice != "" {
payload["tool_choice"] = req.ToolChoice
}
for key, value := range req.Extra {
if _, exists := payload[key]; exists {
continue
}
if key == "reasoning" {
if level, ok := value.(string); ok {
payload[key] = map[string]any{"effort": mapReasoningEffort(level)}
continue
}
}
payload[key] = value
}
return payload
}
// buildResponsesInput maps chat-style messages to Responses input items.
// Assistant tool calls become function_call items and tool results become
// function_call_output items, preserving their position in the conversation.
func buildResponsesInput(messages []Message) []any {
items := make([]any, 0, len(messages)+4)
for _, msg := range messages {
switch msg.Role {
case RoleSystem:
// System instructions are sent as the top-level "instructions"
// field and must not be duplicated inside the input array.
continue
case RoleUser:
items = append(items, responsesMessageInput{Role: "user", Content: msg.Content})
case RoleAssistant:
if msg.Content != "" {
items = append(items, responsesMessageInput{Role: "assistant", Content: msg.Content})
}
for _, call := range msg.ToolCalls {
items = append(items, responsesFunctionCallInput{
Type: "function_call",
CallID: call.ID,
Name: call.Function.Name,
Arguments: call.Function.Arguments,
})
}
case RoleTool:
items = append(items, responsesFunctionCallOutput{
Type: "function_call_output",
CallID: msg.ToolCallID,
Output: msg.Content,
})
}
}
return items
}
// responsesInstructions extracts system messages into a single instructions
// string, which is the OpenAI Responses API way to pass developer/system
// guidance (and the most prompt-cache friendly).
func responsesInstructions(messages []Message) string {
var instructions []string
for _, msg := range messages {
if msg.Role != RoleSystem {
continue
}
if text := strings.TrimSpace(msg.Content); text != "" {
instructions = append(instructions, text)
}
}
return strings.Join(instructions, "\n\n")
}
func buildResponsesTools(tools []Tool) []responsesTool {
out := make([]responsesTool, 0, len(tools))
for _, tool := range tools {
out = append(out, responsesTool{
Type: "function",
Name: tool.Function.Name,
Description: tool.Function.Description,
Parameters: tool.Function.Parameters,
})
}
return out
}
func mapReasoningEffort(level string) string {
level = strings.ToLower(strings.TrimSpace(level))
if level == "middle" {
return "medium"
}
return level
}
type responsesMessageInput struct {
Role string `json:"role"`
Content string `json:"content"`
}
type responsesFunctionCallInput struct {
Type string `json:"type"`
CallID string `json:"call_id"`
Name string `json:"name"`
Arguments string `json:"arguments"`
}
type responsesFunctionCallOutput struct {
Type string `json:"type"`
CallID string `json:"call_id"`
Output string `json:"output"`
}
type responsesTool struct {
Type string `json:"type"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters json.RawMessage `json:"parameters,omitempty"`
}
// responsesStreamState tracks in-flight output items while parsing SSE events.
type responsesStreamState struct {
calls map[int]*responsesCallBuilder
done bool
}
type responsesCallBuilder struct {
index int
id string
callType string
name string
arguments strings.Builder
finalized bool
}
func (s *responsesStreamState) handle(payload string, emit func(StreamEvent) error) error {
if s.done {
return nil
}
var event responsesStreamEvent
if err := json.Unmarshal([]byte(payload), &event); err != nil {
return fmt.Errorf("parse responses stream payload: %w", err)
}
switch event.Type {
case "error":
s.done = true
return event.err()
case "response.failed":
s.done = true
if event.Response != nil && event.Response.Error != nil {
return fmt.Errorf("provider error: %s: %s", event.Response.Error.Code, event.Response.Error.Message)
}
return fmt.Errorf("provider error: response failed")
case "response.incomplete":
s.done = true
return s.finish(emit, "length", nil)
case "response.completed":
s.done = true
var usage *Usage
if event.Response != nil && event.Response.Usage != nil {
usage = &Usage{
PromptTokens: event.Response.Usage.InputTokens,
CompletionTokens: event.Response.Usage.OutputTokens,
TotalTokens: event.Response.Usage.TotalTokens,
}
}
return s.finish(emit, "stop", usage)
case "response.output_text.delta":
if event.Delta != "" {
return emit(StreamEvent{Content: event.Delta})
}
case "response.output_item.added":
if event.Item != nil && event.Item.Type == "function_call" {
s.builder(event.OutputIndex).applyAdded(*event.Item)
}
case "response.function_call_arguments.delta":
if event.Delta != "" {
s.builder(event.OutputIndex).arguments.WriteString(event.Delta)
}
case "response.function_call_arguments.done":
builder := s.builder(event.OutputIndex)
if event.Name != "" {
builder.name = event.Name
}
if event.Arguments != "" {
builder.arguments.Reset()
builder.arguments.WriteString(event.Arguments)
}
case "response.output_item.done":
if event.Item != nil && event.Item.Type == "function_call" {
builder := s.builder(event.OutputIndex)
builder.applyDone(*event.Item)
return s.emitCall(builder, emit)
}
}
return nil
}
// finish flushes any not-yet-finalized function calls and reports the
// terminal finish reason.
func (s *responsesStreamState) finish(emit func(StreamEvent) error, finishReason string, usage *Usage) error {
if len(s.calls) == 0 {
return emit(StreamEvent{FinishReason: finishReason, Usage: usage})
}
indexes := make([]int, 0, len(s.calls))
for index := range s.calls {
indexes = append(indexes, index)
}
sort.Ints(indexes)
var deltas []ToolCallDelta
for _, index := range indexes {
builder := s.calls[index]
if builder.finalized {
continue
}
builder.finalized = true
deltas = append(deltas, builder.delta())
}
if len(deltas) > 0 {
if err := emit(StreamEvent{ToolCalls: deltas}); err != nil {
return err
}
}
return emit(StreamEvent{FinishReason: finishReason, Usage: usage})
}
func (s *responsesStreamState) emitCall(builder *responsesCallBuilder, emit func(StreamEvent) error) error {
if builder.finalized {
return nil
}
builder.finalized = true
return emit(StreamEvent{ToolCalls: []ToolCallDelta{builder.delta()}})
}
func (s *responsesStreamState) builder(index int) *responsesCallBuilder {
if s.calls == nil {
s.calls = make(map[int]*responsesCallBuilder)
}
builder := s.calls[index]
if builder == nil {
builder = &responsesCallBuilder{index: index}
s.calls[index] = builder
}
return builder
}
func (b *responsesCallBuilder) applyAdded(item responsesStreamItem) {
if b.id == "" {
b.id = item.CallID
if b.id == "" {
b.id = item.ID
}
}
if b.name == "" {
b.name = item.Name
}
if item.Arguments != "" {
b.arguments.Reset()
b.arguments.WriteString(item.Arguments)
}
}
func (b *responsesCallBuilder) applyDone(item responsesStreamItem) {
if item.CallID != "" {
b.id = item.CallID
} else if item.ID != "" {
b.id = item.ID
}
if item.Name != "" {
b.name = item.Name
}
if item.Arguments != "" {
b.arguments.Reset()
b.arguments.WriteString(item.Arguments)
}
}
func (b *responsesCallBuilder) delta() ToolCallDelta {
id := b.id
if id == "" {
id = fmt.Sprintf("call_%d", b.index)
}
callType := b.callType
if callType == "" {
callType = "function"
}
return ToolCallDelta{
Index: b.index,
ID: id,
Type: callType,
Name: b.name,
Arguments: b.arguments.String(),
}
}
type responsesStreamEvent struct {
Type string `json:"type"`
OutputIndex int `json:"output_index"`
Delta string `json:"delta"`
Arguments string `json:"arguments"`
Name string `json:"name"`
Code string `json:"code"`
Message string `json:"message"`
Param string `json:"param"`
Item *responsesStreamItem `json:"item"`
Response *responsesStreamResponse `json:"response"`
Error *responsesStreamError `json:"error"`
}
type responsesStreamItem struct {
Type string `json:"type"`
ID string `json:"id"`
CallID string `json:"call_id"`
Name string `json:"name"`
Arguments string `json:"arguments"`
}
type responsesStreamResponse struct {
Status string `json:"status"`
Error *responsesStreamError `json:"error"`
Usage *responsesStreamUsage `json:"usage"`
}
type responsesStreamUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
}
type responsesStreamError struct {
Type string `json:"type"`
Code string `json:"code"`
Message string `json:"message"`
Param string `json:"param"`
}
func (e responsesStreamEvent) err() error {
code, message, param := e.Code, e.Message, e.Param
if code == "" && message == "" && e.Error != nil {
code, message, param = e.Error.Code, e.Error.Message, e.Error.Param
}
if message == "" {
if param != "" {
return fmt.Errorf("provider error: %s", param)
}
return fmt.Errorf("provider error")
}
if code != "" {
return fmt.Errorf("provider error: %s: %s", code, message)
}
return fmt.Errorf("provider error: %s", message)
}
+337
View File
@@ -0,0 +1,337 @@
package llm
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestBuildResponsesRequest(t *testing.T) {
req := ChatRequest{
Model: "gpt-test",
Messages: []Message{
{Role: RoleSystem, Content: "be concise"},
{Role: RoleUser, Content: "hi"},
{
Role: RoleAssistant,
Content: "let me check",
ToolCalls: []ToolCall{{
ID: "call_1",
Type: "function",
Function: FunctionCall{
Name: "file_read",
Arguments: `{"path":"README.md"}`,
},
}},
},
{Role: RoleTool, ToolCallID: "call_1", Content: "contents"},
{Role: RoleAssistant, Content: "done"},
},
Tools: []Tool{{
Type: "function",
Function: ToolFunction{
Name: "file_read",
Description: "Read a local file",
Parameters: json.RawMessage(`{"type":"object","properties":{}}`),
},
}},
ToolChoice: "auto",
Extra: map[string]any{
"reasoning": "middle",
"thinking": "high",
"temperature": 0.2,
},
}
payload := buildResponsesRequest(req)
data, err := json.Marshal(payload)
if err != nil {
t.Fatal(err)
}
var raw map[string]any
if err := json.Unmarshal(data, &raw); err != nil {
t.Fatal(err)
}
if raw["model"] != "gpt-test" || raw["stream"] != true {
t.Fatalf("model/stream = %#v / %#v", raw["model"], raw["stream"])
}
if raw["instructions"] != "be concise" {
t.Fatalf("instructions = %#v", raw["instructions"])
}
if raw["tool_choice"] != "auto" {
t.Fatalf("tool_choice = %#v", raw["tool_choice"])
}
input, ok := raw["input"].([]any)
if !ok || len(input) != 5 {
t.Fatalf("input = %#v", raw["input"])
}
checkInputItem := func(index int, wantType, wantKey, wantValue string) {
t.Helper()
item, ok := input[index].(map[string]any)
if !ok {
t.Fatalf("input[%d] = %#v", index, input[index])
}
if item[wantKey] != wantValue {
t.Fatalf("input[%d].%s = %#v, want %q", index, wantKey, item[wantKey], wantValue)
}
if wantType != "" && item["type"] != wantType {
t.Fatalf("input[%d].type = %#v, want %q", index, item["type"], wantType)
}
}
checkInputItem(0, "", "role", "user")
checkInputItem(1, "", "role", "assistant")
checkInputItem(2, "function_call", "call_id", "call_1")
checkInputItem(3, "function_call_output", "call_id", "call_1")
checkInputItem(4, "", "role", "assistant")
if item := input[2].(map[string]any); item["name"] != "file_read" || item["arguments"] != `{"path":"README.md"}` {
t.Fatalf("function_call item = %#v", item)
}
if item := input[3].(map[string]any); item["output"] != "contents" {
t.Fatalf("function_call_output item = %#v", item)
}
tools, ok := raw["tools"].([]any)
if !ok || len(tools) != 1 {
t.Fatalf("tools = %#v", raw["tools"])
}
tool := tools[0].(map[string]any)
if tool["type"] != "function" || tool["name"] != "file_read" || tool["description"] != "Read a local file" {
t.Fatalf("tool = %#v", tool)
}
if _, ok := tool["parameters"].(map[string]any); !ok {
t.Fatalf("tool parameters = %#v", tool["parameters"])
}
reasoning, ok := raw["reasoning"].(map[string]any)
if !ok || reasoning["effort"] != "medium" {
t.Fatalf("reasoning = %#v", raw["reasoning"])
}
if raw["thinking"] != "high" {
t.Fatalf("thinking = %#v", raw["thinking"])
}
if raw["temperature"] != 0.2 {
t.Fatalf("temperature = %#v", raw["temperature"])
}
}
func TestResponsesInstructionsCombinesSystemMessages(t *testing.T) {
payload := buildResponsesRequest(ChatRequest{
Model: "gpt-test",
Messages: []Message{
{Role: RoleSystem, Content: "first instruction"},
{Role: RoleUser, Content: "hi"},
{Role: RoleSystem, Content: "second instruction"},
},
})
data, err := json.Marshal(payload)
if err != nil {
t.Fatal(err)
}
var raw map[string]any
if err := json.Unmarshal(data, &raw); err != nil {
t.Fatal(err)
}
if raw["instructions"] != "first instruction\n\nsecond instruction" {
t.Fatalf("instructions = %#v", raw["instructions"])
}
input, ok := raw["input"].([]any)
if !ok || len(input) != 1 {
t.Fatalf("input = %#v", raw["input"])
}
item, ok := input[0].(map[string]any)
if !ok || item["role"] != "user" {
t.Fatalf("input[0] = %#v", input[0])
}
}
func TestOpenAIResponsesClientStreamsContentAndToolCalls(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/responses" {
t.Fatalf("path = %s", r.URL.Path)
}
if got := r.Header.Get("Authorization"); got != "Bearer sk-test" {
t.Fatalf("authorization = %q", got)
}
body := mustReadBody(t, r)
var raw map[string]any
if err := json.NewDecoder(strings.NewReader(body)).Decode(&raw); err != nil {
t.Fatal(err)
}
if raw["stream"] != true {
t.Fatal("request did not enable stream")
}
if raw["tool_choice"] != "auto" {
t.Fatalf("tool_choice = %#v", raw["tool_choice"])
}
input, ok := raw["input"].([]any)
if !ok || len(input) == 0 {
t.Fatalf("input = %#v", raw["input"])
}
w.Header().Set("Content-Type", "text/event-stream")
events := []string{
`{"type":"response.output_item.added","output_index":0,"item":{"type":"message","id":"msg_1","role":"assistant","content":[]}}`,
`{"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hel"}`,
`{"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"lo"}`,
`{"type":"response.output_item.added","output_index":1,"item":{"type":"function_call","id":"fc_1","call_id":"call_1","name":"file_read","arguments":""}}`,
`{"type":"response.function_call_arguments.delta","item_id":"fc_1","output_index":1,"delta":"{\"path\""}`,
`{"type":"response.function_call_arguments.delta","item_id":"fc_1","output_index":1,"delta":":\"README.md\"}"}`,
`{"type":"response.function_call_arguments.done","item_id":"fc_1","output_index":1,"name":"file_read","arguments":"{\"path\":\"README.md\"}"}`,
`{"type":"response.output_item.done","output_index":1,"item":{"type":"function_call","id":"fc_1","call_id":"call_1","name":"file_read","arguments":"{\"path\":\"README.md\"}"}}`,
`{"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_1","role":"assistant","content":[{"type":"output_text","text":"Hello"}]}}`,
`{"type":"response.completed","response":{"id":"resp_1","status":"completed","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}}`,
}
for _, event := range events {
_, _ = w.Write([]byte("data: " + event + "\n\n"))
}
}))
defer server.Close()
client := NewOpenAIResponsesClient(server.URL, "sk-test", server.Client())
var content strings.Builder
var calls []ToolCallDelta
var finishes []string
var usage *Usage
err := client.ChatStream(context.Background(), ChatRequest{
Model: "gpt-test",
Messages: []Message{{Role: RoleUser, Content: "hi"}},
ToolChoice: "auto",
Tools: []Tool{{
Type: "function",
Function: ToolFunction{
Name: "file_read",
Description: "Read a local file",
Parameters: json.RawMessage(`{"type":"object"}`),
},
}},
}, func(event StreamEvent) error {
content.WriteString(event.Content)
calls = append(calls, event.ToolCalls...)
if event.Usage != nil {
usage = event.Usage
}
if event.FinishReason != "" {
finishes = append(finishes, event.FinishReason)
}
return nil
})
if err != nil {
t.Fatal(err)
}
if content.String() != "Hello" {
t.Fatalf("content = %q", content.String())
}
if len(calls) != 1 {
t.Fatalf("calls = %#v", calls)
}
call := calls[0]
if call.Index != 1 || call.ID != "call_1" || call.Name != "file_read" {
t.Fatalf("call = %#v", call)
}
if call.Arguments != `{"path":"README.md"}` {
t.Fatalf("arguments = %q", call.Arguments)
}
if len(finishes) != 1 || finishes[0] != "stop" {
t.Fatalf("finishes = %#v", finishes)
}
if usage == nil || usage.PromptTokens != 10 || usage.CompletionTokens != 20 || usage.TotalTokens != 30 {
t.Fatalf("usage = %#v", usage)
}
}
func TestOpenAIResponsesClientErrorEvent(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"error","code":"invalid_request_error","message":"bad request"}` + "\n\n"))
}))
defer server.Close()
client := NewOpenAIResponsesClient(server.URL, "sk-test", server.Client())
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(StreamEvent) error {
return nil
})
if err == nil || !strings.Contains(err.Error(), "bad request") {
t.Fatalf("err = %v", err)
}
}
func TestOpenAIResponsesClientFailedEvent(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"server_error","message":"boom"}}}` + "\n\n"))
}))
defer server.Close()
client := NewOpenAIResponsesClient(server.URL, "sk-test", server.Client())
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(StreamEvent) error {
return nil
})
if err == nil || !strings.Contains(err.Error(), "boom") {
t.Fatalf("err = %v", err)
}
}
func TestOpenAIResponsesClientIncomplete(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"partial"}` + "\n\n"))
_, _ = w.Write([]byte(`data: {"type":"response.incomplete","response":{"id":"resp_1","status":"incomplete","incomplete_details":{"reason":"max_tokens"}}}` + "\n\n"))
}))
defer server.Close()
client := NewOpenAIResponsesClient(server.URL, "sk-test", server.Client())
var content strings.Builder
var finishes []string
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
content.WriteString(event.Content)
if event.FinishReason != "" {
finishes = append(finishes, event.FinishReason)
}
return nil
})
if err != nil {
t.Fatal(err)
}
if content.String() != "partial" {
t.Fatalf("content = %q", content.String())
}
if len(finishes) != 1 || finishes[0] != "length" {
t.Fatalf("finishes = %#v", finishes)
}
}
func TestOpenAIResponsesClientRetriesTransientHTTPError(t *testing.T) {
attempts := 0
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
attempts++
if attempts == 1 {
http.Error(w, "temporary", http.StatusBadGateway)
return
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(`data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"ok"}` + "\n\n"))
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed"}}` + "\n\n"))
}))
defer server.Close()
client := NewOpenAIResponsesClient(server.URL, "sk-test", server.Client())
var content strings.Builder
err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
content.WriteString(event.Content)
return nil
})
if err != nil {
t.Fatal(err)
}
if attempts != 2 {
t.Fatalf("attempts = %d, want 2", attempts)
}
if content.String() != "ok" {
t.Fatalf("content = %q", content.String())
}
}