From ceb50b9aeeeb07d87375f2f154434a0b412c4c3e Mon Sep 17 00:00:00 2001
From: Cynthia J <cynthiajoan@gmail.com>
Date: Mon, 6 Oct 2025 19:28:35 -0700
Subject: [PATCH] make sure willContinue will only set when the parsed response
 has it.

---
 .../firebase_ai/lib/src/content.dart          | 24 ++++++---
 .../firebase_ai/test/content_test.dart        | 54 +++++++++++++++++++
 2 files changed, 72 insertions(+), 6 deletions(-)

diff --git a/packages/firebase_ai/firebase_ai/lib/src/content.dart b/packages/firebase_ai/firebase_ai/lib/src/content.dart
index bd8501a119e9..fac91c3650a0 100644
--- a/packages/firebase_ai/firebase_ai/lib/src/content.dart
+++ b/packages/firebase_ai/firebase_ai/lib/src/content.dart
@@ -141,22 +141,34 @@ Part parsePart(Object? jsonObject) {
       throw unhandledFormat('codeExecutionResult', codeExecutionResult);
     }
   }
+
+  if (jsonObject.containsKey('inlineData')) {
+    final inlineDataResult = jsonObject['inlineData'];
+    if (inlineDataResult is Map &&
+        inlineDataResult.containsKey('mimeType') &&
+        inlineDataResult.containsKey('data')) {
+      return InlineDataPart._(
+        inlineDataResult['mimeType'] as String,
+        base64Decode(inlineDataResult['data'] as String),
+        willContinue: inlineDataResult['willContinue'] as bool?,
+        isThought: isThought,
+        thoughtSignature: thoughtSignature,
+      );
+    } else {
+      throw unhandledFormat('inlineData', inlineDataResult);
+    }
+  }
   return switch (jsonObject) {
     {'text': final String text} => TextPart._(text,
         isThought: isThought, thoughtSignature: thoughtSignature),
     {
       'file_data': {
         'file_uri': final String fileUri,
-        'mime_type': final String mimeType
+        'mime_type': final String mimeType,
       }
     } =>
       FileData._(mimeType, fileUri,
           isThought: isThought, thoughtSignature: thoughtSignature),
-    {'inlineData': {'mimeType': String mimeType, 'data': String bytes}} =>
-      InlineDataPart._(mimeType, base64Decode(bytes),
-          willContinue: false,
-          isThought: isThought,
-          thoughtSignature: thoughtSignature),
     _ => () {
         log('unhandled part format: $jsonObject');
         return UnknownPart(jsonObject);
diff --git a/packages/firebase_ai/firebase_ai/test/content_test.dart b/packages/firebase_ai/firebase_ai/test/content_test.dart
index 7b949f2fc429..6b7255f0d5db 100644
--- a/packages/firebase_ai/firebase_ai/test/content_test.dart
+++ b/packages/firebase_ai/firebase_ai/test/content_test.dart
@@ -93,6 +93,27 @@ void main() {
       final json = part.toJson();
       expect((json as Map)['inlineData']['mimeType'], 'image/png');
       expect(json['inlineData']['data'], '');
+      expect(json['inlineData'].containsKey('willContinue'), false);
+    });
+
+    test('DataPart with false willContinue toJson', () {
+      final part =
+          InlineDataPart('image/png', Uint8List(0), willContinue: false);
+      final json = part.toJson();
+      expect((json as Map)['inlineData']['mimeType'], 'image/png');
+      expect(json['inlineData']['data'], '');
+      expect(json['inlineData'].containsKey('willContinue'), true);
+      expect(json['inlineData']['willContinue'], false);
+    });
+
+    test('DataPart with true willContinue toJson', () {
+      final part =
+          InlineDataPart('image/png', Uint8List(0), willContinue: true);
+      final json = part.toJson();
+      expect((json as Map)['inlineData']['mimeType'], 'image/png');
+      expect(json['inlineData']['data'], '');
+      expect(json['inlineData'].containsKey('willContinue'), true);
+      expect(json['inlineData']['willContinue'], true);
     });
 
     test('FunctionCall toJson', () {
@@ -178,6 +199,38 @@ void main() {
     });
 
     test('parses InlineDataPart correctly', () {
+      final json = {
+        'inlineData': {
+          'mimeType': 'image/png',
+          'data': base64Encode([1, 2, 3]),
+          'willContinue': true
+        }
+      };
+      final result = parsePart(json);
+      expect(result, isA<InlineDataPart>());
+      final inlineData = result as InlineDataPart;
+      expect(inlineData.mimeType, 'image/png');
+      expect(inlineData.bytes, [1, 2, 3]);
+      expect(inlineData.willContinue, true);
+    });
+
+    test('parses InlineDataPart with false willContinue', () {
+      final json = {
+        'inlineData': {
+          'mimeType': 'image/png',
+          'data': base64Encode([1, 2, 3]),
+          'willContinue': false
+        }
+      };
+      final result = parsePart(json);
+      expect(result, isA<InlineDataPart>());
+      final inlineData = result as InlineDataPart;
+      expect(inlineData.mimeType, 'image/png');
+      expect(inlineData.bytes, [1, 2, 3]);
+      expect(inlineData.willContinue, false);
+    });
+
+    test('parses InlineDataPart without willContinue', () {
       final json = {
         'inlineData': {
           'mimeType': 'image/png',
@@ -189,6 +242,7 @@ void main() {
       final inlineData = result as InlineDataPart;
       expect(inlineData.mimeType, 'image/png');
       expect(inlineData.bytes, [1, 2, 3]);
+      expect(inlineData.willContinue, null);
     });
 
     test('returns UnknownPart for functionResponse', () {