feat(acp): SessionService.Create issues system MCP token in tx

- 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 <noreply@anthropic.com>
This commit is contained in:
2026-05-08 13:01:16 +08:00
parent 146dca7738
commit 18f6c8bd2e
5 changed files with 365 additions and 17 deletions
+1
View File
@@ -152,6 +152,7 @@ func TestHandler_Session_CreateAndWS_E2E(t *testing.T) {
repo, wsStub, nil, nil, paStub, sup, repo, wsStub, nil, nil, paStub, sup,
audit.NewPostgresRecorder(pool), audit.NewPostgresRecorder(pool),
acp.SessionServiceConfig{MaxActiveGlobal: 10, MaxActivePerUser: 5}, acp.SessionServiceConfig{MaxActiveGlobal: 10, MaxActivePerUser: 5},
nil, // mcpTokens — not exercised by handler e2e
) )
h := acp.NewHandler( h := acp.NewHandler(
+40 -8
View File
@@ -10,17 +10,21 @@ import (
"time" "time"
"github.com/google/uuid" "github.com/google/uuid"
"github.com/jackc/pgx/v5"
"github.com/yan1h/agent-coding-workflow/internal/audit" "github.com/yan1h/agent-coding-workflow/internal/audit"
"github.com/yan1h/agent-coding-workflow/internal/infra/errs" "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/project"
"github.com/yan1h/agent-coding-workflow/internal/workspace" "github.com/yan1h/agent-coding-workflow/internal/workspace"
) )
// SessionServiceConfig 是 service 的限流配置(从 AcpConfig 派生)。 // SessionServiceConfig 是 service 的限流 + MCP 配置(从 AcpConfig / MCPConfig 派生)。
type SessionServiceConfig struct { type SessionServiceConfig struct {
MaxActiveGlobal int MaxActiveGlobal int
MaxActivePerUser int MaxActivePerUser int
SystemTokenTTL time.Duration
MCPPublicURL string
} }
// ProjectAccess 是 acp 模块对 project 模块的窄接口;adapter 在 app.go 实现。 // ProjectAccess 是 acp 模块对 project 模块的窄接口;adapter 在 app.go 实现。
@@ -39,16 +43,17 @@ type sessionService struct {
sup *Supervisor sup *Supervisor
audit audit.Recorder audit audit.Recorder
cfg SessionServiceConfig cfg SessionServiceConfig
mcpTokens mcp.TokenService // spec §6.1 — system token issuance
} }
// NewSessionService 构造 SessionService。 // NewSessionService 构造 SessionService。
func NewSessionService(repo Repository, wsSvc workspace.WorkspaceService, func NewSessionService(repo Repository, wsSvc workspace.WorkspaceService,
wtSvc workspace.WorktreeService, wsRepo workspace.Repository, wtSvc workspace.WorktreeService, wsRepo workspace.Repository,
pa ProjectAccess, sup *Supervisor, rec audit.Recorder, pa ProjectAccess, sup *Supervisor, rec audit.Recorder,
cfg SessionServiceConfig) SessionService { cfg SessionServiceConfig, mcpTokens mcp.TokenService) SessionService {
return &sessionService{ return &sessionService{
repo: repo, wsSvc: wsSvc, wtSvc: wtSvc, wsRepo: wsRepo, 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{ sess := &Session{
ID: sessionID, WorkspaceID: ws.ID, ProjectID: ws.ProjectID, ID: sessionID, WorkspaceID: ws.ID, ProjectID: ws.ProjectID,
AgentKindID: kind.ID, UserID: c.UserID, 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, Branch: branch, CwdPath: cwdPath, IsMainWorktree: isMain,
Status: SessionStarting, Status: SessionStarting,
} }
out, err := s.repo.InsertSession(ctx, sess) 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 { 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) s.releaseOnFailure(ctx, sess, sessCaller, worktreeID)
return nil, err return nil, txErr
} }
// 8. audit // 8. audit
@@ -215,11 +242,16 @@ func (s *sessionService) Create(ctx context.Context, c Caller, in CreateSessionI
"requirement_id": reqID, "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() { go func() {
spawnCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second) spawnCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel() 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) s.releaseOnFailure(spawnCtx, out, sessCaller, worktreeID)
} }
}() }()
+311
View File
@@ -2,13 +2,18 @@ package acp_test
import ( import (
"context" "context"
"fmt"
"sync"
"testing" "testing"
"time"
"github.com/google/uuid" "github.com/google/uuid"
"github.com/jackc/pgx/v5"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
"github.com/yan1h/agent-coding-workflow/internal/acp" "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/project"
"github.com/yan1h/agent-coding-workflow/internal/workspace" "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.Contains(t, got, "acp-auto/")
assert.False(t, isMain) 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()
}
+1
View File
@@ -185,6 +185,7 @@ func newACPIntegrationSuite(t *testing.T) *acpIntegrationSuite {
sessSvc := acp.NewSessionService( sessSvc := acp.NewSessionService(
acpRepo, wsStub, nil, nil, paStub, sup, auditRec, acpRepo, wsStub, nil, nil, paStub, sup, auditRec,
acp.SessionServiceConfig{MaxActiveGlobal: 20, MaxActivePerUser: 3}, acp.SessionServiceConfig{MaxActiveGlobal: 20, MaxActivePerUser: 3},
nil, // mcpTokens — not exercised by integration tests
) )
r := chi.NewRouter() r := chi.NewRouter()
+3
View File
@@ -351,7 +351,10 @@ func New(ctx context.Context, cfg *config.Config, dist embed.FS) (*App, error) {
acp.SessionServiceConfig{ acp.SessionServiceConfig{
MaxActiveGlobal: cfg.Acp.MaxActiveGlobal, MaxActiveGlobal: cfg.Acp.MaxActiveGlobal,
MaxActivePerUser: cfg.Acp.MaxActivePerUser, 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{ acp.NewHandler(akSvc, sessSvc, acpSup, acpRepo, userSvc, userAdminAdapter{svc: userSvc}, enc, acp.WSConfig{
PingInterval: cfg.Acp.WSPingInterval, PingInterval: cfg.Acp.WSPingInterval,