Extend exec.Request with capability fields and enforce budgets end to end
Why: the frozen exec port could not express per-request model, tool allowlist, skills, budget, or identifiers, so squad roles and product profiles had no way to steer one execution; contract v0.1.0 semantics (budget/cancel) stayed unconsumed (parity-matrix #4/#10/#16). Evidence: go vet ./... clean; go test ./... all ok, including new loop budget caps (tool calls, prompt tokens, deadline), cancel propagation into a blocking tool, registry Restrict fail-closed policy tests, native adapter v2 tests (model override, allowlist, deadline-bound context, execution-id log header, skills prompt line), and settle classification marking budget_exceeded permanent. Impact: exec.Request gains optional Model/ToolAllowlist/Skills/Budget/ SessionScope/ExecutionID/Env (zero value = legacy behavior); runtime loop enforces a shared budget and emits budget_exceeded events; adapters/native consumes all new fields. Consumers on the old port keep compiling and behaving identically. Refs: docs/parity-matrix.md rows #4/#10/#16 (Batch2 scope); deps/ development-guidelines standards/agent-execution.md (shared budget across retry layers).
This commit is contained in:
@@ -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(".")
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 过程帧的统一投影,
|
||||
|
||||
@@ -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, "治理拦截")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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) }
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user