agentic-coding-workflow/internal/acp/session_service_test.go

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"
)

// fakeProjectAccess 满足 acp.ProjectAccess 接口；测试时按需返回。
type fakeProjectAccess struct {
	pj    *project.Project
	issue *project.Issue
	req   *project.Requirement
	err   error
}

func (f fakeProjectAccess) GetProjectByWorkspace(_ context.Context, _ uuid.UUID) (*project.Project, error) {
	return f.pj, f.err
}
func (f fakeProjectAccess) GetIssueByNumber(_ context.Context, _ uuid.UUID, _ int) (*project.Issue, error) {
	if f.issue == nil {
		return nil, f.err
	}
	return f.issue, nil
}
func (f fakeProjectAccess) GetRequirementByNumber(_ context.Context, _ uuid.UUID, _ int) (*project.Requirement, error) {
	if f.req == nil {
		return nil, f.err
	}
	return f.req, nil
}

func TestResolveBranch_Explicit(t *testing.T) {
	t.Parallel()
	ws := &workspace.Workspace{ID: uuid.New(), DefaultBranch: "main"}
	br := "feat/x"
	got, isMain, err := acp.ResolveBranch(context.Background(), ws, uuid.New(),
		acp.CreateSessionInput{Branch: &br}, "demo", fakeProjectAccess{})
	require.NoError(t, err)
	assert.Equal(t, "feat/x", got)
	assert.False(t, isMain)
}

func TestResolveBranch_ExplicitMain(t *testing.T) {
	t.Parallel()
	ws := &workspace.Workspace{ID: uuid.New(), DefaultBranch: "main"}
	br := "main"
	got, isMain, err := acp.ResolveBranch(context.Background(), ws, uuid.New(),
		acp.CreateSessionInput{Branch: &br}, "demo", fakeProjectAccess{})
	require.NoError(t, err)
	assert.Equal(t, "main", got)
	assert.True(t, isMain)
}

func TestResolveBranch_FromIssue(t *testing.T) {
	t.Parallel()
	ws := &workspace.Workspace{ID: uuid.New(), ProjectID: uuid.New(), DefaultBranch: "main"}
	issueNum := 42
	pa := fakeProjectAccess{issue: &project.Issue{Number: 42, WorkspaceID: nil}}
	got, isMain, err := acp.ResolveBranch(context.Background(), ws, uuid.New(),
		acp.CreateSessionInput{IssueNumber: &issueNum}, "demo", pa)
	require.NoError(t, err)
	assert.Equal(t, "issue/demo-42", got)
	assert.False(t, isMain)
}

func TestResolveBranch_IssueLinkedToOtherWorkspace(t *testing.T) {
	t.Parallel()
	ws := &workspace.Workspace{ID: uuid.New(), ProjectID: uuid.New(), DefaultBranch: "main"}
	other := uuid.New()
	pa := fakeProjectAccess{issue: &project.Issue{Number: 42, WorkspaceID: &other}}
	num := 42
	_, _, err := acp.ResolveBranch(context.Background(), ws, uuid.New(),
		acp.CreateSessionInput{IssueNumber: &num}, "demo", pa)
	require.Error(t, err)
}

func TestResolveBranch_FromRequirement(t *testing.T) {
	t.Parallel()
	ws := &workspace.Workspace{ID: uuid.New(), ProjectID: uuid.New(), DefaultBranch: "main"}
	num := 7
	pa := fakeProjectAccess{req: &project.Requirement{Number: 7, WorkspaceID: nil}}
	got, isMain, err := acp.ResolveBranch(context.Background(), ws, uuid.New(),
		acp.CreateSessionInput{RequirementNumber: &num}, "demo", pa)
	require.NoError(t, err)
	assert.Equal(t, "req/demo-7", got)
	assert.False(t, isMain)
}

func TestResolveBranch_AutoFallback(t *testing.T) {
	t.Parallel()
	ws := &workspace.Workspace{ID: uuid.New(), ProjectID: uuid.New(), DefaultBranch: "main"}
	sid := uuid.New()
	got, isMain, err := acp.ResolveBranch(context.Background(), ws, sid,
		acp.CreateSessionInput{}, "demo", fakeProjectAccess{})
	require.NoError(t, err)
	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()
}