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

//go:build integration

package acp_test

import (
	"context"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	tcpg "github.com/testcontainers/testcontainers-go/modules/postgres"

	"github.com/yan1h/agent-coding-workflow/internal/acp"
	"github.com/yan1h/agent-coding-workflow/internal/infra/db"
	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
)

func setupRepo(t *testing.T) (acp.Repository, *pgxpool.Pool) {
	t.Helper()
	ctx := context.Background()
	container, err := tcpg.Run(ctx, "postgres:16-alpine",
		tcpg.WithDatabase("acw"),
		tcpg.WithUsername("acw"),
		tcpg.WithPassword("acw"),
	)
	require.NoError(t, err)
	t.Cleanup(func() { _ = container.Terminate(ctx) })

	dsn, err := container.ConnectionString(ctx, "sslmode=disable")
	require.NoError(t, err)
	require.NoError(t, db.Migrate(ctx, dsn, "file://../../migrations"))

	pool, err := db.NewPool(ctx, dsn)
	require.NoError(t, err)
	t.Cleanup(pool.Close)

	return acp.NewPostgresRepository(pool), pool
}

func TestPgRepo_AgentKind_CRUD(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "admin@local", true)

	// Create
	k, err := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "claude_code", DisplayName: "Claude Code",
		BinaryPath: "claude-code", Args: []string{"--acp"},
		EncryptedEnv: []byte("ciphertext"), Enabled: true, CreatedBy: uid,
	})
	require.NoError(t, err)
	assert.Equal(t, "claude_code", k.Name)

	// Get by ID
	got, err := repo.GetAgentKindByID(ctx, k.ID)
	require.NoError(t, err)
	assert.Equal(t, k.Name, got.Name)

	// Get by Name
	got2, err := repo.GetAgentKindByName(ctx, "claude_code")
	require.NoError(t, err)
	assert.Equal(t, k.ID, got2.ID)

	// List enabled
	enabled, err := repo.ListEnabledAgentKinds(ctx)
	require.NoError(t, err)
	require.Len(t, enabled, 1)

	// Update（enabled=false）
	k.DisplayName = "Claude Code v2"
	k.Enabled = false
	updated, err := repo.UpdateAgentKind(ctx, k)
	require.NoError(t, err)
	assert.Equal(t, "Claude Code v2", updated.DisplayName)
	assert.False(t, updated.Enabled)

	enabled, _ = repo.ListEnabledAgentKinds(ctx)
	assert.Empty(t, enabled)

	// Delete
	require.NoError(t, repo.DeleteAgentKind(ctx, k.ID))
	_, err = repo.GetAgentKindByID(ctx, k.ID)
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeAcpAgentKindNotFound, ae.Code)
}

func TestPgRepo_AgentKind_DeleteInUse(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "admin2@local", true)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)

	k, err := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "kind-in-use", DisplayName: "x",
		BinaryPath: "x", Enabled: true, CreatedBy: uid,
	})
	require.NoError(t, err)

	_, err = repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "main", CwdPath: "/tmp", IsMainWorktree: true, Status: acp.SessionRunning,
	})
	require.NoError(t, err)

	err = repo.DeleteAgentKind(ctx, k.ID)
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeAcpAgentKindInUse, ae.Code)
}

func TestPgRepo_Session_Lifecycle(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "u3@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)
	k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "kind3", DisplayName: "x",
		BinaryPath: "x", Enabled: true, CreatedBy: uid,
	})

	sess, err := repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "feat-x", CwdPath: "/tmp", IsMainWorktree: false, Status: acp.SessionStarting,
	})
	require.NoError(t, err)

	// Count active
	n, _ := repo.CountActiveSessions(ctx)
	assert.Equal(t, int64(1), n)
	n2, _ := repo.CountActiveSessionsByUser(ctx, uid)
	assert.Equal(t, int64(1), n2)

	// Update running
	require.NoError(t, repo.UpdateSessionRunning(ctx, sess.ID, "agent-sid-abc", 12345))
	got, _ := repo.GetSessionByID(ctx, sess.ID)
	assert.Equal(t, acp.SessionRunning, got.Status)
	require.NotNil(t, got.AgentSessionID)
	assert.Equal(t, "agent-sid-abc", *got.AgentSessionID)
	require.NotNil(t, got.PID)
	assert.Equal(t, int32(12345), *got.PID)

	// Update finished
	ec := int32(0)
	require.NoError(t, repo.UpdateSessionFinished(ctx, sess.ID, acp.SessionExited, &ec, nil))
	got, _ = repo.GetSessionByID(ctx, sess.ID)
	assert.Equal(t, acp.SessionExited, got.Status)
	assert.False(t, got.Status.IsActive())

	// Reset stuck（应无影响：本 session 已 exited）
	rows, _ := repo.ResetStuckSessionsOnRestart(ctx)
	assert.Empty(t, rows)
}

