diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 4203b881..3578c1da 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -11,6 +11,7 @@ import ( "neo-code/internal/provider" "neo-code/internal/provider/builtin" agentruntime "neo-code/internal/runtime" + "neo-code/internal/security" "neo-code/internal/tools" "neo-code/internal/tools/bash" "neo-code/internal/tools/filesystem" @@ -27,6 +28,10 @@ func NewProgram(ctx context.Context) (*tea.Program, error) { } toolRegistry := buildToolRegistry(cfg) + toolManager, err := buildToolManager(toolRegistry) + if err != nil { + return nil, err + } providerRegistry, err := builtin.NewRegistry() if err != nil { @@ -35,7 +40,13 @@ func NewProgram(ctx context.Context) (*tea.Program, error) { providerService := provider.NewService(manager, providerRegistry) sessionStore := agentruntime.NewSessionStore(loader.BaseDir()) - runtimeSvc := agentruntime.NewWithFactory(manager, toolRegistry, sessionStore, providerService, agentcontext.NewBuilder()) + runtimeSvc := agentruntime.NewWithFactory( + manager, + toolManager, + sessionStore, + providerService, + agentcontext.NewBuilder(), + ) tuiApp, err := tui.New(&cfg, manager, runtimeSvc, providerService) if err != nil { @@ -63,3 +74,11 @@ func buildToolRegistry(cfg config.Config) *tools.Registry { })) return toolRegistry } + +func buildToolManager(registry *tools.Registry) (tools.Manager, error) { + engine, err := security.NewStaticGateway(security.DecisionAllow, nil) + if err != nil { + return nil, err + } + return tools.NewManager(registry, engine, nil) +} diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index 51dc688a..14281848 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -70,3 +70,48 @@ func TestBuildToolRegistryUsesWebFetchConfig(t *testing.T) { t.Fatalf("expected formatted webfetch content") } } + +func TestBuildToolManagerWrapsRegistry(t *testing.T) { + t.Parallel() + + registry := tools.NewRegistry() + registry.Register(stubToolForBootstrap{name: "bash", content: "ok"}) + manager, err := buildToolManager(registry) + if err != nil { + t.Fatalf("buildToolManager() error = %v", err) + } + if manager == nil { + t.Fatalf("expected tool manager") + } + + specs, err := manager.ListAvailableSpecs(context.Background(), tools.SpecListInput{}) + if err != nil { + t.Fatalf("ListAvailableSpecs() error = %v", err) + } + if len(specs) != 1 { + t.Fatalf("expected 1 spec, got %+v", specs) + } + + result, execErr := manager.Execute(context.Background(), tools.ToolCallInput{ + Name: "bash", + Arguments: []byte(`{"command":"echo hi"}`), + }) + if execErr != nil { + t.Fatalf("Execute() error = %v", execErr) + } + if result.Content != "ok" { + t.Fatalf("expected ok result, got %+v", result) + } +} + +type stubToolForBootstrap struct { + name string + content string +} + +func (s stubToolForBootstrap) Name() string { return s.name } +func (s stubToolForBootstrap) Description() string { return "stub" } +func (s stubToolForBootstrap) Schema() map[string]any { return map[string]any{"type": "object"} } +func (s stubToolForBootstrap) Execute(ctx context.Context, call tools.ToolCallInput) (tools.ToolResult, error) { + return tools.ToolResult{Name: s.name, Content: s.content}, nil +} diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 6e0da509..65aa67b7 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -49,7 +49,7 @@ type ProviderFactory interface { type Service struct { configManager *config.Manager // 配置管理器，提供当前选中的 provider、model、workdir 等配置读取能力。 sessionStore Store // 会话持久化接口，负责保存和加载聊天会话。 - toolRegistry *tools.Registry // 工具注册表，维护所有工具的 schema。 + toolManager tools.Manager // 工具管理器，统一工具 schema 暴露与执行入口。 providerFactory ProviderFactory // Provider 工厂接口，根据配置动态创建具体的 provider 实例。 contextBuilder agentcontext.Builder // 上下文构建器，负责组装 system prompt 与本轮发给模型的消息上下文。 events chan RuntimeEvent // 事件通道，Runtime 在运行过程中产生的事件都通过该通道发送给 TUI 层消费和展示。 @@ -61,7 +61,7 @@ type Service struct { func NewWithFactory( configManager *config.Manager, - toolRegistry *tools.Registry, + toolManager tools.Manager, sessionStore Store, providerFactory ProviderFactory, contextBuilder agentcontext.Builder, @@ -69,6 +69,9 @@ func NewWithFactory( if providerFactory == nil { providerFactory = provider.NewRegistry() } + if toolManager == nil { + toolManager = tools.NewRegistry() + } if contextBuilder == nil { contextBuilder = agentcontext.NewBuilder() } @@ -76,7 +79,7 @@ func NewWithFactory( return &Service{ configManager: configManager, sessionStore: sessionStore, - toolRegistry: toolRegistry, + toolManager: toolManager, providerFactory: providerFactory, contextBuilder: contextBuilder, events: make(chan RuntimeEvent, 128), @@ -136,12 +139,18 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { return s.handleRunError(ctx, input.RunID, session.ID, err) } - resp, err := s.callProviderWithRetry(ctx, input.RunID, session.ID, provider.ChatRequest{ + toolSpecs, err := s.toolManager.ListAvailableSpecs(ctx, tools.SpecListInput{ + SessionID: session.ID, + }) + if err != nil { + return s.handleRunError(ctx, input.RunID, session.ID, err) + } + resp, err := s.callProviderWithRetry(ctx, input.RunID, session.ID, provider.ChatRequest{ Model: cfg.CurrentModel, SystemPrompt: builtContext.SystemPrompt, Messages: builtContext.Messages, - Tools: s.toolRegistry.GetSpecs(), + Tools: toolSpecs, }) if err != nil { return s.handleRunError(ctx, input.RunID, session.ID, err) @@ -178,7 +187,7 @@ 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.toolRegistry.Execute(runCtx, tools.ToolCallInput{ + result, execErr := s.toolManager.Execute(runCtx, tools.ToolCallInput{ ID: call.ID, Name: call.Name, Arguments: []byte(call.Arguments), diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index 4f3eb869..9547607f 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -7,6 +7,7 @@ import ( "os" "strings" "testing" + "time" "neo-code/internal/config" agentcontext "neo-code/internal/context" @@ -192,6 +193,37 @@ func (b *stubContextBuilder) Build(ctx context.Context, input agentcontext.Build }, nil } +type stubToolManager struct { + specs []provider.ToolSpec + result tools.ToolResult + err error + listErr error + listCalls int + executeCalls int + lastInput tools.ToolCallInput +} + +func (m *stubToolManager) ListAvailableSpecs(ctx context.Context, input tools.SpecListInput) ([]provider.ToolSpec, error) { + m.listCalls++ + if err := ctx.Err(); err != nil { + return nil, err + } + if m.listErr != nil { + return nil, m.listErr + } + return append([]provider.ToolSpec(nil), m.specs...), nil +} + +func (m *stubToolManager) Execute(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + m.executeCalls++ + m.lastInput = input + result := m.result + if result.Name == "" { + result.Name = input.Name + } + return result, m.err +} + func TestServiceRun(t *testing.T) { tests := []struct { name string @@ -430,6 +462,125 @@ func TestServiceRunDelegatesToContextBuilder(t *testing.T) { } } +func TestServiceRunUsesToolManager(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + toolManager := &stubToolManager{ + specs: []provider.ToolSpec{ + {Name: "filesystem_edit", Description: "stub", Schema: map[string]any{"type": "object"}}, + }, + result: tools.ToolResult{ + Name: "filesystem_edit", + Content: "tool manager output", + }, + } + + scripted := &scriptedProvider{ + responses: []provider.ChatResponse{ + { + Message: provider.Message{ + Role: "assistant", + ToolCalls: []provider.ToolCall{ + {ID: "call-manager", Name: "filesystem_edit", Arguments: `{"path":"main.go"}`}, + }, + }, + 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{RunID: "run-tool-manager", Content: "edit file"}); err != nil { + t.Fatalf("Run() error = %v", err) + } + + if toolManager.listCalls != 2 { + t.Fatalf("expected 2 spec list calls, got %d", toolManager.listCalls) + } + if toolManager.executeCalls != 1 { + t.Fatalf("expected 1 execute call, got %d", toolManager.executeCalls) + } + if toolManager.lastInput.ID != "call-manager" { + t.Fatalf("expected forwarded tool call id, got %q", toolManager.lastInput.ID) + } + if len(scripted.requests) == 0 || len(scripted.requests[0].Tools) != 1 || scripted.requests[0].Tools[0].Name != "filesystem_edit" { + t.Fatalf("expected tool specs from tool manager, got %+v", scripted.requests) + } + + session := onlySession(t, store) + foundToolMessage := false + for _, message := range session.Messages { + if message.Role == provider.RoleTool && message.Content == "tool manager output" { + foundToolMessage = true + break + } + } + if !foundToolMessage { + t.Fatalf("expected tool manager result in session messages, got %+v", session.Messages) + } +} + +func TestServiceRunHandlesToolManagerSpecError(t *testing.T) { + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + toolManager := &stubToolManager{ + listErr: errors.New("tool specs unavailable"), + } + + service := NewWithFactory(manager, toolManager, store, &scriptedProviderFactory{ + provider: &scriptedProvider{}, + }, nil) + input := UserInput{RunID: "run-tool-spec-error", Content: "hello"} + err := service.Run(context.Background(), input) + if err == nil || !containsError(err, "tool specs unavailable") { + t.Fatalf("expected tool spec error, got %v", err) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventUserMessage, EventError}) + assertNoEventType(t, events, EventAgentDone) + assertEventsRunID(t, events, input.RunID) + + session := onlySession(t, store) + if len(session.Messages) != 1 || session.Messages[0].Role != provider.RoleUser { + t.Fatalf("expected only user message to persist, got %+v", session.Messages) + } +} + +func TestServiceNewWithFactoryDefaultsToolManager(t *testing.T) { + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + service := NewWithFactory(manager, nil, store, &scriptedProviderFactory{ + provider: &scriptedProvider{ + responses: []provider.ChatResponse{ + { + Message: provider.Message{ + Role: provider.RoleAssistant, + Content: "done", + }, + FinishReason: "stop", + }, + }, + }, + }, nil) + + if err := service.Run(context.Background(), UserInput{RunID: "run-default-tool-manager", Content: "hello"}); err != nil { + t.Fatalf("Run() error = %v", err) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventUserMessage, EventAgentDone}) +} + func TestServiceRunErrorPaths(t *testing.T) { tests := []struct { name string @@ -1213,16 +1364,41 @@ func TestIsRetryableProviderError(t *testing.T) { func TestProviderRetryBackoff(t *testing.T) { t.Parallel() - first := providerRetryBackoff(1) - second := providerRetryBackoff(2) - - if second <= first { - t.Fatalf("expected backoff to increase: first=%v second=%v", first, second) + tests := []struct { + name string + attempt int + min time.Duration + max time.Duration + }{ + { + name: "first retry stays within jittered base window", + attempt: 1, + min: 500 * time.Millisecond, + max: 1500 * time.Millisecond, + }, + { + name: "second retry stays within jittered doubled window", + attempt: 2, + min: 1 * time.Second, + max: 3 * time.Second, + }, + { + name: "large retry is capped at max wait", + attempt: 20, + min: providerRetryMaxWait, + max: providerRetryMaxWait, + }, } - // 验证不超过 MaxWait - large := providerRetryBackoff(20) - if large > providerRetryMaxWait { - t.Fatalf("expected backoff <= MaxWait, got %v > %v", large, providerRetryMaxWait) + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got := providerRetryBackoff(tt.attempt) + if got < tt.min || got > tt.max { + t.Fatalf("providerRetryBackoff(%d) = %v, want within [%v, %v]", tt.attempt, got, tt.min, tt.max) + } + }) } } diff --git a/internal/security/gateway.go b/internal/security/gateway.go new file mode 100644 index 00000000..d4a698c1 --- /dev/null +++ b/internal/security/gateway.go @@ -0,0 +1,91 @@ +package security + +import ( + "context" + "strings" +) + +// PermissionEngine evaluates actions deterministically. +type PermissionEngine interface { + Check(ctx context.Context, action Action) (CheckResult, error) +} + +// StaticGateway evaluates actions against an in-memory ordered rule list. +type StaticGateway struct { + defaultDecision Decision + rules []Rule +} + +// NewStaticGateway creates a permission engine with a default decision and ordered rules. +func NewStaticGateway(defaultDecision Decision, rules []Rule) (*StaticGateway, error) { + if defaultDecision == "" { + defaultDecision = DecisionAllow + } + if err := defaultDecision.Validate(); err != nil { + return nil, err + } + + cloned := make([]Rule, 0, len(rules)) + for _, rule := range rules { + if err := rule.Validate(); err != nil { + return nil, err + } + cloned = append(cloned, rule) + } + + return &StaticGateway{ + defaultDecision: defaultDecision, + rules: cloned, + }, nil +} + +// Check returns the first matching rule result, or the default decision. +func (g *StaticGateway) 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 + } + + for idx := range g.rules { + rule := g.rules[idx] + if !matchesRule(rule, action) { + continue + } + return CheckResult{ + Decision: rule.Decision, + Action: action, + Rule: &g.rules[idx], + Reason: strings.TrimSpace(rule.Reason), + }, nil + } + + return CheckResult{ + Decision: g.defaultDecision, + Action: action, + }, nil +} + +func matchesRule(rule Rule, action Action) bool { + if rule.Type != "" && rule.Type != action.Type { + return false + } + if !matchesResource(rule.Resource, action.Payload.Resource) { + return false + } + + prefix := strings.TrimSpace(rule.TargetPrefix) + if prefix == "" { + return true + } + return strings.HasPrefix(strings.TrimSpace(action.Payload.Target), prefix) +} + +func matchesResource(ruleResource string, actionResource string) bool { + trimmed := strings.TrimSpace(ruleResource) + if trimmed == "" || trimmed == "*" { + return true + } + return strings.EqualFold(trimmed, strings.TrimSpace(actionResource)) +} diff --git a/internal/security/gateway_test.go b/internal/security/gateway_test.go new file mode 100644 index 00000000..c53c4b21 --- /dev/null +++ b/internal/security/gateway_test.go @@ -0,0 +1,242 @@ +package security + +import ( + "context" + "strings" + "testing" +) + +func TestNewStaticGateway(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + defaultDecision Decision + rules []Rule + wantErr string + }{ + { + name: "default allow with no rules", + defaultDecision: DecisionAllow, + }, + { + name: "empty default falls back to allow", + defaultDecision: "", + }, + { + name: "rejects invalid default decision", + defaultDecision: Decision("block"), + wantErr: "invalid decision", + }, + { + name: "rejects invalid rule action", + defaultDecision: DecisionAllow, + rules: []Rule{ + {Resource: "bash", Type: ActionType("boom"), Decision: DecisionDeny}, + }, + wantErr: "invalid action type", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + _, err := NewStaticGateway(tt.defaultDecision, tt.rules) + if tt.wantErr != "" { + if err == nil || !containsText(err, tt.wantErr) { + t.Fatalf("expected error containing %q, got %v", tt.wantErr, err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + } +} + +func TestStaticGatewayCheck(t *testing.T) { + t.Parallel() + + gateway, err := NewStaticGateway(DecisionAllow, []Rule{ + {ID: "deny-bash", Resource: "bash", Type: ActionTypeBash, Decision: DecisionDeny, Reason: "shell is blocked"}, + {ID: "ask-private", Resource: "webfetch", Type: ActionTypeRead, TargetPrefix: "http://127.", Decision: DecisionAsk, Reason: "requires approval"}, + {ID: "deny-all-writes", Type: ActionTypeWrite, Decision: DecisionDeny, Reason: "writes disabled"}, + {ID: "ask-mcp-tool", Resource: "mcp.github.create_issue", Type: ActionTypeMCP, Decision: DecisionAsk, Reason: "mcp approval"}, + }) + if err != nil { + t.Fatalf("new gateway: %v", err) + } + + tests := []struct { + name string + ctx func() context.Context + action Action + wantDecision Decision + wantRuleID string + wantReason string + wantErr string + }{ + { + name: "default allow", + ctx: context.Background, + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + Operation: "read_file", + TargetType: TargetTypePath, + Target: "main.go", + }, + }, + wantDecision: DecisionAllow, + }, + { + name: "matched deny by tool and action", + ctx: context.Background, + action: Action{ + Type: ActionTypeBash, + Payload: ActionPayload{ + ToolName: "bash", + Resource: "bash", + Operation: "command", + TargetType: TargetTypeCommand, + Target: "rm -rf .", + }, + }, + wantDecision: DecisionDeny, + wantRuleID: "deny-bash", + wantReason: "shell is blocked", + }, + { + name: "matched ask by target prefix", + ctx: context.Background, + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + Operation: "fetch", + TargetType: TargetTypeURL, + Target: "http://127.0.0.1:8080", + }, + }, + wantDecision: DecisionAsk, + wantRuleID: "ask-private", + wantReason: "requires approval", + }, + { + name: "matched wildcard action rule", + ctx: context.Background, + action: Action{ + Type: ActionTypeWrite, + Payload: ActionPayload{ + ToolName: "filesystem_write_file", + Resource: "filesystem_write_file", + Operation: "write_file", + TargetType: TargetTypePath, + Target: "notes.txt", + }, + }, + wantDecision: DecisionDeny, + wantRuleID: "deny-all-writes", + wantReason: "writes disabled", + }, + { + name: "matched mcp resource rule", + ctx: context.Background, + action: Action{ + Type: ActionTypeMCP, + Payload: ActionPayload{ + ToolName: "mcp.github.create_issue", + Resource: "mcp.github.create_issue", + Operation: "invoke", + TargetType: TargetTypeMCP, + Target: "github", + }, + }, + wantDecision: DecisionAsk, + wantRuleID: "ask-mcp-tool", + wantReason: "mcp approval", + }, + { + name: "invalid action type", + ctx: context.Background, + action: Action{ + Type: ActionType("invalid"), + Payload: ActionPayload{ + ToolName: "bash", + Resource: "bash", + }, + }, + wantErr: "invalid action type", + }, + { + name: "empty payload tool name", + ctx: context.Background, + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + Resource: "filesystem_read_file", + }, + }, + wantErr: "tool_name is empty", + }, + { + name: "context canceled", + ctx: func() context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + action: Action{ + Type: ActionTypeBash, + Payload: ActionPayload{ + ToolName: "bash", + Resource: "bash", + }, + }, + wantErr: context.Canceled.Error(), + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + result, err := gateway.Check(tt.ctx(), tt.action) + if tt.wantErr != "" { + if err == nil || !containsText(err, tt.wantErr) { + t.Fatalf("expected error containing %q, got %v", tt.wantErr, err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + 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) + } + } else { + if result.Rule == nil || result.Rule.ID != tt.wantRuleID { + t.Fatalf("expected rule id %q, got %+v", tt.wantRuleID, result.Rule) + } + } + if result.Reason != tt.wantReason { + t.Fatalf("expected reason %q, got %q", tt.wantReason, result.Reason) + } + }) + } +} + +func containsText(err error, text string) bool { + return err != nil && strings.Contains(err.Error(), text) +} diff --git a/internal/security/types.go b/internal/security/types.go new file mode 100644 index 00000000..95a32e6a --- /dev/null +++ b/internal/security/types.go @@ -0,0 +1,120 @@ +package security + +import ( + "errors" + "fmt" + "strings" +) + +// Decision is the deterministic outcome returned by the permission engine. +type Decision string + +const ( + DecisionAllow Decision = "allow" + DecisionDeny Decision = "deny" + DecisionAsk Decision = "ask" +) + +// Validate reports whether the decision is supported. +func (d Decision) Validate() error { + switch d { + case DecisionAllow, DecisionDeny, DecisionAsk: + return nil + default: + return fmt.Errorf("security: invalid decision %q", d) + } +} + +// ActionType is the top-level normalized security action category. +type ActionType string + +const ( + ActionTypeBash ActionType = "bash" + ActionTypeRead ActionType = "read" + ActionTypeWrite ActionType = "write" + ActionTypeMCP ActionType = "mcp" +) + +// Validate reports whether the action type is supported. +func (a ActionType) Validate() error { + switch a { + case ActionTypeBash, ActionTypeRead, ActionTypeWrite, ActionTypeMCP: + return nil + default: + return fmt.Errorf("security: invalid action type %q", a) + } +} + +// TargetType describes the resource type an action intends to touch. +type TargetType string + +const ( + TargetTypePath TargetType = "path" + TargetTypeDirectory TargetType = "directory" + TargetTypeCommand TargetType = "command" + TargetTypeURL TargetType = "url" + TargetTypeMCP TargetType = "mcp" +) + +// ActionPayload is the normalized structured context used by policy and sandbox. +type ActionPayload struct { + ToolName string + Resource string + Operation string + SessionID string + Workdir string + TargetType TargetType + Target string +} + +// Action is the unified security input for one tool execution request. +type Action struct { + Type ActionType + Payload ActionPayload +} + +// Validate reports whether the action is usable by the permission engine. +func (a Action) Validate() error { + if err := a.Type.Validate(); err != nil { + return err + } + if strings.TrimSpace(a.Payload.ToolName) == "" { + return errors.New("security: action payload tool_name is empty") + } + if strings.TrimSpace(a.Payload.Resource) == "" { + return errors.New("security: action payload resource is empty") + } + return nil +} + +// Rule is one deterministic permission rule. Empty Type, Resource, or TargetPrefix act +// as wildcards. +type Rule struct { + ID string + Type ActionType + Resource string + TargetPrefix string + Decision Decision + Reason string +} + +// Validate reports whether the rule is well formed. +func (r Rule) Validate() error { + if err := r.Decision.Validate(); err != nil { + return err + } + if r.Type != "" { + if err := r.Type.Validate(); err != nil { + return err + } + } + return nil +} + +// CheckResult is the permission engine response for one action. +type CheckResult struct { + Decision Decision + Action Action + Rule *Rule + Reason string +} diff --git a/internal/security/types_test.go b/internal/security/types_test.go new file mode 100644 index 00000000..93f11a77 --- /dev/null +++ b/internal/security/types_test.go @@ -0,0 +1,116 @@ +package security + +import "testing" + +func TestActionValidate(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + action Action + wantErr string + }{ + { + name: "valid action", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + }, + }, + }, + { + name: "missing resource", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_read_file", + }, + }, + wantErr: "resource is empty", + }, + { + name: "missing tool name", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + Resource: "filesystem_read_file", + }, + }, + wantErr: "tool_name is empty", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := tt.action.Validate() + if tt.wantErr != "" { + if err == nil || err.Error() == "" || !containsText(err, tt.wantErr) { + t.Fatalf("expected error containing %q, got %v", tt.wantErr, err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + } +} + +func TestRuleValidate(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + rule Rule + wantErr string + }{ + { + name: "valid rule", + rule: Rule{ + Type: ActionTypeBash, + Resource: "bash", + Decision: DecisionDeny, + }, + }, + { + name: "invalid decision", + rule: Rule{ + Resource: "bash", + Decision: Decision("maybe"), + }, + wantErr: "invalid decision", + }, + { + name: "invalid action type", + rule: Rule{ + Type: ActionType("custom"), + Resource: "bash", + Decision: DecisionAllow, + }, + wantErr: "invalid action type", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := tt.rule.Validate() + if tt.wantErr != "" { + if err == nil || !containsText(err, tt.wantErr) { + t.Fatalf("expected error containing %q, got %v", tt.wantErr, err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + } +} diff --git a/internal/tools/manager.go b/internal/tools/manager.go new file mode 100644 index 00000000..d7535f1e --- /dev/null +++ b/internal/tools/manager.go @@ -0,0 +1,230 @@ +package tools + +import ( + "context" + "errors" + "fmt" + "strings" + + "neo-code/internal/provider" + "neo-code/internal/security" +) + +// SpecListInput carries future session and agent context for tool filtering. +type SpecListInput struct { + SessionID string + Agent string +} + +// Manager is the runtime-facing tool execution and schema exposure boundary. +type Manager interface { + ListAvailableSpecs(ctx context.Context, input SpecListInput) ([]provider.ToolSpec, error) + Execute(ctx context.Context, input ToolCallInput) (ToolResult, error) +} + +// Executor is the concrete tool execution layer under the manager. +type Executor interface { + ListAvailableSpecs(ctx context.Context, input SpecListInput) ([]provider.ToolSpec, error) + Execute(ctx context.Context, input ToolCallInput) (ToolResult, error) + Supports(name string) bool +} + +// WorkspaceSandbox enforces workspace-oriented constraints before execution. +type WorkspaceSandbox interface { + Check(ctx context.Context, action security.Action) error +} + +// NoopWorkspaceSandbox keeps the explicit sandbox stage in the execution chain +// without changing current behavior. +type NoopWorkspaceSandbox struct{} + +// Check implements WorkspaceSandbox. +func (NoopWorkspaceSandbox) Check(ctx context.Context, action security.Action) error { + return ctx.Err() +} + +// PermissionDecisionError reports a non-allow permission decision. +type PermissionDecisionError struct { + decision security.Decision + toolName string + action security.Action + reason string +} + +// Error returns a stable error message for the blocked tool call. +func (e *PermissionDecisionError) Error() string { + if e == nil { + return "" + } + + reason := strings.TrimSpace(e.reason) + switch e.decision { + case security.DecisionAsk: + if reason == "" { + reason = "permission approval required" + } + default: + if reason == "" { + reason = "permission denied" + } + } + return "tools: " + reason +} + +// Decision returns the blocking engine decision. +func (e *PermissionDecisionError) Decision() string { + if e == nil { + return "" + } + return string(e.decision) +} + +// ToolName returns the tool that was blocked. +func (e *PermissionDecisionError) ToolName() string { + if e == nil { + return "" + } + return e.toolName +} + +// DefaultManager routes tool calls through the permission engine, workspace +// sandbox, and executor. +type DefaultManager struct { + executor Executor + engine security.PermissionEngine + sandbox WorkspaceSandbox +} + +// NewManager creates a manager that wraps an executor with security checks. +func NewManager(executor Executor, engine security.PermissionEngine, sandbox WorkspaceSandbox) (*DefaultManager, error) { + if executor == nil { + return nil, errors.New("tools: executor is nil") + } + if engine == nil { + defaultEngine, err := security.NewStaticGateway(security.DecisionAllow, nil) + if err != nil { + return nil, err + } + engine = defaultEngine + } + if sandbox == nil { + sandbox = NoopWorkspaceSandbox{} + } + + return &DefaultManager{ + executor: executor, + engine: engine, + sandbox: sandbox, + }, nil +} + +// ListAvailableSpecs returns the currently visible tool specs from the executor. +func (m *DefaultManager) ListAvailableSpecs(ctx context.Context, input SpecListInput) ([]provider.ToolSpec, error) { + if m == nil || m.executor == nil { + return nil, errors.New("tools: manager executor is nil") + } + return m.executor.ListAvailableSpecs(ctx, input) +} + +// Execute runs the tool if the permission engine allows it and the sandbox +// check passes. +func (m *DefaultManager) Execute(ctx context.Context, input ToolCallInput) (ToolResult, error) { + if m == nil || m.executor == nil { + return ToolResult{}, errors.New("tools: manager executor is nil") + } + + if !m.executor.Supports(input.Name) { + return m.executor.Execute(ctx, input) + } + + action, err := buildPermissionAction(input) + if err != nil { + result := NewErrorResult(input.Name, "invalid permission action", err.Error(), nil) + result.ToolCallID = input.ID + return result, err + } + + decision, err := m.engine.Check(ctx, action) + if err != nil { + result := NewErrorResult(input.Name, "permission evaluation failed", err.Error(), nil) + result.ToolCallID = input.ID + return result, err + } + if decision.Decision != security.DecisionAllow { + result := blockedToolResult(input, decision) + return result, permissionErrorFromDecision(decision) + } + + if err := m.sandbox.Check(ctx, action); err != nil { + result := NewErrorResult(input.Name, "workspace sandbox rejected action", err.Error(), actionMetadata(action)) + result.ToolCallID = input.ID + return result, err + } + + return m.executor.Execute(ctx, input) +} + +func blockedToolResult(input ToolCallInput, decision security.CheckResult) ToolResult { + reason := "permission denied" + if decision.Decision == security.DecisionAsk { + reason = "permission approval required" + } + if strings.TrimSpace(decision.Reason) != "" { + reason = strings.TrimSpace(decision.Reason) + } + + result := NewErrorResult(input.Name, reason, permissionDetails(decision), permissionMetadata(decision)) + result.ToolCallID = input.ID + return result +} + +func permissionErrorFromDecision(decision security.CheckResult) error { + return &PermissionDecisionError{ + decision: decision.Decision, + toolName: decision.Action.Payload.ToolName, + action: decision.Action, + reason: decision.Reason, + } +} + +func permissionMetadata(decision security.CheckResult) map[string]any { + metadata := actionMetadata(decision.Action) + metadata["permission_decision"] = string(decision.Decision) + if decision.Rule != nil && strings.TrimSpace(decision.Rule.ID) != "" { + metadata["permission_rule_id"] = decision.Rule.ID + } + return metadata +} + +func actionMetadata(action security.Action) map[string]any { + metadata := map[string]any{ + "permission_action_type": string(action.Type), + "permission_resource": action.Payload.Resource, + "permission_operation": action.Payload.Operation, + } + if action.Payload.TargetType != "" { + metadata["permission_target_type"] = string(action.Payload.TargetType) + } + if action.Payload.Target != "" { + metadata["permission_target"] = action.Payload.Target + } + return metadata +} + +func permissionDetails(decision security.CheckResult) string { + parts := make([]string, 0, 5) + parts = append(parts, "type: "+string(decision.Action.Type)) + if strings.TrimSpace(decision.Action.Payload.Resource) != "" { + parts = append(parts, "resource: "+decision.Action.Payload.Resource) + } + if strings.TrimSpace(decision.Action.Payload.Operation) != "" { + parts = append(parts, "operation: "+decision.Action.Payload.Operation) + } + if decision.Action.Payload.TargetType != "" && strings.TrimSpace(decision.Action.Payload.Target) != "" { + parts = append(parts, fmt.Sprintf("%s: %s", decision.Action.Payload.TargetType, decision.Action.Payload.Target)) + } + if strings.TrimSpace(decision.Reason) != "" { + parts = append(parts, "policy: "+strings.TrimSpace(decision.Reason)) + } + return strings.Join(parts, "\n") +} diff --git a/internal/tools/manager_test.go b/internal/tools/manager_test.go new file mode 100644 index 00000000..7dcda95c --- /dev/null +++ b/internal/tools/manager_test.go @@ -0,0 +1,564 @@ +package tools + +import ( + "context" + "errors" + "strings" + "testing" + + "neo-code/internal/security" +) + +type managerStubTool struct { + name string + content string + err error + callCount int +} + +func (t *managerStubTool) Name() string { return t.name } + +func (t *managerStubTool) Description() string { return "stub tool" } + +func (t *managerStubTool) Schema() map[string]any { return map[string]any{"type": "object"} } + +func (t *managerStubTool) Execute(ctx context.Context, call ToolCallInput) (ToolResult, error) { + t.callCount++ + return ToolResult{ + Name: t.name, + Content: t.content, + }, t.err +} + +type stubSandbox struct { + err error + callCount int + lastAction security.Action +} + +func (s *stubSandbox) Check(ctx context.Context, action security.Action) error { + s.callCount++ + s.lastAction = action + if err := ctx.Err(); err != nil { + return err + } + return s.err +} + +func TestDefaultManagerListAvailableSpecs(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + registry.Register(&managerStubTool{name: "bash"}) + manager, err := NewManager(registry, nil, nil) + if err != nil { + t.Fatalf("new manager: %v", err) + } + + specs, err := manager.ListAvailableSpecs(context.Background(), SpecListInput{SessionID: "s-1"}) + if err != nil { + t.Fatalf("list specs: %v", err) + } + if len(specs) != 1 || specs[0].Name != "bash" { + t.Fatalf("unexpected specs: %+v", specs) + } +} + +func TestDefaultManagerListAvailableSpecsBoundaries(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + manager *DefaultManager + ctx func() context.Context + expectErr string + }{ + { + name: "nil manager executor", + manager: &DefaultManager{}, + ctx: context.Background, + expectErr: "manager executor is nil", + }, + { + name: func() string { return "canceled context" }(), + manager: func() *DefaultManager { + registry := NewRegistry() + registry.Register(&managerStubTool{name: "bash"}) + manager, _ := NewManager(registry, nil, nil) + return manager + }(), + ctx: func() context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + expectErr: context.Canceled.Error(), + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + _, err := tt.manager.ListAvailableSpecs(tt.ctx(), SpecListInput{}) + if err == nil || !strings.Contains(err.Error(), tt.expectErr) { + t.Fatalf("expected error containing %q, got %v", tt.expectErr, err) + } + }) + } +} + +func TestDefaultManagerExecute(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + rules []security.Rule + sandboxErr error + input ToolCallInput + expectErr string + expectContent []string + expectDecision string + expectCalls int + expectSandboxRuns int + }{ + { + name: "allow executes tool", + input: ToolCallInput{ + ID: "call-1", + Name: "bash", + Arguments: []byte(`{"command":"echo hi"}`), + }, + expectContent: []string{"ok"}, + expectCalls: 1, + expectSandboxRuns: 1, + }, + { + name: "deny blocks execution before sandbox", + rules: []security.Rule{ + {ID: "deny-bash", Resource: "bash", Type: security.ActionTypeBash, Decision: security.DecisionDeny, Reason: "bash denied"}, + }, + input: ToolCallInput{ + ID: "call-2", + Name: "bash", + Arguments: []byte(`{"command":"echo hi"}`), + }, + expectErr: "bash denied", + expectContent: []string{"tool error", "tool: bash", "reason: bash denied"}, + expectDecision: "deny", + expectCalls: 0, + expectSandboxRuns: 0, + }, + { + name: "ask blocks execution before sandbox", + rules: []security.Rule{ + {ID: "ask-private", Resource: "webfetch", Type: security.ActionTypeRead, Decision: security.DecisionAsk, Reason: "requires approval"}, + }, + input: ToolCallInput{ + ID: "call-3", + Name: "webfetch", + Arguments: []byte(`{"url":"https://example.com"}`), + }, + expectErr: "requires approval", + expectContent: []string{"tool error", "tool: webfetch", "reason: requires approval"}, + expectDecision: "ask", + expectCalls: 0, + expectSandboxRuns: 0, + }, + { + name: "sandbox blocks after allow", + input: ToolCallInput{ + ID: "call-5", + Name: "filesystem_write_file", + Arguments: []byte(`{"path":"notes.txt","content":"hi"}`), + }, + sandboxErr: errors.New("workspace denied"), + expectErr: "workspace denied", + expectContent: []string{"tool error", "reason: workspace sandbox rejected action"}, + expectCalls: 0, + expectSandboxRuns: 1, + }, + { + name: "unknown tool uses executor error", + input: ToolCallInput{ + ID: "call-4", + Name: "missing", + }, + expectErr: "tool: not found", + expectContent: []string{"tool error", "tool: missing"}, + expectDecision: "", + expectCalls: 0, + expectSandboxRuns: 0, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + bashTool := &managerStubTool{name: "bash", content: "ok"} + webTool := &managerStubTool{name: "webfetch", content: "ok"} + writeTool := &managerStubTool{name: "filesystem_write_file", content: "ok"} + registry.Register(bashTool) + registry.Register(webTool) + registry.Register(writeTool) + + engine, err := security.NewStaticGateway(security.DecisionAllow, tt.rules) + if err != nil { + t.Fatalf("new engine: %v", err) + } + sandbox := &stubSandbox{err: tt.sandboxErr} + manager, err := NewManager(registry, engine, sandbox) + if err != nil { + t.Fatalf("new manager: %v", err) + } + + result, execErr := manager.Execute(context.Background(), tt.input) + if tt.expectErr != "" { + if execErr == nil || !strings.Contains(execErr.Error(), tt.expectErr) { + t.Fatalf("expected error containing %q, got %v", tt.expectErr, execErr) + } + } else if execErr != nil { + t.Fatalf("unexpected error: %v", execErr) + } + + for _, fragment := range tt.expectContent { + if !strings.Contains(result.Content, fragment) { + t.Fatalf("expected content containing %q, got %q", fragment, result.Content) + } + } + if decision, _ := result.Metadata["permission_decision"].(string); decision != tt.expectDecision { + t.Fatalf("expected permission decision %q, got %q", tt.expectDecision, decision) + } + + totalCalls := bashTool.callCount + webTool.callCount + writeTool.callCount + if totalCalls != tt.expectCalls { + t.Fatalf("expected %d tool calls, got %d", tt.expectCalls, totalCalls) + } + if sandbox.callCount != tt.expectSandboxRuns { + t.Fatalf("expected sandbox runs %d, got %d", tt.expectSandboxRuns, sandbox.callCount) + } + }) + } +} + +func TestDefaultManagerExecuteBoundaries(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + manager *DefaultManager + input ToolCallInput + expectErr string + }{ + { + name: "nil manager executor", + manager: &DefaultManager{}, + input: ToolCallInput{Name: "bash"}, + expectErr: "manager executor is nil", + }, + { + name: "invalid permission mapping", + manager: func() *DefaultManager { + registry := NewRegistry() + registry.Register(&managerStubTool{name: "custom_tool"}) + manager, _ := NewManager(registry, nil, nil) + return manager + }(), + input: ToolCallInput{Name: "custom_tool"}, + expectErr: "unsupported permission mapping", + }, + { + name: "canceled evaluation context", + manager: func() *DefaultManager { + registry := NewRegistry() + registry.Register(&managerStubTool{name: "bash"}) + manager, _ := NewManager(registry, nil, nil) + return manager + }(), + input: ToolCallInput{Name: "bash", Arguments: []byte(`{"command":"echo hi"}`)}, + expectErr: context.Canceled.Error(), + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + ctx := context.Background() + if tt.expectErr == context.Canceled.Error() { + canceled, cancel := context.WithCancel(context.Background()) + cancel() + ctx = canceled + } + + _, err := tt.manager.Execute(ctx, tt.input) + if err == nil || !strings.Contains(err.Error(), tt.expectErr) { + t.Fatalf("expected error containing %q, got %v", tt.expectErr, err) + } + }) + } +} + +func TestPermissionDecisionError(t *testing.T) { + t.Parallel() + + err := &PermissionDecisionError{ + decision: security.DecisionAsk, + toolName: "webfetch", + action: security.Action{ + Type: security.ActionTypeRead, + Payload: security.ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + }, + }, + reason: "approval required", + } + if !strings.Contains(err.Error(), "approval required") { + t.Fatalf("expected reason in error, got %q", err.Error()) + } + if err.Decision() != "ask" { + t.Fatalf("expected ask decision, got %q", err.Decision()) + } + if err.ToolName() != "webfetch" { + t.Fatalf("expected tool name webfetch, got %q", err.ToolName()) + } + if errors.Is(err, context.Canceled) { + t.Fatalf("permission error should not match unrelated errors") + } + + denyErr := &PermissionDecisionError{} + if !strings.Contains(denyErr.Error(), "permission denied") { + t.Fatalf("expected default deny message, got %q", denyErr.Error()) + } + if denyErr.Decision() != "" { + t.Fatalf("expected empty decision, got %q", denyErr.Decision()) + } + if denyErr.ToolName() != "" { + t.Fatalf("expected empty tool name, got %q", denyErr.ToolName()) + } + + var nilErr *PermissionDecisionError + if nilErr.Error() != "" || nilErr.Decision() != "" || nilErr.ToolName() != "" { + t.Fatalf("expected nil permission error helpers to be empty") + } +} + +func TestBuildPermissionAction(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input ToolCallInput + wantType security.ActionType + wantResource string + wantTarget string + wantErr string + }{ + { + name: "bash maps to bash action", + input: ToolCallInput{ + Name: "bash", + Arguments: []byte(`{"command":"echo hi"}`), + }, + wantType: security.ActionTypeBash, + wantResource: "bash", + wantTarget: "echo hi", + }, + { + name: "read file maps to read action", + input: ToolCallInput{ + Name: "filesystem_read_file", + Arguments: []byte(`{"path":"main.go"}`), + }, + wantType: security.ActionTypeRead, + wantResource: "filesystem_read_file", + wantTarget: "main.go", + }, + { + name: "grep maps to read action", + input: ToolCallInput{ + Name: "filesystem_grep", + Arguments: []byte(`{"dir":"internal"}`), + }, + wantType: security.ActionTypeRead, + wantResource: "filesystem_grep", + wantTarget: "internal", + }, + { + name: "glob maps to read action", + input: ToolCallInput{ + Name: "filesystem_glob", + Arguments: []byte(`{"dir":"cmd"}`), + }, + wantType: security.ActionTypeRead, + wantResource: "filesystem_glob", + wantTarget: "cmd", + }, + { + name: "write file maps to write action", + input: ToolCallInput{ + Name: "filesystem_write_file", + Arguments: []byte(`{"path":"main.go"}`), + }, + wantType: security.ActionTypeWrite, + wantResource: "filesystem_write_file", + wantTarget: "main.go", + }, + { + name: "webfetch maps to read action", + input: ToolCallInput{ + Name: "webfetch", + Arguments: []byte(`{"url":"https://example.com"}`), + }, + wantType: security.ActionTypeRead, + wantResource: "webfetch", + wantTarget: "https://example.com", + }, + { + name: "write maps to write action", + input: ToolCallInput{ + Name: "filesystem_edit", + Arguments: []byte(`{"path":"main.go"}`), + }, + wantType: security.ActionTypeWrite, + wantResource: "filesystem_edit", + wantTarget: "main.go", + }, + { + name: "mcp tool maps to mcp action", + input: ToolCallInput{ + Name: "mcp.github.create_issue", + Arguments: []byte(`{"title":"hello"}`), + }, + wantType: security.ActionTypeMCP, + wantResource: "mcp.github.create_issue", + wantTarget: "github", + }, + { + name: "unsupported tool returns error", + input: ToolCallInput{ + Name: "custom_tool", + }, + wantErr: "unsupported permission mapping", + }, + { + name: "empty tool name returns error", + input: ToolCallInput{}, + wantErr: "tool name is empty", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + action, err := buildPermissionAction(tt.input) + if tt.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("expected error containing %q, got %v", tt.wantErr, err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if action.Type != tt.wantType { + t.Fatalf("expected type %q, got %q", tt.wantType, action.Type) + } + if action.Payload.Resource != tt.wantResource { + t.Fatalf("expected resource %q, got %q", tt.wantResource, action.Payload.Resource) + } + if action.Payload.Target != tt.wantTarget { + t.Fatalf("expected target %q, got %q", tt.wantTarget, action.Payload.Target) + } + }) + } +} + +func TestPermissionMapperHelpers(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input []byte + key string + want string + serverTool string + serverWant string + }{ + { + name: "extracts string value", + input: []byte(`{"path":"main.go"}`), + key: "path", + want: "main.go", + }, + { + name: "invalid json returns empty", + input: []byte(`{invalid`), + key: "path", + want: "", + }, + { + name: "missing key returns empty", + input: []byte(`{"url":"https://example.com"}`), + key: "path", + want: "", + }, + { + name: "non string returns empty", + input: []byte(`{"path":123}`), + key: "path", + want: "", + }, + { + name: "mcp server target with server and tool", + serverTool: "mcp.github.create_issue", + serverWant: "github", + }, + { + name: "mcp server target without server", + serverTool: "mcp", + serverWant: "", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + if tt.key != "" { + if got := extractStringArgument(tt.input, tt.key); got != tt.want { + t.Fatalf("expected %q, got %q", tt.want, got) + } + } + if tt.serverTool != "" { + if got := mcpServerTarget(tt.serverTool); got != tt.serverWant { + t.Fatalf("expected server %q, got %q", tt.serverWant, got) + } + } + }) + } +} + +func TestNoopWorkspaceSandbox(t *testing.T) { + t.Parallel() + + sandbox := NoopWorkspaceSandbox{} + if err := sandbox.Check(context.Background(), security.Action{}); err != nil { + t.Fatalf("expected nil error, got %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := sandbox.Check(ctx, security.Action{}); !errors.Is(err, context.Canceled) { + t.Fatalf("expected context canceled, got %v", err) + } +} diff --git a/internal/tools/permission_mapper.go b/internal/tools/permission_mapper.go new file mode 100644 index 00000000..04bdf09a --- /dev/null +++ b/internal/tools/permission_mapper.go @@ -0,0 +1,101 @@ +package tools + +import ( + "encoding/json" + "fmt" + "strings" + + "neo-code/internal/security" +) + +func buildPermissionAction(input ToolCallInput) (security.Action, error) { + toolName := strings.TrimSpace(input.Name) + if toolName == "" { + return security.Action{}, fmt.Errorf("tools: tool name is empty") + } + + action := security.Action{ + Payload: security.ActionPayload{ + ToolName: toolName, + Resource: toolName, + Operation: toolName, + SessionID: input.SessionID, + Workdir: input.Workdir, + }, + } + + switch strings.ToLower(toolName) { + case "bash": + action.Type = security.ActionTypeBash + action.Payload.Operation = "command" + action.Payload.TargetType = security.TargetTypeCommand + action.Payload.Target = extractStringArgument(input.Arguments, "command") + case "filesystem_read_file": + action.Type = security.ActionTypeRead + action.Payload.Operation = "read_file" + action.Payload.TargetType = security.TargetTypePath + action.Payload.Target = extractStringArgument(input.Arguments, "path") + case "filesystem_grep": + action.Type = security.ActionTypeRead + action.Payload.Operation = "grep" + action.Payload.TargetType = security.TargetTypeDirectory + action.Payload.Target = extractStringArgument(input.Arguments, "dir") + case "filesystem_glob": + action.Type = security.ActionTypeRead + action.Payload.Operation = "glob" + action.Payload.TargetType = security.TargetTypeDirectory + action.Payload.Target = extractStringArgument(input.Arguments, "dir") + case "webfetch": + action.Type = security.ActionTypeRead + action.Payload.Operation = "fetch" + action.Payload.TargetType = security.TargetTypeURL + action.Payload.Target = extractStringArgument(input.Arguments, "url") + case "filesystem_write_file": + action.Type = security.ActionTypeWrite + action.Payload.Operation = "write_file" + action.Payload.TargetType = security.TargetTypePath + action.Payload.Target = extractStringArgument(input.Arguments, "path") + case "filesystem_edit": + action.Type = security.ActionTypeWrite + action.Payload.Operation = "edit" + action.Payload.TargetType = security.TargetTypePath + action.Payload.Target = extractStringArgument(input.Arguments, "path") + default: + if strings.HasPrefix(strings.ToLower(toolName), "mcp.") { + action.Type = security.ActionTypeMCP + action.Payload.Operation = "invoke" + action.Payload.TargetType = security.TargetTypeMCP + action.Payload.Target = mcpServerTarget(toolName) + action.Payload.Resource = toolName + return action, nil + } + return security.Action{}, fmt.Errorf("tools: unsupported permission mapping for %q", input.Name) + } + + return action, nil +} + +func mcpServerTarget(toolName string) string { + parts := strings.Split(strings.TrimSpace(toolName), ".") + if len(parts) < 2 { + return "" + } + return parts[1] +} + +func extractStringArgument(raw []byte, key string) string { + if len(raw) == 0 { + return "" + } + + var payload map[string]any + if err := json.Unmarshal(raw, &payload); err != nil { + return "" + } + + value, ok := payload[key].(string) + if !ok { + return "" + } + return strings.TrimSpace(value) +} diff --git a/internal/tools/registry.go b/internal/tools/registry.go index c0fc57a1..c057ee86 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -34,6 +34,12 @@ func (r *Registry) Get(name string) (Tool, error) { return tool, nil } +// Supports reports whether a tool is registered. +func (r *Registry) Supports(name string) bool { + _, err := r.Get(name) + return err == nil +} + func (r *Registry) GetSpecs() []provider.ToolSpec { names := make([]string, 0, len(r.tools)) for name := range r.tools { @@ -57,6 +63,14 @@ func (r *Registry) ListSchemas() []provider.ToolSpec { return r.GetSpecs() } +// ListAvailableSpecs returns all registered tool specs. +func (r *Registry) ListAvailableSpecs(ctx context.Context, input SpecListInput) ([]provider.ToolSpec, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + return r.GetSpecs(), nil +} + func (r *Registry) Execute(ctx context.Context, input ToolCallInput) (ToolResult, error) { tool, err := r.Get(input.Name) if err != nil { diff --git a/internal/tools/registry_test.go b/internal/tools/registry_test.go index fd12f3b1..e8d2db35 100644 --- a/internal/tools/registry_test.go +++ b/internal/tools/registry_test.go @@ -145,3 +145,38 @@ func TestRegistryExecute(t *testing.T) { }) } } + +func TestRegistryHelpers(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + registry.Register(nil) + registry.Register(stubTool{name: "a_tool", description: "first", schema: map[string]any{"type": "object"}}) + + if !registry.Supports("a_tool") { + t.Fatalf("expected registry to support a_tool") + } + if registry.Supports("missing") { + t.Fatalf("did not expect registry to support missing tool") + } + + schemas := registry.ListSchemas() + if len(schemas) != 1 || schemas[0].Name != "a_tool" { + t.Fatalf("unexpected schemas: %+v", schemas) + } + + specs, err := registry.ListAvailableSpecs(context.Background(), SpecListInput{SessionID: "s-1"}) + if err != nil { + t.Fatalf("ListAvailableSpecs() error = %v", err) + } + if len(specs) != 1 || specs[0].Name != "a_tool" { + t.Fatalf("unexpected specs: %+v", specs) + } + + canceled, cancel := context.WithCancel(context.Background()) + cancel() + _, err = registry.ListAvailableSpecs(canceled, SpecListInput{}) + if err == nil || !strings.Contains(err.Error(), context.Canceled.Error()) { + t.Fatalf("expected context canceled, got %v", err) + } +}