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>
This commit is contained in:
2026-05-08 13:26:12 +08:00
parent e0e9ae7bda
commit f708e102d8
+203 -4
View File
@@ -39,6 +39,8 @@ import (
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"strings" "strings"
"sync"
"sync/atomic"
"testing" "testing"
"time" "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) { 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) { 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) { 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
} }
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------