Rocky
0d236f57fb
Provide App Attest-bound, one-time onboarding AI access without creating accounts, with durable replay protection and production deployment safeguards.
265 lines
10 KiB
Kotlin
265 lines
10 KiB
Kotlin
package com.osglab.account.features.oobe
|
|
|
|
import com.osglab.account.features.gateway.models.GatewayCapability
|
|
import com.osglab.account.features.gateway.models.GatewayModelProfile
|
|
import com.osglab.account.features.gateway.models.GatewayPrincipal
|
|
import com.osglab.account.features.gateway.models.GatewayReasoningEffort
|
|
import com.osglab.account.features.gateway.models.GatewayRequestPurpose
|
|
import com.osglab.account.features.gateway.models.GatewaySubjectType
|
|
import com.osglab.account.features.gateway.models.GatewayTaskExecutionPolicy
|
|
import com.osglab.account.features.gateway.models.GatewayTaskKind
|
|
import com.osglab.account.features.gateway.models.GatewayThinkingMode
|
|
import com.osglab.account.features.gateway.models.GatewayToolsMode
|
|
import com.osglab.account.features.gateway.models.GatewayWebSearchMode
|
|
import com.osglab.account.features.gateway.models.OobeFeature
|
|
import com.osglab.account.features.gateway.models.ProviderDescriptor
|
|
import com.osglab.account.features.gateway.models.ProviderOutput
|
|
import com.osglab.account.features.gateway.models.ProviderRequest
|
|
import com.osglab.account.features.gateway.models.ProviderUsage
|
|
import com.osglab.account.features.gateway.models.TextProviderRequest
|
|
import com.osglab.account.features.gateway.models.UsageMeter
|
|
import com.osglab.account.features.gateway.ports.CreditReservation
|
|
import com.osglab.account.features.gateway.ports.CreditReservationPort
|
|
import com.osglab.account.features.gateway.ports.GatewayGrantPort
|
|
import com.osglab.account.features.gateway.ports.GatewayUsagePort
|
|
import com.osglab.account.features.gateway.ports.PendingSettlement
|
|
import com.osglab.account.features.gateway.ports.ProviderRequestMetadata
|
|
import com.osglab.account.features.gateway.providers.GatewayProvider
|
|
import com.osglab.account.features.gateway.providers.ProviderCatalog
|
|
import com.osglab.account.features.gateway.services.GatewayAccessDeniedException
|
|
import com.osglab.account.features.gateway.services.GatewayService
|
|
import io.kotest.assertions.throwables.shouldThrow
|
|
import io.kotest.core.spec.style.StringSpec
|
|
import io.kotest.matchers.shouldBe
|
|
import java.time.Instant
|
|
|
|
class OobeGatewayServiceTest : StringSpec({
|
|
"executes each fixed OOBE feature once without touching credits or account audit" {
|
|
val credits = CountingCredits()
|
|
val oobe = FakeOobeExecutionRepository()
|
|
val service = service(credits, oobe)
|
|
|
|
OobeFeature.entries.forEachIndexed { index, feature ->
|
|
service.execute(OOBE_PRINCIPAL, request(feature, "oobe-feature-$index"), DISCARD)
|
|
}
|
|
|
|
credits.calls shouldBe 0
|
|
oobe.consumed.map(OobeRequestClaim::feature).toSet() shouldBe OobeFeature.entries.toSet()
|
|
}
|
|
|
|
"rejects a fifth call and a repeated feature without paid fallback" {
|
|
val credits = CountingCredits()
|
|
val oobe = FakeOobeExecutionRepository()
|
|
val service = service(credits, oobe)
|
|
OobeFeature.entries.forEachIndexed { index, feature ->
|
|
service.execute(OOBE_PRINCIPAL, request(feature, "oobe-once-$index"), DISCARD)
|
|
}
|
|
|
|
shouldThrow<OobeFeatureAlreadyUsedException> {
|
|
service.execute(OOBE_PRINCIPAL, request(OobeFeature.ASK_AI, "oobe-fifth-call"), DISCARD)
|
|
}
|
|
credits.calls shouldBe 0
|
|
}
|
|
|
|
"releases the feature claim when the provider fails" {
|
|
val credits = CountingCredits()
|
|
val oobe = FakeOobeExecutionRepository()
|
|
val service = service(credits, oobe, fail = true)
|
|
|
|
shouldThrow<ProviderFailure> {
|
|
service.execute(OOBE_PRINCIPAL, request(OobeFeature.VOICE_INPUT, "oobe-provider-fail"), DISCARD)
|
|
}
|
|
|
|
oobe.released.map(OobeRequestClaim::feature) shouldBe listOf(OobeFeature.VOICE_INPUT)
|
|
credits.calls shouldBe 0
|
|
}
|
|
|
|
"enforces token boundary and exact feature mapping" {
|
|
val credits = CountingCredits()
|
|
val service = service(credits, FakeOobeExecutionRepository())
|
|
|
|
shouldThrow<IllegalArgumentException> {
|
|
service.execute(
|
|
OOBE_PRINCIPAL,
|
|
request(OobeFeature.ASK_AI, "oobe-wrong-map").copy(
|
|
executionPolicy = policy(GatewayTaskKind.CLIPBOARD_TRANSFORM),
|
|
),
|
|
DISCARD,
|
|
)
|
|
}
|
|
shouldThrow<GatewayAccessDeniedException> {
|
|
service.execute(
|
|
ACCOUNT_PRINCIPAL,
|
|
request(OobeFeature.ASK_AI, "account-oobe-feature"),
|
|
DISCARD,
|
|
)
|
|
}
|
|
credits.calls shouldBe 0
|
|
}
|
|
})
|
|
|
|
private fun service(
|
|
credits: CountingCredits,
|
|
oobe: OobeRepository,
|
|
fail: Boolean = false,
|
|
): GatewayService = GatewayService(
|
|
catalog = ProviderCatalog(listOf(FakeOobeProvider(fail))),
|
|
credits = credits,
|
|
grants = GatewayGrantPort { _, _ -> error("account grant lookup must not run for OOBE") },
|
|
usageRecords = NoAccountUsage,
|
|
oobeRequests = oobe,
|
|
)
|
|
|
|
private class FakeOobeProvider(private val fail: Boolean) : GatewayProvider {
|
|
override val descriptor = ProviderDescriptor(
|
|
id = "oobe-test-provider",
|
|
capabilities = setOf(GatewayCapability.POLISH, GatewayCapability.AI),
|
|
streaming = false,
|
|
usageMeter = UsageMeter.LLM_TOKEN,
|
|
)
|
|
|
|
override suspend fun execute(request: ProviderRequest, output: ProviderOutput): ProviderUsage {
|
|
if (fail) throw ProviderFailure()
|
|
return ProviderUsage(
|
|
meter = UsageMeter.LLM_TOKEN,
|
|
units = 2,
|
|
inputUnits = 1,
|
|
outputUnits = 1,
|
|
)
|
|
}
|
|
}
|
|
|
|
private class CountingCredits : CreditReservationPort {
|
|
var calls = 0
|
|
|
|
override suspend fun reserve(
|
|
accountId: String,
|
|
meter: UsageMeter,
|
|
estimatedUnits: Long,
|
|
requestId: String,
|
|
): CreditReservation {
|
|
calls += 1
|
|
error("credits must not be called")
|
|
}
|
|
|
|
override suspend fun settle(reservationId: String, actualUnits: Long) {
|
|
calls += 1
|
|
error("credits must not be called")
|
|
}
|
|
|
|
override suspend fun release(reservationId: String) {
|
|
calls += 1
|
|
error("credits must not be called")
|
|
}
|
|
}
|
|
|
|
private class FakeOobeExecutionRepository : OobeRepository {
|
|
private val claimedFeatures = mutableSetOf<OobeFeature>()
|
|
val consumed = mutableListOf<OobeRequestClaim>()
|
|
val released = mutableListOf<OobeRequestClaim>()
|
|
|
|
override suspend fun claim(
|
|
request: OobeProviderRequest,
|
|
expiresAt: Instant,
|
|
now: Instant,
|
|
): OobeRequestClaim? {
|
|
if (!claimedFeatures.add(request.feature)) return null
|
|
return OobeRequestClaim(request.subjectId, request.feature, request.requestId)
|
|
}
|
|
|
|
override suspend fun markStarted(claim: OobeRequestClaim) = Unit
|
|
|
|
override suspend fun consume(claim: OobeRequestClaim, usage: ProviderUsage) {
|
|
consumed += claim
|
|
}
|
|
|
|
override suspend fun release(claim: OobeRequestClaim, errorCode: String) {
|
|
claimedFeatures -= claim.feature
|
|
released += claim
|
|
}
|
|
|
|
override suspend fun markManualReview(claim: OobeRequestClaim, errorCode: String) = Unit
|
|
|
|
override suspend fun findOrCreateSubject(
|
|
keyId: String,
|
|
installationHash: String,
|
|
subjectId: String,
|
|
now: Instant,
|
|
): OobeSubject = error("not used")
|
|
|
|
override suspend fun createGrant(grant: NewOobeGrant, now: Instant): StoredOobeRefresh = error("not used")
|
|
|
|
override suspend fun rotateRefresh(
|
|
currentTokenHash: String,
|
|
rotationIdempotencyKey: String,
|
|
newTokenId: String,
|
|
newTokenHash: String,
|
|
newExpiresAt: Instant,
|
|
now: Instant,
|
|
): OobeRefreshRotationResult = error("not used")
|
|
|
|
override suspend fun findActiveGrant(grantId: String, subjectId: String, now: Instant): OobeGrant? =
|
|
error("not used")
|
|
}
|
|
|
|
private object NoAccountUsage : GatewayUsagePort {
|
|
override suspend fun claim(metadata: ProviderRequestMetadata) = error("account audit must not be called")
|
|
override suspend fun markStarted(accountId: String, requestId: String) = error("account audit must not be called")
|
|
override suspend fun markSettlementPending(
|
|
accountId: String,
|
|
requestId: String,
|
|
usage: ProviderUsage,
|
|
) = error("account audit must not be called")
|
|
|
|
override suspend fun markSucceeded(accountId: String, requestId: String, usage: ProviderUsage) =
|
|
error("account audit must not be called")
|
|
|
|
override suspend fun markReleased(accountId: String, requestId: String, errorCode: String) =
|
|
error("account audit must not be called")
|
|
|
|
override suspend fun markManualReview(accountId: String, requestId: String, errorCode: String) =
|
|
error("account audit must not be called")
|
|
|
|
override suspend fun findSettlementPending(limit: Int): List<PendingSettlement> = emptyList()
|
|
}
|
|
|
|
private fun request(feature: OobeFeature, requestId: String): TextProviderRequest {
|
|
val mapping = OobeContract.policy(feature)
|
|
return TextProviderRequest(
|
|
requestId = requestId,
|
|
capability = mapping.capability,
|
|
executionPolicy = policy(mapping.taskKind),
|
|
input = "hello",
|
|
context = null,
|
|
maxOutputTokens = 1,
|
|
temperature = 0.0,
|
|
stream = false,
|
|
requestPurpose = GatewayRequestPurpose.OOBE,
|
|
oobeFeature = feature,
|
|
)
|
|
}
|
|
|
|
private fun policy(taskKind: GatewayTaskKind) = GatewayTaskExecutionPolicy(
|
|
taskKind = taskKind,
|
|
modelProfile = GatewayModelProfile.LOW_LATENCY,
|
|
thinking = GatewayThinkingMode.DISABLED,
|
|
reasoningEffort = null as GatewayReasoningEffort?,
|
|
webSearch = GatewayWebSearchMode.DISABLED,
|
|
tools = GatewayToolsMode.DISABLED,
|
|
allowEmptyContentRetry = false,
|
|
maxOutputTokens = 1,
|
|
)
|
|
|
|
private val OOBE_PRINCIPAL = GatewayPrincipal(
|
|
userId = "20000000-0000-0000-0000-000000000001",
|
|
grantId = "30000000-0000-0000-0000-000000000001",
|
|
scopes = OobeContract.scopes,
|
|
subjectType = GatewaySubjectType.OOBE,
|
|
)
|
|
private val ACCOUNT_PRINCIPAL = GatewayPrincipal(
|
|
userId = "40000000-0000-0000-0000-000000000001",
|
|
scopes = OobeContract.scopes,
|
|
)
|
|
private val DISCARD = ProviderOutput {}
|
|
private class ProviderFailure : RuntimeException()
|