Skip to content

Add limited use tokens for ai logic #1454

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AustinBenoit wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add limited use tokens for ai logic #1454

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c6d44fc
Add limited use tokens for ai logic
AustinBenoit May 12, 2026
e7852d0
include app check in the firebase ai test app
AustinBenoit May 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions app/src/internal/FirebaseInterops.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -442,9 +442,9 @@ internal static async Task AddFirebaseTokensAsync(HttpRequestMessage request, Fi
}

// Adds the other Firebase tokens to the WebSocket, as available.
internal static async Task AddFirebaseTokensAsync(ClientWebSocket socket, FirebaseApp firebaseApp, string authTokenPrefix = "Firebase")
internal static async Task AddFirebaseTokensAsync(ClientWebSocket socket, FirebaseApp firebaseApp, string authTokenPrefix = "Firebase", bool limitedUseAppCheckTokens = false)
{
string appCheckToken = await GetAppCheckTokenAsync(firebaseApp);
string appCheckToken = await GetAppCheckTokenAsync(firebaseApp, limitedUseAppCheckTokens);
if (!string.IsNullOrEmpty(appCheckToken))
{
socket.Options.SetRequestHeader(appCheckHeader, appCheckToken);
Expand Down
33 changes: 22 additions & 11 deletions firebaseai/src/FirebaseAI.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,11 +92,13 @@ public override readonly string ToString()

private readonly FirebaseApp _firebaseApp;
private readonly Backend _backend;
private readonly bool _useLimitedUseAppCheckTokens;

private FirebaseAI(FirebaseApp firebaseApp, Backend backend)
private FirebaseAI(FirebaseApp firebaseApp, Backend backend, bool useLimitedUseAppCheckTokens)
{
_firebaseApp = firebaseApp;
_backend = backend;
_useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens;
}

/// <summary>
Expand All @@ -114,18 +116,22 @@ public static FirebaseAI DefaultInstance
/// Returns a `FirebaseAI` instance with the default `FirebaseApp` and the given Backend.
/// </summary>
/// <param name="backend">The backend AI service to use.</param>
/// <param name="useLimitedUseAppCheckTokens">Whether to use limited use App Check tokens for
/// requests.</param>
/// <returns>A configured instance of `FirebaseAI`.</returns>
public static FirebaseAI GetInstance(Backend? backend = null)
public static FirebaseAI GetInstance(Backend? backend = null, bool useLimitedUseAppCheckTokens = false)
{
return GetInstance(FirebaseApp.DefaultInstance, backend);
return GetInstance(FirebaseApp.DefaultInstance, backend, useLimitedUseAppCheckTokens);
}
/// <summary>
/// Returns a `FirebaseAI` instance with the given `FirebaseApp` and Backend.
/// </summary>
/// <param name="app">The custom `FirebaseApp` used for initialization.</param>
/// <param name="backend">The backend AI service to use.</param>
/// <param name="useLimitedUseAppCheckTokens">Whether to use limited use App Check tokens for
/// requests.</param>
/// <returns>A configured instance of `FirebaseAI`.</returns>
public static FirebaseAI GetInstance(FirebaseApp app, Backend? backend = null)
public static FirebaseAI GetInstance(FirebaseApp app, Backend? backend = null, bool useLimitedUseAppCheckTokens = false)
{
if (app == null)
{
Expand All @@ -135,13 +141,13 @@ public static FirebaseAI GetInstance(FirebaseApp app, Backend? backend = null)
Backend resolvedBackend = backend ?? Backend.GoogleAI();

// FirebaseAI instances are keyed by a combination of the app name and backend.
string key = $"{app.Name}::{resolvedBackend}";
string key = $"{app.Name}::{resolvedBackend}::{useLimitedUseAppCheckTokens}";
if (_instances.ContainsKey(key))
{
return _instances[key];
}

return _instances.GetOrAdd(key, _ => new FirebaseAI(app, resolvedBackend));
return _instances.GetOrAdd(key, _ => new FirebaseAI(app, resolvedBackend, useLimitedUseAppCheckTokens));
}

/// <summary>
Expand Down Expand Up @@ -173,7 +179,8 @@ public GenerativeModel GetGenerativeModel(
{
return new GenerativeModel(_firebaseApp, _backend, modelName,
generationConfig, safetySettings, tools,
toolConfig, systemInstruction, requestOptions);
toolConfig, systemInstruction, requestOptions,
_useLimitedUseAppCheckTokens);
}

/// <summary>
Expand All @@ -200,7 +207,8 @@ public LiveGenerativeModel GetLiveModel(
{
return new LiveGenerativeModel(_firebaseApp, _backend, modelName,
liveGenerationConfig, tools,
systemInstruction, requestOptions);
systemInstruction, requestOptions,
_useLimitedUseAppCheckTokens);
}

/// <summary>
Expand Down Expand Up @@ -230,7 +238,8 @@ public ImagenModel GetImagenModel(
RequestOptions? requestOptions = null)
{
return new ImagenModel(_firebaseApp, _backend, modelName,
generationConfig, safetySettings, requestOptions);
generationConfig, safetySettings, requestOptions,
_useLimitedUseAppCheckTokens);
}

/// <summary>
Expand All @@ -241,7 +250,8 @@ public ImagenModel GetImagenModel(
public TemplateGenerativeModel GetTemplateGenerativeModel(
RequestOptions? requestOptions = null)
{
return new TemplateGenerativeModel(_firebaseApp, _backend, requestOptions);
return new TemplateGenerativeModel(_firebaseApp, _backend, requestOptions,
_useLimitedUseAppCheckTokens);
}

/// <summary>
Expand All @@ -259,7 +269,8 @@ public TemplateGenerativeModel GetTemplateGenerativeModel(
public TemplateImagenModel GetTemplateImagenModel(
RequestOptions? requestOptions = null)
{
return new TemplateImagenModel(_firebaseApp, _backend, requestOptions);
return new TemplateImagenModel(_firebaseApp, _backend, requestOptions,
_useLimitedUseAppCheckTokens);
}
}

Expand Down
11 changes: 7 additions & 4 deletions firebaseai/src/GenerativeModel.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ public class GenerativeModel
private readonly ToolConfig? _toolConfig;
private readonly ModelContent? _systemInstruction;
private readonly RequestOptions? _requestOptions;
private readonly bool _useLimitedUseAppCheckTokens;

private readonly HttpClient _httpClient;
// String prefix to look for when handling streaming a response.
Expand All @@ -62,7 +63,8 @@ internal GenerativeModel(FirebaseApp firebaseApp,
Tool[] tools = null,
ToolConfig? toolConfig = null,
ModelContent? systemInstruction = null,
RequestOptions? requestOptions = null)
RequestOptions? requestOptions = null,
bool useLimitedUseAppCheckTokens = false)
{
_firebaseApp = firebaseApp;
_backend = backend;
Expand All @@ -74,6 +76,7 @@ internal GenerativeModel(FirebaseApp firebaseApp,
// Make sure that the system instructions have the role "system".
_systemInstruction = systemInstruction?.ConvertToSystem();
_requestOptions = requestOptions;
_useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens;

// Create a HttpClient using the timeout requested, or the default one.
_httpClient = new HttpClient()
Expand Down Expand Up @@ -272,7 +275,7 @@ private async Task<GenerateContentResponse> GenerateContentAsyncInternal(
Firebase.AI.Internal.HttpHelpers.GetURL(_firebaseApp, _backend, _modelName) + ":generateContent");

// Set the request headers
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp);
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp, limitedUseAppCheckTokens: _useLimitedUseAppCheckTokens);
Comment thread
AustinBenoit marked this conversation as resolved.

// Set the content
string bodyJson = MakeGenerateContentRequest(content);
Expand Down Expand Up @@ -302,7 +305,7 @@ private async IAsyncEnumerable<GenerateContentResponse> GenerateContentStreamAsy
Firebase.AI.Internal.HttpHelpers.GetURL(_firebaseApp, _backend, _modelName) + ":streamGenerateContent?alt=sse");

// Set the request headers
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp);
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp, limitedUseAppCheckTokens: _useLimitedUseAppCheckTokens);

