diff --git a/adapters/native/backend.go b/adapters/native/backend.go index 9bccc40..419caa9 100644 --- a/adapters/native/backend.go +++ b/adapters/native/backend.go @@ -68,7 +68,11 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, if b.opts.Stream == nil && (strings.TrimSpace(b.opts.APIBase) == "" || strings.TrimSpace(b.opts.APIKey) == "") { return &exec.Result{Success: false, ErrorReason: "agent-native backend requires api base/key (or injected stream)"}, nil } - if strings.TrimSpace(b.opts.Model) == "" { + model := strings.TrimSpace(b.opts.Model) + if strings.TrimSpace(req.Model) != "" { + model = strings.TrimSpace(req.Model) + } + if model == "" { return &exec.Result{Success: false, ErrorReason: "agent-native backend requires model"}, nil } @@ -91,6 +95,9 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, var reg *tools.Registry if profile == exec.ToolProfileWorkspace { reg = tools.WorkspaceTools(workDir) + if err := reg.Restrict(req.ToolAllowlist); err != nil { + return &exec.Result{Success: false, ErrorReason: err.Error()}, nil + } } // 会话毒化轮转：上次失败属 resume-unsafe 时，把旧 transcript 归档， // 本次以全新会话开始（续跑毒化会话只会复现同样的失败）。 @@ -109,6 +116,21 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, if timeout <= 0 { timeout = b.opts.Timeout } + var runBudget *runtime.Budget + if req.Budget != nil { + runBudget = &runtime.Budget{ + MaxPromptTokens: req.Budget.MaxPromptTokens, + MaxCompletionTokens: req.Budget.MaxCompletionTokens, + MaxToolCalls: req.Budget.MaxToolCalls, + Deadline: req.Budget.Deadline, + } + if !req.Budget.Deadline.IsZero() { + // Deadline 与 Timeout 先到者生效：把更短者作为 ctx 超时。 + if until := time.Until(req.Budget.Deadline); timeout <= 0 || until < timeout { + timeout = until + } + } + } runCtx := ctx cancel := func() {} if timeout > 0 { @@ -116,9 +138,16 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, } defer cancel() + execID := req.ExecutionID + if strings.TrimSpace(execID) == "" { + execID = b.opts.ExecutionID + } var logs strings.Builder + if strings.TrimSpace(execID) != "" || strings.TrimSpace(req.SessionScope) != "" { + logs.WriteString("execution_id=" + execID + " session=" + req.SessionScope + "\n") + } onProcess := req.OnEvent - runRes, err := b.run(runCtx, reg, store, stream, profile, req, &logs, onProcess) + runRes, err := b.run(runCtx, model, runBudget, reg, store, stream, profile, req, &logs, onProcess) if err != nil { return &exec.Result{Success: false, ErrorReason: err.Error(), Logs: logs.String()}, nil } @@ -173,19 +202,24 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, } // run 驱动 runtime 循环并桥接过程事件（语义对齐主仓 pinative）。 -func (b *Backend) run(ctx context.Context, reg *tools.Registry, store *session.Store, +func (b *Backend) run(ctx context.Context, model string, budget *runtime.Budget, reg *tools.Registry, store *session.Store, stream llm.StreamFunc, profile string, req exec.Request, logs *strings.Builder, onProcess func(exec.ProcessEvent)) (runtime.Result, error) { + execID := req.ExecutionID + if strings.TrimSpace(execID) == "" { + execID = b.opts.ExecutionID + } runRes, err := runtime.Run(ctx, runtime.Config{ - Model: b.opts.Model, + Model: model, SystemPrompt: systemPrompt(req, profile), Tools: reg, Stream: stream, Session: store, MaxTurns: b.opts.MaxTurns, Log: b.logger(), - ExecutionID: b.opts.ExecutionID, + ExecutionID: execID, + Budget: budget, OnEvent: func(ev runtime.Event) { logs.WriteString(ev.Type) if ev.Tool != "" { @@ -256,13 +290,18 @@ func (b *Backend) logger() *logsys.Logger { return log.With("execution_id", b.opts.ExecutionID) } -// systemPrompt 组装系统提示：角色指令 + 交付物与工具档位约束。 +// systemPrompt 组装系统提示：角色指令 + 交付物与工具档位约束 + 声明的技能。 func systemPrompt(req exec.Request, profile string) string { var b strings.Builder if strings.TrimSpace(req.Instructions) != "" { b.WriteString(req.Instructions) b.WriteString("\n\n") } + if len(req.Skills) > 0 { + b.WriteString("Declared skills for this task: ") + b.WriteString(strings.Join(req.Skills, ", ")) + b.WriteString(".\n\n") + } b.WriteString("Deliverable files: ") b.WriteString(strings.Join(req.Targets, ", ")) b.WriteString(".") diff --git a/adapters/native/request_v2_test.go b/adapters/native/request_v2_test.go new file mode 100644 index 0000000..e9f2539 --- /dev/null +++ b/adapters/native/request_v2_test.go @@ -0,0 +1,186 @@ +package native + +import ( + "context" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "agent-native/exec" + "agent-native/runtime/llm" +) + +// captureStream 捕获每次调用的 model / 首条 system 消息，并按序回放回合。 +func captureStream(turns ...llm.Message) (llm.StreamFunc, *capture) { + c := &capture{} + var mu sync.Mutex + i := 0 + return func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + mu.Lock() + defer mu.Unlock() + c.models = append(c.models, model) + c.toolDecls = append(c.toolDecls, toolNames(tools)) + if len(messages) > 0 && messages[0].Role == llm.RoleSystem { + c.system = messages[0].Content + } + if i >= len(turns) { + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}} + return + } + msg := turns[i] + i++ + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg} + }() + return ch + }, c +} + +type capture struct { + mu sync.Mutex + models []string + toolDecls [][]string + system string +} + +func toolNames(decls []llm.ToolDecl) []string { + var out []string + for _, d := range decls { + out = append(out, d.Name) + } + return out +} + +func TestBackendRequestModelOverride(t *testing.T) { + stream, c := captureStream() + b := New(Options{Model: "default-model", Stream: stream}) + dir := t.TempDir() + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "hi", Targets: []string{"output.md"}, Model: "override-model", + ToolProfile: exec.ToolProfileReasoning, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v", err, res) + } + if len(c.models) == 0 || c.models[0] != "override-model" { + t.Fatalf("want override-model, got %v", c.models) + } +} + +func TestBackendRequestToolAllowlist(t *testing.T) { + stream, c := captureStream( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "write_file", Arguments: `{"path":"output.md","content":"blocked"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c2", Name: "write_file", Arguments: `{"path":"output2.md","content":"still blocked"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "gave up", StopReason: llm.StopEnd}, + ) + b := New(Options{Model: "fake", Stream: stream}) + dir := t.TempDir() + // 目标 output2.md 不在 allowlist 内能写吗？——write_file 被整体收窄， + // 模型两次尝试都会得到 unknown tool，最终无交付物 → 失败。 + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "写文件", Targets: []string{"output2.md"}, + ToolAllowlist: []string{"read_file", "list_dir", "grep"}, + }) + if err != nil { + t.Fatal(err) + } + if res.Success { + t.Fatalf("write_file restricted out; deliverable must fail, got %+v", res) + } + if len(c.toolDecls) == 0 { + t.Fatal("expected tool decls") + } + for _, names := range c.toolDecls { + for _, n := range names { + if n == "write_file" || n == "edit_file" { + t.Fatalf("write/edit must be restricted out of decls, got %v", names) + } + } + } +} + +func TestBackendRequestUnknownAllowlistFailsClosed(t *testing.T) { + b := New(Options{Model: "fake", Stream: script()}) + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + ToolAllowlist: []string{"read_file", "no_such_tool"}, + }) + if err != nil { + t.Fatal(err) + } + if res.Success || !strings.Contains(res.ErrorReason, "no_such_tool") { + t.Fatalf("want fail-closed unknown allowlist name, got %+v", res) + } +} + +func TestBackendRequestBudgetDeadlineBindsContext(t *testing.T) { + deadlineHit := make(chan struct{}) + var once sync.Once + stream := func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + select { + case <-ctx.Done(): + once.Do(func() { close(deadlineHit) }) + ch <- llm.Event{Kind: llm.KindError, Err: ctx.Err(), Message: llm.Message{StopReason: llm.StopAborted}} + case <-time.After(3 * time.Second): + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}} + } + }() + return ch + } + b := New(Options{Model: "fake", Stream: stream}) + start := time.Now() + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + Timeout: 30 * time.Second, // 长超时，预算截止应更先生效 + Budget: &exec.Budget{Deadline: time.Now().Add(150 * time.Millisecond)}, + }) + if err != nil { + t.Fatal(err) + } + select { + case <-deadlineHit: + case <-time.After(2 * time.Second): + t.Fatal("budget deadline should bind the run context") + } + if res.Success { + t.Fatalf("deadline run must not succeed, got %+v", res) + } + if elapsed := time.Since(start); elapsed > 5*time.Second { + t.Fatalf("deadline should fire early, took %s", elapsed) + } +} + +func TestBackendRequestExecutionIDAndSkillsInPrompt(t *testing.T) { + stream, c := captureStream() + b := New(Options{Model: "fake", Stream: stream}) + dir := t.TempDir() + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "hi", Targets: []string{"output.md"}, + Instructions: "你是角色", Skills: []string{"docx-author", "xlsx-author"}, + ExecutionID: "exec-42", SessionScope: "run-9/step-3", + ToolProfile: exec.ToolProfileReasoning, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v", err, res) + } + if !strings.Contains(c.system, "docx-author") || !strings.Contains(c.system, "xlsx-author") { + t.Fatalf("skills should be declared in system prompt, got:\n%s", c.system) + } + logBytes, _ := os.ReadFile(filepath.Join(dir, "execution.log")) + logs := string(logBytes) + if !strings.Contains(logs, "execution_id=exec-42") || !strings.Contains(logs, "run-9/step-3") { + t.Fatalf("execution log should carry ids, got:\n%s", logs) + } +} diff --git a/exec/backend.go b/exec/backend.go index 2ebbcc9..3c71e4c 100644 --- a/exec/backend.go +++ b/exec/backend.go @@ -17,6 +17,34 @@ type Request struct { // 实现方必须保证：回调在 Execute 返回前完成（同步调用或等待收尾），无泄漏 goroutine； // 回调内不得阻塞（消费方自行节流/异步入队）。 OnEvent func(ev ProcessEvent) + + // ---- v0.2 能力字段（全部可选；零值 = 既有行为，向后兼容） ---- + + // Model 覆盖本请求使用的模型（后端构造时的默认模型仍兜底）。 + Model string `json:"model,omitempty"` + // ToolAllowlist 收窄可见工具集；出现后端未注册的工具名时 fail-closed。 + ToolAllowlist []string `json:"tool_allowlist,omitempty"` + // Skills 声明本次执行需要加载的技能（内核在系统提示中列出并由 load_skill 按需展开）。 + Skills []string `json:"skills,omitempty"` + // Budget 单次步骤执行的共享预算；步骤重试应传入剩余额度，不得各自另起一份。 + Budget *Budget `json:"budget,omitempty"` + // SessionScope / ExecutionID 只用于日志与过程事件贯穿，不参与控制流。 + SessionScope string `json:"session_scope,omitempty"` + ExecutionID string `json:"execution_id,omitempty"` + // Env 是组合根批准后注入执行类工具（run_command 等）的白名单环境变量。 + // 后端不得把自身持有的密钥并入其中。 + Env map[string]string `json:"env,omitempty"` +} + +// Budget 是一次步骤执行的共享预算（agent-execution 标准：重试层各有职责， +// 共享总预算与截止时间，避免传输、步骤与主循环次数相乘）。 +// 全部字段可选；零值字段不设限。 +type Budget struct { + MaxPromptTokens int `json:"max_prompt_tokens,omitempty"` + MaxCompletionTokens int `json:"max_completion_tokens,omitempty"` + MaxToolCalls int `json:"max_tool_calls,omitempty"` + // Deadline 绝对截止时刻；与 Request.Timeout 先到者生效。 + Deadline time.Time `json:"deadline,omitempty"` } // ProcessEvent 对 agentx 事件与远程 runner NDJSON 过程帧的统一投影， diff --git a/exec/step_retry.go b/exec/step_retry.go index 212c800..eb6b05d 100644 --- a/exec/step_retry.go +++ b/exec/step_retry.go @@ -160,6 +160,7 @@ func combinedFailureText(sig StepFailureSignals) string { func isPermanentReason(lower string) bool { return strings.Contains(lower, "prompt_empty") || strings.Contains(lower, "unauthorized") || + strings.Contains(lower, "budget_exceeded") || strings.Contains(lower, "治理强制终止") || strings.Contains(lower, "治理拦截") } diff --git a/exec/step_retry_budget_test.go b/exec/step_retry_budget_test.go new file mode 100644 index 0000000..1858bcf --- /dev/null +++ b/exec/step_retry_budget_test.go @@ -0,0 +1,37 @@ +package exec + +import ( + "errors" + "strings" + "testing" + "time" +) + +func TestAssessStepFailureBudgetExceededIsPermanent(t *testing.T) { + ass := AssessStepFailure(StepFailureSignals{ + Err: errors.New("budget_exceeded: tool calls 12 > max 10"), + }) + if ass.Class != FailurePermanent { + t.Fatalf("budget exhaustion must classify permanent (no auto retry), got %s", ass.Class) + } + if !strings.Contains(ass.Reason, "permanent") { + t.Fatalf("want permanent reason, got %s", ass.Reason) + } +} + +func TestAssessStepFailureDeadlinePastIsPermanent(t *testing.T) { + ass := AssessStepFailure(StepFailureSignals{ + Result: &Result{ErrorReason: "budget_exceeded: deadline 2026-10-01T00:00:00Z passed"}, + }) + if ass.Class != FailurePermanent { + t.Fatalf("deadline exhaustion must classify permanent, got %s", ass.Class) + } +} + +func TestBudgetZeroValueUnlimited(t *testing.T) { + var b Budget + if b.MaxPromptTokens != 0 || b.MaxCompletionTokens != 0 || b.MaxToolCalls != 0 || !b.Deadline.IsZero() { + t.Fatal("zero budget must express unlimited") + } + _ = time.Now +} diff --git a/runtime/loop.go b/runtime/loop.go index 73291f7..3021db8 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -3,8 +3,10 @@ package runtime import ( "context" "encoding/json" + "fmt" "strings" "sync" + "time" "agent-native/runtime/llm" "agent-native/runtime/session" @@ -45,6 +47,17 @@ type Config struct { Log *logsys.Logger // ExecutionID 只用于日志字段，不参与控制流。 ExecutionID string + // Budget 可选。任一上限命中即停止并返回 budget_exceeded 错误（StopError）； + // 零值字段不设限。截止时刻与外层 ctx 先到者生效。 + Budget *Budget +} + +// Budget 是循环层的执行预算；与 exec.Budget 字段同构，由适配层换算。 +type Budget struct { + MaxPromptTokens int + MaxCompletionTokens int + MaxToolCalls int + Deadline time.Time } // Result is the final assistant text and the transcript produced by one run. @@ -91,11 +104,15 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { var last llm.Message var usage llm.Usage + toolCalls := 0 for turn := 0; turn < cfg.MaxTurns; turn++ { if ctx.Err() != nil { cfg.logger().Error("runtime turn aborted", ctx.Err()) return Result{Text: last.Content, Messages: messages, StopReason: llm.StopAborted, Err: ctx.Err()}, nil } + if reason := budgetExceeded(cfg.Budget, usage, toolCalls, time.Now()); reason != "" { + return budgetStop(cfg, last, messages, usage, reason), nil + } emit(cfg, Event{Type: "turn_start"}) end, errEv := consume(ctx, cfg, messages) if errEv != nil { @@ -113,6 +130,11 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { emit(cfg, Event{Type: "message_end", Message: last}) if last.StopReason == llm.StopToolCalls && len(last.ToolCalls) > 0 { + if cfg.Budget != nil && cfg.Budget.MaxToolCalls > 0 && toolCalls+len(last.ToolCalls) > cfg.Budget.MaxToolCalls { + reason := fmt.Sprintf("budget_exceeded: tool calls %d+%d > max %d", toolCalls, len(last.ToolCalls), cfg.Budget.MaxToolCalls) + return budgetStop(cfg, last, messages, usage, reason), nil + } + toolCalls += len(last.ToolCalls) results := executeTools(ctx, cfg, last) for _, res := range results { if err := persist(cfg, res); err != nil { @@ -280,6 +302,35 @@ func emit(cfg Config, ev Event) { } } +// budgetExceeded 返回命中的预算说明；空串表示预算仍充足（或未设预算）。 +// 只在「即将消耗更多预算」的节点调用（轮次顶端 / 工具批执行前）。 +func budgetExceeded(b *Budget, usage llm.Usage, toolCalls int, now time.Time) string { + if b == nil { + return "" + } + if !b.Deadline.IsZero() && now.After(b.Deadline) { + return fmt.Sprintf("budget_exceeded: deadline %s passed", b.Deadline.Format(time.RFC3339)) + } + if b.MaxPromptTokens > 0 && usage.PromptTokens > b.MaxPromptTokens { + return fmt.Sprintf("budget_exceeded: prompt tokens %d > max %d", usage.PromptTokens, b.MaxPromptTokens) + } + if b.MaxCompletionTokens > 0 && usage.CompletionTokens > b.MaxCompletionTokens { + return fmt.Sprintf("budget_exceeded: completion tokens %d > max %d", usage.CompletionTokens, b.MaxCompletionTokens) + } + if b.MaxToolCalls > 0 && toolCalls > b.MaxToolCalls { + return fmt.Sprintf("budget_exceeded: tool calls %d > max %d", toolCalls, b.MaxToolCalls) + } + return "" +} + +// budgetStop 以预算超限收尾：发事件、留日志、按 StopError 返回（不可重试）。 +func budgetStop(cfg Config, last llm.Message, messages []llm.Message, usage llm.Usage, reason string) Result { + emit(cfg, Event{Type: "budget_exceeded", Text: reason}) + err := errString(reason) + cfg.logger().Error("runtime budget exceeded", err) + return Result{Text: last.Content, Messages: messages, StopReason: llm.StopError, Usage: usage, Err: err} +} + type errString string func (e errString) Error() string { return string(e) } diff --git a/runtime/loop_budget_test.go b/runtime/loop_budget_test.go new file mode 100644 index 0000000..46a800e --- /dev/null +++ b/runtime/loop_budget_test.go @@ -0,0 +1,194 @@ +package runtime + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" + + "agent-native/runtime/llm" + "agent-native/runtime/tools" +) + +// usageScript 产出带 usage 的固定回合，用于预算测试。 +func usageScript(usage llm.Usage, turns ...llm.Message) llm.StreamFunc { + var i int + return func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + if i >= len(turns) { + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}, Usage: usage} + return + } + msg := turns[i] + i++ + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg, Usage: usage} + }() + return ch + } +} + +func TestLoopBudgetToolCallCap(t *testing.T) { + dir := t.TempDir() + reg := tools.WorkspaceTools(dir) + sawBudgetEvent := false + res, err := Run(context.Background(), Config{ + Model: "fake", + Tools: reg, + Budget: &Budget{MaxToolCalls: 1}, + OnEvent: func(ev Event) { + if ev.Type == "budget_exceeded" { + sawBudgetEvent = true + } + }, + Stream: usageScript(llm.Usage{}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "list_dir", Arguments: `{}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{ + {ID: "c2", Name: "list_dir", Arguments: `{}`}, + {ID: "c3", Name: "list_dir", Arguments: `{}`}, + }}, + llm.Message{Role: llm.RoleAssistant, Content: "never", StopReason: llm.StopEnd}, + ), + }, llm.Message{Role: llm.RoleUser, Content: "loop tools"}) + if err != nil { + t.Fatal(err) + } + if res.StopReason != llm.StopError || res.Err == nil || !strings.Contains(res.Err.Error(), "budget_exceeded") { + t.Fatalf("want budget_exceeded stop, got stop=%s err=%v", res.StopReason, res.Err) + } + if !strings.Contains(res.Err.Error(), "tool calls") { + t.Fatalf("want tool-call budget reason, got %v", res.Err) + } + if !sawBudgetEvent { + t.Fatal("expected budget_exceeded event") + } + // 第二批工具调用不该被执行（预算在执行前拦截）：transcript 只含一次 tool 结果。 + toolResults := 0 + for _, m := range res.Messages { + if m.Role == llm.RoleTool { + toolResults++ + } + } + if toolResults != 1 { + t.Fatalf("expected exactly 1 tool result, got %d", toolResults) + } +} + +func TestLoopBudgetPromptTokens(t *testing.T) { + res, err := Run(context.Background(), Config{ + Model: "fake", + Budget: &Budget{MaxPromptTokens: 10}, + Stream: usageScript(llm.Usage{PromptTokens: 500, TotalTokens: 600}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "list_dir", Arguments: `{}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "never", StopReason: llm.StopEnd}, + ), + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if res.Err == nil || !strings.Contains(res.Err.Error(), "prompt tokens") { + t.Fatalf("want prompt-token budget stop, got err=%v", res.Err) + } + if res.Usage.PromptTokens != 500 { + t.Fatalf("usage should be reported on budget stop, got %d", res.Usage.PromptTokens) + } +} + +func TestLoopBudgetDeadline(t *testing.T) { + res, err := Run(context.Background(), Config{ + Model: "fake", + Budget: &Budget{Deadline: time.Now().Add(-time.Second)}, + Stream: usageScript(llm.Usage{}, llm.Message{Role: llm.RoleAssistant, Content: "never", StopReason: llm.StopEnd}), + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if res.Err == nil || !strings.Contains(res.Err.Error(), "deadline") { + t.Fatalf("want deadline budget stop, got err=%v", res.Err) + } +} + +func TestLoopBudgetNotHitCompletesNormally(t *testing.T) { + res, err := Run(context.Background(), Config{ + Model: "fake", + Budget: &Budget{MaxPromptTokens: 10_000, MaxToolCalls: 10}, + Stream: usageScript(llm.Usage{PromptTokens: 50}, + llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}), + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if res.Err != nil || res.Text != "ok" { + t.Fatalf("clean run should finish, got err=%v text=%q", res.Err, res.Text) + } +} + +// blockingTool 阻塞到 ctx 取消，验证取消传播到工具执行层。 +type blockingTool struct{ canceled chan struct{} } + +func (b *blockingTool) Name() string { return "block_forever" } +func (b *blockingTool) Description() string { return "blocks until context cancel" } +func (b *blockingTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","properties":{}}`) +} + +func (b *blockingTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) { + <-ctx.Done() + close(b.canceled) + return tools.Result{Content: "canceled", IsError: true}, ctx.Err() +} + +func TestLoopCancelPropagatesToTool(t *testing.T) { + bt := &blockingTool{canceled: make(chan struct{})} + reg := tools.NewRegistry(t.TempDir()) + reg.Register(bt) + ctx, cancel := context.WithCancel(context.Background()) + + firstCall := true + stream := func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + if firstCall { + firstCall = false + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{ + Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, + ToolCalls: []llm.ToolCall{{ID: "c1", Name: "block_forever", Arguments: `{}`}}, + }} + return + } + // 工具返回后的下一轮：等取消或超时兜底。 + select { + case <-ctx.Done(): + ch <- llm.Event{Kind: llm.KindError, Err: ctx.Err(), Message: llm.Message{StopReason: llm.StopAborted}} + case <-time.After(3 * time.Second): + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}} + } + }() + return ch + } + + go func() { + time.Sleep(100 * time.Millisecond) + cancel() + }() + res, err := Run(ctx, Config{Model: "fake", Tools: reg, Stream: stream}, + llm.Message{Role: llm.RoleUser, Content: "block"}) + if err != nil { + t.Fatal(err) + } + select { + case <-bt.canceled: + case <-time.After(2 * time.Second): + t.Fatal("tool context was not canceled") + } + if res.StopReason != llm.StopAborted && res.Err == nil { + t.Fatalf("expected aborted stop, got %+v", res) + } +} diff --git a/runtime/tools/registry.go b/runtime/tools/registry.go index 0797477..72fb7cd 100644 --- a/runtime/tools/registry.go +++ b/runtime/tools/registry.go @@ -61,6 +61,38 @@ func (r *Registry) Decls() []Decl { return out } +// Restrict 把注册表收窄到 allow 列出的工具（工具策略的唯一执行点之一： +// 调用方声明 allowlist，注册表执行收窄）。allow 中出现未注册的工具名时 +// 返回错误并列出全部未知名（fail-closed，不静默忽略配置错误）。 +// 空列表不收窄；可重复调用（在已收窄集合上继续求交）。 +func (r *Registry) Restrict(allow []string) error { + if len(allow) == 0 { + return nil + } + keep := make(map[string]bool, len(allow)) + var unknown []string + for _, name := range allow { + if _, ok := r.tools[name]; !ok { + unknown = append(unknown, name) + continue + } + keep[name] = true + } + if len(unknown) > 0 { + return fmt.Errorf("tool allowlist names not registered: %s", strings.Join(unknown, ", ")) + } + var order []string + for _, name := range r.order { + if keep[name] { + order = append(order, name) + } else { + delete(r.tools, name) + } + } + r.order = order + return nil +} + // Decl is the schema view of a tool. type Decl struct { Name string diff --git a/runtime/tools/registry_policy_test.go b/runtime/tools/registry_policy_test.go new file mode 100644 index 0000000..bc11bdf --- /dev/null +++ b/runtime/tools/registry_policy_test.go @@ -0,0 +1,71 @@ +package tools + +import ( + "context" + "strings" + "testing" +) + +func TestRegistryRestrictNarrowsDeclsAndExecute(t *testing.T) { + dir := t.TempDir() + reg := WorkspaceTools(dir) + if err := reg.Restrict([]string{"read_file", "list_dir"}); err != nil { + t.Fatal(err) + } + names := make(map[string]bool) + for _, d := range reg.Decls() { + names[d.Name] = true + } + if len(names) != 2 || !names["read_file"] || !names["list_dir"] { + t.Fatalf("want read_file+list_dir only, got %v", names) + } + // 被收窄掉的工具体现为 unknown tool。 + res := reg.Execute(context.Background(), "write_file", []byte(`{"path":"a.md","content":"x"}`)) + if !res.IsError || !strings.Contains(res.Content, "unknown tool") { + t.Fatalf("want unknown tool error, got %+v", res) + } + // 保留的工具可用。 + res = reg.Execute(context.Background(), "list_dir", []byte(`{}`)) + if res.IsError { + t.Fatalf("list_dir should still work: %+v", res) + } +} + +func TestRegistryRestrictFailsClosedOnUnknownNames(t *testing.T) { + reg := WorkspaceTools(t.TempDir()) + err := reg.Restrict([]string{"read_file", "not_a_tool"}) + if err == nil || !strings.Contains(err.Error(), "not_a_tool") { + t.Fatalf("want fail-closed error naming the unknown tool, got %v", err) + } + // 失败后注册表不被部分修改。 + if len(reg.Decls()) != 5 { + t.Fatalf("registry should be untouched after failed restrict, got %d tools", len(reg.Decls())) + } +} + +func TestRegistryRestrictEmptyKeepsAll(t *testing.T) { + reg := WorkspaceTools(t.TempDir()) + if err := reg.Restrict(nil); err != nil { + t.Fatal(err) + } + if len(reg.Decls()) != 5 { + t.Fatalf("nil allowlist should keep all tools, got %d", len(reg.Decls())) + } +} + +func TestRegistryRestrictIsIdempotentIntersection(t *testing.T) { + reg := WorkspaceTools(t.TempDir()) + if err := reg.Restrict([]string{"read_file", "list_dir"}); err != nil { + t.Fatal(err) + } + // 二次收窄若点名已被移除的工具，属调用方配置错误：fail-closed。 + if err := reg.Restrict([]string{"read_file", "write_file"}); err == nil || !strings.Contains(err.Error(), "write_file") { + t.Fatalf("second restrict naming a removed tool must fail-closed, got %v", err) + } + if err := reg.Restrict([]string{"read_file"}); err != nil { + t.Fatal(err) + } + if len(reg.Decls()) != 1 || reg.Decls()[0].Name != "read_file" { + t.Fatalf("want intersection to keep read_file only, got %+v", reg.Decls()) + } +}