@cortexkit/opencode-antigravity-auth 3.0.4 → 3.0.6

package/dist/src/ga/server/index.js

@@ -1,6 +1,6 @@
 /**
  * OpenCode 2 (2.0.22) server plugin. It sends the host's Google model
- * requests for this plugin's models to Antigravity accounts through a
+ * requests carrying the common-auth placeholder to Antigravity accounts through a
  * loopback bridge, and serves the `antigravity-auth` v1 RPC with each
  * location's account state and commands. `createGaAntigravityPlugin` is the
  * plugin the host loads; `createGaPluginFromServices` builds one over a
@@ -9,12 +9,14 @@
 import { randomUUID } from 'node:crypto';
 import { createServer, } from 'node:http';
 import { buildAntigravityHarnessUserAgent, fetchWithAgyCliTransport, getPublicModelDefinitions, isRetiredAgyModel, loadCommonAuthCommands, } from '@cortexkit/antigravity-auth-core';
+import { placeholderSecret } from '../../common-auth-embedded/opencode2/index.js';
 import { createGaLocationServices } from "../../plugin/ga-location-services.js";
+import { isCanonicalGoogleOrigin } from "../../plugin/google-origin.js";
 import { createLocationRuntime, } from "../../plugin/shared/runtime.js";
 import { redactAccountForSidebar, } from "../../sidebar-state.js";
 import { assertLoopbackProxyGuard, } from "../proxy-guard.js";
 import { ANTIGRAVITY_MENU_COMMAND_NAME, ANTIGRAVITY_RPC_DEFINITION, ANTIGRAVITY_RPC_LIMITS, ANTIGRAVITY_RPC_VERSION, AntigravityApplyInputSchema, AntigravityApplyOutputSchema, AntigravityApplyResultSchema, AntigravityNotificationPayloadSchema, AntigravityStateInputSchema, AntigravityStateOutputSchema, } from "../rpc/protocol.js";
-import { registerGaVaultActivation } from "./auth.js";
+import { registerGaAuthMethods } from "./auth.js";
 // Raw request sender
 /**
  * How long a response body may stall with no bytes before the raw sender
@@ -99,23 +101,35 @@ export function gaPreHeaderFailureResponse(error) {
 export async function startGaLoopbackBridge(options) {
     const queued = new Map();
     const running = new Set();
+    const reads = new Set();
     const sockets = new Set();
-    // Post-header failures by job id, held until the job's adapted response
-    // reads them or the hold expires.
-    const failures = new Map();
-    const recordFailure = (id, error) => {
-        const timer = setTimeout(() => failures.delete(id), GA_UNCLAIMED_JOB_TTL_MS);
-        timer.unref?.();
-        failures.set(id, { error, timer });
+    // Responses not yet wrapped by adaptResponse stay in this map with a bounded
+    // hold after completion. Once wrapped, the body's closure keeps the upstream
+    // Error available to slow readers without an expiring global error record.
+    const unclaimed = new Map();
+    // Weak references let disposal detach its bookkeeping without keeping
+    // abandoned response bodies (and their errors) alive.
+    const lifetimes = new Set();
+    const finalized = new FinalizationRegistry((reference) => lifetimes.delete(reference));
+    const release = (state, discardFailure = true) => {
+        state.released = true;
+        if (discardFailure)
+            state.failure = undefined;
+        const entry = unclaimed.get(state.id);
+        if (entry?.state === state) {
+            clearTimeout(entry.timer);
+            unclaimed.delete(state.id);
+        }
+        if (state.reference)
+            lifetimes.delete(state.reference);
+        finalized.unregister(state);
     };
-    const takeFailure = (id) => {
-        const entry = failures.get(id);
-        if (!entry)
-            return undefined;
-        clearTimeout(entry.timer);
-        failures.delete(id);
-        return entry.error;
+    const takeFailure = (state) => {
+        const failure = state.failure;
+        release(state);
+        return failure;
     };
+    let closing = false;
     let disposal = null;
     const report = (message) => {
         try {
@@ -126,7 +140,7 @@ export async function startGaLoopbackBridge(options) {
         }
     };
     const server = createServer((request, response) => {
-        void handle(request, response).catch((error) => {
+        const pending = handle(request, response).catch((error) => {
             report(`bridge request failed: ${errorMessage(error)}`);
             if (!response.headersSent) {
                 writeResponseHead(response, 500, 'application/json');
@@ -136,6 +150,8 @@ export async function startGaLoopbackBridge(options) {
                 response.destroy();
             }
         });
+        reads.add(pending);
+        void pending.finally(() => reads.delete(pending));
     });
     server.on('connection', (socket) => {
         sockets.add(socket);
@@ -148,7 +164,7 @@ export async function startGaLoopbackBridge(options) {
         const entry = id ? queued.get(id) : undefined;
         // The body is always `{}`; read and ignore it so the socket stays usable.
         request.resume();
-        if (request.method !== 'POST' || !id || !entry || disposal) {
+        if (request.method !== 'POST' || !id || !entry || closing) {
             writeResponseHead(response, 404, 'application/json');
             response.end(JSON.stringify({ error: { message: 'unknown job' } }));
             return;
@@ -156,6 +172,17 @@ export async function startGaLoopbackBridge(options) {
         // A job runs at most once.
         queued.delete(id);
         clearTimeout(entry.timer);
+        const state = {
+            id,
+            failure: undefined,
+            claimed: false,
+            released: false,
+        };
+        const reference = new WeakRef(state);
+        state.reference = reference;
+        lifetimes.add(reference);
+        finalized.register(state, reference, state);
+        unclaimed.set(id, { state });
         const controller = new AbortController();
         running.add(controller);
         response.once('close', () => {
@@ -174,10 +201,22 @@ export async function startGaLoopbackBridge(options) {
             catch (error) {
                 upstream = gaPreHeaderFailureResponse(error);
             }
-            await pipeResponse(upstream, response, (error) => recordFailure(id, error));
+            await pipeResponse(upstream, response, (error) => {
+                // An adapted body owns its error even while disposal drains its reader.
+                // Unadapted jobs must not create retained failures after closing starts.
+                if (!state.released && (state.claimed || !closing))
+                    state.failure = error;
+            });
         }
         finally {
             running.delete(controller);
+            if (!closing && !state.claimed && !state.released) {
+                const entry = unclaimed.get(id);
+                if (entry) {
+                    entry.timer = setTimeout(() => release(state), GA_UNCLAIMED_JOB_TTL_MS);
+                    entry.timer.unref?.();
+                }
+            }
         }
     };
     await new Promise((resolve, reject) => {
@@ -195,7 +234,7 @@ export async function startGaLoopbackBridge(options) {
         port: address.port,
         origin,
         rewrite(original, job, env) {
-            if (disposal)
+            if (closing)
                 throw new Error('The Antigravity bridge is disposed');
             const id = randomUUID();
             const target = `${origin}/agy/${id}`;
@@ -230,6 +269,16 @@ export async function startGaLoopbackBridge(options) {
             if (!source)
                 return null;
             const reader = source.getReader();
+            const entry = unclaimed.get(id);
+            const state = entry?.state ?? {
+                id,
+                failure: undefined,
+                claimed: false,
+                released: true,
+            };
+            state.claimed = true;
+            clearTimeout(entry?.timer);
+            unclaimed.delete(id);
             const body = new ReadableStream({
                 async pull(controller) {
                     try {
@@ -238,7 +287,8 @@ export async function startGaLoopbackBridge(options) {
                             // The bridge ends a failed job's loopback body cleanly after
                             // its last byte and records the failure first, so the end of
                             // the body is where the genuine error is delivered.
-                            const failure = takeFailure(id);
+                            const failure = takeFailure(state);
+                            reader.releaseLock();
                             if (failure)
                                 controller.error(failure);
                             else
@@ -248,12 +298,18 @@ export async function startGaLoopbackBridge(options) {
                         controller.enqueue(value);
                     }
                     catch (error) {
-                        controller.error(takeFailure(id) ?? error);
+                        controller.error(takeFailure(state) ?? error);
+                        reader.releaseLock();
                     }
                 },
-                cancel(reason) {
-                    takeFailure(id);
-                    return reader.cancel(reason);
+                async cancel(reason) {
+                    release(state);
+                    try {
+                        await reader.cancel(reason);
+                    }
+                    finally {
+                        reader.releaseLock();
+                    }
                 },
             });
             return new Response(body, {
@@ -263,15 +319,32 @@ export async function startGaLoopbackBridge(options) {
             });
         },
         pendingJobs: () => queued.size + running.size,
+        retainedResponses() {
+            let claimed = 0;
+            let failures = 0;
+            let timers = 0;
+            for (const reference of lifetimes) {
+                const state = reference.deref();
+                if (!state || state.released)
+                    continue;
+                if (state.claimed)
+                    claimed++;
+                if (state.failure)
+                    failures++;
+            }
+            for (const entry of unclaimed.values()) {
+                if (entry.timer)
+                    timers++;
+            }
+            return { unclaimed: unclaimed.size, claimed, failures, timers };
+        },
         dispose() {
             if (!disposal) {
+                closing = true;
                 disposal = (async () => {
                     for (const { timer } of queued.values())
                         clearTimeout(timer);
                     queued.clear();
-                    for (const { timer } of failures.values())
-                        clearTimeout(timer);
-                    failures.clear();
                     const closed = new Promise((resolve) => {
                         server.close(() => resolve());
                     });
@@ -280,7 +353,17 @@ export async function startGaLoopbackBridge(options) {
                     }
                     for (const socket of sockets)
                         socket.destroy();
+                    await Promise.allSettled([...reads]);
                     await closed;
+                    for (const reference of lifetimes) {
+                        const state = reference.deref();
+                        // Detach every bridge owner, but leave an adapted body's terminal
+                        // error in its own closure until EOF, cancellation or collection.
+                        if (state)
+                            release(state, !state.claimed);
+                    }
+                    lifetimes.clear();
+                    unclaimed.clear();
                 })();
             }
             return disposal;
@@ -319,6 +402,10 @@ function writeResponseHead(response, status, contentType) {
  */
 async function pipeResponse(upstream, response, onFailure) {
     const contentType = upstream.headers.get('content-type') ?? 'application/octet-stream';
+    if (response.destroyed) {
+        await upstream.body?.cancel().catch(() => { });
+        return;
+    }
     writeResponseHead(response, upstream.status, contentType);
     if (!upstream.body) {
         response.end();
@@ -339,7 +426,11 @@ async function pipeResponse(upstream, response, onFailure) {
             if (eventStream) {
                 tail = (tail + Buffer.from(value).toString('latin1')).slice(-2);
             }
-            if (!response.write(value)) {
+            if (response.destroyed) {
+                await reader.cancel().catch(() => { });
+                return;
+            }
+            if (!response.write(value) && !response.destroyed) {
                 await new Promise((resolve) => {
                     response.once('drain', resolve);
                     response.once('close', resolve);
@@ -360,6 +451,9 @@ async function pipeResponse(upstream, response, onFailure) {
         response.end();
         return;
     }
+    finally {
+        reader.releaseLock();
+    }
     if (eventStream && wroteAny && tail !== '\n\n') {
         response.write(tail.endsWith('\n') ? '\n' : '\n\n');
     }
@@ -470,6 +564,19 @@ function routeDto(entry) {
         updatedAt: entry.updatedAt,
     };
 }
+/** Sidebar optionals must be omitted, not sent as own undefined JSON fields. */
+function omitUndefinedFields(value) {
+    if (Array.isArray(value))
+        return value.map(omitUndefinedFields);
+    if (value === null || typeof value !== 'object')
+        return value;
+    const prototype = Object.getPrototypeOf(value);
+    if (prototype !== Object.prototype && prototype !== null)
+        return value;
+    return Object.fromEntries(Object.entries(value)
+        .filter(([, field]) => field !== undefined)
+        .map(([key, field]) => [key, omitUndefinedFields(field)]));
+}
 export function createGaRpcActivation(options) {
     const generation = options.generation ?? createGaGeneration();
     const now = options.now ?? Date.now;
@@ -575,13 +682,13 @@ export function createGaRpcActivation(options) {
         };
         const gemini = entry(quota.gemini);
         const nonGemini = entry(quota['non-gemini']);
-        return {
+        return omitUndefinedFields({
             ...rest,
             quota: {
                 ...(gemini ? { gemini } : {}),
                 ...(nonGemini ? { 'non-gemini': nonGemini } : {}),
             },
-        };
+        });
     };
     const state = async (input, signal) => {
         const parsed = AntigravityStateInputSchema.parse(input);
@@ -744,6 +851,8 @@ export function createGaRpcActivation(options) {
 export const GA_PLUGIN_ID = 'cortexkit.antigravity-auth';
 /** The host provider whose requests the hooks handle. */
 export const GA_PROVIDER_ID = 'google';
+// OpenCode 2 uses `default` to mean that no explicit thinking variant is selected.
+const GA_HOST_DEFAULT_VARIANT = 'default';
 const REQUEST_KIND_MATCHES_HOST = true;
 void REQUEST_KIND_MATCHES_HOST;
 /** Session headers the host sets on a model request. */
@@ -751,21 +860,38 @@ const PARENT_SESSION_HEADERS = [
     'x-opencode-parent-session-id',
     'x-parent-session-id',
 ];
+/** Ownership follows the host-prepared credential, never the active connection. */
+export function isGaOwnedRequest(event) {
+    return (event.request.headers.get('x-goog-api-key') === placeholderSecret('google'));
+}
+export const GA_RETRY_OWNERSHIP_LIMIT = 512;
+export const GA_UNSUPPORTED_REQUEST_MESSAGE = 'This request cannot use Antigravity. Choose a supported Antigravity model.';
+export const GA_CUSTOM_GOOGLE_ORIGIN_MESSAGE = 'Antigravity pool requests need the default Google base URL. Remove the custom baseURL from the Google provider settings.';
 /**
- * Whether the plugin owns a Google request: every title request (the host
- * may pick an off-catalog title model) and any request for a model this
- * plugin registered or withdrew. Withdrawn selections must reach shared
- * admission rather than passing through to the host's Google transport.
+ * The retry API has no attempt identity, so only the latest Google transport
+ * per session is known. Updating an entry makes it newest; eviction makes
+ * ownership unknown rather than assuming a retry belongs to Antigravity.
  */
-export function isGaOwnedRequest(event, isRegisteredModel) {
-    // The hooks are registered for the Google provider only; the provider
-    // check repeats that scope so a host that widened it still passes other
-    // providers through untouched.
-    if (event.model.providerID !== GA_PROVIDER_ID)
-        return false;
-    return (event.kind === 'title' ||
-        isRegisteredModel(event.model.id) ||
-        isRetiredAgyModel(event.model.id));
+export class GaRequestOwnership {
+    latest = new Map();
+    record(sessionID, owned) {
+        this.latest.delete(sessionID);
+        this.latest.set(sessionID, owned);
+        if (this.latest.size > GA_RETRY_OWNERSHIP_LIMIT) {
+            const oldest = this.latest.keys().next().value;
+            if (oldest !== undefined)
+                this.latest.delete(oldest);
+        }
+    }
+    get(sessionID) {
+        return this.latest.get(sessionID);
+    }
+    delete(sessionID) {
+        this.latest.delete(sessionID);
+    }
+    clear() {
+        this.latest.clear();
+    }
 }
 /**
  * Request hook body: record an owned Google content request as a bridge job
@@ -775,17 +901,33 @@ export function isGaOwnedRequest(event, isRegisteredModel) {
  * then records it as the request's error and dispatches nothing.
  */
 export async function rewriteGaHttpRequest(event, input) {
-    if (!isGaOwnedRequest(event, input.isRegisteredModel))
+    if (event.model.providerID !== GA_PROVIDER_ID)
         return;
-    const original = event.request;
-    if (!isGaGoogleContentRequest(original))
+    const owned = isGaOwnedRequest(event);
+    input.ownership.record(event.sessionID, owned);
+    if (!owned)
         return;
+    const original = event.request;
+    if (!isGaGoogleContentRequest(original)) {
+        throw new Error(GA_UNSUPPORTED_REQUEST_MESSAGE);
+    }
+    if (!isCanonicalGoogleOrigin(original.url)) {
+        throw new Error(GA_CUSTOM_GOOGLE_ORIGIN_MESSAGE);
+    }
+    // Titles must reach the engine's title router even for retired or foreign IDs.
+    if (event.kind !== 'title' &&
+        !input.isRegisteredModel(event.model.id) &&
+        !isRetiredAgyModel(event.model.id)) {
+        throw new Error(GA_UNSUPPORTED_REQUEST_MESSAGE);
+    }
     const job = {
         sessionID: event.sessionID,
         parentSessionID: headerValue(original, PARENT_SESSION_HEADERS),
         kind: event.kind,
         modelID: event.model.id,
-        variant: event.model.variant ?? null,
+        variant: event.model.variant === GA_HOST_DEFAULT_VARIANT
+            ? null
+            : (event.model.variant ?? null),
         url: original.url,
         body: await original.clone().text(),
     };
@@ -800,11 +942,12 @@ export function adaptGaHttpResponse(event, bridge) {
 /**
  * Retry hook body. The bridge's job runner rotates accounts and retries
  * endpoints itself; a host retry on top would repeat that whole sequence.
- * Registered and withdrawn models are affected; withdrawn selections must not retry.
+ * Response adaptation cannot clear ownership: the host classifies status and
+ * consumes the body only after its response hook has run.
  */
-export function decideGaRetry(event, isRegisteredModel) {
+export function decideGaRetry(event, ownership) {
     if (event.model.providerID === GA_PROVIDER_ID &&
-        (isRegisteredModel(event.model.id) || isRetiredAgyModel(event.model.id))) {
+        ownership.get(event.sessionID) === true) {
         event.decision = { retry: false };
     }
 }
@@ -893,7 +1036,7 @@ function headerValue(request, names) {
  * the location runtime, the RPC activation and the loopback bridge for the
  * host-owned location directory; registers the `antigravity-auth` RPC, the
  * public catalog models on the Google provider, the `http.request` rewrite,
- * the `http.response` post-header error adapter and the registered-model
+ * the `http.response` post-header error adapter and the latest-owned-session
  * retry policy (all scoped to the Google provider);
  * and returns a Cleanup that removes only this activation's hooks and RPC
  * and stops everything else: producers are stopped and awaited, queued
@@ -925,9 +1068,9 @@ export async function setupGaActivation(context, factory, overrides = {}) {
     // writes drain and the runtime's files and handles close.
     try {
         if (context.integration?.transform) {
-            runtime.scope.add(await registerGaVaultActivation({
+            runtime.scope.add(await registerGaAuthMethods({
                 integration: { transform: context.integration.transform },
-            }), 'producer');
+            }, undefined, { local: services.localLogin, oauth: overrides.oauth }), 'producer');
         }
         let serving = null;
         const requireServing = () => {
@@ -968,6 +1111,36 @@ export async function setupGaActivation(context, factory, overrides = {}) {
         });
         runtime.scope.add(bridge, 'producer');
         runtime.scope.add(activation, 'producer');
+        const ownership = new GaRequestOwnership();
+        const events = new AbortController();
+        const eventSurface = context.event;
+        const eventTask = eventSurface
+            ? (async () => {
+                try {
+                    for await (const event of eventSurface.subscribe({
+                        signal: events.signal,
+                    })) {
+                        if (event.type === 'session.deleted') {
+                            ownership.delete(event.data.sessionID);
+                        }
+                    }
+                }
+                catch {
+                    if (!events.signal.aborted) {
+                        runtime.logger
+                            .createLogger('ownership')
+                            .warn('Session cleanup subscription stopped');
+                    }
+                }
+            })()
+            : undefined;
+        runtime.scope.add({
+            dispose: async () => {
+                events.abort();
+                await eventTask;
+                ownership.clear();
+            },
+        }, 'producer');
         const handlers = {
             state: (input, call) => activation.handlers.state(input, call.signal),
             apply: (input, call) => activation.handlers.apply(input, call.signal),
@@ -999,7 +1172,7 @@ export async function setupGaActivation(context, factory, overrides = {}) {
         await context.command.reload();
         const responseHook = await context.session.hook('http.response', (event) => adaptGaHttpResponse(event, bridge), { providerID: GA_PROVIDER_ID });
         runtime.scope.add(responseHook, 'producer');
-        const retryHook = await context.session.hook('retry', (event) => decideGaRetry(event, isRegisteredModel), { providerID: GA_PROVIDER_ID });
+        const retryHook = await context.session.hook('retry', (event) => decideGaRetry(event, ownership), { providerID: GA_PROVIDER_ID });
         runtime.scope.add(retryHook, 'producer');
         // The `http.request` hook, which rewrites requests to the bridge, is
         // installed last, so every request it rewrites already has the
@@ -1009,6 +1182,7 @@ export async function setupGaActivation(context, factory, overrides = {}) {
             bridge,
             isRegisteredModel,
             env: process.env,
+            ownership,
         }), { providerID: GA_PROVIDER_ID });
         runtime.scope.add(requestHook, 'producer');
     }