// Set the content
string bodyJson = MakeGenerateContentRequest(content);
Expand Down Expand Up @@ -342,7 +345,7 @@ private async Task<CountTokensResponse> CountTokensAsyncInternal(
Firebase.AI.Internal.HttpHelpers.GetURL(_firebaseApp, _backend, _modelName) + ":countTokens");

// Set the request headers
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp);
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp, limitedUseAppCheckTokens: _useLimitedUseAppCheckTokens);

// Set the content
string bodyJson = MakeCountTokensRequest(content);
Expand Down
14 changes: 10 additions & 4 deletions firebaseai/src/Imagen/ImagenModel.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ public class ImagenModel
private readonly ImagenGenerationConfig? _generationConfig;
private readonly ImagenSafetySettings? _safetySettings;
private readonly RequestOptions? _requestOptions;
private readonly bool _useLimitedUseAppCheckTokens;

private readonly HttpClient _httpClient;

Expand All @@ -61,14 +62,16 @@ internal ImagenModel(FirebaseApp firebaseApp,
string modelName,
ImagenGenerationConfig? generationConfig = null,
ImagenSafetySettings? safetySettings = null,
RequestOptions? requestOptions = null)
RequestOptions? requestOptions = null,
bool useLimitedUseAppCheckTokens = false)
{
_firebaseApp = firebaseApp;
_backend = backend;
_modelName = modelName;
_generationConfig = generationConfig;
_safetySettings = safetySettings;
_requestOptions = requestOptions;
_useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens;

// Create a HttpClient using the timeout requested, or the default one.
_httpClient = new HttpClient()
Expand Down Expand Up @@ -101,7 +104,7 @@ private async Task<ImagenGenerationResponse<ImagenInlineImage>> GenerateImagesAs
Firebase.AI.Internal.HttpHelpers.GetURL(_firebaseApp, _backend, _modelName) + ":predict");

