1use std::{sync::Arc, time::Duration};
2
3use futures_util::{Stream, StreamExt};
4use longbridge_httpcli::{HttpClient, Json, Method};
5use serde::{Deserialize, Serialize};
6use tracing::{Subscriber, dispatcher, instrument::WithSubscriber};
7
8use crate::{Config, Result, agent::types::*};
9
10const AGENT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
16
17struct InnerAgentContext {
18 http_cli: HttpClient,
19 log_subscriber: Arc<dyn Subscriber + Send + Sync>,
20}
21
22impl Drop for InnerAgentContext {
23 fn drop(&mut self) {
24 dispatcher::with_default(&self.log_subscriber.clone().into(), || {
25 tracing::info!("agent context dropped");
26 });
27 }
28}
29
30#[derive(Clone)]
34pub struct AgentContext(Arc<InnerAgentContext>);
35
36#[derive(Debug, Deserialize)]
37struct SseEnvelope {
38 event: String,
39 #[serde(default)]
40 data: serde_json::Value,
41 #[serde(default)]
44 tool_name: Option<String>,
45}
46
47fn map_conversation_event(
51 item: longbridge_httpcli::HttpClientResult<longbridge_httpcli::SseEvent>,
52 started: &mut Option<(String, String)>,
53) -> Result<ConversationStreamEvent> {
54 let event = item?;
55 let envelope: SseEnvelope = serde_json::from_str(&event.data)?;
56 Ok(match envelope.event.as_str() {
57 "chat_started" => {
58 let payload: ChatStartedPayload = serde_json::from_value(envelope.data)?;
59 *started = Some((payload.chat_uid.clone(), payload.message_id.clone()));
60 ConversationStreamEvent::ChatStarted(payload)
61 }
62 "message" => ConversationStreamEvent::Message(serde_json::from_value(envelope.data)?),
63 "workflow_started" => {
64 ConversationStreamEvent::WorkflowStarted(serde_json::from_value(envelope.data)?)
65 }
66 "ping" => ConversationStreamEvent::Ping,
67 "thinking_started" => {
68 ConversationStreamEvent::ThinkingStarted(serde_json::from_value(envelope.data)?)
69 }
70 "thinking_finished" => {
71 ConversationStreamEvent::ThinkingFinished(serde_json::from_value(envelope.data)?)
72 }
73 "node_tool_use_started" => {
74 ConversationStreamEvent::NodeToolUseStarted(serde_json::from_value(envelope.data)?)
75 }
76 "node_tool_use_finished" => {
77 ConversationStreamEvent::NodeToolUseFinished(serde_json::from_value(envelope.data)?)
78 }
79 "subagent_started" => {
80 ConversationStreamEvent::SubagentStarted(serde_json::from_value(envelope.data)?)
81 }
82 "subagent_progress" => {
83 ConversationStreamEvent::SubagentProgress(serde_json::from_value(envelope.data)?)
84 }
85 "subagent_finished" => {
86 ConversationStreamEvent::SubagentFinished(serde_json::from_value(envelope.data)?)
87 }
88 "agent_tool_started" => {
89 ConversationStreamEvent::AgentToolStarted(serde_json::from_value(envelope.data)?)
90 }
91 "agent_tool_progress" => {
92 ConversationStreamEvent::AgentToolProgress(serde_json::from_value(envelope.data)?)
93 }
94 "agent_tool_finished" => {
95 ConversationStreamEvent::AgentToolFinished(serde_json::from_value(envelope.data)?)
96 }
97 "human_interaction_required" => {
98 let interrupt: Interrupt = serde_json::from_value(envelope.data)?;
99 ConversationStreamEvent::HumanInteractionRequired(
100 ConversationResponse::from_stream_interrupt(started.clone(), interrupt),
101 )
102 }
103 "query_masked" => {
104 ConversationStreamEvent::QueryMasked(serde_json::from_value(envelope.data)?)
105 }
106 "plan_changed" => {
107 let mut payload: PlanChangedPayload = serde_json::from_value(envelope.data)?;
108 payload.tool_name = envelope.tool_name.clone().unwrap_or_default();
109 ConversationStreamEvent::PlanChanged(payload)
110 }
111 "context_compress_started" => {
112 ConversationStreamEvent::ContextCompressStarted(serde_json::from_value(envelope.data)?)
113 }
114 "context_compress_finished" => {
115 ConversationStreamEvent::ContextCompressFinished(serde_json::from_value(envelope.data)?)
116 }
117 "chat_finished" => {
118 ConversationStreamEvent::ChatFinished(serde_json::from_value(envelope.data)?)
119 }
120 "chat_title_updated" => {
121 ConversationStreamEvent::ChatTitleUpdated(serde_json::from_value(envelope.data)?)
122 }
123 "workflow_finished" => {
124 let payload: WorkflowFinishedPayload = serde_json::from_value(envelope.data)?;
125 ConversationStreamEvent::WorkflowFinished(ConversationResponse::from_stream_parts(
126 started.clone(),
127 payload,
128 ))
129 }
130 _ => ConversationStreamEvent::Other {
131 event: envelope.event,
132 data: envelope.data,
133 },
134 })
135}
136
137impl AgentContext {
138 pub fn new(config: Arc<Config>) -> Self {
140 let log_subscriber = config.create_log_subscriber("agent");
141 dispatcher::with_default(&log_subscriber.clone().into(), || {
142 tracing::info!(language = ?config.language, "creating agent context");
143 });
144 let ctx = Self(Arc::new(InnerAgentContext {
145 http_cli: config.create_http_client(),
146 log_subscriber,
147 }));
148 dispatcher::with_default(&ctx.0.log_subscriber.clone().into(), || {
149 tracing::info!("agent context created");
150 });
151 ctx
152 }
153
154 #[inline]
156 pub fn log_subscriber(&self) -> Arc<dyn Subscriber + Send + Sync> {
157 self.0.log_subscriber.clone()
158 }
159
160 pub async fn workspaces(&self) -> Result<WorkspacesResponse> {
164 Ok(self
165 .0
166 .http_cli
167 .request(Method::GET, "/v1/ai/workspaces")
168 .response::<Json<WorkspacesResponse>>()
169 .send()
170 .with_subscriber(self.0.log_subscriber.clone())
171 .await?
172 .0)
173 }
174
175 pub async fn agents(
179 &self,
180 workspace_id: impl Into<String>,
181 opts: impl Into<Option<GetAgentsOptions>>,
182 ) -> Result<AgentsResponse> {
183 let workspace_id = workspace_id.into();
184 Ok(self
185 .0
186 .http_cli
187 .request(
188 Method::GET,
189 format!("/v1/ai/workspaces/{workspace_id}/agents"),
190 )
191 .query_params(opts.into().unwrap_or_default())
192 .response::<Json<AgentsResponse>>()
193 .send()
194 .with_subscriber(self.0.log_subscriber.clone())
195 .await?
196 .0)
197 }
198
199 pub async fn public_agents(
209 &self,
210 opts: impl Into<Option<GetAgentsOptions>>,
211 ) -> Result<AgentsResponse> {
212 Ok(self
213 .0
214 .http_cli
215 .request(Method::GET, "/v1/ai/agents")
216 .query_params(opts.into().unwrap_or_default())
217 .response::<Json<AgentsResponse>>()
218 .send()
219 .with_subscriber(self.0.log_subscriber.clone())
220 .await?
221 .0)
222 }
223
224 pub async fn conversation(
235 &self,
236 agent_id: impl Into<String>,
237 query: impl Into<String>,
238 chat_uid: impl Into<Option<String>>,
239 parent_message_id: impl Into<Option<String>>,
240 ) -> Result<ConversationResponse> {
241 #[derive(Debug, Serialize)]
242 struct Body {
243 query: String,
244 #[serde(skip_serializing_if = "Option::is_none")]
245 chat_uid: Option<String>,
246 #[serde(skip_serializing_if = "Option::is_none")]
247 parent_message_id: Option<String>,
248 }
249
250 let agent_id = agent_id.into();
251 Ok(self
252 .0
253 .http_cli
254 .request(
255 Method::POST,
256 format!("/v1/ai/agents/{agent_id}/conversations"),
257 )
258 .header("Accept", "application/json")
259 .body(Json(Body {
260 query: query.into(),
261 chat_uid: chat_uid.into(),
262 parent_message_id: parent_message_id.into(),
263 }))
264 .timeout(AGENT_REQUEST_TIMEOUT)
265 .response::<Json<ConversationResponse>>()
266 .send()
267 .with_subscriber(self.0.log_subscriber.clone())
268 .await?
269 .0)
270 }
271
272 pub async fn continue_conversation(
279 &self,
280 agent_id: impl Into<String>,
281 chat_uid: impl Into<String>,
282 message_id: impl Into<String>,
283 answers: AnswersByToolCall,
284 ) -> Result<ConversationResponse> {
285 #[derive(Debug, Serialize)]
286 struct Body {
287 answers_by_tool_call: AnswersByToolCall,
288 }
289
290 let agent_id = agent_id.into();
291 let chat_uid = chat_uid.into();
292 let message_id = message_id.into();
293 Ok(self
294 .0
295 .http_cli
296 .request(
297 Method::POST,
298 format!(
299 "/v1/ai/agents/{agent_id}/conversations/{chat_uid}/messages/{message_id}/continue"
300 ),
301 )
302 .header("Accept", "application/json")
303 .body(Json(Body {
304 answers_by_tool_call: answers,
305 }))
306 .timeout(AGENT_REQUEST_TIMEOUT)
307 .response::<Json<ConversationResponse>>()
308 .send()
309 .with_subscriber(self.0.log_subscriber.clone())
310 .await?
311 .0)
312 }
313
314 pub async fn conversation_streamed(
331 &self,
332 agent_id: impl Into<String>,
333 query: impl Into<String>,
334 chat_uid: impl Into<Option<String>>,
335 parent_message_id: impl Into<Option<String>>,
336 ) -> Result<impl Stream<Item = Result<ConversationStreamEvent>> + Send + 'static> {
337 #[derive(Debug, Serialize)]
338 struct Body {
339 query: String,
340 #[serde(skip_serializing_if = "Option::is_none")]
341 chat_uid: Option<String>,
342 #[serde(skip_serializing_if = "Option::is_none")]
343 parent_message_id: Option<String>,
344 }
345
346 let agent_id = agent_id.into();
347 let raw = self
348 .0
349 .http_cli
350 .request(
351 Method::POST,
352 format!("/v1/ai/agents/{agent_id}/conversations"),
353 )
354 .body(Json(Body {
355 query: query.into(),
356 chat_uid: chat_uid.into(),
357 parent_message_id: parent_message_id.into(),
358 }))
359 .timeout(AGENT_REQUEST_TIMEOUT)
360 .send_events()
361 .with_subscriber(self.0.log_subscriber.clone())
362 .await?;
363
364 let mut started: Option<(String, String)> = None;
365 Ok(raw.map(move |item| map_conversation_event(item, &mut started)))
366 }
367
368 pub async fn continue_conversation_streamed(
375 &self,
376 agent_id: impl Into<String>,
377 chat_uid: impl Into<String>,
378 message_id: impl Into<String>,
379 answers: AnswersByToolCall,
380 ) -> Result<impl Stream<Item = Result<ConversationStreamEvent>> + Send + 'static> {
381 #[derive(Debug, Serialize)]
382 struct Body {
383 answers_by_tool_call: AnswersByToolCall,
384 }
385
386 let agent_id = agent_id.into();
387 let chat_uid = chat_uid.into();
388 let message_id = message_id.into();
389 let mut started = Some((chat_uid.clone(), message_id.clone()));
393 let raw = self
394 .0
395 .http_cli
396 .request(
397 Method::POST,
398 format!(
399 "/v1/ai/agents/{agent_id}/conversations/{chat_uid}/messages/{message_id}/continue"
400 ),
401 )
402 .body(Json(Body {
403 answers_by_tool_call: answers,
404 }))
405 .timeout(AGENT_REQUEST_TIMEOUT)
406 .send_events()
407 .with_subscriber(self.0.log_subscriber.clone())
408 .await?;
409
410 Ok(raw.map(move |item| map_conversation_event(item, &mut started)))
411 }
412}
413
414#[cfg(test)]
415mod tests {
416 use super::*;
417
418 const CHAT_STARTED: &str = r#"{"event":"chat_started","workflow_run_id":"wr_1","data":{"chat_uid":"ct_9f2c1a5b","message_id":42}}"#;
421 const MESSAGE: &str = r#"{"event":"message","workflow_run_id":"wr_1","data":{"text":"Tesla"}}"#;
422 const WORKFLOW_FINISHED: &str = r#"{"event":"workflow_finished","workflow_run_id":"wr_1","data":{"status":"succeeded","elapsed_time":3.21,"outputs":{"answer":"Tesla (TSLA.US) recently..."}}}"#;
423
424 const WORKFLOW_STARTED: &str = r#"{"event":"workflow_started","workflow_run_id":"wr_1","data":{"hit_cache":false,"inputs":{"chat_id":834552,"chat_uid":"ct_9f2c1a5b","message_id":42,"query":"How has Tesla stock performed recently?"},"started_at":1784545150,"workflow_id":176476}}"#;
427 const PING: &str = r#"{"event":"ping","workflow_run_id":"wr_1","data":null}"#;
428 const CHAT_FINISHED: &str = r#"{"event":"chat_finished","workflow_run_id":"wr_1","data":{"chat_id":834552,"chat_uid":"ct_9f2c1a5b","error":"","error_message":"","message_id":42}}"#;
429 const CHAT_TITLE_UPDATED: &str = r#"{"event":"chat_title_updated","workflow_run_id":"wr_1","data":{"chat_id":834552,"chat_uid":"ct_9f2c1a5b","source":"ai_generated","title":"Tesla stock performance","updated_at":1784546957}}"#;
430
431 fn sse(data: &str) -> longbridge_httpcli::HttpClientResult<longbridge_httpcli::SseEvent> {
432 Ok(longbridge_httpcli::SseEvent {
433 event: "message".to_string(),
434 data: data.to_string(),
435 id: String::new(),
436 retry: None,
437 })
438 }
439
440 #[test]
441 fn map_conversation_event_full_sequence() {
442 let mut started = None;
443
444 match map_conversation_event(sse(CHAT_STARTED), &mut started).unwrap() {
445 ConversationStreamEvent::ChatStarted(payload) => {
446 assert_eq!(payload.chat_uid, "ct_9f2c1a5b");
447 assert_eq!(payload.message_id, "42");
448 }
449 other => panic!("unexpected event: {other:?}"),
450 }
451 assert_eq!(started, Some(("ct_9f2c1a5b".to_string(), "42".to_string())));
452
453 match map_conversation_event(sse(WORKFLOW_STARTED), &mut started).unwrap() {
458 ConversationStreamEvent::WorkflowStarted(payload) => {
459 assert!(!payload.hit_cache);
460 assert_eq!(payload.inputs.chat_uid, "ct_9f2c1a5b");
461 assert_eq!(payload.inputs.message_id, "42");
462 assert_eq!(payload.workflow_id, 176476);
463 }
464 other => panic!("unexpected event: {other:?}"),
465 }
466
467 match map_conversation_event(sse(MESSAGE), &mut started).unwrap() {
468 ConversationStreamEvent::Message(payload) => assert_eq!(payload.text, "Tesla"),
469 other => panic!("unexpected event: {other:?}"),
470 }
471
472 match map_conversation_event(sse(PING), &mut started).unwrap() {
473 ConversationStreamEvent::Ping => {}
474 other => panic!("unexpected event: {other:?}"),
475 }
476
477 match map_conversation_event(sse(CHAT_FINISHED), &mut started).unwrap() {
478 ConversationStreamEvent::ChatFinished(payload) => {
479 assert_eq!(payload.chat_uid, "ct_9f2c1a5b");
480 assert_eq!(payload.message_id, "42");
481 assert_eq!(payload.error, "");
482 assert_eq!(payload.error_message, "");
483 }
484 other => panic!("unexpected event: {other:?}"),
485 }
486
487 match map_conversation_event(sse(WORKFLOW_FINISHED), &mut started).unwrap() {
488 ConversationStreamEvent::WorkflowFinished(resp) => {
489 assert_eq!(resp.chat_uid, "ct_9f2c1a5b");
490 assert_eq!(resp.message_id, "42");
491 assert_eq!(resp.status, ConversationStatus::Succeeded);
492 assert_eq!(resp.answer, "Tesla (TSLA.US) recently...");
493 }
494 other => panic!("unexpected event: {other:?}"),
495 }
496
497 match map_conversation_event(sse(CHAT_TITLE_UPDATED), &mut started).unwrap() {
499 ConversationStreamEvent::ChatTitleUpdated(payload) => {
500 assert_eq!(payload.chat_uid, "ct_9f2c1a5b");
501 assert_eq!(payload.source, "ai_generated");
502 assert_eq!(payload.title, "Tesla stock performance");
503 }
504 other => panic!("unexpected event: {other:?}"),
505 }
506 }
507
508 #[test]
509 fn map_conversation_event_unknown_type_falls_back_to_other() {
510 let mut started = None;
511 let json = r#"{"event":"some_future_event","data":{"foo":"bar"}}"#;
512 match map_conversation_event(sse(json), &mut started).unwrap() {
513 ConversationStreamEvent::Other { event, data } => {
514 assert_eq!(event, "some_future_event");
515 assert_eq!(data["foo"], "bar");
516 }
517 other => panic!("unexpected event: {other:?}"),
518 }
519 }
520
521 const HUMAN_INTERACTION_REQUIRED: &str = r#"{"event":"human_interaction_required","workflow_run_id":"wr_1","data":{"node_id":"n_ask_human","tool_call_id":"call_abc123","questions":[{"question":"Which time range would you like to check?","options":[{"description":"Past week"},{"description":"Past month"}],"multi_select":false}],"message_id":43,"chat_id":1001}}"#;
525
526 #[test]
527 fn map_conversation_event_interrupted_sequence_has_no_workflow_finished() {
528 let mut started = None;
529 map_conversation_event(sse(CHAT_STARTED), &mut started).unwrap();
530 map_conversation_event(sse(WORKFLOW_STARTED), &mut started).unwrap();
531
532 match map_conversation_event(sse(HUMAN_INTERACTION_REQUIRED), &mut started).unwrap() {
533 ConversationStreamEvent::HumanInteractionRequired(resp) => {
534 assert_eq!(resp.chat_uid, "ct_9f2c1a5b");
535 assert_eq!(resp.message_id, "42");
536 assert_eq!(resp.status, ConversationStatus::Interrupted);
537 let interrupt = resp.interrupt.expect("interrupt");
538 assert_eq!(interrupt.node_id, "n_ask_human");
539 assert_eq!(interrupt.tool_call_id, "call_abc123");
540 }
541 other => panic!("unexpected event: {other:?}"),
542 }
543
544 match map_conversation_event(sse(CHAT_FINISHED), &mut started).unwrap() {
547 ConversationStreamEvent::ChatFinished(_) => {}
548 other => panic!("unexpected event: {other:?}"),
549 }
550 }
551
552 #[test]
553 fn map_conversation_event_plan_changed_picks_up_sibling_tool_name() {
554 let mut started = None;
555 let json = r#"{"event":"plan_changed","workflow_run_id":"wr_1","tool_name":"planner","data":{"node_id":"n_plan","started_at":1752048000}}"#;
558 match map_conversation_event(sse(json), &mut started).unwrap() {
559 ConversationStreamEvent::PlanChanged(payload) => {
560 assert_eq!(payload.node_id, "n_plan");
561 assert_eq!(payload.tool_name, "planner");
562 }
563 other => panic!("unexpected event: {other:?}"),
564 }
565 }
566}