diff --git a/internal/acp/handler_e2e_test.go b/internal/acp/handler_e2e_test.go index 1c0b5c9..57431ae 100644 --- a/internal/acp/handler_e2e_test.go +++ b/internal/acp/handler_e2e_test.go @@ -130,6 +130,7 @@ func TestHandler_Session_CreateAndWS_E2E(t *testing.T) { audit.NewPostgresRecorder(pool), disp, nil, nil, // wtSvc/wsRepo: IsMainWorktree=true path → not called + nil, // mcpTokens enc, acp.SupervisorConfig{ SpawnTimeout: 5 * time.Second, diff --git a/internal/acp/supervisor_test.go b/internal/acp/supervisor_test.go index d70a756..add3a35 100644 --- a/internal/acp/supervisor_test.go +++ b/internal/acp/supervisor_test.go @@ -23,12 +23,13 @@ import ( "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" ) // 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) *acp.Supervisor { +func newSupervisorForTest(t *testing.T, pool *pgxpool.Pool, repo acp.Repository, withAudit bool, mcpTokens mcp.TokenService) *acp.Supervisor { t.Helper() var rec audit.Recorder if withAudit { @@ -38,6 +39,7 @@ func newSupervisorForTest(t *testing.T, pool *pgxpool.Pool, repo acp.Repository, return acp.NewSupervisor( repo, rec, disp, nil, nil, // wtSvc / wsRepo: tests 用 IsMain=true,走 repo.ReleaseMainWorktree + mcpTokens, testEncryptor(t), acp.SupervisorConfig{ SpawnTimeout: 5 * time.Second, @@ -85,7 +87,7 @@ func TestSupervisor_HappyPath_Spawn_Handshake(t *testing.T) { require.NoError(t, err) require.True(t, ok) - sup := newSupervisorForTest(t, pool, repo, true) + sup := newSupervisorForTest(t, pool, repo, true, nil) proc, err := sup.Spawn(ctx, sess, k, nil) require.NoError(t, err) require.NotNil(t, proc) @@ -150,7 +152,7 @@ func TestSupervisor_HangAgent_HandshakeTimeout(t *testing.T) { disp := notify.NewDispatcher(notifyRepoStub{}, slogDiscard()) sup := acp.NewSupervisor( repo, audit.NewPostgresRecorder(pool), disp, - nil, nil, testEncryptor(t), + nil, nil, nil, testEncryptor(t), acp.SupervisorConfig{ SpawnTimeout: 1 * time.Second, KillGrace: 500 * time.Millisecond, @@ -214,7 +216,7 @@ func TestSupervisor_AgentCrashes(t *testing.T) { require.NoError(t, err) require.True(t, ok) - sup := newSupervisorForTest(t, pool, repo, true) + sup := newSupervisorForTest(t, pool, repo, true, nil) proc, err := sup.Spawn(ctx, sess, k, nil) require.NoError(t, err) require.NotNil(t, proc) @@ -243,3 +245,86 @@ func TestSupervisor_AgentCrashes(t *testing.T) { } } } + +// 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) 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) + 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") +} diff --git a/internal/app/acp_integration_test.go b/internal/app/acp_integration_test.go index 1194807..ba5ca63 100644 --- a/internal/app/acp_integration_test.go +++ b/internal/app/acp_integration_test.go @@ -175,7 +175,7 @@ func newACPIntegrationSuite(t *testing.T) *acpIntegrationSuite { EventTruncateField: 4096, ShutdownGrace: 5 * time.Second, } - sup := acp.NewSupervisor(acpRepo, auditRec, disp, nil, nil, enc, supCfg, log) + sup := acp.NewSupervisor(acpRepo, auditRec, disp, nil, nil, nil, enc, supCfg, log) // Minimal stubs for SessionService dependencies — sufficient for branch=main // path (no wtSvc / wsRepo calls) and tests that exercise quota/conflict logic. diff --git a/internal/app/app.go b/internal/app/app.go index 0130d5e..0b73dfb 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -333,7 +333,7 @@ func New(ctx context.Context, cfg *config.Config, dist embed.FS) (*App, error) { acpProjAccess := acpProjectAccessAdapter{wsRepo: wsRepo, projectRepo: projectRepo} acpSup := acp.NewSupervisor( acpRepo, auditRec, notifyDispatcher, - wtSvc, wsRepo, nil, enc, + wtSvc, wsRepo, mcpTokenSvc, enc, acp.SupervisorConfig{ SpawnTimeout: cfg.Acp.SpawnTimeout, KillGrace: cfg.Acp.KillGrace,