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:
@@ -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")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user