diff --git a/internal/app/mcp_integration_test.go b/internal/app/mcp_integration_test.go index 5796a4b..e2c5e8a 100644 --- a/internal/app/mcp_integration_test.go +++ b/internal/app/mcp_integration_test.go @@ -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 } // ---------------------------------------------------------------------------