diff --git a/docs/advanced-setup.md b/docs/advanced-setup.md index 6d23c8622d..3ca8ffe607 100644 --- a/docs/advanced-setup.md +++ b/docs/advanced-setup.md @@ -938,3 +938,29 @@ project's `disabledMcpServers` key to restore inheritance. Builtin servers that are disabled by default keep their separate project-level `enabledMcpServers` allowlist. Existing running sessions are not forcibly disconnected by a toggle in another project. + +## Needle local planner + +Needle is an optional local pre-query planner and is disabled by default. Install Python 3 and Git first, then install the upstream package from source. + +Linux and macOS: + +```bash +git clone https://github.com/cactus-compute/needle.git +cd needle +python3 -m pip install --user -e '.[jax]' +``` + +Windows PowerShell: + +```powershell +git clone https://github.com/cactus-compute/needle.git +Set-Location needle +py -3 -m pip install --user -e ".[jax]" +``` + +OpenClaude uses `python3` by default on Linux/macOS and `py -3` on Windows. Set `OPENCLAUDE_NEEDLE_PYTHON` to an interpreter path when your setup differs. + +Enable it in an OpenClaude session with `/needle enable`. OpenClaude verifies that the selected interpreter can import `needle`; if not, enablement is refused with an installation message. Use `/needle mode worker` for AgentTool workers or `/needle mode main` for the main conversation; `/needle status` reports the current settings, and `/needle disable` turns it off. In either mode, Needle proposes at most one tool call. OpenClaude executes it using normal permissions and tool handling, then continues the same conversation with the configured provider. The provider remains responsible for the final response, summaries, interpreting WebFetch results, redirects, and any subsequent tool calls. Needle does not summarize the conversation or execute tools itself. + +The `needle-worker` model route selects Needle for AgentTool workers when enabled; main-conversation use is controlled by the `main` mode. The TypeScript bridge exchanges JSONL requests with Python. If Needle fails after enablement, OpenClaude skips the local prelude and continues with the configured provider. Set `OPENCLAUDE_NEEDLE_TIMEOUT_MS` to change the bounded planner timeout. diff --git a/src/commands.ts b/src/commands.ts index ea0e53113d..e5e1806e18 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -26,6 +26,7 @@ import ctx_viz from './commands/ctx_viz/index.js' import doctor from './commands/doctor/index.js' import onboardGithub from './commands/onboard-github/index.js' import knowledge from './commands/knowledge/index.js' +import needle from './commands/needle/index.js' import memory from './commands/memory/index.js' import repomap from './commands/repomap/index.js' import help from './commands/help/index.js' @@ -326,6 +327,7 @@ const COMMANDS = memoize((): Command[] => [ init, keybindings, knowledge, + needle, lsp, installGitHubApp, installSlackApp, diff --git a/src/commands/needle/index.ts b/src/commands/needle/index.ts new file mode 100644 index 0000000000..51044958cd --- /dev/null +++ b/src/commands/needle/index.ts @@ -0,0 +1,12 @@ +import type { Command } from '../../commands.js' + +const needle: Command = { + type: 'local', + name: 'needle', + description: 'Enable or disable the local Needle tool-calling planner', + supportsNonInteractive: true, + argumentHint: 'enable | disable | status | mode ', + load: () => import('./needle.js'), +} + +export default needle diff --git a/src/commands/needle/needle.test.ts b/src/commands/needle/needle.test.ts new file mode 100644 index 0000000000..9cf434676b --- /dev/null +++ b/src/commands/needle/needle.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, it, beforeEach, afterEach, spyOn } from 'bun:test' +import { mkdtempSync, rmSync } from 'fs' +import { tmpdir } from 'os' +import { join } from 'path' +import { call as needleCall } from './needle.js' +import * as needleClient from '../../services/needle/needleClient.js' +import { getGlobalConfig, saveGlobalConfig } from '../../utils/config.js' +import { setClaudeConfigHomeDirForTesting } from '../../utils/envUtils.js' +import type { LocalJSXCommandContext } from '../../types/command.js' +import { + acquireSharedMutationLock, + releaseSharedMutationLock, +} from '../../test/sharedMutationLock.js' + +describe('needle command', () => { + let configDir: string + beforeEach(async () => { + await acquireSharedMutationLock('commands/needle.test.ts') + configDir = mkdtempSync(join(tmpdir(), 'openclaude-needle-command-')) + process.env.CLAUDE_CONFIG_DIR = configDir + setClaudeConfigHomeDirForTesting(configDir) + }) + afterEach(() => { + saveGlobalConfig(current => ({ ...current, needleEnabled: false, needleMode: 'worker' })) + setClaudeConfigHomeDirForTesting(undefined) + delete process.env.CLAUDE_CONFIG_DIR + rmSync(configDir, { recursive: true, force: true }) + releaseSharedMutationLock() + }) + + it('refuses enablement when Needle is unavailable', async () => { + const previous = process.env.OPENCLAUDE_NEEDLE_PYTHON + process.env.OPENCLAUDE_NEEDLE_PYTHON = process.platform === 'win32' ? 'missing-needle-python.exe' : '/bin/false' + const context = {} as unknown as LocalJSXCommandContext + try { + const result = await needleCall('enable', context) + expect(getGlobalConfig().needleEnabled).toBe(false) + expect(result.type).toBe('text') + if (result.type === 'text') expect(result.value).toContain('was not enabled') + } finally { + if (previous === undefined) delete process.env.OPENCLAUDE_NEEDLE_PYTHON + else process.env.OPENCLAUDE_NEEDLE_PYTHON = previous + } + }) + + it('toggles the worker and persists its mode', async () => { + const availability = spyOn(needleClient, 'checkNeedleAvailability').mockResolvedValue({ available: true }) + const context = {} as unknown as LocalJSXCommandContext + try { + await needleCall('enable', context) + expect(getGlobalConfig().needleEnabled).toBe(true) + expect(getGlobalConfig().needleMode).toBe('worker') + + await needleCall('mode main', context) + expect(getGlobalConfig().needleMode).toBe('main') + + const status = await needleCall('status', context) + expect(status.type).toBe('text') + if (status.type === 'text') expect(status.value).toContain('main') + } finally { + availability.mockRestore() + } + }) +}) diff --git a/src/commands/needle/needle.ts b/src/commands/needle/needle.ts new file mode 100644 index 0000000000..04dd233828 --- /dev/null +++ b/src/commands/needle/needle.ts @@ -0,0 +1,62 @@ +import type { LocalCommandCall } from '../../types/command.js' +import { getGlobalConfig, saveGlobalConfig } from '../../utils/config.js' +import chalk from 'chalk' +import { checkNeedleAvailability } from '../../services/needle/needleClient.js' + +export const call: LocalCommandCall = async args => { + const parts = String(args ?? '').trim().toLowerCase().split(/\s+/).filter(Boolean) + const command = parts[0] ?? 'status' + const config = getGlobalConfig() + const enabled = config.needleEnabled === true + const mode = config.needleMode ?? 'worker' + + if (command === 'status') { + return { + type: 'text', + value: `Needle: ${enabled ? chalk.green('enabled') : chalk.red('disabled')} (${mode} mode)`, + } + } + + if (command === 'enable' || command === 'disable') { + const nextEnabled = command === 'enable' + if (nextEnabled) { + const availability = await checkNeedleAvailability() + if (!availability.available) { + return { + type: 'text', + value: `Needle was not enabled because its Python package is unavailable. Install it from source, then retry /needle enable. (${availability.error ?? 'import failed'})`, + } + } + } + saveGlobalConfig(current => ({ + ...current, + needleEnabled: nextEnabled, + needleMode: current.needleMode ?? 'worker', + })) + return { + type: 'text', + value: `Needle ${nextEnabled ? chalk.green('enabled') : chalk.red('disabled')} (${mode} mode).`, + } + } + + if (command === 'mode') { + const nextMode = parts[1] + if (nextMode !== 'worker' && nextMode !== 'main') { + return { type: 'text', value: 'Usage: /needle mode ' } + } + saveGlobalConfig(current => ({ ...current, needleMode: nextMode })) + return { + type: 'text', + value: `Needle mode set to ${chalk.cyan(nextMode)}. ${ + nextMode === 'main' + ? 'The main conversation will use Needle when enabled.' + : 'AgentTool workers will use Needle when enabled.' + }`, + } + } + + return { + type: 'text', + value: 'Usage: /needle ', + } +} diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index be7a3a4306..e5e8a03815 100644 --- a/src/screens/REPL.tsx +++ b/src/screens/REPL.tsx @@ -148,6 +148,7 @@ import { useMailboxBridge } from '../hooks/useMailboxBridge.js'; import { queryCheckpoint, logQueryProfileReport, clearQueryProfile } from '../utils/queryProfiler.js'; import type { Message as MessageType, UserMessage, ProgressMessage, HookResultMessage, PartialCompactDirection } from '../types/message.js'; import { query, type QueryTurnBudget } from '../query.js'; +import { closeNeedleClient, runNeedleWorkerTurn } from '../services/needle/runNeedleWorkerTurn.js'; import type { Terminal as QueryTerminal } from '../query/transitions.js'; import type { AutoCompactTrackingState } from '../services/compact/autoCompact.js'; import { mergeClients, useMergedClients } from '../hooks/useMergedClients.js'; @@ -681,6 +682,7 @@ export function REPL({ // without requiring a REPL remount. CLI prop still wins over env/config. const isRemoteSession = !!remoteSessionConfig; const foregroundTurnBudgetRef = useRef(null); + const needleClientOwnerRef = useRef({}); const backgroundHandoffStartedRef = useRef(false); const [backgroundHandoffPreparing, setBackgroundHandoffPreparing] = useState(false); @@ -1017,6 +1019,10 @@ export function REPL({ // REPL bridge to abort the active query when a remote interrupt arrives. const abortControllerRef = useRef(null); abortControllerRef.current = abortController; + useEffect(() => () => { + abortControllerRef.current?.abort(); + closeNeedleClient(needleClientOwnerRef.current); + }, []); // Ref for the bridge result callback — set after useReplBridge initializes, // read in the onQuery finally block to notify mobile clients that a turn ended. @@ -3292,7 +3298,8 @@ export function REPL({ resetTurnToolDuration(); resetTurnClassifierDuration(); let expectedAutoCompactTracking = queryAutoCompactTracking; - const queryGenerator = query({ + const needleConfig = getGlobalConfig(); + const queryParams = { messages: messagesIncludingNewMessages, requestOnlyMessages, onModelRequestStart: interruptionCorrectionQueryId @@ -3301,8 +3308,8 @@ export function REPL({ shouldQuery, isInterruptionCorrectionEligible: true, queryId: interruptionCorrectionQueryId, - }) - onModelRequestStart?.() + }); + onModelRequestStart?.(); } : undefined, systemPrompt, @@ -3318,13 +3325,34 @@ export function REPL({ if (setAutoCompactTrackingForSessionIfUnchanged(querySessionId, expectedAutoCompactTracking, tracking)) { expectedAutoCompactTracking = tracking; } - } - }); + }, + }; + const currentUserMessage = newMessages.findLast(selectableUserMessagesFilter); + const currentUserContent = currentUserMessage?.message.content; + const needlePrompt = currentUserContent ? getContentText(currentUserContent) ?? '' : ''; + const hasImageContent = + Array.isArray(currentUserContent) && currentUserContent.some(block => block.type === 'image'); + const queryGenerator = + needleConfig.needleEnabled === true && + needleConfig.needleMode === 'main' && + needlePrompt && + !hasImageContent + ? runNeedleWorkerTurn({ + prompt: needlePrompt, + queryParams, + tools: freshTools, + toolUseContext, + clientOwner: needleClientOwnerRef.current, + canUseTool, + signal: abortController.signal, + }) + : query(queryParams); const queryTerminal = await driveQueryEvents( queryGenerator, reason => queryGuard.registerActivity(reason, queryGeneration), onQueryEvent, ); + if (isBuddyEnabled()) { void fireCompanionObserver(messagesRef.current, reaction => setAppState(prev => prev.companionReaction === reaction ? prev : { ...prev, diff --git a/src/services/needle/needleClient.test.ts b/src/services/needle/needleClient.test.ts new file mode 100644 index 0000000000..541b467bbe --- /dev/null +++ b/src/services/needle/needleClient.test.ts @@ -0,0 +1,108 @@ +import { describe, expect, test } from 'bun:test' +import { chmodSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { + buildNeedleToolPayload, + DEFAULT_NEEDLE_TIMEOUT_MS, + getNeedlePythonCommand, + getNeedleTimeoutMs, + NEEDLE_WORKER_MODEL, +} from './needleClient.js' + +describe('Needle client configuration', () => { + test('uses the dedicated worker model identifier', () => { + expect(NEEDLE_WORKER_MODEL).toBe('needle-worker') + }) + + test('uses platform-appropriate Python launchers', () => { + const previous = process.env.OPENCLAUDE_NEEDLE_PYTHON + delete process.env.OPENCLAUDE_NEEDLE_PYTHON + try { + expect(getNeedlePythonCommand('linux')).toEqual({ command: 'python3', args: [] }) + expect(getNeedlePythonCommand('darwin')).toEqual({ command: 'python3', args: [] }) + expect(getNeedlePythonCommand('win32')).toEqual({ command: 'py', args: ['-3'] }) + } finally { + if (previous === undefined) delete process.env.OPENCLAUDE_NEEDLE_PYTHON + else process.env.OPENCLAUDE_NEEDLE_PYTHON = previous + } + }) + + test('converts Anthropic-style tool schemas to Needle payloads', () => { + expect(buildNeedleToolPayload([ + { name: 'read', description: 'Read a file', input_schema: { type: 'object', properties: {} } }, + ])).toEqual([ + { name: 'read', description: 'Read a file', parameters: { type: 'object', properties: {} } }, + ]) + }) + + test.skipIf(process.platform === 'win32')('restarts the worker after a child exit before serving later requests', async () => { + const previous = process.env.OPENCLAUDE_NEEDLE_PYTHON + process.env.OPENCLAUDE_NEEDLE_PYTHON = '/bin/false' + const client = new (await import('./needleClient.js')).NeedleClient(100) + try { + await expect(client.complete({ prompt: 'first', tools: [] })).rejects.toThrow('exited with code 1') + await expect(client.complete({ prompt: 'second', tools: [] })).rejects.toThrow('exited with code 1') + } finally { + client.close() + if (previous === undefined) delete process.env.OPENCLAUDE_NEEDLE_PYTHON + else process.env.OPENCLAUDE_NEEDLE_PYTHON = previous + } + }) + + test.skipIf(process.platform === 'win32')('restarts the worker after a zero child exit before serving later requests', async () => { + const previous = process.env.OPENCLAUDE_NEEDLE_PYTHON + process.env.OPENCLAUDE_NEEDLE_PYTHON = '/bin/true' + const client = new (await import('./needleClient.js')).NeedleClient(100) + try { + await expect(client.complete({ prompt: 'first', tools: [] })).rejects.toThrow('exited with code 0') + await expect(client.complete({ prompt: 'second', tools: [] })).rejects.toThrow('exited with code 0') + } finally { + client.close() + if (previous === undefined) delete process.env.OPENCLAUDE_NEEDLE_PYTHON + else process.env.OPENCLAUDE_NEEDLE_PYTHON = previous + } + }) + + test.skipIf(process.platform === 'win32')('rejects malformed matching responses instead of leaving requests pending', async () => { + const directory = mkdtempSync(join(tmpdir(), 'needle-runtime-')) + const runtime = join(directory, 'runtime.py') + const command = join(directory, 'python') + writeFileSync(runtime, `import json, sys\nfor line in sys.stdin:\n request = json.loads(line)\n print(json.dumps({"id": request["id"]}), flush=True)\n`) + writeFileSync(command, `#!/bin/sh\nexec python3 /${runtime.replace(/^\//, '')}\n`) + chmodSync(command, 0o755) + const previous = process.env.OPENCLAUDE_NEEDLE_PYTHON + process.env.OPENCLAUDE_NEEDLE_PYTHON = command + const client = new (await import('./needleClient.js')).NeedleClient(1_000) + try { + await expect(client.complete({ prompt: 'malformed', tools: [] })).rejects.toThrow('Needle worker returned no result') + } finally { + client.close() + if (previous === undefined) delete process.env.OPENCLAUDE_NEEDLE_PYTHON + else process.env.OPENCLAUDE_NEEDLE_PYTHON = previous + rmSync(directory, { recursive: true, force: true }) + } + }) + + test.skipIf(process.platform === 'win32')('restarts the worker after a timeout before serving later requests', async () => { + const directory = mkdtempSync(join(tmpdir(), 'needle-runtime-')) + const runtime = join(directory, 'runtime.py') + const command = join(directory, 'python') + writeFileSync(runtime, `import json, sys, time\nfor line in sys.stdin:\n request = json.loads(line)\n if request["prompt"] == "delay": time.sleep(0.2)\n print(json.dumps({"id": request["id"], "result": {"function_calls": [{"name": "read", "arguments": {"prompt": request["prompt"]}}]}}), flush=True)\n`) + writeFileSync(command, `#!/bin/sh\nexec python3 /${runtime.replace(/^\//, '')}\n`) + chmodSync(command, 0o755) + const previous = process.env.OPENCLAUDE_NEEDLE_PYTHON + process.env.OPENCLAUDE_NEEDLE_PYTHON = command + const client = new (await import('./needleClient.js')).NeedleClient(30) + try { + await expect(client.complete({ prompt: 'delay', tools: [] })).rejects.toThrow('timed out') + const response = await client.complete({ prompt: 'fresh', tools: [] }) + expect(response.functionCalls[0]?.arguments.prompt).toBe('fresh') + } finally { + client.close() + if (previous === undefined) delete process.env.OPENCLAUDE_NEEDLE_PYTHON + else process.env.OPENCLAUDE_NEEDLE_PYTHON = previous + rmSync(directory, { recursive: true, force: true }) + } + }) +}) diff --git a/src/services/needle/needleClient.ts b/src/services/needle/needleClient.ts new file mode 100644 index 0000000000..0380bd1540 --- /dev/null +++ b/src/services/needle/needleClient.ts @@ -0,0 +1,278 @@ +import { spawn, type ChildProcessWithoutNullStreams } from 'node:child_process' +import { createInterface, type Interface } from 'node:readline' +import { randomUUID } from 'node:crypto' + +export const NEEDLE_WORKER_MODEL = 'needle-worker' +export const DEFAULT_NEEDLE_TIMEOUT_MS = 5_000 + +const NEEDLE_RUNTIME = String.raw` +import json, sys, needle + +def main(): + agent = None + agent_tools = None + for line in sys.stdin: + request = None + try: + request = json.loads(line) + if agent is None or agent_tools != request["tools"]: + agent_tools = request["tools"] + agent = needle.Needle(tools=agent_tools, generation=3) + result = agent.complete(request["prompt"], max_new_tokens=request.get("max_new_tokens", 256)) + print(json.dumps({"id": request["id"], "result": result}, separators=(",", ":")), flush=True) + except Exception as error: + request_id = request.get("id") if isinstance(request, dict) else None + print(json.dumps({"id": request_id, "error": str(error)}, separators=(",", ":")), flush=True) + +main() +` + +export type NeedleTool = { + name: string + description?: string + parameters: Record +} + +export type NeedleRequest = { + prompt: string + tools: NeedleTool[] + maxNewTokens?: number + signal?: AbortSignal +} + +export type NeedleFunctionCall = { + name: string + arguments: Record +} + +export type NeedleResponse = { + functionCalls: NeedleFunctionCall[] + text?: string + decodeTps?: number + prefillTps?: number +} + +type PendingRequest = { + resolve: (response: NeedleResponse) => void + reject: (error: Error) => void + timer: ReturnType + abortCleanup?: () => void +} + +type RuntimeResponse = { + id: string + error?: string + result?: { + function_calls?: Array<{ name?: unknown; arguments?: unknown }> + text?: unknown + decode_tps?: unknown + prefill_tps?: unknown + } +} + +export function getNeedlePythonCommand( + platform: NodeJS.Platform = process.platform, +): { command: string; args: string[] } { + const configured = process.env.OPENCLAUDE_NEEDLE_PYTHON?.trim() + if (configured) return { command: configured, args: [] } + return platform === 'win32' + ? { command: 'py', args: ['-3'] } + : { command: 'python3', args: [] } +} + +export async function checkNeedleAvailability(): Promise<{ + available: boolean + error?: string +}> { + const { command, args } = getNeedlePythonCommand() + return new Promise(resolve => { + const child = spawn(command, [...args, '-c', 'import needle'], { + stdio: ['ignore', 'ignore', 'pipe'], + windowsHide: true, + }) + let stderr = '' + const timer = setTimeout(() => { + child.kill() + resolve({ available: false, error: 'Needle availability check timed out' }) + }, 5_000) + child.stderr.on('data', chunk => { + stderr += String(chunk) + }) + child.once('error', error => { + clearTimeout(timer) + resolve({ available: false, error: error.message }) + }) + child.once('exit', code => { + clearTimeout(timer) + resolve(code === 0 + ? { available: true } + : { available: false, error: stderr.trim() || `Python exited with code ${code ?? 'unknown'}` }) + }) + }) +} + +export function getNeedleTimeoutMs(): number { + const value = Number(process.env.OPENCLAUDE_NEEDLE_TIMEOUT_MS) + return Number.isFinite(value) && value > 0 + ? Math.min(Math.floor(value), 120_000) + : DEFAULT_NEEDLE_TIMEOUT_MS +} + +export function buildNeedleToolPayload( + tools: Array<{ + name: string + description?: string + input_schema?: Record + parameters?: Record + }>, +): NeedleTool[] { + return tools.map(tool => ({ + name: tool.name, + ...(tool.description ? { description: tool.description } : {}), + parameters: tool.parameters ?? tool.input_schema ?? { type: 'object', properties: {} }, + })) +} + +function normalizeResponse(response: RuntimeResponse): NeedleResponse { + if (response.error) throw new Error(`Needle worker error: ${response.error}`) + if (!response.result || typeof response.result !== 'object') { + throw new Error('Needle worker returned no result') + } + const functionCalls = (response.result.function_calls ?? []).map(call => { + if (typeof call.name !== 'string' || !call.name) { + throw new Error('Needle worker returned a function call without a name') + } + if (!call.arguments || typeof call.arguments !== 'object' || Array.isArray(call.arguments)) { + throw new Error(`Needle worker returned invalid arguments for ${call.name}`) + } + return { name: call.name, arguments: call.arguments as Record } + }) + return { + functionCalls, + ...(typeof response.result.text === 'string' ? { text: response.result.text } : {}), + ...(typeof response.result.decode_tps === 'number' ? { decodeTps: response.result.decode_tps } : {}), + ...(typeof response.result.prefill_tps === 'number' ? { prefillTps: response.result.prefill_tps } : {}), + } +} + +export class NeedleClient { + private process?: ChildProcessWithoutNullStreams + private lines?: Interface + private readonly pending = new Map() + private readonly timeoutMs: number + private toolsKey = '' + + constructor(timeoutMs = getNeedleTimeoutMs()) { + this.timeoutMs = timeoutMs + } + + async complete(request: NeedleRequest): Promise { + this.ensureProcess(request.tools) + const id = randomUUID() + const payload = { + id, + prompt: request.prompt, + tools: request.tools, + max_new_tokens: request.maxNewTokens ?? 256, + } + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + this.pending.delete(id) + pending.abortCleanup?.() + this.restart(new Error(`Needle worker timed out after ${this.timeoutMs}ms`)) + reject(new Error(`Needle worker timed out after ${this.timeoutMs}ms`)) + }, this.timeoutMs) + const pending: PendingRequest = { resolve, reject, timer } + if (request.signal) { + const abort = () => { + this.pending.delete(id) + clearTimeout(timer) + this.restart(new Error('Needle worker request aborted')) + reject(new Error('Needle worker request aborted')) + } + if (request.signal.aborted) return abort() + request.signal.addEventListener('abort', abort, { once: true }) + pending.abortCleanup = () => request.signal?.removeEventListener('abort', abort) + } + this.pending.set(id, pending) + this.process!.stdin.write(`${JSON.stringify(payload)}\n`, error => { + if (!error) return + this.restart(new Error(`Unable to write to Needle worker: ${error.message}`)) + }) + }) + } + + close(): void { + this.lines?.close() + this.process?.kill() + this.rejectAll(new Error('Needle worker closed')) + this.lines = undefined + this.process = undefined + this.toolsKey = '' + } + + private restart(error: Error): void { + const child = this.process + this.lines?.close() + this.process = undefined + this.lines = undefined + this.toolsKey = '' + child?.kill() + this.rejectAll(error) + } + + private ensureProcess(tools: NeedleTool[]): void { + const key = JSON.stringify(tools) + if (this.process && this.toolsKey === key) return + if (this.process) this.close() + const { command, args } = getNeedlePythonCommand() + const child = spawn(command, [...args, '-u', '-c', NEEDLE_RUNTIME], { + stdio: ['pipe', 'pipe', 'pipe'], + windowsHide: true, + }) + this.process = child + this.toolsKey = key + this.lines = createInterface({ input: child.stdout }) + child.stderr.resume() + this.lines.on('line', line => { + try { + const response = JSON.parse(line) as RuntimeResponse + const pending = response.id ? this.pending.get(response.id) : undefined + if (!pending) return + if (response.error) { + this.pending.delete(response.id) + clearTimeout(pending.timer) + pending.abortCleanup?.() + pending.reject(new Error(`Needle worker error: ${response.error}`)) + } else { + const normalized = normalizeResponse(response) + this.pending.delete(response.id) + clearTimeout(pending.timer) + pending.abortCleanup?.() + pending.resolve(normalized) + } + } catch (error) { + this.rejectAll(error instanceof Error ? error : new Error(String(error))) + } + }) + const fail = (error: Error) => { + if (this.process === child) this.restart(error) + } + child.once('error', fail) + child.stdin.on('error', fail) + child.once('exit', code => { + if (this.process === child) { + fail(new Error(`Needle worker exited with code ${code ?? 'unknown'}`)) + } + }) + } + + private rejectAll(error: Error): void { + for (const [id, pending] of this.pending) { + clearTimeout(pending.timer) + pending.abortCleanup?.() + pending.reject(error) + this.pending.delete(id) + } + } +} diff --git a/src/services/needle/runNeedleWorkerTurn.test.ts b/src/services/needle/runNeedleWorkerTurn.test.ts new file mode 100644 index 0000000000..dc53f8a094 --- /dev/null +++ b/src/services/needle/runNeedleWorkerTurn.test.ts @@ -0,0 +1,211 @@ +import { describe, expect, test } from 'bun:test' +import { z } from 'zod/v4' +import type { QueryParams } from '../../query.js' +import type { Message } from '../../types/message.js' +import { createUserMessage } from '../../utils/messages.js' +import { asSystemPrompt } from '../../utils/systemPromptType.js' +import { createToolFixture } from '../../test/toolFixtures.js' +import { runNeedleWorkerTurn } from './runNeedleWorkerTurn.js' +import type { NeedleFunctionCall, NeedleResponse } from './needleClient.js' + +function makeContext(tools: QueryParams['toolUseContext']['options']['tools']) { + let inProgress = new Set() + const context = { + abortController: new AbortController(), + getAppState: () => ({ toolPermissionContext: { mode: 'default' } }), + options: { + tools, + commands: [], + debug: false, + thinkingConfig: { type: 'disabled' }, + verbose: false, + mcpClients: [], + mcpResources: {}, + isNonInteractiveSession: false, + agentDefinitions: { activeAgents: [], allowedAgentTypes: undefined }, + mainLoopModel: 'provider-model', + }, + setInProgressToolUseIDs: (update: (previous: Set) => Set) => { + inProgress = update(inProgress) + }, + } as unknown as QueryParams['toolUseContext'] + return { context, getInProgress: () => inProgress } +} + +function makeParams(toolUseContext: QueryParams['toolUseContext']): QueryParams { + return { + messages: [createUserMessage({ content: 'original prompt' })], + systemPrompt: asSystemPrompt([]), + userContext: {}, + systemContext: {}, + canUseTool: async () => ({ behavior: 'allow' }), + toolUseContext, + querySource: 'agent:builtin:general-purpose', + } +} + +function fakeClient(response: NeedleResponse) { + return { + complete: async () => response, + close: () => {}, + } +} + +async function collect(generator: AsyncGenerator) { + const yielded: Array = [] + while (true) { + const next = await generator.next() + if (next.done) return yielded + yielded.push(next.value) + } +} + +describe('runNeedleWorkerTurn', () => { + test.each<{ functionCalls: NeedleFunctionCall[] }>([ + { functionCalls: [] }, + { functionCalls: [ + { name: 'Read', arguments: {} }, + { name: 'Write', arguments: {} }, + ] }, + ])('delegates unchanged when Needle returns zero or multiple calls', async ({ functionCalls }) => { + const { context } = makeContext([]) + const params = makeParams(context) + let continuedWith: QueryParams | undefined + const continueQuery = async function* (next: QueryParams) { + continuedWith = next + return { reason: 'completed' as const } + } + + await collect(runNeedleWorkerTurn({ + prompt: 'plan', + queryParams: params, + tools: [], + toolUseContext: context, + canUseTool: params.canUseTool, + client: fakeClient({ functionCalls }), + continueQuery, + })) + + expect(continuedWith?.messages).toBe(params.messages) + expect(continuedWith?.toolUseContext).toBe(context) + }) + + test('falls back to the provider when the Needle runtime fails', async () => { + const { context } = makeContext([]) + const params = makeParams(context) + let continued = false + const continueQuery = async function* () { + continued = true + return { reason: 'completed' as const } + } + const client = { + complete: async () => { throw new Error('missing runtime') }, + close: () => {}, + } + + await collect(runNeedleWorkerTurn({ + prompt: 'plan', + queryParams: params, + tools: [], + toolUseContext: context, + canUseTool: params.canUseTool, + client, + continueQuery, + })) + + expect(continued).toBe(true) + }) + + test('does not execute or continue after aborting after the planner prelude', async () => { + const tool = createToolFixture(z.object({ path: z.string() }), { + name: 'Read', + async call() { + throw new Error('tool should not execute') + }, + }) + const { context, getInProgress } = makeContext([tool]) + const params = makeParams(context) + const controller = new AbortController() + let continued = false + const continueQuery = async function* () { + continued = true + return { reason: 'completed' as const } + } + const generator = runNeedleWorkerTurn({ + prompt: 'plan', + queryParams: params, + tools: [tool], + toolUseContext: context, + canUseTool: params.canUseTool, + signal: controller.signal, + client: fakeClient({ functionCalls: [{ name: 'Read', arguments: { path: '/tmp/a' } }] }), + continueQuery, + }) + const prelude = (await generator.next()).value + expect(prelude).toMatchObject({ type: 'assistant' }) + controller.abort() + const result = await generator.next() + expect(result.done).toBe(false) + expect(result.value).toMatchObject({ type: 'user' }) + const cancellation = result.value as Message + const cancellationContent = cancellation.type === 'user' ? cancellation.message.content : [] + expect(cancellationContent).toContainEqual(expect.objectContaining({ + type: 'tool_result', + content: 'Cancelled before tool execution.', + is_error: true, + })) + const cancellationResult = Array.isArray(cancellationContent) + ? cancellationContent.find(block => block.type === 'tool_result') + : undefined + const preludeContent = prelude && 'type' in prelude && prelude.type === 'assistant' ? prelude.message.content : [] + const preludeToolUse = Array.isArray(preludeContent) ? preludeContent.find(block => block.type === 'tool_use') : undefined + expect(cancellationResult?.tool_use_id).toBe(preludeToolUse?.id) + const terminal = await generator.next() + expect(terminal.done).toBe(true) + expect(terminal.value).toEqual({ reason: 'aborted_streaming' }) + expect(continued).toBe(false) + expect(getInProgress().size).toBe(0) + }) + + test('yields a paired prelude and continues the provider with normalized results', async () => { + const tool = createToolFixture(z.object({ path: z.string() }), { + name: 'Read', + async call({ path }) { + return { data: `contents:${path}` } + }, + mapToolResultToToolResultBlockParam(content, toolUseID) { + return { type: 'tool_result', tool_use_id: toolUseID, content: String(content) } + }, + }) + const { context, getInProgress } = makeContext([tool]) + const params = makeParams(context) + let providerMessages: Message[] = [] + const continueQuery = async function* (next: QueryParams) { + providerMessages = next.messages + return { reason: 'completed' as const } + } + + const yielded = await collect(runNeedleWorkerTurn({ + prompt: 'plan', + queryParams: params, + tools: [tool], + toolUseContext: context, + canUseTool: params.canUseTool, + client: fakeClient({ functionCalls: [{ name: 'Read', arguments: { path: '/tmp/a' } }] }), + continueQuery, + })) + + expect(yielded[0]).toMatchObject({ type: 'assistant' }) + expect(yielded.some(message => 'type' in message && message.type === 'user')).toBe(true) + expect(providerMessages[0]).toBe(params.messages[0]) + expect(providerMessages[1]).toMatchObject({ type: 'assistant' }) + expect(providerMessages).toHaveLength(3) + expect(providerMessages[2]).toMatchObject({ type: 'user' }) + const assistantContent = providerMessages[1]?.type === 'assistant' ? providerMessages[1].message.content : [] + const userContent = providerMessages[2]?.type === 'user' ? providerMessages[2].message.content : [] + const toolUse = Array.isArray(assistantContent) ? assistantContent.find(block => block.type === 'tool_use') : undefined + const toolResult = Array.isArray(userContent) ? userContent.find(block => block.type === 'tool_result') : undefined + expect(toolResult?.tool_use_id).toBe(toolUse?.id) + expect(getInProgress().size).toBe(0) + }) +}) diff --git a/src/services/needle/runNeedleWorkerTurn.ts b/src/services/needle/runNeedleWorkerTurn.ts new file mode 100644 index 0000000000..6ce78d0900 --- /dev/null +++ b/src/services/needle/runNeedleWorkerTurn.ts @@ -0,0 +1,150 @@ +import { randomUUID } from 'node:crypto' +import type { ToolUseBlock } from '@anthropic-ai/sdk/resources/index.mjs' +import { createAssistantMessage, createUserMessage } from '../../utils/messages/factories.js' +import { normalizeMessagesForAPI } from '../../utils/messages.js' +import { runToolUse } from '../tools/toolExecution.js' +import type { AssistantMessage, Message, RequestStartEvent, StreamEvent, TombstoneMessage, ToolUseSummaryMessage } from '../../types/message.js' +import type { Terminal } from '../../query/transitions.js' +import { zodToJsonSchema } from '../../utils/zodToJsonSchema.js' +import type { CanUseToolFn } from '../../hooks/useCanUseTool.js' +import type { ToolUseContext, Tools } from '../../Tool.js' +import { buildNeedleToolPayload, NeedleClient } from './needleClient.js' +import { registerCleanup } from '../../utils/cleanupRegistry.js' +import { query, type QueryParams } from '../../query.js' + +type NeedleClientOwner = object + +const needleClients = new WeakMap() +const activeNeedleClients = new Set() +registerCleanup(() => { + for (const client of activeNeedleClients) client.close() + activeNeedleClients.clear() +}) + +function getNeedleClient(context: NeedleClientOwner): NeedleClient { + const existing = needleClients.get(context) + if (existing) return existing + const client = new NeedleClient() + needleClients.set(context, client) + activeNeedleClients.add(client) + return client +} + +export function closeNeedleClient(context: NeedleClientOwner): void { + const client = needleClients.get(context) + if (!client) return + client.close() + needleClients.delete(context) + activeNeedleClients.delete(client) +} + +export async function* runNeedleWorkerTurn({ + prompt, + queryParams, + tools, + toolUseContext, + clientOwner = toolUseContext, + canUseTool, + signal, + client, + continueQuery = query, +}: { + prompt: string + queryParams: QueryParams + tools: Tools + toolUseContext: ToolUseContext + clientOwner?: NeedleClientOwner + canUseTool: CanUseToolFn + signal?: AbortSignal + client?: Pick + continueQuery?: typeof query +}): AsyncGenerator { + const needleClient = client ?? getNeedleClient(clientOwner) + try { + let context = toolUseContext + const executableTools = tools.filter(tool => tool.isEnabled()) + let response + try { + const needleTools = await Promise.all(executableTools.map(async tool => ({ + name: tool.name, + description: await tool.description({}, { + isNonInteractiveSession: context.options.isNonInteractiveSession ?? false, + toolPermissionContext: context.getAppState().toolPermissionContext, + tools: executableTools, + }), + input_schema: tool.inputJSONSchema ?? zodToJsonSchema(tool.inputSchema), + }))) + response = await needleClient.complete({ prompt, tools: buildNeedleToolPayload(needleTools), signal }) + } catch { + if (signal?.aborted) return { reason: 'aborted_streaming' } + return yield* continueQuery({ ...queryParams, toolUseContext: context }) + } + if (signal?.aborted) return { reason: 'aborted_streaming' } + if (response.functionCalls.length !== 1) { + return yield* continueQuery({ ...queryParams, toolUseContext: context }) + } + + const call = response.functionCalls[0]! + const toolUse = { + type: 'tool_use' as const, + id: `toolu_needle_${randomUUID()}`, + name: call.name, + input: call.arguments, + caller: { type: 'direct' as const }, + } satisfies ToolUseBlock + const assistantMessage = createAssistantMessage({ content: [toolUse] }) as AssistantMessage + context.setInProgressToolUseIDs(prev => new Set(prev).add(toolUse.id)) + yield assistantMessage + if (signal?.aborted) { + context.setInProgressToolUseIDs(prev => { + const next = new Set(prev) + next.delete(toolUse.id) + return next + }) + yield createUserMessage({ + content: [{ + type: 'tool_result', + tool_use_id: toolUse.id, + content: 'Cancelled before tool execution.', + is_error: true, + }], + }) + return { reason: 'aborted_streaming' } + } + + const resultMessages: Message[] = [] + let shouldPreventContinuation = false + try { + for await (const update of runToolUse(toolUse, assistantMessage, canUseTool, context)) { + if (update.contextModifier) context = update.contextModifier.modifyContext(context) + yield update.message + if (update.message.type === 'user') resultMessages.push(update.message) + if ( + update.message.type === 'attachment' && + update.message.attachment.type === 'hook_stopped_continuation' + ) { + shouldPreventContinuation = true + } + } + } finally { + context.setInProgressToolUseIDs(prev => { + const next = new Set(prev) + next.delete(toolUse.id) + return next + }) + } + + if (shouldPreventContinuation) return { reason: 'hook_stopped' } + if (signal?.aborted) return { reason: 'aborted_streaming' } + + const normalizedResults = normalizeMessagesForAPI(resultMessages, executableTools) + .filter(message => message.type === 'user') + return yield* continueQuery({ + ...queryParams, + messages: [...queryParams.messages, assistantMessage, ...normalizedResults], + toolUseContext: context, + }) + } finally { + client?.close() + } +} diff --git a/src/tools/AgentTool/runAgent.ts b/src/tools/AgentTool/runAgent.ts index 72bb87e868..ff4518858c 100644 --- a/src/tools/AgentTool/runAgent.ts +++ b/src/tools/AgentTool/runAgent.ts @@ -18,6 +18,7 @@ import type { Terminal } from '../../query/transitions.js' import { getFeatureValue_CACHED_MAY_BE_STALE } from '../../services/analytics/growthbook.js' import type { AppState } from '../../state/AppStateStore.js' import { getDumpPromptsPath } from '../../services/api/dumpPrompts.js' +import { getGlobalConfig } from '../../utils/config.js' import { cleanupAgentTracking } from '../../services/api/promptCacheBreakDetection.js' import { connectToServer, @@ -58,7 +59,8 @@ import { import { registerFrontmatterHooks } from '../../utils/hooks/registerFrontmatterHooks.js' import { clearSessionHooks } from '../../utils/hooks/sessionHooks.js' import { executeSubagentStartHooks } from '../../utils/hooks.js' -import { createUserMessage } from '../../utils/messages.js' +import { createUserMessage, getContentText } from '../../utils/messages.js' +import { selectableUserMessagesFilter } from '../../utils/messageFilters.js' import { getAgentModel } from '../../utils/model/agent.js' import { isModelAllowed } from '../../utils/model/modelAllowlist.js' import { resolveAgentRunModelRouting, shouldEnforceModelAllowlist } from '../../services/api/agentRouting.js' @@ -81,6 +83,8 @@ import type { ContentReplacementState } from '../../utils/toolResultStorage.js' import { createAgentId } from '../../utils/uuid.js' import { resolveAgentTools } from './agentToolUtils.js' import { type AgentDefinition, isBuiltInAgent } from './loadAgentsDir.js' +import { closeNeedleClient, runNeedleWorkerTurn } from '../../services/needle/runNeedleWorkerTurn.js' +import { NEEDLE_WORKER_MODEL } from '../../services/needle/needleClient.js' import { canDescendantShowPermissionPrompts, shouldAvoidAgentPermissionPrompts, @@ -390,7 +394,7 @@ export async function* runAgent({ // Resolve per-agent provider routing from settings const settings = getInitialSettings() - const { mainLoopModel: effectiveModel, providerOverride } = + const { mainLoopModel: routedModel, providerOverride } = resolveAgentRunModelRouting({ resolvedAgentModel, parentModel: toolUseContext.options.mainLoopModel, @@ -401,6 +405,12 @@ export async function* runAgent({ settings, permissionMode, }) + const useNeedlePlanner = routedModel === NEEDLE_WORKER_MODEL + const effectiveModel = useNeedlePlanner + ? resolvedAgentModel === NEEDLE_WORKER_MODEL + ? toolUseContext.options.mainLoopModel + : resolvedAgentModel + : routedModel if ( shouldEnforceModelAllowlist( @@ -873,7 +883,7 @@ export async function* runAgent({ agentDefinition.maxSteps! > 0 ? agentDefinition.maxSteps : undefined - const queryIterator = query({ + const queryParams = { messages: initialMessages, systemPrompt: agentSystemPrompt, userContext: resolvedUserContext, @@ -889,7 +899,28 @@ export async function* runAgent({ agentType: agentDefinition.agentType, } : undefined, - })[Symbol.asyncIterator]() + } + const needleConfig = getGlobalConfig() + const latestUserMessage = initialMessages.findLast(selectableUserMessagesFilter) + const latestUserContent = latestUserMessage?.message.content + const needlePrompt = latestUserContent ? getContentText(latestUserContent) ?? '' : '' + const hasImageContent = + Array.isArray(latestUserContent) && latestUserContent.some(block => block.type === 'image') + const queryIterator = + useNeedlePlanner && + needleConfig.needleEnabled === true && + needlePrompt.length > 0 && + !hasImageContent && + (needleConfig.needleMode ?? 'worker') === 'worker' + ? runNeedleWorkerTurn({ + prompt: needlePrompt, + queryParams, + tools: allTools, + toolUseContext: agentToolUseContext, + canUseTool, + signal: agentAbortController.signal, + }) + : query(queryParams)[Symbol.asyncIterator]() try { while (true) { @@ -964,6 +995,7 @@ export async function* runAgent({ agentDefinition.callback() } } finally { + closeNeedleClient(agentToolUseContext) // Clean up agent-specific MCP servers (runs on normal completion, abort, or error) await mcpCleanup() // Clean up agent's session hooks diff --git a/src/utils/config.ts b/src/utils/config.ts index 709e8851ef..c4446d4ee0 100644 --- a/src/utils/config.ts +++ b/src/utils/config.ts @@ -689,6 +689,10 @@ export type GlobalConfig = { // Knowledge Graph configuration knowledgeGraphEnabled: boolean + // Needle local planner configuration. + needleEnabled?: boolean + needleMode?: 'worker' | 'main' + // Startup splash logo color scheme — set via /logo. See // src/components/StartupScreen.palettes.ts for valid values. Stored as a // plain string (validated on read) to avoid pulling a UI module into the @@ -760,6 +764,8 @@ function createDefaultGlobalConfig(): GlobalConfig { providerProfiles: [], openaiAdditionalModelOptionsCacheByProfile: {}, knowledgeGraphEnabled: true, + needleEnabled: false, + needleMode: 'worker', // Omitted by default so callers can distinguish "unset" from an explicit // persisted "off"; normalizeMaxMessagesCompactionThreshold resolves an // unset value to the effective default of '200' (message-count compaction @@ -819,6 +825,8 @@ export const GLOBAL_CONFIG_KEYS = [ 'remoteControlAtStartup', 'remoteDialogSeen', 'knowledgeGraphEnabled', + 'needleEnabled', + 'needleMode', 'logoColor', 'maxMessagesCompactionThreshold', 'compactModel', @@ -932,6 +940,8 @@ function getTestGlobalConfigForTesting(): GlobalConfig { ...DEFAULT_GLOBAL_CONFIG, autoUpdates: false, knowledgeGraphEnabled: true, + needleEnabled: false, + needleMode: 'worker', } } return testGlobalConfigForTesting