feat(ai) Live Service Going Away notifications (#9461)

### Discussion

Support the handling of Going Away notifications in the Live Service. The service may send this event when the session is coming to a close due to timeout or nearing a token limit.

### Testing

Unit tests added.
Manual testing.

### API Changes

See internal [API proposal](https://goto.google.com/fal-firebase-api-live-api-goaway-message).

.changeset/long-monkeys-prove.md

@@ -0,0 +1,7 @@
+---
+'@firebase/ai': minor
+'firebase': minor
+---
+
+AI Logic: Extend the `LiveSession.receive()` return types to include a `LiveServiceGoingAwayNotice`
+which signifies that the service will soon close the session.

common/api-review/ai.api.md

@@ -990,6 +990,7 @@ export const LiveResponseType: {
     SERVER_CONTENT: string;
     TOOL_CALL: string;
     TOOL_CALL_CANCELLATION: string;
+    GOING_AWAY_NOTICE: string;
 };
 
 // @beta
@@ -1006,6 +1007,13 @@ export interface LiveServerContent {
     type: 'serverContent';
 }
 
+// @beta
+export interface LiveServerGoingAwayNotice {
+    timeLeft: number;
+    // (undocumented)
+    type: 'goingAwayNotice';
+}
+
 // @beta
 export interface LiveServerToolCall {
     functionCalls: FunctionCall[];
@@ -1027,7 +1035,7 @@ export class LiveSession {
     close(): Promise<void>;
     inConversation: boolean;
     isClosed: boolean;
-    receive(): AsyncGenerator<LiveServerContent | LiveServerToolCall | LiveServerToolCallCancellation>;
+    receive(): AsyncGenerator<LiveServerContent | LiveServerToolCall | LiveServerToolCallCancellation | LiveServerGoingAwayNotice>;
     send(request: string | Array<string | Part>, turnComplete?: boolean): Promise<void>;
     sendAudioRealtime(blob: GenerativeContentBlob): Promise<void>;
     sendFunctionResponses(functionResponses: FunctionResponse[]): Promise<void>;

docs-devsite/_toc.yaml

@@ -146,6 +146,8 @@ toc:
     path: /docs/reference/js/ai.livemodelparams.md
   - title: LiveServerContent
     path: /docs/reference/js/ai.liveservercontent.md
+  - title: LiveServerGoingAwayNotice
+    path: /docs/reference/js/ai.liveservergoingawaynotice.md
   - title: LiveServerToolCall
     path: /docs/reference/js/ai.liveservertoolcall.md
   - title: LiveServerToolCallCancellation

docs-devsite/ai.liveservergoingawaynotice.md

@@ -0,0 +1,53 @@
+Project: /docs/reference/js/_project.yaml
+Book: /docs/reference/_book.yaml
+page_type: reference
+
+{% comment %}
+DO NOT EDIT THIS FILE!
+This is generated by the JS SDK team, and any local changes will be
+overwritten. Changes should be made in the source code at
+https://github.com/firebase/firebase-js-sdk
+{% endcomment %}
+
+# LiveServerGoingAwayNotice interface
+> This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
+> 
+
+Notification that the server will not be able to service the client soon.
+
+<b>Signature:</b>
+
+```typescript
+export interface LiveServerGoingAwayNotice 
+```
+
+## Properties
+
+|  Property | Type | Description |
+|  --- | --- | --- |
+|  [timeLeft](./ai.liveservergoingawaynotice.md#liveservergoingawaynoticetimeleft) | number | <b><i>(Public Preview)</i></b> The remaining time (in seconds) before the connection will be terminated. |
+|  [type](./ai.liveservergoingawaynotice.md#liveservergoingawaynoticetype) | 'goingAwayNotice' | <b><i>(Public Preview)</i></b> |
+
+## LiveServerGoingAwayNotice.timeLeft
+
+> This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
+> 
+
+The remaining time (in seconds) before the connection will be terminated.
+
+<b>Signature:</b>
+
+```typescript
+timeLeft: number;
+```
+
+## LiveServerGoingAwayNotice.type
+
+> This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
+> 
+
+<b>Signature:</b>
+
+```typescript
+type: 'goingAwayNotice';
+```

docs-devsite/ai.livesession.md

@@ -98,11 +98,11 @@ Yields messages received from the server. This can only be used by one consumer
 <b>Signature:</b>
 
 ```typescript
-receive(): AsyncGenerator<LiveServerContent | LiveServerToolCall | LiveServerToolCallCancellation>;
+receive(): AsyncGenerator<LiveServerContent | LiveServerToolCall | LiveServerToolCallCancellation | LiveServerGoingAwayNotice>;
 ```
 <b>Returns:</b>
 
-AsyncGenerator&lt;[LiveServerContent](./ai.liveservercontent.md#liveservercontent_interface) \| [LiveServerToolCall](./ai.liveservertoolcall.md#liveservertoolcall_interface) \| [LiveServerToolCallCancellation](./ai.liveservertoolcallcancellation.md#liveservertoolcallcancellation_interface)<!-- -->&gt;
+AsyncGenerator&lt;[LiveServerContent](./ai.liveservercontent.md#liveservercontent_interface) \| [LiveServerToolCall](./ai.liveservertoolcall.md#liveservertoolcall_interface) \| [LiveServerToolCallCancellation](./ai.liveservertoolcallcancellation.md#liveservertoolcallcancellation_interface) \| [LiveServerGoingAwayNotice](./ai.liveservergoingawaynotice.md#liveservergoingawaynotice_interface)<!-- -->&gt;
 
 An `AsyncGenerator` that yields server messages as they arrive.
 

docs-devsite/ai.md

@@ -115,6 +115,7 @@ The Firebase AI Web SDK.
 |  [LiveGenerationConfig](./ai.livegenerationconfig.md#livegenerationconfig_interface) | <b><i>(Public Preview)</i></b> Configuration parameters used by [LiveGenerativeModel](./ai.livegenerativemodel.md#livegenerativemodel_class) to control live content generation. |
 |  [LiveModelParams](./ai.livemodelparams.md#livemodelparams_interface) | <b><i>(Public Preview)</i></b> Params passed to [getLiveGenerativeModel()](./ai.md#getlivegenerativemodel_f2099ac)<!-- -->. |
 |  [LiveServerContent](./ai.liveservercontent.md#liveservercontent_interface) | <b><i>(Public Preview)</i></b> An incremental content update from the model. |
+|  [LiveServerGoingAwayNotice](./ai.liveservergoingawaynotice.md#liveservergoingawaynotice_interface) | <b><i>(Public Preview)</i></b> Notification that the server will not be able to service the client soon. |
 |  [LiveServerToolCall](./ai.liveservertoolcall.md#liveservertoolcall_interface) | <b><i>(Public Preview)</i></b> A request from the model for the client to execute one or more functions. |
 |  [LiveServerToolCallCancellation](./ai.liveservertoolcallcancellation.md#liveservertoolcallcancellation_interface) | <b><i>(Public Preview)</i></b> Notification to cancel a previous function call triggered by [LiveServerToolCall](./ai.liveservertoolcall.md#liveservertoolcall_interface)<!-- -->. |
 |  [ModalityTokenCount](./ai.modalitytokencount.md#modalitytokencount_interface) | Represents token counting info for a single modality. |
@@ -748,6 +749,7 @@ LiveResponseType: {
     SERVER_CONTENT: string;
     TOOL_CALL: string;
     TOOL_CALL_CANCELLATION: string;
+    GOING_AWAY_NOTICE: string;
 }
 ```
 

packages/ai/integration/live.test.ts

@@ -22,6 +22,7 @@ import {
   LiveServerContent,
   LiveServerToolCall,
   LiveServerToolCallCancellation,
+  LiveServerGoingAwayNotice,
   ResponseModality
 } from '../src';
 import { liveTestConfigs } from './constants';
@@ -30,7 +31,10 @@ import { HELLO_AUDIO_PCM_BASE64 } from './sample-data/hello-audio';
 // A helper function to consume the generator and collect text parts from one turn.
 async function nextTurnData(
   stream: AsyncGenerator<
-    LiveServerContent | LiveServerToolCall | LiveServerToolCallCancellation
+    | LiveServerContent
+    | LiveServerToolCall
+    | LiveServerToolCallCancellation
+    | LiveServerGoingAwayNotice
   >
 ): Promise<{
   text: string;
@@ -48,7 +52,8 @@ async function nextTurnData(
     const chunk = result.value as
       | LiveServerContent
       | LiveServerToolCall
-      | LiveServerToolCallCancellation;
+      | LiveServerToolCallCancellation
+      | LiveServerGoingAwayNotice;
     switch (chunk.type) {
       case 'serverContent':
         if (chunk.turnComplete) {
@@ -81,6 +86,10 @@ async function nextTurnData(
           });
         }
         break;
+      case 'goingAwayNotice':
+        // Ignore going away notices for now. It can take 15 minutes before the service
+        // generates one, which is much too long to test.
+        break;
       default:
         throw new Error(`Unexpected chunk type '${(chunk as any).type}'`);
     }
@@ -323,7 +332,8 @@ describe('Live', function () {
             const chunk = result.value as
               | LiveServerContent
               | LiveServerToolCall
-              | LiveServerToolCallCancellation;
+              | LiveServerToolCallCancellation
+              | LiveServerGoingAwayNotice;
             if (chunk.type === 'serverContent') {
               if (chunk.turnComplete) {
                 break;

packages/ai/src/methods/live-session.test.ts

@@ -23,6 +23,7 @@ import {
   FunctionResponse,
   LiveResponseType,
   LiveServerContent,
+  LiveServerGoingAwayNotice,
   LiveServerToolCall,
   LiveServerToolCallCancellation
 } from '../types';
@@ -238,6 +239,9 @@ describe('LiveSession', () => {
       mockHandler.simulateServerMessage({
         toolCallCancellation: { functionIds: ['123'] }
       });
+      mockHandler.simulateServerMessage({
+        goAway: { timeLeft: '30s' }
+      });
       mockHandler.simulateServerMessage({
         serverContent: { turnComplete: true }
       });
@@ -245,7 +249,7 @@ describe('LiveSession', () => {
       mockHandler.endStream();
 
       const responses = await receivePromise;
-      expect(responses).to.have.lengthOf(4);
+      expect(responses).to.have.lengthOf(5);
       expect(responses[0]).to.deep.equal({
         type: LiveResponseType.SERVER_CONTENT,
         modelTurn: { parts: [{ text: 'response 1' }] }
@@ -258,6 +262,60 @@ describe('LiveSession', () => {
         type: LiveResponseType.TOOL_CALL_CANCELLATION,
         functionIds: ['123']
       } as LiveServerToolCallCancellation);
+      expect(responses[3]).to.deep.equal({
+        type: LiveResponseType.GOING_AWAY_NOTICE,
+        timeLeft: 30
+      } as LiveServerGoingAwayNotice);
+      expect(responses[4]).to.deep.equal({
+        type: LiveResponseType.SERVER_CONTENT,
+        turnComplete: true
+      } as LiveServerContent);
+    });
+
+    it('should correctly parse high precision duration in LiveServerGoingAwayNotice', async () => {
+      const receivePromise = (async () => {
+        const responses = [];
+        for await (const response of session.receive()) {
+          responses.push(response);
+        }
+        return responses;
+      })();
+
+      mockHandler.simulateServerMessage({
+        goAway: { timeLeft: '3.000000001s' }
+      });
+      await new Promise<void>(r => setTimeout(() => r(), 10));
+      mockHandler.endStream();
+
+      const responses = await receivePromise;
+      expect(responses).to.have.lengthOf(1);
+      expect(responses[0]).to.deep.equal({
+        type: LiveResponseType.GOING_AWAY_NOTICE,
+        timeLeft: 3.000000001
+      } as LiveServerGoingAwayNotice);
+    });
+
+    it('should default timeLeft to 0 if format is invalid', async () => {
+      const receivePromise = (async () => {
+        const responses = [];
+        for await (const response of session.receive()) {
+          responses.push(response);
+        }
+        return responses;
+      })();
+
+      mockHandler.simulateServerMessage({
+        goAway: { timeLeft: 'invalid' }
+      });
+      await new Promise<void>(r => setTimeout(() => r(), 10));
+      mockHandler.endStream();
+
+      const responses = await receivePromise;
+      expect(responses).to.have.lengthOf(1);
+      expect(responses[0]).to.deep.equal({
+        type: LiveResponseType.GOING_AWAY_NOTICE,
+        timeLeft: 0
+      } as LiveServerGoingAwayNotice);
     });
 
     it('should log a warning and skip messages that are not objects', async () => {

packages/ai/src/methods/live-session.ts

@@ -21,6 +21,7 @@ import {
   GenerativeContentBlob,
   LiveResponseType,
   LiveServerContent,
+  LiveServerGoingAwayNotice,
   LiveServerToolCall,
   LiveServerToolCallCancellation,
   Part
@@ -227,7 +228,10 @@ export class LiveSession {
    * @beta
    */
   async *receive(): AsyncGenerator<
-    LiveServerContent | LiveServerToolCall | LiveServerToolCallCancellation
+    | LiveServerContent
+    | LiveServerToolCall
+    | LiveServerToolCallCancellation
+    | LiveServerGoingAwayNotice
   > {
     if (this.isClosed) {
       throw new AIError(
@@ -261,6 +265,16 @@ export class LiveSession {
               }
             ).toolCallCancellation
           } as LiveServerToolCallCancellation;
+        } else if ('goAway' in message) {
+          const notice = (
+            message as {
+              goAway: { timeLeft: string };
+            }
+          ).goAway;
+          yield {
+            type: LiveResponseType.GOING_AWAY_NOTICE,
+            timeLeft: parseDuration(notice.timeLeft)
+          } as LiveServerGoingAwayNotice;
         } else {
           logger.warn(
             `Received an unknown message type from the server: ${JSON.stringify(
@@ -360,3 +374,16 @@ export class LiveSession {
     }
   }
 }
+
+/**
+ * Parses a duration string (e.g. "3.000000001s") into a number of seconds.
+ *
+ * @param duration - The duration string to parse.
+ * @returns The duration in seconds.
+ */
+function parseDuration(duration: string | undefined): number {
+  if (!duration || !duration.endsWith('s')) {
+    return 0;
+  }
+  return Number(duration.slice(0, -1)); // slice removes the trailing 's'.
+}

packages/ai/src/types/responses.ts

@@ -616,6 +616,19 @@ export interface LiveServerToolCallCancellation {
   functionIds: string[];
 }
 
+/**
+ * Notification that the server will not be able to service the client soon.
+ *
+ * @beta
+ */
+export interface LiveServerGoingAwayNotice {
+  type: 'goingAwayNotice';
+  /**
+   * The remaining time (in seconds) before the connection will be terminated.
+   */
+  timeLeft: number;
+}
+
 /**
  * The types of responses that can be returned by {@link LiveSession.receive}.
  *
@@ -624,7 +637,8 @@ export interface LiveServerToolCallCancellation {
 export const LiveResponseType = {
   SERVER_CONTENT: 'serverContent',
   TOOL_CALL: 'toolCall',
-  TOOL_CALL_CANCELLATION: 'toolCallCancellation'
+  TOOL_CALL_CANCELLATION: 'toolCallCancellation',
+  GOING_AWAY_NOTICE: 'goingAwayNotice'
 };
 
 /**