1use std::collections::BTreeMap;
13use std::sync::Arc;
14
15use gate4agent_wake::hosted::{plan_hosted_wake, HostedSessions, HostedTarget, HostedWake};
16use gate4agent_wake::{ProviderKind, WakeLetter};
17
18use crate::wake_pipeline::{deliver_routine, launch_spawn_profile, route, Route, RoutineLetter, SessionShape};
19
20use super::*;
21
22pub struct NodeLetter {
24 pub nick: String,
25 pub from: String,
26 pub event_id: String,
27 pub body: String,
28 pub own_spawn: bool,
31}
32
33#[derive(Clone)]
34pub struct NodeWake {
35 shared: Arc<NodeShared>,
36 token: Arc<str>,
37 policy: Arc<crate::wake_pipeline::WakePolicy>,
38}
39
40struct Snapshot(BTreeMap<String, HostedTarget>);
41
42impl HostedSessions for Snapshot {
43 fn lookup(&self, nick: &str) -> Option<HostedTarget> {
44 self.0.get(nick).cloned()
45 }
46}
47
48fn kind_of(agent: &AgentId) -> Option<ProviderKind> {
49 ProviderKind::parse(crate::m4a_accounts::m4a_provider(agent.as_str())?)
50}
51
52impl NodeWake {
53 pub(super) fn new(shared: Arc<NodeShared>, token: String, policy: Arc<crate::wake_pipeline::WakePolicy>) -> Self {
54 Self { shared, token: token.into(), policy }
55 }
56
57 pub fn authorized(&self, token: &str) -> bool {
59 let (a, b) = (self.token.as_bytes(), token.as_bytes());
60 a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
61 }
62
63 fn snapshot(&self) -> (Snapshot, BTreeMap<String, SessionAddress>) {
64 let live = self.shared.m4a_live.lock().unwrap_or_else(|p| p.into_inner()).clone();
65 let bindings = self.shared.session_bindings.lock().unwrap_or_else(|p| p.into_inner());
66 let mut targets = BTreeMap::new();
67 let mut addresses = BTreeMap::new();
68 for (nick, (instance_id, agent)) in live {
69 let Some(provider) = kind_of(&agent) else { continue };
70 let binding = bindings.get(&instance_id);
71 if let Some(b) = binding {
72 addresses.insert(
73 nick.clone(),
74 SessionAddress { workspace_id: b.workspace_id.clone(), session: SessionKey { instance_id, generation: b.generation } },
75 );
76 }
77 let spawn_allowed = self.policy.spawn_profile(&nick).is_some();
78 targets.insert(nick.clone(), HostedTarget { nick, provider, live: binding.is_some(), spawn_allowed });
79 }
80 (Snapshot(targets), addresses)
81 }
82
83 async fn resume_with_wake(&self, nick: &str, address: &SessionAddress, prompt: String) -> Result<&'static str, String> {
84 let record_id = self
85 .shared
86 .session_bindings
87 .lock()
88 .unwrap_or_else(|p| p.into_inner())
89 .get(&address.session.instance_id)
90 .and_then(|b| b.record_id.clone())
91 .ok_or("exited session has no record to resume")?;
92 let size = gate4agent_types::TerminalSize { rows: 40, columns: 120 };
93 match self.shared.resume_session_record(&record_id, size, Some(prompt)).await {
94 Ok(_) => {
95 tracing::info!(%nick, "letter resumed an exited session");
96 Ok("resumed")
97 }
98 Err(e) => {
99 tracing::warn!(%nick, code = ?e.code, "exited session not resumed");
100 Err("exited session could not be resumed".into())
101 }
102 }
103 }
104
105 pub async fn wake_spawned(&self, nick: String, from: String, event_id: String, body: String) -> Result<&'static str, String> {
107 self.wake(NodeLetter { nick, from, event_id, body, own_spawn: true }).await
108 }
109
110 pub async fn wake(&self, letter: NodeLetter) -> Result<&'static str, String> {
113 let (sessions, addresses) = self.snapshot();
114 let wake = WakeLetter { body: &letter.body, from_nick: &letter.from, event_id: &letter.event_id, room: None, reaction: None };
115 match plan_hosted_wake(&sessions, &letter.nick, &wake) {
116 HostedWake::Submit { nick, prompt } => {
117 let prompt = if self.policy.deliver_body { prompt } else { crate::wake_pipeline::wake_line(&letter.from, &letter.event_id) };
118 let address = addresses.get(&nick).cloned().ok_or("session vanished")?;
119 let wake_line = crate::wake_pipeline::wake_line(&letter.from, &letter.event_id);
120 let agent = match self.shared.validate_address(&address) {
121 Ok(agent) => agent,
122 Err(_) if letter.own_spawn => return self.resume_with_wake(&nick, &address, wake_line).await,
123 Err(_) => return Err("session is not live".to_owned()),
124 };
125 let refused = |e: NodeFailure| {
126 tracing::warn!(nick = %nick, code = ?e.code, "letter wake input refused");
127 "session did not take the turn".to_owned()
128 };
129 let shape = if self.shared.require_session_runtime_policy(&address, ProviderRuntimeRequirement::SemanticPrompt).is_ok() {
130 SessionShape::Semantic
131 } else {
132 SessionShape::RawTerminal
133 };
134 let routed = match route(&self.policy, &nick, shape) {
135 Route::Refuse(_) if letter.own_spawn && shape == SessionShape::RawTerminal => Route::TerminalTyping,
136 Route::Refuse(_) if letter.own_spawn && shape == SessionShape::Semantic => Route::Deliver,
138 other => other,
139 };
140 let prompt = match (letter.own_spawn, shape) {
143 (true, SessionShape::Semantic) => crate::wake_pipeline::letter_prompt(&letter.from, &letter.event_id, &letter.body),
144 (true, _) => crate::wake_pipeline::wake_line(&letter.from, &letter.event_id),
145 _ => prompt,
146 };
147 match routed {
148 Route::Deliver => {
149 let action = InputAction::SubmitPrompt(PromptPayload { text: prompt, framing: prompt_framing(&agent) });
150 match self.shared.dispatch_input_bounded(&address, action).await {
151 Err(_) if letter.own_spawn => return self.resume_with_wake(&nick, &address, wake_line).await,
154 other => other.map_err(refused)?,
155 };
156 }
157 Route::TerminalTyping => {
158 let text = prompt.replace(['\n', '\r'], " ");
161 match self.shared.dispatch_input_bounded(&address, InputAction::TerminalText(TerminalText { text })).await {
162 Err(_) if letter.own_spawn => return self.resume_with_wake(&nick, &address, wake_line).await,
165 other => other.map_err(refused)?,
166 };
167 tokio::time::sleep(std::time::Duration::from_millis(400)).await;
169 self.shared.dispatch_input_bounded(&address, InputAction::TerminalControl(TerminalControl::ControlM)).await.map_err(refused)?;
170 }
171 Route::Refuse(why) => {
172 tracing::info!(nick = %nick, "letter not typed: no terminal grant");
173 return Err(why);
174 }
175 Route::Routine(_) => return Err("unexpected routine route for a live session".into()),
176 Route::Spawn(_) => return Err("unexpected spawn route for a live session".into()),
177 }
178 tracing::info!(nick = %nick, "letter woke a live session");
179 Ok("submitted")
180 }
181 HostedWake::Spawn { nick, .. } => match self.policy.spawn_profile(&nick).cloned() {
182 Some(profile) => {
183 tokio::task::spawn_blocking(move || launch_spawn_profile(&profile))
184 .await
185 .map_err(|_| "spawn profile task failed".to_owned())?
186 .map_err(|e| e)?;
187 tracing::info!(nick = %nick, "letter woke an offline session via spawn profile");
188 Ok("spawned")
189 }
190 None => Err("spawn on letter is not enabled on this node".into()),
191 },
192 HostedWake::Refuse(why) => match route(&self.policy, &letter.nick, SessionShape::None) {
193 Route::Spawn(profile) => {
194 let nick = letter.nick.clone();
195 tokio::task::spawn_blocking(move || launch_spawn_profile(&profile))
196 .await
197 .map_err(|_| "spawn profile task failed".to_owned())?
198 .map_err(|e| e)?;
199 tracing::info!(nick = %nick, "letter woke an offline session via spawn profile");
200 Ok("spawned")
201 }
202 Route::Routine(target) => {
203 let routine_letter = RoutineLetter {
204 to: letter.nick.clone(),
205 from: letter.from.clone(),
206 event_id: letter.event_id.clone(),
207 subject: None,
208 room: None,
209 body: letter.body.clone(),
210 reply: None,
211 note: None,
212 };
213 let nick = letter.nick.clone();
214 match tokio::task::spawn_blocking(move || deliver_routine(&target, &routine_letter)).await {
215 Ok(Ok(_)) => {
216 tracing::info!(nick = %nick, "letter delivered to a routine");
217 Ok("routine")
218 }
219 Ok(Err(e)) => Err(e),
220 Err(_) => Err("routine delivery task failed".into()),
221 }
222 }
223 _ => Err(why),
224 },
225 }
226 }
227}