|
| 1 | +import {ChildProcess, spawn} from 'child_process'; |
| 2 | +import {join, relative} from 'path'; |
| 3 | +import {existsSync} from 'fs'; |
| 4 | +import assert from 'assert'; |
| 5 | +import { |
| 6 | + LlmConstrainedOutputGenerateResponse, |
| 7 | + LlmGenerateFilesRequestOptions, |
| 8 | + LlmGenerateFilesResponse, |
| 9 | + LlmGenerateTextResponse, |
| 10 | +} from './llm-runner.js'; |
| 11 | +import {DirectorySnapshot} from './directory-snapshot.js'; |
| 12 | +import {LlmResponseFile} from '../shared-interfaces.js'; |
| 13 | +import {UserFacingError} from '../utils/errors.js'; |
| 14 | + |
| 15 | +/** Base class for a command-line-based runner. */ |
| 16 | +export abstract class BaseCliAgentRunner { |
| 17 | + abstract readonly displayName: string; |
| 18 | + protected abstract readonly binaryName: string; |
| 19 | + protected abstract readonly ignoredFilePatterns: string[]; |
| 20 | + protected abstract getCommandLineFlags(options: LlmGenerateFilesRequestOptions): string[]; |
| 21 | + protected abstract writeAgentFiles(options: LlmGenerateFilesRequestOptions): Promise<void>; |
| 22 | + |
| 23 | + private pendingTimeouts = new Set<ReturnType<typeof setTimeout>>(); |
| 24 | + private pendingProcesses = new Set<ChildProcess>(); |
| 25 | + private binaryPath: string | null = null; |
| 26 | + private commonIgnoredPatterns = ['**/node_modules/**', '**/dist/**', '**/.angular/**']; |
| 27 | + |
| 28 | + async generateFiles(options: LlmGenerateFilesRequestOptions): Promise<LlmGenerateFilesResponse> { |
| 29 | + const {context} = options; |
| 30 | + |
| 31 | + // TODO: Consider removing these assertions when we have better types. |
| 32 | + assert( |
| 33 | + context.buildCommand, |
| 34 | + 'Expected a `buildCommand` to be set in the LLM generate request context', |
| 35 | + ); |
| 36 | + assert( |
| 37 | + context.packageManager, |
| 38 | + 'Expected a `packageManager` to be set in the LLM generate request context', |
| 39 | + ); |
| 40 | + |
| 41 | + const ignoredPatterns = [...this.commonIgnoredPatterns, ...this.ignoredFilePatterns]; |
| 42 | + const initialSnapshot = await DirectorySnapshot.forDirectory( |
| 43 | + context.directory, |
| 44 | + ignoredPatterns, |
| 45 | + ); |
| 46 | + |
| 47 | + await this.writeAgentFiles(options); |
| 48 | + |
| 49 | + const reasoning = await this.runAgentProcess(options, 2, 10); |
| 50 | + const finalSnapshot = await DirectorySnapshot.forDirectory(context.directory, ignoredPatterns); |
| 51 | + |
| 52 | + const diff = finalSnapshot.getChangedOrAddedFiles(initialSnapshot); |
| 53 | + const files: LlmResponseFile[] = []; |
| 54 | + |
| 55 | + for (const [absolutePath, code] of diff) { |
| 56 | + files.push({ |
| 57 | + filePath: relative(context.directory, absolutePath), |
| 58 | + code, |
| 59 | + }); |
| 60 | + } |
| 61 | + |
| 62 | + return {files, reasoning, toolLogs: []}; |
| 63 | + } |
| 64 | + |
| 65 | + generateText(): Promise<LlmGenerateTextResponse> { |
| 66 | + // Technically we can make this work, but we don't need it at the time of writing. |
| 67 | + throw new UserFacingError(`Generating text with ${this.displayName} is not supported.`); |
| 68 | + } |
| 69 | + |
| 70 | + generateConstrained(): Promise<LlmConstrainedOutputGenerateResponse<any>> { |
| 71 | + // We can't support this, because there's no straightforward |
| 72 | + // way to tell the agent to follow a schema. |
| 73 | + throw new UserFacingError(`Constrained output with ${this.displayName} is not supported.`); |
| 74 | + } |
| 75 | + |
| 76 | + async dispose(): Promise<void> { |
| 77 | + for (const timeout of this.pendingTimeouts) { |
| 78 | + clearTimeout(timeout); |
| 79 | + } |
| 80 | + |
| 81 | + for (const childProcess of this.pendingProcesses) { |
| 82 | + childProcess.kill('SIGKILL'); |
| 83 | + } |
| 84 | + |
| 85 | + this.pendingTimeouts.clear(); |
| 86 | + this.pendingProcesses.clear(); |
| 87 | + } |
| 88 | + |
| 89 | + private resolveBinaryPath(binaryName: string): string { |
| 90 | + let dir = import.meta.dirname; |
| 91 | + let closestRoot: string | null = null; |
| 92 | + |
| 93 | + // Attempt to resolve the agent CLI binary by starting at the current file and going up until |
| 94 | + // we find the closest `node_modules`. Note that we can't rely on `import.meta.resolve` here, |
| 95 | + // because that'll point us to the agent bundle, but not its binary. In some package |
| 96 | + // managers (pnpm specifically) the `node_modules` in which the file is installed is different |
| 97 | + // from the one in which the binary is placed. |
| 98 | + while (dir.length > 1) { |
| 99 | + if (existsSync(join(dir, 'node_modules'))) { |
| 100 | + closestRoot = dir; |
| 101 | + break; |
| 102 | + } |
| 103 | + |
| 104 | + const parent = join(dir, '..'); |
| 105 | + |
| 106 | + if (parent === dir) { |
| 107 | + // We've reached the root, stop traversing. |
| 108 | + break; |
| 109 | + } else { |
| 110 | + dir = parent; |
| 111 | + } |
| 112 | + } |
| 113 | + |
| 114 | + const binaryPath = closestRoot ? join(closestRoot, `node_modules/.bin/${binaryName}`) : null; |
| 115 | + |
| 116 | + if (!binaryPath || !existsSync(binaryPath)) { |
| 117 | + throw new UserFacingError(`${this.displayName} is not installed inside the current project`); |
| 118 | + } |
| 119 | + |
| 120 | + return binaryPath; |
| 121 | + } |
| 122 | + |
| 123 | + private runAgentProcess( |
| 124 | + options: LlmGenerateFilesRequestOptions, |
| 125 | + inactivityTimeoutMins: number, |
| 126 | + totalRequestTimeoutMins: number, |
| 127 | + ): Promise<string> { |
| 128 | + return new Promise<string>(resolve => { |
| 129 | + let stdoutBuffer = ''; |
| 130 | + let stdErrBuffer = ''; |
| 131 | + let isDone = false; |
| 132 | + const msPerMin = 1000 * 60; |
| 133 | + const finalize = (finalMessage: string) => { |
| 134 | + if (isDone) { |
| 135 | + return; |
| 136 | + } |
| 137 | + |
| 138 | + isDone = true; |
| 139 | + |
| 140 | + if (inactivityTimeout) { |
| 141 | + clearTimeout(inactivityTimeout); |
| 142 | + this.pendingTimeouts.delete(inactivityTimeout); |
| 143 | + } |
| 144 | + |
| 145 | + clearTimeout(globalTimeout); |
| 146 | + childProcess.kill('SIGKILL'); |
| 147 | + this.pendingTimeouts.delete(globalTimeout); |
| 148 | + this.pendingProcesses.delete(childProcess); |
| 149 | + |
| 150 | + const separator = '\n--------------------------------------------------\n'; |
| 151 | + |
| 152 | + if (stdErrBuffer.length > 0) { |
| 153 | + stdoutBuffer += separator + 'Stderr output:\n' + stdErrBuffer; |
| 154 | + } |
| 155 | + |
| 156 | + stdoutBuffer += separator + finalMessage; |
| 157 | + resolve(stdoutBuffer); |
| 158 | + }; |
| 159 | + |
| 160 | + const noOutputCallback = () => { |
| 161 | + finalize( |
| 162 | + `There was no output from ${this.displayName} for ${inactivityTimeoutMins} minute(s). ` + |
| 163 | + `Stopping the process...`, |
| 164 | + ); |
| 165 | + }; |
| 166 | + |
| 167 | + // The agent can get into a state where it stops outputting code, but it also doesn't exit |
| 168 | + // the process. Stop if there hasn't been any output for a certain amount of time. |
| 169 | + let inactivityTimeout = setTimeout(noOutputCallback, inactivityTimeoutMins * msPerMin); |
| 170 | + this.pendingTimeouts.add(inactivityTimeout); |
| 171 | + |
| 172 | + // Also add a timeout for the entire codegen process. |
| 173 | + const globalTimeout = setTimeout(() => { |
| 174 | + finalize( |
| 175 | + `${this.displayName} didn't finish within ${totalRequestTimeoutMins} minute(s). ` + |
| 176 | + `Stopping the process...`, |
| 177 | + ); |
| 178 | + }, totalRequestTimeoutMins * msPerMin); |
| 179 | + |
| 180 | + this.binaryPath ??= this.resolveBinaryPath(this.binaryName); |
| 181 | + |
| 182 | + const childProcess = spawn(this.binaryPath, this.getCommandLineFlags(options), { |
| 183 | + cwd: options.context.directory, |
| 184 | + env: {...process.env}, |
| 185 | + }); |
| 186 | + |
| 187 | + childProcess.on('close', code => |
| 188 | + finalize( |
| 189 | + `${this.displayName} process has exited` + (code == null ? '.' : ` with ${code} code.`), |
| 190 | + ), |
| 191 | + ); |
| 192 | + childProcess.stdout.on('data', data => { |
| 193 | + if (inactivityTimeout) { |
| 194 | + this.pendingTimeouts.delete(inactivityTimeout); |
| 195 | + clearTimeout(inactivityTimeout); |
| 196 | + } |
| 197 | + |
| 198 | + stdoutBuffer += data.toString(); |
| 199 | + inactivityTimeout = setTimeout(noOutputCallback, inactivityTimeoutMins * msPerMin); |
| 200 | + this.pendingTimeouts.add(inactivityTimeout); |
| 201 | + }); |
| 202 | + childProcess.stderr.on('data', data => { |
| 203 | + stdErrBuffer += data.toString(); |
| 204 | + }); |
| 205 | + }); |
| 206 | + } |
| 207 | +} |
0 commit comments