feat(pi): add BYOK Exa search to Create PR

.github/workflows/pi-search-smoke.yml

@@ -0,0 +1,33 @@
+name: Pi Search Staging Smoke
+
+on:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  smoke:
+    name: Verify deployed Pi search boundary
+    runs-on: ubuntu-latest
+    environment: staging
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6
+
+      - name: Setup Bun
+        uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2
+        with:
+          bun-version: 1.3.13
+
+      - name: Install dependencies
+        run: bun install --frozen-lockfile
+
+      - name: Run remote Pi search smoke
+        env:
+          E2B_API_KEY: ${{ secrets.E2B_API_KEY }}
+          E2B_PI_TEMPLATE_ID: ${{ secrets.E2B_PI_TEMPLATE_ID }}
+          INTERNAL_API_SECRET: ${{ secrets.INTERNAL_API_SECRET }}
+          PI_EXA_BROKER_BASE_URL: ${{ vars.PI_EXA_BROKER_BASE_URL }}
+          PI_SEARCH_SMOKE_CAPABILITY: ${{ secrets.PI_SEARCH_SMOKE_CAPABILITY }}
+        run: bun run ./apps/sim/scripts/check-pi-sandbox-contract.ts --remote

.github/workflows/test-build.yml

@@ -12,6 +12,20 @@ jobs:
     name: Lint and Test
     runs-on: ${{ (vars.CI_PROVIDER == '' || vars.CI_PROVIDER == 'blacksmith') && 'blacksmith-8vcpu-ubuntu-2404' || 'ubuntu-latest' }}
     timeout-minutes: 15
+    services:
+      postgres:
+        image: pgvector/pgvector:pg17
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_DB: simstudio
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd "pg_isready -U postgres"
+          --health-interval 5s
+          --health-timeout 5s
+          --health-retries 10
 
     steps:
       - name: Checkout code
@@ -55,6 +69,18 @@ jobs:
       - name: Install dependencies
         run: bun install --frozen-lockfile
 
+      - name: Migrate test database
+        working-directory: packages/db
+        env:
+          DATABASE_URL: 'postgresql://postgres:postgres@localhost:5432/simstudio'
+        run: bun run db:migrate
+
+      - name: Run Pi search capability integration tests
+        working-directory: apps/sim
+        env:
+          DATABASE_URL: 'postgresql://postgres:postgres@localhost:5432/simstudio'
+        run: bun run test:pi-search-integration
+
       # Surfaces known CVEs in the dependency tree. Non-blocking until the
       # existing advisory backlog is triaged, then flip to a required gate by
       # removing continue-on-error.
@@ -143,6 +169,9 @@ jobs:
       - name: Verify agent stream capability docs are in sync
         run: bun run agent-stream-docs:check
 
+      - name: Verify Pi sandbox extension contract
+        run: bun run ./apps/sim/scripts/check-pi-sandbox-contract.ts
+
       - name: Migration safety (zero-downtime) audit
         run: |
           if [ "${{ github.event_name }}" = "pull_request" ]; then
@@ -197,11 +226,11 @@ jobs:
   # harmless), but the Turbo cache gets its own key: with a shared key, only
   # the last committer's new entries survive each run, so the test and build
   # Turbo entries would evict each other nondeterministically.
-  # Same next build as the app image, so it needs the same larger runner in
-  # GitHub mode: it peaks ~21.9 GB and SIGKILLs on the free 16 GB ones.
+  # Next static generation can exceed 32 GB on large app/docs diffs, so use the
+  # 64 GB Blacksmith runner; GitHub mode already uses the largest hosted option.
   build:
     name: Build App
-    runs-on: ${{ (vars.CI_PROVIDER == '' || vars.CI_PROVIDER == 'blacksmith') && 'blacksmith-8vcpu-ubuntu-2404' || 'linux-x64-8-core' }}
+    runs-on: ${{ (vars.CI_PROVIDER == '' || vars.CI_PROVIDER == 'blacksmith') && 'blacksmith-16vcpu-ubuntu-2404' || 'linux-x64-8-core' }}
     timeout-minutes: 15
 
     steps:

apps/docs/content/docs/en/workflows/blocks/pi.mdx

