Use logger class and send App Check dummy token in header (#8591)

Author: hsubox76

.changeset/tall-peas-tell.md

@@ -0,0 +1,5 @@
+---
+'@firebase/vertexai': patch
+---
+
+Send App Check dummy token in header if there is an App Check getToken error.

packages/vertexai/src/logger.ts

@@ -0,0 +1,20 @@
+/**
+ * @license
+ * Copyright 2024 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 { Logger } from '@firebase/logger';
+
+export const logger = new Logger('@firebase/vertexai');

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

@@ -87,7 +87,7 @@ describe('ChatSession', () => {
         match.any
       );
       await clock.runAllAsync();
-      expect(consoleStub.args[0][0].toString()).to.include(
+      expect(consoleStub.args[0][1].toString()).to.include(
         // Firefox has different wording when a property is undefined
         'undefined'
       );

packages/vertexai/src/methods/chat-session.ts

@@ -29,6 +29,7 @@ import { formatBlockErrorMessage } from '../requests/response-helpers';
 import { validateChatHistory } from './chat-session-helpers';
 import { generateContent, generateContentStream } from './generate-content';
 import { ApiSettings } from '../types/internal';
+import { logger } from '../logger';
 
 /**
  * Do not log a message for this error.
@@ -112,7 +113,7 @@ export class ChatSession {
         } else {
           const blockErrorMessage = formatBlockErrorMessage(result.response);
           if (blockErrorMessage) {
-            console.warn(
+            logger.warn(
               `sendMessage() was unsuccessful. ${blockErrorMessage}. Inspect response object for details.`
             );
           }
@@ -169,7 +170,7 @@ export class ChatSession {
         } else {
           const blockErrorMessage = formatBlockErrorMessage(response);
           if (blockErrorMessage) {
-            console.warn(
+            logger.warn(
               `sendMessageStream() was unsuccessful. ${blockErrorMessage}. Inspect response object for details.`
             );
           }
@@ -182,7 +183,7 @@ export class ChatSession {
         if (e.message !== SILENT_ERROR) {
           // Users do not have access to _sendPromise to catch errors
           // downstream from streamPromise, so they should not throw.
-          console.error(e);
+          logger.error(e);
         }
       });
     return streamPromise;

packages/vertexai/src/requests/request.test.ts

@@ -16,7 +16,7 @@
  */
 
 import { expect, use } from 'chai';
-import { restore, stub } from 'sinon';
+import { match, restore, stub } from 'sinon';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
 import { RequestUrl, Task, getHeaders, makeRequest } from './request';
@@ -169,13 +169,18 @@ describe('request methods', () => {
           project: 'myproject',
           location: 'moon',
           getAppCheckToken: () =>
-            Promise.resolve({ token: 'token', error: Error('oops') })
+            Promise.resolve({ token: 'dummytoken', error: Error('oops') })
         },
         true,
         {}
       );
+      const warnStub = stub(console, 'warn');
       const headers = await getHeaders(fakeUrl);
-      expect(headers.has('X-Firebase-AppCheck')).to.be.false;
+      expect(headers.get('X-Firebase-AppCheck')).to.equal('dummytoken');
+      expect(warnStub).to.be.calledWith(
+        match(/vertexai/),
+        match(/App Check.*oops/)
+      );
     });
     it('adds auth token if it exists', async () => {
       const headers = await getHeaders(fakeUrl);

packages/vertexai/src/requests/request.ts

@@ -24,6 +24,7 @@ import {
   LANGUAGE_TAG,
   PACKAGE_VERSION
 } from '../constants';
+import { logger } from '../logger';
 
 export enum Task {
   GENERATE_CONTENT = 'generateContent',
@@ -83,8 +84,13 @@ export async function getHeaders(url: RequestUrl): Promise<Headers> {
   headers.append('x-goog-api-key', url.apiSettings.apiKey);
   if (url.apiSettings.getAppCheckToken) {
     const appCheckToken = await url.apiSettings.getAppCheckToken();
-    if (appCheckToken && !appCheckToken.error) {
+    if (appCheckToken) {
       headers.append('X-Firebase-AppCheck', appCheckToken.token);
+      if (appCheckToken.error) {
+        logger.warn(
+          `Unable to obtain a valid App Check token: ${appCheckToken.error.message}`
+        );
+      }
     }
   }
 

packages/vertexai/src/requests/response-helpers.ts

@@ -24,6 +24,7 @@ import {
   VertexAIErrorCode
 } from '../types';
 import { VertexAIError } from '../errors';
+import { logger } from '../logger';
 
 /**
  * Creates an EnhancedGenerateContentResponse object that has helper functions and
@@ -56,7 +57,7 @@ export function addHelpers(
   (response as EnhancedGenerateContentResponse).text = () => {
     if (response.candidates && response.candidates.length > 0) {
       if (response.candidates.length > 1) {
-        console.warn(
+        logger.warn(
           `This response had ${response.candidates.length} ` +
             `candidates. Returning text from the first candidate only. ` +
             `Access response.candidates directly to use the other candidates.`
@@ -88,7 +89,7 @@ export function addHelpers(
   (response as EnhancedGenerateContentResponse).functionCalls = () => {
     if (response.candidates && response.candidates.length > 0) {
       if (response.candidates.length > 1) {
-        console.warn(
+        logger.warn(
           `This response had ${response.candidates.length} ` +
             `candidates. Returning function calls from the first candidate only. ` +
             `Access response.candidates directly to use the other candidates.`