perf(api): write routes stop re-resolving the active company (#1928)

* perf(api): write routes stop re-resolving the active company

withRouteContext resolves the active company (one resolve_active_company
RPC, ~40 ms p50 on prod) and then, for the 256 routes that pass
requireWrite: true, called requireWritePermission(), which resolved it a
second time before its role select. Two sequential round trips repeating
work the wrapper had just done, on every mutating request.

requireWritePermission() and getCompanyRole() now accept an optional
`known` context; the wrapper passes { companyId }, so the helper goes
straight to the membership select. Callers that pass nothing behave
exactly as before, and the shared selectRole() keeps both helpers on the
same query. The role is still looked up, never trusted from the caller.

Tests: known companyId skips resolution, known role skips the select, a
known viewer is still 403, a known company without a membership row is
still 403, legacy calls unchanged; new lib/api/__tests__/with-route-
context.test.ts pins that the wrapper resolves the company exactly once,
hands it to the guard, never calls the guard on read routes, passes the
guard's 403 through with a request id, and emits Server-Timing.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* test(customers): viewer gate expects the wrapper to hand over the resolved company

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Jakob Wennberg <311770904+jakobwennberg-oss@users.noreply.github.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

app/api/customers/__tests__/viewer.test.ts

@@ -82,6 +82,10 @@ describe('POST /api/customers: viewer role gate', () => {
 
     await POST(request)
 
-    expect(requireWritePermissionMock).toHaveBeenCalledWith(mockSupabase, 'user-1')
+    // The wrapper hands over the company it already resolved so the guard
+    // does not repeat the resolve_active_company round trip.
+    expect(requireWritePermissionMock).toHaveBeenCalledWith(mockSupabase, 'user-1', {
+      companyId: 'company-1',
+    })
   })
 })

lib/api/__tests__/with-route-context.test.ts

@@ -0,0 +1,120 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest'
+import { NextResponse } from 'next/server'
+import { createMockSupabase } from '@/tests/helpers'
+
+const authState = vi.hoisted(() => ({
+  user: { id: 'user-1' } as { id: string } | null,
+}))
+
+const requireWriteMock = vi.hoisted(() => vi.fn())
+
+vi.mock('@/lib/auth/require-auth', () => ({
+  requireAuth: vi.fn(async () => {
+    if (!authState.user) {
+      return { error: NextResponse.json({ error: 'unauthorized' }, { status: 401 }) }
+    }
+    return { user: authState.user, supabase: supabaseRef.supabase }
+  }),
+}))
+
+vi.mock('@/lib/company/context', () => ({
+  getActiveCompanyId: vi.fn(),
+}))
+
+vi.mock('@/lib/auth/require-write', () => ({
+  requireWritePermission: (...args: unknown[]) => requireWriteMock(...args),
+}))
+
+const supabaseRef = vi.hoisted(() => ({ supabase: null as unknown }))
+
+import { withRouteContext } from '../with-route-context'
+import { getActiveCompanyId } from '@/lib/company/context'
+
+const EMPTY_PARAMS = { params: Promise.resolve({}) }
+
+describe('withRouteContext', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    authState.user = { id: 'user-1' }
+    supabaseRef.supabase = createMockSupabase().supabase
+    vi.mocked(getActiveCompanyId).mockResolvedValue('company-1')
+    requireWriteMock.mockResolvedValue({ ok: true })
+  })
+
+  it('resolves the company once and hands it to the write guard on write routes', async () => {
+    const handler = vi.fn(async () => NextResponse.json({ ok: true }))
+    const route = withRouteContext('test.write', handler, { requireWrite: true })
+
+    const res = await route(new Request('http://localhost/api/test', { method: 'POST' }), EMPTY_PARAMS)
+
+    expect(res.status).toBe(200)
+    expect(getActiveCompanyId).toHaveBeenCalledTimes(1)
+    expect(requireWriteMock).toHaveBeenCalledTimes(1)
+    expect(requireWriteMock).toHaveBeenCalledWith(supabaseRef.supabase, 'user-1', {
+      companyId: 'company-1',
+    })
+    expect(handler).toHaveBeenCalledWith(
+      expect.any(Request),
+      expect.objectContaining({ companyId: 'company-1', user: { id: 'user-1' } }),
+      EMPTY_PARAMS,
+    )
+  })
+
+  it('never invokes the write guard on read routes', async () => {
+    const route = withRouteContext('test.read', async () => NextResponse.json({ ok: true }))
+
+    const res = await route(new Request('http://localhost/api/test'), EMPTY_PARAMS)
+
+    expect(res.status).toBe(200)
+    expect(getActiveCompanyId).toHaveBeenCalledTimes(1)
+    expect(requireWriteMock).not.toHaveBeenCalled()
+  })
+
+  it('passes the guard 403 through with a request id and skips the handler', async () => {
+    requireWriteMock.mockResolvedValue({
+      ok: false,
+      response: NextResponse.json({ error: 'viewer' }, { status: 403 }),
+    })
+    const handler = vi.fn(async () => NextResponse.json({ ok: true }))
+    const route = withRouteContext('test.write', handler, { requireWrite: true })
+
+    const res = await route(new Request('http://localhost/api/test', { method: 'POST' }), EMPTY_PARAMS)
+
+    expect(res.status).toBe(403)
+    expect(res.headers.get('X-Request-Id')).toMatch(/^req_/)
+    expect(handler).not.toHaveBeenCalled()
+  })
+
+  it('returns COMPANY_CONTEXT_MISSING before the guard when no company resolves', async () => {
+    vi.mocked(getActiveCompanyId).mockResolvedValue(null)
+    const route = withRouteContext('test.write', async () => NextResponse.json({ ok: true }), {
+      requireWrite: true,
+    })
+
+    const res = await route(new Request('http://localhost/api/test', { method: 'POST' }), EMPTY_PARAMS)
+
+    expect(res.status).toBe(400)
+    expect(requireWriteMock).not.toHaveBeenCalled()
+  })
+
+  it('returns 401 from requireAuth untouched except for the request id', async () => {
+    authState.user = null
+    const route = withRouteContext('test.read', async () => NextResponse.json({ ok: true }))
+
+    const res = await route(new Request('http://localhost/api/test'), EMPTY_PARAMS)
+
+    expect(res.status).toBe(401)
+    expect(res.headers.get('X-Request-Id')).toMatch(/^req_/)
+    expect(getActiveCompanyId).not.toHaveBeenCalled()
+  })
+
+  it('emits a Server-Timing header with the auth, company and handler phases', async () => {
+    const route = withRouteContext('test.read', async () => NextResponse.json({ ok: true }))
+
+    const res = await route(new Request('http://localhost/api/test'), EMPTY_PARAMS)
+
+    expect(res.headers.get('Server-Timing')).toMatch(
+      /^auth;dur=\d+, company;dur=\d+, handler;dur=\d+$/,
+    )
+  })
+})

