1use std::fmt;
4use std::future::Future;
5use std::path::PathBuf;
6use std::pin::Pin;
7use std::sync::Arc;
8
9use ag_protocol::{AgentResponse, ProtocolRequestProfile, TurnPrompt};
10use tokio::sync::mpsc;
11
12use crate::model::agent::ReasoningLevel;
13use crate::model::session::SpeedMode;
14
15pub type AgentFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
17
18pub trait LiveTranscript: fmt::Debug + Send + Sync {
20 fn replay_text(&self) -> Option<String>;
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
26pub enum AgentRequestKind {
27 SessionStart,
29 SessionResume,
31 UtilityPrompt,
37 AccountRead,
39}
40
41impl AgentRequestKind {
42 #[must_use]
44 pub fn protocol_profile(&self) -> ProtocolRequestProfile {
45 match self {
46 Self::SessionStart | Self::SessionResume => ProtocolRequestProfile::SessionTurn,
47 Self::UtilityPrompt | Self::AccountRead => ProtocolRequestProfile::UtilityPrompt,
48 }
49 }
50
51 #[must_use]
53 pub fn is_resume(&self) -> bool {
54 matches!(self, Self::SessionResume)
55 }
56}
57
58#[derive(Clone, Debug)]
63pub struct TurnContinuation {
64 kind: TurnContinuationKind,
65}
66
67#[derive(Clone, Debug, Default, Eq, PartialEq)]
69pub struct PersonalityPrompt {
70 current: Option<String>,
71 update: PersonalityPromptUpdate,
72}
73
74impl PersonalityPrompt {
75 #[must_use]
78 pub fn active(prompt: String, changed: bool) -> Self {
79 let update = if changed {
80 PersonalityPromptUpdate::Set(prompt.clone())
81 } else {
82 PersonalityPromptUpdate::Unchanged
83 };
84
85 Self {
86 current: Some(prompt),
87 update,
88 }
89 }
90
91 #[must_use]
96 pub fn cleared(changed: bool) -> Self {
97 Self {
98 current: None,
99 update: if changed {
100 PersonalityPromptUpdate::Clear
101 } else {
102 PersonalityPromptUpdate::Unchanged
103 },
104 }
105 }
106
107 #[must_use]
109 pub fn current(&self) -> Option<&str> {
110 self.current.as_deref()
111 }
112
113 pub(crate) fn update(&self) -> &PersonalityPromptUpdate {
114 &self.update
115 }
116}
117
118#[derive(Clone, Debug, Default, Eq, PartialEq)]
120pub(crate) enum PersonalityPromptUpdate {
121 Clear,
123 Set(String),
125 #[default]
127 Unchanged,
128}
129
130impl TurnContinuation {
131 #[must_use]
133 pub fn fresh() -> Self {
134 Self {
135 kind: TurnContinuationKind::Fresh,
136 }
137 }
138
139 #[must_use]
141 pub fn replaying(replay_transcript: String) -> Self {
142 Self {
143 kind: TurnContinuationKind::Replay { replay_transcript },
144 }
145 }
146
147 #[must_use]
150 pub fn provider(
151 live_transcript: Option<Arc<dyn LiveTranscript>>,
152 persisted_instruction_conversation_id: Option<String>,
153 provider_conversation_id: Option<String>,
154 replay_transcript: Option<String>,
155 ) -> Self {
156 Self {
157 kind: TurnContinuationKind::Provider {
158 live_transcript,
159 persisted_instruction_conversation_id,
160 provider_conversation_id,
161 replay_transcript,
162 },
163 }
164 }
165
166 #[must_use]
168 pub fn replay_transcript(&self) -> Option<&str> {
169 match &self.kind {
170 TurnContinuationKind::Fresh => None,
171 TurnContinuationKind::Provider {
172 replay_transcript, ..
173 } => replay_transcript.as_deref(),
174 TurnContinuationKind::Replay { replay_transcript } => Some(replay_transcript.as_str()),
175 }
176 }
177
178 #[must_use]
180 pub fn provider_conversation_id(&self) -> Option<&str> {
181 match &self.kind {
182 TurnContinuationKind::Provider {
183 provider_conversation_id,
184 ..
185 } => provider_conversation_id.as_deref(),
186 TurnContinuationKind::Fresh | TurnContinuationKind::Replay { .. } => None,
187 }
188 }
189
190 #[must_use]
193 pub fn persisted_instruction_conversation_id(&self) -> Option<&str> {
194 match &self.kind {
195 TurnContinuationKind::Provider {
196 persisted_instruction_conversation_id,
197 ..
198 } => persisted_instruction_conversation_id.as_deref(),
199 TurnContinuationKind::Fresh | TurnContinuationKind::Replay { .. } => None,
200 }
201 }
202
203 pub(crate) fn into_parts(self) -> TurnContinuationParts {
204 match self.kind {
205 TurnContinuationKind::Fresh => TurnContinuationParts::default(),
206 TurnContinuationKind::Replay { replay_transcript } => TurnContinuationParts {
207 replay_transcript: Some(replay_transcript),
208 ..TurnContinuationParts::default()
209 },
210 TurnContinuationKind::Provider {
211 live_transcript,
212 persisted_instruction_conversation_id,
213 provider_conversation_id,
214 replay_transcript,
215 } => TurnContinuationParts {
216 live_transcript,
217 persisted_instruction_conversation_id,
218 provider_conversation_id,
219 replay_transcript,
220 },
221 }
222 }
223}
224
225#[derive(Clone, Debug)]
226enum TurnContinuationKind {
227 Fresh,
228 Provider {
229 live_transcript: Option<Arc<dyn LiveTranscript>>,
230 persisted_instruction_conversation_id: Option<String>,
231 provider_conversation_id: Option<String>,
232 replay_transcript: Option<String>,
233 },
234 Replay {
235 replay_transcript: String,
236 },
237}
238
239#[derive(Default)]
240pub(crate) struct TurnContinuationParts {
241 pub(crate) live_transcript: Option<Arc<dyn LiveTranscript>>,
242 pub(crate) persisted_instruction_conversation_id: Option<String>,
243 pub(crate) provider_conversation_id: Option<String>,
244 pub(crate) replay_transcript: Option<String>,
245}
246
247#[derive(Debug, Clone)]
249pub struct TurnRequest {
250 pub continuation: TurnContinuation,
252 pub folder: PathBuf,
254 pub main_checkout_root: Option<PathBuf>,
257 pub model: String,
259 pub personality: PersonalityPrompt,
261 pub prompt: TurnPrompt,
263 pub reasoning_level: ReasoningLevel,
267 pub request_kind: AgentRequestKind,
270 pub speed_mode: SpeedMode,
272}
273
274#[derive(Clone, Debug, PartialEq)]
280pub enum TurnEvent {
281 ThoughtDelta(String),
284 Completed {
286 context_reset: bool,
288 input_tokens: u64,
290 output_tokens: u64,
292 },
293 Failed(String),
295 PidUpdate(Option<u32>),
301}
302
303#[derive(Debug)]
305pub struct TurnResult {
306 pub assistant_message: AgentResponse,
308 pub context_reset: bool,
310 pub input_tokens: u64,
312 pub output_tokens: u64,
314 pub provider_conversation_id: Option<String>,
319}
320
321pub struct SessionRef {
323 pub session_id: String,
325}
326
327pub struct StartSessionRequest {
329 pub folder: PathBuf,
331 pub session_id: String,
333}
334
335#[derive(Debug, thiserror::Error)]
340pub enum AgentError {
341 #[error(transparent)]
344 AppServer(#[from] crate::app_server::AppServerError),
345
346 #[error("{0}")]
348 Backend(String),
349
350 #[error("{0}")]
352 InterruptedByUser(String),
353
354 #[error("{0}")]
356 Io(String),
357}
358
359#[cfg_attr(any(test, feature = "test-utils"), mockall::automock)]
365pub trait AgentChannel: Send + Sync {
366 fn start_session(
371 &self,
372 req: StartSessionRequest,
373 ) -> AgentFuture<Result<SessionRef, AgentError>>;
374
375 fn run_turn(
385 &self,
386 session_id: String,
387 req: TurnRequest,
388 events: mpsc::UnboundedSender<TurnEvent>,
389 ) -> AgentFuture<Result<TurnResult, AgentError>>;
390
391 fn shutdown_session(&self, session_id: String) -> AgentFuture<Result<(), AgentError>>;
396}
397
398#[cfg(test)]
399mod tests {
400 use super::*;
401
402 #[test]
403 fn test_personality_prompt_tracks_active_change_and_clear_state() {
404 let changed = PersonalityPrompt::active("Review carefully.".to_string(), true);
406 let unchanged = PersonalityPrompt::active("Review carefully.".to_string(), false);
407 let cleared = PersonalityPrompt::cleared(true);
408 let empty = PersonalityPrompt::cleared(false);
409
410 assert_eq!(changed.current(), Some("Review carefully."));
412 assert_eq!(
413 changed.update(),
414 &PersonalityPromptUpdate::Set("Review carefully.".to_string())
415 );
416 assert_eq!(unchanged.update(), &PersonalityPromptUpdate::Unchanged);
417 assert_eq!(cleared.current(), None);
418 assert_eq!(cleared.update(), &PersonalityPromptUpdate::Clear);
419 assert_eq!(empty.update(), &PersonalityPromptUpdate::Unchanged);
420 }
421
422 #[test]
423 fn test_turn_continuation_fresh_has_no_context() {
424 let continuation = TurnContinuation::fresh();
426
427 assert_eq!(continuation.replay_transcript(), None);
429 assert_eq!(continuation.provider_conversation_id(), None);
430 assert_eq!(continuation.persisted_instruction_conversation_id(), None);
431 }
432
433 #[test]
434 fn test_turn_continuation_replaying_exposes_transcript_only() {
435 let continuation = TurnContinuation::replaying("prior turn".to_string());
437
438 let parts = continuation.clone().into_parts();
440
441 assert_eq!(continuation.replay_transcript(), Some("prior turn"));
443 assert_eq!(continuation.provider_conversation_id(), None);
444 assert!(parts.live_transcript.is_none());
445 assert_eq!(parts.persisted_instruction_conversation_id, None);
446 assert_eq!(parts.provider_conversation_id, None);
447 assert_eq!(parts.replay_transcript.as_deref(), Some("prior turn"));
448 }
449
450 #[test]
451 fn test_turn_continuation_provider_exposes_persisted_context() {
452 let continuation = TurnContinuation::provider(
454 None,
455 Some("instruction-1".to_string()),
456 Some("thread-1".to_string()),
457 Some("prior turn".to_string()),
458 );
459
460 assert_eq!(continuation.replay_transcript(), Some("prior turn"));
462 assert_eq!(continuation.provider_conversation_id(), Some("thread-1"));
463 assert_eq!(
464 continuation.persisted_instruction_conversation_id(),
465 Some("instruction-1")
466 );
467 }
468
469 #[test]
470 fn test_agent_request_kind_session_variants_use_session_protocol_profile() {
473 let start = AgentRequestKind::SessionStart;
475 let resume = AgentRequestKind::SessionResume;
476
477 let start_profile = start.protocol_profile();
479 let resume_profile = resume.protocol_profile();
480
481 assert_eq!(start_profile, ProtocolRequestProfile::SessionTurn);
483 assert_eq!(resume_profile, ProtocolRequestProfile::SessionTurn);
484 }
485
486 #[test]
487 fn test_agent_request_kind_utility_prompt_uses_utility_protocol_profile() {
489 let request_kind = AgentRequestKind::UtilityPrompt;
491
492 let protocol_profile = request_kind.protocol_profile();
494
495 assert_eq!(protocol_profile, ProtocolRequestProfile::UtilityPrompt);
497 }
498
499 #[test]
500 fn test_agent_request_kind_account_read_uses_utility_protocol_profile() {
502 let request_kind = AgentRequestKind::AccountRead;
504
505 let protocol_profile = request_kind.protocol_profile();
507
508 assert_eq!(protocol_profile, ProtocolRequestProfile::UtilityPrompt);
510 }
511}