grpc_client.rs

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! gRPC client for fetching sandbox policy, provider environment, and inference
//! route bundles from OpenShell server.

use std::collections::HashMap;
use std::time::Duration;

use miette::{IntoDiagnostic, Result, WrapErr};
use openshell_core::proto::{
    DenialSummary, GetInferenceBundleRequest, GetInferenceBundleResponse, GetSandboxPolicyRequest,
    GetSandboxProviderEnvironmentRequest, PolicyStatus, ReportPolicyStatusRequest,
    SandboxPolicy as ProtoSandboxPolicy, SubmitPolicyAnalysisRequest, UpdateSandboxPolicyRequest,
    inference_client::InferenceClient, open_shell_client::OpenShellClient,
};
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity};
use tracing::debug;

/// Create a channel to the OpenShell server.
///
/// When the endpoint uses `https://`, mTLS is configured using these env vars:
/// - `OPENSHELL_TLS_CA` -- path to the CA certificate
/// - `OPENSHELL_TLS_CERT` -- path to the client certificate
/// - `OPENSHELL_TLS_KEY` -- path to the client private key
///
/// When the endpoint uses `http://`, a plaintext connection is used (for
/// deployments where TLS is disabled, e.g. behind a Cloudflare Tunnel).
async fn connect_channel(endpoint: &str) -> Result<Channel> {
    let mut ep = Endpoint::from_shared(endpoint.to_string())
        .into_diagnostic()
        .wrap_err("invalid gRPC endpoint")?
        .connect_timeout(Duration::from_secs(10))
        .http2_keep_alive_interval(Duration::from_secs(10))
        .keep_alive_while_idle(true)
        .keep_alive_timeout(Duration::from_secs(20));

    let tls_enabled = endpoint.starts_with("https://");

    if tls_enabled {
        let ca_path = std::env::var("OPENSHELL_TLS_CA")
            .into_diagnostic()
            .wrap_err("OPENSHELL_TLS_CA is required")?;
        let cert_path = std::env::var("OPENSHELL_TLS_CERT")
            .into_diagnostic()
            .wrap_err("OPENSHELL_TLS_CERT is required")?;
        let key_path = std::env::var("OPENSHELL_TLS_KEY")
            .into_diagnostic()
            .wrap_err("OPENSHELL_TLS_KEY is required")?;

        let ca_pem = std::fs::read(&ca_path)
            .into_diagnostic()
            .wrap_err_with(|| format!("failed to read CA cert from {ca_path}"))?;
        let cert_pem = std::fs::read(&cert_path)
            .into_diagnostic()
            .wrap_err_with(|| format!("failed to read client cert from {cert_path}"))?;
        let key_pem = std::fs::read(&key_path)
            .into_diagnostic()
            .wrap_err_with(|| format!("failed to read client key from {key_path}"))?;

        let tls_config = ClientTlsConfig::new()
            .ca_certificate(Certificate::from_pem(ca_pem))
            .identity(Identity::from_pem(cert_pem, key_pem));

        ep = ep
            .tls_config(tls_config)
            .into_diagnostic()
            .wrap_err("failed to configure TLS")?;
    }

    ep.connect()
        .await
        .into_diagnostic()
        .wrap_err("failed to connect to OpenShell server")
}

/// Connect to the OpenShell server (mTLS or plaintext based on endpoint scheme).
async fn connect(endpoint: &str) -> Result<OpenShellClient<Channel>> {
    let channel = connect_channel(endpoint).await?;
    Ok(OpenShellClient::new(channel))
}

/// Fetch sandbox policy from OpenShell server via gRPC.
///
/// Returns `Ok(Some(policy))` when the server has a policy configured,
/// or `Ok(None)` when the sandbox was created without a policy (the sandbox
/// should discover one from disk or use the restrictive default).
pub async fn fetch_policy(endpoint: &str, sandbox_id: &str) -> Result<Option<ProtoSandboxPolicy>> {
    debug!(endpoint = %endpoint, sandbox_id = %sandbox_id, "Connecting to OpenShell server");

    let mut client = connect(endpoint).await?;

    debug!("Connected, fetching sandbox policy");

    fetch_policy_with_client(&mut client, sandbox_id).await
}