func TestPgRepo_Session_ResetStuck(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "u4@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)
	k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "kind4", DisplayName: "x",
		BinaryPath: "x", Enabled: true, CreatedBy: uid,
	})

	for i := 0; i < 2; i++ {
		_, _ = repo.InsertSession(ctx, &acp.Session{
			ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
			AgentKindID: k.ID, UserID: uid,
			Branch: "stuck", CwdPath: "/tmp", IsMainWorktree: false, Status: acp.SessionRunning,
		})
	}

	rows, err := repo.ResetStuckSessionsOnRestart(ctx)
	require.NoError(t, err)
	assert.Len(t, rows, 2)
	for _, s := range rows {
		assert.Equal(t, acp.SessionCrashed, s.Status)
	}
	rows2, _ := repo.ResetStuckSessionsOnRestart(ctx)
	assert.Empty(t, rows2)
}

func TestPgRepo_Event_OrderAndPurge(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "u5@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)
	k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "kind5", DisplayName: "x",
		BinaryPath: "x", Enabled: true, CreatedBy: uid,
	})
	sess, _ := repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "main", CwdPath: "/tmp", IsMainWorktree: true, Status: acp.SessionRunning,
	})

	for i := 0; i < 3; i++ {
		method := "session/update"
		_, err := repo.InsertEvent(ctx, &acp.Event{
			SessionID: sess.ID, Direction: acp.DirectionOut, RPCKind: acp.RPCNotification,
			Method: &method, Payload: []byte(`{"foo":"bar"}`), PayloadSize: 13, Truncated: false,
		})
		require.NoError(t, err)
	}

	// ListSince 0 → 全部 3 条
	all, err := repo.ListEventsSince(ctx, sess.ID, 0, 100)
	require.NoError(t, err)
	require.Len(t, all, 3)
	assert.True(t, all[0].ID < all[1].ID && all[1].ID < all[2].ID)

	// Since 第 1 条 ID → 仅 2 条
	rest, _ := repo.ListEventsSince(ctx, sess.ID, all[0].ID, 100)
	assert.Len(t, rest, 2)

	// Purge before now+1m → 全删
	n, _ := repo.PurgeEventsBefore(ctx, time.Now().Add(time.Minute))
	assert.Equal(t, int64(3), n)
}

func TestPgRepo_Turn_Lifecycle(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "u6@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)
	k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "kind6", DisplayName: "x",
		BinaryPath: "x", Enabled: true, CreatedBy: uid,
	})
	sess, _ := repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "main", CwdPath: "/tmp", IsMainWorktree: true, Status: acp.SessionRunning,
	})

	// 新 session 无回合
	latest, err := repo.GetLatestTurnBySession(ctx, sess.ID)
	require.NoError(t, err)
	assert.Nil(t, latest)

	// 插入回合 0
	pidStr := "client-init"
	turn, err := repo.InsertTurn(ctx, &acp.Turn{
		SessionID: sess.ID, TurnIndex: 0, PromptRequestID: &pidStr, Status: acp.TurnInProgress,
	})
	require.NoError(t, err)
	assert.Equal(t, 0, turn.TurnIndex)
	assert.Equal(t, acp.TurnInProgress, turn.Status)

	// IncrementTurnUpdateCount ×2（开放回合）
	require.NoError(t, repo.IncrementTurnUpdateCount(ctx, sess.ID, 0))
	require.NoError(t, repo.IncrementTurnUpdateCount(ctx, sess.ID, 0))

	// 完成回合：MarkTurnCompleted 用内存计数覆盖 update_count
	require.NoError(t, repo.MarkTurnCompleted(ctx, sess.ID, 0, "end_turn", 5))
	require.NoError(t, repo.UpdateSessionLastStopReason(ctx, sess.ID, "end_turn"))

	ts, err := repo.ListTurnsBySession(ctx, sess.ID)
	require.NoError(t, err)
	require.Len(t, ts, 1)
	assert.Equal(t, acp.TurnCompleted, ts[0].Status)
	require.NotNil(t, ts[0].StopReason)
	assert.Equal(t, "end_turn", *ts[0].StopReason)
	assert.Equal(t, 5, ts[0].UpdateCount)
	require.NotNil(t, ts[0].CompletedAt)

	// last_stop_reason round-trips on session SELECT
	got, _ := repo.GetSessionByID(ctx, sess.ID)
	require.NotNil(t, got.LastStopReason)
	assert.Equal(t, "end_turn", *got.LastStopReason)

	// GetLatestTurnBySession 返回 index 最大的回合
	latest, err = repo.GetLatestTurnBySession(ctx, sess.ID)
	require.NoError(t, err)
	require.NotNil(t, latest)
	assert.Equal(t, 0, latest.TurnIndex)

	// UNIQUE(session_id, turn_index) 约束
	_, err = repo.InsertTurn(ctx, &acp.Turn{SessionID: sess.ID, TurnIndex: 0, Status: acp.TurnInProgress})
	require.Error(t, err, "duplicate (session_id, turn_index) must violate UNIQUE")

	// 插入并中止回合 1
	_, err = repo.InsertTurn(ctx, &acp.Turn{SessionID: sess.ID, TurnIndex: 1, Status: acp.TurnInProgress})
	require.NoError(t, err)
	require.NoError(t, repo.MarkTurnAborted(ctx, sess.ID, 1))
	latest, _ = repo.GetLatestTurnBySession(ctx, sess.ID)
	assert.Equal(t, 1, latest.TurnIndex)
	assert.Equal(t, acp.TurnAborted, latest.Status)

	// Purge before now+1m → 全删
	purged, err := repo.PurgeTurnsBefore(ctx, time.Now().Add(time.Minute))
	require.NoError(t, err)
	assert.Equal(t, int64(2), purged)
	// 幂等
	purged2, _ := repo.PurgeTurnsBefore(ctx, time.Now().Add(time.Minute))
	assert.Equal(t, int64(0), purged2)
}

