use super::{
protocol::{ServiceErrorCode, ServiceRequest},
runtime::ServiceRuntime,
};
use crate::sessions::{Session, SessionEvent, SessionEventKind, SessionWriterLease};
use serde::Deserialize;
use serde_json::{Value, json};
use std::{collections::HashMap, sync::Arc};
pub(super) const MAX_ATTACHMENTS: usize = 16;
use crate::sessions::FRONTEND_PAGE_LIMIT as MAX_PAGE_ITEMS;
enum AttachmentKind {
ImplicitTurn,
ExplicitOpen,
DaemonClaim,
}
struct Attachment {
session: Session,
_writer: SessionWriterLease,
explicit: bool,
closing: bool,
}
pub(super) struct ServiceSessions {
runtime: Arc<ServiceRuntime>,
attachments: HashMap<String, Attachment>,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct PageParams {
#[serde(default)]
after: Option<String>,
#[serde(default = "default_limit")]
limit: usize,
}
fn default_limit() -> usize {
MAX_PAGE_ITEMS
}
impl ServiceSessions {
pub(super) fn new(runtime: Arc<ServiceRuntime>) -> Self {
Self {
runtime,
attachments: HashMap::new(),
}
}
pub(super) fn for_turn(&mut self, id: Option<&str>) -> Result<Session, ServiceErrorCode> {
self.attach(id, AttachmentKind::ImplicitTurn)
}
fn attach(
&mut self,
id: Option<&str>,
kind: AttachmentKind,
) -> Result<Session, ServiceErrorCode> {
let explicit = !matches!(kind, AttachmentKind::ImplicitTurn);
if let Some(id) = id {
validate_id(id)?;
if let Some(attachment) = self.attachments.get_mut(id) {
if attachment.closing {
return Err(ServiceErrorCode::SessionBusy);
}
attachment.explicit |= explicit;
return Ok(attachment.session.clone());
}
}
if self.attachments.len() >= MAX_ATTACHMENTS {
return Err(ServiceErrorCode::LimitExceeded);
}
let session = match id {
Some(id) => self.runtime.session_manager.open_existing(id),
None => self.runtime.session_manager.create(),
}
.map_err(|_| ServiceErrorCode::SessionUnavailable)?;
let writer = match kind {
AttachmentKind::DaemonClaim => session.try_daemon_writer(),
AttachmentKind::ImplicitTurn | AttachmentKind::ExplicitOpen => {
session.try_frontend_writer()
}
}
.map_err(|_| ServiceErrorCode::SessionUnavailable)?
.ok_or(ServiceErrorCode::SessionBusy)?;
let session = session
.activate()
.map_err(|_| ServiceErrorCode::SessionBusy)?;
if id.is_none() {
let event = SessionEvent::new_kind(
SessionEventKind::Diagnostic,
session.id().to_owned(),
self.runtime.cwd.clone(),
json!({"lifecycle": "created"}),
);
session
.append_with_outcome(&event)
.map_err(|_| ServiceErrorCode::SessionUnavailable)?;
}
self.attachments.insert(
session.id().to_owned(),
Attachment {
session: session.clone(),
_writer: writer,
explicit,
closing: false,
},
);
Ok(session)
}
pub(super) fn operation(
&mut self,
request: &ServiceRequest,
active: bool,
) -> Result<Value, ServiceErrorCode> {
match request.method.as_str() {
"session.list" => {
if request.session_id.is_some() {
return Err(ServiceErrorCode::InvalidPayload);
}
let params = page_params(&request.payload)?;
if let Some(after) = ¶ms.after {
validate_id(after)?;
}
let page: crate::sessions::FrontendSessionList = self
.runtime
.session_manager
.frontend_list(params.after.as_deref(), params.limit)
.map_err(|error| {
if error
.downcast_ref::<crate::sessions::BoundedReadError>()
.is_some()
{
ServiceErrorCode::LimitExceeded
} else {
ServiceErrorCode::SessionUnavailable
}
})?;
serde_json::to_value(page).map_err(|_| ServiceErrorCode::InternalError)
}
"session.create" => {
empty_params(&request.payload)?;
if request.session_id.is_some() {
return Err(ServiceErrorCode::InvalidPayload);
}
let session = self.attach(None, AttachmentKind::ExplicitOpen)?;
Ok(json!({"session_id": session.id(), "status": "open"}))
}
"session.open" => {
empty_params(&request.payload)?;
let id = request_id(request)?;
let session = self.attach(Some(id), AttachmentKind::ExplicitOpen)?;
Ok(json!({"session_id": session.id(), "status": "open"}))
}
"session.replay" => {
let id = request_id(request)?;
let params = page_params(&request.payload)?;
let session = self
.runtime
.session_manager
.open_existing(id)
.map_err(|_| ServiceErrorCode::SessionUnavailable)?;
let page: crate::sessions::FrontendReplayPage = session
.frontend_replay(params.after.as_deref(), params.limit)
.map_err(|error| {
if error
.downcast_ref::<crate::sessions::FrontendSnapshotBusy>()
.is_some()
{
ServiceErrorCode::SessionBusy
} else {
ServiceErrorCode::SessionUnavailable
}
})?;
serde_json::to_value(page).map_err(|_| ServiceErrorCode::InternalError)
}
"session.close" => {
empty_params(&request.payload)?;
let id = request_id(request)?;
if active {
let attachment = self
.attachments
.get_mut(id)
.ok_or(ServiceErrorCode::SessionUnavailable)?;
attachment.closing = true;
Ok(json!({"session_id": id, "status": "closing"}))
} else {
self.attachments.remove(id);
Ok(json!({"session_id": id, "status": "closed"}))
}
}
_ => Err(ServiceErrorCode::UnsupportedOperation),
}
}
pub(super) fn claim(&mut self, id: &str) -> Result<(), ServiceErrorCode> {
let already_attached = self.attachments.contains_key(id);
let session = self.attach(Some(id), AttachmentKind::DaemonClaim)?;
if !already_attached {
let replay = session.read_events_tolerant_bounded(
crate::context::REPLAY_JSONL_MAX_LINES,
crate::context::REPLAY_JSONL_MAX_BYTES,
);
if !matches!(replay, Ok(ref replay) if replay.diagnostics.is_empty()) {
self.attachments.remove(id);
return Err(ServiceErrorCode::SessionUnavailable);
}
}
Ok(())
}
pub(super) fn detach(&mut self, id: &str, active: bool) {
if active {
if let Some(attachment) = self.attachments.get_mut(id) {
attachment.explicit = false;
}
} else {
self.attachments.remove(id);
}
}
pub(super) fn finish_turn(&mut self, id: &str) {
if self
.attachments
.get(id)
.is_some_and(|attachment| !attachment.explicit || attachment.closing)
{
self.attachments.remove(id);
}
}
pub(super) fn clear(&mut self) {
self.attachments.clear();
}
}
fn validate_id(id: &str) -> Result<(), ServiceErrorCode> {
if id.len() > super::protocol::MAX_ID_BYTES
|| crate::sessions::validate_session_id(id.to_owned()).is_err()
{
return Err(ServiceErrorCode::InvalidSessionId);
}
Ok(())
}
fn request_id(request: &ServiceRequest) -> Result<&str, ServiceErrorCode> {
let id = request
.session_id
.as_deref()
.ok_or(ServiceErrorCode::InvalidPayload)?;
validate_id(id)?;
Ok(id)
}
fn empty_params(value: &Value) -> Result<(), ServiceErrorCode> {
serde_json::from_value::<super::protocol::EmptyParams>(value.clone())
.map(|_| ())
.map_err(|_| ServiceErrorCode::InvalidPayload)
}
fn page_params(value: &Value) -> Result<PageParams, ServiceErrorCode> {
let params: PageParams =
serde_json::from_value(value.clone()).map_err(|_| ServiceErrorCode::InvalidPayload)?;
if !(1..=MAX_PAGE_ITEMS).contains(¶ms.limit)
|| params.after.as_ref().is_some_and(|after| after.len() > 128)
{
return Err(ServiceErrorCode::LimitExceeded);
}
Ok(params)
}