From 679501d4d67aacd673f721d2f76bd511ec602821 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 29 Apr 2026 16:10:02 -0700 Subject: [PATCH 1/7] Add Gravity Index search tool --- agents/base2/base2.ts | 2 + agents/context-pruner.ts | 6 + agents/e2e/gravity-index.e2e.test.ts | 76 ++++++ agents/types/tools.ts | 10 + common/src/constants/analytics-events.ts | 5 + .../initial-agents-dir/types/tools.ts | 10 + common/src/tools/constants.ts | 2 + common/src/tools/list.ts | 2 + .../tools/params/tool/gravity-index-search.ts | 62 +++++ .../gravity-index-search-tool.test.ts | 222 ++++++++++++++++++ .../src/llm-api/codebuff-web-api.ts | 45 +++- .../agent-runtime/src/tools/handlers/list.ts | 2 + .../handlers/tool/gravity-index-search.ts | 136 +++++++++++ .../__tests__/gravity-index-search.test.ts | 197 ++++++++++++++++ .../app/api/v1/gravity-index/search/_post.ts | 178 ++++++++++++++ .../app/api/v1/gravity-index/search/route.ts | 21 ++ 16 files changed, 974 insertions(+), 2 deletions(-) create mode 100644 agents/e2e/gravity-index.e2e.test.ts create mode 100644 common/src/tools/params/tool/gravity-index-search.ts create mode 100644 packages/agent-runtime/src/__tests__/gravity-index-search-tool.test.ts create mode 100644 packages/agent-runtime/src/tools/handlers/tool/gravity-index-search.ts create mode 100644 web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts create mode 100644 web/src/app/api/v1/gravity-index/search/_post.ts create mode 100644 web/src/app/api/v1/gravity-index/search/route.ts diff --git a/agents/base2/base2.ts b/agents/base2/base2.ts index 1a81f948bf..6d32921aa7 100644 --- a/agents/base2/base2.ts +++ b/agents/base2/base2.ts @@ -70,6 +70,7 @@ export function createBase2( 'set_output', 'list_directory', 'glob', + 'gravity_index_search', ), spawnableAgents: buildArray( !isMax && 'file-picker', @@ -115,6 +116,7 @@ export function createBase2( - **Conventions:** Rigorously adhere to existing project conventions when reading or modifying code. Analyze surrounding code, tests, and configuration first. - **Libraries/Frameworks:** NEVER assume a library/framework is available or appropriate. Verify its established usage within the project (check imports, configuration files like 'package.json', 'Cargo.toml', 'requirements.txt', 'build.gradle', etc., or observe neighboring files) before employing it. +- **Third-party services:** When the user needs an external developer service (database, auth, email, payments, AI provider, vector store, observability, hosting, etc.), use the gravity_index_search tool before recommending or integrating a provider. Present its conversion_url prominently as "Get your {service.name} API key" when credentials are needed. - **Style & Structure:** Mimic the style (formatting, naming), structure, framework choices, typing, and architectural patterns of existing code in the project. - **Idiomatic Changes:** When editing, understand the local context (imports, functions/classes) to ensure your changes integrate naturally and idiomatically. - **Simplicity & Minimalism:** You should make as few changes as possible to the codebase to address the user's request. Only do what the user has asked for and no more. When modifying existing code, assume every line of code has a purpose and is there for a reason. Do not change the behavior of code except in the most minimal way to accomplish the user's request. diff --git a/agents/context-pruner.ts b/agents/context-pruner.ts index 804f3cebb5..5244668fd0 100644 --- a/agents/context-pruner.ts +++ b/agents/context-pruner.ts @@ -291,6 +291,12 @@ const definition: AgentDefinition = { const query = input.query as string | undefined return query ? `Web search: "${query}"` : 'Web search' } + case 'gravity_index_search': { + const query = input.query as string | undefined + return query + ? `Gravity Index search: "${query}"` + : 'Gravity Index search' + } 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 0000000000..d4cd0d50f7 --- /dev/null +++ b/agents/e2e/gravity-index.e2e.test.ts @@ -0,0 +1,76 @@ +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( + 'base2-free uses gravity_index_search 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[] = [] + + try { + const client = new CodebuffClient({ + apiKey, + cwd: tmpDir, + projectFiles: { + 'package.json': JSON.stringify({ + scripts: {}, + dependencies: { next: '^15.0.0' }, + }), + }, + agentDefinitions: [base2Free as AgentDefinition], + handleEvent: (event) => { + events.push(event) + }, + }) + + const run = await client.run({ + agent: base2Free.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_search', + ), + ).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 9ff49e007e..d58f6a37c3 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_search' | '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_search: GravityIndexSearchParams list_directory: ListDirectoryParams lookup_agent_info: LookupAgentInfoParams propose_str_replace: ProposeStrReplaceParams @@ -156,6 +158,14 @@ export interface GlobParams { cwd?: string } +/** + * Search the Gravity Index for a recommended developer service. + */ +export interface GravityIndexSearchParams { + /** A natural-language description of the third-party developer service the user needs, including stack and constraints when known. */ + query: 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 5df0f2809d..7c54e22d33 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_SEARCH_REQUEST = 'api.gravity_index_search_request', + GRAVITY_INDEX_SEARCH_AUTH_ERROR = 'api.gravity_index_search_auth_error', + GRAVITY_INDEX_SEARCH_VALIDATION_ERROR = 'api.gravity_index_search_validation_error', + GRAVITY_INDEX_SEARCH_ERROR = 'api.gravity_index_search_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 9ff49e007e..d58f6a37c3 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_search' | '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_search: GravityIndexSearchParams list_directory: ListDirectoryParams lookup_agent_info: LookupAgentInfoParams propose_str_replace: ProposeStrReplaceParams @@ -156,6 +158,14 @@ export interface GlobParams { cwd?: string } +/** + * Search the Gravity Index for a recommended developer service. + */ +export interface GravityIndexSearchParams { + /** A natural-language description of the third-party developer service the user needs, including stack and constraints when known. */ + query: string +} + /** * List files and directories in the specified path. Returns separate arrays of file names and directory names. */ diff --git a/common/src/tools/constants.ts b/common/src/tools/constants.ts index f4a6d2ad4e..4ba789e466 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_search', 'list_directory', 'lookup_agent_info', 'propose_str_replace', @@ -62,6 +63,7 @@ export const publishedTools = [ 'end_turn', 'find_files', 'glob', + 'gravity_index_search', 'list_directory', 'lookup_agent_info', 'propose_str_replace', diff --git a/common/src/tools/list.ts b/common/src/tools/list.ts index 2671376ef6..5a1db9901a 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 { gravityIndexSearchParams } from './params/tool/gravity-index-search' 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_search: gravityIndexSearchParams, list_directory: listDirectoryParams, lookup_agent_info: lookupAgentInfoParams, propose_str_replace: proposeStrReplaceParams, diff --git a/common/src/tools/params/tool/gravity-index-search.ts b/common/src/tools/params/tool/gravity-index-search.ts new file mode 100644 index 0000000000..e289b16564 --- /dev/null +++ b/common/src/tools/params/tool/gravity-index-search.ts @@ -0,0 +1,62 @@ +import z from 'zod/v4' + +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' +import { jsonObjectSchema } from '../../../types/json' + +import type { $ToolParams } from '../../constants' + +const toolName = 'gravity_index_search' +const endsAgentStep = true +const inputSchema = z + .object({ + query: z + .string() + .min(1, 'Query cannot be empty') + .describe( + `A natural-language description of the third-party developer service the user needs, including stack and constraints when known.`, + ), + }) + .describe(`Search the Gravity Index for a recommended developer service.`) + +const description = ` +Purpose: Search the Gravity Index, a curated catalog of developer services, when the user needs to choose or wire up a third-party integration such as a database, auth provider, email API, payments provider, AI provider, vector store, observability service, or similar SaaS. + +Use this before proposing a third-party service or integration stack. Include the user's app framework, language, budget/free-tier constraints, hosting environment, and required capabilities in the query when available. + +The result includes a recommended service, reasoning, install guidance, environment variables, and a tracked conversion URL. When you present the result to the user, surface the conversion URL prominently as "Get your {service.name} API key" so the user can create credentials. + +Example: +${$getNativeToolCallExampleString({ + toolName, + inputSchema, + input: { + query: + 'transactional email API with a generous free tier for a Next.js app', + }, + endsAgentStep, +})} + +${$getNativeToolCallExampleString({ + toolName, + inputSchema, + input: { + query: 'managed Postgres with database branching for preview environments', + }, + endsAgentStep, +})} +`.trim() + +export const gravityIndexSearchParams = { + toolName, + endsAgentStep, + description, + inputSchema, + outputSchema: jsonToolResultSchema( + z.union([ + jsonObjectSchema, + z.object({ + errorMessage: z.string(), + }), + ]), + ), +} satisfies $ToolParams diff --git a/packages/agent-runtime/src/__tests__/gravity-index-search-tool.test.ts b/packages/agent-runtime/src/__tests__/gravity-index-search-tool.test.ts new file mode 100644 index 0000000000..434b0397bd --- /dev/null +++ b/packages/agent-runtime/src/__tests__/gravity-index-search-tool.test.ts @@ -0,0 +1,222 @@ +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_search', 'end_turn'], + systemPrompt: 'Use Gravity Index when choosing developer services.', +} + +describe('gravity_index_search 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, 'callGravityIndexSearchAPI').mockResolvedValue({ + result: { + search_id: 'search-1', + recommendation: { name: 'SendGrid', slug: 'sendgrid' }, + conversion_url: 'https://index.trygravity.ai/go/test', + }, + }) + + mockAgentStream([ + createToolCallChunk('gravity_index_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({ query: 'transactional email for Next.js' }), + ) + }) + + test('stores recommendation and conversion URL in tool output', async () => { + spyOn(webApi, 'callGravityIndexSearchAPI').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_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_search', + ) + 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, 'callGravityIndexSearchAPI').mockResolvedValue({ + error: 'Gravity Index is not configured', + }) + + mockAgentStream([ + createToolCallChunk('gravity_index_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_search', + ) + const last = JSON.stringify(toolMsgs[toolMsgs.length - 1].content) + expect(last).toContain('errorMessage') + expect(last).toContain('Gravity Index is not configured') + }) +}) 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 61b77fd752..95d7bc121d 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/search' 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-search' }): 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 callGravityIndexSearchAPI(params: { + query: string + fetch: typeof globalThis.fetch + logger: Logger + env: CodebuffWebApiEnv + baseUrl?: string + apiKey?: string +}): Promise<{ + result?: JSONObject + error?: string + creditsUsed?: number +}> { + const { query, fetch, logger, env } = params + + const res = await callCodebuffV1({ + endpoint: '/api/v1/gravity-index/search', + payload: { query }, + fetch, + logger, + env, + baseUrl: params.baseUrl, + apiKey: params.apiKey, + requestName: 'gravity-index-search', + }) + 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 148be8438a..036903f251 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 { handleGravityIndexSearch } from './tool/gravity-index-search' 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_search: handleGravityIndexSearch, list_directory: handleListDirectory, lookup_agent_info: handleLookupAgentInfo, propose_str_replace: handleProposeStrReplace, diff --git a/packages/agent-runtime/src/tools/handlers/tool/gravity-index-search.ts b/packages/agent-runtime/src/tools/handlers/tool/gravity-index-search.ts new file mode 100644 index 0000000000..66950f1264 --- /dev/null +++ b/packages/agent-runtime/src/tools/handlers/tool/gravity-index-search.ts @@ -0,0 +1,136 @@ +import { jsonToolResult } from '@codebuff/common/util/messages' + +import { callGravityIndexSearchAPI } 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 { Logger } from '@codebuff/common/types/contracts/logger' + +export const handleGravityIndexSearch = (async (params: { + previousToolCallFinished: Promise + toolCall: CodebuffToolCall<'gravity_index_search'> + 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_search'> + creditsUsed: number +}> => { + const { + previousToolCallFinished, + toolCall, + agentStepId, + apiKey, + clientSessionId, + fingerprintId, + logger, + repoId, + userId, + userInputId, + fetch, + clientEnv, + ciEnv, + } = params + const { query } = toolCall.input + + const startedAt = Date.now() + const searchContext = { + toolCallId: toolCall.toolCallId, + query, + userId, + agentStepId, + clientSessionId, + fingerprintId, + userInputId, + repoId, + } + + await previousToolCallFinished + + let creditsUsed = 0 + try { + const webApi = await callGravityIndexSearchAPI({ + query, + fetch, + logger, + apiKey, + env: { clientEnv, ciEnv }, + }) + + if (webApi.error || !webApi.result) { + logger.warn( + { + ...searchContext, + durationMs: Date.now() - startedAt, + success: false, + error: webApi.error, + }, + 'Gravity Index search returned error', + ) + return { + output: jsonToolResult({ + errorMessage: webApi.error ?? 'Invalid Gravity Index response', + }), + creditsUsed, + } + } + + if (typeof webApi.creditsUsed === 'number') { + creditsUsed = webApi.creditsUsed + } + + logger.info( + { + ...searchContext, + durationMs: Date.now() - startedAt, + recommendation: + typeof webApi.result.recommendation === 'object' + ? webApi.result.recommendation + : undefined, + creditsUsed, + success: true, + }, + 'Gravity Index search completed via web API', + ) + + return { + output: jsonToolResult(webApi.result), + creditsUsed, + } + } catch (error) { + const errorMessage = `Error searching Gravity Index for "${query}": ${ + error instanceof Error ? error.message : 'Unknown error' + }` + logger.error( + { + ...searchContext, + error: + error instanceof Error + ? { + name: error.name, + message: error.message, + stack: error.stack, + } + : error, + durationMs: Date.now() - startedAt, + success: false, + }, + 'Gravity Index search failed with error', + ) + return { output: jsonToolResult({ errorMessage }), creditsUsed } + } +}) satisfies CodebuffToolHandlerFunction<'gravity_index_search'> diff --git a/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts b/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts new file mode 100644 index 0000000000..09bbba79eb --- /dev/null +++ b/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts @@ -0,0 +1,197 @@ +import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test' +import { NextRequest } from 'next/server' + +import { postGravityIndexSearch } 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/search POST endpoint', () => { + let mockLogger: Logger + let mockLoggerWithContext: LoggerWithContextFn + let mockTrackEvent: TrackEventFn + let mockGetUserInfoFromApiKey: GetUserInfoFromApiKeyFn + let mockFetch: typeof globalThis.fetch + + beforeEach(() => { + mockLogger = { + error: mock(() => {}), + warn: mock(() => {}), + 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/search', + { + method: 'POST', + body: JSON.stringify({ query: 'transactional email' }), + }, + ) + + const res = await postGravityIndexSearch({ + 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/search', + { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ query: 'transactional email' }), + }, + ) + + const res = await postGravityIndexSearch({ + req, + getUserInfoFromApiKey: mockGetUserInfoFromApiKey, + logger: mockLogger, + loggerWithContext: mockLoggerWithContext, + trackEvent: mockTrackEvent, + fetch: mockFetch, + serverEnv: {}, + }) + + expect(res.status).toBe(503) + expect(mockFetch).not.toHaveBeenCalled() + }) + + test('sends Gravity API key only from server env', async () => { + const req = new NextRequest( + 'http://localhost:3000/api/v1/gravity-index/search', + { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ + query: 'transactional email', + platform_api_key: 'user-supplied-key', + }), + }, + ) + + const res = await postGravityIndexSearch({ + 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/search', + { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ query: 'transactional email' }), + }, + ) + + const res = await postGravityIndexSearch({ + 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('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/search', + { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ query: 'transactional email' }), + }, + ) + + const res = await postGravityIndexSearch({ + 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' }) + }) +}) diff --git a/web/src/app/api/v1/gravity-index/search/_post.ts b/web/src/app/api/v1/gravity-index/search/_post.ts new file mode 100644 index 0000000000..9760ec7e47 --- /dev/null +++ b/web/src/app/api/v1/gravity-index/search/_post.ts @@ -0,0 +1,178 @@ +import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events' +import { NextResponse } from 'next/server' +import { z } from 'zod/v4' + +import { parseJsonBody, requireUserFromApiKey } from '../../_helpers' + +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 bodySchema = z.object({ + query: z.string().min(1, 'query is required'), +}) + +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 +} + +export async function postGravityIndexSearch(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: bodySchema, + logger: baseLogger, + trackEvent, + validationErrorEvent: AnalyticsEvent.GRAVITY_INDEX_SEARCH_VALIDATION_ERROR, + }) + if (!parsedBody.ok) return parsedBody.response + + const authed = await requireUserFromApiKey({ + req, + getUserInfoFromApiKey, + logger: baseLogger, + loggerWithContext, + trackEvent, + authErrorEvent: AnalyticsEvent.GRAVITY_INDEX_SEARCH_AUTH_ERROR, + }) + if (!authed.ok) return authed.response + + const { userId, logger } = authed.data + const { query } = parsedBody.data + const publisherKey = serverEnv.GRAVITY_API_KEY + + trackEvent({ + event: AnalyticsEvent.GRAVITY_INDEX_SEARCH_REQUEST, + userId, + properties: { queryLength: query.length }, + logger, + }) + + if (!publisherKey) { + logger.error('GRAVITY_API_KEY is not configured') + trackEvent({ + event: AnalyticsEvent.GRAVITY_INDEX_SEARCH_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(`${GRAVITY_INDEX_BASE_URL}/search`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + query, + platform_api_key: publisherKey, + }), + signal: controller.signal, + }) + + const text = await response.text() + const json = tryParseJson(text) + + if (!response.ok) { + const error = (getErrorMessage(json) ?? text) || 'Gravity Index failed' + logger.warn( + { + status: response.status, + statusText: response.statusText, + body: text.slice(0, 500), + }, + 'Gravity Index upstream request failed', + ) + trackEvent({ + event: AnalyticsEvent.GRAVITY_INDEX_SEARCH_ERROR, + userId, + properties: { status: response.status, error }, + logger, + }) + return NextResponse.json({ error }, { status: 502 }) + } + + if (!json || typeof json !== 'object' || Array.isArray(json)) { + logger.warn({ body: text.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 searching Gravity Index' + logger.error( + { + error: + error instanceof Error + ? { name: error.name, message: error.message, stack: error.stack } + : error, + }, + message, + ) + trackEvent({ + event: AnalyticsEvent.GRAVITY_INDEX_SEARCH_ERROR, + userId, + properties: { + 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/search/route.ts b/web/src/app/api/v1/gravity-index/search/route.ts new file mode 100644 index 0000000000..dd244852ad --- /dev/null +++ b/web/src/app/api/v1/gravity-index/search/route.ts @@ -0,0 +1,21 @@ +import { trackEvent } from '@codebuff/common/analytics' +import { env } from '@codebuff/internal/env' + +import { postGravityIndexSearch } 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 postGravityIndexSearch({ + req, + getUserInfoFromApiKey, + logger, + loggerWithContext, + trackEvent, + fetch, + serverEnv: { GRAVITY_API_KEY: env.GRAVITY_API_KEY }, + }) +} From 47347ed89eb6719052f23fe63058a7b7e1c72692 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 29 Apr 2026 17:39:03 -0700 Subject: [PATCH 2/7] Redact Gravity API key from upstream errors --- .../__tests__/gravity-index-search.test.ts | 48 ++++++++++++++++++- .../app/api/v1/gravity-index/search/_post.ts | 17 +++++-- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts b/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts index 09bbba79eb..815c79f309 100644 --- a/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts +++ b/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts @@ -18,11 +18,13 @@ describe('/api/v1/gravity-index/search POST endpoint', () => { let mockTrackEvent: TrackEventFn let mockGetUserInfoFromApiKey: GetUserInfoFromApiKeyFn let mockFetch: typeof globalThis.fetch + let mockWarn: ReturnType beforeEach(() => { + mockWarn = mock(() => {}) mockLogger = { error: mock(() => {}), - warn: mock(() => {}), + warn: mockWarn, info: mock(() => {}), debug: mock(() => {}), } @@ -194,4 +196,48 @@ describe('/api/v1/gravity-index/search POST endpoint', () => { 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/search', + { + method: 'POST', + headers: { Authorization: 'Bearer valid' }, + body: JSON.stringify({ query: 'transactional email' }), + }, + ) + + const res = await postGravityIndexSearch({ + 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/search/_post.ts b/web/src/app/api/v1/gravity-index/search/_post.ts index 9760ec7e47..f61c946d7c 100644 --- a/web/src/app/api/v1/gravity-index/search/_post.ts +++ b/web/src/app/api/v1/gravity-index/search/_post.ts @@ -34,6 +34,9 @@ const getErrorMessage = (value: unknown): string | undefined => { return typeof message === 'string' ? message : undefined } +const redactGravityApiKey = (text: string, gravityApiKey: string) => + text.split(gravityApiKey).join('[redacted]') + export async function postGravityIndexSearch(params: { req: NextRequest getUserInfoFromApiKey: GetUserInfoFromApiKeyFn @@ -116,15 +119,20 @@ export async function postGravityIndexSearch(params: { }) const text = await response.text() + const redactedText = redactGravityApiKey(text, publisherKey) const json = tryParseJson(text) if (!response.ok) { - const error = (getErrorMessage(json) ?? text) || 'Gravity Index failed' + const upstreamError = getErrorMessage(json) + const error = + (upstreamError + ? redactGravityApiKey(upstreamError, publisherKey) + : redactedText) || 'Gravity Index failed' logger.warn( { status: response.status, statusText: response.statusText, - body: text.slice(0, 500), + body: redactedText.slice(0, 500), }, 'Gravity Index upstream request failed', ) @@ -138,7 +146,10 @@ export async function postGravityIndexSearch(params: { } if (!json || typeof json !== 'object' || Array.isArray(json)) { - logger.warn({ body: text.slice(0, 500) }, 'Invalid Gravity Index JSON') + logger.warn( + { body: redactedText.slice(0, 500) }, + 'Invalid Gravity Index JSON', + ) return NextResponse.json( { error: 'Invalid Gravity Index response' }, { status: 502 }, From 7b7f0d00a334a1e7018046e2c1fa49280110c3be Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 29 Apr 2026 18:20:38 -0700 Subject: [PATCH 3/7] Consolidate Gravity Index tool actions --- agents/base2/base2.ts | 4 +- agents/context-pruner.ts | 10 +- agents/e2e/gravity-index.e2e.test.ts | 4 +- agents/types/tools.ts | 47 ++- common/src/constants/analytics-events.ts | 8 +- .../initial-agents-dir/types/tools.ts | 47 ++- common/src/tools/constants.ts | 4 +- common/src/tools/list.ts | 4 +- .../tools/params/tool/gravity-index-search.ts | 62 --- common/src/tools/params/tool/gravity-index.ts | 156 +++++++ ...ool.test.ts => gravity-index-tool.test.ts} | 78 +++- .../src/llm-api/codebuff-web-api.ts | 16 +- .../agent-runtime/src/tools/handlers/list.ts | 4 +- ...avity-index-search.ts => gravity-index.ts} | 35 +- .../__tests__/gravity-index.test.ts | 398 ++++++++++++++++++ .../v1/gravity-index/{search => }/_post.ts | 139 ++++-- .../v1/gravity-index/{search => }/route.ts | 4 +- .../__tests__/gravity-index-search.test.ts | 243 ----------- 18 files changed, 859 insertions(+), 404 deletions(-) delete mode 100644 common/src/tools/params/tool/gravity-index-search.ts create mode 100644 common/src/tools/params/tool/gravity-index.ts rename packages/agent-runtime/src/__tests__/{gravity-index-search-tool.test.ts => gravity-index-tool.test.ts} (75%) rename packages/agent-runtime/src/tools/handlers/tool/{gravity-index-search.ts => gravity-index.ts} (76%) create mode 100644 web/src/app/api/v1/gravity-index/__tests__/gravity-index.test.ts rename web/src/app/api/v1/gravity-index/{search => }/_post.ts (51%) rename web/src/app/api/v1/gravity-index/{search => }/route.ts (84%) delete mode 100644 web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts diff --git a/agents/base2/base2.ts b/agents/base2/base2.ts index 6d32921aa7..809c9a0892 100644 --- a/agents/base2/base2.ts +++ b/agents/base2/base2.ts @@ -70,7 +70,7 @@ export function createBase2( 'set_output', 'list_directory', 'glob', - 'gravity_index_search', + 'gravity_index', ), spawnableAgents: buildArray( !isMax && 'file-picker', @@ -116,7 +116,7 @@ export function createBase2( - **Conventions:** Rigorously adhere to existing project conventions when reading or modifying code. Analyze surrounding code, tests, and configuration first. - **Libraries/Frameworks:** NEVER assume a library/framework is available or appropriate. Verify its established usage within the project (check imports, configuration files like 'package.json', 'Cargo.toml', 'requirements.txt', 'build.gradle', etc., or observe neighboring files) before employing it. -- **Third-party services:** When the user needs an external developer service (database, auth, email, payments, AI provider, vector store, observability, hosting, etc.), use the gravity_index_search tool before recommending or integrating a provider. Present its conversion_url prominently as "Get your {service.name} API key" when credentials are needed. +- **Third-party services:** When the user needs an external developer service (database, auth, email, payments, AI provider, vector store, observability, hosting, etc.), use the gravity_index tool before recommending or integrating a provider. Present its conversion_url prominently as "Get your {service.name} API key" when credentials are needed. - **Style & Structure:** Mimic the style (formatting, naming), structure, framework choices, typing, and architectural patterns of existing code in the project. - **Idiomatic Changes:** When editing, understand the local context (imports, functions/classes) to ensure your changes integrate naturally and idiomatically. - **Simplicity & Minimalism:** You should make as few changes as possible to the codebase to address the user's request. Only do what the user has asked for and no more. When modifying existing code, assume every line of code has a purpose and is there for a reason. Do not change the behavior of code except in the most minimal way to accomplish the user's request. diff --git a/agents/context-pruner.ts b/agents/context-pruner.ts index 5244668fd0..c92687887c 100644 --- a/agents/context-pruner.ts +++ b/agents/context-pruner.ts @@ -291,11 +291,13 @@ const definition: AgentDefinition = { const query = input.query as string | undefined return query ? `Web search: "${query}"` : 'Web search' } - case 'gravity_index_search': { + case 'gravity_index': { const query = input.query as string | undefined - return query - ? `Gravity Index search: "${query}"` - : 'Gravity Index search' + 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 diff --git a/agents/e2e/gravity-index.e2e.test.ts b/agents/e2e/gravity-index.e2e.test.ts index d4cd0d50f7..72c750419c 100644 --- a/agents/e2e/gravity-index.e2e.test.ts +++ b/agents/e2e/gravity-index.e2e.test.ts @@ -12,7 +12,7 @@ import type { PrintModeEvent } from '@codebuff/common/types/print-mode' describe('Gravity Index SDK E2E', () => { it( - 'base2-free uses gravity_index_search for third-party service selection', + 'base2-free uses gravity_index for third-party service selection', async () => { const apiKey = process.env[API_KEY_ENV_VAR] if (!apiKey) { @@ -58,7 +58,7 @@ describe('Gravity Index SDK E2E', () => { expect( toolCalls.some( (event) => - 'toolName' in event && event.toolName === 'gravity_index_search', + 'toolName' in event && event.toolName === 'gravity_index', ), ).toBe(true) diff --git a/agents/types/tools.ts b/agents/types/tools.ts index d58f6a37c3..d5ad314150 100644 --- a/agents/types/tools.ts +++ b/agents/types/tools.ts @@ -9,7 +9,7 @@ export type ToolName = | 'end_turn' | 'find_files' | 'glob' - | 'gravity_index_search' + | 'gravity_index' | 'list_directory' | 'lookup_agent_info' | 'propose_str_replace' @@ -42,7 +42,7 @@ export interface ToolParamsMap { end_turn: EndTurnParams find_files: FindFilesParams glob: GlobParams - gravity_index_search: GravityIndexSearchParams + gravity_index: GravityIndexParams list_directory: ListDirectoryParams lookup_agent_info: LookupAgentInfoParams propose_str_replace: ProposeStrReplaceParams @@ -159,12 +159,45 @@ export interface GlobParams { } /** - * Search the Gravity Index for a recommended developer service. + * Search, browse, inspect, or report integrations in the Gravity Index. */ -export interface GravityIndexSearchParams { - /** A natural-language description of the third-party developer service the user needs, including stack and constraints when known. */ - query: string -} +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 7c54e22d33..5db705be58 100644 --- a/common/src/constants/analytics-events.ts +++ b/common/src/constants/analytics-events.ts @@ -124,10 +124,10 @@ export enum AnalyticsEvent { DOCS_SEARCH_INSUFFICIENT_CREDITS = 'api.docs_search_insufficient_credits', DOCS_SEARCH_ERROR = 'api.docs_search_error', - GRAVITY_INDEX_SEARCH_REQUEST = 'api.gravity_index_search_request', - GRAVITY_INDEX_SEARCH_AUTH_ERROR = 'api.gravity_index_search_auth_error', - GRAVITY_INDEX_SEARCH_VALIDATION_ERROR = 'api.gravity_index_search_validation_error', - GRAVITY_INDEX_SEARCH_ERROR = 'api.gravity_index_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', diff --git a/common/src/templates/initial-agents-dir/types/tools.ts b/common/src/templates/initial-agents-dir/types/tools.ts index d58f6a37c3..d5ad314150 100644 --- a/common/src/templates/initial-agents-dir/types/tools.ts +++ b/common/src/templates/initial-agents-dir/types/tools.ts @@ -9,7 +9,7 @@ export type ToolName = | 'end_turn' | 'find_files' | 'glob' - | 'gravity_index_search' + | 'gravity_index' | 'list_directory' | 'lookup_agent_info' | 'propose_str_replace' @@ -42,7 +42,7 @@ export interface ToolParamsMap { end_turn: EndTurnParams find_files: FindFilesParams glob: GlobParams - gravity_index_search: GravityIndexSearchParams + gravity_index: GravityIndexParams list_directory: ListDirectoryParams lookup_agent_info: LookupAgentInfoParams propose_str_replace: ProposeStrReplaceParams @@ -159,12 +159,45 @@ export interface GlobParams { } /** - * Search the Gravity Index for a recommended developer service. + * Search, browse, inspect, or report integrations in the Gravity Index. */ -export interface GravityIndexSearchParams { - /** A natural-language description of the third-party developer service the user needs, including stack and constraints when known. */ - query: string -} +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/constants.ts b/common/src/tools/constants.ts index 4ba789e466..452ba09b88 100644 --- a/common/src/tools/constants.ts +++ b/common/src/tools/constants.ts @@ -30,7 +30,7 @@ export const toolNames = [ 'end_turn', 'find_files', 'glob', - 'gravity_index_search', + 'gravity_index', 'list_directory', 'lookup_agent_info', 'propose_str_replace', @@ -63,7 +63,7 @@ export const publishedTools = [ 'end_turn', 'find_files', 'glob', - 'gravity_index_search', + '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 5a1db9901a..7834ebd514 100644 --- a/common/src/tools/list.ts +++ b/common/src/tools/list.ts @@ -11,7 +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 { gravityIndexSearchParams } from './params/tool/gravity-index-search' +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' @@ -50,7 +50,7 @@ export const toolParams = { end_turn: endTurnParams, find_files: findFilesParams, glob: globParams, - gravity_index_search: gravityIndexSearchParams, + gravity_index: gravityIndexParams, list_directory: listDirectoryParams, lookup_agent_info: lookupAgentInfoParams, propose_str_replace: proposeStrReplaceParams, diff --git a/common/src/tools/params/tool/gravity-index-search.ts b/common/src/tools/params/tool/gravity-index-search.ts deleted file mode 100644 index e289b16564..0000000000 --- a/common/src/tools/params/tool/gravity-index-search.ts +++ /dev/null @@ -1,62 +0,0 @@ -import z from 'zod/v4' - -import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' -import { jsonObjectSchema } from '../../../types/json' - -import type { $ToolParams } from '../../constants' - -const toolName = 'gravity_index_search' -const endsAgentStep = true -const inputSchema = z - .object({ - query: z - .string() - .min(1, 'Query cannot be empty') - .describe( - `A natural-language description of the third-party developer service the user needs, including stack and constraints when known.`, - ), - }) - .describe(`Search the Gravity Index for a recommended developer service.`) - -const description = ` -Purpose: Search the Gravity Index, a curated catalog of developer services, when the user needs to choose or wire up a third-party integration such as a database, auth provider, email API, payments provider, AI provider, vector store, observability service, or similar SaaS. - -Use this before proposing a third-party service or integration stack. Include the user's app framework, language, budget/free-tier constraints, hosting environment, and required capabilities in the query when available. - -The result includes a recommended service, reasoning, install guidance, environment variables, and a tracked conversion URL. When you present the result to the user, surface the conversion URL prominently as "Get your {service.name} API key" so the user can create credentials. - -Example: -${$getNativeToolCallExampleString({ - toolName, - inputSchema, - input: { - query: - 'transactional email API with a generous free tier for a Next.js app', - }, - endsAgentStep, -})} - -${$getNativeToolCallExampleString({ - toolName, - inputSchema, - input: { - query: 'managed Postgres with database branching for preview environments', - }, - endsAgentStep, -})} -`.trim() - -export const gravityIndexSearchParams = { - toolName, - endsAgentStep, - description, - inputSchema, - outputSchema: jsonToolResultSchema( - z.union([ - jsonObjectSchema, - z.object({ - errorMessage: z.string(), - }), - ]), - ), -} satisfies $ToolParams 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 0000000000..196fa75c58 --- /dev/null +++ b/common/src/tools/params/tool/gravity-index.ts @@ -0,0 +1,156 @@ +import z from 'zod/v4' + +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' +import { jsonObjectSchema } from '../../../types/json' + +import type { $ToolParams } from '../../constants' + +const toolName = 'gravity_index' +const endsAgentStep = true + +const inputSchema = 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: z + .record(z.string(), z.unknown()) + .optional() + .describe( + 'Optional structured 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.`) + +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, + input: { + action: 'search', + query: + 'transactional email API with a generous free tier for a Next.js app', + }, + endsAgentStep, +})} + +${$getNativeToolCallExampleString({ + toolName, + inputSchema, + input: { + action: 'browse', + category: 'Email', + q: 'send', + }, + endsAgentStep, +})} + +${$getNativeToolCallExampleString({ + toolName, + inputSchema, + input: { + action: 'get_service', + slug: 'sendgrid', + }, + endsAgentStep, +})} + +${$getNativeToolCallExampleString({ + toolName, + inputSchema, + input: { + action: 'report_integration', + search_id: 'search_id_from_previous_search', + integrated_slug: 'sendgrid', + }, + endsAgentStep, +})} +`.trim() + +export const gravityIndexParams = { + toolName, + endsAgentStep, + description, + inputSchema, + outputSchema: jsonToolResultSchema( + z.union([ + jsonObjectSchema, + z.object({ + errorMessage: z.string(), + }), + ]), + ), +} satisfies $ToolParams diff --git a/packages/agent-runtime/src/__tests__/gravity-index-search-tool.test.ts b/packages/agent-runtime/src/__tests__/gravity-index-tool.test.ts similarity index 75% rename from packages/agent-runtime/src/__tests__/gravity-index-search-tool.test.ts rename to packages/agent-runtime/src/__tests__/gravity-index-tool.test.ts index 434b0397bd..3b87b475f0 100644 --- a/packages/agent-runtime/src/__tests__/gravity-index-search-tool.test.ts +++ b/packages/agent-runtime/src/__tests__/gravity-index-tool.test.ts @@ -43,11 +43,11 @@ const gravityTestAgent = { id: 'gravity-test-agent', displayName: 'Gravity Test Agent', model: 'openai/gpt-4o-mini', - toolNames: ['gravity_index_search', 'end_turn'], + toolNames: ['gravity_index', 'end_turn'], systemPrompt: 'Use Gravity Index when choosing developer services.', } -describe('gravity_index_search tool', () => { +describe('gravity_index tool', () => { beforeEach(() => { agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL, @@ -90,7 +90,7 @@ describe('gravity_index_search tool', () => { }) test('calls Gravity Index facade with the query', async () => { - const spy = spyOn(webApi, 'callGravityIndexSearchAPI').mockResolvedValue({ + const spy = spyOn(webApi, 'callGravityIndexAPI').mockResolvedValue({ result: { search_id: 'search-1', recommendation: { name: 'SendGrid', slug: 'sendgrid' }, @@ -99,7 +99,8 @@ describe('gravity_index_search tool', () => { }) mockAgentStream([ - createToolCallChunk('gravity_index_search', { + createToolCallChunk('gravity_index', { + action: 'search', query: 'transactional email for Next.js', }), createToolCallChunk('end_turn', {}), @@ -126,12 +127,17 @@ describe('gravity_index_search tool', () => { }) expect(spy).toHaveBeenCalledWith( - expect.objectContaining({ query: 'transactional email for Next.js' }), + expect.objectContaining({ + input: { + action: 'search', + query: 'transactional email for Next.js', + }, + }), ) }) test('stores recommendation and conversion URL in tool output', async () => { - spyOn(webApi, 'callGravityIndexSearchAPI').mockResolvedValue({ + spyOn(webApi, 'callGravityIndexAPI').mockResolvedValue({ result: { search_id: 'search-1', recommendation: { @@ -145,7 +151,8 @@ describe('gravity_index_search tool', () => { }) mockAgentStream([ - createToolCallChunk('gravity_index_search', { + createToolCallChunk('gravity_index', { + action: 'search', query: 'transactional email for Next.js', }), createToolCallChunk('end_turn', {}), @@ -172,7 +179,7 @@ describe('gravity_index_search tool', () => { }) const toolMsgs = newAgentState.messageHistory.filter( - (m) => m.role === 'tool' && m.toolName === 'gravity_index_search', + (m) => m.role === 'tool' && m.toolName === 'gravity_index', ) expect(toolMsgs.length).toBeGreaterThan(0) const last = JSON.stringify(toolMsgs[toolMsgs.length - 1].content) @@ -181,12 +188,13 @@ describe('gravity_index_search tool', () => { }) test('surfaces API errors in tool output', async () => { - spyOn(webApi, 'callGravityIndexSearchAPI').mockResolvedValue({ + spyOn(webApi, 'callGravityIndexAPI').mockResolvedValue({ error: 'Gravity Index is not configured', }) mockAgentStream([ - createToolCallChunk('gravity_index_search', { + createToolCallChunk('gravity_index', { + action: 'search', query: 'transactional email for Next.js', }), createToolCallChunk('end_turn', {}), @@ -213,10 +221,58 @@ describe('gravity_index_search tool', () => { }) const toolMsgs = newAgentState.messageHistory.filter( - (m) => m.role === 'tool' && m.toolName === 'gravity_index_search', + (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 95d7bc121d..a4b81c9971 100644 --- a/packages/agent-runtime/src/llm-api/codebuff-web-api.ts +++ b/packages/agent-runtime/src/llm-api/codebuff-web-api.ts @@ -40,14 +40,14 @@ const callCodebuffV1 = async (params: { endpoint: | '/api/v1/web-search' | '/api/v1/docs-search' - | '/api/v1/gravity-index/search' + | '/api/v1/gravity-index' payload: unknown fetch: typeof globalThis.fetch logger: Logger env: CodebuffWebApiEnv baseUrl?: string apiKey?: string - requestName: 'web-search' | 'docs-search' | 'gravity-index-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 @@ -230,8 +230,8 @@ export async function callDocsSearchAPI(params: { return { error: error ?? 'Invalid response format' } } -export async function callGravityIndexSearchAPI(params: { - query: string +export async function callGravityIndexAPI(params: { + input: JSONObject fetch: typeof globalThis.fetch logger: Logger env: CodebuffWebApiEnv @@ -242,17 +242,17 @@ export async function callGravityIndexSearchAPI(params: { error?: string creditsUsed?: number }> { - const { query, fetch, logger, env } = params + const { input, fetch, logger, env } = params const res = await callCodebuffV1({ - endpoint: '/api/v1/gravity-index/search', - payload: { query }, + endpoint: '/api/v1/gravity-index', + payload: input, fetch, logger, env, baseUrl: params.baseUrl, apiKey: params.apiKey, - requestName: 'gravity-index-search', + requestName: 'gravity-index', }) if (res.error) return { error: res.error } diff --git a/packages/agent-runtime/src/tools/handlers/list.ts b/packages/agent-runtime/src/tools/handlers/list.ts index 036903f251..6543669963 100644 --- a/packages/agent-runtime/src/tools/handlers/list.ts +++ b/packages/agent-runtime/src/tools/handlers/list.ts @@ -8,7 +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 { handleGravityIndexSearch } from './tool/gravity-index-search' +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' @@ -55,7 +55,7 @@ export const codebuffToolHandlers = { end_turn: handleEndTurn, find_files: handleFindFiles, glob: handleGlob, - gravity_index_search: handleGravityIndexSearch, + 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-search.ts b/packages/agent-runtime/src/tools/handlers/tool/gravity-index.ts similarity index 76% rename from packages/agent-runtime/src/tools/handlers/tool/gravity-index-search.ts rename to packages/agent-runtime/src/tools/handlers/tool/gravity-index.ts index 66950f1264..97aa88860c 100644 --- a/packages/agent-runtime/src/tools/handlers/tool/gravity-index-search.ts +++ b/packages/agent-runtime/src/tools/handlers/tool/gravity-index.ts @@ -1,6 +1,6 @@ import { jsonToolResult } from '@codebuff/common/util/messages' -import { callGravityIndexSearchAPI } from '../../../llm-api/codebuff-web-api' +import { callGravityIndexAPI } from '../../../llm-api/codebuff-web-api' import type { CodebuffToolHandlerFunction } from '../handler-function-type' import type { @@ -8,11 +8,12 @@ import type { 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 handleGravityIndexSearch = (async (params: { +export const handleGravityIndex = (async (params: { previousToolCallFinished: Promise - toolCall: CodebuffToolCall<'gravity_index_search'> + toolCall: CodebuffToolCall<'gravity_index'> logger: Logger apiKey: string @@ -27,7 +28,7 @@ export const handleGravityIndexSearch = (async (params: { clientEnv: ClientEnv ciEnv: CiEnv }): Promise<{ - output: CodebuffToolOutput<'gravity_index_search'> + output: CodebuffToolOutput<'gravity_index'> creditsUsed: number }> => { const { @@ -45,12 +46,12 @@ export const handleGravityIndexSearch = (async (params: { clientEnv, ciEnv, } = params - const { query } = toolCall.input + const { action } = toolCall.input const startedAt = Date.now() - const searchContext = { + const gravityContext = { toolCallId: toolCall.toolCallId, - query, + action, userId, agentStepId, clientSessionId, @@ -63,8 +64,8 @@ export const handleGravityIndexSearch = (async (params: { let creditsUsed = 0 try { - const webApi = await callGravityIndexSearchAPI({ - query, + const webApi = await callGravityIndexAPI({ + input: toolCall.input as JSONObject, fetch, logger, apiKey, @@ -74,12 +75,12 @@ export const handleGravityIndexSearch = (async (params: { if (webApi.error || !webApi.result) { logger.warn( { - ...searchContext, + ...gravityContext, durationMs: Date.now() - startedAt, success: false, error: webApi.error, }, - 'Gravity Index search returned error', + 'Gravity Index returned error', ) return { output: jsonToolResult({ @@ -95,7 +96,7 @@ export const handleGravityIndexSearch = (async (params: { logger.info( { - ...searchContext, + ...gravityContext, durationMs: Date.now() - startedAt, recommendation: typeof webApi.result.recommendation === 'object' @@ -104,7 +105,7 @@ export const handleGravityIndexSearch = (async (params: { creditsUsed, success: true, }, - 'Gravity Index search completed via web API', + 'Gravity Index request completed via web API', ) return { @@ -112,12 +113,12 @@ export const handleGravityIndexSearch = (async (params: { creditsUsed, } } catch (error) { - const errorMessage = `Error searching Gravity Index for "${query}": ${ + const errorMessage = `Error calling Gravity Index action "${action}": ${ error instanceof Error ? error.message : 'Unknown error' }` logger.error( { - ...searchContext, + ...gravityContext, error: error instanceof Error ? { @@ -129,8 +130,8 @@ export const handleGravityIndexSearch = (async (params: { durationMs: Date.now() - startedAt, success: false, }, - 'Gravity Index search failed with error', + 'Gravity Index request failed with error', ) return { output: jsonToolResult({ errorMessage }), creditsUsed } } -}) satisfies CodebuffToolHandlerFunction<'gravity_index_search'> +}) 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 0000000000..079fb1a843 --- /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/search/_post.ts b/web/src/app/api/v1/gravity-index/_post.ts similarity index 51% rename from web/src/app/api/v1/gravity-index/search/_post.ts rename to web/src/app/api/v1/gravity-index/_post.ts index f61c946d7c..8f89b8fa9e 100644 --- a/web/src/app/api/v1/gravity-index/search/_post.ts +++ b/web/src/app/api/v1/gravity-index/_post.ts @@ -2,7 +2,7 @@ import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events' import { NextResponse } from 'next/server' import { z } from 'zod/v4' -import { parseJsonBody, requireUserFromApiKey } from '../../_helpers' +import { parseJsonBody, requireUserFromApiKey } from '../_helpers' import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics' import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database' @@ -15,9 +15,31 @@ import type { NextRequest } from 'next/server' const GRAVITY_INDEX_BASE_URL = 'https://index.trygravity.ai' const FETCH_TIMEOUT_MS = 30_000 -const bodySchema = z.object({ - query: z.string().min(1, 'query is required'), -}) +const bodySchema = z.discriminatedUnion('action', [ + z.object({ + action: z.literal('search'), + query: z.string().min(1, 'query is required').max(1000), + search_id: z.string().optional(), + context: z.record(z.string(), z.unknown()).optional(), + }), + z.object({ + action: z.literal('browse'), + category: z.string().optional(), + q: z.string().optional(), + }), + z.object({ + action: z.literal('list_categories'), + }), + z.object({ + action: z.literal('get_service'), + slug: z.string().min(1, 'slug is required'), + }), + z.object({ + action: z.literal('report_integration'), + search_id: z.string().min(1, 'search_id is required'), + integrated_slug: z.string().min(1, 'integrated_slug is required'), + }), +]) const tryParseJson = (text: string): unknown => { try { @@ -34,10 +56,24 @@ const getErrorMessage = (value: unknown): string | undefined => { return typeof message === 'string' ? message : undefined } -const redactGravityApiKey = (text: string, gravityApiKey: string) => - text.split(gravityApiKey).join('[redacted]') +const redactGravityApiKey = ( + text: string, + gravityApiKey: string | undefined, +) => (gravityApiKey ? text.split(gravityApiKey).join('[redacted]') : text) -export async function postGravityIndexSearch(params: { +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 +} + +export async function postGravityIndex(params: { req: NextRequest getUserInfoFromApiKey: GetUserInfoFromApiKeyFn logger: Logger @@ -63,7 +99,7 @@ export async function postGravityIndexSearch(params: { schema: bodySchema, logger: baseLogger, trackEvent, - validationErrorEvent: AnalyticsEvent.GRAVITY_INDEX_SEARCH_VALIDATION_ERROR, + validationErrorEvent: AnalyticsEvent.GRAVITY_INDEX_VALIDATION_ERROR, }) if (!parsedBody.ok) return parsedBody.response @@ -73,25 +109,28 @@ export async function postGravityIndexSearch(params: { logger: baseLogger, loggerWithContext, trackEvent, - authErrorEvent: AnalyticsEvent.GRAVITY_INDEX_SEARCH_AUTH_ERROR, + authErrorEvent: AnalyticsEvent.GRAVITY_INDEX_AUTH_ERROR, }) if (!authed.ok) return authed.response const { userId, logger } = authed.data - const { query } = parsedBody.data + const input = parsedBody.data const publisherKey = serverEnv.GRAVITY_API_KEY trackEvent({ - event: AnalyticsEvent.GRAVITY_INDEX_SEARCH_REQUEST, + event: AnalyticsEvent.GRAVITY_INDEX_REQUEST, userId, - properties: { queryLength: query.length }, + properties: { action: input.action }, logger, }) - if (!publisherKey) { + const needsPublisherKey = + input.action === 'search' || input.action === 'report_integration' + + if (needsPublisherKey && !publisherKey) { logger.error('GRAVITY_API_KEY is not configured') trackEvent({ - event: AnalyticsEvent.GRAVITY_INDEX_SEARCH_ERROR, + event: AnalyticsEvent.GRAVITY_INDEX_ERROR, userId, properties: { reason: 'missing_gravity_api_key' }, logger, @@ -106,17 +145,58 @@ export async function postGravityIndexSearch(params: { const timeout = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS) try { - const response = await fetch(`${GRAVITY_INDEX_BASE_URL}/search`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - body: JSON.stringify({ - query, - platform_api_key: publisherKey, - }), - signal: controller.signal, - }) + let response: Response + if (input.action === 'search') { + if (!publisherKey) { + throw new Error('GRAVITY_API_KEY is not configured') + } + response = await fetch(`${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: publisherKey, + }), + signal: controller.signal, + }) + } else if (input.action === 'browse') { + response = await fetch( + `${GRAVITY_INDEX_BASE_URL}${withQuery('/services', { + category: input.category, + q: input.q, + })}`, + { signal: controller.signal }, + ) + } else if (input.action === 'list_categories') { + response = await fetch(`${GRAVITY_INDEX_BASE_URL}/categories`, { + signal: controller.signal, + }) + } else if (input.action === 'get_service') { + response = await fetch( + `${GRAVITY_INDEX_BASE_URL}/services/${encodeURIComponent(input.slug)}`, + { signal: controller.signal }, + ) + } else { + if (!publisherKey) { + throw new Error('GRAVITY_API_KEY is not configured') + } + response = await fetch(`${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: publisherKey, + }), + signal: controller.signal, + }) + } const text = await response.text() const redactedText = redactGravityApiKey(text, publisherKey) @@ -137,9 +217,9 @@ export async function postGravityIndexSearch(params: { 'Gravity Index upstream request failed', ) trackEvent({ - event: AnalyticsEvent.GRAVITY_INDEX_SEARCH_ERROR, + event: AnalyticsEvent.GRAVITY_INDEX_ERROR, userId, - properties: { status: response.status, error }, + properties: { action: input.action, status: response.status, error }, logger, }) return NextResponse.json({ error }, { status: 502 }) @@ -164,7 +244,7 @@ export async function postGravityIndexSearch(params: { const message = error instanceof Error && error.name === 'AbortError' ? 'Gravity Index request timed out' - : 'Error searching Gravity Index' + : 'Error calling Gravity Index' logger.error( { error: @@ -175,9 +255,10 @@ export async function postGravityIndexSearch(params: { message, ) trackEvent({ - event: AnalyticsEvent.GRAVITY_INDEX_SEARCH_ERROR, + event: AnalyticsEvent.GRAVITY_INDEX_ERROR, userId, properties: { + action: input.action, error: error instanceof Error ? error.message : 'Unknown error', }, logger, diff --git a/web/src/app/api/v1/gravity-index/search/route.ts b/web/src/app/api/v1/gravity-index/route.ts similarity index 84% rename from web/src/app/api/v1/gravity-index/search/route.ts rename to web/src/app/api/v1/gravity-index/route.ts index dd244852ad..dbcfb7d73c 100644 --- a/web/src/app/api/v1/gravity-index/search/route.ts +++ b/web/src/app/api/v1/gravity-index/route.ts @@ -1,7 +1,7 @@ import { trackEvent } from '@codebuff/common/analytics' import { env } from '@codebuff/internal/env' -import { postGravityIndexSearch } from './_post' +import { postGravityIndex } from './_post' import type { NextRequest } from 'next/server' @@ -9,7 +9,7 @@ import { getUserInfoFromApiKey } from '@/db/user' import { logger, loggerWithContext } from '@/util/logger' export async function POST(req: NextRequest) { - return postGravityIndexSearch({ + return postGravityIndex({ req, getUserInfoFromApiKey, logger, diff --git a/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts b/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts deleted file mode 100644 index 815c79f309..0000000000 --- a/web/src/app/api/v1/gravity-index/search/__tests__/gravity-index-search.test.ts +++ /dev/null @@ -1,243 +0,0 @@ -import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test' -import { NextRequest } from 'next/server' - -import { postGravityIndexSearch } 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/search 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/search', - { - method: 'POST', - body: JSON.stringify({ query: 'transactional email' }), - }, - ) - - const res = await postGravityIndexSearch({ - 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/search', - { - method: 'POST', - headers: { Authorization: 'Bearer valid' }, - body: JSON.stringify({ query: 'transactional email' }), - }, - ) - - const res = await postGravityIndexSearch({ - req, - getUserInfoFromApiKey: mockGetUserInfoFromApiKey, - logger: mockLogger, - loggerWithContext: mockLoggerWithContext, - trackEvent: mockTrackEvent, - fetch: mockFetch, - serverEnv: {}, - }) - - expect(res.status).toBe(503) - expect(mockFetch).not.toHaveBeenCalled() - }) - - test('sends Gravity API key only from server env', async () => { - const req = new NextRequest( - 'http://localhost:3000/api/v1/gravity-index/search', - { - method: 'POST', - headers: { Authorization: 'Bearer valid' }, - body: JSON.stringify({ - query: 'transactional email', - platform_api_key: 'user-supplied-key', - }), - }, - ) - - const res = await postGravityIndexSearch({ - 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/search', - { - method: 'POST', - headers: { Authorization: 'Bearer valid' }, - body: JSON.stringify({ query: 'transactional email' }), - }, - ) - - const res = await postGravityIndexSearch({ - 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('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/search', - { - method: 'POST', - headers: { Authorization: 'Bearer valid' }, - body: JSON.stringify({ query: 'transactional email' }), - }, - ) - - const res = await postGravityIndexSearch({ - 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/search', - { - method: 'POST', - headers: { Authorization: 'Bearer valid' }, - body: JSON.stringify({ query: 'transactional email' }), - }, - ) - - const res = await postGravityIndexSearch({ - 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]') - }) -}) From da5c25e534c79bbaa4b294b44c5723bf33ba3aa1 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 29 Apr 2026 21:31:45 -0700 Subject: [PATCH 4/7] Share Gravity Index input schema --- common/src/tools/params/tool/gravity-index.ts | 78 +------- common/src/types/gravity-index.ts | 75 ++++++++ web/src/app/api/v1/gravity-index/_post.ts | 169 +++++++++--------- 3 files changed, 162 insertions(+), 160 deletions(-) create mode 100644 common/src/types/gravity-index.ts diff --git a/common/src/tools/params/tool/gravity-index.ts b/common/src/tools/params/tool/gravity-index.ts index 196fa75c58..24ce9dbb5e 100644 --- a/common/src/tools/params/tool/gravity-index.ts +++ b/common/src/tools/params/tool/gravity-index.ts @@ -1,5 +1,6 @@ import z from 'zod/v4' +import { gravityIndexInputSchema } from '../../../types/gravity-index' import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import { jsonObjectSchema } from '../../../types/json' @@ -8,73 +9,6 @@ import type { $ToolParams } from '../../constants' const toolName = 'gravity_index' const endsAgentStep = true -const inputSchema = 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: z - .record(z.string(), z.unknown()) - .optional() - .describe( - 'Optional structured 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.`) - 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. @@ -98,7 +32,7 @@ Implementation guidance: Examples: ${$getNativeToolCallExampleString({ toolName, - inputSchema, + inputSchema: gravityIndexInputSchema, input: { action: 'search', query: @@ -109,7 +43,7 @@ ${$getNativeToolCallExampleString({ ${$getNativeToolCallExampleString({ toolName, - inputSchema, + inputSchema: gravityIndexInputSchema, input: { action: 'browse', category: 'Email', @@ -120,7 +54,7 @@ ${$getNativeToolCallExampleString({ ${$getNativeToolCallExampleString({ toolName, - inputSchema, + inputSchema: gravityIndexInputSchema, input: { action: 'get_service', slug: 'sendgrid', @@ -130,7 +64,7 @@ ${$getNativeToolCallExampleString({ ${$getNativeToolCallExampleString({ toolName, - inputSchema, + inputSchema: gravityIndexInputSchema, input: { action: 'report_integration', search_id: 'search_id_from_previous_search', @@ -144,7 +78,7 @@ export const gravityIndexParams = { toolName, endsAgentStep, description, - inputSchema, + inputSchema: gravityIndexInputSchema, outputSchema: jsonToolResultSchema( z.union([ jsonObjectSchema, diff --git a/common/src/types/gravity-index.ts b/common/src/types/gravity-index.ts new file mode 100644 index 0000000000..f0d8c2aeba --- /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/web/src/app/api/v1/gravity-index/_post.ts b/web/src/app/api/v1/gravity-index/_post.ts index 8f89b8fa9e..0bd4da00f7 100644 --- a/web/src/app/api/v1/gravity-index/_post.ts +++ b/web/src/app/api/v1/gravity-index/_post.ts @@ -1,9 +1,13 @@ import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events' +import { + gravityIndexActionRequiresApiKey, + gravityIndexInputSchema, +} from '@codebuff/common/types/gravity-index' import { NextResponse } from 'next/server' -import { z } from 'zod/v4' 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 { @@ -15,32 +19,6 @@ import type { NextRequest } from 'next/server' const GRAVITY_INDEX_BASE_URL = 'https://index.trygravity.ai' const FETCH_TIMEOUT_MS = 30_000 -const bodySchema = z.discriminatedUnion('action', [ - z.object({ - action: z.literal('search'), - query: z.string().min(1, 'query is required').max(1000), - search_id: z.string().optional(), - context: z.record(z.string(), z.unknown()).optional(), - }), - z.object({ - action: z.literal('browse'), - category: z.string().optional(), - q: z.string().optional(), - }), - z.object({ - action: z.literal('list_categories'), - }), - z.object({ - action: z.literal('get_service'), - slug: z.string().min(1, 'slug is required'), - }), - z.object({ - action: z.literal('report_integration'), - search_id: z.string().min(1, 'search_id is required'), - integrated_slug: z.string().min(1, 'integrated_slug is required'), - }), -]) - const tryParseJson = (text: string): unknown => { try { return JSON.parse(text) @@ -73,6 +51,74 @@ const withQuery = ( 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 @@ -96,7 +142,7 @@ export async function postGravityIndex(params: { const parsedBody = await parseJsonBody({ req, - schema: bodySchema, + schema: gravityIndexInputSchema, logger: baseLogger, trackEvent, validationErrorEvent: AnalyticsEvent.GRAVITY_INDEX_VALIDATION_ERROR, @@ -115,7 +161,7 @@ export async function postGravityIndex(params: { const { userId, logger } = authed.data const input = parsedBody.data - const publisherKey = serverEnv.GRAVITY_API_KEY + const gravityApiKey = serverEnv.GRAVITY_API_KEY trackEvent({ event: AnalyticsEvent.GRAVITY_INDEX_REQUEST, @@ -124,10 +170,7 @@ export async function postGravityIndex(params: { logger, }) - const needsPublisherKey = - input.action === 'search' || input.action === 'report_integration' - - if (needsPublisherKey && !publisherKey) { + if (gravityIndexActionRequiresApiKey(input.action) && !gravityApiKey) { logger.error('GRAVITY_API_KEY is not configured') trackEvent({ event: AnalyticsEvent.GRAVITY_INDEX_ERROR, @@ -145,68 +188,18 @@ export async function postGravityIndex(params: { const timeout = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS) try { - let response: Response - if (input.action === 'search') { - if (!publisherKey) { - throw new Error('GRAVITY_API_KEY is not configured') - } - response = await fetch(`${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: publisherKey, - }), - signal: controller.signal, - }) - } else if (input.action === 'browse') { - response = await fetch( - `${GRAVITY_INDEX_BASE_URL}${withQuery('/services', { - category: input.category, - q: input.q, - })}`, - { signal: controller.signal }, - ) - } else if (input.action === 'list_categories') { - response = await fetch(`${GRAVITY_INDEX_BASE_URL}/categories`, { - signal: controller.signal, - }) - } else if (input.action === 'get_service') { - response = await fetch( - `${GRAVITY_INDEX_BASE_URL}/services/${encodeURIComponent(input.slug)}`, - { signal: controller.signal }, - ) - } else { - if (!publisherKey) { - throw new Error('GRAVITY_API_KEY is not configured') - } - response = await fetch(`${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: publisherKey, - }), - signal: controller.signal, - }) - } - + const response = await fetch( + ...buildGravityIndexRequest(input, gravityApiKey, controller.signal), + ) const text = await response.text() - const redactedText = redactGravityApiKey(text, publisherKey) + const redactedText = redactGravityApiKey(text, gravityApiKey) const json = tryParseJson(text) if (!response.ok) { const upstreamError = getErrorMessage(json) const error = (upstreamError - ? redactGravityApiKey(upstreamError, publisherKey) + ? redactGravityApiKey(upstreamError, gravityApiKey) : redactedText) || 'Gravity Index failed' logger.warn( { From 0add1f8322c1e7de76159c3393b6851f230cbc33 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 29 Apr 2026 22:12:33 -0700 Subject: [PATCH 5/7] Fix tool definition generation for union schemas --- .../compile-tool-definitions.test.ts | 20 +++++++++++++++++ common/src/tools/compile-tool-definitions.ts | 22 +++++++++++++++++-- 2 files changed, 40 insertions(+), 2 deletions(-) create mode 100644 common/src/tools/__tests__/compile-tool-definitions.test.ts 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 0000000000..a4766d8363 --- /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 a2dc2c372e..b84a49f955 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(' | ') From f3f61daa30c0e8523a58c0db6d029c5da70dce5e Mon Sep 17 00:00:00 2001 From: James Grugett Date: Thu, 30 Apr 2026 11:07:01 -0700 Subject: [PATCH 6/7] Gate Gravity Index from base agent --- agents/base2/base2.ts | 222 ++++++++++++++------------- agents/e2e/gravity-index.e2e.test.ts | 18 ++- 2 files changed, 129 insertions(+), 111 deletions(-) diff --git a/agents/base2/base2.ts b/agents/base2/base2.ts index 809c9a0892..9be93c8e12 100644 --- a/agents/base2/base2.ts +++ b/agents/base2/base2.ts @@ -30,11 +30,13 @@ export function createBase2( return { publisher, model, - providerOptions: isFree ? { - data_collection: 'deny', - } : { - only: ['amazon-bedrock'], - }, + providerOptions: isFree + ? { + data_collection: 'deny', + } + : { + only: ['amazon-bedrock'], + }, displayName: 'Buffy the Orchestrator', spawnerPrompt: 'Advanced base agent that orchestrates planning, editing, and reviewing for complex coding tasks', @@ -70,7 +72,6 @@ export function createBase2( 'set_output', 'list_directory', 'glob', - 'gravity_index', ), spawnableAgents: buildArray( !isMax && 'file-picker', @@ -103,11 +104,12 @@ export function createBase2( - **Spawn mentioned agents:** If the user uses "@AgentName" in their message, you must spawn that agent. - **Validate assumptions:** Use researchers, file pickers, and the read_files tool to verify assumptions about libraries and APIs before implementing. - **Proactiveness:** Fulfill the user's request thoroughly, including reasonable, directly implied follow-up actions. -- **Confirm Ambiguity/Expansion:** Do not take significant actions beyond the clear scope of the request without confirming with the user. If asked *how* to do something, explain first, don't just do it.${noAskUser +- **Confirm Ambiguity/Expansion:** Do not take significant actions beyond the clear scope of the request without confirming with the user. If asked *how* to do something, explain first, don't just do it.${ + noAskUser ? '' : ` - **Ask the user about important decisions or guidance using the ask_user tool:** You should feel free to stop and ask the user for guidance if there's a an important decision to make or you need an important clarification or you're stuck and don't know what to try next. Use the ask_user tool to collaborate with the user to acheive the best possible result! Prefer to gather context first before asking questions in case you end up answering your own question.` - } + } - **Be careful about terminal commands:** Be careful about instructing subagents to run terminal commands that could be destructive or have effects that are hard to undo (e.g. git push, git commit, running any scripts -- especially ones that could alter production environments (!), installing packages globally, etc). Don't run any of these effectful commands unless the user explicitly asks you to. - **Do what the user asks:** If the user asks you to do something, even running a risky terminal command, do it. - **Don't use set_output:** The set_output tool is for spawned subagents to report results. Don't use it yourself. @@ -116,7 +118,6 @@ export function createBase2( - **Conventions:** Rigorously adhere to existing project conventions when reading or modifying code. Analyze surrounding code, tests, and configuration first. - **Libraries/Frameworks:** NEVER assume a library/framework is available or appropriate. Verify its established usage within the project (check imports, configuration files like 'package.json', 'Cargo.toml', 'requirements.txt', 'build.gradle', etc., or observe neighboring files) before employing it. -- **Third-party services:** When the user needs an external developer service (database, auth, email, payments, AI provider, vector store, observability, hosting, etc.), use the gravity_index tool before recommending or integrating a provider. Present its conversion_url prominently as "Get your {service.name} API key" when credentials are needed. - **Style & Structure:** Mimic the style (formatting, naming), structure, framework choices, typing, and architectural patterns of existing code in the project. - **Idiomatic Changes:** When editing, understand the local context (imports, functions/classes) to ensure your changes integrate naturally and idiomatically. - **Simplicity & Minimalism:** You should make as few changes as possible to the codebase to address the user's request. Only do what the user has asked for and no more. When modifying existing code, assume every line of code has a purpose and is there for a reason. Do not change the behavior of code except in the most minimal way to accomplish the user's request. @@ -143,24 +144,25 @@ Use the spawn_agents tool to spawn specialized agents to help you complete the u - **Spawn multiple agents in parallel:** This increases the speed of your response **and** allows you to be more comprehensive by spawning more total agents to synthesize the best response. - **Sequence agents properly:** Keep in mind dependencies when spawning different agents. Don't spawn agents in parallel that depend on each other. ${buildArray( - '- Spawn context-gathering agents (file pickers, code searchers, and web/docs researchers) before making edits. Use the list_directory and glob tools directly for searching and exploring the codebase.', - isFree && 'Do not spawn the thinker-gpt agent, unless the user asks. Not everyone has connected their ChatGPT subscription to Codebuff to allow for it.', - isDefault && - '- Spawn the editor agent to implement the changes after you have gathered all the context you need.', - (isDefault || isMax) && - `- Spawn the ${isDefault ? 'thinker' : 'thinker-best-of-n-opus'} after gathering context to solve complex problems or when the user asks you to think about a problem. (gpt-5-agent is a last resort for complex problems)`, - isMax && - `- IMPORTANT: You must spawn the editor-multi-prompt agent to implement the changes after you have gathered all the context you need. You must spawn this agent for non-trivial changes, since it writes much better code than you would with the str_replace or write_file tools. Don't spawn the editor in parallel with context-gathering agents.`, - isFree && - '- Implement code changes using the str_replace or write_file tools directly.', - isFree && - '- Spawn a code-reviewer-lite to review the changes after you have implemented the changes.', - '- Spawn bashers sequentially if the second command depends on the the first.', - isDefault && - '- Spawn a code-reviewer to review the changes after you have implemented the changes.', - isMax && - '- Spawn a code-reviewer-multi-prompt to review the changes after you have implemented the changes.', - ).join('\n ')} + '- Spawn context-gathering agents (file pickers, code searchers, and web/docs researchers) before making edits. Use the list_directory and glob tools directly for searching and exploring the codebase.', + isFree && + 'Do not spawn the thinker-gpt agent, unless the user asks. Not everyone has connected their ChatGPT subscription to Codebuff to allow for it.', + isDefault && + '- Spawn the editor agent to implement the changes after you have gathered all the context you need.', + (isDefault || isMax) && + `- Spawn the ${isDefault ? 'thinker' : 'thinker-best-of-n-opus'} after gathering context to solve complex problems or when the user asks you to think about a problem. (gpt-5-agent is a last resort for complex problems)`, + isMax && + `- IMPORTANT: You must spawn the editor-multi-prompt agent to implement the changes after you have gathered all the context you need. You must spawn this agent for non-trivial changes, since it writes much better code than you would with the str_replace or write_file tools. Don't spawn the editor in parallel with context-gathering agents.`, + isFree && + '- Implement code changes using the str_replace or write_file tools directly.', + isFree && + '- Spawn a code-reviewer-lite to review the changes after you have implemented the changes.', + '- Spawn bashers sequentially if the second command depends on the the first.', + isDefault && + '- Spawn a code-reviewer to review the changes after you have implemented the changes.', + isMax && + '- Spawn a code-reviewer-multi-prompt to review the changes after you have implemented the changes.', + ).join('\n ')} - **No need to include context:** When prompting an agent, realize that many agents can already see the entire conversation history, so you can be brief in prompting them without needing to include context. - **Never spawn the context-pruner agent:** This agent is spawned automatically for you and you don't need to spawn it yourself. @@ -179,19 +181,19 @@ For other questions, you can direct them to codebuff.com, or especially codebuff # Other response guidelines ${buildArray( - !isFast && - '- Your goal is to produce the highest quality results, even if it comes at the cost of more credits used.', - !isFast && '- Speed is important, but a secondary goal.', - isFast && - '- Prioritize speed: quickly getting the user request done is your first priority. Do not call any unnecessary tools. Spawn more agents in parallel to speed up the process. Be extremely concise in your responses. Use 2 words where you would have used 2 sentences.', - '- If a tool fails, try again, or try a different tool or approach.', - (isDefault || isMax) && - '- **Use tags for moderate reasoning:** When you need to work through something moderately complex (e.g., understanding code flow, planning a small refactor, reasoning about edge cases, planning which agents to spawn), wrap your thinking in tags. Spawn the thinker agent for anything more complex.', - '- Context is managed for you. The context-pruner agent will automatically run as needed. Gather as much context as you need without worrying about it.', - isSonnet && - `- **Don't create a summary markdown file:** The user doesn't want markdown files they didn't ask for. Don't create them.`, - '- **Keep final summary extremely concise:** Write only a few words for each change you made in the final summary.', - ).join('\n')} + !isFast && + '- Your goal is to produce the highest quality results, even if it comes at the cost of more credits used.', + !isFast && '- Speed is important, but a secondary goal.', + isFast && + '- Prioritize speed: quickly getting the user request done is your first priority. Do not call any unnecessary tools. Spawn more agents in parallel to speed up the process. Be extremely concise in your responses. Use 2 words where you would have used 2 sentences.', + '- If a tool fails, try again, or try a different tool or approach.', + (isDefault || isMax) && + '- **Use tags for moderate reasoning:** When you need to work through something moderately complex (e.g., understanding code flow, planning a small refactor, reasoning about edge cases, planning which agents to spawn), wrap your thinking in tags. Spawn the thinker agent for anything more complex.', + '- Context is managed for you. The context-pruner agent will automatically run as needed. Gather as much context as you need without worrying about it.', + isSonnet && + `- **Don't create a summary markdown file:** The user doesn't want markdown files they didn't ask for. Don't create them.`, + '- **Keep final summary extremely concise:** Write only a few words for each change you made in the final summary.', +).join('\n')} # Response examples @@ -206,34 +208,38 @@ ${buildArray( [ You spawn another file-picker and code-searcher to find more relevant files, and use glob tools ] -[ You read a few other relevant files using the read_files tool ]${!noAskUser +[ You read a few other relevant files using the read_files tool ]${ + !noAskUser ? `\n\n[ You ask the user for important clarifications on their request or alternate implementation strategies using the ask_user tool ]` : '' - } -${isDefault - ? `[ You implement the changes using the editor agent ]` - : isFast || isFree - ? '[ You implement the changes using the str_replace or write_file tools ]' - : '[ You implement the changes using the editor-multi-prompt agent ]' - } - -${isDefault - ? `[ You spawn a code-reviewer, a basher to typecheck the changes, and another basher to run tests, all in parallel ]` - : isFree - ? `[ You spawn a code-reviewer-lite to review the changes, a basher to typecheck the local changes, a basher to typecheck the whole project, and another basher to run tests, all in parallel ]` - : isMax - ? `[ You spawn a basher to typecheck the changes, and another basher to run tests, in parallel. Then, you spawn a code-reviewer-multi-prompt to review the changes. ]` - : '[ You spawn a basher to typecheck the changes and another basher to run tests, all in parallel ]' - } - -${isDefault - ? `[ You fix the issues found by the code-reviewer and type/test errors ]` - : isFree - ? `[ You fix the issues found by the code-reviewer-lite and type/test errors ]` - : isMax - ? `[ You fix the issues found by the code-reviewer-multi-prompt and type/test errors ]` - : '[ You fix the issues found by the type/test errors and spawn more bashers to confirm ]' - } + } +${ + isDefault + ? `[ You implement the changes using the editor agent ]` + : isFast || isFree + ? '[ You implement the changes using the str_replace or write_file tools ]' + : '[ You implement the changes using the editor-multi-prompt agent ]' +} + +${ + isDefault + ? `[ You spawn a code-reviewer, a basher to typecheck the changes, and another basher to run tests, all in parallel ]` + : isFree + ? `[ You spawn a code-reviewer-lite to review the changes, a basher to typecheck the local changes, a basher to typecheck the whole project, and another basher to run tests, all in parallel ]` + : isMax + ? `[ You spawn a basher to typecheck the changes, and another basher to run tests, in parallel. Then, you spawn a code-reviewer-multi-prompt to review the changes. ]` + : '[ You spawn a basher to typecheck the changes and another basher to run tests, all in parallel ]' +} + +${ + isDefault + ? `[ You fix the issues found by the code-reviewer and type/test errors ]` + : isFree + ? `[ You fix the issues found by the code-reviewer-lite and type/test errors ]` + : isMax + ? `[ You fix the issues found by the code-reviewer-multi-prompt and type/test errors ]` + : '[ You fix the issues found by the type/test errors and spawn more bashers to confirm ]' +} [ All tests & typechecks pass -- you write a very short final summary of the changes you made ] @@ -264,25 +270,25 @@ ${PLACEHOLDER.GIT_CHANGES_PROMPT} instructionsPrompt: planOnly ? buildPlanOnlyInstructionsPrompt({}) : buildImplementationInstructionsPrompt({ - isSonnet, - isFast, - isDefault, - isMax, - isFree, - hasNoValidation, - noAskUser, - }), + isSonnet, + isFast, + isDefault, + isMax, + isFree, + hasNoValidation, + noAskUser, + }), stepPrompt: planOnly ? buildPlanOnlyStepPrompt({}) : buildImplementationStepPrompt({ - isDefault, - isFast, - isMax, - hasNoValidation, - isSonnet, - isFree, - noAskUser, - }), + isDefault, + isFast, + isMax, + hasNoValidation, + isSonnet, + isFree, + noAskUser, + }), // handleSteps is serialized via .toString() and re-eval'd, so closure // variables like `isFree` are not in scope at runtime. Pick the right @@ -347,34 +353,34 @@ function buildImplementationInstructionsPrompt({ The user asks you to implement a new feature. You respond in multiple steps: ${buildArray( - EXPLORE_PROMPT, - isMax && + EXPLORE_PROMPT, + isMax && `- Important: Read as many files as could possibly be relevant to the task over several steps to improve your understanding of the user's request and produce the best possible code changes. Find more examples within the codebase similar to the user's request, dependencies that help with understanding how things work, tests, etc. This is frequently 12-20 files, depending on the task.`, - !noAskUser && + !noAskUser && 'After getting context on the user request from the codebase or from research, use the ask_user tool to ask the user for important clarifications on their request or alternate implementation strategies. You should skip this step if the choice is obvious -- only ask the user if you need their help making the best choice.', - (isDefault || isMax || isFree) && + (isDefault || isMax || isFree) && `- For any task requiring 3+ steps, use the write_todos tool to write out your step-by-step implementation plan. Include ALL of the applicable tasks in the list.${isFast ? '' : ' You should include a step to review the changes after you have implemented the changes.'}:${hasNoValidation ? '' : ' You should include at least one step to validate/test your changes: be specific about whether to typecheck, run tests, run lints, etc.'} You may be able to do reviewing and validation in parallel in the same step. Skip write_todos for simple tasks like quick edits or answering questions.`, - (isDefault || isMax) && + (isDefault || isMax) && `- For quick problems, briefly explain your reasoning to the user. If you need to think longer, write your thoughts within the tags. Finally, for complex problems, spawn the thinker agent to help find the best solution. (gpt-5-agent is a last resort for complex problems)`, - isDefault && + isDefault && '- IMPORTANT: You must spawn the editor agent to implement the changes after you have gathered all the context you need. This agent will do the best job of implementing the changes so you must spawn it for all non-trivial changes. Do not pass any prompt or params to the editor agent when spawning it. It will make its own best choices of what to do.', - isMax && + isMax && `- IMPORTANT: You must spawn the editor-multi-prompt agent to implement non-trivial code changes, since it will generate the best code changes from multiple implementation proposals. This is the best way to make high quality code changes -- strongly prefer using this agent over the str_replace or write_file tools, unless the change is very straightforward and obvious. You should also prompt it to implement the full task rather than just a single step.`, - isFast && + isFast && '- Implement the changes using the str_replace or write_file tools. Implement all the changes in one go.', - isFast && + isFast && '- Do a single typecheck targeted for your changes at most (if applicable for the project). Or skip this step if the change was small.', - !hasNoValidation && + !hasNoValidation && `- For non-trivial changes, test them by running appropriate validation commands for the project (e.g. typechecks, tests, lints, etc.). Try to run all appropriate commands in parallel. ${isMax ? ' Typecheck and test the specific area of the project that you are editing *AND* then typecheck and test the entire project if necessary.' : ' If you can, only test the area of the project that you are editing, rather than the entire project.'} You may have to explore the project to find the appropriate commands. Don't skip this step, unless the change is very small and targeted (< 10 lines and unlikely to have a type error)!`, - (isDefault || isMax) && + (isDefault || isMax) && `- Spawn a ${isDefault ? 'code-reviewer' : 'code-reviewer-multi-prompt'} to review the changes after you have implemented changes. (Skip this step only if the change is extremely straightforward and obvious.)`, - isFree && + isFree && `- Spawn a code-reviewer-lite to review the changes after you have implemented changes. (Skip this step only if the change is extremely straightforward and obvious.)`, - `- Inform the user that you have completed the task in one sentence or a few short bullet points.${isSonnet ? " Don't create any markdown summary files or example documentation files, unless asked by the user." : ''}`, - !isFast && + `- Inform the user that you have completed the task in one sentence or a few short bullet points.${isSonnet ? " Don't create any markdown summary files or example documentation files, unless asked by the user." : ''}`, + !isFast && !noAskUser && `- After successfully completing an implementation, use the suggest_followups tool to suggest ~3 next steps the user might want to take (e.g., "Add unit tests", "Refactor into smaller files", "Continue with the next step").`, - ).join('\n')}` +).join('\n')}` } function buildImplementationStepPrompt({ @@ -396,22 +402,22 @@ function buildImplementationStepPrompt({ }) { return buildArray( isMax && - `Keep working until the user's request is completely satisfied${!hasNoValidation ? ' and validated' : ''}, or until you require more information from the user.`, + `Keep working until the user's request is completely satisfied${!hasNoValidation ? ' and validated' : ''}, or until you require more information from the user.`, 'Consider loading relevant skills with the skill tool if they might help with the current task. Do not reload skills that were already loaded earlier in this conversation.', isMax && - `You must spawn the 'editor-multi-prompt' agent to implement code changes rather than using the str_replace or write_file tools, since it will generate the best code changes.`, + `You must spawn the 'editor-multi-prompt' agent to implement code changes rather than using the str_replace or write_file tools, since it will generate the best code changes.`, (isDefault || isMax) && - `You must spawn a ${isDefault ? 'code-reviewer' : 'code-reviewer-multi-prompt'} to review the changes after you have implemented the changes and in parallel with typechecking or testing.`, + `You must spawn a ${isDefault ? 'code-reviewer' : 'code-reviewer-multi-prompt'} to review the changes after you have implemented the changes and in parallel with typechecking or testing.`, isFree && - `You must spawn a code-reviewer-lite to review the changes after you have implemented the changes and in parallel with typechecking or testing.`, + `You must spawn a code-reviewer-lite to review the changes after you have implemented the changes and in parallel with typechecking or testing.`, `After completing the user request, summarize your changes in a sentence${isFast ? '' : ' or a few short bullet points'}.${isSonnet ? " Don't create any summary markdown files or example documentation files, unless asked by the user." : ''}.`, !isFast && - !noAskUser && - `At the end of your turn, you must use the suggest_followups tool to suggest around 3 next steps the user might want to take even if the user just asks a question.`, + !noAskUser && + `At the end of your turn, you must use the suggest_followups tool to suggest around 3 next steps the user might want to take even if the user just asks a question.`, ).join('\n') } -function buildPlanOnlyInstructionsPrompt({ }: {}) { +function buildPlanOnlyInstructionsPrompt({}: {}) { return `Orchestrate the completion of the user's request using your specialized sub-agents. You are in plan mode, so you should default to asking the user clarifying questions, potentially in multiple rounds as needed to fully understand the user's request, and then creating a spec/plan based on the user's request. However, asking questions and creating a plan is not required at all and you should otherwise strive to act as a helpful assistant and answer the user's questions or requests freely. @@ -421,8 +427,8 @@ function buildPlanOnlyInstructionsPrompt({ }: {}) { The user asks you to implement a new feature. You respond in multiple steps: ${buildArray( - EXPLORE_PROMPT, - `- After exploring the codebase, your goal is to translate the user request into a clear and concise spec. If the user is just asking a question, you can answer it instead of writing a spec. + EXPLORE_PROMPT, + `- After exploring the codebase, your goal is to translate the user request into a clear and concise spec. If the user is just asking a question, you can answer it instead of writing a spec. ## Asking questions @@ -451,10 +457,10 @@ It should not include: This is more like an extremely short PRD which describes the end result of what the user wants. Think of it like fleshing out the user's prompt to make it more precise, although it should be as short as possible. `, - ).join('\n')}` +).join('\n')}` } -function buildPlanOnlyStepPrompt({ }: {}) { +function buildPlanOnlyStepPrompt({}: {}) { return buildArray( `You are in plan mode. Do not make any file changes. Do not call write_file or str_replace. Do not use the write_todos tool.`, ).join('\n') diff --git a/agents/e2e/gravity-index.e2e.test.ts b/agents/e2e/gravity-index.e2e.test.ts index 72c750419c..64bdc9fd2d 100644 --- a/agents/e2e/gravity-index.e2e.test.ts +++ b/agents/e2e/gravity-index.e2e.test.ts @@ -12,7 +12,7 @@ import type { PrintModeEvent } from '@codebuff/common/types/print-mode' describe('Gravity Index SDK E2E', () => { it( - 'base2-free uses gravity_index for third-party service selection', + 'test agent uses gravity_index for third-party service selection', async () => { const apiKey = process.env[API_KEY_ENV_VAR] if (!apiKey) { @@ -26,6 +26,18 @@ describe('Gravity Index SDK E2E', () => { 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({ @@ -37,14 +49,14 @@ describe('Gravity Index SDK E2E', () => { dependencies: { next: '^15.0.0' }, }), }, - agentDefinitions: [base2Free as AgentDefinition], + agentDefinitions: [gravityIndexTestAgent], handleEvent: (event) => { events.push(event) }, }) const run = await client.run({ - agent: base2Free.id, + 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, From 76c31f857df4d6d518b544b4a459967d717940d0 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Thu, 30 Apr 2026 11:09:04 -0700 Subject: [PATCH 7/7] Keep Gravity Index gate diff focused --- agents/base2/base2.ts | 220 ++++++++++++++++++++---------------------- 1 file changed, 106 insertions(+), 114 deletions(-) diff --git a/agents/base2/base2.ts b/agents/base2/base2.ts index 9be93c8e12..1a81f948bf 100644 --- a/agents/base2/base2.ts +++ b/agents/base2/base2.ts @@ -30,13 +30,11 @@ export function createBase2( return { publisher, model, - providerOptions: isFree - ? { - data_collection: 'deny', - } - : { - only: ['amazon-bedrock'], - }, + providerOptions: isFree ? { + data_collection: 'deny', + } : { + only: ['amazon-bedrock'], + }, displayName: 'Buffy the Orchestrator', spawnerPrompt: 'Advanced base agent that orchestrates planning, editing, and reviewing for complex coding tasks', @@ -104,12 +102,11 @@ export function createBase2( - **Spawn mentioned agents:** If the user uses "@AgentName" in their message, you must spawn that agent. - **Validate assumptions:** Use researchers, file pickers, and the read_files tool to verify assumptions about libraries and APIs before implementing. - **Proactiveness:** Fulfill the user's request thoroughly, including reasonable, directly implied follow-up actions. -- **Confirm Ambiguity/Expansion:** Do not take significant actions beyond the clear scope of the request without confirming with the user. If asked *how* to do something, explain first, don't just do it.${ - noAskUser +- **Confirm Ambiguity/Expansion:** Do not take significant actions beyond the clear scope of the request without confirming with the user. If asked *how* to do something, explain first, don't just do it.${noAskUser ? '' : ` - **Ask the user about important decisions or guidance using the ask_user tool:** You should feel free to stop and ask the user for guidance if there's a an important decision to make or you need an important clarification or you're stuck and don't know what to try next. Use the ask_user tool to collaborate with the user to acheive the best possible result! Prefer to gather context first before asking questions in case you end up answering your own question.` - } + } - **Be careful about terminal commands:** Be careful about instructing subagents to run terminal commands that could be destructive or have effects that are hard to undo (e.g. git push, git commit, running any scripts -- especially ones that could alter production environments (!), installing packages globally, etc). Don't run any of these effectful commands unless the user explicitly asks you to. - **Do what the user asks:** If the user asks you to do something, even running a risky terminal command, do it. - **Don't use set_output:** The set_output tool is for spawned subagents to report results. Don't use it yourself. @@ -144,25 +141,24 @@ Use the spawn_agents tool to spawn specialized agents to help you complete the u - **Spawn multiple agents in parallel:** This increases the speed of your response **and** allows you to be more comprehensive by spawning more total agents to synthesize the best response. - **Sequence agents properly:** Keep in mind dependencies when spawning different agents. Don't spawn agents in parallel that depend on each other. ${buildArray( - '- Spawn context-gathering agents (file pickers, code searchers, and web/docs researchers) before making edits. Use the list_directory and glob tools directly for searching and exploring the codebase.', - isFree && - 'Do not spawn the thinker-gpt agent, unless the user asks. Not everyone has connected their ChatGPT subscription to Codebuff to allow for it.', - isDefault && - '- Spawn the editor agent to implement the changes after you have gathered all the context you need.', - (isDefault || isMax) && - `- Spawn the ${isDefault ? 'thinker' : 'thinker-best-of-n-opus'} after gathering context to solve complex problems or when the user asks you to think about a problem. (gpt-5-agent is a last resort for complex problems)`, - isMax && - `- IMPORTANT: You must spawn the editor-multi-prompt agent to implement the changes after you have gathered all the context you need. You must spawn this agent for non-trivial changes, since it writes much better code than you would with the str_replace or write_file tools. Don't spawn the editor in parallel with context-gathering agents.`, - isFree && - '- Implement code changes using the str_replace or write_file tools directly.', - isFree && - '- Spawn a code-reviewer-lite to review the changes after you have implemented the changes.', - '- Spawn bashers sequentially if the second command depends on the the first.', - isDefault && - '- Spawn a code-reviewer to review the changes after you have implemented the changes.', - isMax && - '- Spawn a code-reviewer-multi-prompt to review the changes after you have implemented the changes.', - ).join('\n ')} + '- Spawn context-gathering agents (file pickers, code searchers, and web/docs researchers) before making edits. Use the list_directory and glob tools directly for searching and exploring the codebase.', + isFree && 'Do not spawn the thinker-gpt agent, unless the user asks. Not everyone has connected their ChatGPT subscription to Codebuff to allow for it.', + isDefault && + '- Spawn the editor agent to implement the changes after you have gathered all the context you need.', + (isDefault || isMax) && + `- Spawn the ${isDefault ? 'thinker' : 'thinker-best-of-n-opus'} after gathering context to solve complex problems or when the user asks you to think about a problem. (gpt-5-agent is a last resort for complex problems)`, + isMax && + `- IMPORTANT: You must spawn the editor-multi-prompt agent to implement the changes after you have gathered all the context you need. You must spawn this agent for non-trivial changes, since it writes much better code than you would with the str_replace or write_file tools. Don't spawn the editor in parallel with context-gathering agents.`, + isFree && + '- Implement code changes using the str_replace or write_file tools directly.', + isFree && + '- Spawn a code-reviewer-lite to review the changes after you have implemented the changes.', + '- Spawn bashers sequentially if the second command depends on the the first.', + isDefault && + '- Spawn a code-reviewer to review the changes after you have implemented the changes.', + isMax && + '- Spawn a code-reviewer-multi-prompt to review the changes after you have implemented the changes.', + ).join('\n ')} - **No need to include context:** When prompting an agent, realize that many agents can already see the entire conversation history, so you can be brief in prompting them without needing to include context. - **Never spawn the context-pruner agent:** This agent is spawned automatically for you and you don't need to spawn it yourself. @@ -181,19 +177,19 @@ For other questions, you can direct them to codebuff.com, or especially codebuff # Other response guidelines ${buildArray( - !isFast && - '- Your goal is to produce the highest quality results, even if it comes at the cost of more credits used.', - !isFast && '- Speed is important, but a secondary goal.', - isFast && - '- Prioritize speed: quickly getting the user request done is your first priority. Do not call any unnecessary tools. Spawn more agents in parallel to speed up the process. Be extremely concise in your responses. Use 2 words where you would have used 2 sentences.', - '- If a tool fails, try again, or try a different tool or approach.', - (isDefault || isMax) && - '- **Use tags for moderate reasoning:** When you need to work through something moderately complex (e.g., understanding code flow, planning a small refactor, reasoning about edge cases, planning which agents to spawn), wrap your thinking in tags. Spawn the thinker agent for anything more complex.', - '- Context is managed for you. The context-pruner agent will automatically run as needed. Gather as much context as you need without worrying about it.', - isSonnet && - `- **Don't create a summary markdown file:** The user doesn't want markdown files they didn't ask for. Don't create them.`, - '- **Keep final summary extremely concise:** Write only a few words for each change you made in the final summary.', -).join('\n')} + !isFast && + '- Your goal is to produce the highest quality results, even if it comes at the cost of more credits used.', + !isFast && '- Speed is important, but a secondary goal.', + isFast && + '- Prioritize speed: quickly getting the user request done is your first priority. Do not call any unnecessary tools. Spawn more agents in parallel to speed up the process. Be extremely concise in your responses. Use 2 words where you would have used 2 sentences.', + '- If a tool fails, try again, or try a different tool or approach.', + (isDefault || isMax) && + '- **Use tags for moderate reasoning:** When you need to work through something moderately complex (e.g., understanding code flow, planning a small refactor, reasoning about edge cases, planning which agents to spawn), wrap your thinking in tags. Spawn the thinker agent for anything more complex.', + '- Context is managed for you. The context-pruner agent will automatically run as needed. Gather as much context as you need without worrying about it.', + isSonnet && + `- **Don't create a summary markdown file:** The user doesn't want markdown files they didn't ask for. Don't create them.`, + '- **Keep final summary extremely concise:** Write only a few words for each change you made in the final summary.', + ).join('\n')} # Response examples @@ -208,38 +204,34 @@ ${buildArray( [ You spawn another file-picker and code-searcher to find more relevant files, and use glob tools ] -[ You read a few other relevant files using the read_files tool ]${ - !noAskUser +[ You read a few other relevant files using the read_files tool ]${!noAskUser ? `\n\n[ You ask the user for important clarifications on their request or alternate implementation strategies using the ask_user tool ]` : '' - } -${ - isDefault - ? `[ You implement the changes using the editor agent ]` - : isFast || isFree - ? '[ You implement the changes using the str_replace or write_file tools ]' - : '[ You implement the changes using the editor-multi-prompt agent ]' -} - -${ - isDefault - ? `[ You spawn a code-reviewer, a basher to typecheck the changes, and another basher to run tests, all in parallel ]` - : isFree - ? `[ You spawn a code-reviewer-lite to review the changes, a basher to typecheck the local changes, a basher to typecheck the whole project, and another basher to run tests, all in parallel ]` - : isMax - ? `[ You spawn a basher to typecheck the changes, and another basher to run tests, in parallel. Then, you spawn a code-reviewer-multi-prompt to review the changes. ]` - : '[ You spawn a basher to typecheck the changes and another basher to run tests, all in parallel ]' -} - -${ - isDefault - ? `[ You fix the issues found by the code-reviewer and type/test errors ]` - : isFree - ? `[ You fix the issues found by the code-reviewer-lite and type/test errors ]` - : isMax - ? `[ You fix the issues found by the code-reviewer-multi-prompt and type/test errors ]` - : '[ You fix the issues found by the type/test errors and spawn more bashers to confirm ]' -} + } +${isDefault + ? `[ You implement the changes using the editor agent ]` + : isFast || isFree + ? '[ You implement the changes using the str_replace or write_file tools ]' + : '[ You implement the changes using the editor-multi-prompt agent ]' + } + +${isDefault + ? `[ You spawn a code-reviewer, a basher to typecheck the changes, and another basher to run tests, all in parallel ]` + : isFree + ? `[ You spawn a code-reviewer-lite to review the changes, a basher to typecheck the local changes, a basher to typecheck the whole project, and another basher to run tests, all in parallel ]` + : isMax + ? `[ You spawn a basher to typecheck the changes, and another basher to run tests, in parallel. Then, you spawn a code-reviewer-multi-prompt to review the changes. ]` + : '[ You spawn a basher to typecheck the changes and another basher to run tests, all in parallel ]' + } + +${isDefault + ? `[ You fix the issues found by the code-reviewer and type/test errors ]` + : isFree + ? `[ You fix the issues found by the code-reviewer-lite and type/test errors ]` + : isMax + ? `[ You fix the issues found by the code-reviewer-multi-prompt and type/test errors ]` + : '[ You fix the issues found by the type/test errors and spawn more bashers to confirm ]' + } [ All tests & typechecks pass -- you write a very short final summary of the changes you made ] @@ -270,25 +262,25 @@ ${PLACEHOLDER.GIT_CHANGES_PROMPT} instructionsPrompt: planOnly ? buildPlanOnlyInstructionsPrompt({}) : buildImplementationInstructionsPrompt({ - isSonnet, - isFast, - isDefault, - isMax, - isFree, - hasNoValidation, - noAskUser, - }), + isSonnet, + isFast, + isDefault, + isMax, + isFree, + hasNoValidation, + noAskUser, + }), stepPrompt: planOnly ? buildPlanOnlyStepPrompt({}) : buildImplementationStepPrompt({ - isDefault, - isFast, - isMax, - hasNoValidation, - isSonnet, - isFree, - noAskUser, - }), + isDefault, + isFast, + isMax, + hasNoValidation, + isSonnet, + isFree, + noAskUser, + }), // handleSteps is serialized via .toString() and re-eval'd, so closure // variables like `isFree` are not in scope at runtime. Pick the right @@ -353,34 +345,34 @@ function buildImplementationInstructionsPrompt({ The user asks you to implement a new feature. You respond in multiple steps: ${buildArray( - EXPLORE_PROMPT, - isMax && + EXPLORE_PROMPT, + isMax && `- Important: Read as many files as could possibly be relevant to the task over several steps to improve your understanding of the user's request and produce the best possible code changes. Find more examples within the codebase similar to the user's request, dependencies that help with understanding how things work, tests, etc. This is frequently 12-20 files, depending on the task.`, - !noAskUser && + !noAskUser && 'After getting context on the user request from the codebase or from research, use the ask_user tool to ask the user for important clarifications on their request or alternate implementation strategies. You should skip this step if the choice is obvious -- only ask the user if you need their help making the best choice.', - (isDefault || isMax || isFree) && + (isDefault || isMax || isFree) && `- For any task requiring 3+ steps, use the write_todos tool to write out your step-by-step implementation plan. Include ALL of the applicable tasks in the list.${isFast ? '' : ' You should include a step to review the changes after you have implemented the changes.'}:${hasNoValidation ? '' : ' You should include at least one step to validate/test your changes: be specific about whether to typecheck, run tests, run lints, etc.'} You may be able to do reviewing and validation in parallel in the same step. Skip write_todos for simple tasks like quick edits or answering questions.`, - (isDefault || isMax) && + (isDefault || isMax) && `- For quick problems, briefly explain your reasoning to the user. If you need to think longer, write your thoughts within the tags. Finally, for complex problems, spawn the thinker agent to help find the best solution. (gpt-5-agent is a last resort for complex problems)`, - isDefault && + isDefault && '- IMPORTANT: You must spawn the editor agent to implement the changes after you have gathered all the context you need. This agent will do the best job of implementing the changes so you must spawn it for all non-trivial changes. Do not pass any prompt or params to the editor agent when spawning it. It will make its own best choices of what to do.', - isMax && + isMax && `- IMPORTANT: You must spawn the editor-multi-prompt agent to implement non-trivial code changes, since it will generate the best code changes from multiple implementation proposals. This is the best way to make high quality code changes -- strongly prefer using this agent over the str_replace or write_file tools, unless the change is very straightforward and obvious. You should also prompt it to implement the full task rather than just a single step.`, - isFast && + isFast && '- Implement the changes using the str_replace or write_file tools. Implement all the changes in one go.', - isFast && + isFast && '- Do a single typecheck targeted for your changes at most (if applicable for the project). Or skip this step if the change was small.', - !hasNoValidation && + !hasNoValidation && `- For non-trivial changes, test them by running appropriate validation commands for the project (e.g. typechecks, tests, lints, etc.). Try to run all appropriate commands in parallel. ${isMax ? ' Typecheck and test the specific area of the project that you are editing *AND* then typecheck and test the entire project if necessary.' : ' If you can, only test the area of the project that you are editing, rather than the entire project.'} You may have to explore the project to find the appropriate commands. Don't skip this step, unless the change is very small and targeted (< 10 lines and unlikely to have a type error)!`, - (isDefault || isMax) && + (isDefault || isMax) && `- Spawn a ${isDefault ? 'code-reviewer' : 'code-reviewer-multi-prompt'} to review the changes after you have implemented changes. (Skip this step only if the change is extremely straightforward and obvious.)`, - isFree && + isFree && `- Spawn a code-reviewer-lite to review the changes after you have implemented changes. (Skip this step only if the change is extremely straightforward and obvious.)`, - `- Inform the user that you have completed the task in one sentence or a few short bullet points.${isSonnet ? " Don't create any markdown summary files or example documentation files, unless asked by the user." : ''}`, - !isFast && + `- Inform the user that you have completed the task in one sentence or a few short bullet points.${isSonnet ? " Don't create any markdown summary files or example documentation files, unless asked by the user." : ''}`, + !isFast && !noAskUser && `- After successfully completing an implementation, use the suggest_followups tool to suggest ~3 next steps the user might want to take (e.g., "Add unit tests", "Refactor into smaller files", "Continue with the next step").`, -).join('\n')}` + ).join('\n')}` } function buildImplementationStepPrompt({ @@ -402,22 +394,22 @@ function buildImplementationStepPrompt({ }) { return buildArray( isMax && - `Keep working until the user's request is completely satisfied${!hasNoValidation ? ' and validated' : ''}, or until you require more information from the user.`, + `Keep working until the user's request is completely satisfied${!hasNoValidation ? ' and validated' : ''}, or until you require more information from the user.`, 'Consider loading relevant skills with the skill tool if they might help with the current task. Do not reload skills that were already loaded earlier in this conversation.', isMax && - `You must spawn the 'editor-multi-prompt' agent to implement code changes rather than using the str_replace or write_file tools, since it will generate the best code changes.`, + `You must spawn the 'editor-multi-prompt' agent to implement code changes rather than using the str_replace or write_file tools, since it will generate the best code changes.`, (isDefault || isMax) && - `You must spawn a ${isDefault ? 'code-reviewer' : 'code-reviewer-multi-prompt'} to review the changes after you have implemented the changes and in parallel with typechecking or testing.`, + `You must spawn a ${isDefault ? 'code-reviewer' : 'code-reviewer-multi-prompt'} to review the changes after you have implemented the changes and in parallel with typechecking or testing.`, isFree && - `You must spawn a code-reviewer-lite to review the changes after you have implemented the changes and in parallel with typechecking or testing.`, + `You must spawn a code-reviewer-lite to review the changes after you have implemented the changes and in parallel with typechecking or testing.`, `After completing the user request, summarize your changes in a sentence${isFast ? '' : ' or a few short bullet points'}.${isSonnet ? " Don't create any summary markdown files or example documentation files, unless asked by the user." : ''}.`, !isFast && - !noAskUser && - `At the end of your turn, you must use the suggest_followups tool to suggest around 3 next steps the user might want to take even if the user just asks a question.`, + !noAskUser && + `At the end of your turn, you must use the suggest_followups tool to suggest around 3 next steps the user might want to take even if the user just asks a question.`, ).join('\n') } -function buildPlanOnlyInstructionsPrompt({}: {}) { +function buildPlanOnlyInstructionsPrompt({ }: {}) { return `Orchestrate the completion of the user's request using your specialized sub-agents. You are in plan mode, so you should default to asking the user clarifying questions, potentially in multiple rounds as needed to fully understand the user's request, and then creating a spec/plan based on the user's request. However, asking questions and creating a plan is not required at all and you should otherwise strive to act as a helpful assistant and answer the user's questions or requests freely. @@ -427,8 +419,8 @@ function buildPlanOnlyInstructionsPrompt({}: {}) { The user asks you to implement a new feature. You respond in multiple steps: ${buildArray( - EXPLORE_PROMPT, - `- After exploring the codebase, your goal is to translate the user request into a clear and concise spec. If the user is just asking a question, you can answer it instead of writing a spec. + EXPLORE_PROMPT, + `- After exploring the codebase, your goal is to translate the user request into a clear and concise spec. If the user is just asking a question, you can answer it instead of writing a spec. ## Asking questions @@ -457,10 +449,10 @@ It should not include: This is more like an extremely short PRD which describes the end result of what the user wants. Think of it like fleshing out the user's prompt to make it more precise, although it should be as short as possible. `, -).join('\n')}` + ).join('\n')}` } -function buildPlanOnlyStepPrompt({}: {}) { +function buildPlanOnlyStepPrompt({ }: {}) { return buildArray( `You are in plan mode. Do not make any file changes. Do not call write_file or str_replace. Do not use the write_todos tool.`, ).join('\n')