// ===== test helpers =====

func mustInsertUser(t *testing.T, ctx context.Context, pool *pgxpool.Pool, email string, admin bool) uuid.UUID {
	t.Helper()
	id := uuid.New()
	_, err := pool.Exec(ctx, `INSERT INTO users (id, email, password_hash, display_name, is_admin)
		VALUES ($1, $2, 'h', 'n', $3)`, id, email, admin)
	require.NoError(t, err)
	return id
}

func mustInsertProjectWorkspace(t *testing.T, ctx context.Context, pool *pgxpool.Pool, owner uuid.UUID) (uuid.UUID, uuid.UUID) {
	t.Helper()
	pid := uuid.New()
	wsid := uuid.New()
	slug := "p-" + pid.String()[:8]
	wsslug := "ws-" + wsid.String()[:8]
	_, err := pool.Exec(ctx, `INSERT INTO projects (id, slug, name, owner_id) VALUES ($1, $2, 'p', $3)`,
		pid, slug, owner)
	require.NoError(t, err)
	_, err = pool.Exec(ctx, `INSERT INTO workspaces (id, project_id, slug, name, git_remote_url, main_path)
		VALUES ($1, $2, $3, 'w', 'git@x:y.git', '/tmp')`, wsid, pid, wsslug)
	require.NoError(t, err)
	return pid, wsid
}

// TestPgRepo_UsagePersist_BumpsDashboardColumns 验证 AddSessionUsageTotals 同时
// 把 dashboard rollup 列（cost_usd / tokens_total）与运行时累加器对齐。
func TestPgRepo_UsagePersist_BumpsDashboardColumns(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	uid := mustInsertUser(t, ctx, pool, "dash-usage@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid)
	k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "kind-dash-usage", DisplayName: "x",
		BinaryPath: "x", Enabled: true, CreatedBy: uid,
	})
	sess, _ := repo.InsertSession(ctx, &acp.Session{
		ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
		AgentKindID: k.ID, UserID: uid,
		Branch: "main", CwdPath: "/tmp", IsMainWorktree: true, Status: acp.SessionRunning,
	})

	// 两次回合累加。
	_, _, _, _, err := repo.AddSessionUsageTotals(ctx, sess.ID, 100, 50, 10, 0.30)
	require.NoError(t, err)
	_, _, _, totCost, err := repo.AddSessionUsageTotals(ctx, sess.ID, 200, 40, 0, 0.20)
	require.NoError(t, err)
	assert.InDelta(t, 0.50, totCost, 1e-9)

	got, err := repo.GetSessionByID(ctx, sess.ID)
	require.NoError(t, err)
	assert.InDelta(t, 0.50, got.TotalCostUSD, 1e-9)
	// 运行时累加器 totals。
	assert.Equal(t, int64(300), got.PromptTokens)
	assert.Equal(t, int64(90), got.CompletionTokens)
	assert.Equal(t, int64(10), got.ThinkingTokens)

	// dashboard 汇总应读到与运行时一致的 cost / tokens。
	rollup, err := repo.SessionRollup(ctx, acp.DashboardFilter{
		UserID: &uid, From: time.Now().Add(-time.Hour), To: time.Now().Add(time.Hour),
	})
	require.NoError(t, err)
	assert.Equal(t, int64(1), rollup.Total)
	assert.Equal(t, int64(1), rollup.Active)
	assert.InDelta(t, 0.50, rollup.TotalCostUSD, 1e-9)
	assert.Equal(t, int64(400), rollup.TotalTokens) // 100+50+10 + 200+40+0
}

