fix(ai): Fix node tests

Author: hsubox76

packages/ai/src/factory-node.ts

@@ -0,0 +1,51 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import {
+  ComponentContainer,
+  InstanceFactoryOptions
+} from '@firebase/component';
+import { AIError } from './errors';
+import { decodeInstanceIdentifier } from './helpers';
+import { AIService } from './service';
+import { AIErrorCode } from './types';
+
+export function factory(
+  container: ComponentContainer,
+  { instanceIdentifier }: InstanceFactoryOptions
+): AIService {
+  if (!instanceIdentifier) {
+    throw new AIError(
+      AIErrorCode.ERROR,
+      'AIService instance identifier is undefined.'
+    );
+  }
+
+  const backend = decodeInstanceIdentifier(instanceIdentifier);
+
+  // getImmediate for FirebaseApp will always succeed
+  const app = container.getProvider('app').getImmediate();
+  const auth = container.getProvider('auth-internal');
+  const appCheckProvider = container.getProvider('app-check-internal');
+
+  return new AIService(
+    app,
+    backend,
+    auth,
+    appCheckProvider
+  );
+}

packages/ai/src/index.node.ts

@@ -22,36 +22,16 @@
  */
 
 import { registerVersion, _registerComponent } from '@firebase/app';
-import { AIService } from './service';
 import { AI_TYPE } from './constants';
 import { Component, ComponentType } from '@firebase/component';
 import { name, version } from '../package.json';
