agentic-coding-workflow/internal/acp/supervisor_test.go

//go:build integration

// supervisor_test.go 集成测试 supervisor 的 spawn / handshake / kill / monitor /
// onExit 全链路。复用 cmd/fake-acp-agent 子进程：
//   - HappyPath：正常 spawn → handshake → 正常 Kill → 写 SessionExited
//   - HangAgent：FAKE_ACP_HANG=true 让 agent 永不响应 → handshake 超时 → Spawn 返错
//   - AgentCrashes：FAKE_ACP_CRASH_AFTER_PROMPTS=1 让 agent 处理一次 prompt 后 exit(1)
//
// 全部走 testcontainers PG，每测试一个独立容器（setupRepo 内 t.Cleanup）。
package acp_test

import (
	"context"
	"encoding/json"
	"sync"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/yan1h/agent-coding-workflow/internal/acp"
	"github.com/yan1h/agent-coding-workflow/internal/audit"
	"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
	"github.com/yan1h/agent-coding-workflow/internal/mcp"
	"github.com/yan1h/agent-coding-workflow/internal/workspace"
)

// newSupervisorForTest 构造一个绑定到 PG 真实 audit Recorder 的 Supervisor，
// 使用静音 logger + 默认 SupervisorConfig（5s SpawnTimeout 给 fake agent 留 buffer）。
// withAudit=false 时 Recorder 为 nil，对应 onExit 跳过 audit 路径。
func newSupervisorForTest(t *testing.T, pool *pgxpool.Pool, repo acp.Repository, withAudit bool, mcpTokens mcp.TokenService, wtSvc workspace.WorktreeService) *acp.Supervisor {
	t.Helper()
	var rec audit.Recorder
	if withAudit {
		rec = audit.NewPostgresRecorder(pool)
	}
	disp := notify.NewDispatcher(notifyRepoStub{}, slogDiscard())
	return acp.NewSupervisor(
		repo, rec, disp,
		wtSvc, nil, // wtSvc: IsMain=false 测试注入 fake；wsRepo: 测试不走该路径
		mcpTokens,
		testEncryptor(t),
		acp.SupervisorConfig{
			SpawnTimeout:       5 * time.Second,
			KillGrace:          1 * time.Second,
			StderrBufferLines:  50,
			StderrTailBytes:    1000,
			EventMaxPayload:    65536,
			EventTruncateField: 4096,
			ShutdownGrace:      5 * time.Second,
		},
		slogDiscard(),
	)
}

// TestSupervisor_HappyPath_Spawn_Handshake 验证：
//  1. Spawn 启动 fake-acp-agent 并完成 initialize + session/new
//  2. UpdateSessionRunning 把 agent_session_id="fake-session-id" + pid 写入 DB
//  3. Kill 触发 monitorProcess + onExit，session 落到 SessionExited
func TestSupervisor_HappyPath_Spawn_Handshake(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "sup1@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)

	bin := fakeAgentBinary(t)
	k, err := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "fake1", DisplayName: "Fake",
		BinaryPath: bin, Args: []string{},
		Enabled: true, CreatedBy: uid,
	})
	require.NoError(t, err)

	cwd := t.TempDir()
	sess, err := repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "main", CwdPath: cwd, IsMainWorktree: true, Status: acp.SessionStarting,
	})
	require.NoError(t, err)

	// onExit 会调 ReleaseMainWorktree → 必须先 Acquire 才能 Release（CAS 语义）。
	ok, err := repo.AcquireMainWorktree(ctx, sess.ID, wsid)
	require.NoError(t, err)
	require.True(t, ok)

	sup := newSupervisorForTest(t, pool, repo, true, nil, nil)
	proc, err := sup.Spawn(ctx, sess, k, nil)
	require.NoError(t, err)
	require.NotNil(t, proc)

	got, err := repo.GetSessionByID(ctx, sess.ID)
	require.NoError(t, err)
	assert.Equal(t, acp.SessionRunning, got.Status)
	require.NotNil(t, got.AgentSessionID)
	assert.Equal(t, "fake-session-id", *got.AgentSessionID)
	require.NotNil(t, got.PID)
	assert.NotZero(t, *got.PID)

	sup.Kill(ctx, sess.ID, 1*time.Second)

	// monitorProcess + onExit 在 goroutine 里跑；轮询直到 DB 反映终态。
	deadline := time.After(5 * time.Second)
	for {
		got, err = repo.GetSessionByID(ctx, sess.ID)
		require.NoError(t, err)
		if got.Status == acp.SessionExited || got.Status == acp.SessionCrashed {
			break
		}
		select {
		case <-deadline:
			t.Fatalf("session never reached terminal state, last status=%s", got.Status)
		case <-time.After(50 * time.Millisecond):
		}
	}
	assert.Equal(t, acp.SessionExited, got.Status)
}

