diff --git a/packages/firebase_ai/firebase_ai/lib/src/live_api.dart b/packages/firebase_ai/firebase_ai/lib/src/live_api.dart index 83c105b708f2..86817d61a593 100644 --- a/packages/firebase_ai/firebase_ai/lib/src/live_api.dart +++ b/packages/firebase_ai/firebase_ai/lib/src/live_api.dart @@ -209,6 +209,15 @@ class LiveServerToolCallCancellation implements LiveServerMessage { final List? functionIds; } +/// A message indicating that the live server is going away. +class LiveServerGoAway implements LiveServerMessage { + /// Creates a [LiveServerGoAway] instance. + const LiveServerGoAway({this.timeLeft}); + + /// The time left in seconds for the live session to be terminated. + final String? timeLeft; +} + /// A single response chunk received during a live content generation. /// /// It can contain generated content, function calls to be executed, or @@ -435,6 +444,9 @@ LiveServerMessage _parseServerMessage(Object jsonObject) { return LiveServerToolCallCancellation(functionIds: toolCancelJson['ids']); } else if (json.containsKey('setupComplete')) { return LiveServerSetupComplete(); + } else if (json.containsKey('goAway')) { + final goAwayJson = json['goAway'] as Map; + return LiveServerGoAway(timeLeft: goAwayJson['timeLeft'] as String?); } else { throw unhandledFormat('LiveServerMessage', json); } diff --git a/packages/firebase_ai/firebase_ai/test/live_test.dart b/packages/firebase_ai/firebase_ai/test/live_test.dart index ea1ebb382b80..03f440f3aba2 100644 --- a/packages/firebase_ai/firebase_ai/test/live_test.dart +++ b/packages/firebase_ai/firebase_ai/test/live_test.dart @@ -228,6 +228,16 @@ void main() { expect(response.message, isA()); }); + test('parseServerMessage parses goAway message correctly', () { + final jsonObject = { + 'goAway': {'timeLeft': '50s'} + }; + final response = parseServerResponse(jsonObject); + expect(response.message, isA()); + final goAwayMessage = response.message as LiveServerGoAway; + expect(goAwayMessage.timeLeft, '50s'); + }); + test('parseServerMessage throws VertexAIException for error message', () { final jsonObject = {'error': {}}; expect(() => parseServerResponse(jsonObject),