Compare commits
29
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
33573d0d75 | ||
|
|
4a10818535 | ||
|
|
d9b0f8f178 | ||
|
|
855fe41152 | ||
|
|
11d1a76bdd | ||
|
|
17f2a4ecfa | ||
|
|
b1f89bf969 | ||
|
|
56848f3cbd | ||
|
|
91c1556184 | ||
|
|
0773a25ab9 | ||
|
|
694b08cdbe | ||
|
|
7e0374ff40 | ||
|
|
e047d90985 | ||
|
|
c6d867cde7 | ||
|
|
1a999c344d | ||
|
|
a34689c78f | ||
|
|
e2815bbf94 | ||
|
|
35f65a3c49 | ||
|
|
886e86c6a8 | ||
|
|
6b7d2432b3 | ||
|
|
23b5e3f336 | ||
|
|
dfd0d23bf3 | ||
|
|
2dedc0d20c | ||
|
|
e7ab48e8ea | ||
|
|
6bc1b08540 | ||
|
|
ec731addef | ||
|
|
cfc840c11d | ||
|
|
951c141f95 | ||
|
|
f9ab2ef698 |
+38
-5
@@ -6,15 +6,48 @@ on:
|
||||
branches: [main]
|
||||
jobs:
|
||||
vet-and-test:
|
||||
runs-on: nightly
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
GOPROXY: https://goproxy.cn,direct
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- run: go vet ./... && go test -count=1 ./...
|
||||
- name: Checkout
|
||||
env:
|
||||
TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
SUBMODULE_TOKEN: ${{ secrets.SUBMODULE_TOKEN }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
git init -q
|
||||
git remote add origin "${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}.git"
|
||||
git -c http.extraheader="AUTHORIZATION: token ${TOKEN}" fetch --depth=1 origin "${GITHUB_SHA}"
|
||||
git checkout -q --detach FETCH_HEAD
|
||||
if [ -z "${SUBMODULE_TOKEN}" ]; then
|
||||
echo "SUBMODULE_TOKEN is not configured" >&2
|
||||
exit 1
|
||||
fi
|
||||
git config --global url."https://weichen:${SUBMODULE_TOKEN}@git-home.u1680832.nyat.app:30319/".insteadOf "ssh://git@git-ssh.u1680832.nyat.app:31411/"
|
||||
git -c fetch.recurseSubmodules=false submodule update --init --recursive
|
||||
- name: Install Go 1.27.1
|
||||
run: |
|
||||
set -euo pipefail
|
||||
curl -fsSL -o "$RUNNER_TEMP/go.tgz" https://mirrors.aliyun.com/golang/go1.27.1.linux-amd64.tar.gz
|
||||
tar -C /usr/local -xzf "$RUNNER_TEMP/go.tgz"
|
||||
ln -sfn /usr/local/go/bin/go /usr/local/bin/go
|
||||
ln -sfn /usr/local/go/bin/gofmt /usr/local/bin/gofmt
|
||||
- name: Vet, test, and race
|
||||
run: go vet ./... && go test -count=1 ./... && go test -count=1 -race ./...
|
||||
docker-checkout-hostgateway:
|
||||
runs-on: ubuntu-latest
|
||||
options: --add-host gitea:host-gateway
|
||||
env:
|
||||
GITHUB_SERVER_URL: http://gitea:3000
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- run: test -f go.mod && echo "docker-mode checkout via host-gateway OK"
|
||||
- name: Checkout without GitHub actions
|
||||
env:
|
||||
TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
git init -q
|
||||
git remote add origin "${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}.git"
|
||||
git -c http.extraheader="AUTHORIZATION: token ${TOKEN}" fetch --depth=1 origin "${GITHUB_SHA}"
|
||||
git checkout -q --detach FETCH_HEAD
|
||||
test -f go.mod && echo "docker-mode checkout via host-gateway OK"
|
||||
|
||||
@@ -1,56 +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}
|
||||
}
|
||||
|
||||
func (g *gated) Name() string { return g.next.Name() }
|
||||
|
||||
func (g *gated) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
|
||||
res, err := g.rec.Recognize(ctx, req.Prompt, g.snap)
|
||||
if err != nil {
|
||||
// 识别系统错误 fail-closed:不静默放行可能应拦截的内容,也不吞错。
|
||||
return nil, err
|
||||
}
|
||||
switch res.Branch {
|
||||
case intent.BranchGuardHit, intent.BranchMatched, intent.BranchExpectationHit:
|
||||
default:
|
||||
return g.next.Execute(ctx, req) // 无裁决:透传原后端
|
||||
}
|
||||
out, herr, handled := g.handler(ctx, req, res)
|
||||
if !handled {
|
||||
return g.next.Execute(ctx, req)
|
||||
}
|
||||
return out, herr
|
||||
}
|
||||
@@ -1,166 +0,0 @@
|
||||
package intentgate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"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")
|
||||
}
|
||||
}
|
||||
+102
-40
@@ -25,37 +25,48 @@ 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/。
|
||||
Stream llm.StreamFunc
|
||||
}
|
||||
|
||||
// ToolOptions 是可选工具装配。nil 字段不注册对应工具。
|
||||
type ToolOptions struct {
|
||||
Search tools.SearchProvider
|
||||
Exec *tools.ExecConfig
|
||||
Fetch *tools.FetchConfig
|
||||
MCP *mcp.Manager
|
||||
SkillsRoots []string
|
||||
// Compact 可选:转录超过阈值时把头部历史压成摘要(长会话上下文治理)。
|
||||
Compact *CompactConfig
|
||||
// Subagent 可选:提供时注册 delegate_task(委派由宿主实现,内核只开端口)。
|
||||
Subagent subagent.Factory
|
||||
// Log 可选。nil 时丢弃日志。
|
||||
Log *logsys.Logger
|
||||
// ExecutionID 写入日志,不参与执行控制。
|
||||
Subagent subagent.Factory
|
||||
HostTools []tools.Tool
|
||||
}
|
||||
|
||||
// SessionOptions 是单次执行的会话预算与压缩。
|
||||
type SessionOptions struct {
|
||||
Compact *CompactConfig
|
||||
MaxToolCalls int
|
||||
}
|
||||
|
||||
// ObserveOptions 是日志。Log 为 nil 时丢弃。
|
||||
type ObserveOptions struct {
|
||||
Log *logsys.Logger
|
||||
ExecutionID string
|
||||
}
|
||||
|
||||
// Options 构造后端所需的模型接入与执行参数(由组合根注入,无产品依赖)。
|
||||
// 字段经嵌入提升,读取仍写 opts.Model;复合字面量要写子结构名。
|
||||
type Options struct {
|
||||
ModelOptions
|
||||
ToolOptions
|
||||
SessionOptions
|
||||
ObserveOptions
|
||||
}
|
||||
|
||||
// CompactConfig 声明长会话压缩策略(优化性质:压缩失败不阻断执行,但留日志)。
|
||||
type CompactConfig struct {
|
||||
// ThresholdChars 触发阈值(按字节数近似),默认 300_000。
|
||||
@@ -75,6 +86,19 @@ type Backend struct {
|
||||
lastFailure map[string]string
|
||||
}
|
||||
|
||||
type workDirKey struct{}
|
||||
|
||||
// WithWorkDir 把本步工作目录放进上下文,供宿主工具读取内核刚写出的文件。
|
||||
func WithWorkDir(ctx context.Context, dir string) context.Context {
|
||||
return context.WithValue(ctx, workDirKey{}, dir)
|
||||
}
|
||||
|
||||
// WorkDir 返回 WithWorkDir 写入的目录。没有时返回空串。
|
||||
func WorkDir(ctx context.Context) string {
|
||||
dir, _ := ctx.Value(workDirKey{}).(string)
|
||||
return dir
|
||||
}
|
||||
|
||||
// New 创建 native 后端。
|
||||
func New(opts Options) *Backend {
|
||||
return &Backend{opts: opts, lastFailure: map[string]string{}}
|
||||
@@ -168,6 +192,12 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
||||
}
|
||||
}
|
||||
subagent.Register(reg, b.opts.Subagent)
|
||||
for _, tool := range b.opts.HostTools {
|
||||
if reg.Has(tool.Name()) {
|
||||
return &exec.Result{Success: false, ErrorReason: "host tool name conflicts with a builtin: " + tool.Name()}, nil
|
||||
}
|
||||
reg.Register(tool)
|
||||
}
|
||||
if err := reg.Restrict(req.ToolAllowlist); err != nil {
|
||||
return &exec.Result{Success: false, ErrorReason: err.Error()}, nil
|
||||
}
|
||||
@@ -197,6 +227,9 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
||||
timeout = b.opts.Timeout
|
||||
}
|
||||
var runBudget *runtime.Budget
|
||||
if b.opts.MaxToolCalls > 0 {
|
||||
runBudget = &runtime.Budget{MaxToolCalls: b.opts.MaxToolCalls}
|
||||
}
|
||||
if req.Budget != nil {
|
||||
runBudget = &runtime.Budget{
|
||||
MaxPromptTokens: req.Budget.MaxPromptTokens,
|
||||
@@ -204,6 +237,9 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
||||
MaxToolCalls: req.Budget.MaxToolCalls,
|
||||
Deadline: req.Budget.Deadline,
|
||||
}
|
||||
if runBudget.MaxToolCalls == 0 {
|
||||
runBudget.MaxToolCalls = b.opts.MaxToolCalls
|
||||
}
|
||||
if !req.Budget.Deadline.IsZero() {
|
||||
// Deadline 与 Timeout 先到者生效:把更短者作为 ctx 超时。
|
||||
if until := time.Until(req.Budget.Deadline); timeout <= 0 || until < timeout {
|
||||
@@ -211,10 +247,10 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
||||
}
|
||||
}
|
||||
}
|
||||
runCtx := ctx
|
||||
runCtx := WithWorkDir(ctx, workDir)
|
||||
cancel := func() {}
|
||||
if timeout > 0 {
|
||||
runCtx, cancel = context.WithTimeout(ctx, timeout)
|
||||
runCtx, cancel = context.WithTimeout(runCtx, timeout)
|
||||
}
|
||||
defer cancel()
|
||||
|
||||
@@ -247,9 +283,7 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
||||
|
||||
if runRes.Err != nil {
|
||||
reason := runRes.Err.Error()
|
||||
b.mu.Lock()
|
||||
b.lastFailure[workDir] = reason
|
||||
b.mu.Unlock()
|
||||
b.noteFailure(workDir, reason)
|
||||
assessment := exec.AssessStepFailure(exec.StepFailureSignals{
|
||||
Err: runRes.Err,
|
||||
WorkDir: workDir,
|
||||
@@ -264,9 +298,7 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
b.mu.Lock()
|
||||
delete(b.lastFailure, workDir)
|
||||
b.mu.Unlock()
|
||||
b.clearFailure(workDir)
|
||||
|
||||
patches, matErr := exec.MaterializeTargets(workDir, req.Targets, nil)
|
||||
if matErr != nil {
|
||||
@@ -290,7 +322,7 @@ func (b *Backend) run(ctx context.Context, model string, budget *runtime.Budget,
|
||||
if strings.TrimSpace(execID) == "" {
|
||||
execID = b.opts.ExecutionID
|
||||
}
|
||||
runRes, err := runtime.Run(ctx, runtime.Config{
|
||||
agent := runtime.NewAgentSession(runtime.Config{
|
||||
Model: model,
|
||||
SystemPrompt: b.systemPrompt(req, model, profile, workDir, skillIndex),
|
||||
Tools: reg,
|
||||
@@ -313,15 +345,17 @@ func (b *Backend) run(ctx context.Context, model string, budget *runtime.Budget,
|
||||
logs.WriteByte('\n')
|
||||
if onProcess != nil {
|
||||
onProcess(exec.ProcessEvent{
|
||||
Type: ev.Type,
|
||||
Text: ev.Text,
|
||||
Tool: ev.Tool,
|
||||
Detail: ev.Detail,
|
||||
Type: ev.Type,
|
||||
Text: ev.Text,
|
||||
Tool: ev.Tool,
|
||||
Detail: ev.Detail,
|
||||
PromptTokens: ev.PromptTokens,
|
||||
CompletionTokens: ev.CompletionTokens,
|
||||
})
|
||||
}
|
||||
},
|
||||
}, llm.Message{Role: llm.RoleUser, Content: req.Prompt})
|
||||
return runRes, err
|
||||
})
|
||||
return agent.Prompt(ctx, llm.Message{Role: llm.RoleUser, Content: req.Prompt})
|
||||
}
|
||||
|
||||
// maybeCompact 在转录超过阈值时压缩头部历史。压缩是优化性质:
|
||||
@@ -402,9 +436,7 @@ var resumeUnsafeMarkers = []string{
|
||||
|
||||
// rotateIfPoisoned 上次失败属 resume-unsafe 且存在旧会话时归档旧会话。
|
||||
func (b *Backend) rotateIfPoisoned(workDir string) {
|
||||
b.mu.Lock()
|
||||
reason := strings.ToLower(b.lastFailure[workDir])
|
||||
b.mu.Unlock()
|
||||
reason := strings.ToLower(b.failureReason(workDir))
|
||||
if reason == "" {
|
||||
return
|
||||
}
|
||||
@@ -438,6 +470,36 @@ func (b *Backend) logger() *logsys.Logger {
|
||||
return log.With("execution_id", b.opts.ExecutionID)
|
||||
}
|
||||
|
||||
const lastFailureFile = ".native-last-failure"
|
||||
|
||||
func (b *Backend) noteFailure(workDir, reason string) {
|
||||
b.mu.Lock()
|
||||
b.lastFailure[workDir] = reason
|
||||
b.mu.Unlock()
|
||||
_ = os.WriteFile(filepath.Join(workDir, lastFailureFile), []byte(reason), 0o644)
|
||||
}
|
||||
|
||||
func (b *Backend) clearFailure(workDir string) {
|
||||
b.mu.Lock()
|
||||
delete(b.lastFailure, workDir)
|
||||
b.mu.Unlock()
|
||||
_ = os.Remove(filepath.Join(workDir, lastFailureFile))
|
||||
}
|
||||
|
||||
func (b *Backend) failureReason(workDir string) string {
|
||||
b.mu.Lock()
|
||||
reason := b.lastFailure[workDir]
|
||||
b.mu.Unlock()
|
||||
if reason != "" {
|
||||
return reason
|
||||
}
|
||||
raw, err := os.ReadFile(filepath.Join(workDir, lastFailureFile))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
// systemPrompt 经 promptx 组装:角色指令 + 环境事实 + 技能索引 + 交付物约束。
|
||||
func (b *Backend) systemPrompt(req exec.Request, model, profile, workDir string, skillIndex []skills.Skill) string {
|
||||
var index []promptx.IndexEntry
|
||||
|
||||
@@ -39,12 +39,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工具写入的交付物内容。","reason":"产出交付物"}`,
|
||||
}}},
|
||||
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 +86,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,12 +119,21 @@ 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)
|
||||
|
||||
@@ -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"},
|
||||
@@ -47,8 +51,11 @@ func TestBackendCompactBeforeRun(t *testing.T) {
|
||||
t.Fatal("compaction should have run before the loop")
|
||||
}
|
||||
data, _ := os.ReadFile(transcript)
|
||||
if !strings.Contains(string(data), "[conversation summary]") {
|
||||
t.Fatalf("transcript should contain the summary marker, got:\n%.200s", data)
|
||||
if !strings.Contains(string(data), "预置历史摘要") || !strings.Contains(string(data), `"type":"compaction"`) {
|
||||
t.Fatalf("transcript should keep history and append a compaction entry, got:\n%.400s", data)
|
||||
}
|
||||
if !strings.Contains(string(data), `"role":"system"`) {
|
||||
t.Fatal("compaction must not rewrite the original transcript")
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
@@ -190,7 +222,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 +269,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 +305,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 +333,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 +372,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()
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
package native
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"agent-native/exec"
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/tools"
|
||||
)
|
||||
|
||||
type hostTool struct{ name string }
|
||||
|
||||
func (h hostTool) Name() string { return h.name }
|
||||
func (h hostTool) Description() string { return "host tool" }
|
||||
func (h hostTool) Parameters() json.RawMessage {
|
||||
return json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`)
|
||||
}
|
||||
func (h hostTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
|
||||
return tools.Result{Content: "registered " + h.name}, nil
|
||||
}
|
||||
|
||||
func TestHostToolAndUsageEvent(t *testing.T) {
|
||||
var events []exec.ProcessEvent
|
||||
b := New(Options{
|
||||
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: "登记",
|
||||
OnEvent: func(ev exec.ProcessEvent) { events = append(events, ev) },
|
||||
})
|
||||
if err != nil || !res.Success {
|
||||
t.Fatalf("execute: err=%v res=%+v logs=%s", err, res, logsOf(res))
|
||||
}
|
||||
sawText, sawStart, sawEnd, sawUsage := false, false, false, false
|
||||
for _, ev := range events {
|
||||
switch ev.Type {
|
||||
case "text_delta":
|
||||
if ev.Text != "" {
|
||||
sawText = true
|
||||
}
|
||||
case "tool_execution_start":
|
||||
if ev.Tool == "register_artifact" {
|
||||
sawStart = true
|
||||
}
|
||||
case "tool_execution_end":
|
||||
if ev.Tool == "register_artifact" {
|
||||
sawEnd = true
|
||||
}
|
||||
case "usage":
|
||||
if ev.PromptTokens == 11 && ev.CompletionTokens == 7 {
|
||||
sawUsage = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if !sawText || !sawStart || !sawEnd || !sawUsage {
|
||||
t.Fatalf("events missing text/tool/usage: %+v", events)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHostToolSeesStepWorkDir(t *testing.T) {
|
||||
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: "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"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if seen != dir {
|
||||
t.Fatalf("workdir = %q, want %s", seen, dir)
|
||||
}
|
||||
}
|
||||
|
||||
type peekTool struct{ got *string }
|
||||
|
||||
func (p peekTool) Name() string { return "peek" }
|
||||
func (p peekTool) Description() string { return "peek" }
|
||||
func (p peekTool) Parameters() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
|
||||
func (p peekTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
|
||||
*p.got = WorkDir(ctx)
|
||||
return tools.Result{Content: "ok"}, nil
|
||||
}
|
||||
|
||||
func TestHostToolNameConflictFailsClosed(t *testing.T) {
|
||||
b := New(Options{
|
||||
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"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.Success || res.ErrorReason == "" {
|
||||
t.Fatalf("conflict must fail closed, got %+v", res)
|
||||
}
|
||||
}
|
||||
|
||||
func usageScript() llm.StreamFunc {
|
||||
var mu sync.Mutex
|
||||
step := 0
|
||||
return func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
ch := make(chan llm.Event, 4)
|
||||
go func() {
|
||||
defer close(ch)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
step++
|
||||
switch step {
|
||||
case 1:
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
||||
Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
|
||||
ToolCalls: []llm.ToolCall{{ID: "c1", Name: "register_artifact", Arguments: `{"path":"a.docx"}`}},
|
||||
}, Usage: llm.Usage{PromptTokens: 4, CompletionTokens: 2}}
|
||||
case 2:
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
||||
Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
|
||||
ToolCalls: []llm.ToolCall{{ID: "c2", Name: "write_file", Arguments: `{"path":"output.md","content":"done"}`}},
|
||||
}, Usage: llm.Usage{PromptTokens: 4, CompletionTokens: 2}}
|
||||
default:
|
||||
ch <- llm.Event{Kind: llm.KindTextDelta, Text: "已登记"}
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
||||
Role: llm.RoleAssistant, Content: "已登记", StopReason: llm.StopEnd,
|
||||
}, Usage: llm.Usage{PromptTokens: 3, CompletionTokens: 3}}
|
||||
}
|
||||
}()
|
||||
return ch
|
||||
}
|
||||
}
|
||||
|
||||
func logsOf(res *exec.Result) string {
|
||||
if res == nil {
|
||||
return ""
|
||||
}
|
||||
return res.Logs
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package native
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"os/exec"
|
||||
"testing"
|
||||
|
||||
"agent-native/rpc"
|
||||
"agent-native/runtime"
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/session"
|
||||
)
|
||||
|
||||
func TestEventAndRPCFixtures(t *testing.T) {
|
||||
var types []string
|
||||
agent := runtime.NewAgentSession(runtime.Config{
|
||||
Model: "fake",
|
||||
Stream: func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}}
|
||||
close(ch)
|
||||
return ch
|
||||
},
|
||||
OnEvent: func(ev runtime.Event) { types = append(types, ev.Type) },
|
||||
})
|
||||
if _, err := agent.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "hi"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, _ := json.Marshal(types)
|
||||
for _, want := range []string{"message_start", "message_end", "agent_settled"} {
|
||||
if !contains(types, want) {
|
||||
t.Fatalf("fixture missing %s in %s", want, raw)
|
||||
}
|
||||
}
|
||||
dir := t.TempDir()
|
||||
store, err := session.Open(dir + "/s.jsonl")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
srv := &rpc.Server{Agent: agent, Store: store}
|
||||
resp, err := srv.Handle(context.Background(), []byte(`{"type":"get_session_stats"}`))
|
||||
if err != nil || !json.Valid(resp) {
|
||||
t.Fatalf("rpc fixture: %s %v", resp, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPiCrossCompareOptional(t *testing.T) {
|
||||
if _, err := exec.LookPath("pi"); err != nil {
|
||||
t.Skip("pi is not installed")
|
||||
}
|
||||
if os.Getenv("PI_CROSS") != "1" {
|
||||
t.Skip("set PI_CROSS=1 to compare a live pi process")
|
||||
}
|
||||
}
|
||||
|
||||
func contains(list []string, want string) bool {
|
||||
for _, item := range list {
|
||||
if item == want {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package native
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestR2PoisonMarkerSurvivesNewProcess(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
if err := os.MkdirAll(filepath.Join(dir, ".agentx"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
transcript := filepath.Join(dir, ".agentx", "transcript.jsonl")
|
||||
if err := os.WriteFile(transcript, []byte("{}\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(dir, lastFailureFile), []byte("context length exceeded"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b := New(Options{})
|
||||
b.rotateIfPoisoned(dir)
|
||||
if _, err := os.Stat(transcript); !os.IsNotExist(err) {
|
||||
t.Fatal("poisoned transcript still at the live path")
|
||||
}
|
||||
entries, err := os.ReadDir(filepath.Join(dir, ".agentx"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var renamed bool
|
||||
for _, e := range entries {
|
||||
if strings.Contains(e.Name(), "transcript.poisoned.") {
|
||||
renamed = true
|
||||
}
|
||||
}
|
||||
if !renamed {
|
||||
t.Fatal("expected archived transcript")
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"agent-native/rpc"
|
||||
"agent-native/runtime"
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/session"
|
||||
)
|
||||
|
||||
func main() {
|
||||
mode := flag.String("mode", "print", "print, json, or rpc")
|
||||
model := flag.String("model", "", "provider/model[:thinking]")
|
||||
prompt := flag.String("prompt", "", "user message for print and json modes")
|
||||
sessionPath := flag.String("session", "", "transcript path")
|
||||
noSession := flag.Bool("no-session", false, "do not read or write a transcript")
|
||||
export := flag.String("export", "", "write the active session as HTML and exit")
|
||||
flag.Parse()
|
||||
|
||||
if *sessionPath == "" && !*noSession {
|
||||
home, _ := os.UserHomeDir()
|
||||
*sessionPath = filepath.Join(home, ".agent-native", "session.jsonl")
|
||||
}
|
||||
var store *session.Store
|
||||
if !*noSession {
|
||||
st, err := session.Open(*sessionPath)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
store = st
|
||||
}
|
||||
if *export != "" {
|
||||
if store == nil {
|
||||
fmt.Fprintln(os.Stderr, "export requires a session")
|
||||
os.Exit(1)
|
||||
}
|
||||
if err := store.ExportHTML(*export); err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
ref, err := llm.ParseModelRef(*model)
|
||||
if *model != "" && err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
agent := runtime.NewAgentSession(runtime.Config{
|
||||
Model: ref.Model,
|
||||
ThinkingLevel: ref.Thinking,
|
||||
Session: store,
|
||||
Stream: liveStream(ref),
|
||||
OnEvent: func(ev runtime.Event) {
|
||||
if *mode != "json" {
|
||||
return
|
||||
}
|
||||
raw, _ := json.Marshal(map[string]any{"type": ev.Type, "text": ev.Text, "tool": ev.Tool})
|
||||
fmt.Println(string(raw))
|
||||
},
|
||||
})
|
||||
srv := &rpc.Server{Agent: agent, Store: store, NoStore: *noSession}
|
||||
ctx := context.Background()
|
||||
switch *mode {
|
||||
case "rpc":
|
||||
in := bufio.NewScanner(os.Stdin)
|
||||
in.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
|
||||
for in.Scan() {
|
||||
resp, err := srv.Handle(ctx, in.Bytes())
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
os.Stdout.Write(resp)
|
||||
}
|
||||
case "print", "json":
|
||||
if *prompt == "" {
|
||||
*prompt = readAll(os.Stdin)
|
||||
}
|
||||
res, err := agent.Prompt(ctx, llm.Message{Role: llm.RoleUser, Content: *prompt})
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if res.Err != nil {
|
||||
fmt.Fprintln(os.Stderr, res.Err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if *mode == "print" {
|
||||
fmt.Println(res.Text)
|
||||
}
|
||||
default:
|
||||
fmt.Fprintln(os.Stderr, "unknown mode")
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func liveStream(ref llm.ModelRef) llm.StreamFunc {
|
||||
provider, ok := llm.Lookup(ref.Provider, nil)
|
||||
if !ok {
|
||||
provider = llm.Provider{Protocol: llm.ProtocolOpenAIChat, BaseURL: os.Getenv("OPENAI_BASE_URL"), APIKey: os.Getenv("OPENAI_API_KEY")}
|
||||
}
|
||||
client, err := llm.ClientFor(provider)
|
||||
if err != nil {
|
||||
return func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindError, Err: err, Message: llm.Message{StopReason: llm.StopError}}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
}
|
||||
return llm.WithRetry(client.Stream, 2, nil)
|
||||
}
|
||||
|
||||
func readAll(r io.Reader) string {
|
||||
raw, _ := io.ReadAll(r)
|
||||
return string(raw)
|
||||
}
|
||||
+23
-23
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
Vendored
+1
-1
Submodule deps/agent-intent updated: 3674dc4bfd...65aaebca1e
Vendored
+1
-1
Submodule deps/error-code-system updated: dbeed2f1b6...a3ee67bba2
Vendored
+1
-1
Submodule deps/logging-system updated: 2c5716fdcf...49b5d8be60
@@ -89,3 +89,19 @@ baseline_fixed: true
|
||||
|
||||
安全声明(不冒充已解决):run_command **不是 OS 沙箱**(宿主负责隔离);技能正文与 MCP
|
||||
server 描述是提示注入面(只信受信根 + 长度上限);fetch 的 DNS 解析存在 rebinding 窗口。
|
||||
扩展进程、安装的包和外部 CLI 都以调用者的完整权限运行,同样不是沙箱。
|
||||
|
||||
## R16 无界面全量(pi-headless,2026-10-01)
|
||||
|
||||
对标口径改为 pi 的无界面能力。不做 TUI、主题、快捷键,也不做订阅制 OAuth。
|
||||
|
||||
| 域 | 现状 |
|
||||
|---|---|
|
||||
| 会话树 | available:id/parentId、fork/clone、叶子切换、continue/resume、追加式压缩、旧线性文件仍可 Load、export_html |
|
||||
| 供应商 | available:Anthropic、OpenAI Chat 与 Responses、Google、Azure、兼容预设、models.json、`provider/id:thinking`、图片、缓存 token、自动重试;上下文溢出不重试 |
|
||||
| AgentSession | available:steer/follow-up、one 与 all、abort、clear_queue、中途换模型;事件含 message_start、tool_execution_update、queue_update、agent_settled |
|
||||
| 资源 | available:pi 工具名与原名并存、上下文文件、提示词模板、`/skill:name`、项目信任 |
|
||||
| 扩展与包 | available:stdio JSONL 子进程;install/remove/update/list/config |
|
||||
| 模式 | available:`cmd/agent-native` 的 print/json/rpc,以及 Go `rpc.Client` |
|
||||
| exec.Backend | available:`adapters/native` 的 Execute 经 AgentSession,请求/结果字段不变 |
|
||||
|
||||
|
||||
+7
-4
@@ -16,7 +16,7 @@ type Request struct {
|
||||
// OnEvent 可选过程事件回调(nil = 调用方不需要过程,后端行为与现状完全一致)。
|
||||
// 实现方必须保证:回调在 Execute 返回前完成(同步调用或等待收尾),无泄漏 goroutine;
|
||||
// 回调内不得阻塞(消费方自行节流/异步入队)。
|
||||
OnEvent func(ev ProcessEvent)
|
||||
OnEvent func(ev ProcessEvent) `json:"-"`
|
||||
|
||||
// ---- v0.2 能力字段(全部可选;零值 = 既有行为,向后兼容) ----
|
||||
|
||||
@@ -40,9 +40,9 @@ type Request struct {
|
||||
// 共享总预算与截止时间,避免传输、步骤与主循环次数相乘)。
|
||||
// 全部字段可选;零值字段不设限。
|
||||
type Budget struct {
|
||||
MaxPromptTokens int `json:"max_prompt_tokens,omitempty"`
|
||||
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
|
||||
MaxToolCalls int `json:"max_tool_calls,omitempty"`
|
||||
MaxPromptTokens int `json:"max_prompt_tokens,omitempty"`
|
||||
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
|
||||
MaxToolCalls int `json:"max_tool_calls,omitempty"`
|
||||
// Deadline 绝对截止时刻;与 Request.Timeout 先到者生效。
|
||||
Deadline time.Time `json:"deadline,omitempty"`
|
||||
}
|
||||
@@ -54,6 +54,9 @@ type ProcessEvent struct {
|
||||
Text string // delta 文本(Type 为 *_delta 时)
|
||||
Tool string // 工具名(tool_* 时)
|
||||
Detail string // 工具参数/结果摘要(tool_execution_* 时)
|
||||
// PromptTokens / CompletionTokens 在 Type=usage 时填写。
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
}
|
||||
|
||||
const (
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -3,13 +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
|
||||
)
|
||||
|
||||
replace agent-intent => ./deps/agent-intent
|
||||
|
||||
replace error-code-system => ./deps/error-code-system
|
||||
|
||||
replace logging-system => ./deps/logging-system
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
)
|
||||
|
||||
// Client is the Go SDK for the JSONL RPC mode.
|
||||
type Client struct {
|
||||
in io.Writer
|
||||
out *bufio.Reader
|
||||
}
|
||||
|
||||
// NewClient wraps a bidirectional pipe. Commands are one JSON object per line.
|
||||
func NewClient(in io.Writer, out io.Reader) *Client {
|
||||
return &Client{in: in, out: bufio.NewReader(out)}
|
||||
}
|
||||
|
||||
// Call sends one command and reads the matching response.
|
||||
func (c *Client) Call(ctx context.Context, req Request) (Response, error) {
|
||||
raw, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
if _, err := c.in.Write(append(raw, '\n')); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
type result struct {
|
||||
resp Response
|
||||
err error
|
||||
}
|
||||
ch := make(chan result, 1)
|
||||
go func() {
|
||||
line, err := c.out.ReadBytes('\n')
|
||||
if err != nil {
|
||||
ch <- result{err: err}
|
||||
return
|
||||
}
|
||||
var resp Response
|
||||
if err := json.Unmarshal(line, &resp); err != nil {
|
||||
ch <- result{err: err}
|
||||
return
|
||||
}
|
||||
ch <- result{resp: resp}
|
||||
}()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return Response{}, ctx.Err()
|
||||
case got := <-ch:
|
||||
return got.resp, got.err
|
||||
}
|
||||
}
|
||||
+193
@@ -0,0 +1,193 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"agent-native/runtime"
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/session"
|
||||
)
|
||||
|
||||
// Request is one RPC command. Type matches the pi headless command set.
|
||||
type Request struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type"`
|
||||
Message string `json:"message,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
ThinkingLevel string `json:"thinkingLevel,omitempty"`
|
||||
Path string `json:"path,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Leaf string `json:"leaf,omitempty"`
|
||||
Keep int `json:"keep,omitempty"`
|
||||
}
|
||||
|
||||
// Response is the synchronous result of one command.
|
||||
type Response struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type"`
|
||||
OK bool `json:"ok"`
|
||||
Text string `json:"text,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Result any `json:"result,omitempty"`
|
||||
}
|
||||
|
||||
// Server serves RPC commands against one agent session and its transcript.
|
||||
type Server struct {
|
||||
Agent *runtime.AgentSession
|
||||
Store *session.Store
|
||||
NoStore bool
|
||||
}
|
||||
|
||||
// Handle runs one command and returns the JSON response line.
|
||||
func (s *Server) Handle(ctx context.Context, line []byte) ([]byte, error) {
|
||||
var req Request
|
||||
if err := json.Unmarshal(line, &req); err != nil {
|
||||
return encode(Response{Type: "response", OK: false, Error: err.Error()})
|
||||
}
|
||||
resp := s.dispatch(ctx, req)
|
||||
resp.ID = req.ID
|
||||
if resp.Type == "" {
|
||||
resp.Type = "response"
|
||||
}
|
||||
return encode(resp)
|
||||
}
|
||||
|
||||
func (s *Server) dispatch(ctx context.Context, req Request) Response {
|
||||
switch req.Type {
|
||||
case "prompt":
|
||||
if s.Agent == nil {
|
||||
return fail("no agent")
|
||||
}
|
||||
res, err := s.Agent.Prompt(ctx, llm.Message{Role: llm.RoleUser, Content: req.Message})
|
||||
if err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
if res.Err != nil {
|
||||
return fail(res.Err.Error())
|
||||
}
|
||||
return Response{OK: true, Text: res.Text}
|
||||
case "steer":
|
||||
s.Agent.Steer(llm.Message{Role: llm.RoleUser, Content: req.Message})
|
||||
return Response{OK: true}
|
||||
case "follow_up":
|
||||
s.Agent.FollowUp(llm.Message{Role: llm.RoleUser, Content: req.Message})
|
||||
return Response{OK: true}
|
||||
case "abort":
|
||||
s.Agent.Abort()
|
||||
return Response{OK: true}
|
||||
case "clear_queue":
|
||||
s.Agent.ClearQueue()
|
||||
return Response{OK: true}
|
||||
case "set_model":
|
||||
if err := s.Agent.SetModel(req.Model, req.ThinkingLevel); err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
return Response{OK: true}
|
||||
case "compact":
|
||||
keep := req.Keep
|
||||
if keep <= 0 {
|
||||
keep = 4
|
||||
}
|
||||
done, err := s.Agent.Compact(keep, func(messages []llm.Message) (string, error) {
|
||||
return fmt.Sprintf("%d messages", len(messages)), nil
|
||||
})
|
||||
if err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
return Response{OK: true, Result: done}
|
||||
case "get_state":
|
||||
leaf, _ := s.Store.Leaf()
|
||||
return Response{OK: true, Result: map[string]any{"leaf": leaf, "name": s.Store.Name(), "path": s.Store.Path()}}
|
||||
case "get_tree":
|
||||
entries, err := s.Store.Entries()
|
||||
if err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
return Response{OK: true, Result: entries}
|
||||
case "get_entries":
|
||||
entries, err := s.Store.Entries()
|
||||
if err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
return Response{OK: true, Result: entries}
|
||||
case "get_last_assistant_text":
|
||||
msgs, err := s.Store.Load()
|
||||
if err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
text := ""
|
||||
for _, m := range msgs {
|
||||
if m.Role == llm.RoleAssistant {
|
||||
text = m.Content
|
||||
}
|
||||
}
|
||||
return Response{OK: true, Text: text}
|
||||
case "set_session_name":
|
||||
if err := s.Store.SetName(req.Name); err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
return Response{OK: true}
|
||||
case "fork":
|
||||
dest := req.Path
|
||||
if dest == "" {
|
||||
return fail("path is required")
|
||||
}
|
||||
if _, err := s.Store.Fork(dest, req.Leaf); err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
return Response{OK: true, Text: dest}
|
||||
case "clone":
|
||||
if req.Path == "" {
|
||||
return fail("path is required")
|
||||
}
|
||||
if _, err := s.Store.Clone(req.Path); err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
return Response{OK: true, Text: req.Path}
|
||||
case "switch_session":
|
||||
if req.Path == "" {
|
||||
return fail("path is required")
|
||||
}
|
||||
st, err := session.Open(req.Path)
|
||||
if err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
s.Store = st
|
||||
return Response{OK: true}
|
||||
case "new_session":
|
||||
if req.Path == "" {
|
||||
return fail("path is required")
|
||||
}
|
||||
st, err := session.Open(req.Path)
|
||||
if err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
s.Store = st
|
||||
return Response{OK: true, Text: req.Path}
|
||||
case "export_html":
|
||||
if req.Path == "" {
|
||||
return fail("path is required")
|
||||
}
|
||||
if err := s.Store.ExportHTML(req.Path); err != nil {
|
||||
return fail(err.Error())
|
||||
}
|
||||
return Response{OK: true, Text: req.Path}
|
||||
case "get_session_stats":
|
||||
st := s.Store.Stats()
|
||||
return Response{OK: true, Result: st}
|
||||
default:
|
||||
return fail("unknown command " + req.Type)
|
||||
}
|
||||
}
|
||||
|
||||
func fail(msg string) Response { return Response{OK: false, Error: msg} }
|
||||
|
||||
func encode(resp Response) ([]byte, error) {
|
||||
raw, err := json.Marshal(resp)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return append(raw, '\n'), nil
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"agent-native/runtime"
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/session"
|
||||
)
|
||||
|
||||
func TestRPCSessionCommandsAndClient(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
store, err := session.Open(filepath.Join(dir, "s.jsonl"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := store.Append(llm.Message{Role: llm.RoleAssistant, Content: "ready"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
agent := runtime.NewAgentSession(runtime.Config{
|
||||
Model: "fake",
|
||||
Stream: func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "echo", StopReason: llm.StopEnd}}
|
||||
close(ch)
|
||||
return ch
|
||||
},
|
||||
})
|
||||
srv := &Server{Agent: agent, Store: store}
|
||||
clientIn, serverOut := io.Pipe()
|
||||
serverIn, clientOut := io.Pipe()
|
||||
go func() {
|
||||
defer serverOut.Close()
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, err := serverIn.Read(buf)
|
||||
if n > 0 {
|
||||
resp, herr := srv.Handle(context.Background(), buf[:n])
|
||||
if herr == nil {
|
||||
_, _ = serverOut.Write(resp)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
client := NewClient(clientOut, clientIn)
|
||||
resp, err := client.Call(context.Background(), Request{ID: "1", Type: "get_last_assistant_text"})
|
||||
if err != nil || !resp.OK || resp.Text != "ready" {
|
||||
t.Fatalf("%+v %v", resp, err)
|
||||
}
|
||||
fork := filepath.Join(dir, "fork.jsonl")
|
||||
resp, err = client.Call(context.Background(), Request{Type: "fork", Path: fork})
|
||||
if err != nil || !resp.OK {
|
||||
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" {
|
||||
t.Fatalf("prompt %+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 {
|
||||
t.Fatalf("export %+v %v", resp, err)
|
||||
}
|
||||
_ = clientOut.Close()
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
package pkg
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Manifest is the agent-native.json discovered in a package.
|
||||
type Manifest struct {
|
||||
Name string `json:"name"`
|
||||
Extensions []string `json:"extensions,omitempty"`
|
||||
Config map[string]string `json:"config,omitempty"`
|
||||
}
|
||||
|
||||
// Manager installs packages under Dir.
|
||||
type Manager struct {
|
||||
Dir string
|
||||
}
|
||||
|
||||
func (m Manager) root() string {
|
||||
if m.Dir != "" {
|
||||
return m.Dir
|
||||
}
|
||||
home, _ := os.UserHomeDir()
|
||||
return filepath.Join(home, ".agent-native", "packages")
|
||||
}
|
||||
|
||||
// Install copies a local directory, clones a git URL, or runs npm pack
|
||||
// when src starts with "npm:". The package name comes from its manifest.
|
||||
func (m Manager) Install(src string) (Manifest, error) {
|
||||
dir, err := os.MkdirTemp("", "agent-native-pkg-*")
|
||||
if err != nil {
|
||||
return Manifest{}, err
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
switch {
|
||||
case strings.HasPrefix(src, "npm:"):
|
||||
if _, err := exec.LookPath("npm"); err != nil {
|
||||
return Manifest{}, fmt.Errorf("pkg: npm is not installed")
|
||||
}
|
||||
cmd := exec.Command("npm", "pack", strings.TrimPrefix(src, "npm:"))
|
||||
cmd.Dir = dir
|
||||
if out, err := cmd.CombinedOutput(); err != nil {
|
||||
return Manifest{}, fmt.Errorf("pkg: npm pack: %s", out)
|
||||
}
|
||||
case strings.HasPrefix(src, "git+") || strings.HasSuffix(src, ".git") || strings.Contains(src, "://"):
|
||||
cmd := exec.Command("git", "clone", "--depth", "1", src, filepath.Join(dir, "src"))
|
||||
if out, err := cmd.CombinedOutput(); err != nil {
|
||||
return Manifest{}, fmt.Errorf("pkg: git clone: %s", out)
|
||||
}
|
||||
dir = filepath.Join(dir, "src")
|
||||
default:
|
||||
if err := copyDir(src, filepath.Join(dir, "src")); err != nil {
|
||||
return Manifest{}, err
|
||||
}
|
||||
dir = filepath.Join(dir, "src")
|
||||
}
|
||||
manifest, err := ReadManifest(dir)
|
||||
if err != nil {
|
||||
return Manifest{}, err
|
||||
}
|
||||
dest := filepath.Join(m.root(), manifest.Name)
|
||||
if err := os.RemoveAll(dest); err != nil {
|
||||
return Manifest{}, err
|
||||
}
|
||||
if err := copyDir(dir, dest); err != nil {
|
||||
return Manifest{}, err
|
||||
}
|
||||
return manifest, nil
|
||||
}
|
||||
|
||||
// Remove deletes an installed package.
|
||||
func (m Manager) Remove(name string) error {
|
||||
return os.RemoveAll(filepath.Join(m.root(), name))
|
||||
}
|
||||
|
||||
// Update reinstalls from the recorded source when present, otherwise no-ops.
|
||||
func (m Manager) Update(name string) (Manifest, error) {
|
||||
manifest, err := ReadManifest(filepath.Join(m.root(), name))
|
||||
if err != nil {
|
||||
return Manifest{}, err
|
||||
}
|
||||
src := manifest.Config["source"]
|
||||
if src == "" {
|
||||
return manifest, nil
|
||||
}
|
||||
return m.Install(src)
|
||||
}
|
||||
|
||||
// List returns installed manifests.
|
||||
func (m Manager) List() ([]Manifest, error) {
|
||||
entries, err := os.ReadDir(m.root())
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
var out []Manifest
|
||||
for _, ent := range entries {
|
||||
if !ent.IsDir() {
|
||||
continue
|
||||
}
|
||||
manifest, err := ReadManifest(filepath.Join(m.root(), ent.Name()))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, manifest)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Config sets one package config key.
|
||||
func (m Manager) Config(name, key, value string) error {
|
||||
path := filepath.Join(m.root(), name)
|
||||
manifest, err := ReadManifest(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if manifest.Config == nil {
|
||||
manifest.Config = map[string]string{}
|
||||
}
|
||||
manifest.Config[key] = value
|
||||
raw, err := json.MarshalIndent(manifest, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(filepath.Join(path, "agent-native.json"), raw, 0o644)
|
||||
}
|
||||
|
||||
// ReadManifest loads agent-native.json from a directory.
|
||||
func ReadManifest(dir string) (Manifest, error) {
|
||||
raw, err := os.ReadFile(filepath.Join(dir, "agent-native.json"))
|
||||
if err != nil {
|
||||
return Manifest{}, err
|
||||
}
|
||||
var manifest Manifest
|
||||
if err := json.Unmarshal(raw, &manifest); err != nil {
|
||||
return Manifest{}, err
|
||||
}
|
||||
if manifest.Name == "" {
|
||||
return Manifest{}, fmt.Errorf("pkg: manifest name is empty")
|
||||
}
|
||||
return manifest, nil
|
||||
}
|
||||
|
||||
func copyDir(src, dst string) error {
|
||||
return filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rel, err := filepath.Rel(src, path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
target := filepath.Join(dst, rel)
|
||||
if info.IsDir() {
|
||||
return os.MkdirAll(target, 0o755)
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
in, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer in.Close()
|
||||
out, err := os.OpenFile(target, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, info.Mode())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer out.Close()
|
||||
_, err = io.Copy(out, in)
|
||||
return err
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,235 @@
|
||||
package ext
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os/exec"
|
||||
"sync"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/tools"
|
||||
)
|
||||
|
||||
// Message is one JSONL frame on the extension pipe.
|
||||
type Message struct {
|
||||
Type string `json:"type"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Parameters json.RawMessage `json:"parameters,omitempty"`
|
||||
Arguments json.RawMessage `json:"arguments,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
IsError bool `json:"isError,omitempty"`
|
||||
Event string `json:"event,omitempty"`
|
||||
Block bool `json:"block,omitempty"`
|
||||
Reason string `json:"reason,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
Protocol string `json:"protocol,omitempty"`
|
||||
BaseURL string `json:"baseURL,omitempty"`
|
||||
CustomType string `json:"customType,omitempty"`
|
||||
Data json.RawMessage `json:"data,omitempty"`
|
||||
Command string `json:"command,omitempty"`
|
||||
}
|
||||
|
||||
// UIRequest is a question an extension wants a human to answer.
|
||||
type UIRequest struct {
|
||||
ID string
|
||||
Message string
|
||||
}
|
||||
|
||||
// Process is one extension subprocess.
|
||||
type Process struct {
|
||||
cmd *exec.Cmd
|
||||
stdin io.WriteCloser
|
||||
mu sync.Mutex
|
||||
tools []registeredTool
|
||||
commands []string
|
||||
providers []llm.Provider
|
||||
pending map[string]chan Message
|
||||
UI chan UIRequest
|
||||
OnCustom func(customType string, data json.RawMessage)
|
||||
closed bool
|
||||
}
|
||||
|
||||
type registeredTool struct {
|
||||
name string
|
||||
description string
|
||||
parameters json.RawMessage
|
||||
proc *Process
|
||||
}
|
||||
|
||||
func (t registeredTool) Name() string { return t.name }
|
||||
func (t registeredTool) Description() string { return t.description }
|
||||
func (t registeredTool) Parameters() json.RawMessage { return t.parameters }
|
||||
func (t registeredTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
|
||||
return t.proc.callTool(ctx, t.name, args)
|
||||
}
|
||||
|
||||
// Start launches command and reads its registration frames until it goes idle
|
||||
// or the context ends. command is argv, not a shell line.
|
||||
func Start(ctx context.Context, command []string) (*Process, error) {
|
||||
if len(command) == 0 {
|
||||
return nil, fmt.Errorf("ext: empty command")
|
||||
}
|
||||
cmd := exec.CommandContext(ctx, command[0], command[1:]...)
|
||||
stdin, err := cmd.StdinPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p := &Process{
|
||||
cmd: cmd, stdin: stdin, pending: map[string]chan Message{},
|
||||
UI: make(chan UIRequest, 4),
|
||||
}
|
||||
go p.read(stdout)
|
||||
if err := p.send(Message{Type: "init"}); err != nil {
|
||||
_ = p.Close()
|
||||
return nil, err
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// Tools returns the tools the extension registered.
|
||||
func (p *Process) Tools() []tools.Tool {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
out := make([]tools.Tool, len(p.tools))
|
||||
for i := range p.tools {
|
||||
out[i] = p.tools[i]
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Commands returns slash-command names.
|
||||
func (p *Process) Commands() []string {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return append([]string(nil), p.commands...)
|
||||
}
|
||||
|
||||
// Providers returns endpoints the extension registered.
|
||||
func (p *Process) Providers() []llm.Provider {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return append([]llm.Provider(nil), p.providers...)
|
||||
}
|
||||
|
||||
// Hook tells the extension about an event. A block result stops the host action.
|
||||
func (p *Process) Hook(ctx context.Context, event, name string, args json.RawMessage) (bool, string) {
|
||||
id := name + ":" + event
|
||||
ch := make(chan Message, 1)
|
||||
p.mu.Lock()
|
||||
p.pending[id] = ch
|
||||
p.mu.Unlock()
|
||||
if err := p.send(Message{Type: "event", ID: id, Event: event, Name: name, Arguments: args}); err != nil {
|
||||
return false, ""
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false, ctx.Err().Error()
|
||||
case msg := <-ch:
|
||||
return msg.Block, msg.Reason
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Process) callTool(ctx context.Context, name string, args json.RawMessage) (tools.Result, error) {
|
||||
id := name
|
||||
ch := make(chan Message, 1)
|
||||
p.mu.Lock()
|
||||
p.pending[id] = ch
|
||||
p.mu.Unlock()
|
||||
if err := p.send(Message{Type: "tool_call", ID: id, Name: name, Arguments: args}); err != nil {
|
||||
return tools.Result{IsError: true, Content: err.Error()}, err
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return tools.Result{IsError: true, Content: ctx.Err().Error()}, ctx.Err()
|
||||
case msg := <-ch:
|
||||
return tools.Result{Content: msg.Content, IsError: msg.IsError}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Process) send(msg Message) error {
|
||||
raw, err := json.Marshal(msg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.closed {
|
||||
return fmt.Errorf("ext: closed")
|
||||
}
|
||||
_, err = p.stdin.Write(append(raw, '\n'))
|
||||
return err
|
||||
}
|
||||
|
||||
func (p *Process) read(r io.Reader) {
|
||||
sc := bufio.NewScanner(r)
|
||||
sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
||||
for sc.Scan() {
|
||||
var msg Message
|
||||
if err := json.Unmarshal(sc.Bytes(), &msg); err != nil {
|
||||
continue
|
||||
}
|
||||
p.dispatch(msg)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Process) dispatch(msg Message) {
|
||||
switch msg.Type {
|
||||
case "register_tool":
|
||||
p.mu.Lock()
|
||||
p.tools = append(p.tools, registeredTool{
|
||||
name: msg.Name, description: msg.Description, parameters: msg.Parameters, proc: p,
|
||||
})
|
||||
p.mu.Unlock()
|
||||
case "register_command":
|
||||
p.mu.Lock()
|
||||
p.commands = append(p.commands, msg.Name)
|
||||
p.mu.Unlock()
|
||||
case "register_provider":
|
||||
p.mu.Lock()
|
||||
p.providers = append(p.providers, llm.Provider{
|
||||
Name: msg.Name, Protocol: llm.Protocol(msg.Protocol), BaseURL: msg.BaseURL,
|
||||
})
|
||||
p.mu.Unlock()
|
||||
case "tool_result", "hook_result":
|
||||
p.mu.Lock()
|
||||
ch := p.pending[msg.ID]
|
||||
delete(p.pending, msg.ID)
|
||||
p.mu.Unlock()
|
||||
if ch != nil {
|
||||
ch <- msg
|
||||
}
|
||||
case "ui_request":
|
||||
select {
|
||||
case p.UI <- UIRequest{ID: msg.ID, Message: msg.Message}:
|
||||
default:
|
||||
}
|
||||
case "custom_entry":
|
||||
if p.OnCustom != nil {
|
||||
p.OnCustom(msg.CustomType, msg.Data)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close stops the subprocess.
|
||||
func (p *Process) Close() error {
|
||||
p.mu.Lock()
|
||||
p.closed = true
|
||||
_ = p.stdin.Close()
|
||||
p.mu.Unlock()
|
||||
if p.cmd.Process != nil {
|
||||
_ = p.cmd.Process.Kill()
|
||||
}
|
||||
return p.cmd.Wait()
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
package ext
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"agent-native/runtime/ext/pkg"
|
||||
)
|
||||
|
||||
func TestExtensionRegistersAndHooks(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
script := filepath.Join(dir, "ext.sh")
|
||||
body := `#!/bin/sh
|
||||
while IFS= read -r line; do
|
||||
case "$line" in
|
||||
*'"type":"init"'*)
|
||||
echo '{"type":"register_tool","name":"ping","description":"p","parameters":{"type":"object"}}'
|
||||
echo '{"type":"register_command","name":"hi"}'
|
||||
echo '{"type":"register_provider","name":"local","protocol":"openai-chat","baseURL":"http://127.0.0.1:9"}'
|
||||
echo '{"type":"ui_request","id":"u1","message":"choose"}'
|
||||
;;
|
||||
*'"type":"tool_call"'*)
|
||||
echo '{"type":"tool_result","id":"ping","content":"pong"}'
|
||||
;;
|
||||
*'"type":"event"'*)
|
||||
echo '{"type":"hook_result","id":"ping:tool_call","block":true,"reason":"no"}'
|
||||
;;
|
||||
esac
|
||||
done
|
||||
`
|
||||
if err := os.WriteFile(script, []byte(body), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
proc, err := Start(ctx, []string{"/bin/sh", script})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer proc.Close()
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for len(proc.Tools()) == 0 && time.Now().Before(deadline) {
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
tools := proc.Tools()
|
||||
if len(tools) != 1 || tools[0].Name() != "ping" {
|
||||
t.Fatalf("tools %+v", tools)
|
||||
}
|
||||
res, err := tools[0].Execute(ctx, []byte(`{}`))
|
||||
if err != nil || res.Content != "pong" {
|
||||
t.Fatalf("tool: %+v %v", res, err)
|
||||
}
|
||||
block, reason := proc.Hook(ctx, "tool_call", "ping", []byte(`{}`))
|
||||
if !block || reason != "no" {
|
||||
t.Fatalf("hook %v %s", block, reason)
|
||||
}
|
||||
if len(proc.Commands()) != 1 || len(proc.Providers()) != 1 {
|
||||
t.Fatalf("commands %v providers %v", proc.Commands(), proc.Providers())
|
||||
}
|
||||
select {
|
||||
case ui := <-proc.UI:
|
||||
if ui.Message != "choose" {
|
||||
t.Fatal(ui)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("missing ui request")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPackageInstallListConfigRemove(t *testing.T) {
|
||||
src := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(src, "agent-native.json"), []byte(`{"name":"demo","extensions":["./ext.sh"]}`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
mgr := pkg.Manager{Dir: t.TempDir()}
|
||||
manifest, err := mgr.Install(src)
|
||||
if err != nil || manifest.Name != "demo" {
|
||||
t.Fatalf("%+v %v", manifest, err)
|
||||
}
|
||||
if err := mgr.Config("demo", "source", src); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
list, err := mgr.List()
|
||||
if err != nil || len(list) != 1 || list[0].Config["source"] != src {
|
||||
t.Fatalf("list %+v %v", list, err)
|
||||
}
|
||||
if _, err := mgr.Update("demo"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := mgr.Remove("demo"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
list, _ = mgr.List()
|
||||
if len(list) != 0 {
|
||||
t.Fatalf("still installed: %+v", list)
|
||||
}
|
||||
}
|
||||
@@ -53,6 +53,9 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool
|
||||
if opts.MaxTokens > 0 {
|
||||
body["max_tokens"] = opts.MaxTokens
|
||||
}
|
||||
if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" {
|
||||
body["reasoning_effort"] = effort
|
||||
}
|
||||
if len(tools) > 0 {
|
||||
body["tools"] = toProviderTools(tools)
|
||||
body["tool_choice"] = "auto"
|
||||
@@ -125,7 +128,7 @@ func toProviderMessages(messages []Message) []map[string]any {
|
||||
default:
|
||||
out = append(out, map[string]any{
|
||||
"role": string(m.Role),
|
||||
"content": m.Content,
|
||||
"content": messageContent(m),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,447 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Thinking levels, from off to max. Unknown values are rejected by NormalizeThinking.
|
||||
var ThinkingLevels = []string{"off", "minimal", "low", "medium", "high", "xhigh", "max"}
|
||||
|
||||
// NormalizeThinking returns the canonical level. Empty stays empty.
|
||||
func NormalizeThinking(level string) (string, error) {
|
||||
if level == "" {
|
||||
return "", nil
|
||||
}
|
||||
level = strings.ToLower(strings.TrimSpace(level))
|
||||
for _, known := range ThinkingLevels {
|
||||
if level == known {
|
||||
return level, nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("llm: unknown thinking level %q", level)
|
||||
}
|
||||
|
||||
// OpenAIReasoningEffort maps a thinking level onto OpenAI reasoning_effort.
|
||||
// off and empty produce no parameter.
|
||||
func OpenAIReasoningEffort(level string) string {
|
||||
switch level {
|
||||
case "minimal":
|
||||
return "minimal"
|
||||
case "low":
|
||||
return "low"
|
||||
case "medium":
|
||||
return "medium"
|
||||
case "high", "xhigh", "max":
|
||||
return "high"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// AnthropicThinkingBudget maps a thinking level onto a token budget.
|
||||
// off and empty disable thinking.
|
||||
func AnthropicThinkingBudget(level string) int {
|
||||
switch level {
|
||||
case "minimal":
|
||||
return 1024
|
||||
case "low":
|
||||
return 2048
|
||||
case "medium":
|
||||
return 4096
|
||||
case "high":
|
||||
return 8192
|
||||
case "xhigh":
|
||||
return 16384
|
||||
case "max":
|
||||
return 32768
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
// GoogleThinkingBudget maps a thinking level onto Gemini thinkingBudget.
|
||||
func GoogleThinkingBudget(level string) int {
|
||||
switch level {
|
||||
case "minimal", "low":
|
||||
return 1024
|
||||
case "medium":
|
||||
return 4096
|
||||
case "high":
|
||||
return 8192
|
||||
case "xhigh", "max":
|
||||
return 24576
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
// ModelRef is provider/model with an optional :thinking suffix.
|
||||
type ModelRef struct {
|
||||
Provider string
|
||||
Model string
|
||||
Thinking string
|
||||
}
|
||||
|
||||
// ParseModelRef accepts "model", "provider/model", and "provider/model:level".
|
||||
func ParseModelRef(spec string) (ModelRef, error) {
|
||||
spec = strings.TrimSpace(spec)
|
||||
if spec == "" {
|
||||
return ModelRef{}, fmt.Errorf("llm: empty model ref")
|
||||
}
|
||||
thinking := ""
|
||||
if i := strings.LastIndex(spec, ":"); i >= 0 {
|
||||
thinking = spec[i+1:]
|
||||
spec = spec[:i]
|
||||
norm, err := NormalizeThinking(thinking)
|
||||
if err != nil {
|
||||
return ModelRef{}, err
|
||||
}
|
||||
thinking = norm
|
||||
}
|
||||
provider, model := "", spec
|
||||
if i := strings.Index(spec, "/"); i >= 0 {
|
||||
provider = spec[:i]
|
||||
model = spec[i+1:]
|
||||
}
|
||||
if model == "" {
|
||||
return ModelRef{}, fmt.Errorf("llm: model id missing in %q", spec)
|
||||
}
|
||||
return ModelRef{Provider: provider, Model: model, Thinking: thinking}, nil
|
||||
}
|
||||
|
||||
// Protocol names the wire format a provider speaks.
|
||||
type Protocol string
|
||||
|
||||
const (
|
||||
ProtocolOpenAIChat Protocol = "openai-chat"
|
||||
ProtocolOpenAIResponses Protocol = "openai-responses"
|
||||
ProtocolAnthropic Protocol = "anthropic"
|
||||
ProtocolGoogle Protocol = "google"
|
||||
ProtocolAzure Protocol = "azure"
|
||||
)
|
||||
|
||||
// Provider is one configured endpoint.
|
||||
type Provider struct {
|
||||
Name string `json:"name"`
|
||||
Protocol Protocol `json:"protocol"`
|
||||
BaseURL string `json:"baseURL"`
|
||||
APIKeyEnv string `json:"apiKeyEnv,omitempty"`
|
||||
APIKey string `json:"-"`
|
||||
}
|
||||
|
||||
// Presets are OpenAI-compatible or native endpoints that only need an API key.
|
||||
func Presets() []Provider {
|
||||
return []Provider{
|
||||
{Name: "anthropic", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.anthropic.com/v1", APIKeyEnv: "ANTHROPIC_API_KEY"},
|
||||
{Name: "openai", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"},
|
||||
{Name: "openai-responses", Protocol: ProtocolOpenAIResponses, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"},
|
||||
{Name: "google", Protocol: ProtocolOpenAIChat, BaseURL: "https://generativelanguage.googleapis.com/v1beta/openai", APIKeyEnv: "GEMINI_API_KEY"},
|
||||
{Name: "azure", Protocol: ProtocolAzure, BaseURL: "", APIKeyEnv: "AZURE_OPENAI_API_KEY"},
|
||||
{Name: "deepseek", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.deepseek.com/v1", APIKeyEnv: "DEEPSEEK_API_KEY"},
|
||||
{Name: "minimax", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.minimaxi.com/v1", APIKeyEnv: "MINIMAX_API_KEY"},
|
||||
{Name: "openrouter", Protocol: ProtocolOpenAIChat, BaseURL: "https://openrouter.ai/api/v1", APIKeyEnv: "OPENROUTER_API_KEY"},
|
||||
{Name: "groq", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.groq.com/openai/v1", APIKeyEnv: "GROQ_API_KEY"},
|
||||
}
|
||||
}
|
||||
|
||||
// modelsFile is the on-disk shape of models.json.
|
||||
type modelsFile struct {
|
||||
Providers []Provider `json:"providers"`
|
||||
}
|
||||
|
||||
// LoadModelsJSON reads custom providers. Missing file yields an empty list.
|
||||
func LoadModelsJSON(path string) ([]Provider, error) {
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
var file modelsFile
|
||||
if err := json.Unmarshal(raw, &file); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for i := range file.Providers {
|
||||
if file.Providers[i].APIKey == "" && file.Providers[i].APIKeyEnv != "" {
|
||||
file.Providers[i].APIKey = os.Getenv(file.Providers[i].APIKeyEnv)
|
||||
}
|
||||
}
|
||||
return file.Providers, nil
|
||||
}
|
||||
|
||||
// Lookup finds a preset or a custom provider by name.
|
||||
func Lookup(name string, custom []Provider) (Provider, bool) {
|
||||
for _, p := range custom {
|
||||
if p.Name == name {
|
||||
return p, true
|
||||
}
|
||||
}
|
||||
for _, p := range Presets() {
|
||||
if p.Name == name {
|
||||
if p.APIKey == "" {
|
||||
p.APIKey = os.Getenv(p.APIKeyEnv)
|
||||
}
|
||||
return p, true
|
||||
}
|
||||
}
|
||||
return Provider{}, false
|
||||
}
|
||||
|
||||
// ClientFor 返回会说该协议的客户端。responses、azure 以及原生 Anthropic/Google
|
||||
// 协议当前没有接到客户端上,返回明确错误而不是发到错误的路径。
|
||||
func ClientFor(p Provider) (*Client, error) {
|
||||
switch p.Protocol {
|
||||
case ProtocolOpenAIChat:
|
||||
return &Client{BaseURL: p.BaseURL, APIKey: p.APIKey}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("llm: protocol %q is not implemented", p.Protocol)
|
||||
}
|
||||
}
|
||||
|
||||
func messageContent(m Message) any {
|
||||
if len(m.Images) == 0 {
|
||||
return m.Content
|
||||
}
|
||||
parts := []map[string]any{}
|
||||
if m.Content != "" {
|
||||
parts = append(parts, map[string]any{"type": "text", "text": m.Content})
|
||||
}
|
||||
for _, img := range m.Images {
|
||||
url := img.URL
|
||||
if url == "" && img.Data != "" {
|
||||
mime := img.MIME
|
||||
if mime == "" {
|
||||
mime = "image/png"
|
||||
}
|
||||
url = "data:" + mime + ";base64," + img.Data
|
||||
}
|
||||
parts = append(parts, map[string]any{
|
||||
"type": "image_url",
|
||||
"image_url": map[string]any{"url": url},
|
||||
})
|
||||
}
|
||||
return parts
|
||||
}
|
||||
|
||||
// BuildBody returns the JSON body for one native protocol.
|
||||
func BuildBody(protocol Protocol, model string, messages []Message, tools []ToolDecl, opts StreamOptions) ([]byte, error) {
|
||||
switch protocol {
|
||||
case ProtocolOpenAIChat, ProtocolAzure:
|
||||
body := map[string]any{
|
||||
"model": model,
|
||||
"messages": toProviderMessages(messages),
|
||||
"stream": true,
|
||||
}
|
||||
if opts.MaxTokens > 0 {
|
||||
body["max_tokens"] = opts.MaxTokens
|
||||
}
|
||||
if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" {
|
||||
body["reasoning_effort"] = effort
|
||||
}
|
||||
if len(tools) > 0 {
|
||||
body["tools"] = toProviderTools(tools)
|
||||
}
|
||||
return json.Marshal(body)
|
||||
case ProtocolOpenAIResponses:
|
||||
body := map[string]any{
|
||||
"model": model,
|
||||
"input": responsesInput(messages),
|
||||
"stream": true,
|
||||
}
|
||||
if opts.MaxTokens > 0 {
|
||||
body["max_output_tokens"] = opts.MaxTokens
|
||||
}
|
||||
if effort := OpenAIReasoningEffort(opts.ThinkingLevel); effort != "" {
|
||||
body["reasoning"] = map[string]any{"effort": effort}
|
||||
}
|
||||
return json.Marshal(body)
|
||||
// ProtocolAnthropic 与 ProtocolGoogle 只供离线组包。
|
||||
// 当前 Client 只发 OpenAI chat completions,这两支没有接到客户端上。
|
||||
// 原生协议是后续阶段的事;预设表里的 anthropic / google 已改走兼容端点。
|
||||
case ProtocolAnthropic:
|
||||
sys, msgs := anthropicMessages(messages)
|
||||
body := map[string]any{
|
||||
"model": model,
|
||||
"messages": msgs,
|
||||
"max_tokens": opts.MaxTokens,
|
||||
"stream": true,
|
||||
}
|
||||
// Anthropic 要求 max_tokens。0 表示调用方不限制:用足够大的默认值,
|
||||
// 避免 4096 把长回复截断。真正的上限由调用方 StreamOptions.MaxTokens 决定。
|
||||
if body["max_tokens"] == 0 {
|
||||
body["max_tokens"] = 128000
|
||||
}
|
||||
if sys != "" {
|
||||
body["system"] = sys
|
||||
}
|
||||
if budget := AnthropicThinkingBudget(opts.ThinkingLevel); budget > 0 {
|
||||
body["thinking"] = map[string]any{"type": "enabled", "budget_tokens": budget}
|
||||
}
|
||||
return json.Marshal(body)
|
||||
case ProtocolGoogle:
|
||||
body := map[string]any{
|
||||
"contents": googleContents(messages),
|
||||
}
|
||||
if budget := GoogleThinkingBudget(opts.ThinkingLevel); budget > 0 {
|
||||
body["generationConfig"] = map[string]any{
|
||||
"thinkingConfig": map[string]any{"thinkingBudget": budget},
|
||||
}
|
||||
}
|
||||
return json.Marshal(body)
|
||||
default:
|
||||
return nil, fmt.Errorf("llm: unknown protocol %q", protocol)
|
||||
}
|
||||
}
|
||||
|
||||
func responsesInput(messages []Message) []map[string]any {
|
||||
out := make([]map[string]any, 0, len(messages))
|
||||
for _, m := range messages {
|
||||
role := string(m.Role)
|
||||
if role == "" {
|
||||
role = "user"
|
||||
}
|
||||
out = append(out, map[string]any{
|
||||
"role": role,
|
||||
"content": messageContent(m),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func anthropicMessages(messages []Message) (string, []map[string]any) {
|
||||
var system strings.Builder
|
||||
var out []map[string]any
|
||||
for _, m := range messages {
|
||||
if m.Role == RoleSystem {
|
||||
if system.Len() > 0 {
|
||||
system.WriteString("\n")
|
||||
}
|
||||
system.WriteString(m.Content)
|
||||
continue
|
||||
}
|
||||
role := "user"
|
||||
if m.Role == RoleAssistant {
|
||||
role = "assistant"
|
||||
}
|
||||
content := []map[string]any{}
|
||||
if m.Content != "" {
|
||||
content = append(content, map[string]any{"type": "text", "text": m.Content})
|
||||
}
|
||||
for _, img := range m.Images {
|
||||
content = append(content, map[string]any{
|
||||
"type": "image",
|
||||
"source": map[string]any{
|
||||
"type": "base64",
|
||||
"media_type": firstNonEmpty(img.MIME, "image/png"),
|
||||
"data": img.Data,
|
||||
},
|
||||
})
|
||||
}
|
||||
out = append(out, map[string]any{"role": role, "content": content})
|
||||
}
|
||||
return system.String(), out
|
||||
}
|
||||
|
||||
func googleContents(messages []Message) []map[string]any {
|
||||
var out []map[string]any
|
||||
for _, m := range messages {
|
||||
if m.Role == RoleSystem {
|
||||
continue
|
||||
}
|
||||
role := "user"
|
||||
if m.Role == RoleAssistant {
|
||||
role = "model"
|
||||
}
|
||||
parts := []map[string]any{}
|
||||
if m.Content != "" {
|
||||
parts = append(parts, map[string]any{"text": m.Content})
|
||||
}
|
||||
for _, img := range m.Images {
|
||||
parts = append(parts, map[string]any{
|
||||
"inlineData": map[string]any{
|
||||
"mimeType": firstNonEmpty(img.MIME, "image/png"),
|
||||
"data": img.Data,
|
||||
},
|
||||
})
|
||||
}
|
||||
out = append(out, map[string]any{"role": role, "parts": parts})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func firstNonEmpty(a, b string) string {
|
||||
if a != "" {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// 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"))
|
||||
}
|
||||
|
||||
// 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.
|
||||
func WithRetry(base StreamFunc, attempts int, onEvent func(kind, detail string)) StreamFunc {
|
||||
if attempts < 1 {
|
||||
attempts = 1
|
||||
}
|
||||
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
|
||||
for try := 0; try <= attempts; try++ {
|
||||
if try > 0 {
|
||||
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())
|
||||
return
|
||||
case <-time.After(delay):
|
||||
}
|
||||
}
|
||||
last = nil
|
||||
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 ev.Kind == KindError && ev.Retryable {
|
||||
retryable = true
|
||||
}
|
||||
}
|
||||
if overflow || !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:
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
return out
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestParseModelRefAndThinking(t *testing.T) {
|
||||
ref, err := ParseModelRef("anthropic/claude-sonnet:high")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ref.Provider != "anthropic" || ref.Model != "claude-sonnet" || ref.Thinking != "high" {
|
||||
t.Fatalf("%+v", ref)
|
||||
}
|
||||
if _, err := ParseModelRef("openai/gpt:nope"); err == nil {
|
||||
t.Fatal("unknown thinking level must fail")
|
||||
}
|
||||
if AnthropicThinkingBudget("high") != 8192 || OpenAIReasoningEffort("off") != "" {
|
||||
t.Fatal("thinking map")
|
||||
}
|
||||
p, ok := Lookup("groq", nil)
|
||||
if !ok || !strings.Contains(p.BaseURL, "groq") || p.Protocol != ProtocolOpenAIChat {
|
||||
t.Fatalf("preset: %+v", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelsJSONAndBodies(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "models.json")
|
||||
raw := []byte(`{"providers":[{"name":"local","protocol":"openai-chat","baseURL":"http://127.0.0.1:9","apiKeyEnv":"UNUSED_KEY"}]}`)
|
||||
if err := os.WriteFile(path, raw, 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := LoadModelsJSON(path)
|
||||
if err != nil || len(got) != 1 || got[0].Name != "local" {
|
||||
t.Fatalf("%+v %v", got, err)
|
||||
}
|
||||
msgs := []Message{{
|
||||
Role: RoleUser, Content: "see",
|
||||
Images: []Image{{MIME: "image/png", Data: "aaaa"}},
|
||||
}}
|
||||
body, err := BuildBody(ProtocolAnthropic, "claude", append([]Message{{Role: RoleSystem, Content: "sys"}}, msgs...), nil, StreamOptions{ThinkingLevel: "low"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(body), "budget_tokens") || !strings.Contains(string(body), "base64") {
|
||||
t.Fatalf("anthropic body: %s", body)
|
||||
}
|
||||
gbody, err := BuildBody(ProtocolGoogle, "gemini", msgs, nil, StreamOptions{ThinkingLevel: "max"})
|
||||
if err != nil || !strings.Contains(string(gbody), "inlineData") || !strings.Contains(string(gbody), "thinkingBudget") {
|
||||
t.Fatalf("google: %s %v", gbody, err)
|
||||
}
|
||||
rbody, err := BuildBody(ProtocolOpenAIResponses, "gpt", msgs, nil, StreamOptions{ThinkingLevel: "medium", MaxTokens: 20})
|
||||
if err != nil || !strings.Contains(string(rbody), "max_output_tokens") || !strings.Contains(string(rbody), "image_url") {
|
||||
t.Fatalf("responses: %s %v", rbody, err)
|
||||
}
|
||||
var decoded map[string]any
|
||||
if err := json.Unmarshal(rbody, &decoded); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPresetsAndClientForChatCompletions(t *testing.T) {
|
||||
want := map[string]string{
|
||||
"anthropic": "https://api.anthropic.com/v1",
|
||||
"google": "https://generativelanguage.googleapis.com/v1beta/openai",
|
||||
}
|
||||
for _, p := range Presets() {
|
||||
base, ok := want[p.Name]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if p.Protocol != ProtocolOpenAIChat || p.BaseURL != base {
|
||||
t.Fatalf("%s preset: %+v", p.Name, p)
|
||||
}
|
||||
}
|
||||
for _, name := range []string{"openai-responses", "azure"} {
|
||||
p, ok := Lookup(name, nil)
|
||||
if !ok {
|
||||
t.Fatal(name)
|
||||
}
|
||||
if _, err := ClientFor(p); err == nil {
|
||||
t.Fatalf("%s should fail closed", name)
|
||||
}
|
||||
}
|
||||
var path string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
path = r.URL.Path
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||
}))
|
||||
defer srv.Close()
|
||||
client, err := ClientFor(Provider{Protocol: ProtocolOpenAIChat, BaseURL: srv.URL, APIKey: "k"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for ev := range client.Stream(context.Background(), "m", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}) {
|
||||
if ev.Kind == KindError {
|
||||
t.Fatal(ev.Err)
|
||||
}
|
||||
}
|
||||
if path != "/chat/completions" {
|
||||
t.Fatalf("path %q", path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetryThenSuccess(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++
|
||||
if n == 1 {
|
||||
out <- Event{Kind: KindError, Retryable: true, ErrorMessage: "429"}
|
||||
} else {
|
||||
out <- Event{Kind: KindTextDelta, Text: "ok"}
|
||||
out <- Event{Kind: KindMessageEnd, Message: Message{Role: RoleAssistant, Content: "ok"}}
|
||||
}
|
||||
close(out)
|
||||
return out
|
||||
})
|
||||
var events []string
|
||||
wrapped := WithRetry(base, 2, func(kind, detail string) { events = append(events, kind) })
|
||||
var texts []string
|
||||
for ev := range wrapped(context.Background(), "m", nil, nil, StreamOptions{}) {
|
||||
if ev.Kind == KindTextDelta {
|
||||
texts = append(texts, ev.Text)
|
||||
}
|
||||
}
|
||||
if len(texts) != 1 || texts[0] != "ok" {
|
||||
t.Fatalf("texts %v calls %d", texts, n)
|
||||
}
|
||||
if len(events) == 0 || events[0] != "auto_retry_start" {
|
||||
t.Fatalf("events %v", events)
|
||||
}
|
||||
}
|
||||
|
||||
func TestContextOverflowIsNotRetried(t *testing.T) {
|
||||
if !ContextOverflow("maximum context length exceeded") {
|
||||
t.Fatal("expected overflow")
|
||||
}
|
||||
n := 0
|
||||
base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
|
||||
out := make(chan Event, 1)
|
||||
n++
|
||||
out <- Event{Kind: KindError, Retryable: true, ErrorMessage: "context length overflow"}
|
||||
close(out)
|
||||
return out
|
||||
})
|
||||
for range WithRetry(base, 3, nil)(context.Background(), "m", nil, nil, StreamOptions{}) {
|
||||
}
|
||||
if n != 1 {
|
||||
t.Fatalf("overflow must not retry, calls=%d", n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestR2SSEJoinsMultilineData(t *testing.T) {
|
||||
body := strings.Join([]string{
|
||||
`data: {"choices":[{"delta":{"content":"hi"}`,
|
||||
`data: }]}`,
|
||||
``,
|
||||
`data: [DONE]`,
|
||||
``,
|
||||
}, "\n")
|
||||
var got string
|
||||
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
|
||||
if ev.Kind == KindTextDelta {
|
||||
got += ev.Text
|
||||
}
|
||||
if ev.Kind == KindError {
|
||||
t.Fatal(ev.Err)
|
||||
}
|
||||
}
|
||||
if got != "hi" {
|
||||
t.Fatalf("joined content %q", got)
|
||||
}
|
||||
}
|
||||
+70
-12
@@ -26,12 +26,19 @@ type sseChoice struct {
|
||||
|
||||
type sseChunk struct {
|
||||
Choices []sseChoice `json:"choices"`
|
||||
Usage *Usage `json:"usage"`
|
||||
Usage *usageWire `json:"usage"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
|
||||
type usageWire struct {
|
||||
Usage
|
||||
Details struct {
|
||||
Cached int `json:"cached_tokens"`
|
||||
} `json:"prompt_tokens_details"`
|
||||
}
|
||||
|
||||
// ParseSSE reads an OpenAI-compatible server-sent event body and emits
|
||||
// deltas plus one terminal message_end or error event.
|
||||
func ParseSSE(ctx context.Context, r io.Reader) <-chan Event {
|
||||
@@ -51,19 +58,16 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
|
||||
finished := ""
|
||||
sc := bufio.NewScanner(r)
|
||||
sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
||||
for sc.Scan() {
|
||||
reader := &sseEventReader{sc: sc}
|
||||
for {
|
||||
data, ok := reader.next()
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
emitAbort(out, content.String(), ctx.Err())
|
||||
return
|
||||
}
|
||||
line := strings.TrimSpace(sc.Text())
|
||||
if line == "" || strings.HasPrefix(line, ":") {
|
||||
continue
|
||||
}
|
||||
if !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if data == "[DONE]" {
|
||||
break
|
||||
}
|
||||
@@ -77,7 +81,10 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
|
||||
return
|
||||
}
|
||||
if chunk.Usage != nil {
|
||||
usage = *chunk.Usage
|
||||
usage = chunk.Usage.Usage
|
||||
if chunk.Usage.Details.Cached > 0 {
|
||||
usage.CacheReadTokens = chunk.Usage.Details.Cached
|
||||
}
|
||||
}
|
||||
if len(chunk.Choices) == 0 {
|
||||
continue
|
||||
@@ -154,10 +161,61 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
emitAbort(out, content.String(), ctx.Err())
|
||||
case out <- Event{Kind: KindMessageEnd, Message: msg, Usage: usage}:
|
||||
case out <- Event{Kind: KindMessageEnd, Message: msg, Usage: usage, FinishReason: finished}:
|
||||
}
|
||||
}
|
||||
|
||||
// sseEventReader 按空行切事件,并把同一事件里的多行 data: 用换行拼回。
|
||||
// 已经是完整 JSON 的 data 行单独成事件,这样单行流不用空行也能解析。
|
||||
type sseEventReader struct {
|
||||
sc *bufio.Scanner
|
||||
carry string
|
||||
}
|
||||
|
||||
func (r *sseEventReader) next() (string, bool) {
|
||||
var parts []string
|
||||
nextPiece := func() (string, bool) {
|
||||
if r.carry != "" {
|
||||
piece := r.carry
|
||||
r.carry = ""
|
||||
return piece, true
|
||||
}
|
||||
if !r.sc.Scan() {
|
||||
return "", false
|
||||
}
|
||||
return r.sc.Text(), true
|
||||
}
|
||||
for {
|
||||
line, ok := nextPiece()
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
if strings.TrimSpace(line) == "" {
|
||||
if len(parts) == 0 {
|
||||
continue
|
||||
}
|
||||
return strings.Join(parts, "\n"), true
|
||||
}
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, ":") {
|
||||
continue
|
||||
}
|
||||
if !strings.HasPrefix(trimmed, "data:") {
|
||||
continue
|
||||
}
|
||||
piece := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
|
||||
if len(parts) > 0 && json.Valid([]byte(strings.Join(parts, "\n"))) {
|
||||
r.carry = line
|
||||
return strings.Join(parts, "\n"), true
|
||||
}
|
||||
parts = append(parts, piece)
|
||||
}
|
||||
if len(parts) > 0 {
|
||||
return strings.Join(parts, "\n"), true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func mapFinish(reason string, hasTools bool) StopReason {
|
||||
switch reason {
|
||||
case "tool_calls", "function_call":
|
||||
|
||||
@@ -38,6 +38,29 @@ func TestParseSSETextAndUsage(t *testing.T) {
|
||||
if end.Usage.TotalTokens != 5 {
|
||||
t.Fatalf("usage %#v", end.Usage)
|
||||
}
|
||||
if end.FinishReason != "stop" {
|
||||
t.Fatalf("finish reason %q, want stop", end.FinishReason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSSEFinishReasonLengthKeepsStopEnd(t *testing.T) {
|
||||
body := strings.Join([]string{
|
||||
`data: {"choices":[{"delta":{"content":"半"},"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.FinishReason != "length" {
|
||||
t.Fatalf("finish reason %q, want length", end.FinishReason)
|
||||
}
|
||||
if end.Message.StopReason != StopEnd {
|
||||
t.Fatalf("stop reason %s, want end so existing callers stay unchanged", end.Message.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSSEToolCallIncrements(t *testing.T) {
|
||||
|
||||
+19
-3
@@ -48,6 +48,7 @@ type Message struct {
|
||||
Role Role `json:"role"`
|
||||
Content string `json:"content,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Images []Image `json:"images,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
StopReason StopReason `json:"stop_reason,omitempty"`
|
||||
@@ -61,11 +62,20 @@ type ToolDecl struct {
|
||||
Parameters json.RawMessage `json:"parameters"`
|
||||
}
|
||||
|
||||
// Image is one picture attached to a user message.
|
||||
type Image struct {
|
||||
MIME string `json:"mime,omitempty"`
|
||||
URL string `json:"url,omitempty"`
|
||||
Data string `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
// Usage counts tokens reported by the provider. 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"`
|
||||
}
|
||||
|
||||
// Event is one streamed update. Failures are KindError plus a final assistant
|
||||
@@ -80,9 +90,12 @@ type Event struct {
|
||||
ArgumentsDelta string
|
||||
Message Message
|
||||
Usage Usage
|
||||
Err error
|
||||
Retryable bool
|
||||
ErrorMessage string
|
||||
// FinishReason 是供应商原始结束原因(stop、length、tool_calls 等)。
|
||||
// 空表示这条事件不携带结束原因。StopReason 的映射不受此字段影响。
|
||||
FinishReason string
|
||||
Err error
|
||||
Retryable bool
|
||||
ErrorMessage string
|
||||
}
|
||||
|
||||
// StreamOptions tunes one provider request.
|
||||
@@ -90,6 +103,9 @@ type StreamOptions struct {
|
||||
MaxTokens int
|
||||
APIBase string
|
||||
APIKey string
|
||||
// ThinkingLevel is off, minimal, low, medium, high, xhigh, or max.
|
||||
// Empty leaves the provider default unchanged.
|
||||
ThinkingLevel string
|
||||
}
|
||||
|
||||
// StreamFunc is the provider boundary. It must return a channel that is closed
|
||||
|
||||
+315
-10
@@ -18,8 +18,15 @@ import (
|
||||
const (
|
||||
ExecSequential = "sequential"
|
||||
ExecParallel = "parallel"
|
||||
|
||||
maxStopContinuations = 3
|
||||
maxLengthContinuations = 3
|
||||
maxRepeatReminders = 3
|
||||
)
|
||||
|
||||
const repeatReminderText = "你已连续 3 次以相同参数调用同一工具。如果结果不符合预期,请更换参数或做法,不要再重复这一调用。"
|
||||
const lengthContinuePrompt = "上一段回复因为达到长度上限被截断。请从断点接着写,不要重复已经写出的内容。"
|
||||
|
||||
// Event is a lifecycle notice for logs and UI.
|
||||
type Event struct {
|
||||
Type string
|
||||
@@ -29,6 +36,11 @@ type Event struct {
|
||||
// Detail carries tool arguments/results or reasoning deltas for
|
||||
// process-visualization consumers; log writers may ignore it.
|
||||
Detail string
|
||||
// PromptTokens / CompletionTokens / TotalTokens 在 type=usage 或
|
||||
// type=message_end 时有意义,供调用方累计用量。
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
TotalTokens int
|
||||
}
|
||||
|
||||
// Config drives one agent session.
|
||||
@@ -42,7 +54,13 @@ type Config struct {
|
||||
ToolExecution string
|
||||
GetSteering func() []llm.Message
|
||||
GetFollowUp func() []llm.Message
|
||||
OnEvent func(Event)
|
||||
// OnEvent 观察循环事件。同一次 Run 内不会并发调用
|
||||
// (与 BeforeTool、AfterTool 共用一把锁串行)。
|
||||
OnEvent func(Event)
|
||||
// ThinkingLevel is forwarded on each provider request.
|
||||
ThinkingLevel string
|
||||
// Resolve, when set, overrides Model and ThinkingLevel at the start of each turn.
|
||||
Resolve func() (model, thinking string)
|
||||
// Log 可选。nil 时丢弃日志,不改变循环结果。
|
||||
Log *logsys.Logger
|
||||
// ExecutionID 只用于日志字段,不参与控制流。
|
||||
@@ -50,6 +68,46 @@ type Config struct {
|
||||
// Budget 可选。任一上限命中即停止并返回 budget_exceeded 错误(StopError);
|
||||
// 零值字段不设限。截止时刻与外层 ctx 先到者生效。
|
||||
Budget *Budget
|
||||
// History 是会话转录为空时的初始转录。以 system 开头时,不再前置
|
||||
// SystemPrompt 和工具清单。不回写进 Session。
|
||||
History []llm.Message
|
||||
// Stop 在模型不再调用工具、且本轮没有 follow-up 时询问是否续跑。
|
||||
// nil 表示不询问。内核最多接受 3 次 Continue。
|
||||
Stop func(last llm.Message, info StopInfo) StopDecision
|
||||
// BeforeTool 返回错误时不执行工具,但仍写一条 IsError 的 tool 结果。
|
||||
// 同一次 Run 内与 OnEvent、AfterTool 串行,不会并发调用。
|
||||
BeforeTool func(call llm.ToolCall) error
|
||||
// UnknownTool 处理注册表里没有的调用。nil 时仍返回 unknown tool。
|
||||
UnknownTool func(call llm.ToolCall) tools.Result
|
||||
// AfterTool 只观察工具结果,不能改写。
|
||||
// 同一次 Run 内与 OnEvent、BeforeTool 串行,不会并发调用。
|
||||
AfterTool func(call llm.ToolCall, res llm.Message)
|
||||
// Reminders 在每次调用模型前提供提醒。正文包在 system-reminder 里,
|
||||
// 只进入本次请求,不写进 Session。
|
||||
Reminders func(turn int) []string
|
||||
// RepeatGuard 为真时,同名且参数按键排序后相同的调用连续 3 次,注入一条提醒。
|
||||
// 每个 run 最多 3 条,不中断循环。
|
||||
RepeatGuard bool
|
||||
// ContinueOnLength 为真时,FinishReason=length 且没有工具调用就续写。
|
||||
// 续写提示不写进 Session。次数由 MaxLengthContinuations 决定。
|
||||
ContinueOnLength bool
|
||||
// MaxLengthContinuations 是 length 续写的次数上限。
|
||||
// 0 表示默认 3 次;负数表示不限制。
|
||||
MaxLengthContinuations int
|
||||
}
|
||||
|
||||
// StopInfo 是 Stop 钩子看到的本 run 进度。
|
||||
type StopInfo struct {
|
||||
Continuations int
|
||||
ToolCalls int
|
||||
}
|
||||
|
||||
// StopDecision 决定模型停下之后是否再跑一轮。
|
||||
type StopDecision struct {
|
||||
Continue bool
|
||||
Replacement string
|
||||
Reason string
|
||||
Notice string
|
||||
}
|
||||
|
||||
// Budget 是循环层的执行预算;与 exec.Budget 字段同构,由适配层换算。
|
||||
@@ -66,6 +124,7 @@ type Result struct {
|
||||
Messages []llm.Message
|
||||
StopReason llm.StopReason
|
||||
Usage llm.Usage
|
||||
Notice string
|
||||
Err error
|
||||
Retryable bool
|
||||
}
|
||||
@@ -81,6 +140,7 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
if cfg.ToolExecution == "" {
|
||||
cfg.ToolExecution = ExecSequential
|
||||
}
|
||||
serializeHostCallbacks(&cfg)
|
||||
var messages []llm.Message
|
||||
if cfg.Session != nil {
|
||||
loaded, err := cfg.Session.Load()
|
||||
@@ -89,6 +149,9 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
}
|
||||
messages = loaded
|
||||
}
|
||||
if len(messages) == 0 && len(cfg.History) > 0 {
|
||||
messages = append([]llm.Message{}, cfg.History...)
|
||||
}
|
||||
if len(messages) == 0 || messages[0].Role != llm.RoleSystem {
|
||||
sys := llm.Message{Role: llm.RoleSystem, Content: systemText(cfg)}
|
||||
messages = append([]llm.Message{sys}, messages...)
|
||||
@@ -105,6 +168,10 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
var last llm.Message
|
||||
var usage llm.Usage
|
||||
toolCalls := 0
|
||||
stopContinues := 0
|
||||
lengthContinues := 0
|
||||
lengthNudge := false
|
||||
var repeats repeatState
|
||||
for turn := 0; turn < cfg.MaxTurns; turn++ {
|
||||
if ctx.Err() != nil {
|
||||
cfg.logger().Error("runtime turn aborted", ctx.Err())
|
||||
@@ -114,7 +181,9 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
return budgetStop(cfg, last, messages, usage, reason), nil
|
||||
}
|
||||
emit(cfg, Event{Type: "turn_start"})
|
||||
end, errEv := consume(ctx, cfg, messages)
|
||||
emit(cfg, Event{Type: "message_start"})
|
||||
request := withEphemeral(messages, turnReminders(cfg, turn, &repeats, &lengthNudge))
|
||||
end, errEv := consume(ctx, cfg, request)
|
||||
if errEv != nil {
|
||||
cfg.logger().Error("runtime turn failed", errEv, "retryable", end.Retryable)
|
||||
return Result{Text: end.Message.Content, Messages: messages, StopReason: end.Message.StopReason, Err: errEv, Retryable: end.Retryable}, nil
|
||||
@@ -127,14 +196,30 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
return Result{}, err
|
||||
}
|
||||
messages = append(messages, last)
|
||||
emit(cfg, Event{Type: "message_end", Message: last})
|
||||
emit(cfg, Event{Type: "message_end", Message: last, PromptTokens: end.Usage.PromptTokens, CompletionTokens: end.Usage.CompletionTokens, TotalTokens: end.Usage.TotalTokens})
|
||||
|
||||
if last.StopReason == llm.StopToolCalls && len(last.ToolCalls) > 0 {
|
||||
if reason := budgetExceeded(cfg.Budget, usage, toolCalls, time.Now()); reason != "" {
|
||||
var err error
|
||||
messages, err = appendSkippedToolResults(cfg, last.ToolCalls, messages)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
return budgetStop(cfg, last, messages, usage, reason), nil
|
||||
}
|
||||
if cfg.Budget != nil && cfg.Budget.MaxToolCalls > 0 && toolCalls+len(last.ToolCalls) > cfg.Budget.MaxToolCalls {
|
||||
reason := fmt.Sprintf("budget_exceeded: tool calls %d+%d > max %d", toolCalls, len(last.ToolCalls), cfg.Budget.MaxToolCalls)
|
||||
var err error
|
||||
messages, err = appendSkippedToolResults(cfg, last.ToolCalls, messages)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
return budgetStop(cfg, last, messages, usage, reason), nil
|
||||
}
|
||||
toolCalls += len(last.ToolCalls)
|
||||
if cfg.RepeatGuard {
|
||||
repeats.observe(last.ToolCalls)
|
||||
}
|
||||
results := executeTools(ctx, cfg, last)
|
||||
for _, res := range results {
|
||||
if err := persist(cfg, res); err != nil {
|
||||
@@ -154,6 +239,20 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
continue
|
||||
}
|
||||
emit(cfg, Event{Type: "turn_end", Message: last})
|
||||
if cfg.ContinueOnLength && end.FinishReason == "length" && len(last.ToolCalls) == 0 {
|
||||
lengthCap := cfg.MaxLengthContinuations
|
||||
if lengthCap == 0 {
|
||||
lengthCap = maxLengthContinuations
|
||||
}
|
||||
if lengthCap > 0 && lengthContinues >= lengthCap {
|
||||
errCap := errString("length continuation exceeded")
|
||||
cfg.logger().Error("runtime length continuation exceeded", errCap)
|
||||
return Result{Text: last.Content, Messages: messages, StopReason: llm.StopError, Usage: usage, Err: errCap}, nil
|
||||
}
|
||||
lengthContinues++
|
||||
lengthNudge = true
|
||||
continue
|
||||
}
|
||||
if follow := drain(cfg.GetFollowUp); len(follow) > 0 {
|
||||
for _, m := range follow {
|
||||
if err := persist(cfg, m); err != nil {
|
||||
@@ -163,8 +262,31 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
||||
}
|
||||
continue
|
||||
}
|
||||
emit(cfg, Event{Type: "agent_end", Message: last})
|
||||
return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: usage}, nil
|
||||
if cfg.Stop != nil {
|
||||
decision := cfg.Stop(last, StopInfo{Continuations: stopContinues, ToolCalls: toolCalls})
|
||||
if decision.Continue && stopContinues < maxStopContinuations {
|
||||
stopContinues++
|
||||
if decision.Replacement != "" && len(messages) > 0 {
|
||||
messages[len(messages)-1].Content = decision.Replacement
|
||||
last.Content = decision.Replacement
|
||||
note := llm.Message{Role: llm.RoleUser, Content: decision.Replacement}
|
||||
if err := persist(cfg, note); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
messages = append(messages, note)
|
||||
}
|
||||
if decision.Reason != "" {
|
||||
reasonMsg := llm.Message{Role: llm.RoleUser, Content: decision.Reason}
|
||||
if err := persist(cfg, reasonMsg); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
messages = append(messages, reasonMsg)
|
||||
}
|
||||
continue
|
||||
}
|
||||
return settle(cfg, last, messages, usage, decision.Notice), nil
|
||||
}
|
||||
return settle(cfg, last, messages, usage, ""), nil
|
||||
}
|
||||
errCap := errString("max turns exceeded")
|
||||
cfg.logger().Error("runtime turn cap hit", errCap)
|
||||
@@ -218,9 +340,21 @@ func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event
|
||||
decls = append(decls, llm.ToolDecl{Name: d.Name, Description: d.Description, Parameters: d.Parameters})
|
||||
}
|
||||
}
|
||||
model := cfg.Model
|
||||
opts := llm.StreamOptions{ThinkingLevel: cfg.ThinkingLevel}
|
||||
if cfg.Resolve != nil {
|
||||
if m, th := cfg.Resolve(); m != "" || th != "" {
|
||||
if m != "" {
|
||||
model = m
|
||||
}
|
||||
if th != "" {
|
||||
opts.ThinkingLevel = th
|
||||
}
|
||||
}
|
||||
}
|
||||
var end llm.Event
|
||||
sawEnd := false
|
||||
for ev := range cfg.Stream(ctx, cfg.Model, messages, decls, llm.StreamOptions{}) {
|
||||
for ev := range cfg.Stream(ctx, model, messages, decls, opts) {
|
||||
switch ev.Kind {
|
||||
case llm.KindTextDelta:
|
||||
emit(cfg, Event{Type: "text_delta", Text: ev.Text})
|
||||
@@ -241,12 +375,66 @@ func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event
|
||||
return end, nil
|
||||
}
|
||||
|
||||
const toolBudgetSkipContent = "not executed: tool budget exceeded"
|
||||
|
||||
// appendSkippedToolResults 为未执行的 tool_call 各写一条错误结果,保持调用与结果成对。
|
||||
// 不发工具事件:这些调用没有进入执行。
|
||||
func appendSkippedToolResults(cfg Config, calls []llm.ToolCall, messages []llm.Message) ([]llm.Message, error) {
|
||||
for _, call := range calls {
|
||||
msg := llm.Message{
|
||||
Role: llm.RoleTool,
|
||||
Name: call.Name,
|
||||
ToolCallID: call.ID,
|
||||
Content: toolBudgetSkipContent,
|
||||
IsError: true,
|
||||
}
|
||||
if err := persist(cfg, msg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
messages = append(messages, msg)
|
||||
}
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
// serializeHostCallbacks 让同一次 Run 的宿主回调不会并发进入。
|
||||
func serializeHostCallbacks(cfg *Config) {
|
||||
var mu sync.Mutex
|
||||
if cfg.OnEvent != nil {
|
||||
inner := cfg.OnEvent
|
||||
cfg.OnEvent = func(ev Event) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
inner(ev)
|
||||
}
|
||||
}
|
||||
if cfg.BeforeTool != nil {
|
||||
inner := cfg.BeforeTool
|
||||
cfg.BeforeTool = func(call llm.ToolCall) error {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
return inner(call)
|
||||
}
|
||||
}
|
||||
if cfg.AfterTool != nil {
|
||||
inner := cfg.AfterTool
|
||||
cfg.AfterTool = func(call llm.ToolCall, res llm.Message) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
inner(call, res)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func executeTools(ctx context.Context, cfg Config, assistant llm.Message) []llm.Message {
|
||||
if cfg.ToolExecution == ExecParallel && len(assistant.ToolCalls) > 1 {
|
||||
return executeParallel(ctx, cfg, assistant)
|
||||
}
|
||||
out := make([]llm.Message, 0, len(assistant.ToolCalls))
|
||||
for _, call := range assistant.ToolCalls {
|
||||
if err := ctx.Err(); err != nil {
|
||||
out = append(out, cancelledToolMessage(call, err))
|
||||
continue
|
||||
}
|
||||
out = append(out, runOne(ctx, cfg, call))
|
||||
}
|
||||
return out
|
||||
@@ -266,20 +454,133 @@ func executeParallel(ctx context.Context, cfg Config, assistant llm.Message) []l
|
||||
return out
|
||||
}
|
||||
|
||||
func cancelledToolMessage(call llm.ToolCall, err error) llm.Message {
|
||||
return llm.Message{
|
||||
Role: llm.RoleTool,
|
||||
Name: call.Name,
|
||||
ToolCallID: call.ID,
|
||||
Content: "not executed: " + err.Error(),
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message {
|
||||
emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name})
|
||||
if err := ctx.Err(); err != nil {
|
||||
return cancelledToolMessage(call, err)
|
||||
}
|
||||
emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name, Detail: call.Arguments})
|
||||
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})
|
||||
if cfg.AfterTool != nil {
|
||||
cfg.AfterTool(call, msg)
|
||||
}
|
||||
return msg
|
||||
}
|
||||
}
|
||||
emit(cfg, Event{Type: "tool_execution_update", Tool: call.Name, Detail: call.Arguments})
|
||||
res := tools.Result{Content: "no tools configured", IsError: true}
|
||||
if cfg.Tools != nil {
|
||||
known := cfg.Tools != nil && cfg.Tools.Has(call.Name)
|
||||
switch {
|
||||
case known:
|
||||
res = cfg.Tools.Execute(ctx, call.Name, json.RawMessage(call.Arguments))
|
||||
case cfg.UnknownTool != nil:
|
||||
res = cfg.UnknownTool(call)
|
||||
case cfg.Tools != nil:
|
||||
res = cfg.Tools.Execute(ctx, call.Name, json.RawMessage(call.Arguments))
|
||||
}
|
||||
emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: res.Content})
|
||||
return llm.Message{
|
||||
msg := llm.Message{
|
||||
Role: llm.RoleTool,
|
||||
Name: call.Name,
|
||||
ToolCallID: call.ID,
|
||||
Content: res.Content,
|
||||
IsError: res.IsError,
|
||||
}
|
||||
emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: res.Content})
|
||||
if cfg.AfterTool != nil {
|
||||
cfg.AfterTool(call, msg)
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
func settle(cfg Config, last llm.Message, messages []llm.Message, usage llm.Usage, notice string) Result {
|
||||
emit(cfg, Event{Type: "agent_end", Message: last})
|
||||
emit(cfg, Event{Type: "agent_settled", Message: last})
|
||||
emitUsage(cfg, usage)
|
||||
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 turnReminders(cfg Config, turn int, repeats *repeatState, lengthNudge *bool) []llm.Message {
|
||||
var extra []llm.Message
|
||||
if cfg.Reminders != nil {
|
||||
for _, text := range cfg.Reminders(turn) {
|
||||
if strings.TrimSpace(text) == "" {
|
||||
continue
|
||||
}
|
||||
extra = append(extra, reminderMessage(text))
|
||||
}
|
||||
}
|
||||
if lengthNudge != nil && *lengthNudge {
|
||||
extra = append(extra, llm.Message{Role: llm.RoleUser, Content: lengthContinuePrompt})
|
||||
*lengthNudge = false
|
||||
}
|
||||
if cfg.RepeatGuard && repeats != nil {
|
||||
due := repeats.streak/3 - repeats.reminders
|
||||
room := maxRepeatReminders - repeats.reminders
|
||||
if due > room {
|
||||
due = room
|
||||
}
|
||||
for i := 0; i < due; i++ {
|
||||
extra = append(extra, reminderMessage(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
|
||||
reminders int
|
||||
}
|
||||
|
||||
func (s *repeatState) observe(calls []llm.ToolCall) {
|
||||
for _, call := range calls {
|
||||
sig := call.Name + "\n" + canonicalArgs(call.Arguments)
|
||||
if sig == s.sig {
|
||||
s.streak++
|
||||
continue
|
||||
}
|
||||
s.sig = sig
|
||||
s.streak = 1
|
||||
}
|
||||
}
|
||||
|
||||
func canonicalArgs(raw string) string {
|
||||
var v any
|
||||
if err := json.Unmarshal([]byte(raw), &v); err != nil {
|
||||
return raw
|
||||
}
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return raw
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func persist(cfg Config, msg llm.Message) error {
|
||||
@@ -296,6 +597,10 @@ func drain(fn func() []llm.Message) []llm.Message {
|
||||
return fn()
|
||||
}
|
||||
|
||||
func emitUsage(cfg Config, usage llm.Usage) {
|
||||
emit(cfg, Event{Type: "usage", PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens, TotalTokens: usage.TotalTokens})
|
||||
}
|
||||
|
||||
func emit(cfg Config, ev Event) {
|
||||
if cfg.OnEvent != nil {
|
||||
cfg.OnEvent(ev)
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -66,15 +67,108 @@ func TestLoopBudgetToolCallCap(t *testing.T) {
|
||||
if !sawBudgetEvent {
|
||||
t.Fatal("expected budget_exceeded event")
|
||||
}
|
||||
// 第二批工具调用不该被执行(预算在执行前拦截):transcript 只含一次 tool 结果。
|
||||
// 第一批执行 1 次;第二批超限,两条合成错误结果,不执行。
|
||||
toolResults := 0
|
||||
skipped := 0
|
||||
for _, m := range res.Messages {
|
||||
if m.Role == llm.RoleTool {
|
||||
toolResults++
|
||||
if m.Role != llm.RoleTool {
|
||||
continue
|
||||
}
|
||||
toolResults++
|
||||
if m.IsError && m.Content == toolBudgetSkipContent {
|
||||
skipped++
|
||||
}
|
||||
}
|
||||
if toolResults != 1 {
|
||||
t.Fatalf("expected exactly 1 tool result, got %d", toolResults)
|
||||
if toolResults != 3 || skipped != 2 {
|
||||
t.Fatalf("expected 3 tool results including 2 skipped, got results=%d skipped=%d", toolResults, skipped)
|
||||
}
|
||||
assertToolCallsPaired(t, res.Messages)
|
||||
}
|
||||
|
||||
func assertToolCallsPaired(t *testing.T, messages []llm.Message) {
|
||||
t.Helper()
|
||||
pending := map[string]string{}
|
||||
for _, m := range messages {
|
||||
if m.Role == llm.RoleAssistant {
|
||||
for _, call := range m.ToolCalls {
|
||||
if _, ok := pending[call.ID]; ok {
|
||||
t.Fatalf("duplicate tool call id %s", call.ID)
|
||||
}
|
||||
pending[call.ID] = call.Name
|
||||
}
|
||||
}
|
||||
if m.Role == llm.RoleTool {
|
||||
name, ok := pending[m.ToolCallID]
|
||||
if !ok {
|
||||
t.Fatalf("tool result %s has no matching call", m.ToolCallID)
|
||||
}
|
||||
if name != m.Name {
|
||||
t.Fatalf("tool result %s name %q != call %q", m.ToolCallID, m.Name, name)
|
||||
}
|
||||
delete(pending, m.ToolCallID)
|
||||
}
|
||||
}
|
||||
if len(pending) != 0 {
|
||||
t.Fatalf("unpaired tool calls: %v", pending)
|
||||
}
|
||||
}
|
||||
|
||||
type pauseTool struct{}
|
||||
|
||||
func (pauseTool) Name() string { return "pause" }
|
||||
func (pauseTool) Description() string { return "slow tool" }
|
||||
func (pauseTool) Parameters() json.RawMessage {
|
||||
return json.RawMessage(`{"type":"object","properties":{}}`)
|
||||
}
|
||||
func (pauseTool) Execute(context.Context, json.RawMessage) (tools.Result, error) {
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
return tools.Result{Content: "ok"}, nil
|
||||
}
|
||||
|
||||
func TestParallelHostCallbacksAreSerial(t *testing.T) {
|
||||
reg := tools.NewRegistry(t.TempDir())
|
||||
reg.Register(pauseTool{})
|
||||
var mu sync.Mutex
|
||||
inFlight, maxInFlight := 0, 0
|
||||
enter := func() {
|
||||
mu.Lock()
|
||||
inFlight++
|
||||
if inFlight > maxInFlight {
|
||||
maxInFlight = inFlight
|
||||
}
|
||||
mu.Unlock()
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
mu.Lock()
|
||||
inFlight--
|
||||
mu.Unlock()
|
||||
}
|
||||
_, err := Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
Tools: reg,
|
||||
ToolExecution: ExecParallel,
|
||||
OnEvent: func(Event) {
|
||||
enter()
|
||||
},
|
||||
BeforeTool: func(llm.ToolCall) error {
|
||||
enter()
|
||||
return nil
|
||||
},
|
||||
AfterTool: func(llm.ToolCall, llm.Message) {
|
||||
enter()
|
||||
},
|
||||
Stream: usageScript(llm.Usage{},
|
||||
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{
|
||||
{ID: "a", Name: "pause", Arguments: `{}`},
|
||||
{ID: "b", Name: "pause", Arguments: `{}`},
|
||||
}},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
|
||||
),
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if maxInFlight != 1 {
|
||||
t.Fatalf("host callbacks overlapped: max in flight %d", maxInFlight)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,372 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/session"
|
||||
"agent-native/runtime/tools"
|
||||
)
|
||||
|
||||
func TestLoopStopContinuesUntilCap(t *testing.T) {
|
||||
var calls int
|
||||
res, err := Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
MaxTurns: 8,
|
||||
Stream: script(
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "forged-claim", StopReason: llm.StopEnd},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "forged-claim", StopReason: llm.StopEnd},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "forged-claim", StopReason: llm.StopEnd},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "forged-claim", StopReason: llm.StopEnd},
|
||||
),
|
||||
Stop: func(last llm.Message, info StopInfo) StopDecision {
|
||||
calls++
|
||||
if info.Continuations != calls-1 {
|
||||
t.Errorf("continuations = %d, want %d", info.Continuations, calls-1)
|
||||
}
|
||||
return StopDecision{Continue: true, Replacement: "已作废", Reason: "再试", Notice: "到顶了"}
|
||||
},
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if calls != 4 {
|
||||
t.Fatalf("stop calls = %d, want 4", calls)
|
||||
}
|
||||
if res.Err != nil || res.Notice != "到顶了" {
|
||||
t.Fatalf("notice %q err %v", res.Notice, res.Err)
|
||||
}
|
||||
replaced, forged := 0, 0
|
||||
for _, m := range res.Messages {
|
||||
if m.Role != llm.RoleAssistant {
|
||||
continue
|
||||
}
|
||||
if m.Content == "已作废" {
|
||||
replaced++
|
||||
}
|
||||
if m.Content == "forged-claim" {
|
||||
forged++
|
||||
}
|
||||
}
|
||||
if replaced != 3 || forged != 1 {
|
||||
t.Fatalf("replaced=%d forged=%d transcript=%s", replaced, forged, transcriptText(res.Messages))
|
||||
}
|
||||
if strings.Count(transcriptText(res.Messages), "再试") != 3 {
|
||||
t.Fatalf("reasons = %s", transcriptText(res.Messages))
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoopBeforeToolRejectPairsResult(t *testing.T) {
|
||||
hits := 0
|
||||
reg := tools.NewRegistry(t.TempDir())
|
||||
reg.Register(hitTool{hits: &hits})
|
||||
var seen llm.Message
|
||||
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: "echo", Arguments: `{}`,
|
||||
}}},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
|
||||
),
|
||||
BeforeTool: func(call llm.ToolCall) error {
|
||||
return errors.New("不允许")
|
||||
},
|
||||
AfterTool: func(call llm.ToolCall, msg llm.Message) {
|
||||
seen = msg
|
||||
},
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if hits != 0 {
|
||||
t.Fatalf("tool executed %d times", hits)
|
||||
}
|
||||
if !seen.IsError || seen.Content != "不允许" || seen.ToolCallID != "c1" {
|
||||
t.Fatalf("after = %#v", seen)
|
||||
}
|
||||
paired := false
|
||||
for _, m := range res.Messages {
|
||||
if m.Role == llm.RoleTool && m.ToolCallID == "c1" && m.IsError && m.Content == "不允许" {
|
||||
paired = true
|
||||
}
|
||||
}
|
||||
if !paired {
|
||||
t.Fatalf("missing paired tool result: %s", transcriptText(res.Messages))
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoopRemindersStayOutOfSession(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "transcript.jsonl")
|
||||
st, err := session.Open(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var saw bool
|
||||
_, err = Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
Session: st,
|
||||
Stream: func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
for _, m := range messages {
|
||||
if strings.Contains(m.Content, "<system-reminder>") && strings.Contains(m.Content, "先核对交付物") {
|
||||
saw = true
|
||||
}
|
||||
}
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}}
|
||||
close(ch)
|
||||
return ch
|
||||
},
|
||||
Reminders: func(turn int) []string {
|
||||
if turn != 0 {
|
||||
t.Errorf("turn %d", turn)
|
||||
}
|
||||
return []string{"先核对交付物"}
|
||||
},
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "ping"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !saw {
|
||||
t.Fatal("reminder missing from the provider request")
|
||||
}
|
||||
loaded, err := st.Load()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Contains(transcriptText(loaded), "先核对交付物") {
|
||||
t.Fatalf("reminder persisted: %#v", loaded)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoopRepeatGuardFiresAndCaps(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var requests []string
|
||||
turns := make([]llm.Message, 0, 13)
|
||||
for i := 0; i < 12; i++ {
|
||||
args := `{"b":1,"a":2}`
|
||||
if i%2 == 1 {
|
||||
args = `{"a":2,"b":1}`
|
||||
}
|
||||
turns = append(turns, llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
|
||||
ID: "c", Name: "echo", Arguments: args,
|
||||
}}})
|
||||
}
|
||||
turns = append(turns, llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd})
|
||||
var i int
|
||||
stream := func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
mu.Lock()
|
||||
requests = append(requests, transcriptText(messages))
|
||||
n := i
|
||||
i++
|
||||
mu.Unlock()
|
||||
ch := make(chan llm.Event, 1)
|
||||
msg := turns[len(turns)-1]
|
||||
if n < len(turns) {
|
||||
msg = turns[n]
|
||||
}
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
res, err := Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
MaxTurns: 20,
|
||||
RepeatGuard: true,
|
||||
Stream: stream,
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.Text != "done" {
|
||||
t.Fatalf("text %q", res.Text)
|
||||
}
|
||||
reminders := 0
|
||||
for _, req := range requests {
|
||||
reminders += strings.Count(req, repeatReminderText)
|
||||
}
|
||||
if reminders != 3 {
|
||||
t.Fatalf("reminders = %d, want 3; requests=%d", reminders, len(requests))
|
||||
}
|
||||
if strings.Contains(transcriptText(res.Messages), repeatReminderText) {
|
||||
t.Fatal("repeat reminder was written into the transcript")
|
||||
}
|
||||
if !strings.Contains(requests[3], repeatReminderText) {
|
||||
t.Fatal("first reminder should arrive on the request after the third identical call")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoopContinueOnLengthCaps(t *testing.T) {
|
||||
var requests []string
|
||||
var n int
|
||||
stream := func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
requests = append(requests, transcriptText(messages))
|
||||
n++
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{
|
||||
Kind: llm.KindMessageEnd,
|
||||
FinishReason: "length",
|
||||
Message: llm.Message{Role: llm.RoleAssistant, Content: "part", StopReason: llm.StopEnd},
|
||||
}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
res, err := Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
MaxTurns: 8,
|
||||
ContinueOnLength: true,
|
||||
Stream: stream,
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.Err == nil || !strings.Contains(res.Err.Error(), "length continuation exceeded") {
|
||||
t.Fatalf("err = %v", res.Err)
|
||||
}
|
||||
if n != 4 {
|
||||
t.Fatalf("model calls = %d, want 4", n)
|
||||
}
|
||||
if len(res.Messages) == 0 || !strings.Contains(transcriptText(res.Messages), "part") {
|
||||
t.Fatal("truncated transcript was dropped")
|
||||
}
|
||||
if strings.Contains(transcriptText(res.Messages), lengthContinuePrompt) {
|
||||
t.Fatal("continuation prompt was persisted")
|
||||
}
|
||||
prompts := 0
|
||||
for _, req := range requests {
|
||||
prompts += strings.Count(req, lengthContinuePrompt)
|
||||
}
|
||||
if prompts != 3 {
|
||||
t.Fatalf("continuation prompts = %d, want 3", prompts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoopContinueOnLengthUnlimited(t *testing.T) {
|
||||
var n int
|
||||
stream := func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
n++
|
||||
ch := make(chan llm.Event, 1)
|
||||
finish := "length"
|
||||
text := "part"
|
||||
if n >= 5 {
|
||||
finish = "stop"
|
||||
text = "done"
|
||||
}
|
||||
ch <- llm.Event{
|
||||
Kind: llm.KindMessageEnd,
|
||||
FinishReason: finish,
|
||||
Message: llm.Message{Role: llm.RoleAssistant, Content: text, StopReason: llm.StopEnd},
|
||||
}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
res, err := Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
MaxTurns: 8,
|
||||
ContinueOnLength: true,
|
||||
MaxLengthContinuations: -1,
|
||||
Stream: stream,
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.Err != nil {
|
||||
t.Fatalf("err = %v", res.Err)
|
||||
}
|
||||
if n != 5 {
|
||||
t.Fatalf("model calls = %d, want 5", n)
|
||||
}
|
||||
if res.Text != "done" {
|
||||
t.Fatalf("text = %q", res.Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoopHistorySkipsGeneratedSystem(t *testing.T) {
|
||||
var first string
|
||||
_, err := Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
SystemPrompt: "LEAK",
|
||||
History: []llm.Message{
|
||||
{Role: llm.RoleSystem, Content: "FROM_HISTORY"},
|
||||
{Role: llm.RoleUser, Content: "earlier"},
|
||||
},
|
||||
Stream: func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
if len(messages) > 0 {
|
||||
first = messages[0].Content
|
||||
}
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}}
|
||||
close(ch)
|
||||
return ch
|
||||
},
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "now"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if first != "FROM_HISTORY" || strings.Contains(first, "LEAK") {
|
||||
t.Fatalf("system = %q", first)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoopUnknownToolFallsThrough(t *testing.T) {
|
||||
res, err := Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
UnknownTool: func(call llm.ToolCall) tools.Result {
|
||||
return tools.Result{Content: "via " + call.Name}
|
||||
},
|
||||
Stream: script(
|
||||
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
|
||||
ID: "c1", Name: "make_report", Arguments: `{}`,
|
||||
}}},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
|
||||
),
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, m := range res.Messages {
|
||||
if m.Role == llm.RoleTool && m.Content == "via make_report" && !m.IsError {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("unknown tool was not delegated: %s", transcriptText(res.Messages))
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoopTurnCapKeepsTranscript(t *testing.T) {
|
||||
res, err := Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
MaxTurns: 1,
|
||||
Stream: script(llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
|
||||
ID: "c1", Name: "missing", Arguments: `{}`,
|
||||
}}}),
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "ping"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.Err == nil || !strings.Contains(res.Err.Error(), "max turns exceeded") {
|
||||
t.Fatalf("err = %v", res.Err)
|
||||
}
|
||||
if !strings.Contains(transcriptText(res.Messages), "ping") {
|
||||
t.Fatalf("transcript dropped: %s", transcriptText(res.Messages))
|
||||
}
|
||||
}
|
||||
|
||||
type hitTool struct{ hits *int }
|
||||
|
||||
func (h hitTool) Name() string { return "echo" }
|
||||
func (h hitTool) Description() string { return "echo" }
|
||||
func (h hitTool) Parameters() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
|
||||
func (h hitTool) Execute(context.Context, json.RawMessage) (tools.Result, error) {
|
||||
*h.hits++
|
||||
return tools.Result{Content: "ran"}, nil
|
||||
}
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"time"
|
||||
)
|
||||
|
||||
// RunFakeServer 在 stdin/stdout 上运行最小 MCP server(echo / boom / 未知方法),
|
||||
@@ -67,6 +68,9 @@ func fakeServerHandle(req rpcRequest) rpcResponse {
|
||||
Content: []contentBlock{{Type: "text", Text: "boom failed"}},
|
||||
IsError: true,
|
||||
})
|
||||
case "slow":
|
||||
time.Sleep(2 * time.Second)
|
||||
return respond(callToolResult{Content: []contentBlock{{Type: "text", Text: "slow done"}}})
|
||||
default:
|
||||
return fail(-32602, "unknown tool "+p.Name)
|
||||
}
|
||||
@@ -87,6 +91,11 @@ func fakeServerTools() []toolInfo {
|
||||
Description: "Always fails",
|
||||
InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
|
||||
},
|
||||
{
|
||||
Name: "slow",
|
||||
Description: "Responds after a delay",
|
||||
InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+25
-1
@@ -92,7 +92,31 @@ func (c *httpClient) CallTool(ctx context.Context, name string, args json.RawMes
|
||||
return res, err
|
||||
}
|
||||
|
||||
func (c *httpClient) Close() error { return nil }
|
||||
func (c *httpClient) Close() error {
|
||||
c.mu.Lock()
|
||||
sid := c.sessionID
|
||||
url := c.cfg.URL
|
||||
client := c.client
|
||||
c.mu.Unlock()
|
||||
if sid == "" || url == "" || client == nil {
|
||||
return nil
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodDelete, url, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Mcp-Session-Id", sid)
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 2048))
|
||||
if resp.StatusCode >= 400 {
|
||||
return fmt.Errorf("mcp session terminate: http %d", resp.StatusCode)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *httpClient) roundTrip(ctx context.Context, method string, params any, out any) error {
|
||||
c.mu.Lock()
|
||||
|
||||
+25
-6
@@ -96,7 +96,7 @@ func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error {
|
||||
}
|
||||
desc := ti.Description
|
||||
if len(desc) > maxDescriptionBytes {
|
||||
desc = desc[:maxDescriptionBytes] + "…[truncated]"
|
||||
desc = tools.CutUTF8(desc, maxDescriptionBytes) + "…[truncated]"
|
||||
}
|
||||
schema := ti.InputSchema
|
||||
if len(schema) == 0 {
|
||||
@@ -109,6 +109,12 @@ func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error {
|
||||
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)
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -130,7 +136,7 @@ func (m *Manager) Close() {
|
||||
m.clients = map[string]Client{}
|
||||
}
|
||||
|
||||
// bridgeTool 把一个 MCP 工具适配为 tools.Tool;调用失败重连一次后重试。
|
||||
// bridgeTool 把一个 MCP 工具适配为 tools.Tool。CallTool 失败时重连一次再试。
|
||||
type bridgeTool struct {
|
||||
client Client
|
||||
serverCfg ServerConfig
|
||||
@@ -138,6 +144,7 @@ type bridgeTool struct {
|
||||
declName string
|
||||
description string
|
||||
schema json.RawMessage
|
||||
redial func(ctx context.Context) (Client, error)
|
||||
}
|
||||
|
||||
func (t *bridgeTool) Name() string { return t.declName }
|
||||
@@ -147,9 +154,15 @@ func (t *bridgeTool) Parameters() json.RawMessage {
|
||||
}
|
||||
|
||||
func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
|
||||
callCtx, cancel := context.WithTimeout(ctx, requestTimeout(t.serverCfg))
|
||||
defer cancel()
|
||||
res, err := t.client.CallTool(callCtx, t.toolName, args)
|
||||
res, err := t.call(ctx, args)
|
||||
if err != nil && t.redial != nil {
|
||||
redial := t.redial
|
||||
t.redial = nil
|
||||
if c, rerr := redial(ctx); rerr == nil {
|
||||
t.client = c
|
||||
res, err = t.call(ctx, args)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return tools.Result{Content: fmt.Sprintf("mcp tool %s failed: %v", t.declName, err), IsError: true}, nil
|
||||
}
|
||||
@@ -162,10 +175,16 @@ func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.R
|
||||
}
|
||||
out := b.String()
|
||||
if len(out) > maxTextResultBytes {
|
||||
out = out[:maxTextResultBytes] + "\n[truncated]"
|
||||
out = tools.CutUTF8(out, maxTextResultBytes) + "\n[truncated]"
|
||||
}
|
||||
if res.IsError {
|
||||
return tools.Result{Content: out, IsError: true}, nil
|
||||
}
|
||||
return tools.Result{Content: out}, nil
|
||||
}
|
||||
|
||||
func (t *bridgeTool) call(ctx context.Context, args json.RawMessage) (callToolResult, error) {
|
||||
callCtx, cancel := context.WithTimeout(ctx, requestTimeout(t.serverCfg))
|
||||
defer cancel()
|
||||
return t.client.CallTool(callCtx, t.toolName, args)
|
||||
}
|
||||
|
||||
+29
-3
@@ -45,7 +45,7 @@ func TestStdioClientHandshakeListAndCall(t *testing.T) {
|
||||
defer c.Close()
|
||||
|
||||
list, err := c.ListTools(ctx)
|
||||
if err != nil || len(list) != 2 || list[0].Name != "echo" {
|
||||
if err != nil || len(list) != 3 || list[0].Name != "echo" {
|
||||
t.Fatalf("list failed: %v %+v", err, list)
|
||||
}
|
||||
res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"hi"}`))
|
||||
@@ -75,6 +75,32 @@ func TestStdioClientServerRejectsBadSchema(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestStdioSlowCancelThenEcho(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
c, err := DialStdio(ctx, fakeStdioConfig(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer c.Close()
|
||||
|
||||
slowCtx, cancel := context.WithTimeout(ctx, 80*time.Millisecond)
|
||||
defer cancel()
|
||||
_, err = c.CallTool(slowCtx, "slow", nil)
|
||||
if err == nil {
|
||||
t.Fatal("slow call should time out")
|
||||
}
|
||||
|
||||
echoCtx, echoCancel := context.WithTimeout(ctx, 3*time.Second)
|
||||
defer echoCancel()
|
||||
res, err := c.CallTool(echoCtx, "echo", json.RawMessage(`{"text":"after-slow"}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(res.Content) != 1 || res.Content[0].Text != "echo: after-slow" {
|
||||
t.Fatalf("echo must receive its own result, got %+v", res)
|
||||
}
|
||||
}
|
||||
|
||||
func fakeHTTPHandler(w http.ResponseWriter, r *http.Request) {
|
||||
var req rpcRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
@@ -104,7 +130,7 @@ func TestHTTPClientJSONMode(t *testing.T) {
|
||||
t.Fatalf("session id must be captured, got %q", c.sessionID)
|
||||
}
|
||||
list, err := c.ListTools(ctx)
|
||||
if err != nil || len(list) != 2 {
|
||||
if err != nil || len(list) != 3 {
|
||||
t.Fatalf("list failed: %v %+v", err, list)
|
||||
}
|
||||
res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"net"}`))
|
||||
@@ -130,7 +156,7 @@ func TestHTTPClientSSEMode(t *testing.T) {
|
||||
}
|
||||
defer c.Close()
|
||||
list, err := c.ListTools(ctx)
|
||||
if err != nil || len(list) != 2 {
|
||||
if err != nil || len(list) != 3 {
|
||||
t.Fatalf("sse list failed: %v %+v", err, list)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type flakyClient struct {
|
||||
fail bool
|
||||
calls int
|
||||
}
|
||||
|
||||
func (f *flakyClient) Initialize(context.Context) (initializeResult, error) {
|
||||
return initializeResult{}, nil
|
||||
}
|
||||
func (f *flakyClient) ListTools(context.Context) ([]toolInfo, error) { return nil, nil }
|
||||
func (f *flakyClient) Close() error { return nil }
|
||||
func (f *flakyClient) CallTool(context.Context, string, json.RawMessage) (callToolResult, error) {
|
||||
f.calls++
|
||||
if f.fail {
|
||||
return callToolResult{}, fmt.Errorf("down")
|
||||
}
|
||||
return callToolResult{Content: []contentBlock{{Type: "text", Text: "back"}}}, nil
|
||||
}
|
||||
|
||||
func TestR2MCPRetriesOnceAfterReconnect(t *testing.T) {
|
||||
first := &flakyClient{fail: true}
|
||||
second := &flakyClient{}
|
||||
tool := &bridgeTool{
|
||||
client: first,
|
||||
toolName: "echo",
|
||||
declName: "mcp__s__echo",
|
||||
redial: func(context.Context) (Client, error) {
|
||||
return second, nil
|
||||
},
|
||||
}
|
||||
res, err := tool.Execute(context.Background(), []byte(`{}`))
|
||||
if err != nil || res.IsError {
|
||||
t.Fatalf("retry: %+v %v", res, err)
|
||||
}
|
||||
if first.calls != 1 || second.calls != 1 || !strings.Contains(res.Content, "back") {
|
||||
t.Fatalf("calls first=%d second=%d content=%q", first.calls, second.calls, res.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestR2HTTPCloseTerminatesSession(t *testing.T) {
|
||||
var method, sid string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
method = r.Method
|
||||
sid = r.Header.Get("Mcp-Session-Id")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
c := &httpClient{cfg: ServerConfig{URL: srv.URL}, client: &http.Client{}, sessionID: "sess-1"}
|
||||
if err := c.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if method != http.MethodDelete || sid != "sess-1" {
|
||||
t.Fatalf("terminate %s session %q", method, sid)
|
||||
}
|
||||
}
|
||||
+94
-46
@@ -12,14 +12,23 @@ import (
|
||||
)
|
||||
|
||||
// stdioClient 通过「子进程 stdin/stdout 按行分隔的 JSON-RPC」与 server 通信。
|
||||
// 同一时刻只允许一个在途请求(串行锁),匹配简单且足以支撑智能体工具调用。
|
||||
// 一个长驻读循环按 id 把响应分发给各请求;写 stdin 仍串行。
|
||||
// 取消会摘掉等待项,迟到的回复直接丢弃。读循环退出时,在等的请求全部失败,客户端标记为已死。
|
||||
type stdioClient struct {
|
||||
cmd *exec.Cmd
|
||||
stdin io.WriteCloser
|
||||
stdout *bufio.Reader
|
||||
mu sync.Mutex
|
||||
nextID int64
|
||||
closed bool
|
||||
cmd *exec.Cmd
|
||||
stdin io.WriteCloser
|
||||
stdout *bufio.Reader
|
||||
writeMu sync.Mutex
|
||||
mu sync.Mutex
|
||||
pending map[int64]chan stdioResult
|
||||
nextID int64
|
||||
closed bool
|
||||
readerOnce sync.Once
|
||||
}
|
||||
|
||||
type stdioResult struct {
|
||||
resp rpcResponse
|
||||
err error
|
||||
}
|
||||
|
||||
// DialStdio 启动 server 子进程并完成 initialize 握手。
|
||||
@@ -43,7 +52,13 @@ func DialStdio(ctx context.Context, cfg ServerConfig) (*stdioClient, error) {
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, fmt.Errorf("mcp server %q: start %s: %w", cfg.Name, cfg.Command, err)
|
||||
}
|
||||
c := &stdioClient{cmd: cmd, stdin: stdin, stdout: bufio.NewReaderSize(stdout, 1<<20)}
|
||||
c := &stdioClient{
|
||||
cmd: cmd,
|
||||
stdin: stdin,
|
||||
stdout: bufio.NewReaderSize(stdout, 1<<20),
|
||||
pending: map[int64]chan stdioResult{},
|
||||
}
|
||||
c.startReader()
|
||||
if err := c.handshake(ctx); err != nil {
|
||||
_ = c.Close()
|
||||
return nil, err
|
||||
@@ -104,11 +119,12 @@ func (c *stdioClient) CallTool(ctx context.Context, name string, args json.RawMe
|
||||
|
||||
func (c *stdioClient) Close() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.closed {
|
||||
c.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
c.closed = true
|
||||
c.mu.Unlock()
|
||||
_ = c.stdin.Close()
|
||||
if c.cmd.Process != nil {
|
||||
_ = c.cmd.Process.Kill()
|
||||
@@ -117,29 +133,43 @@ func (c *stdioClient) Close() error {
|
||||
}
|
||||
|
||||
func (c *stdioClient) roundTrip(ctx context.Context, method string, params any, out any) error {
|
||||
c.startReader()
|
||||
id := atomic.AddInt64(&c.nextID, 1)
|
||||
ch := make(chan stdioResult, 1)
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.closed {
|
||||
c.mu.Unlock()
|
||||
return fmt.Errorf("mcp client closed")
|
||||
}
|
||||
id := atomic.AddInt64(&c.nextID, 1)
|
||||
c.pending[id] = ch
|
||||
c.mu.Unlock()
|
||||
|
||||
req := rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params}
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
c.dropPending(id)
|
||||
return err
|
||||
}
|
||||
body = append(body, '\n')
|
||||
if _, err := c.stdin.Write(body); err != nil {
|
||||
c.writeMu.Lock()
|
||||
_, err = c.stdin.Write(body)
|
||||
c.writeMu.Unlock()
|
||||
if err != nil {
|
||||
c.dropPending(id)
|
||||
return fmt.Errorf("mcp write %s: %w", method, err)
|
||||
}
|
||||
line, err := c.readResponseLine(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("mcp read %s: %w", method, err)
|
||||
|
||||
var got stdioResult
|
||||
select {
|
||||
case got = <-ch:
|
||||
case <-ctx.Done():
|
||||
c.dropPending(id)
|
||||
return fmt.Errorf("mcp read %s: %w", method, ctx.Err())
|
||||
}
|
||||
var resp rpcResponse
|
||||
if err := json.Unmarshal(line, &resp); err != nil {
|
||||
return fmt.Errorf("mcp %s: bad response frame: %w", method, err)
|
||||
if got.err != nil {
|
||||
return fmt.Errorf("mcp read %s: %w", method, got.err)
|
||||
}
|
||||
resp := got.resp
|
||||
if resp.ID == nil || *resp.ID != id {
|
||||
return fmt.Errorf("mcp %s: response id mismatch", method)
|
||||
}
|
||||
@@ -152,40 +182,56 @@ func (c *stdioClient) roundTrip(ctx context.Context, method string, params any,
|
||||
return nil
|
||||
}
|
||||
|
||||
// readResponseLine 逐行读取,跳过通知帧(无 id),带 ctx 取消感知。
|
||||
func (c *stdioClient) readResponseLine(ctx context.Context) ([]byte, error) {
|
||||
type result struct {
|
||||
line []byte
|
||||
err error
|
||||
}
|
||||
done := make(chan result, 1)
|
||||
go func() {
|
||||
for {
|
||||
line, err := c.stdout.ReadBytes('\n')
|
||||
if err != nil {
|
||||
done <- result{nil, err}
|
||||
return
|
||||
}
|
||||
var probe rpcResponse
|
||||
if json.Unmarshal(line, &probe) == nil && probe.ID != nil {
|
||||
done <- result{line, nil}
|
||||
return
|
||||
}
|
||||
// 通知或日志行:忽略。
|
||||
func (c *stdioClient) dropPending(id int64) {
|
||||
c.mu.Lock()
|
||||
delete(c.pending, id)
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func (c *stdioClient) startReader() {
|
||||
c.readerOnce.Do(func() {
|
||||
go c.readLoop()
|
||||
})
|
||||
}
|
||||
|
||||
func (c *stdioClient) readLoop() {
|
||||
var fatal error
|
||||
for {
|
||||
line, err := c.stdout.ReadBytes('\n')
|
||||
if err != nil {
|
||||
fatal = err
|
||||
break
|
||||
}
|
||||
}()
|
||||
select {
|
||||
case r := <-done:
|
||||
return r.line, r.err
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
var resp rpcResponse
|
||||
if json.Unmarshal(line, &resp) != nil || resp.ID == nil {
|
||||
continue
|
||||
}
|
||||
c.mu.Lock()
|
||||
ch := c.pending[*resp.ID]
|
||||
if ch != nil {
|
||||
delete(c.pending, *resp.ID)
|
||||
}
|
||||
c.mu.Unlock()
|
||||
if ch == nil {
|
||||
continue
|
||||
}
|
||||
ch <- stdioResult{resp: resp}
|
||||
}
|
||||
c.mu.Lock()
|
||||
c.closed = true
|
||||
pending := c.pending
|
||||
c.pending = map[int64]chan stdioResult{}
|
||||
c.mu.Unlock()
|
||||
for _, ch := range pending {
|
||||
ch <- stdioResult{err: fatal}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *stdioClient) notify(method string) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.closed {
|
||||
closed := c.closed
|
||||
c.mu.Unlock()
|
||||
if closed {
|
||||
return fmt.Errorf("mcp client closed")
|
||||
}
|
||||
body, err := json.Marshal(rpcNotification{JSONRPC: "2.0", Method: method})
|
||||
@@ -193,6 +239,8 @@ func (c *stdioClient) notify(method string) error {
|
||||
return err
|
||||
}
|
||||
body = append(body, '\n')
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
_, err = c.stdin.Write(body)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/session"
|
||||
"agent-native/runtime/tools"
|
||||
)
|
||||
|
||||
type countingTool struct {
|
||||
name string
|
||||
calls int
|
||||
hook func()
|
||||
}
|
||||
|
||||
func (c *countingTool) Name() string { return c.name }
|
||||
func (c *countingTool) Description() string { return c.name }
|
||||
func (c *countingTool) Parameters() json.RawMessage {
|
||||
return json.RawMessage(`{"type":"object","properties":{}}`)
|
||||
}
|
||||
func (c *countingTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
|
||||
c.calls++
|
||||
if c.hook != nil {
|
||||
c.hook()
|
||||
}
|
||||
return tools.Result{Content: "ok"}, nil
|
||||
}
|
||||
|
||||
func TestR2TokenBudgetSkipsTools(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,
|
||||
Budget: &Budget{MaxPromptTokens: 10},
|
||||
Stream: usageScript(llm.Usage{PromptTokens: 500},
|
||||
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
|
||||
ID: "c1", Name: "list_dir", Arguments: `{}`,
|
||||
}}},
|
||||
),
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tool.calls != 0 {
|
||||
t.Fatalf("tool ran after token budget: %d", tool.calls)
|
||||
}
|
||||
if res.Err == nil || !strings.Contains(res.Err.Error(), "prompt tokens") {
|
||||
t.Fatalf("budget error: %v", res.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestR2StopReplacementIsPersisted(t *testing.T) {
|
||||
store, err := session.Open(filepath.Join(t.TempDir(), "t.jsonl"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = Run(context.Background(), Config{
|
||||
Model: "fake",
|
||||
MaxTurns: 3,
|
||||
Session: store,
|
||||
Stream: script(
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "draft", StopReason: llm.StopEnd},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "final", StopReason: llm.StopEnd},
|
||||
),
|
||||
Stop: func(last llm.Message, info StopInfo) StopDecision {
|
||||
if info.Continuations == 0 {
|
||||
return StopDecision{Continue: true, Replacement: "corrected-claim"}
|
||||
}
|
||||
return StopDecision{}
|
||||
},
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
msgs, err := store.Load()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var saw bool
|
||||
for _, m := range msgs {
|
||||
if m.Content == "corrected-claim" {
|
||||
saw = true
|
||||
}
|
||||
}
|
||||
if !saw {
|
||||
t.Fatalf("replacement missing from session: %+v", msgs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestR2SequentialToolsStopAfterCancel(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
first := &countingTool{name: "first", hook: cancel}
|
||||
second := &countingTool{name: "second"}
|
||||
reg := tools.NewRegistry(t.TempDir())
|
||||
reg.Register(first)
|
||||
reg.Register(second)
|
||||
res, err := Run(ctx, Config{
|
||||
Model: "fake",
|
||||
Tools: reg,
|
||||
Stream: script(
|
||||
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{
|
||||
{ID: "1", Name: "first", Arguments: `{}`},
|
||||
{ID: "2", Name: "second", Arguments: `{}`},
|
||||
}},
|
||||
llm.Message{Role: llm.RoleAssistant, Content: "after", StopReason: llm.StopEnd},
|
||||
),
|
||||
}, llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if first.calls != 1 || second.calls != 0 {
|
||||
t.Fatalf("calls first=%d second=%d", first.calls, second.calls)
|
||||
}
|
||||
var paired int
|
||||
for _, m := range res.Messages {
|
||||
if m.Role == llm.RoleTool {
|
||||
paired++
|
||||
}
|
||||
}
|
||||
if paired != 2 {
|
||||
t.Fatalf("tool results = %d, want a pair for each call", paired)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
package resources
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ContextFiles walks from dir toward the filesystem root and concatenates
|
||||
// AGENTS.md and CLAUDE.md (nearest first). SYSTEM.md and APPEND_SYSTEM.md
|
||||
// are read from dir only.
|
||||
func ContextFiles(dir string) (string, error) {
|
||||
var parts []string
|
||||
cur := dir
|
||||
for {
|
||||
for _, name := range []string{"AGENTS.md", "CLAUDE.md"} {
|
||||
raw, err := os.ReadFile(filepath.Join(cur, name))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
parts = append(parts, "# "+name+" ("+cur+")\n"+string(raw))
|
||||
}
|
||||
parent := filepath.Dir(cur)
|
||||
if parent == cur {
|
||||
break
|
||||
}
|
||||
cur = parent
|
||||
}
|
||||
for _, name := range []string{"SYSTEM.md", "APPEND_SYSTEM.md"} {
|
||||
raw, err := os.ReadFile(filepath.Join(dir, name))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
parts = append(parts, "# "+name+"\n"+string(raw))
|
||||
}
|
||||
return strings.Join(parts, "\n\n"), nil
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package resources
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Expand replaces {{var}} placeholders and /name includes.
|
||||
// An include reads <dir>/<name>.md. Unknown variables are left intact.
|
||||
func Expand(text, dir string, vars map[string]string) (string, error) {
|
||||
out := text
|
||||
for key, value := range vars {
|
||||
out = strings.ReplaceAll(out, "{{"+key+"}}", value)
|
||||
}
|
||||
for {
|
||||
i := strings.Index(out, "/")
|
||||
if i < 0 {
|
||||
break
|
||||
}
|
||||
rest := out[i+1:]
|
||||
j := 0
|
||||
for j < len(rest) && isName(rest[j]) {
|
||||
j++
|
||||
}
|
||||
if j == 0 || (i > 0 && out[i-1] != ' ' && out[i-1] != '\n') {
|
||||
break
|
||||
}
|
||||
name := rest[:j]
|
||||
raw, err := os.ReadFile(filepath.Join(dir, name+".md"))
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
out = out[:i] + string(raw) + out[i+1+j:]
|
||||
}
|
||||
for key, value := range vars {
|
||||
out = strings.ReplaceAll(out, "{{"+key+"}}", value)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func isName(b byte) bool {
|
||||
return b == '-' || b == '_' || (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') || (b >= '0' && b <= '9')
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package resources
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"agent-native/runtime/skills"
|
||||
"agent-native/runtime/tools"
|
||||
)
|
||||
|
||||
func TestContextAndPromptAndTrust(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
proj := filepath.Join(root, "proj")
|
||||
if err := os.MkdirAll(proj, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(root, "AGENTS.md"), []byte("root-agents"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(proj, "SYSTEM.md"), []byte("sys"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
text, err := ContextFiles(proj)
|
||||
if err != nil || !strings.Contains(text, "root-agents") || !strings.Contains(text, "sys") {
|
||||
t.Fatalf("context: %s %v", text, err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(proj, "greet.md"), []byte("hello {{name}}"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := Expand("say /greet", proj, map[string]string{"name": "Ada"})
|
||||
if err != nil || got != "say hello Ada" {
|
||||
t.Fatalf("expand: %q %v", got, err)
|
||||
}
|
||||
if ProjectTrusted(TrustAsk, false, false) {
|
||||
t.Fatal("ask is not trusted without approve")
|
||||
}
|
||||
if !ProjectTrusted(TrustNever, true, false) {
|
||||
t.Fatal("approve overrides never")
|
||||
}
|
||||
if ProjectTrusted(TrustAlways, false, true) {
|
||||
t.Fatal("no-approve wins")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPiAliasesFilterAndSkillCommand(t *testing.T) {
|
||||
reg := tools.WorkspaceTools(t.TempDir())
|
||||
reg.RegisterPiAliases()
|
||||
if !reg.Has("read") || !reg.Has("ls") || !reg.Has("read_file") {
|
||||
t.Fatal("aliases missing")
|
||||
}
|
||||
if err := reg.Filter(tools.FilterOptions{Exclude: []string{"read"}, NoBuiltin: false}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if reg.Has("read") || !reg.Has("read_file") {
|
||||
t.Fatal("exclude should drop only the alias")
|
||||
}
|
||||
dir := t.TempDir()
|
||||
skillDir := filepath.Join(dir, "review")
|
||||
if err := os.MkdirAll(skillDir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte("be thorough"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
store := skills.NewStore(dir)
|
||||
out, err := store.ExpandCommands("use /skill:review now")
|
||||
if err != nil || !strings.Contains(out, "be thorough") {
|
||||
t.Fatalf("skill expand: %q %v", out, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
package resources
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
// Trust is how a project directory is treated in non-interactive runs.
|
||||
type Trust string
|
||||
|
||||
const (
|
||||
TrustAsk Trust = "ask"
|
||||
TrustAlways Trust = "always"
|
||||
TrustNever Trust = "never"
|
||||
)
|
||||
|
||||
// Settings is the global plus project configuration.
|
||||
type Settings struct {
|
||||
Trust Trust `json:"trust,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
Thinking string `json:"thinking,omitempty"`
|
||||
Tools []string `json:"tools,omitempty"`
|
||||
Exclude []string `json:"exclude,omitempty"`
|
||||
Vars map[string]string `json:"vars,omitempty"`
|
||||
}
|
||||
|
||||
// LoadSettings reads global (optional) then project settings.
|
||||
// Project keys override global keys. A missing file is an empty settings.
|
||||
func LoadSettings(globalPath, projectPath string) (Settings, error) {
|
||||
global, err := readSettings(globalPath)
|
||||
if err != nil {
|
||||
return Settings{}, err
|
||||
}
|
||||
project, err := readSettings(projectPath)
|
||||
if err != nil {
|
||||
return Settings{}, err
|
||||
}
|
||||
return mergeSettings(global, project), nil
|
||||
}
|
||||
|
||||
// ProjectTrusted reports whether project resources may be loaded.
|
||||
// ask and never refuse in non-interactive mode. approve forces trust.
|
||||
func ProjectTrusted(trust Trust, approve, noApprove bool) bool {
|
||||
if noApprove {
|
||||
return false
|
||||
}
|
||||
if approve {
|
||||
return true
|
||||
}
|
||||
return trust == TrustAlways
|
||||
}
|
||||
|
||||
func readSettings(path string) (Settings, error) {
|
||||
if path == "" {
|
||||
return Settings{}, nil
|
||||
}
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return Settings{}, nil
|
||||
}
|
||||
return Settings{}, err
|
||||
}
|
||||
var s Settings
|
||||
if err := json.Unmarshal(raw, &s); err != nil {
|
||||
return Settings{}, err
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func mergeSettings(base, over Settings) Settings {
|
||||
if over.Trust != "" {
|
||||
base.Trust = over.Trust
|
||||
}
|
||||
if over.Model != "" {
|
||||
base.Model = over.Model
|
||||
}
|
||||
if over.Thinking != "" {
|
||||
base.Thinking = over.Thinking
|
||||
}
|
||||
if over.Tools != nil {
|
||||
base.Tools = over.Tools
|
||||
}
|
||||
if over.Exclude != nil {
|
||||
base.Exclude = over.Exclude
|
||||
}
|
||||
if len(over.Vars) > 0 {
|
||||
if base.Vars == nil {
|
||||
base.Vars = map[string]string{}
|
||||
}
|
||||
for k, v := range over.Vars {
|
||||
base.Vars[k] = v
|
||||
}
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
// DefaultPaths returns the global and project settings paths for a cwd.
|
||||
func DefaultPaths(home, cwd string) (string, string) {
|
||||
return filepath.Join(home, ".agent-native", "settings.json"), filepath.Join(cwd, ".agent-native", "settings.json")
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Info describes one session file in a directory.
|
||||
type Info struct {
|
||||
ID string
|
||||
Path string
|
||||
Name string
|
||||
CWD string
|
||||
Leaf string
|
||||
Updated string
|
||||
}
|
||||
|
||||
// List returns session files in dir, newest first.
|
||||
func List(dir string) ([]Info, error) {
|
||||
ents, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
var out []Info
|
||||
for _, ent := range ents {
|
||||
name := ent.Name()
|
||||
if ent.IsDir() || !strings.HasSuffix(name, ".jsonl") {
|
||||
continue
|
||||
}
|
||||
path := filepath.Join(dir, name)
|
||||
st, err := Open(path)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
info, err := st.Info()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, info)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Updated > out[j].Updated })
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Continue opens the newest session whose cwd matches. A second return
|
||||
// value of false means there is nothing to continue.
|
||||
func Continue(dir, cwd string) (*Store, bool, error) {
|
||||
list, err := List(dir)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
for _, info := range list {
|
||||
if cwd == "" || info.CWD == cwd {
|
||||
st, err := Open(info.Path)
|
||||
return st, err == nil, err
|
||||
}
|
||||
}
|
||||
return nil, false, nil
|
||||
}
|
||||
|
||||
// Resume opens a session by id (filename stem) or by name.
|
||||
func Resume(dir, idOrName string) (*Store, error) {
|
||||
list, err := List(dir)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, info := range list {
|
||||
if info.ID == idOrName || info.Name == idOrName {
|
||||
return Open(info.Path)
|
||||
}
|
||||
}
|
||||
return nil, os.ErrNotExist
|
||||
}
|
||||
|
||||
// SetCWD records the working directory used by Continue.
|
||||
func (s *Store) SetCWD(cwd string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.writeMetaKeyLocked("cwd", cwd)
|
||||
}
|
||||
|
||||
// Info reads the catalog fields for this file.
|
||||
func (s *Store) Info() (Info, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
entries, err := s.readEntriesLocked()
|
||||
if err != nil {
|
||||
return Info{}, err
|
||||
}
|
||||
meta := s.readMetaLocked()
|
||||
updated := ""
|
||||
if len(entries) > 0 {
|
||||
updated = entries[len(entries)-1].Timestamp
|
||||
}
|
||||
id := strings.TrimSuffix(filepath.Base(s.path), ".jsonl")
|
||||
return Info{
|
||||
ID: id, Path: s.path, Name: meta["name"], CWD: meta["cwd"],
|
||||
Leaf: s.leafLocked(entries), Updated: updated,
|
||||
}, nil
|
||||
}
|
||||
+38
-42
@@ -3,6 +3,7 @@ package session
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
@@ -66,73 +67,68 @@ func (s *Store) ReadMeta() map[string]string {
|
||||
return meta
|
||||
}
|
||||
|
||||
// Compact 把转录改写为「系统消息 + 摘要消息 + 最近 keepLastN 条」。
|
||||
// 摘要由 summarize 生成(由适配层注入,通常是一次无工具模型调用);
|
||||
// 改写原子落盘(临时文件 + rename)。
|
||||
// 返回是否发生了压缩(消息数不足/已是紧凑形态时返回 false 且不动文件)。
|
||||
// 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.
|
||||
func (s *Store) Compact(keepLastN int, summarize func(messages []llm.Message) (string, error)) (bool, error) {
|
||||
messages, err := s.Load()
|
||||
s.mu.Lock()
|
||||
entries, err := s.readEntriesLocked()
|
||||
if err != nil {
|
||||
s.mu.Unlock()
|
||||
return false, err
|
||||
}
|
||||
if len(messages) <= keepLastN+1 { // 系统消息 + 尾部,无可压缩头部
|
||||
leaf := s.leafLocked(entries)
|
||||
generation := len(entries)
|
||||
s.mu.Unlock()
|
||||
path := pathTo(entries, leaf)
|
||||
var messages []llm.Message
|
||||
var msgEntries []Entry
|
||||
for _, e := range path {
|
||||
if e.Type == TypeMessage && e.Message != nil {
|
||||
messages = append(messages, *e.Message)
|
||||
msgEntries = append(msgEntries, e)
|
||||
}
|
||||
}
|
||||
if len(messages) <= keepLastN+1 {
|
||||
return false, nil
|
||||
}
|
||||
var system *llm.Message
|
||||
rest := messages
|
||||
restEntries := msgEntries
|
||||
if len(rest) > 0 && rest[0].Role == llm.RoleSystem {
|
||||
copyOf := rest[0]
|
||||
system = ©Of
|
||||
rest = rest[1:]
|
||||
restEntries = restEntries[1:]
|
||||
}
|
||||
tailStart := len(rest) - keepLastN
|
||||
if tailStart <= 0 {
|
||||
return false, nil
|
||||
}
|
||||
tailStart = alignToolTail(rest, tailStart)
|
||||
head := rest[:tailStart]
|
||||
tail := rest[tailStart:]
|
||||
summary, err := summarize(head)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
summaryMsg := llm.Message{
|
||||
Role: llm.RoleUser,
|
||||
Content: "[conversation summary]\n" + summary + "\n[recent messages follow]",
|
||||
var tailIDs []string
|
||||
for _, e := range restEntries[tailStart:] {
|
||||
tailIDs = append(tailIDs, e.ID)
|
||||
}
|
||||
var rebuilt []llm.Message
|
||||
if system != nil {
|
||||
rebuilt = append(rebuilt, *system)
|
||||
}
|
||||
rebuilt = append(rebuilt, summaryMsg)
|
||||
rebuilt = append(rebuilt, tail...)
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
tmp := s.path + ".compact.tmp"
|
||||
f, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
|
||||
now, err := s.readEntriesLocked()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
w := bufio.NewWriter(f)
|
||||
for _, m := range rebuilt {
|
||||
raw, err := json.Marshal(m)
|
||||
if err != nil {
|
||||
_ = f.Close()
|
||||
_ = os.Remove(tmp)
|
||||
return false, err
|
||||
}
|
||||
w.Write(raw)
|
||||
w.WriteByte('\n')
|
||||
if len(now) != generation {
|
||||
return false, fmt.Errorf("compaction conflict: session changed during summary")
|
||||
}
|
||||
if err := w.Flush(); err != nil {
|
||||
_ = f.Close()
|
||||
_ = os.Remove(tmp)
|
||||
return false, err
|
||||
_, err = s.appendEntryLocked(Entry{Type: TypeCompaction, Summary: summary, TailIDs: tailIDs})
|
||||
return err == nil, err
|
||||
}
|
||||
|
||||
// alignToolTail 让保留尾巴从一条非 tool 消息开始,避免把 tool_calls 和结果拆开。
|
||||
func alignToolTail(messages []llm.Message, tailStart int) int {
|
||||
for tailStart > 0 && messages[tailStart].Role == llm.RoleTool {
|
||||
tailStart--
|
||||
}
|
||||
_ = f.Close()
|
||||
if err := os.Rename(tmp, s.path); err != nil {
|
||||
_ = os.Remove(tmp)
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
return tailStart
|
||||
}
|
||||
|
||||
@@ -114,9 +114,12 @@ func TestCompactRewritesTranscript(t *testing.T) {
|
||||
if msgs[2].Content != "reply" || msgs[3].Content != "earlyj" || msgs[4].Content != "reply" {
|
||||
t.Fatalf("tail order wrong: %+v", msgs[2:])
|
||||
}
|
||||
// 无 tmp 残留。
|
||||
if _, err := os.Stat(s.Path() + ".compact.tmp"); !os.IsNotExist(err) {
|
||||
t.Fatal("compact tmp file must not remain")
|
||||
entries, err := s.Entries()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(entries) <= 5 {
|
||||
t.Fatalf("compaction must append, not rewrite; got %d entries", len(entries))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"html"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ExportHTML writes the active branch as a standalone HTML document.
|
||||
func (s *Store) ExportHTML(path string) error {
|
||||
msgs, err := s.Load()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("<!DOCTYPE html><html><head><meta charset=\"utf-8\"><title>")
|
||||
title := s.Name()
|
||||
if title == "" {
|
||||
title = "session"
|
||||
}
|
||||
b.WriteString(html.EscapeString(title))
|
||||
b.WriteString("</title></head><body>\n")
|
||||
for _, m := range msgs {
|
||||
fmt.Fprintf(&b, "<article data-role=\"%s\"><h2>%s</h2><pre>%s</pre></article>\n",
|
||||
html.EscapeString(string(m.Role)),
|
||||
html.EscapeString(string(m.Role)),
|
||||
html.EscapeString(m.Content))
|
||||
}
|
||||
b.WriteString("</body></html>\n")
|
||||
return os.WriteFile(path, []byte(b.String()), 0o644)
|
||||
}
|
||||
+12
-44
@@ -1,17 +1,16 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
)
|
||||
|
||||
// Store is an append-only JSONL transcript. Each line is one llm.Message.
|
||||
// Store is an append-only JSONL transcript. New sessions are a tree of
|
||||
// entries (id / parentId). Load still returns the model-visible message
|
||||
// path, and files written as one llm.Message per line remain readable.
|
||||
type Store struct {
|
||||
path string
|
||||
mu sync.Mutex
|
||||
@@ -31,56 +30,25 @@ func Open(path string) (*Store, error) {
|
||||
|
||||
func (s *Store) Path() string { return s.path }
|
||||
|
||||
// Append records one message on the active branch.
|
||||
func (s *Store) Append(msg llm.Message) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
raw, err := json.Marshal(msg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f, err := os.OpenFile(s.path, os.O_APPEND|os.O_WRONLY, 0o644)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
raw = append(raw, '\n')
|
||||
_, err = f.Write(raw)
|
||||
_, err := s.appendEntryLocked(Entry{Type: TypeMessage, Message: &msg})
|
||||
return err
|
||||
}
|
||||
|
||||
// Load 全量读取转录。崩溃残留容错:Append 恒以换行收尾,因此无换行结尾的
|
||||
// 尾部若不是完整 JSON 即为崩溃半行,丢弃而不是让整个会话不可用;
|
||||
// 行中间的坏行仍报错(那是文件损坏,不是崩溃残留)。
|
||||
// Load returns the model-visible messages on the active branch.
|
||||
// A compaction entry replaces earlier turns with its summary plus the
|
||||
// retained tail; the file itself is never rewritten.
|
||||
// Crash residue: a trailing line without a newline that is not valid JSON
|
||||
// is dropped. A bad line in the middle is still an error.
|
||||
func (s *Store) Load() ([]llm.Message, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
raw, err := os.ReadFile(s.path)
|
||||
entries, err := s.readEntriesLocked()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
text := string(raw)
|
||||
if text != "" && !strings.HasSuffix(text, "\n") {
|
||||
if idx := strings.LastIndexByte(text, '\n'); idx >= 0 {
|
||||
if !json.Valid([]byte(text[idx+1:])) {
|
||||
text = text[:idx+1]
|
||||
}
|
||||
} else if !json.Valid([]byte(text)) {
|
||||
text = ""
|
||||
}
|
||||
}
|
||||
var out []llm.Message
|
||||
sc := bufio.NewScanner(strings.NewReader(text))
|
||||
sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
|
||||
for sc.Scan() {
|
||||
line := sc.Bytes()
|
||||
if len(line) == 0 {
|
||||
continue
|
||||
}
|
||||
var msg llm.Message
|
||||
if err := json.Unmarshal(line, &msg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, msg)
|
||||
}
|
||||
return out, sc.Err()
|
||||
return visibleMessages(entries, s.leafLocked(entries)), nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
)
|
||||
|
||||
func TestR2CompactKeepsToolPair(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: "ask"},
|
||||
{Role: llm.RoleAssistant, Content: "", ToolCalls: []llm.ToolCall{{ID: "c", Name: "read"}}},
|
||||
{Role: llm.RoleTool, ToolCallID: "c", Content: "file-body"},
|
||||
}
|
||||
for _, m := range msgs {
|
||||
if err := s.Append(m); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ok, err := s.Compact(1, func(messages []llm.Message) (string, error) {
|
||||
return "summary", nil
|
||||
})
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("compact: %v %v", ok, err)
|
||||
}
|
||||
loaded, err := s.Load()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var sawTool bool
|
||||
for _, m := range loaded {
|
||||
if m.Role == llm.RoleTool && m.Content == "file-body" {
|
||||
sawTool = true
|
||||
}
|
||||
}
|
||||
if !sawTool {
|
||||
t.Fatalf("tool result dropped: %+v", loaded)
|
||||
}
|
||||
}
|
||||
|
||||
func TestR2CompactRejectsAppendDuringSummary(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)
|
||||
}
|
||||
}
|
||||
_, err = s.Compact(1, func(messages []llm.Message) (string, error) {
|
||||
return "summary", s.Append(llm.Message{Role: llm.RoleUser, Content: "raced"})
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "compaction conflict") {
|
||||
t.Fatalf("want conflict, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestR2AppendVisibleAfterReopen(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "s.jsonl")
|
||||
s, err := Open(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "durable"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
again, err := Open(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
msgs, err := again.Load()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(msgs) != 1 || msgs[0].Content != "durable" {
|
||||
t.Fatalf("reopen: %+v", msgs)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,401 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
)
|
||||
|
||||
// EntryType is the kind of one JSONL record.
|
||||
type EntryType string
|
||||
|
||||
const (
|
||||
TypeHeader EntryType = "header"
|
||||
TypeMessage EntryType = "message"
|
||||
TypeModelChange EntryType = "model_change"
|
||||
TypeThinkingLevelChange EntryType = "thinking_level_change"
|
||||
TypeCompaction EntryType = "compaction"
|
||||
TypeLabel EntryType = "label"
|
||||
TypeCustom EntryType = "custom"
|
||||
TypeSessionInfo EntryType = "session_info"
|
||||
)
|
||||
|
||||
// Entry is one node in the session tree. Files written before the tree
|
||||
// format (a bare llm.Message per line) are loaded as message entries.
|
||||
type Entry struct {
|
||||
ID string `json:"id"`
|
||||
ParentID string `json:"parentId,omitempty"`
|
||||
Timestamp string `json:"timestamp"`
|
||||
Type EntryType `json:"type"`
|
||||
Message *llm.Message `json:"message,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
ThinkingLevel string `json:"thinkingLevel,omitempty"`
|
||||
Label string `json:"label,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
CustomType string `json:"customType,omitempty"`
|
||||
Data json.RawMessage `json:"data,omitempty"`
|
||||
Summary string `json:"summary,omitempty"`
|
||||
TailIDs []string `json:"tailIds,omitempty"`
|
||||
CWD string `json:"cwd,omitempty"`
|
||||
}
|
||||
|
||||
// Entries returns every record in file order, including branches that are
|
||||
// not on the active path.
|
||||
func (s *Store) Entries() ([]Entry, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.readEntriesLocked()
|
||||
}
|
||||
|
||||
// Leaf is the id of the current branch tip. Empty when the file has no entries.
|
||||
func (s *Store) Leaf() (string, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
entries, err := s.readEntriesLocked()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return s.leafLocked(entries), nil
|
||||
}
|
||||
|
||||
// SetLeaf moves the active branch to an existing entry.
|
||||
func (s *Store) SetLeaf(id string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
entries, err := s.readEntriesLocked()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !hasID(entries, id) {
|
||||
return fmt.Errorf("session: unknown entry %s", id)
|
||||
}
|
||||
return s.writeMetaKeyLocked("leaf", id)
|
||||
}
|
||||
|
||||
// AppendEntry adds a non-message record (label, model change, custom, …)
|
||||
// as a child of the current leaf and moves the leaf to it.
|
||||
func (s *Store) AppendEntry(e Entry) (Entry, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.appendEntryLocked(e)
|
||||
}
|
||||
|
||||
// Fork copies the tree into a new file and sets that file's leaf to atID.
|
||||
// The source file is not modified. atID empty means the current leaf.
|
||||
func (s *Store) Fork(dest, atID string) (*Store, error) {
|
||||
return s.copyTo(dest, atID, true)
|
||||
}
|
||||
|
||||
// Clone copies the tree and keeps the same leaf.
|
||||
func (s *Store) Clone(dest string) (*Store, error) {
|
||||
return s.copyTo(dest, "", false)
|
||||
}
|
||||
|
||||
func (s *Store) copyTo(dest, atID string, moveLeaf bool) (*Store, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
entries, err := s.readEntriesLocked()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := copyFile(s.path, dest); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out, err := Open(dest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
meta := s.readMetaLocked()
|
||||
if meta == nil {
|
||||
meta = map[string]string{}
|
||||
}
|
||||
leaf := s.leafLocked(entries)
|
||||
if moveLeaf && atID != "" {
|
||||
if !hasID(entries, atID) {
|
||||
return nil, fmt.Errorf("session: unknown entry %s", atID)
|
||||
}
|
||||
leaf = atID
|
||||
}
|
||||
if leaf != "" {
|
||||
meta["leaf"] = leaf
|
||||
}
|
||||
if err := out.SaveMeta(meta); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// SetName records a human name for continue/resume.
|
||||
func (s *Store) SetName(name string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if err := s.writeMetaKeyLocked("name", name); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := s.appendEntryLocked(Entry{Type: TypeSessionInfo, Name: name})
|
||||
return err
|
||||
}
|
||||
|
||||
// Name returns the session name, or empty.
|
||||
func (s *Store) Name() string {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.readMetaLocked()["name"]
|
||||
}
|
||||
|
||||
func (s *Store) appendEntryLocked(e Entry) (Entry, error) {
|
||||
entries, err := s.readEntriesLocked()
|
||||
if err != nil {
|
||||
return Entry{}, err
|
||||
}
|
||||
if e.ID == "" {
|
||||
e.ID = newID()
|
||||
}
|
||||
if e.Timestamp == "" {
|
||||
e.Timestamp = time.Now().UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
if e.Type == "" {
|
||||
e.Type = TypeCustom
|
||||
}
|
||||
if e.ParentID == "" {
|
||||
e.ParentID = s.leafLocked(entries)
|
||||
}
|
||||
raw, err := json.Marshal(e)
|
||||
if err != nil {
|
||||
return Entry{}, err
|
||||
}
|
||||
f, err := os.OpenFile(s.path, os.O_APPEND|os.O_WRONLY, 0o644)
|
||||
if err != nil {
|
||||
return Entry{}, err
|
||||
}
|
||||
defer f.Close()
|
||||
if _, err := f.Write(append(raw, '\n')); err != nil {
|
||||
return Entry{}, err
|
||||
}
|
||||
// 单写者:锁只覆盖本进程。Sync 把这一行落到存储,崩溃后重开仍能读到。
|
||||
if err := f.Sync(); err != nil {
|
||||
return Entry{}, err
|
||||
}
|
||||
if err := s.writeMetaKeyLocked("leaf", e.ID); err != nil {
|
||||
return Entry{}, err
|
||||
}
|
||||
return e, nil
|
||||
}
|
||||
|
||||
func (s *Store) readEntriesLocked() ([]Entry, error) {
|
||||
raw, err := os.ReadFile(s.path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
text := string(raw)
|
||||
if text != "" && !strings.HasSuffix(text, "\n") {
|
||||
if idx := strings.LastIndexByte(text, '\n'); idx >= 0 {
|
||||
if !json.Valid([]byte(text[idx+1:])) {
|
||||
text = text[:idx+1]
|
||||
}
|
||||
} else if !json.Valid([]byte(text)) {
|
||||
text = ""
|
||||
}
|
||||
}
|
||||
var out []Entry
|
||||
sc := bufio.NewScanner(strings.NewReader(text))
|
||||
sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
|
||||
legacyParent := ""
|
||||
n := 0
|
||||
for sc.Scan() {
|
||||
line := sc.Bytes()
|
||||
if len(line) == 0 {
|
||||
continue
|
||||
}
|
||||
var probe map[string]json.RawMessage
|
||||
if err := json.Unmarshal(line, &probe); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if _, ok := probe["type"]; ok {
|
||||
var e Entry
|
||||
if err := json.Unmarshal(line, &e); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, e)
|
||||
legacyParent = e.ID
|
||||
n++
|
||||
continue
|
||||
}
|
||||
var msg llm.Message
|
||||
if err := json.Unmarshal(line, &msg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
id := fmt.Sprintf("legacy-%d", n)
|
||||
out = append(out, Entry{
|
||||
ID: id, ParentID: legacyParent, Type: TypeMessage, Message: &msg,
|
||||
})
|
||||
legacyParent = id
|
||||
n++
|
||||
}
|
||||
return out, sc.Err()
|
||||
}
|
||||
|
||||
func (s *Store) leafLocked(entries []Entry) string {
|
||||
if leaf := s.readMetaLocked()["leaf"]; leaf != "" && hasID(entries, leaf) {
|
||||
return leaf
|
||||
}
|
||||
if len(entries) == 0 {
|
||||
return ""
|
||||
}
|
||||
return entries[len(entries)-1].ID
|
||||
}
|
||||
|
||||
func visibleMessages(entries []Entry, leaf string) []llm.Message {
|
||||
path := pathTo(entries, leaf)
|
||||
var last *Entry
|
||||
for i := range path {
|
||||
if path[i].Type == TypeCompaction {
|
||||
last = &path[i]
|
||||
}
|
||||
}
|
||||
if last == nil {
|
||||
var msgs []llm.Message
|
||||
for _, e := range path {
|
||||
if e.Type == TypeMessage && e.Message != nil {
|
||||
msgs = append(msgs, *e.Message)
|
||||
}
|
||||
}
|
||||
return msgs
|
||||
}
|
||||
byID := map[string]Entry{}
|
||||
for _, e := range entries {
|
||||
byID[e.ID] = e
|
||||
}
|
||||
var msgs []llm.Message
|
||||
for _, e := range path {
|
||||
if e.ID == last.ID {
|
||||
break
|
||||
}
|
||||
if e.Type == TypeMessage && e.Message != nil && e.Message.Role == llm.RoleSystem {
|
||||
msgs = append(msgs, *e.Message)
|
||||
}
|
||||
}
|
||||
msgs = append(msgs, llm.Message{
|
||||
Role: llm.RoleUser,
|
||||
Content: "[conversation summary]\n" + last.Summary + "\n[recent messages follow]",
|
||||
})
|
||||
for _, id := range last.TailIDs {
|
||||
e, ok := byID[id]
|
||||
if ok && e.Message != nil {
|
||||
msgs = append(msgs, *e.Message)
|
||||
}
|
||||
}
|
||||
after := false
|
||||
for _, e := range path {
|
||||
if e.ID == last.ID {
|
||||
after = true
|
||||
continue
|
||||
}
|
||||
if after && e.Type == TypeMessage && e.Message != nil {
|
||||
msgs = append(msgs, *e.Message)
|
||||
}
|
||||
}
|
||||
return msgs
|
||||
}
|
||||
|
||||
func pathTo(entries []Entry, leaf string) []Entry {
|
||||
byID := map[string]Entry{}
|
||||
for _, e := range entries {
|
||||
byID[e.ID] = e
|
||||
}
|
||||
var rev []Entry
|
||||
seen := map[string]bool{}
|
||||
for id := leaf; id != "" && !seen[id]; {
|
||||
e, ok := byID[id]
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
seen[id] = true
|
||||
rev = append(rev, e)
|
||||
id = e.ParentID
|
||||
}
|
||||
for i, j := 0, len(rev)-1; i < j; i, j = i+1, j-1 {
|
||||
rev[i], rev[j] = rev[j], rev[i]
|
||||
}
|
||||
return rev
|
||||
}
|
||||
|
||||
func hasID(entries []Entry, id string) bool {
|
||||
for _, e := range entries {
|
||||
if e.ID == id {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *Store) readMetaLocked() map[string]string {
|
||||
raw, err := os.ReadFile(s.path + ".meta.json")
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
var meta map[string]string
|
||||
if err := json.Unmarshal(raw, &meta); err != nil {
|
||||
return nil
|
||||
}
|
||||
return meta
|
||||
}
|
||||
|
||||
func (s *Store) writeMetaKeyLocked(key, value string) error {
|
||||
meta := s.readMetaLocked()
|
||||
if meta == nil {
|
||||
meta = map[string]string{}
|
||||
}
|
||||
meta[key] = value
|
||||
raw, err := json.MarshalIndent(meta, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tmp := s.path + ".meta.json.tmp"
|
||||
if err := os.WriteFile(tmp, raw, 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, s.path+".meta.json")
|
||||
}
|
||||
|
||||
func newID() string {
|
||||
var b [8]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
return fmt.Sprintf("%d", time.Now().UnixNano())
|
||||
}
|
||||
return hex.EncodeToString(b[:])
|
||||
}
|
||||
|
||||
func copyFile(src, dst string) error {
|
||||
in, err := os.Open(src)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer in.Close()
|
||||
if err := os.MkdirAll(dirOf(dst), 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer out.Close()
|
||||
_, err = io.Copy(out, in)
|
||||
return err
|
||||
}
|
||||
|
||||
func dirOf(path string) string {
|
||||
i := strings.LastIndex(path, "/")
|
||||
if i < 0 {
|
||||
return "."
|
||||
}
|
||||
return path[:i]
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
)
|
||||
|
||||
func TestTreeForkCloneAndLeaf(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
s, err := Open(filepath.Join(dir, "a.jsonl"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "one"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Append(llm.Message{Role: llm.RoleAssistant, Content: "two"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
entries, err := s.Entries()
|
||||
if err != nil || len(entries) != 2 || entries[1].ParentID != entries[0].ID {
|
||||
t.Fatalf("tree links: %+v %v", entries, err)
|
||||
}
|
||||
fork, err := s.Fork(filepath.Join(dir, "fork.jsonl"), entries[0].ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := fork.Append(llm.Message{Role: llm.RoleUser, Content: "branch"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
msgs, err := fork.Load()
|
||||
if err != nil || len(msgs) != 2 || msgs[1].Content != "branch" {
|
||||
t.Fatalf("fork path: %+v %v", msgs, err)
|
||||
}
|
||||
orig, _ := s.Load()
|
||||
if len(orig) != 2 || orig[1].Content != "two" {
|
||||
t.Fatalf("source branch changed: %+v", orig)
|
||||
}
|
||||
clone, err := s.Clone(filepath.Join(dir, "clone.jsonl"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cmsgs, _ := clone.Load()
|
||||
if len(cmsgs) != 2 || cmsgs[1].Content != "two" {
|
||||
t.Fatalf("clone: %+v", cmsgs)
|
||||
}
|
||||
if err := s.SetLeaf(entries[0].ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
switched, _ := s.Load()
|
||||
if len(switched) != 1 || switched[0].Content != "one" {
|
||||
t.Fatalf("set leaf: %+v", switched)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLegacyTranscriptStillLoads(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "old.jsonl")
|
||||
body := "{\"role\":\"system\",\"content\":\"sys\"}\n{\"role\":\"user\",\"content\":\"hi\"}\n"
|
||||
if err := os.WriteFile(path, []byte(body), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s, err := Open(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
msgs, err := s.Load()
|
||||
if err != nil || len(msgs) != 2 || msgs[1].Content != "hi" {
|
||||
t.Fatalf("legacy: %+v %v", msgs, err)
|
||||
}
|
||||
if err := s.Append(llm.Message{Role: llm.RoleAssistant, Content: "next"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
msgs, err = s.Load()
|
||||
if err != nil || len(msgs) != 3 || msgs[2].Content != "next" {
|
||||
t.Fatalf("legacy append: %+v %v", msgs, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCatalogContinueResumeAndHTML(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
s, err := Open(filepath.Join(dir, "sess.jsonl"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.SetCWD("/work"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.SetName("notes"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "hello"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, ok, err := Continue(dir, "/work")
|
||||
if err != nil || !ok || got.Path() != s.Path() {
|
||||
t.Fatalf("continue: %v %v %v", got, ok, err)
|
||||
}
|
||||
resumed, err := Resume(dir, "notes")
|
||||
if err != nil || resumed.Name() != "notes" {
|
||||
t.Fatalf("resume: %v %v", resumed, err)
|
||||
}
|
||||
htmlPath := filepath.Join(dir, "out.html")
|
||||
if err := s.ExportHTML(htmlPath); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, _ := os.ReadFile(htmlPath)
|
||||
if !strings.Contains(string(raw), "hello") || !strings.Contains(string(raw), "notes") {
|
||||
t.Fatalf("html: %s", raw)
|
||||
}
|
||||
if _, err := s.AppendEntry(Entry{Type: TypeLabel, Label: "checkpoint"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
entries, _ := s.Entries()
|
||||
var labels int
|
||||
for _, e := range entries {
|
||||
if e.Type == TypeLabel && e.Label == "checkpoint" {
|
||||
labels++
|
||||
}
|
||||
}
|
||||
if labels != 1 {
|
||||
t.Fatalf("label missing: %+v", entries)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
"agent-native/runtime/session"
|
||||
)
|
||||
|
||||
// Delivery selects how a queued steer or follow-up is handed to the loop.
|
||||
type Delivery string
|
||||
|
||||
const (
|
||||
DeliverOne Delivery = "one"
|
||||
DeliverAll Delivery = "all"
|
||||
)
|
||||
|
||||
// AgentSession is the headless session object: queues, abort, and
|
||||
// mid-run model changes around Run.
|
||||
type AgentSession struct {
|
||||
cfg Config
|
||||
|
||||
mu sync.Mutex
|
||||
steer []llm.Message
|
||||
follow []llm.Message
|
||||
steerMode Delivery
|
||||
followMode Delivery
|
||||
model string
|
||||
thinking string
|
||||
cancel context.CancelFunc
|
||||
runID int
|
||||
}
|
||||
|
||||
// NewAgentSession copies cfg. Queue modes default to one-at-a-time.
|
||||
func NewAgentSession(cfg Config) *AgentSession {
|
||||
if cfg.Model == "" {
|
||||
cfg.Model = "default"
|
||||
}
|
||||
return &AgentSession{
|
||||
cfg: cfg,
|
||||
steerMode: DeliverOne,
|
||||
followMode: DeliverOne,
|
||||
model: cfg.Model,
|
||||
thinking: cfg.ThinkingLevel,
|
||||
}
|
||||
}
|
||||
|
||||
// SetQueueMode sets the delivery mode for steer and follow-up.
|
||||
func (a *AgentSession) SetQueueMode(steer, follow Delivery) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
if steer != "" {
|
||||
a.steerMode = steer
|
||||
}
|
||||
if follow != "" {
|
||||
a.followMode = follow
|
||||
}
|
||||
}
|
||||
|
||||
// Steer queues a message to be delivered after the current tool batch.
|
||||
func (a *AgentSession) Steer(msg llm.Message) {
|
||||
a.mu.Lock()
|
||||
a.steer = append(a.steer, msg)
|
||||
a.mu.Unlock()
|
||||
a.emitQueue("steer")
|
||||
}
|
||||
|
||||
// FollowUp queues a message for after the assistant stops.
|
||||
func (a *AgentSession) FollowUp(msg llm.Message) {
|
||||
a.mu.Lock()
|
||||
a.follow = append(a.follow, msg)
|
||||
a.mu.Unlock()
|
||||
a.emitQueue("follow_up")
|
||||
}
|
||||
|
||||
// ClearQueue drops both queues.
|
||||
func (a *AgentSession) ClearQueue() {
|
||||
a.mu.Lock()
|
||||
a.steer = nil
|
||||
a.follow = nil
|
||||
a.mu.Unlock()
|
||||
a.emitQueue("clear")
|
||||
}
|
||||
|
||||
// Abort cancels the in-flight Prompt.
|
||||
func (a *AgentSession) Abort() {
|
||||
a.mu.Lock()
|
||||
cancel := a.cancel
|
||||
a.mu.Unlock()
|
||||
if cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// SetModel changes the model and thinking level used on the next turn.
|
||||
// A session store, when present, records the change.
|
||||
func (a *AgentSession) SetModel(model, thinking string) error {
|
||||
if thinking != "" {
|
||||
norm, err := llm.NormalizeThinking(thinking)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
thinking = norm
|
||||
}
|
||||
a.mu.Lock()
|
||||
if model != "" {
|
||||
a.model = model
|
||||
}
|
||||
if thinking != "" {
|
||||
a.thinking = thinking
|
||||
}
|
||||
store := a.cfg.Session
|
||||
a.mu.Unlock()
|
||||
if store == nil {
|
||||
return nil
|
||||
}
|
||||
if model != "" {
|
||||
if _, err := store.AppendEntry(session.Entry{Type: session.TypeModelChange, Model: model}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if thinking != "" {
|
||||
if _, err := store.AppendEntry(session.Entry{Type: session.TypeThinkingLevelChange, ThinkingLevel: thinking}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Prompt runs one user message until the agent settles.
|
||||
func (a *AgentSession) Prompt(ctx context.Context, user llm.Message) (Result, error) {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
a.mu.Lock()
|
||||
a.runID++
|
||||
id := a.runID
|
||||
a.cancel = cancel
|
||||
cfg := a.cfg
|
||||
a.mu.Unlock()
|
||||
defer func() {
|
||||
cancel()
|
||||
a.mu.Lock()
|
||||
if a.runID == id {
|
||||
a.cancel = nil
|
||||
}
|
||||
a.mu.Unlock()
|
||||
}()
|
||||
cfg.GetSteering = func() []llm.Message { return a.take("steer") }
|
||||
cfg.GetFollowUp = func() []llm.Message { return a.take("follow") }
|
||||
cfg.Resolve = func() (string, string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
return a.model, a.thinking
|
||||
}
|
||||
return Run(ctx, cfg, user)
|
||||
}
|
||||
|
||||
// Compact appends a compaction entry and emits compaction events.
|
||||
func (a *AgentSession) Compact(keep int, summarize func([]llm.Message) (string, error)) (bool, error) {
|
||||
if a.cfg.Session == nil {
|
||||
return false, nil
|
||||
}
|
||||
emit(a.cfg, Event{Type: "compaction_start"})
|
||||
done, err := a.cfg.Session.Compact(keep, summarize)
|
||||
text := "noop"
|
||||
if done {
|
||||
text = "compacted"
|
||||
}
|
||||
if err != nil {
|
||||
text = err.Error()
|
||||
}
|
||||
emit(a.cfg, Event{Type: "compaction_end", Text: text})
|
||||
return done, err
|
||||
}
|
||||
|
||||
func (a *AgentSession) take(kind string) []llm.Message {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
q := &a.steer
|
||||
mode := a.steerMode
|
||||
if kind == "follow" {
|
||||
q = &a.follow
|
||||
mode = a.followMode
|
||||
}
|
||||
if len(*q) == 0 {
|
||||
return nil
|
||||
}
|
||||
if mode == DeliverAll {
|
||||
out := append([]llm.Message(nil), (*q)...)
|
||||
*q = nil
|
||||
return out
|
||||
}
|
||||
out := []llm.Message{(*q)[0]}
|
||||
*q = (*q)[1:]
|
||||
return out
|
||||
}
|
||||
|
||||
func (a *AgentSession) emitQueue(text string) {
|
||||
emit(a.cfg, Event{Type: "queue_update", Text: text})
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"agent-native/runtime/llm"
|
||||
)
|
||||
|
||||
func TestAgentSessionPromptForwardsStop(t *testing.T) {
|
||||
called := false
|
||||
sess := NewAgentSession(Config{
|
||||
Model: "m",
|
||||
Stream: func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "hi", StopReason: llm.StopEnd}}
|
||||
close(ch)
|
||||
return ch
|
||||
},
|
||||
Stop: func(last llm.Message, info StopInfo) StopDecision {
|
||||
called = true
|
||||
return StopDecision{}
|
||||
},
|
||||
})
|
||||
res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !called || res.Text != "hi" {
|
||||
t.Fatalf("called=%v text=%q", called, res.Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentSessionSteerFollowAndAbort(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
step := 0
|
||||
stream := llm.StreamFunc(func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
ch := make(chan llm.Event, 2)
|
||||
go func() {
|
||||
defer close(ch)
|
||||
mu.Lock()
|
||||
n := step
|
||||
step++
|
||||
mu.Unlock()
|
||||
if opts.ThinkingLevel != "low" && n == 0 {
|
||||
t.Errorf("thinking level not forwarded: %s", opts.ThinkingLevel)
|
||||
}
|
||||
if n == 0 {
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
||||
Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
|
||||
ToolCalls: []llm.ToolCall{{ID: "1", Name: "missing", Arguments: "{}"}},
|
||||
}}
|
||||
return
|
||||
}
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
||||
Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd,
|
||||
}}
|
||||
}()
|
||||
return ch
|
||||
})
|
||||
var types []string
|
||||
sess := NewAgentSession(Config{
|
||||
Model: "m",
|
||||
ThinkingLevel: "low",
|
||||
Stream: stream,
|
||||
MaxTurns: 6,
|
||||
OnEvent: func(ev Event) {
|
||||
types = append(types, ev.Type)
|
||||
},
|
||||
})
|
||||
sess.SetQueueMode(DeliverOne, DeliverAll)
|
||||
sess.Steer(llm.Message{Role: llm.RoleUser, Content: "steer-me"})
|
||||
sess.FollowUp(llm.Message{Role: llm.RoleUser, Content: "later"})
|
||||
res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.StopReason != llm.StopEnd {
|
||||
t.Fatalf("stop %s err %v", res.StopReason, res.Err)
|
||||
}
|
||||
found := map[string]bool{}
|
||||
for _, typ := range types {
|
||||
found[typ] = true
|
||||
}
|
||||
for _, want := range []string{"message_start", "queue_update", "tool_execution_update", "agent_settled"} {
|
||||
if !found[want] {
|
||||
t.Fatalf("missing %s in %v", want, types)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestClearQueueAndSetModel(t *testing.T) {
|
||||
sess := NewAgentSession(Config{Model: "old", Stream: func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||
ch := make(chan llm.Event, 1)
|
||||
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: model, StopReason: llm.StopEnd}}
|
||||
close(ch)
|
||||
return ch
|
||||
}})
|
||||
sess.Steer(llm.Message{Role: llm.RoleUser, Content: "drop"})
|
||||
sess.ClearQueue()
|
||||
if err := sess.SetModel("new-model", "high"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "hi"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res.Text != "new-model" {
|
||||
t.Fatalf("model switch did not apply: %q", res.Text)
|
||||
}
|
||||
}
|
||||
@@ -107,6 +107,35 @@ func (s *Store) Load(name string) (Skill, error) {
|
||||
return Skill{}, fmt.Errorf("unknown skill %q (not installed in any configured root)", name)
|
||||
}
|
||||
|
||||
// ExpandCommands replaces /skill:name with that skill's body.
|
||||
// An unknown name returns an error and leaves the text unchanged.
|
||||
func (s *Store) ExpandCommands(text string) (string, error) {
|
||||
const marker = "/skill:"
|
||||
var b strings.Builder
|
||||
rest := text
|
||||
for {
|
||||
i := strings.Index(rest, marker)
|
||||
if i < 0 {
|
||||
b.WriteString(rest)
|
||||
break
|
||||
}
|
||||
b.WriteString(rest[:i])
|
||||
name := rest[i+len(marker):]
|
||||
end := 0
|
||||
for end < len(name) && name[end] != ' ' && name[end] != '\n' && name[end] != '\t' {
|
||||
end++
|
||||
}
|
||||
skillName := name[:end]
|
||||
skill, err := s.Load(skillName)
|
||||
if err != nil {
|
||||
return text, err
|
||||
}
|
||||
b.WriteString(skill.Body)
|
||||
rest = name[end:]
|
||||
}
|
||||
return b.String(), nil
|
||||
}
|
||||
|
||||
// parse 解析一份 SKILL.md:目录名权威,frontmatter 可选。
|
||||
func parse(dirName, path string, data []byte) (Skill, error) {
|
||||
if !validSkillName(dirName) {
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
// aliasTool exposes an existing tool under a pi-style name.
|
||||
type aliasTool struct {
|
||||
name string
|
||||
inner Tool
|
||||
}
|
||||
|
||||
func (a aliasTool) Name() string { return a.name }
|
||||
func (a aliasTool) Description() string { return a.inner.Description() }
|
||||
func (a aliasTool) Parameters() json.RawMessage { return a.inner.Parameters() }
|
||||
func (a aliasTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) {
|
||||
return a.inner.Execute(ctx, args)
|
||||
}
|
||||
|
||||
// piAliases maps the headless pi tool name onto the name already registered.
|
||||
var piAliases = map[string]string{
|
||||
"read": "read_file",
|
||||
"write": "write_file",
|
||||
"edit": "edit_file",
|
||||
"bash": "run_command",
|
||||
"grep": "grep",
|
||||
"find": "glob",
|
||||
"ls": "list_dir",
|
||||
}
|
||||
|
||||
// RegisterPiAliases adds pi tool names beside the existing ones.
|
||||
// A name is skipped when the underlying tool is not registered.
|
||||
func (r *Registry) RegisterPiAliases() {
|
||||
for alias, inner := range piAliases {
|
||||
if alias == inner || r.Has(alias) || !r.Has(inner) {
|
||||
continue
|
||||
}
|
||||
r.Register(aliasTool{name: alias, inner: r.tools[inner]})
|
||||
}
|
||||
}
|
||||
|
||||
// FilterOptions selects which tools stay visible.
|
||||
type FilterOptions struct {
|
||||
Tools []string
|
||||
Exclude []string
|
||||
NoBuiltin bool
|
||||
}
|
||||
|
||||
// Filter applies an allow list, then exclusions. NoBuiltin drops the
|
||||
// workspace tool names and their pi aliases.
|
||||
func (r *Registry) Filter(opts FilterOptions) error {
|
||||
if len(opts.Tools) > 0 {
|
||||
if err := r.Restrict(opts.Tools); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
drop := map[string]bool{}
|
||||
for _, name := range opts.Exclude {
|
||||
drop[name] = true
|
||||
}
|
||||
if opts.NoBuiltin {
|
||||
for _, name := range []string{"read_file", "write_file", "edit_file", "multi_edit", "list_dir", "grep", "glob", "run_command", "run_python"} {
|
||||
drop[name] = true
|
||||
}
|
||||
for alias := range piAliases {
|
||||
drop[alias] = true
|
||||
}
|
||||
}
|
||||
if len(drop) == 0 {
|
||||
return nil
|
||||
}
|
||||
var order []string
|
||||
for _, name := range r.order {
|
||||
if drop[name] {
|
||||
delete(r.tools, name)
|
||||
continue
|
||||
}
|
||||
order = append(order, name)
|
||||
}
|
||||
r.order = order
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package tools
|
||||
|
||||
import "os"
|
||||
|
||||
func atomicWriteFile(path string, data []byte, mode os.FileMode) error {
|
||||
tmp := path + ".tmp"
|
||||
f, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := f.Write(data); err != nil {
|
||||
_ = f.Close()
|
||||
_ = os.Remove(tmp)
|
||||
return err
|
||||
}
|
||||
if err := f.Sync(); err != nil {
|
||||
_ = f.Close()
|
||||
_ = os.Remove(tmp)
|
||||
return err
|
||||
}
|
||||
if err := f.Close(); err != nil {
|
||||
_ = os.Remove(tmp)
|
||||
return err
|
||||
}
|
||||
if err := os.Rename(tmp, path); err != nil {
|
||||
_ = os.Remove(tmp)
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -39,7 +39,7 @@ func (t *readFile) Execute(_ context.Context, args json.RawMessage) (Result, err
|
||||
if obj.Offset <= 0 && obj.Limit <= 0 {
|
||||
truncated := false
|
||||
if len(b) > maxReadBytes {
|
||||
b = b[:maxReadBytes]
|
||||
b = []byte(CutUTF8(string(b), maxReadBytes))
|
||||
truncated = true
|
||||
}
|
||||
content := string(b)
|
||||
@@ -93,7 +93,7 @@ func (t *writeFile) Execute(_ context.Context, args json.RawMessage) (Result, er
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
if err := os.WriteFile(path, []byte(obj.Content), 0o644); err != nil {
|
||||
if err := atomicWriteFile(path, []byte(obj.Content), 0o644); err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
return Result{Content: "wrote " + obj.Path}, nil
|
||||
@@ -134,7 +134,7 @@ func (t *editFile) Execute(_ context.Context, args json.RawMessage) (Result, err
|
||||
return Result{Content: fmt.Sprintf("expected one match, found %d", n), IsError: true}, nil
|
||||
}
|
||||
updated := strings.Replace(text, obj.Old, obj.New, 1)
|
||||
if err := os.WriteFile(path, []byte(updated), 0o644); err != nil {
|
||||
if err := atomicWriteFile(path, []byte(updated), 0o644); err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
return Result{Content: "edited " + obj.Path}, nil
|
||||
@@ -279,7 +279,7 @@ type multiEditSpan struct {
|
||||
|
||||
func (t *multiEditTool) Execute(_ context.Context, args json.RawMessage) (Result, error) {
|
||||
var obj struct {
|
||||
Path string `json:"path"`
|
||||
Path string `json:"path"`
|
||||
Edits []multiEditSpan `json:"edits"`
|
||||
}
|
||||
if err := json.Unmarshal(args, &obj); err != nil {
|
||||
@@ -307,7 +307,7 @@ func (t *multiEditTool) Execute(_ context.Context, args json.RawMessage) (Result
|
||||
}
|
||||
text = strings.Replace(text, e.Old, e.New, 1)
|
||||
}
|
||||
if err := os.WriteFile(path, []byte(text), 0o644); err != nil {
|
||||
if err := atomicWriteFile(path, []byte(text), 0o644); err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
}
|
||||
return Result{Content: fmt.Sprintf("applied %d edits to %s", len(obj.Edits), obj.Path)}, nil
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// ExecConfig 声明执行类工具的运行条件,由组合根显式提供;
|
||||
@@ -176,7 +177,12 @@ func formatProcessOutput(out []byte, timedOut bool, err error) string {
|
||||
if len(out) > execMaxOutput {
|
||||
half := execMaxOutput / 2
|
||||
dropped := len(out) - execMaxOutput
|
||||
text = string(out[:half]) + fmt.Sprintf("\n[... truncated %d bytes ...]\n", dropped) + string(out[len(out)-half:])
|
||||
head := CutUTF8(string(out), half)
|
||||
tail := string(out[len(out)-half:])
|
||||
for len(tail) > 0 && !utf8.RuneStart(tail[0]) {
|
||||
tail = tail[1:]
|
||||
}
|
||||
text = head + fmt.Sprintf("\n[... truncated %d bytes ...]\n", dropped) + tail
|
||||
}
|
||||
switch {
|
||||
case timedOut:
|
||||
|
||||
+63
-13
@@ -17,6 +17,10 @@ import (
|
||||
// FetchConfig 声明 fetch_url 的运行条件;未注册(组合根未提供)则工具不存在。
|
||||
type FetchConfig struct {
|
||||
// Client 可选注入(测试替身);缺省构造 20s 超时客户端。
|
||||
// Transport 为 nil 或 *http.Transport 时,Clone 后替换 DialContext:
|
||||
// 拨号时再次解析,只连接通过 SSRF 校验的 IP。
|
||||
// 自定义 RoundTripper 无法替换拨号,只做请求前预校验;实现不得在
|
||||
// RoundTrip 里重新解析到未校验的地址。
|
||||
Client *http.Client
|
||||
// MaxBytes 响应体上限(默认 2 MiB)。
|
||||
MaxBytes int
|
||||
@@ -89,19 +93,24 @@ func (t *fetchTool) Execute(ctx context.Context, args json.RawMessage) (Result,
|
||||
req.Header.Set("User-Agent", "agent-native-fetch/1.0")
|
||||
var finalURL string
|
||||
resolve := t.cfg.resolver()
|
||||
client2 := client
|
||||
if client2.CheckRedirect == nil {
|
||||
client2 = &http.Client{
|
||||
Transport: client.Transport,
|
||||
Jar: client.Jar,
|
||||
Timeout: client.Timeout,
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= 5 {
|
||||
return fmt.Errorf("stopped after 5 redirects")
|
||||
}
|
||||
return validateFetchTarget(req.URL, allowed, resolve)
|
||||
},
|
||||
}
|
||||
hostRedirect := client.CheckRedirect
|
||||
transport, _ := pinFetchTransport(client.Transport, resolve)
|
||||
client2 := &http.Client{
|
||||
Transport: transport,
|
||||
Jar: client.Jar,
|
||||
Timeout: client.Timeout,
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= 5 {
|
||||
return fmt.Errorf("stopped after 5 redirects")
|
||||
}
|
||||
if err := validateFetchTarget(req.URL, allowed, resolve); err != nil {
|
||||
return err
|
||||
}
|
||||
if hostRedirect != nil {
|
||||
return hostRedirect(req, via)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
if err := validateFetchTarget(req.URL, allowed, resolve); err != nil {
|
||||
return Result{Content: err.Error(), IsError: true}, nil
|
||||
@@ -127,6 +136,47 @@ func (t *fetchTool) Execute(ctx context.Context, args json.RawMessage) (Result,
|
||||
return Result{Content: "url: " + finalURL + "\ncontent-type: " + ct + "\n\n" + text}, nil
|
||||
}
|
||||
|
||||
// pinFetchTransport 在默认或 *http.Transport 上安装「解析并拨号已校验 IP」的 DialContext。
|
||||
// 自定义 RoundTripper 原样返回,调用方保留预校验。
|
||||
func pinFetchTransport(base http.RoundTripper, resolve func(string) ([]net.IP, error)) (http.RoundTripper, bool) {
|
||||
var tr *http.Transport
|
||||
switch t := base.(type) {
|
||||
case nil:
|
||||
tr = http.DefaultTransport.(*http.Transport).Clone()
|
||||
case *http.Transport:
|
||||
tr = t.Clone()
|
||||
default:
|
||||
return base, false
|
||||
}
|
||||
// 不走环境变量里的代理:代理拨号会绕开「只连接已校验 IP」。
|
||||
tr.Proxy = nil
|
||||
dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
|
||||
tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ips, err := resolve(host)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("resolve %q: %w", host, err)
|
||||
}
|
||||
var chosen net.IP
|
||||
for _, ip := range ips {
|
||||
if ipIsPrivate(ip) {
|
||||
return nil, fmt.Errorf("target %q resolves to private/loopback address %s (SSRF guard)", host, ip)
|
||||
}
|
||||
if chosen == nil {
|
||||
chosen = ip
|
||||
}
|
||||
}
|
||||
if chosen == nil {
|
||||
return nil, fmt.Errorf("resolve %q: no addresses", host)
|
||||
}
|
||||
return dialer.DialContext(ctx, network, net.JoinHostPort(chosen.String(), port))
|
||||
}
|
||||
return tr, true
|
||||
}
|
||||
|
||||
// validateFetchTarget 校验 scheme、主机策略与 SSRF 硬边界(私网/回环/链路本地)。
|
||||
func validateFetchTarget(u *url.URL, allowed map[string]bool, resolve func(string) ([]net.IP, error)) error {
|
||||
if u.Scheme != "http" && u.Scheme != "https" {
|
||||
|
||||
@@ -144,6 +144,45 @@ func (s redirectTransport) RoundTrip(req *http.Request) (*http.Response, error)
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func TestFetchResolveChangesToLoopbackRejected(t *testing.T) {
|
||||
var n int
|
||||
reg := fetchReg(FetchConfig{
|
||||
Resolve: func(host string) ([]net.IP, error) {
|
||||
n++
|
||||
if n == 1 {
|
||||
return []net.IP{net.ParseIP("93.184.216.34")}, nil
|
||||
}
|
||||
return []net.IP{net.ParseIP("127.0.0.1")}, nil
|
||||
},
|
||||
})
|
||||
res := runTool(t, reg, "fetch_url", `{"url":"http://flip.example/x"}`)
|
||||
if !res.IsError || !strings.Contains(res.Content, "private/loopback") {
|
||||
t.Fatalf("second resolve to loopback must be rejected, got %+v", res)
|
||||
}
|
||||
if n < 2 {
|
||||
t.Fatalf("dial should resolve again, got %d calls", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchHostCheckRedirectStillRejectsPrivate(t *testing.T) {
|
||||
reg := fetchReg(FetchConfig{
|
||||
Client: &http.Client{
|
||||
Transport: redirectTransport{to: "http://127.0.0.1/evil"},
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
return nil
|
||||
},
|
||||
},
|
||||
Resolve: fakeResolve(map[string][]net.IP{
|
||||
"public.example": {net.ParseIP("93.184.216.34")},
|
||||
"127.0.0.1": {net.ParseIP("127.0.0.1")},
|
||||
}),
|
||||
})
|
||||
res := runTool(t, reg, "fetch_url", `{"url":"http://public.example/redirect"}`)
|
||||
if !res.IsError || !strings.Contains(res.Content, "private/loopback") {
|
||||
t.Fatalf("host CheckRedirect must not bypass SSRF guard, got %+v", res)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchNotRegisteredByDefault(t *testing.T) {
|
||||
reg := WorkspaceTools(t.TempDir())
|
||||
for _, d := range reg.Decls() {
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func TestR2WriteFileIsAtomic(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
reg := NewRegistry(dir)
|
||||
for _, tool := range []Tool{&writeFile{root: dir}} {
|
||||
reg.Register(tool)
|
||||
}
|
||||
res := reg.Execute(context.Background(), "write_file", []byte(`{"path":"note.txt","content":"完整"}`))
|
||||
if res.IsError {
|
||||
t.Fatal(res.Content)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, "note.txt.tmp")); !os.IsNotExist(err) {
|
||||
t.Fatal("temp file left behind")
|
||||
}
|
||||
b, err := os.ReadFile(filepath.Join(dir, "note.txt"))
|
||||
if err != nil || string(b) != "完整" {
|
||||
t.Fatalf("file: %q %v", b, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestR2CutUTF8DoesNotSplitRunes(t *testing.T) {
|
||||
s := strings.Repeat("界", 5)
|
||||
cut := CutUTF8(s, 4)
|
||||
if !utf8.ValidString(cut) {
|
||||
t.Fatalf("invalid utf-8 %q", cut)
|
||||
}
|
||||
if cut != "界" {
|
||||
t.Fatalf("cut %q", cut)
|
||||
}
|
||||
}
|
||||
|
||||
func TestR2SchemaRejectsUnknownKeys(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
reg := NewRegistry(dir)
|
||||
reg.Register(&writeFile{root: dir})
|
||||
res := reg.Execute(context.Background(), "write_file", []byte(`{"path":"a.txt","content":"x","extra":1}`))
|
||||
if !res.IsError || !strings.Contains(res.Content, "unknown argument") {
|
||||
t.Fatalf("result: %+v", res)
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
@@ -37,6 +38,8 @@ type Registry struct {
|
||||
Before BeforeHook
|
||||
After AfterHook
|
||||
Audit io.Writer
|
||||
// hookMu 串行 Before、After 与 Audit,避免并行工具调用交错写审计或重入钩子。
|
||||
hookMu sync.Mutex
|
||||
}
|
||||
|
||||
func NewRegistry(root string) *Registry {
|
||||
@@ -45,6 +48,11 @@ func NewRegistry(root string) *Registry {
|
||||
|
||||
func (r *Registry) Root() string { return r.root }
|
||||
|
||||
func (r *Registry) Has(name string) bool {
|
||||
_, ok := r.tools[name]
|
||||
return ok
|
||||
}
|
||||
|
||||
func (r *Registry) Register(t Tool) {
|
||||
if _, ok := r.tools[t.Name()]; !ok {
|
||||
r.order = append(r.order, t.Name())
|
||||
@@ -154,7 +162,10 @@ func (r *Registry) Execute(ctx context.Context, name string, args json.RawMessag
|
||||
return res
|
||||
}
|
||||
if r.Before != nil {
|
||||
if block, reason := r.Before(ctx, name, args); block {
|
||||
r.hookMu.Lock()
|
||||
block, reason := r.Before(ctx, name, args)
|
||||
r.hookMu.Unlock()
|
||||
if block {
|
||||
if reason == "" {
|
||||
reason = "blocked"
|
||||
}
|
||||
@@ -168,13 +179,17 @@ func (r *Registry) Execute(ctx context.Context, name string, args json.RawMessag
|
||||
res = Result{Content: err.Error(), IsError: true}
|
||||
}
|
||||
if r.After != nil {
|
||||
r.hookMu.Lock()
|
||||
res = r.After(ctx, name, args, res)
|
||||
r.hookMu.Unlock()
|
||||
}
|
||||
r.audit(name, false, res.IsError, "")
|
||||
return res
|
||||
}
|
||||
|
||||
func (r *Registry) audit(name string, blocked, isError bool, note string) {
|
||||
r.hookMu.Lock()
|
||||
defer r.hookMu.Unlock()
|
||||
if r.Audit == nil {
|
||||
return
|
||||
}
|
||||
@@ -223,6 +238,13 @@ func validateArgs(schema, args json.RawMessage) error {
|
||||
return fmt.Errorf("argument %s must be %s", key, prop.Type)
|
||||
}
|
||||
}
|
||||
if len(spec.Properties) > 0 {
|
||||
for key := range obj {
|
||||
if _, ok := spec.Properties[key]; !ok {
|
||||
return fmt.Errorf("unknown argument %s", key)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -7,7 +7,9 @@ import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSchemaRejectsMissingAndWrongType(t *testing.T) {
|
||||
@@ -91,6 +93,47 @@ func TestBeforeHookBlocksAndAudits(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
type slowHookTool struct{}
|
||||
|
||||
func (slowHookTool) Name() string { return "slow_hook" }
|
||||
func (slowHookTool) Description() string { return "noop" }
|
||||
func (slowHookTool) Parameters() json.RawMessage { return json.RawMessage(`{"type":"object","properties":{}}`) }
|
||||
func (slowHookTool) Execute(context.Context, json.RawMessage) (Result, error) {
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
|
||||
func TestHooksDoNotRunConcurrently(t *testing.T) {
|
||||
reg := NewRegistry(t.TempDir())
|
||||
reg.Register(slowHookTool{})
|
||||
var mu sync.Mutex
|
||||
inFlight, maxInFlight := 0, 0
|
||||
reg.Before = func(context.Context, string, json.RawMessage) (bool, string) {
|
||||
mu.Lock()
|
||||
inFlight++
|
||||
if inFlight > maxInFlight {
|
||||
maxInFlight = inFlight
|
||||
}
|
||||
mu.Unlock()
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
mu.Lock()
|
||||
inFlight--
|
||||
mu.Unlock()
|
||||
return false, ""
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 2; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
reg.Execute(context.Background(), "slow_hook", []byte(`{}`))
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if maxInFlight != 1 {
|
||||
t.Fatalf("hooks overlapped: max in flight %d", maxInFlight)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAfterHookRewrites(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
reg := WorkspaceTools(dir)
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
package tools
|
||||
|
||||
import "unicode/utf8"
|
||||
|
||||
// CutUTF8 按字节上限截断,并回退到完整 rune,避免切开 UTF-8。
|
||||
func CutUTF8(s string, maxBytes int) string {
|
||||
if maxBytes <= 0 || len(s) <= maxBytes {
|
||||
return s
|
||||
}
|
||||
end := maxBytes
|
||||
for end > 0 && !utf8.ValidString(s[:end]) {
|
||||
end--
|
||||
}
|
||||
return s[:end]
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"work_dir": "/workspace/demo",
|
||||
"prompt": "写一份说明",
|
||||
"instructions": "只写目标文件",
|
||||
"targets": ["output.md"],
|
||||
"tool_profile": "workspace"
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"success": true,
|
||||
"patches": {"output.md": "ZGVtbw=="},
|
||||
"logs": "done"
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user