lib/api/with-route-context.ts

@@ -136,8 +136,11 @@ export function withRouteContext<P extends DynamicParams = { params: Promise<Rec
       if (requireWrite) {
         // Delegate to the existing helper so tests that already mock it
         // continue to work. The helper returns its own 403 NextResponse;
-        // we wrap it in our request-id header for traceability.
-        const writeCheck = await requireWritePermission(supabase, user.id)
+        // we wrap it in our request-id header for traceability. The
+        // company id resolved above is handed over so the helper does not
+        // repeat the resolve_active_company round trip (measured ~40 ms p50
+        // on prod) on every write route.
+        const writeCheck = await requireWritePermission(supabase, user.id, { companyId })
         if (!writeCheck.ok) {
           userLog.warn('write permission denied')
           if (!writeCheck.response.headers.get('X-Request-Id')) {

lib/auth/__tests__/require-write.test.ts

@@ -80,6 +80,66 @@ describe('requireWritePermission', () => {
   })
 })
 
+describe('requireWritePermission with a known route context', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('skips the active-company resolution when companyId is known', async () => {
+    const { supabase, mockResult } = createMockSupabase()
+    mockResult({ data: { role: 'member' } })
+
+    const result = await requireWritePermission(supabase, 'user-1', { companyId: 'company-9' })
+    expect(result.ok).toBe(true)
+    expect(getActiveCompanyId).not.toHaveBeenCalled()
+    expect(supabase.from).toHaveBeenCalledTimes(1)
+    expect(supabase.from).toHaveBeenCalledWith('company_members')
+  })
+
+  it('skips the membership select when the role is known too', async () => {
+    const { supabase } = createMockSupabase()
+
+    const result = await requireWritePermission(supabase, 'user-1', {
+      companyId: 'company-9',
+      role: 'admin',
+    })
+    expect(result.ok).toBe(true)
+    expect(getActiveCompanyId).not.toHaveBeenCalled()
+    expect(supabase.from).not.toHaveBeenCalled()
+  })
+
+  it('still rejects a known viewer role with 403', async () => {
+    const { supabase } = createMockSupabase()
+
+    const result = await requireWritePermission(supabase, 'user-1', {
+      companyId: 'company-9',
+      role: 'viewer',
+    })
+    expect(result.ok).toBe(false)
+    if (!result.ok) expect(result.response.status).toBe(403)
+    expect(supabase.from).not.toHaveBeenCalled()
+  })
+
+  it('a known company with no membership row is rejected, not trusted', async () => {
+    const { supabase, mockResult } = createMockSupabase()
+    mockResult({ data: null })
+
+    const result = await requireWritePermission(supabase, 'user-1', { companyId: 'company-9' })
+    expect(result.ok).toBe(false)
+    if (!result.ok) expect(result.response.status).toBe(403)
+  })
+
+  it('falls back to resolution when no context is passed (legacy callers)', async () => {
+    const { supabase, mockResult } = createMockSupabase()
+    vi.mocked(getActiveCompanyId).mockResolvedValue('company-1')
+    mockResult({ data: { role: 'owner' } })
+
+    const result = await requireWritePermission(supabase, 'user-1', undefined)
+    expect(result.ok).toBe(true)
+    expect(getActiveCompanyId).toHaveBeenCalledTimes(1)
+  })
+})
+
 describe('getCompanyRole', () => {
   beforeEach(() => {
     vi.clearAllMocks()

lib/auth/require-write.ts

@@ -27,11 +27,40 @@ type WritePermissionResult =
   | { ok: true }
   | { ok: false; response: NextResponse }
 
+/**
+ * Facts the caller has already established for this request. Passing
+ * `companyId` skips the `resolve_active_company` round trip that
+ * `getActiveCompanyId` would otherwise repeat (withRouteContext resolves it
+ * two awaits earlier for every route); passing `role` skips the membership
+ * select as well. Only ever pass values that came from `getActiveCompanyId`
+ * / `company_members` for the same user in the same request: this is a
+ * dedupe, not a trust boundary.
+ */
+export interface KnownRouteContext {
+  companyId: string
+  role?: CompanyRole
+}
+
+async function selectRole(
+  supabase: SupabaseClient,
+  companyId: string,
+  userId: string,
+): Promise<CompanyRole | null> {
+  const { data: membership } = await supabase
+    .from('company_members')
+    .select('role')
+    .eq('company_id', companyId)
+    .eq('user_id', userId)
+    .maybeSingle()
+  return membership ? (membership.role as CompanyRole) : null
+}
+
 export async function requireWritePermission(
   supabase: SupabaseClient,
   userId: string,
+  known?: KnownRouteContext,
 ): Promise<WritePermissionResult> {
-  const companyId = await getActiveCompanyId(supabase, userId)
+  const companyId = known?.companyId ?? (await getActiveCompanyId(supabase, userId))
 
   if (!companyId) {
     return {
@@ -43,14 +72,9 @@ export async function requireWritePermission(
     }
   }
 
-  const { data: membership } = await supabase
-    .from('company_members')
-    .select('role')
-    .eq('company_id', companyId)
-    .eq('user_id', userId)
-    .maybeSingle()
+  const role = known?.role ?? (await selectRole(supabase, companyId, userId))
 
-  if (!membership || membership.role === 'viewer') {
+  if (!role || role === 'viewer') {
     return {
       ok: false,
       response: NextResponse.json(
@@ -81,8 +105,9 @@ export type CompanyRoleResult =
 export async function getCompanyRole(
   supabase: SupabaseClient,
   userId: string,
+  known?: Pick<KnownRouteContext, 'companyId'>,
 ): Promise<CompanyRoleResult> {
-  const companyId = await getActiveCompanyId(supabase, userId)
+  const companyId = known?.companyId ?? (await getActiveCompanyId(supabase, userId))
 
   if (!companyId) {
     return {
@@ -94,14 +119,9 @@ export async function getCompanyRole(
     }
   }
 
-  const { data: membership } = await supabase
-    .from('company_members')
-    .select('role')
-    .eq('company_id', companyId)
-    .eq('user_id', userId)
-    .maybeSingle()
+  const role = await selectRole(supabase, companyId, userId)
 
-  if (!membership) {
+  if (!role) {
     return {
       ok: false,
       response: NextResponse.json(
@@ -111,5 +131,5 @@ export async function getCompanyRole(
     }
   }
 
-  return { ok: true, role: membership.role as CompanyRole, companyId }
+  return { ok: true, role, companyId }
 }