kontext_dev_sdk/

orchestrator.rs

use std::collections::{BTreeMap, HashMap, HashSet};
use std::future::Future;
use std::sync::Arc;
use std::time::{Duration, Instant};

use serde_json::{Map, Value};
use tokio::sync::{Mutex, RwLock};

use crate::KontextDevError;
use crate::client::{
    ClientState, ConnectSessionResult, IntegrationInfo, KontextClient, KontextClientConfig,
    ToolResult, create_kontext_client,
};
use crate::mcp::{
    KontextMcp, KontextMcpConfig, KontextTool, RuntimeIntegrationCategory,
    RuntimeIntegrationRecord, extract_text_content,
};
use crate::prompt_guidance::{KontextPromptGuidance, build_kontext_prompt_guidance};

const RETRY_DELAY_MS: u64 = 250;
const INTERNAL_TOKEN_SKEW_SECS: u64 = 30;

#[derive(Clone, Debug)]
pub struct ToolRouteRecord {
    pub tool_id: String,
    pub backend_id: String,
    pub source: BackendSource,
    pub backend_tool_id: String,
    pub integration_id: Option<String>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BackendSource {
    Gateway,
    Internal,
}

#[derive(Clone, Debug)]
pub struct RouteConflictRecord {
    pub tool_id: String,
    pub kept: ToolRouteRecord,
    pub dropped: ToolRouteRecord,
}

#[derive(Clone, Debug)]
struct RouteInventoryCandidate {
    tool: KontextTool,
    route: ToolRouteRecord,
}

#[derive(Clone, Debug)]
struct RouteInventorySnapshot {
    version: u64,
    tools: Vec<KontextTool>,
    routes: HashMap<String, ToolRouteRecord>,
    #[allow(dead_code)]
    conflicts: Vec<RouteConflictRecord>,
}

impl RouteInventorySnapshot {
    fn empty() -> Self {
        Self {
            version: 0,
            tools: Vec::new(),
            routes: HashMap::new(),
            conflicts: Vec::new(),
        }
    }
}

#[derive(Clone, Debug)]
struct InternalResourceToken {
    access_token: String,
    expires_at: Option<Instant>,
}

impl InternalResourceToken {
    fn is_fresh(&self) -> bool {
        match self.expires_at {
            Some(expires_at) => {
                let skew = Duration::from_secs(INTERNAL_TOKEN_SKEW_SECS);
                Instant::now()
                    .checked_add(skew)
                    .is_some_and(|v| v < expires_at)
            }
            None => false,
        }
    }
}

pub type KontextOrchestratorState = ClientState;
pub type KontextOrchestratorConfig = KontextClientConfig;

#[derive(Clone)]
pub struct KontextOrchestrator {
    state: Arc<RwLock<KontextOrchestratorState>>,
    config: KontextOrchestratorConfig,
    gateway_client: KontextClient,
    route_inventory: Arc<RwLock<RouteInventorySnapshot>>,
    build_lock: Arc<Mutex<()>>,
    last_build_key: Arc<RwLock<Option<String>>>,
    version_counter: Arc<Mutex<u64>>,
    reset_generation: Arc<Mutex<u64>>,
    build_run_counter: Arc<Mutex<u64>>,
    latest_committed_build_run: Arc<Mutex<u64>>,
    runtime_integrations: Arc<RwLock<HashMap<String, RuntimeIntegrationRecord>>>,
    internal_mcps: Arc<RwLock<HashMap<String, KontextMcp>>>,
    internal_token_cache: Arc<Mutex<HashMap<String, InternalResourceToken>>>,
    internal_token_locks: Arc<Mutex<HashMap<String, Arc<Mutex<()>>>>>,
    managed_internal_ops: Arc<Mutex<HashMap<String, u32>>>,
    resolved_internal_listings: Arc<RwLock<HashSet<String>>>,
    integration_refresh_lock: Arc<Mutex<()>>,
}

impl KontextOrchestrator {
    pub fn new(config: KontextOrchestratorConfig) -> Self {
        let gateway_client = create_kontext_client(config.clone());

        Self {
            state: Arc::new(RwLock::new(ClientState::Idle)),
            config,
            gateway_client,
            route_inventory: Arc::new(RwLock::new(RouteInventorySnapshot::empty())),
            build_lock: Arc::new(Mutex::new(())),
            last_build_key: Arc::new(RwLock::new(None)),
            version_counter: Arc::new(Mutex::new(0)),
            reset_generation: Arc::new(Mutex::new(0)),
            build_run_counter: Arc::new(Mutex::new(0)),
            latest_committed_build_run: Arc::new(Mutex::new(0)),
            runtime_integrations: Arc::new(RwLock::new(HashMap::new())),
            internal_mcps: Arc::new(RwLock::new(HashMap::new())),
            internal_token_cache: Arc::new(Mutex::new(HashMap::new())),
            internal_token_locks: Arc::new(Mutex::new(HashMap::new())),
            managed_internal_ops: Arc::new(Mutex::new(HashMap::new())),
            resolved_internal_listings: Arc::new(RwLock::new(HashSet::new())),
            integration_refresh_lock: Arc::new(Mutex::new(())),
        }
    }

