* # Then /login google-antigravity
*
* Or drop the directory into ~/.prime/agent/extensions/ for auto-discovery.
*/
import { readFileSync } from "node:fs";
import { readFileSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
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 } from "@earendil-works/pi-coding-agent";
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";
}
return undefined;
}
export default async function (pi: ExtensionAPI) {
const baseUrl = ANTIGRAVITY_DAILY_ENDPOINT;
let models = STATIC_PROVIDER_MODELS;
const DISCOVERY_CACHE_FILE = "models-cache.json";
const DISCOVERY_CACHE_TTL_MS = 6 * 60 * 60 * 1000;
const accessToken = readExistingAccessToken();
if (accessToken) {
try {
const discovered = await fetchAntigravityModels(accessToken);
if (discovered && discovered.length > 0) {
models = discovered;
}
} catch {
// Keep the static fallback list.
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,
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.
}
})();
}