From 99f5d2eda02254ed04661bc4954787e35a6dd4d7 Mon Sep 17 00:00:00 2001
From: Aashish Patil <aashishp@google.com>
Date: Mon, 23 Feb 2026 14:24:27 -0800
Subject: [PATCH 1/2] Update maxAge handling to match Dart implementation

---
 Sources/Cache/Cache.swift             | 15 ++++++++++++++-
 Sources/Cache/ResultTree.swift        |  8 ++++++--
 Sources/Queries/GenericQueryRef.swift |  3 +--
 Tests/Unit/CacheTests.swift           |  7 ++++---
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/Sources/Cache/Cache.swift b/Sources/Cache/Cache.swift
index cca56dc..c145145 100644
--- a/Sources/Cache/Cache.swift
+++ b/Sources/Cache/Cache.swift
@@ -89,7 +89,7 @@ actor Cache {
     return encoded
   }
 
-  func resultTree(queryId: String) -> ResultTree? {
+  func resultTree(queryId: String, allowStale: Bool = false) -> ResultTree? {
     // result trees are stored dehydrated in the cache
     // retrieve cache, hydrate it and then return
     guard let cacheProvider else {
@@ -101,6 +101,16 @@ actor Cache {
       return nil
     }
 
+    /*
+     isStale = true, allowStale = false: false || false => return nil
+     isStale = true, allowStale = true: false || true => proceed
+     isStale = false, allowStale = false: true || false => proceed
+     isStale = false, allowStale = true: true || true => proceed
+     */
+    guard !dehydratedTree.isStale || allowStale else {
+      return nil
+    }
+
     do {
       let resultsProcessor = ResultTreeProcessor()
       let (hydratedResults, rootObj) = try resultsProcessor.hydrateResults(
@@ -112,6 +122,7 @@ actor Cache {
         data: hydratedResults,
         cachedAt: dehydratedTree.cachedAt,
         lastAccessed: dehydratedTree.lastAccessed,
+        maxAge: dehydratedTree.maxAge,
         rootObject: rootObj
       )
 
@@ -142,6 +153,7 @@ actor Cache {
                                                                                         [:],
                                                                                       cacheProvider: cacheProvider)
 
+      let maxAge = response.extensions?.maxAge ?? config.maxAge
       cacheProvider
         .setResultTree(
           queryId: queryId,
@@ -149,6 +161,7 @@ actor Cache {
             data: dehydratedResults,
             cachedAt: Date(),
             lastAccessed: Date(),
+            maxAge: maxAge,
             rootObject: rootObj
           )
         )
diff --git a/Sources/Cache/ResultTree.swift b/Sources/Cache/ResultTree.swift
index 0fc93b4..6ccf9ec 100644
--- a/Sources/Cache/ResultTree.swift
+++ b/Sources/Cache/ResultTree.swift
@@ -25,16 +25,20 @@ struct ResultTree: Codable {
   // Local time when the entry was read or updated
   var lastAccessed: Date
 
+  // Validity for the results
+  var maxAge: TimeInterval
+
   var rootObject: EntityNode?
 
-  func isStale(_ ttl: TimeInterval) -> Bool {
+  var isStale: Bool {
     let now = Date()
-    return now.timeIntervalSince(cachedAt) > ttl
+    return now.timeIntervalSince(cachedAt) > maxAge
   }
 
   enum CodingKeys: String, CodingKey {
     case cachedAt = "ca" // cached at
     case lastAccessed = "la" // last accessed
     case data = "d" // data cached
+    case maxAge = "ma" // max age
   }
 }
diff --git a/Sources/Queries/GenericQueryRef.swift b/Sources/Queries/GenericQueryRef.swift
index 2bb2217..89e0ca0 100644
--- a/Sources/Queries/GenericQueryRef.swift
+++ b/Sources/Queries/GenericQueryRef.swift
@@ -136,8 +136,7 @@ actor GenericQueryRef<ResultData: Decodable & Sendable, Variable: OperationVaria
       return OperationResult(data: nil, source: .cache)
     }
 
-    if let cacheEntry = await cache.resultTree(queryId: request.requestId),
-       (cacheEntry.isStale(maxAge) && allowStale) || !cacheEntry.isStale(maxAge) {
+    if let cacheEntry = await cache.resultTree(queryId: request.requestId, allowStale: allowStale) {
       let decoder = JSONDecoder()
       let decodedData = try decoder.decode(
         ResultData.self,
diff --git a/Tests/Unit/CacheTests.swift b/Tests/Unit/CacheTests.swift
index c7e35b0..b853752 100644
--- a/Tests/Unit/CacheTests.swift
+++ b/Tests/Unit/CacheTests.swift
@@ -483,16 +483,17 @@ final class CacheTests: XCTestCase {
     let resultTree = ResultTree(
       data: resultTreeData.data(using: .utf8)!,
       cachedAt: Date(),
-      lastAccessed: Date()
+      lastAccessed: Date(),
+      maxAge: ttl
     )
 
     // Within TTL
-    XCTAssertFalse(resultTree.isStale(ttl))
+    XCTAssertFalse(resultTree.isStale)
 
     // Outside TTL
     let expectation = XCTestExpectation(description: "Wait for TTL to expire")
     DispatchQueue.main.asyncAfter(deadline: .now() + ttl + 0.1) {
-      XCTAssertTrue(resultTree.isStale(ttl))
+      XCTAssertTrue(resultTree.isStale)
       expectation.fulfill()
     }
 

From de1003923a367b79c2322033b3f66a4530d3cb22 Mon Sep 17 00:00:00 2001
From: Aashish Patil <aashishp@google.com>
Date: Mon, 23 Feb 2026 14:53:50 -0800
Subject: [PATCH 2/2] Address Gemini feedback

---
 Sources/Cache/Cache.swift             |  6 ------
 Sources/Queries/GenericQueryRef.swift | 20 --------------------
 2 files changed, 26 deletions(-)

diff --git a/Sources/Cache/Cache.swift b/Sources/Cache/Cache.swift
index c145145..9ee6b90 100644
--- a/Sources/Cache/Cache.swift
+++ b/Sources/Cache/Cache.swift
@@ -101,12 +101,6 @@ actor Cache {
       return nil
     }
 
-    /*
-     isStale = true, allowStale = false: false || false => return nil
-     isStale = true, allowStale = true: false || true => proceed
-     isStale = false, allowStale = false: true || false => proceed
-     isStale = false, allowStale = true: true || true => proceed
-     */
     guard !dehydratedTree.isStale || allowStale else {
       return nil
     }
diff --git a/Sources/Queries/GenericQueryRef.swift b/Sources/Queries/GenericQueryRef.swift
index 89e0ca0..62b8246 100644
--- a/Sources/Queries/GenericQueryRef.swift
+++ b/Sources/Queries/GenericQueryRef.swift
@@ -30,25 +30,6 @@ actor GenericQueryRef<ResultData: Decodable & Sendable, Variable: OperationVaria
 
   private let cache: Cache?
 
-  // maxAge received from server in query response
-  private var serverMaxAge: TimeInterval? = nil
-
-  // maxAge computed based on server or cache settings
-  // server is given priority over cache settings
-  private var maxAge: TimeInterval {
-    if let serverMaxAge {
-      DataConnectLogger.debug("Using maxAge specified from server \(serverMaxAge)")
-      return serverMaxAge
-    }
-
-    if let ma = cache?.config.maxAge {
-      DataConnectLogger.debug("Using maxAge specified from cache settings \(ma)")
-      return ma
-    }
-
-    return 0
-  }
-
   // Ideally we would like this to be part of the QueryRef protocol
   // Not adding for now since the protocol is Public
   // This property is for now an internal property.
@@ -111,7 +92,6 @@ actor GenericQueryRef<ResultData: Decodable & Sendable, Variable: OperationVaria
 
     do {
       if let cache {
-        serverMaxAge = response.extensions?.maxAge
         await cache.update(queryId: operationId, response: response, requestor: self)
       }
     }