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::permission::PermissionMode;
14use crate::model::session::SpeedMode;
15
16pub type AgentFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
18
19pub trait LiveTranscript: fmt::Debug + Send + Sync {
21 fn replay_text(&self) -> Option<String>;
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
27pub enum AgentRequestKind {
28 SessionStart,
30 SessionResume,
32 UtilityPrompt,
38 AccountRead,
40}
41
42impl AgentRequestKind {
43 #[must_use]
45 pub fn protocol_profile(&self) -> ProtocolRequestProfile {
46 match self {
47 Self::SessionStart | Self::SessionResume => ProtocolRequestProfile::SessionTurn,
48 Self::UtilityPrompt | Self::AccountRead => ProtocolRequestProfile::UtilityPrompt,
49 }
50 }
51
52 #[must_use]
54 pub fn is_resume(&self) -> bool {
55 matches!(self, Self::SessionResume)
56 }
57}
58
59#[derive(Clone, Debug)]
64pub struct TurnContinuation {
65 kind: TurnContinuationKind,
66}
67
68#[derive(Clone, Debug, Default, Eq, PartialEq)]
70pub struct PersonalityPrompt {
71 current: Option<String>,
72 update: PersonalityPromptUpdate,
73}
74
75impl PersonalityPrompt {
76 #[must_use]
79 pub fn active(prompt: String, changed: bool) -> Self {
80 let update = if changed {
81 PersonalityPromptUpdate::Set(prompt.clone())
82 } else {
83 PersonalityPromptUpdate::Unchanged
84 };
85
86 Self {
87 current: Some(prompt),
88 update,
89 }
90 }
91
92 #[must_use]
97 pub fn cleared(changed: bool) -> Self {
98 Self {
99 current: None,
100 update: if changed {
101 PersonalityPromptUpdate::Clear
102 } else {
103 PersonalityPromptUpdate::Unchanged
104 },
105 }
106 }
107
108 #[must_use]
110 pub fn current(&self) -> Option<&str> {
111 self.current.as_deref()
112 }
113
114 pub(crate) fn update(&self) -> &PersonalityPromptUpdate {
115 &self.update
116 }
117}
118
119#[derive(Clone, Debug, Default, Eq, PartialEq)]
121pub(crate) enum PersonalityPromptUpdate {
122 Clear,
124 Set(String),
126 #[default]
128 Unchanged,
129}
130
131impl TurnContinuation {
132 #[must_use]
134 pub fn fresh() -> Self {
135 Self {
136 kind: TurnContinuationKind::Fresh,
137 }
138 }
139
140 #[must_use]
142 pub fn replaying(replay_transcript: String) -> Self {
143 Self {
144 kind: TurnContinuationKind::Replay { replay_transcript },
145 }
146 }
147
148 #[must_use]
151 pub fn provider(
152 live_transcript: Option<Arc<dyn LiveTranscript>>,
153 persisted_instruction_conversation_id: Option<String>,
154 provider_conversation_id: Option<String>,
155 replay_transcript: Option<String>,
156 ) -> Self {
157 Self {
158 kind: TurnContinuationKind::Provider {
159 live_transcript,
160 persisted_instruction_conversation_id,
161 provider_conversation_id,
162 replay_transcript,
163 },
164 }
165 }
166
167 #[must_use]
169 pub fn replay_transcript(&self) -> Option<&str> {
170 match &self.kind {
171 TurnContinuationKind::Fresh => None,
172 TurnContinuationKind::Provider {
173 replay_transcript, ..
174 } => replay_transcript.as_deref(),
175 TurnContinuationKind::Replay { replay_transcript } => Some(replay_transcript.as_str()),
176 }
177 }
178
179 #[must_use]
181 pub fn provider_conversation_id(&self) -> Option<&str> {
182 match &self.kind {
183 TurnContinuationKind::Provider {
184 provider_conversation_id,
185 ..
186 } => provider_conversation_id.as_deref(),
187 TurnContinuationKind::Fresh | TurnContinuationKind::Replay { .. } => None,
188 }
189 }
190
191 #[must_use]
194 pub fn persisted_instruction_conversation_id(&self) -> Option<&str> {
195 match &self.kind {
196 TurnContinuationKind::Provider {
197 persisted_instruction_conversation_id,
198 ..
199 } => persisted_instruction_conversation_id.as_deref(),
200 TurnContinuationKind::Fresh | TurnContinuationKind::Replay { .. } => None,
201 }
202 }
203
204 pub(crate) fn into_parts(self) -> TurnContinuationParts {
205 match self.kind {
206 TurnContinuationKind::Fresh => TurnContinuationParts::default(),
207 TurnContinuationKind::Replay { replay_transcript } => TurnContinuationParts {
208 replay_transcript: Some(replay_transcript),
209 ..TurnContinuationParts::default()
210 },
211 TurnContinuationKind::Provider {
212 live_transcript,
213 persisted_instruction_conversation_id,
214 provider_conversation_id,
215 replay_transcript,
216 } => TurnContinuationParts {
217 live_transcript,
218 persisted_instruction_conversation_id,
219 provider_conversation_id,
220 replay_transcript,
221 },
222 }
223 }
224}
225
226#[derive(Clone, Debug)]
227enum TurnContinuationKind {
228 Fresh,
229 Provider {
230 live_transcript: Option<Arc<dyn LiveTranscript>>,
231 persisted_instruction_conversation_id: Option<String>,
232 provider_conversation_id: Option<String>,
233 replay_transcript: Option<String>,
234 },
235 Replay {
236 replay_transcript: String,
237 },
238}
239
240#[derive(Default)]
241pub(crate) struct TurnContinuationParts {
242 pub(crate) live_transcript: Option<Arc<dyn LiveTranscript>>,
243 pub(crate) persisted_instruction_conversation_id: Option<String>,
244 pub(crate) provider_conversation_id: Option<String>,
245 pub(crate) replay_transcript: Option<String>,
246}
247
248#[derive(Debug, Clone)]
250pub struct TurnRequest {
251 pub continuation: TurnContinuation,
253 pub folder: PathBuf,
255 pub main_checkout_root: Option<PathBuf>,
258 pub model: String,
260 pub permission_mode: PermissionMode,
262 pub personality: PersonalityPrompt,
264 pub prompt: TurnPrompt,
266 pub reasoning_level: ReasoningLevel,
270 pub request_kind: AgentRequestKind,
273 pub speed_mode: SpeedMode,
275}
276
277#[derive(Clone, Debug, PartialEq)]
283pub enum TurnEvent {
284 ThoughtDelta(String),
287 Completed {
289 context_reset: bool,
291 input_tokens: u64,
293 output_tokens: u64,
295 },
296 Failed(String),
298 PidUpdate(Option<u32>),
304}
305
306#[derive(Debug)]
308pub struct TurnResult {
309 pub assistant_message: AgentResponse,
311 pub context_reset: bool,
313 pub input_tokens: u64,
315 pub output_tokens: u64,
317 pub provider_conversation_id: Option<String>,
322}
323
324pub struct SessionRef {
326 pub session_id: String,
328}
329
330pub struct StartSessionRequest {
332 pub folder: PathBuf,
334 pub session_id: String,
336}
337
338#[derive(Debug, thiserror::Error)]
343pub enum AgentError {
344 #[error(transparent)]
347 AppServer(#[from] crate::app_server::AppServerError),
348
349 #[error("{0}")]
351 Backend(String),
352
353 #[error("{0}")]
355 InterruptedByUser(String),
356
357 #[error("{0}")]
359 Io(String),
360}
361
362#[cfg_attr(any(test, feature = "test-utils"), mockall::automock)]
368pub trait AgentChannel: Send + Sync {
369 fn start_session(
374 &self,
375 req: StartSessionRequest,
376 ) -> AgentFuture<Result<SessionRef, AgentError>>;
377
378 fn run_turn(
388 &self,
389 session_id: String,
390 req: TurnRequest,
391 events: mpsc::UnboundedSender<TurnEvent>,
392 ) -> AgentFuture<Result<TurnResult, AgentError>>;
393
394 fn shutdown_session(&self, session_id: String) -> AgentFuture<Result<(), AgentError>>;
399}
400
401#[cfg(test)]
402mod tests {
403 use super::*;
404
405 #[test]
406 fn test_personality_prompt_tracks_active_change_and_clear_state() {
407 let changed = PersonalityPrompt::active("Review carefully.".to_string(), true);
409 let unchanged = PersonalityPrompt::active("Review carefully.".to_string(), false);
410 let cleared = PersonalityPrompt::cleared(true);
411 let empty = PersonalityPrompt::cleared(false);
412
413 assert_eq!(changed.current(), Some("Review carefully."));
415 assert_eq!(
416 changed.update(),
417 &PersonalityPromptUpdate::Set("Review carefully.".to_string())
418 );
419 assert_eq!(unchanged.update(), &PersonalityPromptUpdate::Unchanged);
420 assert_eq!(cleared.current(), None);
421 assert_eq!(cleared.update(), &PersonalityPromptUpdate::Clear);
422 assert_eq!(empty.update(), &PersonalityPromptUpdate::Unchanged);
423 }
424
425 #[test]
426 fn test_turn_continuation_fresh_has_no_context() {
427 let continuation = TurnContinuation::fresh();
429
430 assert_eq!(continuation.replay_transcript(), None);
432 assert_eq!(continuation.provider_conversation_id(), None);
433 assert_eq!(continuation.persisted_instruction_conversation_id(), None);
434 }
435
436 #[test]
437 fn test_turn_continuation_replaying_exposes_transcript_only() {
438 let continuation = TurnContinuation::replaying("prior turn".to_string());
440
441 let parts = continuation.clone().into_parts();
443
444 assert_eq!(continuation.replay_transcript(), Some("prior turn"));
446 assert_eq!(continuation.provider_conversation_id(), None);
447 assert!(parts.live_transcript.is_none());
448 assert_eq!(parts.persisted_instruction_conversation_id, None);
449 assert_eq!(parts.provider_conversation_id, None);
450 assert_eq!(parts.replay_transcript.as_deref(), Some("prior turn"));
451 }
452
453 #[test]
454 fn test_turn_continuation_provider_exposes_persisted_context() {
455 let continuation = TurnContinuation::provider(
457 None,
458 Some("instruction-1".to_string()),
459 Some("thread-1".to_string()),
460 Some("prior turn".to_string()),
461 );
462
463 assert_eq!(continuation.replay_transcript(), Some("prior turn"));
465 assert_eq!(continuation.provider_conversation_id(), Some("thread-1"));
466 assert_eq!(
467 continuation.persisted_instruction_conversation_id(),
468 Some("instruction-1")
469 );
470 }
471
472 #[test]
473 fn test_agent_request_kind_session_variants_use_session_protocol_profile() {
476 let start = AgentRequestKind::SessionStart;
478 let resume = AgentRequestKind::SessionResume;
479
480 let start_profile = start.protocol_profile();
482 let resume_profile = resume.protocol_profile();
483
484 assert_eq!(start_profile, ProtocolRequestProfile::SessionTurn);
486 assert_eq!(resume_profile, ProtocolRequestProfile::SessionTurn);
487 }
488
489 #[test]
490 fn test_agent_request_kind_utility_prompt_uses_utility_protocol_profile() {
492 let request_kind = AgentRequestKind::UtilityPrompt;
494
495 let protocol_profile = request_kind.protocol_profile();
497
498 assert_eq!(protocol_profile, ProtocolRequestProfile::UtilityPrompt);
500 }
501
502 #[test]
503 fn test_agent_request_kind_account_read_uses_utility_protocol_profile() {
505 let request_kind = AgentRequestKind::AccountRead;
507
508 let protocol_profile = request_kind.protocol_profile();
510
511 assert_eq!(protocol_profile, ProtocolRequestProfile::UtilityPrompt);
513 }
514}