pygienium/tests/provider-forwarding.test.ts

/**
 * provider-forwarding.test.ts — the provider hand-off from the invoking session
 * into a sub-agent session.
 *
 * A sub-agent session is created with `noExtensions: true`, so a provider that
 * exists only because an extension called `pi.registerProvider` is absent from
 * its `ModelRuntime` and its credentials cannot be resolved — pinning the
 * invoking session's model then fails with `No API key found for <provider>`
 * before a single token, on a perfectly healthy login. These tests pin the two
 * halves of the fix: narrowing the host registry at the boundary, and the inline
 * extension that re-registers the pinned model's provider inside the sub-session
 * (inline factories load even under `noExtensions`).
 */
import { describe, expect, it } from "bun:test";
import {
	asProviderRegistry,
	providerForwardingExtension,
	type ProviderRegistry,
} from "../src/agent-runner.js";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

/** Record `registerProvider` calls so the relay can be asserted on. */
function recorder(): {
	pi: ExtensionAPI;
	calls: unknown[][];
} {
	const calls: unknown[][] = [];
	const pi = {
		registerProvider: (...args: unknown[]) => {
			calls.push(args);
		},
	} as unknown as ExtensionAPI;
	return { pi, calls };
}

/** Registry stub holding one provider id, native and/or config form. */
function registry(entry: {
	native?: unknown;
	config?: unknown;
}): ProviderRegistry {
	return {
		getRegisteredNativeProvider: (_id) =>
			entry.native === undefined ? undefined : entry.native,
		getRegisteredProviderConfig: (_id) =>
			entry.config === undefined ? undefined : entry.config,
	};
}

describe("asProviderRegistry", () => {
	it("accepts a registry that exposes both registration lookups", () => {
		const host = registry({ native: { id: "neuralwatt" } });
		expect(asProviderRegistry(host)).toBe(host);
	});

	it("rejects values that cannot answer both lookups", () => {
		expect(asProviderRegistry(undefined)).toBeUndefined();
		expect(asProviderRegistry(null)).toBeUndefined();
		expect(asProviderRegistry("modelRegistry")).toBeUndefined();
		// A registry from a host that never records extension providers, such as
		// the omp port, degrades to "nothing to forward" instead of throwing.
		expect(asProviderRegistry({})).toBeUndefined();
		expect(
			asProviderRegistry({ getRegisteredNativeProvider: () => undefined }),
		).toBeUndefined();
	});
});

describe("providerForwardingExtension", () => {
	it("is a no-op without a registry, a pinned model, or a registered provider", () => {
		// Built-in and models.json providers need no forwarding: pi's own
		// catalogue already answers for them.
		const empty = registry({});
		expect(providerForwardingExtension(undefined, "neuralwatt")).toBeUndefined();
		expect(providerForwardingExtension(empty, undefined)).toBeUndefined();
		expect(providerForwardingExtension(empty, "neuralwatt")).toBeUndefined();
	});

	it("re-registers a native provider, keeping its own stream and auth wiring", () => {
		const native = { id: "neuralwatt" };
		const extension = providerForwardingExtension(
			registry({ native, config: { baseUrl: "ignored-when-native" } }),
			"neuralwatt",
		);
		if (!extension) throw new Error("expected a forwarding extension");
		expect(extension.name).toBe("pygienium-provider");
		// Hidden: the sub-agent's extension list must not advertise an extension
		// the user never installed for it.
		expect(extension.hidden).toBe(true);

		const { pi, calls } = recorder();
		extension.factory(pi);
		// The native provider wins over a config-only registration, which would
		// only augment a built-in provider and lose the wrapped stream.
		expect(calls).toEqual([[native]]);
	});

	it("re-registers a config-only provider under its provider id", () => {
		const config = { baseUrl: "config-only-provider" };
		const extension = providerForwardingExtension(registry({ config }), "probe");
		if (!extension) throw new Error("expected a forwarding extension");
		const { pi, calls } = recorder();
		extension.factory(pi);
		expect(calls).toEqual([["probe", config]]);
	});
});