@nijaru/pi-provider-compaction 0.0.8 → 0.0.10

package/protocol.ts

@@ -208,7 +208,7 @@ async function collectProviderUsage(stream: AssistantMessageEventStream, signal:
 	let usage: Usage | undefined;
 	for await (const event of stream) {
 		if (signal.aborted) throw new DOMException("Compaction aborted", "AbortError");
-		if (event.type === "error") throw new Error(event.error.errorMessage ?? "Native compaction provider request failed");
+		if (event.type === "error") throw new NativeCompactionValidationError("Native compaction provider stream failed", event.error.usage);
 		if (event.type === "done") usage = event.message.usage;
 	}
 	if (!usage) throw new Error("Native compaction provider stream ended without usage");
@@ -227,10 +227,17 @@ interface TransportRequest {
 	env?: Record<string, string>;
 	/** Cost-relevant request option captured with the original request. */
 	serviceTier?: string;
+	sessionId?: string;
 	signal: AbortSignal;
+	/** Revalidate request ownership after asynchronous provider preparation. */
+	isCurrent?: () => boolean;
 	fetch?: typeof globalThis.fetch;
 }
 
+function assertCurrent(request: TransportRequest): void {
+	if (request.signal.aborted || request.isCurrent?.() === false) throw new DOMException("Compaction invalidated", "AbortError");
+}
+
 /** The provider answered and was billed, but its window cannot be committed. */
 export class NativeCompactionValidationError extends Error {
 	constructor(message: string, readonly usage?: Usage) {
@@ -246,12 +253,12 @@ export interface TransportResponse {
 
 async function requestResponsesCompact(request: TransportRequest): Promise<TransportResponse> {
 	let preparedPayload: JsonObject | undefined;
-	let compacted: { response: JsonObject; output: ResponseItem[] } | undefined;
+	let compacted: ResponseItem[] | undefined;
 	let bridgeError: unknown;
 	let dispatches = 0;
 	const baseFetch = request.fetch ?? globalThis.fetch;
 	const bridgeFetch = (async (input: RequestInfo | URL, init?: RequestInit) => {
-		if (request.signal.aborted) throw new DOMException("Compaction aborted", "AbortError");
+		assertCurrent(request);
 		dispatches += 1;
 		if (dispatches !== 1 || !preparedPayload) {
 			bridgeError = new Error("Responses Compact provider dispatched an unexpected additional request");
@@ -268,13 +275,13 @@ async function requestResponsesCompact(request: TransportRequest): Promise<Trans
 			});
 			if (!response.ok) return response;
 			const parsed = await readJsonBounded(response, request.signal);
-			try {
-				if (!isObject(parsed)) throw new Error("Responses Compact returned a non-object response");
-				compacted = { response: parsed, output: validateCompactedOutput(parsed) };
-			} catch (error) {
-				throw new NativeCompactionValidationError(error instanceof Error ? error.message : String(error));
-			}
-			return syntheticCompletion(parsed, preparedPayload);
+			if (!isObject(parsed)) throw new Error("Responses Compact returned a non-object response");
+			// Let the SDK normalize and price reported usage even when the returned
+			// checkpoint is invalid. A synthetic HTTP error would discard that bill.
+			const completion = syntheticCompletion(parsed, preparedPayload);
+			try { compacted = validateCompactedOutput(parsed); }
+			catch (error) { bridgeError = new NativeCompactionValidationError(error instanceof Error ? error.message : String(error)); }
+			return completion;
 		} catch (error) {
 			bridgeError = error;
 			return errorResponse(error);
@@ -290,8 +297,10 @@ async function requestResponsesCompact(request: TransportRequest): Promise<Trans
 		transport: "sse",
 		cacheRetention: "none",
 		maxRetries: 0,
+		sessionId: request.sessionId,
 		fetch: bridgeFetch,
 		onPayload: (payload: unknown) => {
+			assertCurrent(request);
 			if (!isObject(payload)) throw new Error("Responses provider exposed a non-object payload");
 			if (request.preparedPayload && request.preparedPayload.model !== payload.model) throw new Error("Native route model changed since capture");
 			preparedPayload = { ...structuredClone(request.preparedPayload ?? payload), input: structuredClone(request.input) };
@@ -300,13 +309,13 @@ async function requestResponsesCompact(request: TransportRequest): Promise<Trans
 	});
 
 	const usage = await collectProviderUsage(stream, request.signal).catch((error) => {
-		if (bridgeError) throw bridgeError;
+		if (bridgeError) throw new NativeCompactionValidationError("Native compaction transport or validation failed", error instanceof NativeCompactionValidationError ? error.usage : undefined);
 		throw error;
 	});
 	if (bridgeError instanceof NativeCompactionValidationError) throw new NativeCompactionValidationError(bridgeError.message, usage);
 	if (bridgeError) throw bridgeError;
 	if (!compacted || dispatches !== 1) throw new Error("Responses Compact did not complete exactly one request");
-	return { output: compacted.output, usage };
+	return { output: compacted, usage };
 }
 
 interface CollectedSse {
@@ -408,6 +417,7 @@ async function requestRemoteV2(request: TransportRequest): Promise<TransportResp
 	const signal = AbortSignal.any([request.signal, controller.signal]);
 	const inspections: Promise<{ result: CollectedSse } | { error: unknown }>[] = [];
 	const inspectedFetch = (async (input: RequestInfo | URL, init?: RequestInit) => {
+		assertCurrent(request);
 		const response = await baseFetch(input, init);
 		if (!response.ok || !response.body) return response;
 		const [providerBody, inspectionBody] = response.body.tee();
@@ -430,8 +440,10 @@ async function requestRemoteV2(request: TransportRequest): Promise<TransportResp
 		transport: "sse",
 		cacheRetention: "none",
 		maxRetries: 0,
+		sessionId: request.sessionId,
 		fetch: inspectedFetch,
 		onPayload: (payload: unknown) => {
+			assertCurrent(request);
 			if (!isObject(payload)) throw new Error("Codex provider exposed a non-object payload");
 			sentInput = structuredClone(request.input);
 			if (request.preparedPayload && request.preparedPayload.model !== payload.model) throw new Error("Native route model changed since capture");
@@ -440,9 +452,9 @@ async function requestRemoteV2(request: TransportRequest): Promise<TransportResp
 	});
 	try {
 		const usage = await collectProviderUsage(stream, request.signal);
-		if (!sentInput || inspections.length !== 1) throw new Error(`Remote V2 observed ${inspections.length} successful provider responses; expected exactly one`);
+		if (!sentInput || inspections.length !== 1) throw new NativeCompactionValidationError(`Remote V2 observed ${inspections.length} successful provider responses; expected exactly one`, usage);
 		const inspection = await inspections[0]!;
-		if ("error" in inspection) throw inspection.error;
+		if ("error" in inspection) throw new NativeCompactionValidationError("Remote V2 checkpoint validation failed", usage);
 		return { output: buildRemoteV2History(sentInput, inspection.result.item), usage };
 	} finally {
 		controller.abort();
@@ -452,6 +464,7 @@ async function requestRemoteV2(request: TransportRequest): Promise<TransportResp
 export async function requestProviderCompaction(
 	request: TransportRequest & { protocol: NativeProtocol },
 ): Promise<TransportResponse> {
+	assertCurrent(request);
 	const result = await (request.protocol === "remote-v2" ? requestRemoteV2(request) : requestResponsesCompact(request));
 	try {
 		return { ...result, output: validateCompactedOutput(result) };

package/provider.ts

@@ -1,10 +1,27 @@
-import type { Provider, StreamOptions, Model, TranscriptContext } from "@earendil-works/pi-ai";
+import { createAssistantMessageEventStream, type AssistantMessage, type AssistantMessageEventStream, type CacheRetention, type Provider, type StreamOptions, type Model, type TranscriptContext } from "@earendil-works/pi-ai";
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+interface UsageScope { sessionId?: string; }
+const USAGE_OWNERS = Symbol.for("@nijaru/pi-provider-compaction/usage-owners");
+const usageGlobal = globalThis as typeof globalThis & { [USAGE_OWNERS]?: WeakMap<object, string> };
+const usageOwners = usageGlobal[USAGE_OWNERS] ??= new WeakMap<object, string>();
 
-export type ObserveRequest = (model: Model<any>, context: TranscriptContext, options: StreamOptions) => ((payload: unknown) => unknown | Promise<unknown>) | undefined;
+/** Associate the SDK-owned header record without reading or changing any headers. */
+export function bindUsageHeaders(headers: object, sessionId: string): void {
+	usageOwners.set(headers, sessionId);
+}
+
+export type PayloadObserver = ((payload: unknown) => unknown | Promise<unknown>) & {
+	complete?: (message: AssistantMessage) => void;
+	/** Adjust only usage before any completion observer or consumer sees it. */
+	priceUsage?: (message: AssistantMessage) => void;
+	cacheRetention?: CacheRetention;
+};
+export type ObserveRequest = (model: Model<any>, context: TranscriptContext, options: StreamOptions, scope: UsageScope) => PayloadObserver | undefined;
 
 export interface InstalledObserver {
 	provider: string;
+	/** installObserver leases one dispatch API; lower-level leases can be unscoped. */
+	api?: string;
 	isInstalled(): boolean;
 	remove(): void;
 }
@@ -13,26 +30,148 @@ export interface InstalledObserver {
 // streams: this observer delegates unchanged and can release only its own slot.
 export const STREAM_OBSERVER = Symbol.for("@nijaru/pi-provider-compaction/stream-observer");
 
-function optionsWithObserver<T extends StreamOptions>(model: Model<any>, context: TranscriptContext, options: T | undefined, observe: ObserveRequest): T | undefined {
+function observedStream<T extends StreamOptions>(model: Model<any>, context: TranscriptContext, options: T | undefined, observe: ObserveRequest, dispatch: (options: T | undefined) => AssistantMessageEventStream): AssistantMessageEventStream {
+	const headerOwner = options?.headers ? usageOwners.get(options.headers) : undefined;
+	const scope: UsageScope = { sessionId: headerOwner ?? options?.sessionId };
 	let handler: ReturnType<ObserveRequest>;
-	try { handler = observe(model, context, options ?? {}); }
-	catch { return options; } // An expired or failing optional observer must not veto dispatch.
-	if (!handler) return options;
-	return {
+	try { handler = observe(model, context, options ?? {}, scope); }
+	catch { return dispatch(options); } // An expired optional observer must not veto dispatch.
+	if (!handler) return dispatch(options);
+	const run = () => dispatch({
 		...options,
+		...(handler.cacheRetention === undefined ? {} : { cacheRetention: handler.cacheRetention }),
 		onPayload: async (payload: unknown, requestModel: Model<any>) => {
 			const replacement = await options?.onPayload?.(payload, requestModel);
 			const prepared = replacement === undefined ? payload : replacement;
 			// Observation failures must not suppress or alter the portable request.
 			try { return await handler(prepared); } catch { return prepared; }
 		},
-	} as T;
+	} as T);
+	const stream = run();
+	if (handler.priceUsage) return pricedStream(stream, model, handler);
+	// Pure capture remains transparent, including the original stream identity.
+	if (handler.complete) void stream.result().then(handler.complete).catch(() => {});
+	return stream;
+}
+
+/** Pricing must precede both terminal-event delivery and result() resolution. */
+function pricedStream(stream: AssistantMessageEventStream, model: Model<any>, handler: PayloadObserver): AssistantMessageEventStream {
+	const out = createAssistantMessageEventStream();
+	void (async () => {
+		let final: AssistantMessage | undefined;
+		let partial: AssistantMessage | undefined;
+		try {
+			for await (const event of stream) {
+				if ("partial" in event) partial = event.partial;
+				if (event.type === "done" || event.type === "error") {
+					final = event.type === "done" ? event.message : event.error;
+					try { handler.priceUsage?.(final); } catch { /* Optional estimates never veto completion. */ }
+					try { handler.complete?.(final); } catch { /* Capture failure never vetoes completion. */ }
+				}
+				out.push(event);
+			}
+			if (final) { out.end(final); return; }
+		} catch { /* Malformed plugin streams still need a terminal outcome. */ }
+		const message: AssistantMessage = {
+			...(partial ?? {
+				role: "assistant", api: model.api, provider: model.provider, model: model.id,
+				content: [], timestamp: Date.now(),
+				usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
+			}),
+			stopReason: "error", errorMessage: "Provider stream ended before a terminal event",
+		};
+		try { handler.priceUsage?.(message); } catch { /* Preserve the known usage on pricing failure. */ }
+		out.push({ type: "error", reason: "error", error: message });
+		out.end(message);
+	})();
+	return out;
+}
+
+/** One dispatch layer, with session-local selection and no released lease history. */
+function observerLeases(provider: string, owned: () => boolean, restore: () => void) {
+	const leases = new Map<symbol, { observe: ObserveRequest; current: () => boolean }>();
+	let closed = false;
+	const live = (token: symbol) => !closed && leases.has(token);
+	const active = (token: symbol) => live(token) && leases.get(token)?.current() === true && owned();
+	const release = () => {
+		if (closed) return;
+		closed = true;
+		leases.clear();
+		restore();
+	};
+	const observe: ObserveRequest = (...args) => {
+		const handlers: PayloadObserver[] = [];
+		for (const [token, lease] of leases) {
+			try {
+				if (!active(token)) continue;
+				const handler = lease.observe(...args);
+				if (!handler) continue;
+				handlers.push(Object.assign((payload: unknown) => active(token) ? handler(payload) : payload, {
+					cacheRetention: handler.cacheRetention,
+					// A catalog/API recomposition may detach a provider after dispatch.
+					// Its in-flight bill still uses the sampled policy; explicit lease
+					// removal/disposal, unlike recomposition, revokes that callback.
+					priceUsage: handler.priceUsage ? (message: AssistantMessage) => { if (live(token)) handler.priceUsage!(message); } : undefined,
+					complete: handler.complete ? (message: AssistantMessage) => { if (active(token)) handler.complete!(message); } : undefined,
+				}));
+			} catch { /* A failed optional observer must not disable the other sessions. */ }
+		}
+		if (!handlers.length) return;
+		return Object.assign(async (payload: unknown) => {
+			for (const handler of handlers) {
+				try {
+					const replacement = await handler(payload);
+					if (replacement !== undefined) payload = replacement;
+				} catch { /* Keep the last valid request on optional observation failure. */ }
+			}
+			return payload;
+		}, {
+			cacheRetention: [...handlers].reverse().find((handler) => handler.cacheRetention !== undefined)?.cacheRetention,
+			priceUsage: handlers.some((handler) => handler.priceUsage) ? (message: AssistantMessage) => {
+				for (const handler of handlers) {
+					try { handler.priceUsage?.(message); } catch { /* Isolate independent pricing owners. */ }
+				}
+			} : undefined,
+			complete: handlers.some((handler) => handler.complete) ? (message: AssistantMessage) => {
+				for (const handler of handlers) {
+					try { handler.complete?.(message); } catch { /* Isolate optional completion observers too. */ }
+				}
+			} : undefined,
+		});
+	};
+	return {
+		observe, release, isOpen: () => !closed,
+		add(observe: ObserveRequest, current: () => boolean = () => true): InstalledObserver {
+			const token = Symbol();
+			leases.set(token, { observe, current });
+			return { provider, isInstalled: () => active(token), remove() {
+				leases.delete(token);
+				if (!leases.size) release();
+			} };
+		},
+	};
 }
 
+type ObserverLeases = ReturnType<typeof observerLeases>;
+// Pi's uncached module loader can instantiate this extension separately for
+// sessions sharing one model runtime. Weak ownership must span those instances.
+const LEASE_TABLES = Symbol.for("@nijaru/pi-provider-compaction/observer-leases-v2");
+interface LeaseTables {
+	effectiveLeases: WeakMap<Provider, { owned(): boolean; leases: ObserverLeases }>;
+	effectiveRestores: WeakMap<Function, { base: Provider; restore(): void }>;
+	nativeLeases: WeakMap<Provider, ObserverLeases>;
+	legacyLeases: WeakMap<Provider["streamSimple"], ObserverLeases>;
+}
+const shared = globalThis as typeof globalThis & { [LEASE_TABLES]?: LeaseTables };
+const { effectiveLeases, effectiveRestores, nativeLeases, legacyLeases } = shared[LEASE_TABLES] ??= {
+	effectiveLeases: new WeakMap(), effectiveRestores: new WeakMap(),
+	nativeLeases: new WeakMap(), legacyLeases: new WeakMap(),
+};
+
 /** Preserve method receivers, getters, optional members and dynamic catalogs. */
 export function decorateProvider(base: Provider, observe: ObserveRequest): Provider {
-	const stream: Provider["stream"] = (model, context, options) => base.stream(model, context, optionsWithObserver(model, context, options, observe));
-	const streamSimple: Provider["streamSimple"] = (model, context, options) => base.streamSimple(model, context, optionsWithObserver(model, context, options, observe));
+	const stream: Provider["stream"] = (model, context, options) => observedStream(model, context, options, observe, (next) => base.stream(model, context, next));
+	const streamSimple: Provider["streamSimple"] = (model, context, options) => observedStream(model, context, options, observe, (next) => base.streamSimple(model, context, next));
 	const bound = new Map<PropertyKey, unknown>();
 	return new Proxy(base, {
 		get(target, key) {
@@ -46,63 +185,106 @@ export function decorateProvider(base: Provider, observe: ObserveRequest): Provi
 	});
 }
 
+/** Observe a builtin/composer in place, preserving its auth, catalogs and dispatch closure. */
+export function leaseEffectiveProvider(base: Provider | undefined, current: () => Provider | undefined, observe: ObserveRequest): InstalledObserver | undefined {
+	if (!base) return;
+	const existing = effectiveLeases.get(base);
+	if (existing?.leases.isOpen() && existing.owned()) return existing.leases.add(observe, () => current() === base);
+	// External replacement invalidated the old pair. Restore its still-owned
+	// method before capturing delegates, rather than burying a released wrapper.
+	existing?.leases.release();
+	// A third party may have restored saved method values after their release.
+	// They remain inert, but must not become delegates of another observation layer.
+	const keys = ["stream", "streamSimple"] as const;
+	for (const key of keys) {
+		const saved = effectiveRestores.get(base[key]);
+		if (saved?.base === base) saved.restore();
+	}
+	const descriptors = keys.map((key) => Object.getOwnPropertyDescriptor(base, key));
+	if (descriptors.some((descriptor) => descriptor && (!("value" in descriptor) || !descriptor.configurable && !descriptor.writable)) || !Object.isExtensible(base) && descriptors.some((descriptor) => !descriptor)) return;
+	const originalStream = base.stream, originalSimple = base.streamSimple;
+	const owned = () => base.stream === stream && base.streamSimple === streamSimple;
+	const restore = () => {
+		for (const [index, key] of keys.entries()) if (base[key] === (key === "stream" ? stream : streamSimple)) {
+			const descriptor = descriptors[index];
+			if (descriptor) Object.defineProperty(base, key, descriptor);
+			else Reflect.deleteProperty(base, key);
+		}
+	};
+	const leases = observerLeases(base.id, owned, restore);
+	const stream: Provider["stream"] = (model, context, options) => observedStream(model, context, options, leases.observe, (next) => originalStream.call(base, model, context, next));
+	const streamSimple: Provider["streamSimple"] = (model, context, options) => observedStream(model, context, options, leases.observe, (next) => originalSimple.call(base, model, context, next));
+	try {
+		for (const [index, key] of keys.entries()) Object.defineProperty(base, key, { ...descriptors[index], value: key === "stream" ? stream : streamSimple, ...(descriptors[index] ? {} : { configurable: true, writable: true, enumerable: true }) });
+	} catch { leases.release(); return; }
+	for (const method of [stream, streamSimple]) effectiveRestores.set(method, { base, restore });
+	effectiveLeases.set(base, { owned, leases });
+	return leases.add(observe, () => current() === base);
+}
+
 /**
 	* Pi 0.87 has replacement registration, not middleware. Native BASE providers can
 	* be decorated directly. Legacy stream overlays retain their configuration and
 	* delegate their original callback (including /fast). Never register an effective
 	* composed provider as a native base: that loses model headers and freezes config.
 	*/
-export function installObserver(pi: ExtensionAPI, ctx: ExtensionContext, observe: ObserveRequest): InstalledObserver | undefined {
+export function installObserver(pi: ExtensionAPI, ctx: Pick<ExtensionContext, "model" | "modelRegistry">, observe: ObserveRequest): InstalledObserver | undefined {
 	const model = ctx.model;
 	if (!model) return;
 	const registry = ctx.modelRegistry;
 	const id = model.provider;
+	const observeApi: ObserveRequest = (requestModel, ...args) => requestModel.api === model.api ? observe(requestModel, ...args) : undefined;
+	const tagged = (lease: InstalledObserver | undefined) => lease ? Object.assign(lease, { api: model.api }) : undefined;
+	const effective = () => {
+		const base = registry.getProvider(id);
+		if (!base?.getModels().some((entry) => entry.api === model.api)) return;
+		return tagged(leaseEffectiveProvider(base, () => registry.getProvider(id), observeApi));
+	};
 	const native = registry.getRegisteredNativeProvider(id);
 	if (native) {
 		// Pi composes a native base only when its catalog serves the active API.
 		// Otherwise the request goes to the global API provider and bypasses this
 		// wrapper; claiming an observer there would report native compaction as
 		// installed while never capturing anything.
-		if (typeof native.getModels !== "function" || !native.getModels().some((entry) => entry.api === model.api)) return;
-		const isInstalled = () => registry.getRegisteredNativeProvider(id) === wrapper;
-		const guarded: ObserveRequest = (...args) => {
-			if (!isInstalled()) return;
-			const handler = observe(...args);
-			return handler ? (payload) => isInstalled() ? handler(payload) : payload : undefined;
-		};
-		const wrapper = decorateProvider(native, guarded);
+		if (typeof native.getModels !== "function" || !native.getModels().some((entry) => entry.api === model.api)) return effective();
+		const existing = nativeLeases.get(native);
+		if (existing) return tagged(existing.add(observeApi));
+		const leases = observerLeases(id, () => registry.getRegisteredNativeProvider(id) === wrapper, () => {
+			// Keep cleanup usable after bare SDK disposal invalidates the extension API.
+			if (registry.getRegisteredNativeProvider(id) === wrapper) registry.registerProvider(native);
+			nativeLeases.delete(wrapper);
+		});
+		const wrapper = decorateProvider(native, leases.observe);
+		nativeLeases.set(wrapper, leases);
 		pi.registerProvider(wrapper);
-		return {
-			provider: id, isInstalled,
-			remove() {
-				// Keep cleanup usable after bare SDK disposal invalidates the extension API.
-				if (isInstalled()) registry.registerProvider(native);
-			},
-		};
+		return tagged(leases.add(observeApi));
 	}
 	const original = registry.getRegisteredProviderConfig(id);
 	// Refresh callbacks own mutable catalog state inside the composer. Replacing
 	// that composer cannot promise preservation, so leave it alone.
-	if (!original?.streamSimple || original.refreshModels || original.oauth?.modifyModels) return;
-	if (original.api !== model.api) return;
+	if (!original?.streamSimple || original.refreshModels || original.oauth?.modifyModels || original.api !== model.api) {
+		// Lease the effective object's methods, NOT its registry slot. Re-registering
+		// a composed provider freezes catalogs and loses overlay/header semantics.
+		return effective();
+	}
 	// An injected legacy streamSimple overlay would also intercept full stream()
 	// calls and change reasoning semantics. Only decorate an existing stream owner.
+	const existing = legacyLeases.get(original.streamSimple);
+	if (existing) return tagged(existing.add(observeApi));
 	const delegate = original.streamSimple.bind(original);
-	const isInstalled = () => registry.getRegisteredProviderConfig(id)?.streamSimple === wrapper && !registry.getRegisteredNativeProvider(id);
-	const guarded: ObserveRequest = (...args) => {
-		if (!isInstalled()) return;
-		const handler = observe(...args);
-		return handler ? (payload) => isInstalled() ? handler(payload) : payload : undefined;
-	};
-	const wrapper: Provider["streamSimple"] = (requestModel, context, options) => delegate(requestModel, context, optionsWithObserver(requestModel, context, options, guarded));
-	const remove = () => {
-		if (!isInstalled()) return;
+	const owned = () => registry.getRegisteredProviderConfig(id)?.streamSimple === wrapper && !registry.getRegisteredNativeProvider(id);
+	const leases = observerLeases(id, owned, () => {
+		legacyLeases.delete(wrapper);
+		if (!owned()) return;
 		// registerProvider merges undefined properties: remove our layer before restore.
 		const restored = { ...registry.getRegisteredProviderConfig(id)!, streamSimple: original.streamSimple };
 		registry.unregisterProvider(id);
 		registry.registerProvider(id, restored);
-	};
-	Object.defineProperty(wrapper, STREAM_OBSERVER, { value: { delegate: original.streamSimple, release: remove } });
+	});
+	const wrapper: Provider["streamSimple"] = (requestModel, context, options) => observedStream(requestModel, context, options, leases.observe, (next) => delegate(requestModel, context, next));
+	// /fast owns the delegate: its cleanup must release all observing sessions.
+	Object.defineProperty(wrapper, STREAM_OBSERVER, { value: { delegate: original.streamSimple, release: leases.release } });
+	legacyLeases.set(wrapper, leases);
 	pi.registerProvider(id, { ...original, streamSimple: wrapper });
-	return { provider: id, isInstalled, remove };
+	return tagged(leases.add(observeApi));
 }

package/usage-observer.ts

@@ -0,0 +1,52 @@
+import type { Usage } from "@earendil-works/pi-ai";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { installObserver, type InstalledObserver } from "./provider";
+
+const REQUEST = "pi-provider-compaction:observe-usage";
+const READY = "pi-provider-compaction:usage-observer-ready";
+type ModelKey = { provider: string; id: string };
+type PriceUsage = (timestamp: number, usage: Usage) => void;
+interface UsageRequest {
+	context: ExtensionContext;
+	model: NonNullable<ExtensionContext["model"]>;
+	priceFor(model: ModelKey): PriceUsage;
+	respond(observer: InstalledObserver | undefined): void;
+	handled?: boolean;
+}
+
+/** Share the existing stream owner, rather than stacking competing pricing wrappers. */
+export function serveUsageObservers(pi: ExtensionAPI, current: () => ExtensionContext | undefined) {
+	const leases = new Set<InstalledObserver>();
+	const remove = pi.events.on(REQUEST, (value) => {
+		const request = value as Partial<UsageRequest> | undefined;
+		const owner = current();
+		if (!owner || !request || request.handled || !request.context || !request.model || typeof request.priceFor !== "function" || typeof request.respond !== "function") return;
+		try {
+			if (request.context.modelRegistry !== owner.modelRegistry || request.context.sessionManager.getSessionId() !== owner.sessionManager.getSessionId()) return;
+			request.handled = true;
+			const sessionId = owner.sessionManager.getSessionId();
+			const installed = installObserver(pi, { model: request.model, modelRegistry: owner.modelRegistry }, (model, _context, _options, scope) => {
+				// The SDK header event identifies the owner even when a branch
+				// summary has an independent routing id. Sample policy at dispatch.
+				const price = request.priceFor!({ provider: model.provider, id: model.id });
+				return Object.assign((payload: unknown) => payload, { priceUsage: (message: { timestamp: number; usage: Usage }) => {
+					if (scope.sessionId === sessionId) price(message.timestamp, message.usage);
+				} });
+			});
+			if (!installed) { request.respond(undefined); return; }
+			const lease: InstalledObserver = {
+				provider: installed.provider,
+				api: installed.api,
+				isInstalled: () => installed.isInstalled(),
+				remove() { leases.delete(lease); installed.remove(); },
+			};
+			leases.add(lease);
+			request.respond(lease);
+		} catch { /* Optional cost estimates must not veto a provider request. */ }
+	});
+	return {
+		ready(context: ExtensionContext) { pi.events.emit(READY, { context }); },
+		clear() { for (const lease of leases) lease.remove(); },
+		dispose() { remove(); for (const lease of leases) lease.remove(); },
+	};
+}