diff --git a/.mcp.json b/.mcp.json new file mode 100644 index 0000000000..32d5a74110 --- /dev/null +++ b/.mcp.json @@ -0,0 +1,8 @@ +{ + "mcpServers": { + "trigger": { + "command": "bunx", + "args": ["trigger.dev@latest", "mcp"] + } + } +} diff --git a/apps/api/caBundleExtension.ts b/apps/api/caBundleExtension.ts new file mode 100644 index 0000000000..0e602e0b15 --- /dev/null +++ b/apps/api/caBundleExtension.ts @@ -0,0 +1,50 @@ +import type { BuildContext, BuildExtension, BuildManifest } from '@trigger.dev/build'; +import { existsSync } from 'node:fs'; +import { cp, mkdir } from 'node:fs/promises'; +import { dirname, join, resolve } from 'node:path'; + +// Path relative to the monorepo root (apps/api or apps/app → ../../packages/db/certs/...) +const BUNDLE_RELATIVE_FROM_APP = '../../packages/db/certs/rds-global-bundle.pem'; +const BUNDLE_DEST_REL = 'certs/rds-global-bundle.pem'; + +function findBundleSrc(workingDir: string): string | undefined { + // Walk up from workingDir to find the cert — handles both normal checkouts and git worktrees + // where workspaceDir points to the main worktree root (wrong for us). + const candidates = [ + resolve(workingDir, BUNDLE_RELATIVE_FROM_APP), + resolve(workingDir, '../packages/db/certs/rds-global-bundle.pem'), + resolve(workingDir, 'packages/db/certs/rds-global-bundle.pem'), + ]; + + return candidates.find((c) => existsSync(c)); +} + +export function caBundleExtension(): BuildExtension { + return { + name: 'CABundleExtension', + onBuildStart: (context) => { + // Real OS env var at task spawn time — verified flow: + // addLayer.deploy.env → manifest.deploy.sync.env → syncEnvVarsWithServer → + // taskRunProcessProvider injects into worker env before Node TLS init. + context.addLayer({ + id: 'ca-bundle-env', + deploy: { + env: { NODE_EXTRA_CA_CERTS: `/app/${BUNDLE_DEST_REL}` }, + override: true, + }, + }); + }, + onBuildComplete: async (context: BuildContext, manifest: BuildManifest) => { + const src = findBundleSrc(context.workingDir); + if (!src) { + throw new Error( + `CABundleExtension: rds-global-bundle.pem not found. Searched relative to ${context.workingDir}`, + ); + } + const dest = join(manifest.outputPath, BUNDLE_DEST_REL); + await mkdir(dirname(dest), { recursive: true }); + await cp(src, dest); + context.logger.log(`Copied RDS CA bundle to ${BUNDLE_DEST_REL}`); + }, + }; +} diff --git a/apps/api/prisma/client.js b/apps/api/prisma/client.js index 47ab4f329a..f0ef56f567 100644 --- a/apps/api/prisma/client.js +++ b/apps/api/prisma/client.js @@ -2,8 +2,66 @@ Object.defineProperty(exports, "__esModule", { value: true }); exports.db = void 0; const client_1 = require("@prisma/client"); +const adapter_pg_1 = require("@prisma/adapter-pg"); const globalForPrisma = global; -exports.db = globalForPrisma.prisma || new client_1.PrismaClient(); +const LOCAL_HOSTNAMES = new Set(['localhost', '127.0.0.1', '::1']); +function stripSslMode(connectionString) { + const url = new URL(connectionString); + url.searchParams.delete('sslmode'); + return url.toString(); +} +function isLocalhostUrl(connectionString) { + try { + const { hostname } = new URL(connectionString); + // Strip square brackets from IPv6 host form (e.g. [::1] → ::1) + const stripped = hostname.replace(/^\[/, '').replace(/\]$/, ''); + return LOCAL_HOSTNAMES.has(stripped); + } + catch { + // Malformed URL — be conservative and treat as remote so we don't + // accidentally disable TLS verification. + return false; + } +} +function createPrismaClient() { + const rawUrl = process.env.DATABASE_URL; + const isLocalhost = isLocalhostUrl(rawUrl); + // Strategy: + // - Localhost: TLS off (typical dev Postgres has no cert). + // - Remote with NODE_EXTRA_CA_CERTS set: verified TLS using that bundle + // (e.g. Docker with the RDS CA bundle baked in). + // - Remote in explicit opt-out mode (PRISMA_ALLOW_INSECURE_TLS=1): + // unverified TLS — used by Trigger.dev / Vercel envs that connect via + // a tunneled proxy whose cert can't be pinned. Must be set deliberately; + // the previous default ("just turn off verification") silently exposed + // prod connections to MITM. (Cubic finding #1 on PR #2671.) + // - Remote with neither: throw at boot — surface the misconfig instead of + // silently downgrading. + const hasCABundle = !!process.env.NODE_EXTRA_CA_CERTS; + const allowInsecure = process.env.PRISMA_ALLOW_INSECURE_TLS === '1'; + let ssl; + if (isLocalhost) { + ssl = undefined; + } + else if (hasCABundle) { + ssl = true; + } + else if (allowInsecure) { + ssl = { rejectUnauthorized: false }; + } + else { + throw new Error('Refusing to connect to a non-local Postgres without TLS verification. Set NODE_EXTRA_CA_CERTS to a CA bundle, or set PRISMA_ALLOW_INSECURE_TLS=1 if you intentionally want unverified TLS.'); + } + // Strip sslmode from the connection string to avoid conflicts with the explicit ssl option + const url = ssl !== undefined ? stripSslMode(rawUrl) : rawUrl; + const adapter = new adapter_pg_1.PrismaPg({ connectionString: url, ssl }); + return new client_1.PrismaClient({ + adapter, + transactionOptions: { + timeout: 60000, + }, + }); +} +exports.db = globalForPrisma.prisma || createPrismaClient(); if (process.env.NODE_ENV !== 'production') globalForPrisma.prisma = exports.db; -//# sourceMappingURL=client.js.map \ No newline at end of file diff --git a/apps/api/prisma/client.ts b/apps/api/prisma/client.ts index 1bd8069f56..5f3c1738d2 100644 --- a/apps/api/prisma/client.ts +++ b/apps/api/prisma/client.ts @@ -1,7 +1,9 @@ import { PrismaClient } from '@prisma/client'; import { PrismaPg } from '@prisma/adapter-pg'; -const globalForPrisma = global as unknown as { prisma: PrismaClient }; +const globalForPrisma = global as unknown as { prisma?: PrismaClient }; + +const LOCAL_HOSTNAMES = new Set(['localhost', '127.0.0.1', '::1']); function stripSslMode(connectionString: string): string { const url = new URL(connectionString); @@ -9,13 +11,54 @@ function stripSslMode(connectionString: string): string { return url.toString(); } +function isLocalhostUrl(connectionString: string): boolean { + try { + const { hostname } = new URL(connectionString); + // Strip square brackets from IPv6 host form (e.g. [::1] → ::1) + const stripped = hostname.replace(/^\[/, '').replace(/\]$/, ''); + return LOCAL_HOSTNAMES.has(stripped); + } catch { + // Malformed URL — be conservative and treat as remote so we don't + // accidentally disable TLS verification. + return false; + } +} + function createPrismaClient(): PrismaClient { const rawUrl = process.env.DATABASE_URL!; - const isLocalhost = /localhost|127\.0\.0\.1|::1/.test(rawUrl); - // Use verified SSL when NODE_EXTRA_CA_CERTS is set (Docker with RDS CA bundle), - // otherwise fall back to unverified SSL (Trigger.dev, Vercel, other environments). + const isLocalhost = isLocalhostUrl(rawUrl); + // Strategy: + // - Localhost: TLS off (typical dev Postgres has no cert). + // - Remote with NODE_EXTRA_CA_CERTS set: verified TLS using that bundle + // (e.g. Docker with the RDS CA bundle baked in). + // - Remote in explicit opt-out mode (PRISMA_ALLOW_INSECURE_TLS=1): + // unverified TLS — used by Trigger.dev / Vercel envs that connect via + // a tunneled proxy whose cert can't be pinned. Must be set deliberately; + // the previous default ("just turn off verification") silently exposed + // prod connections to MITM. (Cubic finding #1 on PR #2671.) + // - Remote with neither: throw at boot — surface the misconfig instead of + // silently downgrading. const hasCABundle = !!process.env.NODE_EXTRA_CA_CERTS; - const ssl = isLocalhost ? undefined : hasCABundle ? true : { rejectUnauthorized: false }; + const allowInsecure = process.env.PRISMA_ALLOW_INSECURE_TLS === '1'; + let ssl: + | undefined + | { checkServerIdentity: () => undefined } + | { rejectUnauthorized: false }; + if (isLocalhost) { + ssl = undefined; + } else if (hasCABundle) { + // Verified TLS: rely on Node's TLS context (NODE_EXTRA_CA_CERTS adds the AWS + // RDS CA to the trust store). Skip hostname check because connections may + // traverse an AWS NLB whose hostname isn't in the RDS Proxy cert's SAN list. + // The chain check still rejects forged or wrong-CA certs. + ssl = { checkServerIdentity: () => undefined }; + } else if (allowInsecure) { + ssl = { rejectUnauthorized: false }; + } else { + throw new Error( + 'Refusing to connect to a non-local Postgres without TLS verification. Set NODE_EXTRA_CA_CERTS to a CA bundle, or set PRISMA_ALLOW_INSECURE_TLS=1 if you intentionally want unverified TLS.', + ); + } // Strip sslmode from the connection string to avoid conflicts with the explicit ssl option const url = ssl !== undefined ? stripSslMode(rawUrl) : rawUrl; const adapter = new PrismaPg({ connectionString: url, ssl }); @@ -27,6 +70,21 @@ function createPrismaClient(): PrismaClient { }); } -export const db = globalForPrisma.prisma || createPrismaClient(); +// Lazy initialization. Importing this module does NOT construct a Prisma client +// — that only happens on first property access on `db`. Critical so that +// Next.js `next build` (which imports every route handler to analyze it) does +// not trigger the strict TLS check at build time when no actual queries run. +function getClient(): PrismaClient { + if (!globalForPrisma.prisma) { + globalForPrisma.prisma = createPrismaClient(); + } + return globalForPrisma.prisma; +} -if (process.env.NODE_ENV !== 'production') globalForPrisma.prisma = db; +export const db = new Proxy({} as PrismaClient, { + get(_target, prop, _receiver) { + const client = getClient(); + const value = Reflect.get(client, prop, client); + return typeof value === 'function' ? value.bind(client) : value; + }, +}); diff --git a/apps/api/prisma/index.js b/apps/api/prisma/index.js index a818c99668..b6014e58b0 100644 --- a/apps/api/prisma/index.js +++ b/apps/api/prisma/index.js @@ -18,4 +18,3 @@ exports.db = void 0; __exportStar(require("@prisma/client"), exports); var client_1 = require("./client"); Object.defineProperty(exports, "db", { enumerable: true, get: function () { return client_1.db; } }); -//# sourceMappingURL=index.js.map \ No newline at end of file diff --git a/apps/api/src/app.module.ts b/apps/api/src/app.module.ts index 318ba6afde..bf6d8b15fb 100644 --- a/apps/api/src/app.module.ts +++ b/apps/api/src/app.module.ts @@ -14,6 +14,7 @@ import { awsConfig } from './config/aws.config'; import { betterAuthConfig } from './config/better-auth.config'; import { HealthModule } from './health/health.module'; import { OrganizationModule } from './organization/organization.module'; +import { OrganizationAccessModule } from './organization-access/organization-access.module'; import { PoliciesModule } from './policies/policies.module'; import { RisksModule } from './risks/risks.module'; import { TasksModule } from './tasks/tasks.module'; @@ -74,6 +75,7 @@ import { BillingModule } from './billing/billing.module'; ]), AuthModule, OrganizationModule, + OrganizationAccessModule, PeopleModule, RisksModule, VendorsModule, diff --git a/apps/api/src/auth/permission.guard.ts b/apps/api/src/auth/permission.guard.ts index fee0efac3b..033922d499 100644 --- a/apps/api/src/auth/permission.guard.ts +++ b/apps/api/src/auth/permission.guard.ts @@ -187,12 +187,13 @@ export class PermissionGuard implements CanActivate { // the schema rejects every request with `[body] Invalid input`, the // catch in canActivate turns that into a generic "Unable to verify // permissions" 403, and EVERY cookie-authenticated request returns 403. - // Reproduced repo-side via `bun run zod-repro.mjs`. Discovered on - // ENG-221 and the same fix applies here. - const result = await auth.api.hasPermission({ - headers, - body: { permissions, permission: undefined }, - }); + // + // Spell the body out via a separate variable so TypeScript's excess- + // property check (only applied to fresh object literals) doesn't + // reject the extra `permission` key — the runtime accepts the wider + // shape per the union schema. + const body = { permissions, permission: undefined }; + const result = await auth.api.hasPermission({ headers, body }); return result.success === true; } diff --git a/apps/api/src/billing/billing.service.spec.ts b/apps/api/src/billing/billing.service.spec.ts index 05084a2b17..d2cc35915d 100644 --- a/apps/api/src/billing/billing.service.spec.ts +++ b/apps/api/src/billing/billing.service.spec.ts @@ -208,6 +208,87 @@ describe('BillingService', () => { ); }); + it('aggregates wallet credit balances per product in getStatus', async () => { + // The customer-facing /v1/billing/status response is the only + // surface the pentest + BG-check UIs read from. If we ever stop + // including creditBalances here, both UIs silently regress to + // paywalling users whose admin-granted credits would be consumed + // by the create endpoint. Lock the contract with this test. + const listBalances = jest.fn().mockResolvedValue([ + { + id: 'bcb_1', + productKey: 'pentest', + skuKey: null, + balance: 3, + totalGranted: 5, + totalConsumed: 2, + totalRefunded: 0, + lastSource: 'manual', + updatedAt: '2026-05-01T00:00:00.000Z', + }, + { + id: 'bcb_2', + productKey: 'pentest', + skuKey: 'pentest_monthly_1', + balance: 2, + totalGranted: 2, + totalConsumed: 0, + totalRefunded: 0, + lastSource: 'topup', + updatedAt: '2026-05-01T00:00:00.000Z', + }, + { + id: 'bcb_3', + productKey: 'background_check', + skuKey: null, + balance: 4, + totalGranted: 4, + totalConsumed: 0, + totalRefunded: 0, + lastSource: 'manual', + updatedAt: '2026-05-01T00:00:00.000Z', + }, + ]); + const service = new BillingService( + mockStripeService({ + invoices: { list: jest.fn().mockResolvedValue({ data: [] }) }, + customers: { retrieve: jest.fn().mockResolvedValue({}) }, + paymentMethods: { retrieve: jest.fn() }, + }), + { syncSubscriptionItem: jest.fn() } as never, + { listBalances } as never, + ); + + const result = await service.getStatus('org_1'); + + expect(listBalances).toHaveBeenCalledWith('org_1'); + expect(result.creditBalances).toEqual( + expect.arrayContaining([ + { productKey: 'pentest', balance: 5 }, + { productKey: 'background_check', balance: 4 }, + ]), + ); + expect(result.creditBalances).toHaveLength(2); + }); + + it('returns an empty creditBalances array when no credits service is wired in', async () => { + // BillingCreditsService is @Optional() so unit tests can keep + // hand-constructing BillingService without it. Verify the absent- + // dependency branch produces a typesafe empty array, not undefined. + const service = new BillingService( + mockStripeService({ + invoices: { list: jest.fn().mockResolvedValue({ data: [] }) }, + customers: { retrieve: jest.fn().mockResolvedValue({}) }, + paymentMethods: { retrieve: jest.fn() }, + }), + { syncSubscriptionItem: jest.fn() } as never, + ); + + const result = await service.getStatus('org_1'); + + expect(result.creditBalances).toEqual([]); + }); + it('marks trial eligibility false after any product subscription history', async () => { organizationBillingSubscriptionFindMany.mockResolvedValue([ { diff --git a/apps/api/src/billing/billing.service.ts b/apps/api/src/billing/billing.service.ts index 43fa2d84cb..7fe8792bf4 100644 --- a/apps/api/src/billing/billing.service.ts +++ b/apps/api/src/billing/billing.service.ts @@ -2,6 +2,7 @@ import { BadRequestException, Injectable, NotFoundException, + Optional, } from '@nestjs/common'; import { db } from '@db'; import { @@ -12,6 +13,7 @@ import { isSubscriptionBillingSkuKey, } from '@trycompai/billing'; import { StripeService } from '../stripe/stripe.service'; +import { BillingCreditsService } from './billing-credits.service'; import { findOrCreateBillingCustomer } from './billing-customer'; import { BillingEntitlementsService } from './billing-entitlements.service'; import { listBillingInvoices } from './billing-invoices'; @@ -38,6 +40,10 @@ export class BillingService { constructor( private readonly stripeService: StripeService, private readonly entitlements: BillingEntitlementsService, + // Optional so existing unit tests that hand-construct BillingService + // (without a credits service) keep working. In production the + // BillingModule always provides it. + @Optional() private readonly credits?: BillingCreditsService, ) {} async getStatus(organizationId: string): Promise { @@ -67,18 +73,31 @@ export class BillingService { db.backgroundCheckRequest.count({ where: { organizationId } }), db.securityPenetrationTestRun.count({ where: { organizationId } }), ]); - const [invoices, preferences, usageRows] = await Promise.all([ - listBillingInvoices({ - stripeService: this.stripeService, - stripeCustomerId: billing?.stripeCustomerId ?? null, - }), - getBillingPreferences({ - stripeService: this.stripeService, - stripeCustomerId: billing?.stripeCustomerId ?? null, - fallbackCompanyName: organization.name, - }), - listBillingUsageRows({ organizationId, subscriptions }), - ]); + const [invoices, preferences, usageRows, creditBalances] = + await Promise.all([ + listBillingInvoices({ + stripeService: this.stripeService, + stripeCustomerId: billing?.stripeCustomerId ?? null, + }), + getBillingPreferences({ + stripeService: this.stripeService, + stripeCustomerId: billing?.stripeCustomerId ?? null, + fallbackCompanyName: organization.name, + }), + listBillingUsageRows({ organizationId, subscriptions }), + // Sum wallet balances per product. There can be multiple + // BillingCreditBalance rows per (org, product) when grants are + // scoped to a specific SKU, so we aggregate before returning. + this.credits + ? this.credits.listBalances(organizationId) + : Promise.resolve([]), + ]); + + const creditBalancesByProduct = new Map(); + for (const balance of creditBalances) { + const current = creditBalancesByProduct.get(balance.productKey) ?? 0; + creditBalancesByProduct.set(balance.productKey, current + balance.balance); + } return { hasBilling: !!billing, @@ -98,6 +117,9 @@ export class BillingService { currentPeriodEnd: subscription.currentPeriodEnd?.toISOString() ?? null, cancelAtPeriodEnd: subscription.cancelAtPeriodEnd, })), + creditBalances: Array.from(creditBalancesByProduct.entries()).map( + ([productKey, balance]) => ({ productKey, balance }), + ), invoices, }; } diff --git a/apps/api/src/billing/billing.types.ts b/apps/api/src/billing/billing.types.ts index f9db394ae9..fd3b4e13c4 100644 --- a/apps/api/src/billing/billing.types.ts +++ b/apps/api/src/billing/billing.types.ts @@ -1,3 +1,4 @@ +import type { BillingProductKey } from '@trycompai/billing'; import type { BillingInvoice } from './billing-invoices'; import type { BillingPreferences } from './billing-preferences'; @@ -24,6 +25,15 @@ export interface BillingStatus { currentPeriodEnd: string | null; cancelAtPeriodEnd: boolean; }>; + // Aggregated wallet balance per product. Mirrors what + // `BillingEntitlementsService.tryConsumeIncludedUsageForProduct` falls + // back to when a Stripe subscription is missing or exhausted, so the UI + // can keep its allowance display in sync with the backend's actual + // consumption decision. + creditBalances: Array<{ + productKey: BillingProductKey; + balance: number; + }>; invoices: BillingInvoice[]; } diff --git a/apps/api/src/evidence-forms/evidence-forms.controller.ts b/apps/api/src/evidence-forms/evidence-forms.controller.ts index e52d0fbf2f..cec36356d0 100644 --- a/apps/api/src/evidence-forms/evidence-forms.controller.ts +++ b/apps/api/src/evidence-forms/evidence-forms.controller.ts @@ -86,7 +86,7 @@ export class EvidenceFormsController { } @Get('my-submissions') - @RequirePermission('evidence', 'read') + @RequirePermission('portal', 'update') @ApiOperation({ summary: 'Get current user submissions', description: @@ -105,7 +105,7 @@ export class EvidenceFormsController { } @Get('my-submissions/pending-count') - @RequirePermission('evidence', 'read') + @RequirePermission('portal', 'update') @ApiOperation({ summary: 'Get pending submission count for current user', description: @@ -188,7 +188,7 @@ export class EvidenceFormsController { } @Post(':formType/submissions') - @RequirePermission('evidence', 'create') + @RequirePermission('portal', 'update') @ApiOperation({ summary: 'Submit evidence form entry', description: @@ -253,7 +253,7 @@ export class EvidenceFormsController { } @Post('uploads') - @RequirePermission('evidence', 'create') + @RequirePermission('portal', 'update') @ApiOperation({ summary: 'Upload evidence form file', description: diff --git a/apps/api/src/frameworks/frameworks-scores.helper.ts b/apps/api/src/frameworks/frameworks-scores.helper.ts index 1671598d12..c4cfca9cf1 100644 --- a/apps/api/src/frameworks/frameworks-scores.helper.ts +++ b/apps/api/src/frameworks/frameworks-scores.helper.ts @@ -20,7 +20,11 @@ export async function getOverviewScores(organizationId: string) { }), db.task.findMany({ where: { organizationId, archivedAt: null } }), db.member.findMany({ - where: { organizationId, deactivated: false }, + where: { + organizationId, + deactivated: false, + isActive: true, + }, include: { user: true }, }), db.onboarding.findUnique({ diff --git a/apps/api/src/organization-access/organization-access.controller.ts b/apps/api/src/organization-access/organization-access.controller.ts new file mode 100644 index 0000000000..67d38850fa --- /dev/null +++ b/apps/api/src/organization-access/organization-access.controller.ts @@ -0,0 +1,44 @@ +import { + Controller, + HttpCode, + HttpStatus, + Post, + UseGuards, +} from '@nestjs/common'; +import { ApiOperation, ApiTags } from '@nestjs/swagger'; +import { AuthContext, OrganizationId } from '../auth/auth-context.decorator'; +import { HybridAuthGuard } from '../auth/hybrid-auth.guard'; +import { PermissionGuard } from '../auth/permission.guard'; +import { RequirePermission } from '../auth/require-permission.decorator'; +import type { AuthContext as AuthContextType } from '../auth/types'; +import { + AutoApproveResult, + OrganizationAccessService, +} from './organization-access.service'; + +@ApiTags('Organization') +@Controller({ path: 'organization-access', version: '1' }) +@UseGuards(HybridAuthGuard, PermissionGuard) +export class OrganizationAccessController { + constructor( + private readonly organizationAccessService: OrganizationAccessService, + ) {} + + @Post('auto-approve') + @HttpCode(HttpStatus.OK) + @RequirePermission('organization', 'update') + @ApiOperation({ + summary: 'Auto-approve organization access via domain or self-hosted check', + description: + 'Grants hasAccess on the active organization if the requesting user is an internal trycomp.ai user, the deployment is self-hosted, or the user email domain matches the organization website domain and is an active Stripe customer.', + }) + async autoApprove( + @OrganizationId() organizationId: string, + @AuthContext() authContext: AuthContextType, + ): Promise { + return this.organizationAccessService.autoApproveAccess({ + organizationId, + userEmail: authContext.userEmail, + }); + } +} diff --git a/apps/api/src/organization-access/organization-access.module.ts b/apps/api/src/organization-access/organization-access.module.ts new file mode 100644 index 0000000000..378b96789e --- /dev/null +++ b/apps/api/src/organization-access/organization-access.module.ts @@ -0,0 +1,11 @@ +import { Module } from '@nestjs/common'; +import { AuthModule } from '../auth/auth.module'; +import { OrganizationAccessController } from './organization-access.controller'; +import { OrganizationAccessService } from './organization-access.service'; + +@Module({ + imports: [AuthModule], + controllers: [OrganizationAccessController], + providers: [OrganizationAccessService], +}) +export class OrganizationAccessModule {} diff --git a/apps/api/src/organization-access/organization-access.service.spec.ts b/apps/api/src/organization-access/organization-access.service.spec.ts new file mode 100644 index 0000000000..211b60779f --- /dev/null +++ b/apps/api/src/organization-access/organization-access.service.spec.ts @@ -0,0 +1,239 @@ +import { NotFoundException } from '@nestjs/common'; +import { db } from '@db'; +import { OrganizationAccessService } from './organization-access.service'; + +jest.mock('@db', () => ({ + db: { + organization: { + findUnique: jest.fn(), + update: jest.fn(), + }, + }, +})); + +const mockedDb = db as unknown as { + organization: { + findUnique: jest.Mock; + update: jest.Mock; + }; +}; + +const buildService = ( + overrides: Partial<{ isDomainActiveCustomer: jest.Mock }> = {}, +) => { + const isDomainActiveCustomer = + overrides.isDomainActiveCustomer ?? jest.fn().mockResolvedValue(false); + const stripeService = { isDomainActiveCustomer } as never; + return { + service: new OrganizationAccessService(stripeService), + isDomainActiveCustomer, + }; +}; + +describe('OrganizationAccessService', () => { + const ORIGINAL_SELF_HOSTED = process.env.SELF_HOSTED; + const ORIGINAL_NEXT_SELF_HOSTED = process.env.NEXT_PUBLIC_SELF_HOSTED; + + beforeEach(() => { + jest.clearAllMocks(); + delete process.env.SELF_HOSTED; + delete process.env.NEXT_PUBLIC_SELF_HOSTED; + mockedDb.organization.update.mockResolvedValue({}); + }); + + afterAll(() => { + if (ORIGINAL_SELF_HOSTED === undefined) { + delete process.env.SELF_HOSTED; + } else { + process.env.SELF_HOSTED = ORIGINAL_SELF_HOSTED; + } + if (ORIGINAL_NEXT_SELF_HOSTED === undefined) { + delete process.env.NEXT_PUBLIC_SELF_HOSTED; + } else { + process.env.NEXT_PUBLIC_SELF_HOSTED = ORIGINAL_NEXT_SELF_HOSTED; + } + }); + + it('throws NotFoundException when org does not exist', async () => { + mockedDb.organization.findUnique.mockResolvedValue(null); + const { service } = buildService(); + + await expect( + service.autoApproveAccess({ + organizationId: 'org_x', + userEmail: 'a@b.com', + }), + ).rejects.toBeInstanceOf(NotFoundException); + }); + + it('returns already-has-access without writing when org has access', async () => { + mockedDb.organization.findUnique.mockResolvedValue({ + id: 'org_1', + hasAccess: true, + website: 'acme.com', + }); + const { service } = buildService(); + + const result = await service.autoApproveAccess({ + organizationId: 'org_1', + userEmail: 'user@acme.com', + }); + + expect(result).toEqual({ + hasAccess: true, + autoApproved: false, + reason: 'already-has-access', + }); + expect(mockedDb.organization.update).not.toHaveBeenCalled(); + }); + + it('grants on self-hosted via SELF_HOSTED env', async () => { + process.env.SELF_HOSTED = 'true'; + mockedDb.organization.findUnique.mockResolvedValue({ + id: 'org_1', + hasAccess: false, + website: null, + }); + const { service, isDomainActiveCustomer } = buildService(); + + const result = await service.autoApproveAccess({ + organizationId: 'org_1', + userEmail: 'user@gmail.com', + }); + + expect(result).toEqual({ + hasAccess: true, + autoApproved: true, + reason: 'self-hosted', + }); + expect(mockedDb.organization.update).toHaveBeenCalledWith({ + where: { id: 'org_1' }, + data: { hasAccess: true }, + }); + expect(isDomainActiveCustomer).not.toHaveBeenCalled(); + }); + + it('grants on @trycomp.ai email without consulting Stripe', async () => { + mockedDb.organization.findUnique.mockResolvedValue({ + id: 'org_1', + hasAccess: false, + website: 'acme.com', + }); + const { service, isDomainActiveCustomer } = buildService(); + + const result = await service.autoApproveAccess({ + organizationId: 'org_1', + userEmail: 'tofik@trycomp.ai', + }); + + expect(result).toEqual({ + hasAccess: true, + autoApproved: true, + reason: 'trycomp-email', + }); + expect(isDomainActiveCustomer).not.toHaveBeenCalled(); + expect(mockedDb.organization.update).toHaveBeenCalled(); + }); + + it('grants when user email domain matches org website AND is an active Stripe customer', async () => { + mockedDb.organization.findUnique.mockResolvedValue({ + id: 'org_1', + hasAccess: false, + website: 'https://acme.com', + }); + const { service, isDomainActiveCustomer } = buildService({ + isDomainActiveCustomer: jest.fn().mockResolvedValue(true), + }); + + const result = await service.autoApproveAccess({ + organizationId: 'org_1', + userEmail: 'cfo@acme.com', + }); + + expect(result).toEqual({ + hasAccess: true, + autoApproved: true, + reason: 'stripe-customer', + }); + expect(isDomainActiveCustomer).toHaveBeenCalledWith('acme.com'); + expect(mockedDb.organization.update).toHaveBeenCalled(); + }); + + it('does not grant when domain matches but Stripe says not an active customer', async () => { + mockedDb.organization.findUnique.mockResolvedValue({ + id: 'org_1', + hasAccess: false, + website: 'acme.com', + }); + const { service } = buildService({ + isDomainActiveCustomer: jest.fn().mockResolvedValue(false), + }); + + const result = await service.autoApproveAccess({ + organizationId: 'org_1', + userEmail: 'cfo@acme.com', + }); + + expect(result).toEqual({ + hasAccess: false, + autoApproved: false, + reason: 'not-eligible', + }); + expect(mockedDb.organization.update).not.toHaveBeenCalled(); + }); + + it('does not grant when user email domain mismatches org website', async () => { + mockedDb.organization.findUnique.mockResolvedValue({ + id: 'org_1', + hasAccess: false, + website: 'acme.com', + }); + const { service, isDomainActiveCustomer } = buildService(); + + const result = await service.autoApproveAccess({ + organizationId: 'org_1', + userEmail: 'cfo@example.com', + }); + + expect(result.autoApproved).toBe(false); + expect(result.reason).toBe('not-eligible'); + expect(isDomainActiveCustomer).not.toHaveBeenCalled(); + }); + + it('does not grant when user domain is a public mailbox provider', async () => { + mockedDb.organization.findUnique.mockResolvedValue({ + id: 'org_1', + hasAccess: false, + website: 'gmail.com', // pathological — even if website "matches", we refuse + }); + const { service, isDomainActiveCustomer } = buildService(); + + const result = await service.autoApproveAccess({ + organizationId: 'org_1', + userEmail: 'someone@gmail.com', + }); + + expect(result.autoApproved).toBe(false); + expect(isDomainActiveCustomer).not.toHaveBeenCalled(); + }); + + it('does not grant when user email is missing', async () => { + mockedDb.organization.findUnique.mockResolvedValue({ + id: 'org_1', + hasAccess: false, + website: 'acme.com', + }); + const { service } = buildService(); + + const result = await service.autoApproveAccess({ + organizationId: 'org_1', + userEmail: undefined, + }); + + expect(result).toEqual({ + hasAccess: false, + autoApproved: false, + reason: 'not-eligible', + }); + }); +}); diff --git a/apps/api/src/organization-access/organization-access.service.ts b/apps/api/src/organization-access/organization-access.service.ts new file mode 100644 index 0000000000..818acecfc9 --- /dev/null +++ b/apps/api/src/organization-access/organization-access.service.ts @@ -0,0 +1,115 @@ +import { Injectable, Logger, NotFoundException } from '@nestjs/common'; +import { db } from '@db'; +import { extractDomain, isPublicEmailDomain } from '../stripe/domain.utils'; +import { StripeService } from '../stripe/stripe.service'; + +export type AutoApproveReason = + | 'already-has-access' + | 'self-hosted' + | 'trycomp-email' + | 'stripe-customer' + | 'not-eligible'; + +export interface AutoApproveResult { + hasAccess: boolean; + autoApproved: boolean; + reason: AutoApproveReason; +} + +const isSelfHosted = (): boolean => + process.env.SELF_HOSTED === 'true' || + process.env.NEXT_PUBLIC_SELF_HOSTED === 'true'; + +@Injectable() +export class OrganizationAccessService { + private readonly logger = new Logger(OrganizationAccessService.name); + + constructor(private readonly stripeService: StripeService) {} + + /** + * Decide whether to grant `hasAccess` to the org for the given user, and + * persist the flag if so. Membership is enforced upstream by HybridAuthGuard + * (the auth guard rejects requests where the user isn't an active member of + * the active organization), but we re-fetch the organization here for the + * website + current hasAccess values. + * + * Decision rules (in order): + * 1. Org already has access → no-op. + * 2. Self-hosted instance → grant. + * 3. User email at @trycomp.ai → grant (internal team). + * 4. User email domain matches the org website domain AND that domain has + * an active Stripe customer → grant. Public mailbox domains are + * excluded. + * 5. Otherwise → not eligible. + */ + async autoApproveAccess(input: { + organizationId: string; + userEmail: string | undefined; + }): Promise { + const { organizationId, userEmail } = input; + + const organization = await db.organization.findUnique({ + where: { id: organizationId }, + select: { id: true, hasAccess: true, website: true }, + }); + + if (!organization) { + throw new NotFoundException('Organization not found'); + } + + if (organization.hasAccess) { + return { + hasAccess: true, + autoApproved: false, + reason: 'already-has-access', + }; + } + + if (isSelfHosted()) { + await this.grantAccess(organizationId); + this.logger.log( + `Auto-approved org ${organizationId} (reason: self-hosted)`, + ); + return { hasAccess: true, autoApproved: true, reason: 'self-hosted' }; + } + + const userEmailDomain = extractDomain(userEmail); + const orgWebsiteDomain = extractDomain(organization.website); + + if (!userEmailDomain) { + return { hasAccess: false, autoApproved: false, reason: 'not-eligible' }; + } + + const isTrycompEmail = userEmailDomain === 'trycomp.ai'; + + const canAutoApproveViaDomain = + !isTrycompEmail && + Boolean(orgWebsiteDomain) && + userEmailDomain === orgWebsiteDomain && + !isPublicEmailDomain(userEmailDomain); + + const isStripeCustomer = canAutoApproveViaDomain + ? await this.stripeService.isDomainActiveCustomer(userEmailDomain) + : false; + + if (isTrycompEmail || isStripeCustomer) { + await this.grantAccess(organizationId); + const reason: AutoApproveReason = isTrycompEmail + ? 'trycomp-email' + : 'stripe-customer'; + this.logger.log( + `Auto-approved org ${organizationId} (reason: ${reason}, domain: ${userEmailDomain})`, + ); + return { hasAccess: true, autoApproved: true, reason }; + } + + return { hasAccess: false, autoApproved: false, reason: 'not-eligible' }; + } + + private async grantAccess(organizationId: string): Promise { + await db.organization.update({ + where: { id: organizationId }, + data: { hasAccess: true }, + }); + } +} diff --git a/apps/api/src/policies/policies.controller.ts b/apps/api/src/policies/policies.controller.ts index 105260e843..67a4587c40 100644 --- a/apps/api/src/policies/policies.controller.ts +++ b/apps/api/src/policies/policies.controller.ts @@ -1319,7 +1319,7 @@ Keep responses helpful and focused on the policy editing task.`; ]; const result = streamText({ - model: openai('gpt-5.1'), + model: openai('gpt-5.5'), system: systemPrompt, messages: convertToModelMessages(messages), }); diff --git a/apps/api/src/risks/risks.service.ts b/apps/api/src/risks/risks.service.ts index 37c920d4ec..e4d29cff92 100644 --- a/apps/api/src/risks/risks.service.ts +++ b/apps/api/src/risks/risks.service.ts @@ -8,6 +8,7 @@ import { db, Prisma } from '@db'; import { CreateRiskDto } from './dto/create-risk.dto'; import { GetRisksQueryDto } from './dto/get-risks-query.dto'; import { UpdateRiskDto } from './dto/update-risk.dto'; +import { resolveStrategyDescriptionUpdate } from './strategy-descriptions'; export interface PaginatedRisksResult { data: Prisma.RiskGetPayload<{ @@ -94,6 +95,10 @@ export class RisksService { }, }, }, + // Linked task statuses are needed by the table to compute the + // current (interpolated) severity score so the badge reflects + // treatment progress, not just inherent risk. + tasks: { select: { id: true, status: true } }, }, }), db.risk.count({ where }), @@ -128,6 +133,14 @@ export class RisksService { user: true, }, }, + tasks: { + select: { + id: true, + title: true, + status: true, + controls: { select: { id: true, name: true } }, + }, + }, }, }); @@ -182,8 +195,8 @@ export class RisksService { updateRiskDto: UpdateRiskDto, ) { try { - // First check if the risk exists in the organization - await this.findById(id, organizationId); + // Need the existing row to resolve strategy/description swaps below. + const existing = await this.findById(id, organizationId); if (updateRiskDto.assigneeId) { await this.validateAssigneeNotPlatformAdmin( @@ -192,9 +205,18 @@ export class RisksService { ); } + // Keep per-strategy descriptions independent in the strategyDescriptions + // JSON map: a Mitigate plan, an Accept rationale, and a Transfer + // rationale all live alongside each other, swapped in/out of the active + // `treatmentStrategyDescription` field as the user changes strategy. + const resolvedStrategyFields = resolveStrategyDescriptionUpdate( + existing, + updateRiskDto, + ); + const updatedRisk = await db.risk.update({ where: { id }, - data: updateRiskDto, + data: { ...updateRiskDto, ...resolvedStrategyFields }, }); this.logger.log(`Updated risk: ${updatedRisk.title} (${id})`); diff --git a/apps/api/src/risks/strategy-descriptions.spec.ts b/apps/api/src/risks/strategy-descriptions.spec.ts new file mode 100644 index 0000000000..6da074d226 --- /dev/null +++ b/apps/api/src/risks/strategy-descriptions.spec.ts @@ -0,0 +1,124 @@ +import { describe, expect, it, jest } from '@jest/globals'; + +// Mock @db before importing the helper so we don't pull in the real Prisma +// client (which requires DATABASE_URL at module load). +jest.mock('@db', () => ({ + RiskTreatmentType: { + accept: 'accept', + avoid: 'avoid', + mitigate: 'mitigate', + transfer: 'transfer', + }, +})); + +import { RiskTreatmentType } from '@db'; +import { resolveStrategyDescriptionUpdate } from './strategy-descriptions'; + +const baseExisting = { + treatmentStrategy: RiskTreatmentType.mitigate, + treatmentStrategyDescription: 'Current mitigate plan', + strategyDescriptions: {} as unknown, +}; + +describe('resolveStrategyDescriptionUpdate', () => { + it('returns empty when neither strategy nor description is changing', () => { + expect(resolveStrategyDescriptionUpdate(baseExisting, {})).toEqual({}); + expect( + resolveStrategyDescriptionUpdate(baseExisting, { assigneeId: 'mbr_1' } as never), + ).toEqual({}); + }); + + it('on strategy change: saves current description into the OLD slot', () => { + const result = resolveStrategyDescriptionUpdate(baseExisting, { + treatmentStrategy: RiskTreatmentType.accept, + }); + expect(result.treatmentStrategy).toBe(RiskTreatmentType.accept); + expect(result.strategyDescriptions).toEqual({ + mitigate: 'Current mitigate plan', + }); + // Active text becomes empty (no Accept rationale saved yet) + expect(result.treatmentStrategyDescription).toBeNull(); + }); + + it('on strategy change: loads NEW strategy slot into active text when present', () => { + const result = resolveStrategyDescriptionUpdate( + { + ...baseExisting, + strategyDescriptions: { + accept: 'We accept this risk because of cost-benefit analysis.', + }, + }, + { treatmentStrategy: RiskTreatmentType.accept }, + ); + expect(result.treatmentStrategyDescription).toBe( + 'We accept this risk because of cost-benefit analysis.', + ); + // Mitigate plan still preserved in the map + expect(result.strategyDescriptions).toEqual({ + mitigate: 'Current mitigate plan', + accept: 'We accept this risk because of cost-benefit analysis.', + }); + }); + + it('on description change without strategy change: mirrors into active strategy slot', () => { + const result = resolveStrategyDescriptionUpdate(baseExisting, { + treatmentStrategyDescription: 'Updated mitigate plan', + }); + expect(result.treatmentStrategy).toBeUndefined(); + expect(result.treatmentStrategyDescription).toBe('Updated mitigate plan'); + expect(result.strategyDescriptions).toEqual({ + mitigate: 'Updated mitigate plan', + }); + }); + + it('on both change: explicit description wins as the new active text', () => { + const result = resolveStrategyDescriptionUpdate( + { + ...baseExisting, + strategyDescriptions: { + accept: 'Old accept rationale', + }, + }, + { + treatmentStrategy: RiskTreatmentType.accept, + treatmentStrategyDescription: 'Brand-new accept rationale', + }, + ); + expect(result.treatmentStrategy).toBe(RiskTreatmentType.accept); + expect(result.treatmentStrategyDescription).toBe('Brand-new accept rationale'); + expect(result.strategyDescriptions).toEqual({ + mitigate: 'Current mitigate plan', + accept: 'Brand-new accept rationale', + }); + }); + + it('clears the slot when description is set to empty', () => { + const result = resolveStrategyDescriptionUpdate( + { + ...baseExisting, + strategyDescriptions: { mitigate: 'old text' }, + }, + { treatmentStrategyDescription: '' }, + ); + expect(result.strategyDescriptions).toEqual({}); + }); + + it('handles malformed strategyDescriptions gracefully', () => { + const result = resolveStrategyDescriptionUpdate( + { ...baseExisting, strategyDescriptions: 'not an object' as unknown }, + { treatmentStrategy: RiskTreatmentType.accept }, + ); + // Falls back to empty map; saves current Mitigate text + expect(result.strategyDescriptions).toEqual({ + mitigate: 'Current mitigate plan', + }); + }); + + it('does not save an empty old description into the slot', () => { + const result = resolveStrategyDescriptionUpdate( + { ...baseExisting, treatmentStrategyDescription: '' }, + { treatmentStrategy: RiskTreatmentType.accept }, + ); + expect(result.strategyDescriptions).toEqual({}); + }); +}); diff --git a/apps/api/src/risks/strategy-descriptions.ts b/apps/api/src/risks/strategy-descriptions.ts new file mode 100644 index 0000000000..1511392286 --- /dev/null +++ b/apps/api/src/risks/strategy-descriptions.ts @@ -0,0 +1,91 @@ +import { RiskTreatmentType } from '@db'; + +interface EntityWithStrategy { + treatmentStrategy: RiskTreatmentType; + treatmentStrategyDescription: string | null; + strategyDescriptions: unknown; +} + +interface UpdateInput { + treatmentStrategy?: RiskTreatmentType; + treatmentStrategyDescription?: string | null; +} + +export interface ResolvedStrategyUpdate { + treatmentStrategy?: RiskTreatmentType; + treatmentStrategyDescription?: string | null; + strategyDescriptions?: Record; +} + +/** + * Computes the data fields that need to be persisted when an update touches + * `treatmentStrategy` and/or `treatmentStrategyDescription`. The logic keeps + * each strategy's description independent in the `strategyDescriptions` JSON + * column so a Mitigate plan, an Accept rationale, and a Transfer rationale + * can coexist on the same risk/vendor. + * + * Cases: + * - Strategy changes → save the current description into the OLD strategy + * slot, then load the NEW strategy's saved text (if any) back into the + * active `treatmentStrategyDescription`. + * - Description changes (no strategy change) → mirror it into the active + * strategy's slot in `strategyDescriptions`. + * - Both change → strategy swap runs first; explicit description in the + * update wins as the new active text and is mirrored into the new + * strategy's slot. + */ +export function resolveStrategyDescriptionUpdate( + existing: EntityWithStrategy, + update: UpdateInput, +): ResolvedStrategyUpdate { + if ( + update.treatmentStrategy === undefined && + update.treatmentStrategyDescription === undefined + ) { + return {}; + } + + const map = parseStrategyMap(existing.strategyDescriptions); + const oldStrategy = existing.treatmentStrategy; + const newStrategy = update.treatmentStrategy ?? oldStrategy; + const isStrategyChange = + update.treatmentStrategy !== undefined && update.treatmentStrategy !== oldStrategy; + + const result: ResolvedStrategyUpdate = {}; + + if (isStrategyChange) { + if ((existing.treatmentStrategyDescription ?? '').length > 0) { + map[oldStrategy] = existing.treatmentStrategyDescription as string; + } else { + delete map[oldStrategy]; + } + result.treatmentStrategy = newStrategy; + if (update.treatmentStrategyDescription === undefined) { + result.treatmentStrategyDescription = map[newStrategy] ?? null; + } + } + + if (update.treatmentStrategyDescription !== undefined) { + const next = update.treatmentStrategyDescription ?? ''; + if (next.length > 0) { + map[newStrategy] = next; + } else { + delete map[newStrategy]; + } + result.treatmentStrategyDescription = update.treatmentStrategyDescription; + } + + result.strategyDescriptions = map; + return result; +} + +function parseStrategyMap(raw: unknown): Record { + if (raw && typeof raw === 'object' && !Array.isArray(raw)) { + const out: Record = {}; + for (const [k, v] of Object.entries(raw as Record)) { + if (typeof v === 'string' && v.length > 0) out[k] = v; + } + return out; + } + return {}; +} diff --git a/apps/api/src/stripe/domain.utils.spec.ts b/apps/api/src/stripe/domain.utils.spec.ts new file mode 100644 index 0000000000..b9df4964a0 --- /dev/null +++ b/apps/api/src/stripe/domain.utils.spec.ts @@ -0,0 +1,46 @@ +import { extractDomain, isPublicEmailDomain } from './domain.utils'; + +describe('extractDomain', () => { + it('returns null for empty input', () => { + expect(extractDomain(null)).toBeNull(); + expect(extractDomain(undefined)).toBeNull(); + expect(extractDomain('')).toBeNull(); + }); + + it('extracts domain from email', () => { + expect(extractDomain('user@acme.com')).toBe('acme.com'); + expect(extractDomain('USER@ACME.COM')).toBe('acme.com'); + }); + + it('extracts hostname from a URL and strips www.', () => { + expect(extractDomain('https://www.acme.com/foo')).toBe('acme.com'); + expect(extractDomain('http://acme.com')).toBe('acme.com'); + }); + + it('adds protocol when missing', () => { + expect(extractDomain('acme.com')).toBe('acme.com'); + expect(extractDomain('www.acme.com/path')).toBe('acme.com'); + }); + + it('returns null on garbage input', () => { + expect(extractDomain(' ')).toBeNull(); + }); +}); + +describe('isPublicEmailDomain', () => { + it('flags common free providers', () => { + expect(isPublicEmailDomain('gmail.com')).toBe(true); + expect(isPublicEmailDomain('outlook.com')).toBe(true); + expect(isPublicEmailDomain('icloud.com')).toBe(true); + }); + + it('is case- and trailing-dot-insensitive', () => { + expect(isPublicEmailDomain('GMAIL.COM')).toBe(true); + expect(isPublicEmailDomain('gmail.com.')).toBe(true); + }); + + it('returns false for a real company domain', () => { + expect(isPublicEmailDomain('acme.com')).toBe(false); + expect(isPublicEmailDomain('trycomp.ai')).toBe(false); + }); +}); diff --git a/apps/api/src/stripe/domain.utils.ts b/apps/api/src/stripe/domain.utils.ts new file mode 100644 index 0000000000..954f5646f6 --- /dev/null +++ b/apps/api/src/stripe/domain.utils.ts @@ -0,0 +1,57 @@ +// Public/free mailbox providers — domain ownership of these does NOT imply +// company affiliation, so they must never be used for domain-based auto-approval. +const PUBLIC_EMAIL_DOMAINS = new Set([ + // Google + 'gmail.com', + 'googlemail.com', + // Microsoft + 'outlook.com', + 'hotmail.com', + 'live.com', + 'msn.com', + // Yahoo + 'yahoo.com', + 'ymail.com', + // Apple + 'icloud.com', + 'me.com', + 'mac.com', + // Proton + 'proton.me', + 'protonmail.com', + 'pm.me', + // AOL + 'aol.com', +]); + +export const isPublicEmailDomain = (domain: string): boolean => { + const normalized = domain.toLowerCase().trim().replace(/\.$/, ''); + return PUBLIC_EMAIL_DOMAINS.has(normalized); +}; + +/** + * Extract a normalized domain from either a website URL or an email address. + * Returns null on empty/invalid input. + */ +export const extractDomain = ( + input: string | null | undefined, +): string | null => { + if (!input) return null; + + try { + if (input.includes('@') && !input.includes('://')) { + const domain = input.split('@')[1]?.toLowerCase().trim(); + return domain || null; + } + + let url = input.trim().toLowerCase(); + if (!url.startsWith('http://') && !url.startsWith('https://')) { + url = `https://${url}`; + } + + const parsed = new URL(url); + return parsed.hostname.replace(/^www\./, ''); + } catch { + return null; + } +}; diff --git a/apps/api/src/stripe/stripe.service.ts b/apps/api/src/stripe/stripe.service.ts index 3eac26d69a..274a29ee8e 100644 --- a/apps/api/src/stripe/stripe.service.ts +++ b/apps/api/src/stripe/stripe.service.ts @@ -1,5 +1,6 @@ import { Injectable, Logger } from '@nestjs/common'; import Stripe from 'stripe'; +import { isPublicEmailDomain } from './domain.utils'; @Injectable() export class StripeService { @@ -28,4 +29,102 @@ export class StripeService { isConfigured(): boolean { return this.client !== null; } + + /** + * Look up a Stripe customer by company domain. Prefers an exact match against + * the customer's `domain` metadata, then falls back to scanning customers + * whose primary email belongs to that domain. Public mailbox domains + * (gmail.com, etc.) are never matched — domain ownership of those does not + * imply company affiliation. + */ + async findCustomerByDomain( + domain: string, + ): Promise<{ customerId: string; customerName: string | null } | null> { + if (!this.client || !domain) { + return null; + } + + const normalizedDomain = domain.toLowerCase().trim().replace(/\.$/, ''); + if (isPublicEmailDomain(normalizedDomain)) { + return null; + } + + try { + const byMetadata = await this.client.customers.search({ + query: `metadata["domain"]:"${normalizedDomain}"`, + limit: 1, + }); + + if (byMetadata.data.length > 0) { + const customer = byMetadata.data[0]; + return { + customerId: customer.id, + customerName: customer.name ?? null, + }; + } + + // `email~` is substring matching — re-filter for exact email-domain match. + const byEmail = await this.client.customers.search({ + query: `email~"@${normalizedDomain}"`, + limit: 25, + }); + + const exactMatch = byEmail.data.find((customer) => { + const email = customer.email ?? ''; + const emailDomain = email.split('@')[1]?.toLowerCase().trim() ?? ''; + return emailDomain === normalizedDomain; + }); + + if (exactMatch) { + return { + customerId: exactMatch.id, + customerName: exactMatch.name ?? null, + }; + } + + return null; + } catch (error) { + this.logger.error( + `Error searching Stripe customers by domain "${normalizedDomain}"`, + error instanceof Error ? error.stack : undefined, + ); + return null; + } + } + + /** + * Returns true when the given domain belongs to a Stripe customer with at + * least one active subscription. Used for domain-based auto-approval of + * organization access. + */ + async isDomainActiveCustomer(domain: string): Promise { + if (!this.client) { + return false; + } + + const normalizedDomain = domain.toLowerCase().trim().replace(/\.$/, ''); + if (!normalizedDomain || isPublicEmailDomain(normalizedDomain)) { + return false; + } + + const customer = await this.findCustomerByDomain(normalizedDomain); + if (!customer) { + return false; + } + + try { + const subscriptions = await this.client.subscriptions.list({ + customer: customer.customerId, + status: 'active', + limit: 1, + }); + return subscriptions.data.length > 0; + } catch (error) { + this.logger.error( + `Error checking active subscriptions for customer "${customer.customerId}"`, + error instanceof Error ? error.stack : undefined, + ); + return false; + } + } } diff --git a/apps/api/src/trigger/vendor/vendor-risk-assessment-task.ts b/apps/api/src/trigger/vendor/vendor-risk-assessment-task.ts index 23abf82887..fd664a6e54 100644 --- a/apps/api/src/trigger/vendor/vendor-risk-assessment-task.ts +++ b/apps/api/src/trigger/vendor/vendor-risk-assessment-task.ts @@ -1,20 +1,17 @@ import { db, - Impact, - Likelihood, TaskItemPriority, TaskItemStatus, VendorStatus, type TaskItemEntityType, } from '@db'; -import { openai } from '@ai-sdk/openai'; import type { Prisma } from '@db'; import type { Task } from '@trigger.dev/sdk'; import { logger, metadata, queue, schemaTask, tags } from '@trigger.dev/sdk'; -import { generateObject } from 'ai'; import { z } from 'zod'; import { resolveTaskCreatorAndAssignee } from './vendor-risk-assessment/assignee'; +import { extractInherentRisk } from './vendor-risk-assessment/assessment-output'; import { VENDOR_RISK_ASSESSMENT_TASK_ID } from './vendor-risk-assessment/constants'; import { buildRiskAssessmentDescription, @@ -174,124 +171,6 @@ function parseRiskAssessmentJson(value: string): Prisma.InputJsonValue { return parsed; } -const riskLevelSchema = z - .object({ - riskLevel: z.string().optional(), - }) - .passthrough(); - -function extractRiskLevel(value: Prisma.InputJsonValue): string | null { - const parsed = riskLevelSchema.safeParse(value); - if (!parsed.success) { - return null; - } - return parsed.data.riskLevel ?? null; -} - -/** - * Risk level categories that map to database enums: - * - critical → Likelihood.very_likely / Impact.severe (highest) - * - high → Likelihood.likely / Impact.major - * - medium → Likelihood.possible / Impact.moderate - * - low → Likelihood.unlikely / Impact.minor - * - very_low → Likelihood.very_unlikely / Impact.insignificant (lowest) - */ -type NormalizedRiskLevel = 'critical' | 'high' | 'medium' | 'low' | 'very_low'; - -const normalizedRiskLevelSchema = z.object({ - riskLevel: z - .enum(['critical', 'high', 'medium', 'low', 'very_low']) - .describe( - 'The normalized risk level - must be exactly one of these values', - ), -}); - -/** - * Use AI to normalize any risk level string to one of our exact enum values. - * Uses gpt-4o-mini (fast and cheap) with structured output to ensure valid values. - */ -async function normalizeRiskLevel( - rawRiskLevel: string | null | undefined, -): Promise { - if (!rawRiskLevel?.trim()) { - return null; - } - - try { - const result = await generateObject({ - model: openai('gpt-5.2'), - schema: normalizedRiskLevelSchema, - prompt: `Classify this vendor security risk level into exactly one of these 5 categories. - -Risk level from assessment: "${rawRiskLevel}" - -Categories (highest to lowest risk): -- critical: Highest risk (severe, extreme, very high, critical concerns) -- high: Significant risk (high, major issues) -- medium: Moderate risk (medium, moderate, average) -- low: Low risk (low, minimal, minor) -- very_low: Minimal risk (very low, negligible, none) - -Rules: -- Return exactly one of: critical, high, medium, low, very_low -- If ambiguous (e.g., "Low to Moderate"), pick the HIGHER risk to be conservative`, - }); - - logger.info('Normalized risk level', { - rawRiskLevel, - normalizedRiskLevel: result.object.riskLevel, - }); - - return result.object.riskLevel; - } catch (error) { - logger.warn('Failed to normalize risk level, defaulting to medium', { - rawRiskLevel, - error: error instanceof Error ? error.message : String(error), - }); - return 'medium'; - } -} - -function mapRiskLevelToLikelihood( - normalizedLevel: NormalizedRiskLevel | null, -): Likelihood { - switch (normalizedLevel) { - case 'critical': - return Likelihood.very_likely; - case 'high': - return Likelihood.likely; - case 'medium': - return Likelihood.possible; - case 'low': - return Likelihood.unlikely; - case 'very_low': - return Likelihood.very_unlikely; - default: - // Default to medium (safer than lowest) - return Likelihood.possible; - } -} - -function mapRiskLevelToImpact( - normalizedLevel: NormalizedRiskLevel | null, -): Impact { - switch (normalizedLevel) { - case 'critical': - return Impact.severe; - case 'high': - return Impact.major; - case 'medium': - return Impact.moderate; - case 'low': - return Impact.minor; - case 'very_low': - return Impact.insignificant; - default: - // Default to medium (safer than lowest) - return Impact.moderate; - } -} - /** * Valid compliance badge types for trust portal */ @@ -857,7 +736,9 @@ export const vendorRiskAssessmentTask: Task< verifiedCertifications: verifiedCount, links: linkCount, hasAssessment: Boolean(result.securityAssessment), - riskLevel: result.riskLevel ?? 'none', + likelihood: result.likelihood ?? 'none', + impact: result.impact ?? 'none', + hasRationale: Boolean(result.rationale), }); // Report each finding individually with delays so the UI @@ -1043,22 +924,23 @@ export const vendorRiskAssessmentTask: Task< }); // Extract risk level and badges - logger.info('🎯 Normalizing risk level', { + logger.info('🎯 Extracting inherent risk from assessment payload', { vendor: payload.vendorName, }); - const rawRiskLevel = extractRiskLevel(data); - const normalizedRiskLvl = await normalizeRiskLevel(rawRiskLevel); - const inherentProbability = mapRiskLevelToLikelihood(normalizedRiskLvl); - const inherentImpact = mapRiskLevelToImpact(normalizedRiskLvl); - const residualProbability = mapRiskLevelToLikelihood(normalizedRiskLvl); - const residualImpact = mapRiskLevelToImpact(normalizedRiskLvl); + const inherentRisk = extractInherentRisk(data); + if (!inherentRisk) { + logger.warn( + '⚠️ Assessment payload missing likelihood/impact — preserving existing vendor scores', + { vendor: payload.vendorName }, + ); + } const complianceBadges = extractComplianceBadges(data); const logoUrl = generateLogoUrl(vendor.website); logger.info('📊 Risk level and badges extracted', { vendor: payload.vendorName, - rawRiskLevel, - normalizedRiskLevel: normalizedRiskLvl, + inherentLikelihood: inherentRisk?.likelihood ?? null, + inherentImpact: inherentRisk?.impact ?? null, hasBadges: Boolean(complianceBadges), badgeCount: Array.isArray(complianceBadges) ? complianceBadges.length @@ -1078,14 +960,23 @@ export const vendorRiskAssessmentTask: Task< ), }); - // Update vendor with core data (keep status in_progress — news may still be loading) + // Update vendor with core data (keep status in_progress — news may still be loading). + // Only write risk fields when the assessment payload produced a valid + // inherentRisk — otherwise preserve whatever the vendor row already has + // (e.g. pre-ENG-221 globalVendors.riskAssessmentData that lacks + // likelihood/impact). Residual defaults to inherent until a human + // mitigates. Preserved behavior. await db.vendor.update({ where: { id: vendor.id }, data: { - inherentProbability, - inherentImpact, - residualProbability, - residualImpact, + ...(inherentRisk + ? { + inherentProbability: inherentRisk.likelihood, + inherentImpact: inherentRisk.impact, + residualProbability: inherentRisk.likelihood, + residualImpact: inherentRisk.impact, + } + : {}), ...(complianceBadges ? { complianceBadges } : {}), ...(logoUrl ? { logoUrl } : {}), }, diff --git a/apps/api/src/trigger/vendor/vendor-risk-assessment/agent-schema.ts b/apps/api/src/trigger/vendor/vendor-risk-assessment/agent-schema.ts index f9f1f9541e..fcf888b62f 100644 --- a/apps/api/src/trigger/vendor/vendor-risk-assessment/agent-schema.ts +++ b/apps/api/src/trigger/vendor/vendor-risk-assessment/agent-schema.ts @@ -11,6 +11,26 @@ const dateStringOrEmptySchema = z .nullable(); export const vendorRiskAssessmentAgentSchema = z.object({ + /** + * ENG-221: replaces risk_level. Likelihood and impact are scored + * independently so vendors can land on any cell of the 5x5 matrix + * instead of pooling on the diagonal. + */ + likelihood: z + .enum(['very_unlikely', 'unlikely', 'possible', 'likely', 'very_likely']) + .optional() + .nullable(), + impact: z + .enum(['insignificant', 'minor', 'moderate', 'major', 'severe']) + .optional() + .nullable(), + rationale: z.string().optional().nullable(), + /** + * Legacy single-dimension score retained as optional so stored payloads + * from before ENG-221 still parse. New assessments should set + * `likelihood` + `impact` + `rationale` instead. + */ + // TODO(ENG-221 follow-up): remove once globalVendors.riskAssessmentData backfill ships. risk_level: z.string().optional().nullable(), security_assessment: z.string().optional().nullable(), last_researched_at: dateStringOrEmptySchema, diff --git a/apps/api/src/trigger/vendor/vendor-risk-assessment/agent-types.ts b/apps/api/src/trigger/vendor/vendor-risk-assessment/agent-types.ts index b934550b98..b95060eeb0 100644 --- a/apps/api/src/trigger/vendor/vendor-risk-assessment/agent-types.ts +++ b/apps/api/src/trigger/vendor/vendor-risk-assessment/agent-types.ts @@ -31,11 +31,43 @@ export type VendorRiskAssessmentNewsItem = { sentiment?: VendorRiskAssessmentNewsSentiment | null; }; +/** + * Likelihood enum string values — mirrors the Prisma `Likelihood` enum. + * Kept as a literal union here so this module stays independent from + * `@db` / `@prisma/client`. + */ +export type VendorRiskAssessmentLikelihood = + | 'very_unlikely' + | 'unlikely' + | 'possible' + | 'likely' + | 'very_likely'; + +/** + * Impact enum string values — mirrors the Prisma `Impact` enum. + */ +export type VendorRiskAssessmentImpact = + | 'insignificant' + | 'minor' + | 'moderate' + | 'major' + | 'severe'; + export type VendorRiskAssessmentDataV1 = { kind: 'vendorRiskAssessmentV1'; vendorName?: string | null; vendorWebsite?: string | null; lastResearchedAt?: string | null; + /** + * ENG-221: two independent dimensions. Preferred over the legacy + * single-bucket `riskLevel`. New assessments always set these; legacy + * payloads may only have `riskLevel`. + */ + likelihood?: VendorRiskAssessmentLikelihood | null; + impact?: VendorRiskAssessmentImpact | null; + rationale?: string | null; + /** Legacy single-bucket score. Retained for pre-ENG-221 payloads. */ + // TODO(ENG-221 follow-up): remove once globalVendors.riskAssessmentData backfill ships. riskLevel?: string | null; securityAssessment?: string | null; certifications?: VendorRiskAssessmentCertification[] | null; diff --git a/apps/api/src/trigger/vendor/vendor-risk-assessment/assessment-output.spec.ts b/apps/api/src/trigger/vendor/vendor-risk-assessment/assessment-output.spec.ts new file mode 100644 index 0000000000..417c49030d --- /dev/null +++ b/apps/api/src/trigger/vendor/vendor-risk-assessment/assessment-output.spec.ts @@ -0,0 +1,99 @@ +import { Impact, Likelihood } from '@prisma/client'; + +import { + assessmentOutputSchema, + extractInherentRisk, +} from './assessment-output'; + +describe('assessmentOutputSchema', () => { + it('accepts independent likelihood and impact combinations across the full matrix', () => { + const offDiagonal = assessmentOutputSchema.safeParse({ + likelihood: 'very_likely', + impact: 'insignificant', + rationale: 'Motivated adversary but public-only data.', + }); + expect(offDiagonal.success).toBe(true); + if (offDiagonal.success) { + expect(offDiagonal.data.likelihood).toBe(Likelihood.very_likely); + expect(offDiagonal.data.impact).toBe(Impact.insignificant); + } + }); + + it('accepts every enum value on each dimension', () => { + const likelihoods: Likelihood[] = [ + Likelihood.very_unlikely, + Likelihood.unlikely, + Likelihood.possible, + Likelihood.likely, + Likelihood.very_likely, + ]; + const impacts: Impact[] = [ + Impact.insignificant, + Impact.minor, + Impact.moderate, + Impact.major, + Impact.severe, + ]; + for (const likelihood of likelihoods) { + for (const impact of impacts) { + const result = assessmentOutputSchema.safeParse({ + likelihood, + impact, + rationale: 'test rationale long enough.', + }); + expect(result.success).toBe(true); + } + } + }); + + it('rejects a missing dimension', () => { + const missingImpact = assessmentOutputSchema.safeParse({ + likelihood: 'possible', + rationale: 'missing impact on purpose.', + }); + expect(missingImpact.success).toBe(false); + }); + + it('rejects out-of-enum values', () => { + const result = assessmentOutputSchema.safeParse({ + likelihood: 'sometimes', + impact: 'catastrophic', + rationale: 'garbage in.', + }); + expect(result.success).toBe(false); + }); + + it('rejects too-short rationale', () => { + const result = assessmentOutputSchema.safeParse({ + likelihood: 'possible', + impact: 'moderate', + rationale: 'short', + }); + expect(result.success).toBe(false); + }); +}); + +describe('extractInherentRisk', () => { + it('threads likelihood + impact from a clean assessment payload', () => { + const payload = { + likelihood: 'very_likely', + impact: 'minor', + rationale: 'High adversary motivation, low blast radius due to scope.', + }; + const extracted = extractInherentRisk(payload); + expect(extracted).toEqual({ + likelihood: Likelihood.very_likely, + impact: Impact.minor, + }); + }); + + it('returns null when the payload is missing dimensions', () => { + const extracted = extractInherentRisk({ riskLevel: 'high' }); + expect(extracted).toBeNull(); + }); + + it('returns null when the payload is garbage', () => { + const extracted = extractInherentRisk({ foo: 'bar' }); + expect(extracted).toBeNull(); + }); +}); diff --git a/apps/api/src/trigger/vendor/vendor-risk-assessment/assessment-output.ts b/apps/api/src/trigger/vendor/vendor-risk-assessment/assessment-output.ts new file mode 100644 index 0000000000..da51a689c6 --- /dev/null +++ b/apps/api/src/trigger/vendor/vendor-risk-assessment/assessment-output.ts @@ -0,0 +1,39 @@ +// Import enums + types from `@prisma/client` directly (not `@db`) so this +// module has no runtime dependency on the PrismaClient singleton. That lets +// the spec import it without needing to mock `@db`. +import { Impact, Likelihood, type Prisma } from '@prisma/client'; +import { z } from 'zod'; + +/** + * Output of the vendor AI risk assessment — two independent dimensions. + * + * Replaces the previous single-level → diagonal-cell mapping, which pooled + * every vendor into 5 of 25 possible cells (5 of 10 score buckets). + * + * `likelihood` = probability of a security incident originating from or + * involving this vendor (adversary motivation + exposure + data handled). + * `impact` = blast radius if the vendor is compromised (regulatory + * exposure, customer-data sensitivity, operational criticality). + */ +export const assessmentOutputSchema = z.object({ + likelihood: z.nativeEnum(Likelihood), + impact: z.nativeEnum(Impact), + rationale: z.string().min(20), +}); + +export type AssessmentOutput = z.infer; + +/** + * Extract the inherent likelihood + impact from a vendor risk assessment payload. + * Returns null if the payload doesn't conform to `assessmentOutputSchema` + * (e.g. old-format payloads from globalVendors created before ENG-221). + */ +export function extractInherentRisk( + payload: Prisma.InputJsonValue, +): { likelihood: Likelihood; impact: Impact } | null { + const parsed = assessmentOutputSchema.safeParse(payload); + if (!parsed.success) { + return null; + } + return { likelihood: parsed.data.likelihood, impact: parsed.data.impact }; +} diff --git a/apps/api/src/trigger/vendor/vendor-risk-assessment/description.ts b/apps/api/src/trigger/vendor/vendor-risk-assessment/description.ts index d19f6aae86..906f9131d1 100644 --- a/apps/api/src/trigger/vendor/vendor-risk-assessment/description.ts +++ b/apps/api/src/trigger/vendor/vendor-risk-assessment/description.ts @@ -18,6 +18,9 @@ export function buildRiskAssessmentDescription(params: { vendorName, vendorWebsite, lastResearchedAt: null, + likelihood: null, + impact: null, + rationale: null, riskLevel: null, securityAssessment: null, certifications: null, diff --git a/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-core.ts b/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-core.ts index 93683039e3..c48d2e131b 100644 --- a/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-core.ts +++ b/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-core.ts @@ -207,7 +207,10 @@ export async function firecrawlResearchCore(params: { verifiedAgentCertCount: certifications.filter( (c) => c.status === 'verified', ).length, - agentRiskLevel: parsed.data.risk_level ?? null, + agentLikelihood: parsed.data.likelihood ?? null, + agentImpact: parsed.data.impact ?? null, + agentRationaleLength: parsed.data.rationale?.length ?? 0, + agentRiskLevelLegacy: parsed.data.risk_level ?? null, }); const deepScrapeSourceUrl = pickDeepScrapeSourceUrl({ @@ -287,6 +290,9 @@ export async function firecrawlResearchCore(params: { lastResearchedAt: normalizeIso(parsed.data.last_researched_at ?? null) ?? new Date().toISOString(), + likelihood: parsed.data.likelihood ?? null, + impact: parsed.data.impact ?? null, + rationale: parsed.data.rationale ?? null, riskLevel: parsed.data.risk_level ?? null, securityAssessment: parsed.data.security_assessment ?? null, certifications: diff --git a/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-prompt.ts b/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-prompt.ts index 4cd4059b50..947298b700 100644 --- a/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-prompt.ts +++ b/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-prompt.ts @@ -48,7 +48,9 @@ Only return a certification when the page explicitly names a framework as curren - links.privacy_policy_url, links.terms_of_service_url, links.security_page_url, links.soc2_report_url — return only when confirmed; otherwise empty. - certifications — may be an empty array. Do NOT pad it. - security_assessment — one paragraph summarising what you observed. If the trust portal was SPA-only and you could not read content, say so explicitly ("Trust portal at appears to be a JavaScript SPA; deep-scrape will extract content"). -- risk_level — your best estimate among critical/high/medium/low/very_low based on what you found. +- likelihood — probability of a security incident originating from or involving this vendor. Score INDEPENDENTLY of impact based on adversary motivation, exposure surface, and the kind of data handled. Choose exactly one of: very_unlikely, unlikely, possible, likely, very_likely. +- impact — blast radius if this vendor is compromised. Score INDEPENDENTLY of likelihood based on regulatory exposure, customer-data sensitivity, and operational criticality. Choose exactly one of: insignificant, minor, moderate, major, severe. +- rationale — two to four sentences justifying BOTH scores. Name the data handled and the adversary motivation. Do not merge the two dimensions into a single "risk" verdict. Focus on ${vendorWebsite} and its trust/security/compliance paths. Only cite URLs on ${vendorDomain}, its subdomains, or a recognised third-party portal hosting this vendor's trust page.`; } diff --git a/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-schema-json.ts b/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-schema-json.ts index 310ee19892..3efc452c34 100644 --- a/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-schema-json.ts +++ b/apps/api/src/trigger/vendor/vendor-risk-assessment/firecrawl-agent-schema-json.ts @@ -10,10 +10,28 @@ export const firecrawlAgentJsonSchema = { type: 'object', properties: { - risk_level: { + likelihood: { type: 'string', + enum: [ + 'very_unlikely', + 'unlikely', + 'possible', + 'likely', + 'very_likely', + ], description: - 'Overall vendor risk level: critical, high, medium, low, or very_low', + 'Probability of an incident originating from or involving this vendor. Score independently of impact.', + }, + impact: { + type: 'string', + enum: ['insignificant', 'minor', 'moderate', 'major', 'severe'], + description: + 'Blast radius if the vendor is compromised. Score independently of likelihood.', + }, + rationale: { + type: 'string', + description: + 'Two to four sentences justifying both scores. Name the data handled and the adversary motivation.', }, security_assessment: { type: 'string', @@ -91,5 +109,5 @@ export const firecrawlAgentJsonSchema = { }, }, }, - required: ['security_assessment'], + required: ['security_assessment', 'likelihood', 'impact', 'rationale'], } as const; diff --git a/apps/api/src/vendors/dto/update-vendor.dto.spec.ts b/apps/api/src/vendors/dto/update-vendor.dto.spec.ts index 53641b820f..b0f91c13c0 100644 --- a/apps/api/src/vendors/dto/update-vendor.dto.spec.ts +++ b/apps/api/src/vendors/dto/update-vendor.dto.spec.ts @@ -136,6 +136,56 @@ describe('UpdateVendorDto', () => { expect(errors[0].property).toBe('status'); }); + // ── treatment strategy fields (ENG-221) ─────────────────────────── + it('should accept valid treatmentStrategy enum values', async () => { + for (const strategy of ['accept', 'avoid', 'mitigate', 'transfer']) { + const dto = toDto({ treatmentStrategy: strategy }); + const errors = await validate(dto, { + whitelist: true, + forbidNonWhitelisted: true, + }); + expect(errors).toHaveLength(0); + } + }); + + it('should reject invalid treatmentStrategy enum value', async () => { + const dto = toDto({ treatmentStrategy: 'ignore' }); + const errors = await validate(dto, { + whitelist: true, + forbidNonWhitelisted: true, + }); + expect(errors.length).toBeGreaterThan(0); + expect(errors[0].property).toBe('treatmentStrategy'); + }); + + it('should accept treatmentStrategyDescription as a string', async () => { + const dto = toDto({ + treatmentStrategy: 'mitigate', + treatmentStrategyDescription: + 'We isolated the vendor to a dedicated VPC.', + }); + const errors = await validate(dto, { + whitelist: true, + forbidNonWhitelisted: true, + }); + expect(errors).toHaveLength(0); + expect(dto.treatmentStrategyDescription).toBe( + 'We isolated the vendor to a dedicated VPC.', + ); + }); + + it('should reject treatmentStrategyDescription longer than 20,000 chars', async () => { + const dto = toDto({ + treatmentStrategyDescription: 'x'.repeat(20_001), + }); + const errors = await validate(dto, { + whitelist: true, + forbidNonWhitelisted: true, + }); + expect(errors.length).toBeGreaterThan(0); + expect(errors[0].property).toBe('treatmentStrategyDescription'); + }); + // ── forbidNonWhitelisted ────────────────────────────────────────── it('should reject unknown properties', async () => { const dto = toDto({ name: 'Acronis', unknownField: 'value' }); diff --git a/apps/api/src/vendors/dto/update-vendor.dto.ts b/apps/api/src/vendors/dto/update-vendor.dto.ts index c7186b3f70..2c4f2efe63 100644 --- a/apps/api/src/vendors/dto/update-vendor.dto.ts +++ b/apps/api/src/vendors/dto/update-vendor.dto.ts @@ -6,9 +6,16 @@ import { IsEnum, IsUrl, IsBoolean, + MaxLength, } from 'class-validator'; import { Transform } from 'class-transformer'; -import { VendorCategory, VendorStatus, Likelihood, Impact } from '@db'; +import { + VendorCategory, + VendorStatus, + Likelihood, + Impact, + RiskTreatmentType, +} from '@db'; /** * DTO for PATCH /vendors/:id @@ -67,6 +74,25 @@ export class UpdateVendorDto { @IsEnum(Impact) residualImpact?: Impact; + @ApiPropertyOptional({ + description: 'Risk treatment strategy', + enum: RiskTreatmentType, + default: RiskTreatmentType.accept, + example: RiskTreatmentType.mitigate, + }) + @IsOptional() + @IsEnum(RiskTreatmentType) + treatmentStrategy?: RiskTreatmentType; + + @ApiPropertyOptional({ + description: 'Description of the treatment strategy', + example: 'We isolated the vendor to a dedicated VPC.', + }) + @IsOptional() + @IsString() + @MaxLength(20_000) + treatmentStrategyDescription?: string | null; + @ApiPropertyOptional({ description: 'Vendor website URL' }) @IsOptional() @IsUrl() diff --git a/apps/api/src/vendors/vendors.service.ts b/apps/api/src/vendors/vendors.service.ts index 47c49b0d17..349a2cff9a 100644 --- a/apps/api/src/vendors/vendors.service.ts +++ b/apps/api/src/vendors/vendors.service.ts @@ -11,6 +11,7 @@ import { tasks } from '@trigger.dev/sdk'; import { Prisma } from '@db'; import type { TriggerVendorRiskAssessmentVendorDto } from './dto/trigger-vendor-risk-assessment.dto'; import { resolveTaskCreatorAndAssignee } from '../trigger/vendor/vendor-risk-assessment/assignee'; +import { resolveStrategyDescriptionUpdate } from '../risks/strategy-descriptions'; const normalizeWebsite = ( website: string | null | undefined, @@ -106,6 +107,11 @@ export class VendorsService { }, }, }, + // Linked task statuses are needed by the vendors table to compute + // the current (interpolated) severity score so the residual badge + // reflects treatment progress, not just the static residual + // probability/impact. Mirrors the risks service. + tasks: { select: { id: true, status: true } }, }, }); @@ -142,6 +148,14 @@ export class VendorsService { }, }, }, + tasks: { + select: { + id: true, + title: true, + status: true, + controls: { select: { id: true, name: true } }, + }, + }, }, }); @@ -626,9 +640,16 @@ export class VendorsService { ); } + // Keep per-strategy descriptions independent across treatment + // strategies — see `apps/api/src/risks/strategy-descriptions.ts`. + const resolvedStrategyFields = resolveStrategyDescriptionUpdate( + existing, + updateVendorDto, + ); + const updatedVendor = await db.vendor.update({ where: { id }, - data: updateVendorDto, + data: { ...updateVendorDto, ...resolvedStrategyFields }, }); this.logger.log(`Updated vendor: ${updatedVendor.name} (${id})`); diff --git a/apps/api/trigger.config.ts b/apps/api/trigger.config.ts index 8c6e1b7700..44f4e88364 100644 --- a/apps/api/trigger.config.ts +++ b/apps/api/trigger.config.ts @@ -1,4 +1,5 @@ import { defineConfig } from '@trigger.dev/sdk'; +import { caBundleExtension } from './caBundleExtension'; import { prismaExtension } from './customPrismaExtension'; import { emailExtension } from './emailExtension'; import { integrationPlatformExtension } from './integrationPlatformExtension'; @@ -10,6 +11,7 @@ export default defineConfig({ maxDuration: 300, // 5 minutes build: { extensions: [ + caBundleExtension(), prismaExtension({ version: '7.6.0', dbPackageVersion: '^2.0.0', diff --git a/apps/app/caBundleExtension.ts b/apps/app/caBundleExtension.ts new file mode 100644 index 0000000000..0e602e0b15 --- /dev/null +++ b/apps/app/caBundleExtension.ts @@ -0,0 +1,50 @@ +import type { BuildContext, BuildExtension, BuildManifest } from '@trigger.dev/build'; +import { existsSync } from 'node:fs'; +import { cp, mkdir } from 'node:fs/promises'; +import { dirname, join, resolve } from 'node:path'; + +// Path relative to the monorepo root (apps/api or apps/app → ../../packages/db/certs/...) +const BUNDLE_RELATIVE_FROM_APP = '../../packages/db/certs/rds-global-bundle.pem'; +const BUNDLE_DEST_REL = 'certs/rds-global-bundle.pem'; + +function findBundleSrc(workingDir: string): string | undefined { + // Walk up from workingDir to find the cert — handles both normal checkouts and git worktrees + // where workspaceDir points to the main worktree root (wrong for us). + const candidates = [ + resolve(workingDir, BUNDLE_RELATIVE_FROM_APP), + resolve(workingDir, '../packages/db/certs/rds-global-bundle.pem'), + resolve(workingDir, 'packages/db/certs/rds-global-bundle.pem'), + ]; + + return candidates.find((c) => existsSync(c)); +} + +export function caBundleExtension(): BuildExtension { + return { + name: 'CABundleExtension', + onBuildStart: (context) => { + // Real OS env var at task spawn time — verified flow: + // addLayer.deploy.env → manifest.deploy.sync.env → syncEnvVarsWithServer → + // taskRunProcessProvider injects into worker env before Node TLS init. + context.addLayer({ + id: 'ca-bundle-env', + deploy: { + env: { NODE_EXTRA_CA_CERTS: `/app/${BUNDLE_DEST_REL}` }, + override: true, + }, + }); + }, + onBuildComplete: async (context: BuildContext, manifest: BuildManifest) => { + const src = findBundleSrc(context.workingDir); + if (!src) { + throw new Error( + `CABundleExtension: rds-global-bundle.pem not found. Searched relative to ${context.workingDir}`, + ); + } + const dest = join(manifest.outputPath, BUNDLE_DEST_REL); + await mkdir(dirname(dest), { recursive: true }); + await cp(src, dest); + context.logger.log(`Copied RDS CA bundle to ${BUNDLE_DEST_REL}`); + }, + }; +} diff --git a/apps/app/next.config.ts b/apps/app/next.config.ts index c83ff486f7..bcce976cff 100644 --- a/apps/app/next.config.ts +++ b/apps/app/next.config.ts @@ -76,6 +76,9 @@ const config: NextConfig = { webpackMemoryOptimizations: true, }, outputFileTracingRoot: workspaceRoot, + outputFileTracingIncludes: { + '/**/*': ['../../packages/db/certs/rds-global-bundle.pem'], + }, // Reduce memory usage during production build productionBrowserSourceMaps: false, diff --git a/apps/app/package.json b/apps/app/package.json index f39fce8fd3..ac2bc13498 100644 --- a/apps/app/package.json +++ b/apps/app/package.json @@ -79,6 +79,7 @@ "@uiw/react-json-view": "^2.0.0-alpha.40", "@uploadthing/react": "^7.3.0", "@upstash/ratelimit": "^2.0.5", + "@upstash/vector": "^1.2.2", "@vercel/analytics": "^1.5.0", "@vercel/sandbox": "^0.0.21", "@vercel/sdk": "^1.7.1", diff --git a/apps/app/prisma/client.ts b/apps/app/prisma/client.ts index 1bd8069f56..d7b33cfb60 100644 --- a/apps/app/prisma/client.ts +++ b/apps/app/prisma/client.ts @@ -1,7 +1,9 @@ import { PrismaClient } from '@prisma/client'; import { PrismaPg } from '@prisma/adapter-pg'; -const globalForPrisma = global as unknown as { prisma: PrismaClient }; +const globalForPrisma = global as unknown as { prisma?: PrismaClient }; + +const LOCAL_HOSTNAMES = new Set(['localhost', '127.0.0.1', '::1']); function stripSslMode(connectionString: string): string { const url = new URL(connectionString); @@ -9,14 +11,42 @@ function stripSslMode(connectionString: string): string { return url.toString(); } +function isLocalhostUrl(connectionString: string): boolean { + try { + const { hostname } = new URL(connectionString); + const stripped = hostname.replace(/^\[/, '').replace(/\]$/, ''); + return LOCAL_HOSTNAMES.has(stripped); + } catch { + return false; + } +} + function createPrismaClient(): PrismaClient { const rawUrl = process.env.DATABASE_URL!; - const isLocalhost = /localhost|127\.0\.0\.1|::1/.test(rawUrl); - // Use verified SSL when NODE_EXTRA_CA_CERTS is set (Docker with RDS CA bundle), - // otherwise fall back to unverified SSL (Trigger.dev, Vercel, other environments). + const isLocalhost = isLocalhostUrl(rawUrl); const hasCABundle = !!process.env.NODE_EXTRA_CA_CERTS; - const ssl = isLocalhost ? undefined : hasCABundle ? true : { rejectUnauthorized: false }; - // Strip sslmode from the connection string to avoid conflicts with the explicit ssl option + const allowInsecure = process.env.PRISMA_ALLOW_INSECURE_TLS === '1'; + + let ssl: + | undefined + | { checkServerIdentity: () => undefined } + | { rejectUnauthorized: false }; + if (isLocalhost) { + ssl = undefined; + } else if (hasCABundle) { + // Verified TLS: rely on Node's TLS context (NODE_EXTRA_CA_CERTS adds the AWS + // RDS CA to the trust store). Skip hostname check because connections may + // traverse an AWS NLB whose hostname isn't in the RDS Proxy cert's SAN list. + // The chain check still rejects forged or wrong-CA certs. + ssl = { checkServerIdentity: () => undefined }; + } else if (allowInsecure) { + ssl = { rejectUnauthorized: false }; + } else { + throw new Error( + 'Refusing to connect to a non-local Postgres without TLS verification. Set NODE_EXTRA_CA_CERTS to a CA bundle, or set PRISMA_ALLOW_INSECURE_TLS=1 if you intentionally want unverified TLS.', + ); + } + const url = ssl !== undefined ? stripSslMode(rawUrl) : rawUrl; const adapter = new PrismaPg({ connectionString: url, ssl }); return new PrismaClient({ @@ -27,6 +57,21 @@ function createPrismaClient(): PrismaClient { }); } -export const db = globalForPrisma.prisma || createPrismaClient(); +// Lazy initialization. Importing this module does NOT construct a Prisma client +// — that only happens on first property access on `db`. Critical so that +// Next.js `next build` (which imports every route handler to analyze it) does +// not trigger the strict TLS check at build time when no actual queries run. +function getClient(): PrismaClient { + if (!globalForPrisma.prisma) { + globalForPrisma.prisma = createPrismaClient(); + } + return globalForPrisma.prisma; +} -if (process.env.NODE_ENV !== 'production') globalForPrisma.prisma = db; +export const db = new Proxy({} as PrismaClient, { + get(_target, prop, _receiver) { + const client = getClient(); + const value = Reflect.get(client, prop, client); + return typeof value === 'function' ? value.bind(client) : value; + }, +}); diff --git a/apps/app/prisma/index.ts b/apps/app/prisma/index.ts index b329db54e3..f661a0b1d3 100644 --- a/apps/app/prisma/index.ts +++ b/apps/app/prisma/index.ts @@ -1 +1,65 @@ -export * from '@prisma/client'; +// `export *` from `@prisma/client` (CommonJS) makes Turbopack emit a warning +// on every compile because it can't statically resolve the export list. +// Listing the runtime values explicitly + `export type *` for types gives the +// same public surface with a clean compile. Keep enum names alphabetized so +// additions are obvious in diffs. +export type * from '@prisma/client'; +export { + Prisma, + PrismaClient, + AttachmentEntityType, + AttachmentType, + AuditLogEntityType, + BrowserAutomationEvaluationStatus, + BrowserAutomationRunStatus, + CommentEntityType, + Departments, + DevicePlatform, + EvidenceAutomationEvaluationStatus, + EvidenceAutomationRunStatus, + EvidenceAutomationTrigger, + EvidenceFormType, + FindingArea, + FindingSeverity, + FindingStatus, + FindingType, + FrameworkStatus, + FrameworkSyncOperationKind, + Frequency, + Impact, + IntegrationConnectionStatus, + IntegrationFindingSeverity, + IntegrationFindingStatus, + IntegrationRunJobType, + IntegrationRunStatus, + IntegrationSyncLogStatus, + KnowledgeBaseDocumentProcessingStatus, + Likelihood, + PhaseCompletionType, + PolicyDisplayFormat, + PolicyStatus, + PolicyVisibility, + QuestionnaireAnswerStatus, + QuestionnaireStatus, + RiskCategory, + RiskStatus, + RiskTreatmentType, + Role, + SOAAnswerStatus, + SOADocumentStatus, + TaskAutomationStatus, + TaskFrequency, + TaskItemEntityType, + TaskItemPriority, + TaskItemStatus, + TaskStatus, + TimelinePhaseStatus, + TimelineStatus, + TrustAccessGrantStatus, + TrustAccessRequestStatus, + TrustFramework, + TrustNDAStatus, + TrustStatus, + VendorCategory, + VendorStatus, +} from '@prisma/client'; diff --git a/apps/app/prisma/server.ts b/apps/app/prisma/server.ts index 54d1c4b9c9..b04e8bc13d 100644 --- a/apps/app/prisma/server.ts +++ b/apps/app/prisma/server.ts @@ -1,2 +1,65 @@ -export * from '@prisma/client'; +// `export *` from `@prisma/client` (CommonJS) makes Turbopack warn on every +// compile. Listing runtime values explicitly + `export type *` for types +// gives the same public surface with a clean compile. Keep enums in sync +// with `apps/app/prisma/index.ts`. +export type * from '@prisma/client'; +export { + Prisma, + PrismaClient, + AttachmentEntityType, + AttachmentType, + AuditLogEntityType, + BrowserAutomationEvaluationStatus, + BrowserAutomationRunStatus, + CommentEntityType, + Departments, + DevicePlatform, + EvidenceAutomationEvaluationStatus, + EvidenceAutomationRunStatus, + EvidenceAutomationTrigger, + EvidenceFormType, + FindingArea, + FindingSeverity, + FindingStatus, + FindingType, + FrameworkStatus, + FrameworkSyncOperationKind, + Frequency, + Impact, + IntegrationConnectionStatus, + IntegrationFindingSeverity, + IntegrationFindingStatus, + IntegrationRunJobType, + IntegrationRunStatus, + IntegrationSyncLogStatus, + KnowledgeBaseDocumentProcessingStatus, + Likelihood, + PhaseCompletionType, + PolicyDisplayFormat, + PolicyStatus, + PolicyVisibility, + QuestionnaireAnswerStatus, + QuestionnaireStatus, + RiskCategory, + RiskStatus, + RiskTreatmentType, + Role, + SOAAnswerStatus, + SOADocumentStatus, + TaskAutomationStatus, + TaskFrequency, + TaskItemEntityType, + TaskItemPriority, + TaskItemStatus, + TaskStatus, + TimelinePhaseStatus, + TimelineStatus, + TrustAccessGrantStatus, + TrustAccessRequestStatus, + TrustFramework, + TrustNDAStatus, + TrustStatus, + VendorCategory, + VendorStatus, +} from '@prisma/client'; export { db } from './client'; diff --git a/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/backgroundCheckTypes.ts b/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/backgroundCheckTypes.ts index 44347d001f..448129631b 100644 --- a/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/backgroundCheckTypes.ts +++ b/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/backgroundCheckTypes.ts @@ -43,6 +43,13 @@ export interface BackgroundCheckBillingStatus { currentPeriodEnd: string | null; cancelAtPeriodEnd: boolean; }>; + // Wallet credits per product. Optional because older API responses + // and existing test fixtures don't set it; treated as zero balance + // when absent. + creditBalances?: Array<{ + productKey: 'pentest' | 'background_check'; + balance: number; + }>; } export function isCompletedBackgroundCheck(status: BackgroundCheckStatus): boolean { diff --git a/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/useEmployeeBackgroundCheckData.test.ts b/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/useEmployeeBackgroundCheckData.test.ts new file mode 100644 index 0000000000..b42b63ac34 --- /dev/null +++ b/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/useEmployeeBackgroundCheckData.test.ts @@ -0,0 +1,108 @@ +import { describe, expect, it } from 'vitest'; +import type { BackgroundCheckBillingStatus } from './backgroundCheckTypes'; +import { getBackgroundChecksRemaining } from './useEmployeeBackgroundCheckData'; + +function status( + overrides: Partial = {}, +): BackgroundCheckBillingStatus { + return { + hasPaymentMethod: false, + setupAt: null, + ...overrides, + }; +} + +describe('getBackgroundChecksRemaining', () => { + it('returns null when billing status has not loaded', () => { + expect(getBackgroundChecksRemaining({ billingStatus: undefined })).toBeNull(); + }); + + it('returns null when no subscription and no wallet credits exist', () => { + expect(getBackgroundChecksRemaining({ billingStatus: status() })).toBeNull(); + }); + + it('returns subscription remainder when an active subscription exists', () => { + expect( + getBackgroundChecksRemaining({ + billingStatus: status({ + subscriptions: [ + { + skuKey: 'background_checks_monthly_3', + status: 'active', + includedQuantity: 3, + usedQuantity: 1, + currentPeriodStart: null, + currentPeriodEnd: null, + cancelAtPeriodEnd: false, + }, + ], + }), + }), + ).toBe(2); + }); + + it('returns wallet balance when no subscription but wallet credits exist', () => { + // Admin grant flow: organization has no subscription, platform + // admin granted 5 BG-check credits. Backend would consume from + // wallet on POST, so the wizard must allow the request. + expect( + getBackgroundChecksRemaining({ + billingStatus: status({ + creditBalances: [{ productKey: 'background_check', balance: 5 }], + }), + }), + ).toBe(5); + }); + + it('sums subscription remainder and wallet credits', () => { + expect( + getBackgroundChecksRemaining({ + billingStatus: status({ + subscriptions: [ + { + skuKey: 'background_checks_monthly_3', + status: 'trialing', + includedQuantity: 3, + usedQuantity: 3, + currentPeriodStart: null, + currentPeriodEnd: null, + cancelAtPeriodEnd: false, + }, + ], + creditBalances: [{ productKey: 'background_check', balance: 5 }], + }), + }), + ).toBe(5); + }); + + it('ignores pentest wallet credits when computing background-check allowance', () => { + expect( + getBackgroundChecksRemaining({ + billingStatus: status({ + creditBalances: [{ productKey: 'pentest', balance: 10 }], + }), + }), + ).toBeNull(); + }); + + it('ignores cancelled subscriptions', () => { + expect( + getBackgroundChecksRemaining({ + billingStatus: status({ + subscriptions: [ + { + skuKey: 'background_checks_monthly_3', + status: 'canceled', + includedQuantity: 3, + usedQuantity: 0, + currentPeriodStart: null, + currentPeriodEnd: null, + cancelAtPeriodEnd: false, + }, + ], + creditBalances: [{ productKey: 'background_check', balance: 2 }], + }), + }), + ).toBe(2); + }); +}); diff --git a/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/useEmployeeBackgroundCheckData.ts b/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/useEmployeeBackgroundCheckData.ts index 348047cdfe..c55edecb6a 100644 --- a/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/useEmployeeBackgroundCheckData.ts +++ b/apps/app/src/app/(app)/[orgId]/people/[employeeId]/components/useEmployeeBackgroundCheckData.ts @@ -54,11 +54,31 @@ export function getBackgroundChecksRemaining({ }: { billingStatus: BackgroundCheckBillingStatus | undefined; }): number | null { - const subscription = (billingStatus?.subscriptions ?? []).find( + if (!billingStatus) return null; + const subscription = (billingStatus.subscriptions ?? []).find( (item) => getBillingSkuProductKey(item.skuKey) === 'background_check' && (item.status === 'active' || item.status === 'trialing'), ); - if (!subscription) return null; - return Math.max(subscription.includedQuantity - subscription.usedQuantity, 0); + // Wallet credits granted by platform admins. The backend's + // `BillingEntitlementsService.tryConsumeIncludedUsageForProduct` + // falls back to this wallet when no active subscription exists or + // the subscription's included usage is exhausted, so we mirror that + // logic here — otherwise the wizard paywalls users whose admin- + // granted credits would actually be consumed by the create call. + const walletBalance = + (billingStatus.creditBalances ?? []).find( + (entry) => entry.productKey === 'background_check', + )?.balance ?? 0; + if (!subscription) { + // Returning `null` keeps the existing "no allowance — go pick a + // plan" wizard path. We only have a positive allowance if there + // are wallet credits to consume. + return walletBalance > 0 ? walletBalance : null; + } + const subscriptionRemaining = Math.max( + subscription.includedQuantity - subscription.usedQuantity, + 0, + ); + return subscriptionRemaining + walletBalance; } diff --git a/apps/app/src/app/(app)/[orgId]/risk/(overview)/RisksTable.test.tsx b/apps/app/src/app/(app)/[orgId]/risk/(overview)/RisksTable.test.tsx index 22c95d34d2..c8be570601 100644 --- a/apps/app/src/app/(app)/[orgId]/risk/(overview)/RisksTable.test.tsx +++ b/apps/app/src/app/(app)/[orgId]/risk/(overview)/RisksTable.test.tsx @@ -70,6 +70,34 @@ vi.mock('@/hooks/use-risks', () => ({ // Mock @db vi.mock('@db', () => ({ Risk: {}, + // Enums used by suggested-residual / risk-score helpers. Values mirror + // the real Prisma enum string values. + Likelihood: { + very_unlikely: 'very_unlikely', + unlikely: 'unlikely', + possible: 'possible', + likely: 'likely', + very_likely: 'very_likely', + }, + Impact: { + insignificant: 'insignificant', + minor: 'minor', + moderate: 'moderate', + major: 'major', + severe: 'severe', + }, + RiskTreatmentType: { + accept: 'accept', + avoid: 'avoid', + mitigate: 'mitigate', + transfer: 'transfer', + }, + TaskStatus: { + todo: 'todo', + in_progress: 'in_progress', + done: 'done', + not_relevant: 'not_relevant', + }, })); // Mock onboarding hooks @@ -116,6 +144,11 @@ vi.mock('@trycompai/design-system', () => ({ InputGroup: ({ children }: any) =>
{children}
, InputGroupAddon: ({ children }: any) => {children}, InputGroupInput: (props: any) => , + Select: ({ children }: any) =>
{children}
, + SelectContent: ({ children }: any) =>
{children}
, + SelectItem: ({ children, value }: any) => , + SelectTrigger: ({ children }: any) =>
{children}
, + SelectValue: ({ placeholder }: any) => {placeholder}, Spinner: () => , Stack: ({ children }: any) =>
{children}
, Table: ({ children, pagination }: any) => {children}
, @@ -202,11 +235,27 @@ describe('RisksTable permission gating', () => { expect(screen.getByText('RISK')).toBeInTheDocument(); expect(screen.getByText('SEVERITY')).toBeInTheDocument(); + expect(screen.getByText('RISK SCORE')).toBeInTheDocument(); expect(screen.getByText('STATUS')).toBeInTheDocument(); expect(screen.getByText('OWNER')).toBeInTheDocument(); expect(screen.getByText('UPDATED')).toBeInTheDocument(); }); + it('renders SEVERITY label + RISK SCORE number, both from current state', () => { + setMockPermissions({}); + + render(); + + // Fixture: possible × moderate, mitigate, no linked tasks. + // inherent: 3 × 3 = 9 raw → ceil(9/2.5) = 4 + // coverage gate (no tasks) → target = inherent → current = 4 + // → severity label "Low" (score 4 → low band) + numeric "4/10". + // The severity filter dropdown also contains "Low" as an option, so + // we expect at least one (row + dropdown option). + expect(screen.getAllByText('Low').length).toBeGreaterThanOrEqual(1); + expect(screen.getByText('4/10')).toBeInTheDocument(); + }); + it('renders search bar regardless of permissions', () => { setMockPermissions({}); diff --git a/apps/app/src/app/(app)/[orgId]/risk/(overview)/RisksTable.tsx b/apps/app/src/app/(app)/[orgId]/risk/(overview)/RisksTable.tsx index 31f3c5f948..b748b39dc6 100644 --- a/apps/app/src/app/(app)/[orgId]/risk/(overview)/RisksTable.tsx +++ b/apps/app/src/app/(app)/[orgId]/risk/(overview)/RisksTable.tsx @@ -1,5 +1,6 @@ 'use client'; +import { RiskScoreBadge } from '@/components/risks/RiskScoreBadge'; import { usePermissions } from '@/hooks/use-permissions'; import { useRiskActions, @@ -9,6 +10,13 @@ import { type RisksQueryParams, } from '@/hooks/use-risks'; import { getSortingStateParser } from '@/lib/parsers'; +import { getRiskLevelFromScore, getRiskScore, LEVEL_LABEL } from '@/lib/risk-score'; +import { + interpolatedResidualScore, + previewResidual, + suggestedResidual, +} from '@/lib/suggested-residual'; +import { TaskStatus } from '@db'; import type { Member, User } from '@db'; import { Risk as RiskType } from '@db'; import { @@ -34,6 +42,11 @@ import { InputGroup, InputGroupAddon, InputGroupInput, + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, Spinner, Stack, Table, @@ -66,33 +79,47 @@ const ACTIVE_STATUSES: Array<'pending' | 'processing' | 'created' | 'assessing'> 'assessing', ]; -function getSeverityBadge(likelihood: string, impact: string) { - // Calculate severity based on likelihood and impact - const likelihoodScore: Record = { - very_unlikely: 1, - unlikely: 2, - possible: 3, - likely: 4, - very_likely: 5, - }; - const impactScore: Record = { - insignificant: 1, - minor: 2, - moderate: 3, - major: 4, - severe: 5, - }; +/** + * The risk's current severity score (1-10), interpolated by linked-task + * completion the same way the Treatment Plan hero does it. Returns the + * inherent score when there's no linked work or the strategy doesn't + * project a reduction. Falls back to inherent on malformed input. + */ +function currentSeverityScore(risk: { + likelihood: ApiRisk['likelihood']; + impact: ApiRisk['impact']; + treatmentStrategy: ApiRisk['treatmentStrategy']; + tasks?: Array<{ status: TaskStatus }>; +}): number { + const inherent = getRiskScore(risk.likelihood, risk.impact); + const tasks = risk.tasks ?? []; + const target = previewResidual({ + inherentLikelihood: risk.likelihood, + inherentImpact: risk.impact, + strategy: risk.treatmentStrategy, + hasLinkedWork: tasks.length > 0, + }); + const targetScore = getRiskScore(target.likelihood, target.impact).score; + const completion = suggestedResidual({ + likelihood: risk.likelihood, + impact: risk.impact, + strategy: risk.treatmentStrategy, + tasks, + }).completion; + return interpolatedResidualScore({ + inherentScore: inherent.score, + targetScore, + completion, + }); +} - const score = (likelihoodScore[likelihood] || 1) * (impactScore[impact] || 1); - if (score >= 15) { - return High; - } - if (score >= 8) { - return Medium; - } - return Low; -} +const STATUS_LABEL: Record = { + open: 'Open', + pending: 'Pending', + closed: 'Closed', + archived: 'Archived', +}; function getStatusBadge(status: string) { switch (status) { @@ -146,24 +173,48 @@ export const RisksTable = ({ // Read current search params from URL const [title, setTitle] = useQueryState('title', parseAsString.withDefault('')); + const [statusFilter, setStatusFilter] = useQueryState( + 'status', + parseAsString.withDefault(''), + ); + const [assigneeFilter, setAssigneeFilter] = useQueryState( + 'assignee', + parseAsString.withDefault(''), + ); + // Severity is computed from the current treatment-aware score, so it's + // filtered client-side after fetch (the API can't query a derived value). + const [severityFilter, setSeverityFilter] = useQueryState( + 'severity', + parseAsString.withDefault(''), + ); const [sort, setSort] = useQueryState( 'sort', getSortingStateParser().withDefault([{ id: 'title', desc: false }]), ); - // Build query params for the API + // Build query params for the API. Status and assignee are server-side; + // severity is computed from the current treatment-aware score (which the + // API can't query directly), so when severity is filtered we fetch + // the org's full risk set in one page and paginate after the filter. + // Without this, severity filtering would only see whatever happened to + // be on the current server page — causing wrong/missing results and an + // incorrect total count. (Cubic finding #27 on PR #2671.) + const FILTER_ALL_PAGE_SIZE = 1000; + const fetchAllForSeverity = Boolean(severityFilter); const queryParams = useMemo(() => { const currentSort = sort[0]; return { - page, - perPage, + page: fetchAllForSeverity ? 1 : page, + perPage: fetchAllForSeverity ? FILTER_ALL_PAGE_SIZE : perPage, ...(title && { title }), + ...(statusFilter && { status: statusFilter }), + ...(assigneeFilter && { assigneeId: assigneeFilter }), ...(currentSort && { sort: currentSort.id, sortDirection: currentSort.desc ? 'desc' as const : 'asc' as const, }), }; - }, [page, perPage, title, sort]); + }, [page, perPage, title, statusFilter, assigneeFilter, sort, fetchAllForSeverity]); // Use the useRisks hook with query params const { data: risksData, mutate: mutateRisks } = useRisks({ @@ -173,12 +224,34 @@ export const RisksTable = ({ keepPreviousData: true, }); - const risks = useMemo(() => { + // Full result set from the API — already narrowed by status/assignee/title. + const fullList = useMemo(() => { const apiData = risksData?.data?.data; return Array.isArray(apiData) ? apiData : initialRisks; }, [risksData, initialRisks]); - const pageCount = risksData?.data?.pageCount ?? initialPageCount; + // After server filtering, apply the (derived) severity filter and then + // re-paginate client-side so totals reflect the filtered result. + const filteredList = useMemo(() => { + if (!severityFilter) return fullList; + return fullList.filter((risk) => { + const score = currentSeverityScore(risk); + return getRiskLevelFromScore(score) === severityFilter; + }); + }, [fullList, severityFilter]); + + const risks = useMemo(() => { + if (!fetchAllForSeverity) return filteredList; + const start = (page - 1) * perPage; + return filteredList.slice(start, start + perPage); + }, [filteredList, fetchAllForSeverity, page, perPage]); + + const pageCount = useMemo(() => { + if (fetchAllForSeverity) { + return Math.max(1, Math.ceil(filteredList.length / perPage)); + } + return risksData?.data?.pageCount ?? initialPageCount; + }, [fetchAllForSeverity, filteredList.length, perPage, risksData, initialPageCount]); // Check if all risks are done assessing const allRisksDoneAssessing = useMemo(() => { @@ -380,18 +453,107 @@ export const RisksTable = ({ return ( - {/* Search Bar */} -
- - - - - setTitle(e.target.value || null)} - /> - + {/* Search + Filters. Severity is client-side (derived score), + Status and Owner are server-side via the risks API. Each + filter is URL-backed so links are shareable. */} +
+
+ + + + + setTitle(e.target.value || null)} + /> + +
+
+ +
+
+ +
+
+ +
+ {(severityFilter || statusFilter || assigneeFilter || title) && ( + + )}
{/* Onboarding Progress Banner */} @@ -463,6 +625,8 @@ export const RisksTable = ({ SEVERITY + INHERENT RISK + CURRENT RISK STATUS OWNER @@ -494,7 +658,36 @@ export const RisksTable = ({ {risk.title} - {getSeverityBadge(risk.likelihood, risk.impact)} + {(() => { + // Three score columns paint the before-vs-now picture: + // SEVERITY = current treatment-aware level (text). + // INHERENT = raw score before treatment, fixed. + // CURRENT = treatment-aware score interpolated by + // linked-task completion. Named "Current" + // (not "Residual") because the canonical + // residual is the *target* score at 100% + // completion — what's shown here moves + // with progress and matches the hero's + // "Currently X/10" subline. + // SEVERITY is plain text and CURRENT carries the + // colored chip so we don't double-paint the band. + const inherentScore = getRiskScore(risk.likelihood, risk.impact).score; + const score = currentSeverityScore(risk); + const level = getRiskLevelFromScore(score); + return ( + <> + + {LEVEL_LABEL[level]} + + + + + + + + + ); + })()} {getStatusBadge(risk.status)} {risk.assignee?.user?.name || 'Unassigned'} diff --git a/apps/app/src/app/(app)/[orgId]/risk/(overview)/page.tsx b/apps/app/src/app/(app)/[orgId]/risk/(overview)/page.tsx index 55e499a40c..61a5231678 100644 --- a/apps/app/src/app/(app)/[orgId]/risk/(overview)/page.tsx +++ b/apps/app/src/app/(app)/[orgId]/risk/(overview)/page.tsx @@ -69,9 +69,30 @@ export default async function RiskRegisterPage(props: { const risks = risksResult.data?.data ?? []; const pageCount = risksResult.data?.pageCount ?? 0; - // Transform people response to assignees format expected by CreateRiskSheet + // Transform people response to assignees format expected by CreateRiskSheet. + // Risk owners need write access to risks (risk:update) and the compliance + // app (app:read). The built-in roles that grant both are `owner` and + // `admin`; `auditor` is read-only (no update); `employee` and `contractor` + // are portal-only. Custom roles can grant any combination of permissions, + // so anyone with a non-built-in role is included — we can't resolve + // permissions client-side and customers control their own role design. + const PORTAL_ONLY_OR_READ_ROLES = new Set([ + 'auditor', + 'employee', + 'contractor', + ]); + const canOwnRisks = (roleField: string): boolean => { + const roles = roleField.split(',').map((r) => r.trim()).filter(Boolean); + if (roles.length === 0) return false; + return roles.some((r) => { + if (r === 'owner' || r === 'admin') return true; + if (PORTAL_ONLY_OR_READ_ROLES.has(r)) return false; + // Custom role — trust the org's role design. + return true; + }); + }; const assignees = (peopleResult.data?.data ?? []) - .filter((p) => !p.deactivated && !['employee', 'contractor'].includes(p.role)) + .filter((p) => !p.deactivated && canOwnRisks(p.role)) .map((p) => ({ id: p.id, role: p.role, diff --git a/apps/app/src/app/(app)/[orgId]/risk/[riskId]/components/RiskPageClient.test.tsx b/apps/app/src/app/(app)/[orgId]/risk/[riskId]/components/RiskPageClient.test.tsx index fbdb7a5f6d..70ab66f3a4 100644 --- a/apps/app/src/app/(app)/[orgId]/risk/[riskId]/components/RiskPageClient.test.tsx +++ b/apps/app/src/app/(app)/[orgId]/risk/[riskId]/components/RiskPageClient.test.tsx @@ -40,11 +40,48 @@ vi.mock('@/hooks/use-risks', () => ({ useRisk: () => ({ risk: mockRisk, }), + useRiskActions: () => ({ + updateRisk: vi.fn(), + regenerateMitigation: vi.fn().mockResolvedValue({ + runId: 'run_test', + publicAccessToken: 'tok_test', + }), + suggestRiskLinks: vi.fn(), + applyRiskLinks: vi.fn(), + fetchActiveRiskAutoLinkRun: vi.fn().mockResolvedValue(null), + discardRiskAutoLinkRun: vi.fn(), + }), })); // Mock @db vi.mock('@db', () => ({ CommentEntityType: { risk: 'risk' }, + Likelihood: { + very_unlikely: 'very_unlikely', + unlikely: 'unlikely', + possible: 'possible', + likely: 'likely', + very_likely: 'very_likely', + }, + Impact: { + insignificant: 'insignificant', + minor: 'minor', + moderate: 'moderate', + major: 'major', + severe: 'severe', + }, + RiskTreatmentType: { + accept: 'accept', + avoid: 'avoid', + mitigate: 'mitigate', + transfer: 'transfer', + }, + TaskStatus: { + todo: 'todo', + in_progress: 'in_progress', + done: 'done', + not_relevant: 'not_relevant', + }, })); // Mock design system diff --git a/apps/app/src/app/(app)/[orgId]/risk/[riskId]/components/RiskPageClient.tsx b/apps/app/src/app/(app)/[orgId]/risk/[riskId]/components/RiskPageClient.tsx index a5858b4cec..e58fa9b684 100644 --- a/apps/app/src/app/(app)/[orgId]/risk/[riskId]/components/RiskPageClient.tsx +++ b/apps/app/src/app/(app)/[orgId]/risk/[riskId]/components/RiskPageClient.tsx @@ -5,16 +5,16 @@ import { RecentAuditLogs } from '@/components/RecentAuditLogs'; import { InherentRiskChart } from '@/components/risks/charts/InherentRiskChart'; import { ResidualRiskChart } from '@/components/risks/charts/ResidualRiskChart'; import { RiskOverview } from '@/components/risks/risk-overview'; +import { TreatmentPlanTab } from '@/components/risks/treatment-plan/TreatmentPlanTab'; import { TaskItems } from '@/components/task-items/TaskItems'; import { useAuditLogs } from '@/hooks/use-audit-logs'; -import { useRisk, useRiskActions, type RiskResponse } from '@/hooks/use-risks'; +import { useRisk, useRiskActions, type RiskLinkedTask, type RiskResponse } from '@/hooks/use-risks'; import { useTaskItems, useTaskItemActions } from '@/hooks/use-task-items'; import { usePermissions } from '@/hooks/use-permissions'; import { CommentEntityType } from '@db'; -import type { Member, Risk, User } from '@db'; +import type { Member, Risk, RiskTreatmentType, User } from '@db'; import { Breadcrumb, - Button, HStack, Stack, Tabs, @@ -24,8 +24,8 @@ import { Text, } from '@trycompai/design-system'; import Link from 'next/link'; -import { useSearchParams } from 'next/navigation'; -import { useMemo, useState } from 'react'; +import { useQueryState } from 'nuqs'; +import { useCallback, useMemo, useState } from 'react'; import { toast } from 'sonner'; type RiskWithAssignee = Risk & { @@ -62,10 +62,25 @@ export function RiskPageClient({ taskItemId, }: RiskPageClientProps) { const { risk: swrRisk, mutate: mutateRisk } = useRisk(riskId); - const { updateRisk, regenerateMitigation } = useRiskActions(); + const { + updateRisk, + regenerateMitigation, + suggestRiskLinks, + applyRiskLinks, + fetchActiveRiskAutoLinkRun, + discardRiskAutoLinkRun, + } = useRiskActions(); const { hasPermission } = usePermissions(); - const searchParams = useSearchParams(); - const defaultTab = searchParams.get('tab') || 'overview'; + // URL-backed tab state — bookmarks the active tab in `?tab=...` so a + // refresh keeps the same view. Control the value so we can conditionally + // render only the active panel below: base-ui's Tabs.Panel keeps the + // outgoing panel mounted at full opacity for the duration of the + // incoming panel's `fade-in-0 duration-200` animation, which produces + // a visible "both panels stacked" flash. Mounting only the active panel + // sidesteps the transition window entirely. + const [activeTab, setActiveTab] = useQueryState('tab', { + defaultValue: 'overview', + }); const isViewingTask = Boolean(taskItemId); const canUpdate = hasPermission('risk', 'update'); const canUpdateTask = hasPermission('task', 'update'); @@ -82,6 +97,10 @@ export function RiskPageClient({ const [isEditingDescription, setIsEditingDescription] = useState(false); const [descriptionValue, setDescriptionValue] = useState(''); const [isRegenerating, setIsRegenerating] = useState(false); + const [regenRun, setRegenRun] = useState<{ + runId: string; + publicAccessToken: string; + } | null>(null); const risk = useMemo(() => { if (swrRisk) return normalizeRisk(swrRisk); @@ -141,20 +160,80 @@ export function RiskPageClient({ } }; + const handleUpdateStrategy = async (strategy: RiskTreatmentType) => { + await updateRisk(riskId, { treatmentStrategy: strategy }); + mutateRisk(); + }; + + const handleUpdateDescription = async (description: string) => { + await updateRisk(riskId, { treatmentStrategyDescription: description }); + mutateRisk(); + }; + const handleRegenerateMitigation = async () => { setIsRegenerating(true); - toast.info('Regenerating risk mitigation...'); try { - await regenerateMitigation(riskId); - toast.success('Regeneration triggered. This may take a moment.'); - mutateRisk(); + const handle = await regenerateMitigation(riskId); + // Hand control to the realtime subscription in DescriptionEditor. + // It calls back via onRegenSettled when the run terminates, at which + // point we clear isRegenerating + refetch the risk for the new prose. + setRegenRun(handle); } catch { toast.error('Failed to trigger mitigation regeneration'); - } finally { setIsRegenerating(false); } }; + const handleRegenSettled = useCallback( + (result: { success: boolean; reason?: string }) => { + setRegenRun(null); + setIsRegenerating(false); + if (result.success) { + toast.success('Treatment plan regenerated.'); + void mutateRisk(); + } else { + toast.error(result.reason ?? 'Failed to regenerate the treatment plan.'); + } + }, + [mutateRisk], + ); + + const handleSuggest = async () => { + return suggestRiskLinks(riskId); + }; + + const handleApply = async (params: { taskIds: string[]; replace: boolean }) => { + await applyRiskLinks(riskId, params); + await mutateRisk(); + }; + + const handleUnlinkTask = async (taskId: string) => { + try { + const res = await fetch(`/api/risks/${riskId}/tasks/${taskId}`, { + method: 'DELETE', + credentials: 'include', + }); + if (!res.ok) throw new Error('unlink failed'); + await mutateRisk(); + } catch { + toast.error('Failed to unlink task'); + } + }; + + // Memoize so AutoLinkSuggestions' resume effect (which depends on the + // callback identity) doesn't re-fire on every parent re-render. + const handleResumeAutoLink = useCallback( + () => fetchActiveRiskAutoLinkRun(riskId), + [fetchActiveRiskAutoLinkRun, riskId], + ); + + const handleDiscardAutoLinkRun = useCallback( + async () => { + await discardRiskAutoLinkRun(riskId); + }, + [discardRiskAutoLinkRun, riskId], + ); + return ( <> ) : ( - + void setActiveTab(String(next))}> Overview + Treatment Plan Risk Matrix Tasks Comments @@ -247,52 +327,91 @@ export function RiskPageClient({ Settings - - - + {activeTab === 'overview' && ( + + + + )} + + {activeTab === 'treatment-plan' && ( + + > + | null + | undefined ?? null, + // Fall through to the server-rendered initial risk so the + // Linked Work column doesn't blink empty between SSR and + // the first SWR resolution. (Cubic finding #28.) + tasks: + swrRisk?.tasks ?? + (initialRisk as unknown as { tasks?: RiskLinkedTask[] }) + .tasks ?? + [], + }} + canUpdate={canUpdate} + onUpdateStrategy={handleUpdateStrategy} + onUpdateDescription={handleUpdateDescription} + onRegenerate={handleRegenerateMitigation} + regenerating={isRegenerating} + onSuggest={handleSuggest} + onApply={handleApply} + // Gate the unlink affordance behind risk:update so the trash + // button doesn't render for read-only users. The Next API + // route enforces the same check server-side as defense in + // depth. (Cubic finding on PR #2671.) + onUnlinkTask={canUpdate ? handleUnlinkTask : undefined} + onResumeAutoLink={handleResumeAutoLink} + onDiscardAutoLinkRun={handleDiscardAutoLinkRun} + regenRun={regenRun} + onRegenSettled={handleRegenSettled} + /> + + )} - - - - - - + {activeTab === 'risk-matrix' && ( + + + + + + + )} - - - + {activeTab === 'tasks' && ( + + + + )} - - - + {activeTab === 'comments' && ( + + + + )} - - t.id) || []} /> - + {activeTab === 'activity' && ( + + t.id) || []} /> + + )} - - - {canUpdate && ( - - - Regenerate Risk Mitigation - - Generate a fresh mitigation comment for this risk using AI - - - - - )} - - + {activeTab === 'settings' && ( + + No settings yet. + + )} )} diff --git a/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/SplitView.tsx b/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/SplitView.tsx index 79147c1014..0846203fff 100644 --- a/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/SplitView.tsx +++ b/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/SplitView.tsx @@ -6,13 +6,16 @@ import type { PentestIssue, PentestRun, } from '@/lib/security/penetration-tests-client'; -import { getBillingSkuProductKey } from '@trycompai/billing'; import { cn } from '@trycompai/design-system/cn'; import { ArrowLeft } from '@trycompai/design-system/icons'; import { useRouter } from 'next/navigation'; import { useState } from 'react'; import { toast } from 'sonner'; import useSWR from 'swr'; +import { + getPentestAllowance, + type PentestBillingStatusInput, +} from './pentest-allowance'; import { useCreatePenetrationTest, usePenetrationTest, @@ -27,15 +30,6 @@ import { OverviewPane } from './OverviewPane'; import { RunList } from './RunList'; import './pentest-tokens.css'; -interface BillingStatus { - subscriptions?: Array<{ - skuKey: string; - status: string; - includedQuantity: number; - usedQuantity: number; - }>; -} - interface SplitViewProps { orgId: string; selectedRunId: string | null; @@ -61,28 +55,20 @@ export function SplitView({ orgId, selectedRunId, mode = 'default' }: SplitViewP const { issues } = usePenetrationTestIssues(orgId, selectedRunId ?? '', selectedRun?.status); const { events } = usePenetrationTestEvents(orgId, selectedRunId ?? '', selectedRun?.status); const { createReport, isCreating } = useCreatePenetrationTest(orgId); - const { data: billingStatus } = useSWR( + const { data: billingStatus } = useSWR( orgId ? (['/v1/billing/status', orgId] as const) : null, async ([endpoint, organizationId]: readonly [string, string]) => { - const response = await api.get(endpoint, organizationId); + const response = await api.get( + endpoint, + organizationId, + ); if (response.status < 200 || response.status >= 300) { throw new Error(response.error ?? 'Failed to load billing status'); } return response.data ?? {}; }, ); - const pentestSubscription = (billingStatus?.subscriptions ?? []).find( - (subscription) => - getBillingSkuProductKey(subscription.skuKey) === 'pentest' && - (subscription.status === 'active' || subscription.status === 'trialing'), - ); - const subscriptionBalance = pentestSubscription - ? Math.max(pentestSubscription.includedQuantity - pentestSubscription.usedQuantity, 0) - : null; - // Keep `balance` undefined while billing is loading so the page does not - // flash a blocked state before subscription allowance is known. - const balance = subscriptionBalance ?? (billingStatus === undefined ? undefined : 0); - const planRequired = subscriptionBalance === null && billingStatus !== undefined; + const { balance, planRequired } = getPentestAllowance(billingStatus); const quotaLabel = 'Plan'; const showEmptyState = diff --git a/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/pentest-allowance.test.ts b/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/pentest-allowance.test.ts new file mode 100644 index 0000000000..22e1be6e2c --- /dev/null +++ b/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/pentest-allowance.test.ts @@ -0,0 +1,107 @@ +import { describe, expect, it } from 'vitest'; +import { getPentestAllowance } from './pentest-allowance'; + +describe('getPentestAllowance', () => { + it('returns undefined balance while billing status is loading', () => { + expect(getPentestAllowance(undefined)).toEqual({ + balance: undefined, + planRequired: false, + }); + }); + + it('reports plan required when neither subscription nor wallet credits exist', () => { + expect(getPentestAllowance({})).toEqual({ + balance: 0, + planRequired: true, + }); + }); + + it('uses subscription remainder when a trial is active', () => { + expect( + getPentestAllowance({ + subscriptions: [ + { + skuKey: 'pentest_monthly_1', + status: 'trialing', + includedQuantity: 1, + usedQuantity: 0, + }, + ], + }), + ).toEqual({ balance: 1, planRequired: false }); + }); + + it('uses wallet credits when no subscription exists', () => { + // The "platform admin granted credits to a paywalled org" case — + // backend falls back to wallet, so the UI must too. + expect( + getPentestAllowance({ + creditBalances: [{ productKey: 'pentest', balance: 5 }], + }), + ).toEqual({ balance: 5, planRequired: false }); + }); + + it('sums subscription remainder and wallet credits', () => { + // Trial used (0 remaining) + 5 wallet credits => 5 total. Without + // this fix the UI would compute balance=0 and reroute to billing. + expect( + getPentestAllowance({ + subscriptions: [ + { + skuKey: 'pentest_monthly_1', + status: 'trialing', + includedQuantity: 1, + usedQuantity: 1, + }, + ], + creditBalances: [{ productKey: 'pentest', balance: 5 }], + }), + ).toEqual({ balance: 5, planRequired: false }); + }); + + it('clamps negative subscription remainder to zero before adding wallet credits', () => { + // Defensive: if Stripe ever reports usedQuantity > includedQuantity + // (e.g. proration glitch), the UI should not subtract from wallet. + expect( + getPentestAllowance({ + subscriptions: [ + { + skuKey: 'pentest_monthly_1', + status: 'active', + includedQuantity: 1, + usedQuantity: 5, + }, + ], + creditBalances: [{ productKey: 'pentest', balance: 3 }], + }), + ).toEqual({ balance: 3, planRequired: false }); + }); + + it('ignores wallet credits for unrelated products', () => { + expect( + getPentestAllowance({ + creditBalances: [ + { productKey: 'background_check', balance: 10 }, + ], + }), + ).toEqual({ balance: 0, planRequired: true }); + }); + + it('ignores cancelled subscriptions', () => { + // Match the backend's access-status filter — only active/trialing + // subscriptions count toward subscription remainder. + expect( + getPentestAllowance({ + subscriptions: [ + { + skuKey: 'pentest_monthly_1', + status: 'canceled', + includedQuantity: 1, + usedQuantity: 0, + }, + ], + creditBalances: [{ productKey: 'pentest', balance: 2 }], + }), + ).toEqual({ balance: 2, planRequired: false }); + }); +}); diff --git a/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/pentest-allowance.ts b/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/pentest-allowance.ts new file mode 100644 index 0000000000..a39384ef39 --- /dev/null +++ b/apps/app/src/app/(app)/[orgId]/security/penetration-tests/_components/pentest-allowance.ts @@ -0,0 +1,71 @@ +import { getBillingSkuProductKey } from '@trycompai/billing'; + +/** + * Shape of `/v1/billing/status` fields the pentest UI consumes. We only + * declare what we read so the type doesn't drift from the broader + * BillingStatus contract every time an unrelated billing field + * changes. + */ +export interface PentestBillingStatusInput { + subscriptions?: Array<{ + skuKey: string; + status: string; + includedQuantity: number; + usedQuantity: number; + }>; + creditBalances?: Array<{ + productKey: 'pentest' | 'background_check'; + balance: number; + }>; +} + +export interface PentestAllowance { + /** + * Total runs the user can start right now: subscription remainder + + * admin-granted wallet credits. `undefined` while billing is still + * loading so the page does not flash a paywalled state. + */ + balance: number | undefined; + /** + * True only when neither an active subscription nor a wallet credit + * exists. The UI uses this to decide whether to show "+ New Scan" or + * "View plans". + */ + planRequired: boolean; +} + +/** + * Compute the user-visible pentest allowance from billing status. + * + * Mirrors `BillingEntitlementsService.tryConsumeIncludedUsageForProduct`: + * the backend consumes from the active Stripe subscription first and + * falls back to the per-product wallet (BillingCreditBalance). Surfacing + * both here keeps the UI's gating in lockstep with what the create + * endpoint will actually accept — without this, admin-granted credits + * are invisible to the customer and the New Scan button incorrectly + * redirects to billing. + */ +export function getPentestAllowance( + billingStatus: PentestBillingStatusInput | undefined, +): PentestAllowance { + if (billingStatus === undefined) { + return { balance: undefined, planRequired: false }; + } + const subscription = (billingStatus.subscriptions ?? []).find( + (item) => + getBillingSkuProductKey(item.skuKey) === 'pentest' && + (item.status === 'active' || item.status === 'trialing'), + ); + const subscriptionRemaining = subscription + ? Math.max(subscription.includedQuantity - subscription.usedQuantity, 0) + : null; + const walletBalance = + (billingStatus.creditBalances ?? []).find( + (entry) => entry.productKey === 'pentest', + )?.balance ?? 0; + + return { + balance: (subscriptionRemaining ?? 0) + walletBalance, + planRequired: subscriptionRemaining === null && walletBalance === 0, + }; +} diff --git a/apps/app/src/app/(app)/[orgId]/settings/billing/emptyBillingStatus.ts b/apps/app/src/app/(app)/[orgId]/settings/billing/emptyBillingStatus.ts index 635758166c..defe12071e 100644 --- a/apps/app/src/app/(app)/[orgId]/settings/billing/emptyBillingStatus.ts +++ b/apps/app/src/app/(app)/[orgId]/settings/billing/emptyBillingStatus.ts @@ -9,6 +9,7 @@ export const emptyBillingStatus: BackgroundCheckBillingStatus = { }, invoices: [], subscriptions: [], + creditBalances: [], trialEligibility: { pentest: false, background_check: false, diff --git a/apps/app/src/app/(app)/[orgId]/settings/billing/types.ts b/apps/app/src/app/(app)/[orgId]/settings/billing/types.ts index 71f1b1885d..6f8fc66c26 100644 --- a/apps/app/src/app/(app)/[orgId]/settings/billing/types.ts +++ b/apps/app/src/app/(app)/[orgId]/settings/billing/types.ts @@ -34,6 +34,14 @@ export interface BackgroundCheckBillingStatus { currentPeriodEnd: string | null; cancelAtPeriodEnd: boolean; }>; + // Aggregated wallet credit balance per product. The API guarantees + // one entry per product key; an absent product means zero balance. + // Optional on the client so older payloads (and existing test + // fixtures) keep typechecking. + creditBalances?: Array<{ + productKey: 'pentest' | 'background_check'; + balance: number; + }>; invoices?: BillingInvoice[]; } diff --git a/apps/app/src/app/(app)/[orgId]/tasks/components/ModernSingleStatusTaskList.test.tsx b/apps/app/src/app/(app)/[orgId]/tasks/components/ModernSingleStatusTaskList.test.tsx index 037bf8306b..403f5c1127 100644 --- a/apps/app/src/app/(app)/[orgId]/tasks/components/ModernSingleStatusTaskList.test.tsx +++ b/apps/app/src/app/(app)/[orgId]/tasks/components/ModernSingleStatusTaskList.test.tsx @@ -142,6 +142,7 @@ const mockTasks = [ approvedAt: null, approvalComment: null, organizationId: 'org_1', + embeddingHash: null, controls: [], }, { @@ -168,6 +169,7 @@ const mockTasks = [ approvedAt: null, approvalComment: null, organizationId: 'org_1', + embeddingHash: null, controls: [], }, ]; diff --git a/apps/app/src/app/(app)/[orgId]/tasks/components/TaskList.test.tsx b/apps/app/src/app/(app)/[orgId]/tasks/components/TaskList.test.tsx index 459f1714c7..e3d0c90c25 100644 --- a/apps/app/src/app/(app)/[orgId]/tasks/components/TaskList.test.tsx +++ b/apps/app/src/app/(app)/[orgId]/tasks/components/TaskList.test.tsx @@ -150,6 +150,7 @@ const baseMockTask = { approverId: null, approvedAt: null, approvalComment: null, + embeddingHash: null, controls: [] as { id: string; name: string }[], }; diff --git a/apps/app/src/app/(app)/[orgId]/trust/portal-settings/components/TrustPortalSwitch.tsx b/apps/app/src/app/(app)/[orgId]/trust/portal-settings/components/TrustPortalSwitch.tsx index 5c6d49369d..40cf35037c 100644 --- a/apps/app/src/app/(app)/[orgId]/trust/portal-settings/components/TrustPortalSwitch.tsx +++ b/apps/app/src/app/(app)/[orgId]/trust/portal-settings/components/TrustPortalSwitch.tsx @@ -1,9 +1,8 @@ 'use client'; -import { useDebounce } from '@/hooks/useDebounce'; import { usePermissions } from '@/hooks/use-permissions'; import { useTrustPortalSettings } from '@/hooks/use-trust-portal-settings'; -import { Form } from '@trycompai/ui/form'; +import { useDebounce } from '@/hooks/useDebounce'; import { zodResolver } from '@hookform/resolvers/zod'; import { Button, @@ -25,20 +24,12 @@ import { TooltipContent, TooltipTrigger, } from '@trycompai/design-system'; +import { Form } from '@trycompai/ui/form'; import { Download, Eye, FileCheck2, Upload } from 'lucide-react'; import { useCallback, useEffect, useRef, useState } from 'react'; import { useForm } from 'react-hook-form'; import { toast } from 'sonner'; import { z } from 'zod'; -import { - TrustPortalAdditionalDocumentsSection, - type TrustPortalDocument, -} from './TrustPortalAdditionalDocumentsSection'; -import { TrustPortalCustomLinks } from './TrustPortalCustomLinks'; -import { TrustPortalFaqBuilder } from './TrustPortalFaqBuilder'; -import { TrustPortalOverview } from './TrustPortalOverview'; -import { TrustPortalVendors } from './TrustPortalVendors'; -import { UpdateTrustFavicon } from './UpdateTrustFavicon'; import { BrandSettings } from './BrandSettings'; import { GDPR, @@ -62,6 +53,15 @@ import { SOC3, SOC3InProgress, } from './logos'; +import { + TrustPortalAdditionalDocumentsSection, + type TrustPortalDocument, +} from './TrustPortalAdditionalDocumentsSection'; +import { TrustPortalCustomLinks } from './TrustPortalCustomLinks'; +import { TrustPortalFaqBuilder } from './TrustPortalFaqBuilder'; +import { TrustPortalOverview } from './TrustPortalOverview'; +import { TrustPortalVendors } from './TrustPortalVendors'; +import { UpdateTrustFavicon } from './UpdateTrustFavicon'; // Client-side form schema (includes all fields for form state) const trustPortalSwitchSchema = z.object({ @@ -521,7 +521,10 @@ export function TrustPortalSwitch({ }); toast.success('ISO 27001 status updated'); } catch (error) { - toast.error('Failed to update ISO 27001 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update ISO 27001 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { @@ -531,7 +534,10 @@ export function TrustPortalSwitch({ }); toast.success('ISO 27001 status updated'); } catch (error) { - toast.error('Failed to update ISO 27001 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update ISO 27001 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.iso27001} @@ -554,7 +560,10 @@ export function TrustPortalSwitch({ }); toast.success('ISO 42001 status updated'); } catch (error) { - toast.error('Failed to update ISO 42001 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update ISO 42001 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { @@ -564,7 +573,10 @@ export function TrustPortalSwitch({ }); toast.success('ISO 42001 status updated'); } catch (error) { - toast.error('Failed to update ISO 42001 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update ISO 42001 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.iso42001} @@ -587,7 +599,10 @@ export function TrustPortalSwitch({ }); toast.success('GDPR status updated'); } catch (error) { - toast.error('Failed to update GDPR status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update GDPR status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { @@ -597,7 +612,10 @@ export function TrustPortalSwitch({ }); toast.success('GDPR status updated'); } catch (error) { - toast.error('Failed to update GDPR status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update GDPR status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.gdpr} @@ -620,7 +638,10 @@ export function TrustPortalSwitch({ }); toast.success('HIPAA status updated'); } catch (error) { - toast.error('Failed to update HIPAA status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update HIPAA status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { @@ -630,7 +651,10 @@ export function TrustPortalSwitch({ }); toast.success('HIPAA status updated'); } catch (error) { - toast.error('Failed to update HIPAA status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update HIPAA status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.hipaa} @@ -653,7 +677,10 @@ export function TrustPortalSwitch({ }); toast.success('SOC 2 Type 1 status updated'); } catch (error) { - toast.error('Failed to update SOC 2 Type 1 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update SOC 2 Type 1 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { @@ -663,7 +690,10 @@ export function TrustPortalSwitch({ }); toast.success('SOC 2 Type 1 status updated'); } catch (error) { - toast.error('Failed to update SOC 2 Type 1 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update SOC 2 Type 1 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.soc2type1} @@ -686,19 +716,23 @@ export function TrustPortalSwitch({ }); toast.success('SOC 2 Type 2 status updated'); } catch (error) { - toast.error('Failed to update SOC 2 Type 2 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update SOC 2 Type 2 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { try { await updateFrameworkSettings( - checked - ? { soc2type2: true } - : { soc2: false, soc2type2: false }, + checked ? { soc2type2: true } : { soc2: false, soc2type2: false }, ); toast.success('SOC 2 Type 2 status updated'); } catch (error) { - toast.error('Failed to update SOC 2 Type 2 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update SOC 2 Type 2 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.soc2type2} @@ -721,7 +755,10 @@ export function TrustPortalSwitch({ }); toast.success('SOC 3 status updated'); } catch (error) { - toast.error('Failed to update SOC 3 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update SOC 3 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { @@ -731,7 +768,10 @@ export function TrustPortalSwitch({ }); toast.success('SOC 3 status updated'); } catch (error) { - toast.error('Failed to update SOC 3 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update SOC 3 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.soc3} @@ -754,7 +794,10 @@ export function TrustPortalSwitch({ }); toast.success('PCI DSS status updated'); } catch (error) { - toast.error('Failed to update PCI DSS status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update PCI DSS status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { @@ -764,7 +807,10 @@ export function TrustPortalSwitch({ }); toast.success('PCI DSS status updated'); } catch (error) { - toast.error('Failed to update PCI DSS status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update PCI DSS status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.pcidss} @@ -787,7 +833,10 @@ export function TrustPortalSwitch({ }); toast.success('NEN 7510 status updated'); } catch (error) { - toast.error('Failed to update NEN 7510 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update NEN 7510 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { @@ -797,7 +846,10 @@ export function TrustPortalSwitch({ }); toast.success('NEN 7510 status updated'); } catch (error) { - toast.error('Failed to update NEN 7510 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update NEN 7510 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.nen7510} @@ -820,7 +872,10 @@ export function TrustPortalSwitch({ }); toast.success('ISO 9001 status updated'); } catch (error) { - toast.error('Failed to update ISO 9001 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update ISO 9001 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} onToggle={async (checked) => { @@ -830,7 +885,10 @@ export function TrustPortalSwitch({ }); toast.success('ISO 9001 status updated'); } catch (error) { - toast.error('Failed to update ISO 9001 status'); + console.error('[trust framework update] failed', error); + toast.error('Failed to update ISO 9001 status', { + description: error instanceof Error ? error.message : undefined, + }); } }} fileName={certificateFiles.iso9001} @@ -890,14 +948,21 @@ export function TrustPortalSwitch({ />
- ); } -function ComplianceFrameworkLogo({ title, status, enabled }: { title: string; status: string; enabled: boolean }) { +function ComplianceFrameworkLogo({ + title, + status, + enabled, +}: { + title: string; + status: string; + enabled: boolean; +}) { const isInProgress = status === 'in_progress'; let LogoComponent: React.ElementType | null = null; @@ -1047,9 +1112,7 @@ function ComplianceFramework({
{title}
- - {description} - + {description}
@@ -1059,20 +1122,31 @@ function ComplianceFramework({
{isEnabled ? ( - { + if (!value) return; + const prev = status; + setStatus(value); + try { + await onStatusChange(value); + } catch { + setStatus(prev); + } + }} + > - - {status === 'compliant' ? 'Compliant' : status === 'in_progress' ? 'In Progress' : 'Started'} + + {status === 'compliant' + ? 'Compliant' + : status === 'in_progress' + ? 'In Progress' + : 'Started'} @@ -1103,14 +1177,18 @@ function ComplianceFramework({ )}
- { - setIsEnabled(checked); - try { - await onToggle(checked); - } catch { - setIsEnabled(!checked); - } - }} /> + { + setIsEnabled(checked); + try { + await onToggle(checked); + } catch { + setIsEnabled(!checked); + } + }} + />
diff --git a/apps/app/src/app/(app)/[orgId]/vendors/(overview)/components/VendorsTable.test.tsx b/apps/app/src/app/(app)/[orgId]/vendors/(overview)/components/VendorsTable.test.tsx index 05baa58be4..f8cd902fee 100644 --- a/apps/app/src/app/(app)/[orgId]/vendors/(overview)/components/VendorsTable.test.tsx +++ b/apps/app/src/app/(app)/[orgId]/vendors/(overview)/components/VendorsTable.test.tsx @@ -257,6 +257,43 @@ describe('VendorsTable', () => { expect(screen.getByText('4/10')).toBeInTheDocument(); }); + it('renders the RESIDUAL RISK column immediately after INHERENT RISK', () => { + setMockPermissions({}); + + render( + , + ); + + expect(screen.getByText('RESIDUAL RISK')).toBeInTheDocument(); + + const headers = screen + .getAllByRole('columnheader') + .map((h) => (h.textContent || '').toUpperCase()); + const inherentIdx = headers.findIndex((h) => h.includes('INHERENT RISK')); + const residualIdx = headers.findIndex((h) => h.includes('RESIDUAL RISK')); + expect(inherentIdx).toBeGreaterThanOrEqual(0); + expect(residualIdx).toBe(inherentIdx + 1); + }); + + it('renders a residual score badge for assessed vendors', () => { + setMockPermissions({}); + + render( + , + ); + + // Acme Corp residual (unlikely × minor) → raw 4 → score 2/10 + expect(screen.getByText('2/10')).toBeInTheDocument(); + }); + it('shows an em-dash for vendors that have not been assessed', () => { setMockPermissions({}); @@ -275,7 +312,8 @@ describe('VendorsTable', () => { />, ); - expect(screen.getByText('—')).toBeInTheDocument(); + // One em-dash per risk column (inherent + residual) for not_assessed vendors. + expect(screen.getAllByText('—').length).toBe(2); expect(screen.queryByText('1/10')).not.toBeInTheDocument(); }); }); diff --git a/apps/app/src/app/(app)/[orgId]/vendors/(overview)/components/VendorsTable.tsx b/apps/app/src/app/(app)/[orgId]/vendors/(overview)/components/VendorsTable.tsx index 376b6b731a..806fabd745 100644 --- a/apps/app/src/app/(app)/[orgId]/vendors/(overview)/components/VendorsTable.tsx +++ b/apps/app/src/app/(app)/[orgId]/vendors/(overview)/components/VendorsTable.tsx @@ -6,6 +6,12 @@ import { VendorStatus } from '@/components/vendor-status'; import { usePermissions } from '@/hooks/use-permissions'; import { useVendors, useVendorActions, type Vendor } from '@/hooks/use-vendors'; import { getRiskScore } from '@/lib/risk-score'; +import { + interpolatedResidualScore, + previewResidual, + suggestedResidual, +} from '@/lib/suggested-residual'; +import type { TaskStatus } from '@db'; import { AlertDialog, AlertDialogAction, @@ -54,6 +60,40 @@ export type VendorRow = Vendor & { isAssessing?: boolean; }; +/** + * Mirrors `currentSeverityScore` in the risks table — projects the vendor's + * inherent + treatment-strategy + linked-task completion into the same + * interpolated 1–10 score the Treatment Plan hero shows. Falls back to + * inherent when there's no linked work or strategy doesn't reduce. + */ +function currentVendorSeverityScore(vendor: { + inherentProbability: VendorRow['inherentProbability']; + inherentImpact: VendorRow['inherentImpact']; + treatmentStrategy: VendorRow['treatmentStrategy']; + tasks?: Array<{ status: TaskStatus }>; +}): number { + const inherent = getRiskScore(vendor.inherentProbability, vendor.inherentImpact); + const tasks = vendor.tasks ?? []; + const target = previewResidual({ + inherentLikelihood: vendor.inherentProbability, + inherentImpact: vendor.inherentImpact, + strategy: vendor.treatmentStrategy, + hasLinkedWork: tasks.length > 0, + }); + const targetScore = getRiskScore(target.likelihood, target.impact).score; + const completion = suggestedResidual({ + likelihood: vendor.inherentProbability, + impact: vendor.inherentImpact, + strategy: vendor.treatmentStrategy, + tasks, + }).completion; + return interpolatedResidualScore({ + inherentScore: inherent.score, + targetScore, + completion, + }); +} + type AssigneeMember = { id: string; role: string; @@ -171,7 +211,7 @@ export function VendorsTable({ // Local state for search, sorting, and pagination const [searchQuery, setSearchQuery] = useState(''); const [sort, setSort] = useState<{ - id: 'name' | 'updatedAt' | 'inherentRisk'; + id: 'name' | 'updatedAt' | 'inherentRisk' | 'residualRisk'; desc: boolean; }>({ id: 'name', @@ -264,6 +304,8 @@ export function VendorsTable({ inherentImpact: 'insignificant' as const, residualProbability: 'very_unlikely' as const, residualImpact: 'insignificant' as const, + treatmentStrategy: 'accept' as const, + treatmentStrategyDescription: null, website: null, isSubProcessor: false, logoUrl: null, @@ -290,6 +332,8 @@ export function VendorsTable({ inherentImpact: 'insignificant' as const, residualProbability: 'very_unlikely' as const, residualImpact: 'insignificant' as const, + treatmentStrategy: 'accept' as const, + treatmentStrategyDescription: null, website: null, isSubProcessor: false, logoUrl: null, @@ -334,6 +378,25 @@ export function VendorsTable({ const comparison = aScore - bScore; return sort.desc ? -comparison : comparison; } + if (sort.id === 'residualRisk') { + // Unassessed vendors carry default residual values (very_unlikely + // × insignificant = 1). Without this branch they'd cluster at the + // bottom (or top, when desc) of the residual sort even though we + // render them as `—` and they have no real residual yet. Force + // them to the end of the list regardless of sort direction so + // assessed vendors are always grouped together. (Cubic finding + // on PR #2671.) + const aAssessed = a.status === 'assessed'; + const bAssessed = b.status === 'assessed'; + if (aAssessed !== bAssessed) return aAssessed ? -1 : 1; + // Sort by the SAME interpolated score the badge renders so the + // sort order matches what the user sees (treatment-progress + // aware), not the static residual fields. + const aScore = currentVendorSeverityScore(a); + const bScore = currentVendorSeverityScore(b); + const comparison = aScore - bScore; + return sort.desc ? -comparison : comparison; + } const comparison = new Date(a.updatedAt).getTime() - new Date(b.updatedAt).getTime(); return sort.desc ? -comparison : comparison; }); @@ -392,7 +455,7 @@ export function VendorsTable({ router.push(`/${orgId}/vendors/${vendorId}`); }; - const handleSort = (columnId: 'name' | 'updatedAt' | 'inherentRisk') => { + const handleSort = (columnId: 'name' | 'updatedAt' | 'inherentRisk' | 'residualRisk') => { if (sort.id === columnId) { setSort({ id: columnId, desc: !sort.desc }); } else { @@ -400,7 +463,7 @@ export function VendorsTable({ } }; - const getSortIcon = (columnId: 'name' | 'updatedAt' | 'inherentRisk') => { + const getSortIcon = (columnId: 'name' | 'updatedAt' | 'inherentRisk' | 'residualRisk') => { if (sort.id !== columnId) { return ; } @@ -545,6 +608,16 @@ export function VendorsTable({ {getSortIcon('inherentRisk')} + + + CATEGORY OWNER {hasPermission('vendor', 'delete') && ACTIONS} @@ -576,6 +649,18 @@ export function VendorsTable({ /> )} + + {vendor.status === 'not_assessed' ? ( + + ) : ( + // Show the current (interpolated) score that + // reflects how far the linked tasks have driven + // the residual down — same logic the risks table + // uses. Static residualProbability / Impact alone + // can't reflect mid-treatment progress. + + )} + {CATEGORY_MAP[vendor.category] || vendor.category} diff --git a/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorDetailTabs.tsx b/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorDetailTabs.tsx index a247af7428..1fc0f69235 100644 --- a/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorDetailTabs.tsx +++ b/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorDetailTabs.tsx @@ -17,7 +17,8 @@ import { VendorResearchBadges, VendorResearchLinks } from './VendorResearchSecti import { VendorResearchFeed } from './VendorResearchFeed'; import { VendorInherentRiskChart } from './VendorInherentRiskChart'; import { VendorResidualRiskChart } from './VendorResidualRiskChart'; -import type { Member, User, Vendor } from '@db'; +import { TreatmentPlanTab } from '@/components/risks/treatment-plan/TreatmentPlanTab'; +import type { Member, RiskTreatmentType, User, Vendor } from '@db'; import { CommentEntityType } from '@db'; import type { Prisma } from '@db'; import { useRealtimeRun } from '@trigger.dev/react-hooks'; @@ -35,6 +36,7 @@ import { } from '@trycompai/design-system'; import Link from 'next/link'; import { useSearchParams } from 'next/navigation'; +import { useQueryState } from 'nuqs'; import { useCallback, useEffect, useMemo, useState } from 'react'; import { toast } from 'sonner'; @@ -75,11 +77,18 @@ export function VendorDetailTabs({ isViewingTask, }: VendorDetailTabsProps) { const searchParams = useSearchParams(); - const defaultTab = searchParams.get('tab') || 'overview'; const taskItemId = searchParams.get('taskItemId'); const { vendor: swrVendor, mutate: refreshVendor } = useVendor(vendorId); - const { updateVendor, triggerAssessment, regenerateMitigation } = useVendorActions(); + const { + updateVendor, + triggerAssessment, + regenerateMitigation, + suggestVendorLinks, + applyVendorLinks, + fetchActiveVendorAutoLinkRun, + discardVendorAutoLinkRun, + } = useVendorActions(); const { hasPermission } = usePermissions(); const canUpdate = hasPermission('vendor', 'update'); const canUpdateTask = hasPermission('task', 'update'); @@ -92,9 +101,20 @@ export function VendorDetailTabs({ const [isEditingDescription, setIsEditingDescription] = useState(false); const [descriptionValue, setDescriptionValue] = useState(''); const [isMitigationLoading, setIsMitigationLoading] = useState(false); + const [regenRun, setRegenRun] = useState<{ + runId: string; + publicAccessToken: string; + } | null>(null); const [isAssessmentLoading, setIsAssessmentLoading] = useState(false); const [isRegenerating, setIsRegenerating] = useState(false); - const [activeTab, setActiveTab] = useState(defaultTab); + // URL-backed tab state — bookmarks the active tab in `?tab=...` so a + // refresh keeps the same view. Conditional rendering below mounts only + // the active panel, sidestepping base-ui's transition window where the + // outgoing panel stays at full opacity during the incoming fade-in + // (visible "both panels stacked" flash on tab switch). + const [activeTab, setActiveTab] = useQueryState('tab', { + defaultValue: 'overview', + }); const { data: taskItemsData, mutate: refreshTaskItems } = useTaskItems( vendorId, 'vendor', 1, 50, 'createdAt', 'desc', {}, @@ -248,18 +268,75 @@ export function VendorDetailTabs({ const handleRegenerateMitigation = async () => { setIsMitigationLoading(true); - toast.info('Regenerating vendor risk mitigation...'); try { - await regenerateMitigation(vendorId); - toast.success('Mitigation regeneration triggered.'); - refreshVendor(); + const handle = await regenerateMitigation(vendorId); + setRegenRun(handle); } catch { toast.error('Failed to trigger mitigation regeneration'); - } finally { setIsMitigationLoading(false); } }; + const handleRegenSettled = useCallback( + (result: { success: boolean; reason?: string }) => { + setRegenRun(null); + setIsMitigationLoading(false); + if (result.success) { + toast.success('Treatment plan regenerated.'); + void refreshVendor(); + } else { + toast.error(result.reason ?? 'Failed to regenerate the treatment plan.'); + } + }, + [refreshVendor], + ); + + const handleUpdateStrategy = async (strategy: RiskTreatmentType) => { + await updateVendor(vendorId, { treatmentStrategy: strategy }); + refreshVendor(); + }; + + const handleSuggest = async () => { + return suggestVendorLinks(vendorId); + }; + + const handleApply = async (params: { taskIds: string[]; replace: boolean }) => { + await applyVendorLinks(vendorId, params); + await refreshVendor(); + }; + + const handleUnlinkTask = async (taskId: string) => { + try { + const res = await fetch(`/api/vendors/${vendorId}/tasks/${taskId}`, { + method: 'DELETE', + credentials: 'include', + }); + if (!res.ok) throw new Error('unlink failed'); + await refreshVendor(); + } catch { + toast.error('Failed to unlink task'); + } + }; + + // Memoize so AutoLinkSuggestions' resume effect (which depends on the + // callback identity) doesn't re-fire on every parent re-render. + const handleResumeAutoLink = useCallback( + () => fetchActiveVendorAutoLinkRun(vendorId), + [fetchActiveVendorAutoLinkRun, vendorId], + ); + + const handleDiscardAutoLinkRun = useCallback( + async () => { + await discardVendorAutoLinkRun(vendorId); + }, + [discardVendorAutoLinkRun, vendorId], + ); + + const handleUpdateDescription = async (description: string) => { + await updateVendor(vendorId, { treatmentStrategyDescription: description }); + refreshVendor(); + }; + const handleRegenerateAssessment = async () => { setIsAssessmentLoading(true); toast.info('Regenerating vendor risk assessment...'); @@ -268,7 +345,7 @@ export function VendorDetailTabs({ toast.success('Assessment regeneration triggered.'); if (result.runId && result.publicAccessToken) { setIsRegenerating(true); - setActiveTab('risk-assessment'); + void setActiveTab('risk-assessment'); handleAssessmentTriggered(result.runId, result.publicAccessToken); } refreshVendor(); @@ -402,10 +479,11 @@ export function VendorDetailTabs({ {isViewingTask ? ( ) : ( - + void setActiveTab(String(next))}> Overview + Treatment Plan Risk Matrix Risk Assessment Tasks @@ -414,17 +492,61 @@ export function VendorDetailTabs({ Settings - - + {activeTab === 'overview' && ( + + + + )} + + {activeTab === 'treatment-plan' && ( + + > + | null + | undefined ?? null, + tasks: swrVendor?.tasks ?? [], + }} + canUpdate={canUpdate} + onUpdateStrategy={handleUpdateStrategy} + onUpdateDescription={handleUpdateDescription} + onRegenerate={handleRegenerateMitigation} + regenerating={isMitigationLoading} + onSuggest={handleSuggest} + onApply={handleApply} + // Gate the unlink affordance behind vendor:update so it + // doesn't render for read-only users. The Next API route + // also enforces this check server-side. + onUnlinkTask={canUpdate ? handleUnlinkTask : undefined} + onResumeAutoLink={handleResumeAutoLink} + onDiscardAutoLinkRun={handleDiscardAutoLinkRun} + regenRun={regenRun} + onRegenSettled={handleRegenSettled} + /> + )} + {activeTab === 'risk-matrix' && ( + )} + {activeTab === 'risk-assessment' && ( @@ -472,64 +594,51 @@ export function VendorDetailTabs({ - - - - - - - - - - - t.id) || []} /> - - + )} + + {activeTab === 'tasks' && ( + + + + )} + + {activeTab === 'comments' && ( + + + + )} + + {activeTab === 'activity' && ( + + t.id) || []} /> + + )} + + {activeTab === 'settings' && ( {canUpdate && ( - <> - - - Regenerate Risk Assessment - - Generate or regenerate the AI risk assessment for this vendor - - - - - -
- - - - Regenerate Mitigation - - Generate a fresh risk mitigation comment for this vendor - - - - - + + + Regenerate Risk Assessment + + Generate or regenerate the AI risk assessment for this vendor + + + + )} + )} )} diff --git a/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorInherentRiskChart.tsx b/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorInherentRiskChart.tsx index 14cbb25fe3..75fadc494b 100644 --- a/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorInherentRiskChart.tsx +++ b/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorInherentRiskChart.tsx @@ -18,6 +18,7 @@ export function VendorInherentRiskChart({ vendor }: InherentRiskChartProps) { ({ })); // Mock useVendor and useVendorActions +const mockTriggerAssessment = vi.fn(); vi.mock('@/hooks/use-vendors', () => ({ useVendor: () => ({ vendor: null, @@ -23,9 +24,19 @@ vi.mock('@/hooks/use-vendors', () => ({ }), useVendorActions: () => ({ updateVendor: vi.fn(), + triggerAssessment: mockTriggerAssessment, }), })); +// Mock toast +vi.mock('sonner', () => ({ + toast: { + success: vi.fn(), + error: vi.fn(), + info: vi.fn(), + }, +})); + // Capture props passed to RiskMatrixChart let capturedProps: any = null; vi.mock('@/components/risks/charts/RiskMatrixChart', () => ({ @@ -39,14 +50,20 @@ import { VendorResidualRiskChart } from './VendorResidualRiskChart'; const mockVendor: any = { id: 'vendor-1', + status: 'assessed', + inherentProbability: 'possible', + inherentImpact: 'moderate', residualProbability: 'unlikely', residualImpact: 'minor', + treatmentStrategy: 'accept', + tasks: [], }; describe('VendorResidualRiskChart', () => { beforeEach(() => { setMockPermissions({}); capturedProps = null; + mockTriggerAssessment.mockReset(); }); it('passes readOnly=true when user lacks vendor:update permission', () => { @@ -83,7 +100,7 @@ describe('VendorResidualRiskChart', () => { expect(capturedProps.title).toBe('Residual Risk'); expect(capturedProps.description).toBe( - 'Select the residual risk level for this vendor', + 'Risk level after the treatment plan is applied. The dashed cell is the suggestion computed from your strategy and linked task completion.', ); expect(capturedProps.riskId).toBe('vendor-1'); expect(capturedProps.activeLikelihood).toBe('unlikely'); @@ -98,4 +115,37 @@ describe('VendorResidualRiskChart', () => { expect(capturedProps).not.toBeNull(); expect(capturedProps.readOnly).toBe(true); }); + + describe('status branches', () => { + it('renders NotAssessedState when status is not_assessed', () => { + setMockPermissions(ADMIN_PERMISSIONS); + const vendor = { ...mockVendor, status: 'not_assessed' }; + + render(); + + expect(screen.queryByTestId('risk-matrix-chart')).toBeNull(); + expect(screen.getByRole('button', { name: /Run risk assessment/i })).toBeInTheDocument(); + expect(capturedProps).toBeNull(); + }); + + it('passes preliminary=true to RiskMatrixChart when status is in_progress', () => { + setMockPermissions(ADMIN_PERMISSIONS); + const vendor = { ...mockVendor, status: 'in_progress' }; + + render(); + + expect(capturedProps).not.toBeNull(); + expect(capturedProps.preliminary).toBe(true); + }); + + it('passes preliminary=false to RiskMatrixChart when status is assessed', () => { + setMockPermissions(ADMIN_PERMISSIONS); + const vendor = { ...mockVendor, status: 'assessed' }; + + render(); + + expect(capturedProps).not.toBeNull(); + expect(capturedProps.preliminary).toBe(false); + }); + }); }); diff --git a/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorResidualRiskChart.tsx b/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorResidualRiskChart.tsx index ec2062dd83..55cabf9d44 100644 --- a/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorResidualRiskChart.tsx +++ b/apps/app/src/app/(app)/[orgId]/vendors/[vendorId]/components/VendorResidualRiskChart.tsx @@ -1,27 +1,70 @@ 'use client'; +import { RiskMatrixChart } from '@/components/risks/charts/RiskMatrixChart'; +import { NotAssessedState } from '@/components/risks/treatment-plan/NotAssessedState'; import { usePermissions } from '@/hooks/use-permissions'; import { useVendor, useVendorActions } from '@/hooks/use-vendors'; -import { RiskMatrixChart } from '@/components/risks/charts/RiskMatrixChart'; -import type { Vendor } from '@db'; +import { suggestedResidual } from '@/lib/suggested-residual'; +import { VendorStatus, type TaskStatus, type Vendor } from '@db'; +import { toast } from 'sonner'; interface ResidualRiskChartProps { - vendor: Vendor; + vendor: Vendor & { tasks?: { status: TaskStatus }[] }; } export function VendorResidualRiskChart({ vendor }: ResidualRiskChartProps) { - const { updateVendor } = useVendorActions(); + const { updateVendor, triggerAssessment } = useVendorActions(); const { mutate } = useVendor(vendor.id); const { hasPermission } = usePermissions(); + const canUpdate = hasPermission('vendor', 'update'); + + if (vendor.status === VendorStatus.not_assessed) { + return ( + { + try { + await triggerAssessment(vendor.id); + toast.success('Risk assessment started. This may take a moment.'); + await mutate(); + } catch { + toast.error('Failed to start risk assessment'); + } + }} + /> + ); + } + + // Only compute a suggestion when tasks are actually loaded — falling back to + // [] would render a misleading "0% complete" ghost cell on vendors that + // haven't hydrated yet. + const suggestion = vendor.tasks + ? suggestedResidual({ + likelihood: vendor.inherentProbability, + impact: vendor.inherentImpact, + strategy: vendor.treatmentStrategy, + tasks: vendor.tasks, + }) + : undefined; + + const preliminary = vendor.status === VendorStatus.in_progress; + return ( { await updateVendor(id, { residualProbability: probability, diff --git a/apps/app/src/app/(app)/upgrade/[orgId]/page.tsx b/apps/app/src/app/(app)/upgrade/[orgId]/page.tsx index d33ae89074..533604bd6f 100644 --- a/apps/app/src/app/(app)/upgrade/[orgId]/page.tsx +++ b/apps/app/src/app/(app)/upgrade/[orgId]/page.tsx @@ -1,6 +1,6 @@ -import { extractDomain, isDomainActiveStripeCustomer, isPublicEmailDomain } from '@/lib/stripe'; -import { auth } from '@/utils/auth'; import { env } from '@/env.mjs'; +import { serverApi } from '@/lib/api-server'; +import { auth } from '@/utils/auth'; import { db } from '@db/server'; import { headers } from 'next/headers'; import { redirect } from 'next/navigation'; @@ -13,6 +13,12 @@ interface PageProps { }>; } +interface AutoApproveResponse { + hasAccess: boolean; + autoApproved: boolean; + reason: string; +} + export default async function UpgradePage({ params }: PageProps) { const { orgId } = await params; @@ -44,7 +50,9 @@ export default async function UpgradePage({ params }: PageProps) { redirect('/'); } - // Sync activeOrganizationId only after membership is verified + // Sync activeOrganizationId only after membership is verified. + // Required so the API's HybridAuthGuard resolves the right org from session + // when we call /v1/organization-access/auto-approve below. const currentActiveOrgId = authSession.session.activeOrganizationId; if (!currentActiveOrgId || currentActiveOrgId !== orgId) { try { @@ -61,44 +69,32 @@ export default async function UpgradePage({ params }: PageProps) { let hasAccess = member.organization.hasAccess; - // Auto-approve based on user's email domain or self-hosted instance if (!hasAccess) { - // Auto-approve for self-hosted/OSS instances - const isSelfHosted = env.NEXT_PUBLIC_SELF_HOSTED === 'true'; - - if (isSelfHosted) { + // Self-hosted instances auto-approve every org. The flag is a Next.js + // build-time env var (NEXT_PUBLIC_SELF_HOSTED) that the OSS Docker + // deployment sets on the app container only — the API container does NOT + // have this env, so the check stays on the page. The DB write here is the + // single exception to "all mutations through the API" — it's gated on a + // build-time deploy flag, not user input. + if (env.NEXT_PUBLIC_SELF_HOSTED === 'true') { await db.organization.update({ where: { id: orgId }, data: { hasAccess: true }, }); hasAccess = true; } else { - const userEmail = authSession.user.email; - const userEmailDomain = extractDomain(userEmail ?? ''); - const orgWebsiteDomain = extractDomain(member.organization.website ?? ''); - - if (userEmailDomain) { - // Auto-approve for trycomp.ai emails (internal team) - const isTrycompEmail = userEmailDomain === 'trycomp.ai'; - - const canAutoApproveViaDomain = - !isTrycompEmail && - Boolean(orgWebsiteDomain) && - userEmailDomain === orgWebsiteDomain && - !isPublicEmailDomain(userEmailDomain); - - // Check Stripe for other domains - const isStripeCustomer = canAutoApproveViaDomain - ? await isDomainActiveStripeCustomer(userEmailDomain) - : false; - - if (isTrycompEmail || isStripeCustomer) { - await db.organization.update({ - where: { id: orgId }, - data: { hasAccess: true }, - }); - hasAccess = true; - } + // Stripe-domain auto-approval (and the @trycomp.ai shortcut) live in the + // API so STRIPE_SECRET_KEY only has to exist on the API and the + // hasAccess flip is RBAC-checked + audit-logged. Soft-fail so a transient + // API error never blocks the booking step from rendering. + const response = await serverApi.post( + '/v1/organization-access/auto-approve', + ); + + if (response.data?.hasAccess) { + hasAccess = true; + } else if (response.error) { + console.error('[UpgradePage] auto-approve API error:', response.error); } } } diff --git a/apps/app/src/app/api/risks/[riskId]/auto-link/active/route.ts b/apps/app/src/app/api/risks/[riskId]/auto-link/active/route.ts new file mode 100644 index 0000000000..f7edbfe67e --- /dev/null +++ b/apps/app/src/app/api/risks/[riskId]/auto-link/active/route.ts @@ -0,0 +1,122 @@ +import { requireApiPermission } from '@/lib/permissions.server'; +import { db } from '@db/server'; +import { auth as triggerAuth } from '@trigger.dev/sdk'; +import { NextRequest, NextResponse } from 'next/server'; + +/** + * Decide whether a `createPublicToken` failure is the "run was purged" + * terminal case (drop the stored runId — the run is permanently gone) or + * something transient (a network blip, rate limit, auth issue) where we + * should keep the runId so the next attempt can recover. + * + * Trigger.dev surfaces purged/missing runs as 404 with a "not found" body. + * Anything else — including unauthorized, timeouts, 5xx — gets treated as + * transient and bubbles up as a 502 from the route. See Cubic finding + * #4/#5 on PR #2671. + */ +function isRunGoneError(err: unknown): boolean { + if (!err) return false; + const e = err as { status?: number; statusCode?: number; message?: string }; + const status = e.status ?? e.statusCode; + if (status === 404) return true; + const msg = (e.message ?? String(err)).toLowerCase(); + return msg.includes('not found') || msg.includes('purged'); +} + +/** + * GET /api/risks/[riskId]/auto-link/active + * + * Resumes an in-flight or completed-but-unreviewed AI suggestion run after a + * page reload. The runId is persisted on the Risk row by `/auto-link`; this + * endpoint mints a fresh public-access token (the previous one expires after + * 15 minutes) so the UI can re-subscribe via `useRealtimeRun`. + * + * Returns `{ runId: null }` when no active run exists. Also returns null when + * the trigger.dev run has been purged (TTL elapsed) — the caller treats both + * the same: drop the stale runId and start fresh on the next user action. + * + * DELETE on the same path discards the active run (user clicked Discard). + */ +export async function GET( + req: NextRequest, + { params }: { params: Promise<{ riskId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'risk', 'read'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { riskId } = await params; + + const risk = await db.risk.findUnique({ + where: { id: riskId }, + select: { id: true, organizationId: true, autoLinkRunId: true }, + }); + if (!risk || risk.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + if (!risk.autoLinkRunId) { + return NextResponse.json({ runId: null }); + } + + try { + const publicAccessToken = await triggerAuth.createPublicToken({ + scopes: { read: { runs: [risk.autoLinkRunId] } }, + expirationTime: '15m', + }); + return NextResponse.json({ runId: risk.autoLinkRunId, publicAccessToken }); + } catch (err) { + if (isRunGoneError(err)) { + // Run was purged by trigger.dev (retention TTL or never existed). + // Drop the stale id so the next /auto-link call starts cleanly. + console.warn('[auto-link/active] run gone, clearing stored runId', err); + await db.risk.update({ + where: { id: riskId }, + data: { autoLinkRunId: null, autoLinkRunStartedAt: null }, + }); + return NextResponse.json({ runId: null }); + } + // Transient failure (network, 5xx, rate limit). Keep the runId so a + // retry can recover; the UI treats 502 as "try again later". + console.error('[auto-link/active] transient token mint failure', err); + return NextResponse.json( + { error: 'Failed to mint access token; try again' }, + { status: 502 }, + ); + } + } catch (error) { + console.error('Error reading active risk auto-link run:', error); + return NextResponse.json({ error: 'Failed to read active run' }, { status: 500 }); + } +} + +export async function DELETE( + req: NextRequest, + { params }: { params: Promise<{ riskId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'risk', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { riskId } = await params; + + const risk = await db.risk.findUnique({ + where: { id: riskId }, + select: { id: true, organizationId: true }, + }); + if (!risk || risk.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + await db.risk.update({ + where: { id: riskId }, + data: { autoLinkRunId: null, autoLinkRunStartedAt: null }, + }); + return NextResponse.json({ ok: true }); + } catch (error) { + console.error('Error discarding risk auto-link run:', error); + return NextResponse.json({ error: 'Failed to discard run' }, { status: 500 }); + } +} diff --git a/apps/app/src/app/api/risks/[riskId]/auto-link/apply/route.ts b/apps/app/src/app/api/risks/[riskId]/auto-link/apply/route.ts new file mode 100644 index 0000000000..01e57cf621 --- /dev/null +++ b/apps/app/src/app/api/risks/[riskId]/auto-link/apply/route.ts @@ -0,0 +1,101 @@ +import { requireApiPermission } from '@/lib/permissions.server'; +import { db } from '@db/server'; +import { NextRequest, NextResponse } from 'next/server'; +import { z } from 'zod'; + +const ApplyBodySchema = z.object({ + taskIds: z.array(z.string()).max(100), + replace: z.boolean().optional().default(false), +}); + +/** + * POST /api/risks/[riskId]/auto-link/apply + * + * Persists the user-confirmed task selection from the AI-suggestion review UI. + * + * - `replace: true` → re-assess flow (sync semantics: connect-only-these tasks). + * - `replace: false` → fresh suggest flow (additive: connect these to whatever's + * already linked). + * + * Mutating endpoints elsewhere in the app go through the NestJS API, but the + * task↔risk join already lives in this Next.js layer (see `tasks/[taskId]` + * DELETE), and the AI scan / suggestion plumbing all lives here too. Keeping + * apply alongside avoids a round trip and stays consistent with the existing + * unlink endpoint. + */ +export async function POST( + req: NextRequest, + { params }: { params: Promise<{ riskId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'risk', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { riskId } = await params; + if (!riskId) { + return NextResponse.json({ error: 'Risk ID is required' }, { status: 400 }); + } + + const risk = await db.risk.findUnique({ + where: { id: riskId }, + select: { id: true, organizationId: true }, + }); + if (!risk || risk.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + const body = await req.json().catch(() => null); + const parsed = ApplyBodySchema.safeParse(body); + if (!parsed.success) { + return NextResponse.json( + { error: 'Invalid body', details: parsed.error.flatten() }, + { status: 400 }, + ); + } + const { taskIds, replace } = parsed.data; + + // Validate every taskId belongs to the active organization. Without this, + // a malicious caller could connect another org's tasks to this risk by + // crafting the request — Prisma's `connect`/`set` doesn't check tenancy. + if (taskIds.length > 0) { + const ownedTasks = await db.task.findMany({ + where: { id: { in: taskIds }, organizationId }, + select: { id: true }, + }); + if (ownedTasks.length !== taskIds.length) { + return NextResponse.json( + { error: 'One or more tasks do not belong to this organization' }, + { status: 400 }, + ); + } + } + + if (replace) { + await db.risk.update({ + where: { id: riskId }, + data: { + tasks: { set: taskIds.map((id) => ({ id })) }, + autoLinkRunId: null, + autoLinkRunStartedAt: null, + }, + }); + } else { + await db.risk.update({ + where: { id: riskId }, + data: { + ...(taskIds.length > 0 + ? { tasks: { connect: taskIds.map((id) => ({ id })) } } + : {}), + autoLinkRunId: null, + autoLinkRunStartedAt: null, + }, + }); + } + + return NextResponse.json({ linked: taskIds.length }); + } catch (error) { + console.error('Error applying risk auto-link:', error); + return NextResponse.json({ error: 'Failed to apply auto-link' }, { status: 500 }); + } +} diff --git a/apps/app/src/app/api/risks/[riskId]/auto-link/route.ts b/apps/app/src/app/api/risks/[riskId]/auto-link/route.ts new file mode 100644 index 0000000000..27cbce1e41 --- /dev/null +++ b/apps/app/src/app/api/risks/[riskId]/auto-link/route.ts @@ -0,0 +1,68 @@ +import { requireApiPermission } from '@/lib/permissions.server'; +import type { linkRisksAndVendorsToWork } from '@/trigger/tasks/onboarding/link-risks-and-vendors-to-work'; +import { db } from '@db/server'; +import { auth as triggerAuth, tasks } from '@trigger.dev/sdk'; +import { NextRequest, NextResponse } from 'next/server'; + +/** + * POST /api/risks/[riskId]/auto-link + * + * Triggers the linkage task in `suggestionsOnly` mode for one risk and returns + * a public-access token so the frontend can subscribe via `useRealtimeRun`, + * display live progress, and read `run.output.suggestions` once complete. + * + * No DB writes happen here — the user reviews the suggestions in the UI and + * the apply endpoint (`/auto-link/apply`) persists their final selection. + */ +export async function POST( + req: NextRequest, + { params }: { params: Promise<{ riskId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'risk', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { riskId } = await params; + if (!riskId) { + return NextResponse.json( + { error: 'Risk ID is required' }, + { status: 400 }, + ); + } + + const risk = await db.risk.findUnique({ + where: { id: riskId }, + select: { id: true, organizationId: true }, + }); + + if (!risk || risk.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + const handle = await tasks.trigger( + 'link-risks-and-vendors-to-work', + { organizationId, riskId, suggestionsOnly: true }, + ); + + // Persist the runId so the UI can resume an in-flight scan after a page + // reload. Cleared by /apply or /discard once the user reviews the result. + await db.risk.update({ + where: { id: riskId }, + data: { + autoLinkRunId: handle.id, + autoLinkRunStartedAt: new Date(), + }, + }); + + const publicAccessToken = await triggerAuth.createPublicToken({ + scopes: { read: { runs: [handle.id] } }, + expirationTime: '15m', + }); + + return NextResponse.json({ runId: handle.id, publicAccessToken }); + } catch (error) { + console.error('Error triggering risk auto-link:', error); + return NextResponse.json({ error: 'Failed to trigger auto-link' }, { status: 500 }); + } +} diff --git a/apps/app/src/app/api/risks/[riskId]/regenerate-mitigation/route.ts b/apps/app/src/app/api/risks/[riskId]/regenerate-mitigation/route.ts index 38ef486948..26d9157312 100644 --- a/apps/app/src/app/api/risks/[riskId]/regenerate-mitigation/route.ts +++ b/apps/app/src/app/api/risks/[riskId]/regenerate-mitigation/route.ts @@ -1,8 +1,8 @@ import { generateRiskMitigation } from '@/trigger/tasks/onboarding/generate-risk-mitigation'; import type { PolicyContext } from '@/trigger/tasks/onboarding/onboard-organization-helpers'; import { serverApi } from '@/lib/api-server'; -import { auth } from '@/utils/auth'; -import { tasks } from '@trigger.dev/sdk'; +import { requireApiPermission } from '@/lib/permissions.server'; +import { auth as triggerAuth, tasks } from '@trigger.dev/sdk'; import { NextRequest, NextResponse } from 'next/server'; interface PeopleApiResponse { @@ -27,13 +27,9 @@ export async function POST( { params }: { params: Promise<{ riskId: string }> }, ) { try { - const session = await auth.api.getSession({ - headers: req.headers, - }); - - if (!session?.session?.activeOrganizationId) { - return NextResponse.json({ error: 'Unauthorized' }, { status: 401 }); - } + const ctx = await requireApiPermission(req, 'risk', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; const { riskId } = await params; if (!riskId) { @@ -43,8 +39,6 @@ export async function POST( ); } - const organizationId = session.session.activeOrganizationId; - const [peopleResult, policiesResult] = await Promise.all([ serverApi.get('/v1/people'), serverApi.get('/v1/policies'), @@ -71,7 +65,7 @@ export async function POST( description: policy.description, })); - await tasks.trigger( + const handle = await tasks.trigger( 'generate-risk-mitigation', { organizationId, @@ -81,17 +75,27 @@ export async function POST( }, ); - return NextResponse.json({ success: true }); + // The run is now in flight server-side. Mint a 15-min public token so + // the UI can subscribe via useRealtimeRun. If the mint fails, do NOT + // throw — that would let a client retry start a duplicate run. Return + // the runId with a null token; the UI's `/active` endpoint can mint a + // fresh token on the next render. (Cubic finding #29 on PR #2671.) + let publicAccessToken: string | null = null; + try { + publicAccessToken = await triggerAuth.createPublicToken({ + scopes: { read: { runs: [handle.id] } }, + expirationTime: '15m', + }); + } catch (mintErr) { + console.error( + '[regenerate-mitigation] run triggered but token mint failed; client must resume via /active', + { runId: handle.id, mintErr }, + ); + } + + return NextResponse.json({ runId: handle.id, publicAccessToken }); } catch (error) { console.error('Error regenerating risk mitigation:', error); - return NextResponse.json( - { - error: - error instanceof Error - ? error.message - : 'Failed to regenerate mitigation', - }, - { status: 500 }, - ); + return NextResponse.json({ error: 'Failed to regenerate mitigation' }, { status: 500 }); } } diff --git a/apps/app/src/app/api/risks/[riskId]/relink/route.ts b/apps/app/src/app/api/risks/[riskId]/relink/route.ts new file mode 100644 index 0000000000..30aec86054 --- /dev/null +++ b/apps/app/src/app/api/risks/[riskId]/relink/route.ts @@ -0,0 +1,57 @@ +import { requireApiPermission } from '@/lib/permissions.server'; +import type { linkRisksAndVendorsToWork } from '@/trigger/tasks/onboarding/link-risks-and-vendors-to-work'; +import { db } from '@db/server'; +import { auth as triggerAuth, tasks } from '@trigger.dev/sdk'; +import { NextRequest, NextResponse } from 'next/server'; + +/** + * POST /api/risks/[riskId]/relink + * + * Wipes ALL current task links on the risk and rebuilds linkage from the + * top-K embedding-similar tasks. Destructive — clears user's manual unlinks. + * Frontend confirms before calling. + * + * Returns { runId, publicAccessToken } so the UI can subscribe via useRealtimeRun. + */ +export async function POST( + req: NextRequest, + { params }: { params: Promise<{ riskId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'risk', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { riskId } = await params; + if (!riskId) { + return NextResponse.json( + { error: 'Risk ID is required' }, + { status: 400 }, + ); + } + + const risk = await db.risk.findUnique({ + where: { id: riskId }, + select: { id: true, organizationId: true }, + }); + + if (!risk || risk.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + const handle = await tasks.trigger( + 'link-risks-and-vendors-to-work', + { organizationId, riskId, replace: true }, + ); + + const publicAccessToken = await triggerAuth.createPublicToken({ + scopes: { read: { runs: [handle.id] } }, + expirationTime: '15m', + }); + + return NextResponse.json({ runId: handle.id, publicAccessToken }); + } catch (error) { + console.error('Error triggering risk relink:', error); + return NextResponse.json({ error: 'Failed to trigger relink' }, { status: 500 }); + } +} diff --git a/apps/app/src/app/api/risks/[riskId]/tasks/[taskId]/route.ts b/apps/app/src/app/api/risks/[riskId]/tasks/[taskId]/route.ts new file mode 100644 index 0000000000..008f809f06 --- /dev/null +++ b/apps/app/src/app/api/risks/[riskId]/tasks/[taskId]/route.ts @@ -0,0 +1,107 @@ +import { generateRiskMitigation } from '@/trigger/tasks/onboarding/generate-risk-mitigation'; +import type { PolicyContext } from '@/trigger/tasks/onboarding/onboard-organization-helpers'; +import { serverApi } from '@/lib/api-server'; +import { requireApiPermission } from '@/lib/permissions.server'; +import { db } from '@db/server'; +import { tasks as triggerTasks } from '@trigger.dev/sdk'; +import { NextRequest, NextResponse } from 'next/server'; + +interface PoliciesApiResponse { + data: Array<{ + id: string; + name: string; + description: string | null; + }>; +} + +/** + * Best-effort fan-out: re-trigger the risk mitigation generator so the saved + * treatment plan reflects the now-changed task linkage. We deliberately swallow + * errors here — the unlink itself already succeeded and refreshing the plan is + * not load-bearing for the user-facing operation. + */ +async function refreshTreatmentPlan(organizationId: string, riskId: string): Promise { + try { + const policiesResult = await serverApi.get('/v1/policies'); + const policyRows = policiesResult.data?.data ?? []; + const policies: PolicyContext[] = policyRows.map((policy) => ({ + name: policy.name, + description: policy.description, + })); + + await triggerTasks.trigger('generate-risk-mitigation', { + organizationId, + riskId, + policies, + }); + } catch (err) { + console.warn('Unlink succeeded but plan refresh failed to enqueue', { riskId, err }); + } +} + +/** + * DELETE /api/risks/[riskId]/tasks/[taskId] + * + * Soft-removes the link between a risk and a task by disconnecting the + * many-to-many join row. The task itself is not deleted. Controls remain + * derived through the remaining tasks, so removing the last task that + * references a given control implicitly removes it from the risk's view. + * + * After a successful unlink, fire-and-forget a re-generation of the + * treatment plan so the persisted citations reflect the new linkage. + */ +export async function DELETE( + req: NextRequest, + { params }: { params: Promise<{ riskId: string; taskId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'risk', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { riskId, taskId } = await params; + if (!riskId || !taskId) { + return NextResponse.json( + { error: 'Risk ID and Task ID are required' }, + { status: 400 }, + ); + } + + // Verify the risk + the link in one query, scoped to the active org. + // Without this, calling DELETE for a non-linked or wrong-tenant task + // would let Prisma's `disconnect` no-op or throw a 500 depending on + // the case — neither is a useful client signal. (Cubic #22.) + const risk = await db.risk.findUnique({ + where: { id: riskId }, + select: { + id: true, + organizationId: true, + tasks: { where: { id: taskId }, select: { id: true } }, + }, + }); + if (!risk || risk.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + if (risk.tasks.length === 0) { + return NextResponse.json( + { error: 'Task is not linked to this risk' }, + { status: 404 }, + ); + } + + await db.risk.update({ + where: { id: riskId }, + data: { tasks: { disconnect: { id: taskId } } }, + }); + + // Fire-and-forget: do NOT await. The unlink itself already succeeded; + // we don't want the response to wait on (or fail because of) the + // background plan-refresh trigger. (Cubic #30.) + void refreshTreatmentPlan(organizationId, riskId); + + return NextResponse.json({ ok: true }); + } catch (error) { + console.error('Error unlinking task from risk:', error); + return NextResponse.json({ error: 'Failed to unlink task' }, { status: 500 }); + } +} diff --git a/apps/app/src/app/api/vendors/[vendorId]/auto-link/active/route.ts b/apps/app/src/app/api/vendors/[vendorId]/auto-link/active/route.ts new file mode 100644 index 0000000000..90d5979a65 --- /dev/null +++ b/apps/app/src/app/api/vendors/[vendorId]/auto-link/active/route.ts @@ -0,0 +1,103 @@ +import { requireApiPermission } from '@/lib/permissions.server'; +import { db } from '@db/server'; +import { auth as triggerAuth } from '@trigger.dev/sdk'; +import { NextRequest, NextResponse } from 'next/server'; + +/** + * GET /api/vendors/[vendorId]/auto-link/active — see Risk counterpart. + * DELETE clears the active run (Discard). + */ + +/** + * Decide whether a `createPublicToken` failure means the run is permanently + * gone (drop the runId) or just transient (keep it, surface a 502). See the + * risk counterpart and Cubic findings #4/#5 on PR #2671. + */ +function isRunGoneError(err: unknown): boolean { + if (!err) return false; + const e = err as { status?: number; statusCode?: number; message?: string }; + const status = e.status ?? e.statusCode; + if (status === 404) return true; + const msg = (e.message ?? String(err)).toLowerCase(); + return msg.includes('not found') || msg.includes('purged'); +} + +export async function GET( + req: NextRequest, + { params }: { params: Promise<{ vendorId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'vendor', 'read'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { vendorId } = await params; + + const vendor = await db.vendor.findUnique({ + where: { id: vendorId }, + select: { id: true, organizationId: true, autoLinkRunId: true }, + }); + if (!vendor || vendor.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + if (!vendor.autoLinkRunId) { + return NextResponse.json({ runId: null }); + } + + try { + const publicAccessToken = await triggerAuth.createPublicToken({ + scopes: { read: { runs: [vendor.autoLinkRunId] } }, + expirationTime: '15m', + }); + return NextResponse.json({ runId: vendor.autoLinkRunId, publicAccessToken }); + } catch (err) { + if (isRunGoneError(err)) { + console.warn('[auto-link/active] run gone, clearing stored runId', err); + await db.vendor.update({ + where: { id: vendorId }, + data: { autoLinkRunId: null, autoLinkRunStartedAt: null }, + }); + return NextResponse.json({ runId: null }); + } + console.error('[auto-link/active] transient token mint failure', err); + return NextResponse.json( + { error: 'Failed to mint access token; try again' }, + { status: 502 }, + ); + } + } catch (error) { + console.error('Error reading active vendor auto-link run:', error); + return NextResponse.json({ error: 'Failed to read active run' }, { status: 500 }); + } +} + +export async function DELETE( + req: NextRequest, + { params }: { params: Promise<{ vendorId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'vendor', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { vendorId } = await params; + + const vendor = await db.vendor.findUnique({ + where: { id: vendorId }, + select: { id: true, organizationId: true }, + }); + if (!vendor || vendor.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + await db.vendor.update({ + where: { id: vendorId }, + data: { autoLinkRunId: null, autoLinkRunStartedAt: null }, + }); + return NextResponse.json({ ok: true }); + } catch (error) { + console.error('Error discarding vendor auto-link run:', error); + return NextResponse.json({ error: 'Failed to discard run' }, { status: 500 }); + } +} diff --git a/apps/app/src/app/api/vendors/[vendorId]/auto-link/apply/route.ts b/apps/app/src/app/api/vendors/[vendorId]/auto-link/apply/route.ts new file mode 100644 index 0000000000..03b93c236e --- /dev/null +++ b/apps/app/src/app/api/vendors/[vendorId]/auto-link/apply/route.ts @@ -0,0 +1,93 @@ +import { requireApiPermission } from '@/lib/permissions.server'; +import { db } from '@db/server'; +import { NextRequest, NextResponse } from 'next/server'; +import { z } from 'zod'; + +const ApplyBodySchema = z.object({ + taskIds: z.array(z.string()).max(100), + replace: z.boolean().optional().default(false), +}); + +/** + * POST /api/vendors/[vendorId]/auto-link/apply + * + * Persists the user-confirmed task selection from the AI-suggestion review UI. + * + * - `replace: true` → re-assess flow (sync semantics: connect-only-these tasks). + * - `replace: false` → fresh suggest flow (additive). + */ +export async function POST( + req: NextRequest, + { params }: { params: Promise<{ vendorId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'vendor', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { vendorId } = await params; + if (!vendorId) { + return NextResponse.json({ error: 'Vendor ID is required' }, { status: 400 }); + } + + const vendor = await db.vendor.findUnique({ + where: { id: vendorId }, + select: { id: true, organizationId: true }, + }); + if (!vendor || vendor.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + const body = await req.json().catch(() => null); + const parsed = ApplyBodySchema.safeParse(body); + if (!parsed.success) { + return NextResponse.json( + { error: 'Invalid body', details: parsed.error.flatten() }, + { status: 400 }, + ); + } + const { taskIds, replace } = parsed.data; + + // Validate tenancy of every taskId — see risks/auto-link/apply/route.ts + // for the same rationale (Prisma `connect`/`set` doesn't enforce it). + if (taskIds.length > 0) { + const ownedTasks = await db.task.findMany({ + where: { id: { in: taskIds }, organizationId }, + select: { id: true }, + }); + if (ownedTasks.length !== taskIds.length) { + return NextResponse.json( + { error: 'One or more tasks do not belong to this organization' }, + { status: 400 }, + ); + } + } + + if (replace) { + await db.vendor.update({ + where: { id: vendorId }, + data: { + tasks: { set: taskIds.map((id) => ({ id })) }, + autoLinkRunId: null, + autoLinkRunStartedAt: null, + }, + }); + } else { + await db.vendor.update({ + where: { id: vendorId }, + data: { + ...(taskIds.length > 0 + ? { tasks: { connect: taskIds.map((id) => ({ id })) } } + : {}), + autoLinkRunId: null, + autoLinkRunStartedAt: null, + }, + }); + } + + return NextResponse.json({ linked: taskIds.length }); + } catch (error) { + console.error('Error applying vendor auto-link:', error); + return NextResponse.json({ error: 'Failed to apply auto-link' }, { status: 500 }); + } +} diff --git a/apps/app/src/app/api/vendors/[vendorId]/auto-link/route.ts b/apps/app/src/app/api/vendors/[vendorId]/auto-link/route.ts new file mode 100644 index 0000000000..7791a6d062 --- /dev/null +++ b/apps/app/src/app/api/vendors/[vendorId]/auto-link/route.ts @@ -0,0 +1,68 @@ +import { requireApiPermission } from '@/lib/permissions.server'; +import type { linkRisksAndVendorsToWork } from '@/trigger/tasks/onboarding/link-risks-and-vendors-to-work'; +import { db } from '@db/server'; +import { auth as triggerAuth, tasks } from '@trigger.dev/sdk'; +import { NextRequest, NextResponse } from 'next/server'; + +/** + * POST /api/vendors/[vendorId]/auto-link + * + * Triggers the linkage task in `suggestionsOnly` mode for one vendor and + * returns a public-access token so the frontend can subscribe via + * `useRealtimeRun` and read `run.output.suggestions` once complete. + * + * No DB writes happen here — the user reviews the suggestions in the UI and + * the apply endpoint (`/auto-link/apply`) persists their final selection. + */ +export async function POST( + req: NextRequest, + { params }: { params: Promise<{ vendorId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'vendor', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { vendorId } = await params; + if (!vendorId) { + return NextResponse.json( + { error: 'Vendor ID is required' }, + { status: 400 }, + ); + } + + const vendor = await db.vendor.findUnique({ + where: { id: vendorId }, + select: { id: true, organizationId: true }, + }); + + if (!vendor || vendor.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + const handle = await tasks.trigger( + 'link-risks-and-vendors-to-work', + { organizationId, vendorId, suggestionsOnly: true }, + ); + + // Persist the runId so the UI can resume an in-flight scan after a page + // reload. Cleared by /apply or /discard once the user reviews the result. + await db.vendor.update({ + where: { id: vendorId }, + data: { + autoLinkRunId: handle.id, + autoLinkRunStartedAt: new Date(), + }, + }); + + const publicAccessToken = await triggerAuth.createPublicToken({ + scopes: { read: { runs: [handle.id] } }, + expirationTime: '15m', + }); + + return NextResponse.json({ runId: handle.id, publicAccessToken }); + } catch (error) { + console.error('Error triggering vendor auto-link:', error); + return NextResponse.json({ error: 'Failed to trigger auto-link' }, { status: 500 }); + } +} diff --git a/apps/app/src/app/api/vendors/[vendorId]/regenerate-mitigation/route.ts b/apps/app/src/app/api/vendors/[vendorId]/regenerate-mitigation/route.ts index 8ea0b7dd54..ac604ba614 100644 --- a/apps/app/src/app/api/vendors/[vendorId]/regenerate-mitigation/route.ts +++ b/apps/app/src/app/api/vendors/[vendorId]/regenerate-mitigation/route.ts @@ -1,8 +1,8 @@ import { generateVendorMitigation } from '@/trigger/tasks/onboarding/generate-vendor-mitigation'; import type { PolicyContext } from '@/trigger/tasks/onboarding/onboard-organization-helpers'; import { serverApi } from '@/lib/api-server'; -import { auth } from '@/utils/auth'; -import { tasks } from '@trigger.dev/sdk'; +import { requireApiPermission } from '@/lib/permissions.server'; +import { auth as triggerAuth, tasks } from '@trigger.dev/sdk'; import { NextRequest, NextResponse } from 'next/server'; interface PeopleApiResponse { @@ -27,13 +27,9 @@ export async function POST( { params }: { params: Promise<{ vendorId: string }> }, ) { try { - const session = await auth.api.getSession({ - headers: req.headers, - }); - - if (!session?.session?.activeOrganizationId) { - return NextResponse.json({ error: 'Unauthorized' }, { status: 401 }); - } + const ctx = await requireApiPermission(req, 'vendor', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; const { vendorId } = await params; if (!vendorId) { @@ -43,8 +39,6 @@ export async function POST( ); } - const organizationId = session.session.activeOrganizationId; - const [peopleResult, policiesResult] = await Promise.all([ serverApi.get('/v1/people'), serverApi.get('/v1/policies'), @@ -71,7 +65,7 @@ export async function POST( description: policy.description, })); - await tasks.trigger( + const handle = await tasks.trigger( 'generate-vendor-mitigation', { organizationId, @@ -81,17 +75,24 @@ export async function POST( }, ); - return NextResponse.json({ success: true }); + // See risks/regenerate-mitigation: don't fail the request when only the + // token mint fails, since the run is already in flight. (Cubic #29.) + let publicAccessToken: string | null = null; + try { + publicAccessToken = await triggerAuth.createPublicToken({ + scopes: { read: { runs: [handle.id] } }, + expirationTime: '15m', + }); + } catch (mintErr) { + console.error( + '[regenerate-mitigation] vendor run triggered but token mint failed; client must resume via /active', + { runId: handle.id, mintErr }, + ); + } + + return NextResponse.json({ runId: handle.id, publicAccessToken }); } catch (error) { console.error('Error regenerating vendor mitigation:', error); - return NextResponse.json( - { - error: - error instanceof Error - ? error.message - : 'Failed to regenerate mitigation', - }, - { status: 500 }, - ); + return NextResponse.json({ error: 'Failed to regenerate mitigation' }, { status: 500 }); } } diff --git a/apps/app/src/app/api/vendors/[vendorId]/relink/route.ts b/apps/app/src/app/api/vendors/[vendorId]/relink/route.ts new file mode 100644 index 0000000000..bed99e4c51 --- /dev/null +++ b/apps/app/src/app/api/vendors/[vendorId]/relink/route.ts @@ -0,0 +1,57 @@ +import { requireApiPermission } from '@/lib/permissions.server'; +import type { linkRisksAndVendorsToWork } from '@/trigger/tasks/onboarding/link-risks-and-vendors-to-work'; +import { db } from '@db/server'; +import { auth as triggerAuth, tasks } from '@trigger.dev/sdk'; +import { NextRequest, NextResponse } from 'next/server'; + +/** + * POST /api/vendors/[vendorId]/relink + * + * Wipes ALL current task links on the vendor and rebuilds linkage from the + * top-K embedding-similar tasks. Destructive — clears user's manual unlinks. + * Frontend confirms before calling. + * + * Returns { runId, publicAccessToken } so the UI can subscribe via useRealtimeRun. + */ +export async function POST( + req: NextRequest, + { params }: { params: Promise<{ vendorId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'vendor', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { vendorId } = await params; + if (!vendorId) { + return NextResponse.json( + { error: 'Vendor ID is required' }, + { status: 400 }, + ); + } + + const vendor = await db.vendor.findUnique({ + where: { id: vendorId }, + select: { id: true, organizationId: true }, + }); + + if (!vendor || vendor.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + + const handle = await tasks.trigger( + 'link-risks-and-vendors-to-work', + { organizationId, vendorId, replace: true }, + ); + + const publicAccessToken = await triggerAuth.createPublicToken({ + scopes: { read: { runs: [handle.id] } }, + expirationTime: '15m', + }); + + return NextResponse.json({ runId: handle.id, publicAccessToken }); + } catch (error) { + console.error('Error triggering vendor relink:', error); + return NextResponse.json({ error: 'Failed to trigger relink' }, { status: 500 }); + } +} diff --git a/apps/app/src/app/api/vendors/[vendorId]/tasks/[taskId]/route.ts b/apps/app/src/app/api/vendors/[vendorId]/tasks/[taskId]/route.ts new file mode 100644 index 0000000000..ddd390670a --- /dev/null +++ b/apps/app/src/app/api/vendors/[vendorId]/tasks/[taskId]/route.ts @@ -0,0 +1,104 @@ +import { generateVendorMitigation } from '@/trigger/tasks/onboarding/generate-vendor-mitigation'; +import type { PolicyContext } from '@/trigger/tasks/onboarding/onboard-organization-helpers'; +import { serverApi } from '@/lib/api-server'; +import { requireApiPermission } from '@/lib/permissions.server'; +import { db } from '@db/server'; +import { tasks as triggerTasks } from '@trigger.dev/sdk'; +import { NextRequest, NextResponse } from 'next/server'; + +interface PoliciesApiResponse { + data: Array<{ + id: string; + name: string; + description: string | null; + }>; +} + +/** + * Best-effort fan-out: re-trigger the vendor mitigation generator so the saved + * treatment plan reflects the now-changed task linkage. We deliberately swallow + * errors here — the unlink itself already succeeded. + */ +async function refreshVendorTreatmentPlan( + organizationId: string, + vendorId: string, +): Promise { + try { + const policiesResult = await serverApi.get('/v1/policies'); + const policyRows = policiesResult.data?.data ?? []; + const policies: PolicyContext[] = policyRows.map((policy) => ({ + name: policy.name, + description: policy.description, + })); + + await triggerTasks.trigger('generate-vendor-mitigation', { + organizationId, + vendorId, + policies, + }); + } catch (err) { + console.warn('Vendor unlink succeeded but plan refresh failed to enqueue', { vendorId, err }); + } +} + +/** + * DELETE /api/vendors/[vendorId]/tasks/[taskId] + * + * Soft-removes the link between a vendor and a task by disconnecting the + * many-to-many join row. The task itself is not deleted. Controls remain + * derived through the remaining tasks. + * + * After a successful unlink, fire-and-forget a re-generation of the + * treatment plan so the persisted citations reflect the new linkage. + */ +export async function DELETE( + req: NextRequest, + { params }: { params: Promise<{ vendorId: string; taskId: string }> }, +) { + try { + const ctx = await requireApiPermission(req, 'vendor', 'update'); + if (ctx instanceof NextResponse) return ctx; + const { organizationId } = ctx; + + const { vendorId, taskId } = await params; + if (!vendorId || !taskId) { + return NextResponse.json( + { error: 'Vendor ID and Task ID are required' }, + { status: 400 }, + ); + } + + // Verify the vendor + the link in one query, scoped to the active org. + // (Cubic #22.) + const vendor = await db.vendor.findUnique({ + where: { id: vendorId }, + select: { + id: true, + organizationId: true, + tasks: { where: { id: taskId }, select: { id: true } }, + }, + }); + if (!vendor || vendor.organizationId !== organizationId) { + return NextResponse.json({ error: 'Not found' }, { status: 404 }); + } + if (vendor.tasks.length === 0) { + return NextResponse.json( + { error: 'Task is not linked to this vendor' }, + { status: 404 }, + ); + } + + await db.vendor.update({ + where: { id: vendorId }, + data: { tasks: { disconnect: { id: taskId } } }, + }); + + // Fire-and-forget — see risks counterpart. (Cubic #31.) + void refreshVendorTreatmentPlan(organizationId, vendorId); + + return NextResponse.json({ ok: true }); + } catch (error) { + console.error('Error unlinking task from vendor:', error); + return NextResponse.json({ error: 'Failed to unlink task' }, { status: 500 }); + } +} diff --git a/apps/app/src/components/risks/RiskScoreBadge.tsx b/apps/app/src/components/risks/RiskScoreBadge.tsx index 5cb0f52811..b0778822c5 100644 --- a/apps/app/src/components/risks/RiskScoreBadge.tsx +++ b/apps/app/src/components/risks/RiskScoreBadge.tsx @@ -1,42 +1,59 @@ import { cn } from '@/lib/utils'; +import { + LEVEL_COLOR, + LEVEL_LABEL, + getRiskLevelFromScore, + getRiskScore, +} from '@/lib/risk-score'; import type { Impact, Likelihood } from '@db'; -import { getRiskScore, type RiskLevel } from '@/lib/risk-score'; - -const LEVEL_CLASSES: Record = { - 'very-low': - 'bg-emerald-500/15 border-emerald-500/40 text-emerald-700 dark:text-emerald-300', - low: 'bg-green-500/15 border-green-500/40 text-green-700 dark:text-green-300', - medium: 'bg-yellow-500/15 border-yellow-600/40 text-yellow-700 dark:text-yellow-300', - high: 'bg-orange-500/15 border-orange-500/40 text-orange-700 dark:text-orange-300', - 'very-high': 'bg-red-500/15 border-red-500/40 text-red-700 dark:text-red-300', -}; - -const LEVEL_LABEL: Record = { - 'very-low': 'Very low', - low: 'Low', - medium: 'Medium', - high: 'High', - 'very-high': 'Very high', -}; export interface RiskScoreBadgeProps { - likelihood: Likelihood; - impact: Impact; + /** + * Provide a precomputed 1-10 score directly, or pass `likelihood` + `impact` + * to have the badge derive it via `getRiskScore`. The score-derived path is + * what callers use when they want a current/interpolated value. + */ + score?: number; + likelihood?: Likelihood; + impact?: Impact; + /** + * When true, renders the level label (Low / Medium / High / etc.) instead + * of the score numeral. Same color treatment either way. Useful for the + * Risks list "Severity" column where a qualitative label scans faster + * than a number. + */ + labelOnly?: boolean; className?: string; } -export function RiskScoreBadge({ likelihood, impact, className }: RiskScoreBadgeProps) { - const { score, level } = getRiskScore(likelihood, impact); +export function RiskScoreBadge({ + score, + likelihood, + impact, + labelOnly, + className, +}: RiskScoreBadgeProps) { + const resolvedScore = + score ?? (likelihood && impact ? getRiskScore(likelihood, impact).score : 1); + const level = getRiskLevelFromScore(resolvedScore); return ( - {score}/10 + {labelOnly ? LEVEL_LABEL[level] : `${resolvedScore}/10`} ); } diff --git a/apps/app/src/components/risks/charts/AxisTooltip.tsx b/apps/app/src/components/risks/charts/AxisTooltip.tsx new file mode 100644 index 0000000000..25a4b1943b --- /dev/null +++ b/apps/app/src/components/risks/charts/AxisTooltip.tsx @@ -0,0 +1,29 @@ +'use client'; + +// TODO(design-system): migrate to @trycompai/design-system when Tooltip ships. +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '@trycompai/ui/tooltip'; +import type { ReactNode } from 'react'; + +interface AxisTooltipProps { + label: ReactNode; + definition: string; +} + +/** + * Wraps an axis tier label (or info icon) with a hover definition, + * e.g. "Insignificant: no material impact on operations, customers, or compliance." + */ +export function AxisTooltip({ label, definition }: AxisTooltipProps) { + return ( + + + + + {label} + + + {definition} + + + ); +} diff --git a/apps/app/src/components/risks/charts/InherentRiskChart.tsx b/apps/app/src/components/risks/charts/InherentRiskChart.tsx index e8cdc9c050..56cec0ae7b 100644 --- a/apps/app/src/components/risks/charts/InherentRiskChart.tsx +++ b/apps/app/src/components/risks/charts/InherentRiskChart.tsx @@ -19,6 +19,7 @@ export function InherentRiskChart({ risk }: InherentRiskChartProps) { = { + 'Very Likely': 'Expected to occur multiple times per year without mitigations.', + Likely: 'Expected to occur at least once per year without mitigations.', + Possible: 'Could occur in the next 1–3 years.', + Unlikely: 'Could occur but not expected in the next 3 years.', + 'Very Unlikely': 'Rare event; theoretical possibility only.', +}; + +const IMPACT_DEFINITIONS: Record = { + Insignificant: 'No material impact on operations, customers, or compliance.', + Minor: 'Contained, short-term operational issue; limited customer or compliance exposure.', + Moderate: 'Operational disruption, customer complaint, or contained compliance finding.', + Major: 'Significant customer impact, regulator scrutiny, or material revenue loss.', + Severe: 'Existential or multi-year impact: breach, enforcement action, or loss of trust.', +}; + +interface RiskCell { + probability: string; + impact: string; + level: 'very-low' | 'low' | 'medium' | 'high' | 'very-high'; + value?: number; +} + +const getRiskColor = (level: string, readOnly?: boolean) => { + switch (level) { + case 'very-low': + return `bg-emerald-500/20 border-emerald-500/30${readOnly ? '' : ' hover:bg-emerald-500/30'}`; + case 'low': + return `bg-green-500/20 border-green-500/30${readOnly ? '' : ' hover:bg-green-500/30'}`; + case 'medium': + return `bg-yellow-500/20 border-yellow-500/30${readOnly ? '' : ' hover:bg-yellow-500/30'}`; + case 'high': + return `bg-orange-500/20 border-orange-500/30${readOnly ? '' : ' hover:bg-orange-500/30'}`; + case 'very-high': + return `bg-red-500/20 border-red-500/30${readOnly ? '' : ' hover:bg-red-500/30'}`; + default: + return 'bg-slate-500/20 border-slate-500/30'; + } +}; + +export function buildRiskData( + activeLikelihood: Likelihood, + activeImpact: Impact, +): RiskCell[] { + const activeProbability = probabilityLevels[VISUAL_LIKELIHOOD_ORDER.indexOf(activeLikelihood)]; + const activeImpactLevel = impactLevels[VISUAL_IMPACT_ORDER.indexOf(activeImpact)]; + + return probabilityLevels.flatMap((probability) => + impactLevels.map((impact) => { + const likelihoodScore = + LIKELIHOOD_SCORES[VISUAL_LIKELIHOOD_ORDER[probabilityLevels.indexOf(probability)]]; + const impactScore = IMPACT_SCORES[VISUAL_IMPACT_ORDER[impactLevels.indexOf(impact)]]; + const level = getRiskLevel(likelihoodScore * impactScore); + + return { + probability, + impact, + level, + value: probability === activeProbability && impact === activeImpactLevel ? 1 : undefined, + }; + }), + ); +} + +interface MatrixBodyProps { + readOnly?: boolean; + riskData: RiskCell[]; + handleCellClick: (probability: string, impact: string) => void; + suggestedLikelihood?: Likelihood; + suggestedImpact?: Impact; +} + +export function MatrixBody({ + readOnly, + riskData, + handleCellClick, + suggestedLikelihood, + suggestedImpact, +}: MatrixBodyProps) { + return ( +
+
+
+ Impact +
+
+ {probabilityLevels.map((probability, rowIdx) => ( +
+
+ + + +
+ {impactLevels.map((impact, colIdx) => { + const cell = riskData.find( + (item) => item.probability === probability && item.impact === impact, + ); + let rounding = ''; + if (rowIdx === 0 && colIdx === 0) rounding = 'rounded-tl-lg'; + if (rowIdx === 0 && colIdx === impactLevels.length - 1) + rounding = 'rounded-tr-lg'; + if (rowIdx === probabilityLevels.length - 1 && colIdx === 0) + rounding = 'rounded-bl-lg'; + if ( + rowIdx === probabilityLevels.length - 1 && + colIdx === impactLevels.length - 1 + ) + rounding = 'rounded-br-lg'; + const isSuggested = + suggestedLikelihood !== undefined && + suggestedImpact !== undefined && + VISUAL_LIKELIHOOD_ORDER[rowIdx] === suggestedLikelihood && + VISUAL_IMPACT_ORDER[colIdx] === suggestedImpact; + return ( +
handleCellClick(probability, impact)} + > + {cell?.value && ( +
+ )} + {isSuggested && !cell?.value && ( +
+ )} +
+ ); + })} +
+ ))} +
+
+
+ {impactLevels.map((impact) => ( +
+ + + +
+ ))} +
+
+
+ Likelihood +
+
+ ); +} diff --git a/apps/app/src/components/risks/charts/MatrixLegend.tsx b/apps/app/src/components/risks/charts/MatrixLegend.tsx new file mode 100644 index 0000000000..6bd5995f29 --- /dev/null +++ b/apps/app/src/components/risks/charts/MatrixLegend.tsx @@ -0,0 +1,38 @@ +'use client'; + +import { HStack, Stack, Text } from '@trycompai/design-system'; + +interface LegendItem { + level: 'very-low' | 'low' | 'medium' | 'high' | 'very-high'; + label: string; + range: string; + color: string; +} + +const ITEMS: LegendItem[] = [ + { level: 'very-low', label: 'Very Low', range: 'Raw 1', color: 'bg-emerald-500/30' }, + { level: 'low', label: 'Low', range: 'Raw 2–4', color: 'bg-green-500/30' }, + { level: 'medium', label: 'Medium', range: 'Raw 5–9', color: 'bg-yellow-500/30' }, + { level: 'high', label: 'High', range: 'Raw 10–16', color: 'bg-orange-500/30' }, + { level: 'very-high', label: 'Very High', range: 'Raw 17–25', color: 'bg-red-500/30' }, +]; + +export function MatrixLegend() { + return ( + + {ITEMS.map((item) => ( +
+ + + + {item.label} + + + {item.range} + + +
+ ))} +
+ ); +} diff --git a/apps/app/src/components/risks/charts/ResidualRiskChart.test.tsx b/apps/app/src/components/risks/charts/ResidualRiskChart.test.tsx index 537a970450..ca8c62ab2c 100644 --- a/apps/app/src/components/risks/charts/ResidualRiskChart.test.tsx +++ b/apps/app/src/components/risks/charts/ResidualRiskChart.test.tsx @@ -42,8 +42,12 @@ import { ResidualRiskChart } from './ResidualRiskChart'; const mockRisk: any = { id: 'risk-1', + likelihood: 'possible', + impact: 'moderate', residualLikelihood: 'unlikely', residualImpact: 'minor', + treatmentStrategy: 'accept', + tasks: [], }; describe('ResidualRiskChart permission gating', () => { @@ -86,7 +90,7 @@ describe('ResidualRiskChart permission gating', () => { expect(capturedProps.title).toBe('Residual Risk'); expect(capturedProps.description).toBe( - 'Remaining risk level after controls are applied', + 'Risk level after the treatment plan is applied. The dashed cell is the suggestion computed from your strategy and linked task completion.', ); expect(capturedProps.riskId).toBe('risk-1'); expect(capturedProps.activeLikelihood).toBe('unlikely'); diff --git a/apps/app/src/components/risks/charts/ResidualRiskChart.tsx b/apps/app/src/components/risks/charts/ResidualRiskChart.tsx index 14c26290db..00c986aa0d 100644 --- a/apps/app/src/components/risks/charts/ResidualRiskChart.tsx +++ b/apps/app/src/components/risks/charts/ResidualRiskChart.tsx @@ -2,12 +2,13 @@ import { usePermissions } from '@/hooks/use-permissions'; import { useRiskActions } from '@/hooks/use-risks'; -import type { Risk } from '@db'; +import { suggestedResidual } from '@/lib/suggested-residual'; +import type { Risk, TaskStatus } from '@db'; import { useSWRConfig } from 'swr'; import { RiskMatrixChart } from './RiskMatrixChart'; interface ResidualRiskChartProps { - risk: Risk; + risk: Risk & { tasks?: { status: TaskStatus }[] }; } export function ResidualRiskChart({ risk }: ResidualRiskChartProps) { @@ -15,24 +16,39 @@ export function ResidualRiskChart({ risk }: ResidualRiskChartProps) { const { mutate: globalMutate } = useSWRConfig(); const { hasPermission } = usePermissions(); + // Only compute a suggestion when tasks are actually loaded — falling back to + // [] would render a misleading "0% complete" ghost cell on orgs that haven't + // hydrated yet. + const suggestion = risk.tasks + ? suggestedResidual({ + likelihood: risk.likelihood, + impact: risk.impact, + strategy: risk.treatmentStrategy, + tasks: risk.tasks, + }) + : undefined; + return ( { await updateRisk(id, { residualLikelihood: probability, residualImpact: impact, }); - globalMutate( - (key) => Array.isArray(key) && key[0]?.includes('/v1/risks'), - undefined, - { revalidate: true }, - ); + globalMutate((key) => Array.isArray(key) && key[0]?.includes('/v1/risks'), undefined, { + revalidate: true, + }); }} /> ); diff --git a/apps/app/src/components/risks/charts/RiskMatrixChart.spec.tsx b/apps/app/src/components/risks/charts/RiskMatrixChart.spec.tsx new file mode 100644 index 0000000000..82c6a49d58 --- /dev/null +++ b/apps/app/src/components/risks/charts/RiskMatrixChart.spec.tsx @@ -0,0 +1,117 @@ +import { Impact, Likelihood } from '@db'; +import { fireEvent, render, screen } from '@testing-library/react'; +import { describe, expect, it, vi } from 'vitest'; +import { RiskMatrixChart } from './RiskMatrixChart'; + +describe('RiskMatrixChart', () => { + it('renders the legend', () => { + render( + , + ); + expect(screen.getByText('Very Low')).toBeInTheDocument(); + expect(screen.getByText('Very High')).toBeInTheDocument(); + }); + + it('renders the Accept-suggested button only when a suggestion differs from active', () => { + const saveAction = vi.fn(); + const { rerender } = render( + , + ); + expect(screen.getByRole('button', { name: /Accept suggested residual/i })).toBeInTheDocument(); + + rerender( + , + ); + expect(screen.queryByRole('button', { name: /Accept suggested residual/i })).toBeNull(); + }); + + it('Accept-suggested snaps the active cell to the suggested coords', () => { + const saveAction = vi.fn(); + render( + , + ); + fireEvent.click(screen.getByRole('button', { name: /Accept suggested residual/i })); + expect(screen.getByRole('button', { name: 'Save' })).toBeInTheDocument(); + }); + + it('renders no ghost marker when suggestion matches active cell', () => { + render( + , + ); + expect(document.querySelectorAll('.border-dashed').length).toBe(0); + }); + + describe('preliminary subtitle', () => { + it('does not render the preliminary subtitle by default', () => { + render( + , + ); + expect(screen.queryByText(/Preliminary/i)).toBeNull(); + }); + + it('renders the preliminary subtitle when preliminary=true', () => { + render( + , + ); + expect(screen.getByText(/Preliminary — assessment still running/i)).toBeInTheDocument(); + }); + }); +}); diff --git a/apps/app/src/components/risks/charts/RiskMatrixChart.tsx b/apps/app/src/components/risks/charts/RiskMatrixChart.tsx index 7853115fd3..7f48b2cc15 100644 --- a/apps/app/src/components/risks/charts/RiskMatrixChart.tsx +++ b/apps/app/src/components/risks/charts/RiskMatrixChart.tsx @@ -1,60 +1,19 @@ 'use client'; -import { IMPACT_SCORES, LIKELIHOOD_SCORES, getRiskLevel } from '@/lib/risk-score'; import { Impact, Likelihood } from '@db'; -import { Button, HStack, Section } from '@trycompai/design-system'; +import { Button, HStack, Section, Text } from '@trycompai/design-system'; +import { Information } from '@trycompai/design-system/icons'; import { useEffect, useState } from 'react'; - -const VISUAL_LIKELIHOOD_ORDER: Likelihood[] = [ - Likelihood.very_likely, - Likelihood.likely, - Likelihood.possible, - Likelihood.unlikely, - Likelihood.very_unlikely, -]; -const VISUAL_IMPACT_ORDER: Impact[] = [ - Impact.insignificant, - Impact.minor, - Impact.moderate, - Impact.major, - Impact.severe, -]; - -interface RiskCell { - probability: string; - impact: string; - level: 'very-low' | 'low' | 'medium' | 'high' | 'very-high'; - value?: number; -} - -const getRiskColor = (level: string, readOnly?: boolean) => { - switch (level) { - case 'very-low': - return `bg-emerald-500/20 border-emerald-500/30${readOnly ? '' : ' hover:bg-emerald-500/30'}`; - case 'low': - return `bg-green-500/20 border-green-500/30${readOnly ? '' : ' hover:bg-green-500/30'}`; - case 'medium': - return `bg-yellow-500/20 border-yellow-500/30${readOnly ? '' : ' hover:bg-yellow-500/30'}`; - case 'high': - return `bg-orange-500/20 border-orange-500/30${readOnly ? '' : ' hover:bg-orange-500/30'}`; - case 'very-high': - return `bg-red-500/20 border-red-500/30${readOnly ? '' : ' hover:bg-red-500/30'}`; - default: - return 'bg-slate-500/20 border-slate-500/30'; - } -}; - -const probabilityLevels = ['Very Likely', 'Likely', 'Possible', 'Unlikely', 'Very Unlikely']; -const probabilityNumbers = ['5', '4', '3', '2', '1']; -const probabilityLabels = [ - 'Very Likely (5)', - 'Likely (4)', - 'Possible (3)', - 'Unlikely (2)', - 'Very Unlikely (1)', -]; -const impactLevels = ['Insignificant', 'Minor', 'Moderate', 'Major', 'Severe']; -const impactNumbers = ['1', '2', '3', '4', '5']; +import { AxisTooltip } from './AxisTooltip'; +import { + MatrixBody, + VISUAL_IMPACT_ORDER, + VISUAL_LIKELIHOOD_ORDER, + buildRiskData, + impactLevels, + probabilityLevels, +} from './MatrixBody'; +import { MatrixLegend } from './MatrixLegend'; interface RiskMatrixChartProps { title: string; @@ -62,8 +21,16 @@ interface RiskMatrixChartProps { riskId: string; activeLikelihood: Likelihood; activeImpact: Impact; - saveAction: (data: { id: string; probability: Likelihood; impact: Impact }) => Promise; + saveAction: (data: { id: string; probability: Likelihood; impact: Impact }) => Promise; readOnly?: boolean; + /** If set, renders a pulsing dashed-outline cell as the "suggested" residual. */ + suggestedLikelihood?: Likelihood; + /** If set, renders a pulsing dashed-outline cell as the "suggested" residual. */ + suggestedImpact?: Impact; + /** Tooltip body shown on the title info icon. */ + titleInfo?: string; + /** When true, render a small "Preliminary — assessment still running" subtitle below the matrix. */ + preliminary?: boolean; } export function RiskMatrixChart({ @@ -74,6 +41,10 @@ export function RiskMatrixChart({ activeImpact: initialImpactProp, saveAction, readOnly, + suggestedLikelihood, + suggestedImpact, + titleInfo, + preliminary, }: RiskMatrixChartProps) { const [initialLikelihood, setInitialLikelihood] = useState(initialLikelihoodProp); const [initialImpact, setInitialImpact] = useState(initialImpactProp); @@ -90,34 +61,14 @@ export function RiskMatrixChart({ setActiveImpact(initialImpactProp); }, [initialImpactProp]); - const activeProbability = probabilityLevels[VISUAL_LIKELIHOOD_ORDER.indexOf(activeLikelihood)]; - const activeImpactLevel = impactLevels[VISUAL_IMPACT_ORDER.indexOf(activeImpact)]; - - // Create risk data - const riskData: RiskCell[] = probabilityLevels.flatMap((probability) => - impactLevels.map((impact) => { - const likelihoodScore = - LIKELIHOOD_SCORES[VISUAL_LIKELIHOOD_ORDER[probabilityLevels.indexOf(probability)]]; - const impactScore = IMPACT_SCORES[VISUAL_IMPACT_ORDER[impactLevels.indexOf(impact)]]; - const level = getRiskLevel(likelihoodScore * impactScore); - - return { - probability, - impact, - level, - value: probability === activeProbability && impact === activeImpactLevel ? 1 : undefined, - }; - }), - ); + const riskData = buildRiskData(activeLikelihood, activeImpact); const handleCellClick = (probability: string, impact: string) => { if (readOnly) return; const likelihoodIdx = probabilityLevels.indexOf(probability); const impactIdx = impactLevels.indexOf(impact); - const newLikelihood = VISUAL_LIKELIHOOD_ORDER[likelihoodIdx]; - const newImpact = VISUAL_IMPACT_ORDER[impactIdx]; - setActiveLikelihood(newLikelihood); - setActiveImpact(newImpact); + setActiveLikelihood(VISUAL_LIKELIHOOD_ORDER[likelihoodIdx]); + setActiveImpact(VISUAL_IMPACT_ORDER[impactIdx]); }; const hasChanges = activeLikelihood !== initialLikelihood || activeImpact !== initialImpact; @@ -132,81 +83,90 @@ export function RiskMatrixChart({ }); setInitialLikelihood(activeLikelihood); setInitialImpact(activeImpact); - } catch (e) { + } catch (_e) { } finally { setLoading(false); } }; - return ( -
- Save + const hasSuggestion = suggestedLikelihood !== undefined && suggestedImpact !== undefined; + const suggestionDiffers = + hasSuggestion && + (suggestedLikelihood !== activeLikelihood || suggestedImpact !== activeImpact); + + const handleAcceptSuggestion = () => { + if (!hasSuggestion) return; + setActiveLikelihood(suggestedLikelihood); + setActiveImpact(suggestedImpact); + }; + + const sectionActions = + !readOnly && hasChanges ? ( + + ) : undefined; + + const body = ( + <> + + {suggestionDiffers && !readOnly && ( +
+ - ) : undefined - } - > -
-
-
- Impact -
-
- {probabilityLevels.map((probability, rowIdx) => ( -
-
- {probabilityNumbers[rowIdx]} -
- {impactLevels.map((impact, colIdx) => { - const cell = riskData.find( - (item) => item.probability === probability && item.impact === impact, - ); - let rounding = ''; - if (rowIdx === 0 && colIdx === 0) rounding = 'rounded-tl-lg'; - if (rowIdx === 0 && colIdx === impactLevels.length - 1) - rounding = 'rounded-tr-lg'; - if (rowIdx === probabilityLevels.length - 1 && colIdx === 0) - rounding = 'rounded-bl-lg'; - if ( - rowIdx === probabilityLevels.length - 1 && - colIdx === impactLevels.length - 1 - ) - rounding = 'rounded-br-lg'; - return ( -
handleCellClick(probability, impact)} - > - {cell?.value && ( -
- )} -
- ); - })} -
- ))} -
-
-
- {impactLevels.map((impact) => ( -
- {impact} -
- ))} -
-
-
- Likelihood -
+ + Ghost cell shows the residual suggested by this entity's treatment plan. + +
+ )} + {preliminary && ( +
+ + Preliminary — assessment still running + +
+ )} +
+ +
+ + ); + + // When titleInfo is present, we render the title manually so we can inject an info icon. + // The DS Section's `title` prop only accepts `string`, so when titleInfo is set we omit + // `title` and render our own heading row. When no titleInfo, fall back to the normal prop. + if (titleInfo) { + return ( +
+
+ +

{title}

+ + } + definition={titleInfo} + /> +
+ {body} +
+ ); + } + + return ( +
+ {body}
); } diff --git a/apps/app/src/components/risks/treatment-plan/AutoLinkButton.spec.tsx b/apps/app/src/components/risks/treatment-plan/AutoLinkButton.spec.tsx new file mode 100644 index 0000000000..6e8c49aee8 --- /dev/null +++ b/apps/app/src/components/risks/treatment-plan/AutoLinkButton.spec.tsx @@ -0,0 +1,97 @@ +import { fireEvent, render, screen, waitFor } from '@testing-library/react'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { AutoLinkButton } from './AutoLinkButton'; + +vi.mock('sonner', () => ({ toast: { info: vi.fn(), success: vi.fn(), error: vi.fn() } })); + +const realtimeRunMock = vi.fn(); +vi.mock('@trigger.dev/react-hooks', () => ({ + useRealtimeRun: (runId: string, opts: { accessToken?: string; enabled?: boolean }) => + realtimeRunMock(runId, opts), +})); + +afterEach(() => { + realtimeRunMock.mockReset(); +}); + +describe('AutoLinkButton', () => { + it('shows "generate plan" label when description is empty', () => { + realtimeRunMock.mockReturnValue({ run: null }); + render( + , + ); + expect( + screen.getByRole('button', { name: /Auto-link tasks & generate plan/i }), + ).toBeInTheDocument(); + }); + + it('shows "refresh plan" label when description is non-empty', () => { + realtimeRunMock.mockReturnValue({ run: null }); + render( + , + ); + expect( + screen.getByRole('button', { name: /Auto-link tasks & refresh plan/i }), + ).toBeInTheDocument(); + }); + + it('replaces the button with a progress card after click', async () => { + realtimeRunMock.mockReturnValue({ + run: { status: 'EXECUTING', metadata: { phase: 'embedding-tasks', current: 5, total: 10 } }, + }); + const onAutoLink = vi.fn().mockResolvedValue({ runId: 'r1', publicAccessToken: 't1' }); + render(); + fireEvent.click(screen.getByRole('button')); + await waitFor(() => { + expect(screen.getByRole('status')).toBeInTheDocument(); + }); + expect(screen.getByText(/Embedding tasks/i)).toBeInTheDocument(); + expect(screen.getByText(/\(5\/10\)/i)).toBeInTheDocument(); + }); + + it('chains onAfterLink and toasts when run completes with links', async () => { + realtimeRunMock.mockReturnValue({ + run: { status: 'COMPLETED', metadata: { phase: 'done', riskLinks: 3, vendorLinks: 0 } }, + }); + const onAutoLink = vi.fn().mockResolvedValue({ runId: 'r1', publicAccessToken: 't1' }); + const onAfterLink = vi.fn().mockResolvedValue(undefined); + render(); + fireEvent.click(screen.getByRole('button')); + await waitFor(() => { + expect(onAfterLink).toHaveBeenCalled(); + }); + }); + + it('skips onAfterLink and toasts info when run completes with zero links', async () => { + realtimeRunMock.mockReturnValue({ + run: { status: 'COMPLETED', metadata: { phase: 'done', riskLinks: 0, vendorLinks: 0 } }, + }); + const onAutoLink = vi.fn().mockResolvedValue({ runId: 'r1', publicAccessToken: 't1' }); + const onAfterLink = vi.fn(); + render(); + fireEvent.click(screen.getByRole('button')); + await waitFor(() => { + // Component returns to idle (button visible again). + expect(screen.getByRole('button')).toBeInTheDocument(); + }); + expect(onAfterLink).not.toHaveBeenCalled(); + }); + + it('toasts error and returns to idle when run fails', async () => { + realtimeRunMock.mockReturnValue({ + run: { status: 'FAILED', metadata: {} }, + }); + const onAutoLink = vi.fn().mockResolvedValue({ runId: 'r1', publicAccessToken: 't1' }); + render(); + fireEvent.click(screen.getByRole('button')); + await waitFor(() => { + expect(screen.getByRole('button')).toBeInTheDocument(); + }); + }); +}); diff --git a/apps/app/src/components/risks/treatment-plan/AutoLinkButton.tsx b/apps/app/src/components/risks/treatment-plan/AutoLinkButton.tsx new file mode 100644 index 0000000000..1fef7b3f42 --- /dev/null +++ b/apps/app/src/components/risks/treatment-plan/AutoLinkButton.tsx @@ -0,0 +1,177 @@ +'use client'; + +import { Button, Stack, Text } from '@trycompai/design-system'; +import { Link as LinkIcon } from '@trycompai/design-system/icons'; +import { useRealtimeRun } from '@trigger.dev/react-hooks'; +import { useEffect, useState } from 'react'; +import { toast } from 'sonner'; + +interface AutoLinkButtonProps { + /** Whether the entity already has a saved treatment-plan description. */ + hasDescription: boolean; + disabled?: boolean; + onAutoLink: () => Promise<{ runId: string; publicAccessToken: string }>; + onAfterLink?: () => Promise; +} + +type RunState = + | { kind: 'idle' } + | { kind: 'running'; runId: string; publicAccessToken: string }; + +const PHASE_LABEL: Record = { + starting: 'Starting…', + 'embedding-tasks': 'Embedding tasks', + 'embedding-risks': 'Embedding risks', + 'embedding-vendors': 'Embedding vendors', + 'waiting-for-index': 'Waiting for the index', + 'matching-risks': 'Matching tasks to risks', + 'matching-vendors': 'Matching tasks to vendors', + done: 'Finishing up…', +}; + +export function AutoLinkButton({ + hasDescription, + disabled, + onAutoLink, + onAfterLink, +}: AutoLinkButtonProps) { + const [state, setState] = useState({ kind: 'idle' }); + const [submitting, setSubmitting] = useState(false); + + const handleClick = async () => { + setSubmitting(true); + try { + const { runId, publicAccessToken } = await onAutoLink(); + setState({ kind: 'running', runId, publicAccessToken }); + } catch { + toast.error('Auto-link failed. Try again or link manually.'); + } finally { + setSubmitting(false); + } + }; + + if (state.kind === 'running') { + return ( + { + if (linked === 0) { + toast.info('No matching tasks found. Link manually from the Tasks tab.'); + } else { + toast.success( + `Linked ${linked} task${linked === 1 ? '' : 's'}${ + hasDescription ? ' · refreshing treatment plan' : ' · generating treatment plan' + }`, + ); + // Post-link refresh (mitigation regen, swr revalidation) is + // separate from the link itself succeeding. If it fails we + // surface a distinct toast so the user knows the link landed + // but the plan refresh did not. Without this branch a + // refresh failure could either be silent or — worse — get + // mistaken for a link failure. (Cubic finding #23.) + if (onAfterLink) { + try { + await onAfterLink(); + } catch (err) { + console.error('[AutoLinkButton] post-link refresh failed', err); + toast.warning( + 'Linked the tasks, but refreshing the treatment plan failed. Try regenerating manually.', + ); + } + } + } + setState({ kind: 'idle' }); + }} + onFailed={() => { + toast.error('Auto-link failed. Try again or link manually.'); + setState({ kind: 'idle' }); + }} + /> + ); + } + + return ( + + ); +} + +interface RunProgressProps { + runId: string; + accessToken: string; + hasDescription: boolean; + onComplete: (linked: number) => void; + onFailed: () => void; +} + +function RunProgress({ runId, accessToken, onComplete, onFailed }: RunProgressProps) { + const { run } = useRealtimeRun(runId, { accessToken, enabled: true }); + + const meta = (run?.metadata ?? {}) as Record; + const phase = typeof meta.phase === 'string' ? meta.phase : 'starting'; + const current = typeof meta.current === 'number' ? meta.current : null; + const total = typeof meta.total === 'number' ? meta.total : null; + + const status = run?.status; + useEffect(() => { + if (!status) return; + if (status === 'COMPLETED') { + const riskLinks = typeof meta.riskLinks === 'number' ? meta.riskLinks : 0; + const vendorLinks = typeof meta.vendorLinks === 'number' ? meta.vendorLinks : 0; + onComplete(riskLinks + vendorLinks); + return; + } + if ( + status === 'FAILED' || + status === 'CANCELED' || + status === 'CRASHED' || + status === 'SYSTEM_FAILURE' || + status === 'EXPIRED' || + status === 'TIMED_OUT' + ) { + onFailed(); + } + }, [status]); // eslint-disable-line react-hooks/exhaustive-deps + + // Show progress phases (skip 'done' — onComplete handles transition). + const showProgress = phase !== 'done' && phase in PHASE_LABEL; + const label = showProgress ? PHASE_LABEL[phase] : PHASE_LABEL.starting; + const countSuffix = + current !== null && total !== null && total > 0 && phase !== 'starting' && phase !== 'done' + ? ` (${current}/${total})` + : ''; + + return ( +
+ +
+
+ + {label} + {countSuffix} + +
+
+ ); +} diff --git a/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.parts.tsx b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.parts.tsx new file mode 100644 index 0000000000..fd5e730a67 --- /dev/null +++ b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.parts.tsx @@ -0,0 +1,459 @@ +'use client'; + +import { Button } from '@trycompai/design-system'; +import { MagicWandFilled } from '@trycompai/design-system/icons'; +import { useRealtimeRun } from '@trigger.dev/react-hooks'; +import { useEffect, useMemo, useRef } from 'react'; +import { ControlsSection, TasksSection } from './AutoLinkSuggestions.sections'; +import { + isControlDerived, + type SuggestedControl, + type SuggestedTask, +} from './AutoLinkSuggestions.types'; + +export function EmptyState({ + canUpdate, + submitting, + onSuggest, + onStartFromScratch, + variant = 'default', +}: { + canUpdate: boolean; + submitting: boolean; + onSuggest: () => void; + /** Required when variant === 'kickoff' — dismisses the kickoff so the user + * can write the plan manually. Ignored for other variants. */ + onStartFromScratch?: () => void; + /** + * `'kickoff'` — the wide centered panel ("Let AI kick this off") shown when + * plan and tasks are both empty. Has a primary "Draft plan & suggest + * links" button plus a "Start from scratch" escape hatch. + * `'kickoff-with-plan'` — same wide panel, copy adapted for the case where + * a plan already exists ("AI will only suggest tasks/controls; your plan + * stays as-is unless you regenerate"). + * `'default'` — the smaller per-column empty CTA shown when the user has a + * plan but no linked tasks yet. + */ + variant?: 'default' | 'kickoff' | 'kickoff-with-plan'; +}) { + // Adapt kickoff copy based on whether a plan already exists. The same + // wide-panel layout applies in both cases. + const isKickoff = variant === 'kickoff' || variant === 'kickoff-with-plan'; + if (isKickoff) { + const hasPlan = variant === 'kickoff-with-plan'; + const title = hasPlan ? 'Let AI suggest tasks and controls' : 'Let AI kick this off'; + const description = hasPlan + ? "Based on your treatment plan, AI can scan your library and suggest the tasks and controls most likely to drive this risk down. You'll review everything before anything is linked." + : "Based on the strategy above, AI can draft a treatment plan and suggest the tasks and controls most likely to drive this risk down. You'll review everything before anything is saved or linked."; + const primaryLabel = hasPlan ? 'Suggest tasks & controls' : 'Draft plan & suggest links'; + const escapeLabel = hasPlan ? 'Edit plan manually' : 'Start from scratch'; + + return ( +
+
+
+

{title}

+

+ {description} +

+
+ + +
+
+
+
+ 02 · Plan +
+
+ {hasPlan + ? 'Your existing plan stays as-is unless you regenerate.' + : 'A concrete plan grounded in the selected tasks and controls.'} +
+
+
+
+ 03 · Links +
+
+ Tasks and framework controls AI ranks for this risk. +
+
+
+
+ ); + } + + return ( +
+
+
+
No tasks or controls linked yet
+
+ Have AI scan your library and suggest the tasks and controls most likely to drive this risk + down. +
+
+ +
+
+ You'll review before anything is linked. +
+
+ ); +} + +/** + * Statuses we treat as "in flight, render the loading UI". Anything outside + * this set is either terminal-success (COMPLETED) or terminal-failure (which + * we surface via FailedState). + */ +const IN_FLIGHT_STATUSES = new Set([ + 'WAITING_FOR_DEPLOY', + 'DELAYED', + 'QUEUED', + 'EXECUTING', + 'INTERRUPTED', + 'WAITING_TO_RESUME', +]); + +/** + * Terminal failures we route to FailedState. EXPIRED happens when a run sits + * in the queue past its TTL; TIMED_OUT when execution exceeds maxDuration. + */ +const TERMINAL_FAILURE_STATUSES = new Set([ + 'FAILED', + 'CANCELED', + 'CRASHED', + 'SYSTEM_FAILURE', + 'EXPIRED', + 'TIMED_OUT', +]); + +/** + * Phases written to `metadata.phase` by the trigger task wrapper. Mirrored + * from `LinkagePhase` in `lib/embedding/run-linkage.ts`. + */ +type PhaseName = + | 'starting' + | 'embedding-tasks' + | 'embedding-risks' + | 'embedding-vendors' + | 'waiting-for-index' + | 'matching-risks' + | 'matching-vendors' + | 'done'; + +function describeStatus( + status: string | undefined, + phase: PhaseName | undefined, + current: number | undefined, + total: number | undefined, +): { headline: string; sub: string } { + if (!status || status === 'WAITING_FOR_DEPLOY') { + return { + headline: 'Waiting for a worker…', + sub: 'Allocating compute capacity.', + }; + } + if (status === 'QUEUED' || status === 'DELAYED') { + return { + headline: 'Queued — waiting to start…', + sub: 'Your scan will pick up momentarily.', + }; + } + if (status === 'INTERRUPTED' || status === 'WAITING_TO_RESUME') { + return { + headline: 'Resuming…', + sub: 'Picking up where the run left off.', + }; + } + // EXECUTING from here down — describe by phase metadata. + if (!phase || phase === 'starting') { + return { + headline: 'Starting AI scan…', + sub: 'Initializing the suggestion run.', + }; + } + if (phase === 'embedding-tasks') { + const progress = total ? ` (${current ?? 0} / ${total})` : ''; + return { + headline: 'Embedding your task library…', + sub: `Indexing tasks for semantic search${progress}.`, + }; + } + if (phase === 'embedding-risks' || phase === 'embedding-vendors') { + return { + headline: 'Embedding the subject…', + sub: 'Preparing the query vector.', + }; + } + if (phase === 'waiting-for-index') { + return { + headline: 'Waiting for the index…', + sub: 'Letting the vector store finish ingesting before we query it.', + }; + } + if (phase === 'matching-risks' || phase === 'matching-vendors') { + return { + headline: 'Scanning your library…', + sub: 'Matching candidates and reranking with AI.', + }; + } + // 'done' — the status itself will flip to COMPLETED imminently. + return { + headline: 'Wrapping up…', + sub: 'Compiling the final suggestions.', + }; +} + +export function LoadingState({ + runId, + accessToken, + onReady, + onFailed, +}: { + runId: string; + accessToken: string; + onReady: (s: { tasks: SuggestedTask[]; controls: SuggestedControl[] }) => void; + onFailed: (reason: string) => void; +}) { + const { run } = useRealtimeRun(runId, { accessToken, enabled: true }); + const status = run?.status; + const output = run?.output as + | { suggestions?: { tasks?: SuggestedTask[]; controls?: SuggestedControl[] } } + | undefined; + const meta = (run?.metadata ?? {}) as { + phase?: PhaseName; + current?: number; + total?: number; + }; + + // Stash callbacks in refs so the effect can fire only on `status` changes + // (we never want to re-fire on every parent render that changes callback + // identity), but always read the freshest `output` and handlers — avoids + // the stale-closure race Cubic flagged (#32 on PR #2671). + const outputRef = useRef(output); + outputRef.current = output; + const onReadyRef = useRef(onReady); + onReadyRef.current = onReady; + const onFailedRef = useRef(onFailed); + onFailedRef.current = onFailed; + + useEffect(() => { + if (!status) return; + if (status === 'COMPLETED') { + const sugg = outputRef.current?.suggestions; + onReadyRef.current({ + tasks: Array.isArray(sugg?.tasks) ? sugg.tasks : [], + controls: Array.isArray(sugg?.controls) ? sugg.controls : [], + }); + return; + } + if (TERMINAL_FAILURE_STATUSES.has(status)) { + const reasons: Record = { + FAILED: 'The AI scan hit an error.', + CANCELED: 'The AI scan was canceled.', + CRASHED: 'The worker crashed mid-scan.', + SYSTEM_FAILURE: 'A system error stopped the scan.', + EXPIRED: 'The scan expired before it could start.', + TIMED_OUT: 'The scan took too long and timed out.', + }; + onFailedRef.current(reasons[status] ?? 'The AI scan failed.'); + } + }, [status]); + + // Treat unknown / pre-subscribe state as in-flight so the spinner stays up + // until the realtime stream catches up. + const inFlight = !status || IN_FLIGHT_STATUSES.has(status); + if (!inFlight) { + // Either COMPLETED (the effect above will transition state) or terminal + // failure (same). Render a quiet placeholder for the brief gap. + return ( +
+
Finishing up…
+
+ ); + } + + const { headline, sub } = describeStatus(status, meta.phase, meta.current, meta.total); + const showProgress = meta.phase === 'embedding-tasks' && typeof meta.total === 'number' && meta.total > 0; + const pct = showProgress + ? Math.min(100, Math.round(((meta.current ?? 0) / (meta.total ?? 1)) * 100)) + : null; + + return ( +
+
+
+
+ {sub} +
+ {pct !== null && ( +
+ + )} + {[0, 1, 2, 3].map((i) => ( +
+
+
+
+ ))} +
+ ); +} + +export function FailedState({ + reason, + retrying, + onRetry, + onDiscard, +}: { + reason: string; + retrying: boolean; + onRetry: () => void; + onDiscard: () => void; +}) { + return ( +
+
+
Auto-link failed
+
{reason}
+
+
+ + +
+
+ ); +} + +export function SuggestionsState({ + tasks, + controls, + checkedTaskIds, + applying, + submitting, + onToggle, + onDiscard, + onApply, + onRerun, +}: { + tasks: SuggestedTask[]; + controls: SuggestedControl[]; + checkedTaskIds: Set; + applying: boolean; + submitting: boolean; + onToggle: (id: string) => void; + onDiscard: () => void; + onApply: () => void; + onRerun: () => void; +}) { + const derivedControlsCount = useMemo( + () => controls.filter((c) => isControlDerived(c, checkedTaskIds)).length, + [controls, checkedTaskIds], + ); + const taskCount = tasks.length; + const controlCount = controls.length; + + return ( +
+
+
+ + + + {controls.length > 0 && ( + + )} + +
+ + + + +
+
+ ); +} diff --git a/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.sections.tsx b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.sections.tsx new file mode 100644 index 0000000000..83b0b37508 --- /dev/null +++ b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.sections.tsx @@ -0,0 +1,229 @@ +'use client'; + +import { cn } from '@/lib/utils'; +import { Text } from '@trycompai/design-system'; +import { ChevronLeft, ChevronRight, Checkmark } from '@trycompai/design-system/icons'; +import { useEffect, useState } from 'react'; +import { + isControlDerived, + type SuggestedControl, + type SuggestedTask, +} from './AutoLinkSuggestions.types'; + +// Match the LinkedWork list pagination so the selection UI feels consistent +// with the post-apply view (also 4 per page). +const PAGE_SIZE = 4; + +function Pagination({ + page, + pageCount, + onPageChange, + total, +}: { + page: number; + pageCount: number; + onPageChange: (next: number) => void; + total: number; +}) { + if (pageCount <= 1) return null; + const start = (page - 1) * PAGE_SIZE + 1; + const end = Math.min(page * PAGE_SIZE, total); + return ( +
+ + {start}–{end} of {total} + +
+ + + {page} / {pageCount} + + +
+
+ ); +} + +export function ConfidencePill({ score }: { score: number }) { + if (score <= 0) { + return ( + + — + + ); + } + // Cap at 100%; raw similarity + boost can exceed 1 in theory. + const pct = Math.min(100, Math.round(score * 100)); + const tier = pct >= 85 ? 'high' : pct >= 70 ? 'med' : 'low'; + const colorClass = + tier === 'high' + ? 'text-green-600 dark:text-green-400' + : tier === 'med' + ? 'text-amber-600 dark:text-amber-400' + : 'text-muted-foreground'; + return ( + + {pct}% + + ); +} + +export function TasksSection({ + tasks, + checkedTaskIds, + onToggle, +}: { + tasks: SuggestedTask[]; + checkedTaskIds: Set; + onToggle: (id: string) => void; +}) { + const pageCount = Math.max(1, Math.ceil(tasks.length / PAGE_SIZE)); + const [page, setPage] = useState(1); + // Re-clamp the page when the list shrinks (e.g. after a re-run returns + // fewer tasks). Keeps the user on a valid page rather than rendering + // empty space. + useEffect(() => { + if (page > pageCount) setPage(pageCount); + }, [page, pageCount]); + const start = (page - 1) * PAGE_SIZE; + const visible = tasks.slice(start, start + PAGE_SIZE); + + return ( +
+
+ Tasks + + {checkedTaskIds.size} / {tasks.length} selected + +
+ {tasks.length === 0 ? ( +

+ No tasks suggested. Try rerunning or link manually. +

+ ) : ( + <> + {visible.map((t) => { + const checked = checkedTaskIds.has(t.id); + return ( + + ); + })} + + + )} +
+ ); +} + +export function ControlsSection({ + controls, + checkedTaskIds, + derivedControlsCount, +}: { + controls: SuggestedControl[]; + checkedTaskIds: Set; + derivedControlsCount: number; +}) { + const pageCount = Math.max(1, Math.ceil(controls.length / PAGE_SIZE)); + const [page, setPage] = useState(1); + useEffect(() => { + if (page > pageCount) setPage(pageCount); + }, [page, pageCount]); + const start = (page - 1) * PAGE_SIZE; + const visible = controls.slice(start, start + PAGE_SIZE); + + return ( +
+
+ Controls + + {derivedControlsCount} {derivedControlsCount === 1 ? 'control' : 'controls'} via tasks + +
+ + These controls will be linked through the selected tasks. + +
+ {visible.map((c) => { + const isDerived = isControlDerived(c, checkedTaskIds); + // Show "code · name" when the framework provides a code + // (e.g. "1.2.7 · Credential Management"); fall back to just + // the name when code is empty so we don't render a leading + // "· " orphan separator. + const heading = c.code ? `${c.code} · ${c.name}` : c.name; + return ( +
+
+
{heading}
+
+ {c.framework} +
+
+ +
+ ); + })} + +
+
+ ); +} diff --git a/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.spec.tsx b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.spec.tsx new file mode 100644 index 0000000000..04c8cf3c9d --- /dev/null +++ b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.spec.tsx @@ -0,0 +1,238 @@ +import { TaskStatus } from '@db'; +import { fireEvent, render, screen, waitFor } from '@testing-library/react'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { AutoLinkSuggestions, type LinkedTask } from './AutoLinkSuggestions'; + +vi.mock('sonner', () => ({ toast: { info: vi.fn(), success: vi.fn(), error: vi.fn() } })); + +const realtimeRunMock = vi.fn(); +vi.mock('@trigger.dev/react-hooks', () => ({ + useRealtimeRun: (runId: string, opts: { accessToken?: string; enabled?: boolean }) => + realtimeRunMock(runId, opts), +})); + +afterEach(() => { + realtimeRunMock.mockReset(); +}); + +const linkedTask: LinkedTask = { + id: 'tsk_existing', + title: 'Existing task', + status: TaskStatus.todo, + controls: [{ id: 'ctl_existing', name: 'Existing control' }], +}; + +function defaultProps() { + return { + orgId: 'org_1', + canUpdate: true, + tasks: [] as LinkedTask[], + onSuggest: vi.fn().mockResolvedValue({ runId: 'r1', publicAccessToken: 't1' }), + onApply: vi.fn().mockResolvedValue(undefined), + onAfterApply: vi.fn().mockResolvedValue(undefined), + onUnlinkTask: vi.fn().mockResolvedValue(undefined), + }; +} + +describe('AutoLinkSuggestions', () => { + it('renders the empty state when no tasks are linked', () => { + realtimeRunMock.mockReturnValue({ run: null }); + render(); + expect(screen.getByRole('button', { name: /Suggest with AI/i })).toBeInTheDocument(); + expect(screen.getByText(/No tasks or controls linked yet/i)).toBeInTheDocument(); + }); + + it('renders the linked state with Re-assess affordance when tasks exist', () => { + realtimeRunMock.mockReturnValue({ run: null }); + render(); + expect(screen.getByText('Existing task')).toBeInTheDocument(); + expect(screen.getByRole('button', { name: /Re-assess/i })).toBeInTheDocument(); + }); + + it('calls onSuggest when "Suggest with AI" is clicked', async () => { + realtimeRunMock.mockReturnValue({ run: null }); + const props = defaultProps(); + render(); + fireEvent.click(screen.getByRole('button', { name: /Suggest with AI/i })); + await waitFor(() => { + expect(props.onSuggest).toHaveBeenCalledTimes(1); + }); + }); + + it('transitions to suggestions when the run completes with output', async () => { + realtimeRunMock.mockReturnValue({ + run: { + status: 'COMPLETED', + output: { + suggestions: { + tasks: [ + { id: 'tsk_a', title: 'Suggested A', status: 'todo', score: 0.9 }, + { id: 'tsk_b', title: 'Suggested B', status: 'todo', score: 0.7 }, + ], + controls: [], + }, + }, + }, + }); + const props = defaultProps(); + render(); + fireEvent.click(screen.getByRole('button', { name: /Suggest with AI/i })); + + await waitFor(() => { + expect(screen.getByText(/AI found/i)).toBeInTheDocument(); + }); + expect(screen.getByText('Suggested A')).toBeInTheDocument(); + expect(screen.getByText('Suggested B')).toBeInTheDocument(); + // Both default-checked → Link 2. + expect(screen.getByRole('button', { name: /^Link 2$/i })).toBeInTheDocument(); + }); + + it('updates the Link N count when a task is unchecked', async () => { + realtimeRunMock.mockReturnValue({ + run: { + status: 'COMPLETED', + output: { + suggestions: { + tasks: [ + { id: 'tsk_a', title: 'Suggested A', status: 'todo', score: 0.9 }, + { id: 'tsk_b', title: 'Suggested B', status: 'todo', score: 0.7 }, + ], + controls: [], + }, + }, + }, + }); + const props = defaultProps(); + render(); + fireEvent.click(screen.getByRole('button', { name: /Suggest with AI/i })); + + await waitFor(() => { + expect(screen.getByRole('button', { name: /^Link 2$/i })).toBeInTheDocument(); + }); + + // Click the row for "Suggested A" — the row itself is the toggle. + fireEvent.click(screen.getByRole('button', { name: /Uncheck task Suggested A/i })); + expect(screen.getByRole('button', { name: /^Link 1$/i })).toBeInTheDocument(); + }); + + it('calls onApply with selected ids and replace=false in fresh mode', async () => { + realtimeRunMock.mockReturnValue({ + run: { + status: 'COMPLETED', + output: { + suggestions: { + tasks: [ + { id: 'tsk_a', title: 'Suggested A', status: 'todo', score: 0.9 }, + ], + controls: [], + }, + }, + }, + }); + const props = defaultProps(); + render(); + fireEvent.click(screen.getByRole('button', { name: /Suggest with AI/i })); + + await waitFor(() => { + expect(screen.getByRole('button', { name: /^Link 1$/i })).toBeInTheDocument(); + }); + + fireEvent.click(screen.getByRole('button', { name: /^Link 1$/i })); + await waitFor(() => { + expect(props.onApply).toHaveBeenCalledWith({ taskIds: ['tsk_a'], replace: false }); + }); + }); + + it('calls onApply with replace=true when launched from the Re-assess button', async () => { + realtimeRunMock.mockReturnValue({ + run: { + status: 'COMPLETED', + output: { + suggestions: { + tasks: [{ id: 'tsk_a', title: 'A', status: 'todo', score: 0.9 }], + controls: [], + }, + }, + }, + }); + const props = defaultProps(); + render(); + + fireEvent.click(screen.getByRole('button', { name: /Re-assess/i })); + await waitFor(() => { + expect(props.onSuggest).toHaveBeenCalled(); + }); + // After the realtime "completed" event, suggestions render with mode=reassess. + await waitFor(() => { + expect(screen.getByText(/AI found/i)).toBeInTheDocument(); + }); + + // 1 AI-suggested + 1 currently-linked (merged) → 2 checked. + fireEvent.click(screen.getByRole('button', { name: /^Link 2$/i })); + await waitFor(() => { + expect(props.onApply).toHaveBeenCalledWith({ + taskIds: expect.arrayContaining(['tsk_a', 'tsk_existing']), + replace: true, + }); + }); + }); + + it('renders read-only controls section with derived dimming', async () => { + realtimeRunMock.mockReturnValue({ + run: { + status: 'COMPLETED', + output: { + suggestions: { + tasks: [{ id: 'tsk_a', title: 'A', status: 'todo', score: 0.9 }], + controls: [ + { + id: 'ctl_1', + code: 'CC1.1', + name: 'Awareness', + framework: 'SOC 2', + score: 0.9, + viaTaskIds: ['tsk_a'], + }, + ], + }, + }, + }, + }); + const props = defaultProps(); + render(); + fireEvent.click(screen.getByRole('button', { name: /Suggest with AI/i })); + + await waitFor(() => { + expect( + screen.getByText(/These controls will be linked through the selected tasks/i), + ).toBeInTheDocument(); + }); + expect(screen.getByText(/CC1.1 · Awareness/i)).toBeInTheDocument(); + }); + + it('returns to the prior state when Discard is clicked', async () => { + realtimeRunMock.mockReturnValue({ + run: { + status: 'COMPLETED', + output: { + suggestions: { + tasks: [{ id: 'tsk_a', title: 'A', status: 'todo', score: 0.9 }], + controls: [], + }, + }, + }, + }); + const props = defaultProps(); + render(); + fireEvent.click(screen.getByRole('button', { name: /Suggest with AI/i })); + + await waitFor(() => { + expect(screen.getByRole('button', { name: /^Link 1$/i })).toBeInTheDocument(); + }); + + fireEvent.click(screen.getByRole('button', { name: /^Discard$/i })); + await waitFor(() => { + expect(screen.getByRole('button', { name: /Suggest with AI/i })).toBeInTheDocument(); + }); + }); +}); diff --git a/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.tsx b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.tsx new file mode 100644 index 0000000000..30fd232d8d --- /dev/null +++ b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.tsx @@ -0,0 +1,295 @@ +'use client'; + +import { MagicWandFilled } from '@trycompai/design-system/icons'; +import { useEffect, useState } from 'react'; +import { toast } from 'sonner'; +import { LinkedWork } from './LinkedWork'; +import { + EmptyState, + FailedState, + LoadingState, + SuggestionsState, +} from './AutoLinkSuggestions.parts'; +import type { + LinkedTask, + Mode, + State, + SuggestedControl, + SuggestedTask, +} from './AutoLinkSuggestions.types'; + +export type { + SuggestedTask, + SuggestedControl, + LinkedTask, +} from './AutoLinkSuggestions.types'; + +export interface AutoLinkSuggestionsProps { + orgId: string; + /** Existing linked tasks (rendered in `linked` state via ). */ + tasks: LinkedTask[]; + canUpdate: boolean; + /** Triggers the AI scan; returns a runId + token for realtime subscription. */ + onSuggest: () => Promise<{ runId: string; publicAccessToken: string }>; + /** Persists the user-confirmed selection. `replace` is true for re-assess. */ + onApply: (params: { taskIds: string[]; replace: boolean }) => Promise; + /** Called after apply succeeds — typically the parent's onRegenerate. */ + onAfterApply?: () => Promise; + /** Per-task unlink, plumbed through to in linked state. */ + onUnlinkTask?: (taskId: string) => Promise; + /** + * `'kickoff'` — wide panel for the truly-fresh case. + * `'kickoff-with-plan'` — wide panel with copy adapted to "AI will only + * suggest tasks/controls; your plan stays". + * `'default'` — small per-column CTA (currently unused). + */ + emptyVariant?: 'default' | 'kickoff' | 'kickoff-with-plan'; + /** Called when the user clicks "Start from scratch" in the kickoff panel. + * Parent should dismiss the kickoff state so the editor renders. */ + onStartFromScratch?: () => void; + /** + * Resumes an in-flight or completed-but-unreviewed AI scan after a page + * reload. Parent fetches `GET /auto-link/active`. When this returns a run, + * the component jumps straight into the loading state and re-subscribes to + * the trigger.dev run via `useRealtimeRun`. + */ + onResume?: () => Promise<{ runId: string; publicAccessToken: string } | null>; + /** + * Clears the persisted runId server-side. Called when the user clicks + * Discard so the next reload doesn't re-resume a discarded run. + */ + onDiscardRun?: () => Promise; +} + +export function AutoLinkSuggestions({ + orgId, + tasks, + canUpdate, + onSuggest, + onApply, + onAfterApply, + onUnlinkTask, + emptyVariant = 'default', + onStartFromScratch, + onResume, + onDiscardRun, +}: AutoLinkSuggestionsProps) { + const [state, setState] = useState(() => + tasks.length > 0 ? { kind: 'linked' } : { kind: 'empty' }, + ); + const [submitting, setSubmitting] = useState(false); + const [applying, setApplying] = useState(false); + + // Keep linked/empty state in sync with parent task list. Don't override + // mid-flow loading/suggestions/failed states. + useEffect(() => { + setState((prev) => { + if (prev.kind === 'linked' || prev.kind === 'empty') { + return tasks.length > 0 ? { kind: 'linked' } : { kind: 'empty' }; + } + return prev; + }); + }, [tasks.length]); + + // On mount, if a run is already in flight (or completed-but-unreviewed) on + // the server, jump straight to the loading state and re-subscribe. Mode is + // unknown after a reload; default 'fresh' so apply stays additive (no + // accidental destructive replace on user side). + // + // Guard against overwriting a newer transition: the user might click + // "Suggest" before /active resolves, in which case `state` will already + // be `loading` (with their own runId) and we must NOT clobber it with + // the server-side resume payload. Same goes for `review` / `confirming` + // — only apply the resume when we're still in the initial render state. + // (Cubic finding #33 on PR #2671.) + useEffect(() => { + if (!onResume) return; + let cancelled = false; + void onResume().then((active) => { + if (cancelled || !active) return; + setState((prev) => { + if (prev.kind !== 'empty' && prev.kind !== 'linked') return prev; + return { + kind: 'loading', + runId: active.runId, + publicAccessToken: active.publicAccessToken, + mode: 'fresh', + }; + }); + }); + return () => { + cancelled = true; + }; + }, [onResume]); + + const handleSuggest = async (mode: Mode) => { + setSubmitting(true); + try { + const { runId, publicAccessToken } = await onSuggest(); + setState({ kind: 'loading', runId, publicAccessToken, mode }); + } catch { + toast.error('Suggest failed. Try again.'); + } finally { + setSubmitting(false); + } + }; + + const handleSuggestionsReady = (suggestions: { + tasks: SuggestedTask[]; + controls: SuggestedControl[]; + }) => { + setState((prev) => { + const mode: Mode = prev.kind === 'loading' ? prev.mode : 'fresh'; + const checked = new Set(suggestions.tasks.map((t) => t.id)); + const merged: SuggestedTask[] = [...suggestions.tasks]; + if (mode === 'reassess') { + for (const t of tasks) checked.add(t.id); + const seen = new Set(suggestions.tasks.map((t) => t.id)); + for (const t of tasks) { + if (seen.has(t.id)) continue; + merged.push({ id: t.id, title: t.title, status: t.status, score: 0 }); + } + } + return { + kind: 'suggestions', + mode, + tasks: merged, + controls: suggestions.controls, + checkedTaskIds: checked, + }; + }); + }; + + const handleSuggestionsFailed = (reason: string) => { + setState((prev) => ({ + kind: 'failed', + reason, + mode: prev.kind === 'loading' ? prev.mode : 'fresh', + })); + }; + + const handleToggleTask = (id: string) => { + setState((prev) => { + if (prev.kind !== 'suggestions') return prev; + const next = new Set(prev.checkedTaskIds); + if (next.has(id)) next.delete(id); + else next.add(id); + return { ...prev, checkedTaskIds: next }; + }); + }; + + const handleDiscard = () => { + if (onDiscardRun) { + void onDiscardRun(); + } + setState(tasks.length > 0 ? { kind: 'linked' } : { kind: 'empty' }); + }; + + const handleApply = async () => { + if (state.kind !== 'suggestions') return; + setApplying(true); + try { + const taskIds = [...state.checkedTaskIds]; + const isReassess = state.mode === 'reassess'; + await onApply({ taskIds, replace: isReassess }); + if (onAfterApply) { + try { + await onAfterApply(); + } catch { + /* parent surfaces its own errors */ + } + } + toast.success(`Linked ${taskIds.length} task${taskIds.length === 1 ? '' : 's'}`); + // Reassess (replace=true) wipes the existing set, so the post-apply + // state is exactly taskIds. Additive (replace=false) keeps existing + // linked tasks, so we should land in `linked` if EITHER the existing + // set has any rows OR we just added any. Without this, applying 0 + // new tasks in additive mode wrongly swaps the linked-state UI back + // to the kickoff "empty" state. (Cubic finding on PR #2671.) + const willHaveLinked = isReassess + ? taskIds.length > 0 + : tasks.length > 0 || taskIds.length > 0; + setState(willHaveLinked ? { kind: 'linked' } : { kind: 'empty' }); + } catch { + toast.error('Failed to apply suggestions.'); + } finally { + setApplying(false); + } + }; + + const handleRerun = () => { + if (state.kind !== 'suggestions') return; + void handleSuggest(state.mode); + }; + + if (state.kind === 'empty') { + return ( + void handleSuggest('fresh')} + onStartFromScratch={onStartFromScratch} + variant={emptyVariant} + /> + ); + } + + if (state.kind === 'loading') { + return ( + + ); + } + + if (state.kind === 'suggestions') { + return ( + void handleApply()} + onRerun={handleRerun} + /> + ); + } + + if (state.kind === 'failed') { + return ( + void handleSuggest(state.mode)} + onDiscard={handleDiscard} + /> + ); + } + + // linked + return ( +
+ {canUpdate && ( +
+ +
+ )} + +
+ ); +} diff --git a/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.types.ts b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.types.ts new file mode 100644 index 0000000000..d67c10c825 --- /dev/null +++ b/apps/app/src/components/risks/treatment-plan/AutoLinkSuggestions.types.ts @@ -0,0 +1,46 @@ +import type { TaskStatus } from '@db'; + +export interface SuggestedTask { + id: string; + title: string; + status: string; + score: number; +} + +export interface SuggestedControl { + id: string; + code: string; + name: string; + framework: string; + score: number; + viaTaskIds: string[]; +} + +export interface LinkedTask { + id: string; + title: string; + status: TaskStatus; + controls: { id: string; name: string }[]; +} + +export type Mode = 'fresh' | 'reassess'; + +export type State = + | { kind: 'linked' } + | { kind: 'empty' } + | { kind: 'loading'; runId: string; publicAccessToken: string; mode: Mode } + | { + kind: 'suggestions'; + mode: Mode; + tasks: SuggestedTask[]; + controls: SuggestedControl[]; + checkedTaskIds: Set; + } + | { kind: 'failed'; reason: string; mode: Mode }; + +export function isControlDerived( + c: SuggestedControl, + checkedTaskIds: Set, +): boolean { + return c.viaTaskIds.some((id) => checkedTaskIds.has(id)); +} diff --git a/apps/app/src/components/risks/treatment-plan/DescriptionEditor.tsx b/apps/app/src/components/risks/treatment-plan/DescriptionEditor.tsx new file mode 100644 index 0000000000..becb4fd583 --- /dev/null +++ b/apps/app/src/components/risks/treatment-plan/DescriptionEditor.tsx @@ -0,0 +1,406 @@ +'use client'; + +import { Button } from '@trycompai/design-system'; +import { Edit, Renew } from '@trycompai/design-system/icons'; +import { useRealtimeRun } from '@trigger.dev/react-hooks'; +import { useEffect, useLayoutEffect, useRef, useState } from 'react'; +import ReactMarkdown from 'react-markdown'; +import remarkGfm from 'remark-gfm'; +import { cn } from '@/lib/utils'; + +interface DescriptionEditorProps { + value: string; + onSave: (next: string) => Promise; + onRegenerate: () => Promise; + regenerating: boolean; + disabled?: boolean; + /** + * The trigger.dev run handle for an in-flight regeneration. When set, the + * editor subscribes via `useRealtimeRun`, renders status-specific progress + * copy, and notifies the parent via `onRegenSettled` when the run reaches + * a terminal state. Null/undefined while no regen is active. + */ + regenRun?: { runId: string; publicAccessToken: string } | null; + /** Called once the regeneration run terminates (success or failure). */ + onRegenSettled?: (result: { success: boolean; reason?: string }) => void; +} + +const TERMINAL_FAILURE_STATUSES = new Set([ + 'FAILED', + 'CANCELED', + 'CRASHED', + 'SYSTEM_FAILURE', + 'EXPIRED', + 'TIMED_OUT', +]); + +/** + * Cap (in px) for both the markdown preview and the auto-growing textarea. + * Past this height, the body scrolls internally so the Treatment plan column + * stays roughly aligned with the Strategy and Linked Work columns instead + * of pushing the whole row downward when AI emits a long plan. + */ +const TEXTAREA_MAX_PX = 480; + +function regenStatusCopy(status: string | undefined): { headline: string; sub: string } { + if (!status || status === 'WAITING_FOR_DEPLOY') { + return { + headline: 'Starting AI scan…', + sub: 'Allocating compute capacity.', + }; + } + if (status === 'QUEUED' || status === 'DELAYED') { + return { + headline: 'Queued — waiting to start…', + sub: 'Your regeneration will begin in a moment.', + }; + } + if (status === 'INTERRUPTED' || status === 'WAITING_TO_RESUME') { + return { + headline: 'Resuming…', + sub: 'Picking up where the run left off.', + }; + } + return { + headline: 'AI is drafting your treatment plan…', + sub: 'Reading linked controls and tasks, then writing each citation.', + }; +} + +function countWords(text: string): number { + const trimmed = text.trim(); + if (trimmed === '') return 0; + return trimmed.split(/\s+/).length; +} + +export function DescriptionEditor({ + value, + onSave, + onRegenerate, + regenerating, + disabled, + regenRun, + onRegenSettled, +}: DescriptionEditorProps) { + const [draft, setDraft] = useState(value); + const [saving, setSaving] = useState(false); + // Mode: 'preview' renders markdown, 'edit' shows the auto-growing textarea. + // We default to 'edit' when the value is empty (nothing to preview yet) and + // stay in 'edit' when an AI regeneration completes with new content so the + // user immediately sees what was drafted. + const [mode, setMode] = useState<'preview' | 'edit'>( + value.trim().length > 0 ? 'preview' : 'edit', + ); + const textareaRef = useRef(null); + + // Resync the draft from upstream `value` ONLY when the user isn't + // actively editing. Without the `mode === 'edit'` guard, a background + // SWR revalidation, AI regeneration, or any other prop change would + // wipe whatever the user was typing. (Cubic finding on PR #2671.) + useEffect(() => { + if (saving) return; + if (mode === 'edit') return; + setDraft(value); + }, [value, saving, mode]); + + // When a fresh value arrives from upstream (regenerate, server update) and + // we're not actively editing, drop back to preview. + useEffect(() => { + if (mode === 'edit' || saving) return; + if (value.trim().length === 0) setMode('edit'); + }, [value, mode, saving]); + + // Regenerate-with-AI bypasses the in-edit guard above. When a regen run + // terminates (`regenRun` flips from set → null), the user explicitly + // asked to overwrite whatever they had — keeping the stale draft and + // requiring a refresh to see the new prose was confusing. + // + // The new prose may already be in `value` at the moment regenRun + // clears (sync write before the parent flips the run handle), or it + // may arrive in a later render after SWR refetches. Both paths are + // handled: + // + // 1. Sync arrival: when regenRun flips set→null, immediately apply + // the current value and force preview. + // 2. Async arrival: capture the value-at-clear-time. The next render + // where `value` differs from the captured snapshot is the AI prose + // landing — apply it, force preview, and clear the latch. + // + // Without (2), a regen that completes BEFORE the SWR refetch would + // sync-apply the OLD value, and the new prose arriving moments later + // would be ignored because the in-edit guard skips resync while + // mode === 'edit'. + const prevRegenRunRef = useRef(regenRun); + const valueAtRegenClearRef = useRef(null); + useEffect(() => { + const wasRunning = prevRegenRunRef.current != null; + const isRunning = regenRun != null; + prevRegenRunRef.current = regenRun; + if (wasRunning && !isRunning) { + // Path 1: sync arrival — value has already updated. + valueAtRegenClearRef.current = value; + setDraft(value); + if (value.trim().length > 0) setMode('preview'); + } + }, [regenRun, value]); + + useEffect(() => { + const captured = valueAtRegenClearRef.current; + if (captured === null) return; + if (value === captured) return; + // Path 2: async arrival — value just changed since regen cleared, + // so this is the AI prose landing. Overwrite even if user is in + // edit mode (they explicitly opted into the overwrite by clicking + // Regenerate). + valueAtRegenClearRef.current = null; + setDraft(value); + if (value.trim().length > 0) setMode('preview'); + }, [value]); + + // Auto-grow the textarea to fit content, but cap at TEXTAREA_MAX_PX so a + // long draft doesn't stretch the Treatment plan column past the Strategy + // / Linked Work columns. Internal scroll kicks in past the cap. + useLayoutEffect(() => { + if (mode !== 'edit') return; + const el = textareaRef.current; + if (!el) return; + el.style.height = 'auto'; + const next = Math.max(Math.min(el.scrollHeight, TEXTAREA_MAX_PX), 200); + el.style.height = `${next}px`; + el.style.overflowY = el.scrollHeight > TEXTAREA_MAX_PX ? 'auto' : 'hidden'; + }, [draft, mode]); + + const isDirty = draft.trim() !== (value ?? '').trim(); + const wordCount = countWords(draft); + const charCount = draft.length; + const hasValue = value.trim().length > 0; + + const handleSave = async () => { + if (!isDirty) { + setMode('preview'); + return; + } + setSaving(true); + try { + await onSave(draft.trim()); + setMode('preview'); + } finally { + setSaving(false); + } + }; + + const handleCancelEdit = () => { + setDraft(value); + setMode('preview'); + }; + + return ( +
+ {mode === 'preview' && hasValue ? ( +
+ +
+ ) : ( +