@sourcegraph/the-orb-is-awake
Version:
TypeScript SDK for Amp CLI - Build custom AI agents with Amp's capabilities
325 lines • 11.9 kB
JavaScript
/**
* Amp TypeScript SDK
*
* This SDK provides a TypeScript interface to the Amp CLI, allowing you to
* run Amp programmatically in Node.js applications. It wraps the Amp CLI
* with the --stream-json flag to provide structured output.
*/
import { spawn } from 'node:child_process';
import fs from 'node:fs';
import { createRequire } from 'node:module';
import path from 'node:path';
import { createInterface } from 'node:readline';
import { AmpOptionsSchema, UserInputMessageSchema, } from './types.js';
// Platform detection
const isWindows = process.platform === 'win32';
// ============================================================================
// Main Execute Function
// ============================================================================
/**
* Execute a command with Amp CLI and return an async iterator of messages
*
* @param options Execute configuration including prompt and options
* @returns Async iterator of stream messages from Amp CLI
*/
export async function* execute(options) {
const { prompt, options: ampOptions = {}, signal } = options;
// Validate options
const validatedOptions = AmpOptionsSchema.parse(ampOptions);
// Handle different prompt types
const isStreamingInput = typeof prompt !== 'string' && Symbol.asyncIterator in prompt;
// Build CLI arguments
const args = buildCliArgs(validatedOptions);
// Add flag for streaming input if needed
if (isStreamingInput) {
args.push('--stream-json-input');
}
// Log the full CLI command for debugging (only in debug mode)
if (process.env.AMP_DEBUG || validatedOptions.logLevel === 'debug') {
const ampCommand = findAmpCommand();
console.debug(`Executing Amp CLI: ${ampCommand.command} ${[...ampCommand.args, ...args].join(' ')}`);
}
// Build environment variables
const env = buildEnvironmentVariables(validatedOptions);
// Spawn Amp CLI process
const ampProcess = spawnAmpCli(args, {
env,
cwd: validatedOptions.cwd,
signal,
});
try {
// Handle input via stdin
if (isStreamingInput) {
handleStreamingInput(ampProcess, prompt);
}
else if (typeof prompt === 'string' && prompt) {
handleStringInput(ampProcess, prompt);
}
// Process output stream
yield* processOutputStream(ampProcess.stdout);
// Wait for process to complete and check exit code
const { exitCode, stderr, signal: processSignal } = await waitForProcess(ampProcess, signal);
// Handle exit code properly - null indicates signal termination
if (exitCode === null) {
// Process was killed by signal
if (signal?.aborted) {
// Expected termination due to AbortController
throw new Error('Amp CLI process was aborted');
}
else if (processSignal) {
// Killed by system signal
throw new Error(`Amp CLI process was killed by signal ${processSignal}`);
}
else {
// Unexpected null exit code
throw new Error('Amp CLI process terminated unexpectedly');
}
}
else if (exitCode !== 0) {
// Normal exit with error code
const errorDetails = stderr ? `: ${stderr}` : '';
throw new Error(`Amp CLI process exited with code ${exitCode}${errorDetails}`);
}
}
catch (error) {
// Clean up process on error
if (!ampProcess.killed) {
ampProcess.kill(isWindows ? 'SIGKILL' : 'SIGTERM');
}
throw error;
}
}
// ============================================================================
// Environment Variable Management
// ============================================================================
/**
* Build environment variables for the Amp CLI process
* Combines system environment variables with user-provided options
*/
function buildEnvironmentVariables(options) {
const env = {
...process.env,
};
// Set AMP_TOOLBOX if toolbox option is provided
if (options.toolbox) {
env.AMP_TOOLBOX = options.toolbox;
}
// User-provided env variables override system ones
if (options.env) {
Object.assign(env, options.env);
}
return env;
}
// ============================================================================
// CLI Process Management
// ============================================================================
function spawnAmpCli(args, options) {
const { env, cwd, signal } = options;
// Find amp executable
const ampCommand = findAmpCommand();
const childProcess = spawn(ampCommand.command, [...ampCommand.args, ...args], {
cwd: cwd || process.cwd(),
env: env || process.env,
stdio: ['pipe', 'pipe', 'pipe'],
signal,
});
// Add error handler immediately to prevent unhandled error events
childProcess.on('error', () => {
// Error will be handled by waitForProcess
});
return childProcess;
}
function findAmpCommand() {
// Try to find the local CLI (npm already resolved the correct version)
try {
const require = createRequire(import.meta.url);
const pkgJsonPath = require.resolve('@sourcegraph/amp/package.json');
const pkgJsonRaw = fs.readFileSync(pkgJsonPath, 'utf8');
const pkgJson = JSON.parse(pkgJsonRaw);
if (pkgJson.bin?.amp) {
const binPath = path.join(path.dirname(pkgJsonPath), pkgJson.bin.amp);
return {
command: 'node',
args: [binPath],
};
}
throw new Error('Local @sourcegraph/amp package found but no bin entry for Amp CLI');
}
catch (error) {
if (error instanceof Error && error.message.includes('Local @sourcegraph/amp')) {
throw error;
}
throw new Error('Could not find local @sourcegraph/amp package. Make sure it is installed.');
}
}
function buildCliArgs(options) {
const args = [];
// Handle continue option first as it changes the base command structure
if (typeof options.continue === 'string') {
args.push('threads', 'continue', options.continue);
}
else if (options.continue === true) {
args.push('threads', 'continue');
}
// Add execute and stream-json flags
args.push('--execute', '--stream-json');
// Optional flags
if (options.dangerouslyAllowAll) {
args.push('--dangerously-allow-all');
}
if (options.visibility) {
args.push('--visibility', options.visibility);
}
if (options.settingsFile) {
args.push('--settings-file', options.settingsFile);
}
if (options.logLevel) {
args.push('--log-level', options.logLevel);
}
if (options.logFile) {
args.push('--log-file', options.logFile);
}
if (options.mcpConfig) {
const mcpConfigValue = typeof options.mcpConfig === 'string'
? options.mcpConfig
: JSON.stringify(options.mcpConfig);
args.push('--mcp-config', mcpConfigValue);
}
return args;
}
// ============================================================================
// Stream Processing
// ============================================================================
async function* processOutputStream(stdout) {
if (!stdout) {
throw new Error('No stdout stream available from Amp CLI process');
}
const readline = createInterface({
input: stdout,
crlfDelay: Number.POSITIVE_INFINITY,
});
try {
for await (const line of readline) {
if (!line.trim())
continue;
try {
const message = JSON.parse(line);
yield message;
}
catch (parseError) {
throw new Error(`Failed to parse JSON response, raw line: ${line}`);
}
}
}
finally {
readline.close();
}
}
async function handleStreamingInput(process, prompt, signal) {
if (!process.stdin) {
throw new Error('No stdin stream available for Amp CLI process');
}
// Check if already aborted before starting
signal?.throwIfAborted();
try {
for await (const message of prompt) {
// Check for abort before processing each message
signal?.throwIfAborted();
// Validate message format
const validatedMessage = UserInputMessageSchema.parse(message);
// Send the complete JSON message
const jsonMessage = JSON.stringify(validatedMessage) + '\n';
if (!process.stdin.write(jsonMessage)) {
// Wait for drain if buffer is full, but allow cancellation during wait
await new Promise((resolve, reject) => {
const onAbort = () => reject(signal?.reason || new Error('Aborted'));
signal?.addEventListener('abort', onAbort, { once: true });
process.stdin.once('drain', () => {
// Clean up abort listener when drain completes
signal?.removeEventListener('abort', onAbort);
resolve();
});
});
}
}
}
finally {
// Close stdin to signal end of input
process.stdin.end();
}
}
function handleStringInput(process, prompt) {
if (!process.stdin) {
throw new Error('No stdin stream available for Amp CLI process');
}
// Write the prompt to stdin and close it
process.stdin.write(prompt + '\n');
process.stdin.end();
}
function waitForProcess(process, signal) {
return new Promise((resolve, reject) => {
let resolved = false;
let stderrData = '';
const cleanup = () => {
resolved = true;
};
// Capture stderr output
if (process.stderr) {
process.stderr.on('data', (data) => {
stderrData += data.toString();
});
}
process.on('exit', (code, processSignal) => {
if (!resolved) {
cleanup();
resolve({ exitCode: code, stderr: stderrData, signal: processSignal });
}
});
process.on('error', (error) => {
if (!resolved) {
cleanup();
reject(error);
}
});
// Handle abort signal
if (signal?.aborted) {
if (!resolved) {
cleanup();
reject(new Error('Operation was aborted'));
}
}
else if (signal) {
const abortHandler = () => {
if (!resolved) {
cleanup();
reject(new Error('Operation was aborted'));
}
};
signal.addEventListener('abort', abortHandler);
// Clean up listener
process.on('exit', () => {
signal.removeEventListener('abort', abortHandler);
});
}
});
}
// ============================================================================
// Utility Functions
// ============================================================================
/**
* Helper function to create streaming input messages
*/
export function createUserMessage(text) {
return {
type: 'user',
message: {
role: 'user',
content: [{ type: 'text', text }],
},
};
}
// ============================================================================
// Re-export types for convenience
// ============================================================================
export * from './types.js';
//# sourceMappingURL=index.js.map