use crate::client::ClaudeClient;
use crate::errors::{ClaudeError, Result};
use crate::types::config::ClaudeAgentOptions;
use crate::types::messages::Message;
use futures::StreamExt;
use std::sync::Arc;
use tokio::sync::Mutex;
use super::types::SessionOptions;
pub struct Session {
pub id: String,
pub options: SessionOptions,
client: Arc<Mutex<ClaudeClient>>,
}
impl Session {
fn new(id: String, options: SessionOptions, client: ClaudeClient) -> Self {
Self {
id,
options,
client: Arc::new(Mutex::new(client)),
}
}
pub async fn send(&mut self, message: impl Into<String>) -> Result<()> {
let message_text = message.into();
if message_text.trim().is_empty() {
return Err(ClaudeError::InvalidInput(
"Message cannot be empty".to_string(),
));
}
let mut client = self.client.lock().await;
client.query(&message_text).await?;
Ok(())
}
pub async fn receive(&self) -> Result<Vec<V2Message>> {
let client = self.client.lock().await;
let mut stream = client.receive_response();
let mut messages = Vec::new();
while let Some(result) = stream.next().await {
let msg = result?;
match msg {
Message::Assistant(assist_msg) => {
let content = assist_msg
.message
.content
.iter()
.filter_map(|block| match block {
crate::types::messages::ContentBlock::Text(text) => {
Some(text.text.clone())
}
_ => None,
})
.collect::<Vec<_>>()
.join("\n");
messages.push(V2Message::Assistant { content });
}
Message::Result(_) => {
break;
}
_ => {
}
}
}
Ok(messages)
}
pub async fn model(&self) -> Option<String> {
self.options.model.clone()
}
pub async fn is_connected(&self) -> bool {
true
}
pub async fn close(self) -> Result<()> {
Ok(())
}
}
#[derive(Debug, Clone)]
pub enum V2Message {
Assistant {
content: String,
},
}
impl V2Message {
pub fn as_text(&self) -> Option<&str> {
match self {
V2Message::Assistant { content } => Some(content),
}
}
}
pub async fn create_session(options: SessionOptions) -> Result<Session> {
let opts: ClaudeAgentOptions = options.clone().into();
let mut client = ClaudeClient::new(opts);
client.connect().await?;
let id = uuid::Uuid::new_v4().to_string();
Ok(Session::new(id, options, client))
}
pub async fn resume_session(
session_id: &str,
options: SessionOptions,
) -> Result<Session> {
let opts: ClaudeAgentOptions = options.clone().into();
let mut client = ClaudeClient::new(opts);
client.connect().await?;
Ok(Session::new(
session_id.to_string(),
options,
client,
))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_v2_message_as_text() {
let msg = V2Message::Assistant {
content: "Hello!".to_string(),
};
assert_eq!(msg.as_text(), Some("Hello!"));
}
}