    pub async fn state(&self) -> KontextOrchestratorState {
        *self.state.read().await
    }

    pub fn mcp(&self) -> &KontextMcp {
        self.gateway_client.mcp()
    }

    pub async fn connect(&self) -> Result<(), KontextDevError> {
        self.ensure_connected().await
    }

    pub async fn disconnect(&self) {
        self.dispose_internal_clients().await;
        self.runtime_integrations.write().await.clear();
        self.reset_inventory().await;
        self.gateway_client.disconnect().await;
        self.internal_token_cache.lock().await.clear();
        *self.state.write().await = ClientState::Idle;
    }

    pub async fn get_connect_page_url(&self) -> Result<ConnectSessionResult, KontextDevError> {
        self.ensure_connected().await?;
        self.gateway_client.get_connect_page_url().await
    }

    pub async fn sign_in(&self) -> Result<(), KontextDevError> {
        // Mirrors TS semantics: sign-in in mixed mode is equivalent to connect.
        self.ensure_connected().await
    }

    pub async fn sign_out(&self) -> Result<(), KontextDevError> {
        self.dispose_internal_clients().await;
        self.runtime_integrations.write().await.clear();
        self.reset_inventory().await;
        self.internal_token_cache.lock().await.clear();
        self.gateway_client.sign_out().await?;
        *self.state.write().await = ClientState::Idle;
        Ok(())
    }

    pub async fn integrations_list(&self) -> Result<Vec<IntegrationInfo>, KontextDevError> {
        self.ensure_connected().await?;

        let discovered = self.refresh_integration_inventory().await?;

        let gateway_statuses = self
            .gateway_client
            .integrations_list()
            .await
            .unwrap_or_default();
        let gateway_by_id = gateway_statuses
            .into_iter()
            .map(|item| (item.id.clone(), item))
            .collect::<HashMap<_, _>>();

        let needs_internal_connect_link = discovered.iter().any(|integration| {
            integration.category == RuntimeIntegrationCategory::InternalMcpCredentials
                && !integration
                    .connection
                    .as_ref()
                    .map(|c| c.connected)
                    .unwrap_or(false)
        });

        let mut shared_connect_url = None;
        if needs_internal_connect_link
            && let Ok(connect) = self.gateway_client.get_connect_page_url().await
        {
            shared_connect_url = Some(connect.connect_url);
        }

        Ok(discovered
            .into_iter()
            .map(|integration| {
                let gateway_status = gateway_by_id.get(&integration.id);
                let connected = gateway_status
                    .map(|item| item.connected)
                    .unwrap_or_else(|| {
                        integration
                            .connection
                            .as_ref()
                            .map(|c| c.connected)
                            .unwrap_or(false)
                    });
                let connect_url = gateway_status
                    .and_then(|item| item.connect_url.clone())
                    .or_else(|| {
                        (!connected
                            && integration.category
                                == RuntimeIntegrationCategory::InternalMcpCredentials)
                            .then(|| shared_connect_url.clone())
                            .flatten()
                    });
                let reason = gateway_status
                    .and_then(|item| item.reason.clone())
                    .or_else(|| {
                        (!connected
                            && integration.category
                                == RuntimeIntegrationCategory::InternalMcpCredentials)
                            .then_some("credentials_required".to_string())
                    });

                IntegrationInfo {
                    id: integration.id,
                    name: integration.name,
                    connected,
                    connect_url,
                    reason,
                }
            })
            .collect())
    }

