From f510e7a11154b10a166dc1a1f79e1f8ec6208332 Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Sun, 18 Oct 2026 14:50:50 +0000 Subject: [PATCH 01/10] feat: integrate Needle local planner bridge & fallback --- docs/advanced-setup.md | 26 ++ src/commands.ts | 2 + src/commands/needle/index.ts | 12 + src/commands/needle/needle.test.ts | 57 ++++ src/commands/needle/needle.ts | 62 ++++ src/screens/REPL.tsx | 24 +- src/services/needle/needleClient.test.ts | 88 ++++++ src/services/needle/needleClient.ts | 270 ++++++++++++++++++ .../needle/runNeedleWorkerTurn.test.ts | 160 +++++++++++ src/services/needle/runNeedleWorkerTurn.ts | 92 ++++++ src/tools/AgentTool/runAgent.ts | 29 +- src/utils/config.ts | 10 + 12 files changed, 824 insertions(+), 8 deletions(-) create mode 100644 src/commands/needle/index.ts create mode 100644 src/commands/needle/needle.test.ts create mode 100644 src/commands/needle/needle.ts create mode 100644 src/services/needle/needleClient.test.ts create mode 100644 src/services/needle/needleClient.ts create mode 100644 src/services/needle/runNeedleWorkerTurn.test.ts create mode 100644 src/services/needle/runNeedleWorkerTurn.ts 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..94e00c682f --- /dev/null +++ b/src/commands/needle/needle.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, it, beforeEach, afterEach } from 'bun:test' +import { mkdtempSync, rmSync } from 'fs' +import { tmpdir } from 'os' +import { join } from 'path' +import { call as needleCall } from './needle.js' +import { getGlobalConfig } 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(() => { + 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 context = {} as unknown as LocalJSXCommandContext + 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') + }) +}) 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..b091f28a2e 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 { 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'; @@ -3292,7 +3293,8 @@ export function REPL({ resetTurnToolDuration(); resetTurnClassifierDuration(); let expectedAutoCompactTracking = queryAutoCompactTracking; - const queryGenerator = query({ + const needleConfig = getGlobalConfig(); + const queryParams = { messages: messagesIncludingNewMessages, requestOnlyMessages, onModelRequestStart: interruptionCorrectionQueryId @@ -3301,8 +3303,8 @@ export function REPL({ shouldQuery, isInterruptionCorrectionEligible: true, queryId: interruptionCorrectionQueryId, - }) - onModelRequestStart?.() + }); + onModelRequestStart?.(); } : undefined, systemPrompt, @@ -3318,13 +3320,25 @@ export function REPL({ if (setAutoCompactTrackingForSessionIfUnchanged(querySessionId, expectedAutoCompactTracking, tracking)) { expectedAutoCompactTracking = tracking; } - } - }); + }, + }; + const queryGenerator = + needleConfig.needleEnabled === true && needleConfig.needleMode === 'main' + ? runNeedleWorkerTurn({ + prompt: JSON.stringify(messagesIncludingNewMessages), + queryParams, + tools: freshTools, + toolUseContext, + 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..d719b96bc6 --- /dev/null +++ b/src/services/needle/needleClient.test.ts @@ -0,0 +1,88 @@ +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('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('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('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..b7eb52e123 --- /dev/null +++ b/src/services/needle/needleClient.ts @@ -0,0 +1,270 @@ +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: + 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, + env: { ...process.env, NEEDLE_TELEMETRY: '0' }, + }) + 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.rejectAll(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, + env: { ...process.env, NEEDLE_TELEMETRY: '0' }, + }) + this.process = child + this.toolsKey = key + this.lines = createInterface({ input: child.stdout }) + 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 + this.pending.delete(response.id) + clearTimeout(pending.timer) + pending.abortCleanup?.() + if (response.error) pending.reject(new Error(`Needle worker error: ${response.error}`)) + else pending.resolve(normalizeResponse(response)) + } 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.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..f02a7d601c --- /dev/null +++ b/src/services/needle/runNeedleWorkerTurn.test.ts @@ -0,0 +1,160 @@ +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('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..170ba4da3f --- /dev/null +++ b/src/services/needle/runNeedleWorkerTurn.ts @@ -0,0 +1,92 @@ +import { randomUUID } from 'node:crypto' +import type { ToolUseBlock } from '@anthropic-ai/sdk/resources/index.mjs' +import { createAssistantMessage } 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 { query, type QueryParams } from '../../query.js' + +export async function* runNeedleWorkerTurn({ + prompt, + queryParams, + tools, + toolUseContext, + canUseTool, + signal, + client = new NeedleClient(), + continueQuery = query, +}: { + prompt: string + queryParams: QueryParams + tools: Tools + toolUseContext: ToolUseContext + canUseTool: CanUseToolFn + signal?: AbortSignal + client?: Pick + continueQuery?: typeof query +}): AsyncGenerator { + try { + let context = toolUseContext + const executableTools = tools.filter(tool => tool.isEnabled()) + 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), + }))) + let response + try { + response = await client.complete({ prompt, tools: buildNeedleToolPayload(needleTools), signal }) + } catch { + return yield* continueQuery({ ...queryParams, toolUseContext: context }) + } + 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 + + const resultMessages: Message[] = [] + 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) + } + } finally { + context.setInProgressToolUseIDs(prev => { + const next = new Set(prev) + next.delete(toolUse.id) + return next + }) + } + + 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..5ff30c3180 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, @@ -81,6 +82,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 { runNeedleWorkerTurn } from '../../services/needle/runNeedleWorkerTurn.js' +import { NEEDLE_WORKER_MODEL } from '../../services/needle/needleClient.js' import { canDescendantShowPermissionPrompts, shouldAvoidAgentPermissionPrompts, @@ -390,7 +393,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 +404,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 +882,7 @@ export async function* runAgent({ agentDefinition.maxSteps! > 0 ? agentDefinition.maxSteps : undefined - const queryIterator = query({ + const queryParams = { messages: initialMessages, systemPrompt: agentSystemPrompt, userContext: resolvedUserContext, @@ -889,7 +898,21 @@ export async function* runAgent({ agentType: agentDefinition.agentType, } : undefined, - })[Symbol.asyncIterator]() + } + const needleConfig = getGlobalConfig() + const queryIterator = + effectiveModel === NEEDLE_WORKER_MODEL && + needleConfig.needleEnabled === true && + (needleConfig.needleMode ?? 'worker') === 'worker' + ? runNeedleWorkerTurn({ + prompt: JSON.stringify(initialMessages), + queryParams, + tools: allTools, + toolUseContext: agentToolUseContext, + canUseTool, + signal: override?.abortController?.signal, + }) + : query(queryParams)[Symbol.asyncIterator]() try { while (true) { 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 From 9b48fade94b92666119f3ec63dc23bb07b1c085d Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Sun, 18 Oct 2026 16:49:50 +0000 Subject: [PATCH 02/10] fix: harden Needle worker lifecycle Keep session workers reusable while releasing completed agent workers, and make routing, cancellation, subprocess failures, and focused tests reliable. --- src/commands/needle/needle.test.ts | 33 ++++++++++++++-------- src/screens/REPL.tsx | 6 +++- src/services/needle/needleClient.test.ts | 6 ++-- src/services/needle/needleClient.ts | 5 +++- src/services/needle/runNeedleWorkerTurn.ts | 32 +++++++++++++++++++-- src/tools/AgentTool/runAgent.ts | 16 +++++++---- 6 files changed, 74 insertions(+), 24 deletions(-) diff --git a/src/commands/needle/needle.test.ts b/src/commands/needle/needle.test.ts index 94e00c682f..9e78ebfb44 100644 --- a/src/commands/needle/needle.test.ts +++ b/src/commands/needle/needle.test.ts @@ -1,9 +1,9 @@ import { describe, expect, it, beforeEach, afterEach } from 'bun:test' -import { mkdtempSync, rmSync } from 'fs' +import { chmodSync, mkdtempSync, rmSync, writeFileSync } from 'fs' import { tmpdir } from 'os' import { join } from 'path' import { call as needleCall } from './needle.js' -import { getGlobalConfig } from '../../utils/config.js' +import { getGlobalConfig, saveGlobalConfig } from '../../utils/config.js' import { setClaudeConfigHomeDirForTesting } from '../../utils/envUtils.js' import type { LocalJSXCommandContext } from '../../types/command.js' import { @@ -20,6 +20,7 @@ describe('needle command', () => { setClaudeConfigHomeDirForTesting(configDir) }) afterEach(() => { + saveGlobalConfig(current => ({ ...current, needleEnabled: false, needleMode: 'worker' })) setClaudeConfigHomeDirForTesting(undefined) delete process.env.CLAUDE_CONFIG_DIR rmSync(configDir, { recursive: true, force: true }) @@ -41,17 +42,27 @@ describe('needle command', () => { } }) - it('toggles the worker and persists its mode', async () => { + it.skipIf(process.platform === 'win32')('toggles the worker and persists its mode', async () => { + const pythonStub = join(configDir, 'python-stub') + writeFileSync(pythonStub, '#!/bin/sh\nexit 0\n') + chmodSync(pythonStub, 0o755) + const previous = process.env.OPENCLAUDE_NEEDLE_PYTHON + process.env.OPENCLAUDE_NEEDLE_PYTHON = pythonStub const context = {} as unknown as LocalJSXCommandContext - await needleCall('enable', context) - expect(getGlobalConfig().needleEnabled).toBe(true) - expect(getGlobalConfig().needleMode).toBe('worker') + 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') + 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') + const status = await needleCall('status', context) + expect(status.type).toBe('text') + if (status.type === 'text') expect(status.value).toContain('main') + } finally { + if (previous === undefined) delete process.env.OPENCLAUDE_NEEDLE_PYTHON + else process.env.OPENCLAUDE_NEEDLE_PYTHON = previous + } }) }) diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index b091f28a2e..33b6224e54 100644 --- a/src/screens/REPL.tsx +++ b/src/screens/REPL.tsx @@ -3322,10 +3322,14 @@ export function REPL({ } }, }; + const latestUserMessage = messagesIncludingNewMessages.findLast(selectableUserMessagesFilter); + const needlePrompt = latestUserMessage + ? getContentText(latestUserMessage.message.content) ?? '' + : ''; const queryGenerator = needleConfig.needleEnabled === true && needleConfig.needleMode === 'main' ? runNeedleWorkerTurn({ - prompt: JSON.stringify(messagesIncludingNewMessages), + prompt: needlePrompt, queryParams, tools: freshTools, toolUseContext, diff --git a/src/services/needle/needleClient.test.ts b/src/services/needle/needleClient.test.ts index d719b96bc6..bc70c5c21d 100644 --- a/src/services/needle/needleClient.test.ts +++ b/src/services/needle/needleClient.test.ts @@ -36,7 +36,7 @@ describe('Needle client configuration', () => { ]) }) - test('restarts the worker after a child exit before serving later requests', async () => { + 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) @@ -50,7 +50,7 @@ describe('Needle client configuration', () => { } }) - test('restarts the worker after a zero child exit before serving later requests', async () => { + 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) @@ -64,7 +64,7 @@ describe('Needle client configuration', () => { } }) - test('restarts the worker after a timeout before serving later requests', async () => { + 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') diff --git a/src/services/needle/needleClient.ts b/src/services/needle/needleClient.ts index b7eb52e123..152106d78d 100644 --- a/src/services/needle/needleClient.ts +++ b/src/services/needle/needleClient.ts @@ -12,6 +12,7 @@ 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"]: @@ -197,7 +198,7 @@ export class NeedleClient { this.pending.set(id, pending) this.process!.stdin.write(`${JSON.stringify(payload)}\n`, error => { if (!error) return - this.rejectAll(new Error(`Unable to write to Needle worker: ${error.message}`)) + this.restart(new Error(`Unable to write to Needle worker: ${error.message}`)) }) }) } @@ -234,6 +235,7 @@ export class NeedleClient { 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 @@ -252,6 +254,7 @@ export class NeedleClient { 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'}`)) diff --git a/src/services/needle/runNeedleWorkerTurn.ts b/src/services/needle/runNeedleWorkerTurn.ts index 170ba4da3f..f8c326a6aa 100644 --- a/src/services/needle/runNeedleWorkerTurn.ts +++ b/src/services/needle/runNeedleWorkerTurn.ts @@ -9,8 +9,33 @@ 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' +const needleClients = new WeakMap() +const activeNeedleClients = new Set() +registerCleanup(() => { + for (const client of activeNeedleClients) client.close() + activeNeedleClients.clear() +}) + +function getNeedleClient(context: ToolUseContext): 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: ToolUseContext): void { + const client = needleClients.get(context) + if (!client) return + client.close() + needleClients.delete(context) + activeNeedleClients.delete(client) +} + export async function* runNeedleWorkerTurn({ prompt, queryParams, @@ -18,7 +43,7 @@ export async function* runNeedleWorkerTurn({ toolUseContext, canUseTool, signal, - client = new NeedleClient(), + client, continueQuery = query, }: { prompt: string @@ -30,6 +55,7 @@ export async function* runNeedleWorkerTurn({ client?: Pick continueQuery?: typeof query }): AsyncGenerator { + const needleClient = client ?? getNeedleClient(toolUseContext) try { let context = toolUseContext const executableTools = tools.filter(tool => tool.isEnabled()) @@ -44,7 +70,7 @@ export async function* runNeedleWorkerTurn({ }))) let response try { - response = await client.complete({ prompt, tools: buildNeedleToolPayload(needleTools), signal }) + response = await needleClient.complete({ prompt, tools: buildNeedleToolPayload(needleTools), signal }) } catch { return yield* continueQuery({ ...queryParams, toolUseContext: context }) } @@ -87,6 +113,6 @@ export async function* runNeedleWorkerTurn({ toolUseContext: context, }) } finally { - client.close() + client?.close() } } diff --git a/src/tools/AgentTool/runAgent.ts b/src/tools/AgentTool/runAgent.ts index 5ff30c3180..7251051b09 100644 --- a/src/tools/AgentTool/runAgent.ts +++ b/src/tools/AgentTool/runAgent.ts @@ -59,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' @@ -82,7 +83,7 @@ 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 { runNeedleWorkerTurn } from '../../services/needle/runNeedleWorkerTurn.js' +import { closeNeedleClient, runNeedleWorkerTurn } from '../../services/needle/runNeedleWorkerTurn.js' import { NEEDLE_WORKER_MODEL } from '../../services/needle/needleClient.js' import { canDescendantShowPermissionPrompts, @@ -900,17 +901,21 @@ export async function* runAgent({ : undefined, } const needleConfig = getGlobalConfig() + const latestUserMessage = initialMessages.findLast(selectableUserMessagesFilter) + const needlePrompt = latestUserMessage + ? getContentText(latestUserMessage.message.content) ?? '' + : '' const queryIterator = - effectiveModel === NEEDLE_WORKER_MODEL && + useNeedlePlanner && needleConfig.needleEnabled === true && (needleConfig.needleMode ?? 'worker') === 'worker' ? runNeedleWorkerTurn({ - prompt: JSON.stringify(initialMessages), + prompt: needlePrompt, queryParams, tools: allTools, toolUseContext: agentToolUseContext, canUseTool, - signal: override?.abortController?.signal, + signal: agentAbortController.signal, }) : query(queryParams)[Symbol.asyncIterator]() @@ -987,6 +992,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 From 01b4bad5858eed2c33e6256bb7b20464b025741a Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Sun, 18 Oct 2026 17:41:42 +0000 Subject: [PATCH 03/10] fix: stabilize Needle REPL sessions Reuse one planner worker per REPL session, bypass textless turns, and keep command coverage platform-neutral. --- src/commands/needle/needle.test.ts | 16 ++++++---------- src/screens/REPL.tsx | 8 ++++++-- src/services/needle/runNeedleWorkerTurn.ts | 12 ++++++++---- 3 files changed, 20 insertions(+), 16 deletions(-) diff --git a/src/commands/needle/needle.test.ts b/src/commands/needle/needle.test.ts index 9e78ebfb44..9cf434676b 100644 --- a/src/commands/needle/needle.test.ts +++ b/src/commands/needle/needle.test.ts @@ -1,8 +1,9 @@ -import { describe, expect, it, beforeEach, afterEach } from 'bun:test' -import { chmodSync, mkdtempSync, rmSync, writeFileSync } from 'fs' +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' @@ -42,12 +43,8 @@ describe('needle command', () => { } }) - it.skipIf(process.platform === 'win32')('toggles the worker and persists its mode', async () => { - const pythonStub = join(configDir, 'python-stub') - writeFileSync(pythonStub, '#!/bin/sh\nexit 0\n') - chmodSync(pythonStub, 0o755) - const previous = process.env.OPENCLAUDE_NEEDLE_PYTHON - process.env.OPENCLAUDE_NEEDLE_PYTHON = pythonStub + 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) @@ -61,8 +58,7 @@ describe('needle command', () => { expect(status.type).toBe('text') if (status.type === 'text') expect(status.value).toContain('main') } finally { - if (previous === undefined) delete process.env.OPENCLAUDE_NEEDLE_PYTHON - else process.env.OPENCLAUDE_NEEDLE_PYTHON = previous + availability.mockRestore() } }) }) diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index 33b6224e54..05ae442871 100644 --- a/src/screens/REPL.tsx +++ b/src/screens/REPL.tsx @@ -148,7 +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 { runNeedleWorkerTurn } from '../services/needle/runNeedleWorkerTurn.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'; @@ -682,9 +682,12 @@ 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); + useEffect(() => () => closeNeedleClient(needleClientOwnerRef.current), []); + useEffect(() => { if ( shouldShowReplMaxTurnsUnlimitedWarning(maxTurnsProp, { @@ -3327,12 +3330,13 @@ export function REPL({ ? getContentText(latestUserMessage.message.content) ?? '' : ''; const queryGenerator = - needleConfig.needleEnabled === true && needleConfig.needleMode === 'main' + needleConfig.needleEnabled === true && needleConfig.needleMode === 'main' && needlePrompt ? runNeedleWorkerTurn({ prompt: needlePrompt, queryParams, tools: freshTools, toolUseContext, + clientOwner: needleClientOwnerRef.current, canUseTool, signal: abortController.signal, }) diff --git a/src/services/needle/runNeedleWorkerTurn.ts b/src/services/needle/runNeedleWorkerTurn.ts index f8c326a6aa..7017ff6446 100644 --- a/src/services/needle/runNeedleWorkerTurn.ts +++ b/src/services/needle/runNeedleWorkerTurn.ts @@ -12,14 +12,16 @@ import { buildNeedleToolPayload, NeedleClient } from './needleClient.js' import { registerCleanup } from '../../utils/cleanupRegistry.js' import { query, type QueryParams } from '../../query.js' -const needleClients = new WeakMap() +type NeedleClientOwner = object + +const needleClients = new WeakMap() const activeNeedleClients = new Set() registerCleanup(() => { for (const client of activeNeedleClients) client.close() activeNeedleClients.clear() }) -function getNeedleClient(context: ToolUseContext): NeedleClient { +function getNeedleClient(context: NeedleClientOwner): NeedleClient { const existing = needleClients.get(context) if (existing) return existing const client = new NeedleClient() @@ -28,7 +30,7 @@ function getNeedleClient(context: ToolUseContext): NeedleClient { return client } -export function closeNeedleClient(context: ToolUseContext): void { +export function closeNeedleClient(context: NeedleClientOwner): void { const client = needleClients.get(context) if (!client) return client.close() @@ -41,6 +43,7 @@ export async function* runNeedleWorkerTurn({ queryParams, tools, toolUseContext, + clientOwner = toolUseContext, canUseTool, signal, client, @@ -50,12 +53,13 @@ export async function* runNeedleWorkerTurn({ queryParams: QueryParams tools: Tools toolUseContext: ToolUseContext + clientOwner?: NeedleClientOwner canUseTool: CanUseToolFn signal?: AbortSignal client?: Pick continueQuery?: typeof query }): AsyncGenerator { - const needleClient = client ?? getNeedleClient(toolUseContext) + const needleClient = client ?? getNeedleClient(clientOwner) try { let context = toolUseContext const executableTools = tools.filter(tool => tool.isEnabled()) From 2a30e985294ce54887d262068ba008a1ccc1dd72 Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Sun, 18 Oct 2026 18:30:20 +0000 Subject: [PATCH 04/10] fix: scope Needle planning to current turn Avoid reusing stale transcript prompts when speculative or textless turns continue directly with the provider. --- src/screens/REPL.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index 05ae442871..23aa562936 100644 --- a/src/screens/REPL.tsx +++ b/src/screens/REPL.tsx @@ -3325,9 +3325,9 @@ export function REPL({ } }, }; - const latestUserMessage = messagesIncludingNewMessages.findLast(selectableUserMessagesFilter); - const needlePrompt = latestUserMessage - ? getContentText(latestUserMessage.message.content) ?? '' + const currentUserMessage = newMessages.findLast(selectableUserMessagesFilter); + const needlePrompt = currentUserMessage + ? getContentText(currentUserMessage.message.content) ?? '' : ''; const queryGenerator = needleConfig.needleEnabled === true && needleConfig.needleMode === 'main' && needlePrompt From 6b2d82755b43988e611352d391b7624601f5636a Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Sun, 18 Oct 2026 19:14:45 +0000 Subject: [PATCH 05/10] fix: bypass Needle for image turns Route image-bearing and empty-prompt turns directly to the provider so planner tool choices never omit required visual context. --- src/screens/REPL.tsx | 12 ++++++++---- src/tools/AgentTool/runAgent.ts | 9 ++++++--- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index 23aa562936..4466f7d391 100644 --- a/src/screens/REPL.tsx +++ b/src/screens/REPL.tsx @@ -3326,11 +3326,15 @@ export function REPL({ }, }; const currentUserMessage = newMessages.findLast(selectableUserMessagesFilter); - const needlePrompt = currentUserMessage - ? getContentText(currentUserMessage.message.content) ?? '' - : ''; + 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 + needleConfig.needleEnabled === true && + needleConfig.needleMode === 'main' && + needlePrompt && + !hasImageContent ? runNeedleWorkerTurn({ prompt: needlePrompt, queryParams, diff --git a/src/tools/AgentTool/runAgent.ts b/src/tools/AgentTool/runAgent.ts index 7251051b09..ff4518858c 100644 --- a/src/tools/AgentTool/runAgent.ts +++ b/src/tools/AgentTool/runAgent.ts @@ -902,12 +902,15 @@ export async function* runAgent({ } const needleConfig = getGlobalConfig() const latestUserMessage = initialMessages.findLast(selectableUserMessagesFilter) - const needlePrompt = latestUserMessage - ? getContentText(latestUserMessage.message.content) ?? '' - : '' + 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, From 8ce756c62cc6e4199ec86838315f13aef9395db3 Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Sun, 18 Oct 2026 19:54:46 +0000 Subject: [PATCH 06/10] fix: preserve Needle turn safety Honor hook-stop results, reject malformed worker replies promptly, and preserve user telemetry configuration. --- src/services/needle/needleClient.test.ts | 20 ++++++++++++++++++++ src/services/needle/needleClient.ts | 19 ++++++++++++------- src/services/needle/runNeedleWorkerTurn.ts | 9 +++++++++ 3 files changed, 41 insertions(+), 7 deletions(-) diff --git a/src/services/needle/needleClient.test.ts b/src/services/needle/needleClient.test.ts index bc70c5c21d..0c72a4dca5 100644 --- a/src/services/needle/needleClient.test.ts +++ b/src/services/needle/needleClient.test.ts @@ -64,6 +64,26 @@ describe('Needle client configuration', () => { } }) + 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() + } 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') diff --git a/src/services/needle/needleClient.ts b/src/services/needle/needleClient.ts index 152106d78d..0380bd1540 100644 --- a/src/services/needle/needleClient.ts +++ b/src/services/needle/needleClient.ts @@ -89,7 +89,6 @@ export async function checkNeedleAvailability(): Promise<{ const child = spawn(command, [...args, '-c', 'import needle'], { stdio: ['ignore', 'ignore', 'pipe'], windowsHide: true, - env: { ...process.env, NEEDLE_TELEMETRY: '0' }, }) let stderr = '' const timer = setTimeout(() => { @@ -230,7 +229,6 @@ export class NeedleClient { const child = spawn(command, [...args, '-u', '-c', NEEDLE_RUNTIME], { stdio: ['pipe', 'pipe', 'pipe'], windowsHide: true, - env: { ...process.env, NEEDLE_TELEMETRY: '0' }, }) this.process = child this.toolsKey = key @@ -241,11 +239,18 @@ export class NeedleClient { const response = JSON.parse(line) as RuntimeResponse const pending = response.id ? this.pending.get(response.id) : undefined if (!pending) return - this.pending.delete(response.id) - clearTimeout(pending.timer) - pending.abortCleanup?.() - if (response.error) pending.reject(new Error(`Needle worker error: ${response.error}`)) - else pending.resolve(normalizeResponse(response)) + 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))) } diff --git a/src/services/needle/runNeedleWorkerTurn.ts b/src/services/needle/runNeedleWorkerTurn.ts index 7017ff6446..b894a1832c 100644 --- a/src/services/needle/runNeedleWorkerTurn.ts +++ b/src/services/needle/runNeedleWorkerTurn.ts @@ -95,11 +95,18 @@ export async function* runNeedleWorkerTurn({ yield assistantMessage 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 => { @@ -109,6 +116,8 @@ export async function* runNeedleWorkerTurn({ }) } + if (shouldPreventContinuation) return { reason: 'hook_stopped' } + const normalizedResults = normalizeMessagesForAPI(resultMessages, executableTools) .filter(message => message.type === 'user') return yield* continueQuery({ From de606928167c14e4bdb1f61676fceea6cbd0c5ec Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Sun, 18 Oct 2026 20:32:01 +0000 Subject: [PATCH 07/10] fix: complete Needle fallback coverage Fall back when tool metadata cannot be built and ensure malformed-response tests assert the intended rejection. --- src/services/needle/needleClient.test.ts | 2 +- src/services/needle/runNeedleWorkerTurn.ts | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/services/needle/needleClient.test.ts b/src/services/needle/needleClient.test.ts index 0c72a4dca5..541b467bbe 100644 --- a/src/services/needle/needleClient.test.ts +++ b/src/services/needle/needleClient.test.ts @@ -75,7 +75,7 @@ describe('Needle client configuration', () => { 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() + 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 diff --git a/src/services/needle/runNeedleWorkerTurn.ts b/src/services/needle/runNeedleWorkerTurn.ts index b894a1832c..b7ed3c0f65 100644 --- a/src/services/needle/runNeedleWorkerTurn.ts +++ b/src/services/needle/runNeedleWorkerTurn.ts @@ -63,17 +63,17 @@ export async function* runNeedleWorkerTurn({ try { let context = toolUseContext const executableTools = tools.filter(tool => tool.isEnabled()) - 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), - }))) 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 { return yield* continueQuery({ ...queryParams, toolUseContext: context }) From 3d0810fedb2e137963cbe874153ff07a693d5ee3 Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Mon, 19 Oct 2026 05:41:27 +0000 Subject: [PATCH 08/10] fix: stop Needle turns cleanly on shutdown Abort the active turn before closing its worker and suppress provider fallback after cancellation, including post-yield races. --- src/screens/REPL.tsx | 6 ++-- .../needle/runNeedleWorkerTurn.test.ts | 33 +++++++++++++++++++ src/services/needle/runNeedleWorkerTurn.ts | 4 +++ 3 files changed, 41 insertions(+), 2 deletions(-) diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index 4466f7d391..e5e8a03815 100644 --- a/src/screens/REPL.tsx +++ b/src/screens/REPL.tsx @@ -686,8 +686,6 @@ export function REPL({ const backgroundHandoffStartedRef = useRef(false); const [backgroundHandoffPreparing, setBackgroundHandoffPreparing] = useState(false); - useEffect(() => () => closeNeedleClient(needleClientOwnerRef.current), []); - useEffect(() => { if ( shouldShowReplMaxTurnsUnlimitedWarning(maxTurnsProp, { @@ -1021,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. diff --git a/src/services/needle/runNeedleWorkerTurn.test.ts b/src/services/needle/runNeedleWorkerTurn.test.ts index f02a7d601c..4602219492 100644 --- a/src/services/needle/runNeedleWorkerTurn.test.ts +++ b/src/services/needle/runNeedleWorkerTurn.test.ts @@ -116,6 +116,39 @@ describe('runNeedleWorkerTurn', () => { 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 } = 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, + }) + expect((await generator.next()).value).toMatchObject({ type: 'assistant' }) + controller.abort() + const result = await generator.next() + expect(result.done).toBe(true) + expect(result.value).toEqual({ reason: 'aborted_streaming' }) + expect(continued).toBe(false) + }) + test('yields a paired prelude and continues the provider with normalized results', async () => { const tool = createToolFixture(z.object({ path: z.string() }), { name: 'Read', diff --git a/src/services/needle/runNeedleWorkerTurn.ts b/src/services/needle/runNeedleWorkerTurn.ts index b7ed3c0f65..aac8bbd92e 100644 --- a/src/services/needle/runNeedleWorkerTurn.ts +++ b/src/services/needle/runNeedleWorkerTurn.ts @@ -76,8 +76,10 @@ export async function* runNeedleWorkerTurn({ }))) 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 }) } @@ -93,6 +95,7 @@ export async function* runNeedleWorkerTurn({ const assistantMessage = createAssistantMessage({ content: [toolUse] }) as AssistantMessage context.setInProgressToolUseIDs(prev => new Set(prev).add(toolUse.id)) yield assistantMessage + if (signal?.aborted) return { reason: 'aborted_streaming' } const resultMessages: Message[] = [] let shouldPreventContinuation = false @@ -117,6 +120,7 @@ export async function* runNeedleWorkerTurn({ } if (shouldPreventContinuation) return { reason: 'hook_stopped' } + if (signal?.aborted) return { reason: 'aborted_streaming' } const normalizedResults = normalizeMessagesForAPI(resultMessages, executableTools) .filter(message => message.type === 'user') From ada08043abaf2ecc846c55516621922b86f432a8 Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Mon, 19 Oct 2026 06:07:03 +0000 Subject: [PATCH 09/10] fix: clear Needle state on early abort Remove the pending tool-use ID when cancellation lands after the planner prelude, keeping subsequent turns free of unmatched work. --- src/services/needle/runNeedleWorkerTurn.test.ts | 3 ++- src/services/needle/runNeedleWorkerTurn.ts | 9 ++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/src/services/needle/runNeedleWorkerTurn.test.ts b/src/services/needle/runNeedleWorkerTurn.test.ts index 4602219492..73a2634c5e 100644 --- a/src/services/needle/runNeedleWorkerTurn.test.ts +++ b/src/services/needle/runNeedleWorkerTurn.test.ts @@ -123,7 +123,7 @@ describe('runNeedleWorkerTurn', () => { throw new Error('tool should not execute') }, }) - const { context } = makeContext([tool]) + const { context, getInProgress } = makeContext([tool]) const params = makeParams(context) const controller = new AbortController() let continued = false @@ -147,6 +147,7 @@ describe('runNeedleWorkerTurn', () => { expect(result.done).toBe(true) expect(result.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 () => { diff --git a/src/services/needle/runNeedleWorkerTurn.ts b/src/services/needle/runNeedleWorkerTurn.ts index aac8bbd92e..8e6a1aaa1d 100644 --- a/src/services/needle/runNeedleWorkerTurn.ts +++ b/src/services/needle/runNeedleWorkerTurn.ts @@ -95,7 +95,14 @@ export async function* runNeedleWorkerTurn({ const assistantMessage = createAssistantMessage({ content: [toolUse] }) as AssistantMessage context.setInProgressToolUseIDs(prev => new Set(prev).add(toolUse.id)) yield assistantMessage - if (signal?.aborted) return { reason: 'aborted_streaming' } + if (signal?.aborted) { + context.setInProgressToolUseIDs(prev => { + const next = new Set(prev) + next.delete(toolUse.id) + return next + }) + return { reason: 'aborted_streaming' } + } const resultMessages: Message[] = [] let shouldPreventContinuation = false From ce76b4315e3f3cc2558bc0021e6c0577ed11da00 Mon Sep 17 00:00:00 2001 From: OpenClaude Contributor Date: Mon, 19 Oct 2026 06:52:50 +0000 Subject: [PATCH 10/10] fix: pair cancelled Needle tool prelude Emit a matching error tool result when shutdown aborts after the planner prelude, and assert transcript pairing in the regression test. --- .../needle/runNeedleWorkerTurn.test.ts | 23 ++++++++++++++++--- src/services/needle/runNeedleWorkerTurn.ts | 10 +++++++- 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/src/services/needle/runNeedleWorkerTurn.test.ts b/src/services/needle/runNeedleWorkerTurn.test.ts index 73a2634c5e..dc53f8a094 100644 --- a/src/services/needle/runNeedleWorkerTurn.test.ts +++ b/src/services/needle/runNeedleWorkerTurn.test.ts @@ -141,11 +141,28 @@ describe('runNeedleWorkerTurn', () => { client: fakeClient({ functionCalls: [{ name: 'Read', arguments: { path: '/tmp/a' } }] }), continueQuery, }) - expect((await generator.next()).value).toMatchObject({ type: 'assistant' }) + const prelude = (await generator.next()).value + expect(prelude).toMatchObject({ type: 'assistant' }) controller.abort() const result = await generator.next() - expect(result.done).toBe(true) - expect(result.value).toEqual({ reason: 'aborted_streaming' }) + 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) }) diff --git a/src/services/needle/runNeedleWorkerTurn.ts b/src/services/needle/runNeedleWorkerTurn.ts index 8e6a1aaa1d..6ce78d0900 100644 --- a/src/services/needle/runNeedleWorkerTurn.ts +++ b/src/services/needle/runNeedleWorkerTurn.ts @@ -1,6 +1,6 @@ import { randomUUID } from 'node:crypto' import type { ToolUseBlock } from '@anthropic-ai/sdk/resources/index.mjs' -import { createAssistantMessage } from '../../utils/messages/factories.js' +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' @@ -101,6 +101,14 @@ export async function* runNeedleWorkerTurn({ 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' } }