@@ -74,6 +74,17 @@ Your key for the chosen provider. On hosted Sim it is optional for Local Dev and
 - **Branch Name** *(advanced)* — the branch to push. Auto-generated when blank.
 - **Open as Draft PR** *(advanced)* — opens the PR as a draft. On by default.
 - **PR Title / PR Body** *(advanced)* — generated from the run when blank.
+- **Enable internet search** *(advanced)* — gives the Create PR agent one bounded `exa_search` operation. It is off by default and requires a workspace Exa key under **Settings → BYOK**.
+
+#### Create PR internet search
+
+Internet search is workspace-BYOK only. Sim checks for the workspace Exa key before starting a sandbox and never falls back to a Sim-hosted `EXA_API_KEY`. The Exa key stays in Sim; the sandbox receives a short-lived capability that can call only the bounded search route.
+
+The model can provide only a query and result count. It cannot use Exa research, content fetching, arbitrary Exa parameters, arbitrary Sim integrations, or the Exa key. Returned titles, URLs, snippets, and dates are marked as untrusted reference data and cannot change system instructions or credential boundaries.
+
+Search-enabled Create PR uses a fresh finalizer for committing and pushing, so the model process never receives the GitHub token. Its coding shell otherwise retains the normal Create PR network access. Agent-generated search queries are sent to Exa, so do not enable search for repositories whose data-handling policy prohibits that external processing.
+
+Credential filtering covers exact and URL-encoded active model, GitHub, Exa, and broker-capability values plus common credential prefixes. It cannot guarantee detection of arbitrarily transformed or encrypted secrets. Each run allows at most 10 sequential searches, 10 results per search, bounded snippets, and 500 KiB cumulative normalized search output.
 
 ### Review Code fields
 
@@ -149,6 +160,13 @@ Create PR runs in a sandbox image with the Pi CLI and git baked in.
 3. **Create a GitHub token** with permission to clone, push, and open a PR:
    - *Fine-grained:* select the repo, then **Contents: Read and write** + **Pull requests: Read and write**.
    - *Classic:* the **`repo`** scope. For org repos, authorize the token for SSO.
+4. **Optional: enable Exa search.** Store an Exa API key in the workspace under **Settings → BYOK → Exa**, then turn on **Enable internet search** in the block. Self-hosted deployments must set `PI_CREATE_PR_SEARCH_ENABLED=true`, `NEXT_PUBLIC_PI_CREATE_PR_SEARCH_ENABLED=true`, and `PI_EXA_BROKER_BASE_URL` to a canonical HTTPS app origin reachable from the cloud sandbox.
+
+Before enabling the server flag in a hosted environment, mint a short-lived smoke capability and run the protected **Pi Search Staging Smoke** workflow. The capability expires with the normal Pi execution window and should be provisioned immediately before dispatch rather than stored long-term.
+
+Operators must also schedule an authenticated hourly `GET /api/cron/cleanup-pi-search-capabilities` request. The Helm chart includes this CronJob; non-Helm hosted deployments must wire the equivalent external scheduler before enabling the feature.
+
+If search is enabled without an Exa BYOK key, the run stops before the coding agent starts and directs you to the BYOK settings. If the branch push succeeds but PR creation fails, Sim reports the pushed branch and does not delete it automatically.
 
 ### Review Code [#setup-cloud-code-review]
 

apps/sim/.env.example

@@ -17,6 +17,11 @@ BETTER_AUTH_URL=http://localhost:3000
 
 # NextJS (Required)
 NEXT_PUBLIC_APP_URL=http://localhost:3000
+# Optional Pi Create PR web search rollout. The broker base must be a public origin
+# reachable from the cloud sandbox; local development may use an HTTPS tunnel.
+# PI_CREATE_PR_SEARCH_ENABLED=false
+# NEXT_PUBLIC_PI_CREATE_PR_SEARCH_ENABLED=false
+# PI_EXA_BROKER_BASE_URL=https://app.example.com
 # INTERNAL_API_BASE_URL=http://sim-app.default.svc.cluster.local:3000  # Optional: internal URL for server-side /api self-calls; defaults to NEXT_PUBLIC_APP_URL
 # TRUSTED_ORIGINS=https://www.example.com,https://app.example.com  # Optional: comma-separated additional public origins to trust for auth (apex+www, alias domains). Merged into Better Auth trustedOrigins.
 # AUTH_TRUSTED_PROXIES=10.0.0.0/24,192.0.2.10  # Optional: reverse-proxy IPs/CIDRs in front of the app. Better Auth walks x-forwarded-for right to left, skips these hops, and uses the first untrusted address as the client IP (prevents forwarded-header spoofing). Use your proxies' actual addresses, not broad private ranges that also cover clients.