/// Fetch sandbox policy using an existing client connection.
async fn fetch_policy_with_client(
    client: &mut OpenShellClient<Channel>,
    sandbox_id: &str,
) -> Result<Option<ProtoSandboxPolicy>> {
    let response = client
        .get_sandbox_policy(GetSandboxPolicyRequest {
            sandbox_id: sandbox_id.to_string(),
        })
        .await
        .into_diagnostic()?;

    let inner = response.into_inner();

    // version 0 with no policy means the sandbox was created without one.
    if inner.version == 0 && inner.policy.is_none() {
        return Ok(None);
    }

    Ok(Some(inner.policy.ok_or_else(|| {
        miette::miette!("Server returned non-zero version but empty policy")
    })?))
}

/// Sync a locally-discovered policy using an existing client connection.
async fn sync_policy_with_client(
    client: &mut OpenShellClient<Channel>,
    sandbox: &str,
    policy: &ProtoSandboxPolicy,
) -> Result<()> {
    client
        .update_sandbox_policy(UpdateSandboxPolicyRequest {
            name: sandbox.to_string(),
            policy: Some(policy.clone()),
        })
        .await
        .into_diagnostic()
        .wrap_err("failed to sync policy to server")?;

    Ok(())
}

/// Discover and sync policy using a single gRPC connection.
///
/// Performs the full discovery flow (fetch → sync → re-fetch) over one
/// channel instead of establishing three separate connections.
pub async fn discover_and_sync_policy(
    endpoint: &str,
    sandbox_id: &str,
    sandbox: &str,
    discovered_policy: &ProtoSandboxPolicy,
) -> Result<ProtoSandboxPolicy> {
    debug!(
        endpoint = %endpoint,
        sandbox_id = %sandbox_id,
        sandbox = %sandbox,
        "Syncing discovered policy and re-fetching canonical version"
    );

    let mut client = connect(endpoint).await?;

    // Sync the discovered policy to the gateway.
    sync_policy_with_client(&mut client, sandbox, discovered_policy).await?;

    // Re-fetch from the gateway to get the canonical version/hash.
    fetch_policy_with_client(&mut client, sandbox_id)
        .await?
        .ok_or_else(|| {
            miette::miette!("Server still returned no policy after sync — this is a bug")
        })
}

/// Sync an enriched policy back to the gateway.
///
/// Used by the supervisor to push baseline-path-enriched policies so the
/// gateway stores the effective policy users see via `openshell sandbox get`.
pub async fn sync_policy(endpoint: &str, sandbox: &str, policy: &ProtoSandboxPolicy) -> Result<()> {
    debug!(endpoint = %endpoint, sandbox = %sandbox, "Syncing enriched policy to gateway");
    let mut client = connect(endpoint).await?;
    sync_policy_with_client(&mut client, sandbox, policy).await
}

/// Provider environment fetched from the server, indexed by provider type.
pub struct ProviderEnvironment {
    /// Env vars indexed by provider type (e.g. `"anthropic"` -> `{"ANTHROPIC_API_KEY": "sk-..."}`).
    pub by_type: HashMap<String, HashMap<String, String>>,
}

impl ProviderEnvironment {
    /// Flatten all provider env vars into a single map for injection into the
    /// child process. When two different provider types set the same env var,
    /// one value wins arbitrarily (iteration order over `HashMap` keys is
    /// nondeterministic).
    pub fn flatten(self) -> HashMap<String, String> {
        let mut flat = HashMap::new();
        for (_provider_type, env) in self.by_type {
            for (key, value) in env {
                flat.entry(key).or_insert(value);
            }
        }
        flat
    }

    /// Returns the set of provider types present.
    pub fn provider_types(&self) -> Vec<String> {
        self.by_type.keys().cloned().collect()
    }

    /// Check if a specific provider type is present.
    pub fn has_provider_type(&self, provider_type: &str) -> bool {
        self.by_type.contains_key(provider_type)
    }
}