    pub async fn tools_list(&self) -> Result<Vec<KontextTool>, KontextDevError> {
        self.tools_list_with_options(None).await
    }

    pub async fn prompt_guidance(&self) -> Result<KontextPromptGuidance, KontextDevError> {
        let tools = self.tools_list().await?;
        let integrations = self.integrations_list().await?;
        let tool_names = tools
            .into_iter()
            .map(|tool| tool.name)
            .collect::<Vec<String>>();

        Ok(build_kontext_prompt_guidance(
            tool_names.as_slice(),
            integrations.as_slice(),
        ))
    }

    pub async fn tools_list_with_options(
        &self,
        options: Option<ToolsListOptions>,
    ) -> Result<Vec<KontextTool>, KontextDevError> {
        self.ensure_connected().await?;

        let mut inventory = self.build_tool_inventory(options.clone(), None).await?;

        let missing_internal = self.get_internal_integrations_missing_tools().await;
        if !missing_internal.is_empty() {
            let refreshed = self.refresh_integration_inventory().await?;
            inventory = self.build_tool_inventory(options, Some(refreshed)).await?;
        }

        let unresolved_internal = self.get_internal_integrations_missing_tools().await;
        if self.state().await == ClientState::NeedsAuth
            && self.gateway_client.state().await == ClientState::Ready
            && unresolved_internal.is_empty()
        {
            *self.state.write().await = ClientState::Ready;
        }

        Ok(inventory.tools)
    }

    pub async fn tools_execute(
        &self,
        tool_id: &str,
        args: Option<Map<String, Value>>,
    ) -> Result<ToolResult, KontextDevError> {
        self.ensure_connected().await?;

        let route = self.route_for_execution(tool_id).await?;

        match route.source {
            BackendSource::Gateway => {
                self.gateway_client
                    .tools_execute(route.backend_tool_id.as_str(), args)
                    .await
            }
            BackendSource::Internal => {
                let Some(integration_id) = route.integration_id else {
                    return Err(KontextDevError::ConnectSession {
                        message: format!(
                            "route for tool `{tool_id}` is missing integration metadata"
                        ),
                    });
                };

                let integration =
                    self.integration_by_id(&integration_id)
                        .await?
                        .ok_or_else(|| KontextDevError::ConnectSession {
                            message: format!(
                                "internal integration `{integration_id}` is no longer attached"
                            ),
                        })?;

                let internal_mcp = self.internal_mcp_for(&integration).await.ok_or_else(|| {
                    KontextDevError::ConnectSession {
                        message: format!(
                            "internal integration `{integration_id}` has no MCP client"
                        ),
                    }
                })?;

                let token = match self
                    .ensure_internal_resource_token(&integration, false)
                    .await
                {
                    Ok(token) => token,
                    Err(err) if is_auth_recovery_required(&err) => {
                        self.ensure_internal_resource_token(&integration, true)
                            .await?
                    }
                    Err(err) => return Err(err),
                };

                let integration_id_for_exec = integration_id.clone();
                let execute_once = || {
                    let integration_id = integration_id_for_exec.clone();
                    let internal_mcp = internal_mcp.clone();
                    let token = token.clone();
                    let backend_tool_id = route.backend_tool_id.clone();
                    let args = args.clone();
                    async move {
                        self.run_managed_internal_op(integration_id.as_str(), || async {
                            let raw = internal_mcp
                                .call_tool_with_access_token(
                                    token.as_str(),
                                    backend_tool_id.as_str(),
                                    args,
                                )
                                .await?;
                            Ok(ToolResult {
                                content: extract_text_content(&raw),
                                raw,
                            })
                        })
                        .await
                    }
                };

                match with_transient_retry(execute_once, 1).await {
                    Ok(result) => Ok(result),
                    Err(err) => {
                        if is_auth_recovery_required(&err) {
                            *self.state.write().await = ClientState::NeedsAuth;
                        }
                        Err(err)
                    }
                }
            }
        }
    }

