Add runtime provider controls and searchable AI routing

Manage provider keys at runtime, route current-information questions through server-side search with safe fallback, and scope OOBE usage claims to grants.
2026-08-22 16:33:17 +08:00
parent f8fa93dc48
commit 9fb947aa7d
43 changed files with 2079 additions and 80 deletions
@@ -88,6 +88,12 @@ import com.osglab.account.features.gateway.adapters.CreditReservationAdapter
 import com.osglab.account.features.gateway.adapters.SessionIdentityAdapter
 import com.osglab.account.features.gateway.GatewaySettings
 import com.osglab.account.features.gateway.asr.AsrStreamingService
+import com.osglab.account.features.gateway.credentials.DatabaseProviderApiKeyResolver
+import com.osglab.account.features.gateway.credentials.EnvironmentProviderCredentials
+import com.osglab.account.features.gateway.credentials.ExposedGatewayCredentialRepository
+import com.osglab.account.features.gateway.credentials.GatewayCredentialRepository
+import com.osglab.account.features.gateway.credentials.GatewayCredentialService
+import com.osglab.account.features.gateway.credentials.ProviderApiKeyResolver
 import com.osglab.account.features.gateway.ports.CreditReservationPort
 import com.osglab.account.features.gateway.ports.ComplimentaryRequestPort
 import com.osglab.account.features.gateway.ports.GatewayAccessTokenPort
@@ -272,7 +278,11 @@ fun Application.module() {
        val providerConfig = appConfig.providers.volcengine.toProviderConfig()
        AsrStreamingService(
            gateway = koin.get(),
-            upstream = KtorVolcengineAsrTransport(koin.get(), providerConfig),
+            upstream = KtorVolcengineAsrTransport(
+                client = koin.get(),
+                config = providerConfig,
+                credentialResolver = koin.get(),
+            ),
            scope = this,
        )
    } else {
@@ -367,6 +377,7 @@ fun Application.module() {
                    grantService = koin.get(),
                    operatorService = koin.get(),
                    auditService = koin.get(),
+                    credentialService = koin.get(),
                    contentService = koin.get(),
                    hintFeedService = koin.get(),
                )
@@ -399,6 +410,19 @@ fun accountServerModule(config: AppConfig): Module = module {
    single { SessionJwt(config.session) }
    single { FieldEncryptor(config.encryption.key) }
    single { IdentityFingerprint(config.antiAbuse.identityHmacKey) }
+    single<GatewayCredentialRepository> { ExposedGatewayCredentialRepository(get()) }
+    single {
+        EnvironmentProviderCredentials(
+            deepSeekApiKey = config.providers.deepSeek.apiKey,
+            volcengineApiKey = config.providers.volcengine.apiKey,
+            volcengineLegacyConfigured =
+                !config.providers.volcengine.appId.isNullOrBlank() &&
+                    !config.providers.volcengine.accessToken.isNullOrBlank(),
+        )
+    }
+    single { DatabaseProviderApiKeyResolver(get(), get(), get()) }
+    single<ProviderApiKeyResolver> { get<DatabaseProviderApiKeyResolver>() }
+    single { GatewayCredentialService(get(), get(), get()) }
    single<AdminRepository> { ExposedAdminRepository(get()) }
    single<AdminPasswordHasher> { BouncyCastleArgon2idPasswordHasher() }
    single<AdminTotpVerifier> { HmacTotpVerifier() }
@@ -667,7 +691,7 @@ fun accountServerModule(config: AppConfig): Module = module {
        )
    }
    single {
-        ProviderCatalog(configuredProviders(config, get()))
+        ProviderCatalog(configuredProviders(config, get(), get()))
    }
    single { GatewayService(get(), get(), get(), get(), get(), get()) }
    single { GatewayReconciliationService(get(), get()) }
@@ -680,7 +704,11 @@ fun accountServerModule(config: AppConfig): Module = module {
    }
 }

-private fun configuredProviders(config: AppConfig, client: HttpClient): List<GatewayProvider> =
+private fun configuredProviders(
+    config: AppConfig,
+    client: HttpClient,
+    credentialResolver: ProviderApiKeyResolver,
+): List<GatewayProvider> =
    buildList {
        config.providers.deepSeek.apiKey?.let { apiKey ->
            add(
@@ -692,12 +720,21 @@ private fun configuredProviders(config: AppConfig, client: HttpClient): List<Gat
                        model = config.providers.deepSeek.model,
                        reasoningModel = config.providers.deepSeek.reasoningModel,
                    ),
+                    credentialResolver = credentialResolver,
                ),
            )
        }
        if (config.providers.volcengine.credentialsAvailable) {
            val providerConfig = config.providers.volcengine.toProviderConfig()
-            add(VolcengineAsrProvider(KtorVolcengineAsrTransport(client, providerConfig)))
+            add(
+                VolcengineAsrProvider(
+                    KtorVolcengineAsrTransport(
+                        client = client,
+                        config = providerConfig,
+                        credentialResolver = credentialResolver,
+                    ),
+                ),
+            )
        }
    }

