1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3//! A slow listener never paces the model: stream pieces it has no room for
4. 4//! are dropped (each carries the message so far), so a long reasoning reply
5. 5//! finishes at the provider's speed, not the listener's.
6. 6//!
7. 7//! Found live: the eval runner handed the loop an event channel nobody read,
8. 8//! and a long reasoning reply waited on it until the step watchdog although
9. 9//! the provider had finished.
10. 10
11. 11use std::sync::Arc;
12. 12use std::time::Duration;
13. 13
14. 14use async_trait::async_trait;
15. 15use tokio_util::sync::CancellationToken;
16. 16
17. 17use tempfile::tempdir;
18. 18
19. 19use vak_agent::{Agent, AgentConfig, TurnOutcome};
20. 20use vak_llm::stream;
21. 21use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
22. 22use vak_llm::{EventStream, LlmError, Provider};
23. 23use vak_session::types::{FrozenContract, SessionHeader};
24. 24use vak_session::{SessionLog, SessionPath};
25. 25
26. 26/// Streams two thousand reasoning pieces, then the answer.
27. 27struct LongReasoning;
28. 28
29. 29#[async_trait]
30. 30impl Provider for LongReasoning {
31. 31 fn name(&self) -> &str {
32. 32 "long-reasoning"
33. 33 }
34. 34
35. 35 async fn stream(
36. 36 &self,
37. 37 _request: ChatRequest,
38. 38 _cancel: CancellationToken,
39. 39 ) -> Result<EventStream, LlmError> {
40. 40 let (mut sink, rx) = stream::channel(4096);
41. 41 tokio::spawn(async move {
42. 42 let mut partial = AssistantMessage::empty("test-model");
43. 43 for n in 0..2000 {
44. 44 let delta = format!("step {n} ");
45. 45 partial.content = vec![ContentBlock::Thinking {
46. 46 text: format!(
47. 47 "{}{delta}",
48. 48 match partial.content.first() {
49. 49 Some(ContentBlock::Thinking { text, .. }) => text.clone(),
50. 50 _ => String::new(),
51. 51 }
52. 52 ),
53. 53 signature: None,
54. 54 }];
55. 55 sink.push(stream::StreamEvent::ThinkingDelta {
56. 56 delta,
57. 57 partial: partial.clone(),
58. 58 });
59. 59 }
60. 60 sink.close_message(AssistantMessage {
61. 61 content: vec![ContentBlock::text("done")],
62. 62 stop_reason: StopReason::EndTurn,
63. 63 usage: Usage {
64. 64 input_tokens: 1,
65. 65 output_tokens: 1,
66. 66 ..Default::default()
67. 67 },
68. 68 model: "test-model".into(),
69. 69 response_id: None,
70. 70 })
71. 71 .await;
72. 72 });
73. 73 Ok(rx)
74. 74 }
75. 75}
76. 76
77. 77#[tokio::test]
78. 78async fn a_slow_listener_does_not_pace_a_long_reply() {
79. 79 let dir = tempdir().unwrap();
80. 80 let home = dir.path().join("home");
81. 81 std::fs::create_dir_all(&home).unwrap();
82. 82 let header = SessionHeader {
83. 83 agent: None,
84. 84 session_id: "stalled".into(),
85. 85 created_at: chrono::Utc::now(),
86. 86 cwd: dir.path().to_path_buf(),
87. 87 parent_session_id: None,
88. 88 contract_id: None,
89. 89 work_item_id: None,
90. 90 conversation: None,
91. 91 contract: FrozenContract {
92. 92 app_version: "0".into(),
93. 93 provider: "long-reasoning".into(),
94. 94 model: "test-model".into(),
95. 95 route_ladder: Vec::new(),
96. 96 route_objective: String::new(),
97. 97 route_annotations: Vec::new(),
98. 98 system_prompt: "sys".into(),
99. 99 permission_mode: "read-only".into(),
100. 100 capabilities: Vec::new(),
101. 101 prompt_layers: Vec::new(),
102. 102 },
103. 103 };
104. 104 let log = SessionLog::create(
105. 105 SessionPath::new_session_file(&home, dir.path(), "stalled"),
106. 106 header,
107. 107 )
108. 108 .unwrap();
109. 109 let mut cfg = AgentConfig::new("sys");
110. 110 cfg.model = "test-model".into();
111. 111 let mut agent = Agent::new(Arc::new(LongReasoning), log, cfg);
112. 112
113. 113 // 20ms per event: pacing the two thousand pieces would take 40s.
114. 114 let (events, mut rx) = tokio::sync::mpsc::channel(64);
115. 115 tokio::spawn(async move {
116. 116 while rx.recv().await.is_some() {
117. 117 tokio::time::sleep(Duration::from_millis(20)).await;
118. 118 }
119. 119 });
120. 120 let outcome = tokio::time::timeout(
121. 121 Duration::from_secs(10),
122. 122 agent.run(
123. 123 "think it through",
124. 124 &Default::default(),
125. 125 CancellationToken::new(),
126. 126 events,
127. 127 ),
128. 128 )
129. 129 .await
130. 130 .expect("the reply is not paced by the listener");
131. 131 match outcome {
132. 132 TurnOutcome::Completed { response } => assert_eq!(response.text_content(), "done"),
133. 133 other => panic!("expected completed, got {other:?}"),
134. 134 }
135. 135}
136. 136