test(mcp): K4 — rate limit + ACP lifecycle + startup reaper

- RateLimit_PerToken: 110 burst calls -> at least 1 rate-limit error - ACPSession_TokenLifecycle: create session -> token appears; terminate -> revoked - StartupReaper: crashed session + live token -> ReapStaleSystemTokens revokes Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-05-08 13:26:12 +08:00
parent e0e9ae7bda
commit f708e102d8
1 changed files with 203 additions and 4 deletions
@@ -39,6 +39,8 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"

@@ -693,19 +695,216 @@ func (s *mcpIntegrationSuite) testWriteToolAudit(t *testing.T) {
 }

 // ---------------------------------------------------------------------------
-// K4 stubs
+// K4: Rate limit + ACP lifecycle + startup reaper (3 scenarios)
 // ---------------------------------------------------------------------------

+// testRateLimitPerToken issues a fresh token (isolated bucket from suite tokens),
+// then fires 110 concurrent requests against a rate-limited endpoint (default 100/min).
+// At least 1 request must be rejected with 429.
 func (s *mcpIntegrationSuite) testRateLimitPerToken(t *testing.T) {
-	t.Skip("filled in K4")
+	plain := s.issueRateLimitTestToken(t)
+
+	// Build an isolated HTTP server with its own rate limiter so the bucket is fresh.
+	rl := mcp.NewRateLimiter(mcp.RateLimitConfig{
+		PerTokenPerMinute: 100,
+		GlobalPerMinute:   10000,
+	}, nil)
+
+	mcpRepo := mcp.NewPostgresRepository(s.pool)
+	ts := mcp.NewTokenService(mcpRepo, s.auditRec, nil)
+
+	r := chi.NewRouter()
+	r.Use(middleware.RequestID)
+	r.Use(mcp.NewAuthMiddleware(ts))
+	r.Use(mcp.NewRateLimitMiddleware(rl))
+	r.Post("/mcp", func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	srv := httptest.NewServer(r)
+	defer srv.Close()
+
+	const total = 110
+	var (
+		wg       sync.WaitGroup
+		rejected atomic.Int32
+	)
+
+	for i := 0; i < total; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			req, _ := http.NewRequest(http.MethodPost, srv.URL+"/mcp", nil)
+			req.Header.Set("Authorization", "Bearer "+plain)
+			resp, err := http.DefaultClient.Do(req)
+			if err != nil {
+				return
+			}
+			defer resp.Body.Close()
+			if resp.StatusCode == http.StatusTooManyRequests {
+				rejected.Add(1)
+			}
+		}()
+	}
+	wg.Wait()
+
+	assert.GreaterOrEqual(t, rejected.Load(), int32(1),
+		"at least 1 of %d requests should be rate-limited", total)
 }

+// testACPSessionTokenLifecycle verifies the full token lifecycle tied to an ACP session:
+//  1. Seed acp_sessions row (status='running')
+//  2. Create system token bound to session -> issuer='system', revoked_at IS NULL
+//  3. Revoke by session -> revoked_at gets written
+//  4. Audit log contains mcp.token.revoke_by_session
 func (s *mcpIntegrationSuite) testACPSessionTokenLifecycle(t *testing.T) {
-	t.Skip("filled in K4")
+	ctx := s.ctx
+	sessionID := uuid.New()
+
+	// 1. Seed prerequisite rows: workspace + acp_agent_kind + acp_session.
+	wsID, akID := s.seedWorkspaceAndAgentKind(t, "lifecycle")
+	err := s.seedRunningSession(ctx, t, sessionID, wsID, akID)
+	require.NoError(t, err)
+
+	// 2. Create system token bound to this session via direct SQL.
+	_, hash, _ := mcp.NewToken()
+	tokenID := uuid.New()
+	expires := time.Now().Add(time.Hour)
+	_, err = s.pool.Exec(ctx, `INSERT INTO mcp_tokens
+		(id, token_hash, user_id, name, issuer, scope, acp_session_id, expires_at, created_by)
+		VALUES ($1, $2, $3, 'acp:' || $4::text, 'system', '{}'::jsonb, $4, $5, $6)`,
+		tokenID, hash, s.adminID, sessionID, expires, s.adminID)
+	require.NoError(t, err)
+
+	// Verify token is active.
+	tok, err := mcp.NewPostgresRepository(s.pool).GetByID(ctx, tokenID)
+	require.NoError(t, err)
+	assert.Equal(t, mcp.IssuerSystem, tok.Issuer)
+	assert.Nil(t, tok.RevokedAt, "newly created token should not be revoked")
+
+	// 3. Simulate session termination - revoke tokens bound to session.
+	err = s.mcpTokenSvc.RevokeBySession(ctx, sessionID)
+	require.NoError(t, err)
+
+	// 4. Verify token revoked_at gets written.
+	tok, err = mcp.NewPostgresRepository(s.pool).GetByID(ctx, tokenID)
+	require.NoError(t, err)
+	assert.NotNil(t, tok.RevokedAt, "token should be revoked after session termination")
+
+	// 5. Verify audit_logs has mcp.token.revoke_by_session entry.
+	var auditCount int
+	err = s.pool.QueryRow(ctx,
+		`SELECT COUNT(*) FROM audit_logs WHERE action = 'mcp.token.revoke_by_session' AND target_id = $1`,
+		tokenID.String(),
+	).Scan(&auditCount)
+	require.NoError(t, err)
+	assert.Equal(t, 1, auditCount, "expected 1 revoke_by_session audit entry")
 }

