From c468b5f99d14897f741e8657c9620b4de278a1c2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B6=E5=85=AC?= Date: Tue, 14 Jul 2026 21:16:57 +0800 Subject: [PATCH 1/9] feat(core): propagate invocation context --- docs/design/invocation-context/README.md | 196 +++ .../invocation-context/v1-conformance.json | 125 ++ .../cli/invocation-context-v1.test.ts | 464 ++++++++ .../fixtures/invocation-context-echo.mjs | 49 + packages/acp-bridge/src/bridge.test.ts | 109 +- packages/acp-bridge/src/bridge.ts | 49 +- packages/acp-bridge/src/internal/testUtils.ts | 4 +- packages/channels/base/src/AcpBridge.test.ts | 28 +- packages/channels/base/src/AcpBridge.ts | 25 +- .../channels/base/src/ChannelAgentBridge.ts | 10 +- .../channels/base/src/ChannelBase.test.ts | 13 +- packages/channels/base/src/ChannelBase.ts | 5 +- .../channels/base/src/DaemonChannelBridge.ts | 10 +- .../cli/src/acp-integration/acpAgent.test.ts | 210 ++++ packages/cli/src/acp-integration/acpAgent.ts | 95 +- .../acp-integration/session/Session.test.ts | 104 +- .../src/acp-integration/session/Session.ts | 90 +- .../channel/memory-intent-classifier.test.ts | 2 +- .../channel/memory-intent-classifier.ts | 2 +- packages/cli/src/gemini.test.tsx | 79 +- packages/cli/src/gemini.tsx | 25 +- packages/cli/src/nonInteractiveCli.test.ts | 127 +- packages/cli/src/nonInteractiveCli.ts | 98 +- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 28 +- packages/cli/src/ui/hooks/useGeminiStream.ts | 41 +- packages/cli/src/utils/relaunch.test.ts | 23 + packages/cli/src/utils/relaunch.ts | 11 +- packages/cli/src/utils/sandbox.ts | 7 + .../agents/background-agent-resume.test.ts | 28 +- .../src/agents/background-agent-resume.ts | 16 +- .../src/agents/runtime/agent-core.test.ts | 57 + .../core/src/agents/runtime/agent-core.ts | 20 +- .../core/src/core/coreToolScheduler.test.ts | 65 + packages/core/src/core/coreToolScheduler.ts | 19 +- packages/core/src/index.ts | 1 + packages/core/src/tools/mcp-client.test.ts | 163 +++ packages/core/src/tools/mcp-client.ts | 18 + packages/core/src/tools/mcp-tool.test.ts | 269 +++++ packages/core/src/tools/mcp-tool.ts | 26 +- .../core/src/utils/invocation-context.test.ts | 138 +++ packages/core/src/utils/invocation-context.ts | 104 ++ packages/sdk-typescript/scripts/build.js | 21 +- .../src/daemon-mcp/serve-bridge/README.md | 4 + .../src/daemon-mcp/serve-bridge/bin.ts | 31 +- .../src/daemon-mcp/serve-bridge/bindings.ts | 206 ++++ .../createServeBridgeMcpServer.ts | 41 +- .../src/daemon-mcp/serve-bridge/helpers.ts | 23 +- .../src/daemon-mcp/serve-bridge/index.ts | 1 + .../src/daemon-mcp/serve-bridge/shutdown.ts | 44 + .../src/daemon-mcp/serve-bridge/sse.ts | 77 +- .../daemon-mcp/serve-bridge/tools/agent.ts | 270 +++-- .../daemon-mcp/serve-bridge/tools/session.ts | 160 ++- .../src/daemon-mcp/serve-bridge/types.ts | 20 +- .../sdk-typescript/src/daemon/DaemonClient.ts | 82 +- .../src/daemon/DaemonHttpError.ts | 12 + .../src/daemon/DaemonSessionClient.ts | 32 +- packages/sdk-typescript/src/daemon/index.ts | 2 + packages/sdk-typescript/src/index.ts | 7 +- .../test/unit/DaemonClient.test.ts | 191 ++- .../test/unit/serve-bridge-shutdown.test.ts | 76 ++ .../test/unit/serve-bridge.test.ts | 1049 ++++++++++++++++- 61 files changed, 4831 insertions(+), 471 deletions(-) create mode 100644 docs/design/invocation-context/README.md create mode 100644 docs/design/invocation-context/v1-conformance.json create mode 100644 integration-tests/cli/invocation-context-v1.test.ts create mode 100644 integration-tests/fixtures/invocation-context-echo.mjs create mode 100644 packages/core/src/utils/invocation-context.test.ts create mode 100644 packages/core/src/utils/invocation-context.ts create mode 100644 packages/sdk-typescript/src/daemon-mcp/serve-bridge/bindings.ts create mode 100644 packages/sdk-typescript/src/daemon-mcp/serve-bridge/shutdown.ts create mode 100644 packages/sdk-typescript/test/unit/serve-bridge-shutdown.test.ts diff --git a/docs/design/invocation-context/README.md b/docs/design/invocation-context/README.md new file mode 100644 index 00000000000..89090fc5375 --- /dev/null +++ b/docs/design/invocation-context/README.md @@ -0,0 +1,196 @@ +# Invocation Context v1 + +## Status + +Implementation design for Qwen Code v0.19.10 and later. + +## Problem + +Qwen Code can enter a session through the CLI, ACP, the HTTP daemon, +messaging channels, scheduled work, or another MCP server. Once the prompt is +inside `Session`, those origins are no longer available to a local MCP tool. +Passing caller-provided metadata through unchanged would be incorrect: public +ACP and daemon clients could forge session, prompt, or client identities, and +client-hosted MCP transports could receive data intended only for a process +started by Qwen Code. + +Invocation Context v1 provides correlation metadata with three constraints: + +1. Native entry points create or rebuild the context; request payloads do not + become authority by themselves. +2. The context exists only for the lifetime of an async execution tree and is + not written into chat history. +3. Only MCP servers launched through an actual stdio transport receive it. + +The context is correlation metadata, not an authorization credential. + +## Wire contract + +```ts +type InvocationIngress = + | 'cli' + | 'acp' + | 'daemon' + | 'channel' + | 'scheduler' + | 'external_mcp' + | 'internal'; + +interface InvocationContextV1 { + version: 1; + ingress: InvocationIngress; + sessionId: string; + promptId: string; + originatorClientId?: string; +} +``` + +The parser accepts only this shape. Version, ingress, and IDs are not coerced; +unknown properties and blank IDs are rejected. `promptId` is an opaque root +correlation ID. It does not replace the durable prompt ID used by history, +rewind, or file snapshots. + +The reserved metadata keys are: + +| Key | Use | +| ------------------------------------- | ------------------------------------------------------------------------- | +| `qwen-code/invocation` | Complete v1 context on a trusted ACP request or MCP tool request | +| `qwen-code/invocation-ingress` | Advisory ingress supplied to the daemon or a trusted standalone ACP child | +| `qwen-code/private-parent-capability` | Per-spawn capability in the first ACP initialize request | + +## Runtime propagation + +Core owns a dedicated `AsyncLocalStorage`. Root entry +points call `runWithInvocationContext`; normal in-process async work and +subagents inherit the store. Deferred approval callbacks restore the snapshot +captured with the request. Timer-driven and persisted work must not rely on +implicit inheritance: cron, notifications, and background resume each create a +new root context and omit `originatorClientId`. + +The sources are fixed: + +| Entry point | ingress | promptId | client identity | +| ------------------------------------------- | -------------- | ----------------------- | ---------------- | +| CLI | `cli` | current CLI turn ID | omitted | +| Public ACP | `acp` | child durable prompt ID | omitted | +| Daemon prompt or public continue | `daemon` | daemon admission UUID | validated client | +| qwen-serve-mcp | `external_mcp` | daemon admission UUID | validated client | +| User channel message | `channel` | current root request ID | omitted | +| Channel loop, cron, unattended webhook | `scheduler` | new execution root ID | omitted | +| Notification, background resume, classifier | `internal` | new execution root ID | omitted | + +## Daemon authority + +Daemon prompt metadata is advisory. The daemon first performs its existing +client attachment and identity validation. It then removes all three reserved +metadata keys and constructs a complete context from the route session ID, +daemon admission UUID, normalized ingress, and validated client ID. + +`channel`, `scheduler`, and `internal` deliberately drop the client ID. +`daemon` and `external_mcp` retain the validated client ID. Missing or unknown +advisory values normalize to `daemon`. Consequently, forging ingress can only +change a label or attenuate an already validated identity; it cannot grant an +identity. + +SDK defaults and per-prompt overrides are stamped after copying caller +metadata. The precedence is per-prompt override, `DaemonClient` default, then +`daemon`. Both blocking and non-blocking prompt paths use the same stamping +logic. + +## Private ACP capability + +The daemon bridge and standalone channel bridge each create a 32-byte random +base64url nonce for a private ACP child. The parent forces the nonce into the +child environment as `QWEN_CODE_PRIVATE_ACP_CAPABILITY` and sends it in the +first `InitializeRequest._meta` under +`qwen-code/private-parent-capability`. + +The ACP entry point captures and deletes the environment capability before +settings parsing, application initialization, or MCP discovery. A regular or +sandbox relaunch passes it only in the final child's spawn environment and +never restores it in the parent. The final process passes the captured value +in memory to the ACP agent. The child decides trust exactly once on the first +initialize request: + +- no capability in the environment means the connection remains untrusted, + even if initialize or prompt metadata contains reserved keys; +- an environment capability with a missing or different initialize value + fails initialization rather than downgrading silently; +- the environment value is never visible to MCP servers or later subprocesses; + later initialize calls cannot upgrade either a trusted, untrusted, or + rejected connection. + +A trusted daemon parent may send a complete context. A trusted standalone +channel parent sends only advisory ingress, because the child owns its native +session and prompt IDs. An untrusted ACP connection strips all reserved prompt +metadata and creates an `acp` root. + +This prevents request-level spoofing by ordinary or remote ACP peers. It does +not attempt to defend against a hostile process running as the same OS user +that deliberately emulates the private launcher; that user already controls +the child process and local MCP processes. + +## MCP disclosure policy + +Disclosure is based on transport ownership, not on a user-configurable server +type or URL. The internal allow bit is set only where Qwen Code constructs a +`StdioClientTransport`. Streamable HTTP, SSE, WebSocket, SDK/reverse transport, +and unknown transports deny disclosure, including loopback URLs. + +The policy follows a discovered tool through qualification and trust wrappers. +Reconnect takes the policy from the newly discovered tool; a missing policy is +deny. It is used only by the raw `DiscoveredMCPToolInvocation` path. +Computer Use, A2UI, and callable-tool fallbacks do not inject v1 metadata. + +For an allowed call, Qwen Code merges the context into +`CallToolRequest.params._meta['qwen-code/invocation']`. Existing request +metadata such as `progressToken` remains intact. Tool arguments, including an +`arguments._meta` property generated by a model, cannot overwrite request +metadata. When disclosure is denied or no runtime context exists, only the +reserved invocation property is removed. + +## qwen-serve-mcp bindings + +qwen-serve-mcp keeps one immutable binding record per attached session. The +record contains the session ID, optional legacy client ID, SSE stream, and its +single release promise. Changing the default session changes only a pointer; +other attached sessions stay live. + +Load, resume, and close operations are serialized per known session. Create is +committed under the returned session ID after the server responds. A binding +with an active prompt collector cannot be replaced. A successful replacement +installs the new record before releasing one acquisition of the old record; +failure preserves the old record. + +Asynchronous stream finalizers and error handlers may remove state only when +their captured record is still current. Prompt cancellation captures the same +record as the prompt and does not replay on an invalid client ID. Idle cleanup +skips active collectors. + +Release performs at most one detach attempt per acquisition. Pending detach +promises remain tracked after a record leaves the map. Session lifecycle +operations, including the acquisition phase before create has a session ID, +are tracked as well. The exported idempotent `dispose()` rejects new lifecycle +work, waits for in-flight lifecycle work to settle, releases any resulting +current records, and then waits for all pending releases. + +An active prompt handler is tracked through its abort cancellation and +collector cleanup. Shutdown drains that handler before releasing its binding, +so detach cannot overtake cancellation and leave the daemon turn running. + +For an accepted asynchronous daemon prompt, `DaemonClient` owns cancellation +until prompt admission has settled and awaits that cancel before returning an +abort. The bridge owns cancellation only after admission has settled while it +is still collecting events. This single-owner boundary prevents a delayed +session-level cancel from reaching the next turn. MCP close, SIGINT, SIGTERM, +and stdin close share one shutdown promise. + +## Compatibility and exclusions + +There is no feature flag or persisted migration. Older MCP servers ignore the +namespaced request metadata, and older daemons continue to use the anonymous +client path. v1 does not change opencode, remote MCP transports, Computer Use, +A2UI, or the durable prompt history schema. + +Normative examples, including rejected payloads, are in +`v1-conformance.json`. diff --git a/docs/design/invocation-context/v1-conformance.json b/docs/design/invocation-context/v1-conformance.json new file mode 100644 index 00000000000..d9d6d02c3cc --- /dev/null +++ b/docs/design/invocation-context/v1-conformance.json @@ -0,0 +1,125 @@ +{ + "schemaVersion": 1, + "metaKeys": { + "invocation": "qwen-code/invocation", + "ingress": "qwen-code/invocation-ingress", + "privateParentCapability": "qwen-code/private-parent-capability" + }, + "contextCases": [ + { + "name": "valid-with-originator", + "valid": true, + "input": { + "version": 1, + "ingress": "external_mcp", + "sessionId": "session-1", + "promptId": "prompt-1", + "originatorClientId": "client-1" + } + }, + { + "name": "valid-without-originator", + "valid": true, + "input": { + "version": 1, + "ingress": "cli", + "sessionId": "session-2", + "promptId": "prompt-2" + } + }, + { + "name": "missing-context", + "valid": false + }, + { + "name": "future-version", + "valid": false, + "input": { + "version": 2, + "ingress": "cli", + "sessionId": "session-3", + "promptId": "prompt-3" + } + }, + { + "name": "unknown-ingress", + "valid": false, + "input": { + "version": 1, + "ingress": "webhook", + "sessionId": "session-4", + "promptId": "prompt-4" + } + }, + { + "name": "blank-session-id", + "valid": false, + "input": { + "version": 1, + "ingress": "daemon", + "sessionId": " ", + "promptId": "prompt-5" + } + }, + { + "name": "empty-prompt-id", + "valid": false, + "input": { + "version": 1, + "ingress": "daemon", + "sessionId": "session-6", + "promptId": "" + } + }, + { + "name": "blank-originator", + "valid": false, + "input": { + "version": 1, + "ingress": "daemon", + "sessionId": "session-7", + "promptId": "prompt-7", + "originatorClientId": " " + } + }, + { + "name": "extra-field", + "valid": false, + "input": { + "version": 1, + "ingress": "daemon", + "sessionId": "session-8", + "promptId": "prompt-8", + "trusted": true + } + } + ], + "metadataCases": [ + { + "name": "legacy-metadata-coexists", + "existing": { + "progressToken": "progress-1", + "vendor.example/trace": "legacy-trace" + }, + "expectPreserved": ["progressToken", "vendor.example/trace"], + "expectInvocation": true + }, + { + "name": "denied-transport-removes-only-reserved-key", + "existing": { + "progressToken": "progress-2", + "qwen-code/invocation": { + "forged": true + } + }, + "expectPreserved": ["progressToken"], + "expectInvocation": false + } + ], + "consumerBehavior": { + "missing": "omit-reserved-metadata", + "malformedTrustedPrivateAcp": "fail-prompt", + "malformedUntrustedAcp": "strip-and-create-acp-root", + "unknownDaemonIngress": "normalize-to-daemon" + } +} diff --git a/integration-tests/cli/invocation-context-v1.test.ts b/integration-tests/cli/invocation-context-v1.test.ts new file mode 100644 index 00000000000..0d7063cc29d --- /dev/null +++ b/integration-tests/cli/invocation-context-v1.test.ts @@ -0,0 +1,464 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { spawn } from 'node:child_process'; +import { randomUUID } from 'node:crypto'; +import { mkdirSync, mkdtempSync, readFileSync, rmSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import * as path from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'; +import { describe, expect, it } from 'vitest'; +import { fakeToolCall, startFakeOpenAIServer } from '../fake-openai-server.js'; +import { + approveWorkspaceMcpServers, + spawnDaemon, + type SpawnedDaemon, + writeWorkspaceSettings, +} from './_daemon-harness.js'; + +const __dirname = path.dirname(fileURLToPath(import.meta.url)); +const REPO_ROOT = path.resolve(__dirname, '../..'); +const CLI_BIN = + process.env['TEST_CLI_PATH'] ?? path.join(REPO_ROOT, 'dist', 'cli.js'); +const ECHO_SERVER = path.join( + REPO_ROOT, + 'integration-tests', + 'fixtures', + 'invocation-context-echo.mjs', +); +const SERVE_BRIDGE_BIN = path.join( + REPO_ROOT, + 'packages', + 'sdk-typescript', + 'dist', + 'daemon-mcp', + 'serve-bridge', + 'bin.js', +); +const ECHO_TOOL = 'mcp__invocation-echo__capture_invocation_context'; +const INVOCATION_META_KEY = 'qwen-code/invocation'; +const DIRECT_CLI_SENTINEL = 'CAPTURE_DIRECT_CLI_INVOCATION_CONTEXT'; +const DAEMON_A_SENTINEL = 'CAPTURE_DAEMON_INVOCATION_CONTEXT_A'; +const DAEMON_B_SENTINEL = 'CAPTURE_DAEMON_INVOCATION_CONTEXT_B'; +const DAEMON_B_AFTER_CLOSE_SENTINEL = + 'CAPTURE_DAEMON_INVOCATION_CONTEXT_B_AFTER_A_CLOSE'; +const FINAL_ASSISTANT_TEXT = 'INVOCATION_CONTEXT_CAPTURED'; +const UUID_PATTERN = + /^[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i; + +type CaptureRecord = { + arguments: Record; + metadata: Record | null; + privateCapabilityInEnv: boolean; +}; + +type BridgeSession = { + sessionId: string; + clientId?: string; +}; + +function runProcess( + command: string, + args: string[], + options: { cwd: string; env: NodeJS.ProcessEnv }, +): Promise<{ stdout: string; stderr: string }> { + return new Promise((resolve, reject) => { + const child = spawn(command, args, { + cwd: options.cwd, + env: options.env, + stdio: ['ignore', 'pipe', 'pipe'], + }); + let stdout = ''; + let stderr = ''; + child.stdout.on('data', (chunk: Buffer) => { + stdout += chunk.toString(); + }); + child.stderr.on('data', (chunk: Buffer) => { + stderr += chunk.toString(); + }); + let timedOut = false; + let forceKillDeadline: NodeJS.Timeout | undefined; + const deadline = setTimeout(() => { + timedOut = true; + child.kill('SIGTERM'); + forceKillDeadline = setTimeout(() => child.kill('SIGKILL'), 5_000); + forceKillDeadline.unref(); + }, 45_000); + deadline.unref(); + child.once('error', (error) => { + clearTimeout(deadline); + if (forceKillDeadline) clearTimeout(forceKillDeadline); + reject(error); + }); + child.once('exit', (code, signal) => { + clearTimeout(deadline); + if (forceKillDeadline) clearTimeout(forceKillDeadline); + if (code === 0 && !timedOut) { + resolve({ stdout, stderr }); + return; + } + reject( + new Error( + `${command} ${timedOut ? 'timed out' : 'exited'} with code=${String(code)} signal=${String(signal)}\n` + + `stdout=${stdout}\nstderr=${stderr}`, + ), + ); + }); + }); +} + +function readCaptures(file: string): CaptureRecord[] { + return readFileSync(file, 'utf8') + .trim() + .split('\n') + .filter(Boolean) + .map((line) => JSON.parse(line) as CaptureRecord); +} + +function cleanupCaptureRoot(root: string, captureFile: string): void { + if (process.env['KEEP_INVOCATION_CONTEXT_CAPTURE'] === '1') { + process.stderr.write(`Invocation context capture: ${captureFile}\n`); + return; + } + rmSync(root, { recursive: true, force: true }); +} + +function invocationFrom(record: CaptureRecord): Record { + const invocation = record.metadata?.[INVOCATION_META_KEY]; + expect(invocation).toBeTypeOf('object'); + return invocation as Record; +} + +function stringEnv(env: NodeJS.ProcessEnv): Record { + return Object.fromEntries( + Object.entries(env).filter( + (entry): entry is [string, string] => typeof entry[1] === 'string', + ), + ); +} + +function pendingSentinel( + body: Record, + sentinels: string[], +): string | undefined { + const messages = body['messages']; + if (!Array.isArray(messages)) { + return undefined; + } + + const userIndex = messages.findLastIndex( + (message) => + typeof message === 'object' && + message !== null && + (message as Record)['role'] === 'user', + ); + if (userIndex === -1) { + return undefined; + } + + const userMessage = JSON.stringify(messages[userIndex]); + const sentinel = sentinels.find((candidate) => + userMessage.includes(candidate), + ); + if (!sentinel) { + return undefined; + } + + const hasToolResult = messages + .slice(userIndex + 1) + .some( + (message) => + typeof message === 'object' && + message !== null && + (message as Record)['role'] === 'tool', + ); + return hasToolResult ? undefined : sentinel; +} + +async function callBridgeTool( + client: Client, + name: string, + args: Record, +): Promise { + const result = await client.callTool({ name, arguments: args }); + const content = result.content as Array<{ type?: string; text?: string }>; + const text = content.find((item) => item.type === 'text')?.text; + if (!text) { + throw new Error(`${name} returned no text: ${JSON.stringify(result)}`); + } + if (result.isError) { + throw new Error(`${name} failed: ${text}`); + } + return JSON.parse(text) as T; +} + +describe('Invocation Context v1 runtime propagation', () => { + it('marks a bundled direct CLI turn as cli', async () => { + const root = mkdtempSync(path.join(tmpdir(), 'qwen-invocation-cli-')); + const workspace = path.join(root, 'workspace'); + const home = path.join(root, 'home'); + const captureFile = path.join(root, 'capture.jsonl'); + mkdirSync(workspace, { recursive: true }); + mkdirSync(home, { recursive: true }); + + const fakeModel = await startFakeOpenAIServer(({ body }) => { + const sentinel = pendingSentinel(body, [DIRECT_CLI_SENTINEL]); + if (sentinel) { + return { + toolCalls: [fakeToolCall(ECHO_TOOL, { probe: sentinel })], + }; + } + return { content: FINAL_ASSISTANT_TEXT }; + }); + + try { + const sessionId = randomUUID(); + const mcpConfig = JSON.stringify({ + mcpServers: { + 'invocation-echo': { + command: process.execPath, + args: [ECHO_SERVER], + env: { INVOCATION_CONTEXT_ECHO_FILE: captureFile }, + }, + }, + }); + const result = await runProcess( + process.execPath, + [ + CLI_BIN, + '--prompt', + DIRECT_CLI_SENTINEL, + '--session-id', + sessionId, + '--mcp-config', + mcpConfig, + '--allowed-mcp-server-names', + 'invocation-echo', + '--allowed-tools', + ECHO_TOOL, + '--approval-mode', + 'yolo', + '--auth-type', + 'openai', + '--model', + 'fake-model', + '--output-format', + 'json', + '--no-chat-recording', + ], + { + cwd: workspace, + env: { + ...process.env, + HOME: home, + QWEN_HOME: path.join(home, '.qwen'), + QWEN_SANDBOX: 'false', + QWEN_CODE_NO_RELAUNCH: 'true', + QWEN_CODE_LEGACY_MCP_BLOCKING: '1', + OPENAI_API_KEY: 'fake-key', + OPENAI_BASE_URL: fakeModel.baseUrl, + OPENAI_MODEL: 'fake-model', + QWEN_MODEL: 'fake-model', + NO_PROXY: '127.0.0.1,localhost', + no_proxy: '127.0.0.1,localhost', + }, + }, + ); + + expect(result.stdout).toContain(FINAL_ASSISTANT_TEXT); + const captures = readCaptures(captureFile); + expect(captures).toHaveLength(1); + expect(captures[0]?.privateCapabilityInEnv).toBe(false); + const invocation = invocationFrom(captures[0]!); + expect(invocation).toEqual({ + version: 1, + ingress: 'cli', + sessionId, + promptId: expect.any(String), + }); + expect(String(invocation['promptId']).trim()).not.toBe(''); + } finally { + await fakeModel.close(); + cleanupCaptureRoot(root, captureFile); + } + }, 60_000); + + it('keeps external_mcp session markings isolated through the daemon', async () => { + const root = mkdtempSync(path.join(tmpdir(), 'qwen-invocation-daemon-')); + const workspace = path.join(root, 'workspace'); + const home = path.join(root, 'home'); + const captureFile = path.join(root, 'capture.jsonl'); + mkdirSync(workspace, { recursive: true }); + mkdirSync(home, { recursive: true }); + + const sentinels = [ + DAEMON_A_SENTINEL, + DAEMON_B_AFTER_CLOSE_SENTINEL, + DAEMON_B_SENTINEL, + ]; + const fakeModel = await startFakeOpenAIServer(({ body }) => { + const sentinel = pendingSentinel(body, sentinels); + if (sentinel) { + return { + toolCalls: [fakeToolCall(ECHO_TOOL, { probe: sentinel })], + }; + } + return { content: FINAL_ASSISTANT_TEXT }; + }); + + let daemon: SpawnedDaemon | undefined; + let bridgeClient: Client | undefined; + try { + const mcpServers = { + 'invocation-echo': { + command: process.execPath, + args: [ECHO_SERVER], + env: { INVOCATION_CONTEXT_ECHO_FILE: captureFile }, + trust: true, + alwaysLoadTools: true, + }, + }; + writeWorkspaceSettings(workspace, { + tools: { approvalMode: 'yolo' }, + mcpServers, + }); + const approvalEnv = approveWorkspaceMcpServers(workspace, mcpServers); + const runtimeEnv = { + ...process.env, + ...approvalEnv, + HOME: home, + QWEN_HOME: path.join(home, '.qwen'), + QWEN_SANDBOX: 'false', + QWEN_CODE_NO_RELAUNCH: 'true', + QWEN_CODE_LEGACY_MCP_BLOCKING: '1', + QWEN_CODE_SUPPRESS_YOLO_WARNING: '1', + OPENAI_API_KEY: 'fake-key', + OPENAI_BASE_URL: fakeModel.baseUrl, + OPENAI_MODEL: 'fake-model', + QWEN_MODEL: 'fake-model', + NO_PROXY: '127.0.0.1,localhost', + no_proxy: '127.0.0.1,localhost', + }; + daemon = await spawnDaemon({ + workspaceCwd: workspace, + bootTimeoutMs: 20_000, + env: runtimeEnv, + }); + + bridgeClient = new Client({ + name: 'invocation-context-runtime-test', + version: '1.0.0', + }); + const bridgeTransport = new StdioClientTransport({ + command: process.execPath, + args: [SERVE_BRIDGE_BIN], + env: stringEnv({ + ...runtimeEnv, + QWEN_DAEMON_URL: daemon.base, + QWEN_DAEMON_TOKEN: daemon.token, + QWEN_WORKSPACE_CWD: workspace, + }), + }); + await bridgeClient.connect(bridgeTransport); + + const [sessionA, sessionB] = await Promise.all([ + callBridgeTool(bridgeClient, 'session_create', { + workspace_cwd: workspace, + session_scope: 'thread', + }), + callBridgeTool(bridgeClient, 'session_create', { + workspace_cwd: workspace, + session_scope: 'thread', + }), + ]); + expect(sessionA.sessionId).not.toBe(sessionB.sessionId); + expect(sessionA.clientId).toBeTypeOf('string'); + expect(sessionB.clientId).toBeTypeOf('string'); + + await Promise.all([ + callBridgeTool(bridgeClient, 'prompt', { + session_id: sessionA.sessionId, + prompt: DAEMON_A_SENTINEL, + }), + callBridgeTool(bridgeClient, 'prompt', { + session_id: sessionB.sessionId, + prompt: DAEMON_B_SENTINEL, + }), + ]); + + const firstCaptures = readCaptures(captureFile); + expect(firstCaptures).toHaveLength(2); + const captureA = firstCaptures.find( + (record) => record.arguments['probe'] === DAEMON_A_SENTINEL, + ); + const captureB = firstCaptures.find( + (record) => record.arguments['probe'] === DAEMON_B_SENTINEL, + ); + expect(captureA).toBeDefined(); + expect(captureB).toBeDefined(); + expect(captureA?.privateCapabilityInEnv).toBe(false); + expect(captureB?.privateCapabilityInEnv).toBe(false); + + const invocationA = invocationFrom(captureA!); + const invocationB = invocationFrom(captureB!); + expect(invocationA).toEqual({ + version: 1, + ingress: 'external_mcp', + sessionId: sessionA.sessionId, + promptId: expect.stringMatching(UUID_PATTERN), + originatorClientId: sessionA.clientId, + }); + expect(invocationB).toEqual({ + version: 1, + ingress: 'external_mcp', + sessionId: sessionB.sessionId, + promptId: expect.stringMatching(UUID_PATTERN), + originatorClientId: sessionB.clientId, + }); + expect(invocationA['promptId']).not.toBe(invocationB['promptId']); + + await callBridgeTool(bridgeClient, 'session_close', { + session_id: sessionA.sessionId, + }); + await callBridgeTool(bridgeClient, 'session_context', { + session_id: sessionB.sessionId, + }); + await callBridgeTool(bridgeClient, 'prompt', { + session_id: sessionB.sessionId, + prompt: DAEMON_B_AFTER_CLOSE_SENTINEL, + }); + + const allCaptures = readCaptures(captureFile); + expect(allCaptures).toHaveLength(3); + const captureBAfterClose = allCaptures.find( + (record) => record.arguments['probe'] === DAEMON_B_AFTER_CLOSE_SENTINEL, + ); + expect(captureBAfterClose).toBeDefined(); + const invocationBAfterClose = invocationFrom(captureBAfterClose!); + expect(invocationBAfterClose).toEqual({ + version: 1, + ingress: 'external_mcp', + sessionId: sessionB.sessionId, + promptId: expect.stringMatching(UUID_PATTERN), + originatorClientId: sessionB.clientId, + }); + expect(invocationBAfterClose['promptId']).not.toBe( + invocationB['promptId'], + ); + + await callBridgeTool(bridgeClient, 'session_close', { + session_id: sessionB.sessionId, + }); + } finally { + await bridgeClient?.close().catch(() => undefined); + await daemon?.dispose(); + await fakeModel.close(); + cleanupCaptureRoot(root, captureFile); + } + }, 120_000); +}); diff --git a/integration-tests/fixtures/invocation-context-echo.mjs b/integration-tests/fixtures/invocation-context-echo.mjs new file mode 100644 index 00000000000..8ae6c499338 --- /dev/null +++ b/integration-tests/fixtures/invocation-context-echo.mjs @@ -0,0 +1,49 @@ +import { appendFile } from 'node:fs/promises'; +import process from 'node:process'; + +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; +import { + CallToolRequestSchema, + ListToolsRequestSchema, +} from '@modelcontextprotocol/sdk/types.js'; + +const outputFile = process.env['INVOCATION_CONTEXT_ECHO_FILE']; +if (!outputFile) { + throw new Error('INVOCATION_CONTEXT_ECHO_FILE is required'); +} + +const server = new Server( + { name: 'invocation-context-echo', version: '1.0.0' }, + { capabilities: { tools: {} } }, +); + +server.setRequestHandler(ListToolsRequestSchema, async () => ({ + tools: [ + { + name: 'capture_invocation_context', + description: 'Capture this tool request metadata.', + inputSchema: { + type: 'object', + properties: { probe: { type: 'string' } }, + required: ['probe'], + additionalProperties: false, + }, + }, + ], +})); + +server.setRequestHandler(CallToolRequestSchema, async (request) => { + const record = { + arguments: request.params.arguments ?? {}, + metadata: request.params._meta ?? null, + privateCapabilityInEnv: + process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY'] !== undefined, + }; + await appendFile(outputFile, `${JSON.stringify(record)}\n`); + return { + content: [{ type: 'text', text: JSON.stringify(record) }], + }; +}); + +await server.connect(new StdioServerTransport()); diff --git a/packages/acp-bridge/src/bridge.test.ts b/packages/acp-bridge/src/bridge.test.ts index 7437c7720e0..cb95605fbd5 100644 --- a/packages/acp-bridge/src/bridge.test.ts +++ b/packages/acp-bridge/src/bridge.test.ts @@ -467,17 +467,21 @@ describe('createAcpSessionBridge', () => { // inside each bridge — so each bridge's spawn factory sees // ITS own overrides, regardless of what other daemons did. const seenEnvs: Array | undefined> = []; + const handles: ChannelHandle[] = []; const factory: ChannelFactory = async (_cwd, env) => { // Snapshot the override map so later iterations don't // accidentally mutate the recorded value. seenEnvs.push(env ? { ...env } : env); - return makeChannel().channel; + const handle = makeChannel(); + handles.push(handle); + return handle.channel; }; const bridge1 = makeBridge({ channelFactory: factory, childEnvOverrides: { QWEN_SERVE_MCP_CLIENT_BUDGET: '5', QWEN_SERVE_MCP_BUDGET_MODE: 'enforce', + QWEN_CODE_PRIVATE_ACP_CAPABILITY: 'caller-must-not-win', }, }); const bridge2 = makeBridge({ @@ -490,14 +494,31 @@ describe('createAcpSessionBridge', () => { await bridge1.spawnOrAttach({ workspaceCwd: WS_A }); await bridge2.spawnOrAttach({ workspaceCwd: WS_A }); expect(seenEnvs).toHaveLength(2); - expect(seenEnvs[0]).toEqual({ + expect(seenEnvs[0]).toMatchObject({ QWEN_SERVE_MCP_CLIENT_BUDGET: '5', QWEN_SERVE_MCP_BUDGET_MODE: 'enforce', }); - expect(seenEnvs[1]).toEqual({ + expect(seenEnvs[1]).toMatchObject({ QWEN_SERVE_MCP_CLIENT_BUDGET: '20', QWEN_SERVE_MCP_BUDGET_MODE: 'warn', }); + const capabilities = seenEnvs.map( + (env) => env?.['QWEN_CODE_PRIVATE_ACP_CAPABILITY'], + ); + expect(capabilities[0]).toMatch(/^[A-Za-z0-9_-]{43}$/u); + expect(capabilities[0]).not.toBe('caller-must-not-win'); + expect(capabilities[1]).toMatch(/^[A-Za-z0-9_-]{43}$/u); + expect(capabilities[1]).not.toBe(capabilities[0]); + expect( + handles[0]?.agent.initializeCalls[0]?._meta?.[ + 'qwen-code/private-parent-capability' + ], + ).toBe(capabilities[0]); + expect( + handles[1]?.agent.initializeCalls[0]?._meta?.[ + 'qwen-code/private-parent-capability' + ], + ).toBe(capabilities[1]); await bridge1.shutdown(); await bridge2.shutdown(); }); @@ -4100,6 +4121,88 @@ describe('createAcpSessionBridge', () => { await bridge.shutdown(); }); + it('validates the client before rebuilding authoritative invocation metadata', async () => { + const handle = makeChannel(); + const bridge = makeBridge({ channelFactory: async () => handle.channel }); + const session = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + const forgedInvocation = { + version: 1, + ingress: 'internal', + sessionId: 'forged-session', + promptId: 'forged-prompt', + originatorClientId: 'forged-client', + }; + + expect(() => + bridge.sendPrompt( + session.sessionId, + { + sessionId: session.sessionId, + prompt: [{ type: 'text', text: 'rejected' }], + _meta: { 'qwen-code/invocation': forgedInvocation }, + }, + undefined, + { clientId: 'invalid-client', promptId: 'rejected-prompt' }, + ), + ).toThrow(InvalidClientIdError); + expect(handle.agent.promptCalls).toHaveLength(0); + + const cases = [ + ['external_mcp', 'external_mcp', session.clientId], + ['channel', 'channel', undefined], + ['scheduler', 'scheduler', undefined], + ['internal', 'internal', undefined], + ['unknown-ingress', 'daemon', session.clientId], + ] as const; + for (const [advisory, ingress, originatorClientId] of cases) { + const promptId = `server-${advisory}`; + await bridge.sendPrompt( + session.sessionId, + { + sessionId: session.sessionId, + prompt: [{ type: 'text', text: advisory }], + _meta: { + keep: true, + 'qwen-code/invocation': forgedInvocation, + 'qwen-code/invocation-ingress': advisory, + 'qwen-code/private-parent-capability': 'forged-capability', + }, + }, + undefined, + { clientId: session.clientId, promptId }, + ); + expect(handle.agent.promptCalls.at(-1)?._meta).toMatchObject({ + keep: true, + 'qwen-code/invocation': { + version: 1, + ingress, + sessionId: session.sessionId, + promptId, + ...(originatorClientId ? { originatorClientId } : {}), + }, + }); + expect( + handle.agent.promptCalls.at(-1)?._meta?.[ + 'qwen-code/invocation-ingress' + ], + ).toBeUndefined(); + expect( + handle.agent.promptCalls.at(-1)?._meta?.[ + 'qwen-code/private-parent-capability' + ], + ).toBeUndefined(); + expect( + ( + handle.agent.promptCalls.at(-1)?._meta?.[ + 'qwen-code/invocation' + ] as { originatorClientId?: string } + ).originatorClientId, + ).toBe(originatorClientId); + } + + await bridge.shutdown(); + }); + it('ignores client retry when no turn_error made the session retryable', async () => { const handle = makeChannel(); const bridge = makeBridge({ channelFactory: async () => handle.channel }); diff --git a/packages/acp-bridge/src/bridge.ts b/packages/acp-bridge/src/bridge.ts index 71670d3a8e5..7e86b2e910d 100644 --- a/packages/acp-bridge/src/bridge.ts +++ b/packages/acp-bridge/src/bridge.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { randomUUID } from 'node:crypto'; +import { randomBytes, randomUUID } from 'node:crypto'; import * as path from 'node:path'; import { ClientSideConnection, @@ -24,10 +24,15 @@ import type { import { DAEMON_TRACEPARENT_META_KEY, DAEMON_TRACESTATE_META_KEY, + INVOCATION_CONTEXT_META_KEY, + INVOCATION_INGRESS_META_KEY, + PRIVATE_PARENT_CAPABILITY_META_KEY, SESSION_ARTIFACT_PERSISTENCE_VERSION, TrustGateError, normalizeSnapshotPayload, + runWithInvocationContext, ShellExecutionService, + type InvocationContextV1, type ShellOutputEvent, } from '@qwen-code/qwen-code-core'; import type { ShellCommandResult } from './bridgeTypes.js'; @@ -138,6 +143,8 @@ import { type SessionArtifactMutationResult, } from './sessionArtifacts.js'; +const PRIVATE_ACP_CAPABILITY_ENV = 'QWEN_CODE_PRIVATE_ACP_CAPABILITY'; + const NOOP_BRIDGE_TELEMETRY: BridgeTelemetry = { captureContext: () => undefined, runWithContext(_captured, fn) { @@ -1754,13 +1761,18 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { if (inFlightChannelSpawn) return await inFlightChannelSpawn; const promise = (async () => { + const privateParentCapability = randomBytes(32).toString('base64url'); const channel = await telemetry.withSpan( 'channel.spawn', { 'qwen-code.daemon.bridge.operation': 'channel.spawn', 'qwen-code.daemon.channel.reused': false, }, - async () => await channelFactory(boundWorkspace, childEnvOverrides), + async () => + await channelFactory(boundWorkspace, { + ...childEnvOverrides, + [PRIVATE_ACP_CAPABILITY_ENV]: privateParentCapability, + }), ); const sessionIds = new Set(); const client = new BridgeClient( @@ -1994,6 +2006,9 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { fs: { readTextFile: true, writeTextFile: true }, }, clientInfo: { name: 'qwen-serve-bridge', version: '0' }, + _meta: { + [PRIVATE_PARENT_CAPABILITY_META_KEY]: privateParentCapability, + }, }), initTimeoutMs, 'initialize', @@ -4260,6 +4275,26 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { // the first prompt on an idle session starts immediately and // doesn't need a queue event. const promptId = context?.promptId ?? randomUUID(); + const advisoryIngress = (() => { + const value = req._meta?.[INVOCATION_INGRESS_META_KEY]; + return value === 'channel' || + value === 'scheduler' || + value === 'internal' || + value === 'external_mcp' + ? value + : 'daemon'; + })(); + const invocationContext: InvocationContextV1 = { + version: 1, + ingress: advisoryIngress, + sessionId, + promptId, + ...((advisoryIngress === 'daemon' || + advisoryIngress === 'external_mcp') && + originatorClientId + ? { originatorClientId } + : {}), + }; const isQueued = entry.pendingPromptCount > 1; const pendingAbort = new AbortController(); if (signal) { @@ -4297,7 +4332,7 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { // Force the body's sessionId to match the routing id — a client that // sent a stale id in the body would otherwise be dispatched to the // wrong agent process. - const result = entry.promptQueue.then(() => + const executePrompt = () => telemetry.runWithContext(capturedContext, async () => { const queueWaitMs = Date.now() - queuedAt; telemetry.metrics?.promptQueueWait(queueWaitMs); @@ -4363,6 +4398,9 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { ? { ...copy._meta } : {}; delete meta[DAEMON_RETRY_META_KEY]; + delete meta[INVOCATION_CONTEXT_META_KEY]; + delete meta[INVOCATION_INGRESS_META_KEY]; + delete meta[PRIVATE_PARENT_CAPABILITY_META_KEY]; // External prompt callers cannot self-trigger a continuation; // only `continueSession` (via the trusted `isContinue` flag // below) re-arms it after this strip. @@ -4373,6 +4411,7 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { if (isContinue) { meta[DAEMON_CONTINUE_META_KEY] = true; } + meta[INVOCATION_CONTEXT_META_KEY] = invocationContext; if (Object.keys(meta).length > 0) { copy._meta = meta; } else { @@ -4556,7 +4595,9 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { } finally { telemetry.metrics?.promptDuration(Date.now() - dispatchStartMs); } - }), + }); + const result = entry.promptQueue.then(() => + runWithInvocationContext(invocationContext, executePrompt), ); result.then( (promptResult) => { diff --git a/packages/acp-bridge/src/internal/testUtils.ts b/packages/acp-bridge/src/internal/testUtils.ts index c11555fd754..bfc8619ed24 100644 --- a/packages/acp-bridge/src/internal/testUtils.ts +++ b/packages/acp-bridge/src/internal/testUtils.ts @@ -141,6 +141,7 @@ export interface FakeAgentOpts { } export class FakeAgent implements Agent { + initializeCalls: InitializeRequest[] = []; newSessionCalls: NewSessionRequest[] = []; loadSessionCalls: LoadSessionRequest[] = []; resumeSessionCalls: ResumeSessionRequest[] = []; @@ -150,7 +151,8 @@ export class FakeAgent implements Agent { []; constructor(private readonly opts: FakeAgentOpts = {}) {} - async initialize(_p: InitializeRequest): Promise { + async initialize(p: InitializeRequest): Promise { + this.initializeCalls.push(p); if (this.opts.initializeThrows) throw this.opts.initializeThrows; if (this.opts.initializeDelayMs) { await new Promise((r) => setTimeout(r, this.opts.initializeDelayMs)); diff --git a/packages/channels/base/src/AcpBridge.test.ts b/packages/channels/base/src/AcpBridge.test.ts index 8cd7cee64b3..ab5e768e884 100644 --- a/packages/channels/base/src/AcpBridge.test.ts +++ b/packages/channels/base/src/AcpBridge.test.ts @@ -57,7 +57,7 @@ const child = vi.hoisted(() => { cancel: ReturnType; }>, MockChild, - spawn: vi.fn(() => { + spawn: vi.fn((_command?: unknown, _args?: unknown, _options?: unknown) => { const instance = new MockChild(); child.instances.push(instance); return instance; @@ -580,6 +580,32 @@ describe('AcpBridge', () => { expect(disconnected).toHaveBeenCalledWith(null, 'SIGKILL'); }); + it('uses the same fresh private capability for child env and initialize', async () => { + vi.useFakeTimers(); + try { + const bridge = new AcpBridge({ + cliEntryPath: '/tmp/qwen', + cwd: '/tmp', + }); + const started = bridge.start(); + await vi.advanceTimersByTimeAsync(1000); + await started; + + const env = child.spawn.mock.calls[0]?.[2]?.env as + | Record + | undefined; + const capability = env?.['QWEN_CODE_PRIVATE_ACP_CAPABILITY']; + expect(capability).toMatch(/^[A-Za-z0-9_-]{43}$/u); + expect( + child.connections[0]?.initialize.mock.calls[0]?.[0]?._meta?.[ + 'qwen-code/private-parent-capability' + ], + ).toBe(capability); + } finally { + vi.useRealTimers(); + } + }); + it('kills the ACP child when a stall line is coalesced with prior stderr', async () => { const bridge = new AcpBridge({ cliEntryPath: '/tmp/qwen', diff --git a/packages/channels/base/src/AcpBridge.ts b/packages/channels/base/src/AcpBridge.ts index aaf530ef6d2..e6c49415a42 100644 --- a/packages/channels/base/src/AcpBridge.ts +++ b/packages/channels/base/src/AcpBridge.ts @@ -1,6 +1,6 @@ import { spawn } from 'node:child_process'; import type { ChildProcess } from 'node:child_process'; -import { randomUUID } from 'node:crypto'; +import { randomBytes, randomUUID } from 'node:crypto'; import { Readable, Writable } from 'node:stream'; import { EventEmitter } from 'node:events'; import { @@ -18,6 +18,7 @@ import type { AvailableCommand, ChannelLoopToolHandler, ChannelAgentBridge, + ChannelPromptOptions, ToolCallEvent, } from './ChannelAgentBridge.js'; import { @@ -32,6 +33,9 @@ import { sanitizeLogText } from './sanitize.js'; export type { AvailableCommand, ToolCallEvent } from './ChannelAgentBridge.js'; const MID_TURN_QUEUE_DRAIN_METHOD = 'craft/drainMidTurnQueue'; +const PRIVATE_ACP_CAPABILITY_ENV = 'QWEN_CODE_PRIVATE_ACP_CAPABILITY'; +const PRIVATE_PARENT_CAPABILITY_META_KEY = + 'qwen-code/private-parent-capability'; export interface AcpBridgeOptions { cliEntryPath: string; @@ -99,6 +103,7 @@ export class AcpBridge extends EventEmitter implements ChannelAgentBridge { async start(): Promise { const { cliEntryPath, cwd } = this.options; + const privateParentCapability = randomBytes(32).toString('base64url'); const args = [ ...process.execArgv.filter((a) => !/^--inspect(-brk)?($|=)/.test(a)), @@ -112,7 +117,11 @@ export class AcpBridge extends EventEmitter implements ChannelAgentBridge { this.child = spawn(process.execPath, args, { cwd, stdio: ['pipe', 'pipe', 'pipe'], - env: { ...process.env, QWEN_CODE_DISABLE_CRON: '1' }, + env: { + ...process.env, + QWEN_CODE_DISABLE_CRON: '1', + [PRIVATE_ACP_CAPABILITY_ENV]: privateParentCapability, + }, shell: false, }); @@ -174,6 +183,9 @@ export class AcpBridge extends EventEmitter implements ChannelAgentBridge { await this.connection.initialize({ protocolVersion: PROTOCOL_VERSION, clientCapabilities: {}, + _meta: { + [PRIVATE_PARENT_CAPABILITY_META_KEY]: privateParentCapability, + }, }); await this.registerChannelLoopMcpServer(); } @@ -214,7 +226,7 @@ export class AcpBridge extends EventEmitter implements ChannelAgentBridge { async prompt( sessionId: string, text: string, - options?: { imageBase64?: string; imageMimeType?: string }, + options?: ChannelPromptOptions, ): Promise { const conn = this.ensureConnection(); @@ -250,6 +262,13 @@ export class AcpBridge extends EventEmitter implements ChannelAgentBridge { await conn.prompt({ sessionId, prompt: prompt as Array<{ type: 'text'; text: string }>, + ...(options?.invocationIngress + ? { + _meta: { + 'qwen-code/invocation-ingress': options.invocationIngress, + }, + } + : {}), }); } finally { this.off('textChunk', onChunk); diff --git a/packages/channels/base/src/ChannelAgentBridge.ts b/packages/channels/base/src/ChannelAgentBridge.ts index 928e4a7df52..7aedd171fb0 100644 --- a/packages/channels/base/src/ChannelAgentBridge.ts +++ b/packages/channels/base/src/ChannelAgentBridge.ts @@ -82,6 +82,14 @@ export interface ChannelAgentBridgeSessionOptions { approvalMode?: string; } +export type ChannelInvocationIngress = 'channel' | 'scheduler' | 'internal'; + +export interface ChannelPromptOptions { + imageBase64?: string; + imageMimeType?: string; + invocationIngress?: ChannelInvocationIngress; +} + export interface ChannelAgentBridge { readonly availableCommands: AvailableCommand[]; getAvailableCommands?(sessionId: string): AvailableCommand[]; @@ -107,7 +115,7 @@ export interface ChannelAgentBridge { prompt( sessionId: string, text: string, - options?: { imageBase64?: string; imageMimeType?: string }, + options?: ChannelPromptOptions, ): Promise; cancelSession(sessionId: string): Promise; /** Release a bridge-owned session that will not be routed to a caller. */ diff --git a/packages/channels/base/src/ChannelBase.test.ts b/packages/channels/base/src/ChannelBase.test.ts index 7fdbcb2384d..b91550ace56 100644 --- a/packages/channels/base/src/ChannelBase.test.ts +++ b/packages/channels/base/src/ChannelBase.test.ts @@ -3821,7 +3821,11 @@ describe('ChannelBase', () => { await ch.handleInbound(envelope({ text: '/schedule list' })); - expect(bridge.prompt).toHaveBeenCalledWith('s-1', '/schedule list', {}); + expect(bridge.prompt).toHaveBeenCalledWith('s-1', '/schedule list', { + imageBase64: undefined, + imageMimeType: undefined, + invocationIngress: 'channel', + }); expect(ch.sent).toEqual([{ chatId: 'chat1', text: 'agent response' }]); }); @@ -8660,6 +8664,7 @@ describe('ChannelBase', () => { expect(bridge.prompt).toHaveBeenCalledWith('s-1', '!echo hello', { imageBase64: undefined, imageMimeType: undefined, + invocationIngress: 'channel', }); }); @@ -10685,7 +10690,7 @@ describe('ChannelBase', () => { expect.stringContaining( '[External event "ci_failed" from github-ci]', ), - {}, + { invocationIngress: 'scheduler' }, ); expect(ch.proactive).toEqual([ { chatId: 'group-1', text: 'CI failed because lint broke.' }, @@ -11262,7 +11267,7 @@ describe('ChannelBase', () => { expect(bridge.prompt).toHaveBeenLastCalledWith( expect.any(String), '[Loop "daily summary" created by Alice] Scheduled task running unattended: no one is present to answer questions, and your final response is delivered to this chat automatically — do whatever work the task requires, then put the result in your final response instead of trying to deliver it to this chat yourself.\n\npost summary', - {}, + { invocationIngress: 'scheduler' }, ); expect(ch.proactive).toEqual([ { chatId: 'group-1', text: 'loop response' }, @@ -12288,7 +12293,7 @@ describe('ChannelBase', () => { expect(bridge.prompt).toHaveBeenLastCalledWith( 's-1', '[Loop "daily summary" created by Alice] Scheduled task running unattended: no one is present to answer questions, and your final response is delivered to this chat automatically — do whatever work the task requires, then put the result in your final response instead of trying to deliver it to this chat yourself.\n\npost again', - {}, + { invocationIngress: 'scheduler' }, ); expect(ch.proactive).toEqual([ { chatId: 'chat1', text: 'second response' }, diff --git a/packages/channels/base/src/ChannelBase.ts b/packages/channels/base/src/ChannelBase.ts index 93cd2a69a9c..794fbef226f 100644 --- a/packages/channels/base/src/ChannelBase.ts +++ b/packages/channels/base/src/ChannelBase.ts @@ -1247,7 +1247,9 @@ export abstract class ChannelBase { jobId: string, timeoutMs: number | undefined, ): Promise { - const prompt = promptBridge.prompt(sessionId, promptText, {}); + const prompt = promptBridge.prompt(sessionId, promptText, { + invocationIngress: 'scheduler', + }); prompt.catch(() => {}); if (timeoutMs === undefined) { return prompt; @@ -3872,6 +3874,7 @@ export abstract class ChannelBase { const response = await promptBridge.prompt(sessionId, promptToSend, { imageBase64, imageMimeType, + invocationIngress: 'channel', }); await this.settleCancelRequested(promptState); diff --git a/packages/channels/base/src/DaemonChannelBridge.ts b/packages/channels/base/src/DaemonChannelBridge.ts index c0034feab81..c2cf38f4ac9 100644 --- a/packages/channels/base/src/DaemonChannelBridge.ts +++ b/packages/channels/base/src/DaemonChannelBridge.ts @@ -7,6 +7,7 @@ import type { AvailableCommand, BridgeSessionInfo, ChannelAgentBridge, + ChannelPromptOptions, ToolCallEvent, } from './ChannelAgentBridge.js'; import { readAvailableCommandAltNames } from './AcpBridge.js'; @@ -31,6 +32,7 @@ export interface DaemonChannelSessionClient { prompt: Array>; }, signal?: AbortSignal, + options?: { invocationIngress?: ChannelPromptOptions['invocationIngress'] }, ): Promise<{ stopReason?: string; [key: string]: unknown }>; events(opts?: { signal?: AbortSignal; @@ -302,7 +304,7 @@ export class DaemonChannelBridge async prompt( sessionId: string, text: string, - options?: { imageBase64?: string; imageMimeType?: string }, + options?: ChannelPromptOptions, ): Promise { const session = this.ensureSession(sessionId); if (this.activePrompts.has(sessionId)) { @@ -360,7 +362,11 @@ export class DaemonChannelBridge prompt.push({ type: 'text', text }); try { - const result = await session.prompt({ prompt }, controller.signal); + const result = options?.invocationIngress + ? await session.prompt({ prompt }, controller.signal, { + invocationIngress: options.invocationIngress, + }) + : await session.prompt({ prompt }, controller.signal); // Prefer turn_complete for deterministic chunk collection (SSE path). // Fall back to one event-loop tick for non-SSE prompt paths (blocking // HTTP, non-202 responses) where turn_complete never arrives. diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index f618b85faed..ec5f5f2fadd 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -128,6 +128,12 @@ vi.mock('node:stream', async (importOriginal) => { // Mock core dependencies vi.mock('@qwen-code/qwen-code-core', () => ({ + INVOCATION_CONTEXT_META_KEY: 'qwen-code/invocation', + INVOCATION_INGRESS_META_KEY: 'qwen-code/invocation-ingress', + PRIVATE_PARENT_CAPABILITY_META_KEY: 'qwen-code/private-parent-capability', + parseInvocationContext: vi.fn((value: unknown) => + typeof value === 'object' && value !== null ? value : undefined, + ), SESSION_ARTIFACT_PERSISTENCE_VERSION: 2, normalizeEventPayload: vi.fn((payload: unknown) => typeof payload === 'object' && @@ -748,6 +754,7 @@ describe('runAcpAgent shutdown cleanup', () => { beforeEach(() => { vi.clearAllMocks(); + delete process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']; // Reset mockConfig after clearAllMocks mockConfig = { initialize: vi.fn().mockResolvedValue(undefined), @@ -803,6 +810,7 @@ describe('runAcpAgent shutdown cleanup', () => { }); afterEach(() => { + delete process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']; processExitSpy.mockRestore(); stdinDestroySpy.mockRestore(); stdoutDestroySpy.mockRestore(); @@ -938,6 +946,31 @@ describe('runAcpAgent shutdown cleanup', () => { expect(dispose).toHaveBeenCalledTimes(1); }); + it('removes the private parent capability before config initialization', async () => { + const envName = 'QWEN_CODE_PRIVATE_ACP_CAPABILITY'; + const previous = process.env[envName]; + process.env[envName] = 'private-capability'; + let observedDuringInitialize: string | undefined; + mockConfig.initialize = vi.fn().mockImplementation(async () => { + observedDuringInitialize = process.env[envName]; + }); + + try { + const agentPromise = runAcpAgent(mockConfig, mockSettings, mockArgv); + await vi.waitFor(() => + expect(mockConfig.initialize).toHaveBeenCalledTimes(1), + ); + expect(observedDuringInitialize).toBeUndefined(); + expect(process.env[envName]).toBeUndefined(); + + mockConnectionState.resolve(); + await agentPromise; + } finally { + if (previous === undefined) delete process.env[envName]; + else process.env[envName] = previous; + } + }); + it('disposes the event loop monitor when connection setup fails', async () => { const dispose = vi.fn(); vi.mocked(startEventLoopLagMonitor).mockReturnValueOnce({ @@ -1291,6 +1324,7 @@ describe('QwenAgent MCP SSE/HTTP support', () => { type AgentLike = { initialize: (args: Record) => Promise; newSession: (args: Record) => Promise; + prompt: (args: Record) => Promise; extMethod: ( method: string, args: Record, @@ -1305,6 +1339,7 @@ describe('QwenAgent MCP SSE/HTTP support', () => { restoreHistory: ReturnType; rewindToTurn: ReturnType; getRewindableUserTurnCount: ReturnType; + prompt: ReturnType; } | undefined; let processExitSpy: MockInstance; @@ -1316,6 +1351,7 @@ describe('QwenAgent MCP SSE/HTTP support', () => { beforeEach(() => { vi.clearAllMocks(); + delete process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']; mockConnectionState.reset(); mockRunExitCleanup.mockResolvedValue(undefined); mockExtensionManagerState.extensions = []; @@ -1368,6 +1404,7 @@ describe('QwenAgent MCP SSE/HTTP support', () => { }); afterEach(() => { + delete process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']; processExitSpy.mockRestore(); stdinDestroySpy.mockRestore(); stdoutDestroySpy.mockRestore(); @@ -1410,6 +1447,178 @@ describe('QwenAgent MCP SSE/HTTP support', () => { await agentPromise; }); + it('rejects every initialize retry after a private capability mismatch', async () => { + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + { + privateParentCapability: 'expected-capability', + }, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + const agent = capturedAgentFactory!({ + get closed() { + return mockConnectionState.promise; + }, + }) as AgentLike; + + await expect( + agent.initialize({ + clientCapabilities: {}, + _meta: { + 'qwen-code/private-parent-capability': 'wrong-capability', + }, + }), + ).rejects.toThrow('Invalid private ACP parent capability'); + await expect( + agent.initialize({ + clientCapabilities: {}, + _meta: { + 'qwen-code/private-parent-capability': 'expected-capability', + }, + }), + ).rejects.toThrow('Invalid private ACP parent capability'); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('rejects every initialize retry after a private capability is missing', async () => { + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + { + privateParentCapability: 'expected-capability', + }, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + const agent = capturedAgentFactory!({ + get closed() { + return mockConnectionState.promise; + }, + }) as AgentLike; + + await expect(agent.initialize({ clientCapabilities: {} })).rejects.toThrow( + 'Invalid private ACP parent capability', + ); + await expect( + agent.initialize({ + clientCapabilities: {}, + _meta: { + 'qwen-code/private-parent-capability': 'expected-capability', + }, + }), + ).rejects.toThrow('Invalid private ACP parent capability'); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('preserves trusted parent invocation metadata and strips the advisory', async () => { + await setupSessionMocks('trusted-session'); + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + { + privateParentCapability: 'expected-capability', + }, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + const agent = capturedAgentFactory!({ + get closed() { + return mockConnectionState.promise; + }, + }) as AgentLike; + await agent.initialize({ + clientCapabilities: {}, + _meta: { + 'qwen-code/private-parent-capability': 'expected-capability', + }, + }); + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + const invocation = { + version: 1, + ingress: 'channel', + sessionId: 'trusted-session', + promptId: 'trusted-prompt', + }; + + await agent.prompt({ + sessionId: 'trusted-session', + prompt: [{ type: 'text', text: 'hello' }], + _meta: { + keep: true, + 'qwen-code/invocation': invocation, + 'qwen-code/invocation-ingress': 'scheduler', + 'qwen-code/private-parent-capability': 'must-not-propagate', + }, + }); + + expect(lastSessionMock?.prompt).toHaveBeenCalledWith({ + sessionId: 'trusted-session', + prompt: [{ type: 'text', text: 'hello' }], + _meta: { + keep: true, + 'qwen-code/invocation': invocation, + }, + }); + mockConnectionState.resolve(); + await agentPromise; + }); + + it('strips forged invocation metadata from an untrusted ACP client', async () => { + const envName = 'QWEN_CODE_PRIVATE_ACP_CAPABILITY'; + process.env[envName] = 'late-env-capability'; + await setupSessionMocks('untrusted-session'); + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + { privateParentCapability: undefined }, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + expect(process.env[envName]).toBeUndefined(); + const agent = capturedAgentFactory!({ + get closed() { + return mockConnectionState.promise; + }, + }) as AgentLike; + await agent.initialize({ + clientCapabilities: {}, + _meta: { + 'qwen-code/private-parent-capability': 'late-env-capability', + }, + }); + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + + await agent.prompt({ + sessionId: 'untrusted-session', + prompt: [{ type: 'text', text: 'hello' }], + _meta: { + keep: true, + 'qwen-code/invocation': { + version: 1, + ingress: 'daemon', + sessionId: 'forged-session', + promptId: 'forged-prompt', + }, + 'qwen-code/invocation-ingress': 'scheduler', + 'qwen-code/private-parent-capability': 'must-not-propagate', + }, + }); + + expect(lastSessionMock?.prompt).toHaveBeenCalledWith({ + sessionId: 'untrusted-session', + prompt: [{ type: 'text', text: 'hello' }], + _meta: { keep: true }, + }); + mockConnectionState.resolve(); + await agentPromise; + }); + it('configures ACP file system fallback roots for read_file allowed local roots', async () => { const previousRoots = process.env[acpLocalReadRootsEnv]; delete process.env[acpLocalReadRootsEnv]; @@ -1901,6 +2110,7 @@ describe('QwenAgent MCP SSE/HTTP support', () => { .fn() .mockReturnValue({ targetTurnIndex: 1, apiTruncateIndex: 2 }), getRewindableUserTurnCount: vi.fn().mockReturnValue(1), + prompt: vi.fn().mockResolvedValue({ stopReason: 'end_turn' }), }; lastSessionMock = sessionMock; return sessionMock as unknown as InstanceType; diff --git a/packages/cli/src/acp-integration/acpAgent.ts b/packages/cli/src/acp-integration/acpAgent.ts index 2db659c8ab1..d36c55adf15 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -40,6 +40,10 @@ import { MCPServerStatus, McpTransportPool, POOLED_TRANSPORTS_DEFAULT, + INVOCATION_CONTEXT_META_KEY, + INVOCATION_INGRESS_META_KEY, + PRIVATE_PARENT_CAPABILITY_META_KEY, + parseInvocationContext, findExistingProviderModels, ExtensionManager, ExtensionSettingScope, @@ -2450,7 +2454,14 @@ export async function runAcpAgent( config: Config, settings: LoadedSettings, argv: CliArgs, + options?: { privateParentCapability?: string }, ) { + const privateParentCapability = + options === undefined + ? process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY'] + : options.privateParentCapability; + delete process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']; + // Reverse tool channel (issue #5626, Phase 2). Runtime-MCP-add targets the // BOOTSTRAP (workspace-level) config's `McpClientManager` — `this.config` in // the `workspaceMcpRuntimeAdd` handler — so a client-hosted MCP server's SDK @@ -2494,7 +2505,13 @@ export async function runAcpAgent( const stream = ndJsonStream(stdout, stdin); connection = new AgentSideConnection((conn) => { acpConnection = conn; - agentInstance = new QwenAgent(config, settings, argv, conn); + agentInstance = new QwenAgent( + config, + settings, + argv, + conn, + privateParentCapability, + ); return agentInstance; }, stream); } catch (err) { @@ -2778,6 +2795,11 @@ class QwenAgent implements Agent { TranscriptReplayConfigCacheEntry >(); private clientCapabilities: ClientCapabilities | undefined; + private privateParentState: + | 'uninitialized' + | 'trusted' + | 'untrusted' + | 'rejected' = 'uninitialized'; // CPU-usage delta baseline for the daemon's `workspaceResource` extMethod // (Daemon Status child-resource chart). The daemon polls this at a fixed // cadence, so successive calls form a clean delta window independent of tool @@ -2952,6 +2974,7 @@ class QwenAgent implements Agent { private settings: LoadedSettings, private argv: CliArgs, private connection: AgentSideConnection, + private readonly expectedPrivateParentCapability?: string, ) { // Pool kill switch via env var so operators can A/B compare or // roll back without rebuilding. `run-qwen-serve.ts` sets this when @@ -3027,6 +3050,29 @@ class QwenAgent implements Agent { } async initialize(args: InitializeRequest): Promise { + if (this.privateParentState === 'rejected') { + throw RequestError.invalidParams( + undefined, + 'Invalid private ACP parent capability', + ); + } + if (this.privateParentState === 'uninitialized') { + const expectedCapability = this.expectedPrivateParentCapability; + if (expectedCapability !== undefined) { + const receivedCapability = + args._meta?.[PRIVATE_PARENT_CAPABILITY_META_KEY]; + if (receivedCapability !== expectedCapability) { + this.privateParentState = 'rejected'; + throw RequestError.invalidParams( + undefined, + 'Invalid private ACP parent capability', + ); + } + this.privateParentState = 'trusted'; + } else { + this.privateParentState = 'untrusted'; + } + } this.clientCapabilities = args.clientCapabilities; const authMethods = buildAuthMethods(); const version = process.env['CLI_VERSION'] || process.version; @@ -3419,7 +3465,52 @@ class QwenAgent implements Agent { if (!session) { throw new Error(`Session not found: ${params.sessionId}`); } - return session.prompt(params); + return session.prompt(this.sanitizeInvocationMetadata(params)); + } + + private sanitizeInvocationMetadata(params: PromptRequest): PromptRequest { + const copy = { ...params }; + const meta = + params._meta && typeof params._meta === 'object' + ? { ...params._meta } + : {}; + const suppliedContext = meta[INVOCATION_CONTEXT_META_KEY]; + const suppliedIngress = meta[INVOCATION_INGRESS_META_KEY]; + delete meta[INVOCATION_CONTEXT_META_KEY]; + delete meta[INVOCATION_INGRESS_META_KEY]; + delete meta[PRIVATE_PARENT_CAPABILITY_META_KEY]; + + if (this.privateParentState === 'trusted') { + if (suppliedContext !== undefined) { + const context = parseInvocationContext(suppliedContext); + if (!context) { + throw RequestError.invalidParams( + undefined, + 'Invalid trusted ACP invocation context', + ); + } + meta[INVOCATION_CONTEXT_META_KEY] = context; + } else if (suppliedIngress !== undefined) { + if ( + suppliedIngress !== 'channel' && + suppliedIngress !== 'scheduler' && + suppliedIngress !== 'internal' + ) { + throw RequestError.invalidParams( + undefined, + 'Invalid trusted ACP invocation ingress', + ); + } + meta[INVOCATION_INGRESS_META_KEY] = suppliedIngress; + } + } + + if (Object.keys(meta).length > 0) { + copy._meta = meta; + } else { + delete copy._meta; + } + return copy; } async cancel(params: CancelNotification): Promise { diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 2bad043012b..f0d825f7cb5 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -2011,6 +2011,59 @@ describe('Session', () => { }); describe('prompt', () => { + it('creates a native ACP invocation root with the durable prompt id', async () => { + let observed: core.InvocationContextV1 | undefined; + mockChat.sendMessageStream = vi.fn().mockImplementation(() => { + observed = core.getInvocationContext(); + return Promise.resolve(createEmptyStream()); + }); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'direct ACP prompt' }], + }); + + expect(observed).toEqual({ + version: 1, + ingress: 'acp', + sessionId: 'test-session-id', + promptId: 'test-session-id########1', + }); + expect(mockChat.sendMessageStream).toHaveBeenCalledWith( + expect.any(String), + expect.any(Object), + 'test-session-id########1', + ); + }); + + it('keeps the native durable prompt id under a trusted parent root', async () => { + const trustedContext: core.InvocationContextV1 = { + version: 1, + ingress: 'daemon', + sessionId: 'test-session-id', + promptId: 'daemon-request-id', + originatorClientId: 'client-1', + }; + let observed: core.InvocationContextV1 | undefined; + mockChat.sendMessageStream = vi.fn().mockImplementation(() => { + observed = core.getInvocationContext(); + return Promise.resolve(createEmptyStream()); + }); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'daemon prompt' }], + _meta: { 'qwen-code/invocation': trustedContext }, + }); + + expect(observed).toEqual(trustedContext); + expect(mockChat.sendMessageStream).toHaveBeenCalledWith( + expect.any(String), + expect.any(Object), + 'test-session-id########1', + ); + }); + it('records the latest file history snapshot after makeSnapshot', async () => { const latestSnapshot = { promptId: 'test-session-id########1', @@ -2125,25 +2178,29 @@ describe('Session', () => { }); it('drains background task notifications through ACP after the prompt is idle', async () => { + let notificationInvocation: core.InvocationContextV1 | undefined; mockChat.sendMessageStream = vi .fn() .mockResolvedValueOnce(createEmptyStream()) - .mockResolvedValueOnce( - createStreamWithChunks([ - { - type: core.StreamEventType.CHUNK, - value: { - candidates: [ - { - content: { - parts: [{ text: 'I saw the background result.' }], + .mockImplementationOnce(() => { + notificationInvocation = core.getInvocationContext(); + return Promise.resolve( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + candidates: [ + { + content: { + parts: [{ text: 'I saw the background result.' }], + }, }, - }, - ], + ], + }, }, - }, - ]), - ); + ]), + ); + }); await session.prompt({ sessionId: 'test-session-id', @@ -2171,6 +2228,13 @@ describe('Session', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); }); + expect(notificationInvocation).toMatchObject({ + version: 1, + ingress: 'internal', + sessionId: 'test-session-id', + }); + expect(notificationInvocation?.promptId).not.toMatch(/notification/u); + expect(mockChat.sendMessageStream).toHaveBeenNthCalledWith( 2, 'qwen3-code-plus', @@ -6342,6 +6406,7 @@ describe('Session', () => { }); it('runs automatic compression before cron-fired ACP prompt sends', async () => { + let cronInvocation: core.InvocationContextV1 | undefined; const scheduler = { size: 1, hasPendingWork: true, @@ -6356,7 +6421,10 @@ describe('Session', () => { mockChat.sendMessageStream = vi .fn() .mockResolvedValueOnce(createEmptyStream()) - .mockResolvedValueOnce(createEmptyStream()); + .mockImplementationOnce(() => { + cronInvocation = core.getInvocationContext(); + return Promise.resolve(createEmptyStream()); + }); await session.prompt({ sessionId: 'test-session-id', @@ -6368,6 +6436,12 @@ describe('Session', () => { }); expect(scheduler.start).toHaveBeenCalledTimes(1); + expect(cronInvocation).toMatchObject({ + version: 1, + ingress: 'scheduler', + sessionId: 'test-session-id', + }); + expect(cronInvocation?.promptId).not.toMatch(/cron/u); expect(mockGeminiClient.tryCompressChat).toHaveBeenNthCalledWith( 1, 'test-session-id########1', diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 8db40a97e11..53c7ba553df 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -5,6 +5,7 @@ */ import { Buffer } from 'node:buffer'; +import { randomUUID } from 'node:crypto'; import * as os from 'node:os'; import * as path from 'node:path'; import type { @@ -129,6 +130,11 @@ import { shouldRunVisionBridge, splitImageParts, approxBase64Bytes, + INVOCATION_CONTEXT_META_KEY, + INVOCATION_INGRESS_META_KEY, + PRIVATE_PARENT_CAPABILITY_META_KEY, + parseInvocationContext, + runWithInvocationContext, } from '@qwen-code/qwen-code-core'; import { NOT_CURRENTLY_GENERATING_CANCEL_MESSAGE } from '@qwen-code/acp-bridge/bridgeErrors'; // Single source of truth shared with the daemon-side answerer (BridgeClient), @@ -1659,27 +1665,73 @@ export class Session implements SessionContext { params: PromptRequest, pendingSend: AbortController, ): Promise { + this.turn += 1; + const sessionId = this.config.getSessionId(); + const promptId = sessionId + '########' + this.turn; + const promptMeta = + params._meta && typeof params._meta === 'object' + ? { ...params._meta } + : {}; + const suppliedContext = promptMeta[INVOCATION_CONTEXT_META_KEY]; + const suppliedIngress = promptMeta[INVOCATION_INGRESS_META_KEY]; + delete promptMeta[INVOCATION_CONTEXT_META_KEY]; + delete promptMeta[INVOCATION_INGRESS_META_KEY]; + delete promptMeta[PRIVATE_PARENT_CAPABILITY_META_KEY]; + const sanitizedParams = { ...params }; + if (Object.keys(promptMeta).length > 0) { + sanitizedParams._meta = promptMeta; + } else { + delete sanitizedParams._meta; + } + const trustedContext = + suppliedContext === undefined + ? undefined + : parseInvocationContext(suppliedContext); + if (suppliedContext !== undefined && !trustedContext) { + throw RequestError.invalidParams( + undefined, + 'Invalid trusted ACP invocation context', + ); + } + const trustedIngress = + suppliedIngress === 'channel' || + suppliedIngress === 'scheduler' || + suppliedIngress === 'internal' + ? suppliedIngress + : undefined; + if (suppliedIngress !== undefined && !trustedIngress) { + throw RequestError.invalidParams( + undefined, + 'Invalid trusted ACP invocation ingress', + ); + } + const invocationContext = trustedContext ?? { + version: 1 as const, + ingress: trustedIngress ?? ('acp' as const), + sessionId, + promptId, + }; + // Bind this turn to the session's ID via AsyncLocalStorage so shell // subprocesses (and hooks) read the CURRENT session's ID instead of // the process-global env slot, which in daemon mode only ever holds // the first session created in this process. - return sessionIdContext.run(this.config.getSessionId(), () => - this.#executePromptInner(params, pendingSend), + return runWithInvocationContext(invocationContext, () => + sessionIdContext.run(sessionId, () => + this.#executePromptInner(sanitizedParams, pendingSend, promptId), + ), ); } async #executePromptInner( params: PromptRequest, pendingSend: AbortController, + promptId: string, ): Promise { return Storage.runWithRuntimeBaseDir( this.runtimeBaseDir, this.config.getWorkingDir(), async () => { - // Increment turn counter for each user prompt - this.turn += 1; - - const promptId = this.config.getSessionId() + '########' + this.turn; const parentContext = extractDaemonTraceContext(params); return await withInteractionSpan( @@ -3083,8 +3135,17 @@ export class Session implements SessionContext { */ async #executeCronPrompt(item: CronQueueItem): Promise { // Same session-ID binding rationale as #executePrompt. - return sessionIdContext.run(this.config.getSessionId(), () => - this.#executeCronPromptInner(item), + return runWithInvocationContext( + { + version: 1, + ingress: 'scheduler', + sessionId: this.config.getSessionId(), + promptId: randomUUID(), + }, + () => + sessionIdContext.run(this.config.getSessionId(), () => + this.#executeCronPromptInner(item), + ), ); } @@ -3533,8 +3594,17 @@ export class Session implements SessionContext { // toolUseId) used in display and response _meta. Merging would // misattribute the combined response to a single task. const item = this.notificationQueue.shift()!; - await sessionIdContext.run(this.config.getSessionId(), () => - this.#executeBackgroundNotificationPromptInner(item), + await runWithInvocationContext( + { + version: 1, + ingress: 'internal', + sessionId: this.config.getSessionId(), + promptId: randomUUID(), + }, + () => + sessionIdContext.run(this.config.getSessionId(), () => + this.#executeBackgroundNotificationPromptInner(item), + ), ); } } finally { diff --git a/packages/cli/src/commands/channel/memory-intent-classifier.test.ts b/packages/cli/src/commands/channel/memory-intent-classifier.test.ts index 07f24a02e01..f76b4ee415d 100644 --- a/packages/cli/src/commands/channel/memory-intent-classifier.test.ts +++ b/packages/cli/src/commands/channel/memory-intent-classifier.test.ts @@ -32,7 +32,7 @@ describe('BridgeChannelMemoryIntentClassifier', () => { expect(bridge.prompt).toHaveBeenCalledWith( 'classifier-session', expect.stringContaining('"你记一下以后回复前说 1122"'), - {}, + { invocationIngress: 'internal' }, ); expect(bridge.cancelSession).toHaveBeenCalledWith('classifier-session'); }); diff --git a/packages/cli/src/commands/channel/memory-intent-classifier.ts b/packages/cli/src/commands/channel/memory-intent-classifier.ts index 5076cfdef4b..51a14dde527 100644 --- a/packages/cli/src/commands/channel/memory-intent-classifier.ts +++ b/packages/cli/src/commands/channel/memory-intent-classifier.ts @@ -94,7 +94,7 @@ export class BridgeChannelMemoryIntentClassifier const response = await bridge.prompt( sessionId, `${CLASSIFIER_PROMPT}${JSON.stringify(text)}`, - {}, + { invocationIngress: 'internal' }, ); return normalizeClassifierResult(extractJsonObject(response)); } finally { diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index c991d08c025..fa87af87acd 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/packages/cli/src/gemini.test.tsx @@ -793,6 +793,7 @@ describe('gemini.tsx main function', () => { const runSandboxRelaunch = async ( argv: string[], sessionId = '123e4567-e89b-12d3-a456-426614174000', + acp = false, ): Promise => { const originalArgv = process.argv; process.argv = argv; @@ -814,6 +815,7 @@ describe('gemini.tsx main function', () => { debug: true, prompt: 'hello', extensions: [], + acp, } as unknown as CliArgs); vi.mocked(loadSettings).mockReturnValue({ errors: [], @@ -865,6 +867,31 @@ describe('gemini.tsx main function', () => { expect(sandboxArgs).not.toContain('--session-id'); }); + it('passes the private ACP capability only to the sandbox child', async () => { + const envName = 'QWEN_CODE_PRIVATE_ACP_CAPABILITY'; + const originalCapability = process.env[envName]; + process.env[envName] = 'private-capability'; + + try { + await runSandboxRelaunch( + ['node', 'script.js', '--acp'], + '123e4567-e89b-12d3-a456-426614174000', + true, + ); + const { start_sandbox } = await import('./utils/sandbox.js'); + expect(vi.mocked(start_sandbox).mock.calls[0]?.[4]).toEqual({ + [envName]: 'private-capability', + }); + expect(process.env[envName]).toBeUndefined(); + } finally { + if (originalCapability === undefined) { + delete process.env[envName]; + } else { + process.env[envName] = originalCapability; + } + } + }); + it('does not pass an empty session ID into the sandbox child process', async () => { const sandboxArgs = await runSandboxRelaunch( ['node', 'script.js', '--debug', '-p', 'hello'], @@ -1545,10 +1572,15 @@ describe('gemini.tsx main function kitty protocol', () => { }); it('should not defer IDE connection when Zed integration is enabled', async () => { + const privateCapability = 'private-capability'; + const originalPrivateCapability = + process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']; + process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY'] = privateCapability; const { loadCliConfig, parseArguments } = await import( './config/config.js' ); const { loadSettings } = await import('./config/settings.js'); + const { relaunchAppInChildProcess } = await import('./utils/relaunch.js'); const initializerModule = await import('./core/initializer.js'); const initializeAppSpy = vi .spyOn(initializerModule, 'initializeApp') @@ -1578,19 +1610,22 @@ describe('gemini.tsx main function kitty protocol', () => { getUsageStatisticsEnabled: () => true, getSessionId: () => 'test-session-id', } as unknown as Config); - vi.mocked(loadSettings).mockReturnValue({ - errors: [], - merged: { - advanced: {}, - security: { auth: {} }, - ui: {}, - }, - setValue: vi.fn(), - forScope: () => ({ settings: {}, originalSettings: {}, path: '' }), - migrationWarnings: [], - getUserHooks: () => undefined, - getProjectHooks: () => undefined, - } as never); + vi.mocked(loadSettings).mockImplementation(() => { + expect(process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']).toBeUndefined(); + return { + errors: [], + merged: { + advanced: {}, + security: { auth: {} }, + ui: {}, + }, + setValue: vi.fn(), + forScope: () => ({ settings: {}, originalSettings: {}, path: '' }), + migrationWarnings: [], + getUserHooks: () => undefined, + getProjectHooks: () => undefined, + } as never; + }); vi.mocked(parseArguments).mockResolvedValue({ model: undefined, sandbox: undefined, @@ -1613,7 +1648,7 @@ describe('gemini.tsx main function kitty protocol', () => { allowedMcpServerNames: undefined, mcpConfig: undefined, allowedTools: undefined, - acp: undefined, + acp: true, experimentalAcp: undefined, extensions: undefined, listExtensions: undefined, @@ -1656,6 +1691,12 @@ describe('gemini.tsx main function kitty protocol', () => { if (!(e instanceof MockProcessExitError)) throw e; } finally { process.exit = originalExit; + if (originalPrivateCapability === undefined) { + delete process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']; + } else { + process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY'] = + originalPrivateCapability; + } } expect(initializeAppSpy).toHaveBeenCalledWith( @@ -1669,6 +1710,16 @@ describe('gemini.tsx main function kitty protocol', () => { expect.any(Object), expect.any(Object), expect.any(Object), + { privateParentCapability: privateCapability }, + ); + expect(relaunchAppInChildProcess).toHaveBeenCalledWith( + expect.any(Array), + [], + expect.objectContaining({ + childEnv: { + QWEN_CODE_PRIVATE_ACP_CAPABILITY: privateCapability, + }, + }), ); expect(mockStartEarlyStartupPrefetches).toHaveBeenCalledWith( expect.any(Object), diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 0896b9c0cbf..687a16f4c06 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -80,6 +80,7 @@ import { getHeadlessYoloSafetyWarning } from './utils/headlessSafetyWarnings.js' import { initializeLlmOutputLanguage } from './utils/languageUtils.js'; const debugLogger = createDebugLogger('STARTUP'); +const PRIVATE_ACP_CAPABILITY_ENV = 'QWEN_CODE_PRIVATE_ACP_CAPABILITY'; interface RuntimeLspReinitializeResult { reconcile: { @@ -234,6 +235,9 @@ export async function main() { setupUnhandledRejectionHandler(); initializeWarningHandler(); + const privateAcpParentCapability = process.env[PRIVATE_ACP_CAPABILITY_ENV]; + delete process.env[PRIVATE_ACP_CAPABILITY_ENV]; + if (process.argv.includes('--bare')) { process.env[QWEN_CODE_SIMPLE_ENV_VAR] = '1'; } @@ -244,6 +248,11 @@ export async function main() { let argv = await parseArguments(); profileCheckpoint('after_parse_arguments'); + const isAcpMode = argv.acp || argv.experimentalAcp; + const privateAcpChildEnv = + isAcpMode && privateAcpParentCapability !== undefined + ? { [PRIVATE_ACP_CAPABILITY_ENV]: privateAcpParentCapability } + : undefined; if (isBareMode(argv.bare)) { process.env[QWEN_CODE_SIMPLE_ENV_VAR] = '1'; @@ -378,7 +387,6 @@ export async function main() { // protocol data (not a user prompt) and should be forwarded to the sandbox // intact via stdio: 'inherit'. const inputFormat = argv.inputFormat as string | undefined; - const isAcpMode = argv.acp || argv.experimentalAcp; let stdinData = ''; if (!process.stdin.isTTY && inputFormat !== 'stream-json' && !isAcpMode) { stdinData = await readStdin(); @@ -443,7 +451,13 @@ export async function main() { : injectStdinIntoArgs(process.argv, stdinData); await relaunchOnExitCode(() => - start_sandbox(sandboxConfig, memoryArgs, partialConfig, sandboxArgs), + start_sandbox( + sandboxConfig, + memoryArgs, + partialConfig, + sandboxArgs, + privateAcpChildEnv, + ), ); process.exit(0); } else { @@ -451,6 +465,7 @@ export async function main() { // restarted if needed. await relaunchAppInChildProcess(memoryArgs, [], { afterSpawn: clearCorruptionEnvVars, + childEnv: privateAcpChildEnv, }); } } @@ -809,7 +824,11 @@ export async function main() { if (config.getExperimentalZedIntegration()) { const { runAcpAgent } = await import('./acp-integration/acpAgent.js'); - await runAcpAgent(config, settings, argv); + await runAcpAgent(config, settings, argv, { + privateParentCapability: isAcpMode + ? privateAcpParentCapability + : undefined, + }); // Clean up child processes and force exit, matching other non-interactive modes await runExitCleanup(); process.exit(0); diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index e298e519556..befdd26db4d 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -29,6 +29,7 @@ import { AUTONOMOUS_SENTINEL_DYNAMIC, LOOP_SENTINEL_CRON, LOOP_SENTINEL_DYNAMIC, + getInvocationContext, } from '@qwen-code/qwen-code-core'; import type { Part } from '@google/genai'; import { @@ -448,6 +449,39 @@ describe('runNonInteractive', () => { expect(mockShutdownTelemetry).toHaveBeenCalled(); }); + it('runs a direct notification turn in a fresh internal invocation context', async () => { + setupMetricsMock(); + const contexts: Array> = []; + mockGeminiClient.sendMessageStream.mockReturnValue( + (async function* () { + contexts.push(getInvocationContext()); + yield { + type: GeminiEventType.Finished, + value: { + reason: undefined, + usageMetadata: { totalTokenCount: 1 }, + }, + }; + })(), + ); + + await runNonInteractive( + mockConfig, + mockSettings, + 'background update', + 'durable-notification-prompt', + { sendMessageType: SendMessageType.Notification }, + ); + + expect(contexts).toHaveLength(1); + expect(contexts[0]).toMatchObject({ + version: 1, + ingress: 'internal', + sessionId: 'test-session-id', + }); + expect(contexts[0]?.promptId).not.toBe('durable-notification-prompt'); + }); + describe('continueInterrupted', () => { it('re-submits an orphaned trailing user prompt with Retry semantics', async () => { setupMetricsMock(); @@ -2290,6 +2324,22 @@ describe('runNonInteractive', () => { ); }); + const drainInvocationContexts: Array< + ReturnType + > = []; + const notificationStream = async function* () { + drainInvocationContexts.push(getInvocationContext()); + yield { type: GeminiEventType.Content, value: 'Observed.' }; + drainInvocationContexts.push(getInvocationContext()); + yield { + type: GeminiEventType.Finished, + value: { + reason: undefined, + usageMetadata: { totalTokenCount: 1 }, + }, + }; + }; + mockGeminiClient.sendMessageStream .mockReturnValueOnce( createStreamFromEvents([ @@ -2303,18 +2353,7 @@ describe('runNonInteractive', () => { }, ]), ) - .mockReturnValueOnce( - createStreamFromEvents([ - { type: GeminiEventType.Content, value: 'Observed.' }, - { - type: GeminiEventType.Finished, - value: { - reason: undefined, - usageMetadata: { totalTokenCount: 1 }, - }, - }, - ]), - ); + .mockReturnValueOnce(notificationStream()); await runNonInteractive( mockConfig, @@ -2324,6 +2363,14 @@ describe('runNonInteractive', () => { ); expect(mockGeminiClient.sendMessageStream).toHaveBeenCalledTimes(2); + expect(drainInvocationContexts).toHaveLength(2); + expect(drainInvocationContexts[0]).toMatchObject({ + version: 1, + ingress: 'internal', + sessionId: 'test-session-id', + }); + expect(drainInvocationContexts[0]?.promptId).not.toBe('prompt-monitor'); + expect(drainInvocationContexts[1]).toEqual(drainInvocationContexts[0]); expect(mockGeminiClient.sendMessageStream).toHaveBeenNthCalledWith( 2, [{ text: notificationXml }], @@ -3624,6 +3671,62 @@ describe('runNonInteractive', () => { expect(predicate({ prompt: 'regular cron job' } as CronJob)).toBe(false); }); + it('runs each drained cron stream in a fresh scheduler invocation context', async () => { + setupMetricsMock(); + let fire: ((job: CronJob) => void) | undefined; + const scheduler = { + sessionSize: 0, + setSkipDurableFire: vi.fn(), + enableDurable: vi.fn().mockResolvedValue(undefined), + start: vi.fn((callback: (job: CronJob) => void) => { + fire = callback; + callback({ + prompt: 'scheduled check', + cronExpr: '* * * * *', + } as CronJob); + }), + stop: vi.fn(), + } as unknown as CronScheduler; + mockConfig.isCronEnabled = vi.fn().mockReturnValue(true); + mockConfig.getCronScheduler = vi.fn().mockReturnValue(scheduler); + + const observedContexts: Array> = []; + const cronStream = async function* () { + observedContexts.push(getInvocationContext()); + yield { type: GeminiEventType.Content, value: 'Scheduled result.' }; + }; + mockGeminiClient.sendMessageStream + .mockReturnValueOnce( + createStreamFromEvents([ + { type: GeminiEventType.Content, value: 'Initial result.' }, + ]), + ) + .mockReturnValueOnce(cronStream()); + + await runNonInteractive( + mockConfig, + mockSettings, + 'initial prompt', + 'durable-prompt-id', + ); + + expect(fire).toBeTypeOf('function'); + expect(observedContexts).toHaveLength(1); + expect(observedContexts[0]).toMatchObject({ + version: 1, + ingress: 'scheduler', + sessionId: 'test-session-id', + }); + expect(observedContexts[0]?.promptId).not.toBe('durable-prompt-id'); + expect(mockGeminiClient.sendMessageStream).toHaveBeenNthCalledWith( + 2, + [{ text: 'scheduled check' }], + expect.any(AbortSignal), + 'durable-prompt-id', + expect.objectContaining({ type: SendMessageType.Cron }), + ); + }); + describe('--json-schema structured output', () => { // Helper: walk an emitted event and extract the first tool_use_id when // it represents a tool_result block. Returns undefined for any other diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 0f3a410f394..8d4aea8a3e6 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -4,11 +4,13 @@ * SPDX-License-Identifier: Apache-2.0 */ +import { randomUUID } from 'node:crypto'; import type { BackgroundTaskStatus, Config, CronJob, CronScheduler, + InvocationContextV1, ToolCallRequestInfo, } from '@qwen-code/qwen-code-core'; import { isSlashCommand } from './ui/utils/commandUtils.js'; @@ -21,6 +23,7 @@ import { GeminiEventType, FatalInputError, promptIdContext, + runWithInvocationContext, OutputFormat, InputFormat, LoopType, @@ -77,6 +80,30 @@ const debugLogger = createDebugLogger('NON_INTERACTIVE_CLI'); */ const STRUCTURED_SHUTDOWN_HOLDBACK_MS = 500; +async function* iterateWithInvocationContext( + iterable: AsyncIterable, + context: InvocationContextV1, +): AsyncGenerator { + const iterator = iterable[Symbol.asyncIterator](); + let completed = false; + try { + while (true) { + const result = await runWithInvocationContext(context, () => + iterator.next(), + ); + if (result.done) { + completed = true; + return; + } + yield result.value; + } + } finally { + if (!completed && iterator.return) { + await runWithInvocationContext(context, () => iterator.return!()); + } + } +} + function isHeadlessLoopSentinel(prompt: string): boolean { return ( detectLoopSentinel(prompt) !== null || @@ -308,6 +335,32 @@ export async function runNonInteractive( input: string, prompt_id: string, options: RunNonInteractiveOptions = {}, +): Promise { + const ingress = + options.sendMessageType === SendMessageType.Cron + ? ('scheduler' as const) + : options.sendMessageType === SendMessageType.Notification || + options.sendMessageType === SendMessageType.Teammate + ? ('internal' as const) + : ('cli' as const); + return runWithInvocationContext( + { + version: 1, + ingress, + sessionId: config.getSessionId(), + promptId: ingress === 'cli' ? prompt_id : randomUUID(), + }, + () => + runNonInteractiveWithContext(config, settings, input, prompt_id, options), + ); +} + +async function runNonInteractiveWithContext( + config: Config, + settings: LoadedSettings, + input: string, + prompt_id: string, + options: RunNonInteractiveOptions, ): Promise { return promptIdContext.run(prompt_id, async (): Promise => { // Create output adapter based on format @@ -1585,6 +1638,13 @@ export async function runNonInteractive( modelText: batch.map((i) => i.modelText).join('\n\n'), sendMessageType: targetType, }; + const itemInvocationContext: InvocationContextV1 = { + version: 1, + ingress: + targetType === SendMessageType.Cron ? 'scheduler' : 'internal', + sessionId: config.getSessionId(), + promptId: randomUUID(), + }; turnCount++; if ( @@ -1604,19 +1664,24 @@ export async function runNonInteractive( while (true) { const itemToolCallRequests: ToolCallRequestInfo[] = []; const itemApiStartTime = Date.now(); - const itemStream = geminiClient.sendMessageStream( - itemMessages[0]?.parts || [], - abortController.signal, - prompt_id, - { - type: itemIsFirstTurn - ? item.sendMessageType - : SendMessageType.ToolResult, - modelOverride: itemModelOverride, - ...(itemIsFirstTurn && { - notificationDisplayText: item.displayText, - }), - }, + const itemStream = iterateWithInvocationContext( + runWithInvocationContext(itemInvocationContext, () => + geminiClient.sendMessageStream( + itemMessages[0]?.parts || [], + abortController.signal, + prompt_id, + { + type: itemIsFirstTurn + ? item.sendMessageType + : SendMessageType.ToolResult, + modelOverride: itemModelOverride, + ...(itemIsFirstTurn && { + notificationDisplayText: item.displayText, + }), + }, + ), + ), + itemInvocationContext, ); itemIsFirstTurn = false; @@ -1689,11 +1754,10 @@ export async function runNonInteractive( const { responseParts: itemToolResponseParts, repeatedDuplicateProviderToolCall, - } = await processToolCallBatch( - itemToolCallRequests, - (override) => { + } = await runWithInvocationContext(itemInvocationContext, () => + processToolCallBatch(itemToolCallRequests, (override) => { itemModelOverride = override; - }, + }), ); if (structuredSubmission !== undefined) { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 22f8086c461..84e51710347 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -23,6 +23,7 @@ import type { EditorType, GeminiClient, AnyToolInvocation, + InvocationContextV1, } from '@qwen-code/qwen-code-core'; import { ApprovalMode, @@ -32,6 +33,8 @@ import { SendMessageType, ToolErrorType, ToolConfirmationOutcome, + getInvocationContext, + runWithInvocationContext, } from '@qwen-code/qwen-code-core'; import type { Part, PartListUnion } from '@google/genai'; import type { UseHistoryManagerReturn } from './useHistoryManager.js'; @@ -387,6 +390,17 @@ describe('useGeminiStream', () => { it('queues background shell terminal notifications for the model loop', async () => { const { mockSendMessageStream } = renderTestHook(); + const inheritedContext: InvocationContextV1 = { + version: 1, + ingress: 'cli', + sessionId: 'inherited-session', + promptId: 'test-session-id########5', + }; + let observedContext: InvocationContextV1 | undefined; + mockSendMessageStream.mockImplementationOnce(async function* () { + observedContext = getInvocationContext(); + yield* []; + }); const displayText = 'Background shell "npm test" completed.'; const modelText = '\nshell\ncompleted\n'; @@ -400,9 +414,11 @@ describe('useGeminiStream', () => { const callback = mockBackgroundShellRegistry.setNotificationCallback.mock .calls[0][0] as (displayText: string, modelText: string) => void; - act(() => { - callback(displayText, modelText); - }); + act(() => + runWithInvocationContext(inheritedContext, () => + callback(displayText, modelText), + ), + ); await waitFor(() => { expect(mockAddItem).toHaveBeenCalledWith( @@ -421,6 +437,12 @@ describe('useGeminiStream', () => { }), ); }); + expect(observedContext).toMatchObject({ + version: 1, + ingress: 'internal', + sessionId: mockConfig.getSessionId(), + }); + expect(observedContext?.promptId).not.toBe(inheritedContext.promptId); }); describe('vision bridge gate', () => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 6b22edd681b..d93b7755a9d 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -12,6 +12,7 @@ import { useMemo, useLayoutEffect, } from 'react'; +import { randomUUID } from 'node:crypto'; import { type Config, type EditorType, @@ -28,6 +29,7 @@ import { type StopFailureErrorType, type ActiveGoal, type VisionBridgeResult, + type InvocationContextV1, GeminiEventType as ServerGeminiEventType, SendMessageType, createDebugLogger, @@ -45,6 +47,8 @@ import { ApprovalMode, parseAndFormatApiError, promptIdContext, + getInvocationContext, + runWithInvocationContext, ToolConfirmationOutcome, logApiCancel, ApiCancelEvent, @@ -113,6 +117,16 @@ import process from 'node:process'; const debugLogger = createDebugLogger('GEMINI_STREAM'); +function runWithPromptContexts( + invocation: InvocationContextV1, + promptId: string, + callback: () => T, +): T { + return runWithInvocationContext(invocation, () => + promptIdContext.run(promptId, callback), + ); +} + // The per-turn model override is held in two coupled refs: the model id and a // flag marking whether it came from an explicit inline `/model ` // (which must win over skill-tool overrides for the rest of the turn). The two @@ -2465,7 +2479,32 @@ export const useGeminiStream = ( prompt_id = config.getSessionId() + '########' + getPromptCount(); } - return promptIdContext.run(prompt_id, async () => { + const inheritedInvocation = getInvocationContext(); + const asyncIngress = + submitType === SendMessageType.Cron + ? ('scheduler' as const) + : submitType === SendMessageType.Notification || + submitType === SendMessageType.Teammate + ? ('internal' as const) + : undefined; + const invocation = + asyncIngress !== undefined + ? { + version: 1 as const, + ingress: asyncIngress, + sessionId: config.getSessionId(), + promptId: randomUUID(), + } + : inheritedInvocation?.promptId === prompt_id + ? inheritedInvocation + : { + version: 1 as const, + ingress: 'cli' as const, + sessionId: config.getSessionId(), + promptId: prompt_id, + }; + + return runWithPromptContexts(invocation, prompt_id, async () => { const { queryToSend, shouldProceed } = submitType === SendMessageType.Retry ? { queryToSend: query, shouldProceed: true } diff --git a/packages/cli/src/utils/relaunch.test.ts b/packages/cli/src/utils/relaunch.test.ts index 5c45049b548..b51d918c0d9 100644 --- a/packages/cli/src/utils/relaunch.test.ts +++ b/packages/cli/src/utils/relaunch.test.ts @@ -308,6 +308,29 @@ describe('relaunchAppInChildProcess', () => { expect(afterSpawn).toHaveBeenCalledTimes(1); }); + it('passes child-only environment without restoring it on the parent', async () => { + process.argv = ['/usr/bin/node', '/app/cli.js']; + delete process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']; + const mockChild = createMockChildProcess(0, false); + mockedSpawn.mockReturnValue(mockChild); + + const promise = relaunchAppInChildProcess([], [], { + childEnv: { + QWEN_CODE_PRIVATE_ACP_CAPABILITY: 'private-capability', + }, + }); + + await vi.waitFor(() => expect(mockedSpawn).toHaveBeenCalledOnce()); + expect(mockedSpawn.mock.calls[0]?.[2]?.env).toMatchObject({ + QWEN_CODE_PRIVATE_ACP_CAPABILITY: 'private-capability', + QWEN_CODE_NO_RELAUNCH: 'true', + }); + expect(process.env['QWEN_CODE_PRIVATE_ACP_CAPABILITY']).toBeUndefined(); + + mockChild.emit('close', 0); + await expect(promise).rejects.toThrow('PROCESS_EXIT_CALLED'); + }); + it('should handle null exit code from child process', async () => { process.argv = ['/usr/bin/node', '/app/cli.js']; diff --git a/packages/cli/src/utils/relaunch.ts b/packages/cli/src/utils/relaunch.ts index 07eb634cef2..b957bf224ec 100644 --- a/packages/cli/src/utils/relaunch.ts +++ b/packages/cli/src/utils/relaunch.ts @@ -28,7 +28,10 @@ export async function relaunchOnExitCode(runner: () => Promise) { export async function relaunchAppInChildProcess( additionalNodeArgs: string[], additionalScriptArgs: string[], - options?: { afterSpawn?: () => void }, + options?: { + afterSpawn?: () => void; + childEnv?: Readonly>; + }, ) { if (process.env['QWEN_CODE_NO_RELAUNCH']) { return; @@ -47,7 +50,11 @@ export async function relaunchAppInChildProcess( ...additionalScriptArgs, ...scriptArgs, ]; - const newEnv = { ...process.env, QWEN_CODE_NO_RELAUNCH: 'true' }; + const newEnv = { + ...process.env, + ...options?.childEnv, + QWEN_CODE_NO_RELAUNCH: 'true', + }; // The parent process should not be reading from stdin while the child is running. process.stdin.pause(); diff --git a/packages/cli/src/utils/sandbox.ts b/packages/cli/src/utils/sandbox.ts index 11c9b9e1227..d6ea9a2ff8e 100644 --- a/packages/cli/src/utils/sandbox.ts +++ b/packages/cli/src/utils/sandbox.ts @@ -193,6 +193,7 @@ export async function start_sandbox( nodeArgs: string[] = [], cliConfig?: Config, cliArgs: string[] = [], + childEnv?: Readonly>, ): Promise { if (config.command === 'sandbox-exec') { // disallow BUILD_SANDBOX @@ -342,6 +343,7 @@ export async function start_sandbox( process.stdin.pause(); sandboxProcess = spawn(config.command, args, { stdio: 'inherit', + ...(childEnv ? { env: { ...process.env, ...childEnv } } : {}), }); return new Promise((resolve, reject) => { sandboxProcess?.on('error', reject); @@ -759,6 +761,10 @@ export async function start_sandbox( } } + for (const name of Object.keys(childEnv ?? {})) { + args.push('--env', name); + } + // copy NODE_OPTIONS const existingNodeOptions = process.env['NODE_OPTIONS'] || ''; const allNodeOptions = [ @@ -896,6 +902,7 @@ export async function start_sandbox( process.stdin.pause(); sandboxProcess = spawn(config.command, args, { stdio: 'inherit', + ...(childEnv ? { env: { ...process.env, ...childEnv } } : {}), }); return new Promise((resolve, reject) => { diff --git a/packages/core/src/agents/background-agent-resume.test.ts b/packages/core/src/agents/background-agent-resume.test.ts index fb157718412..28f8c482889 100644 --- a/packages/core/src/agents/background-agent-resume.test.ts +++ b/packages/core/src/agents/background-agent-resume.test.ts @@ -21,6 +21,11 @@ import { AgentTerminateMode } from './runtime/agent-types.js'; import { AgentEventEmitter } from './runtime/agent-events.js'; import { getCurrentAgentDepth } from './runtime/agent-context.js'; import { AgentHeadless } from './runtime/agent-headless.js'; +import { + getInvocationContext, + runWithInvocationContext, + type InvocationContextV1, +} from '../utils/invocation-context.js'; import { FORK_DEFAULT_MAX_TURNS, FORK_SUBAGENT_TYPE, @@ -437,8 +442,11 @@ describe('BackgroundAgentResumeService', () => { metaPath, }); + let resumedInvocation: InvocationContextV1 | undefined; const execute = vi.fn( - async (_context: { get: (key: string) => unknown }) => undefined, + async (_context: { get: (key: string) => unknown }) => { + resumedInvocation = getInvocationContext(); + }, ); const setExternalMessageProvider = vi.fn(); const subagent = { @@ -463,7 +471,16 @@ describe('BackgroundAgentResumeService', () => { getAdditionalContext: () => 'resume-context', }); - const resumed = await service.resumeBackgroundAgent(agentId, 'continue'); + const resumed = await runWithInvocationContext( + { + version: 1, + ingress: 'daemon', + sessionId, + promptId: 'stale-daemon-prompt', + originatorClientId: 'stale-client', + }, + () => service.resumeBackgroundAgent(agentId, 'continue'), + ); expect(resumed).toBeDefined(); expect(hookSystem.fireSubagentStartEvent).toHaveBeenCalledWith( @@ -482,6 +499,13 @@ describe('BackgroundAgentResumeService', () => { } expect(contextArg.get('hook_context')).toBe('resume-context'); expect(contextArg.get('task_prompt')).toBe('continue'); + expect(resumedInvocation).toMatchObject({ + version: 1, + ingress: 'internal', + sessionId, + }); + expect(resumedInvocation?.promptId).not.toBe('stale-daemon-prompt'); + expect(resumedInvocation?.originatorClientId).toBeUndefined(); await vi.waitFor(() => { expect(registry.get(agentId)?.status).toBe('completed'); }); diff --git a/packages/core/src/agents/background-agent-resume.ts b/packages/core/src/agents/background-agent-resume.ts index 9edd0e08c46..ec46cee6885 100644 --- a/packages/core/src/agents/background-agent-resume.ts +++ b/packages/core/src/agents/background-agent-resume.ts @@ -6,6 +6,7 @@ import * as fs from 'node:fs/promises'; import * as path from 'node:path'; +import { randomUUID } from 'node:crypto'; import type { Content, FunctionDeclaration, Part } from '@google/genai'; import type { Config } from '../config/config.js'; import * as jsonl from '../utils/jsonl-utils.js'; @@ -28,6 +29,7 @@ import type { ChatRecord } from '../services/chatRecordingService.js'; import { buildOrderedUuidChain } from '../utils/conversation-chain.js'; import { getInitialChatHistory } from '../utils/environmentContext.js'; import { getGitBranch } from '../utils/gitUtils.js'; +import { runWithInvocationContext } from '../utils/invocation-context.js'; import { PermissionMode, type StopHookOutput } from '../hooks/types.js'; import { appendStopHookBlockingCapWarning, @@ -1040,7 +1042,19 @@ export class BackgroundAgentResumeService { runBody, normalizeResumedAgentDepth(meta.depth), ); - void (target.isFork ? runInForkContext(framedRunBody) : framedRunBody()); + const invocationRunBody = () => + runWithInvocationContext( + { + version: 1, + ingress: 'internal', + sessionId: meta.parentSessionId, + promptId: randomUUID(), + }, + framedRunBody, + ); + void (target.isFork + ? runInForkContext(invocationRunBody) + : invocationRunBody()); return entry; } catch (error) { cleanupOwnedMonitorNotifications?.(); diff --git a/packages/core/src/agents/runtime/agent-core.test.ts b/packages/core/src/agents/runtime/agent-core.test.ts index feb68aa74b6..552fa475e7c 100644 --- a/packages/core/src/agents/runtime/agent-core.test.ts +++ b/packages/core/src/agents/runtime/agent-core.test.ts @@ -16,6 +16,11 @@ import { type RuntimeContentGeneratorView, } from './agent-context.js'; import { subagentNameContext } from '../../utils/subagentNameContext.js'; +import { + getInvocationContext, + runWithInvocationContext, + type InvocationContextV1, +} from '../../utils/invocation-context.js'; import { runInForkContext } from '../../tools/agent/fork-subagent.js'; import { ToolNames } from '../../tools/tool-names.js'; import { @@ -138,6 +143,58 @@ describe('AgentCore.runInAgentFrames', () => { expect(onConfirmInvocations[0]!.name).toBe('approval-agent'); }); + it('restores the captured invocation context for deferred approval', async () => { + const core = makeCore('approval-agent'); + const capturedContext: InvocationContextV1 = { + version: 1, + ingress: 'channel', + sessionId: 'session-a', + promptId: 'prompt-a', + }; + const unrelatedContext: InvocationContextV1 = { + version: 1, + ingress: 'daemon', + sessionId: 'session-b', + promptId: 'prompt-b', + }; + let respond: (() => Promise) | undefined; + let observed: InvocationContextV1 | undefined; + + await runWithInvocationContext(capturedContext, async () => { + const inherited = getInvocationContext(); + respond = () => + runWithInvocationContext(inherited, () => + core.runInAgentFrames(async () => { + observed = getInvocationContext(); + }), + ); + }); + + await runWithInvocationContext(unrelatedContext, () => respond!()); + expect(observed).toBe(capturedContext); + }); + + it('clears an unrelated invocation context when deferred approval captured none', async () => { + const core = makeCore('approval-agent'); + const inherited = getInvocationContext(); + const unrelatedContext: InvocationContextV1 = { + version: 1, + ingress: 'daemon', + sessionId: 'session-b', + promptId: 'prompt-b', + }; + let observed: InvocationContextV1 | undefined; + const respond = () => + runWithInvocationContext(inherited, () => + core.runInAgentFrames(async () => { + observed = getInvocationContext(); + }), + ); + + await runWithInvocationContext(unrelatedContext, respond); + expect(observed).toBeUndefined(); + }); + it('still publishes the agent name when no runtime view is set (inheriting agent)', async () => { const core = makeCore('inherit-agent'); diff --git a/packages/core/src/agents/runtime/agent-core.ts b/packages/core/src/agents/runtime/agent-core.ts index ca1fbc0bca6..35f40005769 100644 --- a/packages/core/src/agents/runtime/agent-core.ts +++ b/packages/core/src/agents/runtime/agent-core.ts @@ -20,6 +20,10 @@ import { randomUUID } from 'node:crypto'; import { createChildAbortController } from '../../utils/abortController.js'; import { reportError } from '../../utils/errorReporting.js'; import { subagentNameContext } from '../../utils/subagentNameContext.js'; +import { + getInvocationContext, + runWithInvocationContext, +} from '../../utils/invocation-context.js'; import type { Config } from '../../config/config.js'; import { getCurrentAgentDepth, @@ -1608,6 +1612,7 @@ export class AgentCore { // can restore it. See `runInAgentFrames` for why this matters // (mis-attributed `from="leader"` + leader-guard bypass). const inheritedTeammateIdentity = getTeammateContext(); + const inheritedInvocationContext = getInvocationContext(); this.eventEmitter?.emit(AgentEventType.TOOL_WAITING_APPROVAL, { subagentId: this.subagentId, round: currentRound, @@ -1633,12 +1638,15 @@ export class AgentCore { // tool body runs. See `runInAgentFrames` for rationale. // Also restore the logical owner agent id when present so // approved tools such as Monitor keep owner routing. - await this.runInAgentFrames( - () => waiting.confirmationDetails.onConfirm(outcome, payload), - inheritedView, - inheritedAgentId ?? undefined, - inheritedTeammateIdentity, - inheritedAgentDepth, + await runWithInvocationContext(inheritedInvocationContext, () => + this.runInAgentFrames( + () => + waiting.confirmationDetails.onConfirm(outcome, payload), + inheritedView, + inheritedAgentId ?? undefined, + inheritedTeammateIdentity, + inheritedAgentDepth, + ), ); }, timestamp: Date.now(), diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 78d6164ec9f..98219f47c0a 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -64,6 +64,11 @@ import type { ShellToolParams } from '../tools/shell.js'; import type { ShellExecutionConfig } from '../services/shellExecutionService.js'; import { runWithAgentContext } from '../agents/runtime/agent-context.js'; import { runWithTeammateIdentity } from '../agents/team/identity.js'; +import { + getInvocationContext, + runWithInvocationContext, + type InvocationContextV1, +} from '../utils/invocation-context.js'; type ToolSpanRecord = { name: string; @@ -793,6 +798,66 @@ describe('CoreToolScheduler', () => { }; } + it('restores the invocation context when a delayed confirmation executes', async () => { + const invocationContext: InvocationContextV1 = { + version: 1, + ingress: 'cli', + sessionId: 'session-context', + promptId: 'prompt-context', + }; + const unrelatedContext: InvocationContextV1 = { + ...invocationContext, + sessionId: 'unrelated-session', + promptId: 'unrelated-prompt', + }; + let observedContext: InvocationContextV1 | undefined; + const tool = new MockTool({ + name: 'approval-context-tool', + getDefaultPermission: async () => 'ask', + getConfirmationDetails: async () => ({ + type: 'info' as const, + title: 'Confirm context tool', + prompt: 'Run context tool?', + onConfirm: vi.fn().mockResolvedValue(undefined), + }), + execute: async () => { + observedContext = getInvocationContext(); + return { llmContent: 'ok', returnDisplay: 'ok' }; + }, + }); + const { scheduler, onToolCallsUpdate } = createSchedulerForLegacyToolTests({ + toolsByName: new Map([[tool.name, tool]]), + approvalMode: ApprovalMode.DEFAULT, + }); + + await runWithInvocationContext(invocationContext, () => + scheduler.schedule( + [ + { + callId: 'approval-context-call', + name: tool.name, + args: {}, + isClientInitiated: false, + prompt_id: invocationContext.promptId, + }, + ], + new AbortController().signal, + ), + ); + const waiting = (await waitForStatus( + onToolCallsUpdate, + 'awaiting_approval', + )) as WaitingToolCall; + + await runWithInvocationContext(unrelatedContext, () => + waiting.confirmationDetails.onConfirm( + ToolConfirmationOutcome.ProceedOnce, + ), + ); + + expect(observedContext).toEqual(invocationContext); + }); + it('dispatches legacy tool names through their canonical registered tools', async () => { const canonicalNamesByLegacyName = new Map( Object.entries(ToolNamesMigration), diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index bf62b8104d0..ed0aeafd873 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -142,6 +142,10 @@ import { } from '../telemetry/index.js'; import { safeJsonStringify } from '../utils/safeJsonStringify.js'; import { acquireSleepInhibitor } from '../services/sleepInhibitor.js'; +import { + getInvocationContext, + runWithInvocationContext, +} from '../utils/invocation-context.js'; const debugLogger = createDebugLogger('TOOL_SCHEDULER'); @@ -2658,6 +2662,7 @@ export class CoreToolScheduler { ); const originalOnConfirm = confirmationDetails.onConfirm; + const invocationContext = getInvocationContext(); const wrappedConfirmationDetails: ToolCallConfirmationDetails = { ...confirmationDetails, // When PM has an explicit 'ask' rule, 'always allow' would be @@ -2668,12 +2673,14 @@ export class CoreToolScheduler { outcome: ToolConfirmationOutcome, payload?: ToolConfirmationPayload, ) => - this.handleConfirmationResponse( - reqInfo.callId, - originalOnConfirm, - outcome, - signal, - payload, + runWithInvocationContext(invocationContext, () => + this.handleConfirmationResponse( + reqInfo.callId, + originalOnConfirm, + outcome, + signal, + payload, + ), ), }; this.setStatusInternal( diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index ae7a247524d..eb55d99b7fd 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -470,6 +470,7 @@ export * from './utils/gitDirect.js'; export * from './utils/gitIgnoreParser.js'; export * from './utils/gitUtils.js'; export * from './utils/ignorePatterns.js'; +export * from './utils/invocation-context.js'; export { DEFAULT_QWEN_CUSTOM_IGNORE_FILE_NAMES, QwenIgnoreParser, diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index bcf599743a5..3cc79386c99 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -9,6 +9,7 @@ import * as ClientLib from '@modelcontextprotocol/sdk/client/index.js'; import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'; import * as SdkClientStdioLib from '@modelcontextprotocol/sdk/client/stdio.js'; import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; +import type { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { AuthProviderType, @@ -22,6 +23,11 @@ import { OAuthUtils } from '../mcp/oauth-utils.js'; import type { PromptRegistry } from '../prompts/prompt-registry.js'; import type { ResourceRegistry } from '../resources/resource-registry.js'; import type { WorkspaceContext } from '../utils/workspaceContext.js'; +import { + INVOCATION_CONTEXT_META_KEY, + runWithInvocationContext, + type InvocationContextV1, +} from '../utils/invocation-context.js'; import { addMCPStatusChangeListener, attemptAutomaticMcpOAuth, @@ -1326,6 +1332,163 @@ describe('mcp-client', () => { expect(tools[0].alwaysLoad).toBe(true); }); + it('allows invocation context for the raw client bound to a created stdio transport', async () => { + const callTool = vi.fn().mockResolvedValue({ + content: [{ type: 'text', text: 'ok' }], + }); + const mockedClient = { + connect: vi.fn(), + registerCapabilities: vi.fn(), + setRequestHandler: vi.fn(), + getServerCapabilities: vi.fn().mockReturnValue({}), + listTools: vi.fn().mockResolvedValue({ tools: [] }), + getInstructions: vi.fn(), + callTool, + }; + vi.mocked(ClientLib.Client).mockReturnValue( + mockedClient as unknown as ClientLib.Client, + ); + vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue( + {} as SdkClientStdioLib.StdioClientTransport, + ); + vi.mocked(GenAiLib.mcpToTool).mockReturnValue({ + tool: () => + Promise.resolve({ + functionDeclarations: [{ name: 'local-tool' }], + }), + } as unknown as GenAiLib.CallableTool); + const client = new McpClient( + 'local-server', + { command: 'test-command' }, + {} as ToolRegistry, + {} as PromptRegistry, + {} as WorkspaceContext, + false, + ); + await client.connect(); + const { tools } = await client.discoverAndReturn(cfgWithResources()); + const context: InvocationContextV1 = { + version: 1, + ingress: 'cli', + sessionId: 'session-1', + promptId: 'prompt-1', + }; + + await runWithInvocationContext(context, () => + tools[0].build({ param: 'test' }).execute(new AbortController().signal), + ); + + expect(callTool.mock.calls[0][0]._meta).toEqual({ + [INVOCATION_CONTEXT_META_KEY]: context, + }); + }); + + it('does not trust a stdio-shaped server config without an internally bound client', async () => { + const callTool = vi.fn().mockResolvedValue({ + content: [{ type: 'text', text: 'ok' }], + }); + const arbitraryClient = { + listTools: vi.fn().mockResolvedValue({ tools: [] }), + callTool, + } as unknown as ClientLib.Client; + vi.mocked(GenAiLib.mcpToTool).mockReturnValue({ + tool: () => + Promise.resolve({ + functionDeclarations: [{ name: 'unbound-tool' }], + }), + } as unknown as GenAiLib.CallableTool); + const [unboundTool] = await discoverTools( + 'unbound-server', + { command: 'looks-like-stdio' }, + arbitraryClient, + cfgWithResources(), + ); + const context: InvocationContextV1 = { + version: 1, + ingress: 'daemon', + sessionId: 'session-1', + promptId: 'prompt-1', + }; + const modelArguments = { + _meta: { [INVOCATION_CONTEXT_META_KEY]: { forged: true } }, + }; + + await runWithInvocationContext(context, () => + unboundTool.build(modelArguments).execute(new AbortController().signal), + ); + + expect(callTool.mock.calls[0][0]).toEqual({ + name: 'unbound-tool', + arguments: modelArguments, + }); + }); + + it.each([ + ['Streamable HTTP', { httpUrl: 'http://example.test/mcp' }, undefined], + ['SSE', { url: 'http://example.test/sse' }, undefined], + [ + 'SDK/reverse', + { type: 'sdk' as const }, + async (_serverName: string, message: JSONRPCMessage) => message, + ], + ])( + 'denies invocation context for %s transport clients', + async (_transportName, serverConfig, sendSdkMcpMessage) => { + const callTool = vi.fn().mockResolvedValue({ + content: [{ type: 'text', text: 'ok' }], + }); + const mockedClient = { + connect: vi.fn(), + registerCapabilities: vi.fn(), + setRequestHandler: vi.fn(), + getServerCapabilities: vi.fn().mockReturnValue({}), + listTools: vi.fn().mockResolvedValue({ tools: [] }), + getInstructions: vi.fn(), + callTool, + }; + vi.mocked(ClientLib.Client).mockReturnValue( + mockedClient as unknown as ClientLib.Client, + ); + vi.mocked(MCPOAuthTokenStorage).mockImplementation( + () => + ({ + getCredentials: vi.fn().mockResolvedValue(null), + }) as unknown as MCPOAuthTokenStorage, + ); + vi.mocked(GenAiLib.mcpToTool).mockReturnValue({ + tool: () => + Promise.resolve({ + functionDeclarations: [{ name: 'remote-tool' }], + }), + } as unknown as GenAiLib.CallableTool); + const client = new McpClient( + 'remote-server', + serverConfig, + {} as ToolRegistry, + {} as PromptRegistry, + {} as WorkspaceContext, + false, + sendSdkMcpMessage, + ); + await client.connect(); + const { tools } = await client.discoverAndReturn(cfgWithResources()); + const context: InvocationContextV1 = { + version: 1, + ingress: 'daemon', + sessionId: 'session-1', + promptId: 'prompt-1', + }; + + await runWithInvocationContext(context, () => + tools[0] + .build({ param: 'test' }) + .execute(new AbortController().signal), + ); + + expect(Object.hasOwn(callTool.mock.calls[0][0], '_meta')).toBe(false); + }, + ); + it('discoverAndReturn with { applyConfigFilters: false } ignores config filters and trust for shared pool snapshots', async () => { const mockedClient = { connect: vi.fn(), diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index f62d0c8036e..7adf6079c7a 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -78,6 +78,19 @@ export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes const debugLogger = createDebugLogger('MCP'); const AUTOMATIC_MCP_OAUTH_TIMEOUT_MS = 60_000; +const invocationContextTransports = new WeakSet(); +const invocationContextClients = new WeakSet(); + +function bindInvocationContextPolicy( + client: Client, + transport: Transport, +): void { + invocationContextClients.delete(client); + if (invocationContextTransports.has(transport)) { + invocationContextClients.add(client); + } +} + const STREAMABLE_HTTP_GET_SSE_FALLBACK_STATUSES = new Set([400]); const STREAMABLE_HTTP_GET_SSE_ERROR_BODY_LIMIT = 512; @@ -313,6 +326,7 @@ export class McpClient { */ async connect(): Promise { this.isDisconnecting = false; + invocationContextClients.delete(this.client); clearMcpOAuthRequirement(this.serverName, this.serverConfig); // clear stale upstream error from // any prior connect/disconnect cycle. The silent-drop reader @@ -367,6 +381,7 @@ export class McpClient { timeout: this.serverConfig.timeout, }); this.instructions = this.client.getInstructions(); + bindInvocationContextPolicy(this.client, this.transport); this.updateStatus(MCPServerStatus.CONNECTED); } catch (error) { @@ -1286,6 +1301,7 @@ export async function discoverTools( cliConfig?.getMcpToolIdleTimeoutMs?.(), annotationsMap.get(funcDecl.name!), mcpServerConfig.alwaysLoadTools === true, + invocationContextClients.has(mcpClient), ), ); } catch (error) { @@ -1626,6 +1642,7 @@ export async function connectToMcpServer( await mcpClient.connect(transport, { timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, }); + bindInvocationContextPolicy(mcpClient, transport); clearMcpOAuthRequirement(mcpServerName, mcpServerConfig); return mcpClient; } catch (error) { @@ -2132,6 +2149,7 @@ export async function createTransport( cwd: mcpServerConfig.cwd, stderr: 'pipe', }); + invocationContextTransports.add(transport); if (debugMode) { transport.stderr!.on('data', (data) => { const stderrStr = data.toString().trim(); diff --git a/packages/core/src/tools/mcp-tool.test.ts b/packages/core/src/tools/mcp-tool.test.ts index a695654a243..10d554ad3eb 100644 --- a/packages/core/src/tools/mcp-tool.test.ts +++ b/packages/core/src/tools/mcp-tool.test.ts @@ -16,12 +16,23 @@ import { import type { ToolResult } from './tools.js'; import { ToolConfirmationOutcome } from './tools.js'; import type { CallableTool, Part } from '@google/genai'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; +import { + LATEST_PROTOCOL_VERSION, + type JSONRPCMessage, +} from '@modelcontextprotocol/sdk/types.js'; import { ToolErrorType } from './tool-error.js'; import { MCPServerStatus, removeMCPServerStatus, updateMCPServerStatus, } from './mcp-client.js'; +import { + INVOCATION_CONTEXT_META_KEY, + runWithInvocationContext, + type InvocationContextV1, +} from '../utils/invocation-context.js'; vi.mock('node:fs/promises'); @@ -1086,6 +1097,197 @@ describe('DiscoveredMCPTool', () => { }); }); + describe('invocation context metadata', () => { + const invocationContext: InvocationContextV1 = { + version: 1, + ingress: 'daemon', + sessionId: 'session-1', + promptId: 'prompt-1', + originatorClientId: 'client-1', + }; + + const createDirectTool = ( + mcpClient: McpDirectClient, + allowInvocationContext: boolean, + ) => + new DiscoveredMCPTool( + mockCallableToolInstance, + serverName, + serverToolName, + baseDescription, + inputSchema, + undefined, + undefined, + undefined, + mcpClient, + undefined, + undefined, + undefined, + false, + allowInvocationContext, + ); + + const successfulClient = () => + ({ + callTool: vi.fn(async () => ({ + content: [{ type: 'text', text: 'ok' }], + })), + }) satisfies McpDirectClient; + + it('injects request metadata only for an allowed tool with a context', async () => { + const mcpClient = successfulClient(); + const modelArguments = { + param: 'test', + _meta: { + [INVOCATION_CONTEXT_META_KEY]: { forged: true }, + }, + }; + + await runWithInvocationContext(invocationContext, () => + createDirectTool(mcpClient, true) + .build(modelArguments) + .execute(new AbortController().signal), + ); + + expect(mcpClient.callTool).toHaveBeenCalledWith( + { + name: serverToolName, + arguments: modelArguments, + _meta: { + [INVOCATION_CONTEXT_META_KEY]: invocationContext, + }, + }, + undefined, + expect.objectContaining({ + onprogress: expect.any(Function), + }), + ); + }); + + it.each([ + { allowInvocationContext: false, runWithContext: true }, + { allowInvocationContext: true, runWithContext: false }, + ])( + 'omits reserved request metadata for $allowInvocationContext/$runWithContext', + async ({ allowInvocationContext, runWithContext }) => { + const mcpClient = successfulClient(); + const execute = () => + createDirectTool(mcpClient, allowInvocationContext) + .build({ param: 'test' }) + .execute(new AbortController().signal); + + if (runWithContext) { + await runWithInvocationContext(invocationContext, execute); + } else { + await execute(); + } + + expect(mcpClient.callTool).toHaveBeenCalled(); + expect( + Object.hasOwn( + vi.mocked(mcpClient.callTool).mock.calls[0][0], + '_meta', + ), + ).toBe(false); + }, + ); + + it('preserves the policy through qualification and trust clones', async () => { + const mcpClient = successfulClient(); + const clonedTool = createDirectTool(mcpClient, true) + .asFullyQualifiedTool() + .withTrust(true); + + await runWithInvocationContext(invocationContext, () => + clonedTool + .build({ param: 'test' }) + .execute(new AbortController().signal), + ); + + expect(vi.mocked(mcpClient.callTool).mock.calls[0][0]._meta).toEqual({ + [INVOCATION_CONTEXT_META_KEY]: invocationContext, + }); + }); + + it('coexists with MCP SDK progress metadata on allowed and denied calls', async () => { + const sentMessages: JSONRPCMessage[] = []; + const transport: Transport = { + start: vi.fn(async () => undefined), + close: vi.fn(async () => undefined), + send: vi.fn(async (message: JSONRPCMessage) => { + sentMessages.push(message); + if (!('id' in message) || !('method' in message)) { + return; + } + + if (message.method === 'initialize') { + queueMicrotask(() => + transport.onmessage?.({ + jsonrpc: '2.0', + id: message.id, + result: { + protocolVersion: LATEST_PROTOCOL_VERSION, + capabilities: { tools: {} }, + serverInfo: { name: 'test-server', version: '1.0.0' }, + }, + }), + ); + } else if (message.method === 'tools/call') { + queueMicrotask(() => + transport.onmessage?.({ + jsonrpc: '2.0', + id: message.id, + result: { + content: [{ type: 'text', text: 'ok' }], + }, + }), + ); + } + }), + }; + const client = new Client({ name: 'test-client', version: '1.0.0' }); + await client.connect(transport); + + await runWithInvocationContext(invocationContext, () => + createDirectTool(client, true) + .build({ param: 'test' }) + .execute(new AbortController().signal), + ); + await runWithInvocationContext(invocationContext, () => + createDirectTool(client, false) + .build({ param: 'test' }) + .execute(new AbortController().signal), + ); + + const callRequests = sentMessages.filter( + (message) => 'method' in message && message.method === 'tools/call', + ); + expect(callRequests[0]).toMatchObject({ + params: { + _meta: { + [INVOCATION_CONTEXT_META_KEY]: invocationContext, + progressToken: expect.any(Number), + }, + }, + }); + expect(callRequests[1]).toMatchObject({ + params: { + _meta: { + progressToken: expect.any(Number), + }, + }, + }); + expect(callRequests[1]).not.toMatchObject({ + params: { + _meta: { + [INVOCATION_CONTEXT_META_KEY]: expect.anything(), + }, + }, + }); + await client.close(); + }); + }); + describe('streaming progress for long-running MCP tools', () => { it('should have canUpdateOutput set to true so the scheduler creates liveOutputCallback', () => { // For long-running MCP tools (e.g., browseruse), the scheduler needs @@ -1304,6 +1506,73 @@ describe('DiscoveredMCPTool', () => { expect(result.llmContent).toEqual([{ text: 'Success after reconnect' }]); }); + it('uses the newly discovered tool policy after reconnect', async () => { + const invocationContext: InvocationContextV1 = { + version: 1, + ingress: 'daemon', + sessionId: 'session-1', + promptId: 'prompt-1', + }; + const oldMcpClient: McpDirectClient = { + callTool: vi.fn().mockRejectedValueOnce(new Error('Connection closed')), + }; + const newMcpClient: McpDirectClient = { + callTool: vi.fn().mockResolvedValueOnce({ + content: [{ type: 'text', text: 'reconnected' }], + }), + }; + const newTool = new DiscoveredMCPTool( + mockCallableToolInstance, + serverName, + serverToolName, + baseDescription, + inputSchema, + undefined, + undefined, + undefined, + newMcpClient, + ); + const mockConfig = { + isTrustedFolder: () => true, + getToolRegistry: () => ({ + discoverToolsForServer: vi.fn().mockResolvedValue(undefined), + ensureTool: vi.fn().mockResolvedValue(newTool), + }), + getTruncateToolOutputThreshold: () => 0, + getTruncateToolOutputLines: () => 0, + }; + const oldTool = new DiscoveredMCPTool( + mockCallableToolInstance, + serverName, + serverToolName, + baseDescription, + inputSchema, + undefined, + undefined, + mockConfig as any, + oldMcpClient, + undefined, + undefined, + undefined, + false, + true, + ); + + await runWithInvocationContext(invocationContext, () => + oldTool.build({ param: 'test' }).execute(new AbortController().signal), + ); + + expect(vi.mocked(oldMcpClient.callTool).mock.calls[0][0]._meta).toEqual({ + [INVOCATION_CONTEXT_META_KEY]: invocationContext, + }); + expect( + Object.hasOwn( + vi.mocked(newMcpClient.callTool).mock.calls[0][0], + '_meta', + ), + ).toBe(false); + }); + it('should not retry on non-connection errors', async () => { const params = { param: 'test' }; const mockMcpClient: McpDirectClient = { diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index 47ea64426cb..e09cb3ad027 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -24,6 +24,10 @@ import { truncateToolOutput } from '../utils/truncation.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import { getErrorMessage, isAbortError } from '../utils/errors.js'; import { getMCPServerStatus, MCPServerStatus } from './mcp-status.js'; +import { + getInvocationContext, + INVOCATION_CONTEXT_META_KEY, +} from '../utils/invocation-context.js'; const debugLogger = createDebugLogger('MCP_TOOL'); @@ -47,7 +51,11 @@ type ToolParams = Record; */ export interface McpDirectClient { callTool( - params: { name: string; arguments?: Record }, + params: { + name: string; + arguments?: Record; + _meta?: Record; + }, resultSchema?: unknown, options?: { onprogress?: (progress: { @@ -138,6 +146,7 @@ class DiscoveredMCPToolInvocation extends BaseToolInvocation< private readonly mcpTimeout?: number, private readonly mcpToolIdleTimeoutMs?: number, private readonly annotations?: McpToolAnnotations, + private readonly allowInvocationContext: boolean = false, private readonly retryCount: number = 0, ) { super(params); @@ -267,6 +276,7 @@ class DiscoveredMCPToolInvocation extends BaseToolInvocation< this.mcpTimeout, this.mcpToolIdleTimeoutMs, this.annotations, + newTool['allowInvocationContext'] === true, this.retryCount + 1, ); return newInvocation.execute(signal, updateOutput); @@ -351,10 +361,20 @@ class DiscoveredMCPToolInvocation extends BaseToolInvocation< // Start the idle timeout resetIdleTimeout(); + const invocationContext = this.allowInvocationContext + ? getInvocationContext() + : undefined; const callToolResult = await this.mcpClient!.callTool( { name: this.serverToolName, arguments: this.params as Record, + ...(invocationContext + ? { + _meta: { + [INVOCATION_CONTEXT_META_KEY]: invocationContext, + }, + } + : {}), }, undefined, { @@ -554,6 +574,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< private readonly mcpToolIdleTimeoutMs?: number, readonly annotations?: McpToolAnnotations, alwaysLoad = false, + private readonly allowInvocationContext: boolean = false, ) { super( nameOverride ?? @@ -588,6 +609,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< this.mcpToolIdleTimeoutMs, this.annotations, this.alwaysLoad, + this.allowInvocationContext, ); } @@ -625,6 +647,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< this.mcpToolIdleTimeoutMs, this.annotations, this.alwaysLoad, + this.allowInvocationContext, ); } @@ -643,6 +666,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< this.mcpTimeout, this.mcpToolIdleTimeoutMs, this.annotations, + this.allowInvocationContext, ); } } diff --git a/packages/core/src/utils/invocation-context.test.ts b/packages/core/src/utils/invocation-context.test.ts new file mode 100644 index 00000000000..85c15f41bbf --- /dev/null +++ b/packages/core/src/utils/invocation-context.test.ts @@ -0,0 +1,138 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { + getInvocationContext, + INVOCATION_CONTEXT_META_KEY, + INVOCATION_INGRESS_META_KEY, + parseInvocationContext, + PRIVATE_PARENT_CAPABILITY_META_KEY, + runWithInvocationContext, + type InvocationContextV1, +} from './invocation-context.js'; + +const context: InvocationContextV1 = { + version: 1, + ingress: 'external_mcp', + sessionId: 'session-1', + promptId: 'prompt-1', + originatorClientId: 'client-1', +}; + +describe('invocation context wire contract', () => { + it('exports the reserved metadata keys', () => { + expect(INVOCATION_CONTEXT_META_KEY).toBe('qwen-code/invocation'); + expect(INVOCATION_INGRESS_META_KEY).toBe('qwen-code/invocation-ingress'); + expect(PRIVATE_PARENT_CAPABILITY_META_KEY).toBe( + 'qwen-code/private-parent-capability', + ); + }); + + it('parses valid contexts with and without an originator', () => { + expect(parseInvocationContext(context)).toEqual(context); + expect( + parseInvocationContext({ + version: 1, + ingress: 'cli', + sessionId: 'session-2', + promptId: 'prompt-2', + }), + ).toEqual({ + version: 1, + ingress: 'cli', + sessionId: 'session-2', + promptId: 'prompt-2', + }); + }); + + it.each([ + undefined, + null, + [], + { + version: 2, + ingress: 'cli', + sessionId: 'session-3', + promptId: 'prompt-3', + }, + { + version: 1, + ingress: 'webhook', + sessionId: 'session-4', + promptId: 'prompt-4', + }, + { + version: 1, + ingress: 'daemon', + sessionId: ' ', + promptId: 'prompt-5', + }, + { + version: 1, + ingress: 'daemon', + sessionId: 'session-6', + promptId: '', + }, + { + version: 1, + ingress: 'daemon', + sessionId: 'session-7', + promptId: 'prompt-7', + originatorClientId: ' ', + }, + { + version: 1, + ingress: 'daemon', + sessionId: 'session-8', + promptId: 'prompt-8', + trusted: true, + }, + ])('rejects malformed context %#', (value) => { + expect(parseInvocationContext(value)).toBeUndefined(); + }); +}); + +describe('invocation context async storage', () => { + it('is undefined outside a context and restores nested contexts', () => { + const nestedContext: InvocationContextV1 = { + ...context, + promptId: 'prompt-2', + }; + + expect(getInvocationContext()).toBeUndefined(); + runWithInvocationContext(context, () => { + expect(getInvocationContext()).toBe(context); + runWithInvocationContext(nestedContext, () => { + expect(getInvocationContext()).toBe(nestedContext); + }); + expect(getInvocationContext()).toBe(context); + }); + expect(getInvocationContext()).toBeUndefined(); + }); + + it('can explicitly clear and then restore an outer context', () => { + runWithInvocationContext(context, () => { + expect(getInvocationContext()).toBe(context); + runWithInvocationContext(undefined, () => { + expect(getInvocationContext()).toBeUndefined(); + }); + expect(getInvocationContext()).toBe(context); + }); + }); + + it('isolates concurrent async execution trees', async () => { + const readPromptId = (promptId: string) => + runWithInvocationContext({ ...context, promptId }, async () => { + await Promise.resolve(); + return getInvocationContext()?.promptId; + }); + + await expect( + Promise.all([readPromptId('prompt-a'), readPromptId('prompt-b')]), + ).resolves.toEqual(['prompt-a', 'prompt-b']); + }); +}); diff --git a/packages/core/src/utils/invocation-context.ts b/packages/core/src/utils/invocation-context.ts new file mode 100644 index 00000000000..1777b9a8a15 --- /dev/null +++ b/packages/core/src/utils/invocation-context.ts @@ -0,0 +1,104 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { AsyncLocalStorage } from 'node:async_hooks'; + +export const INVOCATION_CONTEXT_META_KEY = 'qwen-code/invocation'; +export const INVOCATION_INGRESS_META_KEY = 'qwen-code/invocation-ingress'; +export const PRIVATE_PARENT_CAPABILITY_META_KEY = + 'qwen-code/private-parent-capability'; + +export type InvocationIngress = + | 'cli' + | 'acp' + | 'daemon' + | 'channel' + | 'scheduler' + | 'external_mcp' + | 'internal'; + +export interface InvocationContextV1 { + readonly version: 1; + readonly ingress: InvocationIngress; + readonly sessionId: string; + readonly promptId: string; + readonly originatorClientId?: string; +} + +const invocationIngresses = new Set([ + 'cli', + 'acp', + 'daemon', + 'channel', + 'scheduler', + 'external_mcp', + 'internal', +]); + +const invocationContextKeys = new Set([ + 'version', + 'ingress', + 'sessionId', + 'promptId', + 'originatorClientId', +]); + +const invocationContextStorage = new AsyncLocalStorage< + InvocationContextV1 | undefined +>(); + +function isNonBlankString(value: unknown): value is string { + return typeof value === 'string' && value.trim().length > 0; +} + +export function parseInvocationContext( + value: unknown, +): InvocationContextV1 | undefined { + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + return undefined; + } + + const record = value as Record; + if (Object.keys(record).some((key) => !invocationContextKeys.has(key))) { + return undefined; + } + if ( + record['version'] !== 1 || + typeof record['ingress'] !== 'string' || + !invocationIngresses.has(record['ingress'] as InvocationIngress) || + !isNonBlankString(record['sessionId']) || + !isNonBlankString(record['promptId']) + ) { + return undefined; + } + + let originatorClientId: string | undefined; + if (Object.hasOwn(record, 'originatorClientId')) { + if (!isNonBlankString(record['originatorClientId'])) { + return undefined; + } + originatorClientId = record['originatorClientId']; + } + + return { + version: 1, + ingress: record['ingress'] as InvocationIngress, + sessionId: record['sessionId'], + promptId: record['promptId'], + ...(originatorClientId ? { originatorClientId } : {}), + }; +} + +export function runWithInvocationContext( + context: InvocationContextV1 | undefined, + callback: () => T, +): T { + return invocationContextStorage.run(context, callback); +} + +export function getInvocationContext(): InvocationContextV1 | undefined { + return invocationContextStorage.getStore(); +} diff --git a/packages/sdk-typescript/scripts/build.js b/packages/sdk-typescript/scripts/build.js index e79d45b92ad..4b01cb63abe 100755 --- a/packages/sdk-typescript/scripts/build.js +++ b/packages/sdk-typescript/scripts/build.js @@ -207,18 +207,35 @@ await esbuild.build({ }); // Build serve-bridge CLI bin entry +const serveBridgeBinPath = join( + rootDir, + 'dist', + 'daemon-mcp', + 'serve-bridge', + 'bin.js', +); await esbuild.build({ entryPoints: [join(rootDir, 'src', 'daemon-mcp', 'serve-bridge', 'bin.ts')], bundle: true, format: 'esm', platform: 'node', target: 'node22', - outfile: join(rootDir, 'dist', 'daemon-mcp', 'serve-bridge', 'bin.js'), + outfile: serveBridgeBinPath, external: ['@modelcontextprotocol/sdk'], sourcemap: false, - banner: { js: '#!/usr/bin/env node' }, }); +const serveBridgeBin = readFileSync(serveBridgeBinPath, 'utf8'); +const nodeShebang = '#!/usr/bin/env node'; +const shebangCount = serveBridgeBin + .split(/\r?\n/) + .filter((line) => line === nodeShebang).length; +if (!serveBridgeBin.startsWith(`${nodeShebang}\n`) || shebangCount !== 1) { + throw new Error( + `qwen-serve-mcp bundle must start with exactly one Node shebang; found ${shebangCount}`, + ); +} + // Copy LICENSE from root directory to dist const licenseSource = join(rootDir, '..', '..', 'LICENSE'); const licenseTarget = join(rootDir, 'dist', 'LICENSE'); diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/README.md b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/README.md index ae3e5b12f3a..3a7683363e1 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/README.md +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/README.md @@ -116,6 +116,10 @@ const server = createServeBridgeMcpServer({ const transport = new StdioServerTransport(); await server.instance.connect(transport); + +// During shutdown, close the protocol first, then release daemon bindings. +await server.instance.close(); +await server.dispose(); ``` ## 提供的工具（共 31 个） diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bin.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bin.ts index e0ada0b959a..c6da9c9a555 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bin.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bin.ts @@ -21,6 +21,7 @@ import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; import { createServeBridgeMcpServer } from './createServeBridgeMcpServer.js'; +import { createShutdownController } from './shutdown.js'; const server = createServeBridgeMcpServer({ daemonUrl: process.env['QWEN_DAEMON_URL'] ?? 'http://127.0.0.1:4170', @@ -31,18 +32,24 @@ const server = createServeBridgeMcpServer({ const transport = new StdioServerTransport(); -// Graceful shutdown on signals -async function shutdown() { - try { - await server.instance.close(); - } catch (e) { - process.stderr.write(`[qwen-serve-bridge] close error: ${e}\n`); - } - process.exit(0); -} +const { shutdown, markInstanceClosed } = createShutdownController({ + close: () => server.instance.close(), + dispose: () => server.dispose(), + exit: (code) => process.exit(code), + reportCloseError: (err) => { + process.stderr.write(`[qwen-serve-bridge] close error: ${err}\n`); + }, +}); + +const previousOnclose = server.instance.server.onclose; +server.instance.server.onclose = () => { + markInstanceClosed(); + previousOnclose?.(); + void shutdown(); +}; -process.on('SIGINT', shutdown); -process.on('SIGTERM', shutdown); +process.on('SIGINT', () => void shutdown()); +process.on('SIGTERM', () => void shutdown()); // Prevent silent crashes from unhandled rejections process.on('unhandledRejection', (err) => { @@ -53,6 +60,6 @@ process.on('unhandledRejection', (err) => { }); // Exit cleanly when stdio pipe closes (parent process gone) -process.stdin.on('close', shutdown); +process.stdin.on('close', () => void shutdown()); await server.instance.connect(transport); diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bindings.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bindings.ts new file mode 100644 index 00000000000..3b8216e7939 --- /dev/null +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bindings.ts @@ -0,0 +1,206 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { isInvalidDaemonClientIdError } from '../../daemon/DaemonHttpError.js'; +import { createEventStream, startEventStream, stopEventStream } from './sse.js'; +import type { BridgeState, SessionBinding } from './types.js'; + +const SESSION_TTL_MS = 30 * 60 * 1000; +const CLEANUP_INTERVAL_MS = 5 * 60 * 1000; + +export function createBinding( + sessionId: string, + clientId?: string, +): SessionBinding { + return { + sessionId, + clientId, + stream: createEventStream(sessionId), + }; +} + +export async function replaceBinding( + state: BridgeState, + next: SessionBinding, +): Promise { + const previous = state.bindings.get(next.sessionId); + if (previous?.stream.activeCollector) { + await releaseBinding(state, next); + throw new Error( + 'Cannot replace a session binding while a prompt is in progress.', + ); + } + if (state.disposed) { + await releaseBinding(state, next); + throw new Error('qwen-serve-bridge is shutting down.'); + } + + state.bindings.set(next.sessionId, next); + state.defaultSessionId = next.sessionId; + startEventStream(state, next, () => { + void withSessionLock(state, next.sessionId, async () => { + if (state.bindings.get(next.sessionId) === next) { + await releaseBinding(state, next); + } + }); + }); + + if (previous) { + await releaseBinding(state, previous); + } +} + +export function releaseBinding( + state: BridgeState, + binding: SessionBinding, + detach = true, +): Promise { + if (binding.releasePromise) return binding.releasePromise; + + if (state.bindings.get(binding.sessionId) === binding) { + state.bindings.delete(binding.sessionId); + if (state.defaultSessionId === binding.sessionId) { + state.defaultSessionId = undefined; + } + } + + const releasePromise = ( + detach && binding.clientId + ? state.client.detachSession(binding.sessionId, binding.clientId) + : Promise.resolve() + ) + .catch((err: unknown) => { + const detail = err instanceof Error ? err.message : String(err); + process.stderr.write( + `[serve-bridge] Failed to detach session ${binding.sessionId}: ${detail}\n`, + ); + }) + .finally(() => { + state.pendingReleases.delete(releasePromise); + }); + binding.releasePromise = releasePromise; + state.pendingReleases.add(releasePromise); + stopEventStream(binding.stream); + return releasePromise; +} + +export async function withSessionLock( + state: BridgeState, + sessionId: string, + fn: () => Promise, +): Promise { + const previous = state.sessionLocks.get(sessionId) ?? Promise.resolve(); + let unlock!: () => void; + const current = new Promise((resolve) => { + unlock = resolve; + }); + const tail = previous.then(() => current); + state.sessionLocks.set(sessionId, tail); + await previous; + try { + return await fn(); + } finally { + unlock(); + if (state.sessionLocks.get(sessionId) === tail) { + state.sessionLocks.delete(sessionId); + } + } +} + +export function trackLifecycle( + state: BridgeState, + fn: () => Promise, +): Promise { + if (state.disposed) { + return Promise.reject(new Error('qwen-serve-bridge is shutting down.')); + } + + const operation = Promise.resolve().then(fn); + const tracked = operation.then( + () => undefined, + () => undefined, + ); + state.pendingLifecycles.add(tracked); + void tracked.then(() => state.pendingLifecycles.delete(tracked)); + return operation; +} + +export function resolveBinding( + state: BridgeState, + sessionId: string, +): SessionBinding { + const binding = state.bindings.get(sessionId); + if (!binding) { + throw new Error( + `No active binding for session ${sessionId}. Call session_resume first.`, + ); + } + binding.stream.lastActivityMs = Date.now(); + return binding; +} + +export async function invalidateBinding( + state: BridgeState, + binding: SessionBinding, +): Promise { + if (state.bindings.get(binding.sessionId) !== binding) return false; + await releaseBinding(state, binding, false); + return true; +} + +export async function rethrowBindingError( + state: BridgeState, + binding: SessionBinding, + err: unknown, +): Promise { + if (!isInvalidDaemonClientIdError(err)) throw err; + + if (await invalidateBinding(state, binding)) { + throw new Error( + `The daemon rejected the client binding for session ${binding.sessionId}. Call session_resume before retrying.`, + { cause: err }, + ); + } + + throw new Error( + `The binding for session ${binding.sessionId} changed while the request was in flight. Retry the request.`, + { cause: err }, + ); +} + +export function startSessionCleanup(state: BridgeState): () => void { + const timer = setInterval(() => { + const now = Date.now(); + for (const binding of state.bindings.values()) { + if ( + !binding.stream.activeCollector && + now - binding.stream.lastActivityMs > SESSION_TTL_MS + ) { + process.stderr.write( + `[serve-bridge] Cleaning up idle session SSE: ${binding.sessionId}\n`, + ); + void releaseBinding(state, binding); + } + } + }, CLEANUP_INTERVAL_MS); + timer.unref(); + return () => clearInterval(timer); +} + +export async function disposeBindings(state: BridgeState): Promise { + state.disposed = true; + while (state.pendingLifecycles.size > 0) { + await Promise.all([...state.pendingLifecycles]); + } + await Promise.all( + [...state.bindings.values()].map((binding) => + releaseBinding(state, binding), + ), + ); + while (state.pendingReleases.size > 0) { + await Promise.all([...state.pendingReleases]); + } +} diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/createServeBridgeMcpServer.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/createServeBridgeMcpServer.ts index 93674e27348..b61cbbedb4c 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/createServeBridgeMcpServer.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/createServeBridgeMcpServer.ts @@ -12,9 +12,13 @@ import { DaemonClient } from '../../daemon/DaemonClient.js'; import { createSdkMcpServer } from '../createSdkMcpServer.js'; import type { McpSdkServerConfigWithInstance } from '../createSdkMcpServer.js'; import type { ServeBridgeMcpServerOptions, BridgeState } from './types.js'; -import { startSessionCleanup, stopEventStream } from './sse.js'; +import { disposeBindings, startSessionCleanup } from './bindings.js'; import { allTools } from './tools/index.js'; +export type ServeBridgeMcpServer = McpSdkServerConfigWithInstance & { + dispose(): Promise; +}; + /** Strip trailing slashes without regex (avoids CodeQL ReDoS flag). */ function stripTrailingSlashes(url: string): string { let end = url.length; @@ -37,21 +41,30 @@ function stripTrailingSlashes(url: string): string { * * const transport = new StdioServerTransport(); * await server.instance.connect(transport); + * + * // During shutdown: + * await server.instance.close(); + * await server.dispose(); * ``` */ export function createServeBridgeMcpServer( opts: ServeBridgeMcpServerOptions, -): McpSdkServerConfigWithInstance { +): ServeBridgeMcpServer { const state: BridgeState = { client: new DaemonClient({ baseUrl: opts.daemonUrl, token: opts.token, + invocationIngress: 'external_mcp', }), daemonUrl: stripTrailingSlashes(opts.daemonUrl), token: opts.token, defaultSessionId: undefined, workspaceCwd: opts.workspaceCwd, - eventStreams: new Map(), + bindings: new Map(), + sessionLocks: new Map(), + pendingLifecycles: new Set(), + pendingReleases: new Set(), + disposed: false, allowGlobalScope: opts.allowGlobalScope ?? false, }; @@ -66,16 +79,20 @@ export function createServeBridgeMcpServer( tools, }); - // Stop cleanup timer and abort all active SSE streams when server closes. - // Use the SDK's onclose lifecycle hook (Protocol.onclose) instead of - // monkey-patching close() — the SDK calls onclose after transport shutdown - // and internal state cleanup, which is the supported extension point. + let disposePromise: Promise | undefined; + const dispose = (): Promise => { + disposePromise ??= (async () => { + stopCleanup(); + await disposeBindings(state); + })(); + return disposePromise; + }; + + // Programmatic callers should await dispose(); onclose is a backstop for + // callers that only close the underlying MCP server. server.instance.server.onclose = () => { - stopCleanup(); - for (const sessionId of [...state.eventStreams.keys()]) { - stopEventStream(state, sessionId); - } + void dispose(); }; - return server; + return { ...server, dispose }; } diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/helpers.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/helpers.ts index fa3721ae67b..524742dd34f 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/helpers.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/helpers.ts @@ -10,6 +10,10 @@ import type { BridgeState } from './types.js'; +export interface HandlerExtra { + signal: AbortSignal | undefined; +} + /** * Resolve the session ID from explicit arg or default state. * Returns the session ID or throws a descriptive error. @@ -25,9 +29,9 @@ export function resolveSessionId( ); } // Bump activity timestamp so workspace operations reset the idle TTL - const stream = state.eventStreams.get(sessionId); - if (stream) { - stream.lastActivityMs = Date.now(); + const binding = state.bindings.get(sessionId); + if (binding) { + binding.stream.lastActivityMs = Date.now(); } return sessionId; } @@ -38,11 +42,18 @@ export function resolveSessionId( */ /* eslint-disable @typescript-eslint/no-explicit-any */ export function handler( - fn: (args: T) => Promise, + fn: (args: T, extra: HandlerExtra) => Promise, ): (args: T, extra: unknown) => Promise { - return async (args: T, _extra: unknown) => { + return async (args: T, extra: unknown) => { try { - return await fn(args); + const signal = + typeof extra === 'object' && + extra !== null && + 'signal' in extra && + extra.signal instanceof AbortSignal + ? extra.signal + : undefined; + return await fn(args, { signal }); } catch (err) { const message = err instanceof Error ? err.message : String(err); // Log full error with stack for debugging diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/index.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/index.ts index e3f9a11235a..07beaa2f26f 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/index.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/index.ts @@ -5,4 +5,5 @@ */ export { createServeBridgeMcpServer } from './createServeBridgeMcpServer.js'; +export type { ServeBridgeMcpServer } from './createServeBridgeMcpServer.js'; export type { ServeBridgeMcpServerOptions, BridgeState } from './types.js'; diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/shutdown.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/shutdown.ts new file mode 100644 index 00000000000..af398af5d5b --- /dev/null +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/shutdown.ts @@ -0,0 +1,44 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +export function createShutdownController(options: { + close: () => Promise; + dispose: () => Promise; + exit: (code: number) => void; + reportCloseError: (err: unknown) => void; +}): { + shutdown: () => Promise; + markInstanceClosed: () => void; +} { + let instanceClosed = false; + let shutdownPromise: Promise | undefined; + + const performShutdown = async (): Promise => { + if (!instanceClosed) { + instanceClosed = true; + try { + await options.close(); + } catch (err) { + options.reportCloseError(err); + } + } + await options.dispose(); + options.exit(0); + }; + + const shutdown = (): Promise => { + // Defer work so a recursive onclose observes the assigned promise. + shutdownPromise ??= Promise.resolve().then(performShutdown); + return shutdownPromise; + }; + + return { + shutdown, + markInstanceClosed: () => { + instanceClosed = true; + }, + }; +} diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/sse.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/sse.ts index 76fe19216d4..5a07b012f1f 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/sse.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/sse.ts @@ -11,6 +11,7 @@ import type { BridgeState, PromptCollector, + SessionBinding, SessionEventStream, } from './types.js'; @@ -43,29 +44,27 @@ export function createPromptCollector(): PromptCollector { * Start a persistent SSE subscription for a session. * Collects agent_message_chunk events into the active PromptCollector. */ -export function startEventStream(state: BridgeState, sessionId: string): void { - // Don't create duplicate streams; allow re-creation if the old stream is dead - if (state.eventStreams.has(sessionId)) { - const existing = state.eventStreams.get(sessionId)!; - if (!existing.abortCtrl.signal.aborted) return; - // Stale entry from a dead stream — clean up and re-create - state.eventStreams.delete(sessionId); - } - +export function createEventStream(sessionId: string): SessionEventStream { const abortCtrl = new AbortController(); - const stream: SessionEventStream = { + return { sessionId, abortCtrl, activeCollector: null, lastActivityMs: Date.now(), }; - state.eventStreams.set(sessionId, stream); +} +export function startEventStream( + state: BridgeState, + binding: SessionBinding, + onEnd: () => void, +): void { + const { sessionId, stream } = binding; // Start consuming SSE in the background (fire-and-forget) (async () => { try { for await (const event of state.client.subscribeEvents(sessionId, { - signal: abortCtrl.signal, + signal: stream.abortCtrl.signal, })) { const data = event.data as Record | undefined; if (!data) continue; @@ -123,14 +122,7 @@ export function startEventStream(state: BridgeState, sessionId: string): void { stream.activeCollector.interrupted = true; stream.activeCollector.resolve(); } - // Only delete if this is still our stream (not replaced by startEventStream) - if (state.eventStreams.get(sessionId) === stream) { - state.eventStreams.delete(sessionId); - // Clear defaultSessionId so callers get a clear error - if (state.defaultSessionId === sessionId) { - state.defaultSessionId = undefined; - } - } + onEnd(); } })(); } @@ -138,44 +130,11 @@ export function startEventStream(state: BridgeState, sessionId: string): void { /** * Stop the persistent SSE subscription for a session. */ -export function stopEventStream(state: BridgeState, sessionId: string): void { - const stream = state.eventStreams.get(sessionId); - if (stream) { - stream.abortCtrl.abort(); - // Resolve any pending collector so prompt doesn't hang - if (stream.activeCollector) { - stream.activeCollector.interrupted = true; - stream.activeCollector.resolve(); - } - state.eventStreams.delete(sessionId); +export function stopEventStream(stream: SessionEventStream): void { + stream.abortCtrl.abort(); + // Resolve any pending collector so prompt doesn't hang + if (stream.activeCollector) { + stream.activeCollector.interrupted = true; + stream.activeCollector.resolve(); } } - -/** Default session idle TTL: 30 minutes. */ -const SESSION_TTL_MS = 30 * 60 * 1000; -/** Cleanup interval: every 5 minutes. */ -const CLEANUP_INTERVAL_MS = 5 * 60 * 1000; - -/** - * Start a periodic cleanup timer that removes idle SSE streams. - * Returns a cleanup function to stop the timer (call on server shutdown). - */ -export function startSessionCleanup(state: BridgeState): () => void { - const timer = setInterval(() => { - const now = Date.now(); - for (const [sessionId, stream] of state.eventStreams) { - if (now - stream.lastActivityMs > SESSION_TTL_MS) { - process.stderr.write( - `[serve-bridge] Cleaning up idle session SSE: ${sessionId}\n`, - ); - stopEventStream(state, sessionId); - if (state.defaultSessionId === sessionId) { - state.defaultSessionId = undefined; - } - } - } - }, CLEANUP_INTERVAL_MS); - // Don't keep the process alive just for cleanup - timer.unref(); - return () => clearInterval(timer); -} diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts index 3fa9169002d..ed45374109d 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts @@ -7,9 +7,16 @@ import { z } from 'zod'; import { tool } from '../../tool.js'; import { formatJsonResult } from '../../formatters.js'; +import { isInvalidDaemonClientIdError } from '../../../daemon/DaemonHttpError.js'; import type { BridgeState } from '../types.js'; import { createPromptCollector } from '../sse.js'; import { handler, resolveSessionId } from '../helpers.js'; +import { + invalidateBinding, + resolveBinding, + rethrowBindingError, + trackLifecycle, +} from '../bindings.js'; /* eslint-disable @typescript-eslint/no-explicit-any */ export function agentTools(state: BridgeState): any[] { @@ -24,116 +31,156 @@ export function agentTools(state: BridgeState): any[] { .optional() .describe('Session ID. Uses default session if omitted.'), }, - handler(async (args) => { - const sessionId = resolveSessionId(state, args.session_id); + handler((args, { signal }) => + trackLifecycle(state, async () => { + const sessionId = resolveSessionId(state, args.session_id); + const binding = resolveBinding(state, sessionId); + const { stream } = binding; - // Use the persistent SSE stream established at session_create. - const stream = state.eventStreams.get(sessionId); - if (!stream) { - throw new Error( - 'No SSE stream for session. Was the session created via session_create?', - ); - } + // Guard against concurrent prompts on the same session + if (stream.activeCollector) { + throw new Error( + 'Another prompt is already in progress for this session. Wait for it to complete or call prompt_cancel first.', + ); + } - // Guard against concurrent prompts on the same session - if (stream.activeCollector) { - throw new Error( - 'Another prompt is already in progress for this session. Wait for it to complete or call prompt_cancel first.', - ); - } + // Install a new collector to capture this prompt's response chunks. + signal?.throwIfAborted(); + stream.lastActivityMs = Date.now(); + const collector = createPromptCollector(); + stream.activeCollector = collector; + let promptSettled = false; + let cancelPromise: Promise | undefined; + const cancelCaptured = (): Promise => { + cancelPromise ??= state.client.cancel(sessionId, binding.clientId); + return cancelPromise; + }; + const awaitBestEffortCancel = async (): Promise => { + try { + await cancelCaptured(); + return false; + } catch (err) { + if (isInvalidDaemonClientIdError(err)) { + return await invalidateBinding(state, binding); + } + return false; + } + }; + const onAbort = () => { + collector.interrupted = true; + collector.resolve(); + if (promptSettled) { + void awaitBestEffortCancel(); + } + }; + signal?.addEventListener('abort', onAbort, { once: true }); - // Install a new collector to capture this prompt's response chunks. - stream.lastActivityMs = Date.now(); - const collector = createPromptCollector(); - stream.activeCollector = collector; + try { + // Send prompt — response text arrives via the persistent SSE stream. + const result = await state.client + .prompt( + sessionId, + { prompt: [{ type: 'text', text: args.prompt }] }, + signal, + binding.clientId, + ) + .catch(async (err: unknown) => + rethrowBindingError(state, binding, err), + ); + promptSettled = true; + if (signal?.aborted) { + await awaitBestEffortCancel(); + } - try { - // Send prompt — response text arrives via the persistent SSE stream. - const result = await state.client.prompt(sessionId, { - prompt: [{ type: 'text', text: args.prompt }], - }); + // Wait for the collector to be resolved by _meta event (with timeout). + const COLLECT_TIMEOUT_MS = 30000; + let timedOut = false; + let timeoutId: ReturnType; + await Promise.race([ + collector.promise, + new Promise((r) => { + timeoutId = setTimeout(() => { + timedOut = true; + r(); + }, COLLECT_TIMEOUT_MS); + }), + ]); + clearTimeout(timeoutId!); - // Wait for the collector to be resolved by _meta event (with timeout). - const COLLECT_TIMEOUT_MS = 30000; - let timedOut = false; - let timeoutId: ReturnType; - await Promise.race([ - collector.promise, - new Promise((r) => { - timeoutId = setTimeout(() => { - timedOut = true; - r(); - }, COLLECT_TIMEOUT_MS); - }), - ]); - clearTimeout(timeoutId!); + // Guard against Promise.race microtask race: only treat as timeout + // if collector was NOT already resolved by _meta + if (timedOut && !collector.resolved) { + const bindingInvalidated = await awaitBestEffortCancel(); + const partialText = collector.texts.join(''); + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify( + { + session_id: sessionId, + stop_reason: 'timeout', + response: partialText || '(no text received)', + warning: + 'Agent response may be incomplete. _meta event not received within 30s.' + + (bindingInvalidated + ? ` Call session_resume with session_id "${sessionId}" before retrying.` + : ''), + }, + null, + 2, + ), + }, + ], + isError: true, + }; + } - // Guard against Promise.race microtask race: only treat as timeout - // if collector was NOT already resolved by _meta - if (timedOut && !collector.resolved) { - try { - await state.client.cancel(sessionId); - } catch { - /* best-effort */ + if (signal?.aborted) { + await awaitBestEffortCancel(); } - const partialText = collector.texts.join(''); - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify( - { - session_id: sessionId, - stop_reason: 'timeout', - response: partialText || '(no text received)', - warning: - 'Agent response may be incomplete. _meta event not received within 30s.', - }, - null, - 2, - ), - }, - ], - isError: true, - }; - } - // SSE disconnect or stopEventStream resolved the collector - if (collector.interrupted) { - return { - content: [ - { - type: 'text' as const, - text: JSON.stringify( - { - session_id: sessionId, - stop_reason: 'interrupted', - response: - collector.texts.join('') || '(no text received)', - warning: - 'SSE stream was closed before the response completed.', - }, - null, - 2, - ), - }, - ], - isError: true, - }; - } + // SSE disconnect or stopEventStream resolved the collector + if (collector.interrupted) { + return { + content: [ + { + type: 'text' as const, + text: JSON.stringify( + { + session_id: sessionId, + stop_reason: 'interrupted', + response: + collector.texts.join('') || '(no text received)', + warning: + 'SSE stream was closed before the response completed.', + }, + null, + 2, + ), + }, + ], + isError: true, + }; + } - const responseText = - collector.texts.join('') || '(task completed, no text output)'; - return formatJsonResult({ - session_id: sessionId, - stop_reason: result.stopReason, - response: responseText, - }); - } finally { - // Clear the collector regardless of outcome. - stream.activeCollector = null; - } - }), + const responseText = + collector.texts.join('') || '(task completed, no text output)'; + return formatJsonResult({ + session_id: sessionId, + stop_reason: result.stopReason, + response: responseText, + }); + } finally { + if (cancelPromise) { + await awaitBestEffortCancel(); + } + signal?.removeEventListener('abort', onAbort); + // Clear the collector regardless of outcome. + stream.activeCollector = null; + } + }), + ), ), tool( @@ -147,17 +194,18 @@ export function agentTools(state: BridgeState): any[] { }, handler(async (args) => { const sessionId = resolveSessionId(state, args.session_id); - const stream = state.eventStreams.get(sessionId); - // Best-effort cancel — must not prevent collector resolution + const binding = resolveBinding(state, sessionId); + const { stream } = binding; try { - await state.client.cancel(sessionId); - } catch { - /* best-effort */ - } - // Resolve active collector so the prompt handler returns immediately - if (stream?.activeCollector) { - stream.activeCollector.interrupted = true; - stream.activeCollector.resolve(); + await state.client.cancel(sessionId, binding.clientId); + } catch (err) { + await rethrowBindingError(state, binding, err); + } finally { + // Resolve active collector even when the daemon rejects cancellation. + if (stream.activeCollector) { + stream.activeCollector.interrupted = true; + stream.activeCollector.resolve(); + } } return formatJsonResult({ ok: true, sessionId }); }), diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/session.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/session.ts index 4336c4c5233..52d0cd1e0b4 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/session.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/session.ts @@ -8,7 +8,15 @@ import { z } from 'zod'; import { tool } from '../../tool.js'; import { formatJsonResult } from '../../formatters.js'; import type { BridgeState } from '../types.js'; -import { startEventStream, stopEventStream } from '../sse.js'; +import { + createBinding, + releaseBinding, + replaceBinding, + resolveBinding, + rethrowBindingError, + trackLifecycle, + withSessionLock, +} from '../bindings.js'; import { handler, resolveSessionId } from '../helpers.js'; /* eslint-disable @typescript-eslint/no-explicit-any */ @@ -31,24 +39,22 @@ export function sessionTools(state: BridgeState): any[] { .optional() .describe('Session scope.'), }, - handler(async (args) => { - const session = await state.client.createOrAttachSession({ - workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, - modelServiceId: args.model_service_id, - sessionScope: args.session_scope, - }); - // Stop old SSE only after new session is confirmed - if ( - state.defaultSessionId && - state.defaultSessionId !== session.sessionId - ) { - stopEventStream(state, state.defaultSessionId); - } - state.defaultSessionId = session.sessionId; - // Start persistent SSE connection for this session - startEventStream(state, session.sessionId); - return formatJsonResult(session); - }), + handler((args) => + trackLifecycle(state, async () => { + const session = await state.client.createOrAttachSession({ + workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, + modelServiceId: args.model_service_id, + sessionScope: args.session_scope, + }); + await withSessionLock(state, session.sessionId, async () => { + await replaceBinding( + state, + createBinding(session.sessionId, session.clientId), + ); + }); + return formatJsonResult(session); + }), + ), ), tool( @@ -58,21 +64,35 @@ export function sessionTools(state: BridgeState): any[] { session_id: z.string().describe('Session ID to restore.'), workspace_cwd: z.string().optional().describe('Workspace path.'), }, - handler(async (args) => { - const result = await state.client.loadSession(args.session_id, { - workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, - }); - // Stop old SSE only after load is confirmed - if ( - state.defaultSessionId && - state.defaultSessionId !== result.sessionId - ) { - stopEventStream(state, state.defaultSessionId); - } - state.defaultSessionId = result.sessionId; - startEventStream(state, result.sessionId); - return formatJsonResult(result); - }), + handler((args) => + trackLifecycle(state, async () => { + const result = await withSessionLock( + state, + args.session_id, + async () => { + const current = state.bindings.get(args.session_id); + if (current?.stream.activeCollector) { + throw new Error( + 'Cannot replace a session binding while a prompt is in progress.', + ); + } + const loaded = await state.client.loadSession( + args.session_id, + { + workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, + }, + current?.clientId, + ); + await replaceBinding( + state, + createBinding(loaded.sessionId, loaded.clientId), + ); + return loaded; + }, + ); + return formatJsonResult(result); + }), + ), ), tool( @@ -82,21 +102,35 @@ export function sessionTools(state: BridgeState): any[] { session_id: z.string().describe('Session ID to resume.'), workspace_cwd: z.string().optional().describe('Workspace path.'), }, - handler(async (args) => { - const result = await state.client.resumeSession(args.session_id, { - workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, - }); - // Stop old SSE only after resume is confirmed - if ( - state.defaultSessionId && - state.defaultSessionId !== result.sessionId - ) { - stopEventStream(state, state.defaultSessionId); - } - state.defaultSessionId = result.sessionId; - startEventStream(state, result.sessionId); - return formatJsonResult(result); - }), + handler((args) => + trackLifecycle(state, async () => { + const result = await withSessionLock( + state, + args.session_id, + async () => { + const current = state.bindings.get(args.session_id); + if (current?.stream.activeCollector) { + throw new Error( + 'Cannot replace a session binding while a prompt is in progress.', + ); + } + const resumed = await state.client.resumeSession( + args.session_id, + { + workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, + }, + current?.clientId, + ); + await replaceBinding( + state, + createBinding(resumed.sessionId, resumed.clientId), + ); + return resumed; + }, + ); + return formatJsonResult(result); + }), + ), ), tool( @@ -108,19 +142,21 @@ export function sessionTools(state: BridgeState): any[] { .optional() .describe('Session ID. Uses default session if omitted.'), }, - handler(async (args) => { - const sessionId = resolveSessionId(state, args.session_id); - try { - await state.client.closeSession(sessionId); - } finally { - // Always clean up SSE even if closeSession throws - stopEventStream(state, sessionId); - if (state.defaultSessionId === sessionId) { - state.defaultSessionId = undefined; - } - } - return formatJsonResult({ ok: true, sessionId }); - }), + handler((args) => + trackLifecycle(state, async () => { + const sessionId = resolveSessionId(state, args.session_id); + await withSessionLock(state, sessionId, async () => { + const binding = resolveBinding(state, sessionId); + try { + await state.client.closeSession(sessionId, binding.clientId); + } catch (err) { + await rethrowBindingError(state, binding, err); + } + await releaseBinding(state, binding, false); + }); + return formatJsonResult({ ok: true, sessionId }); + }), + ), ), tool( diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/types.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/types.ts index 4ee2066e9e3..546f28dd135 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/types.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/types.ts @@ -57,6 +57,15 @@ export interface SessionEventStream { lastActivityMs: number; } +/** One daemon attachment and its matching event stream. */ +export interface SessionBinding { + readonly sessionId: string; + readonly clientId: string | undefined; + readonly stream: SessionEventStream; + /** One-shot release latch. Identity fields remain immutable. */ + releasePromise?: Promise; +} + /** * Mutable bridge state shared across all tool handlers. */ @@ -68,8 +77,15 @@ export interface BridgeState { token: string | undefined; defaultSessionId: string | undefined; workspaceCwd: string | undefined; - /** Persistent SSE connections keyed by sessionId. */ - eventStreams: Map; + /** Live daemon attachments keyed by sessionId. */ + bindings: Map; + /** Promise tails for the few lifecycle operations that must be serialized. */ + sessionLocks: Map>; + /** In-flight session lifecycle operations that disposal must drain. */ + pendingLifecycles: Set>; + /** Releases that may outlive removal from `bindings`. */ + pendingReleases: Set>; + disposed: boolean; /** Whether global scope writes are allowed (default: false). */ allowGlobalScope: boolean; } diff --git a/packages/sdk-typescript/src/daemon/DaemonClient.ts b/packages/sdk-typescript/src/daemon/DaemonClient.ts index c2dc6286dcc..3b932149c6a 100644 --- a/packages/sdk-typescript/src/daemon/DaemonClient.ts +++ b/packages/sdk-typescript/src/daemon/DaemonClient.ts @@ -10,7 +10,10 @@ import { } from '@qwen-code/acp-bridge/mcpTimeouts'; import { CHANNEL_CONTROL_DEFAULT_TIMEOUT_MS } from '@qwen-code/acp-bridge/channelControlTimeouts'; import { DaemonAuthFlow } from './DaemonAuthFlow.js'; -import { DaemonHttpError } from './DaemonHttpError.js'; +import { + DaemonHttpError, + isInvalidDaemonClientIdError, +} from './DaemonHttpError.js'; import type { DaemonTransport } from './DaemonTransport.js'; import { RestSseTransport } from './RestSseTransport.js'; import type { @@ -216,6 +219,18 @@ export interface DaemonClientOptions { * owner routing and strict mutation authentication remain authoritative. */ transport?: DaemonTransport; + /** Default prompt ingress advisory sent to the daemon. */ + invocationIngress?: DaemonInvocationIngress; +} + +export type DaemonInvocationIngress = + | 'channel' + | 'scheduler' + | 'internal' + | 'external_mcp'; + +export interface DaemonPromptOptions { + invocationIngress?: DaemonInvocationIngress; } const DEFAULT_SESSION_LIST_PAGE_SIZE = 20; @@ -230,6 +245,7 @@ const DEFAULT_MAX_PENDING_PROMPTS_PER_SESSION = 5; const MCP_RESTART_DEFAULT_TIMEOUT_MS = MCP_RESTART_SERVER_DEADLINE_MS + MCP_RESTART_CLIENT_HEADROOM_MS; const CLIENT_ID_HEADER = 'X-Qwen-Client-Id'; +const INVOCATION_INGRESS_META_KEY = 'qwen-code/invocation-ingress'; const urlEncode = encodeURIComponent; function normalizePermissionRuleInput(rule: string): string { @@ -438,6 +454,7 @@ export class DaemonClient { private readonly _fetch: typeof globalThis.fetch; private readonly fetchTimeoutMs: number; private readonly promptLimit: number; + private readonly invocationIngress: DaemonInvocationIngress | 'daemon'; private readonly promptCounts: Record = Object.create(null); /** * Pluggable transport layer. Defaults to `RestSseTransport` when @@ -486,6 +503,7 @@ export class DaemonClient { this.promptLimit = normalizePendingPromptLimit( opts.maxPendingPromptsPerSession, ); + this.invocationIngress = opts.invocationIngress ?? 'daemon'; this.transport = opts.transport ?? new RestSseTransport(this.baseUrl, this.token, this._fetch); @@ -495,6 +513,16 @@ export class DaemonClient { return this.promptLimit; } + private withPromptIngress( + req: PromptRequest, + options?: DaemonPromptOptions, + ): PromptRequest { + const meta = req._meta ? { ...req._meta } : {}; + meta[INVOCATION_INGRESS_META_KEY] = + options?.invocationIngress ?? this.invocationIngress; + return { ...req, _meta: meta }; + } + /** @internal */ reservePromptSlot(sessionId: string, limit = this.promptLimit): () => void { if (limit === Infinity) return () => {}; @@ -3196,10 +3224,12 @@ export class DaemonClient { req: PromptRequest, signal?: AbortSignal, clientId?: string, + options?: DaemonPromptOptions, ): Promise { signal?.throwIfAborted(); + const stampedRequest = this.withPromptIngress(req, options); const releasePromptSlot = this.reservePromptSlot(sessionId); - let releaseOnExit = true; + let acceptedAsynchronousPrompt = false; try { const res = await this.transport.fetch( `${this.baseUrl}/session/${urlEncode(sessionId)}/prompt`, @@ -3209,25 +3239,20 @@ export class DaemonClient { { 'Content-Type': 'application/json' }, clientId, ), - body: JSON.stringify(req), + body: JSON.stringify(stampedRequest), signal, }, ); if (res.status === 202) { + acceptedAsynchronousPrompt = true; const accept = (await res.json()) as NonBlockingPromptAccepted; - releaseOnExit = false; - try { - return await this._awaitTurnComplete( - sessionId, - accept.promptId, - accept.lastEventId, - signal, - clientId, - ); - } finally { - releasePromptSlot(); - } + return await this._awaitTurnComplete( + sessionId, + accept.promptId, + accept.lastEventId, + signal, + ); } if (!res.ok) { @@ -3238,8 +3263,18 @@ export class DaemonClient { ); } return (await res.json()) as PromptResult; + } catch (err) { + if (acceptedAsynchronousPrompt && signal?.aborted) { + try { + await this.cancel(sessionId, clientId); + } catch (cancelError) { + if (isInvalidDaemonClientIdError(cancelError)) throw cancelError; + } + throw signal.reason ?? err; + } + throw err; } finally { - if (releaseOnExit) releasePromptSlot(); + releasePromptSlot(); } } @@ -3266,13 +3301,15 @@ export class DaemonClient { req: PromptRequest, signal?: AbortSignal, clientId?: string, + options?: DaemonPromptOptions, ): Promise { + const stampedRequest = this.withPromptIngress(req, options); const res = await this.transport.fetch( `${this.baseUrl}/session/${urlEncode(sessionId)}/prompt`, { method: 'POST', headers: this.headers({ 'Content-Type': 'application/json' }, clientId), - body: JSON.stringify(req), + body: JSON.stringify(stampedRequest), signal, }, ); @@ -3292,7 +3329,6 @@ export class DaemonClient { promptId: string, lastEventId: number, signal?: AbortSignal, - clientId?: string, ): Promise { const sseAbort = new AbortController(); const composedSignal = signal @@ -3309,16 +3345,6 @@ export class DaemonClient { if (result !== undefined) return result; } throw new Error('SSE stream ended'); - } catch (err) { - if ( - signal?.aborted && - err instanceof DOMException && - err.name === 'AbortError' - ) { - this.cancel(sessionId, clientId).catch(() => {}); - throw err; - } - throw err; } finally { if (!sseAbort.signal.aborted) sseAbort.abort(); } diff --git a/packages/sdk-typescript/src/daemon/DaemonHttpError.ts b/packages/sdk-typescript/src/daemon/DaemonHttpError.ts index d863b9374a1..1495b02cd03 100644 --- a/packages/sdk-typescript/src/daemon/DaemonHttpError.ts +++ b/packages/sdk-typescript/src/daemon/DaemonHttpError.ts @@ -23,3 +23,15 @@ export class DaemonHttpError extends Error { this.body = body; } } + +export function isInvalidDaemonClientIdError( + err: unknown, +): err is DaemonHttpError { + return ( + err instanceof DaemonHttpError && + err.status === 400 && + typeof err.body === 'object' && + err.body !== null && + (err.body as { code?: unknown }).code === 'invalid_client_id' + ); +} diff --git a/packages/sdk-typescript/src/daemon/DaemonSessionClient.ts b/packages/sdk-typescript/src/daemon/DaemonSessionClient.ts index c84ceeff912..af27eeefb02 100644 --- a/packages/sdk-typescript/src/daemon/DaemonSessionClient.ts +++ b/packages/sdk-typescript/src/daemon/DaemonSessionClient.ts @@ -5,13 +5,14 @@ */ import type { DaemonClient } from './DaemonClient.js'; -import { DaemonHttpError } from './DaemonHttpError.js'; +import { isInvalidDaemonClientIdError } from './DaemonHttpError.js'; import { isNonBlockingAccepted, matchTurnEvent, normalizePendingPromptLimit, type CreateSessionRequest, type NonBlockingPromptAccepted, + type DaemonPromptOptions, type PromptRequest, type RestoreSessionRequest, type SubscribeOptions, @@ -267,11 +268,12 @@ export class DaemonSessionClient { async prompt( req: PromptRequest, signal?: AbortSignal, + options?: DaemonPromptOptions, ): Promise { signal?.throwIfAborted(); if (!this.subscriptionActive) { return await this.withClientIdSelfHeal(() => - this.client.prompt(this.sessionId, req, signal, this.clientId), + this.client.prompt(this.sessionId, req, signal, this.clientId, options), ); } @@ -287,6 +289,7 @@ export class DaemonSessionClient { req, signal, this.clientId, + options, ), ); if (!isNonBlockingAccepted(accepted)) { @@ -343,10 +346,17 @@ export class DaemonSessionClient { async submitPrompt( req: PromptRequest, signal?: AbortSignal, + options?: DaemonPromptOptions, ): Promise { signal?.throwIfAborted(); const accepted = await this.withClientIdSelfHeal(() => - this.client.promptNonBlocking(this.sessionId, req, signal, this.clientId), + this.client.promptNonBlocking( + this.sessionId, + req, + signal, + this.clientId, + options, + ), ); if (!isNonBlockingAccepted(accepted)) { throw new Error('Expected non-blocking prompt acceptance'); @@ -369,7 +379,7 @@ export class DaemonSessionClient { try { return await fn(); } catch (err) { - if (!isInvalidClientId(err)) throw err; + if (!isInvalidDaemonClientIdError(err)) throw err; await this.reattach(); return await fn(); } @@ -775,17 +785,3 @@ function validateLastEventId( } return lastEventId; } - -/** - * True for the daemon's `400 invalid_client_id` prompt-admission rejection - * (the stale-clientId signal a daemon restart / session reload produces). - */ -function isInvalidClientId(err: unknown): boolean { - return ( - err instanceof DaemonHttpError && - err.status === 400 && - typeof err.body === 'object' && - err.body !== null && - (err.body as { code?: unknown }).code === 'invalid_client_id' - ); -} diff --git a/packages/sdk-typescript/src/daemon/index.ts b/packages/sdk-typescript/src/daemon/index.ts index 1ed142985a0..ad3c142b0b7 100644 --- a/packages/sdk-typescript/src/daemon/index.ts +++ b/packages/sdk-typescript/src/daemon/index.ts @@ -14,6 +14,8 @@ export { matchTurnEvent, type CreateSessionRequest, type DaemonClientOptions, + type DaemonInvocationIngress, + type DaemonPromptOptions, type DaemonTurnError, type NonBlockingPromptAccepted, type PromptRequest, diff --git a/packages/sdk-typescript/src/index.ts b/packages/sdk-typescript/src/index.ts index e6fa6e1fb23..708a289fb4b 100644 --- a/packages/sdk-typescript/src/index.ts +++ b/packages/sdk-typescript/src/index.ts @@ -100,6 +100,8 @@ export { type DaemonStateResyncRequiredData, type DaemonStateResyncRequiredEvent, type DaemonClientOptions, + type DaemonInvocationIngress, + type DaemonPromptOptions, type DaemonContentHash, type DaemonControlEvent, type DaemonEvent, @@ -324,7 +326,10 @@ export type { McpSdkServerConfigWithInstance, } from './daemon-mcp/createSdkMcpServer.js'; -export type { ServeBridgeMcpServerOptions } from './daemon-mcp/serve-bridge/index.js'; +export type { + ServeBridgeMcpServer, + ServeBridgeMcpServerOptions, +} from './daemon-mcp/serve-bridge/index.js'; export type { QueryOptions } from './query/createQuery.js'; export type { LogLevel, LoggerConfig, ScopedLogger } from './utils/logger.js'; diff --git a/packages/sdk-typescript/test/unit/DaemonClient.test.ts b/packages/sdk-typescript/test/unit/DaemonClient.test.ts index e17362274cc..d0a5d93dc62 100644 --- a/packages/sdk-typescript/test/unit/DaemonClient.test.ts +++ b/packages/sdk-typescript/test/unit/DaemonClient.test.ts @@ -1589,6 +1589,66 @@ describe('DaemonClient', () => { expect(calls[0]?.method).toBe('POST'); const body = JSON.parse(calls[0]!.body!); expect(body.prompt).toEqual([{ type: 'text', text: 'hi' }]); + expect(body._meta).toEqual({ + 'qwen-code/invocation-ingress': 'daemon', + }); + }); + + it('stamps the configured ingress on blocking and non-blocking prompts', async () => { + const { fetch, calls } = recordingFetch(() => + jsonResponse(200, { stopReason: 'end_turn' }), + ); + const client = new DaemonClient({ + baseUrl: 'http://daemon', + fetch, + invocationIngress: 'external_mcp', + }); + + const blockingRequest = { + prompt: [{ type: 'text' as const, text: 'blocking' }], + _meta: { keep: true }, + }; + await client.prompt('s-1', blockingRequest); + await client.promptNonBlocking('s-1', { + prompt: [{ type: 'text', text: 'non-blocking' }], + }); + + expect(JSON.parse(calls[0]!.body!)).toMatchObject({ + _meta: { + keep: true, + 'qwen-code/invocation-ingress': 'external_mcp', + }, + }); + expect(JSON.parse(calls[1]!.body!)).toMatchObject({ + _meta: { 'qwen-code/invocation-ingress': 'external_mcp' }, + }); + expect(blockingRequest._meta).toEqual({ keep: true }); + }); + + it('lets an explicit per-prompt ingress override caller meta and the client default', async () => { + const { fetch, calls } = recordingFetch(() => + jsonResponse(200, { stopReason: 'end_turn' }), + ); + const client = new DaemonClient({ + baseUrl: 'http://daemon', + fetch, + invocationIngress: 'channel', + }); + + await client.prompt( + 's-1', + { + prompt: [{ type: 'text', text: 'hi' }], + _meta: { 'qwen-code/invocation-ingress': 'internal' }, + }, + undefined, + undefined, + { invocationIngress: 'scheduler' }, + ); + + expect(JSON.parse(calls[0]!.body!)._meta).toEqual({ + 'qwen-code/invocation-ingress': 'scheduler', + }); }); it('url-encodes the sessionId', async () => { @@ -1621,14 +1681,22 @@ describe('DaemonClient', () => { // signal arg on `sendPrompt`; the SDK had the parameter wired // but no test, so a regression that dropped it on the floor // would silently leave callers unable to cancel. - const fetch = vi.fn( - (_input: RequestInfo | URL, init?: RequestInit) => - new Promise((_res, rej) => { - init?.signal?.addEventListener('abort', () => - rej(new DOMException('aborted', 'AbortError')), - ); - }), - ) as unknown as typeof globalThis.fetch; + const fetch = vi.fn((input: RequestInfo | URL, init?: RequestInit) => { + const url = + typeof input === 'string' + ? input + : input instanceof URL + ? input.toString() + : input.url; + if (url.endsWith('/cancel')) { + return Promise.resolve(new Response(null, { status: 204 })); + } + return new Promise((_res, rej) => { + init?.signal?.addEventListener('abort', () => + rej(new DOMException('aborted', 'AbortError')), + ); + }); + }) as unknown as typeof globalThis.fetch; const client = new DaemonClient({ baseUrl: 'http://daemon', fetch }); const ctrl = new AbortController(); setTimeout(() => ctrl.abort(), 30); @@ -2056,6 +2124,113 @@ describe('DaemonClient', () => { ).resolves.toEqual({ stopReason: 'end_turn' }); expect(calls.filter((c) => c.url.endsWith('/prompt'))).toHaveLength(2); }); + + it('does not cancel an unadmitted legacy prompt after caller abort', async () => { + let promptCount = 0; + const { fetch, calls } = recordingFetch((req) => { + if (req.url.endsWith('/session/s-1/prompt')) { + promptCount++; + if (promptCount > 1) { + return jsonResponse(200, { stopReason: 'end_turn' }); + } + return new Promise((_resolve, reject) => { + req.signal?.addEventListener( + 'abort', + () => reject(req.signal?.reason), + { once: true }, + ); + }); + } + if (req.url.endsWith('/session/s-1/cancel')) { + return new Response(null, { status: 204 }); + } + return jsonResponse(500, { error: `unexpected ${req.url}` }); + }); + const client = new DaemonClient({ + baseUrl: 'http://daemon', + fetch, + maxPendingPromptsPerSession: 1, + }); + const ctrl = new AbortController(); + const pending = client + .prompt( + 's-1', + { prompt: [{ type: 'text', text: 'first' }] }, + ctrl.signal, + ) + .catch((err: unknown) => err); + await vi.waitFor(() => { + expect(calls.filter((c) => c.url.endsWith('/prompt'))).toHaveLength(1); + }); + + ctrl.abort(); + const result = await pending; + + expect(result).toBeInstanceOf(DOMException); + expect((result as DOMException).name).toBe('AbortError'); + expect(calls.filter((c) => c.url.endsWith('/cancel'))).toHaveLength(0); + await expect( + client.prompt('s-1', { + prompt: [{ type: 'text', text: 'next' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + }); + + it('waits for a transient 202 cancellation failure before rejecting caller abort', async () => { + let rejectCancel!: (reason?: unknown) => void; + const { fetch, calls } = recordingFetch((req) => { + if (req.url.endsWith('/session/s-1/prompt')) { + return jsonResponse(202, { promptId: 'p-1', lastEventId: 0 }); + } + if (req.url.endsWith('/session/s-1/events')) { + return pendingSseResponse(); + } + if (req.url.endsWith('/session/s-1/cancel')) { + return new Promise((_resolve, reject) => { + rejectCancel = reject; + }); + } + return jsonResponse(500, { error: `unexpected ${req.url}` }); + }); + const client = new DaemonClient({ + baseUrl: 'http://daemon', + fetch, + maxPendingPromptsPerSession: 1, + }); + const ctrl = new AbortController(); + let settled = false; + const pending = client + .prompt( + 's-1', + { prompt: [{ type: 'text', text: 'first' }] }, + ctrl.signal, + ) + .catch((err: unknown) => err) + .finally(() => { + settled = true; + }); + await vi.waitFor(() => { + expect(calls.filter((c) => c.url.endsWith('/events'))).toHaveLength(1); + }); + + ctrl.abort(); + await vi.waitFor(() => expect(rejectCancel).toBeTypeOf('function')); + expect(calls.filter((c) => c.url.endsWith('/cancel'))).toHaveLength(1); + expect(settled).toBe(false); + await expect( + client.prompt('s-1', { + prompt: [{ type: 'text', text: 'too early' }], + }), + ).rejects.toBeInstanceOf(DaemonPendingPromptLimitError); + expect(calls.filter((c) => c.url.endsWith('/prompt'))).toHaveLength(1); + + rejectCancel(new Error('cancel network failed')); + const result = await pending; + + expect(result).toBeInstanceOf(DOMException); + expect((result as DOMException).name).toBe('AbortError'); + expect(calls.filter((c) => c.url.endsWith('/cancel'))).toHaveLength(1); + }); }); describe('loadSession / resumeSession', () => { diff --git a/packages/sdk-typescript/test/unit/serve-bridge-shutdown.test.ts b/packages/sdk-typescript/test/unit/serve-bridge-shutdown.test.ts new file mode 100644 index 00000000000..383a41d26df --- /dev/null +++ b/packages/sdk-typescript/test/unit/serve-bridge-shutdown.test.ts @@ -0,0 +1,76 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it, vi } from 'vitest'; +import { createShutdownController } from '../../src/daemon-mcp/serve-bridge/shutdown.js'; + +describe('serve-bridge shutdown', () => { + it('shares recursive shutdown and waits for close and disposal before exit', async () => { + let resolveDispose!: () => void; + let recursiveShutdown: Promise | undefined; + const close = vi.fn(async () => { + controller.markInstanceClosed(); + recursiveShutdown = controller.shutdown(); + }); + const dispose = vi.fn( + () => + new Promise((resolve) => { + resolveDispose = resolve; + }), + ); + const exit = vi.fn(); + const controller = createShutdownController({ + close, + dispose, + exit, + reportCloseError: vi.fn(), + }); + + const shutdown = controller.shutdown(); + await vi.waitFor(() => expect(dispose).toHaveBeenCalledOnce()); + + expect(recursiveShutdown).toBe(shutdown); + expect(close).toHaveBeenCalledOnce(); + expect(exit).not.toHaveBeenCalled(); + resolveDispose(); + await shutdown; + + expect(dispose).toHaveBeenCalledOnce(); + expect(exit).toHaveBeenCalledWith(0); + }); + + it('skips duplicate close after peer shutdown and shares disposal', async () => { + let resolveDispose!: () => void; + const close = vi.fn(async () => undefined); + const dispose = vi.fn( + () => + new Promise((resolve) => { + resolveDispose = resolve; + }), + ); + const exit = vi.fn(); + const controller = createShutdownController({ + close, + dispose, + exit, + reportCloseError: vi.fn(), + }); + controller.markInstanceClosed(); + + const first = controller.shutdown(); + const second = controller.shutdown(); + await vi.waitFor(() => expect(dispose).toHaveBeenCalledOnce()); + + expect(second).toBe(first); + expect(close).not.toHaveBeenCalled(); + expect(exit).not.toHaveBeenCalled(); + resolveDispose(); + await Promise.all([first, second]); + + expect(dispose).toHaveBeenCalledOnce(); + expect(exit).toHaveBeenCalledWith(0); + }); +}); diff --git a/packages/sdk-typescript/test/unit/serve-bridge.test.ts b/packages/sdk-typescript/test/unit/serve-bridge.test.ts index 31a7b76267f..92aae651b36 100644 --- a/packages/sdk-typescript/test/unit/serve-bridge.test.ts +++ b/packages/sdk-typescript/test/unit/serve-bridge.test.ts @@ -17,10 +17,14 @@ import { resolveSessionId, handler, } from '../../src/daemon-mcp/serve-bridge/helpers.js'; -import type { - BridgeState, - SessionEventStream, -} from '../../src/daemon-mcp/serve-bridge/types.js'; +import type { BridgeState } from '../../src/daemon-mcp/serve-bridge/types.js'; +import { + createBinding, + disposeBindings, + replaceBinding, + startSessionCleanup, +} from '../../src/daemon-mcp/serve-bridge/bindings.js'; +import { createPromptCollector } from '../../src/daemon-mcp/serve-bridge/sse.js'; import { DaemonClient } from '../../src/daemon/DaemonClient.js'; // --- Helpers --- @@ -37,6 +41,7 @@ interface CapturedRequest { method: string; headers: Record; body: string | null; + signal: AbortSignal | null; } function recordingFetch( @@ -58,7 +63,13 @@ function recordingFetch( h.forEach((v, k) => (headers[k.toLowerCase()] = v)); } const body = typeof init?.body === 'string' ? init.body : null; - const captured: CapturedRequest = { url, method, headers, body }; + const captured: CapturedRequest = { + url, + method, + headers, + body, + signal: init?.signal ?? null, + }; calls.push(captured); return reply(captured); }, @@ -80,18 +91,44 @@ function makeMockState(opts?: { baseUrl: 'http://127.0.0.1:4170', token, fetch, + invocationIngress: 'external_mcp', }), daemonUrl: 'http://127.0.0.1:4170', token, defaultSessionId: opts?.defaultSessionId, workspaceCwd: '/tmp/test-workspace', - eventStreams: new Map(), + bindings: new Map(), + sessionLocks: new Map(), + pendingLifecycles: new Set(), + pendingReleases: new Set(), + disposed: false, allowGlobalScope: false, }; + vi.spyOn(state.client, 'subscribeEvents').mockImplementation( + async function* (_sessionId, subscribeOpts) { + await new Promise((resolve) => { + if (subscribeOpts.signal?.aborted) { + resolve(); + } else { + subscribeOpts.signal?.addEventListener('abort', () => resolve(), { + once: true, + }); + } + }); + yield* []; + }, + ); + return { state, calls }; } +function bindSession(state: BridgeState, sessionId: string, clientId?: string) { + const binding = createBinding(sessionId, clientId); + state.bindings.set(sessionId, binding); + return binding; +} + // --- Tests --- describe('serve-bridge', () => { @@ -106,6 +143,7 @@ describe('serve-bridge', () => { expect(server).toBeDefined(); expect(server.name).toBe('qwen-serve-bridge'); expect(server.instance).toBeDefined(); + expect(server.dispose).toBeTypeOf('function'); }); it('should strip trailing slashes from daemonUrl', () => { @@ -115,6 +153,17 @@ describe('serve-bridge', () => { }); expect(server).toBeDefined(); }); + + it('should return one idempotent dispose promise', async () => { + const server = createServeBridgeMcpServer({ + daemonUrl: 'http://127.0.0.1:4170', + token: 'test', + }); + + const first = server.dispose(); + expect(server.dispose()).toBe(first); + await first; + }); }); describe('resolveSessionId', () => { @@ -143,7 +192,7 @@ describe('serve-bridge', () => { .mockResolvedValue({ content: [{ type: 'text', text: 'ok' }] }); const wrapped = handler(fn); const result = await wrapped({ foo: 'bar' }, {}); - expect(fn).toHaveBeenCalledWith({ foo: 'bar' }); + expect(fn).toHaveBeenCalledWith({ foo: 'bar' }, { signal: undefined }); expect(result).toEqual({ content: [{ type: 'text', text: 'ok' }] }); }); @@ -219,6 +268,208 @@ describe('serve-bridge', () => { expect(result.content[0].text).toContain('new-session-id'); expect(state.defaultSessionId).toBe('new-session-id'); }); + + it('should retain bindings for multiple created sessions', async () => { + let createCount = 0; + const { state, calls } = makeMockState({ + fetchReply: (req) => { + if (req.url.endsWith('/session') && req.method === 'POST') { + createCount++; + return jsonResponse(200, { + sessionId: createCount === 1 ? 'session-a' : 'session-b', + clientId: createCount === 1 ? 'client-a' : 'client-b', + workspaceCwd: '/tmp', + attached: false, + }); + } + return new Response(null, { status: 204 }); + }, + }); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const createTool = sessionTools(state).find( + (t: { name: string }) => t.name === 'session_create', + ); + + await createTool.handler({}, {}); + await createTool.handler({}, {}); + + expect(state.defaultSessionId).toBe('session-b'); + expect([...state.bindings.keys()]).toEqual(['session-a', 'session-b']); + expect(state.bindings.get('session-a')?.clientId).toBe('client-a'); + expect(state.bindings.get('session-b')?.clientId).toBe('client-b'); + expect(calls.some((call) => call.url.endsWith('/detach'))).toBe(false); + }); + + it('should discard a create acquisition when its session has an active prompt', async () => { + const { state, calls } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: (req) => { + if (req.url.endsWith('/session') && req.method === 'POST') { + return jsonResponse(200, { + sessionId: 'session-a', + clientId: 'client-new', + workspaceCwd: '/tmp', + attached: true, + }); + } + return new Response(null, { status: 204 }); + }, + }); + const existing = bindSession(state, 'session-a', 'client-old'); + existing.stream.activeCollector = createPromptCollector(); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const createTool = sessionTools(state).find( + (t: { name: string }) => t.name === 'session_create', + ); + + const result = await createTool.handler({}, {}); + + expect(result.isError).toBe(true); + expect(state.bindings.get('session-a')).toBe(existing); + expect(state.defaultSessionId).toBe('session-a'); + const detachCalls = calls.filter((call) => + call.url.endsWith('/session/session-a/detach'), + ); + expect(detachCalls).toHaveLength(1); + expect(detachCalls[0]?.headers['x-qwen-client-id']).toBe('client-new'); + }); + + it('should keep disposal pending until an in-flight create is released', async () => { + let resolveCreate!: (response: Response) => void; + const { state, calls } = makeMockState({ + fetchReply: (req) => { + if (req.url.endsWith('/session') && req.method === 'POST') { + return new Promise((resolve) => { + resolveCreate = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const createTool = sessionTools(state).find( + (t: { name: string }) => t.name === 'session_create', + ); + + const creating = createTool.handler({}, {}); + await vi.waitFor(() => expect(resolveCreate).toBeTypeOf('function')); + let disposed = false; + const disposing = disposeBindings(state).finally(() => { + disposed = true; + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(disposed).toBe(false); + + resolveCreate( + jsonResponse(200, { + sessionId: 'session-a', + clientId: 'client-new', + workspaceCwd: '/tmp', + attached: false, + }), + ); + const result = await creating; + await disposing; + + expect(result.isError).toBe(true); + expect(result.content[0].text).toContain('shutting down'); + expect(state.bindings.size).toBe(0); + const detachCalls = calls.filter((call) => + call.url.endsWith('/detach'), + ); + expect(detachCalls).toHaveLength(1); + expect(detachCalls[0]?.headers['x-qwen-client-id']).toBe('client-new'); + }); + + it('should reject a create started after disposal without acquiring a session', async () => { + const { state, calls } = makeMockState({ + fetchReply: () => + jsonResponse(200, { + sessionId: 'unexpected', + clientId: 'unexpected-client', + workspaceCwd: '/tmp', + attached: false, + }), + }); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const createTool = sessionTools(state).find( + (t: { name: string }) => t.name === 'session_create', + ); + await disposeBindings(state); + + const result = await createTool.handler({}, {}); + + expect(result.isError).toBe(true); + expect(result.content[0].text).toContain('shutting down'); + expect(calls).toHaveLength(0); + }); + }); + + describe('session_resume', () => { + it('should serialize replacements and detach every prior acquisition', async () => { + let resolveFirst!: (response: Response) => void; + let resumeCount = 0; + const { state, calls } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: (req) => { + if (req.url.endsWith('/resume')) { + resumeCount++; + if (resumeCount === 1) { + return new Promise((resolve) => { + resolveFirst = resolve; + }); + } + return jsonResponse(200, { + sessionId: 'session-a', + clientId: 'client-old', + workspaceCwd: '/tmp', + }); + } + return new Response(null, { status: 204 }); + }, + }); + await replaceBinding(state, createBinding('session-a', 'client-old')); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const resumeTool = sessionTools(state).find( + (t: { name: string }) => t.name === 'session_resume', + ); + + const first = resumeTool.handler({ session_id: 'session-a' }, {}); + const second = resumeTool.handler({ session_id: 'session-a' }, {}); + await vi.waitFor(() => expect(resumeCount).toBe(1)); + resolveFirst( + jsonResponse(200, { + sessionId: 'session-a', + clientId: 'client-old', + workspaceCwd: '/tmp', + }), + ); + await first; + await second; + + const resumeCalls = calls.filter((call) => + call.url.endsWith('/resume'), + ); + expect(resumeCalls).toHaveLength(2); + expect(resumeCalls[0]?.headers['x-qwen-client-id']).toBe('client-old'); + expect(resumeCalls[1]?.headers['x-qwen-client-id']).toBe('client-old'); + const detachIds = calls + .filter((call) => call.url.endsWith('/detach')) + .map((call) => call.headers['x-qwen-client-id']); + expect(detachIds).toEqual(['client-old', 'client-old']); + expect(state.bindings.get('session-a')?.clientId).toBe('client-old'); + expect(state.defaultSessionId).toBe('session-a'); + }); }); describe('session_close', () => { @@ -227,6 +478,10 @@ describe('serve-bridge', () => { defaultSessionId: 'sess-to-close', fetchReply: () => new Response(null, { status: 204 }), }); + state.bindings.set( + 'sess-to-close', + createBinding('sess-to-close', 'client-close'), + ); const { sessionTools } = await import( '../../src/daemon-mcp/serve-bridge/tools/session.js' @@ -243,8 +498,12 @@ describe('serve-bridge', () => { it('should not clear defaultSessionId when closing a different session', async () => { const { state } = makeMockState({ defaultSessionId: 'keep-this', - fetchReply: () => new Response(null, { status: 204 }), + fetchReply: () => jsonResponse(404, { code: 'not_found' }), }); + state.bindings.set( + 'other-session', + createBinding('other-session', 'client-other'), + ); const { sessionTools } = await import( '../../src/daemon-mcp/serve-bridge/tools/session.js' @@ -256,6 +515,133 @@ describe('serve-bridge', () => { await closeTool.handler({ session_id: 'other-session' }, {}); expect(state.defaultSessionId).toBe('keep-this'); + expect(state.bindings.has('other-session')).toBe(false); + }); + + it('should not detach when the SSE stream ends during close', async () => { + let resolveClose!: (response: Response) => void; + let endStream!: () => void; + const { state, calls } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: (req) => { + if (req.method === 'DELETE') { + return new Promise((resolve) => { + resolveClose = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + vi.spyOn(state.client, 'subscribeEvents').mockImplementation( + async function* () { + await new Promise((resolve) => { + endStream = resolve; + }); + yield* []; + }, + ); + await replaceBinding(state, createBinding('session-a', 'client-old')); + await vi.waitFor(() => expect(endStream).toBeTypeOf('function')); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const closeTool = sessionTools(state).find( + (t: { name: string }) => t.name === 'session_close', + ); + + const pending = closeTool.handler({ session_id: 'session-a' }, {}); + await vi.waitFor(() => expect(resolveClose).toBeTypeOf('function')); + endStream(); + await new Promise((resolve) => setTimeout(resolve, 0)); + resolveClose(new Response(null, { status: 204 })); + await pending; + + expect( + calls.filter((call) => call.url.endsWith('/detach')), + ).toHaveLength(0); + }); + + it('should invalidate the current binding on invalid_client_id', async () => { + const { state, calls } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: (req) => + req.method === 'DELETE' + ? jsonResponse(400, { code: 'invalid_client_id' }) + : new Response(null, { status: 204 }), + }); + const binding = bindSession(state, 'session-a', 'client-old'); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const closeTool = sessionTools(state).find( + (t: { name: string }) => t.name === 'session_close', + ); + + const result = await closeTool.handler({ session_id: 'session-a' }, {}); + + expect(result.isError).toBe(true); + expect(result.content[0].text).toContain('Call session_resume'); + expect(state.bindings.has('session-a')).toBe(false); + expect(binding.stream.abortCtrl.signal.aborted).toBe(true); + expect( + calls.filter((call) => call.url.endsWith('/detach')), + ).toHaveLength(0); + }); + + it('should preserve the binding on a transient close failure', async () => { + const { state } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: (req) => + req.method === 'DELETE' + ? jsonResponse(500, { code: 'internal_error' }) + : new Response(null, { status: 204 }), + }); + const binding = bindSession(state, 'session-a', 'client-old'); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const closeTool = sessionTools(state).find( + (t: { name: string }) => t.name === 'session_close', + ); + + const result = await closeTool.handler({ session_id: 'session-a' }, {}); + + expect(result.isError).toBe(true); + expect(state.bindings.get('session-a')).toBe(binding); + expect(binding.stream.abortCtrl.signal.aborted).toBe(false); + }); + + it('should preserve a newer binding on a stale close rejection', async () => { + let resolveClose!: (response: Response) => void; + const { state } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: (req) => { + if (req.method === 'DELETE') { + return new Promise((resolve) => { + resolveClose = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'session-a', 'client-old'); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const closeTool = sessionTools(state).find( + (t: { name: string }) => t.name === 'session_close', + ); + + const pending = closeTool.handler({ session_id: 'session-a' }, {}); + await vi.waitFor(() => expect(resolveClose).toBeTypeOf('function')); + const newer = createBinding('session-a', 'client-new'); + state.bindings.set('session-a', newer); + resolveClose(jsonResponse(400, { code: 'invalid_client_id' })); + const result = await pending; + + expect(result.isError).toBe(true); + expect(result.content[0].text).toContain('Retry the request'); + expect(state.bindings.get('session-a')).toBe(newer); }); }); @@ -335,7 +721,7 @@ describe('serve-bridge', () => { describe('prompt tool with persistent SSE', () => { it('should collect response text via the persistent event stream', async () => { - const { state } = makeMockState({ + const { state, calls } = makeMockState({ defaultSessionId: 'test-session', fetchReply: (req) => { if (req.url.includes('/prompt')) { @@ -343,7 +729,7 @@ describe('serve-bridge', () => { // persistent SSE stream will have populated the collector. // We simulate this by filling the collector just before the // prompt response resolves. - const stream = state.eventStreams.get('test-session')!; + const stream = state.bindings.get('test-session')!.stream; const collector = stream.activeCollector!; collector.texts.push('hello'); collector.texts.push(' world'); @@ -354,14 +740,9 @@ describe('serve-bridge', () => { }, }); - // Set up a fake persistent event stream (normally created by session_create) - const fakeStream: SessionEventStream = { - sessionId: 'test-session', - abortCtrl: new AbortController(), - activeCollector: null, - lastActivityMs: Date.now(), - }; - state.eventStreams.set('test-session', fakeStream); + const binding = createBinding('test-session', 'client-test'); + const fakeStream = binding.stream; + state.bindings.set('test-session', binding); const { agentTools } = await import( '../../src/daemon-mcp/serve-bridge/tools/agent.js' @@ -377,6 +758,11 @@ describe('serve-bridge', () => { expect(parsed.stop_reason).toBe('end_turn'); expect(parsed.session_id).toBe('test-session'); expect(parsed.response).toBe('hello world'); + expect( + calls.find((call) => call.url.endsWith('/prompt'))?.headers[ + 'x-qwen-client-id' + ], + ).toBe('client-test'); // Collector should be cleared after prompt completes expect(fakeStream.activeCollector).toBeNull(); }); @@ -396,7 +782,7 @@ describe('serve-bridge', () => { const result = await promptTool.handler({ prompt: 'test' }, {}); expect(result.isError).toBe(true); - expect(result.content[0].text).toContain('No SSE stream'); + expect(result.content[0].text).toContain('No active binding'); }); it('should reject concurrent prompts on the same session', async () => { @@ -408,13 +794,10 @@ describe('serve-bridge', () => { const { createPromptCollector } = await import( '../../src/daemon-mcp/serve-bridge/sse.js' ); - const fakeStream: SessionEventStream = { - sessionId: 'test-session', - abortCtrl: new AbortController(), - activeCollector: createPromptCollector(), // already has an active collector - lastActivityMs: Date.now(), - }; - state.eventStreams.set('test-session', fakeStream); + const binding = createBinding('test-session', 'client-test'); + const fakeStream = binding.stream; + fakeStream.activeCollector = createPromptCollector(); + state.bindings.set('test-session', binding); const { agentTools } = await import( '../../src/daemon-mcp/serve-bridge/tools/agent.js' @@ -439,13 +822,10 @@ describe('serve-bridge', () => { '../../src/daemon-mcp/serve-bridge/sse.js' ); const collector = createPromptCollector(); - const fakeStream: SessionEventStream = { - sessionId: 'test-session', - abortCtrl: new AbortController(), - activeCollector: collector, - lastActivityMs: Date.now(), - }; - state.eventStreams.set('test-session', fakeStream); + const binding = createBinding('test-session', 'client-test'); + const fakeStream = binding.stream; + fakeStream.activeCollector = collector; + state.bindings.set('test-session', binding); const { agentTools } = await import( '../../src/daemon-mcp/serve-bridge/tools/agent.js' @@ -459,6 +839,609 @@ describe('serve-bridge', () => { expect(collector.resolved).toBe(true); expect(collector.interrupted).toBe(true); }); + + it('should keep legacy anonymous bindings usable', async () => { + const { state, calls } = makeMockState({ + defaultSessionId: 'legacy-session', + fetchReply: (req) => { + if (req.url.endsWith('/prompt')) { + const collector = + state.bindings.get('legacy-session')!.stream.activeCollector!; + collector.resolve(); + return jsonResponse(200, { stopReason: 'end_turn' }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'legacy-session'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + const result = await promptTool.handler({ prompt: 'test' }, {}); + + expect(result.isError).toBeUndefined(); + expect( + calls.find((call) => call.url.endsWith('/prompt'))?.headers[ + 'x-qwen-client-id' + ], + ).toBeUndefined(); + }); + + it('should not cancel an unadmitted legacy prompt after abort', async () => { + const controller = new AbortController(); + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/prompt')) { + return new Promise((_resolve, reject) => { + req.signal?.addEventListener( + 'abort', + () => reject(req.signal?.reason), + { once: true }, + ); + }); + } + if (req.url.endsWith('/cancel')) { + return new Response(null, { status: 204 }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'test-session', 'client-test'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + const pending = promptTool.handler( + { prompt: 'test' }, + { signal: controller.signal }, + ); + await vi.waitFor(() => + expect(calls.some((call) => call.url.endsWith('/prompt'))).toBe(true), + ); + controller.abort(); + const result = await pending; + + expect(result.isError).toBe(true); + expect(result.content[0].text).toContain('aborted'); + expect(state.bindings.has('test-session')).toBe(true); + expect(calls.filter((call) => call.url.endsWith('/cancel'))).toHaveLength( + 0, + ); + }); + + it('should wait for active prompt cancellation before detaching on disposal', async () => { + let resolveCancel!: (response: Response) => void; + const controller = new AbortController(); + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/prompt')) { + return jsonResponse(200, { stopReason: 'end_turn' }); + } + if (req.url.endsWith('/cancel')) { + return new Promise((resolve) => { + resolveCancel = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'test-session', 'client-test'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + const prompting = promptTool.handler( + { prompt: 'test' }, + { signal: controller.signal }, + ); + await vi.waitFor(() => + expect( + state.bindings.get('test-session')?.stream.activeCollector, + ).not.toBeNull(), + ); + controller.abort(); + await vi.waitFor(() => expect(resolveCancel).toBeTypeOf('function')); + + let disposalSettled = false; + const disposing = disposeBindings(state).finally(() => { + disposalSettled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + const detachedBeforeCancel = calls.some((call) => + call.url.endsWith('/detach'), + ); + const settledBeforeCancel = disposalSettled; + + resolveCancel(new Response(null, { status: 204 })); + await Promise.all([prompting, disposing]); + + expect(detachedBeforeCancel).toBe(false); + expect(settledBeforeCancel).toBe(false); + expect(calls.filter((call) => call.url.endsWith('/detach'))).toHaveLength( + 1, + ); + }); + + it('should invalidate the current binding when abort cancellation is rejected', async () => { + const controller = new AbortController(); + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/prompt')) { + return jsonResponse(202, { + promptId: 'prompt-accepted', + lastEventId: 0, + }); + } + if (req.url.endsWith('/events')) { + return pendingSseResponse(); + } + if (req.url.endsWith('/cancel')) { + return jsonResponse(400, { code: 'invalid_client_id' }); + } + return new Response(null, { status: 204 }); + }, + }); + const binding = bindSession(state, 'test-session', 'client-stale'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + const pending = promptTool.handler( + { prompt: 'test' }, + { signal: controller.signal }, + ); + await vi.waitFor(() => + expect(state.client.subscribeEvents).toHaveBeenCalled(), + ); + controller.abort(); + const result = await pending; + + expect(result.isError).toBe(true); + expect(result.content[0].text).toContain('session_resume'); + expect(state.bindings.has('test-session')).toBe(false); + expect(binding.stream.abortCtrl.signal.aborted).toBe(true); + expect(calls.filter((call) => call.url.endsWith('/cancel'))).toHaveLength( + 1, + ); + }); + + it('should cancel an aborted settled prompt exactly once with its captured binding', async () => { + const controller = new AbortController(); + let resolveCancel!: (response: Response) => void; + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/prompt')) { + return jsonResponse(200, { stopReason: 'end_turn' }); + } + if (req.url.endsWith('/cancel')) { + return new Promise((resolve) => { + resolveCancel = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'test-session', 'client-test'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + let settled = false; + const pending = promptTool + .handler( + { prompt: 'test' }, + { signal: controller.signal, _meta: { forged: true } }, + ) + .finally(() => { + settled = true; + }); + await vi.waitFor(() => + expect(calls.some((call) => call.url.endsWith('/prompt'))).toBe(true), + ); + controller.abort(); + await vi.waitFor(() => expect(resolveCancel).toBeTypeOf('function')); + expect(settled).toBe(false); + expect(calls.filter((call) => call.url.endsWith('/cancel'))).toHaveLength( + 1, + ); + + resolveCancel(new Response(null, { status: 204 })); + const result = await pending; + + expect(result.isError).toBe(true); + const cancelCalls = calls.filter((call) => call.url.endsWith('/cancel')); + expect(cancelCalls).toHaveLength(1); + expect(cancelCalls[0]?.headers['x-qwen-client-id']).toBe('client-test'); + }); + + it('should preserve 202 SSE abort after a transient cancellation failure', async () => { + const controller = new AbortController(); + let promptCount = 0; + let rejectCancel!: (reason?: unknown) => void; + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/prompt')) { + promptCount++; + if (promptCount === 1) { + return jsonResponse(202, { + promptId: 'prompt-1', + lastEventId: 0, + }); + } + const collector = + state.bindings.get('test-session')!.stream.activeCollector!; + collector.texts.push('second response'); + collector.resolve(); + return jsonResponse(200, { stopReason: 'end_turn' }); + } + if (req.url.endsWith('/cancel')) { + return new Promise((_resolve, reject) => { + rejectCancel = reject; + }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'test-session', 'client-test'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + let settled = false; + const first = promptTool + .handler({ prompt: 'first' }, { signal: controller.signal }) + .finally(() => { + settled = true; + }); + await vi.waitFor(() => + expect( + calls.filter((call) => call.url.endsWith('/prompt')), + ).toHaveLength(1), + ); + controller.abort(); + await vi.waitFor(() => expect(rejectCancel).toBeTypeOf('function')); + + expect(settled).toBe(false); + expect(calls.filter((call) => call.url.endsWith('/cancel'))).toHaveLength( + 1, + ); + rejectCancel(new Error('cancel network failed')); + const firstResult = await first; + expect(firstResult.isError).toBe(true); + expect(firstResult.content[0].text).toContain('aborted'); + + const secondResult = await promptTool.handler({ prompt: 'second' }, {}); + + expect(secondResult.isError).toBeUndefined(); + expect(JSON.parse(secondResult.content[0].text).response).toBe( + 'second response', + ); + expect(calls.filter((call) => call.url.endsWith('/cancel'))).toHaveLength( + 1, + ); + }); + + it('should use the captured binding for timeout cancellation', async () => { + vi.useFakeTimers(); + try { + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/prompt')) { + return jsonResponse(200, { stopReason: 'end_turn' }); + } + if (req.url.endsWith('/cancel')) { + return jsonResponse(400, { code: 'invalid_client_id' }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'test-session', 'client-test'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + const pending = promptTool.handler({ prompt: 'test' }, {}); + await vi.advanceTimersByTimeAsync(30000); + const result = await pending; + + expect(result.isError).toBe(true); + const parsed = JSON.parse(result.content[0].text); + expect(parsed.stop_reason).toBe('timeout'); + expect(parsed.warning).toContain('session_resume'); + expect(parsed.warning).toContain('test-session'); + const cancelCalls = calls.filter((call) => + call.url.endsWith('/cancel'), + ); + expect(cancelCalls).toHaveLength(1); + expect(cancelCalls[0]?.headers['x-qwen-client-id']).toBe('client-test'); + expect(state.bindings.has('test-session')).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + it('should preserve a newer binding when stale timeout cancellation is rejected', async () => { + vi.useFakeTimers(); + try { + let resolveCancel!: (response: Response) => void; + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/prompt')) { + return jsonResponse(200, { stopReason: 'end_turn' }); + } + if (req.url.endsWith('/cancel')) { + return new Promise((resolve) => { + resolveCancel = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'test-session', 'client-stale'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + const pending = promptTool.handler({ prompt: 'test' }, {}); + await vi.advanceTimersByTimeAsync(30000); + expect(resolveCancel).toBeTypeOf('function'); + const newer = createBinding('test-session', 'client-new'); + state.bindings.set('test-session', newer); + resolveCancel(jsonResponse(400, { code: 'invalid_client_id' })); + const result = await pending; + + const parsed = JSON.parse(result.content[0].text); + expect(parsed.stop_reason).toBe('timeout'); + expect(parsed.warning).not.toContain('session_resume'); + expect(state.bindings.get('test-session')).toBe(newer); + const cancelCalls = calls.filter((call) => + call.url.endsWith('/cancel'), + ); + expect(cancelCalls).toHaveLength(1); + expect(cancelCalls[0]?.headers['x-qwen-client-id']).toBe( + 'client-stale', + ); + } finally { + vi.useRealTimers(); + } + }); + + it('should invalidate the current binding on explicit cancel rejection', async () => { + const { state } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => + req.url.endsWith('/cancel') + ? jsonResponse(400, { code: 'invalid_client_id' }) + : new Response(null, { status: 204 }), + }); + const current = bindSession(state, 'test-session', 'client-old'); + current.stream.activeCollector = createPromptCollector(); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const cancelTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt_cancel', + ); + + const currentResult = await cancelTool.handler({}, {}); + expect(currentResult.isError).toBe(true); + expect(currentResult.content[0].text).toContain('Call session_resume'); + expect(current.stream.activeCollector?.resolved).toBe(true); + expect(state.bindings.has('test-session')).toBe(false); + }); + + it('should preserve a newer binding on a stale cancel rejection', async () => { + let resolveCancel!: (response: Response) => void; + const { state } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/cancel')) { + return new Promise((resolve) => { + resolveCancel = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + const old = bindSession(state, 'test-session', 'client-stale'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const cancelTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt_cancel', + ); + + const pending = cancelTool.handler({}, {}); + await vi.waitFor(() => expect(resolveCancel).toBeTypeOf('function')); + const newer = createBinding('test-session', 'client-new'); + state.bindings.set('test-session', newer); + resolveCancel(jsonResponse(400, { code: 'invalid_client_id' })); + const staleResult = await pending; + + expect(staleResult.isError).toBe(true); + expect(staleResult.content[0].text).toContain('Retry the request'); + expect(state.bindings.get('test-session')).toBe(newer); + expect(old).not.toBe(newer); + }); + }); + + describe('binding cleanup', () => { + it.each([ + { + toolName: 'session_load', + request: { session_id: 'session-a' }, + matches: (req: CapturedRequest) => req.url.endsWith('/load'), + response: () => + jsonResponse(200, { + sessionId: 'session-a', + clientId: 'client-new', + workspaceCwd: '/tmp', + }), + expectedDetachIds: ['client-new', 'client-old'], + }, + { + toolName: 'session_resume', + request: { session_id: 'session-a' }, + matches: (req: CapturedRequest) => req.url.endsWith('/resume'), + response: () => + jsonResponse(200, { + sessionId: 'session-a', + clientId: 'client-new', + workspaceCwd: '/tmp', + }), + expectedDetachIds: ['client-new', 'client-old'], + }, + { + toolName: 'session_close', + request: { session_id: 'session-a' }, + matches: (req: CapturedRequest) => req.method === 'DELETE', + response: () => new Response(null, { status: 204 }), + expectedDetachIds: [], + }, + ])( + 'should drain a pending $toolName before disposal releases bindings', + async ({ toolName, request, matches, response, expectedDetachIds }) => { + let resolveOperation!: (response: Response) => void; + const { state, calls } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: (req) => { + if (matches(req)) { + return new Promise((resolve) => { + resolveOperation = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'session-a', 'client-old'); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const lifecycleTool = sessionTools(state).find( + (tool: { name: string }) => tool.name === toolName, + ); + + const operation = lifecycleTool.handler(request, {}); + await vi.waitFor(() => expect(resolveOperation).toBeTypeOf('function')); + let disposed = false; + const disposing = disposeBindings(state).finally(() => { + disposed = true; + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(disposed).toBe(false); + + resolveOperation(response()); + await operation; + await disposing; + + const detachIds = calls + .filter((call) => call.url.endsWith('/detach')) + .map((call) => call.headers['x-qwen-client-id']); + expect(detachIds).toEqual(expectedDetachIds); + expect(new Set(detachIds).size).toBe(detachIds.length); + expect(state.bindings.size).toBe(0); + }, + ); + + it('should skip active prompts during idle cleanup, then detach once idle', async () => { + vi.useFakeTimers(); + try { + const { state, calls } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: () => new Response(null, { status: 204 }), + }); + const binding = bindSession(state, 'session-a', 'client-a'); + binding.stream.lastActivityMs = 0; + binding.stream.activeCollector = createPromptCollector(); + const stopCleanup = startSessionCleanup(state); + + await vi.advanceTimersByTimeAsync(5 * 60 * 1000); + expect(state.bindings.get('session-a')).toBe(binding); + expect( + calls.filter((call) => call.url.endsWith('/detach')), + ).toHaveLength(0); + + binding.stream.activeCollector = null; + await vi.advanceTimersByTimeAsync(5 * 60 * 1000); + expect(state.bindings.has('session-a')).toBe(false); + expect( + calls.filter((call) => call.url.endsWith('/detach')), + ).toHaveLength(1); + stopCleanup(); + } finally { + vi.useRealTimers(); + } + }); + + it('should await an in-flight detach across repeated disposal', async () => { + let resolveDetach!: (response: Response) => void; + const { state, calls } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: (req) => { + if (req.url.endsWith('/detach')) { + return new Promise((resolve) => { + resolveDetach = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'session-a', 'client-a'); + + const first = disposeBindings(state); + const second = disposeBindings(state); + await vi.waitFor(() => expect(resolveDetach).toBeTypeOf('function')); + expect(calls.filter((call) => call.url.endsWith('/detach'))).toHaveLength( + 1, + ); + expect(state.pendingReleases.size).toBe(1); + resolveDetach(new Response(null, { status: 204 })); + await Promise.all([first, second]); + + expect(state.pendingReleases.size).toBe(0); + }); + + it('should dispose a legacy anonymous binding without a detach request', async () => { + const { state, calls } = makeMockState(); + bindSession(state, 'legacy-session'); + + await disposeBindings(state); + + expect(state.bindings.size).toBe(0); + expect(calls.some((call) => call.url.endsWith('/detach'))).toBe(false); + }); }); describe('safety guards', () => { From b932b84c65e2c446db8e6e9be5fba59c48fd20d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B6=E5=85=AC?= Date: Wed, 15 Jul 2026 01:47:43 +0800 Subject: [PATCH 2/9] fix(sdk): preserve completed prompt responses --- .../invocation-context/v1-conformance.json | 28 +++++++++++++ .../cli/invocation-context-v1.test.ts | 5 ++- .../daemon-mcp/serve-bridge/tools/agent.ts | 6 ++- .../test/unit/serve-bridge.test.ts | 42 +++++++++++++++++-- 4 files changed, 75 insertions(+), 6 deletions(-) diff --git a/docs/design/invocation-context/v1-conformance.json b/docs/design/invocation-context/v1-conformance.json index d9d6d02c3cc..1c9a7611d92 100644 --- a/docs/design/invocation-context/v1-conformance.json +++ b/docs/design/invocation-context/v1-conformance.json @@ -31,6 +31,34 @@ "name": "missing-context", "valid": false }, + { + "name": "missing-session-id", + "valid": false, + "input": { + "version": 1, + "ingress": "cli", + "promptId": "prompt-missing-session" + } + }, + { + "name": "missing-prompt-id", + "valid": false, + "input": { + "version": 1, + "ingress": "cli", + "sessionId": "session-missing-prompt" + } + }, + { + "name": "string-version", + "valid": false, + "input": { + "version": "1", + "ingress": "cli", + "sessionId": "session-string-version", + "promptId": "prompt-string-version" + } + }, { "name": "future-version", "valid": false, diff --git a/integration-tests/cli/invocation-context-v1.test.ts b/integration-tests/cli/invocation-context-v1.test.ts index 0d7063cc29d..798723c5bca 100644 --- a/integration-tests/cli/invocation-context-v1.test.ts +++ b/integration-tests/cli/invocation-context-v1.test.ts @@ -288,7 +288,10 @@ describe('Invocation Context v1 runtime propagation', () => { } }, 60_000); - it('keeps external_mcp session markings isolated through the daemon', async () => { + it('keeps external_mcp session markings isolated through the daemon', async ({ + skip, + }) => { + skip(process.platform === 'win32', 'requires the POSIX daemon harness'); const root = mkdtempSync(path.join(tmpdir(), 'qwen-invocation-daemon-')); const workspace = path.join(root, 'workspace'); const home = path.join(root, 'home'); diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts index ed45374109d..922977b67b5 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts @@ -67,8 +67,10 @@ export function agentTools(state: BridgeState): any[] { } }; const onAbort = () => { - collector.interrupted = true; - collector.resolve(); + if (!collector.resolved) { + collector.interrupted = true; + collector.resolve(); + } if (promptSettled) { void awaitBestEffortCancel(); } diff --git a/packages/sdk-typescript/test/unit/serve-bridge.test.ts b/packages/sdk-typescript/test/unit/serve-bridge.test.ts index 92aae651b36..ecf8be4f70a 100644 --- a/packages/sdk-typescript/test/unit/serve-bridge.test.ts +++ b/packages/sdk-typescript/test/unit/serve-bridge.test.ts @@ -871,6 +871,45 @@ describe('serve-bridge', () => { ).toBeUndefined(); }); + it('should preserve a completed response when abort arrives after collection', async () => { + const controller = new AbortController(); + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/prompt')) { + const collector = + state.bindings.get('test-session')!.stream.activeCollector!; + collector.texts.push('completed response'); + collector.resolve(); + controller.abort(); + return jsonResponse(200, { stopReason: 'end_turn' }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'test-session', 'client-test'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + const result = await promptTool.handler( + { prompt: 'test' }, + { signal: controller.signal }, + ); + + expect(result.isError).toBeUndefined(); + expect(JSON.parse(result.content[0].text)).toMatchObject({ + stop_reason: 'end_turn', + response: 'completed response', + }); + expect(calls.filter((call) => call.url.endsWith('/cancel'))).toHaveLength( + 1, + ); + }); + it('should not cancel an unadmitted legacy prompt after abort', async () => { const controller = new AbortController(); const { state, calls } = makeMockState({ @@ -985,9 +1024,6 @@ describe('serve-bridge', () => { lastEventId: 0, }); } - if (req.url.endsWith('/events')) { - return pendingSseResponse(); - } if (req.url.endsWith('/cancel')) { return jsonResponse(400, { code: 'invalid_client_id' }); } From aaa18347cc10b9333d066e69e83aa646c020dfc6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B6=E5=85=AC?= Date: Wed, 15 Jul 2026 04:10:35 +0800 Subject: [PATCH 3/9] fix(sdk): drain prompt cancellation during shutdown --- .../cli/src/acp-integration/acpAgent.test.ts | 47 +++++++++++++++++-- .../daemon-mcp/serve-bridge/tools/agent.ts | 34 +++++++------- .../test/unit/serve-bridge.test.ts | 39 +++++++++++++++ 3 files changed, 101 insertions(+), 19 deletions(-) diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index ec5f5f2fadd..30e074d3c3a 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -127,12 +127,13 @@ vi.mock('node:stream', async (importOriginal) => { }); // Mock core dependencies -vi.mock('@qwen-code/qwen-code-core', () => ({ +vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => ({ INVOCATION_CONTEXT_META_KEY: 'qwen-code/invocation', INVOCATION_INGRESS_META_KEY: 'qwen-code/invocation-ingress', PRIVATE_PARENT_CAPABILITY_META_KEY: 'qwen-code/private-parent-capability', - parseInvocationContext: vi.fn((value: unknown) => - typeof value === 'object' && value !== null ? value : undefined, + parseInvocationContext: vi.fn( + (await importOriginal()) + .parseInvocationContext, ), SESSION_ARTIFACT_PERSISTENCE_VERSION: 2, normalizeEventPayload: vi.fn((payload: unknown) => @@ -1569,6 +1570,46 @@ describe('QwenAgent MCP SSE/HTTP support', () => { await agentPromise; }); + it('rejects invalid invocation metadata from a trusted parent', async () => { + await setupSessionMocks('trusted-session'); + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + { + privateParentCapability: 'expected-capability', + }, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + const agent = capturedAgentFactory!({ + get closed() { + return mockConnectionState.promise; + }, + }) as AgentLike; + await agent.initialize({ + clientCapabilities: {}, + _meta: { + 'qwen-code/private-parent-capability': 'expected-capability', + }, + }); + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + await expect( + agent.prompt({ + sessionId: 'trusted-session', + prompt: [{ type: 'text', text: 'hello' }], + _meta: { + 'qwen-code/invocation': { + version: '1', + }, + }, + }), + ).rejects.toThrow('Invalid trusted ACP invocation context'); + expect(lastSessionMock?.prompt).not.toHaveBeenCalled(); + + mockConnectionState.resolve(); + await agentPromise; + }); + it('strips forged invocation metadata from an untrusted ACP client', async () => { const envName = 'QWEN_CODE_PRIVATE_ACP_CAPABILITY'; process.env[envName] = 'late-env-capability'; diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts index 922977b67b5..31d6abd6854 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/agent.ts @@ -194,23 +194,25 @@ export function agentTools(state: BridgeState): any[] { .optional() .describe('Session ID. Uses default session if omitted.'), }, - handler(async (args) => { - const sessionId = resolveSessionId(state, args.session_id); - const binding = resolveBinding(state, sessionId); - const { stream } = binding; - try { - await state.client.cancel(sessionId, binding.clientId); - } catch (err) { - await rethrowBindingError(state, binding, err); - } finally { - // Resolve active collector even when the daemon rejects cancellation. - if (stream.activeCollector) { - stream.activeCollector.interrupted = true; - stream.activeCollector.resolve(); + handler((args) => + trackLifecycle(state, async () => { + const sessionId = resolveSessionId(state, args.session_id); + const binding = resolveBinding(state, sessionId); + const { stream } = binding; + try { + await state.client.cancel(sessionId, binding.clientId); + } catch (err) { + await rethrowBindingError(state, binding, err); + } finally { + // Resolve active collector even when the daemon rejects cancellation. + if (stream.activeCollector) { + stream.activeCollector.interrupted = true; + stream.activeCollector.resolve(); + } } - } - return formatJsonResult({ ok: true, sessionId }); - }), + return formatJsonResult({ ok: true, sessionId }); + }), + ), ), ]; } diff --git a/packages/sdk-typescript/test/unit/serve-bridge.test.ts b/packages/sdk-typescript/test/unit/serve-bridge.test.ts index ecf8be4f70a..f746f1dcfb4 100644 --- a/packages/sdk-typescript/test/unit/serve-bridge.test.ts +++ b/packages/sdk-typescript/test/unit/serve-bridge.test.ts @@ -840,6 +840,45 @@ describe('serve-bridge', () => { expect(collector.interrupted).toBe(true); }); + it('should drain prompt_cancel before disposal releases its binding', async () => { + let resolveCancel!: (response: Response) => void; + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => { + if (req.url.endsWith('/cancel')) { + return new Promise((resolve) => { + resolveCancel = resolve; + }); + } + return new Response(null, { status: 204 }); + }, + }); + bindSession(state, 'test-session', 'client-test'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const cancelTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt_cancel', + ); + + const cancelling = cancelTool.handler({}, {}); + await vi.waitFor(() => expect(resolveCancel).toBeTypeOf('function')); + let disposalSettled = false; + const disposing = disposeBindings(state).finally(() => { + disposalSettled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(disposalSettled).toBe(false); + expect(calls.some((call) => call.url.endsWith('/detach'))).toBe(false); + + resolveCancel(new Response(null, { status: 204 })); + await Promise.all([cancelling, disposing]); + expect(calls.filter((call) => call.url.endsWith('/detach'))).toHaveLength( + 1, + ); + }); + it('should keep legacy anonymous bindings usable', async () => { const { state, calls } = makeMockState({ defaultSessionId: 'legacy-session', From 1ccbf886dd769c3dc51d6cf1d6a5c65eb73be0e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B6=E5=85=AC?= Date: Wed, 15 Jul 2026 06:03:55 +0800 Subject: [PATCH 4/9] test(cli): cover Docker child environment passthrough --- packages/cli/src/utils/sandbox.test.ts | 86 ++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/packages/cli/src/utils/sandbox.test.ts b/packages/cli/src/utils/sandbox.test.ts index 5d807872786..d2160378903 100644 --- a/packages/cli/src/utils/sandbox.test.ts +++ b/packages/cli/src/utils/sandbox.test.ts @@ -5,7 +5,11 @@ */ import fs from 'node:fs'; +import os from 'node:os'; import path from 'node:path'; +import { EventEmitter } from 'node:events'; +import { PassThrough } from 'node:stream'; +import type { ChildProcess } from 'node:child_process'; import { pathToFileURL } from 'node:url'; import { FatalSandboxError, QWEN_DIR } from '@qwen-code/qwen-code-core'; import { afterEach, describe, expect, it, vi } from 'vitest'; @@ -14,7 +18,21 @@ import { resolveSeatbeltProfileFile, start_sandbox } from './sandbox.js'; import { parseSandboxImageName } from './sandboxImageName.js'; import { parseSandboxMountSpec } from './sandboxMounts.js'; +const { spawnMock } = vi.hoisted(() => ({ + spawnMock: vi.fn(), +})); + +vi.mock('node:child_process', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + default: { ...actual, spawn: spawnMock }, + spawn: spawnMock, + }; +}); + afterEach(() => { + spawnMock.mockReset(); vi.restoreAllMocks(); vi.unstubAllEnvs(); }); @@ -76,6 +94,74 @@ describe('resolveSeatbeltProfileFile', () => { }); }); +describe('start_sandbox', () => { + it('passes child-only environment into a Docker sandbox', async () => { + const envName = 'QWEN_CODE_PRIVATE_ACP_CAPABILITY'; + const originalCapability = process.env[envName]; + const tempDir = fs.mkdtempSync( + path.join(os.tmpdir(), 'qwen-sandbox-test-'), + ); + delete process.env[envName]; + vi.stubEnv('QWEN_CODE_INTEGRATION_TEST', 'true'); + vi.stubEnv('SANDBOX_SET_UID_GID', 'false'); + vi.stubEnv('QWEN_HOME', path.join(tempDir, 'settings')); + vi.stubEnv('QWEN_RUNTIME_DIR', path.join(tempDir, 'runtime')); + vi.spyOn(process.stdin, 'pause').mockImplementation(() => process.stdin); + vi.spyOn(process.stdin, 'resume').mockImplementation(() => process.stdin); + spawnMock.mockImplementation((_command, args) => { + const isImageCheck = args?.[0] === 'images'; + const stdout = isImageCheck ? new PassThrough() : null; + const child = Object.assign(new EventEmitter(), { + stdout, + stderr: null, + connected: false, + disconnect: vi.fn(), + }) as unknown as ChildProcess; + queueMicrotask(() => { + stdout?.end('image-id\n'); + child.emit('close', 0, null); + }); + return child; + }); + + try { + await expect( + start_sandbox( + { command: 'docker', image: 'test-image' }, + [], + undefined, + ['node', 'script.js'], + { [envName]: 'private-capability' }, + ), + ).resolves.toBe(0); + + const runCall = spawnMock.mock.calls.find( + ([, args]) => args?.[0] === 'run', + ); + expect(runCall).toBeDefined(); + const runArgs = runCall?.[1] as readonly string[] | undefined; + const runOptions = runCall?.[2] as + | { env?: NodeJS.ProcessEnv } + | undefined; + expect( + runArgs?.findIndex( + (arg, index) => arg === '--env' && runArgs[index + 1] === envName, + ), + ).toBeGreaterThanOrEqual(0); + expect(runArgs).not.toContain(`${envName}=private-capability`); + expect(runOptions?.env?.[envName]).toBe('private-capability'); + expect(process.env[envName]).toBeUndefined(); + } finally { + fs.rmSync(tempDir, { recursive: true, force: true }); + if (originalCapability === undefined) { + delete process.env[envName]; + } else { + process.env[envName] = originalCapability; + } + } + }); +}); + describe('isContainerPathWithinWorkdir', () => { it('allows the workdir itself', () => { expect(isContainerPathWithinWorkdir('/repo/app', '/repo/app')).toBe(true); From b2a37fc44f3dcd024ec96bbdf39ff8bac3cb6cae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B6=E5=85=AC?= Date: Wed, 15 Jul 2026 10:22:53 +0800 Subject: [PATCH 5/9] refactor(cli): centralize invocation ingress mapping --- packages/cli/src/nonInteractiveCli.ts | 12 ++------ packages/cli/src/ui/hooks/useGeminiStream.ts | 9 ++---- .../cli/src/utils/invocation-ingress.test.ts | 24 +++++++++++++++ packages/cli/src/utils/invocation-ingress.ts | 29 +++++++++++++++++++ 4 files changed, 58 insertions(+), 16 deletions(-) create mode 100644 packages/cli/src/utils/invocation-ingress.test.ts create mode 100644 packages/cli/src/utils/invocation-ingress.ts diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 8d4aea8a3e6..6f0fa77d2ba 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -66,6 +66,7 @@ import { settleChatRecording, subscribeToHeadlessChatRecordingFailures, } from './utils/chat-recording-failure.js'; +import { getAsyncInvocationIngress } from './utils/invocation-ingress.js'; const debugLogger = createDebugLogger('NON_INTERACTIVE_CLI'); @@ -336,13 +337,7 @@ export async function runNonInteractive( prompt_id: string, options: RunNonInteractiveOptions = {}, ): Promise { - const ingress = - options.sendMessageType === SendMessageType.Cron - ? ('scheduler' as const) - : options.sendMessageType === SendMessageType.Notification || - options.sendMessageType === SendMessageType.Teammate - ? ('internal' as const) - : ('cli' as const); + const ingress = getAsyncInvocationIngress(options.sendMessageType) ?? 'cli'; return runWithInvocationContext( { version: 1, @@ -1640,8 +1635,7 @@ async function runNonInteractiveWithContext( }; const itemInvocationContext: InvocationContextV1 = { version: 1, - ingress: - targetType === SendMessageType.Cron ? 'scheduler' : 'internal', + ingress: getAsyncInvocationIngress(targetType) ?? 'internal', sessionId: config.getSessionId(), promptId: randomUUID(), }; diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index d93b7755a9d..3e646e96354 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -113,6 +113,7 @@ import { t } from '../../i18n/index.js'; import { useDualOutput } from '../../dualOutput/DualOutputContext.js'; import { recordGoalStatusItem } from '../utils/restoreGoal.js'; import { sanitizeDisplayText } from '../../utils/extension-mention.js'; +import { getAsyncInvocationIngress } from '../../utils/invocation-ingress.js'; import process from 'node:process'; const debugLogger = createDebugLogger('GEMINI_STREAM'); @@ -2480,13 +2481,7 @@ export const useGeminiStream = ( } const inheritedInvocation = getInvocationContext(); - const asyncIngress = - submitType === SendMessageType.Cron - ? ('scheduler' as const) - : submitType === SendMessageType.Notification || - submitType === SendMessageType.Teammate - ? ('internal' as const) - : undefined; + const asyncIngress = getAsyncInvocationIngress(submitType); const invocation = asyncIngress !== undefined ? { diff --git a/packages/cli/src/utils/invocation-ingress.test.ts b/packages/cli/src/utils/invocation-ingress.test.ts new file mode 100644 index 00000000000..b92da1d6941 --- /dev/null +++ b/packages/cli/src/utils/invocation-ingress.test.ts @@ -0,0 +1,24 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { SendMessageType } from '@qwen-code/qwen-code-core'; +import { getAsyncInvocationIngress } from './invocation-ingress.js'; + +describe('getAsyncInvocationIngress', () => { + it.each([ + [SendMessageType.Cron, 'scheduler'], + [SendMessageType.Notification, 'internal'], + [SendMessageType.Teammate, 'internal'], + [SendMessageType.UserQuery, undefined], + [SendMessageType.ToolResult, undefined], + [SendMessageType.Retry, undefined], + [SendMessageType.Hook, undefined], + [undefined, undefined], + ] as const)('maps %s to %s', (type, expected) => { + expect(getAsyncInvocationIngress(type)).toBe(expected); + }); +}); diff --git a/packages/cli/src/utils/invocation-ingress.ts b/packages/cli/src/utils/invocation-ingress.ts new file mode 100644 index 00000000000..b09c510f2db --- /dev/null +++ b/packages/cli/src/utils/invocation-ingress.ts @@ -0,0 +1,29 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { SendMessageType } from '@qwen-code/qwen-code-core'; + +export function getAsyncInvocationIngress( + type: SendMessageType | undefined, +): 'scheduler' | 'internal' | undefined { + switch (type) { + case SendMessageType.Cron: + return 'scheduler'; + case SendMessageType.Notification: + case SendMessageType.Teammate: + return 'internal'; + case SendMessageType.UserQuery: + case SendMessageType.ToolResult: + case SendMessageType.Retry: + case SendMessageType.Hook: + case undefined: + return undefined; + default: { + const exhaustiveType: never = type; + return exhaustiveType; + } + } +} From aa31ea3a37a5ca20f79b714b377cc618a1508c28 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B6=E5=85=AC?= Date: Wed, 15 Jul 2026 10:23:07 +0800 Subject: [PATCH 6/9] test(cli): cover invalid trusted invocation metadata --- .../acp-integration/session/Session.test.ts | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index dcbe9bc12ec..39d08f3be68 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -2079,6 +2079,36 @@ describe('Session', () => { ); }); + it('rejects an invalid trusted invocation context', async () => { + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'invalid context' }], + _meta: { 'qwen-code/invocation': { version: '1' } }, + }), + ).rejects.toMatchObject({ + code: -32602, + message: 'Invalid params: Invalid trusted ACP invocation context', + }); + + expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + }); + + it('rejects an invalid trusted invocation ingress', async () => { + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'invalid ingress' }], + _meta: { 'qwen-code/invocation-ingress': 'daemon' }, + }), + ).rejects.toMatchObject({ + code: -32602, + message: 'Invalid params: Invalid trusted ACP invocation ingress', + }); + + expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + }); + it('records the latest file history snapshot after makeSnapshot', async () => { const latestSnapshot = { promptId: 'test-session-id########1', From 523d3cd081c9f5a121aa943ba51157beaa44e44a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B6=E5=85=AC?= Date: Wed, 15 Jul 2026 15:15:23 +0800 Subject: [PATCH 7/9] fix(sdk): recover rejected session bindings --- .../daemon-mcp/serve-bridge/tools/session.ts | 42 ++++++++----- .../test/unit/serve-bridge.test.ts | 62 +++++++++++++++++++ 2 files changed, 90 insertions(+), 14 deletions(-) diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/session.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/session.ts index 52d0cd1e0b4..ec81e352761 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/session.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/tools/session.ts @@ -76,13 +76,20 @@ export function sessionTools(state: BridgeState): any[] { 'Cannot replace a session binding while a prompt is in progress.', ); } - const loaded = await state.client.loadSession( - args.session_id, - { - workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, - }, - current?.clientId, - ); + const loaded = await state.client + .loadSession( + args.session_id, + { + workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, + }, + current?.clientId, + ) + .catch(async (err: unknown) => { + if (current) { + return rethrowBindingError(state, current, err); + } + throw err; + }); await replaceBinding( state, createBinding(loaded.sessionId, loaded.clientId), @@ -114,13 +121,20 @@ export function sessionTools(state: BridgeState): any[] { 'Cannot replace a session binding while a prompt is in progress.', ); } - const resumed = await state.client.resumeSession( - args.session_id, - { - workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, - }, - current?.clientId, - ); + const resumed = await state.client + .resumeSession( + args.session_id, + { + workspaceCwd: args.workspace_cwd ?? state.workspaceCwd, + }, + current?.clientId, + ) + .catch(async (err: unknown) => { + if (current) { + return rethrowBindingError(state, current, err); + } + throw err; + }); await replaceBinding( state, createBinding(resumed.sessionId, resumed.clientId), diff --git a/packages/sdk-typescript/test/unit/serve-bridge.test.ts b/packages/sdk-typescript/test/unit/serve-bridge.test.ts index f746f1dcfb4..18790cabbcb 100644 --- a/packages/sdk-typescript/test/unit/serve-bridge.test.ts +++ b/packages/sdk-typescript/test/unit/serve-bridge.test.ts @@ -472,6 +472,68 @@ describe('serve-bridge', () => { }); }); + describe.each([ + { + toolName: 'session_load', + path: '/load', + }, + { + toolName: 'session_resume', + path: '/resume', + }, + ])('$toolName invalid binding recovery', ({ toolName, path }) => { + it('should invalidate a rejected binding before retrying', async () => { + let attempts = 0; + const { state, calls } = makeMockState({ + defaultSessionId: 'session-a', + fetchReply: (req) => { + if (req.url.endsWith(path)) { + attempts++; + return attempts === 1 + ? jsonResponse(400, { code: 'invalid_client_id' }) + : jsonResponse(200, { + sessionId: 'session-a', + clientId: 'client-new', + workspaceCwd: '/tmp', + }); + } + return new Response(null, { status: 204 }); + }, + }); + const old = bindSession(state, 'session-a', 'client-old'); + const { sessionTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/session.js' + ); + const lifecycleTool = sessionTools(state).find( + (tool: { name: string }) => tool.name === toolName, + ); + + const rejected = await lifecycleTool.handler( + { session_id: 'session-a' }, + {}, + ); + + expect(rejected.isError).toBe(true); + expect(rejected.content[0].text).toContain('Call session_resume'); + expect(state.bindings.has('session-a')).toBe(false); + expect(old.stream.abortCtrl.signal.aborted).toBe(true); + + const retried = await lifecycleTool.handler( + { session_id: 'session-a' }, + {}, + ); + + expect(retried.isError).toBeUndefined(); + const lifecycleCalls = calls.filter((call) => call.url.endsWith(path)); + expect(lifecycleCalls).toHaveLength(2); + expect(lifecycleCalls[0]?.headers['x-qwen-client-id']).toBe( + 'client-old', + ); + expect(lifecycleCalls[1]?.headers['x-qwen-client-id']).toBeUndefined(); + expect(state.bindings.get('session-a')?.clientId).toBe('client-new'); + }); + }); + describe('session_close', () => { it('should clear defaultSessionId when closing the default session', async () => { const { state } = makeMockState({ From 9ea26857b67301d9b2d2c689c500a8032249f1d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B6=E5=85=AC?= Date: Thu, 16 Jul 2026 10:58:28 +0800 Subject: [PATCH 8/9] fix(sdk): account for invocation ingress bundle size --- packages/sdk-typescript/scripts/build.js | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/sdk-typescript/scripts/build.js b/packages/sdk-typescript/scripts/build.js index 4b01cb63abe..dc784cc4d9f 100755 --- a/packages/sdk-typescript/scripts/build.js +++ b/packages/sdk-typescript/scripts/build.js @@ -55,7 +55,8 @@ const rootDir = join(__dirname, '..'); // Bumped from 154KB to 155KB after merging workspace skill-toggle APIs. // Bumped from 155KB to 160KB to accommodate recent growth and reduce churn // from repeated 1KB bumps as new daemon APIs are added. -const MAX_DAEMON_BROWSER_BUNDLE_BYTES = 160 * 1024; +// Bumped from 160KB to 161KB for daemon prompt invocation ingress metadata. +const MAX_DAEMON_BROWSER_BUNDLE_BYTES = 161 * 1024; // The opt-in `daemon/transports` browser bundle legitimately ships the concrete // ACP transports (AcpHttpTransport/AcpWsTransport/AutoReconnect + negotiate), so // it's larger than the default barrel — but still budgeted so a future PR can't From 8084e28e058b373d6ad45feabdbe4825d4197eb1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8F=B6=E5=85=AC?= Date: Thu, 16 Jul 2026 13:26:37 +0800 Subject: [PATCH 9/9] fix: address invocation context review feedback --- packages/channels/base/src/AcpBridge.test.ts | 20 +++++++++ .../acp-integration/session/Session.test.ts | 2 + .../src/acp-integration/session/Session.ts | 4 +- .../src/daemon-mcp/serve-bridge/bindings.ts | 9 ++++ .../test/unit/serve-bridge.test.ts | 42 +++++++++++++++++++ 5 files changed, 75 insertions(+), 2 deletions(-) diff --git a/packages/channels/base/src/AcpBridge.test.ts b/packages/channels/base/src/AcpBridge.test.ts index ab5e768e884..bace90648ab 100644 --- a/packages/channels/base/src/AcpBridge.test.ts +++ b/packages/channels/base/src/AcpBridge.test.ts @@ -289,6 +289,26 @@ describe('AcpBridge', () => { ); }); + it('forwards invocation ingress metadata to the ACP prompt', async () => { + const prompt = vi.fn().mockResolvedValue(undefined); + const bridge = new AcpBridge({ + cliEntryPath: '/tmp/qwen', + cwd: '/tmp', + }) as unknown as TestableAcpBridge; + bridge.child = { killed: false, exitCode: null }; + bridge.connection = { extMethod: vi.fn(), prompt }; + + await bridge.prompt('s-1', 'question', { + invocationIngress: 'channel', + }); + + expect(prompt).toHaveBeenCalledWith({ + sessionId: 's-1', + prompt: [{ type: 'text', text: 'question' }], + _meta: { 'qwen-code/invocation-ingress': 'channel' }, + }); + }); + it('excludes nested subagent text from the final response', async () => { const bridge = new AcpBridge({ cliEntryPath: '/tmp/qwen', diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index b94ac0d925e..63bb3abde0d 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -2092,6 +2092,7 @@ describe('Session', () => { }); expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + expect(session.getTurnCount()).toBe(0); }); it('rejects an invalid trusted invocation ingress', async () => { @@ -2107,6 +2108,7 @@ describe('Session', () => { }); expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + expect(session.getTurnCount()).toBe(0); }); it('records the latest file history snapshot after makeSnapshot', async () => { diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 813238359f8..5cd5c56a33e 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -1688,9 +1688,7 @@ export class Session implements SessionContext { params: PromptRequest, pendingSend: AbortController, ): Promise { - this.turn += 1; const sessionId = this.config.getSessionId(); - const promptId = sessionId + '########' + this.turn; const promptMeta = params._meta && typeof params._meta === 'object' ? { ...params._meta } @@ -1728,6 +1726,8 @@ export class Session implements SessionContext { 'Invalid trusted ACP invocation ingress', ); } + this.turn += 1; + const promptId = sessionId + '########' + this.turn; const invocationContext = trustedContext ?? { version: 1 as const, ingress: trustedIngress ?? ('acp' as const), diff --git a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bindings.ts b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bindings.ts index 3b8216e7939..41d953616c5 100644 --- a/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bindings.ts +++ b/packages/sdk-typescript/src/daemon-mcp/serve-bridge/bindings.ts @@ -192,6 +192,15 @@ export function startSessionCleanup(state: BridgeState): () => void { export async function disposeBindings(state: BridgeState): Promise { state.disposed = true; + // Let already-tracked handlers start so any collector they create is visible. + await Promise.resolve(); + for (const binding of state.bindings.values()) { + const collector = binding.stream.activeCollector; + if (collector) { + collector.interrupted = true; + collector.resolve(); + } + } while (state.pendingLifecycles.size > 0) { await Promise.all([...state.pendingLifecycles]); } diff --git a/packages/sdk-typescript/test/unit/serve-bridge.test.ts b/packages/sdk-typescript/test/unit/serve-bridge.test.ts index 18790cabbcb..4f5230fc919 100644 --- a/packages/sdk-typescript/test/unit/serve-bridge.test.ts +++ b/packages/sdk-typescript/test/unit/serve-bridge.test.ts @@ -1114,6 +1114,48 @@ describe('serve-bridge', () => { ); }); + it('should interrupt collector waits before draining prompt disposal', async () => { + const { state, calls } = makeMockState({ + defaultSessionId: 'test-session', + fetchReply: (req) => + req.url.endsWith('/prompt') + ? jsonResponse(200, { stopReason: 'end_turn' }) + : new Response(null, { status: 204 }), + }); + const binding = bindSession(state, 'test-session', 'client-test'); + const { agentTools } = await import( + '../../src/daemon-mcp/serve-bridge/tools/agent.js' + ); + const promptTool = agentTools(state).find( + (t: { name: string }) => t.name === 'prompt', + ); + + const prompting = promptTool.handler({ prompt: 'test' }, {}); + let disposalSettled = false; + const disposing = disposeBindings(state).finally(() => { + disposalSettled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + const settledBeforeCollectorTimeout = disposalSettled; + if (!settledBeforeCollectorTimeout) { + const collector = binding.stream.activeCollector; + if (collector) { + collector.interrupted = true; + collector.resolve(); + } + } + const [result] = await Promise.all([prompting, disposing]); + + expect(settledBeforeCollectorTimeout).toBe(true); + expect(JSON.parse(result.content[0].text).stop_reason).toBe( + 'interrupted', + ); + expect(calls.some((call) => call.url.endsWith('/cancel'))).toBe(false); + expect(calls.filter((call) => call.url.endsWith('/detach'))).toHaveLength( + 1, + ); + }); + it('should invalidate the current binding when abort cancellation is rejected', async () => { const controller = new AbortController(); const { state, calls } = makeMockState({