-import { decodeInstanceIdentifier } from './helpers';
-import { AIError } from './errors';
-import { AIErrorCode } from './public-types';
+import { factory } from './factory-node';
 
 function registerAI(): void {
   _registerComponent(
-    new Component(
-      AI_TYPE,
-      (container, { instanceIdentifier }) => {
-        if (!instanceIdentifier) {
-          throw new AIError(
-            AIErrorCode.ERROR,
-            'AIService instance identifier is undefined.'
-          );
-        }
-
-        const backend = decodeInstanceIdentifier(instanceIdentifier);
-
-        // getImmediate for FirebaseApp will always succeed
-        const app = container.getProvider('app').getImmediate();
-        const auth = container.getProvider('auth-internal');
-        const appCheckProvider = container.getProvider('app-check-internal');
-        return new AIService(app, backend, auth, appCheckProvider);
-      },
-      ComponentType.PUBLIC
-    ).setMultipleInstances(true)
+    new Component(AI_TYPE, factory, ComponentType.PUBLIC).setMultipleInstances(
+      true
+    )
   );
 
   registerVersion(name, version, 'node');

packages/ai/src/methods/chat-session.test.ts

@@ -20,11 +20,11 @@ import { match, restore, stub, useFakeTimers } from 'sinon';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
 import * as generateContentMethods from './generate-content';
-import { Content, GenerateContentStreamResult, InferenceMode } from '../types';
+import { Content, GenerateContentStreamResult } from '../types';
 import { ChatSession } from './chat-session';
 import { ApiSettings } from '../types/internal';
 import { VertexAIBackend } from '../backend';
-import { ChromeAdapterImpl } from './chrome-adapter';
+import { fakeChromeAdapter } from '../../test-utils/get-fake-firebase-services';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -37,12 +37,6 @@ const fakeApiSettings: ApiSettings = {
   backend: new VertexAIBackend()
 };
 
-const fakeChromeAdapter = new ChromeAdapterImpl(
-  // @ts-expect-error
-  undefined,
-  InferenceMode.PREFER_ON_DEVICE
-);
-
 describe('ChatSession', () => {
   afterEach(() => {
     restore();

packages/ai/src/methods/count-tokens.test.ts

@@ -27,7 +27,7 @@ import { ApiSettings } from '../types/internal';
 import { Task } from '../requests/request';
 import { mapCountTokensRequest } from '../googleai-mappers';
 import { GoogleAIBackend, VertexAIBackend } from '../backend';
-import { ChromeAdapterImpl } from './chrome-adapter';
+import { fakeChromeAdapter } from '../../test-utils/get-fake-firebase-services';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -52,12 +52,6 @@ const fakeRequestParams: CountTokensRequest = {
   contents: [{ parts: [{ text: 'hello' }], role: 'user' }]
 };
 
-const fakeChromeAdapter = new ChromeAdapterImpl(
-  // @ts-expect-error
-  undefined,
-  InferenceMode.PREFER_ON_DEVICE
-);
-
 describe('countTokens()', () => {
   afterEach(() => {
     restore();
@@ -197,11 +191,10 @@ describe('countTokens()', () => {
     });
   });
   it('throws if mode is ONLY_ON_DEVICE', async () => {
-    const chromeAdapter = new ChromeAdapterImpl(
-      // @ts-expect-error
-      undefined,
-      InferenceMode.ONLY_ON_DEVICE
-    );
+    const chromeAdapter = {
+      ...fakeChromeAdapter,
+      mode: InferenceMode.ONLY_ON_DEVICE
+    };
     await expect(
       countTokens(fakeApiSettings, 'model', fakeRequestParams, chromeAdapter)
     ).to.be.rejectedWith(

packages/ai/src/methods/count-tokens.ts

@@ -28,7 +28,6 @@ import { ApiSettings } from '../types/internal';
 import * as GoogleAIMapper from '../googleai-mappers';
 import { BackendType } from '../public-types';
 import { ChromeAdapter } from '../types/chrome-adapter';
-import { ChromeAdapterImpl } from './chrome-adapter';
 
 export async function countTokensOnCloud(
   apiSettings: ApiSettings,
@@ -61,9 +60,7 @@ export async function countTokens(
   chromeAdapter?: ChromeAdapter,
   requestOptions?: RequestOptions
 ): Promise<CountTokensResponse> {
-  if (
-    (chromeAdapter as ChromeAdapterImpl)?.mode === InferenceMode.ONLY_ON_DEVICE
-  ) {
+  if (chromeAdapter?.mode === InferenceMode.ONLY_ON_DEVICE) {
     throw new AIError(
       AIErrorCode.UNSUPPORTED,
       'countTokens() is not supported for on-device models.'

packages/ai/src/methods/generate-content.test.ts

@@ -28,7 +28,6 @@ import {
   HarmBlockMethod,
   HarmBlockThreshold,
   HarmCategory,
-  InferenceMode,
   Language,
   Outcome
 } from '../types';
@@ -37,17 +36,11 @@ import { Task } from '../requests/request';
 import { AIError } from '../api';
 import { mapGenerateContentRequest } from '../googleai-mappers';
 import { GoogleAIBackend, VertexAIBackend } from '../backend';
-import { ChromeAdapterImpl } from './chrome-adapter';
+import { fakeChromeAdapter } from '../../test-utils/get-fake-firebase-services';
 
 use(sinonChai);
 use(chaiAsPromised);
 
-const fakeChromeAdapter = new ChromeAdapterImpl(
-  // @ts-expect-error
-  undefined,
-  InferenceMode.PREFER_ON_DEVICE
-);
-
 const fakeApiSettings: ApiSettings = {
   apiKey: 'key',
   project: 'my-project',

packages/ai/src/models/generative-model.test.ts

@@ -20,7 +20,8 @@ import {
   FunctionCallingMode,
   AI,
   InferenceMode,
-  AIErrorCode
+  AIErrorCode,
+  ChromeAdapter
 } from '../public-types';
 import * as request from '../requests/request';
 import { SinonStub, match, restore, stub } from 'sinon';
@@ -30,9 +31,9 @@ import {
 } from '../../test-utils/mock-response';
 import sinonChai from 'sinon-chai';
 import { VertexAIBackend } from '../backend';
-import { ChromeAdapterImpl } from '../methods/chrome-adapter';
 import { AIError } from '../errors';
 import chaiAsPromised from 'chai-as-promised';
+import { fakeChromeAdapter } from '../../test-utils/get-fake-firebase-services';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -51,12 +52,6 @@ const fakeAI: AI = {
   location: 'us-central1'
 };
 
-const fakeChromeAdapter = new ChromeAdapterImpl(
-  // @ts-expect-error
-  undefined,
-  InferenceMode.PREFER_ON_DEVICE
-);
-
 describe('GenerativeModel', () => {
   it('passes params through to generateContent', async () => {
     const genModel = new GenerativeModel(
@@ -436,7 +431,7 @@ describe('GenerativeModel', () => {
 
 describe('GenerativeModel dispatch logic', () => {
   let makeRequestStub: SinonStub;
-  let mockChromeAdapter: ChromeAdapterImpl;
+  let mockChromeAdapter: ChromeAdapter;
 
   function stubMakeRequest(stream?: boolean): void {
     if (stream) {

packages/ai/src/requests/hybrid-helpers.test.ts

@@ -18,17 +18,21 @@
 import { use, expect } from 'chai';
 import { SinonStub, SinonStubbedInstance, restore, stub } from 'sinon';
 import { callCloudOrDevice } from './hybrid-helpers';
-import { GenerateContentRequest, InferenceMode, AIErrorCode } from '../types';
+import {
+  GenerateContentRequest,
+  InferenceMode,
+  AIErrorCode,
+  ChromeAdapter
+} from '../types';
 import { AIError } from '../errors';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
-import { ChromeAdapterImpl } from '../methods/chrome-adapter';
 
 use(sinonChai);
 use(chaiAsPromised);
 
 describe('callCloudOrDevice', () => {
-  let chromeAdapter: SinonStubbedInstance<ChromeAdapterImpl>;
+  let chromeAdapter: SinonStubbedInstance<ChromeAdapter>;
   let onDeviceCall: SinonStub;
   let inCloudCall: SinonStub;
   let request: GenerateContentRequest;

packages/ai/src/requests/hybrid-helpers.ts

@@ -22,7 +22,6 @@ import {
   AIErrorCode,
   ChromeAdapter
 } from '../types';
-import { ChromeAdapterImpl } from '../methods/chrome-adapter';
 
 const errorsCausingFallback: AIErrorCode[] = [
   // most network errors
@@ -52,7 +51,7 @@ export async function callCloudOrDevice<Response>(
   if (!chromeAdapter) {
     return inCloudCall();
   }
-  switch ((chromeAdapter as ChromeAdapterImpl).mode) {
+  switch (chromeAdapter.mode) {
     case InferenceMode.ONLY_ON_DEVICE:
       if (await chromeAdapter.isAvailable(request)) {
         return onDeviceCall();
@@ -80,9 +79,7 @@ export async function callCloudOrDevice<Response>(
     default:
       throw new AIError(
         AIErrorCode.ERROR,
-        `Unexpected infererence mode: ${
-          (chromeAdapter as ChromeAdapterImpl).mode
-        }`
+        `Unexpected infererence mode: ${chromeAdapter.mode}`
       );
   }
 }

packages/ai/src/service.ts

@@ -16,7 +16,13 @@
  */
 
 import { FirebaseApp, _FirebaseService } from '@firebase/app';
-import { AI, AIOptions, InferenceMode, OnDeviceParams } from './public-types';
+import {
+  AI,
+  AIOptions,
+  ChromeAdapter,
+  InferenceMode,
+  OnDeviceParams
+} from './public-types';
 import {
   AppCheckInternalComponentName,
   FirebaseAppCheckInternal
@@ -27,7 +33,6 @@ import {
   FirebaseAuthInternalName
 } from '@firebase/auth-interop-types';
 import { Backend, VertexAIBackend } from './backend';
-import { ChromeAdapterImpl } from './methods/chrome-adapter';
 
 export class AIService implements AI, _FirebaseService {
   auth: FirebaseAuthInternal | null;
@@ -44,7 +49,7 @@ export class AIService implements AI, _FirebaseService {
       mode: InferenceMode,
       window?: Window,
       params?: OnDeviceParams
-    ) => ChromeAdapterImpl | undefined
+    ) => ChromeAdapter | undefined
   ) {
     const appCheck = appCheckProvider?.getImmediate({ optional: true });
     const auth = authProvider?.getImmediate({ optional: true });