diff --git a/apps/opik-backend/config.yml b/apps/opik-backend/config.yml index bba61f5d02b..c6709b14057 100644 --- a/apps/opik-backend/config.yml +++ b/apps/opik-backend/config.yml @@ -1189,6 +1189,26 @@ partitionMetrics: lwdTables: ${PARTITION_METRICS_LWD_TABLES:-traces,spans} # LLM providers client configuration +# Configuration for dynamic token auth on custom LLM providers (auth_config recipes) +llmProviderTokenAuth: + # Default: 0.25 + # Description: Refresh the cached token once its remaining lifetime drops below this fraction of the total lifetime + refreshFraction: ${LLM_PROVIDER_TOKEN_AUTH_REFRESH_FRACTION:-0.25} + # Default: 10s + # Description: Timeout for the token fetch call to the configured auth service + fetchTimeout: ${LLM_PROVIDER_TOKEN_AUTH_FETCH_TIMEOUT:-10s} + # Default: 15s + # Description: Lease time for the cross-pod single-flight lock around a token fetch + lockTimeout: ${LLM_PROVIDER_TOKEN_AUTH_LOCK_TIMEOUT:-15s} + # Default: 1000000 + # Description: Maximum accepted size (in characters) of a token endpoint reply. Capped at 10000000 + maxResponseChars: ${LLM_PROVIDER_TOKEN_AUTH_MAX_RESPONSE_CHARS:-1000000} + # Default: strict + # Description: SSRF guard on the token URL. "strict" refuses non-HTTPS and private/internal + # destinations; "relaxed" allows internal gateways. Strict by default so a missing override fails + # loudly instead of exposing the deployment + destinationGuard: ${LLM_PROVIDER_TOKEN_AUTH_DESTINATION_GUARD:-strict} + llmProviderClient: # Default: 3 # Description: Max amount of attempts to reach the LLM provider before giving up diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKey.java b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKey.java index 648d96aae2c..33319e8ef02 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKey.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKey.java @@ -8,6 +8,7 @@ import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import com.fasterxml.jackson.databind.annotation.JsonNaming; import io.swagger.v3.oas.annotations.media.Schema; +import jakarta.validation.Valid; import jakarta.validation.constraints.NotNull; import jakarta.validation.constraints.Pattern; import jakarta.validation.constraints.Size; @@ -41,6 +42,10 @@ public record ProviderApiKey( @JsonView({View.Public.class, View.Write.class}) Map configuration, @JsonView({View.Public.class, View.Write.class}) @Pattern(regexp = NULL_OR_NOT_BLANK, message = "must not be blank") String baseUrl, + @JsonView({View.Public.class, + View.Write.class}) @Valid @Schema(description = "Dynamic token auth recipe. When set, Opik fetches a short-lived bearer from the configured auth service instead of using a static api_key. " + + + "Only supported for custom providers. Secret credential values read back masked.") ProviderAuthConfig authConfig, @JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) Instant createdAt, @JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) String createdBy, @JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) Instant lastUpdatedAt, @@ -57,6 +62,7 @@ public String toString() { ", name='" + name + '\'' + ", providerName='" + providerName + '\'' + ", headers=" + headers + + ", authConfig=" + authConfig + ", baseUrl='" + baseUrl + '\'' + ", createdAt=" + createdAt + ", createdBy='" + createdBy + '\'' + diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKeyUpdate.java b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKeyUpdate.java index ffcc4e5b80c..67ba001384b 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKeyUpdate.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKeyUpdate.java @@ -7,6 +7,7 @@ import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import com.fasterxml.jackson.databind.annotation.JsonNaming; import io.swagger.v3.oas.annotations.media.Schema; +import jakarta.validation.Valid; import jakarta.validation.constraints.Size; import lombok.Builder; @@ -27,7 +28,11 @@ public record ProviderApiKeyUpdate( @JsonView({ProviderApiKey.View.Public.class, ProviderApiKey.View.Write.class}) Map headers, @JsonView({ProviderApiKey.View.Public.class, ProviderApiKey.View.Write.class}) Map configuration, - @JsonView({ProviderApiKey.View.Public.class, ProviderApiKey.View.Write.class}) String baseUrl) { + @JsonView({ProviderApiKey.View.Public.class, ProviderApiKey.View.Write.class}) String baseUrl, + @JsonView({ProviderApiKey.View.Public.class, + ProviderApiKey.View.Write.class}) @Valid @Schema(description = "Dynamic token auth recipe. Send the '" + + ProviderAuthConfig.SECRET_SENTINEL + + "' sentinel as a credential value to keep the stored secret; send an empty object to clear the auth config.") ProviderAuthConfig authConfig) { @Override public String toString() { diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderAuthCheck.java b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderAuthCheck.java new file mode 100644 index 00000000000..9628f40dbca --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderAuthCheck.java @@ -0,0 +1,31 @@ +package com.comet.opik.api; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.databind.PropertyNamingStrategies; +import com.fasterxml.jackson.databind.annotation.JsonNaming; +import io.swagger.v3.oas.annotations.media.Schema; +import jakarta.validation.Valid; +import lombok.Builder; + +import java.util.UUID; + +/** + * Test-connection request for a provider's dynamic token auth. Two modes, both executed by the + * backend: by {@code provider_id} alone the stored recipe is used (secrets never transit the + * browser); with an {@code auth_config} the submitted values are used, resolving + * {@code __SECRET__} sentinels against the stored recipe when {@code provider_id} is also given. + */ +@Builder(toBuilder = true) +@JsonIgnoreProperties(ignoreUnknown = true) +@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class) +public record ProviderAuthCheck( + @Schema(description = "Test the stored auth config of this provider; also the sentinel-resolution target when auth_config is sent") UUID providerId, + @Valid @Schema(description = "Auth config to test as-submitted; omit to test the stored one") ProviderAuthConfig authConfig) { + + @Builder(toBuilder = true) + @JsonIgnoreProperties(ignoreUnknown = true) + @JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class) + public record Result( + @Schema(description = "Lifetime of the fetched token in seconds; the token itself is never returned") long lifetimeSeconds) { + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderAuthConfig.java b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderAuthConfig.java new file mode 100644 index 00000000000..3ffa335756e --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderAuthConfig.java @@ -0,0 +1,142 @@ +package com.comet.opik.api; + +import com.comet.opik.utils.ValidationUtils; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonValue; +import com.fasterxml.jackson.annotation.JsonView; +import com.fasterxml.jackson.databind.PropertyNamingStrategies; +import com.fasterxml.jackson.databind.annotation.JsonNaming; +import io.swagger.v3.oas.annotations.media.Schema; +import jakarta.validation.Valid; +import jakarta.validation.constraints.Max; +import jakarta.validation.constraints.Min; +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.Size; +import lombok.Builder; +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import org.apache.commons.collections4.CollectionUtils; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +/** + * Recipe describing how to fetch a short-lived bearer token before calling a custom LLM provider. + * Stored as a single AES-GCM-encrypted JSON document in {@code llm_provider_api_key.auth_config}; + * a row without one behaves exactly as before (static api_key auth). + */ +@Builder(toBuilder = true) +@JsonIgnoreProperties(ignoreUnknown = true) +@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class) +public record ProviderAuthConfig( + @JsonView({ + ProviderApiKey.View.Public.class, + ProviderApiKey.View.Write.class}) @Schema(description = "Auth service URL the credentials are sent to", example = "https://developer.api.example.com/authentication/v1/token") String tokenUrl, + @JsonView({ProviderApiKey.View.Public.class, + ProviderApiKey.View.Write.class}) @Schema(description = "How credentials are sent: form body (default), JSON body, or basic auth (id/secret in an HTTP Basic header, remaining fields in the form body)") SendAs sendAs, + @JsonView({ProviderApiKey.View.Public.class, + ProviderApiKey.View.Write.class}) @Valid @Schema(description = "Fields sent to the token URL. Values flagged as secret are write-only: they read back as the '" + + ProviderAuthConfig.SECRET_SENTINEL + "' sentinel") List credentials, + @JsonView({ProviderApiKey.View.Public.class, + ProviderApiKey.View.Write.class}) @Size(max = 250) @Schema(description = "Field holding the token in the reply; dot-path for nested replies", example = "access_token") String tokenField, + @JsonView({ProviderApiKey.View.Public.class, + ProviderApiKey.View.Write.class}) @Size(max = 250) @Schema(description = "Field holding the token lifetime in seconds in the reply; dot-path for nested replies", example = "expires_in") String expiresField, + @JsonView({ProviderApiKey.View.Public.class, + ProviderApiKey.View.Write.class}) @Min(0) @Max(31_536_000) @Schema(description = "Lifetime in seconds assumed when the reply doesn't state one, capped at one year; 0 means such tokens are not cached (fetched per call)") Long fallbackTtlSeconds) { + + public static final String SECRET_SENTINEL = "__SECRET__"; + private static final ProviderAuthConfig EMPTY = ProviderAuthConfig.builder().build(); + + @Builder(toBuilder = true) + @JsonIgnoreProperties(ignoreUnknown = true) + @JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class) + public record Credential( + @JsonView({ + ProviderApiKey.View.Public.class, + ProviderApiKey.View.Write.class}) @NotBlank @Size(max = 250) String key, + @JsonView({ProviderApiKey.View.Public.class, ProviderApiKey.View.Write.class}) String value, + @JsonView({ProviderApiKey.View.Public.class, + ProviderApiKey.View.Write.class}) @Schema(description = "Secret values are encrypted at rest and never read back; once true it cannot be unset") boolean secret) { + + @Override + public String toString() { + return "Credential{key='" + key + "', value='*******', secret=" + secret + '}'; + } + } + + @Getter + @RequiredArgsConstructor + public enum SendAs { + FORM("form"), + JSON("json"), + BASIC("basic"), + ; + + @JsonValue + private final String value; + + @JsonCreator + public static SendAs fromString(String value) { + return Arrays.stream(values()) + .filter(sendAs -> sendAs.value.equals(value)) + .findFirst() + .orElseThrow(() -> new IllegalArgumentException("Unknown send_as '%s'".formatted(value))); + } + } + + /** + * A literal empty object ({@code {}}) is the API convention for clearing the auth config on + * update, mirroring how an empty headers map clears headers. + */ + @JsonIgnore + public boolean isEmpty() { + return equals(EMPTY); + } + + /** + * Requiredness rules shared by the create validator and the update path. Field-level Jakarta + * annotations can't carry these because the empty clear-object must pass bean validation. + */ + @JsonIgnore + public List validationErrors() { + var errors = new ArrayList(); + if (!ValidationUtils.isAbsoluteUri(tokenUrl)) { + errors.add("auth_config.token_url must be a valid absolute URI"); + } + if (CollectionUtils.isEmpty(credentials)) { + errors.add("auth_config.credentials must not be empty"); + } + return errors; + } + + /** + * Copy safe to return from the API: secret values are replaced with {@link #SECRET_SENTINEL}. + */ + public ProviderAuthConfig mask() { + if (CollectionUtils.isEmpty(credentials)) { + return this; + } + return toBuilder() + .credentials(credentials.stream() + .map(credential -> credential.secret() + ? credential.toBuilder().value(SECRET_SENTINEL).build() + : credential) + .toList()) + .build(); + } + + @Override + public String toString() { + return "ProviderAuthConfig{" + + "tokenUrl='" + tokenUrl + '\'' + + ", sendAs=" + sendAs + + ", credentials=" + credentials + + ", tokenField='" + tokenField + '\'' + + ", expiresField='" + expiresField + '\'' + + ", fallbackTtlSeconds=" + fallbackTtlSeconds + + '}'; + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResource.java index daa4170fc2f..a08cde547ae 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResource.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResource.java @@ -4,11 +4,13 @@ import com.comet.opik.api.BatchDelete; import com.comet.opik.api.ProviderApiKey; import com.comet.opik.api.ProviderApiKeyUpdate; +import com.comet.opik.api.ProviderAuthCheck; import com.comet.opik.api.error.ErrorMessage; import com.comet.opik.domain.LlmProviderApiKeyService; import com.comet.opik.infrastructure.auth.RequestContext; import com.comet.opik.infrastructure.auth.RequiredPermissions; import com.comet.opik.infrastructure.auth.WorkspaceUserPermission; +import com.comet.opik.infrastructure.ratelimit.RateLimited; import com.fasterxml.jackson.annotation.JsonView; import io.swagger.v3.oas.annotations.Operation; import io.swagger.v3.oas.annotations.headers.Header; @@ -71,6 +73,9 @@ public Response find() { .apiKey(providerApiKey.apiKey() != null ? maskApiKey(decrypt(providerApiKey.apiKey())) : "null") + .authConfig(providerApiKey.authConfig() != null + ? providerApiKey.authConfig().mask() + : null) .build()) .toList(); @@ -101,6 +106,7 @@ public Response getById(@PathParam("id") UUID id) { return Response.ok().entity(providerApiKey.toBuilder() .apiKey(providerApiKey.apiKey() != null ? maskApiKey(decrypt(providerApiKey.apiKey())) : null) + .authConfig(providerApiKey.authConfig() != null ? providerApiKey.authConfig().mask() : null) .build()).build(); } @@ -147,6 +153,30 @@ public Response updateApiKey(@PathParam("id") UUID id, return Response.noContent().build(); } + @POST + @Path("/auth-config/test") + @RateLimited + @RequiredPermissions(WorkspaceUserPermission.AI_PROVIDER_UPDATE) + @Operation(operationId = "testLlmProviderAuthConfig", summary = "Test a provider's dynamic token auth", description = "Runs the token fetch once, backend-side, and reports the token lifetime. The token itself is never returned. " + + + "Send provider_id to test the stored config, auth_config to test submitted values, or both to resolve secret sentinels against the stored config.", responses = { + @ApiResponse(responseCode = "200", description = "Token fetched", content = @Content(schema = @Schema(implementation = ProviderAuthCheck.Result.class))), + @ApiResponse(responseCode = "400", description = "Bad Request", content = @Content(schema = @Schema(implementation = ErrorMessage.class))), + @ApiResponse(responseCode = "403", description = "Access forbidden", content = @Content(schema = @Schema(implementation = ErrorMessage.class))), + @ApiResponse(responseCode = "404", description = "Not found", content = @Content(schema = @Schema(implementation = ErrorMessage.class))) + }) + public Response testAuthConfig( + @NotNull @RequestBody(content = @Content(schema = @Schema(implementation = ProviderAuthCheck.class))) @Valid ProviderAuthCheck providerAuthTest) { + String workspaceId = requestContext.get().getWorkspaceId(); + + log.info("Testing LLM provider auth config on workspace_id '{}'", workspaceId); + ProviderAuthCheck.Result result = llmProviderApiKeyService.testAuthConfig(providerAuthTest, workspaceId); + log.info("Tested LLM provider auth config on workspace_id '{}': token received, lifetime '{}'s", + workspaceId, result.lifetimeSeconds()); + + return Response.ok(result).build(); + } + @POST @Path("/delete") @RequiredPermissions(WorkspaceUserPermission.AI_PROVIDER_UPDATE) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/validation/ProviderApiKeyValidator.java b/apps/opik-backend/src/main/java/com/comet/opik/api/validation/ProviderApiKeyValidator.java index f458de5c767..155be668c85 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/validation/ProviderApiKeyValidator.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/validation/ProviderApiKeyValidator.java @@ -1,10 +1,14 @@ package com.comet.opik.api.validation; import com.comet.opik.api.ProviderApiKey; +import com.comet.opik.api.ProviderAuthConfig; import com.comet.opik.infrastructure.EncryptionUtils; import jakarta.validation.ConstraintValidator; import jakarta.validation.ConstraintValidatorContext; +import java.util.List; +import java.util.Optional; + import static org.apache.commons.lang3.StringUtils.isBlank; public class ProviderApiKeyValidator @@ -37,7 +41,15 @@ public boolean isValid(ProviderApiKey providerApiKey, ConstraintValidatorContext } // If provider supports naming, no need to validate api key - return true; + return isValidAuthConfig(providerApiKey, context); + } + + if (providerApiKey.authConfig() != null) { + context.buildConstraintViolationWithTemplate( + "auth_config is only supported for providers with provider_name (custom LLM, Bedrock, Ollama)") + .addPropertyNode("authConfig") + .addConstraintViolation(); + return false; } // Validate API key for non-custom providers @@ -50,4 +62,40 @@ public boolean isValid(ProviderApiKey providerApiKey, ConstraintValidatorContext return true; } + + private boolean isValidAuthConfig(ProviderApiKey providerApiKey, ConstraintValidatorContext context) { + var authConfig = providerApiKey.authConfig(); + if (authConfig == null) { + return true; + } + + boolean valid = true; + for (String error : authConfig.validationErrors()) { + context.buildConstraintViolationWithTemplate(error) + .addPropertyNode("authConfig") + .addConstraintViolation(); + valid = false; + } + + // The sentinel means "keep the stored secret" — meaningless on create, where nothing is stored yet + boolean hasSentinel = Optional.ofNullable(authConfig.credentials()).orElse(List.of()).stream() + .anyMatch(credential -> ProviderAuthConfig.SECRET_SENTINEL.equals(credential.value())); + if (hasSentinel) { + context.buildConstraintViolationWithTemplate( + "auth_config credential values must not be the '%s' sentinel on create" + .formatted(ProviderAuthConfig.SECRET_SENTINEL)) + .addPropertyNode("authConfig") + .addConstraintViolation(); + valid = false; + } + + if (providerApiKey.apiKey() != null && !isBlank(EncryptionUtils.decrypt(providerApiKey.apiKey()))) { + context.buildConstraintViolationWithTemplate("api_key must not be set when auth_config is set") + .addPropertyNode("apiKey") + .addConstraintViolation(); + valid = false; + } + + return valid; + } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyDAO.java index acffdf62d01..8f6a38c0139 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyDAO.java @@ -2,7 +2,9 @@ import com.comet.opik.api.ProviderApiKey; import com.comet.opik.api.ProviderApiKeyUpdate; +import com.comet.opik.api.ProviderAuthConfig; import com.comet.opik.infrastructure.db.MapFlatArgumentFactory; +import com.comet.opik.infrastructure.db.ProviderAuthConfigArgumentFactory; import com.comet.opik.infrastructure.db.UUIDArgumentFactory; import org.jdbi.v3.sqlobject.config.RegisterArgumentFactory; import org.jdbi.v3.sqlobject.config.RegisterColumnMapper; @@ -21,14 +23,15 @@ @RegisterRowMapper(ProviderApiKeyRowMapper.class) @RegisterArgumentFactory(UUIDArgumentFactory.class) @RegisterArgumentFactory(MapFlatArgumentFactory.class) +@RegisterArgumentFactory(ProviderAuthConfigArgumentFactory.class) @RegisterColumnMapper(MapFlatArgumentFactory.class) public interface LlmProviderApiKeyDAO { String NULL_SENTINEL = "__NULL__"; - @SqlUpdate("INSERT INTO llm_provider_api_key (id, provider, workspace_id, api_key, name, provider_name, created_by, last_updated_by, headers, base_url, configuration) " + @SqlUpdate("INSERT INTO llm_provider_api_key (id, provider, workspace_id, api_key, name, provider_name, created_by, last_updated_by, headers, base_url, configuration, auth_config) " + - "VALUES (:bean.id, :bean.provider, :workspaceId, :bean.apiKey, :bean.name, :providerName, :bean.createdBy, :bean.lastUpdatedBy, :bean.headers, :bean.baseUrl, :bean.configuration)") + "VALUES (:bean.id, :bean.provider, :workspaceId, :bean.apiKey, :bean.name, :providerName, :bean.createdBy, :bean.lastUpdatedBy, :bean.headers, :bean.baseUrl, :bean.configuration, :bean.authConfig)") void saveInternal(@Bind("workspaceId") String workspaceId, @Bind("providerName") String providerName, @BindMethods("bean") ProviderApiKey providerApiKey); @@ -50,12 +53,17 @@ default void save(String workspaceId, ProviderApiKey providerApiKey) { "headers = CASE WHEN :bean.headers IS NULL THEN headers ELSE :bean.headers END, " + "base_url = CASE WHEN :bean.baseUrl IS NULL THEN base_url ELSE :bean.baseUrl END, " + "configuration = CASE WHEN :bean.configuration IS NULL THEN configuration ELSE :bean.configuration END, " + + "auth_config = CASE WHEN :clearAuthConfig THEN NULL " + + "WHEN :authConfig IS NULL THEN auth_config " + + "ELSE :authConfig END, " + "last_updated_by = :lastUpdatedBy " + "WHERE id = :id AND workspace_id = :workspaceId") void update(@Bind("id") UUID id, @Bind("workspaceId") String workspaceId, @Bind("lastUpdatedBy") String lastUpdatedBy, - @BindMethods("bean") ProviderApiKeyUpdate providerApiKeyUpdate); + @BindMethods("bean") ProviderApiKeyUpdate providerApiKeyUpdate, + @Bind("clearAuthConfig") boolean clearAuthConfig, + @Bind("authConfig") ProviderAuthConfig authConfig); @SqlQuery("SELECT * FROM llm_provider_api_key WHERE id = :id AND workspace_id = :workspaceId") ProviderApiKey findById(@Bind("id") UUID id, @Bind("workspaceId") String workspaceId); diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyService.java index 6b6b7268267..3450286c6cd 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyService.java @@ -3,17 +3,24 @@ import com.comet.opik.api.LlmProvider; import com.comet.opik.api.ProviderApiKey; import com.comet.opik.api.ProviderApiKeyUpdate; +import com.comet.opik.api.ProviderAuthCheck; +import com.comet.opik.api.ProviderAuthConfig; import com.comet.opik.api.error.EntityAlreadyExistsException; import com.comet.opik.api.error.ErrorMessage; +import com.comet.opik.infrastructure.EncryptionUtils; import com.comet.opik.infrastructure.OpikConfiguration; +import com.comet.opik.infrastructure.llm.customllm.AuthTokenException; +import com.comet.opik.infrastructure.llm.customllm.AuthTokenProvider; import com.google.inject.ImplementedBy; import jakarta.inject.Inject; import jakarta.inject.Singleton; +import jakarta.ws.rs.BadRequestException; import jakarta.ws.rs.NotFoundException; import jakarta.ws.rs.core.Response; import lombok.NonNull; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; import org.jdbi.v3.core.statement.UnableToExecuteStatementException; import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate; @@ -23,6 +30,8 @@ import java.util.Map; import java.util.Set; import java.util.UUID; +import java.util.function.Function; +import java.util.stream.Collectors; import static com.comet.opik.infrastructure.FreeModelConfig.FREE_MODEL_PROVIDER_ID; import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.READ_ONLY; @@ -40,6 +49,8 @@ public interface LlmProviderApiKeyService { void updateApiKey(UUID id, ProviderApiKeyUpdate providerApiKeyUpdate, String userName, String workspaceId); void delete(Set ids, String workspaceId); + + ProviderAuthCheck.Result testAuthConfig(ProviderAuthCheck request, String workspaceId); } @Slf4j @@ -51,6 +62,7 @@ class LlmProviderApiKeyServiceImpl implements LlmProviderApiKeyService { private final @NonNull IdGenerator idGenerator; private final @NonNull TransactionTemplate template; private final @NonNull OpikConfiguration configuration; + private final @NonNull AuthTokenProvider authTokenProvider; @Override public ProviderApiKey find(@NonNull UUID id, @NonNull String workspaceId) { @@ -147,15 +159,128 @@ public void updateApiKey(@NonNull UUID id, @NonNull ProviderApiKeyUpdate provide ProviderApiKey providerApiKey = repository.fetch(id, workspaceId) .orElseThrow(this::createNotFoundError); + var authConfigUpdate = resolveAuthConfigUpdate(providerApiKeyUpdate, providerApiKey); + repository.update(providerApiKey.id(), workspaceId, userName, - providerApiKeyUpdate); + providerApiKeyUpdate, + authConfigUpdate.clear(), + authConfigUpdate.authConfig()); return null; }); } + @Override + public ProviderAuthCheck.Result testAuthConfig(@NonNull ProviderAuthCheck request, @NonNull String workspaceId) { + ProviderAuthConfig resolved = resolveAuthConfigForTest(request, workspaceId); + try { + return new ProviderAuthCheck.Result(authTokenProvider.testFetch(resolved)); + } catch (AuthTokenException exception) { + // the message is user-facing and redacted by contract; a failed test is the caller's + // configuration problem, not a server error + log.info("Auth config test failed on workspace_id '{}': {}", workspaceId, exception.getMessage()); + throw new BadRequestException(exception.getMessage()); + } + } + + private ProviderAuthConfig resolveAuthConfigForTest(ProviderAuthCheck request, String workspaceId) { + ProviderAuthConfig incoming = request.authConfig(); + if (incoming == null || incoming.isEmpty()) { + if (request.providerId() == null) { + throw new BadRequestException("either provider_id or auth_config must be provided"); + } + ProviderAuthConfig stored = find(request.providerId(), workspaceId).authConfig(); + if (stored == null) { + throw new BadRequestException("the provider has no auth_config to test"); + } + return stored; + } + + var errors = incoming.validationErrors(); + if (!errors.isEmpty()) { + throw new BadRequestException(String.join("; ", errors)); + } + ProviderAuthConfig stored = request.providerId() != null + ? find(request.providerId(), workspaceId).authConfig() + : null; + return mergeSecretSentinels(incoming, stored); + } + + private record AuthConfigUpdate(boolean clear, ProviderAuthConfig authConfig) { + } + + /** + * Maps the update's auth_config to the DAO binds: absent -> keep (null value, no clear), + * empty object -> clear, otherwise validate + resolve {@code __SECRET__} sentinels against + * the stored recipe -> set. + */ + private AuthConfigUpdate resolveAuthConfigUpdate(ProviderApiKeyUpdate update, ProviderApiKey stored) { + ProviderAuthConfig incoming = update.authConfig(); + if (incoming == null) { + validateNoStaticKeyConflict( + update.apiKey() != null ? update.apiKey() : stored.apiKey(), stored.authConfig()); + return new AuthConfigUpdate(false, null); + } + if (incoming.isEmpty()) { + return new AuthConfigUpdate(true, null); + } + if (!stored.provider().supportsProviderName()) { + throw new BadRequestException("auth_config is only supported for providers with provider_name"); + } + var errors = incoming.validationErrors(); + if (!errors.isEmpty()) { + throw new BadRequestException(String.join("; ", errors)); + } + validateNoStaticKeyConflict(update.apiKey() != null ? update.apiKey() : stored.apiKey(), incoming); + var merged = mergeSecretSentinels(incoming, stored.authConfig()); + return new AuthConfigUpdate(false, merged); + } + + private void validateNoStaticKeyConflict(String encryptedApiKey, ProviderAuthConfig authConfig) { + if (authConfig != null && encryptedApiKey != null + && StringUtils.isNotBlank(EncryptionUtils.decrypt(encryptedApiKey))) { + throw new BadRequestException( + "api_key and auth_config cannot both be set; clear one of them (send auth_config as an empty object to clear it)"); + } + } + + /** + * Resolves {@code __SECRET__} sentinel values ("keep the stored secret") and enforces that a + * secret flag can never be unset once saved. Sentinels are only accepted where the stored + * recipe holds a secret under the same key — so a literal password that happens to look masked + * can never be confused with "unchanged". + */ + private ProviderAuthConfig mergeSecretSentinels(ProviderAuthConfig incoming, ProviderAuthConfig stored) { + Map storedByKey = stored == null || stored.credentials() == null + ? Map.of() + : stored.credentials().stream() + .collect(Collectors.toMap( + ProviderAuthConfig.Credential::key, Function.identity(), (first, second) -> second)); + + var mergedCredentials = incoming.credentials().stream() + .map(credential -> { + var storedCredential = storedByKey.get(credential.key()); + boolean lockedInStore = storedCredential != null && storedCredential.secret(); + if (ProviderAuthConfig.SECRET_SENTINEL.equals(credential.value())) { + if (!lockedInStore) { + throw new BadRequestException( + "credential '%s' uses the secret sentinel but no stored secret exists under that key; provide the value" + .formatted(credential.key())); + } + return credential.toBuilder().value(storedCredential.value()).secret(true).build(); + } + // once saved as secret, a credential stays secret + return lockedInStore && !credential.secret() + ? credential.toBuilder().secret(true).build() + : credential; + }) + .toList(); + + return incoming.toBuilder().credentials(mergedCredentials).build(); + } + @Override public void delete(@NonNull Set ids, @NonNull String workspaceId) { if (ids.isEmpty()) { diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/ProviderApiKeyRowMapper.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/ProviderApiKeyRowMapper.java index 62a1b01e433..2caac058b88 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/ProviderApiKeyRowMapper.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/ProviderApiKeyRowMapper.java @@ -2,7 +2,10 @@ import com.comet.opik.api.LlmProvider; import com.comet.opik.api.ProviderApiKey; +import com.comet.opik.api.ProviderAuthConfig; +import com.comet.opik.infrastructure.EncryptionUtils; import com.comet.opik.infrastructure.db.MapFlatArgumentFactory; +import com.comet.opik.utils.JsonUtils; import lombok.NonNull; import org.jdbi.v3.core.mapper.RowMapper; import org.jdbi.v3.core.statement.StatementContext; @@ -36,6 +39,7 @@ public ProviderApiKey map(@NonNull ResultSet rs, @NonNull StatementContext ctx) .providerName(providerName) .headers(mapMapper.map(rs, "headers", ctx)) .configuration(mapMapper.map(rs, "configuration", ctx)) + .authConfig(readAuthConfig(rs.getString("auth_config"))) .baseUrl(rs.getString("base_url")) .createdAt(rs.getTimestamp("created_at").toInstant()) .createdBy(rs.getString("created_by")) @@ -43,4 +47,11 @@ public ProviderApiKey map(@NonNull ResultSet rs, @NonNull StatementContext ctx) .lastUpdatedBy(rs.getString("last_updated_by")) .build(); } + + private static ProviderAuthConfig readAuthConfig(String encryptedAuthConfig) { + if (encryptedAuthConfig == null) { + return null; + } + return JsonUtils.readValue(EncryptionUtils.decryptGcm(encryptedAuthConfig), ProviderAuthConfig.class); + } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtils.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtils.java index 1beb757994e..dd776224f89 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtils.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtils.java @@ -8,18 +8,27 @@ import javax.crypto.Cipher; import javax.crypto.IllegalBlockSizeException; import javax.crypto.NoSuchPaddingException; +import javax.crypto.spec.GCMParameterSpec; import javax.crypto.spec.SecretKeySpec; import java.nio.charset.StandardCharsets; +import java.security.GeneralSecurityException; import java.security.InvalidKeyException; import java.security.Key; import java.security.NoSuchAlgorithmException; +import java.security.SecureRandom; import java.util.Base64; @UtilityClass public class EncryptionUtils { private static final String ALGO = "AES"; + // AES-GCM for larger, structured payloads (e.g. auth_config JSON): the legacy no-IV mode is + // deterministic, which is acceptable for short random keys but not for predictable plaintext. + private static final String GCM_ALGO = "AES/GCM/NoPadding"; + private static final int GCM_IV_LENGTH = 12; + private static final int GCM_TAG_BITS = 128; + private static final SecureRandom SECURE_RANDOM = new SecureRandom(); private static final Base64.Encoder mimeEncoder = Base64.getMimeEncoder(); private static final Base64.Decoder mimeDecoder = Base64.getMimeDecoder(); private static Key key; @@ -54,6 +63,34 @@ public static String decrypt(@NonNull String encryptedData) { } } + public static String encryptGcm(@NonNull String data) { + try { + byte[] iv = new byte[GCM_IV_LENGTH]; + SECURE_RANDOM.nextBytes(iv); + Cipher cipher = Cipher.getInstance(GCM_ALGO); + cipher.init(Cipher.ENCRYPT_MODE, key, new GCMParameterSpec(GCM_TAG_BITS, iv)); + byte[] encrypted = cipher.doFinal(data.getBytes(StandardCharsets.UTF_8)); + byte[] payload = new byte[iv.length + encrypted.length]; + System.arraycopy(iv, 0, payload, 0, iv.length); + System.arraycopy(encrypted, 0, payload, iv.length, encrypted.length); + return Base64.getEncoder().encodeToString(payload); + } catch (GeneralSecurityException ex) { + throw new SecurityException("Failed to encrypt. " + ex.getMessage(), ex); + } + } + + public static String decryptGcm(@NonNull String encryptedData) { + try { + byte[] payload = Base64.getDecoder().decode(encryptedData); + Cipher cipher = Cipher.getInstance(GCM_ALGO); + cipher.init(Cipher.DECRYPT_MODE, key, new GCMParameterSpec(GCM_TAG_BITS, payload, 0, GCM_IV_LENGTH)); + byte[] decrypted = cipher.doFinal(payload, GCM_IV_LENGTH, payload.length - GCM_IV_LENGTH); + return new String(decrypted, StandardCharsets.UTF_8); + } catch (GeneralSecurityException | IllegalArgumentException ex) { + throw new SecurityException("Failed to decrypt. " + ex.getMessage(), ex); + } + } + public static String maskApiKey(@NonNull String apiKey) { return apiKey.length() <= 12 ? StringUtils.repeat('*', apiKey.length()) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/LlmProviderTokenAuthConfig.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/LlmProviderTokenAuthConfig.java new file mode 100644 index 00000000000..6347ed48efd --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/LlmProviderTokenAuthConfig.java @@ -0,0 +1,54 @@ +package com.comet.opik.infrastructure; + +import com.comet.opik.infrastructure.net.DestinationGuard; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.dropwizard.util.Duration; +import io.dropwizard.validation.MinDuration; +import jakarta.validation.constraints.DecimalMax; +import jakarta.validation.constraints.DecimalMin; +import jakarta.validation.constraints.Max; +import jakarta.validation.constraints.Min; +import jakarta.validation.constraints.NotNull; +import lombok.Data; + +import java.util.concurrent.TimeUnit; + +/** + * Tuning for dynamic token auth on custom LLM providers (the {@code auth_config} recipes). + */ +@Data +public class LlmProviderTokenAuthConfig { + + /** + * A cached token is refreshed once its remaining lifetime drops below this fraction of the + * total lifetime. Proportional rather than absolute so the same default behaves correctly for + * a 60-second token and a 25-hour one. + */ + @JsonProperty + @DecimalMin("0.0") @DecimalMax("0.99") private double refreshFraction = 0.25; + + @JsonProperty + @NotNull @MinDuration(value = 1, unit = TimeUnit.MILLISECONDS) + private Duration fetchTimeout = Duration + .seconds(10); + + /** + * Lease time for the cross-pod single-flight lock around a token fetch. Must comfortably cover + * one fetch, so a pod dying mid-fetch doesn't block the others for long. + */ + @JsonProperty + @NotNull @MinDuration(value = 1, unit = TimeUnit.MILLISECONDS) + private Duration lockTimeout = Duration.seconds(15); + + @JsonProperty + @Min(1) @Max(10_000_000) private int maxResponseChars = 1_000_000; + + /** + * SSRF guard on the token URL: {@code strict} refuses non-HTTPS and private/internal + * destinations; {@code relaxed} allows internal gateways. Strict by default so a missing + * override fails loudly instead of exposing the deployment; the self-hosted distributions + * can ship {@code relaxed} explicitly. + */ + @JsonProperty + @NotNull private DestinationGuard.Mode destinationGuard = DestinationGuard.Mode.STRICT; +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/OpikConfiguration.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/OpikConfiguration.java index 8c0d78de332..c10cd327c48 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/OpikConfiguration.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/OpikConfiguration.java @@ -76,6 +76,9 @@ public class OpikConfiguration extends JobConfiguration { @Valid @NotNull @JsonProperty private LlmProviderClientConfig llmProviderClient = new LlmProviderClientConfig(); + @Valid @NotNull @JsonProperty + private LlmProviderTokenAuthConfig llmProviderTokenAuth = new LlmProviderTokenAuthConfig(); + @Valid @NotNull @JsonProperty private CacheConfiguration cacheManager = new CacheConfiguration(); diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/ProviderAuthConfigArgumentFactory.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/ProviderAuthConfigArgumentFactory.java new file mode 100644 index 00000000000..1f19703d573 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/db/ProviderAuthConfigArgumentFactory.java @@ -0,0 +1,33 @@ +package com.comet.opik.infrastructure.db; + +import com.comet.opik.api.ProviderAuthConfig; +import com.comet.opik.infrastructure.EncryptionUtils; +import com.comet.opik.utils.JsonUtils; +import org.jdbi.v3.core.argument.AbstractArgumentFactory; +import org.jdbi.v3.core.argument.Argument; +import org.jdbi.v3.core.config.ConfigRegistry; + +import java.sql.Types; + +/** + * Binds a {@link ProviderAuthConfig} as its AES-GCM-encrypted JSON representation, so the recipe + * (which can contain several secrets) is never stored or logged in plaintext. The read side lives + * in {@code ProviderApiKeyRowMapper}. + */ +public class ProviderAuthConfigArgumentFactory extends AbstractArgumentFactory { + + public ProviderAuthConfigArgumentFactory() { + super(Types.VARCHAR); + } + + @Override + protected Argument build(ProviderAuthConfig value, ConfigRegistry config) { + return (position, statement, ctx) -> { + if (value == null) { + statement.setNull(position, Types.VARCHAR); + } else { + statement.setString(position, EncryptionUtils.encryptGcm(JsonUtils.writeValueAsString(value))); + } + }; + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/LlmProviderClientApiConfig.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/LlmProviderClientApiConfig.java index d310fba8dfa..8eb8b438dd8 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/LlmProviderClientApiConfig.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/LlmProviderClientApiConfig.java @@ -1,13 +1,15 @@ package com.comet.opik.infrastructure.llm; +import com.comet.opik.api.ProviderAuthConfig; import lombok.Builder; import lombok.ToString; import java.util.Map; +import java.util.UUID; @Builder public record LlmProviderClientApiConfig(@ToString.Exclude String apiKey, Map headers, String baseUrl, - Map configuration) { + Map configuration, UUID providerId, String workspaceId, ProviderAuthConfig authConfig) { @Override public String toString() { @@ -16,6 +18,9 @@ public String toString() { ", headers=" + headers + ", baseUrl='" + baseUrl + '\'' + ", configuration=" + configuration + + ", providerId=" + providerId + + ", workspaceId='" + workspaceId + '\'' + + ", authConfig=" + authConfig + '}'; } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/LlmProviderFactoryImpl.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/LlmProviderFactoryImpl.java index 1bac97a41e4..490fe1ed056 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/LlmProviderFactoryImpl.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/LlmProviderFactoryImpl.java @@ -52,7 +52,7 @@ public void register(LlmProvider llmProvider, LlmServiceProvider service) { public LlmProviderService getService(@NonNull String workspaceId, @NonNull String model) { var llmProvider = getLlmProvider(model); var providerConfig = getProviderApiKey(workspaceId, llmProvider, model); - var config = buildConfig(providerConfig); + var config = buildConfig(providerConfig, workspaceId); return Optional.ofNullable(services.get(llmProvider)) .map(provider -> provider.getService(config)) @@ -60,7 +60,7 @@ public LlmProviderService getService(@NonNull String workspaceId, @NonNull Strin "LLM provider not supported: %s".formatted(llmProvider))); } - private LlmProviderClientApiConfig buildConfig(ProviderApiKey providerConfig) { + private LlmProviderClientApiConfig buildConfig(ProviderApiKey providerConfig, String workspaceId) { var configuration = Optional.ofNullable(providerConfig.configuration()).orElse(Map.of()); // For providers that support naming, add provider_name to configuration if present @@ -75,6 +75,9 @@ private LlmProviderClientApiConfig buildConfig(ProviderApiKey providerConfig) { .headers(Optional.ofNullable(providerConfig.headers()).orElse(Map.of())) .baseUrl(providerConfig.baseUrl()) .configuration(configuration) + .providerId(providerConfig.id()) + .workspaceId(workspaceId) + .authConfig(providerConfig.authConfig()) .build(); } @@ -82,7 +85,7 @@ public ChatModel getLanguageModel(@NonNull String workspaceId, @NonNull LlmAsJudgeModelParameters modelParameters) { var llmProvider = getLlmProvider(modelParameters.name()); var providerConfig = getProviderApiKey(workspaceId, llmProvider, modelParameters.name()); - var config = buildConfig(providerConfig); + var config = buildConfig(providerConfig, workspaceId); return Optional.ofNullable(services.get(llmProvider)) .map(provider -> provider.getLanguageModel(config, modelParameters)) diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/AuthTokenException.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/AuthTokenException.java new file mode 100644 index 00000000000..d28cdf393a2 --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/AuthTokenException.java @@ -0,0 +1,17 @@ +package com.comet.opik.infrastructure.llm.customllm; + +/** + * A token fetch for a custom provider's {@code auth_config} recipe failed. The message is + * user-facing by contract: it carries the upstream status/body with credential values and tokens + * already redacted, so callers can surface it verbatim (eval logs, playground, test-connection). + */ +public class AuthTokenException extends RuntimeException { + + public AuthTokenException(String message) { + super(message); + } + + public AuthTokenException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/AuthTokenProvider.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/AuthTokenProvider.java new file mode 100644 index 00000000000..d29d8f7f1ba --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/AuthTokenProvider.java @@ -0,0 +1,450 @@ +package com.comet.opik.infrastructure.llm.customllm; + +import com.comet.opik.api.ProviderAuthConfig; +import com.comet.opik.infrastructure.EncryptionUtils; +import com.comet.opik.infrastructure.LlmProviderTokenAuthConfig; +import com.comet.opik.infrastructure.lock.LockService; +import com.comet.opik.infrastructure.net.DestinationGuard; +import com.comet.opik.infrastructure.net.DestinationGuardException; +import com.comet.opik.infrastructure.redis.StringRedisClient; +import com.comet.opik.utils.JsonUtils; +import com.fasterxml.jackson.databind.JsonNode; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.metrics.LongCounter; +import io.opentelemetry.api.metrics.LongHistogram; +import io.opentelemetry.api.metrics.Meter; +import jakarta.inject.Inject; +import jakarta.inject.Singleton; +import lombok.NonNull; +import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; +import reactor.core.publisher.Mono; +import ru.vyarus.dropwizard.guice.module.yaml.bind.Config; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.net.URI; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.time.Duration; +import java.time.Instant; +import java.util.Base64; +import java.util.HexFormat; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.stream.Collectors; + +import static org.apache.commons.lang3.StringUtils.defaultIfBlank; +import static org.apache.commons.lang3.StringUtils.isBlank; +import static org.apache.commons.lang3.StringUtils.isNotBlank; + +/** + * Executes a custom provider's {@link ProviderAuthConfig} recipe and manages the resulting + * short-lived bearer's lifecycle. + * + *

