1use crate::Error;
2use crate::Result;
3use crate::RunCollectTextError;
4use crate::builder::AgentDefinition;
5use crate::events::AgentEvent;
6use crate::events::TurnItemDelta;
7use crate::events::TurnSummary;
8use crate::input::InputQueue;
9use crate::input::SessionInput;
10use crate::kernel::AgentKernel;
11use crate::processor::SessionTurnProcessor;
12use crate::processor::outcome_summary;
13use crate::snapshot::ActiveTurnSnapshot;
14use crate::snapshot::SessionSnapshot;
15use crate::status::SessionStatus;
16use crate::tools::ToolExecutor;
17use async_channel::Receiver;
18use async_channel::Sender;
19use lha_llm::BaseInstructions;
20use lha_llm::SemanticRuntimeSession;
21use lha_llm::ToolDescriptor;
22use lha_llm::TranscriptItem;
23use lha_llm::TurnRequest;
24use std::collections::VecDeque;
25use std::sync::Arc;
26use std::sync::atomic::AtomicU64;
27use std::sync::atomic::Ordering;
28use tokio::sync::Mutex;
29use tokio_util::sync::CancellationToken;
30
31pub type SessionId = u64;
32pub type SubmissionId = u64;
33
34#[derive(Clone)]
35pub struct AgentSession {
36 inner: Arc<AgentSessionInner>,
37}
38
39pub(crate) struct AgentSessionInner {
40 pub(crate) session_id: SessionId,
41 definition: Arc<AgentDefinition>,
42 runtime_session: Mutex<Box<dyn SemanticRuntimeSession>>,
43 state: Mutex<SessionState>,
44 tx_event: Sender<AgentEvent>,
45 rx_event: Receiver<AgentEvent>,
46 next_submission_id: AtomicU64,
47}
48
49struct SessionState {
50 conversation: Vec<TranscriptItem>,
51 steering_queue: VecDeque<SessionInput>,
52 follow_up_queue: VecDeque<SessionInput>,
53 status: SessionStatus,
54 active_turn: Option<ActiveTurnState>,
55}
56
57struct ActiveTurnState {
58 submission_id: SubmissionId,
59 cancellation_token: CancellationToken,
60}
61
62impl AgentSession {
63 pub(crate) fn new(
64 session_id: SessionId,
65 definition: Arc<AgentDefinition>,
66 conversation: Vec<TranscriptItem>,
67 ) -> Self {
68 let (tx_event, rx_event) = async_channel::unbounded();
69 let inner = Arc::new(AgentSessionInner {
70 session_id,
71 runtime_session: Mutex::new(definition.runtime.new_session()),
72 definition,
73 state: Mutex::new(SessionState {
74 conversation,
75 steering_queue: VecDeque::new(),
76 follow_up_queue: VecDeque::new(),
77 status: SessionStatus::Idle,
78 active_turn: None,
79 }),
80 tx_event,
81 rx_event,
82 next_submission_id: AtomicU64::new(1),
83 });
84
85 let _ = inner
86 .tx_event
87 .try_send(AgentEvent::SessionStarted { session_id });
88 let _ = inner.tx_event.try_send(AgentEvent::SessionStatusChanged {
89 session_id,
90 status: SessionStatus::Idle,
91 });
92
93 Self { inner }
94 }
95
96 pub fn id(&self) -> SessionId {
97 self.inner.session_id
98 }
99
100 pub async fn run(&self, input: SessionInput) -> Result<SubmissionId> {
101 let items = input.into_items();
102 self.inner
103 .spawn_turn_loop(items.clone(), Some((InputQueue::Primary, items)))
104 .await
105 }
106
107 pub async fn continue_turn(&self) -> Result<SubmissionId> {
108 {
109 let state = self.inner.state.lock().await;
110 if state.active_turn.is_some() {
111 return Err(Error::SessionBusy);
112 }
113 if state.conversation.is_empty() && !state.has_bootstrap_input() {
114 return Err(Error::EmptyConversation);
115 }
116 if !state.has_bootstrap_input() && !can_continue_from_history(&state.conversation) {
117 return Err(Error::InvalidContinuation);
118 }
119 }
120 self.inner.spawn_turn_loop(Vec::new(), None).await
121 }
122
123 pub async fn steer(&self, input: SessionInput) {
124 let items = input.items().to_vec();
125 {
126 let mut state = self.inner.state.lock().await;
127 state.steering_queue.push_back(input);
128 }
129 self.inner
130 .emit_event(AgentEvent::InputQueued {
131 session_id: self.inner.session_id,
132 queue: InputQueue::Steering,
133 items,
134 })
135 .await;
136 }
137
138 pub async fn follow_up(&self, input: SessionInput) {
139 let items = input.items().to_vec();
140 {
141 let mut state = self.inner.state.lock().await;
142 state.follow_up_queue.push_back(input);
143 }
144 self.inner
145 .emit_event(AgentEvent::InputQueued {
146 session_id: self.inner.session_id,
147 queue: InputQueue::FollowUp,
148 items,
149 })
150 .await;
151 }
152
153 pub async fn abort_current_turn(&self) -> bool {
154 let cancellation_token = {
155 let mut state = self.inner.state.lock().await;
156 let Some(cancellation_token) = state
157 .active_turn
158 .as_ref()
159 .map(|active_turn| active_turn.cancellation_token.clone())
160 else {
161 return false;
162 };
163 state.status = SessionStatus::Aborting;
164 cancellation_token
165 };
166 self.inner
167 .emit_event(AgentEvent::SessionStatusChanged {
168 session_id: self.inner.session_id,
169 status: SessionStatus::Aborting,
170 })
171 .await;
172 cancellation_token.cancel();
173 true
174 }
175
176 pub async fn next_event(&self) -> Result<AgentEvent> {
177 self.inner
178 .rx_event
179 .recv()
180 .await
181 .map_err(|_| Error::EventChannelClosed)
182 }
183
184 pub async fn run_collect_text(
185 &self,
186 input: impl Into<SessionInput>,
187 ) -> std::result::Result<String, RunCollectTextError> {
188 let submission_id = self.run(input.into()).await?;
189 let mut current_turn_streamed_text = String::new();
190 let mut current_turn_saw_output_text_delta = false;
191 let mut terminal_result = None;
192
193 loop {
194 match self.next_event().await? {
195 AgentEvent::TurnStarted {
196 submission_id: event_submission_id,
197 ..
198 } if event_submission_id == submission_id && terminal_result.is_none() => {
199 current_turn_streamed_text.clear();
200 current_turn_saw_output_text_delta = false;
201 }
202 AgentEvent::OutputItemDelta {
203 submission_id: event_submission_id,
204 delta: TurnItemDelta::OutputText { delta },
205 ..
206 } if event_submission_id == submission_id && terminal_result.is_none() => {
207 current_turn_saw_output_text_delta = true;
208 current_turn_streamed_text.push_str(&delta);
209 }
210 AgentEvent::TurnCompleted {
211 submission_id: event_submission_id,
212 outcome,
213 ..
214 } if event_submission_id == submission_id && terminal_result.is_none() => {
215 if !outcome.needs_follow_up {
216 let final_text = if current_turn_saw_output_text_delta {
217 current_turn_streamed_text.clone()
218 } else {
219 outcome.last_agent_message.unwrap_or_default()
220 };
221 terminal_result = Some(Ok(final_text));
222 }
223 }
224 AgentEvent::TurnFailed {
225 submission_id: event_submission_id,
226 error,
227 ..
228 } if event_submission_id == submission_id && terminal_result.is_none() => {
229 terminal_result = Some(Err(RunCollectTextError::TurnFailed(error)));
230 }
231 AgentEvent::TurnAborted {
232 submission_id: event_submission_id,
233 ..
234 } if event_submission_id == submission_id && terminal_result.is_none() => {
235 terminal_result = Some(Err(Error::Aborted.into()));
236 }
237 AgentEvent::SessionStatusChanged {
238 status: SessionStatus::Idle,
239 ..
240 } if terminal_result.is_some() => {
241 if let Some(result) = terminal_result.take() {
242 return result;
243 }
244 }
245 _ => {}
246 }
247 }
248 }
249
250 pub async fn status(&self) -> SessionStatus {
251 let state = self.inner.state.lock().await;
252 state.status
253 }
254
255 pub async fn snapshot(&self) -> SessionSnapshot {
256 self.inner.snapshot().await
257 }
258}
259
260impl AgentSessionInner {
261 pub(crate) async fn emit_event(&self, event: AgentEvent) {
262 let _ = self.tx_event.send(event).await;
263 }
264
265 pub(crate) async fn push_conversation_item(&self, item: TranscriptItem) {
266 let mut state = self.state.lock().await;
267 state.conversation.push(item);
268 }
269
270 async fn snapshot(&self) -> SessionSnapshot {
271 let state = self.state.lock().await;
272 SessionSnapshot {
273 session_id: self.session_id,
274 status: state.status,
275 conversation: state.conversation.clone(),
276 steering_queue: state
277 .steering_queue
278 .iter()
279 .map(|input| input.items().to_vec())
280 .collect(),
281 follow_up_queue: state
282 .follow_up_queue
283 .iter()
284 .map(|input| input.items().to_vec())
285 .collect(),
286 runtime: self.definition.runtime_metadata.clone(),
287 active_turn: state.active_turn.as_ref().map(|active| ActiveTurnSnapshot {
288 submission_id: active.submission_id,
289 }),
290 }
291 }
292
293 async fn spawn_turn_loop(
294 self: &Arc<Self>,
295 initial_items: Vec<TranscriptItem>,
296 queued_event: Option<(InputQueue, Vec<TranscriptItem>)>,
297 ) -> Result<SubmissionId> {
298 let cancellation_token = CancellationToken::new();
299
300 let submission_id = {
301 let mut state = self.state.lock().await;
302 if state.active_turn.is_some() {
303 return Err(Error::SessionBusy);
304 }
305 let submission_id = self.next_submission_id.fetch_add(1, Ordering::SeqCst);
306 state.status = SessionStatus::Running;
307 state.active_turn = Some(ActiveTurnState {
308 submission_id,
309 cancellation_token: cancellation_token.clone(),
310 });
311 submission_id
312 };
313
314 if let Some((queue, items)) = queued_event {
315 self.emit_event(AgentEvent::InputQueued {
316 session_id: self.session_id,
317 queue,
318 items,
319 })
320 .await;
321 }
322
323 self.emit_event(AgentEvent::SessionStatusChanged {
324 session_id: self.session_id,
325 status: SessionStatus::Running,
326 })
327 .await;
328
329 let session = Arc::clone(self);
330 tokio::spawn(async move {
331 session
332 .drive_turn_loop(submission_id, initial_items, cancellation_token)
333 .await;
334 });
335
336 Ok(submission_id)
337 }
338
339 async fn drive_turn_loop(
340 self: Arc<Self>,
341 submission_id: SubmissionId,
342 mut initial_items: Vec<TranscriptItem>,
343 cancellation_token: CancellationToken,
344 ) {
345 loop {
346 if cancellation_token.is_cancelled() {
347 self.emit_event(AgentEvent::TurnAborted {
348 session_id: self.session_id,
349 submission_id,
350 })
351 .await;
352 break;
353 }
354
355 if !initial_items.is_empty() {
356 self.append_items(std::mem::take(&mut initial_items)).await;
357 } else if !self.drain_bootstrap_inputs().await {
358 let state = self.state.lock().await;
359 if state.conversation.is_empty() {
360 drop(state);
361 self.emit_event(AgentEvent::TurnFailed {
362 session_id: self.session_id,
363 submission_id,
364 error: Error::EmptyConversation.to_string(),
365 })
366 .await;
367 break;
368 }
369 }
370
371 self.emit_event(AgentEvent::TurnStarted {
372 session_id: self.session_id,
373 submission_id,
374 })
375 .await;
376
377 let request = match self.build_turn_request().await {
378 Ok(request) => request,
379 Err(err) => {
380 self.emit_event(AgentEvent::TurnFailed {
381 session_id: self.session_id,
382 submission_id,
383 error: err.to_string(),
384 })
385 .await;
386 break;
387 }
388 };
389 let stream = {
390 let mut runtime_session = self.runtime_session.lock().await;
391 runtime_session.run_turn(&request).await
392 };
393
394 let result = match stream {
395 Ok(stream) => {
396 let processor = SessionTurnProcessor::new(
397 Arc::clone(&self),
398 submission_id,
399 ToolExecutor::new(Arc::clone(&self.definition.tools)),
400 cancellation_token.child_token(),
401 );
402 AgentKernel::new()
403 .run_turn(stream, processor, cancellation_token.child_token())
404 .await
405 .map(outcome_summary)
406 }
407 Err(err) => Err(Error::Runtime(err)),
408 };
409 let result = if cancellation_token.is_cancelled() && result.is_err() {
410 Err(Error::Aborted)
411 } else {
412 result
413 };
414
415 match result {
416 Ok(outcome) => {
417 self.emit_event(AgentEvent::TurnCompleted {
418 session_id: self.session_id,
419 submission_id,
420 outcome: TurnSummary {
421 needs_follow_up: outcome.needs_follow_up,
422 last_agent_message: outcome.last_agent_message.clone(),
423 response_total_tokens: outcome.response_total_tokens,
424 tool_output_tokens: outcome.tool_output_tokens,
425 },
426 })
427 .await;
428
429 if cancellation_token.is_cancelled() {
430 self.emit_event(AgentEvent::TurnAborted {
431 session_id: self.session_id,
432 submission_id,
433 })
434 .await;
435 break;
436 }
437
438 if outcome.needs_follow_up {
439 continue;
440 }
441
442 if self.drain_input_queue(InputQueue::Steering).await {
443 continue;
444 }
445 if self.drain_input_queue(InputQueue::FollowUp).await {
446 continue;
447 }
448 break;
449 }
450 Err(Error::Aborted) => {
451 self.emit_event(AgentEvent::TurnAborted {
452 session_id: self.session_id,
453 submission_id,
454 })
455 .await;
456 break;
457 }
458 Err(err) => {
459 self.emit_event(AgentEvent::TurnFailed {
460 session_id: self.session_id,
461 submission_id,
462 error: err.to_string(),
463 })
464 .await;
465 break;
466 }
467 }
468 }
469
470 let mut state = self.state.lock().await;
471 state.status = SessionStatus::Idle;
472 state.active_turn = None;
473 drop(state);
474
475 self.emit_event(AgentEvent::SessionStatusChanged {
476 session_id: self.session_id,
477 status: SessionStatus::Idle,
478 })
479 .await;
480 }
481
482 async fn build_turn_request(&self) -> Result<TurnRequest> {
483 let (conversation, tools) = {
484 let state = self.state.lock().await;
485 (state.conversation.clone(), self.definition.tools.specs())
486 };
487 let base_instructions = self
488 .base_instructions_for_turn(&conversation, &tools)
489 .await?;
490 Ok(TurnRequest {
491 conversation,
492 tools,
493 parallel_tool_calls: self.definition.tools.any_parallel_tool_calls(),
494 base_instructions,
495 personality: self.definition.personality,
496 output_schema: self.definition.output_schema.clone(),
497 })
498 }
499
500 async fn base_instructions_for_turn(
501 &self,
502 conversation: &[TranscriptItem],
503 tools: &[ToolDescriptor],
504 ) -> Result<BaseInstructions> {
505 if self.definition.skill_providers.is_empty() {
506 return Ok(self.definition.base_instructions.clone());
507 }
508
509 let context = crate::skills::SkillContext {
510 session_id: self.session_id,
511 conversation: conversation.to_vec(),
512 runtime: self.definition.runtime_metadata.clone(),
513 tools: tools.to_vec(),
514 };
515 let mut skills = Vec::new();
516 for provider in &self.definition.skill_providers {
517 skills.extend(provider.skills_for_turn(&context).await?);
518 }
519
520 if skills.is_empty() {
521 return Ok(self.definition.base_instructions.clone());
522 }
523
524 let mut base_instructions = self.definition.base_instructions.clone();
525 base_instructions.text.push_str("\n\n<available_skills>\n");
526 for skill in skills {
527 base_instructions.text.push_str("<skill>\n");
528 base_instructions.text.push_str("id: ");
529 base_instructions.text.push_str(&skill.id);
530 base_instructions.text.push('\n');
531 base_instructions.text.push_str("name: ");
532 base_instructions.text.push_str(&skill.name);
533 base_instructions.text.push('\n');
534 if let Some(description) = skill.description {
535 base_instructions.text.push_str("description: ");
536 base_instructions.text.push_str(&description);
537 base_instructions.text.push('\n');
538 }
539 if !skill.required_tools.is_empty() {
540 base_instructions.text.push_str("required_tools: ");
541 base_instructions
542 .text
543 .push_str(&skill.required_tools.join(", "));
544 base_instructions.text.push('\n');
545 }
546 base_instructions.text.push_str("instructions:\n");
547 base_instructions.text.push_str(&skill.instructions);
548 base_instructions.text.push_str("\n</skill>\n");
549 }
550 base_instructions.text.push_str("</available_skills>");
551 Ok(base_instructions)
552 }
553
554 async fn append_items(&self, items: Vec<TranscriptItem>) {
555 let mut state = self.state.lock().await;
556 state.conversation.extend(items);
557 }
558
559 async fn drain_bootstrap_inputs(&self) -> bool {
560 if self.drain_input_queue(InputQueue::Steering).await {
561 true
562 } else {
563 self.drain_input_queue(InputQueue::FollowUp).await
564 }
565 }
566
567 async fn drain_input_queue(&self, queue: InputQueue) -> bool {
568 let items = {
569 let mut state = self.state.lock().await;
570 let source = match queue {
571 InputQueue::Primary => return false,
572 InputQueue::Steering => &mut state.steering_queue,
573 InputQueue::FollowUp => &mut state.follow_up_queue,
574 };
575 let mut drained = Vec::new();
576 while let Some(input) = source.pop_front() {
577 drained.extend(input.into_items());
578 }
579 drained
580 };
581
582 if items.is_empty() {
583 false
584 } else {
585 self.append_items(items).await;
586 true
587 }
588 }
589}
590
591impl SessionState {
592 fn has_bootstrap_input(&self) -> bool {
593 !self.steering_queue.is_empty() || !self.follow_up_queue.is_empty()
594 }
595}
596
597fn can_continue_from_history(history: &[TranscriptItem]) -> bool {
598 !matches!(
599 history.last(),
600 Some(TranscriptItem::Message { role, .. }) if role == "assistant"
601 )
602}
603
604#[cfg(test)]
605mod tests {
606 use super::*;
607 use crate::AgentBuilder;
608 use crate::skills::Skill;
609 use crate::skills::SkillContext;
610 use crate::skills::SkillError;
611 use crate::skills::SkillProvider;
612 use crate::tools::ToolError;
613 use crate::tools::ToolHandler;
614 use crate::tools::ToolInvocation;
615 use crate::tools::ToolOutput;
616 use async_trait::async_trait;
617 use lha_llm::BaseInstructions;
618 use lha_llm::FunctionToolDescriptor;
619 use lha_llm::RuntimeMetadata;
620 use lha_llm::SemanticConversationCompactor;
621 use lha_llm::SemanticRuntime;
622 use lha_llm::SemanticRuntimeSession;
623 use lha_llm::ToolCallPayload;
624 use lha_llm::ToolCallRequest;
625 use lha_llm::ToolDescriptor;
626 use lha_llm::ToolInputSchema;
627 use lha_llm::ToolResultPayload;
628 use lha_llm::TranscriptItem;
629 use lha_llm::TurnEvent;
630 use lha_llm::TurnEventStream;
631 use lha_llm::types::ContentItem;
632 use lha_llm::types::TokenUsage;
633 use pretty_assertions::assert_eq;
634 use std::collections::BTreeMap;
635 use std::collections::VecDeque;
636 use std::sync::Mutex as StdMutex;
637 use tokio::sync::Notify;
638 use tokio::sync::mpsc;
639 use tokio::sync::oneshot;
640 use tokio::time::Duration;
641 use tokio::time::timeout;
642 use tokio_util::sync::CancellationToken;
643
644 #[derive(Clone)]
645 struct FakeRuntime {
646 scripts: Arc<StdMutex<VecDeque<FakeTurnScript>>>,
647 requests: Arc<StdMutex<Vec<TurnRequest>>>,
648 }
649
650 struct FakeTurnScript {
651 events: Vec<lha_llm::Result<TurnEvent>>,
652 gate: Option<FakeGate>,
653 hold_open: Option<Arc<Notify>>,
654 }
655
656 struct FakeGate {
657 release: Arc<Notify>,
658 ready: Option<oneshot::Sender<()>>,
659 }
660
661 struct FakeRuntimeSession {
662 scripts: Arc<StdMutex<VecDeque<FakeTurnScript>>>,
663 requests: Arc<StdMutex<Vec<TurnRequest>>>,
664 }
665
666 #[async_trait]
667 impl SemanticConversationCompactor for FakeRuntime {
668 async fn compact_conversation_history(
669 &self,
670 input: &TurnRequest,
671 ) -> lha_llm::Result<Vec<TranscriptItem>> {
672 Ok(input.conversation.clone())
673 }
674 }
675
676 #[async_trait]
677 impl SemanticRuntime for FakeRuntime {
678 fn new_session(&self) -> Box<dyn SemanticRuntimeSession> {
679 Box::new(FakeRuntimeSession {
680 scripts: Arc::clone(&self.scripts),
681 requests: Arc::clone(&self.requests),
682 })
683 }
684
685 fn capabilities(&self) -> lha_llm::RuntimeCapabilities {
686 lha_llm::RuntimeCapabilities {
687 supports_parallel_tool_calls: true,
688 enforce_declared_tool_names: false,
689 supports_dynamic_context_window_probe: false,
690 supports_reasoning_summaries: true,
691 supports_output_schema: true,
692 supports_remote_compaction: false,
693 }
694 }
695
696 fn metadata(&self) -> RuntimeMetadata {
697 RuntimeMetadata {
698 endpoint_name: "fake".to_string(),
699 model: "test-model".to_string(),
700 }
701 }
702
703 fn estimated_input_tokens(&self, _input: &TurnRequest) -> Option<i64> {
704 None
705 }
706 }
707
708 #[async_trait]
709 impl SemanticRuntimeSession for FakeRuntimeSession {
710 async fn run_turn(&mut self, input: &TurnRequest) -> lha_llm::Result<TurnEventStream> {
711 self.requests
712 .lock()
713 .unwrap_or_else(std::sync::PoisonError::into_inner)
714 .push(input.clone());
715 let script = self
716 .scripts
717 .lock()
718 .unwrap_or_else(std::sync::PoisonError::into_inner)
719 .pop_front()
720 .expect("script should exist");
721 let (tx, rx) = mpsc::channel(16);
722 tokio::spawn(async move {
723 let FakeTurnScript {
724 events,
725 gate,
726 hold_open,
727 } = script;
728 let mut gate = gate;
729
730 for event in events {
731 if tx.send(event).await.is_err() {
732 return;
733 }
734 if let Some(gate) = gate.take() {
735 if let Some(ready) = gate.ready {
736 let _ = ready.send(());
737 }
738 gate.release.notified().await;
739 }
740 }
741 if let Some(hold_open) = hold_open.as_ref() {
742 hold_open.notified().await;
743 }
744 });
745 Ok(TurnEventStream::from_receiver(rx))
746 }
747 }
748
749 struct EchoTool;
750
751 #[async_trait]
752 impl ToolHandler for EchoTool {
753 fn spec(&self) -> ToolDescriptor {
754 ToolDescriptor::Function(FunctionToolDescriptor {
755 name: "echo_tool".to_string(),
756 description: "echo tool".to_string(),
757 strict: false,
758 parameters: ToolInputSchema::Object {
759 properties: BTreeMap::new(),
760 required: Some(Vec::new()),
761 additional_properties: Some(true.into()),
762 },
763 })
764 }
765
766 async fn handle(
767 &self,
768 invocation: ToolInvocation,
769 _cancellation_token: CancellationToken,
770 ) -> std::result::Result<ToolOutput, ToolError> {
771 Ok(ToolOutput::Function {
772 content: format!("handled {}", invocation.tool_name),
773 content_items: None,
774 success: Some(true),
775 })
776 }
777 }
778
779 fn fake_runtime(scripts: Vec<FakeTurnScript>) -> Arc<dyn SemanticRuntime> {
780 fake_runtime_with_requests(scripts).0
781 }
782
783 fn fake_runtime_with_requests(
784 scripts: Vec<FakeTurnScript>,
785 ) -> (Arc<dyn SemanticRuntime>, Arc<StdMutex<Vec<TurnRequest>>>) {
786 let requests = Arc::new(StdMutex::new(Vec::new()));
787 let runtime: Arc<dyn SemanticRuntime> = Arc::new(FakeRuntime {
788 scripts: Arc::new(StdMutex::new(VecDeque::from(scripts))),
789 requests: Arc::clone(&requests),
790 });
791 (runtime, requests)
792 }
793
794 struct TestSkillProvider {
795 skills: Vec<Skill>,
796 error: Option<String>,
797 contexts: Arc<StdMutex<Vec<SkillContext>>>,
798 }
799
800 #[async_trait]
801 impl SkillProvider for TestSkillProvider {
802 async fn skills_for_turn(
803 &self,
804 context: &SkillContext,
805 ) -> std::result::Result<Vec<Skill>, SkillError> {
806 self.contexts
807 .lock()
808 .unwrap_or_else(std::sync::PoisonError::into_inner)
809 .push(context.clone());
810 if let Some(error) = &self.error {
811 return Err(SkillError::Fatal(error.clone()));
812 }
813 Ok(self.skills.clone())
814 }
815 }
816
817 fn test_skill(id: &str, instructions: &str) -> Skill {
818 Skill {
819 id: id.to_string(),
820 name: format!("{id} skill"),
821 description: Some(format!("{id} description")),
822 instructions: instructions.to_string(),
823 required_tools: vec!["echo_tool".to_string()],
824 }
825 }
826
827 fn completed_script() -> FakeTurnScript {
828 FakeTurnScript {
829 events: vec![Ok(TurnEvent::Completed {
830 response_id: "resp-1".to_string(),
831 token_usage: None,
832 })],
833 gate: None,
834 hold_open: None,
835 }
836 }
837
838 fn lock_requests(requests: &StdMutex<Vec<TurnRequest>>) -> Vec<TurnRequest> {
839 requests
840 .lock()
841 .unwrap_or_else(std::sync::PoisonError::into_inner)
842 .clone()
843 }
844
845 fn assistant_item(text: &str) -> lha_llm::SemanticOutputItem {
846 assistant_item_with_output_texts(&[text])
847 }
848
849 fn assistant_item_with_output_texts(texts: &[&str]) -> lha_llm::SemanticOutputItem {
850 lha_llm::SemanticOutputItem::AssistantMessage {
851 item: TranscriptItem::Message {
852 id: Some("msg-1".to_string()),
853 role: "assistant".to_string(),
854 content: texts
855 .iter()
856 .map(|text| ContentItem::OutputText {
857 text: (*text).to_string(),
858 })
859 .collect(),
860 end_turn: None,
861 },
862 }
863 }
864
865 async fn wait_for_idle(session: &AgentSession) -> Vec<AgentEvent> {
866 let mut events = Vec::new();
867 let mut seen_running = false;
868 loop {
869 let event = timeout(Duration::from_secs(2), session.next_event())
870 .await
871 .expect("event should arrive")
872 .expect("session event should succeed");
873 if matches!(
874 &event,
875 AgentEvent::SessionStatusChanged {
876 status: SessionStatus::Running,
877 ..
878 }
879 ) {
880 seen_running = true;
881 }
882 let is_idle = seen_running
883 && matches!(
884 &event,
885 AgentEvent::SessionStatusChanged {
886 status: SessionStatus::Idle,
887 ..
888 }
889 );
890 events.push(event);
891 if is_idle {
892 return events;
893 }
894 }
895 }
896
897 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
898 async fn run_records_assistant_message_and_completes() {
899 let runtime = fake_runtime(vec![FakeTurnScript {
900 events: vec![
901 Ok(TurnEvent::ItemStarted {
902 handle: "msg-1".to_string(),
903 item: assistant_item("hello"),
904 }),
905 Ok(TurnEvent::OutputTextDelta {
906 handle: "msg-1".to_string(),
907 delta: "hello".to_string(),
908 }),
909 Ok(TurnEvent::ItemCompleted {
910 handle: "msg-1".to_string(),
911 item: assistant_item("hello"),
912 }),
913 Ok(TurnEvent::Completed {
914 response_id: "resp-1".to_string(),
915 token_usage: Some(TokenUsage {
916 input_tokens: 1,
917 cached_input_tokens: 0,
918 output_tokens: 2,
919 reasoning_output_tokens: 0,
920 total_tokens: 3,
921 }),
922 }),
923 ],
924 gate: None,
925 hold_open: None,
926 }]);
927
928 let manager = AgentBuilder::new(runtime).build();
929 let session = manager.create_session();
930
931 session
932 .run(SessionInput::from_user_text("hi"))
933 .await
934 .expect("run should succeed");
935 let events = wait_for_idle(&session).await;
936 let event = events
937 .into_iter()
938 .find(|event| matches!(event, AgentEvent::TurnCompleted { .. }))
939 .expect("turn completion should be emitted");
940
941 match event {
942 AgentEvent::TurnCompleted { outcome, .. } => {
943 assert_eq!(outcome.last_agent_message, Some("hello".to_string()));
944 assert_eq!(outcome.response_total_tokens, Some(3));
945 }
946 other => panic!("expected TurnCompleted, got {other:?}"),
947 }
948
949 let snapshot = session.snapshot().await;
950 assert_eq!(snapshot.status, SessionStatus::Idle);
951 assert_eq!(snapshot.conversation.len(), 2);
952 }
953
954 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
955 async fn turn_completed_summary_joins_output_text_blocks() {
956 let runtime = fake_runtime(vec![FakeTurnScript {
957 events: vec![
958 Ok(TurnEvent::ItemStarted {
959 handle: "msg-1".to_string(),
960 item: assistant_item_with_output_texts(&[]),
961 }),
962 Ok(TurnEvent::ItemCompleted {
963 handle: "msg-1".to_string(),
964 item: assistant_item_with_output_texts(&["hel", "lo"]),
965 }),
966 Ok(TurnEvent::Completed {
967 response_id: "resp-1".to_string(),
968 token_usage: None,
969 }),
970 ],
971 gate: None,
972 hold_open: None,
973 }]);
974
975 let manager = AgentBuilder::new(runtime).build();
976 let session = manager.create_session();
977
978 session
979 .run(SessionInput::from_user_text("hi"))
980 .await
981 .expect("run should succeed");
982 let events = wait_for_idle(&session).await;
983 let event = events
984 .into_iter()
985 .find(|event| matches!(event, AgentEvent::TurnCompleted { .. }))
986 .expect("turn completion should be emitted");
987
988 match event {
989 AgentEvent::TurnCompleted { outcome, .. } => {
990 assert_eq!(outcome.last_agent_message, Some("hello".to_string()));
991 }
992 other => panic!("expected TurnCompleted, got {other:?}"),
993 }
994 }
995
996 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
997 async fn run_collect_text_returns_last_agent_message() {
998 let runtime = fake_runtime(vec![FakeTurnScript {
999 events: vec![
1000 Ok(TurnEvent::ItemStarted {
1001 handle: "msg-1".to_string(),
1002 item: assistant_item("hello"),
1003 }),
1004 Ok(TurnEvent::ItemCompleted {
1005 handle: "msg-1".to_string(),
1006 item: assistant_item("hello"),
1007 }),
1008 Ok(TurnEvent::Completed {
1009 response_id: "resp-1".to_string(),
1010 token_usage: None,
1011 }),
1012 ],
1013 gate: None,
1014 hold_open: None,
1015 }]);
1016 let manager = AgentBuilder::new(runtime).build();
1017 let session = manager.create_session();
1018
1019 let text = session
1020 .run_collect_text(SessionInput::from_user_text("hi"))
1021 .await
1022 .expect("text should be collected");
1023
1024 assert_eq!(text, "hello");
1025 }
1026
1027 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1028 async fn run_collect_text_joins_completed_output_text_blocks_without_deltas() {
1029 let runtime = fake_runtime(vec![FakeTurnScript {
1030 events: vec![
1031 Ok(TurnEvent::ItemStarted {
1032 handle: "msg-1".to_string(),
1033 item: assistant_item_with_output_texts(&[]),
1034 }),
1035 Ok(TurnEvent::ItemCompleted {
1036 handle: "msg-1".to_string(),
1037 item: assistant_item_with_output_texts(&["hel", "lo"]),
1038 }),
1039 Ok(TurnEvent::Completed {
1040 response_id: "resp-1".to_string(),
1041 token_usage: None,
1042 }),
1043 ],
1044 gate: None,
1045 hold_open: None,
1046 }]);
1047 let manager = AgentBuilder::new(runtime).build();
1048 let session = manager.create_session();
1049
1050 let text = session
1051 .run_collect_text(SessionInput::from_user_text("hi"))
1052 .await
1053 .expect("text should be collected");
1054
1055 assert_eq!(text, "hello");
1056 }
1057
1058 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1059 async fn run_collect_text_prefers_streamed_deltas_over_chunked_last_message() {
1060 let runtime = fake_runtime(vec![FakeTurnScript {
1061 events: vec![
1062 Ok(TurnEvent::ItemStarted {
1063 handle: "msg-1".to_string(),
1064 item: assistant_item_with_output_texts(&[]),
1065 }),
1066 Ok(TurnEvent::OutputTextDelta {
1067 handle: "msg-1".to_string(),
1068 delta: "hel".to_string(),
1069 }),
1070 Ok(TurnEvent::OutputTextDelta {
1071 handle: "msg-1".to_string(),
1072 delta: "lo".to_string(),
1073 }),
1074 Ok(TurnEvent::ItemCompleted {
1075 handle: "msg-1".to_string(),
1076 item: assistant_item_with_output_texts(&["hel", "lo"]),
1077 }),
1078 Ok(TurnEvent::Completed {
1079 response_id: "resp-1".to_string(),
1080 token_usage: None,
1081 }),
1082 ],
1083 gate: None,
1084 hold_open: None,
1085 }]);
1086 let manager = AgentBuilder::new(runtime).build();
1087 let session = manager.create_session();
1088
1089 let text = session
1090 .run_collect_text(SessionInput::from_user_text("hi"))
1091 .await
1092 .expect("text should be collected");
1093
1094 assert_eq!(text, "hello");
1095 }
1096
1097 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1098 async fn run_collect_text_falls_back_to_streamed_deltas() {
1099 let runtime = fake_runtime(vec![FakeTurnScript {
1100 events: vec![
1101 Ok(TurnEvent::OutputTextDelta {
1102 handle: "msg-1".to_string(),
1103 delta: "hel".to_string(),
1104 }),
1105 Ok(TurnEvent::OutputTextDelta {
1106 handle: "msg-1".to_string(),
1107 delta: "lo".to_string(),
1108 }),
1109 Ok(TurnEvent::Completed {
1110 response_id: "resp-1".to_string(),
1111 token_usage: None,
1112 }),
1113 ],
1114 gate: None,
1115 hold_open: None,
1116 }]);
1117 let manager = AgentBuilder::new(runtime).build();
1118 let session = manager.create_session();
1119
1120 let text = session
1121 .run_collect_text(SessionInput::from_user_text("hi"))
1122 .await
1123 .expect("text should be collected");
1124
1125 assert_eq!(text, "hello");
1126 }
1127
1128 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1129 async fn run_collect_text_waits_through_follow_up_turn() {
1130 let runtime = fake_runtime(vec![
1131 FakeTurnScript {
1132 events: vec![
1133 Ok(TurnEvent::ToolCall(ToolCallRequest {
1134 id: None,
1135 tool_name: "echo_tool".to_string(),
1136 call_id: "call-1".to_string(),
1137 payload: ToolCallPayload::JsonArguments {
1138 arguments: "{}".to_string(),
1139 },
1140 })),
1141 Ok(TurnEvent::Completed {
1142 response_id: "resp-1".to_string(),
1143 token_usage: None,
1144 }),
1145 ],
1146 gate: None,
1147 hold_open: None,
1148 },
1149 FakeTurnScript {
1150 events: vec![
1151 Ok(TurnEvent::ItemStarted {
1152 handle: "msg-2".to_string(),
1153 item: assistant_item("done"),
1154 }),
1155 Ok(TurnEvent::ItemCompleted {
1156 handle: "msg-2".to_string(),
1157 item: assistant_item("done"),
1158 }),
1159 Ok(TurnEvent::Completed {
1160 response_id: "resp-2".to_string(),
1161 token_usage: None,
1162 }),
1163 ],
1164 gate: None,
1165 hold_open: None,
1166 },
1167 ]);
1168 let manager = AgentBuilder::new(runtime)
1169 .register_tool(Arc::new(EchoTool))
1170 .build();
1171 let session = manager.create_session();
1172
1173 let text = session
1174 .run_collect_text(SessionInput::from_user_text("hi"))
1175 .await
1176 .expect("text should be collected after follow-up");
1177
1178 assert_eq!(text, "done");
1179 }
1180
1181 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1182 async fn run_collect_text_uses_final_turn_streamed_text_after_follow_up() {
1183 let runtime = fake_runtime(vec![
1184 FakeTurnScript {
1185 events: vec![
1186 Ok(TurnEvent::ItemStarted {
1187 handle: "msg-1".to_string(),
1188 item: assistant_item("draft"),
1189 }),
1190 Ok(TurnEvent::OutputTextDelta {
1191 handle: "msg-1".to_string(),
1192 delta: "draft".to_string(),
1193 }),
1194 Ok(TurnEvent::ItemCompleted {
1195 handle: "msg-1".to_string(),
1196 item: assistant_item("draft"),
1197 }),
1198 Ok(TurnEvent::ToolCall(ToolCallRequest {
1199 id: None,
1200 tool_name: "echo_tool".to_string(),
1201 call_id: "call-1".to_string(),
1202 payload: ToolCallPayload::JsonArguments {
1203 arguments: "{}".to_string(),
1204 },
1205 })),
1206 Ok(TurnEvent::Completed {
1207 response_id: "resp-1".to_string(),
1208 token_usage: None,
1209 }),
1210 ],
1211 gate: None,
1212 hold_open: None,
1213 },
1214 FakeTurnScript {
1215 events: vec![
1216 Ok(TurnEvent::ItemStarted {
1217 handle: "msg-2".to_string(),
1218 item: assistant_item_with_output_texts(&[]),
1219 }),
1220 Ok(TurnEvent::OutputTextDelta {
1221 handle: "msg-2".to_string(),
1222 delta: "fi".to_string(),
1223 }),
1224 Ok(TurnEvent::OutputTextDelta {
1225 handle: "msg-2".to_string(),
1226 delta: "nal".to_string(),
1227 }),
1228 Ok(TurnEvent::ItemCompleted {
1229 handle: "msg-2".to_string(),
1230 item: assistant_item_with_output_texts(&["fi", "nal"]),
1231 }),
1232 Ok(TurnEvent::Completed {
1233 response_id: "resp-2".to_string(),
1234 token_usage: None,
1235 }),
1236 ],
1237 gate: None,
1238 hold_open: None,
1239 },
1240 ]);
1241 let manager = AgentBuilder::new(runtime)
1242 .register_tool(Arc::new(EchoTool))
1243 .build();
1244 let session = manager.create_session();
1245
1246 let text = session
1247 .run_collect_text(SessionInput::from_user_text("hi"))
1248 .await
1249 .expect("text should be collected after follow-up");
1250
1251 assert_eq!(text, "final");
1252 }
1253
1254 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1255 async fn run_collect_text_waits_until_session_idle_before_returning() {
1256 let follow_up_gate = Arc::new(Notify::new());
1257 let (follow_up_waiting_tx, follow_up_waiting_rx) = oneshot::channel();
1258 let runtime = fake_runtime(vec![
1259 FakeTurnScript {
1260 events: vec![
1261 Ok(TurnEvent::ItemStarted {
1262 handle: "msg-1".to_string(),
1263 item: assistant_item("first"),
1264 }),
1265 Ok(TurnEvent::ItemCompleted {
1266 handle: "msg-1".to_string(),
1267 item: assistant_item("first"),
1268 }),
1269 Ok(TurnEvent::Completed {
1270 response_id: "resp-1".to_string(),
1271 token_usage: None,
1272 }),
1273 ],
1274 gate: None,
1275 hold_open: None,
1276 },
1277 FakeTurnScript {
1278 events: vec![
1279 Ok(TurnEvent::ItemStarted {
1280 handle: "msg-2".to_string(),
1281 item: assistant_item("queued"),
1282 }),
1283 Ok(TurnEvent::Completed {
1284 response_id: "resp-2".to_string(),
1285 token_usage: None,
1286 }),
1287 ],
1288 gate: Some(FakeGate {
1289 release: Arc::clone(&follow_up_gate),
1290 ready: Some(follow_up_waiting_tx),
1291 }),
1292 hold_open: None,
1293 },
1294 FakeTurnScript {
1295 events: vec![
1296 Ok(TurnEvent::ItemStarted {
1297 handle: "msg-3".to_string(),
1298 item: assistant_item("second"),
1299 }),
1300 Ok(TurnEvent::ItemCompleted {
1301 handle: "msg-3".to_string(),
1302 item: assistant_item("second"),
1303 }),
1304 Ok(TurnEvent::Completed {
1305 response_id: "resp-3".to_string(),
1306 token_usage: None,
1307 }),
1308 ],
1309 gate: None,
1310 hold_open: None,
1311 },
1312 ]);
1313 let manager = AgentBuilder::new(runtime).build();
1314 let session = manager.create_session();
1315 session
1316 .follow_up(SessionInput::from_user_text("queued"))
1317 .await;
1318
1319 let collect_session = session.clone();
1320 let collect_task = tokio::spawn(async move {
1321 collect_session
1322 .run_collect_text(SessionInput::from_user_text("first"))
1323 .await
1324 });
1325
1326 timeout(Duration::from_secs(2), follow_up_waiting_rx)
1327 .await
1328 .expect("follow-up turn should reach gate")
1329 .expect("follow-up gate waiter should be signaled");
1330 assert!(
1331 !collect_task.is_finished(),
1332 "collect-text should wait for queued follow-up cleanup"
1333 );
1334
1335 follow_up_gate.notify_one();
1336 let text = timeout(Duration::from_secs(2), collect_task)
1337 .await
1338 .expect("collect task should finish")
1339 .expect("collect task should not panic")
1340 .expect("text should be collected");
1341
1342 assert_eq!(text, "first");
1343 assert_eq!(session.status().await, SessionStatus::Idle);
1344
1345 let second_text = session
1346 .run_collect_text(SessionInput::from_user_text("second"))
1347 .await
1348 .expect("session should be immediately reusable");
1349 assert_eq!(second_text, "second");
1350 }
1351
1352 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1353 async fn run_collect_text_maps_turn_failed() {
1354 let runtime = fake_runtime(Vec::new());
1355 let manager = AgentBuilder::new(runtime)
1356 .register_skill_provider(Arc::new(TestSkillProvider {
1357 skills: Vec::new(),
1358 error: Some("boom".to_string()),
1359 contexts: Arc::new(StdMutex::new(Vec::new())),
1360 }))
1361 .build();
1362 let session = manager.create_session();
1363
1364 let err = session
1365 .run_collect_text(SessionInput::from_user_text("hi"))
1366 .await
1367 .expect_err("turn failure should be returned");
1368
1369 assert!(matches!(
1370 err,
1371 RunCollectTextError::TurnFailed(message) if message == "boom"
1372 ));
1373 assert_eq!(session.status().await, SessionStatus::Idle);
1374 }
1375
1376 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1377 async fn ask_once_creates_session_and_returns_text() {
1378 let runtime = fake_runtime(vec![FakeTurnScript {
1379 events: vec![
1380 Ok(TurnEvent::ItemStarted {
1381 handle: "msg-1".to_string(),
1382 item: assistant_item("hello"),
1383 }),
1384 Ok(TurnEvent::ItemCompleted {
1385 handle: "msg-1".to_string(),
1386 item: assistant_item("hello"),
1387 }),
1388 Ok(TurnEvent::Completed {
1389 response_id: "resp-1".to_string(),
1390 token_usage: None,
1391 }),
1392 ],
1393 gate: None,
1394 hold_open: None,
1395 }]);
1396 let manager = AgentBuilder::new(runtime).build();
1397
1398 let text = manager
1399 .ask_once("hi")
1400 .await
1401 .expect("text should be collected");
1402
1403 assert_eq!(text, "hello");
1404 }
1405
1406 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1407 async fn run_handles_tool_follow_up_turn() {
1408 let runtime = fake_runtime(vec![
1409 FakeTurnScript {
1410 events: vec![
1411 Ok(TurnEvent::ToolCall(ToolCallRequest {
1412 id: None,
1413 tool_name: "echo_tool".to_string(),
1414 call_id: "call-1".to_string(),
1415 payload: ToolCallPayload::JsonArguments {
1416 arguments: "{}".to_string(),
1417 },
1418 })),
1419 Ok(TurnEvent::Completed {
1420 response_id: "resp-1".to_string(),
1421 token_usage: None,
1422 }),
1423 ],
1424 gate: None,
1425 hold_open: None,
1426 },
1427 FakeTurnScript {
1428 events: vec![
1429 Ok(TurnEvent::ItemStarted {
1430 handle: "msg-2".to_string(),
1431 item: assistant_item("done"),
1432 }),
1433 Ok(TurnEvent::ItemCompleted {
1434 handle: "msg-2".to_string(),
1435 item: assistant_item("done"),
1436 }),
1437 Ok(TurnEvent::Completed {
1438 response_id: "resp-2".to_string(),
1439 token_usage: None,
1440 }),
1441 ],
1442 gate: None,
1443 hold_open: None,
1444 },
1445 ]);
1446
1447 let manager = AgentBuilder::new(runtime)
1448 .register_tool(Arc::new(EchoTool))
1449 .build();
1450 let session = manager.create_session();
1451
1452 session
1453 .run(SessionInput::from_user_text("hi"))
1454 .await
1455 .expect("run should succeed");
1456 let events = wait_for_idle(&session).await;
1457 assert!(
1458 events
1459 .iter()
1460 .filter(|event| matches!(event, AgentEvent::TurnCompleted { .. }))
1461 .count()
1462 >= 2
1463 );
1464 let completed_tool_call = events
1465 .iter()
1466 .find_map(|event| match event {
1467 AgentEvent::ToolCallCompleted { response, .. } => Some(response),
1468 _ => None,
1469 })
1470 .expect("tool completion should be emitted");
1471 assert_eq!(
1472 completed_tool_call,
1473 &lha_llm::ToolResultItem {
1474 call_id: "call-1".to_string(),
1475 tool_name: "echo_tool".to_string(),
1476 payload: ToolResultPayload::Structured {
1477 content: "handled echo_tool".to_string(),
1478 content_items: None,
1479 success: Some(true),
1480 },
1481 }
1482 );
1483
1484 let snapshot = session.snapshot().await;
1485 assert_eq!(snapshot.conversation.len(), 4);
1486 assert_eq!(
1487 snapshot.conversation[2],
1488 TranscriptItem::ToolResult {
1489 call_id: "call-1".to_string(),
1490 tool_name: "echo_tool".to_string(),
1491 payload: ToolResultPayload::Structured {
1492 content: "handled echo_tool".to_string(),
1493 content_items: None,
1494 success: Some(true),
1495 },
1496 }
1497 );
1498 }
1499
1500 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1501 async fn run_reports_unknown_tool_to_model_and_continues() {
1502 let (runtime, requests) = fake_runtime_with_requests(vec![
1503 FakeTurnScript {
1504 events: vec![
1505 Ok(TurnEvent::ToolCall(ToolCallRequest {
1506 id: None,
1507 tool_name: "missing_tool".to_string(),
1508 call_id: "call-1".to_string(),
1509 payload: ToolCallPayload::JsonArguments {
1510 arguments: "{}".to_string(),
1511 },
1512 })),
1513 Ok(TurnEvent::Completed {
1514 response_id: "resp-1".to_string(),
1515 token_usage: None,
1516 }),
1517 ],
1518 gate: None,
1519 hold_open: None,
1520 },
1521 FakeTurnScript {
1522 events: vec![
1523 Ok(TurnEvent::ItemStarted {
1524 handle: "msg-2".to_string(),
1525 item: assistant_item("done"),
1526 }),
1527 Ok(TurnEvent::ItemCompleted {
1528 handle: "msg-2".to_string(),
1529 item: assistant_item("done"),
1530 }),
1531 Ok(TurnEvent::Completed {
1532 response_id: "resp-2".to_string(),
1533 token_usage: None,
1534 }),
1535 ],
1536 gate: None,
1537 hold_open: None,
1538 },
1539 ]);
1540
1541 let manager = AgentBuilder::new(runtime).build();
1542 let session = manager.create_session();
1543
1544 session
1545 .run(SessionInput::from_user_text("hi"))
1546 .await
1547 .expect("run should succeed");
1548 let events = wait_for_idle(&session).await;
1549 assert!(
1550 events
1551 .iter()
1552 .filter(|event| matches!(event, AgentEvent::TurnCompleted { .. }))
1553 .count()
1554 >= 2
1555 );
1556 assert!(
1557 !events
1558 .iter()
1559 .any(|event| matches!(event, AgentEvent::TurnFailed { .. })),
1560 "unknown tool call should not fail the turn"
1561 );
1562 let expected_result = lha_llm::ToolResultItem {
1563 call_id: "call-1".to_string(),
1564 tool_name: "missing_tool".to_string(),
1565 payload: ToolResultPayload::Structured {
1566 content: "unsupported call: missing_tool".to_string(),
1567 content_items: None,
1568 success: Some(false),
1569 },
1570 };
1571 let completed_tool_call = events
1572 .iter()
1573 .find_map(|event| match event {
1574 AgentEvent::ToolCallCompleted { response, .. } => Some(response),
1575 _ => None,
1576 })
1577 .expect("tool completion should be emitted");
1578 assert_eq!(completed_tool_call, &expected_result);
1579
1580 let snapshot = session.snapshot().await;
1581 assert_eq!(snapshot.conversation.len(), 4);
1582 assert_eq!(
1583 snapshot.conversation[1],
1584 TranscriptItem::ToolCall {
1585 id: None,
1586 call_id: "call-1".to_string(),
1587 tool_name: "missing_tool".to_string(),
1588 payload: ToolCallPayload::JsonArguments {
1589 arguments: "{}".to_string(),
1590 },
1591 }
1592 );
1593 assert_eq!(
1594 snapshot.conversation[2],
1595 expected_result.to_transcript_item()
1596 );
1597
1598 let requests = lock_requests(&requests);
1599 assert_eq!(requests.len(), 2);
1600 assert!(
1601 requests[1]
1602 .conversation
1603 .contains(&expected_result.to_transcript_item()),
1604 "follow-up request should carry the failed tool result"
1605 );
1606 }
1607
1608 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1609 async fn no_skill_provider_preserves_turn_request_shape() {
1610 let (runtime, requests) = fake_runtime_with_requests(vec![completed_script()]);
1611 let manager = AgentBuilder::new(runtime)
1612 .with_base_instructions("base")
1613 .register_tool(Arc::new(EchoTool))
1614 .build();
1615 let session = manager.create_session();
1616
1617 session
1618 .run(SessionInput::from_user_text("hi"))
1619 .await
1620 .expect("run should succeed");
1621 wait_for_idle(&session).await;
1622
1623 let requests = lock_requests(&requests);
1624 assert_eq!(requests.len(), 1);
1625 let request = &requests[0];
1626 assert_eq!(
1627 request.conversation,
1628 SessionInput::from_user_text("hi").into_items()
1629 );
1630 assert_eq!(
1631 request.base_instructions,
1632 BaseInstructions {
1633 text: "base".to_string()
1634 }
1635 );
1636 assert_eq!(request.tools, vec![EchoTool.spec()]);
1637 assert!(!request.parallel_tool_calls);
1638 assert_eq!(request.personality, None);
1639 assert_eq!(request.output_schema, None);
1640 }
1641
1642 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1643 async fn skill_provider_appends_instructions_and_receives_context() {
1644 let (runtime, requests) = fake_runtime_with_requests(vec![completed_script()]);
1645 let contexts = Arc::new(StdMutex::new(Vec::new()));
1646 let provider = Arc::new(TestSkillProvider {
1647 skills: vec![test_skill("alpha", "Use alpha carefully.")],
1648 error: None,
1649 contexts: Arc::clone(&contexts),
1650 });
1651 let manager = AgentBuilder::new(runtime)
1652 .with_base_instructions("base")
1653 .register_tool(Arc::new(EchoTool))
1654 .register_skill_provider(provider)
1655 .build();
1656 let session = manager.create_session();
1657
1658 session
1659 .run(SessionInput::from_user_text("hi"))
1660 .await
1661 .expect("run should succeed");
1662 wait_for_idle(&session).await;
1663
1664 let requests = lock_requests(&requests);
1665 assert_eq!(requests.len(), 1);
1666 let instructions = &requests[0].base_instructions.text;
1667 assert!(instructions.contains("<available_skills>"));
1668 assert!(instructions.contains("id: alpha"));
1669 assert!(instructions.contains("name: alpha skill"));
1670 assert!(instructions.contains("Use alpha carefully."));
1671
1672 let contexts = contexts
1673 .lock()
1674 .unwrap_or_else(std::sync::PoisonError::into_inner)
1675 .clone();
1676 assert_eq!(
1677 contexts,
1678 vec![SkillContext {
1679 session_id: session.id(),
1680 conversation: SessionInput::from_user_text("hi").into_items(),
1681 runtime: RuntimeMetadata {
1682 endpoint_name: "fake".to_string(),
1683 model: "test-model".to_string(),
1684 },
1685 tools: vec![EchoTool.spec()],
1686 }]
1687 );
1688 }
1689
1690 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1691 async fn multiple_skill_providers_preserve_registration_order() {
1692 let (runtime, requests) = fake_runtime_with_requests(vec![completed_script()]);
1693 let first_contexts = Arc::new(StdMutex::new(Vec::new()));
1694 let second_contexts = Arc::new(StdMutex::new(Vec::new()));
1695 let manager = AgentBuilder::new(runtime)
1696 .with_base_instructions("base")
1697 .register_skill_provider(Arc::new(TestSkillProvider {
1698 skills: vec![test_skill("first", "First instructions.")],
1699 error: None,
1700 contexts: Arc::clone(&first_contexts),
1701 }))
1702 .register_skill_provider(Arc::new(TestSkillProvider {
1703 skills: vec![test_skill("second", "Second instructions.")],
1704 error: None,
1705 contexts: Arc::clone(&second_contexts),
1706 }))
1707 .build();
1708 let session = manager.create_session();
1709
1710 session
1711 .run(SessionInput::from_user_text("hi"))
1712 .await
1713 .expect("run should succeed");
1714 wait_for_idle(&session).await;
1715
1716 let requests = lock_requests(&requests);
1717 let instructions = &requests[0].base_instructions.text;
1718 let first = instructions
1719 .find("id: first")
1720 .expect("first skill should be present");
1721 let second = instructions
1722 .find("id: second")
1723 .expect("second skill should be present");
1724 assert!(first < second);
1725 }
1726
1727 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1728 async fn skill_provider_error_fails_turn() {
1729 let runtime = fake_runtime(Vec::new());
1730 let contexts = Arc::new(StdMutex::new(Vec::new()));
1731 let manager = AgentBuilder::new(runtime)
1732 .register_skill_provider(Arc::new(TestSkillProvider {
1733 skills: Vec::new(),
1734 error: Some("boom".to_string()),
1735 contexts,
1736 }))
1737 .build();
1738 let session = manager.create_session();
1739
1740 session
1741 .run(SessionInput::from_user_text("hi"))
1742 .await
1743 .expect("run should be accepted");
1744 let events = wait_for_idle(&session).await;
1745 let failed = events
1746 .into_iter()
1747 .find(|event| matches!(event, AgentEvent::TurnFailed { .. }))
1748 .expect("turn failure should be emitted");
1749
1750 match failed {
1751 AgentEvent::TurnFailed { error, .. } => assert_eq!(error, "boom"),
1752 other => panic!("expected TurnFailed, got {other:?}"),
1753 }
1754 }
1755
1756 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1757 async fn continue_turn_rejects_assistant_tail() {
1758 let runtime = fake_runtime(Vec::new());
1759 let manager = AgentBuilder::new(runtime).build();
1760 let session = manager.create_session();
1761
1762 {
1763 let mut state = session.inner.state.lock().await;
1764 state.conversation.push(TranscriptItem::Message {
1765 id: None,
1766 role: "assistant".to_string(),
1767 content: vec![ContentItem::OutputText {
1768 text: "done".to_string(),
1769 }],
1770 end_turn: None,
1771 });
1772 }
1773
1774 let err = session.continue_turn().await.expect_err("should fail");
1775 assert_eq!(err.to_string(), Error::InvalidContinuation.to_string());
1776 }
1777
1778 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1779 async fn abort_current_turn_cancels_active_turn() {
1780 let hold_open = Arc::new(Notify::new());
1781 let runtime = fake_runtime(vec![FakeTurnScript {
1782 events: Vec::new(),
1783 gate: None,
1784 hold_open: Some(Arc::clone(&hold_open)),
1785 }]);
1786
1787 let manager = AgentBuilder::new(runtime).build();
1788 let session = manager.create_session();
1789 session
1790 .run(SessionInput::from_user_text("hi"))
1791 .await
1792 .expect("run should succeed");
1793
1794 tokio::time::sleep(Duration::from_millis(50)).await;
1795 assert!(session.abort_current_turn().await);
1796 hold_open.notify_waiters();
1797
1798 let events = wait_for_idle(&session).await;
1799 assert!(
1800 events
1801 .iter()
1802 .any(|event| matches!(event, AgentEvent::TurnAborted { .. }))
1803 );
1804 assert_eq!(session.status().await, SessionStatus::Idle);
1805 }
1806
1807 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1808 async fn busy_run_rejects_without_primary_input_queued_event() {
1809 let hold_open = Arc::new(Notify::new());
1810 let runtime = fake_runtime(vec![FakeTurnScript {
1811 events: Vec::new(),
1812 gate: None,
1813 hold_open: Some(Arc::clone(&hold_open)),
1814 }]);
1815
1816 let manager = AgentBuilder::new(runtime).build();
1817 let session = manager.create_session();
1818 session
1819 .run(SessionInput::from_user_text("first"))
1820 .await
1821 .expect("first run should be accepted");
1822
1823 loop {
1824 let event = timeout(Duration::from_secs(2), session.next_event())
1825 .await
1826 .expect("event should arrive")
1827 .expect("session event should succeed");
1828 if matches!(event, AgentEvent::TurnStarted { .. }) {
1829 break;
1830 }
1831 }
1832
1833 let err = session
1834 .run(SessionInput::from_user_text("second"))
1835 .await
1836 .expect_err("second run should be rejected while busy");
1837 assert_eq!(err.to_string(), Error::SessionBusy.to_string());
1838 let event = timeout(Duration::from_millis(50), session.next_event()).await;
1839 assert!(
1840 event.is_err(),
1841 "busy run should not emit a primary input event, got {event:?}"
1842 );
1843
1844 hold_open.notify_waiters();
1845 loop {
1846 let event = timeout(Duration::from_secs(2), session.next_event())
1847 .await
1848 .expect("event should arrive")
1849 .expect("session event should succeed");
1850 if matches!(
1851 event,
1852 AgentEvent::SessionStatusChanged {
1853 status: SessionStatus::Idle,
1854 ..
1855 }
1856 ) {
1857 break;
1858 }
1859 }
1860 assert_eq!(session.status().await, SessionStatus::Idle);
1861 }
1862}