use super::{
dispatcher::{
ServiceDispatcher, ServiceOutbound, ServiceSnapshot, ServiceTransportCapabilities,
},
protocol::{
RequestIdentity, ServiceErrorCode, ServiceMessage, ServiceRequest, ServiceResponse,
},
runtime::ServiceRuntime,
turns::{ServiceTurnManager, TurnWorkerMessage},
};
use crossbeam_channel::{Receiver, RecvError, bounded};
use std::sync::Arc;
pub(crate) struct ServiceCoordinator {
dispatcher: ServiceDispatcher,
turns: ServiceTurnManager,
turn_receiver: Receiver<TurnWorkerMessage>,
auth: super::auth::ServiceAuthManager,
runtime: Arc<ServiceRuntime>,
configuration_worker: Option<std::thread::JoinHandle<()>>,
worker_sender: crossbeam_channel::Sender<TurnWorkerMessage>,
}
pub(crate) struct ServiceWorkerOutput {
pub(crate) outbound: ServiceOutbound,
pub(crate) terminal_turn_id: Option<String>,
}
impl ServiceCoordinator {
pub(crate) fn new(
runtime: Arc<ServiceRuntime>,
transport_capabilities: ServiceTransportCapabilities,
) -> Self {
let (turn_sender, turn_receiver) = bounded(super::turns::TURN_EVENT_QUEUE_CAPACITY);
let readiness_runtime = Arc::clone(&runtime);
Self {
dispatcher: ServiceDispatcher::new(
ServiceSnapshot::new(move || readiness_runtime.current_provider_auth_ready()),
transport_capabilities,
),
auth: super::auth::ServiceAuthManager::new(Arc::clone(&runtime)),
runtime: Arc::clone(&runtime),
configuration_worker: None,
worker_sender: turn_sender.clone(),
turns: ServiceTurnManager::new(runtime, turn_sender),
turn_receiver,
}
}
pub(super) fn persistent(runtime: Arc<ServiceRuntime>) -> Self {
let mut coordinator = Self::new(runtime, ServiceTransportCapabilities::new(vec![]));
coordinator.dispatch_request(ServiceRequest {
protocol_version: 1,
kind: super::protocol::MessageKind::Request,
request_id: uuid::Uuid::new_v4().to_string(),
session_id: None,
method: "initialize".into(),
payload: serde_json::json!({"supported_protocol_versions":[1]}),
});
coordinator
}
pub(super) fn claim_session(&mut self, id: &str) -> Result<(), ServiceErrorCode> {
self.turns.claim_session(id)
}
pub(super) fn detach_session(&mut self, id: &str) {
self.turns.detach_session(id);
}
pub(super) fn release_turn_request(&mut self, id: &str) {
self.turns.release_request_guard(id);
}
#[cfg(test)]
pub(super) fn set_turn_worker(&mut self, worker: super::turns::TurnWorker) {
self.turns.set_worker(worker);
}
#[cfg(test)]
pub(super) fn set_login_worker(&mut self, worker: super::auth::LoginWorker) {
self.auth.set_worker(worker);
}
pub(super) fn worker_finished(&self, id: &str) -> bool {
self.turns.worker_finished(id)
}
pub(super) fn cancel_login(&self) {
self.auth.cancel_all();
}
pub(super) fn prepare_logout(
&self,
request: &ServiceRequest,
) -> Result<Option<Arc<std::sync::atomic::AtomicBool>>, ServiceErrorCode> {
let provider = self.auth.validate_logout(request)?;
if provider == crate::providers::OPENAI_CODEX_PROVIDER {
self.auth.cancel_all();
Ok(self.auth.cancellation_barrier())
} else {
Ok(None)
}
}
pub(crate) fn dispatch_request(&mut self, request: ServiceRequest) -> ServiceOutbound {
match request.method.as_str() {
"catalog.providers" | "catalog.models" | "catalog.refresh" | "config.get"
| "config.set" => self.dispatch_turn_request(request, u64::MAX),
"session.list" | "session.create" | "session.open" | "session.replay"
| "session.close" => self.dispatch_turn_request(request, u64::MAX),
super::protocol::TURN_START_METHOD | super::protocol::TURN_CANCEL_METHOD => {
self.dispatch_turn_request(request, u64::MAX)
}
"auth.status"
| "auth.login.start"
| "auth.login.callback"
| "auth.login.cancel"
| "auth.logout" => self.dispatch_turn_request(request, u64::MAX),
_ => self.dispatcher.dispatch(request),
}
}
pub(crate) fn error_outbound(
&mut self,
identity: RequestIdentity,
code: ServiceErrorCode,
) -> ServiceOutbound {
self.dispatcher.error_outbound(identity, code)
}
pub(super) fn configuration_busy(&self) -> bool {
self.configuration_worker.is_some()
}
pub(super) fn dispatch_turn_request(
&mut self,
request: ServiceRequest,
maximum_sequence: u64,
) -> ServiceOutbound {
self.dispatch_prepared_turn_request(request, maximum_sequence, None)
}
pub(super) fn dispatch_prepared_turn_request(
&mut self,
request: ServiceRequest,
maximum_sequence: u64,
prepared_runtime: Option<Arc<ServiceRuntime>>,
) -> ServiceOutbound {
let identity = RequestIdentity::from_request(&request);
let guard = match self.dispatcher.begin_request(&request.request_id) {
Ok(guard) => guard,
Err(code) => {
return ServiceOutbound::unguarded(vec![ServiceMessage::response(
ServiceResponse::error(identity, code),
)]);
}
};
if let Some(code) = self.dispatcher.request_error_code(&request) {
return ServiceOutbound::guarded(
vec![ServiceMessage::response(ServiceResponse::error(
identity, code,
))],
guard,
);
}
match request.method.as_str() {
"catalog.providers" | "catalog.models" | "catalog.refresh" | "config.get"
| "config.set" => {
if !self.dispatcher.is_initialized()
|| request.session_id.is_some()
|| self.configuration_worker.is_some()
{
let code = if !self.dispatcher.is_initialized() {
ServiceErrorCode::NotInitialized
} else if request.session_id.is_some() {
ServiceErrorCode::InvalidPayload
} else {
ServiceErrorCode::ConfigurationBusy
};
return ServiceOutbound::guarded(
vec![ServiceMessage::response(ServiceResponse::error(
identity, code,
))],
guard,
);
}
let paths = self.runtime.config.paths.clone();
let sender = self.worker_sender.clone();
let work = Arc::new(std::sync::Mutex::new(Some((request, guard))));
let worker_work = Arc::clone(&work);
match std::thread::Builder::new()
.name("magi-service-config".into())
.spawn(move || {
let (request, guard) = worker_work
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take()
.expect("configuration worker owns request");
let outbound = super::configuration::dispatch(&paths, request, guard, true);
let _ = sender.send(TurnWorkerMessage::Configuration(outbound));
}) {
Ok(worker) => {
self.configuration_worker = Some(worker);
ServiceOutbound::unguarded(vec![])
}
Err(_) => {
let (_, guard) = work
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take()
.expect("failed spawn retains request");
ServiceOutbound::guarded(
vec![ServiceMessage::response(ServiceResponse::error(
identity,
ServiceErrorCode::InternalError,
))],
guard,
)
}
}
}
"session.list" | "session.create" | "session.open" | "session.replay"
| "session.close" => {
self.turns
.dispatch_session(request, guard, self.dispatcher.is_initialized())
}
super::protocol::TURN_START_METHOD => self.turns.start_with_settings_capture(
request,
guard,
self.dispatcher.is_initialized(),
prepared_runtime,
maximum_sequence,
),
super::protocol::TURN_CANCEL_METHOD => {
self.turns
.cancel(request, guard, self.dispatcher.is_initialized())
}
_ => self
.auth
.dispatch(request, guard, self.dispatcher.is_initialized()),
}
}
pub(crate) fn auth_receiver(&self) -> &Receiver<super::auth::AuthWorkerMessage> {
&self.auth.receiver
}
pub(crate) fn auth_output(
&mut self,
message: super::auth::AuthWorkerMessage,
) -> Option<ServiceOutbound> {
self.auth.output(message)
}
pub(crate) fn worker_receiver(&self) -> &Receiver<TurnWorkerMessage> {
&self.turn_receiver
}
pub(crate) fn worker_output(
&mut self,
message: TurnWorkerMessage,
) -> Option<ServiceWorkerOutput> {
if let TurnWorkerMessage::Configuration(outbound) = message {
if let Some(worker) = self.configuration_worker.take() {
let _ = worker.join();
}
return Some(ServiceWorkerOutput {
outbound,
terminal_turn_id: None,
});
}
let (turn_id, event, is_terminal) = self.turns.handle_worker_message(message)?;
Some(ServiceWorkerOutput {
outbound: ServiceOutbound::unguarded(vec![ServiceMessage::Event(event)]),
terminal_turn_id: is_terminal.then_some(turn_id),
})
}
pub(crate) fn finish_worker_output(&mut self, turn_id: &str) {
self.turns.finish_output(turn_id);
}
pub(crate) fn cancel_all(&self) {
self.turns.cancel_all();
self.auth.cancel_all();
}
pub(crate) fn shutdown_next(&mut self) -> Result<Option<ServiceWorkerOutput>, RecvError> {
self.cancel_all();
while !self.turns.is_empty() || self.auth.is_active() || self.configuration_worker.is_some()
{
crossbeam_channel::select! {
recv(self.turn_receiver) -> message => {
if let Some(output) = self.worker_output(message?) { return Ok(Some(output)); }
}
recv(self.auth.receiver) -> message => {
if let Some(outbound) = self.auth.output(message?) {
return Ok(Some(ServiceWorkerOutput { outbound, terminal_turn_id: None }));
}
}
}
}
self.turns.join_all();
Ok(None)
}
pub(crate) fn join_all(&mut self) {
self.turns.join_all();
}
}