From ad55c22b5e79ac963dec6a695eb05e9d56178c83 Mon Sep 17 00:00:00 2001 From: Trevor Walker Date: Mon, 31 Aug 2026 09:35:18 -0600 Subject: [PATCH] feat: cancel RLM subagent trees when the parent request aborts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cancellation was strictly per-HTTP-request, so Prime Agent's RLM children — independent requests on independent session keys — kept running when the parent was cancelled, holding an SDK permit and a turn lease and billing the subscription until their own sockets closed. A new live registry (sessionTree.ts) records each in-flight keyed request and the immediate parent the client declared in metadata.user_id. A client abort of one request now aborts every live request whose ancestry reaches it, through each child's own request abort controller, so the mapping eviction, permit release, and lease release are the existing abort path's. Scope is deliberate: only an abort propagates (a parent turn that merely completes leaves children running), only live requests are tracked (no persistent tree), and propagation can only reach a request that declared a parent — which is the whole gate, no config flag. Also adds POST /v1/sessions/:key/cancel and sessionTree counters on /telemetry/summary and the dashboard. Fixes #902 --- ARCHITECTURE.md | 41 ++ CLAUDE.md | 2 + README.md | 12 +- docs/agents.md | 26 +- src/__tests__/claude-code-adapter.test.ts | 18 + src/__tests__/prime-adapter.test.ts | 52 ++ .../proxy-session-tree-cancellation.test.ts | 469 ++++++++++++++++++ src/__tests__/session-tree-unit.test.ts | 292 +++++++++++ src/proxy/adapter.ts | 15 + src/proxy/adapters/claudecode.ts | 61 ++- src/proxy/adapters/prime.ts | 19 +- src/proxy/server.ts | 142 +++++- src/proxy/sessionTree.ts | 250 ++++++++++ src/telemetry/dashboard.ts | 7 + src/telemetry/routes.ts | 15 +- src/telemetry/types.ts | 17 + 16 files changed, 1426 insertions(+), 12 deletions(-) create mode 100644 src/__tests__/proxy-session-tree-cancellation.test.ts create mode 100644 src/__tests__/session-tree-unit.test.ts create mode 100644 src/proxy/sessionTree.ts diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index b3f68733..d2571e96 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -38,6 +38,8 @@ src/ ├── proxy/ │ ├── server.ts ← HTTP layer: routes, SSE streaming, concurrency, request orchestration │ ├── concurrency.ts ← Abortable SDK query semaphore and concurrency config parsing +│ ├── requestAbort.ts ← HTTP request abort → SDK query abort bridge +│ ├── sessionTree.ts ← Live parent→child request registry; subtree cancellation (PURE bookkeeping) │ ├── shutdown.ts ← Bounded HTTP drain and connection tracking │ ├── adapter.ts ← AgentAdapter interface (extensibility point for multi-agent support) │ ├── adapters/ @@ -97,6 +99,8 @@ server.ts (HTTP layer) ├── query.ts ──► adapter.ts, mcpTools.ts, passthroughTools.ts ├── errors.ts ├── retryAfter.ts + ├── requestAbort.ts + ├── sessionTree.ts ├── models.ts ├── tools.ts ├── messages.ts @@ -129,6 +133,8 @@ server.ts (HTTP layer) 7. **`query.ts` builds SDK options through the adapter interface**, never importing tool constants directly. +8. **`sessionTree.ts` holds only live-request bookkeeping.** No HTTP, no I/O, no logging: the caller supplies each entry's abort handle and owns the eviction and telemetry discipline that follows an abort. It must not import from `server.ts`, `session/`, or `adapter.ts`. + ## Agent Adapter Pattern Agent-specific behavior is isolated behind the `AgentAdapter` interface (`adapter.ts`). The proxy calls adapter methods instead of hardcoding agent logic. @@ -225,6 +231,41 @@ downgraded every concurrent sibling at once, and the model switch cold-caches each of them — their cached prefixes were built on the 1M model. Clients with no session identity still bench profile-wide; there is nothing narrower to use. +## Cancellation Contract + +Cancellation is per-HTTP-request: `requestAbort.ts` forwards one socket's abort +into that request's SDK abort controller, and the abort path evicts the session +mapping so no interrupted tail stays resumable. + +That is not enough for a client whose subagents are separate requests. Prime +Agent's RLM children arrive on their own session keys, so cancelling the parent +left every child running — holding an SDK permit and a turn lease, billing the +subscription until its own socket closed or the lease watchdog tripped. + +`sessionTree.ts` closes the gap. A client that knows its own tree stamps the +immediate parent alongside the child's session id (`metadata.user_id` → +`{ session_id, parent_session_id }`); `server.ts` registers that link for the +lifetime of the request and, on a client abort, aborts every live request whose +ancestry reaches the aborted key — through each child's own request abort +controller, so the eviction, permit release, and lease release that follow are +the existing abort path's rather than a second implementation. + +Three properties bound it: + +- **Abort, not completion.** A parent turn that finishes normally does not + cancel children; a subagent routinely outlives the turn that spawned it. The + shutdown path already aborts every request directly, and the lease watchdog is + a proxy-side fence rather than a user intent, so neither cascades. +- **Live requests only.** An entry exists between "admitted" and "settled". A + session that was seen once but has nothing in flight is not a cancellation + target, so the registry is bounded by concurrency, not by history. +- **Self-gating.** Propagation can only reach a request that declared a parent, + so every client that does not stamp linkage is unaffected with no flag to set. + +`POST /v1/sessions/:key/cancel` cancels a subtree explicitly, and +`GET /telemetry/summary` reports the live gauges and cumulative counts under +`sessionTree`. + ## Testing Strategy Three tiers, each catching different classes of bugs: diff --git a/CLAUDE.md b/CLAUDE.md index 4513da65..3941d87c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -69,6 +69,8 @@ OpenCode-specific behavior is documented in `ARCHITECTURE.md` under "Agent-Speci ``` server.ts → HTTP routes, SSE streaming, concurrency (orchestration only) concurrency.ts → Abortable SDK query semaphore, max-concurrency config +requestAbort.ts → HTTP request abort → SDK query abort bridge +sessionTree.ts → Live parent→child request registry, subtree cancellation (PURE bookkeeping) shutdown.ts → Bounded HTTP drain, socket tracking, forced close adapter.ts → AgentAdapter interface (extensibility point) adapters/ diff --git a/README.md b/README.md index b1bcfb96..afda51b8 100644 --- a/README.md +++ b/README.md @@ -118,9 +118,15 @@ Prime Agent is reliable through Meridian with one active agent. RLM children hav separate session identities and can execute successfully, but concurrent subagent orchestration is not yet production-safe. Observed failure modes include overload amplification, expensive cache churn after fresh-session replay, loss of child-task -context during recovery, undelivered tool envelopes, and incomplete parent-to-child -cancellation. Use a single active Prime Agent for unattended or usage-sensitive work -until coordinated fixes land in Prime Agent and Meridian. +context during recovery, and undelivered tool envelopes. Use a single active Prime +Agent for unattended or usage-sensitive work until coordinated fixes land in Prime +Agent and Meridian. + +Parent-to-child cancellation is handled on the Meridian side: when the extension +stamps `parent_session_id` alongside the child's session id, aborting a parent's +in-flight request aborts every live request in the subtree below it and evicts +each one's session mapping. See +[Subagent cancellation](docs/agents.md#prime-agent). Prime Agent can keep Opus on the root session while selecting Sol for an individual child. A child inherits its parent's model unless the `rlm` call supplies an exact diff --git a/docs/agents.md b/docs/agents.md index 9c601316..1ffea081 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -364,9 +364,16 @@ export default function (pi: ExtensionAPI) { if (ctx?.model?.provider !== MERIDIAN_PROVIDER) return undefined const sessionId = ctx?.sessionManager?.getSessionId?.() if (typeof sessionId !== "string" || !sessionId) return undefined + const identity: Record = { session_id: sessionId } + // Optional, and only present on newer Prime Agent builds. It is what lets + // Meridian cancel a whole subagent tree — see "Subagent cancellation". + const parentSessionId = ctx?.sessionManager?.getParentSessionId?.() + if (typeof parentSessionId === "string" && parentSessionId) { + identity.parent_session_id = parentSessionId + } return { ...(event.payload as Record), - metadata: { user_id: JSON.stringify({ session_id: sessionId }) }, + metadata: { user_id: JSON.stringify(identity) }, } }) } @@ -405,6 +412,23 @@ design, the model loses track of what it has already run. Stamping real tool calls in its own session state. `getSessionId()` is distinct per agent, so RLM children get their own keys rather than colliding with the parent. +**Subagent cancellation.** RLM children reach Meridian as independent requests +on their own session keys, so cancelling the parent used to leave every child +running — holding an SDK permit and billing the subscription until its own +socket closed. `parent_session_id` closes that: it names the child's *immediate* +parent, Meridian keeps a registry of in-flight requests and their parent links, +and aborting a parent's request aborts every live request in the subtree below +it, evicting each one's session mapping exactly as a direct cancel does. + +Three limits are deliberate. Only an actual abort propagates — a parent turn +that merely finishes leaves its children alone, because a child routinely +outlives the turn that spawned it. Only *live* requests are tracked; a session +with nothing in flight is not remembered. And a client that omits +`parent_session_id` is unaffected, which is the whole gate — there is no config +flag. `POST /v1/sessions//cancel` cancels a subtree explicitly if you want +to stop one without dropping sockets, and `/telemetry/summary` reports the +counts under `sessionTree`. + Detection is by the `x-meridian-agent: prime` header above, or `MERIDIAN_DEFAULT_AGENT=prime`. There is deliberately no User-Agent rule: in API-key mode Prime Agent sends the generic `Anthropic/JS ` that every diff --git a/src/__tests__/claude-code-adapter.test.ts b/src/__tests__/claude-code-adapter.test.ts index 31b2635d..13f85928 100644 --- a/src/__tests__/claude-code-adapter.test.ts +++ b/src/__tests__/claude-code-adapter.test.ts @@ -41,6 +41,24 @@ describe("claudeCodeAdapter.getSessionId", () => { expect(claudeCodeAdapter.getSessionId(ctx as any, body)).toBe("object-session") }) + it("keeps the key equal to session_id when parent linkage is present", () => { + // parent_session_id is additive (#902): it must never change the key a + // client's cached mappings are already stored under. + const ctx = { req: { header: () => undefined } } + const body = { + metadata: { user_id: JSON.stringify({ session_id: "child", parent_session_id: "parent" }) }, + } + expect(claudeCodeAdapter.getSessionId(ctx as any, body)).toBe("child") + expect(claudeCodeAdapter.getParentSessionId!(ctx as any, body)).toBe("parent") + }) + + it("reports no parent for a root session", () => { + const ctx = { req: { header: () => undefined } } + expect(claudeCodeAdapter.getParentSessionId!(ctx as any, { + metadata: { user_id: JSON.stringify({ session_id: "root" }) }, + })).toBeUndefined() + }) + it("falls back to fingerprinting when metadata is absent or malformed", () => { const ctx = { req: { header: () => undefined } } expect(claudeCodeAdapter.getSessionId(ctx as any, {})).toBeUndefined() diff --git a/src/__tests__/prime-adapter.test.ts b/src/__tests__/prime-adapter.test.ts index 1ca62a40..415651d3 100644 --- a/src/__tests__/prime-adapter.test.ts +++ b/src/__tests__/prime-adapter.test.ts @@ -219,6 +219,58 @@ describe("primeAdapter.getSessionId", () => { }) }) +describe("primeAdapter.getParentSessionId", () => { + it("reads the immediate parent out of the same metadata envelope", () => { + // The extension stamps ctx.sessionManager.getParentSessionId() alongside + // the child's own id; the proxy uses it to cancel a live subtree (#902). + const body = { + metadata: { + user_id: JSON.stringify({ + session_id: "019ff7d8-ace2-7060-91dd-0212014a849e", + parent_session_id: "019ff7d8-a616-745d-8cb2-97544a6accac", + }), + }, + } + expect(primeAdapter.getParentSessionId!(ctxWith(), body)) + .toBe("019ff7d8-a616-745d-8cb2-97544a6accac") + // Key derivation is untouched: the child's key is still its own session_id. + expect(primeAdapter.getSessionId(ctxWith(), body)) + .toBe("019ff7d8-ace2-7060-91dd-0212014a849e") + }) + + it("returns undefined for a root session, which carries only session_id", () => { + const body = { metadata: { user_id: JSON.stringify({ session_id: "root-session" }) } } + expect(primeAdapter.getParentSessionId!(ctxWith(), body)).toBeUndefined() + }) + + it("ignores body linkage when an orchestrator owns identity via header", () => { + // x-session-affinity names keys under a different scheme, so a parent id + // read out of the body would point at a key that scheme never produced. + const c = ctxWith({ "x-session-affinity": "orchestrator-key" }) + const body = { + metadata: { user_id: JSON.stringify({ session_id: "child", parent_session_id: "parent" }) }, + } + expect(primeAdapter.getParentSessionId!(c, body)).toBeUndefined() + }) + + it("ignores a self-referential parent", () => { + const body = { + metadata: { user_id: JSON.stringify({ session_id: "same", parent_session_id: "same" }) }, + } + expect(primeAdapter.getParentSessionId!(ctxWith(), body)).toBeUndefined() + }) + + it("ignores malformed linkage without losing the session key", () => { + for (const parent of [null, 42, "", {}, []]) { + const body = { + metadata: { user_id: JSON.stringify({ session_id: "child", parent_session_id: parent }) }, + } + expect(primeAdapter.getParentSessionId!(ctxWith(), body)).toBeUndefined() + expect(primeAdapter.getSessionId(ctxWith(), body)).toBe("child") + } + }) +}) + describe("prime adapter configuration", () => { it("uses its own MCP server name", () => { expect(primeAdapter.getMcpServerName()).toBe("prime") diff --git a/src/__tests__/proxy-session-tree-cancellation.test.ts b/src/__tests__/proxy-session-tree-cancellation.test.ts new file mode 100644 index 00000000..3e258d8f --- /dev/null +++ b/src/__tests__/proxy-session-tree-cancellation.test.ts @@ -0,0 +1,469 @@ +/** + * Parent-to-child cancellation at the server seam (issue #902). + * + * Prime Agent's RLM children are independent HTTP requests on independent + * session keys. Before this, cancelling the parent left every child running — + * holding an SDK permit and a turn lease, billing the subscription — until its + * own socket closed or the lease watchdog tripped. + * + * The child declares its parent in the same `metadata.user_id` envelope that + * carries its session id, so these tests drive the real wire contract: + * client A is the parent, client B is a child stamped with + * `parent_session_id`, and the parent's request is aborted the way a client + * disconnect aborts it. + */ +import { afterEach, beforeEach, describe, expect, it, mock } from "bun:test" +import { resolveMockSdkSessionId, withMockSdkSessionId } from "./helpers" + +type Behavior = "complete" | "hang" + +interface SdkCall { + readonly controller: AbortController | undefined + readonly resume: string | undefined + readonly sessionId: string | undefined +} + +let behaviors: Behavior[] = [] +let calls: SdkCall[] = [] +let notifyQueryStarted: (() => void) | undefined + +function assistantMessage(sessionId: string) { + return { + type: "assistant", + message: { + id: "msg_session_tree", + type: "message", + role: "assistant", + content: [{ type: "text", text: "ok" }], + model: "claude-sonnet-4-6", + stop_reason: "end_turn", + usage: { input_tokens: 1, output_tokens: 1 }, + }, + parent_tool_use_id: null, + uuid: crypto.randomUUID(), + session_id: sessionId, + } +} + +mock.module("@anthropic-ai/claude-agent-sdk", () => ({ + query: (params: { options?: { abortController?: AbortController; resume?: string; sessionId?: string } }) => { + const behavior = behaviors.shift() ?? "complete" + const controller = params.options?.abortController + calls.push({ + controller, + resume: params.options?.resume, + sessionId: params.options?.sessionId, + }) + const notify = notifyQueryStarted + notifyQueryStarted = undefined + notify?.() + return (async function* () { + if (behavior === "complete") { + const sessionId = resolveMockSdkSessionId(params.options, crypto.randomUUID()) + yield withMockSdkSessionId(assistantMessage(sessionId), params.options) + return + } + // Hold the turn open until something aborts it — the shape a live + // subagent turn has when its parent is cancelled. + await new Promise((_resolve, reject) => { + const signal = controller?.signal + if (!signal) return reject(new Error("missing SDK abort controller")) + if (signal.aborted) return reject(new Error("SDK query aborted")) + signal.addEventListener("abort", () => reject(new Error("SDK query aborted")), { once: true }) + }) + })() + }, + createSdkMcpServer: () => ({ type: "sdk", name: "test", instance: { tool: () => {}, registerTool: () => ({}) } }), + tool: () => ({}), +})) + +mock.module("../logger", () => ({ + claudeLog: () => {}, + withClaudeLogContext: (_ctx: unknown, fn: () => unknown) => fn(), +})) + +mock.module("../mcpTools", () => ({ + createOpencodeMcpServer: () => ({ type: "sdk", name: "opencode", instance: {} }), +})) + +const { createProxyServer, clearSessionCache } = await import("../proxy/server") +const { processSessionTree } = await import("../proxy/sessionTree") + +const PARENT = "prime-parent-session" +const CHILD = "prime-child-session" +const GRANDCHILD = "prime-grandchild-session" + +function messagesRequest(options: { + sessionId: string + parentSessionId?: string + stream?: boolean + messages?: Array<{ role: string; content: unknown }> + signal?: AbortSignal +}) { + const identity: Record = { session_id: options.sessionId } + if (options.parentSessionId) identity.parent_session_id = options.parentSessionId + return new Request("http://localhost/v1/messages", { + method: "POST", + headers: { + "Content-Type": "application/json", + // Prime Agent's User-Agent is the generic Anthropic SDK one, so the + // adapter is selected explicitly — exactly as the provider config does. + "x-meridian-agent": "prime", + }, + body: JSON.stringify({ + model: "claude-sonnet-4-6", + max_tokens: 128, + stream: options.stream ?? false, + messages: options.messages ?? [{ role: "user", content: `hello from ${options.sessionId}` }], + metadata: { user_id: JSON.stringify(identity) }, + }), + signal: options.signal, + }) +} + +function queryStarted(): Promise { + return new Promise((resolve) => { notifyQueryStarted = resolve }) +} + +/** Let the cascade's abort listeners and the aborted turns' teardown settle. */ +async function settle(): Promise { + for (let i = 0; i < 5; i++) await new Promise((resolve) => setTimeout(resolve, 0)) +} + +async function drain(response: Response): Promise { + if (!response.body) return "" + const reader = response.body.getReader() + const decoder = new TextDecoder() + let text = "" + while (true) { + const { done, value } = await reader.read() + if (done) break + text += decoder.decode(value, { stream: true }) + } + return text +} + +describe("parent-to-child cancellation", () => { + let originalPassthrough: string | undefined + + beforeEach(() => { + originalPassthrough = process.env.MERIDIAN_PASSTHROUGH + process.env.MERIDIAN_PASSTHROUGH = "1" + behaviors = [] + calls = [] + notifyQueryStarted = undefined + clearSessionCache() + processSessionTree.clear() + }) + + afterEach(() => { + if (originalPassthrough === undefined) delete process.env.MERIDIAN_PASSTHROUGH + else process.env.MERIDIAN_PASSTHROUGH = originalPassthrough + processSessionTree.clear() + }) + + it("aborts a linked child's in-flight SDK query when the parent's request aborts", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + const parentAbort = new AbortController() + + behaviors = ["hang", "hang"] + const parentStarted = queryStarted() + const parentResponse = app.fetch(messagesRequest({ sessionId: PARENT, signal: parentAbort.signal })) + await parentStarted + + const childStarted = queryStarted() + const childResponse = app.fetch(messagesRequest({ sessionId: CHILD, parentSessionId: PARENT })) + await childStarted + + expect(calls).toHaveLength(2) + const childController = calls[1]!.controller! + expect(childController.signal.aborted).toBe(false) + + parentAbort.abort("client hung up") + await settle() + + // The child's own socket is still open; only the parent's was cut. + expect(childController.signal.aborted).toBe(true) + expect(calls[0]!.controller!.signal.aborted).toBe(true) + + expect((await parentResponse).status).toBe(499) + expect((await childResponse).status).toBe(499) + }) + + it("aborts a linked child's stream, closing it with an error frame", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + const parentAbort = new AbortController() + + behaviors = ["hang", "hang"] + const parentStarted = queryStarted() + const parentResponse = app.fetch(messagesRequest({ sessionId: PARENT, signal: parentAbort.signal })) + await parentStarted + + const childStarted = queryStarted() + const childStream = await app.fetch(messagesRequest({ + sessionId: CHILD, + parentSessionId: PARENT, + stream: true, + })) + await childStarted + + parentAbort.abort("client hung up") + const childBody = await drain(childStream) + await parentResponse + + expect(calls[1]!.controller!.signal.aborted).toBe(true) + expect(childBody).toContain("event: error") + }) + + it("cancels children when a streaming parent's response body is cancelled", async () => { + // Prime Agent streams every request, and a cancelled response body is the + // abort path an in-process caller reaches — the request signal never fires. + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + + behaviors = ["hang", "hang"] + const parentStarted = queryStarted() + const parentStream = await app.fetch(messagesRequest({ sessionId: PARENT, stream: true })) + await parentStarted + + const childStarted = queryStarted() + const childResponse = app.fetch(messagesRequest({ sessionId: CHILD, parentSessionId: PARENT })) + await childStarted + + await parentStream.body!.cancel("reader closed") + await settle() + + expect(calls[1]!.controller!.signal.aborted).toBe(true) + expect(processSessionTree.stats().propagations).toBe(1) + expect((await childResponse).status).toBe(499) + }) + + it("propagates once when a teardown trips both client abort paths", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + const parentAbort = new AbortController() + + behaviors = ["hang", "hang"] + const parentStarted = queryStarted() + const parentStream = await app.fetch(messagesRequest({ + sessionId: PARENT, + stream: true, + signal: parentAbort.signal, + })) + await parentStarted + + const childStarted = queryStarted() + const childResponse = app.fetch(messagesRequest({ sessionId: CHILD, parentSessionId: PARENT })) + await childStarted + + parentAbort.abort("client hung up") + await parentStream.body!.cancel("reader closed").catch(() => {}) + await settle() + + expect(calls[1]!.controller!.signal.aborted).toBe(true) + expect(processSessionTree.stats().propagations).toBe(1) + expect(processSessionTree.stats().cancelledDescendants).toBe(1) + await childResponse + }) + + it("walks the whole subtree, not just the immediate children", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + const parentAbort = new AbortController() + + behaviors = ["hang", "hang", "hang"] + const parentStarted = queryStarted() + const parentResponse = app.fetch(messagesRequest({ sessionId: PARENT, signal: parentAbort.signal })) + await parentStarted + + const childStarted = queryStarted() + const childResponse = app.fetch(messagesRequest({ sessionId: CHILD, parentSessionId: PARENT })) + await childStarted + + // The grandchild names its IMMEDIATE parent, per the wire contract. + const grandchildStarted = queryStarted() + const grandchildResponse = app.fetch(messagesRequest({ + sessionId: GRANDCHILD, + parentSessionId: CHILD, + })) + await grandchildStarted + + parentAbort.abort("client hung up") + await settle() + + expect(calls[1]!.controller!.signal.aborted).toBe(true) + expect(calls[2]!.controller!.signal.aborted).toBe(true) + expect((await childResponse).status).toBe(499) + expect((await grandchildResponse).status).toBe(499) + await parentResponse + }) + + it("evicts the cancelled child's session mapping so no interrupted tail is resumable", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + + // Turn 1 for the child completes normally and establishes a mapping. + behaviors = ["complete"] + const firstTurn = await app.fetch(messagesRequest({ sessionId: CHILD, parentSessionId: PARENT })) + expect(firstTurn.status).toBe(200) + await firstTurn.json() + + const continuation = [ + { role: "user", content: `hello from ${CHILD}` }, + { role: "assistant", content: [{ type: "text", text: "ok" }] }, + { role: "user", content: "keep going" }, + ] + + // Turn 2 resumes it, then dies to the parent's abort mid-flight. + behaviors = ["hang", "hang"] + const parentAbort = new AbortController() + const parentStarted = queryStarted() + const parentResponse = app.fetch(messagesRequest({ sessionId: PARENT, signal: parentAbort.signal })) + await parentStarted + + const cancelledStarted = queryStarted() + const cancelledTurn = app.fetch(messagesRequest({ + sessionId: CHILD, + parentSessionId: PARENT, + messages: continuation, + })) + await cancelledStarted + expect(calls[2]!.resume).toBeDefined() + + parentAbort.abort("client hung up") + await settle() + expect((await cancelledTurn).status).toBe(499) + await parentResponse + + // Turn 3 must start fresh: the interrupted turn may have advanced the SDK + // transcript past what the mapping described. + behaviors = ["complete"] + const afterCancel = await app.fetch(messagesRequest({ + sessionId: CHILD, + parentSessionId: PARENT, + messages: continuation, + })) + expect(afterCancel.status).toBe(200) + await afterCancel.json() + expect(calls[3]!.resume).toBeUndefined() + }) + + it("leaves children alone when the parent's turn merely completes", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + + behaviors = ["hang", "complete"] + const childStarted = queryStarted() + const childResponse = app.fetch(messagesRequest({ sessionId: CHILD, parentSessionId: PARENT })) + await childStarted + + const parentTurn = await app.fetch(messagesRequest({ sessionId: PARENT })) + expect(parentTurn.status).toBe(200) + await parentTurn.json() + await settle() + + // A subagent routinely outlives the parent turn that spawned it. + expect(calls[0]!.controller!.signal.aborted).toBe(false) + expect(processSessionTree.stats().cancelledDescendants).toBe(0) + + // Clean up the still-running child. + calls[0]!.controller!.abort("test cleanup") + await childResponse + }) + + it("leaves a request that declared no parent alone", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + const parentAbort = new AbortController() + + behaviors = ["hang", "hang"] + const parentStarted = queryStarted() + const parentResponse = app.fetch(messagesRequest({ sessionId: PARENT, signal: parentAbort.signal })) + await parentStarted + + // Same client, unrelated conversation: no parent_session_id, so propagation + // cannot reach it. This is the gate that keeps every other client inert. + const unlinkedStarted = queryStarted() + const unlinkedResponse = app.fetch(messagesRequest({ sessionId: "prime-unrelated-session" })) + await unlinkedStarted + + parentAbort.abort("client hung up") + await settle() + + expect(calls[1]!.controller!.signal.aborted).toBe(false) + expect(processSessionTree.stats().cancelledDescendants).toBe(0) + await parentResponse + + calls[1]!.controller!.abort("test cleanup") + await unlinkedResponse + }) + + it("counts propagated cancellations on /telemetry/summary", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + const parentAbort = new AbortController() + + behaviors = ["hang", "hang"] + const parentStarted = queryStarted() + const parentResponse = app.fetch(messagesRequest({ sessionId: PARENT, signal: parentAbort.signal })) + await parentStarted + + const childStarted = queryStarted() + const childResponse = app.fetch(messagesRequest({ sessionId: CHILD, parentSessionId: PARENT })) + await childStarted + + const before = await (await app.fetch(new Request("http://localhost/telemetry/summary"))).json() as { + sessionTree: { tracked: number; linked: number; propagations: number; cancelledDescendants: number } + } + expect(before.sessionTree.tracked).toBe(2) + expect(before.sessionTree.linked).toBe(1) + expect(before.sessionTree.propagations).toBe(0) + + parentAbort.abort("client hung up") + await settle() + await parentResponse + await childResponse + + const after = await (await app.fetch(new Request("http://localhost/telemetry/summary"))).json() as { + sessionTree: { tracked: number; propagations: number; cancelledDescendants: number } + } + expect(after.sessionTree.propagations).toBe(1) + expect(after.sessionTree.cancelledDescendants).toBe(1) + // Both requests settled, so the live registry is empty again. + expect(after.sessionTree.tracked).toBe(0) + }) + + it("cancels a subtree explicitly via POST /v1/sessions/:key/cancel", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + + behaviors = ["hang", "hang"] + const parentStarted = queryStarted() + const parentResponse = app.fetch(messagesRequest({ sessionId: PARENT })) + await parentStarted + + const childStarted = queryStarted() + const childResponse = app.fetch(messagesRequest({ sessionId: CHILD, parentSessionId: PARENT })) + await childStarted + + const cancelled = await app.fetch(new Request( + `http://localhost/v1/sessions/${encodeURIComponent(PARENT)}/cancel`, + { method: "POST" }, + )) + expect(cancelled.status).toBe(200) + expect(await cancelled.json()).toMatchObject({ + session: PARENT, + cancelled: { sessions: 2, requests: 2 }, + }) + + await settle() + expect(calls[0]!.controller!.signal.aborted).toBe(true) + expect(calls[1]!.controller!.signal.aborted).toBe(true) + expect((await parentResponse).status).toBe(499) + expect((await childResponse).status).toBe(499) + }) + + it("reports an idle session as nothing to cancel", async () => { + const app = createProxyServer({ port: 0, host: "127.0.0.1", silent: true }).app + const cancelled = await app.fetch(new Request( + "http://localhost/v1/sessions/never-seen/cancel", + { method: "POST" }, + )) + expect(cancelled.status).toBe(200) + expect(await cancelled.json()).toMatchObject({ + cancelled: { sessions: 0, requests: 0 }, + }) + }) +}) diff --git a/src/__tests__/session-tree-unit.test.ts b/src/__tests__/session-tree-unit.test.ts new file mode 100644 index 00000000..1f418e8d --- /dev/null +++ b/src/__tests__/session-tree-unit.test.ts @@ -0,0 +1,292 @@ +/** + * Unit tests for the live session-tree registry (issue #902). + * + * Pure module, no mocks: registration, settle-removal, transitive subtree + * computation, and the hostile shapes that arrive off the wire (cycles, + * self-links, absurd depth). + */ +import { describe, it, expect } from "bun:test" +import { + SessionTreeRegistry, + processSessionTree, + truncateSessionKey, +} from "../proxy/sessionTree" + +interface Aborted { + requestId: string + reason?: unknown +} + +function tracker() { + const aborted: Aborted[] = [] + const entry = (requestId: string, sessionKey: string, parentKey?: string) => ({ + requestId, + sessionKey, + parentKey, + abort: (reason?: unknown) => { aborted.push({ requestId, reason }) }, + }) + return { aborted, entry } +} + +describe("SessionTreeRegistry — registration and settle-removal", () => { + it("tracks live requests and reports how many declared a parent", () => { + const registry = new SessionTreeRegistry() + const { entry } = tracker() + + expect(registry.stats()).toEqual({ + tracked: 0, linked: 0, propagations: 0, cancelledDescendants: 0, + }) + + const root = registry.register(entry("r1", "parent")) + const child = registry.register(entry("r2", "child", "parent")) + + expect(registry.stats().tracked).toBe(2) + expect(registry.stats().linked).toBe(1) + + root.release() + child.release() + expect(registry.stats().tracked).toBe(0) + expect(registry.stats().linked).toBe(0) + }) + + it("leaves nothing behind after every request settles", () => { + const registry = new SessionTreeRegistry() + const { entry } = tracker() + + for (let i = 0; i < 50; i++) { + const a = registry.register(entry(`p${i}`, `parent-${i}`)) + const b = registry.register(entry(`c${i}`, `child-${i}`, `parent-${i}`)) + b.release() + a.release() + } + + expect(registry.stats().tracked).toBe(0) + // The parent index must drain too — a leaked empty Set per session key is + // exactly the unbounded growth this registry exists to avoid. + expect(registry.descendantsOf("parent-7")).toEqual([]) + registry.register(entry("late", "child-7", "parent-7")) + expect(registry.descendantsOf("parent-7").map((e) => e.requestId)).toEqual(["late"]) + }) + + it("release is idempotent and does not disturb a sibling", () => { + const registry = new SessionTreeRegistry() + const { entry } = tracker() + + const first = registry.register(entry("r1", "child", "parent")) + registry.register(entry("r2", "child", "parent")) + + first.release() + first.release() + first.release() + + expect(registry.stats().tracked).toBe(1) + expect(registry.descendantsOf("parent").map((e) => e.requestId)).toEqual(["r2"]) + }) + + it("keeps concurrent requests apart when the client reuses one request id", () => { + const registry = new SessionTreeRegistry() + const { aborted, entry } = tracker() + + // x-request-id is client-supplied, so a collision is a client's choice, not + // a proxy invariant. Keying the registry on it would unregister the wrong + // entry and silently drop a live child from cancellation. + const first = registry.register(entry("same-id", "child-a", "parent")) + registry.register(entry("same-id", "child-b", "parent")) + first.release() + + registry.cancelDescendants("parent") + expect(aborted.map((a) => a.requestId)).toEqual(["same-id"]) + expect(registry.descendantsOf("parent").map((e) => e.sessionKey)).toEqual(["child-b"]) + }) +}) + +describe("SessionTreeRegistry — subtree computation", () => { + it("walks a multi-level tree transitively, nearest level first", () => { + const registry = new SessionTreeRegistry() + const { entry } = tracker() + + registry.register(entry("r-root", "root")) + registry.register(entry("r-a", "a", "root")) + registry.register(entry("r-b", "b", "root")) + registry.register(entry("r-a1", "a1", "a")) + registry.register(entry("r-a1x", "a1x", "a1")) + registry.register(entry("r-unrelated", "elsewhere", "other-root")) + + expect(registry.descendantsOf("root").map((e) => e.requestId)) + .toEqual(["r-a", "r-b", "r-a1", "r-a1x"]) + expect(registry.descendantsOf("a").map((e) => e.requestId)).toEqual(["r-a1", "r-a1x"]) + expect(registry.descendantsOf("a1x")).toEqual([]) + expect(registry.descendantsOf("root").some((e) => e.sessionKey === "elsewhere")).toBe(false) + }) + + it("includes every live request sharing one child key", () => { + const registry = new SessionTreeRegistry() + const { entry } = tracker() + + // The turn coordinator serializes turns per key, so a second request on the + // same child key is queued rather than absent — and it is the one a parent + // abort most needs to reach. + registry.register(entry("running", "child", "parent")) + registry.register(entry("queued", "child", "parent")) + registry.register(entry("grandchild", "grandchild", "child")) + + expect(registry.descendantsOf("parent").map((e) => e.requestId)) + .toEqual(["running", "queued", "grandchild"]) + }) + + it("terminates on a cycle stamped by the client", () => { + const registry = new SessionTreeRegistry() + const { entry } = tracker() + + registry.register(entry("r-a", "a", "b")) + registry.register(entry("r-b", "b", "a")) + + expect(registry.descendantsOf("a").map((e) => e.requestId)).toEqual(["r-b", "r-a"]) + expect(registry.descendantsOf("b").map((e) => e.requestId)).toEqual(["r-a", "r-b"]) + }) + + it("ignores a self-link so a request is never its own descendant", () => { + const registry = new SessionTreeRegistry() + const { aborted, entry } = tracker() + + registry.register(entry("r-self", "same", "same")) + + expect(registry.descendantsOf("same")).toEqual([]) + expect(registry.stats().linked).toBe(1) + registry.cancelDescendants("same") + expect(aborted).toEqual([]) + }) + + it("bounds a pathologically deep chain", () => { + const registry = new SessionTreeRegistry() + const { entry } = tracker() + + for (let i = 1; i <= 500; i++) { + registry.register(entry(`r${i}`, `k${i}`, `k${i - 1}`)) + } + + const reached = registry.descendantsOf("k0") + expect(reached.length).toBeGreaterThan(0) + expect(reached.length).toBeLessThan(500) + }) +}) + +describe("SessionTreeRegistry — cancellation", () => { + it("aborts the whole subtree and leaves the origin request alone", () => { + const registry = new SessionTreeRegistry() + const { aborted, entry } = tracker() + + registry.register(entry("r-root", "root")) + registry.register(entry("r-a", "a", "root")) + registry.register(entry("r-a1", "a1", "a")) + + const reason = new Error("parent cancelled") + const result = registry.cancelDescendants("root", reason) + + expect(aborted.map((a) => a.requestId)).toEqual(["r-a", "r-a1"]) + expect(aborted.every((a) => a.reason === reason)).toBe(true) + expect(result.keys).toEqual(["a", "a1"]) + expect(result.requestIds).toEqual(["r-a", "r-a1"]) + }) + + it("counts propagations and cancelled descendants", () => { + const registry = new SessionTreeRegistry() + const { entry } = tracker() + + registry.register(entry("r-a", "a", "root")) + registry.register(entry("r-a1", "a1", "a")) + registry.cancelDescendants("root") + + expect(registry.stats().propagations).toBe(1) + expect(registry.stats().cancelledDescendants).toBe(2) + + // An idle subtree is not a propagation: there was nothing in flight, which + // is the normal case for a client that spawns no subagents. + registry.cancelDescendants("nobody-here") + expect(registry.stats().propagations).toBe(1) + expect(registry.stats().cancelledDescendants).toBe(2) + }) + + it("cancels every sibling even when one abort handle throws", () => { + const registry = new SessionTreeRegistry() + const aborted: string[] = [] + + registry.register({ + requestId: "r-bad", sessionKey: "bad", parentKey: "root", + abort: () => { throw new Error("teardown exploded") }, + }) + registry.register({ + requestId: "r-good", sessionKey: "good", parentKey: "root", + abort: () => { aborted.push("r-good") }, + }) + + const result = registry.cancelDescendants("root") + expect(aborted).toEqual(["r-good"]) + expect(result.requestIds).toEqual(["r-bad", "r-good"]) + }) + + it("cancelSubtree also aborts live requests on the origin key", () => { + const registry = new SessionTreeRegistry() + const { aborted, entry } = tracker() + + registry.register(entry("r-root", "root")) + registry.register(entry("r-a", "a", "root")) + + const result = registry.cancelSubtree("root") + expect(aborted.map((a) => a.requestId)).toEqual(["r-root", "r-a"]) + expect(result.keys).toEqual(["root", "a"]) + // Only the descendant counts as a propagated child cancellation. + expect(registry.stats().cancelledDescendants).toBe(1) + }) + + it("does not reach a request that already settled", () => { + const registry = new SessionTreeRegistry() + const { aborted, entry } = tracker() + + const child = registry.register(entry("r-a", "a", "root")) + child.release() + + expect(registry.cancelDescendants("root")).toEqual({ keys: [], requestIds: [] }) + expect(aborted).toEqual([]) + }) + + it("clear() resets both the live index and the counters", () => { + const registry = new SessionTreeRegistry() + const { entry } = tracker() + + registry.register(entry("r-a", "a", "root")) + registry.cancelDescendants("root") + registry.clear() + + expect(registry.stats()).toEqual({ + tracked: 0, linked: 0, propagations: 0, cancelledDescendants: 0, + }) + }) +}) + +describe("truncateSessionKey", () => { + it("shortens long keys and leaves short ones alone", () => { + expect(truncateSessionKey("0123456789abcdef")).toBe("01234567…") + expect(truncateSessionKey("short")).toBe("short") + expect(truncateSessionKey("0123456789", 4)).toBe("0123…") + }) +}) + +describe("processSessionTree", () => { + it("is a shared registry so a parent and child on different proxy instances still link", () => { + processSessionTree.clear() + const aborted: string[] = [] + const child = processSessionTree.register({ + requestId: "r-child", + sessionKey: "child", + parentKey: "parent", + abort: () => { aborted.push("r-child") }, + }) + + expect(processSessionTree.cancelDescendants("parent").requestIds).toEqual(["r-child"]) + expect(aborted).toEqual(["r-child"]) + + child.release() + processSessionTree.clear() + }) +}) diff --git a/src/proxy/adapter.ts b/src/proxy/adapter.ts index fd4d8b5a..9e647e42 100644 --- a/src/proxy/adapter.ts +++ b/src/proxy/adapter.ts @@ -32,6 +32,21 @@ export interface AgentIdentity { */ getSessionId(c: Context, body?: unknown): string | undefined + /** + * Optional IMMEDIATE parent session key, for clients that declare a subagent + * tree (Prime Agent's RLM children stamp it in `metadata.user_id`). + * + * Two rules make this safe to consume: + * - the value must be a key `getSessionId` could itself have produced, so + * the returned string is directly comparable to another request's key; + * - it must never alter this request's own key. + * + * Deeper trees name one level each, so consumers walk the chain. Returning + * undefined (the default) means the client declares no lineage, which is what + * keeps parent→child cancellation inert for every other client. + */ + getParentSessionId?(c: Context, body?: unknown): string | undefined + /** * Optional client-declared agent mode. Adapters own their header/protocol * details; the proxy uses the normalized value for model-tier selection. diff --git a/src/proxy/adapters/claudecode.ts b/src/proxy/adapters/claudecode.ts index 6c78d962..5a7fd475 100644 --- a/src/proxy/adapters/claudecode.ts +++ b/src/proxy/adapters/claudecode.ts @@ -51,8 +51,34 @@ function extractClaudeCodeClientCwd(body: any): string | undefined { return match?.[1]?.trim() || undefined } -/** Extract the stable conversation ID embedded by Claude Code in metadata.user_id. */ -export function extractClaudeCodeSessionId(body: unknown): string | undefined { +/** + * Session identity declared in `metadata.user_id`. + * + * `sessionId` is the whole of the session key — nothing is appended, prefixed, + * or normalized — because it is what every cached mapping is already stored + * under. `parentSessionId` is additive: a client that does not stamp it gets + * exactly the identity it got before the field existed. + */ +export interface ClaudeCodeSessionIdentity { + readonly sessionId: string + /** + * The IMMEDIATE parent's session id, when the client declares subagent + * lineage. Deeper trees are expressed by each level naming its own parent, so + * consumers walk the chain rather than expecting a root here. + */ + readonly parentSessionId?: string +} + +/** + * Parse the identity envelope Claude Code (and Prime Agent's extension) embeds + * in `metadata.user_id`. + * + * Strict by design: `user_id` must be, or parse to, an object carrying a + * non-empty string `session_id`. Anything else yields undefined and the caller + * falls back to fingerprint resume, so unrelated Anthropic-API clients that put + * their own value in `user_id` are never mistaken for a keyed session. + */ +export function extractClaudeCodeSessionIdentity(body: unknown): ClaudeCodeSessionIdentity | undefined { if (!body || typeof body !== "object") return undefined const metadata = (body as { metadata?: unknown }).metadata @@ -71,7 +97,28 @@ export function extractClaudeCodeSessionId(body: unknown): string | undefined { if (!userMetadata || typeof userMetadata !== "object") return undefined const sessionId = (userMetadata as { session_id?: unknown }).session_id - return typeof sessionId === "string" && sessionId.length > 0 ? sessionId : undefined + if (typeof sessionId !== "string" || sessionId.length === 0) return undefined + + const parentSessionId = (userMetadata as { parent_session_id?: unknown }).parent_session_id + // A node that names itself as its own parent is not a tree edge, and treating + // it as one would make a request its own cancellation target. + const parent = typeof parentSessionId === "string" + && parentSessionId.length > 0 + && parentSessionId !== sessionId + ? parentSessionId + : undefined + + return parent ? { sessionId, parentSessionId: parent } : { sessionId } +} + +/** Extract the stable conversation ID embedded by Claude Code in metadata.user_id. */ +export function extractClaudeCodeSessionId(body: unknown): string | undefined { + return extractClaudeCodeSessionIdentity(body)?.sessionId +} + +/** Extract the immediate parent session key, when the client declares lineage. */ +export function extractClaudeCodeParentSessionId(body: unknown): string | undefined { + return extractClaudeCodeSessionIdentity(body)?.parentSessionId } export const claudeCodeAdapter: AgentAdapter = { @@ -85,6 +132,14 @@ export const claudeCodeAdapter: AgentAdapter = { return extractClaudeCodeSessionId(body) }, + /** + * Subagent lineage from the same envelope that supplied the session key, so + * a declared parent always names a key derived the same way this one was. + */ + getParentSessionId(_c: Context, body?: unknown): string | undefined { + return extractClaudeCodeParentSessionId(body) + }, + /** * Claude Code is remote relative to the proxy. Do not use its local path * as the SDK subprocess cwd — return undefined so the resolver falls back diff --git a/src/proxy/adapters/prime.ts b/src/proxy/adapters/prime.ts index c1490bee..d1e433ca 100644 --- a/src/proxy/adapters/prime.ts +++ b/src/proxy/adapters/prime.ts @@ -33,7 +33,7 @@ import { type FileChange, extractFileChangesFromBash } from "../fileChanges" import { normalizeContent } from "../messages" import { BLOCKED_BUILTIN_TOOLS, CLAUDE_CODE_ONLY_TOOLS } from "../tools" import { resolvePassthrough } from "../../env" -import { extractClaudeCodeSessionId } from "./claudecode" +import { extractClaudeCodeParentSessionId, extractClaudeCodeSessionId } from "./claudecode" const PRIME_MCP_SERVER_NAME = "prime" @@ -224,6 +224,23 @@ export const primeAdapter: AgentAdapter = { return c.req.header("x-session-affinity") ?? extractClaudeCodeSessionId(body) }, + /** + * RLM subagent lineage: `metadata.user_id` carries + * `{ session_id, parent_session_id }`, where the parent is the IMMEDIATE one + * (`ctx.sessionManager.getParentSessionId()` in the Prime Agent extension). + * The proxy uses it to cancel a whole live subtree when the parent's request + * is aborted — see `sessionTree.ts`. + * + * Reported only when the session key itself came from this envelope. An + * orchestrator that overrides identity with `x-session-affinity` is naming + * keys under a different scheme, so a parent id read out of the body would + * point at a key that scheme never produced. + */ + getParentSessionId(c: Context, body?: unknown): string | undefined { + if (c.req.header("x-session-affinity")) return undefined + return extractClaudeCodeParentSessionId(body) + }, + extractWorkingDirectory(body: any): string | undefined { return extractPrimeCwd(body) }, diff --git a/src/proxy/server.ts b/src/proxy/server.ts index 3f02e210..fb3ba774 100644 --- a/src/proxy/server.ts +++ b/src/proxy/server.ts @@ -9,6 +9,7 @@ import { query } from "@anthropic-ai/claude-agent-sdk" import { rateLimitStore } from "./rateLimitStore" import { guardUpstreamIdle, UpstreamIdleError } from "./streamIdleGuard" import { linkRequestAbort } from "./requestAbort" +import { processSessionTree, truncateSessionKey, type SessionTreeRegistration } from "./sessionTree" import { AbortableSemaphore, getProcessSdkSemaphore, type SemaphoreLease } from "./concurrency" import { closeServerWithGracePeriod, trackServerConnections } from "./shutdown" import { fetchOAuthUsage, fetchOAuthUsageResult } from "./oauthUsage" @@ -259,6 +260,15 @@ interface RequestMeta { } /** Permanently retain the session lease when mandatory durable cleanup fails. */ retainSessionTurnFence?: () => void + /** + * Cancel this request's live session subtree (see `sessionTree.ts`). + * + * Present only for a keyed request. Called from the abort paths a CLIENT can + * reach — a request-signal abort and a cancelled response body — never from a + * turn that merely completed. Latches after the first call, so one client + * teardown that trips both paths propagates once. + */ + cascadeSubtreeCancel?: (source: string) => void } interface PriorityAttemptExposure { @@ -1311,7 +1321,7 @@ export function createProxyServer(config: Partial = {}): ProxyServe status: "ok", service: "meridian", format: "anthropic", - endpoints: ["/v1/messages", "/messages", "/v1/chat/completions", "/v1/responses", "/v1/models", "/v1/design/*", "/design-login", "/telemetry", "/metrics", "/health"] + endpoints: ["/v1/messages", "/messages", "/v1/chat/completions", "/v1/responses", "/v1/models", "/v1/sessions/:key/cancel", "/v1/design/*", "/design-login", "/telemetry", "/metrics", "/health"] }) } return c.html(landingHtml) @@ -6183,6 +6193,11 @@ export function createProxyServer(config: Partial = {}): ProxyServe cancel(reason) { requestAbort.abort(reason) requestAbort.detach() + // A cancelled response body is the other way a client says "stop", + // and the only one an in-process caller can reach. Children are + // cancelled here as well, latched so a real socket teardown — + // which trips both this and the request signal — propagates once. + requestMeta.cascadeSubtreeCancel?.("stream_cancel") if (!isIndependentSession && ( !managedForkTarget || managedForkPublished || clientAssistantContentExposed )) { @@ -6305,6 +6320,35 @@ export function createProxyServer(config: Partial = {}): ProxyServe }) } + /** + * Report one propagated subtree cancellation. + * + * Split out so the two client-reachable abort paths (request signal, cancelled + * response body) log identically and differ only in `source`. + */ + const logSubtreeCancel = ( + parentKey: string, + cancelled: { readonly keys: readonly string[]; readonly requestIds: readonly string[] }, + requestId: string, + source: string, + ): void => { + const children = cancelled.keys.map((key) => truncateSessionKey(key)) + claudeLog("session.tree_cancel_propagated", { + requestId, + source, + parent: truncateSessionKey(parentKey), + children, + requests: cancelled.requestIds.length, + }) + // Named at session level: an autonomous run has nobody watching the + // dashboard, and this is the event that explains why a child turn died. + diagnosticLog.session( + `${requestId} session_tree_cancel source=${source} parent=${truncateSessionKey(parentKey)} ` + + `children=${children.join(",")} requests=${cancelled.requestIds.length}`, + requestId, + ) + } + const handleWithQueue = async (c: Context, endpoint: string) => { // An internal hop carries a request the public route already admitted; // re-checking the gate here would refuse work that is legitimately in @@ -6317,6 +6361,39 @@ export function createProxyServer(config: Partial = {}): ProxyServe claudeLog("request.enter", { requestId, endpoint }) let sessionTurnLease: SessionTurnLease | undefined let crossProcessTurnLease: CrossProcessTurnLease | undefined + let sessionTreeRegistration: SessionTreeRegistration | undefined + let detachSubtreeAbortWatch: (() => void) | undefined + let subtreeSessionKey: string | undefined + let subtreeCascaded = false + /** + * Cancel this request's live session subtree. + * + * Scope is deliberately narrow. Only a CLIENT abort propagates: a parent + * turn that merely completes leaves its children running, because a + * subagent routinely outlives the turn that spawned it. The shutdown path + * already aborts every in-flight request directly, and the lease watchdog + * is a proxy-side fence rather than a user intent, so neither cascades. + * + * Each child is aborted through its OWN request abort controller — the same + * one the lease watchdog and `forceAbortInFlight` use — so the mapping + * eviction, SDK permit release, and turn-lease release that follow are the + * existing abort path's, not a second implementation of it. + * + * Latched unconditionally: one client teardown can trip both the request + * signal and the response-body cancel, and every child that mattered was + * already in flight when the first of them fired. + */ + const cascadeSubtreeCancel = (source: string): void => { + const parentKey = subtreeSessionKey + if (subtreeCascaded || !parentKey) return + subtreeCascaded = true + const cancelled = processSessionTree.cancelDescendants( + parentKey, + new Error(`Parent session ${truncateSessionKey(parentKey)} was cancelled`), + ) + if (cancelled.requestIds.length === 0) return + logSubtreeCancel(parentKey, cancelled, requestId, source) + } const turnWatchdogAbort = new AbortController() activeRequestAborts.add(turnWatchdogAbort) let finished = false @@ -6355,6 +6432,12 @@ export function createProxyServer(config: Partial = {}): ProxyServe } else { releaseSessionTurn(false) } + // The session tree is live requests only: a settled request is no longer + // a cancellation target, and a settled parent no longer cascades. + detachSubtreeAbortWatch?.() + detachSubtreeAbortWatch = undefined + sessionTreeRegistration?.release() + sessionTreeRegistration = undefined activeRequestAborts.delete(turnWatchdogAbort) inFlightRequests-- } @@ -6389,6 +6472,25 @@ export function createProxyServer(config: Partial = {}): ProxyServe routingTurnIdentity = adapter.getRoutingTurnIdentity?.(c, body) const agentSessionId = adapter.getSessionId(c, body) if (agentSessionId) { + // Registered BEFORE the turn lease is acquired: a child queued behind + // its own session's running turn is exactly the request a parent abort + // most needs to reach, and the acquire wait already honors this + // controller's signal. + sessionTreeRegistration = processSessionTree.register({ + requestId, + sessionKey: agentSessionId, + parentKey: adapter.getParentSessionId?.(c, body), + abort: (reason) => turnWatchdogAbort.abort(reason), + }) + subtreeSessionKey = agentSessionId + const clientSignal = c.req.raw.signal + if (clientSignal.aborted) { + cascadeSubtreeCancel("client_abort") + } else { + const onClientAbort = () => cascadeSubtreeCancel("client_abort") + clientSignal.addEventListener("abort", onClientAbort, { once: true }) + detachSubtreeAbortWatch = () => clientSignal.removeEventListener("abort", onClientAbort) + } const arrivalProfileIds = new Set( getEffectiveProfiles(finalConfig.profiles).map((profile) => profile.id), ) @@ -6486,6 +6588,7 @@ export function createProxyServer(config: Partial = {}): ProxyServe sharedSessionRevisionsAtArrival, routingTurnIdentity, retainSessionTurnFence: () => { retainSessionTurnFence = true }, + cascadeSubtreeCancel, } const response = await handleMessages(c, requestMeta, { body, @@ -6510,8 +6613,43 @@ export function createProxyServer(config: Partial = {}): ProxyServe app.post("/v1/messages", (c) => handleWithQueue(c, "/v1/messages")) app.post("/messages", (c) => handleWithQueue(c, "/messages")) + /** + * Cancel a session's live requests and everything live below it. + * + * The same registry that powers abort propagation answers this for free, so a + * harness that wants to stop a subtree without tearing down sockets has a way + * to say so. It cancels only what is IN FLIGHT — there is no persistent tree, + * so an idle session reports `requests: 0` rather than being remembered. + * + * Gated by the `/v1/*` auth middleware like every other `/v1` route. + */ + app.post("/v1/sessions/:key/cancel", (c) => { + const key = c.req.param("key") + if (!key) { + return c.json({ type: "error", error: { type: "invalid_request_error", message: "Session key is required" } }, 400) + } + const cancelled = processSessionTree.cancelSubtree(key, new Error("Session cancelled by request")) + if (cancelled.requestIds.length > 0) { + claudeLog("session.tree_cancel_requested", { + session: truncateSessionKey(key), + keys: cancelled.keys.map((cancelledKey) => truncateSessionKey(cancelledKey)), + requests: cancelled.requestIds.length, + }) + diagnosticLog.session( + `session_tree_cancel_requested session=${truncateSessionKey(key)} requests=${cancelled.requestIds.length}`, + ) + } + return c.json({ + session: key, + cancelled: { sessions: cancelled.keys.length, requests: cancelled.requestIds.length }, + requestIds: cancelled.requestIds, + }) + }) + // Telemetry dashboard and API - app.route("/telemetry", createTelemetryRoutes()) + app.route("/telemetry", createTelemetryRoutes({ + getSessionTree: () => processSessionTree.stats(), + })) // SDK Features settings page and API app.get("/settings", (c) => { diff --git a/src/proxy/sessionTree.ts b/src/proxy/sessionTree.ts new file mode 100644 index 00000000..64c8d63d --- /dev/null +++ b/src/proxy/sessionTree.ts @@ -0,0 +1,250 @@ +/** + * Live session-tree registry: parent→child linkage for in-flight requests. + * + * Cancellation in meridian is per-HTTP-request — `requestAbort.ts` forwards one + * socket's abort into that request's SDK abort controller. A harness that spawns + * subagents (Prime Agent's RLM children) sends each child as an INDEPENDENT + * request on its own session key, so a user cancelling the parent left every + * child running: holding an SDK permit, holding its turn lease, and billing the + * subscription until its own socket closed or the lease watchdog tripped. + * + * This registry is the missing link. A client that knows its own tree stamps the + * immediate parent alongside the child's session id (`metadata.user_id` → + * `{ session_id, parent_session_id }`, read by `adapters/claudecode.ts`), and + * the server registers that linkage for the lifetime of the request. + * + * Three properties are deliberate: + * + * 1. **Live requests only.** An entry exists between "request admitted" and + * "request settled", nothing longer. There is no persistent tree: a session + * that was seen once but has no request in flight is not a cancellation + * target, because there is nothing to cancel. That keeps the registry + * bounded by concurrency rather than by conversation history. + * 2. **Abort, not completion.** Only an actual abort of a node propagates. A + * parent turn that finishes normally leaves its children alone — a child + * routinely outlives the parent turn that spawned it. + * 3. **Self-gating.** Propagation can only reach a request that declared a + * parent, so a client that does not stamp linkage is unaffected with no + * config flag to set. + * + * Pure bookkeeping: no HTTP, no I/O, no logging. The caller supplies the abort + * handle and owns the eviction/telemetry discipline that follows an abort. + */ + +/** Handle returned by `register`; removes the entry when the request settles. */ +export interface SessionTreeRegistration { + /** Remove this request from the registry. Idempotent. */ + release(): void +} + +export interface SessionTreeEntry { + /** Request id, for logging and for the cancellation result. */ + readonly requestId: string + /** This request's client-session key, exactly as the adapter derived it. */ + readonly sessionKey: string + /** The IMMEDIATE parent's session key, when the client stamped linkage. */ + readonly parentKey?: string + /** + * Abort this request. Must route through the same abort path a client + * disconnect uses, so the eviction, permit release, and lease release that + * follow are the existing ones rather than a parallel implementation. + */ + readonly abort: (reason?: unknown) => void +} + +export interface SessionTreeStats { + /** Live requests currently registered. */ + tracked: number + /** Of those, how many declared a parent (i.e. are cancellation targets). */ + linked: number + /** Cancellations that aborted at least one live request, since start. */ + propagations: number + /** Descendant requests aborted by an ancestor's cancellation, since start. */ + cancelledDescendants: number +} + +/** What a cancellation actually reached. Empty when the subtree was idle. */ +export interface SessionTreeCancellation { + /** Session keys whose live requests were aborted, nearest-first. */ + readonly keys: readonly string[] + /** Request ids aborted, in the order they were aborted. */ + readonly requestIds: readonly string[] +} + +const EMPTY_CANCELLATION: SessionTreeCancellation = { keys: [], requestIds: [] } + +/** + * Depth cap for the ancestry walk. + * + * `parentKey` comes off the wire, so the forest is only a forest by convention: + * a buggy or hostile client can stamp a cycle (A→B→A) or an absurdly deep chain. + * The visited set already makes cycles terminate; this bounds the honest-but- + * pathological case so one cancellation can never walk unboundedly. + */ +const MAX_SUBTREE_DEPTH = 64 + +/** Shorten a session key for logs. Keys can be full UUIDs or client-chosen. */ +export function truncateSessionKey(key: string, length = 8): string { + return key.length > length ? `${key.slice(0, length)}…` : key +} + +interface CancelOptions { + /** Also abort live requests running under the origin key itself. */ + readonly includeSelf?: boolean + readonly reason?: unknown +} + +/** + * Registry of live requests and their declared parent links. + * + * Keyed by an internal token rather than by request id: `x-request-id` is + * client-supplied, so two concurrent requests can legitimately arrive carrying + * the same one, and a colliding key would silently unregister the wrong entry. + */ +export class SessionTreeRegistry { + private nextToken = 1 + private readonly entries = new Map() + /** parentKey → tokens of live children. Index for the subtree walk. */ + private readonly childrenByParent = new Map>() + private propagations = 0 + private cancelledDescendants = 0 + + register(entry: SessionTreeEntry): SessionTreeRegistration { + const token = this.nextToken++ + this.entries.set(token, entry) + // A self-link is meaningless and would make a node its own descendant, so + // it is dropped at the index rather than defended against on every walk. + const indexedParent = entry.parentKey && entry.parentKey !== entry.sessionKey + ? entry.parentKey + : undefined + if (indexedParent) { + let siblings = this.childrenByParent.get(indexedParent) + if (!siblings) { + siblings = new Set() + this.childrenByParent.set(indexedParent, siblings) + } + siblings.add(token) + } + let released = false + return { + release: () => { + if (released) return + released = true + this.entries.delete(token) + if (!indexedParent) return + const siblings = this.childrenByParent.get(indexedParent) + if (!siblings) return + siblings.delete(token) + if (siblings.size === 0) this.childrenByParent.delete(indexedParent) + }, + } + } + + /** Live requests whose ancestry chain reaches `sessionKey`, nearest first. */ + descendantsOf(sessionKey: string): SessionTreeEntry[] { + const visitedKeys = new Set([sessionKey]) + let frontier = [sessionKey] + const found: SessionTreeEntry[] = [] + for (let depth = 0; depth < MAX_SUBTREE_DEPTH && frontier.length > 0; depth++) { + const next: string[] = [] + for (const parentKey of frontier) { + const tokens = this.childrenByParent.get(parentKey) + if (!tokens) continue + for (const token of tokens) { + const entry = this.entries.get(token) + if (!entry) continue + found.push(entry) + // Several live requests can share one child key (a queued turn behind + // the running one). Descend through that key only once. + if (visitedKeys.has(entry.sessionKey)) continue + visitedKeys.add(entry.sessionKey) + next.push(entry.sessionKey) + } + } + frontier = next + } + return found + } + + /** Live requests running under `sessionKey` itself. */ + liveRequestsFor(sessionKey: string): SessionTreeEntry[] { + const found: SessionTreeEntry[] = [] + for (const entry of this.entries.values()) { + if (entry.sessionKey === sessionKey) found.push(entry) + } + return found + } + + /** + * Abort every live request below `sessionKey`, transitively. + * + * The origin request is left alone: it is already being torn down by whatever + * triggered this, and aborting it a second time would be a no-op at best. + */ + cancelDescendants(sessionKey: string, reason?: unknown): SessionTreeCancellation { + return this.cancel(sessionKey, { reason }) + } + + /** Abort live requests for `sessionKey` AND everything below it. */ + cancelSubtree(sessionKey: string, reason?: unknown): SessionTreeCancellation { + return this.cancel(sessionKey, { reason, includeSelf: true }) + } + + private cancel(sessionKey: string, options: CancelOptions): SessionTreeCancellation { + const descendants = this.descendantsOf(sessionKey) + const targets = options.includeSelf + ? [...this.liveRequestsFor(sessionKey), ...descendants] + : descendants + if (targets.length === 0) return EMPTY_CANCELLATION + + const keys: string[] = [] + const requestIds: string[] = [] + for (const entry of targets) { + // One entry's abort handle must never strand its siblings. The caller's + // handle runs arbitrary teardown; a throw here is its problem, not the + // rest of the subtree's. + try { + entry.abort(options.reason) + } catch { + // Recorded as cancelled regardless: the entry was targeted, and its + // request settles through its own path either way. + } + if (!keys.includes(entry.sessionKey)) keys.push(entry.sessionKey) + requestIds.push(entry.requestId) + } + this.propagations++ + this.cancelledDescendants += descendants.length + return { keys, requestIds } + } + + stats(): SessionTreeStats { + let linked = 0 + for (const entry of this.entries.values()) { + if (entry.parentKey) linked++ + } + return { + tracked: this.entries.size, + linked, + propagations: this.propagations, + cancelledDescendants: this.cancelledDescendants, + } + } + + /** Reset process-scoped state between tests. */ + clear(): void { + this.entries.clear() + this.childrenByParent.clear() + this.propagations = 0 + this.cancelledDescendants = 0 + } +} + +/** + * Process-wide registry. + * + * Session keys are a process-wide namespace (`processSessionTurns` serializes + * turns on the same basis), and a parent and its child can be served by + * different ProxyServer instances in one process, so per-instance registries + * would lose exactly the links that matter. + */ +export const processSessionTree = new SessionTreeRegistry() diff --git a/src/telemetry/dashboard.ts b/src/telemetry/dashboard.ts index a3594ced..ae9e63d5 100644 --- a/src/telemetry/dashboard.ts +++ b/src/telemetry/dashboard.ts @@ -216,6 +216,13 @@ function render(s, reqs, logs) { + card('Median TTFB', ms(s.ttfb.p50), 'p95: ' + ms(s.ttfb.p95)) + card('Proxy Overhead', ms(s.proxyOverhead.p50), 'p95: ' + ms(s.proxyOverhead.p95)) + card('Queue Wait', ms(s.queueWait.p50), 'p95: ' + ms(s.queueWait.p95)) + // Only rendered when a subagent tree has actually been seen: for a + // single-agent client these numbers are permanently zero and would just be + // a dead tile. Counts are cumulative, not windowed. + + ((s.sessionTree && (s.sessionTree.linked > 0 || s.sessionTree.cancelledDescendants > 0)) + ? card('Subtree Cancels', s.sessionTree.cancelledDescendants, + s.sessionTree.linked + ' linked live / ' + s.sessionTree.propagations + ' propagations') + : '') + ''; // Token usage cards diff --git a/src/telemetry/routes.ts b/src/telemetry/routes.ts index bfe9438e..98e74d1a 100644 --- a/src/telemetry/routes.ts +++ b/src/telemetry/routes.ts @@ -13,12 +13,22 @@ import { fileURLToPath } from "node:url" import { Hono } from "hono" import { telemetryStore, diagnosticLog } from "./index" import { dashboardHtml } from "./dashboard" +import type { SessionTreeSummary } from "./types" // Read once at module load — src/telemetry/ is two levels below the package root const _iconPath = resolve(dirname(fileURLToPath(import.meta.url)), "..", "..", "assets", "icon.svg") const _iconSvg = existsSync(_iconPath) ? readFileSync(_iconPath, "utf-8") : null -export function createTelemetryRoutes() { +export interface TelemetryRouteDeps { + /** + * Live session-tree cancellation counters, injected by the proxy so this + * module keeps depending only on the telemetry store. Optional: embedders and + * unit tests mount the routes standalone and simply get no `sessionTree`. + */ + getSessionTree?: () => SessionTreeSummary +} + +export function createTelemetryRoutes(deps: TelemetryRouteDeps = {}) { const routes = new Hono() // Dashboard @@ -55,7 +65,8 @@ export function createTelemetryRoutes() { const windowMs = Number.parseInt(c.req.query("window") || "3600000", 10) // default 1 hour const summary = telemetryStore.summarize(windowMs) - return c.json(summary) + const sessionTree = deps.getSessionTree?.() + return c.json(sessionTree ? { ...summary, sessionTree } : summary) }) // Diagnostic logs diff --git a/src/telemetry/types.ts b/src/telemetry/types.ts index 96198129..448a1318 100644 --- a/src/telemetry/types.ts +++ b/src/telemetry/types.ts @@ -206,6 +206,23 @@ export interface TelemetrySummary { costEstimate: CostEstimate } +/** + * Live parent→child cancellation counters, mirrored from + * `proxy/sessionTree.ts` onto `GET /telemetry/summary`. + * + * Gauges (`tracked`, `linked`) are instantaneous; counts (`propagations`, + * `cancelledDescendants`) are cumulative since proxy start, so they are NOT + * scoped to the summary's time window. Orphaned subagents were previously + * invisible — a nonzero `cancelledDescendants` is the proof teardown reached + * them. + */ +export interface SessionTreeSummary { + tracked: number + linked: number + propagations: number + cancelledDescendants: number +} + /** Storage backend for request metrics. */ export interface ITelemetryStore { /** Record a completed request metric. */