Compare commits

..
Author SHA1 Message Date
ParthSareen 4784290c9b agent: use platform shell tool names 2026-06-26 17:26:38 -07:00
ParthSareen b839fa49f0 agent: stabilize windows shell tool 2026-06-26 16:02:00 -07:00
ParthSareen 93bb382f9a agent: improve shell and tui resilience 2026-06-26 15:45:34 -07:00
ParthSareen 8ed7b66d31 ci: restore golangci-lint PR filtering 2026-06-26 13:25:25 -07:00
ParthSareen 276afbf3a7 chore(agent): tidy module dependencies 2026-06-26 12:26:59 -07:00
ParthSareen 2ae9c05fbd ci: run golangci-lint without PR diff filtering 2026-06-26 12:09:04 -07:00
ParthSareen 6aa08b2c94 cmd/tui/chat: guard chatRunDoneMsg send against context cancellation 2026-06-25 17:04:25 -07:00
ParthSareen b2bc6c822c chore(agent): address modelfile review feedback 2026-06-25 14:30:35 -07:00
ParthSareen 7703a00ade refactor(agent): introduce RunPolicy and drop --hidethinking
Replaces the scattered auto-approve plumbing (ApprovalManager.AutoApprove,
chatPermissionMode, opts.AutoApproveTools, and the approvalHandlerAutoApproves
type-switch) with a single RunPolicy value object carrying ToolMode
(Review/FullAccess/Disabled), an ApprovalPolicy strategy, and MaxToolRounds.
Mode is derived once per surface via resolveAgentRunPolicy, and the TUI toggles
it in place via RunPolicyState.

These two changes land together because they share struct-field plumbing in
agent_tui.go: the --hidethinking removal drops HideThinking from the same
gofmt-realigned AgentTUIOptions/runOptions literals that gain the Policy field,
so splitting them would produce intermediate states that do not build.

Also drops the now-unused PermissionNotice() and RunPolicyState.SetPolicy,
and simplifies the headless tool-resolution block.
2026-06-25 13:20:14 -07:00
ParthSareen 3052b16971 feat(chat): save raw request with /raw > filename
Adds /raw > <file> to write the JSON request preview to a file in the
working directory instead of opening the popup. Routes /raw through a
new handleRawCommand that distinguishes the popup form from the
redirect form, and extracts rawRequestPreviewJSON so the popup and the
save path share one preview build.
2026-06-25 13:19:24 -07:00
ParthSareen 9eed9466fb chore(tui): clean up chat lint 2026-06-24 16:25:21 -07:00
ParthSareen a0d99b4a02 refactor(agent): move chat store out of app 2026-06-24 16:11:45 -07:00
ParthSareen dd82c333ec build(ci): gate app store to desktop platforms 2026-06-24 16:00:44 -07:00
ParthSareen f8750975c1 app/store: repair agent schema migration 2026-06-24 15:43:49 -07:00
ParthSareen 64d3680ab6 cmd/tui/chat: wire cloud auth tick/poll dispatch, auto-open browser, use PollCloudAuth callback 2026-06-24 15:42:00 -07:00
ParthSareen 4c7d01cd4f cmd/tui/chat: skip checking spinner, use pre-computed badge to go straight to sign-in or upgrade 2026-06-24 15:38:47 -07:00
ParthSareen eacaa79706 cmd/tui/chat: show availability badge beside model name instead of plan in meta 2026-06-24 15:34:04 -07:00
ParthSareen 1e004ef209 cmd/tui/chat: add cloud model sign-in and plan-upgrade prompts to model picker 2026-06-24 15:23:57 -07:00
ParthSareen 4a34477557 cmd/tui/chat: add raw request debug view 2026-06-24 14:59:45 -07:00
ParthSareen eb76282573 app/store: add source column to hide agent chats from desktop app 2026-06-24 14:58:35 -07:00
ParthSareen c47ad4f228 cmd: extract multimodal compat check into applyMultiModalCompat helper 2026-06-24 14:51:55 -07:00
ParthSareen e785d4e8b9 agent: require approval for web_search and web_fetch tools 2026-06-24 14:48:59 -07:00
ParthSareen 2e9b631176 cmd/tui/chat: consolidate model and resume pickers into generic chatPicker 2026-06-24 14:48:50 -07:00
ParthSareen 087ffc697f cmd/tui: satisfy chat lint 2026-06-23 18:22:49 -07:00
ParthSareen 11f5c77b4b cmd: present cloud-tagged models with arch details in agent picker 2026-06-23 18:14:48 -07:00
ParthSareen e76f830221 cmd: restore local model arch details in picker description 2026-06-23 18:06:18 -07:00
ParthSareen 3b54956946 cmd/tui: split chat into subpackage and simplify status display 2026-06-23 18:02:33 -07:00
ParthSareen 8f882b054a agent: consolidate compaction helpers and drop trace sink 2026-06-23 14:24:36 -07:00
ParthSareen e098a6afaa refactor(agent): consolidate chat persistence 2026-06-23 13:59:32 -07:00
ParthSareen 183f03d997 agent: address review feedback 2026-06-23 13:31:18 -07:00
ParthSareen 5e9c48cc40 agent: ignore image bytes in prompt estimates 2026-06-23 10:51:58 -07:00
ParthSareen 6b389711f0 chore(agent): satisfy lint checks 2026-06-22 17:01:47 -07:00
ParthSareen f5a50a7c3f agent: harden compaction and tui state 2026-06-22 16:49:41 -07:00
ParthSareen 5104df202d chore(agent): polish chat runtime and rendering 2026-06-22 16:48:38 -07:00
ParthSareen c1ba011f64 agent: add event trace sink 2026-06-22 16:45:01 -07:00
ParthSareen 6d2bbd68e3 feat(agent): refine chat tui flow 2026-06-20 17:23:56 -07:00
ParthSareen 42467c2431 feat(agent): polish chat tui loading and approvals 2026-06-18 21:59:45 -07:00
ParthSareen 9c04e04e6b agent: address review feedback 2026-06-18 16:24:43 -07:00
ParthSareen 98b022f4f5 agent: collapse large pasted TUI input 2026-06-18 14:29:19 -07:00
ParthSareen d60a21fcaa chore(agent): compact small contexts aggressively 2026-06-18 02:16:20 -07:00
ParthSareen 97c93a93b4 chore(agent): harden compaction and runtime cleanup 2026-06-18 01:52:36 -07:00
ParthSareen a058a75099 agent: harden resume and context preflight 2026-06-18 01:23:15 -07:00
ParthSareen c359d22a08 chore: note empty compaction summaries 2026-06-18 00:26:24 -07:00
ParthSareen 0034967cfd feat: refine agent cli flow 2026-06-17 21:30:38 -07:00
ParthSareen 66757a973d chore(agent): handle windows file urls 2026-06-17 16:39:40 -07:00
ParthSareen a718ed68a1 chore(agent): make windows paths portable 2026-06-17 16:24:44 -07:00
ParthSareen 204f719b77 feat(agent): add root onboarding flow 2026-06-17 15:42:21 -07:00
ParthSareen 9618a45286 feat(agent): improve chat tui workflow 2026-06-17 14:02:21 -07:00
ParthSareen ed6a3e6c4e chore: untrack generated docs 2026-06-15 16:53:05 -07:00
ParthSareen 280b85f8a4 agent: harden tool loop and polish tui 2026-06-15 15:33:47 -07:00
ParthSareen 821a45ca6e agent: update compaction and context tracking 2026-06-12 16:28:58 -07:00
ParthSareen 676facf043 agent: update chat tui controls 2026-06-12 16:28:58 -07:00
ParthSareen 50184444da agent: add chat tui controls and legacy aliases 2026-06-12 16:28:58 -07:00
ParthSareen 632bd88937 agent: improve chat input controls 2026-06-12 16:28:58 -07:00
ParthSareen 2ba0f9b96f agent: add output copy commands 2026-06-12 16:28:58 -07:00
ParthSareen 5625904153 agent: window slash command suggestions 2026-06-12 16:28:58 -07:00
ParthSareen 96f82b5ff2 agent: add model picker slash command 2026-06-12 16:28:58 -07:00
ParthSareen 9da4d0c6f8 agent: restore web tool auto approval and TUI scrolling 2026-06-12 16:28:58 -07:00
ParthSareen edd30d7194 docs: add agent improvements report 2026-06-12 16:28:58 -07:00
ParthSareen f417732279 agent: bound markdown and compaction caches 2026-06-12 16:28:58 -07:00
ParthSareen 72d18c99af agent: validate bash working directory 2026-06-12 16:28:58 -07:00
ParthSareen 02befa963c agent: require approval for web tools 2026-06-12 16:28:58 -07:00
ParthSareen a743721ec8 agent: harden bash approval classification 2026-06-12 16:28:58 -07:00
ParthSareen 2b54b72207 tui: split chat implementation 2026-06-12 16:28:58 -07:00
ParthSareen 30f206e235 agent: merge model-aware chat store 2026-06-12 16:28:58 -07:00
ParthSareen 13006826f8 agent: batch streaming persistence 2026-06-12 16:28:58 -07:00
ParthSareen 76eb5a9dc9 tui: track live working directory 2026-06-12 16:28:58 -07:00
ParthSareen cb8af634b3 tui: surface approvals over modals 2026-06-12 16:28:58 -07:00
ParthSareen 66cdf46e7d cmd: remove orphaned repl 2026-06-12 16:28:58 -07:00
ParthSareen a7c8484bea agent: persist chat images 2026-06-12 16:28:58 -07:00
ParthSareen 3be6794929 agent: unify approval dispatch 2026-06-12 16:28:58 -07:00
ParthSareen 444af2a712 feat: promote agent loop to canonical run path 2026-06-12 16:28:58 -07:00
130 changed files with 28162 additions and 7979 deletions

No files matched your search

@@ -1,98 +0,0 @@
name: test-darwin-xcode-pin
on:
workflow_dispatch:
push:
branches:
- test/darwin-xcode-pin
pull_request:
paths:
- '.github/workflows/test-darwin-xcode-pin.yaml'
- 'scripts/build_darwin.sh'
- 'MLX_VERSION'
- 'MLX_C_VERSION'
- 'cmake/**'
- 'x/mlxrunner/**'
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
env:
CGO_CFLAGS: '-O3'
CGO_CXXFLAGS: '-O3'
PINNED_DEVELOPER_DIR: /Applications/Xcode_26.4.1.app/Contents/Developer
jobs:
darwin-build:
runs-on: macos-26-xlarge
env:
CGO_CFLAGS: '-mmacosx-version-min=14.0 -O3'
CGO_CXXFLAGS: '-mmacosx-version-min=14.0 -O3'
CGO_LDFLAGS: '-mmacosx-version-min=14.0 -O3'
steps:
- uses: actions/checkout@v4
- name: Set build environment
shell: bash
run: |
set -euo pipefail
VERSION="0.0.0-xcode-pin-${GITHUB_SHA::7}"
{
echo "VERSION=${VERSION}"
echo "GOFLAGS='-ldflags=-w -s \"-X=github.com/ollama/ollama/version.Version=${VERSION}\" \"-X=github.com/ollama/ollama/server.mode=release\"'"
} >>"${GITHUB_ENV}"
- name: Select Xcode 26.4.1
shell: bash
run: |
set -euo pipefail
if [ ! -d "${PINNED_DEVELOPER_DIR}" ]; then
echo "Missing ${PINNED_DEVELOPER_DIR}"
ls -1 /Applications | grep '^Xcode' || true
exit 1
fi
sudo xcode-select -s "${PINNED_DEVELOPER_DIR}"
echo "DEVELOPER_DIR=${PINNED_DEVELOPER_DIR}" >>"${GITHUB_ENV}"
sw_vers
xcodebuild -version
xcrun --sdk macosx --show-sdk-version
xcrun --find metal
- uses: actions/setup-go@v5
with:
go-version-file: go.mod
cache-dependency-path: |
go.sum
LLAMA_CPP_VERSION
MLX_VERSION
MLX_C_VERSION
- name: Build unsigned Darwin runtime
run: ./scripts/build_darwin.sh build package
- name: Verify MLX payload
shell: bash
run: |
set -euo pipefail
test -f dist/darwin/lib/ollama/mlx_metal_v3/libmlxc.dylib
test -f dist/darwin/lib/ollama/mlx_metal_v3/mlx.metallib
test -f dist/darwin/lib/ollama/mlx_metal_v4/libmlxc.dylib
test -f dist/darwin/lib/ollama/mlx_metal_v4/mlx.metallib
find dist/darwin/lib/ollama -maxdepth 3 -type f \( -name 'libmlx*.dylib' -o -name '*.metallib' \) -print
lipo -archs dist/darwin/lib/ollama/mlx_metal_v3/libmlxc.dylib
lipo -archs dist/darwin/lib/ollama/mlx_metal_v4/libmlxc.dylib
- name: Log build results
run: ls -l dist/
- uses: actions/upload-artifact@v4
with:
name: ollama-darwin-xcode-pin
path: dist/ollama-darwin.tgz
compression-level: 0
+1 -1
View File
@@ -1 +1 @@
b9672
b9509
+905
View File
@@ -0,0 +1,905 @@
package agent
import (
"bytes"
"context"
"fmt"
"os"
"path/filepath"
"regexp"
"slices"
"strings"
"sync"
"unicode"
)
type ApprovalDecision string
const (
ApprovalAllowOnce ApprovalDecision = "allow_once"
ApprovalAllowSession ApprovalDecision = "allow_session"
ApprovalDeny ApprovalDecision = "deny"
)
type ApprovalRisk string
const (
ApprovalRiskLow ApprovalRisk = "low"
ApprovalRiskMedium ApprovalRisk = "medium"
ApprovalRiskHigh ApprovalRisk = "high"
)
type ApprovalRequest struct {
ToolCallID string
ToolName string
Args map[string]any
WorkingDir string
ToolApprovalRequired bool
Summary string
Risk ApprovalRisk
Reasons []string
}
type ApprovalResult struct {
Decision ApprovalDecision
Reason string
}
type ApprovalHandler interface {
RequiresApproval(context.Context, Tool, ApprovalRequest) bool
Approve(context.Context, ApprovalRequest) (ApprovalResult, error)
}
type ApprovalPrompter interface {
PromptApproval(context.Context, ApprovalRequest) (ApprovalResult, error)
}
type ApprovalPolicy interface {
EvaluateApproval(context.Context, ApprovalRequest) ApprovalEvaluation
}
type ApprovalEvaluation struct {
Decision ApprovalDecision
RequirePrompt bool
Risk ApprovalRisk
Summary string
Reasons []string
SessionKey string
}
type AutoAllowApproval struct{}
func (AutoAllowApproval) RequiresApproval(context.Context, Tool, ApprovalRequest) bool {
return false
}
func (AutoAllowApproval) Approve(context.Context, ApprovalRequest) (ApprovalResult, error) {
return ApprovalResult{Decision: ApprovalAllowOnce}, nil
}
type ApprovalManagerOptions struct {
Policy ApprovalPolicy
Prompter ApprovalPrompter
}
type ApprovalManager struct {
policy ApprovalPolicy
prompter ApprovalPrompter
mu *sync.Mutex
sessionAllowed map[string]struct{}
}
func NewApprovalManager(opts ApprovalManagerOptions) *ApprovalManager {
policy := opts.Policy
if policy == nil {
policy = DefaultApprovalPolicy{}
}
return &ApprovalManager{
policy: policy,
prompter: opts.Prompter,
mu: &sync.Mutex{},
sessionAllowed: make(map[string]struct{}),
}
}
func (m *ApprovalManager) WithPrompter(prompter ApprovalPrompter) *ApprovalManager {
if m == nil {
return NewApprovalManager(ApprovalManagerOptions{Prompter: prompter})
}
return &ApprovalManager{
policy: m.policy,
prompter: prompter,
mu: m.mu,
sessionAllowed: m.sessionAllowed,
}
}
func (m *ApprovalManager) RequiresApproval(ctx context.Context, tool Tool, req ApprovalRequest) bool {
if m == nil {
return false
}
req.ToolApprovalRequired = req.ToolApprovalRequired || ToolRequiresApproval(tool, req.Args)
evaluation := applyToolApprovalRequirement(req, m.evaluate(ctx, req))
if evaluation.Decision == ApprovalDeny {
return true
}
if evaluation.RequirePrompt {
return !m.sessionAllowedFor(evaluation.SessionKey)
}
return false
}
func (m *ApprovalManager) Approve(ctx context.Context, req ApprovalRequest) (ApprovalResult, error) {
if m == nil {
return ApprovalResult{Decision: ApprovalAllowOnce}, nil
}
evaluation := applyToolApprovalRequirement(req, m.evaluate(ctx, req))
req = approvalRequestWithEvaluation(req, evaluation)
if evaluation.Decision == ApprovalDeny {
reason := strings.Join(evaluation.Reasons, "; ")
if reason == "" {
reason = "Tool execution denied."
}
return ApprovalResult{Decision: ApprovalDeny, Reason: reason}, nil
}
if !evaluation.RequirePrompt || m.sessionAllowedFor(evaluation.SessionKey) {
return ApprovalResult{Decision: ApprovalAllowOnce}, nil
}
if m.prompter == nil {
return ApprovalResult{
Decision: ApprovalDeny,
Reason: "Tool execution requires approval, but no approval prompter is available. Re-run with --auto-approve-tools or --yolo to allow tool execution.",
}, nil
}
result, err := m.prompter.PromptApproval(ctx, req)
if err != nil {
return ApprovalResult{}, err
}
if result.Decision == "" {
result.Decision = ApprovalDeny
}
if result.Decision == ApprovalAllowSession {
m.allowSession(evaluation.SessionKey)
}
return result, nil
}
func (m *ApprovalManager) evaluate(ctx context.Context, req ApprovalRequest) ApprovalEvaluation {
if m == nil || m.policy == nil {
return DefaultApprovalPolicy{}.EvaluateApproval(ctx, req)
}
evaluation := m.policy.EvaluateApproval(ctx, req)
if evaluation.Risk == "" {
evaluation.Risk = ApprovalRiskLow
}
if evaluation.SessionKey == "" {
evaluation.SessionKey = approvalSessionKey(req)
}
return evaluation
}
func applyToolApprovalRequirement(req ApprovalRequest, evaluation ApprovalEvaluation) ApprovalEvaluation {
if !req.ToolApprovalRequired || evaluation.Decision == ApprovalDeny {
return evaluation
}
evaluation.RequirePrompt = true
if evaluation.Summary == "" {
evaluation.Summary = fmt.Sprintf("%s wants to run", ToolDisplayName(req.ToolName))
}
if len(evaluation.Reasons) == 0 {
evaluation.Reasons = []string{"tool requires approval"}
}
return evaluation
}
func (m *ApprovalManager) sessionAllowedFor(key string) bool {
if m == nil || key == "" {
return false
}
m.mu.Lock()
defer m.mu.Unlock()
_, ok := m.sessionAllowed[key]
return ok
}
func (m *ApprovalManager) allowSession(key string) {
if m == nil || key == "" {
return
}
m.mu.Lock()
defer m.mu.Unlock()
m.sessionAllowed[key] = struct{}{}
}
func approvalRequestWithEvaluation(req ApprovalRequest, evaluation ApprovalEvaluation) ApprovalRequest {
req.Summary = evaluation.Summary
req.Risk = evaluation.Risk
req.Reasons = slices.Clone(evaluation.Reasons)
return req
}
func approvalSessionKey(req ApprovalRequest) string {
if IsShellToolName(req.ToolName) {
if command, ok := stringApprovalArg(req.Args, "command"); ok {
return req.ToolName + ":" + command
}
}
switch req.ToolName {
case "edit":
if path, ok := stringApprovalArg(req.Args, "path"); ok {
return "edit:" + path
}
case "web_search":
if query, ok := stringApprovalArg(req.Args, "query"); ok {
return "web_search:" + query
}
case "web_fetch":
if targetURL, ok := stringApprovalArg(req.Args, "url"); ok {
return "web_fetch:" + targetURL
}
}
return req.ToolName + ":" + stableApprovalArgs(req.Args)
}
type DefaultApprovalPolicy struct{}
func (DefaultApprovalPolicy) EvaluateApproval(_ context.Context, req ApprovalRequest) ApprovalEvaluation {
switch req.ToolName {
case "read":
if path, ok := stringApprovalArg(req.Args, "path"); ok {
if reason := approvalPathEscapeReason(req.WorkingDir, path); reason != "" {
return denyApproval(req.ToolName, ApprovalRiskHigh, reason)
}
}
return ApprovalEvaluation{Decision: ApprovalAllowOnce, Risk: ApprovalRiskLow, Summary: fmt.Sprintf("%s can run without approval", ToolDisplayName(req.ToolName))}
case "web_search", "web_fetch":
return evaluateWebApproval(req)
case "edit":
return evaluateEditApproval(req)
case "bash", "powershell":
return evaluateShellApproval(req)
default:
return ApprovalEvaluation{
RequirePrompt: true,
Risk: ApprovalRiskMedium,
Summary: fmt.Sprintf("%s wants to run", ToolDisplayName(req.ToolName)),
Reasons: []string{"unknown tool effects"},
SessionKey: approvalSessionKey(req),
}
}
}
func evaluateEditApproval(req ApprovalRequest) ApprovalEvaluation {
path, ok := stringApprovalArg(req.Args, "path")
if !ok {
return denyApproval("edit", ApprovalRiskHigh, "missing path argument")
}
if reason := approvalPathEscapeReason(req.WorkingDir, path); reason != "" {
return denyApproval("edit", ApprovalRiskHigh, reason)
}
reasons := []string{"writes to a file"}
if replaceAll, _ := req.Args["replace_all"].(bool); replaceAll {
reasons = append(reasons, "may replace multiple matches")
}
return ApprovalEvaluation{
RequirePrompt: true,
Risk: ApprovalRiskMedium,
Summary: fmt.Sprintf("Edit wants to modify %s", sanitizeApprovalDisplay(path)),
Reasons: reasons,
SessionKey: "edit:" + path,
}
}
func evaluateWebApproval(req ApprovalRequest) ApprovalEvaluation {
switch req.ToolName {
case "web_search":
query, ok := stringApprovalArg(req.Args, "query")
if !ok {
return denyApproval("web_search", ApprovalRiskHigh, "missing query argument")
}
return ApprovalEvaluation{
RequirePrompt: true,
Risk: ApprovalRiskMedium,
Summary: fmt.Sprintf("Web Search wants to search for %q", sanitizeApprovalDisplay(query)),
Reasons: []string{"searches the web"},
SessionKey: "web_search:" + query,
}
case "web_fetch":
targetURL, ok := stringApprovalArg(req.Args, "url")
if !ok {
return denyApproval("web_fetch", ApprovalRiskHigh, "missing url argument")
}
return ApprovalEvaluation{
RequirePrompt: true,
Risk: ApprovalRiskMedium,
Summary: fmt.Sprintf("Web Fetch wants to fetch %s", sanitizeApprovalDisplay(targetURL)),
Reasons: []string{"fetches web content"},
SessionKey: "web_fetch:" + targetURL,
}
}
return ApprovalEvaluation{
RequirePrompt: true,
Risk: ApprovalRiskMedium,
Summary: fmt.Sprintf("%s wants to run", ToolDisplayName(req.ToolName)),
Reasons: []string{"accesses the web"},
SessionKey: approvalSessionKey(req),
}
}
func sanitizeApprovalDisplay(value string) string {
value = approvalANSIEscapePattern.ReplaceAllString(value, "")
value = strings.Map(func(r rune) rune {
switch r {
case '\n', '\r', '\t':
return ' '
}
if unicode.IsControl(r) {
return -1
}
return r
}, value)
value = strings.Join(strings.Fields(value), " ")
if value == "" {
return "(empty)"
}
return value
}
func evaluateShellApproval(req ApprovalRequest) ApprovalEvaluation {
command, ok := stringApprovalArg(req.Args, "command")
if !ok || strings.TrimSpace(command) == "" {
return denyApproval(req.ToolName, ApprovalRiskHigh, "missing command argument")
}
risk, reasons := classifyBashCommand(command)
if len(reasons) == 0 {
reasons = append(reasons, "runs shell commands")
}
return ApprovalEvaluation{
RequirePrompt: true,
Risk: risk,
Summary: fmt.Sprintf("%s wants to run a command", ToolDisplayName(req.ToolName)),
Reasons: reasons,
SessionKey: req.ToolName + ":" + command,
}
}
func denyApproval(toolName string, risk ApprovalRisk, reason string) ApprovalEvaluation {
return ApprovalEvaluation{
Decision: ApprovalDeny,
Risk: risk,
Summary: fmt.Sprintf("%s cannot run", ToolDisplayName(toolName)),
Reasons: []string{reason},
SessionKey: approvalSessionKey(ApprovalRequest{ToolName: toolName}),
}
}
// Shell approval is static analysis of model-generated commands, not a sandbox.
// Globs, aliases, environment, and runtime shell state are intentionally out of scope.
func classifyBashCommand(command string) (ApprovalRisk, []string) {
classifier := bashClassifier{risk: ApprovalRiskMedium}
tokens, scannerReasons, scannerHigh := scanBashTokens(command)
for _, reason := range scannerReasons {
classifier.addReason(reason)
}
if scannerHigh {
classifier.high = true
}
if bashFunctionDeclPattern.MatchString(command) {
classifier.addReason("defines shell functions")
classifier.high = true
}
if bashSubshellPattern.MatchString(command) {
classifier.addReason("uses a subshell")
classifier.high = true
}
for _, token := range tokens {
if token.kind != bashTokenOperator {
continue
}
switch token.value {
case "&&", "||", "|":
classifier.addReason("uses shell control operator " + token.value)
classifier.high = true
case ";", "\n":
classifier.addReason("uses shell statement separator")
classifier.high = true
case "&":
classifier.addReason("runs a command in the background")
classifier.high = true
case "$(", "`":
classifier.addReason("uses command substitution")
classifier.high = true
case "<(", ">(":
classifier.addReason("uses process substitution")
classifier.high = true
case ">", ">>", ">|", ">&", "&>":
classifier.addReason("writes or redirects files")
classifier.high = true
case "<", "<<", "<&":
classifier.addReason("uses shell redirection")
}
}
for _, args := range bashCommandCalls(tokens) {
classifier.addCall(args)
}
if classifier.high {
classifier.risk = ApprovalRiskHigh
}
return classifier.risk, classifier.reasons
}
type bashClassifier struct {
risk ApprovalRisk
high bool
reasons []string
}
func (c *bashClassifier) addReason(reason string) {
if reason == "" || slices.Contains(c.reasons, reason) {
return
}
c.reasons = append(c.reasons, reason)
}
func (c *bashClassifier) addCall(args []string) {
args = shellCommandArgs(args)
if len(args) == 0 {
return
}
if isDynamicCommandName(args[0]) {
c.addReason("uses dynamic command name")
c.high = true
}
name := shellCommandBase(args[0])
switch name {
case "cd":
c.addReason("changes directory")
c.high = true
case "eval":
c.addReason("evaluates shell code")
c.high = true
case "source", ".":
c.addReason("sources shell code")
c.high = true
case "exec":
c.addReason("replaces the shell process")
c.high = true
case "sudo":
c.addReason("runs with elevated privileges")
c.high = true
case "rm":
if hasAnyFlag(args[1:], "r", "R", "recursive") || hasAnyFlag(args[1:], "f", "force") {
c.addReason("removes files destructively")
c.high = true
}
case "git":
if isGitResetHard(args) {
c.addReason("runs destructive git reset")
c.high = true
}
if isGitCleanDestructive(args) {
c.addReason("runs destructive git clean")
c.high = true
}
case "find":
c.addFindReasons(args)
case "chmod", "chown":
if hasAnyFlag(args[1:], "R", "recursive") {
c.addReason("changes permissions or ownership recursively")
c.high = true
}
case "dd", "diskutil":
c.addReason("can write directly to disks")
c.high = true
case "mkfs", "newfs":
c.addReason("formats filesystems")
c.high = true
case "curl", "wget":
c.addReason("downloads remote content")
case "sh", "bash", "zsh":
if len(args) > 1 {
c.addReason("runs a shell interpreter")
c.high = true
}
}
}
func (c *bashClassifier) addFindReasons(args []string) {
for i := 1; i < len(args); i++ {
switch args[i] {
case "-delete":
c.addReason("deletes files via find")
c.high = true
case "-exec", "-execdir":
c.addReason("executes commands via find")
c.high = true
end := i + 1
for end < len(args) && args[end] != ";" && args[end] != `\;` && args[end] != "+" {
end++
}
if end > i+1 {
c.addCall(args[i+1 : end])
}
i = end
}
}
}
func hasAnyFlag(args []string, flags ...string) bool {
for _, arg := range args {
if arg == "--" {
return false
}
if !strings.HasPrefix(arg, "-") {
continue
}
for _, flag := range flags {
short := len(flag) == 1
if short && strings.HasPrefix(arg, "-") && !strings.HasPrefix(arg, "--") && strings.Contains(arg[1:], flag) {
return true
}
if arg == "--"+flag {
return true
}
}
}
return false
}
func isGitResetHard(args []string) bool {
subcommand := gitSubcommandArgs(args)
return len(subcommand) >= 2 && subcommand[0] == "reset" && slices.Contains(subcommand[1:], "--hard")
}
func isGitCleanDestructive(args []string) bool {
subcommand := gitSubcommandArgs(args)
if len(subcommand) < 1 || subcommand[0] != "clean" {
return false
}
return hasAnyFlag(subcommand[1:], "f", "force") && (hasAnyFlag(subcommand[1:], "d") || hasAnyFlag(subcommand[1:], "x", "X"))
}
func gitSubcommandArgs(args []string) []string {
for i := 1; i < len(args); i++ {
arg := args[i]
if arg == "--" {
return args[i+1:]
}
switch {
case arg == "-C", arg == "-c", arg == "--git-dir", arg == "--work-tree":
i++
continue
case strings.HasPrefix(arg, "-C"), strings.HasPrefix(arg, "-c"):
continue
case strings.HasPrefix(arg, "--git-dir="), strings.HasPrefix(arg, "--work-tree="):
continue
case strings.HasPrefix(arg, "-"):
continue
default:
return args[i:]
}
}
return nil
}
type bashTokenKind int
const (
bashTokenWord bashTokenKind = iota
bashTokenOperator
)
type bashToken struct {
kind bashTokenKind
value string
}
var (
approvalANSIEscapePattern = regexp.MustCompile(`\x1b(?:\[[0-?]*[ -/]*[@-~]|\][^\x07]*(?:\x07|\x1b\\))`)
bashFunctionDeclPattern = regexp.MustCompile(`(?m)(^|[;&|[:space:]])(?:function[[:space:]]+)?[A-Za-z_][A-Za-z0-9_]*[[:space:]]*(?:\(\)[[:space:]]*)?\{`)
bashSubshellPattern = regexp.MustCompile(`(?m)(^|[;&|[:space:]])\(`)
)
func scanBashTokens(command string) ([]bashToken, []string, bool) {
var tokens []bashToken
var reasons []string
var word strings.Builder
var quote byte
escaped := false
high := false
flushWord := func() {
if word.Len() == 0 {
return
}
tokens = append(tokens, bashToken{kind: bashTokenWord, value: word.String()})
word.Reset()
}
addOperator := func(op string) {
tokens = append(tokens, bashToken{kind: bashTokenOperator, value: op})
}
for i := 0; i < len(command); i++ {
ch := command[i]
if escaped {
word.WriteByte(ch)
escaped = false
continue
}
if ch == '\\' && quote != '\'' {
escaped = true
continue
}
if quote == '\'' {
if ch == '\'' {
quote = 0
} else {
word.WriteByte(ch)
}
continue
}
if quote == '"' {
switch {
case ch == '"':
quote = 0
case ch == '`':
addOperator("`")
word.WriteByte(ch)
case ch == '$' && i+1 < len(command) && command[i+1] == '(':
addOperator("$(")
word.WriteString("$(")
i++
default:
word.WriteByte(ch)
}
continue
}
if ch == '\'' || ch == '"' {
quote = ch
continue
}
if ch == '`' {
addOperator("`")
word.WriteByte(ch)
continue
}
if ch == '$' && i+1 < len(command) && command[i+1] == '(' {
addOperator("$(")
word.WriteString("$(")
i++
continue
}
if (ch == '<' || ch == '>') && i+1 < len(command) && command[i+1] == '(' {
flushWord()
addOperator(string([]byte{ch, '('}))
i++
continue
}
if ch == '\n' {
flushWord()
addOperator("\n")
continue
}
if ch == ' ' || ch == '\t' || ch == '\r' {
flushWord()
continue
}
switch ch {
case '&':
flushWord()
switch {
case i+1 < len(command) && command[i+1] == '&':
addOperator("&&")
i++
case i+1 < len(command) && command[i+1] == '>':
addOperator("&>")
i++
default:
addOperator("&")
}
case '|':
flushWord()
if i+1 < len(command) && command[i+1] == '|' {
addOperator("||")
i++
} else {
addOperator("|")
}
case ';':
flushWord()
addOperator(";")
case '<':
flushWord()
if i+1 < len(command) && command[i+1] == '<' {
addOperator("<<")
i++
} else if i+1 < len(command) && command[i+1] == '&' {
addOperator("<&")
i++
} else {
addOperator("<")
}
case '>':
flushWord()
if i+1 < len(command) && command[i+1] == '>' {
addOperator(">>")
i++
} else if i+1 < len(command) && command[i+1] == '|' {
addOperator(">|")
i++
} else if i+1 < len(command) && command[i+1] == '&' {
addOperator(">&")
i++
} else {
addOperator(">")
}
default:
word.WriteByte(ch)
}
}
if escaped || quote != 0 {
reasons = append(reasons, "could not parse shell command")
high = true
}
flushWord()
return tokens, reasons, high
}
func bashCommandCalls(tokens []bashToken) [][]string {
var calls [][]string
var current []string
flush := func() {
if len(current) == 0 {
return
}
calls = append(calls, current)
current = nil
}
for _, token := range tokens {
if token.kind == bashTokenWord {
current = append(current, token.value)
continue
}
if isBashCommandBoundary(token.value) {
flush()
}
}
flush()
return calls
}
func isBashCommandBoundary(op string) bool {
switch op {
case "&&", "||", "|", ";", "&", "\n":
return true
default:
return false
}
}
func shellCommandArgs(args []string) []string {
for len(args) > 0 && isShellAssignment(args[0]) {
args = args[1:]
}
return args
}
func isShellAssignment(word string) bool {
name, _, ok := strings.Cut(word, "=")
if !ok || name == "" {
return false
}
for i := range len(name) {
ch := name[i]
if i == 0 {
if !((ch >= 'A' && ch <= 'Z') || (ch >= 'a' && ch <= 'z') || ch == '_') {
return false
}
continue
}
if !((ch >= 'A' && ch <= 'Z') || (ch >= 'a' && ch <= 'z') || (ch >= '0' && ch <= '9') || ch == '_') {
return false
}
}
return true
}
func isDynamicCommandName(name string) bool {
return strings.HasPrefix(name, "$") || strings.HasPrefix(name, "`") || strings.Contains(name, "$(") || strings.Contains(name, "`")
}
func shellCommandBase(name string) string {
name = strings.TrimSpace(name)
if i := strings.LastIndexAny(name, `/\`); i >= 0 && i+1 < len(name) {
name = name[i+1:]
}
return name
}
func approvalPathEscapeReason(workingDir, path string) string {
if strings.TrimSpace(path) == "" {
return ""
}
if filepath.IsAbs(path) {
return "absolute paths are not allowed"
}
base := workingDir
if base == "" {
var err error
base, err = os.Getwd()
if err != nil {
return "could not determine working directory"
}
}
base, err := canonicalApprovalPath(base)
if err != nil {
return "could not resolve working directory"
}
resolved := filepath.Clean(filepath.Join(base, path))
resolvedForCheck := resolved
if canonical, err := canonicalApprovalPath(resolved); err == nil {
resolvedForCheck = canonical
}
rel, err := filepath.Rel(base, resolvedForCheck)
if err != nil {
return "could not resolve path"
}
if rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
return "path escapes working directory"
}
return ""
}
func canonicalApprovalPath(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 stringApprovalArg(args map[string]any, key string) (string, bool) {
value, ok := args[key].(string)
return value, ok && strings.TrimSpace(value) != ""
}
func stableApprovalArgs(args map[string]any) string {
if len(args) == 0 {
return ""
}
keys := make([]string, 0, len(args))
for key := range args {
keys = append(keys, key)
}
slices.Sort(keys)
var b bytes.Buffer
for _, key := range keys {
if b.Len() > 0 {
b.WriteByte(',')
}
fmt.Fprintf(&b, "%s=%v", key, args[key])
}
return b.String()
}
+384
View File
@@ -0,0 +1,384 @@
package agent
import (
"context"
"strings"
"testing"
"github.com/ollama/ollama/api"
)
type recordingApprovalPrompter struct {
requests []ApprovalRequest
results []ApprovalResult
}
type allowWithoutPromptPolicy struct{}
type approvalRequiredTestTool struct{}
func (p *recordingApprovalPrompter) PromptApproval(_ context.Context, request ApprovalRequest) (ApprovalResult, error) {
p.requests = append(p.requests, request)
if len(p.results) == 0 {
return ApprovalResult{Decision: ApprovalAllowOnce}, nil
}
result := p.results[0]
p.results = p.results[1:]
return result, nil
}
func (allowWithoutPromptPolicy) EvaluateApproval(context.Context, ApprovalRequest) ApprovalEvaluation {
return ApprovalEvaluation{Decision: ApprovalAllowOnce, Risk: ApprovalRiskLow}
}
func (approvalRequiredTestTool) Name() string {
return "approval_required"
}
func (approvalRequiredTestTool) Description() string {
return "requires approval"
}
func (approvalRequiredTestTool) Schema() api.ToolFunction {
return api.ToolFunction{Name: "approval_required"}
}
func (approvalRequiredTestTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) {
return ToolResult{Content: "ok"}, nil
}
func (approvalRequiredTestTool) RequiresApproval(map[string]any) bool {
return true
}
func TestApprovalManagerAllowsSafeToolsWithoutPrompt(t *testing.T) {
prompter := &recordingApprovalPrompter{}
manager := NewApprovalManager(ApprovalManagerOptions{Prompter: prompter})
result, err := manager.Approve(context.Background(), ApprovalRequest{
ToolName: "read",
Args: map[string]any{"path": "README.md"},
WorkingDir: t.TempDir(),
})
if err != nil {
t.Fatal(err)
}
if result.Decision != ApprovalAllowOnce {
t.Fatalf("decision = %q, want allow_once", result.Decision)
}
if len(prompter.requests) != 0 {
t.Fatalf("safe tool prompted: %#v", prompter.requests)
}
}
func TestApprovalManagerToolRequiredOverridePromptsInApprove(t *testing.T) {
prompter := &recordingApprovalPrompter{}
manager := NewApprovalManager(ApprovalManagerOptions{Policy: allowWithoutPromptPolicy{}, Prompter: prompter})
tool := approvalRequiredTestTool{}
request := ApprovalRequest{
ToolName: tool.Name(),
Args: map[string]any{},
ToolApprovalRequired: ToolRequiresApproval(tool, nil),
}
if !manager.RequiresApproval(context.Background(), tool, request) {
t.Fatal("tool-required approval should require a prompt")
}
result, err := manager.Approve(context.Background(), request)
if err != nil {
t.Fatal(err)
}
if result.Decision != ApprovalAllowOnce {
t.Fatalf("decision = %q, want allow_once", result.Decision)
}
if len(prompter.requests) != 1 {
t.Fatalf("prompts = %d, want 1", len(prompter.requests))
}
}
func TestApprovalManagerDeniesEscapingPath(t *testing.T) {
manager := NewApprovalManager(ApprovalManagerOptions{})
result, err := manager.Approve(context.Background(), ApprovalRequest{
ToolName: "edit",
Args: map[string]any{"path": "../outside.txt"},
WorkingDir: t.TempDir(),
})
if err != nil {
t.Fatal(err)
}
if result.Decision != ApprovalDeny {
t.Fatalf("decision = %q, want deny", result.Decision)
}
if !strings.Contains(result.Reason, "path escapes working directory") {
t.Fatalf("reason = %q", result.Reason)
}
}
func TestApprovalManagerSanitizesEditSummary(t *testing.T) {
evaluation := evaluateEditApproval(ApprovalRequest{
ToolName: "edit",
Args: map[string]any{"path": "notes/\x1b[31mred\nfile.txt"},
WorkingDir: t.TempDir(),
})
if strings.ContainsAny(evaluation.Summary, "\n\r\x1b") {
t.Fatalf("summary contains control characters: %q", evaluation.Summary)
}
if !strings.Contains(evaluation.Summary, "notes/red file.txt") {
t.Fatalf("summary = %q, want sanitized path", evaluation.Summary)
}
}
func TestApprovalManagerPromptsForEdit(t *testing.T) {
prompter := &recordingApprovalPrompter{}
manager := NewApprovalManager(ApprovalManagerOptions{Prompter: prompter})
result, err := manager.Approve(context.Background(), ApprovalRequest{
ToolName: "edit",
Args: map[string]any{"path": "note.txt", "old_text": "old", "new_text": "new"},
WorkingDir: t.TempDir(),
})
if err != nil {
t.Fatal(err)
}
if result.Decision != ApprovalAllowOnce {
t.Fatalf("decision = %q, want allow_once", result.Decision)
}
if len(prompter.requests) != 1 {
t.Fatalf("prompts = %d, want 1", len(prompter.requests))
}
request := prompter.requests[0]
if request.Risk != ApprovalRiskMedium {
t.Fatalf("risk = %q, want medium", request.Risk)
}
if !strings.Contains(strings.Join(request.Reasons, " "), "writes to a file") {
t.Fatalf("reasons = %#v", request.Reasons)
}
}
func TestApprovalManagerHeadlessDeniesPromptRequiredTools(t *testing.T) {
manager := NewApprovalManager(ApprovalManagerOptions{})
result, err := manager.Approve(context.Background(), ApprovalRequest{
ToolName: "bash",
Args: map[string]any{"command": "pwd"},
})
if err != nil {
t.Fatal(err)
}
if result.Decision != ApprovalDeny {
t.Fatalf("decision = %q, want deny", result.Decision)
}
if !strings.Contains(result.Reason, "--auto-approve-tools") {
t.Fatalf("reason = %q", result.Reason)
}
}
func TestApprovalManagerSessionAllowList(t *testing.T) {
prompter := &recordingApprovalPrompter{
results: []ApprovalResult{{Decision: ApprovalAllowSession}},
}
manager := NewApprovalManager(ApprovalManagerOptions{Prompter: prompter})
request := ApprovalRequest{
ToolName: "bash",
Args: map[string]any{"command": "go test ./agent"},
}
result, err := manager.Approve(context.Background(), request)
if err != nil {
t.Fatal(err)
}
if result.Decision != ApprovalAllowSession {
t.Fatalf("decision = %q, want allow_session", result.Decision)
}
result, err = manager.Approve(context.Background(), request)
if err != nil {
t.Fatal(err)
}
if result.Decision != ApprovalAllowOnce {
t.Fatalf("second decision = %q, want allow_once", result.Decision)
}
if len(prompter.requests) != 1 {
t.Fatalf("prompts = %d, want 1", len(prompter.requests))
}
}
func TestBashApprovalClassifiesHighRiskShell(t *testing.T) {
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
ToolName: "bash",
Args: map[string]any{"command": "cd / && rm -rf tmp"},
})
if !evaluation.RequirePrompt {
t.Fatal("bash should require prompt")
}
if evaluation.Risk != ApprovalRiskHigh {
t.Fatalf("risk = %q, want high", evaluation.Risk)
}
reasons := strings.Join(evaluation.Reasons, " ")
for _, want := range []string{"changes directory", "control operator", "removes files"} {
if !strings.Contains(reasons, want) {
t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, want)
}
}
}
func TestPowerShellApprovalUsesShellPolicy(t *testing.T) {
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
ToolName: "powershell",
Args: map[string]any{"command": "Remove-Item -Recurse tmp"},
})
if !evaluation.RequirePrompt {
t.Fatal("powershell should require prompt")
}
if evaluation.Summary != "PowerShell wants to run a command" {
t.Fatalf("summary = %q", evaluation.Summary)
}
if evaluation.SessionKey != "powershell:Remove-Item -Recurse tmp" {
t.Fatalf("session key = %q", evaluation.SessionKey)
}
}
func TestBashApprovalClassifiesDynamicShellEvasions(t *testing.T) {
tests := []struct {
name string
cmd string
reason string
}{
{name: "function declaration", cmd: "f() { rm -rf /; } && f", reason: "defines shell functions"},
{name: "eval", cmd: `eval "$cmd"`, reason: "evaluates shell code"},
{name: "variable command name", cmd: "$DANGER --flag", reason: "dynamic command name"},
{name: "command substitution command name", cmd: "$(echo rm) -rf /", reason: "dynamic command name"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
ToolName: "bash",
Args: map[string]any{"command": tt.cmd},
})
if !evaluation.RequirePrompt {
t.Fatal("bash evasion should require prompt")
}
if evaluation.Risk != ApprovalRiskHigh {
t.Fatalf("risk = %q, want high", evaluation.Risk)
}
if reasons := strings.Join(evaluation.Reasons, " "); !strings.Contains(reasons, tt.reason) {
t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, tt.reason)
}
})
}
}
func TestBashApprovalClassifiesDestructiveGitWithGlobalOptions(t *testing.T) {
tests := []struct {
name string
cmd string
reason string
}{
{name: "git reset hard after cwd", cmd: "git -C /tmp reset --hard", reason: "runs destructive git reset"},
{name: "git reset hard after config", cmd: "git -c core.autocrlf=false reset --hard", reason: "runs destructive git reset"},
{name: "git clean after work tree", cmd: "git --work-tree=/tmp clean -fdx", reason: "runs destructive git clean"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
ToolName: "bash",
Args: map[string]any{"command": tt.cmd},
})
if evaluation.Risk != ApprovalRiskHigh {
t.Fatalf("risk = %q, want high", evaluation.Risk)
}
if reasons := strings.Join(evaluation.Reasons, " "); !strings.Contains(reasons, tt.reason) {
t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, tt.reason)
}
})
}
}
func TestBashApprovalClassifiesFindMutations(t *testing.T) {
tests := []struct {
name string
cmd string
reason string
}{
{name: "delete", cmd: "find . -name '*.tmp' -delete", reason: "deletes files via find"},
{name: "exec", cmd: `find . -name '*.tmp' -exec rm -rf {} \;`, reason: "executes commands via find"},
{name: "exec nested destructive command", cmd: `find . -name '*.tmp' -exec rm -rf {} \;`, reason: "removes files destructively"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
ToolName: "bash",
Args: map[string]any{"command": tt.cmd},
})
if evaluation.Risk != ApprovalRiskHigh {
t.Fatalf("risk = %q, want high", evaluation.Risk)
}
if reasons := strings.Join(evaluation.Reasons, " "); !strings.Contains(reasons, tt.reason) {
t.Fatalf("reasons = %#v, want %q", evaluation.Reasons, tt.reason)
}
})
}
}
func TestWebApprovalRequiresPrompt(t *testing.T) {
tests := []struct {
name string
tool string
args map[string]any
summary string
}{
{
name: "search",
tool: "web_search",
args: map[string]any{"query": "Ollama agents"},
summary: "Web Search wants to search for \"Ollama agents\"",
},
{
name: "fetch",
tool: "web_fetch",
args: map[string]any{"url": "https://ollama.com"},
summary: "Web Fetch wants to fetch https://ollama.com",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
ToolName: tt.tool,
Args: tt.args,
})
if !evaluation.RequirePrompt {
t.Fatal("web tool should require prompt")
}
if evaluation.Decision != "" && evaluation.Decision != ApprovalAllowOnce {
t.Fatalf("decision = %q, want allow once", evaluation.Decision)
}
if evaluation.Risk != ApprovalRiskMedium {
t.Fatalf("risk = %q, want medium", evaluation.Risk)
}
if evaluation.Summary != tt.summary {
t.Fatalf("summary = %q, want %q", evaluation.Summary, tt.summary)
}
})
}
}
func TestWebApprovalDeniesMissingArgs(t *testing.T) {
for _, tool := range []string{"web_search", "web_fetch"} {
evaluation := DefaultApprovalPolicy{}.EvaluateApproval(context.Background(), ApprovalRequest{
ToolName: tool,
Args: map[string]any{},
})
if evaluation.Decision != ApprovalDeny {
t.Fatalf("%s missing args decision = %q, want deny", tool, evaluation.Decision)
}
if evaluation.Risk != ApprovalRiskHigh {
t.Fatalf("%s missing args risk = %q, want high", tool, evaluation.Risk)
}
}
}
+636
View File
@@ -0,0 +1,636 @@
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; both
// the in-memory compactor (this package) and the on-disk chat store
// (package store) build and detect summaries through them.
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 CLI 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 CompactionStore interface {
ArchiveForCompaction(context.Context, string, int, string, bool) 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
Store CompactionStore
Options CompactionOptions
}
func NewSimpleCompactor(client ChatClient, store CompactionStore, opts CompactionOptions) *SimpleCompactor {
return &SimpleCompactor{Client: client, Store: store, Options: opts}
}
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, keptUserTurns, 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 == "" {
// TODO(parthsareen): Investigate models that stream compaction output
// without final content, such as thinking-only summaries.
result.Reason = "summary was empty"
return result, nil
}
if c.Store != nil && req.ChatID != "" {
if err := c.Store.ArchiveForCompaction(ctx, req.ChatID, keptUserTurns, summary, req.ContinueTask); err != nil {
result.Reason = err.Error()
return result, err
}
}
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
}
// TODO(parthsareen): If the newest kept user turn contains the oversized
// tool output, compaction can remove older history but still leave the next
// prompt above the safety threshold. Pair this estimate trigger with
// context-aware tool-output paging/range reads so the kept suffix can shrink.
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. This is the canonical builder used by both the compactor
// and the chat store.
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)
}
// EstimateTokens returns the agent's lightweight token estimate for UI hints.
func EstimateTokens(text string) int {
return estimateCompactionTokens(text)
}
// EstimatePromptTokens returns the agent's lightweight estimate for the prompt
// payload sent to /api/chat.
func EstimatePromptTokens(systemPrompt string, messages []api.Message, tools api.Tools, format string) int {
return estimateCompactionRequestTokens(CompactionRequest{
SystemPrompt: systemPrompt,
Messages: messages,
Tools: tools,
Format: format,
})
}
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 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
}
}
+816
View File
@@ -0,0 +1,816 @@
package agent
import (
"context"
"net/http"
"strings"
"testing"
"github.com/ollama/ollama/api"
)
type compactionStore struct {
chatID string
keepUserTurns int
summary string
continueTask bool
}
func (s *compactionStore) ArchiveForCompaction(_ context.Context, chatID string, keepUserTurns int, summary string, continueTask bool) error {
s.chatID = chatID
s.keepUserTurns = keepUserTurns
s.summary = summary
s.continueTask = continueTask
return nil
}
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"}},
}},
}
store := &compactionStore{}
compactor := NewSimpleCompactor(client, store, 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 store.chatID != "chat-1" || store.keepUserTurns != 2 || store.summary != "summary" || store.continueTask {
t.Fatalf("archive call = %#v", store)
}
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"}},
}},
}
store := &compactionStore{}
compactor := NewSimpleCompactor(client, store, 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 store.keepUserTurns != 0 {
t.Fatalf("keepUserTurns = %d, want 0 for small context", store.keepUserTurns)
}
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"}},
}},
}
store := &compactionStore{}
compactor := NewSimpleCompactor(client, store, 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)
}
if !store.continueTask || store.summary != "summary" {
t.Fatalf("archive call = %#v", store)
}
}
func TestSimpleCompactorTruncatesOversizedSummary(t *testing.T) {
longSummary := strings.Repeat("x", maxCompactionSummaryBytes+1024)
client := &fakeClient{
responses: [][]api.ChatResponse{{
{Message: api.Message{Role: "assistant", Content: longSummary}},
}},
}
store := &compactionStore{}
compactor := NewSimpleCompactor(client, store, 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 store.summary != result.Summary {
t.Fatalf("stored summary mismatch")
}
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 := NewSimpleCompactor(client, nil, 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 := NewSimpleCompactor(client, nil, 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 := NewSimpleCompactor(client, nil, 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"}},
}},
}
store := &compactionStore{}
compactor := NewSimpleCompactor(client, store, 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 store.keepUserTurns != 1 {
t.Fatalf("kept user turns = %d, want 1", store.keepUserTurns)
}
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"}},
}},
}
store := &compactionStore{}
compactor := NewSimpleCompactor(client, store, 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 store.keepUserTurns != 0 {
t.Fatalf("kept user turns = %d, want 0", store.keepUserTurns)
}
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 := NewSimpleCompactor(client, nil, 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 := NewSimpleCompactor(client, nil, 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 := NewSimpleCompactor(nil, nil, 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 := NewSimpleCompactor(client, nil, 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"}},
}},
}
store := &compactionStore{}
compactor := NewSimpleCompactor(client, store, 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")
}
if store.keepUserTurns != 3 {
t.Fatalf("keepUserTurns = %d, want 3", store.keepUserTurns)
}
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 := NewSimpleCompactor(client, nil, 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 := NewSimpleCompactor(client, nil, 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)
}
}
+96
View File
@@ -0,0 +1,96 @@
package agent
import (
"context"
"time"
"github.com/ollama/ollama/api"
)
type EventType string
const (
EventMessageStarted EventType = "message_started"
EventMessageDelta EventType = "message_delta"
EventThinkingDelta EventType = "thinking_delta"
EventToolCallDetected EventType = "tool_call_detected"
EventToolStarted EventType = "tool_started"
EventToolFinished EventType = "tool_finished"
EventToolsUnavailable EventType = "tools_unavailable"
EventCompactionStarted EventType = "compaction_started"
EventCompactionProgress EventType = "compaction_progress"
EventCompacted EventType = "compacted"
EventCompactionSkipped EventType = "compaction_skipped"
EventLoopStep EventType = "loop_step"
EventRequestBuilt EventType = "request_built"
EventModelStreamDone EventType = "model_stream_done"
EventRunFinished EventType = "run_finished"
EventError EventType = "error"
)
type Event struct {
Type EventType `json:"type"`
RunID string `json:"runId,omitempty"`
ChatID string `json:"chatId,omitempty"`
MessageID string `json:"messageId,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"`
PromptTokens int `json:"promptTokens,omitempty"`
Error string `json:"error,omitempty"`
StartedAt time.Time `json:"startedAt,omitempty"`
FinishedAt time.Time `json:"finishedAt,omitempty"`
Response *api.ChatResponse `json:"-"`
}
type EventSink interface {
Emit(Event) error
}
type MultiEventSink []EventSink
func (s MultiEventSink) Emit(event Event) error {
var firstErr error
for _, sink := range s {
if sink == nil {
continue
}
if err := sink.Emit(event); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}
type EventSinkFunc func(Event) error
func (fn EventSinkFunc) Emit(event Event) error {
if fn == nil {
return nil
}
return fn(event)
}
func emit(sink EventSink, event Event) error {
if sink == nil {
return nil
}
return sink.Emit(event)
}
func emitIgnoringCanceled(ctx context.Context, sink EventSink, event Event) error {
err := emit(sink, 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
}
+92
View File
@@ -0,0 +1,92 @@
package agent
import "sync"
type ToolMode int
const (
ToolModeReview ToolMode = iota
ToolModeFullAccess
ToolModeDisabled
)
type RunPolicy struct {
ToolMode ToolMode
ApprovalPolicy ApprovalPolicy
// MaxToolRounds limits consecutive model/tool cycles.
// Zero uses the default guard; negative disables the guard for tests or
// special callers.
MaxToolRounds int
}
func (p RunPolicy) UsesTools() bool {
switch p.ToolMode {
case ToolModeReview, ToolModeFullAccess:
return true
default:
return false
}
}
func (p RunPolicy) Tools(registry *Registry) *Registry {
if !p.UsesTools() {
return nil
}
return registry
}
func (p RunPolicy) ApprovalHandler(prompter ApprovalPrompter) ApprovalHandler {
if p.ToolMode == ToolModeFullAccess {
return AutoAllowApproval{}
}
policy := p.ApprovalPolicy
if policy == nil {
policy = DefaultApprovalPolicy{}
}
return NewApprovalManager(ApprovalManagerOptions{
Policy: policy,
Prompter: prompter,
})
}
func (p RunPolicy) ReviewApprovalHandler(prompter ApprovalPrompter) ApprovalHandler {
policy := p.ApprovalPolicy
if policy == nil {
policy = DefaultApprovalPolicy{}
}
return NewApprovalManager(ApprovalManagerOptions{
Policy: policy,
Prompter: prompter,
})
}
type RunPolicyState struct {
mu sync.Mutex
policy RunPolicy
}
func NewRunPolicyState(policy RunPolicy) *RunPolicyState {
return &RunPolicyState{policy: policy}
}
func (s *RunPolicyState) Policy() RunPolicy {
if s == nil {
return RunPolicy{}
}
s.mu.Lock()
defer s.mu.Unlock()
return s.policy
}
func (s *RunPolicyState) ToolMode() ToolMode {
return s.Policy().ToolMode
}
func (s *RunPolicyState) SetToolMode(mode ToolMode) {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
s.policy.ToolMode = mode
}
+103
View File
@@ -0,0 +1,103 @@
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 NewRegistry() *Registry {
return &Registry{tools: make(map[string]Tool)}
}
func (r *Registry) Register(tool Tool) {
if r == nil || tool == nil {
return
}
r.tools[tool.Name()] = tool
}
func (r *Registry) Has(name string) bool {
if r == nil {
return false
}
_, ok := r.tools[name]
return ok
}
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
}
+64
View File
@@ -0,0 +1,64 @@
package agent
import (
"encoding/json"
"strings"
"github.com/ollama/ollama/api"
)
// ChatRequestPreview is the request body plus the estimated prompt tokens for it.
type ChatRequestPreview struct {
Request api.ChatRequest
PromptTokens int
}
// BuildChatRequestPreview builds the chat request shape used for a run and its estimated prompt tokens.
func BuildChatRequestPreview(opts RunOptions, messages []api.Message, tools api.Tools) ChatRequestPreview {
return ChatRequestPreview{
Request: buildChatRequest(opts, messages, tools),
PromptTokens: EstimateChatRequestPromptTokens(opts, messages, tools),
}
}
// EstimateChatRequestPromptTokens estimates the prompt tokens for a chat request before sending it.
func EstimateChatRequestPromptTokens(opts RunOptions, messages []api.Message, tools api.Tools) int {
return estimateCompactionRequestTokens(CompactionRequest{
SystemPrompt: opts.SystemPrompt,
Messages: sanitizeMessagesForRequest(messages),
Tools: tools,
Format: opts.Format,
Options: opts.Options,
})
}
func buildChatRequest(opts RunOptions, messages []api.Message, tools api.Tools) api.ChatRequest {
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...)
}
req := api.ChatRequest{
Model: opts.Model,
Messages: requestMessages,
Format: json.RawMessage(chatRequestFormat(opts.Format)),
Options: opts.Options,
Think: opts.Think,
}
if opts.KeepAlive != nil {
req.KeepAlive = opts.KeepAlive
}
if len(tools) > 0 {
req.Tools = tools
}
return req
}
func chatRequestFormat(format string) string {
if format == "json" {
return `"` + format + `"`
}
return format
}
+46
View File
@@ -0,0 +1,46 @@
package agent
import (
"testing"
"github.com/ollama/ollama/api"
)
func TestBuildChatRequestPreviewBuildsModelRequest(t *testing.T) {
tools := api.Tools{{
Type: "function",
Function: api.ToolFunction{
Name: "read",
Description: "read a file",
Parameters: api.ToolFunctionParameters{Type: "object"},
},
}}
preview := BuildChatRequestPreview(RunOptions{
Model: "llama3.2",
SystemPrompt: "You are Ollama.",
Format: "json",
Options: map[string]any{"temperature": 0.1},
}, []api.Message{{Role: "user", Content: "hello"}}, tools)
if preview.Request.Model != "llama3.2" {
t.Fatalf("model = %q, want llama3.2", preview.Request.Model)
}
if got := string(preview.Request.Format); got != `"json"` {
t.Fatalf("format = %q, want quoted json", got)
}
if len(preview.Request.Messages) != 2 {
t.Fatalf("messages = %d, want 2", len(preview.Request.Messages))
}
if preview.Request.Messages[0].Role != "system" || preview.Request.Messages[0].Content != "You are Ollama." {
t.Fatalf("system message = %#v", preview.Request.Messages[0])
}
if preview.Request.Messages[1].Role != "user" || preview.Request.Messages[1].Content != "hello" {
t.Fatalf("user message = %#v", preview.Request.Messages[1])
}
if len(preview.Request.Tools) != 1 {
t.Fatalf("tools = %d, want 1", len(preview.Request.Tools))
}
if preview.PromptTokens <= 0 {
t.Fatalf("prompt tokens = %d, want positive", preview.PromptTokens)
}
}
+1033
View File
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
+165
View File
@@ -0,0 +1,165 @@
package skills
import (
"errors"
"fmt"
"os"
"path/filepath"
"slices"
"strings"
)
type ImportResult struct {
Source string
Skill Skill
From string
To string
Skipped bool
Error string
}
type skillDirCandidate struct {
Dir string
Skipped bool
Error string
}
func Import(source string, force bool) ([]ImportResult, error) {
dest, err := DefaultDir()
if err != nil {
return nil, err
}
return ImportToDir(source, dest, force)
}
func ImportToDir(source, dest string, force bool) ([]ImportResult, error) {
roots, err := SourceDirs(source)
if err != nil {
return nil, err
}
if err := os.MkdirAll(dest, 0o755); err != nil {
return nil, fmt.Errorf("create skills directory: %w", err)
}
var results []ImportResult
for _, root := range roots {
candidates, err := skillDirs(root)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
continue
}
results = append(results, ImportResult{Source: source, From: root, Skipped: true, Error: err.Error()})
continue
}
for _, candidate := range candidates {
result := ImportResult{Source: source, From: candidate.Dir}
if candidate.Skipped {
result.Skipped = true
result.Error = candidate.Error
results = append(results, result)
continue
}
skill, err := ReadMetadata(filepath.Join(candidate.Dir, SkillFile))
if err != nil {
result.Skipped = true
result.Error = err.Error()
results = append(results, result)
continue
}
result.Skill = skill
result.To = filepath.Join(dest, skill.Name)
copyResult, err := copyDir(candidate.Dir, result.To, force)
if errors.Is(err, os.ErrExist) {
result.Skipped = true
result.Error = "already exists"
} else if err != nil {
result.Skipped = true
result.Error = err.Error()
} else if len(copyResult.Skipped) > 0 {
result.Error = "skipped symlinks: " + strings.Join(copyResult.Skipped, ", ")
}
results = append(results, result)
}
}
slices.SortFunc(results, func(a, b ImportResult) int {
return strings.Compare(a.Skill.Name+a.From, b.Skill.Name+b.From)
})
return results, nil
}
func SourceDirs(source string) ([]string, error) {
source = strings.ToLower(strings.TrimSpace(source))
if source == "" {
source = "all"
}
home, err := os.UserHomeDir()
if err != nil {
return nil, fmt.Errorf("resolve home directory: %w", err)
}
dirs := map[string][]string{
"claude": {filepath.Join(home, ".claude", "skills")},
"codex": {filepath.Join(home, ".codex", "skills")},
"pi": {filepath.Join(home, ".pi", "skills"), filepath.Join(home, ".agents", "skills")},
"agents": {filepath.Join(home, ".agents", "skills")},
}
if source == "all" {
var all []string
for _, name := range []string{"claude", "codex", "pi"} {
all = append(all, dirs[name]...)
}
return uniqueStrings(all), nil
}
if roots, ok := dirs[source]; ok {
return roots, nil
}
return nil, fmt.Errorf("unknown skill source %q (use claude, codex, pi, agents, or all)", source)
}
func skillDirs(root string) ([]skillDirCandidate, error) {
entries, err := os.ReadDir(root)
if err != nil {
return nil, err
}
var dirs []skillDirCandidate
for _, entry := range entries {
if strings.HasPrefix(entry.Name(), ".") {
continue
}
dir := filepath.Join(root, entry.Name())
if entry.Type()&os.ModeSymlink != 0 {
dirs = append(dirs, skillDirCandidate{
Dir: dir,
Skipped: true,
Error: "symlinked skill directories are not supported",
})
continue
}
if !entry.IsDir() {
continue
}
if _, err := os.Stat(filepath.Join(dir, SkillFile)); err == nil {
dirs = append(dirs, skillDirCandidate{Dir: dir})
}
}
slices.SortFunc(dirs, func(a, b skillDirCandidate) int {
return strings.Compare(a.Dir, b.Dir)
})
return dirs, nil
}
func uniqueStrings(values []string) []string {
seen := make(map[string]bool)
var out []string
for _, value := range values {
if seen[value] {
continue
}
seen[value] = true
out = append(out, value)
}
return out
}
+319
View File
@@ -0,0 +1,319 @@
package skills
import (
"errors"
"fmt"
"io/fs"
"os"
"path/filepath"
"regexp"
"slices"
"strings"
"gopkg.in/yaml.v3"
)
const (
SkillFile = "SKILL.md"
maxSkillFileBytes = 1 << 20
)
var validName = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,63}$`)
type Skill struct {
Name string
Description string
Dir string
File string
}
type Catalog struct {
Dir string
Skills []Skill
Warnings []string
}
type frontmatter struct {
Name string `yaml:"name"`
Description string `yaml:"description"`
}
func DefaultDir() (string, error) {
home, err := os.UserHomeDir()
if err != nil {
return "", fmt.Errorf("resolve home directory: %w", err)
}
return filepath.Join(home, ".ollama", "skills"), nil
}
func LoadDefault() (*Catalog, error) {
dir, err := DefaultDir()
if err != nil {
return nil, err
}
return Load(dir)
}
func Load(dir string) (*Catalog, error) {
catalog := &Catalog{Dir: dir}
entries, err := os.ReadDir(dir)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
return catalog, nil
}
return nil, fmt.Errorf("read skills directory: %w", err)
}
seen := make(map[string]string)
for _, entry := range entries {
if !entry.IsDir() || strings.HasPrefix(entry.Name(), ".") {
continue
}
skillDir := filepath.Join(dir, entry.Name())
skill, err := ReadMetadata(filepath.Join(skillDir, SkillFile))
if err != nil {
catalog.Warnings = append(catalog.Warnings, fmt.Sprintf("%s: %v", skillDir, err))
continue
}
skill.Dir = skillDir
skill.File = filepath.Join(skillDir, SkillFile)
if previous, ok := seen[skill.Name]; ok {
catalog.Warnings = append(catalog.Warnings, fmt.Sprintf("%s: duplicate skill name %q already loaded from %s", skillDir, skill.Name, previous))
continue
}
seen[skill.Name] = skillDir
catalog.Skills = append(catalog.Skills, skill)
}
slices.SortFunc(catalog.Skills, func(a, b Skill) int {
return strings.Compare(a.Name, b.Name)
})
return catalog, nil
}
func ReadMetadata(path string) (Skill, error) {
data, err := readSkillFile(path)
if err != nil {
return Skill{}, err
}
meta, _, err := parseSkillFile(data)
if err != nil {
return Skill{}, err
}
if err := validateMetadata(meta); err != nil {
return Skill{}, err
}
return Skill{Name: meta.Name, Description: meta.Description}, nil
}
func (c *Catalog) Empty() bool {
return c == nil || len(c.Skills) == 0
}
func (c *Catalog) Find(name string) (Skill, bool) {
if c == nil {
return Skill{}, false
}
name = NormalizeName(name)
for _, skill := range c.Skills {
if skill.Name == name {
return skill, true
}
}
return Skill{}, false
}
func (c *Catalog) SummaryMarkdown() string {
if c.Empty() {
return "No skills are installed."
}
var b strings.Builder
b.WriteString("Installed skills:\n\n")
for _, skill := range c.Skills {
b.WriteString("- **")
b.WriteString(skill.Name)
b.WriteString("**: ")
b.WriteString(skill.Description)
b.WriteByte('\n')
}
return strings.TrimRight(b.String(), "\n")
}
func (c *Catalog) SystemPrompt(toolAvailable bool) string {
if c.Empty() {
return ""
}
var b strings.Builder
b.WriteString("Agent skills are available. Skills are reusable instruction packages stored under ")
b.WriteString(c.Dir)
b.WriteString(".\n")
b.WriteString("Use a skill when its description matches the user's task. Load only metadata up front; load full instructions only when needed.\n")
if toolAvailable {
b.WriteString("To load a skill, call the skill tool with the skill name. After loading SKILL.md, follow it. Resolve relative references from the returned skill directory.\n")
} else {
b.WriteString("This model cannot call tools in this session. Follow any skill instructions that are explicitly provided by the user or system.\n")
}
b.WriteString("\nAvailable skills:\n")
for _, skill := range c.Skills {
b.WriteString("- ")
b.WriteString(skill.Name)
b.WriteString(": ")
b.WriteString(skill.Description)
b.WriteByte('\n')
}
return strings.TrimRight(b.String(), "\n")
}
func (s Skill) Read() (string, error) {
if s.File == "" {
return "", fmt.Errorf("skill %q has no %s path", s.Name, SkillFile)
}
data, err := readSkillFile(s.File)
if err != nil {
return "", err
}
return string(data), nil
}
func readSkillFile(path string) ([]byte, error) {
info, err := os.Lstat(path)
if err != nil {
return nil, err
}
if info.Mode()&os.ModeSymlink != 0 {
return nil, fmt.Errorf("%s must not be a symlink", SkillFile)
}
if !info.Mode().IsRegular() {
return nil, fmt.Errorf("%s must be a regular file", SkillFile)
}
if info.Size() > maxSkillFileBytes {
return nil, fmt.Errorf("%s exceeds %d bytes", SkillFile, maxSkillFileBytes)
}
return os.ReadFile(path)
}
func NormalizeName(name string) string {
name = strings.TrimSpace(name)
name = strings.TrimPrefix(name, "/")
return strings.ToLower(name)
}
func parseSkillFile(data []byte) (frontmatter, string, error) {
text := strings.ReplaceAll(string(data), "\r\n", "\n")
if !strings.HasPrefix(text, "---\n") {
return frontmatter{}, "", fmt.Errorf("%s must start with YAML frontmatter", SkillFile)
}
rest := text[len("---\n"):]
end := strings.Index(rest, "\n---")
if end < 0 {
return frontmatter{}, "", fmt.Errorf("%s frontmatter is not closed", SkillFile)
}
var meta frontmatter
if err := yaml.Unmarshal([]byte(rest[:end]), &meta); err != nil {
return frontmatter{}, "", fmt.Errorf("parse frontmatter: %w", err)
}
body := rest[end+len("\n---"):]
body = strings.TrimPrefix(body, "\n")
return meta, body, nil
}
func validateMetadata(meta frontmatter) error {
if !validName.MatchString(meta.Name) {
return fmt.Errorf("invalid skill name %q", meta.Name)
}
if strings.TrimSpace(meta.Description) == "" {
return fmt.Errorf("skill %q has empty description", meta.Name)
}
if len([]rune(meta.Description)) > 1024 {
return fmt.Errorf("skill %q description exceeds 1024 characters", meta.Name)
}
return nil
}
type copyDirResult struct {
Skipped []string
}
func copyDir(src, dst string, force bool) (copyDirResult, error) {
if _, err := os.Stat(dst); err == nil && !force {
return copyDirResult{}, fs.ErrExist
} else if err != nil && !errors.Is(err, os.ErrNotExist) {
return copyDirResult{}, err
}
parent := filepath.Dir(dst)
if err := os.MkdirAll(parent, 0o755); err != nil {
return copyDirResult{}, err
}
tmp, err := os.MkdirTemp(parent, "."+filepath.Base(dst)+".tmp-*")
if err != nil {
return copyDirResult{}, err
}
moved := false
defer func() {
if !moved {
_ = os.RemoveAll(tmp)
}
}()
var result copyDirResult
if err := filepath.WalkDir(src, func(path string, d fs.DirEntry, err error) error {
if err != nil {
return err
}
rel, err := filepath.Rel(src, path)
if err != nil {
return err
}
if rel == "." {
return nil
}
target := filepath.Join(tmp, rel)
if d.Type()&os.ModeSymlink != 0 {
result.Skipped = append(result.Skipped, rel)
return nil
}
if d.IsDir() {
info, err := d.Info()
if err != nil {
return err
}
return os.MkdirAll(target, info.Mode().Perm())
}
info, err := d.Info()
if err != nil {
return err
}
data, err := os.ReadFile(path)
if err != nil {
return err
}
return os.WriteFile(target, data, info.Mode().Perm())
}); err != nil {
return result, err
}
if force {
if err := os.RemoveAll(dst); err != nil {
return result, err
}
}
if _, err := os.Stat(dst); err == nil {
return result, fs.ErrExist
} else if err != nil && !errors.Is(err, os.ErrNotExist) {
return result, err
}
if err := os.Rename(tmp, dst); err != nil {
return result, err
}
moved = true
return result, nil
}
+157
View File
@@ -0,0 +1,157 @@
package skills
import (
"os"
"path/filepath"
"runtime"
"strings"
"testing"
)
func TestLoadCatalogReadsSkillMetadata(t *testing.T) {
dir := t.TempDir()
writeSkill(t, filepath.Join(dir, "go-code"), "go-code", "Write idiomatic Go code.")
catalog, err := Load(dir)
if err != nil {
t.Fatal(err)
}
if len(catalog.Skills) != 1 {
t.Fatalf("skills = %d, want 1: %#v", len(catalog.Skills), catalog)
}
if got := catalog.Skills[0].Name; got != "go-code" {
t.Fatalf("skill name = %q", got)
}
if prompt := catalog.SystemPrompt(true); !strings.Contains(prompt, "go-code: Write idiomatic Go code.") || !strings.Contains(prompt, "call the skill tool") {
t.Fatalf("system prompt missing skill metadata: %q", prompt)
}
}
func TestLoadCatalogSkipsInvalidSkills(t *testing.T) {
dir := t.TempDir()
writeSkill(t, filepath.Join(dir, "bad"), "Bad_Name", "bad")
catalog, err := Load(dir)
if err != nil {
t.Fatal(err)
}
if len(catalog.Skills) != 0 {
t.Fatalf("skills = %#v, want none", catalog.Skills)
}
if len(catalog.Warnings) == 0 {
t.Fatal("expected invalid skill warning")
}
}
func TestImportToDirCopiesCanonicalSkill(t *testing.T) {
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("USERPROFILE", home)
src := filepath.Join(home, ".claude", "skills", "go-code")
writeSkill(t, src, "go-code", "Write idiomatic Go code.")
if err := os.WriteFile(filepath.Join(src, "notes.md"), []byte("notes"), 0o644); err != nil {
t.Fatal(err)
}
dest := filepath.Join(home, ".ollama", "skills")
results, err := ImportToDir("claude", dest, false)
if err != nil {
t.Fatal(err)
}
if len(results) != 1 || results[0].Skipped {
t.Fatalf("results = %#v", results)
}
if _, err := os.Stat(filepath.Join(dest, "go-code", "notes.md")); err != nil {
t.Fatal(err)
}
}
func TestReadMetadataRejectsSymlink(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("symlink permissions vary on Windows")
}
dir := t.TempDir()
real := filepath.Join(dir, "real.md")
if err := os.WriteFile(real, []byte("---\nname: go-code\ndescription: Write idiomatic Go code.\n---\n"), 0o644); err != nil {
t.Fatal(err)
}
link := filepath.Join(dir, SkillFile)
if err := os.Symlink(real, link); err != nil {
t.Fatal(err)
}
if _, err := ReadMetadata(link); err == nil || !strings.Contains(err.Error(), "must not be a symlink") {
t.Fatalf("ReadMetadata error = %v, want symlink rejection", err)
}
}
func TestImportToDirReportsSymlinkedSkillDirectory(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("symlink permissions vary on Windows")
}
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("USERPROFILE", home)
srcRoot := filepath.Join(home, ".claude", "skills")
real := filepath.Join(home, "elsewhere", "go-code")
writeSkill(t, real, "go-code", "Write idiomatic Go code.")
if err := os.MkdirAll(srcRoot, 0o755); err != nil {
t.Fatal(err)
}
if err := os.Symlink(real, filepath.Join(srcRoot, "go-code")); err != nil {
t.Fatal(err)
}
results, err := ImportToDir("claude", filepath.Join(home, ".ollama", "skills"), false)
if err != nil {
t.Fatal(err)
}
if len(results) != 1 || !results[0].Skipped {
t.Fatalf("results = %#v, want one skipped symlink directory", results)
}
if !strings.Contains(results[0].Error, "symlinked skill directories") {
t.Fatalf("error = %q, want symlink directory warning", results[0].Error)
}
}
func TestImportToDirReportsSkippedSymlinkEntries(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("symlink permissions vary on Windows")
}
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("USERPROFILE", home)
src := filepath.Join(home, ".claude", "skills", "go-code")
writeSkill(t, src, "go-code", "Write idiomatic Go code.")
target := filepath.Join(home, "outside.md")
if err := os.WriteFile(target, []byte("outside"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.Symlink(target, filepath.Join(src, "outside.md")); err != nil {
t.Fatal(err)
}
results, err := ImportToDir("claude", filepath.Join(home, ".ollama", "skills"), false)
if err != nil {
t.Fatal(err)
}
if len(results) != 1 || results[0].Skipped {
t.Fatalf("results = %#v, want one imported skill", results)
}
if !strings.Contains(results[0].Error, "skipped symlinks: outside.md") {
t.Fatalf("error = %q, want skipped symlink warning", results[0].Error)
}
if _, err := os.Lstat(filepath.Join(home, ".ollama", "skills", "go-code", "outside.md")); !os.IsNotExist(err) {
t.Fatalf("copied symlink err = %v, want missing symlink", err)
}
}
func writeSkill(t *testing.T, dir, name, description string) {
t.Helper()
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
content := "---\nname: " + name + "\ndescription: " + description + "\n---\n\n# " + name + "\n\nUse this skill.\n"
if err := os.WriteFile(filepath.Join(dir, SkillFile), []byte(content), 0o644); err != nil {
t.Fatal(err)
}
}
+991
View File
@@ -0,0 +1,991 @@
package store
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"runtime"
"strings"
"sync"
"time"
"github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
_ "github.com/mattn/go-sqlite3"
)
type Store struct {
DBPath string
dbMu sync.Mutex
db *database
}
type database struct {
conn *sql.DB
}
type AgentChat struct {
ID string
Title string
Model string
CreatedAt time.Time
Messages []api.Message
}
type ChatSummary struct {
ID string
Title string
Model string
CreatedAt time.Time
UpdatedAt time.Time
MessageCount int
ApproxBytes int64
}
func New(path string) (*Store, error) {
store := &Store{DBPath: path}
if err := store.ensureDB(); err != nil {
return nil, err
}
return store, nil
}
func (s *Store) ensureDB() error {
if s.db != nil {
return nil
}
s.dbMu.Lock()
defer s.dbMu.Unlock()
if s.db != nil {
return nil
}
dbPath := s.DBPath
if dbPath == "" {
dbPath = defaultDBPath()
}
if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil {
return fmt.Errorf("create database directory: %w", err)
}
db, err := newDatabase(dbPath)
if err != nil {
return err
}
s.db = db
return nil
}
func (s *Store) Close() error {
if s == nil || s.db == nil {
return nil
}
err := s.db.Close()
s.db = nil
return err
}
func defaultDBPath() string {
switch runtime.GOOS {
case "windows":
return filepath.Join(os.Getenv("LOCALAPPDATA"), "Ollama", "db.sqlite")
case "darwin":
return filepath.Join(os.Getenv("HOME"), "Library", "Application Support", "Ollama", "db.sqlite")
default:
return filepath.Join(os.Getenv("HOME"), ".ollama", "db.sqlite")
}
}
func newDatabase(dbPath string) (*database, error) {
conn, err := sql.Open("sqlite3", dbPath+"?_foreign_keys=on&_journal_mode=WAL&_busy_timeout=5000&_txlock=immediate")
if err != nil {
return nil, fmt.Errorf("open database: %w", err)
}
if err := conn.Ping(); err != nil {
conn.Close()
return nil, fmt.Errorf("ping database: %w", err)
}
db := &database{conn: conn}
if err := db.init(); err != nil {
conn.Close()
return nil, fmt.Errorf("initialize database: %w", err)
}
return db, nil
}
func (db *database) Close() error {
_, _ = db.conn.Exec("PRAGMA wal_checkpoint(TRUNCATE);")
return db.conn.Close()
}
func (db *database) init() error {
if _, err := db.conn.Exec("PRAGMA foreign_keys = ON"); err != nil {
return fmt.Errorf("enable foreign keys: %w", err)
}
if _, err := db.conn.Exec(`
CREATE TABLE IF NOT EXISTS chats (
id TEXT PRIMARY KEY,
title TEXT NOT NULL DEFAULT '',
model_name TEXT NOT NULL DEFAULT '',
source TEXT NOT NULL DEFAULT 'app',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
browser_state TEXT
);
CREATE TABLE IF NOT EXISTS messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
chat_id TEXT NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL DEFAULT '',
thinking TEXT NOT NULL DEFAULT '',
images TEXT NOT NULL DEFAULT '[]',
stream BOOLEAN NOT NULL DEFAULT 0,
model_name TEXT,
model_cloud BOOLEAN,
model_ollama_host BOOLEAN,
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
thinking_time_start TIMESTAMP,
thinking_time_end TIMESTAMP,
tool_result TEXT,
tool_name TEXT NOT NULL DEFAULT '',
tool_call_id TEXT NOT NULL DEFAULT '',
archived BOOLEAN NOT NULL DEFAULT 0,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS tool_calls (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
type TEXT NOT NULL,
tool_call_id TEXT NOT NULL DEFAULT '',
function_name TEXT NOT NULL,
function_arguments TEXT NOT NULL,
function_result TEXT,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS compactions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
chat_id TEXT NOT NULL,
summary TEXT NOT NULL,
archived_message_ids TEXT NOT NULL DEFAULT '[]',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
`); err != nil {
return err
}
return db.ensureAgentSchema()
}
func (db *database) ensureAgentSchema() error {
for _, stmt := range []struct {
sql string
msg string
}{
{`ALTER TABLE chats ADD COLUMN model_name TEXT NOT NULL DEFAULT ''`, "add chats.model_name"},
{`ALTER TABLE chats ADD COLUMN source TEXT NOT NULL DEFAULT 'app'`, "add chats.source"},
{`ALTER TABLE messages ADD COLUMN images TEXT NOT NULL DEFAULT '[]'`, "add messages.images"},
{`ALTER TABLE messages ADD COLUMN tool_name TEXT NOT NULL DEFAULT ''`, "add messages.tool_name"},
{`ALTER TABLE messages ADD COLUMN tool_call_id TEXT NOT NULL DEFAULT ''`, "add messages.tool_call_id"},
{`ALTER TABLE messages ADD COLUMN archived BOOLEAN NOT NULL DEFAULT 0`, "add messages.archived"},
{`ALTER TABLE tool_calls ADD COLUMN tool_call_id TEXT NOT NULL DEFAULT ''`, "add tool_calls.tool_call_id"},
} {
_, err := db.conn.Exec(stmt.sql)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("%s: %w", stmt.msg, err)
}
}
_, err := db.conn.Exec(`
CREATE INDEX IF NOT EXISTS idx_messages_chat_id ON messages(chat_id);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_id ON messages(chat_id, id);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_archived ON messages(chat_id, archived, id);
CREATE INDEX IF NOT EXISTS idx_tool_calls_message_id ON tool_calls(message_id);
CREATE TABLE IF NOT EXISTS compactions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
chat_id TEXT NOT NULL,
summary TEXT NOT NULL,
archived_message_ids TEXT NOT NULL DEFAULT '[]',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_compactions_chat_id ON compactions(chat_id, id);
`)
if err != nil {
return fmt.Errorf("create agent chat persistence tables: %w", err)
}
return nil
}
func duplicateColumnError(err error) bool {
return err != nil && strings.Contains(strings.ToLower(err.Error()), "duplicate column")
}
func (s *Store) EnsureChat(ctx context.Context, id string, title string) error {
if id == "" {
return fmt.Errorf("chat id is required")
}
if err := s.ensureDB(); err != nil {
return err
}
_, err := s.db.conn.ExecContext(ctx, `
INSERT INTO chats (id, title, created_at, source)
VALUES (?, ?, ?, 'agent')
ON CONFLICT(id) DO UPDATE SET
title = CASE
WHEN excluded.title != '' THEN excluded.title
ELSE chats.title
END,
source = 'agent'
`, id, title, time.Now())
if err != nil {
return fmt.Errorf("ensure chat: %w", err)
}
return nil
}
func (s *Store) SetChatModel(ctx context.Context, chatID string, model string) error {
chatID = strings.TrimSpace(chatID)
model = strings.TrimSpace(model)
if chatID == "" {
return fmt.Errorf("chat id is required")
}
if model == "" {
return fmt.Errorf("model is required")
}
if err := s.EnsureChat(ctx, chatID, ""); err != nil {
return err
}
if _, err := s.db.conn.ExecContext(ctx, `UPDATE chats SET model_name = ? WHERE id = ?`, model, chatID); err != nil {
return fmt.Errorf("set chat model: %w", err)
}
return nil
}
func (s *Store) AppendAgentMessage(ctx context.Context, chatID string, msg api.Message, model string) error {
if err := s.EnsureChat(ctx, chatID, ""); err != nil {
return err
}
tx, err := s.db.conn.BeginTx(ctx, nil)
if err != nil {
return fmt.Errorf("begin transaction: %w", err)
}
defer tx.Rollback()
messageID, err := insertAgentMessage(ctx, tx, chatID, msg, model)
if err != nil {
return err
}
for _, toolCall := range msg.ToolCalls {
if err := insertAgentToolCall(ctx, tx, messageID, toolCall); err != nil {
return err
}
}
if msg.Role == "user" && strings.TrimSpace(msg.Content) != "" {
if err := maybeSetAgentTitle(ctx, tx, chatID, msg.Content); err != nil {
return err
}
}
return tx.Commit()
}
func (s *Store) UpdateLastAgentMessage(ctx context.Context, chatID string, msg api.Message, model string) error {
if err := s.ensureDB(); err != nil {
return err
}
tx, err := s.db.conn.BeginTx(ctx, nil)
if err != nil {
return fmt.Errorf("begin transaction: %w", err)
}
defer tx.Rollback()
var messageID int64
if err := tx.QueryRowContext(ctx, `SELECT COALESCE(MAX(id), 0) FROM messages WHERE chat_id = ? AND archived = 0`, chatID).Scan(&messageID); err != nil {
return fmt.Errorf("get last message id: %w", err)
}
if messageID == 0 {
return fmt.Errorf("no message found to update")
}
now := time.Now()
modelName := sql.NullString{}
if model != "" {
modelName = sql.NullString{String: model, Valid: true}
}
imagesJSON, err := marshalAgentMessageImages(msg.Images)
if err != nil {
return err
}
_, err = tx.ExecContext(ctx, `
UPDATE messages
SET role = ?, content = ?, thinking = ?, images = ?, tool_name = ?, tool_call_id = ?, model_name = ?, updated_at = ?
WHERE id = ?
`, msg.Role, msg.Content, msg.Thinking, imagesJSON, msg.ToolName, msg.ToolCallID, modelName, now, messageID)
if err != nil {
return fmt.Errorf("update last message: %w", err)
}
if _, err := tx.ExecContext(ctx, `DELETE FROM tool_calls WHERE message_id = ?`, messageID); err != nil {
return fmt.Errorf("delete old tool calls: %w", err)
}
for _, toolCall := range msg.ToolCalls {
if err := insertAgentToolCall(ctx, tx, messageID, toolCall); err != nil {
return err
}
}
return tx.Commit()
}
func (s *Store) AgentChat(ctx context.Context, id string) (*AgentChat, error) {
if err := s.ensureDB(); err != nil {
return nil, err
}
var chat AgentChat
var chatModel string
if err := s.db.conn.QueryRowContext(ctx, `
SELECT id, title, model_name, created_at FROM chats WHERE id = ?
`, id).Scan(&chat.ID, &chat.Title, &chatModel, &chat.CreatedAt); err != nil {
return nil, err
}
if strings.TrimSpace(chatModel) != "" {
chat.Model = chatModel
} else {
model, err := latestAgentModelForChat(ctx, s.db.conn, id)
if err != nil {
return nil, err
}
chat.Model = model
}
rows, err := s.db.conn.QueryContext(ctx, `
SELECT id, role, content, thinking, images, tool_name, tool_call_id FROM messages WHERE chat_id = ? AND archived = 0 ORDER BY id ASC
`, id)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var messageID int64
var msg api.Message
var imagesJSON string
if err := rows.Scan(&messageID, &msg.Role, &msg.Content, &msg.Thinking, &imagesJSON, &msg.ToolName, &msg.ToolCallID); err != nil {
return nil, err
}
images, err := unmarshalAgentMessageImages(imagesJSON)
if err != nil {
return nil, err
}
msg.Images = images
toolCalls, err := getAgentToolCalls(ctx, s.db.conn, messageID)
if err != nil {
return nil, err
}
msg.ToolCalls = toolCalls
chat.Messages = append(chat.Messages, msg)
}
if err := rows.Err(); err != nil {
return nil, err
}
summary, err := latestCompactionSummary(ctx, s.db.conn, id)
if err != nil {
return nil, err
}
if summary != "" && !messagesContainCompactionSummary(chat.Messages) {
chat.Messages = insertCompactionSummaryAfterLeadingSystemMessages(chat.Messages, agent.CompactionSummaryMessages(summary, false))
} else {
chat.Messages = moveCompactionSummaryBeforeKeptMessages(chat.Messages)
}
chat.Messages = repairDanglingToolCalls(chat.Messages)
return &chat, nil
}
func (s *Store) LatestChat(ctx context.Context) (*AgentChat, error) {
if err := s.ensureDB(); err != nil {
return nil, err
}
var chatID string
query := fmt.Sprintf(`
SELECT c.id
FROM chats c
JOIN messages m ON m.chat_id = c.id AND m.archived = 0
WHERE c.source = 'agent'
GROUP BY c.id
HAVING %[1]s IS NOT NULL
ORDER BY MAX(m.updated_at) DESC, MAX(m.id) DESC
LIMIT 1
`, currentAgentModelSelectExpr("c"))
if err := s.db.conn.QueryRowContext(ctx, query).Scan(&chatID); err != nil {
return nil, err
}
return s.AgentChat(ctx, chatID)
}
func (s *Store) LatestChatForModel(ctx context.Context, model string) (*AgentChat, error) {
if strings.TrimSpace(model) == "" {
return nil, fmt.Errorf("model is required")
}
if err := s.ensureDB(); err != nil {
return nil, err
}
var chatID string
query := fmt.Sprintf(`
SELECT c.id
FROM chats c
JOIN messages m ON m.chat_id = c.id AND m.archived = 0
WHERE c.source = 'agent'
GROUP BY c.id
HAVING %[1]s = ?
ORDER BY MAX(m.updated_at) DESC, MAX(m.id) DESC
LIMIT 1
`, currentAgentModelSelectExpr("c"))
if err := s.db.conn.QueryRowContext(ctx, query, model).Scan(&chatID); err != nil {
return nil, err
}
return s.AgentChat(ctx, chatID)
}
func (s *Store) ListChats(ctx context.Context, limit int) ([]ChatSummary, error) {
if err := s.ensureDB(); err != nil {
return nil, err
}
if limit <= 0 {
limit = 50
}
query := fmt.Sprintf(`
SELECT
c.id,
c.title,
c.created_at,
MAX(m.updated_at) AS updated_at,
COUNT(m.id) AS message_count,
COALESCE(SUM(
LENGTH(m.role) +
LENGTH(m.content) +
LENGTH(m.thinking) +
LENGTH(m.tool_name) +
LENGTH(m.tool_call_id)
), 0) AS approx_bytes,
%[1]s AS current_model
FROM chats c
JOIN messages m ON m.chat_id = c.id AND m.archived = 0
WHERE c.source = 'agent'
GROUP BY c.id
ORDER BY updated_at DESC, MAX(m.id) DESC
LIMIT ?
`, currentAgentModelSelectExpr("c"))
rows, err := s.db.conn.QueryContext(ctx, query, limit)
if err != nil {
return nil, fmt.Errorf("list chats: %w", err)
}
defer rows.Close()
var summaries []ChatSummary
for rows.Next() {
var summary ChatSummary
var updatedAt string
var modelName sql.NullString
if err := rows.Scan(&summary.ID, &summary.Title, &summary.CreatedAt, &updatedAt, &summary.MessageCount, &summary.ApproxBytes, &modelName); err != nil {
return nil, fmt.Errorf("scan chat summary: %w", err)
}
if modelName.Valid {
summary.Model = modelName.String
}
summary.UpdatedAt, err = parseAgentSQLiteTime(updatedAt)
if err != nil {
return nil, fmt.Errorf("parse chat updated_at: %w", err)
}
summaries = append(summaries, summary)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("read chat summaries: %w", err)
}
return summaries, nil
}
func (s *Store) ListUserMessages(ctx context.Context, limit int) ([]string, error) {
if err := s.ensureDB(); err != nil {
return nil, err
}
if limit <= 0 {
limit = 50
}
rows, err := s.db.conn.QueryContext(ctx, `
SELECT content
FROM (
SELECT id, content
FROM messages
WHERE role = 'user'
AND archived = 0
AND TRIM(content) != ''
AND content NOT LIKE ?
ORDER BY id DESC
LIMIT ?
)
ORDER BY id ASC
`, agent.CompactionSummaryMessagePrefix+"%", limit)
if err != nil {
return nil, fmt.Errorf("list user messages: %w", err)
}
defer rows.Close()
var messages []string
for rows.Next() {
var content string
if err := rows.Scan(&content); err != nil {
return nil, fmt.Errorf("scan user message: %w", err)
}
messages = append(messages, content)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("read user messages: %w", err)
}
return messages, nil
}
func (s *Store) ArchiveForCompaction(ctx context.Context, chatID string, keepUserTurns int, summary string, continueTask bool) error {
return s.archiveForCompaction(ctx, chatID, keepUserTurns, summary, continueTask)
}
func (s *Store) archiveForCompaction(ctx context.Context, chatID string, keepUserTurns int, summary string, continueTask bool) error {
if err := s.ensureDB(); err != nil {
return err
}
if chatID == "" {
return fmt.Errorf("chat id is required")
}
if keepUserTurns < 0 {
return fmt.Errorf("keep user turns must be non-negative")
}
if strings.TrimSpace(summary) == "" {
return fmt.Errorf("summary is required")
}
tx, err := s.db.conn.BeginTx(ctx, nil)
if err != nil {
return fmt.Errorf("begin transaction: %w", err)
}
defer tx.Rollback()
var keepStartID int64
if keepUserTurns == 0 {
if err := tx.QueryRowContext(ctx, `
SELECT COALESCE(MAX(id) + 1, 0)
FROM messages
WHERE chat_id = ? AND archived = 0
`, chatID).Scan(&keepStartID); err != nil {
return fmt.Errorf("find compaction boundary: %w", err)
}
if keepStartID == 0 {
return nil
}
} else {
if err := tx.QueryRowContext(ctx, `
SELECT id
FROM messages
WHERE chat_id = ? AND archived = 0 AND role = 'user'
ORDER BY id DESC
LIMIT 1 OFFSET ?
`, chatID, keepUserTurns-1).Scan(&keepStartID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil
}
return fmt.Errorf("find compaction boundary: %w", err)
}
}
rows, err := tx.QueryContext(ctx, `
SELECT m.id
FROM messages m
WHERE m.chat_id = ? AND m.archived = 0 AND (
m.id < ?
OR m.tool_name = ?
OR EXISTS (
SELECT 1 FROM tool_calls tc
WHERE tc.message_id = m.id AND tc.function_name = ?
)
)
ORDER BY id ASC
`, chatID, keepStartID, agent.CompactionToolName, agent.CompactionToolName)
if err != nil {
return fmt.Errorf("list archived messages: %w", err)
}
var archivedIDs []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
rows.Close()
return fmt.Errorf("scan archived message id: %w", err)
}
archivedIDs = append(archivedIDs, id)
}
if err := rows.Err(); err != nil {
rows.Close()
return fmt.Errorf("read archived message ids: %w", err)
}
rows.Close()
if len(archivedIDs) == 0 {
return nil
}
idsJSON, err := json.Marshal(archivedIDs)
if err != nil {
return fmt.Errorf("marshal archived message ids: %w", err)
}
if _, err := tx.ExecContext(ctx, `
INSERT INTO compactions (chat_id, summary, archived_message_ids, created_at)
VALUES (?, ?, ?, ?)
`, chatID, summary, string(idsJSON), time.Now()); err != nil {
return fmt.Errorf("insert compaction: %w", err)
}
if _, err := tx.ExecContext(ctx, `
UPDATE messages
SET archived = 1
WHERE chat_id = ? AND archived = 0 AND (
id < ?
OR tool_name = ?
OR EXISTS (
SELECT 1 FROM tool_calls
WHERE tool_calls.message_id = messages.id AND tool_calls.function_name = ?
)
)
`, chatID, keepStartID, agent.CompactionToolName, agent.CompactionToolName); err != nil {
return fmt.Errorf("archive messages: %w", err)
}
for _, msg := range agent.CompactionSummaryMessages(summary, continueTask) {
messageID, err := insertAgentMessage(ctx, tx, chatID, msg, "")
if err != nil {
return err
}
for _, toolCall := range msg.ToolCalls {
if err := insertAgentToolCall(ctx, tx, messageID, toolCall); err != nil {
return err
}
}
}
return tx.Commit()
}
func parseAgentSQLiteTime(value string) (time.Time, error) {
value = strings.TrimSpace(value)
if value == "" {
return time.Time{}, nil
}
for _, layout := range []string{
time.RFC3339Nano,
"2006-01-02 15:04:05.999999999-07:00",
"2006-01-02 15:04:05.999999999Z07:00",
"2006-01-02 15:04:05.999999999",
"2006-01-02 15:04:05-07:00",
"2006-01-02 15:04:05Z07:00",
"2006-01-02 15:04:05",
} {
t, err := time.Parse(layout, value)
if err == nil {
return t, nil
}
}
return time.Time{}, fmt.Errorf("unsupported time format %q", value)
}
func latestAgentModelForChat(ctx context.Context, db *sql.DB, chatID string) (string, error) {
var modelName string
if err := db.QueryRowContext(ctx, `
SELECT model_name
FROM messages
WHERE chat_id = ? AND archived = 0 AND model_name IS NOT NULL AND model_name != ''
ORDER BY updated_at DESC, id DESC
LIMIT 1
`, chatID).Scan(&modelName); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return "", nil
}
return "", err
}
return modelName, nil
}
func currentAgentModelSelectExpr(chatAlias string) string {
return fmt.Sprintf(`COALESCE(
NULLIF(%[1]s.model_name, ''),
(
SELECT lm.model_name
FROM messages lm
WHERE lm.chat_id = %[1]s.id AND lm.archived = 0 AND lm.model_name IS NOT NULL AND lm.model_name != ''
ORDER BY lm.updated_at DESC, lm.id DESC
LIMIT 1
)
)`, chatAlias)
}
func messagesContainCompactionSummary(messages []api.Message) bool {
for _, msg := range messages {
if agent.IsCompactionSummary(msg) {
return true
}
}
return false
}
func moveCompactionSummaryBeforeKeptMessages(messages []api.Message) []api.Message {
start := -1
end := -1
for i, msg := range messages {
if agent.IsCompactionToolCall(msg) {
start = i
end = i + 1
if end < len(messages) && agent.IsCompactionToolResult(messages[end]) {
end++
}
}
if start >= 0 {
break
}
}
if start <= 0 || end <= start {
return messages
}
insertAt := leadingSystemMessageCount(messages[:start])
reordered := make([]api.Message, 0, len(messages))
reordered = append(reordered, messages[:insertAt]...)
reordered = append(reordered, messages[start:end]...)
reordered = append(reordered, messages[insertAt:start]...)
reordered = append(reordered, messages[end:]...)
return reordered
}
func insertCompactionSummaryAfterLeadingSystemMessages(messages, summary []api.Message) []api.Message {
insertAt := leadingSystemMessageCount(messages)
reordered := make([]api.Message, 0, len(messages)+len(summary))
reordered = append(reordered, messages[:insertAt]...)
reordered = append(reordered, summary...)
reordered = append(reordered, messages[insertAt:]...)
return reordered
}
func leadingSystemMessageCount(messages []api.Message) int {
for i, msg := range messages {
if msg.Role != "system" {
return i
}
}
return len(messages)
}
type pendingToolCall struct {
key string
call api.ToolCall
}
func repairDanglingToolCalls(messages []api.Message) []api.Message {
var pending []pendingToolCall
pendingByKey := map[string]struct{}{}
repaired := make([]api.Message, 0, len(messages))
flushPending := func() {
for _, pendingCall := range pending {
if _, ok := pendingByKey[pendingCall.key]; !ok {
continue
}
repaired = append(repaired, api.Message{
Role: "tool",
Content: "Tool execution interrupted before a result was recorded.",
ToolName: pendingCall.call.Function.Name,
ToolCallID: pendingCall.call.ID,
})
}
pending = nil
pendingByKey = map[string]struct{}{}
}
for _, msg := range messages {
if len(pendingByKey) > 0 && msg.Role != "tool" {
flushPending()
}
repaired = append(repaired, msg)
switch msg.Role {
case "assistant":
for _, call := range msg.ToolCalls {
key := agentToolCallKey(call, len(pending))
pending = append(pending, pendingToolCall{key: key, call: call})
pendingByKey[key] = struct{}{}
}
case "tool":
if key := msg.ToolCallID; key != "" {
delete(pendingByKey, key)
} else if msg.ToolName != "" {
for _, pendingCall := range pending {
if pendingCall.call.ID == "" && pendingCall.call.Function.Name == msg.ToolName {
delete(pendingByKey, pendingCall.key)
break
}
}
}
if len(pendingByKey) == 0 {
pending = nil
}
}
}
if len(pendingByKey) > 0 {
flushPending()
}
return repaired
}
func agentToolCallKey(call api.ToolCall, index int) string {
if call.ID != "" {
return call.ID
}
return fmt.Sprintf("#%d:%s", index, call.Function.Name)
}
func insertAgentMessage(ctx context.Context, tx *sql.Tx, chatID string, msg api.Message, model string) (int64, error) {
now := time.Now()
modelName := sql.NullString{}
if model != "" {
modelName = sql.NullString{String: model, Valid: true}
}
imagesJSON, err := marshalAgentMessageImages(msg.Images)
if err != nil {
return 0, err
}
result, err := tx.ExecContext(ctx, `
INSERT INTO messages (chat_id, role, content, thinking, images, tool_name, tool_call_id, model_name, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`, chatID, msg.Role, msg.Content, msg.Thinking, imagesJSON, msg.ToolName, msg.ToolCallID, modelName, now, now)
if err != nil {
return 0, fmt.Errorf("insert message: %w", err)
}
id, err := result.LastInsertId()
if err != nil {
return 0, fmt.Errorf("get message id: %w", err)
}
return id, nil
}
func marshalAgentMessageImages(images []api.ImageData) (string, error) {
if len(images) == 0 {
return "[]", nil
}
data, err := json.Marshal(images)
if err != nil {
return "", fmt.Errorf("marshal message images: %w", err)
}
return string(data), nil
}
func unmarshalAgentMessageImages(value string) ([]api.ImageData, error) {
value = strings.TrimSpace(value)
if value == "" || value == "null" {
return nil, nil
}
var images []api.ImageData
if err := json.Unmarshal([]byte(value), &images); err != nil {
return nil, fmt.Errorf("unmarshal message images: %w", err)
}
return images, nil
}
func insertAgentToolCall(ctx context.Context, tx *sql.Tx, messageID int64, call api.ToolCall) error {
args, err := json.Marshal(call.Function.Arguments)
if err != nil {
return fmt.Errorf("marshal tool arguments: %w", err)
}
_, err = tx.ExecContext(ctx, `
INSERT INTO tool_calls (message_id, type, tool_call_id, function_name, function_arguments)
VALUES (?, ?, ?, ?, ?)
`, messageID, "function", call.ID, call.Function.Name, string(args))
if err != nil {
return fmt.Errorf("insert tool call: %w", err)
}
return nil
}
func getAgentToolCalls(ctx context.Context, db *sql.DB, messageID int64) ([]api.ToolCall, error) {
rows, err := db.QueryContext(ctx, `
SELECT tool_call_id, function_name, function_arguments FROM tool_calls WHERE message_id = ? ORDER BY id ASC
`, messageID)
if err != nil {
return nil, err
}
defer rows.Close()
var calls []api.ToolCall
for rows.Next() {
var id, name, argsJSON string
if err := rows.Scan(&id, &name, &argsJSON); err != nil {
return nil, err
}
var args api.ToolCallFunctionArguments
if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
return nil, err
}
calls = append(calls, api.ToolCall{
ID: id,
Function: api.ToolCallFunction{
Name: name,
Arguments: args,
},
})
}
return calls, rows.Err()
}
func latestCompactionSummary(ctx context.Context, db *sql.DB, chatID string) (string, error) {
var summary string
if err := db.QueryRowContext(ctx, `
SELECT summary
FROM compactions
WHERE chat_id = ?
ORDER BY id DESC
LIMIT 1
`, chatID).Scan(&summary); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return "", nil
}
return "", fmt.Errorf("get latest compaction summary: %w", err)
}
return summary, nil
}
func maybeSetAgentTitle(ctx context.Context, tx *sql.Tx, chatID string, content string) error {
title := strings.TrimSpace(content)
if len([]rune(title)) > 64 {
title = string([]rune(title)[:64])
}
_, err := tx.ExecContext(ctx, `
UPDATE chats
SET title = CASE WHEN title = '' THEN ? ELSE title END
WHERE id = ?
`, title, chatID)
return err
}
+466
View File
@@ -0,0 +1,466 @@
package store
import (
"bytes"
"context"
"database/sql"
"encoding/json"
"errors"
"path/filepath"
"slices"
"strings"
"testing"
"time"
"github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
)
func newTestAgentStore(t *testing.T) *Store {
t.Helper()
t.Setenv("HOME", t.TempDir())
t.Setenv("LOCALAPPDATA", t.TempDir())
store, err := New(filepath.Join(t.TempDir(), "db.sqlite"))
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if err := store.Close(); err != nil {
t.Fatal(err)
}
})
return store
}
func TestAgentStoreWritesSharedChatRows(t *testing.T) {
store := newTestAgentStore(t)
ctx := context.Background()
if err := store.EnsureChat(ctx, "chat-1", ""); err != nil {
t.Fatal(err)
}
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "user", Content: "hello from cli"}, ""); err != nil {
t.Fatal(err)
}
args := api.NewToolCallFunctionArguments()
args.Set("command", "pwd")
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{
Role: "assistant",
Content: "I'll check.",
ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "bash",
Arguments: args,
},
}},
}, "llama3.2"); err != nil {
t.Fatal(err)
}
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{
Role: "tool",
Content: "cwd",
ToolName: "bash",
ToolCallID: "call-1",
}, ""); err != nil {
t.Fatal(err)
}
agentChat, err := store.AgentChat(ctx, "chat-1")
if err != nil {
t.Fatal(err)
}
if len(agentChat.Messages) != 3 {
t.Fatalf("messages = %d, want 3", len(agentChat.Messages))
}
if agentChat.Title != "hello from cli" {
t.Fatalf("title = %q, want %q", agentChat.Title, "hello from cli")
}
if got := agentChat.Messages[1].ToolCalls[0].Function.Name; got != "bash" {
t.Fatalf("tool name = %q, want bash", got)
}
if got := agentChat.Messages[1].ToolCalls[0].ID; got != "call-1" {
t.Fatalf("tool call id = %q, want call-1", got)
}
if agentChat.Messages[2].Role != "tool" || agentChat.Messages[2].ToolCallID != "call-1" {
t.Fatalf("tool result = %#v", agentChat.Messages[2])
}
var source string
if err := store.db.conn.QueryRowContext(ctx, `SELECT source FROM chats WHERE id = ?`, "chat-1").Scan(&source); err != nil {
t.Fatal(err)
}
if source != "agent" {
t.Fatalf("source = %q, want agent", source)
}
}
func TestAgentStoreRepairsDanglingToolCallsOnResume(t *testing.T) {
store := newTestAgentStore(t)
ctx := context.Background()
args := api.NewToolCallFunctionArguments()
args.Set("command", "pwd")
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "user", Content: "start"}, ""); err != nil {
t.Fatal(err)
}
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{
Role: "assistant",
ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "bash",
Arguments: args,
},
}},
}, "llama3.2"); err != nil {
t.Fatal(err)
}
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "user", Content: "after restart"}, ""); err != nil {
t.Fatal(err)
}
agentChat, err := store.AgentChat(ctx, "chat-1")
if err != nil {
t.Fatal(err)
}
if len(agentChat.Messages) != 4 {
t.Fatalf("messages = %#v, want synthetic tool result inserted", agentChat.Messages)
}
repair := agentChat.Messages[2]
if repair.Role != "tool" || repair.ToolName != "bash" || repair.ToolCallID != "call-1" || !strings.Contains(repair.Content, "interrupted") {
t.Fatalf("repair message = %#v", repair)
}
if agentChat.Messages[3].Role != "user" || agentChat.Messages[3].Content != "after restart" {
t.Fatalf("message after repair = %#v", agentChat.Messages[3])
}
}
func TestAgentStoreRoundTripsToolMetadataAndImages(t *testing.T) {
store := newTestAgentStore(t)
ctx := context.Background()
first := api.ImageData([]byte("first image"))
second := api.ImageData([]byte{0, 1, 2, 3})
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{
Role: "tool",
Content: "tool output",
Images: []api.ImageData{first, second},
ToolName: "bash",
ToolCallID: "call-1",
}, ""); err != nil {
t.Fatal(err)
}
chat, err := store.AgentChat(ctx, "chat-1")
if err != nil {
t.Fatal(err)
}
if len(chat.Messages) != 1 {
t.Fatalf("messages = %d, want 1", len(chat.Messages))
}
msg := chat.Messages[0]
if msg.ToolName != "bash" || msg.ToolCallID != "call-1" {
t.Fatalf("tool metadata = %#v", msg)
}
if len(msg.Images) != 2 || !bytes.Equal(msg.Images[0], first) || !bytes.Equal(msg.Images[1], second) {
t.Fatalf("images = %#v, want %#v", msg.Images, []api.ImageData{first, second})
}
updated := api.ImageData([]byte("updated image"))
if err := store.UpdateLastAgentMessage(ctx, "chat-1", api.Message{
Role: "user",
Content: "updated",
Images: []api.ImageData{updated},
}, ""); err != nil {
t.Fatal(err)
}
chat, err = store.AgentChat(ctx, "chat-1")
if err != nil {
t.Fatal(err)
}
if len(chat.Messages) != 1 || len(chat.Messages[0].Images) != 1 || !bytes.Equal(chat.Messages[0].Images[0], updated) {
t.Fatalf("updated images = %#v", chat.Messages)
}
}
func TestAgentStoreLatestAndListChats(t *testing.T) {
store := newTestAgentStore(t)
ctx := context.Background()
if err := store.AppendAgentMessage(ctx, "chat-old", api.Message{Role: "user", Content: "old topic"}, ""); err != nil {
t.Fatal(err)
}
if err := store.AppendAgentMessage(ctx, "chat-old", api.Message{Role: "assistant", Content: "old answer"}, "llama3.2"); err != nil {
t.Fatal(err)
}
if err := store.AppendAgentMessage(ctx, "chat-new", api.Message{Role: "user", Content: "new topic"}, ""); err != nil {
t.Fatal(err)
}
if err := store.AppendAgentMessage(ctx, "chat-new", api.Message{Role: "assistant", Content: "new answer"}, "qwen3"); err != nil {
t.Fatal(err)
}
chat, err := store.LatestChat(ctx)
if err != nil {
t.Fatal(err)
}
if chat.ID != "chat-new" || chat.Model != "qwen3" {
t.Fatalf("latest chat = %#v, want chat-new with qwen3", chat)
}
chat, err = store.LatestChatForModel(ctx, "llama3.2")
if err != nil {
t.Fatal(err)
}
if chat.ID != "chat-old" {
t.Fatalf("llama latest chat = %q, want chat-old", chat.ID)
}
if _, err := store.LatestChatForModel(ctx, "missing"); !errors.Is(err, sql.ErrNoRows) {
t.Fatalf("missing model err = %v, want sql.ErrNoRows", err)
}
summaries, err := store.ListChats(ctx, 10)
if err != nil {
t.Fatal(err)
}
if len(summaries) != 2 {
t.Fatalf("summaries = %d, want 2", len(summaries))
}
if summaries[0].ID != "chat-new" || summaries[0].Title != "new topic" || summaries[0].Model != "qwen3" {
t.Fatalf("newest summary = %#v", summaries[0])
}
if summaries[1].ID != "chat-old" || summaries[1].Model != "llama3.2" {
t.Fatalf("older summary = %#v", summaries[1])
}
future := time.Date(2099, 1, 1, 12, 0, 0, 0, time.UTC)
if _, err := store.db.conn.ExecContext(ctx, `
INSERT INTO chats (id, title, created_at)
VALUES (?, ?, ?)
`, "chat-archived", "archived topic", future); err != nil {
t.Fatal(err)
}
if _, err := store.db.conn.ExecContext(ctx, `
INSERT INTO messages (chat_id, role, content, model_name, created_at, updated_at, archived)
VALUES (?, ?, ?, ?, ?, ?, 1)
`, "chat-archived", "assistant", "archived answer", "ghost-model", future, future); err != nil {
t.Fatal(err)
}
chat, err = store.LatestChat(ctx)
if err != nil {
t.Fatal(err)
}
if chat.ID != "chat-new" {
t.Fatalf("latest chat = %q, want chat-new after archived future row", chat.ID)
}
if _, err := store.LatestChatForModel(ctx, "ghost-model"); !errors.Is(err, sql.ErrNoRows) {
t.Fatalf("archived model err = %v, want sql.ErrNoRows", err)
}
summaries, err = store.ListChats(ctx, 10)
if err != nil {
t.Fatal(err)
}
if len(summaries) != 2 {
t.Fatalf("summaries = %d, want archived-only chat hidden", len(summaries))
}
}
func TestAgentStoreUpdateLastMessageIgnoresArchivedRows(t *testing.T) {
store := newTestAgentStore(t)
ctx := context.Background()
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "assistant", Content: "active"}, "llama3.2"); err != nil {
t.Fatal(err)
}
future := time.Date(2099, 1, 1, 12, 0, 0, 0, time.UTC)
if _, err := store.db.conn.ExecContext(ctx, `
INSERT INTO messages (chat_id, role, content, created_at, updated_at, archived)
VALUES (?, ?, ?, ?, ?, 1)
`, "chat-1", "assistant", "archived", future, future); err != nil {
t.Fatal(err)
}
if err := store.UpdateLastAgentMessage(ctx, "chat-1", api.Message{Role: "assistant", Content: "active updated"}, "llama3.2"); err != nil {
t.Fatal(err)
}
chat, err := store.AgentChat(ctx, "chat-1")
if err != nil {
t.Fatal(err)
}
if len(chat.Messages) != 1 || chat.Messages[0].Content != "active updated" {
t.Fatalf("active messages = %#v, want updated active message only", chat.Messages)
}
var archivedContent string
if err := store.db.conn.QueryRowContext(ctx, `
SELECT content
FROM messages
WHERE chat_id = ? AND archived = 1
ORDER BY id DESC
LIMIT 1
`, "chat-1").Scan(&archivedContent); err != nil {
t.Fatal(err)
}
if archivedContent != "archived" {
t.Fatalf("archived content = %q, want archived", archivedContent)
}
}
func TestAgentStoreListUserMessages(t *testing.T) {
store := newTestAgentStore(t)
ctx := context.Background()
for _, msg := range []api.Message{
{Role: "user", Content: "old prompt"},
{Role: "assistant", Content: "not user"},
{Role: "user", Content: "middle prompt"},
{Role: "user", Content: " "},
{Role: "user", Content: agent.CompactionSummaryMessagePrefix + "old context"},
{Role: "user", Content: "new prompt"},
} {
if err := store.AppendAgentMessage(ctx, "chat-1", msg, ""); err != nil {
t.Fatal(err)
}
}
if _, err := store.db.conn.ExecContext(ctx, `UPDATE messages SET archived = 1 WHERE content = ?`, "middle prompt"); err != nil {
t.Fatal(err)
}
messages, err := store.ListUserMessages(ctx, 2)
if err != nil {
t.Fatal(err)
}
want := []string{"old prompt", "new prompt"}
if !slices.Equal(messages, want) {
t.Fatalf("messages = %#v, want %#v", messages, want)
}
}
func TestAgentStoreRepairsPreAgentSchema(t *testing.T) {
path := filepath.Join(t.TempDir(), "db.sqlite")
db, err := sql.Open("sqlite3", path)
if err != nil {
t.Fatal(err)
}
if _, err := db.Exec(`
CREATE TABLE chats (
id TEXT PRIMARY KEY,
title TEXT NOT NULL DEFAULT '',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
browser_state TEXT
);
CREATE TABLE messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
chat_id TEXT NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL DEFAULT '',
thinking TEXT NOT NULL DEFAULT '',
stream BOOLEAN NOT NULL DEFAULT 0,
model_name TEXT,
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE TABLE tool_calls (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
type TEXT NOT NULL,
function_name TEXT NOT NULL,
function_arguments TEXT NOT NULL,
function_result TEXT,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
);
`); err != nil {
t.Fatal(err)
}
if err := db.Close(); err != nil {
t.Fatal(err)
}
store, err := New(path)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
if err := store.Close(); err != nil {
t.Fatal(err)
}
})
ctx := context.Background()
if err := store.AppendAgentMessage(ctx, "chat-1", api.Message{Role: "user", Content: "hello"}, "llama3.2"); err != nil {
t.Fatal(err)
}
chat, err := store.AgentChat(ctx, "chat-1")
if err != nil {
t.Fatal(err)
}
if chat.Model != "llama3.2" || len(chat.Messages) != 1 || chat.Messages[0].Content != "hello" {
t.Fatalf("chat = %#v", chat)
}
var source string
if err := store.db.conn.QueryRowContext(ctx, `SELECT source FROM chats WHERE id = ?`, "chat-1").Scan(&source); err != nil {
t.Fatal(err)
}
if source != "agent" {
t.Fatalf("source = %q, want agent", source)
}
}
func TestAgentStoreArchivesCompactedMessages(t *testing.T) {
store := newTestAgentStore(t)
ctx := context.Background()
for _, msg := range []api.Message{
{Role: "user", Content: "old request"},
{Role: "assistant", Content: "old answer"},
{Role: "user", Content: "recent request"},
} {
if err := store.AppendAgentMessage(ctx, "chat-1", msg, ""); err != nil {
t.Fatal(err)
}
}
if err := store.ArchiveForCompaction(ctx, "chat-1", 1, "summary", true); err != nil {
t.Fatal(err)
}
agentChat, err := store.AgentChat(ctx, "chat-1")
if err != nil {
t.Fatal(err)
}
if len(agentChat.Messages) != 3 {
t.Fatalf("active messages = %#v, want compaction pair plus latest request", agentChat.Messages)
}
if agentChat.Messages[0].Role != "assistant" || len(agentChat.Messages[0].ToolCalls) != 1 || agentChat.Messages[0].ToolCalls[0].Function.Name != agent.CompactionToolName {
t.Fatalf("summary tool call = %#v", agentChat.Messages[0])
}
content := agentChat.Messages[1].Content
if !strings.Contains(content, agent.CompactionContinueInstruction) {
t.Fatalf("summary tool result missing continuation instruction: %q", content)
}
if agentChat.Messages[2].Content != "recent request" {
t.Fatalf("kept message = %#v, want recent request", agentChat.Messages[2])
}
var idsJSON string
if err := store.db.conn.QueryRowContext(ctx, `
SELECT archived_message_ids FROM compactions WHERE chat_id = ?
`, "chat-1").Scan(&idsJSON); err != nil {
t.Fatal(err)
}
var ids []int64
if err := json.Unmarshal([]byte(idsJSON), &ids); err != nil {
t.Fatal(err)
}
if len(ids) != 2 {
t.Fatalf("archived ids = %v, want 2 ids", ids)
}
}
+91
View File
@@ -0,0 +1,91 @@
package agent
import (
"fmt"
"sort"
"strconv"
"strings"
)
const maxToolInvocationCommandRunes = 100
// ToolDisplayName returns the user-facing label for a tool name.
func ToolDisplayName(name string) string {
switch name {
case "web_search":
return "Web Search"
case "web_fetch":
return "Web Fetch"
case "bash":
return "Bash"
case "powershell":
return "PowerShell"
case "read":
return "Read"
case "list":
return "List"
case "edit":
return "Edit"
case "skill":
return "Skill"
default:
if name == "" {
return "Tool"
}
return name
}
}
// ToolInvocationLabel returns a compact user-facing label for a tool call.
func ToolInvocationLabel(name string, args map[string]any) string {
displayName := ToolDisplayName(name)
for _, key := range []string{"query", "url", "command", "path", "name"} {
if value, ok := displayStringArg(args, key); ok {
if IsShellToolName(name) && key == "command" {
value = truncateDisplayRunes(value, maxToolInvocationCommandRunes)
}
return fmt.Sprintf("%s(%s)", displayName, strconv.Quote(value))
}
}
if len(args) == 0 {
return displayName
}
return fmt.Sprintf("%s(%s)", displayName, formatDisplayArgs(args))
}
// IsShellToolName reports whether name identifies a platform shell tool.
func IsShellToolName(name string) bool {
return name == "bash" || name == "powershell"
}
func displayStringArg(args map[string]any, key string) (string, bool) {
value, ok := args[key].(string)
if !ok || strings.TrimSpace(value) == "" {
return "", false
}
return value, true
}
func truncateDisplayRunes(value string, limit int) string {
runes := []rune(value)
if limit <= 0 || len(runes) <= limit {
return value
}
return string(runes[:limit]) + "..."
}
func formatDisplayArgs(args map[string]any) string {
keys := make([]string, 0, len(args))
for key := range args {
keys = append(keys, key)
}
sort.Strings(keys)
parts := make([]string, 0, len(keys))
for _, key := range keys {
value := fmt.Sprintf("%v", args[key])
value = truncateDisplayRunes(value, 100)
parts = append(parts, fmt.Sprintf("%s=%s", key, strconv.Quote(value)))
}
return strings.Join(parts, ", ")
}
+33
View File
@@ -0,0 +1,33 @@
package agent
import (
"strings"
"testing"
)
func TestToolInvocationLabelTruncatesLongBashCommand(t *testing.T) {
command := strings.Repeat("a", 101)
label := ToolInvocationLabel("bash", map[string]any{"command": command})
want := `Bash("` + strings.Repeat("a", 100) + `...")`
if label != want {
t.Fatalf("label = %q, want %q", label, want)
}
}
func TestToolInvocationLabelCountsRunes(t *testing.T) {
command := strings.Repeat("界", 101)
label := ToolInvocationLabel("bash", map[string]any{"command": command})
want := `Bash("` + strings.Repeat("界", 100) + `...")`
if label != want {
t.Fatalf("label = %q, want %q", label, want)
}
}
func TestToolInvocationLabelTruncatesLongPowerShellCommand(t *testing.T) {
command := strings.Repeat("a", 101)
label := ToolInvocationLabel("powershell", map[string]any{"command": command})
want := `PowerShell("` + strings.Repeat("a", 100) + `...")`
if label != want {
t.Fatalf("label = %q, want %q", label, want)
}
}
+201
View File
@@ -0,0 +1,201 @@
package tools
import (
"context"
"fmt"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"time"
"github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
)
const (
bashTimeout = 3 * time.Minute
maxBashOutputBytes = 60_000
)
type Bash struct{}
func NewBash() *Bash {
return &Bash{}
}
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")
}
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.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: sb.String() + "\n\nError: command timed out after " + bashTimeout.String(), WorkingDir: finalWorkingDir}, nil
}
if ctx.Err() == context.Canceled {
return agent.ToolResult{Content: sb.String() + "\n\nError: command was canceled", WorkingDir: finalWorkingDir}, nil
}
if exitErr, ok := err.(*exec.ExitError); ok {
return agent.ToolResult{Content: sb.String() + fmt.Sprintf("\n\nExit 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 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')
}
func shellQuote(value string) string {
return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'"
}
type boundedOutput struct {
Limit int
buf strings.Builder
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 - b.buf.Len()
if remaining <= 0 {
b.omitted += len(p)
return len(p), nil
}
if len(p) <= remaining {
b.buf.Write(p)
return len(p), nil
}
b.buf.Write(p[:remaining])
b.omitted += len(p) - remaining
return len(p), nil
}
func (b *boundedOutput) Len() int {
return b.buf.Len() + b.omitted
}
func (b *boundedOutput) String(label string) string {
content := b.buf.String()
if b.omitted == 0 {
return content
}
return content + fmt.Sprintf("\n\n[%s truncated: omitted ~%d tokens]", label, approximateTokensFromBytes(b.omitted))
}
func approximateTokensFromBytes(n int) int {
if n <= 0 {
return 0
}
return max(1, (n+3)/4)
}
+128
View File
@@ -0,0 +1,128 @@
package tools
import (
"context"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"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 := NewBash().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 := NewBash().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 TestBashReportsCanceledCommand(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
result, err := NewBash().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 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)
}
}
+44
View File
@@ -0,0 +1,44 @@
//go:build !windows
package tools
import (
"context"
"os/exec"
"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 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
}
+16
View File
@@ -0,0 +1,16 @@
//go:build !windows
package tools
import (
"os/exec"
"testing"
)
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")
}
}
+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",
"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)
}
}
+577
View File
@@ -0,0 +1,577 @@
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 NewRead() *Read {
return &Read{}
}
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_line", api.ToolProperty{
Type: api.PropertyType{"integer"},
Description: "Optional 1-based line to start reading from.",
})
props.Set("end_line", api.ToolProperty{
Type: api.PropertyType{"integer"},
Description: "Optional 1-based inclusive line to stop reading at.",
})
props.Set("line_count", api.ToolProperty{
Type: api.PropertyType{"integer"},
Description: "Optional maximum number of lines to read, starting at start_line or line 1.",
})
props.Set("line_range", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: `Optional 1-based inclusive range like "10-40", "10:40", "10..40", "10-", or "10".`,
})
return api.ToolFunction{
Name: r.Name(),
Description: r.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: props,
Required: []string{"path"},
},
}
}
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 = io.ReadAll(file)
content = string(contentBytes)
}
if err != nil {
return agent.ToolResult{}, err
}
return agent.ToolResult{Content: content}, nil
}
type Edit struct{}
func NewEdit() *Edit {
return &Edit{}
}
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 := io.ReadAll(file)
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()
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 info.IsDir() {
file.Close()
return nil, nil, fmt.Errorf("%s is a directory", path)
}
return file, info, 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)
}
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 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}
var startSet, endSet bool
for _, key := range []string{"line_range", "range", "lines"} {
if lineRange, ok := stringReadArg(args, key); ok {
start, end, err := parseLineRange(lineRange)
if err != nil {
return readSelection{}, err
}
selection.enabled = true
if start > 0 {
selection.start = start
startSet = true
}
if end > 0 {
selection.end = end
endSet = true
}
break
}
}
if start, ok, err := intReadArg(args, "start_line"); err != nil {
return readSelection{}, err
} else if ok {
selection.enabled = true
selection.start = start
startSet = true
}
if end, ok, err := intReadArg(args, "end_line"); err != nil {
return readSelection{}, err
} else if ok {
selection.enabled = true
selection.end = end
endSet = true
}
lineCount, countSet, err := readLineCountArg(args)
if err != nil {
return readSelection{}, err
}
if countSet {
selection.enabled = true
if !startSet {
selection.start = 1
}
if !endSet {
selection.end = selection.start + lineCount - 1
}
}
if !selection.enabled {
return selection, nil
}
if selection.start < 1 {
return readSelection{}, fmt.Errorf("start_line must be greater than 0")
}
if selection.end > 0 && selection.end < selection.start {
return readSelection{}, fmt.Errorf("end_line must be greater than or equal to start_line")
}
return selection, nil
}
func readLineCountArg(args map[string]any) (int, bool, error) {
for _, key := range []string{"line_count", "num_lines"} {
value, ok, err := intReadArg(args, key)
if err != nil || ok {
if ok && value < 1 {
return 0, false, fmt.Errorf("%s must be greater than 0", key)
}
return value, ok, err
}
}
return 0, false, nil
}
func parseLineRange(value string) (int, int, error) {
value = strings.TrimSpace(value)
if value == "" {
return 0, 0, nil
}
value = strings.TrimPrefix(value, "lines")
value = strings.TrimPrefix(value, "line")
value = strings.TrimSpace(value)
for _, sep := range []string{"..", ":", ","} {
value = strings.ReplaceAll(value, sep, "-")
}
parts := strings.Split(value, "-")
if len(parts) > 2 {
return 0, 0, fmt.Errorf("line_range must look like 10-40, 10:40, 10..40, 10-, or 10")
}
start, end := 0, 0
var err error
if strings.TrimSpace(parts[0]) != "" {
start, err = strconv.Atoi(strings.TrimSpace(parts[0]))
if err != nil || start < 1 {
return 0, 0, fmt.Errorf("line_range start must be a positive line number")
}
}
if len(parts) == 1 {
return start, start, nil
}
if strings.TrimSpace(parts[1]) != "" {
end, err = strconv.Atoi(strings.TrimSpace(parts[1]))
if err != nil || end < 1 {
return 0, 0, fmt.Errorf("line_range end must be a positive line number")
}
}
if start == 0 && end == 0 {
return 0, 0, fmt.Errorf("line_range must include at least one line number")
}
if start == 0 {
start = 1
}
if end > 0 && end < start {
return 0, 0, fmt.Errorf("line_range end must be greater than or equal to start")
}
return start, end, nil
}
func readLineSelection(file *os.File, selection readSelection) (string, error) {
reader := bufio.NewReader(file)
var b strings.Builder
for lineNo := 1; ; lineNo++ {
line, err := reader.ReadString('\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.WriteString(line)
}
if err != nil {
if err == io.EOF {
break
}
return "", err
}
if selection.end > 0 && lineNo >= selection.end {
break
}
}
return b.String(), nil
}
func stringReadArg(args map[string]any, key string) (string, bool) {
value, ok := args[key].(string)
return value, ok && strings.TrimSpace(value) != ""
}
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"
}
+255
View File
@@ -0,0 +1,255 @@
package tools
import (
"context"
"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 := NewEdit().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 := NewEdit().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 := NewEdit().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 := NewEdit().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 := NewEdit().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 := NewRead().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 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 := NewRead().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 TestReadLineRange(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 := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"line_range": "2-3",
})
if err != nil {
t.Fatal(err)
}
if result.Content != "two\nthree\n" {
t.Fatalf("content = %q", result.Content)
}
}
func TestReadLinesAliasAsRange(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 := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"lines": "2-3",
})
if err != nil {
t.Fatal(err)
}
if result.Content != "two\nthree\n" {
t.Fatalf("content = %q", result.Content)
}
}
func TestReadLineCountFromStartLine(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 := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"start_line": 3,
"line_count": 2,
})
if err != nil {
t.Fatal(err)
}
if result.Content != "three\nfour\n" {
t.Fatalf("content = %q", result.Content)
}
}
func TestReadRejectsInvalidLineRange(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 := NewRead().Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{
"path": "note.txt",
"line_range": "4-2",
})
if err == nil {
t.Fatal("expected invalid range to fail")
}
if !strings.Contains(err.Error(), "line_range end") {
t.Fatalf("err = %v", err)
}
}
+137
View File
@@ -0,0 +1,137 @@
package tools
import (
"context"
"fmt"
"strings"
"github.com/ollama/ollama/agent"
"github.com/ollama/ollama/agent/skills"
"github.com/ollama/ollama/api"
)
type Skill struct {
catalog *skills.Catalog
}
func NewSkill(catalog *skills.Catalog) *Skill {
return &Skill{catalog: catalog}
}
func (s *Skill) Name() string {
return "skill"
}
func (s *Skill) Description() string {
return "Load the full SKILL.md instructions for an installed agent skill by name."
}
func (s *Skill) Schema() api.ToolFunction {
props := api.NewToolPropertiesMap()
props.Set("name", api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "Name of the skill to load.",
})
return api.ToolFunction{
Name: s.Name(),
Description: s.Description(),
Parameters: api.ToolFunctionParameters{
Type: "object",
Required: []string{"name"},
Properties: props,
},
}
}
func (s *Skill) Execute(_ context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) {
name, _ := args["name"].(string)
name = skills.NormalizeName(name)
if name == "" {
return agent.ToolResult{}, fmt.Errorf("name parameter is required")
}
if s.catalog == nil || s.catalog.Empty() {
return agent.ToolResult{}, fmt.Errorf("no skills are installed")
}
skill, ok := s.catalog.Find(name)
if !ok {
return agent.ToolResult{}, fmt.Errorf("unknown skill: %s", name)
}
content, err := SkillResultContent(skill)
if err != nil {
return agent.ToolResult{}, err
}
return agent.ToolResult{Content: content}, nil
}
func ManualSkillMessages(skill skills.Skill, request string, ordinal int) ([]api.Message, error) {
content, err := SkillResultContent(skill)
if err != nil {
return nil, err
}
args := api.NewToolCallFunctionArguments()
args.Set("name", skill.Name)
callID := manualSkillToolCallID(skill.Name, ordinal)
userContent := strings.TrimSpace(request)
if userContent == "" {
userContent = fmt.Sprintf("Use the %s skill.", skill.Name)
}
return []api.Message{
{Role: "user", Content: userContent},
{
Role: "assistant",
ToolCalls: []api.ToolCall{{
ID: callID,
Function: api.ToolCallFunction{
Name: "skill",
Arguments: args,
},
}},
},
{Role: "tool", ToolName: "skill", ToolCallID: callID, Content: content},
}, nil
}
func SkillResultContent(skill skills.Skill) (string, error) {
content, err := skill.Read()
if err != nil {
return "", err
}
var b strings.Builder
b.WriteString("Loaded skill: ")
b.WriteString(skill.Name)
b.WriteByte('\n')
b.WriteString("Skill directory: ")
b.WriteString(skill.Dir)
b.WriteString("\nResolve relative file references from the skill directory.\n\n")
b.WriteString(content)
return b.String(), nil
}
func manualSkillToolCallID(skillName string, ordinal int) string {
name := strings.Trim(strings.Map(func(r rune) rune {
switch {
case r >= 'a' && r <= 'z':
return r
case r >= 'A' && r <= 'Z':
return r
case r >= '0' && r <= '9':
return r
case r == '-' || r == '_':
return r
default:
return '-'
}
}, skillName), "-")
if name == "" {
name = "skill"
}
if ordinal <= 0 {
return "manual-skill-" + name
}
return fmt.Sprintf("manual-skill-%d-%s", ordinal, name)
}
+81
View File
@@ -0,0 +1,81 @@
package tools
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
"github.com/ollama/ollama/agent"
"github.com/ollama/ollama/agent/skills"
)
func TestSkillToolLoadsSkill(t *testing.T) {
dir := t.TempDir()
skillDir := filepath.Join(dir, "go-code")
if err := os.MkdirAll(skillDir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(skillDir, skills.SkillFile), []byte("---\nname: go-code\ndescription: Write Go code.\n---\n\n# Go Code\n"), 0o644); err != nil {
t.Fatal(err)
}
catalog, err := skills.Load(dir)
if err != nil {
t.Fatal(err)
}
result, err := NewSkill(catalog).Execute(context.Background(), agent.ToolContext{}, map[string]any{"name": "go-code"})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(result.Content, "Loaded skill: go-code") || !strings.Contains(result.Content, "# Go Code") {
t.Fatalf("content = %q", result.Content)
}
}
func TestManualSkillMessagesUseToolCallShape(t *testing.T) {
dir := t.TempDir()
skillDir := filepath.Join(dir, "go-code")
if err := os.MkdirAll(skillDir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(skillDir, skills.SkillFile), []byte("---\nname: go-code\ndescription: Write Go code.\n---\n\n# Go Code\n"), 0o644); err != nil {
t.Fatal(err)
}
catalog, err := skills.Load(dir)
if err != nil {
t.Fatal(err)
}
skill, ok := catalog.Find("go-code")
if !ok {
t.Fatal("skill not found")
}
messages, err := ManualSkillMessages(skill, "write a test", 7)
if err != nil {
t.Fatal(err)
}
if len(messages) != 3 {
t.Fatalf("messages = %d, want 3", len(messages))
}
if messages[0].Role != "user" || messages[0].Content != "write a test" {
t.Fatalf("user message = %#v", messages[0])
}
if messages[1].Role != "assistant" || len(messages[1].ToolCalls) != 1 {
t.Fatalf("assistant tool call = %#v", messages[1])
}
call := messages[1].ToolCalls[0]
if call.ID != "manual-skill-7-go-code" || call.Function.Name != "skill" {
t.Fatalf("tool call = %#v", call)
}
if name, _ := call.Function.Arguments.Get("name"); name != "go-code" {
t.Fatalf("tool args = %s", call.Function.Arguments.String())
}
if messages[2].Role != "tool" || messages[2].ToolName != "skill" || messages[2].ToolCallID != call.ID {
t.Fatalf("tool result metadata = %#v", messages[2])
}
if !strings.Contains(messages[2].Content, "Loaded skill: go-code") || !strings.Contains(messages[2].Content, "# Go Code") {
t.Fatalf("tool result = %q", messages[2].Content)
}
}
+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 NewWebSearch() *WebSearch {
return &WebSearch{}
}
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) 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 NewWebFetch() *WebFetch {
return &WebFetch{}
}
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) 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 TestWebToolsDoNotRequireApproval(t *testing.T) {
if coreagent.ToolRequiresApproval(NewWebSearch(), map[string]any{"query": "ollama"}) {
t.Fatal("web search should not require approval")
}
if coreagent.ToolRequiresApproval(NewWebFetch(), map[string]any{"url": "https://ollama.com"}) {
t.Fatal("web fetch should not 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 := NewWebFetch().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
}
// 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.
func (c *Client) Signout(ctx context.Context) error {
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)
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+30
View File
@@ -868,6 +868,36 @@ type StatusResponse struct {
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].
type GenerateResponse struct {
// Model is the model name that generated the response.
+183 -25
View File
@@ -14,7 +14,7 @@ import (
// currentSchemaVersion defines the current database schema version.
// Increment this when making schema changes that require migrations.
const currentSchemaVersion = 16
const currentSchemaVersion = 17
// database wraps the SQLite connection.
// SQLite handles its own locking for concurrent access:
@@ -97,6 +97,8 @@ func (db *database) init() error {
CREATE TABLE IF NOT EXISTS chats (
id TEXT PRIMARY KEY,
title TEXT NOT NULL DEFAULT '',
model_name TEXT NOT NULL DEFAULT '',
source TEXT NOT NULL DEFAULT 'app',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
browser_state TEXT
);
@@ -107,6 +109,7 @@ func (db *database) init() error {
role TEXT NOT NULL,
content TEXT NOT NULL DEFAULT '',
thinking TEXT NOT NULL DEFAULT '',
images TEXT NOT NULL DEFAULT '[]',
stream BOOLEAN NOT NULL DEFAULT 0,
model_name TEXT,
model_cloud BOOLEAN, -- deprecated
@@ -116,15 +119,21 @@ func (db *database) init() error {
thinking_time_start TIMESTAMP,
thinking_time_end TIMESTAMP,
tool_result TEXT,
tool_name TEXT NOT NULL DEFAULT '',
tool_call_id TEXT NOT NULL DEFAULT '',
archived BOOLEAN NOT NULL DEFAULT 0,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id ON messages(chat_id);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_id ON messages(chat_id, id);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_archived ON messages(chat_id, archived, id);
CREATE TABLE IF NOT EXISTS tool_calls (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
type TEXT NOT NULL,
tool_call_id TEXT NOT NULL DEFAULT '',
function_name TEXT NOT NULL,
function_arguments TEXT NOT NULL,
function_result TEXT,
@@ -133,6 +142,17 @@ func (db *database) init() error {
CREATE INDEX IF NOT EXISTS idx_tool_calls_message_id ON tool_calls(message_id);
CREATE TABLE IF NOT EXISTS compactions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
chat_id TEXT NOT NULL,
summary TEXT NOT NULL,
archived_message_ids TEXT NOT NULL DEFAULT '[]',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_compactions_chat_id ON compactions(chat_id, id);
CREATE TABLE IF NOT EXISTS attachments (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
@@ -271,6 +291,12 @@ func (db *database) migrate() error {
return fmt.Errorf("migrate v15 to v16: %w", err)
}
version = 16
case 16:
// add agent chat metadata, message archiving, and compaction tables
if err := db.migrateV16ToV17(); err != nil {
return fmt.Errorf("migrate v16 to v17: %w", err)
}
version = 17
default:
// If we have a version we don't recognize, just set it to current
// This might happen during development
@@ -278,6 +304,10 @@ func (db *database) migrate() error {
}
}
if err := db.ensureCurrentSchema(); err != nil {
return fmt.Errorf("ensure current schema: %w", err)
}
return nil
}
@@ -540,6 +570,128 @@ func (db *database) migrateV15ToV16() error {
return nil
}
// migrateV16ToV17 adds the agent chat persistence fields to the app database.
func (db *database) migrateV16ToV17() error {
for _, stmt := range []struct {
sql string
msg string
}{
{`ALTER TABLE chats ADD COLUMN model_name TEXT NOT NULL DEFAULT ''`, "add chats.model_name"},
{`ALTER TABLE chats ADD COLUMN source TEXT NOT NULL DEFAULT 'app'`, "add chats.source"},
{`ALTER TABLE messages ADD COLUMN images TEXT NOT NULL DEFAULT '[]'`, "add messages.images"},
{`ALTER TABLE messages ADD COLUMN tool_name TEXT NOT NULL DEFAULT ''`, "add messages.tool_name"},
{`ALTER TABLE messages ADD COLUMN tool_call_id TEXT NOT NULL DEFAULT ''`, "add messages.tool_call_id"},
{`ALTER TABLE messages ADD COLUMN archived BOOLEAN NOT NULL DEFAULT 0`, "add messages.archived"},
{`ALTER TABLE tool_calls ADD COLUMN tool_call_id TEXT NOT NULL DEFAULT ''`, "add tool_calls.tool_call_id"},
} {
_, err := db.conn.Exec(stmt.sql)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("%s: %w", stmt.msg, err)
}
}
_, err := db.conn.Exec(`
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_id ON messages(chat_id, id);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_archived ON messages(chat_id, archived, id);
CREATE TABLE IF NOT EXISTS compactions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
chat_id TEXT NOT NULL,
summary TEXT NOT NULL,
archived_message_ids TEXT NOT NULL DEFAULT '[]',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_compactions_chat_id ON compactions(chat_id, id);
UPDATE settings SET schema_version = 17;
`)
if err != nil {
return fmt.Errorf("create agent chat persistence tables: %w", err)
}
return nil
}
func (db *database) ensureCurrentSchema() error {
complete, err := db.agentPersistenceSchemaComplete()
if err != nil {
return err
}
if complete {
return nil
}
if err := db.migrateV16ToV17(); err != nil {
return err
}
complete, err = db.agentPersistenceSchemaComplete()
if err != nil {
return err
}
if !complete {
return fmt.Errorf("agent persistence schema is incomplete")
}
return nil
}
func (db *database) agentPersistenceSchemaComplete() (bool, error) {
for _, table := range []string{"compactions"} {
exists, err := db.tableExists(table)
if err != nil || !exists {
return false, err
}
}
for _, column := range []struct {
table string
name string
}{
{"chats", "model_name"},
{"chats", "source"},
{"messages", "images"},
{"messages", "tool_name"},
{"messages", "tool_call_id"},
{"messages", "archived"},
{"tool_calls", "tool_call_id"},
} {
exists, err := db.columnExists(column.table, column.name)
if err != nil || !exists {
return false, err
}
}
return true, nil
}
func (db *database) tableExists(table string) (bool, error) {
var count int
if err := db.conn.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = ?`, table).Scan(&count); err != nil {
return false, err
}
return count > 0, nil
}
func (db *database) columnExists(table, column string) (bool, error) {
rows, err := db.conn.Query(fmt.Sprintf("PRAGMA table_info(%s)", table))
if err != nil {
return false, err
}
defer rows.Close()
for rows.Next() {
var cid int
var name, dataType sql.NullString
var notNull, primaryKey int
var defaultValue sql.NullString
if err := rows.Scan(&cid, &name, &dataType, &notNull, &defaultValue, &primaryKey); err != nil {
return false, err
}
if name.String == column {
return true, nil
}
}
if err := rows.Err(); err != nil {
return false, err
}
return false, nil
}
// cleanupOrphanedData removes orphaned records that may exist due to the foreign key bug
func (db *database) cleanupOrphanedData() error {
_, err := db.conn.Exec(`
@@ -584,18 +736,21 @@ func (db *database) getAllChats() ([]Chat, error) {
c.id,
c.title,
c.created_at,
COALESCE(first_msg.content, '') as first_user_content,
COALESCE(datetime(MAX(m.updated_at)), datetime(c.created_at)) as last_updated
COALESCE((
SELECT fm.content
FROM messages fm
WHERE fm.chat_id = c.id
AND fm.role = 'user'
AND fm.archived = 0
ORDER BY fm.id ASC
LIMIT 1
), '') as first_user_content,
COALESCE(MAX(m.updated_at), c.created_at) as last_updated
FROM chats c
LEFT JOIN (
SELECT chat_id, content, MIN(id) as min_id
FROM messages
WHERE role = 'user'
GROUP BY chat_id
) first_msg ON c.id = first_msg.chat_id
LEFT JOIN messages m ON c.id = m.chat_id
GROUP BY c.id, c.title, c.created_at, first_msg.content
ORDER BY last_updated DESC
LEFT JOIN messages m ON c.id = m.chat_id AND m.archived = 0
WHERE c.source = 'app'
GROUP BY c.id, c.title, c.created_at
ORDER BY last_updated DESC, COALESCE(MAX(m.id), 0) DESC, c.created_at DESC, c.id DESC
`
rows, err := db.conn.Query(query)
@@ -618,25 +773,27 @@ func (db *database) getAllChats() ([]Chat, error) {
&firstUserContent,
&lastUpdatedStr,
)
// Parse the last updated time
lastUpdated, _ := time.Parse("2006-01-02 15:04:05", lastUpdatedStr)
if err != nil {
return nil, fmt.Errorf("scan chat: %w", err)
}
lastUpdated, err := parseAgentSQLiteTime(lastUpdatedStr)
if err != nil {
return nil, fmt.Errorf("parse chat updated_at: %w", err)
}
chat.CreatedAt = createdAt
// Add a dummy first user message for the UI to display
// This is just for the excerpt, full messages are loaded when needed
chat.Messages = []Message{}
if firstUserContent != "" {
chat.Messages = append(chat.Messages, Message{
Role: "user",
Content: firstUserContent,
UpdatedAt: lastUpdated,
})
// Add a summary message for the UI to display the excerpt and latest update.
// Full messages are loaded when a chat is opened.
summary := Message{
UpdatedAt: lastUpdated,
}
if firstUserContent != "" {
summary.Role = "user"
summary.Content = firstUserContent
}
chat.Messages = []Message{summary}
chats = append(chats, chat)
}
@@ -780,6 +937,7 @@ func (db *database) updateLastMessage(chatID string, msg Message) error {
var messageID int64
err = tx.QueryRow(`
SELECT MAX(id) FROM messages WHERE chat_id = ?
AND archived = 0
`, chatID).Scan(&messageID)
if err != nil {
return fmt.Errorf("get last message id: %w", err)
@@ -887,7 +1045,7 @@ func (db *database) getMessages(chatID string, loadAttachmentData bool) ([]Messa
query := `
SELECT id, role, content, thinking, stream, model_name, created_at, updated_at, thinking_time_start, thinking_time_end, tool_result
FROM messages
WHERE chat_id = ?
WHERE chat_id = ? AND archived = 0
ORDER BY id ASC
`
+103
View File
@@ -174,6 +174,100 @@ func TestMigrationV15ToV16LastHomeViewDefaultsToLaunch(t *testing.T) {
}
}
func TestMigrationV16ToV17AddsAgentSchema(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "test.db")
db, err := newDatabase(dbPath)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
if _, err := db.conn.Exec(`
DROP INDEX IF EXISTS idx_messages_chat_id_id;
DROP INDEX IF EXISTS idx_messages_chat_id_archived;
DROP INDEX IF EXISTS idx_compactions_chat_id;
DROP TABLE IF EXISTS compactions;
ALTER TABLE chats DROP COLUMN model_name;
ALTER TABLE chats DROP COLUMN source;
ALTER TABLE messages DROP COLUMN images;
ALTER TABLE messages DROP COLUMN tool_name;
ALTER TABLE messages DROP COLUMN tool_call_id;
ALTER TABLE messages DROP COLUMN archived;
ALTER TABLE tool_calls DROP COLUMN tool_call_id;
UPDATE settings SET schema_version = 16;
`); err != nil {
t.Fatalf("failed to seed v16 schema: %v", err)
}
if err := db.migrate(); err != nil {
t.Fatalf("migration from v16 to v17 failed: %v", err)
}
version, err := db.getSchemaVersion()
if err != nil {
t.Fatalf("failed to get schema version: %v", err)
}
if version != 17 {
t.Fatalf("expected schema version 17, got %d", version)
}
columns := columnMap(db)
for _, want := range []struct {
table string
column string
}{
{"chats", "model_name TEXT NOT NULL DEFAULT ''"},
{"chats", "source TEXT NOT NULL DEFAULT 'app'"},
{"messages", "images TEXT NOT NULL DEFAULT '[]'"},
{"messages", "archived BOOLEAN NOT NULL DEFAULT 0"},
{"tool_calls", "tool_call_id TEXT NOT NULL DEFAULT ''"},
} {
if !containsString(columns[want.table], want.column) {
t.Fatalf("%s columns missing %q: %#v", want.table, want.column, columns[want.table])
}
}
if _, ok := columns["compactions"]; !ok {
t.Fatalf("compactions table was not created: %#v", columns)
}
}
func TestMigrationRepairsIncompleteCurrentSchema(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "test.db")
db, err := newDatabase(dbPath)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
if _, err := db.conn.Exec(`
ALTER TABLE chats DROP COLUMN source;
UPDATE settings SET schema_version = 17;
`); err != nil {
t.Fatalf("failed to seed incomplete current schema: %v", err)
}
if err := db.migrate(); err != nil {
t.Fatalf("migration repair failed: %v", err)
}
version, err := db.getSchemaVersion()
if err != nil {
t.Fatalf("failed to get schema version: %v", err)
}
if version != currentSchemaVersion {
t.Fatalf("expected schema version %d, got %d", currentSchemaVersion, version)
}
columns := columnMap(db)
if !containsString(columns["chats"], "source TEXT NOT NULL DEFAULT 'app'") {
t.Fatalf("chats.source was not repaired: %#v", columns["chats"])
}
}
func TestChatDeletionWithCascade(t *testing.T) {
t.Run("chat deletion cascades to related messages", func(t *testing.T) {
tmpDir := t.TempDir()
@@ -369,6 +463,15 @@ func countRowsWithCondition(t *testing.T, db *database, table, condition string,
return count
}
func containsString(values []string, want string) bool {
for _, value := range values {
if value == want {
return true
}
}
return false
}
// Test helpers for schema migration testing
// schemaMap returns both tables/columns and indexes (ignoring order)
+1 -1
View File
@@ -41,7 +41,7 @@ func ImgBytes(path string) ([]byte, error) {
func (s *Store) ImgDir() string {
dbPath := s.DBPath
if dbPath == "" {
dbPath = defaultDBPath
dbPath = defaultDBPath()
}
storeDir := filepath.Dir(dbPath)
return filepath.Join(storeDir, "cache", "images")
+3 -2
View File
@@ -107,6 +107,7 @@ type Chat struct {
ID string `json:"id"`
Messages []Message `json:"messages"`
Title string `json:"title"`
Model string `json:"model,omitempty"`
CreatedAt time.Time `json:"created_at"`
BrowserState json.RawMessage `json:"browser_state,omitempty" ts_type:"BrowserStateData"`
}
@@ -192,7 +193,7 @@ var defaultDBPath = func() string {
default:
return filepath.Join(os.Getenv("HOME"), ".ollama", "db.sqlite")
}
}()
}
// legacyConfigPath is the path to the old config.json file
var legacyConfigPath = func() string {
@@ -229,7 +230,7 @@ func (s *Store) ensureDB() error {
dbPath := s.DBPath
if dbPath == "" {
dbPath = defaultDBPath
dbPath = defaultDBPath()
}
// Ensure directory exists
+91
View File
@@ -5,6 +5,7 @@ package store
import (
"path/filepath"
"testing"
"time"
)
func TestStore(t *testing.T) {
@@ -227,6 +228,96 @@ func TestStore(t *testing.T) {
})
}
func TestStoreChatSummariesUseLatestActivity(t *testing.T) {
s, cleanup := setupTestStore(t)
defer cleanup()
base := time.Date(2026, 6, 23, 15, 30, 45, 0, time.UTC)
oldChat := NewChat("chat-old")
oldChat.Title = "Old Chat"
oldChat.CreatedAt = base
oldChat.Messages = []Message{
{
Role: "user",
Content: "older first prompt",
CreatedAt: base.Add(100 * time.Millisecond),
UpdatedAt: base.Add(100 * time.Millisecond),
},
{
Role: "user",
Content: "older second prompt",
CreatedAt: base.Add(200 * time.Millisecond),
UpdatedAt: base.Add(200 * time.Millisecond),
},
}
if err := s.SetChat(*oldChat); err != nil {
t.Fatalf("failed to save old chat: %v", err)
}
newChat := NewChat("chat-new")
newChat.Title = "New Chat"
newChat.CreatedAt = base
newChat.Messages = []Message{
{
Role: "user",
Content: "newer prompt",
CreatedAt: base.Add(900 * time.Millisecond),
UpdatedAt: base.Add(900 * time.Millisecond),
},
}
if err := s.SetChat(*newChat); err != nil {
t.Fatalf("failed to save new chat: %v", err)
}
activityOnlyChat := NewChat("chat-activity-only")
activityOnlyChat.Title = "Activity Only Chat"
activityOnlyChat.CreatedAt = base
activityOnlyChat.Messages = []Message{
{
Role: "assistant",
Content: "recent assistant activity",
CreatedAt: base.Add(1500 * time.Millisecond),
UpdatedAt: base.Add(1500 * time.Millisecond),
},
}
if err := s.SetChat(*activityOnlyChat); err != nil {
t.Fatalf("failed to save activity-only chat: %v", err)
}
chats, err := s.Chats()
if err != nil {
t.Fatalf("failed to list chats: %v", err)
}
if len(chats) != 3 {
t.Fatalf("expected 3 chats, got %d", len(chats))
}
if chats[0].ID != "chat-activity-only" {
t.Fatalf("expected chat-activity-only first, got %s", chats[0].ID)
}
for _, chat := range chats {
if len(chat.Messages) != 1 {
t.Fatalf("expected summary message for %s, got %d messages", chat.ID, len(chat.Messages))
}
}
if !chats[0].Messages[0].UpdatedAt.Equal(activityOnlyChat.Messages[0].UpdatedAt) {
t.Fatalf("expected latest activity updated_at %s, got %s", activityOnlyChat.Messages[0].UpdatedAt, chats[0].Messages[0].UpdatedAt)
}
if chats[0].Messages[0].Role != "" || chats[0].Messages[0].Content != "" {
t.Fatalf("expected activity-only chat to have no user excerpt, got role=%q content=%q", chats[0].Messages[0].Role, chats[0].Messages[0].Content)
}
if chats[1].ID != "chat-new" {
t.Fatalf("expected chat-new second, got %s", chats[1].ID)
}
if !chats[1].Messages[0].UpdatedAt.Equal(newChat.Messages[0].UpdatedAt) {
t.Fatalf("expected precise updated_at %s, got %s", newChat.Messages[0].UpdatedAt, chats[1].Messages[0].UpdatedAt)
}
if chats[2].Messages[0].Content != "older first prompt" {
t.Fatalf("expected first user prompt excerpt, got %q", chats[2].Messages[0].Content)
}
}
// setupTestStore creates a temporary store for testing
func setupTestStore(t *testing.T) (*Store, func()) {
t.Helper()
+31
View File
@@ -0,0 +1,31 @@
//go:build windows || darwin
package store
import (
"fmt"
"strings"
"time"
)
func parseAgentSQLiteTime(value string) (time.Time, error) {
value = strings.TrimSpace(value)
if value == "" {
return time.Time{}, nil
}
for _, layout := range []string{
time.RFC3339Nano,
"2006-01-02 15:04:05.999999999-07:00",
"2006-01-02 15:04:05.999999999Z07:00",
"2006-01-02 15:04:05.999999999",
"2006-01-02 15:04:05-07:00",
"2006-01-02 15:04:05Z07:00",
"2006-01-02 15:04:05",
} {
t, err := time.Parse(layout, value)
if err == nil {
return t, nil
}
}
return time.Time{}, fmt.Errorf("unsupported time format %q", value)
}
+885
View File
@@ -0,0 +1,885 @@
package cmd
import (
"context"
"database/sql"
"errors"
"fmt"
"net/http"
"os"
"os/signal"
"runtime"
"slices"
"strconv"
"strings"
"syscall"
"time"
"github.com/google/uuid"
"github.com/spf13/cobra"
coreagent "github.com/ollama/ollama/agent"
"github.com/ollama/ollama/agent/skills"
agentstore "github.com/ollama/ollama/agent/store"
agenttools "github.com/ollama/ollama/agent/tools"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/cmd/config"
"github.com/ollama/ollama/cmd/internal/filedata"
"github.com/ollama/ollama/cmd/launch"
agentchat "github.com/ollama/ollama/cmd/tui/chat"
"github.com/ollama/ollama/format"
internalcloud "github.com/ollama/ollama/internal/cloud"
"github.com/ollama/ollama/internal/modelref"
"github.com/ollama/ollama/types/model"
)
type AgentTUIOptions struct {
Model string
Prompt string
Messages []api.Message
System string
Images []api.ImageData
Format string
Options map[string]any
Think *api.ThinkValue
KeepAlive *api.Duration
ContextWindowTokens int
Resume bool
AutoApproveTools bool
Policy coreagent.RunPolicy
Verbose bool
MultiModal bool
Skill string
Skills *skills.Catalog
}
func agentOptionsFromRunOptions(opts runOptions) AgentTUIOptions {
return AgentTUIOptions{
Model: opts.Model,
Prompt: opts.Prompt,
Messages: opts.Messages,
System: opts.System,
Images: opts.Images,
Format: opts.Format,
Options: opts.Options,
Think: opts.Think,
KeepAlive: opts.KeepAlive,
ContextWindowTokens: opts.ContextWindowTokens,
Resume: opts.Resume,
AutoApproveTools: opts.AutoApproveTools,
Verbose: opts.Verbose,
MultiModal: opts.MultiModal,
}
}
type agentSurface int
const (
agentSurfaceTUI agentSurface = iota
agentSurfaceHeadless
)
func resolveAgentRunPolicy(opts AgentTUIOptions, surface agentSurface) coreagent.RunPolicy {
policy := opts.Policy
if surface == agentSurfaceHeadless {
policy.ToolMode = coreagent.ToolModeDisabled
}
if opts.AutoApproveTools {
policy.ToolMode = coreagent.ToolModeFullAccess
}
return policy
}
type agentRunSetup struct {
opts AgentTUIOptions
client *api.Client
cwd string
store *agentstore.Store
newChatID func(context.Context) (string, error)
chatID string
messages []api.Message
skills *skills.Catalog
registry *coreagent.Registry
approval coreagent.ApprovalHandler
}
func (s *agentRunSetup) close() {
if s != nil && s.store != nil {
_ = s.store.Close()
s.store = nil
}
}
func newAgentRunSetup(cmd *cobra.Command, opts AgentTUIOptions, resumeLatestWithoutModel bool) (*agentRunSetup, error) {
client, err := api.ClientFromEnvironment()
if err != nil {
return nil, err
}
cwd, err := os.Getwd()
if err != nil {
cwd = ""
}
var store *agentstore.Store
if openedStore, err := agentstore.New(""); err != nil {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m chat persistence unavailable: %v\n", err)
} else {
store = openedStore
}
newChatID := func(ctx context.Context) (string, error) {
u, err := uuid.NewV7()
if err != nil {
u = uuid.Must(uuid.NewRandom())
}
chatID := u.String()
if store != nil {
if err := store.EnsureChat(ctx, chatID, ""); err != nil {
return "", err
}
}
return chatID, nil
}
chatID := ""
var resumedMessages []api.Message
if opts.Resume {
if store == nil {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m chat resume unavailable: persistence is disabled\n")
} else {
chat, err := resumeAgentChat(cmd.Context(), store, opts.Model, resumeLatestWithoutModel)
if err == nil {
chatID = chat.ID
if opts.Model == "" {
opts.Model = chat.Model
}
resumedMessages = chat.Messages
} else if errors.Is(err, sql.ErrNoRows) {
if resumeLatestWithoutModel && opts.Model == "" {
if store != nil {
_ = store.Close()
}
return nil, errors.New("no saved chat to resume; pass a model to start a new chat")
}
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m no saved chat for %s; starting a new chat\n", opts.Model)
} else if resumeLatestWithoutModel {
if store != nil {
_ = store.Close()
}
return nil, fmt.Errorf("could not resume chat: %w", err)
} else {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not resume chat: %v\n", err)
}
}
}
if strings.TrimSpace(opts.Model) == "" {
if store != nil {
_ = store.Close()
}
return nil, errors.New("model is required")
}
if chatID == "" {
var err error
chatID, err = newChatID(cmd.Context())
if err != nil {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not create persistent chat: %v\n", err)
if store != nil {
_ = store.Close()
}
store = nil
chatID, _ = newChatID(cmd.Context())
}
}
if store != nil {
if err := store.SetChatModel(cmd.Context(), chatID, opts.Model); err != nil {
_ = store.Close()
return nil, err
}
}
skillCatalog := opts.Skills
if skillCatalog == nil {
skillCatalog = loadAgentSkills()
}
registry := agentToolsRegistry(cmd.Context(), client, opts.Model, skillCatalog)
opts.ContextWindowTokens = contextWindowTokensForRun(cmd.Context(), client, opts.Model, opts.ContextWindowTokens)
approval := opts.Policy.ReviewApprovalHandler(nil)
messages := slices.Clone(resumedMessages)
messages = append(messages, opts.Messages...)
return &agentRunSetup{
opts: opts,
client: client,
cwd: cwd,
store: store,
newChatID: newChatID,
chatID: chatID,
messages: messages,
skills: skillCatalog,
registry: registry,
approval: approval,
}, nil
}
func resumeAgentChat(ctx context.Context, store *agentstore.Store, modelName string, latestWithoutModel bool) (*agentstore.AgentChat, error) {
if latestWithoutModel && modelName == "" {
return store.LatestChat(ctx)
}
return store.LatestChatForModel(ctx, modelName)
}
func GenerateAgentTUI(cmd *cobra.Command, opts AgentTUIOptions) error {
opts.Policy = resolveAgentRunPolicy(opts, agentSurfaceTUI)
setup, err := newAgentRunSetup(cmd, opts, false)
if err != nil {
return err
}
defer setup.close()
opts = setup.opts
_, err = agentchat.Run(cmd.Context(), agentchat.Options{
Model: opts.Model,
ChatID: setup.chatID,
Messages: setup.messages,
Client: setup.client,
Store: setup.store,
Tools: setup.registry,
ToolRegistryForModel: func(ctx context.Context, model string) *coreagent.Registry {
return agentToolsRegistry(ctx, setup.client, model, setup.skills)
},
MultiModalForModel: func(ctx context.Context, model string) bool {
return agentModelSupportsMultimodal(ctx, setup.client, model)
},
ModelOptions: func(ctx context.Context) ([]agentchat.ModelOption, error) {
return agentModelOptions(ctx, setup.client)
},
OnModelSelected: func(_ context.Context, model string) error {
return config.SetLastModel(model)
},
SystemPromptForModel: func(ctx context.Context, model string, registry *coreagent.Registry) string {
modelSystem := opts.System
if strings.TrimSpace(model) != strings.TrimSpace(opts.Model) {
modelSystem = agentSystemFromShow(ctx, setup.client, model)
}
return agentSystemPrompt(model, setup.skills, registry != nil && registry.Has("skill"), modelSystem, "")
},
Approval: setup.approval,
Policy: opts.Policy,
Skills: setup.skills,
SystemPrompt: agentSystemPrompt(opts.Model, setup.skills, setup.registry != nil && setup.registry.Has("skill"), opts.System, ""),
WorkingDir: setup.cwd,
Format: opts.Format,
Options: opts.Options,
Think: opts.Think,
KeepAlive: opts.KeepAlive,
Images: slices.Clone(opts.Images),
MultiModal: opts.MultiModal,
Verbose: opts.Verbose,
Compactor: coreagent.NewSimpleCompactor(setup.client, setup.store, coreagent.CompactionOptions{
ContextWindowTokens: opts.ContextWindowTokens,
}),
ContextWindowTokens: opts.ContextWindowTokens,
ContextWindowTokensForModel: func(ctx context.Context, model string, fallback int) int {
return contextWindowTokensForRun(ctx, setup.client, model, fallback)
},
PreloadModel: func(ctx context.Context, model string, think *api.ThinkValue) error {
preloadOpts := opts
preloadOpts.Think = think
return preloadAgentModelIfLocal(ctx, setup.client, preloadOpts, model)
},
CheckCloudModel: func(ctx context.Context, model, requiredPlan string) error {
return ensureCloudModelAccess(ctx, setup.client, model, requiredPlan)
},
OpenBrowser: func(url string) {
launch.OpenBrowser(url)
},
PollCloudAuth: func(ctx context.Context) (string, bool) {
user, err := setup.client.Whoami(ctx)
if err != nil || user == nil || user.Name == "" {
return "", false
}
return user.Name, true
},
NewChat: setup.newChatID,
})
if err != nil {
return err
}
return nil
}
func GenerateAgentHeadless(cmd *cobra.Command, opts AgentTUIOptions) error {
if strings.TrimSpace(opts.Prompt) == "" {
return errors.New("agent headless mode requires a prompt or stdin")
}
opts.Policy = resolveAgentRunPolicy(opts, agentSurfaceHeadless)
setup, err := newAgentRunSetup(cmd, opts, true)
if err != nil {
return err
}
defer setup.close()
opts = setup.opts
if opts.Model == "" {
return errors.New("model is required")
}
prompt := opts.Prompt
images := slices.Clone(opts.Images)
if opts.MultiModal {
var files []filedata.File
prompt, files, err = filedata.ExtractWithFiles(prompt)
if err != nil {
fmt.Fprintf(os.Stderr, "Couldn't process file: %q\n", err)
return err
}
imgs := make([]api.ImageData, 0, len(files))
for _, file := range files {
switch filedata.Kind(file.Path) {
case "audio":
fmt.Fprintf(os.Stderr, "Added audio '%s'\n", file.Path)
default:
fmt.Fprintf(os.Stderr, "Added image '%s'\n", file.Path)
}
imgs = append(imgs, file.Data)
}
images = imgs
}
tools := opts.Policy.Tools(setup.registry)
toolPrompt := ""
if tools == nil {
toolPrompt = "Tools are unavailable in this headless run because --auto-approve-tools was not passed. Answer directly without tool calls."
}
systemPrompt := agentSystemPrompt(opts.Model, setup.skills, tools != nil && tools.Has("skill"), opts.System, toolPrompt)
newMessages := []api.Message{{Role: "user", Content: prompt, Images: images}}
if strings.TrimSpace(opts.Skill) == "" {
if skill, request, ok := skillFromPrompt(setup.skills, prompt); ok {
opts.Skill = skill.Name
prompt = request
}
}
if strings.TrimSpace(opts.Skill) != "" {
skill, ok := setup.skills.Find(opts.Skill)
if !ok {
return fmt.Errorf("unknown skill: %s", opts.Skill)
}
manualMessages, err := agenttools.ManualSkillMessages(skill, prompt, len(setup.messages)+1)
if err != nil {
return err
}
manualMessages[0].Images = images
newMessages = manualMessages
}
runCtx, cancel := context.WithCancel(cmd.Context())
defer cancel()
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT)
defer signal.Stop(sigChan)
go func() {
select {
case <-sigChan:
cancel()
case <-runCtx.Done():
}
}()
headlessSink := &agentHeadlessEventSink{}
eventSink := coreagent.EventSink(headlessSink)
session := &coreagent.Session{
Client: setup.client,
Store: setup.store,
Events: eventSink,
Tools: tools,
Approval: opts.Policy.ApprovalHandler(nil),
WorkingDir: setup.cwd,
Compactor: coreagent.NewSimpleCompactor(setup.client, setup.store, coreagent.CompactionOptions{
ContextWindowTokens: opts.ContextWindowTokens,
}),
}
result, err := session.Run(runCtx, coreagent.RunOptions{
ChatID: setup.chatID,
Model: opts.Model,
SystemPrompt: systemPrompt,
Messages: setup.messages,
NewMessages: newMessages,
Format: opts.Format,
Options: opts.Options,
Think: opts.Think,
KeepAlive: opts.KeepAlive,
Policy: opts.Policy,
})
if err != nil {
return err
}
headlessSink.Finish()
if headlessSink.denied {
return errors.New("tool execution denied")
}
verbose := opts.Verbose
if cmd != nil && cmd.Flags().Lookup("verbose") != nil {
flagVerbose, err := cmd.Flags().GetBool("verbose")
if err != nil {
return err
}
verbose = verbose || flagVerbose
}
if verbose {
result.Latest.Summary()
}
return nil
}
func skillFromPrompt(catalog *skills.Catalog, prompt string) (skills.Skill, string, bool) {
if catalog == nil || catalog.Empty() {
return skills.Skill{}, "", false
}
prompt = strings.TrimSpace(prompt)
if !strings.HasPrefix(prompt, "/") {
return skills.Skill{}, "", false
}
command, rest, _ := strings.Cut(prompt, " ")
skill, ok := catalog.Find(command)
return skill, strings.TrimSpace(rest), ok
}
type agentHeadlessEventSink struct {
wroteContent bool
contentEndedWithNewline bool
denied bool
}
func (s *agentHeadlessEventSink) Emit(event coreagent.Event) error {
switch event.Type {
case coreagent.EventThinkingDelta:
case coreagent.EventMessageDelta:
if event.Content != "" {
fmt.Fprint(os.Stdout, event.Content)
s.wroteContent = true
s.contentEndedWithNewline = strings.HasSuffix(event.Content, "\n")
}
case coreagent.EventToolFinished:
s.ensureContentNewline()
status := "done"
if event.Status != "done" || event.Error != "" {
status = "failed"
}
fmt.Fprintf(os.Stderr, "• %s %s\n", coreagent.ToolInvocationLabel(event.ToolName, event.Args), status)
case coreagent.EventToolsUnavailable:
fmt.Fprintln(os.Stderr, "Tools are unavailable for this model.")
case coreagent.EventRunFinished:
if event.Status == "denied" {
s.denied = true
}
case coreagent.EventCompactionSkipped:
if event.Content != "" {
fmt.Fprintf(os.Stderr, "%s\n", event.Content)
}
case coreagent.EventError:
if event.Error != "" {
fmt.Fprintf(os.Stderr, "error: %s\n", event.Error)
}
}
return nil
}
func (s *agentHeadlessEventSink) Finish() {
s.ensureContentNewline()
}
func (s *agentHeadlessEventSink) ensureContentNewline() {
if s.wroteContent && !s.contentEndedWithNewline {
fmt.Fprintln(os.Stdout)
s.contentEndedWithNewline = true
}
}
func loadAgentSkills() *skills.Catalog {
catalog, err := skills.LoadDefault()
if err != nil {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not load skills: %v\n", err)
return nil
}
for _, warning := range catalog.Warnings {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m skill ignored: %s\n", warning)
}
return catalog
}
func agentSystemPrompt(modelName string, catalog *skills.Catalog, skillToolAvailable bool, modelSystem string, extra string) string {
return agentSystemPromptAt(time.Now(), modelName, catalog, skillToolAvailable, modelSystem, extra)
}
func agentSystemPromptAt(now time.Time, modelName string, catalog *skills.Catalog, skillToolAvailable bool, modelSystem string, extra string) string {
var parts []string
parts = append(parts, agentDefaultSystemPrompt(now, modelName))
if strings.TrimSpace(modelSystem) != "" {
parts = append(parts, strings.TrimSpace(modelSystem))
}
if catalogPrompt := catalog.SystemPrompt(skillToolAvailable); strings.TrimSpace(catalogPrompt) != "" {
parts = append(parts, catalogPrompt)
}
if strings.TrimSpace(extra) != "" {
parts = append(parts, strings.TrimSpace(extra))
}
return strings.Join(parts, "\n\n")
}
func agentSystemFromShow(ctx context.Context, client *api.Client, modelName string) string {
if client == nil || strings.TrimSpace(modelName) == "" {
return ""
}
resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName})
if err != nil {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not load model system prompt: %v\n", err)
return ""
}
return resp.System
}
func agentDefaultSystemPrompt(now time.Time, modelName string) string {
date := now.Format("Monday, January 2, 2006")
shellName := "bash"
if runtime.GOOS == "windows" {
shellName = "PowerShell"
}
return strings.Join([]string{
"You are running in Ollama, in a harness to help the user accomplish tasks, and the model is " + modelName + ".",
"",
"Current date: " + date + ".",
"",
"Be concise, practical, and action-oriented. Use tools when they materially help. Verify current or fast-changing facts with web tools when available; otherwise state uncertainty.",
"",
"Use " + shellName + " carefully. Prefer read-only inspection first. Stay within the current working directory unless explicitly asked. Surface intent before risky actions such as writes, deletes, moves, installs, git state changes, service changes, sudo, secrets access, network scripts, or commands outside the working directory. Request approval when required and do not work around denied approvals.",
"",
"Tell the user about meaningful changes, verification, failures, blockers, assumptions, and risks. Summarize routine tool output instead of dumping it.",
}, "\n")
}
func agentModelOptions(ctx context.Context, client *api.Client) ([]agentchat.ModelOption, error) {
if client == nil {
return nil, errors.New("model picker requires an API client")
}
list, err := client.List(ctx)
if err != nil {
return nil, err
}
seen := make(map[string]struct{})
var options []agentchat.ModelOption
add := func(name, description string, recommended bool, requiredPlan string, cloud bool) {
name = strings.TrimSpace(name)
if name == "" {
return
}
key := strings.ToLower(name)
if _, ok := seen[key]; ok {
return
}
seen[key] = struct{}{}
options = append(options, agentchat.ModelOption{
Name: name,
Description: strings.TrimSpace(description),
Recommended: recommended,
RequiredPlan: requiredPlan,
Cloud: cloud,
})
}
if disabled, known := agentCloudStatusDisabled(ctx, client); !known || !disabled {
if recs, err := client.ModelRecommendationsExperimental(ctx); err == nil {
for _, rec := range recs.Recommendations {
name := strings.TrimSpace(rec.Model)
if !modelref.HasExplicitCloudSource(name) {
continue
}
add(name, agentRecommendationDescription(rec), true, strings.TrimSpace(rec.RequiredPlan), true)
}
}
}
local := slices.Clone(list.Models)
slices.SortStableFunc(local, func(a, b api.ListModelResponse) int {
return strings.Compare(strings.ToLower(a.Name), strings.ToLower(b.Name))
})
for _, model := range local {
name := strings.TrimSpace(model.Name)
if name == "" {
name = strings.TrimSpace(model.Model)
}
name = strings.TrimSuffix(name, ":latest")
if modelref.HasExplicitCloudSource(name) {
// Cloud-sourced models in the local tags list keep their ":cloud"
// suffix in the name and show arch details without a "cloud" marker.
// If a curated rec already added this model, the dedup above keeps
// the rec entry.
add(name, agentCloudModelDescription(model), false, "", true)
continue
}
add(name, agentLocalModelDescription(model), false, "", false)
}
// Compute availability badges for cloud models based on account state.
badges, signInURLs := cloudAvailabilityBadges(ctx, client, options)
for i := range options {
options[i].AvailabilityBadge = badges[options[i].Name]
options[i].SignInURL = signInURLs[options[i].Name]
}
return options, nil
}
// cloudAvailabilityBadges returns a map of model name → availability badge
// for cloud models that require sign-in or a plan upgrade. It also returns
// a map of model name → sign-in URL for models that require sign-in.
func cloudAvailabilityBadges(ctx context.Context, client *api.Client, options []agentchat.ModelOption) (map[string]string, map[string]string) {
badges := make(map[string]string)
signInURLs := make(map[string]string)
hasCloud := false
for _, opt := range options {
if opt.Cloud {
hasCloud = true
break
}
}
if !hasCloud {
return badges, signInURLs
}
if disabled, known := agentCloudStatusDisabled(ctx, client); known && disabled {
return badges, signInURLs
}
whoamiCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
defer cancel()
user, err := client.Whoami(whoamiCtx)
if err != nil {
// Whoami failed — likely not signed in. Extract the sign-in URL
// from the authorization error so we can show it immediately.
var authErr api.AuthorizationError
signInURL := ""
if errors.As(err, &authErr) && authErr.SigninURL != "" {
signInURL = authErr.SigninURL
}
for _, opt := range options {
if opt.Cloud {
badges[opt.Name] = "Sign in required"
if signInURL != "" {
signInURLs[opt.Name] = signInURL
}
}
}
return badges, signInURLs
}
signedIn := user != nil && user.Name != ""
for _, opt := range options {
if !opt.Cloud {
continue
}
if !signedIn {
badges[opt.Name] = "Sign in required"
} else if opt.RequiredPlan != "" && !launch.PlanSatisfies(user.Plan, opt.RequiredPlan) {
badges[opt.Name] = "Upgrade required"
}
}
return badges, signInURLs
}
func agentRecommendationDescription(rec api.ModelRecommendation) string {
var parts []string
if description := strings.TrimSpace(rec.Description); description != "" {
parts = append(parts, description)
} else {
parts = append(parts, "cloud")
}
if rec.ContextLength > 0 {
parts = append(parts, format.HumanNumber(uint64(rec.ContextLength))+" ctx")
}
return strings.Join(parts, " · ")
}
// agentModelArchDescription builds the shared arch-details + context segment
// used by both local and cloud-sourced tag entries: "<family> <params> <quant>
// · N ctx>". Parameter sizes that are raw numbers (common for cloud stubs) are
// humanized, e.g. "27000000000" -> "27B".
func agentModelArchDescription(model api.ListModelResponse) string {
var details []string
if model.Details.Family != "" {
details = append(details, model.Details.Family)
}
if ps := humanizedParameterSize(model.Details.ParameterSize); ps != "" {
details = append(details, ps)
}
if model.Details.QuantizationLevel != "" {
details = append(details, model.Details.QuantizationLevel)
}
var parts []string
if len(details) > 0 {
parts = append(parts, strings.Join(details, " "))
}
if model.Details.ContextLength > 0 {
parts = append(parts, format.HumanNumber(uint64(model.Details.ContextLength))+" ctx")
}
return strings.Join(parts, " · ")
}
func agentLocalModelDescription(model api.ListModelResponse) string {
desc := agentModelArchDescription(model)
if desc == "" {
return "local"
}
return "local · " + desc
}
// agentCloudModelDescription describes a cloud-sourced model present in the
// local tags list. It omits a "cloud" marker (the model name keeps its
// ":cloud"/"-cloud" suffix) and shows the same arch details as local models.
func agentCloudModelDescription(model api.ListModelResponse) string {
return agentModelArchDescription(model)
}
func humanizedParameterSize(s string) string {
s = strings.TrimSpace(s)
if s == "" {
return ""
}
if f, err := strconv.ParseFloat(s, 64); err == nil {
return format.HumanNumber(uint64(f))
}
return s
}
func agentToolsRegistry(ctx context.Context, client *api.Client, modelName string, catalog *skills.Catalog) *coreagent.Registry {
supportsTools, err := agentModelSupportsTools(ctx, client, modelName)
if err != nil {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not check model capabilities: %v\n", err)
}
if !supportsTools {
return nil
}
registry := coreagent.NewRegistry()
if os.Getenv("OLLAMA_AGENT_DISABLE_SHELL") == "" {
registry.Register(agenttools.NewBash())
}
registry.Register(agenttools.NewRead())
registry.Register(agenttools.NewEdit())
if !catalog.Empty() {
registry.Register(agenttools.NewSkill(catalog))
}
if os.Getenv("OLLAMA_AGENT_DISABLE_WEBSEARCH") == "" {
if disabled, known := agentCloudStatusDisabled(ctx, client); !known || !disabled {
registry.Register(agenttools.NewWebSearch())
registry.Register(agenttools.NewWebFetch())
} else {
fmt.Fprintf(os.Stderr, "%s\n", internalcloud.DisabledError("web search is unavailable"))
}
}
return registry
}
func preloadAgentModelIfLocal(ctx context.Context, client *api.Client, opts AgentTUIOptions, modelName string) error {
modelName = strings.TrimSpace(modelName)
if client == nil || modelName == "" {
return nil
}
info, err := client.Show(ctx, &api.ShowRequest{Model: modelName})
if err != nil {
return err
}
if info.RemoteHost != "" || modelref.HasExplicitCloudSource(modelName) {
return nil
}
return preloadLocalModel(ctx, client, runOptions{
Model: modelName,
KeepAlive: opts.KeepAlive,
Think: opts.Think,
})
}
func agentModelSupportsTools(ctx context.Context, client *api.Client, modelName string) (bool, error) {
resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName})
if err != nil {
return false, err
}
return slices.Contains(resp.Capabilities, model.CapabilityTools), nil
}
func agentModelSupportsMultimodal(ctx context.Context, client *api.Client, modelName string) bool {
if client == nil || strings.TrimSpace(modelName) == "" {
return false
}
resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName})
if err != nil {
fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not check model capabilities: %v\n", err)
return false
}
if slices.Contains(resp.Capabilities, model.CapabilityVision) || slices.Contains(resp.Capabilities, model.CapabilityAudio) {
return true
}
if len(resp.ProjectorInfo) != 0 {
return true
}
for key := range resp.ModelInfo {
if strings.Contains(key, ".vision.") {
return true
}
}
return false
}
func agentCloudStatusDisabled(ctx context.Context, client *api.Client) (disabled bool, known bool) {
if internalcloud.Disabled() {
return true, true
}
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
}
// ensureCloudModelAccess checks whether the user is signed in and has a
// sufficient plan to use a cloud model. It returns an AuthorizationError
// (with SigninURL) when sign-in is needed, or a plan error when upgrade is
// needed.
func ensureCloudModelAccess(ctx context.Context, client *api.Client, model, requiredPlan string) error {
if client == nil {
return errors.New("no API client available")
}
if disabled, known := agentCloudStatusDisabled(ctx, client); known && disabled {
return errors.New("remote inference is unavailable")
}
user, err := client.Whoami(ctx)
if err != nil {
return err
}
if user != nil && user.Name != "" {
if requiredPlan != "" && !launch.PlanSatisfies(user.Plan, requiredPlan) {
return fmt.Errorf("plan upgrade required: %s needs plan %s, you have %s", model, requiredPlan, user.Plan)
}
return nil
}
return fmt.Errorf("%s requires sign in", model)
}
+481
View File
@@ -0,0 +1,481 @@
package cmd
import (
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"net/url"
"os"
"runtime"
"slices"
"strings"
"testing"
"time"
coreagent "github.com/ollama/ollama/agent"
"github.com/ollama/ollama/agent/skills"
agenttools "github.com/ollama/ollama/agent/tools"
"github.com/ollama/ollama/api"
agentchat "github.com/ollama/ollama/cmd/tui/chat"
"github.com/ollama/ollama/envconfig"
modelpkg "github.com/ollama/ollama/types/model"
)
func setAgentTUITestCloudEnabled(t *testing.T) {
t.Helper()
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("USERPROFILE", home)
t.Setenv("OLLAMA_NO_CLOUD", "")
envconfig.ReloadServerConfig()
}
func TestAgentSystemPromptIncludesModel(t *testing.T) {
prompt := agentSystemPromptAt(time.Date(2026, time.June, 12, 9, 30, 0, 0, time.UTC), "llama3.2", nil, false, "", "")
shellName := "bash"
if runtime.GOOS == "windows" {
shellName = "PowerShell"
}
for _, want := range []string{
"You are running in Ollama, in a harness to help the user accomplish tasks, and the model is llama3.2.",
"Current date: Friday, June 12, 2026.",
"Be concise, practical, and action-oriented.",
"Use " + shellName + " carefully.",
"Tell the user about meaningful changes",
} {
if !strings.Contains(prompt, want) {
t.Fatalf("prompt missing %q:\n%s", want, prompt)
}
}
}
func TestAgentSystemPromptIncludesModelSystem(t *testing.T) {
prompt := agentSystemPromptAt(time.Date(2026, time.June, 12, 9, 30, 0, 0, time.UTC), "llama3.2", nil, false, "You are a pirate.", "")
if !strings.Contains(prompt, "You are a pirate.") {
t.Fatalf("prompt missing model system:\n%s", prompt)
}
}
func TestAgentHeadlessEventSinkSuppressesThinking(t *testing.T) {
output := captureStdout(t, func() {
sink := &agentHeadlessEventSink{}
if err := sink.Emit(coreagent.Event{Type: coreagent.EventThinkingDelta, Thinking: "thinking"}); err != nil {
t.Fatal(err)
}
if err := sink.Emit(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "answer"}); err != nil {
t.Fatal(err)
}
})
if output != "answer" {
t.Fatalf("output = %q, want answer only", output)
}
}
func TestAgentHeadlessEventSinkPrintsOnlyFinishedToolEvents(t *testing.T) {
output := captureStderr(t, func() {
sink := &agentHeadlessEventSink{}
if err := sink.Emit(coreagent.Event{
Type: coreagent.EventToolStarted,
ToolName: "bash",
Args: map[string]any{"command": "pwd"},
}); err != nil {
t.Fatal(err)
}
if err := sink.Emit(coreagent.Event{
Type: coreagent.EventToolFinished,
Status: "done",
ToolName: "bash",
Args: map[string]any{"command": "pwd"},
}); err != nil {
t.Fatal(err)
}
if err := sink.Emit(coreagent.Event{
Type: coreagent.EventToolFinished,
Status: "denied",
ToolName: "edit",
Args: map[string]any{"path": "main.go"},
}); err != nil {
t.Fatal(err)
}
if err := sink.Emit(coreagent.Event{
Type: coreagent.EventToolFinished,
Status: "done",
ToolName: "web_fetch",
Args: map[string]any{"url": "https://example.com"},
Error: "timeout",
}); err != nil {
t.Fatal(err)
}
})
if strings.Contains(output, "in progress") {
t.Fatalf("headless output should not include in-progress tool events:\n%s", output)
}
for _, want := range []string{
`• Bash("pwd") done`,
`• Edit("main.go") failed`,
`• Web Fetch("https://example.com") failed`,
} {
if !strings.Contains(output, want) {
t.Fatalf("headless output missing %q:\n%s", want, output)
}
}
}
func TestAgentHeadlessEventSinkPrintsToolEventsAfterContentNewline(t *testing.T) {
var stdout string
stderr := captureStderr(t, func() {
stdout = captureStdout(t, func() {
sink := &agentHeadlessEventSink{}
if err := sink.Emit(coreagent.Event{Type: coreagent.EventMessageDelta, Content: "checking"}); err != nil {
t.Fatal(err)
}
if err := sink.Emit(coreagent.Event{
Type: coreagent.EventToolFinished,
Status: "done",
ToolName: "bash",
Args: map[string]any{"command": "pwd"},
}); err != nil {
t.Fatal(err)
}
})
})
if stdout != "checking\n" {
t.Fatalf("stdout = %q, want content newline before tool event", stdout)
}
if stderr != "• Bash(\"pwd\") done\n" {
t.Fatalf("stderr = %q, want compact tool status", stderr)
}
}
func captureStdout(t *testing.T, fn func()) string {
t.Helper()
oldStdout := os.Stdout
r, w, err := os.Pipe()
if err != nil {
t.Fatal(err)
}
os.Stdout = w
t.Cleanup(func() {
os.Stdout = oldStdout
})
fn()
if err := w.Close(); err != nil {
t.Fatal(err)
}
out, err := io.ReadAll(r)
if err != nil {
t.Fatal(err)
}
if err := r.Close(); err != nil {
t.Fatal(err)
}
os.Stdout = oldStdout
return string(out)
}
func captureStderr(t *testing.T, fn func()) string {
t.Helper()
oldStderr := os.Stderr
r, w, err := os.Pipe()
if err != nil {
t.Fatal(err)
}
os.Stderr = w
t.Cleanup(func() {
os.Stderr = oldStderr
})
fn()
if err := w.Close(); err != nil {
t.Fatal(err)
}
out, err := io.ReadAll(r)
if err != nil {
t.Fatal(err)
}
if err := r.Close(); err != nil {
t.Fatal(err)
}
os.Stderr = oldStderr
return string(out)
}
func TestAgentToolsRegistryNoCloudDisablesWebTools(t *testing.T) {
t.Setenv("OLLAMA_NO_CLOUD", "1")
t.Setenv("OLLAMA_AGENT_DISABLE_SHELL", "")
t.Setenv("OLLAMA_AGENT_DISABLE_WEBSEARCH", "")
statusCalls := 0
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/show":
_ = json.NewEncoder(w).Encode(api.ShowResponse{
Capabilities: []modelpkg.Capability{modelpkg.CapabilityTools},
})
case "/api/status":
statusCalls++
_ = json.NewEncoder(w).Encode(api.StatusResponse{})
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
baseURL, err := url.Parse(srv.URL)
if err != nil {
t.Fatal(err)
}
registry := agentToolsRegistry(context.Background(), api.NewClient(baseURL, srv.Client()), "test-model", nil)
if registry == nil {
t.Fatal("registry = nil, want local tools")
}
if !registry.Has(agenttools.NewBash().Name()) || !registry.Has("read") || !registry.Has("edit") {
t.Fatalf("local tools missing: %v", registry.Names())
}
if registry.Has("list") {
t.Fatalf("list tool should not be registered; got %v", registry.Names())
}
if registry.Has("web_search") || registry.Has("web_fetch") {
t.Fatalf("web tools should be disabled when OLLAMA_NO_CLOUD is set: %v", registry.Names())
}
if statusCalls != 0 {
t.Fatalf("/api/status calls = %d, want local no-cloud short-circuit", statusCalls)
}
}
func TestAgentToolsRegistryRegistersSkillTool(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/show":
_ = json.NewEncoder(w).Encode(api.ShowResponse{
Capabilities: []modelpkg.Capability{modelpkg.CapabilityTools},
})
case "/api/status":
_ = json.NewEncoder(w).Encode(api.StatusResponse{})
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
baseURL, err := url.Parse(srv.URL)
if err != nil {
t.Fatal(err)
}
catalog := &skills.Catalog{Skills: []skills.Skill{{Name: "go-code", Description: "Write Go code."}}}
registry := agentToolsRegistry(context.Background(), api.NewClient(baseURL, srv.Client()), "test-model", catalog)
if registry == nil || !registry.Has("skill") {
t.Fatalf("registry missing skill tool: %#v", registry)
}
}
func TestAgentModelOptionsIncludesCloudRecommendationsAndLocalModels(t *testing.T) {
setAgentTUITestCloudEnabled(t)
var recommendationsCalled bool
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/tags":
_ = json.NewEncoder(w).Encode(api.ListResponse{
Models: []api.ListModelResponse{{
Name: "llama3.2:latest",
Details: api.ModelDetails{
Family: "llama",
ParameterSize: "3B",
QuantizationLevel: "Q4_K_M",
ContextLength: 131072,
},
Size: 2_000_000_000,
}, {
Name: "gemma3:27b-cloud",
Details: api.ModelDetails{
Family: "gemma3",
ParameterSize: "27000000000",
QuantizationLevel: "bf16",
ContextLength: 131072,
},
}},
})
case "/api/status":
_ = json.NewEncoder(w).Encode(api.StatusResponse{})
case "/api/experimental/model-recommendations":
recommendationsCalled = true
_ = json.NewEncoder(w).Encode(api.ModelRecommendationsResponse{
Recommendations: []api.ModelRecommendation{
{Model: "qwen3.5:cloud", Description: "cloud reasoning", ContextLength: 262144, RequiredPlan: "pro"},
{Model: "gemma4", Description: "local recommendation should be ignored"},
},
})
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
baseURL, err := url.Parse(srv.URL)
if err != nil {
t.Fatal(err)
}
options, err := agentModelOptions(context.Background(), api.NewClient(baseURL, srv.Client()))
if err != nil {
t.Fatal(err)
}
if !recommendationsCalled {
t.Fatal("expected recommendations endpoint to be called")
}
if got, want := modelOptionNames(options), []string{"qwen3.5:cloud", "gemma3:27b-cloud", "llama3.2"}; !slices.Equal(got, want) {
t.Fatalf("model options = %#v, want %#v", got, want)
}
if options[0].Description == "" || options[1].Description == "" || options[2].Description == "" {
t.Fatalf("expected descriptions for model options: %#v", options)
}
// A cloud-sourced model in the local tags list keeps its ":cloud" suffix in
// the name, shows humanized arch details + ctx, and carries no "cloud" or
// "local" marker (the name conveys the cloud source).
cloudOpt := options[1]
if cloudOpt.Name != "gemma3:27b-cloud" {
t.Fatalf("cloud-tagged model name = %q, want :cloud suffix kept", cloudOpt.Name)
}
if strings.Contains(cloudOpt.Description, "cloud") || strings.Contains(cloudOpt.Description, "local") {
t.Fatalf("cloud-tagged model description = %q, should not include cloud/local marker", cloudOpt.Description)
}
if !strings.Contains(cloudOpt.Description, "27B") || !strings.Contains(cloudOpt.Description, "bf16") {
t.Fatalf("cloud-tagged model description = %q, want humanized params + quant", cloudOpt.Description)
}
if strings.Contains(cloudOpt.Description, "27000000000") {
t.Fatalf("cloud-tagged model description = %q, should not leak raw param size", cloudOpt.Description)
}
if !options[0].Recommended || options[1].Recommended {
t.Fatalf("recommended flags = %#v, want only cloud recommendation marked", options)
}
if strings.Contains(options[0].Description, "plan") || strings.Contains(options[0].Description, "pro") {
t.Fatalf("recommendation description should not include plan type: %q", options[0].Description)
}
}
func TestAgentModelOptionsNoCloudSkipsCloudRecommendations(t *testing.T) {
t.Setenv("OLLAMA_NO_CLOUD", "1")
var recommendationsCalled bool
var statusCalled bool
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/tags":
_ = json.NewEncoder(w).Encode(api.ListResponse{
Models: []api.ListModelResponse{{Name: "llama3.2:latest"}},
})
case "/api/status":
statusCalled = true
_ = json.NewEncoder(w).Encode(api.StatusResponse{})
case "/api/experimental/model-recommendations":
recommendationsCalled = true
_ = json.NewEncoder(w).Encode(api.ModelRecommendationsResponse{
Recommendations: []api.ModelRecommendation{{Model: "qwen3.5:cloud"}},
})
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
baseURL, err := url.Parse(srv.URL)
if err != nil {
t.Fatal(err)
}
options, err := agentModelOptions(context.Background(), api.NewClient(baseURL, srv.Client()))
if err != nil {
t.Fatal(err)
}
if recommendationsCalled {
t.Fatal("recommendations endpoint should not be called when no-cloud is set")
}
if statusCalled {
t.Fatal("status endpoint should not be called when local no-cloud short-circuits")
}
if got, want := modelOptionNames(options), []string{"llama3.2"}; !slices.Equal(got, want) {
t.Fatalf("model options = %#v, want %#v", got, want)
}
}
func TestPreloadAgentModelIfLocalLoadsLocalModel(t *testing.T) {
var generateReq api.GenerateRequest
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/show":
_ = json.NewEncoder(w).Encode(api.ShowResponse{})
case "/api/generate":
if err := json.NewDecoder(r.Body).Decode(&generateReq); err != nil {
t.Fatal(err)
}
_ = json.NewEncoder(w).Encode(api.GenerateResponse{Done: true})
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
baseURL, err := url.Parse(srv.URL)
if err != nil {
t.Fatal(err)
}
keepAlive := api.Duration{Duration: 5 * time.Minute}
think := &api.ThinkValue{Value: "low"}
err = preloadAgentModelIfLocal(context.Background(), api.NewClient(baseURL, srv.Client()), AgentTUIOptions{
KeepAlive: &keepAlive,
Think: think,
}, "llama3.2")
if err != nil {
t.Fatal(err)
}
if generateReq.Model != "llama3.2" {
t.Fatalf("generate model = %q, want llama3.2", generateReq.Model)
}
if generateReq.KeepAlive == nil || generateReq.KeepAlive.Duration != 5*time.Minute {
t.Fatalf("generate keepalive = %#v, want 5m", generateReq.KeepAlive)
}
if generateReq.Think == nil || generateReq.Think.String() != "low" {
t.Fatalf("generate think = %#v, want low", generateReq.Think)
}
}
func TestPreloadAgentModelIfLocalSkipsCloudModel(t *testing.T) {
generateCalled := false
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/show":
_ = json.NewEncoder(w).Encode(api.ShowResponse{RemoteHost: "https://ollama.com"})
case "/api/generate":
generateCalled = true
t.Fatal("cloud model should not be preloaded with generate")
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
baseURL, err := url.Parse(srv.URL)
if err != nil {
t.Fatal(err)
}
if err := preloadAgentModelIfLocal(context.Background(), api.NewClient(baseURL, srv.Client()), AgentTUIOptions{}, "kimi-k2:cloud"); err != nil {
t.Fatal(err)
}
if generateCalled {
t.Fatal("generate was called for cloud model")
}
}
func modelOptionNames(options []agentchat.ModelOption) []string {
names := make([]string, 0, len(options))
for _, option := range options {
names = append(names, option.Name)
}
return names
}
+603 -518
View File
File diff suppressed because it is too large. Load diff
+276
View File
@@ -2,6 +2,9 @@ package cmd
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/spf13/cobra"
@@ -9,12 +12,14 @@ import (
"github.com/ollama/ollama/cmd/config"
"github.com/ollama/ollama/cmd/launch"
"github.com/ollama/ollama/cmd/tui"
"github.com/ollama/ollama/envconfig"
)
func setCmdTestHome(t *testing.T, dir string) {
t.Helper()
t.Setenv("HOME", dir)
t.Setenv("USERPROFILE", dir)
envconfig.ReloadServerConfig()
}
func unexpectedRunModelResolution(t *testing.T) func(context.Context, launch.RunModelRequest) (string, error) {
@@ -41,6 +46,277 @@ func unexpectedModelLaunch(t *testing.T) func(*cobra.Command, string) error {
}
}
func TestRunAgentModelPickerUsesSavedModelWhenAvailable(t *testing.T) {
setCmdTestHome(t, t.TempDir())
if err := config.SetAgentSignInPromptSeen(true); err != nil {
t.Fatal(err)
}
var gotReq launch.RunModelRequest
var launched string
prefetchedAccount := &launch.AccountState{}
accountUpdates := func(context.Context) <-chan *launch.AccountState { return nil }
deps := agentModelPickerDeps{
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
gotReq = req
return "qwen3:8b", nil
},
runModel: func(cmd *cobra.Command, model string) error {
launched = model
return nil
},
accountState: func() *launch.AccountState {
return prefetchedAccount
},
accountStateUpdates: accountUpdates,
}
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
if err := runAgentModelPickerWithDeps(cmd, deps); err != nil {
t.Fatalf("runAgentModelPickerWithDeps error: %v", err)
}
if gotReq.ForcePicker {
t.Fatal("expected root agent flow to reuse a saved model when available")
}
if gotReq.AccountState != prefetchedAccount {
t.Fatal("expected prefetched account state to be passed to model picker")
}
if gotReq.AccountStateProvider == nil || gotReq.AccountStateUpdates == nil {
t.Fatal("expected account state callbacks to be passed to model picker")
}
if launched != "qwen3:8b" {
t.Fatalf("launched model = %q, want qwen3:8b", launched)
}
}
func TestRunAgentModelPickerFallsBackToPickerWhenPlanVerificationFails(t *testing.T) {
setCmdTestHome(t, t.TempDir())
var requests []launch.RunModelRequest
var launched string
deps := agentModelPickerDeps{
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
requests = append(requests, req)
if len(requests) == 1 {
return "", launch.ErrPlanVerificationUnavailable
}
return "llama3.2", nil
},
runModel: func(cmd *cobra.Command, model string) error {
launched = model
return nil
},
accountState: func() *launch.AccountState {
return &launch.AccountState{}
},
}
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
if err := runAgentModelPickerWithDeps(cmd, deps); err != nil {
t.Fatalf("runAgentModelPickerWithDeps error: %v", err)
}
if len(requests) != 2 {
t.Fatalf("resolve calls = %d, want 2", len(requests))
}
if requests[0].ForcePicker {
t.Fatal("first request should try the saved model path")
}
if !requests[1].ForcePicker {
t.Fatal("second request should force the model picker")
}
if requests[1].AccountStateProvider != nil {
t.Fatal("retry should not keep using the stale account-state provider")
}
if requests[1].AccountState == nil {
t.Fatal("retry should pass an explicit unknown account state")
}
if launched != "llama3.2" {
t.Fatalf("launched model = %q, want llama3.2", launched)
}
}
func TestRunAgentModelPickerReturnsPlanVerificationErrorWhenPickerRetryFails(t *testing.T) {
setCmdTestHome(t, t.TempDir())
var calls int
deps := agentModelPickerDeps{
resolveRunModel: func(ctx context.Context, req launch.RunModelRequest) (string, error) {
calls++
return "", launch.ErrPlanVerificationUnavailable
},
runModel: unexpectedModelLaunch(t),
}
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
err := runAgentModelPickerWithDeps(cmd, deps)
if !errors.Is(err, launch.ErrPlanVerificationUnavailable) {
t.Fatalf("error = %v, want ErrPlanVerificationUnavailable", err)
}
if calls != 2 {
t.Fatalf("resolve calls = %d, want 2", calls)
}
}
func TestMaybeRunAgentOnboarding(t *testing.T) {
t.Run("prompts once and saves seen state", func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
oldPrompt := agentOnboardingPrompt
oldSignedIn := agentOnboardingSignedInStatus
t.Cleanup(func() {
agentOnboardingPrompt = oldPrompt
agentOnboardingSignedInStatus = oldSignedIn
})
var prompts int
agentOnboardingPrompt = func() (bool, error) {
prompts++
return false, nil
}
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
return false, true
}
signIn, err := maybeRunAgentOnboarding(context.Background())
if err != nil {
t.Fatalf("maybeRunAgentOnboarding error: %v", err)
}
if signIn {
t.Fatal("signIn = true, want false")
}
if prompts != 1 {
t.Fatalf("prompts = %d, want 1", prompts)
}
if !config.AgentSignInPromptSeen() {
t.Fatal("expected onboarding state to be saved")
}
signIn, err = maybeRunAgentOnboarding(context.Background())
if err != nil {
t.Fatalf("second maybeRunAgentOnboarding error: %v", err)
}
if signIn {
t.Fatal("second signIn = true, want false")
}
if prompts != 1 {
t.Fatalf("prompt should not run again, prompts = %d", prompts)
}
})
t.Run("skips prompt when already signed in", func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
oldPrompt := agentOnboardingPrompt
oldSignedIn := agentOnboardingSignedInStatus
t.Cleanup(func() {
agentOnboardingPrompt = oldPrompt
agentOnboardingSignedInStatus = oldSignedIn
})
var prompts int
agentOnboardingPrompt = func() (bool, error) {
prompts++
return false, nil
}
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
return true, true
}
signIn, err := maybeRunAgentOnboarding(context.Background())
if err != nil {
t.Fatalf("maybeRunAgentOnboarding error: %v", err)
}
if signIn {
t.Fatal("signIn = true, want false")
}
if prompts != 0 {
t.Fatalf("prompts = %d, want 0", prompts)
}
if !config.AgentSignInPromptSeen() {
t.Fatal("expected onboarding state to be saved")
}
})
t.Run("skips prompt when signed-in check is unknown", func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
oldPrompt := agentOnboardingPrompt
oldSignedIn := agentOnboardingSignedInStatus
t.Cleanup(func() {
agentOnboardingPrompt = oldPrompt
agentOnboardingSignedInStatus = oldSignedIn
})
var prompts int
agentOnboardingPrompt = func() (bool, error) {
prompts++
return false, nil
}
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
return false, false
}
signIn, err := maybeRunAgentOnboarding(context.Background())
if err != nil {
t.Fatalf("maybeRunAgentOnboarding error: %v", err)
}
if signIn {
t.Fatal("signIn = true, want false")
}
if prompts != 0 {
t.Fatalf("prompts = %d, want 0", prompts)
}
if config.AgentSignInPromptSeen() {
t.Fatal("unknown auth state should not save onboarding state")
}
})
t.Run("cancel does not save seen state", func(t *testing.T) {
setCmdTestHome(t, t.TempDir())
oldPrompt := agentOnboardingPrompt
oldSignedIn := agentOnboardingSignedInStatus
t.Cleanup(func() {
agentOnboardingPrompt = oldPrompt
agentOnboardingSignedInStatus = oldSignedIn
})
agentOnboardingPrompt = func() (bool, error) {
return false, tui.ErrCancelled
}
agentOnboardingSignedInStatus = func(context.Context) (bool, bool) {
return false, true
}
_, err := maybeRunAgentOnboarding(context.Background())
if !errors.Is(err, launch.ErrCancelled) {
t.Fatalf("error = %v, want launch.ErrCancelled", err)
}
if config.AgentSignInPromptSeen() {
t.Fatal("cancel should not save onboarding state")
}
})
}
func TestRunAgentOnboardingSignInEmptyWhoamiDoesNotSilentlySucceed(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/me" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{}`))
}))
defer server.Close()
t.Setenv("OLLAMA_HOST", server.URL)
err := runAgentOnboardingSignIn(context.Background())
if !errors.Is(err, errAgentOnboardingNotSignedIn) {
t.Fatalf("error = %v, want errAgentOnboardingNotSignedIn", err)
}
}
func TestRunInteractiveTUI_RunModelActionsUseResolveRunModel(t *testing.T) {
tests := []struct {
name string
+901 -209
View File
File diff suppressed because it is too large. Load diff
+37
View File
@@ -19,6 +19,14 @@ type integration struct {
Onboarded bool `json:"onboarded,omitempty"`
}
type onboarding struct {
Agent *agentOnboarding `json:"agent,omitempty"`
}
type agentOnboarding struct {
SignInPromptSeen bool `json:"sign_in_prompt_seen,omitempty"`
}
// IntegrationConfig is the persisted config for one integration.
type IntegrationConfig = integration
@@ -26,6 +34,7 @@ type config struct {
Integrations map[string]*integration `json:"integrations"`
LastModel string `json:"last_model,omitempty"`
LastSelection string `json:"last_selection,omitempty"` // "run" or integration name
Onboarding *onboarding `json:"onboarding,omitempty"`
}
func configPath() (string, error) {
@@ -230,6 +239,34 @@ func SetLastSelection(selection string) error {
return save(cfg)
}
// AgentSignInPromptSeen reports whether the root agent sign-in onboarding prompt
// has already been shown.
func AgentSignInPromptSeen() bool {
cfg, err := load()
if err != nil {
return false
}
return cfg.Onboarding != nil &&
cfg.Onboarding.Agent != nil &&
cfg.Onboarding.Agent.SignInPromptSeen
}
// SetAgentSignInPromptSeen persists the root agent sign-in onboarding prompt state.
func SetAgentSignInPromptSeen(seen bool) error {
cfg, err := load()
if err != nil {
return err
}
if cfg.Onboarding == nil {
cfg.Onboarding = &onboarding{}
}
if cfg.Onboarding.Agent == nil {
cfg.Onboarding.Agent = &agentOnboarding{}
}
cfg.Onboarding.Agent.SignInPromptSeen = seen
return save(cfg)
}
// LoadIntegration returns the saved config for one integration.
func LoadIntegration(appName string) (*integration, error) {
cfg, err := load()
+30
View File
@@ -302,6 +302,36 @@ func TestLoad(t *testing.T) {
})
}
func TestAgentSignInPromptSeen(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
if AgentSignInPromptSeen() {
t.Fatal("new config should not have agent sign-in onboarding marked seen")
}
if err := SetAgentSignInPromptSeen(true); err != nil {
t.Fatal(err)
}
if !AgentSignInPromptSeen() {
t.Fatal("agent sign-in onboarding seen state was not saved")
}
path, err := configPath()
if err != nil {
t.Fatal(err)
}
data, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(data), `"onboarding"`) ||
!strings.Contains(string(data), `"agent"`) ||
!strings.Contains(string(data), `"sign_in_prompt_seen": true`) {
t.Fatalf("config does not include onboarding state: %s", data)
}
}
func TestMigrateConfig(t *testing.T) {
t.Run("migrates legacy file to new location", func(t *testing.T) {
tmpDir := t.TempDir()
-735
View File
@@ -1,735 +0,0 @@
package cmd
import (
"cmp"
"errors"
"fmt"
"io"
"net/http"
"os"
"os/exec"
"path/filepath"
"regexp"
"slices"
"strings"
"github.com/spf13/cobra"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/envconfig"
"github.com/ollama/ollama/internal/modelref"
"github.com/ollama/ollama/readline"
"github.com/ollama/ollama/types/errtypes"
"github.com/ollama/ollama/types/model"
)
type MultilineState int
const (
MultilineNone MultilineState = iota
MultilinePrompt
MultilineSystem
)
func generateInteractive(cmd *cobra.Command, opts runOptions) error {
usage := func() {
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 <model> Load a session or model")
fmt.Fprintln(os.Stderr, " /save <model> Save your current session")
fmt.Fprintln(os.Stderr, " /clear Clear session context")
fmt.Fprintln(os.Stderr, " /bye Exit")
fmt.Fprintln(os.Stderr, " /?, /help Help for a command")
fmt.Fprintln(os.Stderr, " /? shortcuts Help for keyboard shortcuts")
fmt.Fprintln(os.Stderr, "")
fmt.Fprintln(os.Stderr, "Use \"\"\" to begin a multi-line message.")
if opts.MultiModal {
fmt.Fprintf(os.Stderr, "Use %s to include .jpg, .png, .webp images, or .wav audio files.\n", filepath.FromSlash("/path/to/file"))
}
fmt.Fprintln(os.Stderr, "")
}
usageSet := func() {
fmt.Fprintln(os.Stderr, "Available Commands:")
fmt.Fprintln(os.Stderr, " /set parameter ... Set a parameter")
fmt.Fprintln(os.Stderr, " /set system <string> Set system message")
fmt.Fprintln(os.Stderr, " /set history Enable history")
fmt.Fprintln(os.Stderr, " /set nohistory Disable history")
fmt.Fprintln(os.Stderr, " /set wordwrap Enable wordwrap")
fmt.Fprintln(os.Stderr, " /set nowordwrap Disable wordwrap")
fmt.Fprintln(os.Stderr, " /set format json Enable JSON mode")
fmt.Fprintln(os.Stderr, " /set noformat Disable formatting")
fmt.Fprintln(os.Stderr, " /set verbose Show LLM stats")
fmt.Fprintln(os.Stderr, " /set quiet Disable LLM stats")
fmt.Fprintln(os.Stderr, " /set think Enable thinking")
fmt.Fprintln(os.Stderr, " /set nothink Disable thinking")
fmt.Fprintln(os.Stderr, "")
}
usageShortcuts := func() {
fmt.Fprintln(os.Stderr, "Available keyboard shortcuts:")
fmt.Fprintln(os.Stderr, " Ctrl + a Move to the beginning of the line (Home)")
fmt.Fprintln(os.Stderr, " Ctrl + e Move to the end of the line (End)")
fmt.Fprintln(os.Stderr, " Alt + b Move back (left) one word")
fmt.Fprintln(os.Stderr, " Alt + f Move forward (right) one word")
fmt.Fprintln(os.Stderr, " Ctrl + k Delete the sentence after the cursor")
fmt.Fprintln(os.Stderr, " Ctrl + u Delete the sentence before the cursor")
fmt.Fprintln(os.Stderr, " Ctrl + w Delete the word before the cursor")
fmt.Fprintln(os.Stderr, "")
fmt.Fprintln(os.Stderr, " Ctrl + l Clear the screen")
fmt.Fprintln(os.Stderr, " Ctrl + g Open default editor to compose a prompt")
fmt.Fprintln(os.Stderr, " Ctrl + c Stop the model from responding")
fmt.Fprintln(os.Stderr, " Ctrl + d Exit ollama (/bye)")
fmt.Fprintln(os.Stderr, "")
}
usageShow := func() {
fmt.Fprintln(os.Stderr, "Available Commands:")
fmt.Fprintln(os.Stderr, " /show info Show details for this model")
fmt.Fprintln(os.Stderr, " /show license Show model license")
fmt.Fprintln(os.Stderr, " /show modelfile Show Modelfile for this model")
fmt.Fprintln(os.Stderr, " /show parameters Show parameters for this model")
fmt.Fprintln(os.Stderr, " /show system Show system message")
fmt.Fprintln(os.Stderr, " /show template Show prompt template")
fmt.Fprintln(os.Stderr, "")
}
// only list out the most common parameters
usageParameters := func() {
fmt.Fprintln(os.Stderr, "Available Parameters:")
fmt.Fprintln(os.Stderr, " /set parameter seed <int> Random number seed")
fmt.Fprintln(os.Stderr, " /set parameter num_predict <int> Max number of tokens to predict")
fmt.Fprintln(os.Stderr, " /set parameter top_k <int> Pick from top k num of tokens")
fmt.Fprintln(os.Stderr, " /set parameter top_p <float> Pick token based on sum of probabilities")
fmt.Fprintln(os.Stderr, " /set parameter min_p <float> Pick token based on top token probability * min_p")
fmt.Fprintln(os.Stderr, " /set parameter num_ctx <int> Set the context size")
fmt.Fprintln(os.Stderr, " /set parameter temperature <float> Set creativity level")
fmt.Fprintln(os.Stderr, " /set parameter repeat_penalty <float> How strongly to penalize repetitions")
fmt.Fprintln(os.Stderr, " /set parameter repeat_last_n <int> Set how far back to look for repetitions")
fmt.Fprintln(os.Stderr, " /set parameter num_gpu <int> The number of layers to send to the GPU")
fmt.Fprintln(os.Stderr, " /set parameter stop <string> <string> ... Set the stop parameters")
fmt.Fprintln(os.Stderr, "")
}
scanner, err := readline.New(readline.Prompt{
Prompt: ">>> ",
AltPrompt: "... ",
Placeholder: "Send a message (/? for help)",
AltPlaceholder: "Press Enter to send",
})
if err != nil {
return err
}
if envconfig.NoHistory() {
scanner.HistoryDisable()
}
fmt.Print(readline.StartBracketedPaste)
defer fmt.Printf(readline.EndBracketedPaste)
var sb strings.Builder
var multiline MultilineState
var thinkExplicitlySet bool = opts.Think != nil
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 errors.Is(err, readline.ErrEditPrompt):
sb.Reset()
content, err := editInExternalEditor(line)
if err != nil {
fmt.Fprintf(os.Stderr, "error: %v\n", err)
continue
}
if strings.TrimSpace(content) == "" {
continue
}
scanner.Prefill = content
continue
case err != nil:
return err
}
switch {
case multiline != MultilineNone:
// check if there's a multiline terminating string
before, ok := strings.CutSuffix(line, `"""`)
sb.WriteString(before)
if !ok {
fmt.Fprintln(&sb)
scanner.Prompt.UseAlt = true
continue
}
switch multiline {
case MultilineSystem:
opts.System = sb.String()
opts.Messages = append(opts.Messages, api.Message{Role: "system", Content: opts.System})
fmt.Println("Set system message.")
sb.Reset()
}
multiline = MultilineNone
scanner.Prompt.UseAlt = false
case strings.HasPrefix(line, `"""`):
line := strings.TrimPrefix(line, `"""`)
line, ok := strings.CutSuffix(line, `"""`)
sb.WriteString(line)
if !ok {
// no multiline terminating string; need more input
fmt.Fprintln(&sb)
multiline = MultilinePrompt
scanner.Prompt.UseAlt = true
}
case scanner.Pasting:
fmt.Fprintln(&sb, line)
continue
case strings.HasPrefix(line, "/list"):
args := strings.Fields(line)
if err := ListHandler(cmd, args[1:]); err != nil {
return err
}
case strings.HasPrefix(line, "/load"):
args := strings.Fields(line)
if len(args) != 2 {
fmt.Println("Usage:\n /load <modelname>")
continue
}
origOpts := opts.Copy()
client, err := api.ClientFromEnvironment()
if err != nil {
fmt.Println("error: couldn't connect to ollama server")
return err
}
opts.Model = args[1]
opts.Messages = []api.Message{}
opts.LoadedMessages = nil
fmt.Printf("Loading model '%s'\n", opts.Model)
info, err := client.Show(cmd.Context(), &api.ShowRequest{Model: opts.Model})
if err != nil {
if strings.Contains(err.Error(), "not found") {
fmt.Printf("Couldn't find model '%s'\n", opts.Model)
opts = origOpts.Copy()
continue
}
return err
}
applyShowResponseToRunOptions(&opts, info)
opts.Think, err = inferThinkingOption(&info.Capabilities, &opts, thinkExplicitlySet)
if err != nil {
return err
}
if err := loadOrUnloadModel(cmd, &opts); err != nil {
if strings.Contains(err.Error(), "not found") {
fmt.Printf("Couldn't find model '%s'\n", opts.Model)
opts = origOpts.Copy()
continue
}
if strings.Contains(err.Error(), "does not support thinking") {
fmt.Printf("error: %v\n", err)
continue
}
return err
}
continue
case strings.HasPrefix(line, "/save"):
args := strings.Fields(line)
if len(args) != 2 {
fmt.Println("Usage:\n /save <modelname>")
continue
}
client, err := api.ClientFromEnvironment()
if err != nil {
fmt.Println("error: couldn't connect to ollama server")
return err
}
req := NewCreateRequest(args[1], opts)
fn := func(resp api.ProgressResponse) error { return nil }
err = client.Create(cmd.Context(), req, fn)
if err != nil {
if strings.Contains(err.Error(), errtypes.InvalidModelNameErrMsg) {
fmt.Printf("error: The model name '%s' is invalid\n", args[1])
continue
}
return err
}
fmt.Printf("Created new model '%s'\n", args[1])
continue
case strings.HasPrefix(line, "/clear"):
opts.Messages = []api.Message{}
if opts.System != "" {
newMessage := api.Message{Role: "system", Content: opts.System}
opts.Messages = append(opts.Messages, newMessage)
}
fmt.Println("Cleared session context")
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":
opts.WordWrap = true
fmt.Println("Set 'wordwrap' mode.")
case "nowordwrap":
opts.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 != "" {
// TODO(drifkin): validate the level, could be model dependent
// though... It will also be validated on the server once a call is
// made.
thinkValue.Value = maybeLevel
}
opts.Think = &thinkValue
thinkExplicitlySet = true
if client, err := api.ClientFromEnvironment(); err == nil {
ensureThinkingSupport(cmd.Context(), client, opts.Model)
}
if maybeLevel != "" {
fmt.Printf("Set 'think' mode to '%s'.\n", maybeLevel)
} else {
fmt.Println("Set 'think' mode.")
}
case "nothink":
opts.Think = &api.ThinkValue{Value: false}
thinkExplicitlySet = true
if client, err := api.ClientFromEnvironment(); err == nil {
ensureThinkingSupport(cmd.Context(), client, opts.Model)
}
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 {
opts.Format = args[2]
fmt.Printf("Set format to '%s' mode.\n", args[2])
}
case "noformat":
opts.Format = ""
fmt.Println("Disabled format.")
case "parameter":
if len(args) < 4 {
usageParameters()
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, ", "))
opts.Options[args[2]] = fp[args[2]]
case "system":
if len(args) < 3 {
usageSet()
continue
}
multiline = MultilineSystem
line := strings.Join(args[2:], " ")
line, ok := strings.CutPrefix(line, `"""`)
if !ok {
multiline = MultilineNone
} else {
// only cut suffix if the line is multiline
line, ok = strings.CutSuffix(line, `"""`)
if ok {
multiline = MultilineNone
}
}
sb.WriteString(line)
if multiline != MultilineNone {
scanner.Prompt.UseAlt = true
continue
}
opts.System = sb.String() // for display in modelfile
newMessage := api.Message{Role: "system", Content: sb.String()}
// Check if the slice is not empty and the last message is from 'system'
if len(opts.Messages) > 0 && opts.Messages[len(opts.Messages)-1].Role == "system" {
// Replace the last message
opts.Messages[len(opts.Messages)-1] = newMessage
} else {
opts.Messages = append(opts.Messages, newMessage)
}
fmt.Println("Set system message.")
sb.Reset()
continue
default:
fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1])
}
} else {
usageSet()
}
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")
return err
}
req := &api.ShowRequest{
Name: opts.Model,
System: opts.System,
Options: opts.Options,
}
resp, err := client.Show(cmd.Context(), req)
if err != nil {
fmt.Println("error: couldn't get model")
return err
}
switch args[1] {
case "info":
_ = showInfo(resp, false, 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 for this model.")
} else {
for _, l := range strings.Split(resp.Parameters, "\n") {
fmt.Printf(" %s\n", l)
}
}
fmt.Println()
if len(opts.Options) > 0 {
fmt.Println("User defined parameters:")
for k, v := range opts.Options {
fmt.Printf(" %-*s %v\n", 30, k, v)
}
fmt.Println()
}
case "system":
switch {
case opts.System != "":
fmt.Println(opts.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 {
usageShow()
}
case strings.HasPrefix(line, "/help"), strings.HasPrefix(line, "/?"):
args := strings.Fields(line)
if len(args) > 1 {
switch args[1] {
case "set", "/set":
usageSet()
case "show", "/show":
usageShow()
case "shortcut", "shortcuts":
usageShortcuts()
}
} else {
usage()
}
case strings.HasPrefix(line, "/exit"), strings.HasPrefix(line, "/bye"):
return nil
case strings.HasPrefix(line, "/"):
args := strings.Fields(line)
isFile := false
if opts.MultiModal {
for _, f := range extractFileNames(line) {
if strings.HasPrefix(f, args[0]) {
isFile = true
break
}
}
}
if !isFile {
fmt.Printf("Unknown command '%s'. Type /? for help\n", args[0])
continue
}
sb.WriteString(line)
default:
sb.WriteString(line)
}
if sb.Len() > 0 && multiline == MultilineNone {
newMessage := api.Message{Role: "user", Content: sb.String()}
if opts.MultiModal {
msg, images, err := extractFileData(sb.String())
if err != nil {
return err
}
newMessage.Content = msg
newMessage.Images = images
}
opts.Messages = append(opts.Messages, newMessage)
assistant, err := chat(cmd, opts)
if err != nil {
if strings.Contains(err.Error(), "does not support thinking") ||
strings.Contains(err.Error(), "invalid think value") {
fmt.Printf("error: %v\n", err)
sb.Reset()
continue
}
return err
}
if assistant != nil {
opts.Messages = append(opts.Messages, *assistant)
}
sb.Reset()
}
}
}
func NewCreateRequest(name string, opts runOptions) *api.CreateRequest {
parentModel := opts.ParentModel
modelName := model.ParseName(parentModel)
if !modelName.IsValid() {
parentModel = ""
}
// Preserve explicit cloud intent for sessions started with `:cloud`.
// Cloud model metadata can return a source-less parent_model (for example
// "qwen3.5"), which would otherwise make `/save` create a local derivative.
if modelref.HasExplicitCloudSource(opts.Model) && !modelref.HasExplicitCloudSource(parentModel) {
parentModel = ""
}
req := &api.CreateRequest{
Model: name,
From: cmp.Or(parentModel, opts.Model),
}
if opts.System != "" {
req.System = opts.System
}
if len(opts.Options) > 0 {
req.Parameters = opts.Options
}
messages := slices.Clone(opts.LoadedMessages)
messages = append(messages, opts.Messages...)
if len(messages) > 0 {
req.Messages = messages
}
return req
}
func normalizeFilePath(fp string) string {
return strings.NewReplacer(
"\\ ", " ", // Escaped space
"\\(", "(", // Escaped left parenthesis
"\\)", ")", // Escaped right parenthesis
"\\[", "[", // Escaped left square bracket
"\\]", "]", // Escaped right square bracket
"\\{", "{", // Escaped left curly brace
"\\}", "}", // Escaped right curly brace
"\\$", "$", // Escaped dollar sign
"\\&", "&", // Escaped ampersand
"\\;", ";", // Escaped semicolon
"\\'", "'", // Escaped single quote
"\\\\", "\\", // Escaped backslash
"\\*", "*", // Escaped asterisk
"\\?", "?", // Escaped question mark
"\\~", "~", // Escaped tilde
).Replace(fp)
}
func extractFileNames(input string) []string {
// Regex to match file paths starting with optional drive letter, / ./ \ or .\ and include escaped or unescaped spaces (\ or %20)
// and followed by more characters and a file extension
// This will capture non filename strings, but we'll check for file existence to remove mismatches
regexPattern := `(?:[a-zA-Z]:)?(?:\./|/|\\)[\S\\ ]+?\.(?i:jpg|jpeg|png|webp|wav)\b`
re := regexp.MustCompile(regexPattern)
return re.FindAllString(input, -1)
}
func extractFileData(input string) (string, []api.ImageData, error) {
filePaths := extractFileNames(input)
var imgs []api.ImageData
for _, fp := range filePaths {
nfp := normalizeFilePath(fp)
data, err := getImageData(nfp)
if errors.Is(err, os.ErrNotExist) {
continue
} else if err != nil {
fmt.Fprintf(os.Stderr, "Couldn't process file: %q\n", err)
return "", imgs, err
}
ext := strings.ToLower(filepath.Ext(nfp))
switch ext {
case ".wav":
fmt.Fprintf(os.Stderr, "Added audio '%s'\n", nfp)
default:
fmt.Fprintf(os.Stderr, "Added image '%s'\n", nfp)
}
input = strings.ReplaceAll(input, "'"+nfp+"'", "")
input = strings.ReplaceAll(input, "'"+fp+"'", "")
input = strings.ReplaceAll(input, fp, "")
imgs = append(imgs, data)
}
return strings.TrimSpace(input), imgs, nil
}
func editInExternalEditor(content string) (string, error) {
editor := envconfig.Editor()
if editor == "" {
editor = os.Getenv("VISUAL")
}
if editor == "" {
editor = os.Getenv("EDITOR")
}
if editor == "" {
editor = defaultEditor
}
// Check that the editor binary exists
name := strings.Fields(editor)[0]
if _, err := exec.LookPath(name); err != nil {
return "", fmt.Errorf("editor %q not found, set OLLAMA_EDITOR to the path of your preferred editor", name)
}
tmpFile, err := os.CreateTemp("", "ollama-prompt-*.txt")
if err != nil {
return "", fmt.Errorf("creating temp file: %w", err)
}
defer os.Remove(tmpFile.Name())
if content != "" {
if _, err := tmpFile.WriteString(content); err != nil {
tmpFile.Close()
return "", fmt.Errorf("writing to temp file: %w", err)
}
}
tmpFile.Close()
args := strings.Fields(editor)
args = append(args, tmpFile.Name())
cmd := exec.Command(args[0], args[1:]...)
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
if err := cmd.Run(); err != nil {
return "", fmt.Errorf("editor exited with error: %w", err)
}
data, err := os.ReadFile(tmpFile.Name())
if err != nil {
return "", fmt.Errorf("reading temp file: %w", err)
}
return strings.TrimRight(string(data), "\n"), nil
}
func getImageData(filePath string) ([]byte, error) {
file, err := os.Open(filePath)
if err != nil {
return nil, err
}
defer file.Close()
buf := make([]byte, 512)
_, err = file.Read(buf)
if err != nil {
return nil, err
}
contentType := http.DetectContentType(buf)
allowedTypes := []string{"image/jpeg", "image/jpg", "image/png", "image/webp", "audio/wave"}
if !slices.Contains(allowedTypes, contentType) {
return nil, fmt.Errorf("invalid file type: %s", contentType)
}
info, err := file.Stat()
if err != nil {
return nil, err
}
var maxSize int64 = 100 * 1024 * 1024 // 100MB
if info.Size() > maxSize {
return nil, errors.New("file size exceeds maximum limit (100MB)")
}
buf = make([]byte, info.Size())
_, err = file.Seek(0, 0)
if err != nil {
return nil, err
}
_, err = io.ReadFull(file, buf)
if err != nil {
return nil, err
}
return buf, nil
}
-116
View File
@@ -1,116 +0,0 @@
package cmd
import (
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
)
func TestExtractFilenames(t *testing.T) {
// Unix style paths
input := ` some preamble
./relative\ path/one.png inbetween1 ./not a valid two.jpg inbetween2 ./1.svg
/unescaped space /three.jpeg inbetween3 /valid\ path/dir/four.png "./quoted with spaces/five.JPG
/unescaped space /six.webp inbetween6 /valid\ path/dir/seven.WEBP`
res := extractFileNames(input)
assert.Len(t, res, 7)
assert.Contains(t, res[0], "one.png")
assert.Contains(t, res[1], "two.jpg")
assert.Contains(t, res[2], "three.jpeg")
assert.Contains(t, res[3], "four.png")
assert.Contains(t, res[4], "five.JPG")
assert.Contains(t, res[5], "six.webp")
assert.Contains(t, res[6], "seven.WEBP")
assert.NotContains(t, res[4], '"')
assert.NotContains(t, res, "inbetween1")
assert.NotContains(t, res, "./1.svg")
// Windows style paths
input = ` some preamble
c:/users/jdoe/one.png inbetween1 c:/program files/someplace/two.jpg inbetween2
/absolute/nospace/three.jpeg inbetween3 /absolute/with space/four.png inbetween4
./relative\ path/five.JPG inbetween5 "./relative with/spaces/six.png inbetween6
d:\path with\spaces\seven.JPEG inbetween7 c:\users\jdoe\eight.png inbetween8
d:\program files\someplace\nine.png inbetween9 "E:\program files\someplace\ten.PNG
c:/users/jdoe/eleven.webp inbetween11 c:/program files/someplace/twelve.WebP inbetween12
d:\path with\spaces\thirteen.WEBP some ending
`
res = extractFileNames(input)
assert.Len(t, res, 13)
assert.NotContains(t, res, "inbetween2")
assert.Contains(t, res[0], "one.png")
assert.Contains(t, res[0], "c:")
assert.Contains(t, res[1], "two.jpg")
assert.Contains(t, res[1], "c:")
assert.Contains(t, res[2], "three.jpeg")
assert.Contains(t, res[3], "four.png")
assert.Contains(t, res[4], "five.JPG")
assert.Contains(t, res[5], "six.png")
assert.Contains(t, res[6], "seven.JPEG")
assert.Contains(t, res[6], "d:")
assert.Contains(t, res[7], "eight.png")
assert.Contains(t, res[7], "c:")
assert.Contains(t, res[8], "nine.png")
assert.Contains(t, res[8], "d:")
assert.Contains(t, res[9], "ten.PNG")
assert.Contains(t, res[9], "E:")
assert.Contains(t, res[10], "eleven.webp")
assert.Contains(t, res[10], "c:")
assert.Contains(t, res[11], "twelve.WebP")
assert.Contains(t, res[11], "c:")
assert.Contains(t, res[12], "thirteen.WEBP")
assert.Contains(t, res[12], "d:")
}
// Ensure that file paths wrapped in single quotes are removed with the quotes.
func TestExtractFileDataRemovesQuotedFilepath(t *testing.T) {
dir := t.TempDir()
fp := filepath.Join(dir, "img.jpg")
data := make([]byte, 600)
copy(data, []byte{
0xff, 0xd8, 0xff, 0xe0, 0x00, 0x10, 'J', 'F', 'I', 'F',
0x00, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0xff, 0xd9,
})
if err := os.WriteFile(fp, data, 0o600); err != nil {
t.Fatalf("failed to write test image: %v", err)
}
input := "before '" + fp + "' after"
cleaned, imgs, err := extractFileData(input)
assert.NoError(t, err)
assert.Len(t, imgs, 1)
assert.Equal(t, cleaned, "before after")
}
func TestExtractFileDataWAV(t *testing.T) {
dir := t.TempDir()
fp := filepath.Join(dir, "sample.wav")
data := make([]byte, 600)
copy(data[:44], []byte{
'R', 'I', 'F', 'F',
0x58, 0x02, 0x00, 0x00, // file size - 8
'W', 'A', 'V', 'E',
'f', 'm', 't', ' ',
0x10, 0x00, 0x00, 0x00, // fmt chunk size
0x01, 0x00, // PCM
0x01, 0x00, // mono
0x80, 0x3e, 0x00, 0x00, // 16000 Hz
0x00, 0x7d, 0x00, 0x00, // byte rate
0x02, 0x00, // block align
0x10, 0x00, // 16-bit
'd', 'a', 't', 'a',
0x34, 0x02, 0x00, 0x00, // data size
})
if err := os.WriteFile(fp, data, 0o600); err != nil {
t.Fatalf("failed to write test audio: %v", err)
}
input := "before " + fp + " after"
cleaned, imgs, err := extractFileData(input)
assert.NoError(t, err)
assert.Len(t, imgs, 1)
assert.Equal(t, "before after", cleaned)
}
+186
View File
@@ -0,0 +1,186 @@
package filedata
import (
"errors"
"fmt"
"io"
"net/http"
"net/url"
"os"
"path/filepath"
"regexp"
"slices"
"strings"
"github.com/ollama/ollama/api"
)
type File struct {
Path string
Data api.ImageData
}
func NormalizePath(fp string) string {
fp = strings.Trim(fp, "\"")
fp = strings.NewReplacer(
"\\ ", " ",
"\\(", "(",
"\\)", ")",
"\\[", "[",
"\\]", "]",
"\\{", "{",
"\\}", "}",
"\\$", "$",
"\\&", "&",
"\\;", ";",
"\\'", "'",
"\\\\", "\\",
"\\*", "*",
"\\?", "?",
"\\~", "~",
).Replace(fp)
if u, err := url.Parse(fp); err == nil && strings.EqualFold(u.Scheme, "file") {
return normalizeFileURL(u)
} else if normalized, ok := normalizeMalformedFileURL(fp); ok {
return normalized
}
return fp
}
func ExtractNames(input string) []string {
regexPattern := `(?:file://\S+?\.(?i:jpg|jpeg|png|webp|wav)\b)|(?:(?:[a-zA-Z]:)?(?:\./|\.\\|/|\\)[\S\\ ]+?\.(?i:jpg|jpeg|png|webp|wav)\b)`
re := regexp.MustCompile(regexPattern)
return re.FindAllString(input, -1)
}
func Extract(input string) (string, []api.ImageData, error) {
cleaned, files, err := ExtractWithFiles(input)
if err != nil {
return "", nil, err
}
data := make([]api.ImageData, 0, len(files))
for _, file := range files {
data = append(data, file.Data)
}
return cleaned, data, nil
}
func ExtractWithFiles(input string) (string, []File, error) {
filePaths := ExtractNames(input)
var files []File
for _, fp := range filePaths {
nfp := NormalizePath(fp)
data, err := GetData(nfp)
if errors.Is(err, os.ErrNotExist) {
continue
} else if err != nil {
return "", files, fmt.Errorf("couldn't process file %q: %w", nfp, err)
}
input = strings.ReplaceAll(input, "'"+nfp+"'", "")
input = strings.ReplaceAll(input, "'"+fp+"'", "")
input = strings.ReplaceAll(input, `"`+nfp+`"`, "")
input = strings.ReplaceAll(input, `"`+fp+`"`, "")
input = strings.ReplaceAll(input, fp, "")
files = append(files, File{Path: nfp, Data: data})
}
return strings.TrimSpace(input), files, nil
}
func GetData(filePath string) ([]byte, error) {
file, err := os.Open(filePath)
if err != nil {
return nil, err
}
defer file.Close()
buf := make([]byte, 512)
_, err = file.Read(buf)
if err != nil {
return nil, err
}
contentType := http.DetectContentType(buf)
allowedTypes := []string{"image/jpeg", "image/jpg", "image/png", "image/webp", "audio/wave"}
if !slices.Contains(allowedTypes, contentType) {
return nil, fmt.Errorf("invalid file type: %s", contentType)
}
info, err := file.Stat()
if err != nil {
return nil, err
}
var maxSize int64 = 100 * 1024 * 1024
if info.Size() > maxSize {
return nil, errors.New("file size exceeds maximum limit (100MB)")
}
buf = make([]byte, info.Size())
_, err = file.Seek(0, 0)
if err != nil {
return nil, err
}
_, err = io.ReadFull(file, buf)
if err != nil {
return nil, err
}
return buf, nil
}
func Kind(path string) string {
if strings.EqualFold(filepath.Ext(path), ".wav") {
return "audio"
}
return "image"
}
func normalizeFileURL(u *url.URL) string {
path := u.Path
if unescaped, err := url.PathUnescape(path); err == nil {
path = unescaped
}
host := u.Host
if unescaped, err := url.PathUnescape(host); err == nil {
host = unescaped
}
if len(host) >= 2 && host[1] == ':' && isASCIIAlpha(host[0]) {
return filepath.Clean(filepath.FromSlash(host + path))
}
if len(path) >= 4 && path[0] == '/' && path[2] == ':' && isASCIIAlpha(path[1]) {
path = path[1:]
}
if u.Host != "" && !strings.EqualFold(u.Host, "localhost") {
return `\\` + u.Host + filepath.FromSlash(path)
}
return filepath.FromSlash(path)
}
func normalizeMalformedFileURL(raw string) (string, bool) {
const prefix = "file://"
if !strings.HasPrefix(strings.ToLower(raw), prefix) {
return "", false
}
path := raw[len(prefix):]
if unescaped, err := url.PathUnescape(path); err == nil {
path = unescaped
}
path = strings.TrimPrefix(path, "localhost")
if len(path) >= 3 && path[0] == '/' && path[2] == ':' && isASCIIAlpha(path[1]) {
path = path[1:]
}
if len(path) >= 2 && path[1] == ':' && isASCIIAlpha(path[0]) {
return filepath.Clean(filepath.FromSlash(path)), true
}
return "", false
}
func isASCIIAlpha(b byte) bool {
return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z')
}
+223
View File
@@ -0,0 +1,223 @@
package filedata
import (
"net/url"
"os"
"path/filepath"
"strings"
"testing"
)
func TestNormalizePathMalformedWindowsFileURL(t *testing.T) {
got := NormalizePath(`file://C:%5CUsers%5Cjdoe%5CPictures%5Cimg.png`)
want := filepath.Clean(`C:\Users\jdoe\Pictures\img.png`)
if got != want {
t.Fatalf("path = %q, want %q", got, want)
}
}
func TestNormalizePathTwoSlashWindowsFileURL(t *testing.T) {
got := NormalizePath(`file://C:/Users/jdoe/Pictures/img.png`)
want := filepath.Clean(`C:/Users/jdoe/Pictures/img.png`)
if got != want {
t.Fatalf("path = %q, want %q", got, want)
}
}
func TestNormalizePathLocalhostWindowsFileURL(t *testing.T) {
got := NormalizePath(`file://localhost/C:/Users/jdoe/Pictures/img.png`)
want := filepath.Clean(`C:/Users/jdoe/Pictures/img.png`)
if got != want {
t.Fatalf("path = %q, want %q", got, want)
}
}
func TestExtractNames(t *testing.T) {
// Unix style paths
input := ` some preamble
./relative\ path/one.png inbetween1 ./not a valid two.jpg inbetween2 ./1.svg
/unescaped space /three.jpeg inbetween3 /valid\ path/dir/four.png "./quoted with spaces/five.JPG
/unescaped space /six.webp inbetween6 /valid\ path/dir/seven.WEBP`
res := ExtractNames(input)
if len(res) != 7 {
t.Fatalf("len = %d, want 7", len(res))
}
assertContains(t, res[0], "one.png")
assertContains(t, res[1], "two.jpg")
assertContains(t, res[2], "three.jpeg")
assertContains(t, res[3], "four.png")
assertContains(t, res[4], "five.JPG")
assertContains(t, res[5], "six.webp")
assertContains(t, res[6], "seven.WEBP")
assertNotContains(t, res[4], "\"")
for _, r := range res {
assertNotContains(t, r, "inbetween1")
}
assertNotContainsSlice(t, res, "./1.svg")
}
func TestExtractNamesWindowsPaths(t *testing.T) {
input := ` some preamble
c:/users/jdoe/one.png inbetween1 c:/program files/someplace/two.jpg inbetween2
/absolute/nospace/three.jpeg inbetween3 /absolute/with space/four.png inbetween4
./relative\ path/five.JPG inbetween5 "./relative with/spaces/six.png inbetween6
d:\path with\spaces\seven.JPEG inbetween7 c:\users\jdoe\eight.png inbetween8
d:\program files\someplace\nine.png inbetween9 "E:\program files\someplace\ten.PNG
c:/users/jdoe/eleven.webp inbetween11 c:/program files/someplace/twelve.WebP inbetween12
d:\path with\spaces\thirteen.WEBP some ending
`
res := ExtractNames(input)
if len(res) != 13 {
t.Fatalf("len = %d, want 13", len(res))
}
assertNotContainsSlice(t, res, "inbetween2")
assertContains(t, res[0], "one.png")
assertContains(t, res[0], "c:")
assertContains(t, res[1], "two.jpg")
assertContains(t, res[1], "c:")
assertContains(t, res[2], "three.jpeg")
assertContains(t, res[3], "four.png")
assertContains(t, res[4], "five.JPG")
assertContains(t, res[5], "six.png")
assertContains(t, res[6], "seven.JPEG")
assertContains(t, res[6], "d:")
assertContains(t, res[7], "eight.png")
assertContains(t, res[7], "c:")
assertContains(t, res[8], "nine.png")
assertContains(t, res[8], "d:")
assertContains(t, res[9], "ten.PNG")
assertContains(t, res[9], "E:")
assertContains(t, res[10], "eleven.webp")
assertContains(t, res[10], "c:")
assertContains(t, res[11], "twelve.WebP")
assertContains(t, res[11], "c:")
assertContains(t, res[12], "thirteen.WEBP")
assertContains(t, res[12], "d:")
}
func TestExtractNamesDragDropPaths(t *testing.T) {
input := `file:///Users/jdoe/Pictures/one.png file://localhost/C:/Users/jdoe/Pictures/two.webp file:///C:/Users/jdoe/Pictures/three.jpg .\relative\four.png`
res := ExtractNames(input)
if len(res) != 4 {
t.Fatalf("len = %d, want 4", len(res))
}
assertContains(t, res[0], "file:///Users/jdoe/Pictures/one.png")
assertContains(t, res[1], "file://localhost/C:/Users/jdoe/Pictures/two.webp")
assertContains(t, res[2], "file:///C:/Users/jdoe/Pictures/three.jpg")
assertContains(t, res[3], `.\relative\four.png`)
}
func TestNormalizePathFileURL(t *testing.T) {
got := NormalizePath("file:///C:/Users/jdoe/Pictures/img.png")
want := filepath.FromSlash("C:/Users/jdoe/Pictures/img.png")
if got != want {
t.Fatalf("path = %q, want %q", got, want)
}
}
func TestExtractRemovesQuotedFilepath(t *testing.T) {
dir := t.TempDir()
fp := filepath.Join(dir, "img.jpg")
data := make([]byte, 600)
copy(data, []byte{
0xff, 0xd8, 0xff, 0xe0, 0x00, 0x10, 'J', 'F', 'I', 'F',
0x00, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0xff, 0xd9,
})
if err := os.WriteFile(fp, data, 0o600); err != nil {
t.Fatalf("failed to write test image: %v", err)
}
input := "before '" + fp + "' after"
cleaned, imgs, err := Extract(input)
if err != nil {
t.Fatalf("err: %v", err)
}
if len(imgs) != 1 {
t.Fatalf("imgs = %d, want 1", len(imgs))
}
if cleaned != "before after" {
t.Fatalf("cleaned = %q, want %q", cleaned, "before after")
}
}
func TestExtractFileURL(t *testing.T) {
dir := t.TempDir()
fp := filepath.Join(dir, "img.png")
data := make([]byte, 600)
copy(data, []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'})
if err := os.WriteFile(fp, data, 0o600); err != nil {
t.Fatalf("failed to write test image: %v", err)
}
fileURL := (&url.URL{Scheme: "file", Path: fp}).String()
cleaned, imgs, err := Extract("before " + fileURL + " after")
if err != nil {
t.Fatalf("err: %v", err)
}
if len(imgs) != 1 {
t.Fatalf("imgs = %d, want 1", len(imgs))
}
if cleaned != "before after" {
t.Fatalf("cleaned = %q, want %q", cleaned, "before after")
}
}
func TestExtractWAV(t *testing.T) {
dir := t.TempDir()
fp := filepath.Join(dir, "sample.wav")
data := make([]byte, 600)
copy(data[:44], []byte{
'R', 'I', 'F', 'F',
0x58, 0x02, 0x00, 0x00,
'W', 'A', 'V', 'E',
'f', 'm', 't', ' ',
0x10, 0x00, 0x00, 0x00,
0x01, 0x00,
0x01, 0x00,
0x80, 0x3e, 0x00, 0x00,
0x00, 0x7d, 0x00, 0x00,
0x02, 0x00,
0x10, 0x00,
'd', 'a', 't', 'a',
0x34, 0x02, 0x00, 0x00,
})
if err := os.WriteFile(fp, data, 0o600); err != nil {
t.Fatalf("failed to write test audio: %v", err)
}
input := "before " + fp + " after"
cleaned, imgs, err := Extract(input)
if err != nil {
t.Fatalf("err: %v", err)
}
if len(imgs) != 1 {
t.Fatalf("imgs = %d, want 1", len(imgs))
}
if cleaned != "before after" {
t.Fatalf("cleaned = %q, want %q", cleaned, "before after")
}
}
func assertContains(t *testing.T, s, want string) {
t.Helper()
if !strings.Contains(s, want) {
t.Fatalf("%q does not contain %q", s, want)
}
}
func assertNotContains(t *testing.T, s, want string) {
t.Helper()
if strings.Contains(s, want) {
t.Fatalf("%q unexpectedly contains %q", s, want)
}
}
func assertNotContainsSlice(t *testing.T, ss []string, want string) {
t.Helper()
for _, s := range ss {
if strings.Contains(s, want) {
t.Fatalf("slice unexpectedly contains %q in %q", want, s)
}
}
}
+1 -1
View File
@@ -713,7 +713,7 @@ func (c *launcherClient) resolveRunModel(ctx context.Context, req RunModelReques
}
}
model, err := c.selectSingleModelWithSelector(ctx, "Select model to run:", current, DefaultSingleSelector)
model, err := c.selectSingleModelWithSelector(ctx, "Select model to chat and code with:", current, DefaultSingleSelector)
if err != nil {
return "", err
}
+3 -2
View File
@@ -29,8 +29,9 @@ var DefaultConfirmPrompt func(prompt string, options ConfirmOptions) (bool, erro
// ConfirmOptions customizes labels for confirmation prompts.
type ConfirmOptions struct {
YesLabel string
NoLabel string
YesLabel string
NoLabel string
PlainPrompt bool
}
// SingleSelector is a function type for single item selection.
+324
View File
@@ -0,0 +1,324 @@
package chat
import (
"context"
"fmt"
"strings"
"time"
tea "github.com/charmbracelet/bubbletea"
coreagent "github.com/ollama/ollama/agent"
)
type chatApprovalChoice struct {
label string
key string
decision coreagent.ApprovalDecision
reason string
}
var chatApprovalChoices = []chatApprovalChoice{
{label: "Approve once", key: "1", decision: coreagent.ApprovalAllowOnce},
{label: "Approve session", key: "2", decision: coreagent.ApprovalAllowSession},
{label: "Deny", key: "3", decision: coreagent.ApprovalDeny, reason: "Tool execution denied."},
}
type chatApprovalPrompt struct {
request coreagent.ApprovalRequest
reply chan<- coreagent.ApprovalResult
cursor int
}
func (m chatModel) approvalHandlerForRun(events chan<- tea.Msg) coreagent.ApprovalHandler {
return chatPolicyApprovalHandler{
policy: m.policyState,
review: approvalHandlerForRun(m.reviewApproval, events),
}
}
func approvalHandlerForRun(handler coreagent.ApprovalHandler, events chan<- tea.Msg) coreagent.ApprovalHandler {
prompter := chatApprovalPrompter{ch: events}
if manager, ok := handler.(*coreagent.ApprovalManager); ok {
return manager.WithPrompter(prompter)
}
if handler != nil {
return handler
}
return coreagent.NewApprovalManager(coreagent.ApprovalManagerOptions{Prompter: prompter})
}
func chatReviewApprovalHandler(handler coreagent.ApprovalHandler, policy coreagent.RunPolicy) coreagent.ApprovalHandler {
if handler != nil {
return handler
}
return policy.ReviewApprovalHandler(nil)
}
func (m *chatModel) openApprovalPrompt(msg chatApprovalPromptMsg) {
m.approvalPrompt = &chatApprovalPrompt{request: msg.request, reply: msg.reply}
m.status = "approval required"
m.thinking = false
m.thinkingTokens = 0
m.upsertApprovalToolEntry(msg.request)
}
func (m *chatModel) togglePermissionMode() (tea.Model, tea.Cmd) {
m.ensureRunPolicy()
nextMode := coreagent.ToolModeFullAccess
if m.currentPolicy().ToolMode == coreagent.ToolModeFullAccess {
nextMode = coreagent.ToolModeReview
}
m.policyState.SetToolMode(nextMode)
m.opts.Policy = m.currentPolicy()
if nextMode == coreagent.ToolModeFullAccess {
m.permissionNotice = "full access enabled"
m.status = "full access enabled"
if m.approvalPrompt != nil {
updated, cmd := m.resolveApprovalPrompt(coreagent.ApprovalAllowOnce, "")
if model, ok := updated.(chatModel); ok {
model.permissionNotice = "full access enabled"
model.status = "full access enabled"
return model, cmd
}
return updated, cmd
}
return *m, nil
}
m.permissionNotice = "review mode enabled"
m.status = "review mode enabled"
return *m, nil
}
func (m *chatModel) ensureRunPolicy() {
if m.policyState == nil {
m.policyState = coreagent.NewRunPolicyState(m.opts.Policy)
}
if m.reviewApproval == nil {
m.reviewApproval = chatReviewApprovalHandler(m.opts.Approval, m.currentPolicy())
}
}
func (m chatModel) currentPolicy() coreagent.RunPolicy {
if m.policyState != nil {
return m.policyState.Policy()
}
return m.opts.Policy
}
func (m chatModel) autoApproveTools() bool {
return m.currentPolicy().ToolMode == coreagent.ToolModeFullAccess
}
func (m *chatModel) upsertApprovalToolEntry(request coreagent.ApprovalRequest) {
idx := m.findToolEntry(request.ToolCallID)
if idx < 0 {
m.groupCompletedToolHistory()
m.entries = append(m.entries, newChatEntry(chatEntry{role: "tool"}))
idx = len(m.entries) - 1
}
m.entries[idx].detail = request.ToolName
m.entries[idx].label = toolInvocationLabel(request.ToolName, request.Args)
m.entries[idx].status = "approval"
m.entries[idx].toolID = request.ToolCallID
m.entries[idx].args = request.Args
m.entries[idx].startedAt = time.Now()
m.applyToolOutputModeTo(idx)
m.markEntryDirty(idx)
}
func (m chatModel) updateApprovalPrompt(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
switch msg.Type {
case tea.KeyLeft, tea.KeyUp:
m.moveApprovalChoice(-1)
case tea.KeyRight, tea.KeyDown, tea.KeyTab:
m.moveApprovalChoice(1)
case tea.KeyRunes:
switch string(msg.Runes) {
case "1", "2", "3":
choice := chatApprovalChoices[int(msg.Runes[0]-'1')]
return m.resolveApprovalPrompt(choice.decision, choice.reason)
}
case tea.KeyEnter:
choice := chatApprovalChoices[clamp(m.approvalPrompt.cursor, 0, len(chatApprovalChoices)-1)]
return m.resolveApprovalPrompt(choice.decision, choice.reason)
case tea.KeyEsc, tea.KeyCtrlC:
return m.resolveApprovalPrompt(coreagent.ApprovalDeny, "Tool execution denied.")
}
return m, nil
}
func (m *chatModel) moveApprovalChoice(delta int) {
if m.approvalPrompt == nil {
return
}
m.approvalPrompt.cursor = (m.approvalPrompt.cursor + delta) % len(chatApprovalChoices)
if m.approvalPrompt.cursor < 0 {
m.approvalPrompt.cursor += len(chatApprovalChoices)
}
m.markApprovalPromptEntryDirty()
}
func (m *chatModel) markApprovalPromptEntryDirty() {
if m.approvalPrompt == nil {
return
}
if idx := m.findToolEntry(m.approvalPrompt.request.ToolCallID); idx >= 0 {
m.markEntryDirty(idx)
}
}
func (m chatModel) resolveApprovalPrompt(decision coreagent.ApprovalDecision, reason string) (tea.Model, tea.Cmd) {
if m.approvalPrompt == nil {
return m, nil
}
prompt := m.approvalPrompt
m.approvalPrompt = nil
m.status = "running"
if decision == coreagent.ApprovalDeny {
m.status = "denied"
}
if idx := m.findToolEntry(prompt.request.ToolCallID); idx >= 0 && m.entries[idx].status == "approval" {
if decision == coreagent.ApprovalDeny {
m.entries[idx].status = "error"
m.entries[idx].err = reason
if m.entries[idx].err == "" {
m.entries[idx].err = "Tool execution denied."
}
} else {
m.entries[idx].status = "queued"
}
m.markEntryDirty(idx)
}
prompt.reply <- coreagent.ApprovalResult{Decision: decision, Reason: reason}
return m, waitForChatMsg(m.events)
}
func (m chatModel) renderApprovalPromptLines(width int) []string {
prompt := m.approvalPrompt
if prompt == nil {
return nil
}
if width <= 0 {
width = 80
}
bodyWidth := max(20, width-2)
request := prompt.request
var lines []string
detail := approvalRequestDetail(request, bodyWidth)
if detail == "" && request.Summary != "" {
lines = append(lines, wrapChatText(request.Summary, width)...)
} else if detail == "" {
lines = append(lines, wrapChatText(fmt.Sprintf("%s wants to run", toolDisplayName(request.ToolName)), width)...)
}
if detail != "" {
lines = append(lines, indentLines(splitRenderedBody(detail), " ")...)
}
lines = append(lines, "")
lines = append(lines, indentLines(renderApprovalChoices(prompt.cursor, bodyWidth), " ")...)
return lines
}
func approvalRequestDetail(request coreagent.ApprovalRequest, width int) string {
if coreagent.IsShellToolName(request.ToolName) {
command, ok := rawStringArg(request.Args, "command")
if !ok {
return ""
}
return strings.Join(wrapChatText(shellPromptPrefix(request.ToolName)+command, width), "\n")
}
switch request.ToolName {
case "edit":
path, ok := rawStringArg(request.Args, "path")
if !ok {
return ""
}
var lines []string
lines = append(lines, "path: "+path)
if oldText, ok := rawStringArg(request.Args, "old_text"); ok {
lines = append(lines, fmt.Sprintf("old_text: %d chars", len([]rune(oldText))))
}
if newText, ok := rawStringArg(request.Args, "new_text"); ok {
lines = append(lines, fmt.Sprintf("new_text: %d chars", len([]rune(newText))))
}
return chatMetaStyle.Render(strings.Join(lines, "\n"))
default:
if len(request.Args) == 0 {
return ""
}
return strings.Join(renderToolCallArgs(request.Args, width), "\n")
}
}
func renderApprovalChoices(cursor int, width int) []string {
var lines []string
for i, choice := range chatApprovalChoices {
label := choice.key + ". " + choice.label
wrapped := wrapChatText(label, max(20, width-2))
if i == clamp(cursor, 0, len(chatApprovalChoices)-1) {
for j, line := range wrapped {
if j == 0 {
lines = append(lines, chatResumeSelectedStyle.Render(" "+line))
} else {
lines = append(lines, chatResumeSelectedStyle.Render(" "+line))
}
}
} else {
for j, line := range wrapped {
if j == 0 {
lines = append(lines, chatResumeTextStyle.Render(" "+line))
} else {
lines = append(lines, chatResumeTextStyle.Render(" "+line))
}
}
}
}
return lines
}
type chatApprovalPrompter struct {
ch chan<- tea.Msg
}
type chatPolicyApprovalHandler struct {
policy *coreagent.RunPolicyState
review coreagent.ApprovalHandler
}
func (h chatPolicyApprovalHandler) RequiresApproval(ctx context.Context, tool coreagent.Tool, req coreagent.ApprovalRequest) bool {
if h.policy != nil && h.policy.ToolMode() == coreagent.ToolModeFullAccess {
return false
}
if h.review != nil {
return h.review.RequiresApproval(ctx, tool, req)
}
return req.ToolApprovalRequired || coreagent.ToolRequiresApproval(tool, req.Args)
}
func (h chatPolicyApprovalHandler) Approve(ctx context.Context, req coreagent.ApprovalRequest) (coreagent.ApprovalResult, error) {
if h.policy != nil && h.policy.ToolMode() == coreagent.ToolModeFullAccess {
return coreagent.ApprovalResult{Decision: coreagent.ApprovalAllowOnce}, nil
}
if h.review != nil {
return h.review.Approve(ctx, req)
}
return coreagent.ApprovalResult{Decision: coreagent.ApprovalDeny, Reason: "Tool execution requires approval."}, nil
}
func (p chatApprovalPrompter) PromptApproval(ctx context.Context, request coreagent.ApprovalRequest) (coreagent.ApprovalResult, error) {
reply := make(chan coreagent.ApprovalResult, 1)
select {
case p.ch <- chatApprovalPromptMsg{request: request, reply: reply}:
case <-ctx.Done():
return coreagent.ApprovalResult{Decision: coreagent.ApprovalDeny, Reason: "Tool approval canceled."}, nil
}
select {
case result := <-reply:
return result, nil
case <-ctx.Done():
return coreagent.ApprovalResult{Decision: coreagent.ApprovalDeny, Reason: "Tool approval canceled."}, nil
}
}
+329
View File
@@ -0,0 +1,329 @@
package chat
import (
"context"
"strings"
"testing"
tea "github.com/charmbracelet/bubbletea"
coreagent "github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
)
func TestChatApprovalPromptRendersAndApprovesOnce(t *testing.T) {
reply := make(chan coreagent.ApprovalResult, 1)
request := coreagent.ApprovalRequest{
ToolCallID: "call-1",
ToolName: "bash",
Args: map[string]any{"command": "git status"},
Summary: "Bash wants to run a command",
Risk: coreagent.ApprovalRiskMedium,
Reasons: []string{"runs shell commands"},
}
m := chatModel{
width: 100,
height: 20,
events: make(chan tea.Msg),
}
m.openApprovalPrompt(chatApprovalPromptMsg{request: request, reply: reply})
view := stripANSI(m.View())
if !strings.Contains(view, "$ git status") ||
!strings.Contains(view, "1. Approve once") ||
!strings.Contains(view, "2. Approve session") ||
!strings.Contains(view, "3. Deny") ||
!strings.Contains(view, "") {
t.Fatalf("approval view missing content: %q", view)
}
if strings.Contains(view, " █") {
t.Fatalf("approval view should hide the input cursor: %q", view)
}
if strings.Contains(view, "1/2/3 choose • enter select • esc deny") {
t.Fatalf("approval view should not render shortcut helper: %q", view)
}
if strings.Contains(view, "Bash wants to run a command") {
t.Fatalf("approval view should not repeat tool summary when command detail is shown: %q", view)
}
if strings.Contains(view, "waiting for approval") {
t.Fatalf("approval view should not render waiting spinner: %q", view)
}
if strings.Contains(view, "Approve once (o)") || strings.Contains(view, "same command in this chat") {
t.Fatalf("approval choices should be minimal: %q", view)
}
if strings.Contains(view, "risk:") {
t.Fatalf("approval view should not render risk level: %q", view)
}
approvalIdx := strings.LastIndex(view, "1. Approve once")
inputIdx := strings.LastIndex(view, "")
if approvalIdx < 0 || inputIdx < 0 || approvalIdx > inputIdx {
t.Fatalf("approval picker should render above the input box:\n%s", view)
}
if len(m.entries) != 1 || m.entries[0].status != "approval" {
t.Fatalf("approval tool entry = %#v", m.entries)
}
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = updated.(chatModel)
if cmd == nil {
t.Fatal("approval should resume waiting for agent events")
}
result := <-reply
if result.Decision != coreagent.ApprovalAllowOnce {
t.Fatalf("decision = %q, want allow_once", result.Decision)
}
if m.approvalPrompt != nil {
t.Fatal("approval prompt should close")
}
if m.entries[0].status != "queued" {
t.Fatalf("tool status = %q, want queued", m.entries[0].status)
}
}
func TestChatApprovalPromptCtrlOExpandsToolDetailsInline(t *testing.T) {
reply := make(chan coreagent.ApprovalResult, 1)
m := chatModel{
width: 100,
height: 24,
boundedFrame: true,
fullScreen: true,
events: make(chan tea.Msg),
}
m.openApprovalPrompt(chatApprovalPromptMsg{
request: coreagent.ApprovalRequest{
ToolCallID: "call-1",
ToolName: "bash",
Args: map[string]any{"command": "git status --short --branch --untracked-files=all"},
Summary: "Bash wants to run a command",
},
reply: reply,
})
transcript := stripANSI(m.renderTranscript(100))
if strings.Contains(transcript, "$ git status") {
t.Fatalf("tool details should start collapsed: %q", transcript)
}
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyCtrlO})
m = updated.(chatModel)
if !m.fullScreen || !m.boundedFrame {
t.Fatal("ctrl+o should keep managed fullscreen rendering")
}
transcript = stripANSI(m.renderTranscript(100))
if !strings.Contains(transcript, "Bash") || !strings.Contains(transcript, "$ git status --short --branch --untracked-files=all") {
t.Fatalf("approval tool details should show inline: %q", transcript)
}
}
func TestChatApprovalPromptClosesHistoryPopup(t *testing.T) {
reply := make(chan coreagent.ApprovalResult, 1)
m := chatModel{
width: 100,
height: 24,
historyPopup: &chatHistoryPopup{messages: []api.Message{{Role: "user", Content: "old"}}},
}
updated, cmd := m.Update(chatApprovalPromptMsg{
request: coreagent.ApprovalRequest{
ToolCallID: "call-1",
ToolName: "bash",
Args: map[string]any{"command": "git status"},
Summary: "Bash wants to run a command",
},
reply: reply,
})
if cmd != nil {
t.Fatal("opening approval prompt should not return a command")
}
m = updated.(chatModel)
if m.historyPopup != nil {
t.Fatal("approval prompt should close history popup")
}
if m.approvalPrompt == nil {
t.Fatal("approval prompt should open")
}
view := stripANSI(m.View())
if strings.Contains(view, "Message history") {
t.Fatalf("history popup should not render over approval prompt:\n%s", view)
}
for _, want := range []string{"$ git status", "1. Approve once", "2. Approve session", "3. Deny"} {
if !strings.Contains(view, want) {
t.Fatalf("approval view missing %q:\n%s", want, view)
}
}
if strings.Contains(view, "1/2/3 choose • enter select • esc deny") {
t.Fatalf("approval view should not render shortcut helper:\n%s", view)
}
if strings.Contains(view, "Bash wants to run a command") {
t.Fatalf("approval view should not repeat tool summary when command detail is shown:\n%s", view)
}
if strings.Contains(view, "Approve once (o)") || strings.Contains(view, "same command in this chat") {
t.Fatalf("approval choices should be minimal:\n%s", view)
}
}
func TestChatApprovalPromptDenyNumberShortcut(t *testing.T) {
reply := make(chan coreagent.ApprovalResult, 1)
m := chatModel{
events: make(chan tea.Msg),
}
m.openApprovalPrompt(chatApprovalPromptMsg{
request: coreagent.ApprovalRequest{ToolCallID: "call-1", ToolName: "edit", Args: map[string]any{"path": "note.txt"}},
reply: reply,
})
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("3")})
m = updated.(chatModel)
result := <-reply
if result.Decision != coreagent.ApprovalDeny {
t.Fatalf("decision = %q, want deny", result.Decision)
}
if m.entries[0].status != "error" {
t.Fatalf("tool status = %q, want error", m.entries[0].status)
}
}
func TestChatShiftTabTogglesPermissionMode(t *testing.T) {
m := chatModel{}
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyShiftTab})
if cmd != nil {
t.Fatal("permission toggle should not start a command")
}
m = updated.(chatModel)
if !m.autoApproveTools() {
t.Fatal("shift+tab should enable auto-approve mode")
}
if m.status != "full access enabled" || m.notificationLine() != "" {
t.Fatalf("status = %q notification = %q, want footer-only full access notice", m.status, m.notificationLine())
}
if view := stripANSI(m.View()); !strings.Contains(view, "full access enabled") {
t.Fatalf("visible footer missing full access notice:\n%s", view)
}
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyShiftTab})
m = updated.(chatModel)
if m.autoApproveTools() {
t.Fatal("second shift+tab should return to review mode")
}
if m.status != "review mode enabled" || m.notificationLine() != "" {
t.Fatalf("status = %q notification = %q, want footer-only review notice", m.status, m.notificationLine())
}
if view := stripANSI(m.View()); !strings.Contains(view, "review mode enabled") {
t.Fatalf("visible footer missing review notice:\n%s", view)
}
}
func TestChatAutoApproveInitialFooterShowsFullAccess(t *testing.T) {
m := chatModel{
width: 100,
height: 12,
opts: Options{
Model: "llama3.2",
Policy: coreagent.RunPolicy{ToolMode: coreagent.ToolModeFullAccess},
},
status: "ready",
}
if m.notificationLine() != "" {
t.Fatalf("notification = %q, want footer-only full access notice", m.notificationLine())
}
if view := stripANSI(m.View()); !strings.Contains(view, "full access enabled") {
t.Fatalf("visible footer missing initial full access notice:\n%s", view)
}
}
func TestChatShiftTabApprovesPendingPrompt(t *testing.T) {
reply := make(chan coreagent.ApprovalResult, 1)
m := chatModel{
events: make(chan tea.Msg),
}
m.openApprovalPrompt(chatApprovalPromptMsg{
request: coreagent.ApprovalRequest{ToolCallID: "call-1", ToolName: "bash", Args: map[string]any{"command": "pwd"}},
reply: reply,
})
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyShiftTab})
m = updated.(chatModel)
if cmd == nil {
t.Fatal("approving pending prompt should resume waiting for agent events")
}
result := <-reply
if result.Decision != coreagent.ApprovalAllowOnce {
t.Fatalf("decision = %q, want allow_once", result.Decision)
}
if !m.autoApproveTools() {
t.Fatal("shift+tab should leave future tool calls in auto-approve mode")
}
if m.approvalPrompt != nil {
t.Fatal("approval prompt should close")
}
if m.entries[0].status != "queued" {
t.Fatalf("tool status = %q, want queued", m.entries[0].status)
}
}
func TestChatShiftTabSkillApprovalKeepsFullAccessFooter(t *testing.T) {
reply := make(chan coreagent.ApprovalResult, 1)
m := chatModel{
width: 100,
height: 20,
events: make(chan tea.Msg),
}
m.openApprovalPrompt(chatApprovalPromptMsg{
request: coreagent.ApprovalRequest{
ToolCallID: "call-1",
ToolName: "skill",
Args: map[string]any{"name": "code-review"},
},
reply: reply,
})
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyShiftTab})
m = updated.(chatModel)
if result := <-reply; result.Decision != coreagent.ApprovalAllowOnce {
t.Fatalf("decision = %q, want allow_once", result.Decision)
}
if view := stripANSI(m.View()); !strings.Contains(view, "full access enabled") {
t.Fatalf("visible footer missing full access notice after approval:\n%s", view)
}
updated, _ = m.Update(chatRunDoneMsg{result: &coreagent.RunResult{
Messages: []api.Message{{Role: "assistant", Content: "done"}},
}})
m = updated.(chatModel)
if m.status != "ready" {
t.Fatalf("status = %q, want ready", m.status)
}
if view := stripANSI(m.View()); !strings.Contains(view, "full access enabled") {
t.Fatalf("visible footer should keep full access notice after run completes:\n%s", view)
}
}
func TestChatPermissionApprovalHandlerReadsModeAtApprovalTime(t *testing.T) {
policy := coreagent.NewRunPolicyState(coreagent.RunPolicy{ToolMode: coreagent.ToolModeReview})
handler := chatPolicyApprovalHandler{
policy: policy,
review: coreagent.NewApprovalManager(coreagent.ApprovalManagerOptions{}),
}
req := coreagent.ApprovalRequest{
ToolName: "bash",
Args: map[string]any{"command": "pwd"},
}
if !handler.RequiresApproval(context.Background(), chatTestTool{}, req) {
t.Fatal("review mode should require approval for bash")
}
policy.SetToolMode(coreagent.ToolModeFullAccess)
if handler.RequiresApproval(context.Background(), chatTestTool{}, req) {
t.Fatal("auto-approve mode should not require approval")
}
result, err := handler.Approve(context.Background(), req)
if err != nil {
t.Fatal(err)
}
if result.Decision != coreagent.ApprovalAllowOnce {
t.Fatalf("decision = %q, want allow_once", result.Decision)
}
}
+1175
View File
File diff suppressed because it is too large. Load diff
+252
View File
@@ -0,0 +1,252 @@
package chat
import (
"context"
"errors"
"fmt"
"strings"
"time"
tea "github.com/charmbracelet/bubbletea"
"github.com/ollama/ollama/api"
)
type cloudAuthKind string
const (
cloudAuthSignIn cloudAuthKind = "signin"
cloudAuthUpgrade cloudAuthKind = "upgrade"
cloudAuthChecking cloudAuthKind = "checking"
)
// cloudAuthPrompt is an inline modal that handles sign-in and plan-upgrade
// flows when a user selects a cloud model from the picker.
type cloudAuthPrompt struct {
modelName string
requiredPlan string
signInURL string
kind cloudAuthKind
spinner int
openNow bool
polling bool
}
type cloudAuthCheckMsg struct {
err error
signInURL string
}
type cloudAuthTickMsg struct{}
type cloudAuthPollMsg struct {
done bool
}
func checkCloudModelCmd(ctx context.Context, check func(context.Context, string, string) error, model, requiredPlan string) tea.Cmd {
if check == nil {
return nil
}
return func() tea.Msg {
if ctx == nil {
ctx = context.Background()
}
err := check(ctx, model, requiredPlan)
var signInURL string
if err != nil {
var authErr api.AuthorizationError
if errors.As(err, &authErr) && authErr.SigninURL != "" {
signInURL = authErr.SigninURL
}
}
return cloudAuthCheckMsg{err: err, signInURL: signInURL}
}
}
func cloudAuthTickCmd() tea.Cmd {
return tea.Tick(200*time.Millisecond, func(t time.Time) tea.Msg {
return cloudAuthTickMsg{}
})
}
func pollCloudAuthCmd(ctx context.Context, poll func(context.Context) (string, bool)) tea.Cmd {
if poll == nil {
return nil
}
return func() tea.Msg {
if ctx == nil {
ctx = context.Background()
}
pollCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
defer cancel()
_, done := poll(pollCtx)
return cloudAuthPollMsg{done: done}
}
}
func (m *chatModel) startCloudAuthSignIn(modelName, requiredPlan, signInURL string) (tea.Model, tea.Cmd) {
m.cloudAuthPrompt = &cloudAuthPrompt{
modelName: modelName,
requiredPlan: requiredPlan,
kind: cloudAuthSignIn,
signInURL: signInURL,
polling: true,
}
m.status = "cloud-auth"
m.modelPicker = nil
if m.opts.OpenBrowser != nil && signInURL != "" {
m.opts.OpenBrowser(signInURL)
}
if signInURL == "" {
return m, checkCloudModelCmd(m.ctx, m.opts.CheckCloudModel, modelName, requiredPlan)
}
return m, tea.Batch(cloudAuthTickCmd(), pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth))
}
func (m *chatModel) startCloudAuthUpgrade(modelName, requiredPlan string) (tea.Model, tea.Cmd) {
m.cloudAuthPrompt = &cloudAuthPrompt{
modelName: modelName,
requiredPlan: requiredPlan,
kind: cloudAuthUpgrade,
openNow: true,
}
m.status = "cloud-auth"
m.modelPicker = nil
return m, nil
}
func (m chatModel) updateCloudAuthPrompt(msg tea.Msg) (tea.Model, tea.Cmd) {
switch msg := msg.(type) {
case cloudAuthCheckMsg:
if msg.err == nil {
// Auth passed — apply the pending model.
return m.completeCloudAuth()
}
// Determine if sign-in or upgrade is needed.
if msg.signInURL != "" {
m.cloudAuthPrompt.kind = cloudAuthSignIn
m.cloudAuthPrompt.signInURL = msg.signInURL
m.cloudAuthPrompt.polling = true
if m.opts.OpenBrowser != nil {
m.opts.OpenBrowser(msg.signInURL)
}
return m, tea.Batch(cloudAuthTickCmd(), pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth))
}
// Could be a plan upgrade error or unknown error.
m.cloudAuthPrompt = nil
m.status = "ready"
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", msg.err), err: msg.err.Error()}))
return m, nil
case cloudAuthTickMsg:
if m.cloudAuthPrompt == nil {
return m, nil
}
m.cloudAuthPrompt.spinner++
return m, cloudAuthTickCmd()
case cloudAuthPollMsg:
if m.cloudAuthPrompt == nil {
return m, nil
}
if msg.done {
// Signed in — re-check auth to see if plan is satisfied.
m.cloudAuthPrompt.polling = false
return m, checkCloudModelCmd(m.ctx, m.opts.CheckCloudModel, m.cloudAuthPrompt.modelName, m.cloudAuthPrompt.requiredPlan)
}
return m, pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth)
case tea.KeyMsg:
if msg.Type == tea.KeyEsc || msg.Type == tea.KeyCtrlC {
m.cloudAuthPrompt = nil
m.pendingModel = ""
m.status = "ready"
return m, nil
}
if m.cloudAuthPrompt.kind == cloudAuthUpgrade && !m.cloudAuthPrompt.polling {
switch msg.Type {
case tea.KeyLeft, tea.KeyRight, tea.KeyTab:
m.cloudAuthPrompt.openNow = !m.cloudAuthPrompt.openNow
case tea.KeyEnter:
if m.cloudAuthPrompt.openNow {
m.cloudAuthPrompt.polling = true
return m, tea.Batch(cloudAuthTickCmd(), pollCloudAuthCmd(m.ctx, m.opts.PollCloudAuth))
}
m.cloudAuthPrompt = nil
m.pendingModel = ""
m.status = "ready"
return m, nil
}
}
}
return m, nil
}
func (m chatModel) completeCloudAuth() (tea.Model, tea.Cmd) {
pending := m.cloudAuthPrompt.modelName
m.cloudAuthPrompt = nil
m.pendingModel = ""
m.modelPicker = nil
m.status = "ready"
if err := m.applyModelSelection(pending, true); err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", err), err: err.Error()}))
m.status = "error"
return m, nil
}
return m, m.startModelPreload(pending)
}
func (m chatModel) renderCloudAuthPrompt(width int) string {
if m.cloudAuthPrompt == nil {
return ""
}
if width <= 0 {
width = 80
}
p := m.cloudAuthPrompt
spinnerFrames := []string{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"}
frame := spinnerFrames[p.spinner%len(spinnerFrames)]
var b strings.Builder
switch p.kind {
case cloudAuthChecking:
fmt.Fprintf(&b, "%s Checking %s...\n\n", frame, chatResumeSelectedStyle.Render(p.modelName))
b.WriteString(chatResumeMetaStyle.Render("esc cancel"))
case cloudAuthSignIn:
fmt.Fprintf(&b, "To use %s, please sign in.\n\n", chatResumeSelectedStyle.Render(p.modelName))
b.WriteString("Navigate to:\n")
urlWrap := chatResumeTextStyle
if width > 4 {
urlWrap = chatResumeTextStyle.Width(width - 4)
}
b.WriteString(urlWrap.Render(p.signInURL))
b.WriteString("\n\n")
b.WriteString(chatResumeMetaStyle.Render(frame + " Waiting for sign in to complete..."))
b.WriteString("\n\n")
b.WriteString(chatResumeMetaStyle.Render("esc cancel"))
case cloudAuthUpgrade:
fmt.Fprintf(&b, "To use %s, upgrade your Ollama plan.\n\n", chatResumeSelectedStyle.Render(p.modelName))
if !p.polling {
var yesBtn, noBtn string
if p.openNow {
yesBtn = chatResumeSelectedStyle.Render(" Yes ")
noBtn = chatResumeMetaStyle.Render(" No ")
} else {
yesBtn = chatResumeMetaStyle.Render(" Yes ")
noBtn = chatResumeSelectedStyle.Render(" No ")
}
b.WriteString("Open upgrade page now?\n")
b.WriteString(yesBtn + " " + noBtn)
b.WriteString("\n\n")
b.WriteString(chatResumeMetaStyle.Render("←/→ navigate • enter confirm • esc cancel"))
} else {
b.WriteString(chatResumeMetaStyle.Render(frame + " Waiting for upgrade to complete..."))
b.WriteString("\n\n")
b.WriteString(chatResumeMetaStyle.Render("esc cancel"))
}
}
return b.String()
}
+66
View File
@@ -0,0 +1,66 @@
package chat
import (
"context"
"testing"
)
func TestCloudAuthTickDoesNotPoll(t *testing.T) {
polls := 0
m := chatModel{
cloudAuthPrompt: &cloudAuthPrompt{polling: true},
opts: Options{
PollCloudAuth: func(context.Context) (string, bool) {
polls++
return "", false
},
},
}
updated, cmd := m.updateCloudAuthPrompt(cloudAuthTickMsg{})
m = updated.(chatModel)
if m.cloudAuthPrompt.spinner != 1 {
t.Fatalf("spinner = %d, want 1", m.cloudAuthPrompt.spinner)
}
if polls != 0 {
t.Fatalf("polls = %d, want 0 before running returned tick command", polls)
}
if cmd == nil {
t.Fatal("tick should schedule the next tick")
}
if _, ok := cmd().(cloudAuthTickMsg); !ok {
t.Fatal("tick should schedule another tick, not a poll")
}
if polls != 0 {
t.Fatalf("polls = %d, want 0 after running returned tick command", polls)
}
}
func TestCloudAuthPollSchedulesNextPoll(t *testing.T) {
polls := 0
m := chatModel{
cloudAuthPrompt: &cloudAuthPrompt{polling: true},
opts: Options{
PollCloudAuth: func(context.Context) (string, bool) {
polls++
return "", false
},
},
}
_, cmd := m.updateCloudAuthPrompt(cloudAuthPollMsg{})
if cmd == nil {
t.Fatal("poll should schedule the next poll")
}
msg, ok := cmd().(cloudAuthPollMsg)
if !ok {
t.Fatal("poll should schedule another poll, not a tick")
}
if msg.done {
t.Fatal("poll should report not done")
}
if polls != 1 {
t.Fatalf("polls = %d, want 1", polls)
}
}
+101
View File
@@ -0,0 +1,101 @@
package chat
import (
"context"
"slices"
tea "github.com/charmbracelet/bubbletea"
coreagent "github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
)
func (m *chatModel) startManualCompaction() (tea.Model, tea.Cmd) {
if m.running || m.compacting {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: "Wait for the current response to finish before compacting."}))
return *m, nil
}
m.refreshContextWindowTokens(m.opts.Model)
if m.opts.Compactor == nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: coreagent.CompactionSkippedMessage("compaction is unavailable")}))
m.status = "compact skipped"
return *m, nil
}
ctx := m.ctx
if ctx == nil {
ctx = context.Background()
}
runCtx, cancel := context.WithCancel(ctx)
compactor := m.opts.Compactor
events := make(chan tea.Msg, 128)
m.compacting = true
m.compactingTokens = 0
m.cancel = cancel
m.compactEvents = events
m.status = "compacting"
messages := slices.Clone(m.messages)
var tools api.Tools
if m.opts.Tools != nil {
tools = m.opts.Tools.Tools()
}
req := coreagent.CompactionRequest{
ChatID: m.chatID,
Model: m.opts.Model,
SystemPrompt: m.systemPrompt(""),
Messages: messages,
Tools: tools,
Format: m.opts.Format,
Options: m.opts.Options,
KeepAlive: m.opts.KeepAlive,
Force: true,
Progress: func(progress coreagent.CompactionProgress) {
select {
case events <- chatCompactProgressMsg{tokens: progress.Tokens}:
case <-runCtx.Done():
}
},
}
go func() {
defer close(events)
result, err := compactor.MaybeCompact(runCtx, req)
select {
case events <- chatCompactDoneMsg{result: result, err: err}:
case <-runCtx.Done():
}
}()
tickCmd := m.scheduleTick()
return *m, tea.Batch(waitForChatMsg(events), tickCmd)
}
func (m chatModel) finishManualCompaction(msg chatCompactDoneMsg) (tea.Model, tea.Cmd) {
wasCanceling := m.status == "canceling"
m.compacting = false
m.compactEvents = nil
m.cancel = nil
m.compactingTokens = 0
if wasCanceling || isChatContextCanceledError(msg.err) {
m.status = "compact canceled"
return m.withFlowTranscriptFlush(m.startNextQueued())
}
if msg.err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: coreagent.CompactionSkippedMessage(msg.err.Error())}))
m.status = "compact skipped"
return m.withFlowTranscriptFlush(m.startNextQueued())
}
if !msg.result.Compacted {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: coreagent.CompactionSkippedMessage(msg.result.Reason)}))
m.status = "compact skipped"
return m.withFlowTranscriptFlush(m.startNextQueued())
}
m.messages = msg.result.Messages
m.liveMessages = nil
m.entries = entriesFromMessages(m.messages)
m.contextTokens = m.estimatePromptTokens(m.messages, "")
m.contextEstimate = true
m.scroll = 0
m.flowPrintedLines = 0
m.status = "compacted"
return m.withFlowTranscriptFlush(m.startNextQueued())
}
+217
View File
@@ -0,0 +1,217 @@
package chat
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
tea "github.com/charmbracelet/bubbletea"
coreagent "github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
)
type rawRequestDisplay struct {
Model string `json:"model"`
Messages []rawRequestMessage `json:"messages"`
Format any `json:"format,omitempty"`
KeepAlive *api.Duration `json:"keep_alive,omitempty"`
Tools api.Tools `json:"tools,omitempty"`
Options map[string]any `json:"options"`
Think *api.ThinkValue `json:"think,omitempty"`
}
type rawRequestMessage struct {
Role string `json:"role"`
Content string `json:"content"`
Thinking string `json:"thinking,omitempty"`
Images []string `json:"images,omitempty"`
ToolCalls []api.ToolCall `json:"tool_calls,omitempty"`
ToolName string `json:"tool_name,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
}
func (m chatModel) requestPreview(messages []api.Message) coreagent.ChatRequestPreview {
var tools api.Tools
policy := m.currentPolicy()
if registry := policy.Tools(m.opts.Tools); registry != nil {
tools = registry.Tools()
}
return coreagent.BuildChatRequestPreview(coreagent.RunOptions{
Model: m.opts.Model,
SystemPrompt: m.systemPrompt(""),
Format: m.opts.Format,
Options: m.opts.Options,
Think: m.opts.Think,
KeepAlive: m.opts.KeepAlive,
Policy: policy,
}, messages, tools)
}
func (m *chatModel) openRawRequestPopup() (tea.Model, tea.Cmd) {
raw, tokens, err := m.rawRequestPreviewJSON()
if err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
m.status = "error"
return *m, nil
}
m.historyPopup = &chatHistoryPopup{
title: "Raw request",
empty: "No request preview.",
header: m.promptTokenHeader(tokens),
raw: raw,
stickToBottom: false,
}
m.status = "raw"
return *m, tea.EnterAltScreen
}
func (m *chatModel) handleRawCommand(args string) (tea.Model, tea.Cmd) {
if strings.TrimSpace(args) == "" {
return m.openRawRequestPopup()
}
filename, err := rawRedirectFilename(args)
if err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
m.status = "error"
return *m, nil
}
return m.saveRawRequest(filename)
}
func (m *chatModel) saveRawRequest(filename string) (tea.Model, tea.Cmd) {
raw, _, err := m.rawRequestPreviewJSON()
if err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
m.status = "error"
return *m, nil
}
dir := m.currentWorkingDir()
if strings.TrimSpace(dir) == "" {
dir, err = os.Getwd()
if err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
m.status = "error"
return *m, nil
}
}
path := filepath.Join(dir, filename)
if err := os.WriteFile(path, []byte(raw+"\n"), 0o644); err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
m.status = "error"
return *m, nil
}
m.entries = append(m.entries, newSlashEntry(fmt.Sprintf("Saved raw request to %s", filename)))
m.status = "raw saved"
return *m, nil
}
func (m chatModel) rawRequestPreviewJSON() (string, int, error) {
preview := m.requestPreview(m.messages)
raw, err := rawRequestJSON(preview.Request)
if err != nil {
return "", 0, err
}
return raw, preview.PromptTokens, nil
}
func rawRedirectFilename(args string) (string, error) {
args = strings.TrimSpace(args)
if !strings.HasPrefix(args, ">") {
return "", fmt.Errorf("usage: /raw > filename")
}
fields := strings.Fields(strings.TrimSpace(strings.TrimPrefix(args, ">")))
if len(fields) != 1 {
return "", fmt.Errorf("usage: /raw > filename")
}
filename := strings.TrimSpace(fields[0])
if filename == "" || filename == "." || filename == ".." || strings.ContainsAny(filename, `/\`) {
return "", fmt.Errorf("raw filename must be a file name, not a path")
}
if !strings.HasSuffix(strings.ToLower(filename), ".json") {
filename += ".json"
}
return filename, nil
}
func (m chatModel) promptTokenHeader(tokens int) []string {
window := coreagent.ResolveContextWindowTokens(m.opts.Options, m.opts.ContextWindowTokens)
line := "estimated prompt: " + formatTokenCount(tokens)
if window > 0 {
line = fmt.Sprintf("estimated prompt: %d / %d tokens", max(tokens, 0), window)
}
return []string{chatResumeMetaStyle.Render(line)}
}
func rawRequestJSON(req api.ChatRequest) (string, error) {
display := rawRequestDisplay{
Model: req.Model,
Messages: rawRequestMessages(req.Messages),
KeepAlive: req.KeepAlive,
Tools: req.Tools,
Options: req.Options,
Think: req.Think,
}
if format := rawRequestFormat(req.Format); format != nil {
display.Format = format
}
data, err := json.MarshalIndent(display, "", " ")
if err != nil {
return "", err
}
return string(data), nil
}
func rawRequestFormat(format json.RawMessage) any {
if len(format) == 0 {
return nil
}
var value any
if err := json.Unmarshal(format, &value); err == nil {
return value
}
return string(format)
}
func rawRequestMessages(messages []api.Message) []rawRequestMessage {
out := make([]rawRequestMessage, 0, len(messages))
for _, msg := range messages {
out = append(out, rawRequestMessage{
Role: msg.Role,
Content: msg.Content,
Thinking: msg.Thinking,
Images: rawImagePlaceholders(msg.Images),
ToolCalls: msg.ToolCalls,
ToolName: msg.ToolName,
ToolCallID: msg.ToolCallID,
})
}
return out
}
func rawImagePlaceholders(images []api.ImageData) []string {
if len(images) == 0 {
return nil
}
out := make([]string, 0, len(images))
for _, image := range images {
out = append(out, fmt.Sprintf("[image data: %d bytes]", len(image)))
}
return out
}
func renderRawRequestLines(raw string, width int) []string {
raw = strings.TrimRight(raw, "\n")
if raw == "" {
return nil
}
var lines []string
for _, line := range strings.Split(raw, "\n") {
lines = append(lines, renderHistoryCodeLine("", line, width)...)
}
return stringsTrimTrailingEmptyLines(lines)
}
+97
View File
@@ -0,0 +1,97 @@
package chat
import (
"fmt"
"os"
"os/exec"
"strings"
tea "github.com/charmbracelet/bubbletea"
"github.com/ollama/ollama/envconfig"
)
type chatEditorDoneMsg struct {
content string
err error
}
func (m chatModel) openInputEditor() (tea.Model, tea.Cmd) {
cmd, path, cleanup, err := inputEditorCommand(string(m.input))
if err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
m.status = "editor error"
return m, nil
}
return m, tea.ExecProcess(cmd, func(err error) tea.Msg {
defer cleanup()
if err != nil {
return chatEditorDoneMsg{err: fmt.Errorf("editor exited with error: %w", err)}
}
data, err := os.ReadFile(path)
if err != nil {
return chatEditorDoneMsg{err: fmt.Errorf("reading editor content: %w", err)}
}
return chatEditorDoneMsg{content: strings.TrimRight(string(data), "\n")}
})
}
func (m *chatModel) applyEditorResult(msg chatEditorDoneMsg) {
if msg.err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: msg.err.Error(), err: msg.err.Error()}))
m.status = "editor error"
return
}
m.input = []rune(msg.content)
m.inputCursor = len(m.input)
m.inputCursorSet = true
m.complete = 0
m.resetPromptHistoryCursor()
m.syncInputPlaceholders()
m.status = "editor"
}
func inputEditorCommand(content string) (*exec.Cmd, string, func(), error) {
editor := inputEditorName()
args := strings.Fields(editor)
if len(args) == 0 {
return nil, "", nil, fmt.Errorf("editor is empty, set OLLAMA_EDITOR to the path of your preferred editor")
}
if _, err := exec.LookPath(args[0]); err != nil {
return nil, "", nil, fmt.Errorf("editor %q not found, set OLLAMA_EDITOR to the path of your preferred editor", args[0])
}
tmpFile, err := os.CreateTemp("", "ollama-prompt-*.txt")
if err != nil {
return nil, "", nil, fmt.Errorf("creating temp file: %w", err)
}
cleanup := func() { _ = os.Remove(tmpFile.Name()) }
if content != "" {
if _, err := tmpFile.WriteString(content); err != nil {
tmpFile.Close()
cleanup()
return nil, "", nil, fmt.Errorf("writing to temp file: %w", err)
}
}
if err := tmpFile.Close(); err != nil {
cleanup()
return nil, "", nil, fmt.Errorf("closing temp file: %w", err)
}
args = append(args, tmpFile.Name())
return exec.Command(args[0], args[1:]...), tmpFile.Name(), cleanup, nil
}
func inputEditorName() string {
if editor := strings.TrimSpace(envconfig.Editor()); editor != "" {
return editor
}
if editor := strings.TrimSpace(os.Getenv("VISUAL")); editor != "" {
return editor
}
if editor := strings.TrimSpace(os.Getenv("EDITOR")); editor != "" {
return editor
}
return defaultEditor
}
@@ -1,5 +1,5 @@
//go:build !windows
package cmd
package chat
const defaultEditor = "vi"
@@ -1,5 +1,5 @@
//go:build windows
package cmd
package chat
const defaultEditor = "edit"
+334
View File
@@ -0,0 +1,334 @@
package chat
import (
"context"
"slices"
"strings"
"time"
tea "github.com/charmbracelet/bubbletea"
coreagent "github.com/ollama/ollama/agent"
"github.com/ollama/ollama/api"
)
type chatAgentMsg struct {
event coreagent.Event
}
type chatApprovalPromptMsg struct {
request coreagent.ApprovalRequest
reply chan<- coreagent.ApprovalResult
}
type chatRunDoneMsg struct {
result *coreagent.RunResult
err error
newMessagesPersisted bool
persistedMessages []api.Message
}
type chatCompactDoneMsg struct {
result coreagent.CompactionResult
err error
}
type chatCompactProgressMsg struct {
tokens int
}
// resetStreamingState clears the transient streaming flags that every
// non-streaming event resets before applying its own state.
func (m *chatModel) resetStreamingState() {
m.awaitingModel = false
m.thinking = false
m.thinkingTokens = 0
}
// resetRunState clears all run-progress flags (streaming plus compaction
// progress) for terminal events that fully reset the run view.
func (m *chatModel) resetRunState() {
m.awaitingModel = false
m.compacting = false
m.compactingTokens = 0
m.thinking = false
m.thinkingTokens = 0
}
type chatModelPreloadDoneMsg struct {
model string
err error
}
type chatEventsClosedMsg struct{}
type chatTickMsg struct{}
func (m *chatModel) applyAgentEvent(event coreagent.Event) {
contextChanged := false
skipResponseMetrics := false
switch event.Type {
case coreagent.EventRequestBuilt:
m.awaitingModel = true
case coreagent.EventMessageStarted:
m.awaitingModel = m.running
m.eventErrorRendered = false
m.compacting = false
m.compactingTokens = 0
m.thinking = false
m.thinkingTokens = 0
m.refreshContextWindowTokens(m.opts.Model)
case coreagent.EventThinkingDelta:
m.awaitingModel = false
if event.Thinking != "" {
m.thinking = true
m.thinkingTokens = max(m.thinkingTokens, eventEvalCount(event))
if eventEvalCount(event) <= 0 {
m.thinkingTokens += coreagent.EstimateTokens(event.Thinking)
}
idx := m.ensureLiveAssistantMessage()
m.liveMessages[idx].Thinking += event.Thinking
contextChanged = true
}
case coreagent.EventMessageDelta:
m.resetStreamingState()
m.groupCompletedToolHistory()
idx := m.ensureAssistantEntry()
m.entries[idx].content += event.Content
m.markEntryDirty(idx)
msgIdx := m.ensureLiveAssistantMessage()
m.liveMessages[msgIdx].Content += event.Content
contextChanged = true
case coreagent.EventToolCallDetected:
m.awaitingModel = m.running
m.thinking = false
m.thinkingTokens = 0
idx := m.ensureLiveAssistantMessage()
m.liveMessages[idx].ToolCalls = append(m.liveMessages[idx].ToolCalls, event.ToolCalls...)
contextChanged = true
case coreagent.EventToolStarted:
m.resetStreamingState()
m.refreshContextWindowTokens(m.opts.Model)
idx := m.findActiveToolEntry(event.ToolCallID)
if idx < 0 {
m.groupCompletedToolHistory()
m.entries = append(m.entries, newChatEntry(chatEntry{role: "tool"}))
idx = len(m.entries) - 1
}
m.entries[idx].detail = event.ToolName
m.entries[idx].label = toolInvocationLabel(event.ToolName, event.Args)
m.entries[idx].status = "running"
m.entries[idx].toolID = event.ToolCallID
m.entries[idx].args = event.Args
m.entries[idx].startedAt = event.StartedAt
m.applyToolOutputModeTo(idx)
m.markEntryDirty(idx)
case coreagent.EventToolFinished:
m.resetStreamingState()
m.refreshContextWindowTokens(m.opts.Model)
if event.WorkingDir != "" {
m.workingDir = event.WorkingDir
}
startedAt := m.toolStartedAt(event.ToolCallID)
status := "done"
if event.Error != "" {
status = "error"
}
idx := m.findToolEntry(event.ToolCallID)
if idx < 0 {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "tool"}))
idx = len(m.entries) - 1
}
m.entries[idx].content = event.Content
m.entries[idx].label = toolInvocationLabel(event.ToolName, event.Args)
m.entries[idx].detail = event.ToolName
m.entries[idx].status = status
m.entries[idx].err = event.Error
m.entries[idx].toolID = event.ToolCallID
m.entries[idx].args = event.Args
m.entries[idx].startedAt = startedAt
m.entries[idx].finishedAt = event.FinishedAt
m.applyToolOutputModeTo(idx)
m.markEntryDirty(idx)
m.liveMessages = append(m.liveMessages, api.Message{
Role: "tool",
Content: event.Content,
ToolName: event.ToolName,
ToolCallID: event.ToolCallID,
})
contextChanged = true
case coreagent.EventToolsUnavailable:
m.resetStreamingState()
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: "Tools are unavailable for this model."}))
case coreagent.EventCompacted:
m.resetRunState()
skipResponseMetrics = true
if len(event.Messages) > 0 {
m.liveMessages = slices.Clone(event.Messages)
m.messages = slices.Clone(event.Messages)
contextChanged = true
}
m.status = "compacted"
case coreagent.EventCompactionStarted:
m.awaitingModel = false
m.compacting = true
m.compactingTokens = 0
m.thinking = false
m.thinkingTokens = 0
m.status = "compacting"
case coreagent.EventCompactionProgress:
m.awaitingModel = false
m.compacting = true
m.thinking = false
m.thinkingTokens = 0
if event.Tokens > m.compactingTokens {
m.compactingTokens = event.Tokens
}
case coreagent.EventCompactionSkipped:
m.resetRunState()
message := event.Content
if strings.TrimSpace(message) == "" {
message = coreagent.CompactionSkippedMessage(event.Error)
}
m.entries = append(m.entries, newChatEntry(chatEntry{role: "system", content: message}))
m.status = "compact skipped"
case coreagent.EventModelStreamDone:
m.awaitingModel = m.running && m.awaitingToolStart()
case coreagent.EventError:
m.resetRunState()
m.eventErrorRendered = true
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: event.Error, err: event.Error}))
}
if contextChanged {
m.refreshLiveContextEstimate()
}
if event.Type == coreagent.EventCompacted && event.PromptTokens > 0 {
m.contextTokens = event.PromptTokens
m.contextEstimate = true
}
if messagesEndWithCompactionResult(m.liveMessages) || (len(m.liveMessages) == 0 && messagesEndWithCompactionResult(m.messages)) {
skipResponseMetrics = true
}
if !skipResponseMetrics {
m.applyResponseMetrics(event.Response)
}
}
func messagesEndWithCompactionResult(messages []api.Message) bool {
if len(messages) == 0 {
return false
}
return coreagent.IsCompactionToolResult(messages[len(messages)-1])
}
func (m chatModel) awaitingToolStart() bool {
if len(m.liveMessages) == 0 {
return false
}
msg := m.liveMessages[len(m.liveMessages)-1]
if msg.Role != "assistant" || len(msg.ToolCalls) == 0 {
return false
}
for _, call := range msg.ToolCalls {
if call.ID == "" || m.findToolEntry(call.ID) < 0 {
return true
}
}
return false
}
func (m *chatModel) ensureLiveAssistantMessage() int {
if len(m.liveMessages) > 0 && m.liveMessages[len(m.liveMessages)-1].Role == "assistant" {
return len(m.liveMessages) - 1
}
m.liveMessages = append(m.liveMessages, api.Message{Role: "assistant"})
return len(m.liveMessages) - 1
}
func (m *chatModel) refreshLiveContextEstimate() {
messages := m.liveMessages
if len(messages) == 0 {
messages = m.messages
}
m.contextTokens = m.estimatePromptTokens(messages, "")
m.contextEstimate = true
}
//nolint:containedctx // event sinks need the session context to unblock sends on cancellation.
type chatEventSink struct {
ctx context.Context
ch chan<- tea.Msg
newMessagesPersisted *bool
}
func (s chatEventSink) Emit(event coreagent.Event) error {
if event.Type == coreagent.EventLoopStep && s.newMessagesPersisted != nil {
*s.newMessagesPersisted = true
}
select {
case s.ch <- chatAgentMsg{event: event}:
return nil
case <-s.ctx.Done():
return s.ctx.Err()
}
}
func waitForChatMsg(ch <-chan tea.Msg) tea.Cmd {
if ch == nil {
return nil
}
return func() tea.Msg {
msg, ok := <-ch
if !ok {
return chatEventsClosedMsg{}
}
return msg
}
}
func (m *chatModel) scheduleTick() tea.Cmd {
if m.tickActive {
return nil
}
m.tickActive = true
return chatTickCmd()
}
func chatTickCmd() tea.Cmd {
return tea.Tick(120*time.Millisecond, func(time.Time) tea.Msg {
return chatTickMsg{}
})
}
func preloadModelCmd(ctx context.Context, preload func(context.Context, string, *api.ThinkValue) error, model string, think *api.ThinkValue) tea.Cmd {
if preload == nil || strings.TrimSpace(model) == "" {
return nil
}
if think != nil {
copied := *think
think = &copied
}
return func() tea.Msg {
if ctx == nil {
ctx = context.Background()
}
return chatModelPreloadDoneMsg{model: model, err: preload(ctx, model, think)}
}
}
func isUnsupportedThinkingError(err error) bool {
if err == nil {
return false
}
text := strings.ToLower(err.Error())
return strings.Contains(text, "does not support thinking")
}
func thinkRequestsThinking(think *api.ThinkValue) bool {
if think == nil {
return false
}
return think.Bool()
}
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
+297
View File
@@ -0,0 +1,297 @@
package chat
import (
"context"
"fmt"
"sort"
"strings"
"time"
tea "github.com/charmbracelet/bubbletea"
"github.com/ollama/ollama/api"
)
type chatShowClient interface {
Show(context.Context, *api.ShowRequest) (*api.ShowResponse, error)
}
func (m *chatModel) handleLegacySetCommand(input string) (tea.Model, tea.Cmd) {
args := strings.Fields(input)
if len(args) < 2 {
m.entries = append(m.entries, newSlashEntry(legacySetUsage()))
return *m, nil
}
switch args[1] {
case "think":
return m.handleLegacySetThinkCommand(input)
case "nothink":
return m.applyThinkValue("off")
case "verbose":
return m.handleVerboseCommand("/verbose on")
case "quiet":
return m.handleVerboseCommand("/verbose off")
case "format":
if len(args) != 3 || args[2] != "json" {
return m.legacyUsageError("Invalid or missing format. For JSON mode use `/set format json`.")
}
m.opts.Format = "json"
m.refreshContextEstimate()
m.status = "format json"
return *m, nil
case "noformat":
m.opts.Format = ""
m.refreshContextEstimate()
m.status = "format off"
return *m, nil
case "parameter":
if len(args) < 4 {
m.entries = append(m.entries, newSlashEntry(legacyParameterUsage()))
return *m, nil
}
params, err := api.FormatParams(map[string][]string{args[2]: args[3:]})
if err != nil {
return m.legacyUsageError(fmt.Sprintf("Couldn't set parameter: %v", err))
}
if m.opts.Options == nil {
m.opts.Options = make(map[string]any)
}
for key, value := range params {
m.opts.Options[key] = value
}
m.refreshContextEstimate()
m.status = "parameter " + args[2]
return *m, nil
default:
return m.legacyUsageError(fmt.Sprintf("Unknown command `/set %s`.", args[1]))
}
}
func (m *chatModel) handleLegacyLoadCommand(input string) (tea.Model, tea.Cmd) {
args := strings.Fields(input)
if len(args) != 2 {
return m.legacyUsageError("Usage: `/load <model>`")
}
if err := m.applyModelSelection(args[1], true); err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not switch model: %v", err), err: err.Error()}))
m.status = "error"
return *m, nil
}
m.status = "ready"
return *m, nil
}
func (m *chatModel) handleLegacyShowCommand(input string) (tea.Model, tea.Cmd) {
args := strings.Fields(input)
if len(args) != 2 {
m.entries = append(m.entries, newSlashEntry(legacyShowUsage()))
return *m, nil
}
resp, err := m.legacyShowResponse()
if err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: fmt.Sprintf("Could not show model info: %v", err), err: err.Error()}))
m.status = "error"
return *m, nil
}
content, ok := m.legacyShowContent(args[1], resp)
if !ok {
return m.legacyUsageError(fmt.Sprintf("Unknown command `/show %s`.", args[1]))
}
m.entries = append(m.entries, newSlashEntry(content))
m.status = "show " + args[1]
return *m, nil
}
func (m *chatModel) handleLegacyHelpCommand(input string) (tea.Model, tea.Cmd) {
args := strings.Fields(input)
if len(args) != 2 {
m.entries = append(m.entries, newSlashEntry(m.helpSummary()))
return *m, nil
}
switch strings.TrimPrefix(args[1], "/") {
case "set":
m.entries = append(m.entries, newSlashEntry(legacySetUsage()))
case "show":
m.entries = append(m.entries, newSlashEntry(legacyShowUsage()))
case "shortcut", "shortcuts":
m.entries = append(m.entries, newSlashEntry(legacyShortcutUsage()))
default:
m.entries = append(m.entries, newSlashEntry(m.helpSummary()))
}
return *m, nil
}
func (m *chatModel) legacyShowResponse() (*api.ShowResponse, error) {
client, ok := m.opts.Client.(chatShowClient)
if !ok || client == nil {
return nil, fmt.Errorf("model show is unavailable")
}
ctx := m.ctx
if ctx == nil {
ctx = context.Background()
}
return client.Show(ctx, &api.ShowRequest{
Model: m.opts.Model,
Name: m.opts.Model,
Options: m.opts.Options,
})
}
func (m chatModel) legacyShowContent(kind string, resp *api.ShowResponse) (string, bool) {
switch kind {
case "info":
return legacyShowInfo(m.opts.Model, resp), true
case "license":
if strings.TrimSpace(resp.License) == "" {
return "No license was specified for this model.", true
}
return resp.License, true
case "modelfile":
return resp.Modelfile, true
case "parameters":
return legacyShowParameters(resp.Parameters, m.opts.Options), true
case "system":
if strings.TrimSpace(resp.System) == "" {
return "No system message was specified for this model.", true
}
return resp.System, true
case "template":
if strings.TrimSpace(resp.Template) == "" {
return "No prompt template was specified for this model.", true
}
return resp.Template, true
default:
return "", false
}
}
func legacyShowInfo(modelName string, resp *api.ShowResponse) string {
var b strings.Builder
b.WriteString("**Model info**\n\n")
if strings.TrimSpace(modelName) != "" {
fmt.Fprintf(&b, "- model: `%s`\n", modelName)
}
if !resp.ModifiedAt.IsZero() {
fmt.Fprintf(&b, "- modified: `%s`\n", resp.ModifiedAt.Format(time.RFC3339))
}
if resp.Details.Family != "" {
fmt.Fprintf(&b, "- family: `%s`\n", resp.Details.Family)
}
if resp.Details.ParameterSize != "" {
fmt.Fprintf(&b, "- parameters: `%s`\n", resp.Details.ParameterSize)
}
if resp.Details.QuantizationLevel != "" {
fmt.Fprintf(&b, "- quantization: `%s`\n", resp.Details.QuantizationLevel)
}
if resp.Details.ContextLength > 0 {
fmt.Fprintf(&b, "- context length: `%d`\n", resp.Details.ContextLength)
}
if len(resp.Capabilities) > 0 {
parts := make([]string, 0, len(resp.Capabilities))
for _, capability := range resp.Capabilities {
parts = append(parts, string(capability))
}
sort.Strings(parts)
fmt.Fprintf(&b, "- capabilities: `%s`\n", strings.Join(parts, ", "))
}
if b.Len() == len("**Model info**\n\n") {
b.WriteString("No additional model information was returned.")
}
return strings.TrimRight(b.String(), "\n")
}
func legacyShowParameters(modelParameters string, userOptions map[string]any) string {
var b strings.Builder
b.WriteString("**Model defined parameters**\n\n")
if strings.TrimSpace(modelParameters) == "" {
b.WriteString("No additional parameters were specified for this model.\n")
} else {
b.WriteString("```text\n")
b.WriteString(strings.TrimSpace(modelParameters))
b.WriteString("\n```\n")
}
if len(userOptions) > 0 {
b.WriteString("\n**User defined parameters**\n\n")
keys := make([]string, 0, len(userOptions))
for key := range userOptions {
keys = append(keys, key)
}
sort.Strings(keys)
for _, key := range keys {
fmt.Fprintf(&b, "- `%s`: `%v`\n", key, userOptions[key])
}
}
return strings.TrimRight(b.String(), "\n")
}
func (m *chatModel) legacyUsageError(message string) (tea.Model, tea.Cmd) {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: message, err: message}))
m.status = "error"
return *m, nil
}
func (m *chatModel) refreshContextEstimate() {
m.contextTokens = m.estimatePromptTokens(m.messages, "")
m.contextEstimate = true
}
func bullets(title string, items ...string) string {
parts := make([]string, 0, len(items)+2)
parts = append(parts, title, "")
for _, item := range items {
parts = append(parts, "- "+item)
}
return strings.Join(parts, "\n")
}
func legacySetUsage() string {
return bullets("**Legacy set commands**",
"`/set parameter <name> <value...>`",
"`/set format json`",
"`/set noformat`",
"`/set verbose`",
"`/set quiet`",
"`/set think [low|medium|high|max]`",
"`/set nothink`",
)
}
func legacyParameterUsage() string {
return bullets("**Legacy parameters**",
"`/set parameter seed <int>`",
"`/set parameter num_predict <int>`",
"`/set parameter top_k <int>`",
"`/set parameter top_p <float>`",
"`/set parameter min_p <float>`",
"`/set parameter num_ctx <int>`",
"`/set parameter temperature <float>`",
"`/set parameter repeat_penalty <float>`",
"`/set parameter repeat_last_n <int>`",
"`/set parameter num_gpu <int>`",
"`/set parameter stop <string> <string> ...`",
)
}
func legacyShowUsage() string {
return bullets("**Legacy show commands**",
"`/show info`",
"`/show license`",
"`/show modelfile`",
"`/show parameters`",
"`/show system`",
"`/show template`",
)
}
func legacyShortcutUsage() string {
return bullets("**Shortcuts**",
"`ctrl+o`: toggle tool output",
"`shift+enter`: insert a newline",
"`shift+tab`: toggle permission mode",
"`cmd+backspace`, `option+backspace`, `ctrl+w`: delete previous word",
"`ctrl+c`: clear input, cancel current response, or confirm quit",
)
}
+15
View File
@@ -0,0 +1,15 @@
package chat
import "strings"
func renderMarkdownForView(markdown string, width int) string {
return strings.Join(wrapChatText(markdown, width), "\n")
}
func splitRenderedBody(body string) []string {
body = strings.TrimRight(body, "\n")
if body == "" {
return []string{""}
}
return strings.Split(body, "\n")
}
File diff suppressed because it is too large. Load diff
+881
View File
@@ -0,0 +1,881 @@
package chat
import (
"context"
"errors"
"os"
"path/filepath"
"slices"
"strconv"
"strings"
"testing"
"time"
tea "github.com/charmbracelet/bubbletea"
coreagent "github.com/ollama/ollama/agent"
agentstore "github.com/ollama/ollama/agent/store"
"github.com/ollama/ollama/api"
)
func TestChatHistoryCommandShowsPromptMessages(t *testing.T) {
args := api.NewToolCallFunctionArguments()
args.Set("command", "pwd")
m := chatModel{
input: []rune("/history"),
opts: Options{SystemPrompt: "You are Ollama."},
messages: []api.Message{
{Role: "user", Content: "where am i?"},
{
Role: "assistant",
Thinking: "Need to inspect cwd.",
ToolCalls: []api.ToolCall{{
ID: "call-1",
Function: api.ToolCallFunction{
Name: "bash",
Arguments: args,
},
}},
},
{Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: "/tmp/project\n"},
{Role: "assistant", Content: "You are in /tmp/project."},
},
}
updated, cmd := m.handleSubmit()
if cmd == nil {
t.Fatal("history command should enter the history screen")
}
fm := updated.(chatModel)
if len(fm.entries) != 0 {
t.Fatalf("entries = %d, want 0", len(fm.entries))
}
if fm.historyPopup == nil {
t.Fatal("history popup was not opened")
}
if got := len(fm.historyPopup.messages); got != 5 {
t.Fatalf("history messages = %d, want 5", got)
}
if fm.historyPopup.messages[0].Role != "system" || fm.historyPopup.messages[0].Content != "You are Ollama." {
t.Fatalf("system prompt history message = %#v", fm.historyPopup.messages[0])
}
if fm.historyPopup.messages[2].Thinking != "Need to inspect cwd." || len(fm.historyPopup.messages[2].ToolCalls) != 1 {
t.Fatalf("assistant tool history message = %#v", fm.historyPopup.messages[2])
}
fm.width = 120
fm.height = 40
rendered := stripANSI(fm.View())
for _, want := range []string{
"Message history",
"estimated prompt:",
"system",
"content: You are Ollama.",
"user",
"content: where am i?",
"assistant",
"thinking: Need to inspect cwd.",
"tool calls:",
"call-1 Bash",
"args:",
"\"command\": \"pwd\"",
"tool",
"tool: bash",
"tool call: call-1",
"/tmp/project",
"content: You are in /tmp/project.",
} {
if !strings.Contains(rendered, want) {
t.Fatalf("rendered history missing %q:\n%s", want, rendered)
}
}
updated, cmd = fm.Update(tea.KeyMsg{Type: tea.KeyEsc})
fm = updated.(chatModel)
if cmd == nil {
t.Fatal("closing history should exit the history screen")
}
if fm.historyPopup != nil {
t.Fatal("history popup should close on escape")
}
}
func TestChatHistoryCommandHandlesEmptyHistory(t *testing.T) {
m := chatModel{input: []rune("/history")}
updated, cmd := m.handleSubmit()
if cmd == nil {
t.Fatal("history command should enter the history screen")
}
fm := updated.(chatModel)
if len(fm.entries) != 0 || fm.historyPopup == nil || len(fm.historyPopup.messages) != 0 {
t.Fatalf("history output entries=%#v popup=%#v", fm.entries, fm.historyPopup)
}
fm.width = 80
fm.height = 20
if view := stripANSI(fm.View()); !strings.Contains(view, "No messages yet.") {
t.Fatalf("history popup view missing empty state: %q", view)
}
}
func TestChatRawCommandShowsRequestJSON(t *testing.T) {
registry := coreagent.NewRegistry()
registry.Register(chatTestTool{})
m := chatModel{
input: []rune("/raw"),
opts: Options{
Model: "llama3.2",
SystemPrompt: "You are Ollama.",
Format: "json",
Tools: registry,
Options: map[string]any{"temperature": 0.1},
Think: &api.ThinkValue{Value: false},
ContextWindowTokens: 4096,
},
messages: []api.Message{{
Role: "user",
Content: "what is in this image?",
Images: []api.ImageData{[]byte("abc")},
}},
width: 120,
height: 40,
}
updated, cmd := m.handleSubmit()
if cmd == nil {
t.Fatal("raw command should enter the raw request screen")
}
fm := updated.(chatModel)
if fm.historyPopup == nil {
t.Fatal("raw popup was not opened")
}
if fm.historyPopup.title != "Raw request" {
t.Fatalf("raw popup title = %q, want Raw request", fm.historyPopup.title)
}
rendered := stripANSI(fm.View())
for _, want := range []string{
"Raw request",
"estimated prompt:",
"\"model\": \"llama3.2\"",
"\"role\": \"system\"",
"\"content\": \"You are Ollama.\"",
"\"format\": \"json\"",
"\"tools\":",
"\"options\":",
"\"temperature\": 0.1",
"\"think\": false",
"[image data: 3 bytes]",
} {
if !strings.Contains(rendered, want) {
t.Fatalf("rendered raw request missing %q:\n%s", want, rendered)
}
}
if strings.Contains(rendered, "YWJj") {
t.Fatalf("raw request should redact image base64:\n%s", rendered)
}
}
func TestChatRawCommandSavesRequestJSON(t *testing.T) {
dir := t.TempDir()
m := chatModel{
input: []rune("/raw > request"),
workingDir: dir,
opts: Options{
Model: "llama3.2",
SystemPrompt: "You are Ollama.",
},
messages: []api.Message{{Role: "user", Content: "hello"}},
}
updated, cmd := m.handleSubmit()
if cmd != nil {
t.Fatal("raw redirect should write without opening a popup")
}
fm := updated.(chatModel)
if fm.status != "raw saved" {
t.Fatalf("status = %q, want raw saved", fm.status)
}
if fm.historyPopup != nil {
t.Fatal("raw redirect should not open the raw popup")
}
data, err := os.ReadFile(filepath.Join(dir, "request.json"))
if err != nil {
t.Fatal(err)
}
raw := string(data)
for _, want := range []string{
"\"model\": \"llama3.2\"",
"\"role\": \"system\"",
"\"content\": \"You are Ollama.\"",
"\"role\": \"user\"",
"\"content\": \"hello\"",
} {
if !strings.Contains(raw, want) {
t.Fatalf("saved raw request missing %q:\n%s", want, raw)
}
}
if got := fm.entries[len(fm.entries)-1].content; !strings.Contains(got, "request.json") {
t.Fatalf("save entry = %q, want saved filename", got)
}
}
func TestChatRawCommandRejectsPathRedirect(t *testing.T) {
dir := t.TempDir()
m := chatModel{
input: []rune("/raw > ../request"),
workingDir: dir,
opts: Options{Model: "llama3.2"},
}
updated, cmd := m.handleSubmit()
if cmd != nil {
t.Fatal("invalid raw redirect should not open a popup")
}
fm := updated.(chatModel)
if fm.status != "error" {
t.Fatalf("status = %q, want error", fm.status)
}
if _, err := os.Stat(filepath.Join(filepath.Dir(dir), "request.json")); !os.IsNotExist(err) {
t.Fatalf("raw redirect wrote outside working dir, stat err = %v", err)
}
}
func TestChatHistoryCommandStartsAtBottom(t *testing.T) {
messages := make([]api.Message, 0, 18)
for i := range 18 {
messages = append(messages, api.Message{Role: "user", Content: "prompt " + strconv.Itoa(i)})
}
m := chatModel{
input: []rune("/history"),
messages: messages,
width: 80,
height: 10,
}
updated, cmd := m.handleSubmit()
if cmd == nil {
t.Fatal("history command should enter the history screen")
}
fm := updated.(chatModel)
view := stripANSI(fm.View())
if !strings.Contains(view, "prompt 17") {
t.Fatalf("history popup should start at latest messages:\n%s", view)
}
if strings.Contains(view, "prompt 0") {
t.Fatalf("history popup started at oldest messages:\n%s", view)
}
if title := strings.Split(view, "\n")[0]; strings.Contains(title, "/") {
t.Fatalf("history popup should not show scroll counter in title:\n%s", view)
}
}
func TestChatHistoryMouseWheelScrollsPopup(t *testing.T) {
messages := make([]api.Message, 0, 18)
for i := range 18 {
messages = append(messages, api.Message{Role: "user", Content: "prompt " + strconv.Itoa(i)})
}
m := chatModel{
historyPopup: &chatHistoryPopup{messages: messages, stickToBottom: true},
width: 80,
height: 10,
}
maxScroll := m.historyPopupMaxScroll()
if maxScroll == 0 {
t.Fatal("test setup should produce scrollable history")
}
updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseWheelUp})
m = updated.(chatModel)
if m.historyPopup.stickToBottom {
t.Fatal("mouse wheel should detach history popup from bottom")
}
if m.historyPopup.scroll >= maxScroll {
t.Fatalf("mouse wheel up should scroll toward older history, got %d max %d", m.historyPopup.scroll, maxScroll)
}
updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseWheelDown})
m = updated.(chatModel)
if m.historyPopup.scroll != maxScroll {
t.Fatalf("mouse wheel down should scroll back toward latest history, got %d want %d", m.historyPopup.scroll, maxScroll)
}
}
func TestChatHistoryMouseDragSelectsTextWithoutAutoCopy(t *testing.T) {
m := chatModel{
historyPopup: &chatHistoryPopup{
messages: []api.Message{{Role: "user", Content: "alpha beta"}},
stickToBottom: true,
},
width: 80,
height: 10,
}
top, _ := m.historyPopupLayout()
contentY := top + 1
contentX := len(" content: ")
updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionPress, X: contentX, Y: contentY})
m = updated.(chatModel)
updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionMotion, X: contentX + len("alpha"), Y: contentY})
m = updated.(chatModel)
if got := m.selectedHistoryPopupText(); got != "alpha" {
t.Fatalf("selected history text = %q, want alpha", got)
}
if !m.historyPopup.selection.active {
t.Fatal("history selection should stay active during drag")
}
if !m.historyPopup.selection.dragging {
t.Fatal("history selection should track drag before release")
}
updated, cmd := m.Update(tea.MouseMsg{Type: tea.MouseRelease, Action: tea.MouseActionRelease, X: contentX + len("alpha"), Y: contentY})
m = updated.(chatModel)
if cmd != nil {
t.Fatal("mouse release should not auto-copy selected text")
}
if !m.historyPopup.selection.active {
t.Fatal("history selection should stay visible on release")
}
if m.historyPopup.selection.dragging {
t.Fatal("history selection should stop tracking drag on release")
}
if got := m.selectedHistoryPopupText(); got != "alpha" {
t.Fatalf("selected history text after release = %q, want alpha", got)
}
}
func TestChatHistoryMouseDragSelectionUsesDisplayColumns(t *testing.T) {
m := chatModel{
historyPopup: &chatHistoryPopup{
messages: []api.Message{{Role: "user", Content: "a界b"}},
stickToBottom: true,
},
width: 80,
height: 10,
}
top, _ := m.historyPopupLayout()
contentY := top + 1
contentX := len(" content: ")
updated, _ := m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionPress, X: contentX, Y: contentY})
m = updated.(chatModel)
updated, _ = m.Update(tea.MouseMsg{Type: tea.MouseLeft, Button: tea.MouseButtonLeft, Action: tea.MouseActionMotion, X: contentX + 3, Y: contentY})
m = updated.(chatModel)
if got := m.selectedHistoryPopupText(); got != "a界" {
t.Fatalf("selected history text = %q, want a界", got)
}
}
func TestChatHistoryCommandFormatsMultilineContentWithLabel(t *testing.T) {
m := chatModel{
input: []rune("/history"),
messages: []api.Message{{Role: "assistant", Content: "first\nsecond"}},
}
updated, cmd := m.handleSubmit()
if cmd == nil {
t.Fatal("history command should enter the history screen")
}
fm := updated.(chatModel)
if fm.historyPopup == nil {
t.Fatal("history popup was not opened")
}
fm.width = 80
fm.height = 20
view := stripANSI(fm.View())
if !strings.Contains(view, "content:") || !strings.Contains(view, "first") || !strings.Contains(view, "second") {
t.Fatalf("history should label multiline content before block:\n%s", view)
}
}
func TestChatResumeCommandOpensPicker(t *testing.T) {
store := &chatResumeTestStore{
chats: []agentstore.ChatSummary{{
ID: "chat-1",
Title: "Research Parth Sareen online",
Model: "llama3.2",
UpdatedAt: time.Now().Add(-time.Hour),
MessageCount: 2,
ApproxBytes: 18 * 1024,
}},
byID: map[string]*agentstore.AgentChat{},
}
m := chatModel{
ctx: context.Background(),
input: []rune("/resume"),
width: 100,
height: 20,
opts: Options{Store: store},
}
updated, cmd := m.handleSubmit()
if cmd != nil {
t.Fatal("resume command should not return a command")
}
m = updated.(chatModel)
if m.resumePicker == nil {
t.Fatal("resume picker was not opened")
}
view := stripANSI(m.View())
if !strings.Contains(view, "Resume chat: type to filter") ||
!strings.Contains(view, "Research Parth Sareen online") ||
!strings.Contains(view, "llama3.2") {
t.Fatalf("resume picker view missing content: %q", view)
}
if strings.Contains(view, "Search...") {
t.Fatalf("resume picker should render inline without full search box: %q", view)
}
if !strings.Contains(view, " █") {
t.Fatalf("inline resume picker should keep input box visible: %q", view)
}
}
func TestChatResumePickerFallsBackToFullFrameWhenSmall(t *testing.T) {
store := &chatResumeTestStore{
chats: []agentstore.ChatSummary{{
ID: "chat-1",
Title: "Research Parth Sareen online",
Model: "llama3.2",
}},
byID: map[string]*agentstore.AgentChat{},
}
m := chatModel{
ctx: context.Background(),
input: []rune("/resume"),
width: 44,
height: 10,
opts: Options{Store: store},
}
updated, cmd := m.handleSubmit()
if cmd != nil {
t.Fatal("resume command should not return a command")
}
m = updated.(chatModel)
view := stripANSI(m.View())
if !strings.Contains(view, "Resume session") || !strings.Contains(view, "Search...") {
t.Fatalf("small resume picker should use full-frame fallback: %q", view)
}
}
func TestChatModelCommandOpensPicker(t *testing.T) {
m := chatModel{
ctx: context.Background(),
input: []rune("/model"),
width: 100,
height: 20,
opts: Options{
Model: "llama3.2",
ContextWindowTokens: 131072,
ModelOptions: func(context.Context) ([]ModelOption, error) {
return []ModelOption{
{Name: "kimi-k2.6:cloud", Description: "cloud coding"},
{Name: "llama3.2", Description: "local"},
}, nil
},
},
}
updated, cmd := m.handleSubmit()
if cmd != nil {
t.Fatal("model command should not return a command")
}
m = updated.(chatModel)
if m.modelPicker == nil {
t.Fatal("model picker was not opened")
}
view := stripANSI(m.View())
if !strings.Contains(view, "Select model: type to filter") ||
!strings.Contains(view, "kimi-k2.6:cloud") ||
!strings.Contains(view, "llama3.2") ||
!strings.Contains(view, "current") {
t.Fatalf("model picker view missing content: %q", view)
}
if strings.Contains(view, "Search...") {
t.Fatalf("model picker should render inline without full search box: %q", view)
}
if !strings.Contains(view, " █") {
t.Fatalf("inline model picker should keep input box visible: %q", view)
}
}
func TestChatModelPickerFallsBackToFullFrameWhenSmall(t *testing.T) {
m := chatModel{
ctx: context.Background(),
input: []rune("/model"),
width: 44,
height: 10,
opts: Options{
Model: "llama3.2",
ModelOptions: func(context.Context) ([]ModelOption, error) {
return []ModelOption{
{Name: "kimi-k2.6:cloud", Description: "cloud coding"},
{Name: "llama3.2", Description: "local"},
}, nil
},
},
}
updated, cmd := m.handleSubmit()
if cmd != nil {
t.Fatal("model command should not return a command")
}
m = updated.(chatModel)
view := stripANSI(m.View())
if !strings.Contains(view, "Switch model") || !strings.Contains(view, "Search...") {
t.Fatalf("small model picker should use full-frame fallback: %q", view)
}
}
func TestChatModelPickerShowsRecommendedModelsFirst(t *testing.T) {
models := normalizeModelOptions([]ModelOption{
{Name: "llama3.2", Description: "local"},
{Name: "kimi-k2.6:cloud", Description: "cloud coding", Recommended: true},
{Name: "qwen3.5:cloud", Description: "cloud reasoning", Recommended: true},
{Name: "gemma4", Description: "local"},
})
got := make([]string, 0, len(models))
for _, model := range models {
got = append(got, model.Name)
}
want := []string{"kimi-k2.6:cloud", "qwen3.5:cloud", "llama3.2", "gemma4"}
if !slices.Equal(got, want) {
t.Fatalf("model order = %#v, want %#v", got, want)
}
}
func TestChatModelPickerRanksClosestFilteredModelFirst(t *testing.T) {
models := normalizeModelOptions([]ModelOption{
{Name: "gemma3:27b", Description: "recommended but longer", Recommended: true},
{Name: "llama3.2", Description: "mentions gemm in description"},
{Name: "gemma4:27b", Description: "longer local"},
{Name: "gemma4", Description: "short local"},
})
picker := newChatModelPicker(models, "", "gemm")
filtered := picker.filtered()
got := make([]string, 0, len(filtered))
for _, model := range filtered {
got = append(got, model.Name)
}
want := []string{"gemma4", "gemma3:27b", "gemma4:27b", "llama3.2"}
if !slices.Equal(got, want) {
t.Fatalf("filtered model order = %#v, want %#v", got, want)
}
}
func TestChatModelPickerFiltersAndSwitchesModel(t *testing.T) {
var savedModel string
store := &chatResumeTestStore{}
originalMessages := []api.Message{{Role: "user", Content: "keep me"}}
m := chatModel{
ctx: context.Background(),
chatID: "chat-1",
input: []rune("/model qwen"),
width: 100,
height: 20,
messages: slices.Clone(originalMessages),
opts: Options{
Model: "llama3.2",
Store: store,
ModelOptions: func(context.Context) ([]ModelOption, error) {
return []ModelOption{
{Name: "llama3.2", Description: "local"},
{Name: "qwen3.5:cloud", Description: "cloud reasoning"},
}, nil
},
ToolRegistryForModel: func(ctx context.Context, model string) *coreagent.Registry {
if model != "qwen3.5:cloud" {
t.Fatalf("tool registry model = %q, want qwen3.5:cloud", model)
}
registry := coreagent.NewRegistry()
registry.Register(chatTestTool{})
return registry
},
ContextWindowTokensForModel: func(ctx context.Context, model string, fallback int) int {
if model != "qwen3.5:cloud" {
t.Fatalf("context model = %q, want qwen3.5:cloud", model)
}
if fallback != 0 {
t.Fatalf("context fallback = %d, want 0 after model switch", fallback)
}
return 262144
},
SystemPromptForModel: func(ctx context.Context, model string, registry *coreagent.Registry) string {
if model != "qwen3.5:cloud" {
t.Fatalf("system prompt model = %q, want qwen3.5:cloud", model)
}
if registry == nil || !registry.Has("fake_tool") {
t.Fatalf("system prompt registry missing fake tool: %#v", registry)
}
return "system for " + model
},
OnModelSelected: func(ctx context.Context, model string) error {
savedModel = model
return nil
},
},
}
updated, cmd := m.handleSubmit()
if cmd != nil {
t.Fatal("model command should not return a command")
}
m = updated.(chatModel)
if m.modelPicker == nil || m.modelPicker.filter != "qwen" {
t.Fatalf("model picker = %#v, want qwen filter", m.modelPicker)
}
if view := stripANSI(m.View()); !strings.Contains(view, "qwen3.5:cloud") || strings.Contains(view, "llama3.2") {
t.Fatalf("filtered model picker view = %q", view)
}
updated, cmd = m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = updated.(chatModel)
if cmd != nil {
t.Fatal("switching models should not start a command")
}
if m.modelPicker != nil {
t.Fatal("model picker should close after selection")
}
if m.status != "ready" || m.notificationLine() != "" {
t.Fatalf("model switch should not show action status, status=%q notification=%q", m.status, m.notificationLine())
}
if m.opts.Model != "qwen3.5:cloud" {
t.Fatalf("model = %q, want qwen3.5:cloud", m.opts.Model)
}
if m.chatID != "chat-1" {
t.Fatalf("chatID = %q, want chat-1", m.chatID)
}
if len(m.messages) != len(originalMessages) || m.messages[0].Content != originalMessages[0].Content {
t.Fatalf("messages changed on model switch: %#v", m.messages)
}
if len(m.entries) != 0 {
t.Fatalf("model switch should not append transcript entries: %#v", m.entries)
}
if got := store.setModels["chat-1"]; got != "qwen3.5:cloud" {
t.Fatalf("persisted chat model = %q, want qwen3.5:cloud", got)
}
if savedModel != "qwen3.5:cloud" {
t.Fatalf("saved model = %q, want qwen3.5:cloud", savedModel)
}
if m.opts.Tools == nil || !m.opts.Tools.Has("fake_tool") {
t.Fatalf("tools registry was not rebuilt for model: %#v", m.opts.Tools)
}
if m.opts.ContextWindowTokens != 262144 {
t.Fatalf("context window = %d, want 262144", m.opts.ContextWindowTokens)
}
if m.opts.SystemPrompt != "system for qwen3.5:cloud" {
t.Fatalf("system prompt = %q", m.opts.SystemPrompt)
}
}
func TestChatModelSelectionPersistsBeforeSwitching(t *testing.T) {
originalTools := coreagent.NewRegistry()
errStore := errors.New("write failed")
m := chatModel{
ctx: context.Background(),
chatID: "chat-1",
opts: Options{
Model: "llama3.2",
Store: &chatResumeTestStore{setModelErr: errStore},
Tools: originalTools,
SystemPrompt: "system for llama3.2",
ContextWindowTokens: 8192,
ToolRegistryForModel: func(context.Context, string) *coreagent.Registry {
t.Fatal("tool registry should not rebuild when model persistence fails")
return nil
},
SystemPromptForModel: func(context.Context, string, *coreagent.Registry) string {
t.Fatal("system prompt should not rebuild when model persistence fails")
return ""
},
},
messages: []api.Message{{Role: "user", Content: "history"}},
}
err := m.applyModelSelection("qwen3", true)
if !errors.Is(err, errStore) {
t.Fatalf("error = %v, want %v", err, errStore)
}
if m.opts.Model != "llama3.2" {
t.Fatalf("model = %q, want llama3.2", m.opts.Model)
}
if m.opts.Tools != originalTools {
t.Fatal("tools changed after failed model persistence")
}
if m.opts.SystemPrompt != "system for llama3.2" {
t.Fatalf("system prompt = %q, want original", m.opts.SystemPrompt)
}
if m.opts.ContextWindowTokens != 8192 {
t.Fatalf("context window = %d, want 8192", m.opts.ContextWindowTokens)
}
if len(m.messages) != 1 || m.messages[0].Content != "history" {
t.Fatalf("messages changed after failed model persistence: %#v", m.messages)
}
}
func TestChatModelSelectionStartsBackgroundPreload(t *testing.T) {
m := chatModel{
ctx: context.Background(),
opts: Options{
Model: "llama3.2",
PreloadModel: func(context.Context, string, *api.ThinkValue) error {
return nil
},
},
}
if err := m.applyModelSelection("qwen3", false); err != nil {
t.Fatal(err)
}
cmd := m.startModelPreload("qwen3")
if cmd == nil {
t.Fatal("model switch should start background preload when configured")
}
if m.preloadingModel != "qwen3" {
t.Fatalf("preloadingModel = %q, want qwen3", m.preloadingModel)
}
}
func TestChatModelSwitchNextRunKeepsHistory(t *testing.T) {
client := &chatCaptureClient{}
store := &chatResumeTestStore{}
history := []api.Message{
{Role: "user", Content: "old question"},
{Role: "assistant", Content: "old answer"},
}
m := chatModel{
ctx: context.Background(),
chatID: "chat-1",
messages: slices.Clone(history),
input: []rune("continue"),
opts: Options{
Model: "llama3.2",
Store: store,
Client: client,
SystemPromptForModel: func(_ context.Context, model string, _ *coreagent.Registry) string {
return "system for " + model
},
},
}
if err := m.applyModelSelection("qwen3", true); err != nil {
t.Fatal(err)
}
updated, cmd := m.handleSubmit()
m = updated.(chatModel)
if cmd == nil {
t.Fatal("next prompt should start a model run")
}
done := waitForRunDone(t, m.events)
if done.err != nil {
t.Fatal(done.err)
}
if len(client.requests) != 1 {
t.Fatalf("requests = %d, want 1", len(client.requests))
}
req := client.requests[0]
if req.Model != "qwen3" {
t.Fatalf("request model = %q, want qwen3", req.Model)
}
if len(req.Messages) != 4 {
t.Fatalf("request messages = %#v, want system + 2 history + new user", req.Messages)
}
if req.Messages[0].Role != "system" || req.Messages[0].Content != "system for qwen3" {
t.Fatalf("system message = %#v", req.Messages[0])
}
for i, want := range history {
got := req.Messages[i+1]
if got.Role != want.Role || got.Content != want.Content {
t.Fatalf("history message %d = %#v, want %#v", i, got, want)
}
}
if req.Messages[3].Role != "user" || req.Messages[3].Content != "continue" {
t.Fatalf("new user message = %#v", req.Messages[3])
}
if got := store.setModels["chat-1"]; got != "qwen3" {
t.Fatalf("persisted chat model = %q, want qwen3", got)
}
}
func TestChatResumePickerFiltersAndLoadsSelection(t *testing.T) {
store := &chatResumeTestStore{
chats: []agentstore.ChatSummary{
{ID: "chat-1", Title: "First chat", Model: "llama3.2", UpdatedAt: time.Now().Add(-2 * time.Hour), MessageCount: 2, ApproxBytes: 2048},
{ID: "chat-2", Title: "Second chat", Model: "qwen3", UpdatedAt: time.Now().Add(-time.Hour), MessageCount: 2, ApproxBytes: 4096},
},
byID: map[string]*agentstore.AgentChat{
"chat-2": {
ID: "chat-2",
Title: "Second chat",
Model: "qwen3",
Messages: []api.Message{
{Role: "user", Content: "resume this"},
{Role: "assistant", Content: "loaded"},
},
},
},
}
m := chatModel{
ctx: context.Background(),
input: []rune("/resume"),
queued: []string{"old queued prompt"},
width: 100,
height: 20,
opts: Options{
Model: "llama3.2",
Store: store,
ToolRegistryForModel: func(ctx context.Context, model string) *coreagent.Registry {
if model != "qwen3" {
t.Fatalf("tool registry model = %q, want qwen3", model)
}
registry := coreagent.NewRegistry()
registry.Register(chatTestTool{})
return registry
},
},
}
updated, _ := m.handleSubmit()
m = updated.(chatModel)
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("second")})
m = updated.(chatModel)
view := stripANSI(m.View())
if !strings.Contains(view, "Second chat") || strings.Contains(view, "First chat") {
t.Fatalf("filtered resume picker view = %q", view)
}
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = updated.(chatModel)
if cmd == nil {
t.Fatal("loading a saved chat should flush the resumed transcript")
}
if m.resumePicker != nil {
t.Fatal("resume picker should close after selection")
}
if m.chatID != "chat-2" {
t.Fatalf("chatID = %q, want chat-2", m.chatID)
}
if m.opts.Model != "qwen3" {
t.Fatalf("model = %q, want qwen3", m.opts.Model)
}
if m.opts.Tools == nil || !m.opts.Tools.Has("fake_tool") {
t.Fatalf("tools registry was not rebuilt for resumed model: %#v", m.opts.Tools)
}
if len(m.queued) != 0 {
t.Fatalf("queued prompts should be cleared on resume: %#v", m.queued)
}
if len(m.messages) != 2 || m.messages[0].Content != "resume this" {
t.Fatalf("messages = %#v", m.messages)
}
if len(m.entries) != 2 || m.entries[0].content != "resume this" {
t.Fatalf("entries = %#v", m.entries)
}
}
+307
View File
@@ -0,0 +1,307 @@
package chat
import (
"context"
"net/url"
"os"
"path/filepath"
"strings"
"testing"
tea "github.com/charmbracelet/bubbletea"
)
func TestChatStartRunAttachesDroppedImagePath(t *testing.T) {
fp := writeTestPNG(t)
m := chatModel{
ctx: context.Background(),
opts: Options{
Model: "test",
Client: chatTestClient{},
MultiModal: true,
},
}
updated, cmd := m.startRun("describe " + fp)
m = updated.(chatModel)
if cmd == nil {
t.Fatal("startRun should return a command")
}
if len(m.liveMessages) != 1 {
t.Fatalf("liveMessages = %d, want 1", len(m.liveMessages))
}
if got := m.liveMessages[0].Content; got != "describe" {
t.Fatalf("content = %q, want describe", got)
}
if got := len(m.liveMessages[0].Images); got != 1 {
t.Fatalf("images = %d, want 1", got)
}
if len(m.entries) == 0 {
t.Fatal("missing user transcript entry")
}
entry := m.entries[0].content
if strings.Contains(entry, fp) {
t.Fatalf("transcript entry should hide local file path: %q", entry)
}
if !strings.Contains(entry, "describe") || !strings.Contains(entry, "[attached 1 file]") {
t.Fatalf("transcript entry = %q, want prompt plus attachment note", entry)
}
}
func TestChatStartRunAttachesDroppedFileURL(t *testing.T) {
fp := writeTestPNG(t)
fileURL := (&url.URL{Scheme: "file", Path: fp}).String()
m := chatModel{
ctx: context.Background(),
opts: Options{
Model: "test",
Client: chatTestClient{},
MultiModal: true,
},
}
updated, _ := m.startRun(fileURL)
m = updated.(chatModel)
if got := m.liveMessages[0].Content; got != "" {
t.Fatalf("content = %q, want empty prompt after extracting file URL", got)
}
if got := len(m.liveMessages[0].Images); got != 1 {
t.Fatalf("images = %d, want 1", got)
}
if got := m.entries[0].content; got != "[attached 1 file]" {
t.Fatalf("transcript entry = %q, want attachment-only note", got)
}
}
func TestChatPasteImagePathAttachesOnSubmit(t *testing.T) {
fp := writeTestPNG(t)
m := chatModel{
ctx: context.Background(),
opts: Options{
Model: "test",
Client: chatTestClient{},
MultiModal: true,
},
}
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true})
m = updated.(chatModel)
if got := string(m.input); got != "describe [Image #0]" {
t.Fatalf("pasted path input = %q, want placeholder", got)
}
if got := m.notificationLine(); got != "" {
t.Fatalf("notification = %q, want no attachment notification", got)
}
if got := string(m.input); strings.Contains(got, fp) {
t.Fatalf("pasted path should be hidden behind placeholder, input = %q", got)
}
if completions := m.slashCompletions(); len(completions) != 0 {
t.Fatalf("placeholder input should not show slash completions: %#v", completions)
}
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = updated.(chatModel)
if cmd == nil {
t.Fatal("submit should start a run")
}
if got := m.liveMessages[0].Content; got != "describe [Image #0]" {
t.Fatalf("content = %q, want prompt with placeholder", got)
}
if got := len(m.liveMessages[0].Images); got != 1 {
t.Fatalf("images = %d, want 1", got)
}
if strings.Contains(m.entries[0].content, fp) {
t.Fatalf("transcript entry should hide pasted file path: %q", m.entries[0].content)
}
if !strings.Contains(m.entries[0].content, "[Image #0]") {
t.Fatalf("transcript entry should show placeholder: %q", m.entries[0].content)
}
}
func TestChatPasteImagePathAfterSwitchingToMultimodalModel(t *testing.T) {
fp := writeTestPNG(t)
m := chatModel{
ctx: context.Background(),
input: []rune("/model vision"),
opts: Options{
Model: "text",
ModelOptions: func(context.Context) ([]ModelOption, error) {
return []ModelOption{
{Name: "text", Description: "local"},
{Name: "vision", Description: "local vision"},
}, nil
},
MultiModalForModel: func(_ context.Context, model string) bool {
return model == "vision"
},
},
}
updated, cmd := m.handleSubmit()
if cmd != nil {
t.Fatal("model picker should not return a command")
}
m = updated.(chatModel)
updated, cmd = m.Update(tea.KeyMsg{Type: tea.KeyEnter})
if cmd != nil {
t.Fatal("model switch should not preload without a preload hook")
}
m = updated.(chatModel)
if m.opts.Model != "vision" {
t.Fatalf("model = %q, want vision", m.opts.Model)
}
if !m.opts.MultiModal {
t.Fatal("switching to a multimodal model should enable image paste handling")
}
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true})
m = updated.(chatModel)
if got := string(m.input); got != "describe [Image #0]" {
t.Fatalf("pasted path input = %q, want placeholder", got)
}
if strings.Contains(string(m.input), fp) {
t.Fatalf("pasted path should be hidden behind placeholder, input = %q", string(m.input))
}
}
func TestChatImagePlaceholdersUseSessionNumbers(t *testing.T) {
first := writeTestPNG(t)
second := writeTestPNG(t)
m := chatModel{
ctx: context.Background(),
opts: Options{
Model: "test",
Client: chatTestClient{},
MultiModal: true,
},
}
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(first), Paste: true})
m = updated.(chatModel)
if got := string(m.input); got != "[Image #0]" {
t.Fatalf("first placeholder = %q, want [Image #0]", got)
}
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = updated.(chatModel)
if cmd == nil {
t.Fatal("submit should start a run")
}
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(second), Paste: true})
m = updated.(chatModel)
if got := string(m.input); got != "[Image #1]" {
t.Fatalf("second placeholder = %q, want [Image #1]", got)
}
}
func TestChatAbsoluteImagePathBypassesSlashCommandParsing(t *testing.T) {
fp := writeTestPNG(t)
m := chatModel{
ctx: context.Background(),
opts: Options{
Model: "test",
Client: chatTestClient{},
MultiModal: true,
},
}
m.input = []rune(fp)
m.inputCursor = len(m.input)
m.inputCursorSet = true
updated, cmd := m.handleSubmit()
m = updated.(chatModel)
if cmd == nil {
t.Fatal("absolute image path should start a run instead of being parsed as a slash command")
}
if got := len(m.liveMessages[0].Images); got != 1 {
t.Fatalf("images = %d, want 1", got)
}
if got := m.entries[0].role; got != "user" {
t.Fatalf("entry role = %q, want user", got)
}
}
func TestChatDeletingImagePlaceholderRemovesAttachment(t *testing.T) {
fp := writeTestPNG(t)
m := chatModel{
ctx: context.Background(),
opts: Options{
Model: "test",
Client: chatTestClient{},
MultiModal: true,
},
}
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true})
m = updated.(chatModel)
if got := len(m.inputAttachments); got != 1 {
t.Fatalf("input attachments = %d, want 1", got)
}
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyBackspace})
m = updated.(chatModel)
if got := string(m.input); got != "describe " {
t.Fatalf("input after backspace = %q, want image placeholder removed", got)
}
if got := len(m.inputAttachments); got != 0 {
t.Fatalf("input attachments after editing placeholder = %d, want 0", got)
}
m.input = []rune("describe")
m.inputCursor = len(m.input)
m.inputCursorSet = true
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
m = updated.(chatModel)
if cmd == nil {
t.Fatal("submit should start a run")
}
if got := len(m.liveMessages[0].Images); got != 0 {
t.Fatalf("images = %d, want 0 after deleting placeholder", got)
}
if got := m.liveMessages[0].Content; got != "describe" {
t.Fatalf("content = %q, want describe", got)
}
}
func TestChatWordDeletingImagePlaceholderRemovesAttachment(t *testing.T) {
fp := writeTestPNG(t)
m := chatModel{
ctx: context.Background(),
opts: Options{
Model: "test",
Client: chatTestClient{},
MultiModal: true,
},
}
updated, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("describe " + fp), Paste: true})
m = updated.(chatModel)
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeySpace})
m = updated.(chatModel)
updated, _ = m.Update(tea.KeyMsg{Type: tea.KeyBackspace, Alt: true})
m = updated.(chatModel)
if got := string(m.input); got != "describe " {
t.Fatalf("input after word backspace = %q, want image placeholder removed", got)
}
if got := len(m.inputAttachments); got != 0 {
t.Fatalf("input attachments after word backspace = %d, want 0", got)
}
}
func writeTestPNG(t *testing.T) string {
t.Helper()
dir := t.TempDir()
fp := filepath.Join(dir, "dragged image.png")
data := make([]byte, 600)
copy(data, []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'})
if err := os.WriteFile(fp, data, 0o600); err != nil {
t.Fatalf("failed to write test image: %v", err)
}
return fp
}
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
+187
View File
@@ -0,0 +1,187 @@
package chat
import (
"context"
"database/sql"
"regexp"
"slices"
"testing"
"time"
tea "github.com/charmbracelet/bubbletea"
coreagent "github.com/ollama/ollama/agent"
agentstore "github.com/ollama/ollama/agent/store"
"github.com/ollama/ollama/api"
)
type chatTestTool struct{}
type chatTestClient struct{}
type chatCaptureClient struct {
requests []*api.ChatRequest
}
type chatShowTestClient struct {
chatTestClient
resp *api.ShowResponse
req *api.ShowRequest
}
type chatResumeTestStore struct {
chats []agentstore.ChatSummary
byID map[string]*agentstore.AgentChat
prompts []string
setModels map[string]string
setModelErr error
}
type chatTestCompactor struct {
result coreagent.CompactionResult
err error
progress []int
request coreagent.CompactionRequest
}
func (chatTestClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error {
if err := ctx.Err(); err != nil {
return err
}
return fn(api.ChatResponse{
Message: api.Message{Role: "assistant", Content: "ok"},
Done: true,
})
}
func (c *chatCaptureClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error {
if err := ctx.Err(); err != nil {
return err
}
c.requests = append(c.requests, req)
return fn(api.ChatResponse{
Message: api.Message{Role: "assistant", Content: "ok"},
Done: true,
})
}
func (c *chatShowTestClient) Show(ctx context.Context, req *api.ShowRequest) (*api.ShowResponse, error) {
if err := ctx.Err(); err != nil {
return nil, err
}
c.req = req
if c.resp != nil {
return c.resp, nil
}
return &api.ShowResponse{}, nil
}
func (s *chatResumeTestStore) EnsureChat(context.Context, string, string) error {
return nil
}
func (s *chatResumeTestStore) AppendAgentMessage(context.Context, string, api.Message, string) error {
return nil
}
func (s *chatResumeTestStore) UpdateLastAgentMessage(context.Context, string, api.Message, string) error {
return nil
}
func (s *chatResumeTestStore) SetChatModel(_ context.Context, chatID string, model string) error {
if s.setModelErr != nil {
return s.setModelErr
}
if s.setModels == nil {
s.setModels = map[string]string{}
}
s.setModels[chatID] = model
return nil
}
func (s *chatResumeTestStore) ListChats(context.Context, int) ([]agentstore.ChatSummary, error) {
return slices.Clone(s.chats), nil
}
func (s *chatResumeTestStore) AgentChat(_ context.Context, id string) (*agentstore.AgentChat, error) {
chat, ok := s.byID[id]
if !ok {
return nil, sql.ErrNoRows
}
out := *chat
out.Messages = slices.Clone(chat.Messages)
return &out, nil
}
func (s *chatResumeTestStore) ListUserMessages(context.Context, int) ([]string, error) {
return slices.Clone(s.prompts), nil
}
func (c *chatTestCompactor) MaybeCompact(_ context.Context, req coreagent.CompactionRequest) (coreagent.CompactionResult, error) {
c.request = req
for _, tokens := range c.progress {
if req.Progress != nil {
req.Progress(coreagent.CompactionProgress{Tokens: tokens})
}
}
return c.result, c.err
}
func nextChatMsg(t *testing.T, ch <-chan tea.Msg) tea.Msg {
t.Helper()
select {
case msg, ok := <-ch:
if !ok {
t.Fatal("message channel closed")
}
return msg
case <-time.After(time.Second):
t.Fatal("timed out waiting for chat message")
return nil
}
}
func (chatTestTool) Name() string {
return "fake_tool"
}
func (chatTestTool) Description() string {
return "does test work"
}
func (chatTestTool) Schema() api.ToolFunction {
return api.ToolFunction{
Name: "fake_tool",
Description: "does test work",
Parameters: api.ToolFunctionParameters{
Type: "object",
},
}
}
func (chatTestTool) Execute(context.Context, coreagent.ToolContext, map[string]any) (coreagent.ToolResult, error) {
return coreagent.ToolResult{Content: "ok"}, nil
}
func waitForRunDone(t *testing.T, events <-chan tea.Msg) chatRunDoneMsg {
t.Helper()
timeout := time.After(2 * time.Second)
for {
select {
case msg, ok := <-events:
if !ok {
t.Fatal("events closed before run done")
}
if done, ok := msg.(chatRunDoneMsg); ok {
return done
}
case <-timeout:
t.Fatal("timed out waiting for run done")
}
}
}
func stripANSI(s string) string {
re := regexp.MustCompile(`\x1b\[[0-9;:]*[A-Za-z]`)
return re.ReplaceAllString(s, "")
}
+110
View File
@@ -0,0 +1,110 @@
package chat
import "github.com/charmbracelet/lipgloss"
const (
chatAnsiRed = "1"
chatAnsiGreen = "2"
chatAnsiYellow = "3"
chatAnsiBlue = "4"
chatAnsiCyan = "6"
chatAnsiMuted = "8"
)
var (
chatHeaderStyle = lipgloss.NewStyle().
Bold(true)
chatMetaStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiMuted))
chatNotificationStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color("7"))
chatUserStyle = lipgloss.NewStyle()
chatAssistantStyle = lipgloss.NewStyle()
chatToolStyle = lipgloss.NewStyle()
chatToolRunningStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiYellow))
chatToolDoneStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiGreen))
// chatToolMixedStyle marks a tool group with both succeeded and failed
// calls (partial success). Amber/orange is distinct from green (success),
// red (failure), and yellow (running).
chatToolMixedStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color("208"))
chatDiffMetaStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiMuted))
chatDiffFileStyle = lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(chatAnsiCyan))
chatDiffHunkStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiBlue))
chatDiffAddStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiGreen))
chatDiffDeleteStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiRed))
chatErrorStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiRed))
chatFullAccessStyle = lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(chatAnsiRed))
chatCommandNameStyle = lipgloss.NewStyle()
chatResumeTextStyle = lipgloss.NewStyle()
chatResumeTitleStyle = lipgloss.NewStyle().
Bold(true)
chatResumeSelectedStyle = lipgloss.NewStyle().
Bold(true)
chatResumeMetaStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiMuted))
chatResumeBorderStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiMuted))
chatHistoryTitleStyle = lipgloss.NewStyle().
Bold(true)
chatHistorySystemRoleStyle = lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(chatAnsiMuted))
chatHistoryUserRoleStyle = lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(chatAnsiBlue))
chatHistoryAssistantRoleStyle = lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(chatAnsiYellow))
chatHistoryToolRoleStyle = lipgloss.NewStyle().
Bold(true).
Foreground(lipgloss.Color(chatAnsiGreen))
chatHistoryLabelStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiMuted))
chatHistoryTextStyle = lipgloss.NewStyle()
chatHistoryCodeStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color(chatAnsiCyan))
chatSelectionStyle = lipgloss.NewStyle().
Reverse(true)
)
+174
View File
@@ -0,0 +1,174 @@
package chat
import (
"fmt"
"strings"
tea "github.com/charmbracelet/bubbletea"
"github.com/ollama/ollama/api"
)
type chatThinkOption struct {
value string
label string
description string
}
type chatThinkPicker struct {
options []chatThinkOption
cursor int
}
var chatThinkOptions = []chatThinkOption{
{value: "auto", label: "auto", description: "use the model default"},
{value: "on", label: "on", description: "enable thinking"},
{value: "off", label: "off", description: "disable thinking"},
{value: "low", label: "low", description: "use low thinking effort"},
{value: "medium", label: "medium", description: "use medium thinking effort"},
{value: "high", label: "high", description: "use high thinking effort"},
{value: "max", label: "max", description: "use maximum thinking effort"},
}
func (m *chatModel) openThinkPicker() (tea.Model, tea.Cmd) {
m.thinkPicker = newChatThinkPicker(m.opts.Think)
m.status = "think"
return *m, nil
}
func newChatThinkPicker(current *api.ThinkValue) *chatThinkPicker {
picker := &chatThinkPicker{options: append([]chatThinkOption(nil), chatThinkOptions...)}
currentValue := thinkValueLabel(current)
for i, option := range picker.options {
if option.value == currentValue {
picker.cursor = i
break
}
}
return picker
}
func (m chatModel) updateThinkPicker(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
switch msg.Type {
case tea.KeyCtrlC, tea.KeyEsc:
m.thinkPicker = nil
m.status = "ready"
case tea.KeyEnter:
return m.selectThinkOption()
case tea.KeyUp:
m.thinkPicker.move(-1)
case tea.KeyDown:
m.thinkPicker.move(1)
}
return m, nil
}
func (p *chatThinkPicker) move(delta int) {
if p == nil || len(p.options) == 0 || delta == 0 {
return
}
p.cursor = clamp(p.cursor+delta, 0, len(p.options)-1)
}
func (p *chatThinkPicker) selected() (chatThinkOption, bool) {
if p == nil || len(p.options) == 0 {
return chatThinkOption{}, false
}
return p.options[clamp(p.cursor, 0, len(p.options)-1)], true
}
func (m chatModel) selectThinkOption() (tea.Model, tea.Cmd) {
option, ok := m.thinkPicker.selected()
if !ok {
return m, nil
}
m.thinkPicker = nil
return m.applyThinkValue(option.value)
}
func (m *chatModel) handleThinkCommand(value string) (tea.Model, tea.Cmd) {
return m.applyThinkValue(value)
}
func (m *chatModel) handleLegacySetThinkCommand(input string) (tea.Model, tea.Cmd) {
value := strings.TrimSpace(strings.TrimPrefix(input, "/set think"))
if value == "" {
value = "on"
}
return m.applyThinkValue(value)
}
func (m *chatModel) applyThinkValue(value string) (tea.Model, tea.Cmd) {
think, label, err := parseThinkValue(value)
if err != nil {
m.entries = append(m.entries, newChatEntry(chatEntry{role: "error", content: err.Error(), err: err.Error()}))
m.status = "error"
return *m, nil
}
m.opts.Think = think
m.status = "think " + label
return *m, nil
}
func parseThinkValue(value string) (*api.ThinkValue, string, error) {
switch strings.ToLower(strings.TrimSpace(value)) {
case "", "auto", "default", "unset":
return nil, "auto", nil
case "on", "true", "think", "thinking":
return &api.ThinkValue{Value: true}, "on", nil
case "off", "false", "nothink", "no-think":
return &api.ThinkValue{Value: false}, "off", nil
case "low", "medium", "high", "max":
value = strings.ToLower(strings.TrimSpace(value))
return &api.ThinkValue{Value: value}, value, nil
default:
return nil, "", fmt.Errorf("Usage: /think [auto|on|off|low|medium|high|max]")
}
}
func thinkValueLabel(value *api.ThinkValue) string {
if value == nil || value.Value == nil {
return "auto"
}
switch v := value.Value.(type) {
case bool:
if v {
return "on"
}
return "off"
case string:
return strings.ToLower(v)
default:
return "auto"
}
}
func (m chatModel) renderThinkPicker(width int) string {
picker := m.thinkPicker
if picker == nil {
return ""
}
var b strings.Builder
b.WriteString(chatResumeTitleStyle.Render("Thinking mode"))
b.WriteString("\n\n")
for i, option := range picker.options {
selected := i == picker.cursor
if selected {
b.WriteString(chatResumeSelectedStyle.Render(" " + option.label))
} else {
b.WriteString(" ")
b.WriteString(chatResumeTextStyle.Render(option.label))
}
b.WriteByte('\n')
b.WriteString(chatResumeMetaStyle.Render(" " + option.description))
b.WriteByte('\n')
if i < len(picker.options)-1 {
b.WriteByte('\n')
}
}
b.WriteString("\n")
b.WriteString(chatResumeMetaStyle.Render("↑/↓ move • enter select • esc cancel"))
return b.String()
}
+32 -5
View File
@@ -2,6 +2,7 @@ package tui
import (
"fmt"
"strings"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/lipgloss"
@@ -25,6 +26,7 @@ type confirmModel struct {
confirmed bool
cancelled bool
width int
plain bool
}
type ConfirmOptions = launch.ConfirmOptions
@@ -83,16 +85,40 @@ func (m confirmModel) View() string {
noBtn = confirmActiveStyle.Render(" " + noLabel + " ")
}
s := selectorTitleStyle.Render(m.prompt) + "\n\n"
prompt := renderConfirmPrompt(m.prompt, m.width, m.plain)
s := prompt + "\n\n"
s += " " + yesBtn + " " + noBtn + "\n\n"
s += selectorHelpStyle.Render("←/→ navigate • enter confirm • esc cancel")
s += renderConfirmHelp(m.width)
if m.width > 0 {
return lipgloss.NewStyle().MaxWidth(m.width).Render(s)
}
return s
}
func renderConfirmPrompt(prompt string, width int, plain bool) string {
lines := []string{prompt}
if width > 0 {
lines = wrapText(prompt, width)
}
if plain {
return strings.Join(lines, "\n")
}
for i, line := range lines {
lines[i] = selectorTitleStyle.Render(line)
}
return strings.Join(lines, "\n")
}
func renderConfirmHelp(width int) string {
help := "←/→ navigate • enter confirm • esc cancel"
lines := []string{help}
if width > 0 {
lines = wrapText(help, width)
}
for i, line := range lines {
lines[i] = selectorHelpStyle.Render(line)
}
return strings.Join(lines, "\n")
}
// RunConfirm shows a bubbletea yes/no confirmation prompt.
// Returns true if the user confirmed, false if cancelled.
func RunConfirm(prompt string) (bool, error) {
@@ -116,6 +142,7 @@ func RunConfirmWithOptions(prompt string, options ConfirmOptions) (bool, error)
yesLabel: yesLabel,
noLabel: noLabel,
yes: true, // default to yes
plain: options.PlainPrompt,
}
p := tea.NewProgram(m)
+28
View File
@@ -22,6 +22,34 @@ func TestConfirmModel_View_ContainsPrompt(t *testing.T) {
}
}
func TestConfirmModel_View_PlainPromptDoesNotStylePrompt(t *testing.T) {
m := confirmModel{prompt: "Sign in?", yes: true, plain: true}
got := m.View()
promptBlock := strings.SplitN(got, "\n\n", 2)[0]
if strings.Contains(promptBlock, "\x1b[") {
t.Fatalf("plain prompt should not include ANSI styling: %q", promptBlock)
}
}
func TestConfirmModel_View_WrapsPromptToWidth(t *testing.T) {
m := confirmModel{
prompt: "Sign in to use web search and cloud models with Ollama, Claude Code, OpenClaw, Hermes and more?\n\nYou can keep using local models without signing in.",
yes: true,
width: 52,
plain: true,
}
got := stripANSI(m.View())
if !strings.Contains(got, "Hermes and more?") {
t.Fatalf("wrapped prompt lost trailing text:\n%s", got)
}
for _, line := range strings.Split(got, "\n") {
if len(line) > m.width {
t.Fatalf("line width = %d, want <= %d: %q\n%s", len(line), m.width, line, got)
}
}
}
func TestConfirmModel_View_ContainsButtons(t *testing.T) {
m := confirmModel{prompt: "Download?", yes: true}
got := m.View()
+14
View File
@@ -0,0 +1,14 @@
package tui
// RunAgentSignInOnboarding asks whether the user wants to sign in before
// starting the root agent flow.
func RunAgentSignInOnboarding() (bool, error) {
return RunConfirmWithOptions(
"Sign in to use web search and cloud models with Ollama, Claude Code, OpenClaw, Hermes and more?\n\nYou can keep using local models without signing in.",
ConfirmOptions{
YesLabel: "Sign in",
NoLabel: "Not now",
PlainPrompt: true,
},
)
}
+1 -1
View File
@@ -194,7 +194,7 @@ func renderSignIn(modelName, signInURL string, spinner, width int) string {
var s strings.Builder
fmt.Fprintf(&s, "To use %s, please sign in.\n\n", selectorSelectedItemStyle.Render(modelName))
fmt.Fprintf(&s, "To use %s, please sign in.\n\n", modelName)
s.WriteString("Navigate to:\n")
s.WriteString(urlWrap.Render(urlColor.Render(signInURL)))
+8
View File
@@ -18,6 +18,14 @@ func TestRenderSignIn_ContainsModelName(t *testing.T) {
}
}
func TestRenderSignIn_FirstLineIsPlain(t *testing.T) {
got := renderSignIn("glm-4.7:cloud", "https://example.com/signin", 0, 80)
firstLine := strings.SplitN(got, "\n", 2)[0]
if strings.Contains(firstLine, "\x1b[") {
t.Fatalf("sign-in heading should not include ANSI styling: %q", firstLine)
}
}
func TestRenderSignIn_ContainsURL(t *testing.T) {
url := "https://ollama.com/connect?key=abc123"
got := renderSignIn("test:cloud", url, 0, 120)
+8
View File
@@ -0,0 +1,8 @@
package tui
import "regexp"
func stripANSI(s string) string {
re := regexp.MustCompile(`\x1b\[[0-9;:]*[A-Za-z]`)
return re.ReplaceAllString(s, "")
}
+41
View File
@@ -0,0 +1,41 @@
package tui
import (
"strings"
"github.com/mattn/go-runewidth"
)
func wrapText(text string, width int) []string {
if width <= 0 {
return strings.Split(text, "\n")
}
var lines []string
for _, line := range strings.Split(text, "\n") {
for runewidth.StringWidth(line) > width {
cut := textDisplayWidthCut(line, width)
lines = append(lines, line[:cut])
line = line[cut:]
}
lines = append(lines, line)
}
return lines
}
func textDisplayWidthCut(line string, width int) int {
if width <= 0 {
return 0
}
used := 0
for index, r := range line {
w := runewidth.RuneWidth(r)
if used+w > width {
if index == 0 {
return len(string(r))
}
return index
}
used += w
}
return len(line)
}
+2 -2
View File
@@ -48,8 +48,8 @@ type menuItem struct {
const pinnedIntegrationCount = 4
var runModelMenuItem = menuItem{
title: "Chat with a model",
description: "Start an interactive chat with a model",
title: "Chat and Code",
description: "Ollama's built-in agent to chat, code, and do work",
isRunModel: true,
}
+4 -1
View File
@@ -140,11 +140,14 @@ func TestMenuRendersPinnedItemsAndMore(t *testing.T) {
}
view := menu.View()
for _, want := range []string{"Chat with a model", "Launch Claude Code", "Launch Hermes Agent", "Launch OpenClaw", "More..."} {
for _, want := range []string{"Chat and Code", "Launch Claude Code", "Launch Hermes Agent", "Launch OpenClaw", "More..."} {
if !strings.Contains(view, want) {
t.Fatalf("expected menu view to contain %q\n%s", want, view)
}
}
if !strings.Contains(view, "Ollama's built-in agent to chat, code, and do work") {
t.Fatalf("expected menu view to contain agent description\n%s", view)
}
if findMenuCursorByIntegration(menu.items, "codex-app") != -1 && !strings.Contains(view, "Launch Codex App") {
t.Fatalf("expected menu view to contain Codex App\n%s", view)
}
+8 -7
View File
@@ -22,12 +22,12 @@ require (
require (
github.com/agnivade/levenshtein v1.1.1
github.com/charmbracelet/bubbletea v1.3.10
github.com/charmbracelet/lipgloss v1.1.0
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834
github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1
github.com/dlclark/regexp2 v1.11.5
github.com/emirpasic/gods/v2 v2.0.0-alpha
github.com/klauspost/compress v1.18.3
github.com/mattn/go-runewidth v0.0.16
github.com/mattn/go-runewidth v0.0.17
github.com/nlpodyssey/gopickle v0.3.0
github.com/pdevine/tensor v0.0.0-20240510204454-f88f4562727c
github.com/pelletier/go-toml/v2 v2.2.2
@@ -49,8 +49,8 @@ require (
github.com/buger/jsonparser v1.1.1 // indirect
github.com/bytedance/sonic/loader v0.1.1 // indirect
github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect
github.com/charmbracelet/x/ansi v0.10.1 // indirect
github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect
github.com/charmbracelet/x/ansi v0.10.2 // indirect
github.com/charmbracelet/x/cellbuf v0.0.13 // indirect
github.com/charmbracelet/x/term v0.2.1 // indirect
github.com/chewxy/hm v1.0.0 // indirect
github.com/chewxy/math32 v1.11.0 // indirect
@@ -60,8 +60,8 @@ require (
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/google/flatbuffers v24.3.25+incompatible // indirect
github.com/kr/text v0.2.0 // indirect
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
github.com/kr/pretty v0.3.1 // indirect
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
github.com/mailru/easyjson v0.7.7 // indirect
github.com/mattn/go-localereader v0.0.1 // indirect
github.com/mattn/go-pointer v0.0.1 // indirect
@@ -71,6 +71,7 @@ require (
github.com/pkg/errors v0.9.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/rivo/uniseg v0.4.7 // indirect
github.com/rogpeppe/go-internal v1.14.1 // indirect
github.com/tkrajina/go-reflector v0.5.5 // indirect
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
github.com/xtgo/set v1.0.0 // indirect
@@ -96,7 +97,7 @@ require (
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/spf13/pflag v1.0.5 // indirect
github.com/spf13/pflag v1.0.5
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
github.com/ugorji/go/codec v1.2.12 // indirect
golang.org/x/arch v0.8.0 // indirect
+16 -14
View File
@@ -30,12 +30,12 @@ github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlv
github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4=
github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs=
github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk=
github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY=
github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30=
github.com/charmbracelet/x/ansi v0.10.1 h1:rL3Koar5XvX0pHGfovN03f5cxLbCF2YvLeyz7D2jVDQ=
github.com/charmbracelet/x/ansi v0.10.1/go.mod h1:3RQDQ6lDnROptfpWuUVIUG64bD2g2BgntdxH0Ya5TeE=
github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd h1:vy0GVL4jeHEwG5YOXDmi86oYw2yuYUGqz6a8sLwg0X8=
github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs=
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 h1:ZR7e0ro+SZZiIZD7msJyA+NjkCNNavuiPBLgerbOziE=
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834/go.mod h1:aKC/t2arECF6rNOnaKaVU6y4t4ZeHQzqfxedE/VkVhA=
github.com/charmbracelet/x/ansi v0.10.2 h1:ith2ArZS0CJG30cIUfID1LXN7ZFXRCww6RUvAPA+Pzw=
github.com/charmbracelet/x/ansi v0.10.2/go.mod h1:HbLdJjQH4UH4AqA2HpRWuWNluRE6zxJH/yteYEYCFa8=
github.com/charmbracelet/x/cellbuf v0.0.13 h1:/KBBKHuVRbq1lYx5BzEHBAFBP8VcQzJejZ/IA3iR28k=
github.com/charmbracelet/x/cellbuf v0.0.13/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs=
github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ=
github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg=
github.com/chewxy/hm v1.0.0 h1:zy/TSv3LV2nD3dwUEQL2VhXeoXbb9QkpmdRAVUFiA6k=
@@ -156,16 +156,16 @@ github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa02
github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM=
github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M=
github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728 h1:QwWKgMY28TAXaDl+ExRDqGQltzXqN/xypdKP86niVn8=
github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728/go.mod h1:1fEHWurg7pvf5SG6XNE5Q8UZmOwex51Mkx3SLhrW5B4=
github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
@@ -175,8 +175,8 @@ github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+Ei
github.com/mattn/go-pointer v0.0.1 h1:n+XhsuGeVO6MEAp7xyEukFINEa+Quek5psIR/ylA6o0=
github.com/mattn/go-pointer v0.0.1/go.mod h1:2zXcozF6qYGgmsG+SeTZz3oAbFLdD3OWqnUbNvJZAlc=
github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
github.com/mattn/go-runewidth v0.0.17 h1:78v8ZlW0bP43XfmAfPsdXcoNCelfMHsDmd/pkENfrjQ=
github.com/mattn/go-runewidth v0.0.17/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
github.com/mattn/go-sqlite3 v1.14.24 h1:tpSp2G2KyMnnQu99ngJ47EIkWVmliIizyZBfPrBWDRM=
github.com/mattn/go-sqlite3 v1.14.24/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
@@ -204,6 +204,7 @@ github.com/pierrec/lz4/v4 v4.1.8 h1:ieHkV+i2BRzngO4Wd/3HGowuZStgq6QkPsD1eolNAO4=
github.com/pierrec/lz4/v4 v4.1.8/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
@@ -214,8 +215,9 @@ github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJ
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8=
github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE=
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
github.com/ruudk/golang-pdf417 v0.0.0-20181029194003-1af4ab5afa58/go.mod h1:6lfFZQK844Gfx8o5WFuvpxWRwnSoipWe/p622j1v06w=
github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
+3 -5
View File
@@ -17,10 +17,9 @@ import (
func TestLongInputContext(t *testing.T) {
// Setting NUM_PARALLEL to 1 ensures the allocated context is exactly what
// we asked for and there is nothing extra that we could spill over into.
// Context shift happens after a prompt has been admitted to a slot. Initial
// prompts that fill or exceed the slot are still rejected by llama-server.
// Accept a context-limit error here because older runners may truncate this
// prompt while llama-server reports it as too large to admit.
// Older runners silently truncate oversized prompts, while llama-server
// rejects them with a client error. Accept both behaviors so this test can
// run against main and the llama-server branch.
t.Setenv("OLLAMA_NUM_PARALLEL", "1")
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
@@ -75,7 +74,6 @@ func isContextLimitError(err string) bool {
return strings.Contains(err, "context") &&
(strings.Contains(err, "exceed") ||
strings.Contains(err, "too large") ||
strings.Contains(err, "longer") ||
strings.Contains(err, "too long"))
}
+4 -2
View File
@@ -48,8 +48,10 @@ func skipIfRemote(t *testing.T) {
// findHFCLI returns the path to the HuggingFace CLI, or "" if not found.
func findHFCLI() string {
if p, err := exec.LookPath("hf"); err == nil {
return p
for _, name := range []string{"huggingface-cli", "hf"} {
if p, err := exec.LookPath(name); err == nil {
return p
}
}
return ""
}
+3 -6
View File
@@ -23,14 +23,11 @@ intentionally skipped so a developer can iterate on a local llama.cpp tree.
- `llama-ollama-compat-util.h`, `llama-ollama-compat-util.cpp` - helpers for
KV edits, tensor renames, skip-prefix tracking, tensor load operations, and
small tensor repacking primitives.
- `001-llama-cpp-hooks.patch` - small additive call-site edits in llama.cpp files.
- `llama-cpp-hooks.patch` - small additive call-site edits in llama.cpp files.
It currently touches `src/llama-model-loader.cpp` and `tools/mtmd/clip.cpp`.
- `002-llama-cpp-ui-empty-assets.patch` - lets the llama.cpp UI embed helper
generate an empty asset table when no UI assets are present.
- `compat.cmake`, `apply-patch.cmake` - CMake glue and an idempotent applier
(used by `llama/server/CMakeLists.txt`) that applies every `*.patch` under
this directory by numeric filename order — the hooks patch plus each
`models/` architecture patch.
this directory — the hooks patch plus each `models/` architecture patch.
- `models/` - the sibling **new-architecture** layer: implementations of
architectures llama.cpp doesn't support yet, each added via a small
registration patch. (Those files *add* archs; the files above *translate*
@@ -119,7 +116,7 @@ cd /path/to/llama.cpp
git diff -- \
src/llama-model-loader.cpp \
tools/mtmd/clip.cpp \
> /path/to/ollama/llama/compat/001-llama-cpp-hooks.patch
> /path/to/ollama/llama/compat/llama-cpp-hooks.patch
```
## Implementation Notes
+8 -18
View File
@@ -3,11 +3,10 @@
# Invocation (from a CMake PATCH_COMMAND):
# cmake -DPATCH_DIR=<dir of *.patch> -P apply-patch.cmake
#
# Every *.patch under PATCH_DIR is applied in numeric filename order in the
# current working directory (which ExternalProject / FetchContent sets to the
# fetched source's SOURCE_DIR). A patch already applied detected via
# `git apply --reverse --check` is skipped. This makes re-configuring and
# re-building safe.
# Every *.patch under PATCH_DIR is applied in the current working directory
# (which ExternalProject / FetchContent sets to the fetched source's
# SOURCE_DIR). A patch already applied detected via `git apply --reverse
# --check` is skipped. This makes re-configuring and re-building safe.
if(NOT DEFINED PATCH_DIR)
message(FATAL_ERROR "apply-patch.cmake: PATCH_DIR not set")
@@ -20,17 +19,8 @@ get_filename_component(_git_ceiling "${_patch_workdir}" DIRECTORY)
set(_git_apply_env GIT_CEILING_DIRECTORIES=${_git_ceiling})
file(GLOB_RECURSE _patches "${PATCH_DIR}/*.patch")
set(_patch_entries)
list(SORT _patches)
foreach(PATCH_FILE IN LISTS _patches)
get_filename_component(_patch_name "${PATCH_FILE}" NAME)
list(APPEND _patch_entries "${_patch_name}|${PATCH_FILE}")
endforeach()
list(SORT _patch_entries)
foreach(_patch_entry IN LISTS _patch_entries)
string(REGEX REPLACE "^[^|]*\\|" "" PATCH_FILE "${_patch_entry}")
file(RELATIVE_PATH _patch_rel "${PATCH_DIR}" "${PATCH_FILE}")
# If the patch can be REVERSED cleanly, it's already applied. Skip.
execute_process(
COMMAND ${CMAKE_COMMAND} -E env ${_git_apply_env}
@@ -39,7 +29,7 @@ foreach(_patch_entry IN LISTS _patch_entries)
OUTPUT_QUIET ERROR_QUIET
)
if(_reverse_check EQUAL 0)
message(STATUS "llama/compat: ${_patch_rel} already applied, skipping")
message(STATUS "llama/compat: patch already applied, skipping")
continue()
endif()
@@ -51,10 +41,10 @@ foreach(_patch_entry IN LISTS _patch_entries)
)
if(NOT _apply_result EQUAL 0)
message(FATAL_ERROR
"llama/compat: failed to apply ${_patch_rel}\n"
"llama/compat: failed to apply ${PATCH_FILE}\n"
"This usually means the pinned llama.cpp source has changed. "
"Regenerate the patch against the pinned LLAMA_CPP_VERSION and retry.")
endif()
message(STATUS "llama/compat: applied ${_patch_rel}")
message(STATUS "llama/compat: applied patch")
endforeach()
+2 -2
View File
@@ -18,7 +18,7 @@
# The compat layer consists of:
# 1. Ollama-owned compat source files linked into the fetched llama.cpp
# targets from this directory.
# 2. A small ordered patch set that adds call-sites in llama.cpp loaders.
# 2. A small patch file that adds call-sites in llama.cpp loaders.
set(_compat_dir ${CMAKE_CURRENT_LIST_DIR})
@@ -46,7 +46,7 @@ set(OLLAMA_LLAMA_CPP_COMPAT_DIR
# Also export the individual paths in case callers want to do something
# custom (e.g. emit a dependency on the patch so reconfigures re-apply).
set(OLLAMA_LLAMA_CPP_COMPAT_PATCH_FILE
"${_compat_dir}/001-llama-cpp-hooks.patch"
"${_compat_dir}/llama-cpp-hooks.patch"
CACHE INTERNAL "Path to the llama.cpp compat patch")
set(OLLAMA_LLAMA_CPP_COMPAT_SOURCES
File renamed without changes.
+1 -1
View File
@@ -14,7 +14,7 @@ void llama_model_laguna::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
ml.get_key_or_arr("laguna.attention.layer_types", hparams.is_swa_impl, hparams.n_layer(), false);
ml.get_key_or_arr("laguna.attention.layer_types", hparams.is_swa_impl, hparams.n_layer, false);
ml.get_key("laguna.rope.swa.dimension_count", hparams.n_rot_swa, false);
ml.get_key("laguna.rope.swa.freq_base", hparams.rope_freq_base_train_swa, false);
@@ -16,8 +16,8 @@ diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
diff --git a/src/llama-arch.h b/src/llama-arch.h
--- a/src/llama-arch.h
+++ b/src/llama-arch.h
@@ -145,2 +145,3 @@ enum llm_arch {
LLM_ARCH_EAGLE3,
@@ -140,2 +140,3 @@ enum llm_arch {
LLM_ARCH_MELLUM,
+ LLM_ARCH_LAGUNA,
LLM_ARCH_UNKNOWN,
@@ -382,2 +383,3 @@ enum llm_tensor {
+5 -8
View File
@@ -273,18 +273,15 @@ func (s *llamaServerRunner) completionPromptForRequest(ctx context.Context, req
return nil, err
}
// llama-server rejects prompts that fill the entire slot context, while the
// old runner could accept exactly num_ctx prompt tokens. Keep one token of
// headroom so token-level truncation preserves old behavior as closely as
// llama-server allows.
limit := s.options.NumCtx - 1
if len(tokens) <= limit {
return prompt, nil
}
if !s.launch.config.ContextShift {
return nil, api.StatusError{
StatusCode: http.StatusBadRequest,
ErrorMessage: "the prompt is longer than the context length currently available to the model; shorten the prompt, adjust the context length in settings, or use a model with a longer context length",
}
}
nKeep := req.Options.NumKeep
if nKeep < 0 {
nKeep = len(tokens)
@@ -296,7 +293,7 @@ func (s *llamaServerRunner) completionPromptForRequest(ctx context.Context, req
truncated = append(truncated, tokens[:nKeep]...)
truncated = append(truncated, tokens[nKeep+discard:]...)
slog.Warn("truncating input prompt", "limit", limit, "prompt", len(tokens), "keep", nKeep, "new", len(truncated))
slog.Warn("truncating input prompt", "limit", s.options.NumCtx, "prompt", len(tokens), "keep", nKeep, "new", len(truncated))
return truncated, nil
}
+10 -158
View File
@@ -488,16 +488,14 @@ func TestLlamaServerCompletionForwardsRepeatLastNZero(t *testing.T) {
}
}
func TestLlamaServerCompletionRejectsPromptOverContext(t *testing.T) {
const wantError = "the prompt is longer than the context length currently available to the model; shorten the prompt, adjust the context length in settings, or use a model with a longer context length"
func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
var completionReq llamaServerCompletionRequest
var tokenizeReq struct {
Content string `json:"content"`
AddSpecial bool `json:"add_special"`
ParseSpecial *bool `json:"parse_special"`
}
completionCalled := false
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/health":
@@ -509,7 +507,10 @@ func TestLlamaServerCompletionRejectsPromptOverContext(t *testing.T) {
}
fmt.Fprint(w, `{"tokens":[0,1,2,3,4,5,6,7,8,9]}`)
case "/completion":
completionCalled = true
if err := json.NewDecoder(r.Body).Decode(&completionReq); err != nil {
t.Errorf("invalid completion request body: %v", err)
return
}
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprintln(w, `data: {"content":"ok","stop":true,"timings":{"prompt_n":7,"prompt_ms":1,"predicted_n":1,"predicted_ms":1}}`)
default:
@@ -536,15 +537,8 @@ func TestLlamaServerCompletionRejectsPromptOverContext(t *testing.T) {
Options: &opts,
Truncate: true,
}, func(cr CompletionResponse) {})
var statusErr api.StatusError
if !errors.As(err, &statusErr) {
t.Fatalf("Completion error = %T %v, want api.StatusError", err, err)
}
if statusErr.StatusCode != http.StatusBadRequest {
t.Fatalf("StatusCode = %d, want %d", statusErr.StatusCode, http.StatusBadRequest)
}
if statusErr.ErrorMessage != wantError {
t.Fatalf("ErrorMessage = %q, want %q", statusErr.ErrorMessage, wantError)
if err != nil {
t.Fatalf("Completion error: %v", err)
}
if tokenizeReq.Content != strings.Repeat("long prompt ", 2) {
@@ -553,149 +547,10 @@ func TestLlamaServerCompletionRejectsPromptOverContext(t *testing.T) {
if !tokenizeReq.AddSpecial {
t.Fatal("expected tokenize request to add special tokens")
}
if completionCalled {
t.Fatal("completion endpoint was called")
}
}
func TestLlamaServerCompletionContextShiftAllowsPromptWithHeadroom(t *testing.T) {
var capturedReq llamaServerCompletionRequest
var tokenizeReq struct {
Content string `json:"content"`
AddSpecial bool `json:"add_special"`
ParseSpecial *bool `json:"parse_special"`
}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/health":
fmt.Fprint(w, `{"status":"ok"}`)
case "/tokenize":
if err := json.NewDecoder(r.Body).Decode(&tokenizeReq); err != nil {
t.Errorf("invalid tokenize request body: %v", err)
return
}
fmt.Fprint(w, `{"tokens":[0,1,2,3,4,5,6]}`)
case "/completion":
if err := json.NewDecoder(r.Body).Decode(&capturedReq); err != nil {
t.Errorf("invalid completion request body: %v", err)
return
}
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprintln(w, `data: {"content":"ok","stop":true,"timings":{"prompt_n":10,"prompt_ms":1,"predicted_n":1,"predicted_ms":1}}`)
default:
t.Errorf("unexpected path: %s", r.URL.Path)
}
}))
defer srv.Close()
parts := strings.Split(srv.URL, ":")
var portInt int
fmt.Sscanf(parts[len(parts)-1], "%d", &portInt)
runner := &llamaServerRunner{
port: portInt,
cmd: fakeRunningCmd(),
sem: semaphore.NewWeighted(1),
options: api.Options{Runner: api.Runner{NumCtx: 8}},
launch: llamaServerLaunchConfig{
config: LlamaServerConfig{ContextShift: true},
},
}
opts := api.DefaultOptions()
opts.NumKeep = 3
prompt := strings.Repeat("long prompt ", 2)
err := runner.Completion(t.Context(), CompletionRequest{
Prompt: prompt,
Options: &opts,
Truncate: true,
}, func(cr CompletionResponse) {})
if err != nil {
t.Fatalf("Completion error: %v", err)
}
if tokenizeReq.Content != prompt {
t.Fatalf("tokenize content = %q, want %q", tokenizeReq.Content, prompt)
}
if !tokenizeReq.AddSpecial {
t.Fatal("expected tokenize request to add special tokens")
}
if capturedReq.Prompt != prompt {
t.Fatalf("prompt = %q, want %q", capturedReq.Prompt, prompt)
}
if capturedReq.NKeep != opts.NumKeep {
t.Fatalf("n_keep = %d, want %d", capturedReq.NKeep, opts.NumKeep)
}
}
func TestLlamaServerCompletionContextShiftTruncatesPromptOverContext(t *testing.T) {
var capturedReq llamaServerCompletionRequest
var tokenizeReq struct {
Content string `json:"content"`
AddSpecial bool `json:"add_special"`
ParseSpecial *bool `json:"parse_special"`
}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/health":
fmt.Fprint(w, `{"status":"ok"}`)
case "/tokenize":
if err := json.NewDecoder(r.Body).Decode(&tokenizeReq); err != nil {
t.Errorf("invalid tokenize request body: %v", err)
return
}
fmt.Fprint(w, `{"tokens":[0,1,2,3,4,5,6,7,8,9]}`)
case "/completion":
if err := json.NewDecoder(r.Body).Decode(&capturedReq); err != nil {
t.Errorf("invalid completion request body: %v", err)
return
}
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprintln(w, `data: {"content":"ok","stop":true,"timings":{"prompt_n":7,"prompt_ms":1,"predicted_n":1,"predicted_ms":1}}`)
default:
t.Errorf("unexpected path: %s", r.URL.Path)
}
}))
defer srv.Close()
parts := strings.Split(srv.URL, ":")
var portInt int
fmt.Sscanf(parts[len(parts)-1], "%d", &portInt)
runner := &llamaServerRunner{
port: portInt,
cmd: fakeRunningCmd(),
sem: semaphore.NewWeighted(1),
options: api.Options{Runner: api.Runner{NumCtx: 8}},
launch: llamaServerLaunchConfig{
config: LlamaServerConfig{ContextShift: true},
},
}
opts := api.DefaultOptions()
opts.NumKeep = 3
prompt := strings.Repeat("long prompt ", 2)
err := runner.Completion(t.Context(), CompletionRequest{
Prompt: prompt,
Options: &opts,
Truncate: true,
}, func(cr CompletionResponse) {})
if err != nil {
t.Fatalf("Completion error: %v", err)
}
if tokenizeReq.Content != prompt {
t.Fatalf("tokenize content = %q, want %q", tokenizeReq.Content, prompt)
}
if !tokenizeReq.AddSpecial {
t.Fatal("expected tokenize request to add special tokens")
}
got, ok := capturedReq.Prompt.([]any)
got, ok := completionReq.Prompt.([]any)
if !ok {
t.Fatalf("completion prompt = %T, want token array", capturedReq.Prompt)
t.Fatalf("completion prompt = %T, want token array", completionReq.Prompt)
}
want := []int{0, 1, 2, 6, 7, 8, 9}
if len(got) != len(want) {
@@ -707,9 +562,6 @@ func TestLlamaServerCompletionContextShiftTruncatesPromptOverContext(t *testing.
t.Fatalf("token prompt[%d] = %#v, want %d", i, got[i], wantToken)
}
}
if capturedReq.NKeep != opts.NumKeep {
t.Fatalf("n_keep = %d, want %d", capturedReq.NKeep, opts.NumKeep)
}
}
func TestLlamaServerCompletionWithMediaUsesRunnerMarker(t *testing.T) {
-315
View File
@@ -1,315 +0,0 @@
package parsers
import (
"encoding/json"
"log/slog"
"strings"
"unicode"
"github.com/ollama/ollama/api"
)
// CohereParser parses output from Cohere North / Command A 2026 models
// (e.g. North-Mini-Code-1.0). The generation prompt ends with
// <|START_THINKING|> (reasoning on) or <|START_THINKING|><|END_THINKING|>
// (reasoning off), so output begins inside the thinking block when reasoning
// is enabled. After thinking, the model emits either
// <|START_TEXT|>content<|END_TEXT|> or an <|START_ACTION|>[...]<|END_ACTION|>
// tool call array, then <|END_OF_TURN_TOKEN|>.
type CohereParser struct {
state cohereParserState
buffer strings.Builder
callIndex int
}
type cohereParserState int
const (
cohereCollectingThinking cohereParserState = iota
cohereAwaitingBlock
cohereCollectingContent
cohereCollectingAction
)
const (
cohereEndThinking = "<|END_THINKING|>"
cohereStartText = "<|START_TEXT|>"
cohereEndText = "<|END_TEXT|>"
cohereStartAction = "<|START_ACTION|>"
cohereEndAction = "<|END_ACTION|>"
cohereEndOfTurn = "<|END_OF_TURN_TOKEN|>"
// Legacy response markers from the older Command A chat template, which
// also ships in these models' tokenizer_config and shows up in sampled
// output; treated as aliases for START_TEXT / END_TEXT.
cohereStartResponse = "<|START_RESPONSE|>"
cohereEndResponse = "<|END_RESPONSE|>"
)
func (p *CohereParser) HasToolSupport() bool {
return true
}
func (p *CohereParser) HasThinkingSupport() bool {
return true
}
func (p *CohereParser) PreservedTokens() []string {
return []string{
"<|START_THINKING|>", cohereEndThinking,
cohereStartText, cohereEndText,
cohereStartAction, cohereEndAction,
cohereStartResponse, cohereEndResponse,
}
}
func (p *CohereParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
p.buffer.Reset()
p.callIndex = 0
// The template enables reasoning by default; nil means default.
thinkingEnabled := thinkValue == nil || thinkValue.Bool()
assistantPrefill := lastMessage != nil && lastMessage.Role == "assistant" && lastMessage.Content != ""
switch {
case assistantPrefill:
// The prompt left an open <|START_TEXT|> for continuation.
p.state = cohereCollectingContent
case thinkingEnabled:
p.state = cohereCollectingThinking
default:
p.state = cohereAwaitingBlock
}
return tools
}
func (p *CohereParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
p.buffer.WriteString(s)
var contentSb, thinkingSb strings.Builder
for {
c, t, tc, more := p.eat(done)
contentSb.WriteString(c)
thinkingSb.WriteString(t)
calls = append(calls, tc...)
if !more {
break
}
}
for i := range calls {
calls[i].Function.Index = p.callIndex
p.callIndex++
}
return contentSb.String(), thinkingSb.String(), calls, nil
}
// eat consumes what it can from the buffer for the current state. It returns
// more=true when a state transition happened and the remaining buffer should
// be reprocessed.
func (p *CohereParser) eat(done bool) (content string, thinking string, calls []api.ToolCall, more bool) {
buf := p.buffer.String()
if buf == "" {
return "", "", nil, false
}
switch p.state {
case cohereCollectingThinking:
if idx := strings.Index(buf, cohereEndThinking); idx != -1 {
thinking = strings.TrimRightFunc(buf[:idx], unicode.IsSpace)
p.resetBuffer(buf[idx+len(cohereEndThinking):])
p.state = cohereAwaitingBlock
return "", thinking, nil, true
}
// Emit all but a possible partial tag / trailing whitespace.
keep := overlap(buf, cohereEndThinking)
emitEnd := len(buf) - keep
emitEnd -= trailingWhitespaceLen(buf[:emitEnd])
if emitEnd > 0 {
thinking = buf[:emitEnd]
p.resetBuffer(buf[emitEnd:])
}
return "", thinking, nil, false
case cohereAwaitingBlock:
// Between blocks: look for the next text or action opener, skipping
// whitespace and the end-of-turn token.
for _, open := range []string{cohereStartText, cohereStartResponse} {
if idx := strings.Index(buf, open); idx != -1 {
p.resetBuffer(buf[idx+len(open):])
p.state = cohereCollectingContent
return "", "", nil, true
}
}
if idx := strings.Index(buf, cohereStartAction); idx != -1 {
p.resetBuffer(buf[idx+len(cohereStartAction):])
p.state = cohereCollectingAction
return "", "", nil, true
}
trimmed := strings.TrimLeftFunc(buf, unicode.IsSpace)
if strings.HasPrefix(trimmed, cohereEndOfTurn) {
p.resetBuffer(trimmed[len(cohereEndOfTurn):])
return "", "", nil, true
}
if trimmed == "" {
if done {
p.buffer.Reset()
}
return "", "", nil, false
}
// Wait only while the buffer could still grow into one of the
// expected tags. Anything else — bare content or an unrecognized
// tag — streams out as content rather than buffering forever.
if !done && maybePartialTag(trimmed) {
return "", "", nil, false
}
p.buffer.Reset()
p.state = cohereCollectingContent
p.buffer.WriteString(trimmed)
return "", "", nil, true
case cohereCollectingContent:
// END_OF_TURN also closes content for models that skip END_TEXT.
for _, close := range []string{cohereEndText, cohereEndResponse, cohereEndOfTurn} {
if idx := strings.Index(buf, close); idx != -1 {
content = buf[:idx]
p.resetBuffer(buf[idx+len(close):])
p.state = cohereAwaitingBlock
return content, "", nil, true
}
}
keep := max(overlap(buf, cohereEndText), overlap(buf, cohereEndResponse), overlap(buf, cohereEndOfTurn))
emitEnd := len(buf) - keep
if emitEnd > 0 {
content = buf[:emitEnd]
p.resetBuffer(buf[emitEnd:])
}
if done && p.buffer.Len() > 0 {
content += p.buffer.String()
p.buffer.Reset()
}
return content, "", nil, false
case cohereCollectingAction:
if idx := strings.Index(buf, cohereEndAction); idx != -1 {
payload := buf[:idx]
p.resetBuffer(buf[idx+len(cohereEndAction):])
p.state = cohereAwaitingBlock
calls = parseCohereActions(payload)
return "", "", calls, true
}
if done {
// Best effort on truncated output.
calls = parseCohereActions(buf)
p.buffer.Reset()
return "", "", calls, false
}
return "", "", nil, false
}
return "", "", nil, false
}
func (p *CohereParser) resetBuffer(s string) {
p.buffer.Reset()
p.buffer.WriteString(s)
}
// maybePartialTag reports whether s is a proper prefix of one of the tags the
// awaiting-block state recognizes — the only case worth waiting on for more
// output before treating the buffer as content.
func maybePartialTag(s string) bool {
for _, tag := range []string{cohereStartText, cohereStartResponse, cohereStartAction, cohereEndOfTurn} {
if len(s) < len(tag) && strings.HasPrefix(tag, s) {
return true
}
}
return false
}
type cohereToolCall struct {
ToolCallID string `json:"tool_call_id"`
ToolName string `json:"tool_name"`
Parameters api.ToolCallFunctionArguments `json:"parameters"`
}
func (c cohereToolCall) toolCall() api.ToolCall {
return api.ToolCall{
ID: c.ToolCallID,
Function: api.ToolCallFunction{
Name: c.ToolName,
Arguments: c.Parameters,
},
}
}
// parseCohereActions parses the JSON array inside an action block:
// [{"tool_call_id": "0", "tool_name": ..., "parameters": {...}}, ...]
//
// Sampled output occasionally malforms the JSON (a missing comma between
// calls, an unquoted value). When the array as a whole fails to parse, fall
// back to scanning its balanced top-level objects and parse each
// independently, so one bad call doesn't drop its siblings.
func parseCohereActions(payload string) []api.ToolCall {
payload = strings.TrimSpace(payload)
if payload == "" {
return nil
}
var parsed []cohereToolCall
if err := json.Unmarshal([]byte(payload), &parsed); err == nil {
calls := make([]api.ToolCall, 0, len(parsed))
for _, c := range parsed {
calls = append(calls, c.toolCall())
}
return calls
}
var calls []api.ToolCall
for _, obj := range scanJSONObjects(payload) {
var c cohereToolCall
if err := json.Unmarshal([]byte(obj), &c); err != nil {
if len(obj) > 200 {
obj = obj[:200] + "…"
}
slog.Warn("cohere action parsing failed", "error", err, "action", obj)
continue
}
calls = append(calls, c.toolCall())
}
return calls
}
// scanJSONObjects returns the balanced top-level {...} chunks of s, tracking
// strings and escapes so braces inside values don't split objects.
func scanJSONObjects(s string) []string {
var objects []string
depth, start := 0, -1
inString, escaped := false, false
for i := range len(s) {
switch c := s[i]; {
case escaped:
escaped = false
case c == '\\' && inString:
escaped = true
case c == '"':
inString = !inString
case inString:
case c == '{':
if depth == 0 {
start = i
}
depth++
case c == '}':
if depth > 0 {
depth--
if depth == 0 && start >= 0 {
objects = append(objects, s[start:i+1])
start = -1
}
}
}
}
return objects
}
-268
View File
@@ -1,268 +0,0 @@
package parsers
import (
"testing"
"github.com/ollama/ollama/api"
)
// Model output begins inside <|START_THINKING|> when reasoning is on (the
// generation prompt ends with that tag).
func cohereAddAll(t *testing.T, p *CohereParser, chunks []string) (content, thinking string, calls []api.ToolCall) {
t.Helper()
for i, c := range chunks {
done := i == len(chunks)-1
ct, th, tc, err := p.Add(c, done)
if err != nil {
t.Fatal(err)
}
content += ct
thinking += th
calls = append(calls, tc...)
}
return content, thinking, calls
}
func TestCohereParseThinkingThenText(t *testing.T) {
p := &CohereParser{}
p.Init(nil, nil, nil)
content, thinking, calls := cohereAddAll(t, p, []string{
"Let me think", " about this.<|END_THINKING|>",
"<|START_TEXT|>Hello", " world!<|END_TEXT|>",
})
if thinking != "Let me think about this." {
t.Errorf("thinking = %q", thinking)
}
if content != "Hello world!" {
t.Errorf("content = %q", content)
}
if len(calls) != 0 {
t.Errorf("unexpected calls: %v", calls)
}
}
func TestCohereParseSplitTags(t *testing.T) {
p := &CohereParser{}
p.Init(nil, nil, nil)
// Tags split across chunk boundaries must not leak into output.
content, thinking, _ := cohereAddAll(t, p, []string{
"think<|END_TH", "INKING|><|STAR", "T_TEXT|>ans", "wer<|END_", "TEXT|>",
})
if thinking != "think" {
t.Errorf("thinking = %q", thinking)
}
if content != "answer" {
t.Errorf("content = %q", content)
}
}
func TestCohereParseToolCall(t *testing.T) {
p := &CohereParser{}
p.Init(nil, nil, nil)
content, thinking, calls := cohereAddAll(t, p, []string{
"plan<|END_THINKING|><|START_ACTION|>[\n",
` {"tool_call_id": "0", "tool_name": "get_weather", "parameters": {"city": "Paris"}},`,
` {"tool_call_id": "1", "tool_name": "get_time", "parameters": {}}`,
"\n]<|END_ACTION|>",
})
if thinking != "plan" {
t.Errorf("thinking = %q", thinking)
}
if content != "" {
t.Errorf("content = %q", content)
}
if len(calls) != 2 {
t.Fatalf("calls = %d, want 2", len(calls))
}
if calls[0].Function.Name != "get_weather" || calls[1].Function.Name != "get_time" {
t.Errorf("call names = %q, %q", calls[0].Function.Name, calls[1].Function.Name)
}
if v, ok := calls[0].Function.Arguments.Get("city"); !ok || v != "Paris" {
t.Errorf("call 0 city = %v %v", v, ok)
}
if calls[0].Function.Index != 0 || calls[1].Function.Index != 1 {
t.Errorf("call indices = %d, %d", calls[0].Function.Index, calls[1].Function.Index)
}
}
func TestCohereParseReasoningOff(t *testing.T) {
p := &CohereParser{}
think := &api.ThinkValue{Value: false}
p.Init(nil, nil, think)
content, thinking, _ := cohereAddAll(t, p, []string{
"<|START_TEXT|>direct answer<|END_TEXT|>",
})
if thinking != "" {
t.Errorf("thinking = %q", thinking)
}
if content != "direct answer" {
t.Errorf("content = %q", content)
}
}
func TestCohereParseBareContent(t *testing.T) {
// Models occasionally skip the START_TEXT wrapper; treat raw text after
// thinking as content.
p := &CohereParser{}
p.Init(nil, nil, nil)
content, thinking, _ := cohereAddAll(t, p, []string{
"thought<|END_THINKING|>", "Just plain text", " output",
})
if thinking != "thought" {
t.Errorf("thinking = %q", thinking)
}
if content != "Just plain text output" {
t.Errorf("content = %q", content)
}
}
func TestCohereParseEndOfTurnWithoutEndText(t *testing.T) {
p := &CohereParser{}
p.Init(nil, nil, nil)
content, _, _ := cohereAddAll(t, p, []string{
"t<|END_THINKING|><|START_TEXT|>answer<|END_OF_TURN_TOKEN|>",
})
if content != "answer" {
t.Errorf("content = %q", content)
}
}
func TestCohereParsePrefillContinuation(t *testing.T) {
p := &CohereParser{}
last := &api.Message{Role: "assistant", Content: "partial"}
p.Init(nil, last, nil)
content, thinking, _ := cohereAddAll(t, p, []string{" continued<|END_TEXT|>"})
if thinking != "" {
t.Errorf("thinking = %q", thinking)
}
if content != " continued" {
t.Errorf("content = %q", content)
}
}
func TestCohereParserRegistered(t *testing.T) {
p := ParserForName("cohere")
if p == nil {
t.Fatal("cohere parser not registered")
}
if !p.HasToolSupport() || !p.HasThinkingSupport() {
t.Error("cohere parser should support tools and thinking")
}
}
func TestCohereParseMalformedActions(t *testing.T) {
// One malformed call (unquoted value) must not drop its well-formed
// sibling, and a missing comma between calls must not drop either.
p := &CohereParser{}
p.Init(nil, nil, nil)
_, _, calls := cohereAddAll(t, p, []string{
"plan<|END_THINKING|><|START_ACTION|>[\n",
` {"tool_call_id": "0", "tool_name": "set_alarm", "parameters": {"time": 15:30}},`,
` {"tool_call_id": "1", "tool_name": "get_weather", "parameters": {"city": "Oslo"}}`,
"\n]<|END_ACTION|>",
})
if len(calls) != 1 || calls[0].Function.Name != "get_weather" {
t.Fatalf("calls = %v, want the well-formed get_weather call", calls)
}
p = &CohereParser{}
p.Init(nil, nil, nil)
_, _, calls = cohereAddAll(t, p, []string{
`<|END_THINKING|><|START_ACTION|>[`,
`{"tool_call_id": "0", "tool_name": "a", "parameters": {}}`,
`{"tool_call_id": "1", "tool_name": "b", "parameters": {"x": "{not json}"}}`,
`]<|END_ACTION|>`,
})
if len(calls) != 2 || calls[0].Function.Name != "a" || calls[1].Function.Name != "b" {
t.Fatalf("calls = %v, want both calls despite missing comma", calls)
}
if v, ok := calls[1].Function.Arguments.Get("x"); !ok || v != "{not json}" {
t.Fatalf("braces inside string values must not split objects, got %v", v)
}
// Unparseable garbage yields no calls and no panic.
p = &CohereParser{}
p.Init(nil, nil, nil)
_, _, calls = cohereAddAll(t, p, []string{"x<|END_THINKING|><|START_ACTION|>[!!!]<|END_ACTION|>"})
if len(calls) != 0 {
t.Fatalf("calls = %v, want none", calls)
}
}
func TestCohereParseLegacyResponseMarkers(t *testing.T) {
// Models trained on the older Command A template sometimes emit
// <|START_RESPONSE|>/<|END_RESPONSE|> instead of START_TEXT/END_TEXT.
p := &CohereParser{}
p.Init(nil, nil, nil)
content, thinking, _ := cohereAddAll(t, p, []string{
"plan<|END_THINKING|>", "<|START_RESPONSE|>Hello", " there<|END_RESPONSE|>",
})
if thinking != "plan" {
t.Errorf("thinking = %q", thinking)
}
if content != "Hello there" {
t.Errorf("content = %q", content)
}
}
func TestCohereParseStreamsBeforeDone(t *testing.T) {
// Regression test: output after END_THINKING that opens with an
// unrecognized tag must stream as it arrives, not buffer until the
// generation finishes (it previously buffered forever, presenting as a
// hung response).
for _, tc := range []struct {
name string
chunk string
}{
{"legacy response marker", "<|START_RESPONSE|>The answer is 42."},
{"unrecognized tag", "<|TOOL_PLAN|>I should call the tool."},
{"bare content", "Just plain text."},
} {
t.Run(tc.name, func(t *testing.T) {
p := &CohereParser{}
p.Init(nil, nil, nil)
if _, _, _, err := p.Add("x<|END_THINKING|>", false); err != nil {
t.Fatal(err)
}
content, _, _, err := p.Add(tc.chunk, false) // done=false: still streaming
if err != nil {
t.Fatal(err)
}
if content == "" {
t.Fatalf("content did not stream before done for %q", tc.chunk)
}
})
}
}
func TestCohereParseEndOfTurnBetweenBlocks(t *testing.T) {
// A literal end-of-turn marker between blocks is consumed, not shown.
p := &CohereParser{}
p.Init(nil, nil, nil)
content, thinking, _ := cohereAddAll(t, p, []string{
"t<|END_THINKING|><|START_TEXT|>hi<|END_TEXT|><|END_OF_TURN_TOKEN|>",
})
if thinking != "t" || content != "hi" {
t.Errorf("thinking = %q, content = %q", thinking, content)
}
}
func TestCohereParseBareContentBeforeEndOfTurn(t *testing.T) {
// Bare content followed by an end-of-turn marker keeps the content.
p := &CohereParser{}
p.Init(nil, nil, nil)
content, _, _ := cohereAddAll(t, p, []string{
"t<|END_THINKING|>plain answer<|END_OF_TURN_TOKEN|>",
})
if content != "plain answer" {
t.Errorf("content = %q, want %q", content, "plain answer")
}
}
Loaded 100 of 130 files, more files were not shown because too many files have changed in this diff. Show more