+// testStartupReaper seeds a crashed acp_session + active system mcp_token,
+// then calls ReapStaleSystemTokens and verifies the token gets revoked.
 func (s *mcpIntegrationSuite) testStartupReaper(t *testing.T) {
-	t.Skip("filled in K4")
+	ctx := s.ctx
+
+	// 1. Seed a crashed session + active system token.
+	sessionID, tokenID := s.seedCrashedSessionWithLiveToken(ctx, t)
+	_ = sessionID
+
+	// Verify token is not yet revoked.
+	mcpRepo := mcp.NewPostgresRepository(s.pool)
+	tok, err := mcpRepo.GetByID(ctx, tokenID)
+	require.NoError(t, err)
+	assert.Nil(t, tok.RevokedAt, "token should not be revoked before reaper")
+
+	// 2. Run reaper.
+	reaped, err := mcpRepo.ReapStaleSystemTokens(ctx)
+	require.NoError(t, err)
+	assert.Contains(t, reaped, tokenID, "reaped list should contain our token")
+
+	// 3. Verify token gets revoked.
+	tok, err = mcpRepo.GetByID(ctx, tokenID)
+	require.NoError(t, err)
+	assert.NotNil(t, tok.RevokedAt, "token should be revoked after reaper")
+}
+
+// --- K4 helpers ---
+
+// issueRateLimitTestToken creates a fresh admin token isolated from suite tokens.
+func (s *mcpIntegrationSuite) issueRateLimitTestToken(t *testing.T) string {
+	t.Helper()
+	mcpRepo := mcp.NewPostgresRepository(s.pool)
+	ts := mcp.NewTokenService(mcpRepo, s.auditRec, nil)
+	res, err := ts.IssueAdmin(s.ctx, mcp.IssueAdminInput{
+		UserID:    s.adminID,
+		Name:      "rate-limit-test",
+		Scope:     mcp.Scope{},
+		ExpiresAt: time.Now().Add(time.Hour),
+		CreatedBy: s.adminID,
+	})
+	require.NoError(t, err)
+	return res.Plaintext
+}
+
+// seedWorkspaceAndAgentKind creates a workspace and acp_agent_kind row for FK
+// references in acp_sessions. Returns (workspaceID, agentKindID).
+func (s *mcpIntegrationSuite) seedWorkspaceAndAgentKind(t *testing.T, suffix string) (uuid.UUID, uuid.UUID) {
+	t.Helper()
+	ctx := s.ctx
+
+	wsID := uuid.New()
+	cwd := t.TempDir()
+	_, err := s.pool.Exec(ctx, `INSERT INTO workspaces
+		(id, project_id, slug, name, git_remote_url, main_path, default_branch)
+		VALUES ($1, $2, $3, 'w-k4', 'git@x:y.git', $4, 'main')`,
+		wsID, s.projectID, "ws-k4-"+suffix, cwd)
+	require.NoError(t, err)
+
+	akID := uuid.New()
+	_, err = s.pool.Exec(ctx, `INSERT INTO acp_agent_kinds
+		(id, name, display_name, binary_path, created_by)
+		VALUES ($1, $2, 'K4 Agent', '/bin/echo', $3)`,
+		akID, "k4-agent-"+suffix, s.adminID)
+	require.NoError(t, err)
+
+	return wsID, akID
+}
+
+// seedRunningSession inserts an acp_sessions row with status='running'.
+func (s *mcpIntegrationSuite) seedRunningSession(ctx context.Context, t *testing.T, sessionID, wsID, akID uuid.UUID) error {
+	t.Helper()
+	_, err := s.pool.Exec(ctx, `INSERT INTO acp_sessions
+		(id, workspace_id, project_id, agent_kind_id, user_id,
+		 branch, cwd_path, is_main_worktree, status, started_at)
+		VALUES ($1, $2, $3, $4, $5, 'feat/lifecycle-test', $6, false, 'running', now())`,
+		sessionID, wsID, s.projectID, akID, s.adminID, t.TempDir())
+	return err
+}
+
+// seedCrashedSessionWithLiveToken inserts a crashed acp_session + an active
+// system mcp_token bound to it. Returns (sessionID, tokenID).
+func (s *mcpIntegrationSuite) seedCrashedSessionWithLiveToken(ctx context.Context, t *testing.T) (uuid.UUID, uuid.UUID) {
+	t.Helper()
+
+	wsID, akID := s.seedWorkspaceAndAgentKind(t, "reaper")
+
+	sessionID := uuid.New()
+	_, err := s.pool.Exec(ctx, `INSERT INTO acp_sessions
+		(id, workspace_id, project_id, agent_kind_id, user_id,
+		 branch, cwd_path, is_main_worktree, status, started_at)
+		VALUES ($1, $2, $3, $4, $5, 'feat/reap-test', $6, false, 'crashed', now())`,
+		sessionID, wsID, s.projectID, akID, s.adminID, t.TempDir())
+	require.NoError(t, err)
+
+	_, hash, _ := mcp.NewToken()
+	tokenID := uuid.New()
+	expires := time.Now().Add(time.Hour)
+	_, err = s.pool.Exec(ctx, `INSERT INTO mcp_tokens
+		(id, token_hash, user_id, name, issuer, scope, acp_session_id, expires_at, created_by)
+		VALUES ($1, $2, $3, 'acp:' || $4::text, 'system', '{}'::jsonb, $4, $5, $6)`,
+		tokenID, hash, s.adminID, sessionID, expires, s.adminID)
+	require.NoError(t, err)
+
+	return sessionID, tokenID
 }

 // ---------------------------------------------------------------------------