1. 1#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
2. 2
3. 3//! Capacity feedback (docs/design/68-context-engine.md §1 "Feedback"): a
4. 4//! provider that reports no prefill duration of its own (everyone but
5. 5//! Ollama) must still feed `prefill_tps`, from the wall-clock time between
6. 6//! sending the request and the first `StreamEvent` off the wire, measured
7. 7//! by `Agent::complete_with_reliability`'s own stream-consuming loop and
8. 8//! threaded through `StepLedger::last_first_token_ms`.
9. 9
10. 10use std::path::Path;
11. 11use std::sync::Arc;
12. 12use std::time::Duration;
13. 13
14. 14use tempfile::tempdir;
15. 15use tokio::sync::mpsc;
16. 16use tokio_util::sync::CancellationToken;
17. 17
18. 18use vak_agent::{Agent, AgentConfig, AutoApprove, SteeringQueues, TurnOutcome};
19. 19use vak_context::capacity::{
20. 20 CacheBehaviour, CapacityProfile, Horizon, ProbeProvenance, ProfileKey,
21. 21};
22. 22use vak_llm::stream;
23. 23use vak_llm::types::{AssistantMessage, ChatRequest, ContentBlock, StopReason, Usage};
24. 24use vak_llm::{EventStream, LlmError, Provider};
25. 25use vak_permission::{Mode, PermissionEngine};
26. 26use vak_session::types::{FrozenContract, SessionHeader};
27. 27use vak_session::{SessionLog, SessionPath};
28. 28
29. 29/// Answers after a deliberate delay so the first-token latency the harness
30. 30/// measures is unambiguously non-zero, and reports no `prefill_ms` of its
31. 31/// own (unlike Ollama) so the only way `prefill_tps` gets a sample is via
32. 32/// the caller-measured wall-clock latency.
33. 33struct SlowNoPrefillMs;
34. 34
35. 35#[async_trait::async_trait]
36. 36impl Provider for SlowNoPrefillMs {
37. 37 fn name(&self) -> &str {
38. 38 "slow-no-prefill"
39. 39 }
40. 40
41. 41 async fn stream(
42. 42 &self,
43. 43 _request: ChatRequest,
44. 44 _cancel: CancellationToken,
45. 45 ) -> Result<EventStream, LlmError> {
46. 46 let (mut sink, rx) = stream::channel(64);
47. 47 tokio::time::sleep(Duration::from_millis(30)).await;
48. 48 let message = AssistantMessage {
49. 49 content: vec![ContentBlock::text("done")],
50. 50 stop_reason: StopReason::EndTurn,
51. 51 usage: Usage {
52. 52 input_tokens: 500,
53. 53 output_tokens: 3,
54. 54 ..Default::default()
55. 55 },
56. 56 model: "test-model".into(),
57. 57 response_id: None,
58. 58 };
59. 59 sink.push(stream::StreamEvent::Start {
60. 60 partial: message.clone(),
61. 61 });
62. 62 sink.close_message(message).await;
63. 63 Ok(rx)
64. 64 }
65. 65}
66. 66
67. 67fn session_paths(dir: &Path) -> (std::path::PathBuf, std::path::PathBuf) {
68. 68 let cwd = dir.to_path_buf();
69. 69 let home = cwd.join(".vak-home");
70. 70 (
71. 71 home.clone(),
72. 72 SessionPath::new_session_file(&home, &cwd, "capacity-feedback"),
73. 73 )
74. 74}
75. 75
76. 76fn flat_profile(declared_window: u64) -> CapacityProfile {
77. 77 CapacityProfile::from_probe(
78. 78 declared_window,
79. 79 None,
80. 80 Horizon {
81. 81 tokens: declared_window,
82. 82 confidence: 0.9,
83. 83 last_confirmed: std::time::SystemTime::now(),
84. 84 },
85. 85 CacheBehaviour::Unknown,
86. 86 1_024,
87. 87 ProbeProvenance {
88. 88 probed_at: std::time::SystemTime::now(),
89. 89 rungs: Vec::new(),
90. 90 signals: Vec::new(),
91. 91 metadata_digest: "digest".into(),
92. 92 quantisation: None,
93. 93 },
94. 94 )
95. 95}
96. 96
97. 97#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
98. 98async fn non_ollama_provider_still_gets_a_prefill_tps_sample_from_wall_clock_latency() {
99. 99 let dir = tempdir().unwrap();
100. 100 let (home, path) = session_paths(dir.path());
101. 101 std::fs::create_dir_all(&home).unwrap();
102. 102 let provider: Arc<dyn Provider> = Arc::new(SlowNoPrefillMs);
103. 103 let header = SessionHeader {
104. 104 agent: None,
105. 105 session_id: "capacity-feedback".into(),
106. 106 created_at: chrono::Utc::now(),
107. 107 cwd: dir.path().to_path_buf(),
108. 108 parent_session_id: None,
109. 109 contract_id: None,
110. 110 work_item_id: None,
111. 111 conversation: None,
112. 112 contract: FrozenContract {
113. 113 app_version: "0".into(),
114. 114 provider: provider.name().to_string(),
115. 115 model: "test-model".into(),
116. 116 route_ladder: Vec::new(),
117. 117 route_objective: String::new(),
118. 118 route_annotations: Vec::new(),
119. 119 system_prompt: "sys".into(),
120. 120 permission_mode: "workspace-write".into(),
121. 121 capabilities: Vec::new(),
122. 122 prompt_layers: Vec::new(),
123. 123 },
124. 124 };
125. 125 let log = SessionLog::create(path, header).unwrap();
126. 126 let mut cfg = AgentConfig::new("sys");
127. 127 cfg.model = "test-model".into();
128. 128 cfg.mode = Mode::FullAccess;
129. 129 cfg.permission = Some(Arc::new(PermissionEngine::default()));
130. 130 cfg.approver = Some(Arc::new(AutoApprove));
131. 131 cfg.capacity = Some(flat_profile(128_000));
132. 132
133. 133 let mut agent = Agent::new(provider, log, cfg);
134. 134
135. 135 let (ev_tx, mut ev_rx) = mpsc::channel(256);
136. 136 tokio::spawn(async move { while ev_rx.recv().await.is_some() {} });
137. 137 let cancel = CancellationToken::new();
138. 138 let steering = SteeringQueues::new();
139. 139 let outcome = agent.run("hello", &steering, cancel, ev_tx).await;
140. 140 assert!(matches!(outcome, TurnOutcome::Completed { .. }));
141. 141
142. 142 let session = agent.into_session().await;
143. 143 let key = ProfileKey {
144. 144 provider: "slow-no-prefill".into(),
145. 145 model: "test-model".into(),
146. 146 quantisation: None,
147. 147 };
148. 148 let profile: CapacityProfile = session
149. 149 .latest_capacity_profile(&key)
150. 150 .expect("capacity feedback activity recorded with a matching profile");
151. 151 assert_eq!(
152. 152 profile.prefill_tps.samples, 1,
153. 153 "wall-clock first-token latency must feed prefill_tps even though \
154. 154 this provider never reports usage.prefill_ms itself"
155. 155 );
156. 156 assert!(
157. 157 profile.prefill_tps.value > 0.0,
158. 158 "prefill_tps value should be positive: {}",
159. 159 profile.prefill_tps.value
160. 160 );
161. 161}
162. 162