Tokens are cached in Redis — one bucket per provider, AES-GCM-encrypted, shared by every + * backend pod — so a deployment makes roughly one fetch per refresh window regardless of pod count. + * The cache key includes a hash of the recipe: any config edit is an instant deployment-wide cache + * miss, no invalidation signal needed. A token is refreshed lazily, on the request path, once its + * remaining lifetime drops inside the proportional refresh window. Cold-cache bursts collapse into + * a single fetch via the distributed lock; on any Redis failure the provider degrades to a direct + * fetch rather than failing the LLM call. + * + *

Every error message thrown from here is user-facing by contract: upstream status and body are + * surfaced (never a generic "auth failed"), with credential values redacted. + */ +@Singleton +@Slf4j +public class AuthTokenProvider { + + record FetchedToken(String token, long ttlSeconds) { + } + + record CachedToken(String token, long expiresAtEpochMs, long ttlSeconds) { + } + + private static final String CACHE_KEY_FORMAT = "llm_auth_token:%s:%s"; + private static final String DEFAULT_TOKEN_FIELD = "access_token"; + private static final String DEFAULT_EXPIRES_FIELD = "expires_in"; + private static final String CLIENT_ID_KEY = "client_id"; + private static final String CLIENT_SECRET_KEY = "client_secret"; + private static final String REDACTED = "***"; + private static final int ERROR_BODY_SNIPPET_CHARS = 500; + // Upper bound on any token lifetime (1 year): beyond this a reply is malformed or malicious + static final long MAX_TTL_SECONDS = 31_536_000L; + + private static final AttributeKey OUTCOME = AttributeKey.stringKey("outcome"); + private static final AttributeKey WORKSPACE_ID = AttributeKey.stringKey("workspace_id"); + private static final AttributeKey ORIGIN = AttributeKey.stringKey("origin"); + private static final String ORIGIN_REQUEST = "request"; + private static final String ORIGIN_TEST = "test"; + + private final @NonNull StringRedisClient redisClient; + private final @NonNull LockService lockService; + private final @NonNull LlmProviderTokenAuthConfig config; + private final DestinationGuard destinationGuard; + private final HttpClient httpClient; + private final LongCounter tokenRequests; + private final LongHistogram fetchDurationMs; + + @Inject + public AuthTokenProvider(@NonNull StringRedisClient redisClient, @NonNull LockService lockService, + @NonNull @Config("llmProviderTokenAuth") LlmProviderTokenAuthConfig config) { + this.redisClient = redisClient; + this.lockService = lockService; + this.config = config; + this.destinationGuard = new DestinationGuard(config.getDestinationGuard()); + // Same posture as the LLM clients: HTTP/1.1 and no redirect following (a token endpoint + // redirecting elsewhere is a misconfiguration or an attack, never a flow to honor). + this.httpClient = HttpClient.newBuilder() + .version(HttpClient.Version.HTTP_1_1) + .followRedirects(HttpClient.Redirect.NEVER) + .connectTimeout(Duration.ofMillis(config.getFetchTimeout().toMilliseconds())) + .build(); + Meter meter = GlobalOpenTelemetry.get().getMeter("opik.llm_provider_token_auth"); + this.tokenRequests = meter.counterBuilder("llm_provider_token_requests") + .setDescription("Bearer requests against the token cache, by outcome") + .build(); + this.fetchDurationMs = meter.histogramBuilder("llm_provider_token_fetch_duration_ms") + .setDescription("Duration of HTTP fetches against customers' auth services") + .setUnit("ms") + .ofLongs() + .build(); + } + + /** + * Returns a bearer for the given recipe, fetching or refreshing through the shared cache as + * needed. Blocking — meant to be called from the LLM client's request path. + * + * @throws AuthTokenException with a redacted, user-facing message when the fetch fails + */ + public String bearer(@NonNull String workspaceId, @NonNull UUID providerId, + @NonNull ProviderAuthConfig authConfig) { + String cacheKey = cacheKey(providerId, authConfig); + CachedToken cached = readCache(cacheKey); + if (isFresh(cached)) { + recordRequestMetric(workspaceId, "cache_hit"); + return cached.token(); + } + + try { + String token = lockService + .executeWithLockCustomExpire( + new LockService.Lock(cacheKey + ":fetch-lock"), + Mono.fromCallable(() -> refreshUnderLock(cacheKey, authConfig)), + Duration.ofMillis(config.getLockTimeout().toMilliseconds())) + .block(Duration.ofMillis( + config.getLockTimeout().toMilliseconds() + config.getFetchTimeout().toMilliseconds())); + if (token != null) { + recordRequestMetric(workspaceId, "fetched"); + return token; + } + log.warn("Timed out waiting for the token fetch lock for provider '{}'; fetching directly", providerId); + } catch (AuthTokenException exception) { + recordRequestMetric(workspaceId, "failed"); + throw exception; + } catch (RuntimeException exception) { + log.warn("Token cache unavailable for provider '{}'; falling back to a direct fetch", providerId, + exception); + } + + String token = fetchToken(authConfig).token(); + recordRequestMetric(workspaceId, "degraded_direct"); + return token; + } + + /** + * Drops the cached token — for the 401-retry path, so a revocation discovered by one pod is + * seen by all of them at once. Best-effort: a Redis failure here only delays the cleanup until + * the bucket's own expiry. + */ + public void invalidate(@NonNull UUID providerId, @NonNull ProviderAuthConfig authConfig) { + String cacheKey = cacheKey(providerId, authConfig); + try { + redisClient.getBucket(cacheKey).delete(); + } catch (RuntimeException exception) { + log.warn("Failed to invalidate cached token for provider '{}'", providerId, exception); + } + } + + private String refreshUnderLock(String cacheKey, ProviderAuthConfig authConfig) { + // another pod may have refreshed while this one waited on the lock + CachedToken cached = readCache(cacheKey); + if (isFresh(cached)) { + return cached.token(); + } + FetchedToken fetched = fetchToken(authConfig); + // a resolved lifetime of 0 (the fallback for a reply that states none) means: don't cache + if (fetched.ttlSeconds() > 0) { + writeCache(cacheKey, fetched); + } + return fetched.token(); + } + + private boolean isFresh(CachedToken cached) { + if (cached == null) { + return false; + } + long refreshWindowMs = (long) (cached.ttlSeconds() * 1_000 * config.getRefreshFraction()); + return Instant.now().toEpochMilli() < cached.expiresAtEpochMs() - refreshWindowMs; + } + + private CachedToken readCache(String cacheKey) { + try { + String encrypted = redisClient.getBucket(cacheKey).get(); + if (encrypted == null) { + return null; + } + return JsonUtils.readValue(EncryptionUtils.decryptGcm(encrypted), CachedToken.class); + } catch (RuntimeException exception) { + log.warn("Failed to read the cached token for key '{}': {}", cacheKey, exception.getMessage()); + return null; + } + } + + private void writeCache(String cacheKey, FetchedToken fetched) { + try { + var entry = new CachedToken(fetched.token(), + Instant.now().toEpochMilli() + fetched.ttlSeconds() * 1_000, fetched.ttlSeconds()); + redisClient.getBucket(cacheKey) + .set(EncryptionUtils.encryptGcm(JsonUtils.writeValueAsString(entry)), + Duration.ofSeconds(fetched.ttlSeconds())); + } catch (RuntimeException exception) { + log.warn("Failed to cache the token for key '{}'; the token is still returned", cacheKey, exception); + } + } + + private String cacheKey(UUID providerId, ProviderAuthConfig authConfig) { + return CACHE_KEY_FORMAT.formatted(providerId, sha256Hex(JsonUtils.writeValueAsString(authConfig))); + } + + private static String sha256Hex(String value) { + try { + return HexFormat.of() + .formatHex(MessageDigest.getInstance("SHA-256").digest(value.getBytes(StandardCharsets.UTF_8))); + } catch (NoSuchAlgorithmException exception) { + throw new IllegalStateException(exception); + } + } + + /** + * Runs the recipe once for the test-connection endpoint. Returns the resolved token lifetime — + * never the token itself. + * + * @throws AuthTokenException with a redacted, user-facing message when the fetch fails + */ + public long testFetch(@NonNull ProviderAuthConfig authConfig) { + return fetchToken(authConfig, ORIGIN_TEST).ttlSeconds(); + } + + // --- recipe execution --- + + FetchedToken fetchToken(@NonNull ProviderAuthConfig authConfig) { + return fetchToken(authConfig, ORIGIN_REQUEST); + } + + private FetchedToken fetchToken(ProviderAuthConfig authConfig, String origin) { + long startNanos = System.nanoTime(); + try { + destinationGuard.validate(authConfig.tokenUrl()); + } catch (DestinationGuardException exception) { + recordFetchMetric(startNanos, "destination_refused", origin); + throw new AuthTokenException(exception.getMessage() + + " (self-hosted deployments with internal auth services can set LLM_PROVIDER_TOKEN_AUTH_DESTINATION_GUARD=relaxed)", + exception); + } + HttpRequest request = buildRequest(authConfig); + HttpResponse response; + try { + response = httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + } catch (IOException exception) { + recordFetchMetric(startNanos, "unreachable", origin); + throw new AuthTokenException("token fetch failed: could not reach '%s': %s" + .formatted(authConfig.tokenUrl(), redact(authConfig, exception.getMessage())), exception); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + recordFetchMetric(startNanos, "interrupted", origin); + throw new AuthTokenException("token fetch was interrupted", exception); + } + + String body = Optional.ofNullable(response.body()).orElse(""); + if (response.statusCode() < 200 || response.statusCode() >= 300) { + recordFetchMetric(startNanos, "upstream_error", origin); + throw new AuthTokenException("token fetch failed with status '%d' from '%s': %s" + .formatted(response.statusCode(), authConfig.tokenUrl(), redact(authConfig, body))); + } + if (body.length() > config.getMaxResponseChars()) { + recordFetchMetric(startNanos, "oversized_reply", origin); + throw new AuthTokenException("token reply from '%s' exceeds the maximum accepted size" + .formatted(authConfig.tokenUrl())); + } + + JsonNode root; + try { + root = JsonUtils.getJsonNodeFromString(body); + } catch (UncheckedIOException exception) { + recordFetchMetric(startNanos, "non_json_reply", origin); + throw new AuthTokenException("token endpoint at '%s' returned a non-JSON reply (status '%d')" + .formatted(authConfig.tokenUrl(), response.statusCode())); + } + + String tokenField = defaultIfBlank(authConfig.tokenField(), DEFAULT_TOKEN_FIELD); + JsonNode tokenNode = atDotPath(root, tokenField); + if (!tokenNode.isTextual() || isBlank(tokenNode.asText())) { + recordFetchMetric(startNanos, "token_field_missing", origin); + // field names are safe to surface; values never are + throw new AuthTokenException("field '%s' not found in the token reply; top-level fields: %s" + .formatted(tokenField, root.properties().stream().map(Map.Entry::getKey).toList())); + } + + long ttlSeconds = resolveTtlSeconds(root, authConfig, startNanos, origin); + recordFetchMetric(startNanos, "success", origin); + return new FetchedToken(tokenNode.asText(), ttlSeconds); + } + + private HttpRequest buildRequest(ProviderAuthConfig authConfig) { + List credentials = Optional.ofNullable(authConfig.credentials()) + .orElse(List.of()); + var sendAs = Optional.ofNullable(authConfig.sendAs()).orElse(ProviderAuthConfig.SendAs.FORM); + + var builder = HttpRequest.newBuilder(URI.create(authConfig.tokenUrl())) + .timeout(Duration.ofMillis(config.getFetchTimeout().toMilliseconds())); + + switch (sendAs) { + case FORM -> builder + .header("Content-Type", "application/x-www-form-urlencoded") + .POST(HttpRequest.BodyPublishers.ofString(formEncode(credentials))); + case JSON -> builder + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(jsonEncode(credentials))); + case BASIC -> { + // OAuth2-standard mixed mode: client id/secret in the Basic header, every + // remaining field in the form body + String clientId = credentialValue(credentials, CLIENT_ID_KEY); + String clientSecret = credentialValue(credentials, CLIENT_SECRET_KEY); + if (clientId == null || clientSecret == null) { + throw new AuthTokenException( + "basic auth mode requires '%s' and '%s' credentials".formatted(CLIENT_ID_KEY, + CLIENT_SECRET_KEY)); + } + var bodyCredentials = credentials.stream() + .filter(credential -> !CLIENT_ID_KEY.equals(credential.key()) + && !CLIENT_SECRET_KEY.equals(credential.key())) + .toList(); + builder.header("Authorization", "Basic " + Base64.getEncoder().encodeToString( + (clientId + ":" + clientSecret).getBytes(StandardCharsets.UTF_8))) + .header("Content-Type", "application/x-www-form-urlencoded") + .POST(HttpRequest.BodyPublishers.ofString(formEncode(bodyCredentials))); + } + } + return builder.build(); + } + + /** + * Resolves the token lifetime, in order: the reply's lifetime field (number or numeric string; + * a non-positive value is an error), else the recipe's fallback (where 0 means the token is + * served uncached), else an error naming the missing field. + */ + private long resolveTtlSeconds(JsonNode root, ProviderAuthConfig authConfig, long startNanos, String origin) { + String expiresField = defaultIfBlank(authConfig.expiresField(), DEFAULT_EXPIRES_FIELD); + JsonNode expiresNode = atDotPath(root, expiresField); + + Long ttlSeconds = null; + if (expiresNode.isNumber()) { + ttlSeconds = expiresNode.longValue(); + } else if (expiresNode.isTextual() && isNotBlank(expiresNode.asText())) { + try { + ttlSeconds = Long.parseLong(expiresNode.asText().trim()); + } catch (NumberFormatException ignored) { + // fall through to the fallback + } + } + if (ttlSeconds != null) { + if (ttlSeconds <= 0 || ttlSeconds > MAX_TTL_SECONDS) { + recordFetchMetric(startNanos, "lifetime_invalid", origin); + throw new AuthTokenException( + "token reply states a lifetime of '%d' seconds, outside the accepted range of 1 to %d" + .formatted(ttlSeconds, MAX_TTL_SECONDS)); + } + return ttlSeconds; + } + // the configured fallback may legitimately be 0: the fetch-per-call convention + if (authConfig.fallbackTtlSeconds() == null) { + recordFetchMetric(startNanos, "lifetime_missing", origin); + throw new AuthTokenException( + "token reply has no lifetime field '%s' and no fallback lifetime is configured" + .formatted(expiresField)); + } + return authConfig.fallbackTtlSeconds(); + } + + private static String credentialValue(List credentials, String key) { + return credentials.stream() + .filter(credential -> key.equals(credential.key())) + .map(ProviderAuthConfig.Credential::value) + .filter(StringUtils::isNotBlank) + .findFirst() + .orElse(null); + } + + private static String formEncode(List credentials) { + return credentials.stream() + .map(credential -> URLEncoder.encode(credential.key(), StandardCharsets.UTF_8) + "=" + + URLEncoder.encode(Optional.ofNullable(credential.value()).orElse(""), + StandardCharsets.UTF_8)) + .collect(Collectors.joining("&")); + } + + private static String jsonEncode(List credentials) { + var node = JsonUtils.createObjectNode(); + credentials.forEach(credential -> node.put(credential.key(), credential.value())); + return node.toString(); + } + + private static JsonNode atDotPath(JsonNode root, String dotPath) { + JsonNode current = root; + for (String part : dotPath.split("\\.")) { + current = current.path(part); + } + return current; + } + + /** + * Scrubs every credential value out of text destined for error messages and truncates it. + * Applied to upstream bodies and transport errors alike — the one choke point that lets the + * rest of the class surface upstream errors verbatim. + */ + private static String redact(ProviderAuthConfig authConfig, String text) { + if (isBlank(text)) { + return ""; + } + String result = text.length() > ERROR_BODY_SNIPPET_CHARS + ? text.substring(0, ERROR_BODY_SNIPPET_CHARS) + "…" + : text; + for (var credential : Optional.ofNullable(authConfig.credentials()) + .orElse(List.of())) { + if (isNotBlank(credential.value())) { + result = result.replace(credential.value(), REDACTED); + } + } + return result; + } + + private void recordRequestMetric(String workspaceId, String outcome) { + tokenRequests.add(1, Attributes.of(OUTCOME, outcome, WORKSPACE_ID, workspaceId)); + } + + private void recordFetchMetric(long startNanos, String outcome, String origin) { + fetchDurationMs.record((System.nanoTime() - startNanos) / 1_000_000, + Attributes.of(OUTCOME, outcome, ORIGIN, origin)); + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/CustomLlmClientGenerator.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/CustomLlmClientGenerator.java index 4c55dbc7cf7..2cd14a9b441 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/CustomLlmClientGenerator.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/CustomLlmClientGenerator.java @@ -19,12 +19,14 @@ import java.net.http.HttpClient; import java.util.Map; import java.util.Optional; +import java.util.function.Supplier; @RequiredArgsConstructor @Slf4j public class CustomLlmClientGenerator implements LlmProviderClientGenerator { private final @NonNull LlmProviderClientConfig llmProviderClientConfig; + private final @NonNull AuthTokenProvider authTokenProvider; public OpenAiClient newCustomLlmClient(@NonNull LlmProviderClientApiConfig config) { var baseUrl = Optional.ofNullable(config.baseUrl()) @@ -126,7 +128,27 @@ private HttpClientBuilder newHttpClientBuilder(LlmProviderClientApiConfig config if (!requiresInterceptingBuilder(config)) { return jdkHttpClientBuilder; } - return new InterceptingHttpClientBuilder(jdkHttpClientBuilder, config.configuration(), config.apiKey()); + return new InterceptingHttpClientBuilder(jdkHttpClientBuilder, config.configuration(), config.apiKey(), + bearerSupplier(config), tokenInvalidator(config)); + } + + /** + * Per-request bearer source for token-auth providers. The supplier form matters: clients are + * rebuilt per call but requests are what carry auth, so the interceptor asks the shared cache + * on every request and refresh/rotation need no client rebuild. + */ + private Supplier bearerSupplier(LlmProviderClientApiConfig config) { + if (config.authConfig() == null) { + return null; + } + return () -> authTokenProvider.bearer(config.workspaceId(), config.providerId(), config.authConfig()); + } + + private Runnable tokenInvalidator(LlmProviderClientApiConfig config) { + if (config.authConfig() == null) { + return null; + } + return () -> authTokenProvider.invalidate(config.providerId(), config.authConfig()); } /** @@ -144,6 +166,6 @@ private static boolean requiresInterceptingBuilder(LlmProviderClientApiConfig co configuration.get(InterceptingHttpClient.SUPPRESS_DEFAULT_AUTH_CONFIG_KEY)))); boolean hasModelPlaceholder = config.baseUrl() != null && config.baseUrl().contains(InterceptingHttpClient.MODEL_PLACEHOLDER); - return hasNewConfigKeys || hasModelPlaceholder; + return hasNewConfigKeys || hasModelPlaceholder || config.authConfig() != null; } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/CustomLlmModule.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/CustomLlmModule.java index 5fae3c58148..7f2799d091c 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/CustomLlmModule.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/CustomLlmModule.java @@ -15,8 +15,9 @@ public class CustomLlmModule extends AbstractModule { @Singleton @Named("customLlmGenerator") public CustomLlmClientGenerator clientGenerator( - @NonNull @Config("llmProviderClient") LlmProviderClientConfig config) { - return new CustomLlmClientGenerator(config); + @NonNull @Config("llmProviderClient") LlmProviderClientConfig config, + @NonNull AuthTokenProvider authTokenProvider) { + return new CustomLlmClientGenerator(config, authTokenProvider); } @Provides diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClient.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClient.java index c412629d182..78762cfa493 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClient.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClient.java @@ -7,9 +7,12 @@ import dev.langchain4j.http.client.HttpClient; import dev.langchain4j.http.client.HttpRequest; import dev.langchain4j.http.client.SuccessfulHttpResponse; +import dev.langchain4j.http.client.sse.ServerSentEvent; +import dev.langchain4j.http.client.sse.ServerSentEventContext; import dev.langchain4j.http.client.sse.ServerSentEventListener; import dev.langchain4j.http.client.sse.ServerSentEventParser; import lombok.NonNull; +import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; import org.apache.http.client.utils.URIBuilder; @@ -19,6 +22,7 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.function.Supplier; /** * HTTP client decorator for the Custom LLM provider that mutates outgoing @@ -39,6 +43,13 @@ *

  • Removes the default {@code Authorization: Bearer } header when * {@code configuration["suppress_default_auth"]} is {@code "true"}. Used * by gateways whose policy rejects an {@code Authorization} header.
    • + *
  • Injects a dynamically fetched bearer for token-auth providers ({@code auth_config}): + * the token is asked from the shared cache on every request and placed in + * {@code Authorization: Bearer} (or raw under the configured {@code auth_header_name}). + * Precedence: the fetched token wins — a pre-configured static {@code Authorization} + * header is dropped with a warn. A 401/403 from the gateway invalidates the cached token + * and retries the request once; for streamed responses only before the first delivered + * event.
    • * * *

    When none of the relevant config keys are set and no {@code {model}} @@ -62,6 +73,12 @@ class InterceptingHttpClient implements HttpClient { private final @NonNull HttpClient delegate; private final Map configuration; private final String apiKey; + private final Supplier tokenSupplier; + private final Runnable tokenInvalidator; + + InterceptingHttpClient(@NonNull HttpClient delegate, Map configuration, String apiKey) { + this(delegate, configuration, apiKey, null, null); + } /** * Normalizes a null {@code configuration} to an empty map so helpers can @@ -69,26 +86,107 @@ class InterceptingHttpClient implements HttpClient { * {@code {model}}-only providers (which may carry no new config keys yet * still reach {@code mutate()}) safe. */ - InterceptingHttpClient(@NonNull HttpClient delegate, Map configuration, String apiKey) { + InterceptingHttpClient(@NonNull HttpClient delegate, Map configuration, String apiKey, + Supplier tokenSupplier, Runnable tokenInvalidator) { this.delegate = delegate; this.configuration = configuration != null ? configuration : Map.of(); this.apiKey = apiKey; + this.tokenSupplier = tokenSupplier; + this.tokenInvalidator = tokenInvalidator; } @Override public SuccessfulHttpResponse execute(HttpRequest request) throws HttpException, RuntimeException { - return delegate.execute(mutate(request)); + try { + return delegate.execute(mutate(request)); + } catch (HttpException exception) { + if (!shouldRetryWithFreshToken(exception)) { + throw exception; + } + invalidateToken(); + // re-mutating asks the supplier again, which now yields a fresh token + return delegate.execute(mutate(request)); + } } @Override public void execute(HttpRequest request, ServerSentEventParser parser, ServerSentEventListener listener) { - delegate.execute(mutate(request), parser, listener); + if (tokenSupplier == null) { + delegate.execute(mutate(request), parser, listener); + return; + } + delegate.execute(mutate(request), parser, new RetryOnAuthFailureListener(request, parser, listener)); + } + + private boolean shouldRetryWithFreshToken(HttpException exception) { + return tokenSupplier != null + && (exception.statusCode() == 401 || exception.statusCode() == 403); + } + + private void invalidateToken() { + if (tokenInvalidator != null) { + tokenInvalidator.run(); + } + } + + /** + * Retries a streamed request once on 401/403, but only while nothing has been delivered to + * the downstream listener — once content flowed, a retry would splice two half-responses, so + * the failure propagates honestly instead. + */ + @RequiredArgsConstructor + private class RetryOnAuthFailureListener implements ServerSentEventListener { + + private final HttpRequest originalRequest; + private final ServerSentEventParser parser; + private final ServerSentEventListener downstream; + private boolean delivered; + + @Override + public void onOpen(SuccessfulHttpResponse response) { + delivered = true; + downstream.onOpen(response); + } + + @Override + public void onEvent(ServerSentEvent event, ServerSentEventContext context) { + delivered = true; + downstream.onEvent(event, context); + } + + @Override + public void onEvent(ServerSentEvent event) { + delivered = true; + downstream.onEvent(event); + } + + @Override + public void onError(Throwable throwable) { + if (!delivered && throwable instanceof HttpException httpException + && shouldRetryWithFreshToken(httpException)) { + invalidateToken(); + try { + // retry straight to the downstream listener: a second failure must not retry again + delegate.execute(mutate(originalRequest), parser, downstream); + return; + } catch (RuntimeException retryFailure) { + downstream.onError(retryFailure); + return; + } + } + downstream.onError(throwable); + } + + @Override + public void onClose() { + downstream.onClose(); + } } private HttpRequest mutate(HttpRequest request) { boolean hasPlaceholder = request.url() != null && request.url().contains(MODEL_PLACEHOLDER); - if (configuration.isEmpty() && !hasPlaceholder) { + if (configuration.isEmpty() && !hasPlaceholder && tokenSupplier == null) { return request; } @@ -155,6 +253,9 @@ private String applyModelPlaceholder(String url, String body) { } private Map> applyAuthHeaders(Map> headers) { + if (tokenSupplier != null) { + return applyBearerToken(headers); + } String customHeaderName = configuration.get(AUTH_HEADER_NAME_CONFIG_KEY); boolean addCustomHeader = StringUtils.isNotBlank(customHeaderName) && StringUtils.isNotBlank(apiKey); @@ -189,6 +290,31 @@ private Map> applyAuthHeaders(Map> hea return mutated; } + /** + * Token-auth mode: the fetched bearer is the sole credential on the request. It lands in + * {@code Authorization: Bearer } by default, or raw under the configured + * {@code auth_header_name}. A pre-existing {@code Authorization} (e.g. a static header from + * the provider's headers map) loses to the fetched token and is dropped loudly. + */ + private Map> applyBearerToken(Map> headers) { + String token = tokenSupplier.get(); + String customHeaderName = StringUtils.trimToNull(configuration.get(AUTH_HEADER_NAME_CONFIG_KEY)); + String targetHeader = customHeaderName != null ? customHeaderName : AUTHORIZATION_HEADER; + String targetValue = customHeaderName != null ? token : "Bearer " + token; + + var mutated = new LinkedHashMap>(headers.size() + 1); + for (Map.Entry> entry : headers.entrySet()) { + if (AUTHORIZATION_HEADER.equalsIgnoreCase(entry.getKey()) + || targetHeader.equalsIgnoreCase(entry.getKey())) { + log.warn("Dropping pre-configured '{}' header: the fetched token takes precedence", entry.getKey()); + continue; + } + mutated.put(entry.getKey(), entry.getValue()); + } + mutated.put(targetHeader, List.of(targetValue)); + return mutated; + } + private String applyQueryParams(String url) { String raw = configuration.get(URL_QUERY_PARAMS_CONFIG_KEY); if (StringUtils.isBlank(raw) || StringUtils.isBlank(url)) { diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClientBuilder.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClientBuilder.java index 22595d8059a..429848bdab6 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClientBuilder.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClientBuilder.java @@ -6,11 +6,13 @@ import lombok.NonNull; import java.util.Map; +import java.util.function.Supplier; /** * Decorates an {@link HttpClientBuilder} so the produced {@link HttpClient} is * wrapped with {@link InterceptingHttpClient} to apply Custom LLM-specific - * request mutations (query params, auth headers, {@code {model}} substitution). + * request mutations (query params, auth headers, {@code {model}} substitution, + * dynamic bearer injection). * *

    Timeout forwarding is inherited from {@link DelegatingHttpClientBuilder}; only the * {@link #wrap(HttpClient)} step is specific to this builder. @@ -19,16 +21,20 @@ class InterceptingHttpClientBuilder extends DelegatingHttpClientBuilder { private final Map configuration; private final String apiKey; + private final Supplier tokenSupplier; + private final Runnable tokenInvalidator; InterceptingHttpClientBuilder(@NonNull HttpClientBuilder delegate, Map configuration, - String apiKey) { + String apiKey, Supplier tokenSupplier, Runnable tokenInvalidator) { super(delegate); this.configuration = configuration; this.apiKey = apiKey; + this.tokenSupplier = tokenSupplier; + this.tokenInvalidator = tokenInvalidator; } @Override protected HttpClient wrap(HttpClient delegate) { - return new InterceptingHttpClient(delegate, configuration, apiKey); + return new InterceptingHttpClient(delegate, configuration, apiKey, tokenSupplier, tokenInvalidator); } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/net/DestinationGuard.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/net/DestinationGuard.java new file mode 100644 index 00000000000..aa9870ee66b --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/net/DestinationGuard.java @@ -0,0 +1,113 @@ +package com.comet.opik.infrastructure.net; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonValue; +import lombok.NonNull; +import lombok.RequiredArgsConstructor; + +import java.net.Inet6Address; +import java.net.InetAddress; +import java.net.URI; +import java.net.URISyntaxException; +import java.net.UnknownHostException; +import java.util.Arrays; + +import static org.apache.commons.lang3.StringUtils.isBlank; + +/** + * Pre-flight check for outbound calls to user-supplied URLs (SSRF guard). In {@code STRICT} mode + * (cloud) it requires HTTPS and resolves the hostname before anyone connects, refusing addresses + * only our own network can reach: loopback, link-local (including the cloud metadata endpoint at + * 169.254.169.254), RFC 1918 private ranges, IPv6 unique-local, multicast, and unresolvable hosts. + * In {@code RELAXED} mode (self-hosted default) it is a no-op — internal gateways legitimately + * live on private ranges there. + * + *

    Resolve-then-decide is the accepted level of protection here: the later connection resolves + * again, so a DNS-rebinding attacker with a sub-TTL flip could theoretically pass the check. The + * surfaces this guards are admin-configured (not anonymous input), which keeps that residual risk + * acceptable; connection-time pinning would require a custom socket layer. + */ +@RequiredArgsConstructor +public class DestinationGuard { + + public enum Mode { + RELAXED("relaxed"), + STRICT("strict"), + ; + + @JsonValue + private final String value; + + Mode(String value) { + this.value = value; + } + + @JsonCreator + public static Mode fromString(String value) { + return Arrays.stream(values()) + .filter(mode -> mode.value.equalsIgnoreCase(value)) + .findFirst() + .orElseThrow(() -> new IllegalArgumentException( + "Unknown destination guard mode '%s'".formatted(value))); + } + } + + private final @NonNull Mode mode; + + /** + * @throws DestinationGuardException with a user-facing message when the destination is refused + */ + public void validate(@NonNull String url) { + if (mode == Mode.RELAXED) { + return; + } + + URI uri; + try { + uri = new URI(url); + } catch (URISyntaxException exception) { + throw new DestinationGuardException("destination '%s' is not a valid URL".formatted(url)); + } + if (!"https".equalsIgnoreCase(uri.getScheme())) { + throw new DestinationGuardException( + "destination '%s' was refused: only https URLs are allowed".formatted(url)); + } + String host = uri.getHost(); + if (isBlank(host)) { + throw new DestinationGuardException("destination '%s' has no valid host".formatted(url)); + } + + InetAddress[] addresses; + try { + addresses = InetAddress.getAllByName(host); + } catch (UnknownHostException exception) { + throw new DestinationGuardException( + "destination host '%s' could not be resolved".formatted(host)); + } + for (InetAddress address : addresses) { + if (isNonPublic(address)) { + // deliberately not echoing the resolved address: the hostname is the user's own + // input, the address it maps to inside our network is not theirs to learn + throw new DestinationGuardException( + "destination host '%s' was refused: it resolves to a private or internal address" + .formatted(host)); + } + } + } + + private static boolean isNonPublic(InetAddress address) { + return address.isAnyLocalAddress() + || address.isLoopbackAddress() + || address.isLinkLocalAddress() + || address.isSiteLocalAddress() + || address.isMulticastAddress() + || isUniqueLocalIpv6(address); + } + + /** + * fc00::/7 — Java's {@code isSiteLocalAddress} only covers the deprecated fec0::/10 for IPv6. + */ + private static boolean isUniqueLocalIpv6(InetAddress address) { + return address instanceof Inet6Address && (address.getAddress()[0] & 0xFE) == 0xFC; + } +} diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/net/DestinationGuardException.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/net/DestinationGuardException.java new file mode 100644 index 00000000000..6080f7c6c1b --- /dev/null +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/net/DestinationGuardException.java @@ -0,0 +1,12 @@ +package com.comet.opik.infrastructure.net; + +/** + * A user-supplied outbound destination was refused by {@link DestinationGuard}. The message is + * user-facing: it names the URL the user configured, never the address it resolved to. + */ +public class DestinationGuardException extends RuntimeException { + + public DestinationGuardException(String message) { + super(message); + } +} diff --git a/apps/opik-backend/src/main/resources/liquibase/db-app-state/migrations/000095_add_auth_config_to_llm_provider_api_key.sql b/apps/opik-backend/src/main/resources/liquibase/db-app-state/migrations/000095_add_auth_config_to_llm_provider_api_key.sql new file mode 100644 index 00000000000..0483e6d4f73 --- /dev/null +++ b/apps/opik-backend/src/main/resources/liquibase/db-app-state/migrations/000095_add_auth_config_to_llm_provider_api_key.sql @@ -0,0 +1,7 @@ +--liquibase formatted sql +--changeset miguelg:000095_add_auth_config_to_llm_provider_api_key +--comment: Add encrypted auth_config column for dynamic token auth on custom providers + +ALTER TABLE llm_provider_api_key ADD COLUMN auth_config TEXT DEFAULT NULL; + +--rollback ALTER TABLE llm_provider_api_key DROP COLUMN auth_config; diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/ProviderAuthConfigTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/ProviderAuthConfigTest.java new file mode 100644 index 00000000000..12a40d37b17 --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/ProviderAuthConfigTest.java @@ -0,0 +1,85 @@ +package com.comet.opik.api; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.List; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +class ProviderAuthConfigTest { + + private static final ProviderAuthConfig AUTH_CONFIG = ProviderAuthConfig.builder() + .tokenUrl("https://auth.example.com/oauth/token") + .sendAs(ProviderAuthConfig.SendAs.BASIC) + .credentials(List.of( + credential("grant_type", "client_credentials", false), + credential("client_id", "opik-prod", false), + credential("client_secret", "s3cr3t", true))) + .tokenField("access_token") + .expiresField("expires_in") + .build(); + + private static ProviderAuthConfig.Credential credential(String key, String value, boolean secret) { + return ProviderAuthConfig.Credential.builder().key(key).value(value).secret(secret).build(); + } + + @Test + void maskReplacesOnlySecretValues() { + var masked = AUTH_CONFIG.mask(); + + assertThat(masked.credentials()).containsExactly( + credential("grant_type", "client_credentials", false), + credential("client_id", "opik-prod", false), + credential("client_secret", ProviderAuthConfig.SECRET_SENTINEL, true)); + assertThat(masked.tokenUrl()).isEqualTo(AUTH_CONFIG.tokenUrl()); + } + + @Test + void emptyObjectIsTheClearConvention() { + assertThat(ProviderAuthConfig.builder().build().isEmpty()).isTrue(); + assertThat(AUTH_CONFIG.isEmpty()).isFalse(); + assertThat(ProviderAuthConfig.builder().tokenUrl("https://auth.example.com").build().isEmpty()).isFalse(); + } + + // any single field set means "not the clear convention" — the update path must + // route these through validationErrors(), never silently clear the stored recipe + static Stream partialConfigs() { + return Stream.of( + arguments("sendAs", ProviderAuthConfig.builder().sendAs(ProviderAuthConfig.SendAs.BASIC).build()), + arguments("credentials", ProviderAuthConfig.builder() + .credentials(List.of(credential("client_id", "opik", false))) + .build()), + arguments("tokenField", ProviderAuthConfig.builder().tokenField("access_token").build()), + arguments("expiresField", ProviderAuthConfig.builder().expiresField("expires_in").build()), + arguments("fallbackTtlSeconds", ProviderAuthConfig.builder().fallbackTtlSeconds(60L).build())); + } + + @ParameterizedTest(name = "only {0} set") + @MethodSource("partialConfigs") + void partialConfigsAreNotEmptySoTheyValidateInsteadOfClearing(String field, ProviderAuthConfig partial) { + assertThat(partial.isEmpty()).isFalse(); + assertThat(partial.validationErrors()).isNotEmpty(); + } + + @Test + void validationErrorsRequireTokenUrlAndCredentials() { + assertThat(AUTH_CONFIG.validationErrors()).isEmpty(); + + assertThat(AUTH_CONFIG.toBuilder().tokenUrl("not a uri").build().validationErrors()) + .containsExactly("auth_config.token_url must be a valid absolute URI"); + assertThat(AUTH_CONFIG.toBuilder().credentials(List.of()).build().validationErrors()) + .containsExactly("auth_config.credentials must not be empty"); + assertThat(ProviderAuthConfig.builder().build().validationErrors()).hasSize(2); + } + + @Test + void toStringNeverExposesCredentialValues() { + assertThat(AUTH_CONFIG.toString()).doesNotContain("s3cr3t", "opik-prod", "client_credentials"); + assertThat(AUTH_CONFIG.credentials().getLast().toString()).doesNotContain("s3cr3t"); + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/LlmProviderApiKeyResourceClient.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/LlmProviderApiKeyResourceClient.java index af892a5e561..c475be3aca0 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/LlmProviderApiKeyResourceClient.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/LlmProviderApiKeyResourceClient.java @@ -5,6 +5,7 @@ import com.comet.opik.api.Page; import com.comet.opik.api.ProviderApiKey; import com.comet.opik.api.ProviderApiKeyUpdate; +import com.comet.opik.api.ProviderAuthCheck; import com.comet.opik.api.resources.utils.TestUtils; import jakarta.ws.rs.HttpMethod; import jakarta.ws.rs.client.Entity; @@ -71,6 +72,17 @@ public Response callUpdateProviderApiKey(UUID id, ProviderApiKeyUpdate providerA .method(HttpMethod.PATCH, Entity.json(providerApiKeyUpdate)); } + public Response callTestAuthConfig(ProviderAuthCheck providerAuthTest, String apiKey, String workspaceName) { + return client.target(RESOURCE_PATH.formatted(baseURI)) + .path("auth-config") + .path("test") + .request() + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(providerAuthTest)); + } + public Response callDeleteProviderApiKeys(Set ids, String apiKey, String workspaceName) { return client.target(RESOURCE_PATH.formatted(baseURI)) .path("delete") diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java index 3bf914a8749..ce8dd3f6e1f 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java @@ -4,6 +4,8 @@ import com.comet.opik.api.Page; import com.comet.opik.api.ProviderApiKey; import com.comet.opik.api.ProviderApiKeyUpdate; +import com.comet.opik.api.ProviderAuthCheck; +import com.comet.opik.api.ProviderAuthConfig; import com.comet.opik.api.resources.utils.AuthTestUtils; import com.comet.opik.api.resources.utils.ClickHouseContainerUtils; import com.comet.opik.api.resources.utils.ClientSupportUtils; @@ -621,6 +623,407 @@ void createAndGetProviderApiKeyListWithMinimalFields() { assertPage(actualProviderApiKeyPage, List.of(expectedProviderApiKey)); } + @Nested + @DisplayName("Auth config:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class AuthConfigTest { + + private static final String SECRET_VALUE = "super-s3cr3t"; + + private ProviderAuthConfig tokenAuthConfig() { + return ProviderAuthConfig.builder() + .tokenUrl("https://auth.example.com/oauth/token") + .sendAs(ProviderAuthConfig.SendAs.BASIC) + .credentials(List.of( + credential("grant_type", "client_credentials", false), + credential("client_id", "opik-prod", false), + credential("client_secret", SECRET_VALUE, true))) + .tokenField("access_token") + .expiresField("expires_in") + .fallbackTtlSeconds(3600L) + .build(); + } + + private ProviderAuthConfig.Credential credential(String key, String value, boolean secret) { + return ProviderAuthConfig.Credential.builder().key(key).value(value).secret(secret).build(); + } + + private ProviderApiKey customProviderWithAuthConfig() { + return factory.manufacturePojo(ProviderApiKey.class).toBuilder() + .provider(LlmProvider.CUSTOM_LLM) + .providerName(UUID.randomUUID().toString()) + .apiKey(null) + .authConfig(tokenAuthConfig()) + .build(); + } + + private ProviderAuthConfig storedAuthConfig(UUID id, String workspaceId) { + return mySqlTemplate.inTransaction(READ_ONLY, + handle -> handle.attach(LlmProviderApiKeyDAO.class).findById(id, workspaceId).authConfig()); + } + + @Test + @DisplayName("create stores the recipe encrypted and reads it back masked") + void createAndGetMasksSecrets() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var created = llmProviderApiKeyResourceClient.createProviderApiKey( + customProviderWithAuthConfig(), apiKey, workspaceName, HttpStatus.SC_CREATED); + + var actual = llmProviderApiKeyResourceClient.getById(created.id(), workspaceName, apiKey, + HttpStatus.SC_OK); + assertThat(actual.authConfig()).isEqualTo(tokenAuthConfig().mask()); + assertThat(actual.authConfig().credentials().getLast().value()) + .isEqualTo(ProviderAuthConfig.SECRET_SENTINEL); + + var page = llmProviderApiKeyResourceClient.getAll(workspaceName, apiKey); + assertThat(page.content().getFirst().authConfig()).isEqualTo(tokenAuthConfig().mask()); + + String rawColumn = mySqlTemplate.inTransaction(READ_ONLY, handle -> handle + .createQuery("SELECT auth_config FROM llm_provider_api_key WHERE id = :id") + .bind("id", created.id().toString()) + .mapTo(String.class) + .one()); + assertThat(rawColumn).doesNotContain(SECRET_VALUE, "client_id", "token_url"); + assertThat(EncryptionUtils.decryptGcm(rawColumn)).contains(SECRET_VALUE); + } + + @Test + @DisplayName("update with the sentinel keeps the stored secret, and a lock can't be removed") + void updateWithSentinelKeepsStoredSecret() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var created = llmProviderApiKeyResourceClient.createProviderApiKey( + customProviderWithAuthConfig(), apiKey, workspaceName, HttpStatus.SC_CREATED); + + // sentinel for the secret (also trying to unlock it), a new value for client_id + var update = ProviderApiKeyUpdate.builder() + .authConfig(tokenAuthConfig().toBuilder() + .credentials(List.of( + credential("grant_type", "client_credentials", false), + credential("client_id", "rotated-id", false), + credential("client_secret", ProviderAuthConfig.SECRET_SENTINEL, false))) + .build()) + .build(); + llmProviderApiKeyResourceClient.updateProviderApiKey(created.id(), update, apiKey, workspaceName, + HttpStatus.SC_NO_CONTENT); + + var stored = storedAuthConfig(created.id(), workspaceId); + var storedByKey = stored.credentials().stream() + .collect(java.util.stream.Collectors + .toMap(ProviderAuthConfig.Credential::key, Function.identity())); + assertThat(storedByKey.get("client_secret").value()).isEqualTo(SECRET_VALUE); + assertThat(storedByKey.get("client_secret").secret()).isTrue(); + assertThat(storedByKey.get("client_id").value()).isEqualTo("rotated-id"); + } + + @Test + @DisplayName("sentinel on a key without a stored secret is rejected") + void updateSentinelUnknownKeyIsRejected() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var created = llmProviderApiKeyResourceClient.createProviderApiKey( + customProviderWithAuthConfig(), apiKey, workspaceName, HttpStatus.SC_CREATED); + + var update = ProviderApiKeyUpdate.builder() + .authConfig(tokenAuthConfig().toBuilder() + .credentials(List.of( + credential("brand_new_key", ProviderAuthConfig.SECRET_SENTINEL, true))) + .build()) + .build(); + try (var response = llmProviderApiKeyResourceClient.callUpdateProviderApiKey(created.id(), update, + apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + assertThat(response.readEntity(ErrorMessage.class).getMessage()).contains("brand_new_key"); + } + } + + @Test + @DisplayName("empty object clears the auth config, allowing the switch back to a static key") + void clearAuthConfigAndSwitchToStaticKey() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var created = llmProviderApiKeyResourceClient.createProviderApiKey( + customProviderWithAuthConfig(), apiKey, workspaceName, HttpStatus.SC_CREATED); + + var update = ProviderApiKeyUpdate.builder() + .apiKey("brand-new-static-key") + .authConfig(ProviderAuthConfig.builder().build()) + .build(); + llmProviderApiKeyResourceClient.updateProviderApiKey(created.id(), update, apiKey, workspaceName, + HttpStatus.SC_NO_CONTENT); + + var actual = llmProviderApiKeyResourceClient.getById(created.id(), workspaceName, apiKey, + HttpStatus.SC_OK); + assertThat(actual.authConfig()).isNull(); + checkEncryption(created.id(), workspaceId, "brand-new-static-key"); + } + + @Test + @DisplayName("setting a static key while token auth is configured is rejected") + void updateApiKeyWhileTokenModeActiveIsRejected() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var created = llmProviderApiKeyResourceClient.createProviderApiKey( + customProviderWithAuthConfig(), apiKey, workspaceName, HttpStatus.SC_CREATED); + + var update = ProviderApiKeyUpdate.builder().apiKey("some-static-key").build(); + try (var response = llmProviderApiKeyResourceClient.callUpdateProviderApiKey(created.id(), update, + apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + } + } + + @Test + @DisplayName("auth config on a provider without provider_name is rejected") + void authConfigOnStandardProviderIsRejected() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + // create + llmProviderApiKeyResourceClient.createProviderApiKey( + createProviderApiKey().toBuilder().authConfig(tokenAuthConfig()).build(), + apiKey, workspaceName, HttpStatus.SC_UNPROCESSABLE_CONTENT); + + // update + var created = llmProviderApiKeyResourceClient.createProviderApiKey(createProviderApiKey(), apiKey, + workspaceName, HttpStatus.SC_CREATED); + var update = ProviderApiKeyUpdate.builder().authConfig(tokenAuthConfig()).build(); + try (var response = llmProviderApiKeyResourceClient.callUpdateProviderApiKey(created.id(), update, + apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + } + } + + @Test + @DisplayName("create rejects the sentinel, both credentials set, and an invalid token URL") + void createInvalidAuthConfigIsRejected() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + // sentinel on create: nothing stored to keep + var withSentinel = customProviderWithAuthConfig().toBuilder() + .authConfig(tokenAuthConfig().toBuilder() + .credentials(List.of( + credential("client_secret", ProviderAuthConfig.SECRET_SENTINEL, true))) + .build()) + .build(); + llmProviderApiKeyResourceClient.createProviderApiKey(withSentinel, apiKey, workspaceName, + HttpStatus.SC_UNPROCESSABLE_CONTENT); + + // static key and token auth at once + var withBoth = customProviderWithAuthConfig().toBuilder().apiKey("static-key").build(); + llmProviderApiKeyResourceClient.createProviderApiKey(withBoth, apiKey, workspaceName, + HttpStatus.SC_UNPROCESSABLE_CONTENT); + + // token URL that isn't an absolute URI + var withBadUrl = customProviderWithAuthConfig().toBuilder() + .authConfig(tokenAuthConfig().toBuilder().tokenUrl("not a uri").build()) + .build(); + llmProviderApiKeyResourceClient.createProviderApiKey(withBadUrl, apiKey, workspaceName, + HttpStatus.SC_UNPROCESSABLE_CONTENT); + } + } + + @Nested + @DisplayName("Auth config check endpoint:") + @TestInstance(TestInstance.Lifecycle.PER_CLASS) + class AuthConfigCheckEndpoint { + + private static final String SECRET_VALUE = "endpoint-s3cr3t"; + + /** unique path per test: the wiremock server is shared with the auth mocks, never reset here */ + private String stubTokenEndpoint(String body, int status) { + String tokenPath = "/provider-auth-token/" + UUID.randomUUID(); + wireMock.server().stubFor(com.github.tomakehurst.wiremock.client.WireMock.post(urlPathEqualTo(tokenPath)) + .willReturn(com.github.tomakehurst.wiremock.client.WireMock.aResponse() + .withStatus(status) + .withHeader("Content-Type", "application/json") + .withBody(body))); + return tokenPath; + } + + /** plain-http base: baseUrl() prefers the https port, whose self-signed cert the fetcher rejects */ + private String tokenUrl(String tokenPath) { + return "http://localhost:" + wireMock.server().port() + tokenPath; + } + + private ProviderAuthConfig recipe(String tokenUrl, String secretValue) { + return ProviderAuthConfig.builder() + .tokenUrl(tokenUrl) + .credentials(List.of( + credential("grant_type", "client_credentials", false), + credential("client_id", "opik-prod", false), + credential("client_secret", secretValue, true))) + .build(); + } + + private ProviderAuthConfig.Credential credential(String key, String value, boolean secret) { + return ProviderAuthConfig.Credential.builder().key(key).value(value).secret(secret).build(); + } + + private ProviderApiKey createCustomProvider(ProviderAuthConfig authConfig, String apiKey, + String workspaceName) { + var provider = factory.manufacturePojo(ProviderApiKey.class).toBuilder() + .provider(LlmProvider.CUSTOM_LLM) + .providerName(UUID.randomUUID().toString()) + .apiKey(null) + .authConfig(authConfig) + .build(); + return llmProviderApiKeyResourceClient.createProviderApiKey(provider, apiKey, workspaceName, + HttpStatus.SC_CREATED); + } + + @Test + @DisplayName("submitted values are tested and the lifetime is reported, never the token") + void testWithSubmittedAuthConfig() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + String tokenPath = stubTokenEndpoint("{\"access_token\": \"tok-endpoint\", \"expires_in\": 1800}", 200); + var request = ProviderAuthCheck.builder() + .authConfig(recipe(tokenUrl(tokenPath), SECRET_VALUE)) + .build(); + + try (var response = llmProviderApiKeyResourceClient.callTestAuthConfig(request, apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_OK); + String body = response.readEntity(String.class); + var result = JsonUtils.readValue(body, ProviderAuthCheck.Result.class); + assertThat(result.lifetimeSeconds()).isEqualTo(1800); + assertThat(body).doesNotContain("tok-endpoint"); + } + } + + @Test + @DisplayName("by provider id, the stored recipe is used with its real secrets, server-side") + void testWithStoredAuthConfig() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + String tokenPath = stubTokenEndpoint("{\"access_token\": \"tok\", \"expires_in\": 60}", 200); + var created = createCustomProvider(recipe(tokenUrl(tokenPath), SECRET_VALUE), apiKey, workspaceName); + + var request = ProviderAuthCheck.builder().providerId(created.id()).build(); + try (var response = llmProviderApiKeyResourceClient.callTestAuthConfig(request, apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_OK); + } + + wireMock.server().verify(postRequestedFor(urlPathEqualTo(tokenPath)) + .withRequestBody(com.github.tomakehurst.wiremock.client.WireMock + .containing("client_secret=" + SECRET_VALUE))); + } + + @Test + @DisplayName("sentinels in submitted values resolve against the stored recipe when the id is given") + void testResolvesSentinelsAgainstStoredConfig() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + String tokenPath = stubTokenEndpoint("{\"access_token\": \"tok\", \"expires_in\": 60}", 200); + var created = createCustomProvider(recipe(tokenUrl(tokenPath), SECRET_VALUE), apiKey, workspaceName); + + var request = ProviderAuthCheck.builder() + .providerId(created.id()) + .authConfig(recipe(tokenUrl(tokenPath), ProviderAuthConfig.SECRET_SENTINEL)) + .build(); + try (var response = llmProviderApiKeyResourceClient.callTestAuthConfig(request, apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_OK); + } + + wireMock.server().verify(postRequestedFor(urlPathEqualTo(tokenPath)) + .withRequestBody(com.github.tomakehurst.wiremock.client.WireMock + .containing("client_secret=" + SECRET_VALUE))); + } + + @Test + @DisplayName("sentinels without a provider id are rejected: there is nothing stored to resolve against") + void testSentinelWithoutIdIsRejected() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + var request = ProviderAuthCheck.builder() + .authConfig(recipe("https://auth.example.com/token", ProviderAuthConfig.SECRET_SENTINEL)) + .build(); + try (var response = llmProviderApiKeyResourceClient.callTestAuthConfig(request, apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + assertThat(response.readEntity(ErrorMessage.class).getMessage()).contains("client_secret"); + } + } + + @Test + @DisplayName("upstream auth failures surface status and body with credential values redacted") + void testSurfacesUpstreamErrorsRedacted() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + String tokenPath = stubTokenEndpoint( + "{\"error\": \"invalid_client\", \"echo\": \"%s\"}".formatted(SECRET_VALUE), 401); + var request = ProviderAuthCheck.builder() + .authConfig(recipe(tokenUrl(tokenPath), SECRET_VALUE)) + .build(); + + try (var response = llmProviderApiKeyResourceClient.callTestAuthConfig(request, apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + String message = response.readEntity(ErrorMessage.class).getMessage(); + assertThat(message).contains("401").contains("invalid_client").doesNotContain(SECRET_VALUE); + } + } + + @Test + @DisplayName("a request with neither id nor auth config, or an id without a stored recipe, is rejected") + void testInvalidRequestsAreRejected() { + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + try (var response = llmProviderApiKeyResourceClient.callTestAuthConfig( + ProviderAuthCheck.builder().build(), apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + assertThat(response.readEntity(ErrorMessage.class).getMessage()) + .contains("either provider_id or auth_config"); + } + + var staticProvider = llmProviderApiKeyResourceClient.createProviderApiKey( + createProviderApiKey(), apiKey, workspaceName, HttpStatus.SC_CREATED); + try (var response = llmProviderApiKeyResourceClient.callTestAuthConfig( + ProviderAuthCheck.builder().providerId(staticProvider.id()).build(), + apiKey, workspaceName)) { + assertThat(response.getStatus()).isEqualTo(HttpStatus.SC_BAD_REQUEST); + assertThat(response.readEntity(ErrorMessage.class).getMessage()).contains("no auth_config"); + } + } + } + private void getAndAssertProviderApiKey(ProviderApiKey expected, String apiKey, String workspaceName) { var actualEntity = llmProviderApiKeyResourceClient.getById(expected.id(), workspaceName, apiKey, 200); diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/EncryptionUtilsTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/EncryptionUtilsTest.java new file mode 100644 index 00000000000..60dd35b11a6 --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/EncryptionUtilsTest.java @@ -0,0 +1,52 @@ +package com.comet.opik.infrastructure; + +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Base64; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class EncryptionUtilsTest { + + @BeforeAll + static void setUpAll() { + var config = new OpikConfiguration(); + config.getEncryption().setKey("0123456789abcdef"); + EncryptionUtils.setConfig(config); + } + + @Test + void gcmRoundTrip() { + var plaintext = "{\"credentials\":[{\"key\":\"client_secret\",\"value\":\"s3cr3t\"}]}"; + + var encrypted = EncryptionUtils.encryptGcm(plaintext); + + assertThat(encrypted).doesNotContain("s3cr3t"); + assertThat(EncryptionUtils.decryptGcm(encrypted)).isEqualTo(plaintext); + } + + @Test + void gcmUsesARandomIvPerEncryption() { + var plaintext = "same plaintext"; + + assertThat(EncryptionUtils.encryptGcm(plaintext)).isNotEqualTo(EncryptionUtils.encryptGcm(plaintext)); + } + + @Test + void gcmRejectsTamperedCiphertext() { + byte[] payload = Base64.getDecoder().decode(EncryptionUtils.encryptGcm("payload")); + payload[payload.length - 1] ^= 1; + var tampered = Base64.getEncoder().encodeToString(payload); + + assertThatThrownBy(() -> EncryptionUtils.decryptGcm(tampered)).isInstanceOf(SecurityException.class); + } + + @Test + void gcmRejectsLegacyCiphertext() { + var legacy = EncryptionUtils.encrypt("payload"); + + assertThatThrownBy(() -> EncryptionUtils.decryptGcm(legacy)).isInstanceOf(SecurityException.class); + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/LlmProviderFactoryTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/LlmProviderFactoryTest.java index 63d1c49fd1c..713d7b4eeeb 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/LlmProviderFactoryTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/LlmProviderFactoryTest.java @@ -15,6 +15,7 @@ import com.comet.opik.infrastructure.llm.antropic.AnthropicClientGenerator; import com.comet.opik.infrastructure.llm.antropic.AnthropicModelName; import com.comet.opik.infrastructure.llm.antropic.AnthropicModule; +import com.comet.opik.infrastructure.llm.customllm.AuthTokenProvider; import com.comet.opik.infrastructure.llm.customllm.CustomLlmClientGenerator; import com.comet.opik.infrastructure.llm.customllm.CustomLlmModule; import com.comet.opik.infrastructure.llm.gemini.GeminiClientGenerator; @@ -181,7 +182,8 @@ void testCustomLlmProviderMatching_shouldMatch( // Register custom LLM service (required for getService to work) CustomLlmModule customLlmModule = new CustomLlmModule(); - CustomLlmClientGenerator customLlmClientGenerator = customLlmModule.clientGenerator(llmProviderClientConfig); + CustomLlmClientGenerator customLlmClientGenerator = customLlmModule.clientGenerator(llmProviderClientConfig, + mock(AuthTokenProvider.class)); customLlmModule.llmServiceProvider(llmProviderFactory, customLlmClientGenerator); // When & Then - Should successfully get the service @@ -221,7 +223,8 @@ void testCustomLlmProviderMatching_shouldNotMatch( // Register custom LLM service (required for getService to work) CustomLlmModule customLlmModule = new CustomLlmModule(); - CustomLlmClientGenerator customLlmClientGenerator = customLlmModule.clientGenerator(llmProviderClientConfig); + CustomLlmClientGenerator customLlmClientGenerator = customLlmModule.clientGenerator(llmProviderClientConfig, + mock(AuthTokenProvider.class)); customLlmModule.llmServiceProvider(llmProviderFactory, customLlmClientGenerator); // When & Then - Should throw BadRequestException diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/customllm/AuthTokenProviderTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/customllm/AuthTokenProviderTest.java new file mode 100644 index 00000000000..0c58aaaafbd --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/customllm/AuthTokenProviderTest.java @@ -0,0 +1,439 @@ +package com.comet.opik.infrastructure.llm.customllm; + +import com.comet.opik.api.ProviderAuthConfig; +import com.comet.opik.api.resources.utils.RedisContainerUtils; +import com.comet.opik.infrastructure.EncryptionUtils; +import com.comet.opik.infrastructure.LlmProviderTokenAuthConfig; +import com.comet.opik.infrastructure.OpikConfiguration; +import com.comet.opik.infrastructure.lock.LockService; +import com.comet.opik.infrastructure.net.DestinationGuard; +import com.comet.opik.infrastructure.redis.StringRedisClient; +import com.github.tomakehurst.wiremock.WireMockServer; +import com.github.tomakehurst.wiremock.core.WireMockConfiguration; +import com.redis.testcontainers.RedisContainer; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.redisson.Redisson; +import org.redisson.api.RedissonClient; +import org.redisson.api.options.KeysScanOptions; +import org.redisson.config.Config; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.time.Duration; +import java.util.Base64; +import java.util.List; +import java.util.UUID; + +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +import static com.github.tomakehurst.wiremock.client.WireMock.containing; +import static com.github.tomakehurst.wiremock.client.WireMock.equalTo; +import static com.github.tomakehurst.wiremock.client.WireMock.equalToJson; +import static com.github.tomakehurst.wiremock.client.WireMock.exactly; +import static com.github.tomakehurst.wiremock.client.WireMock.okJson; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static com.github.tomakehurst.wiremock.client.WireMock.postRequestedFor; +import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +@DisplayName("Auth Token Provider Test") +class AuthTokenProviderTest { + + private static final String TOKEN_PATH = "/oauth/token"; + private static final String SECRET_VALUE = "super-s3cr3t"; + + /** + * Single-flight correctness belongs to {@code RedissonLockService}'s own tests (which is also + * package-private); here the lock is a passthrough so the provider's orchestration is what's + * under test. + */ + private static final LockService PASSTHROUGH_LOCK_SERVICE = new LockService() { + @Override + public Mono executeWithLock(Lock lock, Mono action) { + return action; + } + + @Override + public Mono executeWithLockCustomExpire(Lock lock, Mono action, Duration duration) { + return action; + } + + @Override + public Flux executeWithLock(Lock lock, Flux action) { + return action; + } + + @Override + public Mono bestEffortLock(Lock lock, Mono action, Mono failToAcquireLockAction, + Duration actionTimeout, Duration lockTimeout) { + throw new UnsupportedOperationException(); + } + + @Override + public Mono bestEffortLock(Lock lock, Mono action, Mono failToAcquireLockAction, + Duration actionTimeout, Duration lockTimeout, boolean holdUntilExpiry) { + throw new UnsupportedOperationException(); + } + + @Override + public Mono lockUsingToken(Lock lock, Duration lockDuration) { + throw new UnsupportedOperationException(); + } + + @Override + public Mono unlockUsingToken(Lock lock) { + throw new UnsupportedOperationException(); + } + + @Override + public Mono tryAcquireSlot(Lock lock, int totalSlots, Duration leaseTime) { + throw new UnsupportedOperationException(); + } + + @Override + public Mono refreshSlot(Lock lock, String permitId, Duration leaseTime) { + throw new UnsupportedOperationException(); + } + + @Override + public Mono releaseSlot(Lock lock, String permitId) { + throw new UnsupportedOperationException(); + } + + @Override + public Mono addSlotPermits(Lock lock, int delta) { + throw new UnsupportedOperationException(); + } + }; + + private final RedisContainer redis = RedisContainerUtils.newRedisContainer(); + private WireMockServer wireMock; + private RedissonClient redisson; + private StringRedisClient stringRedisClient; + private AuthTokenProvider provider; + + @BeforeAll + void setUpAll() { + redis.start(); + var redissonConfig = new Config(); + redissonConfig.useSingleServer().setAddress(redis.getRedisURI()); + redisson = Redisson.create(redissonConfig); + stringRedisClient = new StringRedisClient(redisson); + + wireMock = new WireMockServer(WireMockConfiguration.options().dynamicPort()); + wireMock.start(); + + var opikConfiguration = new OpikConfiguration(); + opikConfiguration.getEncryption().setKey("0123456789abcdef"); + EncryptionUtils.setConfig(opikConfiguration); + + provider = new AuthTokenProvider(stringRedisClient, PASSTHROUGH_LOCK_SERVICE, relaxedConfig()); + } + + private static LlmProviderTokenAuthConfig relaxedConfig() { + var config = new LlmProviderTokenAuthConfig(); + config.setDestinationGuard(DestinationGuard.Mode.RELAXED); + return config; + } + + @AfterAll + void tearDownAll() { + wireMock.stop(); + redisson.shutdown(); + redis.stop(); + } + + @BeforeEach + void setUp() { + wireMock.resetAll(); + } + + private ProviderAuthConfig.ProviderAuthConfigBuilder oauthRecipe() { + return ProviderAuthConfig.builder() + .tokenUrl(wireMock.baseUrl() + TOKEN_PATH) + .credentials(List.of( + credential("grant_type", "client_credentials", false), + credential("client_id", "opik-prod", false), + credential("client_secret", SECRET_VALUE, true))); + } + + private ProviderAuthConfig.Credential credential(String key, String value, boolean secret) { + return ProviderAuthConfig.Credential.builder().key(key).value(value).secret(secret).build(); + } + + private void stubToken(String body) { + wireMock.stubFor(post(urlEqualTo(TOKEN_PATH)).willReturn(okJson(body))); + } + + private List cachedKeys(UUID providerId) { + var keys = new java.util.ArrayList(); + redisson.getKeys() + .getKeys(KeysScanOptions.defaults().pattern("llm_auth_token:%s:*".formatted(providerId))) + .forEach(keys::add); + return keys; + } + + @Test + @DisplayName("form mode fetches, caches encrypted, and serves cache hits") + void formModeFetchesAndCaches() { + stubToken("{\"access_token\": \"tok-1\", \"expires_in\": 3600}"); + var providerId = UUID.randomUUID(); + var recipe = oauthRecipe().build(); + + assertThat(provider.bearer("ws", providerId, recipe)).isEqualTo("tok-1"); + assertThat(provider.bearer("ws", providerId, recipe)).isEqualTo("tok-1"); + + wireMock.verify(exactly(1), postRequestedFor(urlEqualTo(TOKEN_PATH)) + .withHeader("Content-Type", equalTo("application/x-www-form-urlencoded")) + .withRequestBody(containing("grant_type=client_credentials")) + .withRequestBody(containing("client_secret=" + SECRET_VALUE))); + + // the cached value is ciphertext: neither the token nor the recipe is readable in Redis + var keys = cachedKeys(providerId); + assertThat(keys).hasSize(1); + String rawValue = stringRedisClient.getBucket(keys.getFirst()).get(); + assertThat(rawValue).doesNotContain("tok-1", SECRET_VALUE); + } + + @Test + @DisplayName("basic mode sends id/secret in the Basic header and the remaining fields in the body") + void basicModeSendsBasicHeader() { + stubToken("{\"access_token\": \"tok-basic\", \"expires_in\": 3600}"); + var recipe = oauthRecipe() + .sendAs(ProviderAuthConfig.SendAs.BASIC) + .credentials(List.of( + credential("grant_type", "client_credentials", false), + credential("scope", "data:read", false), + credential("client_id", "opik-prod", false), + credential("client_secret", SECRET_VALUE, true))) + .build(); + + assertThat(provider.bearer("ws", UUID.randomUUID(), recipe)).isEqualTo("tok-basic"); + + String expectedBasic = "Basic " + Base64.getEncoder() + .encodeToString(("opik-prod:" + SECRET_VALUE).getBytes()); + wireMock.verify(postRequestedFor(urlEqualTo(TOKEN_PATH)) + .withHeader("Authorization", equalTo(expectedBasic)) + .withRequestBody(equalTo("grant_type=client_credentials&scope=data%3Aread"))); + } + + @Test + @DisplayName("json mode sends the credentials as a JSON body") + void jsonModeSendsJsonBody() { + stubToken("{\"access_token\": \"tok-json\", \"expires_in\": 3600}"); + var recipe = oauthRecipe().sendAs(ProviderAuthConfig.SendAs.JSON).build(); + + assertThat(provider.bearer("ws", UUID.randomUUID(), recipe)).isEqualTo("tok-json"); + + wireMock.verify(postRequestedFor(urlEqualTo(TOKEN_PATH)) + .withHeader("Content-Type", equalTo("application/json")) + .withRequestBody(equalToJson( + "{\"grant_type\": \"client_credentials\", \"client_id\": \"opik-prod\", \"client_secret\": \"%s\"}" + .formatted(SECRET_VALUE)))); + } + + @Test + @DisplayName("dot-path token field and fallback lifetime cover the service-account shape") + void dotPathAndFallbackTtl() { + stubToken("{\"result\": {\"jwt\": \"tok-nested\"}}"); + var recipe = oauthRecipe() + .tokenField("result.jwt") + .expiresField("result.ttl") + .fallbackTtlSeconds(90_000L) + .build(); + + assertThat(provider.bearer("ws", UUID.randomUUID(), recipe)).isEqualTo("tok-nested"); + } + + @Test + @DisplayName("a zero fallback leaves tokens uncached when the reply states no lifetime") + void zeroTtlFetchesPerCall() { + stubToken("{\"access_token\": \"tok-0\"}"); + var providerId = UUID.randomUUID(); + var recipe = oauthRecipe().fallbackTtlSeconds(0L).build(); + + provider.bearer("ws", providerId, recipe); + provider.bearer("ws", providerId, recipe); + + wireMock.verify(exactly(2), postRequestedFor(urlEqualTo(TOKEN_PATH))); + assertThat(cachedKeys(providerId)).isEmpty(); + } + + @Test + @DisplayName("a reply-stated lifetime wins over the fallback, including a zero fallback") + void replyLifetimeWinsOverFallback() { + stubToken("{\"access_token\": \"tok-authoritative\", \"expires_in\": 3600}"); + var providerId = UUID.randomUUID(); + var recipe = oauthRecipe().fallbackTtlSeconds(0L).build(); + + provider.bearer("ws", providerId, recipe); + provider.bearer("ws", providerId, recipe); + + wireMock.verify(exactly(1), postRequestedFor(urlEqualTo(TOKEN_PATH))); + assertThat(cachedKeys(providerId)).hasSize(1); + } + + @Test + @DisplayName("a token inside the refresh window is refetched") + void refreshWindowTriggersRefetch() throws InterruptedException { + // 1s lifetime with the default 0.25 fraction: fresh for ~750ms, inside the window after + stubToken("{\"access_token\": \"tok-refresh\", \"expires_in\": 1}"); + var providerId = UUID.randomUUID(); + var recipe = oauthRecipe().build(); + + provider.bearer("ws", providerId, recipe); + provider.bearer("ws", providerId, recipe); + wireMock.verify(exactly(1), postRequestedFor(urlEqualTo(TOKEN_PATH))); + + Thread.sleep(800); + provider.bearer("ws", providerId, recipe); + + wireMock.verify(exactly(2), postRequestedFor(urlEqualTo(TOKEN_PATH))); + } + + @Test + @DisplayName("editing the recipe changes the cache key, so the old token is never served") + void configChangeChangesCacheKey() { + stubToken("{\"access_token\": \"tok-cfg\", \"expires_in\": 3600}"); + var providerId = UUID.randomUUID(); + + provider.bearer("ws", providerId, oauthRecipe().build()); + provider.bearer("ws", providerId, oauthRecipe() + .credentials(List.of( + credential("grant_type", "client_credentials", false), + credential("client_id", "opik-prod", false), + credential("client_secret", "rotated-secret", true))) + .build()); + + wireMock.verify(exactly(2), postRequestedFor(urlEqualTo(TOKEN_PATH))); + } + + @Test + @DisplayName("invalidate drops the cached token so the next call refetches") + void invalidateForcesRefetch() { + stubToken("{\"access_token\": \"tok-inv\", \"expires_in\": 3600}"); + var providerId = UUID.randomUUID(); + var recipe = oauthRecipe().build(); + + provider.bearer("ws", providerId, recipe); + provider.invalidate(providerId, recipe); + provider.bearer("ws", providerId, recipe); + + wireMock.verify(exactly(2), postRequestedFor(urlEqualTo(TOKEN_PATH))); + } + + @Test + @DisplayName("upstream errors surface status and body with credential values redacted") + void upstreamErrorIsSurfacedRedacted() { + wireMock.stubFor(post(urlEqualTo(TOKEN_PATH)).willReturn(aResponse() + .withStatus(401) + .withBody("{\"error\": \"invalid_client\", \"echo\": \"%s\"}".formatted(SECRET_VALUE)))); + + assertThatThrownBy(() -> provider.bearer("ws", UUID.randomUUID(), oauthRecipe().build())) + .isInstanceOf(AuthTokenException.class) + .hasMessageContaining("401") + .hasMessageContaining("invalid_client") + .satisfies(exception -> assertThat(exception.getMessage()).doesNotContain(SECRET_VALUE)); + } + + @Test + @DisplayName("a missing token field names the reply's top-level fields, never values") + void missingTokenFieldListsTopLevelKeys() { + stubToken("{\"token\": \"the-actual-token\", \"ttl\": 60}"); + + assertThatThrownBy(() -> provider.bearer("ws", UUID.randomUUID(), oauthRecipe().build())) + .isInstanceOf(AuthTokenException.class) + .hasMessageContaining("access_token") + .hasMessageContaining("[token, ttl]") + .satisfies(exception -> assertThat(exception.getMessage()).doesNotContain("the-actual-token")); + } + + @Test + @DisplayName("an absurdly large lifetime is rejected before it can corrupt cache arithmetic") + void oversizedLifetimeIsRejected() { + stubToken("{\"access_token\": \"tok-huge\", \"expires_in\": 9223372036854775807}"); + + assertThatThrownBy(() -> provider.bearer("ws", UUID.randomUUID(), oauthRecipe().build())) + .isInstanceOf(AuthTokenException.class) + .hasMessageContaining("outside the accepted range"); + } + + @Test + @DisplayName("a reply without a lifetime and without a fallback is a clear error") + void missingLifetimeWithoutFallbackIsRejected() { + stubToken("{\"access_token\": \"tok-nolife\"}"); + + assertThatThrownBy(() -> provider.bearer("ws", UUID.randomUUID(), oauthRecipe().build())) + .isInstanceOf(AuthTokenException.class) + .hasMessageContaining("expires_in") + .hasMessageContaining("no fallback lifetime"); + } + + @Test + @DisplayName("a non-JSON reply is a clear error") + void nonJsonReplyIsRejected() { + wireMock.stubFor(post(urlEqualTo(TOKEN_PATH)).willReturn(aResponse() + .withStatus(200) + .withBody("gateway login page"))); + + assertThatThrownBy(() -> provider.bearer("ws", UUID.randomUUID(), oauthRecipe().build())) + .isInstanceOf(AuthTokenException.class) + .hasMessageContaining("non-JSON reply"); + } + + @Test + @DisplayName("an unreachable token URL is a clear error, not a hang") + void unreachableTokenUrlIsRejected() { + var recipe = oauthRecipe().tokenUrl("http://localhost:1/token").build(); + + assertThatThrownBy(() -> provider.bearer("ws", UUID.randomUUID(), recipe)) + .isInstanceOf(AuthTokenException.class) + .hasMessageContaining("could not reach"); + } + + @Test + @DisplayName("a strict destination guard refuses the token URL before any request is made") + void strictGuardRefusesNonPublicTokenUrl() { + var strictConfig = new LlmProviderTokenAuthConfig(); + strictConfig.setDestinationGuard(DestinationGuard.Mode.STRICT); + var strictProvider = new AuthTokenProvider(stringRedisClient, PASSTHROUGH_LOCK_SERVICE, strictConfig); + + assertThatThrownBy(() -> strictProvider.bearer("ws", UUID.randomUUID(), oauthRecipe().build())) + .isInstanceOf(AuthTokenException.class) + .hasMessageContaining("only https"); + wireMock.verify(exactly(0), postRequestedFor(urlEqualTo(TOKEN_PATH))); + } + + @Test + @DisplayName("a Redis outage degrades to a direct fetch instead of failing the call") + void redisOutageDegradesToDirectFetch() { + var flakyRedis = RedisContainerUtils.newRedisContainer(); + flakyRedis.start(); + var flakyConfig = new Config(); + flakyConfig.useSingleServer() + .setAddress(flakyRedis.getRedisURI()) + .setTimeout(500) + .setRetryAttempts(0) + .setConnectTimeout(500); + var flakyRedisson = Redisson.create(flakyConfig); + try { + var flakyProvider = new AuthTokenProvider(new StringRedisClient(flakyRedisson), + PASSTHROUGH_LOCK_SERVICE, relaxedConfig()); + flakyRedis.stop(); + + stubToken("{\"access_token\": \"tok-degraded\", \"expires_in\": 3600}"); + + assertThat(flakyProvider.bearer("ws", UUID.randomUUID(), oauthRecipe().build())) + .isEqualTo("tok-degraded"); + wireMock.verify(exactly(1), postRequestedFor(urlEqualTo(TOKEN_PATH))); + } finally { + flakyRedisson.shutdown(); + } + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClientTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClientTest.java index 34a4f188ca0..84e1e69146b 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClientTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/llm/customllm/InterceptingHttpClientTest.java @@ -1,9 +1,13 @@ package com.comet.opik.infrastructure.llm.customllm; +import dev.langchain4j.exception.HttpException; import dev.langchain4j.http.client.HttpClient; import dev.langchain4j.http.client.HttpMethod; import dev.langchain4j.http.client.HttpRequest; import dev.langchain4j.http.client.SuccessfulHttpResponse; +import dev.langchain4j.http.client.sse.ServerSentEvent; +import dev.langchain4j.http.client.sse.ServerSentEventListener; +import dev.langchain4j.http.client.sse.ServerSentEventParser; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; @@ -14,7 +18,12 @@ import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -211,4 +220,165 @@ void suppressDefaultAuthIgnoredWhenNoCustomHeaderConfigured() { assertThat(captor.getValue().headers()) .containsEntry("Authorization", List.of("Bearer dummy-key")); } + + // --- dynamic token auth --- + + private HttpRequest chatRequest() { + return HttpRequest.builder() + .method(HttpMethod.POST) + .url("https://example.test/chat/completions") + .body("{\"model\":\"gpt-4o\"}") + .build(); + } + + private SuccessfulHttpResponse ok() { + return SuccessfulHttpResponse.builder().statusCode(200).body("{}").build(); + } + + @Test + void tokenIsInjectedAsBearerOnEveryRequest() { + when(delegate.execute(any(HttpRequest.class))).thenReturn(ok()); + var client = new InterceptingHttpClient(delegate, Map.of(), null, () -> "tok-1", () -> { + }); + + client.execute(chatRequest()); + + var captor = ArgumentCaptor.forClass(HttpRequest.class); + verify(delegate).execute(captor.capture()); + assertThat(captor.getValue().headers()).containsEntry("Authorization", List.of("Bearer tok-1")); + } + + @Test + void tokenReplacesPreConfiguredStaticAuthorizationHeader() { + when(delegate.execute(any(HttpRequest.class))).thenReturn(ok()); + var client = new InterceptingHttpClient(delegate, Map.of(), null, () -> "tok-1", () -> { + }); + var request = HttpRequest.builder() + .method(HttpMethod.POST) + .url("https://example.test/chat/completions") + .addHeader("Authorization", "Bearer stale-static-key") + .addHeader("X-Extra", "kept") + .body("{}") + .build(); + + client.execute(request); + + var captor = ArgumentCaptor.forClass(HttpRequest.class); + verify(delegate).execute(captor.capture()); + assertThat(captor.getValue().headers()) + .containsEntry("Authorization", List.of("Bearer tok-1")) + .containsEntry("X-Extra", List.of("kept")); + } + + @Test + void tokenGoesRawUnderConfiguredAuthHeaderName() { + when(delegate.execute(any(HttpRequest.class))).thenReturn(ok()); + var configuration = Map.of("auth_header_name", "api-key"); + var client = new InterceptingHttpClient(delegate, configuration, null, () -> "tok-1", () -> { + }); + var request = HttpRequest.builder() + .method(HttpMethod.POST) + .url("https://example.test/chat/completions") + .addHeader("Authorization", "Bearer stale") + .body("{}") + .build(); + + client.execute(request); + + var captor = ArgumentCaptor.forClass(HttpRequest.class); + verify(delegate).execute(captor.capture()); + assertThat(captor.getValue().headers()) + .containsEntry("api-key", List.of("tok-1")) + .doesNotContainKey("Authorization"); + } + + @Test + void authFailureInvalidatesAndRetriesOnceWithAFreshToken() { + var tokens = new java.util.concurrent.atomic.AtomicInteger(); + var invalidated = new java.util.concurrent.atomic.AtomicBoolean(); + when(delegate.execute(any(HttpRequest.class))) + .thenThrow(new HttpException(401, "token rejected")) + .thenReturn(ok()); + var client = new InterceptingHttpClient(delegate, Map.of(), null, + () -> "tok-" + tokens.incrementAndGet(), () -> invalidated.set(true)); + + client.execute(chatRequest()); + + assertThat(invalidated).isTrue(); + var captor = ArgumentCaptor.forClass(HttpRequest.class); + verify(delegate, times(2)).execute(captor.capture()); + assertThat(captor.getAllValues().getFirst().headers()) + .containsEntry("Authorization", List.of("Bearer tok-1")); + assertThat(captor.getAllValues().getLast().headers()) + .containsEntry("Authorization", List.of("Bearer tok-2")); + } + + @Test + void authFailureIsRetriedExactlyOnce() { + when(delegate.execute(any(HttpRequest.class))).thenThrow(new HttpException(401, "still rejected")); + var client = new InterceptingHttpClient(delegate, Map.of(), null, () -> "tok", () -> { + }); + + assertThatThrownBy(() -> client.execute(chatRequest())) + .isInstanceOf(HttpException.class); + verify(delegate, times(2)).execute(any(HttpRequest.class)); + } + + @Test + void nonAuthFailureIsNotRetried() { + when(delegate.execute(any(HttpRequest.class))).thenThrow(new HttpException(500, "gateway down")); + var client = new InterceptingHttpClient(delegate, Map.of(), null, () -> "tok", () -> { + }); + + assertThatThrownBy(() -> client.execute(chatRequest())) + .isInstanceOf(HttpException.class); + verify(delegate, times(1)).execute(any(HttpRequest.class)); + } + + @Test + void streamedAuthFailureRetriesOnlyBeforeTheFirstDeliveredEvent() { + var downstream = mock(ServerSentEventListener.class); + var parser = mock(ServerSentEventParser.class); + var tokens = new java.util.concurrent.atomic.AtomicInteger(); + var client = new InterceptingHttpClient(delegate, Map.of(), null, + () -> "tok-" + tokens.incrementAndGet(), () -> { + }); + + client.execute(chatRequest(), parser, downstream); + + var requestCaptor = ArgumentCaptor.forClass(HttpRequest.class); + var listenerCaptor = ArgumentCaptor.forClass(ServerSentEventListener.class); + verify(delegate).execute(requestCaptor.capture(), any(), listenerCaptor.capture()); + assertThat(requestCaptor.getValue().headers()).containsEntry("Authorization", List.of("Bearer tok-1")); + + // 401 arrives before any event: the wrapper retries once, straight to the downstream listener + listenerCaptor.getValue().onError(new HttpException(401, "token rejected")); + + verify(delegate).execute(requestCaptor.capture(), any(), same(downstream)); + assertThat(requestCaptor.getValue().headers()).containsEntry("Authorization", List.of("Bearer tok-2")); + verify(downstream, never()).onError(any()); + } + + @Test + void streamedAuthFailureAfterAnEventPropagatesWithoutRetry() { + var downstream = mock(ServerSentEventListener.class); + var parser = mock(ServerSentEventParser.class); + var client = new InterceptingHttpClient(delegate, Map.of(), null, () -> "tok", () -> { + }); + + client.execute(chatRequest(), parser, downstream); + + var listenerCaptor = ArgumentCaptor.forClass(ServerSentEventListener.class); + verify(delegate).execute(any(HttpRequest.class), any(), listenerCaptor.capture()); + + var event = new ServerSentEvent("message", "chunk"); + listenerCaptor.getValue().onEvent(event); + var failure = new HttpException(401, "cut mid-stream"); + listenerCaptor.getValue().onError(failure); + + verify(downstream).onEvent(event); + verify(downstream).onError(failure); + // no second execute: content already flowed, splicing two half-responses would be worse + verify(delegate, times(1)).execute(any(HttpRequest.class), any(), any()); + } } diff --git a/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/net/DestinationGuardTest.java b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/net/DestinationGuardTest.java new file mode 100644 index 00000000000..c53cd0a4b8b --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/infrastructure/net/DestinationGuardTest.java @@ -0,0 +1,74 @@ +package com.comet.opik.infrastructure.net; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@DisplayName("Destination Guard Test") +class DestinationGuardTest { + + private final DestinationGuard strict = new DestinationGuard(DestinationGuard.Mode.STRICT); + private final DestinationGuard relaxed = new DestinationGuard(DestinationGuard.Mode.RELAXED); + + @ParameterizedTest + @ValueSource(strings = { + "http://public.example.com/token", // https only + "https://localhost/token", + "https://127.0.0.1/token", + "https://10.1.2.3/token", + "https://172.16.0.1/token", + "https://192.168.1.1/token", + "https://169.254.169.254/latest/meta-data", // cloud metadata endpoint + "https://[::1]/token", + "https://[fe80::1]/token", + "https://[fc00::1]/token", // IPv6 unique-local + "https://0.0.0.0/token", + "not a url", + }) + @DisplayName("strict mode refuses non-https, private, internal and malformed destinations") + void strictModeRefuses(String url) { + assertThatThrownBy(() -> strict.validate(url)).isInstanceOf(DestinationGuardException.class); + } + + @Test + @DisplayName("strict mode refuses unresolvable hosts") + void strictModeRefusesUnresolvableHosts() { + // .invalid is reserved (RFC 2606): guaranteed NXDOMAIN, no DNS flakiness + assertThatThrownBy(() -> strict.validate("https://token-endpoint.invalid/token")) + .isInstanceOf(DestinationGuardException.class) + .hasMessageContaining("could not be resolved"); + } + + @ParameterizedTest + @ValueSource(strings = { + "https://8.8.8.8/token", + "https://1.1.1.1/oauth/token", + }) + @DisplayName("strict mode allows public https destinations") + void strictModeAllowsPublicHttps(String url) { + assertThatCode(() -> strict.validate(url)).doesNotThrowAnyException(); + } + + @ParameterizedTest + @ValueSource(strings = { + "http://localhost:9876/token", + "https://10.1.2.3/token", + "https://169.254.169.254/latest/meta-data", + }) + @DisplayName("relaxed mode is a no-op: internal gateways are legitimate self-hosted destinations") + void relaxedModeAllowsEverything(String url) { + assertThatCode(() -> relaxed.validate(url)).doesNotThrowAnyException(); + } + + @Test + @DisplayName("refusal messages echo the user's hostname, never the resolved address") + void refusalNeverEchoesTheResolvedAddress() { + assertThatThrownBy(() -> strict.validate("https://localhost/token")) + .hasMessageContaining("localhost") + .hasMessageNotContainingAny("127.0.0.1", "::1"); + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/podam/manufacturer/ProviderApiKeyUpdateManufacturer.java b/apps/opik-backend/src/test/java/com/comet/opik/podam/manufacturer/ProviderApiKeyUpdateManufacturer.java index 7c6ba69787a..65d48f946c5 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/podam/manufacturer/ProviderApiKeyUpdateManufacturer.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/podam/manufacturer/ProviderApiKeyUpdateManufacturer.java @@ -16,22 +16,23 @@ public class ProviderApiKeyUpdateManufacturer extends AbstractTypeManufacturer

    { + return useMutation< + ProviderAuthCheckResult, + AxiosError, + ProviderAuthCheckRequest + >({ + mutationFn: async (request: ProviderAuthCheckRequest) => { + const { data } = await api.post( + `${PROVIDER_KEYS_REST_ENDPOINT}auth-config/test`, + request, + ); + return data; + }, + }); +}; + +export default useProviderKeysAuthCheckMutation; diff --git a/apps/opik-frontend/src/api/provider-keys/useProviderKeysCreateMutation.ts b/apps/opik-frontend/src/api/provider-keys/useProviderKeysCreateMutation.ts index 738933562ea..2c94562f1f3 100644 --- a/apps/opik-frontend/src/api/provider-keys/useProviderKeysCreateMutation.ts +++ b/apps/opik-frontend/src/api/provider-keys/useProviderKeysCreateMutation.ts @@ -31,6 +31,9 @@ const useProviderKeysCreateMutation = () => { configuration: providerKey.configuration, }), ...(providerKey?.headers && { headers: providerKey.headers }), + ...(providerKey.auth_config !== undefined && { + auth_config: providerKey.auth_config, + }), }); return data; diff --git a/apps/opik-frontend/src/api/provider-keys/useProviderKeysUpdateMutation.ts b/apps/opik-frontend/src/api/provider-keys/useProviderKeysUpdateMutation.ts index 9bc6b56c67a..ad4d0ac4912 100644 --- a/apps/opik-frontend/src/api/provider-keys/useProviderKeysUpdateMutation.ts +++ b/apps/opik-frontend/src/api/provider-keys/useProviderKeysUpdateMutation.ts @@ -28,6 +28,10 @@ const useProviderKeysUpdateMutation = () => { configuration: providerKey.configuration, }), ...(providerKey?.headers && { headers: providerKey.headers }), + // {} clears the stored recipe, so the check is on undefined, not truthiness + ...(providerKey.auth_config !== undefined && { + auth_config: providerKey.auth_config, + }), }, ); return data; diff --git a/apps/opik-frontend/src/shared/EyeInput/EyeInput.tsx b/apps/opik-frontend/src/shared/EyeInput/EyeInput.tsx index b2af38ed707..4abe0f293f4 100644 --- a/apps/opik-frontend/src/shared/EyeInput/EyeInput.tsx +++ b/apps/opik-frontend/src/shared/EyeInput/EyeInput.tsx @@ -4,13 +4,16 @@ import { Eye, EyeOff } from "lucide-react"; import { cn } from "@/lib/utils"; import { Button } from "@/ui/button"; -interface EyeInputProps extends InputProps {} +interface EyeInputProps extends InputProps { + revealable?: boolean; +} -const EyeInput = (props: EyeInputProps) => { +const EyeInput = ({ revealable = true, ...props }: EyeInputProps) => { const [hidden, setHidden] = useState(true); const id = useId(); - const Icon = hidden ? Eye : EyeOff; + const isHidden = hidden || !revealable; + const Icon = isHidden ? Eye : EyeOff; return (

    @@ -20,21 +23,23 @@ const EyeInput = (props: EyeInputProps) => { style={ { ...(props?.style || {}), - WebkitTextSecurity: hidden ? "disc" : "none", + WebkitTextSecurity: isHidden ? "disc" : "none", } as React.CSSProperties } - className={cn(props.className, "pr-8")} + className={cn(props.className, revealable && "pr-8")} /> - + {revealable && ( + + )}
    ); }; diff --git a/apps/opik-frontend/src/types/providers.ts b/apps/opik-frontend/src/types/providers.ts index 20403e84bcf..49e7fa0c4a1 100644 --- a/apps/opik-frontend/src/types/providers.ts +++ b/apps/opik-frontend/src/types/providers.ts @@ -853,6 +853,34 @@ export interface ProviderKeyConfiguration { openai_pipeline_mode?: OpenAiPipelineMode; } +/** Read-back sentinel for secret credential values: the backend never returns a stored secret, + * it returns this marker instead. Submitting it back means "keep the stored value". */ +export const AUTH_SECRET_SENTINEL = "__SECRET__"; + +export const AUTH_SEND_AS_VALUES = ["form", "json", "basic"] as const; +export type AuthSendAs = (typeof AUTH_SEND_AS_VALUES)[number]; + +export interface ProviderAuthCredential { + key: string; + value: string; + /** Secret values are encrypted at rest and read back masked; once true it cannot be unset. */ + secret: boolean; +} + +/** Dynamic token auth recipe (OPIK-7940): how to fetch a short-lived bearer before LLM calls. + * Presence on a provider means token mode; an empty object on update clears it (back to static). */ +export interface ProviderAuthConfig { + token_url: string; + send_as?: AuthSendAs; + credentials: ProviderAuthCredential[]; + /** Field holding the token in the reply; dot-path for nested replies (e.g. "result.jwt"). */ + token_field?: string; + /** Field holding the lifetime in seconds in the reply; dot-path supported. */ + expires_field?: string; + /** Lifetime assumed when the reply doesn't state one; 0 means such tokens are not cached. */ + fallback_ttl_seconds?: number; +} + export interface BaseProviderKey { id: string; created_at: string; @@ -860,6 +888,7 @@ export interface BaseProviderKey { ui_composed_provider: COMPOSED_PROVIDER_TYPE; configuration: ProviderKeyConfiguration; headers?: Record; + auth_config?: ProviderAuthConfig; read_only: boolean; } @@ -887,7 +916,7 @@ export type ProviderObject = | OllamaProviderObject; export type PartialProviderKeyUpdate = Partial< - Omit + Omit > & { id?: string; provider?: PROVIDER_TYPE; @@ -895,6 +924,8 @@ export type PartialProviderKeyUpdate = Partial< base_url?: string; provider_name?: string; headers?: Record; + /** Full recipe to set, or an empty object to clear (switch back to static key). */ + auth_config?: ProviderAuthConfig | Record; }; export type ReasoningEffort = diff --git a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/AuthConfigSection.tsx b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/AuthConfigSection.tsx new file mode 100644 index 00000000000..8612582edb4 --- /dev/null +++ b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/AuthConfigSection.tsx @@ -0,0 +1,327 @@ +import React, { useCallback } from "react"; +import { UseFormReturn, useFieldArray } from "react-hook-form"; +import { v4 as uuidv4 } from "uuid"; +import get from "lodash/get"; +import { AxiosError } from "axios"; +import { Lock, LockOpen, Plus, Trash2 } from "lucide-react"; + +import { cn } from "@/lib/utils"; +import { Button } from "@/ui/button"; +import { Label } from "@/ui/label"; +import { FormControl, FormField, FormItem, FormMessage } from "@/ui/form"; +import { Input } from "@/ui/input"; +import { Description } from "@/ui/description"; +import { RadioGroup, RadioGroupItem } from "@/ui/radio-group"; +import { FormFieldCard } from "@/v2/pages-shared/llm/FormFieldCard"; +import { useToast } from "@/ui/use-toast"; +import EyeInput from "@/shared/EyeInput/EyeInput"; +import ExplainerIcon from "@/shared/ExplainerIcon/ExplainerIcon"; +import TooltipWrapper from "@/shared/TooltipWrapper/TooltipWrapper"; +import { AUTH_SECRET_SENTINEL } from "@/types/providers"; +import useProviderKeysAuthCheckMutation from "@/api/provider-keys/useProviderKeysAuthCheckMutation"; +import { AIProviderFormType } from "@/v2/pages-shared/llm/ManageAIProviderDialog/schema"; +import { + AUTH_SECRET_KEY_PATTERN, + AuthMode, + formValuesToAuthConfig, + oauth2CredentialRows, +} from "@/v2/pages-shared/llm/ManageAIProviderDialog/customProviderConfig"; + +type AuthConfigSectionProps = { + form: UseFormReturn; + /** The provider's static-auth fields (API key etc.), rendered only while static mode is selected. */ + staticModeFields: React.ReactNode; +}; + +/** + * The Authentication block of the custom/Bedrock provider form: a mode switch between the classic + * static API key (rendered via {@code staticModeFields}) and dynamic token auth (OPIK-7940). Of + * the backend's general token-auth recipe, the UI surfaces only the OAuth2 client credentials + * flow — other recipe shapes remain API-only. + * Secret credential values are write-only once saved — they load as the backend's sentinel and + * their lock cannot be removed, mirroring the API contract. + */ +const AuthConfigSection: React.FC = ({ + form, + staticModeFields, +}) => { + const { toast } = useToast(); + const { mutate: checkAuthConfig, isPending: isChecking } = + useProviderKeysAuthCheckMutation(); + + const authMode = form.watch("authMode") ?? "api_key"; + const headers = form.watch("headers"); + const hasStaticAuthorizationHeader = (headers ?? []).some( + (header) => header.key.trim().toLowerCase() === "authorization", + ); + + const { fields, append, remove } = useFieldArray({ + control: form.control, + name: "authCredentials", + }); + + const handleModeChange = useCallback( + (value: AuthMode) => { + form.setValue("authMode", value); + if ( + value === "token" && + (form.getValues("authCredentials") ?? []).length === 0 + ) { + form.setValue("authCredentials", oauth2CredentialRows()); + } + }, + [form], + ); + + const handleCheckConnection = useCallback(() => { + const authConfig = formValuesToAuthConfig( + { ...form.getValues(), authMode: "token" }, + { isEditing: false, hadAuthConfig: false }, + ); + const providerId = form.getValues("id"); + + checkAuthConfig( + { + // the id lets the backend resolve __SECRET__ sentinels against the stored recipe + ...(providerId && { provider_id: providerId }), + auth_config: authConfig as never, + }, + { + onSuccess: (result) => { + toast({ + title: "Connection successful", + description: `Token received, valid for ${result.lifetime_seconds} seconds.`, + }); + }, + onError: (error: AxiosError) => { + const message = + get(error, ["response", "data", "message"]) ?? + get(error, ["response", "data", "errors", "0"], error.message); + toast({ + title: "Connection failed", + description: String(message), + variant: "destructive", + }); + }, + }, + ); + }, [form, checkAuthConfig, toast]); + + const handleCredentialKeyChange = useCallback( + (index: number, key: string) => { + form.setValue(`authCredentials.${index}.key`, key); + // auto-lock names that look like secrets, visibly, as the user types + if ( + AUTH_SECRET_KEY_PATTERN.test(key) && + !form.getValues(`authCredentials.${index}.secret`) + ) { + form.setValue(`authCredentials.${index}.secret`, true); + } + }, + [form], + ); + + return ( + + handleModeChange(value as AuthMode)} + className="flex gap-6" + > +
    + + +
    +
    + + + +
    +     + + {authMode === "api_key" && staticModeFields} + + {authMode === "token" && ( + <> + { + const validationErrors = get(formState.errors, ["authTokenUrl"]); + return ( + + + + field.onChange(e.target.value)} + className={cn({ + "border-destructive": Boolean( + validationErrors?.message, + ), + })} + /> + + + + ); + }} + /> + +
    +
    + + +
    + + {fields.map((row, index) => { + const isSecret = form.watch(`authCredentials.${index}.secret`); + const isSaved = row.saved; + const value = form.watch(`authCredentials.${index}.value`); + const isStoredSecret = isSaved && value === AUTH_SECRET_SENTINEL; + const keyErrors = get(form.formState.errors, [ + "authCredentials", + index, + "key", + ]); + + return ( +
    +
    + + handleCredentialKeyChange(index, e.target.value) + } + className={cn({ + "border-destructive": Boolean(keyErrors?.message), + })} + /> + {keyErrors?.message && ( + {String(keyErrors.message)} + )} +
    +
    + {isSecret ? ( + + form.setValue( + `authCredentials.${index}.value`, + e.target.value, + ) + } + /> + ) : ( + + form.setValue( + `authCredentials.${index}.value`, + e.target.value, + ) + } + /> + )} +
    + + + + +
    + ); + })} + +
    + +
    +
    + + {hasStaticAuthorizationHeader && ( + + A custom Authorization header is configured below — + it is ignored while token auth is on; the fetched token takes + precedence. + + )} + + + + + + )} +     + ); +}; + +export default AuthConfigSection; diff --git a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/BedrockProviderDetails.tsx b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/BedrockProviderDetails.tsx index 87063bbe6f7..2441c496ada 100644 --- a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/BedrockProviderDetails.tsx +++ b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/BedrockProviderDetails.tsx @@ -10,6 +10,7 @@ import { Input } from "@/ui/input"; import { Description } from "@/ui/description"; import { Button } from "@/ui/button"; import CustomHeadersField from "./CustomHeadersField"; +import AuthConfigSection from "./AuthConfigSection"; type BedrockProviderDetailsProps = { form: UseFormReturn; @@ -80,50 +81,6 @@ const BedrockProviderDetails: React.FC = ({ }} /> - { - const validationErrors = get(formState.errors, ["apiKey"]); - - return ( - - - - field.onChange(e.target.value)} - className={cn({ - "border-destructive": Boolean(validationErrors?.message), - })} - /> - - - - Click{" "} - {" "} - for instructions on how to create a service account and assign - the correct permissions. - - - ); - }} - /> = ({ /> + + { + const validationErrors = get(formState.errors, ["apiKey"]); + + return ( + + + + field.onChange(e.target.value)} + className={cn({ + "border-destructive": Boolean( + validationErrors?.message, + ), + })} + /> + + + + Click{" "} + {" "} + for instructions on how to create a service account and + assign the correct permissions. + + + ); + }} + /> + } + /> ); }; diff --git a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/CustomProviderDetails.tsx b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/CustomProviderDetails.tsx index 54c02d914cd..7380d1bda4e 100644 --- a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/CustomProviderDetails.tsx +++ b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/CustomProviderDetails.tsx @@ -14,6 +14,7 @@ import { Switch } from "@/ui/switch"; import { PROVIDERS } from "@/constants/providers"; import { PROVIDER_TYPE } from "@/types/providers"; import CustomHeadersField from "./CustomHeadersField"; +import AuthConfigSection from "./AuthConfigSection"; type CustomProviderDetailsProps = { form: UseFormReturn; @@ -93,50 +94,6 @@ const CustomProviderDetails: React.FC = ({ }} /> - { - const validationErrors = get(formState.errors, ["apiKey"]); - - return ( - - - - field.onChange(e.target.value)} - className={cn({ - "border-destructive": Boolean(validationErrors?.message), - })} - /> - - - - Custom providers may not require an API key, depending on your - server setup. Learn more in the{" "} - - . - - - ); - }} - /> = ({ description="Appended to every outgoing request URL. Some gateways require a version parameter such as api-version=2024-08-01-preview." /> - { - const validationErrors = get(formState.errors, ["authHeaderName"]); + + { + const validationErrors = get(formState.errors, ["apiKey"]); - return ( - - - - field.onChange(e.target.value)} - className={cn({ - "border-destructive": Boolean(validationErrors?.message), - })} - /> - - - - If set, the API key is sent as {"{name}: "} in - addition to the default Authorization: Bearer{" "} - header. - - - ); - }} - /> + return ( + + + + field.onChange(e.target.value)} + className={cn({ + "border-destructive": Boolean( + validationErrors?.message, + ), + })} + /> + + + + Custom providers may not require an API key, depending on + your server setup. Learn more in the{" "} + + . + + + ); + }} + /> - ( - -
    - - -
    - - Turn on only if your gateway rejects requests that include{" "} - Authorization: Bearer. - -     - )} + { + const validationErrors = get(formState.errors, [ + "authHeaderName", + ]); + + return ( + + + + field.onChange(e.target.value)} + className={cn({ + "border-destructive": Boolean( + validationErrors?.message, + ), + })} + /> + + + + If set, the API key is sent as{" "} + {"{name}: "} in addition to the default{" "} + Authorization: Bearer header. + + + ); + }} + /> + + ( + +
    + + +
    + + Turn on only if your gateway rejects requests that include{" "} + Authorization: Bearer. + +     + )} + /> + + } /> ); diff --git a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/ManageAIProviderDialog.tsx b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/ManageAIProviderDialog.tsx index a65919c3ac7..7c464c3e604 100644 --- a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/ManageAIProviderDialog.tsx +++ b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/ManageAIProviderDialog.tsx @@ -42,8 +42,12 @@ import ProviderSelectionStep from "@/v2/pages-shared/llm/SetupProviderDialog/Pro import ProviderConfigurationStep from "@/v2/pages-shared/llm/SetupProviderDialog/ProviderConfigurationStep"; import { + AuthConfigFormValues, + EMPTY_AUTH_FORM_VALUES, + authConfigToFormValues, configStringToQueryParamsArray, convertHeadersForAPI, + formValuesToAuthConfig, queryParamsArrayToConfigString, } from "./customProviderConfig"; @@ -139,6 +143,7 @@ const ManageAIProviderDialog: React.FC = ({ openaiPipelineMode: normalizeOpenAiPipelineMode( providerKey?.configuration?.openai_pipeline_mode, ), + ...authConfigToFormValues(providerKey?.auth_config), } as AIProviderFormType, }); @@ -188,6 +193,7 @@ const ManageAIProviderDialog: React.FC = ({ authHeaderName: "", suppressDefaultAuth: false, openaiPipelineMode: DEFAULT_OPENAI_PIPELINE_MODE, + ...EMPTY_AUTH_FORM_VALUES, }); setStep("select"); }, [form]); @@ -248,6 +254,15 @@ const ManageAIProviderDialog: React.FC = ({ ), ); + const authFormValues = authConfigToFormValues(providerData?.auth_config); + form.setValue("authMode", authFormValues.authMode); + form.setValue("authTokenUrl", authFormValues.authTokenUrl); + form.setValue("authSendAs", authFormValues.authSendAs); + form.setValue("authCredentials", authFormValues.authCredentials); + form.setValue("authTokenField", authFormValues.authTokenField); + form.setValue("authExpiresField", authFormValues.authExpiresField); + form.setValue("authFallbackTtl", authFormValues.authFallbackTtl); + form.setValue("provider", providerType); form.setValue("composedProviderType", composedProviderType); form.setValue("apiKey", ""); @@ -262,7 +277,7 @@ const ManageAIProviderDialog: React.FC = ({ resetSelectionState(); }, [resetSelectionState]); - const cloudConfigHandler = useCallback(() => { + const submitProviderHandler = useCallback(() => { const apiKey = form.getValues("apiKey"); const url = form.getValues("url"); const location = form.getValues("location"); @@ -321,14 +336,31 @@ const ManageAIProviderDialog: React.FC = ({ isCustomLike && !!(providerKey || calculatedProviderKey); const headers = convertHeadersForAPI(headersArray, isEditingCustomProvider); + const storedProvider = providerKey ?? calculatedProviderKey; + const authConfig = isCustomLike + ? // the isCustomLike guard means the union's custom member (which carries the + // auth fields) is the live one; TS can't narrow a union by that boolean + formValuesToAuthConfig( + form.getValues() as Partial, + { + isEditing: isEditingCustomProvider, + hadAuthConfig: Boolean(storedProvider?.auth_config), + }, + ) + : undefined; + // token mode and static key are mutually exclusive on the API + const isTokenMode = isCustomLike && form.getValues("authMode") === "token"; + const effectiveApiKey = isTokenMode ? "" : apiKey; + if (providerKey || calculatedProviderKey) { updateMutate({ providerKey: { id: providerKey?.id ?? calculatedProviderKey?.id, - apiKey, + apiKey: effectiveApiKey, base_url: isCustomLike ? url : undefined, ...(configuration && { configuration }), ...(isCustomLike && headers !== undefined && { headers }), + ...(authConfig !== undefined && { auth_config: authConfig }), }, }); } else if (provider) { @@ -338,12 +370,13 @@ const ManageAIProviderDialog: React.FC = ({ createMutate({ providerKey: { - apiKey, + apiKey: effectiveApiKey, provider, base_url: isCustomLike ? url : undefined, provider_name: isCustomLike ? providerName : undefined, ...(configuration && { configuration }), ...(isCustomLike && headers !== undefined && { headers }), + ...(authConfig !== undefined && { auth_config: authConfig }), }, }); } @@ -428,7 +461,7 @@ const ManageAIProviderDialog: React.FC = ({ @@ -471,7 +504,7 @@ const ManageAIProviderDialog: React.FC = ({ )} diff --git a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/customProviderConfig.test.ts b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/customProviderConfig.test.ts index c1eb4a8ae20..1d487ad577c 100644 --- a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/customProviderConfig.test.ts +++ b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/customProviderConfig.test.ts @@ -1,10 +1,14 @@ import { describe, it, expect } from "vitest"; import { + authConfigToFormValues, configStringToQueryParamsArray, convertHeadersForAPI, + formValuesToAuthConfig, + oauth2CredentialRows, queryParamsArrayToConfigString, } from "./customProviderConfig"; +import { ProviderAuthConfig } from "@/types/providers"; describe("customProviderConfig", () => { describe("queryParamsArrayToConfigString", () => { @@ -128,4 +132,189 @@ describe("customProviderConfig", () => { }); }); }); + + describe("authConfigToFormValues", () => { + it("maps an absent config to static mode", () => { + const values = authConfigToFormValues(undefined); + expect(values.authMode).toBe("api_key"); + expect(values.authCredentials).toEqual([]); + }); + + it("loads a stored config with rows marked as saved", () => { + const stored: ProviderAuthConfig = { + token_url: "https://auth.example.com/token", + send_as: "basic", + credentials: [ + { key: "client_id", value: "opik", secret: false }, + { key: "client_secret", value: "__SECRET__", secret: true }, + ], + token_field: "access_token", + expires_field: "expires_in", + fallback_ttl_seconds: 3600, + }; + + const values = authConfigToFormValues(stored); + + expect(values.authMode).toBe("token"); + expect(values.authTokenUrl).toBe("https://auth.example.com/token"); + expect(values.authSendAs).toBe("basic"); + expect(values.authFallbackTtl).toBe("3600"); + expect(values.authCredentials).toHaveLength(2); + expect(values.authCredentials.every((row) => row.saved)).toBe(true); + expect(values.authCredentials[1].value).toBe("__SECRET__"); + }); + + it("hides the injected client_credentials grant row but keeps a custom one visible", () => { + const values = authConfigToFormValues({ + token_url: "https://auth.example.com/token", + credentials: [ + { key: "grant_type", value: "client_credentials", secret: false }, + { key: "client_id", value: "opik", secret: false }, + ], + }); + expect(values.authCredentials.map((row) => row.key)).toEqual([ + "client_id", + ]); + + const custom = authConfigToFormValues({ + token_url: "https://auth.example.com/token", + credentials: [{ key: "grant_type", value: "password", secret: false }], + }); + expect(custom.authCredentials.map((row) => row.key)).toEqual([ + "grant_type", + ]); + }); + + it("stringifies a zero fallback (fetch-per-call) rather than dropping it", () => { + const values = authConfigToFormValues({ + token_url: "https://auth.example.com/token", + credentials: [], + fallback_ttl_seconds: 0, + }); + expect(values.authFallbackTtl).toBe("0"); + }); + }); + + describe("formValuesToAuthConfig", () => { + const tokenValues = { + authMode: "token" as const, + authTokenUrl: " https://auth.example.com/token ", + authSendAs: "form" as const, + authCredentials: [ + { key: "username", value: "svc", secret: false, saved: false, id: "1" }, + { key: "password", value: "p", secret: true, saved: false, id: "2" }, + { key: " ", value: "dropped", secret: false, saved: false, id: "3" }, + ], + authTokenField: "token", + authExpiresField: "", + authFallbackTtl: "90000", + }; + + it("builds the recipe in token mode, trimming and dropping empty-key rows", () => { + const config = formValuesToAuthConfig(tokenValues, { + isEditing: false, + hadAuthConfig: false, + }); + + expect(config).toEqual({ + token_url: "https://auth.example.com/token", + send_as: "form", + credentials: [ + { key: "grant_type", value: "client_credentials", secret: false }, + { key: "username", value: "svc", secret: false }, + { key: "password", value: "p", secret: true }, + ], + token_field: "token", + fallback_ttl_seconds: 90000, + }); + }); + + it("does not inject grant_type when the user provided their own", () => { + const config = formValuesToAuthConfig( + { + ...tokenValues, + authCredentials: [ + { + key: "grant_type", + value: "password", + secret: false, + saved: false, + id: "1", + }, + ], + }, + { isEditing: false, hadAuthConfig: false }, + ); + + expect((config as ProviderAuthConfig).credentials).toEqual([ + { key: "grant_type", value: "password", secret: false }, + ]); + }); + + it("defaults send_as to basic per the OAuth2 client credentials flow", () => { + const config = formValuesToAuthConfig( + { authMode: "token", authTokenUrl: "https://auth.example.com/token" }, + { isEditing: false, hadAuthConfig: false }, + ); + expect(config).toMatchObject({ send_as: "basic" }); + }); + + it("returns {} to clear when switching back to static on an edited provider", () => { + const config = formValuesToAuthConfig( + { ...tokenValues, authMode: "api_key" }, + { isEditing: true, hadAuthConfig: true }, + ); + expect(config).toEqual({}); + }); + + it("returns undefined in static mode when there is nothing to clear", () => { + expect( + formValuesToAuthConfig( + { ...tokenValues, authMode: "api_key" }, + { isEditing: true, hadAuthConfig: false }, + ), + ).toBeUndefined(); + expect( + formValuesToAuthConfig( + { ...tokenValues, authMode: "api_key" }, + { isEditing: false, hadAuthConfig: false }, + ), + ).toBeUndefined(); + }); + + it("round-trips a loaded config, preserving the secret sentinel for unchanged values", () => { + const stored: ProviderAuthConfig = { + token_url: "https://auth.example.com/token", + send_as: "basic", + credentials: [ + { key: "grant_type", value: "client_credentials", secret: false }, + { key: "client_id", value: "opik", secret: false }, + { key: "client_secret", value: "__SECRET__", secret: true }, + ], + token_field: "access_token", + }; + + const roundTripped = formValuesToAuthConfig( + authConfigToFormValues(stored), + { isEditing: true, hadAuthConfig: true }, + ); + + expect(roundTripped).toEqual(stored); + }); + }); + + describe("oauth2CredentialRows", () => { + it("seeds only the user-owned rows, locking the secret; grant_type is injected on save", () => { + expect( + oauth2CredentialRows().map((row) => [row.key, row.secret]), + ).toEqual([ + ["client_id", false], + ["client_secret", true], + ]); + }); + + it("marks seeded rows as unsaved so their locks stay toggleable", () => { + expect(oauth2CredentialRows().every((row) => !row.saved)).toBe(true); + }); + }); }); diff --git a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/customProviderConfig.ts b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/customProviderConfig.ts index 3d5420056e6..bedc922e819 100644 --- a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/customProviderConfig.ts +++ b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/customProviderConfig.ts @@ -1,11 +1,149 @@ import { v4 as uuidv4 } from "uuid"; +import { AuthSendAs, ProviderAuthConfig } from "@/types/providers"; + export type KeyValueEntry = { key: string; value: string; id: string; }; +export const AUTH_MODE_VALUES = ["api_key", "token"] as const; +export type AuthMode = (typeof AUTH_MODE_VALUES)[number]; + +export type AuthCredentialEntry = { + key: string; + value: string; + secret: boolean; + saved: boolean; // Loaded from a stored provider: its lock can never be removed, per the backend contract + id: string; +}; + +export type AuthConfigFormValues = { + authMode: AuthMode; + authTokenUrl: string; + authSendAs: AuthSendAs; + authCredentials: AuthCredentialEntry[]; + authTokenField: string; + authExpiresField: string; + authFallbackTtl: string; +}; + +/** Field names the backend auto-locks; mirrored here so the lock flips visibly as the user types. */ +export const AUTH_SECRET_KEY_PATTERN = /secret|password|key|token|credential/i; + +// The one grant the UI supports at the moment. Hidden from the credentials list and injected on save. +// A grant_type row the user adds themselves — or one stored with a different value — wins over the +// injected default. +const OAUTH2_GRANT_TYPE = { + key: "grant_type", + value: "client_credentials", + secret: false, +}; + +export const EMPTY_AUTH_FORM_VALUES: AuthConfigFormValues = { + authMode: "api_key", + authTokenUrl: "", + authSendAs: "basic", + authCredentials: [], + authTokenField: "", + authExpiresField: "", + authFallbackTtl: "", +}; + +/** Loads a stored provider's auth config into form state (absent config -> static mode). */ +export function authConfigToFormValues( + authConfig: ProviderAuthConfig | undefined, +): AuthConfigFormValues { + if (!authConfig) { + return { ...EMPTY_AUTH_FORM_VALUES }; + } + return { + authMode: "token", + authTokenUrl: authConfig.token_url ?? "", + authSendAs: authConfig.send_as ?? "basic", + authCredentials: (authConfig.credentials ?? []) + .filter( + (credential) => + credential.key !== OAUTH2_GRANT_TYPE.key || + credential.value !== OAUTH2_GRANT_TYPE.value, + ) + .map((credential) => ({ + key: credential.key, + value: credential.value, + secret: credential.secret, + saved: true, + id: uuidv4(), + })), + authTokenField: authConfig.token_field ?? "", + authExpiresField: authConfig.expires_field ?? "", + authFallbackTtl: + authConfig.fallback_ttl_seconds !== undefined && + authConfig.fallback_ttl_seconds !== null + ? String(authConfig.fallback_ttl_seconds) + : "", + }; +} + +/** + * Builds the auth_config payload from form state. + * + * Three cases, mirroring the headers convention: + * 1. Token mode -> the full recipe (empty-key rows dropped) + * 2. Static mode while editing a provider that had a config -> {} to clear it + * 3. Static mode otherwise -> undefined (field omitted) + */ +export function formValuesToAuthConfig( + formValues: Partial, + options: { isEditing: boolean; hadAuthConfig: boolean }, +): ProviderAuthConfig | Record | undefined { + const values: AuthConfigFormValues = { + ...EMPTY_AUTH_FORM_VALUES, + ...formValues, + }; + if (values.authMode !== "token") { + return options.isEditing && options.hadAuthConfig ? {} : undefined; + } + + const credentials = values.authCredentials + .filter((credential) => credential.key.trim().length > 0) + .map((credential) => ({ + key: credential.key.trim(), + value: credential.value, + secret: credential.secret, + })); + if (!credentials.some(({ key }) => key === OAUTH2_GRANT_TYPE.key)) { + credentials.unshift({ ...OAUTH2_GRANT_TYPE }); + } + + const fallbackTtl = values.authFallbackTtl.trim(); + return { + token_url: values.authTokenUrl.trim(), + send_as: values.authSendAs, + credentials, + ...(values.authTokenField.trim() && { + token_field: values.authTokenField.trim(), + }), + ...(values.authExpiresField.trim() && { + expires_field: values.authExpiresField.trim(), + }), + ...(fallbackTtl && { fallback_ttl_seconds: Number(fallbackTtl) }), + }; +} + +export function oauth2CredentialRows(): AuthCredentialEntry[] { + return [ + { key: "client_id", value: "", secret: false, saved: false, id: uuidv4() }, + { + key: "client_secret", + value: "", + secret: true, + saved: false, + id: uuidv4(), + }, + ]; +} + /** * Converts header array from form state to API-compatible object format. * diff --git a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/schema.ts b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/schema.ts index ad42a94fcf2..73b23e56ce5 100644 --- a/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/schema.ts +++ b/apps/opik-frontend/src/v2/pages-shared/llm/ManageAIProviderDialog/schema.ts @@ -2,10 +2,12 @@ import { z } from "zod"; import uniq from "lodash/uniq"; import { + AUTH_SEND_AS_VALUES, OPENAI_PIPELINE_MODE_VALUES, OpenAiPipelineMode, PROVIDER_TYPE, } from "@/types/providers"; +import { AUTH_MODE_VALUES } from "./customProviderConfig"; export type { OpenAiPipelineMode }; export { OPENAI_PIPELINE_MODE_VALUES }; @@ -118,8 +120,81 @@ export const createCustomProviderDetailsFormSchema = ( .optional(), authHeaderName: z.string().max(150).optional(), suppressDefaultAuth: z.boolean().optional(), + authMode: z.enum(AUTH_MODE_VALUES).optional(), + authTokenUrl: z.string().optional(), + authSendAs: z.enum(AUTH_SEND_AS_VALUES).optional(), + authCredentials: z + .array( + z.object({ + // max mirrors the backend's @Size(max = 250) on Credential.key + key: z.string().max(250), + value: z.string(), + secret: z.boolean(), + saved: z.boolean(), + id: z.string(), + }), + ) + .optional(), + authTokenField: z.string().max(250).optional(), + authExpiresField: z.string().max(250).optional(), + authFallbackTtl: z.string().optional(), }) .superRefine((data, ctx) => { + // Token-auth mode requirements: field-level rules stay optional because static mode is the + // default and none of these apply there. + if (data.authMode === "token") { + const tokenUrl = (data.authTokenUrl ?? "").trim(); + if (!tokenUrl || !z.string().url().safeParse(tokenUrl).success) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Token URL must be a valid URL", + path: ["authTokenUrl"], + }); + } + + const credentials = data.authCredentials ?? []; + if (!credentials.some((entry) => entry.key.trim().length > 0)) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "At least one credential is required", + path: ["authTokenUrl"], + }); + } + + const credentialKeys: string[] = []; + credentials.forEach((entry, index) => { + const hasKey = entry.key.trim().length > 0; + if (!hasKey && entry.value.trim().length > 0) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Credential key is required", + path: ["authCredentials", index, "key"], + }); + } + if (hasKey) { + const trimmedKey = entry.key.trim(); + if (credentialKeys.includes(trimmedKey)) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Credential key must be unique", + path: ["authCredentials", index, "key"], + }); + } else { + credentialKeys.push(trimmedKey); + } + } + }); + + const fallbackTtl = (data.authFallbackTtl ?? "").trim(); + // digits-only implies non-negative; no separate sign/number check needed + if (fallbackTtl.length > 0 && !/^\d+$/.test(fallbackTtl)) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Fallback lifetime must be a whole number of seconds", + path: ["authFallbackTtl"], + }); + } + } // Validate headers: if a header has any content, both key and value must be non-empty if (data.headers) { const headerKeys: string[] = []; diff --git a/deployment/docker-compose/docker-compose.yaml b/deployment/docker-compose/docker-compose.yaml index 4470ccfa05f..cb9724b21d6 100644 --- a/deployment/docker-compose/docker-compose.yaml +++ b/deployment/docker-compose/docker-compose.yaml @@ -171,6 +171,9 @@ services: # Ollie / Agent Insights read-only freeform SQL (default off). When true, the backend provisions the restricted # read-only user (provision_agent_insights_readonly_user.sh, after migrations) and connects as it. TOGGLE_OLLIE_ENABLED: ${TOGGLE_OLLIE_ENABLED:-"false"} + # Self-hosted deployments legitimately point custom-provider token auth at internal gateways; + # the application default is "strict" (SSRF guard) for non-self-hosted deployments. + LLM_PROVIDER_TOKEN_AUTH_DESTINATION_GUARD: ${LLM_PROVIDER_TOKEN_AUTH_DESTINATION_GUARD:-"relaxed"} ANALYTICS_DB_READ_ONLY_FREEFORM_SQL_USER: ${ANALYTICS_DB_READ_ONLY_FREEFORM_SQL_USER:-comet_readonly_freeform_sql_user} ANALYTICS_DB_READ_ONLY_FREEFORM_SQL_PASS: ${ANALYTICS_DB_READ_ONLY_FREEFORM_SQL_PASS:-opik} # Traces cutover knobs (data-migrations/traces-local-v2-cutover). Default off / unset; the operator sets these diff --git a/deployment/helm_chart/opik/README.md b/deployment/helm_chart/opik/README.md index 1a40e51d1ec..b332ec7852f 100644 --- a/deployment/helm_chart/opik/README.md +++ b/deployment/helm_chart/opik/README.md @@ -249,6 +249,7 @@ Call opik api on http://localhost:5173/api | component.backend.env.LLM_MODEL_REGISTRY_REFRESH_INTERVAL_SECONDS | string | `"300"` | | | component.backend.env.LLM_MODEL_REGISTRY_REMOTE_ENABLED | string | `"false"` | | | component.backend.env.LLM_MODEL_REGISTRY_REMOTE_URL | string | `""` | | +| component.backend.env.LLM_PROVIDER_TOKEN_AUTH_DESTINATION_GUARD | string | `"relaxed"` | | | component.backend.env.OPIK_OTEL_SDK_ENABLED | bool | `false` | | | component.backend.env.OTEL_EXPERIMENTAL_EXPORTER_OTLP_RETRY_ENABLED | bool | `true` | | | component.backend.env.OTEL_EXPERIMENTAL_RESOURCE_DISABLED_KEYS | string | `"process.command_args"` | | diff --git a/deployment/helm_chart/opik/values.yaml b/deployment/helm_chart/opik/values.yaml index 9393b1121df..07eeb7a10be 100644 --- a/deployment/helm_chart/opik/values.yaml +++ b/deployment/helm_chart/opik/values.yaml @@ -178,6 +178,9 @@ component: ANALYTICS_DB_DATABASE_NAME: "opik" JAVA_OPTS: "-Dliquibase.propertySubstitutionEnabled=true -XX:+UseG1GC -XX:MaxRAMPercentage=80.0 -XX:MinRAMPercentage=75" REDIS_URL: redis://:wFSuJX9nDBdCa25sKZG7bh@opik-redis-master:6379/ + # Self-hosted deployments legitimately point custom-provider token auth at internal gateways; + # the application default is "strict" (SSRF guard) for non-self-hosted deployments. + LLM_PROVIDER_TOKEN_AUTH_DESTINATION_GUARD: "relaxed" ANALYTICS_DB_MIGRATIONS_PASS: opik ANALYTICS_DB_PASS: opik STATE_DB_PASS: opik