lash_core/plugin/session_obj/

directives.rs

use std::sync::Arc;

use super::*;
use crate::session_model::plugin_message_to_message;

#[derive(Clone, Copy)]
struct PluginDirectivePolicy {
    abort_error: Option<&'static str>,
    tool_directive_error: &'static str,
}

impl PluginDirectivePolicy {
    const BEFORE_TURN: Self = Self {
        abort_error: None,
        tool_directive_error: "tool directives are not valid in before_turn",
    };

    const CHECKPOINT: Self = Self {
        abort_error: None,
        tool_directive_error: "checkpoint hooks only support abort, message enqueue, session creation, events, and trace events",
    };

    const AFTER_TURN: Self = Self {
        abort_error: Some("only message enqueue and session creation are valid in after_turn"),
        tool_directive_error: "only message enqueue, session creation, events, and trace events are valid in after_turn",
    };
}

enum DirectiveAction {
    Abort(PluginAbort),
    EnqueueMessages(Vec<PluginMessage>),
    EmitRuntimeEvents(Vec<crate::SessionEvent>),
    None,
}

fn append_plugin_messages(
    messages: &mut crate::MessageSequence,
    plugin_messages: &[PluginMessage],
) {
    let new_messages = plugin_messages
        .iter()
        .filter(|message| matches!(message.role, MessageRole::User | MessageRole::System))
        .map(plugin_message_to_message)
        .collect::<Vec<_>>();
    if !new_messages.is_empty() {
        messages.extend(new_messages);
    }
}

async fn interpret_directive(
    emitted: PluginOwned<PluginDirective>,
    session_lifecycle: &Arc<dyn SessionLifecycleService>,
    session_graph: &Arc<dyn SessionGraphService>,
    policy: PluginDirectivePolicy,
) -> Result<DirectiveAction, PluginError> {
    match emitted.value {
        PluginDirective::AbortTurn { code, message } => {
            if let Some(error) = policy.abort_error {
                return Err(PluginError::Session(error.to_string()));
            }
            Ok(DirectiveAction::Abort(PluginAbort { code, message }))
        }
        PluginDirective::EnqueueMessages { messages } => {
            Ok(DirectiveAction::EnqueueMessages(messages))
        }
        PluginDirective::CreateSession { request } => {
            session_lifecycle
                .create_session(*request)
                .await
                .map_err(|err| PluginError::Session(err.to_string()))?;
            Ok(DirectiveAction::None)
        }
        PluginDirective::EmitRuntimeEvents { events: surface } => {
            Ok(DirectiveAction::EmitRuntimeEvents(
                crate::plugin::plugin_runtime_session_events(&emitted.plugin_id, surface),
            ))
        }
        PluginDirective::EmitTrace {
            name,
            payload,
            context,
        } => {
            session_graph
                .emit_trace_event(
                    *context,
                    lash_trace::TraceEvent::Custom {
                        name: format!("plugin.{}.{}", emitted.plugin_id, name),
                        payload,
                    },
                )
                .await?;
            Ok(DirectiveAction::None)
        }
        PluginDirective::ReplaceToolArgs { .. } | PluginDirective::ShortCircuitTool { .. } => Err(
            PluginError::Session(policy.tool_directive_error.to_string()),
        ),
    }
}

impl PluginSession {
    async fn apply_turn_directives(
        &self,
        directives: Vec<PluginOwned<PluginDirective>>,
        mut messages: crate::MessageSequence,
        session_lifecycle: Arc<dyn SessionLifecycleService>,
        session_graph: Arc<dyn SessionGraphService>,
        policy: PluginDirectivePolicy,
    ) -> Result<TurnPreparation, PluginError> {
        let mut events = Vec::new();
        let mut abort = None;

        for emitted in directives {
            match interpret_directive(emitted, &session_lifecycle, &session_graph, policy).await? {
                DirectiveAction::Abort(next) => abort = Some(next),
                DirectiveAction::EnqueueMessages(plugin_messages) => {
                    append_plugin_messages(&mut messages, &plugin_messages);
                }
                DirectiveAction::EmitRuntimeEvents(next_events) => events.extend(next_events),
                DirectiveAction::None => {}
            }
        }

        Ok(TurnPreparation {
            messages,
            events,
            abort,
        })
    }

    pub async fn prepare_turn(
        &self,
        request: PrepareTurnRequest,
    ) -> Result<TurnPreparation, PluginError> {
        self.prepare_turn_with_phase_probe(request, None).await
    }

