diff --git a/agents/context-pruner.ts b/agents/context-pruner.ts index 804f3cebb..c92687887 100644 --- a/agents/context-pruner.ts +++ b/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 diff --git a/agents/e2e/gravity-index.e2e.test.ts b/agents/e2e/gravity-index.e2e.test.ts new file mode 100644 index 000000000..64bdc9fd2 --- /dev/null +++ b/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 }, + ) +}) diff --git a/agents/types/tools.ts b/agents/types/tools.ts index 9ff49e007..d5ad31415 100644 --- a/agents/types/tools.ts +++ b/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 + } + | { + /** 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. */ diff --git a/common/src/constants/analytics-events.ts b/common/src/constants/analytics-events.ts index 5df0f2809..5db705be5 100644 --- a/common/src/constants/analytics-events.ts +++ b/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', diff --git a/common/src/templates/initial-agents-dir/types/tools.ts b/common/src/templates/initial-agents-dir/types/tools.ts index 9ff49e007..d5ad31415 100644 --- a/common/src/templates/initial-agents-dir/types/tools.ts +++ b/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 + } + | { + /** 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. */ diff --git a/common/src/tools/__tests__/compile-tool-definitions.test.ts b/common/src/tools/__tests__/compile-tool-definitions.test.ts new file mode 100644 index 000000000..a4766d836 --- /dev/null +++ b/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 {') + }) +}) diff --git a/common/src/tools/compile-tool-definitions.ts b/common/src/tools/compile-tool-definitions.ts index a2dc2c372..b84a49f95 100644 --- a/common/src/tools/compile-tool-definitions.ts +++ b/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(' | ') diff --git a/common/src/tools/constants.ts b/common/src/tools/constants.ts index f4a6d2ad4..452ba09b8 100644 --- a/common/src/tools/constants.ts +++ b/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', diff --git a/common/src/tools/list.ts b/common/src/tools/list.ts index 2671376ef..7834ebd51 100644 --- a/common/src/tools/list.ts +++ b/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, diff --git a/common/src/tools/params/tool/gravity-index.ts b/common/src/tools/params/tool/gravity-index.ts new file mode 100644 index 000000000..24ce9dbb5 --- /dev/null +++ b/common/src/tools/params/tool/gravity-index.ts @@ -0,0 +1,90 @@ +import z from 'zod/v4' + +import { gravityIndexInputSchema } from '../../../types/gravity-index' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' +import { jsonObjectSchema } from '../../../types/json' + +import type { $ToolParams } from '../../constants' + +const toolName = 'gravity_index' +const endsAgentStep = true + +const description = ` +Purpose: Use the Gravity Index to discover, inspect, and report integrations for third-party developer services such as databases, auth, payments, hosting, email, cache, monitoring, analytics, AI, storage, CMS, search, realtime, background jobs, infrastructure, CRM, support, productivity, commerce, video, webhooks, and SMS. + +Choose the action: +- \`search\`: Use when the user asks for a recommendation or when you need to choose a provider before integrating it. Returns a reasoned recommendation with install guidance, env vars, and a setup/conversion URL. Include stack and constraints in \`query\`. Pass \`search_id\` from a previous search for follow-up questions. +- \`browse\`: Use to list catalog services by \`category\` and/or keyword \`q\`. Good when the user wants options or a category-scoped picker. +- \`list_categories\`: Use to see available categories and service counts. +- \`get_service\`: Use when you already know a service slug and need full detail, env vars, website, docs URL, and install metadata. +- \`report_integration\`: Use after you have actually completed and verified an integration from a previous search. Pass the original \`search_id\` and the service slug as \`integrated_slug\`. + +Important setup-link behavior: +- Search results include \`conversion_url\`, the setup link the user should visit to create an account and get API credentials. +- Always show this link prominently as "Get your {service.name} API key" when credentials are needed. +- Do not replace it with the vendor homepage and do not auto-follow it. +- Ask the user to paste the resulting credentials back so you can finish setup. + +Implementation guidance: +- Gravity can help select a provider and identify required env vars, but install steps may be high-level. Use the returned \`docs_url\`, existing codebase conventions, and package/docs research to perform the actual integration. +- For browsing results, use \`get_service\` on promising slugs before making a final recommendation if details matter. + +Examples: +${$getNativeToolCallExampleString({ + toolName, + inputSchema: gravityIndexInputSchema, + input: { + action: 'search', + query: + 'transactional email API with a generous free tier for a Next.js app', + }, + endsAgentStep, +})} + +${$getNativeToolCallExampleString({ + toolName, + inputSchema: gravityIndexInputSchema, + input: { + action: 'browse', + category: 'Email', + q: 'send', + }, + endsAgentStep, +})} + +${$getNativeToolCallExampleString({ + toolName, + inputSchema: gravityIndexInputSchema, + input: { + action: 'get_service', + slug: 'sendgrid', + }, + endsAgentStep, +})} + +${$getNativeToolCallExampleString({ + toolName, + inputSchema: gravityIndexInputSchema, + input: { + action: 'report_integration', + search_id: 'search_id_from_previous_search', + integrated_slug: 'sendgrid', + }, + endsAgentStep, +})} +`.trim() + +export const gravityIndexParams = { + toolName, + endsAgentStep, + description, + inputSchema: gravityIndexInputSchema, + outputSchema: jsonToolResultSchema( + z.union([ + jsonObjectSchema, + z.object({ + errorMessage: z.string(), + }), + ]), + ), +} satisfies $ToolParams diff --git a/common/src/types/gravity-index.ts b/common/src/types/gravity-index.ts new file mode 100644 index 000000000..f0d8c2aeb --- /dev/null +++ b/common/src/types/gravity-index.ts @@ -0,0 +1,75 @@ +import z from 'zod/v4' + +import { jsonObjectSchema } from './json' + +export const gravityIndexInputSchema = z + .discriminatedUnion('action', [ + z.object({ + action: z.literal('search').describe('Search for the best service.'), + query: z + .string() + .min(1, 'Query cannot be empty') + .max(1000, 'Query cannot exceed 1000 characters') + .describe( + `What the user needs, including stack, constraints, and required capabilities when known. Example: "serverless database with branching for a Next.js app".`, + ), + search_id: z + .string() + .optional() + .describe('Continue a previous Gravity Index search as a follow-up.'), + context: jsonObjectSchema + .optional() + .describe( + 'Optional structured JSON context about the project, stack, or constraints.', + ), + }), + z.object({ + action: z + .literal('browse') + .describe('Browse catalog services by category and/or keyword.'), + category: z + .string() + .optional() + .describe( + 'Optional category filter, e.g. Database, Auth, Payments, Hosting, Email, Cache, Monitoring, Analytics, AI, Storage, CMS, Search, Realtime, Background Jobs, Infrastructure, CRM, Support, Productivity, Commerce, Video, Webhooks, SMS.', + ), + q: z + .string() + .optional() + .describe('Optional keyword filter, e.g. sendgrid or postgres.'), + }), + z.object({ + action: z + .literal('list_categories') + .describe('List every category with service counts.'), + }), + z.object({ + action: z + .literal('get_service') + .describe('Fetch full detail for a single service by slug.'), + slug: z + .string() + .min(1, 'Slug cannot be empty') + .describe('Service slug, e.g. supabase, stripe, sendgrid.'), + }), + z.object({ + action: z + .literal('report_integration') + .describe('Report that an integration from a prior search was done.'), + search_id: z + .string() + .min(1, 'search_id cannot be empty') + .describe('search_id from the earlier search result.'), + integrated_slug: z + .string() + .min(1, 'integrated_slug cannot be empty') + .describe('Slug of the service that was actually integrated.'), + }), + ]) + .describe(`Use the Gravity Index catalog and conversion API.`) + +export type GravityIndexInput = z.infer + +export const gravityIndexActionRequiresApiKey = ( + action: GravityIndexInput['action'], +) => action === 'search' || action === 'report_integration' diff --git a/packages/agent-runtime/src/__tests__/gravity-index-tool.test.ts b/packages/agent-runtime/src/__tests__/gravity-index-tool.test.ts new file mode 100644 index 000000000..3b87b475f --- /dev/null +++ b/packages/agent-runtime/src/__tests__/gravity-index-tool.test.ts @@ -0,0 +1,278 @@ +import { TEST_USER_ID } from '@codebuff/common/old-constants' +import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime' +import { getInitialSessionState } from '@codebuff/common/types/session-state' +import { promptSuccess } from '@codebuff/common/util/error' +import { + afterEach, + beforeEach, + describe, + expect, + mock, + spyOn, + test, +} from 'bun:test' + +import { createToolCallChunk, mockFileContext } from './test-utils' +import * as webApi from '../llm-api/codebuff-web-api' +import { runAgentStep } from '../run-agent-step' +import { assembleLocalAgentTemplates } from '../templates/agent-registry' + +import type { + AgentRuntimeDeps, + AgentRuntimeScopedDeps, +} from '@codebuff/common/types/contracts/agent-runtime' +import type { ParamsExcluding } from '@codebuff/common/types/function-params' +import type { StreamChunk } from '@codebuff/common/types/contracts/llm' + +let agentRuntimeImpl: AgentRuntimeDeps & AgentRuntimeScopedDeps +let runAgentStepBaseParams: ParamsExcluding< + typeof runAgentStep, + 'localAgentTemplates' | 'agentState' | 'prompt' | 'agentTemplate' +> + +function mockAgentStream(chunks: StreamChunk[]) { + runAgentStepBaseParams.promptAiSdkStream = async function* ({}) { + for (const chunk of chunks) { + yield chunk + } + return promptSuccess('mock-message-id') + } +} + +const gravityTestAgent = { + id: 'gravity-test-agent', + displayName: 'Gravity Test Agent', + model: 'openai/gpt-4o-mini', + toolNames: ['gravity_index', 'end_turn'], + systemPrompt: 'Use Gravity Index when choosing developer services.', +} + +describe('gravity_index tool', () => { + beforeEach(() => { + agentRuntimeImpl = { + ...TEST_AGENT_RUNTIME_IMPL, + } + runAgentStepBaseParams = { + ...agentRuntimeImpl, + additionalToolDefinitions: () => Promise.resolve({}), + agentType: 'gravity-test-agent', + ancestorRunIds: [], + clientSessionId: 'test-session', + fileContext: { + ...mockFileContext, + agentTemplates: { 'gravity-test-agent': gravityTestAgent }, + }, + fingerprintId: 'test-fingerprint', + onResponseChunk: () => {}, + repoId: undefined, + repoUrl: undefined, + runId: 'test-run-id', + signal: new AbortController().signal, + spawnParams: undefined, + system: 'Test system prompt', + tools: {}, + userId: TEST_USER_ID, + userInputId: 'test-input', + } + + runAgentStepBaseParams.requestFiles = async () => ({}) + runAgentStepBaseParams.requestOptionalFile = async () => null + runAgentStepBaseParams.requestToolCall = async () => ({ + output: [{ type: 'json', value: 'Tool call success' }], + }) + runAgentStepBaseParams.promptAiSdk = async function () { + return promptSuccess('Test response') + } + }) + + afterEach(() => { + mock.restore() + }) + + test('calls Gravity Index facade with the query', async () => { + const spy = spyOn(webApi, 'callGravityIndexAPI').mockResolvedValue({ + result: { + search_id: 'search-1', + recommendation: { name: 'SendGrid', slug: 'sendgrid' }, + conversion_url: 'https://index.trygravity.ai/go/test', + }, + }) + + mockAgentStream([ + createToolCallChunk('gravity_index', { + action: 'search', + query: 'transactional email for Next.js', + }), + createToolCallChunk('end_turn', {}), + ]) + + const sessionState = getInitialSessionState( + runAgentStepBaseParams.fileContext, + ) + const agentState = { + ...sessionState.mainAgentState, + agentType: 'gravity-test-agent', + } + const { agentTemplates } = assembleLocalAgentTemplates({ + ...agentRuntimeImpl, + fileContext: runAgentStepBaseParams.fileContext, + }) + + await runAgentStep({ + ...runAgentStepBaseParams, + localAgentTemplates: agentTemplates, + agentTemplate: agentTemplates['gravity-test-agent'], + agentState, + prompt: 'Find an email provider', + }) + + expect(spy).toHaveBeenCalledWith( + expect.objectContaining({ + input: { + action: 'search', + query: 'transactional email for Next.js', + }, + }), + ) + }) + + test('stores recommendation and conversion URL in tool output', async () => { + spyOn(webApi, 'callGravityIndexAPI').mockResolvedValue({ + result: { + search_id: 'search-1', + recommendation: { + name: 'SendGrid', + slug: 'sendgrid', + category: 'Email', + }, + reasoning: 'Good transactional email fit.', + conversion_url: 'https://index.trygravity.ai/go/test', + }, + }) + + mockAgentStream([ + createToolCallChunk('gravity_index', { + action: 'search', + query: 'transactional email for Next.js', + }), + createToolCallChunk('end_turn', {}), + ]) + + const sessionState = getInitialSessionState( + runAgentStepBaseParams.fileContext, + ) + const agentState = { + ...sessionState.mainAgentState, + agentType: 'gravity-test-agent', + } + const { agentTemplates } = assembleLocalAgentTemplates({ + ...agentRuntimeImpl, + fileContext: runAgentStepBaseParams.fileContext, + }) + + const { agentState: newAgentState } = await runAgentStep({ + ...runAgentStepBaseParams, + localAgentTemplates: agentTemplates, + agentTemplate: agentTemplates['gravity-test-agent'], + agentState, + prompt: 'Find an email provider', + }) + + const toolMsgs = newAgentState.messageHistory.filter( + (m) => m.role === 'tool' && m.toolName === 'gravity_index', + ) + expect(toolMsgs.length).toBeGreaterThan(0) + const last = JSON.stringify(toolMsgs[toolMsgs.length - 1].content) + expect(last).toContain('SendGrid') + expect(last).toContain('https://index.trygravity.ai/go/test') + }) + + test('surfaces API errors in tool output', async () => { + spyOn(webApi, 'callGravityIndexAPI').mockResolvedValue({ + error: 'Gravity Index is not configured', + }) + + mockAgentStream([ + createToolCallChunk('gravity_index', { + action: 'search', + query: 'transactional email for Next.js', + }), + createToolCallChunk('end_turn', {}), + ]) + + const sessionState = getInitialSessionState( + runAgentStepBaseParams.fileContext, + ) + const agentState = { + ...sessionState.mainAgentState, + agentType: 'gravity-test-agent', + } + const { agentTemplates } = assembleLocalAgentTemplates({ + ...agentRuntimeImpl, + fileContext: runAgentStepBaseParams.fileContext, + }) + + const { agentState: newAgentState } = await runAgentStep({ + ...runAgentStepBaseParams, + localAgentTemplates: agentTemplates, + agentTemplate: agentTemplates['gravity-test-agent'], + agentState, + prompt: 'Find an email provider', + }) + + const toolMsgs = newAgentState.messageHistory.filter( + (m) => m.role === 'tool' && m.toolName === 'gravity_index', + ) + const last = JSON.stringify(toolMsgs[toolMsgs.length - 1].content) + expect(last).toContain('errorMessage') + expect(last).toContain('Gravity Index is not configured') + }) + + test('passes non-search actions through the unified facade', async () => { + const spy = spyOn(webApi, 'callGravityIndexAPI').mockResolvedValue({ + result: { + services: [{ name: 'SendGrid', slug: 'sendgrid' }], + total: 1, + }, + }) + + mockAgentStream([ + createToolCallChunk('gravity_index', { + action: 'browse', + category: 'Email', + q: 'send', + }), + createToolCallChunk('end_turn', {}), + ]) + + const sessionState = getInitialSessionState( + runAgentStepBaseParams.fileContext, + ) + const agentState = { + ...sessionState.mainAgentState, + agentType: 'gravity-test-agent', + } + const { agentTemplates } = assembleLocalAgentTemplates({ + ...agentRuntimeImpl, + fileContext: runAgentStepBaseParams.fileContext, + }) + + await runAgentStep({ + ...runAgentStepBaseParams, + localAgentTemplates: agentTemplates, + agentTemplate: agentTemplates['gravity-test-agent'], + agentState, + prompt: 'Browse email providers', + }) + + expect(spy).toHaveBeenCalledWith( + expect.objectContaining({ + input: { + action: 'browse', + category: 'Email', + q: 'send', + }, + }), + ) + }) +}) diff --git a/packages/agent-runtime/src/llm-api/codebuff-web-api.ts b/packages/agent-runtime/src/llm-api/codebuff-web-api.ts index 61b77fd75..a4b81c997 100644 --- a/packages/agent-runtime/src/llm-api/codebuff-web-api.ts +++ b/packages/agent-runtime/src/llm-api/codebuff-web-api.ts @@ -1,6 +1,7 @@ import { withTimeout } from '@codebuff/common/util/promise' import type { ClientEnv, CiEnv } from '@codebuff/common/types/contracts/env' +import type { JSONObject } from '@codebuff/common/types/json' import type { Logger } from '@codebuff/common/types/contracts/logger' const FETCH_TIMEOUT_MS = 30_000 @@ -36,14 +37,17 @@ const getNumberField = (value: unknown, key: string): number | undefined => { } const callCodebuffV1 = async (params: { - endpoint: '/api/v1/web-search' | '/api/v1/docs-search' + endpoint: + | '/api/v1/web-search' + | '/api/v1/docs-search' + | '/api/v1/gravity-index' payload: unknown fetch: typeof globalThis.fetch logger: Logger env: CodebuffWebApiEnv baseUrl?: string apiKey?: string - requestName: 'web-search' | 'docs-search' + requestName: 'web-search' | 'docs-search' | 'gravity-index' }): Promise<{ json?: unknown; error?: string; creditsUsed?: number }> => { const { endpoint, payload, fetch, logger, env, requestName } = params const baseUrl = params.baseUrl ?? env.clientEnv.NEXT_PUBLIC_CODEBUFF_APP_URL @@ -226,6 +230,43 @@ export async function callDocsSearchAPI(params: { return { error: error ?? 'Invalid response format' } } +export async function callGravityIndexAPI(params: { + input: JSONObject + fetch: typeof globalThis.fetch + logger: Logger + env: CodebuffWebApiEnv + baseUrl?: string + apiKey?: string +}): Promise<{ + result?: JSONObject + error?: string + creditsUsed?: number +}> { + const { input, fetch, logger, env } = params + + const res = await callCodebuffV1({ + endpoint: '/api/v1/gravity-index', + payload: input, + fetch, + logger, + env, + baseUrl: params.baseUrl, + apiKey: params.apiKey, + requestName: 'gravity-index', + }) + if (res.error) return { error: res.error } + + if (res.json && typeof res.json === 'object' && !Array.isArray(res.json)) { + return { + result: res.json as JSONObject, + creditsUsed: res.creditsUsed, + } + } + + const error = getStringField(res.json, 'error') + return { error: error ?? 'Invalid response format' } +} + export async function callTokenCountAPI(params: { messages: unknown[] system?: string diff --git a/packages/agent-runtime/src/tools/handlers/list.ts b/packages/agent-runtime/src/tools/handlers/list.ts index 148be8438..654366996 100644 --- a/packages/agent-runtime/src/tools/handlers/list.ts +++ b/packages/agent-runtime/src/tools/handlers/list.ts @@ -8,6 +8,7 @@ import { handleCreatePlan } from './tool/create-plan' import { handleEndTurn } from './tool/end-turn' import { handleFindFiles } from './tool/find-files' import { handleGlob } from './tool/glob' +import { handleGravityIndex } from './tool/gravity-index' import { handleListDirectory } from './tool/list-directory' import { handleLookupAgentInfo } from './tool/lookup-agent-info' import { handleProposeStrReplace } from './tool/propose-str-replace' @@ -54,6 +55,7 @@ export const codebuffToolHandlers = { end_turn: handleEndTurn, find_files: handleFindFiles, glob: handleGlob, + gravity_index: handleGravityIndex, list_directory: handleListDirectory, lookup_agent_info: handleLookupAgentInfo, propose_str_replace: handleProposeStrReplace, diff --git a/packages/agent-runtime/src/tools/handlers/tool/gravity-index.ts b/packages/agent-runtime/src/tools/handlers/tool/gravity-index.ts new file mode 100644 index 000000000..97aa88860 --- /dev/null +++ b/packages/agent-runtime/src/tools/handlers/tool/gravity-index.ts @@ -0,0 +1,137 @@ +import { jsonToolResult } from '@codebuff/common/util/messages' + +import { callGravityIndexAPI } from '../../../llm-api/codebuff-web-api' + +import type { CodebuffToolHandlerFunction } from '../handler-function-type' +import type { + CodebuffToolCall, + CodebuffToolOutput, +} from '@codebuff/common/tools/list' +import type { ClientEnv, CiEnv } from '@codebuff/common/types/contracts/env' +import type { JSONObject } from '@codebuff/common/types/json' +import type { Logger } from '@codebuff/common/types/contracts/logger' + +export const handleGravityIndex = (async (params: { + previousToolCallFinished: Promise + toolCall: CodebuffToolCall<'gravity_index'> + logger: Logger + apiKey: string + + agentStepId: string + clientSessionId: string + fingerprintId: string + repoId: string | undefined + userInputId: string + userId: string | undefined + + fetch: typeof globalThis.fetch + clientEnv: ClientEnv + ciEnv: CiEnv +}): Promise<{ + output: CodebuffToolOutput<'gravity_index'> + creditsUsed: number +}> => { + const { + previousToolCallFinished, + toolCall, + agentStepId, + apiKey, + clientSessionId, + fingerprintId, + logger, + repoId, + userId, + userInputId, + fetch, + clientEnv, + ciEnv, + } = params + const { action } = toolCall.input + + const startedAt = Date.now() + const gravityContext = { + toolCallId: toolCall.toolCallId, + action, + userId, + agentStepId, + clientSessionId, + fingerprintId, + userInputId, + repoId, + } + + await previousToolCallFinished + + let creditsUsed = 0 + try { + const webApi = await callGravityIndexAPI({ + input: toolCall.input as JSONObject, + fetch, + logger, + apiKey, + env: { clientEnv, ciEnv }, + }) + + if (webApi.error || !webApi.result) { + logger.warn( + { + ...gravityContext, + durationMs: Date.now() - startedAt, + success: false, + error: webApi.error, + }, + 'Gravity Index returned error', + ) + return { + output: jsonToolResult({ + errorMessage: webApi.error ?? 'Invalid Gravity Index response', + }), + creditsUsed, + } + } + + if (typeof webApi.creditsUsed === 'number') { + creditsUsed = webApi.creditsUsed + } + + logger.info( + { + ...gravityContext, + durationMs: Date.now() - startedAt, + recommendation: + typeof webApi.result.recommendation === 'object' + ? webApi.result.recommendation + : undefined, + creditsUsed, + success: true, + }, + 'Gravity Index request completed via web API', + ) + + return { + output: jsonToolResult(webApi.result), + creditsUsed, + } + } catch (error) { + const errorMessage = `Error calling Gravity Index action "${action}": ${ + error instanceof Error ? error.message : 'Unknown error' + }` + logger.error( + { + ...gravityContext, + error: + error instanceof Error + ? { + name: error.name, + message: error.message, + stack: error.stack, + } + : error, + durationMs: Date.now() - startedAt, + success: false, + }, + 'Gravity Index request failed with error', + ) + return { output: jsonToolResult({ errorMessage }), creditsUsed } + } +}) satisfies CodebuffToolHandlerFunction<'gravity_index'> diff --git a/web/src/app/api/v1/gravity-index/__tests__/gravity-index.test.ts b/web/src/app/api/v1/gravity-index/__tests__/gravity-index.test.ts new file mode 100644 index 000000000..079fb1a84 --- /dev/null +++ b/web/src/app/api/v1/gravity-index/__tests__/gravity-index.test.ts @@ -0,0 +1,398 @@ +import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test' +import { NextRequest } from 'next/server' + +import { postGravityIndex } from '../_post' + +import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics' +import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database' +import type { + Logger, + LoggerWithContextFn, +} from '@codebuff/common/types/contracts/logger' + +const testServerEnv = { GRAVITY_API_KEY: 'gravity-key' } + +describe('/api/v1/gravity-index POST endpoint', () => { + let mockLogger: Logger + let mockLoggerWithContext: LoggerWithContextFn + let mockTrackEvent: TrackEventFn + let mockGetUserInfoFromApiKey: GetUserInfoFromApiKeyFn + let mockFetch: typeof globalThis.fetch + let mockWarn: ReturnType + + beforeEach(() => { + mockWarn = mock(() => {}) + mockLogger = { + error: mock(() => {}), + warn: mockWarn, + info: mock(() => {}), + debug: mock(() => {}), + } + mockLoggerWithContext = mock(() => mockLogger) + mockTrackEvent = mock(() => {}) + mockGetUserInfoFromApiKey = mock(async ({ apiKey }) => + apiKey === 'valid' ? { id: 'user-1' } : null, + ) as GetUserInfoFromApiKeyFn + mockFetch = Object.assign( + mock(async () => + Response.json({ + search_id: 'search-1', + recommendation: { + name: 'SendGrid', + slug: 'sendgrid', + category: 'Email', + website_url: 'https://sendgrid.com', + docs_url: 'https://docs.sendgrid.com', + }, + reasoning: 'Best fit for transactional email.', + install: { + summary: 'Create an API key', + env_vars: ['SENDGRID_API_KEY'], + }, + conversion_url: 'https://index.trygravity.ai/go/test', + }), + ), + { preconnect: () => {} }, + ) as typeof fetch + }) + + afterEach(() => { + mock.restore() + }) + + test('401 when missing API key', async () => { + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + body: JSON.stringify({ + action: 'search', + query: 'transactional email', + }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: testServerEnv, + }) + + expect(res.status).toBe(401) + expect(mockFetch).not.toHaveBeenCalled() + }) + + test('503 when Gravity API key is not configured', async () => { + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ + action: 'search', + query: 'transactional email', + }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: {}, + }) + + expect(res.status).toBe(503) + expect(mockFetch).not.toHaveBeenCalled() + }) + + test('catalog browse does not require Gravity API key', async () => { + mockFetch = Object.assign( + mock(async () => + Response.json({ + services: [{ name: 'SendGrid', slug: 'sendgrid' }], + total: 1, + }), + ), + { preconnect: () => {} }, + ) as typeof fetch + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ action: 'browse', category: 'Email' }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: {}, + }) + + expect(res.status).toBe(200) + expect( + (mockFetch as unknown as ReturnType).mock.calls[0][0], + ).toBe('https://index.trygravity.ai/services?category=Email') + }) + + test('sends Gravity API key only from server env', async () => { + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ + action: 'search', + query: 'transactional email', + platform_api_key: 'user-supplied-key', + }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: testServerEnv, + }) + + expect(res.status).toBe(200) + expect(mockFetch).toHaveBeenCalledTimes(1) + const [, init] = (mockFetch as unknown as ReturnType).mock + .calls[0] as [string, RequestInit] + expect(JSON.parse(String(init.body))).toEqual({ + query: 'transactional email', + platform_api_key: 'gravity-key', + }) + }) + + test('returns Gravity recommendation on success', async () => { + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ + action: 'search', + query: 'transactional email', + }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: testServerEnv, + }) + + expect(res.status).toBe(200) + const body = await res.json() + expect(body.recommendation.name).toBe('SendGrid') + expect(body.conversion_url).toBe('https://index.trygravity.ai/go/test') + expect(body.creditsUsed).toBe(0) + }) + + test('browse maps to GET /services with filters', async () => { + mockFetch = Object.assign( + mock(async () => + Response.json({ + services: [{ name: 'SendGrid', slug: 'sendgrid' }], + total: 1, + categories: ['Email'], + }), + ), + { preconnect: () => {} }, + ) as typeof fetch + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ action: 'browse', category: 'Email', q: 'send' }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: testServerEnv, + }) + + expect(res.status).toBe(200) + expect( + (mockFetch as unknown as ReturnType).mock.calls[0][0], + ).toBe('https://index.trygravity.ai/services?category=Email&q=send') + }) + + test('list_categories maps to GET /categories', async () => { + mockFetch = Object.assign( + mock(async () => Response.json({ categories: [], total: 0 })), + { preconnect: () => {} }, + ) as typeof fetch + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ action: 'list_categories' }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: testServerEnv, + }) + + expect(res.status).toBe(200) + expect( + (mockFetch as unknown as ReturnType).mock.calls[0][0], + ).toBe('https://index.trygravity.ai/categories') + }) + + test('get_service maps to GET /services/{slug}', async () => { + mockFetch = Object.assign( + mock(async () => Response.json({ name: 'SendGrid', slug: 'sendgrid' })), + { preconnect: () => {} }, + ) as typeof fetch + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ action: 'get_service', slug: 'sendgrid' }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: testServerEnv, + }) + + expect(res.status).toBe(200) + expect( + (mockFetch as unknown as ReturnType).mock.calls[0][0], + ).toBe('https://index.trygravity.ai/services/sendgrid') + }) + + test('report_integration maps to POST /integrations/report', async () => { + mockFetch = Object.assign( + mock(async () => + Response.json({ status: 'converted', slug: 'sendgrid' }), + ), + { preconnect: () => {} }, + ) as typeof fetch + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ + action: 'report_integration', + search_id: 'search-1', + integrated_slug: 'sendgrid', + }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: testServerEnv, + }) + + expect(res.status).toBe(200) + const [, init] = (mockFetch as unknown as ReturnType).mock + .calls[0] as [string, RequestInit] + expect(JSON.parse(String(init.body))).toEqual({ + search_id: 'search-1', + integrated_slug: 'sendgrid', + platform_api_key: 'gravity-key', + }) + }) + + test('502 when Gravity upstream fails', async () => { + mockFetch = Object.assign( + mock(async () => + Response.json({ error: 'bad request' }, { status: 400 }), + ), + { preconnect: () => {} }, + ) as typeof fetch + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ + action: 'search', + query: 'transactional email', + }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: testServerEnv, + }) + + expect(res.status).toBe(502) + expect(await res.json()).toEqual({ error: 'bad request' }) + }) + + test('redacts Gravity API key from upstream error responses and logs', async () => { + mockFetch = Object.assign( + mock( + async () => + new Response( + JSON.stringify({ + detail: [ + { + input: { + query: '', + platform_api_key: 'gravity-key', + }, + }, + ], + }), + { status: 422, headers: { 'Content-Type': 'application/json' } }, + ), + ), + { preconnect: () => {} }, + ) as typeof fetch + const req = new NextRequest('http://localhost:3000/api/v1/gravity-index', { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ + action: 'search', + query: 'transactional email', + }), + }) + + const res = await postGravityIndex({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: testServerEnv, + }) + + expect(res.status).toBe(502) + expect(JSON.stringify(await res.json())).not.toContain('gravity-key') + expect(JSON.stringify(mockWarn.mock.calls)).not.toContain('gravity-key') + expect(JSON.stringify(mockWarn.mock.calls)).toContain('[redacted]') + }) +}) diff --git a/web/src/app/api/v1/gravity-index/_post.ts b/web/src/app/api/v1/gravity-index/_post.ts new file mode 100644 index 000000000..0bd4da00f --- /dev/null +++ b/web/src/app/api/v1/gravity-index/_post.ts @@ -0,0 +1,263 @@ +import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events' +import { + gravityIndexActionRequiresApiKey, + gravityIndexInputSchema, +} from '@codebuff/common/types/gravity-index' +import { NextResponse } from 'next/server' + +import { parseJsonBody, requireUserFromApiKey } from '../_helpers' + +import type { GravityIndexInput } from '@codebuff/common/types/gravity-index' +import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics' +import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database' +import type { + Logger, + LoggerWithContextFn, +} from '@codebuff/common/types/contracts/logger' +import type { NextRequest } from 'next/server' + +const GRAVITY_INDEX_BASE_URL = 'https://index.trygravity.ai' +const FETCH_TIMEOUT_MS = 30_000 + +const tryParseJson = (text: string): unknown => { + try { + return JSON.parse(text) + } catch { + return null + } +} + +const getErrorMessage = (value: unknown): string | undefined => { + if (!value || typeof value !== 'object') return undefined + const record = value as Record + const message = record.error ?? record.message + return typeof message === 'string' ? message : undefined +} + +const redactGravityApiKey = ( + text: string, + gravityApiKey: string | undefined, +) => (gravityApiKey ? text.split(gravityApiKey).join('[redacted]') : text) + +const withQuery = ( + path: string, + params: Record, +) => { + const qs = new URLSearchParams() + for (const [key, value] of Object.entries(params)) { + if (value) qs.set(key, value) + } + const query = qs.toString() + return query ? `${path}?${query}` : path +} + +const requireGravityApiKey = (gravityApiKey: string | undefined) => { + if (!gravityApiKey) { + throw new Error('GRAVITY_API_KEY is not configured') + } + return gravityApiKey +} + +const buildGravityIndexRequest = ( + input: GravityIndexInput, + gravityApiKey: string | undefined, + signal: AbortSignal, +): Parameters => { + switch (input.action) { + case 'search': { + const apiKey = requireGravityApiKey(gravityApiKey) + return [ + `${GRAVITY_INDEX_BASE_URL}/search`, + { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + query: input.query, + ...(input.search_id ? { search_id: input.search_id } : {}), + ...(input.context ? { context: input.context } : {}), + platform_api_key: apiKey, + }), + signal, + }, + ] + } + case 'browse': + return [ + `${GRAVITY_INDEX_BASE_URL}${withQuery('/services', { + category: input.category, + q: input.q, + })}`, + { signal }, + ] + case 'list_categories': + return [`${GRAVITY_INDEX_BASE_URL}/categories`, { signal }] + case 'get_service': + return [ + `${GRAVITY_INDEX_BASE_URL}/services/${encodeURIComponent(input.slug)}`, + { signal }, + ] + case 'report_integration': { + const apiKey = requireGravityApiKey(gravityApiKey) + return [ + `${GRAVITY_INDEX_BASE_URL}/integrations/report`, + { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + search_id: input.search_id, + integrated_slug: input.integrated_slug, + platform_api_key: apiKey, + }), + signal, + }, + ] + } + } +} + +export async function postGravityIndex(params: { + req: NextRequest + getUserInfoFromApiKey: GetUserInfoFromApiKeyFn + logger: Logger + loggerWithContext: LoggerWithContextFn + trackEvent: TrackEventFn + fetch: typeof globalThis.fetch + serverEnv: { + GRAVITY_API_KEY?: string + } +}) { + const { + req, + getUserInfoFromApiKey, + loggerWithContext, + trackEvent, + fetch, + serverEnv, + } = params + const baseLogger = params.logger + + const parsedBody = await parseJsonBody({ + req, + schema: gravityIndexInputSchema, + logger: baseLogger, + trackEvent, + validationErrorEvent: AnalyticsEvent.GRAVITY_INDEX_VALIDATION_ERROR, + }) + if (!parsedBody.ok) return parsedBody.response + + const authed = await requireUserFromApiKey({ + req, + getUserInfoFromApiKey, + logger: baseLogger, + loggerWithContext, + trackEvent, + authErrorEvent: AnalyticsEvent.GRAVITY_INDEX_AUTH_ERROR, + }) + if (!authed.ok) return authed.response + + const { userId, logger } = authed.data + const input = parsedBody.data + const gravityApiKey = serverEnv.GRAVITY_API_KEY + + trackEvent({ + event: AnalyticsEvent.GRAVITY_INDEX_REQUEST, + userId, + properties: { action: input.action }, + logger, + }) + + if (gravityIndexActionRequiresApiKey(input.action) && !gravityApiKey) { + logger.error('GRAVITY_API_KEY is not configured') + trackEvent({ + event: AnalyticsEvent.GRAVITY_INDEX_ERROR, + userId, + properties: { reason: 'missing_gravity_api_key' }, + logger, + }) + return NextResponse.json( + { error: 'Gravity Index is not configured' }, + { status: 503 }, + ) + } + + const controller = new AbortController() + const timeout = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS) + + try { + const response = await fetch( + ...buildGravityIndexRequest(input, gravityApiKey, controller.signal), + ) + const text = await response.text() + const redactedText = redactGravityApiKey(text, gravityApiKey) + const json = tryParseJson(text) + + if (!response.ok) { + const upstreamError = getErrorMessage(json) + const error = + (upstreamError + ? redactGravityApiKey(upstreamError, gravityApiKey) + : redactedText) || 'Gravity Index failed' + logger.warn( + { + status: response.status, + statusText: response.statusText, + body: redactedText.slice(0, 500), + }, + 'Gravity Index upstream request failed', + ) + trackEvent({ + event: AnalyticsEvent.GRAVITY_INDEX_ERROR, + userId, + properties: { action: input.action, status: response.status, error }, + logger, + }) + return NextResponse.json({ error }, { status: 502 }) + } + + if (!json || typeof json !== 'object' || Array.isArray(json)) { + logger.warn( + { body: redactedText.slice(0, 500) }, + 'Invalid Gravity Index JSON', + ) + return NextResponse.json( + { error: 'Invalid Gravity Index response' }, + { status: 502 }, + ) + } + + return NextResponse.json({ + ...(json as Record), + creditsUsed: 0, + }) + } catch (error) { + const message = + error instanceof Error && error.name === 'AbortError' + ? 'Gravity Index request timed out' + : 'Error calling Gravity Index' + logger.error( + { + error: + error instanceof Error + ? { name: error.name, message: error.message, stack: error.stack } + : error, + }, + message, + ) + trackEvent({ + event: AnalyticsEvent.GRAVITY_INDEX_ERROR, + userId, + properties: { + action: input.action, + error: error instanceof Error ? error.message : 'Unknown error', + }, + logger, + }) + return NextResponse.json({ error: message }, { status: 502 }) + } finally { + clearTimeout(timeout) + } +} diff --git a/web/src/app/api/v1/gravity-index/route.ts b/web/src/app/api/v1/gravity-index/route.ts new file mode 100644 index 000000000..dbcfb7d73 --- /dev/null +++ b/web/src/app/api/v1/gravity-index/route.ts @@ -0,0 +1,21 @@ +import { trackEvent } from '@codebuff/common/analytics' +import { env } from '@codebuff/internal/env' + +import { postGravityIndex } from './_post' + +import type { NextRequest } from 'next/server' + +import { getUserInfoFromApiKey } from '@/db/user' +import { logger, loggerWithContext } from '@/util/logger' + +export async function POST(req: NextRequest) { + return postGravityIndex({ + req, + getUserInfoFromApiKey, + logger, + loggerWithContext, + trackEvent, + fetch, + serverEnv: { GRAVITY_API_KEY: env.GRAVITY_API_KEY }, + }) +}