You've already forked agentic-coding-workflow
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:
@@ -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
|
||||||
}
|
}
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|||||||
Reference in New Issue
Block a user