7 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
43 changed files with 7711 additions and 563 deletions
+1
View File
@@ -3,3 +3,4 @@ agentu*.yaml
/coverage.out
*.swp
chat-history.txt
dist/
+85 -14
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,26 +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` for creating or replacing whole files
- `file_edit` for surgical line-range or pattern edits to existing files
- `shell_run` for non-interactive shell commands
- `apply_patch` for atomic-style multi-file V4 patches, including dry runs
- `shell_run` for non-interactive shell commands; output streams live into the
TUI while the command runs instead of appearing only after completion
Before a file mutation, agentu stores an in-memory snapshot and includes its
`rb_...` ID in the tool result. Use `/rollback` or `/rollback <id>` to restore
it. Snapshots last for the lifetime of the process. The plain REPL remains
read-only unless `--yolo` is supplied because it has no interactive approval UI.
Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
example:
@@ -147,7 +200,20 @@ 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
@@ -165,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:
+364 -40
View File
@@ -17,12 +17,15 @@ import (
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" +
@@ -39,6 +42,7 @@ type toolExecutionResult struct {
}
type Agent struct {
mu sync.RWMutex
provider llm.Provider
providerName string
model string
@@ -47,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
@@ -63,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
@@ -86,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,
@@ -114,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
}
@@ -130,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,
@@ -141,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) != "" {
@@ -151,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
@@ -160,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
}
@@ -174,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
@@ -182,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)
}
@@ -218,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
@@ -233,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
}
@@ -242,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)},
@@ -253,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)
}
@@ -301,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 {
@@ -323,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,
@@ -365,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,
})
@@ -375,7 +573,7 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
results := a.executeTools(ctx, calls, logs)
for _, result := range results {
a.messages = append(a.messages, llm.Message{
a.appendMessage(llm.Message{
Role: llm.RoleTool,
ToolCallID: result.call.ID,
Content: result.output,
@@ -391,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)
}
@@ -443,33 +651,145 @@ func (a *Agent) toolDefinitions() []llm.Tool {
return a.toolRegistry.Definitions()
}
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall) string {
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) (string, bool) {
name := call.Function.Name
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
}
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 {
@@ -477,7 +797,11 @@ func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
}
for _, call := range calls {
if tool, ok := a.toolRegistry.Get(call.Function.Name); ok {
if tools.IsMutating(tool) {
args := strings.TrimSpace(call.Function.Arguments)
if args == "" {
args = "{}"
}
if tools.IsMutatingCall(tool, json.RawMessage(args)) {
return true
}
}
@@ -498,9 +822,9 @@ func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.
if a.hasMutatingCall(calls) {
// Sequential: avoid racing mutating tools.
for i, call := range calls {
output := a.executeTool(ctx, call)
output, streamed := a.executeTool(ctx, call, logs)
results[i] = toolExecutionResult{call: call, output: output}
if logs != nil {
if logs != nil && !streamed {
logToolOutput(logs, call, output)
}
}
@@ -514,9 +838,9 @@ func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.
wg.Add(1)
go func(idx int, c llm.ToolCall) {
defer wg.Done()
output := a.executeTool(ctx, c)
output, streamed := a.executeTool(ctx, c, logs)
results[idx] = toolExecutionResult{call: c, output: output}
if logs != nil {
if logs != nil && !streamed {
mu.Lock()
logToolOutput(logs, c, output)
mu.Unlock()
+411
View File
@@ -3,8 +3,13 @@ package agent
import (
"context"
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"reflect"
"strings"
"sync"
"testing"
"agentu/internal/tools"
@@ -87,6 +92,86 @@ func TestAgentRunsToolLoop(t *testing.T) {
}
}
type streamToolProvider struct {
calls int
t *testing.T
}
func (p *streamToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
p.calls++
switch p.calls {
case 1:
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
{Index: 0, ID: "call_1", Type: "function", Name: "test_stream", Arguments: `{}`},
}})
case 2:
last := req.Messages[len(req.Messages)-1]
if last.Role != llm.RoleTool || last.Content != "full result" {
p.t.Fatalf("last message = %#v", last)
}
return emit(llm.StreamEvent{Content: "done"})
default:
p.t.Fatalf("unexpected call count %d", p.calls)
return nil
}
}
type streamEchoTool struct{}
func (streamEchoTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "test_stream",
Description: "streaming test",
Parameters: json.RawMessage(`{"type":"object"}`),
},
}
}
func (streamEchoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
return "full result", nil
}
func (streamEchoTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
if err := emit("one\n"); err != nil {
return "", err
}
if err := emit("two\n"); err != nil {
return "", err
}
return "full result", nil
}
func TestAgentStreamsToolOutputToLogs(t *testing.T) {
provider := &streamToolProvider{t: t}
a := New(Options{
Provider: provider,
Model: "test",
SystemPrompt: "system",
ToolRegistry: tools.NewRegistry(streamEchoTool{}),
})
var out strings.Builder
var logs strings.Builder
if err := a.RunTurn(context.Background(), "stream", &out, &logs); err != nil {
t.Fatal(err)
}
if out.String() != "done" {
t.Fatalf("out = %q", out.String())
}
logText := logs.String()
if strings.Count(logText, "[output]") != 2 {
t.Fatalf("logs = %q", logText)
}
if !strings.Contains(logText, "one\n") || !strings.Contains(logText, "two\n") {
t.Fatalf("logs missing streamed chunks: %q", logText)
}
if strings.Contains(logText, "[output]\nfull result") {
t.Fatalf("logs contain duplicate full output: %q", logText)
}
}
type contentProvider struct {
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 {
@@ -504,6 +801,120 @@ type parallelProvider struct {
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) {
+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
}
+13 -1
View File
@@ -1,12 +1,14 @@
package main
package app
import (
"context"
"path/filepath"
"strings"
"testing"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/internal/skills"
"agentu/pkg/config"
"agentu/pkg/llm"
)
@@ -99,3 +101,13 @@ func TestInitializeSessionResumesExplicitID(t *testing.T) {
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...)
}
+44 -16
View File
@@ -13,6 +13,7 @@ import (
"sort"
"strings"
"agentu/internal/safeio"
"agentu/pkg/llm"
)
@@ -83,6 +84,33 @@ func NewFileWriteTool(workingDir string) *FileWriteTool {
}
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{
Type: "function",
@@ -117,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"
+32 -1
View File
@@ -8,6 +8,7 @@ import (
"os"
"strings"
"agentu/internal/safeio"
"agentu/pkg/llm"
)
@@ -21,6 +22,36 @@ func NewFileEditTool(workingDir string) *FileEditTool {
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",
@@ -111,7 +142,7 @@ func (t *FileEditTool) Execute(ctx context.Context, raw json.RawMessage) (string
if err := ctx.Err(); err != nil {
return "", err
}
if err := os.WriteFile(path, []byte(result), perm); err != nil {
if err := safeio.AtomicWriteFile(path, []byte(result), perm); err != nil {
return "", err
}
+244
View File
@@ -0,0 +1,244 @@
package tools
import (
"context"
"encoding/json"
"errors"
"fmt"
"io/fs"
"os"
"path/filepath"
"strings"
"agentu/pkg/llm"
)
const defaultOutlineMaxFiles = 200
// ────────────────────────────────────────────────────────────────────────────
// ProjectOutlineTool — gitignore-aware file tree + symbol outline
// ────────────────────────────────────────────────────────────────────────────
type ProjectOutlineTool struct {
workingDir string
}
func NewProjectOutlineTool(workingDir string) *ProjectOutlineTool {
return &ProjectOutlineTool{workingDir: workingDir}
}
func (t *ProjectOutlineTool) Definition() llm.Tool {
return llm.Tool{
Type: "function",
Function: llm.ToolFunction{
Name: "project_outline",
Description: "Get the project file tree and code symbol outline in one call. Respects .gitignore. Use this first to understand project structure before reading individual files.",
Parameters: JSONSchema(`{
"type": "object",
"properties": {
"path": {"type": "string", "description": "Subdirectory to outline. Defaults to the project root."},
"max_files": {"type": "integer", "description": "Maximum files to include in the outline. Defaults to 200."}
},
"additionalProperties": false
}`),
},
}
}
type outlineArgs struct {
Path string `json:"path"`
MaxFiles int `json:"max_files"`
}
func (t *ProjectOutlineTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
var args outlineArgs
if err := json.Unmarshal(raw, &args); err != nil {
return "", fmt.Errorf("parse arguments: %w", err)
}
root, err := resolvePath(t.workingDir, args.Path)
if err != nil {
return "", err
}
limit := args.MaxFiles
if limit <= 0 {
limit = defaultOutlineMaxFiles
}
// Parse .gitignore files from the root upward.
ignorer := newGitignoreMatcher(root)
var out strings.Builder
out.WriteString(root)
out.WriteString("\n")
count := 0
truncated := false
err = filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
if err != nil {
return nil
}
if ctx.Err() != nil {
return ctx.Err()
}
// Skip always-ignored directories.
if d.IsDir() {
name := d.Name()
// Always skip VCS and common build-artifact directories.
switch name {
case ".git", ".hg", ".svn", "vendor", "node_modules",
"dist", "build", "__pycache__", "target":
return filepath.SkipDir
}
// Skip hidden directories other than .github and similar.
if strings.HasPrefix(name, ".") && name != ".github" && name != ".vscode" && name != ".devcontainer" {
return filepath.SkipDir
}
}
rel, _ := filepath.Rel(root, path)
if rel == "." {
return nil
}
// Check gitignore.
if ignorer.shouldIgnore(rel, d.IsDir()) {
if d.IsDir() {
return filepath.SkipDir
}
return nil
}
if count >= limit {
truncated = true
return filepath.SkipAll
}
// Build indented tree entry.
depth := strings.Count(rel, string(filepath.Separator))
indent := strings.Repeat(" ", depth)
name := d.Name()
if d.IsDir() {
out.WriteString(fmt.Sprintf("%s%s/\n", indent, name))
return nil
}
count++
out.WriteString(indent + name)
// Extract symbols for source files.
ext := filepath.Ext(name)
if sourceFileExtensions[ext] {
data, readErr := os.ReadFile(path)
if readErr == nil && len(data) <= 256*1024 {
symbols, symErr := extractSymbols(path, data)
if symErr == nil && len(symbols) > 0 {
out.WriteString(" — ")
out.WriteString(compactSymbolSummary(symbols))
}
}
}
out.WriteString("\n")
return nil
})
if err != nil && !errors.Is(err, filepath.SkipAll) {
return "", err
}
if truncated {
out.WriteString(fmt.Sprintf("\n... (%d files shown, more available; use file_list or code_symbols for details)\n", count))
}
return FormatResult(out.String()), nil
}
// ────────────────────────────────────────────────────────────────────────────
// Compact symbol summary for inline display in the file tree
// ────────────────────────────────────────────────────────────────────────────
func compactSymbolSummary(symbols []symbolInfo) string {
var parts []string
// Group by kind.
funcs := []string{}
methods := []string{}
types := []string{}
for _, s := range symbols {
switch s.Kind {
case "function":
funcs = append(funcs, s.Name)
case "method":
methods = append(methods, s.Name)
case "class", "interface", "struct", "enum", "trait", "alias", "func", "type", "impl":
types = append(types, s.Name)
}
// Skip imports and package for compact view.
}
if len(funcs) > 0 {
parts = append(parts, fmt.Sprintf("fn: %s", strings.Join(funcs, ", ")))
}
if len(methods) > 0 {
parts = append(parts, fmt.Sprintf("method: %s", strings.Join(methods, ", ")))
}
if len(types) > 0 {
parts = append(parts, fmt.Sprintf("type: %s", strings.Join(types, ", ")))
}
return strings.Join(parts, "; ")
}
// ────────────────────────────────────────────────────────────────────────────
// Simple .gitignore matcher
// ────────────────────────────────────────────────────────────────────────────
type gitignoreMatcher struct {
patterns []string
}
func newGitignoreMatcher(root string) *gitignoreMatcher {
m := &gitignoreMatcher{}
// Walk up from root, collecting .gitignore patterns.
for dir := root; dir != "" && dir != string(filepath.Separator); dir = filepath.Dir(dir) {
giPath := filepath.Join(dir, ".gitignore")
data, err := os.ReadFile(giPath)
if err != nil {
continue
}
for _, line := range strings.Split(string(data), "\n") {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
continue
}
// Remove leading / for anchored patterns; we handle all as relative.
line = strings.TrimPrefix(line, "/")
// Remove trailing / for directory-only patterns.
line = strings.TrimSuffix(line, "/")
m.patterns = append(m.patterns, line)
}
}
return m
}
func (m *gitignoreMatcher) shouldIgnore(path string, isDir bool) bool {
name := filepath.Base(path)
for _, pat := range m.patterns {
// Simple match against full relative path or just the basename.
matched, _ := filepath.Match(pat, path)
if matched {
return true
}
matched, _ = filepath.Match(pat, name)
if matched {
return true
}
}
return false
}
+430
View File
@@ -0,0 +1,430 @@
package tools
import (
"fmt"
"os"
"path/filepath"
"strings"
"agentu/internal/safeio"
)
// ────────────────────────────────────────────────────────────────────────────
// V4 patch format data structures
// ────────────────────────────────────────────────────────────────────────────
type patchAction string
const (
actionAdd patchAction = "add"
actionUpdate patchAction = "update"
actionDelete patchAction = "delete"
)
type patchFile struct {
action patchAction
path string
newContent string // for "add" — full file content
lines []string // for "update" — lines to match/replace; for "delete" — unused
}
type parsedPatch struct {
files []patchFile
}
// ────────────────────────────────────────────────────────────────────────────
// Parser
// ────────────────────────────────────────────────────────────────────────────
// parsePatch parses a V4-format patch string.
//
// Format:
//
// *** Begin Patch
// *** Add File: <path>
// <file content>
// *** Update File: <path>
// @@
// <context line>
// @@
// -<old line>
// +<new line>
// *** Delete File: <path>
// *** End Patch
func parsePatch(patch string) (*parsedPatch, error) {
patch = strings.TrimSpace(patch)
if patch == "" {
return nil, fmt.Errorf("patch is empty")
}
lines := strings.Split(patch, "\n")
if len(lines) == 0 {
return nil, fmt.Errorf("patch is empty")
}
if strings.TrimSpace(lines[0]) != "*** Begin Patch" {
return nil, fmt.Errorf("patch must start with *** Begin Patch")
}
var files []patchFile
seenPaths := make(map[string]struct{})
seenEnd := false
i := 1
for i < len(lines) {
line := strings.TrimSpace(lines[i])
switch {
case line == "*** End Patch":
seenEnd = true
i = len(lines)
case strings.HasPrefix(line, "*** Add File:"):
path := strings.TrimSpace(strings.TrimPrefix(line, "*** Add File:"))
if err := validatePatchPath(path, seenPaths); err != nil {
return nil, err
}
i++
var contentLines []string
for i < len(lines) {
next := strings.TrimSpace(lines[i])
if strings.HasPrefix(next, "***") {
break
}
content := lines[i]
if strings.HasPrefix(content, "+") {
content = strings.TrimPrefix(content, "+")
}
contentLines = append(contentLines, content)
i++
}
files = append(files, patchFile{
action: actionAdd,
path: path,
newContent: strings.Join(contentLines, "\n") + "\n",
})
case strings.HasPrefix(line, "*** Delete File:"):
path := strings.TrimSpace(strings.TrimPrefix(line, "*** Delete File:"))
if err := validatePatchPath(path, seenPaths); err != nil {
return nil, err
}
files = append(files, patchFile{
action: actionDelete,
path: path,
})
i++
case strings.HasPrefix(line, "*** Update File:"):
path := strings.TrimSpace(strings.TrimPrefix(line, "*** Update File:"))
if err := validatePatchPath(path, seenPaths); err != nil {
return nil, err
}
i++
var hunkLines []string
for i < len(lines) {
next := strings.TrimSpace(lines[i])
if strings.HasPrefix(next, "***") {
break
}
hunkLines = append(hunkLines, lines[i])
i++
}
files = append(files, patchFile{
action: actionUpdate,
path: path,
lines: hunkLines,
})
default:
i++
}
}
if len(files) == 0 {
return nil, fmt.Errorf("patch contains no file operations")
}
if !seenEnd {
return nil, fmt.Errorf("patch must end with *** End Patch")
}
return &parsedPatch{files: files}, nil
}
func validatePatchPath(path string, seen map[string]struct{}) error {
if path == "" {
return fmt.Errorf("patch file path is required")
}
if _, ok := seen[path]; ok {
return fmt.Errorf("patch contains duplicate file path: %s", path)
}
seen[path] = struct{}{}
return nil
}
// ────────────────────────────────────────────────────────────────────────────
// Applier
// ────────────────────────────────────────────────────────────────────────────
// applyPatchFile applies a single patch file operation.
func applyPatchFile(baseDir string, pf patchFile) (string, error) {
switch pf.action {
case actionAdd:
return applyAdd(baseDir, pf)
case actionUpdate:
return applyUpdate(baseDir, pf)
case actionDelete:
return applyDelete(baseDir, pf)
default:
return "", fmt.Errorf("unknown patch action: %s", pf.action)
}
}
func applyAdd(baseDir string, pf patchFile) (string, error) {
path, err := resolvePathForWrite(baseDir, pf.path)
if err != nil {
return "", err
}
if _, err := os.Stat(path); err == nil {
return "", fmt.Errorf("cannot add existing file: %s", pf.path)
} else if !os.IsNotExist(err) {
return "", fmt.Errorf("stat new file %s: %w", pf.path, err)
}
if err := safeio.AtomicWriteFile(path, []byte(pf.newContent), 0o644); err != nil {
return "", fmt.Errorf("write new file %s: %w", pf.path, err)
}
return fmt.Sprintf("+ %s", pf.path), nil
}
func applyDelete(baseDir string, pf patchFile) (string, error) {
path, err := resolvePath(baseDir, pf.path)
if err != nil {
return "", err
}
if info, err := os.Stat(path); err != nil {
if os.IsNotExist(err) {
return "", fmt.Errorf("cannot delete non-existent file: %s", pf.path)
}
return "", err
} else if info.IsDir() {
return "", fmt.Errorf("cannot delete directory: %s", pf.path)
}
if err := os.Remove(path); err != nil {
return "", fmt.Errorf("delete file %s: %w", pf.path, err)
}
// Try to remove empty parent directories.
dir := filepath.Dir(path)
for dir != baseDir && dir != "." {
if entries, err := os.ReadDir(dir); err == nil && len(entries) == 0 {
os.Remove(dir)
dir = filepath.Dir(dir)
} else {
break
}
}
return fmt.Sprintf("- %s", pf.path), nil
}
func applyUpdate(baseDir string, pf patchFile) (string, error) {
path, err := resolvePath(baseDir, pf.path)
if err != nil {
return "", fmt.Errorf("cannot update non-existent file: %s: %w", pf.path, err)
}
existing, err := os.ReadFile(path)
if err != nil {
return "", fmt.Errorf("cannot update non-existent file: %s", pf.path)
}
info, err := os.Stat(path)
if err != nil {
return "", fmt.Errorf("stat file %s: %w", pf.path, err)
}
fileLines := strings.Split(string(existing), "\n")
// Parse the hunk lines into: context blocks + old lines + new lines.
result, err := applyHunkLines(fileLines, pf.lines, pf.path)
if err != nil {
return "", err
}
newContent := strings.Join(result, "\n")
if err := safeio.AtomicWriteFile(path, []byte(newContent), info.Mode().Perm()); err != nil {
return "", fmt.Errorf("write updated file %s: %w", pf.path, err)
}
return fmt.Sprintf("~ %s", pf.path), nil
}
// applyHunkLines applies a sequence of hunk lines to the file content.
// The hunk lines are a mix of:
//
// @@ — hunk separator (a context line block starts)
// <plain line> — context line (must appear in both old and new)
// -<line> — line to remove
// +<line> — line to add
//
// Hunks are matched greedily: for each group of contiguous context + - + + lines,
// find the first position in the file where context + - lines match,
// replace with context + + lines.
func applyHunkLines(fileLines []string, hunkLines []string, path string) ([]string, error) {
if len(hunkLines) == 0 {
return fileLines, nil
}
// Split the hunk lines into separate hunks at @@ markers.
type hunk struct {
lines []string
}
var hunks []hunk
var current hunk
for _, hl := range hunkLines {
trimmed := strings.TrimSpace(hl)
if strings.HasPrefix(trimmed, "@@") {
if len(current.lines) > 0 {
hunks = append(hunks, current)
current = hunk{}
}
continue
}
current.lines = append(current.lines, hl)
}
if len(current.lines) > 0 {
hunks = append(hunks, current)
}
if len(hunks) == 0 {
return fileLines, nil
}
result := make([]string, len(fileLines))
copy(result, fileLines)
offset := 0 // tracks net line count change from previous hunks
for hi, h := range hunks {
// Separate the hunk lines into a search pattern and replacement.
searchLines, replaceLines := buildHunkSearchReplace(h.lines)
if len(searchLines) == 0 && len(replaceLines) == 0 {
continue
}
matchStart := findHunkMatch(result, searchLines, offset)
if matchStart < 0 {
return nil, fmt.Errorf("hunk context not found in %s", path)
}
// Replace: remove searchLines, insert replaceLines.
matchEnd := matchStart + len(searchLines)
newResult := make([]string, 0, len(result)-len(searchLines)+len(replaceLines))
newResult = append(newResult, result[:matchStart]...)
newResult = append(newResult, replaceLines...)
newResult = append(newResult, result[matchEnd:]...)
result = newResult
offset = matchStart + len(replaceLines)
// Multiple matches check for this hunk only.
secondMatch := findHunkMatch(result, searchLines, offset)
if secondMatch >= 0 {
return nil, fmt.Errorf("hunk context matched multiple times in %s", path)
}
_ = hi
}
return result, nil
}
// buildHunkSearchReplace splits a hunk's lines into the search pattern
// (context + old lines) and replacement (context + new lines).
func buildHunkSearchReplace(lines []string) (search, replace []string) {
for _, l := range lines {
if strings.HasPrefix(l, "-") {
search = append(search, l[1:])
} else if strings.HasPrefix(l, "+") {
replace = append(replace, l[1:])
} else {
contextLine := strings.TrimPrefix(l, " ")
search = append(search, contextLine)
replace = append(replace, contextLine)
}
}
return
}
// findHunkMatch finds the first position >= offset in fileLines where
// searchLines matches exactly (line-by-line comparison).
func findHunkMatch(fileLines []string, searchLines []string, offset int) int {
for i := offset; i <= len(fileLines)-len(searchLines); i++ {
match := true
for j, sl := range searchLines {
if fileLines[i+j] != sl {
match = false
break
}
}
if match {
return i
}
}
return -1
}
// ────────────────────────────────────────────────────────────────────────────
// Preview
// ────────────────────────────────────────────────────────────────────────────
// summarizePatch returns a human-readable summary of the patch.
func summarizePatch(patch string) (string, error) {
pp, err := parsePatch(patch)
if err != nil {
return "", err
}
var b strings.Builder
added, updated, deleted := 0, 0, 0
for _, pf := range pp.files {
switch pf.action {
case actionAdd:
added++
fmt.Fprintf(&b, "+ add %s\n", pf.path)
case actionUpdate:
updated++
fmt.Fprintf(&b, "~ update %s\n", pf.path)
case actionDelete:
deleted++
fmt.Fprintf(&b, "- delete %s\n", pf.path)
}
}
summary := fmt.Sprintf("patch: %d file(s) changed", len(pp.files))
if added > 0 {
summary += fmt.Sprintf(" (%d added)", added)
}
if updated > 0 {
summary += fmt.Sprintf(" (%d updated)", updated)
}
if deleted > 0 {
summary += fmt.Sprintf(" (%d deleted)", deleted)
}
return summary + "\n" + b.String(), nil
}
// extractTargetPaths returns all file paths referenced in the patch.
func extractTargetPaths(baseDir, patch string) ([]string, error) {
pp, err := parsePatch(patch)
if err != nil {
return nil, err
}
var paths []string
for _, pf := range pp.files {
// Use resolvePathForWrite for all files — TargetPaths needs paths
// even for files that don't exist yet (rollback snapshots).
resolved, err := resolvePathForWrite(baseDir, pf.path)
if err != nil {
return nil, fmt.Errorf("file %s: %w", pf.path, err)
}
paths = append(paths, resolved)
}
return paths, nil
}
+68 -5
View File
@@ -2,22 +2,85 @@ package tools
import (
"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
}
+51 -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"
)
@@ -21,6 +23,17 @@ func NewShellRunTool(workingDir string) *ShellRunTool {
}
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{
Type: "function",
@@ -46,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)
@@ -66,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
}
+55 -5
View File
@@ -19,18 +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
@@ -81,8 +117,10 @@ func ReadOnlyBuiltins(workingDir string) *Registry {
gitStatus := NewGitStatusTool(workingDir)
gitDiff := NewGitDiffTool(workingDir)
gitLog := NewGitLogTool(workingDir)
projectOutline := NewProjectOutlineTool(workingDir)
todo := NewTodoTool()
registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols, gitStatus, gitDiff, gitLog)
registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols, projectOutline, todo, gitStatus, gitDiff, gitLog)
registry.RegisterAlias("file.read", fileRead)
registry.RegisterAlias("file.list", fileList)
registry.RegisterAlias("file.search", fileSearch)
@@ -97,36 +135,48 @@ 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, git_status, git_diff, git_log, web_search, web_fetch"
yoloTools := "file_list, file_read, file_search, code_symbols, file_write, file_edit, shell_run, git_status, git_diff, git_log, web_search, web_fetch"
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.
- When editing existing files, prefer file_edit for small or localized changes; use file_write only for creating files, replacing whole files intentionally, or appending.
- Use project_outline first when the user asks for a broad understanding of a project; prefer file_edit for localized changes, apply_patch for related multi-file changes, and file_write for whole-file creation or replacement.
- Do not use file tools as a substitute for a command execution request.
- 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.
+3 -3
View File
@@ -104,7 +104,7 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
}
}
for _, want := range []string{"file_list", "file_read", "file_search", "code_symbols", "git_status", "git_diff", "git_log"} {
for _, want := range []string{"file_list", "file_read", "file_search", "code_symbols", "project_outline", "todo", "git_status", "git_diff", "git_log"} {
if !slices.Contains(names, want) {
t.Fatalf("definitions missing %s: %#v", want, names)
}
@@ -173,7 +173,7 @@ func TestRegistryCanExposeWebSearch(t *testing.T) {
func TestBuiltinsExposeYoloTools(t *testing.T) {
registry := Builtins(t.TempDir())
for _, name := range []string{"file_write", "file_edit", "shell_run", "file.write", "file.edit", "shell.run"} {
for _, name := range []string{"file_write", "file_edit", "apply_patch", "shell_run", "file.write", "file.edit", "apply.patch", "shell.run"} {
if _, ok := registry.Get(name); !ok {
t.Fatalf("tool missing: %s", name)
}
@@ -198,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, git_status, git_diff, git_log, web_search, web_fetch"} {
for _, want := range []string{"use web_search", "Use web_fetch", "latest information", "source URLs", "Available tools: file_list, file_read, file_search, code_symbols, project_outline, todo, git_status, git_diff, git_log, web_search, web_fetch"} {
if !strings.Contains(withSearch, want) {
t.Fatalf("web search instructions missing %q:\n%s", want, withSearch)
}
+200
View File
@@ -0,0 +1,200 @@
package tools
import (
"context"
"encoding/json"
"os"
"path/filepath"
"strings"
"testing"
)
func TestApplyPatchChangesMultipleFiles(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "old.txt"), []byte("before\n"), 0o600); err != nil {
t.Fatal(err)
}
patch := `*** Begin Patch
*** Update File: old.txt
@@
-before
+after
*** Add File: new.txt
+created
*** End Patch`
out, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch}))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "2 file(s) changed") {
t.Fatalf("output = %q", out)
}
assertFileContent(t, filepath.Join(dir, "old.txt"), "after\n")
assertFileContent(t, filepath.Join(dir, "new.txt"), "created\n")
if info, err := os.Stat(filepath.Join(dir, "old.txt")); err != nil || info.Mode().Perm() != 0o600 {
t.Fatalf("updated mode = %v, err = %v", info.Mode().Perm(), err)
}
}
func TestApplyPatchRollsBackEarlierFilesOnFailure(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "first.txt")
if err := os.WriteFile(path, []byte("original\n"), 0o644); err != nil {
t.Fatal(err)
}
patch := `*** Begin Patch
*** Update File: first.txt
@@
-original
+changed
*** Update File: missing.txt
@@
-missing
+changed
*** End Patch`
_, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch}))
if err == nil {
t.Fatal("expected patch failure")
}
assertFileContent(t, path, "original\n")
}
func TestApplyPatchDryRunDoesNotWrite(t *testing.T) {
dir := t.TempDir()
patch := "*** Begin Patch\n*** Add File: dry.txt\n+dry\n*** End Patch"
out, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch, "dry_run": true}))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, "dry run") {
t.Fatalf("output = %q", out)
}
if _, err := os.Stat(filepath.Join(dir, "dry.txt")); !os.IsNotExist(err) {
t.Fatalf("dry run created file: %v", err)
}
}
func TestApplyPatchRejectsMalformedEnvelopeAndDuplicatePaths(t *testing.T) {
for _, patch := range []string{
"*** Begin Patch\n*** Add File: x.txt\n+x",
"prefix\n*** Begin Patch\n*** Add File: x.txt\n+x\n*** End Patch",
"*** Begin Patch\n*** Add File: x.txt\n+x\n*** Delete File: x.txt\n*** End Patch",
} {
if _, err := parsePatch(patch); err == nil {
t.Fatalf("expected parse failure for %q", patch)
}
}
}
func TestResolvePathRejectsSymlinkEscape(t *testing.T) {
parent := t.TempDir()
base := filepath.Join(parent, "project")
outside := filepath.Join(parent, "outside")
if err := os.MkdirAll(base, 0o755); err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(outside, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(outside, "secret.txt"), []byte("secret"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.Symlink(outside, filepath.Join(base, "link")); err != nil {
t.Skipf("symlinks unavailable: %v", err)
}
if _, err := resolvePath(base, "link/secret.txt"); err == nil {
t.Fatal("expected read symlink escape rejection")
}
if _, err := resolvePathForWrite(base, "link/new.txt"); err == nil {
t.Fatal("expected write symlink escape rejection")
}
}
func TestTodoRejectsInvalidData(t *testing.T) {
tool := NewTodoTool()
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"create","items":[""]}`)); err == nil {
t.Fatal("expected empty task rejection")
}
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"create","items":["one"]}`)); err != nil {
t.Fatal(err)
}
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"update","task":"one","status":"invalid"}`)); err == nil {
t.Fatal("expected invalid status rejection")
}
}
func TestProjectOutlineIncludesGoTypeScriptAndRustSymbols(t *testing.T) {
dir := t.TempDir()
files := map[string]string{
"main.go": "package demo\nfunc Run() {}\n",
"ui.ts": "export class Panel {}\nexport function render() {}\n",
"lib.rs": "pub struct Worker;\npub fn execute() {}\n",
}
for name, content := range files {
if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644); err != nil {
t.Fatal(err)
}
}
out, err := NewProjectOutlineTool(dir).Execute(context.Background(), json.RawMessage(`{}`))
if err != nil {
t.Fatal(err)
}
for _, want := range []string{"main.go", "Run", "ui.ts", "Panel", "render", "lib.rs", "Worker", "execute"} {
if !strings.Contains(out, want) {
t.Fatalf("outline missing %q:\n%s", want, out)
}
}
}
func TestRollbackStoreRestoresInReverseOrder(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "state.txt")
if err := os.WriteFile(path, []byte("one"), 0o644); err != nil {
t.Fatal(err)
}
store := NewRollbackStore()
first, err := store.Capture("edit", "state.txt", path)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("two"), 0o644); err != nil {
t.Fatal(err)
}
second, err := store.Capture("edit", "state.txt", path)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("three"), 0o644); err != nil {
t.Fatal(err)
}
if store.LastID() != second {
t.Fatalf("last ID = %q", store.LastID())
}
if _, err := store.Restore(second); err != nil {
t.Fatal(err)
}
assertFileContent(t, path, "two")
if store.LastID() != first {
t.Fatalf("last ID after restore = %q, want %q", store.LastID(), first)
}
}
func mustJSON(t *testing.T, value any) json.RawMessage {
t.Helper()
data, err := json.Marshal(value)
if err != nil {
t.Fatal(err)
}
return data
}
func assertFileContent(t *testing.T, path, want string) {
t.Helper()
data, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(data) != want {
t.Fatalf("%s content = %q, want %q", path, string(data), want)
}
}
+855 -127
View File
File diff suppressed because it is too large Load Diff
+642 -23
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", "ctrl+p/ctrl+n history"} {
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) {
@@ -133,6 +130,116 @@ func TestRunningStateRendersActivityIndicator(t *testing.T) {
}
}
type tuiBTWProvider struct{}
func (tuiBTWProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
return emit(llm.StreamEvent{Content: "side answer"})
}
func TestBTWCanBeSubmittedWhileMainTurnIsRunning(t *testing.T) {
assistant := agent.New(agent.Options{Provider: tuiBTWProvider{}, Model: "test-model"})
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
m.running = true
m.events = make(chan tea.Msg, 8)
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("/btw why?")})
m = next.(model)
if m.input.Value() != "/btw why?" {
t.Fatalf("running input = %q", m.input.Value())
}
next, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = next.(model)
if !m.btwRunning || m.input.Value() != "" {
t.Fatalf("btw running/input = %v/%q", m.btwRunning, m.input.Value())
}
if got := m.messages[len(m.messages)-1]; got.role != roleBTWUser || got.content != "why?" {
t.Fatalf("btw question message = %+v", got)
}
if cmd == nil {
t.Fatal("btw did not return a start command")
}
_ = cmd()
chunk := <-m.events
done := <-m.events
if _, ok := chunk.(btwChunkMsg); !ok {
t.Fatalf("first btw event = %T", chunk)
}
if _, ok := done.(btwDoneMsg); !ok {
t.Fatalf("final btw event = %T", done)
}
if len(assistant.Messages()) != 0 {
t.Fatalf("btw polluted agent history: %+v", assistant.Messages())
}
}
func TestBTWRequiresQuestion(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.input.SetValue("/btw")
next, _ := m.submit()
m = next.(model)
if m.btwRunning {
t.Fatal("empty btw question started a request")
}
last := m.messages[len(m.messages)-1]
if last.role != roleBTW || !strings.Contains(last.content, "Usage: /btw") {
t.Fatalf("empty btw message = %+v", last)
}
}
func TestBTWMessagesHaveDistinctLabels(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
question := stripANSI(m.messageView(roleBTWUser, "question", 80))
answer := stripANSI(m.messageView(roleBTW, "answer", 80))
if !strings.Contains(question, "BTW · You") || !strings.Contains(answer, "BTW") {
t.Fatalf("btw labels missing:\n%s\n%s", question, answer)
}
}
func TestBTWAndMainStreamChunksStaySeparate(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleUser, content: "main"})
m.appendAssistantChunk("main one ")
m.messages = append(m.messages, message{role: roleBTWUser, content: "side"})
m.appendBTWChunk("side one ")
m.appendAssistantChunk("main two")
m.appendBTWChunk("side two")
var mainAnswers, sideAnswers []string
for _, msg := range m.messages {
switch msg.role {
case roleAssistant:
mainAnswers = append(mainAnswers, msg.content)
case roleBTW:
sideAnswers = append(sideAnswers, msg.content)
}
}
if !reflect.DeepEqual(mainAnswers, []string{"main one main two"}) {
t.Fatalf("main answers = %#v", mainAnswers)
}
if !reflect.DeepEqual(sideAnswers, []string{"side one side two"}) {
t.Fatalf("side answers = %#v", sideAnswers)
}
}
func TestMainCompletionKeepsBTWEventStreamAlive(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
events := make(chan tea.Msg, 1)
m.events = events
m.running = true
m.btwRunning = true
next, _ := m.Update(turnDoneMsg{})
m = next.(model)
if m.running || !m.btwRunning || m.events != events {
t.Fatalf("state after main completion: running=%v btw=%v events=%p", m.running, m.btwRunning, m.events)
}
next, _ = m.Update(btwDoneMsg{})
m = next.(model)
if m.active() || m.events != nil {
t.Fatalf("state after btw completion: active=%v events=%p", m.active(), m.events)
}
}
func TestUserMessageRendersFullWidthBubble(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
@@ -336,6 +443,23 @@ func TestInputHistoryNavigation(t *testing.T) {
}
}
func TestUpArrowRecallsPreviousInput(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.rememberInput("first")
m.rememberInput("second")
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyUp})
m = next.(model)
if got := m.input.Value(); got != "second" {
t.Fatalf("after first up input = %q, want %q", got, "second")
}
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyUp})
m = next.(model)
if got := m.input.Value(); got != "first" {
t.Fatalf("after second up input = %q, want %q", got, "first")
}
}
func TestToolMessageRendersStructuredCard(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`})
@@ -385,6 +509,98 @@ func TestToolLogOutputUpdatesExistingToolCard(t *testing.T) {
}
}
func TestToolLogStreamingChunksAppendToSameCard(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{role: roleTool, content: `shell_run echo hi`})
if !m.mergeToolLog("shell_run echo hi\n[output]\none\n") {
t.Fatal("expected first mergeToolLog to return true")
}
if !m.mergeToolLog("shell_run echo hi\n[output]\ntwo\n") {
t.Fatal("expected second mergeToolLog to return true")
}
if len(m.messages) != 1 {
t.Fatalf("messages count = %d, want 1", len(m.messages))
}
plain := stripANSI(m.renderMessages())
if !strings.Contains(plain, "one") || !strings.Contains(plain, "two") {
t.Fatalf("rendered output missing streamed lines:\n%s", plain)
}
if strings.Contains(plain, "onetwo") {
t.Fatalf("streamed lines were concatenated:\n%s", plain)
}
}
func TestShellRunCardShowsCommandWorkdirAndOutput(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
m.messages = append(m.messages, message{
role: roleTool,
content: "shell_run {\"command\":\"go test ./...\",\"workdir\":\"internal/tui\"}\n[output]\nok agentu/internal/tui",
})
plain := stripANSI(m.renderMessages())
for _, want := range []string{"$ go test ./...", "cwd internal/tui", "output", "ok agentu/internal/tui"} {
if !strings.Contains(plain, want) {
t.Fatalf("shell card missing %q:\n%s", want, plain)
}
}
if strings.Contains(plain, `{"command"`) || strings.Contains(plain, iconTool+" shell_run") {
t.Fatalf("shell card still exposes raw tool arguments:\n%s", plain)
}
if status := toolStatusText(`shell_run {"command":"go test ./...","workdir":"internal/tui"}`); status != "$ go test ./..." {
t.Fatalf("shell status = %q", status)
}
}
func TestShellRunStreamingEventsPreserveLineBreaks(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
args := `{"command":"printf lines"}`
m.messages = append(m.messages, message{role: roleTool, content: "shell_run " + args})
next, _ := m.Update(toolLogMsg("\n[tool] shell_run " + args + "\n[output]\none\n"))
m = next.(model)
next, _ = m.Update(toolLogMsg("\n[tool] shell_run " + args + "\n[output]\n\n"))
m = next.(model)
next, _ = m.Update(toolLogMsg("\n[tool] shell_run " + args + "\n[output]\ntwo\n"))
m = next.(model)
parts := parseToolLog(m.messages[0].content)
if parts.output != "one\n\ntwo" {
t.Fatalf("streamed output = %q", parts.output)
}
if len(m.messages) != 1 {
t.Fatalf("streaming created %d tool cards, want 1", len(m.messages))
}
plain := stripANSI(m.renderMessages())
if strings.Contains(plain, "onetwo") {
t.Fatalf("streamed lines were concatenated:\n%s", plain)
}
}
func TestShellOutputPreviewKeepsHeadAndLatestLines(t *testing.T) {
var lines []string
for i := 1; i <= 12; i++ {
lines = append(lines, fmt.Sprintf("line %d", i))
}
preview := strings.Join(shellOutputPreviewLines(strings.Join(lines, "\n"), 80), "\n")
for _, want := range []string{"line 1", "line 4", "5 lines omitted", "line 10", "line 12"} {
if !strings.Contains(preview, want) {
t.Fatalf("preview missing %q:\n%s", want, preview)
}
}
if strings.Contains(preview, "line 5\n") || strings.Contains(preview, "line 9\n") {
t.Fatalf("preview contains a line that should be omitted:\n%s", preview)
}
}
func TestSanitizeTerminalOutputRemovesEscapesAndFoldsCarriageReturns(t *testing.T) {
got := sanitizeTerminalOutput("\x1b[31mred\x1b[0m\nprogress 10%\rprogress 20%\r\nclean\x00text")
want := "red\nprogress 20%\ncleantext"
if got != want {
t.Fatalf("sanitized output = %q, want %q", got, want)
}
}
type tuiCompactProvider struct{}
func (tuiCompactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
@@ -480,16 +696,60 @@ func runInDir(t *testing.T, dir string, name string, args ...string) {
}
}
func TestBangShellCommandRunsFromInput(t *testing.T) {
dir := t.TempDir()
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true, WorkingDir: dir})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("! printf hello")
next, _ = m.submit()
m = next.(model)
if m.input.Value() != "" {
t.Fatalf("input not cleared after bang command: %q", m.input.Value())
}
rendered := stripANSI(m.renderMessages())
if !strings.Contains(rendered, "$ printf hello") {
t.Fatalf("rendered output missing shell echo:\n%s", rendered)
}
if !strings.Contains(rendered, "hello") {
t.Fatalf("rendered output missing 'hello':\n%s", rendered)
}
}
func TestBangShellCommandRequiresYolo(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: false})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("! printf hello")
next, _ = m.submit()
m = next.(model)
rendered := stripANSI(m.renderMessages())
if !strings.Contains(rendered, "shell command input requires --yolo") {
t.Fatalf("rendered output missing yolo requirement:\n%s", rendered)
}
}
func TestBangShellCommandRejectsEmptyCommand(t *testing.T) {
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("!")
next, _ = m.submit()
m = next.(model)
rendered := stripANSI(m.renderMessages())
if !strings.Contains(rendered, "usage: ! <command>") {
t.Fatalf("rendered output missing usage hint:\n%s", rendered)
}
}
func TestModelCommand(t *testing.T) {
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)
}
}
@@ -508,11 +768,124 @@ func TestModelThinkingCommand(t *testing.T) {
}
}
func TestModelCommandPaletteShowsConfiguredModels(t *testing.T) {
manager := &fakeModelManager{model: "model-a", models: []string{"model-a", "model-b"}}
m := newModel(context.Background(), nil, Options{ModelName: "model-a", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model ")
m.afterInputChanged()
plain := stripANSI(m.View())
for _, want := range []string{iconCommand + " Model candidates", "model-a", "current model", "model-b", "switch model"} {
if !strings.Contains(plain, want) {
t.Fatalf("rendered palette missing %q:\n%s", want, plain)
}
}
}
func TestModelCommandPaletteSelectsConfiguredModel(t *testing.T) {
manager := &fakeModelManager{model: "model-a", models: []string{"model-a", "model-b"}}
m := newModel(context.Background(), nil, Options{ModelName: "model-a", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model ")
m.afterInputChanged()
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyDown})
m = next.(model)
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = next.(model)
if manager.model != "model-b" {
t.Fatalf("model = %q", manager.model)
}
if got := m.input.Value(); got != "" {
t.Fatalf("input = %q", got)
}
rendered := stripANSI(m.renderMessages())
if !strings.Contains(rendered, "model: model-b") {
t.Fatalf("rendered output missing model switch:\n%s", rendered)
}
}
func TestModelCommandPaletteFiltersBareModelPartial(t *testing.T) {
manager := &fakeModelManager{model: "claude-sonnet", models: []string{"claude-sonnet", "deepseek-v4-flash"}}
m := newModel(context.Background(), nil, Options{ModelName: "claude-sonnet", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model deep")
m.afterInputChanged()
palette := stripANSI(m.commandPaletteView())
if !strings.Contains(palette, "deepseek-v4-flash") {
t.Fatalf("expected deepseek-v4-flash in palette:\n%s", palette)
}
if strings.Contains(palette, "claude-sonnet") {
t.Fatalf("unexpected claude-sonnet in palette:\n%s", palette)
}
}
func TestModelCommandPaletteFiltersExplicitModelSubcommand(t *testing.T) {
manager := &fakeModelManager{model: "claude-sonnet", models: []string{"claude-sonnet", "deepseek-v4-flash"}}
m := newModel(context.Background(), nil, Options{ModelName: "claude-sonnet", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model model deep")
m.afterInputChanged()
palette := stripANSI(m.commandPaletteView())
if !strings.Contains(palette, "deepseek-v4-flash") {
t.Fatalf("expected deepseek-v4-flash in palette:\n%s", palette)
}
if strings.Contains(palette, "claude-sonnet") {
t.Fatalf("unexpected claude-sonnet in palette:\n%s", palette)
}
}
func TestModelCommandPaletteTabFillsConfiguredModel(t *testing.T) {
manager := &fakeModelManager{model: "claude-sonnet", models: []string{"claude-sonnet", "deepseek-v4-flash"}}
m := newModel(context.Background(), nil, Options{ModelName: "claude-sonnet", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model deep")
m.afterInputChanged()
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyTab})
m = next.(model)
if got := m.input.Value(); got != "/model model deepseek-v4-flash" {
t.Fatalf("input = %q", got)
}
if manager.model != "claude-sonnet" {
t.Fatalf("model should not change on tab: %q", manager.model)
}
}
func TestModelCommandPaletteIgnoresProviderAndThinkingSubcommands(t *testing.T) {
manager := &fakeModelManager{model: "model-a", models: []string{"model-a", "model-b"}}
m := newModel(context.Background(), nil, Options{ModelName: "model-a", ModelManager: manager})
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
m = next.(model)
m.input.SetValue("/model provider ")
m.afterInputChanged()
if sugs := m.commandSuggestions(); len(sugs) != 0 {
t.Fatalf("provider suggestions should be empty, got %d", len(sugs))
}
if view := m.commandPaletteView(); view != "" {
t.Fatalf("provider palette should be empty:\n%s", view)
}
m.input.SetValue("/model thinking ")
m.afterInputChanged()
if sugs := m.commandSuggestions(); len(sugs) != 0 {
t.Fatalf("thinking suggestions should be empty, got %d", len(sugs))
}
if view := m.commandPaletteView(); view != "" {
t.Fatalf("thinking palette should be empty:\n%s", view)
}
}
type fakeModelManager struct {
provider string
model string
thinking string
sessionName string
models []string
saved bool
}
@@ -583,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,
@@ -604,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"}`
@@ -611,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"))
}
+40 -15
View File
@@ -43,6 +43,7 @@ type theme struct {
Background string
Surface string
SurfaceAlt string
InputBg string
Border string
Text string
Muted string
@@ -60,11 +61,13 @@ type styles struct {
Muted lipgloss.Style
Activity lipgloss.Style
Header lipgloss.Style
ModelMeta 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
@@ -98,8 +101,8 @@ func lightTheme() theme {
return theme{
Mode: ThemeLight,
Background: "#F8FAFC",
Surface: "#FFFFFF",
SurfaceAlt: "#EEF2F7",
InputBg: "#E4E4E4",
Border: "#CBD5E1",
Text: "#111827",
Muted: "#64748B",
@@ -116,8 +119,8 @@ func darkTheme() theme {
return theme{
Mode: ThemeDark,
Background: "#111827",
Surface: "#1F2937",
SurfaceAlt: "#0F172A",
InputBg: "#383838",
Border: "#475569",
Text: "#E5E7EB",
Muted: "#94A3B8",
@@ -130,10 +133,17 @@ 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)),
@@ -146,10 +156,6 @@ func (t theme) styles() styles {
Bold(true).
Foreground(lipgloss.Color(t.Title)),
ModelMeta: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1),
Footer: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Muted)).
Padding(0, 1),
@@ -159,14 +165,31 @@ func (t theme) styles() styles {
InputBox: lipgloss.NewStyle().
Foreground(lipgloss.Color(t.Text)).
Border(lipgloss.ThickBorder()).
BorderForeground(lipgloss.Color(t.Border)).
Padding(0, 2),
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)),
@@ -251,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().
@@ -263,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 -215
View File
@@ -1,228 +1,15 @@
package main
import (
"bufio"
"context"
"errors"
"flag"
"fmt"
"net/http"
"os"
"os/signal"
"path/filepath"
"strings"
"agentu/internal/agent"
"agentu/internal/session"
"agentu/internal/tools"
"agentu/internal/tui"
"agentu/pkg/config"
"agentu/pkg/llm"
"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
}
if err := initializeSession(modelManager, *resumeID); err != nil {
return err
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
defer stop()
if !*plain {
err := tui.Run(ctx, assistant, tui.Options{
ModelName: providerConfig.Model,
Yolo: *yolo,
ThemeMode: themeMode,
ModelManager: modelManager,
WorkingDir: cfg.Agent.WorkingDir,
})
printExitSession(modelManager)
return err
}
return repl(ctx, assistant, modelManager)
}
func initializeSession(modelManager *session.Manager, resumeID string) error {
if resumeID != "" {
if err := modelManager.Resume(resumeID); err != nil {
return fmt.Errorf("resume session: %w", err)
}
return nil
}
if _, err := modelManager.NewSession(); err != nil {
return fmt.Errorf("start new session: %w", err)
}
return nil
}
func bootstrapConfig(configPath string) error {
resolved, err := config.ResolvePath(configPath)
if err != nil {
return fmt.Errorf("resolve config path: %w", err)
}
dir := filepath.Dir(resolved)
if err := os.MkdirAll(dir, 0o755); err != nil {
return fmt.Errorf("create config directory: %w", err)
}
if err := os.WriteFile(resolved, []byte(defaultConfigContent), 0o644); err != nil {
return fmt.Errorf("write default config: %w", err)
}
return nil
}
func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error {
scanner := bufio.NewScanner(os.Stdin)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
fmt.Println("agentu")
fmt.Println("Type /exit to quit, /clear to reset context, /compact to compress context.")
for {
fmt.Print("> ")
if !scanner.Scan() {
if err := scanner.Err(); err != nil {
return err
}
fmt.Println()
printExitSession(modelManager)
return nil
}
input := strings.TrimSpace(scanner.Text())
switch input {
case "":
continue
case "/exit", "/quit":
printExitSession(modelManager)
return nil
case "/clear":
assistant.Clear()
fmt.Println("context cleared")
continue
case "/compact":
if err := assistant.Compact(ctx); err != nil {
fmt.Fprintln(os.Stderr, "compact error:", err)
} else {
fmt.Println("context compacted")
if modelManager != nil {
_ = modelManager.Save()
}
}
continue
}
if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
fmt.Fprintln(os.Stderr, "\nerror:", err)
}
_ = modelManager.Save()
fmt.Println()
}
}
func printExitSession(m *session.Manager) {
if m == nil {
return
}
_ = m.Save()
id := m.CurrentSessionID()
name := m.CurrentSessionName()
if id != "" {
if name != "" && name != id {
fmt.Fprintf(os.Stderr, "session: %s (%s)\n", id, name)
} else {
fmt.Fprintf(os.Stderr, "session: %s\n", id)
}
fmt.Fprintf(os.Stderr, "resume: agentu --resume %s\n", id)
}
}
+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())
}
}