From 951c141f95b4d6ac0f82804dd5ae9b7a6d7ad31e Mon Sep 17 00:00:00 2001 From: weichen Date: Thu, 1 Oct 2026 08:57:21 +0000 Subject: [PATCH] Let hosts register tools and observe token usage Why: product runners need register_artifact and card tools inside the kernel loop, and the UI needs token usage on the same event stream as text and tool calls. Evidence: go test ./... in agent-native; TestHostToolAndUsageEvent and TestHostToolNameConflictFailsClosed. Impact: Options.HostTools fail closed on name clashes; ProcessEvent and runtime events gain a usage type with token counts. Existing callers that ignore unknown event types stay compatible. Refs: deliverable-quality-system B1 Co-authored-by: Cursor --- adapters/native/backend.go | 18 ++++- adapters/native/host_tools_test.go | 119 +++++++++++++++++++++++++++++ exec/backend.go | 9 ++- runtime/loop.go | 8 ++ runtime/tools/registry.go | 5 ++ 5 files changed, 152 insertions(+), 7 deletions(-) create mode 100644 adapters/native/host_tools_test.go diff --git a/adapters/native/backend.go b/adapters/native/backend.go index 7a7fcd6..eb86098 100644 --- a/adapters/native/backend.go +++ b/adapters/native/backend.go @@ -50,6 +50,8 @@ type Options struct { Compact *CompactConfig // Subagent 可选：提供时注册 delegate_task（委派由宿主实现，内核只开端口）。 Subagent subagent.Factory + // HostTools 是宿主产品工具（登记工件、卡片、知识库）。与内核工具同名时失败，不覆盖。 + HostTools []tools.Tool // Log 可选。nil 时丢弃日志。 Log *logsys.Logger // ExecutionID 写入日志，不参与执行控制。 @@ -168,6 +170,12 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, } } subagent.Register(reg, b.opts.Subagent) + for _, tool := range b.opts.HostTools { + if reg.Has(tool.Name()) { + return &exec.Result{Success: false, ErrorReason: "host tool name conflicts with a builtin: " + tool.Name()}, nil + } + reg.Register(tool) + } if err := reg.Restrict(req.ToolAllowlist); err != nil { return &exec.Result{Success: false, ErrorReason: err.Error()}, nil } @@ -313,10 +321,12 @@ func (b *Backend) run(ctx context.Context, model string, budget *runtime.Budget, logs.WriteByte('\n') if onProcess != nil { onProcess(exec.ProcessEvent{ - Type: ev.Type, - Text: ev.Text, - Tool: ev.Tool, - Detail: ev.Detail, + Type: ev.Type, + Text: ev.Text, + Tool: ev.Tool, + Detail: ev.Detail, + PromptTokens: ev.PromptTokens, + CompletionTokens: ev.CompletionTokens, }) } }, diff --git a/adapters/native/host_tools_test.go b/adapters/native/host_tools_test.go new file mode 100644 index 0000000..017fd8e --- /dev/null +++ b/adapters/native/host_tools_test.go @@ -0,0 +1,119 @@ +package native + +import ( + "context" + "encoding/json" + "sync" + "testing" + + "agent-native/exec" + "agent-native/runtime/llm" + "agent-native/runtime/tools" +) + +type hostTool struct{ name string } + +func (h hostTool) Name() string { return h.name } +func (h hostTool) Description() string { return "host tool" } +func (h hostTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`) +} +func (h hostTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) { + return tools.Result{Content: "registered " + h.name}, nil +} + +func TestHostToolAndUsageEvent(t *testing.T) { + var events []exec.ProcessEvent + b := New(Options{ + Model: "fake", + HostTools: []tools.Tool{hostTool{name: "register_artifact"}}, + Stream: usageScript(), + }) + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "登记", + OnEvent: func(ev exec.ProcessEvent) { events = append(events, ev) }, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v logs=%s", err, res, logsOf(res)) + } + sawText, sawStart, sawEnd, sawUsage := false, false, false, false + for _, ev := range events { + switch ev.Type { + case "text_delta": + if ev.Text != "" { + sawText = true + } + case "tool_execution_start": + if ev.Tool == "register_artifact" { + sawStart = true + } + case "tool_execution_end": + if ev.Tool == "register_artifact" { + sawEnd = true + } + case "usage": + if ev.PromptTokens == 11 && ev.CompletionTokens == 7 { + sawUsage = true + } + } + } + if !sawText || !sawStart || !sawEnd || !sawUsage { + t.Fatalf("events missing text/tool/usage: %+v", events) + } +} + +func TestHostToolNameConflictFailsClosed(t *testing.T) { + b := New(Options{ + Model: "fake", + Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "x", StopReason: llm.StopEnd}), + HostTools: []tools.Tool{hostTool{name: "read_file"}}, + }) + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + }) + if err != nil { + t.Fatal(err) + } + if res.Success || res.ErrorReason == "" { + t.Fatalf("conflict must fail closed, got %+v", res) + } +} + +func usageScript() llm.StreamFunc { + var mu sync.Mutex + step := 0 + return func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + mu.Lock() + defer mu.Unlock() + step++ + switch step { + case 1: + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{ + Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, + ToolCalls: []llm.ToolCall{{ID: "c1", Name: "register_artifact", Arguments: `{"path":"a.docx"}`}}, + }, Usage: llm.Usage{PromptTokens: 4, CompletionTokens: 2}} + case 2: + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{ + Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, + ToolCalls: []llm.ToolCall{{ID: "c2", Name: "write_file", Arguments: `{"path":"output.md","content":"done"}`}}, + }, Usage: llm.Usage{PromptTokens: 4, CompletionTokens: 2}} + default: + ch <- llm.Event{Kind: llm.KindTextDelta, Text: "已登记"} + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{ + Role: llm.RoleAssistant, Content: "已登记", StopReason: llm.StopEnd, + }, Usage: llm.Usage{PromptTokens: 3, CompletionTokens: 3}} + } + }() + return ch + } +} + +func logsOf(res *exec.Result) string { + if res == nil { + return "" + } + return res.Logs +} diff --git a/exec/backend.go b/exec/backend.go index 3c71e4c..85432c2 100644 --- a/exec/backend.go +++ b/exec/backend.go @@ -40,9 +40,9 @@ type Request struct { // 共享总预算与截止时间，避免传输、步骤与主循环次数相乘）。 // 全部字段可选；零值字段不设限。 type Budget struct { - MaxPromptTokens int `json:"max_prompt_tokens,omitempty"` - MaxCompletionTokens int `json:"max_completion_tokens,omitempty"` - MaxToolCalls int `json:"max_tool_calls,omitempty"` + 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"` } @@ -54,6 +54,9 @@ type ProcessEvent struct { Text string // delta 文本（Type 为 *_delta 时） Tool string // 工具名（tool_* 时） Detail string // 工具参数/结果摘要（tool_execution_* 时） + // PromptTokens / CompletionTokens 在 Type=usage 时填写。 + PromptTokens int + CompletionTokens int } const ( diff --git a/runtime/loop.go b/runtime/loop.go index 3021db8..88a52a8 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -29,6 +29,9 @@ type Event struct { // Detail carries tool arguments/results or reasoning deltas for // process-visualization consumers; log writers may ignore it. Detail string + // PromptTokens / CompletionTokens 只在 type=usage 时有意义。 + PromptTokens int + CompletionTokens int } // Config drives one agent session. @@ -164,6 +167,7 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { continue } emit(cfg, Event{Type: "agent_end", Message: last}) + emitUsage(cfg, usage) return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: usage}, nil } errCap := errString("max turns exceeded") @@ -296,6 +300,10 @@ func drain(fn func() []llm.Message) []llm.Message { return fn() } +func emitUsage(cfg Config, usage llm.Usage) { + emit(cfg, Event{Type: "usage", PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens}) +} + func emit(cfg Config, ev Event) { if cfg.OnEvent != nil { cfg.OnEvent(ev) diff --git a/runtime/tools/registry.go b/runtime/tools/registry.go index 17ad4c8..f418af5 100644 --- a/runtime/tools/registry.go +++ b/runtime/tools/registry.go @@ -45,6 +45,11 @@ func NewRegistry(root string) *Registry { func (r *Registry) Root() string { return r.root } +func (r *Registry) Has(name string) bool { + _, ok := r.tools[name] + return ok +} + func (r *Registry) Register(t Tool) { if _, ok := r.tools[t.Name()]; !ok { r.order = append(r.order, t.Name())