From 5f31de631b86e53dd1e9566733eaee84465779e7 Mon Sep 17 00:00:00 2001 From: Piotr Mlocek Date: Fri, 13 Mar 2026 14:46:04 -0700 Subject: [PATCH] feat(inference): validate endpoints before saving routes Closes #273 Verify inference endpoints synchronously on the server during set/update, expose a --no-verify escape hatch in the CLI and Python helper, and return actionable failures when validation does not pass. --- Cargo.lock | 2 + architecture/inference-routing.md | 15 +- crates/openshell-cli/src/main.rs | 69 +++- crates/openshell-cli/src/run.rs | 4 + crates/openshell-server/Cargo.toml | 2 + crates/openshell-server/src/inference.rs | 398 ++++++++++++++++++++++- docs/inference/configure.md | 4 +- proto/inference.proto | 2 + python/openshell/sandbox.py | 2 + python/openshell/sandbox_test.py | 34 ++ 10 files changed, 511 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f4bc2e4d5..051d9e25f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2989,6 +2989,7 @@ dependencies = [ "prost-types", "rand 0.9.2", "rcgen", + "reqwest", "russh", "rustls", "rustls-pemfile", @@ -3008,6 +3009,7 @@ dependencies = [ "tracing", "tracing-subscriber", "uuid", + "wiremock", ] [[package]] diff --git a/architecture/inference-routing.md b/architecture/inference-routing.md index b8ce8a805..851823f3b 100644 --- a/architecture/inference-routing.md +++ b/architecture/inference-routing.md @@ -66,8 +66,9 @@ The gateway implements the `Inference` gRPC service defined in `proto/inference. 1. Validates that both fields are non-empty. 2. Fetches the named provider record from the store. 3. Validates the provider by resolving its route (checking that the provider type is supported and has a usable API key). -4. Builds a managed route spec that stores only `provider_name` and `model_id`. The spec intentionally leaves `base_url`, `api_key`, and `protocols` empty -- these are resolved dynamically at bundle time from the provider record. -5. Upserts the route with name `inference.local`. Version starts at 1 and increments monotonically on each update. +4. Unless `skip_validation` is set, performs a lightweight provider-shaped probe against the resolved upstream endpoint (for example, a tiny chat/messages request with `max_tokens: 1`) to confirm the endpoint is reachable and accepts the expected auth/request shape. +5. Builds a managed route spec that stores only `provider_name` and `model_id`. The spec intentionally leaves `base_url`, `api_key`, and `protocols` empty -- these are resolved dynamically at bundle time from the provider record. +6. Upserts the route with name `inference.local`. Version starts at 1 and increments monotonically on each update. `GetClusterInference` returns `provider_name`, `model_id`, and `version` for the managed route. Returns `NOT_FOUND` if cluster inference is not configured. @@ -91,7 +92,7 @@ File: `proto/inference.proto` Key messages: -- `SetClusterInferenceRequest` -- `provider_name` + `model_id` +- `SetClusterInferenceRequest` -- `provider_name` + `model_id` + optional `skip_validation` - `SetClusterInferenceResponse` -- `provider_name` + `model_id` + `version` - `GetInferenceBundleResponse` -- `repeated ResolvedRoute routes` + `revision` + `generated_at_ms` - `ResolvedRoute` -- `name`, `base_url`, `protocols`, `api_key`, `model_id`, `provider_type` @@ -296,10 +297,10 @@ The system route is stored as a separate `InferenceRoute` record in the gateway Cluster inference commands: -- `openshell cluster inference set --provider --model ` -- configures user-facing cluster inference -- `openshell cluster inference set --system --provider --model ` -- configures system inference -- `openshell cluster inference get` -- displays both user and system inference configuration -- `openshell cluster inference get --system` -- displays only the system inference configuration +- `openshell inference set --provider --model ` -- configures user-facing cluster inference +- `openshell inference set --system --provider --model ` -- configures system inference +- `openshell inference get` -- displays both user and system inference configuration +- `openshell inference get --system` -- displays only the system inference configuration The `--provider` flag references a provider record name (not a provider type). The provider must already exist in the cluster and have a supported inference type (`openai`, `anthropic`, or `nvidia`). diff --git a/crates/openshell-cli/src/main.rs b/crates/openshell-cli/src/main.rs index dcca47034..ca82095ab 100644 --- a/crates/openshell-cli/src/main.rs +++ b/crates/openshell-cli/src/main.rs @@ -906,9 +906,6 @@ enum InferenceCommands { system: bool, /// Skip endpoint verification before saving the route. - /// - /// Accepted now so scripts can opt out explicitly ahead of a future - /// default switch to verification. #[arg(long)] no_verify: bool, }, @@ -929,9 +926,6 @@ enum InferenceCommands { system: bool, /// Skip endpoint verification before saving the route. - /// - /// Accepted now so scripts can opt out explicitly ahead of a future - /// default switch to verification. #[arg(long)] no_verify: bool, }, @@ -1810,17 +1804,27 @@ async fn main() -> Result<()> { provider, model, system, +<<<<<<< HEAD no_verify: _, +======= + no_verify, +>>>>>>> 7f0504d8 (feat(inference): validate endpoints before saving routes) } => { let route_name = if system { "sandbox-system" } else { "" }; - run::gateway_inference_set(endpoint, &provider, &model, route_name, &tls) - .await?; + run::gateway_inference_set( + endpoint, &provider, &model, route_name, no_verify, &tls, + ) + .await?; } InferenceCommands::Update { provider, model, system, +<<<<<<< HEAD no_verify: _, +======= + no_verify, +>>>>>>> 7f0504d8 (feat(inference): validate endpoints before saving routes) } => { let route_name = if system { "sandbox-system" } else { "" }; run::gateway_inference_update( @@ -1828,6 +1832,7 @@ async fn main() -> Result<()> { provider.as_deref(), model.as_deref(), route_name, + no_verify, &tls, ) .await?; @@ -2559,6 +2564,54 @@ mod tests { )); } + #[test] + fn inference_set_accepts_no_verify_flag() { + let cli = Cli::try_parse_from([ + "openshell", + "inference", + "set", + "--provider", + "openai-dev", + "--model", + "gpt-4.1", + "--no-verify", + ]) + .expect("inference set should parse --no-verify"); + + assert!(matches!( + cli.command, + Some(Commands::Inference { + command: Some(InferenceCommands::Set { + no_verify: true, + .. + }) + }) + )); + } + + #[test] + fn inference_update_accepts_no_verify_flag() { + let cli = Cli::try_parse_from([ + "openshell", + "inference", + "update", + "--provider", + "openai-dev", + "--no-verify", + ]) + .expect("inference update should parse --no-verify"); + + assert!(matches!( + cli.command, + Some(Commands::Inference { + command: Some(InferenceCommands::Update { + no_verify: true, + .. + }) + }) + )); + } + #[test] fn completion_script_uses_openshell_command_name() { let script = normalize_completion_script( diff --git a/crates/openshell-cli/src/run.rs b/crates/openshell-cli/src/run.rs index 42ecbbb1e..878c8c2ca 100644 --- a/crates/openshell-cli/src/run.rs +++ b/crates/openshell-cli/src/run.rs @@ -3390,6 +3390,7 @@ pub async fn gateway_inference_set( provider_name: &str, model_id: &str, route_name: &str, + skip_validation: bool, tls: &TlsOptions, ) -> Result<()> { let mut client = grpc_inference_client(server, tls).await?; @@ -3398,6 +3399,7 @@ pub async fn gateway_inference_set( provider_name: provider_name.to_string(), model_id: model_id.to_string(), route_name: route_name.to_string(), + skip_validation, }) .await .into_diagnostic()?; @@ -3422,6 +3424,7 @@ pub async fn gateway_inference_update( provider_name: Option<&str>, model_id: Option<&str>, route_name: &str, + skip_validation: bool, tls: &TlsOptions, ) -> Result<()> { if provider_name.is_none() && model_id.is_none() { @@ -3449,6 +3452,7 @@ pub async fn gateway_inference_update( provider_name: provider.to_string(), model_id: model.to_string(), route_name: route_name.to_string(), + skip_validation, }) .await .into_diagnostic()?; diff --git a/crates/openshell-server/Cargo.toml b/crates/openshell-server/Cargo.toml index 7d53abe63..c49350240 100644 --- a/crates/openshell-server/Cargo.toml +++ b/crates/openshell-server/Cargo.toml @@ -61,6 +61,7 @@ serde = { workspace = true } serde_json = { workspace = true } tokio-stream = { workspace = true } sqlx = { workspace = true } +reqwest = { workspace = true } kube = { workspace = true } kube-runtime = { workspace = true } k8s-openapi = { workspace = true } @@ -78,6 +79,7 @@ rcgen = { version = "0.13", features = ["crypto", "pem"] } tempfile = "3" tokio-tungstenite = { workspace = true } futures-util = "0.3" +wiremock = "0.6" [lints] workspace = true diff --git a/crates/openshell-server/src/inference.rs b/crates/openshell-server/src/inference.rs index f7a8427ec..35b704f17 100644 --- a/crates/openshell-server/src/inference.rs +++ b/crates/openshell-server/src/inference.rs @@ -1,12 +1,14 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use openshell_core::inference::AuthHeader; use openshell_core::proto::{ ClusterInferenceConfig, GetClusterInferenceRequest, GetClusterInferenceResponse, GetInferenceBundleRequest, GetInferenceBundleResponse, InferenceRoute, Provider, ResolvedRoute, SetClusterInferenceRequest, SetClusterInferenceResponse, inference_server::Inference, }; use std::sync::Arc; +use std::time::Duration; use tonic::{Request, Response, Status}; use crate::{ @@ -81,6 +83,7 @@ impl Inference for InferenceService { route_name, &req.provider_name, &req.model_id, + req.skip_validation, ) .await?; @@ -111,7 +114,7 @@ impl Inference for InferenceService { .map_err(|e| Status::internal(format!("fetch route failed: {e}")))? .ok_or_else(|| { Status::not_found(format!( - "inference route '{route_name}' is not configured; run 'openshell cluster inference set --provider --model '" + "inference route '{route_name}' is not configured; run 'openshell inference set --provider --model '" )) })?; @@ -140,6 +143,7 @@ async fn upsert_cluster_inference_route( route_name: &str, provider_name: &str, model_id: &str, + skip_validation: bool, ) -> Result { if provider_name.trim().is_empty() { return Err(Status::invalid_argument("provider_name is required")); @@ -156,9 +160,10 @@ async fn upsert_cluster_inference_route( Status::failed_precondition(format!("provider '{provider_name}' not found")) })?; - // Validate provider shape at set time; endpoint/auth are resolved from the - // provider record when generating sandbox bundles. - let _ = resolve_provider_route(&provider)?; + let resolved = resolve_provider_route(&provider)?; + if !skip_validation { + verify_provider_endpoint(&provider.name, model_id, &resolved).await?; + } let config = build_cluster_inference_config(&provider, model_id); @@ -203,6 +208,13 @@ struct ResolvedProviderRoute { base_url: String, protocols: Vec, api_key: String, + auth: AuthHeader, + default_headers: Vec<(String, String)>, +} + +struct ValidationProbe { + path: &'static str, + body: serde_json::Value, } fn resolve_provider_route(provider: &Provider) -> Result { @@ -241,9 +253,198 @@ fn resolve_provider_route(provider: &Provider) -> Result Result { + if route + .protocols + .iter() + .any(|protocol| protocol == "openai_chat_completions") + { + return Ok(ValidationProbe { + path: "/v1/chat/completions", + body: serde_json::json!({ + "model": model_id, + "messages": [{"role": "user", "content": "ping"}], + "max_tokens": 1, + }), + }); + } + + if route + .protocols + .iter() + .any(|protocol| protocol == "anthropic_messages") + { + return Ok(ValidationProbe { + path: "/v1/messages", + body: serde_json::json!({ + "model": model_id, + "messages": [{"role": "user", "content": "ping"}], + "max_tokens": 1, + }), + }); + } + + if route + .protocols + .iter() + .any(|protocol| protocol == "openai_responses") + { + return Ok(ValidationProbe { + path: "/v1/responses", + body: serde_json::json!({ + "model": model_id, + "input": "ping", + "max_output_tokens": 1, + }), + }); + } + + if route + .protocols + .iter() + .any(|protocol| protocol == "openai_completions") + { + return Ok(ValidationProbe { + path: "/v1/completions", + body: serde_json::json!({ + "model": model_id, + "prompt": "ping", + "max_tokens": 1, + }), + }); + } + + Err(Status::failed_precondition(format!( + "provider type '{}' does not expose a writable inference protocol for validation", + route.provider_type + ))) +} + +fn validation_url(base_url: &str, path: &str) -> String { + let base = base_url.trim_end_matches('/'); + if base.ends_with("/v1") && (path == "/v1" || path.starts_with("/v1/")) { + return format!("{base}{}", &path[3..]); + } + + format!("{base}{path}") +} + +fn validation_failure( + provider_name: &str, + model_id: &str, + base_url: &str, + details: &str, + next_steps: &str, +) -> Status { + Status::failed_precondition(format!( + "failed to verify inference endpoint for provider '{provider_name}' and model '{model_id}' at '{base_url}': {details}. Next steps: {next_steps}, or retry with '--no-verify' if the endpoint is not up yet" + )) +} + +async fn verify_provider_endpoint( + provider_name: &str, + model_id: &str, + route: &ResolvedProviderRoute, +) -> Result<(), Status> { + let probe = validation_probe(route, model_id)?; + let url = validation_url(&route.base_url, probe.path); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(5)) + .build() + .map_err(|err| Status::internal(format!("build validation client failed: {err}")))?; + + let mut request = client + .post(&url) + .header(reqwest::header::CONTENT_TYPE, "application/json"); + request = match &route.auth { + AuthHeader::Bearer => request.bearer_auth(&route.api_key), + AuthHeader::Custom(name) => request.header(*name, &route.api_key), + }; + + for (name, value) in &route.default_headers { + request = request.header(name, value); + } + + let response = request.json(&probe.body).send().await.map_err(|err| { + let details = if err.is_timeout() { + format!("request to {url} timed out") + } else if err.is_connect() { + format!("failed to connect to {url}: {err}") + } else { + format!("request to {url} failed: {err}") + }; + + validation_failure( + provider_name, + model_id, + &route.base_url, + &details, + "check that the service is running, confirm the base URL and protocol, and verify credentials", + ) + })?; + + if response.status().is_success() { + return Ok(()); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + let body = body.trim(); + let body_suffix = if body.is_empty() { + String::new() + } else { + format!( + " Response body: {}", + body.chars().take(200).collect::() + ) + }; + + let (details, next_steps) = match status.as_u16() { + 400 | 404 | 405 | 422 => ( + format!("upstream rejected the validation request with HTTP {status}.{body_suffix}"), + "confirm the provider type, base URL, and model identifier", + ), + 401 | 403 => ( + format!("upstream rejected credentials with HTTP {status}.{body_suffix}"), + "verify the provider API key and any required auth headers", + ), + 429 => ( + format!( + "upstream rate-limited the validation request with HTTP {status}.{body_suffix}" + ), + "retry later or verify quota/limits on the upstream provider", + ), + 500..=599 => ( + format!("upstream returned HTTP {status}.{body_suffix}"), + "check whether the endpoint is healthy and serving requests", + ), + _ => ( + format!("upstream returned unexpected HTTP {status}.{body_suffix}"), + "confirm the endpoint URL, protocol, credentials, and model identifier", + ), + }; + + Err(validation_failure( + provider_name, + model_id, + &route.base_url, + &details, + next_steps, + )) +} + fn find_provider_api_key(provider: &Provider, preferred_key_names: &[&str]) -> Option { for key in preferred_key_names { if let Some(value) = provider.credentials.get(*key) @@ -369,6 +570,8 @@ async fn resolve_route_by_name( #[cfg(test)] mod tests { use super::*; + use wiremock::matchers::{body_partial_json, header, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; fn make_route(name: &str, provider_name: &str, model_id: &str) -> InferenceRoute { InferenceRoute { @@ -392,6 +595,20 @@ mod tests { } } + fn make_provider_with_base_url( + name: &str, + provider_type: &str, + key_name: &str, + key_value: &str, + base_url_key: &str, + base_url: &str, + ) -> Provider { + Provider { + config: std::iter::once((base_url_key.to_string(), base_url.to_string())).collect(), + ..make_provider(name, provider_type, key_name, key_value) + } + } + #[tokio::test] async fn upsert_cluster_route_creates_and_increments_version() { let store = Store::connect("sqlite::memory:?cache=shared") @@ -409,6 +626,7 @@ mod tests { CLUSTER_INFERENCE_ROUTE_NAME, "openai-dev", "gpt-4o", + true, ) .await .expect("first set should succeed"); @@ -420,6 +638,7 @@ mod tests { CLUSTER_INFERENCE_ROUTE_NAME, "openai-dev", "gpt-4.1", + true, ) .await .expect("second set should succeed"); @@ -630,6 +849,7 @@ mod tests { SANDBOX_SYSTEM_ROUTE_NAME, "anthropic-dev", "claude-sonnet-4-20250514", + true, ) .await .expect("should succeed"); @@ -715,6 +935,7 @@ mod tests { SANDBOX_SYSTEM_ROUTE_NAME, "openai-dev", "gpt-4o-mini", + true, ) .await .expect("upsert should succeed"); @@ -730,6 +951,175 @@ mod tests { assert_eq!(config.model_id, "gpt-4o-mini"); } + #[test] + fn openai_validation_probe_uses_lightweight_chat_shape() { + let route = ResolvedProviderRoute { + provider_type: "openai".to_string(), + base_url: "https://api.openai.com/v1".to_string(), + protocols: vec!["openai_chat_completions".to_string()], + api_key: "sk-test".to_string(), + auth: AuthHeader::Bearer, + default_headers: Vec::new(), + }; + + let probe = validation_probe(&route, "gpt-4.1").expect("probe should build"); + + assert_eq!(probe.path, "/v1/chat/completions"); + assert_eq!(probe.body["model"], "gpt-4.1"); + assert_eq!(probe.body["max_tokens"], 1); + } + + #[test] + fn anthropic_validation_probe_uses_messages_shape() { + let route = ResolvedProviderRoute { + provider_type: "anthropic".to_string(), + base_url: "https://api.anthropic.com/v1".to_string(), + protocols: vec!["anthropic_messages".to_string()], + api_key: "sk-test".to_string(), + auth: AuthHeader::Custom("x-api-key"), + default_headers: vec![("anthropic-version".to_string(), "2023-06-01".to_string())], + }; + + let probe = + validation_probe(&route, "claude-sonnet-4-20250514").expect("probe should build"); + + assert_eq!(probe.path, "/v1/messages"); + assert_eq!(probe.body["model"], "claude-sonnet-4-20250514"); + assert_eq!(probe.body["max_tokens"], 1); + } + + #[tokio::test] + async fn upsert_cluster_route_verifies_endpoint_by_default() { + let store = Store::connect("sqlite::memory:?cache=shared") + .await + .expect("store"); + let mock_server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(header("authorization", "Bearer sk-test")) + .and(body_partial_json(serde_json::json!({ + "model": "gpt-4o-mini", + "max_tokens": 1, + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "chatcmpl-123", + "object": "chat.completion", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "model": "gpt-4o-mini" + }))) + .mount(&mock_server) + .await; + + let provider = make_provider_with_base_url( + "openai-dev", + "openai", + "OPENAI_API_KEY", + "sk-test", + "OPENAI_BASE_URL", + &mock_server.uri(), + ); + store + .put_message(&provider) + .await + .expect("persist provider"); + + let route = upsert_cluster_inference_route( + &store, + CLUSTER_INFERENCE_ROUTE_NAME, + "openai-dev", + "gpt-4o-mini", + false, + ) + .await + .expect("validation should succeed"); + + assert_eq!(route.version, 1); + } + + #[tokio::test] + async fn upsert_cluster_route_rejects_failed_validation() { + let store = Store::connect("sqlite::memory:?cache=shared") + .await + .expect("store"); + let mock_server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(401).set_body_string("bad key")) + .mount(&mock_server) + .await; + + let provider = make_provider_with_base_url( + "openai-dev", + "openai", + "OPENAI_API_KEY", + "sk-test", + "OPENAI_BASE_URL", + &mock_server.uri(), + ); + store + .put_message(&provider) + .await + .expect("persist provider"); + + let err = upsert_cluster_inference_route( + &store, + CLUSTER_INFERENCE_ROUTE_NAME, + "openai-dev", + "gpt-4o-mini", + false, + ) + .await + .expect_err("validation should fail"); + + assert_eq!(err.code(), tonic::Code::FailedPrecondition); + assert!( + err.message() + .contains("failed to verify inference endpoint") + ); + assert!(err.message().contains("verify the provider API key")); + assert!(err.message().contains("--no-verify")); + + let persisted = store + .get_message_by_name::(CLUSTER_INFERENCE_ROUTE_NAME) + .await + .expect("fetch route") + .is_none(); + assert!(persisted, "route should not persist on failed validation"); + } + + #[tokio::test] + async fn upsert_cluster_route_skips_validation_when_requested() { + let store = Store::connect("sqlite::memory:?cache=shared") + .await + .expect("store"); + let provider = make_provider_with_base_url( + "openai-dev", + "openai", + "OPENAI_API_KEY", + "sk-test", + "OPENAI_BASE_URL", + "http://127.0.0.1:9", + ); + store + .put_message(&provider) + .await + .expect("persist provider"); + + let route = upsert_cluster_inference_route( + &store, + CLUSTER_INFERENCE_ROUTE_NAME, + "openai-dev", + "gpt-4o-mini", + true, + ) + .await + .expect("skip validation should persist route"); + + assert_eq!(route.version, 1); + } + #[test] fn effective_route_name_defaults_empty_to_inference_local() { assert_eq!( diff --git a/docs/inference/configure.md b/docs/inference/configure.md index b4dcd7817..9d64b6221 100644 --- a/docs/inference/configure.md +++ b/docs/inference/configure.md @@ -135,9 +135,9 @@ Use this endpoint when inference should stay local to the host for privacy and s ### Verify the Endpoint from a Sandbox -`openshell inference get` confirms the configuration was saved, but does not verify the upstream endpoint is reachable. The CLI also accepts `--no-verify` on `openshell inference set` and `openshell inference update` so automation can opt out explicitly ahead of a future verify-by-default rollout. +`openshell inference set` and `openshell inference update` verify the upstream endpoint by default before saving the configuration. If the endpoint is not live yet, try again with `--no-verify` to persist the route without the probe. -To confirm end-to-end connectivity, connect to a sandbox and run: +`openshell inference get` confirms the current saved configuration. To confirm end-to-end connectivity from a sandbox, run: ```bash curl https://inference.local/v1/responses \ diff --git a/proto/inference.proto b/proto/inference.proto index 11670c4a9..828927ecf 100644 --- a/proto/inference.proto +++ b/proto/inference.proto @@ -56,6 +56,8 @@ message SetClusterInferenceRequest { // Route name to target. Empty string defaults to "inference.local" (user-facing). // Use "sandbox-system" for the sandbox system-level inference route. string route_name = 3; + // Skip synchronous endpoint validation before persistence. + bool skip_validation = 4; } message SetClusterInferenceResponse { diff --git a/python/openshell/sandbox.py b/python/openshell/sandbox.py index 7b48ab3b9..3399c2032 100644 --- a/python/openshell/sandbox.py +++ b/python/openshell/sandbox.py @@ -398,11 +398,13 @@ class InferenceRouteClient: *, provider_name: str, model_id: str, + no_verify: bool = False, ) -> ClusterInferenceConfig: response = self._stub.SetClusterInference( inference_pb2.SetClusterInferenceRequest( provider_name=provider_name, model_id=model_id, + skip_validation=no_verify, ), timeout=self._timeout, ) diff --git a/python/openshell/sandbox_test.py b/python/openshell/sandbox_test.py index 4c0eebcda..441e69f49 100644 --- a/python/openshell/sandbox_test.py +++ b/python/openshell/sandbox_test.py @@ -10,6 +10,7 @@ from openshell._proto import openshell_pb2 from openshell.sandbox import ( _PYTHON_CLOUDPICKLE_BOOTSTRAP, _SANDBOX_PYTHON_BIN, + InferenceRouteClient, SandboxClient, ) @@ -33,6 +34,22 @@ class _FakeStub: ) +class _FakeInferenceStub: + def __init__(self) -> None: + self.request = None + + def SetClusterInference(self, request: Any, timeout: float | None = None) -> Any: + self.request = request + _ = timeout + + class _Response: + provider_name = request.provider_name + model_id = request.model_id + version = 1 + + return _Response() + + def _client_with_fake_stub(stub: _FakeStub) -> SandboxClient: client = cast("SandboxClient", object.__new__(SandboxClient)) client._timeout = 30.0 @@ -120,3 +137,20 @@ def test_from_active_cluster_prefers_openshell_gateway_env( assert client._cluster_name == gateway_name finally: client.close() + + +def test_inference_set_cluster_forwards_no_verify_flag() -> None: + stub = _FakeInferenceStub() + client = cast("InferenceRouteClient", object.__new__(InferenceRouteClient)) + client._timeout = 30.0 + client._stub = cast("Any", stub) + + result = client.set_cluster( + provider_name="openai-dev", + model_id="gpt-4.1", + no_verify=True, + ) + + assert result.provider_name == "openai-dev" + assert stub.request is not None + assert stub.request.skip_validation is True