1use std::collections::{BTreeMap, VecDeque};
2use std::path::{Path, PathBuf};
3use std::sync::{Arc, Mutex};
4
5use serde_json::Value;
6
7use crate::runtime::background_sessions::BackgroundSessionSubscription;
8use crate::runtime::sub_task_manager::SubTaskManager;
9use crate::runtime::CancellationToken;
10use crate::types::Metadata;
11
12use super::types::{AgentDefinition, AgentRun};
13
14mod events;
15mod handles;
16mod run;
17mod state;
18mod util;
19mod watchers;
20
21pub use handles::{SessionCancellationHandle, SessionSteeringHandle};
22pub use state::{
23 AgentSessionRunRequest, AgentSessionState, SessionEventHandler, SessionListenerId,
24};
25
26use util::{absolutize_workspace, generate_session_id, normalize_session_id};
27
28pub struct AgentSession {
29 execute_run: Arc<dyn Fn(AgentSessionRunRequest) -> Result<AgentRun, String> + Send + Sync>,
30 session_id: String,
31 _agent_name: String,
32 _definition: AgentDefinition,
33 workspace: PathBuf,
34 shared_state: Metadata,
35 messages: Vec<crate::types::Message>,
36 latest_run: Option<AgentRun>,
37 running: bool,
38 listeners: Arc<Mutex<BTreeMap<SessionListenerId, SessionEventHandler>>>,
39 background_command_subscriptions: Arc<Mutex<BTreeMap<String, BackgroundSessionSubscription>>>,
40 next_listener_id: SessionListenerId,
41 steering_queue: Arc<Mutex<VecDeque<String>>>,
42 follow_up_queue: Arc<Mutex<VecDeque<String>>>,
43 active_cancellation_token: Arc<Mutex<Option<CancellationToken>>>,
44 sub_task_manager: SubTaskManager,
45}
46
47impl AgentSession {
48 pub fn new(
49 execute_run: Arc<dyn Fn(String) -> Result<AgentRun, String> + Send + Sync>,
50 agent_name: impl Into<String>,
51 definition: AgentDefinition,
52 workspace: impl Into<PathBuf>,
53 ) -> Self {
54 let execute_run =
55 Arc::new(move |request: AgentSessionRunRequest| execute_run(request.prompt));
56 Self::new_with_context(execute_run, agent_name, definition, workspace)
57 }
58
59 pub fn new_with_session_id(
60 execute_run: Arc<dyn Fn(String) -> Result<AgentRun, String> + Send + Sync>,
61 session_id: impl Into<String>,
62 agent_name: impl Into<String>,
63 definition: AgentDefinition,
64 workspace: impl Into<PathBuf>,
65 ) -> Self {
66 let execute_run =
67 Arc::new(move |request: AgentSessionRunRequest| execute_run(request.prompt));
68 Self::new_with_context_and_session_id(
69 execute_run,
70 session_id,
71 agent_name,
72 definition,
73 workspace,
74 )
75 }
76
77 pub fn new_with_context(
78 execute_run: Arc<dyn Fn(AgentSessionRunRequest) -> Result<AgentRun, String> + Send + Sync>,
79 agent_name: impl Into<String>,
80 definition: AgentDefinition,
81 workspace: impl Into<PathBuf>,
82 ) -> Self {
83 Self::new_with_context_and_session_id(
84 execute_run,
85 generate_session_id(),
86 agent_name,
87 definition,
88 workspace,
89 )
90 }
91
92 pub fn new_with_context_and_shared_state(
93 execute_run: Arc<dyn Fn(AgentSessionRunRequest) -> Result<AgentRun, String> + Send + Sync>,
94 agent_name: impl Into<String>,
95 definition: AgentDefinition,
96 workspace: impl Into<PathBuf>,
97 shared_state: Metadata,
98 ) -> Self {
99 Self::new_with_context_and_session_id_and_shared_state(
100 execute_run,
101 generate_session_id(),
102 agent_name,
103 definition,
104 workspace,
105 shared_state,
106 )
107 }
108
109 pub fn new_with_context_and_session_id(
110 execute_run: Arc<dyn Fn(AgentSessionRunRequest) -> Result<AgentRun, String> + Send + Sync>,
111 session_id: impl Into<String>,
112 agent_name: impl Into<String>,
113 definition: AgentDefinition,
114 workspace: impl Into<PathBuf>,
115 ) -> Self {
116 Self::new_with_context_and_session_id_and_shared_state(
117 execute_run,
118 session_id,
119 agent_name,
120 definition,
121 workspace,
122 Metadata::new(),
123 )
124 }
125
126 pub fn new_with_context_and_session_id_and_shared_state(
127 execute_run: Arc<dyn Fn(AgentSessionRunRequest) -> Result<AgentRun, String> + Send + Sync>,
128 session_id: impl Into<String>,
129 agent_name: impl Into<String>,
130 definition: AgentDefinition,
131 workspace: impl Into<PathBuf>,
132 shared_state: Metadata,
133 ) -> Self {
134 let mut shared_state = shared_state;
135 shared_state
136 .entry("todo_list".to_string())
137 .or_insert_with(|| Value::Array(Vec::new()));
138 let workspace = absolutize_workspace(workspace.into());
139 Self {
140 execute_run,
141 session_id: normalize_session_id(session_id),
142 _agent_name: agent_name.into(),
143 _definition: definition,
144 workspace,
145 shared_state,
146 messages: Vec::new(),
147 latest_run: None,
148 running: false,
149 listeners: Arc::new(Mutex::new(BTreeMap::new())),
150 background_command_subscriptions: Arc::new(Mutex::new(BTreeMap::new())),
151 next_listener_id: 1,
152 steering_queue: Arc::new(Mutex::new(VecDeque::new())),
153 follow_up_queue: Arc::new(Mutex::new(VecDeque::new())),
154 active_cancellation_token: Arc::new(Mutex::new(None)),
155 sub_task_manager: SubTaskManager::default(),
156 }
157 }
158
159 pub fn session_id(&self) -> &str {
160 &self.session_id
161 }
162
163 pub fn agent_name(&self) -> &str {
164 &self._agent_name
165 }
166
167 pub fn definition(&self) -> &AgentDefinition {
168 &self._definition
169 }
170
171 pub fn workspace(&self) -> &Path {
172 &self.workspace
173 }
174
175 pub fn messages(&self) -> Vec<crate::types::Message> {
176 self.messages.clone()
177 }
178
179 pub fn shared_state(&self) -> Metadata {
180 self.shared_state.clone()
181 }
182
183 pub fn latest_run(&self) -> Option<AgentRun> {
184 self.latest_run.clone()
185 }
186
187 pub fn running(&self) -> bool {
188 self.running
189 }
190}