// TestSupervisor_HangAgent_HandshakeTimeout 验证：FAKE_ACP_HANG agent 永远不响应；
// Spawn 在 SpawnTimeout 内返回 error，Kill 在内部触发 monitor → onExit → DB 终态。
func TestSupervisor_HangAgent_HandshakeTimeout(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "sup2@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)

	bin := fakeAgentBinary(t)
	k, err := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "fake-hang", DisplayName: "Fake",
		BinaryPath: bin, Args: []string{},
		EncryptedEnv: encryptForTest(t, map[string]string{"FAKE_ACP_HANG": "true"}),
		Enabled:      true, CreatedBy: uid,
	})
	require.NoError(t, err)

	sess, err := repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "main", CwdPath: t.TempDir(), IsMainWorktree: true, Status: acp.SessionStarting,
	})
	require.NoError(t, err)
	ok, err := repo.AcquireMainWorktree(ctx, sess.ID, wsid)
	require.NoError(t, err)
	require.True(t, ok)

	// 用更短的 SpawnTimeout 加速测试（默认 5s 够，这里 1s 避免 CI 等太久）。
	disp := notify.NewDispatcher(notifyRepoStub{}, slogDiscard())
	sup := acp.NewSupervisor(
		repo, audit.NewPostgresRecorder(pool), disp,
		nil, nil, nil, testEncryptor(t),
		acp.SupervisorConfig{
			SpawnTimeout:       1 * time.Second,
			KillGrace:          500 * time.Millisecond,
			StderrBufferLines:  10,
			StderrTailBytes:    500,
			EventMaxPayload:    65536,
			EventTruncateField: 4096,
			ShutdownGrace:      2 * time.Second,
		},
		slogDiscard(),
	)

	_, err = sup.Spawn(ctx, sess, k, nil)
	require.Error(t, err, "hang agent must trigger handshake timeout")

	deadline := time.After(5 * time.Second)
	for {
		got, gerr := repo.GetSessionByID(ctx, sess.ID)
		require.NoError(t, gerr)
		if got.Status == acp.SessionCrashed || got.Status == acp.SessionExited {
			// SIGKILL on hang agent 不属于 \"agent crashed by itself\"，
			// supervisor 会标 KilledByUs → exited。两者都接受。
			assert.Contains(t, []acp.SessionStatus{acp.SessionCrashed, acp.SessionExited}, got.Status)
			return
		}
		select {
		case <-deadline:
			t.Fatalf("session never reached terminal state, last status=%s", got.Status)
		case <-time.After(50 * time.Millisecond):
		}
	}
}

