Skip to main content

longbridge/agent/
context.rs

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
10/// The shared httpclient default (30s) is tuned for fast REST calls and is
11/// too tight here: in blocking mode the server holds the connection silent
12/// until the whole LLM turn is done, and that can legitimately take longer.
13/// Only agent calls get this longer budget — every other domain keeps the
14/// 30s default.
15const 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/// AI Agent conversation context.
31///
32/// Reference: <https://open.longbridge.com/en/docs/ai/chat/conversation>
33#[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    /// Only present on `plan_changed`, as a sibling of `data` rather than a
42    /// field inside it — see [`PlanChangedPayload::tool_name`].
43    #[serde(default)]
44    tool_name: Option<String>,
45}
46
47/// Parse one raw SSE frame into a [`ConversationStreamEvent`], threading the
48/// `chat_uid`/`message_id` captured from an earlier `chat_started` event (the
49/// `workflow_finished` event doesn't repeat them) through `started`.
50fn 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    /// Create an [`AgentContext`]
139    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    /// Returns the log subscriber
155    #[inline]
156    pub fn log_subscriber(&self) -> Arc<dyn Subscriber + Send + Sync> {
157        self.0.log_subscriber.clone()
158    }
159
160    /// List the Workspaces the current account belongs to.
161    ///
162    /// Path: `GET /v1/ai/workspaces`
163    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    /// List the Agents in the specified Workspace.
176    ///
177    /// Path: `GET /v1/ai/workspaces/{id}/agents`
178    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    /// List all publicly available Agents on the platform — the same catalog
200    /// shown on the Explore page.
201    ///
202    /// Unlike [`agents`](Self::agents), this endpoint is not scoped to a
203    /// Workspace: it returns every Agent that is published and publicly shared.
204    /// The returned [`Agent::uid`] is used as the path parameter of
205    /// [`conversation`](Self::conversation).
206    ///
207    /// Path: `GET /v1/ai/agents`
208    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    /// Start a conversation with the specified Agent, blocking until the run
225    /// succeeds, is interrupted, or fails.
226    ///
227    /// `parent_message_id` is the `message_id` from a previous response. Pass
228    /// it when asking a follow-up in an existing conversation to attach the new
229    /// message after the specified one, keeping the message stream in order. It
230    /// is only valid together with `chat_uid`, the parent message must belong
231    /// to that conversation, and it must not be set for a new conversation.
232    ///
233    /// Path: `POST /v1/ai/agents/{id}/conversations`
234    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    /// Resume an interrupted conversation, blocking until the run succeeds, is
273    /// interrupted again, or fails.
274    ///
275    /// Path: `POST
276    /// /v1/ai/agents/{id}/conversations/{chat_uid}/messages/{message_id}/
277    /// continue`
278    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    /// Start a conversation with the specified Agent, returning a [`Stream`] of
315    /// run-progress events over SSE. The run's outcome is carried by a
316    /// [`ConversationStreamEvent::WorkflowFinished`] event (succeeded, failed,
317    /// or stopped) or, if the Agent needs more input from you, a
318    /// [`ConversationStreamEvent::HumanInteractionRequired`] event instead —
319    /// an interrupted run never emits `WorkflowFinished`. Neither is
320    /// necessarily the last item — the server may still emit a few more
321    /// housekeeping events (e.g.
322    /// [`ConversationStreamEvent::ChatTitleUpdated`]) before actually closing
323    /// the connection, so keep draining the stream until it ends rather than
324    /// stopping as soon as you see one.
325    ///
326    /// `parent_message_id` behaves as in [`Self::conversation`].
327    ///
328    /// Path: `POST /v1/ai/agents/{id}/conversations` (`Accept:
329    /// text/event-stream`)
330    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    /// Resume an interrupted conversation, returning a [`Stream`] of
369    /// run-progress events over SSE.
370    ///
371    /// Path: `POST
372    /// /v1/ai/agents/{id}/conversations/{chat_uid}/messages/{message_id}/
373    /// continue` (`Accept: text/event-stream`)
374    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        // We already know chat_uid/message_id from the caller (unlike a brand-new
390        // conversation) — seed `started` so the final ConversationResponse carries
391        // them even if the server doesn't re-emit a `chat_started` event here.
392        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    // The `data:` payloads of the three example SSE frames from
419    // https://open.longbridge.com/en/docs/ai/chat/conversation
420    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    // The four event types below aren't in the docs — captured verbatim from
425    // real traffic during manual live testing (see conversation history).
426    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        // The real event stream is richer than the docs' three-event example
454        // — this exercises the fuller, real-world sequence, including
455        // `chat_title_updated` arriving *after* `workflow_finished` (observed
456        // live; see the "drain to the stream's natural end" fix).
457        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        // Arrives *after* workflow_finished in this (real, observed) ordering.
498        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    // https://github.com/longbridge/developers/pull/1176 — an interrupted
522    // run's stream never emits `workflow_finished`; `human_interaction_required`
523    // is the terminal event instead.
524    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        // The stream still ends with `chat_finished`, just never emits
545        // `workflow_finished`.
546        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        // `tool_name` sits outside `data`, as a sibling of `event`/`data` in
556        // the raw envelope.
557        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}