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. */