Skip to content

Commit 96d6717

Browse files
feat(sdk): enable code_search in sdk (#254)
Co-authored-by: James Grugett <jahooma@gmail.com>
1 parent f757b55 commit 96d6717

File tree

7 files changed

+151
-9
lines changed

7 files changed

+151
-9
lines changed

bun.lock

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

common/src/tools/constants.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ export const publishedTools = [
6363
] as const
6464

6565
export type ToolName = (typeof toolNames)[number]
66+
export type PublishedToolName = (typeof publishedTools)[number]
6667

6768
export type $ToolParams<T extends ToolName = ToolName> = {
6869
toolName: T

common/src/tools/list.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import { updateSubgoalParams } from './params/tool/update-subgoal'
2323
import { webSearchParams } from './params/tool/web-search'
2424
import { writeFileParams } from './params/tool/write-file'
2525

26-
import type { ToolName, $ToolParams } from './constants'
26+
import type { ToolName, $ToolParams, PublishedToolName } from './constants'
2727
import type { ToolMessage } from '../types/messages/codebuff-message'
2828
import type {
2929
ToolCallPart,
@@ -125,3 +125,5 @@ export const clientToolNames = clientToolCallSchema.def.options.map(
125125
export type ClientToolCall<T extends ClientToolName = ClientToolName> = z.infer<
126126
typeof clientToolCallSchema
127127
> & { toolName: T } & Omit<ToolCallPart, 'type'>
128+
129+
export type PublishedClientToolName = ClientToolName & PublishedToolName

sdk/package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,8 @@
5858
"diff": "8.0.2",
5959
"zod": "^4.0.0",
6060
"@vscode/tree-sitter-wasm": "0.1.4",
61-
"web-tree-sitter": "0.25.6"
61+
"web-tree-sitter": "0.25.6",
62+
"@vscode/ripgrep": "1.15.14"
6263
},
6364
"devDependencies": {
6465
"@types/diff": "8.0.0",

sdk/src/client.ts

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import {
44
type RunState,
55
} from './run-state'
66
import { changeFile } from './tools/change-file'
7+
import { codeSearch } from './tools/code-search'
78
import { getFiles } from './tools/read-files'
89
import { runTerminalCommand } from './tools/run-terminal-command'
910
import { WebSocketHandler } from './websocket-client'
@@ -16,14 +17,18 @@ import { MAX_AGENT_STEPS_DEFAULT } from '../../common/src/constants/agents'
1617
import { toolNames } from '../../common/src/tools/constants'
1718
import {
1819
clientToolCallSchema,
20+
PublishedClientToolName,
1921
type ClientToolCall,
2022
type ClientToolName,
2123
type CodebuffToolOutput,
2224
} from '../../common/src/tools/list'
2325

2426
import type { CustomToolDefinition } from './custom-tool'
2527
import type { AgentDefinition } from '../../common/src/templates/initial-agents-dir/types/agent-definition'
26-
import type { ToolName } from '../../common/src/tools/constants'
28+
import type {
29+
PublishedToolName,
30+
ToolName,
31+
} from '../../common/src/tools/constants'
2732
import type {
2833
ToolResultOutput,
2934
ToolResultPart,
@@ -38,7 +43,7 @@ export type CodebuffClientOptions = {
3843
onError: (error: { message: string }) => void
3944
overrideTools?: Partial<
4045
{
41-
[K in ClientToolName]: (
46+
[K in ClientToolName & PublishedToolName]: (
4247
input: ClientToolCall<K>['input'],
4348
) => Promise<CodebuffToolOutput<K>>
4449
} & {
@@ -321,7 +326,7 @@ export class CodebuffClient {
321326
}
322327

323328
try {
324-
let override = this.overrideTools[toolName as ClientToolName]
329+
let override = this.overrideTools[toolName as PublishedClientToolName]
325330
if (!override && toolName === 'str_replace') {
326331
// Note: write_file and str_replace have the same implementation, so reuse their write_file override.
327332
override = this.overrideTools['write_file']
@@ -337,6 +342,21 @@ export class CodebuffClient {
337342
...input,
338343
cwd: input.cwd ?? this.cwd,
339344
} as Parameters<typeof runTerminalCommand>[0])
345+
} else if (toolName === 'code_search') {
346+
result = await codeSearch({
347+
projectPath: this.cwd,
348+
...input,
349+
} as Parameters<typeof codeSearch>[0])
350+
} else if (toolName === 'run_file_change_hooks') {
351+
// No-op: SDK doesn't run file change hooks
352+
result = [
353+
{
354+
type: 'json',
355+
value: {
356+
message: 'File change hooks are not supported in SDK mode',
357+
},
358+
},
359+
]
340360
} else {
341361
throw new Error(
342362
`Tool not implemented in SDK. Please provide an override or modify your agent to not use this tool: ${toolName}`,

sdk/src/tools/code-search.ts

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
import { spawn } from 'child_process'
2+
import * as path from 'path'
3+
4+
import { rgPath } from '@vscode/ripgrep'
5+
6+
import type { CodebuffToolOutput } from '../../../common/src/tools/list'
7+
8+
export function codeSearch({
9+
projectPath,
10+
pattern,
11+
flags,
12+
cwd,
13+
}: {
14+
projectPath: string
15+
pattern: string
16+
flags?: string
17+
cwd?: string
18+
}): Promise<CodebuffToolOutput<'code_search'>> {
19+
return new Promise((resolve) => {
20+
let stdout = ''
21+
let stderr = ''
22+
23+
const flagsArray = (flags || '').split(' ').filter(Boolean)
24+
let searchCwd = projectPath
25+
if (cwd) {
26+
const requestedPath = path.resolve(projectPath, cwd)
27+
// Ensure the search path is within the project directory
28+
if (!requestedPath.startsWith(projectPath)) {
29+
resolve([
30+
{
31+
type: 'json',
32+
value: {
33+
errorMessage: `Invalid cwd: Path '${cwd}' is outside the project directory.`,
34+
},
35+
},
36+
])
37+
return
38+
}
39+
searchCwd = requestedPath
40+
}
41+
42+
const args = [...flagsArray, pattern, '.']
43+
44+
const childProcess = spawn(rgPath, args, {
45+
cwd: searchCwd,
46+
stdio: ['ignore', 'pipe', 'pipe'],
47+
})
48+
49+
childProcess.stdout.on('data', (data) => {
50+
stdout += data.toString()
51+
})
52+
53+
childProcess.stderr.on('data', (data) => {
54+
stderr += data.toString()
55+
})
56+
57+
childProcess.on('close', (code) => {
58+
// Truncate output to prevent memory issues
59+
const maxLength = 10000
60+
const truncatedStdout =
61+
stdout.length > maxLength
62+
? stdout.substring(0, maxLength) + '\n\n[Output truncated]'
63+
: stdout
64+
65+
const maxErrorLength = 1000
66+
const truncatedStderr =
67+
stderr.length > maxErrorLength
68+
? stderr.substring(0, maxErrorLength) + '\n\n[Error output truncated]'
69+
: stderr
70+
71+
const result = {
72+
stdout: truncatedStdout,
73+
...(truncatedStderr && { stderr: truncatedStderr }),
74+
...(code !== null && { exitCode: code }),
75+
message: 'Code search completed',
76+
}
77+
78+
resolve([
79+
{
80+
type: 'json',
81+
value: result,
82+
},
83+
])
84+
})
85+
86+
childProcess.on('error', (error) => {
87+
resolve([
88+
{
89+
type: 'json',
90+
value: {
91+
errorMessage: `Failed to execute ripgrep: ${error.message}. Make sure ripgrep is installed and available in PATH.`,
92+
},
93+
},
94+
])
95+
})
96+
})
97+
}
Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import type { CodebuffToolOutput } from '../../../common/src/tools/list'
2+
3+
export function runFileChangeHooks({
4+
files,
5+
}: {
6+
files: string[]
7+
}): Promise<CodebuffToolOutput<'run_file_change_hooks'>> {
8+
// In the SDK, we don't have access to codebuff.json configuration
9+
// or the hook running infrastructure, so this is a no-op
10+
11+
return Promise.resolve([
12+
{
13+
type: 'json',
14+
value: [
15+
{
16+
errorMessage: 'No file change hooks were triggered for the specified files. File change hooks are not supported in the SDK environment.',
17+
},
18+
],
19+
},
20+
])
21+
}

0 commit comments

Comments
 (0)