diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index d87241ea..bd4ea261 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -99,7 +99,7 @@ func buildToolRegistry(cfg config.Config) *tools.Registry { } func buildToolManager(registry *tools.Registry) (tools.Manager, error) { - engine, err := security.NewStaticGateway(security.DecisionAllow, nil) + engine, err := security.NewRecommendedPolicyEngine() if err != nil { return nil, err } diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index bb204cbd..beefb244 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -132,16 +132,13 @@ func TestBuildToolManagerWrapsRegistry(t *testing.T) { t.Fatalf("expected 1 spec, got %+v", specs) } - result, execErr := manager.Execute(context.Background(), tools.ToolCallInput{ + _, execErr := manager.Execute(context.Background(), tools.ToolCallInput{ Name: "bash", Arguments: []byte(`{"command":"echo hi"}`), Workdir: workdir, }) - if execErr != nil { - t.Fatalf("Execute() error = %v", execErr) - } - if result.Content != "ok" { - t.Fatalf("expected ok result, got %+v", result) + if execErr == nil { + t.Fatalf("expected bash to require approval by default policy") } _, execErr = manager.Execute(context.Background(), tools.ToolCallInput{ @@ -154,6 +151,29 @@ func TestBuildToolManagerWrapsRegistry(t *testing.T) { } } +func TestBuildToolManagerAllowsWebfetchWhitelist(t *testing.T) { + t.Parallel() + + registry := tools.NewRegistry() + registry.Register(stubToolForBootstrap{name: "webfetch", content: "ok"}) + manager, err := buildToolManager(registry) + if err != nil { + t.Fatalf("buildToolManager() error = %v", err) + } + + result, execErr := manager.Execute(context.Background(), tools.ToolCallInput{ + Name: "webfetch", + Arguments: []byte(`{"url":"https://github.com/1024XEngineer/neo-code"}`), + Workdir: t.TempDir(), + }) + if execErr != nil { + t.Fatalf("expected whitelist webfetch allow, got %v", execErr) + } + if result.Content != "ok" { + t.Fatalf("expected ok result, got %+v", result) + } +} + func TestEnsureConsoleUTF8SetsOutputThenInput(t *testing.T) { originalOutput := setConsoleOutputCodePage originalInput := setConsoleInputCodePage diff --git a/internal/runtime/events.go b/internal/runtime/events.go index 594e8f5e..704ff14a 100644 --- a/internal/runtime/events.go +++ b/internal/runtime/events.go @@ -15,7 +15,10 @@ type RuntimeEvent struct { // PermissionRequestPayload 描述一次需要审批的权限请求上下文。 type PermissionRequestPayload struct { + RequestID string + ToolCallID string ToolName string + ToolCategory string ActionType string Operation string TargetType string @@ -28,7 +31,10 @@ type PermissionRequestPayload struct { // PermissionResolvedPayload 描述权限请求被运行时处理后的最终状态。 type PermissionResolvedPayload struct { + RequestID string + ToolCallID string ToolName string + ToolCategory string ActionType string Operation string TargetType string diff --git a/internal/runtime/permission.go b/internal/runtime/permission.go new file mode 100644 index 00000000..6782e9be --- /dev/null +++ b/internal/runtime/permission.go @@ -0,0 +1,291 @@ +package runtime + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "time" + + "neo-code/internal/provider" + "neo-code/internal/security" + "neo-code/internal/tools" +) + +// PermissionResolutionInput 描述一次来自界面的权限审批决定。 +type PermissionResolutionInput struct { + RequestID string + Decision PermissionResolutionDecision +} + +// PermissionResolutionDecision 表示用户在权限提示中的最终选择。 +type PermissionResolutionDecision string + +const ( + PermissionResolutionAllowOnce PermissionResolutionDecision = "allow_once" + PermissionResolutionAllowSession PermissionResolutionDecision = "allow_session" + PermissionResolutionReject PermissionResolutionDecision = "reject" +) + +type permissionExecutionInput struct { + RunID string + SessionID string + Call provider.ToolCall + Workdir string + ToolTimeout time.Duration +} + +type pendingPermissionRequest struct { + RequestID string + RunID string + SessionID string + Call provider.ToolCall + Action security.Action + ResultCh chan PermissionResolutionDecision + Submitted bool +} + +var runtimePendingPermissions = struct { + mu sync.Mutex + nextID uint64 + byRun map[*Service]*pendingPermissionRequest +}{ + byRun: make(map[*Service]*pendingPermissionRequest), +} + +// ResolvePermission 接收 UI 的审批决定，并唤醒 runtime 中等待的工具调用。 +func (s *Service) ResolvePermission(ctx context.Context, input PermissionResolutionInput) error { + requestID := strings.TrimSpace(input.RequestID) + if requestID == "" { + return errors.New("runtime: permission request id is empty") + } + decision := normalizePermissionResolutionDecision(input.Decision) + if decision == "" { + return fmt.Errorf("runtime: unsupported permission decision %q", input.Decision) + } + if err := ctx.Err(); err != nil { + return err + } + + runtimePendingPermissions.mu.Lock() + pending := runtimePendingPermissions.byRun[s] + if pending == nil || pending.RequestID != requestID { + runtimePendingPermissions.mu.Unlock() + return fmt.Errorf("runtime: permission request %q not found", requestID) + } + // Submitted 标记用于避免重复提交同一个 request 导致阻塞。 + if pending.Submitted { + runtimePendingPermissions.mu.Unlock() + return nil + } + pending.Submitted = true + resultCh := pending.ResultCh + runtimePendingPermissions.mu.Unlock() + + // 非阻塞提交，避免 UI 重复触发导致写满 channel 后长时间卡住。 + select { + case resultCh <- decision: + return nil + default: + return nil + } +} + +// executeToolCallWithPermission 执行工具调用并处理 ask 决策的显式审批闭环。 +func (s *Service) executeToolCallWithPermission(ctx context.Context, input permissionExecutionInput) (tools.ToolResult, error) { + callInput := tools.ToolCallInput{ + ID: input.Call.ID, + Name: input.Call.Name, + Arguments: []byte(input.Call.Arguments), + Workdir: input.Workdir, + SessionID: input.SessionID, + EmitChunk: func(chunk []byte) { + s.emit(ctx, EventToolChunk, input.RunID, input.SessionID, string(chunk)) + }, + } + + runCtx, cancel := context.WithTimeout(ctx, input.ToolTimeout) + result, execErr := s.toolManager.Execute(runCtx, callInput) + cancel() + if execErr == nil { + return result, nil + } + + var permissionErr *tools.PermissionDecisionError + if !errors.As(execErr, &permissionErr) { + return result, execErr + } + if !strings.EqualFold(permissionErr.Decision(), string(security.DecisionAsk)) { + return result, execErr + } + + decision, scope, requestID, err := s.awaitPermissionDecision(ctx, input, permissionErr) + if err != nil { + return result, err + } + + if decision == PermissionResolutionReject { + if scope == tools.SessionPermissionScopeReject { + if rememberErr := s.toolManager.RememberSessionDecision(input.SessionID, permissionErr.Action(), scope); rememberErr != nil { + return tools.ToolResult{}, rememberErr + } + } + s.emit(ctx, EventPermissionResolved, input.RunID, input.SessionID, PermissionResolvedPayload{ + RequestID: requestID, + ToolCallID: input.Call.ID, + ToolName: input.Call.Name, + ToolCategory: permissionToolCategory(permissionErr.Action()), + ActionType: string(permissionErr.Action().Type), + Operation: permissionErr.Action().Payload.Operation, + TargetType: string(permissionErr.Action().Payload.TargetType), + Target: permissionErr.Action().Payload.Target, + Decision: "deny", + Reason: "permission rejected by user", + RuleID: permissionErr.RuleID(), + RememberScope: string(scope), + ResolvedAs: "rejected", + }) + return tools.ToolResult{ + ToolCallID: input.Call.ID, + Name: input.Call.Name, + Content: "tool error\ntool: " + input.Call.Name + "\nreason: permission rejected by user", + IsError: true, + }, errors.New("tools: permission rejected by user") + } + + if rememberErr := s.toolManager.RememberSessionDecision(input.SessionID, permissionErr.Action(), scope); rememberErr != nil { + return tools.ToolResult{}, rememberErr + } + s.emit(ctx, EventPermissionResolved, input.RunID, input.SessionID, PermissionResolvedPayload{ + RequestID: requestID, + ToolCallID: input.Call.ID, + ToolName: input.Call.Name, + ToolCategory: permissionToolCategory(permissionErr.Action()), + ActionType: string(permissionErr.Action().Type), + Operation: permissionErr.Action().Payload.Operation, + TargetType: string(permissionErr.Action().Payload.TargetType), + Target: permissionErr.Action().Payload.Target, + Decision: "allow", + Reason: "permission approved by user", + RuleID: permissionErr.RuleID(), + RememberScope: string(scope), + ResolvedAs: "approved", + }) + + retryCtx, retryCancel := context.WithTimeout(ctx, input.ToolTimeout) + retryResult, retryErr := s.toolManager.Execute(retryCtx, callInput) + retryCancel() + return retryResult, retryErr +} + +// awaitPermissionDecision 发送 permission_request 事件，并等待 UI 回传审批结果。 +func (s *Service) awaitPermissionDecision( + ctx context.Context, + input permissionExecutionInput, + permissionErr *tools.PermissionDecisionError, +) (PermissionResolutionDecision, tools.SessionPermissionScope, string, error) { + request := registerPendingPermission(s, input, permissionErr.Action()) + defer clearPendingPermission(s, request.RequestID) + + s.emit(ctx, EventPermissionRequest, input.RunID, input.SessionID, PermissionRequestPayload{ + RequestID: request.RequestID, + ToolCallID: input.Call.ID, + ToolName: input.Call.Name, + ToolCategory: permissionToolCategory(permissionErr.Action()), + ActionType: string(permissionErr.Action().Type), + Operation: permissionErr.Action().Payload.Operation, + TargetType: string(permissionErr.Action().Payload.TargetType), + Target: permissionErr.Action().Payload.Target, + Decision: permissionErr.Decision(), + Reason: permissionErr.Reason(), + RuleID: permissionErr.RuleID(), + }) + + select { + case <-ctx.Done(): + return "", "", request.RequestID, ctx.Err() + case decision := <-request.ResultCh: + scope, err := rememberScopeFromDecision(decision) + if err != nil { + return "", "", request.RequestID, err + } + return decision, scope, request.RequestID, nil + } +} + +// registerPendingPermission 注册待审批请求并返回 request id。 +func registerPendingPermission(s *Service, input permissionExecutionInput, action security.Action) pendingPermissionRequest { + runtimePendingPermissions.mu.Lock() + defer runtimePendingPermissions.mu.Unlock() + + runtimePendingPermissions.nextID++ + request := pendingPermissionRequest{ + RequestID: fmt.Sprintf("perm-%d", runtimePendingPermissions.nextID), + RunID: input.RunID, + SessionID: input.SessionID, + Call: input.Call, + Action: action, + ResultCh: make(chan PermissionResolutionDecision, 1), + } + runtimePendingPermissions.byRun[s] = &request + return request +} + +// clearPendingPermission 清理待审批请求，避免跨请求误用。 +func clearPendingPermission(s *Service, requestID string) { + runtimePendingPermissions.mu.Lock() + defer runtimePendingPermissions.mu.Unlock() + + current := runtimePendingPermissions.byRun[s] + if current != nil && current.RequestID == requestID { + delete(runtimePendingPermissions.byRun, s) + } +} + +// normalizePermissionResolutionDecision 将审批决定归一化为受支持枚举值。 +func normalizePermissionResolutionDecision(decision PermissionResolutionDecision) PermissionResolutionDecision { + switch strings.ToLower(strings.TrimSpace(string(decision))) { + case "allow_once", "once", "y", "yes": + return PermissionResolutionAllowOnce + case "allow_session", "always", "always_session", "a": + return PermissionResolutionAllowSession + case "reject", "deny", "n", "no", "r": + return PermissionResolutionReject + default: + return "" + } +} + +// rememberScopeFromDecision 将审批决定映射到工具层权限记忆范围。 +func rememberScopeFromDecision(decision PermissionResolutionDecision) (tools.SessionPermissionScope, error) { + switch decision { + case PermissionResolutionAllowOnce: + return tools.SessionPermissionScopeOnce, nil + case PermissionResolutionAllowSession: + return tools.SessionPermissionScopeAlways, nil + case PermissionResolutionReject: + return tools.SessionPermissionScopeReject, nil + default: + return "", fmt.Errorf("runtime: unsupported permission decision %q", decision) + } +} + +// permissionToolCategory 将 action 归一化为工具类别标签，供审批展示和记忆使用。 +func permissionToolCategory(action security.Action) string { + resource := strings.ToLower(strings.TrimSpace(action.Payload.Resource)) + switch action.Type { + case security.ActionTypeRead: + if strings.HasPrefix(resource, "filesystem_") { + return "filesystem_read" + } + case security.ActionTypeWrite: + if strings.HasPrefix(resource, "filesystem_") { + return "filesystem_write" + } + } + if resource != "" { + return resource + } + return strings.ToLower(strings.TrimSpace(action.Payload.ToolName)) +} diff --git a/internal/runtime/permission_test.go b/internal/runtime/permission_test.go new file mode 100644 index 00000000..a32ce0fd --- /dev/null +++ b/internal/runtime/permission_test.go @@ -0,0 +1,337 @@ +package runtime + +import ( + "context" + "errors" + "testing" + "time" + + "neo-code/internal/provider" + "neo-code/internal/security" + "neo-code/internal/tools" +) + +func TestResolvePermissionValidation(t *testing.T) { + t.Parallel() + + service := NewWithFactory( + newRuntimeConfigManager(t), + &stubToolManager{}, + newMemoryStore(), + &scriptedProviderFactory{provider: &scriptedProvider{}}, + nil, + ) + + if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{}); err == nil { + t.Fatalf("expected empty request id error") + } + if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ + RequestID: "perm-1", + Decision: PermissionResolutionDecision("invalid"), + }); err == nil { + t.Fatalf("expected invalid decision error") + } + if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ + RequestID: "perm-not-found", + Decision: PermissionResolutionAllowOnce, + }); err == nil { + t.Fatalf("expected request not found error") + } +} + +func TestResolvePermissionSuccess(t *testing.T) { + t.Parallel() + + service := NewWithFactory( + newRuntimeConfigManager(t), + &stubToolManager{}, + newMemoryStore(), + &scriptedProviderFactory{provider: &scriptedProvider{}}, + nil, + ) + + request := registerPendingPermission(service, permissionExecutionInput{ + RunID: "run-permission", + SessionID: "session-permission", + Call: provider.ToolCall{ + ID: "call-1", + Name: "webfetch", + }, + }, security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + Operation: "fetch", + TargetType: security.TargetTypeURL, + Target: "https://example.com", + }, + }) + defer clearPendingPermission(service, request.RequestID) + + errCh := make(chan error, 1) + go func() { + errCh <- service.ResolvePermission(context.Background(), PermissionResolutionInput{ + RequestID: request.RequestID, + Decision: PermissionResolutionAllowSession, + }) + }() + + select { + case resolved := <-request.ResultCh: + if resolved != PermissionResolutionAllowSession { + t.Fatalf("expected allow session decision, got %q", resolved) + } + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting permission resolution") + } + + if err := <-errCh; err != nil { + t.Fatalf("ResolvePermission() error = %v", err) + } +} + +func TestResolvePermissionDuplicateSubmissionIsNonBlocking(t *testing.T) { + t.Parallel() + + service := NewWithFactory( + newRuntimeConfigManager(t), + &stubToolManager{}, + newMemoryStore(), + &scriptedProviderFactory{provider: &scriptedProvider{}}, + nil, + ) + + request := registerPendingPermission(service, permissionExecutionInput{ + RunID: "run-permission-dup", + SessionID: "session-permission-dup", + Call: provider.ToolCall{ + ID: "call-dup", + Name: "webfetch", + }, + }, security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + }, + }) + defer clearPendingPermission(service, request.RequestID) + + if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ + RequestID: request.RequestID, + Decision: PermissionResolutionAllowOnce, + }); err != nil { + t.Fatalf("first ResolvePermission() error = %v", err) + } + + secondDone := make(chan error, 1) + go func() { + secondDone <- service.ResolvePermission(context.Background(), PermissionResolutionInput{ + RequestID: request.RequestID, + Decision: PermissionResolutionAllowSession, + }) + }() + + select { + case err := <-secondDone: + if err != nil { + t.Fatalf("second ResolvePermission() error = %v", err) + } + case <-time.After(500 * time.Millisecond): + t.Fatalf("second ResolvePermission() should not block") + } +} + +func TestServiceRunPermissionRejectFlow(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + registry := tools.NewRegistry() + tool := &stubTool{name: "webfetch", content: "should-not-run"} + registry.Register(tool) + + engine, err := security.NewStaticGateway(security.DecisionAllow, []security.Rule{ + { + ID: "ask-webfetch", + Type: security.ActionTypeRead, + Resource: "webfetch", + Decision: security.DecisionAsk, + Reason: "requires approval", + }, + }) + if err != nil { + t.Fatalf("new static gateway: %v", err) + } + toolManager, err := tools.NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new tool manager: %v", err) + } + + scripted := &scriptedProvider{ + responses: []provider.ChatResponse{ + { + Message: provider.Message{ + Role: "assistant", + ToolCalls: []provider.ToolCall{ + {ID: "call-ask-reject", Name: "webfetch", Arguments: `{"url":"https://example.com/private"}`}, + }, + }, + FinishReason: "tool_calls", + }, + { + Message: provider.Message{Role: "assistant", Content: "done"}, + FinishReason: "stop", + }, + }, + } + + service := NewWithFactory(manager, toolManager, store, &scriptedProviderFactory{provider: scripted}, nil) + runErrCh := make(chan error, 1) + go func() { + runErrCh <- service.Run(context.Background(), UserInput{RunID: "run-permission-reject", Content: "fetch private"}) + }() + + var requestPayload PermissionRequestPayload +waitRequest: + for { + select { + case <-time.After(3 * time.Second): + t.Fatalf("timed out waiting permission request") + case event := <-service.Events(): + if event.Type != EventPermissionRequest { + continue + } + payload, ok := event.Payload.(PermissionRequestPayload) + if !ok { + t.Fatalf("expected permission request payload, got %#v", event.Payload) + } + requestPayload = payload + break waitRequest + } + } + + if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ + RequestID: requestPayload.RequestID, + Decision: PermissionResolutionReject, + }); err != nil { + t.Fatalf("ResolvePermission() error = %v", err) + } + if err := <-runErrCh; err != nil { + t.Fatalf("Run() error = %v", err) + } + + if tool.callCount != 0 { + t.Fatalf("expected tool not executed after reject, got %d", tool.callCount) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventPermissionResolved, EventToolResult, EventAgentDone}) + + found := false + for _, event := range events { + if event.Type != EventPermissionResolved { + continue + } + payload, ok := event.Payload.(PermissionResolvedPayload) + if !ok { + t.Fatalf("expected permission resolved payload, got %#v", event.Payload) + } + if payload.Decision == "deny" && payload.ResolvedAs == "rejected" { + found = true + } + } + if !found { + t.Fatalf("expected user reject resolved payload") + } +} + +func TestPermissionHelpers(t *testing.T) { + t.Parallel() + + if got := normalizePermissionResolutionDecision(PermissionResolutionDecision("Y")); got != PermissionResolutionAllowOnce { + t.Fatalf("expected Y => allow_once, got %q", got) + } + if got := normalizePermissionResolutionDecision(PermissionResolutionDecision("a")); got != PermissionResolutionAllowSession { + t.Fatalf("expected a => allow_session, got %q", got) + } + if got := normalizePermissionResolutionDecision(PermissionResolutionDecision("n")); got != PermissionResolutionReject { + t.Fatalf("expected n => reject, got %q", got) + } + if got := normalizePermissionResolutionDecision(PermissionResolutionDecision("???")); got != "" { + t.Fatalf("expected unknown => empty, got %q", got) + } + + if scope, err := rememberScopeFromDecision(PermissionResolutionAllowOnce); err != nil || scope != tools.SessionPermissionScopeOnce { + t.Fatalf("expected once scope, got %q / %v", scope, err) + } + if scope, err := rememberScopeFromDecision(PermissionResolutionAllowSession); err != nil || scope != tools.SessionPermissionScopeAlways { + t.Fatalf("expected always scope, got %q / %v", scope, err) + } + if scope, err := rememberScopeFromDecision(PermissionResolutionReject); err != nil || scope != tools.SessionPermissionScopeReject { + t.Fatalf("expected reject scope, got %q / %v", scope, err) + } + if _, err := rememberScopeFromDecision(PermissionResolutionDecision("invalid")); err == nil { + t.Fatalf("expected invalid decision error") + } + + category := permissionToolCategory(security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "filesystem_grep", + Resource: "filesystem_grep", + }, + }) + if category != "filesystem_read" { + t.Fatalf("expected filesystem_read category, got %q", category) + } + + category = permissionToolCategory(security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + }, + }) + if category != "webfetch" { + t.Fatalf("expected webfetch category, got %q", category) + } +} + +func TestResolvePermissionCanceledContext(t *testing.T) { + t.Parallel() + + service := NewWithFactory( + newRuntimeConfigManager(t), + &stubToolManager{}, + newMemoryStore(), + &scriptedProviderFactory{provider: &scriptedProvider{}}, + nil, + ) + request := registerPendingPermission(service, permissionExecutionInput{ + RunID: "run-canceled", + SessionID: "session-canceled", + Call: provider.ToolCall{ + ID: "call-canceled", + Name: "webfetch", + }, + }, security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + }, + }) + defer clearPendingPermission(service, request.RequestID) + request.ResultCh <- PermissionResolutionAllowOnce + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := service.ResolvePermission(ctx, PermissionResolutionInput{ + RequestID: request.RequestID, + Decision: PermissionResolutionAllowOnce, + }); !errors.Is(err, context.Canceled) { + t.Fatalf("expected context canceled, got %v", err) + } +} diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index c95bbcd3..216ce178 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -39,6 +39,7 @@ var runtimeSessionWorkdirs = struct { type Runtime interface { Run(ctx context.Context, input UserInput) error Compact(ctx context.Context, input CompactInput) (CompactResult, error) + ResolvePermission(ctx context.Context, input PermissionResolutionInput) error CancelActiveRun() bool Events() <-chan RuntimeEvent ListSessions(ctx context.Context) ([]SessionSummary, error) @@ -218,18 +219,13 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { } s.emit(ctx, EventToolStart, input.RunID, session.ID, call) - runCtx, cancel := context.WithTimeout(ctx, time.Duration(cfg.ToolTimeoutSec)*time.Second) - result, execErr := s.toolManager.Execute(runCtx, tools.ToolCallInput{ - ID: call.ID, - Name: call.Name, - Arguments: []byte(call.Arguments), - Workdir: activeWorkdir, - SessionID: session.ID, - EmitChunk: func(chunk []byte) { - s.emit(ctx, EventToolChunk, input.RunID, session.ID, string(chunk)) - }, + result, execErr := s.executeToolCallWithPermission(ctx, permissionExecutionInput{ + RunID: input.RunID, + SessionID: session.ID, + Call: call, + Workdir: activeWorkdir, + ToolTimeout: time.Duration(cfg.ToolTimeoutSec) * time.Second, }) - cancel() if s.isRunCanceled(execErr) { return s.handleRunError(ctx, input.RunID, session.ID, execErr) } @@ -243,9 +239,6 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { result.Content = execErr.Error() } if permissionEvent, ok := permissionEventFromError(execErr); ok { - if permissionEvent.decision == "ask" { - s.emit(ctx, EventPermissionRequest, input.RunID, session.ID, permissionEvent.toRequestPayload()) - } s.emit(ctx, EventPermissionResolved, input.RunID, session.ID, permissionEvent.toResolvedPayload()) } @@ -578,6 +571,7 @@ type permissionEventView struct { decision string reason string ruleID string + scope string resolvedAs string } @@ -609,6 +603,7 @@ func permissionEventFromError(err error) (permissionEventView, bool) { decision: decision, reason: reason, ruleID: strings.TrimSpace(permissionErr.RuleID()), + scope: strings.TrimSpace(permissionErr.RememberScope()), resolvedAs: resolvedAs, }, true } @@ -624,7 +619,7 @@ func (v permissionEventView) toRequestPayload() PermissionRequestPayload { Decision: v.decision, Reason: v.reason, RuleID: v.ruleID, - RememberScope: "", + RememberScope: v.scope, } } @@ -639,7 +634,7 @@ func (v permissionEventView) toResolvedPayload() PermissionResolvedPayload { Decision: v.decision, Reason: v.reason, RuleID: v.ruleID, - RememberScope: "", + RememberScope: v.scope, ResolvedAs: v.resolvedAs, } } diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index ff24909e..33e3590b 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -205,6 +205,12 @@ type stubToolManager struct { listCalls int executeCalls int lastInput tools.ToolCallInput + rememberErr error + remembered []struct { + sessionID string + action security.Action + scope tools.SessionPermissionScope + } } func (m *stubToolManager) ListAvailableSpecs(ctx context.Context, input tools.SpecListInput) ([]provider.ToolSpec, error) { @@ -228,6 +234,19 @@ func (m *stubToolManager) Execute(ctx context.Context, input tools.ToolCallInput return result, m.err } +func (m *stubToolManager) RememberSessionDecision(sessionID string, action security.Action, scope tools.SessionPermissionScope) error { + m.remembered = append(m.remembered, struct { + sessionID string + action security.Action + scope tools.SessionPermissionScope + }{ + sessionID: sessionID, + action: action, + scope: scope, + }) + return m.rememberErr +} + func TestServiceRun(t *testing.T) { tests := []struct { name string @@ -431,6 +450,9 @@ func TestServiceRunDelegatesToContextBuilder(t *testing.T) { manager := newRuntimeConfigManager(t) store := newMemoryStore() + session := newSession("memory reject") + session.ID = "session-memory-reject" + store.sessions[session.ID] = cloneSession(session) registry := tools.NewRegistry() registry.Register(&stubTool{name: "filesystem_read_file", content: "default"}) @@ -645,13 +667,16 @@ func TestServiceRunUsesToolManager(t *testing.T) { } } -func TestServiceRunEmitsPermissionRequestAndResolvedForAsk(t *testing.T) { +func TestServiceRunWaitsForPermissionResolutionAndContinues(t *testing.T) { t.Parallel() manager := newRuntimeConfigManager(t) store := newMemoryStore() + session := newSession("memory reject") + session.ID = "session-memory-reject" + store.sessions[session.ID] = cloneSession(session) registry := tools.NewRegistry() - tool := &stubTool{name: "webfetch", content: "should-not-run"} + tool := &stubTool{name: "webfetch", content: "fetched"} registry.Register(tool) engine, err := security.NewStaticGateway(security.DecisionAllow, []security.Rule{ @@ -690,34 +715,65 @@ func TestServiceRunEmitsPermissionRequestAndResolvedForAsk(t *testing.T) { } service := NewWithFactory(manager, toolManager, store, &scriptedProviderFactory{provider: scripted}, nil) - if err := service.Run(context.Background(), UserInput{RunID: "run-permission-ask", Content: "fetch private"}); err != nil { + runErrCh := make(chan error, 1) + go func() { + runErrCh <- service.Run(context.Background(), UserInput{RunID: "run-permission-ask", Content: "fetch private"}) + }() + + var requestPayload PermissionRequestPayload + deadline := time.After(3 * time.Second) +waitRequest: + for { + select { + case <-deadline: + t.Fatalf("timed out waiting permission request event") + case event := <-service.Events(): + if event.Type != EventPermissionRequest { + continue + } + payload, ok := event.Payload.(PermissionRequestPayload) + if !ok { + t.Fatalf("expected PermissionRequestPayload, got %#v", event.Payload) + } + requestPayload = payload + break waitRequest + } + } + + if strings.TrimSpace(requestPayload.RequestID) == "" { + t.Fatalf("expected non-empty permission request id") + } + if strings.TrimSpace(requestPayload.RememberScope) != "" { + t.Fatalf("expected empty remember scope for permission_request, got %q", requestPayload.RememberScope) + } + if requestPayload.ToolName != "webfetch" || requestPayload.Decision != "ask" { + t.Fatalf("unexpected permission request payload: %+v", requestPayload) + } + + if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ + RequestID: requestPayload.RequestID, + Decision: PermissionResolutionAllowSession, + }); err != nil { + t.Fatalf("ResolvePermission() error = %v", err) + } + if err := <-runErrCh; err != nil { t.Fatalf("Run() error = %v", err) } - if tool.callCount != 0 { - t.Fatalf("expected blocked tool not to execute, got %d", tool.callCount) + if tool.callCount != 1 { + t.Fatalf("expected allowed tool to execute once, got %d", tool.callCount) } events := collectRuntimeEvents(service.Events()) assertEventSequence(t, events, []EventType{ - EventPermissionRequest, EventPermissionResolved, EventToolResult, EventAgentDone, }) assertNoEventType(t, events, EventError) - var ( - requestPayload PermissionRequestPayload - resolvedPayload PermissionResolvedPayload - ) + var resolvedPayload PermissionResolvedPayload for _, event := range events { switch event.Type { - case EventPermissionRequest: - payload, ok := event.Payload.(PermissionRequestPayload) - if !ok { - t.Fatalf("expected PermissionRequestPayload, got %#v", event.Payload) - } - requestPayload = payload case EventPermissionResolved: payload, ok := event.Payload.(PermissionResolvedPayload) if !ok { @@ -727,17 +783,14 @@ func TestServiceRunEmitsPermissionRequestAndResolvedForAsk(t *testing.T) { } } - if requestPayload.ToolName != "webfetch" || requestPayload.Decision != "ask" { - t.Fatalf("unexpected permission request payload: %+v", requestPayload) - } - if requestPayload.RuleID != "ask-webfetch" { - t.Fatalf("expected rule id ask-webfetch, got %+v", requestPayload) - } - if resolvedPayload.ToolName != "webfetch" || resolvedPayload.Decision != "ask" { + if resolvedPayload.ToolName != "webfetch" || resolvedPayload.Decision != "allow" { t.Fatalf("unexpected permission resolved payload: %+v", resolvedPayload) } - if resolvedPayload.ResolvedAs != "rejected" { - t.Fatalf("expected resolved_as rejected, got %+v", resolvedPayload) + if resolvedPayload.ResolvedAs != "approved" { + t.Fatalf("expected resolved_as approved, got %+v", resolvedPayload) + } + if resolvedPayload.RememberScope != string(tools.SessionPermissionScopeAlways) { + t.Fatalf("expected remember scope always_session, got %+v", resolvedPayload) } } @@ -821,6 +874,97 @@ func TestServiceRunEmitsPermissionResolvedForDeny(t *testing.T) { t.Fatalf("expected permission resolved event payload") } +func TestServiceRunEmitsRememberScopeWhenSessionRejectMemoryHits(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + session := newSession("memory reject") + session.ID = "session-memory-reject" + store.sessions[session.ID] = cloneSession(session) + registry := tools.NewRegistry() + tool := &stubTool{name: "webfetch", content: "should-not-run"} + registry.Register(tool) + + engine, err := security.NewStaticGateway(security.DecisionAllow, []security.Rule{ + { + ID: "ask-webfetch", + Type: security.ActionTypeRead, + Resource: "webfetch", + Decision: security.DecisionAsk, + Reason: "requires approval", + }, + }) + if err != nil { + t.Fatalf("new static gateway: %v", err) + } + toolManager, err := tools.NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new tool manager: %v", err) + } + if err := toolManager.RememberSessionDecision("session-memory-reject", security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + Operation: "fetch", + TargetType: security.TargetTypeURL, + Target: "https://example.com/private", + }, + }, tools.SessionPermissionScopeReject); err != nil { + t.Fatalf("remember session reject: %v", err) + } + + scripted := &scriptedProvider{ + responses: []provider.ChatResponse{ + { + Message: provider.Message{ + Role: "assistant", + ToolCalls: []provider.ToolCall{ + {ID: "call-memory-reject", Name: "webfetch", Arguments: `{"url":"https://example.com/private"}`}, + }, + }, + FinishReason: "tool_calls", + }, + { + Message: provider.Message{Role: "assistant", Content: "done"}, + FinishReason: "stop", + }, + }, + } + + service := NewWithFactory(manager, toolManager, store, &scriptedProviderFactory{provider: scripted}, nil) + if err := service.Run(context.Background(), UserInput{ + SessionID: "session-memory-reject", + RunID: "run-memory-reject", + Content: "fetch private", + }); err != nil { + t.Fatalf("Run() error = %v", err) + } + if tool.callCount != 0 { + t.Fatalf("expected remembered reject to skip tool execution, got %d", tool.callCount) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventPermissionResolved, EventToolResult, EventAgentDone}) + assertNoEventType(t, events, EventPermissionRequest) + + for _, event := range events { + if event.Type != EventPermissionResolved { + continue + } + payload, ok := event.Payload.(PermissionResolvedPayload) + if !ok { + t.Fatalf("expected PermissionResolvedPayload, got %#v", event.Payload) + } + if payload.RememberScope != string(tools.SessionPermissionScopeReject) { + t.Fatalf("expected remember_scope reject, got %+v", payload) + } + return + } + t.Fatalf("expected permission resolved event payload") +} + func TestServiceRunHandlesToolManagerSpecError(t *testing.T) { manager := newRuntimeConfigManager(t) store := newMemoryStore() @@ -2418,3 +2562,47 @@ func TestProviderRetryBackoff(t *testing.T) { }) } } + +func TestPermissionEventViewPayloadMapping(t *testing.T) { + t.Parallel() + + view := permissionEventView{ + toolName: "webfetch", + actionType: string(security.ActionTypeRead), + operation: "fetch", + targetType: string(security.TargetTypeURL), + target: "https://example.com", + decision: "ask", + reason: "need approval", + ruleID: "rule-1", + scope: string(tools.SessionPermissionScopeAlways), + resolvedAs: "rejected", + } + + requestPayload := view.toRequestPayload() + if requestPayload.ToolName != view.toolName || + requestPayload.ActionType != view.actionType || + requestPayload.Operation != view.operation || + requestPayload.TargetType != view.targetType || + requestPayload.Target != view.target || + requestPayload.Decision != view.decision || + requestPayload.Reason != view.reason || + requestPayload.RuleID != view.ruleID || + requestPayload.RememberScope != view.scope { + t.Fatalf("unexpected request payload: %+v", requestPayload) + } + + resolvedPayload := view.toResolvedPayload() + if resolvedPayload.ToolName != view.toolName || + resolvedPayload.ActionType != view.actionType || + resolvedPayload.Operation != view.operation || + resolvedPayload.TargetType != view.targetType || + resolvedPayload.Target != view.target || + resolvedPayload.Decision != view.decision || + resolvedPayload.Reason != view.reason || + resolvedPayload.RuleID != view.ruleID || + resolvedPayload.RememberScope != view.scope || + resolvedPayload.ResolvedAs != view.resolvedAs { + t.Fatalf("unexpected resolved payload: %+v", resolvedPayload) + } +} diff --git a/internal/security/policy.go b/internal/security/policy.go new file mode 100644 index 00000000..f414655c --- /dev/null +++ b/internal/security/policy.go @@ -0,0 +1,490 @@ +package security + +import ( + "context" + "fmt" + "net/url" + "path/filepath" + "sort" + "strings" +) + +// PolicyRule 描述一条可组合的权限策略规则。 +// 规则按 Priority 从高到低匹配，同优先级保持声明顺序。 +type PolicyRule struct { + ID string + Priority int + Decision Decision + Reason string + + ActionTypes []ActionType + ResourcePatterns []string + ToolCategories []string + TargetTypes []TargetType + + PathPatterns []string + PathSegmentKeywords []string + PathBasenamePatterns []string + RequireSensitivePath bool + + HostPatterns []string + RequireHostMatch bool + RequireHostMissing bool +} + +// PolicyEngine 基于结构化命中条件执行权限决策。 +type PolicyEngine struct { + defaultDecision Decision + rules []PolicyRule +} + +type compiledRule struct { + rule PolicyRule + order int +} + +type actionView struct { + action Action + resource string + toolCategory string + targetType TargetType + target string + targetPath string + host string + sensitive bool +} + +// NewPolicyEngine 创建支持优先级与多条件匹配的权限引擎。 +func NewPolicyEngine(defaultDecision Decision, rules []PolicyRule) (*PolicyEngine, error) { + if defaultDecision == "" { + defaultDecision = DecisionAllow + } + if err := defaultDecision.Validate(); err != nil { + return nil, err + } + + compiled := make([]compiledRule, 0, len(rules)) + for idx := range rules { + rule := rules[idx] + if strings.TrimSpace(rule.ID) == "" { + return nil, fmt.Errorf("security: policy rule id is empty at index %d", idx) + } + if err := rule.Decision.Validate(); err != nil { + return nil, fmt.Errorf("security: policy rule %q: %w", rule.ID, err) + } + for _, actionType := range rule.ActionTypes { + if actionType == "" { + continue + } + if err := actionType.Validate(); err != nil { + return nil, fmt.Errorf("security: policy rule %q: %w", rule.ID, err) + } + } + compiled = append(compiled, compiledRule{ + rule: normalizePolicyRule(rule), + order: idx, + }) + } + + sort.SliceStable(compiled, func(i, j int) bool { + if compiled[i].rule.Priority == compiled[j].rule.Priority { + return compiled[i].order < compiled[j].order + } + return compiled[i].rule.Priority > compiled[j].rule.Priority + }) + + sortedRules := make([]PolicyRule, 0, len(compiled)) + for _, item := range compiled { + sortedRules = append(sortedRules, item.rule) + } + + return &PolicyEngine{ + defaultDecision: defaultDecision, + rules: sortedRules, + }, nil +} + +// Check 返回首条命中规则；若无命中则返回默认决策。 +func (e *PolicyEngine) Check(ctx context.Context, action Action) (CheckResult, error) { + if err := ctx.Err(); err != nil { + return CheckResult{}, err + } + if err := action.Validate(); err != nil { + return CheckResult{}, err + } + + view := newActionView(action) + for _, rule := range e.rules { + if !matchesPolicyRule(rule, view) { + continue + } + matchedRule := Rule{ + ID: rule.ID, + Type: action.Type, + Resource: action.Payload.Resource, + Decision: rule.Decision, + Reason: rule.Reason, + } + return CheckResult{ + Decision: rule.Decision, + Action: action, + Rule: &matchedRule, + Reason: strings.TrimSpace(rule.Reason), + }, nil + } + + return CheckResult{ + Decision: e.defaultDecision, + Action: action, + }, nil +} + +// NewRecommendedPolicyEngine 返回推荐安全策略： +// bash=ask, filesystem write=ask, filesystem read敏感路径=ask/deny, webfetch白名单allow其余ask。 +func NewRecommendedPolicyEngine() (*PolicyEngine, error) { + const ( + reasonAskBash = "bash command requires approval" + reasonAskFilesystemWrite = "filesystem write requires approval" + reasonDenyPrivateKeyRead = "reading private key material is blocked" + reasonAskSensitiveRead = "reading sensitive path requires approval" + reasonAllowWebfetchDomain = "approved web domain" + reasonAskWebfetchDomain = "external web domain requires approval" + ) + + rules := []PolicyRule{ + { + ID: "deny-sensitive-private-keys", + Priority: 1000, + Decision: DecisionDeny, + Reason: reasonDenyPrivateKeyRead, + ActionTypes: []ActionType{ActionTypeRead}, + ToolCategories: []string{"filesystem_read"}, + PathBasenamePatterns: []string{"id_rsa", "id_dsa", "id_ecdsa", "id_ed25519", "*.pem", "*.p12", "*.pfx", "*.key"}, + }, + { + ID: "ask-sensitive-filesystem-read", + Priority: 900, + Decision: DecisionAsk, + Reason: reasonAskSensitiveRead, + ActionTypes: []ActionType{ActionTypeRead}, + ToolCategories: []string{"filesystem_read"}, + RequireSensitivePath: true, + PathSegmentKeywords: []string{"secrets", ".ssh", ".gnupg", ".aws", ".config"}, + PathBasenamePatterns: []string{".env", ".env.*", "*.env", "*.secret", "*.secrets", "*.token"}, + ResourcePatterns: []string{"filesystem_read_*", "filesystem_grep", "filesystem_glob"}, + TargetTypes: []TargetType{TargetTypePath, TargetTypeDirectory}, + RequireHostMissing: false, + RequireHostMatch: false, + }, + { + ID: "ask-all-bash", + Priority: 800, + Decision: DecisionAsk, + Reason: reasonAskBash, + ActionTypes: []ActionType{ActionTypeBash}, + ResourcePatterns: []string{"bash"}, + }, + { + ID: "ask-filesystem-write", + Priority: 780, + Decision: DecisionAsk, + Reason: reasonAskFilesystemWrite, + ActionTypes: []ActionType{ActionTypeWrite}, + ResourcePatterns: []string{"filesystem_write_*", "filesystem_edit"}, + }, + { + ID: "allow-webfetch-whitelist", + Priority: 760, + Decision: DecisionAllow, + Reason: reasonAllowWebfetchDomain, + ActionTypes: []ActionType{ActionTypeRead}, + ResourcePatterns: []string{"webfetch"}, + HostPatterns: []string{"github.com", "*.github.com"}, + RequireHostMatch: true, + }, + { + ID: "ask-webfetch-non-whitelist", + Priority: 740, + Decision: DecisionAsk, + Reason: reasonAskWebfetchDomain, + ActionTypes: []ActionType{ActionTypeRead}, + ResourcePatterns: []string{"webfetch"}, + HostPatterns: []string{"github.com", "*.github.com"}, + RequireHostMissing: true, + }, + } + + return NewPolicyEngine(DecisionAllow, rules) +} + +func normalizePolicyRule(rule PolicyRule) PolicyRule { + rule.ID = strings.TrimSpace(rule.ID) + rule.Reason = strings.TrimSpace(rule.Reason) + rule.ActionTypes = normalizeActionTypes(rule.ActionTypes) + rule.ResourcePatterns = normalizeLowerList(rule.ResourcePatterns) + rule.ToolCategories = normalizeLowerList(rule.ToolCategories) + rule.TargetTypes = normalizeTargetTypes(rule.TargetTypes) + rule.PathPatterns = normalizePathPatterns(rule.PathPatterns) + rule.PathSegmentKeywords = normalizeLowerList(rule.PathSegmentKeywords) + rule.PathBasenamePatterns = normalizeLowerList(rule.PathBasenamePatterns) + rule.HostPatterns = normalizeHostPatterns(rule.HostPatterns) + return rule +} + +func normalizeActionTypes(values []ActionType) []ActionType { + out := make([]ActionType, 0, len(values)) + for _, value := range values { + if strings.TrimSpace(string(value)) == "" { + continue + } + out = append(out, ActionType(strings.TrimSpace(string(value)))) + } + return out +} + +func normalizeTargetTypes(values []TargetType) []TargetType { + out := make([]TargetType, 0, len(values)) + for _, value := range values { + if strings.TrimSpace(string(value)) == "" { + continue + } + out = append(out, TargetType(strings.TrimSpace(string(value)))) + } + return out +} + +func normalizeLowerList(values []string) []string { + out := make([]string, 0, len(values)) + for _, value := range values { + trimmed := strings.ToLower(strings.TrimSpace(value)) + if trimmed == "" { + continue + } + out = append(out, trimmed) + } + return out +} + +func normalizePathPatterns(values []string) []string { + out := make([]string, 0, len(values)) + for _, value := range values { + trimmed := filepath.ToSlash(strings.ToLower(strings.TrimSpace(value))) + if trimmed == "" { + continue + } + out = append(out, trimmed) + } + return out +} + +func normalizeHostPatterns(values []string) []string { + out := make([]string, 0, len(values)) + for _, value := range values { + trimmed := strings.ToLower(strings.TrimSpace(value)) + trimmed = strings.TrimPrefix(trimmed, ".") + if trimmed == "" { + continue + } + out = append(out, trimmed) + } + return out +} + +func newActionView(action Action) actionView { + resource := strings.ToLower(strings.TrimSpace(action.Payload.Resource)) + target := strings.TrimSpace(action.Payload.Target) + host := "" + if action.Payload.TargetType == TargetTypeURL || resource == "webfetch" { + host = parseURLHost(target) + } + targetPath := filepath.ToSlash(strings.ToLower(strings.TrimSpace(target))) + category := deriveToolCategory(action) + sensitive := classifySensitivePath(targetPath) + + return actionView{ + action: action, + resource: resource, + toolCategory: category, + targetType: action.Payload.TargetType, + target: target, + targetPath: targetPath, + host: host, + sensitive: sensitive, + } +} + +func matchesPolicyRule(rule PolicyRule, view actionView) bool { + if len(rule.ActionTypes) > 0 { + matched := false + for _, actionType := range rule.ActionTypes { + if view.action.Type == actionType { + matched = true + break + } + } + if !matched { + return false + } + } + + if len(rule.ResourcePatterns) > 0 && !matchesAnyPattern(view.resource, rule.ResourcePatterns) { + return false + } + if len(rule.ToolCategories) > 0 && !containsString(rule.ToolCategories, view.toolCategory) { + return false + } + + if len(rule.TargetTypes) > 0 { + matched := false + for _, targetType := range rule.TargetTypes { + if view.targetType == targetType { + matched = true + break + } + } + if !matched { + return false + } + } + + if rule.RequireSensitivePath && !view.sensitive { + return false + } + pathMatcherCount := len(rule.PathPatterns) + len(rule.PathSegmentKeywords) + len(rule.PathBasenamePatterns) + if pathMatcherCount > 0 { + pathMatched := false + if len(rule.PathPatterns) > 0 && matchesAnyPattern(view.targetPath, rule.PathPatterns) { + pathMatched = true + } + if len(rule.PathSegmentKeywords) > 0 && matchesPathSegmentKeyword(view.targetPath, rule.PathSegmentKeywords) { + pathMatched = true + } + if len(rule.PathBasenamePatterns) > 0 && matchesPathBasenamePattern(view.targetPath, rule.PathBasenamePatterns) { + pathMatched = true + } + if !pathMatched { + return false + } + } + + hostMatched := len(rule.HostPatterns) == 0 || matchesAnyPattern(view.host, rule.HostPatterns) + if rule.RequireHostMatch && !hostMatched { + return false + } + if rule.RequireHostMissing && hostMatched { + return false + } + + return true +} + +func deriveToolCategory(action Action) string { + resource := strings.ToLower(strings.TrimSpace(action.Payload.Resource)) + switch action.Type { + case ActionTypeRead: + if strings.HasPrefix(resource, "filesystem_") { + return "filesystem_read" + } + case ActionTypeWrite: + if strings.HasPrefix(resource, "filesystem_") { + return "filesystem_write" + } + case ActionTypeBash: + return "bash" + } + if resource != "" { + return resource + } + return strings.ToLower(strings.TrimSpace(action.Payload.ToolName)) +} + +func classifySensitivePath(normalizedTargetPath string) bool { + if normalizedTargetPath == "" { + return false + } + return matchesPathSegmentKeyword(normalizedTargetPath, []string{"secrets", ".ssh", ".gnupg", ".aws", ".config"}) || + matchesPathBasenamePattern(normalizedTargetPath, []string{".env", ".env.*", "*.env", "*.secret", "*.secrets", "*.token", "*.key", "*.pem", "id_rsa", "id_ed25519"}) +} + +func matchesPathSegmentKeyword(normalizedTargetPath string, keywords []string) bool { + if normalizedTargetPath == "" || len(keywords) == 0 { + return false + } + segments := strings.Split(normalizedTargetPath, "/") + for _, segment := range segments { + token := strings.ToLower(strings.TrimSpace(segment)) + if token == "" { + continue + } + for _, keyword := range keywords { + if token == keyword || strings.Contains(token, keyword) { + return true + } + } + } + return false +} + +func matchesPathBasenamePattern(normalizedTargetPath string, patterns []string) bool { + if normalizedTargetPath == "" || len(patterns) == 0 { + return false + } + base := strings.ToLower(filepath.Base(normalizedTargetPath)) + for _, pattern := range patterns { + matched, err := filepath.Match(pattern, base) + if err != nil { + continue + } + if matched { + return true + } + } + return false +} + +func matchesAnyPattern(value string, patterns []string) bool { + if len(patterns) == 0 { + return true + } + normalized := strings.ToLower(strings.TrimSpace(value)) + for _, pattern := range patterns { + p := strings.ToLower(strings.TrimSpace(pattern)) + if p == "" { + continue + } + matched, err := filepath.Match(p, normalized) + if err == nil && matched { + return true + } + if p == normalized { + return true + } + if strings.HasPrefix(p, "*.") && strings.HasSuffix(normalized, p[1:]) { + return true + } + } + return false +} + +func parseURLHost(raw string) string { + trimmed := strings.TrimSpace(raw) + if trimmed == "" { + return "" + } + parsed, err := url.Parse(trimmed) + if err != nil || parsed == nil { + return "" + } + host := strings.ToLower(strings.TrimSpace(parsed.Hostname())) + return strings.TrimPrefix(host, ".") +} + +func containsString(values []string, target string) bool { + target = strings.ToLower(strings.TrimSpace(target)) + for _, value := range values { + if strings.EqualFold(strings.TrimSpace(value), target) { + return true + } + } + return false +} diff --git a/internal/security/policy_test.go b/internal/security/policy_test.go new file mode 100644 index 00000000..ab12d2ee --- /dev/null +++ b/internal/security/policy_test.go @@ -0,0 +1,189 @@ +package security + +import ( + "context" + "testing" +) + +func TestPolicyEngineRecommendedRules(t *testing.T) { + t.Parallel() + + engine, err := NewRecommendedPolicyEngine() + if err != nil { + t.Fatalf("new recommended engine: %v", err) + } + + tests := []struct { + name string + action Action + wantDecision Decision + wantRuleID string + }{ + { + name: "bash always ask", + action: Action{ + Type: ActionTypeBash, + Payload: ActionPayload{ + ToolName: "bash", + Resource: "bash", + Operation: "command", + TargetType: TargetTypeCommand, + Target: "ls -la", + }, + }, + wantDecision: DecisionAsk, + wantRuleID: "ask-all-bash", + }, + { + name: "filesystem write ask", + action: Action{ + Type: ActionTypeWrite, + Payload: ActionPayload{ + ToolName: "filesystem_write_file", + Resource: "filesystem_write_file", + Operation: "write_file", + TargetType: TargetTypePath, + Target: "README.md", + }, + }, + wantDecision: DecisionAsk, + wantRuleID: "ask-filesystem-write", + }, + { + name: "filesystem read sensitive path ask", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + Operation: "read_file", + TargetType: TargetTypePath, + Target: ".env.production", + }, + }, + wantDecision: DecisionAsk, + wantRuleID: "ask-sensitive-filesystem-read", + }, + { + name: "filesystem read private key deny", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + Operation: "read_file", + TargetType: TargetTypePath, + Target: "C:/Users/test/.ssh/id_rsa", + }, + }, + wantDecision: DecisionDeny, + wantRuleID: "deny-sensitive-private-keys", + }, + { + name: "filesystem read normal source allow", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + Operation: "read_file", + TargetType: TargetTypePath, + Target: "internal/runtime/runtime.go", + }, + }, + wantDecision: DecisionAllow, + wantRuleID: "", + }, + { + name: "webfetch whitelist allow", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + Operation: "fetch", + TargetType: TargetTypeURL, + Target: "https://github.com/1024XEngineer/neo-code", + }, + }, + wantDecision: DecisionAllow, + wantRuleID: "allow-webfetch-whitelist", + }, + { + name: "webfetch non-whitelist ask", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + Operation: "fetch", + TargetType: TargetTypeURL, + Target: "https://example.com", + }, + }, + wantDecision: DecisionAsk, + wantRuleID: "ask-webfetch-non-whitelist", + }, + { + name: "webfetch docs wildcard host is not implicitly trusted", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + Operation: "fetch", + TargetType: TargetTypeURL, + Target: "https://docs.attacker.com", + }, + }, + wantDecision: DecisionAsk, + wantRuleID: "ask-webfetch-non-whitelist", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + result, checkErr := engine.Check(context.Background(), tt.action) + if checkErr != nil { + t.Fatalf("Check() error = %v", checkErr) + } + if result.Decision != tt.wantDecision { + t.Fatalf("expected decision %q, got %q", tt.wantDecision, result.Decision) + } + if tt.wantRuleID == "" { + if result.Rule != nil { + t.Fatalf("expected no matched rule, got %+v", result.Rule) + } + return + } + if result.Rule == nil || result.Rule.ID != tt.wantRuleID { + t.Fatalf("expected rule id %q, got %+v", tt.wantRuleID, result.Rule) + } + }) + } +} + +func TestNewPolicyEngineValidation(t *testing.T) { + t.Parallel() + + _, err := NewPolicyEngine(Decision("invalid"), nil) + if err == nil { + t.Fatalf("expected invalid default decision error") + } + + _, err = NewPolicyEngine(DecisionAllow, []PolicyRule{ + {ID: "", Decision: DecisionAsk}, + }) + if err == nil { + t.Fatalf("expected missing rule id error") + } + + _, err = NewPolicyEngine(DecisionAllow, []PolicyRule{ + {ID: "r1", Decision: Decision("invalid")}, + }) + if err == nil { + t.Fatalf("expected invalid rule decision error") + } +} diff --git a/internal/tools/manager.go b/internal/tools/manager.go index 4c099c4c..64ecb01a 100644 --- a/internal/tools/manager.go +++ b/internal/tools/manager.go @@ -20,6 +20,7 @@ type SpecListInput struct { type Manager interface { ListAvailableSpecs(ctx context.Context, input SpecListInput) ([]provider.ToolSpec, error) Execute(ctx context.Context, input ToolCallInput) (ToolResult, error) + RememberSessionDecision(sessionID string, action security.Action, scope SessionPermissionScope) error } // Executor is the concrete tool execution layer under the manager. @@ -50,6 +51,7 @@ type PermissionDecisionError struct { action security.Action reason string ruleID string + scope SessionPermissionScope } // Error returns a stable error message for the blocked tool call. @@ -112,12 +114,21 @@ func (e *PermissionDecisionError) RuleID() string { return strings.TrimSpace(e.ruleID) } +// RememberScope 返回触发该权限结果时命中的会话记忆范围。 +func (e *PermissionDecisionError) RememberScope() string { + if e == nil { + return "" + } + return strings.TrimSpace(string(e.scope)) +} + // DefaultManager routes tool calls through the permission engine, workspace // sandbox, and executor. type DefaultManager struct { - executor Executor - engine security.PermissionEngine - sandbox WorkspaceSandbox + executor Executor + engine security.PermissionEngine + sandbox WorkspaceSandbox + sessionDecisions *sessionPermissionMemory } // NewManager creates a manager that wraps an executor with security checks. @@ -137,9 +148,10 @@ func NewManager(executor Executor, engine security.PermissionEngine, sandbox Wor } return &DefaultManager{ - executor: executor, - engine: engine, - sandbox: sandbox, + executor: executor, + engine: engine, + sandbox: sandbox, + sessionDecisions: newSessionPermissionMemory(), }, nil } @@ -175,6 +187,28 @@ func (m *DefaultManager) Execute(ctx context.Context, input ToolCallInput) (Tool result.ToolCallID = input.ID return result, err } + // deny 规则始终优先，避免 session 记忆覆盖硬性安全策略。 + if decision.Decision == security.DecisionDeny { + result := blockedToolResult(input, decision) + return result, permissionErrorFromDecision(decision) + } + // session 记忆仅用于自动处理 ask，不提升原本已 allow 的策略结果。 + if decision.Decision == security.DecisionAsk && m.sessionDecisions != nil { + if rememberedDecision, rememberedScope, ok := m.sessionDecisions.resolve(input.SessionID, action); ok { + decision = security.CheckResult{ + Decision: rememberedDecision, + Action: action, + Reason: sessionDecisionReason(rememberedScope), + } + if rememberedScope != "" { + decision.Rule = &security.Rule{ + ID: "session-memory:" + string(rememberedScope), + Decision: rememberedDecision, + Reason: decision.Reason, + } + } + } + } if decision.Decision != security.DecisionAllow { result := blockedToolResult(input, decision) return result, permissionErrorFromDecision(decision) @@ -194,6 +228,17 @@ func (m *DefaultManager) Execute(ctx context.Context, input ToolCallInput) (Tool return m.executor.Execute(ctx, input) } +// RememberSessionDecision 记录会话内权限记忆，用于后续同类 action 快速决策。 +func (m *DefaultManager) RememberSessionDecision(sessionID string, action security.Action, scope SessionPermissionScope) error { + if m == nil { + return errors.New("tools: manager is nil") + } + if m.sessionDecisions == nil { + m.sessionDecisions = newSessionPermissionMemory() + } + return m.sessionDecisions.remember(sessionID, action, scope) +} + func blockedToolResult(input ToolCallInput, decision security.CheckResult) ToolResult { reason := "permission denied" if decision.Decision == security.DecisionAsk { @@ -219,6 +264,39 @@ func permissionErrorFromDecision(decision security.CheckResult) error { action: decision.Action, reason: decision.Reason, ruleID: ruleID, + scope: extractRememberScope(decision), + } +} + +// extractRememberScope 从决策规则中提取会话记忆范围。 +func extractRememberScope(decision security.CheckResult) SessionPermissionScope { + if decision.Rule == nil { + return "" + } + ruleID := strings.TrimSpace(decision.Rule.ID) + switch ruleID { + case "session-memory:" + string(SessionPermissionScopeOnce): + return SessionPermissionScopeOnce + case "session-memory:" + string(SessionPermissionScopeAlways): + return SessionPermissionScopeAlways + case "session-memory:" + string(SessionPermissionScopeReject): + return SessionPermissionScopeReject + default: + return "" + } +} + +// sessionDecisionReason 生成会话记忆命中的统一原因文本。 +func sessionDecisionReason(scope SessionPermissionScope) string { + switch scope { + case SessionPermissionScopeOnce: + return "session permission remembered: once" + case SessionPermissionScopeAlways: + return "session permission remembered: always(session)" + case SessionPermissionScopeReject: + return "session permission remembered: reject" + default: + return "session permission remembered" } } diff --git a/internal/tools/manager_test.go b/internal/tools/manager_test.go index a4a2fe98..1292de68 100644 --- a/internal/tools/manager_test.go +++ b/internal/tools/manager_test.go @@ -377,6 +377,9 @@ func TestPermissionDecisionError(t *testing.T) { if err.Action().Type != security.ActionTypeRead { t.Fatalf("expected action type read, got %q", err.Action().Type) } + if err.RememberScope() != "" { + t.Fatalf("expected empty remember scope, got %q", err.RememberScope()) + } if errors.Is(err, context.Canceled) { t.Fatalf("permission error should not match unrelated errors") } @@ -391,13 +394,349 @@ func TestPermissionDecisionError(t *testing.T) { if denyErr.ToolName() != "" { t.Fatalf("expected empty tool name, got %q", denyErr.ToolName()) } + if denyErr.RememberScope() != "" { + t.Fatalf("expected empty remember scope, got %q", denyErr.RememberScope()) + } var nilErr *PermissionDecisionError - if nilErr.Error() != "" || nilErr.Decision() != "" || nilErr.ToolName() != "" { + if nilErr.Error() != "" || nilErr.Decision() != "" || nilErr.ToolName() != "" || nilErr.RememberScope() != "" { t.Fatalf("expected nil permission error helpers to be empty") } } +func TestDefaultManagerSessionPermissionMemory(t *testing.T) { + t.Parallel() + + newAskManager := func(t *testing.T) (*DefaultManager, *managerStubTool) { + t.Helper() + registry := NewRegistry() + webTool := &managerStubTool{name: "webfetch", content: "ok"} + registry.Register(webTool) + engine, err := security.NewStaticGateway(security.DecisionAllow, []security.Rule{ + { + ID: "ask-webfetch", + Type: security.ActionTypeRead, + Resource: "webfetch", + Decision: security.DecisionAsk, + Reason: "requires approval", + }, + }) + if err != nil { + t.Fatalf("new engine: %v", err) + } + manager, err := NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new manager: %v", err) + } + return manager, webTool + } + + t.Run("once allows only first follow-up", func(t *testing.T) { + t.Parallel() + manager, webTool := newAskManager(t) + input := ToolCallInput{ + ID: "call-once", + Name: "webfetch", + Arguments: []byte(`{"url":"https://example.com/once"}`), + SessionID: "session-once", + } + + _, err := manager.Execute(context.Background(), input) + var permissionErr *PermissionDecisionError + if !errors.As(err, &permissionErr) || permissionErr.Decision() != "ask" { + t.Fatalf("expected initial ask decision, got %v", err) + } + if rememberErr := manager.RememberSessionDecision(input.SessionID, permissionErr.Action(), SessionPermissionScopeOnce); rememberErr != nil { + t.Fatalf("remember once: %v", rememberErr) + } + + result, err := manager.Execute(context.Background(), input) + if err != nil { + t.Fatalf("expected remembered once allow, got %v", err) + } + if result.IsError { + t.Fatalf("expected non-error result, got %+v", result) + } + if webTool.callCount != 1 { + t.Fatalf("expected tool call count 1 after once allow, got %d", webTool.callCount) + } + + _, err = manager.Execute(context.Background(), input) + if !errors.As(err, &permissionErr) || permissionErr.Decision() != "ask" { + t.Fatalf("expected ask after once consumed, got %v", err) + } + }) + + t.Run("always(session) keeps allowing in same session", func(t *testing.T) { + t.Parallel() + manager, webTool := newAskManager(t) + input := ToolCallInput{ + ID: "call-always", + Name: "webfetch", + Arguments: []byte(`{"url":"https://example.com/always"}`), + SessionID: "session-always", + } + + _, err := manager.Execute(context.Background(), input) + var permissionErr *PermissionDecisionError + if !errors.As(err, &permissionErr) || permissionErr.Decision() != "ask" { + t.Fatalf("expected initial ask decision, got %v", err) + } + if rememberErr := manager.RememberSessionDecision(input.SessionID, permissionErr.Action(), SessionPermissionScopeAlways); rememberErr != nil { + t.Fatalf("remember always: %v", rememberErr) + } + + for i := 0; i < 2; i++ { + if _, err := manager.Execute(context.Background(), input); err != nil { + t.Fatalf("expected always allow on iteration %d, got %v", i, err) + } + } + if webTool.callCount != 2 { + t.Fatalf("expected tool to execute twice, got %d", webTool.callCount) + } + }) + + t.Run("reject denies in same session and keeps scope metadata", func(t *testing.T) { + t.Parallel() + manager, webTool := newAskManager(t) + input := ToolCallInput{ + ID: "call-reject", + Name: "webfetch", + Arguments: []byte(`{"url":"https://example.com/reject"}`), + SessionID: "session-reject", + } + + _, err := manager.Execute(context.Background(), input) + var permissionErr *PermissionDecisionError + if !errors.As(err, &permissionErr) || permissionErr.Decision() != "ask" { + t.Fatalf("expected initial ask decision, got %v", err) + } + if rememberErr := manager.RememberSessionDecision(input.SessionID, permissionErr.Action(), SessionPermissionScopeReject); rememberErr != nil { + t.Fatalf("remember reject: %v", rememberErr) + } + + _, err = manager.Execute(context.Background(), input) + if !errors.As(err, &permissionErr) { + t.Fatalf("expected permission error, got %v", err) + } + if permissionErr.Decision() != "deny" { + t.Fatalf("expected deny from remembered reject, got %q", permissionErr.Decision()) + } + if permissionErr.RememberScope() != string(SessionPermissionScopeReject) { + t.Fatalf("expected reject remember scope, got %q", permissionErr.RememberScope()) + } + if webTool.callCount != 0 { + t.Fatalf("expected rejected call to skip tool execution, got %d", webTool.callCount) + } + }) + + t.Run("session memory does not leak across sessions", func(t *testing.T) { + t.Parallel() + manager, _ := newAskManager(t) + inputA := ToolCallInput{ + ID: "call-session-a", + Name: "webfetch", + Arguments: []byte(`{"url":"https://example.com/session-a"}`), + SessionID: "session-a", + } + inputB := ToolCallInput{ + ID: "call-session-b", + Name: "webfetch", + Arguments: []byte(`{"url":"https://example.com/session-a"}`), + SessionID: "session-b", + } + + _, err := manager.Execute(context.Background(), inputA) + var permissionErr *PermissionDecisionError + if !errors.As(err, &permissionErr) { + t.Fatalf("expected permission ask on session A, got %v", err) + } + if rememberErr := manager.RememberSessionDecision(inputA.SessionID, permissionErr.Action(), SessionPermissionScopeAlways); rememberErr != nil { + t.Fatalf("remember session A always: %v", rememberErr) + } + if _, err := manager.Execute(context.Background(), inputA); err != nil { + t.Fatalf("expected session A to be allowed, got %v", err) + } + + _, err = manager.Execute(context.Background(), inputB) + if !errors.As(err, &permissionErr) || permissionErr.Decision() != "ask" { + t.Fatalf("expected session B remain ask, got %v", err) + } + }) + + t.Run("category matching shares decision across same tool category", func(t *testing.T) { + t.Parallel() + manager, _ := newAskManager(t) + inputA := ToolCallInput{ + ID: "call-target-a", + Name: "webfetch", + Arguments: []byte(`{"url":"https://example.com/a"}`), + SessionID: "session-target", + } + inputB := ToolCallInput{ + ID: "call-target-b", + Name: "webfetch", + Arguments: []byte(`{"url":"https://example.com/b"}`), + SessionID: "session-target", + } + + _, err := manager.Execute(context.Background(), inputA) + var permissionErr *PermissionDecisionError + if !errors.As(err, &permissionErr) { + t.Fatalf("expected permission ask on target A, got %v", err) + } + if rememberErr := manager.RememberSessionDecision(inputA.SessionID, permissionErr.Action(), SessionPermissionScopeAlways); rememberErr != nil { + t.Fatalf("remember target A: %v", rememberErr) + } + if _, err := manager.Execute(context.Background(), inputA); err != nil { + t.Fatalf("expected target A to be allowed, got %v", err) + } + + if _, err := manager.Execute(context.Background(), inputB); err != nil { + t.Fatalf("expected target B to inherit same-category allow, got %v", err) + } + }) + + t.Run("filesystem read category applies across file/grep/glob", func(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + readTool := &managerStubTool{name: "filesystem_read_file", content: "ok"} + grepTool := &managerStubTool{name: "filesystem_grep", content: "ok"} + globTool := &managerStubTool{name: "filesystem_glob", content: "ok"} + registry.Register(readTool) + registry.Register(grepTool) + registry.Register(globTool) + + engine, err := security.NewStaticGateway(security.DecisionAllow, []security.Rule{ + { + ID: "ask-filesystem-read", + Type: security.ActionTypeRead, + Resource: "filesystem_read_file", + Decision: security.DecisionAsk, + Reason: "requires approval", + }, + { + ID: "ask-filesystem-grep", + Type: security.ActionTypeRead, + Resource: "filesystem_grep", + Decision: security.DecisionAsk, + Reason: "requires approval", + }, + { + ID: "ask-filesystem-glob", + Type: security.ActionTypeRead, + Resource: "filesystem_glob", + Decision: security.DecisionAsk, + Reason: "requires approval", + }, + }) + if err != nil { + t.Fatalf("new engine: %v", err) + } + manager, err := NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new manager: %v", err) + } + + sessionID := "session-fs-read" + readInput := ToolCallInput{ + ID: "call-read", + Name: "filesystem_read_file", + Arguments: []byte(`{"path":"internal/README.md"}`), + SessionID: sessionID, + } + grepInput := ToolCallInput{ + ID: "call-grep", + Name: "filesystem_grep", + Arguments: []byte(`{"dir":"internal","pattern":"TODO"}`), + SessionID: sessionID, + } + globInput := ToolCallInput{ + ID: "call-glob", + Name: "filesystem_glob", + Arguments: []byte(`{"dir":"internal","pattern":"*.go"}`), + SessionID: sessionID, + } + + _, err = manager.Execute(context.Background(), readInput) + var permissionErr *PermissionDecisionError + if !errors.As(err, &permissionErr) || permissionErr.Decision() != "ask" { + t.Fatalf("expected initial read ask, got %v", err) + } + if rememberErr := manager.RememberSessionDecision(sessionID, permissionErr.Action(), SessionPermissionScopeAlways); rememberErr != nil { + t.Fatalf("remember filesystem read category: %v", rememberErr) + } + + if _, err := manager.Execute(context.Background(), grepInput); err != nil { + t.Fatalf("expected grep allow via filesystem_read category, got %v", err) + } + if _, err := manager.Execute(context.Background(), globInput); err != nil { + t.Fatalf("expected glob allow via filesystem_read category, got %v", err) + } + }) + + t.Run("remembered allow does not override hard deny", func(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + readTool := &managerStubTool{name: "filesystem_read_file", content: "ok"} + registry.Register(readTool) + + engine, err := security.NewStaticGateway(security.DecisionAllow, []security.Rule{ + { + ID: "deny-private-key", + Type: security.ActionTypeRead, + Resource: "filesystem_read_file", + Decision: security.DecisionDeny, + Reason: "private key blocked", + }, + }) + if err != nil { + t.Fatalf("new engine: %v", err) + } + manager, err := NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new manager: %v", err) + } + + sessionID := "session-deny-priority" + action := security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + Operation: "read_file", + TargetType: security.TargetTypePath, + Target: "README.md", + }, + } + if err := manager.RememberSessionDecision(sessionID, action, SessionPermissionScopeAlways); err != nil { + t.Fatalf("remember allow: %v", err) + } + + _, execErr := manager.Execute(context.Background(), ToolCallInput{ + ID: "call-deny-priority", + Name: "filesystem_read_file", + Arguments: []byte(`{"path":"C:/Users/test/.ssh/id_rsa"}`), + SessionID: sessionID, + }) + var permissionErr *PermissionDecisionError + if !errors.As(execErr, &permissionErr) { + t.Fatalf("expected permission error, got %v", execErr) + } + if permissionErr.Decision() != "deny" { + t.Fatalf("expected hard deny to win, got %q", permissionErr.Decision()) + } + if permissionErr.RuleID() != "deny-private-key" { + t.Fatalf("expected deny rule id, got %q", permissionErr.RuleID()) + } + if readTool.callCount != 0 { + t.Fatalf("expected blocked call not to execute tool, got %d", readTool.callCount) + } + }) +} + func TestBuildPermissionAction(t *testing.T) { t.Parallel() diff --git a/internal/tools/registry.go b/internal/tools/registry.go index c057ee86..b512c323 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -7,6 +7,7 @@ import ( "strings" "neo-code/internal/provider" + "neo-code/internal/security" ) type Registry struct { @@ -94,3 +95,8 @@ func (r *Registry) Execute(ctx context.Context, input ToolCallInput) (ToolResult } return result, nil } + +// RememberSessionDecision 对纯 Registry 管理器不生效，保留接口以满足 runtime 依赖。 +func (r *Registry) RememberSessionDecision(sessionID string, action security.Action, scope SessionPermissionScope) error { + return errors.New("tools: session permission memory is unsupported by registry manager") +} diff --git a/internal/tools/registry_test.go b/internal/tools/registry_test.go index e8d2db35..c6cd0ad1 100644 --- a/internal/tools/registry_test.go +++ b/internal/tools/registry_test.go @@ -5,6 +5,8 @@ import ( "errors" "strings" "testing" + + "neo-code/internal/security" ) type stubTool struct { @@ -180,3 +182,16 @@ func TestRegistryHelpers(t *testing.T) { t.Fatalf("expected context canceled, got %v", err) } } + +func TestRegistryRememberSessionDecisionUnsupported(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + err := registry.RememberSessionDecision("session-1", security.Action{}, SessionPermissionScopeAlways) + if err == nil { + t.Fatalf("expected unsupported error") + } + if !strings.Contains(err.Error(), "unsupported") { + t.Fatalf("expected unsupported error, got %v", err) + } +} diff --git a/internal/tools/session_memory.go b/internal/tools/session_memory.go new file mode 100644 index 00000000..91006a14 --- /dev/null +++ b/internal/tools/session_memory.go @@ -0,0 +1,209 @@ +package tools + +import ( + "errors" + "fmt" + "net/url" + "path/filepath" + "strings" + "sync" + + "neo-code/internal/security" +) + +// SessionPermissionScope 表示 session 级权限记忆的作用范围。 +type SessionPermissionScope string + +const ( + // SessionPermissionScopeOnce 表示仅当前一次请求放行。 + SessionPermissionScopeOnce SessionPermissionScope = "once" + // SessionPermissionScopeAlways 表示当前会话内同类请求持续放行。 + SessionPermissionScopeAlways SessionPermissionScope = "always_session" + // SessionPermissionScopeReject 表示当前会话内同类请求持续拒绝。 + SessionPermissionScopeReject SessionPermissionScope = "reject" +) + +type sessionPermissionEntry struct { + decision security.Decision + scope SessionPermissionScope + remaining int +} + +// sessionPermissionMemory 管理按 session/action 维度的审批记忆。 +type sessionPermissionMemory struct { + mu sync.Mutex + entries map[string]map[string]sessionPermissionEntry +} + +// newSessionPermissionMemory 创建 session 级权限记忆存储。 +func newSessionPermissionMemory() *sessionPermissionMemory { + return &sessionPermissionMemory{ + entries: make(map[string]map[string]sessionPermissionEntry), + } +} + +// remember 记录一条 session 级权限决策。 +func (m *sessionPermissionMemory) remember(sessionID string, action security.Action, scope SessionPermissionScope) error { + trimmedSessionID := strings.TrimSpace(sessionID) + if trimmedSessionID == "" { + return errors.New("tools: session id is empty") + } + if err := action.Validate(); err != nil { + return err + } + + var entry sessionPermissionEntry + switch scope { + case SessionPermissionScopeOnce: + entry = sessionPermissionEntry{ + decision: security.DecisionAllow, + scope: scope, + remaining: 1, + } + case SessionPermissionScopeAlways: + entry = sessionPermissionEntry{ + decision: security.DecisionAllow, + scope: scope, + remaining: -1, + } + case SessionPermissionScopeReject: + entry = sessionPermissionEntry{ + decision: security.DecisionDeny, + scope: scope, + remaining: -1, + } + default: + return fmt.Errorf("tools: unsupported session permission scope %q", scope) + } + + actionKey := sessionPermissionActionKey(action) + m.mu.Lock() + defer m.mu.Unlock() + sessionEntries, ok := m.entries[trimmedSessionID] + if !ok { + sessionEntries = make(map[string]sessionPermissionEntry) + m.entries[trimmedSessionID] = sessionEntries + } + sessionEntries[actionKey] = entry + return nil +} + +// resolve 查询并按 scope 规则消费 session 级权限记忆。 +func (m *sessionPermissionMemory) resolve(sessionID string, action security.Action) (security.Decision, SessionPermissionScope, bool) { + trimmedSessionID := strings.TrimSpace(sessionID) + if trimmedSessionID == "" { + return "", "", false + } + actionKey := sessionPermissionActionKey(action) + + m.mu.Lock() + defer m.mu.Unlock() + + sessionEntries, ok := m.entries[trimmedSessionID] + if !ok { + return "", "", false + } + entry, ok := sessionEntries[actionKey] + if !ok { + return "", "", false + } + + if entry.scope == SessionPermissionScopeOnce && entry.remaining > 0 { + entry.remaining-- + if entry.remaining <= 0 { + delete(sessionEntries, actionKey) + } else { + sessionEntries[actionKey] = entry + } + } + + if len(sessionEntries) == 0 { + delete(m.entries, trimmedSessionID) + } + + return entry.decision, entry.scope, true +} + +// sessionPermissionActionKey 基于结构化 action 生成稳定匹配键。 +func sessionPermissionActionKey(action security.Action) string { + return strings.Join([]string{ + string(action.Type), + sessionPermissionCategory(action), + sessionPermissionTargetScope(action), + }, "|") +} + +// sessionPermissionCategory 将安全动作归一为稳定的工具类别。 +// 类别用于聚合同类工具，再配合 target scope 控制最小授权范围。 +func sessionPermissionCategory(action security.Action) string { + resource := strings.ToLower(strings.TrimSpace(action.Payload.Resource)) + switch action.Type { + case security.ActionTypeRead: + if strings.HasPrefix(resource, "filesystem_") { + return "filesystem_read" + } + if resource == "webfetch" { + return "webfetch" + } + case security.ActionTypeWrite: + if strings.HasPrefix(resource, "filesystem_") { + return "filesystem_write" + } + case security.ActionTypeBash: + return "bash" + case security.ActionTypeMCP: + target := strings.ToLower(strings.TrimSpace(action.Payload.Target)) + if target != "" { + return "mcp:" + target + } + return "mcp" + } + + toolName := strings.ToLower(strings.TrimSpace(action.Payload.ToolName)) + if toolName != "" { + return toolName + } + return resource +} + +// sessionPermissionTargetScope 基于 action 的 target 生成最小授权范围键。 +func sessionPermissionTargetScope(action security.Action) string { + target := strings.TrimSpace(action.Payload.Target) + if target == "" { + return "*" + } + + switch action.Payload.TargetType { + case security.TargetTypeURL: + return normalizePermissionURLTarget(target) + case security.TargetTypePath: + return normalizePermissionPathTarget(filepath.Dir(target)) + case security.TargetTypeDirectory: + return normalizePermissionPathTarget(target) + default: + return strings.ToLower(target) + } +} + +// normalizePermissionURLTarget 将 URL 归一到 host[:port] 维度。 +func normalizePermissionURLTarget(raw string) string { + parsed, err := url.Parse(strings.TrimSpace(raw)) + if err != nil || strings.TrimSpace(parsed.Host) == "" { + return strings.ToLower(strings.TrimSpace(raw)) + } + + host := strings.ToLower(strings.TrimSpace(parsed.Hostname())) + if port := strings.TrimSpace(parsed.Port()); port != "" { + host += ":" + port + } + return host +} + +// normalizePermissionPathTarget 统一路径分隔并按平台无关形式生成匹配键。 +func normalizePermissionPathTarget(raw string) string { + cleaned := filepath.Clean(strings.TrimSpace(raw)) + if cleaned == "." || cleaned == "" { + return "." + } + return strings.ToLower(filepath.ToSlash(cleaned)) +} diff --git a/internal/tools/session_memory_test.go b/internal/tools/session_memory_test.go new file mode 100644 index 00000000..194b316e --- /dev/null +++ b/internal/tools/session_memory_test.go @@ -0,0 +1,283 @@ +package tools + +import ( + "strings" + "testing" + + "neo-code/internal/security" +) + +func TestSessionPermissionMemoryRememberAndResolve(t *testing.T) { + t.Parallel() + + action := security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + }, + } + + t.Run("once decision is consumed after first resolve", func(t *testing.T) { + t.Parallel() + + memory := newSessionPermissionMemory() + if err := memory.remember("session-1", action, SessionPermissionScopeOnce); err != nil { + t.Fatalf("remember() error = %v", err) + } + + decision, scope, ok := memory.resolve("session-1", action) + if !ok || decision != security.DecisionAllow || scope != SessionPermissionScopeOnce { + t.Fatalf("expected once allow decision, got decision=%q scope=%q ok=%v", decision, scope, ok) + } + + _, _, ok = memory.resolve("session-1", action) + if ok { + t.Fatalf("expected once decision to be consumed") + } + }) + + t.Run("always decision keeps applying", func(t *testing.T) { + t.Parallel() + + memory := newSessionPermissionMemory() + if err := memory.remember("session-2", action, SessionPermissionScopeAlways); err != nil { + t.Fatalf("remember() error = %v", err) + } + + for i := 0; i < 2; i++ { + decision, scope, ok := memory.resolve("session-2", action) + if !ok || decision != security.DecisionAllow || scope != SessionPermissionScopeAlways { + t.Fatalf("expected always allow decision, got decision=%q scope=%q ok=%v", decision, scope, ok) + } + } + }) + + t.Run("reject decision keeps applying", func(t *testing.T) { + t.Parallel() + + memory := newSessionPermissionMemory() + if err := memory.remember("session-3", action, SessionPermissionScopeReject); err != nil { + t.Fatalf("remember() error = %v", err) + } + + decision, scope, ok := memory.resolve("session-3", action) + if !ok || decision != security.DecisionDeny || scope != SessionPermissionScopeReject { + t.Fatalf("expected reject deny decision, got decision=%q scope=%q ok=%v", decision, scope, ok) + } + }) +} + +func TestSessionPermissionMemoryValidationAndMisses(t *testing.T) { + t.Parallel() + + memory := newSessionPermissionMemory() + validAction := security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + }, + } + + if err := memory.remember(" ", validAction, SessionPermissionScopeAlways); err == nil { + t.Fatalf("expected empty session id error") + } + + invalidAction := security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "", + Resource: "webfetch", + }, + } + if err := memory.remember("session", invalidAction, SessionPermissionScopeAlways); err == nil { + t.Fatalf("expected invalid action error") + } + + if err := memory.remember("session", validAction, SessionPermissionScope("bad")); err == nil { + t.Fatalf("expected unsupported scope error") + } + + if _, _, ok := memory.resolve(" ", validAction); ok { + t.Fatalf("expected empty session resolve miss") + } + if _, _, ok := memory.resolve("missing", validAction); ok { + t.Fatalf("expected missing session resolve miss") + } +} + +func TestSessionPermissionCategoryAndActionKey(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + action security.Action + expected string + }{ + { + name: "filesystem read category", + action: security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + Resource: "filesystem_grep", + }, + }, + expected: "filesystem_read", + }, + { + name: "filesystem write category", + action: security.Action{ + Type: security.ActionTypeWrite, + Payload: security.ActionPayload{ + Resource: "filesystem_edit", + }, + }, + expected: "filesystem_write", + }, + { + name: "webfetch category", + action: security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + Resource: "webfetch", + }, + }, + expected: "webfetch", + }, + { + name: "bash category", + action: security.Action{ + Type: security.ActionTypeBash, + }, + expected: "bash", + }, + { + name: "mcp with target", + action: security.Action{ + Type: security.ActionTypeMCP, + Payload: security.ActionPayload{ + Target: "Server-A", + }, + }, + expected: "mcp:server-a", + }, + { + name: "mcp without target", + action: security.Action{ + Type: security.ActionTypeMCP, + }, + expected: "mcp", + }, + { + name: "fallback to tool name", + action: security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "CustomTool", + Resource: "other_resource", + }, + }, + expected: "customtool", + }, + { + name: "fallback to resource", + action: security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + Resource: "custom_resource", + }, + }, + expected: "custom_resource", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got := sessionPermissionCategory(tt.action) + if got != tt.expected { + t.Fatalf("sessionPermissionCategory() = %q, want %q", got, tt.expected) + } + + key := sessionPermissionActionKey(tt.action) + if !strings.Contains(key, "|"+tt.expected+"|") { + t.Fatalf("sessionPermissionActionKey() = %q, expected category token %q", key, "|"+tt.expected+"|") + } + }) + } +} + +func TestSessionPermissionMemoryResolveRequiresTargetScopeMatch(t *testing.T) { + t.Parallel() + + memory := newSessionPermissionMemory() + sessionID := "session-target-scope" + + webAction := security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + TargetType: security.TargetTypeURL, + Target: "https://docs.github.com/en/rest", + }, + } + if err := memory.remember(sessionID, webAction, SessionPermissionScopeAlways); err != nil { + t.Fatalf("remember web action: %v", err) + } + + sameHost := security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + TargetType: security.TargetTypeURL, + Target: "https://docs.github.com/en/actions", + }, + } + if _, _, ok := memory.resolve(sessionID, sameHost); !ok { + t.Fatalf("expected same host/path scope web action to hit memory") + } + + differentHost := security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + TargetType: security.TargetTypeURL, + Target: "https://example.com/en/actions", + }, + } + if _, _, ok := memory.resolve(sessionID, differentHost); ok { + t.Fatalf("expected different host web action to miss memory") + } + + fileAction := security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + TargetType: security.TargetTypePath, + Target: "src/main.go", + }, + } + if err := memory.remember(sessionID, fileAction, SessionPermissionScopeAlways); err != nil { + t.Fatalf("remember file action: %v", err) + } + + otherFile := security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + TargetType: security.TargetTypePath, + Target: "secrets/secret.key", + }, + } + if _, _, ok := memory.resolve(sessionID, otherFile); ok { + t.Fatalf("expected different path file action to miss memory") + } +} diff --git a/internal/tui/app.go b/internal/tui/app.go index 2fdfaf81..eae67f5d 100644 --- a/internal/tui/app.go +++ b/internal/tui/app.go @@ -20,43 +20,44 @@ import ( ) type App struct { - state UIState - configManager *config.Manager - providerSvc ProviderController - runtime agentruntime.Runtime - keys keyMap - help help.Model - spinner spinner.Model - sessions list.Model - commandMenu list.Model - commandMenuMeta commandMenuMeta - providerPicker list.Model - modelPicker list.Model - fileBrowser filepicker.Model - progress progress.Model - transcript viewport.Model - activity viewport.Model - input textarea.Model - markdownRenderer markdownContentRenderer - codeCopyBlocks map[int]string - pendingCopyID int - nowFn func() time.Time - lastInputEditAt time.Time - lastPasteLikeAt time.Time - inputBurstStart time.Time - inputBurstCount int - pasteMode bool - activeMessages []provider.Message - activities []activityEntry - fileCandidates []string - modelRefreshID string - focus panel - runProgressValue float64 - runProgressKnown bool - runProgressLabel string - width int - height int - styles styles + state UIState + configManager *config.Manager + providerSvc ProviderController + runtime agentruntime.Runtime + keys keyMap + help help.Model + spinner spinner.Model + sessions list.Model + commandMenu list.Model + commandMenuMeta commandMenuMeta + providerPicker list.Model + modelPicker list.Model + fileBrowser filepicker.Model + progress progress.Model + transcript viewport.Model + activity viewport.Model + input textarea.Model + markdownRenderer markdownContentRenderer + codeCopyBlocks map[int]string + pendingCopyID int + nowFn func() time.Time + lastInputEditAt time.Time + lastPasteLikeAt time.Time + inputBurstStart time.Time + inputBurstCount int + pasteMode bool + pendingPermission *pendingPermissionPrompt + activeMessages []provider.Message + activities []activityEntry + fileCandidates []string + modelRefreshID string + focus panel + runProgressValue float64 + runProgressKnown bool + runProgressLabel string + width int + height int + styles styles } func New(cfg *config.Config, configManager *config.Manager, runtime agentruntime.Runtime, providerSvc ProviderController) (App, error) { diff --git a/internal/tui/state.go b/internal/tui/state.go index a4200e1c..1f939dad 100644 --- a/internal/tui/state.go +++ b/internal/tui/state.go @@ -58,6 +58,15 @@ type activityEntry struct { IsError bool } +type pendingPermissionPrompt struct { + RequestID string + ToolCallID string + ToolName string + ToolCategory string + Target string + Submitted bool +} + type commandMenuMeta struct { Title string } diff --git a/internal/tui/update.go b/internal/tui/update.go index db804480..4836fd28 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -24,6 +24,11 @@ import ( type RuntimeMsg struct{ Event agentruntime.RuntimeEvent } type RuntimeClosedMsg struct{} type runFinishedMsg struct{ err error } +type permissionResolveResultMsg struct { + requestID string + decision agentruntime.PermissionResolutionDecision + err error +} type modelCatalogRefreshMsg struct { providerID string models []config.ModelDescriptor @@ -91,6 +96,7 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case runFinishedMsg: if typed.err != nil { a.state.IsAgentRunning = false + a.pendingPermission = nil a.clearRunProgress() a.state.StreamingReply = false a.state.CurrentTool = "" @@ -108,6 +114,20 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { _ = a.refreshSessions() a.syncActiveSessionTitle() return a, tea.Batch(cmds...) + case permissionResolveResultMsg: + if typed.err != nil { + if a.pendingPermission != nil && strings.EqualFold(strings.TrimSpace(a.pendingPermission.RequestID), strings.TrimSpace(typed.requestID)) { + a.pendingPermission.Submitted = false + } + a.state.ExecutionError = typed.err.Error() + a.state.StatusText = typed.err.Error() + a.appendActivity("permission", "Submit permission failed", typed.err.Error(), true) + return a, tea.Batch(cmds...) + } + a.state.ExecutionError = "" + a.state.StatusText = "Permission decision submitted" + a.appendActivity("permission", "Submitted permission decision", string(typed.decision), false) + return a, tea.Batch(cmds...) case modelCatalogRefreshMsg: if strings.EqualFold(a.modelRefreshID, typed.providerID) { a.modelRefreshID = "" @@ -236,6 +256,14 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { a.applyComponentLayout(true) return a, tea.Batch(cmds...) } + if a.pendingPermission != nil { + if permissionCmd, handled := a.handlePermissionDecisionKey(typed); handled { + if permissionCmd != nil { + cmds = append(cmds, permissionCmd) + } + return a, tea.Batch(cmds...) + } + } if a.state.IsAgentRunning && key.Matches(typed, a.keys.CancelAgent) { if a.runtime.CancelActiveRun() { a.state.StatusText = statusCanceling @@ -725,6 +753,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { a.state.IsAgentRunning = false a.state.StreamingReply = false a.state.CurrentTool = "" + a.pendingPermission = nil a.clearRunProgress() if strings.TrimSpace(a.state.ExecutionError) == "" { a.state.StatusText = statusReady @@ -737,6 +766,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { a.state.IsAgentRunning = false a.state.StreamingReply = false a.state.CurrentTool = "" + a.pendingPermission = nil a.state.ExecutionError = "" a.state.StatusText = statusCanceled a.clearRunProgress() @@ -746,6 +776,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { a.state.IsAgentRunning = false a.state.StreamingReply = false a.state.CurrentTool = "" + a.pendingPermission = nil a.clearRunProgress() if payload, ok := event.Payload.(string); ok { a.state.ExecutionError = payload @@ -758,6 +789,47 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { a.runProgressKnown = false a.appendActivity("provider", "Retrying provider call", payload, false) } + case agentruntime.EventPermissionRequest: + payload, ok := event.Payload.(agentruntime.PermissionRequestPayload) + if !ok { + return transcriptDirty + } + a.pendingPermission = &pendingPermissionPrompt{ + RequestID: strings.TrimSpace(payload.RequestID), + ToolCallID: strings.TrimSpace(payload.ToolCallID), + ToolName: strings.TrimSpace(payload.ToolName), + ToolCategory: strings.TrimSpace(payload.ToolCategory), + Target: strings.TrimSpace(payload.Target), + } + prompt := fmt.Sprintf( + "[Permission] Tool=%s Category=%s Target=%s | y=once, a=always(session), n=reject(session)", + fallback(payload.ToolName, "-"), + fallback(payload.ToolCategory, "-"), + fallback(payload.Target, "-"), + ) + a.state.StatusText = "Permission required (y/a/n)" + a.appendActivity("permission", "Awaiting permission decision", prompt, false) + a.appendInlineMessage(roleSystem, prompt) + transcriptDirty = true + case agentruntime.EventPermissionResolved: + payload, ok := event.Payload.(agentruntime.PermissionResolvedPayload) + if !ok { + return transcriptDirty + } + if a.pendingPermission != nil && strings.TrimSpace(a.pendingPermission.RequestID) != "" && + strings.EqualFold(a.pendingPermission.RequestID, strings.TrimSpace(payload.RequestID)) { + a.pendingPermission = nil + } + resolved := fmt.Sprintf( + "[Permission] %s %s (%s)", + fallback(payload.ToolName, "-"), + fallback(payload.ResolvedAs, "-"), + fallback(payload.RememberScope, "-"), + ) + a.appendActivity("permission", "Permission resolved", resolved, strings.EqualFold(payload.ResolvedAs, "rejected")) + if strings.EqualFold(payload.ResolvedAs, "approved") { + a.state.StatusText = "Permission approved, executing tool..." + } case agentruntime.EventCompactDone: payload, ok := event.Payload.(agentruntime.CompactDonePayload) if !ok { @@ -1396,6 +1468,36 @@ func ListenForRuntimeEvent(sub <-chan agentruntime.RuntimeEvent) tea.Cmd { } } +// handlePermissionDecisionKey 处理待审批状态下的快捷授权键位。 +func (a *App) handlePermissionDecisionKey(msg tea.KeyMsg) (tea.Cmd, bool) { + if a.pendingPermission == nil { + return nil, false + } + + keyText := strings.ToLower(strings.TrimSpace(msg.String())) + var decision agentruntime.PermissionResolutionDecision + switch keyText { + case "y": + decision = agentruntime.PermissionResolutionAllowOnce + case "a": + decision = agentruntime.PermissionResolutionAllowSession + case "n": + decision = agentruntime.PermissionResolutionReject + default: + return nil, false + } + if a.pendingPermission.Submitted { + a.state.StatusText = "Permission decision already submitted, waiting runtime..." + return nil, true + } + + requestID := strings.TrimSpace(a.pendingPermission.RequestID) + a.pendingPermission.Submitted = true + a.state.StatusText = "Submitting permission decision..." + a.state.ExecutionError = "" + return runResolvePermission(a.runtime, requestID, decision), true +} + func runAgent(runtime agentruntime.Runtime, sessionID string, workdir string, content string) tea.Cmd { return func() tea.Msg { err := runtime.Run(context.Background(), agentruntime.UserInput{ @@ -1407,6 +1509,25 @@ func runAgent(runtime agentruntime.Runtime, sessionID string, workdir string, co } } +// runResolvePermission 在独立命令中提交权限审批决定，避免阻塞 UI 事件循环。 +func runResolvePermission( + runtime agentruntime.Runtime, + requestID string, + decision agentruntime.PermissionResolutionDecision, +) tea.Cmd { + return func() tea.Msg { + err := runtime.ResolvePermission(context.Background(), agentruntime.PermissionResolutionInput{ + RequestID: requestID, + Decision: decision, + }) + return permissionResolveResultMsg{ + requestID: requestID, + decision: decision, + err: err, + } + } +} + func runSessionWorkdirCommand( runtime agentruntime.Runtime, sessionID string, diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index ed307cc0..85f6b1be 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -38,6 +38,8 @@ type stubRuntime struct { setWorkdirErr error setResult *agentruntime.Session setCalls int + resolveInputs []agentruntime.PermissionResolutionInput + resolveErr error cancelCalls int cancelResult bool } @@ -78,6 +80,11 @@ func (r *stubRuntime) Compact(ctx context.Context, input agentruntime.CompactInp return r.compactResult, r.compactErr } +func (r *stubRuntime) ResolvePermission(ctx context.Context, input agentruntime.PermissionResolutionInput) error { + r.resolveInputs = append(r.resolveInputs, input) + return r.resolveErr +} + func (r *stubRuntime) Events() <-chan agentruntime.RuntimeEvent { return r.events } @@ -489,6 +496,131 @@ func TestRunAgentWorkdirForwarding(t *testing.T) { }) } +func TestHandlePermissionDecisionKey(t *testing.T) { + t.Parallel() + + manager := newTestConfigManager(t) + runtime := newStubRuntime() + app, err := New(nil, manager, runtime, newTestProviderService(t, manager)) + if err != nil { + t.Fatalf("New() error = %v", err) + } + app.pendingPermission = &pendingPermissionPrompt{ + RequestID: "perm-1", + ToolName: "webfetch", + } + + tests := []struct { + name string + key tea.KeyMsg + wantSent agentruntime.PermissionResolutionDecision + handled bool + }{ + { + name: "y maps to allow once", + key: tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'y'}}, + wantSent: agentruntime.PermissionResolutionAllowOnce, + handled: true, + }, + { + name: "a maps to allow session", + key: tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'a'}}, + wantSent: agentruntime.PermissionResolutionAllowSession, + handled: true, + }, + { + name: "n maps to reject", + key: tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'n'}}, + wantSent: agentruntime.PermissionResolutionReject, + handled: true, + }, + { + name: "other key ignored", + key: tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'x'}}, + handled: false, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + app.pendingPermission.Submitted = false + runtime.resolveInputs = nil + + cmd, handled := app.handlePermissionDecisionKey(tt.key) + if handled != tt.handled { + t.Fatalf("expected handled=%v, got %v", tt.handled, handled) + } + if !tt.handled { + if cmd != nil { + t.Fatalf("expected nil cmd for unhandled key") + } + return + } + if cmd == nil { + t.Fatalf("expected resolve cmd") + } + msg := cmd() + result, ok := msg.(permissionResolveResultMsg) + if !ok { + t.Fatalf("expected permissionResolveResultMsg, got %T", msg) + } + if result.err != nil { + t.Fatalf("expected nil resolve error, got %v", result.err) + } + if len(runtime.resolveInputs) == 0 { + t.Fatalf("expected runtime resolve inputs") + } + last := runtime.resolveInputs[len(runtime.resolveInputs)-1] + if last.Decision != tt.wantSent { + t.Fatalf("expected decision %q, got %q", tt.wantSent, last.Decision) + } + if strings.TrimSpace(last.RequestID) != "perm-1" { + t.Fatalf("expected request id perm-1, got %q", last.RequestID) + } + }) + } +} + +func TestHandlePermissionDecisionKeyIgnoresRepeatedSubmission(t *testing.T) { + t.Parallel() + + manager := newTestConfigManager(t) + runtime := newStubRuntime() + app, err := New(nil, manager, runtime, newTestProviderService(t, manager)) + if err != nil { + t.Fatalf("New() error = %v", err) + } + app.pendingPermission = &pendingPermissionPrompt{ + RequestID: "perm-repeat", + ToolName: "webfetch", + } + + cmd, handled := app.handlePermissionDecisionKey(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'y'}}) + if !handled || cmd == nil { + t.Fatalf("expected first permission key to be handled with cmd") + } + msg := cmd() + result, ok := msg.(permissionResolveResultMsg) + if !ok || result.err != nil { + t.Fatalf("expected successful permissionResolveResultMsg, got %#v", msg) + } + if len(runtime.resolveInputs) != 1 { + t.Fatalf("expected one resolve call after first submission, got %d", len(runtime.resolveInputs)) + } + + cmd, handled = app.handlePermissionDecisionKey(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'a'}}) + if !handled { + t.Fatalf("expected repeated permission key to be consumed") + } + if cmd != nil { + t.Fatalf("expected repeated submission to skip runtime command") + } + if len(runtime.resolveInputs) != 1 { + t.Fatalf("expected resolve call count unchanged after repeat, got %d", len(runtime.resolveInputs)) + } +} + func TestAppUpdateModelPickerAndRuntimeMessages(t *testing.T) { tests := []struct { name string