apps/sim/app/api/cron/cleanup-pi-search-capabilities/route.ts

@@ -0,0 +1,41 @@
+import { createLogger } from '@sim/logger'
+import { getErrorMessage } from '@sim/utils/errors'
+import { generateShortId } from '@sim/utils/id'
+import { type NextRequest, NextResponse } from 'next/server'
+import { cleanupPiSearchCapabilitiesContract } from '@/lib/api/contracts/pi-search-cleanup'
+import { parseRequest } from '@/lib/api/server'
+import { verifyCronAuth } from '@/lib/auth/internal'
+import { acquireLock, releaseLock } from '@/lib/core/config/redis'
+import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
+import { cleanupExpiredPiSearchCapabilities } from '@/lib/pi/exa-search/capabilities'
+
+const logger = createLogger('PiSearchCapabilityCleanup')
+const LOCK_KEY = 'pi-search-capability-cleanup'
+const LOCK_TTL_SECONDS = 120
+const BATCH_SIZE = 1_000
+const RETENTION_MS = 24 * 60 * 60 * 1000
+
+export const GET = withRouteHandler(async (request: NextRequest) => {
+  const authError = verifyCronAuth(request, 'Pi search capability cleanup')
+  if (authError) return authError
+  const parsed = await parseRequest(cleanupPiSearchCapabilitiesContract, request, {})
+  if (!parsed.success) return parsed.response
+  const requestId = generateShortId()
+  if (!(await acquireLock(LOCK_KEY, requestId, LOCK_TTL_SECONDS))) {
+    return NextResponse.json({ success: true, skipped: true }, { status: 202 })
+  }
+
+  try {
+    const deleted = await cleanupExpiredPiSearchCapabilities({
+      batchSize: BATCH_SIZE,
+      retentionBefore: new Date(Date.now() - RETENTION_MS),
+    })
+    return NextResponse.json({ success: true, deleted })
+  } catch (error) {
+    const message = getErrorMessage(error, 'Capability cleanup failed')
+    logger.error('Pi search capability cleanup failed', { requestId, error: message })
+    return NextResponse.json({ success: false, error: message }, { status: 500 })
+  } finally {
+    await releaseLock(LOCK_KEY, requestId).catch(() => {})
+  }
+})

apps/sim/app/api/internal/pi/exa-search/route.test.ts

