Spaces:
Runtime error
Runtime error
File size: 5,148 Bytes
cd8bd0a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | /**
* Antigravity project bootstrap — loadCodeAssist.
*
* The Google Cloud Code Assist API (/v1internal:models) requires a prior
* /v1internal:loadCodeAssist call to assign a project context to the
* OAuth token. Without this bootstrap, :models returns 404.
*
* This module provides an idempotent ensureAntigravityProjectAssigned()
* helper that is called once per access-token before every discovery
* attempt. Results are memoized per-token for the process lifetime to
* avoid redundant round-trips.
*
* Based on AntigravityService.loadCodeAssist() in
* src/lib/oauth/services/antigravity.ts and the CLIProxyAPI reference
* implementation in internal/runtime/executor/antigravity_executor.go.
*/
import {
getAntigravityHeaders,
getAntigravityLoadCodeAssistMetadata,
} from "./antigravityHeaders.ts";
import {
getAntigravityBootstrapHeaders,
type AntigravityClientProfile,
} from "./antigravityClientProfile.ts";
import { ANTIGRAVITY_BASE_URLS } from "../config/antigravityUpstream.ts";
const LOAD_CODE_ASSIST_PATH = "/v1internal:loadCodeAssist";
const BOOTSTRAP_TIMEOUT_MS = 8_000;
/** Ordered list of loadCodeAssist endpoint URLs (mirrors the models discovery order). */
export function getAntigravityLoadCodeAssistUrls(): string[] {
return ANTIGRAVITY_BASE_URLS.map((base) => `${base}${LOAD_CODE_ASSIST_PATH}`);
}
/** Per-token memoization cache (lives for the process lifetime). */
const projectCache = new Map<string, string>();
type FetchLike = (url: string, init?: RequestInit) => Promise<Response>;
function getProjectCacheKey(accessToken: string, clientProfile: AntigravityClientProfile): string {
return `${clientProfile}:${accessToken}`;
}
/**
* Attempt loadCodeAssist against each known base URL in order.
* Returns the discovered project id, or null if all endpoints fail.
*/
async function tryLoadCodeAssist(
accessToken: string,
fetchImpl: FetchLike,
clientProfile: AntigravityClientProfile
): Promise<string | null> {
const urls = getAntigravityLoadCodeAssistUrls();
const headers =
clientProfile === "harness"
? getAntigravityBootstrapHeaders(clientProfile, accessToken)
: getAntigravityHeaders("loadCodeAssist", accessToken);
for (const url of urls) {
try {
const response = await fetchImpl(url, {
method: "POST",
headers,
body: JSON.stringify({ metadata: getAntigravityLoadCodeAssistMetadata() }),
signal: AbortSignal.timeout(BOOTSTRAP_TIMEOUT_MS),
});
if (!response.ok) {
console.warn(
`[models] antigravity loadCodeAssist failed at ${url} (${response.status}) — trying next`
);
continue;
}
const data = (await response.json()) as Record<string, unknown>;
// cloudaicompanionProject may be a plain string or an object with an id field.
const raw = data.cloudaicompanionProject;
let projectId =
typeof raw === "string"
? raw.trim()
: raw &&
typeof raw === "object" &&
typeof (raw as Record<string, unknown>).id === "string"
? ((raw as Record<string, unknown>).id as string).trim()
: "";
if (projectId) {
return projectId;
}
console.warn(
`[models] antigravity loadCodeAssist at ${url} returned no project id — trying next`
);
} catch (error) {
const msg = error instanceof Error ? error.message : String(error);
console.warn(`[models] antigravity loadCodeAssist threw for ${url}: ${msg} — trying next`);
}
}
return null;
}
/**
* Ensure a project is assigned to the given access token by calling
* loadCodeAssist if not already cached. Idempotent — repeated calls
* for the same token return the cached result without a network round-trip.
*
* Failures are non-fatal: the caller should proceed with the :models
* request regardless (the stored project_id in the DB may still be valid).
*
* @param accessToken The OAuth bearer token for the current connection.
* @param fetchImpl Injected fetch implementation (defaults to globalThis.fetch).
*/
export async function ensureAntigravityProjectAssigned(
accessToken: string,
fetchImpl: FetchLike = fetch,
clientProfile: AntigravityClientProfile = "ide"
): Promise<string | undefined> {
const cacheKey = getProjectCacheKey(accessToken, clientProfile);
if (projectCache.has(cacheKey)) {
return projectCache.get(cacheKey); // already bootstrapped for this token
}
const projectId = await tryLoadCodeAssist(accessToken, fetchImpl, clientProfile);
if (projectId) {
projectCache.set(cacheKey, projectId);
return projectId;
}
// Non-fatal: if all endpoints failed, we proceed without caching.
return undefined;
}
/** Exported for tests. */
export function clearAntigravityProjectCache(): void {
projectCache.clear();
}
/** Exported for tests — inspect cache state. */
export function getAntigravityProjectFromCache(
accessToken: string,
clientProfile: AntigravityClientProfile = "ide"
): string | undefined {
return projectCache.get(getProjectCacheKey(accessToken, clientProfile));
}
|