// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 mod helpers; use helpers::{ EnvVarGuard, build_ca, build_client_cert, build_server_cert, install_rustls_provider, }; use openshell_cli::run; use openshell_cli::tls::TlsOptions; use openshell_core::proto::open_shell_server::{OpenShell, OpenShellServer}; use openshell_core::proto::{ AttachSandboxProviderRequest, AttachSandboxProviderResponse, CreateProviderRequest, CreateSandboxRequest, CreateSshSessionRequest, CreateSshSessionResponse, DeleteProviderRefreshRequest, DeleteProviderRefreshResponse, DeleteProviderRequest, DeleteProviderResponse, DeleteSandboxRequest, DeleteSandboxResponse, DetachSandboxProviderRequest, DetachSandboxProviderResponse, ExecSandboxEvent, ExecSandboxInput, ExecSandboxRequest, GatewayMessage, GetGatewayConfigRequest, GetGatewayConfigResponse, GetProviderRefreshStatusRequest, GetProviderRefreshStatusResponse, GetProviderRequest, GetSandboxConfigRequest, GetSandboxConfigResponse, GetSandboxProviderEnvironmentRequest, GetSandboxProviderEnvironmentResponse, GetSandboxRequest, HealthRequest, HealthResponse, ListProvidersRequest, ListProvidersResponse, ListSandboxProvidersRequest, ListSandboxProvidersResponse, ListSandboxesRequest, ListSandboxesResponse, Provider, ProviderCredentialRefresh, ProviderCredentialRefreshStatus, ProviderCredentialRefreshStrategy, ProviderProfile, ProviderProfileCredential, ProviderProfileDiscovery, ProviderResponse, RevokeSshSessionRequest, RevokeSshSessionResponse, RotateProviderCredentialRequest, RotateProviderCredentialResponse, Sandbox, SandboxResponse, SandboxStreamEvent, ServiceStatus, SettingValue, SupervisorMessage, UpdateProviderRequest, WatchSandboxRequest, setting_value, }; use openshell_core::{ObjectId, ObjectName}; use std::collections::HashMap; use std::sync::Arc; use tempfile::TempDir; use tokio::net::TcpListener; use tokio::sync::{Mutex, mpsc}; use tokio_stream::wrappers::TcpListenerStream; use tonic::transport::{Certificate as TlsCertificate, Identity, Server, ServerTlsConfig}; use tonic::{Response, Status}; #[derive(Clone, Default)] struct ProviderState { providers: Arc>>, profiles: Arc>>, refresh_statuses: Arc>>, refresh_requests: Arc>>, sandbox_providers: Arc>>>, sandbox_provider_requests: Arc>>, global_settings: Arc>>, } #[derive(Clone, Debug, PartialEq, Eq)] enum ProviderRefreshRequestLog { Status { provider_name: String, credential_key: String, }, Configure { provider_name: String, credential_key: String, expires_at_ms: Option, }, Rotate { provider_name: String, credential_key: String, }, Delete { provider_name: String, credential_key: String, }, } #[derive(Clone, Debug, PartialEq, Eq)] enum SandboxProviderRequestLog { List { sandbox_name: String, }, Attach { sandbox_name: String, provider_name: String, }, Detach { sandbox_name: String, provider_name: String, }, } #[derive(Clone, Default)] struct TestOpenShell { state: ProviderState, } #[tonic::async_trait] impl OpenShell for TestOpenShell { async fn health( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(HealthResponse { status: ServiceStatus::Healthy.into(), version: "test".to_string(), })) } async fn create_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(SandboxResponse::default())) } async fn get_sandbox( &self, request: tonic::Request, ) -> Result, Status> { let name = request.into_inner().name; // Return a minimal sandbox with metadata for CAS operations Ok(Response::new(SandboxResponse { sandbox: Some(Sandbox { metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { id: format!("sb-{name}"), name, created_at_ms: 0, labels: HashMap::new(), resource_version: 1, }), spec: None, status: None, }), })) } async fn list_sandboxes( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(ListSandboxesResponse::default())) } async fn list_sandbox_providers( &self, request: tonic::Request, ) -> Result, Status> { let sandbox_name = request.into_inner().sandbox_name; self.state .sandbox_provider_requests .lock() .await .push(SandboxProviderRequestLog::List { sandbox_name: sandbox_name.clone(), }); let provider_names = self .state .sandbox_providers .lock() .await .get(&sandbox_name) .cloned() .unwrap_or_default(); let providers_by_name = self.state.providers.lock().await; let providers = provider_names .iter() .filter_map(|name| providers_by_name.get(name).cloned()) .collect(); Ok(Response::new(ListSandboxProvidersResponse { providers })) } async fn attach_sandbox_provider( &self, request: tonic::Request, ) -> Result, Status> { let request = request.into_inner(); self.state .sandbox_provider_requests .lock() .await .push(SandboxProviderRequestLog::Attach { sandbox_name: request.sandbox_name.clone(), provider_name: request.provider_name.clone(), }); if !self .state .providers .lock() .await .contains_key(&request.provider_name) { return Err(Status::failed_precondition("provider not found")); } let mut sandbox_providers = self.state.sandbox_providers.lock().await; let providers = sandbox_providers .entry(request.sandbox_name.clone()) .or_default(); let attached = if providers.contains(&request.provider_name) { false } else { providers.push(request.provider_name.clone()); true }; let sandbox = Sandbox { metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { name: request.sandbox_name, ..Default::default() }), spec: Some(openshell_core::proto::SandboxSpec { providers: providers.clone(), ..Default::default() }), ..Default::default() }; Ok(Response::new(AttachSandboxProviderResponse { sandbox: Some(sandbox), attached, })) } async fn detach_sandbox_provider( &self, request: tonic::Request, ) -> Result, Status> { let request = request.into_inner(); self.state .sandbox_provider_requests .lock() .await .push(SandboxProviderRequestLog::Detach { sandbox_name: request.sandbox_name.clone(), provider_name: request.provider_name.clone(), }); let mut sandbox_providers = self.state.sandbox_providers.lock().await; let providers = sandbox_providers .entry(request.sandbox_name.clone()) .or_default(); let before_len = providers.len(); providers.retain(|name| name != &request.provider_name); let detached = providers.len() != before_len; let sandbox = Sandbox { metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { name: request.sandbox_name, ..Default::default() }), spec: Some(openshell_core::proto::SandboxSpec { providers: providers.clone(), ..Default::default() }), ..Default::default() }; Ok(Response::new(DetachSandboxProviderResponse { sandbox: Some(sandbox), detached, })) } async fn delete_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(DeleteSandboxResponse { deleted: true })) } async fn get_sandbox_config( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(GetSandboxConfigResponse::default())) } async fn get_gateway_config( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(GetGatewayConfigResponse { settings: self.state.global_settings.lock().await.clone(), settings_revision: 1, })) } async fn get_sandbox_provider_environment( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new( GetSandboxProviderEnvironmentResponse::default(), )) } async fn create_ssh_session( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(CreateSshSessionResponse::default())) } async fn expose_service( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new( openshell_core::proto::ServiceEndpointResponse::default(), )) } async fn get_service( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn list_services( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn delete_service( &self, _: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("unused")) } async fn revoke_ssh_session( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new(RevokeSshSessionResponse::default())) } async fn create_provider( &self, request: tonic::Request, ) -> Result, Status> { let mut provider = request .into_inner() .provider .ok_or_else(|| Status::invalid_argument("provider is required"))?; let mut providers = self.state.providers.lock().await; let provider_name = provider.object_name().to_string(); if providers.contains_key(&provider_name) { return Err(Status::already_exists("provider already exists")); } if provider.object_id().is_empty() && let Some(metadata) = &mut provider.metadata { metadata.id = format!("id-{provider_name}"); } providers.insert(provider_name, provider.clone()); Ok(Response::new(ProviderResponse { provider: Some(provider), })) } async fn get_provider( &self, request: tonic::Request, ) -> Result, Status> { let name = request.into_inner().name; let providers = self.state.providers.lock().await; let provider = providers .get(&name) .cloned() .ok_or_else(|| Status::not_found("provider not found"))?; Ok(Response::new(ProviderResponse { provider: Some(provider), })) } async fn list_providers( &self, _request: tonic::Request, ) -> Result, Status> { let providers = self .state .providers .lock() .await .values() .cloned() .collect::>(); Ok(Response::new(ListProvidersResponse { providers })) } async fn list_provider_profiles( &self, _request: tonic::Request, ) -> Result, Status> { let mut profiles = openshell_providers::default_profiles() .iter() .map(openshell_providers::ProviderTypeProfile::to_proto) .collect::>(); profiles.extend(self.state.profiles.lock().await.values().cloned()); Ok(Response::new( openshell_core::proto::ListProviderProfilesResponse { profiles }, )) } async fn get_provider_profile( &self, request: tonic::Request, ) -> Result, Status> { let id = request.into_inner().id; let profile = if let Some(profile) = openshell_providers::get_default_profile(&id) { profile.to_proto() } else { self.state .profiles .lock() .await .get(&id) .cloned() .ok_or_else(|| Status::not_found("provider profile not found"))? }; Ok(Response::new( openshell_core::proto::ProviderProfileResponse { profile: Some(profile), }, )) } async fn import_provider_profiles( &self, request: tonic::Request, ) -> Result, Status> { let mut profiles = self.state.profiles.lock().await; let imported = request .into_inner() .profiles .into_iter() .filter_map(|item| item.profile) .inspect(|profile| { profiles.insert(profile.id.clone(), profile.clone()); }) .collect::>(); Ok(Response::new( openshell_core::proto::ImportProviderProfilesResponse { diagnostics: Vec::new(), profiles: imported, imported: true, }, )) } async fn lint_provider_profiles( &self, _request: tonic::Request, ) -> Result, Status> { Ok(Response::new( openshell_core::proto::LintProviderProfilesResponse { diagnostics: Vec::new(), valid: true, }, )) } async fn update_provider( &self, request: tonic::Request, ) -> Result, Status> { let provider = request .into_inner() .provider .ok_or_else(|| Status::invalid_argument("provider is required"))?; let mut providers = self.state.providers.lock().await; let existing = providers .get(provider.object_name()) .cloned() .ok_or_else(|| Status::not_found("provider not found"))?; // Merge semantics: empty map = no change, empty value = delete key. let merge = |mut base: HashMap, incoming: HashMap| -> HashMap { if incoming.is_empty() { return base; } for (k, v) in incoming { if v.is_empty() { base.remove(&k); } else { base.insert(k, v); } } base }; let merge_expiry = |mut base: HashMap, incoming: HashMap| { if incoming.is_empty() { return base; } for (k, v) in incoming { if v <= 0 { base.remove(&k); } else { base.insert(k, v); } } base }; let existing_metadata = existing.metadata.clone().unwrap_or_default(); let provider_metadata = provider.metadata.clone().unwrap_or_default(); let updated = Provider { metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { id: existing_metadata.id, name: provider_metadata.name, created_at_ms: existing_metadata.created_at_ms, labels: existing_metadata.labels, resource_version: 0, }), r#type: existing.r#type, credentials: merge(existing.credentials, provider.credentials), config: merge(existing.config, provider.config), credential_expires_at_ms: merge_expiry( existing.credential_expires_at_ms, provider.credential_expires_at_ms, ), }; let updated_name = updated.object_name().to_string(); providers.insert(updated_name, updated.clone()); Ok(Response::new(ProviderResponse { provider: Some(updated), })) } async fn get_provider_refresh_status( &self, request: tonic::Request, ) -> Result, Status> { let request = request.into_inner(); self.state .refresh_requests .lock() .await .push(ProviderRefreshRequestLog::Status { provider_name: request.provider.clone(), credential_key: request.credential_key.clone(), }); let refresh_statuses = self.state.refresh_statuses.lock().await; let credentials = if request.credential_key.is_empty() { refresh_statuses .values() .filter(|status| status.provider_name == request.provider) .cloned() .collect() } else { refresh_statuses .get(&(request.provider, request.credential_key)) .cloned() .into_iter() .collect() }; Ok(Response::new(GetProviderRefreshStatusResponse { credentials, })) } async fn configure_provider_refresh( &self, request: tonic::Request, ) -> Result, Status> { let request = request.into_inner(); self.state .refresh_requests .lock() .await .push(ProviderRefreshRequestLog::Configure { provider_name: request.provider.clone(), credential_key: request.credential_key.clone(), expires_at_ms: request.expires_at_ms, }); let providers = self.state.providers.lock().await; let provider = providers .get(&request.provider) .ok_or_else(|| Status::not_found("provider not found"))?; let status = ProviderCredentialRefreshStatus { provider_name: request.provider.clone(), provider_id: provider.object_id().to_string(), credential_key: request.credential_key.clone(), strategy: request.strategy, status: "configured".to_string(), expires_at_ms: request.expires_at_ms.unwrap_or_default(), next_refresh_at_ms: 0, last_refresh_at_ms: 0, last_error: String::new(), }; drop(providers); self.state .refresh_statuses .lock() .await .insert((request.provider, request.credential_key), status.clone()); Ok(Response::new( openshell_core::proto::ConfigureProviderRefreshResponse { status: Some(status), }, )) } async fn rotate_provider_credential( &self, request: tonic::Request, ) -> Result, Status> { let request = request.into_inner(); self.state .refresh_requests .lock() .await .push(ProviderRefreshRequestLog::Rotate { provider_name: request.provider.clone(), credential_key: request.credential_key.clone(), }); let mut refresh_statuses = self.state.refresh_statuses.lock().await; let status = refresh_statuses .get_mut(&(request.provider, request.credential_key)) .ok_or_else(|| Status::not_found("provider refresh state not found"))?; status.status = "rotation_requested".to_string(); Ok(Response::new(RotateProviderCredentialResponse { status: Some(status.clone()), })) } async fn delete_provider_refresh( &self, request: tonic::Request, ) -> Result, Status> { let request = request.into_inner(); self.state .refresh_requests .lock() .await .push(ProviderRefreshRequestLog::Delete { provider_name: request.provider.clone(), credential_key: request.credential_key.clone(), }); let deleted = self .state .refresh_statuses .lock() .await .remove(&(request.provider, request.credential_key)) .is_some(); Ok(Response::new(DeleteProviderRefreshResponse { deleted })) } async fn delete_provider( &self, request: tonic::Request, ) -> Result, Status> { let name = request.into_inner().name; let deleted = self.state.providers.lock().await.remove(&name).is_some(); Ok(Response::new(DeleteProviderResponse { deleted })) } async fn delete_provider_profile( &self, request: tonic::Request, ) -> Result, Status> { let id = request.into_inner().id; let deleted = self.state.profiles.lock().await.remove(&id).is_some(); Ok(Response::new( openshell_core::proto::DeleteProviderProfileResponse { deleted }, )) } type WatchSandboxStream = tokio_stream::wrappers::ReceiverStream>; type ExecSandboxStream = tokio_stream::wrappers::ReceiverStream>; type ConnectSupervisorStream = tokio_stream::wrappers::ReceiverStream>; async fn watch_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { let (_tx, rx) = mpsc::channel(1); Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new( rx, ))) } async fn exec_sandbox( &self, _request: tonic::Request, ) -> Result, Status> { let (_tx, rx) = mpsc::channel(1); Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new( rx, ))) } type ExecSandboxInteractiveStream = tokio_stream::wrappers::ReceiverStream>; async fn exec_sandbox_interactive( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn update_config( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_sandbox_policy_status( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn list_sandbox_policies( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn report_policy_status( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_sandbox_logs( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn push_sandbox_logs( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn submit_policy_analysis( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_draft_policy( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn approve_draft_chunk( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn reject_draft_chunk( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn approve_all_draft_chunks( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn edit_draft_chunk( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn undo_draft_chunk( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn clear_draft_chunks( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn get_draft_history( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn issue_sandbox_token( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn refresh_sandbox_token( &self, _request: tonic::Request, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } async fn connect_supervisor( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } type RelayStreamStream = tokio_stream::wrappers::ReceiverStream>; async fn relay_stream( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } type ForwardTcpStream = tokio_stream::wrappers::ReceiverStream< Result, >; async fn forward_tcp( &self, _request: tonic::Request>, ) -> Result, Status> { Err(Status::unimplemented("not implemented in test")) } } /// Test fixture: TLS-enabled server with matching client certs. struct TestServer { endpoint: String, tls: TlsOptions, state: ProviderState, _dir: TempDir, } async fn run_server() -> TestServer { install_rustls_provider(); let (ca, ca_key) = build_ca(); let (server_cert, server_key) = build_server_cert(&ca, &ca_key); let (client_cert, client_key) = build_client_cert(&ca, &ca_key); let ca_cert = ca.pem(); let identity = Identity::from_pem(server_cert, server_key); let client_ca = TlsCertificate::from_pem(ca_cert.clone()); let tls_config = ServerTlsConfig::new() .identity(identity) .client_ca_root(client_ca); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let incoming = TcpListenerStream::new(listener); let state = ProviderState::default(); let service = TestOpenShell { state: state.clone(), }; tokio::spawn(async move { Server::builder() .tls_config(tls_config) .unwrap() .add_service(OpenShellServer::new(service)) .serve_with_incoming(incoming) .await .unwrap(); }); let dir = tempfile::tempdir().unwrap(); let ca_path = dir.path().join("ca.crt"); let cert_path = dir.path().join("tls.crt"); let key_path = dir.path().join("tls.key"); std::fs::write(&ca_path, ca_cert).unwrap(); std::fs::write(&cert_path, client_cert).unwrap(); std::fs::write(&key_path, client_key).unwrap(); let tls = TlsOptions::new(Some(ca_path), Some(cert_path), Some(key_path)); let endpoint = format!("https://localhost:{}", addr.port()); TestServer { endpoint, tls, state, _dir: dir, } } async fn enable_providers_v2(ts: &TestServer) { ts.state.global_settings.lock().await.insert( openshell_core::settings::PROVIDERS_V2_ENABLED_KEY.to_string(), SettingValue { value: Some(setting_value::Value::BoolValue(true)), }, ); } #[tokio::test] async fn provider_cli_run_functions_support_full_crud_flow() { let ts = run_server().await; run::provider_create( &ts.endpoint, "my-claude", "claude", false, &["API_KEY=abc".to_string()], &["profile=dev".to_string()], &ts.tls, ) .await .expect("provider create"); run::provider_get(&ts.endpoint, "my-claude", &ts.tls) .await .expect("provider get"); run::provider_list(&ts.endpoint, 100, 0, false, &ts.tls) .await .expect("provider list"); run::provider_update( &ts.endpoint, "my-claude", false, &["API_KEY=rotated".to_string()], &["profile=prod".to_string()], &[], &ts.tls, ) .await .expect("provider update"); run::provider_delete(&ts.endpoint, &["my-claude".to_string()], &ts.tls) .await .expect("provider delete"); } #[tokio::test] async fn provider_list_profiles_cli_uses_profile_browsing_rpc() { let ts = run_server().await; run::provider_list_profiles(&ts.endpoint, "table", &ts.tls) .await .expect("provider list-profiles"); } #[tokio::test] async fn provider_refresh_cli_run_functions_wire_requests() { let ts = run_server().await; run::provider_create( &ts.endpoint, "my-graph", "outlook", false, &["MS_GRAPH_ACCESS_TOKEN=token".to_string()], &[], &ts.tls, ) .await .expect("provider create"); run::provider_refresh_config( &ts.endpoint, run::ProviderRefreshConfigInput { name: "my-graph", credential_key: "MS_GRAPH_ACCESS_TOKEN", strategy: "oauth2_client_credentials", material: &["tenant_id=tenant".to_string()], secret_material_keys: &["client_secret".to_string()], credential_expires_at_ms: Some(1_767_225_600_000), }, &ts.tls, ) .await .expect("provider refresh configure"); run::provider_refresh_status( &ts.endpoint, "my-graph", Some("MS_GRAPH_ACCESS_TOKEN"), &ts.tls, ) .await .expect("provider refresh status"); run::provider_rotate(&ts.endpoint, "my-graph", "MS_GRAPH_ACCESS_TOKEN", &ts.tls) .await .expect("provider refresh rotate"); run::provider_refresh_delete(&ts.endpoint, "my-graph", "MS_GRAPH_ACCESS_TOKEN", &ts.tls) .await .expect("provider refresh delete"); let requests = ts.state.refresh_requests.lock().await.clone(); assert_eq!( requests, vec![ ProviderRefreshRequestLog::Configure { provider_name: "my-graph".to_string(), credential_key: "MS_GRAPH_ACCESS_TOKEN".to_string(), expires_at_ms: Some(1_767_225_600_000), }, ProviderRefreshRequestLog::Status { provider_name: "my-graph".to_string(), credential_key: "MS_GRAPH_ACCESS_TOKEN".to_string(), }, ProviderRefreshRequestLog::Rotate { provider_name: "my-graph".to_string(), credential_key: "MS_GRAPH_ACCESS_TOKEN".to_string(), }, ProviderRefreshRequestLog::Delete { provider_name: "my-graph".to_string(), credential_key: "MS_GRAPH_ACCESS_TOKEN".to_string(), }, ] ); } #[tokio::test] async fn provider_create_allows_empty_credentials_for_gateway_refresh_profiles() { let ts = run_server().await; ts.state.profiles.lock().await.insert( "custom-refresh".to_string(), ProviderProfile { id: "custom-refresh".to_string(), display_name: "Custom Refresh".to_string(), credentials: vec![ProviderProfileCredential { name: "ACCESS_TOKEN".to_string(), required: true, refresh: Some(ProviderCredentialRefresh { strategy: ProviderCredentialRefreshStrategy::Oauth2RefreshToken as i32, ..Default::default() }), ..Default::default() }], ..Default::default() }, ); run::provider_create( &ts.endpoint, "custom-refresh-provider", "custom-refresh", false, &[], &[], &ts.tls, ) .await .expect("provider create"); let stored = ts.state.providers.lock().await; let provider = stored.get("custom-refresh-provider").expect("provider"); assert_eq!(provider.r#type, "custom-refresh"); assert!(provider.credentials.is_empty()); } #[tokio::test] async fn sandbox_provider_cli_run_functions_wire_requests_and_idempotent_results() { let ts = run_server().await; run::provider_create( &ts.endpoint, "work-github", "github", false, &["GITHUB_TOKEN=ghp-test".to_string()], &[], &ts.tls, ) .await .expect("provider create"); run::sandbox_provider_attach(&ts.endpoint, "dev-sandbox", "work-github", &ts.tls) .await .expect("sandbox provider attach"); run::sandbox_provider_attach(&ts.endpoint, "dev-sandbox", "work-github", &ts.tls) .await .expect("sandbox provider attach is idempotent"); run::sandbox_provider_list(&ts.endpoint, "dev-sandbox", &ts.tls) .await .expect("sandbox provider list"); run::sandbox_provider_detach(&ts.endpoint, "dev-sandbox", "work-github", &ts.tls) .await .expect("sandbox provider detach"); run::sandbox_provider_detach(&ts.endpoint, "dev-sandbox", "work-github", &ts.tls) .await .expect("sandbox provider detach is idempotent"); let requests = ts.state.sandbox_provider_requests.lock().await.clone(); assert_eq!( requests, vec![ SandboxProviderRequestLog::Attach { sandbox_name: "dev-sandbox".to_string(), provider_name: "work-github".to_string(), }, SandboxProviderRequestLog::Attach { sandbox_name: "dev-sandbox".to_string(), provider_name: "work-github".to_string(), }, SandboxProviderRequestLog::List { sandbox_name: "dev-sandbox".to_string(), }, SandboxProviderRequestLog::Detach { sandbox_name: "dev-sandbox".to_string(), provider_name: "work-github".to_string(), }, SandboxProviderRequestLog::Detach { sandbox_name: "dev-sandbox".to_string(), provider_name: "work-github".to_string(), }, ] ); let providers = ts.state.sandbox_providers.lock().await; assert!(providers.get("dev-sandbox").is_none_or(Vec::is_empty)); } #[tokio::test] async fn sandbox_provider_attach_cli_surfaces_server_errors() { let ts = run_server().await; let err = run::sandbox_provider_attach(&ts.endpoint, "dev-sandbox", "missing-provider", &ts.tls) .await .expect_err("missing provider should fail"); assert!( err.to_string().contains("provider not found"), "unexpected error: {err}" ); assert_eq!( ts.state.sandbox_provider_requests.lock().await.as_slice(), [SandboxProviderRequestLog::Attach { sandbox_name: "dev-sandbox".to_string(), provider_name: "missing-provider".to_string(), }] ); } #[tokio::test] async fn provider_profile_cli_run_functions_support_custom_profiles() { let ts = run_server().await; let dir = tempfile::tempdir().unwrap(); let profile_path = dir.path().join("custom-api.yaml"); std::fs::write( &profile_path, r" id: custom-api display_name: Custom API category: other credentials: - name: api_key env_vars: [CUSTOM_API_KEY] auth_style: bearer header_name: authorization discovery: credentials: [api_key] endpoints: - host: api.custom.example port: 443 binaries: [/usr/bin/custom] ", ) .unwrap(); run::provider_profile_lint(&ts.endpoint, Some(&profile_path), None, &ts.tls) .await .expect("profile lint"); run::provider_profile_import(&ts.endpoint, Some(&profile_path), None, &ts.tls) .await .expect("profile import"); run::provider_profile_export(&ts.endpoint, "custom-api", "yaml", &ts.tls) .await .expect("profile export"); run::provider_list_profiles(&ts.endpoint, "json", &ts.tls) .await .expect("provider list-profiles json"); run::provider_create( &ts.endpoint, "custom-provider", "custom-api", false, &["CUSTOM_API_KEY=abc".to_string()], &[], &ts.tls, ) .await .expect("custom profile provider create"); let provider = ts .state .providers .lock() .await .get("custom-provider") .cloned() .expect("custom provider should be stored"); assert_eq!(provider.r#type, "custom-api"); run::provider_delete(&ts.endpoint, &["custom-provider".to_string()], &ts.tls) .await .expect("custom provider delete"); run::provider_profile_delete(&ts.endpoint, "custom-api", &ts.tls) .await .expect("profile delete"); } #[tokio::test] async fn provider_create_from_existing_uses_profile_discovery_when_v2_enabled() { let ts = run_server().await; enable_providers_v2(&ts).await; ts.state.profiles.lock().await.insert( "custom-discovery".to_string(), ProviderProfile { id: "custom-discovery".to_string(), display_name: "Custom Discovery".to_string(), credentials: vec![ProviderProfileCredential { name: "api_key".to_string(), env_vars: vec!["CUSTOM_DISCOVERY_API_KEY".to_string()], required: true, ..Default::default() }], discovery: Some(ProviderProfileDiscovery { credentials: vec!["api_key".to_string()], }), ..Default::default() }, ); let _env = EnvVarGuard::set(&[("CUSTOM_DISCOVERY_API_KEY", "profile-secret")]); run::provider_create( &ts.endpoint, "custom-discovered", "custom-discovery", true, &[], &[], &ts.tls, ) .await .expect("profile-backed provider create --from-existing"); let provider = ts .state .providers .lock() .await .get("custom-discovered") .cloned() .expect("custom provider should be stored"); assert_eq!(provider.r#type, "custom-discovery"); assert_eq!( provider.credentials.get("CUSTOM_DISCOVERY_API_KEY"), Some(&"profile-secret".to_string()) ); } #[tokio::test] async fn provider_create_from_existing_uses_registry_discovery_when_v2_disabled() { let ts = run_server().await; let _env = EnvVarGuard::set(&[("OPENAI_API_KEY", "legacy-openai-secret")]); run::provider_create( &ts.endpoint, "legacy-openai", "openai", true, &[], &[], &ts.tls, ) .await .expect("legacy provider create --from-existing"); let provider = ts .state .providers .lock() .await .get("legacy-openai") .cloned() .expect("legacy provider should be stored"); assert_eq!(provider.r#type, "openai"); assert_eq!( provider.credentials.get("OPENAI_API_KEY"), Some(&"legacy-openai-secret".to_string()) ); } #[tokio::test] async fn provider_create_from_existing_requires_profile_when_v2_enabled() { let ts = run_server().await; enable_providers_v2(&ts).await; let _env = EnvVarGuard::set(&[("OPENAI_API_KEY", "legacy-openai-secret")]); let err = run::provider_create(&ts.endpoint, "v2-openai", "openai", true, &[], &[], &ts.tls) .await .expect_err("v2 discovery without a profile should fail"); assert!( err.to_string() .contains("providers v2 discovery requires a provider profile"), "unexpected error: {err}" ); assert!(!ts.state.providers.lock().await.contains_key("v2-openai")); } #[tokio::test] async fn provider_create_from_existing_fails_when_profile_discovery_finds_nothing() { let ts = run_server().await; enable_providers_v2(&ts).await; ts.state.profiles.lock().await.insert( "empty-discovery".to_string(), ProviderProfile { id: "empty-discovery".to_string(), display_name: "Empty Discovery".to_string(), credentials: vec![ProviderProfileCredential { name: "api_key".to_string(), env_vars: vec!["CUSTOM_DISCOVERY_TOKEN_NOT_SET_1460".to_string()], required: false, ..Default::default() }], discovery: Some(ProviderProfileDiscovery { credentials: vec!["api_key".to_string()], }), ..Default::default() }, ); let err = run::provider_create( &ts.endpoint, "empty-discovered", "empty-discovery", true, &[], &[], &ts.tls, ) .await .expect_err("empty profile-backed discovery should fail"); assert!( err.to_string() .contains("no existing local credentials/config found"), "unexpected error: {err}" ); assert!( !ts.state .providers .lock() .await .contains_key("empty-discovered") ); } #[tokio::test] async fn provider_update_from_existing_uses_profile_discovery_when_v2_enabled() { let ts = run_server().await; enable_providers_v2(&ts).await; ts.state.profiles.lock().await.insert( "custom-update-discovery".to_string(), ProviderProfile { id: "custom-update-discovery".to_string(), display_name: "Custom Update Discovery".to_string(), credentials: vec![ProviderProfileCredential { name: "api_key".to_string(), env_vars: vec!["CUSTOM_UPDATE_DISCOVERY_API_KEY".to_string()], required: true, ..Default::default() }], discovery: Some(ProviderProfileDiscovery { credentials: vec!["api_key".to_string()], }), ..Default::default() }, ); ts.state.providers.lock().await.insert( "custom-update".to_string(), Provider { metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { id: "id-custom-update".to_string(), name: "custom-update".to_string(), ..Default::default() }), r#type: "custom-update-discovery".to_string(), credentials: HashMap::new(), config: HashMap::new(), credential_expires_at_ms: HashMap::new(), }, ); let _env = EnvVarGuard::set(&[("CUSTOM_UPDATE_DISCOVERY_API_KEY", "updated-profile-secret")]); run::provider_update(&ts.endpoint, "custom-update", true, &[], &[], &[], &ts.tls) .await .expect("profile-backed provider update --from-existing"); let provider = ts .state .providers .lock() .await .get("custom-update") .cloned() .expect("custom provider should still be stored"); assert_eq!( provider.credentials.get("CUSTOM_UPDATE_DISCOVERY_API_KEY"), Some(&"updated-profile-secret".to_string()) ); } #[tokio::test] async fn provider_profile_import_from_directory_imports_supported_profile_files() { let ts = run_server().await; let dir = tempfile::tempdir().unwrap(); std::fs::write( dir.path().join("custom-yaml.yaml"), r" id: custom-yaml display_name: Custom YAML category: other endpoints: - host: api.yaml.example port: 443 binaries: [/usr/bin/yaml-client] ", ) .unwrap(); std::fs::write( dir.path().join("custom-json.json"), r#"{ "id": "custom-json", "display_name": "Custom JSON", "description": "", "category": "other", "credentials": [], "endpoints": [{"host": "api.json.example", "port": 443}], "binaries": ["/usr/bin/json-client"], "inference_capable": false }"#, ) .unwrap(); std::fs::write(dir.path().join("notes.txt"), "ignored").unwrap(); run::provider_profile_import(&ts.endpoint, None, Some(dir.path()), &ts.tls) .await .expect("profile import --from"); run::provider_profile_export(&ts.endpoint, "custom-yaml", "yaml", &ts.tls) .await .expect("custom-yaml should be imported"); run::provider_profile_export(&ts.endpoint, "custom-json", "json", &ts.tls) .await .expect("custom-json should be imported"); } #[tokio::test] #[allow(deprecated)] async fn provider_profile_import_preserves_advanced_network_policy_fields() { let ts = run_server().await; let dir = tempfile::tempdir().unwrap(); let profile_path = dir.path().join("advanced-api.yaml"); std::fs::write( &profile_path, r" id: advanced-api display_name: Advanced API category: other endpoints: - host: api.advanced.example ports: [443, 8443] protocol: rest tls: terminate enforcement: enforce rules: - allow: method: GET path: /v1/** allowed_ips: [10.0.0.0/24] deny_rules: - method: POST path: /admin/** allow_encoded_slash: true path: /v1 binaries: - path: /usr/bin/advanced harness: true ", ) .unwrap(); run::provider_profile_import(&ts.endpoint, Some(&profile_path), None, &ts.tls) .await .expect("profile import"); let mut client = openshell_cli::tls::grpc_client(&ts.endpoint, &ts.tls) .await .expect("grpc client should connect"); let profile = client .get_provider_profile(openshell_core::proto::GetProviderProfileRequest { id: "advanced-api".to_string(), }) .await .expect("get provider profile") .into_inner() .profile .expect("profile should exist"); let endpoint = profile.endpoints.first().expect("endpoint should exist"); assert_eq!(endpoint.ports, vec![443, 8443]); assert_eq!(endpoint.rules.len(), 1); assert_eq!(endpoint.deny_rules.len(), 1); assert_eq!(endpoint.allowed_ips, vec!["10.0.0.0/24"]); assert!(endpoint.allow_encoded_slash); assert_eq!(endpoint.path, "/v1"); assert!(profile.binaries[0].harness); } #[tokio::test] async fn provider_profile_import_from_directory_parse_error_prevents_partial_import() { let ts = run_server().await; let dir = tempfile::tempdir().unwrap(); std::fs::write( dir.path().join("custom-good.yaml"), r" id: custom-good display_name: Custom Good category: other endpoints: - host: api.good.example port: 443 ", ) .unwrap(); std::fs::write(dir.path().join("broken.yaml"), "id: [\n").unwrap(); let err = run::provider_profile_import(&ts.endpoint, None, Some(dir.path()), &ts.tls) .await .expect_err("profile import --from should fail on parse errors"); assert!( err.to_string().contains("provider profile import failed"), "unexpected error: {err}" ); run::provider_profile_export(&ts.endpoint, "custom-good", "yaml", &ts.tls) .await .expect_err("valid profiles should not be partially imported after local parse errors"); } #[tokio::test] async fn provider_profile_lint_from_directory_reports_parse_errors_without_importing() { let ts = run_server().await; let dir = tempfile::tempdir().unwrap(); std::fs::write( dir.path().join("custom-good.yaml"), r" id: custom-good display_name: Custom Good category: other endpoints: - host: api.good.example port: 443 ", ) .unwrap(); std::fs::write(dir.path().join("broken.yaml"), "id: [\n").unwrap(); let err = run::provider_profile_lint(&ts.endpoint, None, Some(dir.path()), &ts.tls) .await .expect_err("profile lint --from should fail on parse errors"); assert!( err.to_string().contains("provider profile lint failed"), "unexpected error: {err}" ); run::provider_profile_export(&ts.endpoint, "custom-good", "yaml", &ts.tls) .await .expect_err("lint should not import valid profiles"); } #[tokio::test] async fn provider_create_rejects_key_only_credentials_without_local_env_value() { let ts = run_server().await; let err = run::provider_create( &ts.endpoint, "bad-provider", "claude", false, &["INVALID_PAIR".to_string()], &[], &ts.tls, ) .await .expect_err("invalid key=value should fail"); assert!( err.to_string() .contains("requires local env var 'INVALID_PAIR' to be set to a non-empty value"), "unexpected error: {err}" ); } #[tokio::test] async fn provider_create_supports_generic_type_and_env_lookup_credentials() { let ts = run_server().await; let _guard = EnvVarGuard::set(&[("NAV_GENERIC_TEST_KEY", "generic-value")]); run::provider_create( &ts.endpoint, "my-generic", "generic", false, &["NAV_GENERIC_TEST_KEY".to_string()], &[], &ts.tls, ) .await .expect("provider create"); let mut client = openshell_cli::tls::grpc_client(&ts.endpoint, &ts.tls) .await .expect("grpc client should connect"); let response = client .get_provider(GetProviderRequest { name: "my-generic".to_string(), }) .await .expect("get provider should succeed") .into_inner(); let provider = response.provider.expect("provider should exist"); assert_eq!(provider.r#type, "generic"); assert_eq!( provider.credentials.get("NAV_GENERIC_TEST_KEY"), Some(&"generic-value".to_string()) ); } #[tokio::test] async fn provider_create_rejects_combined_from_existing_and_credentials() { let ts = run_server().await; let err = run::provider_create( &ts.endpoint, "bad-provider", "claude", true, &["API_KEY=abc".to_string()], &[], &ts.tls, ) .await .expect_err("from-existing and credentials should be mutually exclusive"); assert!( err.to_string() .contains("--from-existing cannot be combined with --credential"), "unexpected error: {err}" ); } #[tokio::test] async fn provider_create_rejects_empty_env_var_for_key_only_credential() { let ts = run_server().await; let _guard = EnvVarGuard::set(&[("NAV_EMPTY_ENV_KEY", "")]); let err = run::provider_create( &ts.endpoint, "bad-provider", "generic", false, &["NAV_EMPTY_ENV_KEY".to_string()], &[], &ts.tls, ) .await .expect_err("empty env var should be rejected"); assert!( err.to_string() .contains("requires local env var 'NAV_EMPTY_ENV_KEY' to be set to a non-empty value"), "unexpected error: {err}" ); } #[tokio::test] async fn provider_create_supports_nvidia_type_with_nvidia_api_key() { let ts = run_server().await; let _guard = EnvVarGuard::set(&[("NVIDIA_API_KEY", "nvapi-live-test")]); run::provider_create( &ts.endpoint, "my-nvidia", "nvidia", false, &["NVIDIA_API_KEY".to_string()], &[], &ts.tls, ) .await .expect("provider create"); let mut client = openshell_cli::tls::grpc_client(&ts.endpoint, &ts.tls) .await .expect("grpc client should connect"); let response = client .get_provider(GetProviderRequest { name: "my-nvidia".to_string(), }) .await .expect("get provider should succeed") .into_inner(); let provider = response.provider.expect("provider should exist"); assert_eq!(provider.r#type, "nvidia"); assert_eq!( provider.credentials.get("NVIDIA_API_KEY"), Some(&"nvapi-live-test".to_string()) ); }