/**
* Google Antigravity provider extension.
*
* Provides Google Cloud Code Assist (Antigravity) login, live model discovery,
* and streaming. The wire protocol, OAuth flow, header spoofing, and request
* envelope mirror oh-my-pi's google-antigravity implementation.
*
* Usage:
* prime-agent -e ./packages/coding-agent/examples/extensions/google-antigravity
* # Then /login google-antigravity
*
* Or drop the directory into ~/.prime/agent/extensions/ for auto-discovery.
*/
import { readFileSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import type { Api, AssistantMessageEventStream, Context, Model } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ProviderModelConfig } from "@earendil-works/pi-coding-agent";
import { fetchAntigravityModels } from "./discovery.js";
import type { AntigravitySimpleStreamOptions } from "./helpers.js";
import { STATIC_PROVIDER_MODELS } from "./models.js";
import { antigravityOAuth } from "./oauth.js";
import { streamSimpleAntigravity } from "./stream.js";
import { ANTIGRAVITY_DAILY_ENDPOINT } from "./wire.js";
export const PROVIDER_ID = "google-antigravity";
export const PROVIDER_API = "google-antigravity" as Api;
/**
* Mirror prime's getAgentDir() without a runtime import of
* @earendil-works/pi-coding-agent (its dist barrel pulls pi-ai/mcp through
* jiti and fails to load inside an extension).
*/
function getAgentDir(): string {
const envDir = process.env.PRIME_AGENT_CODING_AGENT_DIR ?? process.env.PI_CODING_AGENT_DIR;
if (envDir) {
if (envDir === "~") return homedir();
if (envDir.startsWith("~/")) return join(homedir(), envDir.slice(2));
return envDir;
}
return join(homedir(), ".prime", "agent");
}
interface StoredCredential {
type?: string;
access?: string;
}
/**
* auth.json stores a bare credential per provider, or an array once more than
* one account is logged in (multi-account credential pooling).
*/
function readExistingAccessToken(): string | undefined {
try {
const authPath = join(getAgentDir(), "auth.json");
const auth = JSON.parse(readFileSync(authPath, "utf-8")) as Record<string, StoredCredential | StoredCredential[]>;
const stored = auth[PROVIDER_ID];
const credentials = Array.isArray(stored) ? stored : [stored];
const entry = credentials.find(
(credential) =>
credential?.type === "oauth" && typeof credential.access === "string" && credential.access.length > 0,
);
if (entry) return entry.access;
} catch {
// No auth file or not logged in yet.
}
return undefined;
}
const DISCOVERY_CACHE_FILE = "models-cache.json";
const DISCOVERY_CACHE_TTL_MS = 6 * 60 * 60 * 1000;
function discoveryCachePath(): string {
return join(dirname(fileURLToPath(import.meta.url)), DISCOVERY_CACHE_FILE);
}
/**
* A cached empty list records a failed discovery: it keeps the next start off the
* network for one TTL while leaving the static models in place.
*/
function readDiscoveryCache(): { models?: ProviderModelConfig[]; fresh: boolean } | undefined {
try {
const cached = JSON.parse(readFileSync(discoveryCachePath(), "utf-8")) as {
fetchedAt?: unknown;
models?: unknown;
};
if (typeof cached.fetchedAt !== "string" || !Array.isArray(cached.models)) {
return undefined;
}
const fetchedAt = Date.parse(cached.fetchedAt);
return {
...(cached.models.length > 0 ? { models: cached.models as ProviderModelConfig[] } : {}),
fresh: Number.isFinite(fetchedAt) && Date.now() - fetchedAt < DISCOVERY_CACHE_TTL_MS,
};
} catch {
return undefined;
}
}
function writeDiscoveryCache(models: ProviderModelConfig[]): void {
try {
writeFileSync(
discoveryCachePath(),
`${JSON.stringify({ fetchedAt: new Date().toISOString(), models }, null, "\t")}\n`,
);
} catch (error) {
// The cache is an optimization; a read-only extension dir must not break loading.
console.error(`google-antigravity: could not write ${DISCOVERY_CACHE_FILE}: ${String(error)}`);
}
}
function registerAntigravityProvider(pi: ExtensionAPI, baseUrl: string, models: ProviderModelConfig[]): void {
pi.registerProvider(PROVIDER_ID, {
name: "Google Antigravity",
baseUrl,
api: PROVIDER_API,
models,
oauth: antigravityOAuth,
streamSimple: (model: Model<Api>, context: Context, options?: AntigravitySimpleStreamOptions): AssistantMessageEventStream =>
streamSimpleAntigravity(model, context, options),
});
}
/**
* Live discovery is a request to Cloud Code Assist, so it must never hold up
* extension loading: both the client and every session worker load this module,
* and awaiting it added ~0.8s to each of them. Models come from the disk cache
* when it is fresh, otherwise from the static list, and a stale cache refreshes
* in the background and re-registers (the registry upserts providers).
*/
export default async function (pi: ExtensionAPI) {
const baseUrl = ANTIGRAVITY_DAILY_ENDPOINT;
const cached = readDiscoveryCache();
registerAntigravityProvider(pi, baseUrl, cached?.models ?? STATIC_PROVIDER_MODELS);
const accessToken = readExistingAccessToken();
if (!accessToken || cached?.fresh === true) {
return;
}
void (async () => {
try {
const discovered = await fetchAntigravityModels(accessToken);
writeDiscoveryCache(discovered ?? []);
if (discovered && discovered.length > 0) {
registerAntigravityProvider(pi, baseUrl, discovered);
}
} catch {
// Keep the models already registered. A failed fetch stays uncached so the
// next start retries instead of hiding a transient outage for a whole TTL.
}
})();
}