AutoZero/agent/occupy_test.go

package agent

import (
	"context"
	"errors"
	"testing"

	"git.gopher.ink/mirrors2/AutoZero/llm"
)

// 工作区占用挂在 Run 的进出口上，而不是各入口自己 TryOccupy/defer Release。
// 理由是一条事故形状：WebUI 与 cron 各抄了一对，终端一个都没抄——于是手机上的
// 回合和终端里的回合可以同时改同一个项目，互相覆盖之外还撞 git 的 index.lock。
// 这里钉四件事：抢不到就什么都不开始、开始过就一定归还、归还排在收尾钩子之后
// （checkpoint 的 git 操作也要在占用窗口内）、子 agent 不再抢一次。

func TestOccupyRefusalStartsNothing(t *testing.T) {
	a, sess := newTestAgent(t)
	busy := errors.New("工作区正被别的会话占用")
	prov := &occupyProvider{reply: "不该跑到这里"}
	a.Provider = prov
	releases := 0
	a.Opt.Occupy = func() error { return busy }
	a.Opt.Release = func() { releases++ }
	a.Opt.OnTurnEnd = func(TurnReport) { t.Error("被拒绝的回合不该走到收尾") }

	err := a.Run(context.Background(), "改一下代码")
	if !errors.Is(err, busy) {
		t.Fatalf("Run = %v, want 原样交出 %v", err, busy)
	}
	// 提问没落盘、没请求网关：被拒的一轮在会话里根本不存在，用户重发就是干净的一次。
	if len(sess.Messages) != 0 {
		t.Errorf("被拒的回合往会话里写了 %d 条消息，want 0", len(sess.Messages))
	}
	if prov.calls != 0 {
		t.Errorf("被拒的回合请求了网关 %d 次，want 0", prov.calls)
	}
	if releases != 0 {
		t.Errorf("没占到却归还了 %d 次，want 0（会把别人的登记摘掉）", releases)
	}
}

// 正常一轮：占用 → 回合 → 收尾钩子 → 归还，各恰好一次，顺序不能倒——
// 归还早于 OnTurnEnd 就等于 checkpoint 还在 commit，工作区已经放行给别人了。
func TestOccupyReleaseWrapsTheWholeTurn(t *testing.T) {
	a, _ := newTestAgent(t)
	var events []string
	a.Provider = &occupyProvider{reply: "做完了"}
	a.Opt.Occupy = func() error { events = append(events, "occupy"); return nil }
	a.Opt.Release = func() { events = append(events, "release") }
	a.Opt.OnTurnEnd = func(TurnReport) { events = append(events, "turn_end") }

	if err := a.Run(context.Background(), "把这件事做完"); err != nil {
		t.Fatal(err)
	}
	want := []string{"occupy", "turn_end", "release"}
	if len(events) != len(want) {
		t.Fatalf("事件 = %v, want %v", events, want)
	}
	for i := range want {
		if events[i] != want[i] {
			t.Fatalf("事件 = %v, want %v", events, want)
		}
	}
}

// 崩掉的一轮同样要归还：占用是进程内的登记表，漏一次就把这个工作区永久锁死，
// 之后每一条消息都只会看到"正被占用"。
func TestReleaseHappensOnFailedTurn(t *testing.T) {
	a, _ := newTestAgent(t)
	boom := errors.New("网关炸了")
	a.Provider = &occupyProvider{err: boom}
	occupies, releases := 0, 0
	a.Opt.Occupy = func() error { occupies++; return nil }
	a.Opt.Release = func() { releases++ }

	if err := a.Run(context.Background(), "办不成"); err == nil {
		t.Fatal("该失败")
	}
	if occupies != 1 || releases != 1 {
		t.Errorf("占用 %d 次 / 归还 %d 次，want 1 / 1", occupies, releases)
	}
}

// 子 agent 不重复占用：占用窗口是父级拿的，子级在同一个目录里干活。
// 真去抢一次会自己撞自己的登记（holder 不同就是失败），一次委派直接卡死。
func TestSubAgentSkipsOccupancy(t *testing.T) {
	a, _ := newTestAgent(t)
	calls := 0
	a.Provider = &occupyProvider{reply: "子过程结束"}
	a.Opt.Occupy = func() error { calls++; return nil }
	a.Opt.Release = func() { calls++ }
	a.Depth = 1

	if err := a.Run(context.Background(), "子任务"); err != nil {
		t.Fatal(err)
	}
	if calls != 0 {
		t.Errorf("子 agent 动用了占用登记 %d 次，want 0", calls)
	}
}

// 父回合握着占用时把活派给子 agent：子级不能再去抢同一份登记（抢就是自锁，
// 一次委派能把整个工作区卡死）。这条走**真实的 agent 工具**，不只是设一个 Depth：
// 将来谁把父级 Options 整包继承给子级（OnNotice、Stderr 已经这么继承了），
// 占用钩子也会跟着下去，届时这里会看到第二次 occupy 而不是断言通过。
func TestDelegationDoesNotReacquireWorkspace(t *testing.T) {
	a, _ := newTestAgent(t)
	var seq []string
	a.Opt.Occupy = func() error { seq = append(seq, "occupy"); return nil }
	a.Opt.Release = func() { seq = append(seq, "release") }
	a.MaxDepth = 1
	a.SubAgentConcurrency = 1
	if err := a.Tools.Register(a.SubAgentTool()); err != nil {
		t.Fatal(err)
	}
	a.Provider = &scriptedProvider{fn: func(call int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		switch call {
		case 1: // 父级：委派
			return llm.Message{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{
				{ID: "t1", Name: "agent", Input: `{"prompt":"去查一下"}`},
			}}, nil
		case 2: // 子级：一句话交差
			return llm.Message{Role: llm.RoleAssistant, Content: "子任务做完了"}, nil
		default: // 父级拿到结果后收尾
			return llm.Message{Role: llm.RoleAssistant, Content: "收尾"}, nil
		}
	}}

	if err := a.Run(context.Background(), "派个子任务"); err != nil {
		t.Fatal(err)
	}
	if len(seq) != 2 || seq[0] != "occupy" || seq[1] != "release" {
		t.Errorf("占用事件 = %v, want [occupy release]（子 agent 抢了一次就是自锁）", seq)
	}
}

// occupyProvider 一问一答，带调用次数（被拒的回合必须证明"根本没请求网关"）。
type occupyProvider struct {
	reply string
	err   error
	calls int
}

func (p *occupyProvider) Stream(context.Context, string, []llm.Message, []llm.ToolDef,
	llm.StreamOptions, func(llm.StreamDelta)) (llm.Message, error) {
	p.calls++
	if p.err != nil {
		return llm.Message{}, p.err
	}
	return llm.Message{Role: llm.RoleAssistant, Content: p.reply}, nil
}