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