From 09c7558e26f4c2a62e3275307ca76969569398e9 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 14 Jul 2026 00:46:51 +0100 Subject: [PATCH] fix: release temporary Amp sessions --- CHANGELOG.md | 1 + .../Providers/Amp/AmpUsageFetcher.swift | 58 ++++++-- .../CodexBarTests/AmpUsageFetcherTests.swift | 139 ++++++++++++++++++ 3 files changed, 187 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 053cc0ef21..0bda5745f8 100644 --- a/CHANGELOG.md +++ b/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! diff --git a/Sources/CodexBarCore/Providers/Amp/AmpUsageFetcher.swift b/Sources/CodexBarCore/Providers/Amp/AmpUsageFetcher.swift index 17d9c29e52..ea612f0806 100644 --- a/Sources/CodexBarCore/Providers/Amp/AmpUsageFetcher.swift +++ b/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 diff --git a/Tests/CodexBarTests/AmpUsageFetcherTests.swift b/Tests/CodexBarTests/AmpUsageFetcherTests.swift index 14fb5c9282..b291201d3d 100644 --- a/Tests/CodexBarTests/AmpUsageFetcherTests.swift +++ b/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 = """ + + """ + 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() {} }