use std::fmt;
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use ag_protocol::{AgentResponse, ProtocolRequestProfile};
use tokio::sync::mpsc;
use crate::model::agent::ReasoningLevel;
use crate::model::turn_prompt::TurnPrompt;
pub type AgentFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
pub trait LiveTranscript: fmt::Debug + Send + Sync {
fn replay_text(&self) -> Option<String>;
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AgentRequestKind {
SessionStart,
SessionResume,
UtilityPrompt,
AccountRead,
}
impl AgentRequestKind {
#[must_use]
pub fn protocol_profile(&self) -> ProtocolRequestProfile {
match self {
Self::SessionStart | Self::SessionResume => ProtocolRequestProfile::SessionTurn,
Self::UtilityPrompt | Self::AccountRead => ProtocolRequestProfile::UtilityPrompt,
}
}
#[must_use]
pub fn is_resume(&self) -> bool {
matches!(self, Self::SessionResume)
}
}
#[derive(Debug, Clone)]
pub struct TurnRequest {
pub folder: PathBuf,
pub live_transcript: Option<Arc<dyn LiveTranscript>>,
pub main_checkout_root: Option<PathBuf>,
pub model: String,
pub request_kind: AgentRequestKind,
pub replay_transcript: Option<String>,
pub prompt: TurnPrompt,
pub provider_conversation_id: Option<String>,
pub persisted_instruction_conversation_id: Option<String>,
pub reasoning_level: ReasoningLevel,
}
#[derive(Clone, Debug, PartialEq)]
pub enum TurnEvent {
ThoughtDelta(String),
Completed {
context_reset: bool,
input_tokens: u64,
output_tokens: u64,
},
Failed(String),
PidUpdate(Option<u32>),
}
#[derive(Debug)]
pub struct TurnResult {
pub assistant_message: AgentResponse,
pub context_reset: bool,
pub input_tokens: u64,
pub output_tokens: u64,
pub provider_conversation_id: Option<String>,
}
pub struct SessionRef {
pub session_id: String,
}
pub struct StartSessionRequest {
pub folder: PathBuf,
pub session_id: String,
}
#[derive(Debug, thiserror::Error)]
pub enum AgentError {
#[error(transparent)]
AppServer(#[from] crate::app_server::AppServerError),
#[error("{0}")]
Backend(String),
#[error("{0}")]
InterruptedByUser(String),
#[error("{0}")]
Io(String),
}
#[cfg_attr(any(test, feature = "test-utils"), mockall::automock)]
pub trait AgentChannel: Send + Sync {
fn start_session(
&self,
req: StartSessionRequest,
) -> AgentFuture<Result<SessionRef, AgentError>>;
fn run_turn(
&self,
session_id: String,
req: TurnRequest,
events: mpsc::UnboundedSender<TurnEvent>,
) -> AgentFuture<Result<TurnResult, AgentError>>;
fn shutdown_session(&self, session_id: String) -> AgentFuture<Result<(), AgentError>>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_agent_request_kind_session_variants_use_session_protocol_profile() {
let start = AgentRequestKind::SessionStart;
let resume = AgentRequestKind::SessionResume;
let start_profile = start.protocol_profile();
let resume_profile = resume.protocol_profile();
assert_eq!(start_profile, ProtocolRequestProfile::SessionTurn);
assert_eq!(resume_profile, ProtocolRequestProfile::SessionTurn);
}
#[test]
fn test_agent_request_kind_utility_prompt_uses_utility_protocol_profile() {
let request_kind = AgentRequestKind::UtilityPrompt;
let protocol_profile = request_kind.protocol_profile();
assert_eq!(protocol_profile, ProtocolRequestProfile::UtilityPrompt);
}
#[test]
fn test_agent_request_kind_account_read_uses_utility_protocol_profile() {
let request_kind = AgentRequestKind::AccountRead;
let protocol_profile = request_kind.protocol_profile();
assert_eq!(protocol_profile, ProtocolRequestProfile::UtilityPrompt);
}
}