Add limitedUseToken option to AI SDK

Author: hsubox76

common/api-review/ai.api.md

@@ -15,6 +15,7 @@ export interface AI {
     backend: Backend;
     // @deprecated (undocumented)
     location: string;
+    options?: AIOptions;
 }
 
 // @public
@@ -61,6 +62,7 @@ export abstract class AIModel {
 
 // @public
 export interface AIOptions {
+    appCheck?: AppCheckOptions;
     backend: Backend;
 }
 
@@ -75,6 +77,12 @@ export class AnyOfSchema extends Schema {
     toJSON(): SchemaRequest;
 }
 
+// @public
+export interface AppCheckOptions {
+    // (undocumented)
+    limitedUseTokens?: boolean;
+}
+
 // @public
 export class ArraySchema extends Schema {
     constructor(schemaParams: SchemaParams, items: TypedSchema);

packages/ai/src/api.ts

@@ -72,18 +72,24 @@ declare module '@firebase/component' {
  *
  * @public
  */
-export function getAI(
-  app: FirebaseApp = getApp(),
-  options: AIOptions = { backend: new GoogleAIBackend() }
-): AI {
+export function getAI(app: FirebaseApp = getApp(), options?: AIOptions): AI {
   app = getModularInstance(app);
   // Dependencies
   const AIProvider: Provider<'AI'> = _getProvider(app, AI_TYPE);
 
-  const identifier = encodeInstanceIdentifier(options.backend);
-  return AIProvider.getImmediate({
+  const finalOptions = {
+    backend: options?.backend ?? new GoogleAIBackend(),
+    appCheck: options?.appCheck ?? { limitedUseTokens: false }
+  };
+
+  const identifier = encodeInstanceIdentifier(finalOptions.backend);
+  const aiInstance = AIProvider.getImmediate({
     identifier
   });
+
+  aiInstance.options = finalOptions;
+
+  return aiInstance;
 }
 
 /**

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

@@ -90,8 +90,13 @@ export abstract class AIModel {
           return Promise.resolve({ token });
         };
       } else if ((ai as AIService).appCheck) {
-        this._apiSettings.getAppCheckToken = () =>
-          (ai as AIService).appCheck!.getToken();
+        if (ai.options?.appCheck?.limitedUseTokens) {
+          this._apiSettings.getAppCheckToken = () =>
+            (ai as AIService).appCheck!.getLimitedUseToken();
+        } else {
+          this._apiSettings.getAppCheckToken = () =>
+            (ai as AIService).appCheck!.getToken();
+        }
       }
 
       if ((ai as AIService).auth) {

packages/ai/src/public-types.ts

@@ -38,6 +38,10 @@ export interface AI {
    * Vertex AI Gemini API (using {@link VertexAIBackend}).
    */
   backend: Backend;
+  /**
+   * Options applied to this {@link AI} instance.
+   */
+  options?: AIOptions;
   /**
    * @deprecated use `AI.backend.location` instead.
    *
@@ -92,4 +96,15 @@ export interface AIOptions {
    * The backend configuration to use for the AI service instance.
    */
   backend: Backend;
+  /**
+   * Configures App Check usage for this AI service instance.
+   */
+  appCheck?: AppCheckOptions;
+}
+
+/**
+ * Configures App Check usage for this AI service instance.
+ */
+export interface AppCheckOptions {
+  limitedUseTokens?: boolean;
 }

packages/ai/src/service.ts

@@ -16,7 +16,7 @@
  */
 
 import { FirebaseApp, _FirebaseService } from '@firebase/app';
-import { AI } from './public-types';
+import { AI, AIOptions } from './public-types';
 import {
   AppCheckInternalComponentName,
   FirebaseAppCheckInternal
@@ -31,6 +31,7 @@ import { Backend, VertexAIBackend } from './backend';
 export class AIService implements AI, _FirebaseService {
   auth: FirebaseAuthInternal | null;
   appCheck: FirebaseAppCheckInternal | null;
+  _options?: AIOptions;
   location: string; // This is here for backwards-compatibility
 
   constructor(
@@ -54,4 +55,12 @@ export class AIService implements AI, _FirebaseService {
   _delete(): Promise<void> {
     return Promise.resolve();
   }
+
+  set options(optionsToSet: AIOptions) {
+    this.options = optionsToSet;
+  }
+
+  get options(): AIOptions {
+    return this.options;
+  }
 }