You've already forked agentic-coding-workflow
feat(acp): Supervisor.onExit revokes system MCP token (spec §6.3)
- mcpTokens dep added to Supervisor; nil-safe (test fixtures may skip) - RevokeBySession called after audit Record in onExit - app.go wires mcpTokenSvc (was nil placeholder) - Integration test: fakeSupMCPTokens verifies onExit calls RevokeBySession Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -130,6 +130,7 @@ func TestHandler_Session_CreateAndWS_E2E(t *testing.T) {
|
|||||||
audit.NewPostgresRecorder(pool),
|
audit.NewPostgresRecorder(pool),
|
||||||
disp,
|
disp,
|
||||||
nil, nil, // wtSvc/wsRepo: IsMainWorktree=true path → not called
|
nil, nil, // wtSvc/wsRepo: IsMainWorktree=true path → not called
|
||||||
|
nil, // mcpTokens
|
||||||
enc,
|
enc,
|
||||||
acp.SupervisorConfig{
|
acp.SupervisorConfig{
|
||||||
SpawnTimeout: 5 * time.Second,
|
SpawnTimeout: 5 * time.Second,
|
||||||
|
|||||||
@@ -23,12 +23,13 @@ import (
|
|||||||
"github.com/yan1h/agent-coding-workflow/internal/acp"
|
"github.com/yan1h/agent-coding-workflow/internal/acp"
|
||||||
"github.com/yan1h/agent-coding-workflow/internal/audit"
|
"github.com/yan1h/agent-coding-workflow/internal/audit"
|
||||||
"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
|
"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
|
||||||
|
"github.com/yan1h/agent-coding-workflow/internal/mcp"
|
||||||
)
|
)
|
||||||
|
|
||||||
// newSupervisorForTest 构造一个绑定到 PG 真实 audit Recorder 的 Supervisor,
|
// newSupervisorForTest 构造一个绑定到 PG 真实 audit Recorder 的 Supervisor,
|
||||||
// 使用静音 logger + 默认 SupervisorConfig(5s SpawnTimeout 给 fake agent 留 buffer)。
|
// 使用静音 logger + 默认 SupervisorConfig(5s SpawnTimeout 给 fake agent 留 buffer)。
|
||||||
// withAudit=false 时 Recorder 为 nil,对应 onExit 跳过 audit 路径。
|
// 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()
|
t.Helper()
|
||||||
var rec audit.Recorder
|
var rec audit.Recorder
|
||||||
if withAudit {
|
if withAudit {
|
||||||
@@ -38,6 +39,7 @@ func newSupervisorForTest(t *testing.T, pool *pgxpool.Pool, repo acp.Repository,
|
|||||||
return acp.NewSupervisor(
|
return acp.NewSupervisor(
|
||||||
repo, rec, disp,
|
repo, rec, disp,
|
||||||
nil, nil, // wtSvc / wsRepo: tests 用 IsMain=true,走 repo.ReleaseMainWorktree
|
nil, nil, // wtSvc / wsRepo: tests 用 IsMain=true,走 repo.ReleaseMainWorktree
|
||||||
|
mcpTokens,
|
||||||
testEncryptor(t),
|
testEncryptor(t),
|
||||||
acp.SupervisorConfig{
|
acp.SupervisorConfig{
|
||||||
SpawnTimeout: 5 * time.Second,
|
SpawnTimeout: 5 * time.Second,
|
||||||
@@ -85,7 +87,7 @@ func TestSupervisor_HappyPath_Spawn_Handshake(t *testing.T) {
|
|||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
require.True(t, ok)
|
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)
|
proc, err := sup.Spawn(ctx, sess, k, nil)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
require.NotNil(t, proc)
|
require.NotNil(t, proc)
|
||||||
@@ -150,7 +152,7 @@ func TestSupervisor_HangAgent_HandshakeTimeout(t *testing.T) {
|
|||||||
disp := notify.NewDispatcher(notifyRepoStub{}, slogDiscard())
|
disp := notify.NewDispatcher(notifyRepoStub{}, slogDiscard())
|
||||||
sup := acp.NewSupervisor(
|
sup := acp.NewSupervisor(
|
||||||
repo, audit.NewPostgresRecorder(pool), disp,
|
repo, audit.NewPostgresRecorder(pool), disp,
|
||||||
nil, nil, testEncryptor(t),
|
nil, nil, nil, testEncryptor(t),
|
||||||
acp.SupervisorConfig{
|
acp.SupervisorConfig{
|
||||||
SpawnTimeout: 1 * time.Second,
|
SpawnTimeout: 1 * time.Second,
|
||||||
KillGrace: 500 * time.Millisecond,
|
KillGrace: 500 * time.Millisecond,
|
||||||
@@ -214,7 +216,7 @@ func TestSupervisor_AgentCrashes(t *testing.T) {
|
|||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
require.True(t, ok)
|
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)
|
proc, err := sup.Spawn(ctx, sess, k, nil)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
require.NotNil(t, proc)
|
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")
|
||||||
|
}
|
||||||
|
|||||||
@@ -175,7 +175,7 @@ func newACPIntegrationSuite(t *testing.T) *acpIntegrationSuite {
|
|||||||
EventTruncateField: 4096,
|
EventTruncateField: 4096,
|
||||||
ShutdownGrace: 5 * time.Second,
|
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
|
// Minimal stubs for SessionService dependencies — sufficient for branch=main
|
||||||
// path (no wtSvc / wsRepo calls) and tests that exercise quota/conflict logic.
|
// path (no wtSvc / wsRepo calls) and tests that exercise quota/conflict logic.
|
||||||
|
|||||||
+1
-1
@@ -333,7 +333,7 @@ func New(ctx context.Context, cfg *config.Config, dist embed.FS) (*App, error) {
|
|||||||
acpProjAccess := acpProjectAccessAdapter{wsRepo: wsRepo, projectRepo: projectRepo}
|
acpProjAccess := acpProjectAccessAdapter{wsRepo: wsRepo, projectRepo: projectRepo}
|
||||||
acpSup := acp.NewSupervisor(
|
acpSup := acp.NewSupervisor(
|
||||||
acpRepo, auditRec, notifyDispatcher,
|
acpRepo, auditRec, notifyDispatcher,
|
||||||
wtSvc, wsRepo, nil, enc,
|
wtSvc, wsRepo, mcpTokenSvc, enc,
|
||||||
acp.SupervisorConfig{
|
acp.SupervisorConfig{
|
||||||
SpawnTimeout: cfg.Acp.SpawnTimeout,
|
SpawnTimeout: cfg.Acp.SpawnTimeout,
|
||||||
KillGrace: cfg.Acp.KillGrace,
|
KillGrace: cfg.Acp.KillGrace,
|
||||||
|
|||||||
Reference in New Issue
Block a user