diff --git a/cmd/server/main.go b/cmd/server/main.go index 5671098a..ecbfc2b5 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -2,6 +2,7 @@ package main import ( "fmt" + "path/filepath" "go-llm-demo/configs" "go-llm-demo/internal/server/infra/provider" @@ -20,6 +21,10 @@ func main() { fmt.Printf("设置工作区失败：%v\n", err) return } + if err := initializeSecurity(filepath.Join(workspaceRoot, "configs", "security")); err != nil { + fmt.Printf("初始化安全策略失败：%v\n", err) + return + } if err := configs.LoadAppConfig("config.yaml"); err != nil { fmt.Printf("加载配置失败：%v\n", err) @@ -54,3 +59,13 @@ func main() { fmt.Printf("Server initialized with services: %+v\n", chatGateway) fmt.Println("Note: This is a placeholder. Actual server implementation goes here.") } + +func initializeSecurity(configDir string) error { + securityRepo := repository.NewSecurityConfigRepository() + securitySvc := service.NewSecurityService(securityRepo) + if err := securitySvc.Initialize(configDir); err != nil { + return err + } + tools.SetSecurityChecker(securitySvc) + return nil +} diff --git a/cmd/tui/main.go b/cmd/tui/main.go index 86ce0ce6..7614cd19 100644 --- a/cmd/tui/main.go +++ b/cmd/tui/main.go @@ -7,9 +7,13 @@ import ( "fmt" "io" "os" + "path/filepath" "strings" "go-llm-demo/configs" + "go-llm-demo/internal/server/infra/repository" + "go-llm-demo/internal/server/infra/tools" + "go-llm-demo/internal/server/service" "go-llm-demo/internal/tui/bootstrap" tea "github.com/charmbracelet/bubbletea" @@ -56,6 +60,7 @@ func main() { if err != nil { os.Exit(1) } + _ = loadDotEnv(".env") if err := run(workspaceFlag, os.Stdin, os.Stdout, os.Stderr); err != nil { fmt.Fprintln(os.Stderr, err) @@ -88,6 +93,10 @@ func runWithDeps(workspaceFlag string, deps runDeps) error { return fmt.Errorf("解析工作区失败: %w", err) } + if err := initializeSecurity(filepath.Join(workspaceRoot, "configs", "security")); err != nil { + return fmt.Errorf("安全策略初始化失败：%w", err) + } + scanner := bufio.NewScanner(deps.stdin) ready, err := deps.ensureAPIKeyInteractive(context.Background(), scanner, defaultConfigPath) if err != nil { @@ -121,6 +130,16 @@ func runWithDeps(workspaceFlag string, deps runDeps) error { return nil } +func initializeSecurity(configDir string) error { + securityRepo := repository.NewSecurityConfigRepository() + securitySvc := service.NewSecurityService(securityRepo) + if err := securitySvc.Initialize(configDir); err != nil { + return err + } + tools.SetSecurityChecker(securitySvc) + return nil +} + func loadDotEnv(path string) error { data, err := os.ReadFile(path) if err != nil { diff --git a/cmd/tui/utf8_other.go b/cmd/tui/utf8_other.go index 079648c7..d5c5aaab 100644 --- a/cmd/tui/utf8_other.go +++ b/cmd/tui/utf8_other.go @@ -1,6 +1,7 @@ //go:build !windows package main + // setUTF8Mode 在非 Windows 系统上是空操作。 // Linux 和 macOS 默认使用 UTF-8 编码，无需额外设置。 func setUTF8Mode() {} diff --git a/configs/persona.txt b/configs/persona.txt index bd28d82b..51b9d54a 100644 --- a/configs/persona.txt +++ b/configs/persona.txt @@ -6,7 +6,15 @@ - 对不确定的内容明确说明假设 - 默认使用中文回答 -你可以调用edit,grep,list,read,write工具,规范和opencode保持一致 +安全与工具执行要求（必须遵守）： +- 对可能有副作用的操作（尤其是 bash、写文件、网络请求），先进行风险判断，再执行。 +- 当命中安全策略为 deny 时，必须明确拒绝并解释原因，不得尝试绕过。 +- 当命中安全策略为 ask 时，先向用户确认，再继续执行。 +- 执行 bash 前必须说明命令目的、影响范围与关键参数；优先使用只读、可回滚、最小权限命令。 +- 严禁执行高危破坏性命令（如无边界删除、系统级破坏、恶意下载执行等）。 +- 涉及路径操作时，默认限定在当前工作区内，不对工作区外路径进行读写。 + +你可以调用edit,grep,list,read,write,bash工具,规范和opencode保持一致 当需要查看指定目录下的文件 / 目录结构时，调用 list 工具。 必填参数：path（目标目录路径，如/home/project）； @@ -29,10 +37,16 @@ 必填参数：filePath（目标文件路径）old_str（待替换的旧文本）、new_str（替换后的新文本）； 可选参数：backup（是否创建文件备份，布尔值，默认 true） +当需要在工作区目录内执行任意终端命令、脚本、系统操作时，调用 bash 工具。 +必填参数：command（待执行的 bash 命令字符串，如 grep 'error' ./log） +可选参数：workdir（命令执行目录，字符串，默认工作区根目录），timeout（命令超时时间，整数，单位毫秒，默认 120000），description（命令用途说明，字符串，默认空） 你必须严格按照以下JSON格式返回工具调用指令，不要返回其他内容： { "tool": "工具名称（如list/read/bash）", "params": {"参数名": "参数值"}, "thought": "你调用该工具的原因" -} \ No newline at end of file +} + +调用工具时，你只能严格遵守json格式，不能输出多余字符，一次只能调用一个工具 +输出json完成后，主动终止对话，等待工具响应 \ No newline at end of file diff --git a/configs/persona.txt.example b/configs/persona.txt.example index 491ac36d..0d0b2141 100644 --- a/configs/persona.txt.example +++ b/configs/persona.txt.example @@ -5,3 +5,11 @@ - 回答简洁清晰，必要时给出步骤 - 对不确定的内容明确说明假设 - 默认使用中文回答 + +安全与工具执行要求（必须遵守）： +- 对可能有副作用的操作（尤其是 bash、写文件、网络请求），先进行风险判断，再执行。 +- 当命中安全策略为 deny 时，必须明确拒绝并解释原因，不得尝试绕过。 +- 当命中安全策略为 ask 时，先向用户确认，再继续执行。 +- 执行 bash 前必须说明命令目的、影响范围与关键参数；优先使用只读、可回滚、最小权限命令。 +- 严禁执行高危破坏性命令（如无边界删除、系统级破坏、恶意下载执行等）。 +- 涉及路径操作时，默认限定在当前工作区内，不对工作区外路径进行读写。 diff --git a/internal/server/infra/tools/bash.go b/internal/server/infra/tools/bash.go index dd5b949f..cfc0641e 100644 --- a/internal/server/infra/tools/bash.go +++ b/internal/server/infra/tools/bash.go @@ -5,6 +5,7 @@ import ( "context" "fmt" "os/exec" + "runtime" "time" ) @@ -30,6 +31,9 @@ func (b *BashTool) Run(params map[string]interface{}) *ToolResult { errRes.ToolName = b.Definition().Name return errRes } + if denied := guardToolExecution("Bash", command, b.Definition().Name); denied != nil { + return denied + } timeoutMs, errRes := optionalInt(params, "timeout", 120000) if errRes != nil { errRes.ToolName = b.Definition().Name @@ -52,7 +56,26 @@ func (b *BashTool) Run(params map[string]interface{}) *ToolResult { ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutMs)*time.Millisecond) defer cancel() - cmd := exec.CommandContext(ctx, "bash", "-lc", command) + + var shell string + var shellArgs []string + switch runtime.GOOS { + case "linux", "darwin": + // Linux/macOS: 使用 bash + shell = "bash" + shellArgs = []string{"-lc", command} + case "windows": + // Windows: 使用 PowerShell + shell = "powershell" + shellArgs = []string{"-Command", command} + default: + shell = "bash" + shellArgs = []string{"-lc", command} + } + + // 使用动态选择的 shell 和参数创建命令 + shell, shellArgs = preferredShellCommand(runtime.GOOS, command, exec.LookPath, shell, shellArgs) + cmd := exec.CommandContext(ctx, shell, shellArgs...) cmd.Dir = workdir var stdoutBuf, stderrBuf bytes.Buffer cmd.Stdout = &stdoutBuf @@ -79,3 +102,30 @@ func (b *BashTool) Run(params map[string]interface{}) *ToolResult { } return result } + +type shellLookup func(string) (string, error) + +func preferredShellCommand(goos, command string, lookPath shellLookup, shell string, shellArgs []string) (string, []string) { + switch goos { + case "windows": + for _, candidate := range []struct { + name string + args []string + }{ + {name: "powershell", args: []string{"-Command", command}}, + {name: "pwsh", args: []string{"-Command", command}}, + {name: "cmd.exe", args: []string{"/C", command}}, + {name: "cmd", args: []string{"/C", command}}, + } { + if _, err := lookPath(candidate.name); err == nil { + return candidate.name, candidate.args + } + } + return "cmd", []string{"/C", command} + default: + if _, err := lookPath("bash"); err == nil { + return "bash", []string{"-lc", command} + } + return "sh", []string{"-c", command} + } +} diff --git a/internal/server/infra/tools/bash_security_test.go b/internal/server/infra/tools/bash_security_test.go new file mode 100644 index 00000000..21703a3e --- /dev/null +++ b/internal/server/infra/tools/bash_security_test.go @@ -0,0 +1,95 @@ +package tools + +import ( + "errors" + "strings" + "testing" + + "go-llm-demo/internal/server/domain" +) + +type mockSecurityChecker struct { + action domain.Action +} + +func (m mockSecurityChecker) Check(_ string, _ string) domain.Action { + return m.action +} + +func TestBashTool_Run_DeniedBySecurity(t *testing.T) { + SetSecurityChecker(mockSecurityChecker{action: domain.ActionDeny}) + defer SetSecurityChecker(nil) + + result := (&BashTool{}).Run(map[string]interface{}{ + "command": "echo hello", + }) + + if result == nil { + t.Fatal("result should not be nil") + } + if result.Success { + t.Fatal("expected bash execution to be denied") + } + if !strings.Contains(result.Error, "安全策略拒绝执行") { + t.Fatalf("unexpected error: %s", result.Error) + } + if result.Metadata["securityAction"] != string(domain.ActionDeny) { + t.Fatalf("unexpected security action: %#v", result.Metadata["securityAction"]) + } +} + +func TestBashTool_Run_AskBySecurity(t *testing.T) { + SetSecurityChecker(mockSecurityChecker{action: domain.ActionAsk}) + defer SetSecurityChecker(nil) + + result := (&BashTool{}).Run(map[string]interface{}{ + "command": "go build ./...", + }) + + if result == nil { + t.Fatal("result should not be nil") + } + if result.Success { + t.Fatal("expected bash execution to require confirmation") + } + if !strings.Contains(result.Error, "需要用户确认") { + t.Fatalf("unexpected error: %s", result.Error) + } + if result.Metadata["securityAction"] != string(domain.ActionAsk) { + t.Fatalf("unexpected security action: %#v", result.Metadata["securityAction"]) + } +} + +func TestPreferredShellCommandFallsBackToCmdOnWindows(t *testing.T) { + lookup := func(name string) (string, error) { + if name == "cmd.exe" { + return name, nil + } + return "", errors.New("not found") + } + + shell, args := preferredShellCommand("windows", "echo hello", lookup, "powershell", []string{"-Command", "echo hello"}) + if shell != "cmd.exe" { + t.Fatalf("expected cmd.exe fallback, got %q", shell) + } + if len(args) != 2 || args[0] != "/C" || args[1] != "echo hello" { + t.Fatalf("unexpected cmd.exe args: %#v", args) + } +} + +func TestPreferredShellCommandFallsBackToShWithoutBash(t *testing.T) { + lookup := func(name string) (string, error) { + if name == "sh" { + return name, nil + } + return "", errors.New("not found") + } + + shell, args := preferredShellCommand("linux", "echo hello", lookup, "bash", []string{"-lc", "echo hello"}) + if shell != "sh" { + t.Fatalf("expected sh fallback, got %q", shell) + } + if len(args) != 2 || args[0] != "-c" || args[1] != "echo hello" { + t.Fatalf("unexpected sh args: %#v", args) + } +} diff --git a/internal/server/infra/tools/security.go b/internal/server/infra/tools/security.go new file mode 100644 index 00000000..82df86b9 --- /dev/null +++ b/internal/server/infra/tools/security.go @@ -0,0 +1,122 @@ +package tools + +import ( + "fmt" + "log" + "strings" + "sync" + + "go-llm-demo/internal/server/domain" +) + +var ( + securityCheckerMu sync.RWMutex + securityChecker domain.SecurityChecker + + securityAskApprovalMu sync.Mutex + securityAskApprovals = map[string]int{} +) + +// SetSecurityChecker 设置工具执行前使用的安全检查器。 +// 传入 nil 表示关闭安全检查（默认行为）。 +func SetSecurityChecker(checker domain.SecurityChecker) { + securityCheckerMu.Lock() + securityChecker = checker + securityCheckerMu.Unlock() +} + +func getSecurityChecker() domain.SecurityChecker { + securityCheckerMu.RLock() + checker := securityChecker + securityCheckerMu.RUnlock() + return checker +} + +// ApproveSecurityAsk 为指定的安全询问发放一次性放行许可。 +// 该许可会在下一次匹配到同一 (toolType, target) 时被消费。 +func ApproveSecurityAsk(toolType, target string) { + key, err := securityApprovalKey(toolType, target) + if err != nil { + log.Printf("warning: failed to record security ask approval: %v", err) + return + } + securityAskApprovalMu.Lock() + securityAskApprovals[key]++ + securityAskApprovalMu.Unlock() +} + +func consumeSecurityAskApproval(toolType, target string) bool { + key, err := securityApprovalKey(toolType, target) + if err != nil { + log.Printf("warning: failed to consume security ask approval: %v", err) + return false + } + securityAskApprovalMu.Lock() + defer securityAskApprovalMu.Unlock() + count := securityAskApprovals[key] + if count <= 0 { + return false + } + if count == 1 { + delete(securityAskApprovals, key) + return true + } + securityAskApprovals[key] = count - 1 + return true +} + +func securityApprovalKey(toolType, target string) (string, error) { + normalizedType := strings.ToLower(strings.TrimSpace(toolType)) + normalizedTarget := strings.TrimSpace(target) + if normalizedType == "" || normalizedTarget == "" { + return "", fmt.Errorf("invalid security approval context: toolType=%q target=%q", toolType, target) + } + return normalizedType + "\x00" + normalizedTarget, nil +} + +func guardToolExecution(toolType, target, toolName string) *ToolResult { + checker := getSecurityChecker() + if checker == nil { + return nil + } + + action := checker.Check(toolType, target) + metadata := map[string]interface{}{ + "securityToolType": toolType, + "securityTarget": target, + "securityAction": string(action), + } + + switch action { + case domain.ActionAllow: + return nil + case domain.ActionDeny: + return &ToolResult{ + ToolName: toolName, + Success: false, + Error: fmt.Sprintf("安全策略拒绝执行 %s: %s", toolType, target), + Metadata: metadata, + } + case domain.ActionAsk: + if consumeSecurityAskApproval(toolType, target) { + return nil + } + return &ToolResult{ + ToolName: toolName, + Success: false, + Error: fmt.Sprintf("命中安全策略，执行 %s 前需要用户确认: %s", toolType, target), + Metadata: metadata, + } + default: + metadata["securityAction"] = string(domain.ActionAsk) + if consumeSecurityAskApproval(toolType, target) { + return nil + } + return &ToolResult{ + ToolName: toolName, + Success: false, + Error: fmt.Sprintf("安全策略返回未知动作(%s)，已按需确认处理: %s", action, target), + Metadata: metadata, + } + } +} diff --git a/internal/server/infra/tools/security_test.go b/internal/server/infra/tools/security_test.go new file mode 100644 index 00000000..99f008af --- /dev/null +++ b/internal/server/infra/tools/security_test.go @@ -0,0 +1,72 @@ +package tools + +import ( + "bytes" + "log" + "strings" + "testing" +) + +func TestApproveSecurityAskLogsInvalidContext(t *testing.T) { + resetSecurityApprovals() + + var buf bytes.Buffer + restoreLogs := captureToolLogs(&buf) + defer restoreLogs() + + ApproveSecurityAsk("", "target") + + if !strings.Contains(buf.String(), "invalid security approval context") { + t.Fatalf("expected invalid approval context to be logged, got %q", buf.String()) + } +} + +func TestConsumeSecurityAskApprovalLogsInvalidContext(t *testing.T) { + resetSecurityApprovals() + + var buf bytes.Buffer + restoreLogs := captureToolLogs(&buf) + defer restoreLogs() + + if consumeSecurityAskApproval("Bash", "") { + t.Fatal("expected invalid approval context to be rejected") + } + if !strings.Contains(buf.String(), "invalid security approval context") { + t.Fatalf("expected invalid approval context to be logged, got %q", buf.String()) + } +} + +func TestSecurityApprovalRoundTrip(t *testing.T) { + resetSecurityApprovals() + + ApproveSecurityAsk("Bash", "echo hello") + + if !consumeSecurityAskApproval("Bash", "echo hello") { + t.Fatal("expected approval to be consumed once") + } + if consumeSecurityAskApproval("Bash", "echo hello") { + t.Fatal("expected approval to be one-time") + } +} + +func resetSecurityApprovals() { + securityAskApprovalMu.Lock() + securityAskApprovals = map[string]int{} + securityAskApprovalMu.Unlock() +} + +func captureToolLogs(buf *bytes.Buffer) func() { + previousWriter := log.Writer() + previousFlags := log.Flags() + previousPrefix := log.Prefix() + + log.SetOutput(buf) + log.SetFlags(0) + log.SetPrefix("") + + return func() { + log.SetOutput(previousWriter) + log.SetFlags(previousFlags) + log.SetPrefix(previousPrefix) + } +} diff --git a/internal/server/service/security_service.go b/internal/server/service/security_service.go index 76fb1247..c1765188 100644 --- a/internal/server/service/security_service.go +++ b/internal/server/service/security_service.go @@ -10,25 +10,21 @@ import ( "github.com/bmatcuk/doublestar/v4" ) -type SecurityService interface { - Check(toolType string, target string) domain.Action -} - -// securityServiceImpl 是 SecurityService 的具体实现 -type securityServiceImpl struct { +// SecurityService provides security checks backed by configured rule sets. +type SecurityService struct { configRepo domain.SecurityConfigRepository blackList *domain.Config whiteList *domain.Config yellowList *domain.Config } -func NewSecurityService(configRepo domain.SecurityConfigRepository) SecurityService { - return &securityServiceImpl{ +func NewSecurityService(configRepo domain.SecurityConfigRepository) *SecurityService { + return &SecurityService{ configRepo: configRepo, } } -func (s *securityServiceImpl) Initialize(configDir string) error { +func (s *SecurityService) Initialize(configDir string) error { blackList, whiteList, yellowList, err := s.configRepo.LoadAll(configDir) if err != nil { return err @@ -39,7 +35,7 @@ func (s *securityServiceImpl) Initialize(configDir string) error { return nil } -func (s *securityServiceImpl) Check(toolType string, target string) domain.Action { +func (s *SecurityService) Check(toolType string, target string) domain.Action { normalizedTarget := target // 安全增强：对路径类操作进行规范化处理 @@ -74,7 +70,7 @@ func (s *securityServiceImpl) Check(toolType string, target string) domain.Actio return domain.ActionAsk } -func (s *securityServiceImpl) matchesList(config *domain.Config, toolType, target string) bool { +func (s *SecurityService) matchesList(config *domain.Config, toolType, target string) bool { if config == nil { return false } diff --git a/internal/server/service/security_service_test.go b/internal/server/service/security_service_test.go index 41f4381e..701eeaed 100644 --- a/internal/server/service/security_service_test.go +++ b/internal/server/service/security_service_test.go @@ -61,8 +61,7 @@ func TestSecurityService_Check(t *testing.T) { // 2. 初始化服务，并注入替身仓储 svc := NewSecurityService(mockRepo) - impl := svc.(*securityServiceImpl) - if err := impl.Initialize(""); err != nil { + if err := svc.Initialize(""); err != nil { t.Fatalf("初始化失败: %v", err) } diff --git a/internal/tui/bootstrap/setup.go b/internal/tui/bootstrap/setup.go index 994eafc0..64a22e7b 100644 --- a/internal/tui/bootstrap/setup.go +++ b/internal/tui/bootstrap/setup.go @@ -5,6 +5,7 @@ import ( "context" "errors" "fmt" + "path/filepath" "strings" "go-llm-demo/configs" @@ -22,6 +23,7 @@ const ( var ( resolveWorkspaceRoot = services.ResolveWorkspaceRoot setWorkspaceRoot = services.SetWorkspaceRoot + initializeSecurity = services.InitializeSecurity ensureConfigFile = configs.EnsureConfigFile validateChatAPIKey = services.ValidateChatAPIKey writeAppConfig = configs.WriteAppConfig @@ -35,6 +37,9 @@ func PrepareWorkspace(workspaceFlag string) (string, error) { if err := setWorkspaceRoot(workspaceRoot); err != nil { return "", err } + if err := initializeSecurity(filepath.Join(workspaceRoot, "configs", "security")); err != nil { + return "", err + } return workspaceRoot, nil } diff --git a/internal/tui/bootstrap/setup_test.go b/internal/tui/bootstrap/setup_test.go index a827fca6..0c48f382 100644 --- a/internal/tui/bootstrap/setup_test.go +++ b/internal/tui/bootstrap/setup_test.go @@ -5,6 +5,7 @@ import ( "context" "errors" "io" + "path/filepath" "strings" "testing" @@ -23,6 +24,7 @@ func restoreSetupGlobals(t *testing.T) { origResolveWorkspaceRoot := resolveWorkspaceRoot origSetWorkspaceRoot := setWorkspaceRoot + origInitializeSecurity := initializeSecurity origEnsureConfigFile := ensureConfigFile origValidateChatAPIKey := validateChatAPIKey origWriteAppConfig := writeAppConfig @@ -31,6 +33,7 @@ func restoreSetupGlobals(t *testing.T) { t.Cleanup(func() { resolveWorkspaceRoot = origResolveWorkspaceRoot setWorkspaceRoot = origSetWorkspaceRoot + initializeSecurity = origInitializeSecurity ensureConfigFile = origEnsureConfigFile validateChatAPIKey = origValidateChatAPIKey writeAppConfig = origWriteAppConfig @@ -113,6 +116,7 @@ func TestPrepareWorkspaceResolvesAndSetsWorkspaceRoot(t *testing.T) { restoreSetupGlobals(t) var setRoot string + var initializedConfigDir string resolveWorkspaceRoot = func(workspaceFlag string) (string, error) { if workspaceFlag != "./workspace" { t.Fatalf("expected workspace flag to flow through, got %q", workspaceFlag) @@ -123,6 +127,10 @@ func TestPrepareWorkspaceResolvesAndSetsWorkspaceRoot(t *testing.T) { setRoot = root return nil } + initializeSecurity = func(configDir string) error { + initializedConfigDir = configDir + return nil + } root, err := PrepareWorkspace("./workspace") if err != nil { @@ -134,6 +142,9 @@ func TestPrepareWorkspaceResolvesAndSetsWorkspaceRoot(t *testing.T) { if setRoot != root { t.Fatalf("expected SetWorkspaceRoot to receive %q, got %q", root, setRoot) } + if initializedConfigDir != filepath.Join("D:/neo-code/workspace", "configs", "security") { + t.Fatalf("expected security config dir to be initialized, got %q", initializedConfigDir) + } } func TestPrepareWorkspaceReturnsSetWorkspaceRootError(t *testing.T) { @@ -141,6 +152,7 @@ func TestPrepareWorkspaceReturnsSetWorkspaceRootError(t *testing.T) { resolveWorkspaceRoot = func(string) (string, error) { return "D:/neo-code/workspace", nil } setWorkspaceRoot = func(string) error { return errors.New("set failed") } + initializeSecurity = func(string) error { return nil } _, err := PrepareWorkspace("./workspace") if err == nil || !strings.Contains(err.Error(), "set failed") { @@ -159,6 +171,19 @@ func TestPrepareWorkspaceReturnsResolveError(t *testing.T) { } } +func TestPrepareWorkspaceReturnsInitializeSecurityError(t *testing.T) { + restoreSetupGlobals(t) + + resolveWorkspaceRoot = func(string) (string, error) { return "D:/neo-code/workspace", nil } + setWorkspaceRoot = func(string) error { return nil } + initializeSecurity = func(string) error { return errors.New("security failed") } + + _, err := PrepareWorkspace("./workspace") + if err == nil || !strings.Contains(err.Error(), "security failed") { + t.Fatalf("expected initialize security error, got %v", err) + } +} + func TestReadInteractiveLineReturnsEOF(t *testing.T) { scanner := bufio.NewScanner(strings.NewReader("")) diff --git a/internal/tui/core/msg.go b/internal/tui/core/msg.go index acbf0e77..2377d69c 100644 --- a/internal/tui/core/msg.go +++ b/internal/tui/core/msg.go @@ -20,6 +20,7 @@ func (StreamErrorMsg) isMsg() {} type ToolResultMsg struct { Result *services.ToolResult + Call services.ToolCall } func (ToolResultMsg) isMsg() {} diff --git a/internal/tui/core/update.go b/internal/tui/core/update.go index 7deaee7c..83aecc5f 100644 --- a/internal/tui/core/update.go +++ b/internal/tui/core/update.go @@ -112,7 +112,7 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { mu.Unlock() return ToolErrorMsg{Err: fmt.Errorf("工具执行失败: 空返回")} } - return ToolResultMsg{Result: result} + return ToolResultMsg{Result: result, Call: call} } } } @@ -170,6 +170,21 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.chat.ToolExecuting = false mu.Unlock() // 将结构化工具上下文添加为系统消息，然后重新获取AI响应 + if toolType, target, ok := isSecurityAskResult(msg.Result); ok { + mu := m.mutex() + mu.Lock() + m.chat.PendingApproval = &state.PendingApproval{ + Call: msg.Call, + ToolType: toolType, + Target: target, + } + pending := m.chat.PendingApproval + mu.Unlock() + + m.AddMessage("assistant", formatPendingApprovalMessage(pending)) + m.refreshViewport() + return m, nil + } m.AddMessage("system", formatToolContextMessage(msg.Result)) m.AddMessage("assistant", "") m.chat.Generating = true @@ -281,6 +296,11 @@ func (m *Model) handleSubmit() (tea.Model, tea.Cmd) { return *m, nil } + if m.chat.PendingApproval != nil { + m.AddMessage("assistant", "A security approval is pending. Use /y to allow once or /n to reject before sending a new message.") + return *m, nil + } + m.AddMessage("user", input) m.AddMessage("assistant", "") // 在请求发出前先裁剪原始消息，避免 UI 历史无限扩张并影响短期上下文质量。 @@ -312,6 +332,71 @@ func (m *Model) handleCommand(input string) (tea.Model, tea.Cmd) { switch cmd { case "/help": m.ui.Mode = state.ModeHelp + case "/y": + if len(args) > 0 { + m.AddMessage("assistant", "Usage: /y") + return *m, nil + } + if m.chat.PendingApproval == nil { + m.AddMessage("assistant", "There is no pending security approval.") + return *m, nil + } + if m.chat.ToolExecuting { + m.AddMessage("assistant", "Another tool is still running. Please retry /y after it finishes.") + return *m, nil + } + + pending := *m.chat.PendingApproval + m.chat.PendingApproval = nil + if strings.TrimSpace(pending.Call.Tool) == "" { + m.AddMessage("assistant", "The pending tool request is incomplete and cannot be executed.") + return *m, nil + } + + m.AddMessage("assistant", fmt.Sprintf("Approved. Running tool %s.", pending.Call.Tool)) + m.AddMessage("system", formatToolStatusMessage(pending.Call.Tool, pending.Call.Params)) + + mu := m.mutex() + mu.Lock() + if m.chat.ToolExecuting { + m.chat.PendingApproval = &pending + mu.Unlock() + return *m, nil + } + m.chat.ToolExecuting = true + mu.Unlock() + + m.refreshViewport() + return *m, func() tea.Msg { + services.ApproveSecurityAsk(pending.ToolType, pending.Target) + result := executeToolCall(pending.Call) + if result == nil { + mu := m.mutex() + mu.Lock() + m.chat.ToolExecuting = false + mu.Unlock() + return ToolErrorMsg{Err: fmt.Errorf("tool execution failed: empty result")} + } + return ToolResultMsg{Result: result, Call: pending.Call} + } + case "/n": + if len(args) > 0 { + m.AddMessage("assistant", "Usage: /n") + return *m, nil + } + if m.chat.PendingApproval == nil { + m.AddMessage("assistant", "There is no pending security approval.") + return *m, nil + } + + pending := *m.chat.PendingApproval + m.chat.PendingApproval = nil + toolName := strings.TrimSpace(pending.Call.Tool) + if toolName == "" { + toolName = "unknown" + } + m.AddMessage("assistant", fmt.Sprintf("Rejected tool %s for target %s.", toolName, pending.Target)) + return *m, nil case "/exit", "/quit", "/q": return *m, tea.Quit case "/apikey": @@ -492,7 +577,7 @@ func (m *Model) handleCommand(input string) (tea.Model, tea.Cmd) { func isAPIKeyRecoveryCommand(cmd string) bool { switch cmd { - case "/apikey", "/provider", "/help", "/switch", "/pwd", "/workspace", "/exit", "/quit", "/q": + case "/apikey", "/provider", "/help", "/switch", "/pwd", "/workspace", "/y", "/n", "/exit", "/quit", "/q": return true default: return false @@ -640,6 +725,33 @@ func formatToolStatusMessage(toolName string, params map[string]interface{}) str return fmt.Sprintf("%s tool=%s%s", toolStatusPrefix, strings.TrimSpace(toolName), detail) } +func isSecurityAskResult(result *services.ToolResult) (string, string, bool) { + if result == nil || result.Success || result.Metadata == nil { + return "", "", false + } + action, _ := result.Metadata["securityAction"].(string) + if strings.TrimSpace(strings.ToLower(action)) != "ask" { + return "", "", false + } + toolType, _ := result.Metadata["securityToolType"].(string) + target, _ := result.Metadata["securityTarget"].(string) + if strings.TrimSpace(toolType) == "" || strings.TrimSpace(target) == "" { + return "", "", false + } + return strings.TrimSpace(toolType), strings.TrimSpace(target), true +} + +func formatPendingApprovalMessage(pending *state.PendingApproval) string { + if pending == nil { + return "Security approval is required. Use /y to allow once or /n to reject." + } + toolName := strings.TrimSpace(pending.Call.Tool) + if toolName == "" { + toolName = "unknown" + } + return fmt.Sprintf("Security approval required for %s.\nTarget: %s\nUse /y to allow once, or /n to reject.", toolName, pending.Target) +} + func formatToolContextMessage(result *services.ToolResult) string { if result == nil { return toolContextPrefix + "\n" + "tool=unknown\n" + "success=false\n" + "error:\n工具返回为空" diff --git a/internal/tui/core/update_test.go b/internal/tui/core/update_test.go index 9215d45e..6ed11b68 100644 --- a/internal/tui/core/update_test.go +++ b/internal/tui/core/update_test.go @@ -1309,3 +1309,86 @@ func TestWorkspaceCommandShowsWorkspaceRoot(t *testing.T) { t.Fatalf("expected workspace path in response, got %q", got.chat.Messages[0].Content) } } + +func TestApproveCommandWhileToolExecutingKeepsPendingApproval(t *testing.T) { + m := Model{ + chat: state.ChatState{ + ToolExecuting: true, + PendingApproval: &state.PendingApproval{ + Call: services.ToolCall{ + Tool: "bash", + Params: map[string]interface{}{ + "command": "echo hello", + }, + }, + ToolType: "Bash", + Target: "echo hello", + }, + }, + } + + updated, cmd := m.handleCommand("/y") + if cmd != nil { + t.Fatal("expected no tool execution command while another tool is running") + } + + got := updated.(Model) + if got.chat.PendingApproval == nil { + t.Fatal("expected pending approval to be preserved") + } + if got.chat.PendingApproval.Call.Tool != "bash" { + t.Fatalf("expected pending tool to stay intact, got %+v", got.chat.PendingApproval.Call) + } + if len(got.chat.Messages) != 1 { + t.Fatalf("expected a single assistant warning, got %d", len(got.chat.Messages)) + } + if !strings.Contains(got.chat.Messages[0].Content, "/y") { + t.Fatalf("expected warning message to mention /y, got %q", got.chat.Messages[0].Content) + } +} + +func TestToolResultMsgSecurityAskStoresPendingApproval(t *testing.T) { + client := &fakeChatClient{} + m := newTestModel(t, client) + m.chat.ToolExecuting = true + + result := &services.ToolResult{ + ToolName: "bash", + Success: false, + Metadata: map[string]interface{}{ + "securityAction": "ask", + "securityToolType": "Bash", + "securityTarget": "echo hello", + }, + } + + updated, cmd := m.Update(ToolResultMsg{ + Result: result, + Call: services.ToolCall{ + Tool: "bash", + Params: map[string]interface{}{ + "command": "echo hello", + }, + }, + }) + if cmd != nil { + t.Fatal("expected no follow-up command while waiting for approval") + } + + got := updated.(Model) + if got.chat.ToolExecuting { + t.Fatal("expected tool executing flag to be cleared") + } + if got.chat.PendingApproval == nil { + t.Fatal("expected pending approval to be recorded") + } + if got.chat.PendingApproval.Target != "echo hello" { + t.Fatalf("unexpected pending approval target %q", got.chat.PendingApproval.Target) + } + if len(got.chat.Messages) != 1 { + t.Fatalf("expected one approval prompt message, got %d", len(got.chat.Messages)) + } + if !strings.Contains(got.chat.Messages[0].Content, "/y") { + t.Fatalf("expected approval prompt to mention /y, got %q", got.chat.Messages[0].Content) + } +} diff --git a/internal/tui/services/runtime_services.go b/internal/tui/services/runtime_services.go index 864b95f6..cef89191 100644 --- a/internal/tui/services/runtime_services.go +++ b/internal/tui/services/runtime_services.go @@ -6,7 +6,9 @@ import ( "go-llm-demo/configs" "go-llm-demo/internal/server/domain" serverprovider "go-llm-demo/internal/server/infra/provider" + serverrepo "go-llm-demo/internal/server/infra/repository" servertools "go-llm-demo/internal/server/infra/tools" + serverservice "go-llm-demo/internal/server/service" ) type ToolCall = domain.ToolCall @@ -45,6 +47,20 @@ func ExecuteToolCall(call ToolCall) *ToolResult { return servertools.GlobalRegistry.Execute(call) } +func ApproveSecurityAsk(toolType, target string) { + servertools.ApproveSecurityAsk(toolType, target) +} + +func InitializeSecurity(configDir string) error { + securityRepo := serverrepo.NewSecurityConfigRepository() + securitySvc := serverservice.NewSecurityService(securityRepo) + if err := securitySvc.Initialize(configDir); err != nil { + return err + } + servertools.SetSecurityChecker(securitySvc) + return nil +} + func ValidateChatAPIKey(ctx context.Context, cfg *configs.AppConfiguration) error { return serverprovider.ValidateChatAPIKey(ctx, cfg) } diff --git a/internal/tui/state/chat_state.go b/internal/tui/state/chat_state.go index cdd146f9..213f5f5d 100644 --- a/internal/tui/state/chat_state.go +++ b/internal/tui/state/chat_state.go @@ -13,16 +13,23 @@ type Message struct { Streaming bool } +type PendingApproval struct { + Call services.ToolCall + ToolType string + Target string +} + type ChatState struct { - Messages []Message - HistoryTurns int - Generating bool - ActiveModel string - MemoryStats services.MemoryStats - CommandHistory []string - CmdHistIndex int - WorkspaceRoot string - ToolExecuting bool - APIKeyReady bool - ConfigPath string + Messages []Message + HistoryTurns int + Generating bool + ActiveModel string + MemoryStats services.MemoryStats + CommandHistory []string + CmdHistIndex int + WorkspaceRoot string + ToolExecuting bool + PendingApproval *PendingApproval + APIKeyReady bool + ConfigPath string }