Skip to content

Commit f2a682a

Browse files
cynthiajoancynthiajoan
authored
feat(firebaseai): Add app check limited use token (#17645)
* Add app check limited use option * fix analyzer * fix format * fix format * fix limit token test
1 parent 537a3c3 commit f2a682a

File tree

7 files changed

+100
-20
lines changed

7 files changed

+100
-20
lines changed

‎packages/firebase_ai/firebase_ai/lib/src/base_model.dart‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -169,14 +169,20 @@ abstract class BaseModel {
169169

170170
/// Returns a function that generates Firebase auth tokens.
171171
static FutureOr<Map<String, String>> Function() firebaseTokens(
172-
FirebaseAppCheck? appCheck, FirebaseAuth? auth, FirebaseApp? app) {
172+
FirebaseAppCheck? appCheck,
173+
FirebaseAuth? auth,
174+
FirebaseApp? app,
175+
bool? useLimitedUseAppCheckTokens,
176+
) {
173177
return () async {
174178
Map<String, String> headers = {};
175179
// Override the client name in Google AI SDK
176180
headers['x-goog-api-client'] =
177181
'gl-dart/$packageVersion fire/$packageVersion';
178182
if (appCheck != null) {
179-
final appCheckToken = await appCheck.getToken();
183+
final appCheckToken = useLimitedUseAppCheckTokens == true
184+
? await appCheck.getLimitedUseToken()
185+
: await appCheck.getToken();
180186
if (appCheckToken != null) {
181187
headers['X-Firebase-AppCheck'] = appCheckToken;
182188
}

‎packages/firebase_ai/firebase_ai/lib/src/firebase_ai.dart‎

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,13 +26,14 @@ const _defaultLocation = 'us-central1';
2626

2727
/// The entrypoint for generative models.
2828
class FirebaseAI extends FirebasePluginPlatform {
29-
FirebaseAI._(
30-
{required this.app,
31-
required this.location,
32-
required bool useVertexBackend,
33-
this.appCheck,
34-
this.auth})
35-
: _useVertexBackend = useVertexBackend,
29+
FirebaseAI._({
30+
required this.app,
31+
required this.location,
32+
required bool useVertexBackend,
33+
this.appCheck,
34+
this.auth,
35+
this.useLimitedUseAppCheckTokens = false,
36+
}) : _useVertexBackend = useVertexBackend,
3637
super(app.name, 'plugins.flutter.io/firebase_vertexai');
3738

3839
/// The [FirebaseApp] for this current [FirebaseAI] instance.
@@ -48,6 +49,9 @@ class FirebaseAI extends FirebasePluginPlatform {
4849
/// The service location for this [FirebaseAI] instance.
4950
String location;
5051

52+
/// Whether to use App Check limited use tokens. Defaults to false.
53+
final bool useLimitedUseAppCheckTokens;
54+
5155
final bool _useVertexBackend;
5256

5357
static final Map<String, FirebaseAI> _cachedInstances = {};
@@ -61,6 +65,7 @@ class FirebaseAI extends FirebasePluginPlatform {
6165
FirebaseAppCheck? appCheck,
6266
FirebaseAuth? auth,
6367
String? location,
68+
bool? useLimitedUseAppCheckTokens,
6469
}) {
6570
app ??= Firebase.app();
6671
var instanceKey = '${app.name}::vertexai';
@@ -77,6 +82,7 @@ class FirebaseAI extends FirebasePluginPlatform {
7782
appCheck: appCheck,
7883
auth: auth,
7984
useVertexBackend: true,
85+
useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens ?? false,
8086
);
8187
_cachedInstances[instanceKey] = newInstance;
8288

@@ -91,6 +97,7 @@ class FirebaseAI extends FirebasePluginPlatform {
9197
FirebaseApp? app,
9298
FirebaseAppCheck? appCheck,
9399
FirebaseAuth? auth,
100+
bool? useLimitedUseAppCheckTokens,
94101
}) {
95102
app ??= Firebase.app();
96103
var instanceKey = '${app.name}::googleai';
@@ -105,6 +112,7 @@ class FirebaseAI extends FirebasePluginPlatform {
105112
appCheck: appCheck,
106113
auth: auth,
107114
useVertexBackend: false,
115+
useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens ?? false,
108116
);
109117
_cachedInstances[instanceKey] = newInstance;
110118

@@ -142,6 +150,7 @@ class FirebaseAI extends FirebasePluginPlatform {
142150
tools: tools,
143151
toolConfig: toolConfig,
144152
systemInstruction: systemInstruction,
153+
useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens,
145154
);
146155
}
147156

@@ -162,7 +171,8 @@ class FirebaseAI extends FirebasePluginPlatform {
162171
generationConfig: generationConfig,
163172
safetySettings: safetySettings,
164173
appCheck: appCheck,
165-
auth: auth);
174+
auth: auth,
175+
useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens);
166176
}
167177

168178
/// Create a [LiveGenerativeModel] for real-time interaction.
@@ -185,6 +195,7 @@ class FirebaseAI extends FirebasePluginPlatform {
185195
systemInstruction: systemInstruction,
186196
appCheck: appCheck,
187197
auth: auth,
198+
useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens,
188199
);
189200
}
190201
}

‎packages/firebase_ai/firebase_ai/lib/src/generative_model.dart‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ final class GenerativeModel extends BaseApiClientModel {
3737
required String location,
3838
required FirebaseApp app,
3939
required bool useVertexBackend,
40+
bool? useLimitedUseAppCheckTokens,
4041
FirebaseAppCheck? appCheck,
4142
FirebaseAuth? auth,
4243
List<SafetySetting>? safetySettings,
@@ -60,13 +61,15 @@ final class GenerativeModel extends BaseApiClientModel {
6061
client: HttpApiClient(
6162
apiKey: app.options.apiKey,
6263
httpClient: httpClient,
63-
requestHeaders: BaseModel.firebaseTokens(appCheck, auth, app)));
64+
requestHeaders: BaseModel.firebaseTokens(
65+
appCheck, auth, app, useLimitedUseAppCheckTokens)));
6466

6567
GenerativeModel._constructTestModel({
6668
required String model,
6769
required String location,
6870
required FirebaseApp app,
6971
required useVertexBackend,
72+
bool? useLimitedUseAppCheckTokens,
7073
FirebaseAppCheck? appCheck,
7174
FirebaseAuth? auth,
7275
List<SafetySetting>? safetySettings,
@@ -90,8 +93,8 @@ final class GenerativeModel extends BaseApiClientModel {
9093
client: apiClient ??
9194
HttpApiClient(
9295
apiKey: app.options.apiKey,
93-
requestHeaders:
94-
BaseModel.firebaseTokens(appCheck, auth, app)));
96+
requestHeaders: BaseModel.firebaseTokens(
97+
appCheck, auth, app, useLimitedUseAppCheckTokens)));
9598

9699
final List<SafetySetting> _safetySettings;
97100
final GenerationConfig? _generationConfig;
@@ -199,6 +202,7 @@ GenerativeModel createGenerativeModel({
199202
required String location,
200203
required String model,
201204
required bool useVertexBackend,
205+
bool? useLimitedUseAppCheckTokens,
202206
FirebaseAppCheck? appCheck,
203207
FirebaseAuth? auth,
204208
GenerationConfig? generationConfig,
@@ -212,6 +216,7 @@ GenerativeModel createGenerativeModel({
212216
app: app,
213217
appCheck: appCheck,
214218
useVertexBackend: useVertexBackend,
219+
useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens,
215220
auth: auth,
216221
location: location,
217222
safetySettings: safetySettings,
@@ -230,6 +235,7 @@ GenerativeModel createModelWithClient({
230235
required String model,
231236
required ApiClient client,
232237
required bool useVertexBackend,
238+
bool? useLimitedUseAppCheckTokens,
233239
Content? systemInstruction,
234240
FirebaseAppCheck? appCheck,
235241
FirebaseAuth? auth,
@@ -243,6 +249,7 @@ GenerativeModel createModelWithClient({
243249
app: app,
244250
appCheck: appCheck,
245251
useVertexBackend: useVertexBackend,
252+
useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens,
246253
auth: auth,
247254
location: location,
248255
safetySettings: safetySettings,

‎packages/firebase_ai/firebase_ai/lib/src/imagen/imagen_model.dart‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ final class ImagenModel extends BaseApiClientModel {
3131
required String model,
3232
required String location,
3333
required bool useVertexBackend,
34+
bool? useLimitedUseAppCheckTokens,
3435
FirebaseAppCheck? appCheck,
3536
FirebaseAuth? auth,
3637
ImagenGenerationConfig? generationConfig,
@@ -45,7 +46,8 @@ final class ImagenModel extends BaseApiClientModel {
4546
: _GoogleAIUri(app: app, model: model),
4647
client: HttpApiClient(
4748
apiKey: app.options.apiKey,
48-
requestHeaders: BaseModel.firebaseTokens(appCheck, auth, app)));
49+
requestHeaders: BaseModel.firebaseTokens(
50+
appCheck, auth, app, useLimitedUseAppCheckTokens)));
4951

5052
final ImagenGenerationConfig? _generationConfig;
5153
final ImagenSafetySettings? _safetySettings;
@@ -198,6 +200,7 @@ ImagenModel createImagenModel({
198200
required String location,
199201
required String model,
200202
required bool useVertexBackend,
203+
bool? useLimitedUseAppCheckTokens,
201204
FirebaseAppCheck? appCheck,
202205
FirebaseAuth? auth,
203206
ImagenGenerationConfig? generationConfig,
@@ -210,6 +213,7 @@ ImagenModel createImagenModel({
210213
auth: auth,
211214
location: location,
212215
useVertexBackend: useVertexBackend,
216+
useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens,
213217
safetySettings: safetySettings,
214218
generationConfig: generationConfig,
215219
);

‎packages/firebase_ai/firebase_ai/lib/src/live_model.dart‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ final class LiveGenerativeModel extends BaseModel {
3434
required String location,
3535
required FirebaseApp app,
3636
required bool useVertexBackend,
37+
bool? useLimitedUseAppCheckTokens,
3738
FirebaseAppCheck? appCheck,
3839
FirebaseAuth? auth,
3940
LiveGenerationConfig? liveGenerationConfig,
@@ -47,6 +48,7 @@ final class LiveGenerativeModel extends BaseModel {
4748
_liveGenerationConfig = liveGenerationConfig,
4849
_tools = tools,
4950
_systemInstruction = systemInstruction,
51+
_useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens,
5052
super._(
5153
serializationStrategy: VertexSerialization(),
5254
modelUri: useVertexBackend
@@ -69,6 +71,7 @@ final class LiveGenerativeModel extends BaseModel {
6971
final LiveGenerationConfig? _liveGenerationConfig;
7072
final List<Tool>? _tools;
7173
final Content? _systemInstruction;
74+
final bool? _useLimitedUseAppCheckTokens;
7275

7376
String _vertexAIUri() => 'wss://${_modelUri.baseAuthority}/'
7477
'$_apiUrl.${_modelUri.apiVersion}.$_apiUrlSuffixVertexAI/'
@@ -107,7 +110,12 @@ final class LiveGenerativeModel extends BaseModel {
107110
};
108111

109112
final request = jsonEncode(setupJson);
110-
final headers = await BaseModel.firebaseTokens(_appCheck, _auth, _app)();
113+
final headers = await BaseModel.firebaseTokens(
114+
_appCheck,
115+
_auth,
116+
_app,
117+
_useLimitedUseAppCheckTokens,
118+
)();
111119

112120
var ws = kIsWeb
113121
? WebSocketChannel.connect(Uri.parse(uri))
@@ -126,6 +134,7 @@ LiveGenerativeModel createLiveGenerativeModel({
126134
required String location,
127135
required String model,
128136
required bool useVertexBackend,
137+
bool? useLimitedUseAppCheckTokens,
129138
FirebaseAppCheck? appCheck,
130139
FirebaseAuth? auth,
131140
LiveGenerationConfig? liveGenerationConfig,
@@ -139,6 +148,7 @@ LiveGenerativeModel createLiveGenerativeModel({
139148
auth: auth,
140149
location: location,
141150
useVertexBackend: useVertexBackend,
151+
useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens,
142152
liveGenerationConfig: liveGenerationConfig,
143153
tools: tools,
144154
systemInstruction: systemInstruction,

‎packages/firebase_ai/firebase_ai/test/base_model_test.dart‎

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,10 @@ class MockFirebaseAppCheck extends Mock implements FirebaseAppCheck {
4545
@override
4646
Future<String?> getToken([bool? forceRefresh = false]) async =>
4747
super.noSuchMethod(Invocation.method(#getToken, [forceRefresh]));
48+
49+
@override
50+
Future<String> getLimitedUseToken() async =>
51+
super.noSuchMethod(Invocation.method(#getLimitedUseToken, [])) ?? '';
4852
}
4953

5054
// Mock Firebase Auth
@@ -72,7 +76,7 @@ class MockApiClient extends Mock implements ApiClient {
7276
void main() {
7377
group('BaseModel', () {
7478
test('firebaseTokens returns a function that generates headers', () async {
75-
final tokenFunction = BaseModel.firebaseTokens(null, null, null);
79+
final tokenFunction = BaseModel.firebaseTokens(null, null, null, false);
7680
final headers = await tokenFunction();
7781
expect(headers['x-goog-api-client'], contains('gl-dart'));
7882
expect(headers['x-goog-api-client'], contains('fire'));
@@ -83,7 +87,8 @@ void main() {
8387
final mockAppCheck = MockFirebaseAppCheck();
8488
when(mockAppCheck.getToken())
8589
.thenAnswer((_) async => 'test-app-check-token');
86-
final tokenFunction = BaseModel.firebaseTokens(mockAppCheck, null, null);
90+
final tokenFunction =
91+
BaseModel.firebaseTokens(mockAppCheck, null, null, false);
8792
final headers = await tokenFunction();
8893
expect(headers['X-Firebase-AppCheck'], 'test-app-check-token');
8994
expect(headers['x-goog-api-client'], contains('gl-dart'));
@@ -96,7 +101,8 @@ void main() {
96101
final mockUser = MockUser();
97102
when(mockUser.getIdToken()).thenAnswer((_) async => 'test-id-token');
98103
when(mockAuth.currentUser).thenReturn(mockUser);
99-
final tokenFunction = BaseModel.firebaseTokens(null, mockAuth, null);
104+
final tokenFunction =
105+
BaseModel.firebaseTokens(null, mockAuth, null, false);
100106
final headers = await tokenFunction();
101107
expect(headers['Authorization'], 'Firebase test-id-token');
102108
expect(headers['x-goog-api-client'], contains('gl-dart'));
@@ -109,7 +115,8 @@ void main() {
109115
() async {
110116
final mockApp = MockFirebaseApp();
111117

112-
final tokenFunction = BaseModel.firebaseTokens(null, null, mockApp);
118+
final tokenFunction =
119+
BaseModel.firebaseTokens(null, null, mockApp, false);
113120
final headers = await tokenFunction();
114121
expect(headers['X-Firebase-AppId'], 'test-app-id');
115122
expect(headers['x-goog-api-client'], contains('gl-dart'));
@@ -128,7 +135,7 @@ void main() {
128135
final mockApp = MockFirebaseApp();
129136

130137
final tokenFunction =
131-
BaseModel.firebaseTokens(mockAppCheck, mockAuth, mockApp);
138+
BaseModel.firebaseTokens(mockAppCheck, mockAuth, mockApp, false);
132139
final headers = await tokenFunction();
133140
expect(headers['X-Firebase-AppCheck'], 'test-app-check-token');
134141
expect(headers['Authorization'], 'Firebase test-id-token');
@@ -137,5 +144,20 @@ void main() {
137144
expect(headers['x-goog-api-client'], contains('fire'));
138145
expect(headers.length, 4);
139146
});
147+
148+
test('firebaseTokens includes limited use App Check token if specified',
149+
() async {
150+
final mockAppCheck = MockFirebaseAppCheck();
151+
when(mockAppCheck.getLimitedUseToken())
152+
.thenAnswer((_) async => 'test-limited-use-app-check-token');
153+
final tokenFunction =
154+
BaseModel.firebaseTokens(mockAppCheck, null, null, true);
155+
final headers = await tokenFunction();
156+
expect(
157+
headers['X-Firebase-AppCheck'], 'test-limited-use-app-check-token');
158+
expect(headers['x-goog-api-client'], contains('gl-dart'));
159+
expect(headers['x-goog-api-client'], contains('fire'));
160+
expect(headers.length, 2);
161+
});
140162
});
141163
}

‎packages/firebase_ai/firebase_ai/test/firebase_vertexai_test.dart‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,9 @@ void main() {
2626
// ignore: unused_local_variable
2727
late FirebaseAppCheck appCheck;
2828
late FirebaseApp customApp;
29+
late FirebaseApp limitTokenApp;
2930
late FirebaseAppCheck customAppCheck;
31+
late FirebaseAppCheck limitTokenAppCheck;
3032

3133
group('FirebaseAI Tests', () {
3234
late FirebaseApp app;
@@ -38,8 +40,13 @@ void main() {
3840
name: 'custom-app',
3941
options: Firebase.app().options,
4042
);
43+
limitTokenApp = await Firebase.initializeApp(
44+
name: 'limit-token-app',
45+
options: Firebase.app().options,
46+
);
4147
appCheck = FirebaseAppCheck.instance;
4248
customAppCheck = FirebaseAppCheck.instanceFor(app: customApp);
49+
limitTokenAppCheck = FirebaseAppCheck.instanceFor(app: limitTokenApp);
4350
});
4451

4552
test('Singleton behavior', () {
@@ -76,6 +83,19 @@ void main() {
7683
expect(model, isA<GenerativeModel>());
7784
});
7885

86+
test('Instance creation with useLimitedUseAppCheckTokens', () {
87+
final vertexAIAppCheck = FirebaseAI.vertexAI(
88+
app: limitTokenApp,
89+
appCheck: limitTokenAppCheck,
90+
location: 'limit-token-location',
91+
useLimitedUseAppCheckTokens: true,
92+
);
93+
expect(vertexAIAppCheck.app, equals(limitTokenApp));
94+
expect(vertexAIAppCheck.appCheck, equals(limitTokenAppCheck));
95+
expect(vertexAIAppCheck.location, equals('limit-token-location'));
96+
expect(vertexAIAppCheck.useLimitedUseAppCheckTokens, true);
97+
});
98+
7999
// ... other tests (e.g., with different parameters)
80100
});
81101
}

0 commit comments

Comments
 (0)