Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
065b30a84c | ||
|
|
47671f3098 | ||
|
|
28195519d1 | ||
|
|
425f6ee5ee | ||
|
|
c8d8329dbb | ||
|
|
dabf5bfecc | ||
|
|
9abd5c3b5f | ||
|
|
3d403bd685 | ||
|
|
7b8bad5e62 |
@@ -3,3 +3,4 @@ agentu*.yaml
|
|||||||
/coverage.out
|
/coverage.out
|
||||||
*.swp
|
*.swp
|
||||||
chat-history.txt
|
chat-history.txt
|
||||||
|
dist/
|
||||||
|
|||||||
@@ -1,18 +1,24 @@
|
|||||||
# agentu
|
# agentu
|
||||||
|
|
||||||
agentu is a small local agent written in Go. The current MVP is a terminal chat
|
agentu is a small local coding agent written in Go. It provides a terminal chat
|
||||||
UI with OpenAI-compatible providers, streaming responses, session-level
|
UI with OpenAI-compatible providers, streaming responses, session-level
|
||||||
model/provider switching, and local tools.
|
model/provider switching, skills, approvals, and local tools.
|
||||||
|
|
||||||
## Current MVP
|
## v0.3.1
|
||||||
|
|
||||||
- TUI-first chat experience with a light theme by default.
|
- TUI-first chat experience with a light theme by default.
|
||||||
- OpenAI-compatible `/v1/chat/completions` provider support.
|
- OpenAI-compatible `/v1/chat/completions` provider support.
|
||||||
|
- OpenAI Responses API (`/v1/responses`) provider support.
|
||||||
- Multiple providers from `~/.agentu/config.yaml`.
|
- Multiple providers from `~/.agentu/config.yaml`.
|
||||||
- Session-only provider, model, and thinking-level switching.
|
- Session-only provider, model, and thinking-level switching.
|
||||||
- Read-only project tools enabled by default.
|
- Interactive approval for model-requested file writes and shell commands.
|
||||||
|
- Atomic file writes, per-file rollback snapshots, and multi-file patches.
|
||||||
|
- Project outlines with Go, TypeScript/JavaScript, and Rust symbols.
|
||||||
|
- Global Codex-style skills loaded from `~/.agentu/skills`.
|
||||||
|
- Persistent session index and in-memory task tracking.
|
||||||
- Built-in web search and URL fetching.
|
- Built-in web search and URL fetching.
|
||||||
- `--yolo` mode for file writes and shell execution.
|
- `--yolo` mode for automatic file writes and shell execution.
|
||||||
|
- Temporary `/btw` side questions that can run while the main task continues.
|
||||||
|
|
||||||
## Quick Start
|
## Quick Start
|
||||||
|
|
||||||
@@ -23,7 +29,15 @@ export AGENTU_API_KEY='your-api-key'
|
|||||||
go run .
|
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
|
```sh
|
||||||
go run . --yolo
|
go run . --yolo
|
||||||
@@ -42,6 +56,7 @@ Use `--config <path>` to load another file.
|
|||||||
```yaml
|
```yaml
|
||||||
providers:
|
providers:
|
||||||
loveuer:
|
loveuer:
|
||||||
|
api_type: chat
|
||||||
base_url: https://ai.loveuer.com
|
base_url: https://ai.loveuer.com
|
||||||
api_key: ${AGENTU_API_KEY}
|
api_key: ${AGENTU_API_KEY}
|
||||||
models:
|
models:
|
||||||
@@ -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
|
first provider by name. Use `/model provider <name>` to switch providers for the
|
||||||
current session.
|
current session.
|
||||||
|
|
||||||
|
`api_type` selects which API protocol the provider speaks:
|
||||||
|
|
||||||
|
```text
|
||||||
|
chat OpenAI-compatible /v1/chat/completions (default)
|
||||||
|
responses OpenAI Responses API /v1/responses
|
||||||
|
```
|
||||||
|
|
||||||
|
The Responses API client converts the conversation history to Responses input
|
||||||
|
items (`user`/`assistant` messages, `function_call`, and
|
||||||
|
`function_call_output`), sends the system prompt as the top-level
|
||||||
|
`instructions` field, streams `response.output_text.delta` events, and supports
|
||||||
|
function calling through `response.output_item.added`,
|
||||||
|
`response.function_call_arguments.delta`, and `response.output_item.done`.
|
||||||
|
|
||||||
`models` is required. `/model model <name>` is restricted to that list. Each
|
`models` is required. `/model model <name>` is restricted to that list. Each
|
||||||
entry is an object with:
|
entry is an object with:
|
||||||
|
|
||||||
@@ -77,7 +106,11 @@ none, middle, high, xhigh, max
|
|||||||
```
|
```
|
||||||
|
|
||||||
When configured or changed with `/model thinking ...`, the value is sent as a
|
When configured or changed with `/model thinking ...`, the value is sent as a
|
||||||
top-level OpenAI-compatible request field named `thinking`.
|
top-level request field. The default field depends on `api_type`: `chat` uses
|
||||||
|
`thinking`, while `responses` uses the official `reasoning` field, sent as
|
||||||
|
`{"effort": ...}` with agentu's `middle` level mapped to the official `medium`
|
||||||
|
value (`none`, `high`, `xhigh`, `max` pass through). Override with
|
||||||
|
`thinking_param` if a provider expects a different field name.
|
||||||
|
|
||||||
## CLI Flags
|
## CLI Flags
|
||||||
|
|
||||||
@@ -88,28 +121,38 @@ go run . [flags]
|
|||||||
- `--config <path>` loads a YAML config file. Default: `~/.agentu/config.yaml`.
|
- `--config <path>` loads a YAML config file. Default: `~/.agentu/config.yaml`.
|
||||||
- `--theme light|dark` selects the TUI theme. Default: `light`.
|
- `--theme light|dark` selects the TUI theme. Default: `light`.
|
||||||
- `--plain` uses the simple line-based REPL instead of the TUI.
|
- `--plain` uses the simple line-based REPL instead of the TUI.
|
||||||
- `--yolo` enables automatic file writes and shell execution.
|
- `--yolo` automatically approves model-requested file writes and shell execution.
|
||||||
- `--resume <id>` resumes a saved session from `~/.agentu/sessions`.
|
- `--resume <id>` resumes a saved session from `~/.agentu/sessions`.
|
||||||
|
|
||||||
## TUI Controls
|
## TUI Controls
|
||||||
|
|
||||||
- `enter` sends the message.
|
- `enter` sends the message.
|
||||||
- `ctrl+j` or `alt+enter` inserts a newline.
|
- `ctrl+j` inserts a newline.
|
||||||
|
- `up` recalls the previous input; `ctrl+p` / `ctrl+n` navigate input history.
|
||||||
- The input starts at one line and grows up to six lines.
|
- The input starts at one line and grows up to six lines.
|
||||||
- `page up` / `page down` scrolls the transcript.
|
- `page up` / `page down` scrolls the transcript.
|
||||||
- `esc` cancels a running response.
|
- `esc` cancels a running response.
|
||||||
|
- While a task is running, `/btw <question>` asks a concurrent side question.
|
||||||
|
- Entering only `.` means “continue with the current task.”
|
||||||
|
- At a tool approval prompt, `y` approves and `n` or `esc` denies the operation.
|
||||||
- Type `/` to show available commands.
|
- Type `/` to show available commands.
|
||||||
|
- In `--yolo` mode, type `! <command>` to run a shell command from the configured working directory.
|
||||||
|
|
||||||
Slash commands:
|
Slash commands:
|
||||||
|
|
||||||
|
- `/btw <question>` answers from a snapshot of the current context without adding the question or answer to conversation history. It never exposes tools and can run alongside the main task.
|
||||||
- `/clear` resets conversation history.
|
- `/clear` resets conversation history.
|
||||||
- `/compact` summarizes older conversation history and keeps 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`.
|
- `/status` shows changed files with `git status --short`.
|
||||||
- `/diff [--stat] [path...]` shows the unstaged git diff.
|
- `/diff [--stat] [path...]` shows the unstaged git diff.
|
||||||
- `/test [command...]` runs a project test command; defaults to `go test ./...` for Go modules or `npm test` for Node projects.
|
- `/test [command...]` runs a project test command; defaults to `go test ./...` for Go modules or `npm test` for Node projects.
|
||||||
|
- `/skills` lists global skills.
|
||||||
|
- `/skills install <path-or-url>` installs a skill and activates it in the current conversation.
|
||||||
|
- `/rollback [id]` restores a file snapshot; without an ID it restores the newest snapshot.
|
||||||
- `/model` shows current provider, model, thinking, and switch usage.
|
- `/model` shows current provider, model, thinking, and switch usage.
|
||||||
- `/model provider <name>` switches provider for the current session.
|
- `/model provider <name>` switches provider for the current session.
|
||||||
- `/model model <name>` or `/model <name>` switches model for the current session.
|
- `/model model <name>` or `/model <name>` switches model for the current session.
|
||||||
|
In the TUI, type `/model ` or `/model model ` to choose from configured model IDs in the candidate palette.
|
||||||
- `/model thinking <none|middle|high|xhigh|max>` changes thinking for the current session.
|
- `/model thinking <none|middle|high|xhigh|max>` changes thinking for the current session.
|
||||||
- `/sessions` lists saved sessions.
|
- `/sessions` lists saved sessions.
|
||||||
- `/resume [id]` resumes by ID or opens an interactive picker.
|
- `/resume [id]` resumes by ID or opens an interactive picker.
|
||||||
@@ -119,22 +162,36 @@ Slash commands:
|
|||||||
|
|
||||||
## Local Tools
|
## Local Tools
|
||||||
|
|
||||||
Read-only tools are available by default:
|
These tools are always available:
|
||||||
|
|
||||||
- `file_read`
|
- `file_read`
|
||||||
- `file_list`
|
- `file_list`
|
||||||
- `file_search` with optional `context_lines` and `max_results`
|
- `file_search` with optional `context_lines` and `max_results`
|
||||||
- `code_symbols` for Go packages, imports, types, functions, and methods
|
- `code_symbols` for Go packages, imports, types, functions, and methods
|
||||||
|
- `project_outline` for a gitignore-aware tree with Go, TypeScript/JavaScript, and Rust symbols
|
||||||
|
- `todo` for conversation-local task tracking
|
||||||
|
- `git_status` shows changed files
|
||||||
|
- `git_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_search`
|
||||||
- `web_fetch`
|
- `web_fetch`
|
||||||
|
|
||||||
`--yolo` additionally enables:
|
In the TUI, these tools require approval by default and run automatically with
|
||||||
|
`--yolo`:
|
||||||
|
|
||||||
- `file_write`
|
- `file_write` for creating or replacing whole files
|
||||||
- `shell_run`
|
- `file_edit` for surgical line-range or pattern edits to existing files
|
||||||
|
- `apply_patch` for atomic-style multi-file V4 patches, including dry runs
|
||||||
|
- `shell_run` for non-interactive shell commands; output streams live into the
|
||||||
|
TUI while the command runs instead of appearing only after completion
|
||||||
|
|
||||||
|
Before a file mutation, agentu stores an in-memory snapshot and includes its
|
||||||
|
`rb_...` ID in the tool result. Use `/rollback` or `/rollback <id>` to restore
|
||||||
|
it. Snapshots last for the lifetime of the process. The plain REPL remains
|
||||||
|
read-only unless `--yolo` is supplied because it has no interactive approval UI.
|
||||||
|
|
||||||
Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
|
Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
|
||||||
example:
|
example:
|
||||||
@@ -143,6 +200,21 @@ example:
|
|||||||
请帮我运行 `go test ./...` 并总结结果
|
请帮我运行 `go test ./...` 并总结结果
|
||||||
```
|
```
|
||||||
|
|
||||||
|
## Global Skills
|
||||||
|
|
||||||
|
At startup, agentu loads both Codex-style directories containing `SKILL.md`
|
||||||
|
and legacy flat Markdown files from `~/.agentu/skills`. Skill instructions are
|
||||||
|
added to the system prompt. Install a local directory, Markdown file, or direct
|
||||||
|
HTTP(S) Markdown URL from the TUI:
|
||||||
|
|
||||||
|
```text
|
||||||
|
/skills install ./my-skill
|
||||||
|
/skills install https://example.com/SKILL.md
|
||||||
|
```
|
||||||
|
|
||||||
|
Directory installs preserve supporting files such as `references/`; symlinks
|
||||||
|
inside a source directory are rejected.
|
||||||
|
|
||||||
## Development
|
## Development
|
||||||
|
|
||||||
Run the verification suite:
|
Run the verification suite:
|
||||||
@@ -159,5 +231,10 @@ Local config and build outputs are intentionally ignored by git.
|
|||||||
|
|
||||||
- Provider/model/thinking changes are session-only; agentu does not write them
|
- Provider/model/thinking changes are session-only; agentu does not write them
|
||||||
back to `~/.agentu/config.yaml`.
|
back to `~/.agentu/config.yaml`.
|
||||||
|
- `api_type` is per provider and cannot be changed at runtime; use
|
||||||
|
`/model provider <name>` to switch to a provider configured with a different
|
||||||
|
API type.
|
||||||
- The config schema is intentionally strict during MVP development. Deprecated
|
- The config schema is intentionally strict during MVP development. Deprecated
|
||||||
fields such as `provider` or `active_provider` are rejected.
|
fields such as `provider` or `active_provider` are rejected.
|
||||||
|
- Session metadata is indexed in `~/.agentu/sessions/index.json`; complete
|
||||||
|
conversations remain stored as one JSON file per session.
|
||||||
|
|||||||
@@ -1,5 +1,8 @@
|
|||||||
providers:
|
providers:
|
||||||
loveuer:
|
loveuer:
|
||||||
|
# API protocol: "chat" (OpenAI-compatible /v1/chat/completions, default)
|
||||||
|
# or "responses" (OpenAI Responses API /v1/responses).
|
||||||
|
api_type: chat
|
||||||
base_url: https://ai.loveuer.com
|
base_url: https://ai.loveuer.com
|
||||||
api_key: ${AGENTU_API_KEY}
|
api_key: ${AGENTU_API_KEY}
|
||||||
models:
|
models:
|
||||||
|
|||||||
+448
-45
@@ -7,18 +7,25 @@ import (
|
|||||||
"io"
|
"io"
|
||||||
"sort"
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"agentu/pkg/llm"
|
|
||||||
"agentu/internal/tools"
|
"agentu/internal/tools"
|
||||||
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
defaultMaxToolRounds = 50
|
defaultMaxToolRounds = 50
|
||||||
doomLoopThreshold = 3
|
doomLoopThreshold = 3
|
||||||
compactRecentMessages = 6
|
compactMinRecentMessages = 6
|
||||||
|
compactRecentRetentionPercent = 20
|
||||||
summaryMessageLimit = 8000
|
summaryMessageLimit = 8000
|
||||||
contextSummaryPrefix = "[context summary]"
|
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" +
|
summarizerSystemPrompt = "You are a conversation summarizer. Summarize the following conversation history into a concise but detailed summary. Preserve:\n" +
|
||||||
"- Key facts, decisions, and conclusions\n" +
|
"- Key facts, decisions, and conclusions\n" +
|
||||||
@@ -29,7 +36,13 @@ const (
|
|||||||
"Output only the summary, no preamble."
|
"Output only the summary, no preamble."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
type toolExecutionResult struct {
|
||||||
|
call llm.ToolCall
|
||||||
|
output string
|
||||||
|
}
|
||||||
|
|
||||||
type Agent struct {
|
type Agent struct {
|
||||||
|
mu sync.RWMutex
|
||||||
provider llm.Provider
|
provider llm.Provider
|
||||||
providerName string
|
providerName string
|
||||||
model string
|
model string
|
||||||
@@ -38,6 +51,8 @@ type Agent struct {
|
|||||||
thinkingEnabled bool
|
thinkingEnabled bool
|
||||||
systemPrompt string
|
systemPrompt string
|
||||||
toolRegistry *tools.Registry
|
toolRegistry *tools.Registry
|
||||||
|
authorizer ToolAuthorizer
|
||||||
|
rollbackStore *tools.RollbackStore
|
||||||
toolTimeout time.Duration
|
toolTimeout time.Duration
|
||||||
maxToolRounds int
|
maxToolRounds int
|
||||||
maxContextTokens int
|
maxContextTokens int
|
||||||
@@ -54,6 +69,8 @@ type Options struct {
|
|||||||
ThinkingEnabled bool
|
ThinkingEnabled bool
|
||||||
SystemPrompt string
|
SystemPrompt string
|
||||||
ToolRegistry *tools.Registry
|
ToolRegistry *tools.Registry
|
||||||
|
Authorizer ToolAuthorizer
|
||||||
|
RollbackStore *tools.RollbackStore
|
||||||
ToolTimeout time.Duration
|
ToolTimeout time.Duration
|
||||||
MaxToolIterations int
|
MaxToolIterations int
|
||||||
MaxContextTokens int
|
MaxContextTokens int
|
||||||
@@ -77,6 +94,8 @@ func New(opts Options) *Agent {
|
|||||||
thinkingEnabled: opts.ThinkingEnabled,
|
thinkingEnabled: opts.ThinkingEnabled,
|
||||||
systemPrompt: opts.SystemPrompt,
|
systemPrompt: opts.SystemPrompt,
|
||||||
toolRegistry: opts.ToolRegistry,
|
toolRegistry: opts.ToolRegistry,
|
||||||
|
authorizer: opts.Authorizer,
|
||||||
|
rollbackStore: opts.RollbackStore,
|
||||||
toolTimeout: toolTimeout,
|
toolTimeout: toolTimeout,
|
||||||
maxToolRounds: maxToolRounds,
|
maxToolRounds: maxToolRounds,
|
||||||
maxContextTokens: opts.MaxContextTokens,
|
maxContextTokens: opts.MaxContextTokens,
|
||||||
@@ -105,6 +124,8 @@ type RuntimeInfo struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) SetRuntime(opts RuntimeOptions) {
|
func (a *Agent) SetRuntime(opts RuntimeOptions) {
|
||||||
|
a.mu.Lock()
|
||||||
|
defer a.mu.Unlock()
|
||||||
if opts.Provider != nil {
|
if opts.Provider != nil {
|
||||||
a.provider = opts.Provider
|
a.provider = opts.Provider
|
||||||
}
|
}
|
||||||
@@ -121,6 +142,8 @@ func (a *Agent) SetRuntime(opts RuntimeOptions) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) RuntimeInfo() RuntimeInfo {
|
func (a *Agent) RuntimeInfo() RuntimeInfo {
|
||||||
|
a.mu.RLock()
|
||||||
|
defer a.mu.RUnlock()
|
||||||
return RuntimeInfo{
|
return RuntimeInfo{
|
||||||
ProviderName: a.providerName,
|
ProviderName: a.providerName,
|
||||||
Model: a.model,
|
Model: a.model,
|
||||||
@@ -132,6 +155,12 @@ func (a *Agent) RuntimeInfo() RuntimeInfo {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) Clear() {
|
func (a *Agent) Clear() {
|
||||||
|
a.mu.Lock()
|
||||||
|
defer a.mu.Unlock()
|
||||||
|
a.clear()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) clear() {
|
||||||
a.lastUsage = nil
|
a.lastUsage = nil
|
||||||
a.messages = nil
|
a.messages = nil
|
||||||
if strings.TrimSpace(a.systemPrompt) != "" {
|
if strings.TrimSpace(a.systemPrompt) != "" {
|
||||||
@@ -142,8 +171,27 @@ func (a *Agent) Clear() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// AddSystemInstructions makes newly installed instructions effective for the
|
||||||
|
// current conversation and for future Clear calls.
|
||||||
|
func (a *Agent) AddSystemInstructions(instructions string) {
|
||||||
|
instructions = strings.TrimSpace(instructions)
|
||||||
|
if instructions == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
a.mu.Lock()
|
||||||
|
defer a.mu.Unlock()
|
||||||
|
a.systemPrompt = strings.TrimSpace(a.systemPrompt + "\n\n" + instructions)
|
||||||
|
if len(a.messages) > 0 && a.messages[0].Role == llm.RoleSystem {
|
||||||
|
a.messages[0].Content = a.systemPrompt
|
||||||
|
return
|
||||||
|
}
|
||||||
|
a.messages = append([]llm.Message{{Role: llm.RoleSystem, Content: a.systemPrompt}}, a.messages...)
|
||||||
|
}
|
||||||
|
|
||||||
// Messages returns a copy of the current conversation messages.
|
// Messages returns a copy of the current conversation messages.
|
||||||
func (a *Agent) Messages() []llm.Message {
|
func (a *Agent) Messages() []llm.Message {
|
||||||
|
a.mu.RLock()
|
||||||
|
defer a.mu.RUnlock()
|
||||||
out := make([]llm.Message, len(a.messages))
|
out := make([]llm.Message, len(a.messages))
|
||||||
copy(out, a.messages)
|
copy(out, a.messages)
|
||||||
return out
|
return out
|
||||||
@@ -151,11 +199,23 @@ func (a *Agent) Messages() []llm.Message {
|
|||||||
|
|
||||||
// SetMessages replaces the conversation messages (used for session restore).
|
// SetMessages replaces the conversation messages (used for session restore).
|
||||||
func (a *Agent) SetMessages(messages []llm.Message) {
|
func (a *Agent) SetMessages(messages []llm.Message) {
|
||||||
|
a.mu.Lock()
|
||||||
|
defer a.mu.Unlock()
|
||||||
a.messages = append([]llm.Message(nil), messages...)
|
a.messages = append([]llm.Message(nil), messages...)
|
||||||
|
if strings.TrimSpace(a.systemPrompt) == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if len(a.messages) > 0 && a.messages[0].Role == llm.RoleSystem {
|
||||||
|
a.messages[0].Content = a.systemPrompt
|
||||||
|
return
|
||||||
|
}
|
||||||
|
a.messages = append([]llm.Message{{Role: llm.RoleSystem, Content: a.systemPrompt}}, a.messages...)
|
||||||
}
|
}
|
||||||
|
|
||||||
// LastUsage returns a copy of the most recent provider token usage data.
|
// LastUsage returns a copy of the most recent provider token usage data.
|
||||||
func (a *Agent) LastUsage() *llm.Usage {
|
func (a *Agent) LastUsage() *llm.Usage {
|
||||||
|
a.mu.RLock()
|
||||||
|
defer a.mu.RUnlock()
|
||||||
if a.lastUsage == nil {
|
if a.lastUsage == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -165,6 +225,12 @@ func (a *Agent) LastUsage() *llm.Usage {
|
|||||||
|
|
||||||
// SetLastUsage restores the most recent provider token usage data.
|
// SetLastUsage restores the most recent provider token usage data.
|
||||||
func (a *Agent) SetLastUsage(usage *llm.Usage) {
|
func (a *Agent) SetLastUsage(usage *llm.Usage) {
|
||||||
|
a.mu.Lock()
|
||||||
|
defer a.mu.Unlock()
|
||||||
|
a.setLastUsage(usage)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) setLastUsage(usage *llm.Usage) {
|
||||||
if usage == nil {
|
if usage == nil {
|
||||||
a.lastUsage = nil
|
a.lastUsage = nil
|
||||||
return
|
return
|
||||||
@@ -173,34 +239,90 @@ func (a *Agent) SetLastUsage(usage *llm.Usage) {
|
|||||||
a.lastUsage = ©
|
a.lastUsage = ©
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 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
|
// Compact summarizes older conversation messages while preserving the system
|
||||||
// prompt and recent messages intact.
|
// prompt and recent messages intact.
|
||||||
func (a *Agent) Compact(ctx context.Context) error {
|
func (a *Agent) Compact(ctx context.Context) error {
|
||||||
if len(a.messages) < 4 {
|
a.mu.RLock()
|
||||||
|
messages := append([]llm.Message(nil), a.messages...)
|
||||||
|
provider := a.provider
|
||||||
|
model := a.model
|
||||||
|
a.mu.RUnlock()
|
||||||
|
if len(messages) < 4 {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
if a.provider == nil {
|
if provider == nil {
|
||||||
return fmt.Errorf("provider is not configured")
|
return fmt.Errorf("provider is not configured")
|
||||||
}
|
}
|
||||||
|
|
||||||
prefixEnd := 0
|
prefixEnd := 0
|
||||||
if a.messages[0].Role == llm.RoleSystem {
|
if messages[0].Role == llm.RoleSystem {
|
||||||
prefixEnd = 1
|
prefixEnd = 1
|
||||||
}
|
}
|
||||||
if len(a.messages)-prefixEnd <= compactRecentMessages {
|
recentStart := compactRecentStart(messages, prefixEnd)
|
||||||
|
if recentStart <= prefixEnd || recentStart >= len(messages) {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
middle := append([]llm.Message(nil), messages[prefixEnd:recentStart]...)
|
||||||
recentStart := len(a.messages) - compactRecentMessages
|
|
||||||
if recentStart <= prefixEnd {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
middle := append([]llm.Message(nil), a.messages[prefixEnd:recentStart]...)
|
|
||||||
if len(middle) == 0 {
|
if len(middle) == 0 {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
summary, usage, err := a.summarizeMessages(ctx, middle)
|
summary, usage, err := summarizeMessages(ctx, provider, model, middle)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("compact context: %w", err)
|
return fmt.Errorf("compact context: %w", err)
|
||||||
}
|
}
|
||||||
@@ -209,13 +331,15 @@ func (a *Agent) Compact(ctx context.Context) error {
|
|||||||
return fmt.Errorf("compact context: summary was empty")
|
return fmt.Errorf("compact context: summary was empty")
|
||||||
}
|
}
|
||||||
|
|
||||||
compressed := make([]llm.Message, 0, prefixEnd+1+len(a.messages[recentStart:]))
|
compressed := make([]llm.Message, 0, prefixEnd+1+len(messages[recentStart:]))
|
||||||
compressed = append(compressed, a.messages[:prefixEnd]...)
|
compressed = append(compressed, messages[:prefixEnd]...)
|
||||||
compressed = append(compressed, llm.Message{
|
compressed = append(compressed, llm.Message{
|
||||||
Role: llm.RoleAssistant,
|
Role: llm.RoleAssistant,
|
||||||
Content: contextSummaryPrefix + "\n" + summary,
|
Content: contextSummaryPrefix + "\n" + summary,
|
||||||
})
|
})
|
||||||
compressed = append(compressed, a.messages[recentStart:]...)
|
compressed = append(compressed, messages[recentStart:]...)
|
||||||
|
a.mu.Lock()
|
||||||
|
defer a.mu.Unlock()
|
||||||
a.messages = compressed
|
a.messages = compressed
|
||||||
if usage != nil {
|
if usage != nil {
|
||||||
a.lastUsage = usage
|
a.lastUsage = usage
|
||||||
@@ -224,6 +348,8 @@ func (a *Agent) Compact(ctx context.Context) error {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) shouldCompact() bool {
|
func (a *Agent) shouldCompact() bool {
|
||||||
|
a.mu.RLock()
|
||||||
|
defer a.mu.RUnlock()
|
||||||
if a.maxContextTokens <= 0 {
|
if a.maxContextTokens <= 0 {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
@@ -233,9 +359,9 @@ func (a *Agent) shouldCompact() bool {
|
|||||||
return llm.EstimateTokens(a.messages, a.toolDefinitions()) >= a.maxContextTokens
|
return llm.EstimateTokens(a.messages, a.toolDefinitions()) >= a.maxContextTokens
|
||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (string, *llm.Usage, error) {
|
func summarizeMessages(ctx context.Context, provider llm.Provider, model string, messages []llm.Message) (string, *llm.Usage, error) {
|
||||||
req := llm.ChatRequest{
|
req := llm.ChatRequest{
|
||||||
Model: a.model,
|
Model: model,
|
||||||
Messages: []llm.Message{
|
Messages: []llm.Message{
|
||||||
{Role: llm.RoleSystem, Content: summarizerSystemPrompt},
|
{Role: llm.RoleSystem, Content: summarizerSystemPrompt},
|
||||||
{Role: llm.RoleUser, Content: formatMessagesForSummary(messages)},
|
{Role: llm.RoleUser, Content: formatMessagesForSummary(messages)},
|
||||||
@@ -244,7 +370,7 @@ func (a *Agent) summarizeMessages(ctx context.Context, messages []llm.Message) (
|
|||||||
|
|
||||||
var summary strings.Builder
|
var summary strings.Builder
|
||||||
var usage *llm.Usage
|
var usage *llm.Usage
|
||||||
if err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
|
if err := provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
|
||||||
if event.Content != "" {
|
if event.Content != "" {
|
||||||
summary.WriteString(event.Content)
|
summary.WriteString(event.Content)
|
||||||
}
|
}
|
||||||
@@ -292,12 +418,93 @@ func roundSignature(calls []llm.ToolCall) string {
|
|||||||
return b.String()
|
return b.String()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// BTW answers a temporary side question using a snapshot of the current
|
||||||
|
// conversation. It intentionally exposes no tools and does not mutate the
|
||||||
|
// conversation, usage, or compaction state.
|
||||||
|
func (a *Agent) BTW(ctx context.Context, input string, out io.Writer) error {
|
||||||
|
input = strings.TrimSpace(input)
|
||||||
|
if input == "" {
|
||||||
|
return fmt.Errorf("btw question is required")
|
||||||
|
}
|
||||||
|
|
||||||
|
a.mu.RLock()
|
||||||
|
provider := a.provider
|
||||||
|
model := a.model
|
||||||
|
messages := append([]llm.Message(nil), a.messages...)
|
||||||
|
thinking := a.thinking
|
||||||
|
thinkingParam := a.thinkingParam
|
||||||
|
thinkingEnabled := a.thinkingEnabled
|
||||||
|
a.mu.RUnlock()
|
||||||
|
if provider == nil {
|
||||||
|
return fmt.Errorf("provider is not configured")
|
||||||
|
}
|
||||||
|
|
||||||
|
// A snapshot taken while tools are running may end with unanswered tool
|
||||||
|
// calls, which providers reject. Keep only the latest complete boundary.
|
||||||
|
messages = completeConversationPrefix(messages)
|
||||||
|
messages = append(messages, llm.Message{Role: llm.RoleUser, Content: input})
|
||||||
|
messages = expandContinuationMessages(messages)
|
||||||
|
req := llm.ChatRequest{Model: model, Messages: messages}
|
||||||
|
if thinkingEnabled {
|
||||||
|
if thinkingParam == "" {
|
||||||
|
thinkingParam = "thinking"
|
||||||
|
}
|
||||||
|
req.Extra = map[string]any{thinkingParam: thinking}
|
||||||
|
}
|
||||||
|
|
||||||
|
return provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
|
||||||
|
if event.Content == "" || out == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
_, err := io.WriteString(out, event.Content)
|
||||||
|
return err
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func completeConversationPrefix(messages []llm.Message) []llm.Message {
|
||||||
|
safeEnd := 0
|
||||||
|
pending := make(map[string]struct{})
|
||||||
|
for i, msg := range messages {
|
||||||
|
if len(pending) > 0 && msg.Role != llm.RoleTool {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
switch {
|
||||||
|
case msg.Role == llm.RoleAssistant && len(msg.ToolCalls) > 0:
|
||||||
|
for _, call := range msg.ToolCalls {
|
||||||
|
pending[call.ID] = struct{}{}
|
||||||
|
}
|
||||||
|
case msg.Role == llm.RoleTool:
|
||||||
|
delete(pending, msg.ToolCallID)
|
||||||
|
}
|
||||||
|
if len(pending) == 0 {
|
||||||
|
safeEnd = i + 1
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return append([]llm.Message(nil), messages[:safeEnd]...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func expandContinuationMessages(messages []llm.Message) []llm.Message {
|
||||||
|
expanded := append([]llm.Message(nil), messages...)
|
||||||
|
for i := range expanded {
|
||||||
|
if expanded[i].Role == llm.RoleUser && strings.TrimSpace(expanded[i].Content) == "." {
|
||||||
|
expanded[i].Content = continuationPrompt
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return expanded
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) appendMessage(message llm.Message) {
|
||||||
|
a.mu.Lock()
|
||||||
|
a.messages = append(a.messages, message)
|
||||||
|
a.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
|
func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
|
||||||
input = strings.TrimSpace(input)
|
input = strings.TrimSpace(input)
|
||||||
if input == "" {
|
if input == "" {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input})
|
a.appendMessage(llm.Message{Role: llm.RoleUser, Content: input})
|
||||||
if a.shouldCompact() {
|
if a.shouldCompact() {
|
||||||
if err := a.Compact(ctx); err != nil {
|
if err := a.Compact(ctx); err != nil {
|
||||||
if logs != nil {
|
if logs != nil {
|
||||||
@@ -314,14 +521,14 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if len(calls) == 0 {
|
if len(calls) == 0 {
|
||||||
a.messages = append(a.messages, llm.Message{
|
a.appendMessage(llm.Message{
|
||||||
Role: llm.RoleAssistant,
|
Role: llm.RoleAssistant,
|
||||||
Content: content,
|
Content: content,
|
||||||
})
|
})
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
a.messages = append(a.messages, llm.Message{
|
a.appendMessage(llm.Message{
|
||||||
Role: llm.RoleAssistant,
|
Role: llm.RoleAssistant,
|
||||||
Content: content,
|
Content: content,
|
||||||
ToolCalls: calls,
|
ToolCalls: calls,
|
||||||
@@ -356,7 +563,7 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
|
|||||||
strings.Join(toolNames, ", "),
|
strings.Join(toolNames, ", "),
|
||||||
doomLoopThreshold,
|
doomLoopThreshold,
|
||||||
)
|
)
|
||||||
a.messages = append(a.messages, llm.Message{
|
a.appendMessage(llm.Message{
|
||||||
Role: llm.RoleAssistant,
|
Role: llm.RoleAssistant,
|
||||||
Content: msg,
|
Content: msg,
|
||||||
})
|
})
|
||||||
@@ -364,14 +571,15 @@ func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs i
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, call := range calls {
|
results := a.executeTools(ctx, calls, logs)
|
||||||
result := a.executeTool(ctx, call, logs)
|
for _, result := range results {
|
||||||
a.messages = append(a.messages, llm.Message{
|
a.appendMessage(llm.Message{
|
||||||
Role: llm.RoleTool,
|
Role: llm.RoleTool,
|
||||||
ToolCallID: call.ID,
|
ToolCallID: result.call.ID,
|
||||||
Content: result,
|
Content: result.output,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return fmt.Errorf("tool round limit reached after %d rounds; the task may be too complex for a single turn", a.maxToolRounds)
|
return fmt.Errorf("tool round limit reached after %d rounds; the task may be too complex for a single turn", a.maxToolRounds)
|
||||||
@@ -381,23 +589,33 @@ func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []l
|
|||||||
var content strings.Builder
|
var content strings.Builder
|
||||||
calls := make(map[int]*toolCallBuilder)
|
calls := make(map[int]*toolCallBuilder)
|
||||||
|
|
||||||
|
a.mu.RLock()
|
||||||
|
provider := a.provider
|
||||||
|
model := a.model
|
||||||
|
messages := append([]llm.Message(nil), a.messages...)
|
||||||
|
thinking := a.thinking
|
||||||
|
thinkingParam := a.thinkingParam
|
||||||
|
thinkingEnabled := a.thinkingEnabled
|
||||||
|
a.mu.RUnlock()
|
||||||
|
if provider == nil {
|
||||||
|
return "", nil, fmt.Errorf("provider is not configured")
|
||||||
|
}
|
||||||
req := llm.ChatRequest{
|
req := llm.ChatRequest{
|
||||||
Model: a.model,
|
Model: model,
|
||||||
Messages: append([]llm.Message(nil), a.messages...),
|
Messages: expandContinuationMessages(messages),
|
||||||
Tools: a.toolDefinitions(),
|
Tools: a.toolDefinitions(),
|
||||||
}
|
}
|
||||||
if a.thinkingEnabled {
|
if thinkingEnabled {
|
||||||
thinkingParam := a.thinkingParam
|
|
||||||
if thinkingParam == "" {
|
if thinkingParam == "" {
|
||||||
thinkingParam = "thinking"
|
thinkingParam = "thinking"
|
||||||
}
|
}
|
||||||
req.Extra = map[string]any{thinkingParam: a.thinking}
|
req.Extra = map[string]any{thinkingParam: thinking}
|
||||||
}
|
}
|
||||||
if len(req.Tools) > 0 {
|
if len(req.Tools) > 0 {
|
||||||
req.ToolChoice = "auto"
|
req.ToolChoice = "auto"
|
||||||
}
|
}
|
||||||
|
|
||||||
err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
|
err := provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
|
||||||
if event.Usage != nil {
|
if event.Usage != nil {
|
||||||
a.SetLastUsage(event.Usage)
|
a.SetLastUsage(event.Usage)
|
||||||
}
|
}
|
||||||
@@ -433,39 +651,224 @@ func (a *Agent) toolDefinitions() []llm.Tool {
|
|||||||
return a.toolRegistry.Definitions()
|
return a.toolRegistry.Definitions()
|
||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) string {
|
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) (string, bool) {
|
||||||
name := call.Function.Name
|
name := call.Function.Name
|
||||||
if logs != nil {
|
|
||||||
fmt.Fprintf(logs, "\n[tool] %s %s\n", name, compact(call.Function.Arguments, 400))
|
|
||||||
}
|
|
||||||
|
|
||||||
if a.toolRegistry == nil {
|
if a.toolRegistry == nil {
|
||||||
return tools.FormatError(fmt.Errorf("tool registry is disabled"))
|
return tools.FormatError(fmt.Errorf("tool registry is disabled")), false
|
||||||
}
|
}
|
||||||
tool, ok := a.toolRegistry.Get(name)
|
tool, ok := a.toolRegistry.Get(name)
|
||||||
if !ok {
|
if !ok {
|
||||||
return tools.FormatError(fmt.Errorf("unknown tool: %s", name))
|
return tools.FormatError(fmt.Errorf("unknown tool: %s", name)), false
|
||||||
}
|
}
|
||||||
|
|
||||||
args := strings.TrimSpace(call.Function.Arguments)
|
args := strings.TrimSpace(call.Function.Arguments)
|
||||||
if args == "" {
|
if args == "" {
|
||||||
args = "{}"
|
args = "{}"
|
||||||
}
|
}
|
||||||
|
raw := json.RawMessage(args)
|
||||||
|
|
||||||
|
rollbackNotes, err := a.prepareMutation(ctx, tool, call, raw, args)
|
||||||
|
if err != nil {
|
||||||
|
return tools.FormatError(err), false
|
||||||
|
}
|
||||||
|
|
||||||
toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
|
toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
output, err := tool.Execute(toolCtx, json.RawMessage(args))
|
|
||||||
|
if streaming, ok := tool.(tools.StreamingTool); ok {
|
||||||
|
output, streamed := a.executeStreamingTool(streaming, toolCtx, raw, name, args, logs)
|
||||||
|
return appendRollbackNotes(output, rollbackNotes), streamed
|
||||||
|
}
|
||||||
|
|
||||||
|
output, err := tool.Execute(toolCtx, raw)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if output != "" {
|
if output != "" {
|
||||||
return output + "\n" + tools.FormatError(err)
|
return appendRollbackNotes(output+"\n"+tools.FormatError(err), rollbackNotes), false
|
||||||
}
|
}
|
||||||
return tools.FormatError(err)
|
return appendRollbackNotes(tools.FormatError(err), rollbackNotes), false
|
||||||
}
|
}
|
||||||
if logs != nil {
|
return appendRollbackNotes(output, rollbackNotes), false
|
||||||
fmt.Fprintf(logs, "[tool] %s done\n", name)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (a *Agent) prepareMutation(ctx context.Context, tool tools.Tool, call llm.ToolCall, raw json.RawMessage, args string) ([]string, error) {
|
||||||
|
if !tools.IsMutatingCall(tool, raw) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var targets []string
|
||||||
|
if targetTool, ok := tool.(tools.TargetPathTool); ok {
|
||||||
|
var err error
|
||||||
|
targets, err = targetTool.TargetPaths(raw)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("resolve mutation targets: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
preview := compact(args, toolLogArgumentLimit)
|
||||||
|
if previewTool, ok := tool.(tools.PreviewableTool); ok {
|
||||||
|
value, err := previewTool.Preview(ctx, raw)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("preview mutation: %w", err)
|
||||||
|
}
|
||||||
|
preview = value
|
||||||
|
}
|
||||||
|
|
||||||
|
if a.authorizer != nil {
|
||||||
|
decision, err := a.authorizer.AuthorizeToolCall(ctx, ToolAuthorizationRequest{
|
||||||
|
ID: call.ID,
|
||||||
|
ToolName: call.Function.Name,
|
||||||
|
Arguments: args,
|
||||||
|
Preview: preview,
|
||||||
|
RollbackAvailable: a.rollbackStore != nil && len(targets) > 0,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("authorize %s: %w", call.Function.Name, err)
|
||||||
|
}
|
||||||
|
if decision != ToolAuthorizationAllow {
|
||||||
|
return nil, fmt.Errorf("tool call denied by user: %s", call.Function.Name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if a.rollbackStore == nil {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
notes := make([]string, 0, len(targets))
|
||||||
|
for _, path := range targets {
|
||||||
|
id, err := a.rollbackStore.Capture(call.Function.Name, path, path)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("capture rollback snapshot for %s: %w", path, err)
|
||||||
|
}
|
||||||
|
notes = append(notes, fmt.Sprintf("[rollback] %s %s", id, path))
|
||||||
|
}
|
||||||
|
return notes, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendRollbackNotes(output string, notes []string) string {
|
||||||
|
if len(notes) == 0 {
|
||||||
return output
|
return output
|
||||||
}
|
}
|
||||||
|
if strings.TrimSpace(output) == "" {
|
||||||
|
return strings.Join(notes, "\n")
|
||||||
|
}
|
||||||
|
return output + "\n" + strings.Join(notes, "\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) executeStreamingTool(streaming tools.StreamingTool, ctx context.Context, raw json.RawMessage, name, args string, logs io.Writer) (string, bool) {
|
||||||
|
var streamed int
|
||||||
|
emit := func(chunk string) error {
|
||||||
|
if chunk == "" || logs == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if streamed >= toolStreamLogLimit {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if len(chunk) > toolStreamLogLimit-streamed {
|
||||||
|
chunk = chunk[:toolStreamLogLimit-streamed]
|
||||||
|
streamed = toolStreamLogLimit
|
||||||
|
_, err := fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s\n[streaming output truncated after %d bytes]", name, compact(args, toolLogArgumentLimit), toolLogOutputMarker, chunk, toolStreamLogLimit)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
streamed += len(chunk)
|
||||||
|
_, err := fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s", name, compact(args, toolLogArgumentLimit), toolLogOutputMarker, chunk)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
output, err := streaming.ExecuteStream(ctx, raw, emit)
|
||||||
|
if logs != nil {
|
||||||
|
if err != nil {
|
||||||
|
_ = emit("\n" + tools.FormatError(err))
|
||||||
|
} else if streamed == 0 {
|
||||||
|
_ = emit("(no output)")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
if output != "" {
|
||||||
|
return output + "\n" + tools.FormatError(err), true
|
||||||
|
}
|
||||||
|
return tools.FormatError(err), true
|
||||||
|
}
|
||||||
|
return output, true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
|
||||||
|
if a.toolRegistry == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for _, call := range calls {
|
||||||
|
if tool, ok := a.toolRegistry.Get(call.Function.Name); ok {
|
||||||
|
args := strings.TrimSpace(call.Function.Arguments)
|
||||||
|
if args == "" {
|
||||||
|
args = "{}"
|
||||||
|
}
|
||||||
|
if tools.IsMutatingCall(tool, json.RawMessage(args)) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.Writer) []toolExecutionResult {
|
||||||
|
results := make([]toolExecutionResult, len(calls))
|
||||||
|
|
||||||
|
// Log all tool starts sequentially so TUI gets ordered placeholders.
|
||||||
|
if logs != nil {
|
||||||
|
for _, call := range calls {
|
||||||
|
logToolStart(logs, call)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if a.hasMutatingCall(calls) {
|
||||||
|
// Sequential: avoid racing mutating tools.
|
||||||
|
for i, call := range calls {
|
||||||
|
output, streamed := a.executeTool(ctx, call, logs)
|
||||||
|
results[i] = toolExecutionResult{call: call, output: output}
|
||||||
|
if logs != nil && !streamed {
|
||||||
|
logToolOutput(logs, call, output)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return results
|
||||||
|
}
|
||||||
|
|
||||||
|
// Concurrent: independent read-only tools.
|
||||||
|
var mu sync.Mutex
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for i, call := range calls {
|
||||||
|
wg.Add(1)
|
||||||
|
go func(idx int, c llm.ToolCall) {
|
||||||
|
defer wg.Done()
|
||||||
|
output, streamed := a.executeTool(ctx, c, logs)
|
||||||
|
results[idx] = toolExecutionResult{call: c, output: output}
|
||||||
|
if logs != nil && !streamed {
|
||||||
|
mu.Lock()
|
||||||
|
logToolOutput(logs, c, output)
|
||||||
|
mu.Unlock()
|
||||||
|
}
|
||||||
|
}(i, call)
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
return results
|
||||||
|
}
|
||||||
|
|
||||||
|
func logToolStart(logs io.Writer, call llm.ToolCall) {
|
||||||
|
fmt.Fprintf(logs, "\n[tool] %s %s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit))
|
||||||
|
}
|
||||||
|
|
||||||
|
func logToolOutput(logs io.Writer, call llm.ToolCall, output string) {
|
||||||
|
fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s\n", call.Function.Name, compact(call.Function.Arguments, toolLogArgumentLimit), toolLogOutputMarker, toolOutputPreview(output))
|
||||||
|
}
|
||||||
|
|
||||||
|
func toolOutputPreview(output string) string {
|
||||||
|
output = strings.TrimSpace(output)
|
||||||
|
if output == "" {
|
||||||
|
return "(no output)"
|
||||||
|
}
|
||||||
|
if len(output) <= toolLogOutputPreviewBytes {
|
||||||
|
return output
|
||||||
|
}
|
||||||
|
return output[:toolLogOutputPreviewBytes] + fmt.Sprintf("\n[truncated: showing first %d bytes; omitted %d bytes]", toolLogOutputPreviewBytes, len(output)-toolLogOutputPreviewBytes)
|
||||||
|
}
|
||||||
|
|
||||||
type toolCallBuilder struct {
|
type toolCallBuilder struct {
|
||||||
index int
|
index int
|
||||||
|
|||||||
@@ -3,12 +3,17 @@ package agent
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"reflect"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"agentu/pkg/llm"
|
|
||||||
"agentu/internal/tools"
|
"agentu/internal/tools"
|
||||||
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
type fakeProvider struct {
|
type fakeProvider struct {
|
||||||
@@ -87,6 +92,86 @@ func TestAgentRunsToolLoop(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type streamToolProvider struct {
|
||||||
|
calls int
|
||||||
|
t *testing.T
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *streamToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
p.calls++
|
||||||
|
switch p.calls {
|
||||||
|
case 1:
|
||||||
|
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
|
||||||
|
{Index: 0, ID: "call_1", Type: "function", Name: "test_stream", Arguments: `{}`},
|
||||||
|
}})
|
||||||
|
case 2:
|
||||||
|
last := req.Messages[len(req.Messages)-1]
|
||||||
|
if last.Role != llm.RoleTool || last.Content != "full result" {
|
||||||
|
p.t.Fatalf("last message = %#v", last)
|
||||||
|
}
|
||||||
|
return emit(llm.StreamEvent{Content: "done"})
|
||||||
|
default:
|
||||||
|
p.t.Fatalf("unexpected call count %d", p.calls)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type streamEchoTool struct{}
|
||||||
|
|
||||||
|
func (streamEchoTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "test_stream",
|
||||||
|
Description: "streaming test",
|
||||||
|
Parameters: json.RawMessage(`{"type":"object"}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (streamEchoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
return "full result", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (streamEchoTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
|
||||||
|
if err := emit("one\n"); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if err := emit("two\n"); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return "full result", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentStreamsToolOutputToLogs(t *testing.T) {
|
||||||
|
provider := &streamToolProvider{t: t}
|
||||||
|
a := New(Options{
|
||||||
|
Provider: provider,
|
||||||
|
Model: "test",
|
||||||
|
SystemPrompt: "system",
|
||||||
|
ToolRegistry: tools.NewRegistry(streamEchoTool{}),
|
||||||
|
})
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
var logs strings.Builder
|
||||||
|
if err := a.RunTurn(context.Background(), "stream", &out, &logs); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if out.String() != "done" {
|
||||||
|
t.Fatalf("out = %q", out.String())
|
||||||
|
}
|
||||||
|
logText := logs.String()
|
||||||
|
if strings.Count(logText, "[output]") != 2 {
|
||||||
|
t.Fatalf("logs = %q", logText)
|
||||||
|
}
|
||||||
|
if !strings.Contains(logText, "one\n") || !strings.Contains(logText, "two\n") {
|
||||||
|
t.Fatalf("logs missing streamed chunks: %q", logText)
|
||||||
|
}
|
||||||
|
if strings.Contains(logText, "[output]\nfull result") {
|
||||||
|
t.Fatalf("logs contain duplicate full output: %q", logText)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
type contentProvider struct {
|
type contentProvider struct {
|
||||||
called bool
|
called bool
|
||||||
text string
|
text string
|
||||||
@@ -97,6 +182,173 @@ func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, e
|
|||||||
return emit(llm.StreamEvent{Content: p.text})
|
return emit(llm.StreamEvent{Content: p.text})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type btwProvider struct {
|
||||||
|
req llm.ChatRequest
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *btwProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
p.req = req
|
||||||
|
if p.err != nil {
|
||||||
|
return p.err
|
||||||
|
}
|
||||||
|
return emit(llm.StreamEvent{Content: "side answer", Usage: &llm.Usage{TotalTokens: 999}})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBTWUsesContextWithoutMutatingConversation(t *testing.T) {
|
||||||
|
provider := &btwProvider{}
|
||||||
|
a := New(Options{Provider: provider, Model: "test", Thinking: "high", ThinkingEnabled: true})
|
||||||
|
a.SetMessages([]llm.Message{
|
||||||
|
{Role: llm.RoleUser, Content: "main question"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "main answer"},
|
||||||
|
})
|
||||||
|
a.SetLastUsage(&llm.Usage{TotalTokens: 7})
|
||||||
|
before := a.Messages()
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
if err := a.BTW(context.Background(), "what does that mean?", &out); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if out.String() != "side answer" {
|
||||||
|
t.Fatalf("output = %q", out.String())
|
||||||
|
}
|
||||||
|
if !reflect.DeepEqual(a.Messages(), before) {
|
||||||
|
t.Fatalf("conversation changed:\nbefore: %+v\nafter: %+v", before, a.Messages())
|
||||||
|
}
|
||||||
|
if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 7 {
|
||||||
|
t.Fatalf("usage changed: %+v", usage)
|
||||||
|
}
|
||||||
|
if len(provider.req.Tools) != 0 || provider.req.ToolChoice != "" {
|
||||||
|
t.Fatalf("btw exposed tools: %+v / %q", provider.req.Tools, provider.req.ToolChoice)
|
||||||
|
}
|
||||||
|
last := provider.req.Messages[len(provider.req.Messages)-1]
|
||||||
|
if last.Role != llm.RoleUser || last.Content != "what does that mean?" {
|
||||||
|
t.Fatalf("last request message = %+v", last)
|
||||||
|
}
|
||||||
|
if provider.req.Extra["thinking"] != "high" {
|
||||||
|
t.Fatalf("thinking extra = %+v", provider.req.Extra)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBTWRejectsEmptyQuestionAndReturnsProviderError(t *testing.T) {
|
||||||
|
provider := &btwProvider{err: errors.New("boom")}
|
||||||
|
a := New(Options{Provider: provider, Model: "test"})
|
||||||
|
if err := a.BTW(context.Background(), " ", nil); err == nil {
|
||||||
|
t.Fatal("empty question did not fail")
|
||||||
|
}
|
||||||
|
if err := a.BTW(context.Background(), "question", nil); err == nil || !strings.Contains(err.Error(), "boom") {
|
||||||
|
t.Fatalf("provider error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDotContinuesCurrentTaskWithoutRewritingHistory(t *testing.T) {
|
||||||
|
provider := &btwProvider{}
|
||||||
|
a := New(Options{Provider: provider, Model: "test"})
|
||||||
|
if err := a.RunTurn(context.Background(), " . ", nil, nil); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
requestMessage := provider.req.Messages[len(provider.req.Messages)-1]
|
||||||
|
if requestMessage.Content != continuationPrompt {
|
||||||
|
t.Fatalf("request content = %q, want %q", requestMessage.Content, continuationPrompt)
|
||||||
|
}
|
||||||
|
history := a.Messages()
|
||||||
|
if history[0].Content != "." {
|
||||||
|
t.Fatalf("stored input = %q, want literal dot", history[0].Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestExpandContinuationMessagesOnlyExpandsSingleDot(t *testing.T) {
|
||||||
|
messages := []llm.Message{
|
||||||
|
{Role: llm.RoleUser, Content: "."},
|
||||||
|
{Role: llm.RoleUser, Content: " . "},
|
||||||
|
{Role: llm.RoleUser, Content: ".."},
|
||||||
|
{Role: llm.RoleUser, Content: ". continue"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "."},
|
||||||
|
}
|
||||||
|
expanded := expandContinuationMessages(messages)
|
||||||
|
if expanded[0].Content != continuationPrompt || expanded[1].Content != continuationPrompt {
|
||||||
|
t.Fatalf("single dots were not expanded: %+v", expanded)
|
||||||
|
}
|
||||||
|
for _, index := range []int{2, 3, 4} {
|
||||||
|
if expanded[index].Content != messages[index].Content {
|
||||||
|
t.Fatalf("message %d changed from %q to %q", index, messages[index].Content, expanded[index].Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if messages[0].Content != "." || messages[1].Content != " . " {
|
||||||
|
t.Fatalf("source messages were mutated: %+v", messages)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCompleteConversationPrefixDropsPendingToolRound(t *testing.T) {
|
||||||
|
messages := []llm.Message{
|
||||||
|
{Role: llm.RoleUser, Content: "first"},
|
||||||
|
{Role: llm.RoleAssistant, Content: "answer"},
|
||||||
|
{Role: llm.RoleUser, Content: "main"},
|
||||||
|
{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "call_1"}, {ID: "call_2"}}},
|
||||||
|
{Role: llm.RoleTool, ToolCallID: "call_1", Content: "one"},
|
||||||
|
}
|
||||||
|
got := completeConversationPrefix(messages)
|
||||||
|
if len(got) != 3 || got[len(got)-1].Content != "main" {
|
||||||
|
t.Fatalf("complete prefix = %+v", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type concurrentBTWProvider struct {
|
||||||
|
mainStarted chan struct{}
|
||||||
|
releaseMain chan struct{}
|
||||||
|
once sync.Once
|
||||||
|
mu sync.Mutex
|
||||||
|
btwRequest llm.ChatRequest
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *concurrentBTWProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
last := req.Messages[len(req.Messages)-1].Content
|
||||||
|
if last == "main task" {
|
||||||
|
p.once.Do(func() { close(p.mainStarted) })
|
||||||
|
select {
|
||||||
|
case <-p.releaseMain:
|
||||||
|
return emit(llm.StreamEvent{Content: "main answer"})
|
||||||
|
case <-ctx.Done():
|
||||||
|
return ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
p.mu.Lock()
|
||||||
|
p.btwRequest = req
|
||||||
|
p.mu.Unlock()
|
||||||
|
return emit(llm.StreamEvent{Content: "side answer"})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBTWCanRunDuringMainTurn(t *testing.T) {
|
||||||
|
provider := &concurrentBTWProvider{mainStarted: make(chan struct{}), releaseMain: make(chan struct{})}
|
||||||
|
a := New(Options{Provider: provider, Model: "test"})
|
||||||
|
mainDone := make(chan error, 1)
|
||||||
|
go func() { mainDone <- a.RunTurn(context.Background(), "main task", nil, nil) }()
|
||||||
|
<-provider.mainStarted
|
||||||
|
|
||||||
|
var side strings.Builder
|
||||||
|
if err := a.BTW(context.Background(), "side question", &side); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if side.String() != "side answer" {
|
||||||
|
t.Fatalf("side output = %q", side.String())
|
||||||
|
}
|
||||||
|
provider.mu.Lock()
|
||||||
|
req := provider.btwRequest
|
||||||
|
provider.mu.Unlock()
|
||||||
|
if len(req.Messages) < 2 || req.Messages[len(req.Messages)-2].Content != "main task" {
|
||||||
|
t.Fatalf("btw request did not include active main turn: %+v", req.Messages)
|
||||||
|
}
|
||||||
|
|
||||||
|
close(provider.releaseMain)
|
||||||
|
if err := <-mainDone; err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(a.Messages()) != 2 {
|
||||||
|
t.Fatalf("main conversation contains btw messages: %+v", a.Messages())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
type compactProvider struct {
|
type compactProvider struct {
|
||||||
summaryCalls int
|
summaryCalls int
|
||||||
answerCalls int
|
answerCalls int
|
||||||
@@ -145,6 +397,9 @@ func TestAgentCompact(t *testing.T) {
|
|||||||
if messages[len(messages)-1].Content != "recent-answer-2" {
|
if messages[len(messages)-1].Content != "recent-answer-2" {
|
||||||
t.Fatalf("last recent message = %#v", messages[len(messages)-1])
|
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 {
|
if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 12 {
|
||||||
t.Fatalf("usage = %#v", usage)
|
t.Fatalf("usage = %#v", usage)
|
||||||
}
|
}
|
||||||
@@ -217,6 +472,48 @@ func longConversation() []llm.Message {
|
|||||||
return messages
|
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 {
|
func joinMessageContents(messages []llm.Message) string {
|
||||||
var b strings.Builder
|
var b strings.Builder
|
||||||
for _, msg := range messages {
|
for _, msg := range messages {
|
||||||
@@ -340,3 +637,288 @@ func TestAgentHardLimitStopsVaryingCalls(t *testing.T) {
|
|||||||
t.Fatalf("calls = %d, want %d", provider.calls, limit)
|
t.Fatalf("calls = %d, want %d", provider.calls, limit)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// slowTool blocks on a channel before returning.
|
||||||
|
type slowTool struct {
|
||||||
|
name string
|
||||||
|
started chan<- string
|
||||||
|
release <-chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *slowTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: t.name,
|
||||||
|
Description: "slow test tool",
|
||||||
|
Parameters: json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}}}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *slowTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args struct {
|
||||||
|
Text string `json:"text"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
t.started <- t.name
|
||||||
|
<-t.release
|
||||||
|
return args.Text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// mutatingSlowTool is like slowTool but implements MutatingTool.
|
||||||
|
type mutatingSlowTool struct {
|
||||||
|
slowTool
|
||||||
|
}
|
||||||
|
|
||||||
|
func (*mutatingSlowTool) Mutates() bool { return true }
|
||||||
|
|
||||||
|
func TestAgentExecutesIndependentToolCallsInParallel(t *testing.T) {
|
||||||
|
started := make(chan string, 2)
|
||||||
|
release := make(chan struct{})
|
||||||
|
|
||||||
|
tool1 := &slowTool{name: "slow_one", started: started, release: release}
|
||||||
|
tool2 := &slowTool{name: "slow_two", started: started, release: release}
|
||||||
|
|
||||||
|
provider := ¶llelProvider{
|
||||||
|
responses: [][]llm.ToolCallDelta{
|
||||||
|
{
|
||||||
|
{Index: 0, ID: "call_1", Type: "function", Name: "slow_one", Arguments: `{"text":"one"}`},
|
||||||
|
{Index: 1, ID: "call_2", Type: "function", Name: "slow_two", Arguments: `{"text":"two"}`},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
finalContent: "done",
|
||||||
|
}
|
||||||
|
|
||||||
|
a := New(Options{
|
||||||
|
Provider: provider,
|
||||||
|
Model: "test",
|
||||||
|
SystemPrompt: "system",
|
||||||
|
ToolRegistry: tools.NewRegistry(tool1, tool2),
|
||||||
|
})
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
done <- a.RunTurn(context.Background(), "run both", &out, nil)
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Both tools must start before we release them.
|
||||||
|
names := make(map[string]bool)
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
select {
|
||||||
|
case name := <-started:
|
||||||
|
names[name] = true
|
||||||
|
case <-make(chan struct{}):
|
||||||
|
// Use a timer in real code, but for test simplicity we rely on deadlock detection.
|
||||||
|
t.Fatal("timed out waiting for tool to start")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !names["slow_one"] || !names["slow_two"] {
|
||||||
|
t.Fatalf("expected both tools to start, got %v", names)
|
||||||
|
}
|
||||||
|
close(release)
|
||||||
|
|
||||||
|
if err := <-done; err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if out.String() != "done" {
|
||||||
|
t.Fatalf("out = %q", out.String())
|
||||||
|
}
|
||||||
|
// Verify tool results were in order.
|
||||||
|
msgs := a.Messages()
|
||||||
|
var toolMsgs []string
|
||||||
|
for _, m := range msgs {
|
||||||
|
if m.Role == llm.RoleTool {
|
||||||
|
toolMsgs = append(toolMsgs, m.Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(toolMsgs) != 2 || toolMsgs[0] != "one" || toolMsgs[1] != "two" {
|
||||||
|
t.Fatalf("tool messages = %v, want [one two]", toolMsgs)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentRunsMutatingToolCallsSequentially(t *testing.T) {
|
||||||
|
started := make(chan string, 2)
|
||||||
|
release := make(chan struct{})
|
||||||
|
|
||||||
|
tool1 := &mutatingSlowTool{slowTool{name: "mut_one", started: started, release: release}}
|
||||||
|
tool2 := &mutatingSlowTool{slowTool{name: "mut_two", started: started, release: release}}
|
||||||
|
|
||||||
|
provider := ¶llelProvider{
|
||||||
|
responses: [][]llm.ToolCallDelta{
|
||||||
|
{
|
||||||
|
{Index: 0, ID: "call_1", Type: "function", Name: "mut_one", Arguments: `{"text":"one"}`},
|
||||||
|
{Index: 1, ID: "call_2", Type: "function", Name: "mut_two", Arguments: `{"text":"two"}`},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
finalContent: "done",
|
||||||
|
}
|
||||||
|
|
||||||
|
a := New(Options{
|
||||||
|
Provider: provider,
|
||||||
|
Model: "test",
|
||||||
|
SystemPrompt: "system",
|
||||||
|
ToolRegistry: tools.NewRegistry(tool1, tool2),
|
||||||
|
})
|
||||||
|
|
||||||
|
var out strings.Builder
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
done <- a.RunTurn(context.Background(), "run both", &out, nil)
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Only the first tool should start.
|
||||||
|
select {
|
||||||
|
case name := <-started:
|
||||||
|
if name != "mut_one" {
|
||||||
|
t.Fatalf("expected mut_one to start first, got %q", name)
|
||||||
|
}
|
||||||
|
case <-make(chan struct{}):
|
||||||
|
t.Fatal("timed out waiting for first tool to start")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Second tool should NOT have started yet.
|
||||||
|
select {
|
||||||
|
case name := <-started:
|
||||||
|
t.Fatalf("second tool started too early: %q", name)
|
||||||
|
default:
|
||||||
|
// Expected: second tool hasn't started.
|
||||||
|
}
|
||||||
|
|
||||||
|
close(release)
|
||||||
|
if err := <-done; err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// parallelProvider emits a batch of tool calls on the first request, then returns finalContent.
|
||||||
|
type parallelProvider struct {
|
||||||
|
responses [][]llm.ToolCallDelta
|
||||||
|
finalContent string
|
||||||
|
calls int
|
||||||
|
}
|
||||||
|
|
||||||
|
type authorizationProvider struct {
|
||||||
|
call llm.ToolCallDelta
|
||||||
|
calls int
|
||||||
|
toolOutput string
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *authorizationProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
p.calls++
|
||||||
|
if p.calls == 1 {
|
||||||
|
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{p.call}})
|
||||||
|
}
|
||||||
|
p.toolOutput = req.Messages[len(req.Messages)-1].Content
|
||||||
|
return emit(llm.StreamEvent{Content: "done"})
|
||||||
|
}
|
||||||
|
|
||||||
|
type recordingAuthorizer struct {
|
||||||
|
decision ToolAuthorizationDecision
|
||||||
|
requests []ToolAuthorizationRequest
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *recordingAuthorizer) AuthorizeToolCall(ctx context.Context, req ToolAuthorizationRequest) (ToolAuthorizationDecision, error) {
|
||||||
|
a.requests = append(a.requests, req)
|
||||||
|
return a.decision, ctx.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentAuthorizesMutationAndCreatesRollbackSnapshot(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
provider := &authorizationProvider{call: llm.ToolCallDelta{
|
||||||
|
Index: 0, ID: "write_1", Type: "function", Name: "file_write",
|
||||||
|
Arguments: `{"path":"new.txt","content":"hello"}`,
|
||||||
|
}}
|
||||||
|
authorizer := &recordingAuthorizer{decision: ToolAuthorizationAllow}
|
||||||
|
rollback := tools.NewRollbackStore()
|
||||||
|
a := New(Options{
|
||||||
|
Provider: provider,
|
||||||
|
Model: "test",
|
||||||
|
ToolRegistry: tools.Builtins(dir),
|
||||||
|
Authorizer: authorizer,
|
||||||
|
RollbackStore: rollback,
|
||||||
|
})
|
||||||
|
|
||||||
|
if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(authorizer.requests) != 1 || !authorizer.requests[0].RollbackAvailable {
|
||||||
|
t.Fatalf("authorization requests = %+v", authorizer.requests)
|
||||||
|
}
|
||||||
|
if !strings.Contains(authorizer.requests[0].Preview, "write 5 bytes to new.txt") {
|
||||||
|
t.Fatalf("preview = %q", authorizer.requests[0].Preview)
|
||||||
|
}
|
||||||
|
rollbackID := rollback.LastID()
|
||||||
|
if rollbackID == "" || !strings.Contains(provider.toolOutput, rollbackID) {
|
||||||
|
t.Fatalf("tool output %q missing rollback ID %q", provider.toolOutput, rollbackID)
|
||||||
|
}
|
||||||
|
if _, err := rollback.Restore(rollbackID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(filepath.Join(dir, "new.txt")); !os.IsNotExist(err) {
|
||||||
|
t.Fatalf("new file still exists after rollback: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentDeniedMutationDoesNotExecute(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
provider := &authorizationProvider{call: llm.ToolCallDelta{
|
||||||
|
Index: 0, ID: "write_1", Type: "function", Name: "file_write",
|
||||||
|
Arguments: `{"path":"denied.txt","content":"no"}`,
|
||||||
|
}}
|
||||||
|
authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
|
||||||
|
a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})
|
||||||
|
|
||||||
|
if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(filepath.Join(dir, "denied.txt")); !os.IsNotExist(err) {
|
||||||
|
t.Fatalf("denied file exists: %v", err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(provider.toolOutput, "denied by user") {
|
||||||
|
t.Fatalf("tool output = %q", provider.toolOutput)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentDoesNotAuthorizeApplyPatchDryRun(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
provider := &authorizationProvider{call: llm.ToolCallDelta{
|
||||||
|
Index: 0, ID: "patch_1", Type: "function", Name: "apply_patch",
|
||||||
|
Arguments: `{"patch":"*** Begin Patch\n*** Add File: dry.txt\n+dry\n*** End Patch","dry_run":true}`,
|
||||||
|
}}
|
||||||
|
authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
|
||||||
|
a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})
|
||||||
|
|
||||||
|
if err := a.RunTurn(context.Background(), "preview", nil, nil); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(authorizer.requests) != 0 {
|
||||||
|
t.Fatalf("dry run requested authorization: %+v", authorizer.requests)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(filepath.Join(dir, "dry.txt")); !os.IsNotExist(err) {
|
||||||
|
t.Fatalf("dry run created a file: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSetMessagesRefreshesCurrentSystemInstructions(t *testing.T) {
|
||||||
|
a := New(Options{SystemPrompt: "current instructions"})
|
||||||
|
a.SetMessages([]llm.Message{
|
||||||
|
{Role: llm.RoleSystem, Content: "old instructions"},
|
||||||
|
{Role: llm.RoleUser, Content: "hello"},
|
||||||
|
})
|
||||||
|
messages := a.Messages()
|
||||||
|
if messages[0].Content != "current instructions" || messages[1].Content != "hello" {
|
||||||
|
t.Fatalf("messages = %+v", messages)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *parallelProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
p.calls++
|
||||||
|
if p.calls <= len(p.responses) {
|
||||||
|
return emit(llm.StreamEvent{ToolCalls: p.responses[p.calls-1]})
|
||||||
|
}
|
||||||
|
return emit(llm.StreamEvent{Content: p.finalContent})
|
||||||
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -0,0 +1,113 @@
|
|||||||
|
package app
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agentu/internal/agent"
|
||||||
|
"agentu/internal/session"
|
||||||
|
"agentu/internal/skills"
|
||||||
|
"agentu/pkg/config"
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
type startupProvider struct{}
|
||||||
|
|
||||||
|
func (p *startupProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
return emit(llm.StreamEvent{Content: "ok"})
|
||||||
|
}
|
||||||
|
|
||||||
|
func startupTestManager(t *testing.T) (*session.Manager, *agent.Agent) {
|
||||||
|
t.Helper()
|
||||||
|
cfg := &config.Config{
|
||||||
|
Providers: map[string]config.ProviderConfig{
|
||||||
|
"one": {
|
||||||
|
Name: "one",
|
||||||
|
BaseURL: "https://one.example.com",
|
||||||
|
APIKey: "sk-test",
|
||||||
|
Model: "model-a",
|
||||||
|
Models: []string{"model-a"},
|
||||||
|
Thinking: "none",
|
||||||
|
ThinkingParam: "thinking",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
assistant := agent.New(agent.Options{
|
||||||
|
Provider: &startupProvider{},
|
||||||
|
ProviderName: "one",
|
||||||
|
Model: "model-a",
|
||||||
|
})
|
||||||
|
store, err := session.NewStore(filepath.Join(t.TempDir(), "sessions"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
manager, err := session.NewManager(cfg, assistant, nil, store)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return manager, assistant
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInitializeSessionStartsNewByDefault(t *testing.T) {
|
||||||
|
manager, assistant := startupTestManager(t)
|
||||||
|
if _, err := manager.NewSession(); err != nil {
|
||||||
|
t.Fatalf("new session: %v", err)
|
||||||
|
}
|
||||||
|
savedID := manager.CurrentSessionID()
|
||||||
|
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: "old context"}})
|
||||||
|
if err := manager.Save(); err != nil {
|
||||||
|
t.Fatalf("save: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := initializeSession(manager, ""); err != nil {
|
||||||
|
t.Fatalf("initialize: %v", err)
|
||||||
|
}
|
||||||
|
if manager.CurrentSessionID() == savedID {
|
||||||
|
t.Fatalf("default startup resumed latest session %s", savedID)
|
||||||
|
}
|
||||||
|
for _, msg := range assistant.Messages() {
|
||||||
|
if msg.Content == "old context" {
|
||||||
|
t.Fatal("default startup carried old context into new session")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInitializeSessionResumesExplicitID(t *testing.T) {
|
||||||
|
manager, assistant := startupTestManager(t)
|
||||||
|
if _, err := manager.NewSession(); err != nil {
|
||||||
|
t.Fatalf("new session: %v", err)
|
||||||
|
}
|
||||||
|
savedID := manager.CurrentSessionID()
|
||||||
|
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: "old context"}})
|
||||||
|
if err := manager.Save(); err != nil {
|
||||||
|
t.Fatalf("save: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := initializeSession(manager, savedID); err != nil {
|
||||||
|
t.Fatalf("initialize: %v", err)
|
||||||
|
}
|
||||||
|
if manager.CurrentSessionID() != savedID {
|
||||||
|
t.Fatalf("expected resumed session %s, got %s", savedID, manager.CurrentSessionID())
|
||||||
|
}
|
||||||
|
found := false
|
||||||
|
for _, msg := range assistant.Messages() {
|
||||||
|
if msg.Content == "old context" {
|
||||||
|
found = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !found {
|
||||||
|
t.Fatal("explicit resume did not restore saved context")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSystemPromptIncludesGlobalSkillsAndApprovalMode(t *testing.T) {
|
||||||
|
cfg := &config.Config{Agent: config.AgentConfig{SystemPrompt: "base", WorkingDir: "/tmp/project"}}
|
||||||
|
prompt := systemPrompt(cfg, true, false, []skills.Skill{{Name: "review", Path: "/tmp/review/SKILL.md", Content: "Review carefully."}})
|
||||||
|
for _, want := range []string{"base", "interactive user approval", "### review", "Review carefully."} {
|
||||||
|
if !strings.Contains(prompt, want) {
|
||||||
|
t.Fatalf("prompt missing %q:\n%s", want, prompt)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -68,11 +68,29 @@ func (m *Manager) Resume(id string) error {
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
m.currentSession = sess
|
m.currentSession = sess
|
||||||
|
m.restoreRuntime(sess)
|
||||||
m.agent.SetMessages(sess.Messages)
|
m.agent.SetMessages(sess.Messages)
|
||||||
m.agent.SetLastUsage(sess.LastUsage)
|
m.agent.SetLastUsage(sess.LastUsage)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// restoreRuntime switches the manager and agent back to the provider, model,
|
||||||
|
// and thinking level recorded on a resumed session. Values that are no longer
|
||||||
|
// valid for the current config are skipped and the defaults remain active.
|
||||||
|
func (m *Manager) restoreRuntime(sess *Session) {
|
||||||
|
if sess.Provider != "" && sess.Provider != m.providerName {
|
||||||
|
if _, ok := m.cfg.Providers[sess.Provider]; ok {
|
||||||
|
_, _ = m.SetProvider(sess.Provider)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if sess.Model != "" && sess.Model != m.provider.Model {
|
||||||
|
_, _ = m.SetModel(sess.Model)
|
||||||
|
}
|
||||||
|
if sess.Thinking != "" && sess.Thinking != m.provider.Thinking {
|
||||||
|
_, _ = m.SetThinking(sess.Thinking)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Save persists the current session to disk.
|
// Save persists the current session to disk.
|
||||||
func (m *Manager) Save() error {
|
func (m *Manager) Save() error {
|
||||||
if m.store == nil || m.currentSession == nil {
|
if m.store == nil || m.currentSession == nil {
|
||||||
@@ -82,6 +100,8 @@ func (m *Manager) Save() error {
|
|||||||
m.currentSession.LastUsage = m.agent.LastUsage()
|
m.currentSession.LastUsage = m.agent.LastUsage()
|
||||||
m.currentSession.Provider = m.providerName
|
m.currentSession.Provider = m.providerName
|
||||||
m.currentSession.Model = m.provider.Model
|
m.currentSession.Model = m.provider.Model
|
||||||
|
m.currentSession.APIType = m.provider.APIType
|
||||||
|
m.currentSession.Thinking = m.provider.Thinking
|
||||||
// Auto-name if empty
|
// Auto-name if empty
|
||||||
if m.currentSession.Name == "" {
|
if m.currentSession.Name == "" {
|
||||||
m.currentSession.Name = DefaultName(m.currentSession.Messages)
|
m.currentSession.Name = DefaultName(m.currentSession.Messages)
|
||||||
@@ -150,6 +170,8 @@ func (m *Manager) NewSession() (string, error) {
|
|||||||
m.currentSession.Messages = msgs
|
m.currentSession.Messages = msgs
|
||||||
m.currentSession.Provider = m.providerName
|
m.currentSession.Provider = m.providerName
|
||||||
m.currentSession.Model = m.provider.Model
|
m.currentSession.Model = m.provider.Model
|
||||||
|
m.currentSession.APIType = m.provider.APIType
|
||||||
|
m.currentSession.Thinking = m.provider.Thinking
|
||||||
if m.currentSession.Name == "" {
|
if m.currentSession.Name == "" {
|
||||||
m.currentSession.Name = DefaultName(msgs)
|
m.currentSession.Name = DefaultName(msgs)
|
||||||
}
|
}
|
||||||
@@ -204,10 +226,18 @@ func (m *Manager) CurrentThinking() string {
|
|||||||
return m.provider.Thinking
|
return m.provider.Thinking
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (m *Manager) CurrentAPIType() string {
|
||||||
|
if m.provider.APIType == "" {
|
||||||
|
return config.APITypeChat
|
||||||
|
}
|
||||||
|
return m.provider.APIType
|
||||||
|
}
|
||||||
|
|
||||||
func (m *Manager) Info() string {
|
func (m *Manager) Info() string {
|
||||||
var b strings.Builder
|
var b strings.Builder
|
||||||
fmt.Fprintf(&b, "provider: %s\n", m.CurrentProvider())
|
fmt.Fprintf(&b, "provider: %s\n", m.CurrentProvider())
|
||||||
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
|
fmt.Fprintf(&b, "model: %s\n", m.CurrentModel())
|
||||||
|
fmt.Fprintf(&b, "api: %s\n", m.CurrentAPIType())
|
||||||
fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking())
|
fmt.Fprintf(&b, "thinking: %s\n", m.CurrentThinking())
|
||||||
fmt.Fprintf(&b, "providers: %s\n", strings.Join(m.providerNames(), ", "))
|
fmt.Fprintf(&b, "providers: %s\n", strings.Join(m.providerNames(), ", "))
|
||||||
fmt.Fprintf(&b, "models: %s\n", strings.Join(m.provider.Models, ", "))
|
fmt.Fprintf(&b, "models: %s\n", strings.Join(m.provider.Models, ", "))
|
||||||
@@ -223,7 +253,7 @@ func (m *Manager) SetProvider(name string) (string, error) {
|
|||||||
}
|
}
|
||||||
m.providerName = name
|
m.providerName = name
|
||||||
m.provider = provider
|
m.provider = provider
|
||||||
m.syncAgent(llm.NewOpenAICompatibleClient(provider.BaseURL, provider.APIKey, m.httpClient))
|
m.syncAgent(providerClient(provider, m.httpClient))
|
||||||
return "Switched provider.\n" + m.summary(), nil
|
return "Switched provider.\n" + m.summary(), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -249,7 +279,7 @@ func (m *Manager) SetThinking(level string) (string, error) {
|
|||||||
m.provider.Thinking = level
|
m.provider.Thinking = level
|
||||||
m.provider.ThinkingConfigured = true
|
m.provider.ThinkingConfigured = true
|
||||||
if m.provider.ThinkingParam == "" {
|
if m.provider.ThinkingParam == "" {
|
||||||
m.provider.ThinkingParam = "thinking"
|
m.provider.ThinkingParam = config.DefaultThinkingParam(m.provider.APIType)
|
||||||
}
|
}
|
||||||
m.cfg.Providers[m.providerName] = m.provider
|
m.cfg.Providers[m.providerName] = m.provider
|
||||||
m.syncAgent(nil)
|
m.syncAgent(nil)
|
||||||
@@ -257,7 +287,14 @@ func (m *Manager) SetThinking(level string) (string, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) summary() string {
|
func (m *Manager) summary() string {
|
||||||
return fmt.Sprintf("provider: %s\nmodel: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentThinking())
|
return fmt.Sprintf("provider: %s\nmodel: %s\napi: %s\nthinking: %s", m.CurrentProvider(), m.CurrentModel(), m.CurrentAPIType(), m.CurrentThinking())
|
||||||
|
}
|
||||||
|
|
||||||
|
func providerClient(provider config.ProviderConfig, httpClient *http.Client) llm.Provider {
|
||||||
|
if provider.APIType == config.APITypeResponses {
|
||||||
|
return llm.NewOpenAIResponsesClient(provider.BaseURL, provider.APIKey, httpClient)
|
||||||
|
}
|
||||||
|
return llm.NewOpenAICompatibleClient(provider.BaseURL, provider.APIKey, httpClient)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) maxContextTokens() int {
|
func (m *Manager) maxContextTokens() int {
|
||||||
@@ -288,9 +325,13 @@ func (m *Manager) providerNames() []string {
|
|||||||
return names
|
return names
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (m *Manager) AvailableModels() []string {
|
||||||
|
return append([]string(nil), m.provider.Models...)
|
||||||
|
}
|
||||||
|
|
||||||
func contains(values []string, value string) bool {
|
func contains(values []string, value string) bool {
|
||||||
for _, item := range values {
|
for _, v := range values {
|
||||||
if item == value {
|
if v == value {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,6 +2,8 @@ package session
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
@@ -109,6 +111,76 @@ func TestManagerRejectsUnknownModelAndThinking(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestManagerSwitchesToResponsesProvider(t *testing.T) {
|
||||||
|
var hitPath string
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
hitPath = r.URL.Path
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
_, _ = w.Write([]byte(`data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"ok"}` + "\n\n"))
|
||||||
|
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed"}}` + "\n\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
cfg := &config.Config{
|
||||||
|
Providers: map[string]config.ProviderConfig{
|
||||||
|
"chat": {
|
||||||
|
Name: "chat",
|
||||||
|
APIType: config.APITypeChat,
|
||||||
|
BaseURL: "https://chat.example.com",
|
||||||
|
APIKey: "sk-test",
|
||||||
|
Model: "model-a",
|
||||||
|
Models: []string{"model-a"},
|
||||||
|
},
|
||||||
|
"responses": {
|
||||||
|
Name: "responses",
|
||||||
|
APIType: config.APITypeResponses,
|
||||||
|
BaseURL: server.URL,
|
||||||
|
APIKey: "sk-test",
|
||||||
|
Model: "model-b",
|
||||||
|
Models: []string{"model-b"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
assistant := agent.New(agent.Options{Provider: &recordingProvider{}, Model: "model-a"})
|
||||||
|
store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
manager, err := NewManager(cfg, assistant, server.Client(), store)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := manager.SetProvider("responses"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if manager.CurrentAPIType() != config.APITypeResponses {
|
||||||
|
t.Fatalf("api type = %q", manager.CurrentAPIType())
|
||||||
|
}
|
||||||
|
var out strings.Builder
|
||||||
|
if err := assistant.RunTurn(context.Background(), "hi", &out, nil); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if hitPath != "/v1/responses" {
|
||||||
|
t.Fatalf("path = %q, want /v1/responses", hitPath)
|
||||||
|
}
|
||||||
|
if out.String() != "ok" {
|
||||||
|
t.Fatalf("output = %q", out.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerExposesModelCommandCandidates(t *testing.T) {
|
||||||
|
manager, _, _ := testManager(t)
|
||||||
|
models := manager.AvailableModels()
|
||||||
|
if len(models) != 2 || models[0] != "model-a" || models[1] != "model-b" {
|
||||||
|
t.Fatalf("AvailableModels() = %v", models)
|
||||||
|
}
|
||||||
|
models[0] = "mutated"
|
||||||
|
if got := manager.AvailableModels()[0]; got != "model-a" {
|
||||||
|
t.Fatalf("mutation leaked: %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestManagerSaveAndResume(t *testing.T) {
|
func TestManagerSaveAndResume(t *testing.T) {
|
||||||
manager, assistant, _ := testManager(t)
|
manager, assistant, _ := testManager(t)
|
||||||
|
|
||||||
@@ -331,3 +403,132 @@ func TestManagerUsesModelContextForAgentCompaction(t *testing.T) {
|
|||||||
t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
|
t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestManagerResumeRestoresRuntime(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
_, _ = w.Write([]byte(`data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"ok"}` + "\n\n"))
|
||||||
|
_, _ = w.Write([]byte(`data: {"type":"response.completed","response":{"id":"resp_1","status":"completed"}}` + "\n\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
cfg := &config.Config{
|
||||||
|
Providers: map[string]config.ProviderConfig{
|
||||||
|
"chat": {
|
||||||
|
Name: "chat",
|
||||||
|
APIType: config.APITypeChat,
|
||||||
|
BaseURL: "https://chat.example.com",
|
||||||
|
APIKey: "sk-test",
|
||||||
|
Model: "model-a",
|
||||||
|
Models: []string{"model-a"},
|
||||||
|
Thinking: "none",
|
||||||
|
ThinkingParam: "thinking",
|
||||||
|
ThinkingConfigured: true,
|
||||||
|
},
|
||||||
|
"resp": {
|
||||||
|
Name: "resp",
|
||||||
|
APIType: config.APITypeResponses,
|
||||||
|
BaseURL: server.URL,
|
||||||
|
APIKey: "sk-test",
|
||||||
|
Model: "m1",
|
||||||
|
Models: []string{"m1", "m2"},
|
||||||
|
Thinking: "none",
|
||||||
|
ThinkingParam: "reasoning",
|
||||||
|
ThinkingConfigured: true,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
assistant := agent.New(agent.Options{
|
||||||
|
Provider: &recordingProvider{},
|
||||||
|
Model: "model-a",
|
||||||
|
Thinking: "none",
|
||||||
|
ThinkingParam: "thinking",
|
||||||
|
ThinkingEnabled: true,
|
||||||
|
})
|
||||||
|
store, err := NewStore(filepath.Join(t.TempDir(), "sessions"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
manager, err := NewManager(cfg, assistant, server.Client(), store)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build a session on the responses provider with a custom model and thinking.
|
||||||
|
manager.newSession()
|
||||||
|
if _, err := manager.SetProvider("resp"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := manager.SetModel("m2"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := manager.SetThinking("xhigh"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := manager.Save(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
id := manager.CurrentSessionID()
|
||||||
|
|
||||||
|
// Drift away from the saved runtime.
|
||||||
|
if _, err := manager.SetProvider("chat"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := manager.SetModel("model-a"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := manager.SetThinking("none"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := manager.Resume(id); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if got := manager.CurrentProvider(); got != "resp" {
|
||||||
|
t.Fatalf("provider = %q, want resp", got)
|
||||||
|
}
|
||||||
|
if got := manager.CurrentModel(); got != "m2" {
|
||||||
|
t.Fatalf("model = %q, want m2", got)
|
||||||
|
}
|
||||||
|
if got := manager.CurrentThinking(); got != "xhigh" {
|
||||||
|
t.Fatalf("thinking = %q, want xhigh", got)
|
||||||
|
}
|
||||||
|
if got := manager.CurrentAPIType(); got != config.APITypeResponses {
|
||||||
|
t.Fatalf("api type = %q, want responses", got)
|
||||||
|
}
|
||||||
|
runtime := assistant.RuntimeInfo()
|
||||||
|
if runtime.ProviderName != "resp" || runtime.Model != "m2" || runtime.Thinking != "xhigh" {
|
||||||
|
t.Fatalf("agent runtime = %#v", runtime)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerResumeFallsBackWhenRuntimeMissing(t *testing.T) {
|
||||||
|
manager, assistant, _ := testManager(t)
|
||||||
|
|
||||||
|
sess := &Session{
|
||||||
|
ID: GenerateID(),
|
||||||
|
Provider: "ghost",
|
||||||
|
Model: "ghost-model",
|
||||||
|
Thinking: "xhigh",
|
||||||
|
Messages: []llm.Message{{Role: llm.RoleUser, Content: "hi"}},
|
||||||
|
}
|
||||||
|
if err := manager.store.Save(sess); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := manager.Resume(sess.ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if got := manager.CurrentProvider(); got != "one" {
|
||||||
|
t.Fatalf("provider = %q, want default one", got)
|
||||||
|
}
|
||||||
|
if got := manager.CurrentModel(); got != "model-a" {
|
||||||
|
t.Fatalf("model = %q, want default model-a", got)
|
||||||
|
}
|
||||||
|
if got := manager.CurrentThinking(); got != "xhigh" {
|
||||||
|
t.Fatalf("thinking = %q, want restored xhigh", got)
|
||||||
|
}
|
||||||
|
runtime := assistant.RuntimeInfo()
|
||||||
|
if runtime.ProviderName != "one" || runtime.Model != "model-a" {
|
||||||
|
t.Fatalf("agent runtime = %#v", runtime)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ import (
|
|||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"agentu/internal/safeio"
|
||||||
"agentu/pkg/llm"
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -22,6 +23,8 @@ type Session struct {
|
|||||||
UpdatedAt time.Time `json:"updated_at"`
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
Provider string `json:"provider"`
|
Provider string `json:"provider"`
|
||||||
Model string `json:"model"`
|
Model string `json:"model"`
|
||||||
|
APIType string `json:"api_type,omitempty"`
|
||||||
|
Thinking string `json:"thinking,omitempty"`
|
||||||
Messages []llm.Message `json:"messages"`
|
Messages []llm.Message `json:"messages"`
|
||||||
LastUsage *llm.Usage `json:"last_usage,omitempty"`
|
LastUsage *llm.Usage `json:"last_usage,omitempty"`
|
||||||
}
|
}
|
||||||
@@ -52,8 +55,11 @@ func NewStore(dir string) (*Store, error) {
|
|||||||
|
|
||||||
// Save writes a session to disk.
|
// Save writes a session to disk.
|
||||||
func (s *Store) Save(sess *Session) error {
|
func (s *Store) Save(sess *Session) error {
|
||||||
if sess.ID == "" {
|
if sess == nil {
|
||||||
return fmt.Errorf("session ID is required")
|
return fmt.Errorf("session is required")
|
||||||
|
}
|
||||||
|
if err := validateSessionID(sess.ID); err != nil {
|
||||||
|
return err
|
||||||
}
|
}
|
||||||
sess.UpdatedAt = time.Now().UTC()
|
sess.UpdatedAt = time.Now().UTC()
|
||||||
if sess.CreatedAt.IsZero() {
|
if sess.CreatedAt.IsZero() {
|
||||||
@@ -64,14 +70,17 @@ func (s *Store) Save(sess *Session) error {
|
|||||||
return fmt.Errorf("marshal session: %w", err)
|
return fmt.Errorf("marshal session: %w", err)
|
||||||
}
|
}
|
||||||
path := s.filePath(sess.ID)
|
path := s.filePath(sess.ID)
|
||||||
if err := os.WriteFile(path, data, 0o644); err != nil {
|
if err := safeio.AtomicWriteFile(path, data, 0o644); err != nil {
|
||||||
return fmt.Errorf("write session: %w", err)
|
return fmt.Errorf("write session: %w", err)
|
||||||
}
|
}
|
||||||
return nil
|
return s.updateIndex(metaFromSession(sess))
|
||||||
}
|
}
|
||||||
|
|
||||||
// Load reads a session by ID.
|
// Load reads a session by ID.
|
||||||
func (s *Store) Load(id string) (*Session, error) {
|
func (s *Store) Load(id string) (*Session, error) {
|
||||||
|
if err := validateSessionID(id); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
path := s.filePath(id)
|
path := s.filePath(id)
|
||||||
data, err := os.ReadFile(path)
|
data, err := os.ReadFile(path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -89,13 +98,27 @@ func (s *Store) Load(id string) (*Session, error) {
|
|||||||
|
|
||||||
// List returns metadata for all sessions, sorted by updated_at descending.
|
// List returns metadata for all sessions, sorted by updated_at descending.
|
||||||
func (s *Store) List() ([]SessionMeta, error) {
|
func (s *Store) List() ([]SessionMeta, error) {
|
||||||
|
if metas, err := s.loadIndex(); err == nil {
|
||||||
|
return metas, nil
|
||||||
|
}
|
||||||
|
metas, err := s.scanSessions()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if err := s.saveIndexFromSessions(metas); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return metas, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Store) scanSessions() ([]SessionMeta, error) {
|
||||||
entries, err := os.ReadDir(s.dir)
|
entries, err := os.ReadDir(s.dir)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("read session dir: %w", err)
|
return nil, fmt.Errorf("read session dir: %w", err)
|
||||||
}
|
}
|
||||||
var metas []SessionMeta
|
var metas []SessionMeta
|
||||||
for _, entry := range entries {
|
for _, entry := range entries {
|
||||||
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
|
if entry.IsDir() || entry.Name() == sessionIndexFile || !strings.HasSuffix(entry.Name(), ".json") {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
path := filepath.Join(s.dir, entry.Name())
|
path := filepath.Join(s.dir, entry.Name())
|
||||||
@@ -125,6 +148,9 @@ func (s *Store) List() ([]SessionMeta, error) {
|
|||||||
|
|
||||||
// Delete removes a session file.
|
// Delete removes a session file.
|
||||||
func (s *Store) Delete(id string) error {
|
func (s *Store) Delete(id string) error {
|
||||||
|
if err := validateSessionID(id); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
path := s.filePath(id)
|
path := s.filePath(id)
|
||||||
if err := os.Remove(path); err != nil {
|
if err := os.Remove(path); err != nil {
|
||||||
if os.IsNotExist(err) {
|
if os.IsNotExist(err) {
|
||||||
@@ -132,13 +158,64 @@ func (s *Store) Delete(id string) error {
|
|||||||
}
|
}
|
||||||
return fmt.Errorf("delete session: %w", err)
|
return fmt.Errorf("delete session: %w", err)
|
||||||
}
|
}
|
||||||
return nil
|
metas, err := s.loadIndex()
|
||||||
|
if err != nil {
|
||||||
|
metas, err = s.scanSessions()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
filtered := metas[:0]
|
||||||
|
for _, meta := range metas {
|
||||||
|
if meta.ID != id {
|
||||||
|
filtered = append(filtered, meta)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return s.saveIndexFromSessions(filtered)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *Store) filePath(id string) string {
|
func (s *Store) filePath(id string) string {
|
||||||
return filepath.Join(s.dir, id+".json")
|
return filepath.Join(s.dir, id+".json")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func validateSessionID(id string) error {
|
||||||
|
if id == "" {
|
||||||
|
return fmt.Errorf("session ID is required")
|
||||||
|
}
|
||||||
|
if id == "." || id == ".." {
|
||||||
|
return fmt.Errorf("invalid session ID %q", id)
|
||||||
|
}
|
||||||
|
for _, r := range id {
|
||||||
|
if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-' || r == '_' || r == '.' {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
return fmt.Errorf("invalid session ID %q", id)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Store) updateIndex(meta SessionMeta) error {
|
||||||
|
metas, err := s.loadIndex()
|
||||||
|
if err != nil {
|
||||||
|
metas, err = s.scanSessions()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
found := false
|
||||||
|
for i := range metas {
|
||||||
|
if metas[i].ID == meta.ID {
|
||||||
|
metas[i] = meta
|
||||||
|
found = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !found {
|
||||||
|
metas = append(metas, meta)
|
||||||
|
}
|
||||||
|
return s.saveIndexFromSessions(metas)
|
||||||
|
}
|
||||||
|
|
||||||
// GenerateID creates a random 8-character hex session ID.
|
// GenerateID creates a random 8-character hex session ID.
|
||||||
func GenerateID() string {
|
func GenerateID() string {
|
||||||
b := make([]byte, 4)
|
b := make([]byte, 4)
|
||||||
|
|||||||
@@ -119,6 +119,9 @@ func TestDeleteNotFound(t *testing.T) {
|
|||||||
|
|
||||||
func TestSaveRequiresID(t *testing.T) {
|
func TestSaveRequiresID(t *testing.T) {
|
||||||
store := tempStore(t)
|
store := tempStore(t)
|
||||||
|
if err := store.Save(nil); err == nil {
|
||||||
|
t.Fatal("expected error for nil session")
|
||||||
|
}
|
||||||
err := store.Save(&Session{Name: "no id"})
|
err := store.Save(&Session{Name: "no id"})
|
||||||
if err == nil {
|
if err == nil {
|
||||||
t.Fatal("expected error for empty ID")
|
t.Fatal("expected error for empty ID")
|
||||||
@@ -172,3 +175,40 @@ func TestNewStoreCreatesDir(t *testing.T) {
|
|||||||
t.Error("expected directory")
|
t.Error("expected directory")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestStoreMaintainsSessionIndex(t *testing.T) {
|
||||||
|
store := tempStore(t)
|
||||||
|
if err := store.Save(&Session{ID: "indexed", Name: "Indexed"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(filepath.Join(store.dir, sessionIndexFile)); err != nil {
|
||||||
|
t.Fatalf("index not created: %v", err)
|
||||||
|
}
|
||||||
|
metas, err := store.loadIndex()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(metas) != 1 || metas[0].ID != "indexed" {
|
||||||
|
t.Fatalf("index metadata = %+v", metas)
|
||||||
|
}
|
||||||
|
if err := store.Delete("indexed"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
metas, err = store.loadIndex()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(metas) != 0 {
|
||||||
|
t.Fatalf("index after delete = %+v", metas)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStoreRejectsUnsafeSessionID(t *testing.T) {
|
||||||
|
store := tempStore(t)
|
||||||
|
if _, err := store.Load("../outside"); err == nil {
|
||||||
|
t.Fatal("expected unsafe ID rejection")
|
||||||
|
}
|
||||||
|
if err := store.Save(&Session{ID: "a/b"}); err == nil {
|
||||||
|
t.Fatal("expected unsafe ID rejection")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
})
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,153 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ────────────────────────────────────────────────────────────────────────────
|
||||||
|
// ApplyPatchTool — multi-file patch using V4 format
|
||||||
|
// ────────────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
type ApplyPatchTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewApplyPatchTool(workingDir string) *ApplyPatchTool {
|
||||||
|
return &ApplyPatchTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
// MutatingTool interface.
|
||||||
|
|
||||||
|
func (*ApplyPatchTool) Mutates() bool { return true }
|
||||||
|
|
||||||
|
// CallMutatingTool — dry_run is non-mutating.
|
||||||
|
func (t *ApplyPatchTool) MutatesCall(args json.RawMessage) bool {
|
||||||
|
var a applyPatchArgs
|
||||||
|
if err := json.Unmarshal(args, &a); err != nil {
|
||||||
|
return true // conservative
|
||||||
|
}
|
||||||
|
return !a.DryRun
|
||||||
|
}
|
||||||
|
|
||||||
|
// PreviewableTool interface.
|
||||||
|
|
||||||
|
func (t *ApplyPatchTool) Preview(ctx context.Context, args json.RawMessage) (string, error) {
|
||||||
|
var a applyPatchArgs
|
||||||
|
if err := json.Unmarshal(args, &a); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
return summarizePatch(a.Patch)
|
||||||
|
}
|
||||||
|
|
||||||
|
// TargetPathTool interface.
|
||||||
|
|
||||||
|
func (t *ApplyPatchTool) TargetPaths(args json.RawMessage) ([]string, error) {
|
||||||
|
var a applyPatchArgs
|
||||||
|
if err := json.Unmarshal(args, &a); err != nil {
|
||||||
|
return nil, fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
return extractTargetPaths(t.workingDir, a.Patch)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Definition.
|
||||||
|
|
||||||
|
func (t *ApplyPatchTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "apply_patch",
|
||||||
|
Description: "Apply a multi-file patch to create, update, or delete files. Uses a V4 diff format with *** Begin Patch / *** Add File / *** Update File / *** Delete File markers. Prefer this over multiple file_edit calls when changing 2+ files.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"patch": {"type": "string", "description": "V4 format patch text. Must start with *** Begin Patch and end with *** End Patch."},
|
||||||
|
"dry_run": {"type": "boolean", "description": "Preview the patch without writing any files."}
|
||||||
|
},
|
||||||
|
"required": ["patch"],
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type applyPatchArgs struct {
|
||||||
|
Patch string `json:"patch"`
|
||||||
|
DryRun bool `json:"dry_run"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute.
|
||||||
|
|
||||||
|
func (t *ApplyPatchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args applyPatchArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if strings.TrimSpace(args.Patch) == "" {
|
||||||
|
return "", fmt.Errorf("patch is empty")
|
||||||
|
}
|
||||||
|
|
||||||
|
pp, err := parsePatch(args.Patch)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
if args.DryRun {
|
||||||
|
summary, err := summarizePatch(args.Patch)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return "dry run — " + summary, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
rollback := NewRollbackStore()
|
||||||
|
var snapshotIDs []string
|
||||||
|
for _, pf := range pp.files {
|
||||||
|
path, err := resolvePathForWrite(t.workingDir, pf.path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
id, err := rollback.Capture("apply_patch", pf.path, path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
snapshotIDs = append(snapshotIDs, id)
|
||||||
|
}
|
||||||
|
|
||||||
|
var results []string
|
||||||
|
for _, pf := range pp.files {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
if rollbackErr := restoreSnapshots(rollback, snapshotIDs); rollbackErr != nil {
|
||||||
|
return "", fmt.Errorf("%w; rollback failed: %v", err, rollbackErr)
|
||||||
|
}
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
result, err := applyPatchFile(t.workingDir, pf)
|
||||||
|
if err != nil {
|
||||||
|
if rollbackErr := restoreSnapshots(rollback, snapshotIDs); rollbackErr != nil {
|
||||||
|
return "", fmt.Errorf("apply patch: %w; rollback failed: %v", err, rollbackErr)
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("apply patch: %w", err)
|
||||||
|
}
|
||||||
|
results = append(results, result)
|
||||||
|
}
|
||||||
|
|
||||||
|
summary := fmt.Sprintf("applied patch: %d file(s) changed\n%s", len(results), strings.Join(results, "\n"))
|
||||||
|
return summary, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func restoreSnapshots(store *RollbackStore, ids []string) error {
|
||||||
|
var errs []error
|
||||||
|
for i := len(ids) - 1; i >= 0; i-- {
|
||||||
|
if _, err := store.Restore(ids[i]); err != nil {
|
||||||
|
errs = append(errs, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return errors.Join(errs...)
|
||||||
|
}
|
||||||
+45
-16
@@ -13,6 +13,7 @@ import (
|
|||||||
"sort"
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
|
"agentu/internal/safeio"
|
||||||
"agentu/pkg/llm"
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -81,6 +82,34 @@ type FileWriteTool struct {
|
|||||||
func NewFileWriteTool(workingDir string) *FileWriteTool {
|
func NewFileWriteTool(workingDir string) *FileWriteTool {
|
||||||
return &FileWriteTool{workingDir: workingDir}
|
return &FileWriteTool{workingDir: workingDir}
|
||||||
}
|
}
|
||||||
|
func (*FileWriteTool) Mutates() bool { return true }
|
||||||
|
|
||||||
|
func (t *FileWriteTool) Preview(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args fileWriteArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
action := "write"
|
||||||
|
if args.Append {
|
||||||
|
action = "append"
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%s %d bytes to %s", action, len(args.Content), args.Path), ctx.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *FileWriteTool) TargetPaths(raw json.RawMessage) ([]string, error) {
|
||||||
|
var args fileWriteArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return nil, fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(args.Path) == "" {
|
||||||
|
return nil, errors.New("path is required")
|
||||||
|
}
|
||||||
|
path, err := resolvePathForWrite(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return []string{path}, nil
|
||||||
|
}
|
||||||
|
|
||||||
func (t *FileWriteTool) Definition() llm.Tool {
|
func (t *FileWriteTool) Definition() llm.Tool {
|
||||||
return llm.Tool{
|
return llm.Tool{
|
||||||
@@ -116,28 +145,28 @@ func (t *FileWriteTool) Execute(ctx context.Context, raw json.RawMessage) (strin
|
|||||||
if strings.TrimSpace(args.Path) == "" {
|
if strings.TrimSpace(args.Path) == "" {
|
||||||
return "", errors.New("path is required")
|
return "", errors.New("path is required")
|
||||||
}
|
}
|
||||||
path, err := resolvePath(t.workingDir, args.Path)
|
path, err := resolvePathForWrite(t.workingDir, args.Path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
|
||||||
return "", err
|
|
||||||
}
|
|
||||||
flag := os.O_CREATE | os.O_WRONLY | os.O_TRUNC
|
|
||||||
if args.Append {
|
|
||||||
flag = os.O_CREATE | os.O_WRONLY | os.O_APPEND
|
|
||||||
}
|
|
||||||
file, err := os.OpenFile(path, flag, 0o644)
|
|
||||||
if err != nil {
|
|
||||||
return "", err
|
|
||||||
}
|
|
||||||
defer file.Close()
|
|
||||||
if _, err := file.WriteString(args.Content); err != nil {
|
|
||||||
return "", err
|
|
||||||
}
|
|
||||||
if err := ctx.Err(); err != nil {
|
if err := ctx.Err(); err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
|
content := []byte(args.Content)
|
||||||
|
perm := os.FileMode(0o644)
|
||||||
|
if info, statErr := os.Stat(path); statErr == nil {
|
||||||
|
perm = info.Mode().Perm()
|
||||||
|
}
|
||||||
|
if args.Append {
|
||||||
|
existing, readErr := os.ReadFile(path)
|
||||||
|
if readErr != nil && !os.IsNotExist(readErr) {
|
||||||
|
return "", readErr
|
||||||
|
}
|
||||||
|
content = append(existing, content...)
|
||||||
|
}
|
||||||
|
if err := safeio.AtomicWriteFile(path, content, perm); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
action := "wrote"
|
action := "wrote"
|
||||||
if args.Append {
|
if args.Append {
|
||||||
action = "appended"
|
action = "appended"
|
||||||
|
|||||||
@@ -0,0 +1,209 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agentu/internal/safeio"
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
type FileEditTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewFileEditTool(workingDir string) *FileEditTool {
|
||||||
|
return &FileEditTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (*FileEditTool) Mutates() bool { return true }
|
||||||
|
|
||||||
|
func (t *FileEditTool) Preview(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args fileEditArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
if args.StartLine > 0 {
|
||||||
|
end := args.EndLine
|
||||||
|
if end == 0 {
|
||||||
|
end = args.StartLine
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("edit %s lines %d-%d", args.Path, args.StartLine, end), ctx.Err()
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("replace text in %s", args.Path), ctx.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *FileEditTool) TargetPaths(raw json.RawMessage) ([]string, error) {
|
||||||
|
var args fileEditArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return nil, fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(args.Path) == "" {
|
||||||
|
return nil, errors.New("path is required")
|
||||||
|
}
|
||||||
|
path, err := resolvePath(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return []string{path}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *FileEditTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "file_edit",
|
||||||
|
Description: "Edit an existing UTF-8 text file by replacing an inclusive 1-based line range or an exact old text pattern.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Existing file path relative to the configured agent working directory."},
|
||||||
|
"start_line": {"type": "integer", "description": "1-based first line for a line-range edit."},
|
||||||
|
"end_line": {"type": "integer", "description": "Optional 1-based final line; defaults to start_line when omitted."},
|
||||||
|
"old": {"type": "string", "description": "Exact text to replace for a pattern edit; not regex."},
|
||||||
|
"new": {"type": "string", "description": "Replacement text; may be empty string for deletion."},
|
||||||
|
"replace_all": {"type": "boolean", "description": "When true, replace every exact old occurrence; otherwise require exactly one occurrence."}
|
||||||
|
},
|
||||||
|
"required": ["path", "new"],
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type fileEditArgs struct {
|
||||||
|
Path string `json:"path"`
|
||||||
|
StartLine int `json:"start_line"`
|
||||||
|
EndLine int `json:"end_line"`
|
||||||
|
Old string `json:"old"`
|
||||||
|
New *string `json:"new"`
|
||||||
|
ReplaceAll bool `json:"replace_all"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *FileEditTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args fileEditArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(args.Path) == "" {
|
||||||
|
return "", errors.New("path is required")
|
||||||
|
}
|
||||||
|
if args.New == nil {
|
||||||
|
return "", errors.New("new is required")
|
||||||
|
}
|
||||||
|
path, err := resolvePath(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
info, err := os.Stat(path)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("cannot access %s: %w", args.Path, err)
|
||||||
|
}
|
||||||
|
if info.IsDir() {
|
||||||
|
return "", fmt.Errorf("%s is a directory, not a file", args.Path)
|
||||||
|
}
|
||||||
|
perm := info.Mode().Perm()
|
||||||
|
|
||||||
|
hasLine := args.StartLine > 0 || args.EndLine > 0
|
||||||
|
hasOld := strings.TrimSpace(args.Old) != ""
|
||||||
|
|
||||||
|
if hasLine && hasOld {
|
||||||
|
return "", errors.New("provide either start_line/end_line or old, not both")
|
||||||
|
}
|
||||||
|
if !hasLine && !hasOld {
|
||||||
|
return "", errors.New("provide either start_line/end_line or old")
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
content := string(data)
|
||||||
|
|
||||||
|
var result string
|
||||||
|
if hasOld {
|
||||||
|
result, err = editByPattern(content, args.Path, args.Old, *args.New, args.ReplaceAll)
|
||||||
|
} else {
|
||||||
|
result, err = editByLineRange(content, args.Path, args.StartLine, args.EndLine, *args.New)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if err := safeio.AtomicWriteFile(path, []byte(result), perm); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
switch {
|
||||||
|
case hasOld:
|
||||||
|
count := strings.Count(content, args.Old)
|
||||||
|
if args.ReplaceAll {
|
||||||
|
return fmt.Sprintf("edited %s: replaced %d occurrences", args.Path, count), nil
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("edited %s: replaced 1 occurrence", args.Path), nil
|
||||||
|
default:
|
||||||
|
endLine := args.EndLine
|
||||||
|
if endLine == 0 {
|
||||||
|
endLine = args.StartLine
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("edited %s: replaced lines %d-%d", args.Path, args.StartLine, endLine), nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func editByLineRange(content, displayPath string, startLine, endLine int, replacement string) (string, error) {
|
||||||
|
if startLine < 1 {
|
||||||
|
return "", errors.New("start_line must be >= 1")
|
||||||
|
}
|
||||||
|
if endLine == 0 {
|
||||||
|
endLine = startLine
|
||||||
|
}
|
||||||
|
if endLine < startLine {
|
||||||
|
return "", fmt.Errorf("end_line (%d) must be >= start_line (%d)", endLine, startLine)
|
||||||
|
}
|
||||||
|
|
||||||
|
lines := strings.Split(content, "\n")
|
||||||
|
if startLine > len(lines) {
|
||||||
|
return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
|
||||||
|
}
|
||||||
|
if endLine > len(lines) {
|
||||||
|
return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
|
||||||
|
}
|
||||||
|
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString(strings.Join(lines[:startLine-1], "\n"))
|
||||||
|
if startLine > 1 {
|
||||||
|
b.WriteString("\n")
|
||||||
|
}
|
||||||
|
b.WriteString(replacement)
|
||||||
|
if endLine < len(lines) {
|
||||||
|
if !strings.HasSuffix(replacement, "\n") && replacement != "" {
|
||||||
|
b.WriteString("\n")
|
||||||
|
}
|
||||||
|
b.WriteString(strings.Join(lines[endLine:], "\n"))
|
||||||
|
}
|
||||||
|
|
||||||
|
return b.String(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func editByPattern(content, displayPath, old, replacement string, replaceAll bool) (string, error) {
|
||||||
|
count := strings.Count(content, old)
|
||||||
|
if count == 0 {
|
||||||
|
return "", errors.New("old text not found")
|
||||||
|
}
|
||||||
|
if !replaceAll && count != 1 {
|
||||||
|
return "", fmt.Errorf("old text matched %d occurrences; set replace_all=true or provide a more specific old value", count)
|
||||||
|
}
|
||||||
|
return strings.ReplaceAll(content, old, replacement), nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,136 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestFileEditReplacesLineRange(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("one\ntwo\nthree\n"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","start_line":2,"end_line":2,"new":"TWO\n"}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "replaced lines 2-2") {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
got, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if string(got) != "one\nTWO\nthree\n" {
|
||||||
|
t.Fatalf("file content = %q", string(got))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditReplacesExactPattern(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("alpha beta alpha"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"beta","new":"BETA"}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "replaced 1 occurrence") {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
got, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if string(got) != "alpha BETA alpha" {
|
||||||
|
t.Fatalf("file content = %q", string(got))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditRejectsAmbiguousPattern(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("x x"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"x","new":"y"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for ambiguous pattern")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "matched 2 occurrences") {
|
||||||
|
t.Fatalf("error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditReplaceAll(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("x x"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
out, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"x","new":"y","replace_all":true}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "replaced 2 occurrences") {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
got, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if string(got) != "y y" {
|
||||||
|
t.Fatalf("file content = %q", string(got))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditRejectsEscapingPath(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"../outside.txt","old":"x","new":"y"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for escaping path")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "escapes") {
|
||||||
|
t.Fatalf("error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditRejectsMissingNew(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "notes.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"notes.txt","old":"hello"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for missing new")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "new is required") {
|
||||||
|
t.Fatalf("error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFileEditRejectsNonExistentFile(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
edit := NewFileEditTool(dir)
|
||||||
|
_, err := edit.Execute(context.Background(), json.RawMessage(`{"path":"missing.txt","old":"x","new":"y"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for non-existent file")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,246 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os/exec"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultGitLogMaxCount = 20
|
||||||
|
maxGitLogMaxCount = 100
|
||||||
|
)
|
||||||
|
|
||||||
|
func runGit(ctx context.Context, workingDir string, args ...string) (string, error) {
|
||||||
|
gitArgs := append([]string{"-C", workingDir}, args...)
|
||||||
|
cmd := exec.CommandContext(ctx, "git", gitArgs...)
|
||||||
|
output, err := cmd.CombinedOutput()
|
||||||
|
text := FormatResult(string(output))
|
||||||
|
if err != nil {
|
||||||
|
if strings.Contains(err.Error(), "context deadline exceeded") || strings.Contains(err.Error(), "signal: killed") {
|
||||||
|
if text != "" {
|
||||||
|
return text, fmt.Errorf("git command timed out: %w", err)
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("git command timed out: %w", err)
|
||||||
|
}
|
||||||
|
if text != "" {
|
||||||
|
return text, fmt.Errorf("git command failed: %w", err)
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("git command failed: %w", err)
|
||||||
|
}
|
||||||
|
return text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func gitPathspec(baseDir, path string) (string, error) {
|
||||||
|
if strings.TrimSpace(path) == "" {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
resolved, err := resolvePath(baseDir, path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
rel, err := filepath.Rel(baseDir, resolved)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return filepath.ToSlash(filepath.Clean(rel)), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeGitLogMaxCount(value int) int {
|
||||||
|
if value <= 0 {
|
||||||
|
return defaultGitLogMaxCount
|
||||||
|
}
|
||||||
|
if value > maxGitLogMaxCount {
|
||||||
|
return maxGitLogMaxCount
|
||||||
|
}
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
|
||||||
|
// git_status
|
||||||
|
|
||||||
|
type GitStatusTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewGitStatusTool(workingDir string) *GitStatusTool {
|
||||||
|
return &GitStatusTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitStatusTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "git_status",
|
||||||
|
Description: "Show changed files using git status --short.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Optional path to limit status to. Relative paths are resolved from the configured agent working directory."}
|
||||||
|
},
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type gitStatusArgs struct {
|
||||||
|
Path string `json:"path"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitStatusTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args gitStatusArgs
|
||||||
|
if len(raw) > 0 {
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
gitArgs := []string{"status", "--short"}
|
||||||
|
if args.Path != "" {
|
||||||
|
pathspec, err := gitPathspec(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
gitArgs = append(gitArgs, "--", pathspec)
|
||||||
|
}
|
||||||
|
text, err := runGit(ctx, t.workingDir, gitArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return text, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(text) == "" {
|
||||||
|
return "Working tree clean.", nil
|
||||||
|
}
|
||||||
|
return "Changed files:\n" + text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// git_diff
|
||||||
|
|
||||||
|
type GitDiffTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewGitDiffTool(workingDir string) *GitDiffTool {
|
||||||
|
return &GitDiffTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitDiffTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "git_diff",
|
||||||
|
Description: "Show unstaged or staged git diff output, optionally as --stat.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Optional path to limit diff to. Relative paths are resolved from the configured agent working directory."},
|
||||||
|
"staged": {"type": "boolean", "description": "Show staged diff instead of unstaged."},
|
||||||
|
"stat": {"type": "boolean", "description": "Show diffstat summary instead of full diff."}
|
||||||
|
},
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type gitDiffArgs struct {
|
||||||
|
Path string `json:"path"`
|
||||||
|
Staged bool `json:"staged"`
|
||||||
|
Stat bool `json:"stat"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitDiffTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args gitDiffArgs
|
||||||
|
if len(raw) > 0 {
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
gitArgs := []string{"diff"}
|
||||||
|
if args.Staged {
|
||||||
|
gitArgs = append(gitArgs, "--cached")
|
||||||
|
}
|
||||||
|
if args.Stat {
|
||||||
|
gitArgs = append(gitArgs, "--stat")
|
||||||
|
}
|
||||||
|
if args.Path != "" {
|
||||||
|
pathspec, err := gitPathspec(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
gitArgs = append(gitArgs, "--", pathspec)
|
||||||
|
}
|
||||||
|
text, err := runGit(ctx, t.workingDir, gitArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return text, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(text) == "" {
|
||||||
|
if args.Staged {
|
||||||
|
return "No staged diff.", nil
|
||||||
|
}
|
||||||
|
return "No unstaged diff.", nil
|
||||||
|
}
|
||||||
|
return text, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// git_log
|
||||||
|
|
||||||
|
type GitLogTool struct {
|
||||||
|
workingDir string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewGitLogTool(workingDir string) *GitLogTool {
|
||||||
|
return &GitLogTool{workingDir: workingDir}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitLogTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "git_log",
|
||||||
|
Description: "Show recent git commits with git log --oneline --decorate.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Optional path to limit log to. Relative paths are resolved from the configured agent working directory."},
|
||||||
|
"max_count": {"type": "integer", "description": "Maximum number of commits to show. Defaults to 20, capped at 100."}
|
||||||
|
},
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type gitLogArgs struct {
|
||||||
|
Path string `json:"path"`
|
||||||
|
MaxCount int `json:"max_count"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *GitLogTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args gitLogArgs
|
||||||
|
if len(raw) > 0 {
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
maxCount := normalizeGitLogMaxCount(args.MaxCount)
|
||||||
|
gitArgs := []string{"log", "--oneline", "--decorate", fmt.Sprintf("-n%d", maxCount)}
|
||||||
|
if args.Path != "" {
|
||||||
|
pathspec, err := gitPathspec(t.workingDir, args.Path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
gitArgs = append(gitArgs, "--", pathspec)
|
||||||
|
}
|
||||||
|
text, err := runGit(ctx, t.workingDir, gitArgs...)
|
||||||
|
if err != nil {
|
||||||
|
return text, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(text) == "" {
|
||||||
|
return "No git history.", nil
|
||||||
|
}
|
||||||
|
return text, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,142 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func requireGit(t *testing.T) {
|
||||||
|
t.Helper()
|
||||||
|
if _, err := exec.LookPath("git"); err != nil {
|
||||||
|
t.Skip("git not installed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func runGitForTest(t *testing.T, dir string, args ...string) {
|
||||||
|
t.Helper()
|
||||||
|
gitArgs := append([]string{"-C", dir}, args...)
|
||||||
|
cmd := exec.Command("git", gitArgs...)
|
||||||
|
out, err := cmd.CombinedOutput()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("git %v failed: %v\n%s", args, err, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func initTestRepo(t *testing.T) string {
|
||||||
|
t.Helper()
|
||||||
|
dir := t.TempDir()
|
||||||
|
runGitForTest(t, dir, "init")
|
||||||
|
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "--allow-empty", "-m", "initial")
|
||||||
|
return dir
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitStatusToolShowsChangedFiles(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
path := filepath.Join(dir, "tracked.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
runGitForTest(t, dir, "add", "tracked.txt")
|
||||||
|
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
|
||||||
|
|
||||||
|
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tool := NewGitStatusTool(dir)
|
||||||
|
out, err := tool.Execute(context.Background(), nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "Changed files:") {
|
||||||
|
t.Fatalf("output missing 'Changed files:': %q", out)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "M tracked.txt") {
|
||||||
|
t.Fatalf("output missing 'M tracked.txt': %q", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitDiffToolShowsUnstagedDiff(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
path := filepath.Join(dir, "tracked.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
runGitForTest(t, dir, "add", "tracked.txt")
|
||||||
|
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
|
||||||
|
|
||||||
|
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tool := NewGitDiffTool(dir)
|
||||||
|
out, err := tool.Execute(context.Background(), nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "diff --git") {
|
||||||
|
t.Fatalf("output missing 'diff --git': %q", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitDiffToolShowsStat(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
path := filepath.Join(dir, "tracked.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("hello"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
runGitForTest(t, dir, "add", "tracked.txt")
|
||||||
|
runGitForTest(t, dir, "-c", "user.email=test@example.com", "-c", "user.name=Test", "commit", "-m", "add tracked")
|
||||||
|
|
||||||
|
if err := os.WriteFile(path, []byte("world"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
tool := NewGitDiffTool(dir)
|
||||||
|
out, err := tool.Execute(context.Background(), json.RawMessage(`{"stat":true}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "tracked.txt") {
|
||||||
|
t.Fatalf("output missing 'tracked.txt': %q", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitLogToolShowsCommits(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
tool := NewGitLogTool(dir)
|
||||||
|
out, err := tool.Execute(context.Background(), json.RawMessage(`{"max_count":1}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "initial") {
|
||||||
|
t.Fatalf("output missing 'initial': %q", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGitToolsRejectEscapingPath(t *testing.T) {
|
||||||
|
requireGit(t)
|
||||||
|
dir := initTestRepo(t)
|
||||||
|
|
||||||
|
tool := NewGitStatusTool(dir)
|
||||||
|
_, err := tool.Execute(context.Background(), json.RawMessage(`{"path":"../outside"}`))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for escaping path")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "escapes") {
|
||||||
|
t.Fatalf("error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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
@@ -2,22 +2,85 @@ package tools
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"strings"
|
"strings"
|
||||||
)
|
)
|
||||||
|
|
||||||
func resolvePath(baseDir, path string) (string, error) {
|
func resolvePath(baseDir, path string) (string, error) {
|
||||||
|
resolved, canonicalBase, err := resolveLexicalPath(baseDir, path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
canonical, err := filepath.EvalSymlinks(resolved)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if !isWithinDir(canonical, canonicalBase) {
|
||||||
|
return "", fmt.Errorf("path escapes the project directory through a symbolic link: %s", path)
|
||||||
|
}
|
||||||
|
return canonical, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// resolvePathForWrite resolves a path that may not exist yet while preventing
|
||||||
|
// existing parent-directory symlinks from escaping the project directory.
|
||||||
|
func resolvePathForWrite(baseDir, path string) (string, error) {
|
||||||
|
resolved, canonicalBase, err := resolveLexicalPath(baseDir, path)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
existing := resolved
|
||||||
|
var suffix []string
|
||||||
|
for {
|
||||||
|
_, err := os.Lstat(existing)
|
||||||
|
if err == nil {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if !os.IsNotExist(err) {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
parent := filepath.Dir(existing)
|
||||||
|
if parent == existing {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
suffix = append(suffix, filepath.Base(existing))
|
||||||
|
existing = parent
|
||||||
|
}
|
||||||
|
|
||||||
|
canonical, err := filepath.EvalSymlinks(existing)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
for i := len(suffix) - 1; i >= 0; i-- {
|
||||||
|
canonical = filepath.Join(canonical, suffix[i])
|
||||||
|
}
|
||||||
|
if !isWithinDir(canonical, canonicalBase) {
|
||||||
|
return "", fmt.Errorf("path escapes the project directory through a symbolic link: %s", path)
|
||||||
|
}
|
||||||
|
return canonical, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveLexicalPath(baseDir, path string) (string, string, error) {
|
||||||
if path == "" {
|
if path == "" {
|
||||||
path = "."
|
path = "."
|
||||||
}
|
}
|
||||||
if filepath.IsAbs(path) {
|
if filepath.IsAbs(path) {
|
||||||
return "", fmt.Errorf("absolute paths are not allowed; use a path relative to the project directory")
|
return "", "", fmt.Errorf("absolute paths are not allowed; use a path relative to the project directory")
|
||||||
}
|
}
|
||||||
resolved := filepath.Clean(filepath.Join(baseDir, path))
|
absBase, err := filepath.Abs(baseDir)
|
||||||
if !isWithinDir(resolved, baseDir) {
|
if err != nil {
|
||||||
return "", fmt.Errorf("path escapes the project directory: %s", path)
|
return "", "", err
|
||||||
}
|
}
|
||||||
return resolved, nil
|
canonicalBase, err := filepath.EvalSymlinks(absBase)
|
||||||
|
if err != nil {
|
||||||
|
return "", "", err
|
||||||
|
}
|
||||||
|
resolved := filepath.Clean(filepath.Join(absBase, path))
|
||||||
|
if !isWithinDir(resolved, absBase) {
|
||||||
|
return "", "", fmt.Errorf("path escapes the project directory: %s", path)
|
||||||
|
}
|
||||||
|
return resolved, canonicalBase, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func isWithinDir(path, dir string) bool {
|
func isWithinDir(path, dir string) bool {
|
||||||
|
|||||||
@@ -0,0 +1,117 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/rand"
|
||||||
|
"encoding/hex"
|
||||||
|
"fmt"
|
||||||
|
"io/fs"
|
||||||
|
"os"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agentu/internal/safeio"
|
||||||
|
)
|
||||||
|
|
||||||
|
// RollbackSnapshot captures the state of a file before a mutation.
|
||||||
|
type RollbackSnapshot struct {
|
||||||
|
ID string
|
||||||
|
Tool string
|
||||||
|
Path string
|
||||||
|
ResolvedPath string
|
||||||
|
Existed bool
|
||||||
|
Data []byte
|
||||||
|
Mode fs.FileMode
|
||||||
|
CreatedAt time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// RollbackStore manages in-memory rollback snapshots for the current process.
|
||||||
|
type RollbackStore struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
snapshots []RollbackSnapshot
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewRollbackStore creates an empty rollback store.
|
||||||
|
func NewRollbackStore() *RollbackStore {
|
||||||
|
return &RollbackStore{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// generateSnapshotID creates a rollback ID: "rb_" + 12 lowercase hex chars.
|
||||||
|
func generateSnapshotID() (string, error) {
|
||||||
|
b := make([]byte, 6)
|
||||||
|
if _, err := rand.Read(b); err != nil {
|
||||||
|
return "", fmt.Errorf("generate snapshot id: %w", err)
|
||||||
|
}
|
||||||
|
return "rb_" + hex.EncodeToString(b), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Capture stores the current state of the file at resolvedPath before a tool
|
||||||
|
// modifies it. Returns the snapshot ID.
|
||||||
|
func (s *RollbackStore) Capture(tool string, displayPath string, resolvedPath string) (string, error) {
|
||||||
|
id, err := generateSnapshotID()
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
snap := RollbackSnapshot{
|
||||||
|
ID: id,
|
||||||
|
Tool: tool,
|
||||||
|
Path: displayPath,
|
||||||
|
ResolvedPath: resolvedPath,
|
||||||
|
CreatedAt: time.Now().UTC(),
|
||||||
|
}
|
||||||
|
|
||||||
|
info, err := os.Stat(resolvedPath)
|
||||||
|
if err != nil {
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
snap.Existed = false
|
||||||
|
} else {
|
||||||
|
return "", fmt.Errorf("stat file: %w", err)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
snap.Existed = true
|
||||||
|
snap.Mode = info.Mode().Perm()
|
||||||
|
data, err := os.ReadFile(resolvedPath)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("read file for snapshot: %w", err)
|
||||||
|
}
|
||||||
|
snap.Data = data
|
||||||
|
}
|
||||||
|
|
||||||
|
s.mu.Lock()
|
||||||
|
s.snapshots = append(s.snapshots, snap)
|
||||||
|
s.mu.Unlock()
|
||||||
|
return id, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Restore rolls back a file to the state captured in the given snapshot.
|
||||||
|
func (s *RollbackStore) Restore(id string) (RollbackSnapshot, error) {
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
for i := len(s.snapshots) - 1; i >= 0; i-- {
|
||||||
|
if s.snapshots[i].ID == id {
|
||||||
|
snap := s.snapshots[i]
|
||||||
|
if snap.Existed {
|
||||||
|
if err := safeio.AtomicWriteFile(snap.ResolvedPath, snap.Data, snap.Mode); err != nil {
|
||||||
|
return RollbackSnapshot{}, fmt.Errorf("restore file: %w", err)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
if err := os.Remove(snap.ResolvedPath); err != nil && !os.IsNotExist(err) {
|
||||||
|
return RollbackSnapshot{}, fmt.Errorf("remove created file: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
s.snapshots = append(s.snapshots[:i], s.snapshots[i+1:]...)
|
||||||
|
return snap, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return RollbackSnapshot{}, fmt.Errorf("rollback snapshot %s not found", id)
|
||||||
|
}
|
||||||
|
|
||||||
|
// LastID returns the most recent snapshot ID, or "" if empty.
|
||||||
|
func (s *RollbackStore) LastID() string {
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
if len(s.snapshots) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return s.snapshots[len(s.snapshots)-1].ID
|
||||||
|
}
|
||||||
+52
-5
@@ -1,6 +1,7 @@
|
|||||||
package tools
|
package tools
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
@@ -8,6 +9,7 @@ import (
|
|||||||
"os"
|
"os"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
|
|
||||||
"agentu/pkg/llm"
|
"agentu/pkg/llm"
|
||||||
)
|
)
|
||||||
@@ -19,6 +21,18 @@ type ShellRunTool struct {
|
|||||||
func NewShellRunTool(workingDir string) *ShellRunTool {
|
func NewShellRunTool(workingDir string) *ShellRunTool {
|
||||||
return &ShellRunTool{workingDir: workingDir}
|
return &ShellRunTool{workingDir: workingDir}
|
||||||
}
|
}
|
||||||
|
func (*ShellRunTool) Mutates() bool { return true }
|
||||||
|
|
||||||
|
func (t *ShellRunTool) Preview(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args shellRunArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(args.Command) == "" {
|
||||||
|
return "", errors.New("command is required")
|
||||||
|
}
|
||||||
|
return "$ " + strings.TrimSpace(args.Command), ctx.Err()
|
||||||
|
}
|
||||||
|
|
||||||
func (t *ShellRunTool) Definition() llm.Tool {
|
func (t *ShellRunTool) Definition() llm.Tool {
|
||||||
return llm.Tool{
|
return llm.Tool{
|
||||||
@@ -45,6 +59,10 @@ type shellRunArgs struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
return t.ExecuteStream(ctx, raw, nil)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *ShellRunTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
|
||||||
var args shellRunArgs
|
var args shellRunArgs
|
||||||
if err := json.Unmarshal(raw, &args); err != nil {
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
return "", fmt.Errorf("parse arguments: %w", err)
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
@@ -65,16 +83,45 @@ func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string
|
|||||||
cmd := exec.CommandContext(ctx, shell, "-lc", args.Command)
|
cmd := exec.CommandContext(ctx, shell, "-lc", args.Command)
|
||||||
cmd.Dir = workdir
|
cmd.Dir = workdir
|
||||||
|
|
||||||
output, err := cmd.CombinedOutput()
|
var mu sync.Mutex
|
||||||
text := FormatResult(string(output))
|
var combined bytes.Buffer
|
||||||
|
writer := &shellStreamWriter{mu: &mu, buf: &combined, emit: emit}
|
||||||
|
cmd.Stdout = writer
|
||||||
|
cmd.Stderr = writer
|
||||||
|
|
||||||
|
runErr := cmd.Run()
|
||||||
|
text := FormatResult(combined.String())
|
||||||
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||||
return text, fmt.Errorf("command timed out: %w", ctx.Err())
|
return text, fmt.Errorf("command timed out: %w", ctx.Err())
|
||||||
}
|
}
|
||||||
if err != nil {
|
if runErr != nil {
|
||||||
if text != "" {
|
if text != "" {
|
||||||
return text, fmt.Errorf("command failed: %w", err)
|
return text, fmt.Errorf("command failed: %w", runErr)
|
||||||
}
|
}
|
||||||
return "", fmt.Errorf("command failed: %w", err)
|
return "", fmt.Errorf("command failed: %w", runErr)
|
||||||
}
|
}
|
||||||
return text, nil
|
return text, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// shellStreamWriter captures command output for the final result while
|
||||||
|
// forwarding every write to the streaming callback as it happens.
|
||||||
|
type shellStreamWriter struct {
|
||||||
|
mu *sync.Mutex
|
||||||
|
buf *bytes.Buffer
|
||||||
|
emit func(string) error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *shellStreamWriter) Write(p []byte) (int, error) {
|
||||||
|
w.mu.Lock()
|
||||||
|
n, writeErr := w.buf.Write(p)
|
||||||
|
w.mu.Unlock()
|
||||||
|
if writeErr != nil {
|
||||||
|
return n, writeErr
|
||||||
|
}
|
||||||
|
if w.emit != nil {
|
||||||
|
if err := w.emit(string(p)); err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return len(p), nil
|
||||||
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -0,0 +1,173 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"agentu/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ────────────────────────────────────────────────────────────────────────────
|
||||||
|
// TodoTool — in-memory task tracking
|
||||||
|
// ────────────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
type todoItem struct {
|
||||||
|
Content string `json:"content"`
|
||||||
|
Status string `json:"status"` // "pending", "in_progress", "completed"
|
||||||
|
}
|
||||||
|
|
||||||
|
type TodoTool struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
items []todoItem
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewTodoTool() *TodoTool {
|
||||||
|
return &TodoTool{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TodoTool) Definition() llm.Tool {
|
||||||
|
return llm.Tool{
|
||||||
|
Type: "function",
|
||||||
|
Function: llm.ToolFunction{
|
||||||
|
Name: "todo",
|
||||||
|
Description: "Manage a task list for the current conversation. Create, update, and track multi-step work. Use this to plan complex tasks and show progress.",
|
||||||
|
Parameters: JSONSchema(`{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"action": {"type": "string", "enum": ["list", "create", "update", "start", "complete", "clear"], "description": "Operation to perform."},
|
||||||
|
"items": {"type": "array", "items": {"type": "string"}, "description": "For 'create': list of task descriptions to add."},
|
||||||
|
"task": {"type": "string", "description": "For 'update'/'start'/'complete': the exact task content string to target."},
|
||||||
|
"status": {"type": "string", "enum": ["pending", "in_progress", "completed"], "description": "For 'update': new status for the task."}
|
||||||
|
},
|
||||||
|
"required": ["action"],
|
||||||
|
"additionalProperties": false
|
||||||
|
}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type todoArgs struct {
|
||||||
|
Action string `json:"action"`
|
||||||
|
Items []string `json:"items"`
|
||||||
|
Task string `json:"task"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TodoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
|
var args todoArgs
|
||||||
|
if err := json.Unmarshal(raw, &args); err != nil {
|
||||||
|
return "", fmt.Errorf("parse arguments: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
t.mu.Lock()
|
||||||
|
defer t.mu.Unlock()
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
switch args.Action {
|
||||||
|
case "list":
|
||||||
|
return t.listTasks(), nil
|
||||||
|
case "create":
|
||||||
|
return t.createTasks(args.Items)
|
||||||
|
case "start":
|
||||||
|
return t.startTask(args.Task)
|
||||||
|
case "complete":
|
||||||
|
return t.completeTask(args.Task)
|
||||||
|
case "update":
|
||||||
|
return t.updateTask(args.Task, args.Status)
|
||||||
|
case "clear":
|
||||||
|
t.items = nil
|
||||||
|
return "Tasks cleared.", nil
|
||||||
|
default:
|
||||||
|
return "", fmt.Errorf("unknown todo action: %s", args.Action)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TodoTool) listTasks() string {
|
||||||
|
if len(t.items) == 0 {
|
||||||
|
return "No tasks."
|
||||||
|
}
|
||||||
|
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString("Tasks:\n")
|
||||||
|
for _, item := range t.items {
|
||||||
|
fmt.Fprintf(&b, "[%s] %s\n", item.Status, item.Content)
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TodoTool) createTasks(items []string) (string, error) {
|
||||||
|
if len(items) == 0 {
|
||||||
|
return "", fmt.Errorf("items is required for create")
|
||||||
|
}
|
||||||
|
for _, content := range items {
|
||||||
|
content = strings.TrimSpace(content)
|
||||||
|
if content == "" {
|
||||||
|
return "", fmt.Errorf("task content must not be empty")
|
||||||
|
}
|
||||||
|
t.items = append(t.items, todoItem{
|
||||||
|
Content: content,
|
||||||
|
Status: "pending",
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return t.listTasks(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TodoTool) startTask(task string) (string, error) {
|
||||||
|
idx := -1
|
||||||
|
for i, item := range t.items {
|
||||||
|
if item.Content == task {
|
||||||
|
idx = i
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if idx < 0 {
|
||||||
|
return "", fmt.Errorf("task not found: %s", task)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Set all in_progress items to pending, then mark the target as in_progress.
|
||||||
|
for i := range t.items {
|
||||||
|
if t.items[i].Status == "in_progress" {
|
||||||
|
t.items[i].Status = "pending"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
t.items[idx].Status = "in_progress"
|
||||||
|
return t.listTasks(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TodoTool) completeTask(task string) (string, error) {
|
||||||
|
idx := -1
|
||||||
|
for i, item := range t.items {
|
||||||
|
if item.Content == task {
|
||||||
|
idx = i
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if idx < 0 {
|
||||||
|
return "", fmt.Errorf("task not found: %s", task)
|
||||||
|
}
|
||||||
|
t.items[idx].Status = "completed"
|
||||||
|
return t.listTasks(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TodoTool) updateTask(task, status string) (string, error) {
|
||||||
|
if status != "pending" && status != "in_progress" && status != "completed" {
|
||||||
|
return "", fmt.Errorf("status must be one of: pending, in_progress, completed")
|
||||||
|
}
|
||||||
|
idx := -1
|
||||||
|
for i, item := range t.items {
|
||||||
|
if item.Content == task {
|
||||||
|
idx = i
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if idx < 0 {
|
||||||
|
return "", fmt.Errorf("task not found: %s", task)
|
||||||
|
}
|
||||||
|
t.items[idx].Status = status
|
||||||
|
return t.listTasks(), nil
|
||||||
|
}
|
||||||
+78
-4
@@ -19,6 +19,54 @@ type Tool interface {
|
|||||||
Execute(ctx context.Context, args json.RawMessage) (string, error)
|
Execute(ctx context.Context, args json.RawMessage) (string, error)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// StreamingTool is optionally implemented by tools that can report partial
|
||||||
|
// output while they execute. emit receives chunks of the tool's output as
|
||||||
|
// they are produced; the returned string is still the complete result that
|
||||||
|
// gets sent back to the model.
|
||||||
|
type StreamingTool interface {
|
||||||
|
Tool
|
||||||
|
ExecuteStream(ctx context.Context, args json.RawMessage, emit func(string) error) (string, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// MutatingTool is optionally implemented by tools that modify external state.
|
||||||
|
// 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 {
|
type Registry struct {
|
||||||
tools map[string]Tool
|
tools map[string]Tool
|
||||||
primary []string
|
primary []string
|
||||||
@@ -66,44 +114,69 @@ func ReadOnlyBuiltins(workingDir string) *Registry {
|
|||||||
fileList := NewFileListTool(workingDir)
|
fileList := NewFileListTool(workingDir)
|
||||||
fileSearch := NewFileSearchTool(workingDir)
|
fileSearch := NewFileSearchTool(workingDir)
|
||||||
codeSymbols := NewCodeSymbolsTool(workingDir)
|
codeSymbols := NewCodeSymbolsTool(workingDir)
|
||||||
|
gitStatus := NewGitStatusTool(workingDir)
|
||||||
|
gitDiff := NewGitDiffTool(workingDir)
|
||||||
|
gitLog := NewGitLogTool(workingDir)
|
||||||
|
projectOutline := NewProjectOutlineTool(workingDir)
|
||||||
|
todo := NewTodoTool()
|
||||||
|
|
||||||
registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols)
|
registry := NewRegistry(fileRead, fileList, fileSearch, codeSymbols, projectOutline, todo, gitStatus, gitDiff, gitLog)
|
||||||
registry.RegisterAlias("file.read", fileRead)
|
registry.RegisterAlias("file.read", fileRead)
|
||||||
registry.RegisterAlias("file.list", fileList)
|
registry.RegisterAlias("file.list", fileList)
|
||||||
registry.RegisterAlias("file.search", fileSearch)
|
registry.RegisterAlias("file.search", fileSearch)
|
||||||
registry.RegisterAlias("code.symbols", codeSymbols)
|
registry.RegisterAlias("code.symbols", codeSymbols)
|
||||||
|
registry.RegisterAlias("git.status", gitStatus)
|
||||||
|
registry.RegisterAlias("git.diff", gitDiff)
|
||||||
|
registry.RegisterAlias("git.log", gitLog)
|
||||||
return registry
|
return registry
|
||||||
}
|
}
|
||||||
|
|
||||||
func Builtins(workingDir string) *Registry {
|
func Builtins(workingDir string) *Registry {
|
||||||
registry := ReadOnlyBuiltins(workingDir)
|
registry := ReadOnlyBuiltins(workingDir)
|
||||||
fileWrite := NewFileWriteTool(workingDir)
|
fileWrite := NewFileWriteTool(workingDir)
|
||||||
|
fileEdit := NewFileEditTool(workingDir)
|
||||||
|
applyPatch := NewApplyPatchTool(workingDir)
|
||||||
shellRun := NewShellRunTool(workingDir)
|
shellRun := NewShellRunTool(workingDir)
|
||||||
registry.Register(fileWrite)
|
registry.Register(fileWrite)
|
||||||
|
registry.Register(fileEdit)
|
||||||
|
registry.Register(applyPatch)
|
||||||
registry.Register(shellRun)
|
registry.Register(shellRun)
|
||||||
registry.RegisterAlias("file.write", fileWrite)
|
registry.RegisterAlias("file.write", fileWrite)
|
||||||
|
registry.RegisterAlias("file.edit", fileEdit)
|
||||||
|
registry.RegisterAlias("apply.patch", applyPatch)
|
||||||
registry.RegisterAlias("shell.run", shellRun)
|
registry.RegisterAlias("shell.run", shellRun)
|
||||||
return registry
|
return registry
|
||||||
}
|
}
|
||||||
|
|
||||||
func AgentInstructions(workingDir string, yolo bool) string {
|
func AgentInstructions(workingDir string, yolo bool) string {
|
||||||
|
return AgentInstructionsForMode(workingDir, yolo, yolo)
|
||||||
|
}
|
||||||
|
|
||||||
|
func AgentInstructionsForMode(workingDir string, writesEnabled bool, yolo bool) string {
|
||||||
searchInstructions := webSearchInstructions()
|
searchInstructions := webSearchInstructions()
|
||||||
readOnlyTools := "file_list, file_read, file_search, code_symbols, web_search, web_fetch"
|
readOnlyTools := "file_list, file_read, file_search, code_symbols, project_outline, todo, git_status, git_diff, git_log, web_search, web_fetch"
|
||||||
yoloTools := "file_list, file_read, file_search, code_symbols, file_write, shell_run, web_search, web_fetch"
|
writeTools := "file_list, file_read, file_search, code_symbols, project_outline, todo, file_write, file_edit, apply_patch, shell_run, git_status, git_diff, git_log, web_search, web_fetch"
|
||||||
|
if writesEnabled {
|
||||||
|
approval := "Mutating tools require interactive user approval before execution. Explain the intended change through the tool call and wait for the decision."
|
||||||
if yolo {
|
if yolo {
|
||||||
|
approval = "Mutating tools execute automatically because --yolo is enabled."
|
||||||
|
}
|
||||||
return fmt.Sprintf(`Local project context:
|
return fmt.Sprintf(`Local project context:
|
||||||
- The project working directory is %q.
|
- The project working directory is %q.
|
||||||
- Tools are available, but use them only when needed to satisfy the user's request.
|
- Tools are available, but use them only when needed to satisfy the user's request.
|
||||||
|
- %s
|
||||||
- Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project or command context.
|
- Do not call tools for greetings, small talk, simple Q&A, or general questions that do not require local project or command context.
|
||||||
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
|
- Do not explore the project preemptively. Do not call file_list just to discover context unless the user asks about project files or the answer truly depends on local files.
|
||||||
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
|
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
|
||||||
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
|
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
|
||||||
- When the user asks to run, execute, test, check, debug, or inspect a shell command, use shell_run. This includes local commands written in backticks, such as go, make, git, pwd, ls, ps, and similar terminal commands.
|
- When the user asks to run, execute, test, check, debug, or inspect a shell command, use shell_run. This includes local commands written in backticks, such as go, make, git, pwd, ls, ps, and similar terminal commands.
|
||||||
|
- Use project_outline first when the user asks for a broad understanding of a project; prefer file_edit for localized changes, apply_patch for related multi-file changes, and file_write for whole-file creation or replacement.
|
||||||
- Do not use file tools as a substitute for a command execution request.
|
- Do not use file tools as a substitute for a command execution request.
|
||||||
|
- When the user asks about changed files, diffs, or git history, use git_status, git_diff, or git_log instead of shell_run.
|
||||||
- Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions.
|
- Do not start interactive or long-running sessions. Prefer bounded, non-interactive local commands with clear completion conditions.
|
||||||
- %s
|
- %s
|
||||||
- Tool paths should normally be relative to the project working directory.
|
- Tool paths should normally be relative to the project working directory.
|
||||||
- Available tools: %s.`, workingDir, searchInstructions, yoloTools)
|
- Available tools: %s.`, workingDir, approval, searchInstructions, writeTools)
|
||||||
}
|
}
|
||||||
return fmt.Sprintf(`Local project context:
|
return fmt.Sprintf(`Local project context:
|
||||||
- The project working directory is %q.
|
- The project working directory is %q.
|
||||||
@@ -113,6 +186,7 @@ func AgentInstructions(workingDir string, yolo bool) string {
|
|||||||
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
|
- When the user asks to inspect, list, search, or read project files, use file_list, file_search, and file_read instead of saying you cannot access local files.
|
||||||
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
|
- When the user asks about code structure, symbols, functions, methods, types, or definitions in Go files, use code_symbols before reading whole files.
|
||||||
- Shell command execution is disabled in this mode. When the user asks to run, execute, test, check, debug, or inspect a shell command, do not use file tools as a substitute; explain that they need to restart agentu with --yolo to enable shell_run.
|
- Shell command execution is disabled in this mode. When the user asks to run, execute, test, check, debug, or inspect a shell command, do not use file tools as a substitute; explain that they need to restart agentu with --yolo to enable shell_run.
|
||||||
|
- When the user asks about changed files, diffs, or git history, use git_status, git_diff, or git_log instead of shell_run.
|
||||||
- %s
|
- %s
|
||||||
- Tool paths should normally be relative to the project working directory.
|
- Tool paths should normally be relative to the project working directory.
|
||||||
- Available tools: %s.`, workingDir, searchInstructions, readOnlyTools)
|
- Available tools: %s.`, workingDir, searchInstructions, readOnlyTools)
|
||||||
|
|||||||
@@ -104,12 +104,12 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, want := range []string{"file_list", "file_read", "file_search", "code_symbols"} {
|
for _, want := range []string{"file_list", "file_read", "file_search", "code_symbols", "project_outline", "todo", "git_status", "git_diff", "git_log"} {
|
||||||
if !slices.Contains(names, want) {
|
if !slices.Contains(names, want) {
|
||||||
t.Fatalf("definitions missing %s: %#v", want, names)
|
t.Fatalf("definitions missing %s: %#v", want, names)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
for _, alias := range []string{"file.list", "file.read", "file.search", "code.symbols"} {
|
for _, alias := range []string{"file.list", "file.read", "file.search", "code.symbols", "git.status", "git.diff", "git.log"} {
|
||||||
if _, ok := registry.Get(alias); !ok {
|
if _, ok := registry.Get(alias); !ok {
|
||||||
t.Fatalf("alias missing: %s", alias)
|
t.Fatalf("alias missing: %s", alias)
|
||||||
}
|
}
|
||||||
@@ -117,6 +117,9 @@ func TestReadOnlyBuiltinsExposeOpenAICompatibleNamesAndAliases(t *testing.T) {
|
|||||||
if _, ok := registry.Get("file_write"); ok {
|
if _, ok := registry.Get("file_write"); ok {
|
||||||
t.Fatal("read-only registry should not expose file_write")
|
t.Fatal("read-only registry should not expose file_write")
|
||||||
}
|
}
|
||||||
|
if _, ok := registry.Get("file_edit"); ok {
|
||||||
|
t.Fatal("read-only registry should not expose file_edit")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestCodeSymbolsToolListsGoSymbols(t *testing.T) {
|
func TestCodeSymbolsToolListsGoSymbols(t *testing.T) {
|
||||||
@@ -170,7 +173,7 @@ func TestRegistryCanExposeWebSearch(t *testing.T) {
|
|||||||
|
|
||||||
func TestBuiltinsExposeYoloTools(t *testing.T) {
|
func TestBuiltinsExposeYoloTools(t *testing.T) {
|
||||||
registry := Builtins(t.TempDir())
|
registry := Builtins(t.TempDir())
|
||||||
for _, name := range []string{"file_write", "shell_run", "file.write", "shell.run"} {
|
for _, name := range []string{"file_write", "file_edit", "apply_patch", "shell_run", "file.write", "file.edit", "apply.patch", "shell.run"} {
|
||||||
if _, ok := registry.Get(name); !ok {
|
if _, ok := registry.Get(name); !ok {
|
||||||
t.Fatalf("tool missing: %s", name)
|
t.Fatalf("tool missing: %s", name)
|
||||||
}
|
}
|
||||||
@@ -179,14 +182,14 @@ func TestBuiltinsExposeYoloTools(t *testing.T) {
|
|||||||
|
|
||||||
func TestAgentInstructionsExplainCommandModes(t *testing.T) {
|
func TestAgentInstructionsExplainCommandModes(t *testing.T) {
|
||||||
readOnly := AgentInstructions("/tmp/project", false)
|
readOnly := AgentInstructions("/tmp/project", false)
|
||||||
for _, want := range []string{"Shell command execution is disabled", "--yolo", "Do not call tools for greetings", "Do not explore the project preemptively", "Do not call file_list just to discover context"} {
|
for _, want := range []string{"Shell command execution is disabled", "--yolo", "Do not call tools for greetings", "Do not explore the project preemptively", "Do not call file_list just to discover context", "git_status, git_diff, git_log"} {
|
||||||
if !strings.Contains(readOnly, want) {
|
if !strings.Contains(readOnly, want) {
|
||||||
t.Fatalf("read-only instructions missing %q:\n%s", want, readOnly)
|
t.Fatalf("read-only instructions missing %q:\n%s", want, readOnly)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
yolo := AgentInstructions("/tmp/project", true)
|
yolo := AgentInstructions("/tmp/project", true)
|
||||||
for _, want := range []string{"use shell_run", "commands written in backticks", "Do not start interactive", "non-interactive local commands", "Do not call tools for greetings"} {
|
for _, want := range []string{"use shell_run", "commands written in backticks", "Do not start interactive", "non-interactive local commands", "Do not call tools for greetings", "prefer file_edit"} {
|
||||||
if !strings.Contains(yolo, want) {
|
if !strings.Contains(yolo, want) {
|
||||||
t.Fatalf("yolo instructions missing %q:\n%s", want, yolo)
|
t.Fatalf("yolo instructions missing %q:\n%s", want, yolo)
|
||||||
}
|
}
|
||||||
@@ -195,7 +198,7 @@ func TestAgentInstructionsExplainCommandModes(t *testing.T) {
|
|||||||
|
|
||||||
func TestAgentInstructionsExplainWebSearch(t *testing.T) {
|
func TestAgentInstructionsExplainWebSearch(t *testing.T) {
|
||||||
withSearch := AgentInstructions("/tmp/project", false)
|
withSearch := AgentInstructions("/tmp/project", false)
|
||||||
for _, want := range []string{"use web_search", "Use web_fetch", "latest information", "source URLs", "Available tools: file_list, file_read, file_search, code_symbols, web_search, web_fetch"} {
|
for _, want := range []string{"use web_search", "Use web_fetch", "latest information", "source URLs", "Available tools: file_list, file_read, file_search, code_symbols, project_outline, todo, git_status, git_diff, git_log, web_search, web_fetch"} {
|
||||||
if !strings.Contains(withSearch, want) {
|
if !strings.Contains(withSearch, want) {
|
||||||
t.Fatalf("web search instructions missing %q:\n%s", want, withSearch)
|
t.Fatalf("web search instructions missing %q:\n%s", want, withSearch)
|
||||||
}
|
}
|
||||||
@@ -257,3 +260,16 @@ func TestResolvePathRejectsDotDotThatEscapes(t *testing.T) {
|
|||||||
t.Fatal("expected error for path escaping via ..")
|
t.Fatal("expected error for path escaping via ..")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestIsMutating(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
if IsMutating(NewFileReadTool(dir)) {
|
||||||
|
t.Fatal("FileReadTool should not be mutating")
|
||||||
|
}
|
||||||
|
if !IsMutating(NewFileWriteTool(dir)) {
|
||||||
|
t.Fatal("FileWriteTool should be mutating")
|
||||||
|
}
|
||||||
|
if !IsMutating(NewShellRunTool(dir)) {
|
||||||
|
t.Fatal("ShellRunTool should be mutating")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,200 @@
|
|||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestApplyPatchChangesMultipleFiles(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, "old.txt"), []byte("before\n"), 0o600); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
patch := `*** Begin Patch
|
||||||
|
*** Update File: old.txt
|
||||||
|
@@
|
||||||
|
-before
|
||||||
|
+after
|
||||||
|
*** Add File: new.txt
|
||||||
|
+created
|
||||||
|
*** End Patch`
|
||||||
|
out, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch}))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "2 file(s) changed") {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
assertFileContent(t, filepath.Join(dir, "old.txt"), "after\n")
|
||||||
|
assertFileContent(t, filepath.Join(dir, "new.txt"), "created\n")
|
||||||
|
if info, err := os.Stat(filepath.Join(dir, "old.txt")); err != nil || info.Mode().Perm() != 0o600 {
|
||||||
|
t.Fatalf("updated mode = %v, err = %v", info.Mode().Perm(), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyPatchRollsBackEarlierFilesOnFailure(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "first.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("original\n"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
patch := `*** Begin Patch
|
||||||
|
*** Update File: first.txt
|
||||||
|
@@
|
||||||
|
-original
|
||||||
|
+changed
|
||||||
|
*** Update File: missing.txt
|
||||||
|
@@
|
||||||
|
-missing
|
||||||
|
+changed
|
||||||
|
*** End Patch`
|
||||||
|
_, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch}))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected patch failure")
|
||||||
|
}
|
||||||
|
assertFileContent(t, path, "original\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyPatchDryRunDoesNotWrite(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
patch := "*** Begin Patch\n*** Add File: dry.txt\n+dry\n*** End Patch"
|
||||||
|
out, err := NewApplyPatchTool(dir).Execute(context.Background(), mustJSON(t, map[string]any{"patch": patch, "dry_run": true}))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "dry run") {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(filepath.Join(dir, "dry.txt")); !os.IsNotExist(err) {
|
||||||
|
t.Fatalf("dry run created file: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApplyPatchRejectsMalformedEnvelopeAndDuplicatePaths(t *testing.T) {
|
||||||
|
for _, patch := range []string{
|
||||||
|
"*** Begin Patch\n*** Add File: x.txt\n+x",
|
||||||
|
"prefix\n*** Begin Patch\n*** Add File: x.txt\n+x\n*** End Patch",
|
||||||
|
"*** Begin Patch\n*** Add File: x.txt\n+x\n*** Delete File: x.txt\n*** End Patch",
|
||||||
|
} {
|
||||||
|
if _, err := parsePatch(patch); err == nil {
|
||||||
|
t.Fatalf("expected parse failure for %q", patch)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolvePathRejectsSymlinkEscape(t *testing.T) {
|
||||||
|
parent := t.TempDir()
|
||||||
|
base := filepath.Join(parent, "project")
|
||||||
|
outside := filepath.Join(parent, "outside")
|
||||||
|
if err := os.MkdirAll(base, 0o755); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := os.MkdirAll(outside, 0o755); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(filepath.Join(outside, "secret.txt"), []byte("secret"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := os.Symlink(outside, filepath.Join(base, "link")); err != nil {
|
||||||
|
t.Skipf("symlinks unavailable: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := resolvePath(base, "link/secret.txt"); err == nil {
|
||||||
|
t.Fatal("expected read symlink escape rejection")
|
||||||
|
}
|
||||||
|
if _, err := resolvePathForWrite(base, "link/new.txt"); err == nil {
|
||||||
|
t.Fatal("expected write symlink escape rejection")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTodoRejectsInvalidData(t *testing.T) {
|
||||||
|
tool := NewTodoTool()
|
||||||
|
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"create","items":[""]}`)); err == nil {
|
||||||
|
t.Fatal("expected empty task rejection")
|
||||||
|
}
|
||||||
|
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"create","items":["one"]}`)); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := tool.Execute(context.Background(), json.RawMessage(`{"action":"update","task":"one","status":"invalid"}`)); err == nil {
|
||||||
|
t.Fatal("expected invalid status rejection")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProjectOutlineIncludesGoTypeScriptAndRustSymbols(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
files := map[string]string{
|
||||||
|
"main.go": "package demo\nfunc Run() {}\n",
|
||||||
|
"ui.ts": "export class Panel {}\nexport function render() {}\n",
|
||||||
|
"lib.rs": "pub struct Worker;\npub fn execute() {}\n",
|
||||||
|
}
|
||||||
|
for name, content := range files {
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out, err := NewProjectOutlineTool(dir).Execute(context.Background(), json.RawMessage(`{}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, want := range []string{"main.go", "Run", "ui.ts", "Panel", "render", "lib.rs", "Worker", "execute"} {
|
||||||
|
if !strings.Contains(out, want) {
|
||||||
|
t.Fatalf("outline missing %q:\n%s", want, out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRollbackStoreRestoresInReverseOrder(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "state.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("one"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
store := NewRollbackStore()
|
||||||
|
first, err := store.Capture("edit", "state.txt", path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(path, []byte("two"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
second, err := store.Capture("edit", "state.txt", path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(path, []byte("three"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if store.LastID() != second {
|
||||||
|
t.Fatalf("last ID = %q", store.LastID())
|
||||||
|
}
|
||||||
|
if _, err := store.Restore(second); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
assertFileContent(t, path, "two")
|
||||||
|
if store.LastID() != first {
|
||||||
|
t.Fatalf("last ID after restore = %q, want %q", store.LastID(), first)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func mustJSON(t *testing.T, value any) json.RawMessage {
|
||||||
|
t.Helper()
|
||||||
|
data, err := json.Marshal(value)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return data
|
||||||
|
}
|
||||||
|
|
||||||
|
func assertFileContent(t *testing.T, path, want string) {
|
||||||
|
t.Helper()
|
||||||
|
data, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if string(data) != want {
|
||||||
|
t.Fatalf("%s content = %q, want %q", path, string(data), want)
|
||||||
|
}
|
||||||
|
}
|
||||||
+1120
-93
File diff suppressed because it is too large
Load Diff
+743
-30
@@ -2,13 +2,17 @@ package tui
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"fmt"
|
||||||
"os"
|
"os"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"reflect"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
"agentu/internal/agent"
|
"agentu/internal/agent"
|
||||||
|
"agentu/internal/tools"
|
||||||
"agentu/pkg/llm"
|
"agentu/pkg/llm"
|
||||||
|
|
||||||
tea "github.com/charmbracelet/bubbletea"
|
tea "github.com/charmbracelet/bubbletea"
|
||||||
@@ -21,7 +25,7 @@ func TestModelRendersChatShell(t *testing.T) {
|
|||||||
m = next.(model)
|
m = next.(model)
|
||||||
rendered := m.View()
|
rendered := m.View()
|
||||||
|
|
||||||
for _, want := range []string{"test-model", "yolo tools", "light", "ctrl+j", "alt+enter"} {
|
for _, want := range []string{"test-model"} {
|
||||||
if !strings.Contains(rendered, want) {
|
if !strings.Contains(rendered, want) {
|
||||||
t.Fatalf("rendered view missing %q:\n%s", want, rendered)
|
t.Fatalf("rendered view missing %q:\n%s", want, rendered)
|
||||||
}
|
}
|
||||||
@@ -37,18 +41,16 @@ func TestModelRendersChatShell(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestInputTextDoesNotPaintBackground(t *testing.T) {
|
func TestInputTextBackgroundMatchesInputBox(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
for name, color := range map[string]any{
|
want := lightTheme().InputBg
|
||||||
"input box": m.styles.InputBox.GetBackground(),
|
for name, color := range map[string]lipgloss.TerminalColor{
|
||||||
"focused base": m.input.FocusedStyle.Base.GetBackground(),
|
"focused base": m.input.FocusedStyle.Base.GetBackground(),
|
||||||
"focused cursor line": m.input.FocusedStyle.CursorLine.GetBackground(),
|
"focused cursor line": m.input.FocusedStyle.CursorLine.GetBackground(),
|
||||||
"blurred base": m.input.BlurredStyle.Base.GetBackground(),
|
"blurred base": m.input.BlurredStyle.Base.GetBackground(),
|
||||||
"blurred cursor line": m.input.BlurredStyle.CursorLine.GetBackground(),
|
"blurred cursor line": m.input.BlurredStyle.CursorLine.GetBackground(),
|
||||||
} {
|
} {
|
||||||
if _, ok := color.(lipgloss.NoColor); !ok {
|
requireStyleColor(t, name, color, want)
|
||||||
t.Fatalf("%s should not paint input text background", name)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -58,12 +60,12 @@ func TestInputBoxIsRoomier(t *testing.T) {
|
|||||||
m = next.(model)
|
m = next.(model)
|
||||||
|
|
||||||
if got := m.styles.InputBox.GetVerticalFrameSize(); got != 2 {
|
if got := m.styles.InputBox.GetVerticalFrameSize(); got != 2 {
|
||||||
t.Fatalf("input box vertical frame should stay compact, got %d", got)
|
t.Fatalf("input box vertical frame should be 2 with vertical padding, got %d", got)
|
||||||
}
|
}
|
||||||
if got := m.styles.InputBox.GetHorizontalFrameSize(); got < 6 {
|
if got := m.styles.InputBox.GetHorizontalFrameSize(); got != 3 {
|
||||||
t.Fatalf("input box horizontal frame should include roomier padding, got %d", got)
|
t.Fatalf("input box horizontal frame should be 3 (asymmetric padding), got %d", got)
|
||||||
}
|
}
|
||||||
if got, want := m.input.Width(), 100-m.styles.InputBox.GetHorizontalFrameSize(); got != want {
|
if got, want := m.input.Width(), 100-m.styles.InputBox.GetHorizontalFrameSize()-m.inputModeWidth()-1; got != want {
|
||||||
t.Fatalf("input content width = %d, want %d", got, want)
|
t.Fatalf("input content width = %d, want %d", got, want)
|
||||||
}
|
}
|
||||||
rendered := m.inputView()
|
rendered := m.inputView()
|
||||||
@@ -78,13 +80,8 @@ func TestModelMetaRendersBelowInput(t *testing.T) {
|
|||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
rendered := next.(model).View()
|
rendered := next.(model).View()
|
||||||
|
|
||||||
inputIndex := strings.Index(rendered, "Message agentu")
|
if !strings.Contains(rendered, "test-model") {
|
||||||
modelIndex := strings.Index(rendered, "test-model")
|
t.Fatalf("rendered view missing %q:\n%s", "test-model", rendered)
|
||||||
if inputIndex < 0 || modelIndex < 0 {
|
|
||||||
t.Fatalf("rendered view missing input or model meta:\n%s", rendered)
|
|
||||||
}
|
|
||||||
if modelIndex < inputIndex {
|
|
||||||
t.Fatalf("model meta should render below input:\n%s", rendered)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -93,7 +90,7 @@ func TestModelMetaRendersContextUsage(t *testing.T) {
|
|||||||
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("x", 800)}})
|
assistant.SetMessages([]llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("x", 800)}})
|
||||||
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
|
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
|
||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 120, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 120, Height: 30})
|
||||||
rendered := next.(model).modelMetaView()
|
rendered := next.(model).footerView()
|
||||||
|
|
||||||
for _, want := range []string{"ctx ~", "/1k", "%"} {
|
for _, want := range []string{"ctx ~", "/1k", "%"} {
|
||||||
if !strings.Contains(rendered, want) {
|
if !strings.Contains(rendered, want) {
|
||||||
@@ -123,11 +120,124 @@ func TestRunningStateRendersActivityIndicator(t *testing.T) {
|
|||||||
m.resize()
|
m.resize()
|
||||||
|
|
||||||
rendered := m.View()
|
rendered := m.View()
|
||||||
for _, want := range []string{"Working", "thinking...", "esc cancels"} {
|
for _, want := range []string{iconWorking, "Working", "thinking...", "esc cancels"} {
|
||||||
if !strings.Contains(rendered, want) {
|
if !strings.Contains(rendered, want) {
|
||||||
t.Fatalf("running view missing %q:\n%s", want, rendered)
|
t.Fatalf("running view missing %q:\n%s", want, rendered)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if backgroundPattern.MatchString(m.activityView()) {
|
||||||
|
t.Fatalf("activity view should not paint a background:\n%q", m.activityView())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type tuiBTWProvider struct{}
|
||||||
|
|
||||||
|
func (tuiBTWProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||||
|
return emit(llm.StreamEvent{Content: "side answer"})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBTWCanBeSubmittedWhileMainTurnIsRunning(t *testing.T) {
|
||||||
|
assistant := agent.New(agent.Options{Provider: tuiBTWProvider{}, Model: "test-model"})
|
||||||
|
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
|
||||||
|
m.running = true
|
||||||
|
m.events = make(chan tea.Msg, 8)
|
||||||
|
|
||||||
|
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("/btw why?")})
|
||||||
|
m = next.(model)
|
||||||
|
if m.input.Value() != "/btw why?" {
|
||||||
|
t.Fatalf("running input = %q", m.input.Value())
|
||||||
|
}
|
||||||
|
next, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
|
||||||
|
m = next.(model)
|
||||||
|
if !m.btwRunning || m.input.Value() != "" {
|
||||||
|
t.Fatalf("btw running/input = %v/%q", m.btwRunning, m.input.Value())
|
||||||
|
}
|
||||||
|
if got := m.messages[len(m.messages)-1]; got.role != roleBTWUser || got.content != "why?" {
|
||||||
|
t.Fatalf("btw question message = %+v", got)
|
||||||
|
}
|
||||||
|
if cmd == nil {
|
||||||
|
t.Fatal("btw did not return a start command")
|
||||||
|
}
|
||||||
|
_ = cmd()
|
||||||
|
chunk := <-m.events
|
||||||
|
done := <-m.events
|
||||||
|
if _, ok := chunk.(btwChunkMsg); !ok {
|
||||||
|
t.Fatalf("first btw event = %T", chunk)
|
||||||
|
}
|
||||||
|
if _, ok := done.(btwDoneMsg); !ok {
|
||||||
|
t.Fatalf("final btw event = %T", done)
|
||||||
|
}
|
||||||
|
if len(assistant.Messages()) != 0 {
|
||||||
|
t.Fatalf("btw polluted agent history: %+v", assistant.Messages())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBTWRequiresQuestion(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.input.SetValue("/btw")
|
||||||
|
next, _ := m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
if m.btwRunning {
|
||||||
|
t.Fatal("empty btw question started a request")
|
||||||
|
}
|
||||||
|
last := m.messages[len(m.messages)-1]
|
||||||
|
if last.role != roleBTW || !strings.Contains(last.content, "Usage: /btw") {
|
||||||
|
t.Fatalf("empty btw message = %+v", last)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBTWMessagesHaveDistinctLabels(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
question := stripANSI(m.messageView(roleBTWUser, "question", 80))
|
||||||
|
answer := stripANSI(m.messageView(roleBTW, "answer", 80))
|
||||||
|
if !strings.Contains(question, "BTW · You") || !strings.Contains(answer, "BTW") {
|
||||||
|
t.Fatalf("btw labels missing:\n%s\n%s", question, answer)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBTWAndMainStreamChunksStaySeparate(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.messages = append(m.messages, message{role: roleUser, content: "main"})
|
||||||
|
m.appendAssistantChunk("main one ")
|
||||||
|
m.messages = append(m.messages, message{role: roleBTWUser, content: "side"})
|
||||||
|
m.appendBTWChunk("side one ")
|
||||||
|
m.appendAssistantChunk("main two")
|
||||||
|
m.appendBTWChunk("side two")
|
||||||
|
|
||||||
|
var mainAnswers, sideAnswers []string
|
||||||
|
for _, msg := range m.messages {
|
||||||
|
switch msg.role {
|
||||||
|
case roleAssistant:
|
||||||
|
mainAnswers = append(mainAnswers, msg.content)
|
||||||
|
case roleBTW:
|
||||||
|
sideAnswers = append(sideAnswers, msg.content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !reflect.DeepEqual(mainAnswers, []string{"main one main two"}) {
|
||||||
|
t.Fatalf("main answers = %#v", mainAnswers)
|
||||||
|
}
|
||||||
|
if !reflect.DeepEqual(sideAnswers, []string{"side one side two"}) {
|
||||||
|
t.Fatalf("side answers = %#v", sideAnswers)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMainCompletionKeepsBTWEventStreamAlive(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
events := make(chan tea.Msg, 1)
|
||||||
|
m.events = events
|
||||||
|
m.running = true
|
||||||
|
m.btwRunning = true
|
||||||
|
|
||||||
|
next, _ := m.Update(turnDoneMsg{})
|
||||||
|
m = next.(model)
|
||||||
|
if m.running || !m.btwRunning || m.events != events {
|
||||||
|
t.Fatalf("state after main completion: running=%v btw=%v events=%p", m.running, m.btwRunning, m.events)
|
||||||
|
}
|
||||||
|
next, _ = m.Update(btwDoneMsg{})
|
||||||
|
m = next.(model)
|
||||||
|
if m.active() || m.events != nil {
|
||||||
|
t.Fatalf("state after btw completion: active=%v events=%p", m.active(), m.events)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestUserMessageRendersFullWidthBubble(t *testing.T) {
|
func TestUserMessageRendersFullWidthBubble(t *testing.T) {
|
||||||
@@ -276,19 +386,219 @@ func TestInputStartsAtOneLineAndGrows(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSlashCommandSuggestions(t *testing.T) {
|
func TestCommandPaletteFuzzyMatches(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||||
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
m = next.(model)
|
m = next.(model)
|
||||||
m.input.SetValue("/")
|
m.input.SetValue("/mdl")
|
||||||
m.afterInputChanged()
|
m.afterInputChanged()
|
||||||
|
|
||||||
rendered := m.View()
|
plain := stripANSI(m.View())
|
||||||
for _, want := range []string{"/clear reset context", "/compact compress context", "/exit quit", "/model model/provider/thinking"} {
|
for _, want := range []string{iconCommand + " Commands", "/model model/provider/thinking"} {
|
||||||
if !strings.Contains(rendered, want) {
|
if !strings.Contains(plain, want) {
|
||||||
t.Fatalf("rendered suggestions missing %q:\n%s", want, rendered)
|
t.Fatalf("rendered palette missing %q:\n%s", want, plain)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
matches := commandMatches("m")
|
||||||
|
modelIndex := -1
|
||||||
|
renameIndex := -1
|
||||||
|
for i, match := range matches {
|
||||||
|
switch match.command.Name {
|
||||||
|
case "/model":
|
||||||
|
modelIndex = i
|
||||||
|
case "/rename":
|
||||||
|
renameIndex = i
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if modelIndex < 0 || renameIndex < 0 || modelIndex > renameIndex {
|
||||||
|
t.Fatalf("/model should sort before /rename for /m query: %#v", matches)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInputHistoryNavigation(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.rememberInput("first")
|
||||||
|
m.rememberInput("second")
|
||||||
|
|
||||||
|
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlP})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "second" {
|
||||||
|
t.Fatalf("after ctrl+p input = %q", got)
|
||||||
|
}
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlP})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "first" {
|
||||||
|
t.Fatalf("after second ctrl+p input = %q", got)
|
||||||
|
}
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlN})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "second" {
|
||||||
|
t.Fatalf("after ctrl+n input = %q", got)
|
||||||
|
}
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyCtrlN})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "" {
|
||||||
|
t.Fatalf("after second ctrl+n input = %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpArrowRecallsPreviousInput(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.rememberInput("first")
|
||||||
|
m.rememberInput("second")
|
||||||
|
|
||||||
|
next, _ := m.Update(tea.KeyMsg{Type: tea.KeyUp})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "second" {
|
||||||
|
t.Fatalf("after first up input = %q, want %q", got, "second")
|
||||||
|
}
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyUp})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "first" {
|
||||||
|
t.Fatalf("after second up input = %q, want %q", got, "first")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolMessageRendersStructuredCard(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`})
|
||||||
|
|
||||||
|
plain := stripANSI(m.renderMessages())
|
||||||
|
for _, want := range []string{iconTool + " file_read", "README.md"} {
|
||||||
|
if !strings.Contains(plain, want) {
|
||||||
|
t.Fatalf("tool card missing %q:\n%s", want, plain)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolMessageRendersOutputPreview(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.messages = append(m.messages, message{role: roleTool, content: "file_read {\"path\":\"README.md\"}\n[output]\nfirst line\nsecond line"})
|
||||||
|
|
||||||
|
plain := stripANSI(m.renderMessages())
|
||||||
|
for _, want := range []string{iconTool + " file_read", "README.md", "output", "first line", "second line"} {
|
||||||
|
if !strings.Contains(plain, want) {
|
||||||
|
t.Fatalf("tool card missing %q:\n%s", want, plain)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolLogOutputUpdatesExistingToolCard(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.messages = append(m.messages, message{role: roleTool, content: `file_read {"path":"README.md"}`})
|
||||||
|
|
||||||
|
updated := m.mergeToolLog("file_read {\"path\":\"README.md\"}\n[output]\nhello")
|
||||||
|
if !updated {
|
||||||
|
t.Fatal("expected mergeToolLog to return true")
|
||||||
|
}
|
||||||
|
if len(m.messages) != 1 {
|
||||||
|
t.Fatalf("messages count = %d, want 1", len(m.messages))
|
||||||
|
}
|
||||||
|
plain := stripANSI(m.renderMessages())
|
||||||
|
if !strings.Contains(plain, "hello") {
|
||||||
|
t.Fatalf("rendered output missing 'hello':\n%s", plain)
|
||||||
|
}
|
||||||
|
|
||||||
|
status := toolStatusText("file_read {\"path\":\"README.md\"}\n[output]\nhello world")
|
||||||
|
if strings.Contains(status, "hello") {
|
||||||
|
t.Fatalf("status should not contain output: %q", status)
|
||||||
|
}
|
||||||
|
if !strings.Contains(status, "file_read") {
|
||||||
|
t.Fatalf("status missing tool name: %q", status)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolLogStreamingChunksAppendToSameCard(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.messages = append(m.messages, message{role: roleTool, content: `shell_run echo hi`})
|
||||||
|
|
||||||
|
if !m.mergeToolLog("shell_run echo hi\n[output]\none\n") {
|
||||||
|
t.Fatal("expected first mergeToolLog to return true")
|
||||||
|
}
|
||||||
|
if !m.mergeToolLog("shell_run echo hi\n[output]\ntwo\n") {
|
||||||
|
t.Fatal("expected second mergeToolLog to return true")
|
||||||
|
}
|
||||||
|
if len(m.messages) != 1 {
|
||||||
|
t.Fatalf("messages count = %d, want 1", len(m.messages))
|
||||||
|
}
|
||||||
|
plain := stripANSI(m.renderMessages())
|
||||||
|
if !strings.Contains(plain, "one") || !strings.Contains(plain, "two") {
|
||||||
|
t.Fatalf("rendered output missing streamed lines:\n%s", plain)
|
||||||
|
}
|
||||||
|
if strings.Contains(plain, "onetwo") {
|
||||||
|
t.Fatalf("streamed lines were concatenated:\n%s", plain)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestShellRunCardShowsCommandWorkdirAndOutput(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
m.messages = append(m.messages, message{
|
||||||
|
role: roleTool,
|
||||||
|
content: "shell_run {\"command\":\"go test ./...\",\"workdir\":\"internal/tui\"}\n[output]\nok agentu/internal/tui",
|
||||||
|
})
|
||||||
|
|
||||||
|
plain := stripANSI(m.renderMessages())
|
||||||
|
for _, want := range []string{"$ go test ./...", "cwd internal/tui", "output", "ok agentu/internal/tui"} {
|
||||||
|
if !strings.Contains(plain, want) {
|
||||||
|
t.Fatalf("shell card missing %q:\n%s", want, plain)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if strings.Contains(plain, `{"command"`) || strings.Contains(plain, iconTool+" shell_run") {
|
||||||
|
t.Fatalf("shell card still exposes raw tool arguments:\n%s", plain)
|
||||||
|
}
|
||||||
|
if status := toolStatusText(`shell_run {"command":"go test ./...","workdir":"internal/tui"}`); status != "$ go test ./..." {
|
||||||
|
t.Fatalf("shell status = %q", status)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestShellRunStreamingEventsPreserveLineBreaks(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
args := `{"command":"printf lines"}`
|
||||||
|
m.messages = append(m.messages, message{role: roleTool, content: "shell_run " + args})
|
||||||
|
|
||||||
|
next, _ := m.Update(toolLogMsg("\n[tool] shell_run " + args + "\n[output]\none\n"))
|
||||||
|
m = next.(model)
|
||||||
|
next, _ = m.Update(toolLogMsg("\n[tool] shell_run " + args + "\n[output]\n\n"))
|
||||||
|
m = next.(model)
|
||||||
|
next, _ = m.Update(toolLogMsg("\n[tool] shell_run " + args + "\n[output]\ntwo\n"))
|
||||||
|
m = next.(model)
|
||||||
|
|
||||||
|
parts := parseToolLog(m.messages[0].content)
|
||||||
|
if parts.output != "one\n\ntwo" {
|
||||||
|
t.Fatalf("streamed output = %q", parts.output)
|
||||||
|
}
|
||||||
|
if len(m.messages) != 1 {
|
||||||
|
t.Fatalf("streaming created %d tool cards, want 1", len(m.messages))
|
||||||
|
}
|
||||||
|
plain := stripANSI(m.renderMessages())
|
||||||
|
if strings.Contains(plain, "onetwo") {
|
||||||
|
t.Fatalf("streamed lines were concatenated:\n%s", plain)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestShellOutputPreviewKeepsHeadAndLatestLines(t *testing.T) {
|
||||||
|
var lines []string
|
||||||
|
for i := 1; i <= 12; i++ {
|
||||||
|
lines = append(lines, fmt.Sprintf("line %d", i))
|
||||||
|
}
|
||||||
|
preview := strings.Join(shellOutputPreviewLines(strings.Join(lines, "\n"), 80), "\n")
|
||||||
|
for _, want := range []string{"line 1", "line 4", "5 lines omitted", "line 10", "line 12"} {
|
||||||
|
if !strings.Contains(preview, want) {
|
||||||
|
t.Fatalf("preview missing %q:\n%s", want, preview)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if strings.Contains(preview, "line 5\n") || strings.Contains(preview, "line 9\n") {
|
||||||
|
t.Fatalf("preview contains a line that should be omitted:\n%s", preview)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSanitizeTerminalOutputRemovesEscapesAndFoldsCarriageReturns(t *testing.T) {
|
||||||
|
got := sanitizeTerminalOutput("\x1b[31mred\x1b[0m\nprogress 10%\rprogress 20%\r\nclean\x00text")
|
||||||
|
want := "red\nprogress 20%\ncleantext"
|
||||||
|
if got != want {
|
||||||
|
t.Fatalf("sanitized output = %q, want %q", got, want)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
type tuiCompactProvider struct{}
|
type tuiCompactProvider struct{}
|
||||||
@@ -386,16 +696,60 @@ func runInDir(t *testing.T, dir string, name string, args ...string) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestBangShellCommandRunsFromInput(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true, WorkingDir: dir})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
m.input.SetValue("! printf hello")
|
||||||
|
next, _ = m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
if m.input.Value() != "" {
|
||||||
|
t.Fatalf("input not cleared after bang command: %q", m.input.Value())
|
||||||
|
}
|
||||||
|
rendered := stripANSI(m.renderMessages())
|
||||||
|
if !strings.Contains(rendered, "$ printf hello") {
|
||||||
|
t.Fatalf("rendered output missing shell echo:\n%s", rendered)
|
||||||
|
}
|
||||||
|
if !strings.Contains(rendered, "hello") {
|
||||||
|
t.Fatalf("rendered output missing 'hello':\n%s", rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBangShellCommandRequiresYolo(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: false})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
m.input.SetValue("! printf hello")
|
||||||
|
next, _ = m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
rendered := stripANSI(m.renderMessages())
|
||||||
|
if !strings.Contains(rendered, "shell command input requires --yolo") {
|
||||||
|
t.Fatalf("rendered output missing yolo requirement:\n%s", rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBangShellCommandRejectsEmptyCommand(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
m.input.SetValue("!")
|
||||||
|
next, _ = m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
rendered := stripANSI(m.renderMessages())
|
||||||
|
if !strings.Contains(rendered, "usage: ! <command>") {
|
||||||
|
t.Fatalf("rendered output missing usage hint:\n%s", rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestModelCommand(t *testing.T) {
|
func TestModelCommand(t *testing.T) {
|
||||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||||
m.input.SetValue("/model")
|
m.input.SetValue("/model")
|
||||||
|
|
||||||
next, _ := m.submit()
|
next, _ := m.submit()
|
||||||
rendered := next.(model).renderMessages()
|
rendered := next.(model).renderMessages()
|
||||||
for _, want := range []string{"model: test-model", "mode: yolo tools", "theme: light"} {
|
if !strings.Contains(rendered, "test-model") {
|
||||||
if !strings.Contains(rendered, want) {
|
t.Fatalf("model command output missing %q:\n%s", "test-model", rendered)
|
||||||
t.Fatalf("model command output missing %q:\n%s", want, rendered)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -414,11 +768,124 @@ func TestModelThinkingCommand(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestModelCommandPaletteShowsConfiguredModels(t *testing.T) {
|
||||||
|
manager := &fakeModelManager{model: "model-a", models: []string{"model-a", "model-b"}}
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "model-a", ModelManager: manager})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
m.input.SetValue("/model ")
|
||||||
|
m.afterInputChanged()
|
||||||
|
plain := stripANSI(m.View())
|
||||||
|
for _, want := range []string{iconCommand + " Model candidates", "model-a", "current model", "model-b", "switch model"} {
|
||||||
|
if !strings.Contains(plain, want) {
|
||||||
|
t.Fatalf("rendered palette missing %q:\n%s", want, plain)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelCommandPaletteSelectsConfiguredModel(t *testing.T) {
|
||||||
|
manager := &fakeModelManager{model: "model-a", models: []string{"model-a", "model-b"}}
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "model-a", ModelManager: manager})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
m.input.SetValue("/model ")
|
||||||
|
m.afterInputChanged()
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyDown})
|
||||||
|
m = next.(model)
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyEnter})
|
||||||
|
m = next.(model)
|
||||||
|
if manager.model != "model-b" {
|
||||||
|
t.Fatalf("model = %q", manager.model)
|
||||||
|
}
|
||||||
|
if got := m.input.Value(); got != "" {
|
||||||
|
t.Fatalf("input = %q", got)
|
||||||
|
}
|
||||||
|
rendered := stripANSI(m.renderMessages())
|
||||||
|
if !strings.Contains(rendered, "model: model-b") {
|
||||||
|
t.Fatalf("rendered output missing model switch:\n%s", rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelCommandPaletteFiltersBareModelPartial(t *testing.T) {
|
||||||
|
manager := &fakeModelManager{model: "claude-sonnet", models: []string{"claude-sonnet", "deepseek-v4-flash"}}
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "claude-sonnet", ModelManager: manager})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
m.input.SetValue("/model deep")
|
||||||
|
m.afterInputChanged()
|
||||||
|
palette := stripANSI(m.commandPaletteView())
|
||||||
|
if !strings.Contains(palette, "deepseek-v4-flash") {
|
||||||
|
t.Fatalf("expected deepseek-v4-flash in palette:\n%s", palette)
|
||||||
|
}
|
||||||
|
if strings.Contains(palette, "claude-sonnet") {
|
||||||
|
t.Fatalf("unexpected claude-sonnet in palette:\n%s", palette)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelCommandPaletteFiltersExplicitModelSubcommand(t *testing.T) {
|
||||||
|
manager := &fakeModelManager{model: "claude-sonnet", models: []string{"claude-sonnet", "deepseek-v4-flash"}}
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "claude-sonnet", ModelManager: manager})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
m.input.SetValue("/model model deep")
|
||||||
|
m.afterInputChanged()
|
||||||
|
palette := stripANSI(m.commandPaletteView())
|
||||||
|
if !strings.Contains(palette, "deepseek-v4-flash") {
|
||||||
|
t.Fatalf("expected deepseek-v4-flash in palette:\n%s", palette)
|
||||||
|
}
|
||||||
|
if strings.Contains(palette, "claude-sonnet") {
|
||||||
|
t.Fatalf("unexpected claude-sonnet in palette:\n%s", palette)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelCommandPaletteTabFillsConfiguredModel(t *testing.T) {
|
||||||
|
manager := &fakeModelManager{model: "claude-sonnet", models: []string{"claude-sonnet", "deepseek-v4-flash"}}
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "claude-sonnet", ModelManager: manager})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
m.input.SetValue("/model deep")
|
||||||
|
m.afterInputChanged()
|
||||||
|
next, _ = m.Update(tea.KeyMsg{Type: tea.KeyTab})
|
||||||
|
m = next.(model)
|
||||||
|
if got := m.input.Value(); got != "/model model deepseek-v4-flash" {
|
||||||
|
t.Fatalf("input = %q", got)
|
||||||
|
}
|
||||||
|
if manager.model != "claude-sonnet" {
|
||||||
|
t.Fatalf("model should not change on tab: %q", manager.model)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelCommandPaletteIgnoresProviderAndThinkingSubcommands(t *testing.T) {
|
||||||
|
manager := &fakeModelManager{model: "model-a", models: []string{"model-a", "model-b"}}
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "model-a", ModelManager: manager})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
|
||||||
|
m.input.SetValue("/model provider ")
|
||||||
|
m.afterInputChanged()
|
||||||
|
if sugs := m.commandSuggestions(); len(sugs) != 0 {
|
||||||
|
t.Fatalf("provider suggestions should be empty, got %d", len(sugs))
|
||||||
|
}
|
||||||
|
if view := m.commandPaletteView(); view != "" {
|
||||||
|
t.Fatalf("provider palette should be empty:\n%s", view)
|
||||||
|
}
|
||||||
|
|
||||||
|
m.input.SetValue("/model thinking ")
|
||||||
|
m.afterInputChanged()
|
||||||
|
if sugs := m.commandSuggestions(); len(sugs) != 0 {
|
||||||
|
t.Fatalf("thinking suggestions should be empty, got %d", len(sugs))
|
||||||
|
}
|
||||||
|
if view := m.commandPaletteView(); view != "" {
|
||||||
|
t.Fatalf("thinking palette should be empty:\n%s", view)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
type fakeModelManager struct {
|
type fakeModelManager struct {
|
||||||
provider string
|
provider string
|
||||||
model string
|
model string
|
||||||
thinking string
|
thinking string
|
||||||
sessionName string
|
sessionName string
|
||||||
|
models []string
|
||||||
saved bool
|
saved bool
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -489,6 +956,16 @@ func (f *fakeModelManager) NewSession() (string, error) {
|
|||||||
return "New session.", nil
|
return "New session.", nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (f *fakeModelManager) AvailableModels() []string {
|
||||||
|
if len(f.models) > 0 {
|
||||||
|
return append([]string(nil), f.models...)
|
||||||
|
}
|
||||||
|
if f.model != "" {
|
||||||
|
return []string{f.model}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
func TestParseThemeMode(t *testing.T) {
|
func TestParseThemeMode(t *testing.T) {
|
||||||
for input, want := range map[string]ThemeMode{
|
for input, want := range map[string]ThemeMode{
|
||||||
"": ThemeLight,
|
"": ThemeLight,
|
||||||
@@ -510,6 +987,130 @@ func TestParseThemeMode(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func requireStyleColor(t *testing.T, name string, got lipgloss.TerminalColor, want string) {
|
||||||
|
t.Helper()
|
||||||
|
color, ok := got.(lipgloss.Color)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("%s = %T, want lipgloss.Color", name, got)
|
||||||
|
}
|
||||||
|
if string(color) != want {
|
||||||
|
t.Fatalf("%s = %q, want %q", name, string(color), want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func requireNoStyleColor(t *testing.T, name string, got lipgloss.TerminalColor) {
|
||||||
|
t.Helper()
|
||||||
|
if _, ok := got.(lipgloss.NoColor); !ok {
|
||||||
|
t.Fatalf("%s should not paint a color, got %T", name, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDarkThemePaintsAppBackground(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
rendered := stripANSI(m.View())
|
||||||
|
_ = rendered
|
||||||
|
if m.theme.Mode != ThemeDark {
|
||||||
|
t.Fatalf("theme.Mode = %q, want %q", m.theme.Mode, ThemeDark)
|
||||||
|
}
|
||||||
|
want := darkTheme()
|
||||||
|
requireStyleColor(t, "App.Background", m.styles.App.GetBackground(), want.Background)
|
||||||
|
requireStyleColor(t, "App.Foreground", m.styles.App.GetForeground(), want.Text)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDarkThemeKeepsComponentBackgroundsLightweight(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
requireNoStyleColor(t, "Viewport.Background", m.styles.Viewport.GetBackground())
|
||||||
|
requireStyleColor(t, "InputBox.Background", m.styles.InputBox.GetBackground(), darkTheme().InputBg)
|
||||||
|
requireNoStyleColor(t, "Activity.Background", m.styles.Activity.GetBackground())
|
||||||
|
m.running = true
|
||||||
|
m.status = "thinking..."
|
||||||
|
m.resize()
|
||||||
|
if backgroundPattern.MatchString(m.activityView()) {
|
||||||
|
t.Fatalf("activity view should not paint a background:\n%q", m.activityView())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDarkThemeStylesUseDarkPalette(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", ThemeMode: ThemeDark})
|
||||||
|
want := darkTheme()
|
||||||
|
requireStyleColor(t, "Header.Foreground", m.styles.Header.GetForeground(), want.Title)
|
||||||
|
requireStyleColor(t, "Muted.Foreground", m.styles.Muted.GetForeground(), want.Muted)
|
||||||
|
requireStyleColor(t, "Pill.Background", m.styles.Pill.GetBackground(), want.SurfaceAlt)
|
||||||
|
requireStyleColor(t, "input.FocusedStyle.Base.Foreground", m.input.FocusedStyle.Base.GetForeground(), want.Text)
|
||||||
|
requireStyleColor(t, "input.FocusedStyle.Placeholder.Foreground", m.input.FocusedStyle.Placeholder.GetForeground(), want.Muted)
|
||||||
|
requireStyleColor(t, "InputModeYolo.Foreground", m.styles.InputModeYolo.GetForeground(), want.Error)
|
||||||
|
requireStyleColor(t, "InputModeCommand.Foreground", m.styles.InputModeCommand.GetForeground(), want.User)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInputModeSymbolReflectsYoloAndBangCommand(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
|
||||||
|
rendered := stripANSI(m.inputView())
|
||||||
|
if !strings.Contains(rendered, "Y") {
|
||||||
|
t.Fatalf("YOLO input view should contain Y symbol:\n%s", rendered)
|
||||||
|
}
|
||||||
|
|
||||||
|
m.input.SetValue("! pwd")
|
||||||
|
m.afterInputChanged()
|
||||||
|
rendered = stripANSI(m.inputView())
|
||||||
|
if !strings.Contains(rendered, iconCommand) {
|
||||||
|
t.Fatalf("bang input view should contain command icon:\n%s", rendered)
|
||||||
|
}
|
||||||
|
if strings.Contains(rendered, "Y") {
|
||||||
|
t.Fatalf("bang input view should not contain Y symbol:\n%s", rendered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInputModeYoloUsesWarningColor(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Yolo: true})
|
||||||
|
requireStyleColor(t, "InputModeYolo.Foreground", m.styles.InputModeYolo.GetForeground(), lightTheme().Error)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestThemeSlashCommandSwitchesAtRuntime(t *testing.T) {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
if m.theme.Mode != ThemeLight {
|
||||||
|
t.Fatalf("initial theme = %q, want light", m.theme.Mode)
|
||||||
|
}
|
||||||
|
requireNoStyleColor(t, "initial App.Background", m.styles.App.GetBackground())
|
||||||
|
|
||||||
|
// Switch to dark
|
||||||
|
m.input.SetValue("/theme dark")
|
||||||
|
next, _ := m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
if m.theme.Mode != ThemeDark {
|
||||||
|
t.Fatalf("after /theme dark: Mode = %q, want dark", m.theme.Mode)
|
||||||
|
}
|
||||||
|
requireStyleColor(t, "dark App.Background", m.styles.App.GetBackground(), darkTheme().Background)
|
||||||
|
last := m.messages[len(m.messages)-1]
|
||||||
|
if !strings.Contains(last.content, "theme: dark") {
|
||||||
|
t.Fatalf("expected confirmation message, got %q", last.content)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Switch back to light
|
||||||
|
m.input.SetValue("/theme light")
|
||||||
|
next, _ = m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
if m.theme.Mode != ThemeLight {
|
||||||
|
t.Fatalf("after /theme light: Mode = %q, want light", m.theme.Mode)
|
||||||
|
}
|
||||||
|
requireNoStyleColor(t, "restored App.Background", m.styles.App.GetBackground())
|
||||||
|
|
||||||
|
// No arg shows usage
|
||||||
|
m.input.SetValue("/theme")
|
||||||
|
next, _ = m.submit()
|
||||||
|
m = next.(model)
|
||||||
|
last = m.messages[len(m.messages)-1]
|
||||||
|
if !strings.Contains(last.content, "usage: /theme") {
|
||||||
|
t.Fatalf("expected usage hint, got %q", last.content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestCleanToolLog(t *testing.T) {
|
func TestCleanToolLog(t *testing.T) {
|
||||||
got := cleanToolLog("\n[tool] shell_run {\"command\":\"pwd\"}\n")
|
got := cleanToolLog("\n[tool] shell_run {\"command\":\"pwd\"}\n")
|
||||||
want := `shell_run {"command":"pwd"}`
|
want := `shell_run {"command":"pwd"}`
|
||||||
@@ -517,3 +1118,115 @@ func TestCleanToolLog(t *testing.T) {
|
|||||||
t.Fatalf("got %q want %q", got, want)
|
t.Fatalf("got %q want %q", got, want)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestApprovalPromptAllowsAndDeniesToolCalls(t *testing.T) {
|
||||||
|
for _, tc := range []struct {
|
||||||
|
key tea.KeyMsg
|
||||||
|
want agent.ToolAuthorizationDecision
|
||||||
|
}{
|
||||||
|
{key: tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'y'}}, want: agent.ToolAuthorizationAllow},
|
||||||
|
{key: tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'n'}}, want: agent.ToolAuthorizationDeny},
|
||||||
|
} {
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||||
|
next, _ := m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
|
||||||
|
m = next.(model)
|
||||||
|
responses := make(chan agent.ToolAuthorizationDecision, 1)
|
||||||
|
next, _ = m.Update(toolApprovalMsg{
|
||||||
|
request: agent.ToolAuthorizationRequest{ToolName: "file_write", Preview: "write 5 bytes", RollbackAvailable: true},
|
||||||
|
respond: responses,
|
||||||
|
})
|
||||||
|
m = next.(model)
|
||||||
|
view := stripANSI(m.View())
|
||||||
|
for _, want := range []string{"Approve file_write?", "write 5 bytes", "rollback snapshot", "y approve"} {
|
||||||
|
if !strings.Contains(view, want) {
|
||||||
|
t.Fatalf("approval view missing %q:\n%s", want, view)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
next, _ = m.Update(tc.key)
|
||||||
|
m = next.(model)
|
||||||
|
if m.approval != nil {
|
||||||
|
t.Fatal("approval state was not cleared")
|
||||||
|
}
|
||||||
|
if got := <-responses; got != tc.want {
|
||||||
|
t.Fatalf("decision = %q, want %q", got, tc.want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestApprovalBrokerBridgesAgentAndTUI(t *testing.T) {
|
||||||
|
broker := NewApprovalBroker()
|
||||||
|
events := make(chan tea.Msg, 1)
|
||||||
|
broker.SetEvents(events)
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||||
|
defer cancel()
|
||||||
|
result := make(chan agent.ToolAuthorizationDecision, 1)
|
||||||
|
go func() {
|
||||||
|
decision, _ := broker.AuthorizeToolCall(ctx, agent.ToolAuthorizationRequest{ToolName: "file_edit"})
|
||||||
|
result <- decision
|
||||||
|
}()
|
||||||
|
msg, ok := (<-events).(toolApprovalMsg)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("approval broker emitted unexpected message")
|
||||||
|
}
|
||||||
|
msg.respond <- agent.ToolAuthorizationAllow
|
||||||
|
if decision := <-result; decision != agent.ToolAuthorizationAllow {
|
||||||
|
t.Fatalf("decision = %q", decision)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRollbackCommandRestoresLatestSnapshot(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "state.txt")
|
||||||
|
if err := os.WriteFile(path, []byte("before"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
rollback := tools.NewRollbackStore()
|
||||||
|
if _, err := rollback.Capture("file_edit", "state.txt", path); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(path, []byte("after"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
m := newModel(context.Background(), nil, Options{ModelName: "test-model", Rollback: rollback})
|
||||||
|
out, err := m.handleRollbackCommand(nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, "Rolled back state.txt") {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
data, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if string(data) != "before" {
|
||||||
|
t.Fatalf("content = %q", data)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSkillsInstallCommandActivatesSkill(t *testing.T) {
|
||||||
|
home := t.TempDir()
|
||||||
|
t.Setenv("HOME", home)
|
||||||
|
sourceDir := t.TempDir()
|
||||||
|
source := filepath.Join(sourceDir, "review skill.md")
|
||||||
|
content := "---\nname: review\ndescription: Review changes\n---\n\nCheck every diff.\n"
|
||||||
|
if err := os.WriteFile(source, []byte(content), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
assistant := agent.New(agent.Options{SystemPrompt: "base"})
|
||||||
|
m := newModel(context.Background(), assistant, Options{ModelName: "test-model"})
|
||||||
|
out, err := m.handleLocalCommand("/skills install " + source)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out, `Installed skill "review"`) {
|
||||||
|
t.Fatalf("output = %q", out)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(filepath.Join(home, ".agentu", "skills", "review", "SKILL.md")); err != nil {
|
||||||
|
t.Fatalf("installed skill missing: %v", err)
|
||||||
|
}
|
||||||
|
messages := assistant.Messages()
|
||||||
|
if len(messages) == 0 || !strings.Contains(messages[0].Content, "Check every diff.") {
|
||||||
|
t.Fatalf("skill not activated in system prompt: %+v", messages)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,128 @@
|
|||||||
|
package tui
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
tea "github.com/charmbracelet/bubbletea"
|
||||||
|
|
||||||
|
"agentu/internal/agent"
|
||||||
|
)
|
||||||
|
|
||||||
|
// toolApprovalMsg is sent to the TUI when a mutating tool call needs approval.
|
||||||
|
type toolApprovalMsg struct {
|
||||||
|
request agent.ToolAuthorizationRequest
|
||||||
|
respond chan agent.ToolAuthorizationDecision
|
||||||
|
}
|
||||||
|
|
||||||
|
// ApprovalBroker bridges the agent's ToolAuthorizer interface with the TUI's
|
||||||
|
// event loop, allowing interactive approval of mutating tool calls.
|
||||||
|
type ApprovalBroker struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
events chan<- tea.Msg
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewApprovalBroker creates a new approval broker.
|
||||||
|
func NewApprovalBroker() *ApprovalBroker {
|
||||||
|
return &ApprovalBroker{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetEvents sets the current event channel for sending approval requests.
|
||||||
|
func (b *ApprovalBroker) SetEvents(events chan<- tea.Msg) {
|
||||||
|
b.mu.Lock()
|
||||||
|
defer b.mu.Unlock()
|
||||||
|
b.events = events
|
||||||
|
}
|
||||||
|
|
||||||
|
// ClearEvents removes the current event channel.
|
||||||
|
func (b *ApprovalBroker) ClearEvents() {
|
||||||
|
b.mu.Lock()
|
||||||
|
defer b.mu.Unlock()
|
||||||
|
b.events = nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// AuthorizeToolCall implements agent.ToolAuthorizer by sending an approval
|
||||||
|
// request to the TUI and waiting for the user's decision.
|
||||||
|
func (b *ApprovalBroker) AuthorizeToolCall(ctx context.Context, req agent.ToolAuthorizationRequest) (agent.ToolAuthorizationDecision, error) {
|
||||||
|
b.mu.Lock()
|
||||||
|
events := b.events
|
||||||
|
b.mu.Unlock()
|
||||||
|
|
||||||
|
if events == nil {
|
||||||
|
return "", fmt.Errorf("approval UI is not available")
|
||||||
|
}
|
||||||
|
|
||||||
|
respond := make(chan agent.ToolAuthorizationDecision, 1)
|
||||||
|
msg := toolApprovalMsg{
|
||||||
|
request: req,
|
||||||
|
respond: respond,
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case events <- msg:
|
||||||
|
case <-ctx.Done():
|
||||||
|
return "", ctx.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case decision := <-respond:
|
||||||
|
return decision, nil
|
||||||
|
case <-ctx.Done():
|
||||||
|
return "", ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *model) updateApproval(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||||
|
if m.approval == nil {
|
||||||
|
return *m, nil
|
||||||
|
}
|
||||||
|
switch msg.String() {
|
||||||
|
case "y", "Y":
|
||||||
|
m.respondToApproval(agent.ToolAuthorizationAllow)
|
||||||
|
return *m, nil
|
||||||
|
case "n", "N", "esc":
|
||||||
|
m.respondToApproval(agent.ToolAuthorizationDeny)
|
||||||
|
return *m, nil
|
||||||
|
case "ctrl+c":
|
||||||
|
m.respondToApproval(agent.ToolAuthorizationDeny)
|
||||||
|
if m.cancel != nil {
|
||||||
|
m.cancel()
|
||||||
|
}
|
||||||
|
return *m, tea.Quit
|
||||||
|
default:
|
||||||
|
return *m, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *model) respondToApproval(decision agent.ToolAuthorizationDecision) {
|
||||||
|
request := m.approval
|
||||||
|
m.approval = nil
|
||||||
|
if request != nil {
|
||||||
|
request.respond <- decision
|
||||||
|
}
|
||||||
|
m.status = "working..."
|
||||||
|
m.resize()
|
||||||
|
m.refreshViewport(true)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m model) approvalView() string {
|
||||||
|
if m.approval == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
req := m.approval.request
|
||||||
|
preview := strings.TrimSpace(req.Preview)
|
||||||
|
if preview == "" {
|
||||||
|
preview = strings.TrimSpace(req.Arguments)
|
||||||
|
}
|
||||||
|
lines := []string{
|
||||||
|
m.styles.ErrorLabel.Render("Approve " + req.ToolName + "?"),
|
||||||
|
m.styles.ToolValue.Render(preview),
|
||||||
|
}
|
||||||
|
if req.RollbackAvailable {
|
||||||
|
lines = append(lines, m.styles.Muted.Render("A rollback snapshot will be created."))
|
||||||
|
}
|
||||||
|
lines = append(lines, m.styles.Muted.Render("y approve · n/esc deny"))
|
||||||
|
return m.styles.InputBox.Width(max(20, m.width)).Render(strings.Join(lines, "\n"))
|
||||||
|
}
|
||||||
+110
-19
@@ -15,6 +15,17 @@ const (
|
|||||||
ThemeDark ThemeMode = "dark"
|
ThemeDark ThemeMode = "dark"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
iconReady = "✓"
|
||||||
|
iconWorking = "◆"
|
||||||
|
iconTool = "▣"
|
||||||
|
iconError = "✕"
|
||||||
|
iconCommand = "⌘"
|
||||||
|
iconSession = "◌"
|
||||||
|
iconContext = "◔"
|
||||||
|
iconModel = "◉"
|
||||||
|
)
|
||||||
|
|
||||||
func ParseThemeMode(value string) (ThemeMode, error) {
|
func ParseThemeMode(value string) (ThemeMode, error) {
|
||||||
switch ThemeMode(strings.ToLower(strings.TrimSpace(value))) {
|
switch ThemeMode(strings.ToLower(strings.TrimSpace(value))) {
|
||||||
case "", ThemeLight:
|
case "", ThemeLight:
|
||||||
@@ -32,6 +43,7 @@ type theme struct {
|
|||||||
Background string
|
Background string
|
||||||
Surface string
|
Surface string
|
||||||
SurfaceAlt string
|
SurfaceAlt string
|
||||||
|
InputBg string
|
||||||
Border string
|
Border string
|
||||||
Text string
|
Text string
|
||||||
Muted string
|
Muted string
|
||||||
@@ -48,11 +60,25 @@ type styles struct {
|
|||||||
App lipgloss.Style
|
App lipgloss.Style
|
||||||
Muted lipgloss.Style
|
Muted lipgloss.Style
|
||||||
Activity lipgloss.Style
|
Activity lipgloss.Style
|
||||||
ModelMeta lipgloss.Style
|
Header lipgloss.Style
|
||||||
Footer lipgloss.Style
|
Footer lipgloss.Style
|
||||||
Viewport lipgloss.Style
|
Viewport lipgloss.Style
|
||||||
InputBox lipgloss.Style
|
InputBox lipgloss.Style
|
||||||
CommandHint lipgloss.Style
|
CommandHint lipgloss.Style
|
||||||
|
InputModeReadonly lipgloss.Style
|
||||||
|
InputModeYolo lipgloss.Style
|
||||||
|
InputModeCommand lipgloss.Style
|
||||||
|
StatusReady lipgloss.Style
|
||||||
|
StatusWarn lipgloss.Style
|
||||||
|
StatusError lipgloss.Style
|
||||||
|
StatusInfo lipgloss.Style
|
||||||
|
Pill lipgloss.Style
|
||||||
|
PillAccent lipgloss.Style
|
||||||
|
Palette lipgloss.Style
|
||||||
|
PaletteMatch lipgloss.Style
|
||||||
|
ToolName lipgloss.Style
|
||||||
|
ToolKey lipgloss.Style
|
||||||
|
ToolValue lipgloss.Style
|
||||||
UserLabel lipgloss.Style
|
UserLabel lipgloss.Style
|
||||||
ToolLabel lipgloss.Style
|
ToolLabel lipgloss.Style
|
||||||
ErrorLabel lipgloss.Style
|
ErrorLabel lipgloss.Style
|
||||||
@@ -75,8 +101,8 @@ func lightTheme() theme {
|
|||||||
return theme{
|
return theme{
|
||||||
Mode: ThemeLight,
|
Mode: ThemeLight,
|
||||||
Background: "#F8FAFC",
|
Background: "#F8FAFC",
|
||||||
Surface: "#FFFFFF",
|
|
||||||
SurfaceAlt: "#EEF2F7",
|
SurfaceAlt: "#EEF2F7",
|
||||||
|
InputBg: "#E4E4E4",
|
||||||
Border: "#CBD5E1",
|
Border: "#CBD5E1",
|
||||||
Text: "#111827",
|
Text: "#111827",
|
||||||
Muted: "#64748B",
|
Muted: "#64748B",
|
||||||
@@ -93,8 +119,8 @@ func darkTheme() theme {
|
|||||||
return theme{
|
return theme{
|
||||||
Mode: ThemeDark,
|
Mode: ThemeDark,
|
||||||
Background: "#111827",
|
Background: "#111827",
|
||||||
Surface: "#1F2937",
|
|
||||||
SurfaceAlt: "#0F172A",
|
SurfaceAlt: "#0F172A",
|
||||||
|
InputBg: "#383838",
|
||||||
Border: "#475569",
|
Border: "#475569",
|
||||||
Text: "#E5E7EB",
|
Text: "#E5E7EB",
|
||||||
Muted: "#94A3B8",
|
Muted: "#94A3B8",
|
||||||
@@ -107,24 +133,28 @@ func darkTheme() theme {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func appStyle(t theme) lipgloss.Style {
|
||||||
|
style := lipgloss.NewStyle().Foreground(lipgloss.Color(t.Text))
|
||||||
|
if t.Mode == ThemeDark {
|
||||||
|
style = style.Background(lipgloss.Color(t.Background))
|
||||||
|
}
|
||||||
|
return style
|
||||||
|
}
|
||||||
|
|
||||||
func (t theme) styles() styles {
|
func (t theme) styles() styles {
|
||||||
return styles{
|
return styles{
|
||||||
App: lipgloss.NewStyle().
|
App: appStyle(t),
|
||||||
Foreground(lipgloss.Color(t.Text)),
|
|
||||||
|
|
||||||
Muted: lipgloss.NewStyle().
|
Muted: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Muted)),
|
Foreground(lipgloss.Color(t.Muted)),
|
||||||
|
|
||||||
Activity: lipgloss.NewStyle().
|
Activity: lipgloss.NewStyle().
|
||||||
Bold(true).
|
Foreground(lipgloss.Color(t.Muted)).
|
||||||
Foreground(lipgloss.Color(t.Background)).
|
|
||||||
Background(lipgloss.Color(t.Title)).
|
|
||||||
Padding(0, 1),
|
Padding(0, 1),
|
||||||
|
|
||||||
ModelMeta: lipgloss.NewStyle().
|
Header: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Text)).
|
Bold(true).
|
||||||
Background(lipgloss.Color(t.SurfaceAlt)).
|
Foreground(lipgloss.Color(t.Title)),
|
||||||
Padding(0, 1),
|
|
||||||
|
|
||||||
Footer: lipgloss.NewStyle().
|
Footer: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Muted)).
|
Foreground(lipgloss.Color(t.Muted)).
|
||||||
@@ -134,15 +164,74 @@ func (t theme) styles() styles {
|
|||||||
Foreground(lipgloss.Color(t.Text)),
|
Foreground(lipgloss.Color(t.Text)),
|
||||||
|
|
||||||
InputBox: lipgloss.NewStyle().
|
InputBox: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Text)).
|
||||||
|
Background(lipgloss.Color(t.InputBg)).
|
||||||
|
Padding(1, 2, 1, 1),
|
||||||
|
|
||||||
|
CommandHint: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Muted)).
|
||||||
|
Padding(0, 1),
|
||||||
|
|
||||||
|
InputModeReadonly: lipgloss.NewStyle().
|
||||||
|
Bold(true).
|
||||||
|
Foreground(lipgloss.Color(t.Muted)).
|
||||||
|
Background(lipgloss.Color(t.InputBg)).
|
||||||
|
Padding(0, 1),
|
||||||
|
|
||||||
|
InputModeYolo: lipgloss.NewStyle().
|
||||||
|
Bold(true).
|
||||||
|
Foreground(lipgloss.Color(t.Error)).
|
||||||
|
Background(lipgloss.Color(t.InputBg)).
|
||||||
|
Padding(0, 1),
|
||||||
|
|
||||||
|
InputModeCommand: lipgloss.NewStyle().
|
||||||
|
Bold(true).
|
||||||
|
Foreground(lipgloss.Color(t.User)).
|
||||||
|
Background(lipgloss.Color(t.InputBg)).
|
||||||
|
Padding(0, 1),
|
||||||
|
|
||||||
|
StatusReady: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Assistant)),
|
||||||
|
|
||||||
|
StatusWarn: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Tool)),
|
||||||
|
|
||||||
|
StatusError: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Error)),
|
||||||
|
|
||||||
|
StatusInfo: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.User)),
|
||||||
|
|
||||||
|
Pill: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Text)).
|
||||||
|
Background(lipgloss.Color(t.SurfaceAlt)).
|
||||||
|
Padding(0, 1),
|
||||||
|
|
||||||
|
PillAccent: lipgloss.NewStyle().
|
||||||
|
Bold(true).
|
||||||
|
Foreground(lipgloss.Color(t.Title)).
|
||||||
|
Background(lipgloss.Color(t.SurfaceAlt)).
|
||||||
|
Padding(0, 1),
|
||||||
|
|
||||||
|
Palette: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Text)).
|
Foreground(lipgloss.Color(t.Text)).
|
||||||
Border(lipgloss.ThickBorder()).
|
Border(lipgloss.ThickBorder()).
|
||||||
BorderForeground(lipgloss.Color(t.Border)).
|
BorderForeground(lipgloss.Color(t.Border)).
|
||||||
Padding(0, 2),
|
Padding(0, 2),
|
||||||
|
|
||||||
CommandHint: lipgloss.NewStyle().
|
PaletteMatch: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Muted)).
|
Bold(true).
|
||||||
Background(lipgloss.Color(t.SurfaceAlt)).
|
Foreground(lipgloss.Color(t.User)),
|
||||||
Padding(0, 1),
|
|
||||||
|
ToolName: lipgloss.NewStyle().
|
||||||
|
Bold(true).
|
||||||
|
Foreground(lipgloss.Color(t.Tool)),
|
||||||
|
|
||||||
|
ToolKey: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Muted)),
|
||||||
|
|
||||||
|
ToolValue: lipgloss.NewStyle().
|
||||||
|
Foreground(lipgloss.Color(t.Text)),
|
||||||
|
|
||||||
UserLabel: labelStyle(t.User),
|
UserLabel: labelStyle(t.User),
|
||||||
ToolLabel: labelStyle(t.Tool),
|
ToolLabel: labelStyle(t.Tool),
|
||||||
@@ -185,8 +274,10 @@ func bubbleMessageStyle(textColor, backgroundColor string) lipgloss.Style {
|
|||||||
func applyTextareaTheme(input *textarea.Model, t theme) {
|
func applyTextareaTheme(input *textarea.Model, t theme) {
|
||||||
focused := textarea.Style{
|
focused := textarea.Style{
|
||||||
Base: lipgloss.NewStyle().
|
Base: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Text)),
|
Foreground(lipgloss.Color(t.Text)).
|
||||||
CursorLine: lipgloss.NewStyle(),
|
Background(lipgloss.Color(t.InputBg)),
|
||||||
|
CursorLine: lipgloss.NewStyle().
|
||||||
|
Background(lipgloss.Color(t.InputBg)),
|
||||||
Placeholder: lipgloss.NewStyle().
|
Placeholder: lipgloss.NewStyle().
|
||||||
Foreground(lipgloss.Color(t.Muted)),
|
Foreground(lipgloss.Color(t.Muted)),
|
||||||
Prompt: lipgloss.NewStyle().
|
Prompt: lipgloss.NewStyle().
|
||||||
@@ -197,7 +288,7 @@ func applyTextareaTheme(input *textarea.Model, t theme) {
|
|||||||
Foreground(lipgloss.Color(t.Surface)),
|
Foreground(lipgloss.Color(t.Surface)),
|
||||||
}
|
}
|
||||||
blurred := focused
|
blurred := focused
|
||||||
blurred.CursorLine = lipgloss.NewStyle()
|
blurred.CursorLine = lipgloss.NewStyle().Background(lipgloss.Color(t.InputBg))
|
||||||
|
|
||||||
input.FocusedStyle = focused
|
input.FocusedStyle = focused
|
||||||
input.BlurredStyle = blurred
|
input.BlurredStyle = blurred
|
||||||
|
|||||||
@@ -1,222 +1,15 @@
|
|||||||
package main
|
package main
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"bufio"
|
|
||||||
"context"
|
|
||||||
"errors"
|
|
||||||
"flag"
|
|
||||||
"fmt"
|
"fmt"
|
||||||
"net/http"
|
|
||||||
"os"
|
"os"
|
||||||
"os/signal"
|
|
||||||
"path/filepath"
|
|
||||||
"strings"
|
|
||||||
|
|
||||||
"agentu/internal/agent"
|
"agentu/internal/app"
|
||||||
"agentu/pkg/config"
|
|
||||||
"agentu/pkg/llm"
|
|
||||||
"agentu/internal/session"
|
|
||||||
"agentu/internal/tools"
|
|
||||||
"agentu/internal/tui"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
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() {
|
func main() {
|
||||||
if err := run(); err != nil {
|
if err := app.Run(); err != nil {
|
||||||
fmt.Fprintln(os.Stderr, "agentu:", err)
|
fmt.Fprintln(os.Stderr, "agentu:", err)
|
||||||
os.Exit(1)
|
os.Exit(1)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func run() error {
|
|
||||||
configPath := flag.String("config", config.DefaultPath, "path to YAML config")
|
|
||||||
yolo := flag.Bool("yolo", false, "also enable automatic file writes and shell execution")
|
|
||||||
plain := flag.Bool("plain", false, "use the simple line-based REPL instead of the TUI")
|
|
||||||
themeName := flag.String("theme", "light", "TUI theme: light or dark")
|
|
||||||
resumeID := flag.String("resume", "", "resume a specific session by ID")
|
|
||||||
flag.Parse()
|
|
||||||
|
|
||||||
themeMode, err := tui.ParseThemeMode(*themeName)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
cfg, err := config.Load(*configPath)
|
|
||||||
if err != nil {
|
|
||||||
if errors.Is(err, os.ErrNotExist) {
|
|
||||||
if bootstrapErr := bootstrapConfig(*configPath); bootstrapErr != nil {
|
|
||||||
return bootstrapErr
|
|
||||||
}
|
|
||||||
return fmt.Errorf("created default config at %s; edit it to set your API key and provider, then restart agentu", *configPath)
|
|
||||||
}
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
providerName := cfg.DefaultProviderName()
|
|
||||||
providerConfig := cfg.Providers[providerName]
|
|
||||||
provider := llm.NewOpenAICompatibleClient(providerConfig.BaseURL, providerConfig.APIKey, http.DefaultClient)
|
|
||||||
|
|
||||||
var registry *tools.Registry
|
|
||||||
if *yolo {
|
|
||||||
registry = tools.Builtins(cfg.Agent.WorkingDir)
|
|
||||||
} else {
|
|
||||||
registry = tools.ReadOnlyBuiltins(cfg.Agent.WorkingDir)
|
|
||||||
}
|
|
||||||
registry.Register(tools.NewBuiltinSearchTool(http.DefaultClient))
|
|
||||||
registry.Register(tools.NewBuiltinFetchTool(http.DefaultClient))
|
|
||||||
systemPrompt := strings.TrimSpace(cfg.Agent.SystemPrompt + "\n\n" + tools.AgentInstructions(cfg.Agent.WorkingDir, *yolo))
|
|
||||||
maxContextTokens := cfg.Agent.MaxContextTokens
|
|
||||||
if maxContextTokens == 0 {
|
|
||||||
maxContextTokens = providerConfig.ContextTokens
|
|
||||||
}
|
|
||||||
|
|
||||||
assistant := agent.New(agent.Options{
|
|
||||||
Provider: provider,
|
|
||||||
ProviderName: providerName,
|
|
||||||
Model: providerConfig.Model,
|
|
||||||
Thinking: providerConfig.Thinking,
|
|
||||||
ThinkingParam: providerConfig.ThinkingParam,
|
|
||||||
ThinkingEnabled: providerConfig.ThinkingConfigured,
|
|
||||||
SystemPrompt: systemPrompt,
|
|
||||||
ToolRegistry: registry,
|
|
||||||
ToolTimeout: cfg.Agent.ToolTimeout,
|
|
||||||
MaxContextTokens: maxContextTokens,
|
|
||||||
})
|
|
||||||
|
|
||||||
// Initialize session store
|
|
||||||
home, err := os.UserHomeDir()
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("resolve home dir: %w", err)
|
|
||||||
}
|
|
||||||
sessionsDir := filepath.Join(home, ".agentu", "sessions")
|
|
||||||
store, err := session.NewStore(sessionsDir)
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("init session store: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
modelManager, err := session.NewManager(cfg, assistant, http.DefaultClient, store)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
// Resume session
|
|
||||||
if *resumeID != "" {
|
|
||||||
if err := modelManager.Resume(*resumeID); err != nil {
|
|
||||||
return fmt.Errorf("resume session: %w", err)
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
if err := modelManager.ResumeLatest(); err != nil {
|
|
||||||
return fmt.Errorf("resume latest session: %w", err)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
|
|
||||||
defer stop()
|
|
||||||
|
|
||||||
if !*plain {
|
|
||||||
err := tui.Run(ctx, assistant, tui.Options{
|
|
||||||
ModelName: providerConfig.Model,
|
|
||||||
Yolo: *yolo,
|
|
||||||
ThemeMode: themeMode,
|
|
||||||
ModelManager: modelManager,
|
|
||||||
WorkingDir: cfg.Agent.WorkingDir,
|
|
||||||
})
|
|
||||||
printExitSession(modelManager)
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
return repl(ctx, assistant, modelManager)
|
|
||||||
}
|
|
||||||
|
|
||||||
func bootstrapConfig(configPath string) error {
|
|
||||||
resolved, err := config.ResolvePath(configPath)
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("resolve config path: %w", err)
|
|
||||||
}
|
|
||||||
dir := filepath.Dir(resolved)
|
|
||||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
|
||||||
return fmt.Errorf("create config directory: %w", err)
|
|
||||||
}
|
|
||||||
if err := os.WriteFile(resolved, []byte(defaultConfigContent), 0o644); err != nil {
|
|
||||||
return fmt.Errorf("write default config: %w", err)
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func repl(ctx context.Context, assistant *agent.Agent, modelManager *session.Manager) error {
|
|
||||||
scanner := bufio.NewScanner(os.Stdin)
|
|
||||||
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
|
||||||
|
|
||||||
fmt.Println("agentu")
|
|
||||||
fmt.Println("Type /exit to quit, /clear to reset context, /compact to compress context.")
|
|
||||||
for {
|
|
||||||
fmt.Print("> ")
|
|
||||||
if !scanner.Scan() {
|
|
||||||
if err := scanner.Err(); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
fmt.Println()
|
|
||||||
printExitSession(modelManager)
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
input := strings.TrimSpace(scanner.Text())
|
|
||||||
switch input {
|
|
||||||
case "":
|
|
||||||
continue
|
|
||||||
case "/exit", "/quit":
|
|
||||||
printExitSession(modelManager)
|
|
||||||
return nil
|
|
||||||
case "/clear":
|
|
||||||
assistant.Clear()
|
|
||||||
fmt.Println("context cleared")
|
|
||||||
continue
|
|
||||||
case "/compact":
|
|
||||||
if err := assistant.Compact(ctx); err != nil {
|
|
||||||
fmt.Fprintln(os.Stderr, "compact error:", err)
|
|
||||||
} else {
|
|
||||||
fmt.Println("context compacted")
|
|
||||||
if modelManager != nil {
|
|
||||||
_ = modelManager.Save()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := assistant.RunTurn(ctx, input, os.Stdout, os.Stderr); err != nil {
|
|
||||||
fmt.Fprintln(os.Stderr, "\nerror:", err)
|
|
||||||
}
|
|
||||||
_ = modelManager.Save()
|
|
||||||
fmt.Println()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func printExitSession(m *session.Manager) {
|
|
||||||
if m == nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
_ = m.Save()
|
|
||||||
id := m.CurrentSessionID()
|
|
||||||
name := m.CurrentSessionName()
|
|
||||||
if id != "" {
|
|
||||||
if name != "" && name != id {
|
|
||||||
fmt.Fprintf(os.Stderr, "session: %s (%s)\n", id, name)
|
|
||||||
} else {
|
|
||||||
fmt.Fprintf(os.Stderr, "session: %s\n", id)
|
|
||||||
}
|
|
||||||
fmt.Fprintf(os.Stderr, "resume: agentu --resume %s\n", id)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
+47
-1
@@ -14,6 +14,11 @@ import (
|
|||||||
|
|
||||||
const DefaultPath = "~/.agentu/config.yaml"
|
const DefaultPath = "~/.agentu/config.yaml"
|
||||||
|
|
||||||
|
const (
|
||||||
|
APITypeChat = "chat"
|
||||||
|
APITypeResponses = "responses"
|
||||||
|
)
|
||||||
|
|
||||||
type Config struct {
|
type Config struct {
|
||||||
Providers map[string]ProviderConfig `yaml:"providers"`
|
Providers map[string]ProviderConfig `yaml:"providers"`
|
||||||
Agent AgentConfig `yaml:"agent"`
|
Agent AgentConfig `yaml:"agent"`
|
||||||
@@ -21,6 +26,7 @@ type Config struct {
|
|||||||
|
|
||||||
type ProviderConfig struct {
|
type ProviderConfig struct {
|
||||||
Name string
|
Name string
|
||||||
|
APIType string
|
||||||
BaseURL string
|
BaseURL string
|
||||||
Model string
|
Model string
|
||||||
Models []string
|
Models []string
|
||||||
@@ -56,9 +62,11 @@ type rawConfig struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type rawProviderConfig struct {
|
type rawProviderConfig struct {
|
||||||
|
APIType string `yaml:"api_type"`
|
||||||
BaseURL string `yaml:"base_url"`
|
BaseURL string `yaml:"base_url"`
|
||||||
Models []rawModelConfig `yaml:"models"`
|
Models []rawModelConfig `yaml:"models"`
|
||||||
APIKey string `yaml:"api_key"`
|
APIKey string `yaml:"api_key"`
|
||||||
|
ThinkingParam string `yaml:"thinking_param"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type rawModelConfig struct {
|
type rawModelConfig struct {
|
||||||
@@ -153,6 +161,9 @@ func (c *Config) Validate() error {
|
|||||||
for name, provider := range c.Providers {
|
for name, provider := range c.Providers {
|
||||||
var missing []string
|
var missing []string
|
||||||
prefix := "providers." + name
|
prefix := "providers." + name
|
||||||
|
if !IsAPIType(provider.APIType) {
|
||||||
|
return fmt.Errorf("%s.api_type must be one of: %s", prefix, strings.Join(APITypes(), ", "))
|
||||||
|
}
|
||||||
if strings.TrimSpace(provider.BaseURL) == "" {
|
if strings.TrimSpace(provider.BaseURL) == "" {
|
||||||
missing = append(missing, prefix+".base_url")
|
missing = append(missing, prefix+".base_url")
|
||||||
}
|
}
|
||||||
@@ -220,21 +231,56 @@ func normalizeProvider(name string, raw rawProviderConfig) ProviderConfig {
|
|||||||
models, modelConfigs := normalizeModelConfigs(raw.Models)
|
models, modelConfigs := normalizeModelConfigs(raw.Models)
|
||||||
provider := ProviderConfig{
|
provider := ProviderConfig{
|
||||||
Name: name,
|
Name: name,
|
||||||
|
APIType: normalizeAPIType(raw.APIType),
|
||||||
BaseURL: raw.BaseURL,
|
BaseURL: raw.BaseURL,
|
||||||
Models: models,
|
Models: models,
|
||||||
ModelConfigs: modelConfigs,
|
ModelConfigs: modelConfigs,
|
||||||
APIKey: raw.APIKey,
|
APIKey: raw.APIKey,
|
||||||
|
ThinkingParam: strings.TrimSpace(raw.ThinkingParam),
|
||||||
}
|
}
|
||||||
if len(provider.Models) > 0 {
|
if len(provider.Models) > 0 {
|
||||||
provider.Model = provider.Models[0]
|
provider.Model = provider.Models[0]
|
||||||
}
|
}
|
||||||
applyModelDefaults(&provider)
|
applyModelDefaults(&provider)
|
||||||
if provider.ThinkingConfigured && provider.ThinkingParam == "" {
|
if provider.ThinkingConfigured && provider.ThinkingParam == "" {
|
||||||
provider.ThinkingParam = "thinking"
|
provider.ThinkingParam = DefaultThinkingParam(provider.APIType)
|
||||||
}
|
}
|
||||||
return provider
|
return provider
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func normalizeAPIType(value string) string {
|
||||||
|
value = strings.ToLower(strings.TrimSpace(value))
|
||||||
|
if value == "" {
|
||||||
|
return APITypeChat
|
||||||
|
}
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
|
||||||
|
func APITypes() []string {
|
||||||
|
return []string{APITypeChat, APITypeResponses}
|
||||||
|
}
|
||||||
|
|
||||||
|
func IsAPIType(value string) bool {
|
||||||
|
value = strings.ToLower(strings.TrimSpace(value))
|
||||||
|
for _, apiType := range APITypes() {
|
||||||
|
if value == apiType {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// DefaultThinkingParam returns the default request field used for the
|
||||||
|
// thinking level. Chat-completions providers historically use a top-level
|
||||||
|
// "thinking" field; the OpenAI Responses API uses "reasoning" with an
|
||||||
|
// {"effort": ...} value.
|
||||||
|
func DefaultThinkingParam(apiType string) string {
|
||||||
|
if apiType == APITypeResponses {
|
||||||
|
return "reasoning"
|
||||||
|
}
|
||||||
|
return "thinking"
|
||||||
|
}
|
||||||
|
|
||||||
func normalizeModelConfigs(rawModels []rawModelConfig) ([]string, map[string]ModelConfig) {
|
func normalizeModelConfigs(rawModels []rawModelConfig) ([]string, map[string]ModelConfig) {
|
||||||
models := make([]string, 0, len(rawModels))
|
models := make([]string, 0, len(rawModels))
|
||||||
configs := make(map[string]ModelConfig, len(rawModels))
|
configs := make(map[string]ModelConfig, len(rawModels))
|
||||||
|
|||||||
@@ -370,3 +370,135 @@ providers:
|
|||||||
t.Fatalf("err = %v", err)
|
t.Fatalf("err = %v", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestLoadDefaultsAPITypeToChat(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
t.Setenv("AGENTU_TEST_KEY", "sk-test")
|
||||||
|
path := filepath.Join(dir, "agentu.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(`
|
||||||
|
providers:
|
||||||
|
test:
|
||||||
|
base_url: https://example.com
|
||||||
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
|
models:
|
||||||
|
- id: test-model
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg, err := Load(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if got := cfg.Providers["test"].APIType; got != APITypeChat {
|
||||||
|
t.Fatalf("api type = %q, want %q", got, APITypeChat)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadParsesResponsesAPIType(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
t.Setenv("AGENTU_TEST_KEY", "sk-test")
|
||||||
|
path := filepath.Join(dir, "agentu.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(`
|
||||||
|
providers:
|
||||||
|
openai:
|
||||||
|
api_type: RESPONSES
|
||||||
|
base_url: https://api.openai.com
|
||||||
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
|
models:
|
||||||
|
- id: gpt-test
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg, err := Load(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if got := cfg.Providers["openai"].APIType; got != APITypeResponses {
|
||||||
|
t.Fatalf("api type = %q, want %q", got, APITypeResponses)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadRejectsInvalidAPIType(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
t.Setenv("AGENTU_TEST_KEY", "sk-test")
|
||||||
|
path := filepath.Join(dir, "agentu.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(`
|
||||||
|
providers:
|
||||||
|
test:
|
||||||
|
api_type: completions
|
||||||
|
base_url: https://example.com
|
||||||
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
|
models:
|
||||||
|
- id: test-model
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := Load(path)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "api_type") {
|
||||||
|
t.Fatalf("err = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadDefaultsThinkingParamByAPIType(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
t.Setenv("AGENTU_TEST_KEY", "sk-test")
|
||||||
|
path := filepath.Join(dir, "agentu.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(`
|
||||||
|
providers:
|
||||||
|
chat:
|
||||||
|
api_type: chat
|
||||||
|
base_url: https://chat.example.com
|
||||||
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
|
models:
|
||||||
|
- id: chat-model
|
||||||
|
thinking: none
|
||||||
|
responses:
|
||||||
|
api_type: responses
|
||||||
|
base_url: https://responses.example.com
|
||||||
|
api_key: ${AGENTU_TEST_KEY}
|
||||||
|
models:
|
||||||
|
- id: responses-model
|
||||||
|
thinking: high
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg, err := Load(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if got := cfg.Providers["chat"].ThinkingParam; got != "thinking" {
|
||||||
|
t.Fatalf("chat thinking param = %q, want %q", got, "thinking")
|
||||||
|
}
|
||||||
|
if got := cfg.Providers["responses"].ThinkingParam; got != "reasoning" {
|
||||||
|
t.Fatalf("responses thinking param = %q, want %q", got, "reasoning")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadCustomThinkingParam(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "agentu.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(`
|
||||||
|
providers:
|
||||||
|
custom:
|
||||||
|
api_type: chat
|
||||||
|
base_url: https://example.com
|
||||||
|
api_key: sk-test
|
||||||
|
thinking_param: reasoning_effort
|
||||||
|
models:
|
||||||
|
- id: custom-model
|
||||||
|
thinking: high
|
||||||
|
`), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
cfg, err := Load(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if got := cfg.Providers["custom"].ThinkingParam; got != "reasoning_effort" {
|
||||||
|
t.Fatalf("thinking param = %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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())
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user