mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
feat(app-server): thread/unsubscribe API (#10954)
Adds a new v2 app-server API for a client to be able to unsubscribe to a thread: - New RPC method: `thread/unsubscribe` - New server notification: `thread/closed` Today clients can start/resume/archive threads, but there wasn’t a way to explicitly unload a live thread from memory without archiving it. With `thread/unsubscribe`, a client can indicate it is no longer actively working with a live Thread. If this is the only client subscribed to that given thread, the thread will be automatically closed by app-server, at which point the server will send `thread/closed` and `thread/status/changed` with `status: notLoaded` notifications. This gives clients a way to prevent long-running app-server processes from accumulating too many thread (and related) objects in memory. Closed threads will also be removed from `thread/loaded/list`.
This commit is contained in:
committed by
GitHub
Unverified
parent
d45ffd5830
commit
21f7032dbb
@@ -2390,7 +2390,11 @@ mod tests {
|
||||
let event_turn_id = "complete1".to_string();
|
||||
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(tx));
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(outgoing, vec![ConnectionId(1)]);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing,
|
||||
vec![ConnectionId(1)],
|
||||
ThreadId::new(),
|
||||
);
|
||||
let thread_state = new_thread_state();
|
||||
|
||||
handle_turn_complete(
|
||||
@@ -2431,7 +2435,11 @@ mod tests {
|
||||
.await;
|
||||
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(tx));
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(outgoing, vec![ConnectionId(1)]);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing,
|
||||
vec![ConnectionId(1)],
|
||||
ThreadId::new(),
|
||||
);
|
||||
|
||||
handle_turn_interrupted(
|
||||
conversation_id,
|
||||
@@ -2471,7 +2479,11 @@ mod tests {
|
||||
.await;
|
||||
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(tx));
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(outgoing, vec![ConnectionId(1)]);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing,
|
||||
vec![ConnectionId(1)],
|
||||
ThreadId::new(),
|
||||
);
|
||||
|
||||
handle_turn_complete(
|
||||
conversation_id,
|
||||
@@ -2505,7 +2517,11 @@ mod tests {
|
||||
async fn test_handle_turn_plan_update_emits_notification_for_v2() -> Result<()> {
|
||||
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(tx));
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(outgoing, vec![ConnectionId(1)]);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing,
|
||||
vec![ConnectionId(1)],
|
||||
ThreadId::new(),
|
||||
);
|
||||
let update = UpdatePlanArgs {
|
||||
explanation: Some("need plan".to_string()),
|
||||
plan: vec![
|
||||
@@ -2555,7 +2571,11 @@ mod tests {
|
||||
let turn_id = "turn-123".to_string();
|
||||
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(tx));
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(outgoing, vec![ConnectionId(1)]);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing,
|
||||
vec![ConnectionId(1)],
|
||||
ThreadId::new(),
|
||||
);
|
||||
|
||||
let info = TokenUsageInfo {
|
||||
total_token_usage: TokenUsage {
|
||||
@@ -2639,7 +2659,11 @@ mod tests {
|
||||
let turn_id = "turn-456".to_string();
|
||||
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(tx));
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(outgoing, vec![ConnectionId(1)]);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing,
|
||||
vec![ConnectionId(1)],
|
||||
ThreadId::new(),
|
||||
);
|
||||
|
||||
handle_token_count_event(
|
||||
conversation_id,
|
||||
@@ -2706,7 +2730,11 @@ mod tests {
|
||||
|
||||
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(tx));
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(outgoing, vec![ConnectionId(1)]);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing,
|
||||
vec![ConnectionId(1)],
|
||||
ThreadId::new(),
|
||||
);
|
||||
|
||||
// Turn 1 on conversation A
|
||||
let a_turn1 = "a_turn1".to_string();
|
||||
@@ -2929,7 +2957,11 @@ mod tests {
|
||||
async fn test_handle_turn_diff_emits_v2_notification() -> Result<()> {
|
||||
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(tx));
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(outgoing, vec![ConnectionId(1)]);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing,
|
||||
vec![ConnectionId(1)],
|
||||
ThreadId::new(),
|
||||
);
|
||||
let unified_diff = "--- a\n+++ b\n".to_string();
|
||||
let conversation_id = ThreadId::new();
|
||||
|
||||
@@ -2963,7 +2995,11 @@ mod tests {
|
||||
async fn test_handle_turn_diff_is_noop_for_v1() -> Result<()> {
|
||||
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(tx));
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(outgoing, vec![ConnectionId(1)]);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing,
|
||||
vec![ConnectionId(1)],
|
||||
ThreadId::new(),
|
||||
);
|
||||
let conversation_id = ThreadId::new();
|
||||
|
||||
handle_turn_diff(
|
||||
|
||||
@@ -127,6 +127,7 @@ use codex_app_server_protocol::ThreadArchiveResponse;
|
||||
use codex_app_server_protocol::ThreadArchivedNotification;
|
||||
use codex_app_server_protocol::ThreadBackgroundTerminalsCleanParams;
|
||||
use codex_app_server_protocol::ThreadBackgroundTerminalsCleanResponse;
|
||||
use codex_app_server_protocol::ThreadClosedNotification;
|
||||
use codex_app_server_protocol::ThreadCompactStartParams;
|
||||
use codex_app_server_protocol::ThreadCompactStartResponse;
|
||||
use codex_app_server_protocol::ThreadForkParams;
|
||||
@@ -160,6 +161,9 @@ use codex_app_server_protocol::ThreadStatus;
|
||||
use codex_app_server_protocol::ThreadUnarchiveParams;
|
||||
use codex_app_server_protocol::ThreadUnarchiveResponse;
|
||||
use codex_app_server_protocol::ThreadUnarchivedNotification;
|
||||
use codex_app_server_protocol::ThreadUnsubscribeParams;
|
||||
use codex_app_server_protocol::ThreadUnsubscribeResponse;
|
||||
use codex_app_server_protocol::ThreadUnsubscribeStatus;
|
||||
use codex_app_server_protocol::Turn;
|
||||
use codex_app_server_protocol::TurnInterruptParams;
|
||||
use codex_app_server_protocol::TurnStartParams;
|
||||
@@ -323,6 +327,12 @@ enum AppListLoadResult {
|
||||
Directory(Result<Vec<AppInfo>, String>),
|
||||
}
|
||||
|
||||
enum ThreadShutdownResult {
|
||||
Complete,
|
||||
SubmitFailed,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
fn convert_remote_scope(scope: ApiHazelnutScope) -> RemoteSkillHazelnutScope {
|
||||
match scope {
|
||||
ApiHazelnutScope::WorkspaceShared => RemoteSkillHazelnutScope::WorkspaceShared,
|
||||
@@ -358,6 +368,7 @@ pub(crate) struct CodexMessageProcessor {
|
||||
cli_overrides: Vec<(String, TomlValue)>,
|
||||
cloud_requirements: Arc<RwLock<CloudRequirementsLoader>>,
|
||||
active_login: Arc<Mutex<Option<ActiveLogin>>>,
|
||||
pending_thread_unloads: Arc<Mutex<HashSet<ThreadId>>>,
|
||||
thread_state_manager: ThreadStateManager,
|
||||
thread_watch_manager: ThreadWatchManager,
|
||||
pending_fuzzy_searches: Arc<Mutex<HashMap<String, Arc<AtomicBool>>>>,
|
||||
@@ -430,6 +441,7 @@ impl CodexMessageProcessor {
|
||||
cli_overrides,
|
||||
cloud_requirements,
|
||||
active_login: Arc::new(Mutex::new(None)),
|
||||
pending_thread_unloads: Arc::new(Mutex::new(HashSet::new())),
|
||||
thread_state_manager: ThreadStateManager::new(),
|
||||
thread_watch_manager: ThreadWatchManager::new_with_outgoing(outgoing),
|
||||
pending_fuzzy_searches: Arc::new(Mutex::new(HashMap::new())),
|
||||
@@ -557,6 +569,10 @@ impl CodexMessageProcessor {
|
||||
self.thread_start(to_connection_request_id(request_id), params)
|
||||
.await;
|
||||
}
|
||||
ClientRequest::ThreadUnsubscribe { request_id, params } => {
|
||||
self.thread_unsubscribe(to_connection_request_id(request_id), params)
|
||||
.await;
|
||||
}
|
||||
ClientRequest::ThreadResume { request_id, params } => {
|
||||
self.thread_resume(to_connection_request_id(request_id), params)
|
||||
.await;
|
||||
@@ -2876,6 +2892,23 @@ impl CodexMessageProcessor {
|
||||
}
|
||||
|
||||
async fn thread_resume(&mut self, request_id: ConnectionRequestId, params: ThreadResumeParams) {
|
||||
if let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id)
|
||||
&& self
|
||||
.pending_thread_unloads
|
||||
.lock()
|
||||
.await
|
||||
.contains(&thread_id)
|
||||
{
|
||||
self.send_invalid_request_error(
|
||||
request_id,
|
||||
format!(
|
||||
"thread {thread_id} is closing; retry thread/resume after the thread is closed"
|
||||
),
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
if self
|
||||
.resume_running_thread(request_id.clone(), ¶ms)
|
||||
.await
|
||||
@@ -4729,6 +4762,150 @@ impl CodexMessageProcessor {
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_thread_shutdown(thread: &Arc<CodexThread>) -> ThreadShutdownResult {
|
||||
match thread.submit(Op::Shutdown).await {
|
||||
Ok(_) => {
|
||||
let wait_for_shutdown = async {
|
||||
loop {
|
||||
if matches!(thread.agent_status().await, AgentStatus::Shutdown) {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
}
|
||||
};
|
||||
if tokio::time::timeout(Duration::from_secs(10), wait_for_shutdown)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
ThreadShutdownResult::TimedOut
|
||||
} else {
|
||||
ThreadShutdownResult::Complete
|
||||
}
|
||||
}
|
||||
Err(_) => ThreadShutdownResult::SubmitFailed,
|
||||
}
|
||||
}
|
||||
|
||||
async fn finalize_thread_teardown(&mut self, thread_id: ThreadId) {
|
||||
self.pending_thread_unloads.lock().await.remove(&thread_id);
|
||||
self.outgoing.cancel_requests_for_thread(thread_id).await;
|
||||
self.thread_state_manager
|
||||
.remove_thread_state(thread_id)
|
||||
.await;
|
||||
self.thread_watch_manager
|
||||
.remove_thread(&thread_id.to_string())
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn thread_unsubscribe(
|
||||
&mut self,
|
||||
request_id: ConnectionRequestId,
|
||||
params: ThreadUnsubscribeParams,
|
||||
) {
|
||||
let thread_id = match ThreadId::from_string(¶ms.thread_id) {
|
||||
Ok(id) => id,
|
||||
Err(err) => {
|
||||
self.send_invalid_request_error(request_id, format!("invalid thread id: {err}"))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let Ok(thread) = self.thread_manager.get_thread(thread_id).await else {
|
||||
// Reconcile stale app-server bookkeeping when the thread has already been
|
||||
// removed from the core manager. This keeps loaded-status/subscription state
|
||||
// consistent with the source of truth before reporting NotLoaded.
|
||||
self.finalize_thread_teardown(thread_id).await;
|
||||
self.outgoing
|
||||
.send_response(
|
||||
request_id,
|
||||
ThreadUnsubscribeResponse {
|
||||
status: ThreadUnsubscribeStatus::NotLoaded,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
};
|
||||
|
||||
let was_subscribed = self
|
||||
.thread_state_manager
|
||||
.unsubscribe_connection_from_thread(thread_id, request_id.connection_id)
|
||||
.await;
|
||||
if !was_subscribed {
|
||||
self.outgoing
|
||||
.send_response(
|
||||
request_id,
|
||||
ThreadUnsubscribeResponse {
|
||||
status: ThreadUnsubscribeStatus::NotSubscribed,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
if !self.thread_state_manager.has_subscribers(thread_id).await {
|
||||
// This connection was the last subscriber. Only now do we unload the thread.
|
||||
info!("thread {thread_id} has no subscribers; shutting down");
|
||||
self.pending_thread_unloads.lock().await.insert(thread_id);
|
||||
// Any pending app-server -> client requests for this thread can no longer be
|
||||
// answered; cancel their callbacks before shutdown/unload.
|
||||
self.outgoing.cancel_requests_for_thread(thread_id).await;
|
||||
self.thread_state_manager
|
||||
.remove_thread_state(thread_id)
|
||||
.await;
|
||||
|
||||
let outgoing = self.outgoing.clone();
|
||||
let pending_thread_unloads = self.pending_thread_unloads.clone();
|
||||
let thread_manager = self.thread_manager.clone();
|
||||
let thread_watch_manager = self.thread_watch_manager.clone();
|
||||
tokio::spawn(async move {
|
||||
match Self::wait_for_thread_shutdown(&thread).await {
|
||||
ThreadShutdownResult::Complete => {
|
||||
if thread_manager.remove_thread(&thread_id).await.is_none() {
|
||||
info!(
|
||||
"thread {thread_id} was already removed before unsubscribe finalized"
|
||||
);
|
||||
thread_watch_manager
|
||||
.remove_thread(&thread_id.to_string())
|
||||
.await;
|
||||
pending_thread_unloads.lock().await.remove(&thread_id);
|
||||
return;
|
||||
}
|
||||
thread_watch_manager
|
||||
.remove_thread(&thread_id.to_string())
|
||||
.await;
|
||||
let notification = ThreadClosedNotification {
|
||||
thread_id: thread_id.to_string(),
|
||||
};
|
||||
outgoing
|
||||
.send_server_notification(ServerNotification::ThreadClosed(
|
||||
notification,
|
||||
))
|
||||
.await;
|
||||
pending_thread_unloads.lock().await.remove(&thread_id);
|
||||
}
|
||||
ThreadShutdownResult::SubmitFailed => {
|
||||
pending_thread_unloads.lock().await.remove(&thread_id);
|
||||
warn!("failed to submit Shutdown to thread {thread_id}");
|
||||
}
|
||||
ThreadShutdownResult::TimedOut => {
|
||||
pending_thread_unloads.lock().await.remove(&thread_id);
|
||||
warn!("thread {thread_id} shutdown timed out; leaving thread loaded");
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
self.outgoing
|
||||
.send_response(
|
||||
request_id,
|
||||
ThreadUnsubscribeResponse {
|
||||
status: ThreadUnsubscribeStatus::Unsubscribed,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn archive_thread_common(
|
||||
&mut self,
|
||||
thread_id: ThreadId,
|
||||
@@ -4800,37 +4977,19 @@ impl CodexMessageProcessor {
|
||||
state_db_ctx = Some(ctx);
|
||||
}
|
||||
info!("thread {thread_id} was active; shutting down");
|
||||
// Request shutdown.
|
||||
match conversation.submit(Op::Shutdown).await {
|
||||
Ok(_) => {
|
||||
// Poll agent status rather than consuming events so attached listeners do not block shutdown.
|
||||
let wait_for_shutdown = async {
|
||||
loop {
|
||||
if matches!(conversation.agent_status().await, AgentStatus::Shutdown) {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
}
|
||||
};
|
||||
if tokio::time::timeout(Duration::from_secs(10), wait_for_shutdown)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
warn!("thread {thread_id} shutdown timed out; proceeding with archive");
|
||||
}
|
||||
match Self::wait_for_thread_shutdown(&conversation).await {
|
||||
ThreadShutdownResult::Complete => {}
|
||||
ThreadShutdownResult::SubmitFailed => {
|
||||
error!(
|
||||
"failed to submit Shutdown to thread {thread_id}; proceeding with archive"
|
||||
);
|
||||
}
|
||||
Err(err) => {
|
||||
error!("failed to submit Shutdown to thread {thread_id}: {err}");
|
||||
ThreadShutdownResult::TimedOut => {
|
||||
warn!("thread {thread_id} shutdown timed out; proceeding with archive");
|
||||
}
|
||||
}
|
||||
self.thread_state_manager
|
||||
.remove_thread_state(thread_id)
|
||||
.await;
|
||||
}
|
||||
|
||||
self.thread_watch_manager
|
||||
.remove_thread(&thread_id.to_string())
|
||||
.await;
|
||||
self.finalize_thread_teardown(thread_id).await;
|
||||
|
||||
if state_db_ctx.is_none() {
|
||||
state_db_ctx = get_state_db(&self.config, None).await;
|
||||
@@ -6180,6 +6339,7 @@ impl CodexMessageProcessor {
|
||||
let thread_outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
outgoing_for_task.clone(),
|
||||
subscribed_connection_ids,
|
||||
conversation_id,
|
||||
);
|
||||
apply_bespoke_event_handling(
|
||||
event.clone(),
|
||||
@@ -6482,10 +6642,8 @@ impl CodexMessageProcessor {
|
||||
WindowsSandboxSetupMode::Unelevated => CoreWindowsSandboxSetupMode::Unelevated,
|
||||
};
|
||||
let config = Arc::clone(&self.config);
|
||||
let outgoing = ThreadScopedOutgoingMessageSender::new(
|
||||
Arc::clone(&self.outgoing),
|
||||
vec![request_id.connection_id],
|
||||
);
|
||||
let outgoing = Arc::clone(&self.outgoing);
|
||||
let connection_id = request_id.connection_id;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let setup_request = WindowsSandboxSetupRequest {
|
||||
@@ -6508,9 +6666,10 @@ impl CodexMessageProcessor {
|
||||
error: setup_result.err().map(|err| err.to_string()),
|
||||
};
|
||||
outgoing
|
||||
.send_server_notification(ServerNotification::WindowsSandboxSetupCompleted(
|
||||
notification,
|
||||
))
|
||||
.send_server_notification_to_connections(
|
||||
&[connection_id],
|
||||
ServerNotification::WindowsSandboxSetupCompleted(notification),
|
||||
)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
@@ -87,7 +87,7 @@ impl ExternalAuthRefresher for ExternalAuthRefreshBridge {
|
||||
|
||||
let (request_id, rx) = self
|
||||
.outgoing
|
||||
.send_request_with_id(ServerRequestPayload::ChatgptAuthTokensRefresh(params))
|
||||
.send_request(ServerRequestPayload::ChatgptAuthTokensRefresh(params))
|
||||
.await;
|
||||
|
||||
let result = match timeout(EXTERNAL_AUTH_REFRESH_TIMEOUT, rx).await {
|
||||
|
||||
@@ -9,6 +9,7 @@ use codex_app_server_protocol::Result;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use codex_app_server_protocol::ServerRequest;
|
||||
use codex_app_server_protocol::ServerRequestPayload;
|
||||
use codex_protocol::ThreadId;
|
||||
use serde::Serialize;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::mpsc;
|
||||
@@ -48,23 +49,31 @@ pub(crate) enum OutgoingEnvelope {
|
||||
pub(crate) struct OutgoingMessageSender {
|
||||
next_server_request_id: AtomicI64,
|
||||
sender: mpsc::Sender<OutgoingEnvelope>,
|
||||
request_id_to_callback: Mutex<HashMap<RequestId, oneshot::Sender<ClientRequestResult>>>,
|
||||
request_id_to_callback: Mutex<HashMap<RequestId, PendingCallbackEntry>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ThreadScopedOutgoingMessageSender {
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
connection_ids: Arc<Vec<ConnectionId>>,
|
||||
thread_id: ThreadId,
|
||||
}
|
||||
|
||||
struct PendingCallbackEntry {
|
||||
callback: oneshot::Sender<ClientRequestResult>,
|
||||
thread_id: Option<ThreadId>,
|
||||
}
|
||||
|
||||
impl ThreadScopedOutgoingMessageSender {
|
||||
pub(crate) fn new(
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
connection_ids: Vec<ConnectionId>,
|
||||
thread_id: ThreadId,
|
||||
) -> Self {
|
||||
Self {
|
||||
outgoing,
|
||||
connection_ids: Arc::new(connection_ids),
|
||||
thread_id,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,12 +81,12 @@ impl ThreadScopedOutgoingMessageSender {
|
||||
&self,
|
||||
payload: ServerRequestPayload,
|
||||
) -> oneshot::Receiver<ClientRequestResult> {
|
||||
if self.connection_ids.is_empty() {
|
||||
let (_tx, rx) = oneshot::channel();
|
||||
return rx;
|
||||
}
|
||||
self.outgoing
|
||||
.send_request_to_connections(self.connection_ids.as_slice(), payload)
|
||||
.send_request_to_thread_connections(
|
||||
self.thread_id,
|
||||
self.connection_ids.as_slice(),
|
||||
payload,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -116,35 +125,52 @@ impl OutgoingMessageSender {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn send_request_to_connections(
|
||||
&self,
|
||||
connection_ids: &[ConnectionId],
|
||||
request: ServerRequestPayload,
|
||||
) -> oneshot::Receiver<ClientRequestResult> {
|
||||
let (_id, rx) = self
|
||||
.send_request_with_id_to_connections(connection_ids, request)
|
||||
.await;
|
||||
rx
|
||||
}
|
||||
|
||||
pub(crate) async fn send_request_with_id(
|
||||
pub(crate) async fn send_request(
|
||||
&self,
|
||||
request: ServerRequestPayload,
|
||||
) -> (RequestId, oneshot::Receiver<ClientRequestResult>) {
|
||||
self.send_request_with_id_to_connections(&[], request).await
|
||||
self.send_request_with_id_to_connections(&[], request, None)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_request_to_thread_connections(
|
||||
&self,
|
||||
thread_id: ThreadId,
|
||||
connection_ids: &[ConnectionId],
|
||||
request: ServerRequestPayload,
|
||||
) -> oneshot::Receiver<ClientRequestResult> {
|
||||
if connection_ids.is_empty() {
|
||||
let (_tx, rx) = oneshot::channel();
|
||||
return rx;
|
||||
}
|
||||
let (_request_id, receiver) = self
|
||||
.send_request_with_id_to_connections(connection_ids, request, Some(thread_id))
|
||||
.await;
|
||||
receiver
|
||||
}
|
||||
|
||||
fn next_request_id(&self) -> RequestId {
|
||||
RequestId::Integer(self.next_server_request_id.fetch_add(1, Ordering::Relaxed))
|
||||
}
|
||||
|
||||
async fn send_request_with_id_to_connections(
|
||||
&self,
|
||||
connection_ids: &[ConnectionId],
|
||||
request: ServerRequestPayload,
|
||||
thread_id: Option<ThreadId>,
|
||||
) -> (RequestId, oneshot::Receiver<ClientRequestResult>) {
|
||||
let id = RequestId::Integer(self.next_server_request_id.fetch_add(1, Ordering::Relaxed));
|
||||
let id = self.next_request_id();
|
||||
let outgoing_message_id = id.clone();
|
||||
let (tx_approve, rx_approve) = oneshot::channel();
|
||||
{
|
||||
let mut request_id_to_callback = self.request_id_to_callback.lock().await;
|
||||
request_id_to_callback.insert(id, tx_approve);
|
||||
request_id_to_callback.insert(
|
||||
id,
|
||||
PendingCallbackEntry {
|
||||
callback: tx_approve,
|
||||
thread_id,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let outgoing_message =
|
||||
@@ -191,8 +217,8 @@ impl OutgoingMessageSender {
|
||||
};
|
||||
|
||||
match entry {
|
||||
Some((id, sender)) => {
|
||||
if let Err(err) = sender.send(Ok(result)) {
|
||||
Some((id, entry)) => {
|
||||
if let Err(err) = entry.callback.send(Ok(result)) {
|
||||
warn!("could not notify callback for {id:?} due to: {err:?}");
|
||||
}
|
||||
}
|
||||
@@ -209,9 +235,9 @@ impl OutgoingMessageSender {
|
||||
};
|
||||
|
||||
match entry {
|
||||
Some((id, sender)) => {
|
||||
Some((id, entry)) => {
|
||||
warn!("client responded with error for {id:?}: {error:?}");
|
||||
if let Err(err) = sender.send(Err(error)) {
|
||||
if let Err(err) = entry.callback.send(Err(error)) {
|
||||
warn!("could not notify callback for {id:?} due to: {err:?}");
|
||||
}
|
||||
}
|
||||
@@ -229,6 +255,19 @@ impl OutgoingMessageSender {
|
||||
entry.is_some()
|
||||
}
|
||||
|
||||
pub(crate) async fn cancel_requests_for_thread(&self, thread_id: ThreadId) {
|
||||
let mut request_id_to_callback = self.request_id_to_callback.lock().await;
|
||||
let request_ids = request_id_to_callback
|
||||
.iter()
|
||||
.filter_map(|(request_id, entry)| {
|
||||
(entry.thread_id == Some(thread_id)).then_some(request_id.clone())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
for request_id in request_ids {
|
||||
request_id_to_callback.remove(&request_id);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn send_response<T: Serialize>(
|
||||
&self,
|
||||
request_id: ConnectionRequestId,
|
||||
@@ -657,7 +696,7 @@ mod tests {
|
||||
let outgoing = OutgoingMessageSender::new(tx);
|
||||
|
||||
let (request_id, wait_for_result) = outgoing
|
||||
.send_request_with_id(ServerRequestPayload::ApplyPatchApproval(
|
||||
.send_request(ServerRequestPayload::ApplyPatchApproval(
|
||||
ApplyPatchApprovalParams {
|
||||
conversation_id: ThreadId::new(),
|
||||
call_id: "call-id".to_string(),
|
||||
|
||||
@@ -207,6 +207,50 @@ impl ThreadStateManager {
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) async fn unsubscribe_connection_from_thread(
|
||||
&mut self,
|
||||
thread_id: ThreadId,
|
||||
connection_id: ConnectionId,
|
||||
) -> bool {
|
||||
let Some(thread_state) = self.thread_states.get(&thread_id) else {
|
||||
return false;
|
||||
};
|
||||
|
||||
if !self
|
||||
.thread_ids_by_connection
|
||||
.get(&connection_id)
|
||||
.is_some_and(|thread_ids| thread_ids.contains(&thread_id))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
if let Some(thread_ids) = self.thread_ids_by_connection.get_mut(&connection_id) {
|
||||
thread_ids.remove(&thread_id);
|
||||
if thread_ids.is_empty() {
|
||||
self.thread_ids_by_connection.remove(&connection_id);
|
||||
}
|
||||
}
|
||||
|
||||
self.subscription_state_by_id.retain(|_, state| {
|
||||
!(state.thread_id == thread_id && state.connection_id == connection_id)
|
||||
});
|
||||
|
||||
let mut thread_state = thread_state.lock().await;
|
||||
thread_state.remove_connection(connection_id);
|
||||
true
|
||||
}
|
||||
|
||||
pub(crate) async fn has_subscribers(&self, thread_id: ThreadId) -> bool {
|
||||
let Some(thread_state) = self.thread_states.get(&thread_id) else {
|
||||
return false;
|
||||
};
|
||||
!thread_state
|
||||
.lock()
|
||||
.await
|
||||
.subscribed_connection_ids()
|
||||
.is_empty()
|
||||
}
|
||||
|
||||
pub(crate) async fn set_listener(
|
||||
&mut self,
|
||||
subscription_id: Uuid,
|
||||
|
||||
@@ -300,8 +300,16 @@ impl ThreadWatchState {
|
||||
}
|
||||
|
||||
fn remove_thread(&mut self, thread_id: &str) -> Option<ThreadStatusChangedNotification> {
|
||||
let previous_status = self.status_for(thread_id);
|
||||
self.runtime_by_thread_id.remove(thread_id);
|
||||
None
|
||||
if previous_status.is_some() && previous_status != Some(ThreadStatus::NotLoaded) {
|
||||
Some(ThreadStatusChangedNotification {
|
||||
thread_id: thread_id.to_string(),
|
||||
status: ThreadStatus::NotLoaded,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn update_runtime<F>(
|
||||
@@ -673,6 +681,15 @@ mod tests {
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
manager.remove_thread(INTERACTIVE_THREAD_ID).await;
|
||||
assert_eq!(
|
||||
recv_status_changed_notification(&mut outgoing_rx).await,
|
||||
ThreadStatusChangedNotification {
|
||||
thread_id: INTERACTIVE_THREAD_ID.to_string(),
|
||||
status: ThreadStatus::NotLoaded,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
async fn wait_for_status(
|
||||
|
||||
Reference in New Issue
Block a user