    async fn ensure_connected(&self) -> Result<(), KontextDevError> {
        let state = self.state().await;
        if state == ClientState::Ready {
            return Ok(());
        }

        *self.state.write().await = ClientState::Connecting;

        match self.gateway_client.connect().await {
            Ok(()) => {
                if let Err(err) = self.refresh_integration_inventory().await {
                    if is_auth_recovery_required(&err) {
                        *self.state.write().await = ClientState::NeedsAuth;
                    } else {
                        *self.state.write().await = ClientState::Failed;
                    }
                    return Err(err);
                }
                *self.state.write().await = ClientState::Ready;
                Ok(())
            }
            Err(err) => {
                if is_auth_recovery_required(&err) {
                    *self.state.write().await = ClientState::NeedsAuth;
                } else {
                    *self.state.write().await = ClientState::Failed;
                }
                Err(err)
            }
        }
    }

    async fn refresh_integration_inventory(
        &self,
    ) -> Result<Vec<RuntimeIntegrationRecord>, KontextDevError> {
        let _guard = self.integration_refresh_lock.lock().await;

        let list_once = || async { self.gateway_client.mcp().list_integrations().await };
        let items = with_transient_retry(list_once, 1).await?;
        self.apply_runtime_integrations(&items).await;
        Ok(items)
    }

    async fn apply_runtime_integrations(&self, integrations: &[RuntimeIntegrationRecord]) {
        let mut runtime = self.runtime_integrations.write().await;
        runtime.clear();

        for integration in integrations {
            runtime.insert(integration.id.clone(), integration.clone());
        }
        drop(runtime);

        let desired = integrations
            .iter()
            .filter(|integration| {
                integration.category == RuntimeIntegrationCategory::InternalMcpCredentials
            })
            .map(|integration| (integration.id.clone(), integration.url.clone()))
            .collect::<BTreeMap<_, _>>();

        let mut internal_mcps = self.internal_mcps.write().await;

        internal_mcps.retain(|integration_id, mcp| {
            desired
                .get(integration_id)
                .is_some_and(|url| mcp.mcp_url().ok().as_deref() == Some(url.as_str()))
        });

        for integration in integrations {
            if integration.category != RuntimeIntegrationCategory::InternalMcpCredentials {
                continue;
            }
            if internal_mcps.contains_key(&integration.id) {
                continue;
            }

            internal_mcps.insert(
                integration.id.clone(),
                KontextMcp::new(KontextMcpConfig {
                    client_session_id: self.config.client_session_id.clone(),
                    client_id: self.config.client_id.clone(),
                    redirect_uri: self.config.redirect_uri.clone(),
                    url: Some(integration.url.clone()),
                    server: self.config.server_url.clone(),
                    client_secret: self.config.client_secret.clone(),
                    scope: self.config.scope.clone(),
                    resource: Some(integration.url.clone()),
                    session_key: Some(format!("internal:{}", integration.id)),
                    integration_ui_url: self.config.integration_ui_url.clone(),
                    integration_return_to: self.config.integration_return_to.clone(),
                    auth_timeout_seconds: self.config.auth_timeout_seconds,
                    open_connect_page_on_login: Some(false),
                    token_cache_path: None,
                }),
            );
        }
    }

    async fn internal_mcp_for(&self, integration: &RuntimeIntegrationRecord) -> Option<KontextMcp> {
        self.internal_mcps
            .read()
            .await
            .get(&integration.id)
            .cloned()
    }

    async fn integration_by_id(
        &self,
        integration_id: &str,
    ) -> Result<Option<RuntimeIntegrationRecord>, KontextDevError> {
        if let Some(existing) = self
            .runtime_integrations
            .read()
            .await
            .get(integration_id)
            .cloned()
        {
            return Ok(Some(existing));
        }

        let refreshed = self.refresh_integration_inventory().await?;
        Ok(refreshed.into_iter().find(|item| item.id == integration_id))
    }

    async fn ensure_internal_resource_token(
        &self,
        integration: &RuntimeIntegrationRecord,
        force_exchange: bool,
    ) -> Result<String, KontextDevError> {
        if !force_exchange
            && let Some(existing) = self
                .internal_token_cache
                .lock()
                .await
                .get(&integration.id)
                .cloned()
            && existing.is_fresh()
        {
            return Ok(existing.access_token);
        }

        let integration_lock = {
            let mut locks = self.internal_token_locks.lock().await;
            locks
                .entry(integration.id.clone())
                .or_insert_with(|| Arc::new(Mutex::new(())))
                .clone()
        };

        let _guard = integration_lock.lock().await;

        if !force_exchange
            && let Some(existing) = self
                .internal_token_cache
                .lock()
                .await
                .get(&integration.id)
                .cloned()
            && existing.is_fresh()
        {
            return Ok(existing.access_token);
        }

        let session = self.gateway_client.mcp().authenticate_mcp().await?;
        let exchanged = self
            .gateway_client
            .mcp()
            .client()
            .exchange_for_resource(&session.gateway_token.access_token, &integration.url, None)
            .await?;

        let expires_at = exchanged.expires_in.and_then(|seconds| {
            (seconds > 0)
                .then(|| Instant::now().checked_add(Duration::from_secs(seconds as u64)))
                .flatten()
        });

        let token = InternalResourceToken {
            access_token: exchanged.access_token,
            expires_at,
        };

        let value = token.access_token.clone();
        self.internal_token_cache
            .lock()
            .await
            .insert(integration.id.clone(), token);
        Ok(value)
    }

