use crate::agent::Session;
use crate::agent::agdx::AgdxReceipt;
use crate::context::{ContextMessage, ContextPolicy};
use crate::error::LaserError;
use crate::provenance::AgentTopic;
use crate::types::MintUlid;
use iggy::prelude::{Identifier, StreamClient, TopicClient};
use laser_wire::agent::{
AgentErrorBody, ContextManifest, CorrelationId, Fragment, METADATA_DURATION_MICROS,
METADATA_PROVIDER_NAME, METADATA_REQUEST_MODEL, METADATA_RESPONSE_MODEL, OPERATION_CHAT,
OPERATION_STATE_DELTA, OPERATION_STATE_SNAPSHOT, SessionRef, StateDelta, StateSnapshot,
TokenUsage, apply_json_patch, estimate_tokens,
};
use laser_wire::content::ContentType;
use laser_wire::dispatch::{OPERATION_CONTEXT_ASSEMBLED, OPERATION_EXECUTE_TOOL};
use laser_wire::framing::{decode_named, encode_named};
use laser_wire::graph::{ProducerInfo, SourceRef};
use laser_wire::query::Value;
use laser_wire::session::SessionStateView;
use std::collections::BTreeMap;
use std::time::{Duration, Instant};
const REDACTED_KEYS: [&str; 6] = [
"authorization",
"api_key",
"token",
"password",
"secret",
"cookie",
];
const SNAPSHOT_EVERY_DELTAS: u32 = 64;
pub fn default_redact(value: &mut serde_json::Value) {
match value {
serde_json::Value::Object(fields) => {
for (key, field) in fields.iter_mut() {
if REDACTED_KEYS.contains(&key.to_ascii_lowercase().as_str()) {
*field = serde_json::Value::String("[redacted]".to_owned());
} else {
default_redact(field);
}
}
}
serde_json::Value::Array(items) => items.iter_mut().for_each(default_redact),
_ => {}
}
}
#[derive(Default)]
pub(crate) struct StateCursor {
revision: u64,
document: Option<serde_json::Value>,
since_snapshot: u32,
}
#[derive(Debug, Clone)]
pub struct ModelRequest {
pub model: String,
pub provider: Option<String>,
pub operation: String,
pub body: Vec<u8>,
}
impl ModelRequest {
pub fn new(model: impl Into<String>, body: impl Into<Vec<u8>>) -> Self {
Self {
model: model.into(),
provider: None,
operation: OPERATION_CHAT.to_owned(),
body: body.into(),
}
}
#[must_use]
pub fn provider(mut self, provider: impl Into<String>) -> Self {
self.provider = Some(provider.into());
self
}
#[must_use]
pub fn operation(mut self, operation: impl Into<String>) -> Self {
self.operation = operation.into();
self
}
}
#[derive(Debug, Clone, Default)]
pub struct ModelResponse {
pub body: Vec<u8>,
pub model: Option<String>,
pub finish_reason: Option<String>,
pub usage: Option<TokenUsage>,
pub duration: Option<Duration>,
}
#[derive(Debug, Clone)]
pub struct AssembledContext {
pub fragments: Vec<ContextMessage>,
pub manifest: ContextManifest,
}
impl AssembledContext {
pub fn text(&self) -> String {
self.fragments
.iter()
.map(|fragment| String::from_utf8_lossy(&fragment.payload).into_owned())
.collect::<Vec<_>>()
.join("\n")
}
}
pub struct ModelCall {
session: Session,
correlation: CorrelationId,
operation: String,
started: Instant,
}
impl ModelCall {
pub fn correlation(&self) -> CorrelationId {
self.correlation
}
pub async fn complete(self, response: ModelResponse) -> Result<AgdxReceipt, LaserError> {
let duration = response.duration.unwrap_or_else(|| self.started.elapsed());
self.session
.write_result(
self.correlation,
&self.operation,
None,
response.body,
response.usage,
duration,
response.model,
response.finish_reason,
)
.await
}
pub async fn fail(self, error: AgentErrorBody) -> Result<AgdxReceipt, LaserError> {
self.session
.write_failure(self.correlation, &self.operation, None, &error)
.await
}
}
pub struct ToolCall {
session: Session,
correlation: CorrelationId,
tool: String,
started: Instant,
}
impl ToolCall {
pub fn correlation(&self) -> CorrelationId {
self.correlation
}
pub async fn complete(self, result: impl Into<Vec<u8>>) -> Result<AgdxReceipt, LaserError> {
let duration = self.started.elapsed();
self.session
.write_result(
self.correlation,
OPERATION_EXECUTE_TOOL,
Some(&self.tool),
result.into(),
None,
duration,
None,
None,
)
.await
}
pub async fn fail(self, error: AgentErrorBody) -> Result<AgdxReceipt, LaserError> {
self.session
.write_failure(
self.correlation,
OPERATION_EXECUTE_TOOL,
Some(&self.tool),
&error,
)
.await
}
}
pub struct SessionState {
session: Session,
}
impl SessionState {
pub async fn set(
&self,
key: &str,
value: serde_json::Value,
) -> Result<AgdxReceipt, LaserError> {
let path = format!("/{}", key.replace('~', "~0").replace('/', "~1"));
let operation = serde_json::from_value(serde_json::json!({
"op": "add",
"path": path,
"value": value,
}))
.map_err(|error| LaserError::Invalid(format!("state patch: {error}")))?;
self.patch(vec![operation]).await
}
pub async fn patch(
&self,
patch: Vec<laser_wire::agent::PatchOperation>,
) -> Result<AgdxReceipt, LaserError> {
self.seed().await?;
let (delta, due) = {
let mut cursor = self.session.state.lock().expect("session state lock");
let base = cursor.document.clone().unwrap_or(serde_json::json!({}));
let next = apply_json_patch(&base, &patch)
.map_err(|error| LaserError::Invalid(error.to_string()))?;
let delta = StateDelta {
base_revision: cursor.revision,
patch,
op_id: ulid::Ulid::generate().to_string(),
};
cursor.revision += 1;
cursor.document = Some(next);
cursor.since_snapshot += 1;
(delta, cursor.since_snapshot >= SNAPSHOT_EVERY_DELTAS)
};
let receipt = self
.session
.write_event(OPERATION_STATE_DELTA, encode_named(&delta)?, None)
.await?;
if due {
self.snapshot().await?;
}
Ok(receipt)
}
pub async fn replace(&self, document: serde_json::Value) -> Result<AgdxReceipt, LaserError> {
let serde_json::Value::Object(next) = document else {
return Err(LaserError::Invalid(
"a state document is a JSON object".to_owned(),
));
};
self.seed().await?;
let current = self
.session
.state
.lock()
.expect("session state lock")
.document
.clone()
.unwrap_or(serde_json::json!({}));
let pointer = |key: &str| format!("/{}", key.replace('~', "~0").replace('/', "~1"));
let mut operations = Vec::new();
if let serde_json::Value::Object(current) = ¤t {
for key in current.keys().filter(|key| !next.contains_key(*key)) {
operations.push(serde_json::json!({"op": "remove", "path": pointer(key)}));
}
}
for (key, value) in next {
operations
.push(serde_json::json!({"op": "add", "path": pointer(&key), "value": value}));
}
let patch = serde_json::from_value(serde_json::Value::Array(operations))
.map_err(|error| LaserError::Invalid(format!("state patch: {error}")))?;
self.patch(patch).await
}
pub async fn snapshot(&self) -> Result<AgdxReceipt, LaserError> {
self.seed().await?;
let snapshot = {
let cursor = self.session.state.lock().expect("session state lock");
StateSnapshot {
base_revision: cursor.revision,
document: cursor.document.clone().unwrap_or(serde_json::json!({})),
}
};
let receipt = self
.session
.write_event(OPERATION_STATE_SNAPSHOT, encode_named(&snapshot)?, None)
.await?;
let mut cursor = self.session.state.lock().expect("session state lock");
if cursor.revision == snapshot.base_revision {
cursor.since_snapshot = 0;
}
Ok(receipt)
}
async fn seed(&self) -> Result<(), LaserError> {
if self
.session
.state
.lock()
.expect("session state lock")
.document
.is_some()
{
return Ok(());
}
let view = self.current_view().await?;
let mut cursor = self.session.state.lock().expect("session state lock");
if cursor.document.is_none() {
cursor.revision = view.revision;
cursor.document = Some(view.document);
}
Ok(())
}
pub(crate) async fn current_view(&self) -> Result<SessionStateView, LaserError> {
let lane = self.get().await?;
if lane.complete || !self.session.laser.capabilities().await.sessions {
return Ok(lane);
}
match self
.session
.laser
.sessions()
.state(self.session.conversation(), 0)
.await
{
Ok(view) if view.revision >= lane.revision => Ok(view),
_ => Ok(lane),
}
}
pub(crate) async fn snapshot_if_changed(&self) -> Result<(), LaserError> {
let changed = self
.session
.state
.lock()
.expect("session state lock")
.since_snapshot
> 0;
if changed {
self.snapshot().await?;
}
Ok(())
}
pub async fn get(&self) -> Result<SessionStateView, LaserError> {
let records = self
.session
.scope()
.fetch_with(
vec![AgentTopic::Sessions],
Box::new(crate::context::LastN(usize::MAX)),
)
.await?;
let mut document = serde_json::json!({});
let mut revision = 0u64;
let mut complete = false;
for record in records {
let Some(envelope) = record.envelope else {
continue;
};
match envelope.operation.as_deref() {
Some(OPERATION_STATE_SNAPSHOT) => {
let snapshot: StateSnapshot = decode_named(&envelope.body)?;
if complete && snapshot.base_revision != revision {
continue;
}
document = snapshot.document;
revision = snapshot.base_revision;
complete = true;
}
Some(OPERATION_STATE_DELTA) => {
let delta: StateDelta = decode_named(&envelope.body)?;
if delta.base_revision == 0 && revision == 0 {
complete = true;
}
if delta.base_revision != revision {
continue;
}
if let Ok(next) = apply_json_patch(&document, &delta.patch) {
document = next;
revision += 1;
}
}
_ => {}
}
}
Ok(SessionStateView {
revision,
document,
history: Vec::new(),
frontier: None,
complete,
})
}
}
impl Session {
pub fn reference(&self) -> Result<SessionRef, LaserError> {
Ok(SessionRef {
stream: self.stream()?.to_owned(),
session: self.conversation().into(),
})
}
#[must_use]
pub fn redact(
mut self,
redactor: impl Fn(&mut serde_json::Value) + Send + Sync + 'static,
) -> Self {
self.redactor = std::sync::Arc::new(redactor);
self
}
#[must_use]
pub fn acting_on(mut self, source: SourceRef) -> Self {
self.current = Some(source);
self
}
pub async fn cancel_requested(&self) -> Result<bool, LaserError> {
let records = self
.scope()
.fetch_with(
vec![AgentTopic::Control],
Box::new(crate::context::LastN(usize::MAX)),
)
.await?;
Ok(records.iter().any(|record| {
record.envelope.as_ref().is_some_and(|envelope| {
envelope.operation.as_deref()
== Some(laser_wire::dispatch::OPERATION_SESSION_CANCEL)
|| envelope.task_state == Some(laser_wire::agent::TaskState::Canceled)
})
}))
}
pub async fn pending_control(&self) -> Result<crate::agent::PendingControl, LaserError> {
Ok(self.control_state().await?.flags)
}
pub(crate) async fn control_state(
&self,
) -> Result<crate::agent::control::ControlState, LaserError> {
let conversation = self.conversation();
if let Some(state) = self
.control
.as_ref()
.and_then(|book| book.state(conversation))
{
return Ok(state);
}
let me = self.agent.clone();
let mut records = self
.scope()
.fetch_with(
vec![AgentTopic::Control],
Box::new(crate::context::LastN(usize::MAX)),
)
.await?;
records.sort_by_key(|record| {
(
record.timestamp_micros,
record.id.partition_id,
record.id.offset,
)
});
let route = self.control_route().await?;
let (state, last) = crate::agent::control::fold(
records.iter().filter_map(|record| {
let request =
crate::agent::control::control_request(record.envelope.as_ref()?, me.as_ref())?;
Some(((record.id.partition_id, record.id.offset), request))
}),
route,
);
Ok(match &self.control {
Some(book) => book.install(conversation, state, last),
None => state,
})
}
async fn control_route(&self) -> Result<Option<u32>, LaserError> {
let stream = iggy::prelude::Identifier::named(self.laser.stream_required()?)?;
let topic = iggy::prelude::Identifier::named(laser_wire::topics::AGENT_CONTROL)?;
let details =
iggy::prelude::TopicClient::get_topic(self.laser.client().as_ref(), &stream, &topic)
.await?;
Ok(details.and_then(|details| {
crate::agent::control::control_partition(self.conversation(), details.partitions_count)
}))
}
pub fn state(&self) -> SessionState {
SessionState {
session: self.clone(),
}
}
pub async fn assemble(
&self,
policy: Box<dyn ContextPolicy>,
) -> Result<AssembledContext, LaserError> {
let name = policy.name();
let version = policy.version();
let fragments = self
.scope()
.fetch_with(vec![AgentTopic::Sessions], policy)
.await?;
let (stream_id, topic_id, generation) = self.lane_identity().await?;
let mut tokens = 0u64;
let mut bytes = 0u64;
let mut frontier: BTreeMap<(u32, u32), u64> = BTreeMap::new();
let manifest_fragments = fragments
.iter()
.map(|message| {
let size = message.payload.len();
tokens += estimate_tokens(size);
bytes += size as u64;
let position = frontier
.entry((topic_id, message.id.partition_id))
.or_insert(message.id.offset);
*position = (*position).max(message.id.offset);
Fragment::Message {
at: SourceRef::Message {
stream: stream_id,
topic: topic_id,
partition: message.id.partition_id,
offset: message.id.offset,
generation: Some(generation),
conversation: Some(self.conversation().to_string()),
},
tokens: u32::try_from(estimate_tokens(size)).unwrap_or(u32::MAX),
bytes: u32::try_from(size).unwrap_or(u32::MAX),
}
})
.collect();
Ok(AssembledContext {
fragments,
manifest: ContextManifest {
policy: name,
policy_version: version,
fragments: manifest_fragments,
tokens,
bytes,
frontier: frontier
.into_iter()
.map(|((topic, partition), offset)| (topic, partition, offset))
.collect(),
correlation: None,
},
})
}
pub async fn model(
&self,
request: ModelRequest,
assembled: Option<&AssembledContext>,
) -> Result<ModelCall, LaserError> {
let correlation = CorrelationId::mint();
self.write_model_request(correlation, &request, assembled)
.await?;
Ok(ModelCall {
session: self.clone(),
correlation,
operation: request.operation,
started: Instant::now(),
})
}
pub async fn tool(
&self,
name: impl Into<String>,
mut args: serde_json::Value,
) -> Result<ToolCall, LaserError> {
let tool = name.into();
(self.redactor)(&mut args);
let correlation = CorrelationId::mint();
let agent = self.require_agent()?;
let body = serde_json::to_vec(&args)
.map_err(|error| LaserError::Codec(format!("tool arguments: {error}")))?;
self.lane()?
.command(correlation, body)
.with_operation(OPERATION_EXECUTE_TOOL)
.with_tool(tool.clone())
.with_target(agent)
.content_type(ContentType::Json)
.send()
.await?;
Ok(ToolCall {
session: self.clone(),
correlation,
tool,
started: Instant::now(),
})
}
pub async fn record_model_call(
&self,
request: ModelRequest,
response: ModelResponse,
assembled: Option<&AssembledContext>,
) -> Result<AgdxReceipt, LaserError> {
let correlation = CorrelationId::mint();
self.write_model_request(correlation, &request, assembled)
.await?;
let duration = response.duration;
self.write_result(
correlation,
&request.operation,
None,
response.body,
response.usage,
duration.unwrap_or_default(),
response.model,
response.finish_reason,
)
.await
}
#[cfg(feature = "kv")]
pub fn kv(&self, namespace: impl Into<String>) -> Result<crate::kv::Kv, LaserError> {
Ok(self.laser.kv(namespace).in_session(self.reference()?))
}
#[cfg(feature = "graph")]
pub fn linked_graph(
&self,
name: impl Into<String>,
) -> Result<crate::graph::GraphHandle<'_>, LaserError> {
let mut graph = self
.laser
.graph(name)
.in_session(self.reference()?)
.produced_by(self.producer());
if let Some(source) = &self.current {
graph = graph.sourced_from(source.clone());
}
Ok(graph)
}
pub async fn record_retrieval(
&self,
query: Option<String>,
items: &[crate::memory::MemoryItem],
) -> Result<AgdxReceipt, LaserError> {
let retrieval = laser_wire::agent::ContextRetrieval {
query,
items: items
.iter()
.map(|item| (item.id.to_string(), item.score.unwrap_or_default()))
.collect(),
};
self.write_event(
laser_wire::dispatch::OPERATION_CONTEXT_RETRIEVED,
encode_named(&retrieval)?,
None,
)
.await
}
pub async fn record_compaction(
&self,
compaction: laser_wire::agent::ContextCompaction,
) -> Result<AgdxReceipt, LaserError> {
self.write_event(
laser_wire::dispatch::OPERATION_CONTEXT_COMPACTED,
encode_named(&compaction)?,
None,
)
.await
}
pub fn linked_memory(&self) -> crate::context_scope::ScopedMemory {
self.memory()
.with_lineage(self.current.clone(), Some(self.producer()))
}
fn producer(&self) -> ProducerInfo {
ProducerInfo {
name: match &self.agent {
Some(agent) => format!("sdk:{agent}"),
None => "sdk".to_owned(),
},
version: self.config().sdk_info().version.clone(),
}
}
fn require_agent(&self) -> Result<laser_wire::agent::AgentId, LaserError> {
self.agent
.clone()
.ok_or_else(|| LaserError::Invalid("this session handle has no agent".to_owned()))
}
async fn lane_identity(&self) -> Result<(u32, u32, u64), LaserError> {
let stream = Identifier::named(self.stream()?)?;
let client = self.laser.client();
let stream_id = client
.get_stream(&stream)
.await?
.ok_or_else(|| LaserError::Invalid("the session stream does not exist".to_owned()))?
.id;
let topic = client
.get_topic(&stream, &AgentTopic::Sessions.as_identifier())
.await?
.ok_or_else(|| LaserError::Invalid("agent.sessions does not exist".to_owned()))?;
Ok((stream_id, topic.id, topic.created_at.as_micros()))
}
async fn write_model_request(
&self,
correlation: CorrelationId,
request: &ModelRequest,
assembled: Option<&AssembledContext>,
) -> Result<(), LaserError> {
let agent = self.require_agent()?;
let mut body = request.body.clone();
if let Ok(mut json) = serde_json::from_slice::<serde_json::Value>(&body) {
(self.redactor)(&mut json);
body = serde_json::to_vec(&json)
.map_err(|error| LaserError::Codec(format!("model request: {error}")))?;
}
let lane = self.lane()?;
let mut send = lane
.command(correlation, body)
.with_operation(request.operation.clone())
.with_target(agent)
.with_metadata(METADATA_REQUEST_MODEL, Value::Str(request.model.clone()));
if let Some(provider) = &request.provider {
send = send.with_metadata(METADATA_PROVIDER_NAME, Value::Str(provider.clone()));
}
send.send().await?;
if let Some(assembled) = assembled {
let mut manifest = assembled.manifest.clone();
manifest.correlation = Some(correlation);
self.write_event(
OPERATION_CONTEXT_ASSEMBLED,
encode_named(&manifest)?,
Some(correlation),
)
.await?;
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn write_result(
&self,
correlation: CorrelationId,
operation: &str,
tool: Option<&str>,
body: Vec<u8>,
usage: Option<TokenUsage>,
duration: Duration,
model: Option<String>,
finish_reason: Option<String>,
) -> Result<AgdxReceipt, LaserError> {
let agent = self.require_agent()?;
let lane = self.lane()?;
let mut send = lane
.respond(correlation, body)
.with_operation(operation)
.with_target(agent)
.with_metadata(
METADATA_DURATION_MICROS,
Value::Uint(u64::try_from(duration.as_micros()).unwrap_or(u64::MAX)),
);
if let Some(tool) = tool {
send = send.with_tool(tool);
}
if let Some(usage) = usage {
send = send.with_usage(usage);
}
if let Some(model) = model {
send = send.with_metadata(METADATA_RESPONSE_MODEL, Value::Str(model));
}
if let Some(reason) = finish_reason {
send = send.with_finish_reason(reason);
}
send.send_receipt().await
}
async fn write_failure(
&self,
correlation: CorrelationId,
operation: &str,
tool: Option<&str>,
error: &AgentErrorBody,
) -> Result<AgdxReceipt, LaserError> {
let agent = self.require_agent()?;
let lane = self.lane()?;
let mut send = lane
.fail(correlation, error)?
.with_operation(operation)
.with_target(agent);
if let Some(tool) = tool {
send = send.with_tool(tool);
}
send.send_receipt().await
}
pub(crate) async fn write_event(
&self,
operation: &str,
body: Vec<u8>,
correlation: Option<CorrelationId>,
) -> Result<AgdxReceipt, LaserError> {
let lane = self.lane()?;
let mut send = lane
.emit(body)
.with_operation(operation)
.content_type(ContentType::Cbor);
if let Some(correlation) = correlation {
send = send.with_correlation(correlation);
}
send.send_receipt().await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn given_secret_keys_at_any_depth_when_redacted_then_should_drop_only_their_values() {
let mut args = serde_json::json!({
"query": "status",
"Authorization": "Bearer x",
"nested": {"api_key": "k", "keep": 1},
"list": [{"password": "p"}],
});
default_redact(&mut args);
assert_eq!(args["query"], "status");
assert_eq!(args["Authorization"], "[redacted]");
assert_eq!(args["nested"]["api_key"], "[redacted]");
assert_eq!(args["nested"]["keep"], 1);
assert_eq!(args["list"][0]["password"], "[redacted]");
}
}
#[must_use = "a submission does nothing until send().await"]
pub struct SubmitBuilder {
sessions: crate::agent::Sessions,
agent: laser_wire::agent::AgentId,
input: Vec<u8>,
submitter: Option<laser_wire::agent::AgentId>,
operation: String,
label: Option<String>,
namespace: Option<String>,
budget: Option<laser_wire::agent::Budget>,
tags: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Submitted {
pub session: crate::types::ConversationId,
pub correlation: CorrelationId,
}
impl SubmitBuilder {
pub fn from(mut self, submitter: impl Into<laser_wire::agent::AgentId>) -> Self {
let submitter = submitter.into();
self.submitter = Some(submitter);
self
}
pub fn operation(mut self, operation: impl Into<String>) -> Self {
self.operation = operation.into();
self
}
pub fn label(mut self, label: impl Into<String>) -> Self {
self.label = Some(label.into());
self
}
pub fn namespace(mut self, namespace: impl Into<String>) -> Self {
self.namespace = Some(namespace.into());
self
}
pub fn budget(mut self, budget: laser_wire::agent::Budget) -> Self {
self.budget = Some(budget);
self
}
pub fn tag(mut self, tag: impl Into<String>) -> Self {
self.tags.push(tag.into());
self
}
pub async fn send(self) -> Result<Submitted, LaserError> {
let submitter = self.submitter.ok_or_else(|| {
LaserError::Invalid("a submission needs `.from(submitter)`".to_owned())
})?;
let stream = self.sessions.stream()?.to_owned();
let session = match &self.label {
Some(label) => crate::agent::derive_session_id(
&stream,
self.namespace.as_deref().unwrap_or_default(),
label,
),
None => crate::types::ConversationId::new(),
};
let start = laser_wire::agent::SessionStart {
label: self.label,
namespace: self.namespace,
agent: self.agent.clone(),
sdk: self.sessions.config().sdk_info().clone(),
parent: None,
root: None,
idle_timeout_micros: Some(
u64::try_from(self.sessions.config().idle_timeout_value().as_micros())
.unwrap_or(u64::MAX),
),
budget: self.budget,
tags: self.tags,
};
laser_wire::validate::Validate::validate(&start)?;
let handle = self.sessions.open(session).as_agent(submitter);
let lane = handle.lane()?;
lane.status(laser_wire::agent::OPERATION_SESSION)
.with_task_state(laser_wire::agent::TaskState::Submitted)
.body(encode_named(&start)?)
.content_type(ContentType::Cbor)
.send()
.await?;
let correlation = CorrelationId::mint();
lane.command(correlation, self.input)
.with_operation(self.operation)
.with_target(self.agent)
.with_metadata(laser_wire::agent::METADATA_SUBMITTED, Value::Bool(true))
.send()
.await?;
Ok(Submitted {
session,
correlation,
})
}
}
pub struct SessionControl {
laser: crate::laser::Laser,
session: crate::types::ConversationId,
operator: Option<laser_wire::agent::AgentId>,
participants: Option<Vec<laser_wire::agent::AgentId>>,
#[cfg(feature = "sign")]
key: Option<crate::sign::SigningKey>,
}
impl SessionControl {
#[must_use]
pub fn as_operator(mut self, operator: impl Into<laser_wire::agent::AgentId>) -> Self {
let operator = operator.into();
self.operator = Some(operator);
self
}
#[cfg(feature = "sign")]
#[must_use]
pub fn signed_by(mut self, key: crate::sign::SigningKey) -> Self {
self.key = Some(key);
self
}
#[must_use]
pub fn participants(
mut self,
participants: impl IntoIterator<Item = impl Into<laser_wire::agent::AgentId>>,
) -> Self {
self.participants = Some(participants.into_iter().map(Into::into).collect());
self
}
pub async fn pause(&self) -> Result<AgdxReceipt, LaserError> {
let participants = match &self.participants {
Some(participants) => participants.clone(),
None => crate::agent::pause::lane_participants(&self.laser, self.session).await?,
};
let body = serde_json::to_vec(&laser_wire::agent::SessionPauseRequest { participants })
.map_err(|error| LaserError::Codec(format!("pause request: {error}")))?;
self.request_with(laser_wire::dispatch::OPERATION_SESSION_PAUSE, body)
.await
}
pub async fn resume(&self) -> Result<AgdxReceipt, LaserError> {
self.request(laser_wire::dispatch::OPERATION_SESSION_RESUME)
.await
}
pub async fn cancel(&self) -> Result<AgdxReceipt, LaserError> {
self.request(laser_wire::dispatch::OPERATION_SESSION_CANCEL)
.await
}
pub async fn force_cancel(&self) -> Result<AgdxReceipt, LaserError> {
let end = laser_wire::agent::SessionEnd {
reason: Some("forced".to_owned()),
error: None,
};
let producer = self.producer()?;
let send = producer
.status(laser_wire::agent::OPERATION_SESSION)
.with_task_state(laser_wire::agent::TaskState::Canceled)
.body(encode_named(&end)?)
.content_type(ContentType::Cbor)
.last();
self.sign(send).send_receipt().await
}
async fn request(&self, operation: &str) -> Result<AgdxReceipt, LaserError> {
self.request_with(operation, b"{}".to_vec()).await
}
async fn request_with(
&self,
operation: &str,
body: Vec<u8>,
) -> Result<AgdxReceipt, LaserError> {
let producer = self.producer()?;
let send = producer
.command(CorrelationId::mint(), body)
.with_operation(operation)
.content_type(ContentType::Json);
self.sign(send).send_receipt().await
}
fn sign<'a>(&'a self, send: crate::agent::AgdxSend<'a>) -> crate::agent::AgdxSend<'a> {
#[cfg(feature = "sign")]
if let Some(key) = &self.key {
return send.signed_by(key);
}
send
}
fn producer(&self) -> Result<crate::agent::Agdx, LaserError> {
let operator = self.operator.clone().ok_or_else(|| {
LaserError::Invalid("session control needs `.as_operator(id)`".to_owned())
})?;
Ok(self
.laser
.agdx(AgentTopic::Control, operator, self.session.into()))
}
}
impl crate::agent::Sessions {
pub fn submit(
&self,
agent: impl Into<laser_wire::agent::AgentId>,
input: impl Into<Vec<u8>>,
) -> SubmitBuilder {
SubmitBuilder {
sessions: self.clone(),
agent: agent.into(),
input: input.into(),
submitter: None,
operation: laser_wire::dispatch::OPERATION_INVOKE_AGENT.to_owned(),
label: None,
namespace: None,
budget: None,
tags: Vec::new(),
}
}
pub fn control(&self, stream: &str, id: crate::types::ConversationId) -> SessionControl {
SessionControl {
laser: self.laser().with_default_stream(stream),
session: id,
operator: None,
participants: None,
#[cfg(feature = "sign")]
key: None,
}
}
}