lash-core 0.1.0-alpha.61

use crate::plugin::{DirectCompletion, PluginError};

#[derive(Clone)]
pub struct ToolDirectCompletionClient<'run> {
    pub(super) session_id: String,
    pub(super) tool_call_id: Option<String>,
    pub(super) direct_completions: crate::DirectCompletionClient<'run>,
    pub(super) parent_invocation: Option<crate::RuntimeInvocation>,
    pub(super) parent_tool_batch_is_durable: bool,
}

impl ToolDirectCompletionClient<'_> {
    pub async fn complete(
        &self,
        mut request: crate::DirectRequest,
        usage_source: &str,
    ) -> Result<DirectCompletion, PluginError> {
        if self.parent_invocation.as_ref().is_some_and(|invocation| {
            invocation.effect_kind() == Some(crate::RuntimeEffectKind::ToolBatch)
        }) && self.parent_tool_batch_is_durable
        {
            return Err(PluginError::Session(
                "direct completions are not available inside a tool batch child".to_string(),
            ));
        }
        if request.session_id.is_none() {
            request.session_id = Some(self.session_id.clone());
        }
        if request.caused_by.is_none()
            && let Some(call_id) = self.tool_call_id.clone()
        {
            request.caused_by = Some(crate::CausalRef::ToolCall {
                session_id: self.session_id.clone(),
                call_id,
            });
        }
        self.direct_completions
            .direct_completion(request, usage_source)
            .await
    }
}