    async fn build_inventory_candidates(
        &self,
        options: Option<ToolsListOptions>,
        runtime_integrations_override: Option<Vec<RuntimeIntegrationRecord>>,
    ) -> Result<Vec<RouteInventoryCandidate>, KontextDevError> {
        self.resolved_internal_listings.write().await.clear();

        let mut candidates = Vec::new();

        let list_gateway_tools_once = || {
            let gateway_client = self.gateway_client.clone();
            async move { gateway_client.tools_list().await }
        };
        let gateway_tools = with_transient_retry(list_gateway_tools_once, 1)
            .await?
            .into_iter()
            .take(
                options
                    .as_ref()
                    .and_then(|value| value.limit)
                    .unwrap_or(u32::MAX) as usize,
            )
            .collect::<Vec<_>>();
        for tool in gateway_tools {
            let route = ToolRouteRecord {
                tool_id: tool.id.clone(),
                backend_id: gateway_backend_id(),
                source: BackendSource::Gateway,
                backend_tool_id: tool.id.clone(),
                integration_id: None,
            };
            candidates.push(RouteInventoryCandidate { tool, route });
        }

        let integrations = if let Some(runtime_integrations) = runtime_integrations_override {
            self.apply_runtime_integrations(&runtime_integrations).await;
            runtime_integrations
        } else {
            self.refresh_integration_inventory().await?
        };

        let mut internal_integrations = integrations
            .into_iter()
            .filter(|integration| {
                integration.category == RuntimeIntegrationCategory::InternalMcpCredentials
            })
            .collect::<Vec<_>>();
        internal_integrations.sort_by(|a, b| a.id.cmp(&b.id));

        let mut pending_auth_retries = Vec::new();

        for integration in &internal_integrations {
            match self
                .add_internal_tools_for_integration(&mut candidates, integration, false)
                .await
            {
                Ok(()) => {}
                Err(err) => {
                    if is_authorization_error(&err) {
                        pending_auth_retries.push(integration.clone());
                        continue;
                    }
                    if is_token_exchange_error(&err) {
                        *self.state.write().await = ClientState::NeedsAuth;
                    }
                }
            }
        }

        let mut unresolved_internal_auth = false;
        for integration in pending_auth_retries {
            match self
                .add_internal_tools_for_integration(&mut candidates, &integration, true)
                .await
            {
                Ok(()) => {}
                Err(err) => {
                    if is_auth_recovery_required(&err) {
                        unresolved_internal_auth = true;
                        continue;
                    }
                }
            }
        }

        if unresolved_internal_auth {
            *self.state.write().await = ClientState::NeedsAuth;
        } else if self.state().await == ClientState::NeedsAuth
            && self.gateway_client.state().await == ClientState::Ready
        {
            *self.state.write().await = ClientState::Ready;
        }

        Ok(candidates)
    }

