Fix retained Amp network sessions

CHANGELOG.md

@@ -10,6 +10,7 @@
 - Claude: cache successful CLI version probes for 30 minutes while invalidating on executable changes, avoiding repeated PTY launches without retaining failed or stale wrapper results. Thanks @Yuxin-Qiao!
 - Linux CLI: bootstrap the configured IANA timezone before Foundation startup on non-FHS systems, preventing SIGILL on NixOS (#2127). Thanks @xikhar!
 - Ollama: release temporary dashboard network sessions after each fetch, preventing repeated refreshes from retaining delegates and URL-cache resources. Thanks @astuteprogrammer!
+- Amp: release temporary API and dashboard network sessions after every fetch, preventing repeated refreshes from retaining delegates and URL-cache resources.
 - Linux CLI: prevent usage rendering from crashing in Foundation bundle discovery when formatting rate windows. Thanks @thanthi-del!
 - Menus: keep overview provider-row clicks reliable during live menu rebuilds without stealing nested Copy or plan actions. Thanks @Yuxin-Qiao!
 - Startup: load persisted plan-utilization history away from the main thread so mature histories no longer delay app launch. Thanks @Yuxin-Qiao!

Sources/CodexBarCore/Providers/Amp/AmpUsageFetcher.swift

@@ -105,9 +105,25 @@ public struct AmpUsageFetcher: Sendable {
     @MainActor private static var recentDumps: [String] = []
 
     public let browserDetection: BrowserDetection
+    private let makeURLSession: @Sendable (URLSessionTaskDelegate?) -> URLSession
+    private let finishURLSession: @Sendable (URLSession) -> Void
 
     public init(browserDetection: BrowserDetection) {
         self.browserDetection = browserDetection
+        self.makeURLSession = { delegate in
+            URLSession(configuration: .ephemeral, delegate: delegate, delegateQueue: nil)
+        }
+        self.finishURLSession = { $0.finishTasksAndInvalidate() }
+    }
+
+    init(
+        browserDetection: BrowserDetection,
+        makeURLSession: @escaping @Sendable (URLSessionTaskDelegate?) -> URLSession,
+        finishURLSession: @escaping @Sendable (URLSession) -> Void = { $0.finishTasksAndInvalidate() })
+    {
+        self.browserDetection = browserDetection
+        self.makeURLSession = makeURLSession
+        self.finishURLSession = finishURLSession
     }
 
     public func fetch(
@@ -154,7 +170,8 @@ public struct AmpUsageFetcher: Sendable {
         }
         let request = try Self.makeUsageAPIRequest(apiToken: token)
         let diagnostics = APIRedirectDiagnostics(logger: logger)
-        let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil)
+        let session = self.makeURLSession(diagnostics)
+        defer { self.finishURLSession(session) }
         let httpResponse = try await session.response(for: request)
         logger?("[amp] API response: \(httpResponse.statusCode) " +
             "\(httpResponse.response.url?.absoluteString ?? "unknown")")
@@ -268,7 +285,8 @@ public struct AmpUsageFetcher: Sendable {
             forHTTPHeaderField: "accept")
         Self.applyBrowserHeaders(to: &request)
 
-        let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil)
+        let session = self.makeURLSession(diagnostics)
+        defer { self.finishURLSession(session) }
         let httpResponse = try await session.response(for: request)
         let responseInfo = ResponseInfo(
             statusCode: httpResponse.statusCode,
@@ -331,7 +349,9 @@ public struct AmpUsageFetcher: Sendable {
     }
 
     @MainActor private static func recordDump(_ text: String) {
-        if self.recentDumps.count >= 5 { self.recentDumps.removeFirst() }
+        if self.recentDumps.count >= 5 {
+            self.recentDumps.removeFirst()
+        }
         self.recentDumps.append(text)
     }
 
@@ -463,27 +483,43 @@ public struct AmpUsageFetcher: Sendable {
 
     private static func isAmpHost(_ url: URL?) -> Bool {
         guard let host = url?.host?.lowercased() else { return false }
-        if host == "ampcode.com" || host == "www.ampcode.com" { return true }
+        if host == "ampcode.com" || host == "www.ampcode.com" {
+            return true
+        }
         return host.hasSuffix(".ampcode.com")
     }
 
     static func isLoginRedirect(_ url: URL) -> Bool {
         guard self.isAmpHost(url) else { return false }
-        if url.host?.lowercased() == "auth.ampcode.com" { return true }
+        if url.host?.lowercased() == "auth.ampcode.com" {
+            return true
+        }
 
         let path = url.path.lowercased()
         let components = path.split(separator: "/").map(String.init)
-        if components.contains("login") { return true }
-        if components.contains("signin") { return true }
-        if components.contains("sign-in") { return true }
+        if components.contains("login") {
+            return true
+        }
+        if components.contains("signin") {
+            return true
+        }
+        if components.contains("sign-in") {
+            return true
+        }
 
         // Amp currently redirects to /auth/sign-in?returnTo=... when session is invalid. Keep this slightly broader
         // than one exact path so we keep working if Amp changes auth routes.
         if components.contains("auth") {
             let query = url.query?.lowercased() ?? ""
-            if query.contains("returnto=") { return true }
-            if query.contains("redirect=") { return true }
-            if query.contains("redirectto=") { return true }
+            if query.contains("returnto=") {
+                return true
+            }
+            if query.contains("redirect=") {
+                return true
+            }
+            if query.contains("redirectto=") {
+                return true
+            }
         }
 
         return false

Tests/CodexBarTests/AmpUsageFetcherTests.swift

@@ -2,6 +2,7 @@ import Foundation
 import Testing
 @testable import CodexBarCore
 
+@Suite(.serialized)
 struct AmpUsageFetcherTests {
     private func makeContext(
         sourceMode: ProviderSourceMode,
@@ -159,4 +160,142 @@ struct AmpUsageFetcherTests {
         let evil = try #require(URL(string: "https://ampcode.com.evil.com/auth/sign-in"))
         #expect(!AmpUsageFetcher.isLoginRedirect(evil))
     }
+
+    @Test
+    func `temporary API session is finished after a successful request`() async throws {
+        defer { AmpStubURLProtocol.handler = nil }
+        AmpStubURLProtocol.handler = { request in
+            let displayText = "Amp Free: $8/$10 remaining (replenishes +$0.5/hour)"
+            let data = try JSONSerialization.data(withJSONObject: [
+                "ok": true,
+                "result": ["displayText": displayText],
+            ])
+            return try Self.makeResponse(request: request, data: data)
+        }
+        let recorder = AmpSessionFinishRecorder()
+        let fetcher = self.makeFetcher(recorder: recorder)
+
+        _ = try await fetcher.fetch(apiToken: "test")
+
+        #expect(recorder.count == 1)
+    }
+
+    @Test
+    func `temporary API session is finished after a transport failure`() async {
+        defer { AmpStubURLProtocol.handler = nil }
+        AmpStubURLProtocol.handler = { _ in throw URLError(.notConnectedToInternet) }
+        let recorder = AmpSessionFinishRecorder()
+        let fetcher = self.makeFetcher(recorder: recorder)
+
+        await #expect(throws: URLError.self) {
+            _ = try await fetcher.fetch(apiToken: "test")
+        }
+        #expect(recorder.count == 1)
+    }
+
+    @Test
+    func `temporary web session is finished after a successful request`() async throws {
+        defer { AmpStubURLProtocol.handler = nil }
+        AmpStubURLProtocol.handler = { request in
+            let html = """
+            <script>
+            __sveltekit_x.data = {user:{},
+            freeTierUsage:{bucket:"ubi",quota:1000,hourlyReplenishment:42,windowHours:24,used:338.5}};
+            </script>
+            """
+            return try Self.makeResponse(request: request, data: Data(html.utf8))
+        }
+        let recorder = AmpSessionFinishRecorder()
+        let fetcher = self.makeFetcher(recorder: recorder)
+
+        _ = try await fetcher.fetch(cookieHeaderOverride: "session=test")
+
+        #expect(recorder.count == 1)
+    }
+
+    @Test
+    func `temporary web session is finished after a transport failure`() async {
+        defer { AmpStubURLProtocol.handler = nil }
+        AmpStubURLProtocol.handler = { _ in throw URLError(.notConnectedToInternet) }
+        let recorder = AmpSessionFinishRecorder()
+        let fetcher = self.makeFetcher(recorder: recorder)
+
+        await #expect(throws: URLError.self) {
+            _ = try await fetcher.fetch(cookieHeaderOverride: "session=test")
+        }
+        #expect(recorder.count == 1)
+    }
+
+    private func makeFetcher(recorder: AmpSessionFinishRecorder) -> AmpUsageFetcher {
+        AmpUsageFetcher(
+            browserDetection: BrowserDetection(cacheTTL: 0),
+            makeURLSession: { delegate in
+                let configuration = URLSessionConfiguration.ephemeral
+                configuration.protocolClasses = [AmpStubURLProtocol.self]
+                return URLSession(configuration: configuration, delegate: delegate, delegateQueue: nil)
+            },
+            finishURLSession: { session in
+                recorder.record(session)
+                session.finishTasksAndInvalidate()
+            })
+    }
+
+    private static func makeResponse(
+        request: URLRequest,
+        data: Data,
+        statusCode: Int = 200) throws -> (HTTPURLResponse, Data)
+    {
+        let url = try #require(request.url)
+        let response = try #require(HTTPURLResponse(
+            url: url,
+            statusCode: statusCode,
+            httpVersion: "HTTP/1.1",
+            headerFields: ["Content-Type": "application/json"]))
+        return (response, data)
+    }
+}
+
+private final class AmpSessionFinishRecorder: @unchecked Sendable {
+    private let lock = NSLock()
+    private var sessions: [URLSession] = []
+
+    var count: Int {
+        self.lock.withLock { self.sessions.count }
+    }
+
+    func record(_ session: URLSession) {
+        self.lock.withLock {
+            self.sessions.append(session)
+        }
+    }
+}
+
+private final class AmpStubURLProtocol: URLProtocol {
+    nonisolated(unsafe) static var handler: ((URLRequest) throws -> (HTTPURLResponse, Data))?
+
+    override static func canInit(with request: URLRequest) -> Bool {
+        request.url?.host?.hasSuffix("ampcode.com") == true
+    }
+
+    override static func canonicalRequest(for request: URLRequest) -> URLRequest {
+        request
+    }
+
+    override func startLoading() {
+        guard let handler = Self.handler else {
+            self.client?.urlProtocol(self, didFailWithError: URLError(.badServerResponse))
+            return
+        }
+
+        do {
+            let (response, data) = try handler(self.request)
+            self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
+            self.client?.urlProtocol(self, didLoad: data)
+            self.client?.urlProtocolDidFinishLoading(self)
+        } catch {
+            self.client?.urlProtocol(self, didFailWithError: error)
+        }
+    }
+
+    override func stopLoading() {}
 }