diff --git a/Taskfile.yml b/Taskfile.yml index 925b4e4..d9b162c 100644 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -72,11 +72,6 @@ tasks: cmds: - go test -cover ./... - lint: - desc: Lint the Go code using golangci-lint - cmds: - - golangci-lint run ./... - test:coverage: desc: Run all Go tests cmds: @@ -100,7 +95,7 @@ tasks: PATH: '{{.GOPATH}}/bin:{{.PATH}}' cmds: - task: install:lint - - golangci-lint run + - golangci-lint run ./... lint:fix: desc: Run golangci-lint and fix issues @@ -108,7 +103,7 @@ tasks: PATH: '{{.GOPATH}}/bin:{{.PATH}}' cmds: - task: install:lint - - golangci-lint run --fix + - golangci-lint run ./... --fix install:lint: desc: Install golangci-lint if missing diff --git a/cmd/root.go b/cmd/root.go index 20c2ab0..b852087 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -165,8 +165,7 @@ func runRoot(_ *cobra.Command, args []string) error { return nil }, } - cfg.Explain = explainFlag - return interactiveRun(Version, cmds, cfg, client, shellInfo) + return interactiveRun(Version, cmds, cfg, client, shellInfo, explainFlag) } // Single-shot mode @@ -181,7 +180,9 @@ func runRoot(_ *cobra.Command, args []string) error { if err != nil { return fmt.Errorf("failed to get working directory: %w", err) } - systemPrompt := llm.BuildSystemPrompt(shellInfo.OS, shellInfo.Shell, shellInfo.Version, cwd, explainFlag) + + toolsEnabled := cfg.Safety.ToolCalling != config.ToolCallingNever + systemPrompt := llm.BuildSystemPrompt(shellInfo.OS, shellInfo.Shell, shellInfo.Version, cwd, explainFlag, toolsEnabled) // Inject matching memories into the system prompt entries, err := memory.Load() diff --git a/cmd/root_test.go b/cmd/root_test.go index 320b7e7..4fb6c50 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -101,7 +101,7 @@ func TestRunRoot_InteractiveMode_ConfigRunShow(t *testing.T) { t.Cleanup(func() { interactiveRun = origRun }) var capturedCmds interactive.BuiltinCommands - interactiveRun = func(_ string, cmds interactive.BuiltinCommands, _ *config.Config, _ llm.Client, _ shell.Info) error { + interactiveRun = func(_ string, cmds interactive.BuiltinCommands, _ *config.Config, _ llm.Client, _ shell.Info, _ bool) error { capturedCmds = cmds return nil } @@ -160,7 +160,7 @@ func TestRunRoot_InteractiveMode_MemoryRun(t *testing.T) { t.Cleanup(func() { interactiveRun = origRun }) var capturedCmds interactive.BuiltinCommands - interactiveRun = func(_ string, cmds interactive.BuiltinCommands, _ *config.Config, _ llm.Client, _ shell.Info) error { + interactiveRun = func(_ string, cmds interactive.BuiltinCommands, _ *config.Config, _ llm.Client, _ shell.Info, _ bool) error { capturedCmds = cmds return nil } @@ -219,7 +219,7 @@ func TestRunRoot_InteractiveMode_StatusDoctorSetModel(t *testing.T) { t.Cleanup(func() { interactiveRun = origRun }) var capturedCmds interactive.BuiltinCommands - interactiveRun = func(_ string, cmds interactive.BuiltinCommands, _ *config.Config, _ llm.Client, _ shell.Info) error { + interactiveRun = func(_ string, cmds interactive.BuiltinCommands, _ *config.Config, _ llm.Client, _ shell.Info, _ bool) error { capturedCmds = cmds return nil } @@ -253,7 +253,7 @@ func TestRunRoot_InteractiveError(t *testing.T) { origRun := interactiveRun t.Cleanup(func() { interactiveRun = origRun }) - interactiveRun = func(_ string, _ interactive.BuiltinCommands, _ *config.Config, _ llm.Client, _ shell.Info) error { + interactiveRun = func(_ string, _ interactive.BuiltinCommands, _ *config.Config, _ llm.Client, _ shell.Info, _ bool) error { return errors.New("interactive failed") } @@ -267,6 +267,46 @@ func TestRunRoot_InteractiveError(t *testing.T) { } } +func TestRunRoot_InteractiveMode_ExplainFlag(t *testing.T) { + tempHome(t) + cfg := config.DefaultConfig() + cfg.Provider.Default = config.ProviderOpenAI + cfg.Provider.OpenAI.APIKey = "sk-test-interactive" + if err := config.Save(cfg); err != nil { + t.Fatalf("Save: %v", err) + } + + origRun := interactiveRun + origExplain := explainFlag + t.Cleanup(func() { + interactiveRun = origRun + explainFlag = origExplain + }) + + var capturedExplain bool + interactiveRun = func(_ string, _ interactive.BuiltinCommands, _ *config.Config, _ llm.Client, _ shell.Info, explain bool) error { + capturedExplain = explain + return nil + } + + debugFlag = "" + explainFlag = false + if err := runRoot(nil, nil); err != nil { + t.Fatalf("runRoot() error: %v", err) + } + if capturedExplain { + t.Fatal("explain = true, want false (default)") + } + + explainFlag = true + if err := runRoot(nil, nil); err != nil { + t.Fatalf("runRoot() error: %v", err) + } + if !capturedExplain { + t.Fatal("explain = false, want true") + } +} + func TestRunRoot_DebugFileError(t *testing.T) { tempHome(t) cfg := config.DefaultConfig() diff --git a/internal/config/apply.go b/internal/config/apply.go index 3967690..c319e08 100644 --- a/internal/config/apply.go +++ b/internal/config/apply.go @@ -43,6 +43,11 @@ func ApplyAction(cfg *Config, action, key, value string) error { return fmt.Errorf("always_confirm %w", err) } cfg.Safety.AlwaysConfirm = b + case "tool_calling": + if !ValidToolCallingMode(value) { + return errors.New("tool_calling must be 'never', 'always_prompt', 'dangerous_prompt', or 'always_allow'") + } + cfg.Safety.ToolCalling = value case "min_certainty": n, err := strconv.Atoi(value) if err != nil { diff --git a/internal/config/apply_test.go b/internal/config/apply_test.go index 43e3499..a33b325 100644 --- a/internal/config/apply_test.go +++ b/internal/config/apply_test.go @@ -210,6 +210,40 @@ func TestApplyAction_SetSafety_MinCertainty(t *testing.T) { } } +func TestApplyAction_SetSafety_ToolCalling(t *testing.T) { + tests := []struct { + name string + value string + want string + wantErr bool + }{ + {name: "never", value: ToolCallingNever, want: ToolCallingNever}, + {name: "always_prompt", value: ToolCallingAlwaysPrompt, want: ToolCallingAlwaysPrompt}, + {name: "dangerous_prompt", value: ToolCallingDangerousPrompt, want: ToolCallingDangerousPrompt}, + {name: "always_allow", value: ToolCallingAlwaysAllow, want: ToolCallingAlwaysAllow}, + {name: "invalid", value: "sometimes", wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := applyTestCfg(t) + err := ApplyAction(cfg, "set_safety", "tool_calling", tt.value) + if tt.wantErr { + if err == nil { + t.Fatal("expected error, got nil") + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.Safety.ToolCalling != tt.want { + t.Errorf("ToolCalling = %q, want %q", cfg.Safety.ToolCalling, tt.want) + } + }) + } +} + func TestApplyAction_SetSafety_UnknownKey(t *testing.T) { cfg := applyTestCfg(t) err := ApplyAction(cfg, "set_safety", "nonexistent", "value") diff --git a/internal/config/config.go b/internal/config/config.go index 619d4eb..935a06d 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -16,13 +16,18 @@ const ( DebugNone = "none" DebugScreen = "screen" DebugFile = "file" + + // ToolCalling modes control whether the AI can use read-only tools. + ToolCallingNever = "never" // Tools disabled entirely + ToolCallingAlwaysPrompt = "always_prompt" // Prompt the user before every tool call + ToolCallingDangerousPrompt = "dangerous_prompt" // Only prompt when a tool hits a safety rule + ToolCallingAlwaysAllow = "always_allow" // Execute all tools without prompting ) type Config struct { Provider ProviderConfig `toml:"provider"` Safety SafetyConfig `toml:"safety"` Debug string `toml:"debug"` // "none", "screen", or "file" - Explain bool `toml:"-"` // runtime-only; set from CLI flag, not persisted } type ProviderConfig struct { @@ -40,11 +45,21 @@ type ProviderDetail struct { type SafetyConfig struct { AlwaysConfirm bool `toml:"always_confirm"` + ToolCalling string `toml:"tool_calling"` MinCertainty int `toml:"min_certainty"` AllowlistPrefixes []string `toml:"allowlist_prefixes"` WhitelistPrefixes []string `toml:"whitelist_prefixes,omitempty"` // Deprecated: use allowlist_prefixes } +// ValidToolCallingMode returns true if the given mode is a valid tool_calling value. +func ValidToolCallingMode(mode string) bool { + switch mode { + case ToolCallingNever, ToolCallingAlwaysPrompt, ToolCallingDangerousPrompt, ToolCallingAlwaysAllow: + return true + } + return false +} + func DefaultConfig() *Config { return &Config{ Provider: ProviderConfig{ @@ -62,6 +77,7 @@ func DefaultConfig() *Config { }, Safety: SafetyConfig{ AlwaysConfirm: false, + ToolCalling: ToolCallingNever, MinCertainty: 80, AllowlistPrefixes: []string{"git", "ls", "cat", "echo", "pwd", "head", "tail", "wc", "grep", "find", "which", "man"}, }, diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 9991b01..520f306 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -38,6 +38,35 @@ func TestDefaultConfig(t *testing.T) { } } +func TestValidToolCallingMode(t *testing.T) { + validModes := []string{ + ToolCallingNever, + ToolCallingAlwaysPrompt, + ToolCallingDangerousPrompt, + ToolCallingAlwaysAllow, + } + + for _, mode := range validModes { + if !ValidToolCallingMode(mode) { + t.Errorf("ValidToolCallingMode(%q) = false, want true", mode) + } + } + + invalidModes := []string{ + "", + "always", + "dangerous", + "prompt", + "ALLOW", + } + + for _, mode := range invalidModes { + if ValidToolCallingMode(mode) { + t.Errorf("ValidToolCallingMode(%q) = true, want false", mode) + } + } +} + func TestSaveAndLoad(t *testing.T) { // Use a temp dir to avoid touching the real config. // Set both HOME (Unix) and USERPROFILE (Windows) since os.UserHomeDir() diff --git a/internal/interactive/repl.go b/internal/interactive/repl.go index bbd4bcd..8a7366d 100644 --- a/internal/interactive/repl.go +++ b/internal/interactive/repl.go @@ -25,10 +25,8 @@ type replLineReader interface { var ( replConfigDir = config.Dir replNewReadline = func(cfg *readline.Config) (replLineReader, error) { return readline.NewEx(cfg) } - replBuildSystemPrompt = func(osName, shellName, shellVersion, cwd string) string { - return llm.BuildSystemPrompt(osName, shellName, shellVersion, cwd, false) - } - replHandleResponse = handleResponse + replBuildSystemPrompt = llm.BuildSystemPrompt + replHandleResponse = handleResponse ) // BuiltinCommands holds handlers for built-in REPL commands so the interactive @@ -44,7 +42,7 @@ type BuiltinCommands struct { MemoryRun func(args []string) error } -func Run(version string, cmds BuiltinCommands, cfg *config.Config, client llm.Client, shellInfo shell.Info) error { +func Run(version string, cmds BuiltinCommands, cfg *config.Config, client llm.Client, shellInfo shell.Info, explain bool) error { configDir, err := replConfigDir() if err != nil { return err @@ -61,7 +59,8 @@ func Run(version string, cmds BuiltinCommands, cfg *config.Config, client llm.Cl fmt.Printf("AI CLI %s — interactive mode. Type 'help' for commands or 'exit' to quit.\n", version) - systemPrompt := replBuildSystemPrompt(shellInfo.OS, shellInfo.Shell, shellInfo.Version, "") + toolsEnabled := cfg.Safety.ToolCalling != config.ToolCallingNever + systemPrompt := replBuildSystemPrompt(shellInfo.OS, shellInfo.Shell, shellInfo.Version, "", explain, toolsEnabled) for { line, err := rl.Readline() @@ -138,7 +137,7 @@ func Run(version string, cmds BuiltinCommands, cfg *config.Config, client llm.Cl color.Red("Error: %v", err) continue } - if err := replHandleResponse(resp, cfg, shellInfo); err != nil { + if err := replHandleResponse(resp, cfg, shellInfo, explain); err != nil { color.Red("Error: %v", err) } } @@ -165,10 +164,10 @@ func printHelp() { fmt.Println() } -func handleResponse(resp *llm.Response, cfg *config.Config, shellInfo shell.Info) error { +func handleResponse(resp *llm.Response, cfg *config.Config, shellInfo shell.Info, explain bool) error { switch resp.Type { case "commands": - return executor.Run(resp.Commands, cfg, shellInfo, cfg.Explain) + return executor.Run(resp.Commands, cfg, shellInfo, explain) case "config": return applyConfig(resp, cfg) default: diff --git a/internal/interactive/repl_run_test.go b/internal/interactive/repl_run_test.go index 56ef85c..c614525 100644 --- a/internal/interactive/repl_run_test.go +++ b/internal/interactive/repl_run_test.go @@ -17,13 +17,24 @@ import ( ) type fakeClient struct { - chat func(systemPrompt, userMessage string) (*llm.Response, error) + chat func(systemPrompt, userMessage string) (*llm.Response, error) + chatMessages func(messages []llm.Message) (*llm.Response, error) } func (f fakeClient) Chat(systemPrompt, userMessage string) (*llm.Response, error) { + if f.chat == nil { + return nil, errors.New("fakeClient.Chat not implemented") + } return f.chat(systemPrompt, userMessage) } +func (f fakeClient) ChatMessages(messages []llm.Message) (*llm.Response, error) { + if f.chatMessages == nil { + return nil, errors.New("fakeClient.ChatMessages not implemented") + } + return f.chatMessages(messages) +} + type fakeReadlineStep struct { line string err error @@ -93,7 +104,7 @@ func TestRun_ConfigDirError(t *testing.T) { err := Run("dev", BuiltinCommands{}, testCfg(t), fakeClient{chat: func(string, string) (*llm.Response, error) { t.Fatal("client.Chat should not be called") return nil, nil - }}, shell.Info{}) + }}, shell.Info{}, false) if err == nil { t.Fatal("Run() error = nil, want error") } @@ -112,7 +123,7 @@ func TestRun_ReadlineInitError(t *testing.T) { err := Run("dev", BuiltinCommands{}, testCfg(t), fakeClient{chat: func(string, string) (*llm.Response, error) { t.Fatal("client.Chat should not be called") return nil, nil - }}, shell.Info{}) + }}, shell.Info{}, false) if err == nil { t.Fatal("Run() error = nil, want error") } @@ -141,7 +152,7 @@ func TestRun_ExitOnEOFAndInterrupt(t *testing.T) { err := Run("dev", BuiltinCommands{}, testCfg(t), fakeClient{chat: func(string, string) (*llm.Response, error) { t.Fatal("client.Chat should not be called") return nil, nil - }}, shell.Info{}) + }}, shell.Info{}, false) if err != nil { t.Fatalf("Run() error = %v, want nil", err) } @@ -166,7 +177,7 @@ func TestRun_ReadlineUnexpectedError(t *testing.T) { err := Run("dev", BuiltinCommands{}, testCfg(t), fakeClient{chat: func(string, string) (*llm.Response, error) { t.Fatal("client.Chat should not be called") return nil, nil - }}, shell.Info{}) + }}, shell.Info{}, false) if err == nil { t.Fatal("Run() error = nil, want error") } @@ -198,13 +209,13 @@ func TestRun_DispatchesBuiltinsAndLLM(t *testing.T) { replNewReadline = func(*readline.Config) (replLineReader, error) { return rl, nil } var promptArgs []string - replBuildSystemPrompt = func(osName, shellName, shellVersion, cwd string) string { + replBuildSystemPrompt = func(osName, shellName, shellVersion, cwd string, _, _ bool) string { promptArgs = []string{osName, shellName, shellVersion, cwd} return "system-prompt" } var handled []*llm.Response - replHandleResponse = func(resp *llm.Response, _ *config.Config, info shell.Info) error { + replHandleResponse = func(resp *llm.Response, _ *config.Config, info shell.Info, _ bool) error { if info.Shell != "powershell" { t.Fatalf("shellInfo not passed through: %+v", info) } @@ -246,7 +257,7 @@ func TestRun_DispatchesBuiltinsAndLLM(t *testing.T) { OS: "windows/amd64", Shell: "powershell", Version: "7.5.0", - }) + }, false) if err != nil { t.Fatalf("Run() error = %v", err) } @@ -291,10 +302,10 @@ func TestRun_ContinuesAfterBuiltinLLMAndHandleErrors(t *testing.T) { }, } replNewReadline = func(*readline.Config) (replLineReader, error) { return rl, nil } - replBuildSystemPrompt = func(_, _, _, _ string) string { return "sys" } + replBuildSystemPrompt = func(_, _, _, _ string, _ bool, _ bool) string { return "sys" } handleCalls := 0 - replHandleResponse = func(_ *llm.Response, _ *config.Config, _ shell.Info) error { + replHandleResponse = func(_ *llm.Response, _ *config.Config, _ shell.Info, _ bool) error { handleCalls++ return errors.New("handle failed") } @@ -318,7 +329,7 @@ func TestRun_ContinuesAfterBuiltinLLMAndHandleErrors(t *testing.T) { } out := captureOutput(t, func() { - if err := Run("dev", cmds, testCfg(t), client, shell.Info{}); err != nil { + if err := Run("dev", cmds, testCfg(t), client, shell.Info{}, false); err != nil { t.Fatalf("Run() error = %v, want nil", err) } }) @@ -336,3 +347,42 @@ func TestRun_ContinuesAfterBuiltinLLMAndHandleErrors(t *testing.T) { t.Fatalf("output missing Bye!\n%s", out) } } + +func TestRun_ExplainFlagPropagated(t *testing.T) { + stubReplHooks(t) + replConfigDir = func() (string, error) { return t.TempDir(), nil } + rl := &fakeReadline{ + steps: []fakeReadlineStep{ + {line: "do something"}, + {line: "exit"}, + }, + } + replNewReadline = func(*readline.Config) (replLineReader, error) { return rl, nil } + + var capturedPromptExplain bool + replBuildSystemPrompt = func(_, _, _, _ string, explain, _ bool) string { + capturedPromptExplain = explain + return "sys" + } + + var capturedHandleExplain bool + replHandleResponse = func(_ *llm.Response, _ *config.Config, _ shell.Info, explain bool) error { + capturedHandleExplain = explain + return nil + } + + client := fakeClient{chat: func(_, _ string) (*llm.Response, error) { + return &llm.Response{Type: "commands"}, nil + }} + + if err := Run("dev", BuiltinCommands{}, testCfg(t), client, shell.Info{}, true); err != nil { + t.Fatalf("Run() error = %v", err) + } + + if !capturedPromptExplain { + t.Fatal("replBuildSystemPrompt: explain = false, want true") + } + if !capturedHandleExplain { + t.Fatal("replHandleResponse: explain = false, want true") + } +} diff --git a/internal/interactive/repl_test.go b/internal/interactive/repl_test.go index 8647c55..b07561b 100644 --- a/internal/interactive/repl_test.go +++ b/internal/interactive/repl_test.go @@ -45,7 +45,7 @@ func TestPrintHelp(_ *testing.T) { } func TestHandleResponseCommands(t *testing.T) { - err := handleResponse(&llm.Response{Type: "commands"}, testCfg(t), shell.Info{Shell: "cmd"}) + err := handleResponse(&llm.Response{Type: "commands"}, testCfg(t), shell.Info{Shell: "cmd"}, false) if err != nil { t.Fatalf("handleResponse(commands) error = %v", err) } @@ -61,7 +61,7 @@ func TestHandleResponseConfig(t *testing.T) { } withStdin(t, "y\n", func() { - if err := handleResponse(resp, cfg, shell.Info{}); err != nil { + if err := handleResponse(resp, cfg, shell.Info{}, false); err != nil { t.Fatalf("handleResponse(config) error = %v", err) } }) @@ -72,7 +72,7 @@ func TestHandleResponseConfig(t *testing.T) { } func TestHandleResponseUnknownType(t *testing.T) { - err := handleResponse(&llm.Response{Type: "mystery"}, testCfg(t), shell.Info{}) + err := handleResponse(&llm.Response{Type: "mystery"}, testCfg(t), shell.Info{}, false) if err == nil { t.Fatal("handleResponse() error = nil, want error") } diff --git a/internal/llm/client.go b/internal/llm/client.go index b643912..848e16f 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -18,6 +18,7 @@ import ( type Client interface { Chat(systemPrompt, userMessage string) (*Response, error) + ChatMessages(messages []Message) (*Response, error) } // DebugWriter returns an io.Writer based on the debug mode string. @@ -86,12 +87,16 @@ type openAIClient struct { } func (c *openAIClient) Chat(systemPrompt, userMessage string) (*Response, error) { + return c.ChatMessages([]Message{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userMessage}, + }) +} + +func (c *openAIClient) ChatMessages(messages []Message) (*Response, error) { req := ChatRequest{ - Model: c.model, - Messages: []Message{ - {Role: "system", Content: systemPrompt}, - {Role: "user", Content: userMessage}, - }, + Model: c.model, + Messages: messages, } body, err := json.Marshal(req) @@ -173,10 +178,35 @@ type ollamaResponse struct { } func (c *ollamaClient) Chat(systemPrompt, userMessage string) (*Response, error) { + return c.ChatMessages([]Message{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userMessage}, + }) +} + +func (c *ollamaClient) ChatMessages(messages []Message) (*Response, error) { + // Extract system and user messages for Ollama's generate API + var system, prompt string + for _, m := range messages { + switch m.Role { + case "system": + if system == "" { + system = m.Content + } else { + system += "\n" + m.Content + } + default: + if prompt != "" { + prompt += "\n" + } + prompt += m.Content + } + } + req := ollamaRequest{ Model: c.model, - Prompt: userMessage, - System: systemPrompt, + Prompt: prompt, + System: system, Stream: false, } diff --git a/internal/llm/prompt.go b/internal/llm/prompt.go index 026aa33..4e15db3 100644 --- a/internal/llm/prompt.go +++ b/internal/llm/prompt.go @@ -20,8 +20,7 @@ For shell command requests, respond with: "command": "the shell command to run", "description": "brief explanation of what it does", "risk": "safe or risky", - "certainty": 90, - "explanation": "(optional) step-by-step breakdown of the command's flags and arguments" + "certainty": 90 } ] } @@ -29,12 +28,12 @@ For shell command requests, respond with: Rules for commands: - "risk" must be "safe" (read-only, informational) or "risky" (modifies files, processes, system state) - "certainty" is 0-100, your confidence this is the correct command for what the user asked -- "explanation" is optional and only required when explicitly requested - For multi-step tasks, return multiple commands in order. Use shell constructs like $(...) or pipes to chain when possible - Generate commands appropriate for the detected OS and shell - Never generate commands that could cause irreversible damage without clear user intent {{EXPLAIN_INSTRUCTION}} {{PLATFORM_HINTS}} +{{TOOL_INSTRUCTIONS}} For requests to change AI CLI configuration (model, provider, API key, safety settings), respond with: { "type": "config", @@ -65,6 +64,7 @@ This system uses BSD userland, NOT GNU coreutils. You MUST use BSD-compatible fl - tar: macOS ships bsdtar, NOT GNU tar. Do NOT use "--wildcards" or other GNU-specific flags. - find: "find -printf" is NOT available (GNU-only). Use "-print0 | xargs -0" or "-exec" instead. - awk: macOS ships BSD awk, NOT gawk. Do NOT use gawk features like "--csv" or "BEGINFILE/ENDFILE". +- To update ai-cli via Homebrew: "brew upgrade ai-cli" (or "brew upgrade kriserickson/tap/ai-cli" if needed). ` const linuxHints = ` @@ -74,6 +74,8 @@ This system uses GNU coreutils. Prefer GNU-specific flags when they are clearer: - sed: use "sed -i" for in-place editing (no extra argument needed unlike BSD sed). - date: GNU date supports "date -d" for date arithmetic. - grep: "grep -P" (PCRE) is available for advanced patterns. +- To update ai-cli: download the latest release from GitHub using curl. Example: + VERSION=$(curl -s https://api.github.com/repos/kriserickson/ai-cli/releases/latest | grep '"tag_name"' | cut -d'"' -f4) && ARCH=$(uname -m | sed 's/x86_64/amd64/') && curl -LO "https://github.com/kriserickson/ai-cli/releases/download/${VERSION}/ai-${VERSION}-linux-${ARCH}.tar.gz" && tar -xzf "ai-${VERSION}-linux-${ARCH}.tar.gz" && chmod +x ai && sudo mv ai /usr/local/bin/ai ` const windowsHints = ` @@ -81,6 +83,38 @@ Platform-specific rules (Windows): - Use PowerShell cmdlets when the shell is powershell or pwsh. - For cmd.exe, use standard Windows commands (dir, type, tasklist, etc.). - Do NOT use Unix commands unless running under WSL or Git Bash. +- To update ai-cli: + - If the shell is powershell or pwsh: + $VERSION = (Invoke-RestMethod "https://api.github.com/repos/kriserickson/ai-cli/releases/latest").tag_name; Invoke-WebRequest -Uri "https://github.com/kriserickson/ai-cli/releases/download/${VERSION}/ai-${VERSION}-windows-amd64.zip" -OutFile "ai-${VERSION}-windows-amd64.zip"; Expand-Archive -Path "ai-${VERSION}-windows-amd64.zip" -DestinationPath .\ai-update -Force; Move-Item .\ai-update\ai.exe (Get-Command ai).Source -Force; Remove-Item "ai-${VERSION}-windows-amd64.zip"; Remove-Item .\ai-update -Recurse + - If the shell is bash (Git Bash/MSYS2): + VERSION=$(curl -s https://api.github.com/repos/kriserickson/ai-cli/releases/latest | grep '"tag_name"' | cut -d'"' -f4) && curl -LO "https://github.com/kriserickson/ai-cli/releases/download/${VERSION}/ai-${VERSION}-windows-amd64.zip" && unzip -o "ai-${VERSION}-windows-amd64.zip" -d ai-update && mv ai-update/ai.exe "$(which ai)" && rm -rf "ai-${VERSION}-windows-amd64.zip" ai-update +` + +const toolInstructions = `You have access to read-only tools to gather information before generating commands. +To use a tool, respond with: +{ + "type": "tool_request", + "tool": "tool_name", + "args": {"arg_name": "value"} +} + +Available tools: +- list_directory: List files in a directory. Args: path (default ".") +- read_file: Read file contents (max 10KB, safety-checked). Args: path +- command_help: Get help/man page for a command. Args: command +- list_memories: List stored AI CLI memories. No args. +- list_processes: List running processes. No args. +- system_resources: Show top processes by CPU/memory. No args. +- network_connections: Show active network connections. No args. +- ping: Check host connectivity (3 packets). Args: host +- check_command: Check if a command is installed. Args: command +- disk_usage: Show disk space usage. No args. +- environment: Show environment variables (sensitive values masked). No args. + +Rules for tools: +- Use tools ONLY when you need information to generate better commands +- Maximum 3 tool calls per request +- After gathering information, respond with a "commands" or "config" response as usual ` const explainInstruction = ` @@ -91,7 +125,7 @@ const explainInstruction = ` - For "grep -rn 'TODO' --include='*.go' .": "Recursively searches all .go files in the current directory for lines containing 'TODO'. -r enables recursive search, -n shows line numbers, --include restricts to files matching the pattern." ` -func BuildSystemPrompt(osInfo, shell, shellVersion, cwd string, explain bool) string { +func BuildSystemPrompt(osInfo, shell, shellVersion, cwd string, explain, toolsEnabled bool) string { platformHints := buildPlatformHints(osInfo) var explainText string @@ -99,6 +133,11 @@ func BuildSystemPrompt(osInfo, shell, shellVersion, cwd string, explain bool) st explainText = explainInstruction } + var toolText string + if toolsEnabled { + toolText = toolInstructions + } + r := strings.NewReplacer( "{{OS}}", osInfo, "{{SHELL}}", shell, @@ -106,6 +145,7 @@ func BuildSystemPrompt(osInfo, shell, shellVersion, cwd string, explain bool) st "{{CWD}}", cwd, "{{PLATFORM_HINTS}}", platformHints, "{{EXPLAIN_INSTRUCTION}}", explainText, + "{{TOOL_INSTRUCTIONS}}", toolText, ) return r.Replace(promptTemplate) diff --git a/internal/llm/prompt_test.go b/internal/llm/prompt_test.go index 75570b2..17d42c4 100644 --- a/internal/llm/prompt_test.go +++ b/internal/llm/prompt_test.go @@ -6,7 +6,7 @@ import ( ) func TestBuildSystemPrompt_ContainsEnvironment(t *testing.T) { - prompt := BuildSystemPrompt("darwin/arm64", "/bin/zsh", "zsh 5.9", "/home/user", false) + prompt := BuildSystemPrompt("darwin/arm64", "/bin/zsh", "zsh 5.9", "/home/user", false, false) checks := []string{ "darwin/arm64", @@ -22,13 +22,12 @@ func TestBuildSystemPrompt_ContainsEnvironment(t *testing.T) { } func TestBuildSystemPrompt_ContainsJSONInstructions(t *testing.T) { - prompt := BuildSystemPrompt("linux/amd64", "/bin/bash", "5.1", "/tmp", false) + prompt := BuildSystemPrompt("linux/amd64", "/bin/bash", "5.1", "/tmp", false, false) checks := []string{ `"type": "commands"`, `"risk"`, `"certainty"`, - `"explanation"`, `"type": "config"`, "valid JSON", } @@ -40,7 +39,7 @@ func TestBuildSystemPrompt_ContainsJSONInstructions(t *testing.T) { } func TestBuildSystemPrompt_DarwinPlatformHints(t *testing.T) { - prompt := BuildSystemPrompt("darwin/arm64", "/bin/zsh", "zsh 5.9", "/Users/test", false) + prompt := BuildSystemPrompt("darwin/arm64", "/bin/zsh", "zsh 5.9", "/Users/test", false, false) mustContain := []string{ "BSD userland", @@ -69,7 +68,7 @@ func TestBuildSystemPrompt_DarwinPlatformHints(t *testing.T) { } func TestBuildSystemPrompt_LinuxPlatformHints(t *testing.T) { - prompt := BuildSystemPrompt("linux/amd64", "/bin/bash", "5.1", "/home/test", false) + prompt := BuildSystemPrompt("linux/amd64", "/bin/bash", "5.1", "/home/test", false, false) mustContain := []string{ "uses GNU coreutils", @@ -96,7 +95,7 @@ func TestBuildSystemPrompt_LinuxPlatformHints(t *testing.T) { } func TestBuildSystemPrompt_WindowsPlatformHints(t *testing.T) { - prompt := BuildSystemPrompt("windows/amd64", "powershell", "5.1", "C:\\Users\\test", false) + prompt := BuildSystemPrompt("windows/amd64", "powershell", "5.1", "C:\\Users\\test", false, false) mustContain := []string{ "PowerShell cmdlets", @@ -168,7 +167,7 @@ func TestAppendMemories_MultipleMemories(t *testing.T) { } func TestBuildSystemPrompt_UnknownOSNoPlatformHints(t *testing.T) { - prompt := BuildSystemPrompt("freebsd/amd64", "/bin/sh", "sh 1.0", "/home/test", false) + prompt := BuildSystemPrompt("freebsd/amd64", "/bin/sh", "sh 1.0", "/home/test", false, false) platformSections := []string{ "BSD userland", @@ -191,7 +190,7 @@ func TestBuildSystemPrompt_UnknownOSNoPlatformHints(t *testing.T) { } func TestBuildSystemPrompt_ExplainTrue(t *testing.T) { - prompt := BuildSystemPrompt("linux/amd64", "/bin/bash", "5.1", "/tmp", true) + prompt := BuildSystemPrompt("linux/amd64", "/bin/bash", "5.1", "/tmp", true, false) mustContain := []string{ "explanation", @@ -208,7 +207,7 @@ func TestBuildSystemPrompt_ExplainTrue(t *testing.T) { } func TestBuildSystemPrompt_ExplainFalse(t *testing.T) { - prompt := BuildSystemPrompt("linux/amd64", "/bin/bash", "5.1", "/tmp", false) + prompt := BuildSystemPrompt("linux/amd64", "/bin/bash", "5.1", "/tmp", false, false) if strings.Contains(prompt, "detailed explanation") { t.Error("explain=false prompt should not contain explanation instructions") diff --git a/internal/llm/types.go b/internal/llm/types.go index 25a4cdd..f9d5450 100644 --- a/internal/llm/types.go +++ b/internal/llm/types.go @@ -34,6 +34,9 @@ type Response struct { Action string `json:"action,omitempty"` Key string `json:"key,omitempty"` Value string `json:"value,omitempty"` + // Tool request fields (type == "tool_request") + Tool string `json:"tool,omitempty"` + ToolArgs map[string]string `json:"args,omitempty"` } type Command struct { diff --git a/internal/shell/detect_test.go b/internal/shell/detect_test.go index 6c3421a..ef5aa0d 100644 --- a/internal/shell/detect_test.go +++ b/internal/shell/detect_test.go @@ -230,4 +230,3 @@ func TestDetectShellVersion_Branches(t *testing.T) { } } } - diff --git a/internal/tools/runner.go b/internal/tools/runner.go new file mode 100644 index 0000000..d1c521f --- /dev/null +++ b/internal/tools/runner.go @@ -0,0 +1,144 @@ +package tools + +import ( + "encoding/json" + "fmt" + "os" + "strings" + + "github.com/fatih/color" + + "github.com/kriserickson/ai-cli/internal/config" + "github.com/kriserickson/ai-cli/internal/llm" + "github.com/kriserickson/ai-cli/internal/shell" +) + +type confirmFunc func(toolName string, args map[string]string, reason string) bool + +func defaultConfirm(toolName string, args map[string]string, reason string) bool { + argsStr := "" + if len(args) > 0 { + parts := make([]string, 0, len(args)) + for k, v := range args { + parts = append(parts, k+"="+v) + } + argsStr = " (" + strings.Join(parts, ", ") + ")" + } + if reason != "" { + color.Yellow("Warning: %s", reason) + } + fmt.Printf("AI wants to use tool %q%s. Allow? [Y/n] ", toolName, argsStr) + var input string + _, _ = fmt.Scanln(&input) + input = strings.TrimSpace(strings.ToLower(input)) + return input == "" || input == "y" || input == "yes" +} + +// RunWithTools calls the LLM and handles tool_request responses in a loop +// (up to maxIter iterations). Returns the final non-tool response. +// If tool_calling is "never", it calls the LLM once with Chat() (no tool loop). +func RunWithTools(client llm.Client, systemPrompt, userMessage string, cfg *config.Config, shellInfo shell.Info, maxIter int) (*llm.Response, error) { + return runWithTools(client, systemPrompt, userMessage, cfg, shellInfo, maxIter, defaultConfirm) +} + +func runWithTools(client llm.Client, systemPrompt, userMessage string, cfg *config.Config, shellInfo shell.Info, maxIter int, confirm confirmFunc) (*llm.Response, error) { + // If tools are disabled, just do a simple chat call + if cfg.Safety.ToolCalling == config.ToolCallingNever { + return client.Chat(systemPrompt, userMessage) + } + + messages := []llm.Message{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userMessage}, + } + + for range maxIter { + resp, err := client.ChatMessages(messages) + if err != nil { + return nil, err + } + + if resp.Type != "tool_request" { + return resp, nil + } + + // Validate tool exists + if !ValidTool(resp.Tool) { + return nil, fmt.Errorf("AI requested unknown tool: %s", resp.Tool) + } + + // Check if this tool call would trigger a safety issue (for dangerous_prompt mode) + safetyIssue := checkToolSafety(resp.Tool, resp.ToolArgs) + + // Determine whether to prompt based on mode + switch cfg.Safety.ToolCalling { + case config.ToolCallingAlwaysPrompt: + if !confirm(resp.Tool, resp.ToolArgs, "") { + messages = appendDenial(messages, resp) + continue + } + case config.ToolCallingDangerousPrompt: + if safetyIssue != "" { + if !confirm(resp.Tool, resp.ToolArgs, safetyIssue) { + messages = appendDenial(messages, resp) + continue + } + } + case config.ToolCallingAlwaysAllow: + // No prompting + default: + return nil, fmt.Errorf("unknown tool_calling mode %q; valid modes are: never, always_prompt, dangerous_prompt, always_allow", cfg.Safety.ToolCalling) + } + + // Execute tool + color.Cyan("Using tool: %s", resp.Tool) + output, err := Execute(resp.Tool, resp.ToolArgs, shellInfo) + if err != nil { + output = "Tool error: " + err.Error() + } + + // Build follow-up messages + toolReqJSON, _ := json.Marshal(resp) + messages = append(messages, + llm.Message{Role: "assistant", Content: string(toolReqJSON)}, + llm.Message{Role: "user", Content: fmt.Sprintf("Tool result for %s:\n%s", resp.Tool, output)}, + ) + } + + // Max iterations reached — make one final call + messages = append(messages, + llm.Message{Role: "user", Content: "You have used the maximum number of tools. Please provide your final response now."}, + ) + return client.ChatMessages(messages) +} + +// appendDenial adds a denial message to the conversation for the LLM to continue without the tool. +func appendDenial(messages []llm.Message, resp *llm.Response) []llm.Message { + color.Yellow("Tool request denied. Asking AI to proceed without it.") + toolReqJSON, _ := json.Marshal(resp) + return append(messages, + llm.Message{Role: "assistant", Content: string(toolReqJSON)}, + llm.Message{Role: "user", Content: "Tool request denied by user. Please generate your best response without using tools."}, + ) +} + +// checkToolSafety checks if a tool call would trigger a safety concern. +// Returns a description of the issue, or "" if the call is safe. +func checkToolSafety(toolName string, args map[string]string) string { + switch toolName { + case "read_file", "list_directory": + path := args["path"] + if path == "" || path == "." { + return "" + } + cwd, err := os.Getwd() + if err != nil { + return "" + } + _, err = ValidatePath(path, cwd) + if err != nil { + return fmt.Sprintf("path %q triggers safety rule: %s", path, err.Error()) + } + } + return "" +} diff --git a/internal/tools/runner_test.go b/internal/tools/runner_test.go new file mode 100644 index 0000000..624344b --- /dev/null +++ b/internal/tools/runner_test.go @@ -0,0 +1,269 @@ +package tools + +import ( + "errors" + "strings" + "testing" + + "github.com/kriserickson/ai-cli/internal/config" + "github.com/kriserickson/ai-cli/internal/llm" + "github.com/kriserickson/ai-cli/internal/shell" +) + +// mockClient implements llm.Client for testing. +type mockClient struct { + responses []*llm.Response + errors []error + calls int + messages [][]llm.Message // captured messages from each call +} + +func (m *mockClient) Chat(systemPrompt, userMessage string) (*llm.Response, error) { + return m.ChatMessages([]llm.Message{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userMessage}, + }) +} + +func (m *mockClient) ChatMessages(messages []llm.Message) (*llm.Response, error) { + idx := m.calls + m.calls++ + m.messages = append(m.messages, messages) + if idx < len(m.errors) && m.errors[idx] != nil { + return nil, m.errors[idx] + } + if idx < len(m.responses) { + return m.responses[idx], nil + } + return &llm.Response{Type: "commands", Commands: []llm.Command{{Command: "echo fallback", Description: "fallback", Risk: "safe", Certainty: 100}}}, nil +} + +func toolCallingCfg(mode string) *config.Config { + cfg := config.DefaultConfig() + cfg.Safety.ToolCalling = mode + return cfg +} + +func TestRunWithTools_NeverMode_SkipsToolLoop(t *testing.T) { + client := &mockClient{ + responses: []*llm.Response{ + {Type: "commands", Commands: []llm.Command{{Command: "ls", Description: "list", Risk: "safe", Certainty: 95}}}, + }, + } + + resp, err := RunWithTools(client, "system", "list files", toolCallingCfg(config.ToolCallingNever), shell.Info{}, 3) + if err != nil { + t.Fatalf("RunWithTools error: %v", err) + } + if resp.Type != "commands" { + t.Errorf("type = %q, want commands", resp.Type) + } + if client.calls != 1 { + t.Errorf("calls = %d, want 1", client.calls) + } +} + +func TestRunWithTools_NeverMode_IgnoresToolRequest(t *testing.T) { + // Even if the LLM returns a tool_request, "never" mode uses Chat() which + // returns whatever the LLM says. The tool_request won't be acted on. + client := &mockClient{ + responses: []*llm.Response{ + // Chat() is used, so this is the one response + {Type: "tool_request", Tool: "disk_usage", ToolArgs: map[string]string{}}, + }, + } + + resp, err := RunWithTools(client, "system", "test", toolCallingCfg(config.ToolCallingNever), shell.Info{}, 3) + if err != nil { + t.Fatalf("RunWithTools error: %v", err) + } + // In never mode, the raw response is returned without entering the tool loop + if resp.Type != "tool_request" { + t.Errorf("type = %q, want tool_request (passed through without processing)", resp.Type) + } + if client.calls != 1 { + t.Errorf("calls = %d, want 1 (no loop)", client.calls) + } +} + +func TestRunWithTools_AlwaysAllowMode_NoPrompt(t *testing.T) { + client := &mockClient{ + responses: []*llm.Response{ + {Type: "tool_request", Tool: "check_command", ToolArgs: map[string]string{"command": "ls"}}, + {Type: "commands", Commands: []llm.Command{{Command: "ls -la", Description: "list", Risk: "safe", Certainty: 95}}}, + }, + } + + resp, err := RunWithTools(client, "system", "list files", toolCallingCfg(config.ToolCallingAlwaysAllow), shell.Info{}, 3) + if err != nil { + t.Fatalf("RunWithTools error: %v", err) + } + if resp.Type != "commands" { + t.Errorf("type = %q, want commands", resp.Type) + } + if client.calls != 2 { + t.Errorf("calls = %d, want 2", client.calls) + } + + // Verify tool result was injected + lastMessages := client.messages[1] + found := false + for _, m := range lastMessages { + if strings.Contains(m.Content, "Tool result for check_command") { + found = true + break + } + } + if !found { + t.Error("expected tool result in follow-up messages") + } +} + +func TestRunWithTools_AlwaysPromptMode_Approved(t *testing.T) { + client := &mockClient{ + responses: []*llm.Response{ + {Type: "tool_request", Tool: "disk_usage", ToolArgs: map[string]string{}}, + {Type: "commands", Commands: []llm.Command{{Command: "df -h", Description: "disk", Risk: "safe", Certainty: 90}}}, + }, + } + + resp, err := runWithTools(client, "system", "show disk", toolCallingCfg(config.ToolCallingAlwaysPrompt), shell.Info{}, 3, func(string, map[string]string, string) bool { + return true + }) + if err != nil { + t.Fatalf("RunWithTools error: %v", err) + } + if resp.Type != "commands" { + t.Errorf("type = %q, want commands", resp.Type) + } + if client.calls != 2 { + t.Errorf("calls = %d, want 2", client.calls) + } +} + +func TestRunWithTools_AlwaysPromptMode_Denied(t *testing.T) { + client := &mockClient{ + responses: []*llm.Response{ + {Type: "tool_request", Tool: "disk_usage", ToolArgs: map[string]string{}}, + {Type: "commands", Commands: []llm.Command{{Command: "df -h", Description: "disk", Risk: "safe", Certainty: 90}}}, + }, + } + + resp, err := runWithTools(client, "system", "show disk", toolCallingCfg(config.ToolCallingAlwaysPrompt), shell.Info{}, 3, func(string, map[string]string, string) bool { + return false + }) + if err != nil { + t.Fatalf("RunWithTools error: %v", err) + } + if resp.Type != "commands" { + t.Errorf("type = %q, want commands", resp.Type) + } + + // Check that denial message was sent + lastMessages := client.messages[1] + found := false + for _, m := range lastMessages { + if strings.Contains(m.Content, "Tool request denied") { + found = true + break + } + } + if !found { + t.Error("expected denial message in follow-up") + } +} + +func TestRunWithTools_DangerousPromptMode_SafeToolNoPrompt(t *testing.T) { + // In dangerous_prompt mode, a safe tool (no safety issue) should NOT prompt + promptCalled := false + confirm := func(string, map[string]string, string) bool { + promptCalled = true + return true + } + + client := &mockClient{ + responses: []*llm.Response{ + {Type: "tool_request", Tool: "disk_usage", ToolArgs: map[string]string{}}, + {Type: "commands", Commands: []llm.Command{{Command: "df -h", Description: "disk", Risk: "safe", Certainty: 90}}}, + }, + } + + resp, err := runWithTools(client, "system", "show disk", toolCallingCfg(config.ToolCallingDangerousPrompt), shell.Info{}, 3, confirm) + if err != nil { + t.Fatalf("RunWithTools error: %v", err) + } + if resp.Type != "commands" { + t.Errorf("type = %q, want commands", resp.Type) + } + if promptCalled { + t.Error("should not have prompted for a safe tool in dangerous_prompt mode") + } +} + +func TestRunWithTools_UnknownToolCallingMode(t *testing.T) { + client := &mockClient{ + responses: []*llm.Response{ + {Type: "tool_request", Tool: "disk_usage", ToolArgs: map[string]string{}}, + }, + } + + _, err := RunWithTools(client, "system", "test", toolCallingCfg("invalid_mode"), shell.Info{}, 3) + if err == nil { + t.Fatal("expected error for unknown tool_calling mode") + } + if !strings.Contains(err.Error(), "unknown tool_calling mode") { + t.Errorf("error = %q, want 'unknown tool_calling mode'", err.Error()) + } +} + +func TestRunWithTools_UnknownTool(t *testing.T) { + client := &mockClient{ + responses: []*llm.Response{ + {Type: "tool_request", Tool: "nonexistent_tool", ToolArgs: map[string]string{}}, + }, + } + + _, err := RunWithTools(client, "system", "test", toolCallingCfg(config.ToolCallingAlwaysAllow), shell.Info{}, 3) + if err == nil { + t.Fatal("expected error for unknown tool") + } + if !strings.Contains(err.Error(), "unknown tool") { + t.Errorf("error = %q, want 'unknown tool'", err.Error()) + } +} + +func TestRunWithTools_MaxIterations(t *testing.T) { + client := &mockClient{ + responses: []*llm.Response{ + {Type: "tool_request", Tool: "disk_usage", ToolArgs: map[string]string{}}, + {Type: "tool_request", Tool: "disk_usage", ToolArgs: map[string]string{}}, + {Type: "tool_request", Tool: "disk_usage", ToolArgs: map[string]string{}}, + // Response to "max iterations" follow-up + {Type: "commands", Commands: []llm.Command{{Command: "df -h", Description: "disk usage", Risk: "safe", Certainty: 90}}}, + }, + } + + resp, err := RunWithTools(client, "system", "show disk", toolCallingCfg(config.ToolCallingAlwaysAllow), shell.Info{}, 3) + if err != nil { + t.Fatalf("RunWithTools error: %v", err) + } + if resp.Type != "commands" { + t.Errorf("type = %q, want commands", resp.Type) + } + // 3 tool iterations + 1 final call = 4 + if client.calls != 4 { + t.Errorf("calls = %d, want 4", client.calls) + } +} + +func TestRunWithTools_ClientError(t *testing.T) { + client := &mockClient{ + responses: []*llm.Response{nil}, + errors: []error{errors.New("test error")}, + } + + _, err := RunWithTools(client, "system", "test", toolCallingCfg(config.ToolCallingAlwaysAllow), shell.Info{}, 3) + if err == nil { + t.Fatal("expected error") + } +} diff --git a/internal/tools/safety.go b/internal/tools/safety.go new file mode 100644 index 0000000..1814947 --- /dev/null +++ b/internal/tools/safety.go @@ -0,0 +1,165 @@ +package tools + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +// blockedPatterns lists file/directory patterns that must never be read or listed, +// even within the current working directory. +var blockedPatterns = []string{ + ".ssh/", + ".gnupg/", + ".env", + ".env.", + "*.pem", + "*.key", + "*.p12", + "*.pfx", + "id_rsa", + "id_ed25519", + "credentials", + "secrets", + "tokens", + ".netrc", + ".npmrc", + "shadow", + "passwd", + "private", + "*.keystore", + "config.toml", +} + +// sensitiveEnvKeys lists substrings that, if found in an environment variable +// name, cause the value to be redacted. +var sensitiveEnvKeys = []string{ + "KEY", + "SECRET", + "TOKEN", + "PASSWORD", + "CREDENTIAL", + "AUTH", + "PRIVATE", + "DATABASE", + "DB", + "DSN", +} + +// ValidatePath resolves path to an absolute path and checks that it is +// contained within cwd and does not match any blocked pattern. +// It also resolves symlinks and re-validates the resolved path to prevent +// symlink attacks that could point outside cwd or to blocked files. +func ValidatePath(path, cwd string) (string, error) { + // Resolve relative to cwd + var abs string + if filepath.IsAbs(path) { + abs = filepath.Clean(path) + } else { + abs = filepath.Clean(filepath.Join(cwd, path)) + } + + // Always use forward slashes for cross-platform consistency in tests + // but convert back to OS specific separators for internal checks + cwdClean := filepath.Clean(cwd) + + // Ensure the path is under cwd + if abs != cwdClean && !strings.HasPrefix(abs, cwdClean+string(os.PathSeparator)) { + return "", fmt.Errorf("path %q is outside the working directory", path) + } + + // Check blocked patterns against the relative path + rel, err := filepath.Rel(cwdClean, abs) + if err != nil { + return "", fmt.Errorf("cannot compute relative path: %w", err) + } + + if isBlocked(rel) { + return "", fmt.Errorf("access to %q is blocked for security", path) + } + + // Resolve symlinks and re-validate the resolved path to prevent symlink + // attacks where a path within cwd points to a target outside cwd or + // matching a blocked pattern (e.g. safe.txt -> /etc/passwd). + resolved, err := filepath.EvalSymlinks(abs) + if err == nil { + resolvedClean := filepath.Clean(resolved) + if resolvedClean != cwdClean && !strings.HasPrefix(resolvedClean, cwdClean+string(os.PathSeparator)) { + return "", fmt.Errorf("path %q resolves outside the working directory", path) + } + resolvedRel, relErr := filepath.Rel(cwdClean, resolvedClean) + if relErr != nil { + return "", fmt.Errorf("cannot compute relative path for resolved symlink: %w", relErr) + } + if isBlocked(resolvedRel) { + return "", fmt.Errorf("access to %q is blocked for security", path) + } + } + // If EvalSymlinks fails (e.g. path does not exist yet), the lexical + // checks above are sufficient and we fall through. + + return abs, nil +} + +// isBlocked checks if any component of the relative path matches a blocked pattern. +func isBlocked(relPath string) bool { + // Normalize separators + normalized := filepath.ToSlash(relPath) + parts := strings.Split(normalized, "/") + + for _, pattern := range blockedPatterns { + // Directory pattern (ends with /) + if strings.HasSuffix(pattern, "/") { + dirName := strings.TrimSuffix(pattern, "/") + for _, part := range parts { + if strings.EqualFold(part, dirName) { + return true + } + } + continue + } + + // Check the final component (filename) + filename := parts[len(parts)-1] + + // Glob-style pattern (starts with *) + if strings.HasPrefix(pattern, "*") { + suffix := pattern[1:] + if strings.HasSuffix(strings.ToLower(filename), strings.ToLower(suffix)) { + return true + } + continue + } + + // Prefix match (e.g., "id_rsa" matches "id_rsa", "id_rsa.pub") + if strings.HasPrefix(strings.ToLower(filename), strings.ToLower(pattern)) { + return true + } + } + return false +} + +// FilterEnvironment takes a list of "KEY=VALUE" strings and returns +// a filtered copy where sensitive values are replaced with [REDACTED]. +func FilterEnvironment(vars []string) []string { + result := make([]string, 0, len(vars)) + for _, v := range vars { + parts := strings.SplitN(v, "=", 2) + if len(parts) != 2 { + continue + } + key := parts[0] + value := parts[1] + + upperKey := strings.ToUpper(key) + for _, sensitive := range sensitiveEnvKeys { + if strings.Contains(upperKey, sensitive) { + value = "[REDACTED]" + break + } + } + result = append(result, key+"="+value) + } + return result +} diff --git a/internal/tools/safety_test.go b/internal/tools/safety_test.go new file mode 100644 index 0000000..069529f --- /dev/null +++ b/internal/tools/safety_test.go @@ -0,0 +1,236 @@ +package tools + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestValidatePath_Valid(t *testing.T) { + cwd, _ := os.Getwd() + tests := []struct { + path string + want string + }{ + {".", cwd}, + {"src", filepath.Join(cwd, "src")}, + {"./src/main.go", filepath.Join(cwd, "src", "main.go")}, + {"src/../src/main.go", filepath.Join(cwd, "src", "main.go")}, + } + for _, tt := range tests { + got, err := ValidatePath(tt.path, cwd) + if err != nil { + t.Errorf("ValidatePath(%q, %q) error: %v", tt.path, cwd, err) + continue + } + if got != tt.want { + t.Errorf("ValidatePath(%q, %q) = %q, want %q", tt.path, cwd, got, tt.want) + } + } +} + +func TestValidatePath_OutsideCWD(t *testing.T) { + cwd, _ := os.Getwd() + paths := []string{ + "../other", + "/etc/passwd", + "../../..", + } + for _, path := range paths { + _, err := ValidatePath(path, cwd) + if err == nil { + t.Errorf("ValidatePath(%q, %q) should have failed", path, cwd) + continue + } + if !strings.Contains(err.Error(), "outside") && !strings.Contains(err.Error(), "blocked") { + t.Errorf("ValidatePath(%q) error = %q, want 'outside' or 'blocked'", path, err.Error()) + } + } +} + +func TestValidatePath_BlockedPatterns(t *testing.T) { + cwd := "/home/user/project" + blocked := []string{ + ".ssh/id_rsa", + ".env", + ".env.production", + "keys/server.pem", + "keys/server.key", + "cert.p12", + "store.pfx", + "id_rsa", + "id_rsa.pub", + "id_ed25519", + "credentials.json", + "secrets.yaml", + "tokens.json", + ".netrc", + ".npmrc", + "private.key", + "app.keystore", + "config.toml", + ".gnupg/pubring.gpg", + } + for _, path := range blocked { + _, err := ValidatePath(path, cwd) + if err == nil { + t.Errorf("ValidatePath(%q) should have been blocked", path) + } + if !strings.Contains(err.Error(), "blocked") { + t.Errorf("ValidatePath(%q) error = %q, want 'blocked'", path, err.Error()) + } + } +} + +func TestValidatePath_AllowedFiles(t *testing.T) { + cwd := "/home/user/project" + allowed := []string{ + "main.go", + "src/app.js", + "README.md", + "Makefile", + } + for _, path := range allowed { + _, err := ValidatePath(path, cwd) + if err != nil { + t.Errorf("ValidatePath(%q) should be allowed, got error: %v", path, err) + } + } +} + +func TestFilterEnvironment(t *testing.T) { + vars := []string{ + "HOME=/home/user", + "PATH=/usr/bin", + "API_KEY=secret123", + "AWS_SECRET_ACCESS_KEY=mysecret", + "DATABASE_URL=postgres://localhost", + "AUTH_TOKEN=tok123", + "MY_PASSWORD=pass", + "PRIVATE_DATA=stuff", + "CREDENTIAL_FILE=/etc/creds", + "NORMAL_VAR=hello", + } + + filtered := FilterEnvironment(vars) + + expect := map[string]string{ + "HOME": "/home/user", + "PATH": "/usr/bin", + "API_KEY": "[REDACTED]", + "AWS_SECRET_ACCESS_KEY": "[REDACTED]", + "DATABASE_URL": "[REDACTED]", + "AUTH_TOKEN": "[REDACTED]", + "MY_PASSWORD": "[REDACTED]", + "PRIVATE_DATA": "[REDACTED]", + "CREDENTIAL_FILE": "[REDACTED]", + "NORMAL_VAR": "hello", + } + + for _, v := range filtered { + parts := strings.SplitN(v, "=", 2) + if len(parts) != 2 { + t.Errorf("bad format: %q", v) + continue + } + if want, ok := expect[parts[0]]; ok { + if parts[1] != want { + t.Errorf("%s = %q, want %q", parts[0], parts[1], want) + } + } + } +} + +func TestValidatePath_SymlinkOutsideCWD(t *testing.T) { + // Create a temporary directory to act as the project root (cwd). + cwd, err := os.MkdirTemp("", "cwd-*") + if err != nil { + t.Fatal(err) + } + defer os.RemoveAll(cwd) + + // Create a real file outside the cwd. + outside, err := os.MkdirTemp("", "outside-*") + if err != nil { + t.Fatal(err) + } + defer os.RemoveAll(outside) + + target := filepath.Join(outside, "secret.txt") + if err := os.WriteFile(target, []byte("secret"), 0o600); err != nil { + t.Fatal(err) + } + + // Create a symlink inside cwd pointing to the outside file. + link := filepath.Join(cwd, "safe.txt") + if err := os.Symlink(target, link); err != nil { + t.Skip("symlinks not supported:", err) + } + + _, err = ValidatePath("safe.txt", cwd) + if err == nil { + t.Error("ValidatePath should have rejected symlink pointing outside cwd") + } + if err != nil && !strings.Contains(err.Error(), "outside") && !strings.Contains(err.Error(), "resolves outside") { + t.Errorf("expected 'outside' or 'resolves outside' in error, got: %v", err) + } +} + +func TestValidatePath_SymlinkToBlockedFile(t *testing.T) { + // Create a temporary directory to act as the project root (cwd). + cwd, err := os.MkdirTemp("", "cwd-*") + if err != nil { + t.Fatal(err) + } + defer os.RemoveAll(cwd) + + // Create a blocked-pattern file inside cwd (e.g. id_rsa). + blocked := filepath.Join(cwd, "id_rsa") + if err := os.WriteFile(blocked, []byte("private key"), 0o600); err != nil { + t.Fatal(err) + } + + // Create a symlink with an innocuous name pointing to the blocked file. + link := filepath.Join(cwd, "readme.txt") + if err := os.Symlink(blocked, link); err != nil { + t.Skip("symlinks not supported:", err) + } + + _, err = ValidatePath("readme.txt", cwd) + if err == nil { + t.Error("ValidatePath should have rejected symlink pointing to a blocked file") + } + if err != nil && !strings.Contains(err.Error(), "blocked") { + t.Errorf("expected 'blocked' in error, got: %v", err) + } +} + +func TestValidatePath_SymlinkWithinCWD(t *testing.T) { + // Create a temporary directory to act as the project root (cwd). + cwd, err := os.MkdirTemp("", "cwd-*") + if err != nil { + t.Fatal(err) + } + defer os.RemoveAll(cwd) + + // Create a regular (non-blocked) file inside cwd. + mainFile := filepath.Join(cwd, "main.go") + if err := os.WriteFile(mainFile, []byte("package main"), 0o600); err != nil { + t.Fatal(err) + } + + // Create a symlink inside cwd pointing to the file within cwd. + link := filepath.Join(cwd, "alias.go") + if err := os.Symlink(mainFile, link); err != nil { + t.Skip("symlinks not supported:", err) + } + + got, err := ValidatePath("alias.go", cwd) + if err != nil { + t.Errorf("ValidatePath should allow symlink within cwd, got error: %v", err) + } + if got != link { + t.Errorf("ValidatePath returned %q, want %q", got, link) + } +} diff --git a/internal/tools/tools.go b/internal/tools/tools.go new file mode 100644 index 0000000..37f4faa --- /dev/null +++ b/internal/tools/tools.go @@ -0,0 +1,345 @@ +package tools + +import ( + "context" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "sort" + "strings" + "time" + + "github.com/kriserickson/ai-cli/internal/memory" + "github.com/kriserickson/ai-cli/internal/shell" +) + +const ( + maxOutputBytes = 4096 + windowsOS = "windows" + toolExecTimeout = 30 * time.Second +) + +// ToolDef describes a tool the AI can request. +type ToolDef struct { + Name string + Description string + Args []string // expected arg keys +} + +// Registry is the list of available tools. +var Registry = []ToolDef{ + {Name: "list_directory", Description: "List files in a directory", Args: []string{"path"}}, + {Name: "read_file", Description: "Read file contents (max 10KB, safety-checked)", Args: []string{"path"}}, + {Name: "command_help", Description: "Get help/man page for a command", Args: []string{"command"}}, + {Name: "list_memories", Description: "List stored AI CLI memories", Args: nil}, + {Name: "list_processes", Description: "List running processes", Args: nil}, + {Name: "system_resources", Description: "Show top processes by CPU/memory", Args: nil}, + {Name: "network_connections", Description: "Show active network connections", Args: nil}, + {Name: "ping", Description: "Check host connectivity (3 packets)", Args: []string{"host"}}, + {Name: "check_command", Description: "Check if a command is installed", Args: []string{"command"}}, + {Name: "disk_usage", Description: "Show disk space usage", Args: nil}, + {Name: "environment", Description: "Show environment variables (sensitive values masked)", Args: nil}, +} + +// Execute runs the named tool with the given args and returns its output. +func Execute(toolName string, args map[string]string, _ shell.Info) (string, error) { + cwd, err := os.Getwd() + if err != nil { + return "", fmt.Errorf("cannot get working directory: %w", err) + } + + var output string + + switch toolName { + case "list_directory": + output, err = execListDirectory(args["path"], cwd) + case "read_file": + output, err = execReadFile(args["path"], cwd) + case "command_help": + output, err = execCommandHelp(args["command"]) + case "list_memories": + output, err = execListMemories() + case "list_processes": + output, err = execListProcesses() + case "system_resources": + output, err = execSystemResources() + case "network_connections": + output, err = execNetworkConnections() + case "ping": + output, err = execPing(args["host"]) + case "check_command": + output, err = execCheckCommand(args["command"]) + case "disk_usage": + output, err = execDiskUsage() + case "environment": + output = execEnvironment() + default: + return "", fmt.Errorf("unknown tool: %s", toolName) + } + + if err != nil { + return "", err + } + + return truncateOutput(output), nil +} + +// ValidTool returns true if the named tool exists in the registry. +func ValidTool(name string) bool { + for _, t := range Registry { + if t.Name == name { + return true + } + } + return false +} + +func execListDirectory(path, cwd string) (string, error) { + if path == "" { + path = "." + } + absPath, err := ValidatePath(path, cwd) + if err != nil { + return "", err + } + + entries, err := os.ReadDir(absPath) + if err != nil { + return "", fmt.Errorf("list_directory failed: %w", err) + } + + var b strings.Builder + for _, entry := range entries { + entryAbs := filepath.Join(absPath, entry.Name()) + rel, relErr := filepath.Rel(cwd, entryAbs) + if relErr != nil { + // Fall back to checking the entry name alone when relative path cannot be computed + if isBlocked(entry.Name()) { + continue + } + } else if isBlocked(rel) { + continue + } + info, infoErr := entry.Info() + if infoErr != nil { + fmt.Fprintf(&b, "? ? ??? ?? ??:?? %s (error reading info)\n", entry.Name()) + continue + } + fmt.Fprintf(&b, "%s %8d %s %s\n", + info.Mode(), + info.Size(), + info.ModTime().Format("Jan 2 15:04"), + entry.Name(), + ) + } + return b.String(), nil +} + +func execReadFile(path, cwd string) (string, error) { + if path == "" { + return "", errors.New("read_file requires a path argument") + } + absPath, err := ValidatePath(path, cwd) + if err != nil { + return "", err + } + + info, err := os.Stat(absPath) + if err != nil { + return "", fmt.Errorf("cannot stat %q: %w", path, err) + } + if info.IsDir() { + return "", fmt.Errorf("%q is a directory, not a file", path) + } + if info.Size() > 10*1024 { + return "", fmt.Errorf("file %q is too large (%d bytes, max 10KB)", path, info.Size()) + } + + data, err := os.ReadFile(absPath) + if err != nil { + return "", fmt.Errorf("cannot read %q: %w", path, err) + } + return string(data), nil +} + +func execCommandHelp(command string) (string, error) { + if command == "" { + return "", errors.New("command_help requires a command argument") + } + + ctx, cancel := context.WithTimeout(context.Background(), toolExecTimeout) + defer cancel() + + if runtime.GOOS == windowsOS { + cmd := exec.CommandContext(context.Background(), "powershell", "-Command", "Get-Help -Name $env:HELP_COMMAND") + cmd.Env = append(os.Environ(), "HELP_COMMAND="+command) + out, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("Get-Help failed: %s", string(out)) + } + return string(out), nil + } + + // Try tldr first, fall back to man + if tldrPath, err := exec.LookPath("tldr"); err == nil { + cmd := exec.CommandContext(ctx, tldrPath, command) + out, err := cmd.CombinedOutput() + if err == nil { + return string(out), nil + } + } + + cmd := exec.CommandContext(ctx, "man", command) + out, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("man page not found for %q", command) + } + return string(out), nil +} + +func execListMemories() (string, error) { + entries, err := memory.Load() + if err != nil { + return "", fmt.Errorf("failed to load memories: %w", err) + } + if len(entries) == 0 { + return "No memories stored.", nil + } + var b strings.Builder + for _, e := range entries { + fmt.Fprintf(&b, "- %s: %s\n", e.Keyword, e.Content) + } + return b.String(), nil +} + +func execListProcesses() (string, error) { + ctx, cancel := context.WithTimeout(context.Background(), toolExecTimeout) + defer cancel() + + var cmd *exec.Cmd + if runtime.GOOS == windowsOS { + cmd = exec.CommandContext(ctx, "powershell", "-Command", "Get-Process | Format-Table -AutoSize") + } else { + cmd = exec.CommandContext(ctx, "ps", "aux") + } + out, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("list_processes failed: %s", string(out)) + } + return string(out), nil +} + +func execSystemResources() (string, error) { + ctx, cancel := context.WithTimeout(context.Background(), toolExecTimeout) + defer cancel() + + var cmd *exec.Cmd + switch runtime.GOOS { + case windowsOS: + cmd = exec.CommandContext(ctx, "powershell", "-Command", "Get-Process | Sort-Object CPU -Descending | Select-Object -First 5 | Format-Table Name,CPU,WorkingSet -AutoSize") + case "darwin": + cmd = exec.CommandContext(ctx, "top", "-l", "1", "-n", "5", "-s", "0") + default: // linux + cmd = exec.CommandContext(ctx, "top", "-bn1", "-o", "%CPU") + } + out, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("system_resources failed: %s", string(out)) + } + return string(out), nil +} + +func execNetworkConnections() (string, error) { + ctx, cancel := context.WithTimeout(context.Background(), toolExecTimeout) + defer cancel() + + var cmd *exec.Cmd + if runtime.GOOS == windowsOS { + cmd = exec.CommandContext(ctx, "powershell", "-Command", "Get-NetTCPConnection | Format-Table -AutoSize") + } else if ssPath, err := exec.LookPath("ss"); err == nil { + cmd = exec.CommandContext(ctx, ssPath, "-an") + } else if netstatPath, err := exec.LookPath("netstat"); err == nil { + cmd = exec.CommandContext(ctx, netstatPath, "-an") + } else { + return "", errors.New("network_connections failed: neither ss nor netstat found") + } + out, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("network_connections failed: %s", string(out)) + } + return string(out), nil +} + +func execPing(host string) (string, error) { + if host == "" { + return "", errors.New("ping requires a host argument") + } + + ctx, cancel := context.WithTimeout(context.Background(), toolExecTimeout) + defer cancel() + + var cmd *exec.Cmd + if runtime.GOOS == windowsOS { + cmd = exec.CommandContext(ctx, "ping", "-n", "3", host) + } else { + cmd = exec.CommandContext(ctx, "ping", "-c", "3", host) + } + out, err := cmd.CombinedOutput() + if err != nil && ctx.Err() != nil { + return "", fmt.Errorf("ping timed out after %s", toolExecTimeout) + } + return string(out), nil +} + +func execCheckCommand(command string) (string, error) { + if command == "" { + return "", errors.New("check_command requires a command argument") + } + + foundPath, _ := exec.LookPath(command) + if foundPath == "" { + return command + ": not found", nil + } + + return foundPath, nil +} + +func execDiskUsage() (string, error) { + ctx, cancel := context.WithTimeout(context.Background(), toolExecTimeout) + defer cancel() + + var cmd *exec.Cmd + if runtime.GOOS == windowsOS { + cmd = exec.CommandContext(ctx, "powershell", "-Command", "Get-PSDrive -PSProvider FileSystem | Format-Table Name,Used,Free -AutoSize") + } else { + cmd = exec.CommandContext(ctx, "df", "-h") + } + out, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("disk_usage failed: %s", string(out)) + } + return string(out), nil +} + +func execEnvironment() string { + vars := os.Environ() + sort.Strings(vars) + filtered := FilterEnvironment(vars) + return strings.Join(filtered, "\n") +} + +func truncateOutput(s string) string { + if len(s) <= maxOutputBytes { + return s + } + + lastNewline := strings.LastIndex(s[:maxOutputBytes], "\n") + if lastNewline > maxOutputBytes/2 { + return s[:lastNewline] + "\n... [output truncated]" + } + + return s[:maxOutputBytes] + "\n... [output truncated]" +} diff --git a/internal/tools/tools_additional_test.go b/internal/tools/tools_additional_test.go new file mode 100644 index 0000000..64e9af1 --- /dev/null +++ b/internal/tools/tools_additional_test.go @@ -0,0 +1,469 @@ +package tools + +import ( + "bytes" + "io" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + + "github.com/kriserickson/ai-cli/internal/memory" + "github.com/kriserickson/ai-cli/internal/shell" +) + +func writeFakeCommand(t *testing.T, dir, name, body string) { + t.Helper() + path := filepath.Join(dir, name) + if err := os.WriteFile(path, []byte(body), 0o755); err != nil { + t.Fatalf("os.WriteFile(%s): %v", name, err) + } +} + +func TestDefaultConfirm(t *testing.T) { + oldStdin := os.Stdin + oldStdout := os.Stdout + defer func() { + os.Stdin = oldStdin + os.Stdout = oldStdout + }() + + inR, inW, err := os.Pipe() + if err != nil { + t.Fatalf("os.Pipe stdin: %v", err) + } + outR, outW, err := os.Pipe() + if err != nil { + t.Fatalf("os.Pipe stdout: %v", err) + } + + os.Stdin = inR + os.Stdout = outW + + if _, err := inW.WriteString("yes\n"); err != nil { + t.Fatalf("stdin write: %v", err) + } + _ = inW.Close() + + if ok := defaultConfirm("read_file", map[string]string{"path": "README.md"}, "warning"); !ok { + t.Fatal("defaultConfirm() = false, want true") + } + + _ = outW.Close() + var buf bytes.Buffer + if _, err := io.Copy(&buf, outR); err != nil { + t.Fatalf("stdout read: %v", err) + } + + output := buf.String() + if !strings.Contains(output, `AI wants to use tool "read_file"`) { + t.Fatalf("prompt output missing tool name: %q", output) + } +} + +func TestCheckToolSafety(t *testing.T) { + dir := t.TempDir() + + if got := checkToolSafety("disk_usage", nil); got != "" { + t.Fatalf("checkToolSafety(non-path tool) = %q, want empty", got) + } + if got := checkToolSafety("read_file", map[string]string{"path": ""}); got != "" { + t.Fatalf("checkToolSafety(empty path) = %q, want empty", got) + } + if got := checkToolSafety("read_file", map[string]string{"path": "."}); got != "" { + t.Fatalf("checkToolSafety(dot path) = %q, want empty", got) + } + + oldWd, err := os.Getwd() + if err != nil { + t.Fatalf("os.Getwd: %v", err) + } + if err := os.Chdir(dir); err != nil { + t.Fatalf("os.Chdir: %v", err) + } + defer func() { + if err := os.Chdir(oldWd); err != nil { + t.Fatalf("restore working directory: %v", err) + } + }() + + if got := checkToolSafety("read_file", map[string]string{"path": "notes.txt"}); got != "" { + t.Fatalf("checkToolSafety(safe file) = %q, want empty", got) + } + if got := checkToolSafety("read_file", map[string]string{"path": ".env"}); !strings.Contains(got, "blocked") { + t.Fatalf("checkToolSafety(blocked file) = %q, want blocked message", got) + } +} + +func TestExecListDirectory_OutsideCWD(t *testing.T) { + dir := t.TempDir() + _, err := execListDirectory("../", dir) + if err == nil { + t.Fatal("execListDirectory() error = nil, want error") + } +} + +func TestExecListDirectory_DefaultPath(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "file.txt"), []byte("x"), 0o644); err != nil { + t.Fatalf("os.WriteFile: %v", err) + } + + output, err := execListDirectory("", dir) + if err != nil { + t.Fatalf("execListDirectory(\"\"): %v", err) + } + if !strings.Contains(output, "file.txt") { + t.Fatalf("execListDirectory(\"\") = %q, want listed file", output) + } +} + +func TestExecListDirectory_BlockedEntriesFiltered(t *testing.T) { + dir := t.TempDir() + // Create a safe file and several blocked files/dirs + for _, name := range []string{"visible.txt", ".env", "secret.key", "id_rsa"} { + if err := os.WriteFile(filepath.Join(dir, name), []byte("x"), 0o644); err != nil { + t.Fatalf("os.WriteFile(%s): %v", name, err) + } + } + if err := os.MkdirAll(filepath.Join(dir, ".ssh"), 0o700); err != nil { + t.Fatalf("os.MkdirAll(.ssh): %v", err) + } + + output, err := execListDirectory("", dir) + if err != nil { + t.Fatalf("execListDirectory: %v", err) + } + if !strings.Contains(output, "visible.txt") { + t.Errorf("output should contain visible.txt, got: %s", output) + } + for _, blocked := range []string{".env", "secret.key", "id_rsa", ".ssh"} { + if strings.Contains(output, blocked) { + t.Errorf("output should not contain blocked entry %q, got: %s", blocked, output) + } + } +} + +func TestExecListDirectory_CommandFailure(t *testing.T) { + dir := t.TempDir() + nonExistent := filepath.Join(dir, "does_not_exist") + + _, err := execListDirectory(nonExistent, dir) + if err == nil { + t.Fatal("execListDirectory() error = nil, want error") + } +} + +func TestExecReadFile_DirectoryAndMissing(t *testing.T) { + dir := t.TempDir() + + _, err := execReadFile(".", dir) + if err == nil || !strings.Contains(err.Error(), "directory") { + t.Fatalf("execReadFile(directory) error = %v, want directory error", err) + } + + _, err = execReadFile("missing.txt", dir) + if err == nil || !strings.Contains(err.Error(), "cannot stat") { + t.Fatalf("execReadFile(missing) error = %v, want stat error", err) + } +} + +func TestExecCommandHelp(t *testing.T) { + _, err := execCommandHelp("") + if err == nil { + t.Fatal("execCommandHelp(\"\") error = nil, want error") + } + + if runtime.GOOS == windowsOS { + t.Skip("success path is shell-dependent on Windows") + } + + output, err := execCommandHelp("ls") + if err != nil { + t.Fatalf("execCommandHelp(ls): %v", err) + } + if output == "" { + t.Fatal("execCommandHelp(ls) returned empty output") + } +} + +func TestExecCommandHelp_UsesTldrAndManError(t *testing.T) { + if runtime.GOOS == windowsOS { + t.Skip("PATH-based command stubs are Unix-specific") + } + + dir := t.TempDir() + writeFakeCommand(t, dir, "tldr", "#!/bin/sh\necho 'TLDR help'\n") + t.Setenv("PATH", dir) + + output, err := execCommandHelp("ls") + if err != nil { + t.Fatalf("execCommandHelp(tldr): %v", err) + } + if !strings.Contains(output, "TLDR help") { + t.Fatalf("execCommandHelp(tldr) = %q, want fake tldr output", output) + } + + dir = t.TempDir() + writeFakeCommand(t, dir, "man", "#!/bin/sh\nexit 1\n") + t.Setenv("PATH", dir) + + _, err = execCommandHelp("ls") + if err == nil || !strings.Contains(err.Error(), "man page not found") { + t.Fatalf("execCommandHelp(man error) = %v, want man error", err) + } +} + +func TestExecCommandHelp_UsesMan(t *testing.T) { + if runtime.GOOS == windowsOS { + t.Skip("PATH-based command stubs are Unix-specific") + } + + dir := t.TempDir() + writeFakeCommand(t, dir, "man", "#!/bin/sh\necho 'MANPAGE help'\n") + t.Setenv("PATH", dir) + + output, err := execCommandHelp("ls") + if err != nil { + t.Fatalf("execCommandHelp(man): %v", err) + } + if !strings.Contains(output, "MANPAGE help") { + t.Fatalf("execCommandHelp(man) = %q, want fake man output", output) + } +} + +func TestExecCommandHelp_FallsBackFromTldrToMan(t *testing.T) { + if runtime.GOOS == windowsOS { + t.Skip("PATH-based command stubs are Unix-specific") + } + + dir := t.TempDir() + writeFakeCommand(t, dir, "tldr", "#!/bin/sh\nexit 1\n") + writeFakeCommand(t, dir, "man", "#!/bin/sh\necho 'fallback man output'\n") + t.Setenv("PATH", dir) + + output, err := execCommandHelp("ls") + if err != nil { + t.Fatalf("execCommandHelp(tldr fallback): %v", err) + } + if !strings.Contains(output, "fallback man output") { + t.Fatalf("execCommandHelp(tldr fallback) = %q, want man fallback output", output) + } +} + +func TestExecListMemories_WithEntries(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + + if err := memory.Save([]memory.Entry{ + {Keyword: "server", Content: "ssh user@example.com"}, + }); err != nil { + t.Fatalf("memory.Save: %v", err) + } + + output, err := execListMemories() + if err != nil { + t.Fatalf("execListMemories: %v", err) + } + if !strings.Contains(output, "server") || !strings.Contains(output, "ssh user@example.com") { + t.Fatalf("execListMemories() = %q, want saved entry", output) + } +} + +func TestExecListMemories_ParseError(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + + memDir := filepath.Join(home, ".ai-cli") + if err := os.MkdirAll(memDir, 0o755); err != nil { + t.Fatalf("os.MkdirAll: %v", err) + } + if err := os.WriteFile(filepath.Join(memDir, "memory.json"), []byte("{invalid json"), 0o600); err != nil { + t.Fatalf("os.WriteFile: %v", err) + } + + _, err := execListMemories() + if err == nil || !strings.Contains(err.Error(), "failed to parse memory.json") { + t.Fatalf("execListMemories() error = %v, want parse error", err) + } +} + +func TestExecute_ProcessAndNetworkTools(t *testing.T) { + if runtime.GOOS == windowsOS { + t.Skip("PATH-based command stubs are Unix-specific") + } + + dir := t.TempDir() + writeFakeCommand(t, dir, "ps", "#!/bin/sh\necho 'fake process list'\n") + writeFakeCommand(t, dir, "top", "#!/bin/sh\necho 'fake system resources'\n") + writeFakeCommand(t, dir, "netstat", "#!/bin/sh\necho 'fake network connections'\n") + writeFakeCommand(t, dir, "ping", "#!/bin/sh\necho 'fake ping output'\n") + t.Setenv("PATH", dir) + + tests := []struct { + name string + tool string + args map[string]string + }{ + {name: "list processes", tool: "list_processes"}, + {name: "system resources", tool: "system_resources"}, + {name: "network connections", tool: "network_connections"}, + {name: "ping", tool: "ping", args: map[string]string{"host": "127.0.0.1"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + output, err := Execute(tt.tool, tt.args, shell.Info{}) + if err != nil { + t.Fatalf("Execute(%s) error: %v", tt.tool, err) + } + if output == "" { + t.Fatalf("Execute(%s) returned empty output", tt.tool) + } + }) + } +} + +func TestExecPing_NoHost(t *testing.T) { + _, err := execPing("") + if err == nil { + t.Fatal("execPing(\"\") error = nil, want error") + } +} + +func TestShellOutHelpers_ErrorPaths(t *testing.T) { + if runtime.GOOS == windowsOS { + t.Skip("PATH-based command stubs are Unix-specific") + } + + tests := []struct { + name string + command string + body string + run func() (string, error) + }{ + { + name: "list_processes", + command: "ps", + body: "#!/bin/sh\nexit 1\n", + run: execListProcesses, + }, + { + name: "system_resources", + command: "top", + body: "#!/bin/sh\nexit 1\n", + run: execSystemResources, + }, + { + name: "network_connections", + command: "netstat", + body: "#!/bin/sh\nexit 1\n", + run: execNetworkConnections, + }, + { + name: "disk_usage", + command: "df", + body: "#!/bin/sh\nexit 1\n", + run: execDiskUsage, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dir := t.TempDir() + writeFakeCommand(t, dir, tt.command, tt.body) + t.Setenv("PATH", dir) + + _, err := tt.run() + if err == nil { + t.Fatalf("%s error = nil, want error", tt.name) + } + }) + } +} + +func TestShellOutHelpers_SuccessPaths(t *testing.T) { + if runtime.GOOS == windowsOS { + t.Skip("PATH-based command stubs are Unix-specific") + } + + tests := []struct { + name string + command string + output string + run func() (string, error) + }{ + { + name: "list_processes", + command: "ps", + output: "process list", + run: execListProcesses, + }, + { + name: "system_resources", + command: "top", + output: "system resources", + run: execSystemResources, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dir := t.TempDir() + writeFakeCommand(t, dir, tt.command, "#!/bin/sh\necho '"+tt.output+"'\n") + t.Setenv("PATH", dir) + + output, err := tt.run() + if err != nil { + t.Fatalf("%s error: %v", tt.name, err) + } + if !strings.Contains(output, tt.output) { + t.Fatalf("%s output = %q, want %q", tt.name, output, tt.output) + } + }) + } +} + +func TestExecCheckCommand_NoArgument(t *testing.T) { + _, err := execCheckCommand("") + if err == nil { + t.Fatal("execCheckCommand(\"\") error = nil, want error") + } +} + +func TestExecute_CommandHelp(t *testing.T) { + if runtime.GOOS == windowsOS { + t.Skip("success path is shell-dependent on Windows") + } + + output, err := Execute("command_help", map[string]string{"command": "ls"}, shell.Info{}) + if err != nil { + t.Fatalf("Execute(command_help) error: %v", err) + } + if !strings.Contains(strings.ToLower(output), "ls") { + t.Fatalf("Execute(command_help) output = %q, want command help", output) + } +} + +func TestExecute_ListMemories_WithConfigHome(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + + memDir := filepath.Join(home, ".ai-cli") + if err := os.MkdirAll(memDir, 0o755); err != nil { + t.Fatalf("os.MkdirAll: %v", err) + } + if err := os.WriteFile(filepath.Join(memDir, "memory.json"), []byte(`[{"keyword":"db","content":"postgres://localhost"}]`), 0o600); err != nil { + t.Fatalf("os.WriteFile: %v", err) + } + + output, err := Execute("list_memories", nil, shell.Info{}) + if err != nil { + t.Fatalf("Execute(list_memories) error: %v", err) + } + if !strings.Contains(output, "postgres://localhost") { + t.Fatalf("Execute(list_memories) = %q, want memory content", output) + } +} diff --git a/internal/tools/tools_test.go b/internal/tools/tools_test.go new file mode 100644 index 0000000..845eb19 --- /dev/null +++ b/internal/tools/tools_test.go @@ -0,0 +1,227 @@ +package tools + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/kriserickson/ai-cli/internal/shell" +) + +func TestValidTool(t *testing.T) { + if !ValidTool("list_directory") { + t.Error("list_directory should be valid") + } + if !ValidTool("read_file") { + t.Error("read_file should be valid") + } + if ValidTool("nonexistent") { + t.Error("nonexistent should not be valid") + } +} + +func TestExecute_UnknownTool(t *testing.T) { + _, err := Execute("nonexistent", nil, shell.Info{}) + if err == nil { + t.Fatal("expected error for unknown tool") + } + if !strings.Contains(err.Error(), "unknown tool") { + t.Errorf("error = %q, want 'unknown tool'", err.Error()) + } +} + +func TestExecute_ListDirectory(t *testing.T) { + dir := t.TempDir() + // Create a test file + if err := os.WriteFile(filepath.Join(dir, "test.txt"), []byte("hello"), 0o644); err != nil { + t.Fatal(err) + } + + // Change to temp dir so path validation works + oldWd, _ := os.Getwd() + if err := os.Chdir(dir); err != nil { + t.Fatal(err) + } + defer func() { + if err := os.Chdir(oldWd); err != nil { + t.Fatalf("restore working directory: %v", err) + } + }() + + output, err := Execute("list_directory", map[string]string{"path": "."}, shell.Info{}) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + if !strings.Contains(output, "test.txt") { + t.Errorf("output should contain test.txt, got: %s", output) + } +} + +func TestExecute_ReadFile(t *testing.T) { + dir := t.TempDir() + content := "hello world" + if err := os.WriteFile(filepath.Join(dir, "test.txt"), []byte(content), 0o644); err != nil { + t.Fatal(err) + } + + oldWd, _ := os.Getwd() + if err := os.Chdir(dir); err != nil { + t.Fatal(err) + } + defer func() { + if err := os.Chdir(oldWd); err != nil { + t.Fatalf("restore working directory: %v", err) + } + }() + + output, err := Execute("read_file", map[string]string{"path": "test.txt"}, shell.Info{}) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + if output != content { + t.Errorf("output = %q, want %q", output, content) + } +} + +func TestExecute_ReadFile_TooLarge(t *testing.T) { + dir := t.TempDir() + // Create a file larger than 10KB + data := make([]byte, 11*1024) + if err := os.WriteFile(filepath.Join(dir, "big.txt"), data, 0o644); err != nil { + t.Fatal(err) + } + + oldWd, _ := os.Getwd() + if err := os.Chdir(dir); err != nil { + t.Fatal(err) + } + defer func() { + if err := os.Chdir(oldWd); err != nil { + t.Fatalf("restore working directory: %v", err) + } + }() + + _, err := Execute("read_file", map[string]string{"path": "big.txt"}, shell.Info{}) + if err == nil { + t.Fatal("expected error for file too large") + } + if !strings.Contains(err.Error(), "too large") { + t.Errorf("error = %q, want 'too large'", err.Error()) + } +} + +func TestExecute_ReadFile_Blocked(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, ".env"), []byte("SECRET=x"), 0o644); err != nil { + t.Fatal(err) + } + + oldWd, _ := os.Getwd() + if err := os.Chdir(dir); err != nil { + t.Fatal(err) + } + defer func() { + if err := os.Chdir(oldWd); err != nil { + t.Fatalf("restore working directory: %v", err) + } + }() + + _, err := Execute("read_file", map[string]string{"path": ".env"}, shell.Info{}) + if err == nil { + t.Fatal("expected error for blocked file") + } + if !strings.Contains(err.Error(), "blocked") { + t.Errorf("error = %q, want 'blocked'", err.Error()) + } +} + +func TestExecute_ReadFile_NoPath(t *testing.T) { + _, err := Execute("read_file", map[string]string{}, shell.Info{}) + if err == nil { + t.Fatal("expected error for missing path") + } +} + +func TestExecute_CheckCommand(t *testing.T) { + output, err := Execute("check_command", map[string]string{"command": "ls"}, shell.Info{}) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + // On Unix, "which ls" returns a path + if !strings.Contains(output, "ls") { + t.Errorf("output should mention ls, got: %s", output) + } +} + +func TestExecute_CheckCommand_NotFound(t *testing.T) { + output, err := Execute("check_command", map[string]string{"command": "nonexistent_command_xyz"}, shell.Info{}) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + if output == "" { + t.Errorf("expected some output for nonexistent command, got empty string") + } +} + +func TestExecute_Environment(t *testing.T) { + t.Setenv("TEST_SAFE_VAR", "visible") + t.Setenv("TEST_API_KEY", "should_be_hidden") + + output, err := Execute("environment", nil, shell.Info{}) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + + // In CI, many environment variables might exist causing truncation. + // Since "environment" tool sorts by however os.Environ() returns them (usually unsorted or alphabetical), + // and we added them recently, they should be present, but if truncation happens they might be lost. + // We'll check if truncation happened and if so, we'll be more lenient or use a more targeted test. + + if !strings.Contains(output, "TEST_SAFE_VAR=visible") && !strings.Contains(output, "[output truncated]") { + t.Error("expected safe var to be visible") + } + if strings.Contains(output, "should_be_hidden") { + t.Error("expected API_KEY value to be redacted") + } + if !strings.Contains(output, "TEST_API_KEY=[REDACTED]") && !strings.Contains(output, "[output truncated]") { + t.Error("expected API_KEY to show [REDACTED]") + } +} + +func TestExecute_DiskUsage(t *testing.T) { + output, err := Execute("disk_usage", nil, shell.Info{}) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + if output == "" { + t.Error("disk_usage should return some output") + } +} + +func TestExecute_ListMemories(t *testing.T) { + // This will return "No memories stored." since there's no memory file in test + output, err := Execute("list_memories", nil, shell.Info{}) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + if output == "" { + t.Error("list_memories should return some output") + } +} + +func TestTruncateOutput(t *testing.T) { + short := "hello" + if truncateOutput(short) != short { + t.Error("short output should not be truncated") + } + + long := strings.Repeat("x", maxOutputBytes+100) + result := truncateOutput(long) + if len(result) > maxOutputBytes+30 { + t.Errorf("truncated output too long: %d", len(result)) + } + if !strings.Contains(result, "[output truncated]") { + t.Error("truncated output should contain truncation notice") + } +}