// TestSupervisor_AgentCrashes 验证：fake agent 在处理 1 次 prompt 后 os.Exit(1)；
// monitorProcess 检测到非主动 Kill → status=SessionCrashed + exit_code=1，
// onExit 写 DB + 调 Dispatcher.Dispatch + audit.Record。
func TestSupervisor_AgentCrashes(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "sup3@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)

	bin := fakeAgentBinary(t)
	k, err := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "fake-crash", DisplayName: "Fake",
		BinaryPath: bin, Args: []string{},
		EncryptedEnv: encryptForTest(t, map[string]string{"FAKE_ACP_CRASH_AFTER_PROMPTS": "1"}),
		Enabled:      true, CreatedBy: uid,
	})
	require.NoError(t, err)

	sess, err := repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "main", CwdPath: t.TempDir(), IsMainWorktree: true, Status: acp.SessionStarting,
	})
	require.NoError(t, err)
	ok, err := repo.AcquireMainWorktree(ctx, sess.ID, wsid)
	require.NoError(t, err)
	require.True(t, ok)

	sup := newSupervisorForTest(t, pool, repo, true, nil, nil)
	proc, err := sup.Spawn(ctx, sess, k, nil)
	require.NoError(t, err)
	require.NotNil(t, proc)

	// 通过 relay.SendClient 发一条 prompt → fake agent 处理后 exit(1)。
	idJSON, _ := json.Marshal("client-1")
	err = proc.Relay.SendClient(&acp.Message{
		JSONRPC: "2.0", ID: idJSON, Method: "session/prompt",
		Params: json.RawMessage(`{"sessionId":"fake-session-id","prompt":[{"type":"text","text":"x"}]}`),
	})
	require.NoError(t, err)

	deadline := time.After(10 * time.Second)
	for {
		got, gerr := repo.GetSessionByID(ctx, sess.ID)
		require.NoError(t, gerr)
		if got.Status == acp.SessionCrashed {
			require.NotNil(t, got.ExitCode)
			assert.Equal(t, int32(1), *got.ExitCode)
			return
		}
		select {
		case <-deadline:
			t.Fatalf("crash not detected, last status=%s", got.Status)
		case <-time.After(100 * time.Millisecond):
		}
	}
}

// fakeSupMCPTokens tracks RevokeBySession calls for onExit verification.
type fakeSupMCPTokens struct {
	revoked []uuid.UUID
}

func (f *fakeSupMCPTokens) IssueAdmin(_ context.Context, _ mcp.IssueAdminInput) (mcp.IssueResult, error) {
	return mcp.IssueResult{}, nil
}
func (f *fakeSupMCPTokens) IssueSystemToken(_ context.Context, _ mcp.IssueSystemTokenInput) (mcp.IssueResult, error) {
	return mcp.IssueResult{}, nil
}
func (f *fakeSupMCPTokens) IssueSystemTokenWithTx(_ context.Context, _ pgx.Tx, _ mcp.IssueSystemTokenInput) (mcp.IssueResult, error) {
	return mcp.IssueResult{}, nil
}
func (f *fakeSupMCPTokens) Authenticate(_ context.Context, _ string) (*mcp.Token, error) {
	return nil, nil
}
func (f *fakeSupMCPTokens) Revoke(_ context.Context, _, _ uuid.UUID) error { return nil }
func (f *fakeSupMCPTokens) RevokeBySession(_ context.Context, sid uuid.UUID) error {
	f.revoked = append(f.revoked, sid)
	return nil
}
func (f *fakeSupMCPTokens) List(_ context.Context, _ *uuid.UUID, _ bool) ([]*mcp.Token, error) {
	return nil, nil
}
func (f *fakeSupMCPTokens) GetByID(_ context.Context, _ uuid.UUID) (*mcp.Token, error) {
	return nil, nil
}

// TestSupervisor_OnExit_RevokesMCPToken 验证：onExit 调用 mcpTokens.RevokeBySession，
// 传入正确的 session ID。使用 fake agent 正常 spawn -> kill -> 等终态。
func TestSupervisor_OnExit_RevokesMCPToken(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "sup-mcp@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)

	bin := fakeAgentBinary(t)
	k, err := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "fake-mcp", DisplayName: "Fake",
		BinaryPath: bin, Args: []string{},
		Enabled: true, CreatedBy: uid,
	})
	require.NoError(t, err)

	sess, err := repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "main", CwdPath: t.TempDir(), IsMainWorktree: true, Status: acp.SessionStarting,
	})
	require.NoError(t, err)

	ok, err := repo.AcquireMainWorktree(ctx, sess.ID, wsid)
	require.NoError(t, err)
	require.True(t, ok)

	fakeTokens := &fakeSupMCPTokens{}
	sup := newSupervisorForTest(t, pool, repo, true, fakeTokens, nil)
	proc, err := sup.Spawn(ctx, sess, k, nil)
	require.NoError(t, err)
	require.NotNil(t, proc)

	sup.Kill(ctx, sess.ID, 1*time.Second)

	// wait for terminal state
	deadline := time.After(5 * time.Second)
	for {
		got, gerr := repo.GetSessionByID(ctx, sess.ID)
		require.NoError(t, gerr)
		if got.Status == acp.SessionExited || got.Status == acp.SessionCrashed {
			break
		}
		select {
		case <-deadline:
			t.Fatalf("session never reached terminal state, last status=%s", got.Status)
		case <-time.After(50 * time.Millisecond):
		}
	}

	// Verify RevokeBySession was called with the correct session ID.
	require.Len(t, fakeTokens.revoked, 1, "expected exactly one RevokeBySession call")
	assert.Equal(t, sess.ID, fakeTokens.revoked[0], "RevokeBySession should receive the session ID")
}

