| use std::sync::Arc; |
| use std::sync::Weak; |
|
|
| use codex_code_mode_protocol::CellId; |
| use codex_code_mode_protocol::CodeModeNestedToolCall; |
| use codex_code_mode_protocol::CodeModeSessionDelegate; |
| use codex_code_mode_protocol::NotificationFuture; |
| use codex_code_mode_protocol::ToolInvocationFuture; |
| use codex_code_mode_protocol::grpc as proto; |
| use tokio::sync::oneshot; |
| use tokio_util::sync::CancellationToken; |
| use uuid::Uuid; |
|
|
| use super::session::GrpcSession; |
|
|
| pub(super) struct GrpcDelegate { |
| session: Weak<GrpcSession>, |
| } |
|
|
| impl GrpcDelegate { |
| pub(super) fn new(session: Weak<GrpcSession>) -> Self { |
| Self { session } |
| } |
| } |
|
|
| impl CodeModeSessionDelegate for GrpcDelegate { |
| fn invoke_tool<'a>( |
| &'a self, |
| invocation: CodeModeNestedToolCall, |
| cancellation: CancellationToken, |
| ) -> ToolInvocationFuture<'a> { |
| Box::pin(async move { |
| let session = self |
| .session |
| .upgrade() |
| .ok_or_else(|| "code-mode session is closed".to_string())?; |
| let _permit = session.delegate_permit()?; |
| let execution_id = session |
| .execution_id(invocation.cell_id.as_str(), &cancellation) |
| .await?; |
| let input_json = invocation |
| .input |
| .as_ref() |
| .map(serde_json::to_vec) |
| .transpose() |
| .map_err(|error| format!("failed to encode code-mode tool input: {error}"))?; |
| let invocation_id = Uuid::new_v4(); |
| let (response, receiver) = oneshot::channel(); |
| session |
| .dispatch_tool( |
| invocation, |
| execution_id, |
| invocation_id, |
| input_json, |
| response, |
| &cancellation, |
| ) |
| .await?; |
| let mut pending = PendingToolCall { |
| session: Arc::clone(&session), |
| id: Some(invocation_id), |
| }; |
| let result = tokio::select! { |
| biased; |
| result = receiver => result |
| .map_err(|_| "code-mode client closed before returning tool output".to_string())?, |
| _ = cancellation.cancelled() => { |
| Err("code mode delegate request cancelled".to_string()) |
| } |
| _ = session.closed.cancelled() => { |
| Err("code-mode session closed before returning tool output".to_string()) |
| } |
| }; |
| if result.is_ok() { |
| pending.id = None; |
| } |
| result |
| }) |
| } |
|
|
| fn notify<'a>( |
| &'a self, |
| call_id: String, |
| cell_id: CellId, |
| text: String, |
| cancellation: CancellationToken, |
| ) -> NotificationFuture<'a> { |
| Box::pin(async move { |
| let session = self |
| .session |
| .upgrade() |
| .ok_or_else(|| "code-mode session is closed".to_string())?; |
| let _permit = session.delegate_permit()?; |
| let execution_id = session |
| .execution_id(cell_id.as_str(), &cancellation) |
| .await?; |
| let notification_id = Uuid::new_v4(); |
| session |
| .send_event( |
| proto::session_event::Event::Notification(proto::Notification { |
| notification_id: notification_id.to_string(), |
| execution_id, |
| cell_id: cell_id.to_string(), |
| call_id, |
| text, |
| }), |
| &cancellation, |
| ) |
| .await |
| }) |
| } |
|
|
| fn cell_closed(&self, cell_id: &CellId) { |
| if let Some(session) = self.session.upgrade() { |
| session.close_cell(cell_id.as_str()); |
| } |
| } |
| } |
|
|
| struct PendingToolCall { |
| session: Arc<GrpcSession>, |
| id: Option<Uuid>, |
| } |
|
|
| impl Drop for PendingToolCall { |
| fn drop(&mut self) { |
| if let Some(id) = self.id.take() { |
| self.session.cancel_invocation(id); |
| } |
| } |
| } |
|
|