    async fn add_internal_tools_for_integration(
        &self,
        candidates: &mut Vec<RouteInventoryCandidate>,
        integration: &RuntimeIntegrationRecord,
        force_exchange: bool,
    ) -> Result<(), KontextDevError> {
        let Some(internal_mcp) = self.internal_mcp_for(integration).await else {
            return Ok(());
        };

        let token = self
            .ensure_internal_resource_token(integration, force_exchange)
            .await?;

        let list_once = || {
            let integration_id = integration.id.clone();
            let internal_mcp = internal_mcp.clone();
            let token = token.clone();
            async move {
                self.run_managed_internal_op(&integration_id, || async {
                    internal_mcp
                        .list_tools_with_access_token(token.as_str())
                        .await
                })
                .await
            }
        };

        let tools = with_transient_retry(list_once, 1).await?;

        self.resolved_internal_listings
            .write()
            .await
            .insert(integration.id.clone());

        for tool in tools {
            let backend_tool_id = if tool.id.is_empty() {
                tool.name.clone()
            } else {
                tool.id.clone()
            };
            let unified_id = format!("{}:{}", integration.id, tool.name);

            let mapped = KontextTool {
                id: unified_id.clone(),
                name: tool.name,
                description: tool.description,
                input_schema: tool.input_schema,
                server: Some(crate::mcp::KontextToolServer {
                    id: integration.id.clone(),
                    name: Some(integration.name.clone()),
                }),
            };

            let route = ToolRouteRecord {
                tool_id: unified_id,
                backend_id: internal_backend_id(&integration.id),
                source: BackendSource::Internal,
                backend_tool_id,
                integration_id: Some(integration.id.clone()),
            };

            candidates.push(RouteInventoryCandidate {
                tool: mapped,
                route,
            });
        }

        Ok(())
    }

    async fn run_managed_internal_op<T, F, Fut>(
        &self,
        integration_id: &str,
        operation: F,
    ) -> Result<T, KontextDevError>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = Result<T, KontextDevError>>,
    {
        {
            let mut managed = self.managed_internal_ops.lock().await;
            *managed.entry(integration_id.to_string()).or_insert(0) += 1;
        }

        let result = operation().await;

        {
            let mut managed = self.managed_internal_ops.lock().await;
            if let Some(count) = managed.get_mut(integration_id) {
                *count = count.saturating_sub(1);
                if *count == 0 {
                    managed.remove(integration_id);
                }
            }
        }

        result
    }

    async fn build_tool_inventory(
        &self,
        options: Option<ToolsListOptions>,
        runtime_integrations_override: Option<Vec<RuntimeIntegrationRecord>>,
    ) -> Result<RouteInventorySnapshot, KontextDevError> {
        let key = inventory_build_key(options.clone(), runtime_integrations_override.as_ref());

        let _guard = self.build_lock.lock().await;

        {
            let last_key = self.last_build_key.read().await;
            let snapshot = self.route_inventory.read().await;
            if snapshot.version > 0 && last_key.as_deref() == Some(key.as_str()) {
                return Ok(snapshot.clone());
            }
        }

        let build_generation = *self.reset_generation.lock().await;
        let build_run = {
            let mut counter = self.build_run_counter.lock().await;
            *counter += 1;
            *counter
        };

        let candidates = self
            .build_inventory_candidates(options, runtime_integrations_override)
            .await?;

        if *self.reset_generation.lock().await != build_generation {
            return Ok(self.route_inventory.read().await.clone());
        }

        let latest_committed = *self.latest_committed_build_run.lock().await;
        let snapshot = if build_run >= latest_committed {
            *self.latest_committed_build_run.lock().await = build_run;
            let mut version = self.version_counter.lock().await;
            *version += 1;
            let built = build_route_inventory(*version, candidates);
            *self.route_inventory.write().await = built.clone();
            *self.last_build_key.write().await = Some(key);
            built
        } else {
            let current_version = self.route_inventory.read().await.version;
            build_route_inventory(current_version, candidates)
        };

        Ok(snapshot)
    }

    async fn route_for_execution(&self, tool_id: &str) -> Result<ToolRouteRecord, KontextDevError> {
        if let Some(route) = self
            .route_inventory
            .read()
            .await
            .routes
            .get(tool_id)
            .cloned()
        {
            return Ok(route);
        }

        let _ = self.build_tool_inventory(None, None).await?;

        if let Some(route) = self
            .route_inventory
            .read()
            .await
            .routes
            .get(tool_id)
            .cloned()
        {
            return Ok(route);
        }

        Err(KontextDevError::ConnectSession {
            message: format!(
                "unknown tool `{tool_id}`. Call tools.list() to refresh available tools"
            ),
        })
    }

    async fn get_internal_integrations_missing_tools(&self) -> Vec<String> {
        let resolved = self.resolved_internal_listings.read().await.clone();
        self.internal_mcps
            .read()
            .await
            .keys()
            .filter(|integration_id| !resolved.contains(*integration_id))
            .cloned()
            .collect()
    }

