Add gated Gravity Index tool (#567)

Co-authored-by: James Grugett <jahooma@gmail.com>

Author: brandonkachen

agents/context-pruner.ts

@@ -291,6 +291,14 @@ const definition: AgentDefinition = {
           const query = input.query as string | undefined
           return query ? `Web search: "${query}"` : 'Web search'
         }
+        case 'gravity_index': {
+          const query = input.query as string | undefined
+          const action = input.action as string | undefined
+          if (query) {
+            return `Gravity Index ${action ?? 'search'}: "${query}"`
+          }
+          return action ? `Gravity Index ${action}` : 'Gravity Index'
+        }
         case 'read_docs': {
           const libraryTitle = input.libraryTitle as string | undefined
           const topic = input.topic as string | undefined

agents/e2e/gravity-index.e2e.test.ts

@@ -0,0 +1,88 @@
+import fs from 'fs'
+import os from 'os'
+import path from 'path'
+
+import { API_KEY_ENV_VAR } from '@codebuff/common/constants/paths'
+import { CodebuffClient, type AgentDefinition } from '@codebuff/sdk'
+import { describe, expect, it } from 'bun:test'
+
+import base2Free from '../base2/base2-free'
+
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
+
+describe('Gravity Index SDK E2E', () => {
+  it(
+    'test agent uses gravity_index for third-party service selection',
+    async () => {
+      const apiKey = process.env[API_KEY_ENV_VAR]
+      if (!apiKey) {
+        console.warn(
+          `Skipping Gravity Index E2E: set ${API_KEY_ENV_VAR} to run.`,
+        )
+        return
+      }
+
+      const tmpDir = await fs.promises.mkdtemp(
+        path.join(os.tmpdir(), 'gravity-index-e2e-'),
+      )
+      const events: PrintModeEvent[] = []
+      const gravityIndexTestAgent = {
+        ...(base2Free as AgentDefinition),
+        id: 'base2-free-gravity-index-e2e',
+        displayName: 'Base2 Free Gravity Index E2E',
+        toolNames: [
+          ...((base2Free as AgentDefinition).toolNames ?? []),
+          'gravity_index',
+        ],
+        systemPrompt: `${(base2Free as AgentDefinition).systemPrompt}
+
+For this E2E test, use the gravity_index tool when asked to recommend third-party developer services.`,
+      } satisfies AgentDefinition
+
+      try {
+        const client = new CodebuffClient({
+          apiKey,
+          cwd: tmpDir,
+          projectFiles: {
+            'package.json': JSON.stringify({
+              scripts: {},
+              dependencies: { next: '^15.0.0' },
+            }),
+          },
+          agentDefinitions: [gravityIndexTestAgent],
+          handleEvent: (event) => {
+            events.push(event)
+          },
+        })
+
+        const run = await client.run({
+          agent: gravityIndexTestAgent.id,
+          prompt:
+            'Use the Gravity Index to recommend a transactional email API for a Next.js app. Include the tracked API-key signup URL from the tool result.',
+          maxAgentSteps: 4,
+        })
+
+        if (run.output.type === 'error') {
+          throw new Error(run.output.message)
+        }
+
+        const toolCalls = events.filter((event) => event.type === 'tool_call')
+        expect(
+          toolCalls.some(
+            (event) =>
+              'toolName' in event && event.toolName === 'gravity_index',
+          ),
+        ).toBe(true)
+
+        const outputText = events
+          .filter((event) => event.type === 'text')
+          .map((event) => ('text' in event ? event.text : ''))
+          .join('')
+        expect(outputText).toMatch(/https:\/\/index\.trygravity\.ai\/go\//)
+      } finally {
+        await fs.promises.rm(tmpDir, { recursive: true, force: true })
+      }
+    },
+    { timeout: 300_000 },
+  )
+})

agents/types/tools.ts

@@ -9,6 +9,7 @@ export type ToolName =
   | 'end_turn'
   | 'find_files'
   | 'glob'
+  | 'gravity_index'
   | 'list_directory'
   | 'lookup_agent_info'
   | 'propose_str_replace'
@@ -41,6 +42,7 @@ export interface ToolParamsMap {
   end_turn: EndTurnParams
   find_files: FindFilesParams
   glob: GlobParams
+  gravity_index: GravityIndexParams
   list_directory: ListDirectoryParams
   lookup_agent_info: LookupAgentInfoParams
   propose_str_replace: ProposeStrReplaceParams
@@ -156,6 +158,47 @@ export interface GlobParams {
   cwd?: string
 }
 
+/**
+ * Search, browse, inspect, or report integrations in the Gravity Index.
+ */
+export type GravityIndexParams =
+  | {
+      /** Search for the best service recommendation. */
+      action: 'search'
+      /** What the user needs, including stack, constraints, and required capabilities when known. */
+      query: string
+      /** Continue a previous Gravity Index search as a follow-up. */
+      search_id?: string
+      /** Optional structured context about the project, stack, or constraints. */
+      context?: Record<string, any>
+    }
+  | {
+      /** Browse catalog services by category and/or keyword. */
+      action: 'browse'
+      /** Optional category filter, e.g. Database, Auth, Payments, Hosting, Email, AI. */
+      category?: string
+      /** Optional keyword filter, e.g. sendgrid or postgres. */
+      q?: string
+    }
+  | {
+      /** List every category with service counts. */
+      action: 'list_categories'
+    }
+  | {
+      /** Fetch full detail for a single service by slug. */
+      action: 'get_service'
+      /** Service slug, e.g. supabase, stripe, sendgrid. */
+      slug: string
+    }
+  | {
+      /** Report that an integration from a prior search was completed. */
+      action: 'report_integration'
+      /** search_id from the earlier search result. */
+      search_id: string
+      /** Slug of the service that was actually integrated. */
+      integrated_slug: string
+    }
+
 /**
  * List files and directories in the specified path. Returns separate arrays of file names and directory names.
  */

common/src/constants/analytics-events.ts

@@ -124,6 +124,11 @@ export enum AnalyticsEvent {
   DOCS_SEARCH_INSUFFICIENT_CREDITS = 'api.docs_search_insufficient_credits',
   DOCS_SEARCH_ERROR = 'api.docs_search_error',
 
+  GRAVITY_INDEX_REQUEST = 'api.gravity_index_request',
+  GRAVITY_INDEX_AUTH_ERROR = 'api.gravity_index_auth_error',
+  GRAVITY_INDEX_VALIDATION_ERROR = 'api.gravity_index_validation_error',
+  GRAVITY_INDEX_ERROR = 'api.gravity_index_error',
+
   // Web - Feedback API
   FEEDBACK_SUBMITTED = 'api.feedback_submitted',
   FEEDBACK_AUTH_ERROR = 'api.feedback_auth_error',

common/src/templates/initial-agents-dir/types/tools.ts

@@ -9,6 +9,7 @@ export type ToolName =
   | 'end_turn'
   | 'find_files'
   | 'glob'
+  | 'gravity_index'
   | 'list_directory'
   | 'lookup_agent_info'
   | 'propose_str_replace'
@@ -41,6 +42,7 @@ export interface ToolParamsMap {
   end_turn: EndTurnParams
   find_files: FindFilesParams
   glob: GlobParams
+  gravity_index: GravityIndexParams
   list_directory: ListDirectoryParams
   lookup_agent_info: LookupAgentInfoParams
   propose_str_replace: ProposeStrReplaceParams
@@ -156,6 +158,47 @@ export interface GlobParams {
   cwd?: string
 }
 
+/**
+ * Search, browse, inspect, or report integrations in the Gravity Index.
+ */
+export type GravityIndexParams =
+  | {
+      /** Search for the best service recommendation. */
+      action: 'search'
+      /** What the user needs, including stack, constraints, and required capabilities when known. */
+      query: string
+      /** Continue a previous Gravity Index search as a follow-up. */
+      search_id?: string
+      /** Optional structured context about the project, stack, or constraints. */
+      context?: Record<string, any>
+    }
+  | {
+      /** Browse catalog services by category and/or keyword. */
+      action: 'browse'
+      /** Optional category filter, e.g. Database, Auth, Payments, Hosting, Email, AI. */
+      category?: string
+      /** Optional keyword filter, e.g. sendgrid or postgres. */
+      q?: string
+    }
+  | {
+      /** List every category with service counts. */
+      action: 'list_categories'
+    }
+  | {
+      /** Fetch full detail for a single service by slug. */
+      action: 'get_service'
+      /** Service slug, e.g. supabase, stripe, sendgrid. */
+      slug: string
+    }
+  | {
+      /** Report that an integration from a prior search was completed. */
+      action: 'report_integration'
+      /** search_id from the earlier search result. */
+      search_id: string
+      /** Slug of the service that was actually integrated. */
+      integrated_slug: string
+    }
+
 /**
  * List files and directories in the specified path. Returns separate arrays of file names and directory names.
  */

common/src/tools/__tests__/compile-tool-definitions.test.ts

@@ -0,0 +1,20 @@
+import { describe, expect, test } from 'bun:test'
+
+import { compileToolDefinitions } from '../compile-tool-definitions'
+
+describe('compileToolDefinitions', () => {
+  test('emits type aliases for root union tool schemas', () => {
+    const definitions = compileToolDefinitions()
+
+    expect(definitions).toContain('export type GravityIndexParams =')
+    expect(definitions).not.toContain('export interface GravityIndexParams {')
+    expect(definitions).toContain('"action": "search"')
+    expect(definitions).toContain('"action": "report_integration"')
+  })
+
+  test('keeps object tool schemas as interfaces', () => {
+    const definitions = compileToolDefinitions()
+
+    expect(definitions).toContain('export interface WebSearchParams {')
+  })
+})

common/src/tools/compile-tool-definitions.ts

@@ -18,18 +18,24 @@ export function compileToolDefinitions(): string {
 
       // Convert Zod schema to TypeScript interface using JSON schema
       let typeDefinition: string
+      let jsonSchema: unknown
       try {
-        const jsonSchema = z.toJSONSchema(parameterSchema, { io: 'input' })
+        jsonSchema = z.toJSONSchema(parameterSchema, { io: 'input' })
         typeDefinition = jsonSchemaToTypeScript(jsonSchema)
       } catch (error) {
         console.warn(`Failed to convert schema for ${toolName}:`, error)
         typeDefinition = '{ [key: string]: any }'
       }
 
+      const typeName = `${toPascalCase(toolName)}Params`
+      const declaration = canEmitInterface(jsonSchema)
+        ? `export interface ${typeName} ${typeDefinition}`
+        : `export type ${typeName} = ${typeDefinition}`
+
       return `/**
  * ${parameterSchema.description || `Parameters for ${toolName} tool`}
  */
-export interface ${toPascalCase(toolName)}Params ${typeDefinition}`
+${declaration}`
     })
     .join('\n\n')
 
@@ -89,10 +95,22 @@ function jsonSchemaToTypeScript(schema: any): string {
   return getTypeFromJsonSchema(schema)
 }
 
+function canEmitInterface(schema: any): boolean {
+  return (
+    schema.type === 'object' &&
+    !!schema.properties &&
+    !schema.anyOf &&
+    !schema.oneOf
+  )
+}
+
 /**
  * Gets TypeScript type from JSON Schema property
  */
 function getTypeFromJsonSchema(prop: any): string {
+  if (prop.const !== undefined) {
+    return JSON.stringify(prop.const)
+  }
   if (prop.type === 'string') {
     if (prop.enum) {
       return prop.enum.map((v: string) => `"${v}"`).join(' | ')

common/src/tools/constants.ts

@@ -30,6 +30,7 @@ export const toolNames = [
   'end_turn',
   'find_files',
   'glob',
+  'gravity_index',
   'list_directory',
   'lookup_agent_info',
   'propose_str_replace',
@@ -62,6 +63,7 @@ export const publishedTools = [
   'end_turn',
   'find_files',
   'glob',
+  'gravity_index',
   'list_directory',
   'lookup_agent_info',
   'propose_str_replace',

common/src/tools/list.ts

@@ -11,6 +11,7 @@ import { createPlanParams } from './params/tool/create-plan'
 import { endTurnParams } from './params/tool/end-turn'
 import { findFilesParams } from './params/tool/find-files'
 import { globParams } from './params/tool/glob'
+import { gravityIndexParams } from './params/tool/gravity-index'
 import { listDirectoryParams } from './params/tool/list-directory'
 import { lookupAgentInfoParams } from './params/tool/lookup-agent-info'
 import { proposeStrReplaceParams } from './params/tool/propose-str-replace'
@@ -49,6 +50,7 @@ export const toolParams = {
   end_turn: endTurnParams,
   find_files: findFilesParams,
   glob: globParams,
+  gravity_index: gravityIndexParams,
   list_directory: listDirectoryParams,
   lookup_agent_info: lookupAgentInfoParams,
   propose_str_replace: proposeStrReplaceParams,