@nijaru/pi-provider-compaction 0.0.8 → 0.0.9

package/provider.ts

@@ -1,7 +1,11 @@
-import type { Provider, StreamOptions, Model, TranscriptContext } from "@earendil-works/pi-ai";
+import type { AssistantMessage, AssistantMessageEventStream, CacheRetention, Provider, StreamOptions, Model, TranscriptContext } from "@earendil-works/pi-ai";
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
 
-export type ObserveRequest = (model: Model<any>, context: TranscriptContext, options: StreamOptions) => ((payload: unknown) => unknown | Promise<unknown>) | undefined;
+export type PayloadObserver = ((payload: unknown) => unknown | Promise<unknown>) & {
+	complete?: (message: AssistantMessage) => void;
+	cacheRetention?: CacheRetention;
+};
+export type ObserveRequest = (model: Model<any>, context: TranscriptContext, options: StreamOptions) => PayloadObserver | undefined;
 
 export interface InstalledObserver {
 	provider: string;
@@ -13,26 +17,101 @@ 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 {
 	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 {
+	catch { return dispatch(options); } // An expired optional observer must not veto dispatch.
+	if (!handler) return dispatch(options);
+	const stream = 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);
+	// result() observes completion without consuming the public event iterator.
+	if (handler.complete) void stream.result().then(handler.complete).catch(() => {});
+	return stream;
+}
+
+/** 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 active = (token: symbol) => !closed && 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,
+					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,
+			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-v1");
+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,6 +125,43 @@ 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
@@ -64,45 +180,44 @@ export function installObserver(pi: ExtensionAPI, ctx: ExtensionContext, observe
 		// 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);
+		const existing = nativeLeases.get(native);
+		if (existing) return existing.add(observe);
+		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 leases.add(observe);
 	}
 	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 leaseEffectiveProvider(registry.getProvider(id), () => registry.getProvider(id), observe);
+	}
 	// 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 existing.add(observe);
 	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 leases.add(observe);
 }