//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) }