You've already forked agentic-coding-workflow
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:
@@ -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()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user