    async fn dispose_internal_clients(&self) {
        self.internal_mcps.write().await.clear();
        self.internal_token_cache.lock().await.clear();
        self.internal_token_locks.lock().await.clear();
        self.resolved_internal_listings.write().await.clear();
    }

    async fn reset_inventory(&self) {
        *self.reset_generation.lock().await += 1;
        *self.version_counter.lock().await = 0;

        let build_run = *self.build_run_counter.lock().await;
        *self.latest_committed_build_run.lock().await = build_run;

        *self.route_inventory.write().await = RouteInventorySnapshot::empty();
        *self.last_build_key.write().await = None;
    }
}

#[derive(Clone, Debug)]
pub struct ToolsListOptions {
    pub limit: Option<u32>,
}

pub fn create_kontext_orchestrator(config: KontextOrchestratorConfig) -> KontextOrchestrator {
    KontextOrchestrator::new(config)
}

fn gateway_backend_id() -> String {
    "gateway".to_string()
}

fn internal_backend_id(integration_id: &str) -> String {
    format!("internal:{integration_id}")
}

fn inventory_build_key(
    options: Option<ToolsListOptions>,
    runtime_integrations: Option<&Vec<RuntimeIntegrationRecord>>,
) -> String {
    let limit_key = match options.and_then(|item| item.limit) {
        Some(limit) => limit.to_string(),
        None => "none".to_string(),
    };
    let mut key = format!("limit:{limit_key}");
    key.push_str("|runtime:");
    if let Some(integrations) = runtime_integrations {
        let mut ids = integrations
            .iter()
            .map(|integration| integration.id.clone())
            .collect::<Vec<_>>();
        ids.sort();
        key.push_str(ids.join(",").as_str());
    }
    key
}

fn build_route_inventory(
    version: u64,
    candidates: Vec<RouteInventoryCandidate>,
) -> RouteInventorySnapshot {
    let mut tools = Vec::<KontextTool>::new();
    let mut routes = HashMap::<String, ToolRouteRecord>::new();
    let mut conflicts = Vec::<RouteConflictRecord>::new();

    for candidate in candidates {
        let tool_id = candidate.route.tool_id.clone();
        if !routes.contains_key(&tool_id) {
            routes.insert(tool_id.clone(), candidate.route);
            tools.push(candidate.tool);
            continue;
        }

        let existing = routes
            .get(&tool_id)
            .cloned()
            .unwrap_or_else(|| candidate.route.clone());

        // TS parity default policy is keep_existing.
        conflicts.push(RouteConflictRecord {
            tool_id,
            kept: existing,
            dropped: candidate.route,
        });
    }

    RouteInventorySnapshot {
        version,
        tools,
        routes,
        conflicts,
    }
}

async fn with_transient_retry<T, F, Fut>(
    mut operation: F,
    max_retries: usize,
) -> Result<T, KontextDevError>
where
    F: FnMut() -> Fut,
    Fut: Future<Output = Result<T, KontextDevError>>,
{
    for attempt in 0..=max_retries {
        match operation().await {
            Ok(value) => return Ok(value),
            Err(err) => {
                if attempt >= max_retries || !is_transient_error(&err) {
                    return Err(err);
                }
                tokio::time::sleep(Duration::from_millis(RETRY_DELAY_MS)).await;
            }
        }
    }

    Err(KontextDevError::ConnectSession {
        message: "transient retry loop exhausted unexpectedly".to_string(),
    })
}

fn is_transient_error(err: &KontextDevError) -> bool {
    match err {
        KontextDevError::ConnectSession { message }
        | KontextDevError::TokenExchange { message, .. }
        | KontextDevError::TokenRequest { message, .. } => {
            let lower = message.to_ascii_lowercase();
            lower.contains("429")
                || lower.contains("500")
                || lower.contains("502")
                || lower.contains("503")
                || lower.contains("504")
                || lower.contains("timed out")
                || lower.contains("connection reset")
                || lower.contains("broken pipe")
                || lower.contains("temporarily unavailable")
                || lower.contains("econnreset")
                || lower.contains("network")
        }
        _ => false,
    }
}

fn is_authorization_error(err: &KontextDevError) -> bool {
    match err {
        KontextDevError::OAuthCallbackTimeout { .. }
        | KontextDevError::OAuthCallbackCancelled
        | KontextDevError::MissingAuthorizationCode
        | KontextDevError::OAuthCallbackError { .. }
        | KontextDevError::InvalidOAuthState
        | KontextDevError::TokenRequest { .. }
        | KontextDevError::TokenExchange { .. } => true,
        KontextDevError::ConnectSession { message }
        | KontextDevError::IntegrationOAuthInit { message } => {
            let lower = message.to_ascii_lowercase();
            lower.contains("401")
                || lower.contains("unauthorized")
                || lower.contains("invalid token")
                || lower.contains("authorization")
                || lower.contains("access token expired")
        }
        _ => false,
    }
}

fn is_token_exchange_error(err: &KontextDevError) -> bool {
    matches!(err, KontextDevError::TokenExchange { .. })
}

fn is_auth_recovery_required(err: &KontextDevError) -> bool {
    is_authorization_error(err) || is_token_exchange_error(err)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn route(tool_id: &str, backend_id: &str) -> ToolRouteRecord {
        ToolRouteRecord {
            tool_id: tool_id.to_string(),
            backend_id: backend_id.to_string(),
            source: BackendSource::Gateway,
            backend_tool_id: tool_id.to_string(),
            integration_id: None,
        }
    }

    fn candidate(tool_id: &str, backend_id: &str) -> RouteInventoryCandidate {
        RouteInventoryCandidate {
            tool: KontextTool {
                id: tool_id.to_string(),
                name: tool_id.to_string(),
                description: None,
                input_schema: None,
                server: None,
            },
            route: route(tool_id, backend_id),
        }
    }

    #[test]
    fn build_route_inventory_keeps_first_route_on_conflict() {
        let snapshot = build_route_inventory(
            1,
            vec![
                candidate("tool-a", "gateway"),
                candidate("tool-a", "internal:linear"),
                candidate("tool-b", "gateway"),
            ],
        );

        assert_eq!(snapshot.tools.len(), 2);
        assert_eq!(snapshot.routes.get("tool-a").unwrap().backend_id, "gateway");
        assert_eq!(snapshot.conflicts.len(), 1);
        assert_eq!(snapshot.conflicts[0].kept.backend_id, "gateway");
        assert_eq!(snapshot.conflicts[0].dropped.backend_id, "internal:linear");
    }

    #[test]
    fn inventory_build_key_is_stable_for_same_runtime_integrations() {
        let runtime = vec![
            RuntimeIntegrationRecord {
                id: "b".to_string(),
                name: "B".to_string(),
                url: "https://b".to_string(),
                category: RuntimeIntegrationCategory::InternalMcpCredentials,
                connect_type: crate::mcp::RuntimeIntegrationConnectType::Credentials,
                auth_mode: None,
                credential_schema: None,
                requires_oauth: None,
                connection: None,
            },
            RuntimeIntegrationRecord {
                id: "a".to_string(),
                name: "A".to_string(),
                url: "https://a".to_string(),
                category: RuntimeIntegrationCategory::InternalMcpCredentials,
                connect_type: crate::mcp::RuntimeIntegrationConnectType::Credentials,
                auth_mode: None,
                credential_schema: None,
                requires_oauth: None,
                connection: None,
            },
        ];

        let first = inventory_build_key(Some(ToolsListOptions { limit: Some(10) }), Some(&runtime));
        let second =
            inventory_build_key(Some(ToolsListOptions { limit: Some(10) }), Some(&runtime));
        assert_eq!(first, second);
    }

    #[test]
    fn inventory_build_key_distinguishes_none_from_zero_limit() {
        let no_limit = inventory_build_key(None, None);
        let zero_limit = inventory_build_key(Some(ToolsListOptions { limit: Some(0) }), None);

        assert_ne!(no_limit, zero_limit);
        assert_eq!(no_limit, "limit:none|runtime:");
        assert_eq!(zero_limit, "limit:0|runtime:");
    }

    #[test]
    fn transient_error_detection_matches_network_and_5xx() {
        let network = KontextDevError::ConnectSession {
            message: "connection reset by peer".to_string(),
        };
        let server_500 = KontextDevError::ConnectSession {
            message: "500 internal server error".to_string(),
        };
        let bad_request = KontextDevError::ConnectSession {
            message: "400 bad request".to_string(),
        };

        assert!(is_transient_error(&network));
        assert!(is_transient_error(&server_500));
        assert!(!is_transient_error(&bad_request));
    }
}