agcodex_core/
conversation_manager.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use agcodex_login::CodexAuth;
5use tokio::sync::RwLock;
6use uuid::Uuid;
7
8use crate::codex::Codex;
9use crate::codex::CodexSpawnOk;
10use crate::codex::INITIAL_SUBMIT_ID;
11use crate::codex_conversation::CodexConversation;
12use crate::config::Config;
13use crate::error::CodexErr;
14use crate::error::Result as CodexResult;
15use crate::protocol::Event;
16use crate::protocol::EventMsg;
17use crate::protocol::SessionConfiguredEvent;
18
19/// Represents a newly created Codex conversation, including the first event
20/// (which is [`EventMsg::SessionConfigured`]).
21pub struct NewConversation {
22    pub conversation_id: Uuid,
23    pub conversation: Arc<CodexConversation>,
24    pub session_configured: SessionConfiguredEvent,
25}
26
27/// [`ConversationManager`] is responsible for creating conversations and
28/// maintaining them in memory.
29pub struct ConversationManager {
30    conversations: Arc<RwLock<HashMap<Uuid, Arc<CodexConversation>>>>,
31}
32
33impl Default for ConversationManager {
34    fn default() -> Self {
35        Self {
36            conversations: Arc::new(RwLock::new(HashMap::new())),
37        }
38    }
39}
40
41impl ConversationManager {
42    pub async fn new_conversation(&self, config: Config) -> CodexResult<NewConversation> {
43        let auth = CodexAuth::from_codex_home(&config.codex_home, config.preferred_auth_method)?;
44        self.new_conversation_with_auth(config, auth).await
45    }
46
47    /// Used for integration tests: should not be used by ordinary business
48    /// logic.
49    pub async fn new_conversation_with_auth(
50        &self,
51        config: Config,
52        auth: Option<CodexAuth>,
53    ) -> CodexResult<NewConversation> {
54        let CodexSpawnOk {
55            codex,
56            session_id: conversation_id,
57        } = Codex::spawn(config, auth).await?;
58
59        // The first event must be `SessionInitialized`. Validate and forward it
60        // to the caller so that they can display it in the conversation
61        // history.
62        let event = codex.next_event().await?;
63        let session_configured = match event {
64            Event {
65                id,
66                msg: EventMsg::SessionConfigured(session_configured),
67            } if id == INITIAL_SUBMIT_ID => session_configured,
68            _ => {
69                return Err(CodexErr::SessionConfiguredNotFirstEvent);
70            }
71        };
72
73        let conversation = Arc::new(CodexConversation::new(codex));
74        self.conversations
75            .write()
76            .await
77            .insert(conversation_id, conversation.clone());
78
79        Ok(NewConversation {
80            conversation_id,
81            conversation,
82            session_configured,
83        })
84    }
85
86    pub async fn get_conversation(
87        &self,
88        conversation_id: Uuid,
89    ) -> CodexResult<Arc<CodexConversation>> {
90        let conversations = self.conversations.read().await;
91        conversations
92            .get(&conversation_id)
93            .cloned()
94            .ok_or_else(|| CodexErr::ConversationNotFound(conversation_id))
95    }
96}