Rust Data Placement

[dependencies]
neutron-ai = "0.2"
reqwest = { version = "0.13", default-features = false, features = ["json", "rustls"] }
serde_json = "1.0"
tokio = { version = "1.53", features = ["macros", "rt-multi-thread"] }

use neutron_ai::{ClientOptions, NeutronAIClient};
use reqwest::{Client as HttpClient, Method};
use serde_json::{Value, json};
use std::{env, error::Error};

fn required_env(name: &str) -> Result<String, Box<dyn Error>> {
    let value = env::var(name)?;
    if value.trim().is_empty() {
        return Err(format!("{name} is required").into());
    }
    Ok(value)
}

async fn neutron_request(
    http: &HttpClient,
    base_url: &str,
    token: &str,
    method: Method,
    path: &str,
    body: Option<Value>,
) -> Result<Value, Box<dyn Error>> {
    let mut request = http
        .request(method, format!("{}{}", base_url.trim_end_matches('/'), path))
        .bearer_auth(token);
    if let Some(payload) = body {
        request = request.json(&payload);
    }
    let response = request.send().await?.error_for_status()?;
    Ok(response.json::<Value>().await?)
}

async fn provision_nuclei(
    http: &HttpClient,
    base_url: &str,
    workspace_key: &str,
) -> Result<(), Box<dyn Error>> {
    let profiles = [
        json!({
            "nucleusId": "product-knowledge-global",
            "name": "Global product knowledge",
            "placement": {
                "mode": "global",
                "guarantee": "none",
                "notes": ["Approved global knowledge; no jurisdictional storage guarantee."]
            }
        }),
        json!({
            "nucleusId": "support-operations-weur",
            "name": "Western Europe support operations",
            "placement": {
                "mode": "regional",
                "regionHint": "weur",
                "r2LocationHint": "weur",
                "guarantee": "best_effort_regional",
                "notes": ["Latency preference only; not a legal residency boundary."]
            }
        }),
        json!({
            "nucleusId": "regulated-cases-eu",
            "name": "EU regulated case memory",
            "placement": {
                "mode": "jurisdictional",
                "jurisdiction": "eu",
                "regionalServicesRequired": true,
                "metadataBoundaryRequired": true,
                "geoKeyManagerRequired": false,
                "guarantee": "jurisdictional_storage",
                "notes": [
                    "Requires the EU Durable Object jurisdiction and matching EU archive binding.",
                    "Processing, logs, metadata, keys, retention, and contracts require separate verification."
                ]
            }
        }),
    ];

    for profile in profiles {
        neutron_request(
            http,
            base_url,
            workspace_key,
            Method::POST,
            "/v1/nuclei",
            Some(profile),
        )
        .await?;
    }
    Ok(())
}

struct MemoryRouter {
    global: NeutronAIClient,
    western_europe: NeutronAIClient,
    eu_regulated: NeutronAIClient,
}

fn memory_client(
    base_url: &str,
    token_env: &str,
    nucleus_id: &str,
) -> Result<NeutronAIClient, Box<dyn Error>> {
    Ok(NeutronAIClient::new(ClientOptions {
        base_url: base_url.to_string(),
        token: Some(required_env(token_env)?),
        nucleus_id: Some(nucleus_id.to_string()),
        tenant_id: None,
    })?)
}

impl MemoryRouter {
    fn new(base_url: &str) -> Result<Self, Box<dyn Error>> {
        Ok(Self {
            global: memory_client(
                base_url,
                "NEUTRON_GLOBAL_TOKEN",
                "product-knowledge-global",
            )?,
            western_europe: memory_client(
                base_url,
                "NEUTRON_WEUR_TOKEN",
                "support-operations-weur",
            )?,
            eu_regulated: memory_client(
                base_url,
                "NEUTRON_EU_TOKEN",
                "regulated-cases-eu",
            )?,
        })
    }

    async fn store_global_product_knowledge(&self, text: &str) -> Result<Value, Box<dyn Error>> {
        Ok(self.global.remember(json!({
            "scopeId": "kb:approved-products",
            "type": "knowledge",
            "privacyClass": "public",
            "text": text
        })).await?)
    }

    async fn store_western_europe_support_lesson(
        &self,
        text: &str,
    ) -> Result<Value, Box<dyn Error>> {
        Ok(self.western_europe.remember(json!({
            "scopeId": "history:support-resolutions",
            "type": "tool_lesson",
            "privacyClass": "tenant",
            "text": text
        })).await?)
    }

    async fn store_eu_regulated_case(
        &self,
        case_id: &str,
        minimized_summary: &str,
    ) -> Result<Value, Box<dyn Error>> {
        Ok(self.eu_regulated.remember(json!({
            "scopeId": format!("case:{case_id}"),
            "type": "experience",
            "privacyClass": "user_private",
            "text": minimized_summary,
            "metadata": {
                "dataClass": "eu-regulated-case",
                "source": "approved-case-summary"
            }
        })).await?)
    }
}

impl MemoryRouter {
    async fn eu_case_context(&self) -> Result<Value, Box<dyn Error>> {
        Ok(self.eu_regulated.agent_context(json!({
            "scopeIds": [
                "case:case-4821",
                "policy:eu-case-handling",
                "kb:approved-products"
            ],
            "agentId": "agent:eu-case-support",
            "task": "Prepare the next authorized case-support step.",
            "privacyClasses": ["tenant", "user_private"],
            "tokenBudget": 1_200,
            "cachePolicy": {
                "mode": "prefer_cache",
                "keyMode": "intent_profile",
                "ttlSeconds": 120,
                "includeDynamicRag": true,
                "allowSensitive": false
            }
        })).await?)
    }
}

async fn verify_eu_readiness(
    http: &HttpClient,
    base_url: &str,
    eu_token: &str,
) -> Result<(), Box<dyn Error>> {
    neutron_request(
        http,
        base_url,
        eu_token,
        Method::GET,
        "/v1/nuclei/regulated-cases-eu/placement",
        None,
    )
    .await?;
    let health = neutron_request(
        http,
        base_url,
        eu_token,
        Method::GET,
        "/v1/nuclei/regulated-cases-eu/health",
        None,
    )
    .await?;

    let blocking = [
        "jurisdictional_do_subnamespace_unavailable",
        "jurisdictional_archive_bucket_missing",
    ];
    let warnings = health["warnings"].as_array().cloned().unwrap_or_default();
    if warnings.iter().any(|warning| {
        warning.as_str().is_some_and(|value| blocking.contains(&value))
    }) {
        return Err("EU Nucleus is not ready".into());
    }
    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let base_url = required_env("NEUTRON_API_URL")?;
    let workspace_key = required_env("NEUTRON_WORKSPACE_API_KEY")?;
    let http = HttpClient::new();

    provision_nuclei(&http, &base_url, &workspace_key).await?;
    let router = MemoryRouter::new(&base_url)?;
    router.eu_case_context().await?;
    verify_eu_readiness(
        &http,
        &base_url,
        &required_env("NEUTRON_EU_TOKEN")?,
    )
    .await?;
    Ok(())
}