Something went wrong. Try again.
source dump of claude code forked from oppi.li/claude-code
Something went wrong. Try again.
11 kB · 367 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368/** * Shared command prefix extraction using Haiku LLM * * This module provides a factory for creating command prefix extractors * that can be used by different shell tools. The core logic * (Haiku query, response validation) is shared, while tool-specific * aspects (examples, pre-checks) are configurable. */
import chalk from 'chalk'import type { QuerySource } from '../../constants/querySource.js'import { getFeatureValue_CACHED_MAY_BE_STALE } from '../../services/analytics/growthbook.js'import { type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, logEvent,} from '../../services/analytics/index.js'import { queryHaiku } from '../../services/api/claude.js'import { startsWithApiErrorPrefix } from '../../services/api/errors.js'import { memoizeWithLRU } from '../memoize.js'import { jsonStringify } from '../slowOperations.js'import { asSystemPrompt } from '../systemPromptType.js'
/** * Shell executables that must never be accepted as bare prefixes. * Allowing e.g. "bash:*" would let any command through, defeating * the permission system. Includes Unix shells and Windows equivalents. */const DANGEROUS_SHELL_PREFIXES = new Set([ 'sh', 'bash', 'zsh', 'fish', 'csh', 'tcsh', 'ksh', 'dash', 'cmd', 'cmd.exe', 'powershell', 'powershell.exe', 'pwsh', 'pwsh.exe', 'bash.exe',])
/** * Result of command prefix extraction */export type CommandPrefixResult = { /** The detected command prefix, or null if no prefix could be determined */ commandPrefix: string | null}
/** * Result including subcommand prefixes for compound commands */export type CommandSubcommandPrefixResult = CommandPrefixResult & { subcommandPrefixes: Map<string, CommandPrefixResult>}
/** * Configuration for creating a command prefix extractor */export type PrefixExtractorConfig = { /** Tool name for logging and warning messages */ toolName: string
/** The policy spec containing examples for Haiku */ policySpec: string /** Analytics event name for logging */ eventName: string
/** Query source identifier for the API call */ querySource: QuerySource
/** Optional pre-check function that can short-circuit the Haiku call */ preCheck?: (command: string) => CommandPrefixResult | null}
/** * Creates a memoized command prefix extractor function. * * Uses two-layer memoization: the outer memoized function creates the promise * and attaches a .catch handler that evicts the cache entry on rejection. * This prevents aborted or failed Haiku calls from poisoning future lookups. * * Bounded to 200 entries via LRU to prevent unbounded growth in heavy sessions. * * @param config - Configuration for the extractor * @returns A memoized async function that extracts command prefixes */export function createCommandPrefixExtractor(config: PrefixExtractorConfig) { const { toolName, policySpec, eventName, querySource, preCheck } = config
const memoized = memoizeWithLRU( ( command: string, abortSignal: AbortSignal, isNonInteractiveSession: boolean, ): Promise<CommandPrefixResult | null> => { const promise = getCommandPrefixImpl( command, abortSignal, isNonInteractiveSession, toolName, policySpec, eventName, querySource, preCheck, ) // Evict on rejection so aborted calls don't poison future turns. // Identity guard: after LRU eviction, a newer promise may occupy // this key; a stale rejection must not delete it. promise.catch(() => { if (memoized.cache.get(command) === promise) { memoized.cache.delete(command) } }) return promise }, command => command, // memoize by command only 200, )
return memoized}
/** * Creates a memoized function to get prefixes for compound commands with subcommands. * * Uses the same two-layer memoization pattern as createCommandPrefixExtractor: * a .catch handler evicts the cache entry on rejection to prevent poisoning. * * @param getPrefix - The single-command prefix extractor (from createCommandPrefixExtractor) * @param splitCommand - Function to split a compound command into subcommands * @returns A memoized async function that extracts prefixes for the main command and all subcommands */export function createSubcommandPrefixExtractor( getPrefix: ReturnType<typeof createCommandPrefixExtractor>, splitCommand: (command: string) => string[] | Promise<string[]>,) { const memoized = memoizeWithLRU( ( command: string, abortSignal: AbortSignal, isNonInteractiveSession: boolean, ): Promise<CommandSubcommandPrefixResult | null> => { const promise = getCommandSubcommandPrefixImpl( command, abortSignal, isNonInteractiveSession, getPrefix, splitCommand, ) // Evict on rejection so aborted calls don't poison future turns. // Identity guard: after LRU eviction, a newer promise may occupy // this key; a stale rejection must not delete it. promise.catch(() => { if (memoized.cache.get(command) === promise) { memoized.cache.delete(command) } }) return promise }, command => command, // memoize by command only 200, )
return memoized}
async function getCommandPrefixImpl( command: string, abortSignal: AbortSignal, isNonInteractiveSession: boolean, toolName: string, policySpec: string, eventName: string, querySource: QuerySource, preCheck?: (command: string) => CommandPrefixResult | null,): Promise<CommandPrefixResult | null> { if (process.env.NODE_ENV === 'test') { return null }
// Run pre-check if provided (e.g., isHelpCommand for Bash) if (preCheck) { const preCheckResult = preCheck(command) if (preCheckResult !== null) { return preCheckResult } }
let preflightCheckTimeoutId: NodeJS.Timeout | undefined const startTime = Date.now() let result: CommandPrefixResult | null = null
try { // Log a warning if the pre-flight check takes too long preflightCheckTimeoutId = setTimeout( (tn, nonInteractive) => { const message = `[${tn}Tool] Pre-flight check is taking longer than expected. Run with ANTHROPIC_LOG=debug to check for failed or slow API requests.` if (nonInteractive) { process.stderr.write(jsonStringify({ level: 'warn', message }) + '\n') } else { // biome-ignore lint/suspicious/noConsole: intentional warning console.warn(chalk.yellow(`⚠️ ${message}`)) } }, 10000, // 10 seconds toolName, isNonInteractiveSession, )
const useSystemPromptPolicySpec = getFeatureValue_CACHED_MAY_BE_STALE( 'tengu_cork_m4q', false, )
const response = await queryHaiku({ systemPrompt: asSystemPrompt( useSystemPromptPolicySpec ? [ `Your task is to process ${toolName} commands that an AI coding agent wants to run.\n\n${policySpec}`, ] : [ `Your task is to process ${toolName} commands that an AI coding agent wants to run.\n\nThis policy spec defines how to determine the prefix of a ${toolName} command:`, ], ), userPrompt: useSystemPromptPolicySpec ? `Command: ${command}` : `${policySpec}\n\nCommand: ${command}`, signal: abortSignal, options: { enablePromptCaching: useSystemPromptPolicySpec, querySource, agents: [], isNonInteractiveSession, hasAppendSystemPrompt: false, mcpTools: [], }, })
// Clear the timeout since the query completed clearTimeout(preflightCheckTimeoutId) const durationMs = Date.now() - startTime
const prefix = typeof response.message.content === 'string' ? response.message.content : Array.isArray(response.message.content) ? (response.message.content.find(_ => _.type === 'text')?.text ?? 'none') : 'none'
if (startsWithApiErrorPrefix(prefix)) { logEvent(eventName, { success: false, error: 'API error' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, durationMs, }) result = null } else if (prefix === 'command_injection_detected') { // Haiku detected something suspicious - treat as no prefix available logEvent(eventName, { success: false, error: 'command_injection_detected' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, durationMs, }) result = { commandPrefix: null, } } else if ( prefix === 'git' || DANGEROUS_SHELL_PREFIXES.has(prefix.toLowerCase()) ) { // Never accept bare `git` or shell executables as a prefix logEvent(eventName, { success: false, error: 'dangerous_shell_prefix' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, durationMs, }) result = { commandPrefix: null, } } else if (prefix === 'none') { // No prefix detected logEvent(eventName, { success: false, error: 'prefix "none"' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, durationMs, }) result = { commandPrefix: null, } } else { // Validate that the prefix is actually a prefix of the command
if (!command.startsWith(prefix)) { // Prefix isn't actually a prefix of the command logEvent(eventName, { success: false, error: 'command did not start with prefix' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, durationMs, }) result = { commandPrefix: null, } } else { logEvent(eventName, { success: true, durationMs, }) result = { commandPrefix: prefix, } } }
return result } catch (error) { clearTimeout(preflightCheckTimeoutId) throw error }}
async function getCommandSubcommandPrefixImpl( command: string, abortSignal: AbortSignal, isNonInteractiveSession: boolean, getPrefix: ReturnType<typeof createCommandPrefixExtractor>, splitCommandFn: (command: string) => string[] | Promise<string[]>,): Promise<CommandSubcommandPrefixResult | null> { const subcommands = await splitCommandFn(command)
const [fullCommandPrefix, ...subcommandPrefixesResults] = await Promise.all([ getPrefix(command, abortSignal, isNonInteractiveSession), ...subcommands.map(async subcommand => ({ subcommand, prefix: await getPrefix(subcommand, abortSignal, isNonInteractiveSession), })), ])
if (!fullCommandPrefix) { return null }
const subcommandPrefixes = subcommandPrefixesResults.reduce( (acc, { subcommand, prefix }) => { if (prefix) { acc.set(subcommand, prefix) } return acc }, new Map<string, CommandPrefixResult>(), )
return { ...fullCommandPrefix, subcommandPrefixes, }}