feat(account): add managed credits and cloud gateway
Introduce optional Apple account-backed credits with scoped gateway access while preserving local and BYOK paths. Refresh assistant behavior, tests, privacy disclosures, docs, and the website for the 2.0 experience.
This commit is contained in:
Rocky
@@ -0,0 +1,256 @@
|
||||
// GatewayGrantCoordinator.swift
|
||||
// OSGKeyboard · Shared
|
||||
//
|
||||
// Creates and rotates scope-limited gateway grants. Refresh rotation is merged
|
||||
// inside this actor so concurrent extension requests never replay an old token.
|
||||
|
||||
import CryptoKit
|
||||
import Foundation
|
||||
|
||||
public actor GatewayGrantCoordinator {
|
||||
public static let defaultBaseURL = URL(string: "https://account.osglab.com")!
|
||||
|
||||
private let baseURL: URL
|
||||
private let store: any GatewayGrantCredentialStore
|
||||
private let session: URLSession
|
||||
private let now: @Sendable () -> Date
|
||||
private var refreshTask: Task<ManagedGatewayGrantCredentials, Error>?
|
||||
|
||||
public init(
|
||||
baseURL: URL = GatewayGrantCoordinator.defaultBaseURL,
|
||||
store: any GatewayGrantCredentialStore = GatewayGrantKeychainStore(),
|
||||
session: URLSession = .shared,
|
||||
now: @escaping @Sendable () -> Date = Date.init
|
||||
) {
|
||||
self.baseURL = baseURL
|
||||
self.store = store
|
||||
self.session = session
|
||||
self.now = now
|
||||
}
|
||||
|
||||
/// Host-only integration point. The account access token authorizes grant
|
||||
/// creation but is used only for this request and is never persisted here.
|
||||
@discardableResult
|
||||
public func createGrant(
|
||||
accountAccessToken: String,
|
||||
scopes: Set<ManagedGatewayCapability>,
|
||||
lifetimeSeconds: Int? = nil,
|
||||
idempotencyKey: String = UUID().uuidString
|
||||
) async throws -> ManagedGatewayGrantCredentials {
|
||||
guard !accountAccessToken.isEmpty else { throw ManagedGatewayError.invalidGrant }
|
||||
guard !scopes.isEmpty else {
|
||||
throw ManagedGatewayError.server(
|
||||
code: "invalid_request",
|
||||
status: 400,
|
||||
requestId: nil
|
||||
)
|
||||
}
|
||||
|
||||
struct Body: Encodable {
|
||||
let scopes: [ManagedGatewayCapability]
|
||||
let lifetimeSeconds: Int?
|
||||
}
|
||||
|
||||
var request = URLRequest(url: endpoint("v1/gateway/grants"))
|
||||
request.httpMethod = "POST"
|
||||
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
||||
request.setValue("Bearer \(accountAccessToken)", forHTTPHeaderField: "Authorization")
|
||||
request.setValue(idempotencyKey, forHTTPHeaderField: "Idempotency-Key")
|
||||
request.setValue(UUID().uuidString, forHTTPHeaderField: "X-Request-ID")
|
||||
request.httpBody = try JSONEncoder().encode(
|
||||
Body(scopes: scopes.sorted { $0.rawValue < $1.rawValue }, lifetimeSeconds: lifetimeSeconds)
|
||||
)
|
||||
|
||||
let credentials = try await sendGrantRequest(request)
|
||||
guard credentials.scopes == scopes else {
|
||||
throw ManagedGatewayError.invalidGrant
|
||||
}
|
||||
try await store.save(credentials)
|
||||
return credentials
|
||||
}
|
||||
|
||||
public func accessToken(
|
||||
for scope: ManagedGatewayCapability,
|
||||
forceRefresh: Bool = false
|
||||
) async throws -> String {
|
||||
guard let credentials = try await store.load() else {
|
||||
throw ManagedGatewayError.missingGrant
|
||||
}
|
||||
guard credentials.scopes.contains(scope) else {
|
||||
throw ManagedGatewayError.scopeNotGranted(scope)
|
||||
}
|
||||
if !forceRefresh, credentials.hasUsableAccessToken(for: scope, at: now()) {
|
||||
return credentials.accessToken
|
||||
}
|
||||
guard credentials.hasUsableRefreshToken(at: now()) else {
|
||||
try? await store.delete()
|
||||
throw ManagedGatewayError.invalidGrant
|
||||
}
|
||||
return try await refresh(credentials).accessToken
|
||||
}
|
||||
|
||||
public func clearGrant() async throws {
|
||||
refreshTask?.cancel()
|
||||
refreshTask = nil
|
||||
try await store.delete()
|
||||
}
|
||||
|
||||
private func refresh(
|
||||
_ credentials: ManagedGatewayGrantCredentials
|
||||
) async throws -> ManagedGatewayGrantCredentials {
|
||||
if let refreshTask {
|
||||
return try await refreshTask.value
|
||||
}
|
||||
|
||||
let task = Task {
|
||||
try await requestRefresh(using: credentials)
|
||||
}
|
||||
refreshTask = task
|
||||
do {
|
||||
let refreshed = try await task.value
|
||||
refreshTask = nil
|
||||
return refreshed
|
||||
} catch {
|
||||
refreshTask = nil
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
private func requestRefresh(
|
||||
using credentials: ManagedGatewayGrantCredentials
|
||||
) async throws -> ManagedGatewayGrantCredentials {
|
||||
struct Body: Encodable {
|
||||
let refreshToken: String
|
||||
}
|
||||
|
||||
var request = URLRequest(url: endpoint("v1/gateway/grants/refresh"))
|
||||
request.httpMethod = "POST"
|
||||
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
||||
request.setValue(
|
||||
Self.refreshTokenIdempotencyKey(credentials.refreshToken),
|
||||
forHTTPHeaderField: "Idempotency-Key"
|
||||
)
|
||||
request.setValue(UUID().uuidString, forHTTPHeaderField: "X-Request-ID")
|
||||
request.httpBody = try JSONEncoder().encode(Body(refreshToken: credentials.refreshToken))
|
||||
|
||||
do {
|
||||
let refreshed = try await sendGrantRequest(request)
|
||||
guard refreshed.grantId == credentials.grantId,
|
||||
refreshed.scopes == credentials.scopes else {
|
||||
try? await store.delete()
|
||||
throw ManagedGatewayError.invalidGrant
|
||||
}
|
||||
try await store.save(refreshed)
|
||||
return refreshed
|
||||
} catch let error as ManagedGatewayError where error == .invalidGrant {
|
||||
try? await store.delete()
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
private func sendGrantRequest(
|
||||
_ request: URLRequest
|
||||
) async throws -> ManagedGatewayGrantCredentials {
|
||||
do {
|
||||
let (data, response) = try await session.data(for: request)
|
||||
guard let http = response as? HTTPURLResponse else {
|
||||
throw LLMError.transport("non-HTTP response")
|
||||
}
|
||||
guard (200..<300).contains(http.statusCode) else {
|
||||
throw ManagedGatewayHTTP.error(
|
||||
data: data,
|
||||
status: http.statusCode,
|
||||
requestId: http.value(forHTTPHeaderField: "X-Request-ID")
|
||||
)
|
||||
}
|
||||
do {
|
||||
return try Self.gatewayDecoder()
|
||||
.decode(ManagedGatewayGrantTokenResponse.self, from: data)
|
||||
.credentials(receivedAt: now())
|
||||
} catch {
|
||||
throw LLMError.decoding(String(describing: error))
|
||||
}
|
||||
} catch is CancellationError {
|
||||
throw LLMError.cancelled
|
||||
} catch let error as URLError where error.code == .cancelled {
|
||||
throw LLMError.cancelled
|
||||
} catch let error as URLError where error.code == .timedOut {
|
||||
throw ManagedGatewayError.timeout
|
||||
} catch let error as ManagedGatewayError {
|
||||
throw error
|
||||
} catch let error as LLMError {
|
||||
throw error
|
||||
} catch {
|
||||
throw LLMError.transport(String(describing: error))
|
||||
}
|
||||
}
|
||||
|
||||
private func endpoint(_ path: String) -> URL {
|
||||
baseURL.appending(path: path)
|
||||
}
|
||||
|
||||
static func refreshTokenIdempotencyKey(_ refreshToken: String) -> String {
|
||||
let digest = SHA256.hash(data: Data(refreshToken.utf8))
|
||||
return "gateway-refresh-v1-" + digest.map { String(format: "%02x", $0) }.joined()
|
||||
}
|
||||
|
||||
private static func gatewayDecoder() -> JSONDecoder {
|
||||
let decoder = JSONDecoder()
|
||||
decoder.dateDecodingStrategy = .custom { decoder in
|
||||
let value = try decoder.singleValueContainer().decode(String.self)
|
||||
let fractional = ISO8601DateFormatter()
|
||||
fractional.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
|
||||
if let date = fractional.date(from: value) {
|
||||
return date
|
||||
}
|
||||
let standard = ISO8601DateFormatter()
|
||||
standard.formatOptions = [.withInternetDateTime]
|
||||
guard let date = standard.date(from: value) else {
|
||||
throw DecodingError.dataCorruptedError(
|
||||
in: try decoder.singleValueContainer(),
|
||||
debugDescription: "Invalid ISO-8601 date"
|
||||
)
|
||||
}
|
||||
return date
|
||||
}
|
||||
return decoder
|
||||
}
|
||||
}
|
||||
|
||||
enum ManagedGatewayHTTP {
|
||||
static func error(
|
||||
data: Data,
|
||||
status: Int,
|
||||
requestId: String?
|
||||
) -> ManagedGatewayError {
|
||||
let decoded = decodeError(from: data)
|
||||
let code = decoded?.code ?? HTTPURLResponse.localizedString(forStatusCode: status)
|
||||
let resolvedRequestId = decoded?.requestId ?? requestId
|
||||
|
||||
switch code.lowercased() {
|
||||
case "insufficient_credits", "insufficient_balance", "credit_balance_insufficient":
|
||||
return .insufficientCredits
|
||||
case "unauthorized", "invalid_gateway_refresh", "gateway_grant_denied", "invalid_grant":
|
||||
return .invalidGrant
|
||||
default:
|
||||
return .server(code: code, status: status, requestId: resolvedRequestId)
|
||||
}
|
||||
}
|
||||
|
||||
static func decodeError(from data: Data) -> ManagedGatewayErrorResponse? {
|
||||
if let direct = try? JSONDecoder().decode(ManagedGatewayErrorResponse.self, from: data) {
|
||||
return direct
|
||||
}
|
||||
guard let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
|
||||
let nested = object["error"] as? [String: Any],
|
||||
let code = nested["code"] as? String,
|
||||
let message = nested["message"] as? String else {
|
||||
return nil
|
||||
}
|
||||
return ManagedGatewayErrorResponse(
|
||||
code: code,
|
||||
message: message,
|
||||
requestId: nested["requestId"] as? String ?? ""
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
// GatewayGrantCredentialStore.swift
|
||||
// OSGKeyboard · Shared
|
||||
//
|
||||
// Shared Keychain storage dedicated to scope-limited gateway grants. Account
|
||||
// session credentials intentionally never enter this service or access group item.
|
||||
|
||||
import Foundation
|
||||
import Security
|
||||
|
||||
public protocol GatewayGrantCredentialStore: Sendable {
|
||||
func load() async throws -> ManagedGatewayGrantCredentials?
|
||||
func save(_ credentials: ManagedGatewayGrantCredentials) async throws
|
||||
func delete() async throws
|
||||
}
|
||||
|
||||
public struct GatewayGrantKeychainStore: GatewayGrantCredentialStore, @unchecked Sendable {
|
||||
public enum StoreError: Error, Equatable, Sendable {
|
||||
case unexpectedStatus(OSStatus)
|
||||
case invalidStoredValue
|
||||
}
|
||||
|
||||
private static let service = "com.osgkeyboard.gateway-grant"
|
||||
private static let account = "scope-limited.active"
|
||||
|
||||
public init() {}
|
||||
|
||||
public func load() async throws -> ManagedGatewayGrantCredentials? {
|
||||
var query = Self.baseQuery
|
||||
query[kSecReturnData as String] = true
|
||||
query[kSecMatchLimit as String] = kSecMatchLimitOne
|
||||
|
||||
var result: CFTypeRef?
|
||||
let status = SecItemCopyMatching(query as CFDictionary, &result)
|
||||
switch status {
|
||||
case errSecSuccess:
|
||||
guard let data = result as? Data,
|
||||
let value = try? Self.decoder.decode(
|
||||
ManagedGatewayGrantCredentials.self,
|
||||
from: data
|
||||
) else {
|
||||
throw StoreError.invalidStoredValue
|
||||
}
|
||||
return value
|
||||
case errSecItemNotFound:
|
||||
return nil
|
||||
default:
|
||||
throw StoreError.unexpectedStatus(status)
|
||||
}
|
||||
}
|
||||
|
||||
public func save(_ credentials: ManagedGatewayGrantCredentials) async throws {
|
||||
let data = try Self.encoder.encode(credentials)
|
||||
let updateStatus = SecItemUpdate(
|
||||
Self.baseQuery as CFDictionary,
|
||||
[kSecValueData as String: data] as CFDictionary
|
||||
)
|
||||
switch updateStatus {
|
||||
case errSecSuccess:
|
||||
return
|
||||
case errSecItemNotFound:
|
||||
var query = Self.baseQuery
|
||||
query[kSecValueData as String] = data
|
||||
// The extension can refresh after reboot without making account
|
||||
// credentials readable; this item contains only the limited grant.
|
||||
query[kSecAttrAccessible as String] = kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly
|
||||
let addStatus = SecItemAdd(query as CFDictionary, nil)
|
||||
guard addStatus == errSecSuccess else {
|
||||
throw StoreError.unexpectedStatus(addStatus)
|
||||
}
|
||||
default:
|
||||
throw StoreError.unexpectedStatus(updateStatus)
|
||||
}
|
||||
}
|
||||
|
||||
public func delete() async throws {
|
||||
let status = SecItemDelete(Self.baseQuery as CFDictionary)
|
||||
guard status == errSecSuccess || status == errSecItemNotFound else {
|
||||
throw StoreError.unexpectedStatus(status)
|
||||
}
|
||||
}
|
||||
|
||||
private static var baseQuery: [String: Any] {
|
||||
var query: [String: Any] = [
|
||||
kSecClass as String: kSecClassGenericPassword,
|
||||
kSecAttrService as String: service,
|
||||
kSecAttrAccount as String: account,
|
||||
kSecAttrSynchronizable as String: kCFBooleanFalse!
|
||||
]
|
||||
#if os(macOS)
|
||||
query[kSecUseDataProtectionKeychain as String] = true
|
||||
#endif
|
||||
return query
|
||||
}
|
||||
|
||||
private static var encoder: JSONEncoder {
|
||||
let encoder = JSONEncoder()
|
||||
encoder.dateEncodingStrategy = .iso8601
|
||||
return encoder
|
||||
}
|
||||
|
||||
private static var decoder: JSONDecoder {
|
||||
let decoder = JSONDecoder()
|
||||
decoder.dateDecodingStrategy = .iso8601
|
||||
return decoder
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
// ManagedGatewayModels.swift
|
||||
// OSGKeyboard · Shared
|
||||
//
|
||||
// Scope-limited credentials and transport models for the managed gateway.
|
||||
|
||||
import Foundation
|
||||
|
||||
public enum ManagedGatewayCapability: String, Codable, CaseIterable, Sendable {
|
||||
case polish
|
||||
case assistant = "ai"
|
||||
case agent
|
||||
case asr
|
||||
}
|
||||
|
||||
/// Stable task contract shared with the account gateway. The app identifies
|
||||
/// user intent; the server remains authoritative for model and feature policy.
|
||||
public enum ManagedGatewayTaskKind: String, Codable, CaseIterable, Sendable {
|
||||
case dictationPolish = "dictation_polish"
|
||||
case translation
|
||||
case editLastInput = "edit_last_input"
|
||||
case aiQuestion = "ai_question"
|
||||
case clipboardTransform = "clipboard_transform"
|
||||
case customSkill = "custom_skill"
|
||||
case agentPlanning = "agent_planning"
|
||||
}
|
||||
|
||||
public struct ManagedGatewayGrantCredentials: Codable, Equatable, Sendable {
|
||||
public static let maximumAccessLifetime: TimeInterval = 5 * 60
|
||||
|
||||
public let grantId: String
|
||||
public let scopes: Set<ManagedGatewayCapability>
|
||||
public let accessToken: String
|
||||
public let accessExpiresAt: Date
|
||||
public let refreshToken: String
|
||||
public let refreshExpiresAt: Date
|
||||
public let receivedAt: Date
|
||||
|
||||
public init(
|
||||
grantId: String,
|
||||
scopes: Set<ManagedGatewayCapability>,
|
||||
accessToken: String,
|
||||
accessExpiresAt: Date,
|
||||
refreshToken: String,
|
||||
refreshExpiresAt: Date,
|
||||
receivedAt: Date = Date()
|
||||
) {
|
||||
self.grantId = grantId
|
||||
self.scopes = scopes
|
||||
self.accessToken = accessToken
|
||||
self.accessExpiresAt = accessExpiresAt
|
||||
self.refreshToken = refreshToken
|
||||
self.refreshExpiresAt = refreshExpiresAt
|
||||
self.receivedAt = receivedAt
|
||||
}
|
||||
|
||||
/// Never trust an unexpectedly long access expiry. The gateway contract
|
||||
/// deliberately limits extension-readable bearer credentials to five minutes.
|
||||
public var effectiveAccessExpiresAt: Date {
|
||||
min(accessExpiresAt, receivedAt.addingTimeInterval(Self.maximumAccessLifetime))
|
||||
}
|
||||
|
||||
public func hasUsableAccessToken(
|
||||
for scope: ManagedGatewayCapability,
|
||||
at now: Date = Date(),
|
||||
refreshLeeway: TimeInterval = 30
|
||||
) -> Bool {
|
||||
scopes.contains(scope)
|
||||
&& !accessToken.isEmpty
|
||||
&& effectiveAccessExpiresAt.timeIntervalSince(now) > refreshLeeway
|
||||
}
|
||||
|
||||
public func hasUsableRefreshToken(at now: Date = Date()) -> Bool {
|
||||
!refreshToken.isEmpty && refreshExpiresAt > now
|
||||
}
|
||||
}
|
||||
|
||||
public enum ManagedGatewayError: Error, LocalizedError, Equatable, Sendable {
|
||||
case missingGrant
|
||||
case scopeNotGranted(ManagedGatewayCapability)
|
||||
case invalidGrant
|
||||
case insufficientCredits
|
||||
case timeout
|
||||
case server(code: String, status: Int, requestId: String?)
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .missingGrant:
|
||||
return SharedL10n.string("managed.error.grantUnavailable")
|
||||
case .scopeNotGranted(let scope):
|
||||
return SharedL10n.format(
|
||||
"managed.error.scopeNotGranted",
|
||||
language: nil,
|
||||
scope.rawValue
|
||||
)
|
||||
case .invalidGrant:
|
||||
return SharedL10n.string("managed.error.grantRejected")
|
||||
case .insufficientCredits:
|
||||
return SharedL10n.string("managed.error.insufficientCredits")
|
||||
case .timeout:
|
||||
return SharedL10n.string("managed.error.timeout")
|
||||
case .server(let code, let status, _):
|
||||
return SharedL10n.format(
|
||||
"managed.error.server",
|
||||
language: nil,
|
||||
code,
|
||||
status
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct ManagedGatewayGrantTokenResponse: Decodable {
|
||||
let grantId: String
|
||||
let scopes: Set<ManagedGatewayCapability>
|
||||
let accessToken: String
|
||||
let accessExpiresAt: Date
|
||||
let refreshToken: String
|
||||
let refreshExpiresAt: Date
|
||||
|
||||
func credentials(receivedAt: Date) -> ManagedGatewayGrantCredentials {
|
||||
ManagedGatewayGrantCredentials(
|
||||
grantId: grantId,
|
||||
scopes: scopes,
|
||||
accessToken: accessToken,
|
||||
accessExpiresAt: accessExpiresAt,
|
||||
refreshToken: refreshToken,
|
||||
refreshExpiresAt: refreshExpiresAt,
|
||||
receivedAt: receivedAt
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
struct ManagedGatewayErrorResponse: Decodable, Sendable {
|
||||
let code: String
|
||||
let message: String
|
||||
let requestId: String
|
||||
}
|
||||
|
||||
struct ManagedGatewayTextRequest: Encodable, Sendable {
|
||||
let input: String
|
||||
let context: String?
|
||||
let maxOutputTokens: Int
|
||||
let temperature: Double
|
||||
let stream: Bool
|
||||
let taskKind: ManagedGatewayTaskKind
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
// ManagedGatewayScopePolicy.swift
|
||||
// OSGKeyboard · Shared
|
||||
//
|
||||
// Computes the smallest grant needed by the currently enabled managed paths.
|
||||
|
||||
import Foundation
|
||||
|
||||
public enum ManagedGatewayScopePolicy {
|
||||
public static func scopes(engineMode: String) -> Set<ManagedGatewayCapability> {
|
||||
var scopes: Set<ManagedGatewayCapability> = [.polish, .assistant, .agent]
|
||||
if engineMode == "cloud" {
|
||||
scopes.insert(.asr)
|
||||
}
|
||||
return scopes
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,549 @@
|
||||
// ManagedLLMClient.swift
|
||||
// OSGKeyboard · Shared
|
||||
//
|
||||
// LLMClient implementation for scope-limited managed polish, AI and agent calls.
|
||||
|
||||
import Foundation
|
||||
|
||||
public struct ManagedLLMClient: LLMClient {
|
||||
public enum Capability: String, Sendable {
|
||||
case polish
|
||||
case assistant = "ai"
|
||||
case agent
|
||||
|
||||
var grantScope: ManagedGatewayCapability {
|
||||
switch self {
|
||||
case .polish: .polish
|
||||
case .assistant: .assistant
|
||||
case .agent: .agent
|
||||
}
|
||||
}
|
||||
|
||||
var defaultTaskKind: ManagedGatewayTaskKind {
|
||||
switch self {
|
||||
case .polish: .dictationPolish
|
||||
case .assistant: .aiQuestion
|
||||
case .agent: .agentPlanning
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private struct Attempt {
|
||||
let input: String
|
||||
let context: String?
|
||||
let timeout: TimeInterval?
|
||||
let options: LLMGenerationOptions
|
||||
let requestId: String
|
||||
let forceRefresh: Bool
|
||||
|
||||
func forcingRefresh() -> Self {
|
||||
Self(
|
||||
input: input,
|
||||
context: context,
|
||||
timeout: timeout,
|
||||
options: options,
|
||||
requestId: requestId,
|
||||
forceRefresh: true
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
public let capability: Capability
|
||||
public let taskKind: ManagedGatewayTaskKind
|
||||
public let requestTimeout: TimeInterval
|
||||
|
||||
private let baseURL: URL
|
||||
private let grants: GatewayGrantCoordinator
|
||||
private let session: URLSession
|
||||
private let requestId: @Sendable () -> String
|
||||
|
||||
public init(
|
||||
capability: Capability,
|
||||
taskKind: ManagedGatewayTaskKind? = nil,
|
||||
grants: GatewayGrantCoordinator,
|
||||
baseURL: URL = GatewayGrantCoordinator.defaultBaseURL,
|
||||
session: URLSession = .shared,
|
||||
requestTimeout: TimeInterval = 15,
|
||||
requestId: @escaping @Sendable () -> String = { UUID().uuidString }
|
||||
) {
|
||||
self.capability = capability
|
||||
self.taskKind = taskKind ?? capability.defaultTaskKind
|
||||
self.grants = grants
|
||||
self.baseURL = baseURL
|
||||
self.session = session
|
||||
self.requestTimeout = requestTimeout
|
||||
self.requestId = requestId
|
||||
}
|
||||
|
||||
public func polish(
|
||||
_ text: String,
|
||||
systemPrompt: String,
|
||||
timeout: TimeInterval?
|
||||
) async throws -> String {
|
||||
try await polish(
|
||||
text,
|
||||
systemPrompt: systemPrompt,
|
||||
timeout: timeout,
|
||||
options: .polishDefault
|
||||
)
|
||||
}
|
||||
|
||||
public func polish(
|
||||
_ text: String,
|
||||
systemPrompt: String,
|
||||
timeout: TimeInterval?,
|
||||
options: LLMGenerationOptions
|
||||
) async throws -> String {
|
||||
try await executeBuffered(
|
||||
input: text,
|
||||
context: systemPrompt.nilIfEmpty,
|
||||
timeout: timeout,
|
||||
options: options
|
||||
)
|
||||
}
|
||||
|
||||
public func complete(
|
||||
messages: [LLMRequest.Message],
|
||||
timeout: TimeInterval?,
|
||||
options: LLMGenerationOptions
|
||||
) async throws -> String {
|
||||
let payload = Self.payload(from: messages)
|
||||
return try await executeBuffered(
|
||||
input: payload.input,
|
||||
context: payload.context,
|
||||
timeout: timeout,
|
||||
options: options
|
||||
)
|
||||
}
|
||||
|
||||
public func completeStreaming(
|
||||
messages: [LLMRequest.Message],
|
||||
timeout: TimeInterval?,
|
||||
options: LLMGenerationOptions
|
||||
) -> AsyncThrowingStream<LLMStreamEvent, Error> {
|
||||
AsyncThrowingStream { continuation in
|
||||
let task = Task {
|
||||
do {
|
||||
let payload = Self.payload(from: messages)
|
||||
let logicalRequestId = requestId()
|
||||
let attempt = Attempt(
|
||||
input: payload.input,
|
||||
context: payload.context,
|
||||
timeout: timeout,
|
||||
options: options,
|
||||
requestId: logicalRequestId,
|
||||
forceRefresh: false
|
||||
)
|
||||
var emittedVisibleText = false
|
||||
do {
|
||||
try await streamAttempt(attempt) { chunk in
|
||||
emittedVisibleText = true
|
||||
continuation.yield(.delta(chunk))
|
||||
}
|
||||
} catch ManagedGatewayError.invalidGrant where !emittedVisibleText {
|
||||
try await streamAttempt(attempt.forcingRefresh()) { chunk in
|
||||
emittedVisibleText = true
|
||||
continuation.yield(.delta(chunk))
|
||||
}
|
||||
}
|
||||
continuation.finish()
|
||||
} catch is CancellationError {
|
||||
continuation.finish(throwing: LLMError.cancelled)
|
||||
} catch let error as URLError where error.code == .cancelled {
|
||||
continuation.finish(throwing: LLMError.cancelled)
|
||||
} catch let error as URLError where error.code == .timedOut {
|
||||
continuation.finish(throwing: ManagedGatewayError.timeout)
|
||||
} catch let error as LLMError {
|
||||
continuation.finish(throwing: error)
|
||||
} catch let error as ManagedGatewayError {
|
||||
continuation.finish(throwing: error)
|
||||
} catch {
|
||||
continuation.finish(throwing: LLMError.transport(String(describing: error)))
|
||||
}
|
||||
}
|
||||
continuation.onTermination = { _ in task.cancel() }
|
||||
}
|
||||
}
|
||||
|
||||
private func executeBuffered(
|
||||
input: String,
|
||||
context: String?,
|
||||
timeout: TimeInterval?,
|
||||
options: LLMGenerationOptions
|
||||
) async throws -> String {
|
||||
let logicalRequestId = requestId()
|
||||
let attempt = Attempt(
|
||||
input: input,
|
||||
context: context,
|
||||
timeout: timeout,
|
||||
options: options,
|
||||
requestId: logicalRequestId,
|
||||
forceRefresh: false
|
||||
)
|
||||
do {
|
||||
return try await bufferedAttempt(attempt)
|
||||
} catch ManagedGatewayError.invalidGrant {
|
||||
do {
|
||||
return try await bufferedAttempt(attempt.forcingRefresh())
|
||||
} catch ManagedGatewayError.invalidGrant {
|
||||
try? await grants.clearGrant()
|
||||
throw ManagedGatewayError.invalidGrant
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func bufferedAttempt(_ attempt: Attempt) async throws -> String {
|
||||
let request = try await makeRequest(
|
||||
attempt,
|
||||
stream: false
|
||||
)
|
||||
do {
|
||||
let (data, response) = try await session.data(for: request)
|
||||
guard let http = response as? HTTPURLResponse else {
|
||||
throw LLMError.transport("non-HTTP response")
|
||||
}
|
||||
guard (200..<300).contains(http.statusCode) else {
|
||||
throw ManagedGatewayHTTP.error(
|
||||
data: data,
|
||||
status: http.statusCode,
|
||||
requestId: attempt.requestId
|
||||
)
|
||||
}
|
||||
return try Self.responseText(from: data)
|
||||
} catch is CancellationError {
|
||||
throw LLMError.cancelled
|
||||
} catch let error as URLError where error.code == .cancelled {
|
||||
throw LLMError.cancelled
|
||||
} catch let error as URLError where error.code == .timedOut {
|
||||
throw ManagedGatewayError.timeout
|
||||
} catch let error as ManagedGatewayError {
|
||||
throw error
|
||||
} catch let error as LLMError {
|
||||
throw error
|
||||
} catch {
|
||||
throw LLMError.transport(String(describing: error))
|
||||
}
|
||||
}
|
||||
|
||||
private func streamAttempt(
|
||||
_ attempt: Attempt,
|
||||
onDelta: @escaping (String) -> Void
|
||||
) async throws {
|
||||
guard capability != .agent else {
|
||||
// The server validates agent output as one structured response.
|
||||
let text = try await bufferedAttempt(attempt)
|
||||
if !text.isEmpty {
|
||||
onDelta(text)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
let request = try await makeRequest(
|
||||
attempt,
|
||||
stream: true
|
||||
)
|
||||
for try await payload in ManagedGatewayStreamTransport.payloads(
|
||||
session: session,
|
||||
request: request,
|
||||
requestId: attempt.requestId
|
||||
) {
|
||||
try Task.checkCancellation()
|
||||
if let error = Self.streamingError(from: payload, requestId: attempt.requestId) {
|
||||
throw error
|
||||
}
|
||||
if let delta = Self.streamingDelta(from: payload), !delta.isEmpty {
|
||||
onDelta(delta)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func makeRequest(_ attempt: Attempt, stream: Bool) async throws -> URLRequest {
|
||||
let trimmedInput = attempt.input.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
guard !trimmedInput.isEmpty, trimmedInput.count <= 32_000 else {
|
||||
throw ManagedGatewayError.server(
|
||||
code: "invalid_request",
|
||||
status: 400,
|
||||
requestId: attempt.requestId
|
||||
)
|
||||
}
|
||||
let boundedContext = attempt.context?
|
||||
.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
.nilIfEmpty
|
||||
guard (boundedContext?.count ?? 0) <= 32_000 else {
|
||||
throw ManagedGatewayError.server(
|
||||
code: "invalid_request",
|
||||
status: 400,
|
||||
requestId: attempt.requestId
|
||||
)
|
||||
}
|
||||
|
||||
let token = try await grants.accessToken(
|
||||
for: capability.grantScope,
|
||||
forceRefresh: attempt.forceRefresh
|
||||
)
|
||||
let body = ManagedGatewayTextRequest(
|
||||
input: trimmedInput,
|
||||
context: boundedContext,
|
||||
maxOutputTokens: min(max(attempt.options.maxTokens ?? 512, 1), 4_096),
|
||||
temperature: min(max(attempt.options.temperature ?? 0.2, 0), 1),
|
||||
stream: stream,
|
||||
taskKind: taskKind
|
||||
)
|
||||
|
||||
var request = URLRequest(
|
||||
url: baseURL.appending(path: "v1/gateway/llm/\(capability.rawValue)")
|
||||
)
|
||||
request.httpMethod = "POST"
|
||||
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
||||
request.setValue(stream ? "text/event-stream" : "application/json", forHTTPHeaderField: "Accept")
|
||||
request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
|
||||
request.setValue(attempt.requestId, forHTTPHeaderField: "X-Request-ID")
|
||||
request.timeoutInterval = attempt.timeout ?? requestTimeout
|
||||
request.httpBody = try JSONEncoder().encode(body)
|
||||
return request
|
||||
}
|
||||
|
||||
static func payload(
|
||||
from messages: [LLMRequest.Message]
|
||||
) -> (input: String, context: String?) {
|
||||
guard let inputIndex = messages.lastIndex(where: { $0.role == "user" }) else {
|
||||
return ("", Self.contextText(from: messages))
|
||||
}
|
||||
let input = messages[inputIndex].content
|
||||
var contextMessages = messages
|
||||
contextMessages.remove(at: inputIndex)
|
||||
return (input, Self.contextText(from: contextMessages))
|
||||
}
|
||||
|
||||
private static func contextText(from messages: [LLMRequest.Message]) -> String? {
|
||||
messages
|
||||
.filter { !$0.content.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty }
|
||||
.map { "\($0.role):\n\($0.content)" }
|
||||
.joined(separator: "\n\n")
|
||||
.nilIfEmpty
|
||||
}
|
||||
|
||||
static func responseText(from data: Data) throws -> String {
|
||||
let raw = String(data: data, encoding: .utf8)?
|
||||
.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
|
||||
guard !raw.isEmpty else { throw LLMError.decoding("empty managed gateway response") }
|
||||
guard let json = try? JSONSerialization.jsonObject(with: data) else {
|
||||
return raw
|
||||
}
|
||||
return extractedText(from: json)?.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
.nilIfEmpty ?? raw
|
||||
}
|
||||
|
||||
private static func extractedText(from value: Any) -> String? {
|
||||
if let string = value as? String {
|
||||
return string
|
||||
}
|
||||
if let array = value as? [Any] {
|
||||
let values = array.compactMap(extractedText(from:))
|
||||
return values.isEmpty ? nil : values.joined()
|
||||
}
|
||||
guard let object = value as? [String: Any] else { return nil }
|
||||
|
||||
for key in ["output_text", "text"] {
|
||||
if let text = object[key] as? String, !text.isEmpty {
|
||||
return text
|
||||
}
|
||||
}
|
||||
if let content = object["content"] {
|
||||
if let text = content as? String, !text.isEmpty {
|
||||
return text
|
||||
}
|
||||
if let extracted = extractedText(from: content), !extracted.isEmpty {
|
||||
return extracted
|
||||
}
|
||||
}
|
||||
if let message = object["message"] as? [String: Any],
|
||||
let extracted = extractedText(from: message) {
|
||||
return extracted
|
||||
}
|
||||
if let choices = object["choices"] as? [[String: Any]],
|
||||
let first = choices.first {
|
||||
if let message = first["message"] as? [String: Any],
|
||||
let extracted = extractedText(from: message) {
|
||||
return extracted
|
||||
}
|
||||
if let text = first["text"] as? String {
|
||||
return text
|
||||
}
|
||||
}
|
||||
if let data = object["data"], let extracted = extractedText(from: data) {
|
||||
return extracted
|
||||
}
|
||||
if let output = object["output"], let extracted = extractedText(from: output) {
|
||||
return extracted
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
static func streamingDelta(from data: Data) -> String? {
|
||||
if let value = LLMStreamDeltaParser.responsesOutputTextDelta(from: data)
|
||||
?? LLMStreamDeltaParser.chatCompletionsDelta(from: data)
|
||||
?? LLMStreamDeltaParser.anthropicTextDelta(from: data) {
|
||||
return value
|
||||
}
|
||||
guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
|
||||
return String(data: data, encoding: .utf8)
|
||||
}
|
||||
for key in ["delta", "text", "content"] {
|
||||
if let value = json[key] as? String, !value.isEmpty {
|
||||
return value
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
private static func streamingError(
|
||||
from data: Data,
|
||||
requestId: String
|
||||
) -> ManagedGatewayError? {
|
||||
guard let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
|
||||
let code = object["code"] as? String,
|
||||
object["message"] != nil || code.hasSuffix("_error") else {
|
||||
return nil
|
||||
}
|
||||
if ["insufficient_credits", "insufficient_balance"].contains(code) {
|
||||
return .insufficientCredits
|
||||
}
|
||||
if ["unauthorized", "gateway_grant_denied", "invalid_grant"].contains(code) {
|
||||
return .invalidGrant
|
||||
}
|
||||
return .server(code: code, status: 200, requestId: requestId)
|
||||
}
|
||||
}
|
||||
|
||||
public enum ManagedGatewayLLMClientFactory {
|
||||
public static func polish(
|
||||
taskKind: ManagedGatewayTaskKind = .dictationPolish,
|
||||
grants: GatewayGrantCoordinator,
|
||||
baseURL: URL = GatewayGrantCoordinator.defaultBaseURL,
|
||||
session: URLSession = .shared
|
||||
) -> any LLMClient {
|
||||
ManagedLLMClient(
|
||||
capability: .polish,
|
||||
taskKind: taskKind,
|
||||
grants: grants,
|
||||
baseURL: baseURL,
|
||||
session: session
|
||||
)
|
||||
}
|
||||
|
||||
public static func ai(
|
||||
taskKind: ManagedGatewayTaskKind = .aiQuestion,
|
||||
grants: GatewayGrantCoordinator,
|
||||
baseURL: URL = GatewayGrantCoordinator.defaultBaseURL,
|
||||
session: URLSession = .shared
|
||||
) -> any LLMClient {
|
||||
ManagedLLMClient(
|
||||
capability: .assistant,
|
||||
taskKind: taskKind,
|
||||
grants: grants,
|
||||
baseURL: baseURL,
|
||||
session: session
|
||||
)
|
||||
}
|
||||
|
||||
public static func agent(
|
||||
taskKind: ManagedGatewayTaskKind = .agentPlanning,
|
||||
grants: GatewayGrantCoordinator,
|
||||
baseURL: URL = GatewayGrantCoordinator.defaultBaseURL,
|
||||
session: URLSession = .shared
|
||||
) -> any LLMClient {
|
||||
ManagedLLMClient(
|
||||
capability: .agent,
|
||||
taskKind: taskKind,
|
||||
grants: grants,
|
||||
baseURL: baseURL,
|
||||
session: session
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
private extension String {
|
||||
var nilIfEmpty: String? {
|
||||
isEmpty ? nil : self
|
||||
}
|
||||
}
|
||||
|
||||
enum ManagedGatewayStreamTransport {
|
||||
static func payloads(
|
||||
session: URLSession,
|
||||
request: URLRequest,
|
||||
requestId: String
|
||||
) -> AsyncThrowingStream<Data, Error> {
|
||||
AsyncThrowingStream { continuation in
|
||||
let task = Task {
|
||||
do {
|
||||
let (bytes, response) = try await session.bytes(for: request)
|
||||
guard let http = response as? HTTPURLResponse else {
|
||||
throw LLMError.transport("non-HTTP response")
|
||||
}
|
||||
guard (200..<300).contains(http.statusCode) else {
|
||||
var body = Data()
|
||||
for try await byte in bytes {
|
||||
body.append(byte)
|
||||
}
|
||||
throw ManagedGatewayHTTP.error(
|
||||
data: body,
|
||||
status: http.statusCode,
|
||||
requestId: requestId
|
||||
)
|
||||
}
|
||||
|
||||
let isEventStream = http.value(forHTTPHeaderField: "Content-Type")?
|
||||
.lowercased()
|
||||
.contains("text/event-stream") == true
|
||||
if !isEventStream {
|
||||
var body = Data()
|
||||
for try await byte in bytes {
|
||||
try Task.checkCancellation()
|
||||
body.append(byte)
|
||||
}
|
||||
if !body.isEmpty {
|
||||
continuation.yield(body)
|
||||
}
|
||||
continuation.finish()
|
||||
return
|
||||
}
|
||||
|
||||
var line = Data()
|
||||
for try await byte in bytes {
|
||||
try Task.checkCancellation()
|
||||
if byte == UInt8(ascii: "\n") {
|
||||
yieldSSELine(line, to: continuation)
|
||||
line.removeAll(keepingCapacity: true)
|
||||
} else if byte != UInt8(ascii: "\r") {
|
||||
line.append(byte)
|
||||
}
|
||||
}
|
||||
yieldSSELine(line, to: continuation)
|
||||
continuation.finish()
|
||||
} catch is CancellationError {
|
||||
continuation.finish(throwing: LLMError.cancelled)
|
||||
} catch let error as URLError where error.code == .cancelled {
|
||||
continuation.finish(throwing: LLMError.cancelled)
|
||||
} catch let error as URLError where error.code == .timedOut {
|
||||
continuation.finish(throwing: ManagedGatewayError.timeout)
|
||||
} catch {
|
||||
continuation.finish(throwing: error)
|
||||
}
|
||||
}
|
||||
continuation.onTermination = { _ in task.cancel() }
|
||||
}
|
||||
}
|
||||
|
||||
private static func yieldSSELine(
|
||||
_ line: Data,
|
||||
to continuation: AsyncThrowingStream<Data, Error>.Continuation
|
||||
) {
|
||||
guard let payload = LLMStreamTransport.sseDataPayload(fromLineBytes: line),
|
||||
payload != Data("[DONE]".utf8) else {
|
||||
return
|
||||
}
|
||||
continuation.yield(payload)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user