@@ -0,0 +1,147 @@
+/**
+ * @vitest-environment node
+ */
+
+import { NextRequest } from 'next/server'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const {
+  mockAuthenticate,
+  mockReserve,
+  mockSettle,
+  mockRelease,
+  mockLeaseCurrent,
+  mockContainsSecret,
+  mockResolveByok,
+  mockExecuteSearch,
+} = vi.hoisted(() => ({
+  mockAuthenticate: vi.fn(),
+  mockReserve: vi.fn(),
+  mockSettle: vi.fn(),
+  mockRelease: vi.fn(),
+  mockLeaseCurrent: vi.fn(),
+  mockContainsSecret: vi.fn(),
+  mockResolveByok: vi.fn(),
+  mockExecuteSearch: vi.fn(),
+}))
+
+vi.mock('@/lib/pi/exa-search/capabilities', () => ({
+  authenticatePiSearchCapability: mockAuthenticate,
+  reservePiSearchCall: mockReserve,
+  settlePiSearchCall: mockSettle,
+  releasePiSearchCall: mockRelease,
+  isPiSearchLeaseCurrent: mockLeaseCurrent,
+  queryContainsProtectedSecret: mockContainsSecret,
+}))
+vi.mock('@/lib/api-key/byok', () => ({ resolveBYOKKeyById: mockResolveByok }))
+vi.mock('@/tools/exa/search-client', () => ({ executePiExaSearch: mockExecuteSearch }))
+
+import { POST } from '@/app/api/internal/pi/exa-search/route'
+
+function request(body: unknown, token = 'capability-token') {
+  return new NextRequest('http://localhost/api/internal/pi/exa-search', {
+    method: 'POST',
+    headers: {
+      authorization: `Bearer ${token}`,
+      'content-type': 'application/json',
+    },
+    body: JSON.stringify(body),
+  })
+}
+
+describe('Pi Exa search broker', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockAuthenticate.mockResolvedValue({
+      id: 'cap-1',
+      workspaceId: 'workspace-1',
+      providerKeyId: 'exa-key-id',
+      executionId: 'execution-1',
+      expiresAt: new Date(Date.now() + 60_000),
+      secretFingerprints: [],
+    })
+    mockContainsSecret.mockReturnValue(false)
+    mockReserve.mockResolvedValue({
+      capabilityId: 'cap-1',
+      token: 'lease',
+      generation: 1,
+      workspaceId: 'workspace-1',
+      executionId: 'execution-1',
+      expiresAt: new Date(Date.now() + 30_000),
+    })
+    mockResolveByok.mockResolvedValue({
+      status: 'found',
+      value: { apiKey: 'exa-user-key', isBYOK: true, keyId: 'exa-key-id' },
+    })
+    mockExecuteSearch.mockResolvedValue({
+      results: [{ title: 'Docs', url: 'https://example.com', snippet: 'text' }],
+    })
+    mockSettle.mockResolvedValue(true)
+    mockRelease.mockResolvedValue(true)
+    mockLeaseCurrent.mockResolvedValue(true)
+  })
+
+  it('authenticates before accepting the request body', async () => {
+    mockAuthenticate.mockResolvedValue(null)
+    const response = await POST(request({ query: 'x'.repeat(10_000) }), {} as never)
+    expect(response.status).toBe(401)
+    expect(mockReserve).not.toHaveBeenCalled()
+  })
+
+  it('uses the workspace BYOK key and returns the bounded result', async () => {
+    const response = await POST(request({ query: 'latest docs', numResults: 3 }), {} as never)
+    expect(response.status).toBe(200)
+    expect(mockResolveByok).toHaveBeenCalledWith('workspace-1', 'exa', 'exa-key-id')
+    expect(mockExecuteSearch).toHaveBeenCalledWith(
+      expect.objectContaining({
+        apiKey: 'exa-user-key',
+        query: 'latest docs',
+        numResults: 3,
+      })
+    )
+    expect(await response.json()).toEqual({
+      results: [{ title: 'Docs', url: 'https://example.com', snippet: 'text' }],
+    })
+  })
+
+  it('rejects protected query material before calling Exa', async () => {
+    mockContainsSecret.mockReturnValue(true)
+    const response = await POST(request({ query: 'leak this token' }), {} as never)
+    expect(response.status).toBe(400)
+    expect(mockReserve).not.toHaveBeenCalled()
+    expect(mockExecuteSearch).not.toHaveBeenCalled()
+  })
+
+  it('rejects fine-grained GitHub PAT patterns', async () => {
+    const response = await POST(
+      request({ query: 'search github_pat_abcdefghijklmnopqrstuvwxyz123456' }),
+      {} as never
+    )
+    expect(response.status).toBe(400)
+    expect(mockReserve).not.toHaveBeenCalled()
+  })
+
+  it('fails rather than falling back when workspace BYOK is missing', async () => {
+    mockResolveByok.mockResolvedValue({ status: 'missing' })
+    const response = await POST(request({ query: 'latest docs' }), {} as never)
+    expect(response.status).toBe(412)
+    expect(mockExecuteSearch).not.toHaveBeenCalled()
+    expect(mockRelease).toHaveBeenCalled()
+  })
+
+  it('rejects protected material returned by Exa before it reaches the agent', async () => {
+    mockExecuteSearch.mockResolvedValue({
+      results: [
+        {
+          title: 'leak',
+          url: 'https://example.com',
+          snippet: 'ghp_abcdefghijklmnopqrstuvwxyz123456',
+        },
+      ],
+    })
+    const response = await POST(request({ query: 'latest docs' }), {} as never)
+    expect(response.status).toBe(502)
+    expect(mockSettle).not.toHaveBeenCalled()
+    expect(mockRelease).toHaveBeenCalled()
+  })
+})