[AI] Add live session resumption

FirebaseAI/Sources/Types/Internal/Live/BidiContextWindowCompressionConfig.swift

@@ -0,0 +1,47 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+struct BidiSlidingWindow: Encodable, Sendable {
+  let targetTokens: Int?
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+struct BidiContextWindowCompressionConfig: Encodable, Sendable {
+  let triggerTokens: Int?
+  let slidingWindow: BidiSlidingWindow?
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+extension SlidingWindow {
+  var bidiSlidingWindow: BidiSlidingWindow {
+    BidiSlidingWindow(targetTokens: targetTokens)
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+extension ContextWindowCompressionConfig {
+  var bidiContextWindowCompressionConfig: BidiContextWindowCompressionConfig {
+    BidiContextWindowCompressionConfig(
+      triggerTokens: triggerTokens,
+      slidingWindow: slidingWindow?.bidiSlidingWindow
+    )
+  }
+}

FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentServerMessage.swift

@@ -35,6 +35,9 @@ struct BidiGenerateContentServerMessage: Sendable {
     /// and should be cancelled.
     case toolCallCancellation(BidiGenerateContentToolCallCancellation)
 
+    /// Update with the state needed to resume the session.
+    case sessionResumptionUpdate(BidiSessionResumptionUpdate)
+
     /// Server will disconnect soon.
     case goAway(GoAway)
   }
@@ -56,6 +59,7 @@ extension BidiGenerateContentServerMessage: Decodable {
     case serverContent
     case toolCall
     case toolCallCancellation
+    case sessionResumptionUpdate
     case goAway
     case usageMetadata
   }
@@ -83,6 +87,11 @@ extension BidiGenerateContentServerMessage: Decodable {
       forKey: .toolCallCancellation
     ) {
       messageType = .toolCallCancellation(toolCallCancellation)
+    } else if let sessionResumptionUpdate = try container.decodeIfPresent(
+      BidiSessionResumptionUpdate.self,
+      forKey: .sessionResumptionUpdate
+    ) {
+      messageType = .sessionResumptionUpdate(sessionResumptionUpdate)
     } else if let goAway = try container.decodeIfPresent(GoAway.self, forKey: .goAway) {
       messageType = .goAway(goAway)
     } else {

FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentSetup.swift

@@ -56,20 +56,25 @@ struct BidiGenerateContentSetup: Encodable {
   /// turn.
   let outputAudioTranscription: BidiAudioTranscriptionConfig?
 
+  /// Configuration for the session resumption mechanism.
+  let sessionResumption: BidiSessionResumptionConfig?
+
   init(model: String,
        generationConfig: BidiGenerationConfig? = nil,
        systemInstruction: ModelContent? = nil,
        tools: [Tool]? = nil,
        toolConfig: ToolConfig? = nil,
        inputAudioTranscription: BidiAudioTranscriptionConfig? = nil,
-       outputAudioTranscription: BidiAudioTranscriptionConfig? = nil) {
+       outputAudioTranscription: BidiAudioTranscriptionConfig? = nil,
+       sessionResumption: BidiSessionResumptionConfig? = nil) {
     self.model = model
     self.generationConfig = generationConfig
     self.systemInstruction = systemInstruction
     self.tools = tools
     self.toolConfig = toolConfig
     self.inputAudioTranscription = inputAudioTranscription
     self.outputAudioTranscription = outputAudioTranscription
+    self.sessionResumption = sessionResumption
   }
 }
 

FirebaseAI/Sources/Types/Internal/Live/BidiGenerationConfig.swift

@@ -27,12 +27,14 @@ struct BidiGenerationConfig: Encodable, Sendable {
   let frequencyPenalty: Float?
   let responseModalities: [ResponseModality]?
   let speechConfig: BidiSpeechConfig?
+  let contextWindowCompression: BidiContextWindowCompressionConfig?
 
   init(temperature: Float? = nil, topP: Float? = nil, topK: Int? = nil,
        candidateCount: Int? = nil, maxOutputTokens: Int? = nil,
        presencePenalty: Float? = nil, frequencyPenalty: Float? = nil,
        responseModalities: [ResponseModality]? = nil,
-       speechConfig: BidiSpeechConfig? = nil) {
+       speechConfig: BidiSpeechConfig? = nil,
+       contextWindowCompression: BidiContextWindowCompressionConfig? = nil) {
     self.temperature = temperature
     self.topP = topP
     self.topK = topK
@@ -42,5 +44,6 @@ struct BidiGenerationConfig: Encodable, Sendable {
     self.frequencyPenalty = frequencyPenalty
     self.responseModalities = responseModalities
     self.speechConfig = speechConfig
+    self.contextWindowCompression = contextWindowCompression
   }
 }

FirebaseAI/Sources/Types/Internal/Live/BidiSessionResumption.swift

@@ -0,0 +1,49 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+struct BidiSessionResumptionConfig: Encodable, Sendable {
+  let handle: String?
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+struct BidiSessionResumptionUpdate: Decodable, Sendable {
+  let newHandle: String?
+  let resumable: Bool?
+  let lastConsumedClientMessageIndex: Int?
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+extension SessionResumptionConfig {
+  var bidiSessionResumptionConfig: BidiSessionResumptionConfig {
+    BidiSessionResumptionConfig(handle: handle)
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+extension LiveSessionResumptionUpdate {
+  init(_ msg: BidiSessionResumptionUpdate) {
+    self.init(
+      newHandle: msg.newHandle,
+      resumable: msg.resumable,
+      lastConsumedClientMessageIndex: msg.lastConsumedClientMessageIndex
+    )
+  }
+}

FirebaseAI/Sources/Types/Internal/Live/LiveSessionService.swift

@@ -109,11 +109,11 @@ actor LiveSessionService {
   /// resuming the same session.
   ///
   /// This function will yield until the websocket is ready to communicate with the client.
-  func connect() async throws {
+  func connect(sessionResumption: SessionResumptionConfig? = nil) async throws {
     close()
 
     let stream = try await setupWebsocket()
-    try await waitForSetupComplete(stream: stream)
+    try await waitForSetupComplete(stream: stream, sessionResumption: sessionResumption)
     spawnMessageTasks(stream: stream)
   }
 
@@ -138,10 +138,10 @@ actor LiveSessionService {
   /// - Server sends back `BidiGenerateContentSetupComplete` when it's ready
   ///
   /// This function will yield until the setup is complete.
-  private func waitForSetupComplete(stream: MappedStream<
-    URLSessionWebSocketTask.Message,
-    Data
-  >) async throws {
+  private func waitForSetupComplete(
+    stream: MappedStream<URLSessionWebSocketTask.Message, Data>,
+    sessionResumption: SessionResumptionConfig?
+  ) async throws {
     guard let webSocket else { return }
 
     do {
@@ -152,7 +152,8 @@ actor LiveSessionService {
         tools: tools,
         toolConfig: toolConfig,
         inputAudioTranscription: generationConfig?.inputAudioTranscription,
-        outputAudioTranscription: generationConfig?.outputAudioTranscription
+        outputAudioTranscription: generationConfig?.outputAudioTranscription,
+        sessionResumption: sessionResumption?.bidiSessionResumptionConfig
       )
       let data = try jsonEncoder.encode(BidiGenerateContentClientMessage.setup(setup))
       try await webSocket.send(.data(data))

FirebaseAI/Sources/Types/Public/Live/ContextWindowCompressionConfig.swift

@@ -0,0 +1,56 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+/// Configures the sliding window context compression mechanism.
+///
+/// The context window will be truncated by keeping only a suffix of it.
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+public struct SlidingWindow: Sendable {
+  /// The session reduction target, i.e., how many tokens we should keep.
+  public let targetTokens: Int?
+
+  /// Creates a ``SlidingWindow`` instance.
+  ///
+  /// - Parameter targetTokens: The target number of tokens to keep in the context window.
+  public init(targetTokens: Int? = nil) {
+    self.targetTokens = targetTokens
+  }
+}
+
+/// Enables context window compression to manage the model's context window.
+///
+/// This mechanism prevents the context from exceeding a given length.
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
+@available(watchOS, unavailable)
+public struct ContextWindowCompressionConfig: Sendable {
+  /// The number of tokens (before running a turn) that triggers the context
+  /// window compression.
+  public let triggerTokens: Int?
+
+  /// The sliding window compression mechanism.
+  public let slidingWindow: SlidingWindow?
+
+  /// Creates a ``ContextWindowCompressionConfig`` instance.
+  ///
+  /// - Parameters:
+  ///   - triggerTokens: The number of tokens that triggers the compression mechanism.
+  ///   - slidingWindow: The sliding window compression mechanism to use.
+  public init(triggerTokens: Int? = nil, slidingWindow: SlidingWindow? = nil) {
+    self.triggerTokens = triggerTokens
+    self.slidingWindow = slidingWindow
+  }
+}

FirebaseAI/Sources/Types/Public/Live/LiveGenerationConfig.swift

@@ -21,6 +21,7 @@ public struct LiveGenerationConfig: Sendable {
   let bidiGenerationConfig: BidiGenerationConfig
   let inputAudioTranscription: BidiAudioTranscriptionConfig?
   let outputAudioTranscription: BidiAudioTranscriptionConfig?
+  public let contextWindowCompression: ContextWindowCompressionConfig?
 
   /// Creates a new ``LiveGenerationConfig`` value.
   ///
@@ -125,7 +126,8 @@ public struct LiveGenerationConfig: Sendable {
               responseModalities: [ResponseModality]? = nil,
               speech: SpeechConfig? = nil,
               inputAudioTranscription: AudioTranscriptionConfig? = nil,
-              outputAudioTranscription: AudioTranscriptionConfig? = nil) {
+              outputAudioTranscription: AudioTranscriptionConfig? = nil,
+              contextWindowCompression: ContextWindowCompressionConfig? = nil) {
     self.init(
       BidiGenerationConfig(
         temperature: temperature,
@@ -136,18 +138,22 @@ public struct LiveGenerationConfig: Sendable {
         presencePenalty: presencePenalty,
         frequencyPenalty: frequencyPenalty,
         responseModalities: responseModalities,
-        speechConfig: speech?.speechConfig
+        speechConfig: speech?.speechConfig,
+        contextWindowCompression: contextWindowCompression?.bidiContextWindowCompressionConfig
       ),
       inputAudioTranscription: inputAudioTranscription?.audioTranscriptionConfig,
-      outputAudioTranscription: outputAudioTranscription?.audioTranscriptionConfig
+      outputAudioTranscription: outputAudioTranscription?.audioTranscriptionConfig,
+      contextWindowCompression: contextWindowCompression
     )
   }
 
   init(_ bidiGenerationConfig: BidiGenerationConfig,
        inputAudioTranscription: BidiAudioTranscriptionConfig? = nil,
-       outputAudioTranscription: BidiAudioTranscriptionConfig? = nil) {
+       outputAudioTranscription: BidiAudioTranscriptionConfig? = nil,
+       contextWindowCompression: ContextWindowCompressionConfig? = nil) {
     self.bidiGenerationConfig = bidiGenerationConfig
     self.inputAudioTranscription = inputAudioTranscription
     self.outputAudioTranscription = outputAudioTranscription
+    self.contextWindowCompression = contextWindowCompression
   }
 }

FirebaseAI/Sources/Types/Public/Live/LiveGenerativeModel.swift

@@ -53,8 +53,9 @@ public final class LiveGenerativeModel {
 
   /// Start a ``LiveSession`` with the server for bidirectional streaming.
   ///
+  /// - Parameter sessionResumption: The configuration for session resumption.
   /// - Returns: A new ``LiveSession`` that you can use to stream messages to and from the server.
-  public func connect() async throws -> LiveSession {
+  public func connect(sessionResumption: SessionResumptionConfig? = nil) async throws -> LiveSession {
     let service = LiveSessionService(
       modelResourceName: modelResourceName,
       generationConfig: generationConfig,
@@ -67,7 +68,7 @@ public final class LiveGenerativeModel {
       requestOptions: requestOptions
     )
 
-    try await service.connect()
+    try await service.connect(sessionResumption: sessionResumption)
 
     return LiveSession(service: service)
   }

FirebaseAI/Sources/Types/Public/Live/LiveServerMessage.swift

@@ -34,6 +34,9 @@ public struct LiveServerMessage: Sendable {
     /// cancelled.
     case toolCallCancellation(LiveServerToolCallCancellation)
 
+    /// Update with the session state needed to resume the session.
+    case sessionResumptionUpdate(LiveSessionResumptionUpdate)
+
     /// Server will disconnect soon.
     case goingAwayNotice(LiveServerGoingAwayNotice)
   }
@@ -74,6 +77,8 @@ extension LiveServerMessage.Payload {
       self = .toolCall(LiveServerToolCall(msg))
     case let .toolCallCancellation(msg):
       self = .toolCallCancellation(LiveServerToolCallCancellation(msg))
+    case let .sessionResumptionUpdate(msg):
+      self = .sessionResumptionUpdate(LiveSessionResumptionUpdate(msg))
     case let .goAway(msg):
       self = .goingAwayNotice(LiveServerGoingAwayNotice(msg))
     }

FirebaseAI/Sources/Types/Public/Live/LiveSession.swift

@@ -147,4 +147,18 @@ public final class LiveSession: Sendable {
   }
 
   // TODO: b(445716402) Add a start method when we support session resumption
+
+
+  /// 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.
+  ///
+  /// - Parameter sessionResumption: The configuration for session resumption,
+  ///   such as the handle to the previous session state to restore.
+  public func resumeSession(sessionResumption: SessionResumptionConfig? = nil) async throws {
+    try await service.connect(sessionResumption: sessionResumption)
+  }
 }
+