test(mcp): RevokeBySession + ReapStaleSystemTokens integration

2026-05-08 07:27:59 +08:00
parent d431b8103d
commit 4071778930
1 changed files with 117 additions and 0 deletions
@@ -271,3 +271,120 @@ func TestPgRepo_PurgeExpired(t *testing.T) {
 	require.True(t, ok)
 	assert.Equal(t, errs.CodeMcpTokenNotFound, ae.Code)
 }
+
+// mustInsertACPSetup 准备测 system token 需要的 fixtures：
+// project + workspace + agent_kind + acp_sessions row（status 任选）。
+// 返回 sessionID + projectID（用于 system token 的 ACPSessionID + scope.project_ids）。
+func mustInsertACPSetup(t *testing.T, ctx context.Context, pool *pgxpool.Pool, ownerUID uuid.UUID, status string) (sessionID, projectID uuid.UUID) {
+	t.Helper()
+
+	// project
+	pid := uuid.New()
+	_, err := pool.Exec(ctx, `
+		INSERT INTO projects (id, slug, name, owner_id, visibility)
+		VALUES ($1, $2, $2, $3, 'private')`,
+		pid, "p"+pid.String()[:8], ownerUID)
+	require.NoError(t, err)
+
+	// workspace
+	wid := uuid.New()
+	_, err = pool.Exec(ctx, `
+		INSERT INTO workspaces (id, project_id, slug, name, git_remote_url, default_branch, main_path, sync_status)
+		VALUES ($1, $2, $3, $3, 'git@local:x.git', 'main', '/tmp/'||$1::text, 'idle')`,
+		wid, pid, "w"+wid.String()[:8])
+	require.NoError(t, err)
+
+	// agent_kind
+	akid := uuid.New()
+	_, err = pool.Exec(ctx, `
+		INSERT INTO acp_agent_kinds (id, name, display_name, binary_path, args, enabled, created_by)
+		VALUES ($1, $2, $2, '/bin/echo', '{}', TRUE, $3)`,
+		akid, "ak"+akid.String()[:8], ownerUID)
+	require.NoError(t, err)
+
+	// session
+	sid := uuid.New()
+	_, err = pool.Exec(ctx, `
+		INSERT INTO acp_sessions (id, workspace_id, project_id, agent_kind_id, user_id, branch, cwd_path, is_main_worktree, status)
+		VALUES ($1, $2, $3, $4, $5, 'main', '/tmp/x', TRUE, $6)`,
+		sid, wid, pid, akid, ownerUID, status)
+	require.NoError(t, err)
+
+	return sid, pid
+}
+
+func TestPgRepo_RevokeBySession(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+
+	uid := mustInsertUser(t, ctx, pool, "rs@local", false)
+	sid, pid := mustInsertACPSetup(t, ctx, pool, uid, "running")
+
+	// 插一个 system token
+	_, hash, _ := mcp.NewToken()
+	tk, err := repo.Create(ctx, &mcp.MCPToken{
+		ID: uuid.New(), TokenHash: hash, UserID: uid,
+		Name: "system-x", Issuer: mcp.IssuerSystem,
+		Scope:        mcp.Scope{ProjectIDs: []uuid.UUID{pid}},
+		ACPSessionID: &sid,
+		CreatedBy:    uid,
+	})
+	require.NoError(t, err)
+
+	// 撤销
+	revoked, err := repo.RevokeBySession(ctx, sid)
+	require.NoError(t, err)
+	require.Len(t, revoked, 1)
+	assert.Equal(t, tk.ID, revoked[0])
+
+	// row 上 revoked_at 写入
+	got, err := repo.GetByID(ctx, tk.ID)
+	require.NoError(t, err)
+	require.NotNil(t, got.RevokedAt)
+
+	// 重复调用 → 0 行（已撤销不再返）
+	again, err := repo.RevokeBySession(ctx, sid)
+	require.NoError(t, err)
+	assert.Empty(t, again)
+}
+
+func TestPgRepo_ReapStaleSystemTokens(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+
+	uid := mustInsertUser(t, ctx, pool, "reap@local", false)
+
+	// 一个 crashed session 的 system token
+	sidCrashed, pidA := mustInsertACPSetup(t, ctx, pool, uid, "crashed")
+	_, h1, _ := mcp.NewToken()
+	tk1, err := repo.Create(ctx, &mcp.MCPToken{
+		ID: uuid.New(), TokenHash: h1, UserID: uid,
+		Name: "crashed-tok", Issuer: mcp.IssuerSystem,
+		Scope:        mcp.Scope{ProjectIDs: []uuid.UUID{pidA}},
+		ACPSessionID: &sidCrashed, CreatedBy: uid,
+	})
+	require.NoError(t, err)
+
+	// 一个 running session 的 system token（不该被 reap）
+	sidRunning, pidB := mustInsertACPSetup(t, ctx, pool, uid, "running")
+	_, h2, _ := mcp.NewToken()
+	tk2, err := repo.Create(ctx, &mcp.MCPToken{
+		ID: uuid.New(), TokenHash: h2, UserID: uid,
+		Name: "running-tok", Issuer: mcp.IssuerSystem,
+		Scope:        mcp.Scope{ProjectIDs: []uuid.UUID{pidB}},
+		ACPSessionID: &sidRunning, CreatedBy: uid,
+	})
+	require.NoError(t, err)
+
+	reaped, err := repo.ReapStaleSystemTokens(ctx)
+	require.NoError(t, err)
+	require.Len(t, reaped, 1)
+	assert.Equal(t, tk1.ID, reaped[0])
+
+	got1, _ := repo.GetByID(ctx, tk1.ID)
+	require.NotNil(t, got1.RevokedAt)
+	got2, _ := repo.GetByID(ctx, tk2.ID)
+	assert.Nil(t, got2.RevokedAt)
+}