// TestPgRepo_Dashboard_RollupTimelineByDay 覆盖三个 dashboard 查询：
// SessionRollup（按状态计数 + 成本 + scope）、SessionsByDay、SessionTimeline。
func TestPgRepo_Dashboard_RollupTimelineByDay(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	repo, pool := setupRepo(t)

	owner := mustInsertUser(t, ctx, pool, "dash-owner@local", false)
	other := mustInsertUser(t, ctx, pool, "dash-other@local", false)
	pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, owner)
	k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{
		ID: uuid.New(), Name: "kind-dash", DisplayName: "x",
		BinaryPath: "x", Enabled: true, CreatedBy: owner,
	})

	mk := func(u uuid.UUID, status acp.SessionStatus) *acp.Session {
		s, err := repo.InsertSession(ctx, &acp.Session{
			ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid,
			AgentKindID: k.ID, UserID: u,
			Branch: "b", CwdPath: "/tmp", IsMainWorktree: false, Status: status,
		})
		require.NoError(t, err)
		return s
	}

	exited := mk(owner, acp.SessionRunning)
	crashed := mk(owner, acp.SessionRunning)
	mk(other, acp.SessionRunning) // 另一用户：owner-scope 下不计入

	// owner 的两个会话各累加成本并结束。
	_, _, _, _, err := repo.AddSessionUsageTotals(ctx, exited.ID, 100, 100, 0, 1.00)
	require.NoError(t, err)
	_, _, _, _, err = repo.AddSessionUsageTotals(ctx, crashed.ID, 50, 50, 0, 0.50)
	require.NoError(t, err)
	ec := int32(0)
	require.NoError(t, repo.UpdateSessionFinished(ctx, exited.ID, acp.SessionExited, &ec, nil))
	require.NoError(t, repo.UpdateSessionFinished(ctx, crashed.ID, acp.SessionCrashed, &ec, nil))

	from := time.Now().Add(-time.Hour)
	to := time.Now().Add(time.Hour)

	// owner scope：只统计 owner 的 2 个会话。
	ownerScope := acp.DashboardFilter{UserID: &owner, From: from, To: to}
	rollup, err := repo.SessionRollup(ctx, ownerScope)
	require.NoError(t, err)
	assert.Equal(t, int64(2), rollup.Total)
	assert.Equal(t, int64(1), rollup.Succeeded) // exited
	assert.Equal(t, int64(1), rollup.Crashed)
	assert.Equal(t, int64(0), rollup.Active)
	assert.InDelta(t, 1.50, rollup.TotalCostUSD, 1e-9)
	assert.True(t, rollup.AvgDurationSeconds >= 0)

	// admin scope（UserID nil）：跨用户 3 个会话。
	adminRollup, err := repo.SessionRollup(ctx, acp.DashboardFilter{From: from, To: to})
	require.NoError(t, err)
	assert.Equal(t, int64(3), adminRollup.Total)

	// SessionsByDay：owner scope 当天 1 行，total=2。
	byDay, err := repo.SessionsByDay(ctx, ownerScope)
	require.NoError(t, err)
	require.Len(t, byDay, 1)
	assert.Equal(t, int64(2), byDay[0].Total)
	assert.Equal(t, int64(1), byDay[0].Succeeded)
	assert.InDelta(t, 1.50, byDay[0].TotalCostUSD, 1e-9)

	// SessionTimeline：插入两类事件，按桶聚合。
	method := "session/update"
	for i := 0; i < 3; i++ {
		_, err := repo.InsertEvent(ctx, &acp.Event{
			SessionID: exited.ID, Direction: acp.DirectionOut, RPCKind: acp.RPCNotification,
			Method: &method, Payload: []byte(`{}`), PayloadSize: 2,
		})
		require.NoError(t, err)
	}
	reqMethod := "session/prompt"
	_, err = repo.InsertEvent(ctx, &acp.Event{
		SessionID: exited.ID, Direction: acp.DirectionOut, RPCKind: acp.RPCRequest,
		Method: &reqMethod, Payload: []byte(`{}`), PayloadSize: 2,
	})
	require.NoError(t, err)

	timeline, err := repo.SessionTimeline(ctx, exited.ID)
	require.NoError(t, err)
	require.Len(t, timeline, 2) // (out,notification,session/update) + (out,request,session/prompt)
	total := int64(0)
	for _, b := range timeline {
		total += b.EventCount
		assert.False(t, b.FirstAt.IsZero())
		assert.False(t, b.LastAt.IsZero())
	}
	assert.Equal(t, int64(4), total)
}