From 18f6c8bd2ecd94dd5b04c36732cb1cf29e991c81 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Fri, 8 May 2026 13:01:16 +0800 Subject: [PATCH] feat(acp): SessionService.Create issues system MCP token in tx MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - InTx wraps InsertSession + IssueSystemTokenWithTx (spec §6.1 atomic guarantee) - Tx failure -> releaseOnFailure (worktree rollback) - Spawn called with extraEnv = ACW_MCP_TOKEN + ACW_MCP_URL - Spawn failure -> RevokeBySession (idempotent) + releaseOnFailure - app.go SessionService construction updated with mcpTokens + config - Unit tests: token issuance happy path + rollback on token failure Co-Authored-By: Claude Opus 4.6 --- internal/acp/handler_e2e_test.go | 1 + internal/acp/session_service.go | 66 ++++-- internal/acp/session_service_test.go | 311 +++++++++++++++++++++++++++ internal/app/acp_integration_test.go | 1 + internal/app/app.go | 3 + 5 files changed, 365 insertions(+), 17 deletions(-) diff --git a/internal/acp/handler_e2e_test.go b/internal/acp/handler_e2e_test.go index 57431ae..a6fea03 100644 --- a/internal/acp/handler_e2e_test.go +++ b/internal/acp/handler_e2e_test.go @@ -152,6 +152,7 @@ func TestHandler_Session_CreateAndWS_E2E(t *testing.T) { repo, wsStub, nil, nil, paStub, sup, audit.NewPostgresRecorder(pool), acp.SessionServiceConfig{MaxActiveGlobal: 10, MaxActivePerUser: 5}, + nil, // mcpTokens — not exercised by handler e2e ) h := acp.NewHandler( diff --git a/internal/acp/session_service.go b/internal/acp/session_service.go index 83d2a81..9e4263a 100644 --- a/internal/acp/session_service.go +++ b/internal/acp/session_service.go @@ -10,17 +10,21 @@ import ( "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/yan1h/agent-coding-workflow/internal/audit" "github.com/yan1h/agent-coding-workflow/internal/infra/errs" + "github.com/yan1h/agent-coding-workflow/internal/mcp" "github.com/yan1h/agent-coding-workflow/internal/project" "github.com/yan1h/agent-coding-workflow/internal/workspace" ) -// SessionServiceConfig 是 service 的限流配置（从 AcpConfig 派生）。 +// SessionServiceConfig 是 service 的限流 + MCP 配置（从 AcpConfig / MCPConfig 派生）。 type SessionServiceConfig struct { MaxActiveGlobal int MaxActivePerUser int + SystemTokenTTL time.Duration + MCPPublicURL string } // ProjectAccess 是 acp 模块对 project 模块的窄接口；adapter 在 app.go 实现。 @@ -31,24 +35,25 @@ type ProjectAccess interface { } type sessionService struct { - repo Repository - wsSvc workspace.WorkspaceService - wtSvc workspace.WorktreeService - wsRepo workspace.Repository - pa ProjectAccess - sup *Supervisor - audit audit.Recorder - cfg SessionServiceConfig + repo Repository + wsSvc workspace.WorkspaceService + wtSvc workspace.WorktreeService + wsRepo workspace.Repository + pa ProjectAccess + sup *Supervisor + audit audit.Recorder + cfg SessionServiceConfig + mcpTokens mcp.TokenService // spec §6.1 — system token issuance } // NewSessionService 构造 SessionService。 func NewSessionService(repo Repository, wsSvc workspace.WorkspaceService, wtSvc workspace.WorktreeService, wsRepo workspace.Repository, pa ProjectAccess, sup *Supervisor, rec audit.Recorder, - cfg SessionServiceConfig) SessionService { + cfg SessionServiceConfig, mcpTokens mcp.TokenService) SessionService { return &sessionService{ repo: repo, wsSvc: wsSvc, wtSvc: wtSvc, wsRepo: wsRepo, - pa: pa, sup: sup, audit: rec, cfg: cfg, + pa: pa, sup: sup, audit: rec, cfg: cfg, mcpTokens: mcpTokens, } } @@ -192,7 +197,7 @@ func (s *sessionService) Create(ctx context.Context, c Caller, in CreateSessionI } } - // 7. INSERT + // 7. INSERT acp_sessions + system MCP token（spec §6.1 atomic guarantee） sess := &Session{ ID: sessionID, WorkspaceID: ws.ID, ProjectID: ws.ProjectID, AgentKindID: kind.ID, UserID: c.UserID, @@ -200,10 +205,32 @@ func (s *sessionService) Create(ctx context.Context, c Caller, in CreateSessionI Branch: branch, CwdPath: cwdPath, IsMainWorktree: isMain, Status: SessionStarting, } - out, err := s.repo.InsertSession(ctx, sess) - if err != nil { + var ( + out *Session + tokenRes mcp.IssueResult + ) + txErr := s.repo.InTx(ctx, func(txCtx context.Context, tx pgx.Tx) error { + inserted, err := s.repo.WithTx(tx).InsertSession(txCtx, sess) + if err != nil { + return err + } + out = inserted + + res, err := s.mcpTokens.IssueSystemTokenWithTx(txCtx, tx, mcp.IssueSystemTokenInput{ + UserID: c.UserID, + ACPSessionID: inserted.ID, + ProjectIDs: []uuid.UUID{ws.ProjectID}, + ExpiresIn: s.cfg.SystemTokenTTL, + }) + if err != nil { + return err + } + tokenRes = res + return nil + }) + if txErr != nil { s.releaseOnFailure(ctx, sess, sessCaller, worktreeID) - return nil, err + return nil, txErr } // 8. audit @@ -215,11 +242,16 @@ func (s *sessionService) Create(ctx context.Context, c Caller, in CreateSessionI "requirement_id": reqID, }) - // 9. async spawn（失败时回滚 worktree） + // 9. async spawn（失败时回滚 worktree + revoke MCP token） + extraEnv := map[string]string{ + "ACW_MCP_TOKEN": tokenRes.Plaintext, + "ACW_MCP_URL": s.cfg.MCPPublicURL, + } go func() { spawnCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() - if _, err := s.sup.Spawn(spawnCtx, out, kind, nil); err != nil { + if _, err := s.sup.Spawn(spawnCtx, out, kind, extraEnv); err != nil { + _ = s.mcpTokens.RevokeBySession(spawnCtx, out.ID) s.releaseOnFailure(spawnCtx, out, sessCaller, worktreeID) } }() diff --git a/internal/acp/session_service_test.go b/internal/acp/session_service_test.go index 439aed4..4223473 100644 --- a/internal/acp/session_service_test.go +++ b/internal/acp/session_service_test.go @@ -2,13 +2,18 @@ package acp_test import ( "context" + "fmt" + "sync" "testing" + "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5" "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/mcp" "github.com/yan1h/agent-coding-workflow/internal/project" "github.com/yan1h/agent-coding-workflow/internal/workspace" ) @@ -104,3 +109,309 @@ func TestResolveBranch_AutoFallback(t *testing.T) { assert.Contains(t, got, "acp-auto/") assert.False(t, isMain) } + +// ===== fakes for SessionService.Create tests ===== + +// fakeAcpRepo satisfies acp.Repository for Create-path unit tests. +type fakeAcpRepo struct { + mu sync.Mutex + sessions []*acp.Session + kinds []*acp.AgentKind + inserted *acp.Session + + // txStaging holds sessions inserted during InTx; committed to sessions only + // when fn returns nil (simulating real transaction semantics). + txStaging []*acp.Session + inTx bool +} + +func (r *fakeAcpRepo) InTx(_ context.Context, fn func(context.Context, pgx.Tx) error) error { + r.mu.Lock() + r.inTx = true + r.txStaging = nil + r.mu.Unlock() + + err := fn(context.Background(), nil) + + r.mu.Lock() + defer r.mu.Unlock() + r.inTx = false + if err == nil { + r.sessions = append(r.sessions, r.txStaging...) + } + r.txStaging = nil + return err +} +func (r *fakeAcpRepo) WithTx(_ pgx.Tx) acp.Repository { return r } + +func (r *fakeAcpRepo) InsertSession(_ context.Context, s *acp.Session) (*acp.Session, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.inserted = s + if r.inTx { + r.txStaging = append(r.txStaging, s) + } else { + r.sessions = append(r.sessions, s) + } + return s, nil +} +func (r *fakeAcpRepo) GetSessionByID(_ context.Context, id uuid.UUID) (*acp.Session, error) { + r.mu.Lock() + defer r.mu.Unlock() + for _, s := range r.sessions { + if s.ID == id { + return s, nil + } + } + return nil, fmt.Errorf("not found") +} +func (r *fakeAcpRepo) GetAgentKindByID(_ context.Context, id uuid.UUID) (*acp.AgentKind, error) { + for _, k := range r.kinds { + if k.ID == id { + return k, nil + } + } + return nil, fmt.Errorf("not found") +} +func (r *fakeAcpRepo) CountActiveSessions(_ context.Context) (int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + var n int64 + for _, s := range r.sessions { + if s.Status == acp.SessionStarting || s.Status == acp.SessionRunning { + n++ + } + } + return n, nil +} +func (r *fakeAcpRepo) CountActiveSessionsByUser(_ context.Context, _ uuid.UUID) (int64, error) { + return 0, nil +} +func (r *fakeAcpRepo) AcquireMainWorktree(_ context.Context, _, _ uuid.UUID) (bool, error) { + return true, nil +} +func (r *fakeAcpRepo) ReleaseMainWorktree(_ context.Context, _, _ uuid.UUID) (bool, error) { + return true, nil +} + +// Stub out remaining Repository methods (not used by Create path). +func (r *fakeAcpRepo) CreateAgentKind(_ context.Context, k *acp.AgentKind) (*acp.AgentKind, error) { + return k, nil +} +func (r *fakeAcpRepo) GetAgentKindByName(_ context.Context, _ string) (*acp.AgentKind, error) { + return nil, fmt.Errorf("not found") +} +func (r *fakeAcpRepo) ListAgentKinds(_ context.Context) ([]*acp.AgentKind, error) { return nil, nil } +func (r *fakeAcpRepo) ListEnabledAgentKinds(_ context.Context) ([]*acp.AgentKind, error) { + return r.kinds, nil +} +func (r *fakeAcpRepo) UpdateAgentKind(_ context.Context, k *acp.AgentKind) (*acp.AgentKind, error) { + return k, nil +} +func (r *fakeAcpRepo) DeleteAgentKind(_ context.Context, _ uuid.UUID) error { return nil } +func (r *fakeAcpRepo) CountAgentKindUsage(_ context.Context, _ uuid.UUID) (int64, error) { return 0, nil } +func (r *fakeAcpRepo) ListSessionsByUser(_ context.Context, _ uuid.UUID) ([]*acp.Session, error) { + return nil, nil +} +func (r *fakeAcpRepo) ListAllSessions(_ context.Context) ([]*acp.Session, error) { return nil, nil } +func (r *fakeAcpRepo) UpdateSessionRunning(_ context.Context, _ uuid.UUID, _ string, _ int32) error { + return nil +} +func (r *fakeAcpRepo) UpdateSessionFinished(_ context.Context, _ uuid.UUID, _ acp.SessionStatus, _ *int32, _ *string) error { + return nil +} +func (r *fakeAcpRepo) ResetStuckSessionsOnRestart(_ context.Context) ([]*acp.Session, error) { + return nil, nil +} +func (r *fakeAcpRepo) InsertEvent(_ context.Context, e *acp.Event) (*acp.Event, error) { return e, nil } +func (r *fakeAcpRepo) ListEventsSince(_ context.Context, _ uuid.UUID, _ int64, _ int32) ([]*acp.Event, error) { + return nil, nil +} +func (r *fakeAcpRepo) PurgeEventsBefore(_ context.Context, _ time.Time) (int64, error) { return 0, nil } + +// fakeMCPTokenSvc satisfies mcp.TokenService for unit tests. +type fakeMCPTokenSvc struct { + mu sync.Mutex + issued []mcp.IssueSystemTokenInput + issueErr error + revokedBySess []uuid.UUID + revokeBySessFn func(ctx context.Context, sessionID uuid.UUID) error +} + +func (f *fakeMCPTokenSvc) IssueAdmin(_ context.Context, _ mcp.IssueAdminInput) (mcp.IssueResult, error) { + return mcp.IssueResult{}, nil +} +func (f *fakeMCPTokenSvc) IssueSystemToken(_ context.Context, _ mcp.IssueSystemTokenInput) (mcp.IssueResult, error) { + return mcp.IssueResult{}, nil +} +func (f *fakeMCPTokenSvc) IssueSystemTokenWithTx(_ context.Context, _ pgx.Tx, in mcp.IssueSystemTokenInput) (mcp.IssueResult, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.issueErr != nil { + return mcp.IssueResult{}, f.issueErr + } + f.issued = append(f.issued, in) + return mcp.IssueResult{ID: uuid.New(), Plaintext: "tok_" + in.ACPSessionID.String()[:8]}, nil +} +func (f *fakeMCPTokenSvc) Authenticate(_ context.Context, _ string) (*mcp.Token, error) { + return nil, nil +} +func (f *fakeMCPTokenSvc) Revoke(_ context.Context, _, _ uuid.UUID) error { return nil } +func (f *fakeMCPTokenSvc) RevokeBySession(_ context.Context, sessionID uuid.UUID) error { + f.mu.Lock() + defer f.mu.Unlock() + f.revokedBySess = append(f.revokedBySess, sessionID) + if f.revokeBySessFn != nil { + return f.revokeBySessFn(context.Background(), sessionID) + } + return nil +} +func (f *fakeMCPTokenSvc) List(_ context.Context, _ *uuid.UUID, _ bool) ([]*mcp.Token, error) { + return nil, nil +} +func (f *fakeMCPTokenSvc) GetByID(_ context.Context, _ uuid.UUID) (*mcp.Token, error) { + return nil, nil +} + +// fakeWsSvc satisfies workspace.WorkspaceService for Create-path tests. +type fakeWsSvc struct { + ws *workspace.Workspace +} + +func (s *fakeWsSvc) Get(_ context.Context, _ workspace.Caller, _ uuid.UUID) (*workspace.Workspace, error) { + return s.ws, nil +} +func (s *fakeWsSvc) Create(context.Context, workspace.Caller, string, workspace.CreateWorkspaceInput) (*workspace.Workspace, error) { + return nil, fmt.Errorf("unused") +} +func (s *fakeWsSvc) List(context.Context, workspace.Caller, string) ([]*workspace.Workspace, error) { + return nil, fmt.Errorf("unused") +} +func (s *fakeWsSvc) Update(context.Context, workspace.Caller, uuid.UUID, workspace.UpdateWorkspaceInput) (*workspace.Workspace, error) { + return nil, fmt.Errorf("unused") +} +func (s *fakeWsSvc) Delete(context.Context, workspace.Caller, uuid.UUID) error { + return fmt.Errorf("unused") +} +func (s *fakeWsSvc) Sync(context.Context, workspace.Caller, uuid.UUID) (*workspace.Workspace, error) { + return nil, fmt.Errorf("unused") +} +func (s *fakeWsSvc) UpsertCredential(context.Context, workspace.Caller, uuid.UUID, workspace.UpsertCredentialInput) (*workspace.Credential, error) { + return nil, fmt.Errorf("unused") +} +func (s *fakeWsSvc) GetCredentialMeta(context.Context, workspace.Caller, uuid.UUID) (*workspace.Credential, error) { + return nil, fmt.Errorf("unused") +} + +// ===== Tests ===== + +func TestSessionService_Create_IssuesSystemToken(t *testing.T) { + t.Parallel() + + uid := uuid.New() + wsID := uuid.New() + pid := uuid.New() + akID := uuid.New() + + repo := &fakeAcpRepo{ + kinds: []*acp.AgentKind{ + {ID: akID, Name: "claude", Enabled: true}, + }, + } + wsSvc := &fakeWsSvc{ws: &workspace.Workspace{ + ID: wsID, ProjectID: pid, DefaultBranch: "main", MainPath: "/tmp/main", + }} + pa := fakeProjectAccess{pj: &project.Project{ID: pid, Slug: "test-proj"}} + mcpTokens := &fakeMCPTokenSvc{} + + sup := acp.NewSupervisor( + repo, nil, nil, nil, nil, mcpTokens, nil, + acp.SupervisorConfig{SpawnTimeout: 5 * time.Second, KillGrace: 1 * time.Second}, + nil, // logger + ) + + svc := acp.NewSessionService( + repo, wsSvc, nil, nil, pa, sup, nil, + acp.SessionServiceConfig{ + MaxActiveGlobal: 10, + MaxActivePerUser: 5, + SystemTokenTTL: 24 * time.Hour, + MCPPublicURL: "http://localhost:8080/mcp", + }, + mcpTokens, + ) + + branch := "main" + sess, err := svc.Create(context.Background(), acp.Caller{UserID: uid}, acp.CreateSessionInput{ + WorkspaceID: wsID, + AgentKindID: akID, + Branch: &branch, + }) + require.NoError(t, err) + require.NotNil(t, sess) + assert.Equal(t, wsID, sess.WorkspaceID) + assert.Equal(t, uid, sess.UserID) + + // Verify system token was issued inside the tx + mcpTokens.mu.Lock() + require.Len(t, mcpTokens.issued, 1) + assert.Equal(t, uid, mcpTokens.issued[0].UserID) + assert.Equal(t, sess.ID, mcpTokens.issued[0].ACPSessionID) + assert.Equal(t, []uuid.UUID{pid}, mcpTokens.issued[0].ProjectIDs) + assert.Equal(t, 24*time.Hour, mcpTokens.issued[0].ExpiresIn) + mcpTokens.mu.Unlock() +} + +func TestSessionService_Create_TokenIssueFails_RollsBack(t *testing.T) { + t.Parallel() + + uid := uuid.New() + wsID := uuid.New() + pid := uuid.New() + akID := uuid.New() + + repo := &fakeAcpRepo{ + kinds: []*acp.AgentKind{ + {ID: akID, Name: "claude", Enabled: true}, + }, + } + wsSvc := &fakeWsSvc{ws: &workspace.Workspace{ + ID: wsID, ProjectID: pid, DefaultBranch: "main", MainPath: "/tmp/main", + }} + pa := fakeProjectAccess{pj: &project.Project{ID: pid, Slug: "test-proj"}} + mcpTokens := &fakeMCPTokenSvc{ + issueErr: fmt.Errorf("token issue boom"), + } + + sup := acp.NewSupervisor( + repo, nil, nil, nil, nil, mcpTokens, nil, + acp.SupervisorConfig{SpawnTimeout: 5 * time.Second, KillGrace: 1 * time.Second}, + nil, + ) + + svc := acp.NewSessionService( + repo, wsSvc, nil, nil, pa, sup, nil, + acp.SessionServiceConfig{ + MaxActiveGlobal: 10, + MaxActivePerUser: 5, + SystemTokenTTL: 24 * time.Hour, + MCPPublicURL: "http://localhost:8080/mcp", + }, + mcpTokens, + ) + + branch := "main" + sess, err := svc.Create(context.Background(), acp.Caller{UserID: uid}, acp.CreateSessionInput{ + WorkspaceID: wsID, + AgentKindID: akID, + Branch: &branch, + }) + require.Error(t, err) + assert.Nil(t, sess) + assert.Contains(t, err.Error(), "token issue boom") + + // Verify no session was committed (InTx rolled back) + repo.mu.Lock() + assert.Empty(t, repo.sessions) + repo.mu.Unlock() +} diff --git a/internal/app/acp_integration_test.go b/internal/app/acp_integration_test.go index ba5ca63..d393e70 100644 --- a/internal/app/acp_integration_test.go +++ b/internal/app/acp_integration_test.go @@ -185,6 +185,7 @@ func newACPIntegrationSuite(t *testing.T) *acpIntegrationSuite { sessSvc := acp.NewSessionService( acpRepo, wsStub, nil, nil, paStub, sup, auditRec, acp.SessionServiceConfig{MaxActiveGlobal: 20, MaxActivePerUser: 3}, + nil, // mcpTokens — not exercised by integration tests ) r := chi.NewRouter() diff --git a/internal/app/app.go b/internal/app/app.go index 0b73dfb..52c806d 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -351,7 +351,10 @@ func New(ctx context.Context, cfg *config.Config, dist embed.FS) (*App, error) { acp.SessionServiceConfig{ MaxActiveGlobal: cfg.Acp.MaxActiveGlobal, MaxActivePerUser: cfg.Acp.MaxActivePerUser, + SystemTokenTTL: cfg.MCP.SystemTokenTTL, + MCPPublicURL: cfg.MCP.PublicURL, }, + mcpTokenSvc, ) acp.NewHandler(akSvc, sessSvc, acpSup, acpRepo, userSvc, userAdminAdapter{svc: userSvc}, enc, acp.WSConfig{ PingInterval: cfg.Acp.WSPingInterval,