// fakeWorktreeSvc 实现 workspace.WorktreeService，仅记录 ReleaseByHolder 收到的
// holder 字符串，用于验证 onExit 的子 worktree 释放路径。
type fakeWorktreeSvc struct {
	mu      sync.Mutex
	holders []string
}

func (f *fakeWorktreeSvc) Create(_ context.Context, _ workspace.Caller, _ uuid.UUID, _ workspace.CreateWorktreeInput) (*workspace.Worktree, error) {
	return nil, nil
}
func (f *fakeWorktreeSvc) List(_ context.Context, _ workspace.Caller, _ uuid.UUID) ([]*workspace.Worktree, error) {
	return nil, nil
}
func (f *fakeWorktreeSvc) Delete(_ context.Context, _ workspace.Caller, _ uuid.UUID) error {
	return nil
}
func (f *fakeWorktreeSvc) Acquire(_ context.Context, _ workspace.Caller, _ uuid.UUID) (*workspace.Worktree, error) {
	return nil, nil
}
func (f *fakeWorktreeSvc) Release(_ context.Context, _ workspace.Caller, _ uuid.UUID) (*workspace.Worktree, error) {
	return nil, nil
}
func (f *fakeWorktreeSvc) ReleaseByHolder(_ context.Context, holder string) ([]*workspace.Worktree, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.holders = append(f.holders, holder)
	return nil, nil
}

// holdersSnapshot 拷贝当前已记录的 holder 列表（onExit 在 monitor goroutine 中调用）。
func (f *fakeWorktreeSvc) holdersSnapshot() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]string(nil), f.holders...)
}

// TestSupervisor_OnExit_ReleasesSubWorktree 验证：IsMainWorktree=false 的 session
// 退出后，onExit 通过 wtSvc.ReleaseByHolder("session:<sid>") 释放子 worktree。
func TestSupervisor_OnExit_ReleasesSubWorktree(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "sup-wt@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)

	bin := fakeAgentBinary(t)
	k, err := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "fake-wt", DisplayName: "Fake",
		BinaryPath: bin, Args: []string{},
		Enabled: true, CreatedBy: uid,
	})
	require.NoError(t, err)

	sess, err := repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "feat-x", CwdPath: t.TempDir(), IsMainWorktree: false, Status: acp.SessionStarting,
	})
	require.NoError(t, err)

	fakeWt := &fakeWorktreeSvc{}
	sup := newSupervisorForTest(t, pool, repo, true, nil, fakeWt)
	proc, err := sup.Spawn(ctx, sess, k, nil)
	require.NoError(t, err)
	require.NotNil(t, proc)

	sup.Kill(ctx, sess.ID, 1*time.Second)

	// onExit 在 monitor goroutine 中异步执行；轮询 fake 直到记录到 release 调用。
	deadline := time.After(5 * time.Second)
	for len(fakeWt.holdersSnapshot()) == 0 {
		select {
		case <-deadline:
			t.Fatal("ReleaseByHolder never called after process exit")
		case <-time.After(50 * time.Millisecond):
		}
	}
	holders := fakeWt.holdersSnapshot()
	require.Len(t, holders, 1, "expected exactly one ReleaseByHolder call")
	assert.Equal(t, "session:"+sess.ID.String(), holders[0])
}