diff --git a/backend/src/assistant/application/agentAnalyzeSessionService.ts b/backend/src/assistant/application/agentAnalyzeSessionService.ts index 7df3b0fd..ca05383b 100644 --- a/backend/src/assistant/application/agentAnalyzeSessionService.ts +++ b/backend/src/assistant/application/agentAnalyzeSessionService.ts @@ -55,7 +55,7 @@ export interface AnalyzeSessionRunContext { query: string; startedAt: number; completedAt?: number; - status: 'pending' | 'running' | 'completed' | 'failed'; + status: 'pending' | 'running' | 'completed' | 'failed' | 'quota_exceeded'; error?: string; } diff --git a/backend/src/assistant/application/assistantApplicationService.ts b/backend/src/assistant/application/assistantApplicationService.ts index e6ddbe8e..399f8ac6 100644 --- a/backend/src/assistant/application/assistantApplicationService.ts +++ b/backend/src/assistant/application/assistantApplicationService.ts @@ -9,7 +9,8 @@ export type AssistantSessionStatus = | 'running' | 'awaiting_user' | 'completed' - | 'failed'; + | 'failed' + | 'quota_exceeded'; export interface ManagedAssistantSession { sessionId: string; diff --git a/backend/src/routes/__tests__/enterpriseReportRoutes.test.ts b/backend/src/routes/__tests__/enterpriseReportRoutes.test.ts index da203401..08b2202f 100644 --- a/backend/src/routes/__tests__/enterpriseReportRoutes.test.ts +++ b/backend/src/routes/__tests__/enterpriseReportRoutes.test.ts @@ -32,6 +32,7 @@ interface ReportArtifactRow { content_hash: string; visibility: string; created_by: string | null; + expires_at: number | null; } let tmpDir: string; @@ -85,6 +86,31 @@ function readAuditActions(): string[] { } } +function writeWorkspacePolicies(input: { + retentionPolicy?: Record; +}): void { + const db = openEnterpriseDb(dbPath); + const now = Date.now(); + try { + db.prepare(` + INSERT OR IGNORE INTO organizations (id, name, status, plan, created_at, updated_at) + VALUES ('tenant-a', 'tenant-a', 'active', 'enterprise', ?, ?) + `).run(now, now); + db.prepare(` + INSERT OR REPLACE INTO workspaces + (id, tenant_id, name, retention_policy, quota_policy, created_at, updated_at) + VALUES + ('workspace-a', 'tenant-a', 'workspace-a', ?, NULL, ?, ?) + `).run( + input.retentionPolicy ? JSON.stringify(input.retentionPolicy) : null, + now, + now, + ); + } finally { + db.close(); + } +} + beforeEach(async () => { tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'smartperfetto-enterprise-report-routes-')); dbPath = path.join(tmpDir, 'enterprise.sqlite'); @@ -198,4 +224,30 @@ describe('enterprise report routes', () => { await expect(fs.access(path.dirname(row!.local_path))).rejects.toThrow(); expect(readAuditActions()).toContain('report.deleted'); }); + + it('applies report retention policy and hides expired cached reports', async () => { + const app = makeApp(); + const reportId = 'report-expired'; + writeWorkspacePolicies({ + retentionPolicy: { + reportRetentionDays: 0, + }, + }); + + persistReport(reportId, { + html: 'expired report', + generatedAt: Date.now() - 1, + sessionId: 'session-expired', + runId: 'run-expired', + traceId: 'trace-expired', + tenantId: 'tenant-a', + workspaceId: 'workspace-a', + userId: 'user-a', + visibility: 'private', + }); + + expect(readReportArtifact(reportId)?.expires_at).toBeLessThanOrEqual(Date.now()); + const getRes = await ssoHeaders(request(app).get(`/api/reports/${reportId}`)); + expect(getRes.status).toBe(404); + }); }); diff --git a/backend/src/routes/__tests__/enterpriseTraceMetadataRoutes.test.ts b/backend/src/routes/__tests__/enterpriseTraceMetadataRoutes.test.ts index fa0f39e7..42de6f1f 100644 --- a/backend/src/routes/__tests__/enterpriseTraceMetadataRoutes.test.ts +++ b/backend/src/routes/__tests__/enterpriseTraceMetadataRoutes.test.ts @@ -35,6 +35,7 @@ interface TraceAssetRow { status: string; size_bytes: number; metadata_json: string; + expires_at: number | null; } let tmpDir: string; @@ -172,6 +173,33 @@ function readCount(table: 'trace_assets' | 'trace_processor_leases'): number { } } +function writeWorkspacePolicies(input: { + quotaPolicy?: Record; + retentionPolicy?: Record; +}): void { + const db = openEnterpriseDb(dbPath); + const now = Date.now(); + try { + db.prepare(` + INSERT OR IGNORE INTO organizations (id, name, status, plan, created_at, updated_at) + VALUES ('tenant-a', 'tenant-a', 'active', 'enterprise', ?, ?) + `).run(now, now); + db.prepare(` + INSERT OR REPLACE INTO workspaces + (id, tenant_id, name, retention_policy, quota_policy, created_at, updated_at) + VALUES + ('workspace-a', 'tenant-a', 'workspace-a', ?, ?, ?, ?) + `).run( + input.retentionPolicy ? JSON.stringify(input.retentionPolicy) : null, + input.quotaPolicy ? JSON.stringify(input.quotaPolicy) : null, + now, + now, + ); + } finally { + db.close(); + } +} + beforeEach(async () => { tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'smartperfetto-enterprise-trace-routes-')); dbPath = path.join(tmpDir, 'enterprise.sqlite'); @@ -301,6 +329,51 @@ describe('enterprise trace metadata routes', () => { expect(otherWorkspaceRes.status).toBe(404); }); + it('rejects uploads that exceed workspace trace quota before metadata is committed', async () => { + const app = makeApp(); + writeWorkspacePolicies({ + quotaPolicy: { + maxTraceBytes: 4, + }, + }); + + const res = await ssoHeaders( + request(app) + .post('/api/traces/upload') + .attach('file', Buffer.from('12345'), 'too-large.trace'), + ); + + expect(res.status).toBe(413); + expect(res.body).toEqual(expect.objectContaining({ + success: false, + code: 'TRACE_SIZE_QUOTA_EXCEEDED', + status: 'quota_exceeded', + })); + expect(readCount('trace_assets')).toBe(0); + expect(fakeTraceProcessorService.initializeUploadWithId).not.toHaveBeenCalled(); + }); + + it('applies workspace trace retention policy to uploaded trace metadata', async () => { + const app = makeApp(); + writeWorkspacePolicies({ + retentionPolicy: { + traceRetentionDays: 3, + }, + }); + const beforeUpload = Date.now(); + + const res = await ssoHeaders( + request(app) + .post('/api/traces/upload') + .attach('file', Buffer.from('trace-with-retention'), 'retained.trace'), + ); + + expect(res.status).toBe(200); + const row = readTraceAsset(res.body.trace.id); + expect(row?.expires_at).toBeGreaterThanOrEqual(beforeUpload + 3 * 24 * 60 * 60 * 1000); + expect(row?.expires_at).toBeLessThanOrEqual(Date.now() + 3 * 24 * 60 * 60 * 1000); + }); + it('records observed processor RSS on the frontend lease and exposes RAM budget stats', async () => { const app = makeApp(); const sourceTracePath = path.join(tmpDir, 'rss.trace'); diff --git a/backend/src/routes/agentRoutes.ts b/backend/src/routes/agentRoutes.ts index 637c47f7..4067d797 100644 --- a/backend/src/routes/agentRoutes.ts +++ b/backend/src/routes/agentRoutes.ts @@ -67,6 +67,10 @@ import { buildTraceProcessorLeaseModeDecision, type TraceProcessorLeaseModeDecision, } from '../services/traceProcessorLeaseModeDecision'; +import { + evaluateAnalysisRunQuota, + type EnterpriseQuotaDecision, +} from '../services/enterpriseQuotaPolicyService'; import { estimateTraceProcessorRssBytes } from '../services/traceProcessorRamBudget'; import { TraceProcessorFactory } from '../services/workingTraceProcessor'; import { registerAgentLogsRoutes } from './agentLogsRoutes'; @@ -167,6 +171,26 @@ function buildRunId(sessionId: string, sequence: number): string { return `run-${sessionId}-${sequence}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; } +function sendAgentQuotaDenied( + res: express.Response, + decision: EnterpriseQuotaDecision, +): express.Response { + return res.status(decision.httpStatus).json({ + success: false, + code: decision.code, + status: decision.status, + error: decision.message, + details: decision.details, + }); +} + +function terminalRunStatusForResult( + result: AgentRuntimeAnalysisResult, +): Extract { + if (result.terminationReason === 'max_budget_usd') return 'quota_exceeded'; + return result.success ? 'completed' : 'failed'; +} + function enterpriseLeasesEnabled(): boolean { return resolveFeatureConfig().enterprise; } @@ -302,7 +326,7 @@ function markSessionRunStatus( ): void { if (!session.activeRun) return; session.activeRun.status = status; - if (status === 'completed' || status === 'failed') { + if (status === 'completed' || status === 'failed' || status === 'quota_exceeded') { session.activeRun.completedAt = Date.now(); } session.activeRun.error = error; @@ -331,7 +355,7 @@ interface AnalysisSession { sessionId: string; sseClients: express.Response[]; result?: AgentRuntimeAnalysisResult; - status: 'pending' | 'running' | 'awaiting_user' | 'completed' | 'failed'; + status: 'pending' | 'running' | 'awaiting_user' | 'completed' | 'failed' | 'quota_exceeded'; error?: string; traceId: string; tenantId?: string; @@ -1108,6 +1132,12 @@ async function handleAnalyzeRequest( console.log(`[AgentRoutes] Comparison mode: current=${traceId}, reference=${referenceTraceId}`); } + const quotaDecision = evaluateAnalysisRunQuota(requestContext); + if (!quotaDecision.allowed) { + sendAgentQuotaDenied(res, quotaDecision); + return; + } + // Initialize tools ensureToolsRegistered(); @@ -1393,7 +1423,7 @@ function handleSessionStream(req: express.Request, res: express.Response, sessio // If analysis is already completed, send the result. // Resumed sessions may not have session.result in memory; recover from persisted turn context. - if (session.status === 'completed') { + if (session.status === 'completed' || session.status === 'quota_exceeded') { recoverResultForSessionIfNeeded(sessionId, session); if (session.result) { sendAgentDrivenResult(res, session); @@ -1489,7 +1519,7 @@ router.get('/:sessionId/status', (req, res) => { observability: buildSessionObservability(session), }; - if (session.status === 'completed') { + if (session.status === 'completed' || session.status === 'quota_exceeded') { const recoveredResult = recoverResultForSessionIfNeeded(sessionId, session); if (recoveredResult) { const conclusion = normalizeNarrativeForClient(recoveredResult.conclusion); @@ -2824,8 +2854,11 @@ async function runAgentDrivenAnalysis( if (idx >= 0) session.hypotheses[idx] = h; } } - session.status = result.success ? 'completed' : 'failed'; - markSessionRunStatus(session, result.success ? 'completed' : 'failed'); + const terminalRunStatus = terminalRunStatusForResult(result); + session.status = terminalRunStatus === 'quota_exceeded' + ? 'quota_exceeded' + : result.success ? 'completed' : 'failed'; + markSessionRunStatus(session, terminalRunStatus); // Record conclusion in cross-turn history if (!session.conclusionHistory) session.conclusionHistory = []; diff --git a/backend/src/routes/reportRoutes.ts b/backend/src/routes/reportRoutes.ts index 8d4b7c5d..e4689367 100644 --- a/backend/src/routes/reportRoutes.ts +++ b/backend/src/routes/reportRoutes.ts @@ -25,6 +25,7 @@ import { import { REPORT_CAUSAL_MAP_CSS, REPORT_CAUSAL_MAP_SCRIPT } from '../services/reportCausalMapAssets'; import { localize, parseOutputLanguage } from '../agentv3/outputLanguage'; import { resolveEnterpriseDataRoot } from '../services/traceMetadataStore'; +import { resolveEnterpriseRetentionExpiresAt } from '../services/enterpriseQuotaPolicyService'; import { sendResourceNotFound, type ResourceOwnerFields, @@ -55,6 +56,7 @@ type PersistedReport = ResourceOwnerFields & { runId?: string; traceId?: string; visibility?: string; + expiresAt?: number | null; }; export const reportStore = new Map(); @@ -147,6 +149,10 @@ function fallbackRunId(entry: PersistedReport): string { return entry.runId || `run-${entry.sessionId}-report`; } +function isReportExpired(entry: PersistedReport, now = Date.now()): boolean { + return typeof entry.expiresAt === 'number' && entry.expiresAt <= now; +} + function ensureEnterpriseReportGraph( db: Database.Database, reportId: string, @@ -247,6 +253,17 @@ function persistEnterpriseReport(reportId: string, entry: PersistedReport): void const createdAt = entry.generatedAt || Date.now(); const visibility = entry.visibility || 'private'; const contentHash = reportContentHash(entry.html); + const expiresAt = resolveEnterpriseRetentionExpiresAt( + db, + { + tenantId: entry.tenantId!, + workspaceId: entry.workspaceId!, + ...(entry.userId ? { userId: entry.userId } : {}), + }, + 'report', + createdAt, + ); + entry.expiresAt = expiresAt; db.prepare(` INSERT INTO report_artifacts (id, tenant_id, workspace_id, session_id, run_id, local_path, content_hash, visibility, created_by, created_at, expires_at) @@ -273,7 +290,7 @@ function persistEnterpriseReport(reportId: string, entry: PersistedReport): void visibility, entry.userId ?? null, createdAt, - null, + expiresAt, ); fs.writeFileSync(metadataPath, JSON.stringify({ @@ -287,6 +304,7 @@ function persistEnterpriseReport(reportId: string, entry: PersistedReport): void userId: entry.userId, visibility, contentHash, + expiresAt, }, null, 2)); }); } @@ -304,6 +322,7 @@ function persistLegacyReport(reportId: string, entry: PersistedReport): void { workspaceId: entry.workspaceId, userId: entry.userId, visibility: entry.visibility, + expiresAt: entry.expiresAt, })); } @@ -315,7 +334,8 @@ function loadEnterpriseReport(reportId: string): PersistedReport | null { SELECT * FROM report_artifacts WHERE id = ? - `).get(reportId); + AND (expires_at IS NULL OR expires_at > ?) + `).get(reportId, Date.now()); if (!row || !fs.existsSync(row.local_path)) return null; const html = fs.readFileSync(row.local_path, 'utf-8'); const entry: PersistedReport = { @@ -327,6 +347,7 @@ function loadEnterpriseReport(reportId: string): PersistedReport | null { workspaceId: row.workspace_id, ...(row.created_by ? { userId: row.created_by } : {}), visibility: row.visibility, + expiresAt: row.expires_at, }; reportStore.set(reportId, entry); return entry; @@ -395,6 +416,7 @@ function loadLegacyReportFromDisk(reportId: string): PersistedReport | null { let runId: string | undefined; let traceId: string | undefined; let visibility: string | undefined; + let expiresAt: number | undefined; let owner: ResourceOwnerFields = {}; if (fs.existsSync(metaPath)) { const meta = JSON.parse(fs.readFileSync(metaPath, 'utf-8')); @@ -403,12 +425,16 @@ function loadLegacyReportFromDisk(reportId: string): PersistedReport | null { runId = meta.runId; traceId = meta.traceId; visibility = meta.visibility; + expiresAt = typeof meta.expiresAt === 'number' ? meta.expiresAt : undefined; owner = { tenantId: meta.tenantId, workspaceId: meta.workspaceId, userId: meta.userId, ownerUserId: meta.ownerUserId, }; + if (typeof expiresAt === 'number' && expiresAt <= Date.now()) { + return null; + } } const entry = { @@ -418,6 +444,7 @@ function loadLegacyReportFromDisk(reportId: string): PersistedReport | null { ...(runId ? { runId } : {}), ...(traceId ? { traceId } : {}), ...(visibility ? { visibility } : {}), + ...(typeof expiresAt === 'number' ? { expiresAt } : {}), ...owner, }; // Cache in memory for subsequent access @@ -478,6 +505,10 @@ function deletePersistedReport(reportId: string): boolean { function getReportForContext(reportId: string, req: express.Request): PersistedReport | null { const context = requireRequestContext(req); const report = reportStore.get(reportId) || loadReportFromDisk(reportId); + if (report && isReportExpired(report)) { + reportStore.delete(reportId); + return null; + } if (!report || !canReadReportResource(report, context)) { return null; } diff --git a/backend/src/routes/simpleTraceRoutes.ts b/backend/src/routes/simpleTraceRoutes.ts index f0d7aa1c..f7d8aadb 100644 --- a/backend/src/routes/simpleTraceRoutes.ts +++ b/backend/src/routes/simpleTraceRoutes.ts @@ -33,6 +33,10 @@ import { sharedQueueLengthForTrace, type TraceProcessorLeaseModeDecision, } from '../services/traceProcessorLeaseModeDecision'; +import { + evaluateTraceUploadQuota, + type EnterpriseQuotaDecision, +} from '../services/enterpriseQuotaPolicyService'; import { buildTraceOwnerMetadata, deleteTraceMetadata, @@ -235,6 +239,19 @@ function recordTraceAudit( }); } +function sendTraceQuotaDenied( + res: Response, + decision: EnterpriseQuotaDecision, +): Response { + return res.status(decision.httpStatus).json({ + success: false, + code: decision.code, + status: decision.status, + error: decision.message, + details: decision.details, + }); +} + function summarizeLeaseBlockers(leases: TraceProcessorLeaseRecord[]): Array<{ id: string; traceId: string; @@ -673,6 +690,11 @@ router.post( } const file = req.file; + const quotaDecision = evaluateTraceUploadQuota(context, file.size); + if (!quotaDecision.allowed) { + await cleanupFile(file.path); + return sendTraceQuotaDenied(res, quotaDecision); + } // Generate trace ID upfront for consistency const traceId = uuidv4(); @@ -760,12 +782,19 @@ router.post('/upload-url', async (req, res) => { const contentLength = response.headers.get('content-length'); const uploadLimitBytes = resolveTraceUploadLimitBytes(); - if (contentLength && Number.parseInt(contentLength, 10) > uploadLimitBytes) { + const contentLengthBytes = contentLength ? Number.parseInt(contentLength, 10) : Number.NaN; + if (Number.isFinite(contentLengthBytes) && contentLengthBytes > uploadLimitBytes) { return res.status(413).json({ error: 'Trace file too large', details: `Remote trace exceeds ${uploadLimitBytes} bytes` }); } + if (Number.isFinite(contentLengthBytes)) { + const quotaDecision = evaluateTraceUploadQuota(context, contentLengthBytes); + if (!quotaDecision.allowed) { + return sendTraceQuotaDenied(res, quotaDecision); + } + } const tracesDir = getWritableTraceDirForContext(context); await fs.mkdir(tracesDir, { recursive: true }); @@ -776,6 +805,11 @@ router.post('/upload-url', async (req, res) => { let size = 0; try { size = await streamResponseBodyToTempFile(response, tempPath); + const quotaDecision = evaluateTraceUploadQuota(context, size); + if (!quotaDecision.allowed) { + await cleanupFile(tempPath); + return sendTraceQuotaDenied(res, quotaDecision); + } await renameTraceAtomically(tempPath, finalPath); } catch (streamError) { await cleanupFile(tempPath); diff --git a/backend/src/services/__tests__/analysisRunStore.test.ts b/backend/src/services/__tests__/analysisRunStore.test.ts index 622bdf73..7cfd3e39 100644 --- a/backend/src/services/__tests__/analysisRunStore.test.ts +++ b/backend/src/services/__tests__/analysisRunStore.test.ts @@ -103,6 +103,33 @@ describe('analysis run store', () => { expect(isAnalysisRunHeartbeatFresh(runScope, 'run-failed', 1_777_000_011_000, 60_000)).toBe(false); }); + it('persists quota_exceeded as a terminal run and session state', () => { + const runScope = scope({ runId: 'run-quota', sessionId: 'session-quota' }); + + persistAnalysisRunState(runScope, 'running', { now: 1_777_000_000_000 }); + persistAnalysisRunState(runScope, 'quota_exceeded', { + now: 1_777_000_010_000, + error: 'single-run LLM budget exhausted', + }); + + expect(getAnalysisRunLifecycle(runScope, 'run-quota')).toEqual(expect.objectContaining({ + status: 'quota_exceeded', + completedAt: 1_777_000_010_000, + heartbeatAt: 1_777_000_010_000, + errorJson: JSON.stringify({ message: 'single-run LLM budget exhausted' }), + })); + expect(isAnalysisRunHeartbeatFresh(runScope, 'run-quota', 1_777_000_011_000, 60_000)).toBe(false); + + const db = openEnterpriseDb(); + try { + expect(db.prepare('SELECT status FROM analysis_sessions WHERE id = ?').get('session-quota')).toEqual({ + status: 'quota_exceeded', + }); + } finally { + db.close(); + } + }); + it('fails interrupted nonterminal runs on backend startup while preserving terminal runs', () => { const pendingScope = scope({ sessionId: 'session-pending', runId: 'run-pending' }); const runningScope = scope({ sessionId: 'session-running', runId: 'run-running' }); diff --git a/backend/src/services/__tests__/enterpriseQuotaPolicyService.test.ts b/backend/src/services/__tests__/enterpriseQuotaPolicyService.test.ts new file mode 100644 index 00000000..03b606da --- /dev/null +++ b/backend/src/services/__tests__/enterpriseQuotaPolicyService.test.ts @@ -0,0 +1,231 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later +// Copyright (C) 2024-2026 Gracker (Chris) +// This file is part of SmartPerfetto. See LICENSE for details. + +import { afterEach, beforeEach, describe, expect, it } from '@jest/globals'; +import fs from 'fs/promises'; +import os from 'os'; +import path from 'path'; +import { ENTERPRISE_FEATURE_FLAG_ENV } from '../../config'; +import type { RequestContext } from '../../middleware/auth'; +import { + ENTERPRISE_DB_PATH_ENV, + openEnterpriseDb, +} from '../enterpriseDb'; +import { + evaluateAnalysisRunQuota, + evaluateTraceUploadQuota, + readWorkspaceEnterprisePolicies, + resolveEnterpriseRetentionExpiresAt, +} from '../enterpriseQuotaPolicyService'; + +const originalEnv = { + enterprise: process.env[ENTERPRISE_FEATURE_FLAG_ENV], + enterpriseDbPath: process.env[ENTERPRISE_DB_PATH_ENV], +}; + +let tmpDir: string; + +function restoreEnvValue(key: string, value: string | undefined): void { + if (value === undefined) { + delete process.env[key]; + } else { + process.env[key] = value; + } +} + +function context(): RequestContext { + return { + tenantId: 'tenant-a', + workspaceId: 'workspace-a', + userId: 'user-a', + roles: ['workspace_admin'], + scopes: ['*'], + authType: 'sso', + requestId: 'req-test', + }; +} + +function seedWorkspacePolicies(input: { + quotaPolicy?: Record; + retentionPolicy?: Record; +}): void { + const db = openEnterpriseDb(); + const now = 1_777_000_000_000; + try { + db.prepare(` + INSERT OR IGNORE INTO organizations (id, name, status, plan, created_at, updated_at) + VALUES ('tenant-a', 'tenant-a', 'active', 'enterprise', ?, ?) + `).run(now, now); + db.prepare(` + INSERT OR REPLACE INTO workspaces + (id, tenant_id, name, retention_policy, quota_policy, created_at, updated_at) + VALUES + ('workspace-a', 'tenant-a', 'workspace-a', ?, ?, ?, ?) + `).run( + input.retentionPolicy ? JSON.stringify(input.retentionPolicy) : null, + input.quotaPolicy ? JSON.stringify(input.quotaPolicy) : null, + now, + now, + ); + db.prepare(` + INSERT OR IGNORE INTO users (id, tenant_id, email, display_name, idp_subject, created_at, updated_at) + VALUES ('user-a', 'tenant-a', 'user-a@example.test', 'user-a', 'user-a', ?, ?) + `).run(now, now); + } finally { + db.close(); + } +} + +function seedAnalysisGraph( + db: ReturnType, + sessionId: string, + traceId: string, + now: number, +): void { + db.prepare(` + INSERT OR IGNORE INTO trace_assets + (id, tenant_id, workspace_id, owner_user_id, local_path, size_bytes, status, metadata_json, created_at) + VALUES + (?, 'tenant-a', 'workspace-a', 'user-a', ?, 0, 'metadata_only', '{}', ?) + `).run(traceId, `/tmp/${traceId}.trace`, now); + db.prepare(` + INSERT OR IGNORE INTO analysis_sessions + (id, tenant_id, workspace_id, trace_id, created_by, title, visibility, status, created_at, updated_at) + VALUES + (?, 'tenant-a', 'workspace-a', ?, 'user-a', ?, 'private', 'running', ?, ?) + `).run(sessionId, traceId, sessionId, now, now); +} + +beforeEach(async () => { + tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'smartperfetto-enterprise-quota-')); + process.env[ENTERPRISE_FEATURE_FLAG_ENV] = 'true'; + process.env[ENTERPRISE_DB_PATH_ENV] = path.join(tmpDir, 'enterprise.sqlite'); +}); + +afterEach(async () => { + restoreEnvValue(ENTERPRISE_FEATURE_FLAG_ENV, originalEnv.enterprise); + restoreEnvValue(ENTERPRISE_DB_PATH_ENV, originalEnv.enterpriseDbPath); + await fs.rm(tmpDir, { recursive: true, force: true }); +}); + +describe('enterprise quota and retention policy service', () => { + it('reads workspace policy JSON and resolves artifact retention expiry', () => { + seedWorkspacePolicies({ + quotaPolicy: { maxTraceBytes: 1024 }, + retentionPolicy: { traceRetentionDays: 7, reportRetentionDays: 30 }, + }); + + const db = openEnterpriseDb(); + try { + const policies = readWorkspaceEnterprisePolicies(db, context()); + expect(policies.quotaPolicy.maxTraceBytes).toBe(1024); + expect(resolveEnterpriseRetentionExpiresAt(db, context(), 'trace', 1000)).toBe(1000 + 7 * 24 * 60 * 60 * 1000); + expect(resolveEnterpriseRetentionExpiresAt(db, context(), 'report', 1000)).toBe(1000 + 30 * 24 * 60 * 60 * 1000); + } finally { + db.close(); + } + }); + + it('rejects trace uploads that exceed per-file or workspace storage quota', () => { + seedWorkspacePolicies({ + quotaPolicy: { + maxTraceBytes: 100, + maxWorkspaceTraceBytes: 150, + }, + }); + + let db = openEnterpriseDb(); + try { + db.prepare(` + INSERT INTO trace_assets + (id, tenant_id, workspace_id, owner_user_id, local_path, size_bytes, status, metadata_json, created_at) + VALUES + ('existing', 'tenant-a', 'workspace-a', 'user-a', '/tmp/existing.trace', 80, 'ready', '{}', ?) + `).run(1_777_000_000_000); + } finally { + db.close(); + } + + expect(evaluateTraceUploadQuota(context(), 101)).toEqual(expect.objectContaining({ + allowed: false, + code: 'TRACE_SIZE_QUOTA_EXCEEDED', + status: 'quota_exceeded', + })); + expect(evaluateTraceUploadQuota(context(), 90)).toEqual(expect.objectContaining({ + allowed: false, + code: 'WORKSPACE_TRACE_STORAGE_QUOTA_EXCEEDED', + status: 'quota_exceeded', + })); + expect(evaluateTraceUploadQuota(context(), 70)).toEqual(expect.objectContaining({ + allowed: true, + code: 'OK', + })); + + db = openEnterpriseDb(); + try { + db.prepare(` + UPDATE trace_assets + SET expires_at = ? + WHERE id = 'existing' + `).run(Date.now() - 1); + } finally { + db.close(); + } + expect(evaluateTraceUploadQuota(context(), 90)).toEqual(expect.objectContaining({ + allowed: true, + code: 'OK', + })); + }); + + it('separates concurrent-run pending from monthly quota_exceeded preflight', () => { + seedWorkspacePolicies({ + quotaPolicy: { + maxConcurrentRuns: 1, + monthlyRunLimit: 2, + }, + }); + const now = Date.UTC(2026, 4, 8); + + const db = openEnterpriseDb(); + try { + seedAnalysisGraph(db, 'session-active', 'trace-active', now); + db.prepare(` + INSERT INTO analysis_runs + (id, tenant_id, workspace_id, session_id, mode, status, question, started_at) + VALUES + ('run-active', 'tenant-a', 'workspace-a', 'session-active', 'agent', 'running', '', ?) + `).run(now); + } finally { + db.close(); + } + + expect(evaluateAnalysisRunQuota(context(), { now })).toEqual(expect.objectContaining({ + allowed: false, + code: 'CONCURRENT_RUN_QUOTA_EXCEEDED', + status: 'pending', + })); + + const db2 = openEnterpriseDb(); + try { + db2.prepare(`DELETE FROM analysis_runs`).run(); + for (const id of ['run-a', 'run-b']) { + seedAnalysisGraph(db2, `session-${id}`, `trace-${id}`, now); + db2.prepare(` + INSERT INTO analysis_runs + (id, tenant_id, workspace_id, session_id, mode, status, question, started_at) + VALUES + (?, 'tenant-a', 'workspace-a', ?, 'agent', 'completed', '', ?) + `).run(id, `session-${id}`, now); + } + } finally { + db2.close(); + } + + expect(evaluateAnalysisRunQuota(context(), { now })).toEqual(expect.objectContaining({ + allowed: false, + code: 'MONTHLY_RUN_QUOTA_EXCEEDED', + status: 'quota_exceeded', + })); + }); +}); diff --git a/backend/src/services/analysisRunStore.ts b/backend/src/services/analysisRunStore.ts index 5fd29c85..24f50001 100644 --- a/backend/src/services/analysisRunStore.ts +++ b/backend/src/services/analysisRunStore.ts @@ -12,7 +12,8 @@ export type PersistedAnalysisRunStatus = | 'awaiting_user' | 'completed' | 'failed' - | 'cancelled'; + | 'cancelled' + | 'quota_exceeded'; export interface AnalysisRunPersistenceScope extends EnterpriseRepositoryScope { sessionId: string; @@ -78,7 +79,7 @@ export function resetAnalysisRunStoreForTests(): void { } function isTerminalStatus(status: string): boolean { - return status === 'completed' || status === 'failed' || status === 'cancelled'; + return status === 'completed' || status === 'failed' || status === 'cancelled' || status === 'quota_exceeded'; } function ensureAnalysisRunGraph( diff --git a/backend/src/services/enterpriseQuotaPolicyService.ts b/backend/src/services/enterpriseQuotaPolicyService.ts new file mode 100644 index 00000000..d9c679b4 --- /dev/null +++ b/backend/src/services/enterpriseQuotaPolicyService.ts @@ -0,0 +1,304 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later +// Copyright (C) 2024-2026 Gracker (Chris) +// This file is part of SmartPerfetto. See LICENSE for details. + +import type Database from 'better-sqlite3'; +import { resolveFeatureConfig } from '../config'; +import type { RequestContext } from '../middleware/auth'; +import { openEnterpriseDb } from './enterpriseDb'; +import type { EnterpriseRepositoryScope } from './enterpriseRepository'; + +export type EnterpriseQuotaDecisionCode = + | 'OK' + | 'TRACE_SIZE_QUOTA_EXCEEDED' + | 'WORKSPACE_TRACE_STORAGE_QUOTA_EXCEEDED' + | 'CONCURRENT_RUN_QUOTA_EXCEEDED' + | 'MONTHLY_RUN_QUOTA_EXCEEDED'; + +export type EnterpriseQuotaDecisionStatus = + | 'allowed' + | 'pending' + | 'quota_exceeded'; + +export interface WorkspaceQuotaPolicy { + maxTraceBytes?: number; + maxWorkspaceTraceBytes?: number; + maxConcurrentRuns?: number; + monthlyRunLimit?: number; +} + +export interface WorkspaceRetentionPolicy { + defaultRetentionDays?: number; + traceRetentionDays?: number; + reportRetentionDays?: number; +} + +export interface WorkspaceEnterprisePolicies { + quotaPolicy: WorkspaceQuotaPolicy; + retentionPolicy: WorkspaceRetentionPolicy; +} + +export interface EnterpriseQuotaDecision { + allowed: boolean; + code: EnterpriseQuotaDecisionCode; + status: EnterpriseQuotaDecisionStatus; + httpStatus: number; + message: string; + details: Record; +} + +interface WorkspacePolicyRow { + quota_policy: string | null; + retention_policy: string | null; +} + +interface CountRow { + count: number; +} + +interface SumRow { + total: number | null; +} + +const DAY_MS = 24 * 60 * 60 * 1000; +const ACTIVE_RUN_STATUSES = ['pending', 'running', 'awaiting_user'] as const; + +function allowedDecision(): EnterpriseQuotaDecision { + return { + allowed: true, + code: 'OK', + status: 'allowed', + httpStatus: 200, + message: 'Allowed', + details: {}, + }; +} + +function denyDecision( + code: Exclude, + status: Exclude, + httpStatus: number, + message: string, + details: Record, +): EnterpriseQuotaDecision { + return { + allowed: false, + code, + status, + httpStatus, + message, + details, + }; +} + +function parseObjectJson(value: string | null): Record { + if (!value) return {}; + try { + const parsed = JSON.parse(value); + return parsed && typeof parsed === 'object' && !Array.isArray(parsed) + ? parsed as Record + : {}; + } catch { + return {}; + } +} + +function nonNegativeInteger(value: unknown): number | undefined { + if (typeof value !== 'number' || !Number.isFinite(value) || value < 0) return undefined; + return Math.floor(value); +} + +function normalizeQuotaPolicy(raw: Record): WorkspaceQuotaPolicy { + return { + maxTraceBytes: nonNegativeInteger(raw.maxTraceBytes), + maxWorkspaceTraceBytes: nonNegativeInteger(raw.maxWorkspaceTraceBytes), + maxConcurrentRuns: nonNegativeInteger(raw.maxConcurrentRuns), + monthlyRunLimit: nonNegativeInteger(raw.monthlyRunLimit), + }; +} + +function normalizeRetentionPolicy(raw: Record): WorkspaceRetentionPolicy { + return { + defaultRetentionDays: nonNegativeInteger(raw.defaultRetentionDays), + traceRetentionDays: nonNegativeInteger(raw.traceRetentionDays), + reportRetentionDays: nonNegativeInteger(raw.reportRetentionDays), + }; +} + +function monthStartMs(now: number): number { + const date = new Date(now); + return Date.UTC(date.getUTCFullYear(), date.getUTCMonth(), 1); +} + +function workspaceTraceBytes(db: Database.Database, scope: EnterpriseRepositoryScope, now: number): number { + const row = db.prepare(` + SELECT COALESCE(SUM(size_bytes), 0) AS total + FROM trace_assets + WHERE tenant_id = ? + AND workspace_id = ? + AND status <> 'deleted' + AND (expires_at IS NULL OR expires_at > ?) + `).get(scope.tenantId, scope.workspaceId, now); + return row?.total ?? 0; +} + +function activeRunCount(db: Database.Database, scope: EnterpriseRepositoryScope): number { + const placeholders = ACTIVE_RUN_STATUSES.map(() => '?').join(', '); + const row = db.prepare(` + SELECT COUNT(*) AS count + FROM analysis_runs + WHERE tenant_id = ? + AND workspace_id = ? + AND status IN (${placeholders}) + `).get(scope.tenantId, scope.workspaceId, ...ACTIVE_RUN_STATUSES); + return row?.count ?? 0; +} + +function monthlyRunCount(db: Database.Database, scope: EnterpriseRepositoryScope, now: number): number { + const row = db.prepare(` + SELECT COUNT(*) AS count + FROM analysis_runs + WHERE tenant_id = ? + AND workspace_id = ? + AND started_at >= ? + `).get(scope.tenantId, scope.workspaceId, monthStartMs(now)); + return row?.count ?? 0; +} + +export function readWorkspaceEnterprisePolicies( + db: Database.Database, + scope: EnterpriseRepositoryScope, +): WorkspaceEnterprisePolicies { + const row = db.prepare(` + SELECT quota_policy, retention_policy + FROM workspaces + WHERE tenant_id = ? + AND id = ? + LIMIT 1 + `).get(scope.tenantId, scope.workspaceId); + return { + quotaPolicy: normalizeQuotaPolicy(parseObjectJson(row?.quota_policy ?? null)), + retentionPolicy: normalizeRetentionPolicy(parseObjectJson(row?.retention_policy ?? null)), + }; +} + +export function resolveEnterpriseRetentionExpiresAt( + db: Database.Database, + scope: EnterpriseRepositoryScope, + artifactType: 'trace' | 'report', + createdAt: number, +): number | null { + const { retentionPolicy } = readWorkspaceEnterprisePolicies(db, scope); + const days = artifactType === 'trace' + ? retentionPolicy.traceRetentionDays ?? retentionPolicy.defaultRetentionDays + : retentionPolicy.reportRetentionDays ?? retentionPolicy.defaultRetentionDays; + if (days === undefined) return null; + return createdAt + days * DAY_MS; +} + +export function evaluateTraceUploadQuota( + context: RequestContext, + sizeBytes: number, + options: { now?: number } = {}, +): EnterpriseQuotaDecision { + if (!resolveFeatureConfig().enterprise) return allowedDecision(); + const now = options.now ?? Date.now(); + const requestedBytes = Math.max(0, Math.floor(sizeBytes)); + const db = openEnterpriseDb(); + try { + const scope = { + tenantId: context.tenantId, + workspaceId: context.workspaceId, + userId: context.userId, + }; + const { quotaPolicy } = readWorkspaceEnterprisePolicies(db, scope); + if (quotaPolicy.maxTraceBytes !== undefined && requestedBytes > quotaPolicy.maxTraceBytes) { + return denyDecision( + 'TRACE_SIZE_QUOTA_EXCEEDED', + 'quota_exceeded', + 413, + 'Trace upload exceeds workspace per-trace quota', + { + requestedBytes, + limitBytes: quotaPolicy.maxTraceBytes, + }, + ); + } + + if (quotaPolicy.maxWorkspaceTraceBytes !== undefined) { + const currentBytes = workspaceTraceBytes(db, scope, now); + if (currentBytes + requestedBytes > quotaPolicy.maxWorkspaceTraceBytes) { + return denyDecision( + 'WORKSPACE_TRACE_STORAGE_QUOTA_EXCEEDED', + 'quota_exceeded', + 409, + 'Trace upload exceeds workspace trace storage quota', + { + currentBytes, + requestedBytes, + limitBytes: quotaPolicy.maxWorkspaceTraceBytes, + }, + ); + } + } + + return allowedDecision(); + } finally { + db.close(); + } +} + +export function evaluateAnalysisRunQuota( + context: RequestContext, + options: { now?: number } = {}, +): EnterpriseQuotaDecision { + if (!resolveFeatureConfig().enterprise) return allowedDecision(); + const now = options.now ?? Date.now(); + const db = openEnterpriseDb(); + try { + const scope = { + tenantId: context.tenantId, + workspaceId: context.workspaceId, + userId: context.userId, + }; + const { quotaPolicy } = readWorkspaceEnterprisePolicies(db, scope); + + if (quotaPolicy.monthlyRunLimit !== undefined) { + const usedRuns = monthlyRunCount(db, scope, now); + if (usedRuns >= quotaPolicy.monthlyRunLimit) { + return denyDecision( + 'MONTHLY_RUN_QUOTA_EXCEEDED', + 'quota_exceeded', + 402, + 'Workspace monthly run quota is exhausted', + { + usedRuns, + limitRuns: quotaPolicy.monthlyRunLimit, + windowStartMs: monthStartMs(now), + }, + ); + } + } + + if (quotaPolicy.maxConcurrentRuns !== undefined) { + const activeRuns = activeRunCount(db, scope); + if (activeRuns >= quotaPolicy.maxConcurrentRuns) { + return denyDecision( + 'CONCURRENT_RUN_QUOTA_EXCEEDED', + 'pending', + 429, + 'Workspace concurrent run quota is full', + { + activeRuns, + limitRuns: quotaPolicy.maxConcurrentRuns, + queuePosition: activeRuns - quotaPolicy.maxConcurrentRuns + 1, + }, + ); + } + } + + return allowedDecision(); + } finally { + db.close(); + } +} diff --git a/backend/src/services/traceMetadataStore.ts b/backend/src/services/traceMetadataStore.ts index a8b68100..4dae147a 100644 --- a/backend/src/services/traceMetadataStore.ts +++ b/backend/src/services/traceMetadataStore.ts @@ -17,6 +17,7 @@ import { type ResourceOwnerFields, } from './resourceOwnership'; import { canReadTraceResource } from './rbac'; +import { resolveEnterpriseRetentionExpiresAt } from './enterpriseQuotaPolicyService'; export interface TraceMetadata extends ResourceOwnerFields { id: string; @@ -27,6 +28,7 @@ export interface TraceMetadata extends ResourceOwnerFields { path?: string; port?: number; externalRpc?: boolean; + expiresAt?: number; } interface TraceAssetRow { @@ -158,6 +160,7 @@ function rowToTraceMetadata(row: TraceAssetRow): TraceMetadata { ...(row.owner_user_id ? { userId: row.owner_user_id } : {}), ...(typeof port === 'number' ? { port } : {}), ...(externalRpc ? { externalRpc: true } : {}), + ...(typeof row.expires_at === 'number' ? { expiresAt: row.expires_at } : {}), }; } @@ -241,6 +244,12 @@ function writeEnterpriseTraceMetadata(metadata: TraceMetadata): void { withEnterpriseTraceDb((db) => { ensureEnterpriseTraceOwner(db, tenantId, workspaceId, ownerUserId ?? undefined); + const expiresAt = resolveEnterpriseRetentionExpiresAt( + db, + { tenantId, workspaceId, ...(ownerUserId ? { userId: ownerUserId } : {}) }, + 'trace', + createdAt, + ); db.prepare(` INSERT INTO trace_assets (id, tenant_id, workspace_id, owner_user_id, local_path, size_bytes, status, metadata_json, created_at, expires_at) @@ -265,7 +274,7 @@ function writeEnterpriseTraceMetadata(metadata: TraceMetadata): void { metadata.status, metadataJsonForRow(metadata), createdAt, - null, + expiresAt, ); }); } @@ -292,7 +301,8 @@ function readEnterpriseTraceMetadata(traceId: string): TraceMetadata | null { SELECT * FROM trace_assets WHERE id = ? - `).get(traceId); + AND (expires_at IS NULL OR expires_at > ?) + `).get(traceId, Date.now()); return row ? rowToTraceMetadata(row) : null; }); } @@ -323,8 +333,9 @@ function listEnterpriseTraceMetadata(): TraceMetadata[] { const rows = db.prepare(` SELECT * FROM trace_assets + WHERE expires_at IS NULL OR expires_at > ? ORDER BY created_at DESC - `).all(); + `).all(Date.now()); return rows.map(rowToTraceMetadata); }); } @@ -360,8 +371,9 @@ export async function listTraceMetadataForContext(context: RequestContext): Prom FROM trace_assets WHERE tenant_id = ? AND workspace_id = ? + AND (expires_at IS NULL OR expires_at > ?) ORDER BY created_at DESC - `).all(context.tenantId, context.workspaceId); + `).all(context.tenantId, context.workspaceId, Date.now()); return rows.map(rowToTraceMetadata).filter(metadata => canReadTraceResource(metadata, context)); }); } @@ -415,7 +427,8 @@ export async function readTraceMetadataForContext( WHERE id = ? AND tenant_id = ? AND workspace_id = ? - `).get(traceId, context.tenantId, context.workspaceId); + AND (expires_at IS NULL OR expires_at > ?) + `).get(traceId, context.tenantId, context.workspaceId, Date.now()); return row ? rowToTraceMetadata(row) : null; }); return isTraceMetadataOwnedByContext(metadata, context) ? metadata : null; diff --git a/docs/features/enterprise-multi-tenant/README.md b/docs/features/enterprise-multi-tenant/README.md index 1836092f..9b5cb980 100644 --- a/docs/features/enterprise-multi-tenant/README.md +++ b/docs/features/enterprise-multi-tenant/README.md @@ -77,7 +77,7 @@ ### 0.5 主线 D：控制面与合规（§18） - [ ] 5.1 tenant / workspace / member / provider / quota 管理 UI 与后端 API - [x] 5.2 `audit_events` 表 + 关键操作埋点（trace / report / provider / memory / cleanup / delete / promote） -- [ ] 5.3 配额 / 预算 / retention policy（§16.1，含 quota_exceeded 终态） +- [x] 5.3 配额 / 预算 / retention policy（§16.1，含 quota_exceeded 终态） - [ ] 5.4 Tenant export bundle（§16.2，含 SHA256 + tenant identity proof） - [ ] 5.5 Tenant tombstone + 7 天硬删窗口 + async purge + audit proof（§16.3） - [ ] 5.6 Custom skill v1 处置（§14.3）：禁用 write endpoint 或修 loader 闭环 @@ -961,6 +961,14 @@ v1 要求： | Trace processor RAM budget 满 | pending 排队；heartbeat stale 后才回收 | | 磁盘空间低于阈值 | 拒绝新 trace upload，提示管理员 cleanup 或扩容 | +当前实现： + +- `workspaces.quota_policy` 读取 JSON 字段：`maxTraceBytes`、`maxWorkspaceTraceBytes`、`maxConcurrentRuns`、`monthlyRunLimit`。 +- trace upload preflight 在落库前拒绝单文件或 workspace trace storage 超额。 +- analysis run preflight 区分 `pending` 型并发 cap 与 `quota_exceeded` 型月度 run cap。 +- runtime 返回 `terminationReason: "max_budget_usd"` 时，`analysis_runs.status` / `analysis_sessions.status` 落为终态 `quota_exceeded`。 +- `workspaces.retention_policy` 读取 JSON 字段：`defaultRetentionDays`、`traceRetentionDays`、`reportRetentionDays`，写入 `trace_assets.expires_at` 与 `report_artifacts.expires_at`，读路径过滤过期 artifact。 + ### 16.2 Tenant export 支持 tenant export bundle：