// Set the request headers
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp);
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp, limitedUseAppCheckTokens: _useLimitedUseAppCheckTokens);
Comment thread
AustinBenoit marked this conversation as resolved.

// Set the content
string bodyJson = MakeGenerateImagenRequest(prompt);
Expand Down Expand Up @@ -182,6 +185,7 @@ public class TemplateImagenModel
{
private readonly FirebaseApp _firebaseApp;
private readonly FirebaseAI.Backend _backend;
private readonly bool _useLimitedUseAppCheckTokens;

private readonly HttpClient _httpClient;

Expand All @@ -190,10 +194,12 @@ public class TemplateImagenModel
/// Use `FirebaseAI.GetTemplateImagenModel` instead to ensure proper initialization and configuration of the `TemplateImagenModel`.
/// </summary>
internal TemplateImagenModel(FirebaseApp firebaseApp,
FirebaseAI.Backend backend, RequestOptions? requestOptions = null)
FirebaseAI.Backend backend, RequestOptions? requestOptions = null,
bool useLimitedUseAppCheckTokens = false)
{
_firebaseApp = firebaseApp;
_backend = backend;
_useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens;

// Create a HttpClient using the timeout requested, or the default one.
_httpClient = new HttpClient()
Expand All @@ -217,7 +223,7 @@ public async Task<ImagenGenerationResponse<ImagenInlineImage>> GenerateImagesAsy
Firebase.AI.Internal.HttpHelpers.GetTemplateURL(_firebaseApp, _backend, templateId) + ":templatePredict");

// Set the request headers
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp);
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp, limitedUseAppCheckTokens: _useLimitedUseAppCheckTokens);
Comment thread
AustinBenoit marked this conversation as resolved.

// Set the content
Dictionary<string, object> jsonDict = new()
Expand Down
7 changes: 5 additions & 2 deletions firebaseai/src/LiveGenerativeModel.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ public class LiveGenerativeModel
private readonly Tool[] _tools;
private readonly ModelContent? _systemInstruction;
private readonly RequestOptions? _requestOptions;
private readonly bool _useLimitedUseAppCheckTokens;

/// <summary>
/// Intended for internal use only.
Expand All @@ -60,7 +61,8 @@ internal LiveGenerativeModel(FirebaseApp firebaseApp,
LiveGenerationConfig? liveConfig = null,
Tool[] tools = null,
ModelContent? systemInstruction = null,
RequestOptions? requestOptions = null)
RequestOptions? requestOptions = null,
bool useLimitedUseAppCheckTokens = false)
{
_firebaseApp = firebaseApp;
_backend = backend;
Expand All @@ -69,6 +71,7 @@ internal LiveGenerativeModel(FirebaseApp firebaseApp,
_tools = tools;
_systemInstruction = systemInstruction;
_requestOptions = requestOptions;
_useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens;
}

private string GetURL()
Expand Down Expand Up @@ -126,7 +129,7 @@ private async Task<ClientWebSocket> CreateClientWebSocketAsync(CancellationToken
clientWebSocket.Options.SetRequestHeader("X-Firebase-AppVersion", UnityEngine.Application.version);
}
// Add additional Firebase tokens to the header.
await Firebase.Internal.FirebaseInterops.AddFirebaseTokensAsync(clientWebSocket, _firebaseApp);
await Firebase.Internal.FirebaseInterops.AddFirebaseTokensAsync(clientWebSocket, _firebaseApp, limitedUseAppCheckTokens: _useLimitedUseAppCheckTokens);

// Add a timeout to the initial connection, using the RequestOptions.
using var connectionCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
Expand Down
9 changes: 6 additions & 3 deletions firebaseai/src/TemplateGenerativeModel.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ public class TemplateGenerativeModel
private readonly FirebaseApp _firebaseApp;
private readonly FirebaseAI.Backend _backend;
private readonly RequestOptions? _requestOptions;
private readonly bool _useLimitedUseAppCheckTokens;

private readonly HttpClient _httpClient;

Expand All @@ -46,10 +47,12 @@ public class TemplateGenerativeModel
/// </summary>
internal TemplateGenerativeModel(FirebaseApp firebaseApp,
FirebaseAI.Backend backend,
RequestOptions? requestOptions = null)
RequestOptions? requestOptions = null,
bool useLimitedUseAppCheckTokens = false)
{
_firebaseApp = firebaseApp;
_backend = backend;
_useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens;

// Create a HttpClient using the timeout requested, or the default one.
_httpClient = new HttpClient()
Expand Down Expand Up @@ -148,7 +151,7 @@ private async Task<GenerateContentResponse> GenerateContentAsyncInternal(
Firebase.AI.Internal.HttpHelpers.GetTemplateURL(_firebaseApp, _backend, templateId) + ":templateGenerateContent");

// Set the request headers
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp);
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp, limitedUseAppCheckTokens: _useLimitedUseAppCheckTokens);
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The HttpHelpers.SetRequestHeaders method is called with a new parameter that isn't defined in the provided changes for that class. Using FirebaseInterops.AddFirebaseTokensAsync is a viable alternative that is already updated in this PR to handle HttpRequestMessage objects.

      await Firebase.Internal.FirebaseInterops.AddFirebaseTokensAsync(request, _firebaseApp, limitedUseAppCheckTokens: _useLimitedUseAppCheckTokens);

