1pub mod actions;
2pub mod transport;
3use crate::llm::actions::client_trait::{Client, ClientActionResult};
4use crate::protocol::{ConnectContext, Event};
5use crate::server::loki::codec::{self, PushBatch};
6use crate::state::{
7 client_handles::{ClientCommand, ClientSendOutcome},
8 AccessLogOwner, ClientStatus,
9};
10use crate::{console_error, console_info};
11pub use actions::LokiClientProtocol;
12use anyhow::{Context, Result};
13use serde_json::{json, Value};
14use std::{collections::VecDeque, future::Future, pin::Pin};
15pub const MAX_QUEUED_EVENTS: usize = 32;
16pub const MAX_HANDLER_ACTIONS: usize = 32;
17pub const MAX_FOLLOWUP_DEPTH: usize = 8;
18type Handler = Pin<Box<dyn Future<Output = Result<crate::llm::ClientLlmResult>> + Send>>;
19type Exchange = Pin<Box<dyn Future<Output = Result<transport::PushResponse>> + Send>>;
20struct Inflight {
21 exchange: Exchange,
22 command: Option<ClientCommand>,
23 depth: usize,
24}
25fn handler(ctx: ConnectContext, event: Event) -> Handler {
26 Box::pin(async move {
27 let instruction = ctx
28 .state
29 .get_instruction_for_client(ctx.client_id)
30 .await
31 .unwrap_or_default();
32 let configured = ctx
33 .state
34 .get_client_event_handler_config(ctx.client_id)
35 .await
36 .is_some_and(|c| c.find_handler(event.id()).is_some());
37 if instruction.is_empty() && !configured {
38 ctx.state
39 .record_access_log(
40 AccessLogOwner::Client(ctx.client_id.as_u32()),
41 "Loki",
42 None,
43 event.id(),
44 event.data.clone(),
45 vec![],
46 )
47 .await;
48 return Ok(crate::llm::ClientLlmResult {
49 actions: vec![],
50 memory_updates: None,
51 });
52 }
53 let memory = ctx
54 .state
55 .get_memory_for_client(ctx.client_id)
56 .await
57 .unwrap_or_default();
58 crate::client::llm_budget::call_llm_for_client(
59 &ctx.llm_client,
60 &ctx.state,
61 ctx.client_id.to_string(),
62 &instruction,
63 &memory,
64 Some(&event),
65 &LokiClientProtocol::new(),
66 &ctx.status_tx,
67 )
68 .await
69 })
70}
71pub struct LokiClient;
72impl LokiClient {
73 pub async fn connect(ctx: ConnectContext) -> Result<std::net::SocketAddr> {
74 let origin = transport::Origin::parse(&ctx.remote_addr)?;
75 let token = ctx
76 .startup_params
77 .as_ref()
78 .map(|p| p.get_optional_string("auth_token"))
79 .transpose()?
80 .flatten();
81 if let Some(token) = &token {
82 codec::validate_token(token)?;
83 }
84 let mut commands =
85 crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id)
86 .await;
87 ctx.state
88 .update_client_status(ctx.client_id, ClientStatus::Connected)
89 .await;
90 let registrar = ctx.state.clone();
91 let id = ctx.client_id;
92 console_info!(
93 ctx.status_tx,
94 "Loki HTTP write session ready for {}",
95 origin.authority
96 );
97 registrar.spawn_client_task(id, async move {
98 let protocol = LokiClientProtocol::new();
99 let mut current = Some((
100 handler(
101 ctx.clone(),
102 Event::new(
103 &actions::LOKI_CONNECTED_EVENT,
104 json!({"remote_addr":format!("http://{}",origin.authority)}),
105 ),
106 ),
107 0,
108 ));
109 let mut events = VecDeque::new();
110 let mut actions = VecDeque::new();
111 let mut pending: Option<Inflight> = None;
112 'session: loop {
113 if current.is_none() {
114 if let Some((event, depth)) = events.pop_front() {
115 current = Some((handler(ctx.clone(), event), depth));
116 }
117 }
118 if pending.is_none() {
119 if let Some((action, depth)) = actions.pop_front() {
120 match prepare(&protocol, action, depth, &origin, &token) {
121 Ok(Prepared::Write(exchange)) => {
122 pending = Some(Inflight {
123 exchange,
124 command: None,
125 depth,
126 })
127 }
128 Ok(Prepared::Disconnect) => break,
129 Err(error) => console_error!(
130 ctx.status_tx,
131 "Loki handler action rejected: {}",
132 error
133 ),
134 }
135 continue;
136 }
137 }
138 tokio::select! {
139 result = async { pending.as_mut().unwrap().exchange.as_mut().await }, if pending.is_some() => {
140 let p = pending.take().unwrap();
141 match result {
142 Ok(response) => {
143 if let Some(command) = p.command {
144 finish(&ctx, command, Ok(ClientSendOutcome::Executed {
145 detail: format!("Loki HTTP {}", response.status),
146 })).await;
147 }
148 if events.len() == MAX_QUEUED_EVENTS {
149 console_error!(ctx.status_tx, "Loki response-event queue limit");
150 break;
151 }
152 events.push_back((Event::new(&actions::LOKI_RESPONSE_EVENT,
153 serde_json::to_value(response).unwrap_or(Value::Null)), p.depth));
154 }
155 Err(error) => {
156 if let Some(command) = p.command {
157 finish(&ctx, command, Err(anyhow::anyhow!("Loki HTTP exchange failed"))).await;
158 }
159 console_error!(ctx.status_tx, "Loki HTTP exchange failed: {}", error);
160 break;
161 }
162 }
163 }
164 result = async { current.as_mut().unwrap().0.as_mut().await }, if current.is_some() => {
165 let (_, depth) = current.take().unwrap();
166 match result {
167 Ok(result) => {
168 if let Some(memory) = result.memory_updates {
169 ctx.state.set_memory_for_client(id, memory).await;
170 }
171 if result.actions.len() > MAX_HANDLER_ACTIONS
172 || actions.len() + result.actions.len() > MAX_HANDLER_ACTIONS {
173 console_error!(ctx.status_tx, "Loki handler action count limit");
174 break;
175 }
176 actions.extend(result.actions.into_iter().map(|a| (a, depth + 1)));
177 }
178 Err(error) => console_error!(ctx.status_tx, "Loki handler failed: {}", error),
179 }
180 }
181 command = commands.recv() => {
182 let Some(command) = command else { break; };
183 let prepared = prepare(&protocol, command.action.clone(), 0, &origin, &token);
184 match prepared {
185 Ok(Prepared::Disconnect) => {
186 finish(&ctx, command, Ok(ClientSendOutcome::Disconnected)).await;
187 break 'session;
188 }
189 Ok(Prepared::Write(exchange)) if pending.is_none() => {
190 pending = Some(Inflight { exchange, command: Some(command), depth: 0 });
191 }
192 Ok(Prepared::Write(_)) => finish(&ctx, command, Ok(ClientSendOutcome::Rejected {
193 error: "One write is already in flight".into(),
194 })).await,
195 Err(error) => finish(&ctx, command, Ok(ClientSendOutcome::Rejected {
196 error: error.to_string(),
197 })).await,
198 }
199 }
200 }
201 let _ = ctx.status_tx.send("__UPDATE_UI__".into());
202 }
203 if let Some(p) = pending.take() {
204 if let Some(command) = p.command {
205 finish(&ctx, command, Err(anyhow::anyhow!("Loki write cancelled"))).await;
206 }
207 }
208 ctx.state.remove_client_handle(id).await;
209 ctx.state
210 .update_client_status(id, ClientStatus::Disconnected)
211 .await;
212 let _ = ctx.status_tx.send("__UPDATE_UI__".into());
213 }).await;
214 Ok("0.0.0.0:0".parse().unwrap())
215 }
216}
217enum Prepared {
218 Disconnect,
219 Write(Exchange),
220}
221fn prepare(
222 protocol: &LokiClientProtocol,
223 action: Value,
224 depth: usize,
225 origin: &transport::Origin,
226 token: &Option<String>,
227) -> Result<Prepared> {
228 match protocol.execute_action(action)? {
229 ClientActionResult::Disconnect => Ok(Prepared::Disconnect),
230 ClientActionResult::Custom { name, data } if name == "push_loki_entries" => {
231 anyhow::ensure!(depth <= MAX_FOLLOWUP_DEPTH, "write follow-up depth limit");
232 let batch: PushBatch = serde_json::from_value(data).context("invalid typed write")?;
233 let body = codec::encode_batch(&batch)?;
234 Ok(Prepared::Write(Box::pin(transport::write(
235 origin.clone(),
236 token.clone(),
237 batch,
238 body,
239 ))))
240 }
241 _ => anyhow::bail!("unsupported Loki action result"),
242 }
243}
244async fn finish(ctx: &ConnectContext, command: ClientCommand, result: Result<ClientSendOutcome>) {
245 let outcome = match &result {
246 Ok(outcome) => serde_json::to_value(outcome).unwrap_or(Value::Null),
247 Err(error) => json!({"error":error.to_string()}),
248 };
249 ctx.state
250 .record_access_log(
251 AccessLogOwner::Client(ctx.client_id.as_u32()),
252 "Loki",
253 None,
254 "injected_action",
255 command.action.clone(),
256 vec![outcome],
257 )
258 .await;
259 crate::client::command_support::reply(command, result);
260}