chore: move 3-14 WS push design to done/
This commit is contained in:
@@ -0,0 +1,537 @@
|
||||
/// ACP Session Manager — session lifecycle management.
|
||||
///
|
||||
/// Wraps an [`AcpClient`] to provide strongly-typed session operations:
|
||||
/// create, prompt, cancel, and receive typed events from the agent.
|
||||
///
|
||||
/// Reference:
|
||||
/// - `reference-code/hermes-vscode-main/src/sessionManager.ts` (primary)
|
||||
/// - `reference-code/hermes-vscode-main/src/protocol.ts` (dedup logic)
|
||||
|
||||
use crate::acp_client::AcpClient;
|
||||
use crate::acp_types::{
|
||||
ContentBlock, SessionNewParams, SessionNewResult, SessionPromptParams,
|
||||
SessionPromptResult, SessionUpdate, ToolCallStatus,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
#[cfg(test)]
|
||||
use tokio::time::{sleep, Duration};
|
||||
|
||||
// ── Events ───────────────────────────────────────────
|
||||
|
||||
/// Strongly-typed event emitted by the session manager when a `session/update` arrives.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AcpSessionEvent {
|
||||
/// Streaming text from the agent's response message.
|
||||
TextDelta { text: String },
|
||||
/// Streaming reasoning/thinking text.
|
||||
ThoughtDelta { text: String },
|
||||
/// Tool call started.
|
||||
ToolCall {
|
||||
tool_call_id: String,
|
||||
title: String,
|
||||
kind: String,
|
||||
status: ToolCallStatus,
|
||||
},
|
||||
/// Tool call status update (with optional result content).
|
||||
ToolCallUpdate {
|
||||
tool_call_id: String,
|
||||
status: ToolCallStatus,
|
||||
},
|
||||
/// Context usage update.
|
||||
UsageUpdate { used: u64, size: u64 },
|
||||
/// Session metadata update (e.g. auto-title).
|
||||
SessionInfoUpdate { title: String },
|
||||
/// Plan entries update.
|
||||
PlanUpdate { entries: Vec<String> },
|
||||
/// Connection closed/error.
|
||||
Disconnected { reason: String },
|
||||
}
|
||||
|
||||
/// Handler for session events.
|
||||
pub type SessionEventHandler = Arc<dyn Fn(AcpSessionEvent) + Send + Sync + 'static>;
|
||||
|
||||
// ── Session state ────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum SessionState {
|
||||
Idle,
|
||||
Running,
|
||||
Cancelling,
|
||||
Closed,
|
||||
}
|
||||
|
||||
// ── AcpSessionManager ────────────────────────────────
|
||||
|
||||
/// Manages ACP sessions — create, prompt, cancel, and event dispatch.
|
||||
///
|
||||
/// Currently supports one active session at a time. The internal [`AcpClient`]
|
||||
/// handles the JSON-RPC wire protocol; this layer adds session semantics and
|
||||
/// typed event dispatching.
|
||||
pub struct AcpSessionManager {
|
||||
client: Arc<AcpClient>,
|
||||
session_id: Arc<Mutex<Option<String>>>,
|
||||
state: Arc<Mutex<SessionState>>,
|
||||
event_handler: Arc<Mutex<Option<SessionEventHandler>>>,
|
||||
/// Accumulated text for deduplication (per-turn).
|
||||
accumulated: Arc<Mutex<String>>,
|
||||
/// Whether we're currently inside a prompt (for dedup gating).
|
||||
in_prompt: Arc<Mutex<bool>>,
|
||||
}
|
||||
|
||||
impl AcpSessionManager {
|
||||
/// Create a new session manager wrapping an existing [`AcpClient`].
|
||||
///
|
||||
/// Registers an internal notification handler that dispatches
|
||||
/// `session/update` notifications as typed [`AcpSessionEvent`]s.
|
||||
pub fn new(client: Arc<AcpClient>) -> Self {
|
||||
let session_id: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
|
||||
let state: Arc<Mutex<SessionState>> = Arc::new(Mutex::new(SessionState::Idle));
|
||||
let event_handler: Arc<Mutex<Option<SessionEventHandler>>> =
|
||||
Arc::new(Mutex::new(None));
|
||||
let accumulated: Arc<Mutex<String>> = Arc::new(Mutex::new(String::new()));
|
||||
let in_prompt: Arc<Mutex<bool>> = Arc::new(Mutex::new(false));
|
||||
|
||||
// Wire up the ACP notification handler
|
||||
let session_id_clone = session_id.clone();
|
||||
let event_handler_clone = event_handler.clone();
|
||||
let accumulated_clone = accumulated.clone();
|
||||
let in_prompt_clone = in_prompt.clone();
|
||||
|
||||
client.on_notification(move |method, params| {
|
||||
if method != "session/update" {
|
||||
return;
|
||||
}
|
||||
let sid = session_id_clone.lock().unwrap().clone();
|
||||
if let Some(ref session_id) = sid {
|
||||
if let Some(msg_sid) = params.get("sessionId").and_then(|v| v.as_str()) {
|
||||
if msg_sid != session_id {
|
||||
return; // not our session
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Parse the update
|
||||
let update: SessionUpdate = match serde_json::from_value(
|
||||
params.get("update").cloned().unwrap_or(Value::Null),
|
||||
) {
|
||||
Ok(u) => u,
|
||||
Err(e) => {
|
||||
warn!("ACP session/update parse error: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let is_in_prompt = *in_prompt_clone.lock().unwrap();
|
||||
let event = Self::session_update_to_event(&update, &accumulated_clone, is_in_prompt);
|
||||
|
||||
if let Some(ev) = event {
|
||||
let handler = event_handler_clone.lock().unwrap();
|
||||
if let Some(ref h) = *handler {
|
||||
h(ev);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Self {
|
||||
client,
|
||||
session_id,
|
||||
state,
|
||||
event_handler,
|
||||
accumulated,
|
||||
in_prompt,
|
||||
}
|
||||
}
|
||||
|
||||
/// Register an event handler for session events.
|
||||
/// Only one handler at a time — subsequent calls replace the previous.
|
||||
pub fn on_event<F>(&self, handler: F)
|
||||
where
|
||||
F: Fn(AcpSessionEvent) + Send + Sync + 'static,
|
||||
{
|
||||
let mut guard = self.event_handler.lock().unwrap();
|
||||
*guard = Some(Arc::new(handler));
|
||||
}
|
||||
|
||||
/// Create a new ACP session.
|
||||
///
|
||||
/// Sends `session/new` to the agent and stores the returned `sessionId`.
|
||||
/// The `page_context` is optional metadata about the current document.
|
||||
pub async fn create_session(
|
||||
&self,
|
||||
cwd: Option<&str>,
|
||||
_page_context: Option<Value>,
|
||||
) -> Result<String, crate::acp_client::AcpError> {
|
||||
let params = SessionNewParams {
|
||||
cwd: cwd.map(|s| s.to_string()),
|
||||
mcp_servers: None,
|
||||
};
|
||||
|
||||
let result: SessionNewResult = self.client.request("session/new", params).await?;
|
||||
let mut sid_guard = self.session_id.lock().unwrap();
|
||||
*sid_guard = Some(result.session_id.clone());
|
||||
info!("ACP session created: {}", result.session_id);
|
||||
|
||||
Ok(result.session_id)
|
||||
}
|
||||
|
||||
/// Send a prompt to the agent and stream events.
|
||||
///
|
||||
/// The `prompt` is a list of content blocks (text + optional page context).
|
||||
/// Returns once the agent finishes (stopReason received) or on error.
|
||||
///
|
||||
/// Sets `in_prompt = true` during the call to enable deduplication,
|
||||
/// then resets to `false` and clears accumulated text on completion.
|
||||
pub async fn run_prompt(
|
||||
&self,
|
||||
prompt: Vec<ContentBlock>,
|
||||
) -> Result<SessionPromptResult, crate::acp_client::AcpError> {
|
||||
{
|
||||
let mut in_prompt = self.in_prompt.lock().unwrap();
|
||||
*in_prompt = true;
|
||||
}
|
||||
{
|
||||
let mut acc = self.accumulated.lock().unwrap();
|
||||
acc.clear();
|
||||
}
|
||||
{
|
||||
let mut state = self.state.lock().unwrap();
|
||||
*state = SessionState::Running;
|
||||
}
|
||||
|
||||
let sid = self.session_id.lock().unwrap().clone();
|
||||
let session_id = sid.ok_or_else(|| {
|
||||
crate::acp_client::AcpError::Internal(
|
||||
"no session created yet — call create_session first".into(),
|
||||
)
|
||||
})?;
|
||||
|
||||
let params = SessionPromptParams {
|
||||
session_id: session_id.clone(),
|
||||
prompt,
|
||||
};
|
||||
|
||||
debug!("ACP session/prompt (session={})", session_id);
|
||||
let result: SessionPromptResult = self.client.request("session/prompt", params).await?;
|
||||
debug!(
|
||||
"ACP session/prompt done (session={}, stop_reason={:?})",
|
||||
session_id, result.stop_reason
|
||||
);
|
||||
|
||||
{
|
||||
let mut state = self.state.lock().unwrap();
|
||||
*state = SessionState::Idle;
|
||||
}
|
||||
{
|
||||
let mut in_prompt = self.in_prompt.lock().unwrap();
|
||||
*in_prompt = false;
|
||||
}
|
||||
{
|
||||
let mut acc = self.accumulated.lock().unwrap();
|
||||
acc.clear();
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Cancel the current prompt.
|
||||
///
|
||||
/// Sends `session/cancel` notification to the agent.
|
||||
pub async fn cancel(&self) -> Result<(), crate::acp_client::AcpError> {
|
||||
let sid = self.session_id.lock().unwrap().clone();
|
||||
let session_id = sid.ok_or_else(|| {
|
||||
crate::acp_client::AcpError::Internal("no active session".into())
|
||||
})?;
|
||||
|
||||
{
|
||||
let mut state = self.state.lock().unwrap();
|
||||
*state = SessionState::Cancelling;
|
||||
}
|
||||
|
||||
self.client
|
||||
.notification(
|
||||
"session/cancel",
|
||||
json!({ "sessionId": session_id }),
|
||||
)
|
||||
.await?;
|
||||
|
||||
info!("ACP session cancelled: {}", session_id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Close the session and the underlying ACP client.
|
||||
pub async fn close(&self) {
|
||||
{
|
||||
let mut state = self.state.lock().unwrap();
|
||||
*state = SessionState::Closed;
|
||||
}
|
||||
|
||||
if let Some(handler) = self.event_handler.lock().unwrap().take() {
|
||||
handler(AcpSessionEvent::Disconnected {
|
||||
reason: "session closed".into(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the current session ID, if any.
|
||||
pub async fn session_id(&self) -> Option<String> {
|
||||
self.session_id.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
/// Get the current session state.
|
||||
pub async fn state(&self) -> SessionState {
|
||||
self.state.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
// ── Internal: session/update → event mapping ─────
|
||||
|
||||
/// Convert a parsed [`SessionUpdate`] into an [`AcpSessionEvent`],
|
||||
/// applying text deduplication for streaming text.
|
||||
///
|
||||
/// Reference: `hermes-vscode-main/src/protocol.ts` `extractTextContent()`,
|
||||
/// `deduplicateChunk()`, `parseToolCall()`, `parseToolCallUpdate()`,
|
||||
/// `parseUsageUpdate()`, `parseSessionInfoUpdate()`
|
||||
fn session_update_to_event(
|
||||
update: &SessionUpdate,
|
||||
accumulated: &Arc<Mutex<String>>,
|
||||
is_in_prompt: bool,
|
||||
) -> Option<AcpSessionEvent> {
|
||||
match update {
|
||||
SessionUpdate::AgentMessageChunk { content, .. }
|
||||
| SessionUpdate::AgentThoughtChunk { content, .. } => {
|
||||
let discrim = match update {
|
||||
SessionUpdate::AgentMessageChunk { .. } => "msg",
|
||||
SessionUpdate::AgentThoughtChunk { .. } => "thought",
|
||||
_ => unreachable!(),
|
||||
};
|
||||
|
||||
// Deduplication (reference: protocol.ts deduplicateChunk)
|
||||
if is_in_prompt {
|
||||
let mut acc = accumulated.lock().unwrap();
|
||||
let text = &content.text;
|
||||
let event = if text == acc.as_str() {
|
||||
// Exact full resend → drop
|
||||
None
|
||||
} else if text.len() > 10 && text.starts_with(acc.as_str()) {
|
||||
// Superset resend → emit only the tail
|
||||
let new_part = text[acc.len()..].to_string();
|
||||
if new_part.is_empty() {
|
||||
None
|
||||
} else {
|
||||
*acc = text.clone();
|
||||
Some(new_part)
|
||||
}
|
||||
} else if text.len() > 10 && acc.ends_with(text) {
|
||||
// Partial resend → drop
|
||||
None
|
||||
} else {
|
||||
// Normal delta
|
||||
let new_acc = format!("{}{}", acc, text);
|
||||
*acc = new_acc;
|
||||
Some(text.clone())
|
||||
};
|
||||
|
||||
return event.map(|t| match discrim {
|
||||
"msg" => AcpSessionEvent::TextDelta { text: t },
|
||||
"thought" => AcpSessionEvent::ThoughtDelta { text: t },
|
||||
_ => unreachable!(),
|
||||
});
|
||||
}
|
||||
|
||||
// Not in prompt — emit directly (historical playback)
|
||||
let text = content.text.clone();
|
||||
Some(match discrim {
|
||||
"msg" => AcpSessionEvent::TextDelta { text },
|
||||
"thought" => AcpSessionEvent::ThoughtDelta { text },
|
||||
_ => unreachable!(),
|
||||
})
|
||||
}
|
||||
|
||||
SessionUpdate::ToolCall {
|
||||
tool_call_id,
|
||||
title,
|
||||
kind,
|
||||
status,
|
||||
..
|
||||
} => {
|
||||
let title = title.clone().unwrap_or_else(|| "tool".into());
|
||||
let kind_str = match kind {
|
||||
Some(k) => format!("{:?}", k).to_lowercase(),
|
||||
None => "other".into(),
|
||||
};
|
||||
let status = status.clone().unwrap_or(ToolCallStatus::Pending);
|
||||
Some(AcpSessionEvent::ToolCall {
|
||||
tool_call_id: tool_call_id.clone(),
|
||||
title,
|
||||
kind: kind_str,
|
||||
status,
|
||||
})
|
||||
}
|
||||
|
||||
SessionUpdate::ToolCallUpdate {
|
||||
tool_call_id,
|
||||
status,
|
||||
..
|
||||
} => {
|
||||
let status = status.clone().unwrap_or(ToolCallStatus::Completed);
|
||||
Some(AcpSessionEvent::ToolCallUpdate {
|
||||
tool_call_id: tool_call_id.clone(),
|
||||
status,
|
||||
})
|
||||
}
|
||||
|
||||
SessionUpdate::UsageUpdate { used, size, .. } => {
|
||||
Some(AcpSessionEvent::UsageUpdate {
|
||||
used: *used,
|
||||
size: *size,
|
||||
})
|
||||
}
|
||||
|
||||
SessionUpdate::SessionInfoUpdate { title, .. } => {
|
||||
Some(AcpSessionEvent::SessionInfoUpdate {
|
||||
title: title.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
SessionUpdate::Plan { entries, .. } => {
|
||||
let summaries: Vec<String> =
|
||||
entries.iter().map(|e| e.content.clone()).collect();
|
||||
Some(AcpSessionEvent::PlanUpdate { entries: summaries })
|
||||
}
|
||||
|
||||
SessionUpdate::Unknown { .. } => {
|
||||
warn!("ACP unknown session/update variant");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::acp_client::AcpClient;
|
||||
|
||||
/// Creates a minimal ACP mock server for testing session operations.
|
||||
async fn spawn_mock_acp() -> Arc<AcpClient> {
|
||||
let script = r#"
|
||||
import * as readline from 'node:readline';
|
||||
import { stdin as input, stdout as output } from 'node:process';
|
||||
const rl = readline.createInterface({ input, output, terminal: false });
|
||||
rl.on('line', (line) => {
|
||||
const msg = JSON.parse(line);
|
||||
if (!msg.id) return;
|
||||
if (msg.method === 'session/new') {
|
||||
process.stdout.write(JSON.stringify({
|
||||
jsonrpc: '2.0', id: msg.id,
|
||||
result: { sessionId: 'test_session_1' }
|
||||
}) + '\n');
|
||||
} else if (msg.method === 'session/prompt') {
|
||||
const sessionId = msg.params?.sessionId || 'test';
|
||||
process.stdout.write(JSON.stringify({
|
||||
jsonrpc: '2.0',
|
||||
method: 'session/update',
|
||||
params: {
|
||||
sessionId,
|
||||
update: {
|
||||
sessionUpdate: 'agent_message_chunk',
|
||||
content: { type: 'text', text: 'Hello from mock ACP' }
|
||||
}
|
||||
}
|
||||
}) + '\n');
|
||||
process.stdout.write(JSON.stringify({
|
||||
jsonrpc: '2.0', id: msg.id,
|
||||
result: { stopReason: 'end_turn' }
|
||||
}) + '\n');
|
||||
} else if (msg.method === 'session/cancel') {
|
||||
// No response for notification
|
||||
} else if (msg.method === 'initialize') {
|
||||
process.stdout.write(JSON.stringify({
|
||||
jsonrpc: '2.0', id: msg.id,
|
||||
result: { protocolVersion: 1, agentCapabilities: {}, agentInfo: { name: 'mock', version: '1.0' }, authMethods: [] }
|
||||
}) + '\n');
|
||||
}
|
||||
});
|
||||
"#;
|
||||
let dir = std::env::temp_dir();
|
||||
let script_path = dir.join("acp_session_test_mock.mjs");
|
||||
std::fs::write(&script_path, script).expect("write mock");
|
||||
let client = AcpClient::spawn("node", &[script_path.to_str().unwrap()])
|
||||
.await
|
||||
.expect("spawn");
|
||||
Arc::new(client)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_session() {
|
||||
let client = spawn_mock_acp().await;
|
||||
let mgr = AcpSessionManager::new(client);
|
||||
let sid = mgr
|
||||
.create_session(Some("/test"), None)
|
||||
.await
|
||||
.expect("create_session");
|
||||
assert_eq!(sid, "test_session_1");
|
||||
assert_eq!(mgr.session_id().await, Some("test_session_1".into()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_run_prompt() {
|
||||
let client = spawn_mock_acp().await;
|
||||
let mgr = AcpSessionManager::new(client);
|
||||
mgr.create_session(Some("/test"), None)
|
||||
.await
|
||||
.expect("create_session");
|
||||
|
||||
let prompt = vec![ContentBlock::Text {
|
||||
text: "Hello agent".into(),
|
||||
}];
|
||||
let result = mgr.run_prompt(prompt).await.expect("run_prompt");
|
||||
assert_eq!(
|
||||
format!("{:?}", result.stop_reason),
|
||||
"EndTurn".to_string()
|
||||
);
|
||||
// After prompt, state should be idle again
|
||||
assert_eq!(mgr.state().await, SessionState::Idle);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancel() {
|
||||
let client = spawn_mock_acp().await;
|
||||
let mgr = AcpSessionManager::new(client);
|
||||
mgr.create_session(Some("/test"), None)
|
||||
.await
|
||||
.expect("create_session");
|
||||
mgr.cancel().await.expect("cancel");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_event_handler() {
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
let client = spawn_mock_acp().await;
|
||||
let mgr = AcpSessionManager::new(client);
|
||||
let received = Arc::new(AtomicBool::new(false));
|
||||
let r = received.clone();
|
||||
|
||||
mgr.on_event(move |ev| {
|
||||
if matches!(ev, AcpSessionEvent::TextDelta { .. }) {
|
||||
r.store(true, Ordering::SeqCst);
|
||||
}
|
||||
});
|
||||
|
||||
mgr.create_session(Some("/test"), None)
|
||||
.await
|
||||
.expect("create_session");
|
||||
let prompt = vec![ContentBlock::Text {
|
||||
text: "Hello".into(),
|
||||
}];
|
||||
mgr.run_prompt(prompt).await.expect("run_prompt");
|
||||
|
||||
// Give the notification handler time to process
|
||||
sleep(Duration::from_millis(200)).await;
|
||||
assert!(received.load(Ordering::SeqCst), "should have received TextDelta");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user