/// Fetch provider environment variables for a sandbox from OpenShell server via gRPC.
///
/// Returns provider credentials indexed by provider type. Use
/// [`ProviderEnvironment::flatten`] to merge into a single env var map for
/// injection into the child process.
pub async fn fetch_provider_environment(
    endpoint: &str,
    sandbox_id: &str,
) -> Result<ProviderEnvironment> {
    debug!(endpoint = %endpoint, sandbox_id = %sandbox_id, "Fetching provider environment");

    let mut client = connect(endpoint).await?;

    let response = client
        .get_sandbox_provider_environment(GetSandboxProviderEnvironmentRequest {
            sandbox_id: sandbox_id.to_string(),
        })
        .await
        .into_diagnostic()?;

    let inner = response.into_inner();
    let by_type = inner
        .providers
        .into_iter()
        .map(|(provider_type, entry)| (provider_type, entry.environment))
        .collect();
    Ok(ProviderEnvironment { by_type })
}

/// A reusable gRPC client for the OpenShell service.
///
/// Wraps a tonic channel connected once and reused for policy polling
/// and status reporting, avoiding per-request TLS handshake overhead.
#[derive(Clone)]
pub struct CachedOpenShellClient {
    client: OpenShellClient<Channel>,
}

/// Policy poll result returned by [`CachedOpenShellClient::poll_policy`].
pub struct PolicyPollResult {
    pub policy: ProtoSandboxPolicy,
    pub version: u32,
    pub policy_hash: String,
}

impl CachedOpenShellClient {
    pub async fn connect(endpoint: &str) -> Result<Self> {
        debug!(endpoint = %endpoint, "Connecting openshell gRPC client for policy polling");
        let channel = connect_channel(endpoint).await?;
        let client = OpenShellClient::new(channel);
        Ok(Self { client })
    }

    /// Get a clone of the underlying tonic client for direct RPC calls.
    pub fn raw_client(&self) -> OpenShellClient<Channel> {
        self.client.clone()
    }

    /// Poll for the current sandbox policy version.
    pub async fn poll_policy(&self, sandbox_id: &str) -> Result<PolicyPollResult> {
        let response = self
            .client
            .clone()
            .get_sandbox_policy(GetSandboxPolicyRequest {
                sandbox_id: sandbox_id.to_string(),
            })
            .await
            .into_diagnostic()?;

        let inner = response.into_inner();
        let policy = inner
            .policy
            .ok_or_else(|| miette::miette!("Server returned empty policy"))?;

        Ok(PolicyPollResult {
            policy,
            version: inner.version,
            policy_hash: inner.policy_hash,
        })
    }

    /// Submit denial summaries for policy analysis.
    pub async fn submit_policy_analysis(
        &self,
        sandbox_name: &str,
        summaries: Vec<DenialSummary>,
        proposed_chunks: Vec<openshell_core::proto::PolicyChunk>,
        analysis_mode: &str,
    ) -> Result<()> {
        self.client
            .clone()
            .submit_policy_analysis(SubmitPolicyAnalysisRequest {
                name: sandbox_name.to_string(),
                summaries,
                proposed_chunks,
                analysis_mode: analysis_mode.to_string(),
            })
            .await
            .into_diagnostic()?;

        Ok(())
    }

    /// Report policy load status back to the server.
    pub async fn report_policy_status(
        &self,
        sandbox_id: &str,
        version: u32,
        loaded: bool,
        error_msg: &str,
    ) -> Result<()> {
        let status = if loaded {
            PolicyStatus::Loaded
        } else {
            PolicyStatus::Failed
        };

        self.client
            .clone()
            .report_policy_status(ReportPolicyStatusRequest {
                sandbox_id: sandbox_id.to_string(),
                version,
                status: status.into(),
                load_error: error_msg.to_string(),
            })
            .await
            .into_diagnostic()?;

        Ok(())
    }
}

/// Fetch the resolved inference route bundle from the server.
pub async fn fetch_inference_bundle(endpoint: &str) -> Result<GetInferenceBundleResponse> {
    debug!(endpoint = %endpoint, "Fetching inference route bundle");

    let channel = connect_channel(endpoint).await?;
    let mut client = InferenceClient::new(channel);

    let response = client
        .get_inference_bundle(GetInferenceBundleRequest {})
        .await
        .into_diagnostic()?;

    Ok(response.into_inner())
}