feat: Implement session resumption and context window compression for live sessions.

cynthiajoan · cynthiajoan · commit 19c9b608198b · 2026-03-04T22:57:41.000-08:00
diff --git a/packages/ai/src/methods/live-session.test.ts b/packages/ai/src/methods/live-session.test.ts
@@ -25,7 +25,8 @@ import {
   LiveServerContent,
   LiveServerGoingAwayNotice,
   LiveServerToolCall,
-  LiveServerToolCallCancellation
+  LiveServerToolCallCancellation,
+  LiveSessionResumptionUpdate
 } from '../types';
 import { LiveSession } from './live-session';
 import { WebSocketHandler } from '../websocket';
@@ -83,7 +84,15 @@ describe('LiveSession', () => {
   beforeEach(() => {
     mockHandler = new MockWebSocketHandler();
     serverMessagesGenerator = mockHandler.listen();
-    session = new LiveSession(mockHandler, serverMessagesGenerator);
+    session = new LiveSession(
+      mockHandler,
+      serverMessagesGenerator,
+      async (resumptionConfig) => {
+        // mock reconnector that replaces the handler
+        mockHandler = new MockWebSocketHandler();
+        return mockHandler.listen();
+      }
+    );
   });
 
   describe('send()', () => {
@@ -220,6 +229,27 @@ describe('LiveSession', () => {
     });
   });
 
+  describe('resumeSession()', () => {
+    it('should close existing session and start a new one using reconnector', async () => {
+      expect(session.isClosed).to.be.false;
+
+      const oldServerMessages = (session as any).serverMessages;
+      await session.resumeSession({ handle: 'testHandle' });
+      
+      expect(mockHandler.close).to.have.been.calledOnce;
+      expect(session.isClosed).to.be.false;
+      expect((session as any).serverMessages).to.not.equal(oldServerMessages);
+    });
+
+    it('should throw if reconnector is not provided', async () => {
+      const basicSession = new LiveSession(mockHandler, serverMessagesGenerator);
+      await expect(basicSession.resumeSession()).to.be.rejectedWith(
+        AIError,
+        /resumeSession is not supported on this session/
+      );
+    });
+  });
+
   describe('receive()', () => {
     it('should correctly parse and transform all server message types', async () => {
       const receivePromise = (async () => {
@@ -242,14 +272,17 @@ describe('LiveSession', () => {
       mockHandler.simulateServerMessage({
         goAway: { timeLeft: '30s' }
       });
+      mockHandler.simulateServerMessage({
+        sessionResumptionUpdate: { newHandle: 'test', resumable: true, lastConsumedClientMessageIndex: 5 }
+      });
       mockHandler.simulateServerMessage({
         serverContent: { turnComplete: true }
       });
       await new Promise<void>(r => setTimeout(() => r(), 10)); // Wait for the listener to process messages
       mockHandler.endStream();
 
       const responses = await receivePromise;
-      expect(responses).to.have.lengthOf(5);
+      expect(responses).to.have.lengthOf(6);
       expect(responses[0]).to.deep.equal({
         type: LiveResponseType.SERVER_CONTENT,
         modelTurn: { parts: [{ text: 'response 1' }] }
@@ -267,6 +300,12 @@ describe('LiveSession', () => {
         timeLeft: 30
       } as LiveServerGoingAwayNotice);
       expect(responses[4]).to.deep.equal({
+        type: LiveResponseType.SESSION_RESUMPTION_UPDATE,
+        newHandle: 'test',
+        resumable: true,
+        lastConsumedClientMessageIndex: 5
+      } as LiveSessionResumptionUpdate);
+      expect(responses[5]).to.deep.equal({
         type: LiveResponseType.SERVER_CONTENT,
         turnComplete: true
       } as LiveServerContent);
diff --git a/packages/ai/src/methods/live-session.ts b/packages/ai/src/methods/live-session.ts
@@ -24,7 +24,9 @@ import {
   LiveServerGoingAwayNotice,
   LiveServerToolCall,
   LiveServerToolCallCancellation,
-  Part
+  LiveSessionResumptionUpdate,
+  Part,
+  SessionResumptionConfig
 } from '../public-types';
 import { formatNewContent } from '../requests/request-helpers';
 import { AIError } from '../errors';
@@ -62,9 +64,36 @@ export class LiveSession {
    */
   constructor(
     private webSocketHandler: WebSocketHandler,
-    private serverMessages: AsyncGenerator<unknown>
+    private serverMessages: AsyncGenerator<unknown>,
+    private reconnector?: (sessionResumption?: SessionResumptionConfig) => Promise<AsyncGenerator<unknown>>
   ) {}
 
+  /**
+   * Resumes an existing live session with the server.
+   *
+   * This closes the current WebSocket connection and establishes a new one using
+   * the same configuration (URI, headers, model, system instruction, tools, etc.)
+   * as the original session.
+   *
+   * @param sessionResumption - The configuration for session resumption, such as the handle to the previous session state to restore.
+   * @throws If the session resumption configuration is unsupported.
+   *
+   * @beta
+   */
+  async resumeSession(
+    sessionResumption?: SessionResumptionConfig
+  ): Promise<void> {
+    if (!this.reconnector) {
+      throw new AIError(
+        AIErrorCode.UNSUPPORTED,
+        'resumeSession is not supported on this session.'
+      );
+    }
+    await this.close();
+    this.isClosed = false;
+    this.serverMessages = await this.reconnector(sessionResumption);
+  }
+
   /**
    * Sends content to the server.
    *
@@ -232,6 +261,7 @@ export class LiveSession {
     | LiveServerToolCall
     | LiveServerToolCallCancellation
     | LiveServerGoingAwayNotice
+    | LiveSessionResumptionUpdate
   > {
     if (this.isClosed) {
       throw new AIError(
@@ -275,6 +305,18 @@ export class LiveSession {
             type: LiveResponseType.GOING_AWAY_NOTICE,
             timeLeft: parseDuration(notice.timeLeft)
           } as LiveServerGoingAwayNotice;
+        } else if (LiveResponseType.SESSION_RESUMPTION_UPDATE in message) {
+          yield {
+            type: LiveResponseType.SESSION_RESUMPTION_UPDATE,
+            ...(
+              message as {
+                sessionResumptionUpdate: Omit<
+                  LiveSessionResumptionUpdate,
+                  'type'
+                >;
+              }
+            ).sessionResumptionUpdate
+          } as LiveSessionResumptionUpdate;
         } else {
           logger.warn(
             `Received an unknown message type from the server: ${JSON.stringify(
diff --git a/packages/ai/src/models/live-generative-model.ts b/packages/ai/src/models/live-generative-model.ts
@@ -25,6 +25,7 @@ import {
   Content,
   LiveGenerationConfig,
   LiveModelParams,
+  SessionResumptionConfig,
   Tool,
   ToolConfig
 } from '../public-types';
@@ -70,12 +71,28 @@ export class LiveGenerativeModel extends AIModel {
   /**
    * Starts a {@link LiveSession}.
    *
+   * @param sessionResumption - Optional configuration for session resumption.
    * @returns A {@link LiveSession}.
    * @throws If the connection failed to be established with the server.
    *
    * @beta
    */
-  async connect(): Promise<LiveSession> {
+  async connect(
+    sessionResumption?: SessionResumptionConfig
+  ): Promise<LiveSession> {
+    const serverMessages = await this._internalConnect(sessionResumption);
+    return new LiveSession(
+      this._webSocketHandler,
+      serverMessages,
+      async (resumptionConfig?: SessionResumptionConfig) => {
+        return this._internalConnect(resumptionConfig);
+      }
+    );
+  }
+
+  private async _internalConnect(
+    sessionResumption?: SessionResumptionConfig
+  ): Promise<AsyncGenerator<unknown>> {
     const url = new WebSocketUrl(this._apiSettings);
     await this._webSocketHandler.connect(url.toString());
 
@@ -102,7 +119,8 @@ export class LiveGenerativeModel extends AIModel {
         toolConfig: this.toolConfig,
         systemInstruction: this.systemInstruction,
         inputAudioTranscription,
-        outputAudioTranscription
+        outputAudioTranscription,
+        sessionResumption
       }
     };
 
@@ -125,7 +143,7 @@ export class LiveGenerativeModel extends AIModel {
         );
       }
 
-      return new LiveSession(this._webSocketHandler, serverMessages);
+      return serverMessages;
     } catch (e) {
       // Ensure connection is closed on any setup error
       await this._webSocketHandler.close();
diff --git a/packages/ai/src/types/live-responses.ts b/packages/ai/src/types/live-responses.ts
@@ -24,6 +24,7 @@ import {
 import {
   AudioTranscriptionConfig,
   LiveGenerationConfig,
+  SessionResumptionConfig,
   Tool,
   ToolConfig
 } from './requests';
@@ -88,6 +89,7 @@ export interface _LiveClientSetup {
     systemInstruction?: string | Part | Content;
     inputAudioTranscription?: AudioTranscriptionConfig;
     outputAudioTranscription?: AudioTranscriptionConfig;
+    sessionResumption?: SessionResumptionConfig;
   };
 }
 
diff --git a/packages/ai/src/types/requests.ts b/packages/ai/src/types/requests.ts
@@ -203,6 +203,63 @@ export interface LiveGenerationConfig {
    * "How are you today?", the model may transcribe that output across three messages, broken up as "How a", "re yo", "u today?".
    */
   outputAudioTranscription?: AudioTranscriptionConfig;
+  /**
+   * The context window compression configuration.
+   *
+   * @beta
+   */
+  contextWindowCompression?: ContextWindowCompressionConfig;
+}
+
+/**
+ * Configures the sliding window context compression mechanism.
+ *
+ * The context window will be truncated by keeping only a suffix of it.
+ *
+ * @beta
+ */
+export interface SlidingWindow {
+  /**
+   * The session reduction target, i.e., how many tokens we should keep.
+   */
+  targetTokens?: number;
+}
+
+/**
+ * Enables context window compression to manage the model's context window.
+ *
+ * This mechanism prevents the context from exceeding a given length.
+ *
+ * @beta
+ */
+export interface ContextWindowCompressionConfig {
+  /**
+   * The number of tokens (before running a turn) that triggers the context
+   * window compression.
+   */
+  triggerTokens?: number;
+
+  /**
+   * The sliding window compression mechanism.
+   */
+  slidingWindow?: SlidingWindow;
+}
+
+/**
+ * Configuration for the session resumption mechanism.
+ *
+ * When included in the session setup, the server will send
+ * {@link LiveSessionResumptionUpdate} messages in the response stream.
+ *
+ * @beta
+ */
+export interface SessionResumptionConfig {
+  /**
+   * The session resumption handle of the previous session to restore.
+   *
+   * If not present, a new session will be started.
+   */
+  handle?: string;
 }
 
 /**
diff --git a/packages/ai/src/types/responses.ts b/packages/ai/src/types/responses.ts
@@ -629,6 +629,30 @@ export interface LiveServerGoingAwayNotice {
   timeLeft: number;
 }
 
+/**
+ * An update of the session resumption state.
+ *
+ * This message is only sent if {@link SessionResumptionConfig} was set in the
+ * session setup.
+ *
+ * @beta
+ */
+export interface LiveSessionResumptionUpdate {
+  type: 'sessionResumptionUpdate';
+  /**
+   * The new handle that represents the state that can be resumed. Empty if `resumable` is false.
+   */
+  newHandle?: string;
+  /**
+   * Indicates if the session can be resumed at this point.
+   */
+  resumable?: boolean;
+  /**
+   * The index of the last client message that is included in the state represented by this update.
+   */
+  lastConsumedClientMessageIndex?: number;
+}
+
 /**
  * The types of responses that can be returned by {@link LiveSession.receive}.
  *
@@ -638,7 +662,8 @@ export const LiveResponseType = {
   SERVER_CONTENT: 'serverContent',
   TOOL_CALL: 'toolCall',
   TOOL_CALL_CANCELLATION: 'toolCallCancellation',
-  GOING_AWAY_NOTICE: 'goingAwayNotice'
+  GOING_AWAY_NOTICE: 'goingAwayNotice',
+  SESSION_RESUMPTION_UPDATE: 'sessionResumptionUpdate'
 };
 
 /**
