mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-21 05:02:14 +00:00
Fix OAuth using Google Cloud Code Assist.
This commit is contained in:
parent
498958a92a
commit
c7bac7583c
1 changed files with 155 additions and 44 deletions
|
|
@ -122,13 +122,28 @@ interface LoadCodeAssistPayload {
|
||||||
allowedTiers?: Array<{ id?: string; isDefault?: boolean }>;
|
allowedTiers?: Array<{ id?: string; isDefault?: boolean }>;
|
||||||
}
|
}
|
||||||
|
|
||||||
interface OnboardUserPayload {
|
/**
|
||||||
|
* Long-running operation response from onboardUser
|
||||||
|
*/
|
||||||
|
interface LongRunningOperationResponse {
|
||||||
|
name?: string;
|
||||||
done?: boolean;
|
done?: boolean;
|
||||||
response?: {
|
response?: {
|
||||||
cloudaicompanionProject?: { id?: string };
|
cloudaicompanionProject?: { id?: string };
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Tier IDs as used by the Cloud Code API
|
||||||
|
const TIER_FREE = "free-tier";
|
||||||
|
const TIER_LEGACY = "legacy-tier";
|
||||||
|
const TIER_STANDARD = "standard-tier";
|
||||||
|
|
||||||
|
interface GoogleRpcErrorResponse {
|
||||||
|
error?: {
|
||||||
|
details?: Array<{ reason?: string }>;
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Wait helper for onboarding retries
|
* Wait helper for onboarding retries
|
||||||
*/
|
*/
|
||||||
|
|
@ -137,18 +152,62 @@ function wait(ms: number): Promise<void> {
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Get default tier ID from allowed tiers
|
* Get default tier from allowed tiers
|
||||||
*/
|
*/
|
||||||
function getDefaultTierId(allowedTiers?: Array<{ id?: string; isDefault?: boolean }>): string | undefined {
|
function getDefaultTier(allowedTiers?: Array<{ id?: string; isDefault?: boolean }>): { id?: string } {
|
||||||
if (!allowedTiers || allowedTiers.length === 0) return undefined;
|
if (!allowedTiers || allowedTiers.length === 0) return { id: TIER_LEGACY };
|
||||||
const defaultTier = allowedTiers.find((t) => t.isDefault);
|
const defaultTier = allowedTiers.find((t) => t.isDefault);
|
||||||
return defaultTier?.id ?? allowedTiers[0]?.id;
|
return defaultTier ?? { id: TIER_LEGACY };
|
||||||
|
}
|
||||||
|
|
||||||
|
function isVpcScAffectedUser(payload: unknown): boolean {
|
||||||
|
if (!payload || typeof payload !== "object") return false;
|
||||||
|
if (!("error" in payload)) return false;
|
||||||
|
const error = (payload as GoogleRpcErrorResponse).error;
|
||||||
|
if (!error?.details || !Array.isArray(error.details)) return false;
|
||||||
|
return error.details.some((detail) => detail.reason === "SECURITY_POLICY_VIOLATED");
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Poll a long-running operation until completion
|
||||||
|
*/
|
||||||
|
async function pollOperation(
|
||||||
|
operationName: string,
|
||||||
|
headers: Record<string, string>,
|
||||||
|
onProgress?: (message: string) => void,
|
||||||
|
): Promise<LongRunningOperationResponse> {
|
||||||
|
let attempt = 0;
|
||||||
|
while (true) {
|
||||||
|
if (attempt > 0) {
|
||||||
|
onProgress?.(`Waiting for project provisioning (attempt ${attempt + 1})...`);
|
||||||
|
await wait(5000);
|
||||||
|
}
|
||||||
|
|
||||||
|
const response = await fetch(`${CODE_ASSIST_ENDPOINT}/v1internal/${operationName}`, {
|
||||||
|
method: "GET",
|
||||||
|
headers,
|
||||||
|
});
|
||||||
|
|
||||||
|
if (!response.ok) {
|
||||||
|
throw new Error(`Failed to poll operation: ${response.status} ${response.statusText}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
const data = (await response.json()) as LongRunningOperationResponse;
|
||||||
|
if (data.done) {
|
||||||
|
return data;
|
||||||
|
}
|
||||||
|
|
||||||
|
attempt += 1;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Discover or provision a Google Cloud project for the user
|
* Discover or provision a Google Cloud project for the user
|
||||||
*/
|
*/
|
||||||
async function discoverProject(accessToken: string, onProgress?: (message: string) => void): Promise<string> {
|
async function discoverProject(accessToken: string, onProgress?: (message: string) => void): Promise<string> {
|
||||||
|
// Check for user-provided project ID via environment variable
|
||||||
|
const envProjectId = process.env["GOOGLE_CLOUD_PROJECT"] || process.env["GOOGLE_CLOUD_PROJECT_ID"];
|
||||||
|
|
||||||
const headers = {
|
const headers = {
|
||||||
Authorization: `Bearer ${accessToken}`,
|
Authorization: `Bearer ${accessToken}`,
|
||||||
"Content-Type": "application/json",
|
"Content-Type": "application/json",
|
||||||
|
|
@ -162,62 +221,114 @@ async function discoverProject(accessToken: string, onProgress?: (message: strin
|
||||||
method: "POST",
|
method: "POST",
|
||||||
headers,
|
headers,
|
||||||
body: JSON.stringify({
|
body: JSON.stringify({
|
||||||
|
cloudaicompanionProject: envProjectId,
|
||||||
metadata: {
|
metadata: {
|
||||||
ideType: "IDE_UNSPECIFIED",
|
ideType: "IDE_UNSPECIFIED",
|
||||||
platform: "PLATFORM_UNSPECIFIED",
|
platform: "PLATFORM_UNSPECIFIED",
|
||||||
pluginType: "GEMINI",
|
pluginType: "GEMINI",
|
||||||
|
duetProject: envProjectId,
|
||||||
},
|
},
|
||||||
}),
|
}),
|
||||||
});
|
});
|
||||||
|
|
||||||
if (loadResponse.ok) {
|
let data: LoadCodeAssistPayload;
|
||||||
const data = (await loadResponse.json()) as LoadCodeAssistPayload;
|
|
||||||
|
|
||||||
// If we have an existing project, use it
|
if (!loadResponse.ok) {
|
||||||
|
let errorPayload: unknown;
|
||||||
|
try {
|
||||||
|
errorPayload = await loadResponse.clone().json();
|
||||||
|
} catch {
|
||||||
|
errorPayload = undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (isVpcScAffectedUser(errorPayload)) {
|
||||||
|
data = { currentTier: { id: TIER_STANDARD } };
|
||||||
|
} else {
|
||||||
|
const errorText = await loadResponse.text();
|
||||||
|
throw new Error(`loadCodeAssist failed: ${loadResponse.status} ${loadResponse.statusText}: ${errorText}`);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
data = (await loadResponse.json()) as LoadCodeAssistPayload;
|
||||||
|
}
|
||||||
|
|
||||||
|
// If user already has a current tier and project, use it
|
||||||
|
if (data.currentTier) {
|
||||||
if (data.cloudaicompanionProject) {
|
if (data.cloudaicompanionProject) {
|
||||||
return data.cloudaicompanionProject;
|
return data.cloudaicompanionProject;
|
||||||
}
|
}
|
||||||
|
// User has a tier but no managed project - they need to provide one via env var
|
||||||
// Otherwise, try to onboard with the FREE tier
|
if (envProjectId) {
|
||||||
const tierId = getDefaultTierId(data.allowedTiers) ?? "FREE";
|
return envProjectId;
|
||||||
|
|
||||||
onProgress?.("Provisioning Cloud Code Assist project (this may take a moment)...");
|
|
||||||
|
|
||||||
// Onboard with retries (the API may take time to provision)
|
|
||||||
for (let attempt = 0; attempt < 10; attempt++) {
|
|
||||||
const onboardResponse = await fetch(`${CODE_ASSIST_ENDPOINT}/v1internal:onboardUser`, {
|
|
||||||
method: "POST",
|
|
||||||
headers,
|
|
||||||
body: JSON.stringify({
|
|
||||||
tierId,
|
|
||||||
metadata: {
|
|
||||||
ideType: "IDE_UNSPECIFIED",
|
|
||||||
platform: "PLATFORM_UNSPECIFIED",
|
|
||||||
pluginType: "GEMINI",
|
|
||||||
},
|
|
||||||
}),
|
|
||||||
});
|
|
||||||
|
|
||||||
if (onboardResponse.ok) {
|
|
||||||
const onboardData = (await onboardResponse.json()) as OnboardUserPayload;
|
|
||||||
const projectId = onboardData.response?.cloudaicompanionProject?.id;
|
|
||||||
|
|
||||||
if (onboardData.done && projectId) {
|
|
||||||
return projectId;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Wait before retrying
|
|
||||||
if (attempt < 9) {
|
|
||||||
onProgress?.(`Waiting for project provisioning (attempt ${attempt + 2}/10)...`);
|
|
||||||
await wait(3000);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
throw new Error(
|
||||||
|
"This account requires setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID environment variable. " +
|
||||||
|
"See https://goo.gle/gemini-cli-auth-docs#workspace-gca",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
// User needs to be onboarded - get the default tier
|
||||||
|
const tier = getDefaultTier(data.allowedTiers);
|
||||||
|
const tierId = tier?.id ?? TIER_FREE;
|
||||||
|
|
||||||
|
if (tierId !== TIER_FREE && !envProjectId) {
|
||||||
|
throw new Error(
|
||||||
|
"This account requires setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID environment variable. " +
|
||||||
|
"See https://goo.gle/gemini-cli-auth-docs#workspace-gca",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
onProgress?.("Provisioning Cloud Code Assist project (this may take a moment)...");
|
||||||
|
|
||||||
|
// Build onboard request - for free tier, don't include project ID (Google provisions one)
|
||||||
|
// For other tiers, include the user's project ID if available
|
||||||
|
const onboardBody: Record<string, unknown> = {
|
||||||
|
tierId,
|
||||||
|
metadata: {
|
||||||
|
ideType: "IDE_UNSPECIFIED",
|
||||||
|
platform: "PLATFORM_UNSPECIFIED",
|
||||||
|
pluginType: "GEMINI",
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
if (tierId !== TIER_FREE && envProjectId) {
|
||||||
|
onboardBody["cloudaicompanionProject"] = envProjectId;
|
||||||
|
(onboardBody["metadata"] as Record<string, unknown>)["duetProject"] = envProjectId;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start onboarding - this returns a long-running operation
|
||||||
|
const onboardResponse = await fetch(`${CODE_ASSIST_ENDPOINT}/v1internal:onboardUser`, {
|
||||||
|
method: "POST",
|
||||||
|
headers,
|
||||||
|
body: JSON.stringify(onboardBody),
|
||||||
|
});
|
||||||
|
|
||||||
|
if (!onboardResponse.ok) {
|
||||||
|
const errorText = await onboardResponse.text();
|
||||||
|
throw new Error(`onboardUser failed: ${onboardResponse.status} ${onboardResponse.statusText}: ${errorText}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
let lroData = (await onboardResponse.json()) as LongRunningOperationResponse;
|
||||||
|
|
||||||
|
// If the operation isn't done yet, poll until completion
|
||||||
|
if (!lroData.done && lroData.name) {
|
||||||
|
lroData = await pollOperation(lroData.name, headers, onProgress);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to get project ID from the response
|
||||||
|
const projectId = lroData.response?.cloudaicompanionProject?.id;
|
||||||
|
if (projectId) {
|
||||||
|
return projectId;
|
||||||
|
}
|
||||||
|
|
||||||
|
// If no project ID from onboarding, fall back to env var
|
||||||
|
if (envProjectId) {
|
||||||
|
return envProjectId;
|
||||||
}
|
}
|
||||||
|
|
||||||
throw new Error(
|
throw new Error(
|
||||||
"Could not discover or provision a Google Cloud project. " +
|
"Could not discover or provision a Google Cloud project. " +
|
||||||
"Please ensure you have access to Google Cloud Code Assist (Gemini CLI).",
|
"Try setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID environment variable. " +
|
||||||
|
"See https://goo.gle/gemini-cli-auth-docs#workspace-gca",
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue