lash_core/plugin/
protocol.rs

1//! Protocol-plugin traits and narrow session/runtime context wrappers.
2//!
3//! Protocol plugins register their implementations here; the runtime narrows
4//! what a protocol plugin can poke at so external crates don't need direct access to
5//! `Session` / `LashRuntime` internals.
6//!
7//! Split out of `plugin/mod.rs` for file size; `pub use` there keeps
8//! the outer module path.
9
10use std::collections::BTreeMap;
11use std::sync::Arc;
12
13use serde::de::DeserializeOwned;
14use serde::{Deserialize, Serialize};
15
16use super::SessionAppendNode;
17use crate::runtime::RuntimeSessionState;
18use crate::{
19    ExecRequest, ExecResponse, LlmRequest, PromptUsage, RuntimeExecutionContext, SessionReadView,
20};
21
22/// Session-scoped plugin that initializes, restores, and extends protocol
23/// state across a session's lifecycle. External protocol crates implement
24/// this via context wrappers ([`ProtocolSessionContext`],
25/// [`ProtocolRuntimeContext`]) so they don't need direct access to
26/// `Session`/`LashRuntime` internals — the context narrows what a
27/// plugin can poke at to the capabilities any protocol reasonably needs.
28#[async_trait::async_trait]
29pub trait ProtocolSessionPlugin: Send + Sync {
30    async fn initialize_session(
31        &self,
32        _ctx: ProtocolSessionContext<'_>,
33    ) -> Result<(), crate::SessionError> {
34        Ok(())
35    }
36
37    async fn restore_session(
38        &self,
39        _ctx: ProtocolSessionContext<'_>,
40        _state: &RuntimeSessionState,
41    ) -> Result<(), crate::SessionError> {
42        Ok(())
43    }
44
45    async fn append_session_nodes(
46        &self,
47        _ctx: ProtocolSessionContext<'_>,
48        _nodes: &[SessionAppendNode],
49    ) -> Result<(), crate::SessionError> {
50        Ok(())
51    }
52
53    async fn apply_session_extension(
54        &self,
55        _extension: crate::ProtocolSessionExtensionHandle,
56    ) -> Result<(), crate::SessionError> {
57        Err(crate::SessionError::Protocol(
58            "protocol does not accept session extensions".to_string(),
59        ))
60    }
61
62    async fn validate_turn_extension(
63        &self,
64        _extension: &crate::ProtocolTurnExtensionHandle,
65    ) -> Result<(), crate::SessionError> {
66        Ok(())
67    }
68
69    /// Fires on every session materialization — root/builder open (including
70    /// resume) and child create — so a protocol plugin can apply and default
71    /// its per-session options at open time (apply-at-open semantics).
72    ///
73    /// The [`ProtocolSessionMaterialization`] descriptor carries the
74    /// plugin-keyed options that reached this materialization (builder options
75    /// for root opens, request options for child create) and whether this is a
76    /// root session. The plugin reads/writes durable protocol turn options
77    /// through [`ProtocolRuntimeContext`].
78    fn configure_runtime_on_materialize(
79        &self,
80        _ctx: ProtocolRuntimeContext<'_>,
81        _materialization: ProtocolSessionMaterialization<'_>,
82    ) -> Result<(), crate::SessionError> {
83        Ok(())
84    }
85
86    async fn before_llm_call(
87        &self,
88        _ctx: ProtocolBeforeLlmCallContext<'_>,
89        _request: &LlmRequest,
90    ) -> Result<Option<ProtocolLlmCallAction>, crate::PluginError> {
91        Ok(None)
92    }
93}
94
95/// Narrow wrapper around `Session` that protocol plugins use to
96/// initialize, restore, and extend their per-session state.
97///
98/// Exposes only generic per-session lifecycle capabilities. Protocol-local
99/// execution state is owned by the protocol plugin itself and is accessed
100/// through [`ProtocolSessionPlugin`] callbacks.
101/// Prevents protocol plugins from reaching into unrelated `Session`
102/// internals.
103pub struct ProtocolSessionContext<'a> {
104    session_id: &'a str,
105}
106
107impl<'a> ProtocolSessionContext<'a> {
108    pub(crate) fn new(_session: &'a mut crate::Session, session_id: &'a str) -> Self {
109        Self { session_id }
110    }
111
112    /// ID of the session being initialized/restored. Equivalent to the
113    /// `session_id` previously passed as a separate argument.
114    pub fn session_id(&self) -> &str {
115        self.session_id
116    }
117}
118
119pub struct ProtocolBeforeLlmCallContext<'run> {
120    pub session_id: String,
121    pub sessions: Arc<dyn crate::plugin::SessionStateService>,
122    pub session_graph: Arc<dyn crate::plugin::SessionGraphService>,
123    pub processes: Arc<dyn crate::ProcessService>,
124    pub state: SessionReadView,
125    pub latest_prompt_usage: Option<PromptUsage>,
126    pub(crate) direct_completions: crate::DirectCompletionClient<'run>,
127    pub(crate) process_parent_invocation: crate::RuntimeInvocation,
128    pub(crate) effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
129}
130
131impl ProtocolBeforeLlmCallContext<'_> {
132    pub async fn direct_llm_completion(
133        &self,
134        request: crate::LlmRequest,
135        usage_source: &str,
136    ) -> Result<crate::DirectLlmCompletion, crate::PluginError> {
137        self.direct_completions
138            .direct_llm_completion(request, usage_source)
139            .await
140    }
141
142    pub fn process_scope(&self) -> crate::ProcessOpScope<'_> {
143        crate::ProcessOpScope::new(self.effect_controller.scoped())
144            .with_parent_invocation(Some(self.process_parent_invocation.clone()))
145    }
146}
147
148#[derive(Clone, Debug, PartialEq, Eq)]
149pub enum ProtocolLlmCallAction {
150    SwitchAgentFrame { frame_id: String, task: String },
151}
152
153/// Narrow wrapper around `LashRuntime` that protocol plugins use when
154/// configuring the runtime from a fresh `SessionCreateRequest`.
155///
156/// Exposes only the runtime-level capabilities protocols need to set
157/// (termination contract, etc.) so plugins don't reach into unrelated
158/// runtime internals.
159pub struct ProtocolRuntimeContext<'a> {
160    runtime: &'a mut crate::runtime::LashRuntime,
161}
162
163impl<'a> ProtocolRuntimeContext<'a> {
164    pub(crate) fn new(runtime: &'a mut crate::runtime::LashRuntime) -> Self {
165        Self { runtime }
166    }
167
168    /// The durable protocol turn options currently recorded on the session.
169    /// Protocol plugins read these to preserve fields (e.g. termination) they
170    /// are not overwriting.
171    pub fn protocol_turn_options(&self) -> &crate::ProtocolTurnOptions {
172        self.runtime.protocol_turn_options()
173    }
174
175    /// Set the durable protocol turn options and mirror them to the current
176    /// agent frame only.
177    pub fn set_protocol_turn_options(&mut self, options: crate::ProtocolTurnOptions) {
178        self.runtime.set_protocol_turn_options(options);
179    }
180
181    /// Set the durable protocol turn options and mirror them to **every** agent
182    /// frame. Apply-at-open semantics: the last applied value is recorded on the
183    /// session and all frames.
184    pub fn set_protocol_turn_options_all_frames(&mut self, options: crate::ProtocolTurnOptions) {
185        self.runtime.set_protocol_turn_options_all_frames(options);
186    }
187}
188
189/// Read-only descriptor of a session materialization handed to
190/// [`ProtocolSessionPlugin::configure_runtime_on_materialize`].
191pub struct ProtocolSessionMaterialization<'a> {
192    /// Plugin-keyed options that reached this materialization: builder options
193    /// for a root/builder open, request options for a child create.
194    pub plugin_options: &'a PluginOptions,
195    /// Whether this materialization is a root session (no parent).
196    pub is_root_session: bool,
197}
198
199#[async_trait::async_trait]
200pub trait CodeExecutorPlugin: Send + Sync {
201    async fn execute_code(
202        &self,
203        ctx: RuntimeExecutionContext<'_>,
204        request: ExecRequest,
205    ) -> Result<ExecResponse, crate::SessionError>;
206
207    fn execution_state_dirty(&self) -> bool {
208        false
209    }
210
211    async fn snapshot_execution_state(
212        &self,
213        _ctx: ProtocolSessionContext<'_>,
214    ) -> Result<Option<Vec<u8>>, crate::SessionError> {
215        Ok(None)
216    }
217
218    async fn restore_execution_state(
219        &self,
220        _ctx: ProtocolSessionContext<'_>,
221        _data: &[u8],
222    ) -> Result<(), crate::SessionError> {
223        Ok(())
224    }
225}
226
227pub trait AssistantProseProjectorPlugin: Send + Sync {
228    fn project_assistant_prose(&self, text: &str) -> String;
229}
230
231/// Singleton kernel extension slot that owns the `ProtocolDriverHandle` and
232/// associated preamble (prompt text, tool catalog, sync/async flag) for this
233/// session.
234///
235/// Core owns the slot and the `HostTurnProtocol` state shape so the turn loop
236/// can persist and resume protocol driver state generically. External protocol
237/// crates own the concrete prompt policy and output parser. Plugin stack
238/// construction must install exactly one implementation.
239pub trait ProtocolDriverPlugin: Send + Sync {
240    /// Build the `TurnDriverPreamble` (driver handle + prompt text + tool
241    /// surface metadata) for a turn.
242    fn build_preamble(&self, input: crate::ProtocolBuildInput) -> crate::TurnDriverPreamble;
243}
244
245/// Plugin-owned options carried on a `SessionCreateRequest`.
246#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
247pub struct PluginOptions {
248    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
249    pub plugins: BTreeMap<String, serde_json::Value>,
250}
251
252impl PluginOptions {
253    pub fn empty() -> Self {
254        Self::default()
255    }
256
257    pub fn typed<T>(plugin_id: impl Into<String>, extras: T) -> Result<Self, serde_json::Error>
258    where
259        T: Serialize,
260    {
261        let mut options = Self::default();
262        options.insert_typed(plugin_id, extras)?;
263        Ok(options)
264    }
265
266    pub fn insert_typed<T>(
267        &mut self,
268        plugin_id: impl Into<String>,
269        extras: T,
270    ) -> Result<(), serde_json::Error>
271    where
272        T: Serialize,
273    {
274        self.plugins
275            .insert(plugin_id.into(), serde_json::to_value(extras)?);
276        Ok(())
277    }
278
279    pub fn decode<T>(&self, plugin_id: &str) -> Result<Option<T>, serde_json::Error>
280    where
281        T: DeserializeOwned,
282    {
283        self.plugins
284            .get(plugin_id)
285            .cloned()
286            .map(serde_json::from_value)
287            .transpose()
288    }
289}
lash_core/plugin/protocol.rs

lash_core/plugin/
protocol.rs
