fix prefer_on_cloud fallback criteria

Author: hsubox76

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

@@ -29,7 +29,7 @@ import { ApiSettings } from '../types/internal';
 import * as GoogleAIMapper from '../googleai-mappers';
 import { BackendType } from '../public-types';
 import { ChromeAdapter } from '../types/chrome-adapter';
-import { callCloudOrDevice } from './helpers';
+import { callCloudOrDevice } from '../requests/hybrid-helpers';
 
 async function generateContentStreamOnCloud(
   apiSettings: ApiSettings,

packages/ai/src/methods/helpers.test.ts renamed to packages/ai/src/requests/hybrid-helpers.test.ts

@@ -17,12 +17,12 @@
 
 import { use, expect } from 'chai';
 import { SinonStub, SinonStubbedInstance, restore, stub } from 'sinon';
-import { callCloudOrDevice } from './helpers';
+import { callCloudOrDevice } from './hybrid-helpers';
 import { GenerateContentRequest, InferenceMode, AIErrorCode } from '../types';
 import { AIError } from '../errors';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
-import { ChromeAdapterImpl } from './chrome-adapter';
+import { ChromeAdapterImpl } from '../methods/chrome-adapter';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -159,7 +159,7 @@ describe('callCloudOrDevice', () => {
       expect(onDeviceCall).to.not.have.been.called;
     });
 
-    it('should fall back to onDeviceCall if inCloudCall fails with AIError', async () => {
+    it('should fall back to onDeviceCall if inCloudCall fails with AIErrorCode.FETCH_ERROR', async () => {
       inCloudCall.rejects(
         new AIError(AIErrorCode.FETCH_ERROR, 'Network error')
       );
@@ -175,7 +175,7 @@ describe('callCloudOrDevice', () => {
     });
 
     it('should re-throw other errors from inCloudCall', async () => {
-      const error = new Error('Some other error');
+      const error = new AIError(AIErrorCode.RESPONSE_ERROR, 'safety problem');
       inCloudCall.rejects(error);
       await expect(
         callCloudOrDevice(request, chromeAdapter, onDeviceCall, inCloudCall)

packages/ai/src/methods/helpers.ts renamed to packages/ai/src/requests/hybrid-helpers.ts

@@ -22,7 +22,16 @@ import {
   AIErrorCode,
   ChromeAdapter
 } from '../types';
-import { ChromeAdapterImpl } from './chrome-adapter';
+import { ChromeAdapterImpl } from '../methods/chrome-adapter';
+
+const errorsCausingFallback: AIErrorCode[] = [
+  // most network errors
+  AIErrorCode.FETCH_ERROR,
+  // fallback code for all other errors in makeRequest
+  AIErrorCode.ERROR,
+  // error due to API not being enabled in project
+  AIErrorCode.API_NOT_ENABLED
+];
 
 /**
  * Dispatches a request to the appropriate backend (on-device or in-cloud)
@@ -58,7 +67,7 @@ export async function callCloudOrDevice<Response>(
       try {
         return await inCloudCall();
       } catch (e) {
-        if (e instanceof AIError) {
+        if (e instanceof AIError && errorsCausingFallback.includes(e.code)) {
           return onDeviceCall();
         }
         throw e;