    pub async fn prepare_turn_with_phase_probe(
        &self,
        request: PrepareTurnRequest,
        phase_probe: Option<Arc<dyn crate::runtime::RuntimeTurnPhaseProbe>>,
    ) -> Result<TurnPreparation, PluginError> {
        let PrepareTurnRequest {
            session_id,
            state,
            messages,
            sessions,
            session_lifecycle,
            session_graph,
            turn_context,
        } = request;
        let directives = self
            .before_turn_with_phase_probe(
                TurnHookContext {
                    session_id,
                    state,
                    sessions,
                    turn_context,
                },
                phase_probe.as_ref(),
            )
            .await?;
        self.apply_turn_directives(
            directives,
            messages,
            session_lifecycle,
            session_graph,
            PluginDirectivePolicy::BEFORE_TURN,
        )
        .await
    }

    pub async fn apply_checkpoint(
        &self,
        ctx: CheckpointHookContext,
    ) -> Result<CheckpointApplication, PluginError> {
        let directives = self.at_checkpoint(ctx.clone()).await?;
        let mut messages = Vec::new();
        let mut events = Vec::new();
        let mut abort = None;

        for emitted in directives {
            match interpret_directive(
                emitted,
                &ctx.session_lifecycle,
                &ctx.session_graph,
                PluginDirectivePolicy::CHECKPOINT,
            )
            .await?
            {
                DirectiveAction::Abort(next) => abort = Some(next),
                DirectiveAction::EnqueueMessages(queued) => messages.extend(queued),
                DirectiveAction::EmitRuntimeEvents(next_events) => events.extend(next_events),
                DirectiveAction::None => {}
            }
        }

        Ok(CheckpointApplication {
            messages,
            events,
            abort,
        })
    }

    pub async fn finalize_turn(
        &self,
        turn: AssembledTurn,
        sessions: Arc<dyn SessionStateService>,
        session_lifecycle: Arc<dyn SessionLifecycleService>,
        session_graph: Arc<dyn SessionGraphService>,
    ) -> Result<TurnFinalization, PluginError> {
        self.finalize_turn_with_phase_probe(turn, sessions, session_lifecycle, session_graph, None)
            .await
    }

    pub async fn finalize_turn_with_phase_probe(
        &self,
        mut turn: AssembledTurn,
        sessions: Arc<dyn SessionStateService>,
        session_lifecycle: Arc<dyn SessionLifecycleService>,
        session_graph: Arc<dyn SessionGraphService>,
        phase_probe: Option<Arc<dyn crate::runtime::RuntimeTurnPhaseProbe>>,
    ) -> Result<TurnFinalization, PluginError> {
        let session_id = turn.state.session_id.clone();
        let directives = if self.contributions.after_turn_hooks.is_empty() {
            Vec::new()
        } else {
            self.after_turn_with_phase_probe(
                TurnResultHookContext {
                    session_id: session_id.clone(),
                    turn: Arc::new(crate::plugin::TurnResultSummary::from_assembled(&turn)),
                    sessions,
                },
                phase_probe.as_ref(),
            )
            .await?
        };
        let mut events = Vec::new();
        let mut updated_messages: Option<crate::MessageSequence> = None;
        for emitted in directives {
            match interpret_directive(
                emitted,
                &session_lifecycle,
                &session_graph,
                PluginDirectivePolicy::AFTER_TURN,
            )
            .await?
            {
                DirectiveAction::Abort(_) => unreachable!("after_turn policy rejects abort"),
                DirectiveAction::EnqueueMessages(plugin_messages) => {
                    let messages = updated_messages.get_or_insert_with(|| {
                        crate::MessageSequence::from_base(
                            turn.state.read_view().messages().to_vec().into(),
                        )
                    });
                    append_plugin_messages(messages, &plugin_messages);
                }
                DirectiveAction::EmitRuntimeEvents(next_events) => events.extend(next_events),
                DirectiveAction::None => {}
            }
        }
        if let Some(messages) = updated_messages.as_ref() {
            turn.state.replace_active_read_state(messages.as_slice());
        }

        if self.has_runtime_event_hooks() {
            self.emit_runtime_event_with_phase_probe(
                PluginLifecycleEvent::TurnFinalized(Arc::new(turn.clone())),
                phase_probe,
            )
            .await;
        }

        Ok(TurnFinalization { turn, events })
    }
}