Compare commits

...
Author SHA1 Message Date
weichen 16a164a594 Merge remote-tracking branch 'origin/main' into feature/sillytavern-review
ci / vet-and-test (pull_request) Successful in 49s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s
2026-10-04 19:42:54 +00:00
weichen 3a110a728c Shape provider messages before they are sent
ci / vet-and-test (pull_request) Successful in 51s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s
Why: Some endpoints reject consecutive system messages or a run of the same role. The product already emits several system messages in one request.
Evidence: go test is covered by PrepareMessages tests for merge, strict folding, and tool pairs. Empty MessageShape leaves the transcript unchanged.
Impact: Callers opt in with StreamOptions.MessageShape. Tool calls and tool results are not merged.
Refs: docs/records/2026-10-04-sillytavern-review.md
2026-10-04 18:18:57 +00:00
weichen ecca2d74b9 Merge pull request 'Add a script host, deferred tools, and background MCP connect' (#15) from feature/code-orchestration into main
ci / vet-and-test (push) Successful in 48s
ci / docker-checkout-hostgateway (push) Successful in 1s
2026-10-03 18:07:11 +00:00
weichenandCursor e4020d297c Add a script host, deferred tools, and background MCP connect
ci / vet-and-test (pull_request) Successful in 49s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s
Why: models need a capped way to call several tools from one script, and a slow MCP server must not block the first prompt.
Evidence: GOWORK=off go test ./runtime/codeorch/ ./runtime/tools/ ./runtime/mcp/ ./runtime/policy/ ./adapters/native/
Impact: run_tool_script stays unregistered until the caller turns it on. Deferred tools stay out of the model list. An unready MCP server connects in the background.
Refs: feature/code-orchestration
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-03 18:02:53 +00:00
weichen b6d6708416 Merge pull request 'Record tool intent before execution and add file checkpoints' (#14) from feature/durable-rewind into main
ci / vet-and-test (push) Successful in 45s
ci / docker-checkout-hostgateway (push) Successful in 1s
2026-10-03 17:40:14 +00:00
weichenandCursor fb2b3b796c Record tool intent before execution and add file checkpoints
ci / vet-and-test (pull_request) Successful in 46s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s
Why: a crash after a tool was about to run left the call unpaired, and workspace writes could not be rewound.
Evidence: GOWORK=off go test ./runtime/ ./runtime/llm/ ./runtime/checkpoint/ ./runtime/telemetry/ ./runtime/tools/ ./runtime/session/ ./adapters/native/
Impact: pending calls rerun only when both the stored and current replay declarations are safe; otherwise the result is an error. Empty fallback lists and a nil checkpoint log keep the previous behavior.
Refs: feature/durable-rewind
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-03 17:27:04 +00:00
weichen ba90a7d9ae Merge pull request 'Add a hook bus and permission decisions for tool calls' (#13) from feature/permission-hooks into main
ci / vet-and-test (push) Successful in 46s
ci / docker-checkout-hostgateway (push) Successful in 1s
2026-10-03 17:06:47 +00:00
weichenandCursor 0a52d9f976 Add a hook bus and permission decisions for tool calls
ci / vet-and-test (pull_request) Successful in 46s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s
Why: Squad members never saw the direct-run stop, reminder, and before-tool gates, and there was no shared place to decide deny, ask, or allow.
Evidence: GOWORK=off go test ./runtime/hooks/ ./runtime/policy/ ./runtime/ ./adapters/native/ passed, including an A/B run where hook adapters match the old callbacks.
Impact: An empty bus leaves stop reason, reminder count, and deny count unchanged. Callers opt in by setting GateOptions.
Refs: feature/permission-hooks
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-03 17:02:28 +00:00
weichen 3851bfb8a7 Merge pull request 'Keep tool pairs and fold reminders into the last result' (#12) from feature/context-engineering into main
ci / vet-and-test (push) Successful in 46s
ci / docker-checkout-hostgateway (push) Successful in 1s
2026-10-03 16:25:36 +00:00
weichenandCursor 6fd1d441a9 Keep tool pairs and fold reminders into the last result
ci / vet-and-test (pull_request) Successful in 45s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s
Why: Old tool results were filling the window, compaction could split a tool group or commit a stale summary, and reminders arrived as extra user messages.
Evidence: Compaction retries a conflict and opens the circuit after three failures; clearing keeps the last five tool results; a stale file write is rejected; RPC reports started, queued, or handled.
Impact: Requests over the context window replace older tool results with a placeholder. Rewriting an existing file now requires a read in the same run. Subagent text over 32000 characters is capped.
Refs: feature/context-engineering
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-03 16:19:36 +00:00
weichen 9335fd40e6 Merge pull request 'Retry model streams only before the first delta' (#11) from feature/llm-robustness into main
ci / vet-and-test (push) Successful in 47s
ci / docker-checkout-hostgateway (push) Successful in 1s
2026-10-03 15:56:14 +00:00
weichenandCursor 9f5751c205 Retry model streams only before the first delta
ci / vet-and-test (pull_request) Successful in 47s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s
Why: Buffering a whole attempt hid live output, overflow detection missed common provider errors, and a cut-off tool call could still run.
Evidence: GOWORK=off go test -race -count=1 ./runtime/llm/ ./runtime/ passed.
Impact: WithRetry forwards events immediately and retries only when nothing has been forwarded. Length or an unclean end drops unfinished tool calls. Loop usage keeps a billing sum and the last request separately.
Refs: feature/llm-robustness
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-03 15:52:08 +00:00
weichen e51f2e6d00 Merge pull request 'Drop the intent adapter and fold contract into exec' (#10) from feature/agent-stack-d1 into main
ci / vet-and-test (push) Successful in 51s
ci / docker-checkout-hostgateway (push) Successful in 1s
2026-10-03 14:02:32 +00:00
weichenandCursor d84a67cd24 Treat a cancelled SSE stream as aborted
ci / vet-and-test (pull_request) Successful in 48s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s
Why: the race detector can deliver message_end when the context is already cancelled, because select picks either ready case.
Evidence: GOWORK=off go test -count=50 -race ./runtime/llm/ passed, including TestParseSSECancel.
Impact: a cancelled parse ends with the aborted event.
Refs: D1
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-03 13:51:41 +00:00
weichenandCursor 33573d0d75 Pin agent-intent to the IntentSpec schema sample
ci / vet-and-test (pull_request) Failing after 48s
ci / docker-checkout-hostgateway (pull_request) Successful in 1s
Why: the nested checkout has to match the intent head that publishes the sample, even though native no longer imports it.
Evidence: gitlink only. The schema test already passed in the intent repository.
Impact: no native source change.
Refs: D1
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-03 13:43:32 +00:00
weichenandCursor 4a10818535 Drop the intent adapter and fold contract into exec
Why: native must compile without agent-intent, and the two execution packages were the same contract under two names.
Evidence: GOWORK=off go test -count=1 ./... passed, including schema.TestRequestAndResultRoundTrip and the native adapter tests after Options was split into embedded sub-structs.
Impact: adapters/intentgate is gone, go.mod no longer requires agent-intent, and the nested gitlink stays for the release task. Callers of Options composite literals name the sub-struct.
Refs: D1
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-03 13:31:48 +00:00
weichen d9b0f8f178 Close the native P1 gaps the audit listed for this batch
ci / vet-and-test (push) Successful in 48s
ci / docker-checkout-hostgateway (push) Successful in 0s
Merge via API
2026-10-03 13:22:54 +00:00
78 changed files with 4764 additions and 647 deletions
-81
View File
@@ -1,81 +0,0 @@
// Package intentgate 提供可选的意图闸门装饰器:在把 Prompt 交给被装饰的
// 执行后端之前,用 agent-intent 做一次确定性识别;分支为守卫/命中/预期命中时
// 交由宿主 Handler 解释,Ambiguous/Fallback 或 Handler 放弃时原样透传。
//
// 本包是 Agent-Native 仓内唯一 import agent-intent 的地方。不使用本包时,
// 本仓与 agent-intent 无任何编译关系;未 Wrap 或快照零意图时,被装饰后端的
// 行为与原来完全一致。
package intentgate
import (
"context"
"agent-intent/intent"
"agent-native/exec"
)
// Handler 解释"识别到意图"的执行请求。返回 handled=false 表示宿主放弃处理,
// 装饰器回退为透传原后端。合成结果的语义(拒绝、引导、定向执行)由宿主定义,
// 本包不发明结果语义。
type Handler func(ctx context.Context, req exec.Request, res intent.Result) (result *exec.Result, err error, handled bool)
type gated struct {
next exec.Backend
rec intent.Recognizer
snap *intent.Snapshot
handler Handler
}
// Wrap 返回带意图闸门的后端。next/rec/snap/handler 任一为 nil 时直接返回 next
// (等价于未启用,行为不变)。
func Wrap(next exec.Backend, rec intent.Recognizer, snap *intent.Snapshot, h Handler) exec.Backend {
if next == nil || rec == nil || snap == nil || h == nil {
return next
}
return &gated{next: next, rec: rec, snap: snap, handler: h}
}
// WrapRegistry 持有 Registry,每次执行读取当前快照。Reload 之后不用重建装饰器。
func WrapRegistry(next exec.Backend, rec intent.Recognizer, reg *intent.Registry, h Handler) exec.Backend {
if next == nil || rec == nil || reg == nil || h == nil {
return next
}
return &gatedRegistry{next: next, rec: rec, reg: reg, handler: h}
}
type gatedRegistry struct {
next exec.Backend
rec intent.Recognizer
reg *intent.Registry
handler Handler
}
func (g *gatedRegistry) Name() string { return g.next.Name() }
func (g *gatedRegistry) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
return executeGated(ctx, g.next, g.rec, g.reg.Snapshot(), g.handler, req)
}
func (g *gated) Name() string { return g.next.Name() }
func (g *gated) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
return executeGated(ctx, g.next, g.rec, g.snap, g.handler, req)
}
func executeGated(ctx context.Context, next exec.Backend, rec intent.Recognizer, snap *intent.Snapshot, handler Handler, req exec.Request) (*exec.Result, error) {
res, err := rec.Recognize(ctx, req.Prompt, snap)
if err != nil {
// 识别系统错误 fail-closed:不静默放行可能应拦截的内容,也不吞错。
return nil, err
}
switch res.Branch {
case intent.BranchGuardHit, intent.BranchMatched, intent.BranchExpectationHit:
default:
return next.Execute(ctx, req) // 无裁决:透传原后端
}
out, herr, handled := handler(ctx, req, res)
if !handled {
return next.Execute(ctx, req)
}
return out, herr
}
-197
View File
@@ -1,197 +0,0 @@
package intentgate
import (
"context"
"errors"
"reflect"
"testing"
"testing/fstest"
"agent-intent/intent"
"agent-native/exec"
)
type stubBackend struct {
name string
calls int
}
func (s *stubBackend) Name() string { return s.name }
func (s *stubBackend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
s.calls++
return &exec.Result{Success: true, Logs: "stub:" + req.Prompt}, nil
}
func emptySnapshot(t *testing.T) *intent.Snapshot {
t.Helper()
reg, err := intent.NewRegistryBuilder(intent.Options{}).Build()
if err != nil {
t.Fatal(err)
}
return reg.Snapshot()
}
// 带一个守卫 spec 与一个可路由 spec 的快照。
func loadedSnapshot(t *testing.T) *intent.Snapshot {
t.Helper()
b := intent.NewRegistryBuilder(intent.Options{
GuardDomains: []string{"guard"},
ProtectedDomains: []string{"guard"},
})
if err := b.Register(intent.IntentSpec{
ID: "guard.config_change", Domain: "guard", Version: 1,
Matcher: intent.MatchSpec{IncludesAny: []string{"系统配置", "网关"}},
Risk: intent.RiskHigh,
Handler: "refuse.and_guide",
Metadata: map[string]string{
"guidance": "系统配置请在管理界面修改",
},
}); err != nil {
t.Fatal(err)
}
reg, err := b.Build()
if err != nil {
t.Fatal(err)
}
return reg.Snapshot()
}
func newRec() intent.Recognizer { return intent.NewRecognizer(intent.Options{}, nil, nil) }
func TestZeroSpecSnapshotPassthroughIdentical(t *testing.T) {
// T0 行为解耦:零意图快照 + Wrap 后,与裸后端行为逐字节一致。
bare := &stubBackend{name: "native"}
gatedBe := Wrap(&stubBackend{name: "native"}, newRec(), emptySnapshot(t),
func(ctx context.Context, req exec.Request, res intent.Result) (*exec.Result, error, bool) {
t.Fatal("handler must not fire on empty snapshot")
return nil, nil, false
})
req := exec.Request{Prompt: "帮我改一下系统网关地址"}
want, werr := bare.Execute(context.Background(), req)
got, gerr := gatedBe.Execute(context.Background(), req)
if werr != nil || gerr != nil {
t.Fatalf("execute errors: %v %v", werr, gerr)
}
if !reflect.DeepEqual(want, got) {
t.Fatalf("results differ: %+v vs %+v", want, got)
}
if gatedBe.(*gated).next.(*stubBackend).calls != 1 {
t.Fatal("next backend not invoked exactly once")
}
}
func TestGuardHitInterceptsWithoutModel(t *testing.T) {
next := &stubBackend{name: "native"}
be := Wrap(next, newRec(), loadedSnapshot(t),
func(ctx context.Context, req exec.Request, res intent.Result) (*exec.Result, error, bool) {
if res.Branch != intent.BranchGuardHit || res.Top == nil || res.Top.ID != "guard.config_change" {
t.Fatalf("unexpected result: %+v", res)
}
// 宿主解释守卫意图:拒绝 + 引导(合成结果语义由宿主定义)。
return &exec.Result{Success: false, ErrorReason: "guard refused: " + res.Top.Metadata["guidance"]}, nil, true
})
out, err := be.Execute(context.Background(), exec.Request{Prompt: "帮我改一下系统网关地址"})
if err != nil {
t.Fatal(err)
}
if out.Success || out.ErrorReason == "" {
t.Fatalf("expected guard refusal, got %+v", out)
}
if next.calls != 0 {
t.Fatalf("model backend must not be called on guard hit, got %d calls", next.calls)
}
}
func TestFallbackPassesThroughToBackend(t *testing.T) {
next := &stubBackend{name: "native"}
be := Wrap(next, newRec(), loadedSnapshot(t),
func(ctx context.Context, req exec.Request, res intent.Result) (*exec.Result, error, bool) {
t.Fatal("handler must not fire on fallback")
return nil, nil, false
})
out, err := be.Execute(context.Background(), exec.Request{Prompt: "写一份周报"})
if err != nil || out == nil || !out.Success {
t.Fatalf("fallback must reach backend: %+v %v", out, err)
}
if next.calls != 1 {
t.Fatalf("backend calls = %d", next.calls)
}
}
func TestHandlerDeclineFallsBackToBackend(t *testing.T) {
next := &stubBackend{name: "native"}
be := Wrap(next, newRec(), loadedSnapshot(t),
func(ctx context.Context, req exec.Request, res intent.Result) (*exec.Result, error, bool) {
return nil, nil, false // 宿主放弃处理
})
out, err := be.Execute(context.Background(), exec.Request{Prompt: "帮我改一下系统网关地址"})
if err != nil || out == nil || !out.Success {
t.Fatalf("declined handling must pass through: %+v %v", out, err)
}
if next.calls != 1 {
t.Fatalf("backend calls = %d", next.calls)
}
}
func TestNilComponentsReturnNext(t *testing.T) {
next := &stubBackend{name: "native"}
if got := Wrap(next, nil, nil, nil); got != exec.Backend(next) {
t.Fatal("nil components must return next unchanged")
}
}
func TestRecognizeErrorFailsClosed(t *testing.T) {
next := &stubBackend{name: "native"}
b := intent.NewRegistryBuilder(intent.Options{MaxInputLen: 8})
reg, err := b.Build()
if err != nil {
t.Fatal(err)
}
be := Wrap(next, newRec(), reg.Snapshot(),
func(ctx context.Context, req exec.Request, res intent.Result) (*exec.Result, error, bool) {
t.Fatal("handler must not fire on recognize error")
return nil, nil, false
})
_, err = be.Execute(context.Background(), exec.Request{Prompt: "这一段输入明显超过八个字节"})
if err == nil {
t.Fatal("expected fail-closed error")
}
var ie *intent.Error
if !errors.As(err, &ie) {
t.Fatalf("expected *intent.Error, got %T", err)
}
if next.calls != 0 {
t.Fatal("backend must not run on recognize error")
}
}
func TestR2RegistryReloadIsVisibleWithoutRewrap(t *testing.T) {
fsys := fstest.MapFS{"intents.json": &fstest.MapFile{Data: []byte(`{"intents":[{"id":"guard.block","domain":"guard","version":1,"matcher":{"includes_any":["旧词"]}}]}`)}}
reg, err := intent.NewRegistryBuilder(intent.Options{GuardDomains: []string{"guard"}}).BuildFS(context.Background(), fsys, ".")
if err != nil {
t.Fatal(err)
}
next := &stubBackend{name: "native"}
var hits int
be := WrapRegistry(next, newRec(), reg, func(ctx context.Context, req exec.Request, res intent.Result) (*exec.Result, error, bool) {
hits++
return &exec.Result{Success: false, ErrorReason: "blocked"}, nil, true
})
if _, err := be.Execute(context.Background(), exec.Request{Prompt: "请处理新词"}); err != nil {
t.Fatal(err)
}
if hits != 0 || next.calls != 1 {
t.Fatalf("before reload hits=%d calls=%d", hits, next.calls)
}
fsys["intents.json"].Data = []byte(`{"intents":[{"id":"guard.block","domain":"guard","version":2,"matcher":{"includes_any":["新词"]}}]}`)
if _, err := reg.Reload(context.Background()); err != nil {
t.Fatal(err)
}
if _, err := be.Execute(context.Background(), exec.Request{Prompt: "请处理新词"}); err != nil {
t.Fatal(err)
}
if hits != 1 || next.calls != 1 {
t.Fatalf("after reload hits=%d calls=%d", hits, next.calls)
}
}
+109 -29
View File
@@ -14,6 +14,9 @@ import (
"agent-native/exec"
"agent-native/runtime"
"agent-native/runtime/checkpoint"
"agent-native/runtime/filestate"
"agent-native/runtime/hooks"
"agent-native/runtime/llm"
"agent-native/runtime/mcp"
"agent-native/runtime/promptx"
@@ -25,39 +28,65 @@ import (
"logging-system/logsys"
)
// Options 构造后端所需的模型接入与执行参数(由组合根注入,无产品依赖)。
type Options struct {
// ModelOptions 是模型接入。Stream 非空时不再按 APIBase/APIKey 构造客户端。
type ModelOptions struct {
APIBase string
APIKey string
Model string
MaxTurns int
Timeout time.Duration
// Stream 测试注入的流函数;缺省按 APIBase/APIKey 构造 llm.Client。
Stream llm.StreamFunc
// Search 可选:提供时注册 tavily_search / web_search(密钥由组合根注入)。
Search tools.SearchProvider
// Exec 可选:提供时注册 run_command / run_python(解释器留空项不注册)。
// nil = 不注册执行类工具(fail-closed,与 BashPermission 默认 deny 对齐)。
Exec *tools.ExecConfig
// Fetch 可选:提供时注册 fetch_url;nil = 不注册(fail-closed)。
Fetch *tools.FetchConfig
// MCP 可选:提供时把各 server 工具以 mcp__<server>__<tool> 注册进工作区。
// nil = 不注册(fail-closed:server 不可达时整个步骤失败,不静默跳过)。
MCP *mcp.Manager
// SkillsRoots 可选:全局技能根目录;执行时自动并入工作区 .agent/skills/。
SkillsRoots []string
// Compact 可选:转录超过阈值时把头部历史压成摘要(长会话上下文治理)。
Compact *CompactConfig
// Subagent 可选:提供时注册 delegate_task(委派由宿主实现,内核只开端口)。
Subagent subagent.Factory
// HostTools 是宿主产品工具(登记工件、卡片、知识库)。与内核工具同名时失败,不覆盖。
HostTools []tools.Tool
// MaxToolCalls 是单步工具调用上限。0 表示不限制。请求自带预算时以请求为准。
Stream llm.StreamFunc
}
// ToolOptions 是可选工具装配。nil 字段不注册对应工具。
type ToolOptions struct {
Search tools.SearchProvider
Exec *tools.ExecConfig
Fetch *tools.FetchConfig
MCP *mcp.Manager
SkillsRoots []string
Subagent subagent.Factory
HostTools []tools.Tool
Checkpoints *checkpoint.Log
CheckpointAnchor func() string
// CodeOrchestration 为真时注册 run_tool_script。默认关。
CodeOrchestration bool
}
// SessionOptions 是单次执行的会话预算与压缩。
type SessionOptions struct {
Compact *CompactConfig
MaxToolCalls int
// Log 可选。nil 时丢弃日志。
Log *logsys.Logger
// ExecutionID 写入日志,不参与执行控制。
}
// ObserveOptions 是日志。Log 为 nil 时丢弃。
type ObserveOptions struct {
Log *logsys.Logger
ExecutionID string
// Trace 交给运行循环。nil 不采集。
Trace func(name string, attrs map[string]string) func()
}
// 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
// Ask 在 MCP 工具要求用户输入时调用。nil 则把问题作为错误返回。
Ask func(ctx context.Context, question string) (string, error)
}
// Options 构造后端所需的模型接入与执行参数(由组合根注入,无产品依赖)。
// 字段经嵌入提升,读取仍写 opts.Model;复合字面量要写子结构名。
type Options struct {
ModelOptions
ToolOptions
SessionOptions
ObserveOptions
GateOptions
}
// CompactConfig 声明长会话压缩策略(优化性质:压缩失败不阻断执行,但留日志)。
@@ -122,6 +151,9 @@ 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 b.opts.Ask != nil {
ctx = mcp.WithAsk(ctx, b.opts.Ask)
}
model := strings.TrimSpace(b.opts.Model)
if strings.TrimSpace(req.Model) != "" {
model = strings.TrimSpace(req.Model)
@@ -164,7 +196,7 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
var reg *tools.Registry
if profile == exec.ToolProfileWorkspace {
reg = tools.WorkspaceTools(workDir)
reg = tools.WorkspaceToolsCheckpointed(workDir, filestate.New(), b.opts.Checkpoints, b.opts.CheckpointAnchor)
skills.RegisterLoader(reg, skillStore)
if b.opts.Search != nil {
tools.RegisterSearch(reg, b.opts.Search)
@@ -179,9 +211,19 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
}
tools.RegisterExecTools(reg, workDir, execCfg)
}
if b.opts.CodeOrchestration {
tools.RegisterCodeTool(reg)
}
reg.SetActivationFile(filepath.Join(workDir, ".agentx", "activated-tools.json"))
_ = reg.RestoreActivated()
if b.opts.MCP != nil {
if err := b.opts.MCP.RegisterInto(ctx, reg); err != nil {
return &exec.Result{Success: false, ErrorReason: err.Error()}, nil
switch {
case len(req.ToolAllowlist) > 0 || b.opts.MCP.Ready():
if err := b.opts.MCP.RegisterInto(ctx, reg); err != nil {
return &exec.Result{Success: false, ErrorReason: err.Error()}, nil
}
default:
b.opts.MCP.StartBackground(ctx, reg)
}
}
subagent.Register(reg, b.opts.Subagent)
@@ -325,6 +367,10 @@ 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),
Trace: b.opts.Trace,
OnEvent: func(ev runtime.Event) {
logs.WriteString(ev.Type)
if ev.Tool != "" {
@@ -415,6 +461,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) }
+64 -11
View File
@@ -2,6 +2,7 @@ package native
import (
"context"
"errors"
"os"
"path/filepath"
"strings"
@@ -39,12 +40,17 @@ func script(turns ...llm.Message) llm.StreamFunc {
func TestBackendWorkspaceProfileWritesDeliverable(t *testing.T) {
dir := t.TempDir()
b := New(Options{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":"# 交付\n\n工具写入的交付物内容。"}`,
}}},
llm.Message{Role: llm.RoleAssistant, Content: "已写入 output.md", StopReason: llm.StopEnd},
)})
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":"# 交付\n\n工具写入的交付物内容。"}`,
}}},
llm.Message{Role: llm.RoleAssistant, Content: "已写入 output.md", StopReason: llm.StopEnd},
),
},
})
var events []exec.ProcessEvent
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: dir, Prompt: "写交付物", Instructions: "你是测试角色",
@@ -81,9 +87,14 @@ func TestBackendWorkspaceProfileWritesDeliverable(t *testing.T) {
func TestBackendReasoningProfileWritesText(t *testing.T) {
dir := t.TempDir()
b := New(Options{Model: "fake", Stream: script(
llm.Message{Role: llm.RoleAssistant, Content: "纯文本交付内容(无工具)。", StopReason: llm.StopEnd},
)})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: script(
llm.Message{Role: llm.RoleAssistant, Content: "纯文本交付内容(无工具)。", StopReason: llm.StopEnd},
),
},
})
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: dir, Prompt: "直接回答", ToolProfile: exec.ToolProfileReasoning,
Targets: []string{"notes.md"},
@@ -109,14 +120,56 @@ func TestBackendFailClosedWithoutConfig(t *testing.T) {
if res.Success || !strings.Contains(res.ErrorReason, "requires") {
t.Fatalf("want fail-closed config error, got %+v", res)
}
res2, _ := New(Options{APIBase: "http://x", APIKey: "k"}).Execute(context.Background(),
res2, _ := New(Options{
ModelOptions: ModelOptions{
APIBase: "http://x",
APIKey: "k",
},
}).Execute(context.Background(),
exec.Request{WorkDir: t.TempDir(), Targets: []string{"x.md"}})
if res2.Success || !strings.Contains(res2.ErrorReason, "model") {
t.Fatalf("want model required error, got %+v", res2)
}
res3, _ := New(Options{Stream: script()}).Execute(context.Background(),
res3, _ := New(Options{
ModelOptions: ModelOptions{
Stream: script(),
},
}).Execute(context.Background(),
exec.Request{WorkDir: t.TempDir(), Targets: []string{"x.md"}, ToolProfile: "bogus"})
if res3.Success || !strings.Contains(res3.ErrorReason, "tool profile") {
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")
}
}
+9 -5
View File
@@ -30,11 +30,15 @@ func TestBackendCompactBeforeRun(t *testing.T) {
var sawSummary bool
// 注入式捕获:验证 Run 收到的消息里有摘要标记。
bk2 := New(Options{
Model: "fake",
Stream: captureFirstStream(&sawSummary),
Compact: &CompactConfig{ThresholdChars: 1000, KeepLastN: 3, Summarizer: func(ctx context.Context, messages []llm.Message) (string, error) {
return "预置历史摘要", nil
}},
ModelOptions: ModelOptions{
Model: "fake",
Stream: captureFirstStream(&sawSummary),
},
SessionOptions: SessionOptions{
Compact: &CompactConfig{ThresholdChars: 1000, KeepLastN: 3, Summarizer: func(ctx context.Context, messages []llm.Message) (string, error) {
return "预置历史摘要", nil
}},
},
})
res, err := bk2.Execute(context.Background(), exec.Request{
WorkDir: dir, Prompt: "继续", Targets: []string{"output.md"},
+96 -28
View File
@@ -22,19 +22,23 @@ func TestBackendExecToolsWiredWithOptionsAndRequestEnv(t *testing.T) {
}
dir := t.TempDir()
b := New(Options{
Model: "fake",
Stream: script(
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
ID: "c1", Name: "run_command",
Arguments: `{"command":"echo role=$ROLE_FLAG step=$STEP_FLAG > exec_out.txt"}`,
}}},
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
ID: "c2", Name: "write_file",
Arguments: `{"path":"output.md","content":"done"}`,
}}},
llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd},
),
Exec: &tools.ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin", "ROLE_FLAG=backend"}},
ModelOptions: ModelOptions{
Model: "fake",
Stream: script(
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
ID: "c1", Name: "run_command",
Arguments: `{"command":"echo role=$ROLE_FLAG step=$STEP_FLAG > exec_out.txt"}`,
}}},
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
ID: "c2", Name: "write_file",
Arguments: `{"path":"output.md","content":"done"}`,
}}},
llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd},
),
},
ToolOptions: ToolOptions{
Exec: &tools.ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin", "ROLE_FLAG=backend"}},
},
})
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: dir, Prompt: "run", Targets: []string{"output.md"},
@@ -54,7 +58,12 @@ func TestBackendExecToolsWiredWithOptionsAndRequestEnv(t *testing.T) {
func TestBackendExecToolsAbsentWithoutOptions(t *testing.T) {
stream, c := captureStream()
b := New(Options{Model: "fake", Stream: stream})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
})
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
ToolProfile: exec.ToolProfileReasoning,
@@ -64,7 +73,12 @@ func TestBackendExecToolsAbsentWithoutOptions(t *testing.T) {
}
// reasoning 档本无工具;改用 workspace 档断言 decls 不含执行工具。
stream2, c2 := captureStream()
b2 := New(Options{Model: "fake", Stream: stream2})
b2 := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream2,
},
})
_, _ = b2.Execute(context.Background(), exec.Request{
WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
ToolProfile: exec.ToolProfileWorkspace,
@@ -83,8 +97,13 @@ func TestBackendAllowlistCanDisableExecTools(t *testing.T) {
bash, _ := osexec.LookPath("bash")
stream, c := captureStream()
b := New(Options{
Model: "fake", Stream: stream,
Exec: &tools.ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin"}},
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
ToolOptions: ToolOptions{
Exec: &tools.ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin"}},
},
})
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
@@ -125,7 +144,15 @@ func TestBackendSearchToolsWiredViaOptions(t *testing.T) {
}}},
llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd},
)
b := New(Options{Model: "fake", Stream: stream, Search: provider})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
ToolOptions: ToolOptions{
Search: provider,
},
})
dir := t.TempDir()
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: dir, Prompt: "search", Targets: []string{"output.md"},
@@ -152,7 +179,12 @@ func TestBackendSearchToolsWiredViaOptions(t *testing.T) {
func TestBackendSearchToolsAbsentWithoutOptions(t *testing.T) {
stream, c := captureStream()
b := New(Options{Model: "fake", Stream: stream})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
})
_, _ = b.Execute(context.Background(), exec.Request{
WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
ToolProfile: exec.ToolProfileWorkspace,
@@ -180,6 +212,9 @@ func TestBackendMCPToolsWiredViaOptions(t *testing.T) {
t.Fatal(err)
}
defer mgr.Close()
if err := mgr.RegisterInto(context.Background(), tools.NewRegistry(t.TempDir())); err != nil {
t.Fatal(err)
}
stream, c := captureStream(
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
ID: "c1", Name: "mcp__fake__echo", Arguments: `{"text":"via-backend"}`,
@@ -190,7 +225,15 @@ func TestBackendMCPToolsWiredViaOptions(t *testing.T) {
}}},
llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd},
)
b := New(Options{Model: "fake", Stream: stream, MCP: mgr})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
ToolOptions: ToolOptions{
MCP: mgr,
},
})
dir := t.TempDir()
res, err2 := b.Execute(context.Background(), exec.Request{
WorkDir: dir, Prompt: "call mcp", Targets: []string{"output.md"},
@@ -229,7 +272,15 @@ func TestBackendSkillsIndexAndLoader(t *testing.T) {
}}},
llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd},
)
b := New(Options{Model: "fake", Stream: stream, SkillsRoots: []string{skillsRoot}})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
ToolOptions: ToolOptions{
SkillsRoots: []string{skillsRoot},
},
})
dir := t.TempDir()
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: dir, Prompt: "写", Targets: []string{"output.md"}, Skills: []string{"docx-author"},
@@ -257,7 +308,15 @@ func TestBackendSkillsIndexAndLoader(t *testing.T) {
}
func TestBackendDeclaredSkillMissingFailsClosed(t *testing.T) {
b := New(Options{Model: "fake", Stream: script(), SkillsRoots: []string{t.TempDir()}})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: script(),
},
ToolOptions: ToolOptions{
SkillsRoots: []string{t.TempDir()},
},
})
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
Skills: []string{"ghost-skill"},
@@ -277,7 +336,12 @@ func TestBackendWorkspaceSkillsRootAutoMerged(t *testing.T) {
}}},
llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd},
)
b := New(Options{Model: "fake", Stream: stream})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
})
dir := t.TempDir()
wsSkill := filepath.Join(dir, ".agent", "skills", "ws-skill")
if err := os.MkdirAll(wsSkill, 0o755); err != nil {
@@ -311,11 +375,15 @@ func TestBackendSubagentWiredViaOptions(t *testing.T) {
llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd},
)
b := New(Options{
Model: "fake",
Stream: stream,
Subagent: func(ctx context.Context, req subagent.Request) (subagent.Result, error) {
delegated = &req
return subagent.Result{Success: true, Text: "结论已查明"}, nil
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
ToolOptions: ToolOptions{
Subagent: func(ctx context.Context, req subagent.Request) (subagent.Result, error) {
delegated = &req
return subagent.Result{Success: true, Text: "结论已查明"}, nil
},
},
})
dir := t.TempDir()
+21 -9
View File
@@ -25,9 +25,13 @@ func (h hostTool) Execute(ctx context.Context, args json.RawMessage) (tools.Resu
func TestHostToolAndUsageEvent(t *testing.T) {
var events []exec.ProcessEvent
b := New(Options{
Model: "fake",
HostTools: []tools.Tool{hostTool{name: "register_artifact"}},
Stream: usageScript(),
ModelOptions: ModelOptions{
Model: "fake",
Stream: usageScript(),
},
ToolOptions: ToolOptions{
HostTools: []tools.Tool{hostTool{name: "register_artifact"}},
},
})
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: t.TempDir(), Prompt: "登记",
@@ -65,9 +69,13 @@ func TestHostToolAndUsageEvent(t *testing.T) {
func TestHostToolSeesStepWorkDir(t *testing.T) {
var seen string
b := New(Options{
Model: "fake",
Stream: script(llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ID: "c1", Name: "peek", Arguments: `{}`}}}),
HostTools: []tools.Tool{peekTool{got: &seen}},
ModelOptions: ModelOptions{
Model: "fake",
Stream: script(llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ID: "c1", Name: "peek", Arguments: `{}`}}}),
},
ToolOptions: ToolOptions{
HostTools: []tools.Tool{peekTool{got: &seen}},
},
})
dir := t.TempDir()
_, err := b.Execute(context.Background(), exec.Request{WorkDir: dir, Prompt: "x", Targets: []string{"output.md"}})
@@ -91,9 +99,13 @@ func (p peekTool) Execute(ctx context.Context, args json.RawMessage) (tools.Resu
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"}},
ModelOptions: ModelOptions{
Model: "fake",
Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "x", StopReason: llm.StopEnd}),
},
ToolOptions: ToolOptions{
HostTools: []tools.Tool{hostTool{name: "read_file"}},
},
})
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
+31 -6
View File
@@ -58,7 +58,12 @@ func toolNames(decls []llm.ToolDecl) []string {
func TestBackendRequestModelOverride(t *testing.T) {
stream, c := captureStream()
b := New(Options{Model: "default-model", Stream: stream})
b := New(Options{
ModelOptions: ModelOptions{
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",
@@ -82,7 +87,12 @@ func TestBackendRequestToolAllowlist(t *testing.T) {
}}},
llm.Message{Role: llm.RoleAssistant, Content: "gave up", StopReason: llm.StopEnd},
)
b := New(Options{Model: "fake", Stream: stream})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
})
dir := t.TempDir()
// 目标 output2.md 不在 allowlist 内能写吗?——write_file 被整体收窄,
// 模型两次尝试都会得到 unknown tool,最终无交付物 → 失败。
@@ -109,7 +119,12 @@ func TestBackendRequestToolAllowlist(t *testing.T) {
}
func TestBackendRequestUnknownAllowlistFailsClosed(t *testing.T) {
b := New(Options{Model: "fake", Stream: script()})
b := New(Options{
ModelOptions: ModelOptions{
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"},
@@ -139,7 +154,12 @@ func TestBackendRequestBudgetDeadlineBindsContext(t *testing.T) {
}()
return ch
}
b := New(Options{Model: "fake", Stream: stream})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
})
start := time.Now()
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
@@ -164,12 +184,17 @@ func TestBackendRequestBudgetDeadlineBindsContext(t *testing.T) {
func TestBackendRequestExecutionIDInLogs(t *testing.T) {
stream, _ := captureStream()
b := New(Options{Model: "fake", Stream: stream})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: stream,
},
})
dir := t.TempDir()
res, err := b.Execute(context.Background(), exec.Request{
WorkDir: dir, Prompt: "hi", Targets: []string{"output.md"},
Instructions: "你是角色",
ExecutionID: "exec-42", SessionScope: "run-9/step-3",
ExecutionID: "exec-42", SessionScope: "run-9/step-3",
ToolProfile: exec.ToolProfileReasoning,
})
if err != nil || !res.Success {
+12 -2
View File
@@ -38,7 +38,12 @@ func failThenSucceed(failMsg string) llm.StreamFunc {
func TestSessionRotationAfterPoisonedFailure(t *testing.T) {
dir := t.TempDir()
b := New(Options{Model: "fake", Stream: failThenSucceed("maximum context length exceeded (context_overflow)")})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: failThenSucceed("maximum context length exceeded (context_overflow)"),
},
})
// 第一次执行:失败(上下文溢出),留下毒化会话。
res1, err := b.Execute(context.Background(), execReq(dir, []string{"output.md"}))
@@ -87,7 +92,12 @@ func TestSessionRotationAfterPoisonedFailure(t *testing.T) {
func TestNoRotationOnRetryableNetworkFailure(t *testing.T) {
dir := t.TempDir()
b := New(Options{Model: "fake", Stream: failThenSucceed("connection reset by peer (provider_network)")})
b := New(Options{
ModelOptions: ModelOptions{
Model: "fake",
Stream: failThenSucceed("connection reset by peer (provider_network)"),
},
})
_, _ = b.Execute(context.Background(), execReq(dir, []string{"output.md"}))
transcript := filepath.Join(dir, ".agentx", "transcript.jsonl")
_ = os.MkdirAll(filepath.Dir(transcript), 0o755)
+14 -14
View File
@@ -3,28 +3,28 @@ package catalogcheck
import (
"testing"
"agent-native/contract"
"agent-native/exec"
"error-code-system/errcode"
)
// 常量取值必须与错误码库种子码相同。旧短字符串只作为错误码库里的别名存在。
func TestContractCodesMatchCatalog(t *testing.T) {
pairs := []struct {
got contract.ErrorCode
got exec.ErrorCode
want errcode.Code
}{
{contract.ErrCanceled, errcode.AgentCanceled},
{contract.ErrDeadlineExceeded, errcode.AgentDeadlineExceeded},
{contract.ErrBudgetExceeded, errcode.AgentBudgetExceeded},
{contract.ErrProviderUnavailable, errcode.AgentProviderUnavailable},
{contract.ErrProviderRejected, errcode.AgentProviderRejected},
{contract.ErrToolRejected, errcode.AgentToolRejected},
{contract.ErrPathOutsideWorkspace, errcode.AgentPathOutsideWorkspace},
{contract.ErrMissingCapability, errcode.AgentMissingCapability},
{contract.ErrMissingDeliverable, errcode.AgentMissingDeliverable},
{contract.ErrEmptyOutput, errcode.AgentEmptyOutput},
{contract.ErrUnknownOutcome, errcode.AgentUnknownOutcome},
{contract.ErrInternal, errcode.AgentInternal},
{exec.ErrCanceled, errcode.AgentCanceled},
{exec.ErrDeadlineExceeded, errcode.AgentDeadlineExceeded},
{exec.ErrBudgetExceeded, errcode.AgentBudgetExceeded},
{exec.ErrProviderUnavailable, errcode.AgentProviderUnavailable},
{exec.ErrProviderRejected, errcode.AgentProviderRejected},
{exec.ErrToolRejected, errcode.AgentToolRejected},
{exec.ErrPathOutsideWorkspace, errcode.AgentPathOutsideWorkspace},
{exec.ErrMissingCapability, errcode.AgentMissingCapability},
{exec.ErrMissingDeliverable, errcode.AgentMissingDeliverable},
{exec.ErrEmptyOutput, errcode.AgentEmptyOutput},
{exec.ErrUnknownOutcome, errcode.AgentUnknownOutcome},
{exec.ErrInternal, errcode.AgentInternal},
}
for _, pair := range pairs {
if string(pair.got) != string(pair.want) {
+23 -23
View File
@@ -11,27 +11,27 @@ import (
"testing"
"time"
"agent-native/contract"
"agent-native/exec"
)
// FakeBackend 是确定性假实现,仅用于契约套件自测与消费方单测。
// 不模拟任何模型行为;按 Script 逐项返回。
type FakeBackend struct {
Desc contract.BackendDescriptor
Desc exec.BackendDescriptor
// Result 为每次 Execute 返回的终态(浅拷贝后返回)。
Result contract.Result
Result exec.ContractResult
// BlockUntil 若非 nil,Execute 阻塞直到该 channel 关闭或 ctx 取消,
// 用于验证取消语义。
BlockUntil <-chan struct{}
// CanceledWithCtx 记录 Execute 返回时 ctx 是否已取消。
CanceledWithCtx bool
// LastRequest 记录最近一次请求。
LastRequest contract.Request
LastRequest exec.ContractRequest
}
func (f *FakeBackend) Descriptor() contract.BackendDescriptor { return f.Desc }
func (f *FakeBackend) Descriptor() exec.BackendDescriptor { return f.Desc }
func (f *FakeBackend) Execute(ctx context.Context, req contract.Request) (*contract.Result, error) {
func (f *FakeBackend) Execute(ctx context.Context, req exec.ContractRequest) (*exec.ContractResult, error) {
f.LastRequest = req
if f.BlockUntil != nil {
select {
@@ -39,10 +39,10 @@ func (f *FakeBackend) Execute(ctx context.Context, req contract.Request) (*contr
case <-ctx.Done():
f.CanceledWithCtx = true
res := *(&f.Result)
res.Outcome = contract.OutcomeCanceled
res.Err = &contract.Error{
Code: contract.ErrCanceled,
SideEffects: contract.SideEffectNone,
res.Outcome = exec.OutcomeCanceled
res.Err = &exec.Error{
Code: exec.ErrCanceled,
SideEffects: exec.SideEffectNone,
}
return &res, nil
}
@@ -55,15 +55,15 @@ func (f *FakeBackend) Execute(ctx context.Context, req contract.Request) (*contr
type Case struct {
Name string
// Backend 为被测实现。
Backend contract.Backend
Backend exec.ContractBackend
// Request 为请求模板。
Request contract.Request
Request exec.ContractRequest
// CancelAfter 若非零,在等待该时长后取消 ctx(验证取消路径)。
CancelAfter time.Duration
// WantOutcome 为期望终态。
WantOutcome contract.Outcome
WantOutcome exec.Outcome
// WantErrCode 为期望错误码(Failed/Canceled/Unknown 时)。
WantErrCode contract.ErrorCode
WantErrCode exec.ErrorCode
}
// Run 执行一条契约用例。
@@ -96,23 +96,23 @@ func Run(t *testing.T, c Case) {
t.Fatalf("[%s] Err.SideEffects must be explicit (none|possible), got empty", c.Name)
}
}
if res.Outcome == contract.OutcomeSucceeded && res.Err != nil {
if res.Outcome == exec.OutcomeSucceeded && res.Err != nil {
t.Fatalf("[%s] Succeeded result must not carry Err", c.Name)
}
}
// Suite 是最小契约套件:成功、失败带副作用标记、取消三态。
// 每个真实 Backend 实现都必须通过;套件外按实现能力补充专项用例。
func Suite(t *testing.T, mk func() contract.Backend) {
func Suite(t *testing.T, mk func() exec.ContractBackend) {
t.Helper()
t.Run("succeeded_minimal", func(t *testing.T) {
b := mk()
Run(t, Case{
Name: "succeeded",
Backend: b,
Request: contract.Request{ExecutionID: "e1", Prompt: "hi"},
WantOutcome: contract.OutcomeSucceeded,
Name: "succeeded",
Backend: b,
Request: exec.ContractRequest{ExecutionID: "e1", Prompt: "hi"},
WantOutcome: exec.OutcomeSucceeded,
})
})
@@ -130,13 +130,13 @@ func Suite(t *testing.T, mk func() contract.Backend) {
t.Run("error_sideeffects_explicit", func(t *testing.T) {
b := mk()
// 通过描述能力判定是否有失败路径可测;无失败路径的实现跳过。
if !b.Descriptor().Has(contract.CapText) && !b.Descriptor().Has(contract.CapFiles) {
if !b.Descriptor().Has(exec.CapText) && !b.Descriptor().Has(exec.CapFiles) {
t.Skip("no text/files capability to force failure path")
}
// 空提示在多数实现下是确定性失败或成功空输出;仅验证当 Outcome=Failed
// 时错误字段完整。具体实现应补充自身失败注入用例。
res, _ := b.Execute(context.Background(), contract.Request{ExecutionID: "e2"})
if res != nil && res.Outcome == contract.OutcomeFailed {
res, _ := b.Execute(context.Background(), exec.ContractRequest{ExecutionID: "e2"})
if res != nil && res.Outcome == exec.OutcomeFailed {
if res.Err == nil || res.Err.Code == "" || res.Err.SideEffects == "" {
t.Fatalf("failed result must carry Code and SideEffects, got %+v", res.Err)
}
+41
View File
@@ -0,0 +1,41 @@
package conformance
import (
"testing"
"agent-native/exec"
)
func TestFakeBackendSuite(t *testing.T) {
Suite(t, func() exec.ContractBackend {
return &FakeBackend{
Desc: exec.BackendDescriptor{
Name: "fake",
Capabilities: []exec.Capability{exec.CapText},
},
Result: exec.ContractResult{
Outcome: exec.OutcomeSucceeded,
Text: "ok",
},
}
})
}
func TestFakeCancelSemantics(t *testing.T) {
block := make(chan struct{})
b := &FakeBackend{
Desc: exec.BackendDescriptor{Name: "fake", Capabilities: []exec.Capability{exec.CapCancel}},
BlockUntil: block,
}
Run(t, Case{
Name: "cancel",
Backend: b,
Request: exec.ContractRequest{ExecutionID: "e-cancel"},
CancelAfter: 50_000_000,
WantOutcome: exec.OutcomeCanceled,
WantErrCode: exec.ErrCanceled,
})
if !b.CanceledWithCtx {
t.Fatal("backend must observe ctx cancellation before returning")
}
}
-63
View File
@@ -1,63 +0,0 @@
package contract_test
import (
"testing"
"agent-native/conformance"
"agent-native/contract"
)
// 契约自测:fake 后端通过最小套件,验证套件本身与类型语义一致。
func TestFakeBackendSuite(t *testing.T) {
conformance.Suite(t, func() contract.Backend {
return &conformance.FakeBackend{
Desc: contract.BackendDescriptor{
Name: "fake",
Capabilities: []contract.Capability{contract.CapText},
},
Result: contract.Result{
Outcome: contract.OutcomeSucceeded,
Text: "ok",
},
}
})
}
func TestFakeCancelSemantics(t *testing.T) {
block := make(chan struct{})
b := &conformance.FakeBackend{
Desc: contract.BackendDescriptor{Name: "fake", Capabilities: []contract.Capability{contract.CapCancel}},
BlockUntil: block,
}
conformance.Run(t, conformance.Case{
Name: "cancel",
Backend: b,
Request: contract.Request{ExecutionID: "e-cancel"},
CancelAfter: 50_000_000, // 50ms in ns (avoid time import here)
WantOutcome: contract.OutcomeCanceled,
WantErrCode: contract.ErrCanceled,
})
if !b.CanceledWithCtx {
t.Fatal("backend must observe ctx cancellation before returning")
}
}
func TestErrorDuckTypeForLogs(t *testing.T) {
e := &contract.Error{
Code: contract.ErrProviderUnavailable,
Reason: "down",
Retryable: true,
SideEffects: contract.SideEffectNone,
}
if e.ErrorCode() != "agent.provider_unavailable" {
t.Fatal(e.ErrorCode())
}
fields := e.ErrorFields()
if fields["reason"] != "down" || fields["side_effects"] != "none" || fields["retryable"] != true {
t.Fatalf("%v", fields)
}
var nilErr *contract.Error
if nilErr.ErrorCode() != "" || nilErr.ErrorFields() != nil {
t.Fatal("nil error")
}
}
+1 -1
View File
@@ -16,7 +16,7 @@ type Request struct {
// OnEvent 可选过程事件回调(nil = 调用方不需要过程,后端行为与现状完全一致)。
// 实现方必须保证:回调在 Execute 返回前完成(同步调用或等待收尾),无泄漏 goroutine;
// 回调内不得阻塞(消费方自行节流/异步入队)。
OnEvent func(ev ProcessEvent)
OnEvent func(ev ProcessEvent) `json:"-"`
// ---- v0.2 能力字段(全部可选;零值 = 既有行为,向后兼容) ----
+18 -23
View File
@@ -1,7 +1,8 @@
// Package contract 定义 Agent-Native 的公共执行契约。
// 原 contract 包并入 exec。与本文件同名的执行门面仍是 Request、Result、Backend;
// 设计契约里冲突的四个名字加 Contract 前缀。
//
// 本包是助手(Core)与小队(Squad)共用的唯一执行接口层:Core 注入具体
// Backend 实现,Squad 只表达需求并使用已注入的后端。契约不依赖任何具体
// ContractBackend 实现,Squad 只表达需求并使用已注入的后端。契约不依赖任何具体
// 内核、产品类型、数据库或配置包。
//
// 语义不变量(实现方与消费方都必须遵守):
@@ -20,7 +21,7 @@
// 上下文,不授予跨会话可见性。
// 7. Workspace 的路径约束不是 OS 沙箱;执行不可信第三方工具需要调用方
// 提供的进程/容器隔离,不能由本契约推导。
package contract
package exec
import (
"context"
@@ -50,12 +51,6 @@ const (
CapStructuredOutput Capability = "structured-output"
)
// ToolProfile 是预定义工具策略名;与现有产品配置兼容。
const (
ToolProfileWorkspace = "workspace" // 默认:文档类工具受限集合
ToolProfileReasoning = "reasoning" // 无工具,仅推理
)
// ToolPolicy 描述执行允许的工具范围。实现不得自行扩展;未知工具一律拒绝。
type ToolPolicy struct {
// Profile 为策略名(workspace/reasoning/或产品注册的自定义名)。
@@ -82,8 +77,8 @@ type SessionScope struct {
Session string `json:"session,omitempty"`
}
// Budget 约束一次执行(含共享 ExecutionID 的重试)。
type Budget struct {
// ContractBudget 约束一次执行(含共享 ExecutionID 的重试)。
type ContractBudget struct {
// MaxPromptTokens/MaxCompletionTokens 为 0 时不限制。
MaxPromptTokens int `json:"max_prompt_tokens,omitempty"`
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
@@ -98,9 +93,9 @@ type Message struct {
Content string `json:"content"`
}
// Request 是一次授权执行的统一输入。字段为设计契约;现有 file-only 流程
// ContractRequest 是一次授权执行的统一输入。字段为设计契约;现有 file-only 流程
// (Prompt+Instructions+Targets→reply.md)是其子集 facade。
type Request struct {
type ContractRequest struct {
// ExecutionID 由调用方生成,重试与受控 fallback 复用同值。
ExecutionID string `json:"execution_id"`
// SessionScope 为隔离身份;可为零值(一次性任务)。
@@ -118,8 +113,8 @@ type Request struct {
Workspace WorkspaceHandle `json:"workspace"`
// ToolPolicy 为工具范围;空 Profile 等同 reasoning。
ToolPolicy ToolPolicy `json:"tool_policy"`
// Budget 为本次执行预算。
Budget Budget `json:"budget"`
// ContractBudget 为本次执行预算。
ContractBudget ContractBudget `json:"budget"`
// Required 为必须满足的能力;解析器据此选择实现,缺失即拒绝。
Required []Capability `json:"required,omitempty"`
@@ -235,14 +230,14 @@ type Usage struct {
CompletionTokens int `json:"completion_tokens,omitempty"`
}
// Result 是执行终态。Text 与 Files 依据任务可为空;Outcome 决定解释。
type Result struct {
// ContractResult 是执行终态。Text 与 Files 依据任务可为空;Outcome 决定解释。
type ContractResult struct {
Outcome Outcome `json:"outcome"`
// Text 为文本结果(CapText 任务必填)。
Text string `json:"text,omitempty"`
// Files 为文件候选产物(内容字节),由调用方决定正式落盘与提交。
Files map[string][]byte `json:"files,omitempty"`
// Logs 为执行日志/过程摘要(兼容现有 Result.Logs)。
// Logs 为执行日志/过程摘要(兼容现有 ContractResult.Logs)。
Logs string `json:"logs,omitempty"`
// Err 为 Failed/Canceled/Unknown 时的标准化错误;Succeeded 时为 nil。
Err *Error `json:"err,omitempty"`
@@ -271,19 +266,19 @@ func (d BackendDescriptor) Has(c Capability) bool {
return false
}
// Backend 是统一执行接口。实现必须:
// ContractBackend 是统一执行接口。实现必须:
// - 在 ctx 取消时尽快返回(Canceled 或 Unknown);
// - 不持有跨执行的可变会话状态(连接池等无状态资源可复用);
// - 不自行扩大工具范围或读取 Workspace 之外的路径。
type Backend interface {
type ContractBackend interface {
Descriptor() BackendDescriptor
// Execute 执行一次授权调用。返回的 Result 永不为 nil;ctx 错误映射为
// Execute 执行一次授权调用。返回的 ContractResult 永不为 nil;ctx 错误映射为
// Outcome=Canceled/Unknown,不返回 Go error 表达业务终态。
Execute(ctx context.Context, req Request) (*Result, error)
Execute(ctx context.Context, req ContractRequest) (*ContractResult, error)
}
// Resolver 按名称与必需能力解析后端。未知名必须返回错误,生产路径不得
// 回落到 Stub。
type Resolver interface {
Resolve(name string, required []Capability) (Backend, error)
Resolve(name string, required []Capability) (ContractBackend, error)
}
+23
View File
@@ -0,0 +1,23 @@
package exec
import "testing"
func TestErrorDuckTypeForLogs(t *testing.T) {
e := &Error{
Code: ErrProviderUnavailable,
Reason: "down",
Retryable: true,
SideEffects: SideEffectNone,
}
if e.ErrorCode() != "agent.provider_unavailable" {
t.Fatal(e.ErrorCode())
}
fields := e.ErrorFields()
if fields["reason"] != "down" || fields["side_effects"] != "none" || fields["retryable"] != true {
t.Fatalf("%v", fields)
}
var nilErr *Error
if nilErr.ErrorCode() != "" || nilErr.ErrorFields() != nil {
t.Fatal("nil error")
}
}
-5
View File
@@ -3,15 +3,10 @@ module agent-native
go 1.27
require (
agent-intent v0.0.0
error-code-system v0.0.0
logging-system v0.0.0
)
require golang.org/x/text v0.29.0 // indirect
replace agent-intent => ./deps/agent-intent
replace error-code-system => ./deps/error-code-system
replace logging-system => ./deps/logging-system
-2
View File
@@ -1,2 +0,0 @@
golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk=
golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4=
+24 -4
View File
@@ -31,6 +31,8 @@ type Response struct {
Text string `json:"text,omitempty"`
Error string `json:"error,omitempty"`
Result any `json:"result,omitempty"`
// Status 是命令回执:started 已接手,queued 先排队,handled 已处理完。
Status string `json:"status,omitempty"`
}
// Server serves RPC commands against one agent session and its transcript.
@@ -60,20 +62,31 @@ func (s *Server) dispatch(ctx context.Context, req Request) Response {
if s.Agent == nil {
return fail("no agent")
}
res, err := s.Agent.Prompt(ctx, llm.Message{Role: llm.RoleUser, Content: req.Message})
msg := llm.Message{Role: llm.RoleUser, Content: req.Message}
if s.Agent.Busy() {
s.Agent.FollowUp(msg)
return Response{OK: true, Status: "queued"}
}
res, err := s.Agent.Prompt(ctx, msg)
if err != nil {
return fail(err.Error())
}
if res.Err != nil {
return fail(res.Err.Error())
}
return Response{OK: true, Text: res.Text}
return Response{OK: true, Text: res.Text, Status: "handled"}
case "steer":
if s.Agent == nil {
return fail("no agent")
}
s.Agent.Steer(llm.Message{Role: llm.RoleUser, Content: req.Message})
return Response{OK: true}
return Response{OK: true, Status: queueStatus(s.Agent.Busy())}
case "follow_up":
if s.Agent == nil {
return fail("no agent")
}
s.Agent.FollowUp(llm.Message{Role: llm.RoleUser, Content: req.Message})
return Response{OK: true}
return Response{OK: true, Status: queueStatus(s.Agent.Busy())}
case "abort":
s.Agent.Abort()
return Response{OK: true}
@@ -184,6 +197,13 @@ func (s *Server) dispatch(ctx context.Context, req Request) Response {
func fail(msg string) Response { return Response{OK: false, Error: msg} }
func queueStatus(busy bool) string {
if busy {
return "queued"
}
return "started"
}
func encode(resp Response) ([]byte, error) {
raw, err := json.Marshal(resp)
if err != nil {
+5 -1
View File
@@ -59,9 +59,13 @@ func TestRPCSessionCommandsAndClient(t *testing.T) {
t.Fatalf("fork %+v %v", resp, err)
}
resp, err = client.Call(context.Background(), Request{Type: "prompt", Message: "hi"})
if err != nil || !resp.OK || resp.Text != "echo" {
if err != nil || !resp.OK || resp.Text != "echo" || resp.Status != "handled" {
t.Fatalf("prompt %+v %v", resp, err)
}
resp, err = client.Call(context.Background(), Request{Type: "steer", Message: "later"})
if err != nil || !resp.OK || resp.Status != "started" {
t.Fatalf("steer %+v %v", resp, err)
}
htmlPath := filepath.Join(dir, "out.html")
resp, err = client.Call(context.Background(), Request{Type: "export_html", Path: htmlPath})
if err != nil || !resp.OK {
+95
View File
@@ -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
}
+242
View File
@@ -0,0 +1,242 @@
// Package checkpoint 记录写文件之前的内容,并按用户消息把文件退回去。
// 只追踪显式交给它的写入。shell 与 python 的改动不经过这里。
package checkpoint
import (
"bytes"
"encoding/json"
"os"
"path/filepath"
"sync"
)
// Record 是一次被追踪的写入。Before 是写入前的字节,After 是写入后的字节。
type Record struct {
Seq int `json:"seq"`
Anchor string `json:"anchor,omitempty"`
Root string `json:"root"`
Path string `json:"path"`
Existed bool `json:"existed"`
Before []byte `json:"before,omitempty"`
After []byte `json:"after,omitempty"`
}
// Result 是一次回滚的分类。Skipped 是当前内容既不是写入后也不是写入前的文件。
type Result struct {
Restored []string
Skipped []string
}
// Log 是追加写的检查点文件。
type Log struct {
path string
mu sync.Mutex
}
func Open(path string) (*Log, error) {
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
return nil, err
}
f, err := os.OpenFile(path, os.O_CREATE, 0o644)
if err != nil {
return nil, err
}
_ = f.Close()
return &Log{path: path}, nil
}
func (l *Log) Path() string { return l.path }
// Record 追加一条检查点。rel 是相对 root 的路径。
func (l *Log) Record(anchor, root, rel string, before []byte, existed bool, after []byte) error {
if l == nil {
return nil
}
l.mu.Lock()
defer l.mu.Unlock()
all, err := l.read()
if err != nil {
return err
}
rec := Record{
Seq: len(all) + 1, Anchor: anchor, Root: root, Path: filepath.ToSlash(rel),
Existed: existed, Before: append([]byte(nil), before...), After: append([]byte(nil), after...),
}
raw, err := json.Marshal(rec)
if err != nil {
return err
}
f, err := os.OpenFile(l.path, os.O_APPEND|os.O_WRONLY, 0o644)
if err != nil {
return err
}
defer f.Close()
_, err = f.Write(append(raw, '\n'))
return err
}
// Plan 列出从 anchor 起(含)会被回滚的文件,不改磁盘。
// ordered 是用户消息 id 的时间顺序。
func (l *Log) Plan(anchor string, ordered []string) (Result, error) {
if l == nil {
return Result{}, nil
}
l.mu.Lock()
defer l.mu.Unlock()
all, err := l.read()
if err != nil {
return Result{}, err
}
return classify(all, anchor, ordered), nil
}
// Rewind 把 anchor 及其后的被追踪写入退回写入前。
// 当前内容与最后一次记录的 After 不一致时跳过并列入 Skipped。
func (l *Log) Rewind(anchor string, ordered []string) (Result, error) {
if l == nil {
return Result{}, nil
}
l.mu.Lock()
defer l.mu.Unlock()
all, err := l.read()
if err != nil {
return Result{}, err
}
plan := classify(all, anchor, ordered)
for _, rel := range plan.Restored {
group := groupFor(selected(all, anchor, ordered), rel)
if err := restore(group); err != nil {
return Result{}, err
}
}
return plan, nil
}
func (l *Log) read() ([]Record, error) {
b, err := os.ReadFile(l.path)
if err != nil {
return nil, err
}
var out []Record
for _, line := range bytes.Split(b, []byte("\n")) {
line = bytes.TrimSpace(line)
if len(line) == 0 {
continue
}
var rec Record
if err := json.Unmarshal(line, &rec); err != nil {
return nil, err
}
out = append(out, rec)
}
return out, nil
}
func selected(all []Record, anchor string, ordered []string) []Record {
keep := map[string]bool{}
on := false
for _, id := range ordered {
if id == anchor {
on = true
}
if on {
keep[id] = true
}
}
if anchor != "" && len(ordered) == 0 {
keep[anchor] = true
}
var out []Record
for _, rec := range all {
if keep[rec.Anchor] {
out = append(out, rec)
}
}
return out
}
func classify(all []Record, anchor string, ordered []string) Result {
sel := selected(all, anchor, ordered)
seen := map[string]bool{}
var restored, skipped []string
for _, rec := range sel {
if seen[rec.Path] {
continue
}
seen[rec.Path] = true
group := groupFor(sel, rec.Path)
switch fileState(group) {
case stateRestore:
restored = append(restored, rec.Path)
case stateSkip:
skipped = append(skipped, rec.Path)
}
}
return Result{Restored: restored, Skipped: skipped}
}
func groupFor(sel []Record, rel string) []Record {
var group []Record
for _, rec := range sel {
if rec.Path == rel {
group = append(group, rec)
}
}
return group
}
const (
stateAlready = iota
stateRestore
stateSkip
)
func fileState(group []Record) int {
if len(group) == 0 {
return stateAlready
}
first, last := group[0], group[len(group)-1]
abs := filepath.Join(last.Root, filepath.FromSlash(last.Path))
current, err := os.ReadFile(abs)
missing := os.IsNotExist(err)
if err != nil && !missing {
return stateSkip
}
if missing {
if !last.Existed && len(last.After) == 0 {
return stateAlready
}
if !first.Existed && bytes.Equal(nil, first.Before) && len(last.After) > 0 {
return stateSkip
}
return stateSkip
}
if bytes.Equal(current, last.After) {
if !first.Existed {
return stateRestore
}
if bytes.Equal(current, first.Before) {
return stateAlready
}
return stateRestore
}
if first.Existed && bytes.Equal(current, first.Before) {
return stateAlready
}
return stateSkip
}
func restore(group []Record) error {
if len(group) == 0 {
return nil
}
first := group[0]
abs := filepath.Join(first.Root, filepath.FromSlash(first.Path))
if !first.Existed {
return os.Remove(abs)
}
if err := os.MkdirAll(filepath.Dir(abs), 0o755); err != nil {
return err
}
return os.WriteFile(abs, first.Before, 0o644)
}
+80
View File
@@ -0,0 +1,80 @@
package checkpoint
import (
"os"
"path/filepath"
"testing"
)
func TestRewindRestoresTrackedWriteAndSkipsExternalChange(t *testing.T) {
root := t.TempDir()
log, err := Open(filepath.Join(t.TempDir(), "cp.jsonl"))
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(root, "keep.txt"), []byte("old"), 0o644); err != nil {
t.Fatal(err)
}
if err := log.Record("m1", root, "keep.txt", []byte("old"), true, []byte("new")); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(root, "keep.txt"), []byte("new"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(root, "side.txt"), []byte("from-shell"), 0o644); err != nil {
t.Fatal(err)
}
if err := log.Record("m1", root, "moved.txt", []byte("a"), true, []byte("b")); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(root, "moved.txt"), []byte("someone-else"), 0o644); err != nil {
t.Fatal(err)
}
got, err := log.Rewind("m1", []string{"m1"})
if err != nil {
t.Fatal(err)
}
if len(got.Restored) != 1 || got.Restored[0] != "keep.txt" {
t.Fatalf("restored %#v", got.Restored)
}
if len(got.Skipped) != 1 || got.Skipped[0] != "moved.txt" {
t.Fatalf("skipped %#v", got.Skipped)
}
body, err := os.ReadFile(filepath.Join(root, "keep.txt"))
if err != nil || string(body) != "old" {
t.Fatalf("keep.txt %q %v", body, err)
}
side, err := os.ReadFile(filepath.Join(root, "side.txt"))
if err != nil || string(side) != "from-shell" {
t.Fatalf("untracked file changed: %q %v", side, err)
}
moved, _ := os.ReadFile(filepath.Join(root, "moved.txt"))
if string(moved) != "someone-else" {
t.Fatalf("external edit was reverted: %q", moved)
}
}
func TestRewindRemovesCreatedFile(t *testing.T) {
root := t.TempDir()
log, err := Open(filepath.Join(t.TempDir(), "cp.jsonl"))
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(root, "new.txt"), []byte("hi"), 0o644); err != nil {
t.Fatal(err)
}
if err := log.Record("m2", root, "new.txt", nil, false, []byte("hi")); err != nil {
t.Fatal(err)
}
got, err := log.Rewind("m2", []string{"m1", "m2"})
if err != nil {
t.Fatal(err)
}
if len(got.Restored) != 1 {
t.Fatalf("restored %#v", got.Restored)
}
if _, err := os.Stat(filepath.Join(root, "new.txt")); !os.IsNotExist(err) {
t.Fatalf("created file still present: %v", err)
}
}
+51
View File
@@ -0,0 +1,51 @@
"""工作流脚本回调。用法:import agentwf as wf。
请求写到 fd 3,响应从 fd 4 按行读回。宿主一次只处理一条。
"""
import json
import os
_seq = 0
_out = os.fdopen(3, "w", encoding="utf-8", buffering=1)
_inp = os.fdopen(4, "r", encoding="utf-8", buffering=1)
def _rpc(op, **fields):
global _seq
_seq += 1
msg = {"id": str(_seq), "op": op}
msg.update(fields)
_out.write(json.dumps(msg, ensure_ascii=False) + "\n")
_out.flush()
line = _inp.readline()
if not line:
raise RuntimeError("workflow host closed")
resp = json.loads(line)
if not resp.get("ok"):
raise RuntimeError(resp.get("error") or "workflow call failed")
return resp.get("output", "")
def tool(name, args=None):
return _rpc("tool.call", name=name, args=args or {})
def ask(question, choices=None):
return _rpc("hitl.ask", question=question, choices=choices or [])
def agent(name, task):
return _rpc("agent.spawn", agent=name, task=task)
def artifact(name, content):
return _rpc("artifact.register", name=name, content=content)
def card(title, kind="result"):
return _rpc("card.create", title=title, kind=kind)
def step(step_id, title, status):
return _rpc("step", stepId=step_id, title=title, status=status)
+12
View File
@@ -0,0 +1,12 @@
package codeorch
// Forbidden 是脚本不能再调用的工具:解释器、自身、派发和工作流。
func Forbidden(name string) bool {
switch name {
case "run_python", "run_shell", "run_command", "run_tool_script",
"delegate_task", "squad_orchestrate", "run_workflow", "review_findings":
return true
default:
return false
}
}
+380
View File
@@ -0,0 +1,380 @@
// Package codeorch 在会话目录里跑一段 Python,并按行把工具调用交回宿主。
// 工作流脚本和 run_tool_script 共用这一套进程与限额。
package codeorch
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"os"
"os/exec"
"path/filepath"
"strings"
"sync"
"time"
"unicode/utf8"
_ "embed"
)
//go:embed agentwf.py
var agentwfSource string
const (
defaultTimeout = 120 * time.Second
defaultMaxWait = 600 * time.Second
defaultCalls = 50
defaultResult = 1 << 20
defaultModel = 16000
)
var (
helperOnce sync.Once
helperDir string
helperErr error
)
// Call 是脚本发来的一行请求。Args 与 Choices 保持原始 JSON。
type Call struct {
ID string `json:"id"`
Op string `json:"op"`
Name string `json:"name"`
Args json.RawMessage `json:"args"`
Question string `json:"question"`
Choices json.RawMessage `json:"choices"`
Agent string `json:"agent"`
Task string `json:"task"`
Content string `json:"content"`
Title string `json:"title"`
Kind string `json:"kind"`
Status string `json:"status"`
StepID string `json:"stepId"`
}
// Limits 控制一次脚本运行。零值使用工具路径的缺省:120 秒、最长 600 秒、50 次、1 MiB、16k 字符。
type Limits struct {
Timeout time.Duration
MaxTimeout time.Duration
MaxCalls int
MaxResultBytes int
MaxModelChars int
FullDir string
}
// Outcome 是给模型的文本、已成功的调用,以及超限时的全文路径。
type Outcome struct {
Output string
FullPath string
Calls []string
Err error
}
// Starter 按工作目录和脚本路径准备命令。空则用 python3。
type Starter func(ctx context.Context, workDir, path string) (*exec.Cmd, error)
// Dispatch 处理一条回调。返回的错误会写回脚本并记入最终结果。
type Dispatch func(ctx context.Context, call Call) (string, error)
// ModelText 把截断后的输出、已执行调用和错误拼成给模型的文本。
func (o Outcome) ModelText() string {
text := o.Output
if o.Err != nil {
msg := o.Err.Error()
if len(o.Calls) > 0 {
msg = "已执行:" + strings.Join(o.Calls, "、") + "\n" + msg
}
if text != "" {
text += "\n"
}
text += msg
}
if o.FullPath != "" {
if text != "" {
text += "\n"
}
text += "全文在 " + o.FullPath
}
return text
}
// RunFile 执行已经落盘的脚本。
func RunFile(ctx context.Context, workDir, path string, lim Limits, start Starter, dispatch Dispatch) Outcome {
if ctx == nil {
ctx = context.Background()
}
timeout := lim.bounded()
runCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
helper, err := helperPath()
if err != nil {
return Outcome{Err: err}
}
if start == nil {
start = defaultStart
}
cmd, err := start(runCtx, workDir, path)
if err != nil {
return Outcome{Err: err}
}
cmd.Env = prependPythonPath(cmd.Env, helper)
reqR, reqW, err := os.Pipe()
if err != nil {
return Outcome{Err: err}
}
respR, respW, err := os.Pipe()
if err != nil {
reqR.Close()
reqW.Close()
return Outcome{Err: err}
}
cmd.ExtraFiles = []*os.File{reqW, respR}
var stdout, stderr bytes.Buffer
cmd.Stdout = &stdout
cmd.Stderr = &stderr
if err := cmd.Start(); err != nil {
reqR.Close()
reqW.Close()
respR.Close()
respW.Close()
return Outcome{Err: err}
}
reqW.Close()
respR.Close()
defer reqR.Close()
defer respW.Close()
calls, callErr := serve(runCtx, reqR, respW, lim, dispatch)
if callErr != nil {
cancel()
}
waitErr := cmd.Wait()
out := Outcome{Output: clipJoined(&stdout, &stderr), Calls: calls}
out = clipModel(out, lim)
switch {
case callErr != nil:
out.Err = callErr
case runCtx.Err() == context.DeadlineExceeded:
out.Err = fmt.Errorf("脚本超时")
case runCtx.Err() == context.Canceled:
out.Err = fmt.Errorf("脚本已取消")
case waitErr != nil:
out.Err = fmt.Errorf("脚本退出:%v", waitErr)
}
return out
}
// Run 把脚本体写入工作目录后再执行。
func Run(ctx context.Context, workDir, body string, lim Limits, start Starter, dispatch Dispatch) Outcome {
dir := filepath.Join(workDir, ".agent")
if err := os.MkdirAll(dir, 0o755); err != nil {
return Outcome{Err: err}
}
path := filepath.Join(dir, fmt.Sprintf("orch-%d.py", time.Now().UnixNano()))
if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
return Outcome{Err: err}
}
return RunFile(ctx, workDir, path, lim, start, dispatch)
}
func (l Limits) bounded() time.Duration {
d := l.Timeout
if d <= 0 {
d = defaultTimeout
}
max := l.MaxTimeout
if max <= 0 {
max = defaultMaxWait
}
if d > max {
return max
}
return d
}
func (l Limits) calls() int {
if l.MaxCalls <= 0 {
return defaultCalls
}
return l.MaxCalls
}
func (l Limits) resultBytes() int {
if l.MaxResultBytes <= 0 {
return defaultResult
}
return l.MaxResultBytes
}
func (l Limits) modelChars() int {
if l.MaxModelChars <= 0 {
return defaultModel
}
return l.MaxModelChars
}
func serve(ctx context.Context, reqR, respW *os.File, lim Limits, dispatch Dispatch) ([]string, error) {
sc := bufio.NewScanner(reqR)
sc.Buffer(make([]byte, 64*1024), lim.resultBytes())
var calls []string
n := 0
for sc.Scan() {
line := strings.TrimSpace(sc.Text())
if line == "" {
continue
}
n++
var call Call
if err := json.Unmarshal([]byte(line), &call); err != nil {
_ = writeResp(respW, call.ID, false, "", "请求不是合法 JSON")
continue
}
if n > lim.calls() {
_ = writeResp(respW, call.ID, false, "", "脚本回调超过 50 次")
return calls, fmt.Errorf("脚本回调超过 50 次")
}
if ctx.Err() != nil {
return calls, ctx.Err()
}
if dispatch == nil {
_ = writeResp(respW, call.ID, false, "", "没有回调处理")
continue
}
output, err := dispatch(ctx, call)
if err != nil {
_ = writeResp(respW, call.ID, false, "", err.Error())
if ctx.Err() != nil {
return calls, ctx.Err()
}
continue
}
if call.Op == "tool.call" && call.Name != "" {
calls = append(calls, call.Name)
}
if err := writeResp(respW, call.ID, true, clipBytes(output, lim.resultBytes()), ""); err != nil {
return calls, err
}
}
if err := sc.Err(); err != nil {
if err == bufio.ErrTooLong {
return calls, fmt.Errorf("单行请求超过 1MB")
}
if ctx.Err() != nil {
return calls, ctx.Err()
}
return calls, err
}
return calls, nil
}
func writeResp(w io.Writer, id string, ok bool, output, errText string) error {
raw, err := json.Marshal(struct {
ID string `json:"id"`
OK bool `json:"ok"`
Output string `json:"output,omitempty"`
Error string `json:"error,omitempty"`
}{ID: id, OK: ok, Output: output, Error: errText})
if err != nil {
return err
}
_, err = w.Write(append(raw, '\n'))
return err
}
func clipBytes(s string, n int) string {
if n <= 0 || len(s) <= n {
return s
}
cut := s[:n]
for !utf8.ValidString(cut) && len(cut) > 0 {
cut = cut[:len(cut)-1]
}
return cut
}
func clipJoined(stdout, stderr *bytes.Buffer) string {
text := strings.TrimSpace(stdout.String())
errText := strings.TrimSpace(stderr.String())
if errText != "" {
if text != "" {
text += "\n"
}
text += errText
}
return text
}
func clipModel(out Outcome, lim Limits) Outcome {
limit := lim.modelChars()
if utf8.RuneCountInString(out.Output) <= limit {
return out
}
dir := lim.FullDir
if dir == "" {
dir = os.TempDir()
}
if err := os.MkdirAll(dir, 0o755); err != nil {
out.Output = clipRunes(out.Output, limit)
return out
}
path := filepath.Join(dir, fmt.Sprintf("script-output-%d.txt", time.Now().UnixNano()))
if err := os.WriteFile(path, []byte(out.Output), 0o644); err != nil {
out.Output = clipRunes(out.Output, limit)
return out
}
out.FullPath = path
out.Output = clipRunes(out.Output, limit)
return out
}
func clipRunes(s string, n int) string {
if utf8.RuneCountInString(s) <= n {
return s
}
i := 0
for count := 0; count < n && i < len(s); count++ {
_, size := utf8.DecodeRuneInString(s[i:])
i += size
}
return s[:i] + "…(超长截断)"
}
func defaultStart(ctx context.Context, workDir, path string) (*exec.Cmd, error) {
bin, err := exec.LookPath("python3")
if err != nil {
return nil, fmt.Errorf("找不到 python3: %w", err)
}
cmd := exec.CommandContext(ctx, bin, path)
cmd.Dir = workDir
return cmd, nil
}
func helperPath() (string, error) {
helperOnce.Do(func() {
dir, err := os.MkdirTemp("", "agentwf-*")
if err != nil {
helperErr = err
return
}
if err := os.WriteFile(filepath.Join(dir, "agentwf.py"), []byte(agentwfSource), 0o644); err != nil {
helperErr = err
return
}
helperDir = dir
})
return helperDir, helperErr
}
func prependPythonPath(env []string, dir string) []string {
if env == nil {
env = os.Environ()
}
for i, kv := range env {
if strings.HasPrefix(kv, "PYTHONPATH=") {
env[i] = "PYTHONPATH=" + dir + string(os.PathListSeparator) + strings.TrimPrefix(kv, "PYTHONPATH=")
return env
}
}
return append(env, "PYTHONPATH="+dir)
}
+89
View File
@@ -0,0 +1,89 @@
package codeorch
import (
"context"
"errors"
"os"
"strings"
"testing"
"time"
"unicode/utf8"
)
func TestScriptCallsTwoToolsAndReportsTheOneThatRan(t *testing.T) {
body := "import agentwf as wf\nprint(wf.tool('alpha', {}))\nprint(wf.tool('beta', {}))\n"
out := Run(context.Background(), t.TempDir(), body, Limits{}, nil, func(_ context.Context, call Call) (string, error) {
if call.Name == "beta" {
return "", errors.New("beta 失败")
}
return "A", nil
})
text := out.ModelText()
if len(out.Calls) != 1 || out.Calls[0] != "alpha" {
t.Fatalf("calls %+v", out.Calls)
}
if !strings.Contains(text, "已执行:alpha") || !strings.Contains(text, "beta 失败") {
t.Fatalf("text %q err %v", text, out.Err)
}
}
func TestForbiddenToolIsRejectedBeforeDispatchRunsIt(t *testing.T) {
body := "import agentwf as wf\nprint(wf.tool('read_file', {}))\nprint(wf.tool('run_shell', {}))\n"
var ran []string
out := Run(context.Background(), t.TempDir(), body, Limits{}, nil, func(_ context.Context, call Call) (string, error) {
if Forbidden(call.Name) {
return "", errors.New("不能调用 " + call.Name)
}
ran = append(ran, call.Name)
return "ok", nil
})
if len(ran) != 1 || ran[0] != "read_file" {
t.Fatalf("ran %+v", ran)
}
if !strings.Contains(out.ModelText(), "不能调用 run_shell") {
t.Fatalf("text %q", out.ModelText())
}
}
func TestLongOutputIsClippedAndStored(t *testing.T) {
body := "import agentwf as wf\nprint(wf.tool('big', {}))\n"
dir := t.TempDir()
out := Run(context.Background(), dir, body, Limits{MaxModelChars: 8, FullDir: dir}, nil, func(_ context.Context, call Call) (string, error) {
return strings.Repeat("测", 20), nil
})
if out.FullPath == "" {
t.Fatal("expected a full-text file")
}
full, err := os.ReadFile(out.FullPath)
if err != nil {
t.Fatal(err)
}
if utf8.RuneCount(full) < 20 {
t.Fatalf("full text too short: %q", full)
}
if !strings.HasSuffix(out.Output, "…(超长截断)") || utf8.RuneCountInString(strings.TrimSuffix(out.Output, "…(超长截断)")) != 8 {
t.Fatalf("model output not clipped: %q", out.Output)
}
if !strings.Contains(out.ModelText(), out.FullPath) {
t.Fatalf("model text missing path: %q", out.ModelText())
}
}
func TestResultGivenToScriptIsCapped(t *testing.T) {
got := clipBytes(strings.Repeat("x", 20), 8)
if got != "xxxxxxxx" {
t.Fatalf("clip %q", got)
}
if !utf8.ValidString(clipBytes("测测测", 4)) {
t.Fatal("clip broke utf-8")
}
}
func TestTimeoutIsCapped(t *testing.T) {
if (Limits{Timeout: 900 * time.Second, MaxTimeout: 600 * time.Second}).bounded() != 600*time.Second {
t.Fatal("timeout must clamp to the max")
}
if (Limits{}).bounded() != defaultTimeout {
t.Fatal("empty limits must use the 120s default")
}
}
+210
View File
@@ -0,0 +1,210 @@
package runtime
import (
"context"
"path/filepath"
"strings"
"testing"
"agent-native/runtime/llm"
"agent-native/runtime/session"
"agent-native/runtime/tools"
)
type replayTool struct {
countingTool
replay string
}
func (r *replayTool) Replay() string {
if r.replay == "" {
return tools.ReplayNever
}
return r.replay
}
func TestPendingUnsafeToolIsNotRerun(t *testing.T) {
dir := t.TempDir()
st, err := session.Open(filepath.Join(dir, "s.jsonl"))
if err != nil {
t.Fatal(err)
}
if err := st.Append(llm.Message{
Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
ToolCalls: []llm.ToolCall{{ID: "c1", Name: "write_file", Arguments: `{}`}},
}); err != nil {
t.Fatal(err)
}
if err := st.AppendToolIntent(session.ToolIntent{CallID: "c1", Name: "write_file", Arguments: `{}`, Replay: tools.ReplayNever}); err != nil {
t.Fatal(err)
}
tool := &countingTool{name: "write_file"}
reg := tools.NewRegistry(dir)
reg.Register(tool)
_, err = Run(context.Background(), Config{
Model: "fake", Session: st, Tools: reg, MaxTurns: 2,
Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "later", StopReason: llm.StopEnd}),
}, llm.Message{Role: llm.RoleUser, Content: "next"})
if err != nil {
t.Fatal(err)
}
if tool.calls != 0 {
t.Fatalf("unsafe tool reran: %d", tool.calls)
}
msgs, err := st.Load()
if err != nil {
t.Fatal(err)
}
var paired bool
for _, m := range msgs {
if m.Role == llm.RoleTool && m.ToolCallID == "c1" && m.Content == unknownSideEffect {
paired = true
}
}
if !paired {
t.Fatalf("pair missing: %#v", msgs)
}
}
func TestPendingSafeToolRerunsOnce(t *testing.T) {
dir := t.TempDir()
st, err := session.Open(filepath.Join(dir, "s.jsonl"))
if err != nil {
t.Fatal(err)
}
if err := st.Append(llm.Message{
Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
ToolCalls: []llm.ToolCall{{ID: "c1", Name: "read_file", Arguments: `{}`}},
}); err != nil {
t.Fatal(err)
}
if err := st.AppendToolIntent(session.ToolIntent{CallID: "c1", Name: "read_file", Arguments: `{}`, Replay: tools.ReplaySafe}); err != nil {
t.Fatal(err)
}
tool := &replayTool{countingTool: countingTool{name: "read_file"}, replay: tools.ReplaySafe}
reg := tools.NewRegistry(dir)
reg.Register(tool)
_, err = Run(context.Background(), Config{
Model: "fake", Session: st, Tools: reg, MaxTurns: 2,
Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "later", StopReason: llm.StopEnd}),
}, llm.Message{Role: llm.RoleUser, Content: "next"})
if err != nil {
t.Fatal(err)
}
if tool.calls != 1 {
t.Fatalf("calls %d", tool.calls)
}
msgs, err := st.Load()
if err != nil {
t.Fatal(err)
}
for _, m := range msgs {
if m.Role == llm.RoleTool && m.ToolCallID == "c1" {
if m.Content != "ok" || strings.Contains(m.Content, unknownSideEffect) {
t.Fatalf("result %#v", m)
}
return
}
}
t.Fatal("missing tool result")
}
func TestStoredSafeCurrentNeverDoesNotRerun(t *testing.T) {
dir := t.TempDir()
st, err := session.Open(filepath.Join(dir, "s.jsonl"))
if err != nil {
t.Fatal(err)
}
if err := st.AppendToolIntent(session.ToolIntent{CallID: "c1", Name: "read_file", Arguments: `{}`, Replay: tools.ReplaySafe}); err != nil {
t.Fatal(err)
}
tool := &replayTool{countingTool: countingTool{name: "read_file"}, replay: tools.ReplayNever}
reg := tools.NewRegistry(dir)
reg.Register(tool)
if err := recoverPending(context.Background(), Config{Session: st, Tools: reg}); err != nil {
t.Fatal(err)
}
if tool.calls != 0 {
t.Fatalf("calls %d", tool.calls)
}
}
func TestFollowUpIsPersistedBeforeNextRequest(t *testing.T) {
dir := t.TempDir()
st, err := session.Open(filepath.Join(dir, "s.jsonl"))
if err != nil {
t.Fatal(err)
}
sent := false
var saw bool
stream := func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
ch := make(chan llm.Event, 1)
go func() {
defer close(ch)
for _, m := range messages {
if m.Content == "follow-please" {
loaded, err := st.Load()
if err != nil {
t.Errorf("load %v", err)
}
for _, got := range loaded {
if got.Content == "follow-please" {
saw = true
}
}
}
}
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}}
}()
return ch
}
_, err = Run(context.Background(), Config{
Model: "fake", Session: st, MaxTurns: 3, Stream: stream,
GetFollowUp: func() []llm.Message {
if sent {
return nil
}
sent = true
return []llm.Message{{Role: llm.RoleUser, Content: "follow-please"}}
},
}, llm.Message{Role: llm.RoleUser, Content: "go"})
if err != nil {
t.Fatal(err)
}
if !saw {
t.Fatal("follow-up entered the request before it was stored")
}
}
func TestSpansUseAgreedNamesWithoutPrompt(t *testing.T) {
var names []string
var leaked bool
reg := tools.NewRegistry(t.TempDir())
reg.Register(&countingTool{name: "read_file"})
_, err := Run(context.Background(), Config{
Model: "fake", Tools: reg, MaxTurns: 2,
Trace: func(name string, attrs map[string]string) func() {
names = append(names, name)
if _, ok := attrs["prompt"]; ok || attrs["content"] != "" {
leaked = true
}
return func() {}
},
Stream: script(
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ID: "1", Name: "read_file", Arguments: `{}`}}},
llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
),
}, llm.Message{Role: llm.RoleUser, Content: "用户的秘密提示词"})
if err != nil {
t.Fatal(err)
}
joined := strings.Join(names, ",")
for _, want := range []string{"invoke_agent", "chat", "execute_tool"} {
if !strings.Contains(joined, want) {
t.Fatalf("names %s", joined)
}
}
if leaked {
t.Fatal("prompt leaked into span attrs")
}
}
+63
View File
@@ -0,0 +1,63 @@
// Package filestate 记录本轮读过的文件,并在改写前核对修改时间与大小。
package filestate
import (
"fmt"
"os"
"sync"
"time"
)
type stamp struct {
mod time.Time
size int64
}
// Tracker 是一次 run 内的读取记录。nil 接收者表示不检查。
type Tracker struct {
mu sync.Mutex
seen map[string]stamp
}
func New() *Tracker {
return &Tracker{seen: map[string]stamp{}}
}
// NoteRead 记下当前的修改时间与大小。
func (t *Tracker) NoteRead(path string) error {
if t == nil {
return nil
}
info, err := os.Stat(path)
if err != nil {
return err
}
t.mu.Lock()
t.seen[path] = stamp{mod: info.ModTime(), size: info.Size()}
t.mu.Unlock()
return nil
}
// CheckWrite 允许新建。已存在的文件必须在本轮读过,且修改时间与大小未变。
func (t *Tracker) CheckWrite(path string) error {
if t == nil {
return nil
}
info, err := os.Stat(path)
if err != nil {
if os.IsNotExist(err) {
return nil
}
return err
}
t.mu.Lock()
got, ok := t.seen[path]
t.mu.Unlock()
if !ok {
return fmt.Errorf("请先读取该文件再修改:%s", path)
}
if !got.mod.Equal(info.ModTime()) || got.size != info.Size() {
return fmt.Errorf("文件在读取后已变化,请重新读取:%s", path)
}
return nil
}
+40
View File
@@ -0,0 +1,40 @@
package filestate
import (
"os"
"path/filepath"
"strings"
"testing"
"time"
)
func TestCheckWriteRejectsStaleFile(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "note.txt")
if err := os.WriteFile(path, []byte("v1"), 0o644); err != nil {
t.Fatal(err)
}
tr := New()
if err := tr.CheckWrite(path); err == nil || !strings.Contains(err.Error(), "请先读取") {
t.Fatalf("unread write: %v", err)
}
if err := tr.NoteRead(path); err != nil {
t.Fatal(err)
}
if err := tr.CheckWrite(path); err != nil {
t.Fatal(err)
}
// 改大小,并尽量推开修改时间。
time.Sleep(20 * time.Millisecond)
if err := os.WriteFile(path, []byte("v1-changed"), 0o644); err != nil {
t.Fatal(err)
}
err := tr.CheckWrite(path)
if err == nil || !strings.Contains(err.Error(), "请重新读取") {
t.Fatalf("stale write: %v", err)
}
missing := filepath.Join(dir, "new.txt")
if err := tr.CheckWrite(missing); err != nil {
t.Fatalf("create: %v", err)
}
}
+134
View File
@@ -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{}
}
}
+40
View File
@@ -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) }
+17 -2
View File
@@ -42,6 +42,9 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool
emitError(out, "", false, "agentx llm client is missing base url, api key, or model")
return
}
if opts.MessageShape != "" {
messages = PrepareMessages(messages, opts.MessageShape)
}
body := map[string]any{
"model": model,
"messages": toProviderMessages(messages),
@@ -89,8 +92,20 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool
defer resp.Body.Close()
if resp.StatusCode >= 400 {
slurp, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
retryable := resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500
emitError(out, "", retryable, fmt.Sprintf("provider status %d: %s", resp.StatusCode, strings.TrimSpace(string(slurp))))
msg := fmt.Sprintf("provider status %d: %s", resp.StatusCode, strings.TrimSpace(string(slurp)))
ev := Event{
Kind: KindError,
HTTPStatus: resp.StatusCode,
Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg),
ErrorMessage: msg,
Err: &StreamError{Message: msg, Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg)},
Message: Message{Role: RoleAssistant, StopReason: StopError},
}
if d, ok := ParseRetryAfter(resp.Header.Get("Retry-After"), time.Now()); ok {
ev.HasRetryAfter = true
ev.RetryAfter = d
}
out <- ev
return
}
for ev := range ParseSSE(ctx, resp.Body) {
+79
View File
@@ -0,0 +1,79 @@
package llm
import "context"
// WithFallback 在尚未转发任何增量时,遇到 5xx 就改用链上下一个模型。
// fallbacks 为空时返回原来的 StreamFunc,不改变行为。
func WithFallback(base StreamFunc, fallbacks []string, onSwitch func(from, to string)) StreamFunc {
chain := nonEmpty(fallbacks)
if base == nil || len(chain) == 0 {
return base
}
return func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
out := make(chan Event)
go func() {
defer close(out)
models := append([]string{model}, chain...)
for i, current := range models {
forwarded := false
switchNext := false
for ev := range base(ctx, current, messages, tools, opts) {
if isDelta(ev) {
forwarded = true
}
if ev.Kind == KindError && !forwarded && serverError(ev) && i+1 < len(models) {
switchNext = true
if onSwitch != nil {
onSwitch(current, models[i+1])
}
break
}
select {
case out <- ev:
case <-ctx.Done():
return
}
if ev.Kind == KindMessageEnd || ev.Kind == KindError {
return
}
}
if !switchNext {
return
}
}
}()
return out
}
}
func nonEmpty(in []string) []string {
var out []string
seen := map[string]bool{}
for _, s := range in {
s = trimModel(s)
if s == "" || seen[s] {
continue
}
seen[s] = true
out = append(out, s)
}
return out
}
func trimModel(s string) string {
for len(s) > 0 && (s[0] == ' ' || s[0] == '\t') {
s = s[1:]
}
for len(s) > 0 && (s[len(s)-1] == ' ' || s[len(s)-1] == '\t') {
s = s[:len(s)-1]
}
return s
}
func isDelta(ev Event) bool {
return ev.Kind == KindTextDelta || ev.Kind == KindReasoningDelta || ev.Kind == KindToolCallDelta
}
func serverError(ev Event) bool {
return ev.HTTPStatus >= 500 && ev.HTTPStatus <= 599
}
+74
View File
@@ -0,0 +1,74 @@
package llm
import (
"context"
"testing"
)
func TestFallbackSwitchesOn5xxBeforeOutput(t *testing.T) {
var switched string
base := func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
ch := make(chan Event, 2)
go func() {
defer close(ch)
if model == "primary" {
ch <- Event{Kind: KindError, HTTPStatus: 503, ErrorMessage: "provider status 503", Message: Message{Role: RoleAssistant, StopReason: StopError}}
return
}
ch <- Event{Kind: KindTextDelta, Text: "ok"}
ch <- Event{Kind: KindMessageEnd, Message: Message{Role: RoleAssistant, Content: "ok", StopReason: StopEnd}}
}()
return ch
}
stream := WithFallback(base, []string{"backup"}, func(from, to string) {
switched = from + "->" + to
})
var text string
var sawError bool
for ev := range stream(context.Background(), "primary", nil, nil, StreamOptions{}) {
if ev.Kind == KindTextDelta {
text += ev.Text
}
if ev.Kind == KindError {
sawError = true
}
}
if switched != "primary->backup" || text != "ok" || sawError {
t.Fatalf("switched %q text %q error %v", switched, text, sawError)
}
}
func TestFallbackStaysAfterFirstDelta(t *testing.T) {
calls := 0
base := func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
calls++
ch := make(chan Event, 2)
go func() {
defer close(ch)
ch <- Event{Kind: KindTextDelta, Text: "partial"}
ch <- Event{Kind: KindError, HTTPStatus: 500, Message: Message{Role: RoleAssistant, StopReason: StopError}}
}()
return ch
}
stream := WithFallback(base, []string{"backup"}, func(from, to string) {
t.Fatal("switched after output")
})
n := 0
for range stream(context.Background(), "primary", nil, nil, StreamOptions{}) {
n++
}
if calls != 1 || n != 2 {
t.Fatalf("calls %d events %d", calls, n)
}
}
func TestEmptyFallbackIsIdentity(t *testing.T) {
base := func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
ch := make(chan Event)
close(ch)
return ch
}
if WithFallback(base, nil, nil) == nil {
t.Fatal("nil stream")
}
}
+60
View File
@@ -0,0 +1,60 @@
package llm
// PrepareMessages 按 endpoint 整形消息。
// merge 合并相邻、且不携带工具调用的同角色消息。
// strict 在此之后把第一条之后的 system 收成 user,再合并,满足严格交替。
// 工具消息和带 tool_calls 的助手消息保持原样,避免拆开调用与结果。
func PrepareMessages(msgs []Message, shape string) []Message {
switch shape {
case "", "merge", "strict":
default:
return msgs
}
out := mergeSameRole(msgs)
if shape == "strict" {
out = foldExtraSystems(out)
out = mergeSameRole(out)
}
return out
}
func mergeable(m Message) bool {
return m.Role != RoleTool && len(m.ToolCalls) == 0 && m.ToolCallID == ""
}
func mergeSameRole(msgs []Message) []Message {
if len(msgs) == 0 {
return msgs
}
out := make([]Message, 0, len(msgs))
for _, m := range msgs {
if n := len(out); n > 0 && out[n-1].Role == m.Role && mergeable(out[n-1]) && mergeable(m) {
if out[n-1].Content != "" && m.Content != "" {
out[n-1].Content += "\n\n" + m.Content
} else {
out[n-1].Content += m.Content
}
continue
}
out = append(out, m)
}
return out
}
func foldExtraSystems(msgs []Message) []Message {
seenSystem := false
out := make([]Message, 0, len(msgs))
for _, m := range msgs {
if m.Role == RoleSystem {
if seenSystem {
m.Role = RoleUser
if m.Content != "" {
m.Content = "[system]\n" + m.Content
}
}
seenSystem = true
}
out = append(out, m)
}
return out
}
+38
View File
@@ -0,0 +1,38 @@
package llm
import "testing"
func TestPrepareMergesAdjacentSystems(t *testing.T) {
out := PrepareMessages([]Message{
{Role: RoleSystem, Content: "a"},
{Role: RoleSystem, Content: "b"},
{Role: RoleUser, Content: "hi"},
}, "merge")
if len(out) != 2 || out[0].Content != "a\n\nb" {
t.Fatalf("merged = %#v", out)
}
}
func TestPrepareStrictFoldsLaterSystem(t *testing.T) {
out := PrepareMessages([]Message{
{Role: RoleSystem, Content: "base"},
{Role: RoleUser, Content: "hi"},
{Role: RoleSystem, Content: "note"},
}, "strict")
if len(out) != 2 {
t.Fatalf("len = %d %#v", len(out), out)
}
if out[1].Role != RoleUser || out[1].Content != "hi\n\n[system]\nnote" {
t.Fatalf("folded = %#v", out[1])
}
}
func TestPrepareKeepsToolPairs(t *testing.T) {
out := PrepareMessages([]Message{
{Role: RoleAssistant, Content: "call", ToolCalls: []ToolCall{{ID: "1", Name: "t"}}},
{Role: RoleTool, ToolCallID: "1", Content: "ok"},
}, "merge")
if len(out) != 2 || out[1].Role != RoleTool {
t.Fatalf("tools = %#v", out)
}
}
+151 -23
View File
@@ -4,7 +4,10 @@ import (
"context"
"encoding/json"
"fmt"
"math/rand"
"net/http"
"os"
"strconv"
"strings"
"time"
)
@@ -382,31 +385,136 @@ func firstNonEmpty(a, b string) string {
return b
}
const (
retryBaseDelay = 200 * time.Millisecond
maxRetryAfter = 30 * time.Second
)
// overflowNeedles are substrings that mean the prompt does not fit.
// Matching is case-insensitive. A bare "context" is not enough.
var overflowNeedles = []string{
"status 413",
"payload too large",
"context_length_exceeded",
"context length",
"maximum context",
"context window",
"prompt is too long",
"exceeds the maximum",
"exceeds max",
"too many tokens",
"input is too long",
"input length",
"tokens exceed",
"上下文",
"超出最大",
"输入过长",
"提示过长",
"请求过长",
"请求体过大",
}
// ContextOverflow reports provider errors that ask the caller to compact.
func ContextOverflow(message string) bool {
m := strings.ToLower(message)
return strings.Contains(m, "context") && (strings.Contains(m, "overflow") || strings.Contains(m, "too long") || strings.Contains(m, "length"))
for _, needle := range overflowNeedles {
if strings.Contains(m, strings.ToLower(needle)) {
return true
}
}
return false
}
// CapacityLimited reports an in-stream overload that is worth retrying
// before any delta has been forwarded.
func CapacityLimited(message string) bool {
m := strings.ToLower(message)
return strings.Contains(m, "overloaded") || strings.Contains(m, "at capacity") || strings.Contains(m, "too many requests") || strings.Contains(m, "rate limit")
}
// StatusRetryable is true for timeouts, rate limits, server errors, and 529.
func StatusRetryable(code int) bool {
return code == http.StatusRequestTimeout || code == http.StatusTooManyRequests || code == 529 || code >= 500
}
// ParseRetryAfter reads a delay in seconds or an HTTP date.
// The returned duration is capped. ok is false when the header is missing or garbage.
func ParseRetryAfter(raw string, now time.Time) (time.Duration, bool) {
raw = strings.TrimSpace(raw)
if raw == "" {
return 0, false
}
if secs, err := strconv.Atoi(raw); err == nil {
if secs < 0 {
return 0, false
}
return capRetryAfter(time.Duration(secs) * time.Second), true
}
when, err := http.ParseTime(raw)
if err != nil {
return 0, false
}
d := when.Sub(now)
if d < 0 {
d = 0
}
return capRetryAfter(d), true
}
func capRetryAfter(d time.Duration) time.Duration {
if d > maxRetryAfter {
return maxRetryAfter
}
return d
}
func retryBackoff(extraTry int) time.Duration {
if extraTry < 1 {
extraTry = 1
}
shift := extraTry - 1
if shift > 4 {
shift = 4
}
delay := retryBaseDelay * time.Duration(1<<uint(shift))
jitter := time.Duration(rand.Int63n(int64(delay/2) + 1))
return delay + jitter
}
func forwardedDelta(ev Event) bool {
switch ev.Kind {
case KindTextDelta, KindReasoningDelta, KindToolCallDelta:
return true
default:
return false
}
}
// WithRetry re-invokes base on a retryable terminal error, up to attempts
// extra tries. onEvent receives auto_retry_start and auto_retry_end.
// Context-overflow errors are returned immediately so the caller can compact.
// extra tries. Events are forwarded as they arrive. A retry happens only when
// no delta has been forwarded yet. onEvent receives auto_retry_start and
// auto_retry_end. Context-overflow errors are returned immediately.
// attempts <= 0 disables extra tries.
func WithRetry(base StreamFunc, attempts int, onEvent func(kind, detail string)) StreamFunc {
if attempts < 1 {
attempts = 1
if attempts < 0 {
attempts = 0
}
return func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
out := make(chan Event)
go func() {
defer close(out)
var last []Event
headerWait := time.Duration(0)
useHeader := false
for try := 0; try <= attempts; try++ {
if try > 0 {
delay := retryBackoff(try)
if useHeader {
delay = headerWait
}
detail := fmt.Sprintf("attempt %d", try+1)
if onEvent != nil {
onEvent("auto_retry_start", detail)
}
delay := time.Duration(200*(1<<uint(try-1))) * time.Millisecond
select {
case <-ctx.Done():
emitAbort(out, "", ctx.Err())
@@ -414,32 +522,52 @@ func WithRetry(base StreamFunc, attempts int, onEvent func(kind, detail string))
case <-time.After(delay):
}
}
last = nil
useHeader = false
forwarded := false
var terminal *Event
for ev := range base(ctx, model, messages, tools, opts) {
last = append(last, ev)
}
overflow, retryable := false, false
for _, ev := range last {
if ev.Kind == KindError && ContextOverflow(ev.ErrorMessage) {
overflow = true
if terminal != nil {
continue
}
if ev.Kind == KindError && ev.Retryable {
retryable = true
if ev.Kind == KindError {
cp := ev
terminal = &cp
continue
}
if forwardedDelta(ev) {
forwarded = true
}
select {
case <-ctx.Done():
return
case out <- ev:
}
if ev.Kind == KindMessageEnd {
if try > 0 && onEvent != nil {
onEvent("auto_retry_end", fmt.Sprintf("attempt %d", try+1))
}
return
}
}
if overflow || !retryable || try == attempts {
if terminal == nil {
return
}
overflow := ContextOverflow(terminal.ErrorMessage)
if forwarded || overflow || !terminal.Retryable || try == attempts {
if try > 0 && onEvent != nil {
onEvent("auto_retry_end", fmt.Sprintf("attempt %d", try+1))
}
for _, ev := range last {
select {
case <-ctx.Done():
return
case out <- ev:
}
select {
case <-ctx.Done():
return
case out <- *terminal:
}
return
}
if terminal.HasRetryAfter {
useHeader = true
headerWait = terminal.RetryAfter
}
}
}()
return out
+70
View File
@@ -9,6 +9,7 @@ import (
"path/filepath"
"strings"
"testing"
"time"
)
func TestParseModelRefAndThinking(t *testing.T) {
@@ -141,6 +142,75 @@ func TestRetryThenSuccess(t *testing.T) {
}
}
func TestContextOverflowPatterns(t *testing.T) {
hits := []string{
"provider status 413: payload too large",
"context_length_exceeded",
"prompt is too long",
"input length exceeds the maximum",
"maximum context length exceeded",
"context length overflow",
"too many tokens",
"输入过长,请缩短上下文",
"请求过长",
"超出最大长度",
}
for _, msg := range hits {
if !ContextOverflow(msg) {
t.Fatalf("want overflow: %s", msg)
}
}
if ContextOverflow("context cancelled") || ContextOverflow("temporary network error") {
t.Fatal("unrelated context text must not count as overflow")
}
}
func TestParseRetryAfter(t *testing.T) {
now := time.Date(2026, 10, 3, 12, 0, 0, 0, time.UTC)
if d, ok := ParseRetryAfter("2", now); !ok || d != 2*time.Second {
t.Fatalf("seconds: %v %v", d, ok)
}
if d, ok := ParseRetryAfter("not-a-date", now); ok || d != 0 {
t.Fatalf("garbage should fail: %v %v", d, ok)
}
if _, ok := ParseRetryAfter("", now); ok {
t.Fatal("empty header")
}
far := now.Add(2 * time.Hour).Format(http.TimeFormat)
if d, ok := ParseRetryAfter(far, now); !ok || d != maxRetryAfter {
t.Fatalf("date cap: %v %v", d, ok)
}
past := now.Add(-time.Minute).Format(http.TimeFormat)
if d, ok := ParseRetryAfter(past, now); !ok || d != 0 {
t.Fatalf("past date: %v %v", d, ok)
}
}
func TestRetryDoesNotRepeatDeltas(t *testing.T) {
n := 0
base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
out := make(chan Event, 2)
n++
out <- Event{Kind: KindTextDelta, Text: "a"}
out <- Event{Kind: KindError, Retryable: true, ErrorMessage: "overloaded"}
close(out)
return out
})
var texts []string
var errors int
for ev := range WithRetry(base, 2, nil)(context.Background(), "m", nil, nil, StreamOptions{}) {
if ev.Kind == KindTextDelta {
texts = append(texts, ev.Text)
}
if ev.Kind == KindError {
errors++
}
}
if n != 1 || len(texts) != 1 || texts[0] != "a" || errors != 1 {
t.Fatalf("calls=%d texts=%v errors=%d", n, texts, errors)
}
}
func TestContextOverflowIsNotRetried(t *testing.T) {
if !ContextOverflow("maximum context length exceeded") {
t.Fatal("expected overflow")
+35 -2
View File
@@ -56,6 +56,7 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
order := []int{}
var usage Usage
finished := ""
sawDone := false
sc := bufio.NewScanner(r)
sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
reader := &sseEventReader{sc: sc}
@@ -69,6 +70,7 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
return
}
if data == "[DONE]" {
sawDone = true
break
}
var chunk sseChunk
@@ -77,7 +79,9 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
return
}
if chunk.Error != nil && chunk.Error.Message != "" {
emitError(out, content.String(), false, chunk.Error.Message)
msg := chunk.Error.Message
retryable := CapacityLimited(msg) && !ContextOverflow(msg)
emitError(out, content.String(), retryable, msg)
return
}
if chunk.Usage != nil {
@@ -151,6 +155,16 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
emitError(out, content.String(), true, err.Error())
return
}
// 取消必须压过正常结束。select 在两个分支都就绪时随机挑选,
// 竞态下会把 message_end 发给已经取消的调用方。
if ctx.Err() != nil {
emitAbort(out, content.String(), ctx.Err())
return
}
// length、空结束原因,或没有 [DONE] 的中断,都不把半截工具调用交给执行层。
if !toolFinishExecutable(finished, sawDone) {
order = nil
}
msg := Message{Role: RoleAssistant, Content: content.String(), StopReason: mapFinish(finished, len(order) > 0)}
for _, idx := range order {
msg.ToolCalls = append(msg.ToolCalls, *calls[idx])
@@ -216,13 +230,32 @@ func (r *sseEventReader) next() (string, bool) {
return "", false
}
// toolFinishExecutable 只在供应商明确以工具调用或正常 stop 结束时执行。
// length、缺失的结束原因,以及没有 [DONE] 的断开,都视为未完成。
func toolFinishExecutable(reason string, sawDone bool) bool {
switch reason {
case "tool_calls", "function_call":
return true
case "stop":
return true
case "length":
return false
case "":
return false
default:
return sawDone && reason != "error" && reason != "content_filter"
}
}
func mapFinish(reason string, hasTools bool) StopReason {
switch reason {
case "tool_calls", "function_call":
return StopToolCalls
case "content_filter", "error":
return StopError
case "stop", "length", "":
case "length":
return StopEnd
case "stop", "":
if hasTools {
return StopToolCalls
}
+45
View File
@@ -93,6 +93,51 @@ func TestParseSSEToolCallIncrements(t *testing.T) {
}
}
func TestParseSSELengthDropsUnfinishedToolCall(t *testing.T) {
body := strings.Join([]string{
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":"{\"path\":"}}]}}]}`,
`data: {"choices":[{"delta":{},"finish_reason":"length"}]}`,
`data: [DONE]`,
"",
}, "\n")
var end Event
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
if ev.Kind == KindMessageEnd {
end = ev
}
}
if end.Message.StopReason != StopEnd || len(end.Message.ToolCalls) != 0 {
t.Fatalf("truncated tool call must not be executable: %#v", end.Message)
}
if end.FinishReason != "length" {
t.Fatalf("finish %q", end.FinishReason)
}
}
func TestParseSSEUncleanEndDropsToolCall(t *testing.T) {
body := `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":"{\"path\":"}}]}}]}`
var end Event
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
if ev.Kind == KindMessageEnd {
end = ev
}
}
if end.Kind != KindMessageEnd || end.Message.StopReason != StopEnd || len(end.Message.ToolCalls) != 0 {
t.Fatalf("unclean stream must not execute tools: %#v", end)
}
}
func TestParseSSECapacityErrorIsRetryable(t *testing.T) {
body := "data: {\"error\":{\"message\":\"server is at capacity\"}}\n\n"
var got Event
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
got = ev
}
if got.Kind != KindError || !got.Retryable || !CapacityLimited(got.ErrorMessage) {
t.Fatalf("event %#v", got)
}
}
func TestParseSSEMalformedIsError(t *testing.T) {
body := "data: {not-json}\n\n"
var got Event
+21 -6
View File
@@ -3,6 +3,7 @@ package llm
import (
"context"
"encoding/json"
"time"
)
// Role is a transcript role carried into the provider request.
@@ -69,13 +70,19 @@ type Image struct {
Data string `json:"data,omitempty"`
}
// Usage counts tokens reported by the provider. Zero means the provider omitted usage.
// Usage counts tokens. On one provider response, PromptTokens, CompletionTokens
// and TotalTokens are that request. On a loop Result they are the billing sum
// across requests; Last* is only the latest request, for context observation.
// Zero means the provider omitted usage.
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
CacheReadTokens int `json:"cache_read_tokens,omitempty"`
CacheWriteTokens int `json:"cache_write_tokens,omitempty"`
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
CacheReadTokens int `json:"cache_read_tokens,omitempty"`
CacheWriteTokens int `json:"cache_write_tokens,omitempty"`
LastPromptTokens int `json:"last_prompt_tokens,omitempty"`
LastCompletionTokens int `json:"last_completion_tokens,omitempty"`
LastTotalTokens int `json:"last_total_tokens,omitempty"`
}
// Event is one streamed update. Failures are KindError plus a final assistant
@@ -96,6 +103,12 @@ type Event struct {
Err error
Retryable bool
ErrorMessage string
// HTTPStatus 是提供方的 HTTP 状态。0 表示这条错误不是 HTTP 响应。
HTTPStatus int
// RetryAfter is the provider's requested wait. HasRetryAfter distinguishes
// a real zero from a missing or unparseable header.
RetryAfter time.Duration
HasRetryAfter bool
}
// StreamOptions tunes one provider request.
@@ -106,6 +119,8 @@ type StreamOptions struct {
// ThinkingLevel is off, minimal, low, medium, high, xhigh, or max.
// Empty leaves the provider default unchanged.
ThinkingLevel string
// MessageShape is merge, strict, or empty. Empty leaves the transcript unchanged.
MessageShape string
}
// StreamFunc is the provider boundary. It must return a channel that is closed
+227 -28
View File
@@ -8,7 +8,9 @@ import (
"sync"
"time"
"agent-native/runtime/hooks"
"agent-native/runtime/llm"
"agent-native/runtime/reminders"
"agent-native/runtime/session"
"agent-native/runtime/tools"
@@ -45,11 +47,14 @@ type Event struct {
// Config drives one agent session.
type Config struct {
Model string
SystemPrompt string
Tools *tools.Registry
Stream llm.StreamFunc
Session *session.Store
Model string
SystemPrompt string
Tools *tools.Registry
Stream llm.StreamFunc
Session *session.Store
// Trace 开始一个 span,返回的函数结束它。nil 表示不采集。
// 属性里不要放提示词正文。
Trace func(name string, attrs map[string]string) func()
MaxTurns int
ToolExecution string
GetSteering func() []llm.Message
@@ -83,7 +88,7 @@ type Config struct {
// 同一次 Run 内与 OnEvent、BeforeTool 串行,不会并发调用。
AfterTool func(call llm.ToolCall, res llm.Message)
// Reminders 在每次调用模型前提供提醒。正文包在 system-reminder 里,
// 只进入本次请求,不写进 Session。
// 并入最后一条工具结果或用户消息,只进入本次请求,不写进 Session。
Reminders func(turn int) []string
// RepeatGuard 为真时,同名且参数按键排序后相同的调用连续 3 次,注入一条提醒。
// 每个 run 最多 3 条,不中断循环。
@@ -94,6 +99,17 @@ type Config struct {
// MaxLengthContinuations 是 length 续写的次数上限。
// 0 表示默认 3 次;负数表示不限制。
MaxLengthContinuations int
// ContextWindow 是模型窗口的字符估算。0 表示用保守默认值。
// 超过窗口减去 OutputReserve 与 ContextBuffer 时,清理最近若干条之前的旧工具结果。
ContextWindow int
// OutputReserve 是留给模型输出的字符数。0 表示默认预留。
OutputReserve int
// ContextBuffer 是窗口估算的余量。0 表示默认余量。
ContextBuffer int
// KeepToolResults 是清理时保留的最近工具结果条数。0 表示 5。
KeepToolResults int
// Hooks 是进程内钩子。空总线不改变停止、提醒和拒绝次数。
Hooks *hooks.Bus
}
// StopInfo 是 Stop 钩子看到的本 run 进度。
@@ -141,6 +157,8 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
cfg.ToolExecution = ExecSequential
}
serializeHostCallbacks(&cfg)
endRun := startSpan(cfg, "invoke_agent", map[string]string{"model": cfg.Model})
defer endRun()
var messages []llm.Message
if cfg.Session != nil {
loaded, err := cfg.Session.Load()
@@ -149,6 +167,16 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
}
messages = loaded
}
if err := recoverPending(ctx, cfg); err != nil {
return Result{}, err
}
if cfg.Session != nil {
loaded, err := cfg.Session.Load()
if err != nil {
return Result{}, err
}
messages = loaded
}
if len(messages) == 0 && len(cfg.History) > 0 {
messages = append([]llm.Message{}, cfg.History...)
}
@@ -182,7 +210,7 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
}
emit(cfg, Event{Type: "turn_start"})
emit(cfg, Event{Type: "message_start"})
request := withEphemeral(messages, turnReminders(cfg, turn, &repeats, &lengthNudge))
request := prepareRequest(cfg, messages, turn, &repeats, &lengthNudge)
end, errEv := consume(ctx, cfg, request)
if errEv != nil {
cfg.logger().Error("runtime turn failed", errEv, "retryable", end.Retryable)
@@ -192,6 +220,9 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
usage.PromptTokens += end.Usage.PromptTokens
usage.CompletionTokens += end.Usage.CompletionTokens
usage.TotalTokens += end.Usage.TotalTokens
usage.LastPromptTokens = end.Usage.PromptTokens
usage.LastCompletionTokens = end.Usage.CompletionTokens
usage.LastTotalTokens = end.Usage.TotalTokens
if err := persist(cfg, last); err != nil {
return Result{}, err
}
@@ -262,8 +293,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 {
@@ -334,6 +364,8 @@ func systemText(cfg Config) string {
}
func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event, error) {
endSpan := startSpan(cfg, "chat", map[string]string{"model": cfg.Model})
defer endSpan()
var decls []llm.ToolDecl
if cfg.Tools != nil {
for _, d := range cfg.Tools.Decls() {
@@ -464,15 +496,96 @@ func cancelledToolMessage(call llm.ToolCall, err error) llm.Message {
}
}
func rejectIncompleteToolCall(call llm.ToolCall) string {
if strings.TrimSpace(call.Name) == "" {
return "工具调用未执行:名称不完整"
}
args := strings.TrimSpace(call.Arguments)
if !json.Valid([]byte(args)) || !strings.HasPrefix(args, "{") {
return "工具调用未执行:参数不是完整的 JSON 对象"
}
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)
}
if why := rejectIncompleteToolCall(call); why != "" {
msg := llm.Message{Role: llm.RoleTool, Name: call.Name, ToolCallID: call.ID, Content: why, IsError: true}
emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name, Detail: call.Arguments})
emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: why})
if cfg.AfterTool != nil {
cfg.AfterTool(call, msg)
}
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)
}
@@ -480,6 +593,13 @@ func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message {
}
}
emit(cfg, Event{Type: "tool_execution_update", Tool: call.Name, Detail: call.Arguments})
if err := commitToolIntent(cfg, 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})
return msg
}
endSpan := startSpan(cfg, "execute_tool", map[string]string{"tool": call.Name})
defer endSpan()
res := tools.Result{Content: "no tools configured", IsError: true}
known := cfg.Tools != nil && cfg.Tools.Has(call.Name)
switch {
@@ -498,6 +618,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)
}
@@ -511,22 +632,110 @@ func settle(cfg Config, last llm.Message, messages []llm.Message, usage llm.Usag
return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: usage, Notice: notice}
}
func reminderMessage(text string) llm.Message {
return llm.Message{Role: llm.RoleUser, Content: "<system-reminder>\n" + text + "\n</system-reminder>"}
func prepareRequest(cfg Config, messages []llm.Message, turn int, repeats *repeatState, lengthNudge *bool) []llm.Message {
request := clearOldResults(cfg, persistedProjection(cfg, messages))
return reminders.Attach(request, turnReminderTexts(cfg, turn, repeats, lengthNudge))
}
func turnReminders(cfg Config, turn int, repeats *repeatState, lengthNudge *bool) []llm.Message {
var extra []llm.Message
// persistedProjection 在有会话存储时只采用已经落盘的消息。
// steer 与 follow-up 必须先写入存储,下一轮请求才会看见它们。
func persistedProjection(cfg Config, messages []llm.Message) []llm.Message {
if cfg.Session == nil {
return messages
}
loaded, err := cfg.Session.Load()
if err != nil {
return messages
}
return loaded
}
const unknownSideEffect = "副作用未知,未重试"
func recoverPending(ctx context.Context, cfg Config) error {
if cfg.Session == nil {
return nil
}
pending, err := cfg.Session.PendingIntents()
if err != nil {
return err
}
for _, in := range pending {
current := tools.ReplayNever
if cfg.Tools != nil {
current = cfg.Tools.Replay(in.Name)
}
msg := llm.Message{
Role: llm.RoleTool, Name: in.Name, ToolCallID: in.CallID,
Content: unknownSideEffect, IsError: true,
}
if in.Replay == tools.ReplaySafe && current == tools.ReplaySafe && cfg.Tools != nil {
res := cfg.Tools.Execute(ctx, in.Name, json.RawMessage(in.Arguments))
msg.Content = res.Content
msg.IsError = res.IsError
}
if err := persist(cfg, msg); err != nil {
return err
}
}
return nil
}
func commitToolIntent(cfg Config, call llm.ToolCall) error {
if cfg.Session == nil {
return nil
}
replay := tools.ReplayNever
if cfg.Tools != nil {
replay = cfg.Tools.Replay(call.Name)
}
return cfg.Session.AppendToolIntent(session.ToolIntent{
CallID: call.ID, Name: call.Name, Arguments: call.Arguments, Replay: replay,
})
}
func startSpan(cfg Config, name string, attrs map[string]string) func() {
if cfg.Trace == nil {
return func() {}
}
end := cfg.Trace(name, attrs)
if end == nil {
return func() {}
}
return end
}
func clearOldResults(cfg Config, messages []llm.Message) []llm.Message {
window := cfg.ContextWindow
if window <= 0 {
window = session.DefaultContextWindow
}
reserve := cfg.OutputReserve
if reserve <= 0 {
reserve = session.DefaultOutputReserve
}
buffer := cfg.ContextBuffer
if buffer <= 0 {
buffer = session.DefaultContextBuffer
}
if !session.NeedsClear(session.EstimateChars(messages), window, reserve, buffer) {
return messages
}
return session.ClearOldToolResults(messages, cfg.KeepToolResults)
}
func turnReminderTexts(cfg Config, turn int, repeats *repeatState, lengthNudge *bool) []string {
var extra []string
if cfg.Reminders != nil {
for _, text := range cfg.Reminders(turn) {
if strings.TrimSpace(text) == "" {
continue
}
extra = append(extra, reminderMessage(text))
extra = append(extra, text)
}
}
if lengthNudge != nil && *lengthNudge {
extra = append(extra, llm.Message{Role: llm.RoleUser, Content: lengthContinuePrompt})
extra = append(extra, lengthContinuePrompt)
*lengthNudge = false
}
if cfg.RepeatGuard && repeats != nil {
@@ -536,23 +745,13 @@ func turnReminders(cfg Config, turn int, repeats *repeatState, lengthNudge *bool
due = room
}
for i := 0; i < due; i++ {
extra = append(extra, reminderMessage(repeatReminderText))
extra = append(extra, repeatReminderText)
repeats.reminders++
}
}
return extra
}
func withEphemeral(messages, extra []llm.Message) []llm.Message {
if len(extra) == 0 {
return messages
}
out := make([]llm.Message, 0, len(messages)+len(extra))
out = append(out, messages...)
out = append(out, extra...)
return out
}
type repeatState struct {
sig string
streak int
@@ -616,8 +815,8 @@ func budgetExceeded(b *Budget, usage llm.Usage, toolCalls int, now time.Time) st
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.MaxPromptTokens > 0 && usage.LastPromptTokens > b.MaxPromptTokens {
return fmt.Sprintf("budget_exceeded: prompt tokens %d > max %d", usage.LastPromptTokens, b.MaxPromptTokens)
}
if b.MaxCompletionTokens > 0 && usage.CompletionTokens > b.MaxCompletionTokens {
return fmt.Sprintf("budget_exceeded: completion tokens %d > max %d", usage.CompletionTokens, b.MaxCompletionTokens)
+51
View File
@@ -194,6 +194,57 @@ func TestLoopBudgetPromptTokens(t *testing.T) {
}
}
func TestLoopUsageSplitsLastRequestFromCumulative(t *testing.T) {
res, err := Run(context.Background(), Config{
Model: "fake",
Tools: tools.WorkspaceTools(t.TempDir()),
Budget: &Budget{MaxPromptTokens: 10},
Stream: usageScript(llm.Usage{PromptTokens: 6, CompletionTokens: 1, TotalTokens: 7},
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
ID: "c1", Name: "list_dir", Arguments: `{}`,
}}},
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("last request stays under the window: err=%v text=%q", res.Err, res.Text)
}
if res.Usage.PromptTokens != 12 || res.Usage.LastPromptTokens != 6 {
t.Fatalf("billing sum %d last %d", res.Usage.PromptTokens, res.Usage.LastPromptTokens)
}
}
func TestLoopRejectsIncompleteToolCall(t *testing.T) {
tool := &countingTool{name: "list_dir"}
reg := tools.NewRegistry(t.TempDir())
reg.Register(tool)
res, err := Run(context.Background(), Config{
Model: "fake",
Tools: reg,
Stream: script(llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
ID: "c1", Name: "list_dir", Arguments: `{"path":`,
}}}, llm.Message{Role: llm.RoleAssistant, Content: "stopped", StopReason: llm.StopEnd}),
}, llm.Message{Role: llm.RoleUser, Content: "go"})
if err != nil {
t.Fatal(err)
}
if tool.calls != 0 {
t.Fatalf("incomplete call executed %d times", tool.calls)
}
var paired bool
for _, m := range res.Messages {
if m.Role == llm.RoleTool && m.ToolCallID == "c1" && m.IsError {
paired = true
}
}
if !paired {
t.Fatal("expected a paired error result")
}
}
func TestLoopBudgetDeadline(t *testing.T) {
res, err := Run(context.Background(), Config{
Model: "fake",
+27
View File
@@ -0,0 +1,27 @@
package mcp
import "context"
type askKey struct{}
// WithAsk 把当前回合的 AskUser 放进上下文。MCP 返回 input_required 时使用它。
func WithAsk(ctx context.Context, fn func(context.Context, string) (string, error)) context.Context {
return context.WithValue(ctx, askKey{}, fn)
}
func askFrom(ctx context.Context) func(context.Context, string) (string, error) {
if ctx == nil {
return nil
}
fn, _ := ctx.Value(askKey{}).(func(context.Context, string) (string, error))
return fn
}
func inputQuestion(res callToolResult) (string, bool) {
for _, blk := range res.Content {
if blk.Type == "input_required" {
return blk.Text, true
}
}
return "", false
}
+144 -48
View File
@@ -19,6 +19,8 @@ const (
maxToolNameBytes = 128
callTimeoutDefault = 30 * time.Second
registerTimeoutFixed = 30 * time.Second
// 工具数超过这个值且没有显式暴露级别时,整台服务器改为延迟暴露。
deferredToolThreshold = 8
)
// Manager 维护一组 MCP server 连接,并把它们的工具注册进 tools.Registry。
@@ -72,55 +74,138 @@ func (m *Manager) client(ctx context.Context, cfg ServerConfig) (Client, error)
return c, nil
}
// RegisterInto 连接全部 server、枚举工具并以 mcp__<server>__<tool> 命名注册。
// 任一 server 不可达即返回错误(fail-closed,不静默跳过)。
func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error {
for _, cfg := range m.configs {
c, err := m.client(ctx, cfg)
if err != nil {
// ApplyOverrides 用项目配置覆盖启用状态和暴露级别。空字符串不改。
func (m *Manager) ApplyOverrides(raw string) error {
if m == nil || strings.TrimSpace(raw) == "" {
return nil
}
var over map[string]struct {
Enabled *bool `json:"enabled"`
Exposure string `json:"exposure"`
}
if err := json.Unmarshal([]byte(raw), &over); err != nil {
return err
}
for i, cfg := range m.configs {
o, ok := over[cfg.Name]
if !ok {
continue
}
if o.Enabled != nil {
m.configs[i].Enabled = o.Enabled
}
if o.Exposure != "" {
m.configs[i].Exposure = o.Exposure
}
if err := m.configs[i].Validate(); err != nil {
return err
}
listCtx, cancel := context.WithTimeout(ctx, registerTimeoutFixed)
list, err := c.ListTools(listCtx)
cancel()
if err != nil {
return fmt.Errorf("mcp server %q tools/list failed: %w", cfg.Name, err)
}
if len(list) > maxToolsPerServer {
return fmt.Errorf("mcp server %q exposes %d tools (max %d)", cfg.Name, len(list), maxToolsPerServer)
}
for _, ti := range list {
name := ToolName(cfg.Name, ti.Name)
if len(name) > maxToolNameBytes {
return fmt.Errorf("mcp tool name too long: %s", name)
}
desc := ti.Description
if len(desc) > maxDescriptionBytes {
desc = tools.CutUTF8(desc, maxDescriptionBytes) + "…[truncated]"
}
schema := ti.InputSchema
if len(schema) == 0 {
schema = json.RawMessage(`{"type":"object","properties":{}}`)
}
reg.Register(&bridgeTool{
client: c,
serverCfg: cfg,
toolName: ti.Name,
declName: name,
description: desc,
schema: schema,
redial: func(ctx context.Context) (Client, error) {
m.mu.Lock()
delete(m.clients, cfg.Name)
m.mu.Unlock()
return m.client(ctx, cfg)
},
})
}
}
return nil
}
// Ready 报告已启用的服务器是否都已经连上。已连上时注册只是枚举工具,不再握手。
func (m *Manager) Ready() bool {
if m == nil {
return true
}
m.mu.Lock()
defer m.mu.Unlock()
for _, cfg := range m.configs {
if !cfg.IsEnabled() {
continue
}
if _, ok := m.clients[cfg.Name]; !ok {
return false
}
}
return true
}
func (m *Manager) StartBackground(ctx context.Context, reg *tools.Registry) {
if m == nil || reg == nil {
return
}
for _, cfg := range m.configs {
cfg := cfg
if !cfg.IsEnabled() {
continue
}
go func() { _ = m.registerServer(ctx, reg, cfg) }()
}
}
// RegisterInto 连接全部已启用的 server 并注册工具。未启用的跳过。
func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error {
for _, cfg := range m.configs {
if !cfg.IsEnabled() {
continue
}
if err := m.registerServer(ctx, reg, cfg); err != nil {
return err
}
}
return nil
}
func (m *Manager) registerServer(ctx context.Context, reg *tools.Registry, cfg ServerConfig) error {
c, err := m.client(ctx, cfg)
if err != nil {
return err
}
listCtx, cancel := context.WithTimeout(ctx, registerTimeoutFixed)
list, err := c.ListTools(listCtx)
cancel()
if err != nil {
return fmt.Errorf("mcp server %q tools/list failed: %w", cfg.Name, err)
}
if len(list) > maxToolsPerServer {
return fmt.Errorf("mcp server %q exposes %d tools (max %d)", cfg.Name, len(list), maxToolsPerServer)
}
exposure := chooseExposure(cfg.Exposure, len(list))
for _, ti := range list {
name := ToolName(cfg.Name, ti.Name)
if len(name) > maxToolNameBytes {
return fmt.Errorf("mcp tool name too long: %s", name)
}
desc := ti.Description
if len(desc) > maxDescriptionBytes {
desc = tools.CutUTF8(desc, maxDescriptionBytes) + "…[truncated]"
}
schema := ti.InputSchema
if len(schema) == 0 {
schema = json.RawMessage(`{"type":"object","properties":{}}`)
}
reg.Register(&bridgeTool{
client: c,
serverCfg: cfg,
toolName: ti.Name,
declName: name,
description: desc,
schema: schema,
exposure: exposure,
redial: func(ctx context.Context) (Client, error) {
m.mu.Lock()
delete(m.clients, cfg.Name)
m.mu.Unlock()
return m.client(ctx, cfg)
},
})
}
return nil
}
func chooseExposure(explicit string, n int) string {
switch explicit {
case tools.ExposureDirect, tools.ExposureDeferred:
return explicit
default:
if n > deferredToolThreshold {
return tools.ExposureDeferred
}
return tools.ExposureDirect
}
}
// ToolName 生成桥接工具名(命名空间隔离)。
func ToolName(server, tool string) string {
return "mcp__" + server + "__" + tool
@@ -144,14 +229,14 @@ type bridgeTool struct {
declName string
description string
schema json.RawMessage
exposure string
redial func(ctx context.Context) (Client, error)
}
func (t *bridgeTool) Name() string { return t.declName }
func (t *bridgeTool) Description() string { return t.description }
func (t *bridgeTool) Parameters() json.RawMessage {
return t.schema
}
func (t *bridgeTool) Name() string { return t.declName }
func (t *bridgeTool) Description() string { return t.description }
func (t *bridgeTool) Parameters() json.RawMessage { return t.schema }
func (t *bridgeTool) Exposure() string { return t.exposure }
func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
res, err := t.call(ctx, args)
@@ -166,6 +251,17 @@ func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.R
if err != nil {
return tools.Result{Content: fmt.Sprintf("mcp tool %s failed: %v", t.declName, err), IsError: true}, nil
}
if question, ok := inputQuestion(res); ok {
ask := askFrom(ctx)
if ask == nil {
return tools.Result{Content: "需要用户输入:" + question, IsError: true}, nil
}
answer, aerr := ask(ctx, question)
if aerr != nil {
return tools.Result{Content: aerr.Error(), IsError: true}, nil
}
return tools.Result{Content: answer}, nil
}
var b strings.Builder
for _, blk := range res.Content {
if blk.Type == "text" && blk.Text != "" {
+113
View File
@@ -99,6 +99,9 @@ func TestStdioSlowCancelThenEcho(t *testing.T) {
if len(res.Content) != 1 || res.Content[0].Text != "echo: after-slow" {
t.Fatalf("echo must receive its own result, got %+v", res)
}
if c.Readers() != 1 {
t.Fatalf("cancel must leave a single reader, got %d", c.Readers())
}
}
func fakeHTTPHandler(w http.ResponseWriter, r *http.Request) {
@@ -265,3 +268,113 @@ func TestManagerRegisterTimeoutBounded(t *testing.T) {
t.Fatalf("timeout must bound the dial, took %s", elapsed)
}
}
func TestBackgroundConnectDoesNotBlockTheCaller(t *testing.T) {
started := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
select {
case <-started:
default:
close(started)
}
time.Sleep(400 * time.Millisecond)
fakeHTTPHandler(w, r)
}))
defer srv.Close()
mgr, err := NewManager([]ServerConfig{{Name: "late", URL: srv.URL}})
if err != nil {
t.Fatal(err)
}
defer mgr.Close()
reg := tools.NewRegistry(t.TempDir())
t0 := time.Now()
mgr.StartBackground(context.Background(), reg)
if time.Since(t0) > 150*time.Millisecond {
t.Fatalf("background connect blocked the caller for %s", time.Since(t0))
}
deadline := time.Now().Add(3 * time.Second)
for time.Now().Before(deadline) {
if reg.Has("mcp__late__echo") {
return
}
time.Sleep(20 * time.Millisecond)
}
t.Fatal("tools did not appear after the server became ready")
}
func TestDeferredMCPToolsUseSearchAndCall(t *testing.T) {
mgr, err := NewManager([]ServerConfig{{Name: "fake", Command: fakeStdioConfig(t).Command, Args: fakeStdioConfig(t).Args, Env: fakeStdioConfig(t).Env, Exposure: "deferred"}})
if err != nil {
t.Fatal(err)
}
defer mgr.Close()
reg := tools.NewRegistry(t.TempDir())
if err := mgr.RegisterInto(context.Background(), reg); err != nil {
t.Fatal(err)
}
for _, d := range reg.Decls() {
if strings.HasPrefix(d.Name, "mcp__fake__") {
t.Fatalf("deferred mcp tool listed: %s", d.Name)
}
}
if !strings.Contains(reg.Search("echo"), "mcp__fake__echo") {
t.Fatalf("search %q", reg.Search("echo"))
}
res := reg.Execute(context.Background(), "call_tool", json.RawMessage(`{"name":"mcp__fake__echo","arguments":{"text":"via"}}`))
if res.IsError || !strings.Contains(res.Content, "echo: via") {
t.Fatalf("call %+v", res)
}
}
func TestDisabledServerIsSkipped(t *testing.T) {
off := false
mgr, err := NewManager([]ServerConfig{{Name: "fake", Command: "x", Enabled: &off}})
if err != nil {
t.Fatal(err)
}
reg := tools.NewRegistry(t.TempDir())
if err := mgr.RegisterInto(context.Background(), reg); err != nil {
t.Fatal(err)
}
if reg.Has("mcp__fake__echo") {
t.Fatal("disabled server registered a tool")
}
if err := mgr.ApplyOverrides(`{"fake":{"enabled":true,"exposure":"direct"}}`); err != nil {
t.Fatal(err)
}
if !mgr.Configs()[0].IsEnabled() || mgr.Configs()[0].Exposure != "direct" {
t.Fatalf("override %+v", mgr.Configs()[0])
}
}
func TestManyToolsDefaultToDeferred(t *testing.T) {
if chooseExposure("", 9) != tools.ExposureDeferred {
t.Fatal("many tools should default to deferred")
}
if chooseExposure("direct", 9) != tools.ExposureDirect {
t.Fatal("explicit direct must win")
}
}
type askClient struct{}
func (askClient) Initialize(context.Context) (initializeResult, error) {
return initializeResult{}, nil
}
func (askClient) ListTools(context.Context) ([]toolInfo, error) { return nil, nil }
func (askClient) Close() error { return nil }
func (askClient) CallTool(context.Context, string, json.RawMessage) (callToolResult, error) {
return callToolResult{Content: []contentBlock{{Type: "input_required", Text: "选一个颜色"}}}, nil
}
func TestInputRequiredUsesAsk(t *testing.T) {
tool := &bridgeTool{
client: askClient{}, declName: "mcp__x__ask", toolName: "ask",
schema: json.RawMessage(`{"type":"object","properties":{}}`),
}
ctx := WithAsk(context.Background(), func(context.Context, string) (string, error) { return "蓝色", nil })
res, err := tool.Execute(ctx, json.RawMessage(`{}`))
if err != nil || res.IsError || res.Content != "蓝色" {
t.Fatalf("ask result %+v %v", res, err)
}
}
+14
View File
@@ -25,6 +25,10 @@ type ServerConfig struct {
Headers map[string]string `json:"headers,omitempty"`
// Timeout 单次请求上限(默认 30s)。
Timeout time.Duration `json:"timeout,omitempty"`
// Enabled 为 false 时跳过这台服务器。nil 表示启用。
Enabled *bool `json:"enabled,omitempty"`
// Exposure 为 direct 或 deferred。空则工具多时延迟暴露。
Exposure string `json:"exposure,omitempty"`
}
// Validate 校验配置完整性(fail-closed:命名不规范或双传输/无传输都拒绝)。
@@ -37,9 +41,19 @@ func (c ServerConfig) Validate() error {
if hasCmd == hasURL {
return fmt.Errorf("mcp server %q: exactly one of command or url is required", c.Name)
}
switch c.Exposure {
case "", "direct", "deferred":
default:
return fmt.Errorf("mcp server %q: exposure must be direct or deferred", c.Name)
}
return nil
}
// IsEnabled 报告这台服务器是否要连接。
func (c ServerConfig) IsEnabled() bool {
return c.Enabled == nil || *c.Enabled
}
func validName(s string) bool {
if s == "" {
return false
+5
View File
@@ -22,6 +22,7 @@ type stdioClient struct {
mu sync.Mutex
pending map[int64]chan stdioResult
nextID int64
readers atomic.Int32
closed bool
readerOnce sync.Once
}
@@ -195,6 +196,8 @@ func (c *stdioClient) startReader() {
}
func (c *stdioClient) readLoop() {
c.readers.Add(1)
defer c.readers.Add(-1)
var fatal error
for {
line, err := c.stdout.ReadBytes('\n')
@@ -227,6 +230,8 @@ func (c *stdioClient) readLoop() {
}
}
func (c *stdioClient) Readers() int { return int(c.readers.Load()) }
func (c *stdioClient) notify(method string) error {
c.mu.Lock()
closed := c.closed
+221
View File
@@ -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", "tool_search":
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", "run_tool_script", "call_tool":
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
}
+41
View File
@@ -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)
}
}
+97
View File
@@ -0,0 +1,97 @@
// Package reminders 收集一轮里要带给模型的提醒,并做节流。
// 提醒并入最后一条工具结果或用户消息,不另起一条 user 消息。
package reminders
import (
"strings"
"sync"
"time"
"agent-native/runtime/llm"
)
type Kind string
const (
KindTodo Kind = "todo"
KindPlan Kind = "plan"
KindBackground Kind = "background"
KindDeliverable Kind = "deliverable"
KindBudget Kind = "budget"
)
// Note 是一条待投递的提醒。Key 相同的提醒在本轮只投递一次。
type Note struct {
Kind Kind
Key string
Text string
}
// Bus 按最小间隔、每 run 上限和去重决定是否投递。
type Bus struct {
mu sync.Mutex
minGap time.Duration
cap int
sent int
last time.Time
seen map[string]struct{}
}
func New(minGap time.Duration, cap int) *Bus {
if cap <= 0 {
cap = 3
}
return &Bus{minGap: minGap, cap: cap, seen: map[string]struct{}{}}
}
// Admit 在通过节流时返回真。同一 Kind+Key 只通过一次。
func (b *Bus) Admit(n Note, now time.Time) bool {
if b == nil {
return false
}
b.mu.Lock()
defer b.mu.Unlock()
if b.sent >= b.cap {
return false
}
if b.minGap > 0 && !b.last.IsZero() && now.Sub(b.last) < b.minGap {
return false
}
key := string(n.Kind) + "\n" + n.Key
if _, ok := b.seen[key]; ok {
return false
}
b.seen[key] = struct{}{}
b.last = now
b.sent++
return true
}
// Attach 把提醒正文并入最后一条工具结果或用户消息。
// 找不到这两种消息时原样返回,不追加 user 消息。
func Attach(messages []llm.Message, texts []string) []llm.Message {
var lines []string
for _, text := range texts {
text = strings.TrimSpace(text)
if text != "" {
lines = append(lines, text)
}
}
if len(lines) == 0 || len(messages) == 0 {
return messages
}
block := "<system-reminder>\n" + strings.Join(lines, "\n") + "\n</system-reminder>"
out := append([]llm.Message{}, messages...)
for i := len(out) - 1; i >= 0; i-- {
if out[i].Role != llm.RoleTool && out[i].Role != llm.RoleUser {
continue
}
if out[i].Content != "" {
out[i].Content += "\n" + block
} else {
out[i].Content = block
}
return out
}
return messages
}
+53
View File
@@ -0,0 +1,53 @@
package reminders
import (
"strings"
"testing"
"time"
"agent-native/runtime/llm"
)
func TestAttachDoesNotAddUserAfterTool(t *testing.T) {
msgs := []llm.Message{
{Role: llm.RoleUser, Content: "请改文件"},
{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "c1", Name: "read_file"}}},
{Role: llm.RoleTool, Name: "read_file", ToolCallID: "c1", Content: "原文"},
}
out := Attach(msgs, []string{"计划和待办已过期"})
if len(out) != len(msgs) {
t.Fatalf("len %d, want %d", len(out), len(msgs))
}
for i := range msgs {
if out[i].Role != msgs[i].Role {
t.Fatalf("role changed at %d: %s", i, out[i].Role)
}
}
last := out[len(out)-1]
if last.Role != llm.RoleTool || !strings.Contains(last.Content, "原文") || !strings.Contains(last.Content, "system-reminder") {
t.Fatalf("reminder not folded into tool result: %+v", last)
}
if strings.Contains(out[0].Content, "system-reminder") {
t.Fatal("reminder leaked onto the earlier user message")
}
}
func TestAdmitThrottles(t *testing.T) {
bus := New(time.Minute, 2)
now := time.Now()
if !bus.Admit(Note{Kind: KindBudget, Key: "near", Text: "预算将尽"}, now) {
t.Fatal("first admit")
}
if bus.Admit(Note{Kind: KindBudget, Key: "near", Text: "预算将尽"}, now.Add(time.Second)) {
t.Fatal("duplicate key admitted")
}
if bus.Admit(Note{Kind: KindBackground, Key: "job", Text: "后台结束"}, now.Add(time.Second)) {
t.Fatal("inside min gap admitted")
}
if !bus.Admit(Note{Kind: KindBackground, Key: "job", Text: "后台结束"}, now.Add(time.Minute)) {
t.Fatal("after gap")
}
if bus.Admit(Note{Kind: KindDeliverable, Key: "doc", Text: "交付物待确认"}, now.Add(2*time.Minute)) {
t.Fatal("over cap")
}
}
+27 -7
View File
@@ -67,15 +67,35 @@ func (s *Store) ReadMeta() map[string]string {
return meta
}
const maxCompactAttempts = 3
// Compact appends a compaction entry. The file keeps every earlier record.
// Load hides the summarized prefix behind the summary plus the last
// keepLastN non-system messages. Returns false when there is nothing to fold.
// A session change during summary is retried from the new transcript.
// Three consecutive summary failures or conflicts trip the breaker and
// leave the file unchanged.
func (s *Store) Compact(keepLastN int, summarize func(messages []llm.Message) (string, error)) (bool, error) {
var last error
for attempt := 0; attempt < maxCompactAttempts; attempt++ {
done, err, retry := s.compactOnce(keepLastN, summarize)
if !retry {
return done, err
}
last = err
}
if last == nil {
last = fmt.Errorf("compaction conflict: session changed during summary")
}
return false, fmt.Errorf("compaction circuit open: %w", last)
}
func (s *Store) compactOnce(keepLastN int, summarize func(messages []llm.Message) (string, error)) (bool, error, bool) {
s.mu.Lock()
entries, err := s.readEntriesLocked()
if err != nil {
s.mu.Unlock()
return false, err
return false, err, false
}
leaf := s.leafLocked(entries)
generation := len(entries)
@@ -90,7 +110,7 @@ func (s *Store) Compact(keepLastN int, summarize func(messages []llm.Message) (s
}
}
if len(messages) <= keepLastN+1 {
return false, nil
return false, nil, false
}
rest := messages
restEntries := msgEntries
@@ -100,13 +120,13 @@ func (s *Store) Compact(keepLastN int, summarize func(messages []llm.Message) (s
}
tailStart := len(rest) - keepLastN
if tailStart <= 0 {
return false, nil
return false, nil, false
}
tailStart = alignToolTail(rest, tailStart)
head := rest[:tailStart]
summary, err := summarize(head)
if err != nil {
return false, err
return false, err, true
}
var tailIDs []string
for _, e := range restEntries[tailStart:] {
@@ -116,13 +136,13 @@ func (s *Store) Compact(keepLastN int, summarize func(messages []llm.Message) (s
defer s.mu.Unlock()
now, err := s.readEntriesLocked()
if err != nil {
return false, err
return false, err, false
}
if len(now) != generation {
return false, fmt.Errorf("compaction conflict: session changed during summary")
return false, fmt.Errorf("compaction conflict: session changed during summary"), true
}
_, err = s.appendEntryLocked(Entry{Type: TypeCompaction, Summary: summary, TailIDs: tailIDs})
return err == nil, err
return err == nil, err, false
}
// alignToolTail 让保留尾巴从一条非 tool 消息开始,避免把 tool_calls 和结果拆开。
+57
View File
@@ -0,0 +1,57 @@
package session
import (
"encoding/json"
"agent-native/runtime/llm"
)
// ToolIntent 是一次尚未配对结果的工具调用声明。
type ToolIntent struct {
CallID string `json:"callId"`
Name string `json:"name"`
Arguments string `json:"arguments,omitempty"`
Replay string `json:"replay"`
}
// AppendToolIntent 先落下「即将执行」。结果另写一条 tool 消息。
func (s *Store) AppendToolIntent(in ToolIntent) error {
raw, err := json.Marshal(in)
if err != nil {
return err
}
_, err = s.AppendEntry(Entry{Type: TypeToolIntent, Data: raw})
return err
}
// PendingIntents 返回活动分支上还没有对应 tool 结果的意图。
func (s *Store) PendingIntents() ([]ToolIntent, error) {
s.mu.Lock()
defer s.mu.Unlock()
entries, err := s.readEntriesLocked()
if err != nil {
return nil, err
}
answered := map[string]bool{}
var intents []ToolIntent
for _, e := range pathTo(entries, s.leafLocked(entries)) {
if e.Type == TypeMessage && e.Message != nil && e.Message.Role == llm.RoleTool && e.Message.ToolCallID != "" {
answered[e.Message.ToolCallID] = true
}
if e.Type != TypeToolIntent {
continue
}
var in ToolIntent
if err := json.Unmarshal(e.Data, &in); err != nil {
return nil, err
}
intents = append(intents, in)
}
var pending []ToolIntent
for _, in := range intents {
if !answered[in.CallID] {
pending = append(pending, in)
}
}
return pending, nil
}
+70
View File
@@ -0,0 +1,70 @@
package session
import "agent-native/runtime/llm"
const (
// DefaultKeepToolResults 是运行内清理时保留的最近工具结果条数。
DefaultKeepToolResults = 5
// DefaultContextWindow 是未配置模型窗口时的保守字符数。
DefaultContextWindow = 32000
// DefaultOutputReserve 是留给模型输出的字符预留。
DefaultOutputReserve = 4096
// DefaultContextBuffer 是窗口估算的余量,避免贴着上限才清理。
DefaultContextBuffer = 1024
)
// EstimateChars 用字符数近似这次请求的体积。
func EstimateChars(msgs []llm.Message) int {
n := 0
for _, m := range msgs {
n += len([]rune(m.Content))
for _, c := range m.ToolCalls {
n += len([]rune(c.Name)) + len([]rune(c.Arguments))
}
}
return n
}
// NeedsClear 在估算超过「窗口 − 输出预留 − 缓冲」时返回真。
// window ≤ 0 时用保守默认值。
func NeedsClear(estimate, window, outputReserve, buffer int) bool {
if window <= 0 {
window = DefaultContextWindow
}
room := window - outputReserve - buffer
if room < 1 {
room = 1
}
return estimate > room
}
// ClearOldToolResults 把最近 keep 条之前的工具结果换成占位。
// 消息条数、顺序和 tool_call 配对都不变;占位写明完整结果仍在会话记录里。
// keep ≤ 0 时保留默认的 5 条。返回的切片不与输入共享被改写的消息。
func ClearOldToolResults(msgs []llm.Message, keep int) []llm.Message {
if keep <= 0 {
keep = DefaultKeepToolResults
}
out := append([]llm.Message{}, msgs...)
var toolIdx []int
for i, m := range out {
if m.Role == llm.RoleTool {
toolIdx = append(toolIdx, i)
}
}
if len(toolIdx) <= keep {
return out
}
for _, i := range toolIdx[:len(toolIdx)-keep] {
name := out[i].Name
if name == "" {
name = "工具"
}
loc := "会话记录"
if out[i].ToolCallID != "" {
loc = "会话记录中的调用 " + out[i].ToolCallID
}
out[i].Content = "旧工具结果已清理。完整内容见" + loc + "(" + name + ")。"
}
return out
}
+60
View File
@@ -0,0 +1,60 @@
package session
import (
"strings"
"testing"
"agent-native/runtime/llm"
)
func TestClearOldToolResultsKeepsLastFive(t *testing.T) {
var msgs []llm.Message
for i := 0; i < 8; i++ {
id := string(rune('a' + i))
msgs = append(msgs,
llm.Message{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: id, Name: "read_file"}}},
llm.Message{Role: llm.RoleTool, Name: "read_file", ToolCallID: id, Content: "正文-" + id},
)
}
out := ClearOldToolResults(msgs, 5)
if len(out) != len(msgs) {
t.Fatalf("len %d, want %d", len(out), len(msgs))
}
cleared := 0
for i, m := range out {
if m.Role != llm.RoleTool {
if m.Role != msgs[i].Role || len(m.ToolCalls) != len(msgs[i].ToolCalls) {
t.Fatalf("pairing changed at %d", i)
}
continue
}
if m.ToolCallID != msgs[i].ToolCallID {
t.Fatalf("tool id changed at %d", i)
}
if strings.Contains(m.Content, "旧工具结果已清理") {
cleared++
if !strings.Contains(m.Content, m.ToolCallID) {
t.Fatalf("placeholder missing call id: %s", m.Content)
}
} else if m.Content != msgs[i].Content {
t.Fatalf("recent result changed: %s", m.Content)
}
}
if cleared != 3 {
t.Fatalf("cleared %d, want 3", cleared)
}
for _, m := range out[len(out)-10:] {
if m.Role == llm.RoleTool && strings.Contains(m.Content, "旧工具结果已清理") {
t.Fatalf("last five were cleared: %s", m.ToolCallID)
}
}
}
func TestNeedsClearUsesConservativeWindow(t *testing.T) {
if NeedsClear(10, 0, 0, 0) {
t.Fatal("small estimate must fit the default window")
}
if !NeedsClear(DefaultContextWindow, 0, DefaultOutputReserve, DefaultContextBuffer) {
t.Fatal("estimate at the default window must clear")
}
}
+99
View File
@@ -1,6 +1,7 @@
package session
import (
"fmt"
"path/filepath"
"strings"
"testing"
@@ -62,6 +63,104 @@ func TestR2CompactRejectsAppendDuringSummary(t *testing.T) {
}
}
func TestCompactRetriesThenCommits(t *testing.T) {
s, err := Open(filepath.Join(t.TempDir(), "s.jsonl"))
if err != nil {
t.Fatal(err)
}
for _, text := range []string{"a", "b", "c"} {
if err := s.Append(llm.Message{Role: llm.RoleUser, Content: text}); err != nil {
t.Fatal(err)
}
}
raced := false
ok, err := s.Compact(1, func(messages []llm.Message) (string, error) {
if !raced {
raced = true
if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "raced"}); err != nil {
return "", err
}
return "", fmt.Errorf("compaction conflict: session changed during summary")
}
return "summary", nil
})
if err != nil || !ok {
t.Fatalf("retry compact: %v %v", ok, err)
}
loaded, err := s.Load()
if err != nil {
t.Fatal(err)
}
var saw bool
for _, m := range loaded {
if m.Content == "raced" || m.Content == "summary" {
saw = true
}
}
if !saw {
t.Fatalf("retried summary missing: %+v", loaded)
}
}
func TestCompactLeavesNoOrphanToolCall(t *testing.T) {
s, err := Open(filepath.Join(t.TempDir(), "s.jsonl"))
if err != nil {
t.Fatal(err)
}
msgs := []llm.Message{
{Role: llm.RoleUser, Content: "先看"},
{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "old", Name: "read_file"}}},
{Role: llm.RoleTool, ToolCallID: "old", Name: "read_file", Content: "旧"},
{Role: llm.RoleUser, Content: "再改"},
{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "a", Name: "read_file"}, {ID: "b", Name: "grep"}}},
{Role: llm.RoleTool, ToolCallID: "a", Name: "read_file", Content: "甲"},
{Role: llm.RoleTool, ToolCallID: "b", Name: "grep", Content: "乙"},
}
for _, m := range msgs {
if err := s.Append(m); err != nil {
t.Fatal(err)
}
}
ok, err := s.Compact(2, func(messages []llm.Message) (string, error) {
return "摘要", nil
})
if err != nil || !ok {
t.Fatalf("compact: %v %v", ok, err)
}
loaded, err := s.Load()
if err != nil {
t.Fatal(err)
}
assertToolPairs(t, loaded)
}
func assertToolPairs(t *testing.T, msgs []llm.Message) {
t.Helper()
calls := map[string]int{}
for _, m := range msgs {
if m.Role != llm.RoleAssistant {
continue
}
for _, c := range m.ToolCalls {
calls[c.ID]++
}
}
for _, m := range msgs {
if m.Role != llm.RoleTool {
continue
}
if calls[m.ToolCallID] == 0 {
t.Fatalf("orphan tool result %s in %+v", m.ToolCallID, msgs)
}
calls[m.ToolCallID]--
}
for id, n := range calls {
if n != 0 {
t.Fatalf("orphan tool call %s in %+v", id, msgs)
}
}
}
func TestR2AppendVisibleAfterReopen(t *testing.T) {
path := filepath.Join(t.TempDir(), "s.jsonl")
s, err := Open(path)
+1
View File
@@ -26,6 +26,7 @@ const (
TypeLabel EntryType = "label"
TypeCustom EntryType = "custom"
TypeSessionInfo EntryType = "session_info"
TypeToolIntent EntryType = "tool_intent"
)
// Entry is one node in the session tree. Files written before the tree
+7
View File
@@ -83,6 +83,13 @@ func (a *AgentSession) ClearQueue() {
a.emitQueue("clear")
}
// Busy 报告是否有尚未结束的 Prompt。
func (a *AgentSession) Busy() bool {
a.mu.Lock()
defer a.mu.Unlock()
return a.cancel != nil
}
// Abort cancels the in-flight Prompt.
func (a *AgentSession) Abort() {
a.mu.Lock()
+23 -2
View File
@@ -13,6 +13,7 @@ import (
"sort"
"strings"
"time"
"unicode/utf8"
"agent-native/runtime/tools"
)
@@ -26,8 +27,8 @@ type Request struct {
DoneCriteria string `json:"done_criteria,omitempty"`
Deliverables []string `json:"deliverables,omitempty"`
// Deadline / MaxToolCalls 是宿主注入的预算余量(模型不可自行扩权)。
Deadline time.Time `json:"-"`
MaxToolCalls int `json:"-"`
Deadline time.Time `json:"-"`
MaxToolCalls int `json:"-"`
}
// Result 是子代理的回传:至少包含结论文本;文件由宿主按写范围物化。
@@ -81,6 +82,15 @@ func (t *delegateTool) Execute(ctx context.Context, args json.RawMessage) (tools
if err != nil {
return tools.Result{Content: fmt.Sprintf("sub-agent failed: %v", err), IsError: true}, nil
}
shown, full := CapParentText(res.Text)
if full != nil {
if res.Files == nil {
res.Files = map[string][]byte{}
}
res.Files["subagent-result.txt"] = full
shown += "\n完整结果见 subagent-result.txt"
}
res.Text = shown
var b strings.Builder
if res.Success {
b.WriteString("[sub-agent succeeded]\n")
@@ -101,3 +111,14 @@ func (t *delegateTool) Execute(ctx context.Context, args json.RawMessage) (tools
}
return tools.Result{Content: b.String(), IsError: !res.Success}, nil
}
// ParentResultCap 是回传父会话的默认字符上限。
const ParentResultCap = 32000
// CapParentText 超出上限时返回截断文本和全文,由调用方把全文写成工件或文件。
func CapParentText(text string) (shown string, full []byte) {
if utf8.RuneCountInString(text) <= ParentResultCap {
return text, nil
}
return string([]rune(text)[:ParentResultCap]), []byte(text)
}
+23
View File
@@ -84,3 +84,26 @@ func TestDelegateToolFactoryErrorAndCancel(t *testing.T) {
t.Fatalf("pre-canceled delegation must not start, got %+v", res)
}
}
func TestCapParentTextWritesOverflowPath(t *testing.T) {
shown, full := CapParentText("短")
if shown != "短" || full != nil {
t.Fatalf("short: %q %d", shown, len(full))
}
body := strings.Repeat("字", ParentResultCap+3)
shown, full = CapParentText(body)
if full == nil || string(full) != body {
t.Fatal("full text not returned")
}
if len([]rune(shown)) != ParentResultCap {
t.Fatalf("shown runes %d", len([]rune(shown)))
}
reg := tools.NewRegistry(t.TempDir())
Register(reg, func(ctx context.Context, req Request) (Result, error) {
return Result{Success: true, Text: body}, nil
})
res := reg.Execute(context.Background(), "delegate_task", json.RawMessage(`{"goal":"g","done_criteria":"d"}`))
if res.IsError || !strings.Contains(res.Content, "subagent-result.txt") {
t.Fatalf("capped result: %+v", res)
}
}
+112
View File
@@ -0,0 +1,112 @@
// Package telemetry 记录 GenAI span。默认不采集提示词正文。
package telemetry
import (
"bytes"
"crypto/rand"
"encoding/hex"
"encoding/json"
"net/http"
"os"
"strings"
"time"
"logging-system/logsys"
)
// Span 是一次 invoke_agent、chat 或 execute_tool。
type Span struct {
Name string
TraceID string
SpanID string
Start time.Time
End time.Time
Attrs map[string]string
}
// StartFunc 返回循环使用的追踪函数。设置了 OTEL_EXPORTER_OTLP_ENDPOINT 时,
// 结束 span 后以 OTLP/HTTP JSON 上报,不引入新的依赖。
func StartFunc(log *logsys.Logger) func(name string, attrs map[string]string) func() {
endpoint := strings.TrimRight(os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT"), "/")
return func(name string, attrs map[string]string) func() {
sp := Span{
Name: name, TraceID: newID(16), SpanID: newID(8),
Start: time.Now(), Attrs: publicAttrs(attrs),
}
return func() {
sp.End = time.Now()
if log != nil {
kvs := []any{"span", sp.Name, "trace_id", sp.TraceID, "span_id", sp.SpanID}
for k, v := range sp.Attrs {
kvs = append(kvs, k, v)
}
log.Info("genai span", kvs...)
}
if endpoint != "" {
_ = ExportOTLP(endpoint, sp)
}
}
}
}
// publicAttrs 丢掉提示词正文。调用方即使传入也不会进入 span。
func publicAttrs(in map[string]string) map[string]string {
out := map[string]string{}
for k, v := range in {
if k == "prompt" || k == "messages" || k == "content" {
continue
}
out[k] = v
}
return out
}
// ExportOTLP 把一个 span POST 到端点的 /v1/traces。
func ExportOTLP(endpoint string, sp Span) error {
endpoint = strings.TrimRight(endpoint, "/")
if !strings.HasSuffix(endpoint, "/v1/traces") {
endpoint += "/v1/traces"
}
body, err := json.Marshal(otlpBody(sp))
if err != nil {
return err
}
req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
return nil
}
func otlpBody(sp Span) map[string]any {
var attrs []map[string]any
for k, v := range publicAttrs(sp.Attrs) {
attrs = append(attrs, map[string]any{"key": k, "value": map[string]any{"stringValue": v}})
}
return map[string]any{
"resourceSpans": []any{map[string]any{
"scopeSpans": []any{map[string]any{
"spans": []any{map[string]any{
"traceId": sp.TraceID,
"spanId": sp.SpanID,
"name": sp.Name,
"startTimeUnixNano": sp.Start.UnixNano(),
"endTimeUnixNano": sp.End.UnixNano(),
"attributes": attrs,
}},
}},
}},
}
}
func newID(n int) string {
b := make([]byte, n)
_, _ = rand.Read(b)
return hex.EncodeToString(b)
}
+42
View File
@@ -0,0 +1,42 @@
package telemetry
import (
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
)
func TestSpanOmitsPromptAndUsesAgreedName(t *testing.T) {
var body string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/traces" {
t.Errorf("path %s", r.URL.Path)
}
b, _ := io.ReadAll(r.Body)
body = string(b)
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
sp := Span{
Name: "chat", TraceID: "abc", SpanID: "def",
Start: time.Unix(0, 0), End: time.Unix(1, 0),
Attrs: map[string]string{"model": "m", "prompt": "用户的秘密提示词"},
}
if err := ExportOTLP(srv.URL, sp); err != nil {
t.Fatal(err)
}
if !strings.Contains(body, `"name":"chat"`) {
t.Fatalf("span name missing: %s", body)
}
if strings.Contains(body, "用户的秘密提示词") || strings.Contains(body, "prompt") {
t.Fatalf("prompt leaked: %s", body)
}
got := publicAttrs(map[string]string{"tool": "read_file", "content": "全文"})
if _, ok := got["content"]; ok || got["tool"] != "read_file" {
t.Fatalf("%#v", got)
}
}
+56 -7
View File
@@ -8,11 +8,17 @@ import (
"path/filepath"
"regexp"
"strings"
"agent-native/runtime/filestate"
)
type readFile struct{ root string }
type readFile struct {
root string
files *filestate.Tracker
}
func (t *readFile) Name() string { return "read_file" }
func (t *readFile) Name() string { return "read_file" }
func (t *readFile) Replay() string { return ReplaySafe }
func (t *readFile) Description() string {
return "Read a UTF-8 text file inside the workspace. Without offset/limit returns the whole file (truncated at 200 KiB); with offset/limit returns numbered lines from a 1-based offset."
}
@@ -46,6 +52,7 @@ func (t *readFile) Execute(_ context.Context, args json.RawMessage) (Result, err
if truncated {
content += "\n[truncated at 200 KiB; use offset/limit to page through the file]"
}
_ = t.files.NoteRead(path)
return Result{Content: content}, nil
}
if obj.Offset <= 0 {
@@ -66,10 +73,15 @@ func (t *readFile) Execute(_ context.Context, args json.RawMessage) (Result, err
for i := obj.Offset - 1; i < end; i++ {
out = append(out, fmt.Sprintf("%d\t%s", i+1, lines[i]))
}
_ = t.files.NoteRead(path)
return Result{Content: strings.Join(out, "\n")}, nil
}
type writeFile struct{ root string }
type writeFile struct {
root string
files *filestate.Tracker
mark func(rel string, before []byte, existed bool, after []byte)
}
func (t *writeFile) Name() string { return "write_file" }
func (t *writeFile) Description() string {
@@ -90,16 +102,27 @@ func (t *writeFile) Execute(_ context.Context, args json.RawMessage) (Result, er
if err != nil {
return Result{Content: err.Error(), IsError: true}, nil
}
if msg := rejectStaleWrite(t.files, path); msg != "" {
return Result{Content: msg, IsError: true}, nil
}
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
return Result{Content: err.Error(), IsError: true}, nil
}
before, existed := fileBefore(path)
if err := atomicWriteFile(path, []byte(obj.Content), 0o644); err != nil {
return Result{Content: err.Error(), IsError: true}, nil
}
if t.mark != nil {
t.mark(obj.Path, before, existed, []byte(obj.Content))
}
return Result{Content: "wrote " + obj.Path}, nil
}
type editFile struct{ root string }
type editFile struct {
root string
files *filestate.Tracker
mark func(rel string, before []byte, existed bool, after []byte)
}
func (t *editFile) Name() string { return "edit_file" }
func (t *editFile) Description() string {
@@ -124,6 +147,9 @@ func (t *editFile) Execute(_ context.Context, args json.RawMessage) (Result, err
if err != nil {
return Result{Content: err.Error(), IsError: true}, nil
}
if msg := rejectStaleWrite(t.files, path); msg != "" {
return Result{Content: msg, IsError: true}, nil
}
b, err := os.ReadFile(path)
if err != nil {
return Result{Content: err.Error(), IsError: true}, nil
@@ -137,12 +163,16 @@ func (t *editFile) Execute(_ context.Context, args json.RawMessage) (Result, err
if err := atomicWriteFile(path, []byte(updated), 0o644); err != nil {
return Result{Content: err.Error(), IsError: true}, nil
}
if t.mark != nil {
t.mark(obj.Path, b, true, []byte(updated))
}
return Result{Content: "edited " + obj.Path}, nil
}
type listDir struct{ root string }
func (t *listDir) Name() string { return "list_dir" }
func (t *listDir) Name() string { return "list_dir" }
func (t *listDir) Replay() string { return ReplaySafe }
func (t *listDir) Description() string {
return "List names in a workspace directory."
}
@@ -175,7 +205,8 @@ func (t *listDir) Execute(_ context.Context, args json.RawMessage) (Result, erro
type grepTool struct{ root string }
func (t *grepTool) Name() string { return "grep" }
func (t *grepTool) Name() string { return "grep" }
func (t *grepTool) Replay() string { return ReplaySafe }
func (t *grepTool) Description() string {
return "Search workspace text files. Literal substring by default; regex=true switches to Go regexp, include filters file names (e.g. *.go), ignore_case folds case. Caps at 50 hits."
}
@@ -262,7 +293,11 @@ func (t *grepTool) Execute(_ context.Context, args json.RawMessage) (Result, err
}
// multiEditTool 一次多处替换:全部命中且互不冲突才落盘,任一失败即整体不动。
type multiEditTool struct{ root string }
type multiEditTool struct {
root string
files *filestate.Tracker
mark func(rel string, before []byte, existed bool, after []byte)
}
func (t *multiEditTool) Name() string { return "multi_edit" }
func (t *multiEditTool) Description() string {
@@ -292,6 +327,9 @@ func (t *multiEditTool) Execute(_ context.Context, args json.RawMessage) (Result
if err != nil {
return Result{Content: err.Error(), IsError: true}, nil
}
if msg := rejectStaleWrite(t.files, path); msg != "" {
return Result{Content: msg, IsError: true}, nil
}
b, err := os.ReadFile(path)
if err != nil {
return Result{Content: err.Error(), IsError: true}, nil
@@ -310,5 +348,16 @@ func (t *multiEditTool) Execute(_ context.Context, args json.RawMessage) (Result
if err := atomicWriteFile(path, []byte(text), 0o644); err != nil {
return Result{Content: err.Error(), IsError: true}, nil
}
if t.mark != nil {
t.mark(obj.Path, b, true, []byte(text))
}
return Result{Content: fmt.Sprintf("applied %d edits to %s", len(obj.Edits), obj.Path)}, nil
}
func fileBefore(path string) ([]byte, bool) {
b, err := os.ReadFile(path)
if err != nil {
return nil, false
}
return b, true
}
+42
View File
@@ -0,0 +1,42 @@
package tools
import (
"context"
"os"
"path/filepath"
"testing"
"agent-native/runtime/checkpoint"
)
func TestWriteAndEditRecordCheckpoints(t *testing.T) {
root := t.TempDir()
log, err := checkpoint.Open(filepath.Join(t.TempDir(), "cp.jsonl"))
if err != nil {
t.Fatal(err)
}
reg := WorkspaceToolsCheckpointed(root, nil, log, func() string { return "m1" })
if reg.Execute(context.Background(), "write_file", []byte(`{"path":"a.txt","content":"one"}`)).IsError {
t.Fatal("write")
}
if reg.Execute(context.Background(), "edit_file", []byte(`{"path":"a.txt","old":"one","new":"two"}`)).IsError {
t.Fatal("edit")
}
if err := os.WriteFile(filepath.Join(root, "shell.txt"), []byte("nope"), 0o644); err != nil {
t.Fatal(err)
}
got, err := log.Rewind("m1", []string{"m1"})
if err != nil {
t.Fatal(err)
}
if len(got.Restored) != 1 || got.Restored[0] != "a.txt" {
t.Fatalf("restored %#v skipped %#v", got.Restored, got.Skipped)
}
if _, err := os.Stat(filepath.Join(root, "a.txt")); !os.IsNotExist(err) {
t.Fatal("created file was kept")
}
side, err := os.ReadFile(filepath.Join(root, "shell.txt"))
if err != nil || string(side) != "nope" {
t.Fatalf("shell file %q %v", side, err)
}
}
+59
View File
@@ -0,0 +1,59 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
"agent-native/runtime/codeorch"
)
// RegisterCodeTool 注册 run_tool_script。调用方在开关打开时才注册。
func RegisterCodeTool(reg *Registry) {
if reg == nil {
return
}
reg.Register(&codeScriptTool{reg: reg})
}
type codeScriptTool struct{ reg *Registry }
func (t *codeScriptTool) Name() string { return "run_tool_script" }
func (t *codeScriptTool) Description() string {
return "用一段 Python 依次调用已有工具。禁止再调用解释器、本工具、派发和工作流。失败时会列出已经执行的调用"
}
func (t *codeScriptTool) Parameters() json.RawMessage {
return json.RawMessage(`{"type":"object","properties":{"script":{"type":"string"},"timeout_sec":{"type":"integer"}},"required":["script"]}`)
}
func (t *codeScriptTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) {
var req struct {
Script string `json:"script"`
Timeout int `json:"timeout_sec"`
}
if err := json.Unmarshal(args, &req); err != nil {
return Result{Content: err.Error(), IsError: true}, nil
}
if strings.TrimSpace(req.Script) == "" {
return Result{Content: "script 不能为空", IsError: true}, nil
}
lim := codeorch.Limits{FullDir: t.reg.Root()}
if req.Timeout > 0 {
lim.Timeout = time.Duration(req.Timeout) * time.Second
}
out := codeorch.Run(ctx, t.reg.Root(), req.Script, lim, nil, func(ctx context.Context, call codeorch.Call) (string, error) {
if call.Op != "tool.call" {
return "", fmt.Errorf("不能调用 %s", call.Op)
}
if codeorch.Forbidden(call.Name) {
return "", fmt.Errorf("不能调用 %s", call.Name)
}
res := t.reg.Execute(ctx, call.Name, call.Args)
if res.IsError {
return "", fmt.Errorf("%s", res.Content)
}
return res.Content, nil
})
return Result{Content: out.ModelText(), IsError: out.Err != nil}, nil
}
+164
View File
@@ -0,0 +1,164 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"os"
"strings"
"agent-native/runtime/codeorch"
)
type searchTool struct{ reg *Registry }
func (t *searchTool) Name() string { return "tool_search" }
func (t *searchTool) Description() string {
return "按名称或描述查找延迟暴露的工具,并以文本返回它们的 schema"
}
func (t *searchTool) Parameters() json.RawMessage {
return json.RawMessage(`{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}`)
}
func (t *searchTool) Execute(_ context.Context, args json.RawMessage) (Result, error) {
var q struct {
Query string `json:"query"`
}
if len(args) > 0 {
_ = json.Unmarshal(args, &q)
}
return Result{Content: t.reg.Search(q.Query)}, nil
}
func (t *searchTool) Replay() string { return ReplaySafe }
type callTool struct{ reg *Registry }
func (t *callTool) Name() string { return "call_tool" }
func (t *callTool) Description() string {
return "调用一个延迟暴露的工具。参数 schema 固定,不随工具清单变化"
}
func (t *callTool) Parameters() json.RawMessage {
return json.RawMessage(`{"type":"object","properties":{"name":{"type":"string"},"arguments":{"type":"object"}},"required":["name"]}`)
}
func (t *callTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) {
var req struct {
Name string `json:"name"`
Arguments json.RawMessage `json:"arguments"`
}
if err := json.Unmarshal(args, &req); err != nil {
return Result{Content: err.Error(), IsError: true}, nil
}
if req.Name == "" || req.Name == "call_tool" || req.Name == "tool_search" || codeorch.Forbidden(req.Name) {
return Result{Content: "不能调用 " + req.Name, IsError: true}, nil
}
res := t.reg.Execute(ctx, req.Name, req.Arguments)
if !res.IsError {
t.reg.markActivated(req.Name)
}
return res, nil
}
func (r *Registry) ensureMetaLocked() {
if _, ok := r.tools["tool_search"]; ok {
return
}
r.tools["tool_search"] = &searchTool{reg: r}
r.order = append(r.order, "tool_search")
r.tools["call_tool"] = &callTool{reg: r}
r.order = append(r.order, "call_tool")
}
// Search 把名称或描述命中的延迟工具 schema 写成文本。
func (r *Registry) Search(query string) string {
r.mu.Lock()
defer r.mu.Unlock()
q := strings.ToLower(strings.TrimSpace(query))
var b strings.Builder
for _, name := range r.order {
t := r.tools[name]
if ExposureOf(t) != ExposureDeferred {
continue
}
blob := strings.ToLower(t.Name() + " " + t.Description())
if q != "" && !strings.Contains(blob, q) {
continue
}
fmt.Fprintf(&b, "name: %s\ndescription: %s\nparameters: %s\n\n", t.Name(), t.Description(), t.Parameters())
}
if b.Len() == 0 {
return "没有匹配的延迟工具"
}
return strings.TrimSpace(b.String())
}
// SetActivationFile 指定已激活工具的持久文件。
func (r *Registry) SetActivationFile(path string) {
r.mu.Lock()
r.activatePath = path
r.mu.Unlock()
}
// RestoreActivated 读回上次激活的延迟工具名。清单仍然不包含它们。
func (r *Registry) RestoreActivated() error {
r.mu.Lock()
path := r.activatePath
r.mu.Unlock()
if path == "" {
return nil
}
raw, err := os.ReadFile(path)
if os.IsNotExist(err) {
return nil
}
if err != nil {
return err
}
var names []string
if err := json.Unmarshal(raw, &names); err != nil {
return err
}
r.mu.Lock()
defer r.mu.Unlock()
if r.activated == nil {
r.activated = map[string]bool{}
}
for _, name := range names {
r.activated[name] = true
}
return nil
}
// Activated 返回已恢复或本次调用过的延迟工具名。
func (r *Registry) Activated() []string {
r.mu.Lock()
defer r.mu.Unlock()
var out []string
for name := range r.activated {
out = append(out, name)
}
return out
}
func (r *Registry) markActivated(name string) {
r.mu.Lock()
defer r.mu.Unlock()
t, ok := r.tools[name]
if !ok || ExposureOf(t) != ExposureDeferred {
return
}
if r.activated == nil {
r.activated = map[string]bool{}
}
r.activated[name] = true
if r.activatePath == "" {
return
}
var names []string
for n := range r.activated {
names = append(names, n)
}
raw, err := json.Marshal(names)
if err != nil {
return
}
_ = os.WriteFile(r.activatePath, raw, 0o644)
}
+75
View File
@@ -0,0 +1,75 @@
package tools
import (
"context"
"encoding/json"
"os"
"path/filepath"
"strings"
"testing"
)
type deferredEcho struct{}
func (deferredEcho) Name() string { return "mcp__late__echo" }
func (deferredEcho) Description() string { return "回声延迟工具" }
func (deferredEcho) Parameters() json.RawMessage {
return json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}}}`)
}
func (deferredEcho) Exposure() string { return ExposureDeferred }
func (deferredEcho) Execute(_ context.Context, args json.RawMessage) (Result, error) {
var a struct {
Text string `json:"text"`
}
_ = json.Unmarshal(args, &a)
return Result{Content: "echo " + a.Text}, nil
}
func TestDeferredToolStaysOutOfDeclsUntilCalledByName(t *testing.T) {
reg := NewRegistry(t.TempDir())
reg.Register(deferredEcho{})
for _, d := range reg.Decls() {
if d.Name == "mcp__late__echo" {
t.Fatal("deferred tool must not be listed")
}
}
found := reg.Search("回声")
if !strings.Contains(found, "mcp__late__echo") || !strings.Contains(found, `"text"`) {
t.Fatalf("search %q", found)
}
res := reg.Execute(context.Background(), "call_tool", json.RawMessage(`{"name":"mcp__late__echo","arguments":{"text":"hi"}}`))
if res.IsError || res.Content != "echo hi" {
t.Fatalf("call %+v", res)
}
denied := reg.Execute(context.Background(), "call_tool", json.RawMessage(`{"name":"run_shell"}`))
if !denied.IsError {
t.Fatalf("forbidden call should fail: %+v", denied)
}
}
func TestActivatedNamesRestoreWithoutEnteringDecls(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "activated.json")
reg := NewRegistry(dir)
reg.SetActivationFile(path)
reg.Register(deferredEcho{})
_ = reg.Execute(context.Background(), "call_tool", json.RawMessage(`{"name":"mcp__late__echo","arguments":{"text":"x"}}`))
if _, err := os.Stat(path); err != nil {
t.Fatal(err)
}
next := NewRegistry(dir)
next.SetActivationFile(path)
if err := next.RestoreActivated(); err != nil {
t.Fatal(err)
}
next.Register(deferredEcho{})
names := next.Activated()
if len(names) != 1 || names[0] != "mcp__late__echo" {
t.Fatalf("activated %+v", names)
}
for _, d := range next.Decls() {
if d.Name == "mcp__late__echo" {
t.Fatal("restored tool must stay out of the list")
}
}
}
+38
View File
@@ -0,0 +1,38 @@
package tools
import (
"os"
"path/filepath"
"strings"
"testing"
"agent-native/runtime/filestate"
)
func TestTrackedWriteRejectsUnreadAndStale(t *testing.T) {
root := t.TempDir()
path := filepath.Join(root, "note.txt")
if err := os.WriteFile(path, []byte("v1"), 0o644); err != nil {
t.Fatal(err)
}
reg := WorkspaceToolsTracked(root, filestate.New())
res := runTool(t, reg, "write_file", `{"path":"note.txt","content":"v2"}`)
if !res.IsError || !strings.Contains(res.Content, "请先读取") {
t.Fatalf("unread: %+v", res)
}
res = runTool(t, reg, "read_file", `{"path":"note.txt"}`)
if res.IsError {
t.Fatal(res.Content)
}
if err := os.WriteFile(path, []byte("v1-external"), 0o644); err != nil {
t.Fatal(err)
}
res = runTool(t, reg, "edit_file", `{"path":"note.txt","old":"v1","new":"v2"}`)
if !res.IsError || !strings.Contains(res.Content, "请重新读取") {
t.Fatalf("stale: %+v", res)
}
res = runTool(t, reg, "write_file", `{"path":"fresh.txt","content":"ok"}`)
if res.IsError {
t.Fatalf("create: %+v", res)
}
}
+2 -1
View File
@@ -15,7 +15,8 @@ const maxGlobResults = 500
type globTool struct{ root string }
func (t *globTool) Name() string { return "glob" }
func (t *globTool) Name() string { return "glob" }
func (t *globTool) Replay() string { return ReplaySafe }
func (t *globTool) Description() string {
return "Find workspace files matching a glob pattern (** crosses directories, * within one segment). Returns up to 500 relative paths."
+139 -16
View File
@@ -8,6 +8,9 @@ import (
"strings"
"sync"
"time"
"agent-native/runtime/checkpoint"
"agent-native/runtime/filestate"
)
// Result is the tool payload returned to the model.
@@ -24,12 +27,51 @@ type Tool interface {
Execute(ctx context.Context, args json.RawMessage) (Result, error)
}
const (
// ReplaySafe 表示恢复时可以再执行一次。ReplayNever 是缺省,不重跑。
ReplaySafe = "safe"
ReplayNever = "never"
)
// Replayer 是可选声明。未实现或返回其他值都按 never。
type Replayer interface {
Replay() string
}
// ReplayOf 读取工具的重放声明。
func ReplayOf(t Tool) string {
r, ok := t.(Replayer)
if !ok || r.Replay() != ReplaySafe {
return ReplayNever
}
return ReplaySafe
}
// BeforeHook can block a call after schema validation.
type BeforeHook func(ctx context.Context, name string, args json.RawMessage) (block bool, reason string)
// AfterHook can rewrite a tool result before it is returned.
type AfterHook func(ctx context.Context, name string, args json.RawMessage, result Result) Result
const (
// ExposureDirect 进入模型看到的工具清单。ExposureDeferred 只通过 tool_search 与 call_tool。
ExposureDirect = "direct"
ExposureDeferred = "deferred"
)
// Exposer 声明工具是否直接出现在清单里。未实现视为直接。
type Exposer interface {
Exposure() string
}
// ExposureOf 读取暴露级别。
func ExposureOf(t Tool) string {
if e, ok := t.(Exposer); ok && e.Exposure() == ExposureDeferred {
return ExposureDeferred
}
return ExposureDirect
}
// Registry executes jailed tools and appends an audit line per call.
type Registry struct {
root string
@@ -40,6 +82,13 @@ type Registry struct {
Audit io.Writer
// hookMu 串行 Before、After 与 Audit,避免并行工具调用交错写审计或重入钩子。
hookMu sync.Mutex
mu sync.Mutex
// restricted 为真时,后注册的工具也必须命中 allowExact / allowPrefix。
restricted bool
allowExact map[string]bool
allowPrefix []string
activated map[string]bool
activatePath string
}
func NewRegistry(root string) *Registry {
@@ -49,21 +98,50 @@ func NewRegistry(root string) *Registry {
func (r *Registry) Root() string { return r.root }
func (r *Registry) Has(name string) bool {
r.mu.Lock()
defer r.mu.Unlock()
_, ok := r.tools[name]
return ok
}
// Replay 返回已注册工具的重放声明。未知工具是 never。
func (r *Registry) Replay(name string) string {
if r == nil {
return ReplayNever
}
r.mu.Lock()
t, ok := r.tools[name]
r.mu.Unlock()
if !ok {
return ReplayNever
}
return ReplayOf(t)
}
func (r *Registry) Register(t Tool) {
r.mu.Lock()
defer r.mu.Unlock()
if r.restricted && !r.keepLocked(t.Name()) {
return
}
if _, ok := r.tools[t.Name()]; !ok {
r.order = append(r.order, t.Name())
}
r.tools[t.Name()] = t
if ExposureOf(t) == ExposureDeferred {
r.ensureMetaLocked()
}
}
func (r *Registry) Decls() []Decl {
r.mu.Lock()
defer r.mu.Unlock()
out := make([]Decl, 0, len(r.order))
for _, name := range r.order {
t := r.tools[name]
if ExposureOf(t) == ExposureDeferred {
continue
}
out = append(out, Decl{Name: t.Name(), Description: t.Description(), Parameters: t.Parameters()})
}
return out
@@ -80,6 +158,8 @@ func (r *Registry) Restrict(allow []string) error {
if len(allow) == 0 {
return nil
}
r.mu.Lock()
defer r.mu.Unlock()
exact := make(map[string]bool, len(allow))
prefixes := make([]string, 0)
for _, entry := range allow {
@@ -106,20 +186,12 @@ func (r *Registry) Restrict(allow []string) error {
if len(unknown) > 0 {
return fmt.Errorf("tool allowlist names not registered: %s", strings.Join(unknown, ", "))
}
keep := func(name string) bool {
if exact[name] {
return true
}
for _, p := range prefixes {
if strings.HasPrefix(name, p) {
return true
}
}
return false
}
r.restricted = true
r.allowExact = exact
r.allowPrefix = prefixes
var order []string
for _, name := range r.order {
if keep(name) {
if r.keepLocked(name) {
order = append(order, name)
} else {
delete(r.tools, name)
@@ -129,6 +201,21 @@ func (r *Registry) Restrict(allow []string) error {
return nil
}
func (r *Registry) keepLocked(name string) bool {
if name == "tool_search" || name == "call_tool" {
return true
}
if r.allowExact[name] {
return true
}
for _, p := range r.allowPrefix {
if strings.HasPrefix(name, p) {
return true
}
}
return false
}
func (r *Registry) hasPrefix(prefix string) bool {
for _, name := range r.order {
if strings.HasPrefix(name, prefix) {
@@ -147,7 +234,9 @@ type Decl struct {
// Execute validates arguments, runs hooks, and audits the outcome.
func (r *Registry) Execute(ctx context.Context, name string, args json.RawMessage) Result {
r.mu.Lock()
t, ok := r.tools[name]
r.mu.Unlock()
if !ok {
res := Result{Content: "unknown tool " + name, IsError: true}
r.audit(name, false, true, "unknown")
@@ -269,17 +358,51 @@ func matchesType(kind string, v any) bool {
// WorkspaceTools registers the document workspace tool set.
func WorkspaceTools(root string) *Registry {
return WorkspaceToolsTracked(root, nil)
}
// WorkspaceToolsTracked 与 WorkspaceTools 相同,但改写已存在的文件前要先在本轮读过。
func WorkspaceToolsTracked(root string, files *filestate.Tracker) *Registry {
return WorkspaceToolsCheckpointed(root, files, nil, nil)
}
// WorkspaceToolsCheckpointed 在写、改文件前把改前内容交给检查点。log 为空则不记录。
func WorkspaceToolsCheckpointed(root string, files *filestate.Tracker, log *checkpoint.Log, anchor func() string) *Registry {
mark := writeMark(log, root, anchor)
r := NewRegistry(root)
r.Register(&readFile{root: root})
r.Register(&writeFile{root: root})
r.Register(&editFile{root: root})
r.Register(&multiEditTool{root: root})
r.Register(&readFile{root: root, files: files})
r.Register(&writeFile{root: root, files: files, mark: mark})
r.Register(&editFile{root: root, files: files, mark: mark})
r.Register(&multiEditTool{root: root, files: files, mark: mark})
r.Register(&listDir{root: root})
r.Register(&grepTool{root: root})
r.Register(&globTool{root: root})
return r
}
func writeMark(log *checkpoint.Log, root string, anchor func() string) func(rel string, before []byte, existed bool, after []byte) {
if log == nil {
return nil
}
return func(rel string, before []byte, existed bool, after []byte) {
id := ""
if anchor != nil {
id = anchor()
}
_ = log.Record(id, root, rel, before, existed, after)
}
}
func rejectStaleWrite(files *filestate.Tracker, path string) string {
if files == nil {
return ""
}
if err := files.CheckWrite(path); err != nil {
return err.Error()
}
return ""
}
func mustPath(root, rel string) (string, error) {
return Resolve(root, rel)
}
+66
View File
@@ -0,0 +1,66 @@
package tools
import "strings"
// ApplySelection 按 +name / -name 调整默认工具集。
// 规格为空时返回 base 的副本。出现不带正负号的名字时,先用这些名字替换默认集,再应用增减。
// 分隔符为空格、逗号或换行。
func ApplySelection(base []string, spec string) []string {
spec = strings.TrimSpace(spec)
if spec == "" {
return append([]string{}, base...)
}
fields := strings.FieldsFunc(spec, func(r rune) bool {
return r == ',' || r == ' ' || r == '\n' || r == '\t'
})
hasBare := false
for _, f := range fields {
if f != "" && f[0] != '+' && f[0] != '-' {
hasBare = true
break
}
}
var out []string
if !hasBare {
out = append(out, base...)
}
for _, f := range fields {
if f == "" {
continue
}
switch f[0] {
case '+':
out = addName(out, f[1:])
case '-':
out = dropName(out, f[1:])
default:
out = addName(out, f)
}
}
return out
}
func addName(names []string, name string) []string {
if name == "" {
return names
}
for _, n := range names {
if n == name {
return names
}
}
return append(names, name)
}
func dropName(names []string, name string) []string {
if name == "" {
return names
}
out := names[:0]
for _, n := range names {
if n != name {
out = append(out, n)
}
}
return out
}
+24
View File
@@ -0,0 +1,24 @@
package tools
import "testing"
func TestApplySelectionDeltas(t *testing.T) {
base := []string{"read_file", "write_file", "grep"}
got := ApplySelection(base, "-write_file,+edit_file")
want := []string{"read_file", "grep", "edit_file"}
if len(got) != len(want) {
t.Fatalf("%v", got)
}
for i := range want {
if got[i] != want[i] {
t.Fatalf("got %v want %v", got, want)
}
}
replaced := ApplySelection(base, "read_file, grep")
if len(replaced) != 2 || replaced[0] != "read_file" || replaced[1] != "grep" {
t.Fatalf("bare list: %v", replaced)
}
if same := ApplySelection(base, ""); len(same) != 3 {
t.Fatalf("empty spec: %v", same)
}
}
+7
View File
@@ -0,0 +1,7 @@
{
"work_dir": "/workspace/demo",
"prompt": "写一份说明",
"instructions": "只写目标文件",
"targets": ["output.md"],
"tool_profile": "workspace"
}
+5
View File
@@ -0,0 +1,5 @@
{
"success": true,
"patches": {"output.md": "ZGVtbw=="},
"logs": "done"
}
+55
View File
@@ -0,0 +1,55 @@
package schema
import (
"bytes"
"encoding/json"
"os"
"path/filepath"
"testing"
"agent-native/exec"
)
func TestRequestAndResultRoundTrip(t *testing.T) {
reqRaw, err := os.ReadFile(filepath.Join("request.sample.json"))
if err != nil {
t.Fatal(err)
}
var req exec.Request
if err := json.Unmarshal(reqRaw, &req); err != nil {
t.Fatal(err)
}
if req.Prompt != "写一份说明" || req.ToolProfile != exec.ToolProfileWorkspace {
t.Fatalf("%+v", req)
}
back, err := json.Marshal(req)
if err != nil {
t.Fatal(err)
}
var again exec.Request
if err := json.Unmarshal(back, &again); err != nil {
t.Fatal(err)
}
if again.Prompt != req.Prompt || len(again.Targets) != 1 {
t.Fatalf("%+v", again)
}
resRaw, err := os.ReadFile("result.sample.json")
if err != nil {
t.Fatal(err)
}
var res exec.Result
if err := json.Unmarshal(resRaw, &res); err != nil {
t.Fatal(err)
}
if !res.Success || string(res.Patches["output.md"]) != "demo" {
t.Fatalf("%+v", res)
}
out, err := json.Marshal(res)
if err != nil {
t.Fatal(err)
}
if !bytes.Contains(out, []byte(`"success":true`)) {
t.Fatalf("%s", out)
}
}