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

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

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

Documentation:
- README updated with git tools and file_edit descriptions
2026-06-24 19:47:14 -07:00
loveuer 7b8bad5e62 feat: crush-inspired TUI improvements, activity bar polish, startup fresh session
TUI improvements:
- semantic icon constants (✓ ◆ ▣ ✕ ⌘ ◌ ◔ ◉)
- pill-style metadata bar with model/provider/session/context
- fuzzy command palette with deterministic subsequence scoring
- input history navigation (ctrl+p / ctrl+n)
- structured tool log cards (▣ tool_name + args preview)
- no-background activity indicator and meta bar for cleaner rendering

Startup behavior:
- default startup creates a fresh session
- explicit --resume <id> restores saved context
- added main_test.go for initializeSession coverage
2026-06-24 18:16:53 -07:00
46 changed files with 9238 additions and 524 deletions
+1
View File
@@ -3,3 +3,4 @@ agentu*.yaml
/coverage.out
*.swp
chat-history.txt
dist/
+92 -15
View File
@@ -1,18 +1,24 @@
# agentu
agentu is a small local agent written in Go. The current MVP is a terminal chat
agentu is a small local coding agent written in Go. It provides a terminal chat
UI with OpenAI-compatible providers, streaming responses, session-level
model/provider switching, and local tools.
model/provider switching, skills, approvals, and local tools.
## Current MVP
## v0.3.1
- TUI-first chat experience with a light theme by default.
- OpenAI-compatible `/v1/chat/completions` provider support.
- OpenAI Responses API (`/v1/responses`) provider support.
- Multiple providers from `~/.agentu/config.yaml`.
- Session-only provider, model, and thinking-level switching.
- Read-only project tools enabled by default.
- Interactive approval for model-requested file writes and shell commands.
- Atomic file writes, per-file rollback snapshots, and multi-file patches.
- Project outlines with Go, TypeScript/JavaScript, and Rust symbols.
- Global Codex-style skills loaded from `~/.agentu/skills`.
- Persistent session index and in-memory task tracking.
- Built-in web search and URL fetching.
- `--yolo` mode for 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
@@ -23,7 +29,15 @@ export AGENTU_API_KEY='your-api-key'
go run .
```
Use `--yolo` only when you want agentu to write files or run shell commands:
Use dark mode when your terminal is dark:
```sh
go run . --theme dark
```
By default, the TUI asks for approval before every model-requested file change
or shell command. Use `--yolo` only when those operations should run without
confirmation:
```sh
go run . --yolo
@@ -42,6 +56,7 @@ Use `--config <path>` to load another file.
```yaml
providers:
loveuer:
api_type: chat
base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY}
models:
@@ -58,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
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
entry is an object with:
@@ -77,7 +106,11 @@ none, middle, high, xhigh, max
```
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
@@ -88,28 +121,38 @@ go run . [flags]
- `--config <path>` loads a YAML config file. Default: `~/.agentu/config.yaml`.
- `--theme light|dark` selects the TUI theme. Default: `light`.
- `--plain` uses the simple line-based REPL instead of the TUI.
- `--yolo` enables automatic file writes and shell execution.
- `--yolo` automatically approves model-requested file writes and shell execution.
- `--resume <id>` resumes a saved session from `~/.agentu/sessions`.
## TUI Controls
- `enter` sends the message.
- `ctrl+j` or `alt+enter` inserts a newline.
- `ctrl+j` inserts a newline.
- `up` recalls the previous input; `ctrl+p` / `ctrl+n` navigate input history.
- The input starts at one line and grows up to six lines.
- `page up` / `page down` scrolls the transcript.
- `esc` cancels a running response.
- While a task is running, `/btw <question>` asks a concurrent side question.
- Entering only `.` means “continue with the current task.”
- At a tool approval prompt, `y` approves and `n` or `esc` denies the operation.
- Type `/` to show available commands.
- In `--yolo` mode, type `! <command>` to run a shell command from the configured working directory.
Slash commands:
- `/btw <question>` answers from a snapshot of the current context without adding the question or answer to conversation history. It never exposes tools and can run alongside the main task.
- `/clear` resets conversation history.
- `/compact` summarizes older conversation history and keeps recent context.
- `/compact` summarizes older conversation history and keeps the newest ~20% of context, with a small recent-message floor.
- `/status` shows changed files with `git status --short`.
- `/diff [--stat] [path...]` shows the unstaged git diff.
- `/test [command...]` runs a project test command; defaults to `go test ./...` for Go modules or `npm test` for Node projects.
- `/skills` lists global skills.
- `/skills install <path-or-url>` installs a skill and activates it in the current conversation.
- `/rollback [id]` restores a file snapshot; without an ID it restores the newest snapshot.
- `/model` shows current provider, model, thinking, and switch usage.
- `/model provider <name>` switches provider for the current session.
- `/model model <name>` or `/model <name>` switches model for the current session.
In the TUI, type `/model ` or `/model model ` to choose from configured model IDs in the candidate palette.
- `/model thinking <none|middle|high|xhigh|max>` changes thinking for the current session.
- `/sessions` lists saved sessions.
- `/resume [id]` resumes by ID or opens an interactive picker.
@@ -119,22 +162,36 @@ Slash commands:
## Local Tools
Read-only tools are available by default:
These tools are always available:
- `file_read`
- `file_list`
- `file_search` with optional `context_lines` and `max_results`
- `code_symbols` for Go packages, imports, types, functions, and methods
- `project_outline` for a gitignore-aware tree with Go, TypeScript/JavaScript, and Rust symbols
- `todo` for conversation-local task tracking
- `git_status` shows changed files
- `git_diff` shows unstaged or staged diff output
- `git_log` shows recent commit history
Web tools are also available in read-only mode:
Web tools are also available:
- `web_search`
- `web_fetch`
`--yolo` additionally enables:
In the TUI, these tools require approval by default and run automatically with
`--yolo`:
- `file_write`
- `shell_run`
- `file_write` for creating or replacing whole files
- `file_edit` for surgical line-range or pattern edits to existing files
- `apply_patch` for atomic-style multi-file V4 patches, including dry runs
- `shell_run` for non-interactive shell commands; output streams live into the
TUI while the command runs instead of appearing only after completion
Before a file mutation, agentu stores an in-memory snapshot and includes its
`rb_...` ID in the tool result. Use `/rollback` or `/rollback <id>` to restore
it. Snapshots last for the lifetime of the process. The plain REPL remains
read-only unless `--yolo` is supplied because it has no interactive approval UI.
Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
example:
@@ -143,6 +200,21 @@ example:
请帮我运行 `go test ./...` 并总结结果
```
## Global Skills
At startup, agentu loads both Codex-style directories containing `SKILL.md`
and legacy flat Markdown files from `~/.agentu/skills`. Skill instructions are
added to the system prompt. Install a local directory, Markdown file, or direct
HTTP(S) Markdown URL from the TUI:
```text
/skills install ./my-skill
/skills install https://example.com/SKILL.md
```
Directory installs preserve supporting files such as `references/`; symlinks
inside a source directory are rejected.
## Development
Run the verification suite:
@@ -159,5 +231,10 @@ Local config and build outputs are intentionally ignored by git.
- Provider/model/thinking changes are session-only; agentu does not write them
back to `~/.agentu/config.yaml`.
- `api_type` is per provider and cannot be changed at runtime; use
`/model provider <name>` to switch to a provider configured with a different
API type.
- The config schema is intentionally strict during MVP development. Deprecated
fields such as `provider` or `active_provider` are rejected.
- Session metadata is indexed in `~/.agentu/sessions/index.json`; complete
conversations remain stored as one JSON file per session.
+3
View File
@@ -1,5 +1,8 @@
providers:
loveuer:
# API protocol: "chat" (OpenAI-compatible /v1/chat/completions, default)
# or "responses" (OpenAI Responses API /v1/responses).
api_type: chat
base_url: https://ai.loveuer.com
api_key: ${AGENTU_API_KEY}
models:
+448 -45
View File
@@ -7,18 +7,25 @@ import (
"io"
"sort"
"strings"
"sync"
"time"
"agentu/pkg/llm"
"agentu/internal/tools"
"agentu/pkg/llm"
)
const (
defaultMaxToolRounds = 50
doomLoopThreshold = 3
compactRecentMessages = 6
compactMinRecentMessages = 6
compactRecentRetentionPercent = 20
summaryMessageLimit = 8000
contextSummaryPrefix = "[context summary]"
toolLogArgumentLimit = 400
toolLogOutputPreviewBytes = 4096
toolLogOutputMarker = "[output]"
toolStreamLogLimit = tools.MaxToolOutputBytes
continuationPrompt = "Continue with the current task."
summarizerSystemPrompt = "You are a conversation summarizer. Summarize the following conversation history into a concise but detailed summary. Preserve:\n" +
"- Key facts, decisions, and conclusions\n" +
@@ -29,7 +36,13 @@ const (
"Output only the summary, no preamble."
)
type toolExecutionResult struct {
call llm.ToolCall
output string
}
type Agent struct {
mu sync.RWMutex
provider llm.Provider
providerName string
model string
@@ -38,6 +51,8 @@ type Agent struct {
thinkingEnabled bool
systemPrompt string
toolRegistry *tools.Registry
authorizer ToolAuthorizer
rollbackStore *tools.RollbackStore
toolTimeout time.Duration
maxToolRounds int
maxContextTokens int
@@ -54,6 +69,8 @@ type Options struct {
ThinkingEnabled bool
SystemPrompt string
ToolRegistry *tools.Registry
Authorizer ToolAuthorizer
RollbackStore *tools.RollbackStore
ToolTimeout time.Duration
MaxToolIterations int
MaxContextTokens int
@@ -77,6 +94,8 @@ func New(opts Options) *Agent {
thinkingEnabled: opts.ThinkingEnabled,
systemPrompt: opts.SystemPrompt,
toolRegistry: opts.ToolRegistry,
authorizer: opts.Authorizer,
rollbackStore: opts.RollbackStore,
toolTimeout: toolTimeout,
maxToolRounds: maxToolRounds,
maxContextTokens: opts.MaxContextTokens,
@@ -105,6 +124,8 @@ type RuntimeInfo struct {
}
func (a *Agent) SetRuntime(opts RuntimeOptions) {
a.mu.Lock()
defer a.mu.Unlock()
if opts.Provider != nil {
a.provider = opts.Provider
}
@@ -121,6 +142,8 @@ func (a *Agent) SetRuntime(opts RuntimeOptions) {
}
func (a *Agent) RuntimeInfo() RuntimeInfo {
a.mu.RLock()
defer a.mu.RUnlock()
return RuntimeInfo{
ProviderName: a.providerName,
Model: a.model,
@@ -132,6 +155,12 @@ func (a *Agent) RuntimeInfo() RuntimeInfo {
}
func (a *Agent) Clear() {
a.mu.Lock()
defer a.mu.Unlock()
a.clear()
}
func (a *Agent) clear() {
a.lastUsage = nil
a.messages = nil
if strings.TrimSpace(a.systemPrompt) != "" {
@@ -142,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.
func (a *Agent) Messages() []llm.Message {
a.mu.RLock()
defer a.mu.RUnlock()
out := make([]llm.Message, len(a.messages))
copy(out, a.messages)
return out
@@ -151,11 +199,23 @@ func (a *Agent) Messages() []llm.Message {
// SetMessages replaces the conversation messages (used for session restore).
func (a *Agent) SetMessages(messages []llm.Message) {
a.mu.Lock()
defer a.mu.Unlock()
a.messages = append([]llm.Message(nil), messages...)
if strings.TrimSpace(a.systemPrompt) == "" {
return
}
if len(a.messages) > 0 && a.messages[0].Role == llm.RoleSystem {
a.messages[0].Content = a.systemPrompt
return
}
a.messages = append([]llm.Message{{Role: llm.RoleSystem, Content: a.systemPrompt}}, a.messages...)
}
// LastUsage returns a copy of the most recent provider token usage data.
func (a *Agent) LastUsage() *llm.Usage {
a.mu.RLock()
defer a.mu.RUnlock()
if a.lastUsage == nil {
return nil
}
@@ -165,6 +225,12 @@ func (a *Agent) LastUsage() *llm.Usage {
// SetLastUsage restores the most recent provider token usage data.
func (a *Agent) SetLastUsage(usage *llm.Usage) {
a.mu.Lock()
defer a.mu.Unlock()
a.setLastUsage(usage)
}
func (a *Agent) setLastUsage(usage *llm.Usage) {
if usage == nil {
a.lastUsage = nil
return
@@ -173,34 +239,90 @@ func (a *Agent) SetLastUsage(usage *llm.Usage) {
a.lastUsage = &copy
}
// compactRetainedTokenBudget returns the number of estimated tokens to retain
// based on compactRecentRetentionPercent, using integer ceiling division.
func compactRetainedTokenBudget(totalTokens int) int {
if totalTokens <= 0 {
return 1
}
return (totalTokens*compactRecentRetentionPercent + 99) / 100
}
// compactTurnBoundary walks backward from start until it finds a user message
// or reaches prefixEnd, ensuring the retained suffix starts at a user turn.
func compactTurnBoundary(messages []llm.Message, prefixEnd int, start int) int {
for start > prefixEnd && messages[start].Role != llm.RoleUser {
start--
}
return start
}
// compactRecentStart returns the index where the retained recent suffix begins.
// It balances the 20% token-budget ratio with the compactMinRecentMessages floor
// and aligns to a user-turn boundary.
func compactRecentStart(messages []llm.Message, prefixEnd int) int {
compactableCount := len(messages) - prefixEnd
if compactableCount <= compactMinRecentMessages {
return len(messages)
}
compactable := messages[prefixEnd:]
totalTokens := llm.EstimateTokens(compactable, nil)
retainTokens := compactRetainedTokenBudget(totalTokens)
// Walk backward from the end to find the earliest suffix fitting the budget.
ratioStart := len(messages) - 1
for i := len(messages) - 1; i >= prefixEnd; i-- {
if llm.EstimateTokens(messages[i:], nil) <= retainTokens {
ratioStart = i
} else {
break
}
}
// Enforce the minimum recent-message floor.
minStart := len(messages) - compactMinRecentMessages
if minStart < prefixEnd {
minStart = prefixEnd
}
// Choose the earlier start so the floor is never weakened.
start := minStart
if ratioStart < start {
start = ratioStart
}
return compactTurnBoundary(messages, prefixEnd, start)
}
// Compact summarizes older conversation messages while preserving the system
// prompt and recent messages intact.
func (a *Agent) Compact(ctx context.Context) error {
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
}
if a.provider == nil {
if provider == nil {
return fmt.Errorf("provider is not configured")
}
prefixEnd := 0
if a.messages[0].Role == llm.RoleSystem {
if messages[0].Role == llm.RoleSystem {
prefixEnd = 1
}
if len(a.messages)-prefixEnd <= compactRecentMessages {
recentStart := compactRecentStart(messages, prefixEnd)
if recentStart <= prefixEnd || recentStart >= len(messages) {
return nil
}
recentStart := len(a.messages) - compactRecentMessages
if recentStart <= prefixEnd {
return nil
}
middle := append([]llm.Message(nil), a.messages[prefixEnd:recentStart]...)
middle := append([]llm.Message(nil), messages[prefixEnd:recentStart]...)
if len(middle) == 0 {
return nil
}
summary, usage, err := a.summarizeMessages(ctx, middle)
summary, usage, err := summarizeMessages(ctx, provider, model, middle)
if err != nil {
return fmt.Errorf("compact context: %w", err)
}
@@ -209,13 +331,15 @@ func (a *Agent) Compact(ctx context.Context) error {
return fmt.Errorf("compact context: summary was empty")
}
compressed := make([]llm.Message, 0, prefixEnd+1+len(a.messages[recentStart:]))
compressed = append(compressed, a.messages[:prefixEnd]...)
compressed := make([]llm.Message, 0, prefixEnd+1+len(messages[recentStart:]))
compressed = append(compressed, messages[:prefixEnd]...)
compressed = append(compressed, llm.Message{
Role: llm.RoleAssistant,
Content: contextSummaryPrefix + "\n" + summary,
})
compressed = append(compressed, a.messages[recentStart:]...)
compressed = append(compressed, messages[recentStart:]...)
a.mu.Lock()
defer a.mu.Unlock()
a.messages = compressed
if usage != nil {
a.lastUsage = usage
@@ -224,6 +348,8 @@ func (a *Agent) Compact(ctx context.Context) error {
}
func (a *Agent) shouldCompact() bool {
a.mu.RLock()
defer a.mu.RUnlock()
if a.maxContextTokens <= 0 {
return false
}
@@ -233,9 +359,9 @@ func (a *Agent) shouldCompact() bool {
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{
Model: a.model,
Model: model,
Messages: []llm.Message{
{Role: llm.RoleSystem, Content: summarizerSystemPrompt},
{Role: llm.RoleUser, Content: formatMessagesForSummary(messages)},
@@ -244,7 +370,7 @@ func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (
var summary strings.Builder
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 != "" {
summary.WriteString(event.Content)
}
@@ -292,12 +418,93 @@ func roundSignature(calls []llm.ToolCall) 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 {
input = strings.TrimSpace(input)
if input == "" {
return nil
}
a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input})
a.appendMessage(llm.Message{Role: llm.RoleUser, Content: input})
if a.shouldCompact() {
if err := a.Compact(ctx); err != nil {
if logs != nil {
@@ -314,14 +521,14 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
return err
}
if len(calls) == 0 {
a.messages = append(a.messages, llm.Message{
a.appendMessage(llm.Message{
Role: llm.RoleAssistant,
Content: content,
})
return nil
}
a.messages = append(a.messages, llm.Message{
a.appendMessage(llm.Message{
Role: llm.RoleAssistant,
Content: content,
ToolCalls: calls,
@@ -356,7 +563,7 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
strings.Join(toolNames, ", "),
doomLoopThreshold,
)
a.messages = append(a.messages, llm.Message{
a.appendMessage(llm.Message{
Role: llm.RoleAssistant,
Content: msg,
})
@@ -364,14 +571,15 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
}
}
for _, call := range calls {
result := a.executeTool(ctx, call, logs)
a.messages = append(a.messages, llm.Message{
results := a.executeTools(ctx, calls, logs)
for _, result := range results {
a.appendMessage(llm.Message{
Role: llm.RoleTool,
ToolCallID: call.ID,
Content: result,
ToolCallID: result.call.ID,
Content: result.output,
})
}
}
return fmt.Errorf("tool round limit reached after %d rounds; the task may be too complex for a single turn", a.maxToolRounds)
@@ -381,23 +589,33 @@ func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []l
var content strings.Builder
calls := make(map[int]*toolCallBuilder)
a.mu.RLock()
provider := a.provider
model := a.model
messages := append([]llm.Message(nil), a.messages...)
thinking := a.thinking
thinkingParam := a.thinkingParam
thinkingEnabled := a.thinkingEnabled
a.mu.RUnlock()
if provider == nil {
return "", nil, fmt.Errorf("provider is not configured")
}
req := llm.ChatRequest{
Model: a.model,
Messages: append([]llm.Message(nil), a.messages...),
Model: model,
Messages: expandContinuationMessages(messages),
Tools: a.toolDefinitions(),
}
if a.thinkingEnabled {
thinkingParam := a.thinkingParam
if thinkingEnabled {
if thinkingParam == "" {
thinkingParam = "thinking"
}
req.Extra = map[string]any{thinkingParam: a.thinking}
req.Extra = map[string]any{thinkingParam: thinking}
}
if len(req.Tools) > 0 {
req.ToolChoice = "auto"
}
err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
err := provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
if event.Usage != nil {
a.SetLastUsage(event.Usage)
}
@@ -433,39 +651,224 @@ func (a *Agent) toolDefinitions() []llm.Tool {
return a.toolRegistry.Definitions()
}
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) string {
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) (string, bool) {
name := call.Function.Name
if logs != nil {
fmt.Fprintf(logs, "\n[tool] %s %s\n", name, compact(call.Function.Arguments, 400))
}
if a.toolRegistry == nil {
return tools.FormatError(fmt.Errorf("tool registry is disabled"))
return tools.FormatError(fmt.Errorf("tool registry is disabled")), false
}
tool, ok := a.toolRegistry.Get(name)
if !ok {
return tools.FormatError(fmt.Errorf("unknown tool: %s", name))
return tools.FormatError(fmt.Errorf("unknown tool: %s", name)), false
}
args := strings.TrimSpace(call.Function.Arguments)
if args == "" {
args = "{}"
}
raw := json.RawMessage(args)
rollbackNotes, err := a.prepareMutation(ctx, tool, call, raw, args)
if err != nil {
return tools.FormatError(err), false
}
toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
defer cancel()
output, err := tool.Execute(toolCtx, json.RawMessage(args))
if streaming, ok := tool.(tools.StreamingTool); ok {
output, streamed := a.executeStreamingTool(streaming, toolCtx, raw, name, args, logs)
return appendRollbackNotes(output, rollbackNotes), streamed
}
output, err := tool.Execute(toolCtx, raw)
if err != nil {
if output != "" {
return output + "\n" + tools.FormatError(err)
return appendRollbackNotes(output+"\n"+tools.FormatError(err), rollbackNotes), false
}
return tools.FormatError(err)
return appendRollbackNotes(tools.FormatError(err), rollbackNotes), false
}
if logs != nil {
fmt.Fprintf(logs, "[tool] %s done\n", name)
return appendRollbackNotes(output, rollbackNotes), false
}
func (a *Agent) prepareMutation(ctx context.Context, tool tools.Tool, call llm.ToolCall, raw json.RawMessage, args string) ([]string, error) {
if !tools.IsMutatingCall(tool, raw) {
return nil, nil
}
var targets []string
if targetTool, ok := tool.(tools.TargetPathTool); ok {
var err error
targets, err = targetTool.TargetPaths(raw)
if err != nil {
return nil, fmt.Errorf("resolve mutation targets: %w", err)
}
}
preview := compact(args, toolLogArgumentLimit)
if previewTool, ok := tool.(tools.PreviewableTool); ok {
value, err := previewTool.Preview(ctx, raw)
if err != nil {
return nil, fmt.Errorf("preview mutation: %w", err)
}
preview = value
}
if a.authorizer != nil {
decision, err := a.authorizer.AuthorizeToolCall(ctx, ToolAuthorizationRequest{
ID: call.ID,
ToolName: call.Function.Name,
Arguments: args,
Preview: preview,
RollbackAvailable: a.rollbackStore != nil && len(targets) > 0,
})
if err != nil {
return nil, fmt.Errorf("authorize %s: %w", call.Function.Name, err)
}
if decision != ToolAuthorizationAllow {
return nil, fmt.Errorf("tool call denied by user: %s", call.Function.Name)
}
}
if a.rollbackStore == nil {
return nil, nil
}
notes := make([]string, 0, len(targets))
for _, path := range targets {
id, err := a.rollbackStore.Capture(call.Function.Name, path, path)
if err != nil {
return nil, fmt.Errorf("capture rollback snapshot for %s: %w", path, err)
}
notes = append(notes, fmt.Sprintf("[rollback] %s %s", id, path))
}
return notes, nil
}
func appendRollbackNotes(output string, notes []string) string {
if len(notes) == 0 {
return output
}
if strings.TrimSpace(output) == "" {
return strings.Join(notes, "\n")
}
return output + "\n" + strings.Join(notes, "\n")
}
func (a *Agent) executeStreamingTool(streaming tools.StreamingTool, ctx context.Context, raw json.RawMessage, name, args string, logs io.Writer) (string, bool) {
var streamed int
emit := func(chunk string) error {
if chunk == "" || logs == nil {
return nil
}
if streamed >= toolStreamLogLimit {
return nil
}
if len(chunk) > toolStreamLogLimit-streamed {
chunk = chunk[:toolStreamLogLimit-streamed]
streamed = toolStreamLogLimit
_, err := fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s\n[streaming output truncated after %d bytes]", name, compact(args, toolLogArgumentLimit), toolLogOutputMarker, chunk, toolStreamLogLimit)
return err
}
streamed += len(chunk)
_, err := fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s", name, compact(args, toolLogArgumentLimit), toolLogOutputMarker, chunk)
return err
}
output, err := streaming.ExecuteStream(ctx, raw, emit)
if logs != nil {
if err != nil {
_ = emit("\n" + tools.FormatError(err))
} else if streamed == 0 {
_ = emit("(no output)")
}
}
if err != nil {
if output != "" {
return output + "\n" + tools.FormatError(err), true
}
return tools.FormatError(err), true
}
return output, true
}
func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
if a.toolRegistry == nil {
return false
}
for _, call := range calls {
if tool, ok := a.toolRegistry.Get(call.Function.Name); ok {
args := strings.TrimSpace(call.Function.Arguments)
if args == "" {
args = "{}"
}
if tools.IsMutatingCall(tool, json.RawMessage(args)) {
return true
}
}
}
return false
}
func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.Writer) []toolExecutionResult {
results := make([]toolExecutionResult, len(calls))
// Log all tool starts sequentially so TUI gets ordered placeholders.
if logs != nil {
for _, call := range calls {
logToolStart(logs, call)
}
}
if a.hasMutatingCall(calls) {
// Sequential: avoid racing mutating tools.
for i, call := range calls {
output, streamed := a.executeTool(ctx, call, logs)
results[i] = toolExecutionResult{call: call, output: output}
if logs != nil && !streamed {
logToolOutput(logs, call, output)
}
}
return results
}
// Concurrent: independent read-only tools.
var mu sync.Mutex
var wg sync.WaitGroup
for i, call := range calls {
wg.Add(1)
go func(idx int, c llm.ToolCall) {
defer wg.Done()
output, streamed := a.executeTool(ctx, c, logs)
results[idx] = toolExecutionResult{call: c, output: output}
if logs != nil && !streamed {
mu.Lock()
logToolOutput(logs, c, output)
mu.Unlock()
}
}(i, call)
}
wg.Wait()
return results
}
func logToolStart(logs io.Writer, call llm.ToolCall) {
fmt.Fprintf(logs, "\n[tool] %s %s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit))
}
func logToolOutput(logs io.Writer, call llm.ToolCall, output string) {
fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit), toolLogOutputMarker, toolOutputPreview(output))
}
func toolOutputPreview(output string) string {
output = strings.TrimSpace(output)
if output == "" {
return "(no output)"
}
if len(output) <= toolLogOutputPreviewBytes {
return output
}
return output[:toolLogOutputPreviewBytes] + fmt.Sprintf("\n[truncated: showing first %d bytes; omitted %d bytes]", toolLogOutputPreviewBytes, len(output)-toolLogOutputPreviewBytes)
}
type toolCallBuilder struct {
index int
+583 -1
View File
@@ -3,12 +3,17 @@ package agent
import (
"context"
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"reflect"
"strings"
"sync"
"testing"
"agentu/pkg/llm"
"agentu/internal/tools"
"agentu/pkg/llm"
)
type fakeProvider struct {
@@ -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 {
called bool
text string
@@ -97,6 +182,173 @@ func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, e
return emit(llm.StreamEvent{Content: p.text})
}
type btwProvider struct {
req llm.ChatRequest
err error
}
func (p *btwProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.req = req
if p.err != nil {
return p.err
}
return emit(llm.StreamEvent{Content: "side answer", Usage: &llm.Usage{TotalTokens: 999}})
}
func TestBTWUsesContextWithoutMutatingConversation(t *testing.T) {
provider := &btwProvider{}
a := New(Options{Provider: provider, Model: "test", Thinking: "high", ThinkingEnabled: true})
a.SetMessages([]llm.Message{
{Role: llm.RoleUser, Content: "main question"},
{Role: llm.RoleAssistant, Content: "main answer"},
})
a.SetLastUsage(&llm.Usage{TotalTokens: 7})
before := a.Messages()
var out strings.Builder
if err := a.BTW(context.Background(), "what does that mean?", &out); err != nil {
t.Fatal(err)
}
if out.String() != "side answer" {
t.Fatalf("output = %q", out.String())
}
if !reflect.DeepEqual(a.Messages(), before) {
t.Fatalf("conversation changed:\nbefore: %+v\nafter: %+v", before, a.Messages())
}
if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 7 {
t.Fatalf("usage changed: %+v", usage)
}
if len(provider.req.Tools) != 0 || provider.req.ToolChoice != "" {
t.Fatalf("btw exposed tools: %+v / %q", provider.req.Tools, provider.req.ToolChoice)
}
last := provider.req.Messages[len(provider.req.Messages)-1]
if last.Role != llm.RoleUser || last.Content != "what does that mean?" {
t.Fatalf("last request message = %+v", last)
}
if provider.req.Extra["thinking"] != "high" {
t.Fatalf("thinking extra = %+v", provider.req.Extra)
}
}
func TestBTWRejectsEmptyQuestionAndReturnsProviderError(t *testing.T) {
provider := &btwProvider{err: errors.New("boom")}
a := New(Options{Provider: provider, Model: "test"})
if err := a.BTW(context.Background(), " ", nil); err == nil {
t.Fatal("empty question did not fail")
}
if err := a.BTW(context.Background(), "question", nil); err == nil || !strings.Contains(err.Error(), "boom") {
t.Fatalf("provider error = %v", err)
}
}
func TestDotContinuesCurrentTaskWithoutRewritingHistory(t *testing.T) {
provider := &btwProvider{}
a := New(Options{Provider: provider, Model: "test"})
if err := a.RunTurn(context.Background(), " . ", nil, nil); err != nil {
t.Fatal(err)
}
requestMessage := provider.req.Messages[len(provider.req.Messages)-1]
if requestMessage.Content != continuationPrompt {
t.Fatalf("request content = %q, want %q", requestMessage.Content, continuationPrompt)
}
history := a.Messages()
if history[0].Content != "." {
t.Fatalf("stored input = %q, want literal dot", history[0].Content)
}
}
func TestExpandContinuationMessagesOnlyExpandsSingleDot(t *testing.T) {
messages := []llm.Message{
{Role: llm.RoleUser, Content: "."},
{Role: llm.RoleUser, Content: " . "},
{Role: llm.RoleUser, Content: ".."},
{Role: llm.RoleUser, Content: ". continue"},
{Role: llm.RoleAssistant, Content: "."},
}
expanded := expandContinuationMessages(messages)
if expanded[0].Content != continuationPrompt || expanded[1].Content != continuationPrompt {
t.Fatalf("single dots were not expanded: %+v", expanded)
}
for _, index := range []int{2, 3, 4} {
if expanded[index].Content != messages[index].Content {
t.Fatalf("message %d changed from %q to %q", index, messages[index].Content, expanded[index].Content)
}
}
if messages[0].Content != "." || messages[1].Content != " . " {
t.Fatalf("source messages were mutated: %+v", messages)
}
}
func TestCompleteConversationPrefixDropsPendingToolRound(t *testing.T) {
messages := []llm.Message{
{Role: llm.RoleUser, Content: "first"},
{Role: llm.RoleAssistant, Content: "answer"},
{Role: llm.RoleUser, Content: "main"},
{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "call_1"}, {ID: "call_2"}}},
{Role: llm.RoleTool, ToolCallID: "call_1", Content: "one"},
}
got := completeConversationPrefix(messages)
if len(got) != 3 || got[len(got)-1].Content != "main" {
t.Fatalf("complete prefix = %+v", got)
}
}
type concurrentBTWProvider struct {
mainStarted chan struct{}
releaseMain chan struct{}
once sync.Once
mu sync.Mutex
btwRequest llm.ChatRequest
}
func (p *concurrentBTWProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
last := req.Messages[len(req.Messages)-1].Content
if last == "main task" {
p.once.Do(func() { close(p.mainStarted) })
select {
case <-p.releaseMain:
return emit(llm.StreamEvent{Content: "main answer"})
case <-ctx.Done():
return ctx.Err()
}
}
p.mu.Lock()
p.btwRequest = req
p.mu.Unlock()
return emit(llm.StreamEvent{Content: "side answer"})
}
func TestBTWCanRunDuringMainTurn(t *testing.T) {
provider := &concurrentBTWProvider{mainStarted: make(chan struct{}), releaseMain: make(chan struct{})}
a := New(Options{Provider: provider, Model: "test"})
mainDone := make(chan error, 1)
go func() { mainDone <- a.RunTurn(context.Background(), "main task", nil, nil) }()
<-provider.mainStarted
var side strings.Builder
if err := a.BTW(context.Background(), "side question", &side); err != nil {
t.Fatal(err)
}
if side.String() != "side answer" {
t.Fatalf("side output = %q", side.String())
}
provider.mu.Lock()
req := provider.btwRequest
provider.mu.Unlock()
if len(req.Messages) < 2 || req.Messages[len(req.Messages)-2].Content != "main task" {
t.Fatalf("btw request did not include active main turn: %+v", req.Messages)
}
close(provider.releaseMain)
if err := <-mainDone; err != nil {
t.Fatal(err)
}
if len(a.Messages()) != 2 {
t.Fatalf("main conversation contains btw messages: %+v", a.Messages())
}
}
type compactProvider struct {
summaryCalls int
answerCalls int
@@ -145,6 +397,9 @@ func TestAgentCompact(t *testing.T) {
if messages[len(messages)-1].Content != "recent-answer-2" {
t.Fatalf("last recent message = %#v", messages[len(messages)-1])
}
if messages[2].Content != "recent-user-0" {
t.Fatalf("first retained recent message = %#v", messages[2])
}
if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 12 {
t.Fatalf("usage = %#v", usage)
}
@@ -217,6 +472,48 @@ func longConversation() []llm.Message {
return messages
}
func manyTurnConversation(turns int) []llm.Message {
messages := []llm.Message{{Role: llm.RoleSystem, Content: "system"}}
for i := 0; i < turns; i++ {
messages = append(messages,
llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("user-%02d", i)},
llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("assistant-%02d", i)},
)
}
return messages
}
func TestAgentCompactKeepsTwentyPercentRecentContext(t *testing.T) {
provider := &compactProvider{}
a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
a.SetMessages(manyTurnConversation(25))
if err := a.Compact(context.Background()); err != nil {
t.Fatal(err)
}
messages := a.Messages()
// 1 system + 1 summary + 10 retained recent messages (5 user/assistant turns)
if len(messages) != 12 {
t.Fatalf("messages = %d, want 12; contents: %v", len(messages), joinMessageContents(messages))
}
if provider.summaryCalls != 1 {
t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
}
if messages[2].Content != "user-20" {
t.Fatalf("first retained message = %#v, want user-20", messages[2])
}
if messages[len(messages)-1].Content != "assistant-24" {
t.Fatalf("last retained message = %#v, want assistant-24", messages[len(messages)-1])
}
if !strings.Contains(provider.lastSummary, "user-19") {
t.Fatalf("summary should contain user-19 (part of summarized 80%%): %q", provider.lastSummary)
}
if strings.Contains(provider.lastSummary, "user-20") {
t.Fatalf("summary should not contain user-20 (part of retained 20%%): %q", provider.lastSummary)
}
}
func joinMessageContents(messages []llm.Message) string {
var b strings.Builder
for _, msg := range messages {
@@ -340,3 +637,288 @@ func TestAgentHardLimitStopsVaryingCalls(t *testing.T) {
t.Fatalf("calls = %d, want %d", provider.calls, limit)
}
}
// slowTool blocks on a channel before returning.
type slowTool struct {
name string
started chan<- string
release <-chan struct{}
}
func (t *slowTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: t.name,
Description: "slow test tool",
Parameters: json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}}}`),
},
}
}
func (t *slowTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args struct {
Text string `json:"text"`
}
if err := json.Unmarshal(raw, &args); err != nil {
return "", err
}
t.started <- t.name
<-t.release
return args.Text, nil
}
// mutatingSlowTool is like slowTool but implements MutatingTool.
type mutatingSlowTool struct {
slowTool
}
func (*mutatingSlowTool) Mutates() bool { return true }
func TestAgentExecutesIndependentToolCallsInParallel(t *testing.T) {
started := make(chan string, 2)
release := make(chan struct{})
tool1 := &slowTool{name: "slow_one", started: started, release: release}
tool2 := &slowTool{name: "slow_two", started: started, release: release}
provider := &parallelProvider{
responses: [][]llm.ToolCallDelta{
{
{Index: 0, ID: "call_1", Type: "function", Name: "slow_one", Arguments: `{"text":"one"}`},
{Index: 1, ID: "call_2", Type: "function", Name: "slow_two", Arguments: `{"text":"two"}`},
},
},
finalContent: "done",
}
a := New(Options{
Provider: provider,
Model: "test",
SystemPrompt: "system",
ToolRegistry: tools.NewRegistry(tool1, tool2),
})
var out strings.Builder
done := make(chan error, 1)
go func() {
done <- a.RunTurn(context.Background(), "run both", &out, nil)
}()
// Both tools must start before we release them.
names := make(map[string]bool)
for i := 0; i < 2; i++ {
select {
case name := <-started:
names[name] = true
case <-make(chan struct{}):
// Use a timer in real code, but for test simplicity we rely on deadlock detection.
t.Fatal("timed out waiting for tool to start")
}
}
if !names["slow_one"] || !names["slow_two"] {
t.Fatalf("expected both tools to start, got %v", names)
}
close(release)
if err := <-done; err != nil {
t.Fatal(err)
}
if out.String() != "done" {
t.Fatalf("out = %q", out.String())
}
// Verify tool results were in order.
msgs := a.Messages()
var toolMsgs []string
for _, m := range msgs {
if m.Role == llm.RoleTool {
toolMsgs = append(toolMsgs, m.Content)
}
}
if len(toolMsgs) != 2 || toolMsgs[0] != "one" || toolMsgs[1] != "two" {
t.Fatalf("tool messages = %v, want [one two]", toolMsgs)
}
}
func TestAgentRunsMutatingToolCallsSequentially(t *testing.T) {
started := make(chan string, 2)
release := make(chan struct{})
tool1 := &mutatingSlowTool{slowTool{name: "mut_one", started: started, release: release}}
tool2 := &mutatingSlowTool{slowTool{name: "mut_two", started: started, release: release}}
provider := &parallelProvider{
responses: [][]llm.ToolCallDelta{
{
{Index: 0, ID: "call_1", Type: "function", Name: "mut_one", Arguments: `{"text":"one"}`},
{Index: 1, ID: "call_2", Type: "function", Name: "mut_two", Arguments: `{"text":"two"}`},
},
},
finalContent: "done",
}
a := New(Options{
Provider: provider,
Model: "test",
SystemPrompt: "system",
ToolRegistry: tools.NewRegistry(tool1, tool2),
})
var out strings.Builder
done := make(chan error, 1)
go func() {
done <- a.RunTurn(context.Background(), "run both", &out, nil)
}()
// Only the first tool should start.
select {
case name := <-started:
if name != "mut_one" {
t.Fatalf("expected mut_one to start first, got %q", name)
}
case <-make(chan struct{}):
t.Fatal("timed out waiting for first tool to start")
}
// Second tool should NOT have started yet.
select {
case name := <-started:
t.Fatalf("second tool started too early: %q", name)
default:
// Expected: second tool hasn't started.
}
close(release)
if err := <-done; err != nil {
t.Fatal(err)
}
}
// parallelProvider emits a batch of tool calls on the first request, then returns finalContent.
type parallelProvider struct {
responses [][]llm.ToolCallDelta
finalContent string
calls int
}
type authorizationProvider struct {
call llm.ToolCallDelta
calls int
toolOutput string
}
func (p *authorizationProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.calls++
if p.calls == 1 {
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{p.call}})
}
p.toolOutput = req.Messages[len(req.Messages)-1].Content
return emit(llm.StreamEvent{Content: "done"})
}
type recordingAuthorizer struct {
decision ToolAuthorizationDecision
requests []ToolAuthorizationRequest
}
func (a *recordingAuthorizer) AuthorizeToolCall(ctx context.Context, req ToolAuthorizationRequest) (ToolAuthorizationDecision, error) {
a.requests = append(a.requests, req)
return a.decision, ctx.Err()
}
func TestAgentAuthorizesMutationAndCreatesRollbackSnapshot(t *testing.T) {
dir := t.TempDir()
provider := &authorizationProvider{call: llm.ToolCallDelta{
Index: 0, ID: "write_1", Type: "function", Name: "file_write",
Arguments: `{"path":"new.txt","content":"hello"}`,
}}
authorizer := &recordingAuthorizer{decision: ToolAuthorizationAllow}
rollback := tools.NewRollbackStore()
a := New(Options{
Provider: provider,
Model: "test",
ToolRegistry: tools.Builtins(dir),
Authorizer: authorizer,
RollbackStore: rollback,
})
if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
t.Fatal(err)
}
if len(authorizer.requests) != 1 || !authorizer.requests[0].RollbackAvailable {
t.Fatalf("authorization requests = %+v", authorizer.requests)
}
if !strings.Contains(authorizer.requests[0].Preview, "write 5 bytes to new.txt") {
t.Fatalf("preview = %q", authorizer.requests[0].Preview)
}
rollbackID := rollback.LastID()
if rollbackID == "" || !strings.Contains(provider.toolOutput, rollbackID) {
t.Fatalf("tool output %q missing rollback ID %q", provider.toolOutput, rollbackID)
}
if _, err := rollback.Restore(rollbackID); err != nil {
t.Fatal(err)
}
if _, err := os.Stat(filepath.Join(dir, "new.txt")); !os.IsNotExist(err) {
t.Fatalf("new file still exists after rollback: %v", err)
}
}
func TestAgentDeniedMutationDoesNotExecute(t *testing.T) {
dir := t.TempDir()
provider := &authorizationProvider{call: llm.ToolCallDelta{
Index: 0, ID: "write_1", Type: "function", Name: "file_write",
Arguments: `{"path":"denied.txt","content":"no"}`,
}}
authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})
if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
t.Fatal(err)
}
if _, err := os.Stat(filepath.Join(dir, "denied.txt")); !os.IsNotExist(err) {
t.Fatalf("denied file exists: %v", err)
}
if !strings.Contains(provider.toolOutput, "denied by user") {
t.Fatalf("tool output = %q", provider.toolOutput)
}
}
func TestAgentDoesNotAuthorizeApplyPatchDryRun(t *testing.T) {
dir := t.TempDir()
provider := &authorizationProvider{call: llm.ToolCallDelta{
Index: 0, ID: "patch_1", Type: "function", Name: "apply_patch",
Arguments: `{"patch":"*** Begin Patch\n*** Add File: dry.txt\n+dry\n*** End Patch","dry_run":true}`,
}}
authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})
if err := a.RunTurn(context.Background(), "preview", nil, nil); err != nil {
t.Fatal(err)
}
if len(authorizer.requests) != 0 {
t.Fatalf("dry run requested authorization: %+v", authorizer.requests)
}
if _, err := os.Stat(filepath.Join(dir, "dry.txt")); !os.IsNotExist(err) {
t.Fatalf("dry run created a file: %v", err)
}
}
func TestSetMessagesRefreshesCurrentSystemInstructions(t *testing.T) {
a := New(Options{SystemPrompt: "current instructions"})
a.SetMessages([]llm.Message{
{Role: llm.RoleSystem, Content: "old instructions"},
{Role: llm.RoleUser, Content: "hello"},
})
messages := a.Messages()
if messages[0].Content != "current instructions" || messages[1].Content != "hello" {
t.Fatalf("messages = %+v", messages)
}
}
func (p *parallelProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.calls++
if p.calls <= len(p.responses) {
return emit(llm.StreamEvent{ToolCalls: p.responses[p.calls-1]})
}
return emit(llm.StreamEvent{Content: p.finalContent})
}
+26
View File
@@ -0,0 +1,26 @@
package agent
import "context"
// ToolAuthorizationDecision is the result of an authorization check.
type ToolAuthorizationDecision string
const (
ToolAuthorizationAllow ToolAuthorizationDecision = "allow"
ToolAuthorizationDeny ToolAuthorizationDecision = "deny"
)
// ToolAuthorizationRequest is sent to an authorizer before executing a
// mutating tool call.
type ToolAuthorizationRequest struct {
ID string
ToolName string
Arguments string
Preview string
RollbackAvailable bool
}
// ToolAuthorizer decides whether a mutating tool call should proceed.
type ToolAuthorizer interface {
AuthorizeToolCall(ctx context.Context, req ToolAuthorizationRequest) (ToolAuthorizationDecision, error)
}
+192
View File
@@ -0,0 +1,192 @@
package app
import (
"context"
"errors"
"flag"
"fmt"
"net/http"
"os"
"os/signal"
"path/filepath"
"strings"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/internal/skills"
"agentu/internal/tools"
"agentu/internal/tui"
"agentu/pkg/config"
"agentu/pkg/llm"
)
type cliOptions struct {
configPath string
yolo bool
plain bool
themeName string
resumeID string
}
func Run() error {
opts := parseCLIOptions()
themeMode, err := tui.ParseThemeMode(opts.themeName)
if err != nil {
return err
}
cfg, err := loadConfig(opts.configPath)
if err != nil {
return err
}
globalSkills, err := skills.LoadGlobal()
if err != nil {
return fmt.Errorf("load global skills: %w", err)
}
rollbackStore := tools.NewRollbackStore()
var approvalBroker *tui.ApprovalBroker
var authorizer agent.ToolAuthorizer
if !opts.yolo && !opts.plain {
approvalBroker = tui.NewApprovalBroker()
authorizer = approvalBroker
}
providerName, providerConfig, provider := defaultProvider(cfg, http.DefaultClient)
assistant := newAssistant(cfg, providerName, providerConfig, provider, opts, globalSkills, authorizer, rollbackStore, http.DefaultClient)
modelManager, err := newSessionManager(cfg, assistant, http.DefaultClient)
if err != nil {
return err
}
if err := initializeSession(modelManager, opts.resumeID); err != nil {
return err
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
defer stop()
return runInterface(ctx, assistant, modelManager, providerConfig, cfg.Agent.WorkingDir, opts, themeMode, approvalBroker, rollbackStore)
}
func parseCLIOptions() cliOptions {
configPath := flag.String("config", config.DefaultPath, "path to YAML config")
yolo := flag.Bool("yolo", false, "also enable automatic file writes and shell execution")
plain := flag.Bool("plain", false, "use the simple line-based REPL instead of the TUI")
themeName := flag.String("theme", "light", "TUI theme: light or dark")
resumeID := flag.String("resume", "", "resume a specific session by ID")
flag.Parse()
return cliOptions{
configPath: *configPath,
yolo: *yolo,
plain: *plain,
themeName: *themeName,
resumeID: *resumeID,
}
}
func loadConfig(configPath string) (*config.Config, error) {
cfg, err := config.Load(configPath)
if err == nil {
return cfg, nil
}
if !errors.Is(err, os.ErrNotExist) {
return nil, err
}
if bootstrapErr := bootstrapConfig(configPath); bootstrapErr != nil {
return nil, bootstrapErr
}
return nil, fmt.Errorf("created default config at %s; edit it to set your API key and provider, then restart agentu", configPath)
}
func defaultProvider(cfg *config.Config, httpClient *http.Client) (string, config.ProviderConfig, llm.Provider) {
providerName := cfg.DefaultProviderName()
providerConfig := cfg.Providers[providerName]
provider := newProviderClient(providerConfig, httpClient)
return providerName, providerConfig, provider
}
func newProviderClient(providerConfig config.ProviderConfig, httpClient *http.Client) llm.Provider {
if providerConfig.APIType == config.APITypeResponses {
return llm.NewOpenAIResponsesClient(providerConfig.BaseURL, providerConfig.APIKey, httpClient)
}
return llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, httpClient)
}
func newAssistant(cfg *config.Config, providerName string, providerConfig config.ProviderConfig, provider llm.Provider, opts cliOptions, globalSkills []skills.Skill, authorizer agent.ToolAuthorizer, rollbackStore *tools.RollbackStore, httpClient *http.Client) *agent.Agent {
writesEnabled := opts.yolo || !opts.plain
return agent.New(agent.Options{
Provider: provider,
ProviderName: providerName,
Model: providerConfig.Model,
Thinking: providerConfig.Thinking,
ThinkingParam: providerConfig.ThinkingParam,
ThinkingEnabled: providerConfig.ThinkingConfigured,
SystemPrompt: systemPrompt(cfg, writesEnabled, opts.yolo, globalSkills),
ToolRegistry: newToolRegistry(cfg.Agent.WorkingDir, writesEnabled, httpClient),
Authorizer: authorizer,
RollbackStore: rollbackStore,
ToolTimeout: cfg.Agent.ToolTimeout,
MaxContextTokens: maxContextTokens(cfg, providerConfig),
})
}
func systemPrompt(cfg *config.Config, writesEnabled bool, yolo bool, globalSkills []skills.Skill) string {
parts := []string{cfg.Agent.SystemPrompt, tools.AgentInstructionsForMode(cfg.Agent.WorkingDir, writesEnabled, yolo)}
if formatted := skills.FormatForSystemPrompt(globalSkills); formatted != "" {
parts = append(parts, formatted)
}
return strings.TrimSpace(strings.Join(parts, "\n\n"))
}
func maxContextTokens(cfg *config.Config, providerConfig config.ProviderConfig) int {
if cfg.Agent.MaxContextTokens != 0 {
return cfg.Agent.MaxContextTokens
}
return providerConfig.ContextTokens
}
func newToolRegistry(workingDir string, writesEnabled bool, httpClient *http.Client) *tools.Registry {
var registry *tools.Registry
if writesEnabled {
registry = tools.Builtins(workingDir)
} else {
registry = tools.ReadOnlyBuiltins(workingDir)
}
registry.Register(tools.NewBuiltinSearchTool(httpClient))
registry.Register(tools.NewBuiltinFetchTool(httpClient))
return registry
}
func newSessionManager(cfg *config.Config, assistant *agent.Agent, httpClient *http.Client) (*session.Manager, error) {
home, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("resolve home dir: %w", err)
}
sessionsDir := filepath.Join(home, ".agentu", "sessions")
store, err := session.NewStore(sessionsDir)
if err != nil {
return nil, fmt.Errorf("init session store: %w", err)
}
return session.NewManager(cfg, assistant, httpClient, store)
}
func runInterface(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager, providerConfig config.ProviderConfig, workingDir string, opts cliOptions, themeMode tui.ThemeMode, approvalBroker *tui.ApprovalBroker, rollbackStore *tools.RollbackStore) error {
if opts.plain {
return repl(ctx, assistant, modelManager)
}
err := tui.Run(ctx, assistant, tui.Options{
ModelName: providerConfig.Model,
Yolo: opts.yolo,
ThemeMode: themeMode,
ModelManager: modelManager,
WorkingDir: workingDir,
Approval: approvalBroker,
Rollback: rollbackStore,
})
printExitSession(modelManager)
return err
}
+38
View File
@@ -0,0 +1,38 @@
package app
import (
"fmt"
"os"
"path/filepath"
"agentu/pkg/config"
)
const defaultConfigContent = `# agentu configuration
# Replace the placeholder values below with your actual provider settings.
# See agentu.example.yaml for the full reference.
providers:
default:
base_url: https://api.openai.com
api_key: YOUR_API_KEY_HERE
models:
- id: gpt-4o
context: 128000
- id: gpt-4o-mini
context: 128000
`
func bootstrapConfig(configPath string) error {
resolved, err := config.ResolvePath(configPath)
if err != nil {
return fmt.Errorf("resolve config path: %w", err)
}
dir := filepath.Dir(resolved)
if err := os.MkdirAll(dir, 0o755); err != nil {
return fmt.Errorf("create config directory: %w", err)
}
if err := os.WriteFile(resolved, []byte(defaultConfigContent), 0o644); err != nil {
return fmt.Errorf("write default config: %w", err)
}
return nil
}
+113
View File
@@ -0,0 +1,113 @@
package app
import (
"context"
"path/filepath"
"strings"
"testing"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/internal/skills"
"agentu/pkg/config"
"agentu/pkg/llm"
)
type startupProvider struct{}
func (p *startupProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
return emit(llm.StreamEvent{Content: "ok"})
}
func startupTestManager(t *testing.T) (*session.Manager, *agent.Agent) {
t.Helper()
cfg := &config.Config{
Providers: map[string]config.ProviderConfig{
"one": {
Name: "one",
BaseURL: "https://one.example.com",
APIKey: "sk-test",
Model: "model-a",
Models: []string{"model-a"},
Thinking: "none",
ThinkingParam: "thinking",
},
},
}
assistant := agent.New(agent.Options{
Provider: &startupProvider{},
ProviderName: "one",
Model: "model-a",
})
store, err := session.NewStore(filepath.Join(t.TempDir(), "sessions"))
if err != nil {
t.Fatal(err)
}
manager, err := session.NewManager(cfg, assistant, nil, store)
if err != nil {
t.Fatal(err)
}
return manager, assistant
}
func TestInitializeSessionStartsNewByDefault(t *testing.T) {
manager, assistant := startupTestManager(t)
if _, err := manager.NewSession(); err != nil {
t.Fatalf("new session: %v", err)
}
savedID := manager.CurrentSessionID()
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: "old context"}})
if err := manager.Save(); err != nil {
t.Fatalf("save: %v", err)
}
if err := initializeSession(manager, ""); err != nil {
t.Fatalf("initialize: %v", err)
}
if manager.CurrentSessionID() == savedID {
t.Fatalf("default startup resumed latest session %s", savedID)
}
for _, msg := range assistant.Messages() {
if msg.Content == "old context" {
t.Fatal("default startup carried old context into new session")
}
}
}
func TestInitializeSessionResumesExplicitID(t *testing.T) {
manager, assistant := startupTestManager(t)
if _, err := manager.NewSession(); err != nil {
t.Fatalf("new session: %v", err)
}
savedID := manager.CurrentSessionID()
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: "old context"}})
if err := manager.Save(); err != nil {
t.Fatalf("save: %v", err)
}
if err := initializeSession(manager, savedID); err != nil {
t.Fatalf("initialize: %v", err)
}
if manager.CurrentSessionID() != savedID {
t.Fatalf("expected resumed session %s, got %s", savedID, manager.CurrentSessionID())
}
found := false
for _, msg := range assistant.Messages() {
if msg.Content == "old context" {
found = true
}
}
if !found {
t.Fatal("explicit resume did not restore saved context")
}
}
func TestSystemPromptIncludesGlobalSkillsAndApprovalMode(t *testing.T) {
cfg := &config.Config{Agent: config.AgentConfig{SystemPrompt: "base", WorkingDir: "/tmp/project"}}
prompt := systemPrompt(cfg, true, false, []skills.Skill{{Name: "review", Path: "/tmp/review/SKILL.md", Content: "Review carefully."}})
for _, want := range []string{"base", "interactive user approval", "### review", "Review carefully."} {
if !strings.Contains(prompt, want) {
t.Fatalf("prompt missing %q:\n%s", want, prompt)
}
}
}
+95
View File
@@ -0,0 +1,95 @@
package app
import (
"bufio"
"context"
"fmt"
"os"
"strings"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/internal/skills"
)
func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error {
scanner := bufio.NewScanner(os.Stdin)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
fmt.Println("agentu")
fmt.Println("Type /exit to quit, /clear to reset context, /compact to compress context, /btw <question> for a side question.")
for {
fmt.Print("> ")
if !scanner.Scan() {
if err := scanner.Err(); err != nil {
return err
}
fmt.Println()
printExitSession(modelManager)
return nil
}
input := strings.TrimSpace(scanner.Text())
fields := strings.Fields(input)
switch input {
case "":
continue
case "/exit", "/quit":
printExitSession(modelManager)
return nil
case "/clear":
assistant.Clear()
fmt.Println("context cleared")
continue
case "/compact":
if err := assistant.Compact(ctx); err != nil {
fmt.Fprintln(os.Stderr, "compact error:", err)
} else {
fmt.Println("context compacted")
if modelManager != nil {
_ = modelManager.Save()
}
}
continue
}
if len(fields) > 0 && fields[0] == "/btw" {
question := strings.TrimSpace(strings.TrimPrefix(input, "/btw"))
if question == "" {
fmt.Fprintln(os.Stderr, "usage: /btw <question>")
continue
}
fmt.Print("[btw] ")
if err := assistant.BTW(ctx, question, os.Stdout); err != nil {
fmt.Fprintln(os.Stderr, "\nbtw error:", err)
} else {
fmt.Println()
}
continue
}
if input == "/skills" || input == "/skills list" {
items, err := skills.LoadGlobal()
if err != nil {
fmt.Fprintln(os.Stderr, "skills error:", err)
} else {
fmt.Println(skills.FormatList(items))
}
continue
}
if len(fields) >= 3 && fields[0] == "/skills" && fields[1] == "install" {
result, err := skills.InstallGlobal(ctx, strings.Join(fields[2:], " "))
if err != nil {
fmt.Fprintln(os.Stderr, "skills error:", err)
} else {
assistant.AddSystemInstructions(skills.FormatForSystemPrompt([]skills.Skill{result.Skill}))
fmt.Printf("installed skill %q; active in this conversation\n", result.Skill.Name)
}
continue
}
if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
fmt.Fprintln(os.Stderr, "\nerror:", err)
}
_ = modelManager.Save()
fmt.Println()
}
}
+38
View File
@@ -0,0 +1,38 @@
package app
import (
"fmt"
"os"
"agentu/internal/session"
)
func initializeSession(modelManager *session.Manager, resumeID string) error {
if resumeID != "" {
if err := modelManager.Resume(resumeID); err != nil {
return fmt.Errorf("resume session: %w", err)
}
return nil
}
if _, err := modelManager.NewSession(); err != nil {
return fmt.Errorf("start new session: %w", err)
}
return nil
}
func printExitSession(m *session.Manager) {
if m == nil {
return
}
_ = m.Save()
id := m.CurrentSessionID()
name := m.CurrentSessionName()
if id != "" {
if name != "" && name != id {
fmt.Fprintf(os.Stderr, "session: %s (%s)\n", id, name)
} else {
fmt.Fprintf(os.Stderr, "session: %s\n", id)
}
fmt.Fprintf(os.Stderr, "resume: agentu --resume %s\n", id)
}
}
+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
}
m.currentSession = sess
m.restoreRuntime(sess)
m.agent.SetMessages(sess.Messages)
m.agent.SetLastUsage(sess.LastUsage)
return nil
}
// restoreRuntime switches the manager and agent back to the provider, model,
// and thinking level recorded on a resumed session. Values that are no longer
// valid for the current config are skipped and the defaults remain active.
func (m *Manager) restoreRuntime(sess *Session) {
if sess.Provider != "" && sess.Provider != m.providerName {
if _, ok := m.cfg.Providers[sess.Provider]; ok {
_, _ = m.SetProvider(sess.Provider)
}
}
if sess.Model != "" && sess.Model != m.provider.Model {
_, _ = m.SetModel(sess.Model)
}
if sess.Thinking != "" && sess.Thinking != m.provider.Thinking {
_, _ = m.SetThinking(sess.Thinking)
}
}
// Save persists the current session to disk.
func (m *Manager) Save() error {
if m.store == nil || m.currentSession == nil {
@@ -82,6 +100,8 @@ func (m *Manager) Save() error {
m.currentSession.LastUsage = m.agent.LastUsage()
m.currentSession.Provider = m.providerName
m.currentSession.Model = m.provider.Model
m.currentSession.APIType = m.provider.APIType
m.currentSession.Thinking = m.provider.Thinking
// Auto-name if empty
if m.currentSession.Name == "" {
m.currentSession.Name = DefaultName(m.currentSession.Messages)
@@ -150,6 +170,8 @@ func (m *Manager) NewSession() (string, error) {
m.currentSession.Messages = msgs
m.currentSession.Provider = m.providerName
m.currentSession.Model = m.provider.Model
m.currentSession.APIType = m.provider.APIType
m.currentSession.Thinking = m.provider.Thinking
if m.currentSession.Name == "" {
m.currentSession.Name = DefaultName(msgs)
}
@@ -204,10 +226,18 @@ func (m *Manager) CurrentThinking() string {
return m.provider.Thinking
}
func (m *Manager) CurrentAPIType() string {
if m.provider.APIType == "" {
return config.APITypeChat
}
return m.provider.APIType
}
func (m *Manager) Info() string {
var b strings.Builder
fmt.Fprintf(&b, "provider: %s\n", m.CurrentProvider())
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
fmt.Fprintf(&b, "api: %s\n", m.CurrentAPIType())
fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking())
fmt.Fprintf(&b, "providers: %s\n", strings.Join(m.providerNames(), ", "))
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.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
}
@@ -249,7 +279,7 @@ func (m *Manager) SetThinking(level string) (string, error) {
m.provider.Thinking = level
m.provider.ThinkingConfigured = true
if m.provider.ThinkingParam == "" {
m.provider.ThinkingParam = "thinking"
m.provider.ThinkingParam = config.DefaultThinkingParam(m.provider.APIType)
}
m.cfg.Providers[m.providerName] = m.provider
m.syncAgent(nil)
@@ -257,7 +287,14 @@ func (m *Manager) SetThinking(level string) (string, error) {
}
func (m *Manager) summary() string {
return fmt.Sprintf("provider: %s\nmodel: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentThinking())
return fmt.Sprintf("provider: %s\nmodel: %s\napi: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentAPIType(), m.CurrentThinking())
}
func providerClient(provider config.ProviderConfig, httpClient *http.Client) llm.Provider {
if provider.APIType == config.APITypeResponses {
return llm.NewOpenAIResponsesClient(provider.BaseURL, provider.APIKey, httpClient)
}
return llm.NewOpenAICompatibleClient(provider.BaseURL, provider.APIKey, httpClient)
}
func (m *Manager) maxContextTokens() int {
@@ -288,9 +325,13 @@ func (m *Manager) providerNames() []string {
return names
}
func (m *Manager) AvailableModels() []string {
return append([]string(nil), m.provider.Models...)
}
func contains(values []string, value string) bool {
for _, item := range values {
if item == value {
for _, v := range values {
if v == value {
return true
}
}
+201
View File
@@ -2,6 +2,8 @@ package session
import (
"context"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"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) {
manager, assistant, _ := testManager(t)
@@ -331,3 +403,132 @@ func TestManagerUsesModelContextForAgentCompaction(t *testing.T) {
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"
"time"
"agentu/internal/safeio"
"agentu/pkg/llm"
)
@@ -22,6 +23,8 @@ type Session struct {
UpdatedAt time.Time `json:"updated_at"`
Provider string `json:"provider"`
Model string `json:"model"`
APIType string `json:"api_type,omitempty"`
Thinking string `json:"thinking,omitempty"`
Messages []llm.Message `json:"messages"`
LastUsage *llm.Usage `json:"last_usage,omitempty"`
}
@@ -52,8 +55,11 @@ func NewStore(dir string) (*Store, error) {
// Save writes a session to disk.
func (s *Store) Save(sess *Session) error {
if sess.ID == "" {
return fmt.Errorf("session ID is required")
if sess == nil {
return fmt.Errorf("session is required")
}
if err := validateSessionID(sess.ID); err != nil {
return err
}
sess.UpdatedAt = time.Now().UTC()
if sess.CreatedAt.IsZero() {
@@ -64,14 +70,17 @@ func (s *Store) Save(sess *Session) error {
return fmt.Errorf("marshal session: %w", err)
}
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 nil
return s.updateIndex(metaFromSession(sess))
}
// Load reads a session by ID.
func (s *Store) Load(id string) (*Session, error) {
if err := validateSessionID(id); err != nil {
return nil, err
}
path := s.filePath(id)
data, err := os.ReadFile(path)
if err != nil {
@@ -89,13 +98,27 @@ func (s *Store) Load(id string) (*Session, error) {
// List returns metadata for all sessions, sorted by updated_at descending.
func (s *Store) List() ([]SessionMeta, error) {
if metas, err := s.loadIndex(); err == nil {
return metas, nil
}
metas, err := s.scanSessions()
if err != nil {
return nil, err
}
if err := s.saveIndexFromSessions(metas); err != nil {
return nil, err
}
return metas, nil
}
func (s *Store) scanSessions() ([]SessionMeta, error) {
entries, err := os.ReadDir(s.dir)
if err != nil {
return nil, fmt.Errorf("read session dir: %w", err)
}
var metas []SessionMeta
for _, entry := range entries {
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
if entry.IsDir() || entry.Name() == sessionIndexFile || !strings.HasSuffix(entry.Name(), ".json") {
continue
}
path := filepath.Join(s.dir, entry.Name())
@@ -125,6 +148,9 @@ func (s *Store) List() ([]SessionMeta, error) {
// Delete removes a session file.
func (s *Store) Delete(id string) error {
if err := validateSessionID(id); err != nil {
return err
}
path := s.filePath(id)
if err := os.Remove(path); err != nil {
if os.IsNotExist(err) {
@@ -132,13 +158,64 @@ func (s *Store) Delete(id string) error {
}
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 {
return filepath.Join(s.dir, id+".json")
}
func validateSessionID(id string) error {
if id == "" {
return fmt.Errorf("session ID is required")
}
if id == "." || id == ".." {
return fmt.Errorf("invalid session ID %q", id)
}
for _, r := range id {
if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-' || r == '_' || r == '.' {
continue
}
return fmt.Errorf("invalid session ID %q", id)
}
return nil
}
func (s *Store) updateIndex(meta SessionMeta) error {
metas, err := s.loadIndex()
if err != nil {
metas, err = s.scanSessions()
if err != nil {
return err
}
}
found := false
for i := range metas {
if metas[i].ID == meta.ID {
metas[i] = meta
found = true
break
}
}
if !found {
metas = append(metas, meta)
}
return s.saveIndexFromSessions(metas)
}
// GenerateID creates a random 8-character hex session ID.
func GenerateID() string {
b := make([]byte, 4)
+40
View File
@@ -119,6 +119,9 @@ func TestDeleteNotFound(t *testing.T) {
func TestSaveRequiresID(t *testing.T) {
store := tempStore(t)
if err := store.Save(nil); err == nil {
t.Fatal("expected error for nil session")
}
err := store.Save(&Session{Name: "no id"})
if err == nil {
t.Fatal("expected error for empty ID")
@@ -172,3 +175,40 @@ func TestNewStoreCreatesDir(t *testing.T) {
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...)
}
+45 -16
View File
@@ -13,6 +13,7 @@ import (
"sort"
"strings"
"agentu/internal/safeio"
"agentu/pkg/llm"
)
@@ -81,6 +82,34 @@ type FileWriteTool struct {
func NewFileWriteTool(workingDir string) *FileWriteTool {
return &FileWriteTool{workingDir: workingDir}
}
func (*FileWriteTool) Mutates() bool { return true }
func (t *FileWriteTool) Preview(ctx context.Context, raw json.RawMessage) (string, error) {
var args fileWriteArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
action := "write"
if args.Append {
action = "append"
}
return fmt.Sprintf("%s %d bytes to %s", action, len(args.Content), args.Path), ctx.Err()
}
func (t *FileWriteTool) TargetPaths(raw json.RawMessage) ([]string, error) {
var args fileWriteArgs
if err := json.Unmarshal(raw, &args); err != nil {
return nil, fmt.Errorf("parse arguments: %w", err)
}
if strings.TrimSpace(args.Path) == "" {
return nil, errors.New("path is required")
}
path, err := resolvePathForWrite(t.workingDir, args.Path)
if err != nil {
return nil, err
}
return []string{path}, nil
}
func (t *FileWriteTool) Definition() llm.Tool {
return llm.Tool{
@@ -116,28 +145,28 @@ func (t *FileWriteTool) Execute(ctx context.Context, raw json.RawMessage) (strin
if strings.TrimSpace(args.Path) == "" {
return "", errors.New("path is required")
}
path, err := resolvePath(t.workingDir, args.Path)
path, err := resolvePathForWrite(t.workingDir, args.Path)
if err != nil {
return "", err
}
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
return "", err
}
flag := os.O_CREATE | os.O_WRONLY | os.O_TRUNC
if args.Append {
flag = os.O_CREATE | os.O_WRONLY | os.O_APPEND
}
file, err := os.OpenFile(path, flag, 0o644)
if err != nil {
return "", err
}
defer file.Close()
if _, err := file.WriteString(args.Content); err != nil {
return "", err
}
if err := ctx.Err(); err != nil {
return "", err
}
content := []byte(args.Content)
perm := os.FileMode(0o644)
if info, statErr := os.Stat(path); statErr == nil {
perm = info.Mode().Perm()
}
if args.Append {
existing, readErr := os.ReadFile(path)
if readErr != nil && !os.IsNotExist(readErr) {
return "", readErr
}
content = append(existing, content...)
}
if err := safeio.AtomicWriteFile(path, content, perm); err != nil {
return "", err
}
action := "wrote"
if args.Append {
action = "appended"
+209
View File
@@ -0,0 +1,209 @@
package tools
import (
"context"
"encoding/json"
"errors"
"fmt"
"os"
"strings"
"agentu/internal/safeio"
"agentu/pkg/llm"
)
type FileEditTool struct {
workingDir string
}
func NewFileEditTool(workingDir string) *FileEditTool {
return &FileEditTool{workingDir: workingDir}
}
func (*FileEditTool) Mutates() bool { return true }
func (t *FileEditTool) Preview(ctx context.Context, raw json.RawMessage) (string, error) {
var args fileEditArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
if args.StartLine > 0 {
end := args.EndLine
if end == 0 {
end = args.StartLine
}
return fmt.Sprintf("edit %s lines %d-%d", args.Path, args.StartLine, end), ctx.Err()
}
return fmt.Sprintf("replace text in %s", args.Path), ctx.Err()
}
func (t *FileEditTool) TargetPaths(raw json.RawMessage) ([]string, error) {
var args fileEditArgs
if err := json.Unmarshal(raw, &args); err != nil {
return nil, fmt.Errorf("parse arguments: %w", err)
}
if strings.TrimSpace(args.Path) == "" {
return nil, errors.New("path is required")
}
path, err := resolvePath(t.workingDir, args.Path)
if err != nil {
return nil, err
}
return []string{path}, nil
}
func (t *FileEditTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "file_edit",
Description: "Edit an existing UTF-8 text file by replacing an inclusive 1-based line range or an exact old text pattern.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Existing file path relative to the configured agent working directory."},
"start_line": {"type": "integer", "description": "1-based first line for a line-range edit."},
"end_line": {"type": "integer", "description": "Optional 1-based final line; defaults to start_line when omitted."},
"old": {"type": "string", "description": "Exact text to replace for a pattern edit; not regex."},
"new": {"type": "string", "description": "Replacement text; may be empty string for deletion."},
"replace_all": {"type": "boolean", "description": "When true, replace every exact old occurrence; otherwise require exactly one occurrence."}
},
"required": ["path", "new"],
"additionalProperties": false
}`),
},
}
}
type fileEditArgs struct {
Path string `json:"path"`
StartLine int `json:"start_line"`
EndLine int `json:"end_line"`
Old string `json:"old"`
New *string `json:"new"`
ReplaceAll bool `json:"replace_all"`
}
func (t *FileEditTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args fileEditArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
if strings.TrimSpace(args.Path) == "" {
return "", errors.New("path is required")
}
if args.New == nil {
return "", errors.New("new is required")
}
path, err := resolvePath(t.workingDir, args.Path)
if err != nil {
return "", err
}
info, err := os.Stat(path)
if err != nil {
return "", fmt.Errorf("cannot access %s: %w", args.Path, err)
}
if info.IsDir() {
return "", fmt.Errorf("%s is a directory, not a file", args.Path)
}
perm := info.Mode().Perm()
hasLine := args.StartLine > 0 || args.EndLine > 0
hasOld := strings.TrimSpace(args.Old) != ""
if hasLine && hasOld {
return "", errors.New("provide either start_line/end_line or old, not both")
}
if !hasLine && !hasOld {
return "", errors.New("provide either start_line/end_line or old")
}
if err := ctx.Err(); err != nil {
return "", err
}
data, err := os.ReadFile(path)
if err != nil {
return "", err
}
content := string(data)
var result string
if hasOld {
result, err = editByPattern(content, args.Path, args.Old, *args.New, args.ReplaceAll)
} else {
result, err = editByLineRange(content, args.Path, args.StartLine, args.EndLine, *args.New)
}
if err != nil {
return "", err
}
if err := ctx.Err(); err != nil {
return "", err
}
if err := safeio.AtomicWriteFile(path, []byte(result), perm); err != nil {
return "", err
}
switch {
case hasOld:
count := strings.Count(content, args.Old)
if args.ReplaceAll {
return fmt.Sprintf("edited %s: replaced %d occurrences", args.Path, count), nil
}
return fmt.Sprintf("edited %s: replaced 1 occurrence", args.Path), nil
default:
endLine := args.EndLine
if endLine == 0 {
endLine = args.StartLine
}
return fmt.Sprintf("edited %s: replaced lines %d-%d", args.Path, args.StartLine, endLine), nil
}
}
func editByLineRange(content, displayPath string, startLine, endLine int, replacement string) (string, error) {
if startLine < 1 {
return "", errors.New("start_line must be >= 1")
}
if endLine == 0 {
endLine = startLine
}
if endLine < startLine {
return "", fmt.Errorf("end_line (%d) must be >= start_line (%d)", endLine, startLine)
}
lines := strings.Split(content, "\n")
if startLine > len(lines) {
return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
}
if endLine > len(lines) {
return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
}
var b strings.Builder
b.WriteString(strings.Join(lines[:startLine-1], "\n"))
if startLine > 1 {
b.WriteString("\n")
}
b.WriteString(replacement)
if endLine < len(lines) {
if !strings.HasSuffix(replacement, "\n") && replacement != "" {
b.WriteString("\n")
}
b.WriteString(strings.Join(lines[endLine:], "\n"))
}
return b.String(), nil
}
func editByPattern(content, displayPath, old, replacement string, replaceAll bool) (string, error) {
count := strings.Count(content, old)
if count == 0 {
return "", errors.New("old text not found")
}
if !replaceAll && count != 1 {
return "", fmt.Errorf("old text matched %d occurrences; set replace_all=true or provide a more specific old value", count)
}
return strings.ReplaceAll(content, old, replacement), nil
}
+136
View File
@@ -0,0 +1,136 @@
package tools
import (
"context"
"encoding/json"
"os"
"path/filepath"
"strings"
"testing"
)
func TestFileEditReplacesLineRange(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("one\ntwo\nthree\n"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","start_line":2,"end_line":2,"new":"TWO\n"}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "replaced lines 2-2") {
t.Fatalf("output = %q", out)
}
got, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(got) != "one\nTWO\nthree\n" {
t.Fatalf("file content = %q", string(got))
}
}
func TestFileEditReplacesExactPattern(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("alpha beta alpha"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"beta","new":"BETA"}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "replaced 1 occurrence") {
t.Fatalf("output = %q", out)
}
got, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(got) != "alpha BETA alpha" {
t.Fatalf("file content = %q", string(got))
}
}
func TestFileEditRejectsAmbiguousPattern(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("x x"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"x","new":"y"}`))
if err == nil {
t.Fatal("expected error for ambiguous pattern")
}
if !strings.Contains(err.Error(), "matched 2 occurrences") {
t.Fatalf("error = %v", err)
}
}
func TestFileEditReplaceAll(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("x x"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"x","new":"y","replace_all":true}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "replaced 2 occurrences") {
t.Fatalf("output = %q", out)
}
got, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(got) != "y y" {
t.Fatalf("file content = %q", string(got))
}
}
func TestFileEditRejectsEscapingPath(t *testing.T) {
dir := t.TempDir()
edit := NewFileEditTool(dir)
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"../outside.txt","old":"x","new":"y"}`))
if err == nil {
t.Fatal("expected error for escaping path")
}
if !strings.Contains(err.Error(), "escapes") {
t.Fatalf("error = %v", err)
}
}
func TestFileEditRejectsMissingNew(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "notes.txt")
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
edit := NewFileEditTool(dir)
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"hello"}`))
if err == nil {
t.Fatal("expected error for missing new")
}
if !strings.Contains(err.Error(), "new is required") {
t.Fatalf("error = %v", err)
}
}
func TestFileEditRejectsNonExistentFile(t *testing.T) {
dir := t.TempDir()
edit := NewFileEditTool(dir)
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"missing.txt","old":"x","new":"y"}`))
if err == nil {
t.Fatal("expected error for non-existent file")
}
}
+246
View File
@@ -0,0 +1,246 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"os/exec"
"path/filepath"
"strings"
"agentu/pkg/llm"
)
const (
defaultGitLogMaxCount = 20
maxGitLogMaxCount = 100
)
func runGit(ctx context.Context, workingDir string, args ...string) (string, error) {
gitArgs := append([]string{"-C", workingDir}, args...)
cmd := exec.CommandContext(ctx, "git", gitArgs...)
output, err := cmd.CombinedOutput()
text := FormatResult(string(output))
if err != nil {
if strings.Contains(err.Error(), "context deadline exceeded") || strings.Contains(err.Error(), "signal: killed") {
if text != "" {
return text, fmt.Errorf("git command timed out: %w", err)
}
return "", fmt.Errorf("git command timed out: %w", err)
}
if text != "" {
return text, fmt.Errorf("git command failed: %w", err)
}
return "", fmt.Errorf("git command failed: %w", err)
}
return text, nil
}
func gitPathspec(baseDir, path string) (string, error) {
if strings.TrimSpace(path) == "" {
return "", nil
}
resolved, err := resolvePath(baseDir, path)
if err != nil {
return "", err
}
rel, err := filepath.Rel(baseDir, resolved)
if err != nil {
return "", err
}
return filepath.ToSlash(filepath.Clean(rel)), nil
}
func normalizeGitLogMaxCount(value int) int {
if value <= 0 {
return defaultGitLogMaxCount
}
if value > maxGitLogMaxCount {
return maxGitLogMaxCount
}
return value
}
// git_status
type GitStatusTool struct {
workingDir string
}
func NewGitStatusTool(workingDir string) *GitStatusTool {
return &GitStatusTool{workingDir: workingDir}
}
func (t *GitStatusTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "git_status",
Description: "Show changed files using git status --short.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Optional path to limit status to. Relative paths are resolved from the configured agent working directory."}
},
"additionalProperties": false
}`),
},
}
}
type gitStatusArgs struct {
Path string `json:"path"`
}
func (t *GitStatusTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args gitStatusArgs
if len(raw) > 0 {
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
}
gitArgs := []string{"status", "--short"}
if args.Path != "" {
pathspec, err := gitPathspec(t.workingDir, args.Path)
if err != nil {
return "", err
}
gitArgs = append(gitArgs, "--", pathspec)
}
text, err := runGit(ctx, t.workingDir, gitArgs...)
if err != nil {
return text, err
}
if strings.TrimSpace(text) == "" {
return "Working tree clean.", nil
}
return "Changed files:\n" + text, nil
}
// git_diff
type GitDiffTool struct {
workingDir string
}
func NewGitDiffTool(workingDir string) *GitDiffTool {
return &GitDiffTool{workingDir: workingDir}
}
func (t *GitDiffTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "git_diff",
Description: "Show unstaged or staged git diff output, optionally as --stat.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Optional path to limit diff to. Relative paths are resolved from the configured agent working directory."},
"staged": {"type": "boolean", "description": "Show staged diff instead of unstaged."},
"stat": {"type": "boolean", "description": "Show diffstat summary instead of full diff."}
},
"additionalProperties": false
}`),
},
}
}
type gitDiffArgs struct {
Path string `json:"path"`
Staged bool `json:"staged"`
Stat bool `json:"stat"`
}
func (t *GitDiffTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args gitDiffArgs
if len(raw) > 0 {
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
}
gitArgs := []string{"diff"}
if args.Staged {
gitArgs = append(gitArgs, "--cached")
}
if args.Stat {
gitArgs = append(gitArgs, "--stat")
}
if args.Path != "" {
pathspec, err := gitPathspec(t.workingDir, args.Path)
if err != nil {
return "", err
}
gitArgs = append(gitArgs, "--", pathspec)
}
text, err := runGit(ctx, t.workingDir, gitArgs...)
if err != nil {
return text, err
}
if strings.TrimSpace(text) == "" {
if args.Staged {
return "No staged diff.", nil
}
return "No unstaged diff.", nil
}
return text, nil
}
// git_log
type GitLogTool struct {
workingDir string
}
func NewGitLogTool(workingDir string) *GitLogTool {
return &GitLogTool{workingDir: workingDir}
}
func (t *GitLogTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "git_log",
Description: "Show recent git commits with git log --oneline --decorate.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Optional path to limit log to. Relative paths are resolved from the configured agent working directory."},
"max_count": {"type": "integer", "description": "Maximum number of commits to show. Defaults to 20, capped at 100."}
},
"additionalProperties": false
}`),
},
}
}
type gitLogArgs struct {
Path string `json:"path"`
MaxCount int `json:"max_count"`
}
func (t *GitLogTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args gitLogArgs
if len(raw) > 0 {
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
}
maxCount := normalizeGitLogMaxCount(args.MaxCount)
gitArgs := []string{"log", "--oneline", "--decorate", fmt.Sprintf("-n%d", maxCount)}
if args.Path != "" {
pathspec, err := gitPathspec(t.workingDir, args.Path)
if err != nil {
return "", err
}
gitArgs = append(gitArgs, "--", pathspec)
}
text, err := runGit(ctx, t.workingDir, gitArgs...)
if err != nil {
return text, err
}
if strings.TrimSpace(text) == "" {
return "No git history.", nil
}
return text, nil
}
+142
View File
@@ -0,0 +1,142 @@
package tools
import (
"context"
"encoding/json"
"os"
"os/exec"
"path/filepath"
"strings"
"testing"
)
func requireGit(t *testing.T) {
t.Helper()
if _, err := exec.LookPath("git"); err != nil {
t.Skip("git not installed")
}
}
func runGitForTest(t *testing.T, dir string, args ...string) {
t.Helper()
gitArgs := append([]string{"-C", dir}, args...)
cmd := exec.Command("git", gitArgs...)
out, err := cmd.CombinedOutput()
if err != nil {
t.Fatalf("git %v failed: %v\n%s", args, err, out)
}
}
func initTestRepo(t *testing.T) string {
t.Helper()
dir := t.TempDir()
runGitForTest(t, dir, "init")
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "--allow-empty", "-m", "initial")
return dir
}
func TestGitStatusToolShowsChangedFiles(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
path := filepath.Join(dir, "tracked.txt")
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
runGitForTest(t, dir, "add", "tracked.txt")
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
t.Fatal(err)
}
tool := NewGitStatusTool(dir)
out, err := tool.Execute(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "Changed files:") {
t.Fatalf("output missing 'Changed files:': %q", out)
}
if !strings.Contains(out, "M tracked.txt") {
t.Fatalf("output missing 'M tracked.txt': %q", out)
}
}
func TestGitDiffToolShowsUnstagedDiff(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
path := filepath.Join(dir, "tracked.txt")
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
runGitForTest(t, dir, "add", "tracked.txt")
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
t.Fatal(err)
}
tool := NewGitDiffTool(dir)
out, err := tool.Execute(context.Background(), nil)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "diff --git") {
t.Fatalf("output missing 'diff --git': %q", out)
}
}
func TestGitDiffToolShowsStat(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
path := filepath.Join(dir, "tracked.txt")
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
runGitForTest(t, dir, "add", "tracked.txt")
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
t.Fatal(err)
}
tool := NewGitDiffTool(dir)
out, err := tool.Execute(context.Background(), json.RawMessage(`{"stat":true}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "tracked.txt") {
t.Fatalf("output missing 'tracked.txt': %q", out)
}
}
func TestGitLogToolShowsCommits(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
tool := NewGitLogTool(dir)
out, err := tool.Execute(context.Background(), json.RawMessage(`{"max_count":1}`))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "initial") {
t.Fatalf("output missing 'initial': %q", out)
}
}
func TestGitToolsRejectEscapingPath(t *testing.T) {
requireGit(t)
dir := initTestRepo(t)
tool := NewGitStatusTool(dir)
_, err := tool.Execute(context.Background(), json.RawMessage(`{"path":"../outside"}`))
if err == nil {
t.Fatal("expected error for escaping path")
}
if !strings.Contains(err.Error(), "escapes") {
t.Fatalf("error = %v", err)
}
}
+244
View File
@@ -0,0 +1,244 @@
package tools
import (
"context"
"encoding/json"
"errors"
"fmt"
"io/fs"
"os"
"path/filepath"
"strings"
"agentu/pkg/llm"
)
const defaultOutlineMaxFiles = 200
// ────────────────────────────────────────────────────────────────────────────
// ProjectOutlineTool — gitignore-aware file tree + symbol outline
// ────────────────────────────────────────────────────────────────────────────
type ProjectOutlineTool struct {
workingDir string
}
func NewProjectOutlineTool(workingDir string) *ProjectOutlineTool {
return &ProjectOutlineTool{workingDir: workingDir}
}
func (t *ProjectOutlineTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "project_outline",
Description: "Get the project file tree and code symbol outline in one call. Respects .gitignore. Use this first to understand project structure before reading individual files.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Subdirectory to outline. Defaults to the project root."},
"max_files": {"type": "integer", "description": "Maximum files to include in the outline. Defaults to 200."}
},
"additionalProperties": false
}`),
},
}
}
type outlineArgs struct {
Path string `json:"path"`
MaxFiles int `json:"max_files"`
}
func (t *ProjectOutlineTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args outlineArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
root, err := resolvePath(t.workingDir, args.Path)
if err != nil {
return "", err
}
limit := args.MaxFiles
if limit <= 0 {
limit = defaultOutlineMaxFiles
}
// Parse .gitignore files from the root upward.
ignorer := newGitignoreMatcher(root)
var out strings.Builder
out.WriteString(root)
out.WriteString("\n")
count := 0
truncated := false
err = filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
if err != nil {
return nil
}
if ctx.Err() != nil {
return ctx.Err()
}
// Skip always-ignored directories.
if d.IsDir() {
name := d.Name()
// Always skip VCS and common build-artifact directories.
switch name {
case ".git", ".hg", ".svn", "vendor", "node_modules",
"dist", "build", "__pycache__", "target":
return filepath.SkipDir
}
// Skip hidden directories other than .github and similar.
if strings.HasPrefix(name, ".") && name != ".github" && name != ".vscode" && name != ".devcontainer" {
return filepath.SkipDir
}
}
rel, _ := filepath.Rel(root, path)
if rel == "." {
return nil
}
// Check gitignore.
if ignorer.shouldIgnore(rel, d.IsDir()) {
if d.IsDir() {
return filepath.SkipDir
}
return nil
}
if count >= limit {
truncated = true
return filepath.SkipAll
}
// Build indented tree entry.
depth := strings.Count(rel, string(filepath.Separator))
indent := strings.Repeat(" ", depth)
name := d.Name()
if d.IsDir() {
out.WriteString(fmt.Sprintf("%s%s/\n", indent, name))
return nil
}
count++
out.WriteString(indent + name)
// Extract symbols for source files.
ext := filepath.Ext(name)
if sourceFileExtensions[ext] {
data, readErr := os.ReadFile(path)
if readErr == nil && len(data) <= 256*1024 {
symbols, symErr := extractSymbols(path, data)
if symErr == nil && len(symbols) > 0 {
out.WriteString(" — ")
out.WriteString(compactSymbolSummary(symbols))
}
}
}
out.WriteString("\n")
return nil
})
if err != nil && !errors.Is(err, filepath.SkipAll) {
return "", err
}
if truncated {
out.WriteString(fmt.Sprintf("\n... (%d files shown, more available; use file_list or code_symbols for details)\n", count))
}
return FormatResult(out.String()), nil
}
// ────────────────────────────────────────────────────────────────────────────
// Compact symbol summary for inline display in the file tree
// ────────────────────────────────────────────────────────────────────────────
func compactSymbolSummary(symbols []symbolInfo) string {
var parts []string
// Group by kind.
funcs := []string{}
methods := []string{}
types := []string{}
for _, s := range symbols {
switch s.Kind {
case "function":
funcs = append(funcs, s.Name)
case "method":
methods = append(methods, s.Name)
case "class", "interface", "struct", "enum", "trait", "alias", "func", "type", "impl":
types = append(types, s.Name)
}
// Skip imports and package for compact view.
}
if len(funcs) > 0 {
parts = append(parts, fmt.Sprintf("fn: %s", strings.Join(funcs, ", ")))
}
if len(methods) > 0 {
parts = append(parts, fmt.Sprintf("method: %s", strings.Join(methods, ", ")))
}
if len(types) > 0 {
parts = append(parts, fmt.Sprintf("type: %s", strings.Join(types, ", ")))
}
return strings.Join(parts, "; ")
}
// ────────────────────────────────────────────────────────────────────────────
// Simple .gitignore matcher
// ────────────────────────────────────────────────────────────────────────────
type gitignoreMatcher struct {
patterns []string
}
func newGitignoreMatcher(root string) *gitignoreMatcher {
m := &gitignoreMatcher{}
// Walk up from root, collecting .gitignore patterns.
for dir := root; dir != "" && dir != string(filepath.Separator); dir = filepath.Dir(dir) {
giPath := filepath.Join(dir, ".gitignore")
data, err := os.ReadFile(giPath)
if err != nil {
continue
}
for _, line := range strings.Split(string(data), "\n") {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
continue
}
// Remove leading / for anchored patterns; we handle all as relative.
line = strings.TrimPrefix(line, "/")
// Remove trailing / for directory-only patterns.
line = strings.TrimSuffix(line, "/")
m.patterns = append(m.patterns, line)
}
}
return m
}
func (m *gitignoreMatcher) shouldIgnore(path string, isDir bool) bool {
name := filepath.Base(path)
for _, pat := range m.patterns {
// Simple match against full relative path or just the basename.
matched, _ := filepath.Match(pat, path)
if matched {
return true
}
matched, _ = filepath.Match(pat, name)
if matched {
return true
}
}
return false
}
+430
View File
@@ -0,0 +1,430 @@
package tools
import (
"fmt"
"os"
"path/filepath"
"strings"
"agentu/internal/safeio"
)
// ────────────────────────────────────────────────────────────────────────────
// V4 patch format data structures
// ────────────────────────────────────────────────────────────────────────────
type patchAction string
const (
actionAdd patchAction = "add"
actionUpdate patchAction = "update"
actionDelete patchAction = "delete"
)
type patchFile struct {
action patchAction
path string
newContent string // for "add" — full file content
lines []string // for "update" — lines to match/replace; for "delete" — unused
}
type parsedPatch struct {
files []patchFile
}
// ────────────────────────────────────────────────────────────────────────────
// Parser
// ────────────────────────────────────────────────────────────────────────────
// parsePatch parses a V4-format patch string.
//
// Format:
//
// *** Begin Patch
// *** Add File: <path>
// <file content>
// *** Update File: <path>
// @@
// <context line>
// @@
// -<old line>
// +<new line>
// *** Delete File: <path>
// *** End Patch
func parsePatch(patch string) (*parsedPatch, error) {
patch = strings.TrimSpace(patch)
if patch == "" {
return nil, fmt.Errorf("patch is empty")
}
lines := strings.Split(patch, "\n")
if len(lines) == 0 {
return nil, fmt.Errorf("patch is empty")
}
if strings.TrimSpace(lines[0]) != "*** Begin Patch" {
return nil, fmt.Errorf("patch must start with *** Begin Patch")
}
var files []patchFile
seenPaths := make(map[string]struct{})
seenEnd := false
i := 1
for i < len(lines) {
line := strings.TrimSpace(lines[i])
switch {
case line == "*** End Patch":
seenEnd = true
i = len(lines)
case strings.HasPrefix(line, "*** Add File:"):
path := strings.TrimSpace(strings.TrimPrefix(line, "*** Add File:"))
if err := validatePatchPath(path, seenPaths); err != nil {
return nil, err
}
i++
var contentLines []string
for i < len(lines) {
next := strings.TrimSpace(lines[i])
if strings.HasPrefix(next, "***") {
break
}
content := lines[i]
if strings.HasPrefix(content, "+") {
content = strings.TrimPrefix(content, "+")
}
contentLines = append(contentLines, content)
i++
}
files = append(files, patchFile{
action: actionAdd,
path: path,
newContent: strings.Join(contentLines, "\n") + "\n",
})
case strings.HasPrefix(line, "*** Delete File:"):
path := strings.TrimSpace(strings.TrimPrefix(line, "*** Delete File:"))
if err := validatePatchPath(path, seenPaths); err != nil {
return nil, err
}
files = append(files, patchFile{
action: actionDelete,
path: path,
})
i++
case strings.HasPrefix(line, "*** Update File:"):
path := strings.TrimSpace(strings.TrimPrefix(line, "*** Update File:"))
if err := validatePatchPath(path, seenPaths); err != nil {
return nil, err
}
i++
var hunkLines []string
for i < len(lines) {
next := strings.TrimSpace(lines[i])
if strings.HasPrefix(next, "***") {
break
}
hunkLines = append(hunkLines, lines[i])
i++
}
files = append(files, patchFile{
action: actionUpdate,
path: path,
lines: hunkLines,
})
default:
i++
}
}
if len(files) == 0 {
return nil, fmt.Errorf("patch contains no file operations")
}
if !seenEnd {
return nil, fmt.Errorf("patch must end with *** End Patch")
}
return &parsedPatch{files: files}, nil
}
func validatePatchPath(path string, seen map[string]struct{}) error {
if path == "" {
return fmt.Errorf("patch file path is required")
}
if _, ok := seen[path]; ok {
return fmt.Errorf("patch contains duplicate file path: %s", path)
}
seen[path] = struct{}{}
return nil
}
// ────────────────────────────────────────────────────────────────────────────
// Applier
// ────────────────────────────────────────────────────────────────────────────
// applyPatchFile applies a single patch file operation.
func applyPatchFile(baseDir string, pf patchFile) (string, error) {
switch pf.action {
case actionAdd:
return applyAdd(baseDir, pf)
case actionUpdate:
return applyUpdate(baseDir, pf)
case actionDelete:
return applyDelete(baseDir, pf)
default:
return "", fmt.Errorf("unknown patch action: %s", pf.action)
}
}
func applyAdd(baseDir string, pf patchFile) (string, error) {
path, err := resolvePathForWrite(baseDir, pf.path)
if err != nil {
return "", err
}
if _, err := os.Stat(path); err == nil {
return "", fmt.Errorf("cannot add existing file: %s", pf.path)
} else if !os.IsNotExist(err) {
return "", fmt.Errorf("stat new file %s: %w", pf.path, err)
}
if err := safeio.AtomicWriteFile(path, []byte(pf.newContent), 0o644); err != nil {
return "", fmt.Errorf("write new file %s: %w", pf.path, err)
}
return fmt.Sprintf("+ %s", pf.path), nil
}
func applyDelete(baseDir string, pf patchFile) (string, error) {
path, err := resolvePath(baseDir, pf.path)
if err != nil {
return "", err
}
if info, err := os.Stat(path); err != nil {
if os.IsNotExist(err) {
return "", fmt.Errorf("cannot delete non-existent file: %s", pf.path)
}
return "", err
} else if info.IsDir() {
return "", fmt.Errorf("cannot delete directory: %s", pf.path)
}
if err := os.Remove(path); err != nil {
return "", fmt.Errorf("delete file %s: %w", pf.path, err)
}
// Try to remove empty parent directories.
dir := filepath.Dir(path)
for dir != baseDir && dir != "." {
if entries, err := os.ReadDir(dir); err == nil && len(entries) == 0 {
os.Remove(dir)
dir = filepath.Dir(dir)
} else {
break
}
}
return fmt.Sprintf("- %s", pf.path), nil
}
func applyUpdate(baseDir string, pf patchFile) (string, error) {
path, err := resolvePath(baseDir, pf.path)
if err != nil {
return "", fmt.Errorf("cannot update non-existent file: %s: %w", pf.path, err)
}
existing, err := os.ReadFile(path)
if err != nil {
return "", fmt.Errorf("cannot update non-existent file: %s", pf.path)
}
info, err := os.Stat(path)
if err != nil {
return "", fmt.Errorf("stat file %s: %w", pf.path, err)
}
fileLines := strings.Split(string(existing), "\n")
// Parse the hunk lines into: context blocks + old lines + new lines.
result, err := applyHunkLines(fileLines, pf.lines, pf.path)
if err != nil {
return "", err
}
newContent := strings.Join(result, "\n")
if err := safeio.AtomicWriteFile(path, []byte(newContent), info.Mode().Perm()); err != nil {
return "", fmt.Errorf("write updated file %s: %w", pf.path, err)
}
return fmt.Sprintf("~ %s", pf.path), nil
}
// applyHunkLines applies a sequence of hunk lines to the file content.
// The hunk lines are a mix of:
//
// @@ — hunk separator (a context line block starts)
// <plain line> — context line (must appear in both old and new)
// -<line> — line to remove
// +<line> — line to add
//
// Hunks are matched greedily: for each group of contiguous context + - + + lines,
// find the first position in the file where context + - lines match,
// replace with context + + lines.
func applyHunkLines(fileLines []string, hunkLines []string, path string) ([]string, error) {
if len(hunkLines) == 0 {
return fileLines, nil
}
// Split the hunk lines into separate hunks at @@ markers.
type hunk struct {
lines []string
}
var hunks []hunk
var current hunk
for _, hl := range hunkLines {
trimmed := strings.TrimSpace(hl)
if strings.HasPrefix(trimmed, "@@") {
if len(current.lines) > 0 {
hunks = append(hunks, current)
current = hunk{}
}
continue
}
current.lines = append(current.lines, hl)
}
if len(current.lines) > 0 {
hunks = append(hunks, current)
}
if len(hunks) == 0 {
return fileLines, nil
}
result := make([]string, len(fileLines))
copy(result, fileLines)
offset := 0 // tracks net line count change from previous hunks
for hi, h := range hunks {
// Separate the hunk lines into a search pattern and replacement.
searchLines, replaceLines := buildHunkSearchReplace(h.lines)
if len(searchLines) == 0 && len(replaceLines) == 0 {
continue
}
matchStart := findHunkMatch(result, searchLines, offset)
if matchStart < 0 {
return nil, fmt.Errorf("hunk context not found in %s", path)
}
// Replace: remove searchLines, insert replaceLines.
matchEnd := matchStart + len(searchLines)
newResult := make([]string, 0, len(result)-len(searchLines)+len(replaceLines))
newResult = append(newResult, result[:matchStart]...)
newResult = append(newResult, replaceLines...)
newResult = append(newResult, result[matchEnd:]...)
result = newResult
offset = matchStart + len(replaceLines)
// Multiple matches check for this hunk only.
secondMatch := findHunkMatch(result, searchLines, offset)
if secondMatch >= 0 {
return nil, fmt.Errorf("hunk context matched multiple times in %s", path)
}
_ = hi
}
return result, nil
}
// buildHunkSearchReplace splits a hunk's lines into the search pattern
// (context + old lines) and replacement (context + new lines).
func buildHunkSearchReplace(lines []string) (search, replace []string) {
for _, l := range lines {
if strings.HasPrefix(l, "-") {
search = append(search, l[1:])
} else if strings.HasPrefix(l, "+") {
replace = append(replace, l[1:])
} else {
contextLine := strings.TrimPrefix(l, " ")
search = append(search, contextLine)
replace = append(replace, contextLine)
}
}
return
}
// findHunkMatch finds the first position >= offset in fileLines where
// searchLines matches exactly (line-by-line comparison).
func findHunkMatch(fileLines []string, searchLines []string, offset int) int {
for i := offset; i <= len(fileLines)-len(searchLines); i++ {
match := true
for j, sl := range searchLines {
if fileLines[i+j] != sl {
match = false
break
}
}
if match {
return i
}
}
return -1
}
// ────────────────────────────────────────────────────────────────────────────
// Preview
// ────────────────────────────────────────────────────────────────────────────
// summarizePatch returns a human-readable summary of the patch.
func summarizePatch(patch string) (string, error) {
pp, err := parsePatch(patch)
if err != nil {
return "", err
}
var b strings.Builder
added, updated, deleted := 0, 0, 0
for _, pf := range pp.files {
switch pf.action {
case actionAdd:
added++
fmt.Fprintf(&b, "+ add %s\n", pf.path)
case actionUpdate:
updated++
fmt.Fprintf(&b, "~ update %s\n", pf.path)
case actionDelete:
deleted++
fmt.Fprintf(&b, "- delete %s\n", pf.path)
}
}
summary := fmt.Sprintf("patch: %d file(s) changed", len(pp.files))
if added > 0 {
summary += fmt.Sprintf(" (%d added)", added)
}
if updated > 0 {
summary += fmt.Sprintf(" (%d updated)", updated)
}
if deleted > 0 {
summary += fmt.Sprintf(" (%d deleted)", deleted)
}
return summary + "\n" + b.String(), nil
}
// extractTargetPaths returns all file paths referenced in the patch.
func extractTargetPaths(baseDir, patch string) ([]string, error) {
pp, err := parsePatch(patch)
if err != nil {
return nil, err
}
var paths []string
for _, pf := range pp.files {
// Use resolvePathForWrite for all files — TargetPaths needs paths
// even for files that don't exist yet (rollback snapshots).
resolved, err := resolvePathForWrite(baseDir, pf.path)
if err != nil {
return nil, fmt.Errorf("file %s: %w", pf.path, err)
}
paths = append(paths, resolved)
}
return paths, nil
}
+68 -5
View File
@@ -2,22 +2,85 @@ package tools
import (
"fmt"
"os"
"path/filepath"
"strings"
)
func resolvePath(baseDir, path string) (string, error) {
resolved, canonicalBase, err := resolveLexicalPath(baseDir, path)
if err != nil {
return "", err
}
canonical, err := filepath.EvalSymlinks(resolved)
if err != nil {
return "", err
}
if !isWithinDir(canonical, canonicalBase) {
return "", fmt.Errorf("path escapes the project directory through a symbolic link: %s", path)
}
return canonical, nil
}
// resolvePathForWrite resolves a path that may not exist yet while preventing
// existing parent-directory symlinks from escaping the project directory.
func resolvePathForWrite(baseDir, path string) (string, error) {
resolved, canonicalBase, err := resolveLexicalPath(baseDir, path)
if err != nil {
return "", err
}
existing := resolved
var suffix []string
for {
_, err := os.Lstat(existing)
if err == nil {
break
}
if !os.IsNotExist(err) {
return "", err
}
parent := filepath.Dir(existing)
if parent == existing {
return "", err
}
suffix = append(suffix, filepath.Base(existing))
existing = parent
}
canonical, err := filepath.EvalSymlinks(existing)
if err != nil {
return "", err
}
for i := len(suffix) - 1; i >= 0; i-- {
canonical = filepath.Join(canonical, suffix[i])
}
if !isWithinDir(canonical, canonicalBase) {
return "", fmt.Errorf("path escapes the project directory through a symbolic link: %s", path)
}
return canonical, nil
}
func resolveLexicalPath(baseDir, path string) (string, string, error) {
if path == "" {
path = "."
}
if filepath.IsAbs(path) {
return "", 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))
if !isWithinDir(resolved, baseDir) {
return "", fmt.Errorf("path escapes the project directory: %s", path)
absBase, err := filepath.Abs(baseDir)
if err != nil {
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 {
+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
}
+52 -5
View File
@@ -1,6 +1,7 @@
package tools
import (
"bytes"
"context"
"encoding/json"
"errors"
@@ -8,6 +9,7 @@ import (
"os"
"os/exec"
"strings"
"sync"
"agentu/pkg/llm"
)
@@ -19,6 +21,18 @@ type ShellRunTool struct {
func NewShellRunTool(workingDir string) *ShellRunTool {
return &ShellRunTool{workingDir: workingDir}
}
func (*ShellRunTool) Mutates() bool { return true }
func (t *ShellRunTool) Preview(ctx context.Context, raw json.RawMessage) (string, error) {
var args shellRunArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
if strings.TrimSpace(args.Command) == "" {
return "", errors.New("command is required")
}
return "$ " + strings.TrimSpace(args.Command), ctx.Err()
}
func (t *ShellRunTool) Definition() llm.Tool {
return llm.Tool{
@@ -45,6 +59,10 @@ type shellRunArgs struct {
}
func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
return t.ExecuteStream(ctx, raw, nil)
}
func (t *ShellRunTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
var args shellRunArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
@@ -65,16 +83,45 @@ func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string
cmd := exec.CommandContext(ctx, shell, "-lc", args.Command)
cmd.Dir = workdir
output, err := cmd.CombinedOutput()
text := FormatResult(string(output))
var mu sync.Mutex
var combined bytes.Buffer
writer := &shellStreamWriter{mu: &mu, buf: &combined, emit: emit}
cmd.Stdout = writer
cmd.Stderr = writer
runErr := cmd.Run()
text := FormatResult(combined.String())
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
return text, fmt.Errorf("command timed out: %w", ctx.Err())
}
if err != nil {
if runErr != nil {
if text != "" {
return text, fmt.Errorf("command failed: %w", err)
return text, fmt.Errorf("command failed: %w", runErr)
}
return "", fmt.Errorf("command failed: %w", err)
return "", fmt.Errorf("command failed: %w", runErr)
}
return text, nil
}
// shellStreamWriter captures command output for the final result while
// forwarding every write to the streaming callback as it happens.
type shellStreamWriter struct {
mu *sync.Mutex
buf *bytes.Buffer
emit func(string) error
}
func (w *shellStreamWriter) Write(p []byte) (int, error) {
w.mu.Lock()
n, writeErr := w.buf.Write(p)
w.mu.Unlock()
if writeErr != nil {
return n, writeErr
}
if w.emit != nil {
if err := w.emit(string(p)); err != nil {
return 0, err
}
}
return len(p), nil
}
+71
View File
@@ -0,0 +1,71 @@
package tools
import (
"context"
"encoding/json"
"strings"
"sync"
"testing"
)
func TestShellRunStreamsOutputWhileRunning(t *testing.T) {
tool := NewShellRunTool(t.TempDir())
var mu sync.Mutex
var chunks []string
output, err := tool.ExecuteStream(context.Background(), json.RawMessage(`{"command":"printf 'a\\n'; sleep 0.05; printf 'b\\n'"}`), func(chunk string) error {
mu.Lock()
chunks = append(chunks, chunk)
mu.Unlock()
return nil
})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(output, "a") || !strings.Contains(output, "b") {
t.Fatalf("output = %q", output)
}
mu.Lock()
joined := strings.Join(chunks, "")
mu.Unlock()
if len(chunks) == 0 {
t.Fatal("expected at least one streamed chunk")
}
if !strings.Contains(joined, "a") || !strings.Contains(joined, "b") {
t.Fatalf("streamed chunks = %q", joined)
}
}
func TestShellRunExecuteCollectsFullOutput(t *testing.T) {
tool := NewShellRunTool(t.TempDir())
output, err := tool.Execute(context.Background(), json.RawMessage(`{"command":"printf 'hello world'"}`))
if err != nil {
t.Fatal(err)
}
if strings.TrimSpace(output) != "hello world" {
t.Fatalf("output = %q", output)
}
}
func TestShellRunStreamsStderrAndStdout(t *testing.T) {
tool := NewShellRunTool(t.TempDir())
var mu sync.Mutex
var chunks []string
output, err := tool.ExecuteStream(context.Background(), json.RawMessage(`{"command":"echo out; echo err >&2"}`), func(chunk string) error {
mu.Lock()
chunks = append(chunks, chunk)
mu.Unlock()
return nil
})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(output, "out") || !strings.Contains(output, "err") {
t.Fatalf("output = %q", output)
}
mu.Lock()
joined := strings.Join(chunks, "")
mu.Unlock()
if !strings.Contains(joined, "out") || !strings.Contains(joined, "err") {
t.Fatalf("streamed chunks = %q", joined)
}
}
+530
View File
@@ -0,0 +1,530 @@
package tools
import (
"bufio"
"bytes"
"fmt"
"go/ast"
"go/parser"
"go/printer"
"go/token"
"regexp"
"sort"
"strings"
)
// symbolInfo holds extracted symbol metadata.
type symbolInfo struct {
Name string
Kind string // "import", "function", "method", "class", "type", "interface", "trait", "enum", "struct"
StartLine int
}
// sourceFileExtensions lists extensions that we can extract symbols from.
var sourceFileExtensions = map[string]bool{
".go": true,
".ts": true,
".tsx": true,
".js": true,
".jsx": true,
".mjs": true,
".py": true,
".rs": true,
}
// extractSymbols extracts symbols from a source file using language-specific
// parsing (go/ast for Go, regex for other languages).
func extractSymbols(path string, data []byte) ([]symbolInfo, error) {
ext := ""
if idx := strings.LastIndex(path, "."); idx >= 0 {
ext = path[idx:]
}
switch ext {
case ".go":
return extractGoSymbols(path, data)
case ".py":
return extractPythonsymbols(data)
case ".ts", ".tsx", ".js", ".jsx", ".mjs":
return extractTSSymbols(data)
case ".rs":
return extractRustSymbols(data)
default:
return nil, nil
}
}
// ────────────────────────────────────────────────────────────────────────────
// Go symbol extraction (using go/ast)
// ────────────────────────────────────────────────────────────────────────────
func extractGoSymbols(path string, data []byte) ([]symbolInfo, error) {
fset := token.NewFileSet()
// Parse imports only first for the import list.
importFile, err := parser.ParseFile(fset, path, data, parser.ImportsOnly|parser.ParseComments)
if err != nil {
return nil, err
}
full, err := parser.ParseFile(fset, path, data, 0)
if err != nil {
return nil, err
}
var symbols []symbolInfo
// Package
symbols = append(symbols, symbolInfo{
Name: full.Name.Name,
Kind: "package",
StartLine: fset.Position(full.Name.Pos()).Line,
})
// Imports
for _, imp := range importFile.Imports {
symbols = append(symbols, symbolInfo{
Name: strings.Trim(imp.Path.Value, "\""),
Kind: "import",
StartLine: fset.Position(imp.Pos()).Line,
})
}
// Declarations
for _, decl := range full.Decls {
switch d := decl.(type) {
case *ast.GenDecl:
if d.Tok != token.TYPE {
continue
}
for _, spec := range d.Specs {
if ts, ok := spec.(*ast.TypeSpec); ok {
kind := goTypeKind(ts.Type)
symbols = append(symbols, symbolInfo{
Name: ts.Name.Name,
Kind: kind,
StartLine: fset.Position(ts.Pos()).Line,
})
}
}
case *ast.FuncDecl:
line := fset.Position(d.Pos()).Line
if d.Recv != nil {
symbols = append(symbols, symbolInfo{
Name: fmt.Sprintf("%s.%s", goReceiverName(d.Recv), d.Name.Name),
Kind: "method",
StartLine: line,
})
} else {
symbols = append(symbols, symbolInfo{
Name: d.Name.Name,
Kind: "function",
StartLine: line,
})
}
}
}
sortSymbols(symbols)
return symbols, nil
}
func goTypeKind(expr ast.Expr) string {
switch expr.(type) {
case *ast.StructType:
return "struct"
case *ast.InterfaceType:
return "interface"
case *ast.FuncType:
return "func"
default:
return "alias"
}
}
func goReceiverName(recv *ast.FieldList) string {
if recv == nil || len(recv.List) == 0 {
return "?"
}
var b bytes.Buffer
if err := printer.Fprint(&b, token.NewFileSet(), recv.List[0].Type); err != nil {
return "?"
}
return b.String()
}
// ────────────────────────────────────────────────────────────────────────────
// Python symbol extraction (regex-based)
// ────────────────────────────────────────────────────────────────────────────
var (
pyFuncRe = regexp.MustCompile(`^\s*def\s+(\w+)\s*\(`)
pyClassRe = regexp.MustCompile(`^\s*class\s+(\w+)\s*[(:]`)
pyImportRe = regexp.MustCompile(`^\s*(?:from\s+\S+\s+)?import\s+`)
pyDecoratorRe = regexp.MustCompile(`^\s*@`)
)
func extractPythonsymbols(data []byte) ([]symbolInfo, error) {
var symbols []symbolInfo
scanner := bufio.NewScanner(bytes.NewReader(data))
lineNum := 0
inDecorator := false
for scanner.Scan() {
lineNum++
line := scanner.Text()
// Track decorator stacks
if pyDecoratorRe.MatchString(line) {
inDecorator = true
continue
}
// Import lines
if pyImportRe.MatchString(line) {
// Extract the module name for a simple "import foo" or "from foo import bar"
importPath := strings.TrimSpace(line)
symbols = append(symbols, symbolInfo{
Name: importPath,
Kind: "import",
StartLine: lineNum,
})
inDecorator = false
continue
}
matches := pyClassRe.FindStringSubmatch(line)
if matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "class",
StartLine: lineNum,
})
inDecorator = false
continue
}
matches = pyFuncRe.FindStringSubmatch(line)
if matches != nil {
kind := "function"
if inDecorator {
// Could be a decorated function or method — still mark as function
// (class-method detection requires indentation tracking)
}
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: kind,
StartLine: lineNum,
})
inDecorator = false
continue
}
inDecorator = false
}
sortSymbols(symbols)
return symbols, scanner.Err()
}
// ────────────────────────────────────────────────────────────────────────────
// TypeScript / JavaScript symbol extraction (regex-based)
// ────────────────────────────────────────────────────────────────────────────
var (
tsFuncRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?(?:async\s+)?function\s+(\w+)\s*[<(]`)
tsArrowRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s*)?\s*\(`)
tsClassRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?class\s+(\w+)`)
tsInterfaceRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?interface\s+(\w+)`)
tsTypeRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?type\s+(\w+)\s*=`)
tsEnumRe = regexp.MustCompile(`(?:export\s+(?:default\s+)?)?enum\s+(\w+)`)
tsImportRe = regexp.MustCompile(`^\s*import\s+`)
tsMethodRe = regexp.MustCompile(`^\s*(?:public\s+|private\s+|protected\s+|static\s+|async\s+)*(\w+)\s*[<(][^)]*\)\s*(?::\s*\S+\s*)?\{`)
)
func extractTSSymbols(data []byte) ([]symbolInfo, error) {
var symbols []symbolInfo
scanner := bufio.NewScanner(bytes.NewReader(data))
lineNum := 0
// Track brace depth to detect class body methods.
braceDepth := 0
inClass := false
for scanner.Scan() {
lineNum++
line := scanner.Text()
trimmed := strings.TrimSpace(line)
// Count braces on this line (applied at end of iteration).
opens := strings.Count(line, "{")
closes := strings.Count(line, "}")
// Import lines — skip everything else, just track braces.
if tsImportRe.MatchString(line) {
importPath := trimmed
symbols = append(symbols, symbolInfo{
Name: importPath,
Kind: "import",
StartLine: lineNum,
})
goto applyBraces
}
// Skip empty and comment-only lines.
if trimmed == "" || strings.HasPrefix(trimmed, "//") || strings.HasPrefix(trimmed, "/*") {
goto applyBraces
}
// Interface
if matches := tsInterfaceRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "interface",
StartLine: lineNum,
})
goto applyBraces
}
// Type alias
if matches := tsTypeRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "type",
StartLine: lineNum,
})
goto applyBraces
}
// Enum
if matches := tsEnumRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "enum",
StartLine: lineNum,
})
goto applyBraces
}
// Class
if matches := tsClassRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "class",
StartLine: lineNum,
})
// Set inClass BEFORE applying braces for this line, so that
// methods on the same line (unlikely) or next line are detected.
if opens > 0 {
inClass = true
}
goto applyBraces
}
// Method inside class body (at depth 1 brace level).
if inClass && braceDepth == 1 {
if matches := tsMethodRe.FindStringSubmatch(line); matches != nil {
name := matches[1]
if !isKeyword(name) {
symbols = append(symbols, symbolInfo{
Name: name,
Kind: "method",
StartLine: lineNum,
})
}
}
}
// Function declaration (top-level only, or nested function).
if !inClass || braceDepth > 1 {
if matches := tsFuncRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "function",
StartLine: lineNum,
})
}
}
// Arrow function assigned to const/let/var.
if !inClass || braceDepth > 1 {
if matches := tsArrowRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "function",
StartLine: lineNum,
})
}
}
applyBraces:
braceDepth += opens - closes
if braceDepth <= 0 {
braceDepth = 0
inClass = false
}
}
sortSymbols(symbols)
return symbols, scanner.Err()
}
func isKeyword(s string) bool {
switch s {
case "if", "else", "for", "while", "switch", "case", "return", "throw",
"try", "catch", "finally", "new", "delete", "typeof", "instanceof",
"constructor", "get", "set":
return true
}
return false
}
// ────────────────────────────────────────────────────────────────────────────
// Rust symbol extraction (regex-based)
// ────────────────────────────────────────────────────────────────────────────
var (
rsFnRe = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?(?:async\s+)?(?:unsafe\s+)?fn\s+(\w+)\s*[<(]`)
rsStructRe = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?struct\s+(\w+)`)
rsEnumRe = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?enum\s+(\w+)`)
rsTraitRe = regexp.MustCompile(`^\s*(?:pub(?:\s*\(\s*crate\s*\))?\s+)?(?:unsafe\s+)?trait\s+(\w+)`)
rsImplRe = regexp.MustCompile(`^\s*impl\s+`)
rsUseRe = regexp.MustCompile(`^\s*(?:pub\s+)?use\s+`)
)
func extractRustSymbols(data []byte) ([]symbolInfo, error) {
var symbols []symbolInfo
scanner := bufio.NewScanner(bytes.NewReader(data))
lineNum := 0
for scanner.Scan() {
lineNum++
line := scanner.Text()
trimmed := strings.TrimSpace(line)
// Skip empty and comments
if trimmed == "" || strings.HasPrefix(trimmed, "//") || strings.HasPrefix(trimmed, "/*") {
continue
}
// Use
if rsUseRe.MatchString(line) {
symbols = append(symbols, symbolInfo{
Name: trimmed,
Kind: "import",
StartLine: lineNum,
})
continue
}
// Struct
if matches := rsStructRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "struct",
StartLine: lineNum,
})
continue
}
// Enum
if matches := rsEnumRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "enum",
StartLine: lineNum,
})
continue
}
// Trait
if matches := rsTraitRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "trait",
StartLine: lineNum,
})
continue
}
// Impl
if rsImplRe.MatchString(line) {
// Extract type name from "impl Foo" or "impl<T> Foo" or "impl Trait for Foo"
rest := strings.TrimPrefix(trimmed, "impl")
rest = strings.TrimSpace(rest)
// Remove generics
if idx := strings.Index(rest, "<"); idx >= 0 {
depth := 1
for i := idx + 1; i < len(rest); i++ {
if rest[i] == '<' {
depth++
} else if rest[i] == '>' {
depth--
if depth == 0 {
rest = strings.TrimSpace(rest[:idx] + rest[i+1:])
break
}
}
}
}
// Handle "for Type"
if idx := strings.Index(rest, " for "); idx >= 0 {
rest = strings.TrimSpace(rest[idx+5:])
}
name := strings.Fields(rest)[0]
symbols = append(symbols, symbolInfo{
Name: name,
Kind: "impl",
StartLine: lineNum,
})
continue
}
// Function
if matches := rsFnRe.FindStringSubmatch(line); matches != nil {
symbols = append(symbols, symbolInfo{
Name: matches[1],
Kind: "function",
StartLine: lineNum,
})
}
}
sortSymbols(symbols)
return symbols, scanner.Err()
}
// ────────────────────────────────────────────────────────────────────────────
// Helpers
// ────────────────────────────────────────────────────────────────────────────
func sortSymbols(symbols []symbolInfo) {
sort.Slice(symbols, func(i, j int) bool {
return symbols[i].StartLine < symbols[j].StartLine
})
}
// groupSymbols groups symbols by kind and returns formatted strings.
func groupSymbols(symbols []symbolInfo) map[string][]string {
groups := make(map[string][]string)
for _, s := range symbols {
kind := s.Kind
switch kind {
case "package":
groups["package"] = append(groups["package"], s.Name)
case "import":
groups["imports"] = append(groups["imports"], s.Name)
case "function":
groups["functions"] = append(groups["functions"], s.Name)
case "method":
groups["methods"] = append(groups["methods"], s.Name)
case "class", "interface", "struct", "enum", "trait", "alias", "func":
// All structural types go under "types" for consistency.
groups["types"] = append(groups["types"], fmt.Sprintf("%s %s", s.Name, kind))
case "type":
// TS type alias
groups["types"] = append(groups["types"], fmt.Sprintf("%s type", s.Name))
case "impl":
groups["types"] = append(groups["types"], fmt.Sprintf("impl %s", s.Name))
}
}
return groups
}
+173
View File
@@ -0,0 +1,173 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"strings"
"sync"
"agentu/pkg/llm"
)
// ────────────────────────────────────────────────────────────────────────────
// TodoTool — in-memory task tracking
// ────────────────────────────────────────────────────────────────────────────
type todoItem struct {
Content string `json:"content"`
Status string `json:"status"` // "pending", "in_progress", "completed"
}
type TodoTool struct {
mu sync.Mutex
items []todoItem
}
func NewTodoTool() *TodoTool {
return &TodoTool{}
}
func (t *TodoTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "todo",
Description: "Manage a task list for the current conversation. Create, update, and track multi-step work. Use this to plan complex tasks and show progress.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"action": {"type": "string", "enum": ["list", "create", "update", "start", "complete", "clear"], "description": "Operation to perform."},
"items": {"type": "array", "items": {"type": "string"}, "description": "For 'create': list of task descriptions to add."},
"task": {"type": "string", "description": "For 'update'/'start'/'complete': the exact task content string to target."},
"status": {"type": "string", "enum": ["pending", "in_progress", "completed"], "description": "For 'update': new status for the task."}
},
"required": ["action"],
"additionalProperties": false
}`),
},
}
}
type todoArgs struct {
Action string `json:"action"`
Items []string `json:"items"`
Task string `json:"task"`
Status string `json:"status"`
}
func (t *TodoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args todoArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
t.mu.Lock()
defer t.mu.Unlock()
if err := ctx.Err(); err != nil {
return "", err
}
switch args.Action {
case "list":
return t.listTasks(), nil
case "create":
return t.createTasks(args.Items)
case "start":
return t.startTask(args.Task)
case "complete":
return t.completeTask(args.Task)
case "update":
return t.updateTask(args.Task, args.Status)
case "clear":
t.items = nil
return "Tasks cleared.", nil
default:
return "", fmt.Errorf("unknown todo action: %s", args.Action)
}
}
func (t *TodoTool) listTasks() string {
if len(t.items) == 0 {
return "No tasks."
}
var b strings.Builder
b.WriteString("Tasks:\n")
for _, item := range t.items {
fmt.Fprintf(&b, "[%s] %s\n", item.Status, item.Content)
}
return b.String()
}
func (t *TodoTool) createTasks(items []string) (string, error) {
if len(items) == 0 {
return "", fmt.Errorf("items is required for create")
}
for _, content := range items {
content = strings.TrimSpace(content)
if content == "" {
return "", fmt.Errorf("task content must not be empty")
}
t.items = append(t.items, todoItem{
Content: content,
Status: "pending",
})
}
return t.listTasks(), nil
}
func (t *TodoTool) startTask(task string) (string, error) {
idx := -1
for i, item := range t.items {
if item.Content == task {
idx = i
break
}
}
if idx < 0 {
return "", fmt.Errorf("task not found: %s", task)
}
// Set all in_progress items to pending, then mark the target as in_progress.
for i := range t.items {
if t.items[i].Status == "in_progress" {
t.items[i].Status = "pending"
}
}
t.items[idx].Status = "in_progress"
return t.listTasks(), nil
}
func (t *TodoTool) completeTask(task string) (string, error) {
idx := -1
for i, item := range t.items {
if item.Content == task {
idx = i
break
}
}
if idx < 0 {
return "", fmt.Errorf("task not found: %s", task)
}
t.items[idx].Status = "completed"
return t.listTasks(), nil
}
func (t *TodoTool) updateTask(task, status string) (string, error) {
if status != "pending" && status != "in_progress" && status != "completed" {
return "", fmt.Errorf("status must be one of: pending, in_progress, completed")
}
idx := -1
for i, item := range t.items {
if item.Content == task {
idx = i
break
}
}
if idx < 0 {
return "", fmt.Errorf("task not found: %s", task)
}
t.items[idx].Status = status
return t.listTasks(), nil
}
+78 -4
View File
@@ -19,6 +19,54 @@ type Tool interface {
Execute(ctx context.Context, args json.RawMessage) (string, error)
}
// StreamingTool is optionally implemented by tools that can report partial
// output while they execute. emit receives chunks of the tool's output as
// they are produced; the returned string is still the complete result that
// gets sent back to the model.
type StreamingTool interface {
Tool
ExecuteStream(ctx context.Context, args json.RawMessage, emit func(string) error) (string, error)
}
// MutatingTool is optionally implemented by tools that modify external state.
// It is used by the agent to determine whether tool calls can run in parallel.
type MutatingTool interface {
Mutates() bool
}
// CallMutatingTool can refine a tool's mutation classification per call.
// This is useful for tools such as apply_patch whose dry-run mode is read-only.
type CallMutatingTool interface {
MutatesCall(args json.RawMessage) bool
}
// PreviewableTool describes a pending mutation for interactive approval.
type PreviewableTool interface {
Preview(ctx context.Context, args json.RawMessage) (string, error)
}
// TargetPathTool reports files that should be snapshotted before a mutation.
type TargetPathTool interface {
TargetPaths(args json.RawMessage) ([]string, error)
}
// IsMutating reports whether tool implements MutatingTool and returns true.
func IsMutating(tool Tool) bool {
mutating, ok := tool.(MutatingTool)
return ok && mutating.Mutates()
}
// IsMutatingCall reports whether this invocation modifies external state.
func IsMutatingCall(tool Tool, args json.RawMessage) bool {
if !IsMutating(tool) {
return false
}
if call, ok := tool.(CallMutatingTool); ok {
return call.MutatesCall(args)
}
return true
}
type Registry struct {
tools map[string]Tool
primary []string
@@ -66,44 +114,69 @@ func ReadOnlyBuiltins(workingDir string) *Registry {
fileList := NewFileListTool(workingDir)
fileSearch := NewFileSearchTool(workingDir)
codeSymbols := NewCodeSymbolsTool(workingDir)
gitStatus := NewGitStatusTool(workingDir)
gitDiff := NewGitDiffTool(workingDir)
gitLog := NewGitLogTool(workingDir)
projectOutline := NewProjectOutlineTool(workingDir)
todo := NewTodoTool()
registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols)
registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols, projectOutline, todo, gitStatus, gitDiff, gitLog)
registry.RegisterAlias("file.read", fileRead)
registry.RegisterAlias("file.list", fileList)
registry.RegisterAlias("file.search", fileSearch)
registry.RegisterAlias("code.symbols", codeSymbols)
registry.RegisterAlias("git.status", gitStatus)
registry.RegisterAlias("git.diff", gitDiff)
registry.RegisterAlias("git.log", gitLog)
return registry
}
func Builtins(workingDir string) *Registry {
registry := ReadOnlyBuiltins(workingDir)
fileWrite := NewFileWriteTool(workingDir)
fileEdit := NewFileEditTool(workingDir)
applyPatch := NewApplyPatchTool(workingDir)
shellRun := NewShellRunTool(workingDir)
registry.Register(fileWrite)
registry.Register(fileEdit)
registry.Register(applyPatch)
registry.Register(shellRun)
registry.RegisterAlias("file.write", fileWrite)
registry.RegisterAlias("file.edit", fileEdit)
registry.RegisterAlias("apply.patch", applyPatch)
registry.RegisterAlias("shell.run", shellRun)
return registry
}
func AgentInstructions(workingDir string, yolo bool) string {
return AgentInstructionsForMode(workingDir, yolo, yolo)
}
func AgentInstructionsForMode(workingDir string, writesEnabled bool, yolo bool) string {
searchInstructions := webSearchInstructions()
readOnlyTools := "file_list, file_read, file_search, code_symbols, web_search, web_fetch"
yoloTools := "file_list, file_read, file_search, code_symbols, file_write, shell_run, 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"
writeTools := "file_list, file_read, file_search, code_symbols, project_outline, todo, file_write, file_edit, apply_patch, shell_run, git_status, git_diff, git_log, web_search, web_fetch"
if writesEnabled {
approval := "Mutating tools require interactive user approval before execution. Explain the intended change through the tool call and wait for the decision."
if yolo {
approval = "Mutating tools execute automatically because --yolo is enabled."
}
return fmt.Sprintf(`Local project context:
- The project working directory is %q.
- Tools are available, but use them only when needed to satisfy the user's request.
- %s
- Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project or command context.
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
- When the user asks to run, execute, test, check, debug, or inspect a shell command, use shell_run. This includes local commands written in backticks, such as go, make, git, pwd, ls, ps, and similar terminal commands.
- Use project_outline first when the user asks for a broad understanding of a project; prefer file_edit for localized changes, apply_patch for related multi-file changes, and file_write for whole-file creation or replacement.
- Do not use file tools as a substitute for a command execution request.
- When the user asks about changed files, diffs, or git history, use git_status, git_diff, or git_log instead of shell_run.
- Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions.
- %s
- Tool paths should normally be relative to the project working directory.
- Available tools: %s.`, workingDir, searchInstructions, yoloTools)
- Available tools: %s.`, workingDir, approval, searchInstructions, writeTools)
}
return fmt.Sprintf(`Local project context:
- The project working directory is %q.
@@ -113,6 +186,7 @@ func AgentInstructions(workingDir string, yolo bool) string {
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
- Shell command execution is disabled in this mode. When the user asks to run, execute, test, check, debug, or inspect a shell command, do not use file tools as a substitute; explain that they need to restart agentu with --yolo to enable shell_run.
- When the user asks about changed files, diffs, or git history, use git_status, git_diff, or git_log instead of shell_run.
- %s
- Tool paths should normally be relative to the project working directory.
- Available tools: %s.`, workingDir, searchInstructions, readOnlyTools)
+22 -6
View File
@@ -104,12 +104,12 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
}
}
for _, want := range []string{"file_list", "file_read", "file_search", "code_symbols"} {
for _, want := range []string{"file_list", "file_read", "file_search", "code_symbols", "project_outline", "todo", "git_status", "git_diff", "git_log"} {
if !slices.Contains(names, want) {
t.Fatalf("definitions missing %s: %#v", want, names)
}
}
for _, alias := range []string{"file.list", "file.read", "file.search", "code.symbols"} {
for _, alias := range []string{"file.list", "file.read", "file.search", "code.symbols", "git.status", "git.diff", "git.log"} {
if _, ok := registry.Get(alias); !ok {
t.Fatalf("alias missing: %s", alias)
}
@@ -117,6 +117,9 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
if _, ok := registry.Get("file_write"); ok {
t.Fatal("read-only registry should not expose file_write")
}
if _, ok := registry.Get("file_edit"); ok {
t.Fatal("read-only registry should not expose file_edit")
}
}
func TestCodeSymbolsToolListsGoSymbols(t *testing.T) {
@@ -170,7 +173,7 @@ func TestRegistryCanExposeWebSearch(t *testing.T) {
func TestBuiltinsExposeYoloTools(t *testing.T) {
registry := Builtins(t.TempDir())
for _, name := range []string{"file_write", "shell_run", "file.write", "shell.run"} {
for _, name := range []string{"file_write", "file_edit", "apply_patch", "shell_run", "file.write", "file.edit", "apply.patch", "shell.run"} {
if _, ok := registry.Get(name); !ok {
t.Fatalf("tool missing: %s", name)
}
@@ -179,14 +182,14 @@ func TestBuiltinsExposeYoloTools(t *testing.T) {
func TestAgentInstructionsExplainCommandModes(t *testing.T) {
readOnly := AgentInstructions("/tmp/project", false)
for _, want := range []string{"Shell command execution is disabled", "--yolo", "Do not call tools for greetings", "Do not explore the project preemptively", "Do not call file_list just to discover context"} {
for _, want := range []string{"Shell command execution is disabled", "--yolo", "Do not call tools for greetings", "Do not explore the project preemptively", "Do not call file_list just to discover context", "git_status, git_diff, git_log"} {
if !strings.Contains(readOnly, want) {
t.Fatalf("read-only instructions missing %q:\n%s", want, readOnly)
}
}
yolo := AgentInstructions("/tmp/project", true)
for _, want := range []string{"use shell_run", "commands written in backticks", "Do not start interactive", "non-interactive local commands", "Do not call tools for greetings"} {
for _, want := range []string{"use shell_run", "commands written in backticks", "Do not start interactive", "non-interactive local commands", "Do not call tools for greetings", "prefer file_edit"} {
if !strings.Contains(yolo, want) {
t.Fatalf("yolo instructions missing %q:\n%s", want, yolo)
}
@@ -195,7 +198,7 @@ func TestAgentInstructionsExplainCommandModes(t *testing.T) {
func TestAgentInstructionsExplainWebSearch(t *testing.T) {
withSearch := AgentInstructions("/tmp/project", false)
for _, want := range []string{"use web_search", "Use web_fetch", "latest information", "source URLs", "Available tools: file_list, file_read, file_search, code_symbols, 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) {
t.Fatalf("web search instructions missing %q:\n%s", want, withSearch)
}
@@ -257,3 +260,16 @@ func TestResolvePathRejectsDotDotThatEscapes(t *testing.T) {
t.Fatal("expected error for path escaping via ..")
}
}
func TestIsMutating(t *testing.T) {
dir := t.TempDir()
if IsMutating(NewFileReadTool(dir)) {
t.Fatal("FileReadTool should not be mutating")
}
if !IsMutating(NewFileWriteTool(dir)) {
t.Fatal("FileWriteTool should be mutating")
}
if !IsMutating(NewShellRunTool(dir)) {
t.Fatal("ShellRunTool should be mutating")
}
}
+200
View File
@@ -0,0 +1,200 @@
package tools
import (
"context"
"encoding/json"
"os"
"path/filepath"
"strings"
"testing"
)
func TestApplyPatchChangesMultipleFiles(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "old.txt"), []byte("before\n"), 0o600); err != nil {
t.Fatal(err)
}
patch := `*** Begin Patch
*** Update File: old.txt
@@
-before
+after
*** Add File: new.txt
+created
*** End Patch`
out, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch}))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "2 file(s) changed") {
t.Fatalf("output = %q", out)
}
assertFileContent(t, filepath.Join(dir, "old.txt"), "after\n")
assertFileContent(t, filepath.Join(dir, "new.txt"), "created\n")
if info, err := os.Stat(filepath.Join(dir, "old.txt")); err != nil || info.Mode().Perm() != 0o600 {
t.Fatalf("updated mode = %v, err = %v", info.Mode().Perm(), err)
}
}
func TestApplyPatchRollsBackEarlierFilesOnFailure(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "first.txt")
if err := os.WriteFile(path, []byte("original\n"), 0o644); err != nil {
t.Fatal(err)
}
patch := `*** Begin Patch
*** Update File: first.txt
@@
-original
+changed
*** Update File: missing.txt
@@
-missing
+changed
*** End Patch`
_, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch}))
if err == nil {
t.Fatal("expected patch failure")
}
assertFileContent(t, path, "original\n")
}
func TestApplyPatchDryRunDoesNotWrite(t *testing.T) {
dir := t.TempDir()
patch := "*** Begin Patch\n*** Add File: dry.txt\n+dry\n*** End Patch"
out, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch, "dry_run": true}))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "dry run") {
t.Fatalf("output = %q", out)
}
if _, err := os.Stat(filepath.Join(dir, "dry.txt")); !os.IsNotExist(err) {
t.Fatalf("dry run created file: %v", err)
}
}
func TestApplyPatchRejectsMalformedEnvelopeAndDuplicatePaths(t *testing.T) {
for _, patch := range []string{
"*** Begin Patch\n*** Add File: x.txt\n+x",
"prefix\n*** Begin Patch\n*** Add File: x.txt\n+x\n*** End Patch",
"*** Begin Patch\n*** Add File: x.txt\n+x\n*** Delete File: x.txt\n*** End Patch",
} {
if _, err := parsePatch(patch); err == nil {
t.Fatalf("expected parse failure for %q", patch)
}
}
}
func TestResolvePathRejectsSymlinkEscape(t *testing.T) {
parent := t.TempDir()
base := filepath.Join(parent, "project")
outside := filepath.Join(parent, "outside")
if err := os.MkdirAll(base, 0o755); err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(outside, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(outside, "secret.txt"), []byte("secret"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.Symlink(outside, filepath.Join(base, "link")); err != nil {
t.Skipf("symlinks unavailable: %v", err)
}
if _, err := resolvePath(base, "link/secret.txt"); err == nil {
t.Fatal("expected read symlink escape rejection")
}
if _, err := resolvePathForWrite(base, "link/new.txt"); err == nil {
t.Fatal("expected write symlink escape rejection")
}
}
func TestTodoRejectsInvalidData(t *testing.T) {
tool := NewTodoTool()
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"create","items":[""]}`)); err == nil {
t.Fatal("expected empty task rejection")
}
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"create","items":["one"]}`)); err != nil {
t.Fatal(err)
}
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"update","task":"one","status":"invalid"}`)); err == nil {
t.Fatal("expected invalid status rejection")
}
}
func TestProjectOutlineIncludesGoTypeScriptAndRustSymbols(t *testing.T) {
dir := t.TempDir()
files := map[string]string{
"main.go": "package demo\nfunc Run() {}\n",
"ui.ts": "export class Panel {}\nexport function render() {}\n",
"lib.rs": "pub struct Worker;\npub fn execute() {}\n",
}
for name, content := range files {
if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644); err != nil {
t.Fatal(err)
}
}
out, err := NewProjectOutlineTool(dir).Execute(context.Background(), json.RawMessage(`{}`))
if err != nil {
t.Fatal(err)
}
for _, want := range []string{"main.go", "Run", "ui.ts", "Panel", "render", "lib.rs", "Worker", "execute"} {
if !strings.Contains(out, want) {
t.Fatalf("outline missing %q:\n%s", want, out)
}
}
}
func TestRollbackStoreRestoresInReverseOrder(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "state.txt")
if err := os.WriteFile(path, []byte("one"), 0o644); err != nil {
t.Fatal(err)
}
store := NewRollbackStore()
first, err := store.Capture("edit", "state.txt", path)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("two"), 0o644); err != nil {
t.Fatal(err)
}
second, err := store.Capture("edit", "state.txt", path)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("three"), 0o644); err != nil {
t.Fatal(err)
}
if store.LastID() != second {
t.Fatalf("last ID = %q", store.LastID())
}
if _, err := store.Restore(second); err != nil {
t.Fatal(err)
}
assertFileContent(t, path, "two")
if store.LastID() != first {
t.Fatalf("last ID after restore = %q, want %q", store.LastID(), first)
}
}
func mustJSON(t *testing.T, value any) json.RawMessage {
t.Helper()
data, err := json.Marshal(value)
if err != nil {
t.Fatal(err)
}
return data
}
func assertFileContent(t *testing.T, path, want string) {
t.Helper()
data, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(data) != want {
t.Fatalf("%s content = %q, want %q", path, string(data), want)
}
}
+1120 -93
View File
File diff suppressed because it is too large Load Diff
+743 -30
View File
@@ -2,13 +2,17 @@ package tui
import (
"context"
"fmt"
"os"
"os/exec"
"path/filepath"
"reflect"
"strings"
"testing"
"time"
"agentu/internal/agent"
"agentu/internal/tools"
"agentu/pkg/llm"
tea "github.com/charmbracelet/bubbletea"
@@ -21,7 +25,7 @@ func TestModelRendersChatShell(t *testing.T) {
m = next.(model)
rendered := m.View()
for _, want := range []string{"test-model", "yolo tools", "light", "ctrl+j", "alt+enter"} {
for _, want := range []string{"test-model"} {
if !strings.Contains(rendered, want) {
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"})
for name, color := range map[string]any{
"input box": m.styles.InputBox.GetBackground(),
want := lightTheme().InputBg
for name, color := range map[string]lipgloss.TerminalColor{
"focused base": m.input.FocusedStyle.Base.GetBackground(),
"focused cursor line": m.input.FocusedStyle.CursorLine.GetBackground(),
"blurred base": m.input.BlurredStyle.Base.GetBackground(),
"blurred cursor line": m.input.BlurredStyle.CursorLine.GetBackground(),
} {
if _, ok := color.(lipgloss.NoColor); !ok {
t.Fatalf("%s should not paint input text background", name)
}
requireStyleColor(t, name, color, want)
}
}
@@ -58,12 +60,12 @@ func TestInputBoxIsRoomier(t *testing.T) {
m = next.(model)
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 {
t.Fatalf("input box horizontal frame should include roomier padding, got %d", got)
if got := m.styles.InputBox.GetHorizontalFrameSize(); got != 3 {
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)
}
rendered := m.inputView()
@@ -78,13 +80,8 @@ func TestModelMetaRendersBelowInput(t *testing.T) {
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
rendered := next.(model).View()
inputIndex := strings.Index(rendered, "Message agentu")
modelIndex := strings.Index(rendered, "test-model")
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)
if !strings.Contains(rendered, "test-model") {
t.Fatalf("rendered view missing %q:\n%s", "test-model", rendered)
}
}
@@ -93,7 +90,7 @@ func TestModelMetaRendersContextUsage(t *testing.T) {
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("x", 800)}})
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 120, Height: 30})
rendered := next.(model).modelMetaView()
rendered := next.(model).footerView()
for _, want := range []string{"ctx ~", "/1k", "%"} {
if !strings.Contains(rendered, want) {
@@ -123,11 +120,124 @@ func TestRunningStateRendersActivityIndicator(t *testing.T) {
m.resize()
rendered := m.View()
for _, want := range []string{"Working", "thinking...", "esc cancels"} {
for _, want := range []string{iconWorking, "Working", "thinking...", "esc cancels"} {
if !strings.Contains(rendered, want) {
t.Fatalf("running view missing %q:\n%s", want, rendered)
}
}
if backgroundPattern.MatchString(m.activityView()) {
t.Fatalf("activity view should not paint a background:\n%q", m.activityView())
}
}
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) {
@@ -276,19 +386,219 @@ func TestInputStartsAtOneLineAndGrows(t *testing.T) {
}
}
func TestSlashCommandSuggestions(t *testing.T) {
func TestCommandPaletteFuzzyMatches(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/")
m.input.SetValue("/mdl")
m.afterInputChanged()
rendered := m.View()
for _, want := range []string{"/clear reset context", "/compact compress context", "/exit quit", "/model model/provider/thinking"} {
if !strings.Contains(rendered, want) {
t.Fatalf("rendered suggestions missing %q:\n%s", want, rendered)
plain := stripANSI(m.View())
for _, want := range []string{iconCommand + " Commands", "/model model/provider/thinking"} {
if !strings.Contains(plain, want) {
t.Fatalf("rendered palette missing %q:\n%s", want, plain)
}
}
matches := commandMatches("m")
modelIndex := -1
renameIndex := -1
for i, match := range matches {
switch match.command.Name {
case "/model":
modelIndex = i
case "/rename":
renameIndex = i
}
}
if modelIndex < 0 || renameIndex < 0 || modelIndex > renameIndex {
t.Fatalf("/model should sort before /rename for /m query: %#v", matches)
}
}
func TestInputHistoryNavigation(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.rememberInput("first")
m.rememberInput("second")
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlP})
m = next.(model)
if got := m.input.Value(); got != "second" {
t.Fatalf("after ctrl+p input = %q", got)
}
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlP})
m = next.(model)
if got := m.input.Value(); got != "first" {
t.Fatalf("after second ctrl+p input = %q", got)
}
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlN})
m = next.(model)
if got := m.input.Value(); got != "second" {
t.Fatalf("after ctrl+n input = %q", got)
}
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlN})
m = next.(model)
if got := m.input.Value(); got != "" {
t.Fatalf("after second ctrl+n input = %q", got)
}
}
func 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) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`})
plain := stripANSI(m.renderMessages())
for _, want := range []string{iconTool + " file_read", "README.md"} {
if !strings.Contains(plain, want) {
t.Fatalf("tool card missing %q:\n%s", want, plain)
}
}
}
func TestToolMessageRendersOutputPreview(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleTool, content: "file_read {\"path\":\"README.md\"}\n[output]\nfirst line\nsecond line"})
plain := stripANSI(m.renderMessages())
for _, want := range []string{iconTool + " file_read", "README.md", "output", "first line", "second line"} {
if !strings.Contains(plain, want) {
t.Fatalf("tool card missing %q:\n%s", want, plain)
}
}
}
func TestToolLogOutputUpdatesExistingToolCard(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`})
updated := m.mergeToolLog("file_read {\"path\":\"README.md\"}\n[output]\nhello")
if !updated {
t.Fatal("expected 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, "hello") {
t.Fatalf("rendered output missing 'hello':\n%s", plain)
}
status := toolStatusText("file_read {\"path\":\"README.md\"}\n[output]\nhello world")
if strings.Contains(status, "hello") {
t.Fatalf("status should not contain output: %q", status)
}
if !strings.Contains(status, "file_read") {
t.Fatalf("status missing tool name: %q", status)
}
}
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{}
@@ -386,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) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
m.input.SetValue("/model")
next, _ := m.submit()
rendered := next.(model).renderMessages()
for _, want := range []string{"model: test-model", "mode: yolo tools", "theme: light"} {
if !strings.Contains(rendered, want) {
t.Fatalf("model command output missing %q:\n%s", want, rendered)
}
if !strings.Contains(rendered, "test-model") {
t.Fatalf("model command output missing %q:\n%s", "test-model", rendered)
}
}
@@ -414,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 {
provider string
model string
thinking string
sessionName string
models []string
saved bool
}
@@ -489,6 +956,16 @@ func (f *fakeModelManager) NewSession() (string, error) {
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) {
for input, want := range map[string]ThemeMode{
"": ThemeLight,
@@ -510,6 +987,130 @@ func TestParseThemeMode(t *testing.T) {
}
}
func requireStyleColor(t *testing.T, name string, got lipgloss.TerminalColor, want string) {
t.Helper()
color, ok := got.(lipgloss.Color)
if !ok {
t.Fatalf("%s = %T, want lipgloss.Color", name, got)
}
if string(color) != want {
t.Fatalf("%s = %q, want %q", name, string(color), want)
}
}
func requireNoStyleColor(t *testing.T, name string, got lipgloss.TerminalColor) {
t.Helper()
if _, ok := got.(lipgloss.NoColor); !ok {
t.Fatalf("%s should not paint a color, got %T", name, got)
}
}
func TestDarkThemePaintsAppBackground(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
rendered := stripANSI(m.View())
_ = rendered
if m.theme.Mode != ThemeDark {
t.Fatalf("theme.Mode = %q, want %q", m.theme.Mode, ThemeDark)
}
want := darkTheme()
requireStyleColor(t, "App.Background", m.styles.App.GetBackground(), want.Background)
requireStyleColor(t, "App.Foreground", m.styles.App.GetForeground(), want.Text)
}
func TestDarkThemeKeepsComponentBackgroundsLightweight(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
requireNoStyleColor(t, "Viewport.Background", m.styles.Viewport.GetBackground())
requireStyleColor(t, "InputBox.Background", m.styles.InputBox.GetBackground(), darkTheme().InputBg)
requireNoStyleColor(t, "Activity.Background", m.styles.Activity.GetBackground())
m.running = true
m.status = "thinking..."
m.resize()
if backgroundPattern.MatchString(m.activityView()) {
t.Fatalf("activity view should not paint a background:\n%q", m.activityView())
}
}
func TestDarkThemeStylesUseDarkPalette(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
want := darkTheme()
requireStyleColor(t, "Header.Foreground", m.styles.Header.GetForeground(), want.Title)
requireStyleColor(t, "Muted.Foreground", m.styles.Muted.GetForeground(), want.Muted)
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.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 TestInputModeSymbolReflectsYoloAndBangCommand(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)
rendered := stripANSI(m.inputView())
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) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
if m.theme.Mode != ThemeLight {
t.Fatalf("initial theme = %q, want light", m.theme.Mode)
}
requireNoStyleColor(t, "initial App.Background", m.styles.App.GetBackground())
// Switch to dark
m.input.SetValue("/theme dark")
next, _ := m.submit()
m = next.(model)
if m.theme.Mode != ThemeDark {
t.Fatalf("after /theme dark: Mode = %q, want dark", m.theme.Mode)
}
requireStyleColor(t, "dark App.Background", m.styles.App.GetBackground(), darkTheme().Background)
last := m.messages[len(m.messages)-1]
if !strings.Contains(last.content, "theme: dark") {
t.Fatalf("expected confirmation message, got %q", last.content)
}
// Switch back to light
m.input.SetValue("/theme light")
next, _ = m.submit()
m = next.(model)
if m.theme.Mode != ThemeLight {
t.Fatalf("after /theme light: Mode = %q, want light", m.theme.Mode)
}
requireNoStyleColor(t, "restored App.Background", m.styles.App.GetBackground())
// No arg shows usage
m.input.SetValue("/theme")
next, _ = m.submit()
m = next.(model)
last = m.messages[len(m.messages)-1]
if !strings.Contains(last.content, "usage: /theme") {
t.Fatalf("expected usage hint, got %q", last.content)
}
}
func TestCleanToolLog(t *testing.T) {
got := cleanToolLog("\n[tool] shell_run {\"command\":\"pwd\"}\n")
want := `shell_run {"command":"pwd"}`
@@ -517,3 +1118,115 @@ func TestCleanToolLog(t *testing.T) {
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"))
}
+110 -19
View File
@@ -15,6 +15,17 @@ const (
ThemeDark ThemeMode = "dark"
)
const (
iconReady = "✓"
iconWorking = "◆"
iconTool = "▣"
iconError = "✕"
iconCommand = "⌘"
iconSession = "◌"
iconContext = "◔"
iconModel = "◉"
)
func ParseThemeMode(value string) (ThemeMode, error) {
switch ThemeMode(strings.ToLower(strings.TrimSpace(value))) {
case "", ThemeLight:
@@ -32,6 +43,7 @@ type theme struct {
Background string
Surface string
SurfaceAlt string
InputBg string
Border string
Text string
Muted string
@@ -48,11 +60,25 @@ type styles struct {
App lipgloss.Style
Muted lipgloss.Style
Activity lipgloss.Style
ModelMeta lipgloss.Style
Header lipgloss.Style
Footer lipgloss.Style
Viewport lipgloss.Style
InputBox lipgloss.Style
CommandHint lipgloss.Style
InputModeReadonly lipgloss.Style
InputModeYolo lipgloss.Style
InputModeCommand lipgloss.Style
StatusReady lipgloss.Style
StatusWarn lipgloss.Style
StatusError lipgloss.Style
StatusInfo lipgloss.Style
Pill lipgloss.Style
PillAccent lipgloss.Style
Palette lipgloss.Style
PaletteMatch lipgloss.Style
ToolName lipgloss.Style
ToolKey lipgloss.Style
ToolValue lipgloss.Style
UserLabel lipgloss.Style
ToolLabel lipgloss.Style
ErrorLabel lipgloss.Style
@@ -75,8 +101,8 @@ func lightTheme() theme {
return theme{
Mode: ThemeLight,
Background: "#F8FAFC",
Surface: "#FFFFFF",
SurfaceAlt: "#EEF2F7",
InputBg: "#E4E4E4",
Border: "#CBD5E1",
Text: "#111827",
Muted: "#64748B",
@@ -93,8 +119,8 @@ func darkTheme() theme {
return theme{
Mode: ThemeDark,
Background: "#111827",
Surface: "#1F2937",
SurfaceAlt: "#0F172A",
InputBg: "#383838",
Border: "#475569",
Text: "#E5E7EB",
Muted: "#94A3B8",
@@ -107,24 +133,28 @@ func darkTheme() theme {
}
}
func appStyle(t theme) lipgloss.Style {
style := lipgloss.NewStyle().Foreground(lipgloss.Color(t.Text))
if t.Mode == ThemeDark {
style = style.Background(lipgloss.Color(t.Background))
}
return style
}
func (t theme) styles() styles {
return styles{
App: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)),
App: appStyle(t),
Muted: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)),
Activity: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Background)).
Background(lipgloss.Color(t.Title)).
Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1),
ModelMeta: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.SurfaceAlt)).
Padding(0, 1),
Header: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Title)),
Footer: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)).
@@ -134,15 +164,74 @@ func (t theme) styles() styles {
Foreground(lipgloss.Color(t.Text)),
InputBox: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.InputBg)).
Padding(1, 2, 1, 1),
CommandHint: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1),
InputModeReadonly: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Muted)).
Background(lipgloss.Color(t.InputBg)).
Padding(0, 1),
InputModeYolo: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Error)).
Background(lipgloss.Color(t.InputBg)).
Padding(0, 1),
InputModeCommand: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.User)).
Background(lipgloss.Color(t.InputBg)).
Padding(0, 1),
StatusReady: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Assistant)),
StatusWarn: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Tool)),
StatusError: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Error)),
StatusInfo: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.User)),
Pill: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.SurfaceAlt)).
Padding(0, 1),
PillAccent: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Title)).
Background(lipgloss.Color(t.SurfaceAlt)).
Padding(0, 1),
Palette: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Border(lipgloss.ThickBorder()).
BorderForeground(lipgloss.Color(t.Border)).
Padding(0, 2),
CommandHint: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)).
Background(lipgloss.Color(t.SurfaceAlt)).
Padding(0, 1),
PaletteMatch: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.User)),
ToolName: lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(t.Tool)),
ToolKey: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)),
ToolValue: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)),
UserLabel: labelStyle(t.User),
ToolLabel: labelStyle(t.Tool),
@@ -185,8 +274,10 @@ func bubbleMessageStyle(textColor, backgroundColor string) lipgloss.Style {
func applyTextareaTheme(input *textarea.Model, t theme) {
focused := textarea.Style{
Base: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)),
CursorLine: lipgloss.NewStyle(),
Foreground(lipgloss.Color(t.Text)).
Background(lipgloss.Color(t.InputBg)),
CursorLine: lipgloss.NewStyle().
Background(lipgloss.Color(t.InputBg)),
Placeholder: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)),
Prompt: lipgloss.NewStyle().
@@ -197,7 +288,7 @@ func applyTextareaTheme(input *textarea.Model, t theme) {
Foreground(lipgloss.Color(t.Surface)),
}
blurred := focused
blurred.CursorLine = lipgloss.NewStyle()
blurred.CursorLine = lipgloss.NewStyle().Background(lipgloss.Color(t.InputBg))
input.FocusedStyle = focused
input.BlurredStyle = blurred
+2 -209
View File
@@ -1,222 +1,15 @@
package main
import (
"bufio"
"context"
"errors"
"flag"
"fmt"
"net/http"
"os"
"os/signal"
"path/filepath"
"strings"
"agentu/internal/agent"
"agentu/pkg/config"
"agentu/pkg/llm"
"agentu/internal/session"
"agentu/internal/tools"
"agentu/internal/tui"
"agentu/internal/app"
)
const defaultConfigContent = `# agentu configuration
# Replace the placeholder values below with your actual provider settings.
# See agentu.example.yaml for the full reference.
providers:
default:
base_url: https://api.openai.com
api_key: YOUR_API_KEY_HERE
models:
- id: gpt-4o
context: 128000
- id: gpt-4o-mini
context: 128000
`
func main() {
if err := run(); err != nil {
if err := app.Run(); err != nil {
fmt.Fprintln(os.Stderr, "agentu:", err)
os.Exit(1)
}
}
func run() error {
configPath := flag.String("config", config.DefaultPath, "path to YAML config")
yolo := flag.Bool("yolo", false, "also enable automatic file writes and shell execution")
plain := flag.Bool("plain", false, "use the simple line-based REPL instead of the TUI")
themeName := flag.String("theme", "light", "TUI theme: light or dark")
resumeID := flag.String("resume", "", "resume a specific session by ID")
flag.Parse()
themeMode, err := tui.ParseThemeMode(*themeName)
if err != nil {
return err
}
cfg, err := config.Load(*configPath)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
if bootstrapErr := bootstrapConfig(*configPath); bootstrapErr != nil {
return bootstrapErr
}
return fmt.Errorf("created default config at %s; edit it to set your API key and provider, then restart agentu", *configPath)
}
return err
}
providerName := cfg.DefaultProviderName()
providerConfig := cfg.Providers[providerName]
provider := llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, http.DefaultClient)
var registry *tools.Registry
if *yolo {
registry = tools.Builtins(cfg.Agent.WorkingDir)
} else {
registry = tools.ReadOnlyBuiltins(cfg.Agent.WorkingDir)
}
registry.Register(tools.NewBuiltinSearchTool(http.DefaultClient))
registry.Register(tools.NewBuiltinFetchTool(http.DefaultClient))
systemPrompt := strings.TrimSpace(cfg.Agent.SystemPrompt + "\n\n" + tools.AgentInstructions(cfg.Agent.WorkingDir, *yolo))
maxContextTokens := cfg.Agent.MaxContextTokens
if maxContextTokens == 0 {
maxContextTokens = providerConfig.ContextTokens
}
assistant := agent.New(agent.Options{
Provider: provider,
ProviderName: providerName,
Model: providerConfig.Model,
Thinking: providerConfig.Thinking,
ThinkingParam: providerConfig.ThinkingParam,
ThinkingEnabled: providerConfig.ThinkingConfigured,
SystemPrompt: systemPrompt,
ToolRegistry: registry,
ToolTimeout: cfg.Agent.ToolTimeout,
MaxContextTokens: maxContextTokens,
})
// Initialize session store
home, err := os.UserHomeDir()
if err != nil {
return fmt.Errorf("resolve home dir: %w", err)
}
sessionsDir := filepath.Join(home, ".agentu", "sessions")
store, err := session.NewStore(sessionsDir)
if err != nil {
return fmt.Errorf("init session store: %w", err)
}
modelManager, err := session.NewManager(cfg, assistant, http.DefaultClient, store)
if err != nil {
return err
}
// Resume session
if *resumeID != "" {
if err := modelManager.Resume(*resumeID); err != nil {
return fmt.Errorf("resume session: %w", err)
}
} else {
if err := modelManager.ResumeLatest(); err != nil {
return fmt.Errorf("resume latest session: %w", err)
}
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
defer stop()
if !*plain {
err := tui.Run(ctx, assistant, tui.Options{
ModelName: providerConfig.Model,
Yolo: *yolo,
ThemeMode: themeMode,
ModelManager: modelManager,
WorkingDir: cfg.Agent.WorkingDir,
})
printExitSession(modelManager)
return err
}
return repl(ctx, assistant, modelManager)
}
func bootstrapConfig(configPath string) error {
resolved, err := config.ResolvePath(configPath)
if err != nil {
return fmt.Errorf("resolve config path: %w", err)
}
dir := filepath.Dir(resolved)
if err := os.MkdirAll(dir, 0o755); err != nil {
return fmt.Errorf("create config directory: %w", err)
}
if err := os.WriteFile(resolved, []byte(defaultConfigContent), 0o644); err != nil {
return fmt.Errorf("write default config: %w", err)
}
return nil
}
func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error {
scanner := bufio.NewScanner(os.Stdin)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
fmt.Println("agentu")
fmt.Println("Type /exit to quit, /clear to reset context, /compact to compress context.")
for {
fmt.Print("> ")
if !scanner.Scan() {
if err := scanner.Err(); err != nil {
return err
}
fmt.Println()
printExitSession(modelManager)
return nil
}
input := strings.TrimSpace(scanner.Text())
switch input {
case "":
continue
case "/exit", "/quit":
printExitSession(modelManager)
return nil
case "/clear":
assistant.Clear()
fmt.Println("context cleared")
continue
case "/compact":
if err := assistant.Compact(ctx); err != nil {
fmt.Fprintln(os.Stderr, "compact error:", err)
} else {
fmt.Println("context compacted")
if modelManager != nil {
_ = modelManager.Save()
}
}
continue
}
if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
fmt.Fprintln(os.Stderr, "\nerror:", err)
}
_ = modelManager.Save()
fmt.Println()
}
}
func printExitSession(m *session.Manager) {
if m == nil {
return
}
_ = m.Save()
id := m.CurrentSessionID()
name := m.CurrentSessionName()
if id != "" {
if name != "" && name != id {
fmt.Fprintf(os.Stderr, "session: %s (%s)\n", id, name)
} else {
fmt.Fprintf(os.Stderr, "session: %s\n", id)
}
fmt.Fprintf(os.Stderr, "resume: agentu --resume %s\n", id)
}
}
+47 -1
View File
@@ -14,6 +14,11 @@ import (
const DefaultPath = "~/.agentu/config.yaml"
const (
APITypeChat = "chat"
APITypeResponses = "responses"
)
type Config struct {
Providers map[string]ProviderConfig `yaml:"providers"`
Agent AgentConfig `yaml:"agent"`
@@ -21,6 +26,7 @@ type Config struct {
type ProviderConfig struct {
Name string
APIType string
BaseURL string
Model string
Models []string
@@ -56,9 +62,11 @@ type rawConfig struct {
}
type rawProviderConfig struct {
APIType string `yaml:"api_type"`
BaseURL string `yaml:"base_url"`
Models []rawModelConfig `yaml:"models"`
APIKey string `yaml:"api_key"`
ThinkingParam string `yaml:"thinking_param"`
}
type rawModelConfig struct {
@@ -153,6 +161,9 @@ func (c *Config) Validate() error {
for name, provider := range c.Providers {
var missing []string
prefix := "providers." + name
if !IsAPIType(provider.APIType) {
return fmt.Errorf("%s.api_type must be one of: %s", prefix, strings.Join(APITypes(), ", "))
}
if strings.TrimSpace(provider.BaseURL) == "" {
missing = append(missing, prefix+".base_url")
}
@@ -220,21 +231,56 @@ func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
models, modelConfigs := normalizeModelConfigs(raw.Models)
provider := ProviderConfig{
Name: name,
APIType: normalizeAPIType(raw.APIType),
BaseURL: raw.BaseURL,
Models: models,
ModelConfigs: modelConfigs,
APIKey: raw.APIKey,
ThinkingParam: strings.TrimSpace(raw.ThinkingParam),
}
if len(provider.Models) > 0 {
provider.Model = provider.Models[0]
}
applyModelDefaults(&provider)
if provider.ThinkingConfigured && provider.ThinkingParam == "" {
provider.ThinkingParam = "thinking"
provider.ThinkingParam = DefaultThinkingParam(provider.APIType)
}
return provider
}
func normalizeAPIType(value string) string {
value = strings.ToLower(strings.TrimSpace(value))
if value == "" {
return APITypeChat
}
return value
}
func APITypes() []string {
return []string{APITypeChat, APITypeResponses}
}
func IsAPIType(value string) bool {
value = strings.ToLower(strings.TrimSpace(value))
for _, apiType := range APITypes() {
if value == apiType {
return true
}
}
return false
}
// DefaultThinkingParam returns the default request field used for the
// thinking level. Chat-completions providers historically use a top-level
// "thinking" field; the OpenAI Responses API uses "reasoning" with an
// {"effort": ...} value.
func DefaultThinkingParam(apiType string) string {
if apiType == APITypeResponses {
return "reasoning"
}
return "thinking"
}
func normalizeModelConfigs(rawModels []rawModelConfig) ([]string, map[string]ModelConfig) {
models := make([]string, 0, len(rawModels))
configs := make(map[string]ModelConfig, len(rawModels))
+132
View File
@@ -370,3 +370,135 @@ providers:
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())
}
}