Rocky
2cc81c4628
Add local clipboard intent recommendations, webpage and phone actions, and safer host handoffs. Refine managed gateway, catalog refresh, onboarding, and adaptive polish behavior.
858 lines
32 KiB
Swift
858 lines
32 KiB
Swift
// ManagedGatewayTests.swift
|
|
// OSGKeyboard · Extension Tests
|
|
//
|
|
// Hermetic grant rotation and managed LLM tests using a private URLProtocol stub.
|
|
|
|
import Foundation
|
|
@testable import OSGKeyboardShared
|
|
import XCTest
|
|
|
|
final class ManagedGatewayTests: XCTestCase {
|
|
private let baseURL = URL(string: "https://gateway.test")!
|
|
|
|
override func setUp() {
|
|
super.setUp()
|
|
GatewayStub.shared.reset()
|
|
}
|
|
|
|
override func tearDown() {
|
|
GatewayStub.shared.reset()
|
|
super.tearDown()
|
|
}
|
|
|
|
func testFiveMinuteAccessBoundaryAndScope() {
|
|
let start = Date(timeIntervalSince1970: 1_000)
|
|
let value = credentials(accessExpiresAt: start.addingTimeInterval(3_600), receivedAt: start)
|
|
XCTAssertEqual(value.effectiveAccessExpiresAt, start.addingTimeInterval(300))
|
|
XCTAssertTrue(value.hasUsableAccessToken(for: .polish, at: start.addingTimeInterval(269)))
|
|
XCTAssertFalse(value.hasUsableAccessToken(for: .polish, at: start.addingTimeInterval(270)))
|
|
XCTAssertFalse(value.hasUsableAccessToken(for: .asr, at: start))
|
|
}
|
|
|
|
func testScopePolicyExcludesASRForLocalEngineAndIncludesItForCloud() {
|
|
let localScopes = ManagedGatewayScopePolicy.scopes(engineMode: "local")
|
|
XCTAssertEqual(localScopes, [.polish, .assistant, .agent])
|
|
XCTAssertFalse(localScopes.contains(.asr))
|
|
|
|
let cloudScopes = ManagedGatewayScopePolicy.scopes(engineMode: "cloud")
|
|
XCTAssertEqual(cloudScopes, [.polish, .assistant, .agent, .asr])
|
|
}
|
|
|
|
func testOOBEFeatureConflictIsNotMappedToNetworkFailure() {
|
|
let error = ManagedGatewayHTTP.error(
|
|
data: errorJSON("oobe_feature_already_used"),
|
|
status: 409,
|
|
requestId: "oobe-conflict"
|
|
)
|
|
|
|
XCTAssertEqual(error, .oobeFeatureAlreadyUsed)
|
|
}
|
|
|
|
func testCreateGrantSendsAccountTokenButStoresOnlyGrant() async throws {
|
|
GatewayStub.shared.enqueue(
|
|
201,
|
|
grantJSON(
|
|
access: "grant-access",
|
|
refresh: "grant-refresh",
|
|
scopes: ["polish", "ai"]
|
|
)
|
|
)
|
|
let store = MemoryGrantStore()
|
|
let coordinator = makeCoordinator(store: store)
|
|
|
|
_ = try await coordinator.createGrant(
|
|
accountAccessToken: "account-secret",
|
|
scopes: [.polish, .assistant],
|
|
lifetimeSeconds: 3_600,
|
|
idempotencyKey: "create-idempotency"
|
|
)
|
|
|
|
let stored = await store.value()
|
|
XCTAssertEqual(stored?.refreshToken, "grant-refresh")
|
|
let request = try XCTUnwrap(GatewayStub.shared.requests().first)
|
|
XCTAssertEqual(request.url?.path, "/v1/gateway/grants")
|
|
XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Bearer account-secret")
|
|
XCTAssertEqual(request.value(forHTTPHeaderField: "Idempotency-Key"), "create-idempotency")
|
|
let body = try jsonBody(request)
|
|
XCTAssertEqual(Set(body["scopes"] as? [String] ?? []), Set(["polish", "ai"]))
|
|
XCTAssertEqual(body["lifetimeSeconds"] as? Int, 3_600)
|
|
let persisted = try JSONEncoder().encode(stored)
|
|
XCTAssertFalse(
|
|
(String(bytes: persisted, encoding: .utf8) ?? "").contains("account-secret")
|
|
)
|
|
}
|
|
|
|
func testCreateGrantRejectsBroaderReturnedScopes() async throws {
|
|
GatewayStub.shared.enqueue(
|
|
201,
|
|
grantJSON(
|
|
access: "grant-access",
|
|
refresh: "grant-refresh",
|
|
scopes: ["polish", "ai", "agent"]
|
|
)
|
|
)
|
|
let store = MemoryGrantStore()
|
|
let coordinator = makeCoordinator(store: store)
|
|
|
|
await XCTAssertThrowsManaged(.invalidGrant) {
|
|
_ = try await coordinator.createGrant(
|
|
accountAccessToken: "account-secret",
|
|
scopes: [.polish],
|
|
idempotencyKey: "create-scope-boundary"
|
|
)
|
|
}
|
|
let stored = await store.value()
|
|
XCTAssertNil(stored)
|
|
}
|
|
|
|
func testConcurrentRefreshIsMergedAndHashesOldTokenForIdempotency() async throws {
|
|
let now = Date(timeIntervalSince1970: 2_000)
|
|
let oldRefresh = "old-refresh-token-that-is-at-least-32-characters"
|
|
let store = MemoryGrantStore(
|
|
credentials(
|
|
accessToken: "expired",
|
|
refreshToken: oldRefresh,
|
|
accessExpiresAt: now.addingTimeInterval(-1),
|
|
receivedAt: now.addingTimeInterval(-600)
|
|
)
|
|
)
|
|
GatewayStub.shared.enqueue(
|
|
200,
|
|
grantJSON(access: "new-access", refresh: "new-refresh"),
|
|
delay: 0.05
|
|
)
|
|
let coordinator = makeCoordinator(store: store, now: { now })
|
|
|
|
async let first = coordinator.accessToken(for: .assistant)
|
|
async let second = coordinator.accessToken(for: .assistant)
|
|
let values = try await (first, second)
|
|
|
|
XCTAssertEqual(values.0, "new-access")
|
|
XCTAssertEqual(values.1, "new-access")
|
|
let requests = GatewayStub.shared.requests()
|
|
XCTAssertEqual(requests.count, 1)
|
|
let request = try XCTUnwrap(requests.first)
|
|
let key = try XCTUnwrap(request.value(forHTTPHeaderField: "Idempotency-Key"))
|
|
XCTAssertEqual(key, GatewayGrantCoordinator.refreshTokenIdempotencyKey(oldRefresh))
|
|
XCTAssertFalse(key.contains(oldRefresh))
|
|
XCTAssertEqual(try jsonBody(request)["refreshToken"] as? String, oldRefresh)
|
|
}
|
|
|
|
func testSignedOutPolicyRejectsCachedAccountGrant() async throws {
|
|
let suiteName = "ManagedGatewayTests.signedOut.\(UUID().uuidString)"
|
|
let defaults = try XCTUnwrap(UserDefaults(suiteName: suiteName))
|
|
defer { defaults.removePersistentDomain(forName: suiteName) }
|
|
let now = Date(timeIntervalSince1970: 3_000)
|
|
let store = MemoryGrantStore(
|
|
credentials(
|
|
accessToken: "cached-account-access",
|
|
accessExpiresAt: now.addingTimeInterval(240),
|
|
refreshExpiresAt: now.addingTimeInterval(3_600),
|
|
receivedAt: now
|
|
)
|
|
)
|
|
let accountState = AppGroupStore(defaults: defaults)
|
|
let coordinator = GatewayGrantCoordinator(
|
|
baseURL: baseURL,
|
|
store: store,
|
|
session: stubSession(),
|
|
now: { now },
|
|
accountAccessPolicy: AppGroupManagedGatewayAccountAccessPolicy(
|
|
defaults: defaults
|
|
)
|
|
)
|
|
|
|
await XCTAssertThrowsManaged(.missingGrant) {
|
|
_ = try await coordinator.accessToken(for: .polish)
|
|
}
|
|
XCTAssertTrue(GatewayStub.shared.requests().isEmpty)
|
|
|
|
accountState.setManagedGatewayAccountSessionAvailable(true)
|
|
let token = try await coordinator.accessToken(for: .polish)
|
|
XCTAssertEqual(token, "cached-account-access")
|
|
}
|
|
|
|
func testManagedClientMapsAllCapabilitiesHeadersBodyAndRequestIds() async throws {
|
|
let now = Date(timeIntervalSince1970: 3_000)
|
|
let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
|
|
let coordinator = makeCoordinator(store: store, now: { now })
|
|
let ids = LockedValues(["request-0001", "request-0002", "request-0003"])
|
|
|
|
let cases: [(ManagedLLMClient.Capability, ManagedGatewayTaskKind)] = [
|
|
(.polish, .dictationPolish),
|
|
(.assistant, .aiQuestion),
|
|
(.agent, .agentPlanning)
|
|
]
|
|
for (capability, _) in cases {
|
|
GatewayStub.shared.enqueue(200, Data(#"{"choices":[{"message":{"content":"ok"}}]}"#.utf8))
|
|
let client = ManagedLLMClient(
|
|
capability: capability,
|
|
grants: coordinator,
|
|
baseURL: baseURL,
|
|
session: stubSession(),
|
|
requestId: { ids.next() }
|
|
)
|
|
let result = try await client.polish(
|
|
"input",
|
|
systemPrompt: "context",
|
|
timeout: 7,
|
|
options: LLMGenerationOptions(temperature: 0.4, maxTokens: 321)
|
|
)
|
|
XCTAssertEqual(result, "ok")
|
|
}
|
|
|
|
let requests = GatewayStub.shared.requests()
|
|
XCTAssertEqual(
|
|
requests.compactMap(\.url?.path),
|
|
[
|
|
"/v1/gateway/llm/polish",
|
|
"/v1/gateway/llm/ai",
|
|
"/v1/gateway/llm/agent"
|
|
]
|
|
)
|
|
XCTAssertEqual(Set(requests.compactMap { $0.value(forHTTPHeaderField: "X-Request-ID") }).count, 3)
|
|
for (request, expectedCase) in zip(requests, cases) {
|
|
XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Bearer access")
|
|
let body = try jsonBody(request)
|
|
XCTAssertEqual(body["input"] as? String, "input")
|
|
XCTAssertEqual(body["context"] as? String, "context")
|
|
XCTAssertEqual(body["maxOutputTokens"] as? Int, 321)
|
|
XCTAssertEqual(body["temperature"] as? Double, 0.4)
|
|
XCTAssertEqual(body["stream"] as? Bool, false)
|
|
XCTAssertEqual(body["taskKind"] as? String, expectedCase.1.rawValue)
|
|
XCTAssertEqual(request.timeoutInterval, 7)
|
|
}
|
|
}
|
|
|
|
func testManagedClientSerializesFineGrainedTaskKinds() async throws {
|
|
let now = Date(timeIntervalSince1970: 3_500)
|
|
let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
|
|
let coordinator = makeCoordinator(store: store, now: { now })
|
|
let cases: [(ManagedLLMClient.Capability, ManagedGatewayTaskKind)] = [
|
|
(.polish, .translation),
|
|
(.polish, .editLastInput),
|
|
(.assistant, .currentInformationQuestion),
|
|
(.assistant, .clipboardTransform),
|
|
(.assistant, .customSkill)
|
|
]
|
|
|
|
for (capability, taskKind) in cases {
|
|
GatewayStub.shared.enqueue(200, Data(#"{"output_text":"ok"}"#.utf8))
|
|
let client = ManagedLLMClient(
|
|
capability: capability,
|
|
taskKind: taskKind,
|
|
grants: coordinator,
|
|
baseURL: baseURL,
|
|
session: stubSession()
|
|
)
|
|
_ = try await client.polish("input", systemPrompt: "context")
|
|
}
|
|
|
|
let taskKinds = try GatewayStub.shared.requests().map {
|
|
try XCTUnwrap(jsonBody($0)["taskKind"] as? String)
|
|
}
|
|
XCTAssertEqual(taskKinds, cases.map { $0.1.rawValue })
|
|
}
|
|
|
|
func testManagedHotwordSerializesSourceAndCurrentInformationIntent() async throws {
|
|
let now = Date(timeIntervalSince1970: 3_600)
|
|
let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
|
|
GatewayStub.shared.enqueue(200, Data(#"{"output_text":"ok"}"#.utf8))
|
|
let client = ManagedLLMClient(
|
|
capability: .assistant,
|
|
taskKind: .currentInformationQuestion,
|
|
requestSource: .hotword,
|
|
grants: makeCoordinator(store: store, now: { now }),
|
|
baseURL: baseURL,
|
|
session: stubSession()
|
|
)
|
|
|
|
_ = try await client.polish("today's news", systemPrompt: "context")
|
|
|
|
let request = try XCTUnwrap(GatewayStub.shared.requests().first)
|
|
let body = try jsonBody(request)
|
|
XCTAssertEqual(body["taskKind"] as? String, "current_information_question")
|
|
XCTAssertEqual(body["requestSource"] as? String, "hotword")
|
|
}
|
|
|
|
func testManagedClientSerializesOOBERequestPurpose() async throws {
|
|
let now = Date(timeIntervalSince1970: 3_750)
|
|
let store = MemoryGrantStore(credentials(accessToken: "account-access", receivedAt: now))
|
|
let oobeStore = MemoryGrantStore(
|
|
credentials(
|
|
accessToken: "oobe-access",
|
|
receivedAt: now,
|
|
scopes: [.polish, .assistant]
|
|
)
|
|
)
|
|
let suite = "managed.oobe.\(UUID().uuidString)"
|
|
let defaults = try XCTUnwrap(UserDefaults(suiteName: suite))
|
|
defer { defaults.removePersistentDomain(forName: suite) }
|
|
let practice = try XCTUnwrap(
|
|
KeyboardSetupBridge.beginOOBEPracticeSession(
|
|
expectedFeature: .voiceInput,
|
|
defaults: defaults,
|
|
now: now
|
|
)
|
|
)
|
|
XCTAssertEqual(practice.expectedFeature, .voiceInput)
|
|
GatewayStub.shared.enqueue(200, Data(#"{"output_text":"ok"}"#.utf8))
|
|
let client = ManagedLLMClient(
|
|
capability: .polish,
|
|
taskKind: .dictationPolish,
|
|
requestPurpose: .oobe,
|
|
oobeFeature: .voiceInput,
|
|
grants: makeCoordinator(store: store, now: { now }),
|
|
oobeGrants: OOBEGatewayGrantCoordinator(
|
|
baseURL: baseURL,
|
|
store: oobeStore,
|
|
session: stubSession(),
|
|
now: { now },
|
|
practiceSession: { practice }
|
|
),
|
|
baseURL: baseURL,
|
|
session: stubSession()
|
|
)
|
|
|
|
_ = try await client.polish("input", systemPrompt: "context")
|
|
|
|
let request = try XCTUnwrap(GatewayStub.shared.requests().last)
|
|
let body = try jsonBody(request)
|
|
XCTAssertEqual(body["taskKind"] as? String, "dictation_polish")
|
|
XCTAssertEqual(body["requestPurpose"] as? String, "oobe")
|
|
XCTAssertEqual(body["oobeFeature"] as? String, "voice_input")
|
|
XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Bearer oobe-access")
|
|
let accountCredentials = await store.value()
|
|
XCTAssertEqual(accountCredentials?.accessToken, "account-access")
|
|
}
|
|
|
|
func testOOBEGrantRejectsAgentOrASRScopes() async {
|
|
let now = Date(timeIntervalSince1970: 3_775)
|
|
let coordinator = OOBEGatewayGrantCoordinator(
|
|
baseURL: baseURL,
|
|
store: MemoryGrantStore(),
|
|
session: stubSession(),
|
|
now: { now }
|
|
)
|
|
|
|
await XCTAssertThrowsManaged(.invalidGrant) {
|
|
try await coordinator.install(
|
|
self.credentials(
|
|
receivedAt: now,
|
|
scopes: [.polish, .assistant, .agent]
|
|
)
|
|
)
|
|
}
|
|
await XCTAssertThrowsManaged(.invalidGrant) {
|
|
try await coordinator.install(
|
|
self.credentials(
|
|
receivedAt: now,
|
|
scopes: [.polish, .assistant, .asr]
|
|
)
|
|
)
|
|
}
|
|
}
|
|
|
|
func testOOBERefreshUsesDedicatedEndpoint() async throws {
|
|
let now = Date(timeIntervalSince1970: 3_790)
|
|
let practice = OOBEPracticeSession(
|
|
sessionID: UUID(),
|
|
expectedFeature: .voiceInput,
|
|
startedAt: now.addingTimeInterval(-10),
|
|
expiresAt: now.addingTimeInterval(300)
|
|
)
|
|
let store = MemoryGrantStore(
|
|
credentials(
|
|
accessToken: "expired",
|
|
accessExpiresAt: now.addingTimeInterval(-1),
|
|
receivedAt: now.addingTimeInterval(-600),
|
|
scopes: [.polish, .assistant]
|
|
)
|
|
)
|
|
GatewayStub.shared.enqueue(
|
|
200,
|
|
grantJSON(
|
|
access: "refreshed",
|
|
refresh: "rotated",
|
|
scopes: ["polish", "ai"]
|
|
)
|
|
)
|
|
let coordinator = OOBEGatewayGrantCoordinator(
|
|
baseURL: baseURL,
|
|
store: store,
|
|
session: stubSession(),
|
|
now: { now },
|
|
practiceSession: { practice }
|
|
)
|
|
|
|
let token = try await coordinator.accessToken(
|
|
for: .polish,
|
|
feature: .voiceInput
|
|
)
|
|
|
|
XCTAssertEqual(token, "refreshed")
|
|
XCTAssertEqual(
|
|
GatewayStub.shared.requests().first?.url?.path,
|
|
"/v1/oobe/grants/refresh"
|
|
)
|
|
}
|
|
|
|
func testOOBEClientRejectsMissingFeatureWithoutUsingAccountGrant() async {
|
|
let now = Date(timeIntervalSince1970: 3_800)
|
|
let store = MemoryGrantStore(credentials(accessToken: "account-access", receivedAt: now))
|
|
let client = ManagedLLMClient(
|
|
capability: .polish,
|
|
requestPurpose: .oobe,
|
|
grants: makeCoordinator(store: store, now: { now }),
|
|
baseURL: baseURL,
|
|
session: stubSession()
|
|
)
|
|
|
|
await XCTAssertThrowsManaged(
|
|
.server(code: "missing_oobe_feature", status: 400, requestId: "fixed")
|
|
) {
|
|
_ = try await ManagedLLMClient(
|
|
capability: client.capability,
|
|
requestPurpose: client.requestPurpose,
|
|
grants: makeCoordinator(store: store, now: { now }),
|
|
baseURL: self.baseURL,
|
|
session: self.stubSession(),
|
|
requestId: { "fixed" }
|
|
).polish("input", systemPrompt: "context")
|
|
}
|
|
XCTAssertTrue(GatewayStub.shared.requests().isEmpty)
|
|
}
|
|
|
|
func testManagedTaskKindWireValuesMatchServerContract() {
|
|
XCTAssertEqual(
|
|
ManagedGatewayTaskKind.allCases.map(\.rawValue),
|
|
[
|
|
"dictation_polish",
|
|
"translation",
|
|
"edit_last_input",
|
|
"ai_question",
|
|
"current_information_question",
|
|
"clipboard_transform",
|
|
"custom_skill",
|
|
"agent_planning"
|
|
]
|
|
)
|
|
}
|
|
|
|
func testConversationMappingAndRawTextFallback() async throws {
|
|
let now = Date(timeIntervalSince1970: 4_000)
|
|
let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
|
|
GatewayStub.shared.enqueue(200, Data("raw answer".utf8), contentType: "text/plain")
|
|
let client = ManagedLLMClient(
|
|
capability: .assistant,
|
|
grants: makeCoordinator(store: store, now: { now }),
|
|
baseURL: baseURL,
|
|
session: stubSession()
|
|
)
|
|
|
|
let answer = try await client.complete(
|
|
messages: [.system("system"), .user("old"), .assistant("reply"), .user("latest")],
|
|
timeout: nil,
|
|
options: .polishDefault
|
|
)
|
|
|
|
XCTAssertEqual(answer, "raw answer")
|
|
let request = try XCTUnwrap(GatewayStub.shared.requests().first)
|
|
let body = try jsonBody(request)
|
|
XCTAssertEqual(body["input"] as? String, "latest")
|
|
XCTAssertEqual(body["context"] as? String, "system:\nsystem\n\nuser:\nold\n\nassistant:\nreply")
|
|
}
|
|
|
|
func testUnauthorizedRequestRefreshesOnceWithSameRequestId() async throws {
|
|
let now = Date(timeIntervalSince1970: 5_000)
|
|
let oldRefresh = "old-refresh-value-123456789012345678901"
|
|
let store = MemoryGrantStore(
|
|
credentials(accessToken: "old-access", refreshToken: oldRefresh, receivedAt: now)
|
|
)
|
|
GatewayStub.shared.enqueue(401, errorJSON("unauthorized"))
|
|
GatewayStub.shared.enqueue(200, grantJSON(access: "new-access", refresh: "new-refresh"))
|
|
GatewayStub.shared.enqueue(200, Data(#"{"output_text":"done"}"#.utf8))
|
|
let client = ManagedLLMClient(
|
|
capability: .polish,
|
|
grants: makeCoordinator(store: store, now: { now }),
|
|
baseURL: baseURL,
|
|
session: stubSession(),
|
|
requestId: { "logical-request" }
|
|
)
|
|
|
|
let result = try await client.polish("input", systemPrompt: "prompt")
|
|
XCTAssertEqual(result, "done")
|
|
let requests = GatewayStub.shared.requests()
|
|
XCTAssertEqual(requests.count, 3)
|
|
XCTAssertEqual(requests[0].value(forHTTPHeaderField: "Authorization"), "Bearer old-access")
|
|
XCTAssertEqual(requests[1].url?.path, "/v1/gateway/grants/refresh")
|
|
XCTAssertEqual(requests[2].value(forHTTPHeaderField: "Authorization"), "Bearer new-access")
|
|
XCTAssertEqual(
|
|
requests[0].value(forHTTPHeaderField: "X-Request-ID"),
|
|
requests[2].value(forHTTPHeaderField: "X-Request-ID")
|
|
)
|
|
XCTAssertEqual(
|
|
requests[1].value(forHTTPHeaderField: "Idempotency-Key"),
|
|
GatewayGrantCoordinator.refreshTokenIdempotencyKey(oldRefresh)
|
|
)
|
|
}
|
|
|
|
func testCreditAndInvalidRefreshErrorsAreMapped() async throws {
|
|
let now = Date(timeIntervalSince1970: 6_000)
|
|
let creditStore = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
|
|
GatewayStub.shared.enqueue(402, errorJSON("insufficient_credits"))
|
|
let creditClient = ManagedLLMClient(
|
|
capability: .assistant,
|
|
grants: makeCoordinator(store: creditStore, now: { now }),
|
|
baseURL: baseURL,
|
|
session: stubSession()
|
|
)
|
|
await XCTAssertThrowsManaged(.insufficientCredits) {
|
|
_ = try await creditClient.complete(
|
|
messages: [.user("question")],
|
|
timeout: nil,
|
|
options: .polishDefault
|
|
)
|
|
}
|
|
|
|
let expiredStore = MemoryGrantStore(
|
|
credentials(
|
|
accessToken: "expired",
|
|
accessExpiresAt: now.addingTimeInterval(-1),
|
|
receivedAt: now.addingTimeInterval(-600)
|
|
)
|
|
)
|
|
GatewayStub.shared.enqueue(401, errorJSON("invalid_gateway_refresh"))
|
|
let coordinator = makeCoordinator(store: expiredStore, now: { now })
|
|
await XCTAssertThrowsManaged(.invalidGrant) {
|
|
_ = try await coordinator.accessToken(for: .assistant)
|
|
}
|
|
let remaining = await expiredStore.value()
|
|
XCTAssertNil(remaining)
|
|
}
|
|
|
|
func testProviderFailuresAreMappedToDistinctManagedErrors() {
|
|
XCTAssertEqual(
|
|
ManagedGatewayHTTP.error(
|
|
data: errorJSON("provider_unavailable"),
|
|
status: 503,
|
|
requestId: "provider-unavailable"
|
|
),
|
|
.providerUnavailable(requestId: "request-id")
|
|
)
|
|
XCTAssertEqual(
|
|
ManagedGatewayHTTP.error(
|
|
data: errorJSON("provider_rate_limited"),
|
|
status: 503,
|
|
requestId: "provider-rate-limited"
|
|
),
|
|
.providerRateLimited(requestId: "request-id")
|
|
)
|
|
XCTAssertEqual(
|
|
ManagedGatewayHTTP.error(
|
|
data: errorJSON("provider_timeout"),
|
|
status: 504,
|
|
requestId: "provider-timeout"
|
|
),
|
|
.providerTimeout(requestId: "request-id")
|
|
)
|
|
XCTAssertEqual(
|
|
ManagedGatewayHTTP.error(
|
|
data: errorJSON("provider_invalid_response"),
|
|
status: 502,
|
|
requestId: "provider-invalid"
|
|
),
|
|
.providerFailure(requestId: "request-id")
|
|
)
|
|
XCTAssertEqual(
|
|
ManagedGatewayHTTP.error(
|
|
data: errorJSON("internal_failure"),
|
|
status: 500,
|
|
requestId: "internal-failure"
|
|
),
|
|
.internalFailure(requestId: "request-id")
|
|
)
|
|
}
|
|
|
|
func testTransportTimeoutMapsToManagedTimeout() async {
|
|
let now = Date(timeIntervalSince1970: 6_500)
|
|
let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
|
|
GatewayStub.shared.enqueue(
|
|
0,
|
|
Data(),
|
|
transportError: .timedOut
|
|
)
|
|
let client = ManagedLLMClient(
|
|
capability: .polish,
|
|
grants: makeCoordinator(store: store, now: { now }),
|
|
baseURL: baseURL,
|
|
session: stubSession()
|
|
)
|
|
|
|
await XCTAssertThrowsManaged(.timeout) {
|
|
_ = try await client.polish("input", systemPrompt: "prompt", timeout: 0.1)
|
|
}
|
|
}
|
|
|
|
func testCancellationAndStreaming() async throws {
|
|
let now = Date(timeIntervalSince1970: 7_000)
|
|
let store = MemoryGrantStore(credentials(accessToken: "access", receivedAt: now))
|
|
let coordinator = makeCoordinator(store: store, now: { now })
|
|
GatewayStub.shared.enqueue(200, Data(#"{"output_text":"late"}"#.utf8), delay: 5)
|
|
let cancelledClient = ManagedLLMClient(
|
|
capability: .polish,
|
|
grants: coordinator,
|
|
baseURL: baseURL,
|
|
session: stubSession()
|
|
)
|
|
let task = Task { try await cancelledClient.polish("input", systemPrompt: "prompt") }
|
|
try await Task.sleep(for: .milliseconds(50))
|
|
task.cancel()
|
|
do {
|
|
_ = try await task.value
|
|
XCTFail("expected cancellation")
|
|
} catch let error as LLMError {
|
|
XCTAssertEqual(error, .cancelled)
|
|
} catch is CancellationError {
|
|
// Foundation may surface native cancellation before client mapping.
|
|
}
|
|
|
|
let sse = """
|
|
data: {"choices":[{"delta":{"reasoning_content":"hidden","content":"Hello"}}]}
|
|
|
|
data: {"type":"response.output_text.delta","delta":" world"}
|
|
|
|
data: [DONE]
|
|
|
|
"""
|
|
GatewayStub.shared.enqueue(200, Data(sse.utf8), contentType: "text/event-stream")
|
|
let streamClient = ManagedLLMClient(
|
|
capability: .assistant,
|
|
grants: coordinator,
|
|
baseURL: baseURL,
|
|
session: stubSession()
|
|
)
|
|
var answer = ""
|
|
for try await event in streamClient.completeStreaming(
|
|
messages: [.user("question")],
|
|
timeout: 2,
|
|
options: .polishDefault
|
|
) {
|
|
if case .delta(let value) = event { answer += value }
|
|
}
|
|
XCTAssertEqual(answer, "Hello world")
|
|
let streamRequest = try XCTUnwrap(GatewayStub.shared.requests().last)
|
|
XCTAssertEqual(try jsonBody(streamRequest)["stream"] as? Bool, true)
|
|
}
|
|
|
|
private func makeCoordinator(
|
|
store: MemoryGrantStore,
|
|
now: @escaping @Sendable () -> Date = Date.init
|
|
) -> GatewayGrantCoordinator {
|
|
GatewayGrantCoordinator(
|
|
baseURL: baseURL,
|
|
store: store,
|
|
session: stubSession(),
|
|
now: now,
|
|
accountAccessPolicy: UnrestrictedManagedGatewayAccountAccessPolicy()
|
|
)
|
|
}
|
|
|
|
private func stubSession() -> URLSession {
|
|
let config = URLSessionConfiguration.ephemeral
|
|
config.protocolClasses = [GatewayURLProtocol.self]
|
|
return URLSession(configuration: config)
|
|
}
|
|
|
|
private func credentials(
|
|
accessToken: String = "access-token",
|
|
refreshToken: String = "refresh-token-value-12345678901234567890",
|
|
accessExpiresAt: Date = Date(timeIntervalSince1970: 20_000),
|
|
refreshExpiresAt: Date = Date(timeIntervalSince1970: 40_000),
|
|
receivedAt: Date = Date(timeIntervalSince1970: 10_000),
|
|
scopes: Set<ManagedGatewayCapability> = [.polish, .assistant, .agent]
|
|
) -> ManagedGatewayGrantCredentials {
|
|
ManagedGatewayGrantCredentials(
|
|
grantId: "11111111-1111-1111-1111-111111111111",
|
|
scopes: scopes,
|
|
accessToken: accessToken,
|
|
accessExpiresAt: accessExpiresAt,
|
|
refreshToken: refreshToken,
|
|
refreshExpiresAt: refreshExpiresAt,
|
|
receivedAt: receivedAt
|
|
)
|
|
}
|
|
|
|
private func grantJSON(
|
|
access: String,
|
|
refresh: String,
|
|
scopes: [String] = ["polish", "ai", "agent"]
|
|
) -> Data {
|
|
let encodedScopes = scopes.map { "\"\($0)\"" }.joined(separator: ",")
|
|
return Data(
|
|
"""
|
|
{"grantId":"11111111-1111-1111-1111-111111111111",
|
|
"scopes":[\(encodedScopes)],"accessToken":"\(access)",
|
|
"accessExpiresAt":"2030-01-01T00:05:00Z","refreshToken":"\(refresh)",
|
|
"refreshExpiresAt":"2030-01-02T00:00:00Z"}
|
|
""".utf8
|
|
)
|
|
}
|
|
|
|
private func errorJSON(_ code: String) -> Data {
|
|
Data(#"{"code":"\#(code)","message":"error","requestId":"request-id"}"#.utf8)
|
|
}
|
|
|
|
private func jsonBody(_ request: URLRequest) throws -> [String: Any] {
|
|
try XCTUnwrap(
|
|
JSONSerialization.jsonObject(with: XCTUnwrap(request.httpBody)) as? [String: Any]
|
|
)
|
|
}
|
|
|
|
private func XCTAssertThrowsManaged(
|
|
_ expected: ManagedGatewayError,
|
|
operation: () async throws -> Void
|
|
) async {
|
|
do {
|
|
try await operation()
|
|
XCTFail("expected \(expected)")
|
|
} catch let error as ManagedGatewayError {
|
|
XCTAssertEqual(error, expected)
|
|
} catch {
|
|
XCTFail("unexpected error: \(error)")
|
|
}
|
|
}
|
|
}
|
|
|
|
private actor MemoryGrantStore: GatewayGrantCredentialStore {
|
|
private var credentials: ManagedGatewayGrantCredentials?
|
|
|
|
init(_ credentials: ManagedGatewayGrantCredentials? = nil) {
|
|
self.credentials = credentials
|
|
}
|
|
|
|
func load() async throws -> ManagedGatewayGrantCredentials? { credentials }
|
|
func save(_ credentials: ManagedGatewayGrantCredentials) async throws { self.credentials = credentials }
|
|
func delete() async throws { credentials = nil }
|
|
func value() -> ManagedGatewayGrantCredentials? { credentials }
|
|
}
|
|
|
|
private final class LockedValues: @unchecked Sendable {
|
|
private let lock = NSLock()
|
|
private var values: [String]
|
|
|
|
init(_ values: [String]) { self.values = values }
|
|
|
|
func next() -> String {
|
|
lock.withLock { values.isEmpty ? UUID().uuidString : values.removeFirst() }
|
|
}
|
|
}
|
|
|
|
private final class GatewayStub: @unchecked Sendable {
|
|
struct Response {
|
|
let status: Int
|
|
let body: Data
|
|
let contentType: String
|
|
let delay: TimeInterval
|
|
let transportError: URLError.Code?
|
|
}
|
|
|
|
static let shared = GatewayStub()
|
|
private let lock = NSLock()
|
|
private var responses: [Response] = []
|
|
private var captured: [URLRequest] = []
|
|
|
|
func enqueue(
|
|
_ status: Int,
|
|
_ body: Data,
|
|
contentType: String = "application/json",
|
|
delay: TimeInterval = 0,
|
|
transportError: URLError.Code? = nil
|
|
) {
|
|
lock.withLock {
|
|
responses.append(
|
|
Response(
|
|
status: status,
|
|
body: body,
|
|
contentType: contentType,
|
|
delay: delay,
|
|
transportError: transportError
|
|
)
|
|
)
|
|
}
|
|
}
|
|
|
|
func take(_ request: URLRequest) -> Response {
|
|
lock.withLock {
|
|
captured.append(Self.materialize(request))
|
|
return responses.isEmpty
|
|
? Response(
|
|
status: 500,
|
|
body: Data(),
|
|
contentType: "application/json",
|
|
delay: 0,
|
|
transportError: nil
|
|
)
|
|
: responses.removeFirst()
|
|
}
|
|
}
|
|
|
|
func requests() -> [URLRequest] { lock.withLock { captured } }
|
|
|
|
func reset() {
|
|
lock.withLock {
|
|
responses.removeAll()
|
|
captured.removeAll()
|
|
}
|
|
}
|
|
|
|
private static func materialize(_ source: URLRequest) -> URLRequest {
|
|
var request = source
|
|
guard request.httpBody == nil, let stream = request.httpBodyStream else { return request }
|
|
stream.open()
|
|
defer { stream.close() }
|
|
var data = Data()
|
|
let buffer = UnsafeMutablePointer<UInt8>.allocate(capacity: 4_096)
|
|
defer { buffer.deallocate() }
|
|
while stream.hasBytesAvailable {
|
|
let count = stream.read(buffer, maxLength: 4_096)
|
|
if count <= 0 { break }
|
|
data.append(buffer, count: count)
|
|
}
|
|
request.httpBody = data
|
|
return request
|
|
}
|
|
}
|
|
|
|
private final class GatewayURLProtocol: URLProtocol, @unchecked Sendable {
|
|
private let lock = NSLock()
|
|
private var stopped = false
|
|
|
|
override static func canInit(with request: URLRequest) -> Bool { true }
|
|
override static func canonicalRequest(for request: URLRequest) -> URLRequest { request }
|
|
|
|
override func startLoading() {
|
|
let response = GatewayStub.shared.take(request)
|
|
DispatchQueue.global().asyncAfter(deadline: .now() + response.delay) { [weak self] in
|
|
guard let self, !self.lock.withLock({ self.stopped }) else { return }
|
|
if let errorCode = response.transportError {
|
|
self.client?.urlProtocol(self, didFailWithError: URLError(errorCode))
|
|
return
|
|
}
|
|
let http = HTTPURLResponse(
|
|
url: self.request.url!,
|
|
statusCode: response.status,
|
|
httpVersion: "HTTP/1.1",
|
|
headerFields: ["Content-Type": response.contentType]
|
|
)!
|
|
self.client?.urlProtocol(self, didReceive: http, cacheStoragePolicy: .notAllowed)
|
|
self.client?.urlProtocol(self, didLoad: response.body)
|
|
self.client?.urlProtocolDidFinishLoading(self)
|
|
}
|
|
}
|
|
|
|
override func stopLoading() {
|
|
lock.withLock { stopped = true }
|
|
}
|
|
}
|