Copy link
Copy Markdown
Contributor Author

@AustinBenoit AustinBenoit May 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This has been added in with functions


// Set the content
string bodyJson = MakeGenerateContentRequest(inputs, chatHistory, tools, toolConfig);
Expand Down Expand Up @@ -181,7 +184,7 @@ private async IAsyncEnumerable<GenerateContentResponse> GenerateContentStreamAsy
Firebase.AI.Internal.HttpHelpers.GetTemplateURL(_firebaseApp, _backend, templateId) + ":templateStreamGenerateContent?alt=sse");

// Set the request headers
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp);
await Firebase.Internal.HttpHelpers.SetRequestHeaders(request, _firebaseApp, limitedUseAppCheckTokens: _useLimitedUseAppCheckTokens);
Comment thread
AustinBenoit marked this conversation as resolved.

// Set the content
string bodyJson = MakeGenerateContentRequest(inputs, chatHistory, tools, toolConfig);
Expand Down
3 changes: 2 additions & 1 deletion firebaseai/testapp/Assets/Firebase/Sample/Firebase.FirebaseAI.TestApp.asmdef
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,8 @@
"Firebase.App.dll",
"Firebase.Platform.dll",
"Google.MiniJson.dll",
"Firebase.TaskExtension.dll"
"Firebase.TaskExtension.dll",
"Firebase.AppCheck.dll"
],
"autoReferenced": true,
"defineConstraints": [],
Expand Down
13 changes: 12 additions & 1 deletion firebaseai/testapp/Assets/Firebase/Sample/FirebaseAI/UIHandler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ namespace Firebase.Sample.FirebaseAI {
using Firebase;
using Firebase.AI;
using Firebase.Extensions;
using Firebase.AppCheck;
using System;
using System.Collections;
using System.Collections.Generic;
Expand All @@ -40,8 +41,18 @@ public class UIHandler : MonoBehaviour {

private DependencyStatus dependencyStatus = DependencyStatus.UnavailableOther;

private string appCheckDebugToken = "REPLACE_WITH_APP_CHECK_TOKEN";

protected void InitializeAppCheck() {
DebugLog("Initializing App Check directly in manual UIHandler");
DebugAppCheckProviderFactory.Instance.SetDebugToken(appCheckDebugToken);
FirebaseAppCheck.SetAppCheckProviderFactory(DebugAppCheckProviderFactory.Instance);
}

protected virtual void Start() {
UIEnabled = true;
InitializeAppCheck();
InitializeFirebase();
}

protected void InitializeFirebase() {
Expand All @@ -65,7 +76,7 @@ private GenerativeModel GetModel() {
? FirebaseAI.Backend.GoogleAI()
: FirebaseAI.Backend.VertexAI();

return FirebaseAI.GetInstance(backend).GetGenerativeModel(ModelName);
return FirebaseAI.GetInstance(backend, useLimitedUseAppCheckTokens: true).GetGenerativeModel(ModelName);
}

// Send a single message to the Generative Model, without any history.
Expand Down
Loading
Loading