@@ -119,6 +119,7 @@ enum class AdminAuditAction {
    CONTENT_HINT_PACK_SAVED,
    CONTENT_HINT_FEED_SETTINGS_UPDATED,
    CONTENT_HINT_FEED_GENERATED,
+    PROVIDER_API_KEY_UPDATED,
 }

 enum class AdminAuditOutcome {
@@ -47,6 +47,9 @@ import com.osglab.account.features.credits.domain.InvalidCreditRequest
 import com.osglab.account.features.credits.domain.LedgerEntryType
 import com.osglab.account.features.content.services.ContentService
 import com.osglab.account.features.content.feed.HintFeedService
+import com.osglab.account.features.gateway.credentials.GatewayCredentialProvider
+import com.osglab.account.features.gateway.credentials.GatewayCredentialService
+import com.osglab.account.features.gateway.credentials.InvalidProviderApiKeyException
 import io.ktor.http.Cookie
 import io.ktor.http.HttpHeaders
 import io.ktor.http.HttpStatusCode
@@ -65,6 +68,7 @@ import io.ktor.server.routing.Route
 import io.ktor.server.routing.delete
 import io.ktor.server.routing.get
 import io.ktor.server.routing.post
+import io.ktor.server.routing.put
 import io.ktor.server.routing.route
 import kotlinx.serialization.Serializable
 import java.time.Clock
@@ -93,6 +97,7 @@ fun Route.adminApiRoutes(
    grantService: AdminGrantService,
    operatorService: AdminOperatorService,
    auditService: AdminAuditService,
+    credentialService: GatewayCredentialService? = null,
    contentService: ContentService? = null,
    hintFeedService: HintFeedService? = null,
    clock: Clock = Clock.systemUTC(),
@@ -174,6 +179,46 @@ fun Route.adminApiRoutes(
            adminContentRoutes(config, sessionService, it, hintFeedService)
        }

+        credentialService?.let { service ->
+            get("/providers") {
+                if (
+                    call.requireRole(
+                        config,
+                        sessionService,
+                        setOf(AdminRole.SUPER_ADMIN),
+                    ) == null
+                ) return@get
+                call.respond(service.listStatuses())
+            }
+
+            put("/providers/{providerId}/api-key") {
+                val principal = call.requireMutationPrincipal(config, sessionService) ?: return@put
+                if (principal.role != AdminRole.SUPER_ADMIN) {
+                    call.respond(HttpStatusCode.Forbidden, AdminErrorResponse("INSUFFICIENT_PERMISSION"))
+                    return@put
+                }
+                val provider = GatewayCredentialProvider.fromProviderId(
+                    call.parameters["providerId"],
+                ) ?: run {
+                    call.respond(HttpStatusCode.NotFound, AdminErrorResponse("PROVIDER_NOT_FOUND"))
+                    return@put
+                }
+                val request = call.receiveAdminRequest<ProviderApiKeyUpdateRequest>() ?: return@put
+                val status = try {
+                    service.updateApiKey(
+                        provider = provider,
+                        apiKey = request.apiKey,
+                        operatorId = principal.operatorId,
+                        requestId = call.request.header("X-Request-ID"),
+                    )
+                } catch (_: InvalidProviderApiKeyException) {
+                    call.respond(HttpStatusCode.BadRequest, AdminErrorResponse("VALIDATION_ERROR"))
+                    return@put
+                }
+                call.respond(status)
+            }
+        }
+
        get("/overview") {
            if (call.requirePrincipal(config, sessionService) == null) return@get
            val stats = statsService.getRange(call.request.queryParameters["range"], clock)
@@ -1098,6 +1143,9 @@ private data class AdminGrantRequest(val userId: String, val amount: Long, val r
@Serializable
 private data class AdminGrantResponse(val transactionId: String, val balanceAfter: Long)

+@Serializable
+private data class ProviderApiKeyUpdateRequest(val apiKey: String)
+
@Serializable
 private data class AdminOperatorCreateRequest(
    val username: String,
@@ -0,0 +1,44 @@
+package com.osglab.account.features.gateway.credentials
+
+import kotlinx.serialization.Serializable
+import java.time.Instant
+import java.util.UUID
+
+enum class GatewayCredentialProvider(val providerId: String) {
+    DEEPSEEK("deepseek"),
+    VOLCENGINE("volcengine");
+
+    companion object {
+        fun fromProviderId(value: String?): GatewayCredentialProvider? =
+            entries.firstOrNull { it.providerId == value }
+    }
+}
+
+@Serializable
+enum class GatewayCredentialSource {
+    ENVIRONMENT,
+    RUNTIME_OVERRIDE,
+}
+
+@Serializable
+data class GatewayCredentialStatus(
+    val providerId: String,
+    val configured: Boolean,
+    val source: GatewayCredentialSource,
+    val updatedAt: String? = null,
+)
+
+/**
+ * Persistence-only encrypted value. Its string representation deliberately
+ * excludes ciphertext so authenticated encryption material cannot reach logs.
+ */
+class ProviderApiKeyOverride(
+    val provider: GatewayCredentialProvider,
+    val encryptedApiKey: String,
+    val updatedAt: Instant,
+    val updatedByOperatorId: UUID,
+) {
+    override fun toString(): String =
+        "ProviderApiKeyOverride(provider=${provider.providerId}, updatedAt=$updatedAt, " +
+            "updatedByOperatorId=$updatedByOperatorId)"
+}
@@ -0,0 +1,90 @@
+package com.osglab.account.features.gateway.credentials
+
+import com.osglab.account.config.DatabaseFactory
+import com.osglab.account.features.admin.models.NewAdminAuditEvent
+import com.osglab.account.features.admin.repositories.AdminAuditLogTable
+import org.jetbrains.exposed.v1.core.ResultRow
+import org.jetbrains.exposed.v1.core.Table
+import org.jetbrains.exposed.v1.core.eq
+import org.jetbrains.exposed.v1.javatime.timestamp
+import org.jetbrains.exposed.v1.jdbc.insert
+import org.jetbrains.exposed.v1.jdbc.insertIgnore
+import org.jetbrains.exposed.v1.jdbc.selectAll
+import org.jetbrains.exposed.v1.jdbc.update
+import java.util.UUID
+
+internal object GatewayProviderCredentialsTable : Table("gateway_provider_credentials") {
+    val providerId = varchar("provider_id", 32)
+    val encryptedApiKey = text("encrypted_api_key")
+    val updatedAt = timestamp("updated_at")
+    val updatedByOperatorId = varchar("updated_by_operator_id", 36)
+    override val primaryKey = PrimaryKey(providerId)
+}
+
+interface GatewayCredentialRepository {
+    suspend fun findOverride(provider: GatewayCredentialProvider): ProviderApiKeyOverride?
+
+    /**
+     * Stores the current override and its successful audit event atomically.
+     */
+    suspend fun upsertOverride(
+        credentialOverride: ProviderApiKeyOverride,
+        auditEvent: NewAdminAuditEvent,
+    )
+}
+
+class ExposedGatewayCredentialRepository(
+    private val databaseFactory: DatabaseFactory,
+) : GatewayCredentialRepository {
+    override suspend fun findOverride(
+        provider: GatewayCredentialProvider,
+    ): ProviderApiKeyOverride? = databaseFactory.query {
+        GatewayProviderCredentialsTable.selectAll()
+            .where { GatewayProviderCredentialsTable.providerId eq provider.providerId }
+            .limit(1)
+            .singleOrNull()
+            ?.toOverride()
+    }
+
+    override suspend fun upsertOverride(
+        credentialOverride: ProviderApiKeyOverride,
+        auditEvent: NewAdminAuditEvent,
+    ) {
+        databaseFactory.query {
+            val inserted = GatewayProviderCredentialsTable.insertIgnore {
+                it[providerId] = credentialOverride.provider.providerId
+                it[encryptedApiKey] = credentialOverride.encryptedApiKey
+                it[updatedAt] = credentialOverride.updatedAt
+                it[updatedByOperatorId] = credentialOverride.updatedByOperatorId.toString()
+            }.insertedCount > 0
+            if (!inserted) {
+                GatewayProviderCredentialsTable.update({
+                    GatewayProviderCredentialsTable.providerId eq credentialOverride.provider.providerId
+                }) {
+                    it[encryptedApiKey] = credentialOverride.encryptedApiKey
+                    it[updatedAt] = credentialOverride.updatedAt
+                    it[updatedByOperatorId] = credentialOverride.updatedByOperatorId.toString()
+                }
+            }
+            AdminAuditLogTable.insert {
+                it[id] = auditEvent.id.toString()
+                it[actorOperatorId] = auditEvent.actorOperatorId?.toString()
+                it[action] = auditEvent.action.name
+                it[outcome] = auditEvent.outcome.name
+                it[targetType] = auditEvent.targetType
+                it[targetId] = auditEvent.targetId
+                it[requestId] = auditEvent.requestId
+                it[occurredAt] = auditEvent.occurredAt
+            }
+        }
+    }
+}
+
+private fun ResultRow.toOverride(): ProviderApiKeyOverride = ProviderApiKeyOverride(
+    provider = requireNotNull(
+        GatewayCredentialProvider.fromProviderId(this[GatewayProviderCredentialsTable.providerId]),
+    ),
+    encryptedApiKey = this[GatewayProviderCredentialsTable.encryptedApiKey],
+    updatedAt = this[GatewayProviderCredentialsTable.updatedAt],
+    updatedByOperatorId = UUID.fromString(this[GatewayProviderCredentialsTable.updatedByOperatorId]),
+)
@@ -0,0 +1,60 @@
+package com.osglab.account.features.gateway.credentials
+
+import com.osglab.account.common.security.FieldEncryptor
+import com.osglab.account.features.admin.models.AdminAuditAction
+import com.osglab.account.features.admin.models.AdminAuditOutcome
+import com.osglab.account.features.admin.models.NewAdminAuditEvent
+import java.time.Clock
+import java.util.UUID
+
+class InvalidProviderApiKeyException : IllegalArgumentException("Invalid provider API key")
+
+class GatewayCredentialService(
+    private val repository: GatewayCredentialRepository,
+    private val resolver: DatabaseProviderApiKeyResolver,
+    private val fieldEncryptor: FieldEncryptor,
+    private val clock: Clock = Clock.systemUTC(),
+) {
+    suspend fun listStatuses(): List<GatewayCredentialStatus> =
+        GatewayCredentialProvider.entries.map { resolver.status(it) }
+
+    suspend fun updateApiKey(
+        provider: GatewayCredentialProvider,
+        apiKey: String,
+        operatorId: UUID,
+        requestId: String?,
+    ): GatewayCredentialStatus {
+        val normalized = apiKey.trim()
+        if (
+            normalized.isEmpty() ||
+            normalized.length > MAX_API_KEY_LENGTH ||
+            apiKey.contains('\r') ||
+            apiKey.contains('\n')
+        ) {
+            throw InvalidProviderApiKeyException()
+        }
+        val now = clock.instant()
+        repository.upsertOverride(
+            credentialOverride = ProviderApiKeyOverride(
+                provider = provider,
+                encryptedApiKey = fieldEncryptor.encrypt(normalized, encryptionContext(provider)),
+                updatedAt = now,
+                updatedByOperatorId = operatorId,
+            ),
+            auditEvent = NewAdminAuditEvent(
+                actorOperatorId = operatorId,
+                action = AdminAuditAction.PROVIDER_API_KEY_UPDATED,
+                outcome = AdminAuditOutcome.SUCCESS,
+                targetType = "PROVIDER",
+                targetId = provider.providerId,
+                requestId = requestId,
+                occurredAt = now,
+            ),
+        )
+        return resolver.status(provider)
+    }
+
+    private companion object {
+        const val MAX_API_KEY_LENGTH = 4_096
+    }
+}
@@ -0,0 +1,76 @@
+package com.osglab.account.features.gateway.credentials
+
+import com.osglab.account.common.security.FieldEncryptor
+
+fun interface ProviderApiKeyResolver {
+    /**
+     * Resolves the effective API key at the start of a new upstream request.
+     */
+    suspend fun resolve(provider: GatewayCredentialProvider): String?
+}
+
+class EnvironmentProviderCredentials(
+    val deepSeekApiKey: String?,
+    val volcengineApiKey: String?,
+    val volcengineLegacyConfigured: Boolean,
+) {
+    fun apiKey(provider: GatewayCredentialProvider): String? =
+        when (provider) {
+            GatewayCredentialProvider.DEEPSEEK -> deepSeekApiKey
+            GatewayCredentialProvider.VOLCENGINE -> volcengineApiKey
+        }?.trim()?.takeIf(String::isNotEmpty)
+
+    fun configured(provider: GatewayCredentialProvider): Boolean =
+        apiKey(provider) != null ||
+            (provider == GatewayCredentialProvider.VOLCENGINE && volcengineLegacyConfigured)
+
+    override fun toString(): String =
+        "EnvironmentProviderCredentials(deepSeekConfigured=${!deepSeekApiKey.isNullOrBlank()}, " +
+            "volcengineApiKeyConfigured=${!volcengineApiKey.isNullOrBlank()}, " +
+            "volcengineLegacyConfigured=$volcengineLegacyConfigured)"
+}
+
+class DatabaseProviderApiKeyResolver(
+    private val repository: GatewayCredentialRepository,
+    private val fieldEncryptor: FieldEncryptor,
+    private val environment: EnvironmentProviderCredentials,
+) : ProviderApiKeyResolver {
+    override suspend fun resolve(provider: GatewayCredentialProvider): String? =
+        repository.findOverride(provider)?.let {
+            fieldEncryptor.decrypt(it.encryptedApiKey, encryptionContext(provider))
+        } ?: environment.apiKey(provider)
+
+    suspend fun status(provider: GatewayCredentialProvider): GatewayCredentialStatus {
+        val runtimeOverride = repository.findOverride(provider)
+        return if (runtimeOverride != null) {
+            GatewayCredentialStatus(
+                providerId = provider.providerId,
+                configured = true,
+                source = GatewayCredentialSource.RUNTIME_OVERRIDE,
+                updatedAt = runtimeOverride.updatedAt.toString(),
+            )
+        } else {
+            GatewayCredentialStatus(
+                providerId = provider.providerId,
+                configured = environment.configured(provider),
+                source = GatewayCredentialSource.ENVIRONMENT,
+            )
+        }
+    }
+}
+
+class StaticProviderApiKeyResolver(
+    deepSeekApiKey: String? = null,
+    volcengineApiKey: String? = null,
+) : ProviderApiKeyResolver {
+    private val apiKeys = mapOf(
+        GatewayCredentialProvider.DEEPSEEK to deepSeekApiKey,
+        GatewayCredentialProvider.VOLCENGINE to volcengineApiKey,
+    )
+
+    override suspend fun resolve(provider: GatewayCredentialProvider): String? =
+        apiKeys[provider]?.trim()?.takeIf(String::isNotEmpty)
+}
+
+internal fun encryptionContext(provider: GatewayCredentialProvider): String =
+    "gateway-provider-api-key:${provider.providerId}"
@@ -109,6 +109,9 @@ enum class GatewayTaskKind {
    @SerialName("ai_question")
    AI_QUESTION,

+    @SerialName("current_information_question")
+    CURRENT_INFORMATION_QUESTION,
+
    @SerialName("clipboard_transform")
    CLIPBOARD_TRANSFORM,

@@ -167,6 +170,12 @@ data class GatewayTaskExecutionPolicy(
        ) {
            "reasoning effort must be explicit exactly when thinking is enabled"
        }
+        require(
+            webSearch == GatewayWebSearchMode.DISABLED ||
+                thinking == GatewayThinkingMode.ENABLED,
+        ) {
+            "web search requires an explicit thinking policy"
+        }
    }
 }

@@ -1,10 +1,14 @@
 package com.osglab.account.features.gateway.providers.deepseek

 import com.osglab.account.features.gateway.agent.AgentPlan
+import com.osglab.account.features.gateway.credentials.GatewayCredentialProvider
+import com.osglab.account.features.gateway.credentials.ProviderApiKeyResolver
+import com.osglab.account.features.gateway.credentials.StaticProviderApiKeyResolver
 import com.osglab.account.features.gateway.models.GatewayCapability
 import com.osglab.account.features.gateway.models.GatewayLimits
 import com.osglab.account.features.gateway.models.GatewayModelProfile
 import com.osglab.account.features.gateway.models.GatewayReasoningEffort
+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
@@ -75,6 +79,19 @@ fun interface DeepSeekClient {
    suspend fun complete(request: TextProviderRequest, output: ProviderOutput): ProviderUsage
 }

+private fun createDeepSeekClient(
+    client: HttpClient,
+    config: DeepSeekConfig,
+    json: Json,
+    credentialResolver: ProviderApiKeyResolver,
+): DeepSeekClient {
+    val chat = KtorDeepSeekClient(client, config, json, credentialResolver)
+    return DeepSeekSearchFallbackClient(
+        search = KtorDeepSeekResponsesClient(client, config, json, credentialResolver),
+        fallback = chat,
+    )
+}
+
 class DeepSeekProvider(
    private val upstream: DeepSeekClient,
 ) : GatewayProvider {
@@ -85,7 +102,9 @@ class DeepSeekProvider(
            ignoreUnknownKeys = true
            explicitNulls = false
        },
-    ) : this(KtorDeepSeekClient(client, config, json))
+        credentialResolver: ProviderApiKeyResolver =
+            StaticProviderApiKeyResolver(deepSeekApiKey = config.apiKey),
+    ) : this(createDeepSeekClient(client, config, json, credentialResolver))

    override val descriptor = ProviderDescriptor(
        id = "deepseek",
@@ -145,8 +164,14 @@ class DeepSeekProvider(
        require(request.maxOutputTokens == request.executionPolicy.maxOutputTokens) {
            "maxOutputTokens must match the server execution policy"
        }
-        require(request.executionPolicy.webSearch == GatewayWebSearchMode.DISABLED) {
-            "DeepSeek web search is not configured"
+        require(
+            request.executionPolicy.webSearch == GatewayWebSearchMode.DISABLED ||
+                (
+                    request.capability == GatewayCapability.AI &&
+                        request.executionPolicy.thinking == GatewayThinkingMode.ENABLED
+                    )
+        ) {
+            "DeepSeek web search is supported only for reasoning AI requests"
        }
        require(request.executionPolicy.tools == GatewayToolsMode.DISABLED) {
            "DeepSeek tools are not configured"
@@ -176,11 +201,18 @@ class KtorDeepSeekClient(
        ignoreUnknownKeys = true
        explicitNulls = false
    },
+    private val credentialResolver: ProviderApiKeyResolver =
+        StaticProviderApiKeyResolver(deepSeekApiKey = config.apiKey),
 ) : DeepSeekClient {
    override suspend fun complete(
        request: TextProviderRequest,
        output: ProviderOutput,
    ): ProviderUsage {
+        require(request.executionPolicy.webSearch == GatewayWebSearchMode.DISABLED) {
+            "Chat Completions cannot execute server-side web search"
+        }
+        val apiKey = credentialResolver.resolve(GatewayCredentialProvider.DEEPSEEK)
+            ?: throw DeepSeekConfigurationException("DeepSeek API key is not configured")
        val payload = DeepSeekChatRequest(
            model = config.modelFor(request.executionPolicy.modelProfile),
            messages = controlledMessages(request),
@@ -203,7 +235,7 @@ class KtorDeepSeekClient(
        )

        return client.preparePost("${config.endpoint.trimEnd('/')}/chat/completions") {
-            bearerAuth(config.apiKey)
+            bearerAuth(apiKey)
            contentType(ContentType.Application.Json)
            header(HttpHeaders.Accept, if (request.stream) ContentType.Text.EventStream else ContentType.Application.Json)
            header("X-Request-ID", request.requestId)
@@ -491,34 +523,11 @@ class KtorDeepSeekClient(
        return bytes
    }

-    private fun controlledMessages(request: TextProviderRequest): List<ChatMessage> {
-        val system = when (request.capability) {
-            GatewayCapability.POLISH ->
-                "Polish the user's text while preserving meaning. Return only the polished text."
-
-            GatewayCapability.AI ->
-                "Answer the user's question accurately and concisely. Do not claim actions you did not perform."
-
-            GatewayCapability.AGENT ->
-                """
-                Return only JSON with this schema:
-                {"summary":"string","steps":[{"id":"string","title":"string","description":"string"}],"warnings":["string"]}.
-                Produce a declarative plan only. Never execute actions, invoke tools, include commands or URLs,
-                or claim that any client-side or external side effect occurred.
-                """.trimIndent()
-
-            GatewayCapability.ASR -> error("ASR is not a DeepSeek capability")
-        }
-        val userText = buildString {
-            request.context?.takeIf(String::isNotBlank)?.let {
-                append("Context:\n")
-                append(it)
-                append("\n\n")
-            }
-            append(request.input)
-        }
-        return listOf(ChatMessage("system", system), ChatMessage("user", userText))
-    }
+    private fun controlledMessages(request: TextProviderRequest): List<ChatMessage> =
+        listOf(
+            ChatMessage("system", deepSeekSystemInstruction(request)),
+            ChatMessage("user", deepSeekUserText(request)),
+        )

    private fun GatewayReasoningEffort.toDeepSeekReasoningEffort(): DeepSeekReasoningEffort =
        when (this) {
@@ -539,6 +548,50 @@ class KtorDeepSeekClient(
    }
 }

+internal fun deepSeekSystemInstruction(request: TextProviderRequest): String =
+    when (request.capability) {
+        GatewayCapability.POLISH ->
+            "Polish the user's text while preserving meaning. Return only the polished text."
+
+        GatewayCapability.AI -> buildString {
+            append("Answer the user's question accurately and concisely. ")
+            append("Do not claim actions you did not perform.")
+            if (request.executionPolicy.webSearch == GatewayWebSearchMode.DISABLED &&
+                request.executionPolicy.taskKind in SEARCHABLE_AI_TASKS
+            ) {
+                append(
+                    " Web search is temporarily unavailable. Do not present time-sensitive " +
+                        "information as current; clearly state that it could not be verified.",
+                )
+            }
+        }
+
+        GatewayCapability.AGENT ->
+            """
+            Return only JSON with this schema:
+            {"summary":"string","steps":[{"id":"string","title":"string","description":"string"}],"warnings":["string"]}.
+            Produce a declarative plan only. Never execute actions, invoke tools, include commands or URLs,
+            or claim that any client-side or external side effect occurred.
+            """.trimIndent()
+
+        GatewayCapability.ASR -> error("ASR is not a DeepSeek capability")
+    }
+
+internal fun deepSeekUserText(request: TextProviderRequest): String =
+    buildString {
+        request.context?.takeIf(String::isNotBlank)?.let {
+            append("Context:\n")
+            append(it)
+            append("\n\n")
+        }
+        append(request.input)
+    }
+
+private val SEARCHABLE_AI_TASKS = setOf(
+    GatewayTaskKind.AI_QUESTION,
+    GatewayTaskKind.CURRENT_INFORMATION_QUESTION,
+)
+
@Serializable
 private data class DeepSeekChatRequest(
    val model: String,
@@ -0,0 +1,356 @@
+package com.osglab.account.features.gateway.providers.deepseek
+
+import com.osglab.account.features.gateway.credentials.GatewayCredentialProvider
+import com.osglab.account.features.gateway.credentials.ProviderApiKeyResolver
+import com.osglab.account.features.gateway.credentials.StaticProviderApiKeyResolver
+import com.osglab.account.features.gateway.models.GatewayLimits
+import com.osglab.account.features.gateway.models.GatewayWebSearchMode
+import com.osglab.account.features.gateway.models.ProviderOutput
+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 io.ktor.client.HttpClient
+import io.ktor.client.call.body
+import io.ktor.client.request.bearerAuth
+import io.ktor.client.request.header
+import io.ktor.client.request.preparePost
+import io.ktor.client.request.setBody
+import io.ktor.http.ContentType
+import io.ktor.http.HttpHeaders
+import io.ktor.http.contentType
+import io.ktor.http.isSuccess
+import io.ktor.utils.io.ByteReadChannel
+import io.ktor.utils.io.readRemaining
+import kotlinx.coroutines.CancellationException
+import kotlinx.io.readByteArray
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonElement
+import kotlinx.serialization.json.JsonNull
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.JsonPrimitive
+import kotlinx.serialization.json.buildJsonArray
+import kotlinx.serialization.json.buildJsonObject
+import kotlinx.serialization.json.jsonArray
+import kotlinx.serialization.json.jsonObject
+import kotlinx.serialization.json.jsonPrimitive
+import kotlinx.serialization.json.longOrNull
+import kotlinx.serialization.json.put
+import org.slf4j.LoggerFactory
+
+/**
+ * Uses DeepSeek Responses only when server policy permits web search. The
+ * search attempt is fully buffered so a failure can safely fall back without
+ * mixing two answers in a downstream stream.
+ */
+internal class DeepSeekSearchFallbackClient(
+    private val search: DeepSeekClient,
+    private val fallback: DeepSeekClient,
+) : DeepSeekClient {
+    override suspend fun complete(
+        request: TextProviderRequest,
+        output: ProviderOutput,
+    ): ProviderUsage {
+        if (request.executionPolicy.webSearch == GatewayWebSearchMode.DISABLED) {
+            return fallback.complete(request, output)
+        }
+
+        val buffered = mutableListOf<ByteArray>()
+        var bufferedBytes = 0L
+        val usage = try {
+            search.complete(
+                request,
+                ProviderOutput { bytes ->
+                    bufferedBytes = Math.addExact(bufferedBytes, bytes.size.toLong())
+                    if (bufferedBytes > GatewayLimits.MAX_UPSTREAM_RESPONSE_BYTES) {
+                        throw DeepSeekProviderException("DeepSeek search output exceeded the gateway limit")
+                    }
+                    buffered += bytes
+                },
+            )
+        } catch (failure: CancellationException) {
+            throw failure
+        } catch (failure: Exception) {
+            LOG.warn(
+                "DeepSeek search path failed requestId={} taskKind={} searchMode={} failureType={}",
+                request.requestId,
+                request.executionPolicy.taskKind.name,
+                request.executionPolicy.webSearch.name,
+                failure::class.simpleName ?: "Exception",
+            )
+            return fallback.complete(
+                request.copy(
+                    executionPolicy = request.executionPolicy.copy(
+                        webSearch = GatewayWebSearchMode.DISABLED,
+                    ),
+                ),
+                output,
+            )
+        }
+
+        buffered.forEach { output.emit(it) }
+        return usage
+    }
+
+    private companion object {
+        val LOG = LoggerFactory.getLogger(DeepSeekSearchFallbackClient::class.java)
+    }
+}
+
+internal class KtorDeepSeekResponsesClient(
+    private val client: HttpClient,
+    private val config: DeepSeekConfig,
+    private val json: Json = Json {
+        ignoreUnknownKeys = true
+        explicitNulls = false
+    },
+    private val credentialResolver: ProviderApiKeyResolver =
+        StaticProviderApiKeyResolver(deepSeekApiKey = config.apiKey),
+) : DeepSeekClient {
+    override suspend fun complete(
+        request: TextProviderRequest,
+        output: ProviderOutput,
+    ): ProviderUsage {
+        val searchMode = request.executionPolicy.webSearch
+        require(searchMode != GatewayWebSearchMode.DISABLED) {
+            "Responses search requires an enabled server search policy"
+        }
+        val reasoningEffort = requireNotNull(request.executionPolicy.reasoningEffort) {
+            "Responses search requires explicit reasoning effort"
+        }
+        val apiKey = credentialResolver.resolve(GatewayCredentialProvider.DEEPSEEK)
+            ?: throw DeepSeekConfigurationException("DeepSeek API key is not configured")
+        val payload = DeepSeekResponsesRequest(
+            model = config.modelFor(request.executionPolicy.modelProfile),
+            instructions = deepSeekSystemInstruction(request),
+            input = listOf(DeepSeekResponsesMessage("user", deepSeekUserText(request))),
+            tools = listOf(DeepSeekResponsesTool("web_search")),
+            toolChoice = when (searchMode) {
+                GatewayWebSearchMode.ALLOWED -> JsonPrimitive("auto")
+                GatewayWebSearchMode.REQUIRED -> buildJsonObject { put("type", "web_search") }
+                GatewayWebSearchMode.DISABLED -> error("Search policy changed during request construction")
+            },
+            maxOutputTokens = request.maxOutputTokens,
+            reasoning = DeepSeekResponsesReasoning(reasoningEffort.name.lowercase()),
+        )
+
+        return client.preparePost("${config.endpoint.trimEnd('/')}/responses") {
+            bearerAuth(apiKey)
+            contentType(ContentType.Application.Json)
+            header(HttpHeaders.Accept, ContentType.Application.Json)
+            header("X-Request-ID", request.requestId)
+            setBody(payload)
+        }.execute { response ->
+            if (!response.status.isSuccess()) {
+                runCatching { response.body<ByteReadChannel>().readBounded() }
+                throw DeepSeekProviderException(
+                    "DeepSeek Responses returned HTTP ${response.status.value}",
+                )
+            }
+            val responseContentType = response.headers[HttpHeaders.ContentType]
+                ?.let { runCatching { ContentType.parse(it) }.getOrNull() }
+            if (responseContentType?.match(ContentType.Application.Json) != true) {
+                runCatching { response.body<ByteReadChannel>().readBounded() }
+                throw DeepSeekProviderException(
+                    "DeepSeek Responses returned an unexpected content type",
+                )
+            }
+
+            val result = parseResponse(response.body<ByteReadChannel>().readBounded())
+            if (searchMode == GatewayWebSearchMode.REQUIRED && !result.webSearchUsed) {
+                throw DeepSeekProviderException("DeepSeek omitted required web search")
+            }
+            emitCompatibleResponse(request, result, output)
+            result.usage
+        }
+    }
+
+    private fun parseResponse(bytes: ByteArray): DeepSeekResponsesResult {
+        val root = runCatching { json.parseToJsonElement(bytes.decodeToString()).jsonObject }
+            .getOrElse { throw DeepSeekProviderException("DeepSeek returned malformed Responses JSON") }
+        if (root["error"] != null && root["error"] !is JsonNull) {
+            throw DeepSeekProviderException("DeepSeek Responses returned an error")
+        }
+        val output = runCatching { root["output"]?.jsonArray ?: emptyList() }
+            .getOrElse { throw DeepSeekProviderException("DeepSeek returned invalid Responses output") }
+        val topLevelText = runCatching {
+            root["output_text"]?.jsonPrimitive?.takeIf { it.isString }?.content
+        }.getOrNull()
+        val messageText = buildString {
+            output.forEach { itemElement ->
+                val item = runCatching { itemElement.jsonObject }.getOrNull() ?: return@forEach
+                if (item.string("type") != "message") return@forEach
+                val content = runCatching { item["content"]?.jsonArray ?: emptyList() }
+                    .getOrElse {
+                        throw DeepSeekProviderException("DeepSeek returned invalid message content")
+                    }
+                content.forEach { partElement ->
+                    val part = runCatching { partElement.jsonObject }.getOrNull() ?: return@forEach
+                    if (part.string("type") in RESPONSE_TEXT_TYPES) {
+                        part.string("text")?.let(::append)
+                    }
+                }
+            }
+        }
+        val text = topLevelText?.takeIf(String::isNotBlank) ?: messageText
+        if (text.isBlank()) {
+            throw DeepSeekProviderException("DeepSeek Responses omitted output text")
+        }
+        val webSearchUsed = output.any { item ->
+            runCatching { item.jsonObject.string("type") == "web_search_call" }.getOrDefault(false)
+        }
+        val usageObject = runCatching { root["usage"]?.jsonObject }
+            .getOrNull()
+            ?: throw DeepSeekProviderException("DeepSeek Responses omitted token usage")
+        val inputTokens = usageObject.long("input_tokens")
+            ?: throw DeepSeekProviderException("DeepSeek Responses omitted input token usage")
+        val outputTokens = usageObject.long("output_tokens")
+            ?: throw DeepSeekProviderException("DeepSeek Responses omitted output token usage")
+        val computedTotal = runCatching { Math.addExact(inputTokens, outputTokens) }
+            .getOrElse { throw DeepSeekUsageException("DeepSeek Responses token usage overflowed") }
+        val totalTokens = usageObject.long("total_tokens") ?: computedTotal
+        if (inputTokens < 0 || outputTokens < 0 || totalTokens != computedTotal) {
+            throw DeepSeekUsageException("DeepSeek Responses returned inconsistent token usage")
+        }
+
+        return DeepSeekResponsesResult(
+            text = text,
+            webSearchUsed = webSearchUsed,
+            usage = ProviderUsage(
+                meter = UsageMeter.LLM_TOKEN,
+                units = totalTokens,
+                inputUnits = inputTokens,
+                outputUnits = outputTokens,
+            ),
+        )
+    }
+
+    private suspend fun emitCompatibleResponse(
+        request: TextProviderRequest,
+        result: DeepSeekResponsesResult,
+        output: ProviderOutput,
+    ) {
+        if (!request.stream) {
+            output.emit(bufferedChatPayload(result).encodeToByteArray())
+            return
+        }
+        val content = buildJsonObject {
+            put(
+                "choices",
+                buildJsonArray {
+                    add(
+                        buildJsonObject {
+                            put(
+                                "delta",
+                                buildJsonObject { put("content", result.text) },
+                            )
+                        },
+                    )
+                },
+            )
+        }
+        val terminal = buildJsonObject {
+            put(
+                "choices",
+                buildJsonArray {
+                    add(
+                        buildJsonObject {
+                            put("delta", buildJsonObject {})
+                            put("finish_reason", "stop")
+                        },
+                    )
+                },
+            )
+        }
+        val usage = buildJsonObject {
+            put("choices", buildJsonArray {})
+            put("usage", usageJson(result.usage))
+        }
+        output.emit("data: $content\n\n".encodeToByteArray())
+        output.emit("data: $terminal\n\n".encodeToByteArray())
+        output.emit("data: $usage\n\n".encodeToByteArray())
+        output.emit("data: [DONE]\n\n".encodeToByteArray())
+    }
+
+    private fun bufferedChatPayload(result: DeepSeekResponsesResult): String =
+        buildJsonObject {
+            put(
+                "choices",
+                buildJsonArray {
+                    add(
+                        buildJsonObject {
+                            put(
+                                "message",
+                                buildJsonObject { put("content", result.text) },
+                            )
+                        },
+                    )
+                },
+            )
+            put("usage", usageJson(result.usage))
+        }.toString()
+
+    private fun usageJson(usage: ProviderUsage): JsonObject =
+        buildJsonObject {
+            put("prompt_tokens", requireNotNull(usage.inputUnits))
+            put("completion_tokens", requireNotNull(usage.outputUnits))
+            put("total_tokens", usage.units)
+        }
+
+    private suspend fun ByteReadChannel.readBounded(): ByteArray {
+        val bytes = readRemaining(GatewayLimits.MAX_UPSTREAM_RESPONSE_BYTES.toLong() + 1L)
+            .readByteArray()
+        if (bytes.size > GatewayLimits.MAX_UPSTREAM_RESPONSE_BYTES) {
+            throw DeepSeekProviderException("DeepSeek Responses exceeded the gateway limit")
+        }
+        return bytes
+    }
+
+    private companion object {
+        val RESPONSE_TEXT_TYPES = setOf("output_text", "text")
+    }
+}
+
+private fun JsonObject.string(name: String): String? =
+    runCatching {
+        get(name)?.jsonPrimitive?.takeIf { it.isString }?.content
+    }.getOrNull()
+
+private fun JsonObject.long(name: String): Long? =
+    runCatching { get(name)?.jsonPrimitive?.longOrNull }.getOrNull()
+
+private data class DeepSeekResponsesResult(
+    val text: String,
+    val webSearchUsed: Boolean,
+    val usage: ProviderUsage,
+)
+
+@Serializable
+private data class DeepSeekResponsesRequest(
+    val model: String,
+    val instructions: String,
+    val input: List<DeepSeekResponsesMessage>,
+    val tools: List<DeepSeekResponsesTool>,
+    @SerialName("tool_choice")
+    val toolChoice: JsonElement,
+    @SerialName("max_output_tokens")
+    val maxOutputTokens: Int,
+    val reasoning: DeepSeekResponsesReasoning,
+)
+
+@Serializable
+private data class DeepSeekResponsesMessage(
+    val role: String,
+    val content: String,
+)
+
+@Serializable
+private data class DeepSeekResponsesTool(
+    val type: String,
+)
+
+@Serializable
+private data class DeepSeekResponsesReasoning(
+    val effort: String,
+)
@@ -1,5 +1,8 @@
 package com.osglab.account.features.gateway.providers.volcengine

+import com.osglab.account.features.gateway.credentials.GatewayCredentialProvider
+import com.osglab.account.features.gateway.credentials.ProviderApiKeyResolver
+import com.osglab.account.features.gateway.credentials.StaticProviderApiKeyResolver
 import com.osglab.account.features.gateway.models.AsrGatewayOptions
 import com.osglab.account.features.gateway.models.AsrProviderRequest
 import com.osglab.account.features.gateway.models.GatewayCapability
@@ -103,6 +106,8 @@ class KtorVolcengineAsrTransport(
        ignoreUnknownKeys = true
        explicitNulls = false
    },
+    private val credentialResolver: ProviderApiKeyResolver =
+        StaticProviderApiKeyResolver(volcengineApiKey = config.apiKey),
 ) : VolcengineAsrTransport, VolcengineStreamingClient {
    override suspend fun transcribe(
        request: AsrProviderRequest,
@@ -131,6 +136,9 @@ class KtorVolcengineAsrTransport(
        var outputBytes = 0L
        var frameCount = 0
        val sequenceValidator = SaucSequenceValidator()
+        // Resolve once per connection so an in-flight stream keeps the key it started with.
+        val apiKey = credentialResolver.resolve(GatewayCredentialProvider.VOLCENGINE)
+            ?.takeIf(String::isNotBlank)

        withTimeout(config.responseTimeoutMillis) {
            client.webSocket(
@@ -140,7 +148,6 @@ class KtorVolcengineAsrTransport(
                    headers.append("X-Api-Request-Id", providerRequestId)
                    headers.append("X-Api-Connect-Id", providerRequestId)
                    headers.append("X-Api-Sequence", "-1")
-                    val apiKey = config.apiKey?.takeIf(String::isNotBlank)
                    if (apiKey != null) {
                        headers.append("X-Api-Key", apiKey)
                    } else {
@@ -458,7 +458,7 @@ private suspend fun ApplicationCall.respondGatewayFailure(
        is OobeFeatureAlreadyUsedException -> respondGatewayError(
            HttpStatusCode.Conflict,
            "oobe_feature_already_used",
-            "This OOBE feature has already been used successfully",
+            "This OOBE feature has already been used successfully in this session",
            requestId,
        )

@@ -61,12 +61,21 @@ class GatewayTaskPolicyResolver(
            GatewayTaskKind.AI_QUESTION -> reasoningPolicy(
                taskKind = taskKind,
                effort = config.aiReasoningEffort,
+                webSearch = GatewayWebSearchMode.ALLOWED,
+                maxOutputTokens = minOf(requestedMaxOutputTokens, config.reasoningMaxOutputTokens),
+            )
+
+            GatewayTaskKind.CURRENT_INFORMATION_QUESTION -> reasoningPolicy(
+                taskKind = taskKind,
+                effort = config.aiReasoningEffort,
+                webSearch = GatewayWebSearchMode.REQUIRED,
                maxOutputTokens = minOf(requestedMaxOutputTokens, config.reasoningMaxOutputTokens),
            )

            GatewayTaskKind.AGENT_PLANNING -> reasoningPolicy(
                taskKind = taskKind,
                effort = config.agentReasoningEffort,
+                webSearch = GatewayWebSearchMode.DISABLED,
                maxOutputTokens = minOf(requestedMaxOutputTokens, config.reasoningMaxOutputTokens),
            )
        }
@@ -89,13 +98,14 @@ class GatewayTaskPolicyResolver(
    private fun reasoningPolicy(
        taskKind: GatewayTaskKind,
        effort: GatewayReasoningEffort,
+        webSearch: GatewayWebSearchMode,
        maxOutputTokens: Int,
    ) = GatewayTaskExecutionPolicy(
        taskKind = taskKind,
        modelProfile = GatewayModelProfile.REASONING,
        thinking = GatewayThinkingMode.ENABLED,
        reasoningEffort = effort,
-        webSearch = GatewayWebSearchMode.DISABLED,
+        webSearch = webSearch,
        tools = GatewayToolsMode.DISABLED,
        allowEmptyContentRetry = true,
        maxOutputTokens = maxOutputTokens,
@@ -125,6 +135,7 @@ class GatewayTaskPolicyResolver(
        )
        val AI_TASKS = setOf(
            GatewayTaskKind.AI_QUESTION,
+            GatewayTaskKind.CURRENT_INFORMATION_QUESTION,
            GatewayTaskKind.CLIPBOARD_TRANSFORM,
            GatewayTaskKind.CUSTOM_SKILL,
        )
@@ -1,6 +1,7 @@
 package com.osglab.account.features.gateway.services

 import com.osglab.account.features.gateway.models.AsrProviderRequest
+import com.osglab.account.features.gateway.models.GatewayWebSearchMode
 import com.osglab.account.features.gateway.models.ProviderRequest
 import com.osglab.account.features.gateway.models.TextProviderRequest
 import com.osglab.account.features.gateway.models.UsageMeter
@@ -29,10 +30,18 @@ object ConservativeGatewayUsageEstimator : GatewayUsageEstimator {
        // covers server-controlled system messages and chat framing.
        val inputBytes = request.input.encodeToByteArray().size.toLong()
        val contextBytes = request.context?.encodeToByteArray()?.size?.toLong() ?: 0L
-        val input = Math.addExact(
+        val requestInput = Math.addExact(
            Math.addExact(inputBytes, contextBytes),
            LLM_PROMPT_OVERHEAD_TOKENS,
        )
+        // Responses web search injects provider-controlled result context that
+        // is included in input-token usage. Reserve a bounded allowance so
+        // settlement remains token-based without charging a separate search fee.
+        val input = if (request.executionPolicy.webSearch == GatewayWebSearchMode.DISABLED) {
+            requestInput
+        } else {
+            Math.addExact(requestInput, WEB_SEARCH_INPUT_TOKEN_ALLOWANCE)
+        }
        val output = request.maxOutputTokens.toLong()
        return ProviderUsageEstimate(
            meter = UsageMeter.LLM_TOKEN,
@@ -43,4 +52,5 @@ object ConservativeGatewayUsageEstimator : GatewayUsageEstimator {
    }

    private const val LLM_PROMPT_OVERHEAD_TOKENS = 256L
+    private const val WEB_SEARCH_INPUT_TOKEN_ALLOWANCE = 32_000L
 }
@@ -52,6 +52,7 @@ private object OobeRefreshTokensTable : Table("oobe_gateway_refresh_tokens") {
 }

 private object OobeClaimsTable : Table("oobe_gateway_claims") {
+    val grantId = varchar("grant_id", 36)
    val subjectId = varchar("subject_id", 36)
    val feature = varchar("feature", 32)
    val requestId = varchar("request_id", 64)
@@ -59,7 +60,7 @@ private object OobeClaimsTable : Table("oobe_gateway_claims") {
    val expiresAt = timestamp("expires_at")
    val createdAt = timestamp("created_at")
    val updatedAt = timestamp("updated_at")
-    override val primaryKey = PrimaryKey(subjectId, feature)
+    override val primaryKey = PrimaryKey(grantId, feature)
 }

 private object OobeProviderRequestsTable : Table("oobe_provider_requests") {
@@ -236,7 +237,7 @@ class ExposedOobeRepository(
        expiresAt: Instant,
        now: Instant,
    ): OobeRequestClaim? = databaseFactory.query {
-        val key = claimKey(request.subjectId, request.feature.name)
+        val key = claimKey(request.grantId, request.feature.name)
        val reclaimed = OobeClaimsTable.update({
            key and
                (OobeClaimsTable.status eq CLAIMED) and
@@ -247,6 +248,7 @@ class ExposedOobeRepository(
            it[updatedAt] = now
        } == 1
        val inserted = !reclaimed && OobeClaimsTable.insertIgnore {
+            it[grantId] = request.grantId
            it[subjectId] = request.subjectId
            it[feature] = request.feature.name
            it[requestId] = request.requestId
@@ -269,7 +271,7 @@ class ExposedOobeRepository(
            it[createdAt] = now
        }.insertedCount == 1
        if (!auditInserted) throw OobeRequestAlreadyClaimedException()
-        OobeRequestClaim(request.subjectId, request.feature, request.requestId)
+        OobeRequestClaim(request.subjectId, request.grantId, request.feature, request.requestId)
    }

    override suspend fun markStarted(claim: OobeRequestClaim) {
@@ -280,7 +282,7 @@ class ExposedOobeRepository(
        databaseFactory.query {
            val now = clock.instant()
            val claimChanged = OobeClaimsTable.update({
-                claimKey(claim.subjectId, claim.feature.name) and
+                claimKey(claim.grantId, claim.feature.name) and
                    (OobeClaimsTable.requestId eq claim.requestId) and
                    (OobeClaimsTable.status eq CLAIMED)
            }) {
@@ -308,7 +310,7 @@ class ExposedOobeRepository(
    override suspend fun release(claim: OobeRequestClaim, errorCode: String) {
        databaseFactory.query {
            OobeClaimsTable.deleteWhere {
-                claimKey(claim.subjectId, claim.feature.name) and
+                claimKey(claim.grantId, claim.feature.name) and
                    (OobeClaimsTable.requestId eq claim.requestId) and
                    (OobeClaimsTable.status eq CLAIMED)
            }
@@ -332,7 +334,7 @@ class ExposedOobeRepository(
            // Fail closed: an uncertain provider outcome must never become
            // reclaimable after the temporary claim TTL.
            OobeClaimsTable.update({
-                claimKey(claim.subjectId, claim.feature.name) and
+                claimKey(claim.grantId, claim.feature.name) and
                    (OobeClaimsTable.requestId eq claim.requestId) and
                    (OobeClaimsTable.status eq CLAIMED)
            }) {
@@ -377,8 +379,8 @@ private fun org.jetbrains.exposed.v1.core.ResultRow.toStoredRefresh(grant: OobeG
        expiresAt = this[OobeRefreshTokensTable.expiresAt],
    )

-private fun claimKey(subjectId: String, feature: String) =
-    (OobeClaimsTable.subjectId eq subjectId) and (OobeClaimsTable.feature eq feature)
+private fun claimKey(grantId: String, feature: String) =
+    (OobeClaimsTable.grantId eq grantId) and (OobeClaimsTable.feature eq feature)

 private fun requestKey(claim: OobeRequestClaim) =
    (OobeProviderRequestsTable.subjectId eq claim.subjectId) and
@@ -103,6 +103,7 @@ sealed interface OobeRefreshRotationResult {

 data class OobeRequestClaim(
    val subjectId: String,
+    val grantId: String,
    val feature: OobeFeature,
    val requestId: String,
 )
@@ -39,4 +39,4 @@ class OobeRequestAlreadyClaimedException :
    RuntimeException("The OOBE provider request ID has already been used")

 class OobeFeatureAlreadyUsedException(val feature: com.osglab.account.features.gateway.models.OobeFeature) :
-    RuntimeException("The OOBE feature ${feature.name.lowercase()} has already been used")
+    RuntimeException("The OOBE feature ${feature.name.lowercase()} has already been used in this grant")