orchestral-cli 0.2.0

Command line interface for running Orchestral workflows
use std::collections::HashMap;
use std::sync::Arc;

use orchestral_core::store::Event;
use orchestral_runtime::api::{
    ApiError, ApiService, HistoryEventView, InteractionSubmitRequest, InteractionSubmitResponse,
    RuntimeApi,
};
use thiserror::Error;
use tokio::sync::{broadcast, RwLock};
use tracing::debug;

#[derive(Debug, Error)]
pub enum ChannelError {
    #[error("api error: {0}")]
    Api(#[from] ApiError),
}

struct InMemoryChannelBindingStore {
    bindings: RwLock<HashMap<String, String>>,
}

impl InMemoryChannelBindingStore {
    fn new() -> Self {
        Self {
            bindings: RwLock::new(HashMap::new()),
        }
    }

    async fn get_thread_id(&self, binding_key: &str) -> Option<String> {
        self.bindings.read().await.get(binding_key).cloned()
    }

    async fn set_thread_id(&self, binding_key: &str, thread_id: String) {
        self.bindings
            .write()
            .await
            .insert(binding_key.to_string(), thread_id);
    }
}

struct CliChannel {
    api: Arc<RuntimeApi>,
    binding_store: Arc<InMemoryChannelBindingStore>,
}

impl CliChannel {
    fn new(api: Arc<RuntimeApi>, binding_store: Arc<InMemoryChannelBindingStore>) -> Self {
        Self { api, binding_store }
    }

    async fn submit_input(
        &self,
        session_key: &str,
        request_id: Option<String>,
        input: String,
    ) -> Result<InteractionSubmitResponse, ChannelError> {
        let thread_id = self.bind_session(session_key).await?;
        debug!(
            session_key = %session_key,
            thread_id = %thread_id,
            request_id = ?request_id,
            input_len = input.len(),
            input_preview = %log_preview(&input, 80),
            "submit_chain: cli_channel submit_input"
        );
        let request = InteractionSubmitRequest { request_id, input };
        self.api
            .submit_interaction(&thread_id, request)
            .await
            .map_err(Into::into)
    }

    #[allow(dead_code)]
    async fn query_history(
        &self,
        session_key: &str,
        limit: usize,
    ) -> Result<Vec<HistoryEventView>, ChannelError> {
        let thread_id = self.bind_session(session_key).await?;
        self.api
            .query_history(&thread_id, limit)
            .await
            .map_err(Into::into)
    }

    async fn bind_session(&self, session_key: &str) -> Result<String, ChannelError> {
        if let Some(thread_id) = self.binding_store.get_thread_id(session_key).await {
            debug!(
                session_key = %session_key,
                thread_id = %thread_id,
                "submit_chain: cli_channel reuse thread binding"
            );
            return Ok(thread_id);
        }
        let thread = self.api.create_thread(None).await?;
        debug!(
            session_key = %session_key,
            thread_id = %thread.id,
            "submit_chain: cli_channel created thread binding"
        );
        self.binding_store
            .set_thread_id(session_key, thread.id.clone())
            .await;
        Ok(thread.id)
    }
}

#[derive(Clone)]
pub struct CliRuntime {
    api: Arc<RuntimeApi>,
    channel: Arc<CliChannel>,
    thread_id: String,
    session_key: String,
}

impl CliRuntime {
    pub async fn from_api(
        api: Arc<RuntimeApi>,
        thread_id_override: Option<String>,
    ) -> Result<Self, ChannelError> {
        let binding_store = Arc::new(InMemoryChannelBindingStore::new());
        let thread_id_override_for_log = thread_id_override.clone();
        let thread = match thread_id_override {
            Some(thread_id) => api.create_thread(Some(thread_id)).await?,
            None => api.create_thread(None).await?,
        };
        debug!(
            thread_id = %thread.id,
            thread_id_override = ?thread_id_override_for_log,
            "submit_chain: cli_runtime initialized thread"
        );
        let session_key = format!("cli:{}", thread.id);
        binding_store
            .set_thread_id(&session_key, thread.id.clone())
            .await;
        let channel = Arc::new(CliChannel::new(api.clone(), binding_store));

        Ok(Self {
            api,
            channel,
            thread_id: thread.id,
            session_key,
        })
    }

    pub fn thread_id(&self) -> &str {
        &self.thread_id
    }

    pub async fn subscribe_events(&self) -> Result<broadcast::Receiver<Event>, ChannelError> {
        Ok(self.api.subscribe_events(&self.thread_id).await?)
    }

    pub async fn submit_input(
        &self,
        input: String,
    ) -> Result<InteractionSubmitResponse, ChannelError> {
        debug!(
            thread_id = %self.thread_id,
            session_key = %self.session_key,
            input_len = input.len(),
            input_preview = %log_preview(&input, 80),
            "submit_chain: cli_runtime submit_input"
        );
        self.channel
            .submit_input(&self.session_key, Some("tui".to_string()), input)
            .await
    }

    pub async fn interrupt(&self) -> Result<(), ChannelError> {
        let app = self.api.runtime_app_for_thread(&self.thread_id).await?;
        app.orchestrator.thread_runtime.cancel_all_active().await;
        Ok(())
    }
}

fn log_preview(text: &str, max_chars: usize) -> String {
    text.chars().take(max_chars).collect()
}