Merge pull request 'Add a hook bus and permission decisions for tool calls' (#13) from feature/permission-hooks into main
This commit was merged in pull request #13.
This commit is contained in:
@@ -15,6 +15,7 @@ import (
|
||||
"agent-native/exec"
|
||||
"agent-native/runtime"
|
||||
"agent-native/runtime/filestate"
|
||||
"agent-native/runtime/hooks"
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/mcp"
|
||||
"agent-native/runtime/promptx"
|
||||
@@ -59,6 +60,16 @@ type ObserveOptions struct {
|
||||
ExecutionID string
|
||||
}
|
||||
|
||||
// GateOptions 把直连已经使用的停止、提醒和工具钩子交给小队成员。零值不改变行为。
|
||||
type GateOptions struct {
|
||||
BeforeTool func(call llm.ToolCall) error
|
||||
AfterTool func(call llm.ToolCall, res llm.Message)
|
||||
Stop func(last llm.Message, info runtime.StopInfo) runtime.StopDecision
|
||||
Reminders func(turn int) []string
|
||||
RepeatGuard bool
|
||||
Hooks *hooks.Bus
|
||||
}
|
||||
|
||||
// Options 构造后端所需的模型接入与执行参数（由组合根注入，无产品依赖）。
|
||||
// 字段经嵌入提升，读取仍写 opts.Model；复合字面量要写子结构名。
|
||||
type Options struct {
|
||||
@@ -66,6 +77,7 @@ type Options struct {
|
||||
ToolOptions
|
||||
SessionOptions
|
||||
ObserveOptions
|
||||
GateOptions
|
||||
}
|
||||
|
||||
// CompactConfig 声明长会话压缩策略（优化性质：压缩失败不阻断执行，但留日志）。
|
||||
@@ -333,6 +345,9 @@ func (b *Backend) run(ctx context.Context, model string, budget *runtime.Budget,
|
||||
Log: b.logger(),
|
||||
ExecutionID: execID,
|
||||
Budget: budget,
|
||||
Reminders: b.opts.Reminders,
|
||||
RepeatGuard: b.opts.RepeatGuard,
|
||||
Hooks: gateBus(b.opts.GateOptions),
|
||||
OnEvent: func(ev runtime.Event) {
|
||||
logs.WriteString(ev.Type)
|
||||
if ev.Tool != "" {
|
||||
@@ -423,6 +438,40 @@ func defaultSummarizer(stream llm.StreamFunc, model string) func(ctx context.Con
|
||||
}
|
||||
}
|
||||
|
||||
func gateBus(g GateOptions) *hooks.Bus {
|
||||
bus := g.Hooks
|
||||
if g.BeforeTool == nil && g.AfterTool == nil && g.Stop == nil {
|
||||
return bus
|
||||
}
|
||||
if bus == nil {
|
||||
bus = hooks.New()
|
||||
}
|
||||
if g.BeforeTool != nil {
|
||||
fn := g.BeforeTool
|
||||
bus.On(hooks.AdaptBefore(func(name, args string) error { return fn(llm.ToolCall{Name: name, Arguments: args}) }))
|
||||
}
|
||||
if g.AfterTool != nil {
|
||||
fn := g.AfterTool
|
||||
bus.On(hooks.AdaptAfter(func(name, args, result string, isError bool) {
|
||||
fn(llm.ToolCall{Name: name, Arguments: args}, llm.Message{Content: result, IsError: isError})
|
||||
}))
|
||||
}
|
||||
if g.Stop != nil {
|
||||
fn := g.Stop
|
||||
bus.On(func(c hooks.Call) hooks.Outcome {
|
||||
if c.Event != hooks.Stop {
|
||||
return hooks.Outcome{}
|
||||
}
|
||||
d := fn(llm.Message{Role: llm.RoleAssistant, Content: c.LastText}, runtime.StopInfo{Continuations: c.Continuations, ToolCalls: c.ToolCalls})
|
||||
if d.Continue {
|
||||
return hooks.Outcome{Decision: hooks.Block, Reason: d.Reason, Replacement: d.Replacement, Notice: d.Notice}
|
||||
}
|
||||
return hooks.Outcome{Decision: hooks.Allow, Notice: d.Notice}
|
||||
})
|
||||
}
|
||||
return bus
|
||||
}
|
||||
|
||||
type errString string
|
||||
|
||||
func (e errString) Error() string { return string(e) }
|
||||
|
||||
@@ -2,6 +2,7 @@ package native
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
@@ -139,3 +140,36 @@ func TestBackendFailClosedWithoutConfig(t *testing.T) {
|
||||
t.Fatalf("want profile error, got %+v", res3)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBeforeToolReachesSquadMembers(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
var seen string
|
||||
b := New(Options{
|
||||
ModelOptions: ModelOptions{
|
||||
Model: "fake",
|
||||
Stream: script(
|
||||
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
|
||||
ID: "c1", Name: "write_file", Arguments: `{"path":"output.md","content":"nope"}`,
|
||||
}}},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "blocked", StopReason: llm.StopEnd},
|
||||
),
|
||||
},
|
||||
GateOptions: GateOptions{BeforeTool: func(call llm.ToolCall) error {
|
||||
seen = call.Name
|
||||
return errors.New("计划模式拒绝了写工具")
|
||||
}},
|
||||
})
|
||||
res, err := b.Execute(context.Background(), exec.Request{WorkDir: dir, Prompt: "写", Targets: []string{"output.md"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if seen != "write_file" {
|
||||
t.Fatalf("before tool saw %q", seen)
|
||||
}
|
||||
if res.Success {
|
||||
t.Fatal("denied write still succeeded")
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, "output.md")); err == nil {
|
||||
t.Fatal("write_file ran after deny")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"agent-native/runtime/hooks"
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/tools"
|
||||
)
|
||||
|
||||
func TestHookAdaptersMatchDirectCallbacks(t *testing.T) {
|
||||
direct := runGate(t, false)
|
||||
adapted := runGate(t, true)
|
||||
if direct.stop != adapted.stop || direct.text != adapted.text {
|
||||
t.Fatalf("stop/text direct=%s %q adapted=%s %q", direct.stop, direct.text, adapted.stop, adapted.text)
|
||||
}
|
||||
if direct.denies != adapted.denies || direct.reminders != adapted.reminders || direct.errors != adapted.errors {
|
||||
t.Fatalf("counts direct=%+v adapted=%+v", direct, adapted)
|
||||
}
|
||||
if direct.denies != 1 || direct.reminders == 0 {
|
||||
t.Fatalf("unexpected counts %+v", direct)
|
||||
}
|
||||
}
|
||||
|
||||
type gateStats struct {
|
||||
stop, text string
|
||||
denies, reminders, errors int
|
||||
}
|
||||
|
||||
func runGate(t *testing.T, viaHooks bool) gateStats {
|
||||
t.Helper()
|
||||
var stats gateStats
|
||||
before := func(call llm.ToolCall) error {
|
||||
if call.Name == "read_file" {
|
||||
stats.denies++
|
||||
return errors.New("denied")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
reminders := func(int) []string {
|
||||
stats.reminders++
|
||||
return []string{"记一笔"}
|
||||
}
|
||||
stop := func(last llm.Message, info StopInfo) StopDecision {
|
||||
if info.Continuations < 1 {
|
||||
return StopDecision{Continue: true, Reason: "再看一眼"}
|
||||
}
|
||||
return StopDecision{}
|
||||
}
|
||||
cfg := Config{
|
||||
Model: "fake",
|
||||
Tools: tools.WorkspaceTools(t.TempDir()),
|
||||
Stream: script(
|
||||
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
|
||||
ID: "c1", Name: "read_file", Arguments: `{"path":"a.md"}`,
|
||||
}}},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
|
||||
),
|
||||
Reminders: reminders,
|
||||
}
|
||||
if viaHooks {
|
||||
bus := hooks.New()
|
||||
bus.On(hooks.AdaptBefore(func(name, args string) error {
|
||||
return before(llm.ToolCall{Name: name, Arguments: args})
|
||||
}))
|
||||
bus.On(func(c hooks.Call) hooks.Outcome {
|
||||
if c.Event != hooks.Stop {
|
||||
return hooks.Outcome{}
|
||||
}
|
||||
d := stop(llm.Message{Content: c.LastText}, StopInfo{Continuations: c.Continuations, ToolCalls: c.ToolCalls})
|
||||
if d.Continue {
|
||||
return hooks.Outcome{Decision: hooks.Block, Reason: d.Reason, Replacement: d.Replacement, Notice: d.Notice}
|
||||
}
|
||||
return hooks.Outcome{Decision: hooks.Allow, Notice: d.Notice}
|
||||
})
|
||||
cfg.Hooks = bus
|
||||
} else {
|
||||
cfg.BeforeTool = before
|
||||
cfg.Stop = stop
|
||||
}
|
||||
res, err := Run(context.Background(), cfg, llm.Message{Role: llm.RoleUser, Content: "读文件"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
stats.stop = string(res.StopReason)
|
||||
stats.text = res.Text
|
||||
for _, m := range res.Messages {
|
||||
if m.Role == llm.RoleTool && m.IsError {
|
||||
stats.errors++
|
||||
}
|
||||
}
|
||||
return stats
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
// Package hooks 是进程内钩子总线。决策按 block/deny、ask、modify、allow 取更强的一条。
|
||||
// 没有表态的处理函数不改变结果。
|
||||
package hooks
|
||||
|
||||
import "sync"
|
||||
|
||||
type Event string
|
||||
|
||||
const (
|
||||
SessionStart Event = "SessionStart"
|
||||
UserPromptSubmit Event = "UserPromptSubmit"
|
||||
PreToolUse Event = "PreToolUse"
|
||||
PermissionRequest Event = "PermissionRequest"
|
||||
PostToolUse Event = "PostToolUse"
|
||||
PostToolUseFailure Event = "PostToolUseFailure"
|
||||
Stop Event = "Stop"
|
||||
SubagentStart Event = "SubagentStart"
|
||||
SubagentStop Event = "SubagentStop"
|
||||
PreCompact Event = "PreCompact"
|
||||
)
|
||||
|
||||
type Decision string
|
||||
|
||||
const (
|
||||
Allow Decision = "allow"
|
||||
Deny Decision = "deny"
|
||||
Ask Decision = "ask"
|
||||
Modify Decision = "modify"
|
||||
Block Decision = "block"
|
||||
)
|
||||
|
||||
// Call 是一次事件。工具字段只在工具事件上有意义。
|
||||
type Call struct {
|
||||
Event Event
|
||||
Tool string
|
||||
Args string
|
||||
Result string
|
||||
IsError bool
|
||||
Turn int
|
||||
LastText string
|
||||
Continuations int
|
||||
ToolCalls int
|
||||
}
|
||||
|
||||
// Outcome 是一个处理函数的表态。空 Decision 表示不插手。
|
||||
type Outcome struct {
|
||||
Decision Decision
|
||||
Reason string
|
||||
Args string
|
||||
Notice string
|
||||
Replacement string
|
||||
}
|
||||
|
||||
type Handler func(Call) Outcome
|
||||
|
||||
type Bus struct {
|
||||
mu sync.Mutex
|
||||
hs []Handler
|
||||
}
|
||||
|
||||
func New() *Bus { return &Bus{} }
|
||||
|
||||
func (b *Bus) On(h Handler) {
|
||||
if b == nil || h == nil {
|
||||
return
|
||||
}
|
||||
b.mu.Lock()
|
||||
b.hs = append(b.hs, h)
|
||||
b.mu.Unlock()
|
||||
}
|
||||
|
||||
// Emit 按优先级合并表态。deny 与 block 同级，先登记的保留。
|
||||
func (b *Bus) Emit(c Call) Outcome {
|
||||
if b == nil {
|
||||
return Outcome{}
|
||||
}
|
||||
b.mu.Lock()
|
||||
hs := append([]Handler(nil), b.hs...)
|
||||
b.mu.Unlock()
|
||||
var merged Outcome
|
||||
for _, h := range hs {
|
||||
merged = prefer(merged, h(c))
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
func prefer(cur, next Outcome) Outcome {
|
||||
if rank(next.Decision) > rank(cur.Decision) {
|
||||
return next
|
||||
}
|
||||
return cur
|
||||
}
|
||||
|
||||
func rank(d Decision) int {
|
||||
switch d {
|
||||
case Block, Deny:
|
||||
return 4
|
||||
case Ask:
|
||||
return 3
|
||||
case Modify:
|
||||
return 2
|
||||
case Allow:
|
||||
return 1
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
// AdaptBefore 把「返回错误即拒绝」收成 PreToolUse 处理函数。
|
||||
func AdaptBefore(fn func(name, args string) error) Handler {
|
||||
return func(c Call) Outcome {
|
||||
if c.Event != PreToolUse || fn == nil {
|
||||
return Outcome{}
|
||||
}
|
||||
if err := fn(c.Tool, c.Args); err != nil {
|
||||
return Outcome{Decision: Deny, Reason: err.Error()}
|
||||
}
|
||||
return Outcome{Decision: Allow}
|
||||
}
|
||||
}
|
||||
|
||||
// AdaptAfter 把工具结果观察收成 PostToolUse 或失败事件。它不改写结果。
|
||||
func AdaptAfter(fn func(name, args, result string, isError bool)) Handler {
|
||||
return func(c Call) Outcome {
|
||||
if fn == nil {
|
||||
return Outcome{}
|
||||
}
|
||||
switch c.Event {
|
||||
case PostToolUse, PostToolUseFailure:
|
||||
fn(c.Tool, c.Args, c.Result, c.IsError)
|
||||
}
|
||||
return Outcome{}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package hooks
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestDenyBeatsAllow(t *testing.T) {
|
||||
b := New()
|
||||
b.On(func(Call) Outcome { return Outcome{Decision: Allow} })
|
||||
b.On(func(c Call) Outcome {
|
||||
if c.Event == PreToolUse {
|
||||
return Outcome{Decision: Deny, Reason: "no"}
|
||||
}
|
||||
return Outcome{}
|
||||
})
|
||||
out := b.Emit(Call{Event: PreToolUse, Tool: "write_file"})
|
||||
if out.Decision != Deny || out.Reason != "no" {
|
||||
t.Fatalf("%+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdaptBefore(t *testing.T) {
|
||||
h := AdaptBefore(func(name, args string) error {
|
||||
if name == "write_file" {
|
||||
return errString("blocked")
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if h(Call{Event: PreToolUse, Tool: "read_file"}).Decision != Allow {
|
||||
t.Fatal("read")
|
||||
}
|
||||
if h(Call{Event: PreToolUse, Tool: "write_file"}).Decision != Deny {
|
||||
t.Fatal("write")
|
||||
}
|
||||
if h(Call{Event: Stop}).Decision != "" {
|
||||
t.Fatal("stop should be ignored")
|
||||
}
|
||||
}
|
||||
|
||||
type errString string
|
||||
|
||||
func (e errString) Error() string { return string(e) }
|
||||
+66
-2
@@ -8,6 +8,7 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"agent-native/runtime/hooks"
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/reminders"
|
||||
"agent-native/runtime/session"
|
||||
@@ -104,6 +105,8 @@ type Config struct {
|
||||
ContextBuffer int
|
||||
// KeepToolResults 是清理时保留的最近工具结果条数。0 表示 5。
|
||||
KeepToolResults int
|
||||
// Hooks 是进程内钩子。空总线不改变停止、提醒和拒绝次数。
|
||||
Hooks *hooks.Bus
|
||||
}
|
||||
|
||||
// StopInfo 是 Stop 钩子看到的本 run 进度。
|
||||
@@ -275,8 +278,7 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
}
|
||||
continue
|
||||
}
|
||||
if cfg.Stop != nil {
|
||||
decision := cfg.Stop(last, StopInfo{Continuations: stopContinues, ToolCalls: toolCalls})
|
||||
if decision, hooked := stopDecision(cfg, last, stopContinues, toolCalls); hooked {
|
||||
if decision.Continue && stopContinues < maxStopContinuations {
|
||||
stopContinues++
|
||||
if decision.Replacement != "" && len(messages) > 0 {
|
||||
@@ -488,6 +490,56 @@ func rejectIncompleteToolCall(call llm.ToolCall) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
func stopDecision(cfg Config, last llm.Message, continuations, toolCalls int) (StopDecision, bool) {
|
||||
if cfg.Hooks != nil {
|
||||
out := cfg.Hooks.Emit(hooks.Call{
|
||||
Event: hooks.Stop, LastText: last.Content, Continuations: continuations, ToolCalls: toolCalls,
|
||||
})
|
||||
if out.Decision != "" {
|
||||
return StopDecision{
|
||||
Continue: out.Decision == hooks.Block, Reason: out.Reason,
|
||||
Replacement: out.Replacement, Notice: out.Notice,
|
||||
}, true
|
||||
}
|
||||
}
|
||||
if cfg.Stop != nil {
|
||||
return cfg.Stop(last, StopInfo{Continuations: continuations, ToolCalls: toolCalls}), true
|
||||
}
|
||||
return StopDecision{}, false
|
||||
}
|
||||
|
||||
func gateTool(cfg Config, call llm.ToolCall) (llm.ToolCall, llm.Message, bool) {
|
||||
if cfg.Hooks == nil {
|
||||
return call, llm.Message{}, false
|
||||
}
|
||||
out := cfg.Hooks.Emit(hooks.Call{Event: hooks.PreToolUse, Tool: call.Name, Args: call.Arguments})
|
||||
switch out.Decision {
|
||||
case hooks.Deny, hooks.Block, hooks.Ask:
|
||||
why := out.Reason
|
||||
if why == "" {
|
||||
why = "工具调用被拒绝"
|
||||
}
|
||||
msg := llm.Message{Role: llm.RoleTool, Name: call.Name, ToolCallID: call.ID, Content: why, IsError: true}
|
||||
return call, msg, true
|
||||
case hooks.Modify:
|
||||
if out.Args != "" {
|
||||
call.Arguments = out.Args
|
||||
}
|
||||
}
|
||||
return call, llm.Message{}, false
|
||||
}
|
||||
|
||||
func noteToolHook(cfg Config, call llm.ToolCall, msg llm.Message) {
|
||||
if cfg.Hooks == nil {
|
||||
return
|
||||
}
|
||||
ev := hooks.PostToolUse
|
||||
if msg.IsError {
|
||||
ev = hooks.PostToolUseFailure
|
||||
}
|
||||
cfg.Hooks.Emit(hooks.Call{Event: ev, Tool: call.Name, Args: call.Arguments, Result: msg.Content, IsError: msg.IsError})
|
||||
}
|
||||
|
||||
func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return cancelledToolMessage(call, err)
|
||||
@@ -502,10 +554,21 @@ func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message {
|
||||
return msg
|
||||
}
|
||||
emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name, Detail: call.Arguments})
|
||||
if gated, msg, blocked := gateTool(cfg, call); blocked {
|
||||
emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: msg.Content})
|
||||
noteToolHook(cfg, call, msg)
|
||||
if cfg.AfterTool != nil {
|
||||
cfg.AfterTool(call, msg)
|
||||
}
|
||||
return msg
|
||||
} else {
|
||||
call = gated
|
||||
}
|
||||
if cfg.BeforeTool != nil {
|
||||
if err := cfg.BeforeTool(call); err != nil {
|
||||
msg := llm.Message{Role: llm.RoleTool, Name: call.Name, ToolCallID: call.ID, Content: err.Error(), IsError: true}
|
||||
emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: msg.Content})
|
||||
noteToolHook(cfg, call, msg)
|
||||
if cfg.AfterTool != nil {
|
||||
cfg.AfterTool(call, msg)
|
||||
}
|
||||
@@ -531,6 +594,7 @@ func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message {
|
||||
IsError: res.IsError,
|
||||
}
|
||||
emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: res.Content})
|
||||
noteToolHook(cfg, call, msg)
|
||||
if cfg.AfterTool != nil {
|
||||
cfg.AfterTool(call, msg)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
// Package policy 决定一次工具调用在当前权限模式下是拒绝、询问还是放行。
|
||||
// 默认模式不额外表态，调用方继续用原有画像策略。
|
||||
package policy
|
||||
|
||||
import "strings"
|
||||
|
||||
type Mode string
|
||||
|
||||
const (
|
||||
ModeDefault Mode = "default"
|
||||
ModePlan Mode = "plan"
|
||||
ModeAutoEdit Mode = "auto_edit"
|
||||
)
|
||||
|
||||
type Effect string
|
||||
|
||||
const (
|
||||
EffectNone Effect = ""
|
||||
EffectDeny Effect = "deny"
|
||||
EffectAsk Effect = "ask"
|
||||
EffectAllow Effect = "allow"
|
||||
)
|
||||
|
||||
type Traits struct {
|
||||
ReadOnly bool
|
||||
Edits bool
|
||||
Risk string
|
||||
}
|
||||
|
||||
type Rule struct {
|
||||
Effect Effect
|
||||
Tool string
|
||||
ArgGlob string
|
||||
}
|
||||
|
||||
type Request struct {
|
||||
Mode Mode
|
||||
Name string
|
||||
Args string
|
||||
Traits Traits
|
||||
SessionAllow bool
|
||||
}
|
||||
|
||||
func Normalize(mode string) Mode {
|
||||
switch Mode(mode) {
|
||||
case ModePlan, ModeAutoEdit, ModeDefault:
|
||||
return Mode(mode)
|
||||
default:
|
||||
return ModeDefault
|
||||
}
|
||||
}
|
||||
|
||||
// Decide 的优先级是拒绝规则、计划模式禁写、本会话放行、询问规则、自动改文件、高风险命令。
|
||||
// 都没有命中时返回空，表示沿用画像的拒绝、确认和放行。
|
||||
func Decide(req Request, rules []Rule) Effect {
|
||||
var sawAsk, sawAllow bool
|
||||
for _, rule := range rules {
|
||||
if !matchTool(rule.Tool, req.Name) || !glob(rule.ArgGlob, req.Args) {
|
||||
continue
|
||||
}
|
||||
switch rule.Effect {
|
||||
case EffectDeny:
|
||||
return EffectDeny
|
||||
case EffectAsk:
|
||||
sawAsk = true
|
||||
case EffectAllow:
|
||||
sawAllow = true
|
||||
}
|
||||
}
|
||||
traits := mergeTraits(req.Name, req.Traits)
|
||||
if commandRisk(req.Name, req.Args) == "high" {
|
||||
traits.Risk = "high"
|
||||
}
|
||||
if Normalize(string(req.Mode)) == ModePlan && !planVisible(req.Name, traits.ReadOnly) {
|
||||
return EffectDeny
|
||||
}
|
||||
if req.SessionAllow {
|
||||
return EffectAllow
|
||||
}
|
||||
if sawAsk {
|
||||
return EffectAsk
|
||||
}
|
||||
if Normalize(string(req.Mode)) == ModeAutoEdit && traits.Edits && traits.Risk != "high" {
|
||||
return EffectAllow
|
||||
}
|
||||
if sawAllow {
|
||||
return EffectAllow
|
||||
}
|
||||
if traits.Risk == "high" && (req.Name == "run_shell" || req.Name == "run_command") {
|
||||
return EffectAsk
|
||||
}
|
||||
return EffectNone
|
||||
}
|
||||
|
||||
// Visible 报告计划模式下这个工具能否出现在清单里。其他模式一律可见。
|
||||
func Visible(mode string, name string, readOnly bool) bool {
|
||||
if Normalize(mode) != ModePlan {
|
||||
return true
|
||||
}
|
||||
return planVisible(name, readOnly)
|
||||
}
|
||||
|
||||
func planVisible(name string, readOnly bool) bool {
|
||||
switch name {
|
||||
case "update_plan", "update_todos", "ask_user", "exit_plan_mode":
|
||||
return true
|
||||
default:
|
||||
return readOnly
|
||||
}
|
||||
}
|
||||
|
||||
// TraitsFor 是内置工具的只读、改文件和风险声明。未知工具为零值。
|
||||
func TraitsFor(name string) Traits {
|
||||
switch name {
|
||||
case "read_file", "list_dir", "grep", "glob", "load_skill",
|
||||
"fs_read_file", "fs_list_dir", "time_now", "knowledge_search",
|
||||
"list_conversations", "read_conversation", "quality_review",
|
||||
"web_search", "tavily_search", "fetch_url", "ask_user", "exit_plan_mode",
|
||||
"update_plan", "update_todos":
|
||||
return Traits{ReadOnly: true, Risk: "low"}
|
||||
case "write_file", "edit_file", "multi_edit", "fs_write_file", "write_artifact", "register_artifact":
|
||||
return Traits{Edits: true, Risk: "medium"}
|
||||
case "run_shell", "run_command", "run_python":
|
||||
return Traits{Risk: "medium"}
|
||||
default:
|
||||
return Traits{}
|
||||
}
|
||||
}
|
||||
|
||||
func mergeTraits(name string, given Traits) Traits {
|
||||
base := TraitsFor(name)
|
||||
if given.ReadOnly {
|
||||
base.ReadOnly = true
|
||||
}
|
||||
if given.Edits {
|
||||
base.Edits = true
|
||||
}
|
||||
if given.Risk != "" {
|
||||
base.Risk = given.Risk
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
// ClassifyCommand 按命令文本分级。高风险会在默认模式下改为询问。
|
||||
func ClassifyCommand(cmd string) string {
|
||||
c := strings.ToLower(cmd)
|
||||
for _, n := range []string{"rm -rf", "rm -fr", "sudo ", "mkfs", "shutdown", "reboot", "> /dev/", "| sh", "|sh", "chmod -r 777"} {
|
||||
if strings.Contains(c, n) {
|
||||
return "high"
|
||||
}
|
||||
}
|
||||
for _, n := range []string{"rm ", "mv ", "chmod ", "chown ", "kill "} {
|
||||
if strings.Contains(c, n) {
|
||||
return "medium"
|
||||
}
|
||||
}
|
||||
return "low"
|
||||
}
|
||||
|
||||
func commandRisk(name, args string) string {
|
||||
if name != "run_shell" && name != "run_command" {
|
||||
return ""
|
||||
}
|
||||
return ClassifyCommand(extractCommand(args))
|
||||
}
|
||||
|
||||
func extractCommand(args string) string {
|
||||
for _, key := range []string{`"command"`, `"cmd"`} {
|
||||
i := strings.Index(args, key)
|
||||
if i < 0 {
|
||||
continue
|
||||
}
|
||||
rest := args[i+len(key):]
|
||||
rest = strings.TrimLeft(rest, " \t\n:")
|
||||
if !strings.HasPrefix(rest, `"`) {
|
||||
continue
|
||||
}
|
||||
rest = rest[1:]
|
||||
if j := strings.Index(rest, `"`); j >= 0 {
|
||||
return rest[:j]
|
||||
}
|
||||
}
|
||||
return args
|
||||
}
|
||||
|
||||
func matchTool(pattern, name string) bool {
|
||||
if pattern == "" || pattern == "*" {
|
||||
return true
|
||||
}
|
||||
return glob(pattern, name)
|
||||
}
|
||||
|
||||
func glob(pattern, s string) bool {
|
||||
if pattern == "" || pattern == "*" {
|
||||
return true
|
||||
}
|
||||
parts := strings.Split(pattern, "*")
|
||||
if !strings.HasPrefix(pattern, "*") {
|
||||
if !strings.HasPrefix(s, parts[0]) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
if !strings.HasSuffix(pattern, "*") {
|
||||
last := parts[len(parts)-1]
|
||||
if !strings.HasSuffix(s, last) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
idx := 0
|
||||
for _, part := range parts {
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
j := strings.Index(s[idx:], part)
|
||||
if j < 0 {
|
||||
return false
|
||||
}
|
||||
idx += j + len(part)
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package policy
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestPlanModeDeniesWrites(t *testing.T) {
|
||||
effect := Decide(Request{Mode: ModePlan, Name: "write_file", Args: `{"path":"a.md"}`}, nil)
|
||||
if effect != EffectDeny {
|
||||
t.Fatalf("effect %q", effect)
|
||||
}
|
||||
if !Visible(string(ModePlan), "update_plan", false) {
|
||||
t.Fatal("plan tool hidden")
|
||||
}
|
||||
if Visible(string(ModePlan), "write_file", false) {
|
||||
t.Fatal("write visible in plan mode")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultModeStaysQuiet(t *testing.T) {
|
||||
if got := Decide(Request{Mode: ModeDefault, Name: "write_file"}, nil); got != EffectNone {
|
||||
t.Fatalf("default write %q", got)
|
||||
}
|
||||
if got := Decide(Request{Mode: ModeDefault, Name: "run_shell", Args: `{"command":"rm -rf /"}`}, nil); got != EffectAsk {
|
||||
t.Fatalf("high shell %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRulePrecedenceAndSessionAllow(t *testing.T) {
|
||||
rules := []Rule{{Effect: EffectAllow, Tool: "write_file"}, {Effect: EffectDeny, Tool: "write_*", ArgGlob: "*secret*"}}
|
||||
if got := Decide(Request{Name: "write_file", Args: `{"path":"secret.txt"}`}, rules); got != EffectDeny {
|
||||
t.Fatalf("deny lost: %q", got)
|
||||
}
|
||||
if got := Decide(Request{Name: "http.fetch", SessionAllow: true}, []Rule{{Effect: EffectAsk, Tool: "http.*"}}); got != EffectAllow {
|
||||
t.Fatalf("session allow %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutoEditAllowsWorkspaceWrites(t *testing.T) {
|
||||
if got := Decide(Request{Mode: ModeAutoEdit, Name: "fs_write_file"}, nil); got != EffectAllow {
|
||||
t.Fatalf("auto edit %q", got)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user